feat(train): checkpointable data-source progress for restart

deepmd/dpmodel/train/__init__.py

@@ -12,6 +12,11 @@
     make_task_maps,
     print_data_summaries,
 )
+from .data_progress import (
+    DATA_PROGRESS_CHECKPOINT_KEY,
+    collect_training_data_progress,
+    restore_training_data_progress,
+)
 from .entrypoint import (
     AbstractTrainEntrypoint,
     TrainEntrypointOptions,
@@ -44,6 +49,7 @@
 )
 
 __all__ = [
+    "DATA_PROGRESS_CHECKPOINT_KEY",
     "DEFAULT_TASK_KEY",
     "AbstractTrainEntrypoint",
     "AbstractTrainer",
@@ -64,9 +70,11 @@
     "build_checkpoint_stores",
     "change_model_out_bias",
     "change_model_out_bias_by_task",
+    "collect_training_data_progress",
     "iter_training_task_configs",
     "make_task_maps",
     "print_data_summaries",
     "resolve_keep_ckpt_count",
     "resolve_step_schedule",
+    "restore_training_data_progress",
 ]

deepmd/dpmodel/train/data_progress.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Checkpointable training-data progress for restartable optimizers.
+
+A training restart restores model, optimizer, learning-rate, and EMA state,
+but historically rebuilt every data source from the beginning of epoch zero.
+This module defines the backend-neutral payload that records each task's
+logical data cursor — the next training batch the optimizer has not yet
+consumed — and the process RNG that drives multi-task and directory-system
+selection.
+
+Prefetch and decoder read-ahead are excluded from the payload. Restoring
+progress rebuilds iterators at the saved cursor so the next ``get_batch``
+matches an uninterrupted run without decoding the dataset during seek.
+"""
+
+from __future__ import (
+    annotations,
+)
+
+from collections.abc import (
+    Mapping,
+)
+from typing import (
+    Any,
+)
+
+from deepmd.utils import random as dp_random
+
+DATA_PROGRESS_CHECKPOINT_KEY = "data_progress"
+DATA_PROGRESS_VERSION = 1
+
+__all__ = [
+    "DATA_PROGRESS_CHECKPOINT_KEY",
+    "DATA_PROGRESS_VERSION",
+    "collect_training_data_progress",
+    "restore_training_data_progress",
+]
+
+
+def collect_training_data_progress(
+    training_data_by_task: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Collect per-task data-source progress and the shared selection RNG.
+
+    Parameters
+    ----------
+    training_data_by_task : Mapping[str, Any]
+        Training data sources keyed by task name.
+
+    Returns
+    -------
+    dict[str, Any]
+        Versioned checkpoint payload. Tasks whose data source does not
+        implement ``state_dict`` are omitted; an empty ``tasks`` map still
+        carries the RNG so multi-task selection can resume.
+    """
+    tasks: dict[str, Any] = {}
+    for task_key, data in training_data_by_task.items():
+        state_dict = getattr(data, "state_dict", None)
+        if callable(state_dict):
+            tasks[task_key] = state_dict()
+    return {
+        "version": DATA_PROGRESS_VERSION,
+        "tasks": tasks,
+        "rng": dp_random.get_state(),
+    }
+
+
+def restore_training_data_progress(
+    training_data_by_task: Mapping[str, Any],
+    progress: Mapping[str, Any] | None,
+) -> None:
+    """Restore per-task data progress and the shared selection RNG.
+
+    Parameters
+    ----------
+    training_data_by_task : Mapping[str, Any]
+        Training data sources keyed by task name. Sources without
+        ``load_state_dict`` are left unchanged.
+    progress : Mapping[str, Any] or None
+        Payload previously produced by :func:`collect_training_data_progress`.
+        ``None`` or a missing payload retains the legacy restart behaviour
+        (data sources stay at their freshly constructed cursor).
+
+    Raises
+    ------
+    ValueError
+        If the payload version is unsupported, a task is missing, or a data
+        source rejects the restored configuration.
+    """
+    if progress is None:
+        return
+    version = int(progress.get("version", -1))
+    if version != DATA_PROGRESS_VERSION:
+        raise ValueError(
+            f"Unsupported data-progress checkpoint version {version}; "
+            f"expected {DATA_PROGRESS_VERSION}."
+        )
+    tasks = progress.get("tasks", {})
+    if not isinstance(tasks, Mapping):
+        raise ValueError("data-progress 'tasks' must be a mapping.")
+    for task_key, data in training_data_by_task.items():
+        load_state_dict = getattr(data, "load_state_dict", None)
+        if not callable(load_state_dict):
+            continue
+        if task_key not in tasks:
+            raise ValueError(f"data-progress checkpoint is missing task {task_key!r}.")
+        load_state_dict(tasks[task_key])
+    if "rng" in progress:
+        dp_random.set_state(progress["rng"])

deepmd/dpmodel/utils/lmdb_data.py

@@ -17,6 +17,7 @@
     Callable,
     Iterable,
     Iterator,
+    Mapping,
     Sequence,
 )
 from concurrent.futures import (
@@ -1617,12 +1618,19 @@ def __init__(
         self._reader = reader
         self._sampler = sampler
         self._epoch = 0
-        self._iterator = self._iter_epoch()
+        self._logical_epoch = 0
+        self._logical_batch_index = 0
+        self._logical_epoch_length = 0
+        # Length of a sampler epoch opened by prefetch before the logical
+        # cursor wraps. LmdbBatchSampler advances on ``__iter__``, so
+        # ``len(self._sampler)`` at wrap would observe the *following* epoch.
+        self._pending_logical_epoch_length: int | None = None
         self._num_workers = num_workers
         self._pool: _LmdbPoolEntry | None = None
         self._pending: _PendingBatch | None = None
         self._deferred_indices: list[int] | None = None
         self._closed = False
+        self._iterator = self._iter_epoch()
 
     def __iter__(self) -> "LmdbBatchIterator":
         return self
@@ -1641,8 +1649,91 @@ def __next__(self) -> dict[str, Any]:
             batch = self._decode(self._next_indices())
 
         self._schedule(self._next_indices())
+        self._advance_logical_cursor()
         return batch
 
+    def state_dict(self) -> dict[str, int]:
+        """Return the logical cursor of the next batch to consume.
+
+        Prefetched decoder work is excluded: only batches already returned to
+        the trainer advance this cursor.
+        """
+        return {
+            "epoch": self._logical_epoch,
+            "batch_index": self._logical_batch_index,
+        }
+
+    def load_state_dict(self, state: Mapping[str, Any]) -> None:
+        """Reposition to the saved logical cursor without decoding frames.
+
+        Parameters
+        ----------
+        state : Mapping[str, Any]
+            Progress previously returned by :meth:`state_dict`.
+        """
+        if self._closed:
+            raise RuntimeError(
+                "cannot restore progress on a closed LMDB batch iterator"
+            )
+        epoch = int(state["epoch"])
+        batch_index = int(state["batch_index"])
+        if epoch < 0 or batch_index < 0:
+            raise ValueError(
+                f"data-progress cursor must be non-negative, got epoch={epoch}, "
+                f"batch_index={batch_index}"
+            )
+        self._cancel_prefetch()
+        self._pending_logical_epoch_length = None
+        self._logical_epoch = epoch
+        self._logical_batch_index = batch_index
+        self._epoch = epoch
+        self._iterator = self._iter_epoch()
+        if self._logical_epoch_length == 0:
+            if batch_index != 0:
+                raise ValueError(
+                    "data-progress batch_index must be 0 for an empty epoch, "
+                    f"got {batch_index}"
+                )
+            return
+        if batch_index >= self._logical_epoch_length:
+            raise ValueError(
+                f"data-progress batch_index {batch_index} is out of range for "
+                f"epoch length {self._logical_epoch_length}"
+            )
+        for _ in range(batch_index):
+            self._next_indices()
+
+    def _advance_logical_cursor(self) -> None:
+        """Record that one batch was returned to the trainer.
+
+        Prefetch may already have opened the next sampler epoch. That open
+        stashes the next logical length in ``_pending_logical_epoch_length``;
+        the wrap applies it. Falling back to ``len(self._sampler)`` is only
+        correct when prefetch has not yet opened the successor (samplers such
+        as ``LmdbBatchSampler`` advance on ``__iter__``, so a post-prefetch
+        ``len`` would name the epoch after next).
+        """
+        if self._logical_epoch_length <= 0:
+            return
+        self._logical_batch_index += 1
+        if self._logical_batch_index < self._logical_epoch_length:
+            return
+        self._logical_epoch += 1
+        self._logical_batch_index = 0
+        if self._pending_logical_epoch_length is not None:
+            self._logical_epoch_length = self._pending_logical_epoch_length
+            self._pending_logical_epoch_length = None
+        else:
+            self._logical_epoch_length = len(self._sampler)
+
+    def _cancel_prefetch(self) -> None:
+        """Drop prefetched work so a seek can rebuild the sampler position."""
+        if self._pending is not None:
+            for future in self._pending.futures:
+                future.cancel()
+            self._pending = None
+        self._deferred_indices = None
+
     def _decode(self, indices: list[int]) -> dict[str, Any]:
         """Decode one batch, in the pool when that is worthwhile and possible."""
         futures = self._offer(indices)
@@ -1733,6 +1824,17 @@ def _iter_epoch(self) -> Iterator[list[int]]:
         set_epoch = getattr(self._sampler, "set_epoch", None)
         if callable(set_epoch):
             set_epoch(self._epoch)
+        epoch_length = len(self._sampler)
+        # Prefetch may open the next sampler epoch before the trainer consumes
+        # the last batch of the logical epoch. Bind the active logical length
+        # only when reconstructing the epoch the cursor still names; otherwise
+        # stash the length for the upcoming wrap — ``iter(self._sampler)`` may
+        # already advance the sampler, so a later ``len`` would be wrong.
+        if self._epoch == self._logical_epoch:
+            self._logical_epoch_length = epoch_length
+            self._pending_logical_epoch_length = None
+        else:
+            self._pending_logical_epoch_length = epoch_length
         return iter(self._sampler)
 
     def _submit(self, indices: list[int]) -> list[Future[dict[str, Any]]]:

deepmd/pt_expt/train/training.py

@@ -47,6 +47,7 @@
     fully_shard = None  # type: ignore[assignment]
 
 from deepmd.dpmodel.train import (
+    DATA_PROGRESS_CHECKPOINT_KEY,
     DEFAULT_TASK_KEY,
     AbstractTrainer,
     RankContext,
@@ -59,7 +60,9 @@
     build_checkpoint_stores,
     change_model_out_bias,
     change_model_out_bias_by_task,
+    collect_training_data_progress,
     resolve_step_schedule,
+    restore_training_data_progress,
 )
 from deepmd.dpmodel.utils.batch import (
     normalize_batch,
@@ -2102,6 +2105,7 @@ def initialize_statistics(
         # cannot be copied into sharded parameters.
         ema_state_dict = None
         optimizer_state_dict = None
+        data_progress_state = None
         if resuming:
             log.info(f"Resuming from {resume_model}.")
             is_pte = resume_model.endswith((".pte", ".pt2"))
@@ -2117,17 +2121,23 @@ def initialize_statistics(
                 if "model" in state_dict:
                     # Optimizer and EMA state describe the weights of the run
                     # they were saved by; a finetune starts a new run and keeps
-                    # neither.
+                    # neither. Data-source progress is likewise resume-only.
                     continues_run = self.restart_training and finetune_model is None
                     optimizer_state_dict = (
                         state_dict["optimizer"] if continues_run else None
                     )
                     ema_state_dict = (
                         state_dict.get(EMA_CHECKPOINT_KEY) if continues_run else None
                     )
+                    data_progress_state = (
+                        state_dict.get(DATA_PROGRESS_CHECKPOINT_KEY)
+                        if continues_run
+                        else None
+                    )
                     state_dict = state_dict["model"]
                 else:
                     optimizer_state_dict = None
+                    data_progress_state = None
                 self.start_step = (
                     state_dict["_extra_state"]["train_infos"]["step"]
                     if self.restart_training
@@ -2371,6 +2381,14 @@ def update_finetune_bias(
             else None
         )
 
+        # Data-source progress is restored after requirements and statistics
+        # hooks have run, and before training can draw a batch. Missing progress
+        # keeps the legacy restart behaviour of rebuilding iterators at epoch 0.
+        restore_training_data_progress(
+            self.training_data_by_task,
+            data_progress_state,
+        )
+
         self._configure_neighbor_graph_method(
             training_params.get("neighbor_graph_method", "auto")
         )
@@ -3066,6 +3084,12 @@ def _write_checkpoint(
             state["optimizer"] = optim_state
         if include_ema_state and self.model_ema is not None:
             state[EMA_CHECKPOINT_KEY] = self.model_ema.state_dict()
+        if include_optimizer:
+            # Resume-capable checkpoints record the logical data cursor so a
+            # restart yields the same next batch; deployment snapshots omit it.
+            state[DATA_PROGRESS_CHECKPOINT_KEY] = collect_training_data_progress(
+                self.training_data_by_task
+            )
         ckpt_path.parent.mkdir(parents=True, exist_ok=True)
         torch.save(state, ckpt_path)