merge Feat/corrfunc kdtree#43
Merged
Merged
Conversation
…ling
soft_pair_counts_from_topology materialised the full per-pair soft-bin
weight tensors at once -- w_cross of shape (P, max_leaf, max_leaf, nbins)
and w_within of shape (L, ...) -- so peak memory scaled as
P*max_leaf^2*nbins and OOM'd above N~1e4-2e4 (value_and_grad ~doubled it).
Reduce the near field in fixed-size chunks over the leaf-pair axis via a
rematerialised jax.lax.scan (new _accumulate_block_pairs helper),
accumulating into the (nbins,) output so the (chunk, ml, ml, nbins)
weight tensor is the only large live intermediate; memory is now
O(chunk_size*max_leaf^2*nbins). jax.checkpoint on the scan body is
essential -- without it the reverse pass would stack every chunk's
residual and restore the O(P) ceiling this removes.
chunk_size is a new keyword-only parameter (default 1024) on
soft_pair_counts_from_topology and soft_pair_counts; a single-chunk fast
path keeps small-N results bit-identical to the previous code.
build_pair_topology and the far-field monopole path are unchanged.
Validated: forward + jax.grad bit-identical at N=1000 (theta in
{0,0.3,0.6}); tests/applications all pass; on an A100-40GB the estimator
now runs at N=1e5 (acc 0.5s / vgrad 0.9s) and N=1e6 (acc 7.7s / vgrad
15.7s) with no OOM, and validation errors match the committed 40GB run
at the overlapping sizes (estimator counts within 6e-16 relative).
Bench harness (no estimator-code change): validate_vs_baseline gains
--brute-force-max-n so the O(N^2) reference is skipped where it would
itself OOM (above it the theta=0 tree run, whose near field is exact, is
used as the reference); scaling right-sizes the oversized traversal
capacity buffers (1<<15 -> 1<<13, which preallocated ~8.6GiB and OOM'd
the *build* at 1e6) and warms the build_s timing.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
perf(corrfunc): chunk the near-field reduction to lift the memory ceiling
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.
This pull request introduces significant improvements to the memory efficiency and scalability of the pair-counting estimator, especially for large particle counts. The main changes include chunked accumulation of near-field pairwise interactions to avoid out-of-memory (OOM) errors, configurable chunk sizing, and improved benchmarking and validation scripts to handle larger problem sizes and edge cases. It also adds clearer documentation and configuration options to both the estimator and its benchmarks.
Estimator memory scaling and chunking:
_accumulate_block_pairsto perform chunked, differentiable reduction of soft-bin weights over leaf pairs, dramatically reducing memory usage for large N by bounding intermediates toO(chunk_size * max_leaf^2 * nbins)instead ofO(P * max_leaf^2 * nbins).soft_pair_counts_from_topologyandsoft_pair_countsto use chunked accumulation, add a configurablechunk_sizeparameter, and propagate it through the call stack [1] [2] [3]._DEFAULT_NEAR_CHUNK(default 1024) as the chunk size for near-field reductions, with documentation on its effect on memory and speed.Benchmark and validation improvements:
bench/corrfunc/scaling.py, reduced preallocated capacities for far/near lists to prevent OOM during topology build, and improved timing methodology for topology construction [1] [2].bench/corrfunc/validate_vs_baseline.py, added--brute-force-max-nto switch to a tree-based exact solution as the reference when brute-force is infeasible, and restructured reference computation logic accordingly [1] [2] [3].max_pair_queuein validation to support larger N at tighttheta.Documentation and configuration:
These changes collectively make the estimator and its benchmarks robust to much larger problem sizes, prevent OOM errors, and provide clear configuration and diagnostics for users.