Repository navigation
feat(finetune): initialize native spin from spin-free models - #6094
Conversation
Allow --finetune to initialize native-spin DPA4/SeZM models on the PyTorch and pt-expt backends, and native-spin DPA4C on pt-expt, directly from spin-free checkpoints without a manually converted checkpoint. Construct the target from its fine-tuning input and transfer learned weights after branch selection and type-map alignment. Preserve the target spin scheme, magnetic element mask, and analytical bridging. A same-name multi-task branch undergoing spin-free to native-spin conversion must update its output statistics instead of resuming unchanged. Keep the conversion rules in a shared tensor-transfer module: - Dispatch descriptor-specific expansion and fitting input projection selection by their concrete network classes. - Copy DPA4 geometric weights while retaining fresh zero-output spin routes and the target's configuration-derived magnetic masks. - Insert DPA4C spin features between the geometric readout and trailing neighborhood masses/type embeddings in descriptor statistics and fitting input weights, including trailing frame and atomic parameters. - Extend the pair-FiLM output columns while retaining initialized spin heads and anchors. Keep the spin gate closed and the added fitting rows initialized so the geometric function is preserved with a trainable gradient into the spin branch. - Support MLP and GLU fittings, including linear fittings without hidden layers. Expand each input projection once and reuse the result for all state-dict aliases. - Reject incompatible tensor shapes, fitting families, and first-layer identity or duplicated residuals whose semantics change on expansion. Check residuals independently of resnet_dt before loading target state. - Detach copied values so temporary transfer tensors retain no autograd graph. All conversion work runs during initialization. Leave out_bias, out_std, and bias_atom_e to the target preset_out_bias and fine-tuning statistics. Start optimizer and EMA state afresh. Retain native-to-native transfer validation, existing type-map restrictions, and explicit RANDOM fitting initialization. Validation: - 127 CPU tests and 2 subtests passed across fine-tuning and bridging. - Check initial energy/force/virial parity, zero initial magnetic force, trainable magnetic gradients, and checkpoint/serialization round trips. - Cover all v20260911 DPA4C presets, MLP/GLU input layouts, residual rejection without target mutation, and frame/atomic fitting inputs. - Exercise single-task and multi-task CLI fine-tuning into plain and ZBL models, including target bias presets/statistics and fresh optimizer/EMA. - Repository pre-commit hooks passed. GPU and distributed validation were not performed for this change.
There was a problem hiding this comment.
Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.
An organization owner can choose the repository organization as the billing source for reviews requested by members with a paid Copilot license. Learn about Copilot code review billing. Review the Copilot code review billing settings.
|
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.
 |
| Layer / File(s) | Summary |
|---|---|
Conversion rules and learned paths deepmd/utils/finetune.py, deepmd/pt/utils/finetune.py, doc/train/finetuning.md, source/tests/common/test_finetune_utils.py, source/tests/pt_expt/model/test_get_model_bridging.py |
Rules identify when conversion is required and resolve applicable automatic fitting widths. Shared helpers locate learned-model paths and remap descriptor or fitting keys. Documentation and rule tests cover conversion requirements. |
Native-spin state conversion deepmd/pt/utils/spin_finetune.py, source/tests/pt_expt/test_finetune_native_spin.py |
The transfer utility validates model compatibility and adapts supported descriptor and fitting state. Tests cover preserved outputs and biases, fitting initialization, descriptor variants, and incompatible layouts. |
Trainer integration and end-to-end validation deepmd/pt/train/training.py, deepmd/pt_expt/train/training.py, source/tests/pt_expt/test_finetune_native_spin.py |
Both trainers apply converted state and compute native-spin input statistics before output-bias adjustment. CLI tests cover training modes, saved state, and model reloads. |
Priority:
Estimated code review effort: 4 (Complex) | ~50 minutes
Change: Feature
Sequence Diagram(s)
sequenceDiagram
participant RuleBuilder
participant Trainer
participant get_native_spin_finetune_state
RuleBuilder->>Trainer: provide conversion rule
Trainer->>get_native_spin_finetune_state: convert pretrained branch state
get_native_spin_finetune_state-->>Trainer: return converted parameters
Possibly related issues
- deepmodeling/deepmd-kit#6101 concerns delta-bias statistics during spin-free DPA4 to native-spin fine-tuning. This change adds that conversion path and initializes statistics, but does not add full-data residual scanning.
Suggested labels: new feature
Suggested reviewers: anyangml
Merge Risk: 
1effe
-
No specific issue remains identified for this change; it is mergeable after normal checks.
Pre-merge checks: 
-
Check name Status Explanation Description Check Check skipped - CodeRabbit’s high-level summary is enabled. Title check The title "feat(finetune): initialize native spin from spin-free models" matches the main change. The PR adds a path that lets native-spin models start from compatible spin-free checkpoints through --… Docstring Coverage Docstring coverage is inconclusive: files from earlier reviews still need to be rechecked within the check's analysis limits. Not blocking. Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Full details: Title check
-
Explanation
The title "feat(finetune): initialize native spin from spin-free models" matches the main change. The PR adds a path that lets native-spin models start from compatible spin-free checkpoints through --finetune. The changed files cover the training code, the finetune rule (convert_to_native_spin), the new spin_finetune.py transfer helpers, docs, and tests. The title is concise and specific, and it uses a conventional commit prefix. It does not mention the DPA4C-specific transfer rules or the rejection checks, but the title does not need to cover every detail.
-
Finishing Touches
-
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
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 help to get the list of available commands.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Agent: dot
Reviewed exact head 911ddbb: shared spin-free transfer, both trainer call sites, branch-rule changes and the submitted regression tests. No actionable source defect found in this review.
The transfer is constructed before loading target state; shape/residual incompatibilities fail without partially replacing model parameters. DPA4C inserts spin columns at the geometric readout boundary, shifting masses/type-embedding and fitting-parameter tails together, and expands every state-dict alias of the input projection. It preserves target spin masks, analytical bridging and output bias/scale initialization. MLP/GLU projection selection covers the no-hidden-layer case; rejecting input-width-dependent residuals is appropriate rather than silently changing the inherited function. Same-name multitask spin conversion is correctly kept out of the resuming shortcut, allowing the fine-tuning statistics pass. Existing nonconversion key remapping retains target-owned unmatched state.
Independently executed 13 exact-source backend-neutral tests from test_finetune_utils.py, including single/multitask rules, alias selection, RANDOM fitting, config nonmutation, spin-free acceptance, virtual-spin rejection and magnetic-set constraints; all passed. This isolated run supplied the environment dtype constant and does not execute the tensor transfer or trainers. I inspected the submitted tensor-layout, energy/force/virial, trainable spin-gradient, serialization and CLI tests; the author's reported 115-test CPU result is reused, not claimed as independently rerun. This review environment lacks PyTorch/native DeePMD, so native/GPU/distributed and performance scopes remain unverified here.
Current Python/C++/build CI is still running; actual CUDA test jobs are skipped. No code, branch or merge actions taken.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #6094 +/- ##
==========================================
- Coverage 78.26% 78.05% -0.22%
==========================================
Files 1176 1177 +1
Lines 142590 142753 +163
Branches 5072 5072
==========================================
- Hits 111603 111422 -181
- Misses 29048 29394 +346
+ Partials 1939 1937 -2 ☔ 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.
This PR lets native-spin models be initialized from a compatible spin-free checkpoint through --finetune (PT DPA4/SeZM and PT-expt DPA4/DPA4C). The shared validator now returns convert_to_native_spin instead of rejecting a spin-free pretrained model; both trainers branch on it, and the transfer itself lives in deepmd/pt/utils/spin_finetune.py, which builds a complete target state dict for a strict load, skipping output bias/statistics (left to the normal statistics pass) and, under random_fitting, the non-descriptor keys. DPA4/SeZM need no expansion; DPA4C gets a spin block inserted into descriptor.mean/stddev, the FiLM spin columns and the fitting input projection, with explicit rejection of mismatched families, shapes and residual layouts.
One blocking item and four non-blocking items are attached inline.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Agent: dot
Delta review of b982f46 against 911ddbb, after reading the existing review and author replies. No new actionable defect found in the four-file delta.
The spin-reference omission identified in discussion_r4366877946 is addressed for the reported same-elements conversion path. Both trainers now call the spin-only calibration after successful transfer and before output-bias adjustment, inside their existing rank-0 statistics action. DPA4C uses the existing per-type RMS estimator and setter without rerunning geometric calibration; DPA4/SeZM do not invoke the sampler. The same-name conversion remains excluded from the resuming shortcut. Failure synchronization remains intact, as do PT-expt's parameter/buffer broadcast and PT's subsequent distributed-wrapper synchronization. This corrects an omission in my earlier source-only review; that review did not establish descriptor-statistics initialization end to end.
The four coverage requests are now represented in the submitted tests: descriptor-family rejection; independent learned-part/type-map/spin-direction guards; real trainer RANDOM cases with exact fresh-fitting-parameter comparisons before optimization; and pretrain-script cases that deliberately change target width then check the saved pretrained width plus native-spin configuration. The new DPA4C checkpoint assertion checks the nonunit reference, and the isolated statistics test checks zero-moment/missing-type defaults plus unchanged geometry and closed-gate outputs. Not freezing precedence between multiple invalid inputs is reasonable; I do not regard that as a remaining correctness blocker.
Independent validation here: all four changed Python files parse, and nine NumPy-only checks passed using unmodified AST-extracted helper/estimator/setter functions with lightweight model-lookup carriers. These cover the fixture RMS (0.9433981132), per-type estimates, bridged lookup, model_spin alias, zero/missing-type defaults, unchanged geometric arrays, non-DPA4C sampler avoidance, and error propagation. They do NOT execute PyTorch modules, buffer conversion, trainer/CLI, or distributed/GPU behavior. PyTorch/native DeePMD are unavailable here; the author's 97-test/two-subtest result was inspected, not independently reproduced.
CI at this SHA is not all green: Python/C++ jobs and docs remain pending; CodeQL reports initial_fitting at test line 572. Source control flow indicates a false positive: both its sole assignment and sole use are guarded by the same unchanged mode == "random" condition. The existing alert thread is already resolved, but the check still reports failure; I have not changed or dismissed any check/alert. Actual CUDA test jobs are skipped.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Review of the spin-free -> native-spin initialization at b982f46. Two blocking issues are attached inline: the DPA4C transfer is unusable with the default auto-width fitting configuration, and the spin block inserted into descriptor.mean/descriptor.stddev is never calibrated. One non-blocking follow-up is attached about fine-tuning onto a target type map that adds an element.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
One finding from the re-review at 0771e81: the new automatic-width inheritance does not apply to the canonical linear_ener composition spelling, so the documented inheritance is skipped and the transfer aborts on a shape check.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Two findings from the re-review at 1effe5a, both on the DPA4/SeZM transfer path: a fitting-family mismatch silently leaves the target fitting network randomly initialized, and a target fitting with more hidden layers than the pretrained one loads without any diagnostic. The canonical-composition width inheritance raised in the earlier review is fixed for the single-learned-child case; the remaining composition issues are separate and not blocking.
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Approved at 1effe5a. The earlier blocking items are addressed: the DPA4C spin reference is calibrated before the bias refit in both trainers, the automatic-width inheritance now locates the learned child in a canonical linear_ener composition for the single-learned-child case, and the statistics, type-map and descriptor-family guards have matching tests.
Two non-blocking observations remain on the DPA4/SeZM transfer path and are filed inline in the preceding review: a fitting-family mismatch, or a target fitting with more hidden layers than the pretrained one, leaves part of the target fitting network randomly initialized without a diagnostic. Neither affects the documented workflow and both can be follow-ups.
51505b8
Summary
Allow native-spin models to initialize directly from compatible spin-free checkpoints through
--finetune, without a manually converted checkpoint. This covers PT DPA4/SeZM and PT-expt DPA4/DPA4C.The target is constructed from its fine-tuning configuration first. A shared transfer implementation copies compatible geometric weights while preserving target spin configuration, magnetic-element masks, analytical bridging, output-bias/statistics initialization, and fresh optimizer/EMA state. Same-name multitask conversion performs fine-tuning statistics rather than resuming the spin-free branch unchanged.
Implementation
Validation on this branch
Based on master
d07c097b; this PR contains one cherry-picked commit, with the same patch as the requested source change. It does not depend on the ASE or grid-pair PRs.Python 3.13 / PyTorch 2.13.0+cu130, CPU only:
test_finetune_native_spin.py,test_finetune.py,test_finetune_utils.py, andtest_get_model_bridging.pyin the PT-expt/common suites.GPU and distributed execution were not rerun for this PR. The original commit's validation record is historical; the results above describe this branch's verification.
Summary by CodeRabbit