fix(pt, dpmodel): spin model finetune#5281
Merged
njzjz merged 7 commits intodeepmodeling:masterfrom Mar 6, 2026
Merged
Conversation
for more information, see https://pre-commit.ci
Contributor
|
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:
📝 WalkthroughWalkthroughAdds cached, spin-aware sampler wrappers and centralizes sampler wrapping in BaseAtomicModel; introduces SpinModel wrappers for change_out_bias and change_type_map that preprocess spin and delegate to backbones; trims a training log; and adds comprehensive unit and E2E tests for spin finetuning. Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant SpinModel
participant SpinSampler as "Spin-wrapped Sampler"
participant BaseAtomicModel
participant Backbone as "Backbone Model"
User->>SpinModel: change_out_bias(merged | lazy samples)
SpinModel->>SpinModel: _get_spin_sampled_func(sampled_func)
SpinModel->>SpinSampler: call wrapped sampler (cached)
SpinSampler->>SpinSampler: apply spin preprocessing (duplicate coords/atypes, remove spin), adjust natoms, preserve extras
SpinSampler->>BaseAtomicModel: return per-frame dicts
BaseAtomicModel->>SpinSampler: inject exclusion masks and default fparam per frame (cached)
SpinModel->>Backbone: change_out_bias(spin-aware sampler)
Backbone->>Backbone: compute/apply bias and type-map changes
Backbone-->>SpinModel: updated bias/state
SpinModel-->>User: return
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
source/tests/pt/test_finetune_spin.py (1)
117-120: Make temp model cleanup failure-safe.If an assertion fails before the final
os.unlink(...), the temporary file is leaked. Wrap this section intry/finallyand close the handle explicitly.🧹 Proposed fix
- tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") - torch.jit.save(dp, tmp_model.name) - dp = DeepEval(tmp_model.name) + tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") + tmp_model_path = tmp_model.name + tmp_model.close() + try: + torch.jit.save(dp, tmp_model_path) + dp = DeepEval(tmp_model_path) + # ... existing assertions / checks ... + finally: + os.unlink(tmp_model_path) @@ - os.unlink(tmp_model.name)Also applies to: 159-159
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt/test_finetune_spin.py` around lines 117 - 120, Wrap the temporary model creation and DeepEval loading in a try/finally so the NamedTemporaryFile handle is closed and the file removed even if assertions fail: after creating tmp_model and saving with torch.jit.save (same block that constructs dp = torch.jit.script(model), tmp_model = NamedTemporaryFile(...), torch.jit.save(...), dp = DeepEval(...)), call tmp_model.close() in the try body as appropriate and ensure os.unlink(tmp_model.name) is called in the finally; apply the same pattern to the other occurrence around the DeepEval usage at the later block referenced in the comment.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pt/model/atomic_model/base_atomic_model.py`:
- Around line 187-196: The current code assumes sampled is non-empty and that
every sample lacks "fparam"; update the logic to first return or skip when
sampled is empty, then fetch default_fparam = self.get_default_fparam() (if
self.has_default_fparam() and default_fparam is not None) and iterate over each
sample in sampled, checking per-sample whether "find_fparam" or "fparam" exists
(e.g., if "find_fparam" not in sample and "fparam" not in sample) before
injecting sample["fparam"] = default_fparam.repeat(nframe, 1) using that
sample's nframe (sample["atype"].shape[0]); this ensures no index error on empty
inputs and only fills missing fparam on a per-sample basis.
In `@deepmd/pt/model/model/spin_model.py`:
- Around line 391-456: There are four Ruff issues to fix: replace unused
unpacked variables by prefixing with underscore (change occurrences like
"nframes, nloc, nnei = ..." to "nframes, nloc, _ = ..." or if nframes is unused
change to "_, nloc, nnei = ..." and for the other two unpackings change
"nframes, nloc = ..." to "_, nloc = ..." wherever they occur), and move the
multi-line exception message so it is passed as the exception constructor
argument instead of being a standalone string (i.e., change the TRY003 pattern
to raise <ExceptionType>(multi-line_message) inside the function that currently
builds the message). After applying these renames and the exception fix, run
ruff check deepmd/pt/model/model/spin_model.py to verify no lint errors remain.
---
Nitpick comments:
In `@source/tests/pt/test_finetune_spin.py`:
- Around line 117-120: Wrap the temporary model creation and DeepEval loading in
a try/finally so the NamedTemporaryFile handle is closed and the file removed
even if assertions fail: after creating tmp_model and saving with torch.jit.save
(same block that constructs dp = torch.jit.script(model), tmp_model =
NamedTemporaryFile(...), torch.jit.save(...), dp = DeepEval(...)), call
tmp_model.close() in the try body as appropriate and ensure
os.unlink(tmp_model.name) is called in the finally; apply the same pattern to
the other occurrence around the DeepEval usage at the later block referenced in
the comment.
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
deepmd/dpmodel/model/spin_model.pydeepmd/pt/model/atomic_model/base_atomic_model.pydeepmd/pt/model/atomic_model/dp_atomic_model.pydeepmd/pt/model/atomic_model/linear_atomic_model.pydeepmd/pt/model/model/spin_model.pydeepmd/pt/train/training.pysource/tests/pt/test_finetune_spin.py
Contributor
There was a problem hiding this comment.
Pull request overview
Adds spin-model finetuning support (PT + dpmodel) by ensuring finetune sampling paths apply spin preprocessing and consistently wrap samplers (exclusion types + default fparam), plus introduces new regression tests for the spin finetune workflow.
Changes:
- Add spin finetune test suite covering
change_out_bias, spin-sampler transformation, e2e finetune load, and spin+default_fparam. - Centralize “wrapped sampler” creation in
BaseAtomicModel._make_wrapped_sampler()and reuse it from PT atomic models. - Add spin-model finetune helpers (
get_spin_sampled_func,change_out_bias,change_type_map) in PT and dpmodel spin wrappers; adjust bias-change logging output.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
source/tests/pt/test_finetune_spin.py |
New tests exercising spin finetune workflows including default fparam injection. |
deepmd/pt/train/training.py |
Adjusts output-bias change log formatting (now truncates printed bias to real type-map length). |
deepmd/pt/model/model/spin_model.py |
Adds spin-aware sampler wrapper and routes change_out_bias/change_type_map through spin preprocessing/type-map expansion. |
deepmd/pt/model/atomic_model/base_atomic_model.py |
Introduces _make_wrapped_sampler() (cached sampler + exclusion types + default fparam injection). |
deepmd/pt/model/atomic_model/dp_atomic_model.py |
Reuses _make_wrapped_sampler() for stats computation (removes duplicated wrapping logic). |
deepmd/pt/model/atomic_model/linear_atomic_model.py |
Reuses _make_wrapped_sampler() for output-stat computation (removes duplicated wrapping logic). |
deepmd/dpmodel/model/spin_model.py |
Mirrors PT spin finetune helpers on the dpmodel spin wrapper. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
source/tests/pt/test_finetune_spin.py (1)
118-120: Always clean up the temp model file, including failure paths.If an assertion fails before Line 159, the temp
.pthfile is leaked. Register cleanup immediately after creation.♻️ Suggested refactor
- tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") - torch.jit.save(dp, tmp_model.name) - dp = DeepEval(tmp_model.name) + tmp_model = tempfile.NamedTemporaryFile(delete=False, suffix=".pth") + tmp_model_path = tmp_model.name + tmp_model.close() + self.addCleanup( + lambda: os.path.exists(tmp_model_path) and os.unlink(tmp_model_path) + ) + torch.jit.save(dp, tmp_model_path) + dp = DeepEval(tmp_model_path) @@ - os.unlink(tmp_model.name)Also applies to: 159-159
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt/test_finetune_spin.py` around lines 118 - 120, The temp model file created by tempfile.NamedTemporaryFile (tmp_model) is not guaranteed to be removed if an assertion fails; register cleanup immediately after creation by scheduling removal (e.g., store tmp_model.name and call os.unlink or use pytest request.addfinalizer) so the file is always deleted even on failures; update the block around torch.jit.save and DeepEval(tmp_model.name) to ensure tmp_model.name is unlinked in all paths (referencing tmp_model, tempfile.NamedTemporaryFile, torch.jit.save, and DeepEval).
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/dpmodel/model/spin_model.py`:
- Line 373: The call is using a nonexistent .atomic_model hop and will raise
AttributeError; modify the return in the function containing this line to call
_make_wrapped_sampler directly on backbone_model (i.e., replace the
.atomic_model access), so invoke
backbone_model._make_wrapped_sampler(spin_sampled_func) instead; locate the
offending line in spin_model.py where return
self.backbone_model.atomic_model._make_wrapped_sampler(spin_sampled_func)
appears (inside the sampler/wrapper setup) and change it to call
_make_wrapped_sampler on the DPAtomicModel-typed backbone_model.
---
Nitpick comments:
In `@source/tests/pt/test_finetune_spin.py`:
- Around line 118-120: The temp model file created by
tempfile.NamedTemporaryFile (tmp_model) is not guaranteed to be removed if an
assertion fails; register cleanup immediately after creation by scheduling
removal (e.g., store tmp_model.name and call os.unlink or use pytest
request.addfinalizer) so the file is always deleted even on failures; update the
block around torch.jit.save and DeepEval(tmp_model.name) to ensure
tmp_model.name is unlinked in all paths (referencing tmp_model,
tempfile.NamedTemporaryFile, torch.jit.save, and DeepEval).
ℹ️ Review info
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
deepmd/dpmodel/model/spin_model.pydeepmd/pt/model/atomic_model/dp_atomic_model.pydeepmd/pt/model/atomic_model/linear_atomic_model.pysource/tests/pt/test_finetune_spin.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/pt/model/atomic_model/linear_atomic_model.py
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/pt/model/model/spin_model.py (1)
394-415: Avoid duplicating spin-sample preprocessing logic.This block duplicates the same sample transformation already present in
compute_or_load_stat(Lines 481-503). Consider extracting a shared internal helper to prevent behavioral drift.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@deepmd/pt/model/model/spin_model.py` around lines 394 - 415, The spin-sample preprocessing in spin_sampled_func duplicates logic from compute_or_load_stat; extract a shared helper (e.g., _prepare_spin_sample or prepare_spin_sample) that takes a sample dict and returns the transformed dict by calling self.process_spin_input and handling "natoms" and passthrough keys, then replace the body of spin_sampled_func and the equivalent block in compute_or_load_stat to call this helper; reference functions: spin_sampled_func, compute_or_load_stat, process_spin_input, sampled_func when updating callers.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@deepmd/pt/model/model/spin_model.py`:
- Around line 394-415: The spin-sample preprocessing in spin_sampled_func
duplicates logic from compute_or_load_stat; extract a shared helper (e.g.,
_prepare_spin_sample or prepare_spin_sample) that takes a sample dict and
returns the transformed dict by calling self.process_spin_input and handling
"natoms" and passthrough keys, then replace the body of spin_sampled_func and
the equivalent block in compute_or_load_stat to call this helper; reference
functions: spin_sampled_func, compute_or_load_stat, process_spin_input,
sampled_func when updating callers.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5281 +/- ##
=======================================
Coverage 82.30% 82.31%
=======================================
Files 767 767
Lines 76984 77026 +42
Branches 3659 3659
=======================================
+ Hits 63360 63402 +42
Misses 12453 12453
Partials 1171 1171 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
wanghan-iapcm
requested changes
Mar 3, 2026
wanghan-iapcm
approved these changes
Mar 5, 2026
github-merge-queue
bot
removed this pull request from the merge queue due to failed status checks
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
source/tests/pt/test_finetune_spin.py (1)
316-323:test_finetune_change_type_map()never runs in this file.The test returns immediately when
self.mixed_typesis false, and the onlySpinFinetuneE2ETestsubclass here setsself.mixed_types = False. That leaves the PTchange_type_mapfinetune path introduced in this PR without an actual end-to-end check. Please add a mixed-types fixture/class or parametrize one.Also applies to: 365-406
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/pt/test_finetune_spin.py` around lines 316 - 323, The test test_finetune_change_type_map never executes because SpinFinetuneE2ETest subclasses set self.mixed_types = False; update the test harness so at least one test class/fixture runs with mixed_types = True (or parametrize the test to run both True/False). Specifically, either add a new subclass of SpinFinetuneE2ETest that sets mixed_types = True (or modify the existing test collection to parametrize mixed_types) so test_finetune_change_type_map and the other affected tests (lines ~365-406) exercise the mixed-types PT finetune path; ensure the new class/param uses the same setup/teardown and naming conventions so the test runner picks it up.source/tests/common/dpmodel/test_finetune_spin.py (1)
218-229: This assertion doesn't verify the_spinmap expansion.
DescrptSeA.change_type_map()raisesNotImplementedErrorregardless of howSpinModel.change_type_map()constructs the forwarded map, so this passes even if the wrapper never appends the virtual types or forwards the wrong order. Please assert the expanded backbone map directly, or use a remappable descriptor, so the new delegation logic is actually covered.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@source/tests/common/dpmodel/test_finetune_spin.py` around lines 218 - 229, The test test_change_type_map is currently ineffective because DescrptSeA.change_type_map always raises NotImplementedError, so replace or wrap the backbone used by self.model (SpinModel) with a spy/mocking stub that records the argument passed to SpinModel.change_type_map, then call self.model.change_type_map(new_type_map) and assert the recorded argument equals the expanded map (each original type with a corresponding "_spin" virtual type appended in correct order, e.g., ["O","O_spin","Ni","Ni_spin"]), or alternatively use a remappable descriptor backend that doesn't raise and assert its received map; update references to SpinModel.change_type_map and the backbone stub/descriptor accordingly so the delegation and suffixing are actually verified.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@deepmd/pt/model/model/spin_model.py`:
- Around line 449-458: change_type_map currently updates the backbone_model but
leaves SpinModel's spin-related metadata stale; update ntpes_real,
virtual_scale_mask, and spin_mask inside change_type_map after calling
self.backbone_model.change_type_map so the wrapper stays consistent—preferably
reload these from model_with_new_type_stat when provided, otherwise recompute
them from the updated backbone_model (the same logic used in __init__ that
process_spin_input/process_spin_output expect) to avoid mismatched type ids
after reordering or adding types.
In `@deepmd/pt/train/training.py`:
- Around line 1836-1838: The log currently flattens old_bias/new_bias before
trimming, which drops extra components for multi-output types; instead, convert
biases to numpy with to_numpy_array(old_bias)/to_numpy_array(new_bias), slice
the type axis first using [: len(model_type_map)] (because
SpinModel.get_type_map exposes real types), then reshape/flatten for
logging—apply the same change wherever the f-string referencing model_type_map,
old_bias, new_bias appears so the full per-type multi-output bias components are
preserved in the log.
---
Nitpick comments:
In `@source/tests/common/dpmodel/test_finetune_spin.py`:
- Around line 218-229: The test test_change_type_map is currently ineffective
because DescrptSeA.change_type_map always raises NotImplementedError, so replace
or wrap the backbone used by self.model (SpinModel) with a spy/mocking stub that
records the argument passed to SpinModel.change_type_map, then call
self.model.change_type_map(new_type_map) and assert the recorded argument equals
the expanded map (each original type with a corresponding "_spin" virtual type
appended in correct order, e.g., ["O","O_spin","Ni","Ni_spin"]), or
alternatively use a remappable descriptor backend that doesn't raise and assert
its received map; update references to SpinModel.change_type_map and the
backbone stub/descriptor accordingly so the delegation and suffixing are
actually verified.
In `@source/tests/pt/test_finetune_spin.py`:
- Around line 316-323: The test test_finetune_change_type_map never executes
because SpinFinetuneE2ETest subclasses set self.mixed_types = False; update the
test harness so at least one test class/fixture runs with mixed_types = True (or
parametrize the test to run both True/False). Specifically, either add a new
subclass of SpinFinetuneE2ETest that sets mixed_types = True (or modify the
existing test collection to parametrize mixed_types) so
test_finetune_change_type_map and the other affected tests (lines ~365-406)
exercise the mixed-types PT finetune path; ensure the new class/param uses the
same setup/teardown and naming conventions so the test runner picks it up.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 1ca00612-129e-4380-a405-e2a648336f16
📒 Files selected for processing (8)
deepmd/dpmodel/model/spin_model.pydeepmd/pt/model/atomic_model/base_atomic_model.pydeepmd/pt/model/atomic_model/dp_atomic_model.pydeepmd/pt/model/atomic_model/linear_atomic_model.pydeepmd/pt/model/model/spin_model.pydeepmd/pt/train/training.pysource/tests/common/dpmodel/test_finetune_spin.pysource/tests/pt/test_finetune_spin.py
🚧 Files skipped from review as they are similar to previous changes (3)
- deepmd/pt/model/atomic_model/base_atomic_model.py
- deepmd/dpmodel/model/spin_model.py
- deepmd/pt/model/atomic_model/linear_atomic_model.py
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@source/tests/common/dpmodel/test_finetune_spin.py`:
- Around line 99-107: The fixture currently sets "find_energy" and "find_fparam"
as np.float32 which hides type-sensitive bugs; change those entries in the test
fixture object to boolean sentinels (prefer np.bool_ to match production) so the
wrapped sampler receives booleans — update the dict keys "find_energy" and
"find_fparam" in test_finetune_spin.py to use np.bool_(True/False) (or
True/False) instead of np.float32(1.0)/np.float32(0.0).
- Around line 219-230: The test currently only checks that change_type_map
raises NotImplementedError and misses verifying that SpinModel expands types
with "_spin" before delegating; update test_change_type_map to spy/mock the
backbone change_type_map (e.g., replace self.model.backbone.change_type_map with
a MagicMock or patch on backbone_model.change_type_map), call
self.model.change_type_map(["O","Ni"]), assert the mock was called once with
["O_spin","Ni_spin"], and still assert that a NotImplementedError is raised (or
re-raise the backbone's NotImplementedError from the mock) to preserve the
original exception expectation; reference SpinModel.change_type_map and the
backbone_model.change_type_map when locating code to change.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: d0e69c2a-9883-4529-a6f5-a026cf11ce6b
📒 Files selected for processing (1)
source/tests/common/dpmodel/test_finetune_spin.py
njzjz
approved these changes
Mar 6, 2026
auto-merge was automatically disabled
March 6, 2026 13:34
Pull request was closed
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Tests
Refactor
Chores