From 01fc1726e90ad00d7b45f25ffed7fb690aa9473d Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Wed, 15 Jul 2026 12:32:24 +0100 Subject: [PATCH] feat(search): add batch_size to the multi-start gradient searches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The multi-start searches vmap every start's value_and_grad at once. For a memory-heavy likelihood that OOMs: measured on an A100 80GB with a pixelized (kernel-CDF) lens likelihood at n_starts=16 in float64, the vmapped jvp fusion is f64[16,15361,2,31,512] = 58.13 GiB -> RESOURCE_EXHAUSTED, with no knob to reduce it (only fp32, a science compromise, or fewer starts, which defeats multi-start). batch_size=None (default) keeps the single vmap, so there is no regression. When set, the tiling is handed to jax.lax.map(..., batch_size=), which vmaps within a chunk and scans across chunks, handling a ragged final chunk without a second compile. lax.map is used ONLY when batching is requested — with batch_size=None it degrades to a sequential scan, which would throw away the default path's parallelism. batch_size is numerically inert: verified through _fit that best_fom and best_params are identical for batch_size in {None,1,4,14,100} (incl. ragged 14%4==2 and batch>n_starts). Unlike af.Nautilus's n_batch — Nautilus's own algorithmic knob that autofit forwards — this is purely implementation-level tiling; n_starts remains the algorithm. Resolves #1373. Co-Authored-By: Claude Opus 4.8 --- .../search/mle/multi_start_gradient/search.py | 44 ++++++++++++++++++- .../search/mle/test_multi_start_gradient.py | 18 ++++++++ 2 files changed, 60 insertions(+), 2 deletions(-) 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 88aa7b256..c999d64a4 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -30,6 +30,7 @@ def __init__( n_starts: int = 48, n_steps: int = 300, learning_rate: Optional[float] = None, + batch_size: Optional[int] = None, start_lower_limit: float = 0.15, start_upper_limit: float = 0.85, initializer: Optional[AbstractInitializer] = None, @@ -65,6 +66,20 @@ def __init__( learning_rate The optax learning rate. If ``None``, the rule's default is used (Adam / ADABelief ``1e-2``; the sign-based Lion ``1e-3``). + 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. + + 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 + 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. start_lower_limit, start_upper_limit The unit-cube bounds broad starts are drawn uniformly from. The interior default ``(0.15, 0.85)`` avoids the prior edges where many @@ -85,6 +100,7 @@ def __init__( self.n_starts = n_starts self.n_steps = n_steps + self.batch_size = batch_size self.learning_rate = ( learning_rate if learning_rate is not None else self._default_learning_rate ) @@ -134,8 +150,31 @@ def _fit( convert_to_chi_squared=True, ) - # -2 * log_posterior, to MINIMIZE. value_and_grad batched over starts. - batched_value_and_grad = jax.jit(jax.vmap(jax.value_and_grad(fitness.call))) + # -2 * log_posterior, to MINIMIZE. value_and_grad over every start. + # + # 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. + _value_and_grad = jax.value_and_grad(fitness.call) + _vmapped = jax.jit(jax.vmap(_value_and_grad)) + + if self.batch_size is None: + batched_value_and_grad = _vmapped + 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) try: search_internal = self.paths.load_search_internal() @@ -302,6 +341,7 @@ def samples_via_internal_from( samples_info = { "n_starts": self.n_starts, "n_steps": self.n_steps, + "batch_size": self.batch_size, "total_steps": total_steps, "optax_method": self.optax_method, "learning_rate": self.learning_rate, 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 4269d500e..b4a37325a 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 @@ -19,6 +19,7 @@ def test__explicit_params(): n_starts=16, n_steps=100, learning_rate=0.05, + batch_size=4, start_lower_limit=0.2, start_upper_limit=0.8, initializer=af.InitializerBall(lower_limit=0.2, upper_limit=0.8), @@ -27,6 +28,7 @@ def test__explicit_params(): assert search.n_starts == 16 assert search.n_steps == 100 assert search.learning_rate == 0.05 + assert search.batch_size == 4 assert search.start_lower_limit == 0.2 assert search.start_upper_limit == 0.8 assert search.optax_method == "adam" @@ -51,6 +53,22 @@ def test__per_rule_defaults(): assert default.n_steps == 300 assert default.start_lower_limit == 0.15 assert default.start_upper_limit == 0.85 + # batch_size defaults to None = evaluate all starts in one vmapped call + # (the pre-batch_size behaviour, so no regression). + assert default.batch_size is None + + +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. + """ + for cls in (af.MultiStartAdam, af.MultiStartADABelief, af.MultiStartLion): + assert cls().batch_size is None + assert cls(batch_size=8).batch_size == 8 def test__dict_round_trip():