docs: correct the VRAM batch-size guidance + make batch_size explicit

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",
     ")"
    ],

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",

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",

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",

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,
 )
 

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).

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

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).

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",