feat(tf2): support DPA4 training#5749
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| from deepmd.tf2.descriptor.base_descriptor import ( | ||
| BaseDescriptor, | ||
| ) | ||
| from deepmd.tf2.utils import exclude_mask as _tf2_exclude_mask # noqa: F401 |
| BaseDescriptor, | ||
| ) | ||
| from deepmd.tf2.utils import exclude_mask as _tf2_exclude_mask # noqa: F401 | ||
| from deepmd.tf2.utils import network as _tf2_network # noqa: F401 |
| from deepmd.tf2.fitting.base_fitting import ( | ||
| BaseFitting, | ||
| ) | ||
| from deepmd.tf2.utils import network as _tf2_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 (2)
📝 WalkthroughWalkthroughThis PR adds TF2 support for DPA4/SeZM descriptors and fitting networks, registers the new model names, extends TensorFlow-backed array handling, updates trainer force-shape postprocessing, and adds consistency-test coverage. ChangesTF2 DPA4/SeZM Integration
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.61% 79.48% -0.13%
==========================================
Files 1014 1016 +2
Lines 115360 115802 +442
Branches 4274 4274
==========================================
+ Hits 91843 92050 +207
- Misses 21976 22205 +229
- Partials 1541 1547 +6 ☔ 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 |
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
New Features
Bug Fixes
forceto match label tensor shapes when dimensions are compatible.Tests