feat(jax): support DPA4 training#5748
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| ClassVar, | ||
| ) | ||
|
|
||
| import deepmd.jax.utils.network as _jax_network # noqa: F401 |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (1)
💤 Files with no reviewable changes (1)
📝 WalkthroughWalkthroughAdds JAX/Flax wrappers and registrations for DPA4/SeZM descriptor and fitting components, routes model construction through a SeZM model factory, updates trainer force handling and checkpoint serialization, restores omitted zero-size checkpoint leaves, and extends consistency tests to run JAX backends. ChangesDPA4/SeZM JAX Support
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant get_model
participant get_sezm_model
participant get_standard_model
participant EnergyModel
Caller->>get_model: SeZM or DPA4 configuration
get_model->>get_sezm_model: route configuration
get_sezm_model->>get_sezm_model: validate and normalize fields
get_sezm_model->>get_standard_model: pass normalized configuration
get_standard_model->>EnergyModel: construct registered model
EnergyModel-->>Caller: model instance
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: 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.61% 79.49% -0.12%
==========================================
Files 1014 1017 +3
Lines 115360 116009 +649
Branches 4274 4274
==========================================
+ Hits 91843 92223 +380
- Misses 21976 22238 +262
- Partials 1541 1548 +7 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
|
|
||
| import numpy as np | ||
|
|
||
| import deepmd.jax.utils.exclude_mask as _jax_exclude_mask # noqa: F401 |
| ) | ||
|
|
||
|
|
||
| _TRAINABLE_ATTRS: dict[str, tuple[str, ...]] = { |
There was a problem hiding this comment.
Non-blocking (correctness — silently frozen params): promotion here is strictly allowlist-gated — _promote_trainable_tree promotes only the attributes named in _TRAINABLE_ATTRS / _TRAINABLE_LIST_ATTRS; anything else stays an ArrayAPIVariable (not an nnx.Param), so nnx.Optimizer(..., wrt=nnx.Param) / nnx.grad never update it — it is silently frozen during training. A few trainable dpmodel params appear to be missing from the maps:
FrameContract.weight/FrameExpand.weight— trainable indeepmd/dpmodel/descriptor/dpa4_nn/grid_net.py(pt registers them asnn.Parameter), reachable whennode_wise_so3/message_node_so3instantiate amode=="cross"grid net. Neither class is a key in the maps.SeZMInteractionBlock.adam_ffn_layer_scales— a trainable LayerScale list whenlayer_scale=True(dpa4_nn/block.py), also listed indpa4_nn/lora.py's unfreeze registry. The maps already cover the analogousSO2Convolution.adam_so2_layer_scales, but not this one.
These are opt-in (non-default) features, so a stock DPA4 config is unaffected — but when a user enables them, training silently skips those weights, and the forward-only consistency tests cannot catch it (a leaf gives identical output whether it is a Param or a Variable). Consider driving promotion from the dpmodel trainable flag / mirroring pt's nn.Parameter set rather than a hand-maintained allowlist (which will keep silently dropping future params), and adding a test that asserts the differentiated leaf set equals the trainable dpmodel param set.
| return _NO_ZERO_SIZE_LEAF | ||
|
|
||
|
|
||
| def _restore_missing_zero_size_leaves(template: Any, restored: Any) -> Any: |
There was a problem hiding this comment.
Non-blocking (test coverage): the new stateful glue in this PR ships without dedicated tests — only the two forward-consistency tests were added and source/tests/jax/ was not touched. Per CLAUDE.md ("UTs should cover all code paths, including both branches of boolean conditions"), worth covering:
- the zero-size drop -> restore round-trip:
_drop_zero_size_array_leaves(jax/train/trainer.py) paired with_restore_missing_zero_size_leaves/_zero_size_subtreehere. This mutates checkpoint state on save and reconstructs it on load, so a mismatch would silently corrupt a restored model — and it is trivially testable as a pure round-trip assertion. get_sezm_model's guard branches (spin / bridging / lora / use_compile / preset_out_bias, plus the descriptor/fitting/exclude_types mismatches) —source/tests/jax/test_model_factory.pyis the established home for routing/validation tests._match_label_shapesand the force-reshape branch in_evaluate_model_dict(both the reshape and no-op branches).
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
New Features
Bug Fixes
Tests