fix(data): preserve LMDB label availability across batches

deepmd/pt/train/validation.py

@@ -284,10 +284,10 @@ def __init__(
             self.checkpoint_dir.mkdir(parents=True, exist_ok=True)
             self._initialize_best_checkpoints(restart_training=restart_training)
 
-        # Lazily-populated full test snapshot for LMDB validation. Mixed-nloc
-        # LMDB datasets cannot be stacked as a single (nframes, natoms*3)
-        # tensor, so we materialize frames grouped by nloc the first time
-        # full validation runs and reuse the snapshot on subsequent calls.
+        # Lazily-populated full test snapshot for LMDB validation. Frames are
+        # evaluated in atom-count and label-availability groups so each stack
+        # has consistent shapes and scalar find_* flags. Reuse the decoded
+        # snapshot on subsequent validation calls.
         self._lmdb_test_data: LmdbTestData | None = None
 
     def should_run(self, display_step: int) -> bool:
@@ -430,20 +430,25 @@ def _iter_validation_data_systems(self) -> Iterator[Any]:
           and we forward its underlying ``DeepmdData`` instance.
         - For ``LmdbDataset`` validation data, we lazily materialize a
           :class:`LmdbTestData` snapshot (cached across calls) and yield one
-          :class:`LmdbTestDataNlocView` per ``nloc`` group, so mixed-nloc
-          frames can be stacked and evaluated group by group.
+          :class:`LmdbTestDataNlocView` per atom-count and label-availability
+          group. This keeps scalar ``find_*`` flags valid while excluding
+          default-filled labels from metrics.
         """
         validation_data = self.validation_data
         if isinstance(validation_data, LmdbDataset):
             lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data)
-            for nloc in sorted(lmdb_test_data.nloc_groups.keys()):
-                yield LmdbTestDataNlocView(lmdb_test_data, nloc)
+            for (nloc, _signature), indices in sorted(
+                lmdb_test_data.find_signature_groups.items()
+            ):
+                yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices)
             return
 
         if hasattr(validation_data, "_reader"):
             lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data)
-            for nloc in sorted(lmdb_test_data.nloc_groups.keys()):
-                yield LmdbTestDataNlocView(lmdb_test_data, nloc)
+            for (nloc, _signature), indices in sorted(
+                lmdb_test_data.find_signature_groups.items()
+            ):
+                yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices)
             return
 
         if hasattr(validation_data, "data_systems"):

deepmd/pt/utils/lmdb_dataset.py

@@ -171,24 +171,41 @@ def __init__(
                 collate_fn=_collate_lmdb_batch,
             )
 
-        # Per-nloc-group dataloaders for make_stat_input.
-        # Each group gets its own DataLoader so torch.cat in stat collection
-        # only concatenates same-shape tensors.
-        self._nloc_dataloaders: list[DataLoader] = []
+        # 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.
+        self._nloc_dataloaders: list[DataLoader] | None = None
+
+    def _rebuild_nloc_dataloaders(self) -> None:
+        """Build homogeneous loaders used by model-stat collection."""
+        dataloaders: list[DataLoader] = []
         for nloc in sorted(self._reader.nloc_groups.keys()):
-            indices = self._reader.nloc_groups[nloc]
-            subset = torch.utils.data.Subset(self, indices)
-            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,
-                )
-            self._nloc_dataloaders.append(dl)
+            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)
+        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()
+        dataloaders = self._nloc_dataloaders
+        if dataloaders is None:
+            raise RuntimeError("Failed to initialize LMDB statistics dataloaders")
+        return dataloaders
 
     def __len__(self) -> int:
         return len(self._reader)
@@ -229,6 +246,7 @@ 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()
 
     def preload_and_modify_all_data_torch(self) -> None:
         """No-op: LMDB reads on demand."""
@@ -328,17 +346,17 @@ def batch_sizes(self) -> list[int]:
 
     @property
     def systems(self) -> list:
-        """One 'system' per nloc group for stat collection compatibility."""
-        return [self] * len(self._nloc_dataloaders)
+        """One logical system per stack-compatible statistics group."""
+        return [self] * len(self._get_nloc_dataloaders())
 
     @property
     def dataloaders(self) -> list:
-        """Per-nloc-group dataloaders for make_stat_input.
+        """Homogeneous dataloaders for make_stat_input.
 
-        Each dataloader yields batches with uniform nloc, so torch.cat
-        in stat collection only concatenates same-shape tensors.
+        Each loader has one nloc and one availability signature, so stat
+        collection sees consistent shapes and scalar ``find_*`` flags.
         """
-        return self._nloc_dataloaders
+        return self._get_nloc_dataloaders()
 
     @property
     def sampler_list(self) -> list:

deepmd/pt_expt/utils/lmdb_dataset.py

@@ -100,6 +100,9 @@ def add_data_requirements(
         self, data_requirement: list[DataRequirementItem]
     ) -> None:
         self._reader.add_data_requirement(data_requirement)
+        # Discard any iterator created under the previous availability
+        # signature so the next batch uses the newly registered labels.
+        self._iter = iter(self._sampler)
 
     def get_nsystems(self) -> int:
         """Return 1: pt_expt's stat collection treats LMDB as a single system.

source/tests/common/dpmodel/test_lmdb_data.py

@@ -724,6 +724,38 @@ def test_sampler_with_block_targets(self):
         self.assertGreater(len(all_indices), 600)
         self.assertEqual(len(set(all_indices)), 600)
 
+    def test_sampler_allocates_block_target_across_find_signatures(self):
+        """Independent signature groups must not each round the block target."""
+
+        class TwoSignatureReader:
+            """Minimal reader exposing two one-frame availability groups."""
+
+            def __init__(self):
+                self.nloc_groups = {6: [0, 1]}
+                self.frame_system_ids = [0, 0]
+
+            @staticmethod
+            def group_indices_by_find_signature(indices):
+                self.assertEqual(indices, [0, 1])
+                return {(0.0,): [0], (1.0,): [1]}
+
+            @staticmethod
+            def get_batch_size_for_nloc(nloc):
+                self.assertEqual(nloc, 6)
+                return 1
+
+        sampler = SameNlocBatchSampler(
+            TwoSignatureReader(),
+            shuffle=False,
+            block_targets=[([0], 3)],
+        )
+        batches = list(sampler)
+        counts = [sum(index in batch for batch in batches) for index in (0, 1)]
+
+        self.assertEqual(len(sampler), len(batches))
+        self.assertEqual(sum(map(len, batches)), 3)
+        self.assertEqual(sorted(counts), [1, 2])
+
     def test_sampler_without_block_targets(self):
         reader = LmdbDataReader(self._lmdb_path, ["O", "H"])
         sampler = SameNlocBatchSampler(reader, shuffle=False)

source/tests/pt/test_lmdb_dataloader.py

@@ -24,6 +24,9 @@
     _remap_keys,
     merge_lmdb,
 )
+from deepmd.pt.loss.ener import (
+    EnergyStdLoss,
+)
 from deepmd.pt.utils.lmdb_dataset import (
     LmdbDataset,
     _collate_lmdb_batch,
@@ -82,6 +85,37 @@ def _create_test_lmdb(path: str, nframes: int = 10, natoms: int = 6) -> None:
     env.close()
 
 
+def _create_partially_labeled_lmdb(path: str) -> None:
+    """Create same-nloc frames with complementary energy/force labels."""
+    nframes = 4
+    natoms = 6
+    env = lmdb.open(path, map_size=10 * 1024 * 1024)
+    with env.begin(write=True) as txn:
+        metadata = {
+            "nframes": nframes,
+            "frame_idx_fmt": "012d",
+            "type_map": ["O", "H"],
+            "system_info": {"natoms": [3, 3]},
+            "frame_nlocs": [natoms] * nframes,
+        }
+        txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True))
+        for index in range(nframes):
+            frame = _make_frame(natoms=natoms, seed=index)
+            if index % 2 == 0:
+                frame.pop("forces")
+                frame["energies"]["data"] = np.array([2.0], dtype=np.float64).tobytes()
+            else:
+                frame.pop("energies")
+                frame["forces"]["data"] = np.ones(
+                    (natoms, 3), dtype=np.float64
+                ).tobytes()
+            txn.put(
+                format(index, "012d").encode(),
+                msgpack.packb(frame, use_bin_type=True),
+            )
+    env.close()
+
+
 @pytest.fixture
 def lmdb_dir(tmp_path):
     """Create a temporary LMDB dataset."""
@@ -312,6 +346,106 @@ def test_full_epoch(self, lmdb_dir):
             total_frames = sum(batch["coord"].shape[0] for batch in dl)
         assert total_frames == 10
 
+    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"))
+        path = str(tmp_path / "partial.lmdb")
+        _create_partially_labeled_lmdb(path)
+        ds = LmdbDataset(path, type_map=["O", "H"], batch_size=2)
+        ds.add_data_requirement(
+            [
+                DataRequirementItem("energy", 1, atomic=False, must=False, default=7.0),
+                DataRequirementItem("force", 3, atomic=True, must=False, default=11.0),
+            ]
+        )
+
+        sampler = SameNlocBatchSampler(ds._reader, shuffle=True, seed=11)
+        batches = list(sampler)
+        assert len(sampler) == len(batches) == 2
+        for indices in batches:
+            signatures = {ds._reader.get_find_signature(index) for index in indices}
+            assert len(signatures) == 1
+
+        distributed_batches = []
+        for rank in range(2):
+            distributed = DistributedSameNlocBatchSampler(
+                ds._reader,
+                rank=rank,
+                world_size=2,
+                shuffle=True,
+                seed=11,
+            )
+            rank_batches = list(distributed)
+            assert len(distributed) == len(rank_batches) == 1
+            distributed_batches.extend(rank_batches)
+        assert sorted(index for batch in distributed_batches for index in batch) == [
+            0,
+            1,
+            2,
+            3,
+        ]
+
+        loss_module = EnergyStdLoss(
+            starter_learning_rate=1.0,
+            start_pref_e=1.0,
+            limit_pref_e=1.0,
+        )
+        force_loss_module = EnergyStdLoss(
+            starter_learning_rate=1.0,
+            start_pref_f=1.0,
+            limit_pref_f=1.0,
+        )
+        observed_flags = set()
+        with torch.device("cpu"):
+            batches = list(ds._inner_dataloader)
+            stat_batches = [
+                batch for dataloader in ds.dataloaders for batch in dataloader
+            ]
+        assert {
+            (float(batch["find_energy"]), float(batch["find_force"]))
+            for batch in stat_batches
+        } == {(1.0, 0.0), (0.0, 1.0)}
+        for batch in batches:
+            flags = (float(batch["find_energy"]), float(batch["find_force"]))
+            observed_flags.add(flags)
+
+            def zero_model(_batch=batch, **kwargs):
+                return {
+                    "energy": torch.zeros_like(_batch["energy"]),
+                    "force": torch.zeros_like(_batch["force"]),
+                }
+
+            _, loss, _ = loss_module(
+                {},
+                zero_model,
+                {
+                    "energy": batch["energy"],
+                    "find_energy": batch["find_energy"],
+                },
+                natoms=6,
+                learning_rate=1.0,
+            )
+            if flags[0] == 0.0:
+                assert loss.item() == 0.0
+            else:
+                assert loss.item() > 0.0
+
+            _, force_loss, _ = force_loss_module(
+                {},
+                zero_model,
+                {
+                    "force": batch["force"],
+                    "find_force": batch["find_force"],
+                },
+                natoms=6,
+                learning_rate=1.0,
+            )
+            if flags[1] == 0.0:
+                assert force_loss.item() == 0.0
+            else:
+                assert force_loss.item() > 0.0
+        assert observed_flags == {(1.0, 0.0), (0.0, 1.0)}
+
 
 # ============================================================
 # Collate function
@@ -356,6 +490,14 @@ def test_collate_none_values(self):
         ]
         assert _collate_lmdb_batch(frames)["box"] is None
 
+    def test_collate_rejects_mixed_find_flags(self):
+        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)
+
 
 # ============================================================
 # Type map remapping (PT-specific: LmdbDataset)