fix(data): preserve LMDB label availability across batches (#5839)

Fixes #5636

## Summary

- cache each LMDB frame availability signature from lightweight msgpack
metadata
- partition training, statistics, auto-probability, and distributed
batches by both atom count and scalar `find_*` signature
- allocate each auto-probability block target exactly across all
homogeneous groups with deterministic integer largest-remainder
apportionment
- retain the existing scalar loss contract and reject heterogeneous
collation defensively
- reset pt_expt sampling after data requirements change
- split full LMDB validation by atom count and availability so
default-filled labels are excluded from metrics

Simply stacking `find_*` values would not be safe here. Current PyTorch
and dpmodel losses reduce a whole batch before applying scalar
availability flags, and display/validation paths convert those flags to
booleans. A vector flag would broadcast incorrectly, produce
vector-valued losses, or fail boolean conversion. Homogeneous batching
is the smaller compatible fix.

## Why existing tests missed this

The direct collator test used `find_energy=1` in every frame and did not
assert the flag value. All LMDB factories stored energy and force in
every record, so same-nloc batching was also label-homogeneous by
accident. Reader consistency tests inspect individual frames before
collation, and loss tests construct homogeneous scalar flags without
connecting them to LMDB batches. Auto-probability tests also used one
availability signature, so they could not expose independent
per-signature rounding that changed a block target of 3 into 4 sampled
frames.

Full-validation coverage only varied atom count. `LmdbTestData`
therefore always observed globally present labels, and no test combined
real labels with default-filled frames or checked order-independent
optional-label metrics.

## Validation

- common dpmodel LMDB, PyTorch LMDB, and full-validation suites: 137
passed, 3 subtests passed
- isolated pt_expt LMDB data-system suite: 4 passed
- review-focused common/PT/pt_expt LMDB suites: 36 passed through the
pure-Python path (the locally installed optional PyTorch custom-op
library has a stale ABI)
- regression verifies an auto-probability target of 3 remains exactly 3
across two availability groups and sampler length stays exact
- regression verifies energy and force defaults contribute zero loss
when their corresponding `find_*` flag is zero
- regression consumes the pre-requirement pt_expt iterator and verifies
it is replaced before availability-homogeneous batching
- regression verifies energy and force full-validation metrics use only
their two labeled frames
- `ruff format .`
- `ruff check .`

Coding agent: Codex
Codex version: codex-cli 0.144.4
Model: gpt-5.6-sol
Reasoning effort: xhigh

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **Bug Fixes**
- LMDB training batching now partitions by both atom count and
consistent scalar label-availability flags, ensuring batches don’t mix
incompatible `find_*` settings.
- Batch collation now raises an error when label-availability flags
conflict across frames.
- Full validation now groups evaluation by atom count and label
availability so default-filled/missing labels don’t skew metrics.

- **Performance / Behavior**
- LMDB dataset loaders are now built lazily and refresh immediately
after adding new label requirements.

- **Tests**
- Added coverage for partially labeled LMDB batching, validator
grouping, and block-target distribution across label-availability
signatures.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: njzjz-bot <njzjz.bot@gmail.com>

Author: njzjz-bot

deepmd/dpmodel/utils/lmdb_data.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/train/validation.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/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."""

deepmd/pt_expt/utils/lmdb_dataset.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)