fix(pt): calibrate DPA4 bias and support PyTorch 2.13 compile - #5910
Conversation
Use the production SeZM forward topology for residual statistics so inner potentials, exclusions, and spin expansion are represented consistently. Align minimum-distance filtering across statistics and distributed training.
PyTorch 2.13 still routes symbolic-denominator divisibility through Inductor's non-factoring Mod, the regression introduced in 2.12, so the existing repair applies unchanged and 2.13 needs no workaround of its own. Collapse the version gates into SUPPORTED_COMPILE_TORCH and _DIVISIBILITY_REGRESSION_TORCH so the runtime check, the patch selector and the compile-parity test skip all read one declaration.
|
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 atomic-only SeZM prediction for statistic-based bias adjustment, minimum pair-distance filtering across statistics and distributed training, renames the analytical bridging potential to ChangesAtomic statistics and SeZM prediction
Minimum-distance frame filtering
InnerPotential migration
PyTorch compile compatibility
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: 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
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
deepmd/pt/utils/stat.py (1)
267-291: 🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick winUse the forwarded
atypeshape for the auto-batch size hint.
_compute_model_predict()forwardsmodel_coordandmodel_atype, butmodel_forward_auto_batch_size()still passessystem["atype"].shape[-1]toAutoBatchSize. For models that populatemodel_atypedifferently fromsystem["atype"], the hint can under/overestimate the actual atom count and waste batching iterations. Useatype.shape[-1]for this hint.🤖 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/pt/utils/stat.py` around lines 267 - 291, Update model_forward_auto_batch_size() to pass the forwarded atype.shape[-1] value to AutoBatchSize.execute_all instead of system["atype"].shape[-1]. Keep the existing model_forward arguments and batching flow unchanged.
🧹 Nitpick comments (1)
deepmd/pt/train/training.py (1)
774-792: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate
min_pair_distparsing.This re-derives the exact same value(s) already computed earlier in
__init__for theDataRequirementItem/single_model_statwiring (single-task: line ~502-504; multi-task: line ~578-582, per model_key). Two independent copies of the same parsing logic risk silently diverging if one is edited without the other.♻️ Suggested consolidation
Compute
self.min_pair_dist/self.has_min_pair_filteronce, earlier in__init__(before the single/multi-task data-loading blocks), and reuseself.min_pair_distat the two existing call sites instead of the localmin_pair_distvariables.🤖 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/pt/train/training.py` around lines 774 - 792, Consolidate min_pair_dist parsing in __init__ by computing self.min_pair_dist and self.has_min_pair_filter once before the single-task and multi-task data-loading blocks. Remove the duplicate parsing shown here, and update the existing DataRequirementItem/single_model_stat wiring to reuse self.min_pair_dist instead of local min_pair_dist values for both task modes.
🤖 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/train/training.py`:
- Around line 2456-2481: Update all_ranks_have_valid_frames to select the
collective tensor device based on the active device and supported backend
capabilities, not an exact dist.get_backend() == "nccl" comparison. Preserve
CUDA/NCCL synchronization when available while retaining CPU/Gloo fallback for
CPU-only environments and ensure dist.all_reduce uses the selected device.
In `@deepmd/pt/utils/compile_compat.py`:
- Around line 61-64: Update _DIVISIBILITY_REGRESSION_TORCH to remove the (2, 13)
entry, leaving only releases that still require
patch_inductor_symbolic_divisibility. Adjust the adjacent comments to state that
PyTorch 2.13 and later are fixed and must not receive the workaround.
In `@deepmd/pt/utils/stat.py`:
- Around line 319-332: Update _reduce_model_prediction to handle frames with
zero unmasked atoms before dividing in the intensive path. Ensure zero-count
frames produce a finite, defined result without inf or nan, while preserving the
existing averaging behavior for frames with at least one unmasked atom.
In `@source/tests/pt/model/test_sezm_spin_model.py`:
- Line 497: Update the skip gate for test_compile_matches_eager to use the
shared SUPPORTED_COMPILE_TORCH versions, importing that symbol and deriving the
skip condition and reason from it. Ensure Torch 2.11, 2.12, and 2.13 are
eligible unless the spin compile path is explicitly intended to support only
2.11.
---
Outside diff comments:
In `@deepmd/pt/utils/stat.py`:
- Around line 267-291: Update model_forward_auto_batch_size() to pass the
forwarded atype.shape[-1] value to AutoBatchSize.execute_all instead of
system["atype"].shape[-1]. Keep the existing model_forward arguments and
batching flow unchanged.
---
Nitpick comments:
In `@deepmd/pt/train/training.py`:
- Around line 774-792: Consolidate min_pair_dist parsing in __init__ by
computing self.min_pair_dist and self.has_min_pair_filter once before the
single-task and multi-task data-loading blocks. Remove the duplicate parsing
shown here, and update the existing DataRequirementItem/single_model_stat wiring
to reuse self.min_pair_dist instead of local min_pair_dist values for both task
modes.
🪄 Autofix (Beta)
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: 099ead31-de6c-43d0-8d12-b21758bb8912
📒 Files selected for processing (13)
deepmd/pt/model/atomic_model/base_atomic_model.pydeepmd/pt/model/model/make_model.pydeepmd/pt/model/model/model.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/model/model/sezm_property_model.pydeepmd/pt/model/model/sezm_spin_model.pydeepmd/pt/train/training.pydeepmd/pt/utils/compile_compat.pydeepmd/pt/utils/stat.pydeepmd/utils/argcheck.pysource/tests/pt/model/test_sezm_model.pysource/tests/pt/model/test_sezm_spin_model.pysource/tests/pt/test_training.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5910 +/- ##
==========================================
- Coverage 79.21% 79.00% -0.21%
==========================================
Files 1069 1069
Lines 124070 124232 +162
Branches 4522 4527 +5
==========================================
- Hits 98278 98146 -132
- Misses 24171 24468 +297
+ Partials 1621 1618 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Requesting changes for the inline compatibility regression.
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: 1
🤖 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/atomic_model/inner_potential.py`:
- Around line 269-270: Update InnerPotentialAtomicModel deserialization and its
registry compatibility so serialized models with type "inter_potential" resolve
to the current "inner_potential" implementation. Preserve normal
"inner_potential" lookup while accepting the legacy type during composition
deserialization, using an alias or schema upgrade alongside the existing
BaseAtomicModel registration.
🪄 Autofix (Beta)
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: 8640731c-57ad-4f23-aecf-7441bfece09a
📒 Files selected for processing (19)
deepmd/dpmodel/atomic_model/inner_potential.pydeepmd/dpmodel/model/dp_linear_model.pydeepmd/dpmodel/model/model.pydeepmd/pt/train/training.pydeepmd/pt/utils/stat.pydeepmd/pt_expt/model/get_model.pydeepmd/pt_expt/utils/__init__.pydeepmd/pt_expt/utils/inner_potential.pydeepmd/pt_expt/utils/serialization.pysource/api_cc/tests/test_deeppot_dpa4_zbl_ptexpt.ccsource/api_cc/tests/test_deepspin_dpa4_zbl_ptexpt.ccsource/lmp/tests/test_lammps_dpa4_zbl_pt2.pysource/tests/common/dpmodel/test_inner_potential.pysource/tests/common/dpmodel/test_zbl_bridging.pysource/tests/infer/gen_dpa4_spin_zbl.pysource/tests/infer/gen_dpa4_zbl.pysource/tests/pt/model/test_sezm_spin_model.pysource/tests/pt/test_training.pysource/tests/pt_expt/model/test_zbl_bridging.py
🚧 Files skipped from review as they are similar to previous changes (4)
- source/tests/pt/test_training.py
- source/tests/pt/model/test_sezm_spin_model.py
- deepmd/pt/utils/stat.py
- deepmd/pt/train/training.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
The coordinated InnerPotential rename is internally consistent across PT, dpmodel, pt_expt, serialization identifiers, and tests. Given the stated development-stage API policy for DPA4/SeZM, I have no remaining blocking concern.
Focused validation: 37 inner-potential and ZBL-bridging tests passed.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
The direction of the bias fix is right, and I checked that rather than assuming it: routing the analytical term through the prediction removes a genuine prior double count, where ZBL was absorbed into the fitted bias and then added again at inference. Reasoning through test_zbl_change_out_bias_is_invariant_for_self_labels, self-generated labels include ZBL while the old stats path excluded it, so the bias would shift and the invariance assertion should fail pre-fix - that makes it a real regression test. I did not execute it against the parent, so treat that as reasoning rather than a measurement.
What I did measure: the pre-existing shared-path suites still pass against this branch - test_atomic_model_atomic_stat.py, test_atomic_model_global_stat.py and test_polar_atomic_model_stat.py all green, and the three failures in test_change_bias.py reproduce identically on the parent commit from an unrelated libcupti problem. So "Tests were not run locally, per request" is a process gap here, not a hidden regression. Worth running them before merge all the same, since make_model.py, base_atomic_model.py and stat.py are on the path of every pt model, not just SeZM.
Eight comments inline. The one on base_atomic_model.py is about the API shape; the others are individual defects.
One finding cannot be anchored because the file is not in this diff. deepmd/pt/entrypoints/main.py calls make_stat_input(...) with no min_pair_dist:
deepmd-kit/deepmd/pt/entrypoints/main.py
Lines 620 to 626 in 17a57b4
So dp train now calibrates on filtered data while dp change-bias calibrates on unfiltered data, and the same dataset plus config yields two different biases. It is structural rather than a one-line oversight: change_bias() only receives the frozen model's model_params, never the training input config, so there is no min_pair_dist value available to thread through. Worth deciding explicitly whether change-bias should grow the option or whether the filter belongs somewhere both entry points can see.
On atomic_output_only: it exists because predict_atomic_outputs_for_stat calls forward_common under torch.no_grad(), and forward_common computes force and virial through autograd, so the derivative work has to be suppressed. That is the same design choice as the model_forward argument, one layer down - both follow from reusing the whole model forward for a job that only needs atomic outputs. If the model instead exposed just its analytical atomic contribution, the atomic model's existing predictor plus that term would give the complete output with neither the flag nor the new argument. The flag also has a side effect worth a second look: it gates the dens branch as well (get_active_mode() == "dens" and not atomic_output_only), so a dens-mode model computes its statistics through the ener edge schema - the same "incomplete output path" divergence this PR set out to remove. If dens models are not meant to reach bias calibration, raising would be clearer than silently switching schema.
Smaller notes, not worth their own threads:
-
SUPPORTED_COMPILE_TORCHand_DIVISIBILITY_REGRESSION_TORCHare now independent tuples. Previously the== (2, 12)test was bound to the same literal as the gate, so they could not drift; a future(2, 14)added to the first alone would compile without the divisibility patch and hitCantSplit. Deriving the second as "everything supported except (2, 11)" would be failure-proof. -
The
model_forward is Nonefallback branch inchange_out_bias- the one every non-SeZM model takes - is covered only by the pre-existing suites noted above, none of which are in this diff. -
The PR bundles bias calibration with PyTorch 2.13 compile support. The commits are mostly separable (
ba3705b8abias,c9cad8c99compile), so splitting would make each independently revertable. -
test_change_out_bias_is_invariant_for_self_labelsand its ZBL sibling assert atatol=1e-12. Worth confirming that tolerance was chosen from observed residuals rather than inherited.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/model/model/make_model.py`:
- Around line 374-379: Update the model-level compute_or_load_stat flow and its
underlying output-stat computation so statistics are generated through
predict_atomic_outputs_for_stat rather than directly through the atomic model’s
compute_or_load_out_stat. Ensure model overrides such as SEZM’s predictor
participate in initial calibration and retain the model-level residual-bias
update path, while preserving existing atomic-model behavior otherwise.
In `@deepmd/pt/train/training.py`:
- Around line 793-811: In deepmd/pt/train/training.py:793-811 and 2340-2382,
change the multi-task retry-budget flow so each step uses one attempt limit
shared across ranks for the selected task, rather than independently keyed
per-task limits. Update the initialization around _training_batch_attempts and
its use in _next_training_batch() so the limit is derived and broadcast
consistently for the current task before all_ranks_have_valid_frames() retries,
preventing any rank from exhausting its budget before peers complete the
collective.
🪄 Autofix (Beta)
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: 866658be-6d04-471a-803b-99d3549daeff
📒 Files selected for processing (9)
deepmd/pt/model/atomic_model/base_atomic_model.pydeepmd/pt/model/model/make_model.pydeepmd/pt/model/model/sezm_model.pydeepmd/pt/train/training.pydeepmd/pt/utils/compile_compat.pydeepmd/pt/utils/stat.pysource/tests/pt/model/test_sezm_model.pysource/tests/pt/model/test_sezm_spin_model.pysource/tests/pt/test_training.py
🚧 Files skipped from review as they are similar to previous changes (3)
- source/tests/pt/model/test_sezm_spin_model.py
- deepmd/pt/utils/compile_compat.py
- deepmd/pt/model/model/sezm_model.py
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Went back through all seven threads against HEAD rather than the replies, and they hold up. The atomic change_out_bias signature is back to the three-argument contract that is byte-identical across the other backends, with the residual calibration moved to the model layer -- and routing it through predict_atomic_outputs_for_stat with return rather than falling through to the atomic call is right, since _store_out_stat(..., add=True) accumulates and running both would double-count. The StopIteration restart is restored with an empty-loader guard, the empty-sample case raises, and the compile allowlist no longer claims CI coverage it does not have.
On the intensive-reduction thread, I accept your reasoning and am not pressing it. A silent NaN bias is worse than an abort, and the frame it rejects -- one with no unmasked atoms -- is pathological rather than routine.
The bug-fix reflection also comes out clean. The reason existing tests missed this is specific: pt/test_finetune.py::test_finetune_change_out_bias does cover a ZBL model, but it calls model.atomic_model.change_out_bias(...), entering at the atomic layer and bypassing the model layer where the analytical term is added -- so the incomplete path was invisible to it by construction, and every other bias test uses non-bridged models where the atomic output already is the complete output. test_change_out_bias_is_invariant_for_self_labels closes exactly that gap and genuinely fails pre-fix: it labels the sample with the model's own energy, so any term missing from the statistics path shows up directly as a non-zero delta bias. Pinning the side conditions too (no compile, no derivatives, training mode and caches restored, no leaked grads) is a nice touch.
Two follow-ups below, both non-blocking and both about coverage rather than correctness.
62aa235
Resolutions: keep the capability-based _needs_with_comm_artifact and the _compose_bridging owner (renamed to upstream's InnerPotentialAtomicModel, deepmodeling#5910); drop this branch's pt extended-atype ZBL fix in favor of upstream's equivalent descriptor_atype at the relocated injection site.
…odels (deepmodeling#5966) Close deepmodeling#5927 — resolved **as designed**, with documentation and semantic pin tests instead of a behavior change. ## The definitions (verified against the code) The model energy decomposes as `E = E_model + E_bias`, where `E_model` is everything the model computes (learned network, analytical ZBL term, or a `linear_ener` combination) and `E_bias` is the per-type output bias. - **`set-by-statistic`** assigns `E_bias` directly: user values, or the per-type least-squares statistic of the **raw labels**. It never calls a model forward (`compute_output_stats` with `model_forward=None`, stored with `add=False`) — it is independent of `E_model` **by definition**, ignoring a trained network and the analytical ZBL term alike. - **`change-by-statistic`** fits the residual of the labels against the **complete model prediction** (the bridged predictor since deepmodeling#5910) and adds the delta to the existing bias. Under these definitions the "double count" described in the issue is not a bridging bug: `set` uniformly ignores all of `E_model` for every model kind. A bridged model after `set` carries no compensation for the mean ZBL contribution — exactly as a trained plain model after `set` carries no compensation for its network output. The two modes "disagree" because they are defined to answer different questions; a self-consistent calibration is `change`'s job. Verified conformance of the `linear_ener` composition path in both backends: children compute no output statistics (`compute_or_load_out_stat=False`); the composition level performs one `set` fit on the raw labels. Also verified: the ZBL term contributes exactly zero for isolated atoms, so its own statistics are trivially `bias = 0`. ## Changes - `doc/model/change-bias.md`: precise definitions of the two modes, and the guidance that calibrating a bridged (or any nonzero-`E_model`) model self-consistently requires `change-by-statistic`. - `doc/model/dpa4.md`: note in the ZBL section. - Semantic pin tests (pt `SeZMModel` + dpmodel `LinearEnergyAtomicModel` composition): `set-by-statistic` equals the raw-label least-squares fit exactly — guarding against a future "fix" that would subtract the analytical term and silently create a third, model-dependent mode. ## Known limitations - No behavior change anywhere; the pin tests cover `model.change_out_bias` (the `dp change-bias --mode set` and finetune routes) and the dpmodel composition out-stat; the pt training-init chain was verified in-session and funnels into the same pinned branch. - The spin variants share the machinery but have no dedicated pin. - Fact worth knowing when reading the docs: `InnerPotential` adds the full ZBL over the whole cutoff (not only below `bridging_r_outer`), so the label-side ZBL at equilibrium geometries is small but not strictly zero; the docs state the offset plainly. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Documentation** * Clarified how `set` and `change` statistic modes handle model energy, analytical contributions, labels, residuals, and existing bias. * Documented the interaction between output-bias calibration and ZBL bridging, including guidance for self-consistent calibration. * **Tests** * Added regression coverage confirming that `set-by-statistic` calibration uses raw energy labels independently of learned and ZBL model contributions. * Verified calibration correctly replaces seeded bias and remains consistent when repeated. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>
Summary
min_pair_distfiltering and multitask data handling.torch.compilecompatibility range to PyTorch 2.13 and apply the Inductor symbolic-divisibility workaround to the affected releases.Root cause
Bias calibration previously used an incomplete output path, so model-level contributions and output reduction semantics could be omitted from the fitted bias. Separately, PyTorch 2.13 was rejected before compilation, even though it requires the same Inductor compatibility handling as PyTorch 2.12.
Validation
git diff --checkpassed.Summary by CodeRabbit
New Features
Bug Fixes