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% 78.49% -1.21%
==========================================
Files 1020 1055 +35
Lines 116359 121227 +4868
Branches 4303 4385 +82
==========================================
+ Hits 92736 95155 +2419
- Misses 22076 24520 +2444
- Partials 1547 1552 +5 ☔ 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
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Both threads I raised are addressed and verified against the current HEAD (fix in 68cb12e):
- Silently-frozen params (
dpa4.py): the promotion maps now includeFrameContract.weight,FrameExpand.weight, andSeZMInteractionBlock.adam_ffn_layer_scales, and the newtest_optional_dpa4_weights_are_jax_parametersenableslayer_scale+message_node_so3and asserts these leaves areArrayAPIParam— a test that fails on the pre-fix maps, so it genuinely locks the fix (with a companion frozen-descriptor test for the other side). - Test coverage (
serialization.py/ trainer glue):source/tests/jax/now covers the zero-size checkpoint drop/restore round-trip, theget_sezm_modelguard/mismatch paths, and_match_label_shapes/_evaluate_model_dictbranches.
Approving on that basis.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The PT DPA4 fitting freeze policy is still not preserved across its serialization round trip.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
JAX DPA4 currently routes through a fixed-width standard neighbor-list path, violating DPA4's conservative full-cutoff edge semantics.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
JAX DPA4 currently accepts two default training options, random_gamma=True and use_amp=True, while silently executing neither behavior. The random augmentation needs a JAX-safe PRNG path (or an explicit unsupported-option guard), and AMP must be implemented or surfaced as unsupported instead of being a no-op.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
The default JAX DPA4 training entrypoint fails before model construction because neighbor-stat preprocessing cannot dispatch the model-level DPA4 aliases.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
One additional non-blocking test-isolation issue.
Preserve PT DPA4 fitting freeze policies across serialization, reject inert JAX random-gamma and AMP options, register model aliases for neighbor-stat preprocessing, and keep JAX-only tests collectable without PyTorch. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Use the full extended-atom capacity at the JAX input boundary for DPA4 so the configured sel cannot truncate in-cutoff neighbors, and cover the edge-count regression. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Possible reviewers based on changed lines, exact file history, and exact-file review history:
No review request was made automatically. Coding agent: Codex |
Keep cross-backend DPA4 consistency coverage within the JAX-supported feature subset by disabling the backend-specific AMP policy. Coding-Agent: Codex Codex-Version: codex-cli 0.144.4 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.