From f737342c93735c103441eb9c6bac6cbc0bb5d808 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 4 Jul 2026 19:03:18 +0800 Subject: [PATCH 01/16] fix(loss): recover the per-atom mask in the pt loss so mixed_type ghosts can be excluded pt losses build model_pred internally and drop the model mask; recover it from atype in the TaskLoss base and call it from every pt loss forward; pt_expt already carries the mask; enables per-frame loss normalization in later tasks; no-op for non-mixed. --- deepmd/pt/loss/dos.py | 2 +- deepmd/pt/loss/ener.py | 2 +- deepmd/pt/loss/ener_spin.py | 2 +- deepmd/pt/loss/loss.py | 30 +++ deepmd/pt/loss/property.py | 2 +- deepmd/pt/loss/tensor.py | 2 +- .../tests/common/dpmodel/test_loss_padding.py | 76 ++++++++ source/tests/pt/test_loss_padding.py | 173 ++++++++++++++++++ 8 files changed, 284 insertions(+), 5 deletions(-) create mode 100644 source/tests/common/dpmodel/test_loss_padding.py create mode 100644 source/tests/pt/test_loss_padding.py diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index 64acae0c05..e28c1d8eb7 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -122,7 +122,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) coef = learning_rate / self.starter_learning_rate pref_dos = ( diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 50d83a4ac9..64d801ea51 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -225,7 +225,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) coef = learning_rate / self.starter_learning_rate pref_e = self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * coef pref_f = self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * coef diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 07e3403fdb..1e26ead84e 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -143,7 +143,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) coef = learning_rate / self.starter_learning_rate pref_e = self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * coef pref_fr = self.limit_pref_fr + (self.start_pref_fr - self.limit_pref_fr) * coef diff --git a/deepmd/pt/loss/loss.py b/deepmd/pt/loss/loss.py index 51eb7a9319..2f081d2abe 100644 --- a/deepmd/pt/loss/loss.py +++ b/deepmd/pt/loss/loss.py @@ -40,6 +40,36 @@ def label_requirement(self) -> list[DataRequirementItem]: """Return data label requirements needed for this loss calculation.""" pass + @staticmethod + def _inject_atom_mask( + model_pred: dict[str, torch.Tensor], + input_dict: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Recover the per-atom mask from atype for mixed_type batches. + + The exported forward drops the model's per-atom mask, so reconstruct it + here from ``atype`` (training-only). Ghost atoms have ``atype < 0``. + An all-ones mask is produced for non-mixed batches, keeping the loss + bit-identical to the pre-fix behavior. + + Parameters + ---------- + model_pred : dict[str, torch.Tensor] + Model predictions (modified in-place). + input_dict : dict[str, torch.Tensor] + Model inputs; must contain ``"atype"`` for injection to occur. + + Returns + ------- + dict[str, torch.Tensor] + ``model_pred`` with ``"mask"`` added if not already present. + """ + if "mask" not in model_pred and "atype" in input_dict: + atype = input_dict["atype"] + ref = model_pred.get("energy", input_dict.get("coord", atype)) + model_pred["mask"] = (atype >= 0).to(ref.dtype) + return model_pred + @staticmethod def display_if_exist(loss: torch.Tensor, find_property: float) -> torch.Tensor: """Display NaN if labeled property is not found. diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index af799d1a89..abf5bcfacb 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -100,7 +100,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) var_name = self.var_name nbz = model_pred[var_name].shape[0] assert model_pred[var_name].shape == (nbz, self.task_dim) diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index f329b79b20..3c1eeb8f99 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -101,7 +101,7 @@ def forward( more_loss: dict[str, torch.Tensor] Other losses for display. """ - model_pred = model(**input_dict) + model_pred = self._inject_atom_mask(model(**input_dict), input_dict) del learning_rate, mae if self.enable_atomic_weight: diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py new file mode 100644 index 0000000000..56266d2dc5 --- /dev/null +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Reusable grad-accumulation invariant harness for dpmodel loss tests. + +This module provides ``assert_grad_accum_invariant`` for Tasks 2-5 that +verify the loss on a padded multi-frame batch equals mean(per_frame_loss). + +The dpmodel losses accept numpy arrays (via the array_api_compat backend). + +Constants +--------- +NA = 3 # real atoms in the short frame +NB = 5 # real atoms in the full-width frame (NB == NP) +NP = 5 # padded width (nloc) +""" + +import numpy as np + +# --------------------------------------------------------------------------- +# Constants used by the multi-frame test harness (Tasks 2-5) +# --------------------------------------------------------------------------- + +NA = 3 # real atoms in frame A +NB = 5 # real atoms in frame B (== NP, so frame B is fully real) +NP = 5 # padded width (nloc) + + +# --------------------------------------------------------------------------- +# Reusable harness (used by Tasks 2-5; imported from here) +# --------------------------------------------------------------------------- + + +def assert_grad_accum_invariant( + loss_fn, + make_batch_A, + make_batch_B, + make_padded_batch, + rtol: float = 1e-5, + atol: float = 1e-6, +) -> None: + """Assert that padded-batch loss == mean(per_frame_loss) for two frames. + + The grad-accumulation invariant: a padded batch of [frame_A (NA real atoms + padded to NP) + frame_B (NB==NP real atoms)] must yield the same loss as + processing each frame separately and averaging. + + Parameters + ---------- + loss_fn : callable + Signature ``(model_pred, label, natoms) -> float``. The function + receives numpy-dict inputs and returns a scalar float. + make_batch_A : callable + Returns ``(model_pred, label, natoms)`` for frame A alone (1 frame, NA atoms). + make_batch_B : callable + Returns ``(model_pred, label, natoms)`` for frame B alone (1 frame, NB atoms). + make_padded_batch : callable + Returns ``(model_pred, label, natoms)`` for the 2-frame padded batch + (nf=2, nloc=NP; frame A is padded with NP-NA ghost rows). + rtol : float + Relative tolerance for ``np.isclose``. + atol : float + Absolute tolerance for ``np.isclose``. + """ + pred_A, label_A, natoms_A = make_batch_A() + pred_B, label_B, natoms_B = make_batch_B() + pred_pad, label_pad, natoms_pad = make_padded_batch() + + loss_A = float(loss_fn(pred_A, label_A, natoms_A)) + loss_B = float(loss_fn(pred_B, label_B, natoms_B)) + ref = 0.5 * (loss_A + loss_B) + + loss_pad = float(loss_fn(pred_pad, label_pad, natoms_pad)) + + assert np.isclose(loss_pad, ref, rtol=rtol, atol=atol), ( + f"Grad-accum invariant violated: padded_loss={loss_pad:.8f}, " + f"ref={ref:.8f}, diff={abs(loss_pad - ref):.2e}" + ) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py new file mode 100644 index 0000000000..82b0d35207 --- /dev/null +++ b/source/tests/pt/test_loss_padding.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for mixed_type loss padding-mask support in the pt backend. + +Task 1: verify that TaskLoss._inject_atom_mask correctly recovers the per-atom +mask from atype (ghost atoms have atype < 0) so that later tasks can exclude +them from loss reductions. + +Harness +------- +assert_grad_accum_invariant -- reusable by Tasks 2-5 to check the + grad-accumulation invariant: loss on a padded multi-frame batch must equal + mean_over_frames(per_frame_loss). +""" + +import torch + +from deepmd.pt.loss.loss import ( + TaskLoss, +) + +# --------------------------------------------------------------------------- +# Constants used by the multi-frame test harness (Tasks 2-5) +# --------------------------------------------------------------------------- + +NA = 3 # real atoms in frame A +NB = 5 # real atoms in frame B (== NP, so frame B is fully real) +NP = 5 # padded width (nloc) + + +# --------------------------------------------------------------------------- +# Reusable harness (used by Tasks 2-5; imported from here) +# --------------------------------------------------------------------------- + + +def assert_grad_accum_invariant( + loss_fn, + make_batch_A, + make_batch_B, + make_padded_batch, + rtol: float = 1e-5, + atol: float = 1e-6, +) -> None: + """Assert that padded-batch loss == mean(per_frame_loss) for two frames. + + The grad-accumulation invariant: a padded batch of [frame_A (NA real atoms + padded to NP) + frame_B (NB==NP real atoms)] must yield the same loss as + processing each frame separately and averaging. + + Parameters + ---------- + loss_fn : callable + Signature ``(model_pred, label, natoms) -> scalar torch.Tensor``. + make_batch_A : callable + Returns ``(model_pred, label, natoms)`` for frame A alone (1 frame, NA atoms). + make_batch_B : callable + Returns ``(model_pred, label, natoms)`` for frame B alone (1 frame, NB atoms). + make_padded_batch : callable + Returns ``(model_pred, label, natoms)`` for the 2-frame padded batch + (nf=2, nloc=NP; frame A is padded with NP-NA ghost rows). + rtol : float + Relative tolerance for ``torch.isclose``. + atol : float + Absolute tolerance for ``torch.isclose``. + """ + pred_A, label_A, natoms_A = make_batch_A() + pred_B, label_B, natoms_B = make_batch_B() + pred_pad, label_pad, natoms_pad = make_padded_batch() + + loss_A = loss_fn(pred_A, label_A, natoms_A) + loss_B = loss_fn(pred_B, label_B, natoms_B) + ref = 0.5 * (loss_A + loss_B) + + loss_pad = loss_fn(pred_pad, label_pad, natoms_pad) + + assert torch.isclose(loss_pad, ref, rtol=rtol, atol=atol), ( + f"Grad-accum invariant violated: padded_loss={loss_pad.item():.8f}, " + f"ref={ref.item():.8f}, diff={abs(loss_pad.item() - ref.item()):.2e}" + ) + + +# --------------------------------------------------------------------------- +# Task 1: unit tests for TaskLoss._inject_atom_mask +# --------------------------------------------------------------------------- + + +class TestInjectAtomMask: + """Unit tests for TaskLoss._inject_atom_mask. + + NOTE: source/tests/pt/__init__.py sets torch.set_default_device("cuda:9999999") + to catch tests that forget to specify device. All tensor creations here + must pass device="cpu" explicitly. + """ + + def test_injects_mask_from_atype(self) -> None: + """When model_pred has no mask, recover it from atype (<0 = ghost).""" + # nf=2, nloc=5; frame 0 has ghost atoms at positions 3 and 4 + atype = torch.tensor([[0, 1, 0, -1, -1], [0, 1, 0, 1, 0]], device="cpu") + energy = torch.zeros(2, 1, dtype=torch.float64, device="cpu") + input_dict = {"atype": atype} + model_pred = {"energy": energy} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert "mask" in result, "mask must be injected" + assert result["mask"].shape == (2, 5), f"wrong shape {result['mask'].shape}" + expected = torch.tensor( + [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]], + dtype=torch.float64, + device="cpu", + ) + assert torch.allclose(result["mask"], expected), ( + f"mask values wrong: {result['mask']}" + ) + assert result["mask"].dtype == energy.dtype, ( + "mask dtype must match reference tensor dtype" + ) + + def test_mask_dtype_follows_ref(self) -> None: + """Mask dtype follows the reference tensor (energy or coord or atype).""" + atype = torch.tensor([[0, -1]], dtype=torch.long, device="cpu") + model_pred = {"energy": torch.zeros(1, 1, dtype=torch.float32, device="cpu")} + input_dict = {"atype": atype} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert result["mask"].dtype == torch.float32 + + def test_fallback_to_coord_when_no_energy(self) -> None: + """When energy absent, ref falls back to coord.""" + atype = torch.tensor([[0, -1]], dtype=torch.long, device="cpu") + coord = torch.zeros(1, 2, 3, dtype=torch.float64, device="cpu") + model_pred = {} + input_dict = {"atype": atype, "coord": coord} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert "mask" in result + assert result["mask"].dtype == torch.float64 + + def test_mask_not_overwritten_if_present(self) -> None: + """When mask is already in model_pred, leave it unchanged.""" + existing = torch.ones(2, 5, device="cpu") * 0.5 + model_pred = { + "mask": existing, + "energy": torch.zeros(2, 1, device="cpu"), + } + input_dict = {"atype": torch.zeros(2, 5, dtype=torch.long, device="cpu")} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert result["mask"] is existing, "existing mask must not be overwritten" + + def test_no_atype_no_injection(self) -> None: + """When input_dict has no atype, mask must not be injected.""" + model_pred = {"energy": torch.zeros(2, 1, device="cpu")} + input_dict = {} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert "mask" not in result, "mask must not appear when atype is absent" + + def test_all_real_atoms_gives_all_ones(self) -> None: + """Non-mixed batch (all atype >= 0) yields all-ones mask.""" + atype = torch.tensor([[0, 1, 2], [1, 2, 0]], device="cpu") + energy = torch.zeros(2, 1, dtype=torch.float32, device="cpu") + model_pred = {"energy": energy} + input_dict = {"atype": atype} + + result = TaskLoss._inject_atom_mask(model_pred, input_dict) + + assert torch.all(result["mask"] == 1.0), ( + "all-real atoms must give all-ones mask" + ) From 7bacaeee3af62da92100673a46c782e6732dc30f Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 00:06:30 +0800 Subject: [PATCH 02/16] fix(loss): per-frame normalize dos/tensor losses to exclude mixed_type padding Atomic terms (ados/acdf for DOSLoss; local tensor for TensorLoss) now use per-frame masked mean (idiom 1) instead of a cross-frame pooled masked mean, ensuring each frame is normalized by its own real-atom count. Global terms (dos/cdf for DOSLoss; global tensor for TensorLoss) now use plain mean (idiom 3), dropping the prior atom-count weighting that violated the grad-accumulation invariant. Both meet the invariant: a padded [3+5]-atom batch yields the same loss as processing each frame separately and averaging. Boolean fancy indexing removed from pt atomic branches. Changes are bit-identical for non-mixed batches (all-ones mask). Applies to deepmd/dpmodel/loss/{dos,tensor}.py (serves pt_expt) and deepmd/pt/loss/{dos,tensor}.py (mirror). TDD: invariant tests confirmed RED before and GREEN after (24/24 passed). --- deepmd/dpmodel/loss/dos.py | 50 ++- deepmd/dpmodel/loss/tensor.py | 22 +- deepmd/pt/loss/dos.py | 52 ++- deepmd/pt/loss/tensor.py | 21 +- .../tests/common/dpmodel/test_loss_padding.py | 354 +++++++++++++++ source/tests/pt/test_loss_padding.py | 411 ++++++++++++++++++ 6 files changed, 846 insertions(+), 64 deletions(-) diff --git a/deepmd/dpmodel/loss/dos.py b/deepmd/dpmodel/loss/dos.py index 0036ae1620..953b413dd1 100644 --- a/deepmd/dpmodel/loss/dos.py +++ b/deepmd/dpmodel/loss/dos.py @@ -130,15 +130,17 @@ def call( local_label = xp.reshape( label_dict["atom_dos"], (-1, natoms, self.numb_dos) ) - diff = xp.reshape(local_pred - local_label, (-1, self.numb_dos)) + diff3d = local_pred - local_label # [nf, natoms, numb_dos] if "mask" in model_dict: - mask = xp.reshape(model_dict["mask"], (-1,)) - mask_float = xp.astype(mask, diff.dtype) - diff = diff * mask_float[:, None] - n_valid = xp.sum(mask_float) - l2_local_loss_dos = xp.sum(xp.square(diff)) / (n_valid * self.numb_dos) + # idiom 1: per-frame masked mean, then average over frames + maskf = xp.astype(model_dict["mask"], diff3d.dtype) # [nf, natoms] + nf = diff3d.shape[0] + sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1)) + per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * self.numb_dos # [nf] + l2_local_loss_dos = xp.mean(per_frame_sum / per_frame_dof) else: - l2_local_loss_dos = xp.mean(xp.square(diff)) + l2_local_loss_dos = xp.mean(xp.square(diff3d)) loss += pref_ados * l2_local_loss_dos more_loss["rmse_local_dos"] = self.display_if_exist( xp.sqrt(l2_local_loss_dos), find_local @@ -155,15 +157,17 @@ def call( xp.reshape(label_dict["atom_dos"], (-1, natoms, self.numb_dos)), axis=-1, ) - diff = xp.reshape(local_pred_cdf - local_label_cdf, (-1, self.numb_dos)) + diff3d = local_pred_cdf - local_label_cdf # [nf, natoms, numb_dos] if "mask" in model_dict: - mask = xp.reshape(model_dict["mask"], (-1,)) - mask_float = xp.astype(mask, diff.dtype) - diff = diff * mask_float[:, None] - n_valid = xp.sum(mask_float) - l2_local_loss_cdf = xp.sum(xp.square(diff)) / (n_valid * self.numb_dos) + # idiom 1: per-frame masked mean, then average over frames + maskf = xp.astype(model_dict["mask"], diff3d.dtype) # [nf, natoms] + nf = diff3d.shape[0] + sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1)) + per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * self.numb_dos # [nf] + l2_local_loss_cdf = xp.mean(per_frame_sum / per_frame_dof) else: - l2_local_loss_cdf = xp.mean(xp.square(diff)) + l2_local_loss_cdf = xp.mean(xp.square(diff3d)) loss += pref_acdf * l2_local_loss_cdf more_loss["rmse_local_cdf"] = self.display_if_exist( xp.sqrt(l2_local_loss_cdf), find_local @@ -175,15 +179,14 @@ def call( global_pred = xp.reshape(model_dict["dos"], (-1, self.numb_dos)) global_label = xp.reshape(label_dict["dos"], (-1, self.numb_dos)) diff = global_pred - global_label + # idiom 3: global dos is already padding-invariant; plain mean suffices + l2_global_loss_dos = xp.mean(xp.square(diff)) if "mask" in model_dict: - atom_num = xp.sum(model_dict["mask"], axis=-1, keepdims=True) - l2_global_loss_dos = xp.mean( - xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num) + atom_num = xp.mean( + xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) ) - atom_num = xp.mean(xp.astype(atom_num, diff.dtype)) else: atom_num = natoms - l2_global_loss_dos = xp.mean(xp.square(diff)) loss += pref_dos * l2_global_loss_dos more_loss["rmse_global_dos"] = self.display_if_exist( xp.sqrt(l2_global_loss_dos) / atom_num, find_global @@ -199,15 +202,14 @@ def call( xp.reshape(label_dict["dos"], (-1, self.numb_dos)), axis=-1 ) diff = global_pred_cdf - global_label_cdf + # idiom 3: global cdf is already padding-invariant; plain mean suffices + l2_global_loss_cdf = xp.mean(xp.square(diff)) if "mask" in model_dict: - atom_num = xp.sum(model_dict["mask"], axis=-1, keepdims=True) - l2_global_loss_cdf = xp.mean( - xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num) + atom_num = xp.mean( + xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) ) - atom_num = xp.mean(xp.astype(atom_num, diff.dtype)) else: atom_num = natoms - l2_global_loss_cdf = xp.mean(xp.square(diff)) loss += pref_cdf * l2_global_loss_cdf more_loss["rmse_global_cdf"] = self.display_if_exist( xp.sqrt(l2_global_loss_cdf) / atom_num, find_global diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py index e367457451..4b21175185 100644 --- a/deepmd/dpmodel/loss/tensor.py +++ b/deepmd/dpmodel/loss/tensor.py @@ -103,11 +103,14 @@ def call( diff = xp.reshape(local_pred - local_label, (-1, self.tensor_size)) diff = diff * atomic_weight if "mask" in model_dict: - mask = xp.reshape(model_dict["mask"], (-1,)) - mask_float = xp.astype(mask, diff.dtype) - diff = diff * mask_float[:, None] - n_valid = xp.sum(mask_float) - l2_local_loss = xp.sum(xp.square(diff)) / (n_valid * self.tensor_size) + # idiom 1: per-frame masked mean, then average over frames + maskf = xp.astype(model_dict["mask"], diff.dtype) # [nf, natoms] + nf = local_pred.shape[0] + diff3d = xp.reshape(diff, (nf, natoms, self.tensor_size)) + sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1)) + per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * self.tensor_size # [nf] + l2_local_loss = xp.mean(per_frame_sum / per_frame_dof) else: l2_local_loss = xp.mean(xp.square(diff)) loss += local_weight * l2_local_loss @@ -129,15 +132,14 @@ def call( label_dict[self.label_name], (-1, self.tensor_size) ) diff = global_pred - global_label + # idiom 3: global tensor is already padding-invariant; plain mean suffices + l2_global_loss = xp.mean(xp.square(diff)) if "mask" in model_dict: - atom_num = xp.sum(model_dict["mask"], axis=-1, keepdims=True) - l2_global_loss = xp.mean( - xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num) + atom_num = xp.mean( + xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype) ) - atom_num = xp.mean(xp.astype(atom_num, diff.dtype)) else: atom_num = natoms - l2_global_loss = xp.mean(xp.square(diff)) loss += global_weight * l2_global_loss more_loss[f"rmse_global_{self.tensor_name}"] = self.display_if_exist( xp.sqrt(l2_global_loss) / atom_num, find_global diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py index e28c1d8eb7..bf2d2b9710 100644 --- a/deepmd/pt/loss/dos.py +++ b/deepmd/pt/loss/dos.py @@ -149,12 +149,19 @@ def forward( local_tensor_label_dos = label["atom_dos"].reshape( [-1, natoms, self.numb_dos] ) - diff = (local_tensor_pred_dos - local_tensor_label_dos).reshape( - [-1, self.numb_dos] - ) + diff = ( + local_tensor_pred_dos - local_tensor_label_dos + ) # [nf, natoms, numb_dos] if "mask" in model_pred: - diff = diff[model_pred["mask"].reshape([-1]).bool()] - l2_local_loss_dos = torch.mean(torch.square(diff)) + # idiom 1: per-frame masked mean, then average over frames + nf = diff.shape[0] + maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] + sq = torch.square(diff) * maskf.reshape(nf, natoms, 1) + per_frame_sum = sq.reshape(nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) * self.numb_dos # [nf] + l2_local_loss_dos = (per_frame_sum / per_frame_dof).mean() + else: + l2_local_loss_dos = torch.mean(torch.square(diff)) if not self.inference: more_loss["l2_local_dos_loss"] = self.display_if_exist( l2_local_loss_dos.detach(), find_local @@ -173,12 +180,19 @@ def forward( local_tensor_label_cdf = torch.cumsum( label["atom_dos"].reshape([-1, natoms, self.numb_dos]), dim=-1 ) - diff = (local_tensor_pred_cdf - local_tensor_label_cdf).reshape( - [-1, self.numb_dos] - ) + diff = ( + local_tensor_pred_cdf - local_tensor_label_cdf + ) # [nf, natoms, numb_dos] if "mask" in model_pred: - diff = diff[model_pred["mask"].reshape([-1]).bool()] - l2_local_loss_cdf = torch.mean(torch.square(diff)) + # idiom 1: per-frame masked mean, then average over frames + nf = diff.shape[0] + maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] + sq = torch.square(diff) * maskf.reshape(nf, natoms, 1) + per_frame_sum = sq.reshape(nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) * self.numb_dos # [nf] + l2_local_loss_cdf = (per_frame_sum / per_frame_dof).mean() + else: + l2_local_loss_cdf = torch.mean(torch.square(diff)) if not self.inference: more_loss["l2_local_cdf_loss"] = self.display_if_exist( l2_local_loss_cdf.detach(), find_local @@ -194,15 +208,12 @@ def forward( global_tensor_pred_dos = model_pred["dos"].reshape([-1, self.numb_dos]) global_tensor_label_dos = label["dos"].reshape([-1, self.numb_dos]) diff = global_tensor_pred_dos - global_tensor_label_dos + # idiom 3: global dos is already padding-invariant; plain mean suffices + l2_global_loss_dos = torch.mean(torch.square(diff)) if "mask" in model_pred: - atom_num = model_pred["mask"].sum(-1, keepdim=True) - l2_global_loss_dos = torch.mean( - torch.sum(torch.square(diff) * atom_num, dim=0) / atom_num.sum() - ) - atom_num = torch.mean(atom_num.float()) + atom_num = model_pred["mask"].sum(-1).float().mean() else: atom_num = natoms - l2_global_loss_dos = torch.mean(torch.square(diff)) if not self.inference: more_loss["l2_global_dos_loss"] = self.display_if_exist( l2_global_loss_dos.detach(), find_global @@ -222,15 +233,12 @@ def forward( label["dos"].reshape([-1, self.numb_dos]), dim=-1 ) diff = global_tensor_pred_cdf - global_tensor_label_cdf + # idiom 3: global cdf is already padding-invariant; plain mean suffices + l2_global_loss_cdf = torch.mean(torch.square(diff)) if "mask" in model_pred: - atom_num = model_pred["mask"].sum(-1, keepdim=True) - l2_global_loss_cdf = torch.mean( - torch.sum(torch.square(diff) * atom_num, dim=0) / atom_num.sum() - ) - atom_num = torch.mean(atom_num.float()) + atom_num = model_pred["mask"].sum(-1).float().mean() else: atom_num = natoms - l2_global_loss_cdf = torch.mean(torch.square(diff)) if not self.inference: more_loss["l2_global_cdf_loss"] = self.display_if_exist( l2_global_loss_cdf.detach(), find_global diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py index 3c1eeb8f99..c46c2ec0dd 100644 --- a/deepmd/pt/loss/tensor.py +++ b/deepmd/pt/loss/tensor.py @@ -129,8 +129,16 @@ def forward( ) diff = diff * atomic_weight if "mask" in model_pred: - diff = diff[model_pred["mask"].reshape([-1]).bool()] - l2_local_loss = torch.mean(torch.square(diff)) + # idiom 1: per-frame masked mean, then average over frames + nf = local_tensor_pred.shape[0] + maskf = model_pred["mask"].to(diff.dtype) # [nf, natoms] + diff3d = diff.reshape(nf, natoms, self.tensor_size) + sq = torch.square(diff3d) * maskf.reshape(nf, natoms, 1) + per_frame_sum = sq.reshape(nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) * self.tensor_size # [nf] + l2_local_loss = (per_frame_sum / per_frame_dof).mean() + else: + l2_local_loss = torch.mean(torch.square(diff)) if not self.inference: more_loss[f"l2_local_{self.tensor_name}_loss"] = self.display_if_exist( l2_local_loss.detach(), find_local @@ -152,15 +160,12 @@ def forward( ) global_tensor_label = label[self.label_name].reshape([-1, self.tensor_size]) diff = global_tensor_pred - global_tensor_label + # idiom 3: global tensor is already padding-invariant; plain mean suffices + l2_global_loss = torch.mean(torch.square(diff)) if "mask" in model_pred: - atom_num = model_pred["mask"].sum(-1, keepdim=True) - l2_global_loss = torch.mean( - torch.sum(torch.square(diff) * atom_num, dim=0) / atom_num.sum() - ) - atom_num = torch.mean(atom_num.float()) + atom_num = model_pred["mask"].sum(-1).float().mean() else: atom_num = natoms - l2_global_loss = torch.mean(torch.square(diff)) if not self.inference: more_loss[f"l2_global_{self.tensor_name}_loss"] = self.display_if_exist( l2_global_loss.detach(), find_global diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 56266d2dc5..4faac008ed 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -15,6 +15,13 @@ import numpy as np +from deepmd.dpmodel.loss.dos import ( + DOSLoss, +) +from deepmd.dpmodel.loss.tensor import ( + TensorLoss, +) + # --------------------------------------------------------------------------- # Constants used by the multi-frame test harness (Tasks 2-5) # --------------------------------------------------------------------------- @@ -74,3 +81,350 @@ def assert_grad_accum_invariant( f"Grad-accum invariant violated: padded_loss={loss_pad:.8f}, " f"ref={ref:.8f}, diff={abs(loss_pad - ref):.2e}" ) + + +# --------------------------------------------------------------------------- +# Helpers for constructing test data +# --------------------------------------------------------------------------- + +RNG = np.random.default_rng(42) +NUMB_DOS = 4 +TENSOR_SIZE = 3 + + +def _rnd(*shape): + return RNG.standard_normal(shape).astype(np.float64) + + +# --------------------------------------------------------------------------- +# Task 2: DOSLoss -- atomic (ados / acdf) and global (dos / cdf) +# --------------------------------------------------------------------------- + + +class TestDOSLossAtomicGradAccum: + """Per-frame masked mean (idiom 1) for atomic dos / acdf terms.""" + + def _make_loss(self): + return DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=0.0, + limit_pref_dos=0.0, + start_pref_cdf=0.0, + limit_pref_cdf=0.0, + start_pref_ados=1.0, + limit_pref_ados=1.0, + start_pref_acdf=1.0, + limit_pref_acdf=1.0, + ) + + def _atom_dos_A(self): + return _rnd(NA, NUMB_DOS) + + def _atom_dos_B(self): + return _rnd(NB, NUMB_DOS) + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def _make_batch_A(self): + pred = _rnd(NA, NUMB_DOS) + label = _rnd(NA, NUMB_DOS) + model_pred = { + "atom_dos": pred, + "mask": np.ones((1, NA), dtype=np.float64), + } + label_dict = {"atom_dos": label, "find_atom_dos": 1.0} + return model_pred, label_dict, NA + + def _make_batch_B(self): + pred = _rnd(NB, NUMB_DOS) + label = _rnd(NB, NUMB_DOS) + model_pred = { + "atom_dos": pred, + "mask": np.ones((1, NB), dtype=np.float64), + } + label_dict = {"atom_dos": label, "find_atom_dos": 1.0} + return model_pred, label_dict, NB + + def _make_padded_batch(self, pred_A_data, label_A_data, pred_B_data, label_B_data): + # Frame A: NA real atoms padded to NP; ghost pred/label = 0 + pred_A_pad = np.zeros((NP, NUMB_DOS), dtype=np.float64) + pred_A_pad[:NA] = pred_A_data + label_A_pad = np.zeros((NP, NUMB_DOS), dtype=np.float64) + label_A_pad[:NA] = label_A_data + mask_A = np.array([[1.0] * NA + [0.0] * (NP - NA)], dtype=np.float64) + + # Frame B: NB == NP, all real + mask_B = np.ones((1, NB), dtype=np.float64) + + atom_dos_pad = np.concatenate( + [pred_A_pad, pred_B_data], axis=0 + ) # [NP+NB, ncomp] + atom_dos_label = np.concatenate([label_A_pad, label_B_data], axis=0) + mask_pad = np.concatenate([mask_A, mask_B], axis=0) # [2, NP] + model_pred = {"atom_dos": atom_dos_pad, "mask": mask_pad} + label_dict = {"atom_dos": atom_dos_label, "find_atom_dos": 1.0} + return model_pred, label_dict, NP + + def test_ados_grad_accum_invariant(self): + """Atomic dos per-frame masked mean meets the grad-accum invariant.""" + pred_A = _rnd(NA, NUMB_DOS) + label_A = _rnd(NA, NUMB_DOS) + pred_B = _rnd(NB, NUMB_DOS) + label_B = _rnd(NB, NUMB_DOS) + + def make_A(): + return ( + {"atom_dos": pred_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {"atom_dos": label_A, "find_atom_dos": 1.0}, + NA, + ) + + def make_B(): + return ( + {"atom_dos": pred_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {"atom_dos": label_B, "find_atom_dos": 1.0}, + NB, + ) + + def make_padded(): + return self._make_padded_batch(pred_A, label_A, pred_B, label_B) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_acdf_grad_accum_invariant(self): + """Atomic cdf per-frame masked mean meets the grad-accum invariant.""" + # acdf uses cumsum of atom_dos -- same data pipeline + self.test_ados_grad_accum_invariant() + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd(NB, NUMB_DOS) + label = _rnd(NB, NUMB_DOS) + with_mask = { + "atom_dos": pred, + "mask": np.ones((1, NB), dtype=np.float64), + } + without_mask = {"atom_dos": pred} + label_dict = {"atom_dos": label, "find_atom_dos": 1.0} + loss_m, _ = self._make_loss().call(1.0, NB, with_mask, label_dict) + loss_nm, _ = self._make_loss().call(1.0, NB, without_mask, label_dict) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + +class TestDOSLossGlobalGradAccum: + """Plain mean (idiom 3) for global dos / cdf terms.""" + + def _make_loss(self): + return DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=1.0, + limit_pref_dos=1.0, + start_pref_cdf=1.0, + limit_pref_cdf=1.0, + start_pref_ados=0.0, + limit_pref_ados=0.0, + start_pref_acdf=0.0, + limit_pref_acdf=0.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def test_dos_grad_accum_invariant(self): + """Global dos plain mean meets the grad-accum invariant.""" + pred_A = _rnd(1, NUMB_DOS) + label_A = _rnd(1, NUMB_DOS) + pred_B = _rnd(1, NUMB_DOS) + label_B = _rnd(1, NUMB_DOS) + + def make_A(): + return ( + {"dos": pred_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {"dos": label_A, "find_dos": 1.0}, + NA, + ) + + def make_B(): + return ( + {"dos": pred_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {"dos": label_B, "find_dos": 1.0}, + NB, + ) + + def make_padded(): + pred_pad = np.concatenate([pred_A, pred_B], axis=0) # [2, ncomp] + label_pad = np.concatenate([label_A, label_B], axis=0) + mask_pad = np.array( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64 + ) + return ( + {"dos": pred_pad, "mask": mask_pad}, + {"dos": label_pad, "find_dos": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_cdf_grad_accum_invariant(self): + """Global cdf plain mean meets the grad-accum invariant.""" + # cdf uses cumsum of dos -- reuse the same logic + self.test_dos_grad_accum_invariant() + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd(2, NUMB_DOS) + label = _rnd(2, NUMB_DOS) + with_mask = {"dos": pred, "mask": np.ones((2, NB), dtype=np.float64)} + without_mask = {"dos": pred} + label_dict = {"dos": label, "find_dos": 1.0} + loss_m, _ = self._make_loss().call(1.0, NB, with_mask, label_dict) + loss_nm, _ = self._make_loss().call(1.0, NB, without_mask, label_dict) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + +# --------------------------------------------------------------------------- +# Task 2: TensorLoss -- local and global tensor +# --------------------------------------------------------------------------- + + +class TestTensorLossLocalGradAccum: + """Per-frame masked mean (idiom 1) for local tensor term.""" + + def _make_loss(self): + return TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=1.0, + pref=0.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def test_local_grad_accum_invariant(self): + """Local tensor per-frame masked mean meets the grad-accum invariant.""" + pred_A = _rnd(NA, TENSOR_SIZE) + label_A = _rnd(NA, TENSOR_SIZE) + pred_B = _rnd(NB, TENSOR_SIZE) + label_B = _rnd(NB, TENSOR_SIZE) + + def make_A(): + return ( + {"dipole": pred_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {"atom_dipole": label_A, "find_atom_dipole": 1.0}, + NA, + ) + + def make_B(): + return ( + {"dipole": pred_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {"atom_dipole": label_B, "find_atom_dipole": 1.0}, + NB, + ) + + def make_padded(): + pred_A_pad = np.zeros((NP, TENSOR_SIZE), dtype=np.float64) + pred_A_pad[:NA] = pred_A + label_A_pad = np.zeros((NP, TENSOR_SIZE), dtype=np.float64) + label_A_pad[:NA] = label_A + mask_A = np.array([[1.0] * NA + [0.0] * (NP - NA)], dtype=np.float64) + mask_B = np.ones((1, NB), dtype=np.float64) + dipole_pad = np.concatenate([pred_A_pad, pred_B], axis=0) + label_pad = np.concatenate([label_A_pad, label_B], axis=0) + mask_pad = np.concatenate([mask_A, mask_B], axis=0) + return ( + {"dipole": dipole_pad, "mask": mask_pad}, + {"atom_dipole": label_pad, "find_atom_dipole": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd(NB, TENSOR_SIZE) + label = _rnd(NB, TENSOR_SIZE) + with_mask = {"dipole": pred, "mask": np.ones((1, NB), dtype=np.float64)} + without_mask = {"dipole": pred} + label_dict = {"atom_dipole": label, "find_atom_dipole": 1.0} + loss_m, _ = self._make_loss().call(1.0, NB, with_mask, label_dict) + loss_nm, _ = self._make_loss().call(1.0, NB, without_mask, label_dict) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + +class TestTensorLossGlobalGradAccum: + """Plain mean (idiom 3) for global tensor term.""" + + def _make_loss(self): + return TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=0.0, + pref=1.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss, _ = self._make_loss().call(1.0, natoms, model_pred, label) + return float(loss) + + def test_global_grad_accum_invariant(self): + """Global tensor plain mean meets the grad-accum invariant.""" + pred_A = _rnd(1, TENSOR_SIZE) + label_A = _rnd(1, TENSOR_SIZE) + pred_B = _rnd(1, TENSOR_SIZE) + label_B = _rnd(1, TENSOR_SIZE) + + def make_A(): + return ( + {"global_dipole": pred_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {"dipole": label_A, "find_dipole": 1.0}, + NA, + ) + + def make_B(): + return ( + {"global_dipole": pred_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {"dipole": label_B, "find_dipole": 1.0}, + NB, + ) + + def make_padded(): + pred_pad = np.concatenate([pred_A, pred_B], axis=0) + label_pad = np.concatenate([label_A, label_B], axis=0) + mask_pad = np.array( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64 + ) + return ( + {"global_dipole": pred_pad, "mask": mask_pad}, + {"dipole": label_pad, "find_dipole": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd(2, TENSOR_SIZE) + label = _rnd(2, TENSOR_SIZE) + with_mask = {"global_dipole": pred, "mask": np.ones((2, NB), dtype=np.float64)} + without_mask = {"global_dipole": pred} + label_dict = {"dipole": label, "find_dipole": 1.0} + loss_m, _ = self._make_loss().call(1.0, NB, with_mask, label_dict) + loss_nm, _ = self._make_loss().call(1.0, NB, without_mask, label_dict) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 82b0d35207..48f954ca56 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -12,11 +12,18 @@ mean_over_frames(per_frame_loss). """ +import numpy as np import torch +from deepmd.pt.loss.dos import ( + DOSLoss, +) from deepmd.pt.loss.loss import ( TaskLoss, ) +from deepmd.pt.loss.tensor import ( + TensorLoss, +) # --------------------------------------------------------------------------- # Constants used by the multi-frame test harness (Tasks 2-5) @@ -78,6 +85,410 @@ def assert_grad_accum_invariant( ) +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +RNG = np.random.default_rng(42) +NUMB_DOS = 4 +TENSOR_SIZE = 3 + + +def _rnd_t(*shape): + """Create a random CPU tensor (device must be explicit to avoid default-device trap).""" + return torch.tensor(RNG.standard_normal(shape), dtype=torch.float64, device="cpu") + + +class _MockModel: + """Callable that ignores inputs and returns a fixed model_pred dict. + + Used in pt loss tests to bypass the actual model forward pass while still + exercising the loss computation code inside DOSLoss.forward / TensorLoss.forward. + The mask is pre-populated in ``pred`` so ``_inject_atom_mask`` leaves it alone. + """ + + def __init__(self, pred: dict): + self._pred = pred + + def __call__(self, **kwargs): + return dict(self._pred) # shallow copy; _inject_atom_mask may mutate it + + +# --------------------------------------------------------------------------- +# Task 2: DOSLoss -- atomic (ados / acdf) and global (dos / cdf) +# --------------------------------------------------------------------------- + + +class TestPTDOSLossAtomicGradAccum: + """Per-frame masked mean (idiom 1) for atomic dos / acdf terms. + + _loss_fn calls the ACTUAL pt DOSLoss.forward() via a mock model so that + RED/GREEN transitions directly reflect changes to deepmd/pt/loss/dos.py. + """ + + def _make_loss(self): + return DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=0.0, + limit_pref_dos=0.0, + start_pref_cdf=0.0, + limit_pref_cdf=0.0, + start_pref_ados=1.0, + limit_pref_ados=1.0, + start_pref_acdf=1.0, + limit_pref_acdf=1.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss_obj = self._make_loss() + _, loss, _ = loss_obj.forward( + input_dict={}, # no atype; mask already in model_pred → not re-injected + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_ados_grad_accum_invariant(self): + """Atomic dos per-frame masked mean meets the grad-accum invariant.""" + pred_A = _rnd_t(NA, NUMB_DOS) + label_A = _rnd_t(NA, NUMB_DOS) + pred_B = _rnd_t(NB, NUMB_DOS) + label_B = _rnd_t(NB, NUMB_DOS) + + def make_A(): + return ( + { + "atom_dos": pred_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {"atom_dos": label_A, "find_atom_dos": 1.0}, + NA, + ) + + def make_B(): + return ( + { + "atom_dos": pred_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {"atom_dos": label_B, "find_atom_dos": 1.0}, + NB, + ) + + def make_padded(): + pred_A_pad = torch.zeros(NP, NUMB_DOS, dtype=torch.float64, device="cpu") + pred_A_pad[:NA] = pred_A + label_A_pad = torch.zeros(NP, NUMB_DOS, dtype=torch.float64, device="cpu") + label_A_pad[:NA] = label_A + mask_A = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA)], dtype=torch.float64, device="cpu" + ) + mask_B = torch.ones(1, NB, dtype=torch.float64, device="cpu") + atom_dos_pad = torch.cat([pred_A_pad, pred_B], dim=0) + atom_dos_label = torch.cat([label_A_pad, label_B], dim=0) + mask_pad = torch.cat([mask_A, mask_B], dim=0) + return ( + {"atom_dos": atom_dos_pad, "mask": mask_pad}, + {"atom_dos": atom_dos_label, "find_atom_dos": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd_t(NB, NUMB_DOS) + label = _rnd_t(NB, NUMB_DOS) + with_mask = { + "atom_dos": pred, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + without_mask = {"atom_dos": pred} + label_dict = {"atom_dos": label, "find_atom_dos": 1.0} + loss_m = self._loss_fn(with_mask, label_dict, NB) + loss_nm = self._loss_fn(without_mask, label_dict, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTDOSLossGlobalGradAccum: + """Plain mean (idiom 3) for global dos / cdf terms. + + _loss_fn calls the ACTUAL pt DOSLoss.forward() via a mock model. + """ + + def _make_loss(self): + return DOSLoss( + starter_learning_rate=1.0, + numb_dos=NUMB_DOS, + start_pref_dos=1.0, + limit_pref_dos=1.0, + start_pref_cdf=1.0, + limit_pref_cdf=1.0, + start_pref_ados=0.0, + limit_pref_ados=0.0, + start_pref_acdf=0.0, + limit_pref_acdf=0.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss_obj = self._make_loss() + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_dos_grad_accum_invariant(self): + """Global dos plain mean meets the grad-accum invariant.""" + pred_A = _rnd_t(1, NUMB_DOS) + label_A = _rnd_t(1, NUMB_DOS) + pred_B = _rnd_t(1, NUMB_DOS) + label_B = _rnd_t(1, NUMB_DOS) + + def make_A(): + return ( + { + "dos": pred_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {"dos": label_A, "find_dos": 1.0}, + NA, + ) + + def make_B(): + return ( + { + "dos": pred_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {"dos": label_B, "find_dos": 1.0}, + NB, + ) + + def make_padded(): + pred_pad = torch.cat([pred_A, pred_B], dim=0) + label_pad = torch.cat([label_A, label_B], dim=0) + mask_pad = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", + ) + return ( + {"dos": pred_pad, "mask": mask_pad}, + {"dos": label_pad, "find_dos": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd_t(2, NUMB_DOS) + label = _rnd_t(2, NUMB_DOS) + with_mask = { + "dos": pred, + "mask": torch.ones(2, NB, dtype=torch.float64, device="cpu"), + } + without_mask = {"dos": pred} + label_dict = {"dos": label, "find_dos": 1.0} + loss_m = self._loss_fn(with_mask, label_dict, NB) + loss_nm = self._loss_fn(without_mask, label_dict, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +# --------------------------------------------------------------------------- +# Task 2: TensorLoss -- local and global tensor +# --------------------------------------------------------------------------- + + +class TestPTTensorLossLocalGradAccum: + """Per-frame masked mean (idiom 1) for local tensor term. + + _loss_fn calls the ACTUAL pt TensorLoss.forward() via a mock model. + """ + + def _make_loss(self): + return TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=1.0, + pref=0.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss_obj = self._make_loss() + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_local_grad_accum_invariant(self): + """Local tensor per-frame masked mean meets the grad-accum invariant.""" + pred_A = _rnd_t(NA, TENSOR_SIZE) + label_A = _rnd_t(NA, TENSOR_SIZE) + pred_B = _rnd_t(NB, TENSOR_SIZE) + label_B = _rnd_t(NB, TENSOR_SIZE) + + def make_A(): + return ( + { + "dipole": pred_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {"atom_dipole": label_A, "find_atom_dipole": 1.0}, + NA, + ) + + def make_B(): + return ( + { + "dipole": pred_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {"atom_dipole": label_B, "find_atom_dipole": 1.0}, + NB, + ) + + def make_padded(): + pred_A_pad = torch.zeros(NP, TENSOR_SIZE, dtype=torch.float64, device="cpu") + pred_A_pad[:NA] = pred_A + label_A_pad = torch.zeros( + NP, TENSOR_SIZE, dtype=torch.float64, device="cpu" + ) + label_A_pad[:NA] = label_A + mask_A = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA)], dtype=torch.float64, device="cpu" + ) + mask_B = torch.ones(1, NB, dtype=torch.float64, device="cpu") + dipole_pad = torch.cat([pred_A_pad, pred_B], dim=0) + label_pad = torch.cat([label_A_pad, label_B], dim=0) + mask_pad = torch.cat([mask_A, mask_B], dim=0) + return ( + {"dipole": dipole_pad, "mask": mask_pad}, + {"atom_dipole": label_pad, "find_atom_dipole": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd_t(NB, TENSOR_SIZE) + label = _rnd_t(NB, TENSOR_SIZE) + with_mask = { + "dipole": pred, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + without_mask = {"dipole": pred} + label_dict = {"atom_dipole": label, "find_atom_dipole": 1.0} + loss_m = self._loss_fn(with_mask, label_dict, NB) + loss_nm = self._loss_fn(without_mask, label_dict, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTTensorLossGlobalGradAccum: + """Plain mean (idiom 3) for global tensor term. + + _loss_fn calls the ACTUAL pt TensorLoss.forward() via a mock model. + """ + + def _make_loss(self): + return TensorLoss( + tensor_name="dipole", + tensor_size=TENSOR_SIZE, + label_name="dipole", + pref_atomic=0.0, + pref=1.0, + ) + + def _loss_fn(self, model_pred, label, natoms): + loss_obj = self._make_loss() + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_global_grad_accum_invariant(self): + """Global tensor plain mean meets the grad-accum invariant.""" + pred_A = _rnd_t(1, TENSOR_SIZE) + label_A = _rnd_t(1, TENSOR_SIZE) + pred_B = _rnd_t(1, TENSOR_SIZE) + label_B = _rnd_t(1, TENSOR_SIZE) + + def make_A(): + return ( + { + "global_dipole": pred_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {"dipole": label_A, "find_dipole": 1.0}, + NA, + ) + + def make_B(): + return ( + { + "global_dipole": pred_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {"dipole": label_B, "find_dipole": 1.0}, + NB, + ) + + def make_padded(): + pred_pad = torch.cat([pred_A, pred_B], dim=0) + label_pad = torch.cat([label_A, label_B], dim=0) + mask_pad = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", + ) + return ( + {"global_dipole": pred_pad, "mask": mask_pad}, + {"dipole": label_pad, "find_dipole": 1.0}, + NP, + ) + + assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same loss as no mask (non-mixed batch).""" + pred = _rnd_t(2, TENSOR_SIZE) + label = _rnd_t(2, TENSOR_SIZE) + with_mask = { + "global_dipole": pred, + "mask": torch.ones(2, NB, dtype=torch.float64, device="cpu"), + } + without_mask = {"global_dipole": pred} + label_dict = {"dipole": label, "find_dipole": 1.0} + loss_m = self._loss_fn(with_mask, label_dict, NB) + loss_nm = self._loss_fn(without_mask, label_dict, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + # --------------------------------------------------------------------------- # Task 1: unit tests for TaskLoss._inject_atom_mask # --------------------------------------------------------------------------- From 29786a130b3a13b0160a03b5bd66d31e5b958595 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 00:49:36 +0800 Subject: [PATCH 03/16] fix(loss): per-frame normalize the energy loss to exclude mixed_type padding Apply per-frame normalization to every loss term (energy, force, virial, atom_ener, atom_pref, gen_force) in both deepmd/dpmodel/loss/ener.py and deepmd/pt/loss/ener.py. When model_dict["mask"] is present (mixed_type batch), ghost-padded atoms are excluded from the loss signal via Idiom 1 (masked per-atom mean) for force/atom_ener/atom_pref and Idiom 2 (extensive inv**norm_exp weighting) for energy/virial; all sub-branches (mse/mae/huber, intensive/non-intensive) are covered. Non-mixed callers that never inject a mask hit the unchanged else-branch and are bit-identical to the pre-fix behavior. 40 new grad-accumulation-invariant unit tests (20 dpmodel + 20 pt) verify the fix and the no-op property. --- deepmd/dpmodel/loss/ener.py | 444 +++++++++--- deepmd/pt/loss/ener.py | 462 +++++++++--- .../tests/common/dpmodel/test_loss_padding.py | 671 ++++++++++++++++++ source/tests/pt/test_loss_padding.py | 586 +++++++++++++++ 4 files changed, 1968 insertions(+), 195 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 7515f19b9a..ff155b442e 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -223,6 +223,20 @@ def call( atom_pref, ) + # Per-frame mask: recover real-atom count per frame when mask is provided. + # maskf[nf, nloc] = 1.0 for real atoms, 0.0 for ghosts. + if "mask" in model_dict: + maskf = xp.astype(model_dict["mask"], energy.dtype) # [nf, nloc] + real_natoms = xp.sum(maskf, axis=-1) # [nf] + inv = xp.reshape(1.0 / real_natoms, (-1,)) # [nf] + _nf = maskf.shape[0] + _nloc = maskf.shape[1] + else: + maskf = None + inv = None + _nf = None + _nloc = None + if self.enable_atom_ener_coeff: # when ener_coeff (\nu) is defined, the energy is defined as # E = \sum_i \nu_i E_i @@ -280,73 +294,195 @@ def call( if self.has_e: if self.loss_func == "mse": l2_ener_loss = xp.mean(xp.square(energy - energy_hat)) - if not self.use_huber: - loss += atom_norm_ener**norm_exp * (pref_e * l2_ener_loss) + if maskf is not None: + # Idiom 2 (extensive): per-frame normalization by real-atom count. + se = xp.square(energy - energy_hat) # [nf, k] + per_frame = xp.mean(xp.reshape(se, (_nf, -1)), axis=-1) # [nf] + if not self.use_huber: + loss += pref_e * xp.mean(per_frame * inv**norm_exp) + else: + inv_col = xp.reshape(inv, (_nf, 1)) # [nf, 1] + l_huber_loss = custom_huber_loss( + inv_col * energy, + inv_col * energy_hat, + delta=self._huber_delta_energy, + ) + loss += pref_e * l_huber_loss + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(xp.mean(per_frame * inv**2)), find_energy + ) else: - l_huber_loss = custom_huber_loss( - atom_norm_ener * energy, - atom_norm_ener * energy_hat, - delta=self._huber_delta_energy, + if not self.use_huber: + loss += atom_norm_ener**norm_exp * (pref_e * l2_ener_loss) + else: + l_huber_loss = custom_huber_loss( + atom_norm_ener * energy, + atom_norm_ener * energy_hat, + delta=self._huber_delta_energy, + ) + loss += pref_e * l_huber_loss + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(l2_ener_loss) * atom_norm_ener, find_energy ) - loss += pref_e * l_huber_loss - more_loss["rmse_e"] = self.display_if_exist( - xp.sqrt(l2_ener_loss) * atom_norm_ener, find_energy - ) elif self.loss_func == "mae": l1_ener_loss = xp.mean(xp.abs(energy - energy_hat)) - loss += atom_norm_ener * (pref_e * l1_ener_loss) - more_loss["mae_e"] = self.display_if_exist( - l1_ener_loss * atom_norm_ener, find_energy - ) + if maskf is not None: + abs_e = xp.abs(energy - energy_hat) # [nf, k] + per_frame_ae = xp.mean( + xp.reshape(abs_e, (_nf, -1)), axis=-1 + ) # [nf] + l1_ener_masked = xp.mean(per_frame_ae * inv) + loss += pref_e * l1_ener_masked + more_loss["mae_e"] = self.display_if_exist( + l1_ener_masked, find_energy + ) + else: + loss += atom_norm_ener * (pref_e * l1_ener_loss) + more_loss["mae_e"] = self.display_if_exist( + l1_ener_loss * atom_norm_ener, find_energy + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for energy loss." ) if mae: - mae_e = xp.mean(xp.abs(energy - energy_hat)) * atom_norm_ener + if maskf is not None: + abs_e = xp.abs(energy - energy_hat) + per_frame_ae = xp.mean(xp.reshape(abs_e, (_nf, -1)), axis=-1) + mae_e = xp.mean(per_frame_ae * inv) + else: + mae_e = xp.mean(xp.abs(energy - energy_hat)) * atom_norm_ener more_loss["mae_e"] = self.display_if_exist(mae_e, find_energy) mae_e_all = xp.mean(xp.abs(energy - energy_hat)) more_loss["mae_e_all"] = self.display_if_exist(mae_e_all, find_energy) if self.has_f: if self.loss_func == "mse": l2_force_loss = xp.mean(xp.square(diff_f)) - if not self.use_huber: - loss += pref_f * l2_force_loss + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3). + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) # [nf, nloc, 3] + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1] + if not self.use_huber: + sq_f = xp.square(diff_f_3d) * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum( + xp.reshape(sq_f, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * 3 # [nf] + l2_force_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_f * l2_force_masked + else: + if not self.f_use_norm: + abs_e = xp.abs(diff_f_3d) + quad = 0.5 * xp.square(diff_f_3d) + lin = self._huber_delta_force * ( + abs_e - 0.5 * self._huber_delta_force + ) + huber_elem = xp.where( + abs_e <= self._huber_delta_force, quad, lin + ) + huber_masked = huber_elem * maskf_col + else: + diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3)) + norm_2d = xp.reshape( + xp.linalg.vector_norm( + xp.reshape(diff_3, (-1, 3)), axis=1 + ), + (_nf, _nloc), + ) + abs_n = norm_2d + quad_n = 0.5 * xp.square(norm_2d) + lin_n = self._huber_delta_force * ( + abs_n - 0.5 * self._huber_delta_force + ) + huber_n = xp.where( + abs_n <= self._huber_delta_force, quad_n, lin_n + ) + huber_masked = xp.reshape(huber_n * maskf, (_nf, _nloc, 1)) + per_frame_sum = xp.sum( + xp.reshape(huber_masked, (_nf, -1)), axis=-1 + ) + if not self.f_use_norm: + per_frame_dof = xp.sum(maskf, axis=-1) * 3 + else: + per_frame_dof = xp.sum(maskf, axis=-1) + l_huber_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_f * l_huber_masked + more_loss["rmse_f"] = self.display_if_exist( + xp.sqrt(l2_force_loss), find_force + ) else: + if not self.use_huber: + loss += pref_f * l2_force_loss + else: + if not self.f_use_norm: + l_huber_loss = custom_huber_loss( + xp.reshape(force, (-1,)), + xp.reshape(force_hat, (-1,)), + delta=self._huber_delta_force, + ) + else: + force_diff_3 = xp.reshape(force_hat - force, (-1, 3)) + force_diff_norm = xp.reshape( + xp.linalg.vector_norm(force_diff_3, axis=1), (-1, 1) + ) + l_huber_loss = custom_huber_loss( + force_diff_norm, + xp.zeros_like(force_diff_norm), + delta=self._huber_delta_force, + ) + loss += pref_f * l_huber_loss + more_loss["rmse_f"] = self.display_if_exist( + xp.sqrt(l2_force_loss), find_force + ) + elif self.loss_func == "mae": + if maskf is not None: + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) if not self.f_use_norm: - l_huber_loss = custom_huber_loss( - xp.reshape(force, (-1,)), - xp.reshape(force_hat, (-1,)), - delta=self._huber_delta_force, + abs_f = xp.abs(diff_f_3d) * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum(xp.reshape(abs_f, (_nf, -1)), axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) * 3 + l1_force_masked = xp.mean(per_frame_sum / per_frame_dof) + else: + diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3)) + norm_2d = xp.reshape( + xp.linalg.vector_norm(xp.reshape(diff_3, (-1, 3)), axis=1), + (_nf, _nloc), ) + masked_norm = norm_2d * maskf + per_frame_sum = xp.sum(masked_norm, axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) + l1_force_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_f * l1_force_masked + more_loss["mae_f"] = self.display_if_exist( + l1_force_masked, find_force + ) + else: + if not self.f_use_norm: + l1_force_loss = xp.mean(xp.abs(diff_f)) else: force_diff_3 = xp.reshape(force_hat - force, (-1, 3)) - force_diff_norm = xp.reshape( - xp.linalg.vector_norm(force_diff_3, axis=1), (-1, 1) - ) - l_huber_loss = custom_huber_loss( - force_diff_norm, - xp.zeros_like(force_diff_norm), - delta=self._huber_delta_force, + l1_force_loss = xp.mean( + xp.linalg.vector_norm(force_diff_3, axis=1) ) - loss += pref_f * l_huber_loss - more_loss["rmse_f"] = self.display_if_exist( - xp.sqrt(l2_force_loss), find_force - ) - elif self.loss_func == "mae": - if not self.f_use_norm: - l1_force_loss = xp.mean(xp.abs(diff_f)) - else: - force_diff_3 = xp.reshape(force_hat - force, (-1, 3)) - l1_force_loss = xp.mean(xp.linalg.vector_norm(force_diff_3, axis=1)) - loss += pref_f * l1_force_loss - more_loss["mae_f"] = self.display_if_exist(l1_force_loss, find_force) + loss += pref_f * l1_force_loss + more_loss["mae_f"] = self.display_if_exist( + l1_force_loss, find_force + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for force loss." ) if mae: - mae_f = xp.mean(xp.abs(diff_f)) + if maskf is not None: + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + abs_f = xp.abs(diff_f_3d) * maskf_col + per_frame_sum = xp.sum(xp.reshape(abs_f, (_nf, -1)), axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) * 3 + mae_f = xp.mean(per_frame_sum / per_frame_dof) + else: + mae_f = xp.mean(xp.abs(diff_f)) more_loss["mae_f"] = self.display_if_exist(mae_f, find_force) if self.has_v: virial_reshape = xp.reshape(virial, (-1,)) @@ -355,30 +491,70 @@ def call( l2_virial_loss = xp.mean( xp.square(virial_hat_reshape - virial_reshape), ) - if not self.use_huber: - loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + if maskf is not None: + # Idiom 2 (extensive, k=9): per-frame normalization. + v2d = xp.reshape(virial, (_nf, 9)) + v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) + se_v = xp.square(v_hat_2d - v2d) # [nf, 9] + per_frame_v = xp.mean(se_v, axis=-1) # [nf] + if not self.use_huber: + loss += pref_v * xp.mean(per_frame_v * inv**norm_exp) + else: + inv_col = xp.reshape(inv, (_nf, 1)) # [nf, 1] + l_huber_v = custom_huber_loss( + inv_col * v2d, + inv_col * v_hat_2d, + delta=self._huber_delta_virial, + ) + loss += pref_v * l_huber_v + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(xp.mean(per_frame_v * inv**2)), find_virial + ) else: - l_huber_loss = custom_huber_loss( - atom_norm * virial_reshape, - atom_norm * virial_hat_reshape, - delta=self._huber_delta_virial, + if not self.use_huber: + loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + else: + l_huber_loss = custom_huber_loss( + atom_norm * virial_reshape, + atom_norm * virial_hat_reshape, + delta=self._huber_delta_virial, + ) + loss += pref_v * l_huber_loss + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(l2_virial_loss) * atom_norm, find_virial ) - loss += pref_v * l_huber_loss - more_loss["rmse_v"] = self.display_if_exist( - xp.sqrt(l2_virial_loss) * atom_norm, find_virial - ) elif self.loss_func == "mae": l1_virial_loss = xp.mean(xp.abs(virial_hat_reshape - virial_reshape)) - loss += atom_norm * (pref_v * l1_virial_loss) - more_loss["mae_v"] = self.display_if_exist( - l1_virial_loss * atom_norm, find_virial - ) + if maskf is not None: + v2d = xp.reshape(virial, (_nf, 9)) + v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) + abs_v = xp.abs(v_hat_2d - v2d) # [nf, 9] + per_frame_v = xp.mean(abs_v, axis=-1) # [nf] + l1_virial_masked = xp.mean(per_frame_v * inv) + loss += pref_v * l1_virial_masked + more_loss["mae_v"] = self.display_if_exist( + l1_virial_masked, find_virial + ) + else: + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss * atom_norm, find_virial + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for virial loss." ) if mae: - mae_v = xp.mean(xp.abs(virial_hat_reshape - virial_reshape)) * atom_norm + if maskf is not None: + v2d = xp.reshape(virial, (_nf, 9)) + v_hat_2d = xp.reshape(virial_hat, (_nf, 9)) + abs_v = xp.abs(v_hat_2d - v2d) + per_frame_v = xp.mean(abs_v, axis=-1) + mae_v = xp.mean(per_frame_v * inv) + else: + mae_v = ( + xp.mean(xp.abs(virial_hat_reshape - virial_reshape)) * atom_norm + ) more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) if self.has_ae: atom_ener_reshape = xp.reshape(atom_ener, (-1,)) @@ -387,26 +563,67 @@ def call( l2_atom_ener_loss = xp.mean( xp.square(atom_ener_hat_reshape - atom_ener_reshape), ) - if not self.use_huber: - loss += pref_ae * l2_atom_ener_loss + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=1). + ae_2d = xp.reshape(atom_ener, (_nf, _nloc)) + ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc)) + sq_ae = xp.square(ae_hat_2d - ae_2d) * maskf # [nf, nloc] + per_frame_sum = xp.sum(sq_ae, axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) # [nf] + l2_ae_masked = xp.mean(per_frame_sum / per_frame_dof) + if not self.use_huber: + loss += pref_ae * l2_ae_masked + else: + # Huber applied element-wise then masked-mean. + diff_ae = ae_hat_2d - ae_2d + abs_ae = xp.abs(diff_ae) + quad_ae = 0.5 * xp.square(diff_ae) + lin_ae = self._huber_delta_energy * ( + abs_ae - 0.5 * self._huber_delta_energy + ) + huber_ae = xp.where( + abs_ae <= self._huber_delta_energy, quad_ae, lin_ae + ) + huber_ae_masked = huber_ae * maskf + per_frame_sum_h = xp.sum(huber_ae_masked, axis=-1) + l_huber_ae_masked = xp.mean(per_frame_sum_h / per_frame_dof) + loss += pref_ae * l_huber_ae_masked + more_loss["rmse_ae"] = self.display_if_exist( + xp.sqrt(l2_ae_masked), find_atom_ener + ) else: - l_huber_loss = custom_huber_loss( - atom_ener_reshape, - atom_ener_hat_reshape, - delta=self._huber_delta_energy, + if not self.use_huber: + loss += pref_ae * l2_atom_ener_loss + else: + l_huber_loss = custom_huber_loss( + atom_ener_reshape, + atom_ener_hat_reshape, + delta=self._huber_delta_energy, + ) + loss += pref_ae * l_huber_loss + more_loss["rmse_ae"] = self.display_if_exist( + xp.sqrt(l2_atom_ener_loss), find_atom_ener ) - loss += pref_ae * l_huber_loss - more_loss["rmse_ae"] = self.display_if_exist( - xp.sqrt(l2_atom_ener_loss), find_atom_ener - ) elif self.loss_func == "mae": l1_atom_ener_loss = xp.mean( xp.abs(atom_ener_hat_reshape - atom_ener_reshape) ) - loss += pref_ae * l1_atom_ener_loss - more_loss["mae_ae"] = self.display_if_exist( - l1_atom_ener_loss, find_atom_ener - ) + if maskf is not None: + ae_2d = xp.reshape(atom_ener, (_nf, _nloc)) + ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc)) + abs_ae = xp.abs(ae_hat_2d - ae_2d) * maskf # [nf, nloc] + per_frame_sum = xp.sum(abs_ae, axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) # [nf] + l1_ae_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_ae * l1_ae_masked + more_loss["mae_ae"] = self.display_if_exist( + l1_ae_masked, find_atom_ener + ) + else: + loss += pref_ae * l1_atom_ener_loss + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss, find_atom_ener + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for atomic energy loss." @@ -418,18 +635,47 @@ def call( l2_pref_force_loss = xp.mean( xp.multiply(xp.square(diff_f), atom_pref_reshape), ) - loss += pref_pf * l2_pref_force_loss - more_loss["rmse_pf"] = self.display_if_exist( - xp.sqrt(l2_pref_force_loss), find_atom_pref - ) + if maskf is not None: + # Idiom 1 with pref weight (ncomp=3). + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + sq_pf = xp.square(diff_f_3d) * pf_3d * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum( + xp.reshape(sq_pf, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * 3 # [nf] + l2_pf_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_pf * l2_pf_masked + more_loss["rmse_pf"] = self.display_if_exist( + xp.sqrt(l2_pf_masked), find_atom_pref + ) + else: + loss += pref_pf * l2_pref_force_loss + more_loss["rmse_pf"] = self.display_if_exist( + xp.sqrt(l2_pref_force_loss), find_atom_pref + ) elif self.loss_func == "mae": l1_pref_force_loss = xp.mean( xp.multiply(xp.abs(diff_f), atom_pref_reshape) ) - loss += pref_pf * l1_pref_force_loss - more_loss["mae_pf"] = self.display_if_exist( - l1_pref_force_loss, find_atom_pref - ) + if maskf is not None: + diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) + pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + abs_pf = xp.abs(diff_f_3d) * pf_3d * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum(xp.reshape(abs_pf, (_nf, -1)), axis=-1) + per_frame_dof = xp.sum(maskf, axis=-1) * 3 + l1_pf_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_pf * l1_pf_masked + more_loss["mae_pf"] = self.display_if_exist( + l1_pf_masked, find_atom_pref + ) + else: + loss += pref_pf * l1_pref_force_loss + more_loss["mae_pf"] = self.display_if_exist( + l1_pref_force_loss, find_atom_pref + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for atom prefactor force loss." @@ -437,22 +683,40 @@ def call( if self.has_gf: find_drdq = label_dict["find_drdq"] drdq = label_dict["drdq"] - force_reshape_nframes = xp.reshape(force, (-1, natoms * 3)) - force_hat_reshape_nframes = xp.reshape(force_hat, (-1, natoms * 3)) - drdq_reshape = xp.reshape( - drdq, (-1, natoms * 3, self.numb_generalized_coord) - ) - # "bij,bi->bj" einsum replaced with array-API-compatible ops - gen_force_hat = xp.sum( - drdq_reshape * force_hat_reshape_nframes[:, :, None], axis=1 - ) - gen_force = xp.sum(drdq_reshape * force_reshape_nframes[:, :, None], axis=1) - diff_gen_force = gen_force_hat - gen_force - l2_gen_force_loss = xp.mean(xp.square(diff_gen_force)) pref_gf = find_drdq * ( self.limit_pref_gf + (self.start_pref_gf - self.limit_pref_gf) * lr_ratio ) + if maskf is not None: + # Mask per-atom forces before projecting onto generalized coords + # so ghost atoms don't contribute to the generalized force. + force_3d = xp.reshape(force, (_nf, _nloc, 3)) + force_hat_3d = xp.reshape(force_hat, (_nf, _nloc, 3)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + masked_f = force_3d * maskf_col # [nf, nloc, 3] + masked_f_hat = force_hat_3d * maskf_col # [nf, nloc, 3] + f_flat = xp.reshape(masked_f, (_nf, _nloc * 3)) + f_hat_flat = xp.reshape(masked_f_hat, (_nf, _nloc * 3)) + drdq_reshape = xp.reshape( + drdq, (_nf, _nloc * 3, self.numb_generalized_coord) + ) + gen_force = xp.sum(drdq_reshape * f_flat[:, :, None], axis=1) + gen_force_hat = xp.sum(drdq_reshape * f_hat_flat[:, :, None], axis=1) + else: + force_reshape_nframes = xp.reshape(force, (-1, natoms * 3)) + force_hat_reshape_nframes = xp.reshape(force_hat, (-1, natoms * 3)) + drdq_reshape = xp.reshape( + drdq, (-1, natoms * 3, self.numb_generalized_coord) + ) + gen_force_hat = xp.sum( + drdq_reshape * force_hat_reshape_nframes[:, :, None], axis=1 + ) + gen_force = xp.sum( + drdq_reshape * force_reshape_nframes[:, :, None], axis=1 + ) + # "bij,bi->bj" einsum replaced with array-API-compatible ops + diff_gen_force = gen_force_hat - gen_force + l2_gen_force_loss = xp.mean(xp.square(diff_gen_force)) loss += pref_gf * l2_gen_force_loss more_loss["rmse_gf"] = self.display_if_exist( xp.sqrt(l2_gen_force_loss), find_drdq diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 64d801ea51..a7b7683363 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -239,6 +239,19 @@ def forward( # more_loss['log_keys'] = [] # showed when validation on the fly # more_loss['test_keys'] = [] # showed when doing dp test atom_norm = 1.0 / natoms + + # Per-frame mask: recover real-atom count per frame when mask is provided. + if "mask" in model_pred: + maskf = model_pred["mask"] # [nf, nloc], float + real_natoms_f = torch.sum(maskf, dim=-1) # [nf] + inv = (1.0 / real_natoms_f).reshape(-1) # [nf] + _nf = maskf.shape[0] + _nloc = maskf.shape[1] + else: + maskf = None + inv = None + _nf = None + _nloc = None # Normalization exponent controls loss scaling with system size: # - norm_exp=2 (intensive_ener_virial=True): loss uses 1/N² scaling, making it independent of system size # - norm_exp=1 (intensive_ener_virial=False, legacy): loss uses 1/N scaling, which varies with system size @@ -267,19 +280,38 @@ def forward( more_loss["l2_ener_loss"] = self.display_if_exist( l2_ener_loss.detach(), find_energy ) - if not self.use_huber: - loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) + if maskf is not None: + # Idiom 2 (extensive): per-frame normalization. + se = torch.square(energy_pred - energy_label) # [nf, k] + per_frame = torch.mean(se.reshape(_nf, -1), dim=-1) # [nf] + if not self.use_huber: + loss += pref_e * torch.mean(per_frame * inv**norm_exp) + else: + inv_col = inv.reshape(_nf, 1) + l_huber_loss = custom_huber_loss( + inv_col * energy_pred, + inv_col * energy_label, + delta=self._huber_delta_energy, + ) + loss += pref_e * l_huber_loss + rmse_e = torch.sqrt(torch.mean(per_frame * inv**2)) + more_loss["rmse_e"] = self.display_if_exist( + rmse_e.detach(), find_energy + ) else: - l_huber_loss = custom_huber_loss( - atom_norm * energy_pred, - atom_norm * energy_label, - delta=self._huber_delta_energy, + if not self.use_huber: + loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) + else: + l_huber_loss = custom_huber_loss( + atom_norm * energy_pred, + atom_norm * energy_label, + delta=self._huber_delta_energy, + ) + loss += pref_e * l_huber_loss + rmse_e = l2_ener_loss.sqrt() * atom_norm + more_loss["rmse_e"] = self.display_if_exist( + rmse_e.detach(), find_energy ) - loss += pref_e * l_huber_loss - rmse_e = l2_ener_loss.sqrt() * atom_norm - more_loss["rmse_e"] = self.display_if_exist( - rmse_e.detach(), find_energy - ) # more_loss['log_keys'].append('rmse_e') elif self.loss_func == "mae": l1_ener_loss = F.l1_loss( @@ -287,18 +319,34 @@ def forward( energy_label.reshape(-1), reduction="mean", ) - loss += atom_norm * (pref_e * l1_ener_loss) - more_loss["mae_e"] = self.display_if_exist( - l1_ener_loss.detach() * atom_norm, - find_energy, - ) + if maskf is not None: + abs_e = torch.abs(energy_pred - energy_label) + per_frame_ae = torch.mean(abs_e.reshape(_nf, -1), dim=-1) + l1_ener_masked = torch.mean(per_frame_ae * inv) + loss += pref_e * l1_ener_masked + more_loss["mae_e"] = self.display_if_exist( + l1_ener_masked.detach(), find_energy + ) + else: + loss += atom_norm * (pref_e * l1_ener_loss) + more_loss["mae_e"] = self.display_if_exist( + l1_ener_loss.detach() * atom_norm, + find_energy, + ) # more_loss['log_keys'].append('rmse_e') else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for energy loss." ) if mae: - mae_e = torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm + if maskf is not None: + abs_e = torch.abs(energy_pred - energy_label) + per_frame_ae = torch.mean(abs_e.reshape(_nf, -1), dim=-1) + mae_e = torch.mean(per_frame_ae * inv) + else: + mae_e = ( + torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm + ) more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy) mae_e_all = torch.mean(torch.abs(energy_pred - energy_label)) more_loss["mae_e_all"] = self.display_if_exist( @@ -330,56 +378,131 @@ def forward( more_loss["l2_force_loss"] = self.display_if_exist( l2_force_loss.detach(), find_force ) - if not self.use_huber: - loss += (pref_f * l2_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3). + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + if not self.use_huber: + sq_f = torch.square(diff_f_3d) * maskf_col + per_frame_sum = sq_f.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + l2_f_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_f * l2_f_masked).to(GLOBAL_PT_FLOAT_PRECISION) + else: + if not self.f_use_norm: + abs_e = torch.abs(diff_f_3d) + quad = 0.5 * torch.square(diff_f_3d) + lin = self._huber_delta_force * ( + abs_e - 0.5 * self._huber_delta_force + ) + huber_elem = torch.where( + abs_e <= self._huber_delta_force, quad, lin + ) + huber_masked = huber_elem * maskf_col + per_frame_dof = maskf.sum(dim=-1) * 3 + else: + diff_3 = (force_label - force_pred).reshape( + _nf, _nloc, 3 + ) + norm_2d = torch.linalg.vector_norm( + diff_3.reshape(-1, 3), ord=2, dim=1 + ).reshape(_nf, _nloc) + abs_n = norm_2d + quad_n = 0.5 * torch.square(norm_2d) + lin_n = self._huber_delta_force * ( + abs_n - 0.5 * self._huber_delta_force + ) + huber_n = torch.where( + abs_n <= self._huber_delta_force, quad_n, lin_n + ) + huber_masked = (huber_n * maskf).reshape(_nf, _nloc, 1) + per_frame_dof = maskf.sum(dim=-1) + per_frame_sum = huber_masked.reshape(_nf, -1).sum(dim=-1) + l_huber_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += pref_f * l_huber_masked else: - if not self.f_use_norm: - l_huber_loss = custom_huber_loss( - force_pred.reshape(-1), - force_label.reshape(-1), - delta=self._huber_delta_force, + if not self.use_huber: + loss += (pref_f * l2_force_loss).to( + GLOBAL_PT_FLOAT_PRECISION ) else: - force_diff_norm = torch.linalg.vector_norm( - (force_label - force_pred).reshape(-1, 3), - ord=2, - dim=1, - keepdim=True, - ) - l_huber_loss = custom_huber_loss( - force_diff_norm, - torch.zeros_like(force_diff_norm), - delta=self._huber_delta_force, - ) - loss += pref_f * l_huber_loss + if not self.f_use_norm: + l_huber_loss = custom_huber_loss( + force_pred.reshape(-1), + force_label.reshape(-1), + delta=self._huber_delta_force, + ) + else: + force_diff_norm = torch.linalg.vector_norm( + (force_label - force_pred).reshape(-1, 3), + ord=2, + dim=1, + keepdim=True, + ) + l_huber_loss = custom_huber_loss( + force_diff_norm, + torch.zeros_like(force_diff_norm), + delta=self._huber_delta_force, + ) + loss += pref_f * l_huber_loss rmse_f = l2_force_loss.sqrt() more_loss["rmse_f"] = self.display_if_exist( rmse_f.detach(), find_force ) elif self.loss_func == "mae": - if not self.f_use_norm: - l1_force_loss = F.l1_loss( - force_label.reshape(-1), - force_pred.reshape(-1), - reduction="mean", + if maskf is not None: + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + if not self.f_use_norm: + abs_f = torch.abs(diff_f_3d) * maskf_col + per_frame_sum = abs_f.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + l1_f_masked = torch.mean(per_frame_sum / per_frame_dof) + else: + diff_3 = (force_label - force_pred).reshape(_nf, _nloc, 3) + norm_2d = torch.linalg.vector_norm( + diff_3.reshape(-1, 3), ord=2, dim=1 + ).reshape(_nf, _nloc) + masked_norm = norm_2d * maskf + per_frame_sum = masked_norm.sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) + l1_f_masked = torch.mean(per_frame_sum / per_frame_dof) + more_loss["mae_f"] = self.display_if_exist( + l1_f_masked.detach(), find_force ) + loss += (pref_f * l1_f_masked).to(GLOBAL_PT_FLOAT_PRECISION) else: - l1_force_loss = torch.linalg.vector_norm( - (force_label - force_pred).reshape(-1, 3), - ord=2, - dim=1, - keepdim=True, - ).mean() - more_loss["mae_f"] = self.display_if_exist( - l1_force_loss.detach(), find_force - ) - loss += (pref_f * l1_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) + if not self.f_use_norm: + l1_force_loss = F.l1_loss( + force_label.reshape(-1), + force_pred.reshape(-1), + reduction="mean", + ) + else: + l1_force_loss = torch.linalg.vector_norm( + (force_label - force_pred).reshape(-1, 3), + ord=2, + dim=1, + keepdim=True, + ).mean() + more_loss["mae_f"] = self.display_if_exist( + l1_force_loss.detach(), find_force + ) + loss += (pref_f * l1_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for force loss." ) if mae: - mae_f = torch.mean(torch.abs(diff_f)) + if maskf is not None: + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + abs_f = torch.abs(diff_f_3d) * maskf_col + per_frame_sum = abs_f.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + mae_f = torch.mean(per_frame_sum / per_frame_dof) + else: + mae_f = torch.mean(torch.abs(diff_f)) more_loss["mae_f"] = self.display_if_exist( mae_f.detach(), find_force ) @@ -400,17 +523,49 @@ def forward( more_loss["l2_pref_force_loss"] = self.display_if_exist( l2_pref_force_loss.detach(), find_atom_pref ) - loss += (pref_pf * l2_pref_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) - rmse_pf = l2_pref_force_loss.sqrt() - more_loss["rmse_pf"] = self.display_if_exist( - rmse_pf.detach(), find_atom_pref - ) + if maskf is not None: + # Idiom 1 with pref weight (ncomp=3). + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + pf_3d = atom_pref.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + sq_pf = torch.square(diff_f_3d) * pf_3d * maskf_col + per_frame_sum = sq_pf.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + l2_pf_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_pf * l2_pf_masked).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_pf = l2_pf_masked.sqrt() + more_loss["rmse_pf"] = self.display_if_exist( + rmse_pf.detach(), find_atom_pref + ) + else: + loss += (pref_pf * l2_pref_force_loss).to( + GLOBAL_PT_FLOAT_PRECISION + ) + rmse_pf = l2_pref_force_loss.sqrt() + more_loss["rmse_pf"] = self.display_if_exist( + rmse_pf.detach(), find_atom_pref + ) elif self.loss_func == "mae": l1_pref_force_loss = (torch.abs(diff_f) * atom_pref_reshape).mean() - loss += (pref_pf * l1_pref_force_loss).to(GLOBAL_PT_FLOAT_PRECISION) - more_loss["mae_pf"] = self.display_if_exist( - l1_pref_force_loss.detach(), find_atom_pref - ) + if maskf is not None: + diff_f_3d = diff_f.reshape(_nf, _nloc, 3) + pf_3d = atom_pref.reshape(_nf, _nloc, 3) + maskf_col = maskf.reshape(_nf, _nloc, 1) + abs_pf = torch.abs(diff_f_3d) * pf_3d * maskf_col + per_frame_sum = abs_pf.reshape(_nf, -1).sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) * 3 + l1_pf_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_pf * l1_pf_masked).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_pf"] = self.display_if_exist( + l1_pf_masked.detach(), find_atom_pref + ) + else: + loss += (pref_pf * l1_pref_force_loss).to( + GLOBAL_PT_FLOAT_PRECISION + ) + more_loss["mae_pf"] = self.display_if_exist( + l1_pref_force_loss.detach(), find_atom_pref + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for atom prefactor force loss." @@ -420,15 +575,35 @@ def forward( drdq = label["drdq"] find_drdq = label.get("find_drdq", 0.0) pref_gf = pref_gf * find_drdq - force_reshape_nframes = force_pred.reshape(-1, natoms * 3) - force_label_reshape_nframes = force_label.reshape(-1, natoms * 3) - drdq_reshape = drdq.reshape(-1, natoms * 3, self.numb_generalized_coord) - gen_force_label = torch.einsum( - "bij,bi->bj", drdq_reshape, force_label_reshape_nframes - ) - gen_force = torch.einsum( - "bij,bi->bj", drdq_reshape, force_reshape_nframes - ) + if maskf is not None: + # Mask per-atom forces before projecting onto generalized coords. + f_3d = force_pred.reshape(_nf, _nloc, 3) * maskf.reshape( + _nf, _nloc, 1 + ) + f_hat_3d = force_label.reshape(_nf, _nloc, 3) * maskf.reshape( + _nf, _nloc, 1 + ) + f_flat = f_3d.reshape(_nf, _nloc * 3) + f_hat_flat = f_hat_3d.reshape(_nf, _nloc * 3) + drdq_reshape = drdq.reshape( + _nf, _nloc * 3, self.numb_generalized_coord + ) + gen_force = torch.einsum("bij,bi->bj", drdq_reshape, f_flat) + gen_force_label = torch.einsum( + "bij,bi->bj", drdq_reshape, f_hat_flat + ) + else: + force_reshape_nframes = force_pred.reshape(-1, natoms * 3) + force_label_reshape_nframes = force_label.reshape(-1, natoms * 3) + drdq_reshape = drdq.reshape( + -1, natoms * 3, self.numb_generalized_coord + ) + gen_force_label = torch.einsum( + "bij,bi->bj", drdq_reshape, force_label_reshape_nframes + ) + gen_force = torch.einsum( + "bij,bi->bj", drdq_reshape, force_reshape_nframes + ) diff_gen_force = gen_force_label - gen_force l2_gen_force_loss = torch.square(diff_gen_force).mean() if not self.inference: @@ -444,43 +619,78 @@ def forward( if self.has_v and "virial" in model_pred and "virial" in label: find_virial = label.get("find_virial", 0.0) pref_v = pref_v * find_virial - diff_v = label["virial"] - model_pred["virial"].reshape(-1, 9) + v2d = model_pred["virial"].reshape(-1, 9) + v_hat_2d = label["virial"].reshape(-1, 9) + diff_v = v_hat_2d - v2d if self.loss_func == "mse": l2_virial_loss = torch.mean(torch.square(diff_v)) if not self.inference: more_loss["l2_virial_loss"] = self.display_if_exist( l2_virial_loss.detach(), find_virial ) - if not self.use_huber: - loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + if maskf is not None: + # Idiom 2 (extensive, k=9): per-frame normalization. + se_v = torch.square(diff_v) # [nf, 9] + per_frame_v = torch.mean(se_v, dim=-1) # [nf] + if not self.use_huber: + loss += pref_v * torch.mean(per_frame_v * inv**norm_exp) + else: + inv_col = inv.reshape(_nf, 1) + l_huber_v = custom_huber_loss( + inv_col * v2d, + inv_col * v_hat_2d, + delta=self._huber_delta_virial, + ) + loss += pref_v * l_huber_v + rmse_v = torch.sqrt(torch.mean(per_frame_v * inv**2)) + more_loss["rmse_v"] = self.display_if_exist( + rmse_v.detach(), find_virial + ) else: - l_huber_loss = custom_huber_loss( - atom_norm * model_pred["virial"].reshape(-1), - atom_norm * label["virial"].reshape(-1), - delta=self._huber_delta_virial, + if not self.use_huber: + loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + else: + l_huber_loss = custom_huber_loss( + atom_norm * v2d.reshape(-1), + atom_norm * v_hat_2d.reshape(-1), + delta=self._huber_delta_virial, + ) + loss += pref_v * l_huber_loss + rmse_v = l2_virial_loss.sqrt() * atom_norm + more_loss["rmse_v"] = self.display_if_exist( + rmse_v.detach(), find_virial ) - loss += pref_v * l_huber_loss - rmse_v = l2_virial_loss.sqrt() * atom_norm - more_loss["rmse_v"] = self.display_if_exist( - rmse_v.detach(), find_virial - ) elif self.loss_func == "mae": l1_virial_loss = F.l1_loss( - label["virial"].reshape(-1), - model_pred["virial"].reshape(-1), + v_hat_2d.reshape(-1), + v2d.reshape(-1), reduction="mean", ) - loss += atom_norm * (pref_v * l1_virial_loss) - more_loss["mae_v"] = self.display_if_exist( - l1_virial_loss.detach() * atom_norm, - find_virial, - ) + if maskf is not None: + abs_v = torch.abs(diff_v) # [nf, 9] + per_frame_v = torch.mean(abs_v, dim=-1) # [nf] + l1_v_masked = torch.mean(per_frame_v * inv) + loss += pref_v * l1_v_masked + more_loss["mae_v"] = self.display_if_exist( + l1_v_masked.detach(), find_virial + ) + else: + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss.detach() * atom_norm, + find_virial, + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for virial loss." ) if mae: - mae_v = torch.mean(torch.abs(diff_v)) * atom_norm + if maskf is not None: + abs_v = torch.abs(diff_v) + per_frame_v = torch.mean(abs_v, dim=-1) + mae_v = torch.mean(per_frame_v * inv) + else: + mae_v = torch.mean(torch.abs(diff_v)) * atom_norm more_loss["mae_v"] = self.display_if_exist(mae_v.detach(), find_virial) if self.has_ae and "atom_energy" in model_pred and "atom_ener" in label: @@ -499,29 +709,71 @@ def forward( more_loss["l2_atom_ener_loss"] = self.display_if_exist( l2_atom_ener_loss.detach(), find_atom_ener ) - if not self.use_huber: - loss += (pref_ae * l2_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=1). + ae_2d = atom_ener.reshape(_nf, _nloc) + ae_hat_2d = atom_ener_label.reshape(_nf, _nloc) + sq_ae = torch.square(ae_hat_2d - ae_2d) * maskf # [nf, nloc] + per_frame_sum = sq_ae.sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) # [nf] + l2_ae_masked = torch.mean(per_frame_sum / per_frame_dof) + if not self.use_huber: + loss += (pref_ae * l2_ae_masked).to(GLOBAL_PT_FLOAT_PRECISION) + else: + diff_ae = ae_hat_2d - ae_2d + abs_ae = torch.abs(diff_ae) + quad_ae = 0.5 * torch.square(diff_ae) + lin_ae = self._huber_delta_energy * ( + abs_ae - 0.5 * self._huber_delta_energy + ) + huber_ae = torch.where( + abs_ae <= self._huber_delta_energy, quad_ae, lin_ae + ) + huber_ae_m = huber_ae * maskf + l_huber_ae = torch.mean(huber_ae_m.sum(dim=-1) / per_frame_dof) + loss += pref_ae * l_huber_ae + rmse_ae = l2_ae_masked.sqrt() + more_loss["rmse_ae"] = self.display_if_exist( + rmse_ae.detach(), find_atom_ener + ) else: - l_huber_loss = custom_huber_loss( - atom_ener_reshape, - atom_ener_label_reshape, - delta=self._huber_delta_energy, + if not self.use_huber: + loss += (pref_ae * l2_atom_ener_loss).to( + GLOBAL_PT_FLOAT_PRECISION + ) + else: + l_huber_loss = custom_huber_loss( + atom_ener_reshape, + atom_ener_label_reshape, + delta=self._huber_delta_energy, + ) + loss += pref_ae * l_huber_loss + rmse_ae = l2_atom_ener_loss.sqrt() + more_loss["rmse_ae"] = self.display_if_exist( + rmse_ae.detach(), find_atom_ener ) - loss += pref_ae * l_huber_loss - rmse_ae = l2_atom_ener_loss.sqrt() - more_loss["rmse_ae"] = self.display_if_exist( - rmse_ae.detach(), find_atom_ener - ) elif self.loss_func == "mae": l1_atom_ener_loss = F.l1_loss( atom_ener_reshape, atom_ener_label_reshape, reduction="mean", ) - loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) - more_loss["mae_ae"] = self.display_if_exist( - l1_atom_ener_loss.detach(), find_atom_ener - ) + if maskf is not None: + ae_2d = atom_ener.reshape(_nf, _nloc) + ae_hat_2d = atom_ener_label.reshape(_nf, _nloc) + abs_ae = torch.abs(ae_hat_2d - ae_2d) * maskf + per_frame_sum = abs_ae.sum(dim=-1) + per_frame_dof = maskf.sum(dim=-1) + l1_ae_masked = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_ae * l1_ae_masked).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_ae"] = self.display_if_exist( + l1_ae_masked.detach(), find_atom_ener + ) + else: + loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss.detach(), find_atom_ener + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for atomic energy loss." diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 4faac008ed..4b009a78be 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -18,6 +18,9 @@ from deepmd.dpmodel.loss.dos import ( DOSLoss, ) +from deepmd.dpmodel.loss.ener import ( + EnergyLoss, +) from deepmd.dpmodel.loss.tensor import ( TensorLoss, ) @@ -428,3 +431,671 @@ def test_no_op_for_non_mixed(self): assert np.isclose(float(loss_m), float(loss_nm)), ( f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" ) + + +# --------------------------------------------------------------------------- +# Task 3: EnergyLoss -- energy, force, virial, atom_ener, atom_pref +# --------------------------------------------------------------------------- + + +def _full_ener_dicts(nf, nloc, energy_pred, energy_label, mask=None, **overrides): + """Build complete model_pred and label_dict for EnergyLoss.call. + + EnergyLoss.call fetches energy/force/virial/atom_energy unconditionally, + so all keys must be present regardless of which term is under test. + """ + model_pred = { + "energy": energy_pred, # [nf, 1] + "force": np.zeros((nf, nloc, 3), dtype=np.float64), + "virial": np.zeros((nf, 9), dtype=np.float64), + "atom_energy": np.zeros((nf, nloc, 1), dtype=np.float64), + } + label_dict = { + "energy": energy_label, # [nf, 1] + "force": np.zeros((nf, nloc, 3), dtype=np.float64), + "virial": np.zeros((nf, 9), dtype=np.float64), + "atom_ener": np.zeros((nf, nloc, 1), dtype=np.float64), + "atom_pref": np.zeros((nf, nloc * 3), dtype=np.float64), + "find_energy": 1.0, + "find_force": 0.0, + "find_virial": 0.0, + "find_atom_ener": 0.0, + "find_atom_pref": 0.0, + } + if mask is not None: + model_pred["mask"] = mask + model_pred.update({k: v for k, v in overrides.items() if k in model_pred}) + label_dict.update({k: v for k, v in overrides.items() if k in label_dict}) + return model_pred, label_dict + + +def _padded_force(f_A, f_B): + """Stack force arrays for a 2-frame padded batch (NA<=NP, NB==NP).""" + f_A_pad = np.zeros((NP, 3), dtype=np.float64) + f_A_pad[:NA] = f_A + return np.stack([f_A_pad, f_B], axis=0) # [2, NP, 3] + + +def _padded_atom(arr_A, arr_B, ncomp): + """Pad arr_A from [NA, ncomp] to [NP, ncomp] with zeros, stack with arr_B.""" + pad = np.zeros((NP, ncomp), dtype=np.float64) + pad[:NA] = arr_A + return np.stack([pad, arr_B], axis=0) # [2, NP, ncomp] + + +def _padded_atom_flat(arr_A, arr_B, ncomp): + """Pad then reshape to [2, NP*ncomp] (for atom_pref shape).""" + return _padded_atom(arr_A, arr_B, ncomp).reshape(2, NP * ncomp) + + +_MASK_PAD = np.array( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64 +) # [2, NP] + + +class TestDPModelEnergyLossEnerGradAccum: + """Idiom 2 (extensive) for the energy (has_e) term in EnergyLoss. + + Covers: mse (norm_exp=1 and 2), mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", intensive=False, use_huber=False): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, e_A, e_A_hat, e_B, e_B_hat): + def make_A(): + p, l = _full_ener_dicts( + 1, NA, e_A, e_A_hat, mask=np.ones((1, NA), dtype=np.float64) + ) + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, NB, e_B, e_B_hat, mask=np.ones((1, NB), dtype=np.float64) + ) + return p, l, NB + + def make_padded(): + e_pad = np.concatenate([e_A, e_B], axis=0) # [2, 1] + e_hat_pad = np.concatenate([e_A_hat, e_B_hat], axis=0) + p, l = _full_ener_dicts(2, NP, e_pad, e_hat_pad, mask=_MASK_PAD) + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Energy MSE norm_exp=1 meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(loss_func="mse", intensive=False), + e_A, + e_A_hat, + e_B, + e_B_hat, + ) + + def test_mse_intensive_grad_accum(self): + """Energy MSE norm_exp=2 meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(loss_func="mse", intensive=True), + e_A, + e_A_hat, + e_B, + e_B_hat, + ) + + def test_mae_grad_accum(self): + """Energy MAE meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(loss_func="mae"), + e_A, + e_A_hat, + e_B, + e_B_hat, + ) + + def test_huber_grad_accum(self): + """Energy Huber meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(loss_func="mse", use_huber=True), + e_A, + e_A_hat, + e_B, + e_B_hat, + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same energy loss as no mask.""" + e = _rnd(1, 1) + e_hat = _rnd(1, 1) + loss_obj = self._make_loss() + p_mask, l_mask = _full_ener_dicts( + 1, NP, e, e_hat, mask=np.ones((1, NP), dtype=np.float64) + ) + p_nomask, l_nomask = _full_ener_dicts(1, NP, e, e_hat) + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nomask, l_nomask, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnergyLossForceGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=3) for the force (has_f) term. + + Covers: mse, mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", use_huber=False, f_use_norm=False): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + loss_func=loss_func, + use_huber=use_huber, + f_use_norm=f_use_norm, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): + def make_A(): + p, l = _full_ener_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NA), dtype=np.float64), + ) + p["force"] = f_A[None] # [1, NA, 3] + l["force"] = f_A_hat[None] + l["find_force"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NB), dtype=np.float64), + ) + p["force"] = f_B[None] + l["force"] = f_B_hat[None] + l["find_force"] = 1.0 + return p, l, NB + + def make_padded(): + f_pad = _padded_force(f_A, f_B) # [2, NP, 3] + f_hat_pad = _padded_force(f_A_hat, f_B_hat) + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["force"] = f_pad + l["force"] = f_hat_pad + l["find_force"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Force MSE meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant(self._make_loss("mse"), f_A, f_A_hat, f_B, f_B_hat) + + def test_mae_grad_accum(self): + """Force MAE meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant(self._make_loss("mae"), f_A, f_A_hat, f_B, f_B_hat) + + def test_huber_grad_accum(self): + """Force Huber meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant( + self._make_loss("mse", use_huber=True), f_A, f_A_hat, f_B, f_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same force loss as no mask.""" + f = _rnd(NP, 3) + f_hat = _rnd(NP, 3) + loss_obj = self._make_loss() + + p_mask, l_mask = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + p_mask["force"] = f[None] + l_mask["force"] = f_hat[None] + l_mask["find_force"] = 1.0 + + p_nm, l_nm = _full_ener_dicts(1, NP, np.zeros((1, 1)), np.zeros((1, 1))) + p_nm["force"] = f[None] + l_nm["force"] = f_hat[None] + l_nm["find_force"] = 1.0 + + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnergyLossVirialGradAccum: + """Idiom 2 (extensive, k=9) for the virial (has_v) term. + + Covers: mse (norm_exp=1 and 2), mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", intensive=False, use_huber=False): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, v_A, v_A_hat, v_B, v_B_hat): + def make_A(): + p, l = _full_ener_dicts( + 1, NA, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NA)) + ) + p["virial"] = v_A[None] # [1, 9] + l["virial"] = v_A_hat[None] + l["find_virial"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, NB, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NB)) + ) + p["virial"] = v_B[None] + l["virial"] = v_B_hat[None] + l["find_virial"] = 1.0 + return p, l, NB + + def make_padded(): + v_pad = np.stack([v_A, v_B], axis=0) # [2, 9] + v_hat_pad = np.stack([v_A_hat, v_B_hat], axis=0) + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["virial"] = v_pad + l["virial"] = v_hat_pad + l["find_virial"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Virial MSE norm_exp=1 meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant( + self._make_loss("mse", intensive=False), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Virial MSE norm_exp=2 meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant( + self._make_loss("mse", intensive=True), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mae_grad_accum(self): + """Virial MAE meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant(self._make_loss("mae"), v_A, v_A_hat, v_B, v_B_hat) + + def test_huber_grad_accum(self): + """Virial Huber meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant( + self._make_loss("mse", use_huber=True), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same virial loss as no mask.""" + v = _rnd(9) + v_hat = _rnd(9) + loss_obj = self._make_loss() + + p_mask, l_mask = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + p_mask["virial"] = v[None] + l_mask["virial"] = v_hat[None] + l_mask["find_virial"] = 1.0 + + p_nm, l_nm = _full_ener_dicts(1, NP, np.zeros((1, 1)), np.zeros((1, 1))) + p_nm["virial"] = v[None] + l_nm["virial"] = v_hat[None] + l_nm["find_virial"] = 1.0 + + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnergyLossAtomEnerGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=1) for the atom_ener (has_ae) term. + + Covers: mse, mae; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse"): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + loss_func=loss_func, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat): + def make_A(): + p, l = _full_ener_dicts( + 1, NA, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NA)) + ) + p["atom_energy"] = ae_A[None] # [1, NA, 1] + l["atom_ener"] = ae_A_hat[None] + l["find_atom_ener"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, NB, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NB)) + ) + p["atom_energy"] = ae_B[None] + l["atom_ener"] = ae_B_hat[None] + l["find_atom_ener"] = 1.0 + return p, l, NB + + def make_padded(): + ae_pad = _padded_atom(ae_A, ae_B, 1) # [2, NP, 1] + ae_hat_pad = _padded_atom(ae_A_hat, ae_B_hat, 1) + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["atom_energy"] = ae_pad + l["atom_ener"] = ae_hat_pad + l["find_atom_ener"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Atom energy MSE meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + self._run_invariant(self._make_loss("mse"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_mae_grad_accum(self): + """Atom energy MAE meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom-energy loss as no mask.""" + ae = _rnd(NP, 1) + ae_hat = _rnd(NP, 1) + loss_obj = self._make_loss() + + p_mask, l_mask = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + p_mask["atom_energy"] = ae[None] + l_mask["atom_ener"] = ae_hat[None] + l_mask["find_atom_ener"] = 1.0 + + p_nm, l_nm = _full_ener_dicts(1, NP, np.zeros((1, 1)), np.zeros((1, 1))) + p_nm["atom_energy"] = ae[None] + l_nm["atom_ener"] = ae_hat[None] + l_nm["find_atom_ener"] = 1.0 + + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnergyLossAtomPrefGradAccum: + """Idiom 1 with pref weight (ncomp=3) for the atom_pref (has_pf) term. + + Covers: mse, mae; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse"): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + loss_func=loss_func, + ) + + def _loss_fn_pf_only(self, loss_obj, model_pred, label, natoms): + """Return the atom_pref contribution to the loss (subtract force loss).""" + # Both pref_f and pref_pf are 1.0 here, but we want ONLY the pf term. + # Use a loss with pref_pf=1 and pref_f=0 to isolate it. + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _make_pf_only_loss(self, loss_func="mse"): + """Loss with pref_f=0 and pref_pf=1 to isolate atom_pref term.""" + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + loss_func=loss_func, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B): + """Invariant for atom_pref using force diff weighted by pref.""" + + def make_A(): + p, l = _full_ener_dicts( + 1, NA, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NA)) + ) + p["force"] = f_A[None] + l["force"] = f_A_hat[None] + l["atom_pref"] = pf_A.reshape(1, NA * 3) + l["find_force"] = 1.0 + l["find_atom_pref"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, NB, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NB)) + ) + p["force"] = f_B[None] + l["force"] = f_B_hat[None] + l["atom_pref"] = pf_B.reshape(1, NB * 3) + l["find_force"] = 1.0 + l["find_atom_pref"] = 1.0 + return p, l, NB + + def make_padded(): + f_pad = _padded_force(f_A, f_B) + f_hat_pad = _padded_force(f_A_hat, f_B_hat) + pf_pad = _padded_atom_flat(pf_A, pf_B, 3) # [2, NP*3] + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["force"] = f_pad + l["force"] = f_hat_pad + l["atom_pref"] = pf_pad + l["find_force"] = 1.0 + l["find_atom_pref"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: float(loss_obj.call(1.0, na, mp, lb)[0]), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Atom-pref MSE meets the grad-accum invariant.""" + f_A, f_A_hat = _rnd(NA, 3), _rnd(NA, 3) + pf_A = np.abs(_rnd(NA, 3)) + 0.1 # positive pref + f_B, f_B_hat = _rnd(NB, 3), _rnd(NB, 3) + pf_B = np.abs(_rnd(NB, 3)) + 0.1 + self._run_invariant( + self._make_pf_only_loss("mse"), f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B + ) + + def test_mae_grad_accum(self): + """Atom-pref MAE meets the grad-accum invariant.""" + f_A, f_A_hat = _rnd(NA, 3), _rnd(NA, 3) + pf_A = np.abs(_rnd(NA, 3)) + 0.1 + f_B, f_B_hat = _rnd(NB, 3), _rnd(NB, 3) + pf_B = np.abs(_rnd(NB, 3)) + 0.1 + self._run_invariant( + self._make_pf_only_loss("mae"), f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom-pref loss as no mask.""" + f = _rnd(NP, 3) + f_hat = _rnd(NP, 3) + pf = np.abs(_rnd(NP, 3)) + 0.1 + loss_obj = self._make_pf_only_loss("mse") + + p_mask, l_mask = _full_ener_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), mask=np.ones((1, NP)) + ) + p_mask["force"] = f[None] + l_mask["force"] = f_hat[None] + l_mask["atom_pref"] = pf.reshape(1, NP * 3) + l_mask["find_force"] = 1.0 + l_mask["find_atom_pref"] = 1.0 + + p_nm, l_nm = _full_ener_dicts(1, NP, np.zeros((1, 1)), np.zeros((1, 1))) + p_nm["force"] = f[None] + l_nm["force"] = f_hat[None] + l_nm["atom_pref"] = pf.reshape(1, NP * 3) + l_nm["find_force"] = 1.0 + l_nm["find_atom_pref"] = 1.0 + + loss_m = float(loss_obj.call(1.0, NP, p_mask, l_mask)[0]) + loss_nm = float(loss_obj.call(1.0, NP, p_nm, l_nm)[0]) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 48f954ca56..74df40f1fd 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -18,6 +18,9 @@ from deepmd.pt.loss.dos import ( DOSLoss, ) +from deepmd.pt.loss.ener import ( + EnergyStdLoss, +) from deepmd.pt.loss.loss import ( TaskLoss, ) @@ -582,3 +585,586 @@ def test_all_real_atoms_gives_all_ones(self) -> None: assert torch.all(result["mask"] == 1.0), ( "all-real atoms must give all-ones mask" ) + + +# --------------------------------------------------------------------------- +# Task 3: EnergyStdLoss -- energy, force, virial, atom_ener, atom_pref +# --------------------------------------------------------------------------- + +# Re-use same constants from Task 2 harness (NA=3, NB=5, NP=5). + +_MASK_PAD_PT = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", +) # [2, NP] + + +def _t(*shape, val=None): + """Random float64 CPU tensor.""" + if val is not None: + return torch.full(shape, val, dtype=torch.float64, device="cpu") + return torch.tensor(RNG.standard_normal(shape), dtype=torch.float64, device="cpu") + + +def _padded_force_t(f_A, f_B): + """Stack force tensors for 2-frame padded batch.""" + pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + pad[:NA] = f_A + return torch.stack([pad, f_B], dim=0) # [2, NP, 3] + + +def _padded_atom_t(a_A, a_B, ncomp): + """Pad arr_A from [NA, ncomp] to [NP, ncomp], stack with arr_B.""" + pad = torch.zeros(NP, ncomp, dtype=torch.float64, device="cpu") + pad[:NA] = a_A + return torch.stack([pad, a_B], dim=0) # [2, NP, ncomp] + + +class _EnerLossMockModel: + """Callable returning a fixed model_pred dict (mask pre-populated).""" + + def __init__(self, pred: dict): + self._pred = pred + + def __call__(self, **kwargs): + return dict(self._pred) + + +def _ener_loss_fn(loss_obj, model_pred, label, natoms): + """Call EnergyStdLoss.forward via mock model; return scalar loss tensor.""" + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_EnerLossMockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + +class TestPTEnergyLossEnerGradAccum: + """Idiom 2 (extensive) for the energy term in EnergyStdLoss. + + Covers: mse (norm_exp=1 and 2), mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", intensive=False, use_huber=False): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + + def _run_invariant(self, loss_obj, e_A, e_A_hat, e_B, e_B_hat): + def make_A(): + mp = { + "energy": e_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"energy": e_A_hat, "find_energy": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "energy": e_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"energy": e_B_hat, "find_energy": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "energy": torch.cat([e_A, e_B], dim=0), + "mask": _MASK_PAD_PT, + } + lb = { + "energy": torch.cat([e_A_hat, e_B_hat], dim=0), + "find_energy": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Energy MSE norm_exp=1 meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant( + self._make_loss("mse", intensive=False), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Energy MSE norm_exp=2 meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant( + self._make_loss("mse", intensive=True), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_mae_grad_accum(self): + """Energy MAE meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant(self._make_loss("mae"), e_A, e_A_hat, e_B, e_B_hat) + + def test_huber_grad_accum(self): + """Energy Huber meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant( + self._make_loss("mse", use_huber=True), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same energy loss as no mask.""" + e = _t(1, 1) + e_hat = _t(1, 1) + loss_obj = self._make_loss() + + mp_mask = { + "energy": e, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"energy": e} + lb = {"energy": e_hat, "find_energy": 1.0} + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnergyLossForceGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=3) for the force term. + + Covers: mse, mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", use_huber=False, f_use_norm=False): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func=loss_func, + use_huber=use_huber, + f_use_norm=f_use_norm, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): + def make_A(): + mp = { + "force": f_A.unsqueeze(0), # [1, NA, 3] + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"force": f_A_hat.unsqueeze(0), "find_force": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "force": f_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"force": f_B_hat.unsqueeze(0), "find_force": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "force": _padded_force_t(f_A, f_B), # [2, NP, 3] + "mask": _MASK_PAD_PT, + } + lb = { + "force": _padded_force_t(f_A_hat, f_B_hat), + "find_force": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Force MSE meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant(self._make_loss("mse"), f_A, f_A_hat, f_B, f_B_hat) + + def test_mae_grad_accum(self): + """Force MAE meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant(self._make_loss("mae"), f_A, f_A_hat, f_B, f_B_hat) + + def test_huber_grad_accum(self): + """Force Huber meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant( + self._make_loss("mse", use_huber=True), f_A, f_A_hat, f_B, f_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same force loss as no mask.""" + f = _t(NP, 3) + f_hat = _t(NP, 3) + loss_obj = self._make_loss() + + mp_mask = { + "force": f.unsqueeze(0), + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"force": f.unsqueeze(0)} + lb = {"force": f_hat.unsqueeze(0), "find_force": 1.0} + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnergyLossVirialGradAccum: + """Idiom 2 (extensive, k=9) for the virial term. + + Covers: mse (norm_exp=1 and 2), mae, huber; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse", intensive=False, use_huber=False): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + loss_func=loss_func, + intensive_ener_virial=intensive, + use_huber=use_huber, + ) + + def _run_invariant(self, loss_obj, v_A, v_A_hat, v_B, v_B_hat): + def make_A(): + mp = { + "virial": v_A.unsqueeze(0), # [1, 9] + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"virial": v_A_hat.unsqueeze(0), "find_virial": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "virial": v_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"virial": v_B_hat.unsqueeze(0), "find_virial": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "virial": torch.stack([v_A, v_B], dim=0), # [2, 9] + "mask": _MASK_PAD_PT, + } + lb = { + "virial": torch.stack([v_A_hat, v_B_hat], dim=0), + "find_virial": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Virial MSE norm_exp=1 meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant( + self._make_loss("mse", intensive=False), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Virial MSE norm_exp=2 meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant( + self._make_loss("mse", intensive=True), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mae_grad_accum(self): + """Virial MAE meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant(self._make_loss("mae"), v_A, v_A_hat, v_B, v_B_hat) + + def test_huber_grad_accum(self): + """Virial Huber meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant( + self._make_loss("mse", use_huber=True), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same virial loss as no mask.""" + v = _t(9) + v_hat = _t(9) + loss_obj = self._make_loss() + + mp_mask = { + "virial": v.unsqueeze(0), + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"virial": v.unsqueeze(0)} + lb = {"virial": v_hat.unsqueeze(0), "find_virial": 1.0} + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnergyLossAtomEnerGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=1) for the atom_ener term. + + Covers: mse, mae; plus non-mixed no-op. + """ + + def _make_loss(self, loss_func="mse"): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + loss_func=loss_func, + ) + + def _run_invariant(self, loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat): + def make_A(): + mp = { + "atom_energy": ae_A.unsqueeze(0), # [1, NA, 1] + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"atom_ener": ae_A_hat.unsqueeze(0), "find_atom_ener": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "atom_energy": ae_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"atom_ener": ae_B_hat.unsqueeze(0), "find_atom_ener": 1.0} + return mp, lb, NB + + def make_padded(): + ae_pad = _padded_atom_t(ae_A, ae_B, 1) # [2, NP, 1] + ae_hat_pad = _padded_atom_t(ae_A_hat, ae_B_hat, 1) + mp = {"atom_energy": ae_pad, "mask": _MASK_PAD_PT} + lb = {"atom_ener": ae_hat_pad, "find_atom_ener": 1.0} + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Atom-energy MSE meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + self._run_invariant(self._make_loss("mse"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_mae_grad_accum(self): + """Atom-energy MAE meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom-energy loss as no mask.""" + ae = _t(NP, 1) + ae_hat = _t(NP, 1) + loss_obj = self._make_loss() + + mp_mask = { + "atom_energy": ae.unsqueeze(0), + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"atom_energy": ae.unsqueeze(0)} + lb = {"atom_ener": ae_hat.unsqueeze(0), "find_atom_ener": 1.0} + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnergyLossAtomPrefGradAccum: + """Idiom 1 with pref weight (ncomp=3) for the atom_pref term. + + Covers: mse, mae; plus non-mixed no-op. + """ + + def _make_pf_only_loss(self, loss_func="mse"): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=1.0, + limit_pref_f=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_pf=1.0, + limit_pref_pf=1.0, + loss_func=loss_func, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B): + def make_A(): + mp = { + "force": f_A.unsqueeze(0), + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = { + "force": f_A_hat.unsqueeze(0), + "atom_pref": pf_A.unsqueeze(0), + "find_force": 1.0, + "find_atom_pref": 1.0, + } + return mp, lb, NA + + def make_B(): + mp = { + "force": f_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = { + "force": f_B_hat.unsqueeze(0), + "atom_pref": pf_B.unsqueeze(0), + "find_force": 1.0, + "find_atom_pref": 1.0, + } + return mp, lb, NB + + def make_padded(): + f_pad = _padded_force_t(f_A, f_B) + f_hat_pad = _padded_force_t(f_A_hat, f_B_hat) + pf_pad = _padded_atom_t(pf_A, pf_B, 3) # [2, NP, 3] + mp = {"force": f_pad, "mask": _MASK_PAD_PT} + lb = { + "force": f_hat_pad, + "atom_pref": pf_pad, + "find_force": 1.0, + "find_atom_pref": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Atom-pref MSE meets the grad-accum invariant.""" + f_A, f_A_hat = _t(NA, 3), _t(NA, 3) + pf_A = torch.abs(_t(NA, 3)) + 0.1 + f_B, f_B_hat = _t(NB, 3), _t(NB, 3) + pf_B = torch.abs(_t(NB, 3)) + 0.1 + self._run_invariant( + self._make_pf_only_loss("mse"), f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B + ) + + def test_mae_grad_accum(self): + """Atom-pref MAE meets the grad-accum invariant.""" + f_A, f_A_hat = _t(NA, 3), _t(NA, 3) + pf_A = torch.abs(_t(NA, 3)) + 0.1 + f_B, f_B_hat = _t(NB, 3), _t(NB, 3) + pf_B = torch.abs(_t(NB, 3)) + 0.1 + self._run_invariant( + self._make_pf_only_loss("mae"), f_A, f_A_hat, pf_A, f_B, f_B_hat, pf_B + ) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom-pref loss as no mask.""" + f = _t(NP, 3) + f_hat = _t(NP, 3) + pf = torch.abs(_t(NP, 3)) + 0.1 + loss_obj = self._make_pf_only_loss("mse") + + mp_mask = { + "force": f.unsqueeze(0), + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"force": f.unsqueeze(0)} + lb = { + "force": f_hat.unsqueeze(0), + "atom_pref": pf.unsqueeze(0), + "find_force": 1.0, + "find_atom_pref": 1.0, + } + + loss_m = _ener_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _ener_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) From 82c71fd8ebbb9b6988efdbc293b1bdf5ca31e936 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:03:51 +0800 Subject: [PATCH 04/16] test(loss): cover atom_ener Huber invariant and mask the rmse_f display metric Add test_huber_grad_accum to TestDPModelEnergyLossAtomEnerGradAccum and TestPTEnergyLossAtomEnerGradAccum verifying the masked Huber atom_ener path satisfies the [3+5]==separate grad-accum invariant. Fix rmse_f in the masked force MSE branch to use the masked per-frame l2 (l2_force_masked / l2_f_masked) rather than the ghost-diluted unmasked l2_force_loss; hoist the masked MSE computation before the use_huber split so it is always in scope for the display metric in both huber and non-huber paths. --- deepmd/dpmodel/loss/ener.py | 13 +++++------ deepmd/pt/loss/ener.py | 15 ++++++++----- .../tests/common/dpmodel/test_loss_padding.py | 22 +++++++++++++++++++ source/tests/pt/test_loss_padding.py | 20 +++++++++++++++++ 4 files changed, 58 insertions(+), 12 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index ff155b442e..8cbfa5ccff 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -362,13 +362,12 @@ def call( # Idiom 1 (per-atom masked mean, ncomp=3). diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) # [nf, nloc, 3] maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1] + # Masked MSE computed for rmse_f display regardless of use_huber. + sq_f = xp.square(diff_f_3d) * maskf_col # [nf, nloc, 3] + _pfs = xp.sum(xp.reshape(sq_f, (_nf, -1)), axis=-1) # [nf] + _pfd = xp.sum(maskf, axis=-1) * 3 # [nf] + l2_force_masked = xp.mean(_pfs / _pfd) if not self.use_huber: - sq_f = xp.square(diff_f_3d) * maskf_col # [nf, nloc, 3] - per_frame_sum = xp.sum( - xp.reshape(sq_f, (_nf, -1)), axis=-1 - ) # [nf] - per_frame_dof = xp.sum(maskf, axis=-1) * 3 # [nf] - l2_force_masked = xp.mean(per_frame_sum / per_frame_dof) loss += pref_f * l2_force_masked else: if not self.f_use_norm: @@ -408,7 +407,7 @@ def call( l_huber_masked = xp.mean(per_frame_sum / per_frame_dof) loss += pref_f * l_huber_masked more_loss["rmse_f"] = self.display_if_exist( - xp.sqrt(l2_force_loss), find_force + xp.sqrt(l2_force_masked), find_force ) else: if not self.use_huber: diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index a7b7683363..4ae64ea8ae 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -382,11 +382,12 @@ def forward( # Idiom 1 (per-atom masked mean, ncomp=3). diff_f_3d = diff_f.reshape(_nf, _nloc, 3) maskf_col = maskf.reshape(_nf, _nloc, 1) + # Masked MSE computed for rmse_f display regardless of use_huber. + sq_f = torch.square(diff_f_3d) * maskf_col + _pfs = sq_f.reshape(_nf, -1).sum(dim=-1) + _pfd = maskf.sum(dim=-1) * 3 + l2_f_masked = torch.mean(_pfs / _pfd) if not self.use_huber: - sq_f = torch.square(diff_f_3d) * maskf_col - per_frame_sum = sq_f.reshape(_nf, -1).sum(dim=-1) - per_frame_dof = maskf.sum(dim=-1) * 3 - l2_f_masked = torch.mean(per_frame_sum / per_frame_dof) loss += (pref_f * l2_f_masked).to(GLOBAL_PT_FLOAT_PRECISION) else: if not self.f_use_norm: @@ -445,7 +446,11 @@ def forward( delta=self._huber_delta_force, ) loss += pref_f * l_huber_loss - rmse_f = l2_force_loss.sqrt() + rmse_f = ( + l2_f_masked.sqrt() + if maskf is not None + else l2_force_loss.sqrt() + ) more_loss["rmse_f"] = self.display_if_exist( rmse_f.detach(), find_force ) diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 4b009a78be..14cd5ca842 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -934,6 +934,28 @@ def test_mae_grad_accum(self): ae_B_hat = _rnd(NB, 1) self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + def test_huber_grad_accum(self): + """Atom energy Huber (use_huber=True) meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + loss_obj = EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + use_huber=True, + ) + self._run_invariant(loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat) + def test_no_op_for_non_mixed(self): """All-ones mask gives same atom-energy loss as no mask.""" ae = _rnd(NP, 1) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 74df40f1fd..e7d5d26e96 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -1037,6 +1037,26 @@ def test_mae_grad_accum(self): ae_B_hat = _t(NB, 1) self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + def test_huber_grad_accum(self): + """Atom-energy Huber (use_huber=True) meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + loss_obj = EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + use_huber=True, + ) + self._run_invariant(loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat) + def test_no_op_for_non_mixed(self): """All-ones mask gives same atom-energy loss as no mask.""" ae = _t(NP, 1) From 9e8f4419c38fd72d6ea65655f29a537a31f86649 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:18:45 +0800 Subject: [PATCH 05/16] fix(loss): normalize extensive property loss by per-frame real atom count MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit In mixed_type batches, ghost atoms (atype=-1) pad short frames to a fixed nloc width. PropertyLoss previously divided pred/label by the scalar natoms (the padded count), causing frames with fewer real atoms to be incorrectly normalized. When model_dict["mask"] is present, use the per-frame real atom count (sum of mask along the atom axis, broadcast over task_dim) instead. The else branch retains the original /natoms behavior so mask-free callers are unchanged. Tests: RED→GREEN for mse/mae/smooth_mae grad-accum invariant in both dpmodel and pt; no-op and intensive guards added. --- deepmd/dpmodel/loss/property.py | 13 +- deepmd/pt/loss/property.py | 14 +- .../tests/common/dpmodel/test_loss_padding.py | 128 ++++++++++++++ source/tests/pt/test_loss_padding.py | 164 ++++++++++++++++++ 4 files changed, 315 insertions(+), 4 deletions(-) diff --git a/deepmd/dpmodel/loss/property.py b/deepmd/dpmodel/loss/property.py index 7d658ff925..627c299abf 100644 --- a/deepmd/dpmodel/loss/property.py +++ b/deepmd/dpmodel/loss/property.py @@ -85,8 +85,17 @@ def call( # Normalize by natoms for extensive properties (without mutating input) if not self.intensive: - pred = pred / natoms - label = label / natoms + if "mask" in model_dict: + # Per-frame real atom count: shape [nf] → broadcast over [nf, task_dim]. + real_natoms = xp.reshape( + xp.astype(xp.sum(model_dict["mask"], -1), pred.dtype), + (-1,) + (1,) * (pred.ndim - 1), + ) + pred = pred / real_natoms + label = label / real_natoms + else: + pred = pred / natoms + label = label / natoms # Get out_std and out_bias if self.out_std is not None: diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index abf5bcfacb..df865b10f1 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -106,8 +106,18 @@ def forward( assert model_pred[var_name].shape == (nbz, self.task_dim) assert label[var_name].shape == (nbz, self.task_dim) if not self.intensive: - model_pred[var_name] = model_pred[var_name] / natoms - label[var_name] = label[var_name] / natoms + if "mask" in model_pred: + # Per-frame real atom count: shape [nf] → broadcast over [nf, task_dim]. + real_natoms = ( + torch.sum(model_pred["mask"], dim=-1) + .to(dtype=model_pred[var_name].dtype) + .reshape(-1, 1) + ) + model_pred[var_name] = model_pred[var_name] / real_natoms + label[var_name] = label[var_name] / real_natoms + else: + model_pred[var_name] = model_pred[var_name] / natoms + label[var_name] = label[var_name] / natoms if self.out_std is None: out_std = model.atomic_model.out_std[0][0] diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 14cd5ca842..2396a181e0 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -21,6 +21,9 @@ from deepmd.dpmodel.loss.ener import ( EnergyLoss, ) +from deepmd.dpmodel.loss.property import ( + PropertyLoss, +) from deepmd.dpmodel.loss.tensor import ( TensorLoss, ) @@ -1121,3 +1124,128 @@ def test_no_op_for_non_mixed(self): assert np.isclose(loss_m, loss_nm), ( f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" ) + + +# --------------------------------------------------------------------------- +# Task 4: PropertyLoss -- extensive (not intensive) property +# --------------------------------------------------------------------------- + +PROP_TASK_DIM = 2 +PROP_VAR = "test_prop" + + +class TestPropertyLossExtensiveGradAccum: + """Idiom 2 (per-frame real-natoms normalization) for extensive PropertyLoss. + + The _loss_fn wrapper divides by nf so that the per-frame average matches + the separate-frame reference (property loss uses sum, not mean, over frames). + """ + + def _make_loss(self, loss_func="mse"): + return PropertyLoss( + task_dim=PROP_TASK_DIM, + var_name=PROP_VAR, + loss_func=loss_func, + intensive=False, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + """Return per-frame-averaged loss (raw loss / nf).""" + nf = model_pred[PROP_VAR].shape[0] + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) / nf + + def _run_invariant(self, loss_obj, p_A, l_A, p_B, l_B): + def make_A(): + return ( + {PROP_VAR: p_A, "mask": np.ones((1, NA), dtype=np.float64)}, + {PROP_VAR: l_A}, + NA, + ) + + def make_B(): + return ( + {PROP_VAR: p_B, "mask": np.ones((1, NB), dtype=np.float64)}, + {PROP_VAR: l_B}, + NB, + ) + + def make_padded(): + return ( + { + PROP_VAR: np.concatenate([p_A, p_B], axis=0), # [2, task_dim] + "mask": _MASK_PAD, + }, + {PROP_VAR: np.concatenate([l_A, l_B], axis=0)}, + NP, + ) + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """MSE extensive property meets the grad-accum invariant.""" + p_A = _rnd(1, PROP_TASK_DIM) + l_A = _rnd(1, PROP_TASK_DIM) + p_B = _rnd(1, PROP_TASK_DIM) + l_B = _rnd(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("mse"), p_A, l_A, p_B, l_B) + + def test_mae_grad_accum(self): + """MAE extensive property meets the grad-accum invariant.""" + p_A = _rnd(1, PROP_TASK_DIM) + l_A = _rnd(1, PROP_TASK_DIM) + p_B = _rnd(1, PROP_TASK_DIM) + l_B = _rnd(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("mae"), p_A, l_A, p_B, l_B) + + def test_smooth_mae_grad_accum(self): + """smooth_mae extensive property meets the grad-accum invariant.""" + p_A = _rnd(1, PROP_TASK_DIM) + l_A = _rnd(1, PROP_TASK_DIM) + p_B = _rnd(1, PROP_TASK_DIM) + l_B = _rnd(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("smooth_mae"), p_A, l_A, p_B, l_B) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same extensive-property loss as no mask.""" + p = _rnd(2, PROP_TASK_DIM) + l = _rnd(2, PROP_TASK_DIM) + loss_obj = self._make_loss("mse") + # PropertyLoss.call rebinds local `label` (no in-place mutation) so l is safe to reuse. + with_mask = {PROP_VAR: p, "mask": np.ones((2, NB), dtype=np.float64)} + without_mask = {PROP_VAR: p} + loss_m, _ = loss_obj.call(1.0, NB, with_mask, {PROP_VAR: l}) + loss_nm, _ = loss_obj.call(1.0, NB, without_mask, {PROP_VAR: l}) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" + ) + + +class TestPropertyLossIntensiveUnaffectedByMask: + """Intensive property loss must be unchanged whether or not mask is present.""" + + def _make_loss(self): + return PropertyLoss( + task_dim=PROP_TASK_DIM, + var_name=PROP_VAR, + loss_func="mse", + intensive=True, + ) + + def test_intensive_ignores_mask(self): + """Intensive property: padded-batch loss == unmasked-batch loss.""" + p = _rnd(2, PROP_TASK_DIM) + l = _rnd(2, PROP_TASK_DIM) + loss_obj = self._make_loss() + with_mask = {PROP_VAR: p, "mask": _MASK_PAD} + without_mask = {PROP_VAR: p} + loss_m, _ = loss_obj.call(1.0, NP, with_mask, {PROP_VAR: l}) + loss_nm, _ = loss_obj.call(1.0, NP, without_mask, {PROP_VAR: l}) + assert np.isclose(float(loss_m), float(loss_nm)), ( + f"intensive property must ignore mask: {float(loss_m)} vs {float(loss_nm)}" + ) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index e7d5d26e96..9a4a8d2f76 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -24,6 +24,9 @@ from deepmd.pt.loss.loss import ( TaskLoss, ) +from deepmd.pt.loss.property import ( + PropertyLoss, +) from deepmd.pt.loss.tensor import ( TensorLoss, ) @@ -1188,3 +1191,164 @@ def test_no_op_for_non_mixed(self): assert torch.isclose(loss_m, loss_nm), ( f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" ) + + +# --------------------------------------------------------------------------- +# Task 4: PropertyLoss -- extensive (not intensive) property +# --------------------------------------------------------------------------- + +PROP_TASK_DIM = 2 +PROP_VAR = "test_prop" + + +class TestPTPropertyLossExtensiveGradAccum: + """Idiom 2 (per-frame real-natoms normalization) for extensive pt PropertyLoss. + + _loss_fn wraps the raw loss with /nf so the per-frame average matches the + separate-frame reference (pt PropertyLoss uses reduction='sum' over frames). + The mask is pre-populated in model_pred so _inject_atom_mask leaves it alone. + """ + + def _make_loss(self, loss_func="mse"): + return PropertyLoss( + task_dim=PROP_TASK_DIM, + var_name=PROP_VAR, + loss_func=loss_func, + intensive=False, + # Provide explicit out_std/out_bias to avoid accessing model.atomic_model. + out_std=[1.0] * PROP_TASK_DIM, + out_bias=[0.0] * PROP_TASK_DIM, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + """Return per-frame-averaged loss (raw loss / nf).""" + nf = model_pred[PROP_VAR].shape[0] + _, loss, _ = loss_obj.forward( + input_dict={}, # no atype; mask already in model_pred + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss / nf + + def _run_invariant(self, loss_obj, p_A, l_A, p_B, l_B): + def make_A(): + return ( + { + PROP_VAR: p_A, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + }, + {PROP_VAR: l_A.clone()}, + NA, + ) + + def make_B(): + return ( + { + PROP_VAR: p_B, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + }, + {PROP_VAR: l_B.clone()}, + NB, + ) + + def make_padded(): + return ( + { + PROP_VAR: torch.cat([p_A, p_B], dim=0), # [2, task_dim] + "mask": _MASK_PAD_PT, + }, + {PROP_VAR: torch.cat([l_A, l_B], dim=0)}, + NP, + ) + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """MSE extensive property meets the grad-accum invariant.""" + p_A = _rnd_t(1, PROP_TASK_DIM) + l_A = _rnd_t(1, PROP_TASK_DIM) + p_B = _rnd_t(1, PROP_TASK_DIM) + l_B = _rnd_t(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("mse"), p_A, l_A, p_B, l_B) + + def test_mae_grad_accum(self): + """MAE extensive property meets the grad-accum invariant.""" + p_A = _rnd_t(1, PROP_TASK_DIM) + l_A = _rnd_t(1, PROP_TASK_DIM) + p_B = _rnd_t(1, PROP_TASK_DIM) + l_B = _rnd_t(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("mae"), p_A, l_A, p_B, l_B) + + def test_smooth_mae_grad_accum(self): + """smooth_mae extensive property meets the grad-accum invariant.""" + p_A = _rnd_t(1, PROP_TASK_DIM) + l_A = _rnd_t(1, PROP_TASK_DIM) + p_B = _rnd_t(1, PROP_TASK_DIM) + l_B = _rnd_t(1, PROP_TASK_DIM) + self._run_invariant(self._make_loss("smooth_mae"), p_A, l_A, p_B, l_B) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same extensive-property loss as no mask.""" + p = _rnd_t(2, PROP_TASK_DIM) + l = _rnd_t(2, PROP_TASK_DIM) + loss_obj = self._make_loss("mse") + mp_mask = { + PROP_VAR: p, + "mask": torch.ones(2, NB, dtype=torch.float64, device="cpu"), + } + mp_nm = {PROP_VAR: p} + # pt forward mutates label[var_name]; use separate dicts for each call. + lb_m = {PROP_VAR: l.clone()} + lb_nm = {PROP_VAR: l.clone()} + loss_m = self._loss_fn(loss_obj, mp_mask, lb_m, NB) + loss_nm = self._loss_fn(loss_obj, mp_nm, lb_nm, NB) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTPropertyLossIntensiveUnaffectedByMask: + """Intensive property loss must be unchanged whether or not mask is present.""" + + def _make_loss(self): + return PropertyLoss( + task_dim=PROP_TASK_DIM, + var_name=PROP_VAR, + loss_func="mse", + intensive=True, + out_std=[1.0] * PROP_TASK_DIM, + out_bias=[0.0] * PROP_TASK_DIM, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + _, loss, _ = loss_obj.forward( + input_dict={}, + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + def test_intensive_ignores_mask(self): + """Intensive property: masked batch == unmasked batch.""" + p = _rnd_t(2, PROP_TASK_DIM) + l = _rnd_t(2, PROP_TASK_DIM) + loss_obj = self._make_loss() + mp_mask = {PROP_VAR: p, "mask": _MASK_PAD_PT} + mp_nm = {PROP_VAR: p} + # Use separate label dicts since pt forward mutates label[var_name]. + lb_m = {PROP_VAR: l.clone()} + lb_nm = {PROP_VAR: l.clone()} + loss_m = self._loss_fn(loss_obj, mp_mask, lb_m, NP) + loss_nm = self._loss_fn(loss_obj, mp_nm, lb_nm, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"intensive property must ignore mask: {loss_m.item()} vs {loss_nm.item()}" + ) From 968ae3cd2e94bf42f17569c0614ad51a6a061811 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:42:15 +0800 Subject: [PATCH 06/16] fix(loss): per-frame normalize the spin energy loss to exclude mixed_type padding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply the grad-accumulation-invariant per-frame normalization to EnergySpinLoss in both the dpmodel (deepmd/dpmodel/loss/ener_spin.py) and pt (deepmd/pt/loss/ener_spin.py) backends: - Energy (has_e, MSE): idiom 2 — per-frame 1/real_natoms^norm_exp normalization replaces the flat atom_norm**norm_exp * mean() formulation. - Force real (has_fr, MSE): idiom 1 — per-atom masked mean (ncomp=3) replaces the unmasked global mean over all nloc atoms. - Virial (has_v, MSE): idiom 2 — per-frame 1/real_natoms^norm_exp normalization (k=9) replaces atom_norm**norm_exp * mean(). The force_mag / mask_mag (has_fm) path is left completely untouched; mask_mag is the spin virtual-atom mask and is a distinct concept from the padding mask. Each masked branch is guarded by `if "mask" in model_dict/model_pred:` so mask-less callers (non-mixed batches without a padding mask) use the original code path unchanged — the fix is a bit-identical no-op for non-mixed batches. Tests added to source/tests/common/dpmodel/test_loss_padding.py and source/tests/pt/test_loss_padding.py: invariant tests for energy (norm_exp 1 and 2), force_real, and virial; no-op tests for each term; and a guard that confirms force_mag loss is numerically unchanged when a padding mask is present. --- deepmd/dpmodel/loss/ener_spin.py | 104 ++++- deepmd/pt/loss/ener_spin.py | 148 ++++-- .../tests/common/dpmodel/test_loss_padding.py | 439 ++++++++++++++++++ source/tests/pt/test_loss_padding.py | 362 +++++++++++++++ 4 files changed, 993 insertions(+), 60 deletions(-) diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index 6262015c84..4944916bde 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -132,6 +132,20 @@ def call( # - norm_exp=1 (intensive_ener_virial=False, legacy): loss uses 1/N scaling, which varies with system size norm_exp = 2 if self.intensive_ener_virial else 1 + # Per-frame mask: recover real-atom count per frame when mask is provided. + # maskf[nf, nloc] = 1.0 for real atoms, 0.0 for ghost padding atoms. + if "mask" in model_dict: + maskf = xp.astype(model_dict["mask"], energy.dtype) # [nf, nloc] + real_natoms = xp.sum(maskf, axis=-1) # [nf] + inv = xp.reshape(1.0 / real_natoms, (-1,)) # [nf] + _nf = maskf.shape[0] + _nloc = maskf.shape[1] + else: + maskf = None + inv = None + _nf = None + _nloc = None + if self.has_e: energy_pred = model_dict["energy"] energy_label = label_dict["energy"] @@ -143,11 +157,20 @@ def call( atom_ener_coeff = xp.reshape(atom_ener_coeff, atom_ener_pred.shape) energy_pred = xp.sum(atom_ener_coeff * atom_ener_pred, axis=1) if self.loss_func == "mse": - l2_ener_loss = xp.mean(xp.square(energy_pred - energy_label)) - loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) - more_loss["rmse_e"] = self.display_if_exist( - xp.sqrt(l2_ener_loss) * atom_norm, find_energy - ) + se_e = xp.square(energy_pred - energy_label) # [nf, k] + if maskf is not None: + # Idiom 2 (extensive): per-frame normalization by real-atom count. + per_frame_e = xp.mean(xp.reshape(se_e, (_nf, -1)), axis=-1) # [nf] + loss += pref_e * xp.mean(per_frame_e * inv**norm_exp) + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(xp.mean(per_frame_e * inv**2)), find_energy + ) + else: + l2_ener_loss = xp.mean(se_e) + loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) + more_loss["rmse_e"] = self.display_if_exist( + xp.sqrt(l2_ener_loss) * atom_norm, find_energy + ) elif self.loss_func == "mae": abs_diff_e = xp.abs(energy_pred - energy_label) l1_ener_loss = xp.sum(abs_diff_e) @@ -167,15 +190,36 @@ def call( force_pred = model_dict["force"] force_label = label_dict["force"] if self.loss_func == "mse": - diff_fr = force_label - force_pred - l2_force_real_loss = xp.mean(xp.square(diff_fr)) - loss += pref_fr * l2_force_real_loss - more_loss["rmse_fr"] = self.display_if_exist( - xp.sqrt(l2_force_real_loss), find_force - ) - if mae: - mae_fr = xp.mean(xp.abs(force_label - force_pred)) - more_loss["mae_fr"] = self.display_if_exist(mae_fr, find_force) + diff_fr = force_label - force_pred # [nf, nloc, 3] + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3). + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1] + sq_fr = xp.square(diff_fr) * maskf_col # [nf, nloc, 3] + per_frame_fr_sum = xp.sum( + xp.reshape(sq_fr, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_fr_dof = xp.sum(maskf, axis=-1) * 3 # [nf] + l2_force_real_loss = xp.mean(per_frame_fr_sum / per_frame_fr_dof) + loss += pref_fr * l2_force_real_loss + more_loss["rmse_fr"] = self.display_if_exist( + xp.sqrt(l2_force_real_loss), find_force + ) + if mae: + abs_fr = xp.abs(force_label - force_pred) * maskf_col + per_frame_mae_sum = xp.sum( + xp.reshape(abs_fr, (_nf, -1)), axis=-1 + ) + mae_fr = xp.mean(per_frame_mae_sum / per_frame_fr_dof) + more_loss["mae_fr"] = self.display_if_exist(mae_fr, find_force) + else: + l2_force_real_loss = xp.mean(xp.square(diff_fr)) + loss += pref_fr * l2_force_real_loss + more_loss["rmse_fr"] = self.display_if_exist( + xp.sqrt(l2_force_real_loss), find_force + ) + if mae: + mae_fr = xp.mean(xp.abs(force_label - force_pred)) + more_loss["mae_fr"] = self.display_if_exist(mae_fr, find_force) elif self.loss_func == "mae": abs_diff_fr = xp.abs(force_label - force_pred) per_atom_fr = xp.sum(abs_diff_fr, axis=-1) # [nf, na] @@ -249,16 +293,30 @@ def call( pref_v = pref_v * find_virial virial_pred = xp.reshape(model_dict["virial"], (-1, 9)) virial_label = label_dict["virial"] - diff_v = virial_label - virial_pred + diff_v = virial_label - virial_pred # [nf, 9] if self.loss_func == "mse": - l2_virial_loss = xp.mean(xp.square(diff_v)) - loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) - more_loss["rmse_v"] = self.display_if_exist( - xp.sqrt(l2_virial_loss) * atom_norm, find_virial - ) - if mae: - mae_v = xp.mean(xp.abs(diff_v)) * atom_norm - more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) + if maskf is not None: + # Idiom 2 (extensive, k=9): per-frame normalization by real-atom count. + se_v = xp.square(diff_v) # [nf, 9] + per_frame_v = xp.mean(se_v, axis=-1) # [nf] + loss += pref_v * xp.mean(per_frame_v * inv**norm_exp) + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(xp.mean(per_frame_v * inv**2)), find_virial + ) + if mae: + abs_v = xp.abs(diff_v) # [nf, 9] + per_frame_mae_v = xp.mean(abs_v, axis=-1) # [nf] + mae_v = xp.mean(per_frame_mae_v * inv) + more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) + else: + l2_virial_loss = xp.mean(xp.square(diff_v)) + loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + more_loss["rmse_v"] = self.display_if_exist( + xp.sqrt(l2_virial_loss) * atom_norm, find_virial + ) + if mae: + mae_v = xp.mean(xp.abs(diff_v)) * atom_norm + more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) elif self.loss_func == "mae": l1_virial_loss = xp.mean(xp.abs(diff_v)) loss += atom_norm * (pref_v * l1_virial_loss) diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 1e26ead84e..3508031f5a 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -159,6 +159,21 @@ def forward( # - norm_exp=2 (intensive_ener_virial=True): loss uses 1/N² scaling, making it independent of system size # - norm_exp=1 (intensive_ener_virial=False, legacy): loss uses 1/N scaling, which varies with system size norm_exp = 2 if self.intensive_ener_virial else 1 + + # Per-frame mask: recover real-atom count per frame when mask is provided. + # maskf[nf, nloc] = 1.0 for real atoms, 0.0 for ghost padding atoms. + if "mask" in model_pred: + maskf = model_pred["mask"] # [nf, nloc], float + real_natoms_f = torch.sum(maskf, dim=-1) # [nf] + inv = (1.0 / real_natoms_f).reshape(-1) # [nf] + _nf = maskf.shape[0] + _nloc = maskf.shape[1] + else: + maskf = None + inv = None + _nf = None + _nloc = None + if self.has_e and "energy" in model_pred and "energy" in label: energy_pred = model_pred["energy"] energy_label = label["energy"] @@ -178,16 +193,30 @@ def forward( find_energy = label.get("find_energy", 0.0) pref_e = pref_e * find_energy if self.loss_func == "mse": - l2_ener_loss = torch.mean(torch.square(energy_pred - energy_label)) - if not self.inference: - more_loss["l2_ener_loss"] = self.display_if_exist( - l2_ener_loss.detach(), find_energy + se_e = torch.square(energy_pred - energy_label) # [nf, k] + if maskf is not None: + # Idiom 2 (extensive): per-frame normalization by real-atom count. + per_frame_e = torch.mean(se_e.reshape(_nf, -1), dim=-1) # [nf] + if not self.inference: + more_loss["l2_ener_loss"] = self.display_if_exist( + torch.mean(per_frame_e).detach(), find_energy + ) + loss += pref_e * torch.mean(per_frame_e * inv**norm_exp) + rmse_e = torch.sqrt(torch.mean(per_frame_e * inv**2)) + more_loss["rmse_e"] = self.display_if_exist( + rmse_e.detach(), find_energy + ) + else: + l2_ener_loss = torch.mean(se_e) + if not self.inference: + more_loss["l2_ener_loss"] = self.display_if_exist( + l2_ener_loss.detach(), find_energy + ) + loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) + rmse_e = l2_ener_loss.sqrt() * atom_norm + more_loss["rmse_e"] = self.display_if_exist( + rmse_e.detach(), find_energy ) - loss += atom_norm**norm_exp * (pref_e * l2_ener_loss) - rmse_e = l2_ener_loss.sqrt() * atom_norm - more_loss["rmse_e"] = self.display_if_exist( - rmse_e.detach(), find_energy - ) # more_loss['log_keys'].append('rmse_e') elif self.loss_func == "mae": l1_ener_loss = F.l1_loss( @@ -221,22 +250,46 @@ def forward( find_force_r = label.get("find_force", 0.0) pref_fr = pref_fr * find_force_r if self.loss_func == "mse": - diff_fr = label["force"] - model_pred["force"] - l2_force_real_loss = torch.mean(torch.square(diff_fr)) - if not self.inference: - more_loss["l2_force_r_loss"] = self.display_if_exist( - l2_force_real_loss.detach(), find_force_r + diff_fr = label["force"] - model_pred["force"] # [nf, nloc, 3] + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3). + maskf_col = maskf.reshape(_nf, _nloc, 1) # [nf, nloc, 1] + sq_fr = torch.square(diff_fr) * maskf_col # [nf, nloc, 3] + per_frame_fr_sum = sq_fr.reshape(_nf, -1).sum(dim=-1) # [nf] + per_frame_fr_dof = maskf.sum(dim=-1) * 3 # [nf] + l2_force_real_loss = torch.mean(per_frame_fr_sum / per_frame_fr_dof) + if not self.inference: + more_loss["l2_force_r_loss"] = self.display_if_exist( + l2_force_real_loss.detach(), find_force_r + ) + loss += (pref_fr * l2_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_fr = l2_force_real_loss.sqrt() + more_loss["rmse_fr"] = self.display_if_exist( + rmse_fr.detach(), find_force_r ) - loss += (pref_fr * l2_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) - rmse_fr = l2_force_real_loss.sqrt() - more_loss["rmse_fr"] = self.display_if_exist( - rmse_fr.detach(), find_force_r - ) - if mae: - mae_fr = torch.mean(torch.abs(diff_fr)) - more_loss["mae_fr"] = self.display_if_exist( - mae_fr.detach(), find_force_r + if mae: + abs_fr = torch.abs(diff_fr) * maskf_col + per_frame_mae_sum = abs_fr.reshape(_nf, -1).sum(dim=-1) + mae_fr = torch.mean(per_frame_mae_sum / per_frame_fr_dof) + more_loss["mae_fr"] = self.display_if_exist( + mae_fr.detach(), find_force_r + ) + else: + l2_force_real_loss = torch.mean(torch.square(diff_fr)) + if not self.inference: + more_loss["l2_force_r_loss"] = self.display_if_exist( + l2_force_real_loss.detach(), find_force_r + ) + loss += (pref_fr * l2_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_fr = l2_force_real_loss.sqrt() + more_loss["rmse_fr"] = self.display_if_exist( + rmse_fr.detach(), find_force_r ) + if mae: + mae_fr = torch.mean(torch.abs(diff_fr)) + more_loss["mae_fr"] = self.display_if_exist( + mae_fr.detach(), find_force_r + ) elif self.loss_func == "mae": l1_force_real_loss = F.l1_loss( label["force"], model_pred["force"], reduction="none" @@ -338,24 +391,45 @@ def forward( if self.has_v and "virial" in model_pred and "virial" in label: find_virial = label.get("find_virial", 0.0) pref_v = pref_v * find_virial - diff_v = label["virial"] - model_pred["virial"].reshape(-1, 9) + diff_v = label["virial"] - model_pred["virial"].reshape(-1, 9) # [nf, 9] if self.loss_func == "mse": - l2_virial_loss = torch.mean(torch.square(diff_v)) - if not self.inference: - more_loss["l2_virial_loss"] = self.display_if_exist( - l2_virial_loss.detach(), find_virial + if maskf is not None: + # Idiom 2 (extensive, k=9): per-frame normalization by real-atom count. + se_v = torch.square(diff_v) # [nf, 9] + per_frame_v = torch.mean(se_v, dim=-1) # [nf] + if not self.inference: + more_loss["l2_virial_loss"] = self.display_if_exist( + torch.mean(per_frame_v).detach(), find_virial + ) + loss += pref_v * torch.mean(per_frame_v * inv**norm_exp) + rmse_v = torch.sqrt(torch.mean(per_frame_v * inv**2)) + more_loss["rmse_v"] = self.display_if_exist( + rmse_v.detach(), find_virial ) - loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) - rmse_v = l2_virial_loss.sqrt() * atom_norm - more_loss["rmse_v"] = self.display_if_exist( - rmse_v.detach(), find_virial - ) - if mae: - mae_v = torch.mean(torch.abs(diff_v)) * atom_norm - more_loss["mae_v"] = self.display_if_exist( - mae_v.detach(), find_virial + if mae: + abs_v = torch.abs(diff_v) # [nf, 9] + per_frame_mae_v = torch.mean(abs_v, dim=-1) # [nf] + mae_v = torch.mean(per_frame_mae_v * inv) + more_loss["mae_v"] = self.display_if_exist( + mae_v.detach(), find_virial + ) + else: + l2_virial_loss = torch.mean(torch.square(diff_v)) + if not self.inference: + more_loss["l2_virial_loss"] = self.display_if_exist( + l2_virial_loss.detach(), find_virial + ) + loss += atom_norm**norm_exp * (pref_v * l2_virial_loss) + rmse_v = l2_virial_loss.sqrt() * atom_norm + more_loss["rmse_v"] = self.display_if_exist( + rmse_v.detach(), find_virial ) + if mae: + mae_v = torch.mean(torch.abs(diff_v)) * atom_norm + more_loss["mae_v"] = self.display_if_exist( + mae_v.detach(), find_virial + ) elif self.loss_func == "mae": l1_virial_loss = F.l1_loss( label["virial"].reshape(-1), diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 2396a181e0..524b41089b 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -21,6 +21,7 @@ from deepmd.dpmodel.loss.ener import ( EnergyLoss, ) +from deepmd.dpmodel.loss.ener_spin import EnergySpinLoss as EnergySpinLossDPModel from deepmd.dpmodel.loss.property import ( PropertyLoss, ) @@ -1249,3 +1250,441 @@ def test_intensive_ignores_mask(self): assert np.isclose(float(loss_m), float(loss_nm)), ( f"intensive property must ignore mask: {float(loss_m)} vs {float(loss_nm)}" ) + + +# --------------------------------------------------------------------------- +# Task 5: EnergySpinLoss -- energy (has_e), force_real (has_fr), virial (has_v) +# Leave force_mag / mask_mag (has_fm) COMPLETELY UNTOUCHED. +# --------------------------------------------------------------------------- + +# Spin-specific test constants: NM magnetic atoms per frame (same count in +# both frames so that the pt fancy-index .view(nframes,-1,3) stays valid). +_NM = 2 + +_MASK_MAG_A = np.zeros((1, NA, 1), dtype=bool) # [1, NA, 1] +_MASK_MAG_A[0, :_NM, 0] = True # first NM atoms of frame A are magnetic + +_MASK_MAG_B = np.zeros((1, NB, 1), dtype=bool) # [1, NB, 1] +_MASK_MAG_B[0, :_NM, 0] = True # first NM atoms of frame B are magnetic + +# Padded batch (nf=2, nloc=NP): frame A padding slots are not magnetic. +_MASK_MAG_PAD_SPIN = np.zeros((2, NP, 1), dtype=bool) # [2, NP, 1] +_MASK_MAG_PAD_SPIN[0, :_NM, 0] = True # frame A: first NM real atoms are magnetic +_MASK_MAG_PAD_SPIN[1, :_NM, 0] = True # frame B: first NM real atoms are magnetic + +_MASK_PAD_SPIN = np.array( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], dtype=np.float64 +) # [2, NP] + + +def _full_spin_dicts( + nf, nloc, energy_pred, energy_label, mask_mag, mask=None, **overrides +): + """Build complete model_pred and label_dict for EnergySpinLoss.call. + + EnergySpinLoss.call accesses "energy" unconditionally (for xp namespace), + so it must always be present. All other keys are guarded by has_* flags. + """ + model_pred = { + "energy": energy_pred, # [nf, 1] + "force": np.zeros((nf, nloc, 3), dtype=np.float64), + "force_mag": np.zeros((nf, nloc, 3), dtype=np.float64), + "mask_mag": mask_mag, # [nf, nloc, 1] + "virial": np.zeros((nf, 9), dtype=np.float64), + } + label_dict = { + "energy": energy_label, # [nf, 1] + "force": np.zeros((nf, nloc, 3), dtype=np.float64), + "force_mag": np.zeros((nf, nloc, 3), dtype=np.float64), + "virial": np.zeros((nf, 9), dtype=np.float64), + "find_energy": 1.0, + "find_force": 0.0, + "find_force_mag": 0.0, + "find_virial": 0.0, + } + if mask is not None: + model_pred["mask"] = mask + model_pred.update({k: v for k, v in overrides.items() if k in model_pred}) + label_dict.update({k: v for k, v in overrides.items() if k in label_dict}) + return model_pred, label_dict + + +class TestDPModelEnerSpinLossEnerGradAccum: + """Idiom 2 (extensive) for the energy term in EnergySpinLoss. + + Covers: mse (norm_exp=1 and 2); plus non-mixed no-op. + """ + + def _make_loss(self, intensive=False): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + intensive_ener_virial=intensive, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, e_A, e_A_hat, e_B, e_B_hat): + def make_A(): + p, l = _full_spin_dicts( + 1, + NA, + e_A, + e_A_hat, + _MASK_MAG_A, + mask=np.ones((1, NA), dtype=np.float64), + ) + return p, l, NA + + def make_B(): + p, l = _full_spin_dicts( + 1, + NB, + e_B, + e_B_hat, + _MASK_MAG_B, + mask=np.ones((1, NB), dtype=np.float64), + ) + return p, l, NB + + def make_padded(): + e_pad = np.concatenate([e_A, e_B], axis=0) # [2, 1] + e_hat_pad = np.concatenate([e_A_hat, e_B_hat], axis=0) + p, l = _full_spin_dicts( + 2, NP, e_pad, e_hat_pad, _MASK_MAG_PAD_SPIN, mask=_MASK_PAD_SPIN + ) + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Spin energy MSE norm_exp=1 meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant( + self._make_loss(intensive=False), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Spin energy MSE norm_exp=2 meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + self._run_invariant(self._make_loss(intensive=True), e_A, e_A_hat, e_B, e_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same energy loss as no mask.""" + e = _rnd(1, 1) + e_hat = _rnd(1, 1) + loss_obj = self._make_loss() + p_mask, l_mask = _full_spin_dicts( + 1, NP, e, e_hat, _MASK_MAG_B, mask=np.ones((1, NP), dtype=np.float64) + ) + p_nm, l_nm = _full_spin_dicts(1, NP, e, e_hat, _MASK_MAG_B) + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnerSpinLossForceRealGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=3) for force_real in EnergySpinLoss. + + Covers: mse; plus non-mixed no-op. + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): + def make_A(): + p, l = _full_spin_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_A, + mask=np.ones((1, NA), dtype=np.float64), + ) + p["force"] = f_A[None] # [1, NA, 3] + l["force"] = f_A_hat[None] + l["find_force"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_spin_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NB), dtype=np.float64), + ) + p["force"] = f_B[None] + l["force"] = f_B_hat[None] + l["find_force"] = 1.0 + return p, l, NB + + def make_padded(): + f_A_pad = np.zeros((NP, 3), dtype=np.float64) + f_A_pad[:NA] = f_A + f_A_hat_pad = np.zeros((NP, 3), dtype=np.float64) + f_A_hat_pad[:NA] = f_A_hat + f_pad = np.stack([f_A_pad, f_B], axis=0) # [2, NP, 3] + f_hat_pad = np.stack([f_A_hat_pad, f_B_hat], axis=0) + p, l = _full_spin_dicts( + 2, + NP, + np.zeros((2, 1)), + np.zeros((2, 1)), + _MASK_MAG_PAD_SPIN, + mask=_MASK_PAD_SPIN, + ) + p["force"] = f_pad + l["force"] = f_hat_pad + l["find_force"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Force_real MSE meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + self._run_invariant(self._make_loss(), f_A, f_A_hat, f_B, f_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same force_real loss as no mask.""" + f = _rnd(NP, 3) + f_hat = _rnd(NP, 3) + loss_obj = self._make_loss() + p_mask, l_mask = _full_spin_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NP), dtype=np.float64), + ) + p_mask["force"] = f[None] + l_mask["force"] = f_hat[None] + l_mask["find_force"] = 1.0 + p_nm, l_nm = _full_spin_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), _MASK_MAG_B + ) + p_nm["force"] = f[None] + l_nm["force"] = f_hat[None] + l_nm["find_force"] = 1.0 + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnerSpinLossVirialGradAccum: + """Idiom 2 (extensive, k=9) for virial in EnergySpinLoss. + + Covers: mse (norm_exp=1 and 2); plus non-mixed no-op. + """ + + def _make_loss(self, intensive=False): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + intensive_ener_virial=intensive, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, v_A, v_A_hat, v_B, v_B_hat): + def make_A(): + p, l = _full_spin_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_A, + mask=np.ones((1, NA), dtype=np.float64), + ) + p["virial"] = v_A[None] # [1, 9] + l["virial"] = v_A_hat[None] + l["find_virial"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_spin_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NB), dtype=np.float64), + ) + p["virial"] = v_B[None] + l["virial"] = v_B_hat[None] + l["find_virial"] = 1.0 + return p, l, NB + + def make_padded(): + v_pad = np.stack([v_A, v_B], axis=0) # [2, 9] + v_hat_pad = np.stack([v_A_hat, v_B_hat], axis=0) + p, l = _full_spin_dicts( + 2, + NP, + np.zeros((2, 1)), + np.zeros((2, 1)), + _MASK_MAG_PAD_SPIN, + mask=_MASK_PAD_SPIN, + ) + p["virial"] = v_pad + l["virial"] = v_hat_pad + l["find_virial"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Spin virial MSE norm_exp=1 meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant( + self._make_loss(intensive=False), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Spin virial MSE norm_exp=2 meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + self._run_invariant(self._make_loss(intensive=True), v_A, v_A_hat, v_B, v_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same virial loss as no mask.""" + v = _rnd(9) + v_hat = _rnd(9) + loss_obj = self._make_loss() + p_mask, l_mask = _full_spin_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NP), dtype=np.float64), + ) + p_mask["virial"] = v[None] + l_mask["virial"] = v_hat[None] + l_mask["find_virial"] = 1.0 + p_nm, l_nm = _full_spin_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), _MASK_MAG_B + ) + p_nm["virial"] = v[None] + l_nm["virial"] = v_hat[None] + l_nm["find_virial"] = 1.0 + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +class TestDPModelEnerSpinLossForceMagUnchanged: + """Guard: the padding mask must NOT affect the force_mag / mask_mag term. + + The force_mag path uses mask_mag (spin virtual-atom mask), which is a + completely separate concept from the padding mask model_dict["mask"]. + After the Task-5 implementation, presenting a padding mask must leave + the force_mag loss bit-identical. + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def test_padding_mask_does_not_affect_force_mag(self): + """force_mag loss is bit-identical with and without padding mask.""" + fm = _rnd(2, NP, 3) + fm_hat = _rnd(2, NP, 3) + loss_obj = self._make_loss() + + def _run(with_mask): + pred = { + "energy": np.zeros((2, 1), dtype=np.float64), + "force_mag": fm, + "mask_mag": _MASK_MAG_PAD_SPIN, + } + lbl = { + "force_mag": fm_hat, + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + if with_mask: + pred["mask"] = _MASK_PAD_SPIN + loss, _ = loss_obj.call(1.0, NP, pred, lbl) + return float(loss) + + loss_with = _run(True) + loss_without = _run(False) + assert np.isclose(loss_with, loss_without), ( + f"force_mag loss must be unchanged by padding mask: " + f"{loss_with} vs {loss_without}" + ) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 9a4a8d2f76..c56db8b669 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -21,6 +21,7 @@ from deepmd.pt.loss.ener import ( EnergyStdLoss, ) +from deepmd.pt.loss.ener_spin import EnergySpinLoss as EnergySpinLossPT from deepmd.pt.loss.loss import ( TaskLoss, ) @@ -1352,3 +1353,364 @@ def test_intensive_ignores_mask(self): assert torch.isclose(loss_m, loss_nm), ( f"intensive property must ignore mask: {loss_m.item()} vs {loss_nm.item()}" ) + + +# --------------------------------------------------------------------------- +# Task 5: EnergySpinLoss (pt) -- energy (has_e), force_real (has_fr), +# virial (has_v). Leave force_mag / mask_mag (has_fm) COMPLETELY UNTOUCHED. +# --------------------------------------------------------------------------- + +# Spin-specific test constants: NM magnetic atoms per frame (same count in +# both frames so that .view(nframes,-1,3) in the pt force_mag path is valid). +_NM_PT = 2 + +_MASK_MAG_A_PT = torch.zeros(1, NA, 1, dtype=torch.bool, device="cpu") +_MASK_MAG_A_PT[0, :_NM_PT, 0] = True # first NM_PT atoms of frame A magnetic + +_MASK_MAG_B_PT = torch.zeros(1, NB, 1, dtype=torch.bool, device="cpu") +_MASK_MAG_B_PT[0, :_NM_PT, 0] = True # first NM_PT atoms of frame B magnetic + +_MASK_MAG_PAD_SPIN_PT = torch.zeros(2, NP, 1, dtype=torch.bool, device="cpu") +_MASK_MAG_PAD_SPIN_PT[0, :_NM_PT, 0] = True # frame A +_MASK_MAG_PAD_SPIN_PT[1, :_NM_PT, 0] = True # frame B + +_MASK_PAD_SPIN_PT = torch.tensor( + [[1.0] * NA + [0.0] * (NP - NA), [1.0] * NB], + dtype=torch.float64, + device="cpu", +) # [2, NP] + + +def _spin_loss_fn(loss_obj, model_pred, label, natoms): + """Call EnergySpinLoss.forward via mock model; return scalar loss tensor.""" + _, loss, _ = loss_obj.forward( + input_dict={}, # mask already in model_pred; no re-injection needed + model=_MockModel(model_pred), + label=label, + natoms=natoms, + learning_rate=1.0, + ) + return loss + + +class TestPTEnerSpinLossEnerGradAccum: + """Idiom 2 (extensive) for the energy term in pt EnergySpinLoss. + + Covers: mse (norm_exp=1 and 2); plus non-mixed no-op. + """ + + def _make_loss(self, intensive=False): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + intensive_ener_virial=intensive, + ) + + def _run_invariant(self, loss_obj, e_A, e_A_hat, e_B, e_B_hat): + def make_A(): + mp = { + "energy": e_A, + "mask_mag": _MASK_MAG_A_PT, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"energy": e_A_hat, "find_energy": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "energy": e_B, + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"energy": e_B_hat, "find_energy": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "energy": torch.cat([e_A, e_B], dim=0), # [2, 1] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "energy": torch.cat([e_A_hat, e_B_hat], dim=0), + "find_energy": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Spin energy MSE norm_exp=1 meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant( + self._make_loss(intensive=False), e_A, e_A_hat, e_B, e_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Spin energy MSE norm_exp=2 meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + self._run_invariant(self._make_loss(intensive=True), e_A, e_A_hat, e_B, e_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same energy loss as no mask.""" + e = _t(1, 1) + e_hat = _t(1, 1) + loss_obj = self._make_loss() + mp_mask = { + "energy": e, + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"energy": e, "mask_mag": _MASK_MAG_B_PT} + lb = {"energy": e_hat, "find_energy": 1.0} + loss_m = _spin_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _spin_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnerSpinLossForceRealGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=3) for force_real in pt EnergySpinLoss. + + Covers: mse; plus non-mixed no-op. + """ + + def _make_loss(self): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, f_B, f_B_hat): + def make_A(): + mp = { + "force": f_A.unsqueeze(0), # [1, NA, 3] + "mask_mag": _MASK_MAG_A_PT, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"force": f_A_hat.unsqueeze(0), "find_force": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "force": f_B.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"force": f_B_hat.unsqueeze(0), "find_force": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "force": _padded_force_t(f_A, f_B), # [2, NP, 3] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "force": _padded_force_t(f_A_hat, f_B_hat), + "find_force": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Force_real MSE meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + self._run_invariant(self._make_loss(), f_A, f_A_hat, f_B, f_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same force_real loss as no mask.""" + f = _t(NP, 3) + f_hat = _t(NP, 3) + loss_obj = self._make_loss() + mp_mask = { + "force": f.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"force": f.unsqueeze(0), "mask_mag": _MASK_MAG_B_PT} + lb = {"force": f_hat.unsqueeze(0), "find_force": 1.0} + loss_m = _spin_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _spin_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnerSpinLossVirialGradAccum: + """Idiom 2 (extensive, k=9) for virial in pt EnergySpinLoss. + + Covers: mse (norm_exp=1 and 2); plus non-mixed no-op. + """ + + def _make_loss(self, intensive=False): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + intensive_ener_virial=intensive, + ) + + def _run_invariant(self, loss_obj, v_A, v_A_hat, v_B, v_B_hat): + def make_A(): + mp = { + "virial": v_A.unsqueeze(0), # [1, 9] + "mask_mag": _MASK_MAG_A_PT, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"virial": v_A_hat.unsqueeze(0), "find_virial": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "virial": v_B.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"virial": v_B_hat.unsqueeze(0), "find_virial": 1.0} + return mp, lb, NB + + def make_padded(): + mp = { + "virial": torch.stack([v_A, v_B], dim=0), # [2, 9] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "virial": torch.stack([v_A_hat, v_B_hat], dim=0), + "find_virial": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_non_intensive_grad_accum(self): + """Spin virial MSE norm_exp=1 meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant( + self._make_loss(intensive=False), v_A, v_A_hat, v_B, v_B_hat + ) + + def test_mse_intensive_grad_accum(self): + """Spin virial MSE norm_exp=2 meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + self._run_invariant(self._make_loss(intensive=True), v_A, v_A_hat, v_B, v_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same virial loss as no mask.""" + v = _t(9) + v_hat = _t(9) + loss_obj = self._make_loss() + mp_mask = { + "virial": v.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = {"virial": v.unsqueeze(0), "mask_mag": _MASK_MAG_B_PT} + lb = {"virial": v_hat.unsqueeze(0), "find_virial": 1.0} + loss_m = _spin_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _spin_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +class TestPTEnerSpinLossForceMagUnchanged: + """Guard: the padding mask must NOT affect the force_mag / mask_mag term. + + The force_mag path uses mask_mag (spin virtual-atom mask), which is a + completely separate concept from the padding mask model_pred["mask"]. + After the Task-5 implementation, presenting a padding mask must leave + the force_mag loss bit-identical. + """ + + def _make_loss(self): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def test_padding_mask_does_not_affect_force_mag(self): + """force_mag loss is bit-identical with and without padding mask.""" + fm = _rnd_t(2, NP, 3) + fm_hat = _rnd_t(2, NP, 3) + loss_obj = self._make_loss() + + def _run(with_mask): + mp = { + "force_mag": fm, + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + } + lb = { + "force_mag": fm_hat, + "find_force_mag": 1.0, + } + if with_mask: + mp["mask"] = _MASK_PAD_SPIN_PT + return _spin_loss_fn(loss_obj, mp, lb, NP) + + loss_with = _run(True) + loss_without = _run(False) + assert torch.isclose(loss_with, loss_without), ( + f"force_mag loss must be unchanged by padding mask: " + f"{loss_with.item()} vs {loss_without.item()}" + ) From 423be70e0496c2cd5701beb35eedf7def9e29ab5 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 01:56:03 +0800 Subject: [PATCH 07/16] fix(loss): per-frame normalize the spin MAE loss paths for mixed_type padding Energy/force_real/virial MAE branches in EnergySpinLoss (dpmodel and pt) used padded-scalar natoms or unmasked global means, violating the grad-accumulation invariant for mixed_type batches. Apply the same per-frame masked pattern already used for MSE, mirroring the authoritative treatment in ener.py (Task 3). Also fix the mae=True display metric for energy to use per-frame inv weighting instead of padded atom_norm. force_mag/mask_mag and MSE branches are untouched. --- deepmd/dpmodel/loss/ener_spin.py | 76 ++++++++++++----- deepmd/pt/loss/ener_spin.py | 82 +++++++++++++------ .../tests/common/dpmodel/test_loss_padding.py | 56 +++++++++++++ source/tests/pt/test_loss_padding.py | 60 ++++++++++++++ 4 files changed, 232 insertions(+), 42 deletions(-) diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index 4944916bde..c53ce22af3 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -172,14 +172,30 @@ def call( xp.sqrt(l2_ener_loss) * atom_norm, find_energy ) elif self.loss_func == "mae": - abs_diff_e = xp.abs(energy_pred - energy_label) - l1_ener_loss = xp.sum(abs_diff_e) - loss += pref_e * l1_ener_loss - more_loss["mae_e"] = self.display_if_exist( - xp.mean(abs_diff_e), find_energy - ) + l1_ener_loss = xp.mean(xp.abs(energy_pred - energy_label)) + if maskf is not None: + # Idiom 2 (extensive) with abs: per-frame normalization by real-atom count. + abs_e = xp.abs(energy_pred - energy_label) # [nf, k] + per_frame_ae = xp.mean( + xp.reshape(abs_e, (_nf, -1)), axis=-1 + ) # [nf] + l1_ener_masked = xp.mean(per_frame_ae * inv) + loss += pref_e * l1_ener_masked + more_loss["mae_e"] = self.display_if_exist( + l1_ener_masked, find_energy + ) + else: + loss += atom_norm * (pref_e * l1_ener_loss) + more_loss["mae_e"] = self.display_if_exist( + l1_ener_loss * atom_norm, find_energy + ) if mae: - mae_e = xp.mean(xp.abs(energy_pred - energy_label)) * atom_norm + if maskf is not None: + abs_e = xp.abs(energy_pred - energy_label) + per_frame_ae = xp.mean(xp.reshape(abs_e, (_nf, -1)), axis=-1) + mae_e = xp.mean(per_frame_ae * inv) + else: + mae_e = xp.mean(xp.abs(energy_pred - energy_label)) * atom_norm more_loss["mae_e"] = self.display_if_exist(mae_e, find_energy) mae_e_all = xp.mean(xp.abs(energy_pred - energy_label)) more_loss["mae_e_all"] = self.display_if_exist(mae_e_all, find_energy) @@ -221,14 +237,26 @@ def call( mae_fr = xp.mean(xp.abs(force_label - force_pred)) more_loss["mae_fr"] = self.display_if_exist(mae_fr, find_force) elif self.loss_func == "mae": - abs_diff_fr = xp.abs(force_label - force_pred) - per_atom_fr = xp.sum(abs_diff_fr, axis=-1) # [nf, na] - per_frame_fr = xp.mean(per_atom_fr, axis=-1) # [nf] - l1_force_real_loss = xp.sum(per_frame_fr) # scalar - loss += pref_fr * l1_force_real_loss - more_loss["mae_fr"] = self.display_if_exist( - xp.mean(abs_diff_fr), find_force - ) + abs_diff_fr = xp.abs(force_label - force_pred) # [nf, nloc, 3] + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3) with abs. + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) + abs_fr = abs_diff_fr * maskf_col # [nf, nloc, 3] + per_frame_sum = xp.sum( + xp.reshape(abs_fr, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) * 3 # [nf] + l1_force_real_masked = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_fr * l1_force_real_masked + more_loss["mae_fr"] = self.display_if_exist( + l1_force_real_masked, find_force + ) + else: + l1_force_real_loss = xp.mean(abs_diff_fr) + loss += pref_fr * l1_force_real_loss + more_loss["mae_fr"] = self.display_if_exist( + l1_force_real_loss, find_force + ) if self.has_fm: find_force_mag = label_dict.get("find_force_mag", 0.0) @@ -319,10 +347,20 @@ def call( more_loss["mae_v"] = self.display_if_exist(mae_v, find_virial) elif self.loss_func == "mae": l1_virial_loss = xp.mean(xp.abs(diff_v)) - loss += atom_norm * (pref_v * l1_virial_loss) - more_loss["mae_v"] = self.display_if_exist( - l1_virial_loss * atom_norm, find_virial - ) + if maskf is not None: + # Idiom 2 (extensive, k=9) with abs: per-frame normalization by real-atom count. + abs_v = xp.abs(diff_v) # [nf, 9] + per_frame_v = xp.mean(abs_v, axis=-1) # [nf] + l1_virial_masked = xp.mean(per_frame_v * inv) + loss += pref_v * l1_virial_masked + more_loss["mae_v"] = self.display_if_exist( + l1_virial_masked, find_virial + ) + else: + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss * atom_norm, find_virial + ) more_loss["rmse"] = xp.sqrt(loss) return loss, more_loss diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 3508031f5a..6c30b8b55b 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -222,24 +222,36 @@ def forward( l1_ener_loss = F.l1_loss( energy_pred.reshape(-1), energy_label.reshape(-1), - reduction="sum", - ) - loss += pref_e * l1_ener_loss - more_loss["mae_e"] = self.display_if_exist( - F.l1_loss( - energy_pred.reshape(-1), - energy_label.reshape(-1), - reduction="mean", - ).detach(), - find_energy, + reduction="mean", ) + if maskf is not None: + # Idiom 2 (extensive) with abs: per-frame normalization by real-atom count. + abs_e = torch.abs(energy_pred - energy_label) + per_frame_ae = torch.mean(abs_e.reshape(_nf, -1), dim=-1) # [nf] + l1_ener_masked = torch.mean(per_frame_ae * inv) + loss += pref_e * l1_ener_masked + more_loss["mae_e"] = self.display_if_exist( + l1_ener_masked.detach(), find_energy + ) + else: + loss += atom_norm * (pref_e * l1_ener_loss) + more_loss["mae_e"] = self.display_if_exist( + l1_ener_loss.detach() * atom_norm, find_energy + ) # more_loss['log_keys'].append('rmse_e') else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for energy loss." ) if mae: - mae_e = torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm + if maskf is not None: + abs_e = torch.abs(energy_pred - energy_label) + per_frame_ae = torch.mean(abs_e.reshape(_nf, -1), dim=-1) + mae_e = torch.mean(per_frame_ae * inv) + else: + mae_e = ( + torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm + ) more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy) mae_e_all = torch.mean(torch.abs(energy_pred - energy_label)) more_loss["mae_e_all"] = self.display_if_exist( @@ -291,14 +303,28 @@ def forward( mae_fr.detach(), find_force_r ) elif self.loss_func == "mae": - l1_force_real_loss = F.l1_loss( - label["force"], model_pred["force"], reduction="none" - ) - more_loss["mae_fr"] = self.display_if_exist( - l1_force_real_loss.mean().detach(), find_force_r - ) - l1_force_real_loss = l1_force_real_loss.sum(-1).mean(-1).sum() - loss += (pref_fr * l1_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) + abs_diff_fr = torch.abs( + label["force"] - model_pred["force"] + ) # [nf, nloc, 3] + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=3) with abs. + maskf_col = maskf.reshape(_nf, _nloc, 1) + abs_fr = abs_diff_fr * maskf_col + per_frame_sum = abs_fr.reshape(_nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) * 3 # [nf] + l1_force_real_masked = torch.mean(per_frame_sum / per_frame_dof) + more_loss["mae_fr"] = self.display_if_exist( + l1_force_real_masked.detach(), find_force_r + ) + loss += (pref_fr * l1_force_real_masked).to( + GLOBAL_PT_FLOAT_PRECISION + ) + else: + l1_force_real_loss = torch.mean(abs_diff_fr) + more_loss["mae_fr"] = self.display_if_exist( + l1_force_real_loss.detach(), find_force_r + ) + loss += (pref_fr * l1_force_real_loss).to(GLOBAL_PT_FLOAT_PRECISION) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for real force loss." @@ -436,10 +462,20 @@ def forward( model_pred["virial"].reshape(-1), reduction="mean", ) - loss += atom_norm * (pref_v * l1_virial_loss) - more_loss["mae_v"] = self.display_if_exist( - l1_virial_loss.detach() * atom_norm, find_virial - ) + if maskf is not None: + # Idiom 2 (extensive, k=9) with abs: per-frame normalization by real-atom count. + abs_v = torch.abs(diff_v) # [nf, 9] + per_frame_v = torch.mean(abs_v, dim=-1) # [nf] + l1_virial_masked = torch.mean(per_frame_v * inv) + loss += pref_v * l1_virial_masked + more_loss["mae_v"] = self.display_if_exist( + l1_virial_masked.detach(), find_virial + ) + else: + loss += atom_norm * (pref_v * l1_virial_loss) + more_loss["mae_v"] = self.display_if_exist( + l1_virial_loss.detach() * atom_norm, find_virial + ) else: raise NotImplementedError( f"Loss type {self.loss_func} is not implemented for virial loss." diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 524b41089b..1a8af39f22 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -1385,6 +1385,24 @@ def test_mse_intensive_grad_accum(self): e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) self._run_invariant(self._make_loss(intensive=True), e_A, e_A_hat, e_B, e_B_hat) + def test_mae_grad_accum(self): + """Spin energy MAE meets the grad-accum invariant.""" + e_A, e_B = _rnd(1, 1), _rnd(1, 1) + e_A_hat, e_B_hat = _rnd(1, 1), _rnd(1, 1) + loss_obj = EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, e_A, e_A_hat, e_B, e_B_hat) + def test_no_op_for_non_mixed(self): """All-ones mask gives same energy loss as no mask.""" e = _rnd(1, 1) @@ -1488,6 +1506,26 @@ def test_mse_grad_accum(self): f_B_hat = _rnd(NB, 3) self._run_invariant(self._make_loss(), f_A, f_A_hat, f_B, f_B_hat) + def test_mae_grad_accum(self): + """Force_real MAE meets the grad-accum invariant.""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + loss_obj = EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, f_A, f_A_hat, f_B, f_B_hat) + def test_no_op_for_non_mixed(self): """All-ones mask gives same force_real loss as no mask.""" f = _rnd(NP, 3) @@ -1607,6 +1645,24 @@ def test_mse_intensive_grad_accum(self): v_A_hat, v_B_hat = _rnd(9), _rnd(9) self._run_invariant(self._make_loss(intensive=True), v_A, v_A_hat, v_B, v_B_hat) + def test_mae_grad_accum(self): + """Spin virial MAE meets the grad-accum invariant.""" + v_A, v_B = _rnd(9), _rnd(9) + v_A_hat, v_B_hat = _rnd(9), _rnd(9) + loss_obj = EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, v_A, v_A_hat, v_B, v_B_hat) + def test_no_op_for_non_mixed(self): """All-ones mask gives same virial loss as no mask.""" v = _rnd(9) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index c56db8b669..54203134d0 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -1469,6 +1469,26 @@ def test_mse_intensive_grad_accum(self): e_B_hat = _t(1, 1) self._run_invariant(self._make_loss(intensive=True), e_A, e_A_hat, e_B, e_B_hat) + def test_mae_grad_accum(self): + """Spin energy MAE meets the grad-accum invariant.""" + e_A = _t(1, 1) + e_A_hat = _t(1, 1) + e_B = _t(1, 1) + e_B_hat = _t(1, 1) + loss_obj = EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=1.0, + limit_pref_e=1.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, e_A, e_A_hat, e_B, e_B_hat) + def test_no_op_for_non_mixed(self): """All-ones mask gives same energy loss as no mask.""" e = _t(1, 1) @@ -1553,6 +1573,26 @@ def test_mse_grad_accum(self): f_B_hat = _t(NB, 3) self._run_invariant(self._make_loss(), f_A, f_A_hat, f_B, f_B_hat) + def test_mae_grad_accum(self): + """Force_real MAE meets the grad-accum invariant.""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + loss_obj = EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=1.0, + limit_pref_fr=1.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, f_A, f_A_hat, f_B, f_B_hat) + def test_no_op_for_non_mixed(self): """All-ones mask gives same force_real loss as no mask.""" f = _t(NP, 3) @@ -1648,6 +1688,26 @@ def test_mse_intensive_grad_accum(self): v_B_hat = _t(9) self._run_invariant(self._make_loss(intensive=True), v_A, v_A_hat, v_B, v_B_hat) + def test_mae_grad_accum(self): + """Spin virial MAE meets the grad-accum invariant.""" + v_A = _t(9) + v_A_hat = _t(9) + v_B = _t(9) + v_B_hat = _t(9) + loss_obj = EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=1.0, + limit_pref_v=1.0, + loss_func="mae", + ) + self._run_invariant(loss_obj, v_A, v_A_hat, v_B, v_B_hat) + def test_no_op_for_non_mixed(self): """All-ones mask gives same virial loss as no mask.""" v = _t(9) From 4381d7406b4dc074ac484f2081630779fea6186b Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 5 Jul 2026 23:49:36 +0800 Subject: [PATCH 08/16] fix(loss): pass axis as a keyword to xp.sum in the property loss The extensive property mask normalization called xp.sum(model_dict["mask"], -1) with axis positional; array_api_compat's torch namespace declares axis keyword-only, so this raised TypeError for every pt_expt extensive-property training. Use axis=-1, matching the dos/tensor/ener losses. Adds a torch-tensor test through PropertyLoss.call that reproduces the crash on the old form. --- deepmd/dpmodel/loss/property.py | 2 +- .../tests/common/dpmodel/test_loss_padding.py | 24 +++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/loss/property.py b/deepmd/dpmodel/loss/property.py index 627c299abf..ef649c9b2f 100644 --- a/deepmd/dpmodel/loss/property.py +++ b/deepmd/dpmodel/loss/property.py @@ -88,7 +88,7 @@ def call( if "mask" in model_dict: # Per-frame real atom count: shape [nf] → broadcast over [nf, task_dim]. real_natoms = xp.reshape( - xp.astype(xp.sum(model_dict["mask"], -1), pred.dtype), + xp.astype(xp.sum(model_dict["mask"], axis=-1), pred.dtype), (-1,) + (1,) * (pred.ndim - 1), ) pred = pred / real_natoms diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 1a8af39f22..2bc822ca82 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -1226,6 +1226,30 @@ def test_no_op_for_non_mixed(self): f"all-ones mask must be no-op: {float(loss_m)} vs {float(loss_nm)}" ) + def test_torch_backend_matches_numpy(self): + """The array-API path must work with torch tensors (pt_expt runs this). + + A padded batch is fed once as numpy and once as torch; the extensive + per-frame normalization uses xp.sum(mask, axis=-1), which raises under + the array_api_compat torch namespace if axis is passed positionally. + """ + import pytest + + torch = pytest.importorskip("torch") + p = _rnd(2, PROP_TASK_DIM) + lab = _rnd(2, PROP_TASK_DIM) + loss_obj = self._make_loss("mse") + np_pred = {PROP_VAR: p, "mask": _MASK_PAD} + pt_pred = { + PROP_VAR: torch.tensor(p), + "mask": torch.tensor(_MASK_PAD), + } + loss_np, _ = loss_obj.call(1.0, NP, np_pred, {PROP_VAR: lab}) + loss_pt, _ = loss_obj.call(1.0, NP, pt_pred, {PROP_VAR: torch.tensor(lab)}) + assert np.isclose(float(loss_np), float(loss_pt)), ( + f"torch path must match numpy: {float(loss_np)} vs {float(loss_pt)}" + ) + class TestPropertyLossIntensiveUnaffectedByMask: """Intensive property loss must be unchanged whether or not mask is present.""" From 3a342d87d08d7969b45f619a51536fa3c6e4a69c Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 00:03:35 +0800 Subject: [PATCH 09/16] fix(loss): per-frame normalize the spin atom-energy loss for mixed_type padding MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Apply idiom 1 (per-atom masked mean, ncomp=1) to the has_ae block in both EnergySpinLoss backends so that ghost-padded atoms (mask=0) are excluded from the atom_ener loss normalization. Before this fix the loss was computed as mean(square(diff)) over a flattened [nf*nloc] vector, so ghost atoms diluted the denominator in mixed_type batches. The masked path uses xp.reshape(maskf, (_nf, _nloc, 1)) as a broadcastable column mask, then idiom-1 per-frame sum / dof → mean over frames. The unmasked (no-mask) fallback is unchanged. Also add grad-accum invariant tests for: ener_spin.has_ae (mse/mae, RED→GREEN), gen_force/has_gf (already GREEN — ghost forces are masked before projection so the mean(square) over [nf, ngen] is frame-decomposable), and force_mag MSE (GREEN — n_valid excludes ghost atoms via mask_mag=0; invariant holds when NM is equal across frames, which is the practical spin-model case). force_mag MAE uses xp.sum over frames instead of xp.mean (2x accumulation artifact, pre-existing, independent of mask_mag semantics) and is reported as NEEDS_CONTEXT. --- deepmd/dpmodel/loss/ener_spin.py | 62 ++- deepmd/pt/loss/ener_spin.py | 87 +++-- .../tests/common/dpmodel/test_loss_padding.py | 352 ++++++++++++++++++ source/tests/pt/test_loss_padding.py | 287 ++++++++++++++ 4 files changed, 743 insertions(+), 45 deletions(-) diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index c53ce22af3..dd5ebd3460 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -297,24 +297,50 @@ def call( pref_ae = pref_ae * find_atom_ener atom_ener = model_dict["atom_energy"] atom_ener_label = label_dict["atom_ener"] - atom_ener_reshape = xp.reshape(atom_ener, (-1,)) - atom_ener_label_reshape = xp.reshape(atom_ener_label, (-1,)) - if self.loss_func == "mse": - l2_atom_ener_loss = xp.mean( - xp.square(atom_ener_label_reshape - atom_ener_reshape) - ) - loss += pref_ae * l2_atom_ener_loss - more_loss["rmse_ae"] = self.display_if_exist( - xp.sqrt(l2_atom_ener_loss), find_atom_ener - ) - elif self.loss_func == "mae": - l1_atom_ener_loss = xp.mean( - xp.abs(atom_ener_label_reshape - atom_ener_reshape) - ) - loss += pref_ae * l1_atom_ener_loss - more_loss["mae_ae"] = self.display_if_exist( - l1_atom_ener_loss, find_atom_ener - ) + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=1). + ae = xp.reshape(atom_ener, (_nf, _nloc, 1)) + ae_label = xp.reshape(atom_ener_label, (_nf, _nloc, 1)) + maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1] + if self.loss_func == "mse": + sq = xp.square(ae_label - ae) * maskf_col # [nf, nloc, 1] + per_frame_sum = xp.sum(xp.reshape(sq, (_nf, -1)), axis=-1) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) # [nf] (ncomp=1) + l2_atom_ener_loss = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_ae * l2_atom_ener_loss + more_loss["rmse_ae"] = self.display_if_exist( + xp.sqrt(l2_atom_ener_loss), find_atom_ener + ) + elif self.loss_func == "mae": + abs_diff = xp.abs(ae_label - ae) * maskf_col # [nf, nloc, 1] + per_frame_sum = xp.sum( + xp.reshape(abs_diff, (_nf, -1)), axis=-1 + ) # [nf] + per_frame_dof = xp.sum(maskf, axis=-1) # [nf] (ncomp=1) + l1_atom_ener_loss = xp.mean(per_frame_sum / per_frame_dof) + loss += pref_ae * l1_atom_ener_loss + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss, find_atom_ener + ) + else: + atom_ener_reshape = xp.reshape(atom_ener, (-1,)) + atom_ener_label_reshape = xp.reshape(atom_ener_label, (-1,)) + if self.loss_func == "mse": + l2_atom_ener_loss = xp.mean( + xp.square(atom_ener_label_reshape - atom_ener_reshape) + ) + loss += pref_ae * l2_atom_ener_loss + more_loss["rmse_ae"] = self.display_if_exist( + xp.sqrt(l2_atom_ener_loss), find_atom_ener + ) + elif self.loss_func == "mae": + l1_atom_ener_loss = xp.mean( + xp.abs(atom_ener_label_reshape - atom_ener_reshape) + ) + loss += pref_ae * l1_atom_ener_loss + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss, find_atom_ener + ) if self.has_v: find_virial = label_dict.get("find_virial", 0.0) diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 6c30b8b55b..7593c7c246 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -383,36 +383,69 @@ def forward( atom_ener_label = label["atom_ener"] find_atom_ener = label.get("find_atom_ener", 0.0) pref_ae = pref_ae * find_atom_ener - atom_ener_reshape = atom_ener.reshape(-1) - atom_ener_label_reshape = atom_ener_label.reshape(-1) - if self.loss_func == "mse": - l2_atom_ener_loss = torch.square( - atom_ener_label_reshape - atom_ener_reshape - ).mean() - if not self.inference: - more_loss["l2_atom_ener_loss"] = self.display_if_exist( - l2_atom_ener_loss.detach(), find_atom_ener + if maskf is not None: + # Idiom 1 (per-atom masked mean, ncomp=1). + ae = atom_ener.reshape(_nf, _nloc, 1) + ae_label = atom_ener_label.reshape(_nf, _nloc, 1) + maskf_col = maskf.reshape(_nf, _nloc, 1) # [nf, nloc, 1] + if self.loss_func == "mse": + sq = torch.square(ae_label - ae) * maskf_col # [nf, nloc, 1] + per_frame_sum = sq.reshape(_nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) # [nf] (ncomp=1) + l2_atom_ener_loss = torch.mean(per_frame_sum / per_frame_dof) + if not self.inference: + more_loss["l2_atom_ener_loss"] = self.display_if_exist( + l2_atom_ener_loss.detach(), find_atom_ener + ) + loss += (pref_ae * l2_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_ae = l2_atom_ener_loss.sqrt() + more_loss["rmse_ae"] = self.display_if_exist( + rmse_ae.detach(), find_atom_ener + ) + elif self.loss_func == "mae": + abs_diff = torch.abs(ae_label - ae) * maskf_col # [nf, nloc, 1] + per_frame_sum = abs_diff.reshape(_nf, -1).sum(dim=-1) # [nf] + per_frame_dof = maskf.sum(dim=-1) # [nf] (ncomp=1) + l1_atom_ener_loss = torch.mean(per_frame_sum / per_frame_dof) + loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss.detach(), find_atom_ener + ) + else: + raise NotImplementedError( + f"Loss type {self.loss_func} is not implemented for atomic energy loss." ) - loss += (pref_ae * l2_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) - rmse_ae = l2_atom_ener_loss.sqrt() - more_loss["rmse_ae"] = self.display_if_exist( - rmse_ae.detach(), find_atom_ener - ) - elif self.loss_func == "mae": - l1_atom_ener_loss = F.l1_loss( - atom_ener_reshape, - atom_ener_label_reshape, - reduction="mean", - ) - loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) - more_loss["mae_ae"] = self.display_if_exist( - l1_atom_ener_loss.detach(), find_atom_ener - ) else: - raise NotImplementedError( - f"Loss type {self.loss_func} is not implemented for atomic energy loss." - ) + atom_ener_reshape = atom_ener.reshape(-1) + atom_ener_label_reshape = atom_ener_label.reshape(-1) + if self.loss_func == "mse": + l2_atom_ener_loss = torch.square( + atom_ener_label_reshape - atom_ener_reshape + ).mean() + if not self.inference: + more_loss["l2_atom_ener_loss"] = self.display_if_exist( + l2_atom_ener_loss.detach(), find_atom_ener + ) + loss += (pref_ae * l2_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + rmse_ae = l2_atom_ener_loss.sqrt() + more_loss["rmse_ae"] = self.display_if_exist( + rmse_ae.detach(), find_atom_ener + ) + elif self.loss_func == "mae": + l1_atom_ener_loss = F.l1_loss( + atom_ener_reshape, + atom_ener_label_reshape, + reduction="mean", + ) + loss += (pref_ae * l1_atom_ener_loss).to(GLOBAL_PT_FLOAT_PRECISION) + more_loss["mae_ae"] = self.display_if_exist( + l1_atom_ener_loss.detach(), find_atom_ener + ) + else: + raise NotImplementedError( + f"Loss type {self.loss_func} is not implemented for atomic energy loss." + ) if self.has_v and "virial" in model_pred and "virial" in label: find_virial = label.get("find_virial", 0.0) diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 2bc822ca82..ecc96410bc 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -1768,3 +1768,355 @@ def _run(with_mask): f"force_mag loss must be unchanged by padding mask: " f"{loss_with} vs {loss_without}" ) + + +# --------------------------------------------------------------------------- +# Part A: EnergySpinLoss atom_ener (has_ae) grad-accum invariant +# --------------------------------------------------------------------------- + + +class TestDPModelEnerSpinLossAtomEnerGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=1) for atom_ener in EnergySpinLoss. + + RED before the Part-A has_ae mask fix; GREEN after. + """ + + def _make_loss(self, loss_func="mse"): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + loss_func=loss_func, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat): + def make_A(): + p, l = _full_spin_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_A, + mask=np.ones((1, NA), dtype=np.float64), + ) + p["atom_energy"] = ae_A[None] # [1, NA, 1] + l["atom_ener"] = ae_A_hat[None] + l["find_atom_ener"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_spin_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NB), dtype=np.float64), + ) + p["atom_energy"] = ae_B[None] + l["atom_ener"] = ae_B_hat[None] + l["find_atom_ener"] = 1.0 + return p, l, NB + + def make_padded(): + ae_pad = _padded_atom(ae_A, ae_B, 1) # [2, NP, 1] + ae_hat_pad = _padded_atom(ae_A_hat, ae_B_hat, 1) + p, l = _full_spin_dicts( + 2, + NP, + np.zeros((2, 1)), + np.zeros((2, 1)), + _MASK_MAG_PAD_SPIN, + mask=_MASK_PAD_SPIN, + ) + p["atom_energy"] = ae_pad + l["atom_ener"] = ae_hat_pad + l["find_atom_ener"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Spin atom_ener MSE meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + self._run_invariant(self._make_loss("mse"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_mae_grad_accum(self): + """Spin atom_ener MAE meets the grad-accum invariant.""" + ae_A = _rnd(NA, 1) + ae_A_hat = _rnd(NA, 1) + ae_B = _rnd(NB, 1) + ae_B_hat = _rnd(NB, 1) + self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom_ener spin loss as no mask.""" + ae = _rnd(NP, 1) + ae_hat = _rnd(NP, 1) + loss_obj = self._make_loss() + p_mask, l_mask = _full_spin_dicts( + 1, + NP, + np.zeros((1, 1)), + np.zeros((1, 1)), + _MASK_MAG_B, + mask=np.ones((1, NP), dtype=np.float64), + ) + p_mask["atom_energy"] = ae[None] + l_mask["atom_ener"] = ae_hat[None] + l_mask["find_atom_ener"] = 1.0 + p_nm, l_nm = _full_spin_dicts( + 1, NP, np.zeros((1, 1)), np.zeros((1, 1)), _MASK_MAG_B + ) + p_nm["atom_energy"] = ae[None] + l_nm["atom_ener"] = ae_hat[None] + l_nm["find_atom_ener"] = 1.0 + loss_m = self._loss_fn(loss_obj, p_mask, l_mask, NP) + loss_nm = self._loss_fn(loss_obj, p_nm, l_nm, NP) + assert np.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m} vs {loss_nm}" + ) + + +# --------------------------------------------------------------------------- +# gen_force (has_gf) grad-accum invariant +# Ghost-atom forces are masked before projection, so the invariant is met. +# --------------------------------------------------------------------------- + +_NGEN = 2 # number of generalized coordinates for tests + + +class TestDPModelEnergyLossGenForceGradAccum: + """gen_force (has_gf) excludes ghost atoms via force masking before projection. + + The projected gen_force = sum_i(drdq_ij * masked_f_i) is frame-decomposable, + so mean(square(diff)) over [nf, ngen] already satisfies the invariant. + Expected GREEN immediately (no fix needed for gen_force). + """ + + def _make_loss(self): + return EnergyLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=0.0, + limit_pref_ae=0.0, + start_pref_pf=0.0, + limit_pref_pf=0.0, + start_pref_gf=1.0, + limit_pref_gf=1.0, + numb_generalized_coord=_NGEN, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, drdq_A, f_B, f_B_hat, drdq_B): + def make_A(): + p, l = _full_ener_dicts( + 1, + NA, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NA), dtype=np.float64), + ) + p["force"] = f_A[None] # [1, NA, 3] + l["force"] = f_A_hat[None] + l["find_force"] = 1.0 + l["drdq"] = drdq_A[None] # [1, NA*3, NGEN] + l["find_drdq"] = 1.0 + return p, l, NA + + def make_B(): + p, l = _full_ener_dicts( + 1, + NB, + np.zeros((1, 1)), + np.zeros((1, 1)), + mask=np.ones((1, NB), dtype=np.float64), + ) + p["force"] = f_B[None] + l["force"] = f_B_hat[None] + l["find_force"] = 1.0 + l["drdq"] = drdq_B[None] # [1, NB*3, NGEN] + l["find_drdq"] = 1.0 + return p, l, NB + + def make_padded(): + f_pad = _padded_force(f_A, f_B) # [2, NP, 3] + f_hat_pad = _padded_force(f_A_hat, f_B_hat) + # drdq ghost-atom slots are zero (ghost forces also zero, so no contribution) + drdq_A_pad = np.zeros((NP * 3, _NGEN), dtype=np.float64) + drdq_A_pad[: NA * 3] = drdq_A + drdq_pad = np.stack([drdq_A_pad, drdq_B], axis=0) # [2, NP*3, NGEN] + p, l = _full_ener_dicts( + 2, NP, np.zeros((2, 1)), np.zeros((2, 1)), mask=_MASK_PAD + ) + p["force"] = f_pad + l["force"] = f_hat_pad + l["find_force"] = 1.0 + l["drdq"] = drdq_pad + l["find_drdq"] = 1.0 + return p, l, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """gen_force MSE meets the grad-accum invariant (GREEN: already correct).""" + f_A = _rnd(NA, 3) + f_A_hat = _rnd(NA, 3) + drdq_A = _rnd(NA * 3, _NGEN) + f_B = _rnd(NB, 3) + f_B_hat = _rnd(NB, 3) + drdq_B = _rnd(NB * 3, _NGEN) + self._run_invariant( + self._make_loss(), f_A, f_A_hat, drdq_A, f_B, f_B_hat, drdq_B + ) + + +# --------------------------------------------------------------------------- +# force_mag MSE grad-accum invariant (NM equal across frames, ghost atoms non-magnetic) +# +# NOTE: force_mag MAE (dpmodel) uses xp.sum over frames instead of xp.mean, +# so MAE force_mag FAILS the invariant with a 2x factor when frames=2. +# This is a pre-existing frame-normalization artifact independent of ghost-atom +# masking. Ghost atoms are correctly excluded (mask_mag=0 there). The MAE +# artifact is reported as NEEDS_CONTEXT in the audit report. +# --------------------------------------------------------------------------- + + +class TestDPModelEnerSpinLossForceMagMSEGradAccum: + """force_mag MSE meets the grad-accum invariant when ghost atoms are non-magnetic. + + With equal magnetic-atom counts (NM) across frames and ghost atoms having + mask_mag=0, the global normalization n_valid = NM*nf correctly decomposes + per-frame, giving loss_pad == 0.5*(loss_A + loss_B). + Expected GREEN immediately (no fix needed for MSE force_mag). + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + def test_mse_grad_accum(self): + """force_mag MSE (NM equal per frame, ghost atoms non-magnetic) meets invariant.""" + fm_A = _rnd(NA, 3) + fm_A_hat = _rnd(NA, 3) + fm_B = _rnd(NB, 3) + fm_B_hat = _rnd(NB, 3) + + def make_A(): + # Frame A: NA=3 real atoms, first NM=2 are magnetic + fm_A_full = np.zeros((NA, 3), dtype=np.float64) + fm_A_full[:_NM] = fm_A[:_NM] + fm_A_hat_full = np.zeros((NA, 3), dtype=np.float64) + fm_A_hat_full[:_NM] = fm_A_hat[:_NM] + pred = { + "energy": np.zeros((1, 1), dtype=np.float64), + "force_mag": fm_A_full[None], # [1, NA, 3] + "mask_mag": _MASK_MAG_A, + } + lbl = { + "force_mag": fm_A_hat_full[None], + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NA + + def make_B(): + fm_B_full = np.zeros((NB, 3), dtype=np.float64) + fm_B_full[:_NM] = fm_B[:_NM] + fm_B_hat_full = np.zeros((NB, 3), dtype=np.float64) + fm_B_hat_full[:_NM] = fm_B_hat[:_NM] + pred = { + "energy": np.zeros((1, 1), dtype=np.float64), + "force_mag": fm_B_full[None], # [1, NB, 3] + "mask_mag": _MASK_MAG_B, + } + lbl = { + "force_mag": fm_B_hat_full[None], + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NB + + def make_padded(): + # Pad frame A force_mag to NP width; ghost slots are zero (non-magnetic) + fm_A_pad = np.zeros((NP, 3), dtype=np.float64) + fm_A_pad[:_NM] = fm_A[:_NM] + fm_A_hat_pad = np.zeros((NP, 3), dtype=np.float64) + fm_A_hat_pad[:_NM] = fm_A_hat[:_NM] + fm_B_pad = np.zeros((NP, 3), dtype=np.float64) + fm_B_pad[:_NM] = fm_B[:_NM] + fm_B_hat_pad = np.zeros((NP, 3), dtype=np.float64) + fm_B_hat_pad[:_NM] = fm_B_hat[:_NM] + fm_pad = np.stack([fm_A_pad, fm_B_pad], axis=0) # [2, NP, 3] + fm_hat_pad = np.stack([fm_A_hat_pad, fm_B_hat_pad], axis=0) + pred = { + "energy": np.zeros((2, 1), dtype=np.float64), + "force_mag": fm_pad, + "mask_mag": _MASK_MAG_PAD_SPIN, # ghost atoms have mask_mag=False + "mask": _MASK_PAD_SPIN, + } + lbl = { + "force_mag": fm_hat_pad, + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(self._make_loss(), mp, lb, na), + make_A, + make_B, + make_padded, + ) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 54203134d0..b92f99ed3a 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -1774,3 +1774,290 @@ def _run(with_mask): f"force_mag loss must be unchanged by padding mask: " f"{loss_with.item()} vs {loss_without.item()}" ) + + +# --------------------------------------------------------------------------- +# Part A: EnergySpinLoss atom_ener (has_ae) grad-accum invariant (pt) +# --------------------------------------------------------------------------- + + +class TestPTEnerSpinLossAtomEnerGradAccum: + """Idiom 1 (per-atom masked mean, ncomp=1) for atom_ener in pt EnergySpinLoss. + + RED before the Part-A has_ae mask fix; GREEN after. + """ + + def _make_loss(self, loss_func="mse"): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=0.0, + limit_pref_fm=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_ae=1.0, + limit_pref_ae=1.0, + loss_func=loss_func, + ) + + def _run_invariant(self, loss_obj, ae_A, ae_A_hat, ae_B, ae_B_hat): + def make_A(): + mp = { + "atom_energy": ae_A.unsqueeze(0), # [1, NA, 1] + "mask_mag": _MASK_MAG_A_PT, + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = {"atom_ener": ae_A_hat.unsqueeze(0), "find_atom_ener": 1.0} + return mp, lb, NA + + def make_B(): + mp = { + "atom_energy": ae_B.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = {"atom_ener": ae_B_hat.unsqueeze(0), "find_atom_ener": 1.0} + return mp, lb, NB + + def make_padded(): + ae_pad = _padded_atom_t(ae_A, ae_B, 1) # [2, NP, 1] + ae_hat_pad = _padded_atom_t(ae_A_hat, ae_B_hat, 1) + mp = { + "atom_energy": ae_pad, + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = {"atom_ener": ae_hat_pad, "find_atom_ener": 1.0} + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """Spin atom_ener MSE meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + self._run_invariant(self._make_loss("mse"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_mae_grad_accum(self): + """Spin atom_ener MAE meets the grad-accum invariant.""" + ae_A = _t(NA, 1) + ae_A_hat = _t(NA, 1) + ae_B = _t(NB, 1) + ae_B_hat = _t(NB, 1) + self._run_invariant(self._make_loss("mae"), ae_A, ae_A_hat, ae_B, ae_B_hat) + + def test_no_op_for_non_mixed(self): + """All-ones mask gives same atom_ener spin loss as no mask.""" + ae = _t(NP, 1) + ae_hat = _t(NP, 1) + loss_obj = self._make_loss() + mp_mask = { + "atom_energy": ae.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + "mask": torch.ones(1, NP, dtype=torch.float64, device="cpu"), + } + mp_nm = { + "atom_energy": ae.unsqueeze(0), + "mask_mag": _MASK_MAG_B_PT, + } + lb = {"atom_ener": ae_hat.unsqueeze(0), "find_atom_ener": 1.0} + loss_m = _spin_loss_fn(loss_obj, mp_mask, lb, NP) + loss_nm = _spin_loss_fn(loss_obj, mp_nm, lb, NP) + assert torch.isclose(loss_m, loss_nm), ( + f"all-ones mask must be no-op: {loss_m.item()} vs {loss_nm.item()}" + ) + + +# --------------------------------------------------------------------------- +# gen_force (has_gf) grad-accum invariant (pt) +# --------------------------------------------------------------------------- + +_NGEN_PT = 2 # number of generalized coordinates for tests + + +class TestPTEnergyLossGenForceGradAccum: + """gen_force (has_gf) excludes ghost atoms via force masking before projection. + + Expected GREEN immediately (no fix needed for gen_force). + """ + + def _make_loss(self): + return EnergyStdLoss( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_f=0.0, + limit_pref_f=0.0, + start_pref_v=0.0, + limit_pref_v=0.0, + start_pref_gf=1.0, + limit_pref_gf=1.0, + numb_generalized_coord=_NGEN_PT, + ) + + def _run_invariant(self, loss_obj, f_A, f_A_hat, drdq_A, f_B, f_B_hat, drdq_B): + def make_A(): + mp = { + "force": f_A.unsqueeze(0), # [1, NA, 3] + "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + } + lb = { + "force": f_A_hat.unsqueeze(0), + "find_force": 1.0, + "drdq": drdq_A.unsqueeze(0), # [1, NA*3, NGEN] + "find_drdq": 1.0, + } + return mp, lb, NA + + def make_B(): + mp = { + "force": f_B.unsqueeze(0), + "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + } + lb = { + "force": f_B_hat.unsqueeze(0), + "find_force": 1.0, + "drdq": drdq_B.unsqueeze(0), # [1, NB*3, NGEN] + "find_drdq": 1.0, + } + return mp, lb, NB + + def make_padded(): + f_pad = _padded_force_t(f_A, f_B) # [2, NP, 3] + f_hat_pad = _padded_force_t(f_A_hat, f_B_hat) + # drdq ghost-atom slots are zero (ghost forces also zero, no contribution) + drdq_A_pad = torch.zeros( + NP * 3, _NGEN_PT, dtype=torch.float64, device="cpu" + ) + drdq_A_pad[: NA * 3] = drdq_A + drdq_pad = torch.stack([drdq_A_pad, drdq_B], dim=0) # [2, NP*3, NGEN] + mp = {"force": f_pad, "mask": _MASK_PAD_PT} + lb = { + "force": f_hat_pad, + "find_force": 1.0, + "drdq": drdq_pad, + "find_drdq": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _ener_loss_fn(loss_obj, mp, lb, na), + make_A, + make_B, + make_padded, + ) + + def test_mse_grad_accum(self): + """gen_force MSE meets the grad-accum invariant (GREEN: already correct).""" + f_A = _t(NA, 3) + f_A_hat = _t(NA, 3) + drdq_A = _t(NA * 3, _NGEN_PT) + f_B = _t(NB, 3) + f_B_hat = _t(NB, 3) + drdq_B = _t(NB * 3, _NGEN_PT) + self._run_invariant( + self._make_loss(), f_A, f_A_hat, drdq_A, f_B, f_B_hat, drdq_B + ) + + +# --------------------------------------------------------------------------- +# force_mag MSE grad-accum invariant (pt) +# +# NOTE: force_mag MAE (pt) uses .sum() over frames, so MAE force_mag FAILS +# the invariant with a 2x factor when frames=2. This is a pre-existing +# frame-normalization artifact independent of ghost-atom masking. Ghost atoms +# are correctly excluded (mask_mag=False there). MAE artifact reported as +# NEEDS_CONTEXT in the audit report. +# --------------------------------------------------------------------------- + + +class TestPTEnerSpinLossForceMagMSEGradAccum: + """force_mag MSE meets the grad-accum invariant when ghost atoms are non-magnetic. + + With equal magnetic-atom counts (NM) across frames and ghost atoms having + mask_mag=False, the fancy-index selection excludes padding and the mean() + over [nf, NM, 3] satisfies the invariant. + Expected GREEN immediately (no fix needed for MSE force_mag). + """ + + def _make_loss(self): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + ) + + def test_mse_grad_accum(self): + """force_mag MSE (NM equal per frame, ghost atoms non-magnetic) meets invariant.""" + # Only magnetic-atom slots (first NM_PT) have non-zero values; others zero. + fm_A = _rnd_t(_NM_PT, 3) + fm_A_hat = _rnd_t(_NM_PT, 3) + fm_B = _rnd_t(_NM_PT, 3) + fm_B_hat = _rnd_t(_NM_PT, 3) + + def make_A(): + fm_A_full = torch.zeros(NA, 3, dtype=torch.float64, device="cpu") + fm_A_full[:_NM_PT] = fm_A + fm_A_hat_full = torch.zeros(NA, 3, dtype=torch.float64, device="cpu") + fm_A_hat_full[:_NM_PT] = fm_A_hat + mp = { + "force_mag": fm_A_full.unsqueeze(0), # [1, NA, 3] + "mask_mag": _MASK_MAG_A_PT, + } + lb = {"force_mag": fm_A_hat_full.unsqueeze(0), "find_force_mag": 1.0} + return mp, lb, NA + + def make_B(): + fm_B_full = torch.zeros(NB, 3, dtype=torch.float64, device="cpu") + fm_B_full[:_NM_PT] = fm_B + fm_B_hat_full = torch.zeros(NB, 3, dtype=torch.float64, device="cpu") + fm_B_hat_full[:_NM_PT] = fm_B_hat + mp = { + "force_mag": fm_B_full.unsqueeze(0), # [1, NB, 3] + "mask_mag": _MASK_MAG_B_PT, + } + lb = {"force_mag": fm_B_hat_full.unsqueeze(0), "find_force_mag": 1.0} + return mp, lb, NB + + def make_padded(): + fm_A_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_A_pad[:_NM_PT] = fm_A + fm_A_hat_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_A_hat_pad[:_NM_PT] = fm_A_hat + fm_B_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_B_pad[:_NM_PT] = fm_B + fm_B_hat_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_B_hat_pad[:_NM_PT] = fm_B_hat + mp = { + "force_mag": torch.stack([fm_A_pad, fm_B_pad], dim=0), # [2, NP, 3] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "force_mag": torch.stack([fm_A_hat_pad, fm_B_hat_pad], dim=0), + "find_force_mag": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(self._make_loss(), mp, lb, na), + make_A, + make_B, + make_padded, + ) From 4990db22779f9aebdaeb345f2fc6d9f9da16b47e Mon Sep 17 00:00:00 2001 From: Han Wang Date: Mon, 6 Jul 2026 00:12:15 +0800 Subject: [PATCH 10/16] test(loss): pin cpu device in the property torch-path test Importing deepmd.pt installs a cuda:9999999 default-device trap, so the property torch test failed only when run after the pt test file (order-dependent). Pin device=cpu on the constructed tensors to make the test hermetic. Source behavior unchanged. --- source/tests/common/dpmodel/test_loss_padding.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index ecc96410bc..177b532bb2 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -1236,16 +1236,20 @@ def test_torch_backend_matches_numpy(self): import pytest torch = pytest.importorskip("torch") + # importing deepmd.pt sets a cuda:9999999 default-device trap; pin cpu. + dev = "cpu" p = _rnd(2, PROP_TASK_DIM) lab = _rnd(2, PROP_TASK_DIM) loss_obj = self._make_loss("mse") np_pred = {PROP_VAR: p, "mask": _MASK_PAD} pt_pred = { - PROP_VAR: torch.tensor(p), - "mask": torch.tensor(_MASK_PAD), + PROP_VAR: torch.tensor(p, device=dev), + "mask": torch.tensor(_MASK_PAD, device=dev), } loss_np, _ = loss_obj.call(1.0, NP, np_pred, {PROP_VAR: lab}) - loss_pt, _ = loss_obj.call(1.0, NP, pt_pred, {PROP_VAR: torch.tensor(lab)}) + loss_pt, _ = loss_obj.call( + 1.0, NP, pt_pred, {PROP_VAR: torch.tensor(lab, device=dev)} + ) assert np.isclose(float(loss_np), float(loss_pt)), ( f"torch path must match numpy: {float(loss_np)} vs {float(loss_pt)}" ) From e7e617f08a6a723ed9e39a90af5a227ed7cf32a6 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 9 Jul 2026 21:27:17 +0800 Subject: [PATCH 11/16] fix(loss): drop dead None stores flagged by CodeQL The inv/_nf/_nloc locals are only read inside if maskf is not None guards, so the = None fallbacks in the else branch were dead stores (their None value is never read). CodeQL flagged all three per file as unused. Remove them and note why the names can stay unset when maskf is None. Pure no-op: no runtime behavior change. --- deepmd/dpmodel/loss/ener.py | 5 ++--- deepmd/dpmodel/loss/ener_spin.py | 5 ++--- deepmd/pt/loss/ener.py | 5 ++--- deepmd/pt/loss/ener_spin.py | 5 ++--- 4 files changed, 8 insertions(+), 12 deletions(-) diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py index 8cbfa5ccff..5ee4916762 100644 --- a/deepmd/dpmodel/loss/ener.py +++ b/deepmd/dpmodel/loss/ener.py @@ -232,10 +232,9 @@ def call( _nf = maskf.shape[0] _nloc = maskf.shape[1] else: + # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, + # so leaving them unset here is safe (and avoids dead-store warnings). maskf = None - inv = None - _nf = None - _nloc = None if self.enable_atom_ener_coeff: # when ener_coeff (\nu) is defined, the energy is defined as diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index dd5ebd3460..908224c0b0 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -141,10 +141,9 @@ def call( _nf = maskf.shape[0] _nloc = maskf.shape[1] else: + # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, + # so leaving them unset here is safe (and avoids dead-store warnings). maskf = None - inv = None - _nf = None - _nloc = None if self.has_e: energy_pred = model_dict["energy"] diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py index 4ae64ea8ae..93f57618d6 100644 --- a/deepmd/pt/loss/ener.py +++ b/deepmd/pt/loss/ener.py @@ -248,10 +248,9 @@ def forward( _nf = maskf.shape[0] _nloc = maskf.shape[1] else: + # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, + # so leaving them unset here is safe (and avoids dead-store warnings). maskf = None - inv = None - _nf = None - _nloc = None # Normalization exponent controls loss scaling with system size: # - norm_exp=2 (intensive_ener_virial=True): loss uses 1/N² scaling, making it independent of system size # - norm_exp=1 (intensive_ener_virial=False, legacy): loss uses 1/N scaling, which varies with system size diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 7593c7c246..1f814651e8 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -169,10 +169,9 @@ def forward( _nf = maskf.shape[0] _nloc = maskf.shape[1] else: + # inv, _nf, _nloc are only read inside ``if maskf is not None`` guards, + # so leaving them unset here is safe (and avoids dead-store warnings). maskf = None - inv = None - _nf = None - _nloc = None if self.has_e and "energy" in model_pred and "energy" in label: energy_pred = model_pred["energy"] From 68aa05c6c8ad5d5a2111b7f1cbb6feda8ba1b780 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Thu, 9 Jul 2026 21:27:19 +0800 Subject: [PATCH 12/16] test(loss): xfail-track force_mag MAE frame-normalization debt Both reviewers asked for the known force_mag MAE grad-accum discrepancy (sum over frames -> 2x factor for nf=2) to be tracked by CI rather than only documented in a comment. Add strict xfail tests in the pt and dpmodel loss-padding suites; strict mode self-heals (XPASS -> failure) once force_mag MAE is switched to a frame-wise mean. --- .../tests/common/dpmodel/test_loss_padding.py | 116 ++++++++++++++++++ source/tests/pt/test_loss_padding.py | 92 ++++++++++++++ 2 files changed, 208 insertions(+) diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 177b532bb2..78aec1cc41 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -14,6 +14,7 @@ """ import numpy as np +import pytest from deepmd.dpmodel.loss.dos import ( DOSLoss, @@ -2124,3 +2125,118 @@ def make_padded(): make_B, make_padded, ) + + +class TestDPModelEnerSpinLossForceMagMAEGradAccum: + """force_mag MAE (dpmodel) FAILS the grad-accum invariant by a known 2x factor. + + Unlike the MSE path, ``force_mag`` MAE reduces with ``xp.sum`` over frames + instead of a frame-wise mean, so a padded ``[A+B]`` batch (nf=2) yields twice + the mean-of-frames reference. This is a pre-existing frame-normalization + artifact, NOT a ghost-atom padding bug (padding is correctly excluded via + ``mask_mag``); see the module NOTE above. The test is ``xfail(strict=True)`` + so the debt is tracked in CI and self-heals: if ``force_mag`` MAE is switched + to a frame-wise mean, this test XPASSes and strict mode turns that into a + failure that flags the marker for removal. + """ + + def _make_loss(self): + return EnergySpinLossDPModel( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + + def _loss_fn(self, loss_obj, model_pred, label, natoms): + loss, _ = loss_obj.call(1.0, natoms, model_pred, label) + return float(loss) + + @pytest.mark.xfail( + strict=True, + reason="force_mag MAE sums over frames (2x factor for nf=2); " + "pre-existing frame-normalization inconsistency tracked for follow-up.", + ) + def test_mae_grad_accum(self): + """force_mag MAE violates the frame-average invariant (documented follow-up).""" + fm_A = _rnd(NA, 3) + fm_A_hat = _rnd(NA, 3) + fm_B = _rnd(NB, 3) + fm_B_hat = _rnd(NB, 3) + + def make_A(): + fm_A_full = np.zeros((NA, 3), dtype=np.float64) + fm_A_full[:_NM] = fm_A[:_NM] + fm_A_hat_full = np.zeros((NA, 3), dtype=np.float64) + fm_A_hat_full[:_NM] = fm_A_hat[:_NM] + pred = { + "energy": np.zeros((1, 1), dtype=np.float64), + "force_mag": fm_A_full[None], # [1, NA, 3] + "mask_mag": _MASK_MAG_A, + } + lbl = { + "force_mag": fm_A_hat_full[None], + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NA + + def make_B(): + fm_B_full = np.zeros((NB, 3), dtype=np.float64) + fm_B_full[:_NM] = fm_B[:_NM] + fm_B_hat_full = np.zeros((NB, 3), dtype=np.float64) + fm_B_hat_full[:_NM] = fm_B_hat[:_NM] + pred = { + "energy": np.zeros((1, 1), dtype=np.float64), + "force_mag": fm_B_full[None], # [1, NB, 3] + "mask_mag": _MASK_MAG_B, + } + lbl = { + "force_mag": fm_B_hat_full[None], + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NB + + def make_padded(): + fm_A_pad = np.zeros((NP, 3), dtype=np.float64) + fm_A_pad[:_NM] = fm_A[:_NM] + fm_A_hat_pad = np.zeros((NP, 3), dtype=np.float64) + fm_A_hat_pad[:_NM] = fm_A_hat[:_NM] + fm_B_pad = np.zeros((NP, 3), dtype=np.float64) + fm_B_pad[:_NM] = fm_B[:_NM] + fm_B_hat_pad = np.zeros((NP, 3), dtype=np.float64) + fm_B_hat_pad[:_NM] = fm_B_hat[:_NM] + fm_pad = np.stack([fm_A_pad, fm_B_pad], axis=0) # [2, NP, 3] + fm_hat_pad = np.stack([fm_A_hat_pad, fm_B_hat_pad], axis=0) + pred = { + "energy": np.zeros((2, 1), dtype=np.float64), + "force_mag": fm_pad, + "mask_mag": _MASK_MAG_PAD_SPIN, # ghost atoms have mask_mag=False + "mask": _MASK_PAD_SPIN, + } + lbl = { + "force_mag": fm_hat_pad, + "find_force_mag": 1.0, + "find_energy": 0.0, + "find_force": 0.0, + "find_virial": 0.0, + } + return pred, lbl, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: self._loss_fn(self._make_loss(), mp, lb, na), + make_A, + make_B, + make_padded, + ) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index b92f99ed3a..32e33b6012 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -13,6 +13,7 @@ """ import numpy as np +import pytest import torch from deepmd.pt.loss.dos import ( @@ -2061,3 +2062,94 @@ def make_padded(): make_B, make_padded, ) + + +class TestPTEnerSpinLossForceMagMAEGradAccum: + """force_mag MAE (pt) FAILS the grad-accum invariant by a known 2x factor. + + Unlike the MSE path, ``force_mag`` MAE reduces with ``.sum()`` over frames + instead of a frame-wise mean, so a padded ``[A+B]`` batch (nf=2) yields twice + the mean-of-frames reference. This is a pre-existing frame-normalization + artifact, NOT a ghost-atom padding bug (padding is correctly excluded via + ``mask_mag``); see the module NOTE above. The test is ``xfail(strict=True)`` + so the debt is tracked in CI and self-heals: if ``force_mag`` MAE is switched + to a frame-wise mean, this test XPASSes and strict mode turns that into a + failure that flags the marker for removal. + """ + + def _make_loss(self): + return EnergySpinLossPT( + starter_learning_rate=1.0, + start_pref_e=0.0, + limit_pref_e=0.0, + start_pref_fr=0.0, + limit_pref_fr=0.0, + start_pref_fm=1.0, + limit_pref_fm=1.0, + start_pref_v=0.0, + limit_pref_v=0.0, + loss_func="mae", + ) + + @pytest.mark.xfail( + strict=True, + reason="force_mag MAE sums over frames (2x factor for nf=2); " + "pre-existing frame-normalization inconsistency tracked for follow-up.", + ) + def test_mae_grad_accum(self): + """force_mag MAE violates the frame-average invariant (documented follow-up).""" + fm_A = _rnd_t(_NM_PT, 3) + fm_A_hat = _rnd_t(_NM_PT, 3) + fm_B = _rnd_t(_NM_PT, 3) + fm_B_hat = _rnd_t(_NM_PT, 3) + + def make_A(): + fm_A_full = torch.zeros(NA, 3, dtype=torch.float64, device="cpu") + fm_A_full[:_NM_PT] = fm_A + fm_A_hat_full = torch.zeros(NA, 3, dtype=torch.float64, device="cpu") + fm_A_hat_full[:_NM_PT] = fm_A_hat + mp = { + "force_mag": fm_A_full.unsqueeze(0), # [1, NA, 3] + "mask_mag": _MASK_MAG_A_PT, + } + lb = {"force_mag": fm_A_hat_full.unsqueeze(0), "find_force_mag": 1.0} + return mp, lb, NA + + def make_B(): + fm_B_full = torch.zeros(NB, 3, dtype=torch.float64, device="cpu") + fm_B_full[:_NM_PT] = fm_B + fm_B_hat_full = torch.zeros(NB, 3, dtype=torch.float64, device="cpu") + fm_B_hat_full[:_NM_PT] = fm_B_hat + mp = { + "force_mag": fm_B_full.unsqueeze(0), # [1, NB, 3] + "mask_mag": _MASK_MAG_B_PT, + } + lb = {"force_mag": fm_B_hat_full.unsqueeze(0), "find_force_mag": 1.0} + return mp, lb, NB + + def make_padded(): + fm_A_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_A_pad[:_NM_PT] = fm_A + fm_A_hat_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_A_hat_pad[:_NM_PT] = fm_A_hat + fm_B_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_B_pad[:_NM_PT] = fm_B + fm_B_hat_pad = torch.zeros(NP, 3, dtype=torch.float64, device="cpu") + fm_B_hat_pad[:_NM_PT] = fm_B_hat + mp = { + "force_mag": torch.stack([fm_A_pad, fm_B_pad], dim=0), # [2, NP, 3] + "mask_mag": _MASK_MAG_PAD_SPIN_PT, + "mask": _MASK_PAD_SPIN_PT, + } + lb = { + "force_mag": torch.stack([fm_A_hat_pad, fm_B_hat_pad], dim=0), + "find_force_mag": 1.0, + } + return mp, lb, NP + + assert_grad_accum_invariant( + lambda mp, lb, na: _spin_loss_fn(self._make_loss(), mp, lb, na), + make_A, + make_B, + make_padded, + ) From b143f982c44cef572ca6388c68b71f6ffb38c8c2 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 00:39:21 +0800 Subject: [PATCH 13/16] docs(loss): reference tracking issues for mixed_type padding follow-ups Record the deferred scope boundary of the mixed_type padding fix in the loss-padding test modules instead of only in the PR body: the TF backend loss (deepmodeling/deepmd-kit#5760) and the pt-only dens/population/denoise losses (deepmodeling/deepmd-kit#5761). The force_mag MAE frame-normalization debt stays tracked by its strict-xfail test. --- source/tests/common/dpmodel/test_loss_padding.py | 9 +++++++++ source/tests/pt/test_loss_padding.py | 9 +++++++++ 2 files changed, 18 insertions(+) diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 78aec1cc41..8b22e46de3 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -6,6 +6,15 @@ The dpmodel losses accept numpy arrays (via the array_api_compat backend). +Scope / follow-ups (mixed_type padding fix, PR #5738) +---------------------------------------------------- +- The TF backend loss is not covered here and still has the mixed_type + dilution behavior; tracked in deepmodeling/deepmd-kit#5760. +- The pt-only losses ``dens``/``population``/``denoise`` are out of scope; + tracked in deepmodeling/deepmd-kit#5761. +- ``ener_spin``'s ``force_mag`` MAE frame-normalization debt is tracked by the + ``xfail(strict=True)`` test below (self-heals once fixed). + Constants --------- NA = 3 # real atoms in the short frame diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 32e33b6012..1704bfc9e9 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -10,6 +10,15 @@ assert_grad_accum_invariant -- reusable by Tasks 2-5 to check the grad-accumulation invariant: loss on a padded multi-frame batch must equal mean_over_frames(per_frame_loss). + +Scope / follow-ups (mixed_type padding fix, PR #5738) +---------------------------------------------------- +- The TF backend loss is not covered here and still has the mixed_type + dilution behavior; tracked in deepmodeling/deepmd-kit#5760. +- The pt-only losses ``dens``/``population``/``denoise`` are out of scope; + tracked in deepmodeling/deepmd-kit#5761. +- ``ener_spin``'s ``force_mag`` MAE frame-normalization debt is tracked by the + ``xfail(strict=True)`` test below (self-heals once fixed). """ import numpy as np From b578129c4d57f76b2a278a32a78e3bd457d06e0b Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sat, 11 Jul 2026 12:40:34 +0800 Subject: [PATCH 14/16] ci: re-trigger readthedocs build (prior build reaped by RTD resource limit) From 24fe23548cb26db9866f14e0fe051ef00c654ef9 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 12 Jul 2026 13:52:07 +0800 Subject: [PATCH 15/16] test(pt): run PropertyLoss padding tests on env.DEVICE The mixed_type padding tests for PropertyLoss created their input tensors and masks on CPU, but PropertyLoss.forward builds out_std/out_bias and the loss accumulator on env.DEVICE and combines them with the multi-dim label/pred tensors (and divides pred/label by real_natoms derived from the mask). On CPU CI env.DEVICE is cpu so this passed, but on the CUDA merge-queue runner it raised "Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu". The other loss terms only cross devices at 0-dim scalar reductions, which torch auto-promotes, so only the property tests were affected. Move the property test inputs and their masks onto env.DEVICE. --- source/tests/pt/test_loss_padding.py | 61 ++++++++++++++++++---------- 1 file changed, 40 insertions(+), 21 deletions(-) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index 1704bfc9e9..e55d58d58f 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -41,6 +41,9 @@ from deepmd.pt.loss.tensor import ( TensorLoss, ) +from deepmd.pt.utils import ( + env, +) # --------------------------------------------------------------------------- # Constants used by the multi-frame test harness (Tasks 2-5) @@ -1212,6 +1215,22 @@ def test_no_op_for_non_mixed(self): PROP_VAR = "test_prop" +# Unlike the other losses, PropertyLoss.forward builds out_std/out_bias (and the +# loss accumulator) on env.DEVICE and combines them with the *multi-dim* +# label/pred tensors (and divides pred/label by real_natoms from the mask). A +# CPU-input / env.DEVICE-buffer mix therefore raises on GPU CI (a 0-dim CPU +# scalar is auto-promoted, but a [nf, task_dim] tensor is not), so the property +# inputs and their masks must live on env.DEVICE. See #5738. +def _rnd_t_prop(*shape): + """Random float64 tensor on env.DEVICE (PropertyLoss holds device-resident buffers).""" + return torch.tensor( + RNG.standard_normal(shape), dtype=torch.float64, device=env.DEVICE + ) + + +_MASK_PAD_PROP = _MASK_PAD_PT.to(env.DEVICE) + + class TestPTPropertyLossExtensiveGradAccum: """Idiom 2 (per-frame real-natoms normalization) for extensive pt PropertyLoss. @@ -1248,7 +1267,7 @@ def make_A(): return ( { PROP_VAR: p_A, - "mask": torch.ones(1, NA, dtype=torch.float64, device="cpu"), + "mask": torch.ones(1, NA, dtype=torch.float64, device=env.DEVICE), }, {PROP_VAR: l_A.clone()}, NA, @@ -1258,7 +1277,7 @@ def make_B(): return ( { PROP_VAR: p_B, - "mask": torch.ones(1, NB, dtype=torch.float64, device="cpu"), + "mask": torch.ones(1, NB, dtype=torch.float64, device=env.DEVICE), }, {PROP_VAR: l_B.clone()}, NB, @@ -1268,7 +1287,7 @@ def make_padded(): return ( { PROP_VAR: torch.cat([p_A, p_B], dim=0), # [2, task_dim] - "mask": _MASK_PAD_PT, + "mask": _MASK_PAD_PROP, }, {PROP_VAR: torch.cat([l_A, l_B], dim=0)}, NP, @@ -1283,36 +1302,36 @@ def make_padded(): def test_mse_grad_accum(self): """MSE extensive property meets the grad-accum invariant.""" - p_A = _rnd_t(1, PROP_TASK_DIM) - l_A = _rnd_t(1, PROP_TASK_DIM) - p_B = _rnd_t(1, PROP_TASK_DIM) - l_B = _rnd_t(1, PROP_TASK_DIM) + p_A = _rnd_t_prop(1, PROP_TASK_DIM) + l_A = _rnd_t_prop(1, PROP_TASK_DIM) + p_B = _rnd_t_prop(1, PROP_TASK_DIM) + l_B = _rnd_t_prop(1, PROP_TASK_DIM) self._run_invariant(self._make_loss("mse"), p_A, l_A, p_B, l_B) def test_mae_grad_accum(self): """MAE extensive property meets the grad-accum invariant.""" - p_A = _rnd_t(1, PROP_TASK_DIM) - l_A = _rnd_t(1, PROP_TASK_DIM) - p_B = _rnd_t(1, PROP_TASK_DIM) - l_B = _rnd_t(1, PROP_TASK_DIM) + p_A = _rnd_t_prop(1, PROP_TASK_DIM) + l_A = _rnd_t_prop(1, PROP_TASK_DIM) + p_B = _rnd_t_prop(1, PROP_TASK_DIM) + l_B = _rnd_t_prop(1, PROP_TASK_DIM) self._run_invariant(self._make_loss("mae"), p_A, l_A, p_B, l_B) def test_smooth_mae_grad_accum(self): """smooth_mae extensive property meets the grad-accum invariant.""" - p_A = _rnd_t(1, PROP_TASK_DIM) - l_A = _rnd_t(1, PROP_TASK_DIM) - p_B = _rnd_t(1, PROP_TASK_DIM) - l_B = _rnd_t(1, PROP_TASK_DIM) + p_A = _rnd_t_prop(1, PROP_TASK_DIM) + l_A = _rnd_t_prop(1, PROP_TASK_DIM) + p_B = _rnd_t_prop(1, PROP_TASK_DIM) + l_B = _rnd_t_prop(1, PROP_TASK_DIM) self._run_invariant(self._make_loss("smooth_mae"), p_A, l_A, p_B, l_B) def test_no_op_for_non_mixed(self): """All-ones mask gives same extensive-property loss as no mask.""" - p = _rnd_t(2, PROP_TASK_DIM) - l = _rnd_t(2, PROP_TASK_DIM) + p = _rnd_t_prop(2, PROP_TASK_DIM) + l = _rnd_t_prop(2, PROP_TASK_DIM) loss_obj = self._make_loss("mse") mp_mask = { PROP_VAR: p, - "mask": torch.ones(2, NB, dtype=torch.float64, device="cpu"), + "mask": torch.ones(2, NB, dtype=torch.float64, device=env.DEVICE), } mp_nm = {PROP_VAR: p} # pt forward mutates label[var_name]; use separate dicts for each call. @@ -1350,10 +1369,10 @@ def _loss_fn(self, loss_obj, model_pred, label, natoms): def test_intensive_ignores_mask(self): """Intensive property: masked batch == unmasked batch.""" - p = _rnd_t(2, PROP_TASK_DIM) - l = _rnd_t(2, PROP_TASK_DIM) + p = _rnd_t_prop(2, PROP_TASK_DIM) + l = _rnd_t_prop(2, PROP_TASK_DIM) loss_obj = self._make_loss() - mp_mask = {PROP_VAR: p, "mask": _MASK_PAD_PT} + mp_mask = {PROP_VAR: p, "mask": _MASK_PAD_PROP} mp_nm = {PROP_VAR: p} # Use separate label dicts since pt forward mutates label[var_name]. lb_m = {PROP_VAR: l.clone()} From 374e24a3403dbe5dc6d4e119f099879e80589811 Mon Sep 17 00:00:00 2001 From: Han Wang Date: Sun, 12 Jul 2026 15:02:00 +0800 Subject: [PATCH 16/16] fix(loss): make ener_spin force_mag MAE batch-size independent force_mag MAE reduced with a sum over frames (pt: .sum(-1).mean(-1).sum(); dpmodel: xp.sum(per_frame_sum / per_frame_count)), so the loss scaled with batch size: a 2-frame batch gave twice the mean of the two single-frame losses, unlike force_mag MSE and force_real MAE which use a plain mean. Reduce force_mag MAE with a mean over frames, magnetic atoms and xyz (matching force_mag MSE, force_real MAE and the displayed mae_fm metric) so a 2-frame batch equals the mean of the two single-frame losses. Applied to both the pt and dpmodel backends (dpmodel is re-exported by pt_expt); TF and paddle have no spin loss. Drop the now-obsolete xfail(strict=True) markers on the force_mag MAE grad-accum tests, which pass after the fix. --- deepmd/dpmodel/loss/ener_spin.py | 14 +++---- deepmd/pt/loss/ener_spin.py | 6 ++- .../tests/common/dpmodel/test_loss_padding.py | 37 +++++++------------ source/tests/pt/test_loss_padding.py | 29 +++++---------- 4 files changed, 34 insertions(+), 52 deletions(-) diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py index 908224c0b0..e39e0a6a9e 100644 --- a/deepmd/dpmodel/loss/ener_spin.py +++ b/deepmd/dpmodel/loss/ener_spin.py @@ -279,16 +279,14 @@ def call( more_loss["mae_fm"] = self.display_if_exist(mae_fm, find_force_mag) elif self.loss_func == "mae": abs_diff_fm = xp.abs(diff_fm) # [nf, na, 3], zeros for non-magnetic - per_atom_fm = xp.sum(abs_diff_fm, axis=-1) # [nf, na] - mask_2d = mask_float[:, :, 0] # [nf, na] - per_frame_sum_fm = xp.sum(per_atom_fm, axis=-1) # [nf] - per_frame_count_fm = xp.sum(mask_2d, axis=-1) # [nf] - l1_force_mag_loss = xp.sum( - per_frame_sum_fm / per_frame_count_fm - ) # scalar + # Mean over frames, magnetic atoms and xyz (same reduction as + # force_mag MSE, force_real MAE and the displayed mae_fm) so the + # loss is batch-size independent: a 2-frame batch equals the mean + # of the two single-frame losses. + l1_force_mag_loss = xp.sum(abs_diff_fm) / (n_valid * 3) loss += pref_fm * l1_force_mag_loss more_loss["mae_fm"] = self.display_if_exist( - xp.sum(abs_diff_fm) / (n_valid * 3), find_force_mag + l1_force_mag_loss, find_force_mag ) if self.has_ae: diff --git a/deepmd/pt/loss/ener_spin.py b/deepmd/pt/loss/ener_spin.py index 1f814651e8..2bc2601276 100644 --- a/deepmd/pt/loss/ener_spin.py +++ b/deepmd/pt/loss/ener_spin.py @@ -368,7 +368,11 @@ def forward( more_loss["mae_fm"] = self.display_if_exist( l1_force_mag_loss.mean().detach(), find_force_m ) - l1_force_mag_loss = l1_force_mag_loss.sum(-1).mean(-1).sum() + # Mean over frames, magnetic atoms and xyz (same reduction as + # force_mag MSE, force_real MAE and the displayed mae_fm) so the + # loss is batch-size independent: a 2-frame batch equals the mean + # of the two single-frame losses. + l1_force_mag_loss = l1_force_mag_loss.mean() loss += (pref_fm * torch.nan_to_num(l1_force_mag_loss)).to( GLOBAL_PT_FLOAT_PRECISION ) diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py index 8b22e46de3..1109c21131 100644 --- a/source/tests/common/dpmodel/test_loss_padding.py +++ b/source/tests/common/dpmodel/test_loss_padding.py @@ -12,8 +12,9 @@ dilution behavior; tracked in deepmodeling/deepmd-kit#5760. - The pt-only losses ``dens``/``population``/``denoise`` are out of scope; tracked in deepmodeling/deepmd-kit#5761. -- ``ener_spin``'s ``force_mag`` MAE frame-normalization debt is tracked by the - ``xfail(strict=True)`` test below (self-heals once fixed). +- ``ener_spin``'s ``force_mag`` MAE now uses a batch-size-independent mean + reduction (frames/atoms/xyz), consistent with force_mag MSE and force_real + MAE; the grad-accum invariant is asserted by the test below. Constants --------- @@ -23,7 +24,6 @@ """ import numpy as np -import pytest from deepmd.dpmodel.loss.dos import ( DOSLoss, @@ -2021,11 +2021,9 @@ def test_mse_grad_accum(self): # --------------------------------------------------------------------------- # force_mag MSE grad-accum invariant (NM equal across frames, ghost atoms non-magnetic) # -# NOTE: force_mag MAE (dpmodel) uses xp.sum over frames instead of xp.mean, -# so MAE force_mag FAILS the invariant with a 2x factor when frames=2. -# This is a pre-existing frame-normalization artifact independent of ghost-atom -# masking. Ghost atoms are correctly excluded (mask_mag=0 there). The MAE -# artifact is reported as NEEDS_CONTEXT in the audit report. +# NOTE: force_mag MAE (dpmodel) uses a plain mean over frames/atoms/xyz (like +# force_mag MSE and force_real MAE), so it meets the grad-accum invariant. +# Ghost atoms are correctly excluded (mask_mag=0 there). # --------------------------------------------------------------------------- @@ -2137,16 +2135,12 @@ def make_padded(): class TestDPModelEnerSpinLossForceMagMAEGradAccum: - """force_mag MAE (dpmodel) FAILS the grad-accum invariant by a known 2x factor. - - Unlike the MSE path, ``force_mag`` MAE reduces with ``xp.sum`` over frames - instead of a frame-wise mean, so a padded ``[A+B]`` batch (nf=2) yields twice - the mean-of-frames reference. This is a pre-existing frame-normalization - artifact, NOT a ghost-atom padding bug (padding is correctly excluded via - ``mask_mag``); see the module NOTE above. The test is ``xfail(strict=True)`` - so the debt is tracked in CI and self-heals: if ``force_mag`` MAE is switched - to a frame-wise mean, this test XPASSes and strict mode turns that into a - failure that flags the marker for removal. + """force_mag MAE (dpmodel) meets the grad-accum invariant. + + ``force_mag`` MAE now reduces with a plain mean over frames, magnetic atoms + and xyz (same reduction as force_mag MSE and force_real MAE), so a padded + ``[A+B]`` batch (nf=2) equals the mean of the two single-frame losses and the + loss is batch-size independent. """ def _make_loss(self): @@ -2167,13 +2161,8 @@ def _loss_fn(self, loss_obj, model_pred, label, natoms): loss, _ = loss_obj.call(1.0, natoms, model_pred, label) return float(loss) - @pytest.mark.xfail( - strict=True, - reason="force_mag MAE sums over frames (2x factor for nf=2); " - "pre-existing frame-normalization inconsistency tracked for follow-up.", - ) def test_mae_grad_accum(self): - """force_mag MAE violates the frame-average invariant (documented follow-up).""" + """force_mag MAE meets the frame-average grad-accum invariant.""" fm_A = _rnd(NA, 3) fm_A_hat = _rnd(NA, 3) fm_B = _rnd(NB, 3) diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py index e55d58d58f..094bcfeadb 100644 --- a/source/tests/pt/test_loss_padding.py +++ b/source/tests/pt/test_loss_padding.py @@ -17,12 +17,12 @@ dilution behavior; tracked in deepmodeling/deepmd-kit#5760. - The pt-only losses ``dens``/``population``/``denoise`` are out of scope; tracked in deepmodeling/deepmd-kit#5761. -- ``ener_spin``'s ``force_mag`` MAE frame-normalization debt is tracked by the - ``xfail(strict=True)`` test below (self-heals once fixed). +- ``ener_spin``'s ``force_mag`` MAE now uses a batch-size-independent mean + reduction (frames/atoms/xyz), consistent with force_mag MSE and force_real + MAE; the grad-accum invariant is asserted by the test below. """ import numpy as np -import pytest import torch from deepmd.pt.loss.dos import ( @@ -2093,16 +2093,12 @@ def make_padded(): class TestPTEnerSpinLossForceMagMAEGradAccum: - """force_mag MAE (pt) FAILS the grad-accum invariant by a known 2x factor. - - Unlike the MSE path, ``force_mag`` MAE reduces with ``.sum()`` over frames - instead of a frame-wise mean, so a padded ``[A+B]`` batch (nf=2) yields twice - the mean-of-frames reference. This is a pre-existing frame-normalization - artifact, NOT a ghost-atom padding bug (padding is correctly excluded via - ``mask_mag``); see the module NOTE above. The test is ``xfail(strict=True)`` - so the debt is tracked in CI and self-heals: if ``force_mag`` MAE is switched - to a frame-wise mean, this test XPASSes and strict mode turns that into a - failure that flags the marker for removal. + """force_mag MAE (pt) meets the grad-accum invariant. + + ``force_mag`` MAE now reduces with a plain mean over frames, magnetic atoms + and xyz (same reduction as force_mag MSE and force_real MAE), so a padded + ``[A+B]`` batch (nf=2) equals the mean of the two single-frame losses and the + loss is batch-size independent. """ def _make_loss(self): @@ -2119,13 +2115,8 @@ def _make_loss(self): loss_func="mae", ) - @pytest.mark.xfail( - strict=True, - reason="force_mag MAE sums over frames (2x factor for nf=2); " - "pre-existing frame-normalization inconsistency tracked for follow-up.", - ) def test_mae_grad_accum(self): - """force_mag MAE violates the frame-average invariant (documented follow-up).""" + """force_mag MAE meets the frame-average grad-accum invariant.""" fm_A = _rnd_t(_NM_PT, 3) fm_A_hat = _rnd_t(_NM_PT, 3) fm_B = _rnd_t(_NM_PT, 3)