Repository navigation
perf(pt): add opt-in SO2 edge-cohort inference backward 【hnu】 - #2
Draft
macqueen09 wants to merge 7 commits into
Draft
macqueen09 wants to merge 7 commits into
macqueen09 wants to merge 7 commits into
Conversation
Use a compiled single-cohort branch for small edge sets and retain a streamed runtime path for larger frozen Neo input gradients. Keep native fallback and explicit stride defaults, add 60 CUDA checks and cold-first-peak evidence. The switch remains disabled by default; multi-model and AOT deployment are not validated. N64 cold allocated peak falls 3.62%; no speedup is claimed.
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
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.
Submitted by Hunan University.
Revision measured:
98a65e46d0076319f2fcb49ded77e6a93ecab6beintegrated with upstream4e1479d8. Matrix: 27 processes, 23 completed, 4 OOM (Neo4096 TORCH and Pro4096 in all three arms).Measured Neo results
N1000 passes against both Torch and BASE. At N4096, energy/forces/stress pass against BASE but virial does not; BASE and ECS also fail their respective repeated-run virial check. Torch N4096 runs out of memory.
Summary
Adds an opt-in, default-off SO(2) edge-cohort backward for frozen DPA4-Neo inference. Set
DP_SO2_EDGE_COHORT_INFER=1withDP_TRITON_INFER=2andDP_COMPILE_INFER=1to enable it. Unsupported layouts and protected gradient paths retain the existing implementation.Why it saves scratch memory
The native path advances each backward layer across the full edge set. This path completes the backward layers for a cohort of at most 8192 edges before advancing to the next cohort. The temporary
g,gz, andsigbuffers are reused across cohorts; saved activations and the returned input gradient still span all edges. This is an edge count, not an atom count. It reuses the native GEMM arithmetic and column parallelism, with separate full-graph and scratch strides for partial cohorts. Only owned scratch is overwritten; caller inputs are not mutated.Small runtime-dispatch correction
The latest revision changes two files (35 insertions / 5 deletions): the public backward dispatcher always enters the opaque runtime operator, and a regression test traces 37 edges and then compiles/runs 37, 8193, and 16385 edges. This prevents a small
make_fx/AOT example from fixing the large-input dispatch to the single-cohort trace. It does not change the numerical kernels. The complete feature diff remains two implementation files and one test file; no local benchmark data or experiment documents are included.Validation setup
Baseline: upstream
4e1479d8a23fcc1c9baf36004d5685898ad421a8. The feature branch retains its original base; results below use its clean local integration into that upstream commit. RTX 5090; PyTorch 2.11.0+cu128; Triton 3.6.0; FP32 with TF32 and AMP disabled; expandable CUDA allocator segments; one CPU intra/inter-op thread.TORCH = compiled PyTorch without native inference kernels; BASE = upstream native Triton inference; ECS = the updated feature enabled on the same upstream base. Each case uses the same checkpoint, device, and three deterministically perturbed diamond-carbon frames. There is one fresh process per implementation/case, three evaluations per frame for accuracy, eight timing warmups, and 30 synchronized wall-time samples. Timing includes the ASE energy/forces/stress evaluation, not just the SO(2) kernel. These are screening medians, not confidence intervals. Neo runs on GPU 0; Nano/Mini on GPU 1; Pro on GPU 0.
Checkpoints and actual dispatch
13edeea63448c0f8a0b38e08fd8ba196fb1e2fad5f12486b1cd23320829da268ded047546c2b44a50e2c850680d77279e2bf3e107c10f1c2fe39a5c344e693f68d8b2b02b9575cb9e0af420cc1e8b1d9ffcbc04752c2203d29f48ccd0cd5f07d446a6e893101fc52190cfccabb9d2328863d651aca36981fd8d3401c71a61221Pro's checkpoint metadata uses an older fitting-type alias; the same local load-only compatibility mapping is applied in every arm. Neo's supported narrow layout is the optimization target; Nano, Mini, and Pro are compatibility/fallback checks, not claims of ECS acceleration.
End-to-end latency
Negative change means lower latency. Compare BASE and ECS for incremental PR benefit; Torch is an additional reference.
Peak GPU memory
Decimal GB (10^9 bytes). Full-window allocated includes the first inference/compile window and timed evaluations. Steady allocated resets the peak after warmup. Reserved measures allocator-held capacity, not live tensors. Caches are shared per model, in TORCH → BASE → ECS order: cold compilation/autotuning can affect full-window peaks, so fallback-model first-window differences are not ECS savings. Use the steady allocated comparison to isolate runtime capacity. OOM runs have no valid steady timing or acceptance comparison.
Numerical consistency
Unchanged acceptance rule: every element must satisfy
abs(candidate-reference) <= 1e-5 + 1e-4*abs(reference). The table reports the worst error/tolerance ratio across energy, forces, virial, stress and all same-frame repeat pairs; ≤1 passes. Self-comparisons expose native reduction variability. These test implementation agreement, not physical accuracy against DFT. N1000 is the primary Neo accuracy case; N4096 is also reported without suppressing failures.Per-output absolute and relative errors
Relative error uses denominator
max(abs(reference), 1e-8)and is descriptive, not the acceptance criterion. Energy/virial: eV; force: eV/Å; stress: eV/ų.Regression coverage
Run the focused suite with the project's PyTorch/CUDA test environment:
The feature remains default-off and this PR remains Draft. Standalone AOTInductor/C++ deployment is not validated by the online-Python inference tests above. Multi-model failures/OOM and small-input overhead are reported rather than counted as improvements.
Reproducing the model comparison
Use a CUDA build of this revision and the same checkpoint. Run each arm in a fresh process, setting
DP_COMPILE_INFER=1,DP_CUDA_INFER=0,DP_TF32_INFER=0,DP_AMP_INFER=0,DP_INTRA_OP_PARALLELISM_THREADS=1,DP_INTER_OP_PARALLELISM_THREADS=1, andPYTORCH_ALLOC_CONF=expandable_segments:Truebefore importing DeePMD. SetDP_TRITON_INFER=0for TORCH, or2for BASE/ECS; only ECS setsDP_SO2_EDGE_COHORT_INFER=1.Collect three repeats of all three frames for each arm; compare corresponding frames with the tolerance above. For timing, warm up eight evaluations, synchronize CUDA, then time 30 evaluations with synchronization after each. Record
torch.cuda.max_memory_allocated()andmax_memory_reserved()from before the first evaluation; separately reset peak statistics after warmup and replay 30 evaluations for steady peaks. Raw measurements and checkpoint hashes are retained locally; the tables above are the review evidence, not additional repository files.