Skip to content

feat(pd): add add_chg_spin_ebd parameter to DescrptDPA3 #5333

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
HydrogenSulfate wants to merge 1 commit into
base: master
Choose a base branch
from
Open

feat(pd): add add_chg_spin_ebd parameter to DescrptDPA3 #5333

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
4 changes: 4 additions & 0 deletions deepmd/pd/model/atomic_model/dp_atomic_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,9 @@ def __init__(
self.sel = self.descriptor.get_sel()
self.fitting_net = fitting
super().init_out_stat()
self.add_chg_spin_ebd: bool = getattr(
self.descriptor, "add_chg_spin_ebd", False
)
self.enable_eval_descriptor_hook = False
self.enable_eval_fitting_last_layer_hook = False
self.eval_descriptor_list = []
Expand Down Expand Up @@ -329,6 +332,7 @@ def forward_atomic(
nlist,
mapping=mapping,
comm_dict=comm_dict,
fparam=fparam if self.add_chg_spin_ebd else None,
)
assert descriptor is not None
if self.enable_eval_descriptor_hook:
Expand Down
1 change: 1 addition & 0 deletions deepmd/pd/model/descriptor/dpa1.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -625,6 +625,7 @@ def forward(
nlist: paddle.Tensor,
mapping: paddle.Tensor | None = None,
comm_dict: list[paddle.Tensor] | None = None,
fparam: paddle.Tensor | None = None,
) -> paddle.Tensor:
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

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

forward() is annotated as returning paddle.Tensor, but it actually returns a 5-tuple. Since this method signature was updated, please fix the return type annotation to reflect the tuple return type.

Suggested change
) -> paddle.Tensor:
) -> tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]:

Copilot uses AI. Check for mistakes.
"""Compute the descriptor.

Expand Down
1 change: 1 addition & 0 deletions deepmd/pd/model/descriptor/dpa2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -732,6 +732,7 @@ def forward(
nlist: paddle.Tensor,
mapping: paddle.Tensor | None = None,
comm_dict: list[paddle.Tensor] | None = None,
fparam: paddle.Tensor | None = None,
) -> paddle.Tensor:
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

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

forward() is annotated as returning paddle.Tensor, but it actually returns a 5-tuple (g1, rot_mat, g2, h2, sw). Since this method signature was updated, please fix the return type annotation to reflect the tuple return type.

Suggested change
) -> paddle.Tensor:
) -> tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor, paddle.Tensor]:

Copilot uses AI. Check for mistakes.
"""Compute the descriptor.

Expand Down
83 changes: 78 additions & 5 deletions deepmd/pd/model/descriptor/dpa3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
UpdateSel,
)
from deepmd.pd.utils.utils import (
ActivationFn,
to_numpy_array,
)
from deepmd.utils.data_system import (
Expand Down Expand Up @@ -120,6 +121,7 @@ def __init__(
use_tebd_bias: bool = False,
use_loc_mapping: bool = True,
type_map: list[str] | None = None,
add_chg_spin_ebd: bool = False,
) -> None:
super().__init__()

Expand Down Expand Up @@ -174,6 +176,7 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
)

self.use_econf_tebd = use_econf_tebd
self.add_chg_spin_ebd = add_chg_spin_ebd
self.use_loc_mapping = use_loc_mapping
self.use_tebd_bias = use_tebd_bias
self.type_map = type_map
Expand All @@ -196,6 +199,34 @@ def init_subclass_params(sub_data: dict | Any, sub_class: type) -> Any:
self.concat_output_tebd = concat_output_tebd
self.precision = precision
self.prec = PRECISION_DICT[self.precision]

if self.add_chg_spin_ebd:
self.act = ActivationFn(activation_function)
# -100 ~ 100 is a conservative bound
self.chg_embedding = TypeEmbedNet(
200,
self.tebd_dim,
precision=precision,
seed=child_seed(seed, 3),
)
# 100 is a conservative upper bound
self.spin_embedding = TypeEmbedNet(
100,
Comment on lines +207 to +214
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

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

The charge/spin embedding table sizes look off-by-one given TypeEmbedNet uses a dedicated last index for padding. With TypeEmbedNet(200, ...) and charge = fparam[:,0].int64 + 100, charge==100 maps to index 200 (padding), not a learnable embedding, contradicting the “-100 ~ 100” comment. Consider either (a) resizing to cover the intended inclusive bounds (e.g., 201/101) or (b) tightening the documented/validated allowed ranges and raising a clear error when out of range.

Suggested change
200,
self.tebd_dim,
precision=precision,
seed=child_seed(seed, 3),
)
# 100 is a conservative upper bound
self.spin_embedding = TypeEmbedNet(
100,
201,
self.tebd_dim,
precision=precision,
seed=child_seed(seed, 3),
)
# 100 is a conservative upper bound
self.spin_embedding = TypeEmbedNet(
101,

Copilot uses AI. Check for mistakes.
self.tebd_dim,
precision=precision,
seed=child_seed(seed, 4),
)
self.mix_cs_mlp = MLPLayer(
2 * self.tebd_dim,
self.tebd_dim,
precision=precision,
seed=child_seed(seed, 5),
)
else:
self.chg_embedding = None
self.spin_embedding = None
self.mix_cs_mlp = None

self.exclude_types = exclude_types
self.env_protection = env_protection
self.trainable = trainable
Expand Down Expand Up @@ -433,9 +464,14 @@ def serialize(self) -> dict:
"use_econf_tebd": self.use_econf_tebd,
"use_tebd_bias": self.use_tebd_bias,
"use_loc_mapping": self.use_loc_mapping,
"add_chg_spin_ebd": self.add_chg_spin_ebd,
"type_map": self.type_map,
"type_embedding": self.type_embedding.embedding.serialize(),
}
if self.add_chg_spin_ebd:
data["chg_embedding"] = self.chg_embedding.embedding.serialize()
data["spin_embedding"] = self.spin_embedding.embedding.serialize()
data["mix_cs_mlp"] = self.mix_cs_mlp.serialize()
Comment on lines +467 to +474
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Mar 23, 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 descriptor schema version for the new serialized fields.

This payload now includes a new constructor kwarg plus three extra parameter blobs, but it still advertises @version: 2. Older DescrptDPA3.deserialize() implementations that accept v2 won't reject these checkpoints early; they'll fail later on unexpected fields instead.

Suggested version bump 
-            "@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/pd/model/descriptor/dpa3.py` around lines 467 - 474, The serialized
payload now includes the new constructor kwarg add_chg_spin_ebd and three extra
blobs (chg_embedding, spin_embedding, mix_cs_mlp) but the descriptor still
advertises `@version`: 2; update the schema version in the class's serialize
output to a new integer (e.g., 3) so DescrptDPA3.deserialize() can reject older
v2 checkpoints; locate the serialize method that builds the dict with
"add_chg_spin_ebd", "type_map", and "type_embedding" (and the conditional
chg/spin/mix_cs_mlp entries) and increment the version field, and ensure
DescrptDPA3.deserialize() handles or rejects the new version accordingly.

repflow_variable = {
"edge_embd": repflows.edge_embd.serialize(),
"angle_embd": repflows.angle_embd.serialize(),
Expand All @@ -462,12 +498,24 @@ def deserialize(cls, data: dict) -> "DescrptDPA3":
data.pop("type")
repflow_variable = data.pop("repflow_variable").copy()
type_embedding = data.pop("type_embedding")
chg_embedding = data.pop("chg_embedding", None)
spin_embedding = data.pop("spin_embedding", None)
mix_cs_mlp = data.pop("mix_cs_mlp", None)
data["repflow"] = RepFlowArgs(**data.pop("repflow_args"))
obj = cls(**data)
obj.type_embedding.embedding = TypeEmbedNetConsistent.deserialize(
type_embedding
)

if obj.add_chg_spin_ebd and chg_embedding is not None:
obj.chg_embedding.embedding = TypeEmbedNetConsistent.deserialize(
chg_embedding
)
obj.spin_embedding.embedding = TypeEmbedNetConsistent.deserialize(
spin_embedding
)
obj.mix_cs_mlp = MLPLayer.deserialize(mix_cs_mlp)
Comment on lines +501 to +517
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

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

deserialize() only guards on chg_embedding is not None before deserializing all three add_chg_spin components. If the serialized dict is partially missing spin_embedding or mix_cs_mlp, this will fail later or leave the object in an inconsistent state. Consider checking that all required keys are present when add_chg_spin_ebd is true, and raising a clear error if any are missing.

Copilot uses AI. Check for mistakes.

def t_cvt(xx: Any) -> paddle.Tensor:
return paddle.to_tensor(xx, dtype=obj.repflows.prec, place=env.DEVICE)

Expand All @@ -493,7 +541,14 @@ def forward(
nlist: paddle.Tensor,
mapping: paddle.Tensor | None = None,
comm_dict: list[paddle.Tensor] | None = None,
) -> paddle.Tensor:
fparam: paddle.Tensor | None = None,
) -> tuple[
paddle.Tensor,
paddle.Tensor | None,
paddle.Tensor | None,
paddle.Tensor | None,
paddle.Tensor | None,
]:
"""Compute the descriptor.

Parameters
Expand Down Expand Up @@ -536,6 +591,20 @@ def forward(
node_ebd_ext = self.type_embedding(extended_atype[:, :nloc])
else:
node_ebd_ext = self.type_embedding(extended_atype)

if self.add_chg_spin_ebd:
assert fparam is not None
assert self.chg_embedding is not None
assert self.spin_embedding is not None
charge = fparam[:, 0].to(dtype=paddle.int64) + 100
spin = fparam[:, 1].to(dtype=paddle.int64)
chg_ebd = self.chg_embedding(charge)
spin_ebd = self.spin_embedding(spin)
sys_cs_embd = self.act(
self.mix_cs_mlp(paddle.concat([chg_ebd, spin_ebd], axis=-1))
)
Comment on lines +595 to +605
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

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

Using assert for required runtime inputs (fparam and embedding modules) is fragile because asserts can be disabled with Python optimizations and they don’t provide actionable error messages. Prefer explicit validation (e.g., raise ValueError with expected shape nf x 2 and valid ranges) and also validate mix_cs_mlp is present when add_chg_spin_ebd is enabled.

Copilot uses AI. Check for mistakes.
node_ebd_ext = node_ebd_ext + sys_cs_embd.unsqueeze(1)

node_ebd_inp = node_ebd_ext[:, :nloc, :]
# repflows
node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows(
Expand All @@ -550,10 +619,14 @@ def forward(
node_ebd = paddle.concat([node_ebd, node_ebd_inp], axis=-1)
return (
node_ebd.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION),
rot_mat.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION),
edge_ebd.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION),
h2.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION),
sw.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION),
rot_mat.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION)
if rot_mat is not None
else None,
edge_ebd.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION)
if edge_ebd is not None
else None,
h2.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION) if h2 is not None else None,
sw.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION) if sw is not None else None,
)

@classmethod
Expand Down
1 change: 1 addition & 0 deletions deepmd/pd/model/descriptor/se_a.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,6 +288,7 @@ def forward(
nlist: paddle.Tensor,
mapping: paddle.Tensor | None = None,
comm_dict: list[paddle.Tensor] | None = None,
fparam: paddle.Tensor | None = None,
) -> paddle.Tensor:
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

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

forward() is annotated as returning paddle.Tensor, but it actually returns a 5-tuple (descriptor, rot_mat, g2, h2, sw). Since this method signature was updated, please fix the return type annotation to match the actual return value (see se_atten.forward for an example).

Suggested change
) -> paddle.Tensor:
) -> tuple[paddle.Tensor, paddle.Tensor, paddle.Tensor | None, paddle.Tensor | None, paddle.Tensor]:

Copilot uses AI. Check for mistakes.
"""Compute the descriptor.

Expand Down
1 change: 1 addition & 0 deletions deepmd/pd/model/descriptor/se_t_tebd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -436,6 +436,7 @@ def forward(
nlist: paddle.Tensor,
mapping: paddle.Tensor | None = None,
comm_dict: list[paddle.Tensor] | None = None,
fparam: paddle.Tensor | None = None,
) -> paddle.Tensor:
Copy link

Copilot AI Mar 23, 2026

Choose a reason for hiding this comment

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

forward() is annotated as returning paddle.Tensor, but it actually returns a 5-tuple (g1, None, None, None, sw). Since this method signature was updated, please fix the return type annotation to match the actual return value for better static checking and consistency with other descriptors.

Suggested change
) -> paddle.Tensor:
) -> tuple[paddle.Tensor, None, None, None, paddle.Tensor]:

Copilot uses AI. Check for mistakes.
"""Compute the descriptor.

Expand Down
44 changes: 31 additions & 13 deletions source/tests/pd/model/test_dpa3.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@ def test_consistency(
nme,
prec,
ect,
add_chg_spin,
) in itertools.product(
[True, False], # update_angle
["res_residual"], # update_style
Expand All @@ -65,15 +66,13 @@ def test_consistency(
[1, 2], # n_multi_edge_message
["float64"], # precision
[False], # use_econf_tebd
[False, True], # add_chg_spin_ebd
):
dtype = PRECISION_DICT[prec]
rtol, atol = get_tols(prec)
if prec == "float64":
atol = 1e-8 # marginal GPU test cases...
coord_ext = np.concatenate([self.coord_ext[:1], self.coord_ext[:1]], axis=0)
atype_ext = np.concatenate([self.atype_ext[:1], self.atype_ext[:1]], axis=0)
nlist = np.concatenate([self.nlist[:1], self.nlist[:1]], axis=0)
mapping = np.concatenate([self.mapping[:1], self.mapping[:1]], axis=0)

repflow = RepFlowArgs(
n_dim=20,
e_dim=10,
Expand Down Expand Up @@ -105,24 +104,37 @@ def test_consistency(
precision=prec,
use_econf_tebd=ect,
type_map=["O", "H"] if ect else None,
add_chg_spin_ebd=add_chg_spin,
seed=GLOBAL_SEED,
).to(env.DEVICE)

dd0.repflows.mean = paddle.to_tensor(davg, dtype=dtype, place=env.DEVICE)
dd0.repflows.stddev = paddle.to_tensor(dstd, dtype=dtype, place=env.DEVICE)

# Prepare fparam if needed
fparam = None
fparam_np = None
if add_chg_spin:
fparam = paddle.to_tensor(
[[5, 1]], dtype=dtype, place=env.DEVICE
).expand(nf, -1)
fparam_np = np.array([[5, 1]], dtype=np.float64).repeat(nf, axis=0)

rd0, _, _, _, _ = dd0(
paddle.to_tensor(coord_ext, dtype=dtype, place=env.DEVICE),
paddle.to_tensor(atype_ext, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(nlist, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(mapping, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(self.coord_ext, dtype=dtype, place=env.DEVICE),
paddle.to_tensor(self.atype_ext, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(self.nlist, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(self.mapping, dtype=paddle.int64, place=env.DEVICE),
fparam=fparam,
)
# serialization
dd1 = DescrptDPA3.deserialize(dd0.serialize())
rd1, _, _, _, _ = dd1(
paddle.to_tensor(coord_ext, dtype=dtype, place=env.DEVICE),
paddle.to_tensor(atype_ext, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(nlist, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(mapping, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(self.coord_ext, dtype=dtype, place=env.DEVICE),
paddle.to_tensor(self.atype_ext, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(self.nlist, dtype=paddle.int64, place=env.DEVICE),
paddle.to_tensor(self.mapping, dtype=paddle.int64, place=env.DEVICE),
fparam=fparam,
)
np.testing.assert_allclose(
rd0.numpy(),
Expand All @@ -132,7 +144,13 @@ def test_consistency(
)
# dp impl
dd2 = DPDescrptDPA3.deserialize(dd0.serialize())
rd2, _, _, _, _ = dd2.call(coord_ext, atype_ext, nlist, mapping)
rd2, _, _, _, _ = dd2.call(
self.coord_ext,
self.atype_ext,
self.nlist,
self.mapping,
fparam=fparam_np,
)
np.testing.assert_allclose(
rd0.numpy(),
rd2,
Expand Down
Loading