From 16654d694cbcbb422859e096a0cf93c947728803 Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Wed, 21 Oct 2026 13:07:38 +0800 Subject: [PATCH 1/8] feat(pt_expt): resume checkpointed training data progress Persist compact per-task, per-rank data cursors and selection RNG at optimizer checkpoint boundaries. Restore lazily, isolate diagnostic reads, and preserve rank-local progress when filtering skips batches. Reimplements the reader-owned progress protocol proposed in #6069, without serializing frame permutations or decoded data. Refs deepmodeling/deepmd-kit#5814 Supersedes deepmodeling/deepmd-kit#6069 Co-authored-by: OutisLi_Bot <44916828+OutisLi-Bot@users.noreply.github.com> --- deepmd/dpmodel/train/__init__.py | 10 + deepmd/dpmodel/train/data_progress.py | 82 +++ deepmd/dpmodel/utils/lmdb_data.py | 164 ++++- deepmd/pt_expt/train/training.py | 154 ++++- deepmd/pt_expt/utils/lmdb_dataset.py | 8 + deepmd/utils/data.py | 142 +++- deepmd/utils/data_system.py | 39 ++ deepmd/utils/random.py | 16 +- doc/train/training-advanced.md | 21 + .../common/dpmodel/test_data_progress.py | 74 +++ .../common/dpmodel/test_lmdb_data_progress.py | 336 ++++++++++ .../common/test_directory_data_progress.py | 260 ++++++++ .../pt_expt/test_data_progress_restart.py | 613 ++++++++++++++++++ 13 files changed, 1869 insertions(+), 50 deletions(-) create mode 100644 deepmd/dpmodel/train/data_progress.py create mode 100644 source/tests/common/dpmodel/test_data_progress.py create mode 100644 source/tests/common/dpmodel/test_lmdb_data_progress.py create mode 100644 source/tests/common/test_directory_data_progress.py create mode 100644 source/tests/pt_expt/test_data_progress_restart.py 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..5a9444ee37 --- /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 selection 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 selection 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 0d7ea42a89..3cea3b7284 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, ) @@ -1726,6 +1731,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._last_pass_length: int | None = None self._iterator = self._iter_epoch() @@ -1753,6 +1764,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._next_indices() self._schedule(indices) return batch @@ -1844,6 +1859,9 @@ def _next_indices(self) -> tuple[list[int], int]: def _iter_epoch(self) -> Iterator[tuple[list[int], int]]: """Tag batches with the length of the sampler pass that owns them.""" + 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) @@ -1851,9 +1869,123 @@ def _iter_epoch(self) -> Iterator[tuple[list[int], int]]: # 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) - for indices in 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 + ) + batches = islice(pass_batches, self._resume_batch_index, None) + self._resume_batch_index = 0 + for indices in batches: yield indices, pass_length + 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) + # A diagnostic read can restore this cursor many times during a run. + # Discard its queued result, but keep the shared pool alive. Already + # running read-only tasks may finish; their results are never consumed. + if self._pending is not None: + for future in self._pending.futures: + future.cancel() + self._pending = None + self._deferred_indices = None + 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._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. @@ -3642,11 +3774,15 @@ 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._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. @@ -3657,17 +3793,28 @@ 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 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. @@ -3743,11 +3890,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 + 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. diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index f1e7f27c58..d22b21d999 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -61,6 +61,12 @@ change_model_out_bias_by_task, resolve_step_schedule, ) +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, @@ -146,6 +152,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, min_pair_dist_requirement, @@ -2009,7 +2016,10 @@ def _make_sample( _n: int = _nbatch, _min_pair_dist: float = self.min_pair_dist_by_task[model_key], ) -> list[dict[str, np.ndarray]]: - return make_stat_input(_d, _n, _min_pair_dist) + # Statistics run only on the chief. They must not advance its + # training cursor or the shared task/system selection RNG. + with preserve_training_data_progress({"statistics": _d}): + return make_stat_input(_d, _n, _min_pair_dist) self._sample_funcs[model_key] = _make_sample @@ -2130,6 +2140,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 +2164,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 +2480,10 @@ 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) + def _create_full_validators( self, *, @@ -3056,7 +3076,38 @@ def _broadcast_value_from_rank0(self, value: Any) -> Any: @property def checkpoint_is_collective(self) -> bool: """Whether assembling a checkpoint needs every rank.""" - return self.sharding.enabled + # 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, "ranks": ranks} + + def _restore_data_progress(self, progress: Mapping[str, Any] | None) -> None: + """Restore this rank's cursor, rejecting a changed distributed world.""" + 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.") + if ( + type(progress.get("world_size")) is not int + or progress["world_size"] != self.world_size + ): + raise ValueError("Data-progress checkpoint world size does not match.") + ranks = progress.get("ranks") + if not isinstance(ranks, list) or len(ranks) != self.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.") + restore_training_data_progress(self.training_data_by_task, ranks[self.rank]) def save_checkpoint(self, step: int) -> None: # Abort before writing if any gradient norm since the previous @@ -3080,10 +3131,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, @@ -3105,6 +3158,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. @@ -3119,6 +3173,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): @@ -3129,7 +3186,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, @@ -3138,6 +3195,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 @@ -3160,6 +3218,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: @@ -3167,6 +3227,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) @@ -3276,15 +3337,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( @@ -3309,7 +3371,7 @@ def run(self) -> None: super().run(self.training_tasks) if self.change_bias_after_training and self.num_steps > self.start_step: self._change_bias_after_training() - if self.rank_context.is_chief: + if self._participates_in_checkpoint(): self.save_checkpoint(self.num_steps) finally: self._close_data_systems() @@ -3352,25 +3414,43 @@ 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, - ) + # 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, + ) + finally: + dp_random.set_state(random_state) def select_task(self, tasks: TrainingTaskCollection) -> TrainingTask: """Select a task using DeePMD's seeded random helper.""" if not tasks.is_multitask: return tasks[tasks.keys[0]] - from deepmd.utils import random as dp_random - model_index = dp_random.choice( np.arange(len(tasks), dtype=np.int_), p=tasks.probabilities, @@ -3393,12 +3473,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 9483b3ef66..008692c5db 100644 --- a/deepmd/pt_expt/utils/lmdb_dataset.py +++ b/deepmd/pt_expt/utils/lmdb_dataset.py @@ -150,6 +150,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/data.py b/deepmd/utils/data.py index 4d73e58d87..fb380535fc 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 @@ -330,18 +342,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: @@ -353,6 +373,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, @@ -789,12 +887,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: + 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_system.py b/deepmd/utils/data_system.py index 890224e063..e7df9031c0 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -479,6 +479,45 @@ 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") + if state.get("sys_probs") != np.asarray(self.sys_probs).tolist(): + 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) + 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 31609c09b8..408f0757b5 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -145,6 +145,27 @@ 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, 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. + +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_progress.py b/source/tests/common/dpmodel/test_lmdb_data_progress.py new file mode 100644 index 0000000000..c64c5787f8 --- /dev/null +++ b/source/tests/common/dpmodel/test_lmdb_data_progress.py @@ -0,0 +1,336 @@ +# 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() + 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 + + +@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_directory_data_progress.py b/source/tests/common/test_directory_data_progress.py new file mode 100644 index 0000000000..5973acaaf4 --- /dev/null +++ b/source/tests/common/test_directory_data_progress.py @@ -0,0 +1,260 @@ +# 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) + + +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_data_progress_restart.py b/source/tests/pt_expt/test_data_progress_restart.py new file mode 100644 index 0000000000..bd21048493 --- /dev/null +++ b/source/tests/pt_expt/test_data_progress_restart.py @@ -0,0 +1,613 @@ +# 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 unittest.mock import ( + Mock, + patch, +) + +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, +) +from deepmd.pt_expt.entrypoints.main import ( + get_trainer, +) +from deepmd.pt_expt.train.training import ( + Trainer, +) +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 _config(path: 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) -> 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: + 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", "lmdb", "ragged"]) +def test_restart_matches_uninterrupted_training(tmp_path, monkeypatch, mode) -> 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 == "directory" else "data.lmdb") + if mode == "directory": + _write_directory(path) + else: + _write_lmdb(path, mixed=mode == "ragged") + config = _config(path, mode=mode) + 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)) + + +def test_legacy_init_finetune_and_deployment(tmp_path, monkeypatch) -> 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) + trainer = get_trainer(deepcopy(config)) + 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() + + +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 = [] + + def gather(ranks, local): + gathered.append(deepcopy(local)) + ranks[:] = [deepcopy(local), deepcopy(local)] + + monkeypatch.setattr(dist, "all_gather_object", gather) + 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. + 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, "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, + "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.training_data_by_task = {"Default": data} + trainer.min_pair_dist_by_task = {"Default": 2.0} + 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( + "changed", + [ + {"version": 2}, + {"world_size": 3}, + {"ranks": []}, + {"ranks": {}}, + {"ranks": [None, None]}, + ], +) +def test_malformed_rank_envelope_is_rejected(changed) -> None: + trainer = Trainer.__new__(Trainer) + _set_mock_rank(trainer, 1) + 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) + trainer._restore_data_progress(None) + + +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"), + ) + dp_random.seed(19) + trainer = get_trainer(deepcopy(config)) + expected = _record_training(trainer) + dist.barrier() + dp_random.seed(9876) + restarted = get_trainer( + deepcopy(config), restart_model=str(root / "model.ckpt-10.pt") + ) + actual = _record_training(restarted) + _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, + ) From ea8f9905e2dfe56ce67b317041708651c1541a3d Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:16:47 +0800 Subject: [PATCH 2/8] fix(data): preserve saved probabilities across restart roundoff Accept tightly bounded recomputation noise, then restore the exact saved selection probabilities rather than retaining changed RNG thresholds. Keep zero-probability exclusions, reject malformed values before mutation, and preserve an exactly matching probability dtype for NumPy validation. Add 18 regression cases including a one-ULP threshold counterexample. Validation: 106 tests and 3 subtests pass, including 22 native trainer restart cases; Ruff, isort, and whitespace checks pass. Addresses the probability compatibility finding in PR #6093. --- deepmd/utils/data_system.py | 20 ++- .../common/test_directory_data_progress.py | 124 ++++++++++++++++++ 2 files changed, 143 insertions(+), 1 deletion(-) diff --git a/deepmd/utils/data_system.py b/deepmd/utils/data_system.py index e7df9031c0..efb1c391ca 100644 --- a/deepmd/utils/data_system.py +++ b/deepmd/utils/data_system.py @@ -500,7 +500,18 @@ def load_state_dict(self, state: dict[str, Any]) -> None: 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") - if state.get("sys_probs") != np.asarray(self.sys_probs).tolist(): + 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") @@ -515,6 +526,13 @@ def load_state_dict(self, state: dict[str, Any]) -> None: 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) diff --git a/source/tests/common/test_directory_data_progress.py b/source/tests/common/test_directory_data_progress.py index 5973acaaf4..a57ae77fca 100644 --- a/source/tests/common/test_directory_data_progress.py +++ b/source/tests/common/test_directory_data_progress.py @@ -207,6 +207,130 @@ def test_system_resume_random_choices(tmp_path, hdf5, batch_size): _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) From 3d7e08d380ae1bba506c42718ab69d8bdfefb7ef Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:27:30 +0800 Subject: [PATCH 3/8] fix(pt_expt): preserve restart across changed world sizes Keep existing model, optimizer and step restart semantics when the number of ranks changes. Validate the rank envelope before warning and retaining newly initialized training data and selection RNG; never pretend rank-local cursors can be remapped exactly. Retain the existing one-to-two-rank regression unchanged. Add native checkpoint comparisons with legacy behavior, optimizer-state checks and malformed-envelope tests. Document the data-order limitation. --- deepmd/pt_expt/train/training.py | 26 ++++-- doc/train/training-advanced.md | 7 ++ .../pt_expt/test_data_progress_restart.py | 90 ++++++++++++++++++- 3 files changed, 114 insertions(+), 9 deletions(-) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index d22b21d999..106445b88e 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -3090,23 +3090,35 @@ def _collect_data_progress(self) -> dict[str, Any]: return {"version": 1, "world_size": self.world_size, "ranks": ranks} def _restore_data_progress(self, progress: Mapping[str, Any] | None) -> None: - """Restore this rank's cursor, rejecting a changed distributed world.""" + """Restore rank-local progress, or retain fresh data after a topology change.""" 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.") - if ( - type(progress.get("world_size")) is not int - or progress["world_size"] != self.world_size - ): - raise ValueError("Data-progress checkpoint world size does not match.") + 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) != self.world_size: + 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 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 restore_training_data_progress(self.training_data_by_task, ranks[self.rank]) def save_checkpoint(self, step: int) -> None: diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 408f0757b5..6396c446b6 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -156,6 +156,13 @@ 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 and incompatible data settings at the same world +size remain errors. + 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` diff --git a/source/tests/pt_expt/test_data_progress_restart.py b/source/tests/pt_expt/test_data_progress_restart.py index bd21048493..fa5c9e84d7 100644 --- a/source/tests/pt_expt/test_data_progress_restart.py +++ b/source/tests/pt_expt/test_data_progress_restart.py @@ -310,6 +310,56 @@ def test_legacy_init_finetune_and_deployment(tmp_path, monkeypatch) -> None: fresh._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() + state = torch.load(tmp_path / "model.ckpt-3.pt", weights_only=True) + 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"], state["optimizer"]["state"] + ) + 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) @@ -549,24 +599,60 @@ def gather(ranks, local): 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 + 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 ("Checkpoint world size changed" in caplog.text) == (rank == 0) + + +@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) -> None: +def test_malformed_rank_envelope_is_rejected(changed, world_size, caplog) -> None: trainer = Trainer.__new__(Trainer) - _set_mock_rank(trainer, 1) + 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) + assert "Checkpoint world size changed" not in caplog.text trainer._restore_data_progress(None) From 20a91079d8245b8c59b3aea45a6ded95fcb2735d Mon Sep 17 00:00:00 2001 From: OutisLi <137472077+OutisLi@users.noreply.github.com> Date: Thu, 22 Oct 2026 10:53:58 +0800 Subject: [PATCH 4/8] fix(pt_expt): allow explicit restart data progress reset Keep strict data-source identities by default. Add the opt-in training.reset_data_progress flag for relocated data or intentional data configuration changes, without weakening path guards into basename checks. Validate the rank envelope before retaining all newly initialized task cursors and selection RNG together; preserve model, optimizer, LR, step and EMA restart, and warn that exact data continuation is unavailable. Cover copied directory, HDF5 and LMDB multitask checkpoints against legacy restart, including optimizer moments, EMA, inputs and final weights. Retain strict mismatch and malformed-envelope rejection. Document scope. Validation: 220 CPU tests and 3 subtests passed; 2 real Gloo tests were excluded locally due to socket restrictions. Ruff, format, isort and whitespace checks passed. No GPU qualification claimed. Addresses relocation feedback in deepmodeling/deepmd-kit#6093. --- deepmd/pt_expt/train/training.py | 24 ++- deepmd/utils/argcheck.py | 17 ++ doc/train/training-advanced.md | 27 ++- source/tests/common/test_argcheck_training.py | 29 ++++ .../pt_expt/test_data_progress_restart.py | 155 +++++++++++++++++- 5 files changed, 239 insertions(+), 13 deletions(-) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 106445b88e..9d6a60a845 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -2482,7 +2482,9 @@ def update_finetune_bias( # 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) + self._restore_data_progress( + data_progress, reset=training_params.get("reset_data_progress", False) + ) def _create_full_validators( self, @@ -3089,8 +3091,15 @@ def _collect_data_progress(self) -> dict[str, Any]: dist.all_gather_object(ranks, local_progress) return {"version": 1, "world_size": self.world_size, "ranks": ranks} - def _restore_data_progress(self, progress: Mapping[str, Any] | None) -> None: - """Restore rank-local progress, or retain fresh data after a topology change.""" + 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): @@ -3105,6 +3114,15 @@ def _restore_data_progress(self, progress: Mapping[str, Any] | None) -> None: 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 diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 9f9c98c26c..4fd5b60fe0 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6039,6 +6039,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 " @@ -6237,6 +6247,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/doc/train/training-advanced.md b/doc/train/training-advanced.md index 6396c446b6..5de06912f0 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -147,8 +147,9 @@ 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, 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 +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. @@ -160,8 +161,26 @@ 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 and incompatible data settings at the same world -size remain errors. +Malformed progress envelopes remain errors; by default, so do incompatible data +settings at the same world size. + +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 diff --git a/source/tests/common/test_argcheck_training.py b/source/tests/common/test_argcheck_training.py index cb5d2ffafa..5fab272b62 100644 --- a/source/tests/common/test_argcheck_training.py +++ b/source/tests/common/test_argcheck_training.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Tests for normalized training arguments.""" +import pytest +from dargs.dargs import ( + ArgumentTypeError, +) + from deepmd.utils.argcheck import ( training_args, ) @@ -15,3 +20,27 @@ def test_ema_checkpoint_retention_is_left_for_runtime_inheritance() -> None: training_argument.check_value(normalized, strict=True) 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) diff --git a/source/tests/pt_expt/test_data_progress_restart.py b/source/tests/pt_expt/test_data_progress_restart.py index fa5c9e84d7..c46168a70e 100644 --- a/source/tests/pt_expt/test_data_progress_restart.py +++ b/source/tests/pt_expt/test_data_progress_restart.py @@ -7,11 +7,16 @@ from pathlib import ( Path, ) +from shutil import ( + copyfile, + copytree, +) from unittest.mock import ( Mock, patch, ) +import h5py import lmdb import msgpack import numpy as np @@ -30,6 +35,9 @@ from deepmd.pt_expt.entrypoints.main import ( get_trainer, ) +from deepmd.pt_expt.train.ema import ( + EMA_CHECKPOINT_KEY, +) from deepmd.pt_expt.train.training import ( Trainer, ) @@ -110,7 +118,7 @@ def encode(array: np.ndarray) -> dict: ) -def _config(path: Path, *, mode: str, multitask: bool = True) -> dict: +def _config(path: str | Path, *, mode: str, multitask: bool = True) -> dict: descriptor = { "type": "se_e2_a", "sel": [4, 4], @@ -280,14 +288,20 @@ def test_best_validation_checkpoint_resumes_training(tmp_path, monkeypatch) -> N _assert_records_equal(expected[3:], _record_training(restarted)) -def test_legacy_init_finetune_and_deployment(tmp_path, monkeypatch) -> None: +@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) + 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) @@ -308,6 +322,98 @@ def test_legacy_init_finetune_and_deployment(tmp_path, monkeypatch) -> None: 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) + else: + path = tmp_path / "data" + _write_directory(path) + if mode == "hdf5": + directory = path + path = tmp_path / "data.h5" + 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), + ) + 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) + 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"], state["optimizer"]["state"] + ) + 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", "lmdb"]) @@ -627,6 +733,39 @@ def test_changed_world_size_leaves_data_and_rng_untouched( 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) + 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 ("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", @@ -643,7 +782,9 @@ def test_changed_world_size_leaves_data_and_rng_untouched( {"ranks": [None, None]}, ], ) -def test_malformed_rank_envelope_is_rejected(changed, world_size, caplog) -> 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 @@ -651,9 +792,11 @@ def test_malformed_rank_envelope_is_rejected(changed, world_size, caplog) -> Non 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) + trainer._restore_data_progress(progress | changed, reset=reset) assert "Checkpoint world size changed" not in caplog.text - trainer._restore_data_progress(None) + 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: From f5758d61e11e25442a25f402a54198a53e406b46 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 23 Oct 2026 10:02:14 +0800 Subject: [PATCH 5/8] test(pt_expt): close LMDB pipelines before fixture cleanup --- source/tests/pt_expt/test_lmdb_training.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/source/tests/pt_expt/test_lmdb_training.py b/source/tests/pt_expt/test_lmdb_training.py index 539b72f2c8..443f9de24b 100644 --- a/source/tests/pt_expt/test_lmdb_training.py +++ b/source/tests/pt_expt/test_lmdb_training.py @@ -914,14 +914,12 @@ class TestRaggedTrainingBatches(unittest.TestCase): def setUp(self) -> None: self.tmpdir = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.tmpdir) self.lmdb_path = os.path.join(self.tmpdir, "mixed.lmdb") self.spin_lmdb_path = os.path.join(self.tmpdir, "mixed-spin.lmdb") _create_mixed_nloc_test_lmdb(self.lmdb_path) _create_mixed_nloc_test_lmdb(self.spin_lmdb_path, include_spin=True) - def tearDown(self) -> None: - shutil.rmtree(self.tmpdir, ignore_errors=True) - def _config_for_model( self, model: dict, @@ -1064,6 +1062,7 @@ def _run(self, descriptor: dict, *, compile: bool = False): os.chdir(self.tmpdir) try: trainer = get_trainer(config) + self.addCleanup(trainer._close_data_systems) batch = trainer.training_data.get_batch() trainer.run() return trainer, batch @@ -1081,6 +1080,7 @@ def _run_spin(self, model: dict): os.chdir(self.tmpdir) try: trainer = get_trainer(config) + self.addCleanup(trainer._close_data_systems) batch = trainer.training_data.get_batch() trainer.run() return trainer, batch @@ -1098,7 +1098,9 @@ def test_min_pair_dist_filters_ragged_batches(self) -> None: # rather than its minimum, so the distances that set a splitting # threshold are measured on the geometry itself. A unit pair # threshold makes each margin the minimum pair distance in Å. - batch = get_trainer(copy.deepcopy(config)).training_data.get_batch() + probe = get_trainer(copy.deepcopy(config)) + self.addCleanup(probe._close_data_systems) + batch = probe.training_data.get_batch() distances = compute_min_pair_margin_batch( np.asarray(batch["coord"]), np.asarray(batch["box"]), @@ -1111,6 +1113,7 @@ def test_min_pair_dist_filters_ragged_batches(self) -> None: config["training"]["training_data"]["min_pair_dist"] = threshold trainer = get_trainer(copy.deepcopy(config)) + self.addCleanup(trainer._close_data_systems) self.assertTrue(trainer.training_data._reader.ragged_batches) dropped = False for _ in range(6): @@ -1163,6 +1166,7 @@ def test_loss_contract_can_require_rectangular_batches(self) -> None: config = self._config(self._dpa1()) config["loss"].update(loss_params) trainer = get_trainer(config) + self.addCleanup(trainer._close_data_systems) self.assertFalse(trainer.training_data._reader.ragged_batches) finally: os.chdir(cwd) @@ -1175,6 +1179,7 @@ def test_non_mixing_rule_keeps_the_rectangular_layout(self) -> None: os.chdir(self.tmpdir) try: trainer = get_trainer(config) + self.addCleanup(trainer._close_data_systems) batch = trainer.training_data.get_batch() finally: os.chdir(cwd) @@ -1259,6 +1264,7 @@ def record_layout(trainer_self, model_key: str) -> int: try: with patch.object(Trainer, "_epoch_length", record_layout): trainer = get_trainer(config) + self.addCleanup(trainer._close_data_systems) finally: os.chdir(cwd) From 5237ad0a251153216484cb6d85711aecbf5d9632 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 23 Oct 2026 10:07:37 +0800 Subject: [PATCH 6/8] test(pt_expt): match checkpoint optimizer restore devices --- .../tests/pt_expt/test_data_progress_restart.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/source/tests/pt_expt/test_data_progress_restart.py b/source/tests/pt_expt/test_data_progress_restart.py index c46168a70e..4d163cc8f9 100644 --- a/source/tests/pt_expt/test_data_progress_restart.py +++ b/source/tests/pt_expt/test_data_progress_restart.py @@ -41,6 +41,9 @@ 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, ) @@ -357,6 +360,10 @@ def test_relocated_restart_can_reset_only_data_progress( 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] @@ -395,7 +402,7 @@ def test_relocated_restart_can_reset_only_data_progress( ) torch.testing.assert_close(reset.optimizer.state_dict(), expected_optimizer) torch.testing.assert_close( - reset.optimizer.state_dict()["state"], state["optimizer"]["state"] + 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(): @@ -432,7 +439,11 @@ def test_changed_world_size_retains_legacy_optimizer_restart( config["training"].update(numb_steps=3, save_freq=3) original = get_trainer(deepcopy(config)) original.run() - state = torch.load(tmp_path / "model.ckpt-3.pt", weights_only=True) + 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) @@ -455,7 +466,7 @@ def test_changed_world_size_retains_legacy_optimizer_restart( assert changed.start_step == 3 assert changed.scheduler.get_last_lr() == expected_lr torch.testing.assert_close( - changed.optimizer.state_dict()["state"], state["optimizer"]["state"] + changed.optimizer.state_dict()["state"], checkpoint_optimizer ) for key, value in state["model"].items(): if torch.is_tensor(value): From 6d1ee3f90a1b71980dfc28a1fae1f4d66048341a Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sat, 24 Oct 2026 19:09:59 +0800 Subject: [PATCH 7/8] fix(pt-expt): preserve independent restart sampling streams --- deepmd/dpmodel/train/data_progress.py | 4 +- deepmd/dpmodel/utils/lmdb_data.py | 23 ++-- deepmd/pt_expt/entrypoints/main.py | 4 + deepmd/pt_expt/train/training.py | 41 +++++- doc/train/training-advanced.md | 11 ++ source/tests/common/dpmodel/test_lmdb_data.py | 20 +++ source/tests/pt_expt/test_checkpoint_guard.py | 3 + .../pt_expt/test_data_progress_restart.py | 119 +++++++++++++++--- 8 files changed, 192 insertions(+), 33 deletions(-) diff --git a/deepmd/dpmodel/train/data_progress.py b/deepmd/dpmodel/train/data_progress.py index 5a9444ee37..0c5858fe78 100644 --- a/deepmd/dpmodel/train/data_progress.py +++ b/deepmd/dpmodel/train/data_progress.py @@ -2,7 +2,7 @@ """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 selection RNG, not decoded data or permutations. +adds task names and the shared data-source RNG, not decoded data or permutations. """ from collections.abc import ( @@ -25,7 +25,7 @@ def collect_training_data_progress( training_data_by_task: Mapping[str, Any], ) -> dict[str, Any]: - """Collect supported data-source cursors and the shared selection RNG.""" + """Collect supported data-source cursors and the shared data-source RNG.""" return { "version": DATA_PROGRESS_VERSION, "tasks": { diff --git a/deepmd/dpmodel/utils/lmdb_data.py b/deepmd/dpmodel/utils/lmdb_data.py index 2343a994c2..0aa828dd19 100644 --- a/deepmd/dpmodel/utils/lmdb_data.py +++ b/deepmd/dpmodel/utils/lmdb_data.py @@ -2004,14 +2004,7 @@ def load_state_dict(self, state: dict[str, Any]) -> None: if not isinstance(config, dict): raise ValueError("LMDB data progress is missing its sampling configuration") self._check_restore_config(config, final=False) - # A diagnostic read can restore this cursor many times during a run. - # Discard its queued result, but keep the shared pool alive. Already - # running read-only tasks may finish; their results are never consumed. - if self._pending is not None: - for future in self._pending.futures: - future.cancel() - self._pending = None - self._deferred_indices = None + self._discard_prefetch() self._sampler.set_seed(state["seed"]) self._sampler.set_epoch(state["epoch"]) self._epoch = self._next_epoch = state["epoch"] @@ -2077,11 +2070,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() @@ -2090,6 +2080,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) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 83bf08cc98..6f86edade7 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/deepmd/pt_expt/entrypoints/main.py @@ -37,6 +37,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, ) @@ -195,6 +196,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 62af937be9..39ce82725a 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1812,6 +1812,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 {} @@ -1848,6 +1849,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 ) @@ -3057,6 +3072,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 @@ -3089,7 +3105,13 @@ def _collect_data_progress(self) -> dict[str, Any]: if self.is_distributed: ranks = [None] * self.world_size dist.all_gather_object(ranks, local_progress) - return {"version": 1, "world_size": self.world_size, "ranks": ranks} + 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 @@ -3137,7 +3159,20 @@ def _restore_data_progress( 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.""" @@ -3485,10 +3520,10 @@ def save_with_training_rng( 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]] - model_index = dp_random.choice( + model_index = self._task_rng.choice( np.arange(len(tasks), dtype=np.int_), p=tasks.probabilities, ) diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 5de06912f0..09eb44e748 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -164,6 +164,17 @@ 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. + 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 diff --git a/source/tests/common/dpmodel/test_lmdb_data.py b/source/tests/common/dpmodel/test_lmdb_data.py index 30a49f515d..2c6057a4b0 100644 --- a/source/tests/common/dpmodel/test_lmdb_data.py +++ b/source/tests/common/dpmodel/test_lmdb_data.py @@ -2998,6 +2998,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/pt_expt/test_checkpoint_guard.py b/source/tests/pt_expt/test_checkpoint_guard.py index 0e1bba23aa..54525725ae 100644 --- a/source/tests/pt_expt/test_checkpoint_guard.py +++ b/source/tests/pt_expt/test_checkpoint_guard.py @@ -20,6 +20,7 @@ patch, ) +import numpy as np import pytest import torch import torch.distributed as dist @@ -67,6 +68,8 @@ def _make_trainer(backend, directory, *, rank=0, stage=0, distributed=False): 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") diff --git a/source/tests/pt_expt/test_data_progress_restart.py b/source/tests/pt_expt/test_data_progress_restart.py index 97b32a57e7..0d69d3b15d 100644 --- a/source/tests/pt_expt/test_data_progress_restart.py +++ b/source/tests/pt_expt/test_data_progress_restart.py @@ -31,6 +31,7 @@ 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, @@ -121,6 +122,19 @@ def encode(array: np.ndarray) -> dict: ) +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", @@ -183,7 +197,7 @@ def _config(path: str | Path, *, mode: str, multitask: bool = True) -> dict: return normalize(update_deepmd_input(config, warning=False), multi_task=multitask) -def _record_training(trainer: Trainer) -> list: +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 @@ -193,6 +207,8 @@ def _record_training(trainer: Trainer) -> list: 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, @@ -338,20 +354,12 @@ def test_relocated_restart_can_reset_only_data_progress( 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) - if mode == "hdf5": - directory = path - path = tmp_path / "data.h5" - 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), - ) system = str(path) + ("#/system" if mode == "hdf5" else "") config_mode = "lmdb" if mode == "lmdb" else "directory" config = _config(system, mode=config_mode) @@ -423,6 +431,52 @@ def test_relocated_restart_can_reset_only_data_progress( ) +@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 @@ -601,7 +655,13 @@ def record_save(step): progress = torch.load(checkpoint, weights_only=True)[ DATA_PROGRESS_CHECKPOINT_KEY ] - assert progress == {"version": 1, "world_size": 2, "ranks": [gathered[-1]] * 2} + 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]) @@ -657,6 +717,8 @@ def validate_ema(*, save_checkpoint, **kwargs): 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) @@ -687,6 +749,8 @@ def make_rank(rank): ) 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.min_pair_dist_by_task = {"Default": 2.0} return trainer, data @@ -733,6 +797,8 @@ def test_changed_world_size_leaves_data_and_rng_untouched( 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} @@ -748,6 +814,7 @@ def test_changed_world_size_leaves_data_and_rng_untouched( 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) @@ -756,6 +823,8 @@ def test_explicit_reset_leaves_every_task_and_rng_untouched(caplog, rank) -> Non """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} @@ -779,6 +848,7 @@ def test_explicit_reset_leaves_every_task_and_rng_untouched(caplog, rank) -> Non 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 @@ -828,15 +898,34 @@ def _ddp_restart_worker(rank: int, directory: str, rendezvous: str, mode: str) - 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)) - expected = _record_training(trainer) + 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) + actual = _record_training(restarted, data_rng_draws=rank + 1) _assert_records_equal(expected[10:], actual) finally: dist.destroy_process_group() From 7c745ddc022bcb57d023f27253adbfbc3f8d6715 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Sun, 25 Oct 2026 09:32:45 +0800 Subject: [PATCH 8/8] test(train): select multitask gradient cases explicitly --- source/tests/pt_expt/test_multitask.py | 27 ++++++++------------------ 1 file changed, 8 insertions(+), 19 deletions(-) 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