diff --git a/autofit/non_linear/search/mle/multi_start_gradient/search.py b/autofit/non_linear/search/mle/multi_start_gradient/search.py index cb392e38c..31e84d405 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -85,16 +85,24 @@ def __init__( ell_comps / shear at exactly 0); it does not rescue landscapes with broad non-finite regions (that is the Phase-2 restart-on-death layer). batch_size - The number of starts evaluated per vmapped ``value_and_grad`` call, - via ``jax.lax.map``. ``None`` (default) evaluates all ``n_starts`` in - a single ``jax.vmap`` — fastest, but it allocates the whole batched - jvp at once, which for a memory-heavy likelihood (e.g. a pixelized - source at 16 starts, ~58 GB in float64) exhausts even an 80 GB GPU. - Setting it trades a little speed for a bounded memory footprint. + The number of starts evaluated per compiled ``value_and_grad`` + call. ``None`` (default) evaluates all ``n_starts`` in a single + ``jax.vmap`` — fastest, but it allocates the whole batched jvp at + once, which for a memory-heavy likelihood (e.g. a pixelized source + at 16 starts, ~58 GB in float64) exhausts even an 80 GB GPU. + Setting it sweeps the starts in ``batch_size``-wide vmapped chunks + from a Python loop, bounding both the memory footprint (one + chunk's jvp allocated at a time) and the XLA compile (one + chunk-shaped program, ever). The compile bound is why the sweep is + a Python loop and not an in-XLA scan: ``jax.lax.map`` welds the + whole sweep into one program, which for a multi-band + ``FactorGraphModel`` objective is intractable to compile (>1 hour + cold on CPU), while the same-width chunk alone compiles in + minutes. This is purely an **implementation-level tiling**: it is numerically - inert (identical results, only the allocation changes). That makes it - unlike ``af.Nautilus``'s ``n_batch``, which is Nautilus's own + inert (identical results, only the allocation and dispatch change). + That makes it unlike ``af.Nautilus``'s ``n_batch``, which is Nautilus's own algorithmic knob (how many points it proposes per iteration) that autofit merely forwards. Here ``n_starts`` is the algorithm; this only decides how many of those starts are evaluated at a time. @@ -237,15 +245,18 @@ def _fit( # Unbatched we vmap all starts at once — fastest, but it allocates the # whole batched jvp, which for a memory-heavy likelihood (e.g. a # pixelized source at 16 starts, ~58 GB in float64) exhausts even an - # 80 GB GPU. When `batch_size` is set we hand the tiling to - # `jax.lax.map`, which vmaps *within* each chunk and scans across them, - # handling a ragged final chunk without a second compile. It is - # numerically identical to the vmap; `batch_size` never changes results, - # it only bounds memory. - # - # `lax.map` is only used when batching is requested: with - # `batch_size=None` it degrades to a sequential scan, which would throw - # away the parallelism of the default path. + # 80 GB GPU. When `batch_size` is set we sweep the starts in + # `batch_size`-wide vmapped chunks from a Python loop. Only one + # chunk-shaped function is ever XLA-compiled (the jit cache keys on the + # `(batch_size, ndim)` shape; the ragged final chunk is padded to that + # shape and the padded rows discarded), so the compile cost is that of + # a single chunk regardless of `n_starts`. An in-XLA scan over the + # chunks (`jax.lax.map`) instead welds the whole sweep into one + # program, which for a multi-band `FactorGraphModel` objective is + # intractable to compile (>1 hour cold on CPU, and memory-explosive to + # compile) while the chunk alone compiles in minutes. The tiling is + # numerically identical to the vmap; `batch_size` never changes + # results, it only bounds memory and compile. _value_and_grad = jax.value_and_grad(fitness.call) _vmapped = jax.jit(jax.vmap(_value_and_grad)) @@ -254,9 +265,25 @@ def _fit( else: batch_size = self.batch_size - @jax.jit def batched_value_and_grad(params): - return jax.lax.map(_value_and_grad, params, batch_size=batch_size) + foms_chunks = [] + grads_chunks = [] + for lo, hi, pad in _chunk_slices(params.shape[0], batch_size): + chunk = params[lo:hi] + if pad: + chunk = jnp.concatenate( + [chunk, jnp.tile(chunk[-1:], (pad, 1))] + ) + foms, grads = _vmapped(chunk) + if pad: + foms = foms[:-pad] + grads = grads[:-pad] + foms_chunks.append(foms) + grads_chunks.append(grads) + return ( + jnp.concatenate(foms_chunks), + jnp.concatenate(grads_chunks), + ) # The optax rule (resolved from optax / optax.contrib), guarded by # apply_if_finite, with a jitted per-start (vmapped) update step. Built @@ -284,8 +311,7 @@ def batched_value_and_grad(params): params = self._broad_starts( model=model, - fitness=fitness, - batched_value_and_grad=batched_value_and_grad, + value_and_grad_single=jax.jit(_value_and_grad), jnp=jnp, ) @@ -527,12 +553,18 @@ def merge(old, fresh): return params, opt_state - def _broad_starts(self, model, fitness, batched_value_and_grad, jnp): + def _broad_starts(self, model, value_and_grad_single, jnp): """ Draw ``n_starts`` broad starting points in the unit cube, map them to physical parameters, and keep only those with a finite objective and a finite gradient (degenerate points such as ell_comps / shear at exactly 0 have NaN gradients and must be filtered out). + + ``value_and_grad_single`` is the jitted single-point objective built + once in ``_fit``: one XLA compile (persistently cached) serves every + draw. Evaluating the draws eagerly instead re-pays the un-compiled AD + cost per draw, which on a multi-band ``FactorGraphModel`` objective + dominated the whole fit (~13 minutes for 16 draws, cache or no cache). """ rng = np.random.default_rng(0) @@ -547,7 +579,7 @@ def _broad_starts(self, model, fitness, batched_value_and_grad, jnp): vector = jnp.asarray( model.vector_from_unit_vector(unit_vector=list(unit_vector), xp=jnp) ) - fom, grad = jax_value_and_grad_single(fitness, vector) + fom, grad = value_and_grad_single(vector) if np.isfinite(float(fom)) and np.all(np.isfinite(np.asarray(grad))): starts.append(vector) @@ -645,15 +677,20 @@ def samples_via_internal_from( ) -def jax_value_and_grad_single(fitness, vector): +def _chunk_slices(n_rows, batch_size): """ - Single-point ``value_and_grad`` of the fitness objective, used to filter - broad starts down to those with a finite value and gradient before the - batched loop begins. + The ``(lo, hi, pad)`` chunk bounds the batched ``value_and_grad`` sweep + iterates over: rows ``lo:hi`` of the params array, padded by ``pad`` repeats + of the final row so every chunk presents the same ``(batch_size, ndim)`` + shape to the compiled function (one XLA compile for the whole sweep). Only + the final chunk can be ragged (``pad > 0``) — the broad-start collection may + return fewer than ``n_starts`` rows, so raggedness is not restricted to + ``n_starts % batch_size``. """ - import jax - - return jax.value_and_grad(fitness.call)(vector) + return [ + (lo, min(lo + batch_size, n_rows), max(0, lo + batch_size - n_rows)) + for lo in range(0, n_rows, batch_size) + ] class MultiStartAdam(AbstractMultiStartGradient): diff --git a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py index 460890171..f2d8e98d7 100644 --- a/test_autofit/non_linear/search/mle/test_multi_start_gradient.py +++ b/test_autofit/non_linear/search/mle/test_multi_start_gradient.py @@ -6,6 +6,7 @@ import autofit as af from autofit import example from autofit.non_linear.search import abstract_search +from autofit.non_linear.search.mle.multi_start_gradient.search import _chunk_slices from autonerves.dictable import from_dict, to_dict # The MultiStart gradient searches are JAX-native at fit time, but their @@ -103,9 +104,9 @@ def test__batch_size_is_carried_to_every_rule(): """``batch_size`` is a shared knob on the abstract base, not per-rule. The numerical guarantee it must honour — chunked evaluation is identical to - the unchunked vmap — is a JAX property of ``jax.lax.map(..., batch_size=)`` - and is asserted in autofit_workspace_test, since the library suite is - NumPy-only. + the unchunked vmap — is asserted in autofit_workspace_test, since the + library suite is NumPy-only. The chunk bookkeeping the sweep is built on + (``_chunk_slices``) is pure Python and tested below. """ for cls in ( af.MultiStartAdam, @@ -117,6 +118,38 @@ def test__batch_size_is_carried_to_every_rule(): assert cls(batch_size=8).batch_size == 8 +@pytest.mark.parametrize( + "n_rows, batch_size, expected", + [ + (8, 4, [(0, 4, 0), (4, 8, 0)]), # exact division + (10, 4, [(0, 4, 0), (4, 8, 0), (8, 10, 2)]), # ragged final chunk + (3, 8, [(0, 3, 5)]), # fewer rows than one chunk (short broad-start draw) + (3, 1, [(0, 1, 0), (1, 2, 0), (2, 3, 0)]), # batch_size=1 + (4, 4, [(0, 4, 0)]), # single exact chunk + ], +) +def test__chunk_slices(n_rows, batch_size, expected): + assert _chunk_slices(n_rows, batch_size) == expected + + +@pytest.mark.parametrize("n_rows", [1, 3, 4, 7, 16, 47, 48]) +@pytest.mark.parametrize("batch_size", [1, 3, 4, 16]) +def test__chunk_slices__covers_every_row_at_constant_shape(n_rows, batch_size): + """Every chunk presents the same ``batch_size`` shape to the compiled + function (rows + pad), chunks tile ``0..n_rows`` in order, and only the + final chunk may be padded — the invariants the one-compile sweep rests on. + """ + slices = _chunk_slices(n_rows, batch_size) + + assert slices[0][0] == 0 + assert slices[-1][1] == n_rows + for i, (lo, hi, pad) in enumerate(slices): + assert (hi - lo) + pad == batch_size + assert pad == 0 or i == len(slices) - 1 + if i: + assert lo == slices[i - 1][1] + + def test__convergence_default_is_on_and_carried(): """Auto-convergence is a shared base-class setting, on by default (so users do not hand-tune ``n_steps``), and a custom settings object is carried through."""