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

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions