From 2928483bd533ea5406d2d867cee0d2a8855bdfb4 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 5 Aug 2026 06:02:13 +0000 Subject: [PATCH] feat: warn before an unbatched MultiStartProdigy jvp exhausts memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With batch_size=None the search vmaps every start into one value_and_grad, and the only signal that this will not fit is XLA's RESOURCE_EXHAUSTED raised from inside the compiled program, naming a byte count and nothing else. Nothing points at batch_size, the one knob that fixes it. That cost two nightly release runs in 2026-07, after which two workspace scripts were hand-patched with magic numbers. Adds, before the full vmap is compiled: probe the batched jvp at widths 1 and 2 via XLA's own memory analysis, split the result into fixed and per-start components, and if the projected n_starts footprint exceeds available memory, warn naming a batch_size that would fit. Two probe points rather than one on purpose. Workspace guidance has claimed VRAM 'does not scale with batch size for the persistent buffers, so if it fits at batch_size=1 you can push it up' — true only when the per-start slope is small next to the fixed cost, and for an interferometer jvp it is not: the incident measured ~1.79 GB per start, so the single-start probe that 'fit' predicted nothing. Analysis.print_vram_use gains a flag and now names what it measured. It profiled fitness.call — the likelihood, not the jvp — so its number could not bound a gradient search's footprint, which is how a fit that looked modest reached ~86 GB. Warn only, never auto-apply: batch_size is numerically inert so auto-batching would be safe in principle, but silently changing the execution shape of every existing run on the strength of a projection is not a trade to make unattended. Known limitation, documented at the call site: memory_analysis reports 0 on a CPU-only JAX build, which is where the release harness runs, so the projection is skipped there rather than guessing. The CPU path relies on the now-unfiltered traceback instead. Making it measurable is follow-up work that needs a real run to validate. The projection arithmetic is pure Python and unit-tested here; the JAX measurement it consumes belongs to autofit_workspace_test, per this suite's existing JAX boundary. Refs PyAutoLabs/PyAutoFit#1452 (phase 3 of 3). Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0171voyyTrr91hJ3vU5AjeVz --- autofit/non_linear/analysis/analysis.py | 74 +++++++-- .../search/mle/multi_start_gradient/search.py | 146 ++++++++++++++++++ .../search/mle/test_multi_start_gradient.py | 74 ++++++++- 3 files changed, 278 insertions(+), 16 deletions(-) diff --git a/autofit/non_linear/analysis/analysis.py b/autofit/non_linear/analysis/analysis.py index 3a887c3db..a0527ded6 100644 --- a/autofit/non_linear/analysis/analysis.py +++ b/autofit/non_linear/analysis/analysis.py @@ -334,23 +334,35 @@ def supports_jax_visualization(self) -> bool: def perform_quick_update(self, paths, instance): raise NotImplementedError - def print_vram_use(self, model, batch_size : int) -> str: + def batched_memory_bytes(self, model, batch_size: int, gradient: bool = False): """ - Print JAX VRAM use for a given batch size. + Measure the memory a batched evaluation of this analysis would need. + + Compiles (but never runs) ``jax.vmap`` over ``batch_size`` copies of the + model at its prior medians and reads XLA's own memory analysis, so this + is a lowering-time estimate — it cannot itself OOM the way the run it is + predicting would. Parameters ---------- batch_size - The batch size to profile, which is the number of model evaluations JAX will perform simultaneously. - """ - from autofit.non_linear.test_mode import skip_fit_output - - if skip_fit_output(): - return + Number of model evaluations JAX would perform simultaneously. + gradient + Measure ``jax.value_and_grad`` of the likelihood rather than the + likelihood alone. **A gradient search must pass True.** The jvp + carries the whole forward tape, so a likelihood-only measurement + can under-report a gradient optimizer's footprint by a large + factor — that gap is exactly what let a 48-start interferometer fit + reach ~86 GB while the likelihood-only figure looked modest + (PyAutoFit#1452). + Returns + ------- + The projected bytes, or ``None`` when this analysis is not on the JAX + path (nothing to measure). + """ if not self._use_jax: - print("use_jax=False for this analysis, therefore does not use GPU and VRAM use cannot be profiled.") - return + return None import jax import jax.numpy as jnp @@ -372,16 +384,48 @@ def print_vram_use(self, model, batch_size : int) -> str: parameters = jnp.array(parameters) - batched_call = jax.jit(jax.vmap(fitness.call)) + call = jax.value_and_grad(fitness.call) if gradient else fitness.call + + batched_call = jax.jit(jax.vmap(call)) lowered = batched_call.lower(parameters) compiled = lowered.compile() memory_analysis = compiled.memory_analysis() - vram_bytes = ( - memory_analysis.output_size_in_bytes - + memory_analysis.temp_size_in_bytes + return ( + memory_analysis.output_size_in_bytes + + memory_analysis.temp_size_in_bytes + ) + + def print_vram_use(self, model, batch_size : int, gradient: bool = False) -> str: + """ + Print JAX VRAM use for a given batch size. + + Parameters + ---------- + batch_size + The batch size to profile, which is the number of model evaluations JAX will perform simultaneously. + gradient + Profile ``value_and_grad`` rather than the likelihood alone. Pass + True when the search is a gradient optimizer (e.g. + ``af.MultiStartProdigy``) — see ``batched_memory_bytes``. + """ + from autofit.non_linear.test_mode import skip_fit_output + + if skip_fit_output(): + return + + if not self._use_jax: + print("use_jax=False for this analysis, therefore does not use GPU and VRAM use cannot be profiled.") + return + + vram_bytes = self.batched_memory_bytes( + model=model, batch_size=batch_size, gradient=gradient ) + # Stated so the number is never read as covering more than it measures: + # the likelihood-only figure does not bound a gradient search. + measured = "likelihood + gradient" if gradient else "likelihood only" + if vram_bytes == 0: print( "VRAM USE = 0.000 GB " @@ -389,5 +433,5 @@ def print_vram_use(self, model, batch_size : int) -> str: ) else: print( - f"VRAM USE = {vram_bytes / 1024 ** 3:.3f} GB" + f"VRAM USE = {vram_bytes / 1024 ** 3:.3f} GB ({measured}, batch_size={batch_size})" ) \ No newline at end of file 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 1cb6d902f..aa0837439 100644 --- a/autofit/non_linear/search/mle/multi_start_gradient/search.py +++ b/autofit/non_linear/search/mle/multi_start_gradient/search.py @@ -353,6 +353,96 @@ def _progress_message( return " | ".join(parts) + def _memory_budget_bytes(self): + """ + Bytes available to the batched jvp, or ``None`` if not determinable. + + Prefers the JAX device's own limit (the GPU case this guard mainly + exists for) and falls back to free host RAM. + """ + try: + import jax + + stats = jax.devices()[0].memory_stats() or {} + for key in ("bytes_limit", "bytes_reservable_limit"): + if stats.get(key): + return int(stats[key]) + except Exception: + pass + + try: + import psutil + + return int(psutil.virtual_memory().available) + except Exception: + return None + + def _warn_if_unbatched_exceeds_memory(self, model, analysis): + """ + Warn, before compiling the full ``n_starts`` vmap, when its batched jvp + is projected to exceed available memory — naming the ``batch_size`` that + would fit. + + Without this the only signal is XLA's ``RESOURCE_EXHAUSTED: Out of + memory allocating N bytes``, raised from inside the compiled program + with nothing pointing at the knob that fixes it. That cost two nightly + release runs in 2026-07 (PyAutoFit#1452). + + This only ever *warns*. ``batch_size`` is numerically inert, so + auto-applying the suggestion would be safe in principle, but silently + changing the execution shape of every existing run on the strength of a + projection is not a trade this should make on its own. + + **Known limitation.** ``memory_analysis`` reports 0 on a CPU-only JAX + build, which is where the release harness runs. The projection is + skipped entirely in that case rather than guessing, so this guard + currently helps GPU users and leaves the CPU path to the (now + unfiltered) traceback. Making the CPU path measurable is follow-up + work and needs validating against a real run, not a unit test. + + Any failure here is swallowed: a memory projection must never be the + reason a fit does not start. + """ + try: + probe = getattr(analysis, "batched_memory_bytes", None) + if probe is None or self.n_starts is None or self.n_starts <= 1: + return + + bytes_at_1 = probe(model=model, batch_size=1, gradient=True) + if not bytes_at_1: # None (not on JAX) or 0 (CPU-only: unmeasurable) + return + + bytes_at_2 = probe(model=model, batch_size=2, gradient=True) + if not bytes_at_2: + return + + budget = self._memory_budget_bytes() + if not budget: + return + + fixed, per_start = batch_memory_model(bytes_at_1, bytes_at_2) + suggested = batch_size_within_budget( + fixed, per_start, self.n_starts, budget + ) + if suggested is None: + return + + projected = fixed + self.n_starts * per_start + gb = 1024 ** 3 + self.logger.warning( + f"{type(self).__name__} is set to evaluate all " + f"{self.n_starts} starts in one jax.vmap (batch_size=None). " + f"Its batched value_and_grad is projected to need " + f"~{projected / gb:.1f} GB against ~{budget / gb:.1f} GB " + f"available, so this fit is likely to fail with an XLA " + f"RESOURCE_EXHAUSTED error. Pass batch_size={suggested} (or " + f"lower) to sweep the starts in chunks instead — the result is " + f"numerically identical, only the allocation and dispatch " + f"change." + ) + except Exception: + return + def _fit( self, model: AbstractPriorModel, @@ -416,6 +506,7 @@ def _fit( _vmapped = jax.jit(jax.vmap(_value_and_grad)) if self.batch_size is None: + self._warn_if_unbatched_exceeds_memory(model=model, analysis=analysis) batched_value_and_grad = _vmapped else: batch_size = self.batch_size @@ -879,6 +970,61 @@ def samples_via_internal_from( ) +def batch_memory_model(bytes_at_1, bytes_at_2): + """ + Split a measured batched-jvp footprint into its fixed and per-start parts. + + Two probe points are enough because the footprint is affine in the batch + width: ``bytes(n) = fixed + n * per_start``. ``per_start`` is the slope + (what one extra start costs), ``fixed`` the intercept (buffers shared by + the whole batch, e.g. the dataset and any persistent operator). + + Measuring *both* is the point. Guidance in the workspaces has claimed VRAM + "does not scale with batch size for the persistent buffers, so if it fits + at ``batch_size=1`` you can push it up" — true only when ``per_start`` is + small next to ``fixed``. For an interferometer likelihood under + ``value_and_grad`` it is not: the 2026-07-30 release failure projected to + ~1.79 GB *per start*, so 48 starts wanted ~86 GB and the single-start + probe that "fit" said nothing useful (PyAutoFit#1452). + + Returns ``(fixed, per_start)``, both clamped at 0 — a non-monotonic or + noisy pair of measurements must never yield a negative slope and so a + nonsensically large batch suggestion. + """ + per_start = max(0, bytes_at_2 - bytes_at_1) + fixed = max(0, bytes_at_1 - per_start) + return fixed, per_start + + +def batch_size_within_budget(fixed, per_start, n_starts, budget_bytes): + """ + The largest batch width whose projected footprint fits ``budget_bytes``. + + Returns ``None`` when the full ``n_starts`` batch already fits — the + caller should then leave ``batch_size=None`` and keep today's single-vmap + fast path. Otherwise returns an int in ``[1, n_starts)``. + + A budget that cannot fit even one start still returns 1: tiling to a + single start is the most this knob can do, and a hard failure is the + search's to report, not this projection's to pre-empt. + """ + if budget_bytes <= 0 or n_starts <= 1: + return None + + def projected(n): + return fixed + n * per_start + + if projected(n_starts) <= budget_bytes: + return None + + if per_start <= 0: + # Fixed cost alone blows the budget; tiling cannot help, but a + # single start is still the smallest thing we can ask for. + return 1 + + return max(1, min(n_starts - 1, int((budget_bytes - fixed) // per_start))) + + def _chunk_slices(n_rows, batch_size): """ The ``(lo, hi, pad)`` chunk bounds the batched ``value_and_grad`` sweep 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 359aa931e..cb3803cc7 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,7 +6,11 @@ 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 autofit.non_linear.search.mle.multi_start_gradient.search import ( + _chunk_slices, + batch_memory_model, + batch_size_within_budget, +) from autonerves.dictable import from_dict, to_dict # The MultiStart gradient searches are JAX-native at fit time, but their @@ -691,3 +695,71 @@ def test__samples_info__reports_a_cleared_stop_reason_as_unfinished(): assert samples.samples_info["stop_reason"] is None assert samples.samples_info["converged"] is False + + +# --- Batched-jvp memory projection (PyAutoFit#1452) --------------------------- +# +# Same contract as the _chunk_slices tests above: the projection arithmetic is +# pure Python and lives here, while the JAX measurement it consumes +# (`Analysis.batched_memory_bytes`) is exercised in autofit_workspace_test. + +GB = 1024 ** 3 + +# The 2026-07-30 release failure: XLA reported 85,898,814,480 bytes for a +# 48-start interferometer jvp, i.e. 1,789,558,635 bytes per start. +INCIDENT_PER_START = 1_789_558_635 + + +def test__batch_memory_model__recovers_fixed_and_per_start(): + fixed = 2 * GB + at_1 = fixed + INCIDENT_PER_START + at_2 = fixed + 2 * INCIDENT_PER_START + + assert batch_memory_model(at_1, at_2) == (fixed, INCIDENT_PER_START) + + +@pytest.mark.parametrize( + "at_1, at_2", + [ + (10, 4), # non-monotonic measurements + (10, 10), # identical measurements + ], +) +def test__batch_memory_model__never_yields_a_negative_slope(at_1, at_2): + # A negative slope would invert the budget arithmetic and suggest an + # absurdly large batch, which is worse than saying nothing. + fixed, per_start = batch_memory_model(at_1, at_2) + assert per_start == 0 + assert fixed >= 0 + + +def test__batch_size_within_budget__suggests_the_largest_batch_that_fits(): + fixed, per_start = 2 * GB, INCIDENT_PER_START + + suggested = batch_size_within_budget(fixed, per_start, 48, 16 * GB) + + assert 1 <= suggested < 48 + assert fixed + suggested * per_start <= 16 * GB + assert fixed + (suggested + 1) * per_start > 16 * GB + + +def test__batch_size_within_budget__none_when_the_full_batch_fits(): + # None means "leave batch_size=None" — the single-vmap fast path is kept. + assert batch_size_within_budget(2 * GB, INCIDENT_PER_START, 48, 400 * GB) is None + + +@pytest.mark.parametrize( + "fixed, per_start, n_starts, budget, expected", + [ + (100 * GB, 0, 48, 1 * GB, 1), # fixed cost alone busts it; tiling cannot help + (1, 0, 48, 1 * GB, None), # no per-start cost and it fits + (0, 1, 48, 47, 47), # never returns n_starts itself + (0, 100 * GB, 48, 1 * GB, 1), # not even one start fits -> still 1 + (2 * GB, INCIDENT_PER_START, 48, 0, None), # unknown budget + (2 * GB, INCIDENT_PER_START, 1, 1, None), # nothing to tile + ], +) +def test__batch_size_within_budget__degenerate_inputs( + fixed, per_start, n_starts, budget, expected +): + assert batch_size_within_budget(fixed, per_start, n_starts, budget) == expected