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/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/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/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 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). 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",