Skip to content

feat(pt/dpmodel): add sequential_update for dpa3 #5355

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
iProzd wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat(pt/dpmodel): add sequential_update for dpa3 #5355

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions deepmd/dpmodel/descriptor/dpa3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -170,6 +170,12 @@ class RepFlowArgs:
In the dynamic selection case, neighbor-scale normalization will use `e_sel / sel_reduce_factor`
or `a_sel / sel_reduce_factor` instead of the raw `e_sel` or `a_sel` values,
accommodating larger selection numbers.
sequential_update : bool, optional
Whether to use sequential update mode within each repflow layer.
When True, updates are applied sequentially: edge self → angle self (using updated edge)
→ edge angle (using updated angle) → node (using final edge),
instead of the default parallel mode where all updates use original embeddings.
Currently only supports ``update_style='res_residual'``.
Copy link

Copilot AI Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The sequential_update parameter docstring doesn’t mention that it requires update_angle=True, but RepFlowArgs.__init__ now raises a ValueError when sequential_update is True and update_angle is False. Please reflect this constraint in the docstring to prevent confusing configuration errors.

Suggested change
Currently only supports ``update_style='res_residual'``.
Currently only supports ``update_style='res_residual'`` and requires ``update_angle=True``;
otherwise, a ``ValueError`` will be raised during initialization.

Copilot uses AI. Check for mistakes.
"""

def __init__(
Expand Down Expand Up @@ -201,6 +207,7 @@ def __init__(
use_exp_switch: bool = False,
use_dynamic_sel: bool = False,
sel_reduce_factor: float = 10.0,
sequential_update: bool = False,
) -> None:
self.n_dim = n_dim
self.e_dim = e_dim
Expand Down Expand Up @@ -231,6 +238,15 @@ def __init__(
self.use_exp_switch = use_exp_switch
self.use_dynamic_sel = use_dynamic_sel
self.sel_reduce_factor = sel_reduce_factor
self.sequential_update = sequential_update
if self.sequential_update:
if self.update_style != "res_residual":
raise ValueError(
"sequential_update only supports update_style='res_residual', "
f"got '{self.update_style}'!"
)
if not self.update_angle:
raise ValueError("sequential_update requires update_angle=True!")

def __getitem__(self, key: str) -> Any:
if hasattr(self, key):
Expand Down Expand Up @@ -266,6 +282,7 @@ def serialize(self) -> dict:
"use_exp_switch": self.use_exp_switch,
"use_dynamic_sel": self.use_dynamic_sel,
"sel_reduce_factor": self.sel_reduce_factor,
"sequential_update": self.sequential_update,
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Mar 30, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major

Bump the DPA3 serialization version for the new repflow_args field.

Line 285 adds a new key to the persisted repflow_args payload, but DescrptDPA3.serialize() still advertises @version: 2. A pre-change v2 reader will accept the file and then die on RepFlowArgs(**data) with an unexpected sequential_update kwarg instead of rejecting the newer format cleanly.

Suggested fix 
-            "@version": 2,
+            "@version": 3,
-        check_version_compatibility(version, 2, 1)
+        check_version_compatibility(version, 3, 1)
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@deepmd/dpmodel/descriptor/dpa3.py` at line 285, The serialized payload for
repflow_args now includes a new key "sequential_update" but
DescrptDPA3.serialize() still emits "@version: 2", causing old readers to
attempt RepFlowArgs(**data) and crash; update DescrptDPA3.serialize() to bump
the serialization version (e.g., to 3) so that readers can detect the new
format, and ensure the version constant/metadata that serialize() returns is
updated accordingly (refer to DescrptDPA3.serialize and the repflow_args payload
construction to locate the change).

}

@classmethod
Expand Down Expand Up @@ -404,6 +421,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
use_exp_switch=self.repflow_args.use_exp_switch,
use_dynamic_sel=self.repflow_args.use_dynamic_sel,
sel_reduce_factor=self.repflow_args.sel_reduce_factor,
sequential_update=self.repflow_args.sequential_update,
use_loc_mapping=use_loc_mapping,
exclude_types=exclude_types,
env_protection=env_protection,
Expand Down
Loading
Loading