fix(dpa4): align pt_expt training and native-spin fine-tuning - #5963
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
|
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:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThe pull request updates native-spin forwarding, descriptor migration, configuration-derived checkpoint state, fine-tuning validation, AMP controls, matmul precision policies, and mode-specific compiled execution. ChangesNative spin and descriptor migration
Training runtime
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Training
participant ModelWrapper
participant CompiledGraph
participant Validator
Training->>ModelWrapper: configure matmul precision
ModelWrapper->>CompiledGraph: select training or evaluation cache
CompiledGraph-->>ModelWrapper: run compiled forward
Training->>Validator: pass precision policy
Validator->>CompiledGraph: evaluate with evaluation precision
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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: 4
🤖 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/dpmodel/descriptor/make_base_descriptor.py`:
- Around line 170-181: Update the documentation for reset_spin_parameters to
state that dormant spin parameters may exist in spin-free checkpoints with
arbitrary archived values, and that activation discards this checkpoint state by
resetting them before use. Preserve the description of function-preserving spin
activation and the zero-function behavior.
In `@deepmd/pt/model/descriptor/sezm.py`:
- Around line 2214-2230: Update the use_amp_infer helper docstring in
deepmd/kernels/utils.py to document that inference AMP is controlled
independently by DP_AMP_INFER and applies even when the descriptor’s use_amp is
false, matching the selection logic and
test_train_and_eval_amp_switches_are_independent.
In `@deepmd/pt/model/model/sezm_native_spin_model.py`:
- Around line 83-99: Update prepare_finetune in SeZMNativeSpinModel to call
self.atomic_model.reset_spin_parameters() rather than accessing
self.atomic_model.descriptor directly, so composed atomic models such as ZBL
bridges reset their learned child correctly. Add coverage for
spin-free-to-native fine-tuning with bridging_method="ZBL".
In `@deepmd/utils/argcheck.py`:
- Around line 5958-5959: Update the tf32_infer argument declaration in the
validating configuration, near the existing supported_backends("pt") setting, to
also include the "pt_expt" backend. Keep the existing PyTorch backend support
unchanged so generated documentation exposes this option for both trainers.
🪄 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: 953f6efc-a0fd-4c10-906b-3e90d35eec8d
📒 Files selected for processing (38)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/atomic_model/dp_atomic_model.pydeepmd/dpmodel/atomic_model/linear_atomic_model.pydeepmd/dpmodel/descriptor/dpa4.pydeepmd/dpmodel/descriptor/dpa4_nn/embedding.pydeepmd/dpmodel/descriptor/dpa4_nn/norm.pydeepmd/dpmodel/descriptor/make_base_descriptor.pydeepmd/dpmodel/model/base_model.pydeepmd/dpmodel/model/model.pydeepmd/dpmodel/model/native_spin_model.pydeepmd/dpmodel/model/spin_model.pydeepmd/dpmodel/utils/stat.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt/model/descriptor/sezm_nn/embedding.pydeepmd/pt/model/model/__init__.pydeepmd/pt/model/model/sezm_native_spin_model.pydeepmd/pt/train/training.pydeepmd/pt_expt/common.pydeepmd/pt_expt/descriptor/dpa4.pydeepmd/pt_expt/infer/deep_eval.pydeepmd/pt_expt/model/get_model.pydeepmd/pt_expt/model/native_spin_model.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/train/utils.pydeepmd/pt_expt/train/validation.pydeepmd/pt_expt/train/wrapper.pydeepmd/utils/argcheck.pydeepmd/utils/spin.pysource/tests/common/dpmodel/test_descrpt_dpa4.pysource/tests/common/dpmodel/test_zbl_bridging.pysource/tests/pt/model/test_descriptor_sezm.pysource/tests/pt/model/test_dpa4_ptexpt_grad_parity.pysource/tests/pt/model/test_sezm_model.pysource/tests/pt_expt/descriptor/test_dpa4.pysource/tests/pt_expt/model/test_dpa4_native_spin.pysource/tests/pt_expt/model/test_get_model_dpa4.pysource/tests/pt_expt/test_finetune.pysource/tests/pt_expt/test_training.py
💤 Files with no reviewable changes (1)
- deepmd/dpmodel/model/model.py
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5963 +/- ##
==========================================
- Coverage 79.60% 79.38% -0.23%
==========================================
Files 1085 1085
Lines 126405 126558 +153
Branches 4598 4592 -6
==========================================
- Hits 100631 100464 -167
- Misses 24120 24441 +321
+ Partials 1654 1653 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
A lot of this is solid and I want to name it before the objection. I checked the 1.1 -> 1.2 squaring migration algebraically rather than taking the docstring's word: env_agg is linear in r_tilde, and D = sum_c A[c,i] A[c,j] contracts over the coordinate axis, so it splits into D_geometry + D_spin with no cross terms; the old code's a-scaled spin channels contributed exactly a**2 * D_spin, so a -> a**2 after the move is exact, and sign loss does not matter because the old contribution was already even in a. The AMP/TF32 train-vs-infer split is the right fix and it closes the concern njzjz raised on #5503. Extending the pt/pt_expt grad-parity test to the SO(3) and so2_norm branches is the right way to catch the FrameExpand/FrameContract promotion gap, and it genuinely fails without the fix, because those branches are off by default and no previous parametrization ever constructed those modules.
My objection is to the new reset_spin_parameters / prepare_finetune capability, and it is a design question rather than a list of defects -- which is why I am leading with it. Details inline, but the short version:
The reset is a no-op for every checkpoint this PR can produce. The same commit zero-initializes all four spin routes at construction -- spin_scale from np.ones to np.zeros, mag_layer2.w to zeros, and both adam_spin_*_weight from rng.normal(0, init_std, ...) to zeros. A dormant source built by this code holds zeros, saves zeros and loads zeros, so zeroing it again achieves nothing. The reset only does real work on a pre-1.2 checkpoint, where dormant routes carried spin_scale = 1.0 and random adam_spin_*. That is a legacy-checkpoint hazard, not a fine-tuning one.
And the legacy path already has a hook. _migrate_variables() runs on exactly those checkpoints and already rewrites spin_scale. I checked _variables() in dpa4_nn/embedding.py: spin_scale, mag_layer1.matrix, mag_layer2.matrix, adam_spin_vec_weight and adam_spin_nbr_weight are all in the serialized dict, so a migration can reach every route it would need to zero, and it can discriminate the dormant case because the source descriptor knows its own spin_flags -- which it must, since a legacy checkpoint with genuinely trained routes has to keep them.
I take the docstring's point that a bridging composition owns several children rather than one descriptor, so the model level is the natural place to ask. But that argues the plumbing is needed somewhere, not that it is needed here: the discriminating information (legacy version plus dormant flags) is available at migration time.
Why I think this is worth answering before patching anything: three of the four defects I found are in this plumbing, not in the fix underneath it. If the legacy sanitization moves into the migration that already exists, the new capability and all three go away together, rather than being fixed one predicate at a time. If there is a case I am missing that genuinely needs a model-level reset on a 1.2 checkpoint, say so and I will drop the objection -- but then it should have a test, because none of the current ones exercise a source whose spin routes are non-zero.
Two things independent of that question: a test that cannot fail on unfixed code, and the open CodeQL alert. Both inline.
Green otherwise -- 53 checks pass.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Two blocking fine-tuning compatibility issues remain in the new native-spin handoff. The zero-initialization and migration work look sound, but the handoff currently rejects a plain spin-free source and compares magnetic types by unstable positional masks.
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
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (3)
source/tests/common/dpmodel/test_descrpt_dpa4.py (1)
319-328: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider covering the spin-free
use_spin=Noneretag.A version-1.1 payload built without
use_spinreaches neither migration branch and is still retagged to 1.2. That path is currently untested. One assertion pins the intended no-rewrite retag.♻️ Proposed extra case
def test_pre_spin_versions_keep_their_own_tag(self) -> None:Add after the existing test:
def test_spin_free_config_is_retagged_without_rewrites(self) -> None: """1.1 and 1.2 share their math when no spin route exists.""" dd = make_descriptor() data = dd.serialize() data["`@version`"] = 1.1 migrated = DescrptDPA4.deserialize(data) assert migrated.version == 1.2 coord, atype, nlist = make_inputs() nf = atype.shape[0] np.testing.assert_array_equal( np.asarray(dd.call(coord.reshape(nf, -1), atype, nlist)[0]), np.asarray(migrated.call(coord.reshape(nf, -1), atype, nlist)[0]), )🤖 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/common/dpmodel/test_descrpt_dpa4.py` around lines 319 - 328, Add a regression test after test_pre_spin_versions_keep_their_own_tag named test_spin_free_config_is_retagged_without_rewrites, covering a version-1.1 serialized descriptor with use_spin=None. Assert deserialization retags it to version 1.2 and that its descriptor output matches the original using make_inputs and the existing call path.deepmd/dpmodel/atomic_model/base_atomic_model.py (1)
937-937: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winConsider failing loudly if a spin input reaches the dense route.
Line 937 forwards
spinonly on the graph branch. The dense branch at lines 962-970 drops it silently.uses_graph_lower()returnsFalseafterdisable_graph_lower(), so a native-spin descriptor with the escape hatch pulled would calibrate the output bias from spin-free predictions without any signal.A guard converts that silent miscalibration into an explicit error.
🛡️ Proposed guard
else: + if spin is not None: + raise NotImplementedError( + "native-spin output-bias calibration requires the graph " + "lower; the dense route carries no per-atom moment" + ) ( extended_coord,🤖 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 `@deepmd/dpmodel/atomic_model/base_atomic_model.py` at line 937, Update the dense-route handling in the surrounding atomic-model flow to detect when a non-None spin input would be dropped, and raise an explicit error instead of continuing with spin-free calibration. Preserve the existing spin forwarding on the graph route and allow the dense route only when spin is absent.source/tests/common/test_finetune_utils.py (1)
308-318: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
pytest.raisesover the try/except/else pattern.
pytest.raiseswithmatchexpresses the same intent in fewer lines and reports a clearer failure.♻️ Proposed refactor
def test_finetune_rule_builder_rejects_non_native_spin_pretraining(): - try: + with pytest.raises(ValueError, match="requires a native-spin pretrained model"): finetune.FinetuneRuleBuilder( _model_config(["Fe", "C"]), _native_spin_model_config(["Fe", "C"], [True, False]), change_model_params=False, ).build() - except ValueError as exc: - assert "requires a native-spin pretrained model" in str(exc) - else: - raise AssertionError("expected ValueError") def test_finetune_rule_builder_rejects_changed_magnetic_element_set(): - try: + with pytest.raises(ValueError, match="active magnetic element set"): finetune.FinetuneRuleBuilder( _native_spin_model_config(["Fe", "C"], ["Fe"]), _native_spin_model_config(["Fe", "C"], ["C"]), change_model_params=False, ).build() - except ValueError as exc: - assert "active magnetic element set" in str(exc) - else: - raise AssertionError("expected ValueError")Add
import pytestif the module does not already import it.Also applies to: 321-331
🤖 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/common/test_finetune_utils.py` around lines 308 - 318, Refactor the exception assertions in test_finetune_rule_builder_rejects_non_native_spin_pretraining and the analogous test around lines 321–331 to use pytest.raises with a matching error message instead of try/except/else; add the pytest import if needed and preserve the existing ValueError expectations.
🤖 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/train/training.py`:
- Around line 16-18: Run Ruff formatting on the import in training.py and
collapse the single-name contextlib import to Ruff’s preferred one-line form,
without changing its behavior.
In `@source/tests/pt/model/test_sezm_model.py`:
- Around line 1713-1744: Update
test_migrated_spin_routes_are_trainable_after_activation to load the version-1.1
legacy_state directly into the model built with use_spin=[True, False], ensuring
the magnetic-target migration and zeroing path is exercised; avoid routing
through the spin-free migrated model, or explicitly document that the two-stage
load is the supported behavior.
---
Nitpick comments:
In `@deepmd/dpmodel/atomic_model/base_atomic_model.py`:
- Line 937: Update the dense-route handling in the surrounding atomic-model flow
to detect when a non-None spin input would be dropped, and raise an explicit
error instead of continuing with spin-free calibration. Preserve the existing
spin forwarding on the graph route and allow the dense route only when spin is
absent.
In `@source/tests/common/dpmodel/test_descrpt_dpa4.py`:
- Around line 319-328: Add a regression test after
test_pre_spin_versions_keep_their_own_tag named
test_spin_free_config_is_retagged_without_rewrites, covering a version-1.1
serialized descriptor with use_spin=None. Assert deserialization retags it to
version 1.2 and that its descriptor output matches the original using
make_inputs and the existing call path.
In `@source/tests/common/test_finetune_utils.py`:
- Around line 308-318: Refactor the exception assertions in
test_finetune_rule_builder_rejects_non_native_spin_pretraining and the analogous
test around lines 321–331 to use pytest.raises with a matching error message
instead of try/except/else; add the pytest import if needed and preserve the
existing ValueError expectations.
🪄 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: 84c5a71e-005b-4168-8784-94aae20fe5c6
📒 Files selected for processing (15)
deepmd/dpmodel/atomic_model/base_atomic_model.pydeepmd/dpmodel/descriptor/dpa4.pydeepmd/dpmodel/descriptor/dpa4_nn/embedding.pydeepmd/pt/model/descriptor/sezm.pydeepmd/pt/model/descriptor/sezm_nn/embedding.pydeepmd/pt_expt/descriptor/dpa4.pydeepmd/pt_expt/train/training.pydeepmd/utils/finetune.pysource/tests/common/dpmodel/test_descrpt_dpa4.pysource/tests/common/test_finetune_utils.pysource/tests/pt/model/test_descriptor_sezm.pysource/tests/pt/model/test_sezm_model.pysource/tests/pt_expt/descriptor/test_dpa4.pysource/tests/pt_expt/model/test_dpa4_native_spin.pysource/tests/pt_expt/test_finetune.py
🚧 Files skipped from review as they are similar to previous changes (2)
- source/tests/pt/model/test_descriptor_sezm.py
- source/tests/pt_expt/descriptor/test_dpa4.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Three independent review passes on the current head converge on two blocking execution paths: plain spin-free checkpoints are still rejected before the new initialization machinery can run, and native-spin models cannot be called through the enabled compilation wrapper. I left both findings inline.
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
- carry native spin through pt_expt compiled graph training - honor compiled_infer across regular, full, and EMA validation - reject dense bias calibration that would discard spin
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
All six findings from my earlier review are addressed at 3be0356, each with a test that fails on the pre-fix code:
- The
reset_spin_parameters/prepare_finetunehierarchy is gone; legacy cleanup now happens in the 1.1 -> 1.2 descriptor migration, so the unsafe forwarding loop in the PT linear atomic model went with it. - The native-spin compatibility check moved to
FinetuneRuleBuilderand compares magnetic element sets rather than positional boolean lists, so type-map reordering and added inactive elements no longer abort. - Conversion from a spin-free source is still unsupported, but now fails early with an actionable message instead of a
TypeError. - The AMP test drives
_forward_blockswith a CUDA device double, so the train/eval policy is observable on CPU;use_amp=False, eval, DP_AMP_INFER=1fails on the old gate. - The CodeQL uninitialized-local alert is resolved structurally.
I also checked the replacement mechanism, since it is new since my review. The migration is exact: the quadratic form sums over the coordinate axis only, so D splits into D_geometry + D_spin with no cross terms, an amplitude a contributed a2 * D_spin, and the new linear gate reproduces it with s = a2. Both dpmodel and pt carry the migration and the version bump.
Two things to keep in mind rather than block on: the migration test asserts the stored gate value, not a numeric round-trip of the forward output; and the CUDA jobs report skipping, so the autocast region itself was never executed in CI - only the selection logic was.
…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.
…through pt_expt assembly (deepmodeling#5960) Users reported that compiled DPA4 training runs ~2x slower on `pt_expt` than on `pt`. This PR is the result of chasing that: a performance bug in how the SO3 contractions were lowered, and a correctness bug where a configured `use_amp: false` was lost while `pt_expt` assembled the model. ## Changes **1. Weight broadcast across the node axis (`so3.py`, `lora.py`, `grid_net.py`)** `matmul(x[..., None, :], weight[None, ...])` makes the node count `N` the 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). The fix batches the contraction over the small `(D, F)` axes so `N` stays the GEMM ROW dimension and 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.py` are deliberately NOT changed: their operands are `requires_grad=False` buffers, so no backward reduce exists. Verified rather than assumed. **2. The same lowering for the frame mixers (`_degree_batched_matmul`)** Review found the `FrameContract` / `FrameExpand` mixers still on the broadcast spelling. They now share one helper, `_degree_batched_matmul`, written identically on the dpmodel side (`dpmodel/descriptor/dpa4_nn/grid_net.py`) and the pt side (`pt/model/descriptor/sezm_nn/grid_net.py`). Because it does no reshape, an empty node axis (`N == 0`) flows through unchanged instead of hitting a reshape error — pinned by a test. **3. `use_amp` was lost during `pt_expt` model assembly (correctness)** `use_amp` is a training-runtime policy, not model state, so it stays OUT of the portable serialization record (this is the deepmodeling#5963 position, and the jax deserializer actively rejects records carrying `use_amp: true`). The bug was elsewhere: `pt_expt` assembled its model by converting an already-populated dpmodel instance, and that conversion round-trips the component through `deserialize(serialize())`. Anything that is deliberately not in the portable record — `use_amp` among it — was therefore dropped, and training silently ran under bfloat16 autocast even when the input configured `use_amp: false`. The fix is at the assembly boundary, not in the record: `pt_expt` now constructs the wrapped class directly (`auto_wrapped_class(...)` in `make_model.py`, `get_model.py`, and the bridging composition path), so a live constructor-supplied component keeps its runtime state. The rule is stated once, in the `auto_wrapped_class` docstring; the call sites reference it. An earlier revision of this PR instead added `use_amp` to `serialize()`. That was reverted in review — it put a runtime knob into the portable record and would have broken the jax contract. **Also removed in review: an `enable_tf32` / `DP_TF32_INFER` implementation for `pt_expt`.** It contributes nothing to the speedup measured below (the benchmark card has no TF32 silicon), and deepmodeling#5958 owns the `pt_expt` training-runtime alignment — including the documented position that `pt_expt` always runs at `"highest"` matmul precision. `pt_expt` therefore keeps master's warn-and-ignore behavior for `enable_tf32`. ## 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%. **Provenance: measured at `ae720432b`**, the head at which this PR was opened — i.e. BEFORE the review changes (change 2, the frame-mixer lowering, and change 3's move from `serialize()` to the assembly boundary). Change 1, which is where the entire speedup comes from, is unmodified since. The numbers have not been re-measured on the current head; a re-run is pending and I will post it rather than silently reuse these. | training mode | `pt` (reference) | `pt_expt` at master | `pt_expt` at `ae720432b` | speedup vs master | |---|---|---|---|---| | eager | 0.891 s/step | 1.555 s/step | **0.921 s/step** | **1.69x** | | compiled | 0.545 s/step | 1.611 s/step | **0.535 s/step** | **3.01x** | This reproduces the reported issue at master — `pt_expt` compiled was 3.0x slower than `pt` compiled, 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 fix `pt_expt` is at parity with `pt`: eager within 3.4%, compiled within measurement noise. ## Known limitations - **Backward numerics are covered for the frame mixers, not for the SO3 / LoRA contractions.** `test_dpa4_frame_mixers.py` compares `_degree_batched_matmul`'s weight gradient against the pt module's at rtol/atol 1e-12. The rewritten SO3 and LoRA contractions are pinned on the forward against an explicit `einsum` reference (rtol/atol 1e-12, numpy and torch namespaces); their backward is still exercised only by tracing, not compared by value. - **The `pt` / `pt_expt` TF32 policy gap remains open.** On Ampere+ cards `pt` runs training matmuls under TF32 (`enable_tf32`, default `True`) while `pt_expt` ignores the key with a warning; the two backends are not speed-comparable there. Deferred to the deepmodeling#5958 training-runtime series. - **The residual compiled gap vs `pt` is not stable across sessions.** An earlier session measured `pt_expt` compiled 10.8% slower than `pt` compiled; the 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. - The history contains churn at the GridBranch router (`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. - Unrelated but found while benchmarking: **torch >= 2.11 ships no Volta (CC 7.0) kernels**, and compiled training requires >= 2.11 via `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_dpa4_frame_mixers.py` — `_degree_batched_matmul` vs the pt module: forward parity, the `N == 0` contract, and weight-gradient parity. - `source/tests/common/dpmodel/test_dpa4_lora.py` — new `test_lora_so3_call_matches_einsum_contract`: `LoRASO3.call` against the explicit `einsum("ndfi,difo->ndfo")` reference with a nonzero adapter, on both the numpy and torch namespaces, for `n_focus` 1 and 2. - `source/tests/pt_expt/model/test_get_model_dpa4.py` — `use_amp` survives model assembly (both branches), for the plain and the bridged/composed construction paths. - `source/tests/common/dpmodel/test_descrpt_dpa4.py` — `use_amp` is absent from the portable serialization record and defaults on deserialize. - Existing `test_grid_branch[1]`/`[2]` cover the changed SO3 contraction against the pt implementation at rtol 1e-12. - Run locally: 434 passed / 10 skipped across the dpa4 dpmodel, pt_expt and cross-backend parity suites, plus the pt_expt model suite. CUDA-gated precision-context cases were run on a T4 (29/29). ## Test status caveat — resolved An earlier revision of this description flagged two locally failing `pt_expt` AOTI-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 clean `upstream/master` worktree on the same machine fails both with the identical `InductorError: assert isinstance(index, CppCSEVariable) and index.is_vec` (torch 2.11 CPU-SIMD codegen bug on an `atomic_add` scatter buffer), and both tests pass on this branch with the known workaround `torch._inductor.config.cpp.simdlen = 1` (2 passed). The same bug is already documented in `source/tests/infer/gen_dpa4.py` / `gen_dpa2.py`. --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Summary
Details
The pt_expt wrapper treated the bare NumPy weights in
FrameExpandandFrameContractas buffers, so they never reached the optimizer. It also used NumPy-only operations inReducedEquivariantRMSNorm, coupled inference AMP to the training switch, and reused a train-mode compiled graph during evaluation. This change promotes the missing trainable arrays, makes the norm tensor-safe, separates training and inference precision policies, and caches compiled lowers independently for train and eval mode.Native-spin fine-tuning had two related state-transfer problems. Per-type spin gates derived from
use_spinwere persisted in checkpoints, allowing an all-zero gate from a spin-free pretraining to override the fine-tune configuration. In addition, dormant randomly initialized spin routes became active at full amplitude when a magnetic type was introduced. Configuration-derived gates are now non-persistent and archived copies are ignored when loading older checkpoints; output-bias calibration receives the native moment; and native-spin models no longer fabricate a virtual-atom scale.DPA4 spin routes now initialize at the zero function.
prepare_finetuneresets dormant routes only when activating a spin-free checkpoint, preserves already-trained matching routes, and rejects reassignment to a different magnetic-type set. The environment-spin gate moves after the quadratic form so it retains a nonzero gradient at zero. Versioned migration preserves the function of existing checkpoints, and the reset capability is forwarded through atomic-model compositions such as ZBL bridging.This PR contains no DPA4C-specific or LMDB changes.
Checks
source/tests/pt_expt/test_training.py: 50 passedgit diff --checkSummary by CodeRabbit
New Features
Bug Fixes