Phase 5c + 5d: fused Pallas M2L/near-field + multi-GPU scaling in the distributed FMM#47
Merged
Conversation
Add DistributedFMMConfig.nearfield_backend {auto,pallas,baseline}. On Ampere+
(auto default) route the combined [local;halo] near-field P2P through the fused
leafpair Pallas kernel (the single-GPU fast-lane near-field) instead of the
pure-JAX baseline: intra-leaf self block via
_compute_leaf_p2p_prepared_large_n_self_only_impl + cross-leaf neighbour pairs
via _radix_fast_lane_prepacked_pallas over the combined CSR densified to a padded
[u_leaves, S_near] source-leaf-id table. Baseline path unchanged behind the knob.
Validated on A100 (ndev=2, real+dehnen default): near-field parity vs baseline
P2P = 2.1e-7 (CPU interpret de-risk 2e-16), aggL2 vs direct = 1.78e-4 @ per=8000.
Build-once steady-state whole-eval 10.7 s -> 43.5 ms (~245x) at per=8000/leaf=128
-- the baseline near-field P2P, NOT the dual-tree traversal, was the dominant
cost. jit=True stable with both backends.
The fused real M2L Pallas (JACCPOT_STATIC_STRICT_FUSED_M2L_PALLAS=1) is already
engaged via _apply_real_m2l for both self- and cross-M2L; parity-neutral (4.4e-9),
negligible at this near-field-dominated config.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Near-field Pallas leafpair backend is the ~245x win (10.7s->43.5ms @8k, ndev=2); the baseline near-field P2P -- not the traversal -- was the ~10s cost. jit=True stable; weak-scales flat 2->5 GPUs but default caps overflow at ndev>=4 (5d needs cap calibration first). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
… prep) The cross-domain LET grows with device count, so the fixed default traversal caps overflow at ndev>=4 (measured: cross_max_pair_queue overflow at per=8000). Add auto_scale_caps/cap_scale_factor/max_cap_retries to distributed_fmm_accelerations: on a traversal-buffer overflow, scale the config caps via with_scaled_caps and rebuild/re-run until overflow-free; the returned result carries the config that produced the forces + diagnostics["cap_retries"]. Factor the overflow reduction into _reduce_overflow. Verified (weak scaling, per=8000/GPU): ndev 4/5 need cap_factor=2 to run overflow-free (throughput still rises 5.4e5->6.2e5 part/s). Test test_driver_auto_scale_caps: tiny caps overflow without retry, one x2 retry clears it, forces match direct (aggL2 1.9e-3, CPU 4-dev). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…ted) Weak (per=8000/GPU, cap-calibrated overflow-free): throughput rises 3.8e5->6.2e5 part/s for ndev 2->5 (per-eval 42->64ms). Strong (total=40000) is density-limited: per-GPU N>~8000 explodes the fixed-topology traversal caps (still overflow at x64), 400-720ms truncated -- only ndev=5/per=8000 valid. Healthy regime per-GPU N~8000. CAVEAT: jit=True illegal-address crash is INTERMITTENT (recurred once at ndev=2). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
TobiBu
changed the title
…ase 1)
Add DistributedFMMConfig.local_walk {dual_tree, treecode}. The dual-tree self-walk's
transient pair-queue caps per-GPU N (diagnosed self_queue_overflow: at per=1000/leaf=8
it overflows and gives garbage, aggL2 0.90 vs direct). The device-resident fast-lane
treecode walk (_build_treecode_artifacts_strict_streamed, mac_type=dehnen for accuracy-
profile parity + multi-step stability) streams far/near with no such queue and is a
drop-in for the self-walk: its CompactTaggedFarPairs feed the same real M2L (leaf-only
targets -> L2L no-op) and its self-excluded near CSR feeds the same combined P2P.
Seam adapted: the treecode far pairs are 0-padded (not -1) with the true count in
far_pair_count, so s_active uses far_pair_count on the treecode path (a >=0 test would
count padding -> degenerate delta=0 M2L -> NaN).
CPU-validated (ndev=4): treecode clears the self-overflow (self_ovf 4->0) and matches
direct (per=1000 aggL2 4.0e-3 where dual-tree self-walk overflows to 0.90; per=256
1.3e-3). The cross-domain LET walk is still dual-tree, so auto_scale_caps keeps the
whole eval overflow-free. Test test_driver_local_walk_treecode.
Cross-walk treecode swap + shard_map/GPU scale validation (Phase 2/3) are follow-ups.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Phases 1-2 done + pushed (treecode self-walk under shard_map, multi-GPU parity ndev 2-5 overflow-free). Phase 3 remaining: right-size the treecode near_cap (silent-overflow risk at ~1M/GPU), cross-walk treecode swap or auto_scale for connected ICs, a realistic domain-decomposed galaxy-disk IC, and 100k-1M/GPU scale + weak/strong timing (build-once) vs the single-GPU 4M/GPU curve. Includes resume environment (worktree + finder-repoint), reusable scripts, and gotchas. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…e-length 88) Wrap the >88-char lines pre-commit.ci flagged (black 26.5.1). isort clean; fmm.py already conformant. No behaviour change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…mit hooks) The CI `lint` job installs the [dev] extras and runs `black --check .` / `isort --check-only .`, but [dev] had floating ">=" pins while .pre-commit-config pinned the hooks. CI resolved a newer isort than the pinned 8.0.1 hook; the two disagree on single-import collapsing (no common fixed point), so `isort --check-only .` failed on files the hook considered clean (tests/unit/test_solver_api.py, jaccpot/runtime/_fmm_impl.py). Mirror main's fix: pin [dev] black/isort to the hook revs and bump the isort hook to 9.0.0b1, then reformat to 9.x style (collapse single parenthesized imports). isort + black now agree local/hook/CI. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
for more information, see https://pre-commit.ci
…(Phase 3a/3b) Lift the connected-IC scaling ceiling of the distributed FMM driver from ~50k to ~200k particles/GPU by making capacity self-heal without OOM. Phase 3a -- right-size the treecode local-walk flat buffers + surface overflow: - Add treecode_near_cap / treecode_far_cap to DistributedFMMConfig. When None, the driver sizes the near buffer to max_neighbors_per_leaf * num_leaves and the far buffer to max_interactions_per_node * num_leaves, instead of the builder's fixed 1<<21 (2M) near default -- which both wasted the combined-P2P neighbour build and could SILENTLY truncate at ~1M/GPU (the treecode overflow guard is eager-only, skipped under the shard_map trace). - Plumb an explicit near_cap through _build_treecode_artifacts_strict_streamed; when absent it keeps the env/1<<21 default so the single-GPU fast lane stays byte-identical. - Surface the treecode far/near overflow flags from the true counts vs the caps so _reduce_overflow / auto_scale_caps grow them like the dual-tree caps. Phase 3b -- selective auto-scale (with_selective_scaled_caps): - The retry loop grew EVERY cap uniformly (with_scaled_caps) on any overflow. On a connected IC only the cross-near buffer overflows, but uniform scaling also doubled cross_max_interactions_per_node (the cross-FAR cap) each retry -- and the cross-far list feeds the fused real-M2L, whose Pallas buffer pads to a power of two and OOMs (~6 GiB at N>=200k/2GPU) even though the true cross-far volume is ~10^3 pairs. This coupling was the real ceiling, not cross-far volume. - with_selective_scaled_caps grows only the caps whose per-device overflow flag fired (read from the diagnostic vector); the retry loop now uses it. Validated on 2x A100: per=100k (N=200k, was a hard OOM) and per=200k (N=400k) now run overflow-free; per=400k/GPU hits a separate fused-M2L memory wall (genuine ~1M cross-far pairs), addressed later by M2L chunking. CPU unit test covers the selective-scaling cases (cross-near-only grows only the near cap; self far+near; explicit treecode caps; per-device flags). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…aps as presets Two follow-ups to the Phase 3a/3b self-healing caps, both aimed at the connected-IC high-N path. 1. cross_far_cap -- slice the cross-far M2L input to the actual far volume. The cross walk sizes interaction_sources at t_total * cross_max_interactions_per_node but PACKS the valid far interactions into the front [0, far_pair_count); the tail is -1. The driver was feeding the whole (mostly-invalid) buffer to the fused real-M2L, which pads that length to a power of two and OOMs (~6 GiB at N>=400k/GPU on 2 A100) even though the true far volume is far smaller. Slice to cross_far_cap (default cross_max_interactions_per_node * num_target_leaves, ~2x tighter than the buffer and >= the far volume measured; override via the new config field). A far volume above the cap is OR'd into cross_far_overflow so auto_scale widens it -- the slice never silently drops interactions. Measured: per=100k/200k byte-identical (no regression); per=400k no longer OOMs (now bounded by the retry budget, not memory). 2. cap presets -- cache the converged caps keyed by (per-GPU N, ndev) and reuse them. auto_scale DISCOVERS the right capacities but pays a shard_map recompile per retry (up to 8 at per=200k). New jaccpot/distributed/cap_presets.py persists the converged caps to JSON; distributed_fmm_accelerations gains cap_presets_path= to seed a repeat run at a known size (skip the retries -> single compile) and save the converged caps after an overflow-free run. auto_scale stays the safety net; a preset is a validated starting point that transfers within an IC family. This is the cheap realisation of "size caps from a pre-count": auto_scale is the pre-count, cached by problem size. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
The per-pair rotation blocks [N, Bp, mdp, mdp] built for the fused real M2L dominate its peak memory and are the cross-far OOM at high per-GPU N / domain count. Add DistributedFMMConfig.far_m2l_fp32 (default False): when set, the self + cross far-M2L inputs (multipoles, centres/deltas) are cast to fp32 so the kernel and its rotation blocks are fp32, halving M2L peak memory; the local-expansion accumulation stays coeff_dtype (fp64). Opt-in and distributed-only, so the single-GPU fast lane is untouched. Accuracy-neutral: the far field is bounded by the order-p truncation error (~1e-3), far above fp32 rounding (~1e-7). Measured per=50k/ndev=2 aggL2 1.6904e-3 == the fp64 baseline to 4 sig figs. Capacity: with fp32 + leaf_size=256, per=200k/GPU x4 (800k total) -- which OOM'd in fp64/leaf128 -- now runs overflow-free at aggL2 7.5e-3 (~2x the prior ~100k/GPU ceiling at ndev>=4). per=400k/GPU x4 still OOMs at theta_cross=0.1 (cross-far volume too large even halved); a looser theta_cross=0.3 fits but degrades accuracy to ~1e-1, so the accuracy-preserving path beyond ~200k/GPU is M2L chunking, not looser theta_cross. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The fused real M2L builds per-pair rotation blocks [Npairs, Bp, mdp, mdp]; running the whole cross-far list in one batch is the >200k/GPU OOM (the block tensors, not the pair count, dominate M2L peak memory). Add DistributedFMMConfig.m2l_chunk (default None = one full batch, unchanged). When set, the cross-far M2L is accumulated over ceil(Npairs/chunk) fixed-size blocks via lax.scan: each block builds its own [chunk, Bp, mdp, mdp] tiles, runs the fused kernel, masks, and segment-sums into the local expansions. Peak M2L memory becomes ~chunk pairs, independent of the total -- removing the per-GPU ceiling at some throughput cost (one kernel launch per block). Numerics-identical to the full batch: the per-block contributions are masked and summed associatively into the same accumulator. Validated on 2x A100 (per=50k, m2l_chunk=65536): aggL2 1.6904e-3 == the full-batch fp32 baseline to all printed digits. Opt-in and distributed-only (real basis), so the single-GPU fast lane and the complex path are untouched. Combine with far_m2l_fp32 for the accuracy-preserving path past the ~200k/GPU ceiling. The high-N capacity confirmation (per=400k/GPU x4 @ theta_cross=0.1, which OOMs full-batch) is queued behind GPU availability; the memory bound is by construction. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
…otation guard _chunked_real_m2l_accumulate's lax.scan body was missing param/return annotations, tripping tests/unit/test_type_annotation_guard.py (the only CI failure on the chunking commit; 428/428 other tests pass). Annotate it (carry + 4-tuple block -> (carry, None)). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…M2L ceiling) The cross-far M2L was already chunked; the self-far M2L (_accumulate_m2l_fullbatch) was still full-batch and became the next OOM -- 14.24 GiB at per=1.6M/GPU x4 on the 4-GPU max-N probe (1.2M/GPU = 4.8M total fit; 1.6M/GPU OOMs at the self-far rotation blocks). Route the real self-far M2L through the same _chunked_real_m2l_accumulate (lax.scan over fixed-size blocks) when m2l_chunk is set. Pure refactor of the full-batch path: reproduce its gather exactly -- valid mask (idx < active & src>=0 & tgt>=0), safe src/tgt, and centers cast to m2l_dtype BEFORE the tgt-minus-src delta subtraction -- then accumulate in blocks (associative masked segment-sum). Same helper the cross-far path uses, which was GPU-parity-validated numerics-identical (per=50k aggL2 1.6904e-3 == full batch). Annotation guard passes. Now both far M2Ls are bounded by the block size, so per-GPU N is gated by the near-field / tree buffers rather than the M2L. High-N GPU capacity confirmation (per>=1.6M/GPU x4) is pending shared-node GPU availability. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e >1GB peak) The chunked M2L pre-gathered the full [Npairs, C] source multipoles and [Npairs, 3] deltas before the scan. At high per-GPU N that gather is itself the OOM: the 4-GPU probe cleared the self-far M2L blocks (14.24GiB) but then died on a 1.72GiB _gather (packed_use[s_safe_src], ~14M self-far pairs) on an already-full GPU. Refactor _chunked_real_m2l_accumulate to take the compact multipole array + index/centre arrays and gather [block, C] (and compute the tgt-minus-src delta) INSIDE the scan. Only the per-block rotation blocks are large now; the full [Npairs, *] arrays never materialise. Both far call sites updated (self: single centre array; cross: coarse src tree/centres). Deltas are fp64 tgt-minus-src then cast -- matches the prior cross convention exactly. Parity re-confirmed: per=50k aggL2 1.6904e-3 == baseline. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…arest seed) cap_presets.lookup returned None when no preset was >= the requested per-GPU N, so a new N calibrated from the small defaults -- the full auto_scale doubling ladder = ~10-12 shard_map RECOMPILES (~200s each, GPU idle), ~30-40 min per point on the capacity probe. Add a fallback: when only smaller presets exist, seed from the nearest smaller one SCALED UP by the per-GPU-N ratio (ceil, None-preserving). auto_scale still refines, so the seed only needs to be close -- a new N now converges in ~1 retry instead of ~12. Exact-match and nearest-larger behaviour unchanged. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
…cut at high N) INDEX_DTYPE is int64, so the near-field CSR/densification materialised several [num_edges] int64 temporaries (~0.8GB each at >1M/GPU) plus the bincount scatter (the 1.27GiB alloc that OOM'd the 1.6M/GPU x4 probe). Narrow the pure index temporaries to int32 -- the argsort/bincount target tgt (leaf-row ids <= u_leaves) and the densification e/t_of_e/ rank/flat (< u_leaves*S_near) -- all << 2^31 at any feasible per-GPU N. Kernel-facing arrays (offsets/src_s/counts/sids) stay INDEX_DTYPE, so no kernel dtype risk. Numerically exact (int32 holds the same index values): per=50k parity aggL2 1.6904e-3 == baseline. Frees ~2GiB of the near-field peak; the 1.6M/GPU probe now clears the bincount and reaches a later 5GiB near-field-P2P alloc. Meaningfully exceeding ~1.2M/GPU still needs near-field leaf-block chunking (the whole near-field working set), not per-buffer narrowing. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
…ld peak)
The fused-Pallas near field densified the combined CSR into [u_leaves, S_near] tables +
[num_edges] temporaries and ran the leafpair P2P over all u_leaves at once -- the ~40GB
aggregate wall that capped the driver at ~1.2M particles/GPU (the M2L was already chunked;
the near field was the remaining limiter). Chunk it over blocks of TARGET leaves, mirroring
the M2L chunk (_chunked_real_m2l_accumulate).
- DistributedFMMConfig.nearfield_chunk (Optional[int], default None): target-leaf block size
for the pallas near field; None keeps the full-batch path. Applies only when use_pallas_near.
- nearfield_leafpair_pallas_decoupled (pallas/nearfield_fused_leaf.py): new sibling wrapper that
takes the target rows and the source gather pool as SEPARATE arrays (the kernel already reads
them as separate refs; the old wrapper just bound the same array to both). Kernel body and all
existing callers unchanged.
- _radix_fast_lane_prepacked_pallas_decoupled (nearfield/near_field.py): sibling of the prepacked
path calling the decoupled wrapper.
- _chunked_pallas_nearfield_accumulate (distributed/fmm.py): lax.scan over ceil(n_lloc/block)
target-leaf blocks. Per block: build [block, S_near] sids DIRECTLY from the CSR
(sids[i,j]=src_s[offsets[b0+i]+j] for j<min(counts[b0+i],S_near) -- no [num_edges] temporaries),
run self + decoupled pairs against the FULL resident source pool, scatter-add into the
[cap+halo,3] output. Peak drops from u_leaves*S_near to block*S_near.
Bit-parity (not just associative): only local target rows [0,n_lloc) are kept and each local
particle belongs to exactly one leaf/block, so per-block partials touch disjoint rows. Validated
per=50k/ndev=2: nearfield_chunk in {none, 64 (multi-block), 512 (single block)} all give aggL2
1.6904e-3 identical to the last digit. (An earlier raw-edge-window densification was NOT parity --
raw per-leaf counts can exceed S_near, so the window under-covered dense blocks; the direct
offsets+counts gather fixes it and is simpler.) Single-GPU fast lane provably untouched (new
siblings only). Capacity gain (>1.2M/GPU) measured separately.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
for more information, see https://pre-commit.ci
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
2 participants
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.
Folds the single-GPU fast-lane Pallas kernels into the multi-GPU distributed FMM force path (Phase-5 5c), plus the 5d cap-scaling + scaling study. Stacked on already-merged 5a/5b (real+dehnen default, native-real upward).
5c — Pallas kernels per device
_apply_real_m2l, soJACCPOT_STATIC_STRICT_FUSED_M2L_PALLAS=1engages the fully-fused Pallas M2L.DistributedFMMConfig.nearfield_backend∈ {auto,pallas,baseline}, auto→pallas on sm_80+): intra-leaf self block via_compute_leaf_p2p_prepared_large_n_self_only_impl+ cross-leaf pairs via_radix_fast_lane_prepacked_pallasover the_combined_neighborsCSR densified to a padded source-leaf table. Baseline path unchanged behind the knob.The near-field was the entire bottleneck
Build-once steady-state whole-eval, ndev=2, per=8000/leaf=128: baseline near-field 10.7 s → Pallas 43.5 ms (~245×). Traversal + far-field are only ~40 ms — the ~10 s previously attributed to an overhead-bound traversal was the pure-JAX baseline near-field leaf-pair P2P. M2L-Pallas is negligible at this near-field-dominated config.
Parity (guardrail, all passed)
Dehnen MAC + θ_cross=0.1 kept at defaults.
5d — cap scaling + scaling study
auto_scale_capsondistributed_fmm_accelerations: the cross-domain LET grows with device count, overflowing the fixed default traversal caps at ndev≥4. On overflow, retry withwith_scaled_capsuntil it fits; result carries the winning config +diagnostics["cap_retries"]. Testtest_driver_auto_scale_caps.Known follow-ups (not in this PR)
🤖 Generated with Claude Code