perf(dpa4): stop broadcasting weights across the node axis; fix use_amp serialization - #5960
perf(dpa4): stop broadcasting weights across the node axis; fix use_amp serialization#5960wanghan-iapcm wants to merge 11 commits into
Conversation
The GridBranch router contraction einsum("ngfhc,nfh->ngfc") was written as
a broadcast multiply followed by a reduce over the branch axis. That
materialises the entire (N, G, F, H, C) product -- roughly 0.8 GB at the
grid resolution of examples/water/dpa4 -- writes it to memory and reads it
straight back, and the backward pays the same traffic again.
An op-level CUDA profile of a DPA4 training step measured this single
reduce at 45.6 ms per call over a [1152, 9, 1, 32, 576] operand, three
calls per step: the most expensive kernel in the run. The pt backend
spells the same contraction as torch.einsum and never builds the
intermediate.
Use xp.matmul instead, which is array-API standard (unlike np.einsum,
which is what the broadcast form was avoiding) and contracts H in place so
only the (N, G, F, C) result is written. matmul broadcasts its leading
batch axes, so the router reshapes to (N, 1, F, 1, H) and lines up with
value's (N, G, F, H, C) without any permute -- a permute would reintroduce
the copy this removes.
This reverts commit 7518a41. Measurement did not support it. An op-level profile attributed the 45.6 ms reduce to ExpandBackward0, not to this multiply, and re-benchmarking after the change moved DPA4 eager training by nothing (1.514 -> 1.552 s/step, i.e. run-to-run noise) while the offending kernel stayed byte-identical at 410.7 ms. The GridBranch product is well under the size that would matter. Since matmul is autocast-listed where mul/sum are not, keeping it would have silently moved this contraction into bf16 under the autocast region for no measured gain. The actual site is the broadcast weight in so3.py, fixed separately.
Both so3 channel mixers spelled their einsum as a batched matmul with the
NODE/EDGE axis as the matmul BATCH and the weight carrying a dummy leading
axis:
matmul(x[:, :, :, None, :], weight_expanded[None, ...])
matmul broadcasts batch axes, so this expands the weight to
(N, D, F, Cin, Cout). For examples/water/dpa4 that turns a 165K-element
parameter into 191M elements -- about 0.8 GB -- on every call, and autograd
must then reduce the whole expanded gradient back to the parameter shape.
An op-level CUDA profile of a DPA4 training step attributed 45.6 ms per
call to that ExpandBackward0 reduce over a [1152, 9, 1, 32, 576] operand,
three calls per step, making it the most expensive kernel in the run; the
ChannelLinear twin cost a further ~7-9 ms per call over [102510, 1, 32, 64].
The pt backend spells the same contraction as torch.einsum and never
expands the weight.
Batch over the small (D, F) / (F,) axes instead, which keeps N as matmul
ROWS. The weight is then used in place and its gradient is an ordinary
matmul. The transposes this adds touch only the (N, D, F, C) operands,
which are orders of magnitude smaller than the expanded weight.
Follow-up to the so3.py fix, applying the same correction wherever a contraction was spelled so that the NODE axis becomes the matmul BATCH and a trainable tensor is broadcast across it: * grid_net.GridBranch einsum "ngfhc,nfh->ngfc" -- was a broadcast multiply plus a reduce, materialising an (N, G, F, H, C) product H times the size of its own result. * grid_net.FrameContract / FrameExpand einsum "ndfi,dio->ndfo" -- broadcast the per-degree weight to (N, D, i, o). Both now share _degree_batched_matmul, which batches over the small degree axis. * lora.call einsum "ndfi,difo->ndfo" -- the LoRA twin of the so3.py site. In every case autograd had to reduce the fully expanded gradient back to the parameter shape on each step; batching over the small (D, F) axes keeps N as matmul ROWS so the weight is used in place. The two projection.py sites that share the [None, ...] spelling are left alone deliberately: to_grid_mat / from_grid_mat are registered as BUFFERS with requires_grad=False (verified on a constructed DPA4), so no gradient is taken for them and none of the expensive half applies. Covered by the existing pt-parity gates, which construct these classes directly: test_dpa4_frame_mixers.py (FrameContract/FrameExpand, fp64 weight-copied vs pt), test_dpa4_gridbranch_frames.py, test_dpa4_lora.py, and test_dpa4_dpmodel_parity.py.
…nored The descriptor's use_amp flag was never written to serialize(), and deserialize() feeds config straight into __init__, so any rebuild fell back to the True default. The pt_expt backend rebuilds the descriptor from that dict, so 'use_amp: false' in the input was silently discarded and training stayed in bfloat16 autocast; only the pt backend, which builds once from the config, honoured it. Caught while benchmarking: disabling AMP made pt 23% faster on a Turing GPU (no bf16 tensor cores) while pt-expt did not move at all, and an op-level profile showed pt-expt still spending 45% of its device time in bf16 gemm kernels with use_amp=false. Add the key to both the dpmodel and pt serialize configs so the two stay key-identical and the flag survives a cross-backend round-trip. Records written before this change deserialize unchanged -- the key is simply absent and __init__ supplies the default. The pre-existing round-trip tests compare forward OUTPUTS, which cannot catch this: dpmodel never autocasts, so the outputs agree whatever use_amp says. The new test pins the attribute itself, for both boolean values, and fails on the previous code.
pt_expt accepted `model.enable_tf32` and threw it away with a warning, so
DPA4/SeZM training always ran at "highest" matmul precision while the pt
backend -- reading the same input.json -- ran its training forwards under
`set_float32_matmul_precision("high")`. On Ampere and later that is the
difference between TF32 tensor cores and fp32 CUDA cores for every matmul,
and GEMM is ~60% of compiled device time on this workload, so the two
backends were not comparable on that hardware at all.
Mirror pt's policy exactly: TRAINING forwards follow `enable_tf32`
(argcheck default True), EVAL forwards follow `DP_TF32_INFER` (0/1/2 ->
highest/high/medium, invalid values rejected). Scope matches pt, where
argcheck declares the knob inside the dpa4 model arg block and only the
sezm builders wire it: pt_expt attaches it in `get_sezm_model` and
`get_native_spin_model`, and every other model keeps class defaults that
select full fp32 in both modes.
Ownership: `call_common` is the single owner for eager forwards -- every
pt_expt model's `forward` reaches the backbone through it, and the export
trace roots at `call_common_lower`, so the precision switch never enters an
exported graph. The compiled path needs its own application because
`_CompiledModel.forward` bypasses `call_common` entirely; placing the
context only on the model would have left it dead on exactly the path this
is meant to speed up. The context spans the lazy compile there, since
Inductor picks its GEMM backend while lowering.
Gating on `self.training` is what keeps the existing 1e-12 parity tests
valid: eval and export stay at "highest" unless DP_TF32_INFER asks
otherwise.
The GridBranch router contracts the branch axis H, and H is a handful (1 in
the water example). Spelling it as `matmul(router.reshape(N, 1, F, 1, H),
value)` therefore asks cuBLAS for a batched GEMM with M=1 and K=H, which it
serves from its small-N kernels (`gemmSN_*`, `gemmk1`).
A shape-resolved profile of a compiled DPA4 training step found this to be
the single largest GEMM in the run:
aten::bmm [[119808, 1, 1], [119808, 1, 96]] 0.0249 s/step forward
with its two backward siblings adding 0.0138 s/step -- together ~0.039 s/step
against a total pt-vs-pt_expt compiled gap of 0.055 s/step. The batch is
N(1152) * G(104) and K is 1: no contraction is happening at all, it is a
scalar multiply routed through a GEMM kernel.
Micro-benchmarked fwd+bwd at those exact shapes:
H=1: matmul 7.523 ms mul+sum 1.828 ms (4.1x)
H=3: matmul 4.436 ms mul+sum 4.453 ms (equal)
so the broadcast form is never worse. The comment this replaces claimed the
intermediate costs "H times the size of the result" -- true, but H is small,
and the measurement shows it does not pay for the degenerate GEMM.
This restores the spelling that 7518a41 replaced and 01c58e6 restored
once already; that revert was justified on a different workload (AMP-on
eager, where the site was invisible) and 7545961 then re-applied the
matmul as part of a broader sweep without re-measuring this site. The
numbers above are what was missing both times.
Conflict in deepmd/pt_expt/model/get_model.py, resolved keeping both sides: * imports -- this branch added `os` (for DP_TF32_INFER), upstream added `TYPE_CHECKING`; kept both. * the bridging return -- upstream (deepmodeling#5939) factored the ZBL composition into `_compose_bridging`, while this branch applied `_apply_tf32_policy` at each return site. Took upstream's helper and attached the TF32 policy to whichever model it returns, so both changes keep their behavior.
The contraction and TF32 comments had grown into measurement essays. Keep the part a reader needs -- why the obvious spelling is wrong -- and drop the profiling detail, which belongs in the PR discussion rather than the source. Also drops a `logging` import left unused when the enable_tf32 warn-once test was replaced.
📝 WalkthroughWalkthroughThe PR preserves ChangesDescriptor precision and execution updates
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant get_model
participant GeneratedCM
participant _CompiledModel
participant CUDA
get_model->>GeneratedCM: apply enable_tf32 and DP_TF32_INFER policy
GeneratedCM->>CUDA: set training or inference matmul precision
_CompiledModel->>GeneratedCM: enter tf32_precision_ctx for forward
GeneratedCM-->>_CompiledModel: execute shared forward
GeneratedCM->>CUDA: restore previous matmul precision
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/pt_expt/model/test_get_model_dpa4.py (1)
323-352: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest non-default evaluation precision in the context.
This test runs evaluation only with
DP_TF32_INFERunset. Lines 298-313 verify the stored attribute, but they do not verify thattf32_precision_ctx()uses"high"or"medium".Parameterize this test with
DP_TF32_INFER="1"and"2". This prevents an evaluation branch that always selects"highest"from passing.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt_expt/model/test_get_model_dpa4.py` around lines 323 - 352, Extend test_tf32_precision_ctx_selects_and_restores to parameterize DP_TF32_INFER for evaluation cases, covering "1" and "2" with expected precisions "high" and "medium" respectively, while keeping training cases unset. Set the environment variable per case before entering tf32_precision_ctx so the evaluation branch is verified and precision restoration remains asserted.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/pt_expt/model/make_model.py`:
- Around line 485-507: The shared process-wide precision mutation in
tf32_precision_ctx must not overlap across concurrent forwards. Add and reuse a
model-level lock to serialize the entire precision-setting, yield, and
restoration block, or explicitly document concurrent forwards as unsupported if
that is the intended contract.
---
Nitpick comments:
In `@source/tests/pt_expt/model/test_get_model_dpa4.py`:
- Around line 323-352: Extend test_tf32_precision_ctx_selects_and_restores to
parameterize DP_TF32_INFER for evaluation cases, covering "1" and "2" with
expected precisions "high" and "medium" respectively, while keeping training
cases unset. Set the environment variable per case before entering
tf32_precision_ctx so the evaluation branch is verified and precision
restoration remains asserted.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 4121890c-a131-4838-8ece-1beed8953924
📒 Files selected for processing (10)
deepmd/dpmodel/descriptor/dpa4.pydeepmd/dpmodel/descriptor/dpa4_nn/grid_net.pydeepmd/dpmodel/descriptor/dpa4_nn/lora.pydeepmd/dpmodel/descriptor/dpa4_nn/so3.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt_expt/model/get_model.pydeepmd/pt_expt/model/make_model.pydeepmd/pt_expt/train/training.pysource/tests/common/dpmodel/test_descrpt_dpa4.pysource/tests/pt_expt/model/test_get_model_dpa4.py
The router site ends up byte-identical to master: 7518a41 replaced the broadcast sum with a matmul, 504bb24 put the sum back, and the net diff was a one-line comment swapped for five -- losing master's (N, G, F, C) shape annotation on the way. Restore master's line exactly, so the branch touches this site not at all. The degenerate GEMM that profiling found there was self-inflicted: it existed only on this branch, never on master, so "fixing" it delivered nothing. Also corrects the so3 ChannelLinear comment, which claimed the contraction is batched over the focus axis. What matters is that B stays the GEMM rows; at n_focus=1 -- every shipped config -- both permutes are contiguous views and the whole thing is one (B, Cin) x (Cin, Cout) GEMM at no copy cost.
…backend" This reverts the pt_expt TF32 policy (99d33ea plus its comment edits in ae72043), restoring the warn-and-ignore behavior on master. The knob is unrelated to this PR's measured speedup (the benchmark card has no TF32 silicon; the whole 1.69x/3.01x gain comes from the contraction fix), its benefit was never measured, and PR deepmodeling#5958 owns the pt_expt training runtime alignment -- including the documented position that pt_expt runs at 'highest' matmul precision. Keeping a second, contradicting implementation here would split ownership of the same policy across two PRs.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5960 +/- ##
==========================================
- Coverage 79.59% 78.31% -1.29%
==========================================
Files 1081 1081
Lines 126244 126245 +1
Branches 4592 4598 +6
==========================================
- Hits 100490 98866 -1624
- Misses 24101 25735 +1634
+ Partials 1653 1644 -9 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Users reported that compiled DPA4 training runs ~2x slower on
pt_exptthan onpt. This PR is the result of chasing that: two real bugs (a performance one and a correctness one), benchmarked to parity withpt.Changes
1. Weight broadcast across the node axis (
so3.py,lora.py,grid_net.py)matmul(x[..., None, :], weight[None, ...])makes the node countNthe matmul BATCH, so matmul broadcasts the weight to(N, D, F, Cin, Cout)and autograd then reduces that whole expanded gradient (ExpandBackward0) back to the parameter shape. At the water example's sizes a 165 K-element weight expanded to 191 M elements (~0.8 GB) per call, and the reduce was the single costliest kernel of a training step (45.6 ms, 3x per step). Batching over the small(D, F)axes instead keepsNas matmul ROWS, so the weight is used in place.Micro-benchmark, fwd+bwd at the real shapes: 16.48 ms -> 1.09 ms (15x).
Two lookalike sites in
projection.pyare deliberately NOT changed: their operands arerequires_grad=Falsebuffers, so no backward reduce exists. Verified rather than assumed.2.
use_ampwas never serialized (correctness)deserializedoescls(**config), so a key absent fromserialize()silently reverts to the__init__default on every rebuild.pt_exptrebuilds the descriptor from config, so a configureduse_amp: falsewas reset toTrueand training stayed under bfloat16 autocast. Added toserialize()on both the dpmodel and pt sides so the two backends' contracts stay key-identical.The pre-existing round-trip tests could not catch this: they compare forward outputs, which are identical either way because dpmodel never autocasts. The new test pins the attribute itself.
Removed in review: an
enable_tf32/DP_TF32_INFERimplementation forpt_expt. An earlier revision of this branch madept_expthonormodel.enable_tf32likeptdoes. It was reverted (see theRevert "feat(pt_expt): honor enable_tf32 ..."commit): the knob contributes nothing to the speedup measured below (the benchmark card has no TF32 silicon and the benefit was never measured on one that does), and #5958 owns thept_expttraining-runtime alignment — including the currently documented position thatpt_exptalways runs at"highest"matmul precision. Splitting that policy across two PRs would leave it with two owners.pt_expttherefore keeps master's warn-and-ignore behavior forenable_tf32; closing the gap belongs to the #5958 series.Benchmark
DPA4 water example (
examples/water/dpa4), one Tesla T4, torch 2.11, fp32 (use_amp: false), batch size 6. Steady-state seconds per training step, obtained by differencing the wall time of a 33-step and a 3-step run of the same config, which cancels every one-time cost (import, data load, statistics,torch.compile/ make_fx lowering). All five arms were measured in one session on the same machine; run-to-run variation is about 2-3%.pt(reference)pt_exptat masterpt_exptthis PRThis reproduces the reported issue at master —
pt_exptcompiled was 3.0x slower thanptcompiled, and even slower than its own eager path, because the broadcast-weight contraction lowers to worse code under inductor than under eager cuBLAS. After the fixpt_exptis at parity withpt: eager within 3.4%, compiled within measurement noise. The whole speedup is change 1.Known limitations
pt/pt_exptTF32 policy gap remains open. On Ampere+ cardsptruns training matmuls under TF32 (enable_tf32, defaultTrue) whilept_exptignores the key with a warning; the two backends are not speed-comparable there. Deferred to the feat(pt_expt): align the training runtime with pt #5958 training-runtime series.ptis not stable across sessions. An earlier session measuredpt_exptcompiled 10.8% slower thanptcompiled; the final benchmark above measured it 1.7% faster. Both are within a couple of run-to-run standard deviations, so I treat compiled as at parity and the earlier gap as unconfirmed.7518a417c->01c58e665->75459610a->504bb2430->157444204): a matmul spelling introduced, reverted, reintroduced, and finally restored to master's line. The site is byte-identical to master in the final diff. The degenerate GEMM that profiling found there existed only on this branch, so it is not a fix — I have left the commits rather than rewriting pushed history, and would squash them on request.check_compile_torch_version. Compiled DPA4 training is therefore impossible on V100 with official wheels; T4 (CC 7.5) is the oldest card that works.Tests
source/tests/common/dpmodel/test_descrpt_dpa4.py— newuse_ampround-trip case, both branches.test_grid_branch[1]/[2]cover the changed contraction against the pt implementation at rtol 1e-12.Test status caveat — resolved
An earlier revision of this description flagged two locally failing
pt_exptAOTI-freeze tests (test_zbl_bridging.py::test_native_spin_with_bridging_graph_freeze_and_deep_eval,test_dpa4_zbl_parallel.py::TestBridgedSpinGraphSelfComm::test_freeze_embeds_with_comm_artifact) as unadjudicated. They are now adjudicated as pre-existing and environmental, not caused by this branch: a cleanupstream/masterworktree on the same machine fails both with the identicalInductorError: assert isinstance(index, CppCSEVariable) and index.is_vec(torch 2.11 CPU-SIMD codegen bug on anatomic_addscatter buffer), and both tests pass on this branch with the known workaroundtorch._inductor.config.cpp.simdlen = 1(2 passed). The same bug is already documented insource/tests/infer/gen_dpa4.py/gen_dpa2.py.Summary by CodeRabbit