feat(jax): support DPA4 training#5748
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:
📝 WalkthroughWalkthroughAdds JAX/Flax wrappers and registrations for DPA4/SeZM descriptor and fitting components, routes model construction through SeZM handling, updates force-shape and checkpoint processing, and adds conversion and backend tests. ChangesDPA4/SeZM JAX support
Estimated code review effort: 4 (Complex) | ~60 minutes 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: 3
🧹 Nitpick comments (1)
deepmd/jax/descriptor/dpa4.py (1)
216-249: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNew promotion logic appears untested by the accompanying consistency test.
_promote_trainable/_promote_trainable_lists(and therefore the whole_promote_trainable_treepath wired intoDescrptDPA4.__init__/deserialize) are gated bygetattr(module, "trainable", True)and return immediately whentrainableis falsy. The DPA4 descriptor consistency test fixture (source/tests/consistent/descriptor/test_dpa4.py) constructs the descriptor with"trainable": False, which means this newly added promotion path is never actually exercised by the new JAX consistency coverage introduced in this PR stack.Also applies to: 270-283
🤖 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/jax/descriptor/dpa4.py` around lines 216 - 249, The new trainable-promotion path in _promote_trainable, _promote_trainable_lists, and _promote_trainable_tree is skipped whenever module.trainable is false, so the current DPA4 consistency coverage does not exercise it. Update the descriptor consistency test in test_dpa4.py to include a trainable-true case (or otherwise invoke DescrptDPA4.__init__ and deserialize with trainable enabled) so the promotion logic is actually covered and validated.
🤖 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/jax/fitting/dpa4_ener.py`:
- Around line 52-54: The SeZMEnergyFittingNet.__setattr__ override is currently
a no-op because it only forwards to super without adding any behavior. Either
remove this dead override entirely, or implement the intended
trainable-parameter promotion logic for SeZMEnergyFittingNet and any related
GLUFittingNet internals so this wrapper matches the descriptor-side pattern used
by _promote_trainable_tree.
In `@deepmd/jax/model/model.py`:
- Around line 127-130: The defaults in model/model.py do not handle explicit
null values for descriptor and fitting_net, so the subsequent setdefault calls
can fail on None. Update the initialization in the model config parsing logic
around the data.setdefault usage to normalize both fields with
data.get("descriptor") or {} and data.get("fitting_net") or {} before calling
setdefault on their type keys, keeping the existing dpa4 and dpa4_ener defaults
intact.
In `@deepmd/jax/train/trainer.py`:
- Line 836: The checkpoint save path in trainer.py is dropping zero-size leaves
via _drop_zero_size_array_leaves(state), but the restart flow that rebuilds the
model and calls replace_by_pure_dict() expects those paths to exist. Update the
save/restore logic around the trainer state so zero-size entries are either
preserved in the saved state or explicitly reinserted during restore before
replace_by_pure_dict() runs, using the existing checkpoint/state handling code
paths in trainer.py.
---
Nitpick comments:
In `@deepmd/jax/descriptor/dpa4.py`:
- Around line 216-249: The new trainable-promotion path in _promote_trainable,
_promote_trainable_lists, and _promote_trainable_tree is skipped whenever
module.trainable is false, so the current DPA4 consistency coverage does not
exercise it. Update the descriptor consistency test in test_dpa4.py to include a
trainable-true case (or otherwise invoke DescrptDPA4.__init__ and deserialize
with trainable enabled) so the promotion logic is actually covered and
validated.
🪄 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
Run ID: 5aa70f59-efa8-4ae4-ac9f-cf147f7b123b
📒 Files selected for processing (9)
deepmd/jax/descriptor/__init__.pydeepmd/jax/descriptor/dpa4.pydeepmd/jax/fitting/__init__.pydeepmd/jax/fitting/dpa4_ener.pydeepmd/jax/model/ener_model.pydeepmd/jax/model/model.pydeepmd/jax/train/trainer.pysource/tests/consistent/descriptor/test_dpa4.pysource/tests/consistent/fitting/test_dpa4_ener.py
|
Pushed follow-up commit |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #5748 +/- ##
==========================================
- Coverage 79.69% 71.52% -8.18%
==========================================
Files 1020 996 -24
Lines 116359 110632 -5727
Branches 4303 3532 -771
==========================================
- Hits 92736 79129 -13607
- Misses 22076 30707 +8631
+ Partials 1547 796 -751 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Pushed follow-up commit 68cb12e for the remaining review feedback. Changes:
Validation:
Coding agent: Codex |
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Requesting changes for the blocking DPA4 fitting-trainability correctness issue described inline.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding a second blocking DPA4 correctness finding; the existing changes-requested state remains applicable.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved DPA4 exclusion-routing correctness finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved DPA4 descriptor-freeze correctness finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Additional interoperability finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved charge/spin-conditioning training finding.
Preserve DPA4 freeze policies, support PT SeZM conversion, and apply zero-size filtering to every JAX checkpoint writer. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
There was a problem hiding this comment.
🧹 Nitpick comments (1)
source/tests/jax/test_dpa4.py (1)
88-102: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCover mixed per-layer trainability.
The scalar frozen case would also pass the prior all-or-nothing policy. Add a
[False, True]round trip and assert the policy persists with optimizer-visible parameters.Proposed coverage
+def test_mixed_fitting_trainability_survives_round_trip() -> None: + fitting = SeZMEnergyFittingNet( + ntypes=2, + dim_descrpt=4, + neuron=[4], + trainable=[False, True], + precision="float64", + mixed_types=True, + seed=20260712, + ) + restored = SeZMEnergyFittingNet.deserialize(fitting.serialize()) + + assert restored.nets[0].trainable == [False, True] + assert len(nnx.to_flat_state(nnx.state(restored, nnx.Param))) > 0As per coding guidelines, use
pytestfor the single new test case.🤖 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/jax/test_dpa4.py` around lines 88 - 102, Extend test_frozen_fitting_stays_frozen_after_conversion_round_trip with a mixed per-layer trainability case using [False, True], then serialize and deserialize the fitting network. Assert restored.nets[0].trainable preserves [False, True] and verify nnx.to_flat_state(nnx.state(restored, nnx.Param)) contains the expected optimizer-visible parameters; run this as a single pytest test case.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@source/tests/jax/test_dpa4.py`:
- Around line 88-102: Extend
test_frozen_fitting_stays_frozen_after_conversion_round_trip with a mixed
per-layer trainability case using [False, True], then serialize and deserialize
the fitting network. Assert restored.nets[0].trainable preserves [False, True]
and verify nnx.to_flat_state(nnx.state(restored, nnx.Param)) contains the
expected optimizer-visible parameters; run this as a single pytest test case.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 8151ce1f-2748-4c4c-9dd2-fcb561d862d7
📒 Files selected for processing (11)
deepmd/dpmodel/fitting/dpa4_ener.pydeepmd/jax/descriptor/dpa4.pydeepmd/jax/fitting/dpa4_ener.pydeepmd/jax/model/base_model.pydeepmd/jax/model/model.pydeepmd/jax/train/trainer.pydeepmd/jax/utils/serialization.pysource/tests/jax/test_dpa4.pysource/tests/jax/test_dpa4_conversion.pysource/tests/jax/test_model_factory.pysource/tests/jax/test_training.py
💤 Files with no reviewable changes (1)
- deepmd/jax/fitting/dpa4_ener.py
🚧 Files skipped from review as they are similar to previous changes (2)
- deepmd/jax/model/model.py
- deepmd/jax/descriptor/dpa4.py
Resolve DPA4 JAX training conflicts while retaining the latest shared multi-task and Array API Strict coverage from master. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
for more information, see https://pre-commit.ci
|
Resolved the merge conflicts with the latest master and pushed merge commit 3727480. pre-commit.ci subsequently applied its automatic formatting fix in 180fff0. The resolution retains the JAX DPA4 shape-handling tests from this PR and the latest multi-task/stat-merging and Array API Strict coverage from master. Validation:
Coding agent: Codex |
Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Summary
dpa4_enerfitting.model.type: dpa4/sezmthrough the JAX model factory.source/tests/consistent/descriptor/test_dpa4.pyandsource/tests/consistent/fitting/test_dpa4_ener.pyfiles.Benchmark
1000-step DPA4 water benchmark, batch size 1,
srun --gres=gpu:1 dp, GPU: NVIDIA GeForce RTX 5090.dp --jax train input_jax.json --skip-neighbor-statdp --tf2 train input_tf2.json --skip-neighbor-statPT/pt-expt are not included in the final comparison per follow-up scope.
Validation
ruff format .ruff check .python -m py_compile deepmd/jax/descriptor/dpa4.py deepmd/jax/fitting/dpa4_ener.py deepmd/jax/model/model.py deepmd/jax/train/trainer.py source/tests/consistent/descriptor/test_dpa4.py source/tests/consistent/fitting/test_dpa4_ener.pydp --jax train /tmp/deepmd_dpa4_tiny_1vnghlkx/input_jax.json --skip-neighbor-stat -o /tmp/deepmd_dpa4_tiny_1vnghlkx/out_jax.jsonsrun --gres=gpu:1 dp --jax train /tmp/deepmd_dpa4_bench_1000/input_jax.json --skip-neighbor-stat -o /tmp/deepmd_dpa4_bench_1000/out_jax.jsonNote: collecting the existing DPA4 consistent test module in this local environment segfaults while importing the existing PT DPA4/Triton path before
-k jaxselection is applied, so I validated the new JAX path with the smoke train and benchmark above.Summary by CodeRabbit
BaseModeldeserialization to accept SeZM/DPA4 aliases with automatic descriptor/fitting defaults.