Skip to content

feat(search): add batch_size to the multi-start gradient searches (vmap OOMs on memory-heavy likelihoods) #1373

Description

@Jammy2211

Problem

AbstractMultiStartGradient (PyAutoFit#1369/#1370) vmaps every start at once:

batched_value_and_grad = jax.jit(jax.vmap(jax.value_and_grad(fitness.call)))

For memory-heavy likelihoods this OOMs. Measured on an A100 80GB with a pixelized (kernel-CDF mesh) lens likelihood at n_starts=16, float64: the vmapped jvp fusion is f64[16, 15361, 2, 31, 512] = 58.13 GiBRESOURCE_EXHAUSTED. There is currently no knob to reduce it — the only workarounds are dropping to float32 (a science compromise) or cutting n_starts (defeats multi-start).

Prior A100 work recorded the same class of failure ('OOMs on pix/delaunay via vmap').

Fix

Add batch_size: Optional[int] = None to AbstractMultiStartGradient (inherited by MultiStartAdam / MultiStartADABelief / MultiStartLion): chunk the vmapped value_and_grad over starts in batches of batch_size, concatenating results. None = current all-at-once behaviour (no regression). The optax update still runs on the full stacked array — only the memory-dominant jvp is chunked.

Precedent: af.Nautilus already exposes n_batch, and autolens_profiling maintains a per-(dataset, model, instrument) vmap batch_size table for the A100 (vram/config.py; e.g. imaging/delaunay at HST scale needs batch=16, not the default 100). This makes that table consumable by the gradient searches too.

Test plan

  • numpy-only unit tests: batch_size plumbing/defaults, dict round-trip, and that chunked == unchunked on a simple objective.
  • JAX end-to-end validation in autofit_workspace_test.

Blocks the pixelized multi-start experiment (autolens_workspace_developer#100).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions