Stream mode + Compiled autograd - #2
Merged
Merged
Conversation
use_compile previously had to be ignored for use_autograd_forces=True:
the nested autograd.grad(create_graph=True) double backward cannot be
lowered by torch.compile. New torchnep/compiled_autograd.py takes the
DeepMD/DPA route instead:
1. write (energy, forces, virial) as a pure function of
(params+buffers, batch geometry) — parameters are graph INPUTS, so
loss.backward() still fills model.parameters().grad;
2. make_fx symbolic trace: the first-order dE/drij gradient becomes
ordinary FX ops (no runtime double backward left);
3. strip the make_fx-inserted aten.detach chains, which would silently
cut the second-order path from the force loss to the parameters;
4. torch.compile the traced graph with dynamic shapes — one graph
serves all batch sizes (prime-count synthetic tracing inputs avoid
dim mis-specialization).
The traced energy re-implements the per-type NN dispatch with torch.where
(the eager mask.any() branches are data-dependent) and uses the bmm
contraction backend (the loop backend's per-type-pair masks are too). ZBL
stays outside the graph (typewise cutoffs use .item()) and is added
eagerly, exactly like the cached analytical path.
Verified on PdCuNiP (9615 frames, batch 16, A2000, float32):
outputs and ALL parameter gradients match eager autograd to <=1e-5;
one graph serves batch 16 and 64; fwd+bwd 25.3 -> 5.1 ms/batch (5x).
Training (3 epochs, stream_mode): loss curves identical to both eager
autograd and the analytical path at displayed precision;
autograd eager 25-26 s/epoch, 0.19 GiB
autograd compiled ~6.6-9 s/epoch, 0.29 GiB
analytical compiled ~7.0-7.3 s/epoch, 0.18 GiB
Single-GPU train_nep only (the DDP shim path is unchanged). Tests are
CUDA-gated so CI never runs the compile.
torch.compile on compute_properties_cached was splitting into 4 graph fragments (per-type mask.any() NN branches + the ZBL block's inner autograd.grad), leaving ~1336 kernel launches per training step with the GPU idle half the time — which is why the make_fx autograd path was beating the hand-written analytical gradient despite near-identical pure GPU compute time (8.4 vs 8.1 ms/step). Split the method: _cached_core (descriptors + branchless per-type NN + analytical forces; zero data-dependent Python control flow) compiles to ONE graph; the compute_properties_cached wrapper adds ZBL and assembles the result eagerly and takes an optional core_fn to substitute the compiled core (both trainers pass one). The branchless NN runs every type's net on all atoms with torch.where selection — negligible extra flops, and every parameter stays in the autograd graph on every forward (the DDP dummy-pass bookkeeping becomes unnecessary and is removed). k=16 PdCuNiP-family bench (batch 32, A2000): 451 kernels/step (was 1336), 10.2 ms/step (was 17.5). Full 10-epoch matrix, s/epoch (k=2/4/8/16): 2.44/2.00/2.50/2.70 — was 2.60/2.56/3.28/4.18, and now ahead of the make_fx autograd path (2.30/2.14/2.86/2.92) at k>=4. Loss curves match the previous build digit-for-digit at displayed precision; outputs vs eager and vs autograd agree at the float32 noise level; full suite passes on CPU and CUDA.
Both changes come from the 55-cell benchmark sweep (1-16 element types, 6000-frame subsets, batch 32, A2000): - stream_mode: the batch pipeline (background CPU assembly, pinned async H2D, compiled per-batch basis) hides streaming behind the GPU compute — measured speed parity with preloading under use_compile at every type count (e.g. k=16: 2.70 vs 2.70 s/epoch) and <=a few percent in eager mode, at ~10-15x less GPU memory (5.6 -> 0.46 GiB at k=16). Streaming is the better default; stream_mode=False preloads as before. The stream-vs-default equivalence tests now pin the baseline side to stream_mode=False explicitly. - backend auto: eager loop beats bmm clearly up to 8 types (8.8 vs 14.1 s/epoch) and only ties at 16 (15.6 vs 15.1), so the old ntypes>=8 -> bmm rule picked the slower backend in the 8-19 range. Auto now keeps loop until 20 types in eager mode; under use_compile it still resolves to bmm (fuses best, no data-dependent branches). Full suite passes on CPU and CUDA.
The 55-cell benchmark sweep (1-16 element types, 6000-frame subsets) showed the streamed store at speed parity with full GPU preloading under use_compile at every type count, and within a few percent in eager mode — while using ~10-15x less GPU memory. The pipeline explanation: batch assembly runs in the background prefetch thread, pinned H2D copies ride the copy engine, and the (compiled) per-batch basis costs ~0.4 ms of GPU time — all hidden behind the per-step compute. Preloading had no remaining use case, so: - GPUDataStore and its construction-time chunking helper are deleted; StreamDataStore is THE data store. - The stream_mode parameter is dropped from train_nep and train_nep_sharded; the q-scaler pass always uses the training batch size and the per-batch basis is always compiled under use_compile. - predict/eval/checkpoint paths are unchanged (same collate interface). - tests/test_stream_mode.py now verifies collate against an independently assembled reference (concatenation + offsets + basis straight from the ops functions, torch.equal), metadata/mask consistency, prefetch-vs- direct equality, and 2-rank DDP same-seed reproducibility (opt-in). Other tests migrate their store construction to StreamDataStore. Net -300 lines. Full suite passes on CPU and CUDA; GPU smoke run on the 9615-frame set reproduces the historical loss trajectory at 5.3 s/epoch (compiled) with 0.2 GiB peak.
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.
stream_modeoption are removed — the dataset stays in host memory andbatches are streamed to the device. Same speed, ~10–15x less GPU memory.
backend="auto": eager mode now usesloopbelow 20 element types(was 8);
bmmunderuse_compile.use_autograd_forces=True+use_compile=Truenow works (first-order gradient materialized viamake_fx) — ~4x faster than eager autograd.early_stop: a stage-1 plateau jumps into Stage 2 insteadof ending the run; only a final-stage plateau stops training (kept across
resume).