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.
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