diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py
index 0db1c51603..7f60aaa200 100644
--- a/deepmd/dpmodel/array_api.py
+++ b/deepmd/dpmodel/array_api.py
@@ -11,6 +11,10 @@
     Version,
 )
 
+from deepmd.dpmodel.common import (
+    to_numpy_array,
+)
+
 # Type alias for array_api compatible arrays
 Array = np.ndarray | Any  # Any to support JAX, PyTorch, etc. arrays
 
@@ -27,22 +31,28 @@ def xp_asarray_nodetach(
     ``torch.asarray`` detaches its input from the autograd graph, so calling
     ``xp.asarray`` on a weight attribute that is already a backend tensor
     (e.g. a ``torch.nn.Parameter`` registered by the pt_expt backend)
-    silently breaks gradient flow to that weight.  This helper converts
-    genuine non-backend data (numpy arrays, python scalars/lists) via
-    ``xp.asarray``; backend tensors are returned as-is, with an optional
-    differentiable dtype cast via ``xp.astype``.
-
-    The ``device`` argument only applies to the conversion path: backend
-    tensors are assumed to already live on the working device (they are
-    created together with the inputs).
+    silently breaks gradient flow to that weight. Backend tensors already in
+    ``xp`` are therefore returned as-is, with an optional differentiable dtype
+    cast via ``xp.astype``.
+
+    An array from another namespace cannot retain its autograd graph. It is
+    converted through NumPy before entering ``xp``; this also performs the
+    required device-to-host copy when a CUDA-backed model constant is consumed
+    by a NumPy statistics path.
+
+    The ``device`` argument only applies to the conversion path. Arrays already
+    in ``xp`` are assumed to live on the working device because model buffers
+    and inputs are moved together.
     """
-    if isinstance(obj, np.ndarray) or not array_api_compat.is_array_api_obj(obj):
-        if dtype is None:
-            return xp.asarray(obj, device=device)
-        return xp.asarray(obj, dtype=dtype, device=device)
-    if dtype is not None and obj.dtype != dtype:
-        obj = xp.astype(obj, dtype)
-    return obj
+    if array_api_compat.is_array_api_obj(obj):
+        if array_api_compat.array_namespace(obj) is xp:
+            if dtype is not None and obj.dtype != dtype:
+                obj = xp.astype(obj, dtype)
+            return obj
+        obj = to_numpy_array(obj)
+    if dtype is None:
+        return xp.asarray(obj, device=device)
+    return xp.asarray(obj, dtype=dtype, device=device)
 
 
 # array api adds take_along_axis in https://github.com/data-apis/array-api/pull/816
diff --git a/deepmd/dpmodel/descriptor/dpa4.py b/deepmd/dpmodel/descriptor/dpa4.py
index f1a08123c4..d96e23f84c 100644
--- a/deepmd/dpmodel/descriptor/dpa4.py
+++ b/deepmd/dpmodel/descriptor/dpa4.py
@@ -2186,7 +2186,7 @@ def _canonicalize_charge_spin(
                 raise ValueError("`charge_spin` is required for this SeZM descriptor.")
             charge_spin = xp.reshape(
                 xp_asarray_nodetach(
-                    xp, np.asarray(self.default_chg_spin), dtype=dtype, device=device
+                    xp, self.default_chg_spin, dtype=dtype, device=device
                 ),
                 (1, 2),
             )
diff --git a/deepmd/dpmodel/loss/dos.py b/deepmd/dpmodel/loss/dos.py
index c943ff2aa4..45229640cc 100644
--- a/deepmd/dpmodel/loss/dos.py
+++ b/deepmd/dpmodel/loss/dos.py
@@ -159,7 +159,7 @@ def call(
             )
             diff3d = local_pred - local_label  # [nf, natoms, numb_dos]
             if "mask" in model_dict:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=numb_dos).
                 maskf = xp.astype(model_dict["mask"], diff3d.dtype)  # [nf, natoms]
                 l2_local_loss_dos = masked_atom_mean(
                     xp.square(diff3d), maskf, self.numb_dos
@@ -184,7 +184,7 @@ def call(
             )
             diff3d = local_pred_cdf - local_label_cdf  # [nf, natoms, numb_dos]
             if "mask" in model_dict:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=numb_dos).
                 maskf = xp.astype(model_dict["mask"], diff3d.dtype)  # [nf, natoms]
                 l2_local_loss_cdf = masked_atom_mean(
                     xp.square(diff3d), maskf, self.numb_dos
diff --git a/deepmd/dpmodel/loss/ener.py b/deepmd/dpmodel/loss/ener.py
index 34f7858b83..8459d0847f 100644
--- a/deepmd/dpmodel/loss/ener.py
+++ b/deepmd/dpmodel/loss/ener.py
@@ -16,6 +16,12 @@
     masked_pair_mean,
     per_frame_component_mean,
 )
+from deepmd.dpmodel.utils.neighbor_graph.graph import (
+    frame_id_from_n_node,
+)
+from deepmd.dpmodel.utils.neighbor_graph.segment import (
+    segment_sum,
+)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
@@ -234,6 +240,11 @@ def __init__(
         if self.use_huber and self.has_h:
             raise RuntimeError("Huber loss is not implemented for hessian.")
 
+    @property
+    def supports_ragged_batches(self) -> bool:
+        """Whether the configured terms accept a flat per-node batch axis."""
+        return not (self.has_gf or self.has_h)
+
     def call(
         self,
         learning_rate: float,
@@ -253,45 +264,78 @@ def call(
         including all enabled prefactors and any configured Huber terms.
         """
         energy = model_dict["energy"]
-        force = model_dict["force"]
-        virial = model_dict["virial"]
-        atom_ener = model_dict["atom_energy"]
-        energy_hat = label_dict["energy"]
-        force_hat = label_dict["force"]
-        virial_hat = label_dict["virial"]
-        atom_ener_hat = label_dict["atom_ener"]
-        atom_pref = label_dict["atom_pref"]
-        find_energy = label_dict["find_energy"]
-        find_force = label_dict["find_force"]
-        find_virial = label_dict["find_virial"]
-        find_atom_ener = label_dict["find_atom_ener"]
-        find_atom_pref = (
-            label_dict["find_atom_pref"] if not self.use_default_pf else 1.0
-        )
-        xp = array_api_compat.array_namespace(
-            energy,
-            force,
-            virial,
-            atom_ener,
-            energy_hat,
-            force_hat,
-            virial_hat,
-            atom_ener_hat,
-            atom_pref,
+        xp = array_api_compat.array_namespace(energy)
+
+        force_required = (
+            self.has_f or self.has_pf or self.relative_f is not None or self.has_gf
         )
+        if self.has_e:
+            energy_hat = label_dict["energy"]
+            find_energy = label_dict["find_energy"]
+        if force_required:
+            force = model_dict["force"]
+            force_hat = label_dict["force"]
+            find_force = label_dict["find_force"]
+        if self.has_v:
+            virial = model_dict["virial"]
+            virial_hat = label_dict["virial"]
+            find_virial = label_dict["find_virial"]
+        if self.has_ae:
+            atom_ener = model_dict["atom_energy"]
+            atom_ener_hat = label_dict["atom_ener"]
+            find_atom_ener = label_dict["find_atom_ener"]
+        if self.has_pf:
+            atom_pref = label_dict["atom_pref"]
+            find_atom_pref = (
+                label_dict["find_atom_pref"] if not self.use_default_pf else 1.0
+            )
 
-        # 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]
+        # Two things about a batch decide how its terms reduce, and the node
+        # axis states them differently.
+        #
+        # ``inv``, the reciprocal included-atom count of each frame, is what
+        # the extensive frame-level terms (energy, virial) divide by.
+        # ``maskf`` marks both padded rows and model-excluded atom types for
+        # the per-atom terms. A rectangular batch carries a two-dimensional
+        # mask whose row sums give the counts. A ragged batch carries the same
+        # information on its flat node axis, with ``n_node`` defining the
+        # frame segments.
+        maskf = (
+            xp.astype(model_dict["mask"], energy.dtype)
+            if "mask" in model_dict
+            else None
+        )
+        is_ragged = "n_node" in model_dict
+        frame_id = None
+        included_n_node = None
+        inv = None
+        if is_ragged:
+            n_node = model_dict["n_node"]
+            _nf = n_node.shape[0]
+            if maskf is None:
+                included_n_node = xp.astype(n_node, energy.dtype)
+            else:
+                frame_id = frame_id_from_n_node(n_node, n_total=maskf.shape[0])
+                included_n_node = segment_sum(xp.reshape(maskf, (-1,)), frame_id, _nf)
+        elif maskf is not None:
+            included_n_node = xp.reshape(xp.sum(maskf, axis=-1), (-1,))
             _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
+        if included_n_node is not None:
+            has_included_node = included_n_node > 0
+            safe_included_n_node = xp.where(
+                has_included_node,
+                included_n_node,
+                xp.ones_like(included_n_node),
+            )
+            inv = xp.where(
+                has_included_node,
+                1.0 / safe_included_n_node,
+                xp.zeros_like(included_n_node),
+            )
+        if maskf is not None:
+            _node_shape = maskf.shape
+        if inv is not None:
+            _nf = inv.shape[0]
 
         if self.enable_atom_ener_coeff:
             # when ener_coeff (\nu) is defined, the energy is defined as
@@ -302,10 +346,17 @@ def call(
             # A + B -> C + D
             # E = - E(A) - E(B) + E(C) + E(D)
             # A, B, C, D could be put far away from each other
+            atom_ener = model_dict["atom_energy"]
             atom_ener_coeff = label_dict["atom_ener_coeff"]
             atom_ener_coeff = xp.reshape(atom_ener_coeff, atom_ener.shape)
-            energy = xp.sum(atom_ener_coeff * atom_ener, axis=1)
-        if self.has_f or self.has_pf or self.relative_f or self.has_gf:
+            weighted_atom_ener = atom_ener_coeff * atom_ener
+            if is_ragged:
+                if frame_id is None:
+                    frame_id = frame_id_from_n_node(n_node, n_total=atom_ener.shape[0])
+                energy = segment_sum(weighted_atom_ener, frame_id, _nf)
+            else:
+                energy = xp.sum(weighted_atom_ener, axis=1)
+        if force_required:
             force_reshape = xp.reshape(force, (-1,))
             force_hat_reshape = xp.reshape(force_hat, (-1,))
             diff_f = force_hat_reshape - force_reshape
@@ -325,22 +376,33 @@ def call(
         atom_norm = 1.0 / natoms
         atom_norm_ener = 1.0 / natoms
         lr_ratio = learning_rate / self.starter_learning_rate
-        pref_e = find_energy * (
-            self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * lr_ratio
-        )
-        pref_f = find_force * (
-            self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * lr_ratio
-        )
-        pref_v = find_virial * (
-            self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * lr_ratio
-        )
-        pref_ae = find_atom_ener * (
-            self.limit_pref_ae + (self.start_pref_ae - self.limit_pref_ae) * lr_ratio
-        )
-        pref_pf = find_atom_pref * (
-            self.limit_pref_pf + (self.start_pref_pf - self.limit_pref_pf) * lr_ratio
-        )
-        pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * lr_ratio
+        if self.has_e:
+            pref_e = find_energy * (
+                self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * lr_ratio
+            )
+        if self.has_f:
+            pref_f = find_force * (
+                self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * lr_ratio
+            )
+        if self.has_v:
+            pref_v = find_virial * (
+                self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * lr_ratio
+            )
+        if self.has_ae:
+            pref_ae = find_atom_ener * (
+                self.limit_pref_ae
+                + (self.start_pref_ae - self.limit_pref_ae) * lr_ratio
+            )
+        if self.has_pf:
+            effective_find_pf = find_force * find_atom_pref
+            pref_pf = effective_find_pf * (
+                self.limit_pref_pf
+                + (self.start_pref_pf - self.limit_pref_pf) * lr_ratio
+            )
+        if self.has_h:
+            pref_h = (
+                self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * lr_ratio
+            )
 
         loss = 0
         more_loss = {}
@@ -351,7 +413,7 @@ def call(
         if self.has_e:
             if self.loss_func == "mse":
                 l2_ener_loss = xp.mean(xp.square(energy - energy_hat))
-                if maskf is not None:
+                if inv is not None:
                     # Idiom 2 (extensive): per-frame normalization by real-atom count.
                     se = xp.square(energy - energy_hat)  # [nf, k]
                     per_frame = per_frame_component_mean(se)  # [nf]
@@ -383,7 +445,7 @@ def call(
                     )
             elif self.loss_func == "mae":
                 l1_ener_loss = xp.mean(xp.abs(energy - energy_hat))
-                if maskf is not None:
+                if inv is not None:
                     abs_e = xp.abs(energy - energy_hat)  # [nf, k]
                     per_frame_ae = per_frame_component_mean(abs_e)  # [nf]
                     l1_ener_masked = xp.mean(per_frame_ae * inv)
@@ -401,7 +463,7 @@ def call(
                     f"Loss type {self.loss_func} is not implemented for energy loss."
                 )
             if mae:
-                if maskf is not None:
+                if inv is not None:
                     per_frame_ae = per_frame_component_mean(xp.abs(energy - energy_hat))
                     mae_e = xp.mean(per_frame_ae * inv)
                 else:
@@ -414,13 +476,17 @@ def call(
                 l2_force_loss = xp.mean(xp.square(diff_f))
                 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]
+                    diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3))
                     # Masked MSE computed for rmse_f display regardless of use_huber.
                     l2_force_masked = masked_atom_mean(xp.square(diff_f_3d), maskf, 3)
                     if not self.use_huber:
                         loss += pref_f * l2_force_masked
                     else:
+                        # ``f_use_norm`` selects the residual an atom
+                        # contributes: three independent components, or the
+                        # single L2 norm of its force-error vector. That choice
+                        # sets the label count per atom, which is exactly the
+                        # ``ncomp`` the pooled reduction divides by.
                         if not self.f_use_norm:
                             abs_e = xp.abs(diff_f_3d)
                             quad = 0.5 * xp.square(diff_f_3d)
@@ -429,33 +495,31 @@ def call(
                             )
                             huber_elem = xp.where(
                                 abs_e <= self._huber_delta_force, quad, lin
-                            )
-                            huber_masked = huber_elem * maskf_col
+                            )  # [nf, nloc, 3]
+                            huber_ncomp = 3
                         else:
-                            diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3))
+                            diff_3 = xp.reshape(force_hat - force, (*_node_shape, 3))
                             norm_2d = xp.reshape(
                                 xp.linalg.vector_norm(
                                     xp.reshape(diff_3, (-1, 3)), axis=1
                                 ),
-                                (_nf, _nloc),
+                                _node_shape,
                             )
                             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_elem = xp.reshape(
+                                xp.where(
+                                    abs_n <= self._huber_delta_force, quad_n, lin_n
+                                ),
+                                (*_node_shape, 1),
                             )
-                            huber_masked = xp.reshape(huber_n * maskf, (_nf, _nloc, 1))
-                        per_frame_sum = xp.sum(
-                            xp.reshape(huber_masked, (_nf, -1)), axis=-1
+                            huber_ncomp = 1
+                        l_huber_masked = masked_atom_mean(
+                            huber_elem, maskf, huber_ncomp
                         )
-                        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_masked), find_force
@@ -486,19 +550,19 @@ def call(
                     )
             elif self.loss_func == "mae":
                 if maskf is not None:
-                    diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
+                    diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3))
                     if not self.f_use_norm:
                         l1_force_masked = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3)
                     else:
-                        diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3))
+                        diff_3 = xp.reshape(force_hat - force, (*_node_shape, 3))
                         norm_2d = xp.reshape(
                             xp.linalg.vector_norm(xp.reshape(diff_3, (-1, 3)), axis=1),
-                            (_nf, _nloc),
+                            _node_shape,
+                        )
+                        # One L2 norm per atom, hence one label per atom.
+                        l1_force_masked = masked_atom_mean(
+                            xp.reshape(norm_2d, (*_node_shape, 1)), maskf, 1
                         )
-                        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
@@ -521,7 +585,7 @@ def call(
                 )
             if mae:
                 if maskf is not None:
-                    diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
+                    diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3))
                     mae_f = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3)
                 else:
                     mae_f = xp.mean(xp.abs(diff_f))
@@ -533,7 +597,7 @@ def call(
                 l2_virial_loss = xp.mean(
                     xp.square(virial_hat_reshape - virial_reshape),
                 )
-                if maskf is not None:
+                if inv 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))
@@ -567,7 +631,7 @@ def call(
                     )
             elif self.loss_func == "mae":
                 l1_virial_loss = xp.mean(xp.abs(virial_hat_reshape - virial_reshape))
-                if maskf is not None:
+                if inv is not None:
                     v2d = xp.reshape(virial, (_nf, 9))
                     v_hat_2d = xp.reshape(virial_hat, (_nf, 9))
                     per_frame_v = per_frame_component_mean(
@@ -588,7 +652,7 @@ def call(
                     f"Loss type {self.loss_func} is not implemented for virial loss."
                 )
             if mae:
-                if maskf is not None:
+                if inv is not None:
                     v2d = xp.reshape(virial, (_nf, 9))
                     v_hat_2d = xp.reshape(virial_hat, (_nf, 9))
                     per_frame_v = per_frame_component_mean(xp.abs(v_hat_2d - v2d))
@@ -607,11 +671,10 @@ def call(
                 )
                 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))
-                    per_frame_dof = xp.sum(maskf, axis=-1)  # [nf]
+                    ae_2d = xp.reshape(atom_ener, _node_shape)
+                    ae_hat_2d = xp.reshape(atom_ener_hat, _node_shape)
                     l2_ae_masked = masked_atom_mean(
-                        xp.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1
+                        xp.square(ae_hat_2d - ae_2d)[..., None], maskf, 1
                     )
                     if not self.use_huber:
                         loss += pref_ae * l2_ae_masked
@@ -626,9 +689,9 @@ def call(
                         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)
+                        l_huber_ae_masked = masked_atom_mean(
+                            huber_ae[..., None], maskf, 1
+                        )
                         loss += pref_ae * l_huber_ae_masked
                     more_loss["rmse_ae"] = self.display_if_exist(
                         xp.sqrt(l2_ae_masked), find_atom_ener
@@ -651,10 +714,10 @@ def call(
                     xp.abs(atom_ener_hat_reshape - atom_ener_reshape)
                 )
                 if maskf is not None:
-                    ae_2d = xp.reshape(atom_ener, (_nf, _nloc))
-                    ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc))
+                    ae_2d = xp.reshape(atom_ener, _node_shape)
+                    ae_hat_2d = xp.reshape(atom_ener_hat, _node_shape)
                     l1_ae_masked = masked_atom_mean(
-                        xp.abs(ae_hat_2d - ae_2d)[:, :, None], maskf, 1
+                        xp.abs(ae_hat_2d - ae_2d)[..., None], maskf, 1
                     )
                     loss += pref_ae * l1_ae_masked
                     more_loss["mae_ae"] = self.display_if_exist(
@@ -678,45 +741,57 @@ def call(
                 )
                 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))
+                    diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3))
+                    pf_3d = xp.reshape(atom_pref, (*_node_shape, 3))
                     l2_pf_masked = masked_atom_mean(
                         xp.square(diff_f_3d) * pf_3d, maskf, 3
                     )
                     loss += pref_pf * l2_pf_masked
                     more_loss["rmse_pf"] = self.display_if_exist(
-                        xp.sqrt(l2_pf_masked), find_atom_pref
+                        xp.sqrt(l2_pf_masked), effective_find_pf
                     )
                 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
+                        xp.sqrt(l2_pref_force_loss), effective_find_pf
                     )
             elif self.loss_func == "mae":
                 l1_pref_force_loss = xp.mean(
                     xp.multiply(xp.abs(diff_f), atom_pref_reshape)
                 )
                 if maskf is not None:
-                    diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
-                    pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3))
+                    diff_f_3d = xp.reshape(diff_f, (*_node_shape, 3))
+                    pf_3d = xp.reshape(atom_pref, (*_node_shape, 3))
                     l1_pf_masked = masked_atom_mean(xp.abs(diff_f_3d) * pf_3d, maskf, 3)
                     loss += pref_pf * l1_pf_masked
                     more_loss["mae_pf"] = self.display_if_exist(
-                        l1_pf_masked, find_atom_pref
+                        l1_pf_masked, effective_find_pf
                     )
                 else:
                     loss += pref_pf * l1_pref_force_loss
                     more_loss["mae_pf"] = self.display_if_exist(
-                        l1_pref_force_loss, find_atom_pref
+                        l1_pref_force_loss, effective_find_pf
                     )
             else:
                 raise NotImplementedError(
                     f"Loss type {self.loss_func} is not implemented for atom prefactor force loss."
                 )
         if self.has_gf:
+            if is_ragged:
+                # ``natoms`` below is one number for the whole batch, which a
+                # padded batch can honour and a concatenated one cannot: its
+                # frames differ in atom count, so ``drdq``, stored per frame
+                # against a common atom axis, has no shape to take.
+                raise NotImplementedError(
+                    "the generalized force loss requires every frame of a "
+                    "batch to hold the same number of atoms; a batch whose "
+                    "frames are concatenated cannot provide the common atom "
+                    "axis its ``drdq`` label is stored against"
+                )
             find_drdq = label_dict["find_drdq"]
             drdq = label_dict["drdq"]
-            pref_gf = find_drdq * (
+            effective_find_gf = find_force * find_drdq
+            pref_gf = effective_find_gf * (
                 self.limit_pref_gf
                 + (self.start_pref_gf - self.limit_pref_gf) * lr_ratio
             )
@@ -752,10 +827,14 @@ def call(
             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
+                xp.sqrt(l2_gen_force_loss), effective_find_gf
             )
         hessian = model_dict.get("hessian", model_dict.get("energy_derv_r_derv_r"))
         if self.has_h and hessian is not None and "hessian" in label_dict:
+            if is_ragged:
+                raise NotImplementedError(
+                    "the hessian loss requires a rectangular atom axis"
+                )
             find_hessian = label_dict.get("find_hessian", 0.0)
             if maskf is not None:
                 hessian_shape = (_nf, _nloc * 3, _nloc * 3)
@@ -799,54 +878,60 @@ def call(
     @property
     def label_requirement(self) -> list[DataRequirementItem]:
         """Return data label requirements needed for this loss calculation."""
-        label_requirement = []
-        label_requirement.append(
-            DataRequirementItem(
-                "energy",
-                ndof=1,
-                atomic=False,
-                must=False,
-                high_prec=True,
+        label_requirement: list[DataRequirementItem] = []
+        if self.has_e:
+            label_requirement.append(
+                DataRequirementItem(
+                    "energy",
+                    ndof=1,
+                    atomic=False,
+                    must=False,
+                    high_prec=True,
+                )
             )
-        )
-        label_requirement.append(
-            DataRequirementItem(
-                "force",
-                ndof=3,
-                atomic=True,
-                must=False,
-                high_prec=False,
+        if self.has_f or self.has_pf or self.relative_f is not None or self.has_gf:
+            label_requirement.append(
+                DataRequirementItem(
+                    "force",
+                    ndof=3,
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                )
             )
-        )
-        label_requirement.append(
-            DataRequirementItem(
-                "virial",
-                ndof=9,
-                atomic=False,
-                must=False,
-                high_prec=False,
+        if self.has_v:
+            label_requirement.append(
+                DataRequirementItem(
+                    "virial",
+                    ndof=9,
+                    atomic=False,
+                    must=False,
+                    high_prec=False,
+                )
             )
-        )
-        label_requirement.append(
-            DataRequirementItem(
-                "atom_ener",
-                ndof=1,
-                atomic=True,
-                must=False,
-                high_prec=False,
+        if self.has_ae:
+            label_requirement.append(
+                DataRequirementItem(
+                    "atom_ener",
+                    ndof=1,
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                )
             )
-        )
-        label_requirement.append(
-            DataRequirementItem(
-                "atom_pref",
-                ndof=1,
-                atomic=True,
-                must=False,
-                high_prec=False,
-                repeat=3,
-                default=1.0,
+        if self.has_pf:
+            label_requirement.append(
+                DataRequirementItem(
+                    "atom_pref",
+                    ndof=1,
+                    atomic=True,
+                    must=False,
+                    high_prec=False,
+                    repeat=3,
+                    default=1.0,
+                    source_policy="default" if self.use_default_pf else "tracked",
+                )
             )
-        )
         if self.has_gf > 0:
             label_requirement.append(
                 DataRequirementItem(
@@ -866,6 +951,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     must=False,
                     high_prec=False,
                     default=1.0,
+                    source_policy="default",
                 )
             )
         if self.has_h:
diff --git a/deepmd/dpmodel/loss/ener_spin.py b/deepmd/dpmodel/loss/ener_spin.py
index 30959fa83f..8015def9cd 100644
--- a/deepmd/dpmodel/loss/ener_spin.py
+++ b/deepmd/dpmodel/loss/ener_spin.py
@@ -15,6 +15,12 @@
     masked_atom_mean,
     per_frame_component_mean,
 )
+from deepmd.dpmodel.utils.neighbor_graph.graph import (
+    frame_id_from_n_node,
+)
+from deepmd.dpmodel.utils.neighbor_graph.segment import (
+    segment_sum,
+)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
@@ -135,6 +141,11 @@ def __init__(
         self.has_v = self.start_pref_v != 0.0 or self.limit_pref_v != 0.0
         self.has_ae = self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0
 
+    @property
+    def supports_ragged_batches(self) -> bool:
+        """Whether the configured terms accept a flat per-node batch axis."""
+        return True
+
     def call(
         self,
         learning_rate: float,
@@ -162,18 +173,51 @@ 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:
-            # 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
+        # Rectangular batches describe valid nodes with ``mask``; ragged ones
+        # concatenate the node axis and delimit frames with ``n_node``. The
+        # same included-node counts normalize extensive frame-level terms in
+        # either layout, while the mask drives pooled per-node reductions.
+        maskf = (
+            xp.astype(model_dict["mask"], energy.dtype)
+            if "mask" in model_dict
+            else None
+        )
+        is_ragged = "n_node" in model_dict
+        frame_id = None
+        included_n_node = None
+        inv = None
+        if is_ragged:
+            n_node = model_dict["n_node"]
+            nframes = n_node.shape[0]
+            if maskf is None:
+                included_n_node = xp.astype(n_node, energy.dtype)
+            else:
+                frame_id = frame_id_from_n_node(n_node, n_total=maskf.shape[0])
+                included_n_node = segment_sum(
+                    xp.reshape(maskf, (-1,)), frame_id, nframes
+                )
+        elif maskf is not None:
+            included_n_node = xp.reshape(xp.sum(maskf, axis=-1), (-1,))
+        if included_n_node is not None:
+            has_node = included_n_node > 0
+            safe_n_node = xp.where(
+                has_node,
+                included_n_node,
+                xp.ones_like(included_n_node),
+            )
+            inv = xp.where(
+                has_node,
+                1.0 / safe_n_node,
+                xp.zeros_like(included_n_node),
+            )
+
+        def reshape_atomic(value: Array, ncomp: int) -> Array:
+            """Return one atomic field on the active rectangular or flat axis."""
+            if is_ragged:
+                return xp.reshape(value, (-1, ncomp))
+            if maskf is not None:
+                return xp.reshape(value, (*maskf.shape, ncomp))
+            return xp.reshape(value, (-1, natoms, ncomp))
 
         if self.has_e:
             energy_pred = model_dict["energy"]
@@ -181,14 +225,25 @@ def call(
             find_energy = label_dict.get("find_energy", 0.0)
             pref_e = pref_e * find_energy
             if self.enable_atom_ener_coeff and "atom_energy" in model_dict:
-                atom_ener_pred = model_dict["atom_energy"]
-                atom_ener_coeff = label_dict["atom_ener_coeff"]
-                atom_ener_coeff = xp.reshape(atom_ener_coeff, atom_ener_pred.shape)
-                energy_pred = xp.sum(atom_ener_coeff * atom_ener_pred, axis=1)
+                atom_ener_pred = reshape_atomic(model_dict["atom_energy"], 1)
+                atom_ener_coeff = reshape_atomic(label_dict["atom_ener_coeff"], 1)
+                weighted_atom_ener = atom_ener_coeff * atom_ener_pred
+                if maskf is not None:
+                    weighted_atom_ener = weighted_atom_ener * xp.reshape(
+                        maskf,
+                        (*maskf.shape, 1),
+                    )
+                if is_ragged:
+                    if frame_id is None:
+                        frame_id = frame_id_from_n_node(
+                            n_node, n_total=atom_ener_pred.shape[0]
+                        )
+                    energy_pred = segment_sum(weighted_atom_ener, frame_id, nframes)
+                else:
+                    energy_pred = xp.sum(weighted_atom_ener, axis=1)
             if self.loss_func == "mse":
                 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.
+                if inv is not None:
                     per_frame_e = per_frame_component_mean(se_e)  # [nf]
                     loss += pref_e * xp.mean(per_frame_e * inv**norm_exp)
                     more_loss["rmse_e"] = self.display_if_exist(
@@ -202,8 +257,7 @@ def call(
                     )
             elif self.loss_func == "mae":
                 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.
+                if inv is not None:
                     per_frame_ae = per_frame_component_mean(
                         xp.abs(energy_pred - energy_label)
                     )  # [nf]
@@ -218,7 +272,7 @@ def call(
                         l1_ener_loss * atom_norm, find_energy
                     )
             if mae:
-                if maskf is not None:
+                if inv is not None:
                     per_frame_ae = per_frame_component_mean(
                         xp.abs(energy_pred - energy_label)
                     )
@@ -232,26 +286,18 @@ def call(
         if self.has_fr:
             find_force = label_dict.get("find_force", 0.0)
             pref_fr = pref_fr * find_force
-            # Reshape to the canonical (nf, natoms, 3) atomic shape: the raw
-            # data-loader label is flat (nf, natoms * 3), matching the
-            # ``xp.reshape(label_dict[...], (-1, natoms, ncomp))`` idiom used
-            # by every other atomic-label loss (see ``dpmodel/loss/dos.py``
-            # and ``dpmodel/loss/tensor.py``).
-            force_pred = xp.reshape(model_dict["force"], (-1, natoms, 3))
-            force_label = xp.reshape(label_dict["force"], (-1, natoms, 3))
+            force_pred = reshape_atomic(model_dict["force"], 3)
+            force_label = reshape_atomic(label_dict["force"], 3)
+            diff_fr = force_label - force_pred
             if self.loss_func == "mse":
-                diff_fr = force_label - force_pred  # [nf, nloc, 3]
                 if maskf is not None:
-                    # Idiom 1 (per-atom masked mean, ncomp=3).
                     l2_force_real_loss = masked_atom_mean(xp.square(diff_fr), maskf, 3)
                     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 = masked_atom_mean(
-                            xp.abs(force_label - force_pred), maskf, 3
-                        )
+                        mae_fr = masked_atom_mean(xp.abs(diff_fr), maskf, 3)
                         more_loss["mae_fr"] = self.display_if_exist(mae_fr, find_force)
                 else:
                     l2_force_real_loss = xp.mean(xp.square(diff_fr))
@@ -263,9 +309,8 @@ 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)  # [nf, nloc, 3]
+                abs_diff_fr = xp.abs(diff_fr)
                 if maskf is not None:
-                    # Idiom 1 (per-atom masked mean, ncomp=3) with abs.
                     l1_force_real_masked = masked_atom_mean(abs_diff_fr, maskf, 3)
                     loss += pref_fr * l1_force_real_masked
                     more_loss["mae_fr"] = self.display_if_exist(
@@ -281,18 +326,17 @@ def call(
         if self.has_fm:
             find_force_mag = label_dict.get("find_force_mag", 0.0)
             pref_fm = pref_fm * find_force_mag
-            # Same flat -> (nf, natoms, 3) reshape as the real-force branch above.
-            force_mag_pred = xp.reshape(model_dict["force_mag"], (-1, natoms, 3))
-            force_mag_label = xp.reshape(label_dict["force_mag"], (-1, natoms, 3))
-            mask_mag = model_dict["mask_mag"]
-            # mask_mag: [nframes, natoms, 1], bool -> use mask multiplication
+            force_mag_pred = reshape_atomic(model_dict["force_mag"], 3)
+            force_mag_label = reshape_atomic(label_dict["force_mag"], 3)
+            mask_mag = reshape_atomic(model_dict["mask_mag"], 1)
+            if maskf is not None:
+                mask_mag = xp.logical_and(
+                    mask_mag,
+                    xp.reshape(maskf > 0, (*maskf.shape, 1)),
+                )
             mask_float = xp.astype(mask_mag, force_mag_pred.dtype)
-            # zero out non-magnetic atoms
             diff_fm = (force_mag_label - force_mag_pred) * mask_float
             n_valid = xp.sum(mask_float)
-            # Guard the denominator itself because array backends may evaluate
-            # both branches of ``where``. This is safe under JAX tracing and
-            # makes an all-empty magnetic mask contribute exactly zero.
             safe_n_valid = xp.where(n_valid > 0, n_valid, xp.ones_like(n_valid))
             if self.loss_func == "mse":
                 l2_force_mag_loss = xp.sum(xp.square(diff_fm)) / (safe_n_valid * 3)
@@ -304,11 +348,7 @@ def call(
                     mae_fm = xp.sum(xp.abs(diff_fm)) / (safe_n_valid * 3)
                     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
-                # 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.
+                abs_diff_fm = xp.abs(diff_fm)
                 l1_force_mag_loss = xp.sum(abs_diff_fm) / (safe_n_valid * 3)
                 loss += pref_fm * l1_force_mag_loss
                 more_loss["mae_fm"] = self.display_if_exist(
@@ -318,15 +358,12 @@ def call(
         if self.has_ae:
             find_atom_ener = label_dict.get("find_atom_ener", 0.0)
             pref_ae = pref_ae * find_atom_ener
-            atom_ener = model_dict["atom_energy"]
-            atom_ener_label = label_dict["atom_ener"]
+            atom_ener = reshape_atomic(model_dict["atom_energy"], 1)
+            atom_ener_label = reshape_atomic(label_dict["atom_ener"], 1)
             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))
                 if self.loss_func == "mse":
                     l2_atom_ener_loss = masked_atom_mean(
-                        xp.square(ae_label - ae), maskf, 1
+                        xp.square(atom_ener_label - atom_ener), maskf, 1
                     )
                     loss += pref_ae * l2_atom_ener_loss
                     more_loss["rmse_ae"] = self.display_if_exist(
@@ -334,27 +371,21 @@ def call(
                     )
                 elif self.loss_func == "mae":
                     l1_atom_ener_loss = masked_atom_mean(
-                        xp.abs(ae_label - ae), maskf, 1
+                        xp.abs(atom_ener_label - atom_ener), maskf, 1
                     )
                     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)
-                    )
+                    l2_atom_ener_loss = xp.mean(xp.square(atom_ener_label - atom_ener))
                     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)
-                    )
+                    l1_atom_ener_loss = xp.mean(xp.abs(atom_ener_label - atom_ener))
                     loss += pref_ae * l1_atom_ener_loss
                     more_loss["mae_ae"] = self.display_if_exist(
                         l1_atom_ener_loss, find_atom_ener
@@ -364,11 +395,10 @@ def call(
             find_virial = label_dict.get("find_virial", 0.0)
             pref_v = pref_v * find_virial
             virial_pred = xp.reshape(model_dict["virial"], (-1, 9))
-            virial_label = label_dict["virial"]
+            virial_label = xp.reshape(label_dict["virial"], (-1, 9))
             diff_v = virial_label - virial_pred  # [nf, 9]
             if self.loss_func == "mse":
-                if maskf is not None:
-                    # Idiom 2 (extensive, k=9): per-frame normalization by real-atom count.
+                if inv is not None:
                     per_frame_v = per_frame_component_mean(xp.square(diff_v))  # [nf]
                     loss += pref_v * xp.mean(per_frame_v * inv**norm_exp)
                     more_loss["rmse_v"] = self.display_if_exist(
@@ -391,8 +421,7 @@ 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))
-                if maskf is not None:
-                    # Idiom 2 (extensive, k=9) with abs: per-frame normalization by real-atom count.
+                if inv is not None:
                     per_frame_v = per_frame_component_mean(xp.abs(diff_v))  # [nf]
                     l1_virial_masked = xp.mean(per_frame_v * inv)
                     loss += pref_v * l1_virial_masked
@@ -471,6 +500,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     must=False,
                     high_prec=False,
                     default=1.0,
+                    source_policy="default",
                 )
             )
         return label_requirement
diff --git a/deepmd/dpmodel/loss/loss.py b/deepmd/dpmodel/loss/loss.py
index 3f45b12b8c..efee9c8c55 100644
--- a/deepmd/dpmodel/loss/loss.py
+++ b/deepmd/dpmodel/loss/loss.py
@@ -51,6 +51,11 @@ def call(
     def label_requirement(self) -> list[DataRequirementItem]:
         """Return data label requirements needed for this loss calculation."""
 
+    @property
+    def supports_ragged_batches(self) -> bool:
+        """Whether this objective accepts a flat per-node batch axis."""
+        return False
+
     @staticmethod
     def display_if_exist(loss: Array, find_property: float) -> Array:
         """Display NaN if labeled property is not found.
diff --git a/deepmd/dpmodel/loss/reduction.py b/deepmd/dpmodel/loss/reduction.py
index 2568913538..39980cc707 100644
--- a/deepmd/dpmodel/loss/reduction.py
+++ b/deepmd/dpmodel/loss/reduction.py
@@ -1,19 +1,64 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Shared masked per-frame reduction idioms for the loss modules.
-
-These helpers factor out the three per-frame reduction patterns that the
-mixed_type padding mask (PR #5738) introduced into every loss term (issue
-#5768). They are written with ``array_api_compat`` so both the dpmodel
-(numpy/jax/...) loss backend and the PyTorch loss backend can call them: the
-PyTorch backend passes torch tensors and ``array_api_compat`` dispatches to the
-torch namespace, preserving autograd and producing bit-identical results to the
-previous hand-inlined torch code.
-
-Each helper implements ONLY the masked branch. Callers keep the original
-non-masked expression in the ``else`` branch verbatim, so the "bit-identical
-for non-mixed batches" guarantee from PR #5738 is preserved (defaulting the
-mask to all-ones would change the reduction order at the ULP level and is
-deliberately NOT done here).
+"""Shared masked reduction idioms for the loss modules.
+
+These helpers factor out the reduction patterns every loss term needs once a
+per-atom mask marks which rows of a batch are real. They are written with
+``array_api_compat`` so both the dpmodel (numpy/jax/...) loss backend and the
+PyTorch loss backend can call them: the PyTorch backend passes torch tensors
+and ``array_api_compat`` dispatches to the torch namespace, preserving
+autograd.
+
+Reduction convention
+--------------------
+A batch may hold frames of unequal atom count, padded to a common width. Each
+term must therefore decide how one frame's aggregate error weighs against
+another's, and the answer differs by term because frames do not carry equally
+many labels:
+
+- **Per-atom terms** (force, atomic energy, atomic prefactor force, dos,
+  tensor) carry a number of labels proportional to the frame's atom count.
+  :func:`masked_atom_mean` pools them: it divides the summed contribution of
+  the whole batch by the batch's total label count, so every real label counts
+  once and a frame's weight is proportional to its atom count.
+- **Frame-level terms** (energy, virial) carry a fixed number of labels per
+  frame. Pooling and averaging over frames coincide there, so
+  :func:`per_frame_component_mean` reduces per frame and leaves the frame axis
+  to the caller, which applies the extensive ``1 / natoms`` weighting.
+- **Pair terms** (Hessian) describe one response matrix per frame. Their
+  component count grows quadratically with atom count, so
+  :func:`masked_pair_mean` normalizes each matrix before averaging frames;
+  otherwise a large structure would dominate a batch quadratically.
+
+Pooling is what keeps a frame's weight independent of the company it keeps.
+The alternative -- averaging each frame's own per-label mean -- gives every
+frame the same weight whatever its size, which makes a label in a small frame
+count for more than one in a large frame, by the ratio of their atom counts.
+
+Writing ``S_f`` for the summed contribution of frame ``f``, ``k`` for the
+number of frames and ``d`` for the labels each of them carries, the two
+coincide whenever that count is common to the batch:
+
+    sum_f(S_f) / (k * d)  ==  (1 / k) * sum_f(S_f / d)
+
+so a batch of uniform atom count needs no special case anywhere in this
+module, and the choice between the two is unobservable there. They part
+company only where a batch holds frames of differing real atom count, which
+arises in exactly two places:
+
+- ``mix:N`` LMDB batching, which packs frames of differing atom count by
+  construction.
+- ``mixed_type`` npy data whose ``real_atom_types.npy`` spends a different
+  number of ``-1`` rows on different frames of one system. The format permits
+  this and the documentation describes it as the way to merge frames of
+  unequal atom count, but a system written by dpdata pads every frame equally
+  and is therefore unaffected.
+
+The TensorFlow backend reaches neither case: it drops ``real_natoms_vec``
+before the feed dict and normalizes by the padded width throughout.
+
+Each helper implements ONLY the masked branch. The unmasked branch of each
+caller is a plain mean over the whole batch, which pools by construction, so
+both branches express the same convention.
 """
 
 from typing import (
@@ -28,16 +73,25 @@
 
 
 def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
-    """Idiom 1: per-atom masked mean over ``ncomp`` components, averaged over frames.
+    """Idiom 1: mean of a per-atom contribution over the batch's real labels.
+
+    The contribution of every real atom is pooled across frames before the
+    division, so the reduction is a mean over labels rather than a mean over
+    frames of per-frame means. See the module docstring for why the per-atom
+    terms weigh frames by their label count, and for the identity that makes
+    this reduce to the per-frame mean on a uniform-atom-count batch.
 
     Parameters
     ----------
     elem : Array
-        Non-negative per-element contribution of shape ``[nf, nloc, ncomp]``
-        (already squared or abs'd, and pre-multiplied by any per-atom weight
-        such as ``atom_pref``). NOT yet multiplied by the mask.
+        Non-negative per-element contribution of shape
+        ``[*node_shape, ncomp]`` (already squared or abs'd, and pre-multiplied
+        by any per-atom weight such as ``atom_pref``). NOT yet multiplied by
+        the mask. ``node_shape`` may be rectangular ``(nf, nloc)`` or a flat
+        ragged node axis ``(N,)``.
     maskf : Array
-        Per-atom real/ghost mask of shape ``[nf, nloc]`` (1.0 real, 0.0 ghost).
+        Per-atom inclusion mask of shape ``node_shape`` (1.0 included,
+        0.0 excluded).
     ncomp : int
         Number of components per atom (force: 3, atom energy: 1,
         dos: ``numb_dos``, tensor: ``tensor_size``).
@@ -45,26 +99,22 @@ def masked_atom_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
     Returns
     -------
     Array
-        ``mean_over_frames( sum(elem * mask) / (real_natoms * ncomp) )``.
-        An all-padding frame (zero real atoms) contributes a neutral ``0``
-        instead of ``0/0 = NaN``.
+        ``sum(elem * mask) / (ncomp * sum(mask))`` over the whole batch.
+        A batch holding no real atom contributes a neutral ``0`` instead of
+        ``0/0 = NaN``.
     """
     xp = array_api_compat.array_namespace(elem, maskf)
-    nf = elem.shape[0]
-    masked = elem * maskf[:, :, None]
-    per_frame_sum = xp.sum(xp.reshape(masked, (nf, -1)), axis=-1)
-    per_frame_dof = xp.sum(maskf, axis=-1) * ncomp
-    # An all-padding frame has zero real atoms, so ``per_frame_dof`` is 0 and
-    # the ratio would be 0/0 = NaN -- poisoning the frame mean and, under
-    # autograd, its gradient. Divide by a safe denominator and map those frames
-    # to a neutral per-frame value of 0. Frames with real atoms are untouched,
-    # preserving the bit-identical guarantee.
-    has_dof = per_frame_dof > 0
-    safe_dof = xp.where(has_dof, per_frame_dof, xp.ones_like(per_frame_dof))
-    per_frame = xp.where(
-        has_dof, per_frame_sum / safe_dof, xp.zeros_like(per_frame_sum)
-    )
-    return xp.mean(per_frame)
+    node_elem = xp.reshape(elem, (-1, ncomp))
+    node_mask = xp.reshape(maskf, (-1, 1))
+    total = xp.sum(node_elem * node_mask)
+    total_dof = xp.sum(node_mask) * ncomp
+    # A batch of nothing but padding has no label to average over, and the
+    # ratio would be 0/0 = NaN -- poisoning the whole batch loss and, under
+    # autograd, its gradient. The division still runs on a safe denominator so
+    # that the discarded branch stays differentiable.
+    has_dof = total_dof > 0
+    safe_dof = xp.where(has_dof, total_dof, xp.ones_like(total_dof))
+    return xp.where(has_dof, total / safe_dof, xp.zeros_like(total))
 
 
 def masked_pair_mean(elem: Array, maskf: Array, ncomp: int) -> Array:
diff --git a/deepmd/dpmodel/loss/tensor.py b/deepmd/dpmodel/loss/tensor.py
index e054ab3cc6..9af03a3b85 100644
--- a/deepmd/dpmodel/loss/tensor.py
+++ b/deepmd/dpmodel/loss/tensor.py
@@ -124,7 +124,7 @@ def call(
             diff = xp.reshape(local_pred - local_label, (-1, self.tensor_size))
             diff = diff * atomic_weight
             if "mask" in model_dict:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=tensor_size).
                 maskf = xp.astype(model_dict["mask"], diff.dtype)  # [nf, natoms]
                 diff3d = xp.reshape(
                     diff, (local_pred.shape[0], natoms, self.tensor_size)
@@ -197,6 +197,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     must=False,
                     high_prec=False,
                     default=1.0,
+                    source_policy="default",
                 )
             )
         return label_requirement
diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py
index 73b53ce6ea..b6f765c22a 100644
--- a/deepmd/dpmodel/model/make_model.py
+++ b/deepmd/dpmodel/model/make_model.py
@@ -48,6 +48,8 @@
     NeighborGraph,
     build_neighbor_graph,
     build_neighbor_graph_ase,
+    compact_nodes,
+    expand_node_values,
 )
 from deepmd.utils.path import (
     DPPath,
@@ -543,12 +545,21 @@ def _call_common_graph(
                 )
             xp = array_api_compat.array_namespace(atype)
             nf, nloc = atype.shape[:2]
+            n_padded = nf * nloc
+            atype_flat = xp.reshape(atype, (n_padded,))
+            # A batch of unequal atom counts arrives padded to a common width
+            # with phantom atoms (atype < 0). The builders leave them out of
+            # every edge, so dropping them from the node axis costs nothing and
+            # spares the network from evaluating them. On a batch of uniform
+            # atom count the mask is all true and this is a renumbering by the
+            # identity.
+            ng, node_index = compact_nodes(ng, atype_flat >= 0)
             # OUTPUT-AGNOSTIC standard model dict (``<var>``, ``<var>_redu``,
             # derivative name-holders ``None``, plus int ``mask``), like the
             # dense ``call_common``.  ``call_lower_graph`` masks virtual atoms
             # (atype<0) and sets the real int mask.
             model_predict = self.call_lower_graph(
-                atype=xp.reshape(atype, (nf * nloc,)),
+                atype=xp.take(atype_flat, node_index, axis=0),
                 n_node=ng.n_node,
                 edge_index=ng.edge_index,
                 edge_vec=ng.edge_vec,
@@ -556,25 +567,38 @@ def _call_common_graph(
                 fparam=fp,
                 # graph-lower ABI: aparam is FLAT on the node axis, (N, nda).
                 aparam=(
-                    xp.reshape(ap, (nf * nloc, ap.shape[-1]))
+                    xp.take(
+                        xp.reshape(ap, (n_padded, ap.shape[-1])), node_index, axis=0
+                    )
                     if ap is not None
                     else None
                 ),
-                spin=(xp.reshape(spin, (nf * nloc, 3)) if spin is not None else None),
+                spin=(
+                    xp.take(xp.reshape(spin, (n_padded, 3)), node_index, axis=0)
+                    if spin is not None
+                    else None
+                ),
                 charge_spin=charge_spin,
             )
-            # Public ABI is rectangular (nf, nloc, *); the lower is flat
-            # (N=nf*nloc, *).  Unravel per-atom keys here at the boundary.
-            # public call_common always passes rectangular (nf,nloc) coord/atype (N == nf*nloc), so this unravel always applies; ragged graphs reach call_lower_graph/forward_common_lower_graph directly (no unravel) and stay flat (N,*).
+            # Public ABI is rectangular (nf, nloc, *); the lower is flat over
+            # the real atoms. Scatter per-atom keys back onto the padded width
+            # here at the boundary, where a phantom slot reads zero, which is
+            # what a masked-out atom contributed there before.
+            # Only the rectangular entry reaches this scatter; the ragged
+            # one keeps the flat axis its caller handed over.
+            n_real = node_index.shape[0]
             for k in list(model_predict.keys()):
                 v = model_predict[k]
                 # per-frame reduced keys (..._redu) keep their (nf, *) shape; only node-level (N,*) keys unravel — guards the nloc==1 case where N == nf.
                 if (
                     v is not None
                     and not k.endswith("_redu")
-                    and v.shape[:1] == (nf * nloc,)
+                    and v.shape[:1] == (n_real,)
                 ):
-                    model_predict[k] = xp.reshape(v, (nf, nloc, *v.shape[1:]))
+                    model_predict[k] = xp.reshape(
+                        expand_node_values(v, node_index, n_padded),
+                        (nf, nloc, *v.shape[1:]),
+                    )
             return model_predict
 
         def call_common_lower(
diff --git a/deepmd/dpmodel/model/native_spin_model.py b/deepmd/dpmodel/model/native_spin_model.py
index cd90b763b6..5df83d91df 100644
--- a/deepmd/dpmodel/model/native_spin_model.py
+++ b/deepmd/dpmodel/model/native_spin_model.py
@@ -10,6 +10,7 @@
     Any,
 )
 
+import array_api_compat
 import numpy as np
 
 from deepmd.dpmodel.model.base_model import (
@@ -126,6 +127,8 @@ def translated_output_def(self) -> dict[str, Any]:
                 output_def["virial"].squeeze(-2)
                 output_def["atom_virial"] = deepcopy(out_def_data[f"{var_name}_derv_c"])
                 output_def["atom_virial"].squeeze(-2)
+            if "mask" in out_def_data:
+                output_def["mask"] = out_def_data["mask"]
             return output_def
 
         def _spin_active_mask(self, atype: np.ndarray) -> np.ndarray:
@@ -133,7 +136,21 @@ def _spin_active_mask(self, atype: np.ndarray) -> np.ndarray:
             where the atom type carries spin. Array-api compatible; reused by
             the eager and graph-export translations.
             """
-            return (self.spin_mask[atype] > 0)[..., None]
+            xp = array_api_compat.array_namespace(atype)
+            real_atom = atype >= 0
+            safe_atype = xp.where(real_atom, atype, xp.zeros_like(atype))
+            spin_mask = xp.asarray(
+                self.spin_mask,
+                device=array_api_compat.device(atype),
+            )
+            spin_active = (
+                xp.reshape(
+                    xp.take(spin_mask, xp.reshape(safe_atype, (-1,)), axis=0),
+                    atype.shape,
+                )
+                > 0
+            )
+            return xp.logical_and(spin_active, real_atom)[..., None]
 
         def _translate_eager_call(
             self,
@@ -164,6 +181,9 @@ def _translate_eager_call(
             for kk_src, kk_dst in translated:
                 src = model_ret.get(kk_src)
                 out[kk_dst] = np.squeeze(src, axis=-2) if src is not None else None
+            for key in ("mask", "n_node"):
+                if key in model_ret:
+                    out[key] = model_ret[key]
             return out
 
         def call(
diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py
index 17c575a92b..c7daebef71 100644
--- a/deepmd/dpmodel/model/spin_model.py
+++ b/deepmd/dpmodel/model/spin_model.py
@@ -14,6 +14,7 @@
 
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_asarray_nodetach,
 )
 from deepmd.dpmodel.atomic_model.dp_atomic_model import (
     DPAtomicModel,
@@ -34,6 +35,60 @@
     Spin,
 )
 
+_NATOMS_VEC_KEYS = ("natoms", "real_natoms_vec")
+_SPIN_STAT_RESERVED_KEYS = frozenset(
+    {"coord", "atype", "spin", "natoms", "real_natoms_vec", "aparam"}
+)
+
+
+def _expand_natoms_vec_for_virtual_spin(natoms: Array) -> Array:
+    """Expand a DeePMD natoms vector for the virtual-atom spin layout.
+
+    The leading two entries count local and extended atoms; they are doubled
+    to reflect real/virtual atom pairs. Per-type counts are duplicated so that
+    virtual spin slots inherit the population of their real counterparts.
+
+    Parameters
+    ----------
+    natoms : Array
+        Natoms vector with shape ``(nframes, ntypes_real + 2)``.
+
+    Returns
+    -------
+    Array
+        Expanded vector with shape ``(nframes, 2 * ntypes_real + 2)``.
+    """
+    xp = array_api_compat.array_namespace(natoms)
+    return xp.concat(
+        [2 * natoms[:, :2], natoms[:, 2:], natoms[:, 2:]],
+        axis=-1,
+    )
+
+
+def _pack_spin_stat_sample(
+    spin_model: "SpinModel",
+    sample: dict[str, Any],
+) -> dict[str, Any]:
+    """Transform one statistics sample into the virtual-atom spin layout."""
+    coord_updated, atype_updated, _ = spin_model.process_spin_input(
+        sample["coord"], sample["atype"], sample["spin"]
+    )
+    packed: dict[str, Any] = {
+        "coord": coord_updated,
+        "atype": atype_updated,
+    }
+    if "aparam" in sample:
+        packed["aparam"] = spin_model.expand_aparam(
+            sample["aparam"], atype_updated.shape[1]
+        )
+    for key in _NATOMS_VEC_KEYS:
+        if key in sample:
+            packed[key] = _expand_natoms_vec_for_virtual_spin(sample[key])
+    for key, value in sample.items():
+        if key not in _SPIN_STAT_RESERVED_KEYS:
+            packed[key] = value
+    return packed
+
 
 class SpinModel(NativeOP):
     r"""A spin model wrapper, with spin input preprocess and output split.
@@ -84,8 +139,12 @@ def __init__(
         self.spin_mask = self.spin.get_spin_mask()
 
     def _to_xp(self, arr: Any, xp: Any, ref_arr: Any) -> Any:
-        """Convert a numpy array to the same namespace as ref_arr."""
-        return xp.asarray(arr, device=array_api_compat.device(ref_arr))
+        """Convert an array to the namespace and device of ``ref_arr``."""
+        return xp_asarray_nodetach(
+            xp,
+            arr,
+            device=array_api_compat.device(ref_arr),
+        )
 
     def _lookup_type_values(self, values: Any, atype: Array, ref_arr: Array) -> Array:
         """Gather per-type values while mapping virtual atom types to zero.
@@ -422,37 +481,7 @@ def compute_or_load_stat(
 
         @functools.lru_cache
         def spin_sampled_func() -> list[dict[str, Any]]:
-            sampled = sampled_func()
-            spin_sampled = []
-            for sys in sampled:
-                coord_updated, atype_updated, _ = self.process_spin_input(
-                    sys["coord"], sys["atype"], sys["spin"]
-                )
-                tmp_dict = {
-                    "coord": coord_updated,
-                    "atype": atype_updated,
-                }
-                if "aparam" in sys:
-                    tmp_dict["aparam"] = self.expand_aparam(
-                        sys["aparam"], atype_updated.shape[1]
-                    )
-                if "natoms" in sys:
-                    natoms = sys["natoms"]
-                    xp = array_api_compat.array_namespace(natoms)
-                    tmp_dict["natoms"] = xp.concat(
-                        [2 * natoms[:, :2], natoms[:, 2:], natoms[:, 2:]], axis=-1
-                    )
-                for item_key in sys:
-                    if item_key not in [
-                        "coord",
-                        "atype",
-                        "spin",
-                        "natoms",
-                        "aparam",
-                    ]:
-                        tmp_dict[item_key] = sys[item_key]
-                spin_sampled.append(tmp_dict)
-            return spin_sampled
+            return [_pack_spin_stat_sample(self, sample) for sample in sampled_func()]
 
         self.backbone_model.compute_or_load_stat(
             spin_sampled_func,
@@ -556,27 +585,7 @@ def _get_spin_sampled_func(
 
         @functools.lru_cache
         def spin_sampled_func() -> list[dict]:
-            sampled = sampled_func()
-            spin_sampled = []
-            for sys in sampled:
-                coord_updated, atype_updated, _ = self.process_spin_input(
-                    sys["coord"], sys["atype"], sys["spin"]
-                )
-                tmp_dict = {
-                    "coord": coord_updated,
-                    "atype": atype_updated,
-                }
-                if "natoms" in sys:
-                    natoms = sys["natoms"]
-                    xp = array_api_compat.array_namespace(natoms)
-                    tmp_dict["natoms"] = xp.concat(
-                        [2 * natoms[:, :2], natoms[:, 2:], natoms[:, 2:]], axis=-1
-                    )
-                for item_key in sys.keys():
-                    if item_key not in ["coord", "atype", "spin", "natoms"]:
-                        tmp_dict[item_key] = sys[item_key]
-                spin_sampled.append(tmp_dict)
-            return spin_sampled
+            return [_pack_spin_stat_sample(self, sample) for sample in sampled_func()]
 
         return self.backbone_model.atomic_model._make_wrapped_sampler(spin_sampled_func)
 
@@ -682,6 +691,8 @@ def call_common(
             atomic parameter. nf x nloc x nda
         do_atomic_virial
             If calculate the atomic virial.
+        charge_spin
+            Frame-level charge and spin conditioning.
 
         Returns
         -------
@@ -720,6 +731,8 @@ def call_common(
             model_output_type.pop(model_output_type.index("mask"))
         var_name = model_output_type[0]
         model_ret[f"{var_name}"] = model_ret[f"{var_name}"][:, :nloc]
+        if "mask" in model_ret:
+            model_ret["mask"] = model_ret["mask"][:, :nloc]
         if (
             self.backbone_model.do_grad_r(var_name)
             and model_ret.get(f"{var_name}_derv_r") is not None
@@ -786,6 +799,8 @@ def call(
             atomic parameter. nf x nloc x nda
         do_atomic_virial
             If calculate the atomic virial.
+        charge_spin
+            Frame-level charge and spin conditioning.
 
         Returns
         -------
@@ -813,6 +828,8 @@ def call(
         model_predict[var_name] = model_ret[f"{var_name}_redu"]
         if "mask_mag" in model_ret:
             model_predict["mask_mag"] = model_ret["mask_mag"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
         if (
             self.backbone_model.do_grad_r(var_name)
             and model_ret.get(f"{var_name}_derv_r") is not None
@@ -866,6 +883,10 @@ def call_common_lower(
             atomic parameter. nf x nloc x nda
         do_atomic_virial
             whether calculate atomic virial
+        comm_dict
+            Optional MPI communication metadata for parallel inference.
+        charge_spin
+            Frame-level charge and spin conditioning.
 
         Returns
         -------
@@ -902,6 +923,8 @@ def call_common_lower(
             model_output_type.pop(model_output_type.index("mask"))
         var_name = model_output_type[0]
         model_ret[f"{var_name}"] = model_ret[f"{var_name}"][:, :nloc]
+        if "mask" in model_ret:
+            model_ret["mask"] = model_ret["mask"][:, :nloc]
         if (
             self.backbone_model.do_grad_r(var_name)
             and model_ret.get(f"{var_name}_derv_r") is not None
@@ -974,6 +997,8 @@ def call_lower(
             atomic parameter. nf x nloc x nda
         do_atomic_virial
             whether calculate atomic virial
+        charge_spin
+            Frame-level charge and spin conditioning.
 
         Returns
         -------
@@ -1002,6 +1027,8 @@ def call_lower(
         model_predict[var_name] = model_ret[f"{var_name}_redu"]
         if "mask_mag" in model_ret:
             model_predict["extended_mask_mag"] = model_ret["mask_mag"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
         if (
             self.backbone_model.do_grad_r(var_name)
             and model_ret.get(f"{var_name}_derv_r") is not None
@@ -1050,4 +1077,6 @@ def translated_output_def(self) -> dict[str, Any]:
             output_def["virial"].squeeze(-2)
             output_def["atom_virial"] = deepcopy(out_def_data[f"{var_name}_derv_c"])
             output_def["atom_virial"].squeeze(-2)
+        if "mask" in out_def_data:
+            output_def["mask"] = out_def_data["mask"]
         return output_def
diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py
index 3e439f173e..1f22690d1a 100644
--- a/deepmd/dpmodel/utils/__init__.py
+++ b/deepmd/dpmodel/utils/__init__.py
@@ -10,11 +10,11 @@
     PairExcludeMask,
 )
 from .lmdb_data import (
-    DistributedSameNlocBatchSampler,
+    DistributedLmdbBatchSampler,
+    LmdbBatchSampler,
     LmdbDataReader,
     LmdbTestData,
     LmdbTestDataNlocView,
-    SameNlocBatchSampler,
     is_lmdb,
     make_neighbor_stat_data,
 )
@@ -79,11 +79,12 @@
 __all__ = [
     "AtomExcludeMask",
     "DefaultNeighborList",
-    "DistributedSameNlocBatchSampler",
+    "DistributedLmdbBatchSampler",
     "EmbeddingNet",
     "EnvMat",
     "FittingNet",
     "GraphLayout",
+    "LmdbBatchSampler",
     "LmdbDataReader",
     "LmdbTestData",
     "LmdbTestDataNlocView",
@@ -93,7 +94,6 @@
     "NeighborList",
     "NetworkCollection",
     "PairExcludeMask",
-    "SameNlocBatchSampler",
     "aggregate",
     "apply_pair_exclusion_nlist",
     "build_multiple_neighbor_list",
diff --git a/deepmd/dpmodel/utils/batch.py b/deepmd/dpmodel/utils/batch.py
index 2cbf8a72ff..77416d8dd6 100644
--- a/deepmd/dpmodel/utils/batch.py
+++ b/deepmd/dpmodel/utils/batch.py
@@ -11,7 +11,18 @@
 _DROP_KEYS = {"default_mesh", "sid", "fid"}
 
 # Keys that belong to model input (everything else is label).
-_INPUT_KEYS = {"coord", "atype", "spin", "box", "fparam", "aparam", "charge_spin"}
+# ``n_node`` is an input rather than a label: it states how a ragged batch's
+# flat node axis divides into frames, which the model needs to read it at all.
+_INPUT_KEYS = {
+    "coord",
+    "atype",
+    "spin",
+    "box",
+    "fparam",
+    "aparam",
+    "charge_spin",
+    "n_node",
+}
 
 
 def normalize_batch(batch: dict[str, Any]) -> dict[str, Any]:
diff --git a/deepmd/dpmodel/utils/exclude_mask.py b/deepmd/dpmodel/utils/exclude_mask.py
index 1cc4ea7479..849bd2c2e9 100644
--- a/deepmd/dpmodel/utils/exclude_mask.py
+++ b/deepmd/dpmodel/utils/exclude_mask.py
@@ -5,6 +5,7 @@
 
 from deepmd.dpmodel.array_api import (
     Array,
+    xp_asarray_nodetach,
     xp_take_along_axis,
     xp_take_first_n,
 )
@@ -55,7 +56,11 @@ def build_type_exclude_mask(
         lead = atype.shape  # (nf, natom) dense | (N,) graph
         return xp.reshape(
             xp.take(
-                xp.asarray(self.type_mask[...], device=array_api_compat.device(atype)),
+                xp_asarray_nodetach(
+                    xp,
+                    self.type_mask[...],
+                    device=array_api_compat.device(atype),
+                ),
                 xp.reshape(atype, (-1,)),
                 axis=0,
             ),
@@ -151,7 +156,11 @@ def build_type_exclude_mask(
         type_ij_flat = xp.reshape(type_ij, (-1,))
         mask = xp.reshape(
             xp.take(
-                xp.asarray(self.type_mask[...], device=array_api_compat.device(nlist)),
+                xp_asarray_nodetach(
+                    xp,
+                    self.type_mask[...],
+                    device=array_api_compat.device(nlist),
+                ),
                 type_ij_flat,
             ),
             (nf, nloc, nnei),
@@ -185,7 +194,11 @@ def build_edge_exclude_mask(self, edge_index: Array, atype: Array) -> Array:
         dst_t = xp.take(atype, edge_index[1, :], axis=0)
         type_ij = dst_t * (self.ntypes + 1) + src_t
         return xp.take(
-            xp.asarray(self.type_mask[...], device=array_api_compat.device(atype)),
+            xp_asarray_nodetach(
+                xp,
+                self.type_mask[...],
+                device=array_api_compat.device(atype),
+            ),
             type_ij,
             axis=0,
         )
diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py
index 5cb22ed304..0396640a4a 100644
--- a/deepmd/dpmodel/utils/lmdb_data.py
+++ b/deepmd/dpmodel/utils/lmdb_data.py
@@ -5,12 +5,17 @@
 Backend-specific wrappers (PyTorch Dataset, JAX, etc.) import from here.
 """
 
+import dataclasses
 import logging
 import math
 import multiprocessing
+import os
 import signal
 import threading
+import time
 from collections.abc import (
+    Callable,
+    Iterable,
     Iterator,
     Sequence,
 )
@@ -25,14 +30,12 @@
 from dataclasses import (
     dataclass,
 )
-from itertools import (
-    pairwise,
-)
 from pathlib import (
     Path,
 )
 from typing import (
     Any,
+    cast,
 )
 
 import lmdb
@@ -49,10 +52,17 @@
 from deepmd.utils import random as dp_random
 from deepmd.utils.data import (
     DataRequirementItem,
+    DataRequirementSourcePolicy,
 )
 
 log = logging.getLogger(__name__)
 
+
+def _is_local_rank_zero() -> bool:
+    """Whether this process owns node-local operational logging."""
+    return int(os.environ.get("LOCAL_RANK", "0")) == 0
+
+
 # LMDB key → DeePMD convention
 _KEY_REMAP = {
     "coords": "coord",
@@ -67,20 +77,58 @@
 # (energy is set by Loss DataRequirementItem; reduce() also sets high_prec=True)
 _HIGH_PREC_KEYS = frozenset({"energy"})
 
-# Keys that describe frame geometry or LMDB bookkeeping rather than optional
-# model inputs/labels.  They must not participate in availability signatures.
-_STRUCTURAL_KEYS = frozenset(
+# Frames probed to decide whether availability-sensitive labels partition a
+# dataset. Requirements are registered after reader construction, so the probe
+# is deferred until a consumer first requests sampling groups.
+#
+# The sample makes the partition probabilistic, and a dataset whose odd frame
+# out falls between two probes is grouped as if it were uniform. What such a
+# frame costs is bounded: the batch it lands in reports the disputed label
+# unavailable, so at most ``batch_size - 1`` frames lose that label for that
+# batch, and a reshuffle moves it elsewhere next epoch. No default fill is
+# ever mistaken for a real label, which is the property that matters -- see
+# :func:`_batch_find_flags`.
+#
+# Raising this bound trades startup reads for a tighter guarantee, but never
+# reaches certainty. A dataset that really does interleave labels is better
+# fixed where it is written: recording a per-frame availability id in
+# ``__metadata__``, alongside the ``frame_nlocs`` and ``frame_system_ids``
+# already there, would make the partition exact and free.
+_AVAILABILITY_PROBE_FRAMES = 256
+_AVAILABILITY_SCAN_CHUNK = 4096
+_AVAILABILITY_LOG_SECONDS = 60.0
+
+# Atom type written into the padded slots of a mixed-nloc batch. A phantom
+# atom occupies a tensor slot but no physical site: the neighbor list gives it
+# no neighbors, the atomic model zeroes its output, and the loss masks it out.
+PHANTOM_ATOM_TYPE = -1
+
+# Fields whose leading axis is the atom axis, and which a mixed-nloc batch must
+# therefore pad to the batch-wide atom count. Membership cannot be inferred
+# from array shapes: a frame with ``nloc == 9`` makes ``virial`` (shape ``(9,)``)
+# indistinguishable from a per-atom field, and ``nloc == 2`` does the same for
+# a two-component ``fparam``. The registered data requirements carry the
+# authoritative ``atomic`` flag; these two sets cover the fields that exist
+# without one.
+_STRUCTURAL_PER_ATOM_KEYS = frozenset({"coord", "atype"})
+_OPTIONAL_PER_ATOM_KEYS = frozenset({"aparam", "spin"})
+
+# Frame-level fields that a decoded frame may carry without a registered
+# requirement. They anchor the shape-based fallback below, which classifies any
+# remaining unrecognized field by comparing its leading axis to the frame's
+# atom count.
+_FRAME_LEVEL_KEYS = frozenset(
     {
-        "coord",
         "box",
-        "atype",
+        "energy",
+        "virial",
+        "fparam",
+        "charge_spin",
         "natoms",
         "real_natoms_vec",
-        "fid",
+        "min_pair_dist",
     }
 )
-_LMDB_METADATA_KEYS = frozenset({"atom_numbs", "atom_names", "orig"})
-_OPTIONAL_MODEL_INPUT_KEYS = frozenset({"fparam", "aparam", "spin", "charge_spin"})
 
 # Process-level cache: python-lmdb does not allow opening the same path twice
 # in one process.  We ref-count so the Environment is closed (and freed from
@@ -88,7 +136,7 @@
 _ENV_CACHE: dict[str, tuple[lmdb.Environment, int]] = {}
 
 
-def _open_lmdb(path: str) -> lmdb.Environment:
+def _open_lmdb(path: str, *, sequential: bool = False) -> lmdb.Environment:
     """Open (or reuse) a readonly LMDB environment with reference counting.
 
     The python-lmdb binding raises ``lmdb.Error`` if the same path is opened
@@ -96,6 +144,27 @@ def _open_lmdb(path: str) -> lmdb.Environment:
     and bump a reference count.  Call :func:`_close_lmdb` when done to
     decrement the count; when it reaches zero the environment is closed and
     removed from the cache.
+
+    Parameters
+    ----------
+    path : str
+        Path to the LMDB directory.
+
+    sequential : bool, optional
+        Whether the caller reads frames in ascending key order. Kernel
+        readahead is then left on, which turns the many small faults such a
+        walk would make into few large reads. Measured against NFS: 45 000
+        frames/s with readahead against 1 300 without. A caller reading in
+        shuffled order leaves this false, where readahead instead fetches
+        neighbours it will not use and costs about 12% of the throughput.
+
+        One path admits one environment, so a path already open keeps the
+        setting its first caller asked for.
+
+    Returns
+    -------
+    lmdb.Environment
+        The shared read-only environment for ``path``.
     """
     resolved = str(Path(path).resolve())
     entry = _ENV_CACHE.get(resolved)
@@ -103,7 +172,9 @@ def _open_lmdb(path: str) -> lmdb.Environment:
         env, refcount = entry
         _ENV_CACHE[resolved] = (env, refcount + 1)
         return env
-    env = lmdb.open(path, readonly=True, lock=False, readahead=False, meminit=False)
+    env = lmdb.open(
+        path, readonly=True, lock=False, readahead=sequential, meminit=False
+    )
     _ENV_CACHE[resolved] = (env, 1)
     return env
 
@@ -133,6 +204,43 @@ def _read_metadata(txn: lmdb.Transaction) -> dict:
     return msgpack.unpackb(raw, raw=False)
 
 
+def _read_metadata_of(path: str) -> dict:
+    """Read ``__metadata__`` through an environment that keeps readahead on.
+
+    The value is a single contiguous run of overflow pages -- hundreds of
+    megabytes once the writer records a per-frame table. ``MDB_NORDAHEAD``,
+    which a reader of shuffled frames asks :func:`_open_lmdb` for, advises
+    the kernel ``MADV_RANDOM`` over the whole map, and that turns this
+    sequential run into one synchronous fault per 4 KiB page. Measured
+    against NFS on a 665 MiB table: 379 s without readahead, 8.2 s with it.
+
+    The environment is closed before the caller opens its own, because
+    python-lmdb refuses a second open of one path. A path already open for
+    frame serving is read through that environment instead, its pages being
+    warm by then in the only case that matters.
+
+    Parameters
+    ----------
+    path : str
+        Path to the LMDB directory.
+
+    Returns
+    -------
+    dict
+        The decoded metadata mapping.
+    """
+    entry = _ENV_CACHE.get(str(Path(path).resolve()))
+    if entry is not None:
+        with entry[0].begin() as transaction:
+            return _read_metadata(transaction)
+    env = lmdb.open(path, readonly=True, lock=False, readahead=True, meminit=False)
+    try:
+        with env.begin() as transaction:
+            return _read_metadata(transaction)
+    finally:
+        env.close()
+
+
 def _decode_array(obj: dict, *, copy: bool = True) -> np.ndarray:
     """Reconstruct ndarray from msgpack-encoded dict with {type, shape, data}.
 
@@ -180,36 +288,373 @@ def _decode_frame(
     return result
 
 
+def _canonical_lmdb_key(key: str) -> str:
+    """Return one LMDB field name in the in-memory DeePMD convention."""
+    key = _KEY_REMAP.get(key, key)
+    if key.startswith("find_atomic_"):
+        return "find_atom_" + key.removeprefix("find_atomic_")
+    if key.startswith("atomic_"):
+        return "atom_" + key.removeprefix("atomic_")
+    return key
+
+
 def _remap_keys(frame: dict[str, Any]) -> dict[str, Any]:
     """Remap LMDB key names to the canonical in-memory DeePMD convention."""
-    out = {}
-    for k, v in frame.items():
-        key = _KEY_REMAP.get(k, k)
-        if key.startswith("find_atomic_"):
-            key = "find_atom_" + key.removeprefix("find_atomic_")
-        elif key.startswith("atomic_"):
-            key = "atom_" + key.removeprefix("atomic_")
-        out[key] = v
-    return out
+    return {_canonical_lmdb_key(key): value for key, value in frame.items()}
 
 
-def _availability_signature_keys(
-    frame: dict[str, Any], requirement_keys: Iterator[str]
-) -> list[str]:
-    """Return data keys whose availability can affect frame collation.
+def _requirement_value(requirement: Any, key: str, default: Any) -> Any:
+    """Read one requirement attribute from object or dictionary form."""
+    if isinstance(requirement, dict):
+        return requirement.get(key, default)
+    return getattr(requirement, key, default)
+
+
+def _requirement_source_policy(requirement: Any) -> DataRequirementSourcePolicy:
+    """Return the normalized source-presence policy of one requirement."""
+    policy = str(_requirement_value(requirement, "source_policy", "tracked"))
+    if policy not in {"tracked", "default", "derived"}:
+        raise ValueError(f"Unsupported data requirement source policy {policy!r}")
+    return cast("DataRequirementSourcePolicy", policy)
+
+
+def _requirement_is_mandatory(requirement: Any) -> bool:
+    """Whether unavailable source data violates one requirement."""
+    return bool(_requirement_value(requirement, "must", False))
+
+
+def _availability_requirement_keys(
+    requirements: dict[str, Any],
+) -> tuple[str, ...]:
+    """Return requirements whose source presence affects consumer behavior.
+
+    Only optional ``tracked`` fields participate. Mandatory fields fail at
+    decode time rather than forming a default-filled group. ``default`` and
+    ``derived`` fields have valid per-frame fallbacks and therefore combine
+    safely regardless of source presence.
+
+    Parameters
+    ----------
+    requirements : dict[str, Any]
+        Data requirements keyed by normalized field name.
 
-    In addition to registered requirements and standard optional model inputs,
-    include every label-like field that :class:`LmdbDataReader` exposes from
-    the raw frame.  This keeps sampler/validation grouping consistent with the
-    complete set of ``find_*`` flags checked during collation.
+    Returns
+    -------
+    tuple[str, ...]
+        Sorted availability-sensitive field names.
+    """
+    return tuple(
+        sorted(
+            key
+            for key, requirement in requirements.items()
+            if _requirement_source_policy(requirement) == "tracked"
+            and not _requirement_is_mandatory(requirement)
+        )
+    )
+
+
+def _raw_frame_availability(
+    raw: bytes,
+    key_bits: dict[str, int],
+) -> int:
+    """Return one undecoded frame's selected availability bit mask.
+
+    A selected field is available when the frame carries it unless an
+    explicit ``find_*`` flag says otherwise. Fields outside
+    ``key_bits`` are deliberately ignored, so auxiliary raw data cannot alter
+    a training run's batch partition.
+
+    Only the msgpack map is walked; the arrays it describes stay encoded.
+
+    Parameters
+    ----------
+    raw : bytes
+        The msgpack payload of one LMDB frame.
+    key_bits : dict[str, int]
+        Bit assigned to each normalized availability-sensitive field.
+
+    Returns
+    -------
+    int
+        Availability mask in ``key_bits`` bit positions.
+    """
+    frame = msgpack.unpackb(raw, raw=False)
+    present = 0
+    explicit_known = 0
+    explicit_true = 0
+    for key, value in frame.items():
+        name = _canonical_lmdb_key(key)
+        if name.startswith("find_"):
+            label = name.removeprefix("find_")
+            bit = key_bits.get(label)
+            if bit is not None:
+                explicit_known |= bit
+                if float(np.asarray(_decode_value(value)).item()) != 0.0:
+                    explicit_true |= bit
+        else:
+            bit = key_bits.get(name)
+            if bit is not None:
+                present |= bit
+    return present & (~explicit_known | explicit_true)
+
+
+def _evenly_spaced(keys: Sequence[int], count: int) -> list[int]:
+    """Pick up to ``count`` entries spread evenly over ``keys``.
+
+    Spreading rather than truncating matters for a dataset assembled by
+    concatenating sources: the tail of the key space is as likely to hold
+    the odd one out as the head.
+
+    Parameters
+    ----------
+    keys : Sequence[int]
+        Integer LMDB frame keys to sample from.
+    count : int
+        Upper bound on the number of keys returned.
+
+    Returns
+    -------
+    list[int]
+        The sampled keys, in ascending position order.
     """
-    keys = set(requirement_keys) | set(_OPTIONAL_MODEL_INPUT_KEYS)
-    for frame_key in frame:
-        if frame_key.startswith("find_"):
-            keys.add(frame_key.removeprefix("find_"))
-        elif frame_key not in _STRUCTURAL_KEYS | _LMDB_METADATA_KEYS:
-            keys.add(frame_key)
-    return sorted(keys)
+    total = len(keys)
+    if total <= count:
+        return [int(key) for key in keys]
+    return [int(keys[total * position // count]) for position in range(count)]
+
+
+def _probe_uniform_availability(
+    transaction: lmdb.Transaction,
+    keys: Iterable[int],
+    frame_format: str,
+    availability_keys: Sequence[str],
+) -> bool:
+    """Whether the sampled frames all supply the same labels.
+
+    Parameters
+    ----------
+    transaction : lmdb.Transaction
+        Open read transaction on the LMDB environment.
+    keys : Iterable[int]
+        Integer LMDB frame keys to probe, already reduced to a bounded
+        sample by :func:`_evenly_spaced`.
+    frame_format : str
+        Format specification for integer LMDB frame keys.
+    availability_keys : Sequence[str]
+        Normalized field names used by the active consumer.
+
+    Returns
+    -------
+    bool
+        True when every probed frame supplies the same labels. A dataset
+        larger than the sample may still be mixed, in which case the batch
+        decode reports the disputed label unavailable rather than mixing it.
+    """
+    key_bits = {key: 1 << position for position, key in enumerate(availability_keys)}
+    reference: int | None = None
+    for key in keys:
+        raw = transaction.get(format(int(key), frame_format).encode())
+        if raw is None:
+            continue
+        availability = _raw_frame_availability(raw, key_bits)
+        if reference is None:
+            reference = availability
+        elif availability != reference:
+            return False
+    return True
+
+
+@dataclass(frozen=True)
+class _AvailabilityIndex:
+    """Compact availability signature ID aligned with a frame index domain."""
+
+    ids: np.ndarray
+    signature_count: int
+
+    def groups(
+        self,
+        indices: np.ndarray,
+        *,
+        positions: np.ndarray | None = None,
+    ) -> list[np.ndarray]:
+        """Partition indices by cached signature IDs without source reads."""
+        index_array = np.asarray(indices)
+        signature_ids = (
+            self.ids
+            if positions is None
+            else self.ids[np.asarray(positions, dtype=np.int64)]
+        )
+        if len(index_array) != len(signature_ids):
+            raise ValueError(
+                "availability index and frame indices have different lengths: "
+                f"{len(signature_ids)} != {len(index_array)}"
+            )
+        if len(index_array) == 0:
+            return []
+        if self.signature_count == 1:
+            return [index_array]
+
+        if self.signature_count <= 8:
+            groups: list[np.ndarray] = []
+            for signature_id in range(self.signature_count):
+                mask = signature_ids == signature_id
+                if np.any(mask):
+                    groups.append(index_array[mask])
+            return groups
+
+        order = np.argsort(signature_ids, kind="stable")
+        ordered_ids = signature_ids[order]
+        cuts = np.flatnonzero(ordered_ids[1:] != ordered_ids[:-1]) + 1
+        ordered_indices = index_array[order]
+        return list(np.split(ordered_indices, cuts))
+
+
+def _widen_signature_ids(ids: np.ndarray, written: int) -> np.ndarray:
+    """Return the next wider unsigned ID buffer, preserving written entries."""
+    widths = {
+        np.dtype(np.uint8): np.dtype(np.uint16),
+        np.dtype(np.uint16): np.dtype(np.uint32),
+        np.dtype(np.uint32): np.dtype(np.uint64),
+    }
+    target_dtype = widths.get(ids.dtype)
+    if target_dtype is None:
+        raise OverflowError("LMDB availability signatures exceed uint64 capacity")
+    widened = np.empty(ids.shape, dtype=target_dtype)
+    widened[:written] = ids[:written]
+    return widened
+
+
+def _scan_availability_index(
+    frame_count: int,
+    read_raw: Callable[[int], bytes | None],
+    availability_keys: Sequence[str],
+    dataset: str,
+) -> _AvailabilityIndex:
+    """Build compact signature IDs with one exact source scan.
+
+    The scan stores one unsigned integer per frame and one dictionary entry per
+    distinct signature. It never accumulates frame indices as Python objects.
+
+    Parameters
+    ----------
+    frame_count : int
+        Number of positions in the index domain.
+    read_raw : Callable[[int], bytes or None]
+        Function returning the encoded frame for one domain position.
+    availability_keys : Sequence[str]
+        Normalized field names defining the partition.
+    dataset : str
+        Dataset path included in progress messages.
+
+    Returns
+    -------
+    _AvailabilityIndex
+        Compact signature IDs aligned with domain positions.
+    """
+    report_progress = _is_local_rank_zero()
+    if report_progress:
+        log.info(
+            "LMDB label-availability scan started: dataset=%s, frames=%d, labels=%s",
+            dataset,
+            frame_count,
+            list(availability_keys),
+        )
+
+    key_bits = {key: 1 << position for position, key in enumerate(availability_keys)}
+    signature_ids = np.empty(frame_count, dtype=np.uint8)
+    signature_map: dict[int, int] = {}
+    next_log = time.monotonic() + _AVAILABILITY_LOG_SECONDS if report_progress else 0.0
+    for start in range(0, frame_count, _AVAILABILITY_SCAN_CHUNK):
+        stop = min(start + _AVAILABILITY_SCAN_CHUNK, frame_count)
+        for position in range(start, stop):
+            raw = read_raw(position)
+            if raw is None:
+                raise RuntimeError(
+                    f"LMDB frame at position {position} is missing from {dataset}"
+                )
+            signature = _raw_frame_availability(raw, key_bits)
+            signature_id = signature_map.get(signature)
+            if signature_id is None:
+                signature_id = len(signature_map)
+                signature_map[signature] = signature_id
+                if signature_id > np.iinfo(signature_ids.dtype).max:
+                    signature_ids = _widen_signature_ids(signature_ids, position)
+            signature_ids[position] = signature_id
+
+        if report_progress:
+            now = time.monotonic()
+            if stop < frame_count and now >= next_log:
+                log.info(
+                    "LMDB label-availability scan progress: dataset=%s, "
+                    "frames=%d/%d (%.1f%%)",
+                    dataset,
+                    stop,
+                    frame_count,
+                    100.0 * stop / frame_count,
+                )
+                next_log = now + _AVAILABILITY_LOG_SECONDS
+
+    if report_progress:
+        log.info(
+            "LMDB label-availability scan completed: dataset=%s, frames=%d, "
+            "groups=%d, index_dtype=%s",
+            dataset,
+            frame_count,
+            len(signature_map),
+            signature_ids.dtype,
+        )
+    return _AvailabilityIndex(signature_ids, len(signature_map))
+
+
+def _scan_lmdb_path_sequential(
+    lmdb_path: str,
+    availability_keys: Sequence[str],
+    log_level: int,
+) -> _AvailabilityIndex:
+    """Scan a complete LMDB under a sequential-readahead environment."""
+    logging.basicConfig(level=log_level)
+    logging.getLogger().setLevel(log_level)
+    environment = lmdb.open(
+        lmdb_path,
+        readonly=True,
+        lock=False,
+        readahead=True,
+        meminit=False,
+    )
+    try:
+        with environment.begin() as transaction:
+            metadata = _read_metadata(transaction)
+        frame_count, frame_format, _natoms_per_type = _parse_metadata(metadata)
+        with environment.begin() as transaction:
+
+            def read_raw(position: int) -> bytes | None:
+                return transaction.get(format(position, frame_format).encode())
+
+            return _scan_availability_index(
+                frame_count,
+                read_raw,
+                availability_keys,
+                lmdb_path,
+            )
+    finally:
+        environment.close()
+
+
+def _scan_lmdb_path_in_worker(
+    lmdb_path: str,
+    availability_keys: Sequence[str],
+) -> _AvailabilityIndex:
+    """Run a sequential scan outside a process holding this LMDB open."""
+    executor = _create_lmdb_executor(1)
+    try:
+        return executor.submit(
+            _scan_lmdb_path_sequential,
+            lmdb_path,
+            availability_keys,
+            log.getEffectiveLevel(),
+        ).result()
+    finally:
+        executor.shutdown(wait=True, cancel_futures=True)
 
 
 def _remap_atom_types(atype: np.ndarray, type_remap: np.ndarray) -> np.ndarray:
@@ -243,12 +688,15 @@ class LmdbDecodeConfig:
         Optional LMDB-type to model-type lookup table.
     data_requirements
         Registered data requirements keyed by field name.
+    dataset
+        Dataset identifier used in frame-level diagnostics.
     """
 
     ntypes: int
     natoms: int
     type_remap: np.ndarray | None
     data_requirements: dict[str, Any]
+    dataset: str = "<unknown LMDB>"
 
 
 def _requirement_dtype(requirement: Any) -> np.dtype:
@@ -276,6 +724,104 @@ def _resolve_frame_dtype(config: LmdbDecodeConfig, key: str) -> np.dtype:
     return np.dtype(GLOBAL_NP_FLOAT_PRECISION)
 
 
+def _requirement_is_atomic(requirement: Any) -> bool:
+    """Whether a data requirement describes a per-atom quantity."""
+    if isinstance(requirement, dict):
+        return bool(requirement.get("atomic", False))
+    return bool(getattr(requirement, "atomic", False))
+
+
+def _frame_source_available(frame: dict[str, Any], key: str) -> bool:
+    """Resolve one frame's explicit or inferred source availability."""
+    value = frame.get(key)
+    source_present = _is_encoded_array(value) or isinstance(
+        value, (np.ndarray, np.generic, int, float, bool)
+    )
+    find_key = f"find_{key}"
+    if find_key in frame:
+        return source_present and bool(
+            float(np.asarray(_decode_value(frame[find_key])).item())
+        )
+    return source_present
+
+
+def _raise_if_mandatory_unavailable(
+    frame: dict[str, Any],
+    key: str,
+    requirement: Any,
+    source_available: bool,
+    *,
+    dataset: str,
+    frame_index: int,
+) -> None:
+    """Reject unavailable mandatory data at the shared frame boundary."""
+    if not _requirement_is_mandatory(requirement) or source_available:
+        return
+    find_key = f"find_{key}"
+    if find_key in frame:
+        find_value = float(np.asarray(_decode_value(frame[find_key])).item())
+        reason = (
+            f"explicit {find_key}=0"
+            if find_value == 0.0
+            else f"field is absent or invalid despite {find_key}={find_value:g}"
+        )
+    else:
+        reason = "field is absent or invalid"
+    raise RuntimeError(
+        f"Required LMDB field {key!r} is unavailable in frame {frame_index} "
+        f"of {dataset}: {reason}."
+    )
+
+
+def resolve_per_atom_keys(
+    frame: dict[str, Any],
+    config: LmdbDecodeConfig,
+) -> frozenset[str]:
+    """Return the fields of one frame whose leading axis is the atom axis.
+
+    Classification is authoritative wherever possible: coordinates and atom
+    types are per-atom by construction, and every registered data requirement
+    declares whether it is ``atomic``. Only fields the loader has never been
+    told about fall back to comparing their leading axis against the atom
+    count, and the frame-level fields that DeePMD itself produces are excluded
+    from that fallback so a coincidental shape match cannot misclassify them.
+
+    Parameters
+    ----------
+    frame : dict[str, Any]
+        One decoded frame in DeePMD data-system convention. ``coord`` anchors
+        the atom axis: it is the one field present in every frame whose
+        leading axis is the atom count.
+    config : LmdbDecodeConfig
+        Decoder state holding the registered data requirements.
+
+    Returns
+    -------
+    frozenset[str]
+        Names of the fields to pad along their leading axis.
+    """
+    nloc = frame["coord"].shape[0]
+    keys = set(_STRUCTURAL_PER_ATOM_KEYS)
+    keys |= {
+        key
+        for key, requirement in config.data_requirements.items()
+        if _requirement_is_atomic(requirement)
+    }
+    keys |= _OPTIONAL_PER_ATOM_KEYS
+    for key, value in frame.items():
+        if (
+            key in keys
+            or key in _FRAME_LEVEL_KEYS
+            or key in config.data_requirements
+            or key.startswith("find_")
+            or key == "fid"
+        ):
+            continue
+        if isinstance(value, np.ndarray) and value.ndim >= 1 and value.shape[0] == nloc:
+            keys.add(key)
+    return frozenset(keys & frame.keys())
+
+
 def _compute_frame_natoms(atype: np.ndarray, ntypes: int) -> np.ndarray:
     """Build ``[nloc, nloc, count(type_0), ...]`` for one frame.
 
@@ -294,6 +840,34 @@ def _compute_frame_natoms(atype: np.ndarray, ntypes: int) -> np.ndarray:
     return natoms
 
 
+def _resolve_derived_requirement(
+    frame: dict[str, Any],
+    key: str,
+    requirement: Any,
+    config: LmdbDecodeConfig,
+) -> None:
+    """Resolve one derived field from normalized structural frame data."""
+    if key != "min_pair_dist":
+        raise ValueError(f"Unsupported derived LMDB field {key!r}")
+
+    coord = frame.get("coord")
+    atype = frame.get("atype")
+    if not isinstance(coord, np.ndarray) or not isinstance(atype, np.ndarray):
+        frame.pop(key, None)
+        frame[f"find_{key}"] = np.float32(0.0)
+        return
+
+    box = frame.get("box")
+    if box is not None and np.allclose(box, 0.0):
+        box = None
+    threshold = float(_requirement_value(requirement, "default", 0.0))
+    frame[key] = np.array(
+        [compute_min_pair_dist_single(coord, box, atype, stop_below=threshold)],
+        dtype=_resolve_frame_dtype(config, key),
+    )
+    frame[f"find_{key}"] = np.float32(1.0)
+
+
 def decode_lmdb_frame(
     raw: bytes,
     original_key: int,
@@ -379,34 +953,9 @@ def decode_lmdb_frame(
     frame["real_natoms_vec"] = natoms
 
     requirements = config.data_requirements
-    coord = frame.get("coord")
-    if (
-        "min_pair_dist" in requirements
-        and "min_pair_dist" not in frame
-        and isinstance(coord, np.ndarray)
-        and isinstance(atype, np.ndarray)
-    ):
-        box = frame.get("box")
-        if box is not None and np.allclose(box, 0.0):
-            box = None
-        requirement = requirements["min_pair_dist"]
-        default = (
-            requirement.get("default", 0.0)
-            if isinstance(requirement, dict)
-            else getattr(requirement, "default", 0.0)
-        )
-        frame["find_min_pair_dist"] = np.float32(1.0)
-        frame["min_pair_dist"] = np.array(
-            [
-                compute_min_pair_dist_single(
-                    coord,
-                    box,
-                    atype,
-                    stop_below=float(default),
-                )
-            ],
-            dtype=_resolve_frame_dtype(config, "min_pair_dist"),
-        )
+    for key, requirement in requirements.items():
+        if _requirement_source_policy(requirement) == "derived":
+            _resolve_derived_requirement(frame, key, requirement, config)
 
     structural_keys = frozenset(
         {
@@ -435,16 +984,28 @@ def decode_lmdb_frame(
             atomic = requirement.atomic
             repeat = getattr(requirement, "repeat", 1)
         dtype = _requirement_dtype(requirement)
+        find_key = f"find_{key}"
+        source_available = _frame_source_available(frame, key)
+        _raise_if_mandatory_unavailable(
+            frame,
+            key,
+            requirement,
+            source_available,
+            dataset=config.dataset,
+            frame_index=original_key,
+        )
 
-        if key not in frame:
-            frame[f"find_{key}"] = np.float32(0.0)
+        if not source_available:
+            frame[find_key] = np.float32(
+                1.0 if _requirement_source_policy(requirement) == "default" else 0.0
+            )
             shape = (frame_natoms, ndof) if atomic else (ndof,)
             data = np.full(shape, default, dtype=dtype)
             if repeat != 1:
                 data = np.repeat(data, repeat).reshape(-1)
             frame[key] = data
         else:
-            frame.setdefault(f"find_{key}", np.float32(1.0))
+            frame.setdefault(find_key, np.float32(1.0))
             if repeat != 1 and isinstance(frame[key], np.ndarray):
                 frame[key] = (
                     np.repeat(frame[key], repeat).reshape(-1).astype(dtype, copy=False)
@@ -460,11 +1021,147 @@ def decode_lmdb_frame(
     return frame
 
 
+def _pad_fill_value(key: str) -> int:
+    """Return the value written into the padded tail of one per-atom field.
+
+    Atom types use the phantom sentinel so that downstream code recognizes the
+    slot as unoccupied; every other per-atom field is zeroed, which keeps
+    padded rows neutral in the sums and masked reductions that consume them.
+    """
+    return PHANTOM_ATOM_TYPE if key == "atype" else 0
+
+
+def per_atom_strides(
+    frame: dict[str, Any],
+    per_atom_keys: frozenset[str],
+) -> dict[str, int]:
+    """Return how many leading-axis entries each per-atom field spends per atom.
+
+    Most per-atom fields carry one row per atom, so their leading axis is the
+    atom count itself. A data requirement declared with ``repeat != 1`` is
+    instead stored flat and atom-major, giving a leading axis of
+    ``nloc * repeat``. Padding widens that axis by whole atoms either way, so
+    the factor is all the padding logic needs to tell the two layouts apart.
+
+    Parameters
+    ----------
+    frame : dict[str, Any]
+        One decoded frame; ``coord`` anchors the atom count.
+    per_atom_keys : frozenset[str]
+        Fields to measure, as resolved by :func:`resolve_per_atom_keys`.
+
+    Returns
+    -------
+    dict[str, int]
+        Leading-axis entries per atom, keyed by field name.
+
+    Raises
+    ------
+    ValueError
+        If a field's leading axis is not a whole multiple of the atom count.
+    """
+    nloc = frame["coord"].shape[0]
+    strides: dict[str, int] = {}
+    for key in per_atom_keys:
+        length = np.asarray(frame[key]).shape[0]
+        if nloc == 0 or length % nloc:
+            raise ValueError(
+                f"LMDB field {key!r} has a leading axis of {length}, which is "
+                f"not a whole number of entries per atom in a {nloc}-atom frame"
+            )
+        strides[key] = length // nloc
+    return strides
+
+
+@dataclass(frozen=True)
+class BatchLayout:
+    """Where each frame's per-atom rows sit on a decoded batch's leading axis.
+
+    Two layouts serve the two shapes a model's node axis can take.
+
+    The **rectangular** layout gives every frame a row of the batch-wide atom
+    count and pads the tail of the shorter ones with phantom atoms, which is
+    what a model reading an ``(nf, nloc, ...)`` node axis requires. The
+    **ragged** layout concatenates the frames instead, so nothing is padded and
+    the leading axis is the batch's real atom count; a model reading a flat
+    node axis consumes that directly, paired with ``n_node``. Frame-level
+    fields are stacked on the frame axis either way.
+
+    Attributes
+    ----------
+    n_node : numpy.ndarray
+        Real atom count of each frame of this decode, with shape ``(nf,)``.
+    strides : dict[str, int]
+        Leading-axis entries per atom of each per-atom field, as resolved by
+        :func:`per_atom_strides`.
+    ragged : bool
+        Whether frames are concatenated rather than padded to a common width.
+    width : int
+        Atoms each frame occupies under the rectangular layout. It stays the
+        width of the whole batch even where ``n_node`` covers one chunk of it,
+        since chunks that padded to their own widths would not concatenate.
+    """
+
+    n_node: np.ndarray
+    strides: dict[str, int]
+    ragged: bool
+    width: int
+
+    def __post_init__(self) -> None:
+        # Frame offsets on the ragged axis, in atoms. Held rather than summed
+        # per lookup, since every field of every frame asks for one.
+        object.__setattr__(
+            self, "_offset", np.concatenate([[0], np.cumsum(self.n_node)])
+        )
+
+    @classmethod
+    def over(
+        cls, n_node: np.ndarray, strides: dict[str, int], *, ragged: bool
+    ) -> "BatchLayout":
+        """Return the layout of a batch holding the given per-frame counts."""
+        return cls(
+            n_node=n_node,
+            strides=strides,
+            ragged=ragged,
+            width=int(n_node.max()) if n_node.size else 0,
+        )
+
+    def chunk(self, start: int, stop: int) -> "BatchLayout":
+        """Return the layout of a contiguous run of this batch's frames."""
+        return dataclasses.replace(self, n_node=self.n_node[start:stop])
+
+    def field_length(self, key: str) -> int:
+        """Return the leading-axis length one per-atom field is allocated."""
+        stride = self.strides[key]
+        if self.ragged:
+            return int(self.n_node.sum()) * stride
+        return self.width * stride
+
+    def frame_index(self, row: int, key: str) -> Any:
+        """Return the index selecting one frame's rows of a per-atom field.
+
+        Rectangular batches carry a frame axis, so the rows of frame ``row``
+        are the head of its own row; ragged batches carry none, and the rows
+        are a run at the frame's offset.
+        """
+        stride = self.strides[key]
+        length = int(self.n_node[row]) * stride
+        if not self.ragged:
+            return (row, slice(0, length))
+        start = int(self._offset[row]) * stride
+        return slice(start, start + length)
+
+
 def _allocate_lmdb_batch(
     frame: dict[str, Any],
     batch_size: int,
+    layout: BatchLayout,
 ) -> dict[str, Any]:
-    """Allocate a contiguous NumPy batch from the first decoded frame."""
+    """Allocate a contiguous NumPy batch from the first decoded frame.
+
+    Per-atom fields are allocated at the length :class:`BatchLayout` gives
+    them: one padded row per frame, or one concatenated run over the batch.
+    """
     batch: dict[str, Any] = {}
     for key, value in frame.items():
         if key.startswith("find_"):
@@ -476,6 +1173,20 @@ def _allocate_lmdb_batch(
             continue
         elif value is None:
             batch[key] = None
+        elif key in layout.strides:
+            array = np.asarray(value)
+            head = (
+                (layout.field_length(key),)
+                if layout.ragged
+                else (batch_size, layout.field_length(key))
+            )
+            destination = np.full(
+                (*head, *array.shape[1:]),
+                _pad_fill_value(key),
+                dtype=array.dtype,
+            )
+            destination[layout.frame_index(0, key)] = array
+            batch[key] = destination
         else:
             array = np.asarray(value)
             destination = np.empty((batch_size, *array.shape), dtype=array.dtype)
@@ -484,24 +1195,92 @@ def _allocate_lmdb_batch(
     return batch
 
 
+def _promote_batch_field(
+    batch: dict[str, Any],
+    field: str,
+    layout: BatchLayout,
+    frames_written: int,
+    dtype: np.dtype,
+) -> np.ndarray:
+    """Widen one batch field's dtype in place, preserving what was written.
+
+    The replacement is prefilled with the field's padding value rather than
+    zeroed, so that entries not yet written keep the marker the allocation gave
+    them; for ``atype`` under the rectangular layout that marker is what
+    identifies a phantom atom.
+    """
+    destination = batch[field]
+    promoted = np.full(destination.shape, _pad_fill_value(field), dtype=dtype)
+    # The leading axis counts frames, except for a per-atom field of a ragged
+    # batch, where it counts the atom rows those frames have filled.
+    if field in layout.strides and layout.ragged:
+        written = int(layout._offset[frames_written]) * layout.strides[field]
+    else:
+        written = frames_written
+    promoted[:written] = destination[:written]
+    batch[field] = promoted
+    return promoted
+
+
 def decode_lmdb_batch(
     transaction: lmdb.Transaction,
     original_keys: Sequence[int],
     frame_format: str,
     config: LmdbDecodeConfig,
+    layout: BatchLayout | None = None,
 ) -> dict[str, Any]:
     """Decode LMDB records directly into preallocated contiguous arrays.
 
     The function keeps at most one temporary frame alive. It avoids the
     decode-copy, dtype-copy, Python frame-list, and final ``numpy.stack``
     sequence used by generic collation.
+
+    Parameters
+    ----------
+    transaction : lmdb.Transaction
+        Open read transaction on the LMDB environment.
+    original_keys : Sequence[int]
+        Integer LMDB frame keys in batch order.
+    frame_format : str
+        Format specification for integer LMDB frame keys.
+    config : LmdbDecodeConfig
+        Decoder state independent of the LMDB environment.
+    layout : BatchLayout, optional
+        Where each frame's per-atom rows belong. Defaults to a rectangular
+        layout at the atom count of the first frame, which leaves a batch of
+        uniform atom count untouched.
+
+    Returns
+    -------
+    dict[str, Any]
+        One collated batch of contiguous NumPy arrays. A ragged layout adds
+        ``n_node``, the per-frame atom count its flat axis is read with.
+
+    Notes
+    -----
+    The layout fixes the shape of every field of the result, so a chunked
+    decode must pass each chunk the layout of its own frames, cut from the
+    batch-wide one; :meth:`LmdbDataReader.batch_layout` resolves that once.
+
+    Frames need not expose the same optional fields. A field only some of
+    them carry is left out of the batch, and the matching ``find_*`` flag
+    reports the label unavailable, which matches what
+    :func:`collate_lmdb_frames` produces for the same frames. Registered
+    data requirements are exempt: :func:`decode_lmdb_frame` puts them on
+    every frame.
     """
     if not original_keys:
         raise ValueError("decode_lmdb_batch requires at least one frame key")
 
     batch: dict[str, Any] | None = None
     batch_size = len(original_keys)
-    expected_fields: frozenset[str] | None = None
+    # Frames need not expose the same optional fields. Counting the frames
+    # each field appears on settles, once the batch is read, which fields
+    # have a complete column to stack and which labels the batch may report
+    # as available; see :func:`_batch_find_flags` for the same rule applied
+    # to the generic collation path.
+    field_counts: dict[str, int] = {}
+    find_flags: dict[str, bool] = {}
     for row, original_key in enumerate(original_keys):
         key = format(int(original_key), frame_format).encode()
         raw = transaction.get(key)
@@ -513,47 +1292,83 @@ def decode_lmdb_batch(
             config,
             copy_arrays=False,
         )
+        frame_nloc = frame["coord"].shape[0]
+        if layout is None:
+            layout = BatchLayout.over(
+                np.full(batch_size, frame_nloc, dtype=np.int64),
+                per_atom_strides(frame, resolve_per_atom_keys(frame, config)),
+                ragged=False,
+            )
+        if int(layout.n_node[row]) != frame_nloc:
+            raise ValueError(
+                f"the batch layout gives frame {original_key} "
+                f"{int(layout.n_node[row])} atoms, but it holds {frame_nloc}"
+            )
+        for field in frame:
+            field_counts[field] = field_counts.get(field, 0) + 1
+
         if batch is None:
-            batch = _allocate_lmdb_batch(frame, batch_size)
-            expected_fields = frozenset(frame)
+            batch = _allocate_lmdb_batch(frame, batch_size, layout)
+            find_flags = {
+                field: float(value) != 0.0
+                for field, value in frame.items()
+                if field.startswith("find_")
+            }
             continue
 
-        frame_fields = frozenset(frame)
-        if frame_fields != expected_fields:
-            raise ValueError(
-                "LMDB frames in one same-nloc batch expose inconsistent fields: "
-                f"frame {original_keys[0]} has {sorted(expected_fields)}, while "
-                f"frame {original_key} has {sorted(frame_fields)}"
-            )
         for field, value in frame.items():
             if field.startswith("find_"):
-                if not np.array_equal(batch[field], value):
-                    raise ValueError(
-                        f"LMDB field availability changes within one batch: "
-                        f"{field!r} differs at frame {original_key}"
-                    )
+                find_flags[field] = find_flags.get(field, False) and float(value) != 0.0
                 continue
-            if field == "type" or value is None:
+            if field == "type" or value is None or field not in batch:
                 continue
             if field == "fid":
                 batch[field][row] = value
-            else:
-                destination = batch[field]
-                array = np.asarray(value)
-                if destination.shape[1:] != array.shape:
+                continue
+            destination = batch[field]
+            array = np.asarray(value)
+            stride = layout.strides.get(field)
+            # A per-atom field may differ in its leading axis, which the layout
+            # absorbs as long as the axis stays a whole number of atoms; every
+            # remaining axis must match exactly.
+            if stride is not None:
+                lead_axes = 1 if layout.ragged else 2
+                expected_tail: tuple[int, ...] = destination.shape[lead_axes:]
+                actual_tail = array.shape[1:]
+                if array.shape[0] != frame_nloc * stride:
                     raise ValueError(
-                        f"LMDB field {field!r} changes shape within one batch: "
-                        f"expected {destination.shape[1:]}, got {array.shape} "
-                        f"for frame {original_key}"
+                        f"LMDB field {field!r} spends {array.shape[0]} leading "
+                        f"entries on {frame_nloc} atoms in frame {original_key}, "
+                        f"against {stride} per atom in frame {original_keys[0]}"
                     )
-                result_dtype = np.result_type(destination.dtype, array.dtype)
-                if result_dtype != destination.dtype:
-                    promoted = np.empty(destination.shape, dtype=result_dtype)
-                    promoted[:row] = destination[:row]
-                    batch[field] = destination = promoted
+            else:
+                expected_tail = destination.shape[1:]
+                actual_tail = array.shape
+            if expected_tail != actual_tail:
+                raise ValueError(
+                    f"LMDB field {field!r} changes shape within one batch: "
+                    f"expected {expected_tail}, got {actual_tail} "
+                    f"for frame {original_key}"
+                )
+            result_dtype = np.result_type(destination.dtype, array.dtype)
+            if result_dtype != destination.dtype:
+                destination = _promote_batch_field(
+                    batch, field, layout, row, result_dtype
+                )
+            if stride is not None:
+                destination[layout.frame_index(row, field)] = array
+            else:
                 destination[row] = array
 
-    assert batch is not None
+    assert batch is not None and layout is not None
+    for field, present in find_flags.items():
+        available = present and field_counts.get(field, 0) == batch_size
+        batch[field] = np.float32(1.0 if available else 0.0)
+    for field, count in field_counts.items():
+        if count < batch_size and not field.startswith("find_"):
+            batch.pop(field, None)
+    if layout.ragged:
+        batch["n_node"] = layout.n_node
     batch["sid"] = np.asarray([0], dtype=np.int64)
     return batch
 
@@ -569,8 +1384,14 @@ def _decode_lmdb_worker_chunk(
     frame_format: str,
     config: LmdbDecodeConfig,
     original_keys: list[int],
+    layout: BatchLayout,
 ) -> dict[str, Any]:
-    """Decode one chunk using process-local LMDB state."""
+    """Decode one chunk using process-local LMDB state.
+
+    ``layout`` is this chunk's slice of the batch layout decided by the parent
+    process. Deriving it there rather than per chunk is what lets
+    :func:`_merge_lmdb_chunks` concatenate the results.
+    """
     reader = _WORKER_LMDB_READERS.get(lmdb_path)
     if reader is None:
         environment = lmdb.open(
@@ -587,43 +1408,43 @@ def _decode_lmdb_worker_chunk(
         original_keys,
         frame_format,
         config,
+        layout,
     )
 
 
 def _merge_lmdb_chunks(chunks: list[dict[str, Any]]) -> dict[str, Any]:
-    """Merge ordered worker chunks into one contiguous batch."""
+    """Merge ordered worker chunks into one contiguous batch.
+
+    Each chunk has already reduced its own frames, so merging repeats that
+    reduction one level up: a label is available to the merged batch only
+    where every chunk reports it, and a field only some chunks carry has no
+    complete column to concatenate.
+    """
     if not chunks:
         raise ValueError("cannot merge an empty LMDB chunk list")
     if len(chunks) == 1:
         return chunks[0]
 
     first = chunks[0]
-    expected_fields = frozenset(first)
-    for chunk_index, chunk in enumerate(chunks[1:], start=1):
-        chunk_fields = frozenset(chunk)
-        if chunk_fields != expected_fields:
-            raise ValueError(
-                "LMDB worker chunks expose inconsistent fields: "
-                f"chunk 0 has {sorted(expected_fields)}, while chunk "
-                f"{chunk_index} has {sorted(chunk_fields)}"
-            )
-
     merged: dict[str, Any] = {}
     for key, value in first.items():
         if key.startswith("find_"):
-            for chunk_index, chunk in enumerate(chunks[1:], start=1):
-                if not np.array_equal(value, chunk[key]):
-                    raise ValueError(
-                        "LMDB field availability changes across worker chunks: "
-                        f"{key!r} differs in chunk {chunk_index}"
-                    )
-            merged[key] = value
+            available = float(value) != 0.0 and all(
+                key in chunk and float(chunk[key]) != 0.0 for chunk in chunks[1:]
+            )
+            merged[key] = np.float32(1.0 if available else 0.0)
         elif key == "sid" or value is None:
             merged[key] = value
+        elif any(key not in chunk for chunk in chunks[1:]):
+            continue
         elif key == "fid":
             merged[key] = [frame_id for chunk in chunks for frame_id in chunk[key]]
         else:
             merged[key] = np.concatenate([chunk[key] for chunk in chunks], axis=0)
+    for chunk in chunks[1:]:
+        for key in chunk:
+            if key.startswith("find_"):
+                merged.setdefault(key, np.float32(0.0))
     return merged
 
 
@@ -904,16 +1725,20 @@ def _iter_epoch(self) -> Iterator[list[int]]:
         return iter(self._sampler)
 
     def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]:
-        """Submit one batch as balanced contiguous chunks."""
+        """Submit one batch as balanced contiguous chunks.
+
+        The batch layout is resolved here rather than per chunk, so that every
+        chunk decodes to the same field shapes and the results concatenate.
+        """
         original_keys = self._reader.original_keys(indices)
+        layout = self._reader.batch_layout(indices)
         workers = min(self._num_workers, len(original_keys))
         base_size, remainder = divmod(len(original_keys), workers)
-        chunks: list[list[int]] = []
+        chunks: list[tuple[list[int], BatchLayout]] = []
         start = 0
         for worker_index in range(workers):
-            chunk_size = base_size + int(worker_index < remainder)
-            stop = start + chunk_size
-            chunks.append(original_keys[start:stop])
+            stop = start + base_size + int(worker_index < remainder)
+            chunks.append((original_keys[start:stop], layout.chunk(start, stop)))
             start = stop
         decode_config = self._reader.worker_decode_config()
         return [
@@ -923,8 +1748,9 @@ def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]:
                 self._reader.frame_format,
                 decode_config,
                 chunk,
+                chunk_layout,
             )
-            for chunk in chunks
+            for chunk, chunk_layout in chunks
         ]
 
     def _worth_decoding_in_parallel(self, indices: list[int]) -> bool:
@@ -1019,6 +1845,38 @@ def _scan_frame_nlocs(
     return nlocs
 
 
+def _group_positions_by_value(values: np.ndarray) -> dict[int, np.ndarray]:
+    """Group the positions of ``values`` by the value each holds.
+
+    Sorting once and cutting at the value boundaries keeps the work inside
+    NumPy. Accumulating Python lists instead costs one interpreter iteration
+    and one boxed integer per element, which at the 10^8 frames a large LMDB
+    holds is minutes of runtime and tens of gigabytes of resident memory.
+
+    Parameters
+    ----------
+    values : numpy.ndarray
+        Integer values to group by, with shape ``(n,)``.
+
+    Returns
+    -------
+    dict[int, numpy.ndarray]
+        Value → the ascending positions holding it. Each array is a slice of
+        one shared buffer, so the mapping costs one extra array in total.
+    """
+    if values.size == 0:
+        return {}
+    order = np.argsort(values, kind="stable")
+    ordered = values[order]
+    cuts = np.flatnonzero(ordered[1:] != ordered[:-1]) + 1
+    starts = np.concatenate(([0], cuts))
+    stops = np.concatenate((cuts, [values.size]))
+    return {
+        int(ordered[start]): order[start:stop]
+        for start, stop in zip(starts.tolist(), stops.tolist(), strict=True)
+    }
+
+
 def _compute_batch_size(nloc: int, rule: int) -> int:
     """Compute batch_size for a given nloc using the auto rule."""
     bsi = rule // max(nloc, 1)
@@ -1027,6 +1885,11 @@ def _compute_batch_size(nloc: int, rule: int) -> int:
     return max(bsi, 1)
 
 
+#: Seed of the representative shuffle behind :attr:`LmdbDataReader.total_batch`.
+#: Fixed so that the reported count is reproducible across calls and processes.
+_TOTAL_BATCH_SEED = 0
+
+
 def _parse_positive_rule(spec: str, prefix: str) -> int:
     """Parse the ``N`` in ``<prefix>N`` and require ``N > 0``.
 
@@ -1056,16 +1919,17 @@ class LmdbDataReader:
     Reads LMDB frames and returns dicts of numpy arrays.
     Backend-specific Dataset classes (PyTorch, JAX, etc.) wrap this.
 
-    Datasets are typically mixed-nloc (frames with different atom counts).
-    The ``mixed_batch`` flag controls batching strategy:
+    An LMDB typically holds frames of many different atom counts. The
+    ``batch_size`` rule decides how those frames are grouped:
 
-    - ``mixed_batch=False`` (default, old format): each batch contains only
-      frames with the same nloc. A ``SameNlocBatchSampler`` groups frames
-      by nloc and yields same-nloc batches. Auto batch_size is computed
-      per-nloc-group.
-    - ``mixed_batch=True`` (new format): frames with different nloc can
-      coexist in one batch (requires padding + mask in collate_fn).
-      Currently raises ``NotImplementedError`` at collation time.
+    - Every rule except ``"mix:N"`` keeps a batch homogeneous in atom count,
+      so no padding is ever needed, whichever layout it is decoded in.
+    - ``"mix:N"`` allows one batch to span several atom counts. A consumer
+      reading a flat node axis takes such a batch concatenated; one reading
+      an ``(nf, nloc, ...)`` axis takes it padded to the batch-wide maximum,
+      the padded rows carrying ``atype = -1`` so that the neighbor list, the
+      atomic model and the loss all skip them. The choice is
+      :meth:`use_ragged_batches`.
 
     Parameters
     ----------
@@ -1090,9 +1954,16 @@ class LmdbDataReader:
         - ``"filter:N"``: same per-nloc formula as ``"max:N"`` **and**
           drops every frame whose ``nloc > N`` from the dataset. By
           construction every retained batch has at most ``N`` atoms.
-    mixed_batch : bool
-        If True, allow different nloc in the same batch (future).
-        If False (default), enforce same-nloc-per-batch.
+        - ``"mix:N"``: mixed-nloc batching with an atom-axis budget of
+          ``N``. Frames of different atom counts share a batch, whose atom
+          axis holds at most ``N`` entries: the total atom count under the
+          ragged layout, the padded ``nframes * max_nloc`` under the
+          rectangular one. This is the natural extension of ``"max:N"``:
+          the bound is on the decoded batch, and a lone frame with
+          ``nloc > N`` still forms a batch of its own. Filling batches
+          rather than cutting them at atom-count boundaries also keeps the
+          frame-level loss terms closer to the weighting an atom budget asks
+          for than ``"max:N"`` manages; see :func:`_chop_mixed_nloc`.
     """
 
     def __init__(
@@ -1100,15 +1971,14 @@ def __init__(
         lmdb_path: str,
         type_map: list[str],
         batch_size: int | str = "auto",
-        mixed_batch: bool = False,
     ) -> None:
         self.lmdb_path = str(Path(lmdb_path).resolve())
         self._type_map = type_map
+        # Read before opening the frame-serving environment, which disables
+        # the readahead this one large sequential value depends on. Training
+        # draws frames in shuffled order and so leaves it disabled.
+        meta = _read_metadata_of(self.lmdb_path)
         self._env = _open_lmdb(self.lmdb_path)
-        self.mixed_batch = mixed_batch
-
-        with self._env.begin() as txn:
-            meta = _read_metadata(txn)
 
         self.nframes, self._frame_fmt, self._natoms_per_type = _parse_metadata(meta)
         self._natoms = sum(self._natoms_per_type)
@@ -1139,25 +2009,32 @@ def __init__(
         self._txn = self._env.begin()
         self._closed = False
 
-        # Scan per-frame nloc only when needed for same-nloc batching.
-        # For mixed_batch=True, skip the scan entirely (future: padding handles it).
+        # Per-frame atom counts drive every batching rule: same-nloc grouping,
+        # the ``filter:N`` drop, and the atom-axis layout of a ``mix:N`` batch.
         # ``orig_frame_nlocs`` / ``orig_frame_system_ids`` are indexed by the
         # *original* LMDB frame index. After a potential ``filter:N`` drop we
         # rebuild ``self._frame_nlocs`` / ``self._frame_system_ids`` so they
         # are parallel arrays over the *dataset* index space (0..len(self));
         # the dataset-to-original mapping lives in ``self._retained_keys``.
-        if not mixed_batch:
-            # Fast path: use pre-computed frame_nlocs from metadata if available.
-            # Falls back to scanning each frame's atom_types shape (~10 us/frame).
-            meta_nlocs = meta.get("frame_nlocs")
-            if meta_nlocs is not None:
-                orig_frame_nlocs = [int(n) for n in meta_nlocs]
-            else:
-                orig_frame_nlocs = _scan_frame_nlocs(
-                    self._env, self.nframes, self._frame_fmt, self._natoms
-                )
+        # Metadata carries the counts when the writer recorded them; otherwise
+        # each frame's atom_types shape is scanned (~10 us/frame).
+        # Both tables are held as NumPy arrays rather than Python lists: at
+        # 10^8 frames a list of boxed integers costs tens of gigabytes, while
+        # the arrays cost 4 bytes an entry and let the grouping below run as
+        # array operations. The list msgpack unpacks is released immediately.
+        meta_nlocs = meta.get("frame_nlocs")
+        if meta_nlocs is not None:
+            orig_frame_nlocs = np.fromiter(
+                meta_nlocs, dtype=np.int32, count=len(meta_nlocs)
+            )
         else:
-            orig_frame_nlocs = []
+            orig_frame_nlocs = np.asarray(
+                _scan_frame_nlocs(
+                    self._env, self.nframes, self._frame_fmt, self._natoms
+                ),
+                dtype=np.int32,
+            )
+        del meta_nlocs
 
         # Parse frame_system_ids for auto_prob support. ``_nsystems`` must stay
         # at ``max(original_sid) + 1`` even after filter:N so that user-facing
@@ -1165,18 +2042,22 @@ def __init__(
         # keeps its meaning across filter thresholds.
         meta_sys_ids = meta.get("frame_system_ids")
         if meta_sys_ids is not None:
-            orig_frame_system_ids: list[int] | None = [int(s) for s in meta_sys_ids]
-            self._nsystems = max(orig_frame_system_ids) + 1
+            orig_frame_system_ids: np.ndarray | None = np.fromiter(
+                meta_sys_ids, dtype=np.int32, count=len(meta_sys_ids)
+            )
+            self._nsystems = int(orig_frame_system_ids.max()) + 1
         else:
             orig_frame_system_ids = None
             self._nsystems = 1
+        del meta_sys_ids, meta
 
-        # Parse batch_size spec. ``auto_rule`` and ``max_rule`` are mutually
-        # exclusive; ``filter_rule`` implies ``max_rule`` plus dropping frames
-        # whose nloc exceeds the threshold.
+        # Parse batch_size spec. ``auto_rule``, ``max_rule`` and ``mix_rule``
+        # are mutually exclusive; ``filter_rule`` implies ``max_rule`` plus
+        # dropping frames whose nloc exceeds the threshold.
         self._auto_rule: int | None = None
         self._max_rule: int | None = None
         self._filter_rule: int | None = None
+        self._mix_rule: int | None = None
         if isinstance(batch_size, str):
             if batch_size == "auto":
                 self._auto_rule = 32
@@ -1187,108 +2068,110 @@ def __init__(
             elif batch_size.startswith("filter:"):
                 self._filter_rule = _parse_positive_rule(batch_size, "filter:")
                 self._max_rule = self._filter_rule
+            elif batch_size.startswith("mix:"):
+                self._mix_rule = _parse_positive_rule(batch_size, "mix:")
             else:
                 raise ValueError(
-                    f"Unsupported batch_size {batch_size!r}. "
-                    "Expected int, 'auto', 'auto:N', 'max:N', or 'filter:N'."
+                    f"Unsupported batch_size {batch_size!r}. Expected int, "
+                    "'auto', 'auto:N', 'max:N', 'filter:N', or 'mix:N'."
                 )
 
-        # ``filter:N`` needs per-frame nloc to drop oversized frames; the
-        # ``mixed_batch=True`` fast path skips the nloc scan entirely, so the
-        # two options are incompatible. Fail fast rather than silently
-        # retaining every frame and breaking the documented contract.
-        if self._filter_rule is not None and mixed_batch:
-            raise ValueError(
-                "batch_size='filter:N' is incompatible with mixed_batch=True: "
-                "per-frame nloc is unavailable in the mixed-batch fast path. "
-                "Use mixed_batch=False, or switch to 'max:N' / a fixed int."
-            )
-
         # Determine which original-index frames survive the filter. Without
-        # ``filter:N`` every frame is retained.
+        # ``filter:N`` every frame is retained, and the identity mapping is
+        # left as the ``arange`` so no frame is copied through a mask.
+        retained_keys = np.arange(self.nframes, dtype=np.int64)
         if self._filter_rule is not None:
-            retained_keys = [
-                i for i, n in enumerate(orig_frame_nlocs) if n <= self._filter_rule
-            ]
-            n_dropped = self.nframes - len(retained_keys)
+            n_dropped = int(np.count_nonzero(orig_frame_nlocs > self._filter_rule))
             if n_dropped > 0:
+                retained_keys = retained_keys[orig_frame_nlocs <= self._filter_rule]
                 log.info(
                     f"LMDB filter:{self._filter_rule} drops {n_dropped}/"
                     f"{self.nframes} frames with nloc > {self._filter_rule} "
                     f"({self.lmdb_path})."
                 )
-        else:
-            retained_keys = list(range(self.nframes))
 
         # Dataset-index → original LMDB frame key. ``__getitem__`` looks up
         # this table so that ``reader[i]`` is a valid LMDB read for every
         # ``0 <= i < len(reader)``, no matter how many frames were filtered.
-        self._retained_keys: list[int] = retained_keys
+        self._retained_keys: np.ndarray = retained_keys
 
         # Re-key _frame_nlocs / _frame_system_ids into the dataset-index
         # space so that every downstream consumer (nloc_groups, system_groups,
-        # SameNlocBatchSampler, _expand_indices_by_blocks) operates in a
-        # single, self-consistent indexing scheme.
-        if not mixed_batch:
-            self._frame_nlocs = [orig_frame_nlocs[k] for k in retained_keys]
-        else:
-            self._frame_nlocs = []
+        # LmdbBatchSampler, _expand_indices_by_blocks) operates in a single,
+        # self-consistent indexing scheme.
+        keys_are_identity = retained_keys.size == self.nframes
+        self._keys_are_identity = keys_are_identity
+        self._frame_nlocs = (
+            orig_frame_nlocs if keys_are_identity else orig_frame_nlocs[retained_keys]
+        )
 
-        if orig_frame_system_ids is not None:
-            self._frame_system_ids: list[int] | None = [
-                orig_frame_system_ids[k] for k in retained_keys
-            ]
+        if orig_frame_system_ids is None:
+            self._frame_system_ids: np.ndarray | None = None
+        elif keys_are_identity:
+            self._frame_system_ids = orig_frame_system_ids
         else:
-            self._frame_system_ids = None
+            self._frame_system_ids = orig_frame_system_ids[retained_keys]
 
-        # Group retained frames by nloc using dataset indices (0..len-1).
-        if not mixed_batch:
-            self._nloc_groups: dict[int, list[int]] = {}
-            for ds_idx, nloc in enumerate(self._frame_nlocs):
-                self._nloc_groups.setdefault(nloc, []).append(ds_idx)
-        else:
-            self._nloc_groups = {}
+        # nframes now reflects retained frames; __len__ returns this and the
+        # valid index domain for __getitem__ is [0, self.nframes).
+        self.nframes = int(retained_keys.size)
 
-        # Group retained frames by original system id; the sid numbering is
-        # preserved (no compression) so user-facing auto_prob slices stay
-        # meaningful across filter thresholds. Fully-dropped systems appear
-        # as zero-frame entries in ``_system_nframes``.
+        # Group retained frames by nloc using dataset indices (0..len-1).
+        # Statistics collection consumes these groups in every batching mode,
+        # because per-nloc groups are the largest units that stack without
+        # padding.
+        self._nloc_groups = _group_positions_by_value(self._frame_nlocs)
+
+        # Frames per original system id; the sid numbering is preserved (no
+        # compression) so user-facing auto_prob slices stay meaningful across
+        # filter thresholds. Fully-dropped systems count zero. The per-system
+        # index lists themselves are built only if asked for; see
+        # :attr:`system_groups`.
+        self._system_groups: dict[int, np.ndarray] | None = None
         if self._frame_system_ids is not None:
-            self._system_groups: dict[int, list[int]] = {}
-            for ds_idx, sid in enumerate(self._frame_system_ids):
-                self._system_groups.setdefault(sid, []).append(ds_idx)
-            self._system_nframes: list[int] = [
-                len(self._system_groups.get(i, [])) for i in range(self._nsystems)
-            ]
+            self._system_nframes: list[int] = np.bincount(
+                self._frame_system_ids, minlength=self._nsystems
+            ).tolist()
         else:
-            self._system_groups = {0: list(range(len(retained_keys)))}
-            self._system_nframes = [len(retained_keys)]
+            self._system_nframes = [self.nframes]
 
-        # nframes now reflects retained frames; __len__ returns this and the
-        # valid index domain for __getitem__ is [0, self.nframes).
-        self.nframes = len(retained_keys)
-
-        # Default batch_size used only by the index/total_batch estimate. The
-        # sampler always goes through get_batch_size_for_nloc for real batches.
+        # Nominal batch size, reported to callers that want a single number.
+        # The sampler never uses it: same-nloc modes go through
+        # get_batch_size_for_nloc, and ``mix:N`` sizes each batch by budget.
+        mean_nloc = (
+            float(self._frame_nlocs.mean()) if self._frame_nlocs.size else self._natoms
+        )
         if self._auto_rule is not None:
             self.batch_size = _compute_batch_size(self._natoms, self._auto_rule)
         elif self._max_rule is not None:
             self.batch_size = max(1, self._max_rule // max(self._natoms, 1))
+        elif self._mix_rule is not None:
+            self.batch_size = max(1, int(self._mix_rule / max(mean_nloc, 1.0)))
         else:
             self.batch_size = int(batch_size)
 
         # Data requirements tracking
         self._data_requirements: dict[str, DataRequirementItem] = {}
         self._data_requirements_frozen = False
+        self._data_requirements_revision = 0
+        # Requirements arrive after reader construction. Availability remains
+        # unresolved until the first grouping request so raw fields unused by
+        # the run cause no I/O and cannot influence the partition.
+        self._uniform_availability: bool | None = None
+        self._availability_index: _AvailabilityIndex | None = None
         self._decode_config = LmdbDecodeConfig(
             ntypes=self._ntypes,
             natoms=self._natoms,
             type_remap=self._type_remap,
             data_requirements=self._data_requirements,
+            dataset=self.lmdb_path,
         )
-        # Availability signatures are decoded lazily and reused by every
-        # sampler epoch. Registering new requirements invalidates the cache.
-        self._find_signature_cache: dict[int, tuple[tuple[str, bool], ...]] = {}
+        # Which fields carry an atom axis follows from the registered
+        # requirements, so this cache is invalidated when they change.
+        self._per_atom_strides: dict[str, int] | None = None
+        # Batches are rectangular until a consumer that reads a flat node axis
+        # asks otherwise; see :meth:`use_ragged_batches`.
+        self._ragged_batches = False
 
     def _resolve_dtype(self, key: str) -> np.dtype:
         """Resolve the target numpy dtype for a given key.
@@ -1335,12 +2218,18 @@ def get_batch_size_for_nloc(self, nloc: int) -> int:
         - ``filter:N``: same per-nloc formula as ``max:N``; by
           construction every retained group satisfies ``nloc <= N`` so
           no overshoot occurs.
+        - ``mix:N``: ``max(1, floor(N / nloc))``, the count a batch would
+          hold were every one of its frames this size. Training batches are
+          sized by budget instead; this value serves the per-nloc statistics
+          groups, which stack without padding and therefore batch like
+          ``max:N``.
         - fixed int: the same value for every nloc group.
         """
         if self._auto_rule is not None:
             return _compute_batch_size(nloc, self._auto_rule)
-        if self._max_rule is not None:
-            return max(1, self._max_rule // max(nloc, 1))
+        atom_budget = self._max_rule if self._max_rule is not None else self._mix_rule
+        if atom_budget is not None:
+            return max(1, atom_budget // max(nloc, 1))
         return self.batch_size
 
     def __len__(self) -> int:
@@ -1356,7 +2245,7 @@ def __getitem__(self, index: int) -> dict[str, Any]:
         self._data_requirements_frozen = True
         if index < 0 or index >= self.nframes:
             raise IndexError(f"dataset index {index} out of range [0, {self.nframes})")
-        original_key = self._retained_keys[index]
+        original_key = int(self._retained_keys[index])
         key = format(original_key, self._frame_fmt).encode()
         raw = self._transaction().get(key)
         if raw is None:
@@ -1379,17 +2268,134 @@ def original_keys(self, indices: Sequence[int]) -> list[int]:
                 raise IndexError(
                     f"dataset index {index} out of range [0, {self.nframes})"
                 )
-            keys.append(self._retained_keys[index])
+            keys.append(int(self._retained_keys[index]))
         return keys
 
-    def decode_batch(self, indices: Sequence[int]) -> dict[str, Any]:
-        """Decode a same-nloc batch directly into contiguous NumPy arrays."""
+    def batch_pad_nloc(self, indices: Sequence[int]) -> int:
+        """Return the atom count every frame of one batch is padded to.
+
+        Parameters
+        ----------
+        indices : Sequence[int]
+            Dataset indices forming one batch.
+
+        Returns
+        -------
+        int
+            The largest atom count in the batch. Batches drawn from a single
+            nloc group return that group's atom count, so padding is a no-op.
+        """
+        return int(self._frame_nlocs[np.asarray(indices, dtype=np.int64)].max())
+
+    def batch_layout(
+        self, indices: Sequence[int], *, ragged: bool | None = None
+    ) -> BatchLayout:
+        """Return where one batch's per-atom rows belong once decoded.
+
+        The layout fixes the shape of every field of the decoded batch, so a
+        decode split across worker processes resolves it here once and cuts a
+        chunk's share from it. Resolving it per chunk would let two chunks
+        disagree, both on the padded width and on the field classification,
+        which falls back to comparing a leading axis against the atom count of
+        whichever frame the chunk happens to start with.
+
+        Parameters
+        ----------
+        indices : Sequence[int]
+            Dataset indices forming one batch.
+        ragged : bool, optional
+            Layout to use, overriding the one configured for training batches.
+            Consumers with a layout of their own, such as statistics, name it
+            rather than inherit it.
+
+        Returns
+        -------
+        BatchLayout
+            The per-frame atom counts, the per-atom strides, and whether the
+            frames are concatenated or padded to a common width.
+        """
+        # Widened after gathering, and only to the batch's own size: a layout
+        # carries int64 counts wherever it is built, including in the decode
+        # workers a chunk of it is shipped to.
+        return BatchLayout.over(
+            self._frame_nlocs[np.asarray(indices, dtype=np.int64)].astype(np.int64),
+            self.per_atom_strides(),
+            ragged=self._ragged_batches if ragged is None else ragged,
+        )
+
+    @property
+    def ragged_batches(self) -> bool:
+        """Whether decoded batches concatenate their frames rather than pad."""
+        return self._ragged_batches
+
+    def use_ragged_batches(self, ragged: bool) -> None:
+        """Select the layout mixed-nloc training batches are delivered in.
+
+        The choice belongs to both the batching rule and the model. Only
+        ``mix:N`` may place different atom counts in one batch and therefore
+        needs a layout choice: a model reading a flat node axis takes those
+        frames concatenated, while one reading an ``(nf, nloc, ...)`` axis
+        needs them padded to a common width. Every other batching rule keeps
+        the established rectangular layout. Only the trainer sees both the
+        model and the data, so it requests the model-compatible layout once,
+        before training starts. Consumers with a layout of their own --
+        statistics, validation -- name theirs at the point of use and are
+        unaffected.
+
+        The layout also decides how the sampler packs frames, since padding is
+        what makes a batch's cost depend on its widest frame.
+
+        Parameters
+        ----------
+        ragged : bool
+            Whether the consumer accepts concatenated mixed-nloc frames.
+        """
+        self._ragged_batches = ragged and self.mixed_nloc
+
+    def per_atom_strides(self) -> dict[str, int]:
+        """Return the leading-axis entries per atom of each per-atom field.
+
+        Every frame of one LMDB exposes the same fields, so the classification
+        is a property of the dataset and its registered requirements rather
+        than of a batch, and is resolved from the first frame once.
+
+        Returns
+        -------
+        dict[str, int]
+            Entries per atom, keyed by field name.
+        """
+        if self._per_atom_strides is None:
+            frame = self[0]
+            self._per_atom_strides = per_atom_strides(
+                frame, resolve_per_atom_keys(frame, self._decode_config)
+            )
+        return self._per_atom_strides
+
+    def decode_batch(
+        self, indices: Sequence[int], *, ragged: bool | None = None
+    ) -> dict[str, Any]:
+        """Decode one batch directly into contiguous NumPy arrays.
+
+        Parameters
+        ----------
+        indices : Sequence[int]
+            Dataset indices forming one batch.
+        ragged : bool, optional
+            Layout to decode into, overriding the one configured for training
+            batches. See :meth:`batch_layout`.
+
+        Returns
+        -------
+        dict[str, Any]
+            One collated batch of contiguous NumPy arrays.
+        """
         self._data_requirements_frozen = True
         return decode_lmdb_batch(
             self._transaction(),
             self.original_keys(indices),
             self._frame_fmt,
             self._decode_config,
+            self.batch_layout(indices, ragged=ragged),
         )
 
     @property
@@ -1397,6 +2403,17 @@ def frame_format(self) -> str:
         """Format specification used for integer LMDB frame keys."""
         return self._frame_fmt
 
+    @property
+    def decode_config(self) -> LmdbDecodeConfig:
+        """Decoder state for in-process consumers.
+
+        The returned object shares the reader's live requirement mapping and
+        reading it does not freeze registration, unlike
+        :meth:`worker_decode_config`, which hands the state to another process
+        and so must fix it first.
+        """
+        return self._decode_config
+
     def worker_decode_config(self) -> LmdbDecodeConfig:
         """Freeze and return decoder state for worker serialization."""
         self._data_requirements_frozen = True
@@ -1409,73 +2426,146 @@ def closed(self) -> bool:
 
     # --- Data requirement interface ---
 
+    def _scan_exact_availability(
+        self,
+        availability_keys: Sequence[str],
+    ) -> _AvailabilityIndex:
+        """Build the exact index under sequential kernel readahead."""
+
+        def scan(transaction: lmdb.Transaction) -> _AvailabilityIndex:
+            def read_raw(position: int) -> bytes | None:
+                original_key = int(self._retained_keys[position])
+                return transaction.get(format(original_key, self._frame_fmt).encode())
+
+            return _scan_availability_index(
+                self.nframes,
+                read_raw,
+                availability_keys,
+                self.lmdb_path,
+            )
+
+        environment = self._env
+        if environment is None:
+            raise RuntimeError("cannot scan a closed LMDB reader")
+        if bool(environment.flags().get("readahead", False)):
+            return scan(self._transaction())
+
+        resolved = str(Path(self.lmdb_path).resolve())
+        cache_entry = _ENV_CACHE.get(resolved)
+        owns_environment_exclusively = (
+            cache_entry is not None
+            and cache_entry[0] is environment
+            and cache_entry[1] == 1
+        )
+        if not owns_environment_exclusively:
+            if _is_local_rank_zero():
+                log.info(
+                    "LMDB label-availability scan uses an isolated sequential "
+                    "reader because the random-read environment is shared: %s",
+                    self.lmdb_path,
+                )
+            source_index = _scan_lmdb_path_in_worker(
+                self.lmdb_path,
+                availability_keys,
+            )
+            if self._keys_are_identity:
+                return source_index
+            return _AvailabilityIndex(
+                source_index.ids[self._retained_keys],
+                source_index.signature_count,
+            )
+
+        transaction = self._txn
+        if transaction is not None:
+            transaction.abort()
+        self._txn = None
+        self._env = None
+        _close_lmdb(self.lmdb_path)
+        try:
+            sequential_environment = _open_lmdb(self.lmdb_path, sequential=True)
+            with sequential_environment.begin() as sequential_transaction:
+                return scan(sequential_transaction)
+        finally:
+            _close_lmdb(self.lmdb_path)
+            self._env = _open_lmdb(self.lmdb_path, sequential=False)
+            self._txn = self._env.begin()
+
     def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None:
-        """Register expected keys; missing keys get default fill + find_key=0.0."""
+        """Register the consumer data contract before resolving any frame."""
         if self._data_requirements_frozen:
             raise RuntimeError(
                 "LMDB data requirements must be registered before reading any frame"
             )
         for item in data_requirement:
             self._data_requirements[item["key"]] = item
-        self._find_signature_cache.clear()
+        self._data_requirements_revision += 1
+        self._uniform_availability = None
+        self._availability_index = None
+        self._per_atom_strides = None
 
-    def get_find_signature(self, index: int) -> tuple[tuple[str, bool], ...]:
-        """Return the scalar availability signature for one retained frame.
+    def availability_groups(self, indices: np.ndarray) -> list[np.ndarray]:
+        """Partition dataset indices into label-compatible groups.
 
-        The signature covers registered data requirements and optional model
-        inputs whose ``find_*`` flags are created by :meth:`__getitem__`.
-        Reading only msgpack keys avoids decoding large coordinate and label
-        arrays while the sampler partitions frames.
-        """
-        cached = self._find_signature_cache.get(index)
-        if cached is not None:
-            return cached
-        if index < 0 or index >= self.nframes:
-            raise IndexError(f"dataset index {index} out of range [0, {self.nframes})")
+        A ``find_*`` flag is one scalar per batch. Optional tracked labels
+        therefore remain homogeneous so present labels are not discarded.
+        Mandatory labels fail during frame decoding; default-backed and
+        derived fields do not participate.
 
-        original_key = self._retained_keys[index]
-        key = format(original_key, self._frame_fmt).encode()
-        raw = self._txn.get(key)
-        if raw is None:
-            raise IndexError(
-                f"Frame {original_key} not found in LMDB (dataset index {index})"
-            )
-        raw_frame = msgpack.unpackb(raw, raw=False)
-        frame = {_KEY_REMAP.get(name, name): value for name, value in raw_frame.items()}
-        signature_keys = _availability_signature_keys(
-            frame, iter(self._data_requirements)
-        )
-        signature = []
-        for data_key in signature_keys:
-            find_key = f"find_{data_key}"
-            if find_key in frame:
-                find_value = _decode_value(frame[find_key])
-                available = bool(float(np.asarray(find_value).item()))
-            elif data_key == "min_pair_dist" and data_key in self._data_requirements:
-                # __getitem__ computes this requirement when it is not stored.
-                available = True
-            else:
-                available = data_key in frame
-            signature.append((find_key, available))
+        The bounded uniformity probe is deferred until requirements have been
+        registered. A dataset that probes mixed builds one compact exact index
+        for the entire retained dataset. Statistics and every sampler epoch
+        reuse that index without further LMDB reads.
 
-        result = tuple(signature)
-        self._find_signature_cache[index] = result
-        return result
+        The probe reads a bounded sample, so "uniform" is a finding and not a
+        proof; :data:`_AVAILABILITY_PROBE_FRAMES` states what a missed frame
+        costs and how a dataset can settle the question exactly.
 
-    def group_indices_by_find_signature(
-        self, indices: list[int]
-    ) -> dict[tuple[tuple[str, bool], ...], list[int]]:
-        """Partition dataset indices into scalar-compatible label groups."""
-        groups: dict[tuple[tuple[str, bool], ...], list[int]] = {}
-        for index in indices:
-            groups.setdefault(self.get_find_signature(index), []).append(index)
-        return groups
+        Parameters
+        ----------
+        indices : numpy.ndarray
+            Dataset indices to partition.
+
+        Returns
+        -------
+        list[numpy.ndarray]
+            Non-empty index groups in a stable order, together covering
+            ``indices``.
+        """
+        if len(indices) == 0:
+            return []
+        availability_keys = _availability_requirement_keys(self._data_requirements)
+        if not availability_keys:
+            self._uniform_availability = True
+            return [indices]
+        if self._uniform_availability is None:
+            self._uniform_availability = _probe_uniform_availability(
+                self._transaction(),
+                _evenly_spaced(self._retained_keys, _AVAILABILITY_PROBE_FRAMES),
+                self._frame_fmt,
+                availability_keys,
+            )
+        if self._uniform_availability:
+            return [indices]
+        if self._availability_index is None:
+            self._availability_index = self._scan_exact_availability(
+                availability_keys,
+            )
+        index_array = np.asarray(indices)
+        return self._availability_index.groups(
+            index_array,
+            positions=index_array,
+        )
 
     @property
     def data_requirements(self) -> list[DataRequirementItem]:
         """Registered data requirements in insertion order."""
         return list(self._data_requirements.values())
 
+    @property
+    def data_requirements_revision(self) -> int:
+        """Monotonic revision used to invalidate dependent sampling plans."""
+        return self._data_requirements_revision
+
     def print_summary(self, name: str, prob: Any) -> None:
         """Print basic dataset info."""
         n_groups = len(self._nloc_groups)
@@ -1485,6 +2575,8 @@ def print_summary(self, name: str, prob: Any) -> None:
             bs_str = f"filter:{self._filter_rule}"
         elif self._max_rule is not None:
             bs_str = f"max:{self._max_rule}"
+        elif self._mix_rule is not None:
+            bs_str = f"mix:{self._mix_rule}"
         else:
             bs_str = str(self.batch_size)
 
@@ -1492,7 +2584,7 @@ def print_summary(self, name: str, prob: Any) -> None:
             f"LMDB {name}: {self.lmdb_path}, "
             f"{self.nframes} frames, {n_groups} nloc groups, "
             f"batch_size={bs_str}, "
-            f"mixed_batch={self.mixed_batch}"
+            f"mixed_nloc={self.mixed_nloc}"
         )
         # Print nloc groups in rows of ~10 for readability
         items = [
@@ -1515,18 +2607,39 @@ def index(self) -> list[int]:
 
     @property
     def total_batch(self) -> int:
-        if self.mixed_batch:
-            return math.ceil(self.nframes / self.batch_size) if self.nframes else 0
-        total = 0
-        for nloc, indices in collect_lmdb_sampling_groups(self):
-            bs = self.get_batch_size_for_nloc(nloc)
-            total += (len(indices) + bs - 1) // bs
-        return total
+        """Number of batches in one pass over the dataset.
+
+        Every batching rule but ``mix:N`` fixes the count independently of the
+        order frames are visited in. ``mix:N`` fills batches to an atom budget
+        instead, so its count follows the shuffle: a pass in dataset order
+        groups frames of one original system together, which are already close
+        in atom count, and needs measurably fewer batches than the shuffled
+        pass training actually performs. The count is therefore taken from a
+        shuffled pass under a fixed seed, which is drawn from the same
+        distribution as a training pass while staying reproducible. It remains
+        an estimate: consult the sampler that will actually be iterated when
+        the exact count matters, as the trainers do to derive an epoch length.
+        """
+        return len(LmdbBatchSampler(self, shuffle=True, seed=_TOTAL_BATCH_SEED))
 
     @property
     def batch_sizes(self) -> list[int]:
         return [self.batch_size]
 
+    @property
+    def mixed_nloc(self) -> bool:
+        """Whether one batch may span several atom counts."""
+        return self._mix_rule is not None
+
+    @property
+    def atom_budget(self) -> int | None:
+        """Atom-axis budget of a ``mix:N`` batch, or ``None`` in other modes.
+
+        The axis is measured in the layout the batch will be decoded in: real
+        atoms when the frames are concatenated, padded slots when they are not.
+        """
+        return self._mix_rule
+
     @property
     def mixed_type(self) -> bool:
         """LMDB datasets are always mixed_type (frames may have different compositions)."""
@@ -1538,13 +2651,13 @@ def type_map(self) -> list[str]:
         return self._type_map
 
     @property
-    def nloc_groups(self) -> dict[int, list[int]]:
-        """Nloc → list of frame indices."""
+    def nloc_groups(self) -> dict[int, np.ndarray]:
+        """Atom count → the ascending dataset indices holding it."""
         return self._nloc_groups
 
     @property
-    def frame_nlocs(self) -> list[int]:
-        """Per-frame atom count."""
+    def frame_nlocs(self) -> np.ndarray:
+        """Per-frame atom count, indexed by dataset index."""
         return self._frame_nlocs
 
     @property
@@ -1553,13 +2666,24 @@ def nsystems(self) -> int:
         return self._nsystems
 
     @property
-    def frame_system_ids(self) -> list[int] | None:
-        """Per-frame system index, or None if not available."""
+    def frame_system_ids(self) -> np.ndarray | None:
+        """Per-frame system index, or None when the metadata omits it."""
         return self._frame_system_ids
 
     @property
-    def system_groups(self) -> dict[int, list[int]]:
-        """System index → list of frame indices."""
+    def system_groups(self) -> dict[int, np.ndarray]:
+        """System index → the ascending dataset indices of its frames.
+
+        Built on demand: the per-system index lists cost an array the size of
+        the dataset, which the training and statistics paths never need --
+        they consult :attr:`system_nframes` instead.
+        """
+        if self._system_groups is None:
+            self._system_groups = (
+                _group_positions_by_value(self._frame_system_ids)
+                if self._frame_system_ids is not None
+                else {0: np.arange(self.nframes, dtype=np.int64)}
+            )
         return self._system_groups
 
     @property
@@ -1568,21 +2692,97 @@ def system_nframes(self) -> list[int]:
         return self._system_nframes
 
 
-def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]:
+def _batch_find_flags(frames: list[dict[str, Any]]) -> dict[str, np.float32]:
+    """Reduce the per-frame ``find_*`` flags to the one scalar a batch carries.
+
+    A ``find_*`` flag is a single scalar for the whole batch, so a label
+    counts as available only where every frame supplies it. A frame that
+    lacks the label carries a default fill in its place, and reporting the
+    label as unavailable is what keeps that fill out of the loss.
+
+    Registered data requirements are default-filled on every frame by
+    :func:`decode_lmdb_frame`. Availability-sensitive labels are normally
+    partitioned before collation to preserve their usable frames; this
+    reduction remains the correctness boundary when the bounded probe misses
+    a rare frame. Default-backed and derived requirements may mix
+    deliberately, while mandatory requirements fail during frame resolution.
+
+    Parameters
+    ----------
+    frames : list[dict[str, Any]]
+        Per-frame dicts about to be stacked.
+
+    Returns
+    -------
+    dict[str, numpy.float32]
+        One scalar flag for every ``find_*`` key that any frame carries.
+    """
+    find_keys = sorted(
+        {key for frame in frames for key in frame if key.startswith("find_")}
+    )
+    return {
+        key: np.float32(
+            1.0
+            if all(key in frame and float(frame[key]) != 0.0 for frame in frames)
+            else 0.0
+        )
+        for key in find_keys
+    }
+
+
+def _pad_atom_axis(xp: Any, array: Any, length: int, fill: int, device: Any) -> Any:
+    """Widen one per-atom array's leading axis to ``length``."""
+    if array.shape[0] == length:
+        return array
+    tail = xp.full(
+        (length - array.shape[0], *array.shape[1:]),
+        fill,
+        dtype=array.dtype,
+        device=device,
+    )
+    return xp.concat([array, tail], axis=0)
+
+
+def collate_lmdb_frames(
+    frames: list[dict[str, Any]],
+    per_atom_keys: frozenset[str] = frozenset(),
+) -> dict[str, Any]:
     """Stack a list of per-frame dicts into a single batch dict.
 
     Backend-agnostic via ``array_api_compat``: works for numpy, torch, jax,
     etc. The array library is inferred from the first frame's ``coord``.
 
     Conventions match :func:`deepmd.dpmodel.utils.batch.normalize_batch`:
-    ``find_*`` flags remain scalar and must be constant within a batch;
-    ``fid`` is collected as a list; ``type`` is dropped (callers should
-    already use ``atype``); other arrays are stacked along axis 0. A ``sid``
-    placeholder is appended.
+    ``find_*`` flags remain scalar; ``fid`` is collected as a list; ``type``
+    is dropped (callers should already use ``atype``); other arrays are
+    stacked along axis 0. A ``sid`` placeholder is appended.
+
+    Frames need not agree on label availability. A label only some of them
+    carry is reported unavailable for the whole batch, as
+    :func:`_batch_find_flags` describes, and a field only some of them carry
+    is left out of the batch entirely because it cannot be stacked. Neither
+    can happen to a registered data requirement, which
+    :func:`decode_lmdb_frame` puts on every frame.
 
     The batch keeps the key order of its frames, which is the order
     :func:`decode_lmdb_batch` also produces, so a batch is the same mapping
     whichever of the two decode paths built it.
+
+    Parameters
+    ----------
+    frames : list[dict[str, Any]]
+        Per-frame dicts to stack.
+    per_atom_keys : frozenset[str], optional
+        Fields whose leading axis is the atom axis. When the frames differ in
+        atom count these are padded to the batch maximum, with ``atype``
+        filled by :data:`PHANTOM_ATOM_TYPE` and the rest zeroed. Leave empty
+        for uniform batches, where padding would be a no-op anyway. Resolve
+        the set with :func:`resolve_per_atom_keys`.
+
+    Returns
+    -------
+    dict[str, Any]
+        One collated batch.
     """
     import array_api_compat
 
@@ -1592,36 +2792,38 @@ def collate_lmdb_frames(frames: list[dict[str, Any]]) -> dict[str, Any]:
     xp = array_api_compat.array_namespace(frames[0]["coord"])
     dev = array_api_compat.device(frames[0]["coord"])
 
-    # Availability must agree across the batch before the flags can collapse
-    # to one scalar per key. Frames are checked ahead of collation so a mixed
-    # batch is reported rather than silently reduced to its first frame.
-    find_keys = sorted(
-        {key for frame in frames for key in frame if key.startswith("find_")}
-    )
-    for key in find_keys:
-        if any(key not in frame for frame in frames):
-            raise ValueError(
-                f"LMDB batch has inconsistent availability metadata for {key!r}"
-            )
-        values = [float(frame[key]) for frame in frames]
-        if any(value != values[0] for value in values[1:]):
-            raise ValueError(
-                f"LMDB batch mixes {key!r} values {values}; "
-                "SameNlocBatchSampler must group frames by label availability"
-            )
+    find_flags = _batch_find_flags(frames)
+
+    strides = per_atom_strides(frames[0], per_atom_keys) if per_atom_keys else {}
+    pad_nloc = max(frame["coord"].shape[0] for frame in frames) if strides else 0
 
     out: dict[str, Any] = {}
     for key in frames[0]:
         if key.startswith("find_"):
-            out[key] = frames[0][key]
+            out[key] = find_flags[key]
         elif key == "fid":
             out[key] = [f[key] for f in frames]
         elif key == "type":
             continue
         elif frames[0][key] is None:
             out[key] = None
+        elif any(key not in frame for frame in frames):
+            # A field only some frames carry cannot be stacked. Its ``find_``
+            # flag is false by the same token, so the batch stays coherent
+            # without it.
+            continue
+        elif key in strides:
+            length = pad_nloc * strides[key]
+            fill = _pad_fill_value(key)
+            out[key] = xp.stack(
+                [_pad_atom_axis(xp, f[key], length, fill, dev) for f in frames]
+            )
         else:
             out[key] = xp.stack([f[key] for f in frames])
+    # A flag a later frame raised that the first one never had still belongs
+    # to the batch, reporting the label as unavailable.
+    for key, flag in find_flags.items():
+        out.setdefault(key, flag)
     out["sid"] = xp.asarray([0], dtype=xp.int64, device=dev)
     return out
 
@@ -1664,6 +2866,12 @@ def compute_block_targets(
         stt, end, weight = part.split(":")
         blocks.append((int(stt), int(end), float(weight)))
 
+    # A bare ``prob_sys_size`` names no blocks: it asks for a probability
+    # proportional to system size, which is what sampling the merged frames
+    # uniformly already gives. There is nothing to reweight.
+    if not blocks:
+        return []
+
     # Drop blocks that retain zero frames (can happen when ``filter:N``
     # eliminates every system in a block). prob_sys_size_ext's per-block
     # ``nbatch_block / sum(nbatch_block)`` would otherwise propagate NaN
@@ -1676,8 +2884,9 @@ def compute_block_targets(
     ]
     if not nonempty:
         log.info(
-            "compute_block_targets: all blocks are empty in "
-            f"{auto_prob_style!r}; dataset has no retained frames."
+            "compute_block_targets: every block of "
+            f"{auto_prob_style!r} is empty; the dataset retains no frames in "
+            "any of them, so no reweighting is applied."
         )
         return []
     if len(nonempty) < len(blocks):
@@ -1757,7 +2966,7 @@ def _expand_indices_by_blocks(
     indices : list[int]
         Frame indices in the current nloc group.
     frame_system_ids : np.ndarray
-        Per-frame system id for the entire dataset (int64 array).
+        Per-frame system id for the entire dataset.
     block_targets : list[tuple[list[int], int]]
         Per-block (system_ids, total_target_frames).
     rng : np.random.Generator
@@ -1766,12 +2975,12 @@ def _expand_indices_by_blocks(
         Pre-computed total actual frame count per block (across all nloc
         groups).  When provided, avoids an O(N) scan of frame_system_ids.
     _sid_to_blk_arr : np.ndarray or None
-        Pre-computed system-id to block-index lookup array.  When provided,
+        Pre-computed lookup from :func:`system_block_lookup`. When provided,
         avoids rebuilding the mapping for each call.
     _group_block_targets : list[int] or None
         Exact target for each block in this group. Production samplers
-        allocate these targets globally across all ``(nloc, find-signature)``
-        groups so independent rounding cannot change a block's total size.
+        allocate these targets globally across all groups so independent
+        rounding cannot change a block's total size.
 
     Returns
     -------
@@ -1780,30 +2989,21 @@ def _expand_indices_by_blocks(
     """
     n_blocks = len(block_targets)
 
-    # Build sys_id -> block_idx lookup array
     if _sid_to_blk_arr is None:
-        sys_to_block: dict[int, int] = {}
-        for blk_idx, (sys_ids, _target) in enumerate(block_targets):
-            for sid in sys_ids:
-                sys_to_block[sid] = blk_idx
-        max_sid = max(sys_to_block.keys()) + 1 if sys_to_block else 0
-        _sid_to_blk_arr = np.full(max_sid, -1, dtype=np.int32)
-        for sid, blk in sys_to_block.items():
-            _sid_to_blk_arr[sid] = blk
-
-    # Partition indices by block using numpy for speed
+        _sid_to_blk_arr = system_block_lookup(block_targets)
+
+    sid_arr = np.asarray(frame_system_ids)
     idx_arr = np.asarray(indices, dtype=np.int64)
-    sid_arr = np.asarray(frame_system_ids, dtype=np.int64)
-    # Vectorized lookup: get block id for each index
-    idx_sids = sid_arr[idx_arr]
-    idx_blks = _sid_to_blk_arr[idx_sids]
+    idx_blks = resolve_frame_blocks(idx_arr, sid_arr, _sid_to_blk_arr)
 
     # Pre-compute block_total_actual if not provided
     if _block_total_actual is None and _group_block_targets is None:
-        _block_total_actual = []
-        for sys_ids, _ in block_targets:
-            total = sum(int(np.sum(sid_arr == sid)) for sid in sys_ids)
-            _block_total_actual.append(total)
+        _block_total_actual = count_group_blocks(
+            np.arange(sid_arr.size, dtype=np.int64),
+            sid_arr,
+            _sid_to_blk_arr,
+            n_blocks,
+        ).tolist()
 
     expanded_parts: list[np.ndarray] = []
 
@@ -1863,9 +3063,13 @@ def _expand_indices_by_blocks(
 
 def collect_lmdb_sampling_groups(
     reader: "LmdbDataReader",
-) -> list[tuple[int, list[int]]]:
+) -> list[tuple[int, np.ndarray]]:
     """Collect homogeneous LMDB groups shared by training and statistics.
 
+    Atom count comes from metadata and costs no read. Availability-sensitive
+    requirements subdivide a group only when the bounded probe detects mixed
+    source presence; see :meth:`LmdbDataReader.availability_groups`.
+
     Parameters
     ----------
     reader : LmdbDataReader
@@ -1873,22 +3077,132 @@ def collect_lmdb_sampling_groups(
 
     Returns
     -------
-    list[tuple[int, list[int]]]
+    list[tuple[int, numpy.ndarray]]
         Stable ``(nloc, frame indices)`` groups compatible with collation.
     """
-    groups: list[tuple[int, list[int]]] = []
-    for nloc in sorted(reader.nloc_groups):
-        signature_groups = reader.group_indices_by_find_signature(
-            list(reader.nloc_groups[nloc])
-        )
-        for signature in sorted(signature_groups):
-            groups.append((nloc, list(signature_groups[signature])))
-    return groups
+    return [
+        (nloc, group)
+        for nloc in sorted(reader.nloc_groups)
+        for group in reader.availability_groups(reader.nloc_groups[nloc])
+    ]
+
+
+def _collect_batch_groups(reader: "LmdbDataReader") -> list[np.ndarray]:
+    """Collect the groups a training batch may be drawn from.
+
+    A group is the largest set of frames one batch may span. Atom count
+    partitions the frames in every mode but ``mix:N``, whose decoded batch
+    accommodates unequal counts and so is bounded only by label availability.
+
+    Parameters
+    ----------
+    reader : LmdbDataReader
+        Reader providing the frame grouping and the batching rule.
+
+    Returns
+    -------
+    list[numpy.ndarray]
+        Frame indices per group, in the stable order shared by iteration and
+        length.
+    """
+    if reader.mixed_nloc:
+        return reader.availability_groups(np.arange(len(reader), dtype=np.int64))
+    return [indices for _nloc, indices in collect_lmdb_sampling_groups(reader)]
+
+
+def system_block_lookup(
+    block_targets: list[tuple[list[int], int]],
+) -> np.ndarray:
+    """Build the system-id to block-index lookup a whole group indexes at once.
+
+    The table carries one row beyond the highest system id named by a block,
+    holding the "no block" marker. Clamping a system id into the table then
+    maps everything a block does not name onto that row, which keeps the
+    lookup total without a per-frame membership test.
+
+    Parameters
+    ----------
+    block_targets : list[tuple[list[int], int]]
+        Per-block ``(system_ids, target_frame_count)``.
+
+    Returns
+    -------
+    numpy.ndarray
+        Block index of each system id, ``-1`` where no block names it, with
+        shape ``(max_system_id + 2,)``.
+    """
+    max_sid = max(
+        (sid for system_ids, _target in block_targets for sid in system_ids),
+        default=-1,
+    )
+    lookup = np.full(max_sid + 2, -1, dtype=np.int64)
+    for block_index, (system_ids, _target) in enumerate(block_targets):
+        if system_ids:
+            lookup[np.asarray(system_ids, dtype=np.int64)] = block_index
+    return lookup
+
+
+def resolve_frame_blocks(
+    indices: np.ndarray | list[int],
+    frame_system_ids: np.ndarray,
+    lookup: np.ndarray,
+) -> np.ndarray:
+    """Map each frame of one group to the block it belongs to.
+
+    Parameters
+    ----------
+    indices : numpy.ndarray or list[int]
+        Dataset indices of one group.
+    frame_system_ids : numpy.ndarray
+        Per-frame system id of the whole dataset. Gathered from, never
+        widened: it holds one entry per frame, so converting its dtype would
+        copy the entire dataset for the sake of a single group.
+    lookup : numpy.ndarray
+        System-id to block-index table from :func:`system_block_lookup`.
+
+    Returns
+    -------
+    numpy.ndarray
+        Block index of each frame, ``-1`` where no block claims it.
+    """
+    sids = np.asarray(frame_system_ids)[np.asarray(indices, dtype=np.int64)]
+    return lookup[np.minimum(sids, lookup.size - 1)]
+
+
+def count_group_blocks(
+    indices: np.ndarray | list[int],
+    frame_system_ids: np.ndarray,
+    lookup: np.ndarray,
+    n_blocks: int,
+) -> np.ndarray:
+    """Count how many of one group's frames fall in each block.
+
+    Parameters
+    ----------
+    indices : numpy.ndarray or list[int]
+        Dataset indices of one group.
+    frame_system_ids : numpy.ndarray
+        Per-frame system id of the whole dataset.
+    lookup : numpy.ndarray
+        System-id to block-index table from :func:`system_block_lookup`.
+    n_blocks : int
+        Number of blocks.
+
+    Returns
+    -------
+    numpy.ndarray
+        Frame count per block, with shape ``(n_blocks,)``. Frames belonging
+        to no block are excluded, so the counts need not sum to the group
+        size.
+    """
+    blocks = resolve_frame_blocks(indices, frame_system_ids, lookup)
+    # Shift by one so that the "no block" marker lands in bin zero.
+    return np.bincount(blocks + 1, minlength=n_blocks + 1)[1:]
 
 
 def _allocate_group_block_targets(
-    groups: list[tuple[int, list[int]]],
-    frame_system_ids: list[int] | np.ndarray,
+    groups: list[np.ndarray],
+    frame_system_ids: np.ndarray,
     block_targets: list[tuple[list[int], int]],
 ) -> list[list[int]]:
     """Allocate every block target exactly across homogeneous groups.
@@ -1898,17 +3212,13 @@ def _allocate_group_block_targets(
     remainder method. Stable group order breaks equal-remainder ties, which
     keeps distributed ranks deterministic without floating-point rounding.
     """
-    group_actual = [[0] * len(block_targets) for _ in groups]
-    system_to_block = {
-        system_id: block_index
-        for block_index, (system_ids, _target) in enumerate(block_targets)
-        for system_id in system_ids
-    }
-    for group_index, (_nloc, indices) in enumerate(groups):
-        for index in indices:
-            block_index = system_to_block.get(int(frame_system_ids[index]))
-            if block_index is not None:
-                group_actual[group_index][block_index] += 1
+    lookup = system_block_lookup(block_targets)
+    group_actual = [
+        count_group_blocks(
+            indices, frame_system_ids, lookup, len(block_targets)
+        ).tolist()
+        for indices in groups
+    ]
 
     group_targets = [counts.copy() for counts in group_actual]
     for block_index, (_system_ids, block_target) in enumerate(block_targets):
@@ -1947,24 +3257,125 @@ def _allocate_group_block_targets(
     return group_targets
 
 
+def _chop_same_nloc(
+    reader: "LmdbDataReader", indices: Sequence[int]
+) -> list[list[int]]:
+    """Split one homogeneous group into fixed-size batches."""
+    batch_size = reader.get_batch_size_for_nloc(int(reader.frame_nlocs[indices[0]]))
+    index_list = np.asarray(indices).tolist()
+    return [
+        index_list[start : start + batch_size]
+        for start in range(0, len(index_list), batch_size)
+    ]
+
+
+def _chop_mixed_nloc(reader: "LmdbDataReader", indices: list[int]) -> list[list[int]]:
+    """Split one group into batches under an atom-axis budget.
+
+    ``mix:N`` budgets the length of a batch's atom axis, and the layout the
+    batch will be decoded in decides both how that length is measured and the
+    order the frames are visited in.
+
+    **Ragged.** The frames are concatenated, so the axis is simply their total
+    atom count and nothing is padded. No frame's cost depends on its
+    neighbours, so the group is taken in the caller's order, which training has
+    already shuffled. Sorting would only make each batch homogeneous in system
+    size, correlating the frames of an optimizer step for no gain.
+
+    **Rectangular.** Every frame is padded to the widest one, so the axis is
+    ``nframes * max_nloc`` and a batch pays for atoms it does not hold. Sorting
+    by atom count is what keeps that overhead small: a batch is then a run of
+    frames adjacent in the sorted order, so its padding is bounded by the
+    atom-count spread across that run alone. Ties fall back to the caller's
+    order, so frames of equal atom count still mix freely across epochs.
+
+    Either way the group is cut only where the next frame would push the axis
+    past the budget, which is the fewest batches obtainable **without
+    reordering**: a batch's cost does not fall when frames are dropped from its
+    front, so a batch that starts later can always extend at least as far as
+    one that starts earlier, and by induction on the batch count, cutting as
+    late as possible covers the longest prefix for every count.
+
+    Under the rectangular layout the sort is part of the algorithm and the
+    result is optimal outright, since an exchange argument turns any packing
+    into contiguous runs of the sorted order. Under the ragged layout the order
+    is the caller's, so the count is optimal only for that order; reordering
+    could pack tighter -- the general problem is bin packing -- and is declined
+    to keep the frames of an optimizer step decorrelated in system size.
+    Padding, where it exists, is not separately minimized either.
+
+    That the batches are full is what keeps the gradient weighting faithful. A
+    batch is one optimizer step; its per-atom loss terms pool over the real
+    labels, so a frame's weight there follows its atom count whatever the
+    packing, while its frame-level terms (energy, virial) weigh frames equally,
+    giving a frame the weight ``1 / k_b``. An atom budget makes ``k_b`` follow
+    the atom count, exactly so under the ragged layout and up to the padding
+    under the rectangular one, and an under-filled batch raises the weight of
+    every frame it holds.
+
+    A frame larger than the budget forms a batch of its own, matching how
+    ``max:N`` treats an oversized nloc group.
+
+    Parameters
+    ----------
+    reader : LmdbDataReader
+        Provides the per-frame atom counts, the atom budget and the layout.
+    indices : list[int]
+        Dataset indices of one group.
+
+    Returns
+    -------
+    list[list[int]]
+        Batches whose union is ``indices``.
+    """
+    budget = reader.atom_budget
+    if budget is None:
+        raise ValueError("mixed-nloc batching requires a batch_size of 'mix:N'")
+
+    index_array = np.asarray(indices, dtype=np.int64)
+    # Gathered, not widened: the atom counts are one entry per frame, so a
+    # dtype conversion here would copy the whole dataset for one group.
+    nloc_array = reader.frame_nlocs[index_array]
+    if not reader.ragged_batches:
+        order = np.argsort(nloc_array, kind="stable")
+        index_array, nloc_array = index_array[order], nloc_array[order]
+
+    batches: list[list[int]] = []
+    batch_start = 0
+    real_atoms = 0
+    for position, nloc in enumerate(nloc_array.tolist()):
+        count = position - batch_start + 1
+        # Length of the atom axis this run would occupy once decoded. Under the
+        # rectangular layout the ascending sort makes the current frame the
+        # widest, so it alone sets the padded width.
+        axis = real_atoms + nloc if reader.ragged_batches else count * nloc
+        if count > 1 and axis > budget:
+            batches.append(index_array[batch_start:position].tolist())
+            batch_start, real_atoms = position, 0
+        real_atoms += nloc
+    batches.append(index_array[batch_start:].tolist())
+    return batches
+
+
 def _build_all_batches(
     reader: "LmdbDataReader",
     shuffle: bool,
     rng: np.random.Generator,
     block_targets: list[tuple[list[int], int]] | None = None,
 ) -> list[list[int]]:
-    """Build batches homogeneous in atom count and label availability.
+    """Build the batches of one pass over the dataset.
 
-    This is the shared batch-construction logic used by both
-    SameNlocBatchSampler (single-GPU) and DistributedSameNlocBatchSampler.
+    Groups are chopped into batches, then interleaved round-robin so that
+    consecutive batches come from different groups. Under ``mix:N`` a batch
+    also spans several atom counts; every other mode keeps it uniform.
 
     Parameters
     ----------
     reader : LmdbDataReader
-        Provides nloc_groups and get_batch_size_for_nloc.
+        Provides the frame grouping and the batching rule.
     shuffle : bool
-        Whether to shuffle indices within each nloc group and
-        shuffle the final batch order.
+        Whether to shuffle indices within each group and shuffle the final
+        batch order.
     rng : np.random.Generator
         Random number generator (deterministic for reproducibility).
     block_targets : list[tuple[list[int], int]] or None
@@ -1974,9 +3385,10 @@ def _build_all_batches(
     Returns
     -------
     list[list[int]]
-        Each inner list has one nloc and one scalar ``find_*`` signature.
+        Dataset indices grouped into batches.
     """
-    groups = collect_lmdb_sampling_groups(reader)
+    groups = _collect_batch_groups(reader)
+    chop = _chop_mixed_nloc if reader.mixed_nloc else _chop_same_nloc
 
     # Build per-group batches
     group_batches: list[list[list[int]]] = []
@@ -1986,25 +3398,16 @@ def _build_all_batches(
     sid_to_blk_arr: np.ndarray | None = None
     group_block_targets: list[list[int]] | None = None
     if block_targets and reader.frame_system_ids is not None:
-        # Convert frame_system_ids to numpy once
-        sid_arr = np.array(reader.frame_system_ids, dtype=np.int64)
+        sid_arr = reader.frame_system_ids
         group_block_targets = _allocate_group_block_targets(
             groups, sid_arr, block_targets
         )
-        # Build sys_id -> block_idx lookup array once
-        sys_to_block: dict[int, int] = {}
-        for blk_idx, (sys_ids, _target) in enumerate(block_targets):
-            for sid in sys_ids:
-                sys_to_block[sid] = blk_idx
-        max_sid = max(sys_to_block.keys()) + 1 if sys_to_block else 0
-        sid_to_blk_arr = np.full(max_sid, -1, dtype=np.int32)
-        for sid, blk in sys_to_block.items():
-            sid_to_blk_arr[sid] = blk
-
-    for group_index, (nloc, original_indices) in enumerate(groups):
+        sid_to_blk_arr = system_block_lookup(block_targets)
+
+    for group_index, original_indices in enumerate(groups):
         indices = original_indices
-        # Expand each availability group independently using targets that
-        # were allocated globally, preserving both scalar flags and totals.
+        # Expand each group independently using targets that were allocated
+        # globally, so that a block's total size is preserved exactly.
         if block_targets and sid_arr is not None and group_block_targets is not None:
             indices = _expand_indices_by_blocks(
                 indices,
@@ -2015,12 +3418,10 @@ def _build_all_batches(
                 _group_block_targets=group_block_targets[group_index],
             )
         if shuffle:
-            rng.shuffle(indices)
-        bs = reader.get_batch_size_for_nloc(nloc)
-        batches = []
-        for start in range(0, len(indices), bs):
-            batches.append(indices[start : start + bs])
-        group_batches.append(batches)
+            # ``permutation`` returns a copy, which matters because a group
+            # may alias one of the reader's own index tables.
+            indices = rng.permutation(indices)
+        group_batches.append(chop(reader, indices) if len(indices) else [])
 
     # Interleave groups round-robin
     all_batches: list[list[int]] = []
@@ -2037,26 +3438,34 @@ def _build_all_batches(
     return all_batches
 
 
-class SameNlocBatchSampler:
-    """Batch sampler that groups frames by nloc and ``find_*`` signature.
+class LmdbBatchSampler:
+    """Batch sampler over an LMDB, grouped by atom count.
 
-    For mixed-nloc datasets with mixed_batch=False: each batch contains only
-    frames with the same nloc and label availability. Within each group,
-    frames are shuffled. Groups are interleaved round-robin so training sees
-    diverse nloc and label combinations.
+    Atom count is handled by the reader's batching rule: all rules but
+    ``mix:N`` keep a batch uniform in atom count, while ``mix:N`` fills
+    batches to the atom budget its decoded layout is measured against.
+    Groups are interleaved round-robin and the batch order is then shuffled,
+    so training sees a varied mix.
 
-    When auto batch_size is used, batch_size is computed per-nloc-group.
+    Optional tracked labels subdivide these groups when a bounded probe
+    detects mixed presence. Default-backed and derived fields may mix freely;
+    mandatory fields fail when an unavailable frame is decoded.
 
-    The sampler is deterministic for a fixed seed and epoch. Use
-    :meth:`set_epoch` to select a different reproducible sequence for each
-    training pass.
+    The sampler serves one pass at a time, drawn from ``seed + epoch``. The
+    pending pass is materialized before it is served, which is what lets
+    ``__len__`` report exactly what ``__iter__`` will yield: under ``mix:N``
+    the batch count follows the shuffle, because batches are filled to an atom
+    budget rather than to a fixed frame count. Serving a pass advances the
+    epoch, so a caller that just re-iterates sees a different shuffle every
+    time, and :meth:`set_epoch` repositions that progression for a caller --
+    a distributed run, a resumed one -- that needs to name the pass instead.
 
     Parameters
     ----------
     reader : LmdbDataReader
-        The dataset reader (provides nloc_groups, get_batch_size_for_nloc).
+        The dataset reader providing the frame grouping and batching rule.
     shuffle : bool
-        Whether to shuffle within each nloc group each epoch.
+        Whether to shuffle within each group and shuffle the batch order.
     seed : int or None
         Random seed for reproducibility.
     block_targets : list[tuple[list[int], int]] or None
@@ -2075,61 +3484,64 @@ def __init__(
         self._seed = seed
         self._epoch = 0
         self._block_targets = block_targets
+        self._batches: list[list[int]] | None = None
+        self._data_requirements_revision = -1
+
+    def batches(self) -> list[list[int]]:
+        """Return the batch list of the pending pass, building it if needed.
+
+        Returns
+        -------
+        list[list[int]]
+            Dataset indices grouped into batches.
+        """
+        revision = self._reader.data_requirements_revision
+        if self._batches is None or self._data_requirements_revision != revision:
+            seed = None if self._seed is None else self._seed + self._epoch
+            self._batches = _build_all_batches(
+                self._reader,
+                self._shuffle,
+                np.random.default_rng(seed),
+                self._block_targets,
+            )
+            self._data_requirements_revision = revision
+        return self._batches
 
     def set_epoch(self, epoch: int) -> None:
-        """Set the epoch used to derive the deterministic shuffle state.
+        """Select the pass to serve, discarding any pass still pending.
 
         Parameters
         ----------
         epoch : int
             Zero-based training epoch.
         """
-        self._epoch = epoch
+        if epoch != self._epoch:
+            self._epoch = epoch
+            self._batches = None
 
     def __iter__(self) -> Iterator[list[int]]:
-        """Yield batches of frame indices, all with the same nloc."""
-        seed = None if self._seed is None else self._seed + self._epoch
-        rng = np.random.default_rng(seed)
-        yield from _build_all_batches(
-            self._reader, self._shuffle, rng, self._block_targets
-        )
+        """Yield the pending pass, and move the epoch on to its successor."""
+        batches = self.batches()
+        self.set_epoch(self._epoch + 1)
+        yield from batches
 
     def __len__(self) -> int:
-        """Total batches across nloc and label-availability groups."""
-        groups = collect_lmdb_sampling_groups(self._reader)
-        group_block_targets = None
-        assigned_system_ids: set[int] = set()
-        if self._block_targets and self._reader.frame_system_ids is not None:
-            group_block_targets = _allocate_group_block_targets(
-                groups,
-                self._reader.frame_system_ids,
-                self._block_targets,
-            )
-            assigned_system_ids = {
-                system_id
-                for system_ids, _target in self._block_targets
-                for system_id in system_ids
-            }
+        """Number of batches the pending pass holds."""
+        return len(self.batches())
 
-        total = 0
-        for group_index, (nloc, indices) in enumerate(groups):
-            bs = self._reader.get_batch_size_for_nloc(nloc)
-            n = len(indices)
-            if (
-                group_block_targets is not None
-                and self._reader.frame_system_ids is not None
-            ):
-                unassigned = sum(
-                    int(self._reader.frame_system_ids[index]) not in assigned_system_ids
-                    for index in indices
-                )
-                n = unassigned + sum(group_block_targets[group_index])
-            total += (n + bs - 1) // bs
-        return total
+    @property
+    def total_batches(self) -> int:
+        """Number of batches the pending pass holds over the whole dataset.
+
+        The same count as ``len(self)`` here, and the two part company only in
+        the distributed sampler, so a caller after a global figure need not
+        know which of the two it holds.
+        """
+        return len(self)
 
 
-class DistributedSameNlocBatchSampler:
-    """Distributed wrapper for same-nloc batch sampling.
+class DistributedLmdbBatchSampler:
+    """Distributed wrapper for LMDB batch sampling.
 
     All ranks build the same deterministic global batch list (using
     ``seed + epoch``). The list is padded deterministically when its length is
@@ -2144,8 +3556,7 @@ class DistributedSameNlocBatchSampler:
     Parameters
     ----------
     reader : LmdbDataReader
-        The dataset reader (provides nloc_groups, get_batch_size_for_nloc,
-        frame_nlocs).
+        The dataset reader providing the frame grouping and batching rule.
     rank : int
         Rank of the current process.
     world_size : int
@@ -2174,17 +3585,7 @@ def __init__(
         self._seed = seed if seed is not None else 0
         self._epoch = 0
         self._block_targets = block_targets
-        self.refresh_batch_count()
-
-    def refresh_batch_count(self) -> None:
-        """Refresh the cached global count after sampling groups change."""
-        self._total_batches = len(
-            SameNlocBatchSampler(
-                self._reader,
-                shuffle=False,
-                block_targets=self._block_targets,
-            )
-        )
+        self._global: LmdbBatchSampler | None = None
 
     def set_epoch(self, epoch: int) -> None:
         """Set epoch for deterministic cross-rank shuffling.
@@ -2192,31 +3593,38 @@ def set_epoch(self, epoch: int) -> None:
         Call this before each training epoch/cycle to get different but
         reproducible batch orderings across epochs.
         """
-        self._epoch = epoch
+        if epoch != self._epoch:
+            self._epoch = epoch
+            self._global = None
+
+    def _global_batches(self) -> list[list[int]]:
+        """Return the batch list every rank builds identically."""
+        if self._global is None:
+            self._global = LmdbBatchSampler(
+                self._reader,
+                shuffle=self._shuffle,
+                seed=self._seed + self._epoch,
+                block_targets=self._block_targets,
+            )
+        return self._global.batches()
 
     def __iter__(self) -> Iterator[list[int]]:
         """Yield this rank's partition of the global batch list."""
-        # All ranks build the same global batch list deterministically
-        rng = np.random.default_rng(self._seed + self._epoch)
-        all_batches = _build_all_batches(
-            self._reader, self._shuffle, rng, self._block_targets
-        )
-        # Partition to this rank
-        yield from self._partition_batches(all_batches)
+        yield from self._partition_batches(self._global_batches())
 
     def _partition_batches(self, all_batches: list[list[int]]) -> list[list[int]]:
         """Partition global batches to this rank.
 
         The default pads the global list to a multiple of ``world_size`` and
         then takes ``all_batches[rank::world_size]``. This gives good nloc
-        diversity per rank since batches are interleaved across nloc groups
-        before shuffling, while ensuring that every rank yields the same
-        number of batches.
+        diversity per rank since batches are interleaved across groups before
+        shuffling, while ensuring that every rank yields the same number of
+        batches.
 
         Override this method for custom load-balancing. For example, a
         greedy algorithm could assign batches to ranks based on estimated
-        compute cost (``reader.frame_nlocs[batch[0]]`` gives the nloc of
-        each batch).
+        compute cost (``reader.batch_pad_nloc(batch) * len(batch)`` gives the
+        padded cost of each batch).
         """
         if not all_batches:
             return []
@@ -2233,12 +3641,12 @@ def _partition_batches(self, all_batches: list[list[int]]) -> list[list[int]]:
 
     def __len__(self) -> int:
         """Number of batches for this rank."""
-        return (self._total_batches + self._world_size - 1) // self._world_size
+        return len(self._partition_batches(self._global_batches()))
 
     @property
     def total_batches(self) -> int:
-        """Return the global batch count before distributed padding."""
-        return self._total_batches
+        """Number of batches one full pass holds, before the per-rank split."""
+        return len(self._global_batches())
 
     @property
     def rank(self) -> int:
@@ -2332,13 +3740,16 @@ def __init__(
     ) -> None:
         self.lmdb_path = str(lmdb_path)
         self._type_map = type_map or []
-        self._env = _open_lmdb(self.lmdb_path)
-
-        with self._env.begin() as txn:
-            meta = _read_metadata(txn)
+        meta = _read_metadata_of(self.lmdb_path)
+        # Every read this class serves walks a group in ascending key order,
+        # which is the pattern kernel readahead exists for.
+        self._env = _open_lmdb(self.lmdb_path, sequential=True)
 
         self.nframes, self._frame_fmt, self._natoms_per_type = _parse_metadata(meta)
         self._natoms = sum(self._natoms_per_type)
+        self._ntypes = (
+            len(self._type_map) if self._type_map else len(self._natoms_per_type)
+        )
 
         # Build type remapping if LMDB's type_map differs from model's type_map
         lmdb_type_map = meta.get("type_map")
@@ -2370,6 +3781,18 @@ def __init__(
 
         # Data requirements
         self._requirements: dict[str, dict[str, Any]] = {}
+        self._uniform_availability: bool | None = None
+        self._availability_indices: dict[
+            int, tuple[np.ndarray, _AvailabilityIndex]
+        ] = {}
+        self._full_availability_index: _AvailabilityIndex | None = None
+        self._decode_config = LmdbDecodeConfig(
+            ntypes=self._ntypes,
+            natoms=self._natoms,
+            type_remap=self._type_remap,
+            data_requirements=self._requirements,
+            dataset=self.lmdb_path,
+        )
 
         # Detect PBC from the first retained frame.
         self.pbc = True
@@ -2389,7 +3812,7 @@ def _select_frames(
         meta: dict,
         shuffle_test: bool,
         max_frames: float | None,
-    ) -> dict[int, list[int]]:
+    ) -> dict[int, np.ndarray]:
         """Group the frame indices by atom count, then sample each group.
 
         Parameters
@@ -2405,50 +3828,43 @@ def _select_frames(
 
         Returns
         -------
-        dict[int, list[int]]
+        dict[int, numpy.ndarray]
             The retained LMDB frame indices of each atom count.
         """
         raw_nlocs = meta.get("frame_nlocs")
         if _is_encoded_array(raw_nlocs):
-            nlocs = _decode_array(raw_nlocs).reshape(-1).astype(np.int64)
+            nlocs = _decode_array(raw_nlocs).reshape(-1).astype(np.int32)
         elif raw_nlocs is not None:
-            nlocs = np.asarray(raw_nlocs, dtype=np.int64)
+            nlocs = np.fromiter(raw_nlocs, dtype=np.int32, count=len(raw_nlocs))
         else:
             nlocs = np.asarray(
                 _scan_frame_nlocs(
                     self._env, self.nframes, self._frame_fmt, self._natoms
                 ),
-                dtype=np.int64,
+                dtype=np.int32,
             )
 
-        # Sorting once groups every atom count in a single pass, which matters
-        # for datasets whose frame count reaches into the millions.
-        order = np.argsort(nlocs, kind="stable")
-        starts = np.concatenate(
-            ([0], np.flatnonzero(np.diff(nlocs[order])) + 1, [nlocs.size])
-        )
-
+        groups = _group_positions_by_value(nlocs)
         keep = (
             None
             if max_frames is None or not np.isfinite(max_frames)
             else int(max_frames)
         )
-        groups: dict[int, list[int]] = {}
-        for begin, end in pairwise(starts):
-            indices = order[begin:end].copy()
+        if not shuffle_test and keep is None:
+            return groups
+
+        for nloc, indices in groups.items():
             if shuffle_test:
                 dp_random.shuffle(indices)
-            if keep is not None:
-                indices = indices[:keep]
-            groups[int(nlocs[order[begin]])] = indices.tolist()
+            groups[nloc] = indices if keep is None else indices[:keep]
         return groups
 
-    def _read_frames(self, frame_indices: list[int]) -> list[dict[str, Any]]:
+    def _read_frames(self, frame_indices: Sequence[int]) -> list[dict[str, Any]]:
         """Decode the given LMDB frames, applying the type remapping.
 
         Parameters
         ----------
-        frame_indices : list[int]
+        frame_indices : Sequence[int]
             Indices of the frames to read, as keyed in the LMDB.
 
         Returns
@@ -2457,77 +3873,121 @@ def _read_frames(self, frame_indices: list[int]) -> list[dict[str, Any]]:
             One decoded frame per index that the LMDB holds.
         """
         frames: list[dict[str, Any]] = []
-        with self._env.begin() as txn:
-            for index in frame_indices:
-                raw = txn.get(format(index, self._frame_fmt).encode())
+        with self._env.begin() as transaction:
+            for index in np.asarray(frame_indices).tolist():
+                raw = transaction.get(format(index, self._frame_fmt).encode())
                 if raw is None:
                     continue
-                frame = _remap_keys(_decode_frame(raw))
-                atype = frame.get("atype")
-                if self._type_remap is not None and isinstance(atype, np.ndarray):
-                    frame["atype"] = _remap_atom_types(
-                        atype.reshape(-1), self._type_remap
+                frames.append(
+                    decode_lmdb_frame(
+                        raw,
+                        int(index),
+                        self._decode_config,
+                        copy_arrays=True,
                     )
-                frames.append(frame)
+                )
         return frames
 
     def __del__(self) -> None:
-        """Release the LMDB environment ref-count on garbage collection."""
-        path = getattr(self, "lmdb_path", None)
-        if path is not None:
-            _close_lmdb(path)
+        """Release the LMDB environment ref-count on garbage collection.
+
+        The count is released only once, and only if construction got as far
+        as taking it: an instance that failed earlier holds no reference, and
+        releasing one it never took would close the environment underneath
+        whichever reader does hold it.
+        """
+        if getattr(self, "_env", None) is None:
+            return
+        self._env = None
+        _close_lmdb(self.lmdb_path)
 
     @property
-    def nloc_groups(self) -> dict[int, list[int]]:
+    def nloc_groups(self) -> dict[int, np.ndarray]:
         """Nloc → the LMDB frame indices retained for that atom count."""
         return self._nloc_groups
 
-    @staticmethod
-    def _frame_has_data(frame: dict[str, Any], key: str) -> bool:
-        """Resolve one frame's explicit or inferred ``find_*`` value.
+    def availability_groups(self, frame_indices: np.ndarray) -> list[np.ndarray]:
+        """Partition LMDB frame indices into label-compatible groups.
 
-        The frame may still carry its msgpack payload, so that availability
-        can be settled without decoding the arrays it describes.
-        """
-        find_key = f"find_{key}"
-        if find_key in frame:
-            return bool(float(np.asarray(_decode_value(frame[find_key])).item()))
-        value = frame.get(key)
-        if _is_encoded_array(value):
-            return True
-        return isinstance(value, (np.ndarray, np.generic, int, float, bool))
+        The validation counterpart of
+        :meth:`LmdbDataReader.availability_groups`. Only optional tracked
+        labels participate. Exact signature IDs are cached per retained nloc
+        group and reused by every full-validation pass.
 
-    @property
-    def find_signature_groups(
-        self,
-    ) -> dict[tuple[int, tuple[tuple[str, bool], ...]], list[int]]:
-        """Group the retained frame indices by atom count and label availability.
+        Parameters
+        ----------
+        frame_indices : numpy.ndarray
+            LMDB frame indices to partition.
 
-        Frames that :meth:`_stack_frames` would refuse to stack together land
-        in different groups, because both settle availability with
-        :meth:`_frame_has_data`. Only the msgpack payload of each frame is
-        read, so the grouping does not decode the arrays it separates.
+        Returns
+        -------
+        list[numpy.ndarray]
+            Non-empty index groups in a stable order, together covering the
+            input. An empty group would name no frames to stack, so an empty
+            input yields no group at all.
         """
-        groups: dict[tuple[int, tuple[tuple[str, bool], ...]], list[int]] = {}
-        with self._env.begin() as transaction:
-            for nloc, frame_indices in self._nloc_groups.items():
-                for index in frame_indices:
-                    raw = transaction.get(format(index, self._frame_fmt).encode())
-                    if raw is None:
-                        continue
-                    frame = _remap_keys(msgpack.unpackb(raw, raw=False))
-                    signature = tuple(
-                        (f"find_{key}", self._frame_has_data(frame, key))
-                        for key in _availability_signature_keys(
-                            frame, iter(self._requirements)
-                        )
+        if not len(frame_indices):
+            return []
+        availability_keys = _availability_requirement_keys(self._requirements)
+        if not availability_keys:
+            self._uniform_availability = True
+            return [frame_indices]
+        if self._uniform_availability is None:
+            groups = list(self._nloc_groups.values())
+            per_group = max(1, _AVAILABILITY_PROBE_FRAMES // max(len(groups), 1))
+            with self._env.begin() as transaction:
+                self._uniform_availability = _probe_uniform_availability(
+                    transaction,
+                    (
+                        key
+                        for group in groups
+                        for key in _evenly_spaced(group, per_group)
+                    ),
+                    self._frame_fmt,
+                    availability_keys,
+                )
+        if self._uniform_availability:
+            return [frame_indices]
+        index_array = np.asarray(frame_indices)
+        if not bool(self._env.flags().get("readahead", False)):
+            if self._full_availability_index is None:
+                self._full_availability_index = _scan_lmdb_path_in_worker(
+                    self.lmdb_path,
+                    availability_keys,
+                )
+            return self._full_availability_index.groups(
+                index_array,
+                positions=index_array,
+            )
+
+        cache_key = id(frame_indices)
+        cached = self._availability_indices.get(cache_key)
+        if cached is None or cached[0] is not frame_indices:
+            with self._env.begin() as transaction:
+
+                def read_raw(position: int) -> bytes | None:
+                    frame_index = int(index_array[position])
+                    return transaction.get(
+                        format(frame_index, self._frame_fmt).encode()
                     )
-                    groups.setdefault((nloc, signature), []).append(index)
-        return groups
 
-    def get_test_by_indices(self, frame_indices: list[int]) -> dict[str, Any]:
+                availability_index = _scan_availability_index(
+                    len(index_array),
+                    read_raw,
+                    availability_keys,
+                    self.lmdb_path,
+                )
+            self._availability_indices[cache_key] = (
+                frame_indices,
+                availability_index,
+            )
+        else:
+            availability_index = cached[1]
+        return availability_index.groups(index_array)
+
+    def get_test_by_indices(self, frame_indices: Sequence[int]) -> dict[str, Any]:
         """Stack one homogeneous validation group selected by frame index."""
-        if not frame_indices:
+        if not len(frame_indices):
             raise ValueError("frame_indices must contain at least one frame")
         frames = self._read_frames(frame_indices)
         nlocs = {
@@ -2552,18 +4012,25 @@ def add(
         repeat: int = 1,
         default: float = 0.0,
         dtype: np.dtype | None = None,
+        source_policy: DataRequirementSourcePolicy = "tracked",
         **kwargs: Any,
     ) -> None:
         """Register a data requirement (mirrors DeepmdData.add)."""
-        self._requirements[key] = {
-            "ndof": ndof,
-            "atomic": atomic,
-            "must": must,
-            "high_prec": high_prec,
-            "repeat": repeat,
-            "default": default,
-            "dtype": dtype,
-        }
+        requirement = DataRequirementItem(
+            key,
+            ndof,
+            atomic=atomic,
+            must=must,
+            high_prec=high_prec,
+            repeat=repeat,
+            default=default,
+            dtype=dtype,
+            source_policy=source_policy,
+        )
+        self._requirements[key] = requirement.dict
+        self._uniform_availability = None
+        self._availability_indices.clear()
+        self._full_availability_index = None
 
     def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None:
         """Register expected keys from ``DataRequirementItem`` objects.
@@ -2582,6 +4049,7 @@ def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> N
                 repeat=item["repeat"],
                 default=item["default"],
                 dtype=item["dtype"],
+                source_policy=item["source_policy"],
             )
 
     def _resolve_dtype(self, key: str) -> np.dtype:
@@ -2618,6 +4086,7 @@ def iter_test(
         chunk_atoms: int,
         numb_test: float = float("inf"),
         nloc: int | None = None,
+        frame_indices: Sequence[int] | None = None,
     ) -> Iterator[dict[str, Any]]:
         """Yield the test frames in chunks, reading each chunk on demand.
 
@@ -2634,13 +4103,18 @@ def iter_test(
             serves every frame of the group.
         nloc : int or None, optional
             Atom count selecting the group, resolved as in :meth:`get_test`.
+        frame_indices : Sequence[int] or None, optional
+            Frames to serve in place of the whole group, which is how a
+            label-compatible subgroup names itself. They must all have the
+            atom count ``nloc`` selects.
 
         Yields
         ------
         dict[str, Any]
             One chunk of frames, stacked as :meth:`get_test` stacks them.
         """
-        frame_indices, natoms = self._resolve_group(nloc)
+        group_indices, natoms = self._resolve_group(nloc)
+        frame_indices = group_indices if frame_indices is None else frame_indices
         if np.isfinite(numb_test):
             frame_indices = frame_indices[: int(numb_test)]
         step = max(1, int(chunk_atoms) // max(1, natoms))
@@ -2648,7 +4122,7 @@ def iter_test(
             chunk = frame_indices[begin : begin + step]
             yield self._stack_frames(self._read_frames(chunk), natoms)
 
-    def _resolve_group(self, nloc: int | None) -> tuple[list[int], int]:
+    def _resolve_group(self, nloc: int | None) -> tuple[np.ndarray, int]:
         """Return the retained frame indices and atom count of one group.
 
         Parameters
@@ -2659,7 +4133,7 @@ def _resolve_group(self, nloc: int | None) -> tuple[list[int], int]:
 
         Returns
         -------
-        tuple[list[int], int]
+        tuple[numpy.ndarray, int]
             The LMDB frame indices of the group and its atom count.
 
         Raises
@@ -2688,7 +4162,12 @@ def _resolve_group(self, nloc: int | None) -> tuple[list[int], int]:
     def _stack_frames(
         self, frames: list[dict[str, Any]], natoms: int
     ) -> dict[str, Any]:
-        """Stack a list of same-nloc frames into numpy arrays."""
+        """Stack a list of same-nloc frames into numpy arrays.
+
+        Frames need not agree on label availability; a label only some of
+        them carry is reported unavailable for the whole group, mirroring
+        :func:`_batch_find_flags` on the training path.
+        """
         nframes = len(frames)
         result: dict[str, Any] = {}
 
@@ -2728,7 +4207,9 @@ def _stack_frames(
         # Dynamically discover all data keys from the first frame, plus
         # any registered requirements.  Structural keys (coord, box, type)
         # are excluded — they are already handled above.
-        _structural_keys = frozenset({"coord", "box", "atype"})
+        _structural_keys = frozenset(
+            {"coord", "box", "atype", "natoms", "real_natoms_vec", "fid"}
+        )
         all_keys: dict[str, dict[str, Any]] = {}
         if frames:
             for fk in frames[0]:
@@ -2739,61 +4220,29 @@ def _stack_frames(
         for key, req in self._requirements.items():
             all_keys[key] = req
 
-        for key, req_info in all_keys.items():
-            availability = [self._frame_has_data(frame, key) for frame in frames]
-            if any(flag != availability[0] for flag in availability[1:]):
-                raise ValueError(
-                    f"LMDB validation group mixes find_{key} values {availability}"
-                )
-            has_key = availability[0]
-            if not has_key and req_info.get("must", False):
-                raise RuntimeError(f"Required LMDB test-data field {key!r} is missing.")
-            result[f"find_{key}"] = 1.0 if has_key else 0.0
-
-            # Get repeat factor from registered requirements
-            repeat = 1
-            if key in self._requirements:
-                repeat = self._requirements[key].get("repeat", 1)
-
-            if has_key:
-                arrays = []
-                for frame in frames:
-                    val = frame.get(key)
-                    if isinstance(val, np.ndarray):
-                        arr = val.astype(self._resolve_dtype(key)).ravel()
-                        if repeat != 1:
-                            arr = np.repeat(arr, repeat)
-                        arrays.append(arr)
-                    elif val is not None:
-                        arrays.append(
-                            np.array([float(val)], dtype=self._resolve_dtype(key))
-                        )
-                    else:
-                        ref = next(
-                            (
-                                f[key]
-                                for f in frames
-                                if isinstance(f.get(key), np.ndarray)
-                            ),
-                            None,
-                        )
-                        if ref is not None:
-                            size = ref.size * repeat if repeat != 1 else ref.size
-                            arrays.append(
-                                np.zeros(size, dtype=self._resolve_dtype(key))
-                            )
-                        else:
-                            arrays.append(np.zeros(1, dtype=self._resolve_dtype(key)))
-                result[key] = np.stack(arrays)
-            elif key in self._requirements:
-                ndof = self._requirements[key]["ndof"]
-                atomic = self._requirements[key]["atomic"]
-                default = self._requirements[key]["default"]
-                if atomic:
-                    shape = (nframes, natoms * ndof * repeat)
+        for key in all_keys:
+            has_key = all(_frame_source_available(frame, key) for frame in frames)
+            requirement = self._requirements.get(key)
+            if requirement is None and not has_key:
+                # An unregistered field that only part of the group carries
+                # has no complete column and remains outside the result.
+                continue
+            arrays: list[np.ndarray] = []
+            for frame in frames:
+                value = frame.get(key)
+                if isinstance(value, np.ndarray):
+                    arrays.append(value.astype(self._resolve_dtype(key)).ravel())
+                elif value is not None:
+                    arrays.append(
+                        np.array([float(value)], dtype=self._resolve_dtype(key))
+                    )
                 else:
-                    shape = (nframes, ndof * repeat)
-                result[key] = np.full(shape, default, dtype=self._resolve_dtype(key))
+                    raise RuntimeError(
+                        f"Resolved LMDB field {key!r} is absent in frame "
+                        f"{frame.get('fid', '<unknown>')} of {self.lmdb_path}"
+                    )
+            result[key] = np.stack(arrays)
+            result[f"find_{key}"] = 1.0 if has_key else 0.0
 
         return result
 
@@ -2802,18 +4251,18 @@ class LmdbTestDataNlocView:
     """Expose one stack-compatible subset of :class:`LmdbTestData`.
 
     The underlying :class:`LmdbTestData` groups frames by atom count. This
-    view fixes one ``nloc`` and can additionally select a homogeneous
-    label-availability subgroup. All other attributes (``pbc``,
-    ``mixed_type``, …) are forwarded to the underlying object. It lets
-    downstream consumers that expect a ``DeepmdData``-style system work on
-    mixed-nloc or partially labeled LMDB datasets without vector find flags.
+    view fixes one ``nloc`` and can additionally select a label-compatible
+    subgroup within it. All other attributes (``pbc``, ``mixed_type``, …)
+    are forwarded to the underlying object. It lets downstream consumers
+    that expect a ``DeepmdData``-style system work on mixed-nloc or
+    partially labeled LMDB datasets without vector find flags.
     """
 
     def __init__(
         self,
         lmdb_test_data: "LmdbTestData",
         nloc: int,
-        frame_indices: list[int] | None = None,
+        frame_indices: Sequence[int] | None = None,
     ) -> None:
         self._inner = lmdb_test_data
         self._nloc = nloc
@@ -2833,18 +4282,12 @@ def iter_test(
         chunk_atoms: int,
         numb_test: float = float("inf"),
     ) -> Iterator[dict[str, Any]]:
-        """Yield this group's frames in chunks."""
-        if self._frame_indices is not None:
-            frame_indices = self._frame_indices
-            if np.isfinite(numb_test):
-                frame_indices = frame_indices[: int(numb_test)]
-            step = max(1, int(chunk_atoms) // max(1, self._nloc))
-            return (
-                self._inner.get_test_by_indices(frame_indices[begin : begin + step])
-                for begin in range(0, len(frame_indices), step)
-            )
+        """Yield this group's frames in chunks, as :meth:`get_test` selects them."""
         return self._inner.iter_test(
-            chunk_atoms=chunk_atoms, numb_test=numb_test, nloc=self._nloc
+            chunk_atoms=chunk_atoms,
+            numb_test=numb_test,
+            nloc=self._nloc,
+            frame_indices=self._frame_indices,
         )
 
 
diff --git a/deepmd/dpmodel/utils/neighbor_graph/__init__.py b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
index c047d37e97..32c0905a2d 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/__init__.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/__init__.py
@@ -44,6 +44,8 @@
     GraphLayout,
     NeighborGraph,
     apply_pair_exclusion,
+    compact_nodes,
+    expand_node_values,
     frame_id_from_n_node,
     node_ownership_mask,
     node_validity_mask,
@@ -75,8 +77,10 @@
     "build_neighbor_graph_ase",
     "canonicalize_neighbor_graph",
     "center_edge_pairs",
+    "compact_nodes",
     "edge_env_mat",
     "edge_force_virial",
+    "expand_node_values",
     "frame_id_from_n_node",
     "from_dense_quartet",
     "graph_angle_cos",
diff --git a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py
index 8c7288c7b4..42fae432a8 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/ase_builder.py
@@ -169,7 +169,7 @@ def _to_cpu_numpy(x: Any) -> np.ndarray:
         coord,
         box,
         nframe_all,
-        nloc,
+        np.full(nf, nloc, dtype=np.int64),
         layout=layout,
     )
     if pair_excl is not None:
diff --git a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py
index d8c8f73c27..b210ef48de 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/from_ijs.py
@@ -3,7 +3,7 @@
 
 ``neighbor_graph_from_ijs`` is the canonical sparse converter: it takes an
 already-built sparse edge list -- per-edge center ``i``, neighbor ``j`` (both
-per-frame LOCAL indices in ``[0, nloc)``) and integer periodic-image shift ``S``
+indices within their own frame) and integer periodic-image shift ``S``
 -- and emits a :class:`NeighborGraph` whose ``edge_vec`` is recomputed
 DIFFERENTIABLY from ``coord``/``box`` (it never trusts the builder's distance
 vectors). It is the format-conversion step shared by every O(N) search backend
@@ -46,7 +46,7 @@ def neighbor_graph_from_ijs(
     coord: Array,
     box: Array | None,
     nframe_id: Array,
-    nloc: int,
+    n_node: Array,
     layout: GraphLayout | None = None,
     *,
     with_csr: bool = False,
@@ -61,19 +61,22 @@ def neighbor_graph_from_ijs(
     Parameters
     ----------
     i
-        (E,) int per-edge center, per-frame LOCAL index in ``[0, nloc)``.
+        (E,) int per-edge center, index within its own frame.
     j
-        (E,) int per-edge neighbor, per-frame LOCAL index in ``[0, nloc)``.
+        (E,) int per-edge neighbor, index within its own frame.
     S
         (E, 3) int periodic-image shift: the neighbor sits at ``coord[j] + S @ box``.
     coord
-        (nf, nloc, 3) local coordinates.
+        (N, 3) local coordinates, frame-major over ``n_node``.
     box
         (nf, 3, 3) simulation cell, or ``None`` for non-periodic (``S`` ignored).
     nframe_id
         (E,) int frame index of each edge.
-    nloc
-        number of local atoms per frame (used for the frame-major node offset).
+    n_node
+        (nf,) int atoms per frame. Frames occupy contiguous blocks of the node
+        axis in order, so the prefix sums of this vector are the frame offsets
+        that turn a within-frame index into a node index. A batch padded to a
+        common width is the special case where every entry is that width.
     layout
         edge-axis length policy; ``None`` => dynamic (torch) with ``min_edges`` guards.
     with_csr
@@ -86,24 +89,24 @@ def neighbor_graph_from_ijs(
     Returns
     -------
     NeighborGraph
-        ``edge_index = [j + nframe_id*nloc, i + nframe_id*nloc]`` (src=neighbor,
-        dst=center); ``edge_vec = coord[j] + S@box - coord[i]``; ``n_node`` is
-        ``nloc`` per frame.
+        ``edge_index`` holds node indices (src=neighbor, dst=center) and
+        ``edge_vec = coord[j] + S@box - coord[i]``.
     """
     if layout is None:
         layout = GraphLayout()
     with_csr = with_csr or canonicalize
     xp = array_api_compat.array_namespace(coord)
     dev = array_api_compat.device(coord)
-    nf = coord.shape[0]
-    coord = xp.reshape(coord, (nf, nloc, 3))
+    n_node = xp.astype(xp.asarray(n_node, device=dev), xp.int64)
+    nf = n_node.shape[0]
+    coord_flat = xp.reshape(coord, (-1, 3))
     i = xp.astype(xp.asarray(i, device=dev), xp.int64)
     j = xp.astype(xp.asarray(j, device=dev), xp.int64)
     nframe_id = xp.astype(xp.asarray(nframe_id, device=dev), xp.int64)
-    # flat frame-major node indices
-    i_flat = i + nframe_id * nloc
-    j_flat = j + nframe_id * nloc
-    coord_flat = xp.reshape(coord, (nf * nloc, 3))
+    # Within-frame indices become node indices through the frame offsets.
+    offset = xp.take(xp.cumulative_sum(n_node) - n_node, nframe_id, axis=0)
+    i_flat = i + offset
+    j_flat = j + offset
     r_i = xp.take(coord_flat, i_flat, axis=0)
     r_j = xp.take(coord_flat, j_flat, axis=0)
     edge_vec = r_j - r_i
@@ -120,7 +123,6 @@ def neighbor_graph_from_ijs(
     edge_index, edge_vec, edge_mask = pad_and_guard_edges(
         edge_index, edge_vec, layout.edge_capacity, layout.min_edges
     )
-    n_node = xp.full((nf,), nloc, dtype=xp.int64, device=dev)
     if not with_csr:
         return NeighborGraph(
             n_node=n_node,
@@ -140,7 +142,7 @@ def neighbor_graph_from_ijs(
         edge_index,
         edge_vec,
         edge_mask,
-        nf * nloc,
+        int(coord_flat.shape[0]),
         canonicalize=canonicalize,
     )
     return NeighborGraph(
diff --git a/deepmd/dpmodel/utils/neighbor_graph/graph.py b/deepmd/dpmodel/utils/neighbor_graph/graph.py
index f702c2c6a1..baeef186ef 100644
--- a/deepmd/dpmodel/utils/neighbor_graph/graph.py
+++ b/deepmd/dpmodel/utils/neighbor_graph/graph.py
@@ -396,6 +396,129 @@ def apply_pair_exclusion(
     return out
 
 
+def compact_nodes(
+    graph: NeighborGraph, node_mask: Array
+) -> tuple[NeighborGraph, Array]:
+    """Renumber a graph onto the subset of its nodes selected by a mask.
+
+    A batch whose frames hold unequal atom counts reaches the model as a
+    rectangular tensor padded with phantom atoms, and the graph builders lay
+    their node axis out over that padded shape. Every builder already refuses
+    an edge touching a phantom, so the phantoms survive as isolated nodes
+    whose only effect is to make the network evaluate them. Dropping them here
+    leaves each frame occupying exactly as many nodes as it has real atoms,
+    which is the layout ``n_node`` was always able to express.
+
+    Node order is preserved, so each frame keeps one contiguous block and the
+    frame-major invariant every per-frame reduction relies on still holds.
+
+    Parameters
+    ----------
+    graph : NeighborGraph
+        Graph whose node axis is to be compacted. It must carry no edge
+        incident on a masked-out node, and no multi-rank halo split.
+    node_mask : Array
+        Boolean mask over the flat node axis with shape ``(N,)``, ``True`` for
+        the nodes to retain.
+
+    Returns
+    -------
+    NeighborGraph
+        The graph over the retained nodes, with ``n_node`` recounted and
+        ``edge_index`` renumbered. CSR views, when present, are rebuilt.
+    Array
+        Positions of the retained nodes on the original axis, shape
+        ``(n_kept,)``. Gathering a per-node tensor with it moves that tensor
+        onto the compacted axis; :func:`expand_node_values` inverts it.
+
+    Raises
+    ------
+    ValueError
+        If the graph carries a halo split, or if some edge is incident on a
+        node the mask drops.
+    """
+    import dataclasses
+
+    from deepmd.dpmodel.utils.neighbor_graph.segment import (
+        segment_sum,
+    )
+
+    if graph.n_local is not None:
+        raise ValueError("cannot compact the node axis of a local-plus-halo graph")
+    xp = array_api_compat.array_namespace(graph.n_node, node_mask)
+    device = array_api_compat.device(graph.n_node)
+    n_total = node_mask.shape[0]
+
+    keep_index = xp.reshape(xp.nonzero(node_mask)[0], (-1,))
+    # Position of each retained node on the compacted axis, and -1 for the
+    # nodes that go away. Renumbering by a prefix sum keeps the frame blocks
+    # contiguous and in order.
+    renumber = xp.cumulative_sum(xp.astype(node_mask, xp.int64)) - 1
+    renumber = xp.where(node_mask, renumber, xp.asarray(-1, device=device))
+
+    frame_id = frame_id_from_n_node(graph.n_node, n_total=n_total)
+    n_node = xp.astype(
+        segment_sum(xp.astype(node_mask, xp.int64), frame_id, graph.n_node.shape[0]),
+        graph.n_node.dtype,
+    )
+
+    edge_index = xp.take(renumber, xp.reshape(graph.edge_index, (-1,)), axis=0)
+    edge_index = xp.reshape(edge_index, graph.edge_index.shape)
+    # A dropped node with an edge would leave -1 behind on a real edge, which
+    # would silently address the last node of the compacted axis downstream.
+    if bool(xp.any(xp.logical_and(edge_index < 0, graph.edge_mask[None, :]))):
+        raise ValueError("cannot compact a node that still carries an edge")
+    edge_index = xp.astype(
+        xp.maximum(edge_index, xp.asarray(0, device=device)), graph.edge_index.dtype
+    )
+
+    compacted = dataclasses.replace(
+        graph,
+        n_node=n_node,
+        edge_index=edge_index,
+        destination_order=None,
+        destination_row_ptr=None,
+        source_order=None,
+        source_row_ptr=None,
+        destination_sorted=False,
+    )
+    if graph.destination_row_ptr is not None:
+        from deepmd.dpmodel.utils.neighbor_graph.csr import (
+            attach_edge_csr,
+        )
+
+        compacted = attach_edge_csr(compacted, int(keep_index.shape[0]))
+    return compacted, keep_index
+
+
+def expand_node_values(values: Array, keep_index: Array, n_total: int) -> Array:
+    """Scatter a compacted per-node tensor back onto a padded node axis.
+
+    Inverse of the gather that :func:`compact_nodes` describes. Positions the
+    compaction dropped read as zero, which is what the dropped nodes
+    contributed to every physical quantity in the first place.
+
+    Parameters
+    ----------
+    values : Array
+        Per-node tensor on the compacted axis with shape ``(n_kept, ...)``.
+    keep_index : Array
+        Positions of the retained nodes, as returned by :func:`compact_nodes`.
+    n_total : int
+        Size of the padded node axis to scatter onto.
+
+    Returns
+    -------
+    Array
+        Tensor of shape ``(n_total, ...)``.
+    """
+    from deepmd.dpmodel.utils.neighbor_graph.segment import (
+        segment_sum,
+    )
+
+    return segment_sum(values, keep_index, n_total)
+
+
 def node_validity_mask(n_node: Array, n_total: int) -> Array:
     """Derive the (n_total,) real-vs-padding node mask from per-frame counts.
 
diff --git a/deepmd/pt/loss/dens.py b/deepmd/pt/loss/dens.py
index 03e1c297e4..5c6fb3edb5 100644
--- a/deepmd/pt/loss/dens.py
+++ b/deepmd/pt/loss/dens.py
@@ -4,7 +4,6 @@
 )
 
 import torch
-import torch.nn.functional as F
 
 from deepmd.pt.loss.ener import (
     EnergyStdLoss,
@@ -164,9 +163,18 @@ def _prepare_dens_inputs(
         torch.Tensor,
         torch.Tensor,
         torch.Tensor,
+        torch.Tensor,
         bool,
     ]:
-        """Build noisy coordinates and mixed targets for one forward pass."""
+        """Build noisy coordinates and mixed targets for one forward pass.
+
+        Returns the corrupted and clean atom subsets as two masks. Both are
+        restricted to real atoms: a mixed-nloc batch is padded to a rectangular
+        shape with phantom slots that carry no physical site, and a phantom is
+        neither corrupted nor clean. The two masks are built independently
+        rather than as complements of one another, since complements over the
+        padded width would hand every phantom to the clean subset.
+        """
         atype = input_dict["atype"]
         nf, nloc = atype.shape[:2]
         coord_raw = input_dict["coord"]
@@ -185,22 +193,22 @@ def _prepare_dens_inputs(
             ).item()
             < self.dens_prob
         )
+        real_mask = atype >= 0
         noise_mask = torch.zeros((nf, nloc), dtype=torch.bool, device=coord.device)
         noise_vec = torch.zeros_like(coord)
         if use_dens:
             if self.dens_corrupt_ratio is None:
-                noise_mask = torch.ones(
-                    (nf, nloc), dtype=torch.bool, device=coord.device
-                )
+                noise_mask = real_mask
             else:
                 noise_mask = (
                     torch.rand(
                         (nf, nloc), dtype=GLOBAL_PT_FLOAT_PRECISION, device=coord.device
                     )
                     < self.dens_corrupt_ratio
-                )
+                ) & real_mask
             noise_vec = torch.randn_like(coord) * self.dens_std
             noise_vec = noise_vec * noise_mask.unsqueeze(-1)
+        clean_mask = real_mask & ~noise_mask
         coord_model = coord + noise_vec
 
         # DeNS predicts normalized noise epsilon / sigma for corrupted atoms.
@@ -214,7 +222,14 @@ def _prepare_dens_inputs(
         model_input["noise_mask"] = noise_mask
         if use_dens:
             model_input["force_input"] = force_label
-        return model_input, force_label, noise_target, noise_mask, use_dens
+        return (
+            model_input,
+            force_label,
+            noise_target,
+            noise_mask,
+            clean_mask,
+            use_dens,
+        )
 
     @staticmethod
     def _get_sezm_atomic_model(model: torch.nn.Module) -> Any:
@@ -270,8 +285,14 @@ def forward(
         learning_rate: float,
         mae: bool = False,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor, dict[str, torch.Tensor]]:
-        """Return loss on SeZM `dens` energy and direct-force/noise outputs."""
-        model_input, force_label, noise_target, noise_mask, use_dens = (
+        """Return loss on SeZM `dens` energy and direct-force/noise outputs.
+
+        ``natoms`` is the padded width of the batch and is superseded here by
+        the per-frame real atom count read off ``atype``; it remains in the
+        signature to satisfy the shared loss interface.
+        """
+        del natoms
+        model_input, force_label, noise_target, noise_mask, clean_mask, use_dens = (
             self._prepare_dens_inputs(
                 input_dict,
                 label,
@@ -288,7 +309,14 @@ def forward(
 
         loss = force_label.new_zeros((), dtype=env.GLOBAL_PT_FLOAT_PRECISION)
         more_loss: dict[str, torch.Tensor] = {}
-        atom_norm = 1.0 / natoms
+        # The energy terms are extensive, so a frame's residual is divided by
+        # that frame's own real atom count; a mixed-nloc batch pads to a common
+        # width with phantom atoms that belong to no system. The factor is
+        # therefore applied inside the mean over frames, where on a batch of
+        # uniform atom count it reduces to the scalar ``1 / natoms``.
+        inv_natoms = 1.0 / (input_dict["atype"] >= 0).sum(dim=-1, keepdim=True).to(
+            dtype=force_label.dtype
+        )  # [nf, 1]
 
         if self.has_e and "energy" in model_pred and "energy" in label:
             energy_pred = model_pred.get("energy_norm", model_pred["energy"])
@@ -318,24 +346,23 @@ def forward(
                         l2_ener_loss.detach(),
                         find_energy,
                     )
-                loss += atom_norm * (pref_e * l2_ener_loss)
-                rmse_e = (
-                    torch.mean(torch.square(energy_pred_phys - energy_label)).sqrt()
-                    * atom_norm
+                loss += pref_e * torch.mean(
+                    torch.square(energy_pred - energy_label_norm) * inv_natoms
                 )
+                rmse_e = torch.mean(
+                    torch.square(energy_pred_phys - energy_label) * inv_natoms**2
+                ).sqrt()
                 more_loss["rmse_e"] = self.display_if_exist(
                     rmse_e.detach(),
                     find_energy,
                 )
             elif self.loss_func == "mae":
-                l1_ener_loss = F.l1_loss(
-                    energy_pred.reshape(-1),
-                    energy_label_norm.reshape(-1),
-                    reduction="mean",
+                l1_ener_loss = torch.mean(
+                    torch.abs(energy_pred - energy_label_norm) * inv_natoms
                 )
-                loss += atom_norm * (pref_e * l1_ener_loss)
-                mae_e = (
-                    torch.mean(torch.abs(energy_pred_phys - energy_label)) * atom_norm
+                loss += pref_e * l1_ener_loss
+                mae_e = torch.mean(
+                    torch.abs(energy_pred_phys - energy_label) * inv_natoms
                 )
                 more_loss["mae_e"] = self.display_if_exist(
                     mae_e.detach(),
@@ -346,8 +373,8 @@ def forward(
                     f"Loss type {self.loss_func} is not implemented for `dens` energy loss."
                 )
             if mae:
-                mae_e = (
-                    torch.mean(torch.abs(energy_pred_phys - energy_label)) * atom_norm
+                mae_e = torch.mean(
+                    torch.abs(energy_pred_phys - energy_label) * inv_natoms
                 )
                 more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy)
                 mae_e_all = torch.mean(torch.abs(energy_pred_phys - energy_label))
@@ -386,10 +413,17 @@ def forward(
             else:
                 force_pred_phys = atomic_model.denorm_dens_force(clean_force_pred_norm)
             force_target_norm = atomic_model.norm_dens_force(force_label)
-            clean_mask = ~noise_mask
             noise_only_mask = noise_mask if use_dens else torch.zeros_like(noise_mask)
-            clean_fraction = clean_mask.to(dtype=GLOBAL_PT_FLOAT_PRECISION).mean()
-            noise_fraction = noise_only_mask.to(dtype=GLOBAL_PT_FLOAT_PRECISION).mean()
+            # Each subset loss averages over its own atoms and is then scaled
+            # by that subset's share, so the two recombine into one average
+            # over the real atoms. The share is taken against the real-atom
+            # count rather than the padded width, or phantom slots would
+            # dilute both terms by the padding fraction.
+            real_count = (
+                (input_dict["atype"] >= 0).sum().to(dtype=GLOBAL_PT_FLOAT_PRECISION)
+            )
+            clean_fraction = clean_mask.sum() / real_count
+            noise_fraction = noise_only_mask.sum() / real_count
             clean_force_loss = self._compute_force_subset_loss(
                 clean_force_pred_norm[clean_mask].reshape(-1, 3),
                 force_target_norm[clean_mask].reshape(-1, 3),
diff --git a/deepmd/pt/loss/dos.py b/deepmd/pt/loss/dos.py
index 7046f73687..8f24e94d22 100644
--- a/deepmd/pt/loss/dos.py
+++ b/deepmd/pt/loss/dos.py
@@ -157,7 +157,7 @@ def forward(
                 local_tensor_pred_dos - local_tensor_label_dos
             )  # [nf, natoms, numb_dos]
             if "mask" in model_pred:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=numb_dos).
                 maskf = model_pred["mask"].to(diff.dtype)  # [nf, natoms]
                 l2_local_loss_dos = masked_atom_mean(
                     torch.square(diff), maskf, self.numb_dos
@@ -186,7 +186,7 @@ def forward(
                 local_tensor_pred_cdf - local_tensor_label_cdf
             )  # [nf, natoms, numb_dos]
             if "mask" in model_pred:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=numb_dos).
                 maskf = model_pred["mask"].to(diff.dtype)  # [nf, natoms]
                 l2_local_loss_cdf = masked_atom_mean(
                     torch.square(diff), maskf, self.numb_dos
diff --git a/deepmd/pt/loss/ener.py b/deepmd/pt/loss/ener.py
index 468b95f1eb..7e10e99466 100644
--- a/deepmd/pt/loss/ener.py
+++ b/deepmd/pt/loss/ener.py
@@ -402,7 +402,6 @@ def forward(
                     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)
                         # Masked MSE computed for rmse_f display regardless of use_huber.
                         l2_f_masked = masked_atom_mean(
                             torch.square(diff_f_3d), maskf, 3
@@ -410,6 +409,12 @@ def forward(
                         if not self.use_huber:
                             loss += (pref_f * l2_f_masked).to(GLOBAL_PT_FLOAT_PRECISION)
                         else:
+                            # ``f_use_norm`` selects the residual an atom
+                            # contributes: three independent components, or the
+                            # single L2 norm of its force-error vector. That
+                            # choice sets the label count per atom, which is
+                            # exactly the ``ncomp`` the pooled reduction
+                            # divides by.
                             if not self.f_use_norm:
                                 abs_e = torch.abs(diff_f_3d)
                                 quad = 0.5 * torch.square(diff_f_3d)
@@ -418,9 +423,8 @@ def forward(
                                 )
                                 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
+                                )  # [nf, nloc, 3]
+                                huber_ncomp = 3
                             else:
                                 diff_3 = (force_label - force_pred).reshape(
                                     _nf, _nloc, 3
@@ -433,13 +437,13 @@ def forward(
                                 lin_n = self._huber_delta_force * (
                                     abs_n - 0.5 * self._huber_delta_force
                                 )
-                                huber_n = torch.where(
+                                huber_elem = 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)
+                                ).reshape(_nf, _nloc, 1)
+                                huber_ncomp = 1
+                            l_huber_masked = masked_atom_mean(
+                                huber_elem, maskf, huber_ncomp
+                            )
                             loss += pref_f * l_huber_masked
                     else:
                         if not self.use_huber:
@@ -486,10 +490,10 @@ def forward(
                             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)
+                            # One L2 norm per atom, hence one label per atom.
+                            l1_f_masked = masked_atom_mean(
+                                norm_2d.reshape(_nf, _nloc, 1), maskf, 1
+                            )
                         more_loss["mae_f"] = self.display_if_exist(
                             l1_f_masked.detach(), find_force
                         )
@@ -531,7 +535,8 @@ def forward(
                 find_atom_pref = (
                     label.get("find_atom_pref", 0.0) if not self.use_default_pf else 1.0
                 )
-                pref_pf = pref_pf * find_atom_pref
+                effective_find_pf = find_force * find_atom_pref
+                pref_pf = pref_pf * effective_find_pf
                 atom_pref_reshape = atom_pref.reshape(-1)
 
                 if self.loss_func == "mse":
@@ -540,7 +545,7 @@ def forward(
                     ).mean()
                     if not self.inference:
                         more_loss["l2_pref_force_loss"] = self.display_if_exist(
-                            l2_pref_force_loss.detach(), find_atom_pref
+                            l2_pref_force_loss.detach(), effective_find_pf
                         )
                     if maskf is not None:
                         # Idiom 1 with pref weight (ncomp=3).
@@ -552,7 +557,7 @@ def forward(
                         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
+                            rmse_pf.detach(), effective_find_pf
                         )
                     else:
                         loss += (pref_pf * l2_pref_force_loss).to(
@@ -560,7 +565,7 @@ def forward(
                         )
                         rmse_pf = l2_pref_force_loss.sqrt()
                         more_loss["rmse_pf"] = self.display_if_exist(
-                            rmse_pf.detach(), find_atom_pref
+                            rmse_pf.detach(), effective_find_pf
                         )
                 elif self.loss_func == "mae":
                     l1_pref_force_loss = (torch.abs(diff_f) * atom_pref_reshape).mean()
@@ -572,14 +577,14 @@ def forward(
                         )
                         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
+                            l1_pf_masked.detach(), effective_find_pf
                         )
                     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
+                            l1_pref_force_loss.detach(), effective_find_pf
                         )
                 else:
                     raise NotImplementedError(
@@ -589,7 +594,8 @@ def forward(
             if self.has_gf and "drdq" in label:
                 drdq = label["drdq"]
                 find_drdq = label.get("find_drdq", 0.0)
-                pref_gf = pref_gf * find_drdq
+                effective_find_gf = find_force * find_drdq
+                pref_gf = pref_gf * effective_find_gf
                 if maskf is not None:
                     # Mask per-atom forces before projecting onto generalized coords.
                     f_3d = force_pred.reshape(_nf, _nloc, 3) * maskf.reshape(
@@ -623,12 +629,12 @@ def forward(
                 l2_gen_force_loss = torch.square(diff_gen_force).mean()
                 if not self.inference:
                     more_loss["l2_gen_force_loss"] = self.display_if_exist(
-                        l2_gen_force_loss.detach(), find_drdq
+                        l2_gen_force_loss.detach(), effective_find_gf
                     )
                 loss += (pref_gf * l2_gen_force_loss).to(GLOBAL_PT_FLOAT_PRECISION)
                 rmse_gf = l2_gen_force_loss.sqrt()
                 more_loss["rmse_gf"] = self.display_if_exist(
-                    rmse_gf.detach(), find_drdq
+                    rmse_gf.detach(), effective_find_gf
                 )
 
         if self.has_v and "virial" in model_pred and "virial" in label:
@@ -726,7 +732,6 @@ def forward(
                     # Idiom 1 (per-atom masked mean, ncomp=1).
                     ae_2d = atom_ener.reshape(_nf, _nloc)
                     ae_hat_2d = atom_ener_label.reshape(_nf, _nloc)
-                    per_frame_dof = maskf.sum(dim=-1)  # [nf], kept for huber branch
                     l2_ae_masked = masked_atom_mean(
                         torch.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1
                     )
@@ -742,8 +747,7 @@ def forward(
                         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)
+                        l_huber_ae = masked_atom_mean(huber_ae[:, :, None], maskf, 1)
                         loss += pref_ae * l_huber_ae
                     rmse_ae = l2_ae_masked.sqrt()
                     more_loss["rmse_ae"] = self.display_if_exist(
@@ -889,6 +893,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     high_prec=False,
                     repeat=3,
                     default=1.0,
+                    source_policy="default" if self.use_default_pf else "tracked",
                 )
             )
         if self.has_gf > 0:
@@ -910,6 +915,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     must=False,
                     high_prec=False,
                     default=1.0,
+                    source_policy="default",
                 )
             )
         if self.has_h:
diff --git a/deepmd/pt/loss/tensor.py b/deepmd/pt/loss/tensor.py
index 660cbd0337..701cf62228 100644
--- a/deepmd/pt/loss/tensor.py
+++ b/deepmd/pt/loss/tensor.py
@@ -133,7 +133,7 @@ def forward(
             )
             diff = diff * atomic_weight
             if "mask" in model_pred:
-                # idiom 1: per-frame masked mean, then average over frames
+                # Idiom 1 (per-atom masked mean, ncomp=tensor_size).
                 maskf = model_pred["mask"].to(diff.dtype)  # [nf, natoms]
                 diff3d = diff.reshape(
                     local_tensor_pred.shape[0], natoms, self.tensor_size
@@ -211,6 +211,7 @@ def label_requirement(self) -> list[DataRequirementItem]:
                     must=False,
                     high_prec=False,
                     default=1.0,
+                    source_policy="default",
                 )
             )
         return label_requirement
diff --git a/deepmd/pt/model/model/sezm_model.py b/deepmd/pt/model/model/sezm_model.py
index f7c8337433..5d9d182d79 100644
--- a/deepmd/pt/model/model/sezm_model.py
+++ b/deepmd/pt/model/model/sezm_model.py
@@ -770,6 +770,20 @@ def __init__(
     # Forward Methods
     # =========================================================================
 
+    def _sanitize_atom_types(
+        self,
+        atype: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return lookup-safe atom types and the physical-atom mask.
+
+        Phantom atoms use a negative type sentinel for rectangular mixed-nloc
+        padding. Every type-table lookup receives the safe tensor, while the
+        mask keeps the substituted row from acquiring physical meaning.
+        """
+        real_atom = self.atomic_model.make_atom_mask(atype)
+        safe_atype = torch.where(real_atom, atype, torch.zeros_like(atype))
+        return safe_atype, real_atom
+
     def forward(
         self,
         coord: Float[Tensor, "nf nloc 3"] | Float[Tensor, "nf nloc_x3"],
@@ -1497,11 +1511,36 @@ def core_compute(
         ):
             spin = spin.detach().requires_grad_(True)
 
-        descriptor_atype = extended_atype if comm_dict is not None else atype
-        if descriptor_atype is None:
+        # === Atom mask ===
+        # Phantom atoms (atype < 0) pad a mixed-nloc batch to a rectangular
+        # shape and stand for no physical site. The edge builders keep them out
+        # of every edge, so the networks see them as isolated nodes, and the
+        # mask below zeroes whatever those nodes produce. Their type is still
+        # read as a table index by the type embedding, the fitting bias and the
+        # exclusion masks, none of which has a row for it. ``atype`` is
+        # therefore rebound to a sanitized copy for the remainder of this
+        # method, which is where a phantom would otherwise reach such a table.
+        # The substitute type is immaterial: it reaches no real atom through
+        # any edge, and the mask discards everything it produces.
+        atype, real_atom = self._sanitize_atom_types(atype)
+        atom_mask = real_atom.to(torch.int32)
+        if self.atomic_model.atom_excl is not None:
+            atom_mask = atom_mask * self.atomic_model.atom_excl(atype)
+
+        if comm_dict is None:
+            descriptor_atype = atype
+            descriptor_real_atom = real_atom
+        elif extended_atype is None:
             raise ValueError("`extended_atype` is required with `comm_dict`.")
+        else:
+            # Ghost atoms carry the type of the local atom they image, so this
+            # sanitizes the phantoms among them on the same grounds as above.
+            descriptor_atype, descriptor_real_atom = self._sanitize_atom_types(
+                extended_atype
+            )
         inter_potential_edge_mask = self._make_inter_potential_edge_mask(
             descriptor_atype,
+            descriptor_real_atom,
             edge_index,
             edge_mask,
         )
@@ -1526,11 +1565,6 @@ def core_compute(
                 nloc=nloc,
             )
 
-        # === Atom mask ===
-        atom_mask = self.atomic_model.make_atom_mask(atype).to(torch.int32)
-        if self.atomic_model.atom_excl is not None:
-            atom_mask = atom_mask * self.atomic_model.atom_excl(atype)
-
         # === Step 3. Fitting net ===
         # The same fitting forward serves both modes; ``embedding_only`` only asks
         # it to also return the last hidden activation.
@@ -1555,9 +1589,14 @@ def core_compute(
             structural_feature = (
                 atomic_feature * atom_mask[:, :, None].to(atomic_feature.dtype)
             ).sum(dim=1)
+            # The per-atom outputs are zeroed on the phantom rows alone, so
+            # that a padded batch reports nothing for a slot holding no atom.
+            # Excluded atoms, which ``atom_mask`` also covers, keep their
+            # per-atom embeddings and are dropped from the pooled sum only.
+            real = real_atom[:, :, None]
             return {
-                "descriptor": descriptor,
-                "atomic_feature": atomic_feature,
+                "descriptor": descriptor * real.to(descriptor.dtype),
+                "atomic_feature": atomic_feature * real.to(atomic_feature.dtype),
                 "structural_feature": structural_feature,
             }
 
@@ -2810,7 +2849,7 @@ def build_edge_list_from_nlist(
             1:1 with ``edge_index`` and ``edge_vec``.
         """
         nloc = nlist.shape[1]
-        atype = torch.empty(
+        atype = torch.zeros(
             (nlist.shape[0], nloc),
             dtype=torch.long,
             device=extended_coord.device,
@@ -3206,6 +3245,7 @@ def reset_head_for_mode(self, mode: str) -> None:
     def _make_inter_potential_edge_mask(
         self,
         atype: torch.Tensor,
+        real_atom: torch.Tensor,
         edge_index: torch.Tensor,
         edge_mask: torch.Tensor,
     ) -> torch.Tensor:
@@ -3214,7 +3254,9 @@ def _make_inter_potential_edge_mask(
         Parameters
         ----------
         atype
-            Atom types with shape (nf, nall).
+            Lookup-safe atom types with shape (nf, nall).
+        real_atom
+            Physical-atom mask with shape (nf, nall).
         edge_index
             Edge source and destination indices with shape (2, E).
         edge_mask
@@ -3231,7 +3273,12 @@ def _make_inter_potential_edge_mask(
         src = edge_index[0]
         dst = edge_index[1]
         atype_flat = atype.reshape(-1)
-        keep = edge_mask
+        real_atom_flat = real_atom.reshape(-1)
+        keep = (
+            edge_mask
+            & real_atom_flat.index_select(0, src)
+            & real_atom_flat.index_select(0, dst)
+        )
 
         descriptor = self.atomic_model.descriptor
         if descriptor.exclude_types:
@@ -3239,10 +3286,7 @@ def _make_inter_potential_edge_mask(
 
         atom_excl = self.atomic_model.atom_excl
         if atom_excl is not None:
-            atom_is_present = self.atomic_model.make_atom_mask(atype)
-            safe_atype = torch.where(atom_is_present, atype, 0)
-            atom_is_included = atom_is_present & atom_excl(safe_atype).to(torch.bool)
-            atom_is_included = atom_is_included.reshape(-1)
+            atom_is_included = atom_excl(atype).to(torch.bool).reshape(-1)
             keep = (
                 keep
                 & atom_is_included.index_select(0, src)
diff --git a/deepmd/pt/model/model/sezm_native_spin_model.py b/deepmd/pt/model/model/sezm_native_spin_model.py
index aeec0ae74e..9cc14355b9 100644
--- a/deepmd/pt/model/model/sezm_native_spin_model.py
+++ b/deepmd/pt/model/model/sezm_native_spin_model.py
@@ -84,6 +84,15 @@ def __init__(
     # Forward Methods
     # =========================================================================
 
+    def _make_mask_mag(self, atype: torch.Tensor) -> torch.Tensor:
+        """Return the magnetic-type mask without indexing phantom sentinels."""
+        safe_atype, real_atom = self._sanitize_atom_types(atype)
+        spin_active = self.spin_mask.index_select(
+            0,
+            safe_atype.reshape(-1),
+        ).reshape(*atype.shape, 1)
+        return (spin_active > 0.0) & real_atom.unsqueeze(-1)
+
     def forward(
         self,
         coord: torch.Tensor,
@@ -100,8 +109,9 @@ def forward(
         ``mask_mag`` is built from the per-type spin gate on the local
         ``atype``; non-magnetic atoms already carry a zero magnetic force (the
         descriptor gates the spin embedding by type), so the force itself needs
-        no re-masking. This is the runtime counterpart of the static schema in
-        :meth:`translated_output_def`.
+        no re-masking. Phantom padding uses a lookup-safe type but is intersected
+        with the physical-atom mask, so it always reports ``False``. This is the
+        runtime counterpart of the static schema in :meth:`translated_output_def`.
         """
         model_ret = self.forward_common(
             coord,
@@ -113,14 +123,10 @@ def forward(
             charge_spin=charge_spin,
             spin=spin,
         )
-        nf, nloc = atype.shape[:2]
         model_predict: dict[str, torch.Tensor] = {
             "atom_energy": model_ret["energy"],
             "energy": model_ret["energy_redu"],
-            "mask_mag": self.spin_mask.index_select(0, atype.reshape(-1)).reshape(
-                nf, nloc, 1
-            )
-            > 0.0,
+            "mask_mag": self._make_mask_mag(atype),
         }
         if self.do_grad_r("energy"):
             model_predict["force"] = rearrange(
@@ -397,8 +403,9 @@ def _attach_spin_masks(
             Internal SeZM lower outputs; ``energy_derv_r_mag`` has the
             per-local-atom shape ``(nf, nloc, 1, 3)``.
         atype
-            Local atom types with shape ``(nf, nloc)``, used to build
-            ``mask_mag`` with shape ``(nf, nloc, 1)``.
+            Atom types aligned with the magnetic-force rows before optional
+            ghost padding, used to build ``mask_mag`` with a trailing singleton
+            axis. Negative phantom sentinels are allowed.
         nall
             Extended atom count the magnetic force is padded to.
 
@@ -423,12 +430,7 @@ def _attach_spin_masks(
         nf, nloc = derv_r_mag.shape[:2]
         ghost_pad = derv_r_mag.new_zeros(nf, nall - nloc, *derv_r_mag.shape[2:])
         model_ret["energy_derv_r_mag"] = torch.cat([derv_r_mag, ghost_pad], dim=1)
-        model_ret["mask_mag"] = (
-            self.spin_mask.index_select(0, atype.reshape(-1)).reshape(
-                atype.shape[0], atype.shape[1], 1
-            )
-            > 0.0
-        )
+        model_ret["mask_mag"] = self._make_mask_mag(atype)
         return model_ret
 
     # =========================================================================
diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py
index 4a25ab172e..c0adf618c8 100644
--- a/deepmd/pt/model/model/spin_model.py
+++ b/deepmd/pt/model/model/spin_model.py
@@ -554,6 +554,8 @@ def compute_or_load_stat(
             The lazy sampled function to get data frames from different data systems.
         stat_file_path
             The dictionary of paths to the statistics files.
+        preset_observed_type
+            Atom types known to occur in the sampled data, if already computed.
         """
 
         @functools.lru_cache
@@ -600,6 +602,8 @@ def forward_common(
         model_ret[f"{var_name}"] = torch.split(
             model_ret[f"{var_name}"], [nloc, nloc], dim=1
         )[0]
+        if "mask" in model_ret:
+            model_ret["mask"] = model_ret["mask"][:, :nloc]
         if self.backbone_model.do_grad_r(var_name):
             (
                 model_ret[f"{var_name}_derv_r"],
@@ -665,6 +669,8 @@ def forward_common_lower(
         model_ret[f"{var_name}"] = torch.split(
             model_ret[f"{var_name}"], [nloc, nloc], dim=1
         )[0]
+        if "mask" in model_ret:
+            model_ret["mask"] = model_ret["mask"][:, :nloc]
         if self.backbone_model.do_grad_r(var_name):
             (
                 model_ret[f"{var_name}_derv_r"],
@@ -737,6 +743,8 @@ def translated_output_def(self) -> dict[str, Any]:
             output_def["virial"].squeeze(-2)
             output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"])
             output_def["atom_virial"].squeeze(-2)
+        if "mask" in out_def_data:
+            output_def["mask"] = out_def_data["mask"]
         return output_def
 
     def forward(
@@ -764,6 +772,8 @@ def forward(
         model_predict["atom_energy"] = model_ret["energy"]
         model_predict["energy"] = model_ret["energy_redu"]
         model_predict["mask_mag"] = model_ret["mask_mag"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
         if self.backbone_model.do_grad_r("energy"):
             model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
             model_predict["force_mag"] = model_ret["energy_derv_r_mag"].squeeze(-2)
@@ -804,6 +814,8 @@ def forward_lower(
         model_predict["atom_energy"] = model_ret["energy"]
         model_predict["energy"] = model_ret["energy_redu"]
         model_predict["extended_mask_mag"] = model_ret["mask_mag"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
         if self.backbone_model.do_grad_r("energy"):
             model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
             model_predict["extended_force_mag"] = model_ret[
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
index 5a190fe5af..776bdece43 100644
--- a/deepmd/pt/train/training.py
+++ b/deepmd/pt/train/training.py
@@ -289,46 +289,36 @@ def get_data_loader(
             def get_dataloader_and_iter_lmdb(
                 _data: LmdbDataset,
             ) -> tuple[LmdbBatchDataLoader, Generator[Any, None, None]]:
-                if _data.mixed_batch:
-                    # TODO [mixed_batch=True]: Replace SameNlocBatchSampler with
-                    # RandomSampler(replacement=False) + padding collate_fn.
-                    # Changes needed:
-                    #   1. _collate_lmdb_batch: pad coord/force/atype to max_nloc,
-                    #      add "atom_mask" bool tensor (nframes, max_nloc)
-                    #   2. Use RandomSampler(_data, replacement=False) as sampler
-                    #   3. Use fixed batch_size in DataLoader (not batch_sampler)
-                    #   4. Model forward: apply atom_mask to descriptor/fitting
-                    #   5. Loss: mask out padded atoms in force loss
-                    raise NotImplementedError(
-                        "mixed_batch=True training is not yet supported."
-                    )
-                # mixed_batch=False: group frames by nloc, each batch same nloc.
-                # SameNlocBatchSampler yields list[int] per batch, all same nloc.
-                # Auto batch_size is computed per-nloc-group inside the sampler.
+                # The sampler yields one list[int] per batch. Every batch is
+                # homogeneous in label availability; whether it is also
+                # homogeneous in atom count follows from the dataset's
+                # batch_size rule, which the sampler reads off the reader.
                 from deepmd.dpmodel.utils.lmdb_data import (
-                    SameNlocBatchSampler,
+                    LmdbBatchSampler,
                 )
 
                 _block_targets = getattr(_data, "_block_targets", None)
+                _sampler_kwargs = {
+                    "shuffle": True,
+                    "seed": _training_params.get("seed"),
+                    "block_targets": _block_targets,
+                }
 
                 if self.world_size > 1:
                     from deepmd.dpmodel.utils.lmdb_data import (
-                        DistributedSameNlocBatchSampler,
+                        DistributedLmdbBatchSampler,
                     )
 
-                    _inner_sampler = DistributedSameNlocBatchSampler(
+                    _inner_sampler = DistributedLmdbBatchSampler(
                         _data._reader,
                         rank=self.rank,
                         world_size=self.world_size,
-                        shuffle=True,
-                        seed=_training_params.get("seed"),
-                        block_targets=_block_targets,
+                        **_sampler_kwargs,
                     )
                 else:
-                    _inner_sampler = SameNlocBatchSampler(
+                    _inner_sampler = LmdbBatchSampler(
                         _data._reader,
-                        shuffle=True,
-                        block_targets=_block_targets,
+                        **_sampler_kwargs,
                     )
 
                 _dataloader = LmdbBatchDataLoader(
@@ -519,6 +509,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                         must=False,
                         high_prec=False,
                         default=min_pair_dist,
+                        source_policy="derived",
                     )
                 )
             training_data.add_data_requirement(data_requirement)
@@ -597,6 +588,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR:
                             must=False,
                             high_prec=False,
                             default=min_pair_dist,
+                            source_policy="derived",
                         )
                     )
                 training_data[model_key].add_data_requirement(data_requirement)
@@ -1595,7 +1587,11 @@ def log_loss_valid(_task_key: str = "Default") -> dict:
                             task_key=_task_key,
                         )
                         # more_loss.update({"rmse": math.sqrt(loss)})
-                        natoms = int(input_dict["atype"].shape[-1])
+                        # The metrics are per-atom quantities, so each batch
+                        # weighs by the real atoms it holds summed over its
+                        # frames. Phantom atoms (atype < 0), which pad a
+                        # mixed-nloc batch, contribute to none of them.
+                        natoms = int((input_dict["atype"] >= 0).sum())
                         sum_natoms += natoms
                         for k, v in more_loss.items():
                             if "l2_" not in k:
@@ -2362,6 +2358,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
                 atomic=False,
                 must=not _model.has_default_fparam(),
                 default=_fparam_default,
+                source_policy=("default" if _model.has_default_fparam() else "tracked"),
             )
         ]
         additional_data_requirement += fparam_requirement_items
@@ -2390,6 +2387,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
                 atomic=True,
                 must=not allow_missing_spin,
                 default=0.0,
+                source_policy="default" if allow_missing_spin else "tracked",
             )
         ]
         additional_data_requirement += spin_requirement_items
@@ -2405,6 +2403,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
                 atomic=False,
                 must=not has_default_cs,
                 default=cs_default,
+                source_policy="default" if has_default_cs else "tracked",
             )
         )
     return additional_data_requirement
diff --git a/deepmd/pt/utils/lmdb_dataset.py b/deepmd/pt/utils/lmdb_dataset.py
index 3fed8282dd..1648adcd04 100644
--- a/deepmd/pt/utils/lmdb_dataset.py
+++ b/deepmd/pt/utils/lmdb_dataset.py
@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """PyTorch LMDB dataset — thin wrapper around framework-agnostic LmdbDataReader."""
 
+import functools
 import logging
 from collections.abc import (
     Iterator,
@@ -18,12 +19,17 @@
 
 from deepmd.dpmodel.utils.lmdb_data import (
     LmdbBatchIterator,
+    LmdbBatchSampler,
     LmdbDataReader,
+    LmdbDecodeConfig,
     LmdbTestData,
-    SameNlocBatchSampler,
     collate_lmdb_frames,
+    collect_lmdb_sampling_groups,
     compute_block_targets,
+    count_group_blocks,
     is_lmdb,
+    resolve_per_atom_keys,
+    system_block_lookup,
 )
 from deepmd.env import (
     get_lmdb_num_workers,
@@ -44,7 +50,10 @@
 ]
 
 
-def _collate_lmdb_batch(batch: list[dict[str, Any]]) -> dict[str, Any]:
+def _collate_lmdb_batch(
+    batch: list[dict[str, Any]],
+    config: LmdbDecodeConfig,
+) -> dict[str, Any]:
     """Collate a list of frame dicts into a torch batch dict.
 
     Pre-converts per-frame numpy arrays to CPU torch tensors (zero-copy when
@@ -53,19 +62,24 @@ def _collate_lmdb_batch(batch: list[dict[str, Any]]) -> dict[str, Any]:
     collate yields a torch dict (``sid`` becomes a torch tensor automatically
     via ``array_api_compat``).
 
-    All frames in the batch must have the same nloc (enforced by
-    SameNlocBatchSampler when mixed_batch=False). For mixed_batch=True,
-    raises NotImplementedError.
-    """
-    if len(batch) > 1:
-        atypes = [d.get("atype") for d in batch if d.get("atype") is not None]
-        if atypes and any(len(a) != len(atypes[0]) for a in atypes):
-            raise NotImplementedError(
-                "mixed_batch collation (frames with different atom counts "
-                "in the same batch) is not yet supported. "
-                "Padding + mask in collate_fn needed."
-            )
+    Frames of different atom counts are padded to the batch maximum; the
+    padded slots carry the phantom atom type. Frames need not agree on label
+    availability: a label only some of them carry is reported unavailable
+    for the whole batch.
 
+    Parameters
+    ----------
+    batch : list[dict[str, Any]]
+        Decoded frames to collate.
+    config : LmdbDecodeConfig
+        Decoder state whose data requirements identify the per-atom fields.
+
+    Returns
+    -------
+    dict[str, Any]
+        One collated batch of CPU tensors.
+    """
+    per_atom_keys = resolve_per_atom_keys(batch[0], config)
     with torch.device("cpu"):
         torch_frames: list[dict[str, Any]] = []
         for f in batch:
@@ -78,7 +92,7 @@ def _collate_lmdb_batch(batch: list[dict[str, Any]]) -> dict[str, Any]:
                 else:
                     tf[key] = torch.as_tensor(val)
             torch_frames.append(tf)
-        return collate_lmdb_frames(torch_frames)
+        return collate_lmdb_frames(torch_frames, per_atom_keys)
 
 
 def _lmdb_batch_to_torch(
@@ -100,15 +114,15 @@ def _lmdb_batch_to_torch(
     return converted
 
 
-class _SameNlocBatchSamplerTorch(Sampler):
-    """Torch Sampler adapter around the framework-agnostic SameNlocBatchSampler.
+class _LmdbBatchSamplerTorch(Sampler):
+    """Torch Sampler adapter around the framework-agnostic LmdbBatchSampler.
 
     PyTorch DataLoader with batch_sampler expects a Sampler that yields
-    lists of indices. This wraps SameNlocBatchSampler (or
-    DistributedSameNlocBatchSampler) to satisfy that.
+    lists of indices. This wraps LmdbBatchSampler (or
+    DistributedLmdbBatchSampler) to satisfy that.
     """
 
-    def __init__(self, inner: SameNlocBatchSampler) -> None:
+    def __init__(self, inner: LmdbBatchSampler) -> None:
         self._inner = inner
 
     def __iter__(self) -> Iterator[list[int]]:
@@ -140,7 +154,7 @@ def __init__(
         num_workers: int | None = None,
     ) -> None:
         self.dataset = dataset
-        self.batch_sampler = _SameNlocBatchSamplerTorch(sampler)
+        self.batch_sampler = _LmdbBatchSamplerTorch(sampler)
         self.sampler = sampler
         self._pin_memory = pin_memory
         self._batch_iterator = LmdbBatchIterator(
@@ -188,9 +202,11 @@ class LmdbDataset(Dataset):
         - ``"max:N"``: ``max(1, floor(N / nloc))`` per nloc group.
         - ``"filter:N"``: same per-nloc formula as ``"max:N"`` and drops
           every frame whose ``nloc > N`` from the dataset.
-    mixed_batch : bool
-        If True, allow different nloc in the same batch (future).
-        If False (default), use SameNlocBatchSampler.
+        - ``"mix:N"``: mixed-nloc batching to a padded-slot budget of ``N``;
+          frames of different atom counts share a batch and the shorter ones
+          are padded with phantom atoms.
+    auto_prob_style : str, optional
+        ``auto_prob`` string used to reweight the original systems.
     """
 
     def __init__(
@@ -198,21 +214,16 @@ def __init__(
         lmdb_path: str,
         type_map: list[str],
         batch_size: int | str = "auto",
-        mixed_batch: bool = False,
         auto_prob_style: str | None = None,
     ) -> None:
-        self._reader = LmdbDataReader(
-            lmdb_path, type_map, batch_size, mixed_batch=mixed_batch
+        self._reader = LmdbDataReader(lmdb_path, type_map, batch_size)
+        self._collate = functools.partial(
+            _collate_lmdb_batch, config=self._reader.decode_config
         )
 
-        if mixed_batch:
-            # Future: DataLoader with padding collate_fn
-            raise NotImplementedError(
-                "mixed_batch=True is not yet supported. "
-                "Requires padding + mask in collate_fn."
-            )
-
-        # Compute block_targets from auto_prob_style if provided
+        # Compute block_targets from auto_prob_style if provided. An empty
+        # result means the configured probabilities need no reweighting, which
+        # is the common case and worth no log line of its own.
         self._block_targets = None
         if auto_prob_style is not None and self._reader.frame_system_ids is not None:
             self._block_targets = compute_block_targets(
@@ -220,59 +231,52 @@ def __init__(
                 self._reader.nsystems,
                 self._reader.system_nframes,
             )
-            if self._block_targets is not None:
+            if self._block_targets:
                 log.info(
                     f"LMDB auto_prob: {len(self._block_targets)} blocks, "
                     f"nsystems={self._reader.nsystems}"
                 )
 
-        # Same-nloc batching: use SameNlocBatchSampler
-        sampler = SameNlocBatchSampler(
+        sampler = LmdbBatchSampler(
             self._reader,
             shuffle=True,
             block_targets=self._block_targets,
         )
-        self._batch_sampler = _SameNlocBatchSamplerTorch(sampler)
+        self._batch_sampler = _LmdbBatchSamplerTorch(sampler)
 
         with torch.device("cpu"):
             self._inner_dataloader = DataLoader(
                 self,
                 batch_sampler=self._batch_sampler,
                 num_workers=0,
-                collate_fn=_collate_lmdb_batch,
+                collate_fn=self._collate,
             )
 
-        # Per-nloc and label-availability dataloaders for make_stat_input.
-        # These are rebuilt after requirements are registered so statistics
-        # never collate real labels with default-filled placeholders.
+        # Homogeneous dataloaders for make_stat_input, built on first use and
+        # discarded whenever new requirements change how frames decode.
         self._nloc_dataloaders: list[DataLoader] | None = None
 
-    def _rebuild_nloc_dataloaders(self) -> None:
-        """Build homogeneous loaders used by model-stat collection."""
+    def _build_nloc_dataloaders(self) -> None:
+        """Build the homogeneous loaders used by model-stat collection."""
         dataloaders: list[DataLoader] = []
-        for nloc in sorted(self._reader.nloc_groups.keys()):
-            signature_groups = self._reader.group_indices_by_find_signature(
-                self._reader.nloc_groups[nloc]
-            )
-            for signature in sorted(signature_groups):
-                subset = torch.utils.data.Subset(self, signature_groups[signature])
-                bs = self._reader.get_batch_size_for_nloc(nloc)
-                with torch.device("cpu"):
-                    dl = DataLoader(
-                        subset,
-                        batch_size=bs,
-                        shuffle=False,
-                        num_workers=0,
-                        drop_last=False,
-                        collate_fn=_collate_lmdb_batch,
-                    )
-                dataloaders.append(dl)
+        for nloc, indices in collect_lmdb_sampling_groups(self._reader):
+            subset = torch.utils.data.Subset(self, indices)
+            with torch.device("cpu"):
+                dl = DataLoader(
+                    subset,
+                    batch_size=self._reader.get_batch_size_for_nloc(nloc),
+                    shuffle=False,
+                    num_workers=0,
+                    drop_last=False,
+                    collate_fn=self._collate,
+                )
+            dataloaders.append(dl)
         self._nloc_dataloaders = dataloaders
 
     def _get_nloc_dataloaders(self) -> list[DataLoader]:
         """Materialize statistics loaders lazily when none are registered."""
         if self._nloc_dataloaders is None:
-            self._rebuild_nloc_dataloaders()
+            self._build_nloc_dataloaders()
         dataloaders = self._nloc_dataloaders
         if dataloaders is None:
             raise RuntimeError("Failed to initialize LMDB statistics dataloaders")
@@ -295,8 +299,9 @@ def nframes(self) -> int:
         return self._reader.nframes
 
     @property
-    def mixed_batch(self) -> bool:
-        return self._reader.mixed_batch
+    def mixed_nloc(self) -> bool:
+        """Whether one batch may span several atom counts."""
+        return self._reader.mixed_nloc
 
     @property
     def mixed_type(self) -> bool:
@@ -317,7 +322,10 @@ def data_requirements(self) -> list[DataRequirementItem]:
 
     def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None:
         self._reader.add_data_requirement(data_requirement)
-        self._rebuild_nloc_dataloaders()
+        # Loaders decode through the registered requirements, so any already
+        # built are stale. They are rebuilt on demand, which spares a run
+        # whose statistics come from a stat file the work entirely.
+        self._nloc_dataloaders = None
 
     def close(self) -> None:
         """Release parent-process LMDB resources."""
@@ -358,11 +366,12 @@ def print_summary(self, name: str, prob: Any) -> None:
                     f"{actual}->{target} (x{ratio:.2f})"
                 )
 
-            # Build sys_id -> block_idx mapping
-            sys_to_block: dict[int, int] = {}
-            for blk_idx, (sys_ids, _) in enumerate(self._block_targets):
-                for sid in sys_ids:
-                    sys_to_block[sid] = blk_idx
+            # A whole nloc group's block membership resolves as one indexing
+            # operation, which a dataset of 10^8 frames needs it to be. The
+            # lookup is the one the sampler allocates its targets with, so
+            # both agree on which systems a block claims.
+            n_blocks = len(self._block_targets)
+            lookup = system_block_lookup(self._block_targets)
 
             # Compute expanded nloc counts analytically (no actual expansion)
             expanded_nloc_info = {}
@@ -370,19 +379,12 @@ def print_summary(self, name: str, prob: Any) -> None:
                 if reader.frame_system_ids is None:
                     expanded_nloc_info[nloc] = len(indices)
                     continue
-                # Count indices per block in this nloc group
-                blk_counts: dict[int, int] = {}
-                unassigned = 0
-                for idx in indices:
-                    sid = reader.frame_system_ids[idx]
-                    blk = sys_to_block.get(sid)
-                    if blk is not None:
-                        blk_counts[blk] = blk_counts.get(blk, 0) + 1
-                    else:
-                        unassigned += 1
-                expanded = unassigned
+                counts = count_group_blocks(
+                    indices, reader.frame_system_ids, lookup, n_blocks
+                )
+                expanded = len(indices) - int(counts.sum())
                 for blk_idx, (_, blk_target) in enumerate(self._block_targets):
-                    n_actual = blk_counts.get(blk_idx, 0)
+                    n_actual = int(counts[blk_idx])
                     if n_actual == 0:
                         continue
                     bta = block_total_actual[blk_idx]
@@ -411,8 +413,6 @@ def set_noise(self, noise_settings: dict[str, Any]) -> None:
     @property
     def index(self) -> list[int]:
         """Number of batches per logical LMDB dataset."""
-        if not self._block_targets:
-            return self._reader.index
         return [self.total_batch]
 
     @property
@@ -434,8 +434,8 @@ def systems(self) -> list:
     def dataloaders(self) -> list:
         """Homogeneous dataloaders for make_stat_input.
 
-        Each loader has one nloc and one availability signature, so stat
-        collection sees consistent shapes and scalar ``find_*`` flags.
+        Each loader draws from one atom count and one label availability, so
+        stat collection sees consistent shapes and scalar ``find_*`` flags.
         """
         return self._get_nloc_dataloaders()
 
diff --git a/deepmd/pt/utils/nv_nlist.py b/deepmd/pt/utils/nv_nlist.py
index df5f3dc119..3c68251a7d 100644
--- a/deepmd/pt/utils/nv_nlist.py
+++ b/deepmd/pt/utils/nv_nlist.py
@@ -353,6 +353,21 @@ def _matrix_to_extended_inputs(
     slot_idx = edge_idx % max_neighbors
     zero_shift = torch.all(shift == 0, dim=1)
 
+    # Phantom atoms (atype < 0) pad a mixed-nloc batch and have no physical
+    # site, but the geometric search still pairs them with real atoms near
+    # their placeholder coordinates. Edges touching one are neutralized here,
+    # ahead of the ``sum(sel)`` truncation applied by the caller, so that a
+    # phantom can never displace a genuine neighbor from a real atom's list.
+    # Rather than compacting the edge arrays, which would cost a device sync,
+    # such an edge is rewritten to carry the empty marker and routed through
+    # the direct branch below: it then writes -1 into its own matrix slot,
+    # which is what that slot already holds, and is excluded from the ghost
+    # materialization of Step 3.
+    atype_flat = atype.reshape(-1)
+    real_pair = (atype_flat[dst] >= 0) & (atype_flat[src] >= 0)
+    src_local = torch.where(real_pair, src_local, -1)
+    zero_shift = zero_shift | ~real_pair
+
     # === Step 2. Direct neighbors keep their local extended indices ===
     # Zero-shift neighbors already live in the leading local block of
     # `extended_coord`, so their DeePMD nlist value is simply `src_local`.
diff --git a/deepmd/pt_expt/model/ener_model.py b/deepmd/pt_expt/model/ener_model.py
index 06f57176cd..6aa5f4add4 100644
--- a/deepmd/pt_expt/model/ener_model.py
+++ b/deepmd/pt_expt/model/ener_model.py
@@ -254,6 +254,30 @@ def forward(
             do_atomic_virial=do_atomic_virial,
             neighbor_list=neighbor_list,
         )
+        return self._to_public_keys(model_ret, do_atomic_virial)
+
+    def _to_public_keys(
+        self,
+        model_ret: dict[str, torch.Tensor],
+        do_atomic_virial: bool,
+    ) -> dict[str, torch.Tensor]:
+        """Rename one ``call_common`` result to the public energy-model keys.
+
+        The renaming is a property of the output definition, not of the node
+        axis, so it serves the rectangular and the ragged entry alike.
+
+        Parameters
+        ----------
+        model_ret : dict[str, torch.Tensor]
+            A ``call_common`` result, in internal ``<var>_<derivative>`` keys.
+        do_atomic_virial : bool
+            Whether the per-atom virial was requested and should be carried.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            The same tensors under the public names.
+        """
         model_predict = {}
         model_predict["atom_energy"] = model_ret["energy"]
         model_predict["energy"] = model_ret["energy_redu"]
@@ -263,12 +287,68 @@ def forward(
             model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
             if do_atomic_virial:
                 model_predict["atom_virial"] = model_ret["energy_derv_c"].squeeze(-2)
-        if "mask" in model_ret:
-            model_predict["mask"] = model_ret["mask"]
+        for key in ("mask", "n_node"):
+            if key in model_ret:
+                model_predict[key] = model_ret[key]
         if self.atomic_output_def()["energy"].r_hessian:
             model_predict["hessian"] = model_ret["energy_derv_r_derv_r"].squeeze(-3)
         return model_predict
 
+    def forward_ragged(
+        self,
+        coord: torch.Tensor,
+        atype: torch.Tensor,
+        n_node: torch.Tensor,
+        box: torch.Tensor | None = None,
+        fparam: torch.Tensor | None = None,
+        aparam: torch.Tensor | None = None,
+        do_atomic_virial: bool = False,
+        charge_spin: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Evaluate the energy model over a batch whose node axis is flat.
+
+        The counterpart of :meth:`forward` for frames held concatenated rather
+        than padded to a common atom count. Arguments and results share their
+        meaning with :meth:`call_common_ragged`, whose per-atom outputs keep
+        the flat axis; only the keys are the public ones.
+
+        Parameters
+        ----------
+        coord : torch.Tensor
+            Local coordinates with shape ``(N, 3)``, frame-major over ``n_node``.
+        atype : torch.Tensor
+            Local atom types with shape ``(N,)``.
+        n_node : torch.Tensor
+            Atoms per frame with shape ``(nf,)``.
+        box : torch.Tensor or None, optional
+            Simulation cell with shape ``(nf, 3, 3)``.
+        fparam : torch.Tensor or None, optional
+            Frame parameters with shape ``(nf, ndf)``.
+        aparam : torch.Tensor or None, optional
+            Atomic parameters with shape ``(N, nda)``.
+        do_atomic_virial : bool, default: False
+            Whether to return per-atom virials.
+        charge_spin : torch.Tensor or None, optional
+            Frame-level charge and spin conditioning with shape ``(nf, 2)``.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            Public energy-model keys; per-atom entries have leading dimension
+            ``N`` and per-frame entries ``nf``.
+        """
+        model_ret = self.call_common_ragged(
+            coord,
+            atype,
+            n_node,
+            box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+            charge_spin=charge_spin,
+        )
+        return self._to_public_keys(model_ret, do_atomic_virial)
+
     def forward_lower(
         self,
         extended_coord: torch.Tensor,
@@ -605,6 +685,7 @@ def forward_lower_graph_exportable_with_comm(
         ----------
         atype, n_node, edge_index, edge_vec, edge_mask, fparam, aparam, charge_spin, do_atomic_virial
             As in :meth:`forward_lower_graph_exportable`.
+
         send_list, send_proc, recv_proc, send_num, recv_num, communicator, nlocal, nghost
             The 8 comm tensors (see ``_make_comm_sample_inputs`` in
             ``serialization.py``), packed into ``comm_dict`` inside the
@@ -621,6 +702,7 @@ def forward_lower_graph_exportable_with_comm(
             back with synchronizing D2H reads (``4 * nlayers`` per MD
             step).  The C++ ``run_model_graph_with_comm`` implements this
             placement.
+
         n_local
             (1,) int64 ON THE MODEL DEVICE: the per-frame OWNED node
             count consumed IN-GRAPH by the owned-node energy mask (it
@@ -630,6 +712,7 @@ def forward_lower_graph_exportable_with_comm(
             access).  Carries the same value as the ``nlocal`` comm
             tensor; the two inputs exist precisely to separate the
             device-compute role from the host-MPI-control role.
+
         **make_fx_kwargs
             Extra keyword arguments forwarded to ``make_fx``
             (e.g. ``tracing_mode="symbolic"``).
diff --git a/deepmd/pt_expt/model/make_model.py b/deepmd/pt_expt/model/make_model.py
index f59d507727..7aa82215b2 100644
--- a/deepmd/pt_expt/model/make_model.py
+++ b/deepmd/pt_expt/model/make_model.py
@@ -20,6 +20,10 @@
 from deepmd.dpmodel.output_def import (
     OutputVariableDef,
 )
+from deepmd.dpmodel.utils.neighbor_graph import (
+    compact_nodes,
+    expand_node_values,
+)
 from deepmd.kernels.utils import (
     cuda_infer_level,
 )
@@ -28,6 +32,8 @@
 )
 from deepmd.pt_expt.utils.graph_builder import (
     build_neighbor_graph_for_method,
+    build_ragged_neighbor_graph,
+    resolve_neighbor_graph_method,
 )
 from deepmd.pt_expt.utils.graph_csr import (
     validate_graph_csr_for_export,
@@ -190,6 +196,8 @@ def _cal_hessian_ext(
         Atomic parameters. Shape: [nf, nloc, nap] or None.
     create_graph
         Whether to create graph for higher-order derivatives.
+    charge_spin
+        Frame-level charge and spin conditioning. Shape: [nf, ncs] or None.
 
     Returns
     -------
@@ -311,7 +319,9 @@ def __init__(
         pair_excl: Any,
         rcut: float,
         fparam: torch.Tensor | None,  # (1, ndf) or None
-        aparam: torch.Tensor | None,  # (1, nloc, nda) or None
+        aparam: torch.Tensor | None,  # (nloc, nda) or None
+        spin: torch.Tensor | None,  # (nloc, 3) or None
+        charge_spin: torch.Tensor | None,  # (1, 2) or None
     ) -> None:
         self.model = model
         self.kk = kk
@@ -324,6 +334,8 @@ def __init__(
         self.rcut = rcut
         self.fparam = fparam
         self.aparam = aparam
+        self.spin = spin
+        self.charge_spin = charge_spin
 
     def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor:
         cc = coord_flat.reshape(1, self.nloc, 3)
@@ -334,11 +346,9 @@ def __call__(self, coord_flat: torch.Tensor) -> torch.Tensor:
             ng,
             self.atype.reshape(-1),
             fparam=self.fparam,
-            # graph-lower ABI: aparam is FLAT on the node axis, (N, nda)
-            # (N == nloc for the single-frame carry-all graph).
-            aparam=(
-                self.aparam.reshape(self.nloc, -1) if self.aparam is not None else None
-            ),
+            aparam=self.aparam,
+            spin=self.spin,
+            charge_spin=self.charge_spin,
         )
         # atomic_ret[kk]: flat (N, *def), N == nloc for a single-frame carry-all
         # graph (all nodes owned); reduced output = sum over the node axis.
@@ -355,6 +365,8 @@ def _cal_hessian_ext_graph(
     box: torch.Tensor | None,
     fparam: torch.Tensor | None,
     aparam: torch.Tensor | None,
+    spin: torch.Tensor | None,
+    charge_spin: torch.Tensor | None,
     method: str,
     pair_excl: Any,
     rcut: float,
@@ -363,37 +375,68 @@ def _cal_hessian_ext_graph(
     """Graph twin of :func:`_cal_hessian_ext`.
 
     Computes the Hessian of the reduced output w.r.t. the LOCAL coordinates on
-    the carry-all graph route. Returns shape ``[nf, *vdef.shape, nloc*3,
-    nloc*3]`` -- the local-only counterpart of the dense extended Hessian,
-    already in the same final layout the dense route reaches after
+    the carry-all graph route. Each frame is differentiated only over its real
+    nodes; phantom rows are restored as zero rows and columns at the public
+    rectangular boundary. Returns shape ``[nf, *vdef.shape, nloc*3, nloc*3]``
+    -- the local-only counterpart of the dense extended Hessian, already in
+    the same final layout the dense route reaches after
     ``communicate_extended_output`` folds ``nall -> nloc`` (the graph route
     reduces over owned nodes, so no fold is needed). Node axis is
     atom-major/xyz-minor, matching the dense final reshape.
     """
     nf, nloc, _ = coord.shape
     vsize = math.prod(vdef.shape)
-    coord_flat = coord.reshape(nf, nloc * 3)
+    aparam_by_node = aparam.reshape(nf, nloc, -1) if aparam is not None else None
+    spin_by_node = spin.reshape(nf, nloc, 3) if spin is not None else None
+    charge_spin_by_frame = (
+        charge_spin.reshape(1, -1)
+        if charge_spin is not None and charge_spin.ndim == 1
+        else charge_spin
+    )
     hessians = []
     for ii in range(nf):
+        node_index = torch.nonzero(atype[ii] >= 0, as_tuple=False).reshape(-1)
+        n_real = node_index.shape[0]
+        coord_flat = coord[ii, node_index].reshape(n_real * 3)
+        atype_frame = atype[ii : ii + 1, node_index]
+        aparam_frame = (
+            aparam_by_node[ii, node_index] if aparam_by_node is not None else None
+        )
+        spin_frame = spin_by_node[ii, node_index] if spin_by_node is not None else None
+        charge_spin_frame = None
+        if charge_spin_by_frame is not None:
+            frame_index = 0 if charge_spin_by_frame.shape[0] == 1 else ii
+            charge_spin_frame = charge_spin_by_frame[frame_index : frame_index + 1]
         for ci in range(vsize):
             wrapper = _WrapperForwardEnergyGraph(
-                model,
-                kk,
-                ci,
-                nloc,
-                atype[ii : ii + 1],
-                box[ii : ii + 1] if box is not None else None,
-                method,
-                pair_excl,
-                rcut,
-                fparam[ii : ii + 1] if fparam is not None else None,
-                aparam[ii : ii + 1] if aparam is not None else None,
+                model=model,
+                kk=kk,
+                ci=ci,
+                nloc=n_real,
+                atype=atype_frame,
+                box=box[ii : ii + 1] if box is not None else None,
+                method=method,
+                pair_excl=pair_excl,
+                rcut=rcut,
+                fparam=fparam[ii : ii + 1] if fparam is not None else None,
+                aparam=aparam_frame,
+                spin=spin_frame,
+                charge_spin=charge_spin_frame,
             )
             hess = torch.autograd.functional.hessian(
                 wrapper,
-                coord_flat[ii],
+                coord_flat,
                 create_graph=create_graph,
-            )  # (nloc*3, nloc*3)
+            )  # (n_real*3, n_real*3)
+            if n_real != nloc:
+                component_index = (
+                    node_index[:, None] * 3
+                    + torch.arange(3, dtype=node_index.dtype, device=node_index.device)
+                ).reshape(-1)
+                hess = expand_node_values(hess, component_index, nloc * 3)
+                hess = expand_node_values(
+                    hess.transpose(0, 1), component_index, nloc * 3
+                ).transpose(0, 1)
             hessians.append(hess)
     return torch.stack(hessians).reshape(nf, *vdef.shape, nloc * 3, nloc * 3)
 
@@ -710,6 +753,116 @@ def _resolve_graph_method(
                 return getattr(self, "neighbor_graph_method", "dense")
             return None
 
+        def call_common_ragged(
+            self,
+            coord: torch.Tensor,
+            atype: torch.Tensor,
+            n_node: torch.Tensor,
+            box: torch.Tensor | None = None,
+            fparam: torch.Tensor | None = None,
+            aparam: torch.Tensor | None = None,
+            do_atomic_virial: bool = False,
+            charge_spin: torch.Tensor | None = None,
+            spin: torch.Tensor | None = None,
+        ) -> dict[str, torch.Tensor]:
+            """Model forward over a batch whose node axis is already flat.
+
+            The rectangular :meth:`call_common` pads frames of unequal atom
+            count to a common width and unpads its output again. A caller that
+            holds the frames concatenated skips both: the node axis it passes
+            in is the one the graph lower works on, and the per-atom outputs
+            come back on it.
+
+            Parameters
+            ----------
+            coord : torch.Tensor
+                Local coordinates with shape ``(N, 3)``, frame-major over
+                ``n_node``.
+            atype : torch.Tensor
+                Local atom types with shape ``(N,)``.
+            n_node : torch.Tensor
+                Atoms per frame with shape ``(nf,)``.
+            box : torch.Tensor or None, optional
+                Simulation cell with shape ``(nf, 3, 3)``, or ``None`` for
+                non-periodic.
+            fparam : torch.Tensor or None, optional
+                Frame parameter with shape ``(nf, ndf)``.
+            aparam : torch.Tensor or None, optional
+                Atomic parameter with shape ``(N, nda)``.
+            do_atomic_virial : bool, default: False
+                Whether to compute the atomic virial.
+            charge_spin : torch.Tensor or None, optional
+                Frame-level charge/spin conditioning with shape ``(nf, 2)``.
+            spin : torch.Tensor or None, optional
+                Native spin with shape ``(N, 3)``. Virtual-atom spin models do
+                not expose this ragged entry.
+
+            Returns
+            -------
+            dict[str, torch.Tensor]
+                The standard model dict. Per-atom keys keep the flat ``(N, *)``
+                axis; per-frame keys have leading dimension ``nf``.
+
+            Raises
+            ------
+            NotImplementedError
+                If the model has no graph lower to read a flat node axis with,
+                or if its outputs require a rectangular pair axis.
+            """
+            if not (self.mixed_types() and self.atomic_model.uses_graph_lower()):
+                raise NotImplementedError(
+                    "a flat node axis requires a mixed_types descriptor with a "
+                    "graph lower; this model reads a rectangular one, so its "
+                    "batches must be padded to a common atom count"
+                )
+            if any(
+                vdef.reducible and vdef.r_hessian
+                for vdef in self.atomic_output_def().get_data().values()
+            ):
+                raise NotImplementedError(
+                    "Hessian outputs require a rectangular atom axis; a flat "
+                    "node axis cannot represent their per-frame pair dimensions"
+                )
+            # The trainer resolves ``auto`` once and installs the concrete
+            # builder on the model. A model reached outside it has none, and
+            # resolving against its own device is what keeps that case from
+            # silently taking the CPU builder on a GPU.
+            method = getattr(self, "neighbor_graph_method", None)
+            if method is None:
+                method = resolve_neighbor_graph_method("auto", coord.device)
+            graph = build_ragged_neighbor_graph(
+                method,
+                coord,
+                atype,
+                n_node,
+                box,
+                self.get_rcut(),
+                getattr(self.atomic_model, "pair_excl", None),
+            )
+            predict = self.forward_common_lower_graph(
+                atype,
+                graph.n_node,
+                graph.n_node,
+                graph.edge_index,
+                graph.edge_vec,
+                graph.edge_mask,
+                graph.destination_order,
+                graph.destination_row_ptr,
+                graph.source_order,
+                graph.source_row_ptr,
+                destination_sorted=graph.destination_sorted,
+                do_atomic_virial=do_atomic_virial,
+                fparam=fparam,
+                aparam=aparam,
+                charge_spin=charge_spin,
+                spin=spin,
+            )
+            # The per-atom mask a rectangular batch carries is what tells the
+            # loss each frame's real atom count. Nothing here is padded, so the
+            # counts are stated outright instead.
+            predict["n_node"] = graph.n_node
+            return predict
+
         def _call_common_graph(
             self,
             cc: torch.Tensor,
@@ -788,10 +941,25 @@ def _call_common_graph(
                 method, cc, atype, bb, rcut, pair_excl, with_csr=with_csr
             )
             nf, nloc = atype.shape[:2]
-            atype_flat = atype.reshape(nf * nloc)
+            n_padded = nf * nloc
+            atype_flat = atype.reshape(n_padded)
+            # A batch of unequal atom counts arrives padded to a common width
+            # with phantom atoms (atype < 0). The builders leave them out of
+            # every edge, so dropping them from the node axis costs nothing and
+            # spares the network from evaluating them. On a batch of uniform
+            # atom count the mask is all true and this is a renumbering by the
+            # identity.
+            ng, node_index = compact_nodes(ng, atype_flat >= 0)
+            atype_flat = atype_flat[node_index]
             # graph-lower ABI: aparam/spin are FLAT on the node axis, (N, nda)/(N, 3).
-            ap_flat = ap.reshape(nf * nloc, ap.shape[-1]) if ap is not None else None
-            spin_flat = spin.reshape(nf * nloc, 3) if spin is not None else None
+            ap_flat = (
+                ap.reshape(n_padded, ap.shape[-1])[node_index]
+                if ap is not None
+                else None
+            )
+            spin_flat = (
+                spin.reshape(n_padded, 3)[node_index] if spin is not None else None
+            )
             model_predict = self.forward_common_lower_graph(
                 atype_flat,
                 ng.n_node,
@@ -811,11 +979,14 @@ def _call_common_graph(
                 charge_spin=charge_spin,
             )
             # ``forward_common_lower_graph`` returns flat ``(N, *)`` per-atom
-            # outputs (N = nf * nloc for a carry-all rectangular graph).
-            # Unravel to rectangular ``(nf, nloc, *)`` at the public I/O boundary
-            # so that callers receive the same shape as the dense ``call_common``.
-            N = nf * nloc
-            # public call_common always passes rectangular (nf,nloc) coord/atype (N == nf*nloc), so this unravel always applies; ragged graphs reach call_lower_graph/forward_common_lower_graph directly (no unravel) and stay flat (N,*).
+            # outputs over the real atoms. Scatter them back onto the padded
+            # width and unravel to rectangular ``(nf, nloc, *)`` at the public
+            # I/O boundary, so that callers receive the same shape as the dense
+            # ``call_common``. A phantom slot reads zero, which is what a
+            # masked-out atom contributed there before.
+            N = node_index.shape[0]
+            # Only the rectangular entry reaches this scatter; the ragged
+            # one keeps the flat axis its caller handed over.
             for k in list(model_predict.keys()):
                 v = model_predict[k]
                 # per-frame reduced keys (..._redu) keep their (nf, *) shape; only node-level (N,*) keys unravel — guards the nloc==1 case where N == nf.
@@ -824,11 +995,15 @@ def _call_common_graph(
                     and not k.endswith("_redu")
                     and v.shape[:1] == torch.Size([N])
                 ):
-                    model_predict[k] = v.reshape(nf, nloc, *v.shape[1:])
+                    model_predict[k] = expand_node_values(
+                        v, node_index, n_padded
+                    ).reshape(nf, nloc, *v.shape[1:])
             # Graph-native Hessian (parallel to the dense ``forward_common_atomic``
-            # loop): differentiate the reduced output w.r.t. the LOCAL coords by
-            # rebuilding the graph inside the wrapper. Added AFTER the unravel so
-            # its ``(nf, *def, nloc, 3, nloc, 3)`` shape is returned as-is.
+            # loop): differentiate the reduced output w.r.t. the compact LOCAL
+            # coordinates by rebuilding the graph inside the wrapper, then restore
+            # phantom rows and columns at the public rectangular boundary. Added
+            # AFTER the unravel so its ``(nf, *def, nloc, 3, nloc, 3)`` shape is
+            # returned as-is.
             # Eager-only, like the dense Hessian (autograd.functional.hessian
             # does not export/compile).
             aod = self.atomic_output_def()
@@ -836,17 +1011,19 @@ def _call_common_graph(
                 vdef = aod[kk]
                 if vdef.reducible and vdef.r_hessian:
                     model_predict[get_hessian_name(kk)] = _cal_hessian_ext_graph(
-                        self,
-                        kk,
-                        vdef,
-                        cc,
-                        atype,
-                        bb,
-                        fp,
-                        ap,
-                        method,
-                        pair_excl,
-                        rcut,
+                        model=self,
+                        kk=kk,
+                        vdef=vdef,
+                        coord=cc,
+                        atype=atype,
+                        box=bb,
+                        fparam=fp,
+                        aparam=ap,
+                        spin=spin,
+                        charge_spin=charge_spin,
+                        method=method,
+                        pair_excl=pair_excl,
+                        rcut=rcut,
                         create_graph=self.training,
                     )
             return model_predict
@@ -931,7 +1108,7 @@ def forward_common_lower_exportable(
 
             Parameters
             ----------
-            extended_coord, extended_atype, nlist, mapping, fparam, aparam, do_atomic_virial
+            extended_coord, extended_atype, nlist, mapping, fparam, aparam, do_atomic_virial, charge_spin
                 Sample inputs with representative shapes (used for tracing).
             **make_fx_kwargs
                 Extra keyword arguments forwarded to ``make_fx``
@@ -942,7 +1119,7 @@ def forward_common_lower_exportable(
             torch.nn.Module
                 A traced module whose ``forward`` accepts
                 ``(extended_coord, extended_atype, nlist, mapping,
-                fparam, aparam)`` and returns a dict with the same keys
+                fparam, aparam, charge_spin)`` and returns a dict with the same keys
                 as ``call_common_lower``.
             """
             model = self
diff --git a/deepmd/pt_expt/model/native_spin_model.py b/deepmd/pt_expt/model/native_spin_model.py
index 6b3a87b49e..bba9d104b3 100644
--- a/deepmd/pt_expt/model/native_spin_model.py
+++ b/deepmd/pt_expt/model/native_spin_model.py
@@ -103,6 +103,34 @@ def forward(
             model_ret, atype, do_atomic_virial=do_atomic_virial
         )
 
+    def forward_ragged(
+        self,
+        coord: torch.Tensor,
+        atype: torch.Tensor,
+        n_node: torch.Tensor,
+        spin: torch.Tensor,
+        box: torch.Tensor | None = None,
+        fparam: torch.Tensor | None = None,
+        aparam: torch.Tensor | None = None,
+        do_atomic_virial: bool = False,
+        charge_spin: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Return native-spin predictions on a flat, frame-segmented node axis."""
+        model_ret = self.call_common_ragged(
+            coord,
+            atype,
+            n_node,
+            box=box,
+            fparam=fparam,
+            aparam=aparam,
+            do_atomic_virial=do_atomic_virial,
+            charge_spin=charge_spin,
+            spin=spin,
+        )
+        return self._translate_eager_call(
+            model_ret, atype, do_atomic_virial=do_atomic_virial
+        )
+
     def forward_lower_graph_exportable(
         self,
         atype: torch.Tensor,
diff --git a/deepmd/pt_expt/model/spin_ener_model.py b/deepmd/pt_expt/model/spin_ener_model.py
index 54d0cbb411..178653e821 100644
--- a/deepmd/pt_expt/model/spin_ener_model.py
+++ b/deepmd/pt_expt/model/spin_ener_model.py
@@ -38,6 +38,8 @@ def translated_output_def(self) -> dict[str, Any]:
             output_def["virial"].squeeze(-2)
             output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"])
             output_def["atom_virial"].squeeze(-2)
+        if "mask" in out_def_data:
+            output_def["mask"] = out_def_data["mask"]
         return output_def
 
     def forward(
@@ -65,6 +67,8 @@ def forward(
         model_predict["atom_energy"] = model_ret["energy"]
         model_predict["energy"] = model_ret["energy_redu"]
         model_predict["mask_mag"] = model_ret["mask_mag"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
         if self.backbone_model.do_grad_r("energy"):
             model_predict["force"] = model_ret["energy_derv_r"].squeeze(-2)
             model_predict["force_mag"] = model_ret["energy_derv_r_mag"].squeeze(-2)
@@ -101,6 +105,8 @@ def forward_lower(
         model_predict["atom_energy"] = model_ret["energy"]
         model_predict["energy"] = model_ret["energy_redu"]
         model_predict["extended_mask_mag"] = model_ret["mask_mag"]
+        if "mask" in model_ret:
+            model_predict["mask"] = model_ret["mask"]
         if self.backbone_model.do_grad_r("energy"):
             model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
             model_predict["extended_force_mag"] = model_ret[
@@ -135,7 +141,7 @@ def forward_lower_exportable(
 
         Parameters
         ----------
-        extended_coord, extended_atype, extended_spin, nlist, mapping, fparam, aparam, do_atomic_virial
+        extended_coord, extended_atype, extended_spin, nlist, mapping, fparam, aparam, do_atomic_virial, charge_spin
             Sample inputs with representative shapes (used for tracing).
         **make_fx_kwargs
             Extra keyword arguments forwarded to ``make_fx``
@@ -145,7 +151,7 @@ def forward_lower_exportable(
         -------
         torch.nn.Module
             A traced module whose ``forward`` accepts
-            ``(extended_coord, extended_atype, extended_spin, nlist, mapping, fparam, aparam)``
+            ``(extended_coord, extended_atype, extended_spin, nlist, mapping, fparam, aparam, charge_spin)``
             and returns a dict with the same keys as ``forward_lower``.
         """
         traced = self.forward_common_lower_exportable(
@@ -190,6 +196,8 @@ def fn(
             model_predict["atom_energy"] = model_ret["energy"]
             model_predict["energy"] = model_ret["energy_redu"]
             model_predict["extended_mask_mag"] = model_ret["mask_mag"]
+            if "mask" in model_ret:
+                model_predict["mask"] = model_ret["mask"]
             if do_grad_r:
                 model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
                 model_predict["extended_force_mag"] = model_ret[
diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py
index 805e55cbe9..18cafbff4e 100644
--- a/deepmd/pt_expt/train/training.py
+++ b/deepmd/pt_expt/train/training.py
@@ -358,6 +358,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
                 atomic=False,
                 must=not has_default_fparam,
                 default=fparam_default,
+                source_policy="default" if has_default_fparam else "tracked",
             )
         )
     if _model.get_dim_aparam() > 0:
@@ -380,6 +381,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
                 atomic=True,
                 must=not allow_missing_spin,
                 default=0.0,
+                source_policy="default" if allow_missing_spin else "tracked",
             )
         )
     if _model.has_chg_spin_ebd():
@@ -399,6 +401,7 @@ def get_additional_data_requirement(_model: Any) -> list[DataRequirementItem]:
                 atomic=False,
                 must=not has_default_cs,
                 default=default_cs,
+                source_policy="default" if has_default_cs else "tracked",
             )
         )
     return additional_data_requirement
@@ -460,6 +463,7 @@ def _trace_and_compile(
     model: torch.nn.Module,
     ext_coord: torch.Tensor,
     ext_atype: torch.Tensor,
+    ext_spin: torch.Tensor | None,
     nlist: torch.Tensor,
     mapping: torch.Tensor,
     fparam: torch.Tensor | None,
@@ -474,11 +478,15 @@ def _trace_and_compile(
     ----------
     model : torch.nn.Module
         The (uncompiled) model.
-    ext_coord, ext_atype, nlist, mapping, fparam, aparam
+    ext_coord, ext_atype, ext_spin, nlist, mapping, fparam, aparam
         Sample tensors used to seed the symbolic tracer.
+        ``ext_spin`` is present for virtual-atom spin models and follows the
+        same extended-node axis as coordinates and atom types.
     compile_opts : dict or None
         User-supplied inductor options.  These are merged on top of the
         built-in defaults (user values take precedence).
+    charge_spin : torch.Tensor or None
+        Frame-level charge and spin conditioning used by the traced model.
     task_buffers : dict or None
         Per-task buffers (e.g. ``bias_atom_e``, ``case_embd``, ``out_bias``,
         ``out_std``) detected by ``_detect_task_buffers``.  These are promoted
@@ -523,6 +531,7 @@ def _trace_and_compile(
     def fn(
         extended_coord: torch.Tensor,
         extended_atype: torch.Tensor,
+        extended_spin: torch.Tensor | None,
         nlist: torch.Tensor,
         mapping: torch.Tensor | None,
         fparam: torch.Tensor | None,
@@ -548,6 +557,17 @@ def fn(
                         originals[name] = _fitting._buffers.get(name)
                         _fitting._buffers[name] = val
         try:
+            if extended_spin is not None:
+                return model.forward_lower(
+                    extended_coord,
+                    extended_atype,
+                    extended_spin,
+                    nlist,
+                    mapping,
+                    fparam=fparam,
+                    aparam=aparam,
+                    charge_spin=charge_spin,
+                )
             return model.forward_lower(
                 extended_coord,
                 extended_atype,
@@ -600,6 +620,8 @@ def fn(
     # Pad nf only; nloc and nall retain their real values (no clamping needed).
     ext_coord = _trace_pad_dim(ext_coord[:1], 0, trace_nf)
     ext_atype = _trace_pad_dim(ext_atype[:1], 0, trace_nf)
+    if ext_spin is not None:
+        ext_spin = _trace_pad_dim(ext_spin[:1], 0, trace_nf)
     nlist = _trace_pad_dim(nlist[:1], 0, trace_nf)
     mapping = _trace_pad_dim(mapping[:1], 0, trace_nf)
     if fparam is not None:
@@ -627,6 +649,7 @@ def fn(
     )(
         ext_coord,
         ext_atype,
+        ext_spin,
         nlist,
         mapping,
         fparam,
@@ -1088,6 +1111,123 @@ def __getattr__(self, name: str) -> Any:
         except AttributeError:
             return getattr(self.original_model, name)
 
+    def _forward_eager(
+        self,
+        coord: torch.Tensor,
+        atype: torch.Tensor,
+        box: torch.Tensor | None,
+        fparam: torch.Tensor | None,
+        aparam: torch.Tensor | None,
+        do_atomic_virial: bool,
+        charge_spin: torch.Tensor | None,
+        spin: torch.Tensor | None,
+        n_node: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Delegate one rectangular or ragged evaluation to the original model."""
+        kwargs = {
+            "box": box,
+            "fparam": fparam,
+            "aparam": aparam,
+            "do_atomic_virial": do_atomic_virial,
+            "charge_spin": charge_spin,
+        }
+        if spin is not None:
+            kwargs["spin"] = spin
+        if n_node is not None:
+            return self.original_model.forward_ragged(
+                coord,
+                atype,
+                n_node,
+                **kwargs,
+            )
+        return self.original_model(coord, atype, **kwargs)
+
+    def forward_ragged(
+        self,
+        coord: torch.Tensor,
+        atype: torch.Tensor,
+        n_node: torch.Tensor,
+        box: torch.Tensor | None = None,
+        fparam: torch.Tensor | None = None,
+        aparam: torch.Tensor | None = None,
+        do_atomic_virial: bool = False,
+        charge_spin: torch.Tensor | None = None,
+        spin: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Compiled forward over a batch whose node axis is already flat.
+
+        The compiled lower works on that axis in either case -- its trace keeps
+        the frame count, the node count and the edge count as independent
+        symbols -- so a ragged batch simply skips the padding round trip the
+        rectangular :meth:`forward` performs around it.
+
+        Parameters
+        ----------
+        coord : torch.Tensor
+            Local coordinates with shape ``(N, 3)``, frame-major over ``n_node``.
+        atype : torch.Tensor
+            Local atom types with shape ``(N,)``.
+        n_node : torch.Tensor
+            Atoms per frame with shape ``(nf,)``.
+        box : torch.Tensor or None, optional
+            Simulation cell, ``(nf, 3, 3)`` or ``(nf, 9)``.
+        fparam : torch.Tensor or None, optional
+            Frame parameters with shape ``(nf, ndf)``.
+        aparam : torch.Tensor or None, optional
+            Atomic parameters with shape ``(N, nda)``.
+        do_atomic_virial : bool, default: False
+            Whether to return per-atom virials.
+        charge_spin : torch.Tensor or None, optional
+            Frame-level charge and spin conditioning with shape ``(nf, 2)``.
+        spin : torch.Tensor or None, optional
+            Native spin with shape ``(N, 3)``.
+
+        Returns
+        -------
+        dict[str, torch.Tensor]
+            Public model keys; per-atom entries keep the flat axis.
+
+        Raises
+        ------
+        NotImplementedError
+            If the model reads a rectangular node axis, which cannot represent
+            frames of unequal atom count without padding.
+        """
+        if not self.training and not self._compile_eval:
+            return self._forward_eager(
+                coord,
+                atype,
+                box,
+                fparam,
+                aparam,
+                do_atomic_virial,
+                charge_spin,
+                spin,
+                n_node,
+            )
+        del do_atomic_virial
+        if self._graph_eligible is None:
+            self._graph_eligible = model_uses_graph_lower(self.original_model)
+        if not self._graph_eligible:
+            raise NotImplementedError(
+                "a flat node axis requires a model whose descriptor reads one; "
+                "this model compiles the dense (nlist) lower, whose batches "
+                "must be padded to a common atom count"
+            )
+        return self._forward_graph(
+            coord,
+            atype,
+            box,
+            fparam,
+            aparam,
+            charge_spin,
+            spin,
+            int(n_node.shape[0]),
+            0,
+            self.original_model.get_rcut(),
+            n_node=n_node,
+        )
+
     def forward(
         self,
         coord: torch.Tensor,
@@ -1100,16 +1240,16 @@ def forward(
         spin: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         if not self.training and not self._compile_eval:
-            kwargs = {
-                "box": box,
-                "fparam": fparam,
-                "aparam": aparam,
-                "do_atomic_virial": do_atomic_virial,
-                "charge_spin": charge_spin,
-            }
-            if spin is not None:
-                kwargs["spin"] = spin
-            return self.original_model(coord, atype, **kwargs)
+            return self._forward_eager(
+                coord,
+                atype,
+                box,
+                fparam,
+                aparam,
+                do_atomic_virial,
+                charge_spin,
+                spin,
+            )
 
         from deepmd.dpmodel.utils.nlist import (
             build_neighbor_list,
@@ -1142,13 +1282,7 @@ def forward(
                 rcut,
             )
 
-        if spin is not None:
-            raise NotImplementedError(
-                "model-level spin requires the NeighborGraph lower; this model "
-                "compiles the dense neighbor-list lower"
-            )
-
-        sel = self.original_model.get_sel()
+        nsel = self.original_model.get_nsel()
 
         # coord extension + nlist (data-dependent, run in eager)
         coord_3d = coord.detach().reshape(nframes, nloc, 3)
@@ -1167,20 +1301,30 @@ def forward(
             ext_atype,
             nloc,
             rcut,
-            sel,
+            nsel,
             # Keep the candidate list merged, matching eager training's
             # DefaultNeighborList contract. forward_common_lower always calls
             # model.format_nlist, which performs the type split for non-mixed
-            # descriptors. The shared builder globally truncates to sum(sel)
-            # before either split, so distinguish_types=True would only move
-            # the same layout transform earlier without changing the final
-            # formatted neighbor list consumed by the lower model.
+            # descriptors. The shared builder globally truncates to nsel before
+            # either split, so distinguish_types=True would only move the same
+            # layout transform earlier without changing the final formatted
+            # neighbor list consumed by the lower model. SpinModel.get_nsel()
+            # exposes the real-input count; its lower expands that nlist for
+            # the virtual types together with coordinates and spins.
             distinguish_types=False,
             # model-level pair exclusion is a nlist-BUILD transform (decision
             # #18/A4); the compiled dense lower consumes a pre-excluded nlist.
             pair_excl=self.original_model.atomic_model.pair_excl,
         )
         ext_coord = ext_coord.reshape(nframes, -1, 3)
+        ext_spin = None
+        if spin is not None:
+            spin = spin.reshape(nframes, nloc, 3)
+            ext_spin = torch.gather(
+                spin,
+                1,
+                mapping.unsqueeze(-1).expand(-1, -1, 3),
+            )
 
         # Mirror the uncompiled path's optional-input defaulting (see
         # ``SeZMModel._forward_common`` -> ``convert_fparam_aparam`` /
@@ -1245,6 +1389,7 @@ def forward(
                     self.original_model,
                     ext_coord,
                     ext_atype,
+                    ext_spin,
                     nlist,
                     mapping,
                     fparam,
@@ -1286,6 +1431,7 @@ def forward(
             result = compiled_lower(
                 ext_coord,
                 ext_atype,
+                ext_spin,
                 nlist,
                 mapping,
                 fparam,
@@ -1295,42 +1441,36 @@ def forward(
             )
         self._report_pending_compile()
 
-        # Translate forward_lower keys -> forward keys.  OUTPUT-AGNOSTIC:
-        # every key passes through unchanged (energy models emit
-        # atom_energy/energy/virial/..., property/dos/... models their own
-        # keys -- the hardcoded energy-key copy here used to KeyError on any
-        # non-energy fitting), EXCEPT the extended-region keys
-        # ``extended_force`` (nf, nall, 3) and ``extended_virial``
-        # (nf, nall, 9): their ghost rows are scatter-summed back onto
-        # local owners via ``mapping`` -- the same fold
-        # ``communicate_extended_output`` performs in the uncompiled path
-        # (which exposes them as ``force`` / ``atom_virial``).  Folding
-        # both keeps the compiled and uncompiled outputs key-for-key
-        # consistent, including for a future atom-virial training
-        # objective.
+        # Translate forward_lower keys -> forward keys. OUTPUT-AGNOSTIC:
+        # ordinary keys pass through unchanged; extended derivatives fold
+        # ghost rows back onto their local owners through ``mapping``, exactly
+        # as ``communicate_extended_output`` does in eager execution. A mask
+        # labels nodes rather than additive contributions, so its public value
+        # is the local prefix of the extended axis.
         out: dict[str, torch.Tensor] = {}
-        if "extended_force" in result:
-            ext_force = result["extended_force"]  # (nf, nall, 3)
-            idx = mapping.unsqueeze(-1).expand_as(ext_force)  # (nf, nall, 3)
-            force = torch.zeros(
-                nframes, nloc, 3, dtype=ext_force.dtype, device=ext_force.device
-            )
-            force.scatter_add_(1, idx, ext_force)
-            out["force"] = force
-        if "extended_virial" in result:
-            ext_virial = result["extended_virial"]  # (nf, nall, 9)
-            idx = mapping.unsqueeze(-1).expand_as(ext_virial)  # (nf, nall, 9)
-            atom_virial = torch.zeros(
+        folded_keys = {
+            "extended_force": "force",
+            "extended_force_mag": "force_mag",
+            "extended_virial": "atom_virial",
+        }
+        for extended_key, local_key in folded_keys.items():
+            if extended_key not in result:
+                continue
+            value = result[extended_key]
+            index = mapping.unsqueeze(-1).expand_as(value)
+            local = torch.zeros(
                 nframes,
                 nloc,
-                ext_virial.shape[-1],
-                dtype=ext_virial.dtype,
-                device=ext_virial.device,
+                value.shape[-1],
+                dtype=value.dtype,
+                device=value.device,
             )
-            atom_virial.scatter_add_(1, idx, ext_virial)
-            out["atom_virial"] = atom_virial
+            out[local_key] = local.scatter_add(1, index, value)
+        if "extended_mask_mag" in result:
+            out["mask_mag"] = result["extended_mask_mag"][:, :nloc]
+        translated_extended_keys = {*folded_keys, "extended_mask_mag"}
         for key, val in result.items():
-            if key not in ("extended_force", "extended_virial"):
+            if key not in translated_extended_keys:
                 out[key] = val
         return out
 
@@ -1346,6 +1486,7 @@ def _forward_graph(
         nframes: int,
         nloc: int,
         rcut: float,
+        n_node: torch.Tensor | None = None,
     ) -> dict[str, torch.Tensor]:
         """Carry-all GRAPH forward -> compiled ``forward_common_lower_graph``.
 
@@ -1356,20 +1497,39 @@ def _forward_graph(
         so no extended->local scatter is needed; only the flat ``(N, *)`` node
         keys are unravelled to ``(nf, nloc, *)`` at the I/O boundary.
         """
+        from deepmd.dpmodel.utils.neighbor_graph import (
+            compact_nodes,
+            expand_node_values,
+        )
         from deepmd.pt_expt.utils.graph_builder import (
             build_neighbor_graph_for_method,
+            build_ragged_neighbor_graph,
         )
 
         _model = self.original_model
 
-        coord_3d = coord.detach().reshape(nframes, nloc, 3)
+        # A ragged batch already holds the node axis the lower works on; a
+        # rectangular one is unravelled to it, and its per-atom outputs are
+        # folded back at the end.
+        ragged = n_node is not None
+        n_padded = nframes * nloc
+        # The builders take the shape the layout hands over: a flat node axis,
+        # or the rectangular one a padded batch carries.
+        coord_3d = (
+            coord.detach().reshape(-1, 3)
+            if ragged
+            else coord.detach().reshape(nframes, nloc, 3)
+        )
         box_flat = box.detach().reshape(nframes, 9) if box is not None else None
-        # Graph-lower per-node inputs are flat: aparam is (N, nda) and spin is
-        # (N, 3), matching the synthetic trace schema.
-        if aparam is not None:
-            aparam = aparam.reshape(nframes * nloc, -1)
-        if spin is not None:
-            spin = spin.reshape(nframes * nloc, 3)
+        # Graph-lower per-node inputs are flat: aparam is ``(N, nda)`` and
+        # spin is ``(N, 3)``, matching the synthetic trace schema. A ragged
+        # batch already carries that layout; only rectangular inputs need to
+        # be unfolded from their frame axis.
+        if not ragged:
+            if aparam is not None:
+                aparam = aparam.reshape(n_padded, -1)
+            if spin is not None:
+                spin = spin.reshape(n_padded, 3)
 
         # Mirror the optional-input defaulting of the dense path / eager
         # call_common: a model configured with fparam / charge_spin substitutes
@@ -1406,15 +1566,28 @@ def _forward_graph(
         # into edge_mask here so the compiled lower consumes a pre-excluded graph
         # (the lower no longer re-applies it), matching the eager path exactly.
         pair_excl = _model.atomic_model.pair_excl
-        ng = build_neighbor_graph_for_method(
-            getattr(_model, "neighbor_graph_method", "dense"),
-            coord_3d,
-            atype,
-            box_flat,
-            rcut,
-            pair_excl,
-        )
-        atype_flat = atype.reshape(nframes * nloc)
+        method = getattr(_model, "neighbor_graph_method", "dense")
+        if ragged:
+            ng = build_ragged_neighbor_graph(
+                method, coord_3d, atype, n_node, box_flat, rcut, pair_excl
+            )
+            atype_flat, node_index = atype, None
+        else:
+            ng = build_neighbor_graph_for_method(
+                method, coord_3d, atype, box_flat, rcut, pair_excl
+            )
+            # A rectangular batch of unequal atom counts is padded to a common
+            # width with phantom atoms (atype < 0). The builders leave them out
+            # of every edge, so dropping them from the node axis costs nothing
+            # and spares the network from evaluating them. On a batch of
+            # uniform atom count this is a renumbering by the identity.
+            atype_flat = atype.reshape(n_padded)
+            ng, node_index = compact_nodes(ng, atype_flat >= 0)
+            atype_flat = atype_flat[node_index]
+            if aparam is not None:
+                aparam = aparam[node_index]
+            if spin is not None:
+                spin = spin[node_index]
 
         # Lazy compile of the GRAPH lower (cached per structure key and mode).
         compiled_lower = self._compiled_lower_by_mode.get(self.training)
@@ -1485,9 +1658,14 @@ def _forward_graph(
 
         # The compiled graph lower emits PUBLIC keys on the FLAT node axis
         # (``atom_energy`` / ``force`` are (N, *); ``energy`` / ``virial`` are
-        # (nf, *)).  Unravel the node-level keys to rectangular (nf, nloc, *) so
-        # callers receive the same shapes as the dense path.
-        N = nframes * nloc
+        # (nf, *)). A ragged caller reads that axis directly. A rectangular one
+        # has its node-level keys scattered back onto the padded width and
+        # unravelled to (nf, nloc, *), where a phantom slot reads zero -- what
+        # a masked-out atom contributed there before.
+        if ragged:
+            result["n_node"] = ng.n_node
+            return result
+        N = node_index.shape[0]
         # Node-level (per-atom, lead dim N) public keys emitted by the graph
         # lower; the remaining keys are frame-level (lead dim nf) and must NOT
         # be unravelled. Keying on the NAME rather than the ``N != nframes``
@@ -1509,7 +1687,9 @@ def _forward_graph(
                 and val is not None
                 and val.shape[:1] == torch.Size([N])
             ):
-                out[key] = val.reshape(nframes, nloc, *val.shape[1:])
+                out[key] = expand_node_values(val, node_index, n_padded).reshape(
+                    nframes, nloc, *val.shape[1:]
+                )
             else:
                 out[key] = val
         return out
@@ -1702,6 +1882,11 @@ def __init__(
             else self.valid_numb_batch_by_task[DEFAULT_TASK_KEY]
         )
 
+        # The layout settles before the schedule below, because under
+        # ``mix:N`` it is what decides how many frames a batch holds, and the
+        # schedule counts those batches to turn epochs into steps.
+        self._configure_batch_layout(training_data, validation_data)
+
         # Statistics ----------------------------------------------------------
         self._finetune_update_stat = False
         self._sample_funcs: dict[str, Any] = {}
@@ -2202,6 +2387,58 @@ def _raise_if_full_validation_unsupported(
     # torch.compile helpers
     # ------------------------------------------------------------------
 
+    def _configure_batch_layout(self, *data_maps: Any) -> None:
+        """Ask each LMDB data system for the layout its own model can consume.
+
+        Under ``mix:N``, a model whose descriptor reads a flat node axis takes
+        the frames of a batch concatenated, which spares it the padding that
+        frames of unequal atom count would otherwise need. Every other model
+        reads an ``(nf, nloc, ...)`` axis and needs them padded to a common
+        width. Non-mixing batch rules retain their established rectangular
+        layout. Only the trainer sees both sides, and it settles the question
+        here, once, before any batch is drawn.
+
+        A graph lower is necessary but not sufficient: the model must expose an
+        entry that takes the flat axis, and the configured loss must accept that
+        representation. Composed models (linear, ZBL bridging) have no such
+        model entry. Generalized-force and Hessian losses require rectangular
+        label axes. Traditional virtual-atom spin models expose no flat-axis
+        entry and therefore remain rectangular; native-spin models follow the
+        same graph capability contract once their loss accepts flat magnetic
+        outputs. Requiring both capabilities keeps the complete objective on a
+        layout it can evaluate.
+
+        Each task has its own data system and its own model, so the answer is
+        each task's own: a multi-task run pairing a graph model with a dense
+        one gives the first concatenated batches and the second padded ones.
+
+        This must run before the run length is resolved. Under ``mix:N`` the
+        layout decides where the sampler cuts a batch -- padding is what makes
+        a batch's cost depend on its widest frame -- so a schedule that
+        counted batches first would count a packing training never uses, and
+        would serve its first epoch in that packing.
+
+        Parameters
+        ----------
+        *data_maps : Any
+            The training and validation data systems, either bare or as the
+            per-task mappings a multi-task run builds.
+        """
+        for task_key in self.model_keys:
+            model = self.models[task_key]
+            loss = self.losses[task_key]
+            ragged = (
+                model_uses_graph_lower(model)
+                and hasattr(model, "forward_ragged")
+                and loss.supports_ragged_batches
+            )
+            for data_map in data_maps:
+                data = (
+                    data_map.get(task_key) if isinstance(data_map, dict) else data_map
+                )
+                if hasattr(data, "use_ragged_batches"):
+                    data.use_ragged_batches(ragged)
+
     def _configure_neighbor_graph_method(self, requested: str) -> None:
         """Resolve and install the training graph builder on eligible models."""
         graph_models = [
@@ -3002,7 +3239,10 @@ def evaluate_validation(
                 label=val_label,
                 task_key=task.key,
             )
-            natoms = int(val_input["atype"].shape[-1])
+            # The metrics are per-atom quantities, so each batch weighs by the
+            # real atoms it holds summed over its frames. Phantom atoms
+            # (atype < 0), which pad a mixed-nloc batch, contribute to none.
+            natoms = int((val_input["atype"] >= 0).sum())
             sum_natoms += natoms
             for key, value in vmore.items():
                 if "l2_" not in key:
diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py
index 28b285d434..32af3a0c42 100644
--- a/deepmd/pt_expt/train/validation.py
+++ b/deepmd/pt_expt/train/validation.py
@@ -444,10 +444,10 @@ def _iter_validation_data_systems(self) -> Iterator[Any]:
         - An LMDB-backed dataset owns an ``_reader``. Its frames are lazily
           materialized into a :class:`LmdbTestData` snapshot (cached across
           calls) and yielded as one :class:`LmdbTestDataNlocView` per
-          atom-count and label-availability group. Grouping by atom count lets
-          mixed-nloc frames be stacked, and grouping by label availability
-          keeps the scalar ``find_*`` flags valid so default-filled labels stay
-          out of the metrics.
+          atom-count and label-availability group. Grouping by atom count
+          lets mixed-nloc frames be stacked, and grouping by availability
+          keeps the scalar ``find_*`` flags valid so default-filled labels
+          stay out of the metrics.
         - A ``DeepmdDataSystem`` owns ``data_systems``, which are already
           ``DeepmdData`` instances.
         - A loader set owns ``systems``, each wrapping a ``DeepmdData`` in
@@ -456,10 +456,11 @@ def _iter_validation_data_systems(self) -> Iterator[Any]:
         validation_data = self.validation_data
         if hasattr(validation_data, "_reader"):
             lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data)
-            for (nloc, _signature), indices in sorted(
-                lmdb_test_data.find_signature_groups.items()
-            ):
-                yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices)
+            for nloc in sorted(lmdb_test_data.nloc_groups):
+                for indices in lmdb_test_data.availability_groups(
+                    lmdb_test_data.nloc_groups[nloc]
+                ):
+                    yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices)
             return
 
         if hasattr(validation_data, "data_systems"):
diff --git a/deepmd/pt_expt/train/wrapper.py b/deepmd/pt_expt/train/wrapper.py
index 7e3bc76ff1..42d9aa0710 100644
--- a/deepmd/pt_expt/train/wrapper.py
+++ b/deepmd/pt_expt/train/wrapper.py
@@ -161,6 +161,7 @@ def forward(
         task_key: str | None = None,
         do_atomic_virial: bool = False,
         charge_spin: torch.Tensor | None = None,
+        n_node: torch.Tensor | None = None,
     ) -> tuple[dict[str, torch.Tensor], torch.Tensor | None, dict | None]:
         if not self.multi_task:
             task_key = "Default"
@@ -177,6 +178,7 @@ def forward(
             "fparam": fparam,
             "aparam": aparam,
             "charge_spin": charge_spin,
+            "n_node": n_node,
         }
         # ``spin`` (native or virtual-atom magnetic moment) is only accepted
         # by spin-capable model forward()s; mirrors
@@ -201,6 +203,10 @@ def forward(
             if label is None:
                 return model_pred, None, None
 
+            # The width a rectangular batch pads its frames to. A ragged batch
+            # has none, so this is its whole atom count and the loss reads each
+            # frame's own from ``model_pred["n_node"]``; terms that cannot
+            # express themselves per frame refuse such a batch instead.
             natoms = atype.shape[-1]
             loss, more_loss = self.loss[task_key](
                 cur_lr,
@@ -238,8 +244,17 @@ def _forward_without_loss(
         task_key: str,
         input_dict: dict[str, Any],
     ) -> dict[str, torch.Tensor]:
-        """Return model predictions without constructing a loss."""
-        return self.model[task_key](**input_dict)
+        """Return model predictions without constructing a loss.
+
+        ``n_node`` marks a batch whose frames are concatenated rather than
+        padded to a common atom count. Its node axis is already the one the
+        graph lower works on, so it takes the entry that skips the padding
+        round trip; ``forward`` accepts only the rectangular shape.
+        """
+        model = self.model[task_key]
+        if input_dict.get("n_node") is None:
+            return model(**{k: v for k, v in input_dict.items() if k != "n_node"})
+        return model.forward_ragged(**input_dict)
 
     def set_extra_state(self, state: dict) -> None:
         self.model_params = state.get("model_params", {})
diff --git a/deepmd/pt_expt/utils/edge_schema.py b/deepmd/pt_expt/utils/edge_schema.py
index 4e916abea0..b2cbad736d 100644
--- a/deepmd/pt_expt/utils/edge_schema.py
+++ b/deepmd/pt_expt/utils/edge_schema.py
@@ -18,6 +18,37 @@
 _DUMMY_EDGE_COUNT = 2
 
 
+def real_atom_edge_mask(
+    atype_flat: torch.Tensor,
+    src: torch.Tensor,
+    dst: torch.Tensor,
+) -> torch.Tensor:
+    """Keep only the edges whose two endpoints are both real atoms.
+
+    Phantom atoms (``atype < 0``) pad a mixed-nloc batch to a rectangular
+    shape: they occupy a tensor slot but no physical site. A geometric
+    neighbor search cannot know that, and reports them as neighbors of
+    whichever real atoms lie near their placeholder coordinates. Dropping
+    those edges restores the invariant that a phantom neither carries an
+    environment of its own nor enters a real atom's.
+
+    Parameters
+    ----------
+    atype_flat : torch.Tensor
+        Atom types flattened over the batch, with shape ``(nf * nloc,)``.
+    src : torch.Tensor
+        Source endpoint of each edge, indexing ``atype_flat``.
+    dst : torch.Tensor
+        Destination endpoint of each edge, indexing ``atype_flat``.
+
+    Returns
+    -------
+    torch.Tensor
+        Boolean mask over edges, with shape ``(nedge,)``.
+    """
+    return (atype_flat[src] >= 0) & (atype_flat[dst] >= 0)
+
+
 def _append_dummy_edges(
     edge_index: torch.Tensor,
     edge_vec: torch.Tensor,
@@ -119,8 +150,21 @@ def edge_schema_from_extended(
 
     # No ``edge_len2 <= rcut**2`` upper bound here: ``nlist`` is contractually
     # cutoff-truncated by the caller (see the docstring). Only padding (-1),
-    # ghost-only neighbours, and coincident pairs are dropped.
-    edge_keep = valid_flat & (src_local >= 0) & (src_local < nloc) & (edge_len2 > 1e-10)
+    # ghost-only neighbours, coincident pairs and phantom endpoints are
+    # dropped. Without ``mapping`` the source index spans the extended atoms,
+    # so it is clamped before the type lookup: entries the preceding terms
+    # already reject must still not index out of bounds.
+    edge_keep = (
+        valid_flat
+        & (src_local >= 0)
+        & (src_local < nloc)
+        & (edge_len2 > 1e-10)
+        & real_atom_edge_mask(
+            atype[:, :nloc].reshape(-1),
+            src_actual.clamp(0, nf * nloc - 1),
+            dst_actual,
+        )
+    )
     valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten()
     edge_index = torch.stack(
         [
@@ -206,7 +250,11 @@ def edge_schema_from_neighbor_matrix(
             edge_vec_all.index_add_(0, shifted_idx, shift_cart)
 
     edge_len2 = torch.sum(edge_vec_all * edge_vec_all, dim=-1)
-    edge_keep = (edge_len2 > 1e-10) & (edge_len2 <= float(rcut) * float(rcut))
+    edge_keep = (
+        (edge_len2 > 1e-10)
+        & (edge_len2 <= float(rcut) * float(rcut))
+        & real_atom_edge_mask(atype.reshape(-1), src_actual, dst)
+    )
     valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten()
     schema = _append_dummy_edges(
         torch.stack(
@@ -272,7 +320,11 @@ def edge_schema_from_ij_shifts(
                 (sel_shifts[:, :, None] * cell).sum(1),
             )
     edge_len2 = torch.sum(edge_vec_all * edge_vec_all, dim=-1)
-    edge_keep = (edge_len2 > 1e-10) & (edge_len2 <= float(rcut) * float(rcut))
+    edge_keep = (
+        (edge_len2 > 1e-10)
+        & (edge_len2 <= float(rcut) * float(rcut))
+        & real_atom_edge_mask(atype.reshape(-1), jj, ii)
+    )
     valid_idx = torch.nonzero(edge_keep, as_tuple=False).flatten()
     schema = _append_dummy_edges(
         torch.stack(
diff --git a/deepmd/pt_expt/utils/graph_builder.py b/deepmd/pt_expt/utils/graph_builder.py
index ac8ecb0fab..fc7f12d625 100644
--- a/deepmd/pt_expt/utils/graph_builder.py
+++ b/deepmd/pt_expt/utils/graph_builder.py
@@ -8,6 +8,10 @@
 
 import torch
 
+from deepmd.dpmodel.utils.lmdb_data import (
+    PHANTOM_ATOM_TYPE,
+)
+
 if TYPE_CHECKING:
     from deepmd.dpmodel.utils.exclude_mask import (
         PairExcludeMask,
@@ -39,8 +43,8 @@ def resolve_auto_graph_builder(
     ``nv`` (CUDA) or ``dense`` so ``auto_batch_size`` / ``dp test`` do not
     regress to the per-frame loop.
 
-    Policy
-    ------
+    Notes
+    -----
     * CUDA + ``nvalchemiops``: ``nv`` (any ``nf``).
     * ``nf == 1`` + ``vesin.torch``: ``vesin``.
     * otherwise: ``dense``.
@@ -171,6 +175,82 @@ def resolve_neighbor_graph_method(
     return "nv"
 
 
+def build_ragged_neighbor_graph(
+    method: str,
+    coord: torch.Tensor,
+    atype: torch.Tensor,
+    n_node: torch.Tensor,
+    box: torch.Tensor | None,
+    rcut: float,
+    pair_excl: "PairExcludeMask | None",
+    *,
+    with_csr: bool = False,
+) -> "NeighborGraph":
+    """Build a carry-all graph over a batch whose node axis is already flat.
+
+    The searches all take a rectangular tensor -- ``dense`` compares every pair
+    of one, and the others derive their per-frame bounds from its shape -- so
+    the frames are widened to a common width here and the resulting graph is
+    narrowed back onto the flat axis. Widening is a scatter and narrowing a
+    renumbering that drops no edge, since no builder draws one to a padded
+    slot. ``nv`` additionally withholds the padded slots from the search
+    itself, so for that backend the widening does not reach the geometry.
+
+    Parameters
+    ----------
+    method : str
+        The concrete builder, as resolved by :func:`resolve_neighbor_graph_method`.
+    coord : torch.Tensor
+        Local coordinates with shape ``(N, 3)``, frame-major over ``n_node``.
+    atype : torch.Tensor
+        Local atom types with shape ``(N,)``.
+    n_node : torch.Tensor
+        Atoms per frame with shape ``(nf,)``.
+    box : torch.Tensor or None
+        Simulation cell with shape ``(nf, 3, 3)``, or ``None`` for non-periodic.
+    rcut : float
+        Cutoff radius.
+    pair_excl : PairExcludeMask or None
+        Model-level pair exclusion, folded into the edge mask at build time.
+    with_csr : bool, default: False
+        Whether to attach destination/source CSR views.
+
+    Returns
+    -------
+    NeighborGraph
+        A graph whose node axis is the one described by ``n_node``.
+    """
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        compact_nodes,
+    )
+
+    nf = int(n_node.shape[0])
+    width = int(n_node.max()) if nf else 0
+    # Position of each flat node in the padded batch, built without a pass over
+    # the frames: its frame times the width, plus its rank within that frame.
+    frame = torch.repeat_interleave(
+        torch.arange(nf, dtype=n_node.dtype, device=n_node.device), n_node
+    )
+    offset = torch.cumsum(n_node, 0) - n_node
+    slot = (
+        torch.arange(int(n_node.sum()), dtype=n_node.dtype, device=n_node.device)
+        - offset[frame]
+    )
+    padded_index = frame * width + slot
+
+    padded_coord = coord.new_zeros((nf * width, 3))
+    padded_atype = atype.new_full((nf * width,), PHANTOM_ATOM_TYPE)
+    padded_coord[padded_index] = coord
+    padded_atype[padded_index] = atype
+    padded_coord = padded_coord.reshape(nf, width, 3)
+    padded_atype = padded_atype.reshape(nf, width)
+
+    graph = build_neighbor_graph_for_method(
+        method, padded_coord, padded_atype, box, rcut, pair_excl, with_csr=with_csr
+    )
+    return compact_nodes(graph, padded_atype.reshape(-1) >= 0)[0]
+
+
 def build_neighbor_graph_for_method(
     method: str,
     coord: torch.Tensor,
diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py
index 7608d9eb61..9af91be82c 100644
--- a/deepmd/pt_expt/utils/lmdb_dataset.py
+++ b/deepmd/pt_expt/utils/lmdb_dataset.py
@@ -14,10 +14,10 @@
 )
 
 from deepmd.dpmodel.utils.lmdb_data import (
-    DistributedSameNlocBatchSampler,
+    DistributedLmdbBatchSampler,
     LmdbBatchIterator,
+    LmdbBatchSampler,
     LmdbDataReader,
-    SameNlocBatchSampler,
     collect_lmdb_sampling_groups,
     compute_block_targets,
 )
@@ -39,11 +39,11 @@ class LmdbDataSystem:
     ``get_nsystems()``, and the ``nbatches``/``sys_probs`` pair from which the
     trainer derives an epoch length. The whole LMDB counts as one logical
     system. Internally uses :class:`LmdbDataReader` for I/O and
-    :class:`SameNlocBatchSampler`, or its distributed wrapper, to draw
-    same-nloc batches. Statistics use a separate logical-system view in which
-    every ``(nloc, label-availability)`` group is sampled independently,
-    matching the training sampler without changing the identity of the LMDB as
-    one training dataset.
+    :class:`LmdbBatchSampler`, or its distributed wrapper, to draw batches.
+    Statistics use a separate logical-system view in which every
+    ``(nloc, label-availability)`` group is sampled independently, matching
+    the training sampler without changing the identity of the LMDB as one
+    training dataset.
 
     Parameters
     ----------
@@ -52,12 +52,13 @@ class LmdbDataSystem:
     type_map
         Global type map from the model config.
     batch_size
-        Batch size spec; ``int``, ``"auto"``, or ``"auto:N"``.
+        Batch size spec; ``int``, ``"auto"``, ``"auto:N"``, ``"max:N"``,
+        ``"filter:N"``, or ``"mix:N"`` for mixed-nloc batching.
     auto_prob_style
         Optional ``auto_prob`` string (e.g. ``"prob_sys_size"``) for
         per-system reweighting via :func:`compute_block_targets`.
     seed
-        Optional seed for the shuffle in :class:`SameNlocBatchSampler`.
+        Optional seed for the shuffle in :class:`LmdbBatchSampler`.
     num_workers
         Number of LMDB decoder worker processes. ``None`` selects the
         hardware-aware default; zero or one disables multiprocessing.
@@ -65,7 +66,7 @@ class LmdbDataSystem:
         Rank of this process in distributed training.
     world_size
         Number of distributed training processes. Values greater than one
-        select :class:`DistributedSameNlocBatchSampler`.
+        select :class:`DistributedLmdbBatchSampler`.
     """
 
     def __init__(
@@ -79,9 +80,7 @@ def __init__(
         rank: int = 0,
         world_size: int = 1,
     ) -> None:
-        self._reader = LmdbDataReader(
-            lmdb_path, type_map, batch_size, mixed_batch=False
-        )
+        self._reader = LmdbDataReader(lmdb_path, type_map, batch_size)
 
         block_targets = None
         if auto_prob_style is not None and self._reader.frame_system_ids is not None:
@@ -92,23 +91,23 @@ def __init__(
             )
 
         if world_size > 1:
-            distributed_sampler = DistributedSameNlocBatchSampler(
-                self._reader,
-                rank=rank,
-                world_size=world_size,
-                shuffle=True,
-                seed=seed,
-                block_targets=block_targets,
+            self._sampler: LmdbBatchSampler | DistributedLmdbBatchSampler = (
+                DistributedLmdbBatchSampler(
+                    self._reader,
+                    rank=rank,
+                    world_size=world_size,
+                    shuffle=True,
+                    seed=seed,
+                    block_targets=block_targets,
+                )
             )
-            self._sampler = distributed_sampler
         else:
-            sampler = SameNlocBatchSampler(
+            self._sampler = LmdbBatchSampler(
                 self._reader,
                 shuffle=True,
                 seed=seed,
                 block_targets=block_targets,
             )
-            self._sampler = sampler
         self._refresh_stat_groups()
         num_workers = (
             get_lmdb_num_workers() if num_workers is None else int(num_workers)
@@ -120,7 +119,7 @@ def __init__(
         )
 
     def _refresh_stat_groups(self) -> None:
-        """Rebuild statistical systems from the training sampler's groups."""
+        """Build statistical systems from stack-compatible reader groups."""
         self._stat_groups = collect_lmdb_sampling_groups(self._reader)
         self._stat_offsets = [0] * len(self._stat_groups)
 
@@ -128,6 +127,20 @@ def _refresh_stat_groups(self) -> None:
     # pt_expt trainer surface
     # ------------------------------------------------------------------
 
+    def use_ragged_batches(self, ragged: bool) -> None:
+        """Select the layout :meth:`get_batch` delivers.
+
+        See :meth:`deepmd.dpmodel.utils.lmdb_data.LmdbDataReader.use_ragged_batches`.
+        The trainer calls this once it knows whether the model reads a flat
+        node axis, before training draws its first batch.
+
+        Parameters
+        ----------
+        ragged : bool
+            Whether to concatenate frames instead of padding them.
+        """
+        self._reader.use_ragged_batches(ragged)
+
     def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]:
         """Return one batch as a numpy dict.
 
@@ -143,7 +156,7 @@ def get_stat_batch(self, sys_idx: int) -> dict[str, Any]:
         Parameters
         ----------
         sys_idx : int
-            Index into the ``(nloc, label-availability)`` groups.
+            Index into the per-atom-count groups.
 
         Returns
         -------
@@ -154,6 +167,12 @@ def get_stat_batch(self, sys_idx: int) -> dict[str, Any]:
         ------
         IndexError
             If ``sys_idx`` does not identify an available statistical group.
+
+        Notes
+        -----
+        The batch is rectangular whatever layout training uses: output
+        statistics accumulate over an ``(nf, nloc, ...)`` axis. A group is
+        uniform in atom count, so that layout pads nothing.
         """
         if not 0 <= sys_idx < len(self._stat_groups):
             raise IndexError(
@@ -168,7 +187,7 @@ def get_stat_batch(self, sys_idx: int) -> dict[str, Any]:
             start = 0
         stop = min(start + batch_size, len(group_indices))
         self._stat_offsets[sys_idx] = stop
-        return self._reader.decode_batch(group_indices[start:stop])
+        return self._reader.decode_batch(group_indices[start:stop], ragged=False)
 
     def get_stat_nsystems(self) -> int:
         """Return the number of homogeneous statistical systems."""
@@ -189,13 +208,8 @@ def get_stat_numb_batches(self, sys_idx: int) -> int:
     def add_data_requirements(
         self, data_requirement: list[DataRequirementItem]
     ) -> None:
-        # Batches are partitioned by label availability. The sampler derives
-        # the partition on its first draw; only the distributed batch count is
-        # cached, so it is refreshed after the requirements change.
         self._reader.add_data_requirement(data_requirement)
         self._refresh_stat_groups()
-        if isinstance(self._sampler, DistributedSameNlocBatchSampler):
-            self._sampler.refresh_batch_count()
 
     def close(self) -> None:
         """Cancel prefetched work and release decoder processes."""
@@ -217,9 +231,7 @@ def get_nsystems(self) -> int:
     @property
     def nbatches(self) -> list[int]:
         """Return the global batch count of one full pass."""
-        if isinstance(self._sampler, DistributedSameNlocBatchSampler):
-            return [self._sampler.total_batches]
-        return [len(self._sampler)]
+        return [self._sampler.total_batches]
 
     @property
     def sys_probs(self) -> list[float]:
diff --git a/deepmd/pt_expt/utils/nv_graph_builder.py b/deepmd/pt_expt/utils/nv_graph_builder.py
index e78fea3e16..aacd1383bc 100644
--- a/deepmd/pt_expt/utils/nv_graph_builder.py
+++ b/deepmd/pt_expt/utils/nv_graph_builder.py
@@ -60,6 +60,7 @@ def nv_matrix_to_ijs(
     num_neighbors: torch.Tensor,
     shifts: torch.Tensor,
     nloc: int,
+    node_index: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
     """Decode nvalchemiops' dense neighbor matrix to a sparse edge list.
 
@@ -72,25 +73,30 @@ def nv_matrix_to_ijs(
     ----------
     neighbor_matrix
         (total_atoms, max_neighbors) int; ``neighbor_matrix[dst, slot] = src``,
-        both flattened batch indices in ``[0, total_atoms)``. Frames are
+        both flattened search indices in ``[0, total_atoms)``. Frames are
         batch-isolated: a neighbor always shares its center's frame.
     num_neighbors
         (total_atoms,) int, valid slot count per center.
     shifts
         (total_atoms, max_neighbors, 3) int periodic image shifts per slot.
     nloc
-        Atoms per frame (``total_atoms = nf * nloc``).
+        Atoms per frame of the rectangular batch the indices are reported on.
+    node_index
+        (total_atoms,) int, position of each searched atom in that rectangular
+        batch. Given when the search ran over a subset of the batch, such as
+        its real atoms alone; ``None`` when it ran over every slot, in which
+        case the search index and the batch position coincide.
 
     Returns
     -------
     center_local
-        (E,) int64 per-frame local center index ``i`` (``dst % nloc``).
+        (E,) int64 per-frame local center index ``i``.
     src_local
-        (E,) int64 per-frame local neighbor index ``j`` (``src % nloc``).
+        (E,) int64 per-frame local neighbor index ``j``.
     shift
         (E, 3) int64 periodic image shift ``S``.
     frame_idx
-        (E,) int64 frame of each edge (``dst // nloc``).
+        (E,) int64 frame of each edge.
     """
     device = neighbor_matrix.device
     total_atoms, max_neighbors = neighbor_matrix.shape
@@ -103,6 +109,12 @@ def nv_matrix_to_ijs(
     dst = edge_idx // max_neighbors  # flattened center
     src = neighbor_matrix.reshape(-1).index_select(0, edge_idx).to(torch.int64)
     shift = shifts.reshape(-1, 3).index_select(0, edge_idx).to(torch.int64)
+    if node_index is not None:
+        # Lift both endpoints from the search axis back onto the rectangular
+        # batch, so that the frame and local indices below are the ones every
+        # other builder reports and the caller needs no special case.
+        dst = node_index.index_select(0, dst)
+        src = node_index.index_select(0, src)
     frame_idx = (dst // nloc).to(torch.int64)  # frame of the edge
     center_local = (dst % nloc).to(torch.int64)  # i = center
     src_local = (src % nloc).to(torch.int64)  # j = neighbor
@@ -114,6 +126,7 @@ def nv_search_matrix(
     box: torch.Tensor | None,
     rcut: float,
     start_capacity: int,
+    node_index: torch.Tensor | None = None,
 ) -> tuple[
     torch.Tensor,
     torch.Tensor | None,
@@ -136,12 +149,19 @@ def nv_search_matrix(
     rcut : cutoff radius.
     start_capacity : initial max-neighbor capacity; grown automatically when
         any atom has more neighbors than the current capacity.
+    node_index : (total_atoms,) positions, in the flattened batch, of the atoms
+        to search over, frame-major and ascending. ``None`` searches every slot.
+        The search itself is ragged -- nvalchemiops takes a flat position array
+        with explicit per-frame bounds -- so restricting it to the real atoms of
+        a padded batch keeps the phantom slots out of a cost that grows with the
+        square of the frame width.
 
     Returns
     -------
     coord : (nf, nloc, 3) coordinates, normalized in-cell if periodic.
     cell : (nf, 3, 3) float box, or ``None`` for non-periodic.
-    neighbor_matrix : (total_atoms, capacity) int neighbor matrix.
+    neighbor_matrix : (total_atoms, capacity) int neighbor matrix, indexed on
+        the searched atoms rather than on the batch slots.
     num_neighbors : (total_atoms,) valid neighbor count per center.
     shifts : (total_atoms, capacity, 3) int periodic image shifts.
     """
@@ -163,16 +183,24 @@ def nv_search_matrix(
             cell = None
             pbc = None
 
-        total_atoms = nf * nloc
-        positions = coord.reshape(total_atoms, 3).detach()
-        batch_idx = torch.arange(
-            nf, dtype=torch.int32, device=device
-        ).repeat_interleave(nloc)
-        batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc
-        method = choose_nv_nlist_method(nloc, periodic=periodic, device=device)
+        positions = coord.reshape(nf * nloc, 3).detach()
+        if node_index is None:
+            batch_idx = torch.arange(
+                nf, dtype=torch.int32, device=device
+            ).repeat_interleave(nloc)
+            batch_ptr = torch.arange(nf + 1, dtype=torch.int32, device=device) * nloc
+            widest_frame = nloc
+        else:
+            positions = positions.index_select(0, node_index)
+            batch_idx = (node_index // nloc).to(torch.int32)
+            counts = torch.bincount(batch_idx.to(torch.int64), minlength=nf)
+            batch_ptr = torch.zeros(nf + 1, dtype=torch.int32, device=device)
+            batch_ptr[1:] = torch.cumsum(counts, 0).to(torch.int32)
+            widest_frame = int(counts.max()) if nf > 0 else 0
+        method = choose_nv_nlist_method(widest_frame, periodic=periodic, device=device)
         extra_nl_kwargs: dict[str, Any] = {}
         if method == "batch_naive":
-            extra_nl_kwargs["max_atoms_per_system"] = int(nloc)
+            extra_nl_kwargs["max_atoms_per_system"] = int(widest_frame)
 
         search_capacity = start_capacity
         while True:
@@ -283,7 +311,7 @@ def build_neighbor_graph_nv(
             coord,
             box,
             empty_i,
-            nloc,
+            torch.full((nf,), nloc, dtype=torch.int64, device=device),
             layout=layout,
             with_csr=with_csr,
             canonicalize=canonicalize,
@@ -300,25 +328,24 @@ def build_neighbor_graph_nv(
         64,
         estimate_max_neighbors(float(rcut), atomic_density=0.25),
     )
+    # Virtual atoms (atype < 0) are excluded as centers AND neighbors -- the
+    # World-2 builder contract shared with the dense reference builder. Here
+    # they are withheld from the search rather than filtered out of its result:
+    # the phantom slots of a mixed-nloc batch would otherwise widen every frame
+    # the search sees, and its cost grows with the square of that width.
+    atype_flat = torch.as_tensor(atype, device=device).reshape(nf * nloc)
+    node_index = torch.nonzero(atype_flat >= 0, as_tuple=False).flatten()
     coord, cell, neighbor_matrix, num_neighbors, shifts = nv_search_matrix(
-        coord, box, rcut, start_capacity=initial_capacity
+        coord, box, rcut, start_capacity=initial_capacity, node_index=node_index
     )
     box_out = cell  # edge_vec is recomputed from these (normalized) coords
 
-    # Decode the dense matrix to a sparse (i, j, S) edge list (CPU-testable
-    # helper; see nv_matrix_to_ijs).
+    # Decode the dense matrix to a sparse (i, j, S) edge list, lifted back onto
+    # the rectangular batch (CPU-testable helper; see nv_matrix_to_ijs).
     center_local, src_local, shift, frame_idx = nv_matrix_to_ijs(
-        neighbor_matrix, num_neighbors, shifts, nloc
+        neighbor_matrix, num_neighbors, shifts, nloc, node_index=node_index
     )
 
-    # virtual atoms (atype < 0) are excluded as centers AND neighbors — the
-    # World-2 builder contract shared with the dense reference builder; the
-    # geometric search above cannot know about them.
-    at = torch.as_tensor(atype, device=device).reshape(nf, nloc)
-    keep = (at[frame_idx, center_local] >= 0) & (at[frame_idx, src_local] >= 0)
-    center_local, src_local = center_local[keep], src_local[keep]
-    shift, frame_idx = shift[keep], frame_idx[keep]
-
     graph = neighbor_graph_from_ijs(
         center_local,
         src_local,
@@ -326,12 +353,11 @@ def build_neighbor_graph_nv(
         coord,
         box_out,
         frame_idx,
-        nloc,
+        torch.full((nf,), nloc, dtype=torch.int64, device=device),
         layout=layout,
     )
     if pair_excl is not None:
-        at_flat = torch.as_tensor(atype, device=device).reshape(-1)
-        graph = apply_pair_exclusion(graph, at_flat, pair_excl, compact=compact)
+        graph = apply_pair_exclusion(graph, atype_flat, pair_excl, compact=compact)
     if with_csr or canonicalize:
         graph = attach_edge_csr(graph, nf * nloc, canonicalize=canonicalize)
     return graph
diff --git a/deepmd/pt_expt/utils/vesin_graph_builder.py b/deepmd/pt_expt/utils/vesin_graph_builder.py
index 874e979449..cc999bd1af 100644
--- a/deepmd/pt_expt/utils/vesin_graph_builder.py
+++ b/deepmd/pt_expt/utils/vesin_graph_builder.py
@@ -167,7 +167,7 @@ def build_neighbor_graph_vesin(
             coord,
             box,
             empty_nf,
-            nloc,
+            torch.full((nf,), nloc, dtype=torch.int64, device=dev),
             layout=layout,
             with_csr=with_csr,
             canonicalize=canonicalize,
@@ -224,7 +224,7 @@ def build_neighbor_graph_vesin(
         coord,
         box,
         nf_all,
-        nloc,
+        torch.full((nf,), nloc, dtype=torch.int64, device=dev),
         layout=layout,
     )
     if pair_excl is not None:
diff --git a/deepmd/pt_expt/utils/vesin_neighbor_list.py b/deepmd/pt_expt/utils/vesin_neighbor_list.py
index 2d28e1b090..7453176d32 100644
--- a/deepmd/pt_expt/utils/vesin_neighbor_list.py
+++ b/deepmd/pt_expt/utils/vesin_neighbor_list.py
@@ -231,6 +231,14 @@ def _build_single(
     ii, jj, ss = vesin_search_ijs(
         positions.detach(), cell if periodic else None, periodic, rcut, device
     )
+    # Phantom atoms (atype < 0) pad a mixed-nloc batch and have no physical
+    # site; the geometric search above cannot know that. Their pairs are
+    # dropped before the ``sum(sel)`` truncation below, so that a phantom can
+    # never displace a genuine neighbor from a real atom's list. Compacting
+    # the pair list, rather than marking the pairs empty, also keeps them out
+    # of the ``max_nn`` that sizes the dense candidate matrix.
+    real_pair = (atype[ii] >= 0) & (atype[jj] >= 0)
+    ii, jj, ss = ii[real_pair], jj[real_pair], ss[real_pair]
     # ss is int64 from the helper; cast to the coordinate dtype for the image sum.
     ss = ss.to(positions.dtype)
 
diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
index bbbf15acb4..d9d57dfb1a 100644
--- a/deepmd/utils/argcheck.py
+++ b/deepmd/utils/argcheck.py
@@ -288,6 +288,8 @@ def register(
             the name of a descriptor
         alias : list[str], optional
             the list of aliases of this descriptor
+        doc : str, optional
+            the descriptor documentation prefix
 
         Returns
         -------
@@ -5169,6 +5171,7 @@ def training_data_args() -> list[
 - string "mixed:N": the batch data will be sampled from all systems and merged into a mixed system with the batch size N. Only support the se_atten descriptor for TensorFlow backend.\n\n\
 - string "max:N": automatically determines the batch size so that `batch_size * natoms` is at most `N`. `natoms` is the per-system atom count for npy data and the per-frame nloc for LMDB data. When a single system/frame already has more than `N` atoms, the batch size clamps to 1 and that batch will exceed `N`.\n\n\
 - string "filter:N": the same as `"max:N"` but additionally drops data whose atom count exceeds `N`. For npy data this removes whole systems with natoms > `N`; for LMDB data this removes individual frames with nloc > `N`.\n\n\
+- string "mix:N": LMDB data only. Frames of different atom counts share a batch, filled until the next frame would push the atom axis of the batch past `N`. How that axis is laid out follows from the model and needs no configuration of its own: a descriptor reading a flat node axis (the graph route of the PyTorch Exportable backend) takes the frames of a batch concatenated, so `N` counts real atoms and nothing is padded; every other descriptor takes them padded to the widest frame of the batch, so `N` counts the padded slots `nframes * max_nloc` and the shorter frames carry phantom atoms with `atype = -1` that the neighbor list, the model and the loss all skip. Unlike `"max:N"`, which leaves an under-filled batch whenever an nloc group is small, this keeps every batch close to `N` atoms. A lone frame with nloc > `N` still forms a batch of its own.\n\n\
 If MPI is used, the value should be considered as the batch size per task.'
     doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\
 - "prob_uniform"  : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\
@@ -5264,7 +5267,8 @@ def validation_data_args() -> list[
 - string "auto": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than 32.\n\n\
 - string "auto:N": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than N.\n\n\
 - string "max:N": automatically determines the batch size so that `batch_size * natoms` is at most `N`. `natoms` is the per-system atom count for npy data and the per-frame nloc for LMDB data. When a single system/frame already has more than `N` atoms, the batch size clamps to 1 and that batch will exceed `N`.\n\n\
-- string "filter:N": the same as `"max:N"` but additionally drops data whose atom count exceeds `N`. For npy data this removes whole systems with natoms > `N`; for LMDB data this removes individual frames with nloc > `N`.'
+- string "filter:N": the same as `"max:N"` but additionally drops data whose atom count exceeds `N`. For npy data this removes whole systems with natoms > `N`; for LMDB data this removes individual frames with nloc > `N`.\n\n\
+- string "mix:N": LMDB data only. Frames of different atom counts share a batch, filled until the next frame would push the atom axis of the batch past `N`. How that axis is laid out follows from the model: a descriptor reading a flat node axis takes the frames of a batch concatenated, so `N` counts real atoms and nothing is padded; every other descriptor takes them padded to the widest frame of the batch, so `N` counts the padded slots `nframes * max_nloc` and the shorter frames carry phantom atoms that the neighbor list, the model and the loss all skip.'
     doc_auto_prob_style = 'Determine the probability of systems automatically. The method is assigned by this key and can be\n\n\
 - "prob_uniform"  : the probability all the systems are equal, namely 1.0/self.get_nsystems()\n\n\
 - "prob_sys_size" : the probability of a system is proportional to the number of batches in the system\n\n\
diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py
index 4bde3a8f9e..a232aab393 100644
--- a/deepmd/utils/data.py
+++ b/deepmd/utils/data.py
@@ -18,6 +18,7 @@
 )
 from typing import (
     Any,
+    Literal,
 )
 
 import numpy as np
@@ -35,6 +36,8 @@
 
 log = logging.getLogger(__name__)
 
+DataRequirementSourcePolicy = Literal["tracked", "default", "derived"]
+
 
 class DeepmdData:
     """Class for a data system.
@@ -1198,6 +1201,13 @@ class DataRequirementItem:
     special_shape : str, optional
         Name of a loader-defined non-standard shape contract. ``"hessian"``
         stores one full-frame ``(3 * natoms) x (3 * natoms)`` matrix per frame.
+    source_policy : {"tracked", "default", "derived"}, optional
+        How source availability affects the consumer. ``"tracked"`` keeps
+        source presence in the ``find_*`` contract. ``"default"`` treats the
+        configured default as valid input and reports the resolved field
+        available. ``"derived"`` computes the value from structural frame
+        data. Only optional tracked fields require availability-homogeneous
+        batching.
     """
 
     def __init__(
@@ -1213,7 +1223,18 @@ def __init__(
         dtype: np.dtype | None = None,
         output_natoms_for_type_sel: bool = False,
         special_shape: str | None = None,
+        source_policy: DataRequirementSourcePolicy = "tracked",
     ) -> None:
+        if source_policy not in {"tracked", "default", "derived"}:
+            raise ValueError(
+                "source_policy must be 'tracked', 'default', or 'derived', "
+                f"got {source_policy!r}"
+            )
+        if must and source_policy != "tracked":
+            raise ValueError(
+                f"mandatory data requirement {key!r} cannot use "
+                f"source_policy={source_policy!r}"
+            )
         self.key = key
         self.ndof = ndof
         self.atomic = atomic
@@ -1225,6 +1246,7 @@ def __init__(
         self.dtype = dtype
         self.output_natoms_for_type_sel = output_natoms_for_type_sel
         self.special_shape = special_shape
+        self.source_policy = source_policy
         self.dict = self.to_dict()
 
     def to_dict(self) -> dict:
@@ -1239,6 +1261,7 @@ def to_dict(self) -> dict:
             "default": self.default,
             "dtype": self.dtype,
             "output_natoms_for_type_sel": self.output_natoms_for_type_sel,
+            "source_policy": self.source_policy,
         }
         if self.special_shape is not None:
             data["special_shape"] = self.special_shape
diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py
index dfc4076a9f..cea270bd37 100644
--- a/deepmd/utils/data_system.py
+++ b/deepmd/utils/data_system.py
@@ -197,6 +197,12 @@ def __init__(
                         bsi = 1
                     bs.append(bsi)
                 self.batch_size = bs
+            elif words[0] == "mix":
+                raise RuntimeError(
+                    "the 'mix' batch_size rule packs frames of unequal atom "
+                    "count into one batch and is only available for LMDB "
+                    "datasets on the pt and pt_expt backends"
+                )
             else:
                 raise RuntimeError("unknown batch_size rule " + words[0])
         elif isinstance(self.batch_size, list):
diff --git a/doc/data/system.md b/doc/data/system.md
index bfadb8887f..f11057b55a 100644
--- a/doc/data/system.md
+++ b/doc/data/system.md
@@ -82,4 +82,8 @@ With these edited files, one can put together frames with the same `Natoms`, ins
 
 To put frames with different `Natoms` into the same system, one can pad systems by adding virtual atoms whose type is `-1`. Virtual atoms do not contribute to any fitting property, so the atomic property of virtual atoms (e.g. forces) should be given zero.
 
+:::{note}
+Per-atom loss terms (force, atomic energy, dos, tensor) sum over a batch's real labels and divide by their total count, so a frame counts in proportion to its real atom number. This differs from averaging each frame's own per-label mean only when the frames of one batch pad a *different* number of virtual atoms; where the padding is uniform — including the case of no padding at all — the two are identical. Writing `real_atom_types.npy` with a per-frame varying number of `-1` is therefore the one configuration whose loss values, and hence training trajectory, differ from a per-frame average. This applies to the PyTorch backend; the TensorFlow backend normalizes by the padded width regardless.
+:::
+
 The API to generate or transfer to `mixed_type` format is available on [dpdata](https://github.com/deepmodeling/dpdata) for a more convenient experience.
diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md
index e756f63b6c..84cfd30f1a 100644
--- a/doc/train/training-advanced.md
+++ b/doc/train/training-advanced.md
@@ -81,11 +81,26 @@ The sections {ref}`training_data <training/training_data>` and {ref}`validation_
 
 - The key {ref}`batch_size <training/training_data/batch_size>` specifies the number of frames used to train or validate the model in a training step. It can be set to
   - `list`: the length of which is the same as the {ref}`systems`. The batch size of each system is given by the elements of the list.
+
   - `int`: all systems use the same batch size.
+
   - `"auto"`: the same as `"auto:32"`, see `"auto:N"`
+
   - `"auto:N"`: automatically determines the batch size so that the {ref}`batch_size <training/training_data/batch_size>` times the number of atoms in the system is **no less than** `N`.
+
   - `"max:N"`: automatically determines the batch size so that the {ref}`batch_size <training/training_data/batch_size>` times the number of atoms in the system is **no more than** `N`. The minimum batch size is 1. **Supported backends**: PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }}
+
   - `"filter:N"`: the same as `"max:N"` but removes the systems with the number of atoms larger than `N` from the data set. Throws an error if no system is left in a dataset. **Supported backends**: PyTorch {{ pytorch_icon }}, Paddle {{ paddle_icon }}
+
+  - `"mix:N"`: mixed-nloc batching for LMDB data sets. Frames of different atom counts share a batch, which is closed only when the next frame would push its atom axis past `N`. How that axis is laid out follows from the model and needs no configuration of its own:
+
+    - A descriptor reading a **flat** node axis — the graph route of the PyTorch Exportable backend — takes the frames of a batch concatenated. `N` counts real atoms, nothing is padded, and the frames are packed in the shuffled order they arrive in, so one batch stays mixed in system size.
+
+    - Every other descriptor reads a **rectangular** `(nframes, nloc, ...)` axis, on which the frames of a batch must share one atom count. Frames are therefore sorted by atom count and padded up to the widest one of their batch, which makes `N` the padded-slot count `nframes * max_nloc`. The padded slots hold phantom atoms marked `atype = -1`, which the neighbor list, the model and the loss all skip. Sorting is what keeps their number small, and it also makes the packing optimal: no arrangement of the same frames under the same budget yields fewer batches.
+
+    Either way, a lone frame with more than `N` atoms still forms a batch of its own. Every other rule keeps a batch uniform in atom count, which leaves batches under-filled wherever an atom-count group is small; `"mix:N"` keeps them all close to `N` atoms instead. **Supported backends**: PyTorch {{ pytorch_icon }} (rectangular layout only), PyTorch Exportable
+
+    Filling batches also makes the *sampling weight* more faithful. Per-atom loss terms (force, atomic energy) pool over a batch's real labels, so a frame already counts in proportion to its atom count there. Frame-level terms (energy, virial) weigh the frames of a batch equally, so a frame counts for `1 / nframes`, and an atom budget makes `nframes` follow the atom count. `"max:N"` distorts that wherever an atom-count group is too small to fill a batch — its few frames form a short batch and each is over-weighted many times over — while `"mix:N"` packs them with their neighbours and lands much closer to the intended weighting.
 - The key {ref}`numb_batch <training/validation_data/numb_btch>` in {ref}`validate_data <training/validation_data>` gives the number of batches of model validation. Note that the batches may not be from the same system
 
 The section {ref}`mixed_precision <training/mixed_precision>` specifies the mixed precision settings, which will enable the mixed precision training workflow for DeePMD-kit. The keys are explained below:
diff --git a/source/tests/common/dpmodel/test_from_ijs.py b/source/tests/common/dpmodel/test_from_ijs.py
index adce99c5e5..a3dcad7e58 100644
--- a/source/tests/common/dpmodel/test_from_ijs.py
+++ b/source/tests/common/dpmodel/test_from_ijs.py
@@ -18,7 +18,7 @@ def test_edge_vec_and_index(self) -> None:
         j = np.array([1, 0])  # neighbor
         S = np.array([[0, 0, 0], [0, 0, 0]], dtype=np.int64)
         ng = neighbor_graph_from_ijs(
-            i, j, S, coord, box, nframe_id=np.zeros(2, np.int64), nloc=3
+            i, j, S, coord, box, nframe_id=np.zeros(2, np.int64), n_node=np.array([3])
         )
         np.testing.assert_array_equal(ng.edge_index[0][ng.edge_mask], j)  # src
         np.testing.assert_array_equal(ng.edge_index[1][ng.edge_mask], i)  # dst
@@ -34,7 +34,7 @@ def test_periodic_shift_in_edge_vec(self) -> None:
         j = np.array([1])
         S = np.array([[-1, 0, 0]], dtype=np.int64)
         ng = neighbor_graph_from_ijs(
-            i, j, S, coord, box, nframe_id=np.zeros(1, np.int64), nloc=2
+            i, j, S, coord, box, nframe_id=np.zeros(1, np.int64), n_node=np.array([2])
         )
         # coord[1] + (-1,0,0)@box - coord[0] = 5.5 - 6 - 0.5 = -1.0
         np.testing.assert_allclose(
diff --git a/source/tests/common/dpmodel/test_graph_ragged.py b/source/tests/common/dpmodel/test_graph_ragged.py
index a651d245ea..6b1faae6c9 100644
--- a/source/tests/common/dpmodel/test_graph_ragged.py
+++ b/source/tests/common/dpmodel/test_graph_ragged.py
@@ -81,6 +81,66 @@ def test_forward_common_atomic_graph_ragged():
     assert np.all(np.isfinite(out["energy"]))
 
 
+def test_compact_nodes_recounts_frames_and_renumbers_edges():
+    """A padded graph compacts onto its real nodes, frame blocks intact.
+
+    Two frames padded to a width of four: frame 0 holds three real atoms and
+    frame 1 holds two, so the flat axis carries phantoms at 3, 6 and 7.
+    """
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        NeighborGraph,
+        compact_nodes,
+        expand_node_values,
+    )
+
+    node_mask = np.array([1, 1, 1, 0, 1, 1, 0, 0], dtype=bool)
+    graph = NeighborGraph(
+        n_node=np.array([4, 4], dtype=np.int64),
+        edge_index=np.array([[0, 1, 4], [1, 2, 5]], dtype=np.int64),
+        edge_vec=np.arange(9, dtype=np.float64).reshape(3, 3),
+        edge_mask=np.ones(3, dtype=bool),
+    )
+    compacted, keep_index = compact_nodes(graph, node_mask)
+
+    np.testing.assert_array_equal(compacted.n_node, [3, 2])
+    np.testing.assert_array_equal(keep_index, [0, 1, 2, 4, 5])
+    # Frame 1's nodes 4 and 5 sit at 3 and 4 once frame 0's phantom is gone.
+    np.testing.assert_array_equal(compacted.edge_index, [[0, 1, 3], [1, 2, 4]])
+    np.testing.assert_array_equal(
+        frame_id_from_n_node(compacted.n_node), [0, 0, 0, 1, 1]
+    )
+    # The edge payload is untouched: no edge is dropped, only renumbered.
+    np.testing.assert_array_equal(compacted.edge_vec, graph.edge_vec)
+    np.testing.assert_array_equal(compacted.edge_mask, graph.edge_mask)
+
+    # Gathering onto the compact axis and expanding back is the identity on
+    # the real nodes and zero on the phantoms.
+    values = np.arange(8, dtype=np.float64)[:, None]
+    restored = expand_node_values(values[keep_index], keep_index, 8)
+    np.testing.assert_array_equal(
+        restored.ravel(), [0.0, 1.0, 2.0, 0.0, 4.0, 5.0, 0.0, 0.0]
+    )
+
+
+def test_compact_nodes_refuses_to_drop_a_node_that_carries_an_edge():
+    """Renumbering a node with an edge would silently redirect that edge."""
+    import pytest
+
+    from deepmd.dpmodel.utils.neighbor_graph import (
+        NeighborGraph,
+        compact_nodes,
+    )
+
+    graph = NeighborGraph(
+        n_node=np.array([3], dtype=np.int64),
+        edge_index=np.array([[0], [1]], dtype=np.int64),
+        edge_vec=np.zeros((1, 3), dtype=np.float64),
+        edge_mask=np.ones(1, dtype=bool),
+    )
+    with pytest.raises(ValueError, match="still carries an edge"):
+        compact_nodes(graph, np.array([1, 0, 1], dtype=bool))
+
+
 def test_frame_id_rectangular():
     fid = frame_id_from_n_node(np.array([4, 4], dtype=np.int64))
     np.testing.assert_array_equal(
diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py
index 939d2b8a70..16dd1da1d1 100644
--- a/source/tests/common/dpmodel/test_lmdb_data.py
+++ b/source/tests/common/dpmodel/test_lmdb_data.py
@@ -1,9 +1,10 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Unit tests for LmdbDataReader, LmdbTestData, SameNlocBatchSampler, etc.
+"""Unit tests for LmdbDataReader, LmdbTestData, LmdbBatchSampler, etc.
 
 Pure dpmodel (NumPy/lmdb) tests — no PyTorch dependency.
 """
 
+import gc
 import os
 import signal
 import subprocess
@@ -14,6 +15,9 @@
 from concurrent.futures import (
     Future,
 )
+from itertools import (
+    pairwise,
+)
 from pathlib import (
     Path,
 )
@@ -30,15 +34,19 @@
 
 from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module
 from deepmd.dpmodel.utils.lmdb_data import (
+    DistributedLmdbBatchSampler,
     LmdbBatchIterator,
+    LmdbBatchSampler,
     LmdbDataReader,
     LmdbDecodeConfig,
     LmdbTestData,
     LmdbTestDataNlocView,
-    SameNlocBatchSampler,
+    _chop_mixed_nloc,
     _expand_indices_by_blocks,
     _merge_lmdb_chunks,
     _remap_atom_types,
+    collate_lmdb_frames,
+    collect_lmdb_sampling_groups,
     compute_block_targets,
     decode_lmdb_batch,
     decode_lmdb_frame,
@@ -161,6 +169,46 @@ def _create_mixed_nloc_lmdb(path: str) -> str:
     return path
 
 
+def _create_mix_probe_lmdb(path: str) -> str:
+    """Create an LMDB whose frame-level fields collide in shape with the atom axis.
+
+    Frames alternate between 2 and 9 atoms and carry a two-component
+    ``fparam`` alongside a nine-component ``virials``. A padding rule that
+    guessed the atom axis from a leading dimension would misclassify
+    ``fparam`` on the 2-atom frames and ``virials`` on the 9-atom ones, so this
+    fixture pins the classification down to the exact ambiguous shapes.
+    """
+    nlocs = [2, 9, 2, 9, 2, 9]
+    rng = np.random.RandomState(7)
+    env = lmdb.open(path, map_size=10 * 1024 * 1024)
+    with env.begin(write=True) as txn:
+        meta = {
+            "nframes": len(nlocs),
+            "frame_idx_fmt": "012d",
+            "frame_nlocs": nlocs,
+            "system_info": {"natoms": [1, 1], "formula": "probe"},
+        }
+        txn.put(b"__metadata__", msgpack.packb(meta, use_bin_type=True))
+        for idx, natoms in enumerate(nlocs):
+            frame = _make_frame(natoms=natoms, seed=idx)
+            frame["virials"] = {
+                "type": "<f8",
+                "shape": (9,),
+                "data": rng.randn(9).astype(np.float64).tobytes(),
+            }
+            frame["fparam"] = {
+                "type": "<f8",
+                "shape": (2,),
+                "data": rng.randn(2).astype(np.float64).tobytes(),
+            }
+            txn.put(
+                format(idx, "012d").encode(),
+                msgpack.packb(frame, use_bin_type=True),
+            )
+    env.close()
+    return path
+
+
 def _create_lmdb_with_type_map(
     path: str,
     nframes: int = 6,
@@ -375,10 +423,10 @@ def test_uniform_nloc_single_group(self):
     def test_batch_iterator_advances_epoch_before_prefetch(self):
         """The prefetched successor uses the next epoch's sampler state."""
         reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
-        sampler = SameNlocBatchSampler(reader, shuffle=True, seed=7)
+        sampler = LmdbBatchSampler(reader, shuffle=True, seed=7)
         iterator = LmdbBatchIterator(reader, sampler, num_workers=2)
 
-        expected_sampler = SameNlocBatchSampler(reader, shuffle=True, seed=7)
+        expected_sampler = LmdbBatchSampler(reader, shuffle=True, seed=7)
         expected_sampler.set_epoch(1)
         expected_second = [
             key for indices in expected_sampler for key in reader.original_keys(indices)
@@ -417,6 +465,70 @@ def test_requirements_are_rejected_after_the_first_decode(self):
                     reader.add_data_requirement(requirement)
                 reader.close()
 
+    def test_mandatory_fields_reject_missing_and_explicitly_unavailable_sources(self):
+        """Mandatory LMDB inputs fail at the shared frame-resolution boundary."""
+        path = _create_lmdb(
+            f"{self._tmpdir.name}/mandatory_spin.lmdb",
+            nframes=3,
+            natoms=6,
+        )
+        environment = lmdb.open(path, readonly=False, lock=False)
+        with environment.begin(write=True) as transaction:
+            unavailable_key = format(1, "012d").encode()
+            unavailable_frame = msgpack.unpackb(
+                transaction.get(unavailable_key), raw=False
+            )
+            unavailable_frame["spin"] = {
+                "type": "<f8",
+                "shape": (6, 3),
+                "data": np.zeros((6, 3), dtype=np.float64).tobytes(),
+            }
+            unavailable_frame["find_spin"] = 0.0
+            transaction.put(
+                unavailable_key,
+                msgpack.packb(unavailable_frame, use_bin_type=True),
+            )
+
+            inconsistent_key = format(2, "012d").encode()
+            inconsistent_frame = msgpack.unpackb(
+                transaction.get(inconsistent_key), raw=False
+            )
+            inconsistent_frame["find_spin"] = 1.0
+            transaction.put(
+                inconsistent_key,
+                msgpack.packb(inconsistent_frame, use_bin_type=True),
+            )
+        environment.close()
+
+        requirement = [DataRequirementItem("spin", 3, atomic=True, must=True)]
+        reader = LmdbDataReader(path, self._type_map, batch_size=1)
+        reader.add_data_requirement(requirement)
+        with mock.patch.object(
+            lmdb_data_module,
+            "_raw_frame_availability",
+            side_effect=AssertionError("mandatory fields must fail during decode"),
+        ):
+            self.assertEqual(len(collect_lmdb_sampling_groups(reader)), 1)
+        for frame_index, reason in (
+            (0, "field is absent"),
+            (1, "find_spin=0"),
+            (2, "despite find_spin=1"),
+        ):
+            with self.subTest(frame_index=frame_index):
+                with self.assertRaisesRegex(
+                    RuntimeError,
+                    rf"spin.*frame {frame_index}.*mandatory_spin\.lmdb.*{reason}",
+                ):
+                    reader[frame_index]
+
+        test_data = LmdbTestData(path, type_map=self._type_map, shuffle_test=False)
+        test_data.add_data_requirement(requirement)
+        with self.assertRaisesRegex(
+            RuntimeError,
+            r"spin.*frame 0.*mandatory_spin\.lmdb.*field is absent",
+        ):
+            test_data.get_test_by_indices([0])
+
     def test_close_preserves_other_reader(self):
         """Closing one shared-path reader leaves the other transaction valid."""
         first = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
@@ -507,13 +619,13 @@ def test_parallel_batch_consistency_guards(self):
         }
         frame_cases = (
             (
-                "fields",
-                {**first_frame, "energy": np.zeros(1), "fid": 1},
-                "inconsistent fields",
+                "atom count",
+                {**first_frame, "coord": np.zeros((7, 3)), "fid": 1},
+                "the batch layout gives frame 1 6 atoms, but it holds 7",
             ),
             (
                 "shape",
-                {**first_frame, "coord": np.zeros((7, 3)), "fid": 1},
+                {**first_frame, "coord": np.zeros((6, 4)), "fid": 1},
                 "changes shape within one batch",
             ),
         )
@@ -529,33 +641,35 @@ def test_parallel_batch_consistency_guards(self):
             ):
                 decode_lmdb_batch(transaction, [0, 1], "012d", config)
 
+    def test_merged_chunks_demote_partial_labels(self):
+        """Merging keeps only what every worker chunk supplies."""
         first_chunk = {
             "find_energy": np.float32(1.0),
             "coord": np.zeros((1, 6, 3)),
             "fid": [0],
             "sid": np.array([0], dtype=np.int64),
         }
-        chunk_cases = (
-            (
-                "fields",
-                {key: value for key, value in first_chunk.items() if key != "coord"},
-                "inconsistent fields",
-            ),
-            (
-                "availability",
-                {**first_chunk, "find_energy": np.float32(0.0), "fid": [1]},
-                "availability changes across worker chunks",
-            ),
+
+        without_coord = {
+            key: value for key, value in first_chunk.items() if key != "coord"
+        }
+        merged = _merge_lmdb_chunks([first_chunk, {**without_coord, "fid": [1]}])
+        self.assertNotIn("coord", merged)
+
+        merged = _merge_lmdb_chunks(
+            [first_chunk, {**first_chunk, "find_energy": np.float32(0.0), "fid": [1]}]
         )
-        for guard, second_chunk, expected_error in chunk_cases:
-            with (
-                self.subTest(guard=f"chunk {guard}"),
-                self.assertRaisesRegex(ValueError, expected_error),
-            ):
-                _merge_lmdb_chunks([first_chunk, second_chunk])
+        self.assertEqual(float(merged["find_energy"]), 0.0)
+        self.assertEqual(merged["coord"].shape, (2, 6, 3))
+
+    def test_batch_demotes_mixed_label_availability(self):
+        """A label only some frames carry is unavailable to the whole batch.
 
-    def test_batch_rejects_mixed_label_availability(self):
-        """A scalar find flag cannot represent mixed availability in one batch."""
+        A scalar ``find_*`` flag cannot say "available for frame 0 only", so
+        the batch reports the label absent and falls back to the requirement's
+        default fill. That keeps the fill out of the loss, which is what a
+        partial label would otherwise be mistaken for.
+        """
         path = _create_lmdb(
             f"{self._tmpdir.name}/mixed_availability.lmdb",
             nframes=2,
@@ -584,18 +698,492 @@ def test_batch_rejects_mixed_label_availability(self):
         )
         environment = lmdb.open(path, readonly=True, lock=False)
         with environment.begin() as transaction:
-            with self.assertRaisesRegex(
-                ValueError,
-                "availability changes within one batch",
-            ):
-                decode_lmdb_batch(
-                    transaction,
-                    [0, 1],
-                    "012d",
+            batch = decode_lmdb_batch(transaction, [0, 1], "012d", config)
+            single = decode_lmdb_batch(transaction, [0], "012d", config)
+            frames = [
+                decode_lmdb_frame(
+                    transaction.get(format(index, "012d").encode()),
+                    index,
                     config,
+                    copy_arrays=True,
+                )
+                for index in (0, 1)
+            ]
+        environment.close()
+
+        self.assertEqual(float(batch["find_custom"]), 0.0)
+        # The label survives in a batch where every frame carries it.
+        self.assertEqual(float(single["find_custom"]), 1.0)
+
+        # The generic collation path must reach the same verdict, so that a
+        # batch is the same mapping whichever path decoded it.
+        collated = collate_lmdb_frames(frames)
+        self.assertEqual(float(collated["find_custom"]), 0.0)
+        self.assertEqual(sorted(collated), sorted(batch))
+        np.testing.assert_allclose(collated["custom"], batch["custom"])
+
+    def test_availability_scan_normalizes_atomic_label_prefix(self):
+        """Availability and frame decoding share the atomic-label aliases."""
+        raw = msgpack.packb(
+            {
+                "atomic_dipole": np.ones(3).tolist(),
+                "find_atomic_dipole": 1.0,
+            },
+            use_bin_type=True,
+        )
+
+        availability = lmdb_data_module._raw_frame_availability(raw, {"atom_dipole": 1})
+
+        self.assertEqual(availability, 1)
+
+    def test_availability_probe_is_deferred_until_requirements_are_registered(self):
+        """Reader construction must not inspect labels before their use is known.
+
+        The trainer registers loss and model requirements after constructing
+        the data system. Inspecting raw labels earlier both performs needless
+        I/O and lets fields unused by the run influence the partition.
+        """
+        with mock.patch.object(
+            lmdb_data_module,
+            "_raw_frame_availability",
+            wraps=lmdb_data_module._raw_frame_availability,
+        ) as probe:
+            reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
+            self.assertEqual(probe.call_count, 0)
+            reader.add_data_requirement(
+                [
+                    DataRequirementItem("energy", 1, atomic=False),
+                    DataRequirementItem("force", 3, atomic=True),
+                ]
+            )
+            groups = collect_lmdb_sampling_groups(reader)
+            batches = list(LmdbBatchSampler(reader, shuffle=False))
+
+        self.assertLessEqual(
+            probe.call_count, lmdb_data_module._AVAILABILITY_PROBE_FRAMES
+        )
+        self.assertEqual(
+            sorted(index for _nloc, indices in groups for index in indices),
+            list(range(len(reader))),
+        )
+        self.assertEqual(
+            sorted(index for batch in batches for index in batch),
+            list(range(len(reader))),
+        )
+
+    def test_availability_probe_samples_a_bounded_number_of_frames(self):
+        """The construction-time probe must not scale with the dataset."""
+        nframes = lmdb_data_module._AVAILABILITY_PROBE_FRAMES * 3
+        path = _create_lmdb(
+            f"{self._tmpdir.name}/probe_bound.lmdb", nframes=nframes, natoms=6
+        )
+        with mock.patch.object(
+            lmdb_data_module,
+            "_raw_frame_availability",
+            wraps=lmdb_data_module._raw_frame_availability,
+        ) as probe:
+            reader = LmdbDataReader(path, self._type_map, batch_size=2)
+            self.assertEqual(probe.call_count, 0)
+            reader.add_data_requirement(
+                [DataRequirementItem("energy", 1, atomic=False)]
+            )
+            collect_lmdb_sampling_groups(reader)
+        self.assertEqual(len(reader), nframes)
+        self.assertLessEqual(
+            probe.call_count, lmdb_data_module._AVAILABILITY_PROBE_FRAMES
+        )
+
+    def test_valid_default_does_not_trigger_exact_partition(self):
+        """Present values and valid defaults may share one batch safely."""
+        nframes = lmdb_data_module._AVAILABILITY_PROBE_FRAMES * 3
+        path = _create_lmdb(
+            f"{self._tmpdir.name}/valid_default.lmdb",
+            nframes=nframes,
+            natoms=6,
+        )
+        environment = lmdb.open(path, readonly=False, lock=False)
+        with environment.begin(write=True) as transaction:
+            for index in range(0, nframes, 2):
+                key = format(index, "012d").encode()
+                frame = msgpack.unpackb(transaction.get(key), raw=False)
+                frame["atom_pref"] = {
+                    "type": "<f8",
+                    "shape": (6,),
+                    "data": np.full(6, 2.0, dtype=np.float64).tobytes(),
+                }
+                frame["fparam"] = {
+                    "type": "<f8",
+                    "shape": (2,),
+                    "data": np.array([2.0, 3.0], dtype=np.float64).tobytes(),
+                }
+                frame["charge_spin"] = {
+                    "type": "<f8",
+                    "shape": (2,),
+                    "data": np.array([4.0, 5.0], dtype=np.float64).tobytes(),
+                }
+                if index == 2:
+                    frame["find_atom_pref"] = 0.0
+                transaction.put(key, msgpack.packb(frame, use_bin_type=True))
+        environment.close()
+
+        requirements = [
+            DataRequirementItem("energy", 1, atomic=False),
+            DataRequirementItem("force", 3, atomic=True),
+            DataRequirementItem(
+                "atom_pref",
+                1,
+                atomic=True,
+                repeat=3,
+                default=1.0,
+                source_policy="default",
+            ),
+            DataRequirementItem(
+                "fparam",
+                2,
+                atomic=False,
+                default=np.array([0.0, 1.0]),
+                source_policy="default",
+            ),
+            DataRequirementItem(
+                "charge_spin",
+                2,
+                atomic=False,
+                default=np.array([0.0, 1.0]),
+                source_policy="default",
+            ),
+        ]
+        with mock.patch.object(
+            lmdb_data_module,
+            "_raw_frame_availability",
+            wraps=lmdb_data_module._raw_frame_availability,
+        ) as inspect_availability:
+            reader = LmdbDataReader(path, self._type_map, batch_size=2)
+            self.assertEqual(inspect_availability.call_count, 0)
+            reader.add_data_requirement(requirements)
+            groups = collect_lmdb_sampling_groups(reader)
+            batch = reader.decode_batch([0, 1, 2])
+
+        self.assertLessEqual(
+            inspect_availability.call_count,
+            lmdb_data_module._AVAILABILITY_PROBE_FRAMES,
+        )
+        self.assertEqual(len(groups), 1)
+        np.testing.assert_array_equal(groups[0][1], np.arange(nframes))
+        np.testing.assert_array_equal(batch["atom_pref"][0], np.full(18, 2.0))
+        np.testing.assert_array_equal(batch["atom_pref"][1], np.ones(18))
+        np.testing.assert_array_equal(batch["atom_pref"][2], np.ones(18))
+        np.testing.assert_array_equal(
+            batch["fparam"], [[2.0, 3.0], [0.0, 1.0], [2.0, 3.0]]
+        )
+        np.testing.assert_array_equal(
+            batch["charge_spin"],
+            [[4.0, 5.0], [0.0, 1.0], [4.0, 5.0]],
+        )
+        self.assertEqual(float(batch["find_fparam"]), 1.0)
+        self.assertEqual(float(batch["find_charge_spin"]), 1.0)
+
+        test_data = LmdbTestData(path, type_map=self._type_map, shuffle_test=False)
+        test_data.add_data_requirement(requirements)
+        frame_indices = next(iter(test_data.nloc_groups.values()))
+        self.assertEqual(len(test_data.availability_groups(frame_indices)), 1)
+        stacked = test_data.get_test_by_indices([0, 1, 2])
+        np.testing.assert_array_equal(stacked["atom_pref"][0], np.full(18, 2.0))
+        np.testing.assert_array_equal(stacked["atom_pref"][1], np.ones(18))
+        np.testing.assert_array_equal(stacked["atom_pref"][2], np.ones(18))
+        np.testing.assert_array_equal(
+            stacked["fparam"],
+            [[2.0, 3.0], [0.0, 1.0], [2.0, 3.0]],
+        )
+        np.testing.assert_array_equal(
+            stacked["charge_spin"],
+            [[4.0, 5.0], [0.0, 1.0], [4.0, 5.0]],
+        )
+        self.assertEqual(float(stacked["find_fparam"]), 1.0)
+        self.assertEqual(float(stacked["find_charge_spin"]), 1.0)
+
+    def test_mixed_dataset_is_partitioned_by_label_availability(self):
+        """A dataset the probe finds mixed still groups by availability.
+
+        Demoting a mixed batch keeps it correct but costs the label for
+        training, so a dataset that genuinely interleaves labels is worth
+        partitioning exactly.
+        """
+        path = f"{self._tmpdir.name}/interleaved.lmdb"
+        environment = lmdb.open(path, map_size=10 * 1024 * 1024)
+        nframes = 8
+        with environment.begin(write=True) as transaction:
+            transaction.put(
+                b"__metadata__",
+                msgpack.packb(
+                    {
+                        "nframes": nframes,
+                        "frame_idx_fmt": "012d",
+                        "system_info": {"natoms": [2, 4]},
+                        "frame_nlocs": [6] * nframes,
+                    },
+                    use_bin_type=True,
+                ),
+            )
+            for index in range(nframes):
+                frame = _make_frame(natoms=6, seed=index)
+                # Alternate frames drop the forces, so exactly half the
+                # dataset supplies that label.
+                if index % 2:
+                    frame.pop("forces")
+                    if index == 1:
+                        frame["find_force"] = 1.0
+                transaction.put(
+                    format(index, "012d").encode(),
+                    msgpack.packb(frame, use_bin_type=True),
                 )
         environment.close()
 
+        reader = LmdbDataReader(path, self._type_map, batch_size=2)
+        sampler = LmdbBatchSampler(reader, shuffle=False)
+        sampler.batches()
+        reader.add_data_requirement(
+            [DataRequirementItem("force", 3, atomic=True, must=False)]
+        )
+        opens = []
+        real_open = lmdb.open
+
+        def recording_open(path, **kwargs):
+            opens.append(kwargs.get("readahead"))
+            return real_open(path, **kwargs)
+
+        with mock.patch.object(lmdb_data_module.lmdb, "open", recording_open):
+            with mock.patch.object(
+                lmdb_data_module,
+                "_raw_frame_availability",
+                wraps=lmdb_data_module._raw_frame_availability,
+            ) as inspect_availability:
+                with self.assertLogs(lmdb_data_module.log, level="INFO") as captured:
+                    groups = collect_lmdb_sampling_groups(reader)
+                exact_scan_reads = inspect_availability.call_count
+                batches = list(sampler)
+                self.assertEqual(inspect_availability.call_count, exact_scan_reads)
+        self.assertGreaterEqual(exact_scan_reads, nframes)
+        self.assertLessEqual(
+            exact_scan_reads,
+            nframes + lmdb_data_module._AVAILABILITY_PROBE_FRAMES,
+        )
+        self.assertFalse(reader._uniform_availability)
+        self.assertIsNotNone(reader._availability_index)
+        self.assertEqual(reader._availability_index.ids.dtype, np.uint8)
+        self.assertEqual(reader._availability_index.ids.nbytes, nframes)
+        self.assertEqual(reader._availability_index.signature_count, 2)
+        self.assertEqual(opens, [True, False])
+        self.assertFalse(reader._env.flags()["readahead"])
+        self.assertTrue(
+            any(
+                "label-availability scan started" in message
+                and "frames=8" in message
+                and "force" in message
+                for message in captured.output
+            )
+        )
+        self.assertEqual(
+            sorted(sorted(indices) for _nloc, indices in groups),
+            [[0, 2, 4, 6], [1, 3, 5, 7]],
+        )
+        # Every batch therefore agrees on the label, and none is demoted.
+        for batch in batches:
+            self.assertEqual(len({index % 2 for index in batch}), 1)
+        reader.close()
+
+        test_data = LmdbTestData(path, type_map=self._type_map, shuffle_test=False)
+        test_data.add_data_requirement(
+            [DataRequirementItem("force", 3, atomic=True, must=False)]
+        )
+        frame_indices = next(iter(test_data.nloc_groups.values()))
+        with mock.patch.object(
+            lmdb_data_module,
+            "_raw_frame_availability",
+            wraps=lmdb_data_module._raw_frame_availability,
+        ) as inspect_test_availability:
+            first_groups = test_data.availability_groups(frame_indices)
+            first_scan_reads = inspect_test_availability.call_count
+            second_groups = test_data.availability_groups(frame_indices)
+            self.assertEqual(
+                inspect_test_availability.call_count,
+                first_scan_reads,
+            )
+        self.assertGreaterEqual(first_scan_reads, nframes)
+        self.assertEqual(
+            [indices.tolist() for indices in first_groups],
+            [indices.tolist() for indices in second_groups],
+        )
+
+    def test_failed_exact_scan_restores_random_read_environment(self):
+        """A failed sequential scan leaves the reader usable for training I/O."""
+        path = _create_lmdb(
+            f"{self._tmpdir.name}/failed_exact_scan.lmdb",
+            nframes=2,
+            natoms=6,
+        )
+        environment = lmdb.open(path, readonly=False, lock=False)
+        with environment.begin(write=True) as transaction:
+            key = format(1, "012d").encode()
+            frame = msgpack.unpackb(transaction.get(key), raw=False)
+            frame.pop("forces")
+            transaction.put(key, msgpack.packb(frame, use_bin_type=True))
+        environment.close()
+
+        reader = LmdbDataReader(path, self._type_map, batch_size=1)
+        reader.add_data_requirement(
+            [DataRequirementItem("force", 3, atomic=True, must=False)]
+        )
+        with mock.patch.object(
+            lmdb_data_module,
+            "_scan_availability_index",
+            side_effect=RuntimeError("scan failed"),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "scan failed"):
+                reader.availability_groups(np.arange(len(reader)))
+
+        self.assertIsNotNone(reader._env)
+        self.assertFalse(reader._env.flags()["readahead"])
+        self.assertEqual(reader[0]["coord"].shape, (6, 3))
+
+    def test_exact_scan_logs_only_on_node_local_rank_zero(self):
+        """Nonzero local ranks suppress duplicate node-local scan progress."""
+        raw = msgpack.packb(_make_frame(natoms=6, seed=0), use_bin_type=True)
+        with mock.patch.dict(os.environ, {"LOCAL_RANK": "1"}):
+            with mock.patch.object(lmdb_data_module.log, "info") as log_info:
+                index = lmdb_data_module._scan_availability_index(
+                    1,
+                    lambda _position: raw,
+                    ("force",),
+                    "ranked.lmdb",
+                )
+
+        self.assertEqual(index.signature_count, 1)
+        log_info.assert_not_called()
+
+    def test_metadata_is_read_with_readahead_enabled(self):
+        """``__metadata__`` must not be read under ``MDB_NORDAHEAD``.
+
+        It is one contiguous run of overflow pages that grows with the frame
+        count. Reading it with readahead disabled costs one synchronous fault
+        per 4 KiB page, measured at 1.8 MiB/s against NFS against 81 MiB/s
+        with readahead left on -- minutes rather than seconds once a writer
+        records a per-frame table.
+        """
+        opens = []
+        real_open = lmdb.open
+
+        def recording_open(path, **kwargs):
+            opens.append(kwargs.get("readahead"))
+            return real_open(path, **kwargs)
+
+        with mock.patch.object(lmdb_data_module.lmdb, "open", recording_open):
+            reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
+
+        self.assertEqual(len(reader), 10)
+        # The first open serves the metadata and keeps readahead; the second
+        # serves frames and drops it so random reads stay precise.
+        self.assertEqual(opens, [True, False])
+
+    def test_readahead_follows_the_access_pattern(self):
+        """Each reader asks for the readahead its access pattern wants.
+
+        Training draws frames in shuffled order, where readahead fetches
+        neighbours it will not use. ``LmdbTestData`` walks a group in
+        ascending key order, where the same setting turns many small faults
+        into few large reads -- 45 000 frames/s against 1 300 over NFS.
+        """
+        opens = []
+        real_open = lmdb.open
+
+        def recording_open(path, **kwargs):
+            opens.append(kwargs.get("readahead"))
+            return real_open(path, **kwargs)
+
+        with mock.patch.object(lmdb_data_module.lmdb, "open", recording_open):
+            LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
+            # Metadata first with readahead, then the shuffled-order frames
+            # without it.
+            self.assertEqual(opens, [True, False])
+
+            opens.clear()
+            LmdbTestData(self._lmdb_path, type_map=self._type_map)
+            self.assertEqual(opens, [True, True])
+
+    def test_frame_tables_are_arrays_not_python_lists(self):
+        """The per-frame tables must not be Python lists.
+
+        A list of boxed integers costs about 36 bytes an entry, so the two
+        tables alone reach tens of gigabytes on a dataset of 10^8 frames.
+        """
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
+        self.assertIsInstance(reader.frame_nlocs, np.ndarray)
+        self.assertIsInstance(reader._retained_keys, np.ndarray)
+        for indices in reader.nloc_groups.values():
+            self.assertIsInstance(indices, np.ndarray)
+
+    def test_sampling_does_not_disturb_the_reader_index_tables(self):
+        """Shuffling a pass must not permute the reader's own groups.
+
+        The sampler receives the reader's index arrays directly, so an
+        in-place shuffle would silently corrupt every later pass.
+        """
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
+        before = {nloc: idx.copy() for nloc, idx in reader.nloc_groups.items()}
+        for _ in range(2):
+            list(LmdbBatchSampler(reader, shuffle=True, seed=5))
+        for nloc, indices in before.items():
+            np.testing.assert_array_equal(reader.nloc_groups[nloc], indices)
+
+    def test_availability_groups_skip_an_empty_group(self):
+        """An empty group names no frames to stack, so it is not a group.
+
+        ``LmdbTestData`` retains an empty group when ``max_frames`` rounds a
+        group down to nothing, and full validation stacks whatever group it
+        is handed.
+        """
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
+        self.assertEqual(reader.availability_groups(np.array([], dtype=np.int64)), [])
+
+        empty = LmdbTestData(
+            self._lmdb_path, type_map=self._type_map, max_frames=0, shuffle_test=False
+        )
+        for indices in empty.nloc_groups.values():
+            self.assertEqual(empty.availability_groups(indices), [])
+
+    def test_test_data_releases_only_a_reference_it_took(self):
+        """A failed construction must not close another reader's environment."""
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
+
+        with mock.patch.object(
+            lmdb_data_module,
+            "_read_metadata_of",
+            side_effect=RuntimeError("metadata unreadable"),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "metadata unreadable"):
+                LmdbTestData(self._lmdb_path, type_map=self._type_map)
+        gc.collect()
+
+        # The reader that does hold the environment still reads through it.
+        self.assertIn("coord", reader[0])
+
+    def test_view_iterates_the_frames_it_stacks(self):
+        """``iter_test`` and ``get_test`` must serve one view the same frames."""
+        data = LmdbTestData(
+            self._lmdb_path, type_map=self._type_map, shuffle_test=False
+        )
+        data.add("energy", 1, atomic=False, must=False, high_prec=True)
+        nloc = next(iter(data.nloc_groups))
+        subset = data.nloc_groups[nloc][:3]
+        view = LmdbTestDataNlocView(data, nloc, subset)
+
+        whole = view.get_test()
+        chunked = list(view.iter_test(chunk_atoms=nloc))
+        self.assertEqual(whole["coord"].shape[0], len(subset))
+        self.assertEqual(sum(c["coord"].shape[0] for c in chunked), len(subset))
+        np.testing.assert_allclose(
+            np.concatenate([c["coord"] for c in chunked]), whole["coord"]
+        )
+
     def test_is_lmdb(self):
         self.assertTrue(is_lmdb(self._lmdb_path))
         self.assertTrue(is_lmdb("something.lmdb"))
@@ -625,6 +1213,7 @@ def test_min_pair_dist_requirement_computed(self):
                     atomic=False,
                     must=False,
                     high_prec=False,
+                    source_policy="derived",
                 )
             ]
         )
@@ -641,6 +1230,7 @@ def test_min_pair_dist_requirement_defaults_without_atype(self):
             "min_pair_dist",
             ndof=1,
             default=0.25,
+            source_policy="derived",
         )
         config = LmdbDecodeConfig(
             ntypes=2,
@@ -659,6 +1249,47 @@ def test_min_pair_dist_requirement_defaults_without_atype(self):
         self.assertEqual(frame["find_min_pair_dist"], np.float32(0.0))
         np.testing.assert_allclose(frame["min_pair_dist"], np.array([0.25]))
 
+    def test_derived_min_pair_dist_ignores_raw_presence_for_grouping(self):
+        """Derived fields neither partition frames nor trust stored values."""
+        path = _create_grid_lmdb(
+            f"{self._tmpdir.name}/derived_min_pair.lmdb",
+            nframes=2,
+        )
+        environment = lmdb.open(path, readonly=False, lock=False)
+        with environment.begin(write=True) as transaction:
+            key = format(0, "012d").encode()
+            frame = msgpack.unpackb(transaction.get(key), raw=False)
+            frame["min_pair_dist"] = {
+                "type": "<f8",
+                "shape": (1,),
+                "data": np.array([99.0], dtype=np.float64).tobytes(),
+            }
+            transaction.put(key, msgpack.packb(frame, use_bin_type=True))
+        environment.close()
+
+        reader = LmdbDataReader(path, ["TYPE"], batch_size=1)
+        reader.add_data_requirement(
+            [
+                DataRequirementItem(
+                    "min_pair_dist",
+                    ndof=1,
+                    atomic=False,
+                    default=0.0,
+                    source_policy="derived",
+                )
+            ]
+        )
+        with mock.patch.object(
+            lmdb_data_module,
+            "_raw_frame_availability",
+            side_effect=AssertionError("derived fields must not probe raw presence"),
+        ):
+            groups = collect_lmdb_sampling_groups(reader)
+
+        self.assertEqual(len(groups), 1)
+        np.testing.assert_allclose(reader[0]["min_pair_dist"], [1.0])
+        np.testing.assert_allclose(reader[1]["min_pair_dist"], [1.0])
+
 
 # ============================================================
 # Mixed nloc tests
@@ -666,7 +1297,7 @@ def test_min_pair_dist_requirement_defaults_without_atype(self):
 
 
 class TestMixedNloc(unittest.TestCase):
-    """Tests for mixed-nloc datasets and SameNlocBatchSampler."""
+    """Tests for mixed-nloc datasets and LmdbBatchSampler."""
 
     @classmethod
     def setUpClass(cls):
@@ -702,18 +1333,20 @@ def test_per_frame_shapes(self):
 
     def test_frame_nlocs(self):
         reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
-        self.assertEqual(reader.frame_nlocs, [6, 6, 6, 6, 9, 9, 9, 9, 12, 12])
+        np.testing.assert_array_equal(
+            reader.frame_nlocs, [6, 6, 6, 6, 9, 9, 9, 9, 12, 12]
+        )
 
     def test_sampler_all_batches_same_nloc(self):
         reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
-        sampler = SameNlocBatchSampler(reader, shuffle=False, seed=42)
+        sampler = LmdbBatchSampler(reader, shuffle=False, seed=42)
         for batch_indices in sampler:
             nlocs = [reader.frame_nlocs[i] for i in batch_indices]
             self.assertTrue(all(n == nlocs[0] for n in nlocs))
 
     def test_sampler_covers_all_frames(self):
         reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
-        sampler = SameNlocBatchSampler(reader, shuffle=False, seed=42)
+        sampler = LmdbBatchSampler(reader, shuffle=False, seed=42)
         all_indices = [i for batch in sampler for i in batch]
         self.assertEqual(sorted(all_indices), list(range(10)))
 
@@ -725,15 +1358,413 @@ def test_sampler_auto_batch_size_per_nloc(self):
 
     def test_sampler_shuffle_deterministic(self):
         reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
-        s1 = SameNlocBatchSampler(reader, shuffle=True, seed=123)
-        s2 = SameNlocBatchSampler(reader, shuffle=True, seed=123)
+        s1 = LmdbBatchSampler(reader, shuffle=True, seed=123)
+        s2 = LmdbBatchSampler(reader, shuffle=True, seed=123)
+        # A seed reproduces the whole sequence of passes, not just the first.
+        self.assertEqual(list(s1), list(s2))
         self.assertEqual(list(s1), list(s2))
 
+    def test_sampler_reshuffles_between_passes(self):
+        """Every pass draws a fresh shuffle while still covering each frame once.
+
+        Training re-iterates the sampler whenever it exhausts it, so a pass
+        that repeated the previous one would show the model the same batches
+        for the whole run.
+        """
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
+        sampler = LmdbBatchSampler(reader, shuffle=True, seed=5)
+        first = [list(batch) for batch in sampler]
+        second = [list(batch) for batch in sampler]
+        self.assertNotEqual(first, second)
+        self.assertEqual(
+            sorted(index for batch in first for index in batch),
+            sorted(index for batch in second for index in batch),
+        )
+
     def test_sampler_len(self):
         reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=2)
-        sampler = SameNlocBatchSampler(reader, shuffle=False)
+        sampler = LmdbBatchSampler(reader, shuffle=False)
+        self.assertEqual(len(sampler), len(list(sampler)))
+
+    def test_mix_batches_span_several_nlocs(self):
+        """``mix:N`` puts frames of different atom counts in one batch.
+
+        Frames are sorted by atom count before batches are filled, so a batch
+        spans several counts exactly where a budget boundary falls inside the
+        sorted run. The budget here is chosen to land on such a boundary.
+        """
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:18")
+        self.assertTrue(reader.mixed_nloc)
+        batches = list(LmdbBatchSampler(reader, shuffle=False))
+        spanned = max(len({reader.frame_nlocs[i] for i in batch}) for batch in batches)
+        self.assertGreater(spanned, 1)
+
+    def test_mix_merges_small_neighbouring_atom_count_groups(self):
+        """Nearby atom counts share a batch instead of one tiny batch each.
+
+        This is what ``mix:N`` buys over ``max:N``: an atom-count group too
+        small to fill a batch on its own no longer forces an under-filled
+        step. The gain needs the counts to be close, since padding costs
+        ``nframes * max_nloc``; a lone outlier size still batches alone.
+        """
+        path = f"{self._tmpdir.name}/graded.lmdb"
+        _create_mixed_sid_nloc_lmdb(
+            path,
+            system_specs=[(2, 6), (2, 7), (2, 8)],
+            type_map=self._type_map,
+        )
+        costs = {}
+        for spec in ("max:72", "mix:72"):
+            reader = LmdbDataReader(path, self._type_map, batch_size=spec)
+            costs[spec] = [
+                len(b) * reader.batch_pad_nloc(b)
+                for b in LmdbBatchSampler(reader, shuffle=False)
+            ]
+        # One under-filled batch per atom-count group, against a single full one.
+        self.assertEqual(len(costs["max:72"]), 3)
+        self.assertEqual(costs["mix:72"], [6 * 8])
+
+    def test_mix_respects_the_atom_budget(self):
+        """A batch's padded cost stays within the budget above one frame."""
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36")
+        for batch in LmdbBatchSampler(reader, shuffle=True, seed=3):
+            cost = len(batch) * reader.batch_pad_nloc(batch)
+            self.assertTrue(len(batch) == 1 or cost <= 36, (len(batch), cost))
+
+    def test_mix_covers_every_frame_once(self):
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36")
+        indices = [
+            i for batch in LmdbBatchSampler(reader, shuffle=True, seed=5) for i in batch
+        ]
+        self.assertEqual(sorted(indices), list(range(len(reader))))
+
+    def test_mix_len_matches_iteration(self):
+        """The batch count depends on the shuffle, so ``len`` must reuse it."""
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36")
+        sampler = LmdbBatchSampler(reader, shuffle=True, seed=9)
         self.assertEqual(len(sampler), len(list(sampler)))
 
+    def test_mix_shuffle_deterministic(self):
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36")
+        s1 = LmdbBatchSampler(reader, shuffle=True, seed=17)
+        s2 = LmdbBatchSampler(reader, shuffle=True, seed=17)
+        self.assertEqual(list(s1), list(s2))
+
+    def test_mix_pads_the_atom_axis_only(self):
+        """Padding widens per-atom fields and leaves frame-level ones alone."""
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36")
+        reader.add_data_requirement(
+            [
+                DataRequirementItem("energy", 1, atomic=False, must=False),
+                DataRequirementItem("force", 3, atomic=True, must=False),
+            ]
+        )
+        # Frames 0 and 8 hold 6 and 12 atoms respectively.
+        batch = reader.decode_batch([0, 8])
+        self.assertEqual(batch["coord"].shape, (2, 12, 3))
+        self.assertEqual(batch["atype"].shape, (2, 12))
+        self.assertEqual(batch["force"].shape, (2, 12, 3))
+        self.assertEqual(batch["box"].shape, (2, 9))
+        self.assertEqual(batch["energy"].shape, (2, 1))
+
+        # The short frame keeps its six atoms; the tail is inert.
+        np.testing.assert_array_equal(batch["atype"][0, 6:], -1)
+        self.assertTrue(np.all(batch["atype"][0, :6] >= 0))
+        np.testing.assert_array_equal(batch["coord"][0, 6:], 0.0)
+        np.testing.assert_array_equal(batch["force"][0, 6:], 0.0)
+        np.testing.assert_array_equal(batch["atype"][1], reader[8]["atype"])
+
+        # Atom counts stay the real per-frame values, not the padded width.
+        self.assertEqual(batch["natoms"][0, 0], 6)
+        self.assertEqual(batch["natoms"][1, 0], 12)
+
+    def test_mix_pads_repeated_per_atom_fields(self):
+        """A ``repeat != 1`` requirement is stored flat but still padded by atom.
+
+        ``atom_pref`` spends ``repeat`` leading entries per atom instead of
+        one, so its padded width is ``pad_nloc * repeat``.
+        """
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36")
+        reader.add_data_requirement(
+            [DataRequirementItem("atom_pref", 1, atomic=True, must=False, repeat=3)]
+        )
+        # Frames 0 and 8 hold 6 and 12 atoms respectively.
+        batch = reader.decode_batch([0, 8])
+        self.assertEqual(batch["atom_pref"].shape, (2, 12 * 3))
+        np.testing.assert_array_equal(batch["atom_pref"][0, 6 * 3 :], 0.0)
+
+    def test_mix_chunked_decode_matches_in_process_decode(self):
+        """Splitting a padded batch across workers reproduces it exactly.
+
+        Each chunk pads and lays out its fields independently, so a mixed-nloc
+        batch is the case where the chunks can disagree: they start on frames
+        of different atom counts.
+        """
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36")
+        reader.add_data_requirement(
+            [
+                DataRequirementItem("force", 3, atomic=True, must=False),
+                DataRequirementItem("atom_pref", 1, atomic=True, must=False, repeat=3),
+            ]
+        )
+        # Frames 0 and 8 hold 6 and 12 atoms; one per chunk puts the wide
+        # frame second, so a per-chunk layout would size the chunks differently.
+        indices = [0, 8]
+        layout = reader.batch_layout(indices)
+        chunks = [
+            lmdb_data_module.decode_lmdb_batch(
+                reader._transaction(),
+                [key],
+                reader.frame_format,
+                reader._decode_config,
+                layout.chunk(position, position + 1),
+            )
+            for position, key in enumerate(reader.original_keys(indices))
+        ]
+        merged = lmdb_data_module._merge_lmdb_chunks(chunks)
+        expected = reader.decode_batch(indices)
+        self.assertEqual(sorted(merged), sorted(expected))
+        for field, value in expected.items():
+            np.testing.assert_array_equal(merged[field], value, err_msg=field)
+
+    def test_mix_keeps_frame_fields_of_ambiguous_shape(self):
+        """``fparam`` and ``virial`` are never padded, whatever the atom count."""
+        path = _create_mix_probe_lmdb(f"{self._tmpdir.name}/probe.lmdb")
+        reader = LmdbDataReader(path, self._type_map, batch_size="mix:36")
+        reader.add_data_requirement(
+            [
+                DataRequirementItem("virial", 9, atomic=False, must=False),
+                DataRequirementItem("fparam", 2, atomic=False, must=False),
+                DataRequirementItem("force", 3, atomic=True, must=False),
+            ]
+        )
+        # A 2-atom frame makes fparam (2,) ambiguous; a 9-atom one does the
+        # same for virial (9,). Batching them together exercises both.
+        batch = reader.decode_batch([0, 1])
+        self.assertEqual(batch["fparam"].shape, (2, 2))
+        self.assertEqual(batch["virial"].shape, (2, 9))
+        self.assertEqual(batch["coord"].shape, (2, 9, 3))
+        np.testing.assert_array_equal(batch["fparam"][0], reader[0]["fparam"])
+        np.testing.assert_array_equal(batch["virial"][0], reader[0]["virial"])
+
+    def test_mix_tracks_the_atom_proportional_weighting(self):
+        """Packing to a budget must not distort how much each frame counts.
+
+        A batch is one optimizer step, and its frame-level terms (energy,
+        virial) average over the frames present, so a frame's weight in them
+        is ``1 / k_b``. An atom budget asks for ``k_b ~ budget / nloc``, that
+        is a weight proportional to atom count, matching what the pooled
+        per-atom terms give those frames unconditionally.
+        Same-nloc batching misses that ideal wherever an atom-count group is
+        too small to fill a batch: its few frames form a short batch and each
+        of them is over-weighted many times over. Filling batches across atom
+        counts removes that failure mode, and the residual error is the
+        padding, so both goals improve together.
+        """
+        path = f"{self._tmpdir.name}/skewed.lmdb"
+        # A few dominant atom counts plus a tail of groups too small to fill a
+        # batch, down to one holding a single frame. This is the shape of a
+        # real merged dataset in miniature, and the tail is where same-nloc
+        # batching goes wrong.
+        _create_mixed_sid_nloc_lmdb(
+            path,
+            system_specs=[
+                (120, 6),
+                (80, 8),
+                (60, 10),
+                (3, 12),
+                (2, 14),
+                (2, 16),
+                (1, 20),
+            ],
+            type_map=self._type_map,
+        )
+        budget = 120
+
+        def weighting_error(spec):
+            """Log-spread of the frame weight against the ideal, and its worst case."""
+            reader = LmdbDataReader(path, self._type_map, batch_size=spec)
+            nlocs = np.asarray(reader.frame_nlocs, dtype=np.float64)
+            batches = LmdbBatchSampler(reader, shuffle=True, seed=0).batches()
+            weight = np.empty(len(nlocs))
+            for batch in batches:
+                weight[batch] = 1.0 / len(batch)
+            slots = sum(len(b) * reader.batch_pad_nloc(b) for b in batches)
+            # Weights matter only up to a global scale, which is the learning
+            # rate, so normalize both sides to mean 1 before comparing.
+            ratio = (weight / weight.mean()) / (nlocs / nlocs.mean())
+            return np.std(np.log(ratio)), ratio.max(), nlocs.sum() / slots
+
+        same_spread, same_worst, same_efficiency = weighting_error(f"max:{budget}")
+        mix_spread, mix_worst, mix_efficiency = weighting_error(f"mix:{budget}")
+
+        self.assertEqual(same_efficiency, 1.0)
+        self.assertGreater(mix_efficiency, 0.95)
+        # Both the typical and the worst-case deviation shrink by a wide
+        # margin; the thresholds leave room for the packing to change.
+        self.assertLess(mix_spread, 0.5 * same_spread)
+        self.assertLess(mix_worst, 0.5 * same_worst)
+
+    def test_mix_cuts_a_batch_only_when_the_next_frame_does_not_fit(self):
+        """No batch is closed early, which is what keeps batches full.
+
+        Sorting a sliding window instead of the whole group would break this:
+        each window boundary closes a batch regardless of how full it is, and
+        an under-filled batch over-weights every frame it holds.
+        """
+        budget = 36
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:36")
+        # Chop the group directly, which keeps the batches in the atom-count
+        # order the greedy produced them in.
+        batches = _chop_mixed_nloc(reader, list(range(len(reader))))
+        for batch, following in pairwise(batches):
+            next_nloc = min(reader.frame_nlocs[index] for index in following)
+            self.assertGreater((len(batch) + 1) * next_nloc, budget)
+
+    def _ragged_reader(self, batch_size="mix:36"):
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size=batch_size)
+        reader.add_data_requirement(
+            [
+                DataRequirementItem("force", 3, atomic=True, must=False),
+                DataRequirementItem("energy", 1, atomic=False, must=False),
+            ]
+        )
+        reader.use_ragged_batches(True)
+        return reader
+
+    def test_ragged_layout_concatenates_the_frames(self):
+        """The ragged layout carries the same rows on a flat, unpadded axis."""
+        rectangular = LmdbDataReader(
+            self._lmdb_path, self._type_map, batch_size="mix:36"
+        )
+        rectangular.add_data_requirement(
+            [
+                DataRequirementItem("force", 3, atomic=True, must=False),
+                DataRequirementItem("energy", 1, atomic=False, must=False),
+            ]
+        )
+        ragged = self._ragged_reader()
+
+        # Frames 0 and 8 hold 6 and 12 atoms respectively.
+        indices = [0, 8]
+        padded = rectangular.decode_batch(indices)
+        flat = ragged.decode_batch(indices)
+
+        self.assertEqual(padded["coord"].shape, (2, 12, 3))
+        self.assertEqual(flat["coord"].shape, (18, 3))
+        self.assertEqual(flat["atype"].shape, (18,))
+        self.assertEqual(flat["force"].shape, (18, 3))
+        np.testing.assert_array_equal(flat["n_node"], [6, 12])
+        self.assertTrue((flat["atype"] >= 0).all(), "a ragged batch pads nothing")
+        # Frame-level fields keep their frame axis under either layout.
+        self.assertEqual(flat["energy"].shape, padded["energy"].shape)
+        self.assertEqual(flat["box"].shape, padded["box"].shape)
+
+        offset = 0
+        for row, count in enumerate(flat["n_node"].tolist()):
+            for field in ("coord", "atype", "force"):
+                np.testing.assert_array_equal(
+                    padded[field][row, :count],
+                    flat[field][offset : offset + count],
+                    err_msg=field,
+                )
+            offset += count
+
+    def test_ragged_chunked_decode_matches_in_process_decode(self):
+        """Chunks concatenate on the flat axis, so a split decode is the same."""
+        reader = self._ragged_reader()
+        indices = [0, 8]
+        layout = reader.batch_layout(indices)
+        chunks = [
+            lmdb_data_module.decode_lmdb_batch(
+                reader._transaction(),
+                [key],
+                reader.frame_format,
+                reader._decode_config,
+                layout.chunk(position, position + 1),
+            )
+            for position, key in enumerate(reader.original_keys(indices))
+        ]
+        merged = lmdb_data_module._merge_lmdb_chunks(chunks)
+        expected = reader.decode_batch(indices)
+        self.assertEqual(sorted(merged), sorted(expected))
+        for field, value in expected.items():
+            np.testing.assert_array_equal(merged[field], value, err_msg=field)
+
+    def test_ragged_packing_fills_on_real_atoms_and_keeps_the_caller_order(self):
+        """Without padding the budget counts real atoms, and sorting is dropped.
+
+        Sorting exists only to keep a rectangular batch's padded width close to
+        its frames. A ragged batch has no width, so sorting would merely make
+        each optimizer step homogeneous in system size.
+        """
+        budget = 36
+        ragged = self._ragged_reader(batch_size=f"mix:{budget}")
+        nlocs = np.asarray(ragged.frame_nlocs)
+        # The fixture stores its frames in ascending atom count, so hand them
+        # over reversed: only then does sorting leave a visible trace.
+        group = list(reversed(range(len(ragged))))
+        batches = _chop_mixed_nloc(ragged, group)
+
+        for batch in batches:
+            if len(batch) > 1:
+                self.assertLessEqual(int(nlocs[batch].sum()), budget)
+        self.assertEqual(sorted(index for b in batches for index in b), sorted(group))
+        # The frames arrive in the order they were handed over, not sorted.
+        self.assertEqual([index for batch in batches for index in batch], group)
+
+        rectangular = LmdbDataReader(
+            self._lmdb_path, self._type_map, batch_size=f"mix:{budget}"
+        )
+        padded = _chop_mixed_nloc(rectangular, group)
+        self.assertEqual(
+            [int(nlocs[index]) for batch in padded for index in batch],
+            sorted(nlocs[group].tolist()),
+            "the rectangular layout must still sort by atom count",
+        )
+
+    def test_mix_uses_the_fewest_batches_possible(self):
+        """The packing attains the minimum batch count, not merely a good one.
+
+        The minimum is computed by dynamic programming over the sorted atom
+        counts. Its correctness rests on being free to take every batch
+        contiguous in that order: the batch holding the widest frame may be
+        given the widest frames outright, since its capacity ``budget // nloc``
+        does not depend on which frames fill it. The recurrence then closes
+        each batch at its widest frame,
+
+            steps[i] = 1 + steps[max(0, i - budget // nloc[i - 1])],
+
+        whereas the implementation opens each batch at its narrowest. The two
+        must agree.
+        """
+        for budget in (18, 36, 72):
+            with self.subTest(budget=budget):
+                reader = LmdbDataReader(
+                    self._lmdb_path, self._type_map, batch_size=f"mix:{budget}"
+                )
+                nlocs = sorted(reader.frame_nlocs)
+                steps = [0] * (len(nlocs) + 1)
+                for i in range(1, len(nlocs) + 1):
+                    capacity = max(1, budget // nlocs[i - 1])
+                    steps[i] = 1 + steps[max(0, i - capacity)]
+                batches = _chop_mixed_nloc(reader, list(range(len(reader))))
+                self.assertEqual(len(batches), steps[len(nlocs)])
+
+    def test_mix_distributed_partition_is_disjoint_and_complete(self):
+        reader = LmdbDataReader(self._lmdb_path, self._type_map, batch_size="mix:24")
+        ranks = [
+            DistributedLmdbBatchSampler(
+                reader, rank=rank, world_size=2, shuffle=True, seed=4
+            )
+            for rank in (0, 1)
+        ]
+        batches = [list(sampler) for sampler in ranks]
+        for sampler, own in zip(ranks, batches, strict=True):
+            self.assertEqual(len(sampler), len(own))
+        indices = [i for own in batches for batch in own for i in batch]
+        self.assertEqual(sorted(indices), list(range(len(reader))))
+
     # --- LmdbTestData mixed-nloc tests ---
 
     def test_test_data_nloc_groups(self):
@@ -803,7 +1834,9 @@ def test_test_data_shuffle_uses_global_seed(self):
             shuffle_test=True,
         )
 
-        self.assertEqual(first.nloc_groups, second.nloc_groups)
+        self.assertEqual(set(first.nloc_groups), set(second.nloc_groups))
+        for nloc, first_indices in first.nloc_groups.items():
+            np.testing.assert_array_equal(first_indices, second.nloc_groups[nloc])
 
     def test_test_data_get_test_default_mixed(self):
         td = LmdbTestData(self._lmdb_path, type_map=self._type_map, shuffle_test=False)
@@ -941,7 +1974,7 @@ def test_positive_out_of_range_type_still_raises(self):
 
 class TestAutoProb(unittest.TestCase):
     """Test auto_prob support: frame_system_ids, compute_block_targets,
-    _expand_indices_by_blocks, and SameNlocBatchSampler with block_targets.
+    _expand_indices_by_blocks, and LmdbBatchSampler with block_targets.
     """
 
     @classmethod
@@ -1079,9 +2112,7 @@ def test_sampler_with_block_targets(self):
             nsystems=3,
             system_nframes=[100, 200, 300],
         )
-        sampler = SameNlocBatchSampler(
-            reader, shuffle=True, block_targets=block_targets
-        )
+        sampler = LmdbBatchSampler(reader, shuffle=True, block_targets=block_targets)
         all_indices = [i for batch in sampler for i in batch]
         self.assertGreater(len(all_indices), 600)
         self.assertEqual(len(set(all_indices)), 600)
@@ -1092,21 +2123,25 @@ def test_sampler_allocates_block_target_across_find_signatures(self):
         class TwoSignatureReader:
             """Minimal reader exposing two one-frame availability groups."""
 
+            mixed_nloc = False
+            data_requirements_revision = 0
+
             def __init__(self):
                 self.nloc_groups = {6: [0, 1]}
                 self.frame_system_ids = [0, 0]
+                self.frame_nlocs = [6, 6]
 
             @staticmethod
-            def group_indices_by_find_signature(indices):
+            def availability_groups(indices):
                 self.assertEqual(indices, [0, 1])
-                return {(0.0,): [0], (1.0,): [1]}
+                return [[0], [1]]
 
             @staticmethod
             def get_batch_size_for_nloc(nloc):
                 self.assertEqual(nloc, 6)
                 return 1
 
-        sampler = SameNlocBatchSampler(
+        sampler = LmdbBatchSampler(
             TwoSignatureReader(),
             shuffle=False,
             block_targets=[([0], 3)],
@@ -1118,9 +2153,29 @@ def get_batch_size_for_nloc(nloc):
         self.assertEqual(sum(map(len, batches)), 3)
         self.assertEqual(sorted(counts), [1, 2])
 
+    def test_sampler_tolerates_blocks_that_omit_a_system(self):
+        """A system no auto_prob block names is passed through, not rejected.
+
+        ``prob_sys_size`` slices name a prefix of the systems, so the tail is
+        routinely unnamed. Those frames belong to no block and are neither
+        expanded nor dropped.
+        """
+        reader = LmdbDataReader(self._lmdb_path, ["O", "H"])
+        block_targets = compute_block_targets(
+            "prob_sys_size;0:1:0.5;1:2:0.5",
+            nsystems=3,
+            system_nframes=[100, 200, 300],
+        )
+        sampler = LmdbBatchSampler(reader, shuffle=False, block_targets=block_targets)
+        indices = [index for batch in sampler for index in batch]
+        # Every frame of the unnamed third system survives exactly once.
+        third = [index for index in indices if reader.frame_system_ids[index] == 2]
+        self.assertEqual(sorted(third), sorted(set(third)))
+        self.assertEqual(len(third), 300)
+
     def test_sampler_without_block_targets(self):
         reader = LmdbDataReader(self._lmdb_path, ["O", "H"])
-        sampler = SameNlocBatchSampler(reader, shuffle=False)
+        sampler = LmdbBatchSampler(reader, shuffle=False)
         all_indices = [i for batch in sampler for i in batch]
         self.assertEqual(sorted(all_indices), list(range(600)))
 
@@ -1295,7 +2350,7 @@ def test_filter_dataset_index_is_contiguous_and_live(self):
         )
         self.assertEqual(len(reader), 8)
         self.assertEqual(len(reader._retained_keys), 8)
-        self.assertEqual(reader._retained_keys, [0, 1, 2, 3, 4, 5, 6, 7])
+        np.testing.assert_array_equal(reader._retained_keys, [0, 1, 2, 3, 4, 5, 6, 7])
 
         seen_fids = []
         for i in range(len(reader)):
@@ -1317,11 +2372,11 @@ def test_filter_dataset_index_is_contiguous_and_live(self):
             reader[-1]
 
     def test_sampler_with_filter(self):
-        """SameNlocBatchSampler only emits retained, same-nloc frames."""
+        """LmdbBatchSampler only emits retained, same-nloc frames."""
         reader = LmdbDataReader(
             self._mixed_path, self._type_map, batch_size="filter:10"
         )
-        sampler = SameNlocBatchSampler(reader, shuffle=False, seed=0)
+        sampler = LmdbBatchSampler(reader, shuffle=False, seed=0)
         all_batches = list(sampler)
         all_indices = [idx for batch in all_batches for idx in batch]
 
@@ -1354,22 +2409,6 @@ def test_invalid_batch_size_strings_rejected(self):
                 LmdbDataReader(self._uniform_path, self._type_map, batch_size=spec)
             self.assertIn("positive", str(ctx.exception))
 
-    def test_filter_with_mixed_batch_rejected(self):
-        """``filter:N`` + ``mixed_batch=True`` must fail loudly.
-
-        The mixed-batch fast path skips the per-frame nloc scan, so
-        filter:N cannot honour its documented ``nloc > N`` drop.
-        """
-        with self.assertRaises(ValueError) as ctx:
-            LmdbDataReader(
-                self._mixed_path,
-                self._type_map,
-                batch_size="filter:10",
-                mixed_batch=True,
-            )
-        self.assertIn("filter", str(ctx.exception))
-        self.assertIn("mixed_batch", str(ctx.exception))
-
     def test_auto_prob_with_filter_still_works(self):
         """compute_block_targets + sampler survive a fully-dropped block."""
         path = f"{self._tmpdir.name}/auto_prob_filter.lmdb"
@@ -1399,7 +2438,7 @@ def test_auto_prob_with_filter_still_works(self):
         # Block 0 under-represented relative to weight → expansion needed.
         self.assertGreater(len(block_targets), 0)
 
-        sampler = SameNlocBatchSampler(
+        sampler = LmdbBatchSampler(
             reader, shuffle=False, seed=0, block_targets=block_targets
         )
         all_batches = list(sampler)
diff --git a/source/tests/common/dpmodel/test_loss_ener.py b/source/tests/common/dpmodel/test_loss_ener.py
index 761bc39768..d5c605fb0b 100644
--- a/source/tests/common/dpmodel/test_loss_ener.py
+++ b/source/tests/common/dpmodel/test_loss_ener.py
@@ -66,6 +66,53 @@ def _make_data(
 class TestEnergyLossBasic(TestEnergyLossBase):
     """Test basic energy loss (e, f, v, ae)."""
 
+    def test_inactive_terms_are_absent_from_the_data_contract(self) -> None:
+        """The declared labels must match the terms evaluated by the loss."""
+        loss_fn = EnergyLoss(starter_learning_rate=1.0)
+        required_keys = {item.key for item in loss_fn.label_requirement}
+        self.assertEqual(required_keys, {"energy", "force"})
+
+        model_dict, all_labels, natoms = self._make_data()
+        label_dict = {
+            key: value
+            for key, value in all_labels.items()
+            if key in required_keys
+            or (key.startswith("find_") and key.removeprefix("find_") in required_keys)
+        }
+        loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
+
+        self.assertTrue(np.isfinite(loss))
+        self.assertEqual(set(more_loss), {"rmse_e", "rmse_f", "rmse"})
+
+    def test_prefactor_force_declares_both_inputs(self) -> None:
+        """Atomic force weighting requires force and atom_pref labels."""
+        loss_fn = 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_pf=1.0,
+            limit_pref_pf=1.0,
+        )
+        self.assertEqual(
+            {item.key for item in loss_fn.label_requirement},
+            {"force", "atom_pref"},
+        )
+
+    def test_default_prefactor_is_valid_training_data(self) -> None:
+        """use_default_pf makes a missing atom_pref equivalent to unit weights."""
+        loss_fn = EnergyLoss(
+            starter_learning_rate=1.0,
+            start_pref_pf=1.0,
+            limit_pref_pf=1.0,
+            use_default_pf=True,
+        )
+        atom_pref = next(
+            item for item in loss_fn.label_requirement if item.key == "atom_pref"
+        )
+        self.assertEqual(atom_pref.source_policy, "default")
+
     def test_forward(self) -> None:
         loss_fn = EnergyLoss(
             starter_learning_rate=1.0,
@@ -107,6 +154,32 @@ def test_forward(self) -> None:
         loss, more_loss = loss_fn.call(1.0, natoms, model_dict, label_dict)
         self.assertIsNotNone(loss)
 
+    def test_ragged_frames_are_reduced_independently(self) -> None:
+        """Atomic coefficients retain the frame boundaries of a flat node axis."""
+        loss_fn = 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,
+            enable_atom_ener_coeff=True,
+        )
+        n_node = np.array([2, 3], dtype=np.int64)
+        model_dict = {
+            "energy": np.array([[3.0], [12.0]]),
+            "atom_energy": np.arange(1.0, 6.0).reshape(-1, 1),
+            "n_node": n_node,
+        }
+        label_dict = {
+            "energy": np.array([[3.0], [12.0]]),
+            "find_energy": 1.0,
+            "atom_ener_coeff": np.ones((5, 1)),
+        }
+
+        loss, _ = loss_fn.call(1.0, int(n_node.sum()), model_dict, label_dict)
+
+        self.assertEqual(float(loss), 0.0)
+
 
 class TestEnergyLossGeneralizedForce(TestEnergyLossBase):
     """Test energy loss with generalized force (numb_generalized_coord > 0).
@@ -141,6 +214,56 @@ def test_forward(self) -> None:
         self.assertIn("rmse_gf", more_loss)
         self.assertIn("rmse_pf", more_loss)
 
+    def test_force_derived_terms_require_force_labels(self) -> None:
+        """PF and GF supervision are unavailable when force is unavailable."""
+        cases = (
+            (
+                "prefactor force",
+                {
+                    "start_pref_pf": 1.0,
+                    "limit_pref_pf": 1.0,
+                    "use_default_pf": True,
+                },
+                "rmse_pf",
+                0,
+            ),
+            (
+                "generalized force",
+                {
+                    "start_pref_gf": 1.0,
+                    "limit_pref_gf": 1.0,
+                    "numb_generalized_coord": 2,
+                },
+                "rmse_gf",
+                2,
+            ),
+        )
+        for name, term_kwargs, metric, numb_generalized_coord in cases:
+            with self.subTest(term=name):
+                loss_fn = 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,
+                    **term_kwargs,
+                )
+                model_dict, label_dict, natoms = self._make_data(
+                    numb_generalized_coord=numb_generalized_coord,
+                )
+                label_dict["find_force"] = 0.0
+                loss, more_loss = loss_fn.call(
+                    1.0,
+                    natoms,
+                    model_dict,
+                    label_dict,
+                )
+
+                self.assertEqual(float(loss), 0.0)
+                self.assertTrue(np.isnan(more_loss[metric]))
+
 
 class TestEnergyLossHuber(TestEnergyLossBase):
     """Test energy loss with Huber loss."""
diff --git a/source/tests/common/dpmodel/test_loss_padding.py b/source/tests/common/dpmodel/test_loss_padding.py
index 337c325fdc..e6f9baeae7 100644
--- a/source/tests/common/dpmodel/test_loss_padding.py
+++ b/source/tests/common/dpmodel/test_loss_padding.py
@@ -1,20 +1,22 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Reusable grad-accumulation invariant harness for dpmodel loss tests.
+"""Padding-mask behaviour of the dpmodel losses.
 
-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).
+A batch may hold frames of unequal atom count, padded to a common width with
+slots whose ``atype`` is negative. Every loss term must then reduce over the
+real atoms alone. The dpmodel losses accept numpy arrays through the
+array_api_compat backend.
 
-The dpmodel losses accept numpy arrays (via the array_api_compat backend).
+Harness
+-------
+``assert_grad_accum_invariant`` checks that the loss of a padded two-frame
+batch equals the weighted mean of the frames' individual losses. The weights
+express how much each frame counts: frame-level terms (energy, virial) weigh
+frames equally, while per-atom terms (force, atomic energy, dos, tensor) weigh
+them by their label count, so a caller passes the frames' atom counts there.
 
-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 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.
+Not covered: the TF backend, which still dilutes mixed_type losses
+(deepmodeling/deepmd-kit#5760), and the pt-only ``population`` and ``denoise``
+losses (deepmodeling/deepmd-kit#5761).
 
 Constants
 ---------
@@ -24,6 +26,7 @@
 """
 
 import numpy as np
+import pytest
 
 from deepmd.dpmodel.loss.dos import (
     DOSLoss,
@@ -58,14 +61,24 @@ def assert_grad_accum_invariant(
     make_batch_A,
     make_batch_B,
     make_padded_batch,
+    frame_weights: tuple[float, float] = (1.0, 1.0),
     rtol: float = 1e-5,
     atol: float = 1e-6,
 ) -> None:
-    """Assert that padded-batch loss == mean(per_frame_loss) for two frames.
+    """Assert that a padded batch scores the weighted mean of its 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.
+    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 combining them under the weights the term assigns to a frame:
+
+    - **Frame-level terms** (energy, virial, global dos/cdf, global tensor,
+      property) carry a fixed number of labels per frame, so the frames weigh
+      equally and the reference is the plain mean.  Pass the default weights.
+    - **Per-atom terms** (force, atomic energy, atomic prefactor force, atomic
+      dos/cdf, local tensor) carry labels in proportion to a frame's atom
+      count, and the loss pools them, so the reference is the atom-weighted
+      mean.  Pass ``frame_weights=(NA, NB)``; the per-atom component count
+      cancels out of the ratio.
 
     Parameters
     ----------
@@ -79,6 +92,8 @@ def assert_grad_accum_invariant(
     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).
+    frame_weights : tuple[float, float]
+        Weight of frame A and frame B in the reference combination.
     rtol : float
         Relative tolerance for ``np.isclose``.
     atol : float
@@ -90,7 +105,8 @@ def assert_grad_accum_invariant(
 
     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)
+    weight_A, weight_B = frame_weights
+    ref = (weight_A * loss_A + weight_B * loss_B) / (weight_A + weight_B)
 
     loss_pad = float(loss_fn(pred_pad, label_pad, natoms_pad))
 
@@ -209,7 +225,9 @@ def make_B():
         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)
+        assert_grad_accum_invariant(
+            self._loss_fn, make_A, make_B, make_padded, frame_weights=(NA, NB)
+        )
 
     def test_acdf_grad_accum_invariant(self):
         """Atomic cdf per-frame masked mean meets the grad-accum invariant."""
@@ -366,7 +384,9 @@ def make_padded():
                 NP,
             )
 
-        assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded)
+        assert_grad_accum_invariant(
+            self._loss_fn, make_A, make_B, make_padded, frame_weights=(NA, NB)
+        )
 
     def test_no_op_for_non_mixed(self):
         """All-ones mask gives same loss as no mask (non-mixed batch)."""
@@ -514,6 +534,168 @@ def _padded_atom_flat(arr_A, arr_B, ncomp):
 )  # [2, NP]
 
 
+class TestPerFrameCountSource:
+    """Where the per-frame atom count comes from must not change the loss.
+
+    The extensive terms divide each frame's residual by that frame's atom
+    count. A padded batch states it through its mask, which the per-atom terms
+    also need to skip the padded rows; a concatenated batch has no padded row
+    and states the counts outright in ``n_node``. The two are the same number,
+    so the loss they produce is the same.
+    """
+
+    @staticmethod
+    def _loss(**kwargs):
+        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=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,
+            **kwargs,
+        )
+
+    @pytest.mark.parametrize("loss_func", ["mse", "mae"])
+    @pytest.mark.parametrize("intensive", [False, True])
+    @pytest.mark.parametrize("use_huber", [False, True])
+    def test_mask_and_n_node_agree(self, loss_func, intensive, use_huber):
+        if use_huber and loss_func != "mse":
+            pytest.skip("huber replaces the mse branch only")
+        rng = np.random.default_rng(3)
+        nf, nloc = 2, 3
+        n_node = np.array([2, 3], dtype=np.int64)
+        physical = np.arange(nloc)[None, :] < n_node[:, None]
+        mask = physical.astype(np.float64)
+        # The second physical node is excluded by the model, independently of
+        # the padding slot at the end of the first frame.
+        mask[0, 1] = 0.0
+        pred, label = _full_ener_dicts(
+            nf,
+            nloc,
+            rng.normal(size=(nf, 1)),
+            rng.normal(size=(nf, 1)),
+            mask=mask,
+            force=rng.normal(size=(nf, nloc, 3)),
+            virial=rng.normal(size=(nf, 9)),
+            atom_energy=rng.normal(size=(nf, nloc, 1)),
+            atom_ener=rng.normal(size=(nf, nloc, 1)),
+            atom_pref=rng.random(size=(nf, nloc * 3)),
+            find_force=1.0,
+            find_virial=1.0,
+            find_atom_ener=1.0,
+            find_atom_pref=1.0,
+        )
+        label["force"] = rng.normal(size=(nf, nloc, 3))
+        pref_pf = 0.0 if use_huber else 1.0
+        loss_obj = EnergyLoss(
+            starter_learning_rate=1.0,
+            start_pref_e=1.0,
+            limit_pref_e=1.0,
+            start_pref_f=1.0,
+            limit_pref_f=1.0,
+            start_pref_v=1.0,
+            limit_pref_v=1.0,
+            start_pref_ae=1.0,
+            limit_pref_ae=1.0,
+            start_pref_pf=pref_pf,
+            limit_pref_pf=pref_pf,
+            loss_func=loss_func,
+            intensive_ener_virial=intensive,
+            use_huber=use_huber,
+        )
+        from_mask, more_mask = loss_obj.call(1.0, nloc, pred, label)
+
+        # The same physical nodes on a flat ragged axis. ``n_node`` retains the
+        # frame boundaries while the flat mask retains the model exclusion.
+        by_count = {
+            **pred,
+            "force": pred["force"][physical],
+            "atom_energy": pred["atom_energy"][physical],
+            "mask": mask[physical],
+            "n_node": n_node,
+        }
+        ragged_label = {
+            **label,
+            "force": label["force"][physical],
+            "atom_ener": label["atom_ener"][physical],
+            "atom_pref": label["atom_pref"].reshape(nf, nloc, 3)[physical],
+        }
+        from_counts, more_counts = loss_obj.call(
+            1.0, int(n_node.sum()), by_count, ragged_label
+        )
+
+        np.testing.assert_allclose(
+            float(from_counts), float(from_mask), rtol=1e-14, atol=0
+        )
+        assert sorted(more_counts) == sorted(more_mask)
+        for key, value in more_mask.items():
+            np.testing.assert_allclose(
+                np.asarray(more_counts[key], dtype=np.float64),
+                np.asarray(value, dtype=np.float64),
+                rtol=1e-14,
+                atol=0,
+                err_msg=key,
+            )
+
+    @pytest.mark.parametrize("ragged", [False, True])
+    def test_fully_excluded_frame_is_neutral(self, ragged):
+        """A frame without an included atom contributes no extensive label."""
+        loss_obj = 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,
+        )
+        mask_shape = (1,) if ragged else (1, 1)
+        model = {
+            "energy": np.ones((1, 1), dtype=np.float64),
+            "mask": np.zeros(mask_shape, dtype=np.float64),
+        }
+        if ragged:
+            model["n_node"] = np.ones(1, dtype=np.int64)
+        label = {
+            "energy": np.zeros((1, 1), dtype=np.float64),
+            "find_energy": 1.0,
+        }
+
+        loss, more_loss = loss_obj.call(1.0, 1, model, label)
+
+        assert float(loss) == 0.0
+        assert float(more_loss["rmse_e"]) == 0.0
+
+    def test_generalized_force_refuses_a_concatenated_batch(self):
+        """``drdq`` is stored against a common atom axis, which is gone."""
+        rng = np.random.default_rng(5)
+        nf, nloc = 2, 4
+        pred, label = _full_ener_dicts(
+            nf, nloc, rng.normal(size=(nf, 1)), rng.normal(size=(nf, 1))
+        )
+        pred["n_node"] = np.full(nf, nloc, dtype=np.int64)
+        label["drdq"] = np.zeros((nf, nloc * 3, 2), dtype=np.float64)
+        label["find_drdq"] = 1.0
+        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_gf=1.0,
+            limit_pref_gf=1.0,
+            numb_generalized_coord=2,
+        )
+        with pytest.raises(NotImplementedError, match="same number of atoms"):
+            loss_obj.call(1.0, nloc, pred, label)
+
+
 class TestDPModelEnergyLossEnerGradAccum:
     """Idiom 2 (extensive) for the energy (has_e) term in EnergyLoss.
 
@@ -703,6 +885,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
@@ -731,6 +914,28 @@ def test_huber_grad_accum(self):
             self._make_loss("mse", use_huber=True), f_A, f_A_hat, f_B, f_B_hat
         )
 
+    @pytest.mark.parametrize(
+        ("loss_func", "use_huber"), [("mae", False), ("mse", True)]
+    )
+    def test_f_use_norm_grad_accum(self, loss_func, use_huber):
+        """One L2 norm per atom weighs frames the same way three components do.
+
+        ``f_use_norm`` changes how many labels an atom contributes, so it also
+        changes the divisor of the pooled reduction; the frame weights it
+        produces must still follow the atom counts.
+        """
+        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(loss_func, use_huber=use_huber, f_use_norm=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)
@@ -1230,6 +1435,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
@@ -1385,6 +1591,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
@@ -1834,6 +2041,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
@@ -2030,14 +2238,8 @@ def test_no_op_for_non_mixed(self):
         )
 
 
-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.
-    """
+class TestDPModelEnerSpinLossForceMagConsistentMasks:
+    """A padding mask is idempotent when ``mask_mag`` already excludes it."""
 
     def _make_loss(self):
         return EnergySpinLossDPModel(
@@ -2052,8 +2254,8 @@ def _make_loss(self):
             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."""
+    def test_consistent_padding_mask_is_idempotent(self):
+        """Consistent magnetic and padding masks produce the same loss."""
         fm = _rnd(2, NP, 3)
         fm_hat = _rnd(2, NP, 3)
         loss_obj = self._make_loss()
@@ -2079,8 +2281,7 @@ def _run(with_mask):
         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}"
+            f"consistent masks must be idempotent: {loss_with} vs {loss_without}"
         )
 
 
@@ -2165,6 +2366,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
diff --git a/source/tests/common/dpmodel/test_loss_reduction.py b/source/tests/common/dpmodel/test_loss_reduction.py
index fecb0156c6..ddb59993c3 100644
--- a/source/tests/common/dpmodel/test_loss_reduction.py
+++ b/source/tests/common/dpmodel/test_loss_reduction.py
@@ -14,15 +14,12 @@
 
 
 class TestMaskedAtomMean:
-    """Idiom 1: per-atom masked mean over ncomp components, averaged over frames."""
+    """Idiom 1: mean of a per-atom contribution over the batch's real labels."""
 
     def _ref(self, elem, maskf, ncomp):
-        # reference reduction, numpy
-        nf = elem.shape[0]
+        # reference reduction, numpy: pool over every real label of the batch
         masked = elem * maskf[:, :, None]
-        pfs = masked.reshape(nf, -1).sum(axis=-1)
-        pfd = maskf.sum(axis=-1) * ncomp
-        return (pfs / pfd).mean()
+        return masked.sum() / (maskf.sum() * ncomp)
 
     @pytest.mark.parametrize("ncomp", [1, 3])  # atom-energy (1) and force (3)
     def test_numpy_matches_reference(self, ncomp) -> None:
@@ -32,7 +29,36 @@ def test_numpy_matches_reference(self, ncomp) -> None:
         got = masked_atom_mean(elem, maskf, ncomp)
         np.testing.assert_allclose(got, self._ref(elem, maskf, ncomp), rtol=0, atol=0)
 
-    def test_torch_autograd_and_bit_identical(self) -> None:
+    def test_uniform_atom_count_batch_is_a_frame_mean(self) -> None:
+        """Pooling and the per-frame mean agree when every frame has one size.
+
+        This identity is what lets the pooled convention apply unconditionally:
+        a batch whose frames share an atom count keeps the reduction it had
+        before mixed-nloc batches existed, with no special case for it.
+        """
+        rng = np.random.default_rng(5)
+        elem = rng.random((3, 4, 3))
+        maskf = np.tile(np.array([1.0, 1.0, 1.0, 0.0]), (3, 1))
+        per_frame = (elem * maskf[:, :, None]).reshape(3, -1).sum(axis=-1) / (
+            maskf.sum(axis=-1) * 3
+        )
+        np.testing.assert_allclose(
+            masked_atom_mean(elem, maskf, 3), per_frame.mean(), rtol=1e-14, atol=0
+        )
+
+    def test_frames_weigh_by_their_label_count(self) -> None:
+        """A frame's weight is its share of the batch's labels, not ``1 / nf``."""
+        # Frame 0 keeps one real atom, frame 1 keeps three.
+        maskf = np.array([[1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0]])
+        elem = np.zeros((2, 4, 3))
+        elem[0, :1] = 2.0
+        elem[1, :3] = 6.0
+        # Pooled: (1*3*2 + 3*3*6) / (4*3); a frame mean would give (2+6)/2 = 4.
+        np.testing.assert_allclose(
+            masked_atom_mean(elem, maskf, 3), (6.0 + 54.0) / 12.0, rtol=0, atol=0
+        )
+
+    def test_torch_autograd_and_matches_numpy(self) -> None:
         elem_np = np.random.default_rng(1).random((2, 4, 3))
         maskf_np = np.array([[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0]])
         elem = torch.tensor(elem_np, requires_grad=True)
@@ -40,34 +66,24 @@ def test_torch_autograd_and_bit_identical(self) -> None:
         out = masked_atom_mean(elem, maskf, 3)
         out.backward()
         assert elem.grad is not None
-        # bit-identical to torch-native inline form
-        en = torch.tensor(elem_np)
-        mn = torch.tensor(maskf_np)
-        pfs = (en * mn[:, :, None]).reshape(2, -1).sum(dim=-1)
-        pfd = mn.sum(dim=-1) * 3
-        ref = torch.mean(pfs / pfd)
-        assert out.item() == ref.item()
-
-    def test_all_padding_frame_is_not_nan(self) -> None:
-        # a frame with zero real atoms has per_frame_dof == 0; the ratio must
-        # not become 0/0 = NaN (an independent finiteness invariant -- the
-        # reference formula shares the bug, so equality checks cannot catch it)
+        np.testing.assert_allclose(
+            out.item(), self._ref(elem_np, maskf_np, 3), rtol=1e-14, atol=0
+        )
+
+    def test_all_padding_batch_is_not_nan(self) -> None:
+        # With the pooled reduction only a batch without a single real atom
+        # drives the denominator to zero; the ratio must not become 0/0 = NaN.
         elem = np.random.default_rng(3).random((2, 4, 3))
-        maskf = np.array([[1.0, 1.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]])
+        maskf = np.zeros((2, 4))
         got = masked_atom_mean(elem, maskf, 3)
         assert np.isfinite(got)
-        # the all-padding frame contributes a neutral 0, so the result is the
-        # first frame's masked mean divided by the number of frames
-        pfs0 = (elem[0] * maskf[0][:, None]).reshape(-1).sum()
-        expected = (pfs0 / (maskf[0].sum() * 3)) / 2
-        np.testing.assert_allclose(got, expected, rtol=0, atol=0)
+        np.testing.assert_allclose(got, 0.0, rtol=0, atol=0)
 
-    def test_all_padding_frame_torch_grad_is_not_nan(self) -> None:
+    def test_all_padding_batch_torch_grad_is_not_nan(self) -> None:
         # the guard must keep both the value and the autograd gradient finite
         elem_np = np.random.default_rng(4).random((2, 4, 3))
-        maskf_np = np.array([[1.0, 1.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0]])
         elem = torch.tensor(elem_np, requires_grad=True)
-        maskf = torch.tensor(maskf_np)
+        maskf = torch.zeros((2, 4), dtype=torch.float64)
         out = masked_atom_mean(elem, maskf, 3)
         out.backward()
         assert torch.isfinite(out).item()
diff --git a/source/tests/consistent/test_array_api.py b/source/tests/consistent/test_array_api.py
index a9f2d8eaba..45a6cbe556 100644
--- a/source/tests/consistent/test_array_api.py
+++ b/source/tests/consistent/test_array_api.py
@@ -1,11 +1,16 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import sys
 import unittest
+from unittest.mock import (
+    patch,
+)
 
+import array_api_compat
 import numpy as np
 
 from deepmd.dpmodel.array_api import (
     xp_add_at,
+    xp_asarray_nodetach,
     xp_bincount,
     xp_maximum_at,
     xp_scatter_sum,
@@ -56,6 +61,36 @@ def test_torch_parameter_requires_grad(self) -> None:
         self.assertTrue(param.requires_grad)
         self.assertEqual(param.device, DEVICE)
 
+    @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed")
+    def test_foreign_tensor_is_converted_to_numpy_namespace(self) -> None:
+        tensor = torch.tensor([1.0, 2.0], dtype=torch.float64, device=DEVICE)
+        numpy_namespace = array_api_compat.array_namespace(np.empty(0))
+
+        # A CUDA tensor rejects the direct NumPy protocol. Simulate that
+        # boundary on CPU so the device-to-host fallback is exercised on every
+        # platform.
+        with patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            converted = xp_asarray_nodetach(numpy_namespace, tensor)
+
+        self.assertIsInstance(converted, np.ndarray)
+        np.testing.assert_allclose(converted, np.array([1.0, 2.0]))
+
+    @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed")
+    def test_native_tensor_keeps_its_autograd_graph(self) -> None:
+        tensor = torch.nn.Parameter(
+            torch.tensor([1.0, 2.0], dtype=torch.float64, device=DEVICE)
+        )
+        torch_namespace = array_api_compat.array_namespace(tensor)
+
+        converted = xp_asarray_nodetach(torch_namespace, tensor)
+
+        self.assertIs(converted, tensor)
+        self.assertTrue(converted.requires_grad)
+
 
 class TestXpMaximumAtConsistent(unittest.TestCase):
     """Test maximum-at identities that differ between backend primitives."""
diff --git a/source/tests/consistent/test_lmdb_data.py b/source/tests/consistent/test_lmdb_data.py
index 6e9cecee52..c4f35563ee 100644
--- a/source/tests/consistent/test_lmdb_data.py
+++ b/source/tests/consistent/test_lmdb_data.py
@@ -14,6 +14,8 @@
 
 from deepmd.dpmodel.utils.lmdb_data import (
     LmdbDataReader,
+    collate_lmdb_frames,
+    resolve_per_atom_keys,
 )
 
 try:
@@ -178,7 +180,7 @@ def test_same_properties(self):
         self.assertEqual(self._reader.batch_sizes, self._ds.batch_sizes)
         self.assertEqual(self._reader.nframes, self._ds.nframes)
         self.assertEqual(self._reader.mixed_type, self._ds.mixed_type)
-        self.assertEqual(self._reader.mixed_batch, self._ds.mixed_batch)
+        self.assertEqual(self._reader.mixed_nloc, self._ds.mixed_nloc)
 
     def test_data_requirement(self):
         req = [
@@ -219,8 +221,30 @@ def test_mixed_nloc_same_properties(self):
             ds = LmdbDataset(path, self._type_map, batch_size=2)
             self.assertEqual(reader.nframes, ds.nframes)
             self.assertEqual(reader.batch_sizes, ds.batch_sizes)
-            self.assertEqual(reader.mixed_batch, ds.mixed_batch)
-            self.assertFalse(reader.mixed_batch)
+            self.assertEqual(reader.mixed_nloc, ds.mixed_nloc)
+            self.assertFalse(reader.mixed_nloc)
+
+    def test_mix_batch_size_same_padded_batch(self):
+        """Both decode paths pad a mixed-nloc batch to the same arrays."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            path = _create_mixed_nloc_lmdb(f"{tmpdir}/mixed.lmdb")
+            reader = LmdbDataReader(path, self._type_map, batch_size="mix:24")
+            self.assertTrue(reader.mixed_nloc)
+            # _create_mixed_nloc_lmdb interleaves 6-, 9- and 12-atom frames.
+            indices = [0, 4, 8]
+            per_atom_keys = resolve_per_atom_keys(
+                reader[indices[0]], reader.decode_config
+            )
+            expected = collate_lmdb_frames(
+                [reader[index] for index in indices], per_atom_keys
+            )
+            actual = reader.decode_batch(indices)
+            self.assertEqual(tuple(actual), tuple(expected))
+            for key, expected_value in expected.items():
+                if isinstance(expected_value, np.ndarray):
+                    np.testing.assert_array_equal(actual[key], expected_value)
+                else:
+                    self.assertEqual(actual[key], expected_value)
 
 
 if __name__ == "__main__":
diff --git a/source/tests/pt/model/test_sezm_model.py b/source/tests/pt/model/test_sezm_model.py
index d25b743fae..b9d07df06b 100644
--- a/source/tests/pt/model/test_sezm_model.py
+++ b/source/tests/pt/model/test_sezm_model.py
@@ -1696,6 +1696,20 @@ def _frame(
         )
         return coord, atype, spin, box
 
+    def _mixed_padded_batch(
+        self,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Build two frames whose shorter member carries phantom padding."""
+        coord, atype, spin, box = self._frame()
+        coord = coord.repeat(2, 1, 1)
+        atype = atype.repeat(2, 1)
+        spin = spin.repeat(2, 1, 1)
+        box = box.repeat(2, 1)
+        atype[1, -2:] = -1
+        coord[1, -2:] = 0.0
+        spin[1, -2:] = 0.0
+        return coord, atype, spin, box
+
     def test_spin_routes_initialize_to_zero(self) -> None:
         """Fresh native-spin models start from the spin-free function."""
         model = self._build_model(randomize=False)
@@ -1854,6 +1868,33 @@ def test_joint_rotation_equivariance(self) -> None:
             rtol=1e-6,
         )
 
+    def test_phantom_atoms_are_excluded_from_magnetic_mask(self) -> None:
+        """Mixed-nloc padding never enters the per-type spin lookup."""
+        model = self._build_model()
+        coord, atype, spin, box = self._mixed_padded_batch()
+
+        out = model(coord, atype, spin, box=box)
+        expected_mask = ((atype == 0) & (atype >= 0)).unsqueeze(-1)
+        torch.testing.assert_close(out["mask_mag"], expected_mask)
+        self.assertTrue(torch.all(out["force"][1, -2:] == 0.0))
+        self.assertTrue(torch.all(out["force_mag"][1, -2:] == 0.0))
+
+        lower = model._attach_spin_masks(
+            {
+                "energy_derv_r_mag": torch.zeros(
+                    2,
+                    atype.shape[1],
+                    1,
+                    3,
+                    dtype=coord.dtype,
+                    device=coord.device,
+                )
+            },
+            atype=atype,
+            nall=atype.shape[1],
+        )
+        torch.testing.assert_close(lower["mask_mag"], expected_mask)
+
     @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
     def test_export_matches_forward(self) -> None:
         """The traced ``.pt2`` export reduces to the public forward.
@@ -1962,7 +2003,7 @@ def test_ener_spin_loss_smoke(self) -> None:
     @unittest.skipIf(_SKIP_OFF_COMPILE_TORCH, _SKIP_OFF_COMPILE_TORCH_REASON)
     def test_compile_matches_eager(self) -> None:
         """The compiled native-spin path matches eager force and magnetic force."""
-        coord, atype, spin, box = self._frame()
+        coord, atype, spin, box = self._mixed_padded_batch()
         model_eager = self._build_model(use_compile=False)
         model_cmp = self._build_model(use_compile=True)
         model_cmp.load_state_dict(model_eager.state_dict())
@@ -1980,6 +2021,8 @@ def test_compile_matches_eager(self) -> None:
                 rtol=1.0e-6,
                 msg=f"native-spin compile mismatch on {key}",
             )
+        torch.testing.assert_close(out_c["mask_mag"], out_e["mask_mag"])
+        self.assertTrue(torch.all(~out_c["mask_mag"][1, -2:]))
 
     @staticmethod
     def _extended_spin_inputs(
@@ -2354,6 +2397,147 @@ def test_zbl_respects_exclusions(self) -> None:
                 )
 
 
+class TestSeZMPhantomAtoms(unittest.TestCase):
+    """SeZM must ignore the phantom atoms that pad a mixed-nloc batch.
+
+    A mixed-nloc LMDB batch is rectangular: frames shorter than the batch-wide
+    atom count are padded with slots carrying ``atype = -1`` and zero
+    coordinates. Those slots stand for no physical site, so the padded batch
+    has to reproduce, frame by frame, what the unpadded frames give on their
+    own.
+    """
+
+    def setUp(self) -> None:
+        self.device = env.DEVICE
+        torch.manual_seed(2024)
+
+    def _build_model(self) -> SeZMModel:
+        params = {
+            "type": "SeZM",
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "SeZM",
+                "sel": [20, 20],
+                "rcut": 4.0,
+                "channels": 4,
+                "n_focus": 1,
+                "focus_compete": False,
+                "n_radial": 3,
+                "radial_mlp": [6],
+                "use_env_seed": False,
+                "l_schedule": [1, 0],
+                "mmax": 1,
+                "so2_norm": False,
+                "so2_layers": 1,
+                "n_atten_head": 0,
+                "sandwich_norm": [True, False, True, False],
+                "ffn_neurons": 8,
+                "ffn_blocks": 1,
+                "mlp_bias": True,
+                "layer_scale": True,
+                "use_amp": False,
+                "activation_function": "silu",
+                "glu_activation": True,
+                "precision": "float64",
+                "seed": 7,
+            },
+            "fitting_net": {
+                "neuron": [8],
+                "activation_function": "silu",
+                "precision": "float64",
+                "seed": 7,
+            },
+            "use_compile": False,
+        }
+        model = get_sezm_model(params).to(self.device)
+        torch.manual_seed(1234)
+        with torch.no_grad():
+            for param in model.parameters():
+                param.copy_(torch.randn_like(param) * 0.1)
+        return model.eval()
+
+    @staticmethod
+    def _make_frames() -> tuple[list[np.ndarray], list[np.ndarray], np.ndarray]:
+        """Three frames of unequal atom count sharing one cubic cell."""
+        rng = np.random.default_rng(0)
+        nlocs = (4, 7, 3)
+        coords = [rng.uniform(0.0, 6.0, (nloc, 3)) for nloc in nlocs]
+        atypes = [rng.integers(0, 2, nloc) for nloc in nlocs]
+        return coords, atypes, (np.eye(3) * 10.0).reshape(9)
+
+    def _pad(
+        self, coords: list[np.ndarray], atypes: list[np.ndarray]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        pad_nloc = max(len(a) for a in atypes)
+        coord = np.zeros((len(atypes), pad_nloc, 3))
+        atype = np.full((len(atypes), pad_nloc), -1, dtype=np.int64)
+        for index, (frame_coord, frame_atype) in enumerate(
+            zip(coords, atypes, strict=True)
+        ):
+            coord[index, : len(frame_atype)] = frame_coord
+            atype[index, : len(frame_atype)] = frame_atype
+        return (
+            torch.tensor(coord, dtype=torch.float64, device=self.device),
+            torch.tensor(atype, device=self.device),
+        )
+
+    def test_padded_batch_matches_unpadded_frames(self) -> None:
+        """Energy and real-atom forces are unchanged by the padding."""
+        model = self._build_model()
+        coords, atypes, box = self._make_frames()
+        coord, atype = self._pad(coords, atypes)
+        batched = model(
+            coord,
+            atype,
+            box=torch.tensor(
+                np.tile(box, (len(atypes), 1)),
+                dtype=torch.float64,
+                device=self.device,
+            ),
+        )
+
+        for index, (frame_coord, frame_atype) in enumerate(
+            zip(coords, atypes, strict=True)
+        ):
+            alone = model(
+                torch.tensor(
+                    frame_coord[None], dtype=torch.float64, device=self.device
+                ),
+                torch.tensor(frame_atype[None], device=self.device),
+                box=torch.tensor(box[None], dtype=torch.float64, device=self.device),
+            )
+            torch.testing.assert_close(
+                batched["energy"].reshape(-1)[index],
+                alone["energy"].reshape(-1)[0],
+                atol=1.0e-12,
+                rtol=1.0e-12,
+            )
+            torch.testing.assert_close(
+                batched["force"][index, : len(frame_atype)],
+                alone["force"][0],
+                atol=1.0e-12,
+                rtol=1.0e-12,
+            )
+
+    def test_phantom_atoms_carry_no_force(self) -> None:
+        """Padded slots stay at exactly zero force, so no gradient reaches them."""
+        model = self._build_model()
+        coords, atypes, box = self._make_frames()
+        coord, atype = self._pad(coords, atypes)
+        force = model(
+            coord,
+            atype,
+            box=torch.tensor(
+                np.tile(box, (len(atypes), 1)),
+                dtype=torch.float64,
+                device=self.device,
+            ),
+        )["force"]
+        phantom = atype < 0
+        self.assertTrue(bool(phantom.any()), "fixture must exercise padding")
+        self.assertTrue(bool(torch.all(force[phantom] == 0.0)))
+
+
 class TestSeZMModelModes(unittest.TestCase):
     """Targeted regression tests for SeZM `ener` / `dens` mode routing."""
 
diff --git a/source/tests/pt/test_lmdb_dataloader.py b/source/tests/pt/test_lmdb_dataloader.py
index 4b253aa1ae..d8f70d0dd3 100644
--- a/source/tests/pt/test_lmdb_dataloader.py
+++ b/source/tests/pt/test_lmdb_dataloader.py
@@ -1,7 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """Unit tests for LmdbDataset (PyTorch wrapper) and related PT-specific features.
 
-Pure dpmodel tests (LmdbDataReader, LmdbTestData, SameNlocBatchSampler, type_map
+Pure dpmodel tests (LmdbDataReader, LmdbTestData, LmdbBatchSampler, type_map
 remapping, auto_prob) live in source/tests/common/dpmodel/test_lmdb_data.py.
 Consistency tests (dpmodel vs pt) live in source/tests/consistent/test_lmdb_data.py.
 """
@@ -17,9 +17,11 @@
 )
 from deepmd.dpmodel.utils.lmdb_data import (
     _ENV_CACHE,
-    DistributedSameNlocBatchSampler,
+    PHANTOM_ATOM_TYPE,
+    DistributedLmdbBatchSampler,
+    LmdbBatchSampler,
     LmdbDataReader,
-    SameNlocBatchSampler,
+    LmdbDecodeConfig,
     _decode_frame,
     _read_metadata,
     _remap_keys,
@@ -33,6 +35,9 @@
     LmdbDataset,
     _collate_lmdb_batch,
 )
+from deepmd.pt.utils.stat import (
+    make_stat_input,
+)
 from deepmd.utils.data import (
     DataRequirementItem,
 )
@@ -350,9 +355,7 @@ def test_batch_iteration(self, lmdb_dir):
         )
 
         with torch.device("cpu"):
-            dl = DataLoader(
-                ds, batch_size=2, shuffle=False, collate_fn=_collate_lmdb_batch
-            )
+            dl = DataLoader(ds, batch_size=2, shuffle=False, collate_fn=ds._collate)
             batch = next(iter(dl))
         assert batch["coord"].shape == (2, 6, 3)
         assert batch["energy"].shape == (2, 1)
@@ -368,7 +371,7 @@ def test_inner_dataloader(self, lmdb_dir):
 
     def test_parallel_batch_loader_has_finite_epoch(self, lmdb_dir):
         ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2)
-        sampler = SameNlocBatchSampler(ds._reader, shuffle=False)
+        sampler = LmdbBatchSampler(ds._reader, shuffle=False)
         loader = LmdbBatchDataLoader(
             ds,
             sampler,
@@ -389,13 +392,13 @@ def test_parallel_loaders_share_pool_for_same_dataset(self, lmdb_dir):
         )
         first = LmdbBatchDataLoader(
             first_data,
-            SameNlocBatchSampler(first_data._reader, shuffle=True, seed=1),
+            LmdbBatchSampler(first_data._reader, shuffle=True, seed=1),
             pin_memory=False,
             num_workers=2,
         )
         second = LmdbBatchDataLoader(
             second_data,
-            SameNlocBatchSampler(second_data._reader, shuffle=True, seed=2),
+            LmdbBatchSampler(second_data._reader, shuffle=True, seed=2),
             pin_memory=False,
             num_workers=2,
         )
@@ -415,7 +418,7 @@ def test_small_batch_stays_synchronous(self, lmdb_dir):
         ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2)
         loader = LmdbBatchDataLoader(
             ds,
-            SameNlocBatchSampler(ds._reader, shuffle=False),
+            LmdbBatchSampler(ds._reader, shuffle=False),
             pin_memory=False,
             num_workers=4,
         )
@@ -430,7 +433,7 @@ def test_partial_successor_is_deferred(self, lmdb_dir):
         ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=4)
         loader = LmdbBatchDataLoader(
             ds,
-            SameNlocBatchSampler(ds._reader, shuffle=False),
+            LmdbBatchSampler(ds._reader, shuffle=False),
             pin_memory=False,
             num_workers=4,
         )
@@ -450,7 +453,7 @@ def test_requirements_freeze_after_batch_read(self, lmdb_dir):
         ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2)
         loader = LmdbBatchDataLoader(
             ds,
-            SameNlocBatchSampler(ds._reader, shuffle=False),
+            LmdbBatchSampler(ds._reader, shuffle=False),
             pin_memory=False,
             num_workers=0,
         )
@@ -468,12 +471,73 @@ def test_full_epoch(self, lmdb_dir):
         )
 
         with torch.device("cpu"):
-            dl = DataLoader(
-                ds, batch_size=3, shuffle=False, collate_fn=_collate_lmdb_batch
-            )
+            dl = DataLoader(ds, batch_size=3, shuffle=False, collate_fn=ds._collate)
             total_frames = sum(batch["coord"].shape[0] for batch in dl)
         assert total_frames == 10
 
+    def test_loss_ignores_phantom_atoms(self, lmdb_dir, monkeypatch):
+        """Padding a batch with phantom atoms leaves its loss untouched.
+
+        This closes the loop from the loader to the loss: the slots a
+        mixed-nloc batch adds must enter neither the energy term, which
+        normalizes by atom count, nor the per-atom force mean.
+        """
+        monkeypatch.setattr("deepmd.pt.loss.ener.env.DEVICE", torch.device("cpu"))
+        ds = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=3)
+        ds.add_data_requirement(
+            [
+                DataRequirementItem("energy", 1, atomic=False, must=False),
+                DataRequirementItem("force", 3, atomic=True, must=False),
+            ]
+        )
+        loss_module = EnergyStdLoss(
+            starter_learning_rate=1.0,
+            start_pref_e=1.0,
+            limit_pref_e=1.0,
+            start_pref_f=1.0,
+            limit_pref_f=1.0,
+        )
+
+        def score(batch):
+            """Loss of a constant-zero prediction against this batch's labels."""
+
+            def zero_model(**kwargs):
+                return {
+                    "energy": torch.zeros_like(batch["energy"]),
+                    "force": torch.zeros_like(batch["force"]),
+                }
+
+            _, loss, _ = loss_module(
+                {"atype": batch["atype"]},
+                zero_model,
+                {
+                    "energy": batch["energy"],
+                    "find_energy": batch["find_energy"],
+                    "force": batch["force"],
+                    "find_force": batch["find_force"],
+                },
+                natoms=int(batch["atype"].shape[-1]),
+                learning_rate=1.0,
+            )
+            return float(loss)
+
+        with torch.device("cpu"):
+            batch = ds._collate([ds[index] for index in range(3)])
+            nframes, nloc = batch["atype"].shape
+            padded = dict(batch)
+            padded["atype"] = torch.cat(
+                [
+                    batch["atype"],
+                    torch.full((nframes, 2), PHANTOM_ATOM_TYPE, dtype=torch.int64),
+                ],
+                dim=1,
+            )
+            padded["force"] = torch.cat(
+                [batch["force"], torch.zeros((nframes, 2, 3), dtype=torch.float64)],
+                dim=1,
+            )
+            assert score(padded) == pytest.approx(score(batch), rel=1e-12)
+
     def test_partial_labels_form_homogeneous_loss_batches(self, tmp_path, monkeypatch):
         """Default-filled labels must never share a scalar flag with real ones."""
         monkeypatch.setattr("deepmd.pt.loss.ener.env.DEVICE", torch.device("cpu"))
@@ -487,16 +551,17 @@ def test_partial_labels_form_homogeneous_loss_batches(self, tmp_path, monkeypatc
             ]
         )
 
-        sampler = SameNlocBatchSampler(ds._reader, shuffle=True, seed=11)
+        sampler = LmdbBatchSampler(ds._reader, shuffle=True, seed=11)
         batches = list(sampler)
         assert len(sampler) == len(batches) == 2
+        # Complementary labels must not share a batch: the frames carrying
+        # energy are indices 0 and 2, those carrying force are 1 and 3.
         for indices in batches:
-            signatures = {ds._reader.get_find_signature(index) for index in indices}
-            assert len(signatures) == 1
+            assert len({index % 2 for index in indices}) == 1
 
         distributed_batches = []
         for rank in range(2):
-            distributed = DistributedSameNlocBatchSampler(
+            distributed = DistributedLmdbBatchSampler(
                 ds._reader,
                 rank=rank,
                 world_size=2,
@@ -574,8 +639,8 @@ def zero_model(_batch=batch, **kwargs):
                 assert force_loss.item() > 0.0
         assert observed_flags == {(1.0, 0.0), (0.0, 1.0)}
 
-    def test_unrequested_labels_form_homogeneous_batches(self, tmp_path):
-        """Raw labels outside the loss requirements must still partition batches."""
+    def test_unrequested_labels_do_not_partition_batches(self, tmp_path):
+        """Raw labels outside the active requirements cannot alter sampling."""
         path = str(tmp_path / "partial-virial.lmdb")
         _create_partially_virial_lmdb(path)
         ds = LmdbDataset(path, type_map=["O", "H"], batch_size=2)
@@ -583,10 +648,57 @@ def test_unrequested_labels_form_homogeneous_batches(self, tmp_path):
             [DataRequirementItem("energy", 1, atomic=False, must=False)]
         )
 
-        with torch.device("cpu"):
-            batches = list(ds._inner_dataloader)
-        assert len(batches) == 2
-        assert sorted("find_virial" in batch for batch in batches) == [False, True]
+        groups = ds._reader.availability_groups(np.arange(len(ds), dtype=np.int64))
+        assert len(groups) == 1
+        np.testing.assert_array_equal(groups[0], np.arange(len(ds)))
+
+    def test_default_backed_fparam_survives_statistics(self, lmdb_dir):
+        """Statistics retain explicit and default-resolved frame parameters."""
+        environment = lmdb.open(lmdb_dir, readonly=False, lock=False)
+        with environment.begin(write=True) as transaction:
+            for index in range(0, 10, 2):
+                key = format(index, "012d").encode()
+                frame = msgpack.unpackb(transaction.get(key), raw=False)
+                frame["fparam"] = {
+                    "type": "<f8",
+                    "shape": (2,),
+                    "data": np.array([2.0, 3.0], dtype=np.float64).tobytes(),
+                }
+                transaction.put(key, msgpack.packb(frame, use_bin_type=True))
+        environment.close()
+
+        dataset = LmdbDataset(lmdb_dir, type_map=["O", "H"], batch_size=2)
+        dataset.add_data_requirement(
+            [
+                DataRequirementItem(
+                    "fparam",
+                    2,
+                    atomic=False,
+                    default=np.array([0.0, 1.0]),
+                    source_policy="default",
+                )
+            ]
+        )
+        sampled = make_stat_input(
+            dataset.systems,
+            dataset.dataloaders,
+            nbatches=5,
+        )
+        fparam = sampled[0]["fparam"]
+        explicit = torch.tensor(
+            [2.0, 3.0],
+            dtype=fparam.dtype,
+            device=fparam.device,
+        )
+        default = torch.tensor(
+            [0.0, 1.0],
+            dtype=fparam.dtype,
+            device=fparam.device,
+        )
+
+        assert float(sampled[0]["find_fparam"]) == 1.0
+        assert int(torch.all(fparam == explicit, dim=1).sum()) == 5
+        assert int(torch.all(fparam == default, dim=1).sum()) == 5
 
 
 # ============================================================
@@ -594,6 +706,16 @@ def test_unrequested_labels_form_homogeneous_batches(self, tmp_path):
 # ============================================================
 
 
+_BARE_DECODE_CONFIG = LmdbDecodeConfig(
+    ntypes=2, natoms=0, type_remap=None, data_requirements={}
+)
+
+
+def _collate(frames):
+    """Collate hand-built frames with no registered data requirements."""
+    return _collate_lmdb_batch(frames, _BARE_DECODE_CONFIG)
+
+
 class TestCollate:
     """Test collate function."""
 
@@ -613,7 +735,7 @@ def test_collate_basic(self):
                 "fid": 1,
             },
         ]
-        batch = _collate_lmdb_batch(frames)
+        batch = _collate(frames)
         assert batch["coord"].shape == (2, 4, 3)
         assert batch["fid"] == [0, 1]
         assert batch["sid"] == 0
@@ -623,22 +745,58 @@ def test_collate_skips_type(self):
             {"coord": np.zeros((2, 3)), "type": np.array([0, 1])},
             {"coord": np.zeros((2, 3)), "type": np.array([0, 1])},
         ]
-        assert "type" not in _collate_lmdb_batch(frames)
+        assert "type" not in _collate(frames)
 
     def test_collate_none_values(self):
         frames = [
             {"coord": np.zeros((2, 3)), "box": None},
             {"coord": np.zeros((2, 3)), "box": None},
         ]
-        assert _collate_lmdb_batch(frames)["box"] is None
+        assert _collate(frames)["box"] is None
 
-    def test_collate_rejects_mixed_find_flags(self):
+    def test_collate_demotes_mixed_find_flags(self):
+        """One scalar flag cannot claim a label only some frames supply."""
         frames = [
             {"coord": np.zeros((2, 3)), "find_energy": 1.0},
             {"coord": np.zeros((2, 3)), "find_energy": 0.0},
         ]
-        with pytest.raises(ValueError, match="mixes 'find_energy' values"):
-            _collate_lmdb_batch(frames)
+        assert float(_collate(frames)["find_energy"]) == 0.0
+
+    def test_collate_keeps_unanimous_find_flags(self):
+        frames = [
+            {"coord": np.zeros((2, 3)), "find_energy": 1.0},
+            {"coord": np.zeros((2, 3)), "find_energy": 1.0},
+        ]
+        assert float(_collate(frames)["find_energy"]) == 1.0
+
+    def test_collate_pads_the_atom_axis(self):
+        """Frames of different atom counts stack into one padded batch."""
+        frames = [
+            {
+                "coord": np.ones((2, 3)),
+                "atype": np.zeros(2, dtype=np.int64),
+                "energy": np.array([1.0]),
+            },
+            {
+                "coord": np.ones((4, 3)),
+                "atype": np.zeros(4, dtype=np.int64),
+                "energy": np.array([2.0]),
+            },
+        ]
+        batch = _collate(frames)
+        assert batch["coord"].shape == (2, 4, 3)
+        assert batch["atype"].shape == (2, 4)
+        # The short frame keeps its two atoms and gains two phantom slots.
+        assert batch["atype"][0].tolist() == [
+            0,
+            0,
+            PHANTOM_ATOM_TYPE,
+            PHANTOM_ATOM_TYPE,
+        ]
+        assert batch["atype"][1].tolist() == [0, 0, 0, 0]
+        assert torch.all(batch["coord"][0, 2:] == 0)
+        # Frame-level fields keep their own shape.
+        assert batch["energy"].shape == (2, 1)
 
 
 # ============================================================
@@ -743,15 +901,15 @@ def multi_nloc_lmdb(tmp_path):
     return lmdb_path
 
 
-class TestDistributedSameNlocBatchSampler:
-    """Test DistributedSameNlocBatchSampler (pure logic, no torch.distributed)."""
+class TestDistributedLmdbBatchSampler:
+    """Test DistributedLmdbBatchSampler (pure logic, no torch.distributed)."""
 
     def test_disjoint_batches(self, multi_nloc_lmdb):
         reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=1)
-        s0 = DistributedSameNlocBatchSampler(
+        s0 = DistributedLmdbBatchSampler(
             reader, rank=0, world_size=2, shuffle=True, seed=42
         )
-        s1 = DistributedSameNlocBatchSampler(
+        s1 = DistributedLmdbBatchSampler(
             reader, rank=1, world_size=2, shuffle=True, seed=42
         )
         frames0 = {i for batch in s0 for i in batch}
@@ -760,10 +918,10 @@ def test_disjoint_batches(self, multi_nloc_lmdb):
 
     def test_covers_all_frames(self, multi_nloc_lmdb):
         reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2)
-        s0 = DistributedSameNlocBatchSampler(
+        s0 = DistributedLmdbBatchSampler(
             reader, rank=0, world_size=2, shuffle=True, seed=42
         )
-        s1 = DistributedSameNlocBatchSampler(
+        s1 = DistributedLmdbBatchSampler(
             reader, rank=1, world_size=2, shuffle=True, seed=42
         )
         all_frames = {i for batch in s0 for i in batch} | {
@@ -775,9 +933,9 @@ def test_len(self, multi_nloc_lmdb):
         import math
 
         reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2)
-        total = len(SameNlocBatchSampler(reader, shuffle=False))
+        total = len(LmdbBatchSampler(reader, shuffle=False))
         samplers = [
-            DistributedSameNlocBatchSampler(
+            DistributedLmdbBatchSampler(
                 reader, rank=rank, world_size=4, shuffle=False, seed=0
             )
             for rank in range(4)
@@ -787,17 +945,17 @@ def test_len(self, multi_nloc_lmdb):
 
     def test_deterministic(self, multi_nloc_lmdb):
         reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2)
-        s1 = DistributedSameNlocBatchSampler(
+        s1 = DistributedLmdbBatchSampler(
             reader, rank=0, world_size=2, shuffle=True, seed=42
         )
-        s2 = DistributedSameNlocBatchSampler(
+        s2 = DistributedLmdbBatchSampler(
             reader, rank=0, world_size=2, shuffle=True, seed=42
         )
         assert list(s1) == list(s2)
 
     def test_set_epoch_changes_order(self, multi_nloc_lmdb):
         reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2)
-        s = DistributedSameNlocBatchSampler(
+        s = DistributedLmdbBatchSampler(
             reader, rank=0, world_size=2, shuffle=True, seed=42
         )
         s.set_epoch(0)
@@ -810,12 +968,12 @@ def test_single_gpu_fallback(self, multi_nloc_lmdb):
         reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2)
         single = {
             i
-            for batch in SameNlocBatchSampler(reader, shuffle=True, seed=42)
+            for batch in LmdbBatchSampler(reader, shuffle=True, seed=42)
             for i in batch
         }
         dist = {
             i
-            for batch in DistributedSameNlocBatchSampler(
+            for batch in DistributedLmdbBatchSampler(
                 reader, rank=0, world_size=1, shuffle=True, seed=42
             )
             for i in batch
@@ -824,7 +982,7 @@ def test_single_gpu_fallback(self, multi_nloc_lmdb):
 
     def test_same_nloc_per_batch(self, multi_nloc_lmdb):
         reader = LmdbDataReader(multi_nloc_lmdb, type_map=["O", "H"], batch_size=2)
-        s = DistributedSameNlocBatchSampler(
+        s = DistributedLmdbBatchSampler(
             reader, rank=0, world_size=2, shuffle=True, seed=42
         )
         for batch in s:
@@ -936,14 +1094,14 @@ def test_distributed_len_includes_auto_prob_expansion(self, auto_prob_lmdb):
             auto_prob_style="prob_sys_size;0:1:0.5;1:3:0.5",
         )
         global_batches = len(ds._batch_sampler)
-        dist_sampler_rank0 = DistributedSameNlocBatchSampler(
+        dist_sampler_rank0 = DistributedLmdbBatchSampler(
             ds._reader,
             rank=0,
             world_size=2,
             shuffle=False,
             block_targets=ds._block_targets,
         )
-        dist_sampler_rank1 = DistributedSameNlocBatchSampler(
+        dist_sampler_rank1 = DistributedLmdbBatchSampler(
             ds._reader,
             rank=1,
             world_size=2,
@@ -955,28 +1113,26 @@ def test_distributed_len_includes_auto_prob_expansion(self, auto_prob_lmdb):
         assert len(dist_sampler_rank0) == len(list(dist_sampler_rank0))
         assert len(dist_sampler_rank1) == len(list(dist_sampler_rank1))
 
-    def test_distributed_len_reuses_cached_total(self, auto_prob_lmdb, monkeypatch):
-        calls = 0
-        real_sampler = lmdb_data.SameNlocBatchSampler
+    def test_distributed_builds_batches_once_per_epoch(
+        self, auto_prob_lmdb, monkeypatch
+    ):
+        """Batch construction is shared by ``__len__`` and ``__iter__``."""
+        builds = 0
+        real_build = lmdb_data._build_all_batches
 
-        class CountingSameNlocBatchSampler(real_sampler):
-            def __init__(self, *args, **kwargs):
-                nonlocal calls
-                calls += 1
-                super().__init__(*args, **kwargs)
+        def counting_build(*args, **kwargs):
+            nonlocal builds
+            builds += 1
+            return real_build(*args, **kwargs)
 
-        monkeypatch.setattr(
-            lmdb_data,
-            "SameNlocBatchSampler",
-            CountingSameNlocBatchSampler,
-        )
+        monkeypatch.setattr(lmdb_data, "_build_all_batches", counting_build)
         ds = LmdbDataset(
             auto_prob_lmdb,
             type_map=["O", "H"],
             batch_size=4,
             auto_prob_style="prob_sys_size;0:1:0.5;1:3:0.5",
         )
-        dist_sampler = DistributedSameNlocBatchSampler(
+        dist_sampler = DistributedLmdbBatchSampler(
             ds._reader,
             rank=1,
             world_size=2,
@@ -984,10 +1140,17 @@ def __init__(self, *args, **kwargs):
             block_targets=ds._block_targets,
         )
 
-        assert calls == 1
+        assert builds == 0
         expected_len = (len(ds._batch_sampler) + 1) // 2
+        builds = 0
         assert len(dist_sampler) == expected_len
-        assert calls == 1
+        assert len(list(dist_sampler)) == expected_len
+        assert builds == 1
+
+        # A new epoch reshuffles, and so must rebuild exactly once more.
+        dist_sampler.set_epoch(1)
+        assert len(list(dist_sampler)) == len(dist_sampler)
+        assert builds == 2
 
 
 class TestMergeLmdbSystemIds:
@@ -1151,6 +1314,70 @@ def fail() -> None:
     assert loader.closed
 
 
+def test_numb_epoch_counts_mixed_nloc_batches(multi_nloc_lmdb, tmp_path, monkeypatch):
+    """One epoch of a ``mix:N`` dataset is one pass over its padded batches.
+
+    The batch count of a mixed-nloc pass depends on how the shuffle groups
+    atom counts, so an epoch can only be measured on the sampler the trainer
+    will actually iterate -- not on a nominal batch size.
+    """
+    from deepmd.pt.entrypoints.main import (
+        get_trainer,
+    )
+    from deepmd.utils.argcheck import (
+        normalize,
+    )
+    from deepmd.utils.compat import (
+        update_deepmd_input,
+    )
+
+    config = {
+        "model": {
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "se_e2_a",
+                "sel": [4, 4],
+                "rcut_smth": 0.5,
+                "rcut": 4.0,
+                "neuron": [4, 8],
+                "axis_neuron": 4,
+                "precision": "float64",
+                "seed": 1,
+            },
+            "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 1},
+            "data_stat_nbatch": 1,
+        },
+        "learning_rate": {
+            "type": "exp",
+            "decay_steps": 50,
+            "start_lr": 1e-3,
+            "stop_lr": 1e-8,
+        },
+        "loss": {"type": "ener", "start_pref_e": 1.0, "start_pref_f": 1.0},
+        "training": {
+            "training_data": {
+                "systems": multi_nloc_lmdb,
+                "batch_size": "mix:24",
+            },
+            "numb_epoch": 3,
+            "seed": 10,
+            "disp_file": str(tmp_path / "lcurve.out"),
+            "disp_freq": 100,
+            "save_freq": 100,
+        },
+    }
+    monkeypatch.chdir(tmp_path)
+    config = normalize(update_deepmd_input(config, warning=False))
+    trainer = get_trainer(config)
+    try:
+        batches_per_epoch = len(trainer.training_dataloader)
+        assert trainer.training_dataloader.dataset.mixed_nloc
+        assert batches_per_epoch > 0
+        assert trainer.num_steps == 3 * batches_per_epoch
+    finally:
+        trainer.training_dataloader.close()
+
+
 @pytest.fixture
 def multitask_lmdb_setup(tmp_path):
     """Create two LMDB datasets and a multitask training config."""
diff --git a/source/tests/pt/test_loss_default_pf.py b/source/tests/pt/test_loss_default_pf.py
index c7b837123d..72be87f402 100644
--- a/source/tests/pt/test_loss_default_pf.py
+++ b/source/tests/pt/test_loss_default_pf.py
@@ -73,6 +73,13 @@ def get_batch(system, type_map, data_requirement):
     return np_batch, pt_batch
 
 
+def _is_nan(value: float | torch.Tensor) -> bool:
+    """Return whether a scalar metric is NaN on any tensor device."""
+    if isinstance(value, torch.Tensor):
+        return bool(torch.isnan(value.detach()).item())
+    return bool(np.isnan(value))
+
+
 class TestEnerStdLossDefaultPf(unittest.TestCase):
     """Test use_default_pf feature in EnergyStdLoss."""
 
@@ -176,6 +183,71 @@ def fake_model():
         # The pref_force_loss should be a valid number (not NaN)
         self.assertFalse(np.isnan(pt_more_loss["l2_pref_force_loss"]))
 
+    def test_default_pf_requires_force_labels(self) -> None:
+        """Default atom weights cannot make a missing force label usable."""
+        loss_fn = EnergyStdLoss(
+            self.start_lr,
+            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_pf=1.0,
+            limit_pref_pf=1.0,
+            use_default_pf=True,
+        )
+        label = dict(self.label_without_pref)
+        label["find_force"] = 0.0
+
+        _, loss, more_loss = loss_fn(
+            {},
+            lambda: self.model_pred,
+            label,
+            self.nloc,
+            self.cur_lr,
+        )
+
+        self.assertEqual(float(loss.detach().cpu()), 0.0)
+        self.assertTrue(_is_nan(more_loss["l2_pref_force_loss"]))
+        self.assertTrue(_is_nan(more_loss["rmse_pf"]))
+
+    def test_generalized_force_requires_force_labels(self) -> None:
+        """A present projection matrix cannot validate a missing force label."""
+        numb_generalized_coord = 2
+        loss_fn = EnergyStdLoss(
+            self.start_lr,
+            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=numb_generalized_coord,
+        )
+        label = dict(self.label_with_pref)
+        label["find_force"] = 0.0
+        label["drdq"] = torch.ones(
+            (1, self.nloc * 3 * numb_generalized_coord),
+            dtype=label["force"].dtype,
+            device=label["force"].device,
+        )
+        label["find_drdq"] = 1.0
+
+        _, loss, more_loss = loss_fn(
+            {},
+            lambda: self.model_pred,
+            label,
+            self.nloc,
+            self.cur_lr,
+        )
+
+        self.assertEqual(float(loss.detach().cpu()), 0.0)
+        self.assertTrue(_is_nan(more_loss["l2_gen_force_loss"]))
+        self.assertTrue(_is_nan(more_loss["rmse_gf"]))
+
     def test_default_pf_disabled(self) -> None:
         """With use_default_pf=False (default), pf loss should NOT be computed without find_atom_pref."""
         loss_fn = EnergyStdLoss(
@@ -287,6 +359,7 @@ def test_label_requirement_atom_pref_default(self) -> None:
         label_req = loss_fn.label_requirement
         atom_pref_req = next(r for r in label_req if r.key == "atom_pref")
         self.assertEqual(atom_pref_req.default, 1.0)
+        self.assertEqual(atom_pref_req.source_policy, "default")
 
     def test_serialize_deserialize(self) -> None:
         """Serialization round-trip should preserve use_default_pf."""
diff --git a/source/tests/pt/test_loss_padding.py b/source/tests/pt/test_loss_padding.py
index 094bcfeadb..927c9ba0cd 100644
--- a/source/tests/pt/test_loss_padding.py
+++ b/source/tests/pt/test_loss_padding.py
@@ -1,28 +1,26 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""Tests for mixed_type loss padding-mask support in the pt backend.
+"""Padding-mask behaviour of the pt backend losses.
 
-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.
+A batch may hold frames of unequal atom count, padded to a common width with
+slots whose ``atype`` is negative. ``TaskLoss._inject_atom_mask`` recovers the
+per-atom mask from ``atype``, and every loss term must then reduce over the
+real atoms alone.
 
 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).
-
-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 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.
+``assert_grad_accum_invariant`` checks that the loss of a padded two-frame
+batch equals the weighted mean of the frames' individual losses. The weights
+express how much each frame counts: frame-level terms (energy, virial) weigh
+frames equally, while per-atom terms (force, atomic energy, dos, tensor) weigh
+them by their label count, so a caller passes the frames' atom counts there.
+
+Not covered: the TF backend, which still dilutes mixed_type losses
+(deepmodeling/deepmd-kit#5760), and the pt-only ``population`` and ``denoise``
+losses (deepmodeling/deepmd-kit#5761).
 """
 
 import numpy as np
+import pytest
 import torch
 
 from deepmd.pt.loss.dos import (
@@ -64,14 +62,24 @@ def assert_grad_accum_invariant(
     make_batch_A,
     make_batch_B,
     make_padded_batch,
+    frame_weights: tuple[float, float] = (1.0, 1.0),
     rtol: float = 1e-5,
     atol: float = 1e-6,
 ) -> None:
-    """Assert that padded-batch loss == mean(per_frame_loss) for two frames.
+    """Assert that a padded batch scores the weighted mean of its 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.
+    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 combining them under the weights the term assigns to a frame:
+
+    - **Frame-level terms** (energy, virial, global dos/cdf, global tensor,
+      property) carry a fixed number of labels per frame, so the frames weigh
+      equally and the reference is the plain mean.  Pass the default weights.
+    - **Per-atom terms** (force, atomic energy, atomic prefactor force, atomic
+      dos/cdf, local tensor) carry labels in proportion to a frame's atom
+      count, and the loss pools them, so the reference is the atom-weighted
+      mean.  Pass ``frame_weights=(NA, NB)``; the per-atom component count
+      cancels out of the ratio.
 
     Parameters
     ----------
@@ -84,6 +92,8 @@ def assert_grad_accum_invariant(
     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).
+    frame_weights : tuple[float, float]
+        Weight of frame A and frame B in the reference combination.
     rtol : float
         Relative tolerance for ``torch.isclose``.
     atol : float
@@ -95,7 +105,8 @@ def assert_grad_accum_invariant(
 
     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)
+    weight_A, weight_B = frame_weights
+    ref = (weight_A * loss_A + weight_B * loss_B) / (weight_A + weight_B)
 
     loss_pad = loss_fn(pred_pad, label_pad, natoms_pad)
 
@@ -216,7 +227,9 @@ def make_padded():
                 NP,
             )
 
-        assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded)
+        assert_grad_accum_invariant(
+            self._loss_fn, make_A, make_B, make_padded, frame_weights=(NA, NB)
+        )
 
     def test_no_op_for_non_mixed(self):
         """All-ones mask gives same loss as no mask (non-mixed batch)."""
@@ -404,7 +417,9 @@ def make_padded():
                 NP,
             )
 
-        assert_grad_accum_invariant(self._loss_fn, make_A, make_B, make_padded)
+        assert_grad_accum_invariant(
+            self._loss_fn, make_A, make_B, make_padded, frame_weights=(NA, NB)
+        )
 
     def test_no_op_for_non_mixed(self):
         """All-ones mask gives same loss as no mask (non-mixed batch)."""
@@ -826,6 +841,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
@@ -854,6 +870,28 @@ def test_huber_grad_accum(self):
             self._make_loss("mse", use_huber=True), f_A, f_A_hat, f_B, f_B_hat
         )
 
+    @pytest.mark.parametrize(
+        ("loss_func", "use_huber"), [("mae", False), ("mse", True)]
+    )
+    def test_f_use_norm_grad_accum(self, loss_func, use_huber):
+        """One L2 norm per atom weighs frames the same way three components do.
+
+        ``f_use_norm`` changes how many labels an atom contributes, so it also
+        changes the divisor of the pooled reduction; the frame weights it
+        produces must still follow the atom counts.
+        """
+        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(loss_func, use_huber=use_huber, f_use_norm=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)
@@ -1036,6 +1074,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
@@ -1159,6 +1198,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
@@ -1592,6 +1632,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
@@ -1867,6 +1908,7 @@ def make_padded():
             make_A,
             make_B,
             make_padded,
+            frame_weights=(NA, NB),
         )
 
     def test_mse_grad_accum(self):
diff --git a/source/tests/pt_expt/descriptor/test_dpa4.py b/source/tests/pt_expt/descriptor/test_dpa4.py
index ed405be916..725fa642f4 100644
--- a/source/tests/pt_expt/descriptor/test_dpa4.py
+++ b/source/tests/pt_expt/descriptor/test_dpa4.py
@@ -107,6 +107,30 @@ def test_consistency(self, use_env_seed, use_mapping) -> None:
             err_msg=err_msg,
         )
 
+    def test_default_charge_spin_uses_model_namespace(self) -> None:
+        """A device buffer supplies the default without a NumPy round trip."""
+        dtype = PRECISION_DICT["float64"]
+        descriptor = make_descriptor(
+            self.nt,
+            self.sel_mix,
+            self.rcut,
+            add_chg_spin_ebd=True,
+            default_chg_spin=[0.5, -0.5],
+        ).to(self.device)
+        coord_ext = torch.tensor(self.coord_ext, dtype=dtype, device=self.device)
+        atype_ext = torch.tensor(self.atype_ext, dtype=int, device=self.device)
+        nlist = torch.tensor(self.nlist, dtype=int, device=self.device)
+
+        with mock.patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            output = descriptor(coord_ext, atype_ext, nlist)[0]
+
+        assert output.device == self.device
+        assert torch.isfinite(output).all()
+
     def test_train_and_eval_amp_switches_are_independent(self) -> None:
         """Training follows ``use_amp``, evaluation follows ``DP_AMP_INFER``.
 
diff --git a/source/tests/pt_expt/loss/test_ener_spin.py b/source/tests/pt_expt/loss/test_ener_spin.py
index 2cd49f992a..907d9a2e5b 100644
--- a/source/tests/pt_expt/loss/test_ener_spin.py
+++ b/source/tests/pt_expt/loss/test_ener_spin.py
@@ -293,3 +293,99 @@ def test_no_magnetic_atoms(self, prec, loss_func) -> None:
             torch.testing.assert_close(
                 more_loss["rmse_fm"], torch.zeros_like(more_loss["rmse_fm"])
             )
+
+    @pytest.mark.parametrize("loss_func", ["mse", "mae"])
+    def test_ragged_matches_masked_rectangular(self, loss_func: str) -> None:
+        """Every spin-loss term is invariant to padding the node axis."""
+        dtype = torch.float64
+        device = self.device
+        n_node = torch.tensor([3, 2], dtype=torch.int64, device=device)
+        real_atom = torch.tensor(
+            [[True, True, True], [True, True, False]], device=device
+        )
+
+        def padded(values: torch.Tensor, fill: float = 0.0) -> torch.Tensor:
+            result = torch.full(
+                (2, 3, *values.shape[1:]), fill, dtype=dtype, device=device
+            )
+            result[real_atom] = values
+            return result
+
+        atom_energy = torch.arange(1, 6, dtype=dtype, device=device).reshape(-1, 1)
+        force = torch.arange(15, dtype=dtype, device=device).reshape(-1, 3) / 10.0
+        force_mag = force + 0.25
+        atom_coeff = torch.tensor(
+            [[1.0], [0.5], [-0.5], [2.0], [1.5]], dtype=dtype, device=device
+        )
+        magnetic = torch.tensor(
+            [[True], [False], [True], [False], [True]], device=device
+        )
+        model_rect = {
+            "energy": torch.zeros((2, 1), dtype=dtype, device=device),
+            "force": padded(force),
+            "force_mag": padded(force_mag),
+            "mask_mag": padded(magnetic.to(dtype), fill=1.0).bool(),
+            "virial": torch.arange(18, dtype=dtype, device=device).reshape(2, 9),
+            "atom_energy": padded(atom_energy, fill=100.0),
+            "mask": real_atom.to(dtype),
+        }
+        label_rect = {
+            "energy": torch.tensor([[1.0], [2.0]], dtype=dtype, device=device),
+            "force": padded(force + 0.5, fill=100.0),
+            "force_mag": padded(force_mag - 0.5, fill=100.0),
+            "virial": model_rect["virial"] + 0.5,
+            "atom_ener": padded(atom_energy - 0.25, fill=100.0),
+            "atom_ener_coeff": padded(atom_coeff, fill=3.0),
+            "find_energy": 1.0,
+            "find_force": 1.0,
+            "find_force_mag": 1.0,
+            "find_virial": 1.0,
+            "find_atom_ener": 1.0,
+        }
+        model_ragged = {
+            "energy": model_rect["energy"],
+            "force": force,
+            "force_mag": force_mag,
+            "mask_mag": magnetic,
+            "virial": model_rect["virial"],
+            "atom_energy": atom_energy,
+            "mask": torch.ones(5, dtype=dtype, device=device),
+            "n_node": n_node,
+        }
+        label_ragged = {
+            "energy": label_rect["energy"],
+            "force": force + 0.5,
+            "force_mag": force_mag - 0.5,
+            "virial": label_rect["virial"],
+            "atom_ener": atom_energy - 0.25,
+            "atom_ener_coeff": atom_coeff,
+            "find_energy": 1.0,
+            "find_force": 1.0,
+            "find_force_mag": 1.0,
+            "find_virial": 1.0,
+            "find_atom_ener": 1.0,
+        }
+        loss_fn = EnergySpinLoss(
+            starter_learning_rate=1.0,
+            start_pref_e=1.0,
+            limit_pref_e=1.0,
+            start_pref_fr=1.0,
+            limit_pref_fr=1.0,
+            start_pref_fm=1.0,
+            limit_pref_fm=1.0,
+            start_pref_v=1.0,
+            limit_pref_v=1.0,
+            start_pref_ae=1.0,
+            limit_pref_ae=1.0,
+            enable_atom_ener_coeff=True,
+            loss_func=loss_func,
+            intensive_ener_virial=True,
+        )
+
+        rectangular_loss, rectangular_more = loss_fn(1.0, 3, model_rect, label_rect)
+        ragged_loss, ragged_more = loss_fn(1.0, 5, model_ragged, label_ragged)
+
+        torch.testing.assert_close(ragged_loss, rectangular_loss)
+        assert ragged_more.keys() == rectangular_more.keys()
+        for key in ragged_more:
+            torch.testing.assert_close(ragged_more[key], rectangular_more[key])
diff --git a/source/tests/pt_expt/model/test_dpa2_graph_lower.py b/source/tests/pt_expt/model/test_dpa2_graph_lower.py
index a984fc74bf..d8dc85a43e 100644
--- a/source/tests/pt_expt/model/test_dpa2_graph_lower.py
+++ b/source/tests/pt_expt/model/test_dpa2_graph_lower.py
@@ -200,23 +200,34 @@ def _make_model(
     # 1. routing/parity: proves the Task-7 generic plumbing works for DPA2.
     # ------------------------------------------------------------------
     def test_force_virial_parity_vs_legacy(self) -> None:
-        """Default-flip ``forward_common`` (graph) matches
-        ``neighbor_graph_method="legacy"`` (dense) bit-tight at non-binding
-        sel with repformer attention off.
+        """The compact graph path matches dense and unpadded evaluation.
+
+        A padded input exercises the graph path's compact-gather-scatter
+        round trip. At non-binding ``sel`` with repformer attention disabled,
+        it must match both the legacy dense route on the padded input and the
+        graph route on the equivalent input with the phantom suffix removed.
         """
         model = self._make_model()  # non-binding sel, attention off
         model.eval()
         box = self.cell.reshape(1, 9)
+        padded_atype = self.atype.clone()
+        padded_atype[:, -2:] = -1
+        nreal = self.natoms - 2
 
         graph = model.forward_common(
-            self.coord.clone().requires_grad_(True), self.atype, box
+            self.coord.clone().requires_grad_(True), padded_atype, box
         )
         legacy = model.forward_common(
             self.coord.clone().requires_grad_(True),
-            self.atype,
+            padded_atype,
             box,
             neighbor_graph_method="legacy",
         )
+        unpadded = model.forward_common(
+            self.coord[:, :nreal].clone().requires_grad_(True),
+            self.atype[:, :nreal],
+            box,
+        )
         tol = {"rtol": 1e-10, "atol": 1e-10}
         torch.testing.assert_close(graph["energy_redu"], legacy["energy_redu"], **tol)
         torch.testing.assert_close(
@@ -225,6 +236,21 @@ def test_force_virial_parity_vs_legacy(self) -> None:
         torch.testing.assert_close(
             graph["energy_derv_c_redu"], legacy["energy_derv_c_redu"], **tol
         )
+        torch.testing.assert_close(graph["energy_redu"], unpadded["energy_redu"], **tol)
+        torch.testing.assert_close(
+            graph["energy_derv_r"][:, :nreal],
+            unpadded["energy_derv_r"],
+            **tol,
+        )
+        torch.testing.assert_close(
+            graph["energy_derv_r"][:, nreal:],
+            torch.zeros_like(graph["energy_derv_r"][:, nreal:]),
+            rtol=0,
+            atol=0,
+        )
+        torch.testing.assert_close(
+            graph["energy_derv_c_redu"], unpadded["energy_derv_c_redu"], **tol
+        )
 
     def test_graph_native_hessian_matches_dense(self) -> None:
         """The graph route computes the Hessian natively and it matches the
@@ -255,15 +281,43 @@ def test_graph_native_hessian_matches_dense(self) -> None:
         assert _model_uses_graph_lower(model) is True
 
         box = self.cell.reshape(1, 9)
+        padded_atype = self.atype.clone()
+        padded_atype[:, -2:] = -1
+        nreal = self.natoms - 2
+        nreal_coord = nreal * 3
         # ``EnergyModel.forward`` exposes the Hessian under the ``"hessian"``
         # key (translated from ``energy_derv_r_derv_r``); it would KeyError if
         # the graph route failed to produce it.
         graph_out = model.forward(
-            self.coord.clone().requires_grad_(True), self.atype, box=box
+            self.coord.clone().requires_grad_(True), padded_atype, box=box
         )
         assert "hessian" in graph_out
         assert graph_out["hessian"].shape == (1, self.natoms * 3, self.natoms * 3)
 
+        unpadded_out = model.forward(
+            self.coord[:, :nreal].clone().requires_grad_(True),
+            self.atype[:, :nreal],
+            box=box,
+        )
+        torch.testing.assert_close(
+            graph_out["hessian"][:, :nreal_coord, :nreal_coord],
+            unpadded_out["hessian"],
+            rtol=1e-9,
+            atol=1e-9,
+        )
+        torch.testing.assert_close(
+            graph_out["hessian"][:, nreal_coord:],
+            torch.zeros_like(graph_out["hessian"][:, nreal_coord:]),
+            rtol=0,
+            atol=0,
+        )
+        torch.testing.assert_close(
+            graph_out["hessian"][:, :, nreal_coord:],
+            torch.zeros_like(graph_out["hessian"][:, :, nreal_coord:]),
+            rtol=0,
+            atol=0,
+        )
+
         # dense reference on the SAME weights via the escape hatch.
         ref_model = self._make_model()
         ref_model.eval()
@@ -272,7 +326,7 @@ def test_graph_native_hessian_matches_dense(self) -> None:
         ref_model.enable_hessian()
         assert _model_uses_graph_lower(ref_model) is False
         dense_out = ref_model.forward(
-            self.coord.clone().requires_grad_(True), self.atype, box=box
+            self.coord.clone().requires_grad_(True), padded_atype, box=box
         )
         torch.testing.assert_close(
             graph_out["hessian"], dense_out["hessian"], rtol=1e-9, atol=1e-9
@@ -281,6 +335,21 @@ def test_graph_native_hessian_matches_dense(self) -> None:
         h = graph_out["hessian"][0]
         torch.testing.assert_close(h, h.transpose(-1, -2), rtol=1e-9, atol=1e-9)
 
+    def test_hessian_model_rejects_a_ragged_axis(self) -> None:
+        """A flat node axis cannot encode per-frame Hessian pair dimensions."""
+        model = self._make_model().eval()
+        model.enable_hessian()
+        box = self.cell.reshape(1, 9)
+        n_node = torch.tensor([self.natoms], dtype=torch.int64, device=self.device)
+
+        with pytest.raises(NotImplementedError, match="rectangular atom axis"):
+            model.forward_ragged(
+                self.coord.reshape(-1, 3).requires_grad_(True),
+                self.atype.reshape(-1),
+                n_node,
+                box=box,
+            )
+
     def test_disable_graph_lower_escape_hatch(self) -> None:
         """``descriptor.disable_graph_lower()`` is the documented legacy-dense
         escape hatch: it flips ``uses_graph_lower()`` to ``False`` so the
diff --git a/source/tests/pt_expt/model/test_dpa4_native_spin.py b/source/tests/pt_expt/model/test_dpa4_native_spin.py
index 34119b09dd..35f82b3a95 100644
--- a/source/tests/pt_expt/model/test_dpa4_native_spin.py
+++ b/source/tests/pt_expt/model/test_dpa4_native_spin.py
@@ -190,6 +190,24 @@ def _energy(sp: np.ndarray) -> float:
             atol=1e-6,
         )
 
+    def test_hessian_uses_spin_conditioning(self) -> None:
+        """Coordinate Hessians retain the spin input of the energy surface."""
+        self.model.enable_hessian()
+        hessian = self.model.call_common(
+            self.coord.clone().requires_grad_(True),
+            self.atype,
+            self.box,
+            spin=self.spin,
+        )["energy_derv_r_derv_r"]
+        scaled_hessian = self.model.call_common(
+            self.coord.clone().requires_grad_(True),
+            self.atype,
+            self.box,
+            spin=2.0 * self.spin,
+        )["energy_derv_r_derv_r"]
+
+        assert torch.max(torch.abs(scaled_hessian - hessian)).item() > 1e-6
+
     def test_force_unchanged_by_spin_leaf_wiring(self) -> None:
         """``call_common`` WITHOUT ``spin`` has no ``energy_derv_r_mag`` key,
         and the spin-less forward is deterministic (the new ``spin is not
@@ -322,6 +340,18 @@ def test_forward_keys_and_mask(self) -> None:
         expect_mask = (self.atype == 0).unsqueeze(-1)
         torch.testing.assert_close(out["mask_mag"], expect_mask, rtol=0, atol=0)
 
+    def test_padding_is_not_magnetic_when_last_type_is(self) -> None:
+        """A phantom type never aliases the last magnetic type."""
+        config = copy.deepcopy(NATIVE_SPIN_CONFIG)
+        config["spin"]["use_spin"] = [False, True]
+        model = get_model(config).to(self.device)
+        atype = torch.tensor([[0, -1]], dtype=torch.int64, device=self.device)
+
+        torch.testing.assert_close(
+            model._spin_active_mask(atype),
+            torch.zeros((1, 2, 1), dtype=torch.bool, device=self.device),
+        )
+
     def test_force_mag_zero_on_non_spin_types(self) -> None:
         """Non-spin-type (``atype==1``) rows of ``force_mag`` are exactly
         zero, even though ``self.spin`` feeds nonzero noise there -- the
@@ -342,6 +372,37 @@ def test_atomic_virial_key_present_when_requested(self) -> None:
         assert "atom_virial" in out
         assert out["atom_virial"].shape == (self.nf, self.nloc, 9)
 
+    def test_ragged_matches_rectangular_mixed_size(self) -> None:
+        """Flat and padded node layouts preserve the same spin predictions."""
+        coord = self.coord.repeat(2, 1, 1)
+        atype = self.atype.repeat(2, 1)
+        spin = self.spin.repeat(2, 1, 1)
+        box = self.box.repeat(2, 1)
+        atype[1, -2:] = -1
+        coord[1, -2:] = 0.0
+        spin[1, -2:] = 0.0
+        real_atom = atype >= 0
+        n_node = real_atom.sum(dim=1)
+
+        rectangular = self.model(coord, atype, spin, box=box)
+        ragged = self.model.forward_ragged(
+            coord[real_atom],
+            atype[real_atom],
+            n_node,
+            spin[real_atom],
+            box=box,
+        )
+
+        torch.testing.assert_close(rectangular["mask"].bool(), real_atom)
+        torch.testing.assert_close(
+            ragged["mask"].bool(), torch.ones_like(ragged["mask"], dtype=torch.bool)
+        )
+        torch.testing.assert_close(ragged["n_node"], n_node)
+        for key in ("energy", "virial"):
+            torch.testing.assert_close(ragged[key], rectangular[key])
+        for key in ("atom_energy", "force", "force_mag", "mask_mag"):
+            torch.testing.assert_close(ragged[key], rectangular[key][real_atom])
+
 
 def _pt_native_spin_model(seed: int = 3):
     """Build the pt (reference) ``SeZMNativeSpinModel``, jittered."""
@@ -736,7 +797,7 @@ class TestDPA4NativeSpinCompiledTraining:
     def test_eager_eval_and_compiled_magnetic_force_backward(
         self, monkeypatch: pytest.MonkeyPatch
     ) -> None:
-        """Keep eval eager while retaining compiled magnetic gradients."""
+        """Keep eval eager and compiled mixed-size magnetic gradients correct."""
         from deepmd.pt_expt.train.training import (
             _CompiledModel,
             _get_model_structure_key,
@@ -765,22 +826,53 @@ def test_eager_eval_and_compiled_magnetic_force_backward(
                 ]
             ],
             dtype=torch.float64,
-        )
-        atype = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.int64)
-        spin = torch.arange(18, dtype=torch.float64).reshape(1, 6, 3) / 20.0 + 0.1
-        box = (8.0 * torch.eye(3, dtype=torch.float64)).reshape(1, 9)
+        ).repeat(2, 1, 1)
+        atype = torch.tensor([[0, 0, 0, 1, 1, 1]], dtype=torch.int64).repeat(2, 1)
+        spin = (
+            torch.arange(18, dtype=torch.float64).reshape(1, 6, 3) / 20.0 + 0.1
+        ).repeat(2, 1, 1)
+        box = (8.0 * torch.eye(3, dtype=torch.float64)).reshape(1, 9).repeat(2, 1)
+        atype[1, -2:] = -1
+        coord[1, -2:] = 0.0
+        spin[1, -2:] = 0.0
 
         eager_result = compiled.eval()(coord, atype, box=box, spin=spin)
+        real_atom = atype >= 0
+        n_node = real_atom.sum(dim=1)
+        eager_ragged = compiled.forward_ragged(
+            coord[real_atom],
+            atype[real_atom],
+            n_node,
+            box=box,
+            spin=spin[real_atom],
+        )
         assert compiled._compiled_lower_by_mode == {}
-        assert eager_result["force_mag"].shape == (1, 6, 3)
+        assert eager_result["force_mag"].shape == (2, 6, 3)
+        torch.testing.assert_close(eager_ragged["n_node"], n_node)
 
         result = compiled.train()(coord, atype, box=box, spin=spin)
 
-        assert result["force"].shape == (1, 6, 3)
-        assert result["force_mag"].shape == (1, 6, 3)
-        assert result["mask_mag"].shape == (1, 6, 1)
+        assert result["force"].shape == (2, 6, 3)
+        assert result["force_mag"].shape == (2, 6, 3)
+        assert result["mask_mag"].shape == (2, 6, 1)
+        assert torch.all(result["force"][1, -2:] == 0.0)
+        assert torch.all(result["force_mag"][1, -2:] == 0.0)
+        assert torch.all(~result["mask_mag"][1, -2:])
         assert set(compiled._compiled_lower_by_mode) == {True}
-        result["force_mag"].square().sum().backward()
+
+        ragged = compiled.forward_ragged(
+            coord[real_atom],
+            atype[real_atom],
+            n_node,
+            box=box,
+            spin=spin[real_atom],
+        )
+        torch.testing.assert_close(ragged["n_node"], n_node)
+        for key in ("energy", "virial"):
+            torch.testing.assert_close(ragged[key], result[key])
+        for key in ("atom_energy", "force", "force_mag", "mask", "mask_mag"):
+            torch.testing.assert_close(ragged[key], result[key][real_atom])
+        ragged["force_mag"].square().sum().backward()
 
         spin_parameters = [
             parameter
@@ -1203,6 +1295,7 @@ def test_allow_missing_label_data_requirement(
         spin_req = next(rr for rr in reqs if rr.key == "spin")
         assert spin_req.must is expected_must
         assert spin_req.default == 0.0
+        assert spin_req.source_policy == ("default" if allow_missing else "tracked")
 
 
 class TestPublicBaseModelRoundTrip:
diff --git a/source/tests/pt_expt/model/test_get_model_dpa4.py b/source/tests/pt_expt/model/test_get_model_dpa4.py
index d0990c0aa1..bbb59af1f0 100644
--- a/source/tests/pt_expt/model/test_get_model_dpa4.py
+++ b/source/tests/pt_expt/model/test_get_model_dpa4.py
@@ -92,6 +92,44 @@ def test_get_model_normalized_config(self) -> None:
         self.assertEqual(ret["energy"].shape, (1, 1))
         self.assertEqual(ret["force"].shape, (1, 5, 3))
 
+    def test_graph_hessian_preserves_charge_spin(self) -> None:
+        """The Hessian is evaluated on the explicitly conditioned energy surface."""
+        config = _make_raw_model_config()
+        config["descriptor"]["add_chg_spin_ebd"] = True
+        model = get_model(config).to(self.device).eval()
+        model.enable_hessian()
+        coord = torch.tensor(
+            [[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]],
+            dtype=torch.float64,
+            device=self.device,
+            requires_grad=True,
+        )
+        atype = torch.tensor([[0, 1]], dtype=torch.int64, device=self.device)
+        box = (5.0 * torch.eye(3, dtype=torch.float64, device=self.device)).reshape(
+            1, 9
+        )
+        charge_spin = torch.tensor(
+            [[0.0, 1.0]], dtype=torch.float64, device=self.device
+        )
+
+        graph = model.call_common(coord, atype, box, charge_spin=charge_spin)
+        dense = model.call_common(
+            coord,
+            atype,
+            box,
+            charge_spin=charge_spin,
+            neighbor_graph_method="legacy",
+        )
+
+        graph_hessian = graph["energy_derv_r_derv_r"]
+        self.assertEqual(graph_hessian.shape, (1, 1, 6, 6))
+        torch.testing.assert_close(
+            graph_hessian,
+            dense["energy_derv_r_derv_r"],
+            rtol=1e-10,
+            atol=1e-10,
+        )
+
     def test_get_model_type_aliases(self) -> None:
         """All model-type aliases route to the SeZM path."""
         for alias in ("dpa4", "DPA4", "sezm", "SeZM"):
diff --git a/source/tests/pt_expt/model/test_spin_ener_model.py b/source/tests/pt_expt/model/test_spin_ener_model.py
index 1f600934b5..eaa5b33100 100644
--- a/source/tests/pt_expt/model/test_spin_ener_model.py
+++ b/source/tests/pt_expt/model/test_spin_ener_model.py
@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
+from unittest.mock import (
+    patch,
+)
 
 import numpy as np
 import torch
@@ -143,6 +146,73 @@ def test_output_shapes(self) -> None:
         self.assertEqual(result["mask_mag"].shape, (1, natoms, 1))
         self.assertEqual(result["virial"].shape, (1, 9))
 
+    def test_mixed_size_padding_mask_is_public(self) -> None:
+        """Virtual-atom spin keeps rectangular padding out of every loss."""
+        model = _make_model()
+        generator = torch.Generator(device="cpu").manual_seed(GLOBAL_SEED)
+        cell = 8.0 * torch.eye(3, dtype=dtype)
+        coord = torch.rand((2, 6, 3), dtype=dtype, generator=generator)
+        coord = torch.matmul(coord, cell)
+        atype = torch.tensor(
+            [[0, 0, 1, 0, 1, 1], [0, 0, 1, 0, -1, -1]], dtype=torch.int64
+        )
+        spin = torch.rand((2, 6, 3), dtype=dtype, generator=generator)
+        coord[1, -2:] = 0.0
+        spin[1, -2:] = 0.0
+
+        result = eval_model(model, coord, cell.repeat(2, 1, 1), atype, spin)
+        real_atom = atype.to(env.DEVICE) >= 0
+
+        torch.testing.assert_close(result["mask"].bool(), real_atom)
+        self.assertTrue(torch.all(~result["mask_mag"][~real_atom]))
+        self.assertTrue(torch.all(result["atom_energy"][~real_atom] == 0.0))
+        self.assertTrue(torch.all(result["force"][~real_atom] == 0.0))
+        self.assertTrue(torch.all(result["force_mag"][~real_atom] == 0.0))
+
+    def test_compiled_mixed_size_rectangular_matches_eager(self) -> None:
+        """Dense compilation preserves the virtual-spin rectangular contract."""
+        from deepmd.pt_expt.train.training import (
+            _CompiledModel,
+            _get_model_structure_key,
+        )
+
+        model = _make_model()
+        compiled = _CompiledModel(model, _get_model_structure_key(model)).train()
+        generator = torch.Generator(device="cpu").manual_seed(GLOBAL_SEED)
+        coord = torch.rand((2, 6, 3), dtype=dtype, generator=generator).to(env.DEVICE)
+        atype = torch.tensor(
+            [[0, 0, 1, 0, 1, 1], [0, 0, 1, 0, -1, -1]],
+            dtype=torch.int64,
+            device=env.DEVICE,
+        )
+        spin = torch.rand((2, 6, 3), dtype=dtype, generator=generator).to(env.DEVICE)
+        box = (8.0 * torch.eye(3, dtype=dtype, device=env.DEVICE)).repeat(2, 1, 1)
+        coord[1, -2:] = 0.0
+        spin[1, -2:] = 0.0
+
+        eager = model(coord.clone().requires_grad_(True), atype, spin, box=box)
+        with patch.object(torch, "compile", side_effect=lambda module, **_: module):
+            actual = compiled(coord, atype, box=box, spin=spin)
+
+        for key in (
+            "atom_energy",
+            "energy",
+            "force",
+            "force_mag",
+            "virial",
+            "mask",
+            "mask_mag",
+        ):
+            torch.testing.assert_close(actual[key], eager[key], rtol=1e-10, atol=1e-10)
+        self.assertGreater(float(actual["force_mag"].detach().abs().max()), 0.0)
+        actual["force_mag"].square().sum().backward()
+        self.assertTrue(
+            any(
+                parameter.grad is not None and torch.any(parameter.grad != 0)
+                for parameter in model.parameters()
+            )
+        )
+
 
 class TestSpinEnerModelSerialize(unittest.TestCase):
     def test_serialize_deserialize(self) -> None:
diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py
index 641804b1cd..84c053279a 100644
--- a/source/tests/pt_expt/test_lmdb_training.py
+++ b/source/tests/pt_expt/test_lmdb_training.py
@@ -21,6 +21,7 @@
 import msgpack
 import numpy as np
 
+from deepmd.dpmodel.utils import lmdb_data as lmdb_data_module
 from deepmd.dpmodel.utils.batch import (
     normalize_batch,
     split_batch,
@@ -31,6 +32,12 @@
 from deepmd.pt_expt.entrypoints.main import (
     get_trainer,
 )
+from deepmd.pt_expt.loss import (
+    EnergyLoss,
+)
+from deepmd.pt_expt.train.training import (
+    Trainer,
+)
 from deepmd.pt_expt.utils.lmdb_dataset import (
     LmdbDataSystem,
 )
@@ -47,6 +54,10 @@
     DataRequirementItem,
 )
 
+from .compile_utils import (
+    REQUIRES_SUPPORTED_COMPILE,
+)
+
 
 def _encode_array(arr: np.ndarray) -> dict:
     return {
@@ -58,11 +69,11 @@ def _encode_array(arr: np.ndarray) -> dict:
     }
 
 
-def _make_frame(natoms: int, seed: int) -> dict:
+def _make_frame(natoms: int, seed: int, *, include_spin: bool = False) -> dict:
     """Synthetic LMDB frame matching the on-disk schema used by LmdbDataReader."""
     rng = np.random.RandomState(seed)
     half = natoms // 2
-    return {
+    frame = {
         "atom_numbs": [half, natoms - half],
         "atom_names": ["O", "H"],
         "atom_types": _encode_array(
@@ -74,6 +85,14 @@ def _make_frame(natoms: int, seed: int) -> dict:
         "energies": _encode_array(np.array(rng.randn(), dtype=np.float64)),
         "forces": _encode_array(rng.randn(natoms, 3).astype(np.float64)),
     }
+    if include_spin:
+        spin = rng.randn(natoms, 3).astype(np.float64)
+        spin[half:] = 0.0
+        force_mag = rng.randn(natoms, 3).astype(np.float64)
+        force_mag[half:] = 0.0
+        frame["spin"] = _encode_array(spin)
+        frame["force_mag"] = _encode_array(force_mag)
+    return frame
 
 
 def _create_test_lmdb(path: str, nframes: int, natoms: int) -> None:
@@ -124,7 +143,7 @@ def _create_partially_labeled_lmdb(path: str) -> None:
     env.close()
 
 
-def _create_mixed_nloc_test_lmdb(path: str) -> None:
+def _create_mixed_nloc_test_lmdb(path: str, *, include_spin: bool = False) -> None:
     """Write an LMDB with five six-atom and five nine-atom frames."""
     frame_nlocs = [6] * 5 + [9] * 5
     env = lmdb.open(path, map_size=10 * 1024 * 1024)
@@ -146,7 +165,10 @@ def _create_mixed_nloc_test_lmdb(path: str) -> None:
             key = format(index, fmt).encode()
             txn.put(
                 key,
-                msgpack.packb(_make_frame(nloc, index), use_bin_type=True),
+                msgpack.packb(
+                    _make_frame(nloc, index, include_spin=include_spin),
+                    use_bin_type=True,
+                ),
             )
     env.close()
 
@@ -262,6 +284,27 @@ def test_data_requirements_freeze_after_first_read(self) -> None:
         ):
             ds.add_data_requirements([DataRequirementItem("late_label", ndof=1)])
 
+    def test_missing_mandatory_model_input_is_rejected(self) -> None:
+        """A required model input cannot silently become a zero-filled batch."""
+        ds = LmdbDataSystem(
+            lmdb_path=self.lmdb_path,
+            type_map=["O", "H"],
+            batch_size=2,
+            seed=0,
+            num_workers=0,
+        )
+        ds.add_data_requirements(
+            [DataRequirementItem("spin", 3, atomic=True, must=True)]
+        )
+        try:
+            with self.assertRaisesRegex(
+                RuntimeError,
+                r"spin.*frame \d+.*test\.lmdb.*field is absent",
+            ):
+                ds.get_batch()
+        finally:
+            ds.close()
+
     def test_get_batch_iterates_past_end(self) -> None:
         """get_batch reseeds the sampler at the end of an epoch."""
         ds = LmdbDataSystem(
@@ -327,6 +370,118 @@ def test_add_data_requirements_passthrough(self) -> None:
         self.assertIn("energy", batch)
         self.assertIn("find_energy", batch)
 
+    def test_default_atom_pref_does_not_partition_training_data(self) -> None:
+        """OC20-style unit defaults keep mixed atom_pref frames compatible."""
+        path = os.path.join(self.tmpdir, "default_atom_pref.lmdb")
+        _create_test_lmdb(path, nframes=8, natoms=6)
+        environment = lmdb.open(path, readonly=False, lock=False)
+        with environment.begin(write=True) as transaction:
+            for index in range(0, 8, 2):
+                key = format(index, "012d").encode()
+                frame = msgpack.unpackb(transaction.get(key), raw=False)
+                frame["atom_pref"] = _encode_array(np.full(6, 2.0))
+                transaction.put(key, msgpack.packb(frame, use_bin_type=True))
+        environment.close()
+
+        loss = EnergyLoss(
+            starter_learning_rate=0.002,
+            start_pref_e=20.0,
+            limit_pref_e=20.0,
+            start_pref_f=0.0,
+            limit_pref_f=0.0,
+            start_pref_pf=20.0,
+            limit_pref_pf=20.0,
+            start_pref_v=5.0,
+            limit_pref_v=5.0,
+            loss_func="mae",
+            f_use_norm=True,
+            use_default_pf=True,
+        )
+        with patch.object(
+            lmdb_data_module,
+            "_raw_frame_availability",
+            wraps=lmdb_data_module._raw_frame_availability,
+        ) as inspect_availability:
+            ds = LmdbDataSystem(
+                lmdb_path=path,
+                type_map=["O", "H"],
+                batch_size="mix:30000",
+                seed=0,
+            )
+            self.assertEqual(inspect_availability.call_count, 0)
+            ds.add_data_requirements(loss.label_requirement)
+
+        try:
+            self.assertLessEqual(inspect_availability.call_count, 8)
+            self.assertEqual(len(ds._stat_groups), 1)
+            batch = ds._reader.decode_batch([0, 1])
+            np.testing.assert_array_equal(batch["atom_pref"][0], np.full(18, 2.0))
+            np.testing.assert_array_equal(batch["atom_pref"][1], np.ones(18))
+            self.assertEqual(float(batch["find_atom_pref"]), 1.0)
+            model_output = {
+                "energy": np.zeros_like(batch["energy"]),
+                "force": np.zeros_like(batch["force"]),
+                "virial": np.zeros_like(batch["virial"]),
+                "atom_energy": np.zeros((2, 6, 1), dtype=batch["energy"].dtype),
+            }
+            loss_value, metrics = loss.call(0.002, 6, model_output, batch)
+            self.assertTrue(np.isfinite(loss_value))
+            self.assertIn("mae_pf", metrics)
+        finally:
+            ds.close()
+
+    def test_missing_force_disables_default_prefactor_force_loss(self) -> None:
+        """Default atom weights cannot supervise a frame without force."""
+        path = os.path.join(self.tmpdir, "partial_force_for_pf.lmdb")
+        _create_partially_labeled_lmdb(path)
+        loss = 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_pf=1.0,
+            limit_pref_pf=1.0,
+            use_default_pf=True,
+        )
+        ds = LmdbDataSystem(
+            lmdb_path=path,
+            type_map=["O", "H"],
+            batch_size=2,
+            seed=0,
+            num_workers=0,
+        )
+        ds.add_data_requirements(loss.label_requirement)
+        try:
+            observed: dict[float, float] = {}
+            for _ in range(2):
+                batch = ds.get_batch()
+                nframes, nloc = batch["atype"].shape
+                model_output = {
+                    "energy": np.zeros_like(batch.get("energy", np.zeros(nframes))),
+                    "force": np.zeros_like(batch["force"]),
+                    "virial": np.zeros((nframes, 9), dtype=batch["force"].dtype),
+                    "atom_energy": np.zeros(
+                        (nframes, nloc, 1),
+                        dtype=batch["force"].dtype,
+                    ),
+                }
+                loss_value, metrics = loss.call(
+                    1.0,
+                    nloc,
+                    model_output,
+                    batch,
+                )
+                find_force = float(batch["find_force"])
+                observed[find_force] = float(loss_value)
+                if find_force == 0.0:
+                    self.assertTrue(np.isnan(metrics["rmse_pf"]))
+
+            self.assertEqual(observed[0.0], 0.0)
+            self.assertGreater(observed[1.0], 0.0)
+        finally:
+            ds.close()
+
     def test_partial_labels_are_batched_by_availability(self) -> None:
         from deepmd.utils.data import (
             DataRequirementItem,
@@ -529,6 +684,32 @@ def test_numb_epoch_counts_passes_over_the_lmdb(self) -> None:
         # train.lmdb holds eight frames, read one frame per batch.
         self.assertEqual(trainer.num_steps, 2 * 8)
 
+    def test_numb_epoch_counts_mixed_nloc_batches(self) -> None:
+        """One epoch of a ``mix:N`` dataset is one pass over its padded batches.
+
+        The batch count of a mixed-nloc pass depends on how the shuffle groups
+        atom counts, so an epoch can only be measured on the sampler the
+        trainer will actually draw from, not on a nominal batch size.
+        """
+        config = self._make_lmdb_config()
+        del config["training"]["numb_steps"]
+        config["training"]["numb_epoch"] = 3.0
+        config["training"]["training_data"]["systems"] = self.mixed_lmdb_path
+        config["training"]["training_data"]["batch_size"] = "mix:27"
+        config = update_deepmd_input(config, warning=False)
+        config = normalize(config)
+
+        cwd = os.getcwd()
+        os.chdir(self.tmpdir)
+        try:
+            trainer = get_trainer(config)
+        finally:
+            os.chdir(cwd)
+
+        data = trainer.training_data
+        self.assertTrue(data._reader.mixed_nloc)
+        self.assertEqual(trainer.num_steps, 3 * data.nbatches[0])
+
     def test_training_closes_parallel_lmdb_pipeline(self) -> None:
         """Trainer shutdown releases spawned decoder processes."""
         config = self._make_lmdb_config(numb_steps=2)
@@ -567,6 +748,360 @@ def test_mixed_nloc_statistics_and_training(self) -> None:
             finally:
                 os.chdir(cwd)
 
+    def test_mix_batch_size_trains_on_padded_batches(self) -> None:
+        """``mix:N`` reaches the trainer and its padded batches train."""
+        config = self._make_lmdb_config(numb_steps=2)
+        config["model"]["data_stat_nbatch"] = 10
+        config["training"]["training_data"]["systems"] = self.mixed_lmdb_path
+        # The fixture holds five 6-atom and five 9-atom frames; this budget
+        # puts a batch boundary inside the atom-count-sorted run, so at least
+        # one batch spans both sizes and is padded.
+        config["training"]["training_data"]["batch_size"] = "mix:27"
+        config = update_deepmd_input(config, warning=False)
+        config = normalize(config)
+
+        with tempfile.TemporaryDirectory(dir=self.tmpdir) as run_dir:
+            cwd = os.getcwd()
+            os.chdir(run_dir)
+            try:
+                trainer = get_trainer(config)
+                data = trainer.training_data
+                self.assertTrue(data._reader.mixed_nloc)
+                # Statistics keep their own fixed-nloc view, so padding never
+                # reaches the per-type accumulators.
+                self.assertEqual(data.get_stat_nsystems(), 2)
+                for sys_idx in range(data.get_stat_nsystems()):
+                    stat = data.get_stat_batch(sys_idx)
+                    self.assertTrue((stat["atype"] >= 0).all())
+                padded = next(
+                    batch
+                    for batch in (data.get_batch() for _ in range(10))
+                    if (batch["atype"] < 0).any()
+                )
+                self.assertEqual(padded["atype"].shape[1], 9)
+                self.assertEqual(padded["coord"].shape[1], 9)
+                self.assertEqual(padded["force"].shape[1], 9)
+                np.testing.assert_array_equal(padded["force"][padded["atype"] < 0], 0.0)
+                # Per-frame atom counts stay real, not padded.
+                self.assertEqual(sorted(set(padded["natoms"][:, 0].tolist())), [6, 9])
+                trainer.run()
+            finally:
+                os.chdir(cwd)
+
+
+class TestRaggedTrainingBatches(unittest.TestCase):
+    """A model reading a flat node axis is fed one, with nothing padded."""
+
+    def setUp(self) -> None:
+        self.tmpdir = tempfile.mkdtemp()
+        self.lmdb_path = os.path.join(self.tmpdir, "mixed.lmdb")
+        self.spin_lmdb_path = os.path.join(self.tmpdir, "mixed-spin.lmdb")
+        _create_mixed_nloc_test_lmdb(self.lmdb_path)
+        _create_mixed_nloc_test_lmdb(self.spin_lmdb_path, include_spin=True)
+
+    def tearDown(self) -> None:
+        shutil.rmtree(self.tmpdir, ignore_errors=True)
+
+    def _config_for_model(
+        self,
+        model: dict,
+        loss: dict,
+        *,
+        lmdb_path: str | None = None,
+    ) -> dict:
+        config = {
+            "model": model,
+            "learning_rate": {
+                "type": "exp",
+                "decay_steps": 500,
+                "start_lr": 1e-3,
+                "stop_lr": 3.5e-8,
+            },
+            "loss": loss,
+            "training": {
+                "training_data": {
+                    "systems": self.lmdb_path if lmdb_path is None else lmdb_path,
+                    "batch_size": "mix:27",
+                },
+                "numb_steps": 2,
+                "seed": 10,
+                "disp_file": "lcurve.out",
+                "disp_freq": 1,
+                "save_freq": 100,
+            },
+        }
+        return normalize(update_deepmd_input(config, warning=False))
+
+    def _config(self, descriptor: dict) -> dict:
+        return self._config_for_model(
+            {
+                "type_map": ["O", "H"],
+                "descriptor": descriptor,
+                "fitting_net": {
+                    "neuron": [8, 8],
+                    "precision": "float64",
+                    "seed": 1,
+                },
+                "data_stat_nbatch": 1,
+            },
+            {
+                "type": "ener",
+                "start_pref_e": 0.02,
+                "limit_pref_e": 1,
+                "start_pref_f": 1000,
+                "limit_pref_f": 1,
+                "start_pref_v": 0,
+                "limit_pref_v": 0,
+            },
+        )
+
+    @staticmethod
+    def _dpa1() -> dict:
+        return {
+            "type": "dpa1",
+            "sel": 12,
+            "rcut_smth": 0.5,
+            "rcut": 3.0,
+            "neuron": [8, 16],
+            "axis_neuron": 4,
+            "attn_layer": 0,
+            "precision": "float64",
+            "seed": 1,
+        }
+
+    @staticmethod
+    def _se_e2_a() -> dict:
+        return {
+            "type": "se_e2_a",
+            "sel": [6, 12],
+            "rcut_smth": 0.5,
+            "rcut": 3.0,
+            "neuron": [8, 16],
+            "axis_neuron": 4,
+            "seed": 1,
+        }
+
+    @staticmethod
+    def _dpa4_native_spin_model() -> dict:
+        return {
+            "type": "dpa4",
+            "type_map": ["O", "H"],
+            "descriptor": {
+                "type": "dpa4",
+                "sel": 20,
+                "rcut": 3.0,
+                "channels": 8,
+                "n_radial": 4,
+                "lmax": 1,
+                "mmax": 1,
+                "n_blocks": 1,
+                "precision": "float64",
+                "seed": 1,
+            },
+            "fitting_net": {
+                "type": "dpa4_ener",
+                "neuron": [8],
+                "precision": "float64",
+                "seed": 1,
+            },
+            "spin": {"use_spin": [True, False], "scheme": "native"},
+            "data_stat_nbatch": 1,
+        }
+
+    def _virtual_spin_model(self) -> dict:
+        return {
+            "type_map": ["O", "H"],
+            "descriptor": self._se_e2_a(),
+            "fitting_net": {
+                "neuron": [8, 8],
+                "precision": "float64",
+                "seed": 1,
+            },
+            "spin": {
+                "use_spin": [True, False],
+                "virtual_scale": [0.314],
+            },
+            "data_stat_nbatch": 1,
+        }
+
+    @staticmethod
+    def _spin_loss() -> dict:
+        return {
+            "type": "ener_spin",
+            "start_pref_e": 0.02,
+            "limit_pref_e": 1,
+            "start_pref_fr": 1000,
+            "limit_pref_fr": 1,
+            "start_pref_fm": 1000,
+            "limit_pref_fm": 1,
+        }
+
+    def _run(self, descriptor: dict, *, compile: bool = False):
+        """Train two steps and return the trainer plus one drawn batch."""
+        config = self._config(descriptor)
+        config["training"]["enable_compile"] = compile
+        cwd = os.getcwd()
+        os.chdir(self.tmpdir)
+        try:
+            trainer = get_trainer(config)
+            batch = trainer.training_data.get_batch()
+            trainer.run()
+            return trainer, batch
+        finally:
+            os.chdir(cwd)
+
+    def _run_spin(self, model: dict):
+        """Train a spin model for two mixed-size steps and return one batch."""
+        config = self._config_for_model(
+            model,
+            self._spin_loss(),
+            lmdb_path=self.spin_lmdb_path,
+        )
+        cwd = os.getcwd()
+        os.chdir(self.tmpdir)
+        try:
+            trainer = get_trainer(config)
+            batch = trainer.training_data.get_batch()
+            trainer.run()
+            return trainer, batch
+        finally:
+            os.chdir(cwd)
+
+    def test_graph_model_trains_on_a_flat_node_axis(self) -> None:
+        """DPA1 reads a graph lower, so its batches carry no padded row."""
+        trainer, batch = self._run(self._dpa1())
+        self.assertTrue(trainer.training_data._reader.ragged_batches)
+        self.assertEqual(batch["coord"].ndim, 2)
+        self.assertEqual(batch["atype"].ndim, 1)
+        self.assertEqual(batch["coord"].shape[0], int(batch["n_node"].sum()))
+        self.assertTrue((batch["atype"] >= 0).all(), "nothing is padded")
+        # Frame-level fields keep their frame axis.
+        self.assertEqual(batch["energy"].shape[0], batch["n_node"].shape[0])
+
+    def test_loss_contract_can_require_rectangular_batches(self) -> None:
+        """Generalized-force and Hessian labels keep their rectangular axes."""
+        cases = {
+            "generalized force": {
+                "start_pref_gf": 1.0,
+                "limit_pref_gf": 1.0,
+                "numb_generalized_coord": 2,
+            },
+            "Hessian": {
+                "start_pref_h": 1.0,
+                "limit_pref_h": 1.0,
+            },
+        }
+        cwd = os.getcwd()
+        os.chdir(self.tmpdir)
+        try:
+            for name, loss_params in cases.items():
+                with self.subTest(loss=name):
+                    config = self._config(self._dpa1())
+                    config["loss"].update(loss_params)
+                    trainer = get_trainer(config)
+                    self.assertFalse(trainer.training_data._reader.ragged_batches)
+        finally:
+            os.chdir(cwd)
+
+    def test_non_mixing_rule_keeps_the_rectangular_layout(self) -> None:
+        """Model capability alone does not change established LMDB batches."""
+        config = self._config(self._dpa1())
+        config["training"]["training_data"]["batch_size"] = 2
+        cwd = os.getcwd()
+        os.chdir(self.tmpdir)
+        try:
+            trainer = get_trainer(config)
+            batch = trainer.training_data.get_batch()
+        finally:
+            os.chdir(cwd)
+
+        self.assertFalse(trainer.training_data._reader.ragged_batches)
+        self.assertEqual(batch["coord"].ndim, 3)
+        self.assertEqual(batch["atype"].ndim, 2)
+        self.assertNotIn("n_node", batch)
+
+    @REQUIRES_SUPPORTED_COMPILE
+    def test_compiled_graph_model_trains_on_a_flat_node_axis(self) -> None:
+        """The compiled lower reads the flat axis too, so compiling changes nothing.
+
+        Its trace is taken on a rectangular system, which would bake in
+        ``N == nframes * nloc`` were the frame, node and edge counts not kept
+        as independent symbols.
+        """
+        trainer, batch = self._run(self._dpa1(), compile=True)
+        self.assertEqual(
+            type(trainer.wrapper.model["Default"]).__name__, "_CompiledModel"
+        )
+        self.assertTrue((batch["atype"] >= 0).all())
+        self.assertEqual(batch["coord"].shape[0], int(batch["n_node"].sum()))
+
+    def test_native_spin_trains_on_a_flat_node_axis(self) -> None:
+        """Native spin carries moments and magnetic labels through ragged training."""
+        trainer, batch = self._run_spin(self._dpa4_native_spin_model())
+
+        self.assertTrue(trainer.training_data._reader.ragged_batches)
+        self.assertEqual(batch["coord"].ndim, 2)
+        self.assertEqual(batch["spin"].shape, batch["coord"].shape)
+        self.assertEqual(batch["force_mag"].shape, batch["coord"].shape)
+        self.assertEqual(batch["coord"].shape[0], int(batch["n_node"].sum()))
+        self.assertTrue((batch["atype"] >= 0).all())
+
+    def test_virtual_spin_trains_on_a_rectangular_node_axis(self) -> None:
+        """Virtual-atom spin keeps mixed-size training rectangular and masked."""
+        trainer, batch = self._run_spin(self._virtual_spin_model())
+
+        self.assertFalse(trainer.training_data._reader.ragged_batches)
+        self.assertEqual(batch["coord"].ndim, 3)
+        self.assertEqual(batch["spin"].shape, batch["coord"].shape)
+        self.assertEqual(batch["force_mag"].shape, batch["coord"].shape)
+        self.assertNotIn("n_node", batch)
+        self.assertTrue((batch["atype"] < 0).any())
+
+    def test_dense_model_keeps_padded_batches(self) -> None:
+        """se_e2_a reads a rectangular node axis and must still be padded."""
+        trainer, batch = self._run(self._se_e2_a())
+        self.assertFalse(trainer.training_data._reader.ragged_batches)
+        self.assertEqual(batch["coord"].ndim, 3)
+        self.assertEqual(batch["atype"].ndim, 2)
+        self.assertNotIn("n_node", batch)
+
+    def test_layout_is_settled_before_the_pass_is_counted(self) -> None:
+        """The run length must count the packing training actually uses.
+
+        Under ``mix:N`` the layout decides where the sampler cuts a batch, and
+        the sampler materializes a pass the moment its length is asked for.
+        Settling the layout afterwards would count one packing and train on
+        another, and would serve the first epoch from the counted one.
+
+        The pass is cached once built, so both orderings report the same count
+        after the fact. What distinguishes them is the layout in force at the
+        moment the count is taken, which is what this observes.
+        """
+        config = self._config(self._dpa1())
+        del config["training"]["numb_steps"]
+        config["training"]["numb_epoch"] = 1.0
+
+        observed: dict[str, bool] = {}
+        original = Trainer._epoch_length
+
+        def record_layout(trainer_self, model_key: str) -> int:
+            observed["ragged_when_counted"] = (
+                trainer_self.training_data._reader.ragged_batches
+            )
+            return original(trainer_self, model_key)
+
+        cwd = os.getcwd()
+        os.chdir(self.tmpdir)
+        try:
+            with patch.object(Trainer, "_epoch_length", record_layout):
+                trainer = get_trainer(config)
+        finally:
+            os.chdir(cwd)
+
+        # DPA1 reads a flat node axis, so training runs on ragged batches.
+        self.assertTrue(trainer.training_data._reader.ragged_batches)
+        self.assertTrue(observed["ragged_when_counted"])
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py
index 45eb76fa03..31cbb61df7 100644
--- a/source/tests/pt_expt/test_training.py
+++ b/source/tests/pt_expt/test_training.py
@@ -999,6 +999,7 @@ def has_chg_spin_ebd(self) -> bool:
         self.assertEqual(fparam_req.key, "fparam")
         self.assertEqual(fparam_req.ndof, 2)
         self.assertFalse(fparam_req.must)
+        self.assertEqual(fparam_req.source_policy, "default")
         # default is the model's default_fparam, not 0.0
         self.assertNotIsInstance(fparam_req.default, float)
         import numpy as np
@@ -1036,6 +1037,7 @@ def has_chg_spin_ebd(self) -> bool:
         self.assertEqual(fparam_req.key, "fparam")
         self.assertTrue(fparam_req.must)
         self.assertEqual(fparam_req.default, 0.0)
+        self.assertEqual(fparam_req.source_policy, "tracked")
 
 
 class TestRestart(unittest.TestCase):
diff --git a/source/tests/pt_expt/utils/test_exclusion_mask.py b/source/tests/pt_expt/utils/test_exclusion_mask.py
index cc0671c117..9ea89c12ac 100644
--- a/source/tests/pt_expt/utils/test_exclusion_mask.py
+++ b/source/tests/pt_expt/utils/test_exclusion_mask.py
@@ -1,5 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import unittest
+from unittest.mock import (
+    patch,
+)
 
 import numpy as np
 import torch
@@ -38,6 +41,13 @@ def test_build_type_exclude_mask(self) -> None:
         des = AtomExcludeMask(nt, exclude_types=exclude_types)
         mask = des.build_type_exclude_mask(torch.as_tensor(atype, device=env.DEVICE))
         np.testing.assert_equal(mask.detach().cpu().numpy(), expected_mask)
+        with patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            numpy_mask = des.build_type_exclude_mask(atype)
+        np.testing.assert_equal(numpy_mask, expected_mask)
 
     def test_type_mask_is_buffer(self) -> None:
         des = AtomExcludeMask(3, exclude_types=[0])
@@ -66,6 +76,29 @@ def test_build_type_exclude_mask(self) -> None:
             torch.as_tensor(self.atype_ext, device=env.DEVICE),
         )
         np.testing.assert_equal(mask.detach().cpu().numpy(), expected_mask)
+        with patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            numpy_mask = des.build_type_exclude_mask(self.nlist, self.atype_ext)
+        np.testing.assert_equal(
+            numpy_mask,
+            expected_mask,
+        )
+
+    def test_build_edge_exclude_mask_with_numpy_inputs(self) -> None:
+        des = PairExcludeMask(self.nt, exclude_types=[[0, 1]])
+        edge_index = np.array([[0, 1, 2, 3], [1, 0, 3, 2]], dtype=np.int64)
+        atype = np.array([0, 1, 0, 0], dtype=np.int32)
+
+        with patch.object(
+            torch.Tensor,
+            "numpy",
+            side_effect=TypeError("direct conversion is unavailable"),
+        ):
+            numpy_mask = des.build_edge_exclude_mask(edge_index, atype)
+        np.testing.assert_equal(numpy_mask, np.array([0, 0, 1, 1], dtype=np.int32))
 
     def test_type_mask_is_buffer(self) -> None:
         des = PairExcludeMask(self.nt, exclude_types=[[0, 1]])
diff --git a/source/tests/pt_expt/utils/test_nv_matrix_decode.py b/source/tests/pt_expt/utils/test_nv_matrix_decode.py
index 5eccfa8277..43335230ff 100644
--- a/source/tests/pt_expt/utils/test_nv_matrix_decode.py
+++ b/source/tests/pt_expt/utils/test_nv_matrix_decode.py
@@ -1,16 +1,26 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
-"""CPU unit tests for the nv dense-matrix -> (i, j, S) decode.
+"""Parts of the nv path that run without a GPU.
 
-The GPU ``neighbor_list`` search in ``build_neighbor_graph_nv`` is CUDA-only
-and stays behind the opt-in CUDA suite (test_nv_graph_builder.py); the decode
-(``nv_matrix_to_ijs``) is pure torch index arithmetic, so its regression-prone
-parts (``// max_neighbors``, ``% nloc``, frame isolation, slot-validity mask)
-are pinned here on the default CI with synthetic inputs.
+Performance work on the nv builder is CUDA-bound and stays behind the opt-in
+suite (test_nv_graph_builder.py), but two things are checkable anywhere
+nvalchemiops imports. The decode (``nv_matrix_to_ijs``) is pure torch index
+arithmetic, so its regression-prone parts (``// max_neighbors``, ``% nloc``,
+frame isolation, slot-validity mask, and the lift back onto a padded batch)
+are pinned here with synthetic inputs. The builder's agreement with the dense
+reference builder is pinned here too, since the search itself is device
+agnostic even though it is only worth running on a GPU.
 """
 
 import numpy as np
+import pytest
 import torch
 
+from deepmd.dpmodel.utils.neighbor_graph import (
+    build_neighbor_graph,
+)
+from deepmd.pt.utils.nv_nlist import (
+    is_nv_available,
+)
 from deepmd.pt_expt.utils.nv_graph_builder import (
     nv_matrix_to_ijs,
 )
@@ -66,6 +76,65 @@ def test_empty_no_neighbors(self) -> None:
         assert i.shape == (0,) and j.shape == (0,)
         assert s.shape == (0, 3) and f.shape == (0,)
 
+    def test_search_over_real_atoms_only(self) -> None:
+        """A search restricted to the real atoms still reports batch indices.
+
+        A mixed-nloc batch is padded to a rectangular width, and the search is
+        given only the real slots, so its output is indexed on those. The decode
+        must lift both endpoints back onto the batch, or the frame and local
+        indices it derives would be meaningless.
+
+        Batch: nf=2, nloc=3, real atoms at flat positions 0, 1, 3, 4, 5 -- that
+        is, frame 0 holds two atoms and frame 1 holds three.
+        """
+        node_index = torch.tensor([0, 1, 3, 4, 5], dtype=torch.int64)
+        # Search indices 0..4 address those five atoms in order.
+        neighbor_matrix = torch.tensor(
+            [
+                [1, 9],  # search 0 (batch 0, frame 0 local 0) -> search 1
+                [0, 9],  # search 1 (batch 1, frame 0 local 1) -> search 0
+                [3, 4],  # search 2 (batch 3, frame 1 local 0) -> search 3, 4
+                [2, 9],  # search 3 (batch 4, frame 1 local 1) -> search 2
+                [2, 9],  # search 4 (batch 5, frame 1 local 2) -> search 2
+            ],
+            dtype=torch.int32,
+        )
+        num_neighbors = torch.tensor([1, 1, 2, 1, 1], dtype=torch.int32)
+        shifts = torch.zeros((5, 2, 3), dtype=torch.int32)
+
+        i, j, s, f = nv_matrix_to_ijs(
+            neighbor_matrix, num_neighbors, shifts, 3, node_index=node_index
+        )
+        assert _edge_set(i, j, s, f) == {
+            (0, 0, 1, (0, 0, 0)),
+            (0, 1, 0, (0, 0, 0)),
+            (1, 0, 1, (0, 0, 0)),
+            (1, 0, 2, (0, 0, 0)),
+            (1, 1, 0, (0, 0, 0)),
+            (1, 2, 0, (0, 0, 0)),
+        }
+
+    def test_identity_node_index_matches_the_unrestricted_decode(self) -> None:
+        """Searching every slot is the case where the two indexings coincide."""
+        rng = np.random.default_rng(3)
+        nf, nloc, mn = 2, 4, 3
+        total = nf * nloc
+        num = torch.from_numpy(rng.integers(0, mn + 1, size=total)).to(torch.int32)
+        mat = torch.zeros((total, mn), dtype=torch.int32)
+        for dst in range(total):
+            frame = dst // nloc
+            mat[dst] = torch.from_numpy(
+                rng.integers(frame * nloc, (frame + 1) * nloc, size=mn)
+            ).to(torch.int32)
+        shf = torch.from_numpy(rng.integers(-1, 2, size=(total, mn, 3))).to(torch.int32)
+
+        plain = nv_matrix_to_ijs(mat, num, shf, nloc)
+        lifted = nv_matrix_to_ijs(
+            mat, num, shf, nloc, node_index=torch.arange(total, dtype=torch.int64)
+        )
+        for got, expected in zip(lifted, plain, strict=True):
+            torch.testing.assert_close(got, expected)
+
     def test_random_vs_oracle(self) -> None:
         """Random matrices match a brute-force python oracle."""
         rng = np.random.default_rng(11)
@@ -97,3 +166,81 @@ def test_random_vs_oracle(self) -> None:
             nloc,
         )
         assert _edge_set(i, j, s, f) == oracle
+
+
+def _graph_edges(graph) -> set:
+    """Edges as (src, dst, rounded edge_vec), so two builders can be compared."""
+    keep = graph.edge_mask
+    return {
+        (int(s), int(d), *(round(float(x), 8) for x in v))
+        for s, d, v in zip(
+            graph.edge_index[0][keep],
+            graph.edge_index[1][keep],
+            graph.edge_vec[keep],
+            strict=True,
+        )
+    }
+
+
+@pytest.mark.skipif(not is_nv_available(), reason="nvalchemi-toolkit-ops not installed")
+class TestNvBuilderOnPaddedBatches:
+    """The nv builder withholds phantom atoms from the search itself.
+
+    Its cost grows with the square of the frame width, so a mixed-nloc batch
+    hands it the real atoms alone and lifts the resulting indices back onto the
+    padded batch. Both halves of that have to be exact.
+    """
+
+    @staticmethod
+    def _batch(nlocs, boxlen=9.0, seed=0):
+        rng = np.random.default_rng(seed)
+        width = max(nlocs)
+        coord = np.zeros((len(nlocs), width, 3))
+        atype = np.full((len(nlocs), width), -1, dtype=np.int64)
+        for frame, nloc in enumerate(nlocs):
+            coord[frame, :nloc] = rng.uniform(0.0, boxlen, (nloc, 3))
+            atype[frame, :nloc] = rng.integers(0, 2, nloc)
+        return (
+            torch.tensor(coord),
+            torch.tensor(atype),
+            torch.tensor(np.tile(np.eye(3)[None] * boxlen, (len(nlocs), 1, 1))),
+        )
+
+    @pytest.mark.parametrize(
+        ("nlocs", "boxlen"),
+        [((6, 6, 6), 9.0), ((4, 7, 3), 9.0), ((5, 9, 2), 6.0)],
+        ids=["uniform", "mixed-nloc", "mixed-nloc-dense-box"],
+    )
+    def test_matches_the_dense_builder(self, nlocs, boxlen) -> None:
+        from deepmd.pt_expt.utils.nv_graph_builder import (
+            build_neighbor_graph_nv,
+        )
+
+        coord, atype, box = self._batch(nlocs, boxlen)
+        dense = build_neighbor_graph(coord, atype, box, 4.0)
+        nv = build_neighbor_graph_nv(coord, atype, box, 4.0)
+        assert _graph_edges(nv) == _graph_edges(dense)
+        np.testing.assert_array_equal(nv.n_node.numpy(), dense.n_node.numpy())
+
+    def test_search_enumerates_the_real_atoms_only(self, monkeypatch) -> None:
+        """Padding must not reach the search, only its output indexing."""
+        from deepmd.pt_expt.utils import (
+            nv_graph_builder,
+        )
+
+        nlocs = (4, 7, 3)
+        coord, atype, box = self._batch(nlocs)
+        searched: list[int] = []
+        original = nv_graph_builder.nv_search_matrix
+
+        def spy(*args, **kwargs):
+            result = original(*args, **kwargs)
+            searched.append(int(result[2].shape[0]))
+            return result
+
+        monkeypatch.setattr(nv_graph_builder, "nv_search_matrix", spy)
+        nv_graph_builder.build_neighbor_graph_nv(coord, atype, box, 4.0)
+        assert searched == [sum(nlocs)], (
+            f"the search saw {searched} atoms; the batch holds {sum(nlocs)} real "
+            f"atoms padded to {len(nlocs) * max(nlocs)} slots"
+        )