From d73862b5eb376c48e59bd79ba870f4a7f1446774 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Thu, 16 Jul 2026 12:35:20 +0800 Subject: [PATCH 1/2] fix(pt-expt): group LMDB statistics by atom count Expose fixed-nloc statistical partitions without changing the LMDB's logical dataset identity, preventing mixed-size batches from being concatenated during model-stat initialization. Add regression coverage for statistics packing and end-to-end mixed-nloc training. --- deepmd/pt_expt/utils/lmdb_dataset.py | 67 ++++++++++++++++++-- deepmd/utils/model_stat.py | 44 ++++++++++--- source/tests/pt_expt/test_lmdb_training.py | 72 ++++++++++++++++++++++ 3 files changed, 168 insertions(+), 15 deletions(-) diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index 4b279880e1..1b15c261ba 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -32,7 +32,10 @@ class LmdbDataSystem: Exposes the small surface that pt_expt's trainer touches: ``get_batch(sys_idx=None)``, ``add_data_requirements(list)``, and ``get_nsystems()``. Internally uses :class:`LmdbDataReader` for I/O and - :class:`SameNlocBatchSampler` to draw same-nloc batches. + :class:`SameNlocBatchSampler` to draw same-nloc batches. Statistics use a + separate logical-system view in which every ``nloc`` group is sampled + independently, matching the PyTorch DataLoader adapter without changing + the identity of the LMDB as one training dataset. Parameters ---------- @@ -76,6 +79,8 @@ def __init__( block_targets=block_targets, ) self._iter = iter(self._sampler) + self._stat_nlocs = tuple(sorted(self._reader.nloc_groups)) + self._stat_offsets = [0] * len(self._stat_nlocs) # ------------------------------------------------------------------ # pt_expt trainer surface @@ -93,6 +98,60 @@ def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: except StopIteration: self._iter = iter(self._sampler) indices = next(self._iter) + return self._collate_indices(indices) + + def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: + """Return one batch from a fixed-``nloc`` statistical system. + + Parameters + ---------- + sys_idx : int + Index into the sorted ``nloc`` groups. + + Returns + ------- + dict[str, Any] + A collated NumPy batch whose frames have one atom count. + + Raises + ------ + IndexError + If ``sys_idx`` does not identify an available ``nloc`` group. + """ + if not 0 <= sys_idx < len(self._stat_nlocs): + raise IndexError( + f"Statistical system index {sys_idx} is out of range for " + f"{len(self._stat_nlocs)} nloc groups." + ) + + nloc = self._stat_nlocs[sys_idx] + group_indices = self._reader.nloc_groups[nloc] + batch_size = self._reader.get_batch_size_for_nloc(nloc) + start = self._stat_offsets[sys_idx] + if start >= len(group_indices): + start = 0 + stop = min(start + batch_size, len(group_indices)) + self._stat_offsets[sys_idx] = stop + return self._collate_indices(group_indices[start:stop]) + + def get_stat_nsystems(self) -> int: + """Return the number of fixed-``nloc`` statistical systems.""" + return len(self._stat_nlocs) + + def get_stat_numb_batches(self, sys_idx: int) -> int: + """Return the available batch count for one statistical system.""" + if not 0 <= sys_idx < len(self._stat_nlocs): + raise IndexError( + f"Statistical system index {sys_idx} is out of range for " + f"{len(self._stat_nlocs)} nloc groups." + ) + nloc = self._stat_nlocs[sys_idx] + nframes = len(self._reader.nloc_groups[nloc]) + batch_size = self._reader.get_batch_size_for_nloc(nloc) + return (nframes + batch_size - 1) // batch_size + + def _collate_indices(self, indices: list[int]) -> dict[str, Any]: + """Load and collate the requested dataset indices.""" frames = [self._reader[int(i)] for i in indices] return collate_lmdb_frames(frames) @@ -102,11 +161,7 @@ def add_data_requirements( self._reader.add_data_requirement(data_requirement) def get_nsystems(self) -> int: - """Return 1: pt_expt's stat collection treats LMDB as a single system. - - Per-system sampling within the LMDB is handled by - ``SameNlocBatchSampler`` + ``block_targets``. - """ + """Return one logical LMDB training dataset.""" return 1 # ------------------------------------------------------------------ diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index 33ebbcae57..ca7a2d7070 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -16,11 +16,35 @@ log = logging.getLogger(__name__) +def _get_stat_nsystems(data: Any) -> int: + """Return the number of shape-compatible systems used for statistics.""" + get_stat_nsystems = getattr(data, "get_stat_nsystems", None) + if get_stat_nsystems is not None: + return int(get_stat_nsystems()) + return int(data.get_nsystems()) + + +def _get_stat_numb_batches(data: Any, sys_idx: int, nbatches: int) -> int: + """Limit sampling to the available batches of one statistical system.""" + get_stat_numb_batches = getattr(data, "get_stat_numb_batches", None) + if get_stat_numb_batches is None: + return nbatches + return min(nbatches, int(get_stat_numb_batches(sys_idx))) + + +def _get_stat_batch(data: Any, sys_idx: int) -> dict[str, Any]: + """Return one batch from a shape-compatible statistical system.""" + get_stat_batch = getattr(data, "get_stat_batch", None) + if get_stat_batch is not None: + return get_stat_batch(sys_idx) + return data.get_batch(sys_idx=sys_idx) + + def _make_all_stat_ref(data: Any, nbatches: int) -> dict[str, list[Any]]: all_stat = defaultdict(list) - for ii in range(data.get_nsystems()): - for jj in range(nbatches): - stat_data = data.get_batch(sys_idx=ii) + for ii in range(_get_stat_nsystems(data)): + for jj in range(_get_stat_numb_batches(data, ii, nbatches)): + stat_data = _get_stat_batch(data, ii) for dd in stat_data: if dd == "natoms_vec": stat_data[dd] = stat_data[dd].astype(np.int32) @@ -55,10 +79,10 @@ def collect_batches( all_stat[key][batch_idx][frame_idx] """ all_stat = defaultdict(list) - for ii in range(data.get_nsystems()): + for ii in range(_get_stat_nsystems(data)): sys_stat = defaultdict(list) - for jj in range(nbatches): - stat_data = data.get_batch(sys_idx=ii) + for jj in range(_get_stat_numb_batches(data, ii, nbatches)): + stat_data = _get_stat_batch(data, ii) for dd in stat_data: if dd == "natoms_vec": stat_data[dd] = stat_data[dd].astype(np.int32) @@ -78,8 +102,10 @@ def make_stat_input( ) -> list[dict[str, np.ndarray]]: """Pack data for statistics using DeepmdDataSystem. - Collects *nbatches* batches from each system and concatenates them - into a single dict per system. The returned format + Collects up to *nbatches* batches from each shape-compatible statistical + system and concatenates them into one dictionary per system. Data sources + with variable atom counts may expose dedicated statistical-system methods + so incompatible ``nloc`` groups remain separate. The returned format (``list[dict[str, np.ndarray]]``) is backend-agnostic and can be consumed by ``compute_or_load_stat`` in dpmodel, pt_expt, and jax. @@ -98,7 +124,7 @@ def make_stat_input( """ all_stat = collect_batches(data, nbatches, merge_sys=False) - nsystems = data.get_nsystems() + nsystems = _get_stat_nsystems(data) log.info(f"Packing data for statistics from {nsystems} systems") keys = list(all_stat.keys()) diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 94673a8761..4e3caad233 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -28,6 +28,9 @@ from deepmd.pt_expt.utils.lmdb_dataset import ( LmdbDataSystem, ) +from deepmd.pt_expt.utils.stat import ( + make_stat_input, +) from deepmd.utils.argcheck import ( normalize, ) @@ -85,13 +88,42 @@ def _create_test_lmdb(path: str, nframes: int, natoms: int) -> None: env.close() +def _create_mixed_nloc_test_lmdb(path: str) -> None: + """Write an LMDB with four six-atom and four nine-atom frames.""" + frame_nlocs = [6] * 4 + [9] * 4 + env = lmdb.open(path, map_size=10 * 1024 * 1024) + fmt = "012d" + metadata = { + "nframes": len(frame_nlocs), + "frame_idx_fmt": fmt, + "frame_nlocs": frame_nlocs, + "type_map": ["O", "H"], + "system_info": { + "formula": "mixed", + "natoms": [3, 3], + "nframes": len(frame_nlocs), + }, + } + with env.begin(write=True) as txn: + txn.put(b"__metadata__", msgpack.packb(metadata, use_bin_type=True)) + for index, nloc in enumerate(frame_nlocs): + key = format(index, fmt).encode() + txn.put( + key, + msgpack.packb(_make_frame(nloc, index), use_bin_type=True), + ) + env.close() + + class TestLmdbDataSystemGetBatch(unittest.TestCase): """LmdbDataSystem.get_batch produces a numpy dict that normalize_batch accepts.""" def setUp(self) -> None: self.tmpdir = tempfile.mkdtemp() self.lmdb_path = os.path.join(self.tmpdir, "test.lmdb") + self.mixed_lmdb_path = os.path.join(self.tmpdir, "mixed.lmdb") _create_test_lmdb(self.lmdb_path, nframes=8, natoms=6) + _create_mixed_nloc_test_lmdb(self.mixed_lmdb_path) def tearDown(self) -> None: shutil.rmtree(self.tmpdir, ignore_errors=True) @@ -159,6 +191,25 @@ def test_add_data_requirements_passthrough(self) -> None: self.assertIn("energy", batch) self.assertIn("find_energy", batch) + def test_stat_input_partitions_mixed_nloc_batches(self) -> None: + """Statistics expose each atom-count group as one logical system.""" + ds = LmdbDataSystem( + lmdb_path=self.mixed_lmdb_path, + type_map=["O", "H"], + batch_size=2, + seed=0, + ) + + self.assertEqual(ds.get_nsystems(), 1) + self.assertEqual(ds.get_stat_nsystems(), 2) + sampled = make_stat_input(ds, nbatches=10) + + self.assertEqual(len(sampled), 2) + self.assertEqual( + sorted(sample["atype"].shape for sample in sampled), + [(4, 6), (4, 9)], + ) + class TestLmdbTrainingLoop(unittest.TestCase): """End-to-end: get_trainer routes an LMDB path and runs training steps.""" @@ -167,8 +218,10 @@ class TestLmdbTrainingLoop(unittest.TestCase): def setUpClass(cls) -> None: cls.tmpdir = tempfile.mkdtemp() cls.lmdb_path = os.path.join(cls.tmpdir, "train.lmdb") + cls.mixed_lmdb_path = os.path.join(cls.tmpdir, "mixed.lmdb") cls.val_lmdb_path = os.path.join(cls.tmpdir, "val.lmdb") _create_test_lmdb(cls.lmdb_path, nframes=8, natoms=6) + _create_mixed_nloc_test_lmdb(cls.mixed_lmdb_path) _create_test_lmdb(cls.val_lmdb_path, nframes=4, natoms=6) @classmethod @@ -244,6 +297,25 @@ def test_get_trainer_routes_lmdb(self) -> None: finally: os.chdir(cwd) + def test_mixed_nloc_statistics_and_training(self) -> None: + """Trainer computes statistics and trains across fixed-nloc batches.""" + config = self._make_lmdb_config(numb_steps=2) + config["model"]["data_stat_nbatch"] = 10 + config["training"]["training_data"]["systems"] = self.mixed_lmdb_path + 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) + self.assertEqual(trainer.training_data.get_nsystems(), 1) + self.assertEqual(trainer.training_data.get_stat_nsystems(), 2) + trainer.run() + finally: + os.chdir(cwd) + if __name__ == "__main__": unittest.main() From 77a6511a6b07e18c18999afc2e91692afcf539c1 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 17 Jul 2026 11:06:12 +0800 Subject: [PATCH 2/2] fix --- deepmd/utils/model_stat.py | 13 ++++++-- source/tests/pt_expt/test_lmdb_training.py | 37 ++++++++++++++++++++-- 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/deepmd/utils/model_stat.py b/deepmd/utils/model_stat.py index ca7a2d7070..2ce82ace6b 100644 --- a/deepmd/utils/model_stat.py +++ b/deepmd/utils/model_stat.py @@ -63,7 +63,11 @@ def collect_batches( Parameters ---------- data - The data (must support ``get_nsystems()`` and ``get_batch(sys_idx=)``) + The data source. It must support ``get_nsystems()`` and + ``get_batch(sys_idx=)``. Optional ``get_stat_nsystems()``, + ``get_stat_numb_batches(sys_idx)``, and ``get_stat_batch(sys_idx)`` + hooks may expose shape-compatible logical systems and their available + batches specifically for statistics. nbatches : int The number of batches per system merge_sys : bool (True) @@ -112,8 +116,11 @@ def make_stat_input( Parameters ---------- data - The multi-system data manager - (must support ``get_nsystems()`` and ``get_batch(sys_idx=)``). + The multi-system data manager. It must support ``get_nsystems()`` and + ``get_batch(sys_idx=)``. Optional ``get_stat_nsystems()``, + ``get_stat_numb_batches(sys_idx)``, and ``get_stat_batch(sys_idx)`` + hooks may expose shape-compatible logical systems and their available + batches specifically for statistics. nbatches : int Number of batches to collect per system. diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 4e3caad233..f5aed2f4a8 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -89,8 +89,8 @@ def _create_test_lmdb(path: str, nframes: int, natoms: int) -> None: def _create_mixed_nloc_test_lmdb(path: str) -> None: - """Write an LMDB with four six-atom and four nine-atom frames.""" - frame_nlocs = [6] * 4 + [9] * 4 + """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) fmt = "012d" metadata = { @@ -205,11 +205,44 @@ def test_stat_input_partitions_mixed_nloc_batches(self) -> None: sampled = make_stat_input(ds, nbatches=10) self.assertEqual(len(sampled), 2) + self.assertEqual( + sorted(sample["atype"].shape for sample in sampled), + [(5, 6), (5, 9)], + ) + + def test_stat_input_limits_batches_per_nloc_group(self) -> None: + """Statistics honor the requested batch cap for every nloc group.""" + ds = LmdbDataSystem( + lmdb_path=self.mixed_lmdb_path, + type_map=["O", "H"], + batch_size=2, + seed=0, + ) + + sampled = make_stat_input(ds, nbatches=2) + self.assertEqual( sorted(sample["atype"].shape for sample in sampled), [(4, 6), (4, 9)], ) + def test_stat_batches_restart_after_non_divisible_pass(self) -> None: + """Statistical batches restart after including a partial tail batch.""" + ds = LmdbDataSystem( + lmdb_path=self.mixed_lmdb_path, + type_map=["O", "H"], + batch_size=2, + seed=0, + ) + + for sys_idx, nloc in enumerate((6, 9)): + self.assertEqual(ds.get_stat_numb_batches(sys_idx), 3) + shapes = [ds.get_stat_batch(sys_idx)["atype"].shape for _ in range(4)] + self.assertEqual( + shapes, + [(2, nloc), (2, nloc), (1, nloc), (2, nloc)], + ) + class TestLmdbTrainingLoop(unittest.TestCase): """End-to-end: get_trainer routes an LMDB path and runs training steps."""