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

@@ -159,6 +159,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 one logical LMDB training dataset."""

source/tests/common/dpmodel/test_lmdb_data.py

@@ -780,6 +780,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)