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