feat(tf2): support DPA4 training#5749
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:
📝 WalkthroughWalkthroughThis PR adds TF2 support for DPA4/SeZM descriptors and fitting networks, registers model and component aliases, adapts PT checkpoints, extends TensorFlow-backed array operations, normalizes force shapes, and adds backend and training tests. ChangesTF2 DPA4/SeZM integration
Estimated code review effort: 5 (Critical) | ~120 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: 2
🧹 Nitpick comments (1)
deepmd/tf2/common.py (1)
362-387: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
__getattribute__allocates a newseton every attribute access.This override intercepts every attribute read on the module, and for each read it calls
_tf2_array_variable_attr_names()/_tf2_array_variable_list_attr_names(), both of which materialize a freshset(...)from the class-level tuple. On hot descriptor/fitting paths this adds a per-access allocation. Consider a cheaper membership check against the raw tuple (or a cached frozenset) to avoid rebuilding the set on every access.♻️ Cheaper membership check
def __getattribute__(self, name: str) -> Any: if not name.startswith("_tf2_"): - array_attrs = object.__getattribute__( - self, - "_tf2_array_variable_attr_names", - )() - if name in array_attrs: + if name in object.__getattribute__( + self, "_tf2_array_variable_attrs", () + ):🤖 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/tf2/common.py` around lines 362 - 387, The __getattribute__ override in common.py is doing extra work on every attribute read by calling _tf2_array_variable_attr_names() and _tf2_array_variable_list_attr_names(), which rebuild sets repeatedly. Update __getattribute__ to use a cheaper membership path for the array/list attribute names, such as checking the underlying tuple directly or reusing a cached frozenset, while keeping the existing storage-name lookup and to_tensorflow_array conversion behavior unchanged.
🤖 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/array_api.py`:
- Around line 285-296: The TensorFlow branch in array_api.py is overwriting
prefilled -inf entries because tf.maximum(x_tensor, reduced) replaces
empty-segment sentinels with dtype minimum values. Update the
unsorted-segment-max handling in this branch so empty segments remain -inf,
matching the behavior expected by segment.py and the other backends; use the
existing x_tensor, indices_tensor, values_tensor, and reduced flow, but avoid
applying a blanket maximum that changes sentinel slots.
In `@deepmd/tf2/train/trainer.py`:
- Around line 1490-1493: The shape-iteration guard in the trainer’s
rank-checking helper only handles TypeError, but tf.TensorShape(None) can also
raise ValueError during tracing. Update the try/except around iter(shape) to
return None for both exception types in the same helper path so unknown-rank
shapes are handled safely.
---
Nitpick comments:
In `@deepmd/tf2/common.py`:
- Around line 362-387: The __getattribute__ override in common.py is doing extra
work on every attribute read by calling _tf2_array_variable_attr_names() and
_tf2_array_variable_list_attr_names(), which rebuild sets repeatedly. Update
__getattribute__ to use a cheaper membership path for the array/list attribute
names, such as checking the underlying tuple directly or reusing a cached
frozenset, while keeping the existing storage-name lookup and
to_tensorflow_array conversion behavior unchanged.
🪄 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: 76f9add6-fdb8-4b7b-8908-ab6d1ed973d8
📒 Files selected for processing (11)
deepmd/dpmodel/array_api.pydeepmd/tf2/common.pydeepmd/tf2/descriptor/__init__.pydeepmd/tf2/descriptor/dpa4.pydeepmd/tf2/fitting/__init__.pydeepmd/tf2/fitting/dpa4_ener.pydeepmd/tf2/model/ener_model.pydeepmd/tf2/model/model.pydeepmd/tf2/train/trainer.pysource/tests/consistent/descriptor/test_dpa4.pysource/tests/consistent/fitting/test_dpa4_ener.py
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5749 +/- ##
==========================================
- Coverage 79.69% 79.61% -0.09%
==========================================
Files 1020 1022 +2
Lines 116359 116820 +461
Branches 4303 4307 +4
==========================================
+ Hits 92736 93006 +270
- Misses 22076 22272 +196
+ Partials 1547 1542 -5 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
deepmd/tf2/descriptor/se_atten_v2.py (1)
17-21: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the extra refresh call here.
TF2Modulealready refreshes_refresh_tf2_trackable_lists()fordeepmd.tf2.descriptor.se_atten_v2, so this second call is redundant and can be dropped.🤖 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/tf2/descriptor/se_atten_v2.py` around lines 17 - 21, The DescrptSeAttenV2.deserialize method is calling _refresh_tf2_trackable_lists() twice because TF2Module already performs that refresh for deepmd.tf2.descriptor.se_atten_v2. Remove the explicit refresh call from DescrptSeAttenV2.deserialize and keep the deserialization flow limited to delegating to DescrptSeAttenV2DP.deserialize.__func__(cls, data) and returning the object.
🤖 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 `@deepmd/tf2/descriptor/se_atten_v2.py`:
- Around line 17-21: The DescrptSeAttenV2.deserialize method is calling
_refresh_tf2_trackable_lists() twice because TF2Module already performs that
refresh for deepmd.tf2.descriptor.se_atten_v2. Remove the explicit refresh call
from DescrptSeAttenV2.deserialize and keep the deserialization flow limited to
delegating to DescrptSeAttenV2DP.deserialize.__func__(cls, data) and returning
the object.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: e4a0ef32-c262-4207-8de7-5e2559fd9048
📒 Files selected for processing (6)
deepmd/dpmodel/array_api.pydeepmd/tf2/common.pydeepmd/tf2/descriptor/se_atten_v2.pydeepmd/tf2/train/trainer.pysource/tests/consistent/descriptor/test_dpa4.pysource/tests/consistent/fitting/test_dpa4_ener.py
🚧 Files skipped from review as they are similar to previous changes (3)
- deepmd/tf2/train/trainer.py
- source/tests/consistent/descriptor/test_dpa4.py
- deepmd/tf2/common.py
|
Pushed follow-ups |
Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
|
Pushed follow-up commit 4328329 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 dynamic-shape DPA4 force-loss correctness issue described inline.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved graph-safety correctness finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved fitting-trainability correctness finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved frozen-descriptor tracking correctness finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved normalized-exclusion routing correctness finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved default-random-gamma training semantics finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved DPA4 property-fitting factory compatibility finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved DPA4 parameter-promotion resident-memory finding.
OutisLi
left a comment
OutisLi
left a comment
There was a problem hiding this comment.
Adding the approved PT DPA4 .pt to TF2 conversion integration finding.
Make DPA4 dynamic-shape training graph-safe, preserve frozen state, and support the schema-approved factory and PT conversion paths. 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.
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/descriptor/dpa4_nn/so2.py`:
- Around line 747-762: Update the shape validation around the visible
x_local/radial_feat checks and _project_radial() to use runtime shape values
rather than assuming concrete symbolic rank metadata. Explicitly reject inputs
with rank below 3 before indexing dimensions, compare dimensions through runtime
shape handling, and reshape radial_feat using its runtime batch dimension or -1.
🪄 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: 741870f5-a817-453c-80e8-1695eec6fccb
📒 Files selected for processing (14)
deepmd/dpmodel/descriptor/dpa4.pydeepmd/dpmodel/descriptor/dpa4_nn/so2.pydeepmd/dpmodel/fitting/dpa4_ener.pydeepmd/tf2/common.pydeepmd/tf2/descriptor/dpa4.pydeepmd/tf2/descriptor/se_atten_v2.pydeepmd/tf2/fitting/dpa4_ener.pydeepmd/tf2/model/base_model.pydeepmd/tf2/model/model.pydeepmd/tf2/train/trainer.pysource/tests/tf2/test_dpa4.pysource/tests/tf2/test_dpa4_conversion.pysource/tests/tf2/test_model_factory.pysource/tests/tf2/test_training.py
💤 Files with no reviewable changes (1)
- deepmd/tf2/fitting/dpa4_ener.py
🚧 Files skipped from review as they are similar to previous changes (1)
- deepmd/tf2/common.py
Validate mixer ranks and dimensions at runtime, propagate the proven rank into ndtensorflow, and avoid static batch dimensions in reshapes. Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
Resolve DPA4 TF2 training conflicts while retaining the latest 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 a6e0e1c. pre-commit.ci subsequently applied its automatic formatting fix in 198a832. The resolution retains the TF2 DPA4 coverage from this PR together with the latest Array API Strict coverage from master. Validation:
Coding agent: Codex |
| class DynamicRadialDegreeMixer(DynamicRadialDegreeMixerDP): | ||
| """TF2 mixer with runtime shape checks for generalized edge counts.""" | ||
|
|
||
| def call(self, x_local: Any, radial_feat: Any) -> Any: |
Coding-Agent: Codex Codex-Version: codex-cli 0.144.1 Model: gpt-5.6-sol Reasoning-Effort: xhigh
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 TF2 model factory.training.enable_compile.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-stattraining.enable_compile=true; includes TF/XLA/PTX compile; batch 1 avg 149.4788 s/step.PT/pt-expt are not included in the final comparison per follow-up scope.
Validation
ruff format .ruff check .python -m py_compile deepmd/dpmodel/array_api.py deepmd/tf2/common.py deepmd/tf2/descriptor/dpa4.py deepmd/tf2/fitting/dpa4_ener.py deepmd/tf2/model/model.py deepmd/tf2/train/trainer.py source/tests/consistent/descriptor/test_dpa4.py source/tests/consistent/fitting/test_dpa4_ener.pydp --tf2 train /tmp/deepmd_dpa4_tiny_1vnghlkx/input_tf2.json --skip-neighbor-stat -o /tmp/deepmd_dpa4_tiny_1vnghlkx/out_tf2.jsonwithtraining.enable_compile=true; log confirmedCompiled cluster using XLA!and savedmodel.ckpt.tf2/model.ckpt-1.srun --gres=gpu:1 dp --tf2 train /tmp/deepmd_dpa4_bench_1000/input_tf2.json --skip-neighbor-stat -o /tmp/deepmd_dpa4_bench_1000/out_tf2.json.Note: collecting the existing DPA4 consistent test module in this local environment segfaults while importing the existing PT DPA4/Triton path before backend-specific selection is applied, so I validated the new TF2 path with the checkpoint smoke, XLA train smoke, and benchmark above.
Summary by CodeRabbit
forceshape to labelforceduring training, when possible.