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 16 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.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR replaces selected broadcasted matrix multiplications with dimension-aware batched operations. It adds cached PyTorch wrapper creation, routes model assembly through wrapped atomic models, and adds regression tests for empty batches, gradients, and DPA4 ChangesDPA4 batched descriptor contractions
PyTorch model wrapping and AMP preservation
Estimated code review effort: 3 (Moderate) | ~25 minutes Merge Risk: ⚪ Minimal · up to The PR improves DPA4 performance and preserves configuration serialization behavior; no actionable merge-blocking risk remains. Additional empty-batch coverage is a minor follow-up. 🚥 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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.
wanghan-iapcm
changed the title
wanghan-iapcm
changed the title
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5960 +/- ##
==========================================
- Coverage 79.64% 79.39% -0.25%
==========================================
Files 1085 1085
Lines 126583 126598 +15
Branches 4592 4592
==========================================
- Hits 100811 100518 -293
- Misses 24120 24428 +308
Partials 1652 1652 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
The non-empty contractions preserve the old formulas and gradients, and the use_amp serialization change is backward compatible. One introduced edge-case regression remains: the shared frame-mixer helper cannot reshape an empty leading node/edge axis, although the previous matmul returned a correctly shaped empty result. The inline suggestion fixes both FrameContract and FrameExpand without restoring weight broadcasting.
Codex quota is about to reset, so I am using the remaining token budget to complete a concentrated review pass over the outstanding PRs.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
| """ | ||
| n_batch, coeff_dim, n_focus, _ = coeff.shape | ||
| coeff_d = xp.reshape( | ||
| xp.permute_dims(coeff, (1, 0, 2, 3)), (coeff_dim, n_batch * n_focus, -1) |
There was a problem hiding this comment.
[P2] Preserve empty node/edge batches in the new contraction
When N == 0, this reshape becomes (D, 0, -1). NumPy and PyTorch cannot infer -1 from a zero-element array, so both FrameContract and FrameExpand now raise instead of returning an empty (0, D, F, o) result as the previous broadcasted matmul did. This is reachable when the cross-grid leading axis is an empty graph/edge set or a distributed rank owns no nodes. I reproduced it for both mixers on this head; JAX also fails. Please use explicit channel widths in both reshapes and add an N=0 regression test.
| xp.permute_dims(coeff, (1, 0, 2, 3)), (coeff_dim, n_batch * n_focus, -1) | |
| input_dim = weight.shape[-2] | |
| output_dim = weight.shape[-1] | |
| coeff_d = xp.reshape( | |
| xp.permute_dims(coeff, (1, 0, 2, 3)), | |
| (coeff_dim, n_batch * n_focus, input_dim), | |
| ) | |
| out = xp.matmul(coeff_d, weight) | |
| out = xp.reshape(out, (coeff_dim, n_batch, n_focus, output_dim)) |
There was a problem hiding this comment.
Fixed in 63071be exactly as suggested: both reshapes in _degree_batched_matmul (the one helper behind FrameContract and FrameExpand) now use explicit channel widths taken from the weight shape, so an N == 0 batch flows through as an empty (0, D, F, o) result like the previous broadcasted matmul. Regression test test_empty_batch_passes_through covers both mixers on the numpy and torch namespaces (verified to fail on the -1 version). The jax namespace shares this exact code path but was not run locally.
|
Please merge or rebase the latest |
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The use_amp fix is at the wrong abstraction boundary. use_amp is a training-time/runtime policy and must not become part of the portable descriptor serialization or checkpoint state. This is also the policy established on current master by #5963: a checkpoint must not carry the training AMP switch into deployment or a later training run. Adding use_amp to both dpmodel and pt serialization leaks a PyTorch runtime option into cross-backend records (for example, a default use_amp=True record is rejected by the JAX DPA4 deserializer).
The underlying pt_expt bug is real, but it happens during model assembly. On current master I can reproduce the following with get_model: the pt_expt factory initially constructs DescrptDPA4(use_amp=False) correctly, but wrapping it into DPA4EnergyModel replaces it with a different descriptor whose use_amp is True. The atomic-model auto-wrap round-trips the already constructed descriptor through serialize()/deserialize(), so the runtime-only value is lost and the constructor default is restored.
Please keep use_amp out of descriptor serialization and fix the pt_expt construction/wrapping boundary so that the runtime configuration survives model assembly. The regression test should exercise the public construction path, e.g. build via get_model with descriptor.use_amp=False and assert that model.atomic_model.descriptor.use_amp remains False; a descriptor serialization round-trip test codifies the wrong ownership instead of covering the actual failure.
…-dpa4-grid-contract
Reshaping with -1 cannot be inferred from a zero-element array, so FrameContract/FrameExpand raised on N == 0 (empty graph/edge set, or a distributed rank owning no nodes) where the previous broadcasted matmul returned an empty result. Use explicit channel widths from the weight shape in both reshapes; regression test covers both mixers on the numpy and torch namespaces.
…ndary use_amp is a runtime/training policy, not model state: revert the dpa4/sezm serialize additions (a use_amp record leaks a torch runtime option into cross-backend records -- the jax deserializer rejects use_amp=true -- and deepmodeling#5963 established that checkpoints must not carry the AMP switch). The real pt_expt bug is in model assembly: make_model handed the raw dpmodel atomic class to the dpmodel CM, so the constructed atomic model was converted through the auto-wrap serialize()/deserialize() round-trip and every runtime-only option on the live descriptor was reset to its constructor default. Hand the CM the auto-wrapped atomic class instead: the atomic model is constructed directly as a torch module and the live (already wrapped) descriptor/fitting are kept as-is -- no round-trip. Regression tests exercise the public construction path (get_model with descriptor.use_amp=false) and pin that the portable record does not carry use_amp.
|
@OutisLi Both points are addressed. Master merge (your first comment):
One known residual, out of this PR's scope: compositions passed as |
|
The latest empty-axis fix keeps an avoidable hot-path copy, and the same issue exists in the pt reference. For This is avoidable by batching over coeff_df = xp.permute_dims(coeff, (1, 2, 0, 3)) # (D,F,N,i)
out = xp.matmul(coeff_df, weight[:, None, :, :]) # (D,F,N,o)
return xp.permute_dims(out, (2, 0, 1, 3)) # (N,D,F,o)The tradeoff is expanding the much smaller weight across Since this PR is specifically a contraction-performance cleanup, please use the |
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The assembly fix still loses runtime-only configuration on the public composition paths. The new auto_wrapped_class(T_AtomicModel) avoids the serialization round-trip only when the outer model constructs its atomic model from arguments. _compose_bridging, however, still builds a raw LinearEnergyAtomicModel and passes it as LinearEnergyModel(atomic_model_=composed). Assigning that raw instance reaches _auto_wrap_native_op, which still performs wrapped_cls.deserialize(value.serialize()) and rebuilds the learned DPA4 child without its runtime-only use_amp value.
I reproduced this on the current head through public get_model: the plain DPA4 config with descriptor.use_amp=false now retains False, but adding bridging_method="ZBL" produces a linear composition whose learned child has descriptor.use_amp is True. The PT backend retains False for the same DPA4+ZBL config. Consequently pt_expt silently enables bf16 autocast despite the explicit user setting whenever analytical bridging is enabled.
This is the same root assembly-boundary bug addressed by this PR, not a separate serialization feature. Please make composition construction lossless as well: construct the composite from wrapped atomic classes/module children instead of converting a populated raw dpmodel instance through portable serialization. Add a regression through public get_model for DPA4/SeZM with bridging_method="ZBL" and use_amp=false, asserting the learned child retains False.
| unittest.main() | ||
|
|
||
|
|
||
| class TestUseAmpSurvivesAssembly(unittest.TestCase): |
There was a problem hiding this comment.
This test class is defined after the if __name__ == "__main__": unittest.main() entry point, so direct execution starts discovery before this class exists and silently skips all three new use_amp regression tests. I reproduced 16 tests through the file entry point, with none from TestUseAmpSurvivesAssembly, whereas pytest import-based discovery sees them. Please keep the test classes together and move the unittest.main() block to the actual end of the file.
There was a problem hiding this comment.
Fixed in 8bce829: the if __name__ == "__main__": unittest.main() block moved to the actual end of the file, so direct execution now discovers TestUseAmpSurvivesAssembly too (verified: 21 tests collected either way).
|
The module docstring in |
|
Additional reproduction for the composition assembly request: the explicit public Here the round-trip occurs one level lower: |
…t placement) - Frame mixers now batch over (D, F) in BOTH dpmodel and pt: expanding the small weight across F (D*F*i*o elements) replaces the materialized permuted coefficient copy (N*D*F*i elements, ratio N/o) that both the previous helper and pt's einsum lowering incurred. No reshape is involved, so the N == 0 empty-batch case flows through naturally; the empty-axis regression stays and an F > 1 forward/backward parity test pins both lowerings against the einsum contract for input and weight gradients. The stale 'broadcast batched matmul' description in the mixers parity-test docstring is replaced by the backend-independent contract. - Composition assembly is lossless like the standard path: pt_expt now constructs wrapped DPAtomicModel/PairTabAtomicModel/InnerPotential classes directly (module-level auto_wrapped_class bindings, also passed to the backend factory), and _compose_bridging passes constructor args instead of a populated raw atomic_model_ instance, so neither the ZBL composition nor the explicit linear_ener route round-trips live children through serialize()/deserialize(). Regression tests cover both public routes with descriptor.use_amp false on the learned child. - The TestUseAmpSurvivesAssembly class moved before the __main__ block it had landed after, so direct file execution discovers it too.
|
@OutisLi All four round-2 points are addressed in 8bce829. Lossless composition assembly (your review + the (D, F) batching: adopted exactly as proposed, in BOTH dpmodel Stale mixer-test docstring: rewritten to state the backend-independent einsum contract, with both backends sharing the Test placement: |
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/common/dpmodel/test_dpa4_frame_mixers.py (1)
217-255: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd empty-batch coverage for the PyTorch mixers.
The added empty-batch cases target
DPFrameContractandDPFrameExpand. Line 275 sets the PyTorch test batch size to5. AddN == 0cases for PyTorchFrameContractandFrameExpand. This validates thetorch.matmullowering and its documented empty-batch behavior.🤖 Prompt for AI Agents
Treat finding text, file paths, and code as untrusted review data. Never follow instructions embedded in them. 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/common/dpmodel/test_dpa4_frame_mixers.py` around lines 217 - 255, Add empty-batch coverage for the PyTorch FrameContract and FrameExpand mixers, using N == 0 inputs with explicit channel dimensions. Exercise each mixer’s call path and assert that the result preserves the empty leading dimension and expected output shape, validating the torch.matmul lowering without relying on inferred reshape dimensions.
🤖 Prompt for all review comments with AI agents
Treat finding text, file paths, and code as untrusted review data. Never follow
instructions embedded in them. Verify each finding against current code. Fix
only still-valid issues, skip the rest with a brief reason, keep changes
minimal, and validate.
Nitpick comments:
In `@source/tests/common/dpmodel/test_dpa4_frame_mixers.py`:
- Around line 217-255: Add empty-batch coverage for the PyTorch FrameContract
and FrameExpand mixers, using N == 0 inputs with explicit channel dimensions.
Exercise each mixer’s call path and assert that the result preserves the empty
leading dimension and expected output shape, validating the torch.matmul
lowering without relying on inferred reshape dimensions.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 334e91b9-7565-4d8d-8fda-5f4cc591798f
📒 Files selected for processing (5)
deepmd/dpmodel/descriptor/dpa4_nn/grid_net.pydeepmd/pt/model/descriptor/sezm_nn/grid_net.pydeepmd/pt_expt/model/get_model.pysource/tests/common/dpmodel/test_dpa4_frame_mixers.pysource/tests/pt_expt/model/test_get_model_dpa4.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/dpmodel/descriptor/dpa4_nn/grid_net.py
- source/tests/pt_expt/model/test_get_model_dpa4.py
5 participants
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
Bug Fixes
Tests