feat: wire af.NSS chunk_size kwarg into build_nss (PyAutoFit#1303)#43
Merged
Conversation
Now that PyAutoFit#1303 ships an additive chunk_size knob on af.NSS, switch build_nss to use it: keep num_delete at the SLaM default (50) for convergence, and pass chunk_size=vmap_batch_for_cell(...) to cap the inner-vmap fan-out at the A100-probed budget per cell. Reverts the earlier num_delete=min(default, probe) Band-Aid, which shrank the sampler's per-iteration batch to preserve memory and hurt convergence. With the new chunked path the sampler keeps its preferred num_delete=50 while peak GPU memory becomes chunk_size × per_particle_state instead of num_delete × ... _runner.py:_sampler_config_dict records chunk_size in the metric JSON so each run's sampler_config block faithfully reflects what was constructed. Unblocks the A100 NSS Delaunay + pixelization cells that were OOMing at ~28 GB on every retry (jobs 322592/96/600/602/604 — see #1301 for full diagnosis). Refs PyAutoFit#1301, PyAutoFit#1303 Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Wires the new
chunk_sizekwarg shipped by PyAutoFit#1303 intoautolens_profiling/searches/build_nss. Reverts the earliernum_delete=min(default, probe)Band-Aid that shrank the sampler's per-iteration batch to dodge OOMs; with PyAutoFit#1303's chunked-vmap path the sampler now keeps its preferrednum_delete=50whilechunk_size = vmap_batch_for_cell(...)caps the inner-vmap fan-out per the A100-probed budget.Unblocks the A100 NSS Delaunay + pixelization cells that were OOMing at ~28 GB on every retry (jobs 322592/96/600/602/604).
Scripts Changed
searches/_samplers.py—build_nss(...)now passeschunk_size=vmap_batch_for_cell(...)toaf.NSS, keepsnum_delete=50(default). Comment block explains the memory-budget split.searches/_runner.py—_sampler_config_dictnss branch recordschunk_sizein the metric JSON and reverts the num_delete cap so the JSON shows what was actually constructed.Upstream PR
PyAutoLabs/PyAutoFit#1303 —
feat(nss): chunk_size kwarg for inversion-heavy A100 likelihoodsTest Plan
AUTOLENS_PROFILING_SMOKE=1under the worktree-pinned PyAutoFit.sweep.py --dry-run --only nss/imaging/mgedispatches correctly with the new sampler_config shape.log_Zbetweenchunk_size=Noneandchunk_size=2on the same seed.🤖 Generated with Claude Code