From 5c273e53ad7307e058705f33d583418beca66aa6 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 5 Aug 2026 06:21:38 +0000 Subject: [PATCH 1/3] docs: correct the VRAM batch-size guidance for interferometer fits The __VRAM Use__ section told users "VRAM does not scale with batch size for the persistent buffers, so if the analysis fits within VRAM for batch_size=1 you should be able to push the batch size up (e.g. to 50) to maximise GPU throughput without running out of memory." That is the advice that produces the failure it should prevent. Memory splits into a fixed part and a per-evaluation part, and which dominates is the whole question; for interferometer data under a gradient search the per-evaluation part dominates, because the batched value_and_grad carries the forward tape. A 48-start SDP.81 fit measured ~1.79 GB per start, wanted ~86 GB, and OOMed two nightly release runs - at almost exactly the batch size this text recommended - while a single start fit comfortably. Replaces it with the fixed-vs-per-start split, the measured number, a pointer to start_here.py's batch_size=4, and a note that print_vram_use profiles the likelihood only unless passed gradient=True (Nautilus, used in this script, takes no gradients, so the default is right here). Notebook regenerated: the generator copies each block verbatim and the markdown cell was byte-identical to the docstring before and after. Refs PyAutoLabs/PyAutoFit#1452. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0171voyyTrr91hJ3vU5AjeVz --- notebooks/interferometer/modeling.ipynb | 17 ++++++++++++++--- scripts/interferometer/modeling.py | 17 ++++++++++++++--- 2 files changed, 28 insertions(+), 6 deletions(-) diff --git a/notebooks/interferometer/modeling.ipynb b/notebooks/interferometer/modeling.ipynb index b3be2a5a4..8bfa78182 100644 --- a/notebooks/interferometer/modeling.ipynb +++ b/notebooks/interferometer/modeling.ipynb @@ -556,9 +556,20 @@ "it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits.\n", "\n", "With `TransformerNUFFT` (backed by `nufftax`), the dominant contributor to VRAM is usually the real-space image\n", - "and its transforms inside the likelihood function, rather than the visibility count itself. VRAM does not scale\n", - "with batch size for the persistent buffers, so if the analysis fits within VRAM for `batch_size=1` you should be\n", - "able to push the batch size up (e.g. to 50) to maximise GPU throughput without running out of memory.\n", + "and its transforms inside the likelihood function, rather than the visibility count itself.\n", + "\n", + "**A single-start measurement does not tell you a large batch will fit.** Memory splits into a fixed part (buffers\n", + "shared by the whole batch, e.g. the dataset) and a per-evaluation part that scales with `batch_size`. Whether you\n", + "can raise the batch size depends on which dominates, and for interferometer data under a *gradient* search the\n", + "per-evaluation part dominates: the batched `value_and_grad` carries the whole forward tape, so it is far larger\n", + "than the likelihood alone. A 48-start fit of the SDP.81 dataset measured ~1.79 GB *per start* \u2014 ~86 GB in total,\n", + "which exhausted the machine even though a single start fitted comfortably. This is why\n", + "`interferometer/start_here.py` passes `batch_size=4` rather than leaving it at `None`.\n", + "\n", + "So size the batch against the per-start slope, not against a single measurement. Note that `print_vram_use`\n", + "profiles the likelihood **only** by default; pass `gradient=True` to profile the `value_and_grad` a gradient\n", + "search (e.g. `af.MultiStartProdigy`) actually evaluates. `Nautilus`, used in this script, does not take\n", + "gradients, so the default is the right figure here.\n", "\n", "For an MGE model with the small dataset fitted in this example, VRAM use is modest (~0.3 GB). Larger real-space\n", "masks (finer pixel scales) and higher visibility counts increase VRAM gradually rather than catastrophically, and\n", diff --git a/scripts/interferometer/modeling.py b/scripts/interferometer/modeling.py index a315811b0..56add35af 100644 --- a/scripts/interferometer/modeling.py +++ b/scripts/interferometer/modeling.py @@ -392,9 +392,20 @@ it takes about 20-30 seconds to run so you may want to comment it out once you are familiar with your GPU's VRAM limits. With `TransformerNUFFT` (backed by `nufftax`), the dominant contributor to VRAM is usually the real-space image -and its transforms inside the likelihood function, rather than the visibility count itself. VRAM does not scale -with batch size for the persistent buffers, so if the analysis fits within VRAM for `batch_size=1` you should be -able to push the batch size up (e.g. to 50) to maximise GPU throughput without running out of memory. +and its transforms inside the likelihood function, rather than the visibility count itself. + +**A single-start measurement does not tell you a large batch will fit.** Memory splits into a fixed part (buffers +shared by the whole batch, e.g. the dataset) and a per-evaluation part that scales with `batch_size`. Whether you +can raise the batch size depends on which dominates, and for interferometer data under a *gradient* search the +per-evaluation part dominates: the batched `value_and_grad` carries the whole forward tape, so it is far larger +than the likelihood alone. A 48-start fit of the SDP.81 dataset measured ~1.79 GB *per start* — ~86 GB in total, +which exhausted the machine even though a single start fitted comfortably. This is why +`interferometer/start_here.py` passes `batch_size=4` rather than leaving it at `None`. + +So size the batch against the per-start slope, not against a single measurement. Note that `print_vram_use` +profiles the likelihood **only** by default; pass `gradient=True` to profile the `value_and_grad` a gradient +search (e.g. `af.MultiStartProdigy`) actually evaluates. `Nautilus`, used in this script, does not take +gradients, so the default is the right figure here. For an MGE model with the small dataset fitted in this example, VRAM use is modest (~0.3 GB). Larger real-space masks (finer pixel scales) and higher visibility counts increase VRAM gradually rather than catastrophically, and From 2b712708ccdb2749c62eb5dcf10ea9df4774f27a Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 5 Aug 2026 06:27:04 +0000 Subject: [PATCH 2/3] docs: refresh navigator catalogue for the VRAM prose change The corrected __VRAM Use__ text names interferometer/start_here.py as the script that passes batch_size=4, which the navigator generator extracts as a cross-reference. Adds that entry so the catalogue matches the scripts again. Exactly the one line the Catalogue staleness check asked for; llms-full.txt does not carry this prose and is unchanged. Refs PyAutoLabs/PyAutoFit#1452. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0171voyyTrr91hJ3vU5AjeVz --- workspace_index.json | 1 + 1 file changed, 1 insertion(+) diff --git a/workspace_index.json b/workspace_index.json index 75a2400ab..b39873621 100644 --- a/workspace_index.json +++ b/workspace_index.json @@ -5554,6 +5554,7 @@ ], "cross_refs": [ "/interferometer/data_preparation.ipynb", + "interferometer/start_here.py", "start_here.py" ], "notebook": "notebooks/interferometer/modeling.ipynb", From b860467097341e46491d38ebc152418cd995f924 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 5 Aug 2026 06:42:19 +0000 Subject: [PATCH 3/3] docs: make the multi-start batch_size decision explicit A new static check in PyAutoHands flags any MultiStart* search built without an explicit batch_size, because defaulting to None is how both interferometer/start_here.py files inherited an unbounded 48-way vmap on the same day and OOMed two nightly release runs. These three sites keep their current behaviour exactly - None is already the default, so nothing changes at runtime. What changes is that the choice is now written down where a reviewer sees it, with a pointer to the knob if a reader hits a memory limit. Deliberately not given a bounded value: these are CCD imaging, group and 1D-toy likelihoods rather than interferometer ones, and picking a real number needs a measurement on hardware this could not make. The imaging start_here.py entry points are the ones most worth profiling later, since a user on a laptop GPU meets them first. Refs PyAutoLabs/PyAutoFit#1452. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_0171voyyTrr91hJ3vU5AjeVz --- notebooks/guides/modeling/searches.ipynb | 1 + notebooks/imaging/start_here.ipynb | 1 + notebooks/multi_galaxy/start_here.ipynb | 1 + scripts/guides/modeling/searches.py | 1 + scripts/imaging/start_here.py | 1 + scripts/multi_galaxy/start_here.py | 1 + 6 files changed, 6 insertions(+) diff --git a/notebooks/guides/modeling/searches.ipynb b/notebooks/guides/modeling/searches.ipynb index 956a93660..3d344043f 100644 --- a/notebooks/guides/modeling/searches.ipynb +++ b/notebooks/guides/modeling/searches.ipynb @@ -236,6 +236,7 @@ " path_prefix=Path(\"imaging\", \"searches\"),\n", " name=\"MultiStartProdigy\",\n", " n_starts=50,\n", + " batch_size=None, # Starts evaluated at once: `None` vmaps all 50 together, which is fastest but allocates the whole batched gradient; set an integer (e.g. 4) if you hit an out-of-memory error.\n", " n_steps=500,\n", ")" ], diff --git a/notebooks/imaging/start_here.ipynb b/notebooks/imaging/start_here.ipynb index 2e0c6870b..b3a8f3981 100644 --- a/notebooks/imaging/start_here.ipynb +++ b/notebooks/imaging/start_here.ipynb @@ -402,6 +402,7 @@ " name=\"start_here\", # The name of the fit and folder results are output to.\n", " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", " n_starts=48, # The number of independent optimizations run in parallel, increase for more complex models.\n", + " batch_size=None, # Starts evaluated at once: `None` vmaps all 48 together, which is fastest but allocates the whole batched gradient; set an integer (e.g. 4) if you hit an out-of-memory error.\n", " n_steps=300, # The maximum gradient steps per start; the search stops early once the best fit stops improving.\n", " iterations_per_quick_update=50, # Every N steps the max likelihood model is visualized and output to hard-disk.\n", " live_visual_update=True, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", diff --git a/notebooks/multi_galaxy/start_here.ipynb b/notebooks/multi_galaxy/start_here.ipynb index fe27b7162..161061573 100644 --- a/notebooks/multi_galaxy/start_here.ipynb +++ b/notebooks/multi_galaxy/start_here.ipynb @@ -475,6 +475,7 @@ " name=\"start_here\", # The name of the fit and folder results are output to.\n", " unique_tag=dataset_name, # A unique tag which also defines the folder.\n", " n_starts=48, # The number of independent optimizations run in parallel, increase for more complex models.\n", + " batch_size=None, # Starts evaluated at once: `None` vmaps all 48 together, which is fastest but allocates the whole batched gradient; set an integer (e.g. 4) if you hit an out-of-memory error.\n", " n_steps=300, # The maximum gradient steps per start; the search stops early once the best fit stops improving.\n", " iterations_per_quick_update=50, # Every N steps the max likelihood model is visualized and output to hard-disk.\n", " live_visual_update=True, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).\n", diff --git a/scripts/guides/modeling/searches.py b/scripts/guides/modeling/searches.py index 5a4fe5e49..47f3949ca 100644 --- a/scripts/guides/modeling/searches.py +++ b/scripts/guides/modeling/searches.py @@ -166,6 +166,7 @@ path_prefix=Path("imaging", "searches"), name="MultiStartProdigy", n_starts=50, + batch_size=None, # Starts evaluated at once: `None` vmaps all 50 together, which is fastest but allocates the whole batched gradient; set an integer (e.g. 4) if you hit an out-of-memory error. n_steps=500, ) diff --git a/scripts/imaging/start_here.py b/scripts/imaging/start_here.py index fe295fa95..446050216 100644 --- a/scripts/imaging/start_here.py +++ b/scripts/imaging/start_here.py @@ -315,6 +315,7 @@ name="start_here", # The name of the fit and folder results are output to. unique_tag=dataset_name, # A unique tag which also defines the folder. n_starts=48, # The number of independent optimizations run in parallel, increase for more complex models. + batch_size=None, # Starts evaluated at once: `None` vmaps all 48 together, which is fastest but allocates the whole batched gradient; set an integer (e.g. 4) if you hit an out-of-memory error. n_steps=300, # The maximum gradient steps per start; the search stops early once the best fit stops improving. iterations_per_quick_update=50, # Every N steps the max likelihood model is visualized and output to hard-disk. live_visual_update=True, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook). diff --git a/scripts/multi_galaxy/start_here.py b/scripts/multi_galaxy/start_here.py index c01e1c3bf..d06083273 100644 --- a/scripts/multi_galaxy/start_here.py +++ b/scripts/multi_galaxy/start_here.py @@ -377,6 +377,7 @@ name="start_here", # The name of the fit and folder results are output to. unique_tag=dataset_name, # A unique tag which also defines the folder. n_starts=48, # The number of independent optimizations run in parallel, increase for more complex models. + batch_size=None, # Starts evaluated at once: `None` vmaps all 48 together, which is fastest but allocates the whole batched gradient; set an integer (e.g. 4) if you hit an out-of-memory error. n_steps=300, # The maximum gradient steps per start; the search stops early once the best fit stops improving. iterations_per_quick_update=50, # Every N steps the max likelihood model is visualized and output to hard-disk. live_visual_update=True, # Set True to open a live matplotlib window (script) or refresh a Jupyter cell (notebook).