diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index 1463dad928..25969593d9 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -12,6 +12,12 @@ make_task_maps, print_data_summaries, ) +from .data_progress import ( + DATA_PROGRESS_CHECKPOINT_KEY, + collect_training_data_progress, + preserve_training_data_progress, + restore_training_data_progress, +) from .entrypoint import ( AbstractTrainEntrypoint, TrainEntrypointOptions, @@ -44,6 +50,7 @@ ) __all__ = [ + "DATA_PROGRESS_CHECKPOINT_KEY", "DEFAULT_TASK_KEY", "AbstractTrainEntrypoint", "AbstractTrainer", @@ -64,9 +71,12 @@ "build_checkpoint_stores", "change_model_out_bias", "change_model_out_bias_by_task", + "collect_training_data_progress", "iter_training_task_configs", "make_task_maps", + "preserve_training_data_progress", "print_data_summaries", "resolve_keep_ckpt_count", "resolve_step_schedule", + "restore_training_data_progress", ] diff --git a/deepmd/dpmodel/train/data_progress.py b/deepmd/dpmodel/train/data_progress.py new file mode 100644 index 0000000000..0c5858fe78 --- /dev/null +++ b/deepmd/dpmodel/train/data_progress.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Small, backend-independent checkpoints of the next training batches. + +Data sources own their cursors and deterministic sampler state. This module +adds task names and the shared data-source RNG, not decoded data or permutations. +""" + +from collections.abc import ( + Iterator, + Mapping, +) +from contextlib import ( + contextmanager, +) +from typing import ( + Any, +) + +from deepmd.utils import random as dp_random + +DATA_PROGRESS_CHECKPOINT_KEY = "data_progress" +DATA_PROGRESS_VERSION = 1 + + +def collect_training_data_progress( + training_data_by_task: Mapping[str, Any], +) -> dict[str, Any]: + """Collect supported data-source cursors and the shared data-source RNG.""" + return { + "version": DATA_PROGRESS_VERSION, + "tasks": { + key: data.state_dict() + for key, data in training_data_by_task.items() + if callable(getattr(data, "state_dict", None)) + and callable(getattr(data, "load_state_dict", None)) + }, + "rng": dp_random.get_state(), + } + + +def restore_training_data_progress( + training_data_by_task: Mapping[str, Any], + progress: Mapping[str, Any] | None, +) -> None: + """Restore cursors lazily; missing progress preserves legacy restarts. + + The source implementations validate their sampling configuration and must + not read or decode data here. Restoring the shared RNG last also isolates + it from any source-specific reconstruction. + """ + if progress is None: + return + if progress.get("version") != DATA_PROGRESS_VERSION: + raise ValueError("Unsupported data-progress checkpoint version.") + tasks = progress.get("tasks") + if not isinstance(tasks, Mapping): + raise ValueError("data-progress 'tasks' must be a mapping.") + supported = { + key: data + for key, data in training_data_by_task.items() + if callable(getattr(data, "load_state_dict", None)) + and callable(getattr(data, "state_dict", None)) + } + if set(tasks) != set(supported): + raise ValueError("data-progress tasks do not match the training data sources.") + if "rng" not in progress: + raise ValueError("data-progress checkpoint is missing the selection RNG.") + for key, data in supported.items(): + data.load_state_dict(tasks[key]) + dp_random.set_state(progress["rng"]) + + +@contextmanager +def preserve_training_data_progress( + training_data_by_task: Mapping[str, Any], +) -> Iterator[None]: + """Keep diagnostics and compilation from consuming the training stream.""" + progress = collect_training_data_progress(training_data_by_task) + try: + yield + finally: + restore_training_data_progress(training_data_by_task, progress) diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 8a7da58470..f5e39a98ac 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -5,7 +5,9 @@ Backend-specific wrappers (PyTorch Dataset, JAX, etc.) import from here. """ +import copy import dataclasses +import hashlib import logging import math import multiprocessing @@ -30,6 +32,9 @@ from dataclasses import ( dataclass, ) +from itertools import ( + islice, +) from pathlib import ( Path, ) @@ -1762,6 +1767,12 @@ def __init__( self._reader = reader self._sampler = sampler self._epoch = 0 + # These name the next batch the consumer will see, not the pass that + # the one-ahead decoder (or the sampler itself) has already entered. + self._next_epoch = 0 + self._next_batch_index = 0 + self._resume_batch_index = 0 + self._restore_config: dict[str, Any] | None = None self._queued_pass_length = 0 self._queued_nframes = 0.0 self._last_pass_length: int | None = None @@ -1792,6 +1803,10 @@ def __next__(self) -> dict[str, Any]: batch = self._decode(indices) self._last_pass_length = self._queued_pass_length + self._next_batch_index += 1 + if self._next_batch_index == self._last_pass_length: + self._next_epoch += 1 + self._next_batch_index = 0 indices, self._queued_pass_length, self._queued_nframes = self._next_indices() self._schedule(indices) return batch @@ -1883,6 +1898,9 @@ def _next_indices(self) -> tuple[list[int], int, float]: def _iter_epoch(self) -> Iterator[tuple[list[int], int, float]]: """Snapshot each batch's pass length and rank-averaged frame count.""" + if self._restore_config is not None: + self._check_restore_config(self._restore_config) + self._restore_config = None set_epoch = getattr(self._sampler, "set_epoch", None) if callable(set_epoch): set_epoch(self._epoch) @@ -1890,9 +1908,27 @@ def _iter_epoch(self) -> Iterator[tuple[list[int], int, float]]: # enter another pass before the current batch reaches its consumer. # Snapshot the length lazily, after data requirements and layout are set. pass_length = len(self._sampler) + if pass_length == 0 and self._resume_batch_index == 0: + return + if not 0 <= self._resume_batch_index < pass_length: + raise ValueError( + f"LMDB batch index {self._resume_batch_index} is out of range " + f"for epoch {self._epoch} with {pass_length} batches" + ) + # This iterator already owns epoch advancement. Retain the sampler's + # current pass instead of its standalone __iter__ advancing and + # invalidating it, so diagnostic reads can rewind within the pass. + pass_batches = ( + self._sampler.batches() + if isinstance(self._sampler, LmdbBatchSampler) + else self._sampler + ) frame_counts = getattr(self._sampler, "batch_frame_counts", None) counts = frame_counts() if frame_counts is not None else None - for index, indices in enumerate(self._sampler): + start = self._resume_batch_index + batches = islice(pass_batches, start, None) + self._resume_batch_index = 0 + for index, indices in enumerate(batches, start): nframes = len(indices) if counts is None else counts[index] yield indices, pass_length, float(nframes) @@ -1924,6 +1960,97 @@ def uniform_nominal_mass(self) -> float: return 0.0 return nsystems / float(np.mean(self._sampler.batch_frame_counts())) + def _sampling_config(self) -> dict[str, Any]: + """Describe sampling without building a pass or reading any frames.""" + reader, sampler = self._reader, self._sampler + if not isinstance(sampler, (LmdbBatchSampler, DistributedLmdbBatchSampler)): + raise TypeError("LMDB progress requires an LMDB batch sampler") + return { + "path": reader.lmdb_path, + "transaction": reader._transaction().id(), + "nframes": len(reader), + "type_map": list(reader.type_map), + "batch_rule": [ + reader.batch_size, + reader._auto_rule, + reader._max_rule, + reader._filter_rule, + reader._mix_rule, + ], + "ragged": reader.ragged_batches, + "availability_keys": list( + _availability_requirement_keys(reader._data_requirements) + ), + "shuffle": sampler._shuffle, + "configured_seed": sampler._configured_seed, + "world_size": getattr(sampler, "world_size", 1), + "block_targets": hashlib.sha256( + msgpack.packb(sampler._block_targets, use_bin_type=True) + ).hexdigest(), + } + + def _check_restore_config( + self, config: dict[str, Any], *, final: bool = True + ) -> None: + """Reject changed sampling, deferring trainer-selected configuration.""" + current = self._sampling_config() + for key in current: + if not final and key in {"ragged", "availability_keys"}: + continue + if config.get(key) != current[key]: + raise ValueError(f"LMDB data progress is incompatible: {key} changed") + + def state_dict(self) -> dict[str, Any]: + """Return a compact next-batch checkpoint for this rank. + + No decoded frames, permutations or prefetched work are serialized. + Each rank restores its deterministic partition using its own consumed + cursor. Filtering can advance ranks differently, so distributed + checkpoint writers must retain every rank's state. + """ + config = ( + self._restore_config + if self._restore_config is not None + else self._sampling_config() + ) + return { + "version": 1, + "epoch": self._next_epoch, + "batch_index": self._next_batch_index, + "seed": self._sampler._seed, + "config": copy.deepcopy(config), + } + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore without sampling or decoding, even before trainer setup. + + Data requirements and the ragged layout may still be registered after + this call. Their identity and the saved offset are checked lazily on + the first batch, before decoding it. Rebuilding only the saved pass + skips batch indices, never the frames of earlier batches or epochs. + """ + if self._closed: + raise RuntimeError("cannot restore a closed LMDB batch iterator") + if state.get("version") != 1: + raise ValueError("Unsupported LMDB data progress version") + for key in ("epoch", "batch_index", "seed"): + if type(state.get(key)) is not int or state[key] < 0: + raise ValueError(f"LMDB data progress {key} must be a non-negative int") + config = state.get("config") + if not isinstance(config, dict): + raise ValueError("LMDB data progress is missing its sampling configuration") + self._check_restore_config(config, final=False) + self._discard_prefetch() + self._sampler.set_seed(state["seed"]) + self._sampler.set_epoch(state["epoch"]) + self._epoch = self._next_epoch = state["epoch"] + self._next_batch_index = self._resume_batch_index = state["batch_index"] + self._restore_config = copy.deepcopy(config) + self._queued_pass_length = 0 + self._queued_nframes = 0.0 + self._last_pass_length = None + self._iterator = self._iter_epoch() + def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]: """Submit one batch as balanced contiguous chunks. @@ -1980,11 +2107,8 @@ def closed(self) -> bool: """Whether this iterator has released its resources.""" return self._closed - def close(self) -> None: - """Finish or cancel prefetched work and release the decoder pool.""" - if self._closed: - return - self._closed = True + def _discard_prefetch(self) -> None: + """Discard queued reads while preserving the shared pool's health state.""" if self._pending is not None: running = [ future for future in self._pending.futures if not future.cancel() @@ -1993,6 +2117,13 @@ def close(self) -> None: self._lose_the_pool(running) self._pending = None self._deferred_indices = None + + def close(self) -> None: + """Finish or cancel prefetched work and release the decoder pool.""" + if self._closed: + return + self._closed = True + self._discard_prefetch() if self._pool is not None: self._pool = None _release_lmdb_executor(self._num_workers) @@ -3680,11 +3811,16 @@ def __init__( ) -> None: self._reader = reader self._shuffle = shuffle - self._seed = seed + self._configured_seed = seed + # Draw entropy once, not once per epoch, so even an unseeded run can + # reconstruct its next pass from a compact checkpoint. + self._seed = int(np.random.SeedSequence().entropy) if seed is None else seed self._epoch = 0 self._block_targets = block_targets self._batches: list[list[int]] | None = None + self._batch_frame_counts: np.ndarray | None = None self._data_requirements_revision = -1 + self._ragged_batches: bool | None = None def batches(self) -> list[list[int]]: """Return the batch list of the pending pass, building it if needed. @@ -3695,17 +3831,29 @@ def batches(self) -> 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 + ragged = self._reader.mixed_nloc and self._reader.ragged_batches + if ( + self._batches is None + or self._data_requirements_revision != revision + or self._ragged_batches != ragged + ): self._batches = _build_all_batches( self._reader, self._shuffle, - np.random.default_rng(seed), + np.random.default_rng(self._seed + self._epoch), self._block_targets, ) self._data_requirements_revision = revision + self._ragged_batches = ragged + self._batch_frame_counts = None return self._batches + def set_seed(self, seed: int) -> None: + """Restore the effective seed, invalidating any cached pass.""" + if seed != self._seed: + self._seed = seed + self._batches = None + def set_epoch(self, epoch: int) -> None: """Select the pass to serve, discarding any pass still pending. @@ -3730,7 +3878,12 @@ def __len__(self) -> int: def batch_frame_counts(self) -> np.ndarray: """Return the number of frames in each pending optimizer batch.""" - return np.fromiter((len(batch) for batch in self.batches()), dtype=np.float64) + batches = self.batches() + if self._batch_frame_counts is None: + self._batch_frame_counts = np.fromiter( + (len(batch) for batch in batches), dtype=np.float64 + ) + return self._batch_frame_counts @property def total_batches(self) -> int: @@ -3785,10 +3938,18 @@ def __init__( self._rank = rank self._world_size = world_size self._shuffle = shuffle + self._configured_seed = seed self._seed = seed if seed is not None else 0 self._epoch = 0 self._block_targets = block_targets self._global: LmdbBatchSampler | None = None + self._batch_frame_counts: np.ndarray | None = None + + def set_seed(self, seed: int) -> None: + """Restore the common seed, invalidating any cached global pass.""" + if seed != self._seed: + self._seed = seed + self._global = None def set_epoch(self, epoch: int) -> None: """Set epoch for deterministic cross-rank shuffling. @@ -3809,7 +3970,11 @@ def _global_batches(self) -> list[list[int]]: seed=self._seed + self._epoch, block_targets=self._block_targets, ) - return self._global.batches() + previous = self._global._batches + batches = self._global.batches() + if batches is not previous: + self._batch_frame_counts = None + return batches def __iter__(self) -> Iterator[list[int]]: """Yield this rank's partition of the global batch list.""" @@ -3854,13 +4019,17 @@ def batch_frame_counts(self) -> np.ndarray: Custom partitioning strategies must provide their corresponding rank-averaged counts when used with uniform source sampling. """ - counts = np.fromiter( - (len(batch) for batch in self._global_batches()), dtype=np.float64 - ) - if counts.size == 0: - return counts - padded = np.resize(counts, len(self) * self._world_size) - return padded.reshape(-1, self._world_size).mean(axis=1) + self._global_batches() + if self._batch_frame_counts is None: + counts = self._global.batch_frame_counts() + self._batch_frame_counts = ( + np.resize(counts, len(self) * self._world_size) + .reshape(-1, self._world_size) + .mean(axis=1) + if counts.size + else counts + ) + return self._batch_frame_counts @property def total_batches(self) -> int: diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 4d1fa7d696..eabcc0a8a4 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -34,6 +34,7 @@ from deepmd.pt_expt.utils.lmdb_dataset import ( LmdbDataSystem, ) +from deepmd.utils import random as dp_random from deepmd.utils.data import ( min_pair_dist_requirement, ) @@ -143,6 +144,9 @@ def get_trainer( is_distributed = dist.is_available() and dist.is_initialized() rank = dist.get_rank() if is_distributed else 0 world_size = dist.get_world_size() if is_distributed else 1 + if data_seed is not None: + seed32 = data_seed % (2**32) + dp_random.seed([rank, seed32] if is_distributed else seed32) def factory( task_config: TrainingTaskConfig, diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 84f5e01c7c..c231f428f1 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -64,6 +64,12 @@ from deepmd.dpmodel.train.data import ( add_training_data_requirements, ) +from deepmd.dpmodel.train.data_progress import ( + DATA_PROGRESS_CHECKPOINT_KEY, + collect_training_data_progress, + preserve_training_data_progress, + restore_training_data_progress, +) from deepmd.dpmodel.utils.batch import ( normalize_batch, split_batch, @@ -149,6 +155,7 @@ from deepmd.pt_expt.utils.stat import ( make_stat_input, ) +from deepmd.utils import random as dp_random from deepmd.utils.data import ( DataRequirementItem, ) @@ -1810,6 +1817,7 @@ def __init__( model_params = config["model"] training_params = config["training"] + self._configured_data_seed = training_params.get("seed") optimizer_params = config.get("optimizer", {}) validating_params = config.get("validating", {}) or {} @@ -1846,6 +1854,20 @@ def __init__( self.is_distributed = dist.is_available() and dist.is_initialized() self.rank = dist.get_rank() if self.is_distributed else 0 self.world_size = dist.get_world_size() if self.is_distributed else 1 + # Rank-local data draws must not change the shared task-selection sequence. + self._task_rng = np.random.default_rng( + None + if self._configured_data_seed is None + else self._configured_data_seed % (2**32) + ) + if ( + self.is_distributed + and self.multi_task + and self._configured_data_seed is None + ): + self._task_rng.bit_generator.state = self._broadcast_value_from_rank0( + self._task_rng.bit_generator.state if self.rank == 0 else None + ) self.sharding = ShardingPolicy.from_training_params( training_params, is_distributed=self.is_distributed ) @@ -2009,7 +2031,10 @@ def _make_sample( _n: int = _nbatch, _filter_enabled: bool = self.pair_filter_enabled_by_task[model_key], ) -> list[dict[str, np.ndarray]]: - return make_stat_input(_d, _n, _filter_enabled) + # Statistics run only on the chief. They must not advance its + # training cursors or data-source selection RNGs. + with preserve_training_data_progress({"statistics": _d}): + return make_stat_input(_d, _n, _filter_enabled) self._sample_funcs[model_key] = _make_sample @@ -2130,6 +2155,7 @@ def initialize_statistics( # cannot be copied into sharded parameters. ema_state_dict = None optimizer_state_dict = None + data_progress = None if resuming: log.info(f"Resuming from {resume_model}.") is_pte = resume_model.endswith((".pte", ".pt2")) @@ -2153,6 +2179,11 @@ def initialize_statistics( ema_state_dict = ( state_dict.get(EMA_CHECKPOINT_KEY) if continues_run else None ) + data_progress = ( + state_dict.get(DATA_PROGRESS_CHECKPOINT_KEY) + if continues_run + else None + ) state_dict = state_dict["model"] else: optimizer_state_dict = None @@ -2464,6 +2495,12 @@ def update_finetune_bias( multi_task=self.multi_task, ) + # Data requirements, layout, statistics and compilation may all draw + # batches during setup. Restore only after those steps have finished. + self._restore_data_progress( + data_progress, reset=training_params.get("reset_data_progress", False) + ) + def _create_full_validators( self, *, @@ -3039,6 +3076,7 @@ def _broadcast_model_stat(model: torch.nn.Module) -> None: def _broadcast_value_from_rank0(self, value: Any) -> Any: """Return rank 0's copy of ``value`` on every rank. + Unseeded task-selection RNG state must agree across ranks at startup. Epoch lengths round a quotient of sampling probabilities that is often an exact integer in real arithmetic, so a last-bit difference between ranks -- as reduction kernels dispatched for different CPU features @@ -3059,8 +3097,86 @@ def _broadcast_value_from_rank0(self, value: Any) -> Any: @property def checkpoint_is_collective(self) -> bool: - """Whether checkpoint preflight needs every rank.""" - return self.is_distributed + """Whether checkpoint preflight and data collection need every rank.""" + # Frame filtering can consume different numbers of local batches, so + # even plain DDP must contribute every rank's cursor at save time. + return self.is_distributed or self.sharding.enabled + + def _collect_data_progress(self) -> dict[str, Any]: + """Gather compact rank-local progress only at a checkpoint boundary.""" + local_progress = collect_training_data_progress(self.training_data_by_task) + ranks: list[Any] = [local_progress] + if self.is_distributed: + ranks = [None] * self.world_size + dist.all_gather_object(ranks, local_progress) + return { + "version": 1, + "world_size": self.world_size, + "configured_seed": self._configured_data_seed, + "task_rng": self._task_rng.bit_generator.state, + "ranks": ranks, + } + + def _restore_data_progress( + self, progress: Mapping[str, Any] | None, *, reset: bool = False + ) -> None: + """Restore progress unless an explicit reset or topology change prevents it. + + Validate the rank envelope before either fallback. An explicit reset + retains all newly initialized sources and the selection RNG together, + without partially applying saved task state or bypassing model restart. + """ + if progress is None: + return + if not isinstance(progress, Mapping): + raise ValueError("Data-progress checkpoint must be a mapping.") + if type(progress.get("version")) is not int or progress["version"] != 1: + raise ValueError("Unsupported distributed data-progress version.") + saved_world_size = progress.get("world_size") + if type(saved_world_size) is not int or saved_world_size < 1: + raise ValueError("Data-progress checkpoint world size must be positive.") + ranks = progress.get("ranks") + if not isinstance(ranks, list) or len(ranks) != saved_world_size: + raise ValueError("Data-progress checkpoint must contain every rank.") + if not all(isinstance(state, Mapping) for state in ranks): + raise ValueError("Data-progress checkpoint rank states must be mappings.") + if reset: + if self.rank == 0: + log.warning( + "training.reset_data_progress is true. Model, optimizer and " + "step are restored, but training data order and selection " + "RNG restart from their newly initialized state. Batches " + "may repeat; exact data-order continuation is unavailable." + ) + return + if saved_world_size != self.world_size: + # Model/optimizer/step restart has always allowed a new topology. + # Rank-local cursors cannot be remapped safely, particularly after + # filtering, so retain the freshly initialized data and RNG instead. + if self.rank == 0: + log.warning( + "Checkpoint world size changed from %d to %d. Model, optimizer " + "and step are restored, but training data order and selection " + "RNG restart from their newly initialized state; exact " + "data-order continuation is unavailable.", + saved_world_size, + self.world_size, + ) + return + if "configured_seed" not in progress or "task_rng" not in progress: + raise ValueError( + "Data-progress checkpoint lacks its seed or task RNG identity; " + "set training.reset_data_progress=true to restart data sampling." + ) + if progress["configured_seed"] != self._configured_data_seed: + raise ValueError( + "Training data progress is incompatible: training.seed changed; " + "set training.reset_data_progress=true to restart data sampling." + ) + task_rng = np.random.default_rng(0) + task_rng.bit_generator.state = progress["task_rng"] + restore_training_data_progress(self.training_data_by_task, ranks[self.rank]) + self._task_rng = task_rng def ensure_finite_gradients_for_checkpoint(self) -> None: """Check the accumulated gradient norm on all ranks before publication.""" @@ -3070,10 +3186,11 @@ def ensure_finite_gradients_for_checkpoint(self) -> None: def save_checkpoint(self, step: int) -> None: self.ensure_finite_gradients_for_checkpoint() + data_progress = self._collect_data_progress() if self.rank != 0 and not self.sharding.enabled: return ckpt_path = self.ckpt_store.path_for(step) - self._save_checkpoint_to_path(ckpt_path, step=step) + self._save_checkpoint_to_path(ckpt_path, step=step, data_progress=data_progress) if self.rank == 0: self.ckpt_store.publish(ckpt_path) self.ckpt_store.prune(ckpt_path) @@ -3090,10 +3207,12 @@ def _save_full_validation_checkpoint( save_path: Path, lr: float = 0.0, step: int = 0, + *, + data_progress: dict[str, Any] | None = None, ) -> None: """Save a checkpoint selected by full validation.""" del lr - self._save_checkpoint_to_path(save_path, step=step) + self._save_checkpoint_to_path(save_path, step=step, data_progress=data_progress) def _save_full_validation_ema_checkpoint( self, @@ -3115,6 +3234,7 @@ def _save_checkpoint_to_path( *, step: int, use_ema_weights: bool = False, + data_progress: dict[str, Any] | None = None, ) -> None: """Serialize the current trainer state to an explicit checkpoint path. @@ -3129,6 +3249,9 @@ def _save_checkpoint_to_path( Such a checkpoint is a deployment snapshot: it carries neither the optimizer state nor the EMA state, both of which describe the live weights it does not contain. + data_progress : dict or None, optional + Rank-local progress already gathered before chief-only validation. + Ordinary checkpoints collect it on every rank while being written. """ if use_ema_weights: with self.model_ema.apply_shadow(self.model): @@ -3139,7 +3262,7 @@ def _save_checkpoint_to_path( include_ema_state=False, ) return - self._write_checkpoint(ckpt_path, step=step) + self._write_checkpoint(ckpt_path, step=step, data_progress=data_progress) def _write_checkpoint( self, @@ -3148,6 +3271,7 @@ def _write_checkpoint( step: int, include_optimizer: bool = True, include_ema_state: bool = True, + data_progress: dict[str, Any] | None = None, ) -> None: """Serialize the wrapper, and optionally optimizer and EMA state.""" self._unwrapped.train_infos["step"] = step @@ -3170,6 +3294,8 @@ def _write_checkpoint( finally: for task_key, compiled in compiled_backup.items(): wrapper.model[task_key] = compiled + if include_optimizer and data_progress is None: + data_progress = self._collect_data_progress() # Sharded state is assembled on the chief; the other ranks have played # their part in the collectives above and hold nothing to write. if self.rank != 0: @@ -3177,6 +3303,7 @@ def _write_checkpoint( state: dict[str, Any] = {"model": model_state} if optim_state is not None: state["optimizer"] = optim_state + state[DATA_PROGRESS_CHECKPOINT_KEY] = data_progress if include_ema_state and self.model_ema is not None: state[EMA_CHECKPOINT_KEY] = self.model_ema.state_dict() ckpt_path.parent.mkdir(parents=True, exist_ok=True) @@ -3286,15 +3413,16 @@ def _precompile_outside_collectives(self) -> None: trainable_parameters = tuple( parameter for parameter in inner.parameters() if parameter.requires_grad ) - for task in self.training_tasks: - input_dict, label_dict = self.get_data(is_train=True, task_key=task.key) - _, loss, _ = inner( - **input_dict, - cur_lr=self.scheduler.get_last_lr()[0], - label=label_dict, - task_key=task.key, - ) - torch.autograd.grad(loss, trainable_parameters, allow_unused=True) + with preserve_training_data_progress(self.training_data_by_task): + for task in self.training_tasks: + input_dict, label_dict = self.get_data(is_train=True, task_key=task.key) + _, loss, _ = inner( + **input_dict, + cur_lr=self.scheduler.get_last_lr()[0], + label=label_dict, + task_key=task.key, + ) + torch.autograd.grad(loss, trainable_parameters, allow_unused=True) if torch.cuda.is_available(): torch.cuda.synchronize() log.info( @@ -3361,27 +3489,45 @@ def run_full_validation( learning_rate: float, ) -> None: """Run the active full validation flows for one step.""" - for validator, save_checkpoint in ( - (self.full_validator, self._save_full_validation_checkpoint), - (self.ema_full_validator, self._save_full_validation_ema_checkpoint), + for validator, save_checkpoint, include_progress in ( + (self.full_validator, self._save_full_validation_checkpoint, True), + (self.ema_full_validator, self._save_full_validation_ema_checkpoint, False), ): - if validator is None: + if validator is None or not validator.should_run(display_step): continue - validator.run( - step_id=display_step, - display_step=display_step, - lr=learning_rate, - save_checkpoint=save_checkpoint, - prepare_checkpoint=self.ensure_finite_gradients_for_checkpoint, - ) + # Plain DDP's validator calls its save hook only on the chief. + # Gather before entering it, while every rank is here, and pass + # the result through so that hook never starts another collective. + if include_progress and validator.save_best: + save_checkpoint = functools.partial( + save_checkpoint, data_progress=self._collect_data_progress() + ) + random_state = dp_random.get_state() + + def save_with_training_rng( + *args: Any, _save: Callable = save_checkpoint, **kwargs: Any + ) -> None: + # Best checkpoints are written inside validation. Record the + # training RNG, not the draws made by validation-data reads. + dp_random.set_state(random_state) + _save(*args, **kwargs) + + try: + validator.run( + step_id=display_step, + display_step=display_step, + lr=learning_rate, + save_checkpoint=save_with_training_rng, + prepare_checkpoint=self.ensure_finite_gradients_for_checkpoint, + ) + finally: + dp_random.set_state(random_state) def select_task(self, tasks: TrainingTaskCollection) -> TrainingTask: - """Select a task using DeePMD's seeded random helper.""" + """Select a task independently of rank-local data draws.""" if not tasks.is_multitask: return tasks[tasks.keys[0]] - from deepmd.utils import random as dp_random - - model_index = dp_random.choice( + model_index = self._task_rng.choice( np.arange(len(tasks), dtype=np.int_), p=tasks.probabilities, ) @@ -3403,12 +3549,22 @@ def collect_display_results( self.wrapper.eval() self._display_cur_lr_sched = step_result.payload["cur_lr_sched"] try: - return super().collect_display_results( - tasks, - active_task=active_task, - step=step, - step_result=step_result, - ) + # Only inactive tasks need extra training batches for display. + # The chief alone evaluates them, so those reads must be peeks. + # Preserving the RNG also isolates validation, whose cursor is not + # part of a training checkpoint, from future task/system choices. + display_data = { + task.key: task.training_data + for task in tasks + if task.key != active_task.key and self.metric_accumulator is None + } + with preserve_training_data_progress(display_data): + return super().collect_display_results( + tasks, + active_task=active_task, + step=step, + step_result=step_result, + ) finally: self._display_cur_lr_sched = None self.wrapper.train() diff --git a/deepmd/pt_expt/utils/lmdb_dataset.py b/deepmd/pt_expt/utils/lmdb_dataset.py index 20dcf97010..23d35d9092 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -152,6 +152,14 @@ def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: del sys_idx return next(self._batch_iterator) + def state_dict(self) -> dict[str, Any]: + """Return the next training-batch cursor, excluding statistical reads.""" + return self._batch_iterator.state_dict() + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore training progress without drawing or decoding any batch.""" + self._batch_iterator.load_state_dict(state) + def get_stat_batch(self, sys_idx: int) -> dict[str, Any]: """Return one batch from a homogeneous statistical system. diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index ddfe41ed69..751b6c69b6 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6076,6 +6076,16 @@ def training_args( doc_disp_freq = "The frequency of printing learning curve." doc_save_freq = "The frequency of saving check point." doc_save_ckpt = "The path prefix of saving check point files." + doc_reset_data_progress = ( + "On --restart, discard the saved training-data cursors and task/system " + "selection RNG, retaining their newly initialized state. Model weights, " + "optimizer, learning-rate schedule and training step still resume. " + "Use this explicit reset when relocating data or intentionally changing " + "data settings; batches may repeat and exact data-order continuation " + "is unavailable. By default, saved data progress is restored strictly, " + "including dataset locations. Has no effect on fresh training, " + "initialization, finetuning or checkpoints without data progress." + ) doc_save_dir = ( "The directory in which periodic checkpoint files are written, " "including the regular checkpoints (the `save_ckpt` prefix) and, when " @@ -6277,6 +6287,13 @@ def training_args( Argument( "save_ckpt", str, optional=True, default="model.ckpt", doc=doc_save_ckpt ), + Argument( + "reset_data_progress", + bool, + optional=True, + default=False, + doc=supported_backends("pt_expt") + doc_reset_data_progress, + ), Argument("max_ckpt_keep", int, optional=True, default=5, doc=doc_max_ckpt_keep), Argument( "ckpt_keep_ratio", diff --git a/deepmd/utils/data.py b/deepmd/utils/data.py index 733be3a626..d8e3393a7b 100644 --- a/deepmd/utils/data.py +++ b/deepmd/utils/data.py @@ -4,6 +4,7 @@ import bisect import copy import functools +import hashlib import logging from collections.abc import ( Iterable, @@ -148,6 +149,8 @@ def __init__( # set counters self.set_count = 0 self.iterator = 0 + self._batch_set_size = 0 + self._batch_shuffle_state = None self.shuffle_test = shuffle_test # set modifier self.modifier = modifier @@ -236,6 +239,7 @@ def add( # only non-standard tensors need this extra shape contract. data_config["special_shape"] = special_shape self.data_dict[key] = data_config + self._invalidate_data_cache() return self def reduce(self, key_out: str, key_in: str) -> "DeepmdData": @@ -264,8 +268,16 @@ def reduce(self, key_out: str, key_in: str) -> "DeepmdData": "repeat": 1, "reduce": key_in, } + self._invalidate_data_cache() return self + def _invalidate_data_cache(self) -> None: + """Discard arrays whose data requirements are no longer current.""" + for key in ("batch_set", "_batch_indices", "test_set"): + self.__dict__.pop(key, None) + if hasattr(self, "_modified_frame_cache"): + self._modified_frame_cache.clear() + def get_data_dict(self) -> dict: """Get the `data_dict`.""" return self.data_dict @@ -354,18 +366,26 @@ def get_batch(self, batch_size: int) -> dict: batch_size size of the batch """ - if hasattr(self, "batch_set"): - set_size = self.batch_set["coord"].shape[0] - else: - set_size = 0 - if self.iterator + batch_size > set_size: + if self.iterator + batch_size > self._batch_set_size: self._load_batch_set(self.dirs[self.set_count % self.get_numb_set()]) self.set_count += 1 - set_size = self.batch_set["coord"].shape[0] + elif not hasattr(self, "_batch_indices"): + # Restore the active order only when it is actually needed. An + # exhausted set can advance directly to the next one above. + iterator = self.iterator + set_size = self._batch_set_size + self._load_batch_set( + self.dirs[(self.set_count - 1) % self.get_numb_set()], + shuffle_state=self._batch_shuffle_state, + ) + if self._batch_set_size != set_size: + raise ValueError("The restored data set has a different numb_copy size") + self.iterator = iterator + set_size = self._batch_set_size iterator_1 = self.iterator + batch_size if iterator_1 >= set_size: iterator_1 = set_size - idx = np.arange(self.iterator, iterator_1, dtype=np.int64) + idx = self._batch_indices[self.iterator : iterator_1] self.iterator += batch_size ret = self._get_subdata(self.batch_set, idx) if "pair_margin" in self.data_dict: @@ -377,6 +397,84 @@ def get_batch(self, batch_size: int) -> dict: ) return ret + def _progress_identity(self) -> dict[str, Any]: + """Describe the source using metadata already read at construction.""" + return { + "sets": [str(path) for path in self.dirs], + "prefix_sum": self.prefix_sum, + "atom_types": hashlib.sha256(self.atom_type.tobytes()).hexdigest(), + "type_map": self.type_map, + "sort_atoms": self.sort_atoms, + "mixed_type": self.mixed_type, + "pbc": self.pbc, + } + + def state_dict(self) -> dict[str, Any]: + """Return compact batch progress, without frame arrays or permutations. + + The shared :mod:`deepmd.utils.random` state must be checkpointed + separately, once for all readers. Only the state before the active + set's shuffle is retained here to reconstruct its order lazily. + """ + return copy.deepcopy( + { + "version": 1, + "identity": self._progress_identity(), + "set_count": int(self.set_count), + "iterator": int(self.iterator), + "set_size": int(self._batch_set_size), + "shuffle_state": self._batch_shuffle_state, + } + ) + + def _validate_state_dict(self, state: dict[str, Any]) -> None: + if state.get("version") != 1: + raise ValueError("Unsupported DeepmdData progress version") + if state.get("identity") != self._progress_identity(): + raise ValueError("DeepmdData progress does not match the data source") + for key in ("set_count", "iterator", "set_size"): + if type(state.get(key)) is not int or state[key] < 0: + raise ValueError(f"Invalid DeepmdData progress {key}") + shuffle_state = state.get("shuffle_state") + if state["set_count"] == 0: + if state["iterator"] or state["set_size"] or shuffle_state is not None: + raise ValueError("Invalid unstarted DeepmdData progress") + elif shuffle_state is None: + raise ValueError("Missing DeepmdData shuffle state") + if shuffle_state is not None: + # Validation uses a private generator, never the shared RNG. + try: + np.random.RandomState(0).set_state(shuffle_state) + except (TypeError, ValueError, IndexError) as error: + raise ValueError("Invalid DeepmdData shuffle state") from error + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore logical progress without loading or scanning frame data. + + A same-set canonical cache may be reused when its requirements have + not changed. New readers reconstruct arrays only on the next batch. + """ + self._validate_state_dict(state) + active_set = ( + self.dirs[(state["set_count"] - 1) % self.get_numb_set()] + if state["set_count"] + else None + ) + keep_set = hasattr(self, "batch_set") and self._batch_set_name == active_set + keep_order = ( + keep_set + and self._batch_shuffle_state == state["shuffle_state"] + and self._batch_set_size == state["set_size"] + ) + if not keep_set: + self.__dict__.pop("batch_set", None) + if not keep_order: + self.__dict__.pop("_batch_indices", None) + self.set_count = state["set_count"] + self.iterator = state["iterator"] + self._batch_set_size = state["set_size"] + self._batch_shuffle_state = copy.deepcopy(state["shuffle_state"]) + def _derive_pair_margin( self, coord: np.ndarray, @@ -813,17 +911,32 @@ def _get_subdata( new_data[ii] = dd return new_data - def _load_batch_set(self, set_name: DPPath) -> None: - if not hasattr(self, "batch_set") or self.get_numb_set() > 1: - self.batch_set = self._load_set(set_name) + def _load_batch_set( + self, set_name: DPPath, *, shuffle_state: tuple | None = None + ) -> None: + if not hasattr(self, "batch_set") or self._batch_set_name != set_name: + batch_set = self._load_set(set_name) if self.modifier is not None: - self.modifier.modify_data(self.batch_set, self) - self.batch_set, _ = self._shuffle_data(self.batch_set) + self.modifier.modify_data(batch_set, self) + self.batch_set = batch_set + self._batch_set_name = set_name + # Always index the canonical source order. Shuffling an already + # shuffled cache cannot be reconstructed from just the latest RNG + # state, and repeatedly expands frames with numb_copy > 1. + nframes = self.batch_set["coord"].shape[0] + self._batch_indices = np.repeat( + np.arange(nframes, dtype=np.int64), + self.batch_set["numb_copy"].reshape(nframes), + ) + if shuffle_state is None: + self._batch_shuffle_state = dp_random.get_state() + dp_random.shuffle(self._batch_indices) else: - # Resident frames already contain their requested copies. - indices = np.arange(self.batch_set["coord"].shape[0], dtype=np.int64) - dp_random.shuffle(indices) - self.batch_set = self._get_subdata(self.batch_set, indices) + self._batch_shuffle_state = copy.deepcopy(shuffle_state) + rng = np.random.RandomState(0) + rng.set_state(shuffle_state) + rng.shuffle(self._batch_indices) + self._batch_set_size = len(self._batch_indices) self.reset_get_batch() def reset_get_batch(self) -> None: diff --git a/deepmd/utils/data_group.py b/deepmd/utils/data_group.py index fb4933bb01..df9b15b0bc 100644 --- a/deepmd/utils/data_group.py +++ b/deepmd/utils/data_group.py @@ -6,6 +6,9 @@ ) import logging +from copy import ( + deepcopy, +) from typing import ( TYPE_CHECKING, Any, @@ -205,6 +208,71 @@ def get_batch(self, sys_idx: int | None = None) -> dict[str, Any]: self._remaining -= 1 return self._batch_context(self.children[index].get_batch(), index) + def _sampling_config(self) -> dict[str, Any]: + return { + "probabilities": [ + None if value is None else float(value) + for value in self._assigned_probabilities + ], + "uniform": self._uniform, + } + + def state_dict(self) -> dict[str, Any]: + """Save this node's selection stream and each child's logical cursor.""" + return { + "version": 1, + "config": self._sampling_config(), + "rng": deepcopy(self._rng.bit_generator.state), + "remaining": self._remaining, + "pass_length": self._pass_length, + "probabilities": ( + self.component_probabilities.tolist() if self._remaining else None + ), + "children": [child.state_dict() for child in self.children], + } + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore the hierarchy without drawing batches or decoding frames.""" + if state.get("version") != 1: + raise ValueError("Unsupported grouped-data progress version") + if state.get("config") != self._sampling_config(): + raise ValueError("Grouped data progress has different sampling settings") + children = state.get("children") + if not isinstance(children, list) or len(children) != len(self.children): + raise ValueError("Grouped data progress has different children") + for key in ("remaining", "pass_length"): + if type(state.get(key)) is not int or state[key] < 0: + raise ValueError(f"Grouped data progress {key} must be non-negative") + if state["remaining"] > state["pass_length"]: + raise ValueError("Grouped data progress exceeds its pass length") + probabilities = state.get("probabilities") + if probabilities is not None: + probabilities = np.asarray(probabilities, dtype=np.float64) + if ( + probabilities.shape != (len(self.children),) + or not np.all(np.isfinite(probabilities)) + or np.any(probabilities < 0) + or not np.isclose(probabilities.sum(), 1.0, rtol=1e-12, atol=0) + ): + raise ValueError("Grouped data progress has invalid probabilities") + elif state["remaining"]: + raise ValueError( + "Grouped data progress is missing active-pass probabilities" + ) + if "rng" not in state: + raise ValueError("Grouped data progress is missing its selection RNG") + rng = np.random.default_rng(0) + rng.bit_generator.state = state["rng"] + for child, progress in zip(self.children, children, strict=True): + child.load_state_dict(progress) + self._rng = rng + self._remaining = state["remaining"] + self._pass_length = state["pass_length"] + self._probabilities = probabilities + self._cdf = None if probabilities is None else np.cumsum(probabilities) + if self._cdf is not None: + self._cdf[-1] = 1.0 + def _reader_pass_length(self) -> int: return compute_total_numb_batch( [child.get_batch_pass_length() for child in self.children], diff --git a/deepmd/utils/data_sampling.py b/deepmd/utils/data_sampling.py index 79fe1ea44d..ab22ddb918 100644 --- a/deepmd/utils/data_sampling.py +++ b/deepmd/utils/data_sampling.py @@ -19,18 +19,21 @@ def process_sys_probs( share the remaining probability in proportion to their batch counts. A fully assigned probability vector does not depend on batch counts. """ - probabilities = np.asarray(sys_probs, dtype=np.float64) + probabilities = np.asarray(sys_probs) + if not np.issubdtype(probabilities.dtype, np.floating): + probabilities = probabilities.astype(np.float64) if probabilities.ndim != 1: raise ValueError("Sampling probabilities must be one-dimensional.") if not np.all(np.isfinite(probabilities)): raise ValueError("Sampling probabilities must be finite.") assigned = probabilities >= 0 assigned_sum = float(probabilities[assigned].sum()) - if assigned_sum > 1.0 + 1e-8: + tolerance = max(1e-8, float(np.finfo(probabilities.dtype).eps)) + if assigned_sum > 1.0 + tolerance: raise ValueError("The sum of assigned probabilities must not exceed 1.") result = np.where(assigned, probabilities, 0.0) remaining = max(0.0, 1.0 - assigned_sum) - if remaining > 1e-8: + if remaining > tolerance: sizes = np.asarray(nbatch, dtype=np.float64) if probabilities.shape != sizes.shape: raise ValueError( @@ -119,7 +122,13 @@ def group_probabilities( ): raise ValueError("A group's prob must be non-negative or omitted.") resolved = process_sys_probs( - [-1.0 if probability is None else probability for probability in probabilities], + np.asarray( + [ + -1.0 if probability is None else probability + for probability in probabilities + ], + dtype=np.float64, + ), sizes, ) if np.any((np.asarray(sizes) == 0) & (resolved > 0)): diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index 24c62f7811..65379bca5c 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -520,6 +520,63 @@ def get_batch(self, sys_idx: int | None = None) -> dict: b_data = self.get_batch_mixed() return b_data + def state_dict(self) -> dict[str, Any]: + """Return logical sampling progress for all systems. + + The caller must save the shared ``dp_random`` state separately, + once for all training and validation data sources. + """ + return { + "version": 1, + "batch_size": [int(size) for size in self.batch_size], + "sys_probs": np.asarray(self.sys_probs).tolist(), + "mixed_systems": self.mixed_systems, + "pick_idx": int(self.pick_idx), + "systems": [data.state_dict() for data in self.data_systems], + } + + def load_state_dict(self, state: dict[str, Any]) -> None: + """Restore progress without materializing any frame data.""" + if state.get("version") != 1: + raise ValueError("Unsupported DeepmdDataSystem progress version") + if state.get("batch_size") != [int(size) for size in self.batch_size]: + raise ValueError("DeepmdDataSystem progress has different batch sizes") + try: + saved_probs = np.array(state.get("sys_probs"), dtype=float, copy=True) + except (TypeError, ValueError) as error: + raise ValueError( + "DeepmdDataSystem progress has invalid probabilities" + ) from error + if ( + saved_probs.shape != (self.nsystems,) + or not np.all(np.isfinite(saved_probs)) + or np.any(saved_probs < 0) + or not np.allclose(saved_probs, self.sys_probs, rtol=1e-12, atol=0) + ): + raise ValueError("DeepmdDataSystem progress has different probabilities") + if state.get("mixed_systems") != self.mixed_systems: + raise ValueError("DeepmdDataSystem progress has a different batching mode") + states = state.get("systems", []) + if len(states) != self.nsystems: + raise ValueError("DeepmdDataSystem progress has different systems") + pick_idx = state.get("pick_idx") + if type(pick_idx) is not int or not 0 <= pick_idx < self.nsystems: + raise ValueError("Invalid DeepmdDataSystem progress pick_idx") + # Validate every system before changing any reader's progress. + for data, data_state in zip(self.data_systems, states, strict=True): + data._validate_state_dict(data_state) + for data, data_state in zip(self.data_systems, states, strict=True): + data.load_state_dict(data_state) + # Tolerate recomputation roundoff, but retain the checkpoint's exact + # selection thresholds so the restored RNG resumes the same stream. + # Preserve an exactly matching dtype too: NumPy uses it to choose + # the tolerance for checking whether probabilities sum to one. + if np.array_equal(saved_probs, self.sys_probs): + saved_probs = np.array(self.sys_probs, copy=True) + self.sys_probs = saved_probs + self.pick_idx = pick_idx + self.__dict__.pop("test_data", None) + def get_batch_standard(self, sys_idx: int | None = None) -> dict: """Get a batch of data from the data systems in the standard way. diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index 8710693607..55e23b22f2 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -76,4 +76,18 @@ def shuffle(x: np.ndarray) -> None: _RANDOM_GENERATOR.shuffle(x) -__all__ = ["choice", "random", "seed", "shuffle"] +def get_state() -> tuple: + """Return the generator state using checkpoint-safe Python values.""" + kind, keys, pos, has_gauss, cached_gaussian = _RANDOM_GENERATOR.get_state() + return (kind, keys.tolist(), int(pos), int(has_gauss), float(cached_gaussian)) + + +def set_state(state: tuple) -> None: + """Restore a state returned by :func:`get_state`.""" + kind, keys, pos, has_gauss, cached_gaussian = state + _RANDOM_GENERATOR.set_state( + (kind, np.asarray(keys, dtype=np.uint32), pos, has_gauss, cached_gaussian) + ) + + +__all__ = ["choice", "get_state", "random", "seed", "set_state", "shuffle"] diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index fd12112818..353eda9bb5 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -199,6 +199,69 @@ checkpoints from the same family numbered above `N`. This prevents remnants of a longer earlier run from displacing the newly written checkpoint from the retention window. +PyTorch Exportable optimizer checkpoints also retain each rank's per-task +training-data cursor and task/system selection RNG. With the same immutable +data at the same locations, batching settings, and number of ranks, a restart +continues with the next LMDB or directory/HDF5 training batches rather than +starting their iteration again. Data-source +progress is compact and restored lazily; checkpoints do not store decoded +frames or shuffle permutations. Statistics, compilation warmup, and learning +curve evaluation do not consume this training progression. +Distributed ranks exchange this compact state at checkpoint boundaries, not +on every batch; rank-local filtering may legitimately advance their cursors +by different amounts. + +Restarting with a different number of ranks still restores the model, optimizer, +and training step, but warns and keeps newly initialized data cursors and +selection RNG instead of remapping the saved rank-local progress. Data may be +repeated, and exact data-order continuation is not guaranteed in this case. +Malformed progress envelopes remain errors; by default, so do incompatible data +settings at the same world size. + +PyTorch Exportable initializes rank-local data sampling from `training.seed` +and the rank before constructing data sources, including directory and HDF5 +sources. Task selection has an independent RNG shared across ranks, so local +data filtering cannot change the task sequence. Without a configured seed, +the chief initializes that task RNG and broadcasts its state once at startup. +Both RNG streams resume from checkpoints. The configured seed is also part of +the checkpoint's data-progress identity: changing it requires an explicit reset, +independently of the input format. Earlier experimental progress records without +seed/task-RNG identity likewise require a reset; checkpoints without data +progress retain their legacy restart behavior. + +Grouped inputs also preserve each group's child-selection RNG and every native +source's cursor, including nested groups combining directory/HDF5 and LMDB data. +The group topology, assigned probabilities, and uniform-sampling policy must +remain compatible when restoring progress. + +For PyTorch Exportable, set `training.reset_data_progress` to `true` when +intentionally restarting with relocated data or changed data settings. This +explicitly keeps newly initialized data cursors and task/system selection RNG +for **all** training tasks; model weights, optimizer, learning-rate schedule, +and training step still resume. A warning explains that batches may repeat and +exact data-order continuation is unavailable. The default is `false`, which +retains strict data-source compatibility checks. Malformed progress envelopes +remain errors even when resetting. + +Source paths are identity guards, not content fingerprints: equal frame counts, +atom types, or LMDB transaction numbers do not prove that two datasets contain +the same frames. Consequently, copying an unchanged dataset to another location +requires an explicit data-progress reset; automatic exact continuation after +relocation is not supported. Directory/HDF5 paths must also retain the same path +spelling, including the HDF5 group. `reset_data_progress` has no effect on fresh +training, `--init-model`, finetuning, frozen models, or older checkpoints without +data progress. + +Directory/HDF5 batches now shuffle indices into a canonical, unshuffled set. +For a single-set source, later passes therefore differ from older releases +that shuffled their already shuffled cache. This also prevents `numb_copy` +multiplicities from compounding over repeated passes. + +Older checkpoints without data progress retain the previous restart behavior. +`--init-model`, finetuning, and EMA-only deployment snapshots do not restore data +progress. This is a data-order guarantee, not a guarantee of bitwise-identical +training or of restoring every stochastic data transform. + **`--init-frz-model frozen_model.pb`**, initializes the training with an existing model that is stored in `frozen_model.pb`. **`--skip-neighbor-stat`** will skip calculating neighbor statistics if one is concerned about performance. Some features will be disabled. diff --git a/source/tests/common/dpmodel/test_data_progress.py b/source/tests/common/dpmodel/test_data_progress.py new file mode 100644 index 0000000000..797573c264 --- /dev/null +++ b/source/tests/common/dpmodel/test_data_progress.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Backend-neutral training cursor and selection-RNG checkpoint contract.""" + +import unittest + +from deepmd.dpmodel.train.data_progress import ( + collect_training_data_progress, + preserve_training_data_progress, + restore_training_data_progress, +) +from deepmd.utils import random as dp_random + + +class _Source: + def __init__(self) -> None: + self.cursor = 0 + + def state_dict(self) -> dict: + return {"cursor": self.cursor} + + def load_state_dict(self, state: dict) -> None: + self.cursor = state["cursor"] + + +class TestDataProgress(unittest.TestCase): + def test_task_cursors_and_selection_rng_roundtrip(self) -> None: + sources = {"water": _Source(), "ice": _Source()} + sources["water"].cursor = 2 + sources["ice"].cursor = 5 + dp_random.seed(17) + progress = collect_training_data_progress(sources) + expected = [int(dp_random.choice(4)) for _ in range(20)] + sources["water"].cursor = 20 + sources["ice"].cursor = 50 + restore_training_data_progress(sources, progress) + self.assertEqual(sources["water"].cursor, 2) + self.assertEqual(sources["ice"].cursor, 5) + self.assertEqual([int(dp_random.choice(4)) for _ in range(20)], expected) + + def test_diagnostics_restore_progress_even_on_failure(self) -> None: + source = _Source() + dp_random.seed(11) + before = collect_training_data_progress({"water": source}) + with self.assertRaisesRegex(RuntimeError, "diagnostic failed"): + with preserve_training_data_progress({"water": source}): + source.cursor += 1 + dp_random.choice(10) + raise RuntimeError("diagnostic failed") + self.assertEqual(collect_training_data_progress({"water": source}), before) + + def test_legacy_and_sources_without_protocol(self) -> None: + source = _Source() + source.cursor = 3 + restore_training_data_progress({"water": source}, None) + self.assertEqual(source.cursor, 3) + sources = {"unsupported": object(), "missing": None} + progress = collect_training_data_progress(sources) + self.assertEqual(progress["tasks"], {}) + restore_training_data_progress(sources, progress) + + def test_incompatible_payload_rejected_before_source_restore(self) -> None: + sources = {"water": _Source()} + progress = collect_training_data_progress(sources) + for changed in ( + {"version": 999}, + {"tasks": []}, + {"tasks": {"ice": {"cursor": 5}}}, + ): + with self.subTest(changed=changed), self.assertRaises(ValueError): + restore_training_data_progress(sources, progress | changed) + del progress["rng"] + with self.assertRaises(ValueError): + restore_training_data_progress(sources, progress) + self.assertEqual(sources["water"].cursor, 0) diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 15568f5a3f..27b4c90cdd 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -3067,6 +3067,26 @@ def test_a_second_iterator_does_not_retry_a_lost_pool(self) -> None: self._assert_same_batch(batch, self._reader.decode_batch([0, 1, 2, 3])) self.assertEqual(pool.submissions, submissions) + def test_restore_detects_a_stalled_prefetch(self) -> None: + """Restoring a cursor must retire a decoder that cannot finish its result.""" + iterator = LmdbBatchIterator( + self._reader, LmdbBatchSampler(self._reader, seed=7), num_workers=2 + ) + release = mock.patch.object(lmdb_data_module, "_release_lmdb_executor") + release.start() + self.addCleanup(release.stop) + self.addCleanup(iterator.close) + state = iterator.state_dict() + entry = lmdb_data_module._LmdbPoolEntry(executor=_StalledPool(), users=1) + iterator._pool = entry + running = Future() + self.assertTrue(running.set_running_or_notify_cancel()) + iterator._pending = lmdb_data_module._PendingBatch([0, 1], [running]) + iterator.load_state_dict(state) + self.assertFalse(entry.healthy) + self.assertIsNone(iterator._pending) + self.assertEqual(iterator.state_dict(), state) + @unittest.skipUnless( os.name == "posix" and sys.implementation.name == "cpython", "requires CPython's POSIX process-pool pipe", diff --git a/source/tests/common/dpmodel/test_lmdb_data_progress.py b/source/tests/common/dpmodel/test_lmdb_data_progress.py new file mode 100644 index 0000000000..57a38ca97a --- /dev/null +++ b/source/tests/common/dpmodel/test_lmdb_data_progress.py @@ -0,0 +1,338 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact LMDB restart state through the real pt_expt data adapter.""" + +import copy +import json +from contextlib import ( + ExitStack, + closing, +) +from unittest.mock import ( + patch, +) + +import lmdb +import msgpack +import pytest + +from deepmd.dpmodel.utils import lmdb_data as lmdb_module +from deepmd.utils.data import ( + DataRequirementItem, +) + +from .test_lmdb_data import ( + _create_lmdb, + _create_lmdb_with_system_ids, + _create_mixed_nloc_lmdb, +) + + +def _system(path, **kwargs): + from deepmd.pt_expt.utils.lmdb_dataset import ( + LmdbDataSystem, + ) + + options = {"batch_size": 2, "seed": 7, "num_workers": 0, **kwargs} + return LmdbDataSystem(path, ["O", "H"], **options) + + +def _draw(data, count): + """Record frame identity and the consumed pass, not the prefetched pass.""" + return [ + (data.get_batch()["fid"], data.get_batch_pass_length()) for _ in range(count) + ] + + +@pytest.fixture +def lmdb_path(tmp_path): + return _create_lmdb(str(tmp_path / "progress.lmdb"), nframes=10) + + +@pytest.mark.parametrize("workers", [0, 2]) +@pytest.mark.parametrize("seed", [None, 7]) +@pytest.mark.parametrize("consumed", [0, 1, 4, 5, 6]) +def test_logical_cursor_ignores_prefetch(lmdb_path, workers, seed, consumed): + """Restart before, inside and across a pass, with real process prefetch.""" + options = {"num_workers": workers, "seed": seed} + with closing(_system(lmdb_path, **options)) as original: + _draw(original, consumed) + saved = original.state_dict() + assert saved["epoch"] == consumed // 5 + assert saved["batch_index"] == consumed % 5 + assert isinstance(saved["seed"], int) + # Plain bounded metadata, not arrays, decoded data or permutations. + assert len(json.dumps(saved)) < 2000 + expected = _draw(original, 12) + + with closing(_system(lmdb_path, **options)) as resumed: + resumed.load_state_dict(saved) + assert resumed.state_dict() == saved + assert _draw(resumed, 12) == expected + if workers and consumed: + assert original._batch_iterator.started + assert resumed._batch_iterator.started + + +def test_restore_is_lazy_until_trainer_configuration(tmp_path): + """Reader/iterator construction predates label requirements and layout.""" + path = _create_mixed_nloc_lmdb(str(tmp_path / "mixed.lmdb")) + with lmdb.open(path) as env, env.begin(write=True) as txn: + for index in (1, 4, 7): + key = format(index, "012d").encode() + frame = msgpack.unpackb(txn.get(key), raw=False) + del frame["energies"] + txn.put(key, msgpack.packb(frame, use_bin_type=True)) + requirements = [DataRequirementItem("energy", 1)] + options = {"batch_size": "mix:24", "seed": 1} + build_batches = lmdb_module._build_all_batches + with closing(_system(path, **options)) as original: + original.add_data_requirements(requirements) + original.use_ragged_batches(True) + _draw(original, 3) + saved = original.state_dict() + expected = _draw(original, 15) + + with ( + patch.object( + lmdb_module, "_build_all_batches", side_effect=AssertionError("sampled") + ), + closing(_system(path, **options)) as resumed, + ): + with patch.object( + resumed._reader, "decode_batch", side_effect=AssertionError("decoded") + ): + resumed.load_state_dict(saved) + assert resumed.state_dict() == saved + # Registration must still be legal, and not plan a sampler pass. + resumed.add_data_requirements(requirements) + resumed.use_ragged_batches(True) + assert resumed._sampler._batches is None + # End the build guard before actually asking for the next batch. + with patch.object(lmdb_module, "_build_all_batches", new=build_batches): + assert _draw(resumed, 15) == expected + + +@pytest.mark.parametrize("workers", [0, 2]) +@pytest.mark.parametrize("consumed", [4, 5, 6, 9]) +def test_ragged_mixed_pass_lengths(tmp_path, workers, consumed): + """A five-batch pass followed by a four-batch pass retains its boundary.""" + path = _create_mixed_nloc_lmdb(str(tmp_path / "mixed.lmdb")) + options = {"batch_size": "mix:24", "seed": 1, "num_workers": workers} + with closing(_system(path, **options)) as original: + original.use_ragged_batches(True) + history = _draw(original, consumed) + assert history[0][1] == 5 + if consumed > 5: + assert history[5][1] == 4 + saved = original.state_dict() + expected = _draw(original, 15) + with closing(_system(path, **options)) as resumed: + resumed.load_state_dict(saved) + resumed.use_ragged_batches(True) + assert resumed.get_batch_pass_length() == expected[0][1] + assert _draw(resumed, 15) == expected + + +@pytest.mark.parametrize("seed", [None, 7]) +@pytest.mark.parametrize("consumed", [1, 3, 4]) +def test_equal_rank_cursors_reconstruct_each_shard(lmdb_path, seed, consumed): + """Equal progress reconstructs each rank's partition of the same pass.""" + with ExitStack() as stack: + originals = [ + stack.enter_context( + closing(_system(lmdb_path, seed=seed, rank=rank, world_size=2)) + ) + for rank in range(2) + ] + for data in originals: + _draw(data, consumed) + saved = originals[0].state_dict() + assert originals[1].state_dict() == saved + for rank, original in enumerate(originals): + expected = _draw(original, 10) + resumed = stack.enter_context( + closing(_system(lmdb_path, seed=seed, rank=rank, world_size=2)) + ) + resumed.load_state_dict(saved) + assert _draw(resumed, 10) == expected + + +def test_restore_drops_existing_prefetch(lmdb_path): + with closing(_system(lmdb_path, num_workers=2)) as data: + _draw(data, 3) + saved = data.state_dict() + expected = _draw(data, 9) + pool = data._batch_iterator._pool + with ( + patch.object(lmdb_module, "_build_all_batches") as build, + patch.object(data._reader, "decode_batch") as decode, + patch.object(lmdb_module, "_release_lmdb_executor") as release, + ): + data.load_state_dict(saved) + build.assert_not_called() + decode.assert_not_called() + release.assert_not_called() + assert data._batch_iterator._pool is pool + assert _draw(data, 9) == expected + + +@pytest.mark.parametrize("world_size", [1, 2]) +@pytest.mark.parametrize("diagnostic_reads", [0, 1]) +def test_same_pass_restore_reuses_sampling_plan( + lmdb_path, world_size, diagnostic_reads +): + """Display diagnostics must not rebuild a dataset-sized plan each time.""" + with closing(_system(lmdb_path, batch_size=1, world_size=world_size)) as data: + _draw(data, 1) + saved = data.state_dict() + frame_counts = data._sampler.batch_frame_counts() + with patch.object( + lmdb_module, + "_build_all_batches", + side_effect=AssertionError("rebuilt unchanged sampling plan"), + ): + _draw(data, diagnostic_reads) + data.load_state_dict(saved) + expected = _draw(data, 1) + data.load_state_dict(saved) + assert _draw(data, 1) == expected + assert data._sampler.batch_frame_counts() is frame_counts + + +@pytest.mark.parametrize("world_size", [1, 2]) +def test_cached_plan_before_layout_selection_is_invalidated(tmp_path, world_size): + path = _create_mixed_nloc_lmdb(str(tmp_path / "mixed.lmdb")) + options = {"batch_size": "mix:24", "seed": 1, "world_size": world_size} + with closing(_system(path, **options)) as original: + original.use_ragged_batches(True) + _draw(original, 1) + saved = original.state_dict() + expected = _draw(original, 10) + with closing(_system(path, **options)) as resumed: + # An earlier metadata query is allowed to have cached the default + # rectangular plan before the trainer selects the final layout. + assert resumed.nbatches[0] > 0 + resumed.load_state_dict(saved) + resumed.use_ragged_batches(True) + assert _draw(resumed, 10) == expected + + +def test_unequal_rank_cursors_restore_independently(lmdb_path): + """Local filtering may have consumed a different batch count per rank.""" + with ExitStack() as stack: + saved = [] + for rank in range(2): + original = stack.enter_context( + closing(_system(lmdb_path, rank=rank, world_size=2)) + ) + _draw(original, rank + 1) + saved.append(original.state_dict()) + expected = _draw(original, 10) + resumed = stack.enter_context( + closing(_system(lmdb_path, rank=rank, world_size=2)) + ) + resumed.load_state_dict(saved[-1]) + assert _draw(resumed, 10) == expected + assert saved[0]["batch_index"] != saved[1]["batch_index"] + + +def test_statistics_do_not_advance_training(lmdb_path): + with closing(_system(lmdb_path)) as data: + saved = data.state_dict() + data.get_stat_batch(0) + assert data.state_dict() == saved + + +def test_auto_prob_targets_are_reproduced_and_validated(tmp_path): + path = _create_lmdb_with_system_ids(str(tmp_path / "weighted.lmdb"), [4, 8]) + style = "prob_sys_size;0:1:0.7;1:2:0.3" + with closing(_system(path, auto_prob_style=style)) as original: + _draw(original, 4) + saved = original.state_dict() + expected = _draw(original, 12) + with closing(_system(path, auto_prob_style=style)) as resumed: + resumed.load_state_dict(saved) + assert _draw(resumed, 12) == expected + with closing(_system(path, auto_prob_style="prob_sys_size")) as changed: + with pytest.raises(ValueError, match="block_targets"): + changed.load_state_dict(saved) + + +@pytest.mark.parametrize( + ("changed", "field"), + [ + ({"batch_size": 3}, "batch_rule"), + ({"seed": 12}, "configured_seed"), + ({"world_size": 2}, "world_size"), + ], +) +def test_sampling_configuration_mismatch(lmdb_path, changed, field): + with closing(_system(lmdb_path)) as original: + _draw(original, 2) + saved = original.state_dict() + with closing(_system(lmdb_path, **changed)) as resumed: + with pytest.raises(ValueError, match=field): + resumed.load_state_dict(saved) + + +@pytest.mark.parametrize("field", ["ragged", "availability_keys"]) +def test_late_configuration_mismatch_is_checked_before_decode(tmp_path, field): + path = _create_mixed_nloc_lmdb(str(tmp_path / "mixed.lmdb")) + with closing(_system(path, batch_size="mix:24")) as original: + if field == "ragged": + original.use_ragged_batches(True) + else: + original.add_data_requirements([DataRequirementItem("energy", 1)]) + _draw(original, 2) + saved = original.state_dict() + with closing(_system(path, batch_size="mix:24")) as resumed: + resumed.load_state_dict(saved) + with ( + patch.object(resumed._reader, "decode_batch") as decode, + pytest.raises(ValueError, match=field), + ): + resumed.get_batch() + decode.assert_not_called() + + +@pytest.mark.parametrize( + ("key", "value"), + [("version", 99), ("epoch", -1), ("batch_index", 1.5), ("seed", None)], +) +def test_invalid_state_is_rejected(lmdb_path, key, value): + with closing(_system(lmdb_path)) as data: + saved = data.state_dict() + invalid = copy.deepcopy(saved) + invalid[key] = value + with pytest.raises(ValueError): + data.load_state_dict(invalid) + assert data.state_dict() == saved + + +def test_out_of_range_cursor_is_checked_lazily(lmdb_path): + with closing(_system(lmdb_path)) as data: + saved = data.state_dict() + saved["batch_index"] = 5 + data.load_state_dict(saved) + with pytest.raises(ValueError, match="out of range"): + data.get_batch() + + +def test_dataset_transaction_change_is_rejected(lmdb_path): + with closing(_system(lmdb_path)) as original: + saved = original.state_dict() + with lmdb.open(lmdb_path) as env, env.begin(write=True) as txn: + key = b"000000000000" + frame = msgpack.unpackb(txn.get(key), raw=False) + del frame["energies"] + txn.put(key, msgpack.packb(frame, use_bin_type=True)) + with closing(_system(lmdb_path)) as resumed: + with pytest.raises(ValueError, match="transaction"): + resumed.load_state_dict(saved) + + +def test_different_dataset_is_rejected(lmdb_path, tmp_path): + other = _create_lmdb(str(tmp_path / "other.lmdb"), nframes=10) + with closing(_system(lmdb_path)) as original, closing(_system(other)) as resumed: + with pytest.raises(ValueError, match="path"): + resumed.load_state_dict(original.state_dict()) diff --git a/source/tests/common/test_argcheck_training.py b/source/tests/common/test_argcheck_training.py index 3e783eefda..b954aedcec 100644 --- a/source/tests/common/test_argcheck_training.py +++ b/source/tests/common/test_argcheck_training.py @@ -9,6 +9,7 @@ import pytest from dargs.dargs import ( ArgumentError, + ArgumentTypeError, ) from deepmd.utils.argcheck import ( @@ -28,6 +29,30 @@ def test_ema_checkpoint_retention_is_left_for_runtime_inheritance() -> None: assert normalized["ema_ckpt_keep"] is None +@pytest.mark.parametrize("reset", [None, False, True]) +def test_reset_data_progress_is_an_opt_in_boolean(reset) -> None: + training_argument = training_args() + values = {"numb_steps": 1} + if reset is not None: + values["reset_data_progress"] = reset + + normalized = training_argument.normalize_value(values) + training_argument.check_value(normalized, strict=True) + + assert normalized["reset_data_progress"] is (False if reset is None else reset) + + +@pytest.mark.parametrize("reset", [None, 0, 1, "true", "false", [], {}]) +def test_reset_data_progress_rejects_non_boolean_values(reset) -> None: + training_argument = training_args() + normalized = training_argument.normalize_value( + {"numb_steps": 1, "reset_data_progress": reset} + ) + + with pytest.raises(ArgumentTypeError, match="reset_data_progress"): + training_argument.check_value(normalized, strict=True) + + def _training_input(data: dict[str, Any] | list[dict[str, Any]]) -> dict[str, Any]: return { "model": { diff --git a/source/tests/common/test_directory_data_progress.py b/source/tests/common/test_directory_data_progress.py new file mode 100644 index 0000000000..a57ae77fca --- /dev/null +++ b/source/tests/common/test_directory_data_progress.py @@ -0,0 +1,384 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Compact, lazy batch-progress restoration for directory and HDF5 sources.""" + +import copy +import json +from unittest.mock import ( + patch, +) + +import h5py +import numpy as np +import pytest + +from deepmd.utils import random as dp_random +from deepmd.utils.data import ( + DeepmdData, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) + + +def _write_system(root, *, hdf5=False, nsets=1, copies=False, offset=0, nframes=7): + """Write uniquely identifiable frames with optional zero/repeated copies.""" + arrays = {"type.raw": np.array([0, 0], dtype=np.int32)} + for index in range(nsets): + size = nframes + index + frame_ids = offset + index * 100 + np.arange(size) + prefix = f"set.{index:03d}/" + arrays[prefix + "coord.npy"] = np.repeat(frame_ids[:, None], 6, axis=1) + arrays[prefix + "box.npy"] = np.tile(np.eye(3).reshape(1, 9), (size, 1)) + arrays[prefix + "energy.npy"] = frame_ids[:, None] * 10 + if copies: + arrays[prefix + "numb_copy.npy"] = (np.arange(size) % 4)[:, None] + if hdf5: + filename = root.with_suffix(".h5") + with h5py.File(filename, "w") as handle: + for name, array in arrays.items(): + handle.create_dataset("system/" + name, data=array) + return str(filename) + "#/system" + root.mkdir() + for name, array in arrays.items(): + target = root / name + target.parent.mkdir(exist_ok=True) + if name.endswith(".raw"): + np.savetxt(target, array, fmt="%d") + else: + np.save(target, array) + return str(root) + + +def _assert_batch_equal(actual, expected): + assert actual.keys() == expected.keys() + for name in actual: + np.testing.assert_array_equal(actual[name], expected[name], err_msg=name) + + +@pytest.mark.parametrize("hdf5", [False, True], ids=["directory", "hdf5"]) +@pytest.mark.parametrize("nsets", [1, 3]) +@pytest.mark.parametrize("copies", [False, True]) +@pytest.mark.parametrize("consumed", [0, 1, 3, 11]) +def test_reader_resume_multiple_rounds(tmp_path, hdf5, nsets, copies, consumed): + """Resume before/within/after sets, including previously reshuffled caches.""" + path = _write_system(tmp_path / "data", hdf5=hdf5, nsets=nsets, copies=copies) + original = DeepmdData(path) + original.add("energy", 1) + dp_random.seed(482) + for _ in range(consumed): + original.get_batch(2) + progress = original.state_dict() + rng_state = dp_random.get_state() + expected = [(original.get_batch(2), dp_random.random()) for _ in range(40)] + + restored = DeepmdData(path) + restored.load_state_dict(progress) + # Requirements can be installed after logical state restoration. + restored.add("energy", 1) + dp_random.set_state(rng_state) + for batch, random_value in expected: + _assert_batch_equal(restored.get_batch(2), batch) + assert dp_random.random() == random_value + + +@pytest.mark.parametrize("hdf5", [False, True], ids=["directory", "hdf5"]) +def test_restore_is_lazy_and_invalidates_changed_requirements(tmp_path, hdf5): + path = _write_system(tmp_path / "data", hdf5=hdf5) + data = DeepmdData(path) + data.get_batch(2) + progress = data.state_dict() + data.get_test() + assert hasattr(data, "batch_set") + assert hasattr(data, "test_set") + data.add("energy", 1) + with ( + patch.object(data, "_load_set", side_effect=AssertionError("decoded set")), + patch.object(data, "_get_nframes", side_effect=AssertionError("scanned set")), + patch.object( + type(data.dirs[0]), "load_numpy", side_effect=AssertionError("read data") + ), + ): + data.load_state_dict(progress) + assert not hasattr(data, "batch_set") + assert not hasattr(data, "test_set") + with patch.object(data, "_load_set", wraps=data._load_set) as load: + batch = data.get_batch(2) + load.assert_called_once() + np.testing.assert_array_equal(batch["energy"][:, 0], batch["coord"][:, 0] * 10) + + +@pytest.mark.parametrize("extra_batches", [1, 5]) +def test_in_process_restore_reuses_canonical_cache(tmp_path, extra_batches): + """Display/statistics snapshots must not force another set decode.""" + path = _write_system(tmp_path / "data") + data = DeepmdData(path) + data.get_batch(2) + state = data.state_dict() + expected = data.get_batch(2) + for _ in range(extra_batches): + data.get_batch(2) + cached_set = data.batch_set + with patch.object(data, "_load_set", side_effect=AssertionError("decoded set")): + data.load_state_dict(state) + _assert_batch_equal(data.get_batch(2), expected) + assert data.batch_set is cached_set + + +def test_in_process_restore_discards_other_set_cache(tmp_path): + path = _write_system(tmp_path / "data", nsets=2) + data = DeepmdData(path) + data.get_batch(2) + state = data.state_dict() + expected = data.get_batch(2) + for _ in range(2): + data.get_batch(2) + assert data._batch_set_name == data.dirs[1] + data.load_state_dict(state) + assert not hasattr(data, "batch_set") + _assert_batch_equal(data.get_batch(2), expected) + + +def test_exhausted_set_is_not_decoded_on_resume(tmp_path): + path = _write_system(tmp_path / "data", nsets=2) + original = DeepmdData(path) + for _ in range(3): + original.get_batch(2) + state = original.state_dict() + restored = DeepmdData(path) + restored.load_state_dict(state) + with patch.object(restored, "_load_set", wraps=restored._load_set) as load: + restored.get_batch(2) + assert [call.args[0] for call in load.call_args_list] == [restored.dirs[1]] + + +def test_numb_copy_does_not_compound_between_rounds(tmp_path): + path = _write_system(tmp_path / "data", copies=True) + data = DeepmdData(path) + frame_ids = np.arange(7) + expected_ids = np.repeat(frame_ids, frame_ids % 4) + for _ in range(5): + batch = data.get_batch(len(expected_ids)) + np.testing.assert_array_equal(np.sort(batch["coord"][:, 0]), expected_ids) + assert data.batch_set["coord"].shape[0] == 7 + + +@pytest.mark.parametrize("copies", [False, True]) +def test_batch_order_matches_canonical_shuffle(tmp_path, copies): + """Index-only batching retains the established shuffle/RNG semantics.""" + path = _write_system(tmp_path / "data", nsets=3, copies=copies) + data = DeepmdData(path) + dp_random.seed(716) + for index in range(6): + rng_state = dp_random.get_state() + expected, _ = data._shuffle_data(data._load_set(data.dirs[index % 3])) + expected_rng = dp_random.get_state() + dp_random.set_state(rng_state) + # An oversized batch returns the entire set, preserving legacy + # short-batch behavior rather than crossing into the following set. + _assert_batch_equal(data.get_batch(1000), expected) + assert dp_random.get_state() == expected_rng + + +@pytest.mark.parametrize("hdf5", [False, True], ids=["directory", "hdf5"]) +@pytest.mark.parametrize("batch_size", [2, "mixed:3"]) +def test_system_resume_random_choices(tmp_path, hdf5, batch_size): + paths = [ + _write_system( + tmp_path / f"data{index}", + hdf5=hdf5, + nsets=index + 1, + copies=True, + offset=index * 1000, + ) + for index in range(2) + ] + original = DeepmdDataSystem(paths, batch_size, 1, sys_probs=[0.3, 0.7]) + dp_random.seed(918) + for _ in range(13): + original.get_batch() + progress = original.state_dict() + rng_state = dp_random.get_state() + expected = [original.get_batch() for _ in range(60)] + restored = DeepmdDataSystem(paths, batch_size, 1, sys_probs=[0.3, 0.7]) + with patch.object(DeepmdData, "_load_set", side_effect=AssertionError("decoded")): + restored.load_state_dict(progress) + dp_random.set_state(rng_state) + for batch in expected: + _assert_batch_equal(restored.get_batch(), batch) + + +@pytest.mark.parametrize("batch_size", [2, "mixed:3"]) +def test_system_resume_restores_exact_probabilities_after_roundoff( + tmp_path, batch_size +): + paths = [ + _write_system(tmp_path / f"data{index}", offset=index * 1000) + for index in range(2) + ] + # Put the first random draw exactly on a selection threshold. Even a + # one-ULP probability difference then changes the selected system. + dp_random.seed(1) + threshold = float(dp_random.random()) + saved_probs = np.array([threshold, 1 - threshold]) + rounded_threshold = np.nextafter(threshold, np.inf) + recomputed_probs = np.array([rounded_threshold, 1 - rounded_threshold]) + dp_random.seed(1) + assert dp_random.choice(2, p=saved_probs) == 1 + dp_random.seed(1) + assert dp_random.choice(2, p=recomputed_probs) == 0 + + original = DeepmdDataSystem(paths, batch_size, 1, sys_probs=saved_probs.tolist()) + state = original.state_dict() + dp_random.seed(1) + rng_state = dp_random.get_state() + expected = [original.get_batch() for _ in range(20)] + expected_rng_state = dp_random.get_state() + + restored = DeepmdDataSystem( + paths, batch_size, 1, sys_probs=recomputed_probs.tolist() + ) + with patch.object(DeepmdData, "_load_set", side_effect=AssertionError("decoded")): + restored.load_state_dict(state) + np.testing.assert_array_equal(restored.sys_probs, saved_probs) + # The restored probabilities must not alias the caller's checkpoint. + state["sys_probs"][0] = 0 + np.testing.assert_array_equal(restored.sys_probs, saved_probs) + dp_random.set_state(rng_state) + for batch in expected: + _assert_batch_equal(restored.get_batch(), batch) + assert dp_random.get_state() == expected_rng_state + + +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +def test_system_resume_preserves_exactly_matching_probability_dtype(tmp_path, dtype): + paths = [ + _write_system(tmp_path / f"data{index}", offset=index * 1000) + for index in range(3) + ] + # Float32 thirds are valid for NumPy choice but their float64 sum is + # outside its tighter float64 normalization tolerance. + probabilities = list(np.full(3, 1 / 3, dtype=dtype)) + original = DeepmdDataSystem(paths, 2, 1, sys_probs=probabilities) + state = original.state_dict() + dp_random.seed(193) + rng_state = dp_random.get_state() + expected = [original.get_batch() for _ in range(20)] + + restored = DeepmdDataSystem(paths, 2, 1, sys_probs=probabilities) + restored.load_state_dict(state) + assert restored.sys_probs.dtype == dtype + np.testing.assert_array_equal(restored.sys_probs, probabilities) + dp_random.set_state(rng_state) + for batch in expected: + _assert_batch_equal(restored.get_batch(), batch) + + +@pytest.mark.parametrize( + "saved_probs", + [ + None, + 0.5, + [0.5], + [[0.3, 0.7]], + [0.3, 0.7, 0.0], + [float("nan"), 0.7], + [float("inf"), 0.7], + [-0.3, 1.3], + ["invalid", 0.7], + [0.4, 0.6], + [0.3 + 1e-8, 0.7 - 1e-8], + ], +) +def test_invalid_system_probabilities_are_rejected_without_mutation( + tmp_path, saved_probs +): + paths = [_write_system(tmp_path / f"data{index}") for index in range(2)] + data = DeepmdDataSystem(paths, 2, 1, sys_probs=[0.3, 0.7]) + data.get_batch() + initial = data.state_dict() + invalid = copy.deepcopy(initial) + invalid["sys_probs"] = saved_probs + with pytest.raises(ValueError, match="probabilities"): + data.load_state_dict(invalid) + assert data.state_dict() == initial + + +@pytest.mark.parametrize("saved_zero", [False, True]) +def test_system_probability_roundoff_does_not_change_excluded_systems( + tmp_path, saved_zero +): + paths = [_write_system(tmp_path / f"data{index}") for index in range(2)] + positive = np.nextafter(0.0, 1.0) + current_probs = [positive if saved_zero else 0.0, 1.0] + data = DeepmdDataSystem(paths, 2, 1, sys_probs=current_probs) + initial = data.state_dict() + invalid = copy.deepcopy(initial) + invalid["sys_probs"] = [0.0 if saved_zero else positive, 1.0] + with pytest.raises(ValueError, match="probabilities"): + data.load_state_dict(invalid) + assert data.state_dict() == initial + + +def test_system_probabilities_are_not_restored_before_source_validation(tmp_path): + paths = [_write_system(tmp_path / f"data{index}") for index in range(2)] + data = DeepmdDataSystem(paths, 2, 1, sys_probs=[0.3, 0.7]) + initial = data.state_dict() + invalid = copy.deepcopy(initial) + invalid["sys_probs"][0] = np.nextafter(0.3, 1.0) + invalid["systems"][-1]["version"] = 99 + with pytest.raises(ValueError): + data.load_state_dict(invalid) + assert data.state_dict() == initial + + +def test_checkpoint_is_compact_and_detached(tmp_path): + path = _write_system(tmp_path / "data", nframes=10000) + data = DeepmdData(path) + data.get_batch(2) + state = data.state_dict() + # The fixed-size RNG state is JSON-serializable; no frame permutation or + # label arrays should scale this snapshot with the number of frames. + assert len(json.dumps(state)) < 12000 + assert state["set_size"] == 10000 + state["identity"]["prefix_sum"][0] = -1 + state["shuffle_state"][1][0] = 0 + assert data.state_dict()["identity"]["prefix_sum"][0] == 10000 + assert data.state_dict()["shuffle_state"][1][0] != 0 + + +@pytest.mark.parametrize("change", ["source", "batch_size", "sys_probs", "order"]) +def test_incompatible_system_progress_is_rejected(tmp_path, change): + paths = [_write_system(tmp_path / f"data{index}") for index in range(2)] + original = DeepmdDataSystem(paths, 2, 1, sys_probs=[0.3, 0.7]) + original.get_batch() + state = original.state_dict() + kwargs = {"batch_size": 2, "test_size": 1, "sys_probs": [0.3, 0.7]} + if change == "source": + paths[0] = _write_system(tmp_path / "other") + elif change == "batch_size": + kwargs["batch_size"] = 1 + elif change == "sys_probs": + kwargs["sys_probs"] = [0.4, 0.6] + elif change == "order": + paths.reverse() + restored = DeepmdDataSystem(paths, **kwargs) + initial = restored.state_dict() + with pytest.raises(ValueError): + restored.load_state_dict(state) + assert restored.state_dict() == initial + + +@pytest.mark.parametrize( + "field,value", + [("version", 99), ("iterator", -1), ("set_count", -1), ("shuffle_state", None)], +) +def test_invalid_reader_progress_is_rejected(tmp_path, field, value): + path = _write_system(tmp_path / "data") + data = DeepmdData(path) + data.get_batch(2) + initial = data.state_dict() + invalid = copy.deepcopy(initial) + invalid[field] = value + with pytest.raises(ValueError): + data.load_state_dict(invalid) + assert data.state_dict() == initial diff --git a/source/tests/pt_expt/test_checkpoint_guard.py b/source/tests/pt_expt/test_checkpoint_guard.py index d3955902c2..54525725ae 100644 --- a/source/tests/pt_expt/test_checkpoint_guard.py +++ b/source/tests/pt_expt/test_checkpoint_guard.py @@ -7,6 +7,9 @@ from datetime import ( timedelta, ) +from functools import ( + partial, +) from pathlib import ( Path, ) @@ -17,6 +20,7 @@ patch, ) +import numpy as np import pytest import torch import torch.distributed as dist @@ -62,6 +66,11 @@ def _make_trainer(backend, directory, *, rank=0, stage=0, distributed=False): trainer = trainer_class.__new__(trainer_class) trainer.rank = rank trainer.is_distributed = distributed + if backend == "pt_expt": + trainer.training_data_by_task = {} + trainer._configured_data_seed = None + trainer._task_rng = np.random.default_rng(0) + trainer.world_size = dist.get_world_size() if distributed else 1 trainer.sharding = ShardingPolicy(stage=stage) trainer.model = torch.nn.Linear(2, 1, device="cpu") with torch.no_grad(): @@ -124,6 +133,11 @@ def _best_writer(trainer, ema): def _validate(trainer, validator, *, step, metric, ema=False): + save_checkpoint = _best_writer(trainer, ema) + if isinstance(trainer, PTExptTrainer) and not ema: + save_checkpoint = partial( + save_checkpoint, data_progress=trainer._collect_data_progress() + ) with patch.object( validator, "evaluate_all_systems", return_value={validator.metric_key: metric} ): @@ -131,7 +145,7 @@ def _validate(trainer, validator, *, step, metric, ema=False): step_id=step, display_step=step, lr=1e-3, - save_checkpoint=_best_writer(trainer, ema), + save_checkpoint=save_checkpoint, prepare_checkpoint=trainer.ensure_finite_gradients_for_checkpoint, ) diff --git a/source/tests/pt_expt/test_data_progress_restart.py b/source/tests/pt_expt/test_data_progress_restart.py new file mode 100644 index 0000000000..9ad74a6385 --- /dev/null +++ b/source/tests/pt_expt/test_data_progress_restart.py @@ -0,0 +1,1023 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Training-stream continuity through native pt_expt checkpoints.""" + +from copy import ( + deepcopy, +) +from pathlib import ( + Path, +) +from shutil import ( + copyfile, + copytree, +) +from typing import ( + Any, +) +from unittest.mock import ( + Mock, + patch, +) + +import h5py +import lmdb +import msgpack +import numpy as np +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +from deepmd.dpmodel.train import ( + RankContext, +) +from deepmd.dpmodel.train.data_progress import ( + DATA_PROGRESS_CHECKPOINT_KEY, + collect_training_data_progress, + preserve_training_data_progress, +) +from deepmd.pt_expt.entrypoints.main import ( + get_trainer, +) +from deepmd.pt_expt.train import training as training_module +from deepmd.pt_expt.train.ema import ( + EMA_CHECKPOINT_KEY, +) +from deepmd.pt_expt.train.training import ( + Trainer, +) +from deepmd.pt_expt.utils.env import ( + DEVICE, +) +from deepmd.pt_expt.utils.lmdb_dataset import ( + LmdbDataSystem, +) +from deepmd.utils import random as dp_random +from deepmd.utils.argcheck import ( + normalize, +) +from deepmd.utils.compat import ( + update_deepmd_input, +) +from deepmd.utils.data import ( + DataRequirementItem, +) + +from ..common.dpmodel.test_lmdb_data import ( + _create_lmdb, +) + + +def _write_directory(path: Path) -> None: + path.mkdir() + np.savetxt(path / "type.raw", [0, 1, 1], fmt="%d") + (path / "type_map.raw").write_text("O H\n") + rng = np.random.RandomState(24) + for index in range(2): + target = path / f"set.{index:03d}" + target.mkdir() + np.save(target / "coord.npy", rng.uniform(0.5, 2.5, (2, 9))) + np.save(target / "box.npy", np.tile(np.eye(3).reshape(1, 9) * 8, (2, 1))) + np.save(target / "energy.npy", np.arange(index * 2, index * 2 + 2) * 0.01) + np.save(target / "force.npy", rng.normal(size=(2, 9)) * 0.01) + + +def _write_lmdb(path: Path, *, mixed: bool) -> None: + def encode(array: np.ndarray) -> dict: + return { + "nd": None, + "type": str(array.dtype), + "kind": "", + "shape": list(array.shape), + "data": array.tobytes(), + } + + nlocs = [3, 4, 3, 4, 3, 4] if mixed else [3] * 4 + rng = np.random.RandomState(24) + with lmdb.open(str(path), map_size=10 * 1024 * 1024) as env: + with env.begin(write=True) as txn: + txn.put( + b"__metadata__", + msgpack.packb( + { + "nframes": len(nlocs), + "frame_idx_fmt": "012d", + "frame_nlocs": nlocs, + "type_map": ["O", "H"], + "system_info": {"natoms": [1, 2], "nframes": len(nlocs)}, + }, + use_bin_type=True, + ), + ) + for index, nloc in enumerate(nlocs): + frame = { + "atom_numbs": [1, nloc - 1], + "atom_names": ["O", "H"], + "atom_types": encode(np.array([0] + [1] * (nloc - 1))), + "orig": encode(np.zeros(3)), + "cells": encode(np.eye(3) * 8), + "coords": encode(rng.uniform(0.5, 2.5, (nloc, 3))), + "energies": encode(np.array(index * 0.01)), + "forces": encode(rng.normal(size=(nloc, 3)) * 0.01), + } + txn.put( + f"{index:012d}".encode(), + msgpack.packb(frame, use_bin_type=True), + ) + + +def _write_hdf5(path: Path) -> None: + """Store the directory fixture in a single HDF5 system.""" + directory = path.with_suffix("") + _write_directory(directory) + with h5py.File(path, "w") as handle: + handle.create_dataset("system/type.raw", data=[0, 1, 1]) + handle.create_dataset("system/type_map.raw", data=[b"O", b"H"]) + for array in directory.glob("set.*/*.npy"): + handle.create_dataset( + "system/" + array.relative_to(directory).as_posix(), data=np.load(array) + ) + + +def _config(path: str | Path, *, mode: str, multitask: bool = True) -> dict: + descriptor = { + "type": "se_e2_a", + "sel": [4, 4], + "rcut_smth": 0.5, + "rcut": 3.0, + "neuron": [4, 8], + "axis_neuron": 2, + "seed": 1, + } + if mode == "ragged": + descriptor.update(type="dpa1", sel=8, attn_layer=0) + model = { + "type_map": ["O", "H"], + "descriptor": descriptor, + "fitting_net": {"neuron": [4], "seed": 1}, + "data_stat_nbatch": 1, + } + loss = { + "type": "ener", + "start_pref_e": 1, + "limit_pref_e": 1, + "start_pref_f": 1, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + } + dataset = { + "systems": [str(path)] if mode == "directory" else str(path), + "batch_size": "mix:7" if mode == "ragged" else 1, + } + task_data = { + "training_data": deepcopy(dataset), + "validation_data": {**dataset, "numb_btch": 2}, + } + training = { + "numb_steps": 20, + "seed": 19, + "disp_freq": 3, + "save_freq": 10, + } + config = { + "learning_rate": { + "type": "exp", + "start_lr": 1e-3, + "stop_lr": 1e-5, + "decay_steps": 5, + }, + "training": training, + } + if multitask: + keys = ["first", "second"] + config["model"] = {"model_dict": {key: deepcopy(model) for key in keys}} + config["loss_dict"] = {key: deepcopy(loss) for key in keys} + training["data_dict"] = {key: deepcopy(task_data) for key in keys} + training["model_prob"] = {"first": 0.5, "second": 0.5} + else: + config.update(model=model, loss=loss) + training.update(task_data) + return normalize(update_deepmd_input(config, warning=False), multi_task=multitask) + + +def _record_training(trainer: Trainer, *, data_rng_draws: int = 0) -> list: + """Record real optimizer inputs, excluding diagnostic reads of the data.""" + records = [] + get_data = trainer.get_data + train_step = trainer.train_step + in_train_step = False + + def record_data(is_train: bool = True, task_key: str = "Default") -> tuple: + inputs, labels = get_data(is_train=is_train, task_key=task_key) + if is_train and in_train_step: + if data_rng_draws: + dp_random.random(data_rng_draws) + records.append( + ( + task_key, + { + key: value.detach().cpu().clone() + for key, value in inputs.items() + if torch.is_tensor(value) + }, + labels["energy"].detach().cpu().clone(), + ) + ) + return inputs, labels + + def record_step(task, step): + nonlocal in_train_step + in_train_step = True + try: + return train_step(task, step) + finally: + in_train_step = False + + trainer.get_data = record_data + trainer.train_step = record_step + trainer.run() + return records + + +def _assert_records_equal(expected: list, actual: list) -> None: + assert len(expected) == len(actual) + for (task, inputs, energy), (other_task, other_inputs, other_energy) in zip( + expected, actual, strict=True + ): + assert task == other_task + assert inputs.keys() == other_inputs.keys() + for key, value in inputs.items(): + torch.testing.assert_close(value, other_inputs[key], rtol=0, atol=0) + torch.testing.assert_close(energy, other_energy, rtol=0, atol=0) + + +@pytest.mark.parametrize("mode", ["directory", "hdf5", "grouped"]) +def test_setup_statistics_preserve_training_progress( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: str +) -> None: + """Fresh setup statistics consume samples without consuming training state.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / ("data.h5" if mode == "hdf5" else "data") + if mode == "hdf5": + _write_hdf5(path) + else: + _write_directory(path) + system = str(path) + ("#/system" if mode == "hdf5" else "") + config = _config(system, mode="directory") + if mode == "grouped": + second = tmp_path / "second" + _write_directory(second) + for task in config["training"]["data_dict"].values(): + task["training_data"] = [ + {"systems": [system], "batch_size": 1, "prob": 0.4}, + {"systems": [str(second)], "batch_size": 1, "prob": 0.6}, + ] + for model in config["model"]["model_dict"].values(): + model["data_stat_nbatch"] = 3 + snapshots = [] + make_stat_input = training_module.make_stat_input + + def record_statistics(data: Any, nbatches: int, filter_enabled: bool) -> list[dict]: + before = collect_training_data_progress({"statistics": data}) + assert "statistics" in before["tasks"] + sampled = make_stat_input(data, nbatches, filter_enabled) + assert collect_training_data_progress({"statistics": data}) != before + snapshots.append((data, before)) + return sampled + + monkeypatch.setattr(training_module, "make_stat_input", record_statistics) + trainer = get_trainer(config) + try: + assert len(snapshots) == len(trainer.training_data_by_task) + for data, expected in snapshots: + assert collect_training_data_progress({"statistics": data}) == expected + finally: + trainer._close_data_systems() + + +@pytest.mark.parametrize( + ("mode", "seed"), + [ + ("directory", 19), + ("lmdb", 19), + ("ragged", 19), + ("grouped", 19), + ("grouped", None), + ], +) +def test_restart_matches_uninterrupted_training( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: str, seed: int | None +) -> None: + """Real training matches across passes, display, validation and setup.""" + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + path = tmp_path / ("data" if mode in {"directory", "grouped"} else "data.lmdb") + if mode in {"directory", "grouped"}: + _write_directory(path) + else: + _write_lmdb(path, mixed=mode == "ragged") + config = _config(path, mode="directory" if mode == "grouped" else mode) + config["training"]["seed"] = seed + if mode == "grouped": + lmdb_paths = [tmp_path / f"part_{index}.lmdb" for index in range(2)] + for lmdb_path in lmdb_paths: + _write_lmdb(lmdb_path, mixed=False) + for task in config["training"]["data_dict"].values(): + task["training_data"] = [ + {"systems": [str(path)], "batch_size": 1, "prob": 0.4}, + { + "systems": [str(item) for item in lmdb_paths], + "batch_size": "mix:7", + "auto_prob": "prob_uniform", + "prob": 0.6, + }, + ] + dp_random.seed(19) + uninterrupted = get_trainer(deepcopy(config)) + if mode == "ragged": + assert all( + source._reader.ragged_batches + for source in uninterrupted.training_data_by_task.values() + ) + expected = _record_training(uninterrupted) + checkpoint = tmp_path / "model.ckpt-10.pt" + state = torch.load(checkpoint, weights_only=True) + assert DATA_PROGRESS_CHECKPOINT_KEY in state + # Restart construction and fresh validation readers may consume different + # random draws; none may leak into the resumed training stream. + dp_random.seed(9876) + restarted = get_trainer(deepcopy(config), restart_model=str(checkpoint)) + assert restarted.start_step == 10 + actual = _record_training(restarted) + _assert_records_equal(expected[10:], actual) + for key, value in uninterrupted._unwrapped.state_dict().items(): + if torch.is_tensor(value): + torch.testing.assert_close( + value, restarted._unwrapped.state_dict()[key], rtol=1e-12, atol=1e-12 + ) + + +def test_best_validation_checkpoint_resumes_training(tmp_path, monkeypatch) -> None: + """The real full-validation callback writes a restartable snapshot.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "data" + _write_directory(path) + config = _config(path, mode="directory", multitask=False) + config["training"].update(numb_steps=6, save_freq=0) + config["validating"] = { + "full_validation": True, + "full_val_start": 3, + "validation_freq": 3, + "save_best": True, + "save_best_dir": str(tmp_path / "best"), + "max_best_ckpt": 2, + } + original = get_trainer(deepcopy(config)) + expected = _record_training(original) + checkpoint = next((tmp_path / "best").glob("*-3.t-*.pt")) + restarted = get_trainer(deepcopy(config), restart_model=str(checkpoint)) + assert restarted.start_step == 3 + _assert_records_equal(expected[3:], _record_training(restarted)) + + +@pytest.mark.parametrize("reset", [False, True]) +def test_legacy_init_finetune_and_deployment( + tmp_path, monkeypatch, caplog, reset +) -> None: + """Only a restart restores progress; deployment snapshots omit it.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "data" + _write_directory(path) + config = _config(path, mode="directory", multitask=False) + config["training"].update( + numb_steps=3, save_freq=3, enable_ema=True, reset_data_progress=reset + ) + trainer = get_trainer(deepcopy(config)) + assert trainer.start_step == 0 + trainer.run() + checkpoint = tmp_path / "model.ckpt-3.pt" + state = torch.load(checkpoint, weights_only=True) + assert DATA_PROGRESS_CHECKPOINT_KEY in state + ema_state = torch.load(tmp_path / "model_ema.ckpt-3.pt", weights_only=True) + assert DATA_PROGRESS_CHECKPOINT_KEY not in ema_state + assert "optimizer" not in ema_state + legacy = tmp_path / "legacy.pt" + state.pop(DATA_PROGRESS_CHECKPOINT_KEY) + torch.save(state, legacy) + restarted = get_trainer(deepcopy(config), restart_model=str(legacy)) + assert restarted.start_step == 3 + restarted._close_data_systems() + # A deliberately unsupported version proves these paths do not restore it. + state[DATA_PROGRESS_CHECKPOINT_KEY] = {"version": -1} + torch.save(state, checkpoint) + for option in ("init_model", "finetune_model"): + fresh = get_trainer(deepcopy(config), **{option: str(checkpoint)}) + assert fresh.start_step == 0 + fresh._close_data_systems() + assert "training.reset_data_progress is true" not in caplog.text + + +@pytest.mark.parametrize("mode", ["directory", "hdf5", "lmdb"]) +def test_relocated_restart_can_reset_only_data_progress( + tmp_path, monkeypatch, caplog, mode +) -> None: + """Opting out after a move has exactly the legacy restart semantics.""" + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + if mode == "lmdb": + path = tmp_path / "data.lmdb" + _write_lmdb(path, mixed=False) + elif mode == "hdf5": + path = tmp_path / "data.h5" + _write_hdf5(path) + else: + path = tmp_path / "data" + _write_directory(path) + system = str(path) + ("#/system" if mode == "hdf5" else "") + config_mode = "lmdb" if mode == "lmdb" else "directory" + config = _config(system, mode=config_mode) + config["training"].update(numb_steps=3, save_freq=3, enable_ema=True) + original = get_trainer(deepcopy(config)) + original.run() + checkpoint = tmp_path / "model.ckpt-3.pt" + state = torch.load(checkpoint, weights_only=True) + # Native restart maps every optimizer tensor, including Adam's step. + checkpoint_optimizer = torch.load( + checkpoint, map_location=DEVICE, weights_only=True + )["optimizer"]["state"] + assert DATA_PROGRESS_CHECKPOINT_KEY in state + legacy_state = deepcopy(state) + del legacy_state[DATA_PROGRESS_CHECKPOINT_KEY] + legacy_path = tmp_path / "legacy.pt" + torch.save(legacy_state, legacy_path) + + relocated = path.with_name("relocated" + path.suffix) + if mode == "hdf5": + copyfile(path, relocated) + else: + copytree(path, relocated) + system = str(relocated) + ("#/system" if mode == "hdf5" else "") + config = _config(system, mode=config_mode) + config["training"].update(numb_steps=9, save_freq=3, enable_ema=True) + assert config["training"]["reset_data_progress"] is False + with pytest.raises(ValueError, match=r"[Pp]rogress"): + get_trainer(deepcopy(config), restart_model=str(checkpoint)) + assert "training.reset_data_progress is true" not in caplog.text + + dp_random.seed(9876) + legacy = get_trainer(deepcopy(config), restart_model=str(legacy_path)) + expected_lr = legacy.scheduler.get_last_lr() + expected_optimizer = deepcopy(legacy.optimizer.state_dict()) + expected_progress = collect_training_data_progress(legacy.training_data_by_task) + expected = _record_training(legacy) + dp_random.seed(9876) + config["training"]["reset_data_progress"] = True + reset = get_trainer(deepcopy(config), restart_model=str(checkpoint)) + assert "training.reset_data_progress is true" in caplog.text + assert "Batches may repeat" in caplog.text + assert "exact data-order continuation is unavailable" in caplog.text + assert reset.start_step == legacy.start_step == 3 + assert reset.scheduler.get_last_lr() == expected_lr + assert ( + collect_training_data_progress(reset.training_data_by_task) == expected_progress + ) + torch.testing.assert_close(reset.optimizer.state_dict(), expected_optimizer) + torch.testing.assert_close( + reset.optimizer.state_dict()["state"], checkpoint_optimizer + ) + torch.testing.assert_close(reset.model_ema.state_dict(), state[EMA_CHECKPOINT_KEY]) + for key, value in state["model"].items(): + if torch.is_tensor(value): + torch.testing.assert_close(value, reset._unwrapped.state_dict()[key]) + _assert_records_equal(expected, _record_training(reset)) + assert reset.scheduler.get_last_lr() == legacy.scheduler.get_last_lr() + torch.testing.assert_close( + reset.optimizer.state_dict(), legacy.optimizer.state_dict() + ) + torch.testing.assert_close( + reset.model_ema.state_dict(), legacy.model_ema.state_dict() + ) + for key, value in legacy._unwrapped.state_dict().items(): + if torch.is_tensor(value): + torch.testing.assert_close( + value, reset._unwrapped.state_dict()[key], rtol=1e-12, atol=1e-12 + ) + + +@pytest.mark.parametrize("mode", ["directory", "hdf5", "lmdb"]) +def test_configured_seed_is_format_independent( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, mode: str +) -> None: + """Fresh streams are seeded consistently; exact restarts reject seed changes.""" + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + path = ( + tmp_path / {"directory": "data", "hdf5": "data.h5", "lmdb": "data.lmdb"}[mode] + ) + if mode == "lmdb": + _write_lmdb(path, mixed=False) + elif mode == "hdf5": + _write_hdf5(path) + else: + _write_directory(path) + system = str(path) + ("#/system" if mode == "hdf5" else "") + config = _config( + system, mode="lmdb" if mode == "lmdb" else "directory", multitask=False + ) + config["training"].update(seed=7, numb_steps=1, save_freq=1) + states = [] + for external_seed in (1, 999): + dp_random.seed(external_seed) + trainer = get_trainer(deepcopy(config)) + try: + states.append(collect_training_data_progress(trainer.training_data_by_task)) + finally: + trainer._close_data_systems() + assert states[0] == states[1] + + original = get_trainer(deepcopy(config)) + original.run() + checkpoint = str(tmp_path / "model.ckpt-1.pt") + config["training"]["seed"] = 12 + with pytest.raises(ValueError, match=r"training.seed changed"): + get_trainer(deepcopy(config), restart_model=checkpoint) + config["training"]["reset_data_progress"] = True + restarted = get_trainer(config, restart_model=checkpoint) + try: + assert restarted.start_step == 1 + assert restarted._collect_data_progress()["configured_seed"] == 12 + finally: + restarted._close_data_systems() + + +@pytest.mark.parametrize("mode", ["directory", "lmdb"]) +def test_changed_world_size_retains_legacy_optimizer_restart( + tmp_path, monkeypatch, caplog, mode +) -> None: + """A changed topology keeps native restart state, but not saved data order.""" + monkeypatch.chdir(tmp_path) + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + path = tmp_path / ("data" if mode == "directory" else "data.lmdb") + if mode == "directory": + _write_directory(path) + else: + _write_lmdb(path, mixed=False) + config = _config(path, mode=mode, multitask=False) + config["training"].update(numb_steps=3, save_freq=3) + original = get_trainer(deepcopy(config)) + original.run() + checkpoint = tmp_path / "model.ckpt-3.pt" + state = torch.load(checkpoint, weights_only=True) + checkpoint_optimizer = torch.load( + checkpoint, map_location=DEVICE, weights_only=True + )["optimizer"]["state"] + progress = state.pop(DATA_PROGRESS_CHECKPOINT_KEY) + legacy_path = tmp_path / "legacy.pt" + torch.save(state, legacy_path) + # Exercise the topology fallback at the native checkpoint-load boundary. + # The real 1-to-2-rank integration case remains in TestDDPRestart. + progress.update(world_size=2, ranks=progress["ranks"] * 2) + state[DATA_PROGRESS_CHECKPOINT_KEY] = progress + changed_path = tmp_path / "changed-world.pt" + torch.save(state, changed_path) + config["training"]["numb_steps"] = 6 + + dp_random.seed(9876) + legacy = get_trainer(deepcopy(config), restart_model=str(legacy_path)) + expected_lr = legacy.scheduler.get_last_lr() + expected = _record_training(legacy) + dp_random.seed(9876) + changed = get_trainer(deepcopy(config), restart_model=str(changed_path)) + assert "world size changed from 2 to 1" in caplog.text + assert "exact data-order continuation is unavailable" in caplog.text + assert changed.start_step == 3 + assert changed.scheduler.get_last_lr() == expected_lr + torch.testing.assert_close( + changed.optimizer.state_dict()["state"], checkpoint_optimizer + ) + for key, value in state["model"].items(): + if torch.is_tensor(value): + torch.testing.assert_close(value, changed._unwrapped.state_dict()[key]) + _assert_records_equal(expected, _record_training(changed)) + for key, value in legacy._unwrapped.state_dict().items(): + if torch.is_tensor(value): + torch.testing.assert_close(value, changed._unwrapped.state_dict()[key]) + + +def test_chief_display_preserves_training_progress(tmp_path, monkeypatch) -> None: + """Chief-only inactive-task and validation reads cannot move rank cursors.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "data" + _write_directory(path) + trainer = get_trainer(_config(path, mode="directory")) + active = trainer.training_tasks["first"] + result = trainer.train_step(active, 0) + before = collect_training_data_progress(trainer.training_data_by_task) + trainer.collect_display_results( + trainer.training_tasks, active_task=active, step=0, step_result=result + ) + assert collect_training_data_progress(trainer.training_data_by_task) == before + trainer._close_data_systems() + + +def test_ddp_warmup_preserves_training_progress(tmp_path, monkeypatch) -> None: + """Exercise the warmup body on CPU without invoking a compiler or NCCL.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "data" + _write_directory(path) + trainer = get_trainer(_config(path, mode="directory")) + before = collect_training_data_progress(trainer.training_data_by_task) + + class DDPWrapper: + def __init__(self, module): + self.module = module + + trainer.wrapper = DDPWrapper(trainer.wrapper) + trainer.enable_compile = True + with ( + patch.object(torch.nn.parallel, "DistributedDataParallel", DDPWrapper), + patch.object(dist, "is_initialized", return_value=True), + patch.object(dist, "get_world_size", return_value=2), + patch.object(dist.distributed_c10d, "_get_default_store") as store, + ): + store.return_value.add.return_value = 2 + trainer._precompile_outside_collectives() + assert collect_training_data_progress(trainer.training_data_by_task) == before + trainer._close_data_systems() + + +def test_full_validation_checkpoint_uses_training_rng() -> None: + """A best-checkpoint callback runs before validation returns to the loop.""" + trainer = Trainer.__new__(Trainer) + dp_random.seed(42) + before = dp_random.get_state() + saved = [] + + def save(*args, **kwargs): + saved.append(dp_random.get_state()) + + def validate(*, save_checkpoint, **kwargs): + dp_random.random(100) + save_checkpoint(Path("best.pt"), lr=1e-3, step=1) + dp_random.random(100) + + trainer._collect_data_progress = Mock(return_value={"saved_rng": before}) + trainer.full_validator = Mock(run=validate, save_best=True) + trainer.ema_full_validator = None + trainer._save_full_validation_checkpoint = save + trainer.run_full_validation(step=0, display_step=1, learning_rate=1e-3) + assert saved == [before] + assert dp_random.get_state() == before + + +def _set_mock_rank(trainer: Trainer, rank: int) -> None: + trainer.is_distributed = True + trainer.rank = rank + trainer.world_size = 2 + trainer.rank_context = RankContext(rank=rank, world_size=2) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_every_rank_joins_periodic_final_and_bias_saves( + tmp_path, monkeypatch, rank +) -> None: + """Run the native loop; only checkpoint boundaries gather rank cursors.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "data" + _write_directory(path) + config = _config(path, mode="directory", multitask=False) + config["training"].update( + numb_steps=3, + save_freq=2, + disp_training=False, + change_bias_after_training=True, + enable_ema=True, + ) + trainer = get_trainer(config) + _set_mock_rank(trainer, rank) + trainer._change_bias_after_training = Mock() + gathered = [] + collectives = [] + + def gather(ranks, local): + if isinstance(local, bool): + collectives.append("preflight") + else: + collectives.append("progress") + gathered.append(deepcopy(local)) + ranks[:] = [deepcopy(local), deepcopy(local)] + + monkeypatch.setattr(dist, "all_gather_object", gather) + monkeypatch.setattr(dist, "get_world_size", lambda: 2) + steps = [] + save = trainer.save_checkpoint + + def record_save(step): + steps.append(step) + save(step) + + trainer.save_checkpoint = record_save + trainer.run() + assert steps == [2, 3, 3] + assert len(gathered) == 3 # No per-batch or EMA-snapshot gather. + assert collectives == ["preflight", "progress"] * 3 + trainer._change_bias_after_training.assert_called_once() + checkpoint = tmp_path / "model.ckpt-3.pt" + assert checkpoint.exists() == (rank == 0) + if rank == 0: + progress = torch.load(checkpoint, weights_only=True)[ + DATA_PROGRESS_CHECKPOINT_KEY + ] + assert progress == { + "version": 1, + "world_size": 2, + "configured_seed": config["training"]["seed"], + "task_rng": trainer._task_rng.bit_generator.state, + "ranks": [gathered[-1]] * 2, + } + + +@pytest.mark.parametrize("rank", [0, 1]) +@pytest.mark.parametrize("due,save_best", [(True, True), (True, False), (False, True)]) +def test_full_validation_precollects_before_chief_only_save( + tmp_path, monkeypatch, rank, due, save_best +) -> None: + """The chief callback writes a complete envelope without a nested gather.""" + monkeypatch.chdir(tmp_path) + path = tmp_path / "data" + _write_directory(path) + config = _config(path, mode="directory", multitask=False) + config["training"]["enable_ema"] = True + trainer = get_trainer(config) + _set_mock_rank(trainer, rank) + before = collect_training_data_progress(trainer.training_data_by_task) + gathered = [] + calls = [] + + def gather(ranks, local): + assert not calls + gathered.append(deepcopy(local)) + ranks[:] = [deepcopy(local), deepcopy(local)] + + monkeypatch.setattr(dist, "all_gather_object", gather) + + def validate(*, save_checkpoint, **kwargs): + calls.append("live") + assert len(gathered) == int(save_best) + dp_random.random(30) + if rank == 0 and save_best: + save_checkpoint(tmp_path / "best.pt", step=1) + + def validate_ema(*, save_checkpoint, **kwargs): + calls.append("ema") + if rank == 0 and save_best: + save_checkpoint(tmp_path / "best_ema.pt", step=1) + + trainer.full_validator = Mock( + should_run=Mock(return_value=due), save_best=save_best, run=validate + ) + trainer.ema_full_validator = Mock( + should_run=Mock(return_value=due), save_best=save_best, run=validate_ema + ) + trainer.run_full_validation(step=0, display_step=1, learning_rate=1e-3) + assert len(gathered) == int(due and save_best) + assert calls == (["live", "ema"] if due else []) + assert dp_random.get_state() == before["rng"] + checkpoint = tmp_path / "best.pt" + assert checkpoint.exists() == (rank == 0 and due and save_best) + if checkpoint.exists(): + state = torch.load(checkpoint, weights_only=True) + assert state[DATA_PROGRESS_CHECKPOINT_KEY] == { + "version": 1, + "world_size": 2, + "configured_seed": config["training"]["seed"], + "task_rng": trainer._task_rng.bit_generator.state, + "ranks": [before, before], + } + ema = torch.load(tmp_path / "best_ema.pt", weights_only=True) + assert DATA_PROGRESS_CHECKPOINT_KEY not in ema + trainer._close_data_systems() + + +def test_filtered_rank_progress_restores_its_own_cursor(tmp_path, monkeypatch) -> None: + """Local filter retries make rank zero's cursor wrong for another rank.""" + path = _create_lmdb(str(tmp_path / "filtered.lmdb"), nframes=12, natoms=2) + + def make_rank(rank): + data = LmdbDataSystem( + path, + ["O", "H"], + batch_size=1, + seed=7, + num_workers=0, + rank=rank, + world_size=2, + ) + data.add_data_requirements( + [ + DataRequirementItem( + "pair_margin", 1, default=2.0, source_policy="derived" + ) + ] + ) + trainer = Trainer.__new__(Trainer) + _set_mock_rank(trainer, rank) + trainer._configured_data_seed = 7 + trainer._task_rng = np.random.default_rng(7) + trainer.training_data_by_task = {"Default": data} + trainer.pair_filter_enabled_by_task = {"Default": True} + return trainer, data + + originals = [make_rank(rank) for rank in range(2)] + for trainer, data in originals: + trainer._next_training_batch(data, "Default") + progress = [ + collect_training_data_progress(trainer.training_data_by_task) + for trainer, _ in originals + ] + assert progress[0]["tasks"] != progress[1]["tasks"] + + def gather(ranks, local): + ranks[:] = deepcopy(progress) + + monkeypatch.setattr(dist, "all_gather_object", gather) + checkpoint = tmp_path / "progress.pt" + torch.save(originals[0][0]._collect_data_progress(), checkpoint) + envelope = torch.load(checkpoint, weights_only=True) + for rank, (trainer, data) in enumerate(originals): + expected = [ + trainer._next_training_batch(data, "Default")["coord"] for _ in range(4) + ] + restored, restored_data = make_rank(rank) + restored._restore_data_progress(envelope) + actual = [ + restored._next_training_batch(restored_data, "Default")["coord"] + for _ in range(4) + ] + for want, got in zip(expected, actual, strict=True): + np.testing.assert_array_equal(want, got) + data.close() + restored_data.close() + + +@pytest.mark.parametrize( + "saved_size,current_size,rank", [(1, 2, 0), (1, 2, 1), (2, 1, 0), (2, 3, 2)] +) +def test_changed_world_size_leaves_data_and_rng_untouched( + caplog, saved_size, current_size, rank +) -> None: + """Topology fallback neither remaps a cursor nor restores another rank's RNG.""" + trainer = Trainer.__new__(Trainer) + trainer.rank = rank + trainer.world_size = current_size + trainer._task_rng = np.random.default_rng(7) + task_rng_state = deepcopy(trainer._task_rng.bit_generator.state) + source = Mock() + source.state_dict.return_value = {"cursor": 11} + trainer.training_data_by_task = {"Default": source} + local = collect_training_data_progress(trainer.training_data_by_task) + local["tasks"]["Default"] = {"cursor": 99} + progress = {"version": 1, "world_size": saved_size, "ranks": [local] * saved_size} + dp_random.random(5) + before = collect_training_data_progress(trainer.training_data_by_task) + with patch( + "deepmd.pt_expt.train.training.restore_training_data_progress" + ) as restore: + trainer._restore_data_progress(progress) + restore.assert_not_called() + source.load_state_dict.assert_not_called() + assert collect_training_data_progress(trainer.training_data_by_task) == before + assert trainer._task_rng.bit_generator.state == task_rng_state + assert ("Checkpoint world size changed" in caplog.text) == (rank == 0) + + +@pytest.mark.parametrize("rank", [0, 1]) +def test_explicit_reset_leaves_every_task_and_rng_untouched(caplog, rank) -> None: + """An explicit reset does not even invoke a task or shared RNG restore.""" + trainer = Trainer.__new__(Trainer) + _set_mock_rank(trainer, rank) + trainer._task_rng = np.random.default_rng(7) + task_rng_state = deepcopy(trainer._task_rng.bit_generator.state) + sources = {"first": Mock(), "second": Mock()} + for index, source in enumerate(sources.values()): + source.state_dict.return_value = {"cursor": index + 10} + trainer.training_data_by_task = sources + dp_random.seed(42) + local = collect_training_data_progress(sources) + for task in local["tasks"]: + local["tasks"][task] = {"cursor": 99} + progress = {"version": 1, "world_size": 2, "ranks": [local, local]} + dp_random.random(5) + before = collect_training_data_progress(sources) + with ( + patch( + "deepmd.pt_expt.train.training.restore_training_data_progress" + ) as restore, + patch.object(dp_random, "set_state") as set_rng, + ): + trainer._restore_data_progress(progress, reset=True) + restore.assert_not_called() + set_rng.assert_not_called() + for source in sources.values(): + source.load_state_dict.assert_not_called() + assert collect_training_data_progress(sources) == before + assert trainer._task_rng.bit_generator.state == task_rng_state + assert ("training.reset_data_progress is true" in caplog.text) == (rank == 0) + assert "Checkpoint world size changed" not in caplog.text + + +@pytest.mark.parametrize("reset", [False, True]) +@pytest.mark.parametrize("world_size", [1, 2]) +@pytest.mark.parametrize( + "changed", + [ + {"version": 2}, + {"world_size": 3}, + {"world_size": 0}, + {"world_size": -1}, + {"world_size": True}, + {"world_size": 2.0}, + {"world_size": "2"}, + {"ranks": []}, + {"ranks": {}}, + {"ranks": [None, None]}, + ], +) +def test_malformed_rank_envelope_is_rejected( + changed, world_size, caplog, reset +) -> None: + trainer = Trainer.__new__(Trainer) + trainer.rank = 0 + trainer.world_size = world_size + trainer.training_data_by_task = {} + local = collect_training_data_progress({}) + progress = {"version": 1, "world_size": 2, "ranks": [local, local]} + with pytest.raises(ValueError, match=r"[Dd]ata-progress"): + trainer._restore_data_progress(progress | changed, reset=reset) + assert "Checkpoint world size changed" not in caplog.text + assert "training.reset_data_progress is true" not in caplog.text + trainer._restore_data_progress(None, reset=reset) + assert "training.reset_data_progress is true" not in caplog.text + + +def _ddp_restart_worker(rank: int, directory: str, rendezvous: str, mode: str) -> None: + torch.set_num_threads(1) + root = Path(directory) + dist.init_process_group("gloo", init_method=rendezvous, rank=rank, world_size=2) + try: + path = root / ("data" if mode == "directory" else "data.lmdb") + config = _config(path, mode=mode) + config["training"].update( + save_ckpt=str(root / "model.ckpt"), + disp_file=str(root / "lcurve.out"), + ) + if mode == "directory": + unseeded_config = deepcopy(config) + unseeded_config["training"]["seed"] = None + unseeded = get_trainer(unseeded_config) + choices = [ + unseeded.select_task(unseeded.training_tasks).key for _ in range(8) + ] + ranks = [None, None] + dist.all_gather_object(ranks, choices) + assert ranks[0] == ranks[1] + unseeded._close_data_systems() + dp_random.seed(19) + trainer = get_trainer(deepcopy(config)) + with preserve_training_data_progress(trainer.training_data_by_task): + first = next(iter(trainer.training_data_by_task.values())) + frames = [np.asarray(first.get_batch()["coord"]).tolist() for _ in range(8)] + ranks = [None, None] + dist.all_gather_object(ranks, frames) + assert ranks[0] != ranks[1] + expected = _record_training(trainer, data_rng_draws=rank + 1) + dist.all_gather_object(ranks, [record[0] for record in expected]) + assert ranks[0] == ranks[1] + dist.barrier() + dp_random.seed(9876) + restarted = get_trainer( + deepcopy(config), restart_model=str(root / "model.ckpt-10.pt") + ) + actual = _record_training(restarted, data_rng_draws=rank + 1) + _assert_records_equal(expected[10:], actual) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not dist.is_gloo_available(), reason="Requires Gloo") +@pytest.mark.parametrize("mode", ["directory", "lmdb"]) +def test_two_rank_restart_preserves_each_rank(tmp_path, monkeypatch, mode) -> None: + """Chief-only statistics/display do not desynchronize saved rank cursors.""" + monkeypatch.setenv("DEVICE", "cpu") + monkeypatch.setenv("DP_LMDB_NUM_WORKERS", "0") + if mode == "directory": + _write_directory(tmp_path / "data") + else: + _write_lmdb(tmp_path / "data.lmdb", mixed=False) + mp.spawn( + _ddp_restart_worker, + args=(str(tmp_path), (tmp_path / "rendezvous").as_uri(), mode), + nprocs=2, + join=True, + ) diff --git a/source/tests/pt_expt/test_multitask.py b/source/tests/pt_expt/test_multitask.py index 478f0cf6aa..4c870083b6 100644 --- a/source/tests/pt_expt/test_multitask.py +++ b/source/tests/pt_expt/test_multitask.py @@ -1890,20 +1890,6 @@ def _make_gradient_test_st_config( } -def _deterministic_task_choice(task_sequence: list[int]): - """Return a patched dp_random.choice that forces task selection order.""" - original = dp_random.choice - it = iter(task_sequence) - - def patched(a, size=None, replace=True, p=None): - # Task selection: array with >=2 elements and probability vector - if hasattr(a, "__len__") and len(a) >= 2 and p is not None: - return next(it) - return original(a, size=size, replace=replace, p=p) - - return patched - - def _make_recording_step( trainer, modules_to_record: dict, @@ -2003,9 +1989,12 @@ def test_gradient_accumulation(self) -> None: {"desc": mt_desc, "fit_1": mt_fit_1, "fit_2": mt_fit_2}, mt_grads, ) - with patch( - "deepmd.utils.random.choice", - _deterministic_task_choice([0, 1]), + with patch.object( + mt_trainer, + "select_task", + side_effect=[ + mt_trainer.training_tasks[key] for key in ("model_1", "model_2") + ], ): mt_trainer.run() # 2 steps: task_0 then task_1 @@ -2150,8 +2139,8 @@ class TestCompileCaseEmbdVaryingNframes(unittest.TestCase): The test uses two systems with different atom counts and per-system ``batch_size=[2, 3]`` so every branch's compiled graph sees both nframes values. ``dim_case_embd=2`` is deliberately chosen to also - collide numerically with the nframes=2 runtime case. ``dp_random.choice`` - is mocked so both tasks and both systems are sampled. + collide numerically with the nframes=2 runtime case. Tasks are selected + explicitly, and ``dp_random.choice`` controls data-system selection. """ @classmethod