Skip to content

feat(pt/dpmodel): add use_default_pf #5356

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 3 commits into
base: master
Choose a base branch
from
Open

feat(pt/dpmodel): add use_default_pf #5356

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
e5f7ce5
feat(pt/dpmodel): add use_default_pf
iProzd Mar 30, 2026
28e6fdd
fix ut
iProzd Mar 30, 2026
7d8a307
fix ut
iProzd Mar 31, 2026
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
15 changes: 12 additions & 3 deletions deepmd/dpmodel/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,9 @@ class EnergyLoss(Loss):
The prefactor of generalized force loss at the end of the training.
numb_generalized_coord : int
The dimension of generalized coordinates.
use_default_pf : bool
If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided.
This allows using the prefactor force loss (pf) without requiring atom_pref.npy files.
use_huber : bool
Enables Huber loss calculation for energy/force/virial terms with user-defined threshold delta (D).
The loss function smoothly transitions between L2 and L1 loss:
Expand Down Expand Up @@ -110,6 +113,7 @@ def __init__(
huber_delta: float = 0.01,
loss_func: str = "mse",
f_use_norm: bool = False,
use_default_pf: bool = False,
**kwargs: Any,
) -> None:
# Validate loss_func
Expand Down Expand Up @@ -149,6 +153,7 @@ def __init__(
self.use_huber = use_huber
self.huber_delta = huber_delta
self.f_use_norm = f_use_norm
self.use_default_pf = use_default_pf
if self.f_use_norm and not (self.use_huber or self.loss_func == "mae"):
raise RuntimeError(
"f_use_norm can only be True when use_huber or loss_func='mae'."
Expand Down Expand Up @@ -182,7 +187,9 @@ def call(
find_force = label_dict["find_force"]
find_virial = label_dict["find_virial"]
find_atom_ener = label_dict["find_atom_ener"]
find_atom_pref = label_dict["find_atom_pref"]
find_atom_pref = (
label_dict["find_atom_pref"] if not self.use_default_pf else 1.0
)
Comment on lines +190 to +192
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 use_default_pf behavior change in EnergyLoss.call() (overriding find_atom_pref) isn’t covered by the existing dpmodel loss unit tests (e.g., source/tests/common/dpmodel/test_loss_ener.py). Please add a small test case asserting that with find_atom_pref=0.0 and use_default_pf=True, the pf contribution/metrics are computed, and that with use_default_pf=False they remain suppressed.

Copilot uses AI. Check for mistakes.
xp = array_api_compat.array_namespace(
energy,
force,
Expand Down Expand Up @@ -477,6 +484,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
must=False,
high_prec=False,
repeat=3,
default=1.0,
)
)
if self.has_gf > 0:
Expand Down Expand Up @@ -512,7 +520,7 @@ def serialize(self) -> dict:
"""
return {
"@class": "EnergyLoss",
"@version": 2,
"@version": 3,
"starter_learning_rate": self.starter_learning_rate,
"start_pref_e": self.start_pref_e,
"limit_pref_e": self.limit_pref_e,
Expand All @@ -533,6 +541,7 @@ def serialize(self) -> dict:
"huber_delta": self.huber_delta,
"loss_func": self.loss_func,
"f_use_norm": self.f_use_norm,
"use_default_pf": self.use_default_pf,
}

@classmethod
Expand All @@ -550,6 +559,6 @@ def deserialize(cls, data: dict) -> "Loss":
The deserialized loss module
"""
data = data.copy()
check_version_compatibility(data.pop("@version"), 2, 1)
check_version_compatibility(data.pop("@version"), 3, 1)
data.pop("@class")
return cls(**data)
17 changes: 13 additions & 4 deletions deepmd/pt/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ def __init__(
loss_func: str = "mse",
inference: bool = False,
use_huber: bool = False,
use_default_pf: bool = False,
f_use_norm: bool = False,
huber_delta: float = 0.01,
**kwargs: Any,
Expand Down Expand Up @@ -103,6 +104,9 @@ def __init__(
MAE loss is less sensitive to outliers compared to MSE loss.
inference : bool
If true, it will output all losses found in output, ignoring the pre-factors.
use_default_pf : bool
If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided.
This allows using the prefactor force loss (pf) without requiring atom_pref.npy files.
use_huber : bool
Enables Huber loss calculation for energy/force/virial terms with user-defined threshold delta (D).
The loss function smoothly transitions between L2 and L1 loss:
Expand Down Expand Up @@ -147,6 +151,7 @@ def __init__(
self.limit_pref_pf = limit_pref_pf
self.start_pref_gf = start_pref_gf
self.limit_pref_gf = limit_pref_gf
self.use_default_pf = use_default_pf
self.relative_f = relative_f
self.enable_atom_ener_coeff = enable_atom_ener_coeff
self.numb_generalized_coord = numb_generalized_coord
Expand Down Expand Up @@ -357,7 +362,9 @@ def forward(

if self.has_pf and "atom_pref" in label:
atom_pref = label["atom_pref"]
find_atom_pref = label.get("find_atom_pref", 0.0)
find_atom_pref = (
label.get("find_atom_pref", 0.0) if not self.use_default_pf else 1.0
)
pref_pf = pref_pf * find_atom_pref
atom_pref_reshape = atom_pref.reshape(-1)

Expand Down Expand Up @@ -514,7 +521,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
high_prec=True,
)
)
if self.has_f:
if self.has_f or self.has_pf or self.relative_f is not None or self.has_gf:
label_requirement.append(
DataRequirementItem(
"force",
Expand Down Expand Up @@ -553,6 +560,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
must=False,
high_prec=False,
repeat=3,
default=1.0,
)
)
if self.has_gf > 0:
Expand Down Expand Up @@ -588,7 +596,7 @@ def serialize(self) -> dict:
"""
return {
"@class": "EnergyLoss",
"@version": 2,
"@version": 3,
"starter_learning_rate": self.starter_learning_rate,
"start_pref_e": self.start_pref_e,
"limit_pref_e": self.limit_pref_e,
Expand All @@ -609,6 +617,7 @@ def serialize(self) -> dict:
"huber_delta": self.huber_delta,
"loss_func": self.loss_func,
"f_use_norm": self.f_use_norm,
"use_default_pf": self.use_default_pf,
}

@classmethod
Expand All @@ -626,7 +635,7 @@ def deserialize(cls, data: dict) -> "TaskLoss":
The deserialized loss module
"""
data = data.copy()
check_version_compatibility(data.pop("@version"), 2, 1)
check_version_compatibility(data.pop("@version"), 3, 1)
data.pop("@class")
return cls(**data)

Expand Down
12 changes: 12 additions & 0 deletions deepmd/utils/argcheck.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3189,6 +3189,11 @@ def loss_ener() -> list[Argument]:
"atomic prefactor force", label="atom_pref", abbr="pf"
)
doc_limit_pref_pf = limit_pref("atomic prefactor force")
doc_use_default_pf = (
"If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided. "
"This allows using the prefactor force loss (pf) without requiring atom_pref.npy files in training data. "
"When atom_pref.npy is provided, it will be used as-is regardless of this setting."
Comment on lines 3190 to +3195
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.

use_default_pf is now validated/documented for loss/ener, but TF/Paddle energy losses do not implement this behavior (they accept **kwargs and will ignore it). Consider documenting that this flag is only effective for PyTorch/DPModel (or add an explicit warning/error in non-supporting backends) to avoid a silent no-op configuration. Also, the existing start_pref_pf help text still states atom_pref.npy must be provided when prefactors are non-zero; with use_default_pf=true that statement is no longer universally true and may need to be adjusted.

Suggested change
)
doc_limit_pref_pf = limit_pref("atomic prefactor force")
doc_use_default_pf = (
"If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided. "
"This allows using the prefactor force loss (pf) without requiring atom_pref.npy files in training data. "
"When atom_pref.npy is provided, it will be used as-is regardless of this setting."
) + (
" When `use_default_pf` is true, `atom_pref.npy` is not required even when the "
"prefactor is non-zero; a default per-atom prefactor of 1.0 will be used when "
"`atom_pref.npy` is missing. When `use_default_pf` is false, `atom_pref.npy` "
"must still be provided whenever the prefactor is non-zero."
)
doc_limit_pref_pf = limit_pref("atomic prefactor force")
doc_use_default_pf = (
"If true, use default atom_pref of 1.0 for all atoms when atom_pref data is not provided. "
"This allows using the prefactor force loss (pf) without requiring atom_pref.npy files in training data. "
"When atom_pref.npy is provided, it will be used as-is regardless of this setting. "
"Note: at present this option is only effective for the PyTorch/DPModel backend; "
"other backends (e.g. TensorFlow and Paddle) accept this flag but ignore it."

Copilot uses AI. Check for mistakes.
)
doc_start_pref_gf = start_pref("generalized force", label="drdq", abbr="gf")
doc_limit_pref_gf = limit_pref("generalized force")
doc_numb_generalized_coord = "The dimension of generalized coordinates. Required when generalized force loss is used."
Expand Down Expand Up @@ -3299,6 +3304,13 @@ def loss_ener() -> list[Argument]:
default=0.00,
doc=doc_limit_pref_pf,
),
Argument(
"use_default_pf",
bool,
optional=True,
default=False,
doc=doc_use_default_pf,
),
Comment on lines +3307 to +3313
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

use_default_pf is accepted here but can be a silent no-op on Paddle backend.

This option is now valid in schema, but the Paddle training path (deepmd/pd/train/training.py Line 1316-Line 1329) forwards loss_params into a loss constructor that does not consume use_default_pf as an active field, so users can set it without effect. Please either implement it in the Paddle loss path or explicitly reject it there to avoid backend-inconsistent behavior.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@deepmd/utils/argcheck.py` around lines 3307 - 3313, The schema now accepts
use_default_pf but the Paddle training path silently ignores it; update the
Paddle loss handling in deepmd/pd/train/training.py so that when loss_params
(the dict forwarded into the loss constructor) contains "use_default_pf" it is
either honored by passing it into the Paddle loss implementation or rejected
up-front. Specifically, either (A) extend the Paddle loss constructor(s)
referenced where loss_params is passed (look for the loss creation block around
the code forwarding loss_params at lines ~1316-1329) to accept and act on
use_default_pf, or (B) validate and remove/raise on presence of "use_default_pf"
before forwarding (e.g., check loss_params.get("use_default_pf") and raise a
clear error), and ensure the change references the same loss_params variable and
the Paddle training module functions so behavior is consistent with the other
backends.

Argument("relative_f", [float, None], optional=True, doc=doc_relative_f),
Argument(
"enable_atom_ener_coeff",
Expand Down
16 changes: 16 additions & 0 deletions doc/model/train-se-a-mask.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,22 @@ And the `loss` section in the training input script should be set as follows.
}
```

If `atom_pref.npy` is not provided in the training data, one can set `use_default_pf` to `true` to use a default atom preference of 1.0 for all atoms. This allows using the prefactor force loss (`pf` loss) without requiring `atom_pref.npy` files. When `atom_pref.npy` is provided, it will be used as-is regardless of this setting.
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 new paragraph implies use_default_pf works in general training, but the implementation appears to exist only in the PyTorch/DPModel losses (TF/Paddle EnerStdLoss accept **kwargs and will silently ignore use_default_pf). Please clarify the backend scope here (e.g., explicitly state “PyTorch/DP backend only”) or implement equivalent behavior for the other backends to avoid a configuration option that is accepted but has no effect.

Suggested change
If `atom_pref.npy` is not provided in the training data, one can set `use_default_pf` to `true` to use a default atom preference of 1.0 for all atoms. This allows using the prefactor force loss (`pf` loss) without requiring `atom_pref.npy` files. When `atom_pref.npy` is provided, it will be used as-is regardless of this setting.
If `atom_pref.npy` is not provided in the training data and you are using the PyTorch/DP backend, you can set `use_default_pf` to `true` to use a default atom preference of 1.0 for all atoms. This allows using the prefactor force loss (`pf` loss) without requiring `atom_pref.npy` files. When `atom_pref.npy` is provided, it will be used as-is regardless of this setting. For TensorFlow/Paddle backends (for example, when using `EnerStdLoss`), the `use_default_pf` option is accepted but ignored and therefore has no effect.

Copilot uses AI. Check for mistakes.

```json
"loss": {
"type": "ener",
"start_pref_e": 0.0,
"limit_pref_e": 0.0,
"start_pref_f": 0.0,
"limit_pref_f": 0.0,
"start_pref_pf": 1.0,
"limit_pref_pf": 1.0,
"use_default_pf": true,
"_comment": " that's all"
}
```

## Type embedding

Same as [`se_e2_a`](./train-se-e2-a.md).
Expand Down
Loading
Loading