From 58f7301179fc370df01adbedf7dc7fffab305d95 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 28 Jul 2026 13:47:55 +0800 Subject: [PATCH 1/5] feat(pt_expt): close the checkpointing, EMA and training-report gap with pt A pt_expt run could not be operated under the same conventions as a pt run: it ignored the checkpoint directory and retention options, kept no EMA, never reported its parameter count, and printed progress in its own format with a systematically inflated estimate of the remaining time. Four features are brought over, and everything about them that is not specific to a backend is described once so that both backends execute the same implementation instead of two drifting copies. Checkpointing. `save_dir` and `ckpt_keep_ratio` are honoured, and the on-disk layout -- naming, prefix symlinks, pointer file and retention -- moves into `CheckpointStore` (deepmd/dpmodel/train/checkpoint.py), which both backends now use; pt loses four hand-rolled copies of the publish sequence. The store drops checkpoints numbered above the one being written before it applies the retention window. Those are remnants of a longer earlier run over the same directory, and leaving them in place let the window discard the checkpoint that was just written, so restarting a run in a finished directory kept no result at all. A disabled window (`max_ckpt_keep < 1`) now retains every checkpoint, as the jax and tf2 backends already do, rather than deleting all of them including the current one. `resolve_keep_ckpt_count` also handles `save_freq <= 0`, which previously raised `ZeroDivisionError` when combined with a retention ratio. EMA. `enable_ema`, `ema_decay` and `ema_ckpt_keep` are honoured. The shadow weights are updated after every optimizer step, written as a separate family of checkpoints carrying neither optimizer nor EMA state, and restored on restart. `deepmd/pt/train/ema.py` is reused as is rather than copied. Full validation. `build_full_validators` configures the live-weight and the EMA-weight flow together, since they differ only in the weights they read, the log they write and the prefix of the checkpoints they select. pt_expt thereby gains `ema_full_validation`, and the per-flow eligibility check stays with the backend that knows what it supports. `compiled_infer` and `amp_infer` reach the models through the shared `infer_env_defaults` translation. `tf32_infer` remains unimplemented in pt_expt, which has no TF32 path yet. Training report. The display now prints the losses before the wall-clock line and omits the per-step average, matching pt, and the run ends with the average step time over the representative intervals. The remaining time is extrapolated from the interval that just ended rather than from the average since the run began: the latter carries the one-off cost of the first steps, such as graph compilation, and therefore never stops overestimating. This accounting moves into `TrainingTimer` (deepmd/dpmodel/train/timing.py), replacing pt's three loose counters. The parameter-count report moves to `deepmd/loggers/training.py`, the home of the other training log messages, and reads counts a backend supplies. Relocation. `deepmd/pt/train/{utils,validation,ema}.py` move under `deepmd/pt_expt/train/`. pt is being retired, so the shared training code belongs with the backend that outlives it and the dependency arrow is reversed. While moving, the validator recognizes its validation data by the surface it exposes rather than by pt's dataset types, and reads the environment constants from pt_expt; only `AutoBatchSize` and `to_torch_tensor` still come from `deepmd/pt/utils`, which is the next unit to migrate. --- deepmd/dpmodel/train/__init__.py | 14 + deepmd/dpmodel/train/checkpoint.py | 306 ++++++++++++++ deepmd/dpmodel/train/timing.py | 142 +++++++ deepmd/dpmodel/train/trainer.py | 66 ++- deepmd/loggers/training.py | 52 ++- deepmd/pt/train/training.py | 381 ++++-------------- deepmd/{pt => pt_expt}/train/ema.py | 0 deepmd/pt_expt/train/training.py | 348 ++++++++++------ deepmd/{pt => pt_expt}/train/utils.py | 111 +++-- deepmd/{pt => pt_expt}/train/validation.py | 182 +++++++-- deepmd/utils/argcheck.py | 20 +- doc/train/training-advanced.md | 4 +- .../dpmodel/test_train_abstract_trainer.py | 32 ++ .../common/dpmodel/test_train_checkpoint.py | 179 ++++++++ .../tests/common/dpmodel/test_train_timing.py | 113 ++++++ source/tests/common/test_loggers_training.py | 57 +++ source/tests/pt/test_training.py | 22 +- source/tests/pt/test_validation.py | 12 +- source/tests/pt_expt/test_entrypoint.py | 75 +--- .../tests/{pt => pt_expt}/test_train_utils.py | 20 +- source/tests/pt_expt/test_training.py | 295 +++++++++++--- 21 files changed, 1696 insertions(+), 735 deletions(-) create mode 100644 deepmd/dpmodel/train/checkpoint.py create mode 100644 deepmd/dpmodel/train/timing.py rename deepmd/{pt => pt_expt}/train/ema.py (100%) rename deepmd/{pt => pt_expt}/train/utils.py (77%) rename deepmd/{pt => pt_expt}/train/validation.py (84%) create mode 100644 source/tests/common/dpmodel/test_train_checkpoint.py create mode 100644 source/tests/common/dpmodel/test_train_timing.py create mode 100644 source/tests/common/test_loggers_training.py rename source/tests/{pt => pt_expt}/test_train_utils.py (85%) diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index e6124d8ce9..dc2db6cb4c 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Backend-independent training abstractions.""" +from .checkpoint import ( + CheckpointStore, + build_checkpoint_stores, + resolve_keep_ckpt_count, +) from .data import ( TrainingTaskConfig, iter_training_task_configs, @@ -15,6 +20,10 @@ StepSchedule, resolve_step_schedule, ) +from .timing import ( + DisplayInterval, + TrainingTimer, +) from .trainer import ( DEFAULT_TASK_KEY, AbstractTrainer, @@ -32,6 +41,8 @@ "DEFAULT_TASK_KEY", "AbstractTrainEntrypoint", "AbstractTrainer", + "CheckpointStore", + "DisplayInterval", "LearningCurveWriter", "RankContext", "StepSchedule", @@ -41,10 +52,13 @@ "TrainingTask", "TrainingTaskCollection", "TrainingTaskConfig", + "TrainingTimer", + "build_checkpoint_stores", "change_model_out_bias", "change_model_out_bias_by_task", "iter_training_task_configs", "make_task_maps", "print_data_summaries", + "resolve_keep_ckpt_count", "resolve_step_schedule", ] diff --git a/deepmd/dpmodel/train/checkpoint.py b/deepmd/dpmodel/train/checkpoint.py new file mode 100644 index 0000000000..d9048fd801 --- /dev/null +++ b/deepmd/dpmodel/train/checkpoint.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""On-disk layout and retention policy of periodic training checkpoints. + +A run writes one numbered file per checkpoint and keeps a fixed-size window of +the most recent ones. Both the naming and the pruning are pure filesystem +concerns, independent of how a backend serializes its state, so they are +described once here and shared by every backend. +""" + +from __future__ import ( + annotations, +) + +import logging +from math import ( + ceil, +) +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +from deepmd.common import ( + symlink_prefix_files, +) + +if TYPE_CHECKING: + from collections.abc import ( + Mapping, + ) + +log = logging.getLogger(__name__) + +__all__ = ["CheckpointStore", "build_checkpoint_stores", "resolve_keep_ckpt_count"] + + +class CheckpointStore: + """Naming, publication and retention of a family of checkpoints. + + Numbered checkpoints are written as ``/-``, + where ```` is the file name of ``prefix`` and ```` is + ``save_dir`` when given and the directory of ``prefix`` otherwise. + Publishing a checkpoint points the prefix-named files at it, so a consumer + that only knows the prefix always reaches the newest checkpoint. + + Parameters + ---------- + prefix : str or Path + The checkpoint prefix, such as ``model.ckpt``. Its directory receives + the prefix-named symlinks, and its file name seeds the numbered files. + save_dir : Path, optional + Directory holding the numbered checkpoints. Defaults to the directory + of ``prefix``. + max_keep : int, optional + Number of most recent numbered checkpoints to retain. Values below one + retain every checkpoint. + suffix : str, optional + File suffix of a checkpoint, including the leading dot. + pointer_file : str or Path, optional + File recording the path of the most recently published checkpoint. + ``None`` publishes symlinks only, which is what a secondary family of + checkpoints, such as the EMA one, requires so that it does not claim + the pointer of the primary family. + """ + + def __init__( + self, + prefix: str | Path, + *, + save_dir: Path | None = None, + max_keep: int = 5, + suffix: str = ".pt", + pointer_file: str | Path | None = None, + ) -> None: + self.prefix = Path(prefix) + self.directory = Path(save_dir) if save_dir is not None else self.prefix.parent + self.max_keep = int(max_keep) + self.suffix = suffix + self.pointer_file = Path(pointer_file) if pointer_file is not None else None + + def prepare(self) -> None: + """Create the directories receiving the checkpoints and the symlinks.""" + self.directory.mkdir(parents=True, exist_ok=True) + self.prefix.parent.mkdir(parents=True, exist_ok=True) + + def path_for(self, step: int) -> Path: + """Return the path of the checkpoint recorded at a step. + + Parameters + ---------- + step : int + Training step encoded into the file name. + + Returns + ------- + Path + Path of the numbered checkpoint of this store. + """ + return self.directory / f"{self.prefix.name}-{step}{self.suffix}" + + def step_of(self, path: Path) -> int | None: + """Return the step encoded in a checkpoint name, or ``None``. + + Only the file name is inspected; see :meth:`holds` for membership of + this store. + + Parameters + ---------- + path : Path + Candidate checkpoint path. + + Returns + ------- + int or None + The step of a numbered checkpoint of this store, or ``None`` when + the name does not follow ``-``. + """ + stem_prefix = f"{self.prefix.name}-" + if path.suffix != self.suffix or not path.name.startswith(stem_prefix): + return None + step_text = path.name[len(stem_prefix) : -len(self.suffix)] + if not step_text.isdigit(): + return None + return int(step_text) + + def holds(self, path: Path) -> bool: + """Whether a path is a numbered checkpoint of this store. + + Parameters + ---------- + path : Path + Candidate checkpoint path. + + Returns + ------- + bool + ``True`` when the path lies in this store's directory and its name + encodes a step. + """ + return ( + self.step_of(path) is not None + and path.parent.resolve() == self.directory.resolve() + ) + + def publish(self, path: Path) -> None: + """Point the prefix-named files and the pointer file at a checkpoint. + + Parameters + ---------- + path : Path + Checkpoint the prefix-named files resolve to from now on. + """ + self.prefix.parent.mkdir(parents=True, exist_ok=True) + symlink_prefix_files(str(path.with_suffix("")), str(self.prefix)) + if self.pointer_file is not None: + self.pointer_file.write_text(str(path)) + + def prune(self, current: Path) -> None: + """Drop the checkpoints made obsolete by a fresh one. + + Checkpoints numbered above the current step are remnants of a longer + earlier run over the same directory. They are removed first: leaving + them in place would let the retention window discard the freshly + written checkpoint instead, so a rerun in a finished directory would + keep no result at all. The window then retains the newest ``max_keep`` + checkpoints. The checkpoint just written is never removed. + + Parameters + ---------- + current : Path + Path of the checkpoint that was just written. A path this store + does not hold, such as a checkpoint selected by validation, dates + nothing and claims no slot of the window. + """ + if self.max_keep < 1: + return + current_step = self.step_of(current) if self.holds(current) else None + retained: list[tuple[int, Path]] = [] + for path in self.directory.glob(f"*{self.suffix}"): + step = self.step_of(path) + if step is None or path.is_symlink(): + continue + if current_step is not None and path.name == current.name: + continue + if current_step is not None and step > current_step: + path.unlink(missing_ok=True) + else: + retained.append((step, path)) + retained.sort(key=lambda item: (item[0], item[1].name)) + # The current checkpoint occupies one slot of the window when this + # store holds it. + occupied = 1 if current_step is not None else 0 + for _, path in retained[: len(retained) + occupied - self.max_keep]: + path.unlink(missing_ok=True) + + +def resolve_keep_ckpt_count( + ckpt_keep_ratio: float | None, num_steps: int, save_freq: int +) -> int | None: + """Convert a checkpoint-retention ratio into a sliding-window keep count. + + A checkpoint is written every ``save_freq`` steps and once more at the + final step, so a run of ``num_steps`` produces ``ceil(num_steps / + save_freq)`` of them in total (the terminal checkpoint is off-cadence when + ``num_steps`` is not a multiple of ``save_freq``). Keeping the most recent + ``ceil(ratio * total)`` is equivalent to retaining the final ``ratio`` + fraction of the run by step, without the caller computing the count by + hand. + + Parameters + ---------- + ckpt_keep_ratio : float or None + The fraction of the training run, by step, whose periodic checkpoints + are retained. ``None`` leaves the keep count unchanged. + num_steps : int + The total number of training steps, already resolved (including when + derived from ``numb_epoch``). + save_freq : int + The checkpoint saving frequency in steps. Values below one disable + periodic saving, leaving the final checkpoint as the only one. + + Returns + ------- + int or None + The number of most recent checkpoints to keep (at least one), or + ``None`` when ``ckpt_keep_ratio`` is not set. + """ + if ckpt_keep_ratio is None: + return None + total_ckpts = max(1, ceil(num_steps / save_freq)) if save_freq > 0 else 1 + return max(1, ceil(ckpt_keep_ratio * total_ckpts)) + + +def build_checkpoint_stores( + training_params: Mapping[str, Any], + *, + num_steps: int, + ema_prefix: str | Path, + rank: int = 0, +) -> tuple[CheckpointStore, CheckpointStore]: + """Build the checkpoint stores of a training run. + + A run keeps two families of checkpoints: the periodic ones, which carry + the live weights and the state needed to resume, and the EMA ones, which + carry smoothed weights only. They share a directory and differ in prefix, + retention and whether they own the pointer file. + + Parameters + ---------- + training_params : Mapping[str, Any] + The normalized ``training`` section. ``save_ckpt``, ``save_dir``, + ``save_freq``, ``max_ckpt_keep``, ``ckpt_keep_ratio`` and + ``ema_ckpt_keep`` are read from it. + num_steps : int + The resolved run length, needed to turn ``ckpt_keep_ratio`` into a + keep count. + ema_prefix : str or Path + Checkpoint prefix of the EMA family, derived by the backend from + ``save_ckpt``. + rank : int, optional + Process rank. Only the chief creates directories and reports the + resolved retention. + + Returns + ------- + tuple[CheckpointStore, CheckpointStore] + The periodic store and the EMA store. + """ + save_dir = training_params.get("save_dir") + save_freq = int(training_params.get("save_freq", 1000)) + max_keep = int(training_params.get("max_ckpt_keep", 5)) + ema_max_keep = int(training_params.get("ema_ckpt_keep", 3)) + ckpt_keep_ratio = training_params.get("ckpt_keep_ratio") + + keep_ckpt_count = resolve_keep_ckpt_count(ckpt_keep_ratio, num_steps, save_freq) + if keep_ckpt_count is not None: + max_keep = keep_ckpt_count + ema_max_keep = keep_ckpt_count + if rank == 0: + log.info( + "Resolved checkpoint retention to %d from ckpt_keep_ratio=%s " + "(num_steps=%d, save_freq=%d).", + keep_ckpt_count, + ckpt_keep_ratio, + num_steps, + save_freq, + ) + + directory = Path(save_dir) if save_dir else None + store = CheckpointStore( + training_params.get("save_ckpt", "model.ckpt"), + save_dir=directory, + max_keep=max_keep, + pointer_file="checkpoint", + ) + ema_store = CheckpointStore( + ema_prefix, + save_dir=directory, + max_keep=ema_max_keep, + ) + if rank == 0: + store.prepare() + return store, ema_store diff --git a/deepmd/dpmodel/train/timing.py b/deepmd/dpmodel/train/timing.py new file mode 100644 index 0000000000..603e4f228b --- /dev/null +++ b/deepmd/dpmodel/train/timing.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Wall-clock accounting of a training run. + +Progress is reported once per display interval and summarized once at the end +of the run. Both numbers come from the same bookkeeping, which is independent +of the backend performing the steps. +""" + +from __future__ import ( + annotations, +) + +import datetime +import time +from dataclasses import ( + dataclass, +) + +__all__ = ["DisplayInterval", "TrainingTimer"] + + +@dataclass(frozen=True) +class DisplayInterval: + """Wall-clock summary of the steps since the previous display. + + Attributes + ---------- + display_step : int + One-based index of the step the interval ends at. + wall_time : float + Time elapsed since the previous display, in s. + steps : int + Number of training steps performed during the interval. + eta : int or None + Estimated time left in the run, in s. ``None`` for the opening + interval, whose rate carries the one-off costs of starting the run and + so forecasts nothing. + timestamp : datetime.datetime + Local time at the end of the interval. + """ + + display_step: int + wall_time: float + steps: int + eta: int | None + timestamp: datetime.datetime + + +class TrainingTimer: + """Timer producing per-interval progress and a run-level average. + + The remaining time is extrapolated from the rate of the interval that just + ended rather than from the average since the run began, because the start + of a run also carries one-off costs -- data preparation, graph compilation, + autotuning -- that never recur and would otherwise inflate every estimate. + + The reported average deliberately omits the first ``disp_freq`` steps for + the same reason, unless the run is too short for the exclusion to leave a + meaningful sample. + + Parameters + ---------- + start_step : int + Step the run starts from, non-zero when restarting. + num_steps : int + Step the run ends at. + disp_freq : int + Number of steps between two displays. + """ + + def __init__(self, *, start_step: int, num_steps: int, disp_freq: int) -> None: + self._start_step = int(start_step) + self._num_steps = int(num_steps) + self._disp_freq = max(1, int(disp_freq)) + self._interval_start = time.time() + self._last_display_step = self._start_step + self._timed_time = 0.0 + self._timed_steps = 0 + self._recorded_any = False + + def record(self, display_step: int) -> DisplayInterval: + """Close the current interval and open the next one. + + Parameters + ---------- + display_step : int + One-based index of the step that just completed. + + Returns + ------- + DisplayInterval + Wall-clock summary of the interval that just ended. + """ + now = time.time() + wall_time = now - self._interval_start + steps = max(1, display_step - self._last_display_step) + self._interval_start = now + self._last_display_step = display_step + if self._counts_toward_average(display_step): + self._timed_time += wall_time + self._timed_steps += steps + # The opening interval absorbs the one-off costs of the run -- data + # preparation, graph compilation, autotuning -- and its rate would + # forecast a run many times longer than the real one. + forecasts = self._recorded_any + self._recorded_any = True + return DisplayInterval( + display_step=display_step, + wall_time=wall_time, + steps=steps, + eta=int((self._num_steps - display_step) * wall_time / steps) + if forecasts + else None, + timestamp=datetime.datetime.fromtimestamp( + now, tz=datetime.timezone.utc + ).astimezone(), + ) + + def format_average(self) -> str | None: + """Report the average step time of the run. + + Returns + ------- + str or None + The average time per step over the representative intervals, or + ``None`` when no interval was representative. + """ + if not self._timed_steps: + return None + message = ( + f"average training time: {self._timed_time / self._timed_steps:.4f} s/batch" + ) + excluded = self._num_steps - self._start_step - self._timed_steps + if excluded > 0: + message += f" ({excluded} batches excluded)" + return message + + def _counts_toward_average(self, display_step: int) -> bool: + """Whether an interval ending at ``display_step`` is representative.""" + if self._num_steps - self._start_step <= 2 * self._disp_freq: + return True + return display_step - 1 - self._start_step >= self._disp_freq diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 74f61f486a..62a279eb5f 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -11,9 +11,7 @@ annotations, ) -import datetime import logging -import time from abc import ( ABC, abstractmethod, @@ -50,6 +48,11 @@ format_training_message_per_task, ) +from .timing import ( + DisplayInterval, + TrainingTimer, +) + DEFAULT_TASK_KEY = "Default" log = logging.getLogger(__name__) @@ -532,9 +535,11 @@ def run(self, tasks: TrainingTaskCollection) -> None: try: self.on_train_begin(tasks) fout = self._open_learning_curve() - wall_start = time.time() - last_log_time = wall_start - last_log_step = start_step + timer = TrainingTimer( + start_step=start_step, + num_steps=num_steps, + disp_freq=self.trainer_config.disp_freq, + ) for step in range(start_step, num_steps): task = self.select_task(tasks) step_result = self.train_step(task, step) @@ -548,15 +553,6 @@ def run(self, tasks: TrainingTaskCollection) -> None: step=step, step_result=step_result, ) - current_time = time.time() - interval_wall_time = current_time - last_log_time - interval_steps = max(1, display_step - last_log_step) - self._log_interval( - display_step=display_step, - interval_wall_time=interval_wall_time, - interval_steps=interval_steps, - wall_elapsed=current_time - wall_start, - ) current_lr = self.learning_rate(step) self.lcurve_writer.log_results( step=display_step, @@ -564,6 +560,7 @@ def run(self, tasks: TrainingTaskCollection) -> None: train_results=train_results, valid_results=valid_results, ) + self._log_interval(timer.record(display_step)) if fout is not None: if fout.tell() == 0: self.lcurve_writer.write_header( @@ -578,8 +575,6 @@ def run(self, tasks: TrainingTaskCollection) -> None: train_results=train_results, valid_results=valid_results, ) - last_log_time = current_time - last_log_step = display_step self.run_full_validation( step=step, @@ -596,6 +591,7 @@ def run(self, tasks: TrainingTaskCollection) -> None: if self._should_save_final_checkpoint(): self.save_checkpoint(num_steps) + self._log_average_step_time(timer) finally: if fout is not None: fout.close() @@ -713,37 +709,27 @@ def _should_save_final_checkpoint(self) -> bool: return True return self.trainer_config.num_steps % self.trainer_config.save_freq != 0 - def _log_interval( - self, - *, - display_step: int, - interval_wall_time: float, - interval_steps: int, - wall_elapsed: float, - ) -> None: + def _log_interval(self, interval: DisplayInterval) -> None: if self.trainer_config.timing_in_training: - completed = max(1, display_step - self.trainer_config.start_step) - eta = int( - (self.trainer_config.num_steps - display_step) - / completed - * wall_elapsed - ) log.info( format_training_message( - batch=display_step, - wall_time=interval_wall_time, - eta=eta, - current_time=datetime.datetime.fromtimestamp( - time.time(), - tz=datetime.timezone.utc, - ).astimezone(), - step_time=interval_wall_time / interval_steps, + batch=interval.display_step, + wall_time=interval.wall_time, + eta=interval.eta, + current_time=interval.timestamp, ) ) else: log.info( format_training_message( - batch=display_step, - wall_time=interval_wall_time, + batch=interval.display_step, + wall_time=interval.wall_time, ) ) + + def _log_average_step_time(self, timer: TrainingTimer) -> None: + if not self.rank_context.is_chief or not self.trainer_config.timing_in_training: + return + message = timer.format_average() + if message is not None: + log.info(message) diff --git a/deepmd/loggers/training.py b/deepmd/loggers/training.py index d145f42897..7e24574af1 100644 --- a/deepmd/loggers/training.py +++ b/deepmd/loggers/training.py @@ -1,7 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + import datetime import logging import math +from typing import ( + TYPE_CHECKING, +) + +if TYPE_CHECKING: + from collections.abc import ( + Mapping, + ) log = logging.getLogger(__name__) @@ -33,12 +45,45 @@ def _format_estimated_finish_time( return finish_time.strftime("%Y-%m-%d %H:%M") +def log_parameter_counts( + counts: Mapping[str, tuple[int, int]], + *, + multi_task: bool, +) -> None: + """Log the parameter count of every task. + + Parameters + ---------- + counts : Mapping[str, tuple[int, int]] + The ``(trainable, total)`` parameter count of each task, keyed by task + key and ordered as the trainer orders its tasks. + multi_task : bool + Whether the run trains several task branches. Tasks may share + parameters, which per-task counts cannot express, so a multi-task run + is reported per task and flagged as approximate. + """ + if not multi_task: + trainable, total = next(iter(counts.values())) + log.info( + f"Model Params: {total / 1e6:.3f} M (Trainable: {trainable / 1e6:.3f} M)" + ) + return + log.warning( + "In multitask mode, parameters may be shared across tasks. " + "The following per-task counts may include duplicates." + ) + for task_key, (trainable, total) in counts.items(): + log.info( + f"Model Params [{task_key}]: {total / 1e6:.3f} M " + f"(Trainable: {trainable / 1e6:.3f} M)" + ) + + def format_training_message( batch: int, wall_time: float, eta: int | None = None, current_time: datetime.datetime | None = None, - step_time: float | None = None, ) -> str: """Format the summary message for one training interval. @@ -53,9 +98,6 @@ def format_training_message( current_time : datetime.datetime | None, optional Current local time used to estimate the finish timestamp. This is only used when ``eta`` is provided. - step_time : float | None, optional - Average wall-clock time per training step over this interval, in - seconds. Shown only when provided. Returns ------- @@ -63,8 +105,6 @@ def format_training_message( The formatted training message. """ msg = f"Batch {batch:7d}: total wall time = {wall_time:.2f} s" - if step_time is not None: - msg += f", avg = {step_time:.4f} s/step" if isinstance(eta, int): eta_seconds = int(eta) msg += ( diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index c1699c958b..e37e228924 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import datetime import functools import json import logging -import time from collections.abc import ( Callable, Generator, @@ -26,10 +24,11 @@ import numpy as np import torch -from deepmd.common import ( - symlink_prefix_files, -) from deepmd.dpmodel.train import ( + DEFAULT_TASK_KEY, + CheckpointStore, + TrainingTimer, + build_checkpoint_stores, change_model_out_bias, resolve_step_schedule, ) @@ -39,6 +38,7 @@ from deepmd.loggers.training import ( format_training_message, format_training_message_per_task, + log_parameter_counts, ) from deepmd.pt.loss import ( DenoiseLoss, @@ -69,24 +69,6 @@ KFOptimizerWrapper, LKFOptimizer, ) -from deepmd.pt.train.ema import ( - EMA_CHECKPOINT_KEY, - ModelEMA, - get_ema_checkpoint_prefix, - get_ema_validation_log_path, -) -from deepmd.pt.train.utils import ( - NonFiniteGradGuard, - clip_grad_norm_, - latest_checkpoint_path, - resolve_best_checkpoint_dir, - resolve_keep_ckpt_count, - scoped_env_defaults, -) -from deepmd.pt.train.validation import ( - FullValidator, - resolve_full_validation_start_step, -) from deepmd.pt.train.wrapper import ( ModelWrapper, ) @@ -119,6 +101,23 @@ from deepmd.pt.utils.utils import ( to_numpy_array, ) +from deepmd.pt_expt.train.ema import ( + EMA_CHECKPOINT_KEY, + ModelEMA, + get_ema_checkpoint_prefix, +) +from deepmd.pt_expt.train.utils import ( + NonFiniteGradGuard, + clip_grad_norm_, + count_parameters, + infer_env_defaults, + resolve_best_checkpoint_dir, + scoped_env_defaults, +) +from deepmd.pt_expt.train.validation import ( + FullValidator, + build_full_validators, +) from deepmd.utils.data import ( DataRequirementItem, has_data_requirement, @@ -198,13 +197,7 @@ def __init__( optimizer_params = config.get("optimizer", {}) validating_params = config.get("validating") or {} - infer_env_defaults = {} - if bool(validating_params.get("compiled_infer", False)): - infer_env_defaults["DP_COMPILE_INFER"] = "1" - if bool(validating_params.get("tf32_infer", False)): - infer_env_defaults["DP_TF32_INFER"] = "1" - if bool(validating_params.get("amp_infer", False)): - infer_env_defaults["DP_AMP_INFER"] = "1" + eval_env_defaults = infer_env_defaults(validating_params) self.multi_task = "model_dict" in model_params self.finetune_links = finetune_links finetune_updates_statistics = finetune_links is not None and any( @@ -227,16 +220,9 @@ def __init__( self.disp_freq = training_params.get("disp_freq", 1000) self.disp_avg = training_params.get("disp_avg", False) self.save_ckpt = training_params.get("save_ckpt", "model.ckpt") - save_dir = training_params.get("save_dir") - self.save_dir = Path(save_dir) if save_dir else None - if self.save_dir is not None and self.rank == 0: - self.save_dir.mkdir(parents=True, exist_ok=True) self.save_freq = training_params.get("save_freq", 1000) - self.max_ckpt_keep = training_params.get("max_ckpt_keep", 5) - self.ckpt_keep_ratio = training_params.get("ckpt_keep_ratio") self.enable_ema = bool(training_params.get("enable_ema", False)) self.ema_decay = float(training_params.get("ema_decay", 0.999)) - self.ema_ckpt_keep = int(training_params.get("ema_ckpt_keep", 3)) self.ema_save_ckpt = get_ema_checkpoint_prefix(self.save_ckpt) self.display_in_training = training_params.get("disp_training", True) self.timing_in_training = training_params.get("time_training", True) @@ -475,7 +461,7 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: # Model # SeZMModel samples these eval/inference env vars exactly once inside # __init__; keep config-derived defaults scoped to construction. - with scoped_env_defaults(infer_env_defaults): + with scoped_env_defaults(eval_env_defaults): self.model = get_model_for_wrapper( model_params, resuming=resuming, @@ -716,23 +702,15 @@ def epoch_length(model_key: str) -> int: self.num_steps = schedule.num_steps self.model_prob = schedule.model_prob - # === Derive checkpoint retention from ckpt_keep_ratio === - # num_steps is final here (including when derived from num_epoch), so the - # ratio can be converted into an absolute keep count once. - keep_ckpt_count = resolve_keep_ckpt_count( - self.ckpt_keep_ratio, self.num_steps, self.save_freq + # === Checkpoint layout === + # num_steps is final here (including when derived from num_epoch), so a + # retention ratio can be converted into an absolute keep count once. + self.ckpt_store, self.ema_ckpt_store = build_checkpoint_stores( + training_params, + num_steps=self.num_steps, + ema_prefix=self.ema_save_ckpt, + rank=self.rank, ) - if keep_ckpt_count is not None: - self.max_ckpt_keep = keep_ckpt_count - self.ema_ckpt_keep = keep_ckpt_count - log.info( - "Resolved checkpoint retention to %d from ckpt_keep_ratio=%s " - "(num_steps=%d, save_freq=%d).", - keep_ckpt_count, - self.ckpt_keep_ratio, - self.num_steps, - self.save_freq, - ) # Learning rate self.gradient_max_norm = training_params.get("gradient_max_norm", 0.0) @@ -1174,14 +1152,7 @@ def update_finetune_bias( self.enable_profiler = training_params.get("enable_profiler", False) self.profiling = training_params.get("profiling", False) self.profiling_file = training_params.get("profiling_file", "timeline.json") - self.full_validator = None - self.ema_full_validator = None - - self.full_validator = self._create_full_validator( - validating_params=validating_params, - validation_data=validation_data, - ) - self.ema_full_validator = self._create_ema_full_validator( + self.full_validator, self.ema_full_validator = self._create_full_validators( validating_params=validating_params, validation_data=validation_data, ) @@ -1247,91 +1218,30 @@ def _create_lr_scheduler( last_epoch=start_step - 1, ) - def _create_full_validator( + def _create_full_validators( self, *, validating_params: dict[str, Any], validation_data: DpLoaderSet | None, - ) -> FullValidator | None: - """Create the runtime full validator when it is active.""" - if not self._is_validation_requested(validating_params, "full_validation"): - return None - self._raise_if_full_validation_unsupported(validation_data) - if validation_data is None: - raise RuntimeError( - "validation_data must be available after full validation checks." - ) - return FullValidator( + ) -> tuple[FullValidator | None, FullValidator | None]: + """Create the live-weight and EMA-weight full validators.""" + return build_full_validators( validating_params=validating_params, validation_data=validation_data, model=self.model, state_store=self._get_inner_module().train_infos, num_steps=self.num_steps, rank=self.rank, - zero_stage=self.zero_stage, - restart_training=self.restart_training, - checkpoint_dir=resolve_best_checkpoint_dir( - validating_params, self.save_ckpt - ), - ) - - def _create_ema_full_validator( - self, - *, - validating_params: dict[str, Any], - validation_data: DpLoaderSet | None, - ) -> FullValidator | None: - """Create the runtime EMA full validator when it is active. - - EMA full validation is independent from regular full validation: it - can be enabled on its own to validate only the EMA-smoothed model. - """ - if not self._is_validation_requested(validating_params, "ema_full_validation"): - return None - if self.model_ema is None: - # EMA full validation needs the EMA-smoothed model; when EMA is - # disabled the option is silently ignored rather than failing. - return None - self._raise_if_full_validation_unsupported(validation_data) - if validation_data is None: - raise RuntimeError( - "validation_data must be available after EMA full validation checks." - ) - ema_validating_params = dict(validating_params) - ema_validating_params["full_validation"] = True - return FullValidator( - validating_params=ema_validating_params, - validation_data=validation_data, - model=self.model, - state_store=self.model_ema.validation_state, - num_steps=self.num_steps, - rank=self.rank, - zero_stage=self.zero_stage, restart_training=self.restart_training, checkpoint_dir=resolve_best_checkpoint_dir( validating_params, self.save_ckpt ), - full_val_file=get_ema_validation_log_path( - validating_params.get("full_val_file", "val.log") + ensure_supported=lambda: self._raise_if_full_validation_unsupported( + validation_data ), - best_checkpoint_prefix="best_ema.ckpt", - emit_best_save_log=False, - model_eval_context=lambda: self.model_ema.apply_shadow(self.model), - ) - - def _is_validation_requested( - self, - validating_params: dict[str, Any], - flag_name: str, - ) -> bool: - """Check whether a full validation flow can trigger during this run.""" - if not validating_params.get(flag_name, False): - return False - start_step = resolve_full_validation_start_step( - validating_params.get("full_val_start", 0.5), - self.num_steps, + model_ema=self.model_ema, + zero_stage=self.zero_stage, ) - return start_step is not None and start_step <= self.num_steps def _raise_if_full_validation_unsupported( self, @@ -1362,42 +1272,13 @@ def _raise_if_full_validation_unsupported( "training with training.zero_stage < 2." ) - @staticmethod - def _count_parameters(model: torch.nn.Module) -> tuple[int, int]: - """ - Count model parameters. - - Parameters - ---------- - model : torch.nn.Module - The model to count parameters for. - - Returns - ------- - tuple[int, int] - A tuple of (trainable, total) parameter counts. - """ - trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) - total = sum(p.numel() for p in model.parameters()) - return trainable, total - def _log_parameter_count(self) -> None: """Log model parameter count.""" - if not self.multi_task: - trainable, total = self._count_parameters(self.model) - log.info( - f"Model Params: {total / 1e6:.3f} M (Trainable: {trainable / 1e6:.3f} M)" - ) - else: - log.warning( - "In multitask mode, parameters may be shared across tasks. " - "The following per-task counts may include duplicates." - ) - for model_key in self.model_keys: - trainable, total = self._count_parameters(self.model[model_key]) - log.info( - f"Model Params [{model_key}]: {total / 1e6:.3f} M (Trainable: {trainable / 1e6:.3f} M)" - ) + models = self.model if self.multi_task else {DEFAULT_TASK_KEY: self.model} + log_parameter_counts( + {key: count_parameters(models[key]) for key in self.model_keys}, + multi_task=self.multi_task, + ) def _create_optimizer( self, @@ -1823,39 +1704,16 @@ def log_loss_valid(_task_key: str = "Default") -> dict: self.step_count_in_interval = 0 self.last_display_step = display_step_id - current_time = time.time() - train_time = current_time - self.t0 - self.t0 = current_time + interval = self.step_timer.record(display_step_id) if self.rank == 0 and self.timing_in_training: - eta = int( - (self.num_steps - display_step_id) - / min(self.disp_freq, display_step_id - self.start_step) - * train_time - ) log.info( format_training_message( batch=display_step_id, - wall_time=train_time, - eta=eta, - current_time=datetime.datetime.fromtimestamp( - current_time, - tz=datetime.timezone.utc, - ).astimezone(), + wall_time=interval.wall_time, + eta=interval.eta, + current_time=interval.timestamp, ) ) - if ( - (self.num_steps - self.start_step) - <= 2 * self.disp_freq # not enough steps - or (_step_id - self.start_step) - >= self.disp_freq # skip first disp_freq steps - ): - self.total_train_time += train_time - if display_step_id == 1: - self.timed_steps += 1 - else: - self.timed_steps += min( - self.disp_freq, _step_id - self.start_step - ) if fout: if self.lcurve_should_print_header: @@ -1901,27 +1759,18 @@ def log_loss_valid(_task_key: str = "Default") -> dict: self.zero_stage > 0 or self.rank == 0 or dist.get_rank() == 0 ): # Handle the case if rank 0 aborted and re-assigned - self.latest_model = latest_checkpoint_path( - self.save_ckpt, display_step_id, self.save_dir - ) + self.latest_model = self.ckpt_store.path_for(display_step_id) self.save_model(self.latest_model, lr=cur_lr, step=_step_id) if self.rank == 0 or dist.get_rank() == 0: log.info(f"Saved model to {self.latest_model}") - symlink_prefix_files( - str(self.latest_model.with_suffix("")), self.save_ckpt - ) - with open("checkpoint", "w") as f: - f.write(str(self.latest_model)) + self.ckpt_store.publish(self.latest_model) if self.model_ema is not None: - self.latest_ema_model = latest_checkpoint_path( - self.ema_save_ckpt, display_step_id, self.save_dir + self.latest_ema_model = self.ema_ckpt_store.path_for( + display_step_id ) self.save_ema_model(self.latest_ema_model, lr=cur_lr, step=_step_id) if self.rank == 0 or dist.get_rank() == 0: - symlink_prefix_files( - str(self.latest_ema_model.with_suffix("")), - self.ema_save_ckpt, - ) + self.ema_ckpt_store.publish(self.latest_ema_model) # tensorboard if self.enable_tensorboard and ( @@ -1961,9 +1810,11 @@ def log_loss_valid(_task_key: str = "Default") -> dict: ) self.wrapper.train() - self.t0 = time.time() - self.total_train_time = 0.0 - self.timed_steps = 0 + self.step_timer = TrainingTimer( + start_step=self.start_step, + num_steps=self.num_steps, + disp_freq=self.disp_freq, + ) self._discarded_training_batches = 0 if self.disp_avg: @@ -2006,36 +1857,26 @@ def log_loss_valid(_task_key: str = "Default") -> dict: self.get_sample_func[model_key], _bias_adjust_mode="change-by-statistic", ) - self.latest_model = latest_checkpoint_path( - self.save_ckpt, self.num_steps, self.save_dir - ) + self.latest_model = self.ckpt_store.path_for(self.num_steps) cur_lr = self.lr_schedule.value(self.num_steps - 1) self.save_model(self.latest_model, lr=cur_lr, step=self.num_steps - 1) log.info(f"Saved model to {self.latest_model}") - symlink_prefix_files(str(self.latest_model.with_suffix("")), self.save_ckpt) - with open("checkpoint", "w") as f: - f.write(str(self.latest_model)) + self.ckpt_store.publish(self.latest_model) if self.model_ema is not None: - self.latest_ema_model = latest_checkpoint_path( - self.ema_save_ckpt, self.num_steps, self.save_dir - ) + self.latest_ema_model = self.ema_ckpt_store.path_for(self.num_steps) self.save_ema_model( self.latest_ema_model, lr=cur_lr, step=self.num_steps - 1, ) - symlink_prefix_files( - str(self.latest_ema_model.with_suffix("")), self.ema_save_ckpt - ) + self.ema_ckpt_store.publish(self.latest_ema_model) if self.num_steps == 0 and self.zero_stage > 0: # ZeRO-1 / FSDP: all ranks participate in save_model (collective op) - self.latest_model = latest_checkpoint_path(self.save_ckpt, 0, self.save_dir) + self.latest_model = self.ckpt_store.path_for(0) self.save_model(self.latest_model, lr=0, step=0) if self.model_ema is not None: - self.latest_ema_model = latest_checkpoint_path( - self.ema_save_ckpt, 0, self.save_dir - ) + self.latest_ema_model = self.ema_ckpt_store.path_for(0) self.save_ema_model(self.latest_ema_model, lr=0, step=0) if ( @@ -2044,32 +1885,20 @@ def log_loss_valid(_task_key: str = "Default") -> dict: if self.num_steps == 0: if self.zero_stage == 0: # When num_steps is 0, the checkpoint is never saved in the loop - self.latest_model = latest_checkpoint_path( - self.save_ckpt, 0, self.save_dir - ) + self.latest_model = self.ckpt_store.path_for(0) self.save_model(self.latest_model, lr=0, step=0) if self.model_ema is not None: - self.latest_ema_model = latest_checkpoint_path( - self.ema_save_ckpt, 0, self.save_dir - ) + self.latest_ema_model = self.ema_ckpt_store.path_for(0) self.save_ema_model(self.latest_ema_model, lr=0, step=0) log.info(f"Saved model to {self.latest_model}") - symlink_prefix_files( - str(self.latest_model.with_suffix("")), self.save_ckpt - ) - with open("checkpoint", "w") as f: - f.write(str(self.latest_model)) + self.ckpt_store.publish(self.latest_model) if self.model_ema is not None: - symlink_prefix_files( - str(self.latest_ema_model.with_suffix("")), self.ema_save_ckpt - ) + self.ema_ckpt_store.publish(self.latest_ema_model) - if self.timing_in_training and self.timed_steps: - msg = f"average training time: {self.total_train_time / self.timed_steps:.4f} s/batch" - excluded_steps = self.num_steps - self.start_step - self.timed_steps - if excluded_steps > 0: - msg += f" ({excluded_steps} batches excluded)" - log.info(msg) + if self.timing_in_training: + average_message = self.step_timer.format_average() + if average_message is not None: + log.info(average_message) if JIT: pth_model_path = ( @@ -2180,28 +2009,14 @@ def _collect_checkpoint_states( optim_state = self.optimizer.state_dict() if include_optimizer else None return model_state, optim_state - @staticmethod - def _parse_checkpoint_step(path: Path, prefix_name: str) -> int | None: - """Parse the checkpoint step from ``-.pt`` filenames.""" - checkpoint_prefix = f"{prefix_name}-" - if path.suffix != ".pt" or not path.name.startswith(checkpoint_prefix): - return None - step_text = path.name[len(checkpoint_prefix) : -len(path.suffix)] - if not step_text.isdigit(): - return None - return int(step_text) - def _write_checkpoint( self, save_path: Path, checkpoint_data: dict[str, Any], *, - ckpt_prefix: str, - max_ckpt_keep: int, + store: CheckpointStore, ) -> None: - """Write a checkpoint file and apply prefix-based cleanup.""" - prefix_name = Path(ckpt_prefix).name - + """Write a checkpoint file and apply the store's retention policy.""" # === Only rank 0 writes to disk === if self.rank != 0: return @@ -2210,27 +2025,7 @@ def _write_checkpoint( for item in optim_state["param_groups"]: item["lr"] = float(item["lr"]) torch.save(checkpoint_data, save_path) - checkpoint_dir = save_path.parent - checkpoint_files = [] - for checkpoint_file in checkpoint_dir.glob("*.pt"): - step = self._parse_checkpoint_step(checkpoint_file, prefix_name) - if checkpoint_file.is_symlink() or step is None: - continue - checkpoint_files.append((checkpoint_file, step)) - - current_step = self._parse_checkpoint_step(save_path, prefix_name) - if current_step is not None: - fresh_checkpoint_files = [] - for checkpoint_file, step in checkpoint_files: - if step > current_step: - checkpoint_file.unlink() - else: - fresh_checkpoint_files.append((checkpoint_file, step)) - checkpoint_files = fresh_checkpoint_files - - checkpoint_files.sort(key=lambda item: (item[1], item[0].name)) - while len(checkpoint_files) > max_ckpt_keep: - checkpoint_files.pop(0)[0].unlink() + store.prune(save_path) def save_model( self, @@ -2238,8 +2033,7 @@ def save_model( lr: float = 0.0, step: int = 0, *, - ckpt_prefix: str | None = None, - max_ckpt_keep: int | None = None, + store: CheckpointStore | None = None, use_ema_weights: bool = False, include_ema_state: bool = True, include_optimizer: bool = True, @@ -2259,10 +2053,7 @@ def save_model( self._write_checkpoint( Path(save_path), checkpoint_data, - ckpt_prefix=self.save_ckpt if ckpt_prefix is None else ckpt_prefix, - max_ckpt_keep=( - self.max_ckpt_keep if max_ckpt_keep is None else max_ckpt_keep - ), + store=self.ckpt_store if store is None else store, ) def save_ema_model( @@ -2277,8 +2068,7 @@ def save_ema_model( save_path, lr=lr, step=step, - ckpt_prefix=self.ema_save_ckpt, - max_ckpt_keep=self.ema_ckpt_keep, + store=self.ema_ckpt_store, use_ema_weights=True, include_ema_state=False, include_optimizer=False, @@ -2290,8 +2080,7 @@ def save_model_merged( lr: float = 0.0, step: int = 0, *, - ckpt_prefix: str | None = None, - max_ckpt_keep: int | None = None, + store: CheckpointStore | None = None, use_ema_weights: bool = False, ) -> None: """Save a plain SeZM checkpoint with LoRA adapters folded into base weights. @@ -2331,10 +2120,7 @@ def save_model_merged( self._write_checkpoint( Path(save_path), {"model": merged_state}, - ckpt_prefix=self.save_ckpt if ckpt_prefix is None else ckpt_prefix, - max_ckpt_keep=( - self.max_ckpt_keep if max_ckpt_keep is None else max_ckpt_keep - ), + store=self.ckpt_store if store is None else store, ) def save_ema_model_merged( @@ -2349,8 +2135,7 @@ def save_ema_model_merged( save_path, lr=lr, step=step, - ckpt_prefix=self.ema_save_ckpt, - max_ckpt_keep=self.ema_ckpt_keep, + store=self.ema_ckpt_store, use_ema_weights=True, ) diff --git a/deepmd/pt/train/ema.py b/deepmd/pt_expt/train/ema.py similarity index 100% rename from deepmd/pt/train/ema.py rename to deepmd/pt_expt/train/ema.py diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index cb2fb141a6..6625a1e09c 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -8,7 +8,6 @@ import functools import logging -import os import time from collections.abc import ( Callable, @@ -36,6 +35,7 @@ TrainingTask, TrainingTaskCollection, TrainStepResult, + build_checkpoint_stores, change_model_out_bias, change_model_out_bias_by_task, resolve_step_schedule, @@ -50,16 +50,12 @@ from deepmd.dpmodel.utils.training_utils import ( compute_total_numb_batch, ) +from deepmd.loggers.training import ( + log_parameter_counts, +) from deepmd.pt.optimizer import ( HybridMuonOptimizer, ) -from deepmd.pt.train.utils import ( - resolve_best_checkpoint_dir, -) -from deepmd.pt.train.validation import ( - FullValidator, - resolve_full_validation_start_step, -) from deepmd.pt.utils.compile_compat import ( apply_global_compile_patches, build_inductor_compile_options, @@ -84,6 +80,21 @@ from deepmd.pt_expt.model.graph_lower import ( model_uses_graph_lower, ) +from deepmd.pt_expt.train.ema import ( + EMA_CHECKPOINT_KEY, + ModelEMA, + get_ema_checkpoint_prefix, +) +from deepmd.pt_expt.train.utils import ( + count_parameters, + infer_env_defaults, + resolve_best_checkpoint_dir, + scoped_env_defaults, +) +from deepmd.pt_expt.train.validation import ( + FullValidator, + build_full_validators, +) from deepmd.pt_expt.train.wrapper import ( ModelWrapper, ) @@ -330,14 +341,6 @@ def _as_task_map( return {DEFAULT_TASK_KEY: value} -def _replace_latest_checkpoint_link(latest: Path, ckpt_path: Path) -> None: - """Point latest to ckpt_path using a target relative to latest's directory.""" - if latest.is_symlink() or latest.exists(): - latest.unlink() - latest.parent.mkdir(parents=True, exist_ok=True) - latest.symlink_to(os.path.relpath(ckpt_path, latest.parent)) - - # --------------------------------------------------------------------------- # torch.compile helpers # --------------------------------------------------------------------------- @@ -888,12 +891,14 @@ def __init__( task_buffers: dict[str, torch.Tensor] | None = None, compile_opts: dict[str, Any] | None = None, compiled_by_structure: dict | None = None, + task_key: str = DEFAULT_TASK_KEY, ) -> None: super().__init__() self.original_model = original_model self.compiled_forward_lower: torch.nn.Module | None = None self._task_buf_order = task_buf_order self._structure_key = structure_key + self._task_key = task_key self._compile_opts = compile_opts # Stored only for the first-forward compile call; freed afterwards. self._task_buffers = task_buffers @@ -907,6 +912,45 @@ def __init__( # (graph-eligible mixed_types descriptors) or the dense forward_lower. self._graph_eligible: bool | None = None + def _compiled_lower_for( + self, + path: str, + trace: "Callable[[], tuple[torch.nn.Module, tuple[str, ...]]]", + ) -> tuple[torch.nn.Module, tuple[str, ...]]: + """Return the compiled graph of this model, tracing it at most once. + + Tasks that share a model structure share one compiled graph, so a task + reaching this point second only reports the reuse. + + Parameters + ---------- + path : str + Name of the lowering being compiled, as it appears in the log. + trace : Callable[[], tuple[torch.nn.Module, tuple[str, ...]]] + Traces and compiles the graph, returning it together with the + order of the per-task buffers it expects. + + Returns + ------- + tuple[torch.nn.Module, tuple[str, ...]] + The compiled graph and its buffer order. + """ + attributes = f"task={self._task_key}, path={path}" + cached = self._compiled_by_structure.get(self._structure_key) + if cached is not None: + log.info("Reusing the graph compiled for an earlier task (%s).", attributes) + return cached + log.info("Tracing and compiling the model (%s).", attributes) + started = time.perf_counter() + compiled = trace() + log.info( + "Finished compiling (%s) in %.1f s.", + attributes, + time.perf_counter() - started, + ) + self._compiled_by_structure[self._structure_key] = compiled + return compiled + def __getattr__(self, name: str) -> Any: # Delegate unknown lookups to original_model so that callers such as # share_params (which calls .get_descriptor(), .atomic_model, etc.) and @@ -1043,18 +1087,9 @@ def forward( # Tasks sharing this structure key share the same descriptor / # fitting net and therefore the same dims, so a single compiled # graph is safe to reuse across them. - if self._structure_key in self._compiled_by_structure: - compiled_lower, buf_order = self._compiled_by_structure[ - self._structure_key - ] - log.info("Reusing compiled graph (shared model structure, lazy).") - else: - log.info( - "Lazy compile: tracing model on first forward call " - "(structure_key=%s).", - self._structure_key, - ) - compiled_lower, buf_order = _trace_and_compile( + compiled_lower, buf_order = self._compiled_lower_for( + "neighbor-list", + lambda: _trace_and_compile( self.original_model, ext_coord, ext_atype, @@ -1065,11 +1100,8 @@ def forward( charge_spin=charge_spin, task_buffers=self._task_buffers, compile_opts=self._compile_opts, - ) - self._compiled_by_structure[self._structure_key] = ( - compiled_lower, - buf_order, - ) + ), + ) self.compiled_forward_lower = compiled_lower self._task_buf_order = buf_order self._task_buffers = None # free; no longer needed after compile @@ -1231,29 +1263,17 @@ def _forward_graph( # Lazy compile of the GRAPH lower (cached per structure key). if self.compiled_forward_lower is None: - if self._structure_key in self._compiled_by_structure: - compiled_lower, buf_order = self._compiled_by_structure[ - self._structure_key - ] - log.info("Reusing compiled graph lower (shared structure, lazy).") - else: - log.info( - "Lazy compile (graph lower): tracing on first forward call " - "(structure_key=%s).", - self._structure_key, - ) - compiled_lower, buf_order = _trace_and_compile_graph( + compiled_lower, buf_order = self._compiled_lower_for( + "neighbor-graph", + lambda: _trace_and_compile_graph( _model, fparam, aparam, charge_spin, task_buffers=self._task_buffers, compile_opts=self._compile_opts, - ) - self._compiled_by_structure[self._structure_key] = ( - compiled_lower, - buf_order, - ) + ), + ) self.compiled_forward_lower = compiled_lower self._task_buf_order = buf_order self._task_buffers = None @@ -1422,7 +1442,9 @@ def __init__( self.disp_freq = training_params.get("disp_freq", 1000) self.save_ckpt = training_params.get("save_ckpt", "model.ckpt") self.save_freq = training_params.get("save_freq", 1000) - self.max_ckpt_keep = int(training_params.get("max_ckpt_keep", 5)) + self.enable_ema = bool(training_params.get("enable_ema", False)) + self.ema_decay = float(training_params.get("ema_decay", 0.999)) + self.ema_save_ckpt = get_ema_checkpoint_prefix(self.save_ckpt) self.display_in_training = training_params.get("disp_training", True) self.timing_in_training = training_params.get("time_training", True) self.change_bias_after_training = bool( @@ -1434,12 +1456,16 @@ def __init__( do_case_embd, case_embd_index = ( _get_case_embd_config(model_params) if self.multi_task else (False, {}) ) - for model_key in self.model_keys: - self.models[model_key] = get_model( - deepcopy(self.model_params_by_task[model_key]) - ).to(DEVICE) - if do_case_embd and not resuming: - self.models[model_key].set_case_embd(case_embd_index[model_key]) + # Descriptors sample the eval-time policy variables exactly once, while + # they are being constructed; keep the config-derived defaults scoped to + # construction so they do not leak into the rest of the process. + with scoped_env_defaults(infer_env_defaults(validating_params)): + for model_key in self.model_keys: + self.models[model_key] = get_model( + deepcopy(self.model_params_by_task[model_key]) + ).to(DEVICE) + if do_case_embd and not resuming: + self.models[model_key].set_case_embd(case_embd_index[model_key]) self.model = self.models if self.multi_task else self.models[DEFAULT_TASK_KEY] # Loss ---------------------------------------------------------------- @@ -1554,6 +1580,16 @@ def initialize_statistics( self.num_steps = schedule.num_steps self.model_prob = schedule.model_prob + # Checkpoint layout ---------------------------------------------------- + # num_steps is final here, so a retention ratio can be converted into an + # absolute keep count once. + self.ckpt_store, self.ema_ckpt_store = build_checkpoint_stores( + training_params, + num_steps=self.num_steps, + ema_prefix=self.ema_save_ckpt, + rank=self.rank, + ) + # Learning rate ------------------------------------------------------- self.lr_schedule = make_learning_rate_schedule( config["learning_rate"], self.num_steps @@ -1688,6 +1724,7 @@ def initialize_statistics( ) # Resume -------------------------------------------------------------- + ema_state_dict = None if resuming: log.info(f"Resuming from {resume_model}.") is_pte = resume_model.endswith((".pte", ".pt2")) @@ -1701,10 +1738,15 @@ def initialize_statistics( resume_model, map_location=DEVICE, weights_only=True ) 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. + continues_run = self.restart_training and finetune_model is None optimizer_state_dict = ( - state_dict["optimizer"] - if self.restart_training and finetune_model is None - else None + state_dict["optimizer"] if continues_run else None + ) + ema_state_dict = ( + state_dict.get(EMA_CHECKPOINT_KEY) if continues_run else None ) state_dict = state_dict["model"] else: @@ -1885,6 +1927,15 @@ def update_finetune_bias( last_epoch=self.start_step - 1, ) + # Exponential moving average ------------------------------------------- + # The shadow tracks the raw models, whose parameter tensors the compiled + # graphs keep sharing, so it is unaffected by compilation below. + self.model_ema = ( + ModelEMA(self.model, decay=self.ema_decay, state=ema_state_dict) + if self.enable_ema + else None + ) + self._configure_neighbor_graph_method( training_params.get("neighbor_graph_method", "auto") ) @@ -1894,7 +1945,8 @@ def update_finetune_bias( if self.enable_compile: check_compile_torch_version() compile_opts = training_params.get("compile_options", {}) - log.info("Compiling model with torch.compile (%s)", compile_opts) + if compile_opts: + log.info("torch.compile options: %s", compile_opts) self._compile_model(compile_opts) self.training_tasks = self._make_training_tasks() @@ -1907,53 +1959,41 @@ def update_finetune_bias( ), rank_context=RankContext(rank=self.rank, world_size=self.world_size), ) - self.full_validator = self._create_full_validator( + self.full_validator, self.ema_full_validator = self._create_full_validators( validating_params=validating_params, validation_data=self.validation_data if not self.multi_task else None, ) - def _create_full_validator( + if self.rank == 0: + log_parameter_counts( + {key: count_parameters(self.models[key]) for key in self.model_keys}, + multi_task=self.multi_task, + ) + + def _create_full_validators( self, *, validating_params: dict[str, Any], validation_data: Any | None, - ) -> FullValidator | None: - """Create the runtime full validator when it is active.""" - if not self._is_validation_requested(validating_params, "full_validation"): - return None - self._raise_if_full_validation_unsupported(validation_data) - if validation_data is None: - raise RuntimeError( - "validation_data must be available after full validation checks." - ) - return FullValidator( + ) -> tuple[FullValidator | None, FullValidator | None]: + """Create the live-weight and EMA-weight full validators.""" + return build_full_validators( validating_params=validating_params, validation_data=validation_data, - model=self.models[DEFAULT_TASK_KEY], + model=self.model, state_store=self._unwrapped.train_infos, num_steps=self.num_steps, rank=self.rank, - zero_stage=0, restart_training=self.restart_training, checkpoint_dir=resolve_best_checkpoint_dir( validating_params, self.save_ckpt ), + ensure_supported=lambda: self._raise_if_full_validation_unsupported( + validation_data + ), + model_ema=self.model_ema, ) - def _is_validation_requested( - self, - validating_params: dict[str, Any], - flag_name: str, - ) -> bool: - """Check whether a full validation flow can trigger during this run.""" - if not validating_params.get(flag_name, False): - return False - start_step = resolve_full_validation_start_step( - validating_params.get("full_val_start", 0.5), - self.num_steps, - ) - return start_step is not None and start_step <= self.num_steps - def _raise_if_full_validation_unsupported( self, validation_data: Any | None, @@ -2120,9 +2160,11 @@ def _compile_model(self, compile_opts: dict[str, Any]) -> None: task_buffers=task_bufs if task_bufs else None, compile_opts=compile_opts, compiled_by_structure=_compiled_by_structure, + task_key=task_key, ) log.info( - "Lazy compile registered (task=%s); will trace on first forward call.", + "Compilation enabled (task=%s); the graph is traced and compiled " + "on the first training step.", task_key, ) @@ -2284,12 +2326,16 @@ def _broadcast_value_from_rank0(self, value: Any) -> Any: # ------------------------------------------------------------------ def save_checkpoint(self, step: int) -> None: - ckpt_path = Path(f"{self.save_ckpt}-{step}.pt") + ckpt_path = self.ckpt_store.path_for(step) self._save_checkpoint_to_path(ckpt_path, step=step) - latest = Path(f"{self.save_ckpt}.pt") - _replace_latest_checkpoint_link(latest, ckpt_path) - self._cleanup_old_checkpoints() - log.info(f"Saved checkpoint to {ckpt_path}") + self.ckpt_store.publish(ckpt_path) + self.ckpt_store.prune(ckpt_path) + log.info(f"Saved model to {ckpt_path}") + if self.model_ema is not None: + ema_path = self.ema_ckpt_store.path_for(step) + self._save_checkpoint_to_path(ema_path, step=step, use_ema_weights=True) + self.ema_ckpt_store.publish(ema_path) + self.ema_ckpt_store.prune(ema_path) def _save_full_validation_checkpoint( self, @@ -2301,8 +2347,61 @@ def _save_full_validation_checkpoint( del lr self._save_checkpoint_to_path(save_path, step=step) - def _save_checkpoint_to_path(self, ckpt_path: Path, *, step: int) -> None: - """Serialize the current trainer state to an explicit checkpoint path.""" + def _save_full_validation_ema_checkpoint( + self, + save_path: Path, + lr: float = 0.0, + step: int = 0, + ) -> None: + """Save an EMA-weight checkpoint selected by EMA full validation. + + The validator restores the live weights before selecting a checkpoint, + so the shadow has to be applied again while writing it. + """ + del lr + self._save_checkpoint_to_path(save_path, step=step, use_ema_weights=True) + + def _save_checkpoint_to_path( + self, + ckpt_path: Path, + *, + step: int, + use_ema_weights: bool = False, + ) -> None: + """Serialize the current trainer state to an explicit checkpoint path. + + Parameters + ---------- + ckpt_path : Path + Destination of the checkpoint file. + step : int + Training step recorded in the checkpoint. + use_ema_weights : bool, optional + Whether to substitute the EMA-smoothed weights for the live ones. + 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. + """ + if use_ema_weights: + with self.model_ema.apply_shadow(self.model): + self._write_checkpoint( + ckpt_path, + step=step, + include_optimizer=False, + include_ema_state=False, + ) + return + self._write_checkpoint(ckpt_path, step=step) + + def _write_checkpoint( + self, + ckpt_path: Path, + *, + step: int, + include_optimizer: bool = True, + include_ema_state: bool = True, + ) -> None: + """Serialize the wrapper, and optionally optimizer and EMA state.""" self._unwrapped.train_infos["step"] = step # When compiled, wrapper.model[key] is _CompiledModel whose state_dict # uses keys like "original_model.*". Restart would load into a plain @@ -2317,33 +2416,17 @@ def _save_checkpoint_to_path(self, ckpt_path: Path, *, step: int) -> None: compiled_backup[task_key] = m wrapper.model[task_key] = m.original_model try: - state = { - "model": wrapper.state_dict(), - "optimizer": self.optimizer.state_dict(), - } + state: dict[str, Any] = {"model": wrapper.state_dict()} + if include_optimizer: + state["optimizer"] = self.optimizer.state_dict() finally: for task_key, compiled in compiled_backup.items(): wrapper.model[task_key] = compiled + 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) torch.save(state, ckpt_path) - def _cleanup_old_checkpoints(self) -> None: - """Remove old step checkpoint files beyond the retention limit.""" - if self.max_ckpt_keep <= 0: - return - ckpt_prefix_path = Path(self.save_ckpt) - ckpt_parent = ckpt_prefix_path.parent - ckpt_prefix = ckpt_prefix_path.name - checkpoints: list[tuple[int, Path]] = [] - for path in ckpt_parent.glob(f"{ckpt_prefix}-*.pt"): - if path.is_dir() or path.is_symlink(): - continue - step_text = path.name.removeprefix(f"{ckpt_prefix}-").removesuffix(".pt") - if step_text.isdigit(): - checkpoints.append((int(step_text), path)) - for _, path in sorted(checkpoints)[: -self.max_ckpt_keep]: - path.unlink(missing_ok=True) - # ------------------------------------------------------------------ # Training loop # ------------------------------------------------------------------ @@ -2376,7 +2459,6 @@ def _make_training_tasks(self) -> TrainingTaskCollection: def run(self) -> None: """Run pt_expt training through the backend-independent trainer loop.""" log.info("Start to train %d steps.", self.num_steps) - wall_start = time.time() try: super().run(self.training_tasks) if self.change_bias_after_training and self.num_steps > self.start_step: @@ -2385,7 +2467,8 @@ def run(self) -> None: self.save_checkpoint(self.num_steps) finally: self._close_data_systems() - log.info("Training finished. Total wall time: %.2fs", time.time() - wall_start) + if self.rank_context.is_chief: + log.info(f"Trained model has been saved to: {self.save_ckpt}") def _close_data_systems(self) -> None: """Release asynchronous data pipelines owned by this trainer.""" @@ -2422,16 +2505,19 @@ def run_full_validation( display_step: int, learning_rate: float, ) -> None: - """Run optional full validation for one step.""" - if self.full_validator is None: - return None - self.full_validator.run( - step_id=display_step, - display_step=display_step, - lr=learning_rate, - save_checkpoint=self._save_full_validation_checkpoint, - ) - return 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), + ): + if validator is None: + continue + validator.run( + step_id=display_step, + display_step=display_step, + lr=learning_rate, + save_checkpoint=save_checkpoint, + ) def select_task(self, tasks: TrainingTaskCollection) -> TrainingTask: """Select a task using DeePMD's seeded random helper.""" @@ -2492,6 +2578,8 @@ def train_step(self, task: TrainingTask, step: int) -> TrainStepResult: ) self._optimizer_step() + if self.model_ema is not None: + self.model_ema.update(self.model) return TrainStepResult( task_key=task_key, step=step, diff --git a/deepmd/pt/train/utils.py b/deepmd/pt_expt/train/utils.py similarity index 77% rename from deepmd/pt/train/utils.py rename to deepmd/pt_expt/train/utils.py index fe6d29b24e..4dde1b9579 100644 --- a/deepmd/pt/train/utils.py +++ b/deepmd/pt_expt/train/utils.py @@ -9,9 +9,6 @@ from contextlib import ( contextmanager, ) -from math import ( - ceil, -) from pathlib import ( Path, ) @@ -30,6 +27,25 @@ ) +def count_parameters(module: torch.nn.Module) -> tuple[int, int]: + """ + Count the parameters of a module. + + Parameters + ---------- + module : torch.nn.Module + The module to inspect. + + Returns + ------- + tuple[int, int] + The number of trainable parameters and the total number of parameters. + """ + trainable = sum(p.numel() for p in module.parameters() if p.requires_grad) + total = sum(p.numel() for p in module.parameters()) + return trainable, total + + def clip_grad_norm_( parameters: Iterable[torch.nn.Parameter], max_norm: float, @@ -184,6 +200,35 @@ def raise_nonfinite_gradient_norm( ) +def infer_env_defaults(validating_params: dict[str, Any]) -> dict[str, str]: + """ + Translate the eval-time policy options into environment defaults. + + Models sample these variables once, while they are being constructed, so + the configuration has to reach them through the environment rather than + through a constructor argument. A variable exported by the user takes + precedence; see :func:`scoped_env_defaults`. + + Parameters + ---------- + validating_params : dict[str, Any] + The normalized ``validating`` section. + + Returns + ------- + dict[str, str] + The environment variables requested by the configuration. + """ + flags = { + "compiled_infer": "DP_COMPILE_INFER", + "tf32_infer": "DP_TF32_INFER", + "amp_infer": "DP_AMP_INFER", + } + return { + name: "1" for flag, name in flags.items() if validating_params.get(flag, False) + } + + @contextmanager def scoped_env_defaults(defaults: dict[str, str]) -> Generator[None, None, None]: """Temporarily set missing environment variables and restore them afterward.""" @@ -200,30 +245,6 @@ def scoped_env_defaults(defaults: dict[str, str]) -> Generator[None, None, None] os.environ[key] = value -def latest_checkpoint_path(prefix: str, step_label: int, save_dir: Path | None) -> Path: - """ - Resolve the on-disk path of a periodic checkpoint file. - - Parameters - ---------- - prefix : str - The checkpoint prefix, e.g. ``model.ckpt`` or its EMA counterpart. - step_label : int - The training step encoded into the filename. - save_dir : Path or None - The configured checkpoint directory. When ``None`` the file follows - ``prefix`` relative to the working directory. - - Returns - ------- - Path - ``save_dir/-.pt`` when ``save_dir`` is set, otherwise - ``-.pt`` relative to the working directory. - """ - directory = save_dir if save_dir is not None else Path(prefix).parent - return directory / f"{Path(prefix).name}-{step_label}.pt" - - def resolve_best_checkpoint_dir( validating_params: dict[str, Any], save_ckpt: str ) -> Path: @@ -247,39 +268,3 @@ def resolve_best_checkpoint_dir( if save_best_dir: return Path(save_best_dir) return Path(save_ckpt).parent - - -def resolve_keep_ckpt_count( - ckpt_keep_ratio: float | None, num_steps: int, save_freq: int -) -> int | None: - """ - Convert a checkpoint-retention ratio into a sliding-window keep count. - - A checkpoint is written every ``save_freq`` steps and once more at the final - step, so a run of ``num_steps`` produces ``ceil(num_steps / save_freq)`` of - them in total (the terminal checkpoint is off-cadence when ``num_steps`` is - not a multiple of ``save_freq``). Keeping the most recent - ``ceil(ratio * total)`` is equivalent to retaining the final ``ratio`` - fraction of the run by step, without the caller computing the count by hand. - - Parameters - ---------- - ckpt_keep_ratio : float or None - The fraction of the training run, by step, whose periodic checkpoints - are retained. ``None`` leaves the keep count unchanged. - num_steps : int - The total number of training steps, already resolved (including when - derived from ``numb_epoch``). - save_freq : int - The checkpoint saving frequency in steps. - - Returns - ------- - int or None - The number of most recent checkpoints to keep (at least one), or - ``None`` when ``ckpt_keep_ratio`` is not set. - """ - if ckpt_keep_ratio is None: - return None - total_ckpts = max(1, ceil(num_steps / save_freq)) - return max(1, ceil(ckpt_keep_ratio * total_ckpts)) diff --git a/deepmd/pt/train/validation.py b/deepmd/pt_expt/train/validation.py similarity index 84% rename from deepmd/pt/train/validation.py rename to deepmd/pt_expt/train/validation.py index 784212ecb5..a3a822838c 100644 --- a/deepmd/pt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -35,20 +35,17 @@ from deepmd.pt.utils.auto_batch_size import ( AutoBatchSize, ) -from deepmd.pt.utils.dataset import ( - DeepmdDataSetForLoader, +from deepmd.pt.utils.utils import ( + to_torch_tensor, +) +from deepmd.pt_expt.train.ema import ( + get_ema_validation_log_path, ) -from deepmd.pt.utils.env import ( +from deepmd.pt_expt.utils.env import ( DEVICE, GLOBAL_PT_FLOAT_PRECISION, RESERVED_PRECISION_DICT, ) -from deepmd.pt.utils.lmdb_dataset import ( - LmdbDataset, -) -from deepmd.pt.utils.utils import ( - to_torch_tensor, -) from deepmd.utils.argcheck import ( normalize_full_validation_metric, resolve_full_validation_start_step, @@ -83,6 +80,7 @@ "full_validation_best_records", ) BEST_CKPT_PREFIX = "best.ckpt" +EMA_BEST_CKPT_PREFIX = "best_ema.ckpt" VAL_LOG_SIGNIFICANT_DIGITS = 5 VAL_LOG_COLUMN_GAP = " " VAL_LOG_HEADER_PREFIX = "# " @@ -425,24 +423,22 @@ def evaluate_all_systems(self) -> dict[str, float]: def _iter_validation_data_systems(self) -> Iterator[Any]: """Yield ``DeepmdData``-like systems to evaluate in this run. - - For ``DpLoaderSet``-style validation data, each entry in - ``validation_data.systems`` is a :class:`DeepmdDataSetForLoader`, - and we forward its underlying ``DeepmdData`` instance. - - For ``LmdbDataset`` validation data, we lazily materialize a - :class:`LmdbTestData` snapshot (cached across calls) and yield one - :class:`LmdbTestDataNlocView` per atom-count and label-availability - group. This keeps scalar ``find_*`` flags valid while excluding - default-filled labels from metrics. + The validation data of each backend is recognized by the surface it + exposes rather than by its type, so one validator serves them all: + + - An LMDB-backed dataset owns an ``_reader``. Its frames are lazily + materialized into a :class:`LmdbTestData` snapshot (cached across + calls) and yielded as one :class:`LmdbTestDataNlocView` per + atom-count and label-availability group. Grouping by atom count lets + mixed-nloc frames be stacked, and grouping by label availability + keeps the scalar ``find_*`` flags valid so default-filled labels stay + out of the metrics. + - A ``DeepmdDataSystem`` owns ``data_systems``, which are already + ``DeepmdData`` instances. + - A loader set owns ``systems``, each wrapping a ``DeepmdData`` in + ``data_system``. """ validation_data = self.validation_data - if isinstance(validation_data, LmdbDataset): - lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data) - for (nloc, _signature), indices in sorted( - lmdb_test_data.find_signature_groups.items() - ): - yield LmdbTestDataNlocView(lmdb_test_data, nloc, indices) - return - if hasattr(validation_data, "_reader"): lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data) for (nloc, _signature), indices in sorted( @@ -456,12 +452,13 @@ def _iter_validation_data_systems(self) -> Iterator[Any]: return for dataset in validation_data.systems: - if not isinstance(dataset, DeepmdDataSetForLoader): + data_system = getattr(dataset, "data_system", None) + if data_system is None: raise TypeError( "Full validation expects each dataset in validation_data.systems " - f"to be DeepmdDataSetForLoader, got {type(dataset)!r}." + f"to expose a `data_system`, got {type(dataset)!r}." ) - yield dataset.data_system + yield data_system def _get_lmdb_test_data_snapshot(self, lmdb_dataset: Any) -> LmdbTestData: """Build (once) and return the cached LMDB test snapshot. @@ -896,3 +893,132 @@ def _write_log_file(self, result: FullValidationResult) -> None: f"{result.saved_best_path} ({metric_label} = " f"{format_metric_number_for_log(metric_value)} {metric_unit})\n" ) + + +def _flow_can_trigger( + validating_params: dict[str, Any], + num_steps: int, + flag: str, +) -> bool: + """Whether a full validation flow is enabled and starts within the run.""" + if not validating_params.get(flag, False): + return False + start_step = resolve_full_validation_start_step( + validating_params.get("full_val_start", 0.5), + num_steps, + ) + return start_step is not None and start_step <= num_steps + + +def build_full_validators( + *, + validating_params: dict[str, Any], + validation_data: Any, + model: torch.nn.Module, + state_store: dict[str, Any], + num_steps: int, + rank: int, + restart_training: bool, + checkpoint_dir: Path, + ensure_supported: Callable[[], None], + model_ema: Any | None = None, + zero_stage: int = 0, +) -> tuple[FullValidator | None, FullValidator | None]: + """Build the full validators of a training run. + + A run may validate the live weights, the EMA-smoothed weights, or both. + The two flows share the schedule, the metric and the validation data, and + differ only in the weights they read, the log they write and the prefix of + the checkpoints they select, so they are configured together here. + + Parameters + ---------- + validating_params : dict[str, Any] + The normalized ``validating`` section. + validation_data : Any + The validation data of the run, required by both flows. + model : torch.nn.Module + The single-task model to evaluate. The EMA flow evaluates the same + module with the shadow weights swapped in. A multi-task run passes its + task mapping instead, which is never read because ``ensure_supported`` + rejects the run first. + state_store : dict[str, Any] + Where the live-weight flow records its best-checkpoint bookkeeping, + typically the trainer's ``train_infos``. The EMA flow keeps its own. + num_steps : int + The resolved run length. + rank : int + Process rank. + restart_training : bool + Whether the run continues an earlier one, in which case the validation + logs are appended to rather than truncated. + checkpoint_dir : Path + Directory receiving the best checkpoints of both flows. + ensure_supported : Callable[[], None] + Backend check raising when the run cannot be fully validated. It is + consulted once, and only when a flow would actually trigger. + model_ema : Any, optional + The EMA state of the run. Without it the EMA flow stays inactive, so + that ``ema_full_validation`` is ignored rather than rejected when EMA + itself is disabled. + zero_stage : int, optional + The ZeRO stage of the run, which decides whether checkpoint collection + is a collective operation. + + Returns + ------- + tuple[FullValidator | None, FullValidator | None] + The live-weight validator and the EMA-weight validator, each ``None`` + when its flow is inactive. + + Raises + ------ + RuntimeError + If validation data is missing after the backend check passed. + """ + live_active = _flow_can_trigger(validating_params, num_steps, "full_validation") + ema_active = model_ema is not None and _flow_can_trigger( + validating_params, num_steps, "ema_full_validation" + ) + if not (live_active or ema_active): + return None, None + ensure_supported() + if validation_data is None: + raise RuntimeError( + "validation_data must be available after full validation checks." + ) + + def make(**overrides: Any) -> FullValidator: + return FullValidator( + validation_data=validation_data, + model=model, + num_steps=num_steps, + rank=rank, + zero_stage=zero_stage, + restart_training=restart_training, + checkpoint_dir=checkpoint_dir, + **overrides, + ) + + live_validator = ( + make(validating_params=validating_params, state_store=state_store) + if live_active + else None + ) + if not ema_active: + return live_validator, None + # The EMA flow runs on its own switch, so its schedule must not depend on + # the live-weight one being enabled as well. + ema_params = dict(validating_params) + ema_params["full_validation"] = True + ema_validator = make( + validating_params=ema_params, + state_store=model_ema.validation_state, + full_val_file=get_ema_validation_log_path( + validating_params.get("full_val_file", "val.log") + ), + best_checkpoint_prefix=EMA_BEST_CKPT_PREFIX, + emit_best_save_log=False, + model_eval_context=lambda: model_ema.apply_shadow(model), + ) + return live_validator, ema_validator diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index c6c7fcfd78..a9d57435b4 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5591,7 +5591,7 @@ def training_args( [str, None], optional=True, default=None, - doc=supported_backends("pt") + doc_save_dir, + doc=supported_backends("pt", "pt_expt") + doc_save_dir, ), Argument( "save_ckpt", str, optional=True, default="model.ckpt", doc=doc_save_ckpt @@ -5602,7 +5602,7 @@ def training_args( [float, None], optional=True, default=None, - doc=supported_backends("pt") + doc_ckpt_keep_ratio, + doc=supported_backends("pt", "pt_expt") + doc_ckpt_keep_ratio, extra_check=lambda x: x is None or 0.0 < x < 1.0, extra_check_errmsg="must be a fraction in the open interval (0, 1)", ), @@ -5611,14 +5611,14 @@ def training_args( bool, optional=True, default=False, - doc=supported_backends("pt") + doc_enable_ema, + doc=supported_backends("pt", "pt_expt") + doc_enable_ema, ), Argument( "ema_decay", float, optional=True, default=0.999, - doc=supported_backends("pt") + doc_ema_decay, + doc=supported_backends("pt", "pt_expt") + doc_ema_decay, extra_check=lambda x: 0.0 <= x < 1.0, extra_check_errmsg="must be greater than or equal to 0 and less than 1", ), @@ -5627,7 +5627,7 @@ def training_args( int, optional=True, default=3, - doc=supported_backends("pt") + doc_ema_ckpt_keep, + doc=supported_backends("pt", "pt_expt") + doc_ema_ckpt_keep, extra_check=lambda x: x > 0, extra_check_errmsg="must be greater than 0", ), @@ -5936,7 +5936,9 @@ def validating_args() -> Argument: "flag is translated into `DP_TF32_INFER=1` at trainer startup before any " "model is constructed. A manually exported `DP_TF32_INFER` takes " "precedence over this option. This does not affect training forwards, " - "which are controlled by `model.enable_tf32`." + "which are controlled by `model.enable_tf32`. The PyTorch Exportable " + "backend always runs at full ('highest') matmul precision, so the " + "option has no effect there." ) doc_amp_infer = ( "Whether to enable bf16 automatic mixed precision for eval-time forwards " @@ -5960,7 +5962,7 @@ def validating_args() -> Argument: bool, optional=True, default=False, - doc=supported_backends("pt") + doc_ema_full_validation, + doc=supported_backends("pt", "pt_expt") + doc_ema_full_validation, ), Argument( "validation_freq", @@ -6024,7 +6026,7 @@ def validating_args() -> Argument: bool, optional=True, default=False, - doc=supported_backends("pt") + doc_compiled_infer, + doc=supported_backends("pt", "pt_expt") + doc_compiled_infer, ), Argument( "tf32_infer", @@ -6038,7 +6040,7 @@ def validating_args() -> Argument: bool, optional=True, default=False, - doc=supported_backends("pt") + doc_amp_infer, + doc=supported_backends("pt", "pt_expt") + doc_amp_infer, ), ] return Argument( diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 0112462581..fd34c85c85 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -103,8 +103,8 @@ Other keys in the {ref}`training ` section are explained below: - {ref}`disp_file ` The file for printing learning curve. - {ref}`disp_freq ` The frequency of printing learning curve. Set in the unit of training steps - {ref}`save_freq ` The frequency of saving checkpoint. -- {ref}`save_dir ` The directory where periodic checkpoints are written (PyTorch backend). It is created recursively if missing, while the `model.ckpt.pt` symlinks and the `checkpoint` pointer file stay in the working directory. Defaults to the working directory. -- {ref}`ckpt_keep_ratio ` An alternative to `max_ckpt_keep` (PyTorch backend) that keeps a sliding window of `ceil(ckpt_keep_ratio * ceil(numb_steps / save_freq))` most recent checkpoints, i.e. the final `ckpt_keep_ratio` fraction of the run by step. It overrides `max_ckpt_keep` (and `ema_ckpt_keep`) when set, and works the same whether the run length is given by `numb_steps` or `numb_epoch`. +- {ref}`save_dir ` The directory where periodic checkpoints are written (PyTorch and PyTorch Exportable backends). It is created recursively if missing, while the `model.ckpt.pt` symlinks and the `checkpoint` pointer file stay in the working directory. Defaults to the working directory. +- {ref}`ckpt_keep_ratio ` An alternative to `max_ckpt_keep` (PyTorch and PyTorch Exportable backends) that keeps a sliding window of `ceil(ckpt_keep_ratio * ceil(numb_steps / save_freq))` most recent checkpoints, i.e. the final `ckpt_keep_ratio` fraction of the run by step. It overrides `max_ckpt_keep` (and `ema_ckpt_keep`) when set, and works the same whether the run length is given by `numb_steps` or `numb_epoch`. ## Options and environment variables diff --git a/source/tests/common/dpmodel/test_train_abstract_trainer.py b/source/tests/common/dpmodel/test_train_abstract_trainer.py index 5ac21b71ee..bb450d67b3 100644 --- a/source/tests/common/dpmodel/test_train_abstract_trainer.py +++ b/source/tests/common/dpmodel/test_train_abstract_trainer.py @@ -172,6 +172,38 @@ def test_abstract_trainer_drives_single_task_loop(tmp_path: Path) -> None: assert "rmse_val" in lcurve.read_text() +def test_progress_is_logged_after_the_losses( + tmp_path: Path, + caplog: pytest.LogCaptureFixture, +) -> None: + """Each display reports the losses first and closes with the wall time. + + The order matches the other backends, so a single log parser serves them + all, and the run ends with the average step time. + """ + trainer = DummyTrainer( + TrainerConfig( + num_steps=4, + disp_file=str(tmp_path / "lcurve.out"), + disp_freq=2, + save_freq=4, + ) + ) + tasks = TrainingTaskCollection.single(DummyData([1.0]), DummyData([2.0])) + + with caplog.at_level("INFO", logger="deepmd.dpmodel.train.trainer"): + trainer.run(tasks) + + messages = [record.message for record in caplog.records] + display = [message for message in messages if message.startswith("Batch 2:")] + assert [segment.split(":")[1].strip().split(" ")[0] for segment in display] == [ + "trn", + "val", + "total", + ] + assert any(message.startswith("average training time:") for message in messages) + + def test_non_chief_rank_skips_user_visible_outputs(tmp_path: Path) -> None: lcurve = tmp_path / "lcurve.out" trainer = DummyTrainer( diff --git a/source/tests/common/dpmodel/test_train_checkpoint.py b/source/tests/common/dpmodel/test_train_checkpoint.py new file mode 100644 index 0000000000..788fa05e51 --- /dev/null +++ b/source/tests/common/dpmodel/test_train_checkpoint.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the backend-independent checkpoint layout and retention.""" + +from pathlib import ( + Path, +) + +import pytest + +from deepmd.dpmodel.train import ( + CheckpointStore, + build_checkpoint_stores, + resolve_keep_ckpt_count, +) + + +def _write(path: Path) -> Path: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("") + return path + + +def test_numbered_paths_follow_save_dir(tmp_path: Path) -> None: + store = CheckpointStore(tmp_path / "run" / "model.ckpt") + assert store.path_for(7) == tmp_path / "run" / "model.ckpt-7.pt" + + relocated = CheckpointStore( + tmp_path / "run" / "model.ckpt", save_dir=tmp_path / "ckpts" + ) + assert relocated.path_for(7) == tmp_path / "ckpts" / "model.ckpt-7.pt" + + +def test_publish_links_the_prefix_relative_to_its_directory(tmp_path: Path) -> None: + store = CheckpointStore( + tmp_path / "model.ckpt", pointer_file=tmp_path / "checkpoint" + ) + path = _write(store.path_for(3)) + + store.publish(path) + + latest = tmp_path / "model.ckpt.pt" + assert latest.is_symlink() + assert latest.resolve() == path + assert latest.readlink().as_posix() == "model.ckpt-3.pt" + assert (tmp_path / "checkpoint").read_text() == str(path) + + +def test_publish_reaches_across_save_dir(tmp_path: Path) -> None: + store = CheckpointStore(tmp_path / "model.ckpt", save_dir=tmp_path / "ckpts") + store.prepare() + path = _write(store.path_for(3)) + + store.publish(path) + + latest = tmp_path / "model.ckpt.pt" + assert latest.resolve() == path + assert latest.readlink().as_posix() == "ckpts/model.ckpt-3.pt" + + +def test_prune_keeps_the_newest_checkpoints(tmp_path: Path) -> None: + store = CheckpointStore(tmp_path / "model.ckpt", max_keep=2) + for step in (1, 2, 3): + _write(store.path_for(step)) + store.publish(store.path_for(3)) + + store.prune(store.path_for(3)) + + assert not store.path_for(1).exists() + assert store.path_for(2).exists() + assert store.path_for(3).exists() + assert (tmp_path / "model.ckpt.pt").exists() + + +def test_prune_drops_checkpoints_left_by_a_longer_run(tmp_path: Path) -> None: + """A rerun in a finished directory keeps its own checkpoint. + + Without dropping the higher-numbered remnants first, the retention window + would discard the checkpoint that was just written and leave the run with + no result at all. + """ + store = CheckpointStore(tmp_path / "model.ckpt", max_keep=2) + for step in (900, 950, 1000): + _write(store.path_for(step)) + current = _write(store.path_for(10)) + + store.prune(current) + + assert current.exists() + assert not store.path_for(900).exists() + assert not store.path_for(950).exists() + assert not store.path_for(1000).exists() + + +def test_prune_ignores_foreign_names_and_symlinks(tmp_path: Path) -> None: + store = CheckpointStore(tmp_path / "model.ckpt", max_keep=1) + other = _write(tmp_path / "best.ckpt-5.pt") + unnumbered = _write(tmp_path / "model.ckpt-final.pt") + current = _write(store.path_for(2)) + store.publish(current) + + store.prune(current) + + assert other.exists() + assert unnumbered.exists() + assert (tmp_path / "model.ckpt.pt").is_symlink() + + +def test_prune_without_a_window_keeps_every_checkpoint(tmp_path: Path) -> None: + """A disabled window deletes nothing, not even higher-numbered remnants.""" + store = CheckpointStore(tmp_path / "model.ckpt", max_keep=0) + for step in (1, 2, 900): + _write(store.path_for(step)) + + store.prune(store.path_for(2)) + + assert all(store.path_for(step).exists() for step in (1, 2, 900)) + + +def test_prune_from_a_foreign_path_spares_the_window(tmp_path: Path) -> None: + """A checkpoint outside the store neither dates files nor claims a slot. + + The name alone cannot decide membership: a validation checkpoint written + elsewhere may well parse as ``-``. + """ + store = CheckpointStore(tmp_path / "model.ckpt", max_keep=2) + for step in (100, 200): + _write(store.path_for(step)) + elsewhere = _write(tmp_path / "best" / "model.ckpt-150.pt") + + store.prune(elsewhere) + + assert store.path_for(100).exists() + assert store.path_for(200).exists() + assert elsewhere.exists() + + +def test_retention_ratio_maps_to_a_window_count() -> None: + assert resolve_keep_ckpt_count(None, 1000, 10) is None + # 1000 / 10 = 100 periodic checkpoints; 40% keeps the most recent 40. + assert resolve_keep_ckpt_count(0.4, 1000, 10) == 40 + # 4 periodic checkpoints; ceil(0.4 * 4) = ceil(1.6) = 2. + assert resolve_keep_ckpt_count(0.4, 4, 1) == 2 + # A save frequency above the run length yields a single, final checkpoint. + assert resolve_keep_ckpt_count(0.4, 5, 100) == 1 + + +def test_retention_ratio_handles_disabled_periodic_saving() -> None: + """Without periodic saving the run produces one checkpoint, not zero.""" + assert resolve_keep_ckpt_count(0.5, 1000, 0) == 1 + + +def test_built_stores_share_a_directory_and_split_the_pointer( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Only the periodic family owns the pointer file.""" + monkeypatch.chdir(tmp_path) + store, ema_store = build_checkpoint_stores( + { + "save_ckpt": "model.ckpt", + "save_dir": "ckpts", + "save_freq": 2, + "ckpt_keep_ratio": 0.5, + }, + num_steps=8, + ema_prefix="model_ema.ckpt", + ) + + # 4 periodic checkpoints; ceil(0.5 * 4) = 2 for both families. + assert (store.max_keep, ema_store.max_keep) == (2, 2) + assert store.directory == ema_store.directory == Path("ckpts") + assert store.directory.is_dir() + + ema_store.publish(_write(ema_store.path_for(1))) + assert Path("model_ema.ckpt.pt").is_symlink() + assert not Path("checkpoint").exists() + + store.publish(_write(store.path_for(2))) + assert Path("checkpoint").read_text() == str(store.path_for(2)) diff --git a/source/tests/common/dpmodel/test_train_timing.py b/source/tests/common/dpmodel/test_train_timing.py new file mode 100644 index 0000000000..b8b27a982a --- /dev/null +++ b/source/tests/common/dpmodel/test_train_timing.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the backend-independent training timer.""" + +import types + +import pytest + +from deepmd.dpmodel.train import ( + TrainingTimer, +) +from deepmd.dpmodel.train import timing as timing_module + + +class FakeClock: + """Monotonic clock advanced explicitly by the test.""" + + def __init__(self) -> None: + self.now = 0.0 + + def __call__(self) -> float: + return self.now + + def advance(self, seconds: float) -> None: + self.now += seconds + + +@pytest.fixture +def clock(monkeypatch: pytest.MonkeyPatch) -> FakeClock: + """Replace the clock the timer reads, leaving the stdlib one untouched.""" + fake = FakeClock() + monkeypatch.setattr(timing_module, "time", types.SimpleNamespace(time=fake)) + return fake + + +def test_eta_follows_the_latest_interval(clock: FakeClock) -> None: + """A run that speeds up is not held back by its slow start. + + A start-up cost paid once -- data preparation, compilation, autotuning -- + would dominate an estimate based on the average since the run began, and + forecasts nothing at all while it is still being paid. + """ + timer = TrainingTimer(start_step=0, num_steps=1000, disp_freq=100) + + clock.advance(100.0) + slow = timer.record(100) + clock.advance(10.0) + fast = timer.record(200) + + assert slow.wall_time == pytest.approx(100.0) + assert slow.eta is None + assert fast.wall_time == pytest.approx(10.0) + assert fast.eta == 80 + + +def test_interval_steps_span_the_display_gap(clock: FakeClock) -> None: + """The first display covers one step, so the next one covers the rest.""" + timer = TrainingTimer(start_step=0, num_steps=1000, disp_freq=100) + + clock.advance(1.0) + first = timer.record(1) + clock.advance(99.0) + second = timer.record(100) + + assert (first.display_step, first.steps) == (1, 1) + assert (second.display_step, second.steps) == (100, 99) + # The rate of the second interval, not of the run so far, drives the eta. + assert second.eta == 900 + + +def test_restart_measures_from_the_resumed_step(clock: FakeClock) -> None: + timer = TrainingTimer(start_step=500, num_steps=1000, disp_freq=100) + + clock.advance(50.0) + opening = timer.record(600) + clock.advance(50.0) + following = timer.record(700) + + assert (opening.steps, following.steps) == (100, 100) + # A restart pays the start-up costs again, so only the interval after the + # opening one forecasts: 0.5 s per step over the 300 steps left. + assert opening.eta is None + assert following.eta == 150 + + +def test_average_excludes_the_start_of_the_run(clock: FakeClock) -> None: + timer = TrainingTimer(start_step=0, num_steps=1000, disp_freq=100) + + clock.advance(100.0) + timer.record(100) + clock.advance(10.0) + timer.record(200) + + # Only the second interval is representative: 10 s over 100 steps. + assert timer.format_average() == ( + "average training time: 0.1000 s/batch (900 batches excluded)" + ) + + +def test_short_runs_keep_every_interval(clock: FakeClock) -> None: + timer = TrainingTimer(start_step=0, num_steps=150, disp_freq=100) + + clock.advance(50.0) + timer.record(100) + + assert timer.format_average() == ( + "average training time: 0.5000 s/batch (50 batches excluded)" + ) + + +def test_average_is_absent_without_a_timed_interval(clock: FakeClock) -> None: + timer = TrainingTimer(start_step=0, num_steps=1000, disp_freq=100) + + assert timer.format_average() is None diff --git a/source/tests/common/test_loggers_training.py b/source/tests/common/test_loggers_training.py new file mode 100644 index 0000000000..cfb466cb3e --- /dev/null +++ b/source/tests/common/test_loggers_training.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the training log messages.""" + +import datetime + +import pytest + +from deepmd.loggers.training import ( + format_training_message, + log_parameter_counts, +) + +_LOGGER = "deepmd.loggers.training" + + +def test_progress_message_reports_wall_time_alone() -> None: + assert ( + format_training_message(batch=100, wall_time=18.41) + == "Batch 100: total wall time = 18.41 s" + ) + + +def test_progress_message_appends_the_estimated_finish() -> None: + message = format_training_message( + batch=100, + wall_time=18.41, + eta=100, + current_time=datetime.datetime( + 2026, 6, 7, 5, 21, 29, tzinfo=datetime.timezone.utc + ), + ) + + assert message.startswith("Batch 100: total wall time = 18.41 s, eta = 0:01:40") + + +def test_single_task_parameter_count_is_reported_once( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level("INFO", logger=_LOGGER): + log_parameter_counts({"Default": (1_500_000, 2_000_000)}, multi_task=False) + + assert caplog.records[-1].message == "Model Params: 2.000 M (Trainable: 1.500 M)" + + +def test_multi_task_parameter_counts_are_flagged_as_approximate( + caplog: pytest.LogCaptureFixture, +) -> None: + with caplog.at_level("INFO", logger=_LOGGER): + log_parameter_counts( + {"a": (1_000_000, 1_000_000), "b": (500_000, 500_000)}, + multi_task=True, + ) + + messages = [record.message for record in caplog.records] + assert "may include duplicates" in messages[0] + assert messages[1].startswith("Model Params [a]: 1.000 M") + assert messages[2].startswith("Model Params [b]: 0.500 M") diff --git a/source/tests/pt/test_training.py b/source/tests/pt/test_training.py index a363e54d1a..05caed9b11 100644 --- a/source/tests/pt/test_training.py +++ b/source/tests/pt/test_training.py @@ -33,9 +33,6 @@ get_trainer, ) from deepmd.pt.entrypoints.main import train as train_entry -from deepmd.pt.train.ema import ( - EMA_CHECKPOINT_KEY, -) from deepmd.pt.train.training import ( all_ranks_have_valid_frames, ) @@ -50,6 +47,9 @@ make_stat_input, select_batch_frames, ) +from deepmd.pt_expt.train.ema import ( + EMA_CHECKPOINT_KEY, +) from deepmd.utils.argcheck import ( normalize, ) @@ -1017,7 +1017,7 @@ def tearDown(self) -> None: self._tmpdir.cleanup() @TRAINING_TEST_TIMEOUT - @patch("deepmd.pt.train.validation.FullValidator.evaluate_all_systems") + @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems") def test_full_validation_rotates_best_checkpoint(self, mocked_eval) -> None: mocked_eval.side_effect = [ {"mae_e_per_atom": 1.0}, @@ -1049,7 +1049,7 @@ def test_full_validation_rotates_best_checkpoint(self, mocked_eval) -> None: self.assertEqual(val_lines[1].split()[1], "2000.0") @TRAINING_TEST_TIMEOUT - @patch("deepmd.pt.train.validation.FullValidator.evaluate_all_systems") + @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems") def test_full_validation_save_best_dir(self, mocked_eval) -> None: mocked_eval.side_effect = [ {"mae_e_per_atom": 1.0}, @@ -1071,7 +1071,7 @@ def test_full_validation_save_best_dir(self, mocked_eval) -> None: self.assertEqual(list(Path(".").glob("best.ckpt-*.pt")), []) @TRAINING_TEST_TIMEOUT - @patch("deepmd.pt.train.validation.FullValidator.evaluate_all_systems") + @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems") def test_full_validation_runs_when_start_step_is_final_step( self, mocked_eval ) -> None: @@ -1268,8 +1268,8 @@ def test_ckpt_keep_ratio_overrides_keep_counts(self) -> None: trainer = get_trainer(config) # 4 periodic checkpoints; ceil(0.5 * 4) = 2 overrides both the regular # and EMA keep counts. - self.assertEqual(trainer.max_ckpt_keep, 2) - self.assertEqual(trainer.ema_ckpt_keep, 2) + self.assertEqual(trainer.ckpt_store.max_keep, 2) + self.assertEqual(trainer.ema_ckpt_store.max_keep, 2) save_ckpt = trainer.save_ckpt ema_save_ckpt = trainer.ema_save_ckpt trainer.run() @@ -1341,7 +1341,7 @@ def test_ema_checkpoint_rotation(self) -> None: def test_ema_checkpoint_cleanup_removes_future_steps(self) -> None: trainer = get_trainer(deepcopy(self.config)) - trainer.ema_ckpt_keep = 10 + trainer.ema_ckpt_store.max_keep = 10 ema_prefix = trainer.ema_save_ckpt Path(f"{ema_prefix}-999.pt").touch() @@ -1429,7 +1429,7 @@ def test_restart_restores_ema_state(self) -> None: ) @TRAINING_TEST_TIMEOUT - @patch("deepmd.pt.train.validation.FullValidator.evaluate_all_systems") + @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems") def test_ema_full_validation_writes_separate_outputs(self, mocked_eval) -> None: mocked_eval.side_effect = [ {"mae_e_per_atom": 10.0}, @@ -1460,7 +1460,7 @@ def test_ema_full_validation_writes_separate_outputs(self, mocked_eval) -> None: ) @TRAINING_TEST_TIMEOUT - @patch("deepmd.pt.train.validation.FullValidator.evaluate_all_systems") + @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems") def test_ema_full_validation_ignored_without_full_validation( self, mocked_eval ) -> None: diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py index 3fbf3e00c5..e16913bfbe 100644 --- a/source/tests/pt/test_validation.py +++ b/source/tests/pt/test_validation.py @@ -23,18 +23,18 @@ from deepmd.pt.model.model import ( get_model, ) -from deepmd.pt.train.validation import ( - BEST_METRIC_NAME_INFO_KEY, - TOPK_RECORDS_INFO_KEY, - FullValidator, - resolve_full_validation_start_step, -) from deepmd.pt.utils.env import ( DEVICE, ) from deepmd.pt.utils.lmdb_dataset import ( LmdbDataset, ) +from deepmd.pt_expt.train.validation import ( + BEST_METRIC_NAME_INFO_KEY, + TOPK_RECORDS_INFO_KEY, + FullValidator, + resolve_full_validation_start_step, +) from deepmd.utils.argcheck import ( normalize, ) diff --git a/source/tests/pt_expt/test_entrypoint.py b/source/tests/pt_expt/test_entrypoint.py index 9068df58d3..d6c456fa68 100644 --- a/source/tests/pt_expt/test_entrypoint.py +++ b/source/tests/pt_expt/test_entrypoint.py @@ -7,6 +7,7 @@ from deepmd.pt_expt.entrypoints.main import ( PTExptTrainEntrypoint, _ensure_pt_expt_model_suffix, + _ensure_stat_file_path, train, ) @@ -150,74 +151,10 @@ def test_pt_expt_entrypoint_rejects_random_model_key( ) -def test_pt_expt_checkpoint_cleanup_keeps_newest_steps(tmp_path) -> None: - from deepmd.pt_expt.train.training import ( - Trainer, - ) - - trainer = Trainer.__new__(Trainer) - trainer.save_ckpt = str(tmp_path / "model.ckpt") - trainer.max_ckpt_keep = 2 - for step in (1, 2, 3): - (tmp_path / f"model.ckpt-{step}.pt").write_text("") - (tmp_path / "model.ckpt.pt").symlink_to("model.ckpt-3.pt") - - trainer._cleanup_old_checkpoints() - - assert not (tmp_path / "model.ckpt-1.pt").exists() - assert (tmp_path / "model.ckpt-2.pt").exists() - assert (tmp_path / "model.ckpt-3.pt").exists() - assert (tmp_path / "model.ckpt.pt").exists() - - -def test_pt_expt_latest_checkpoint_link_uses_relative_target(tmp_path) -> None: - from deepmd.pt_expt.train.training import ( - _replace_latest_checkpoint_link, - ) - - ckpt_path = tmp_path / "ckpts" / "model-1.pt" - ckpt_path.parent.mkdir() - ckpt_path.write_text("") - latest = tmp_path / "ckpts" / "model.pt" - - _replace_latest_checkpoint_link(latest, ckpt_path) - - assert latest.is_symlink() - assert latest.resolve() == ckpt_path - assert latest.readlink().as_posix() == "model-1.pt" - - -def test_pt_expt_save_checkpoint_creates_parent_and_latest_link(tmp_path) -> None: - from deepmd.pt_expt.train.training import ( - Trainer, - ) - - class DummyWrapper: - train_infos: dict[str, int] - model: dict[str, object] - - def __init__(self) -> None: - self.train_infos = {} - self.model = {} - - def state_dict(self) -> dict[str, object]: - return {} - - class DummyOptimizer: - def state_dict(self) -> dict[str, object]: - return {} - - trainer = Trainer.__new__(Trainer) - trainer.wrapper = DummyWrapper() - trainer.optimizer = DummyOptimizer() - trainer.save_ckpt = str(tmp_path / "ckpts" / "model") - trainer.max_ckpt_keep = 2 +def test_pt_expt_stat_file_path_creates_hdf5_parent(tmp_path) -> None: + stat_file = tmp_path / "stats" / "model_stat.hdf5" - trainer.save_checkpoint(1) + stat_path = _ensure_stat_file_path(str(stat_file)) - ckpt_path = tmp_path / "ckpts" / "model-1.pt" - latest = tmp_path / "ckpts" / "model.pt" - assert ckpt_path.exists() - assert latest.is_symlink() - assert latest.resolve() == ckpt_path - assert latest.readlink().as_posix() == "model-1.pt" + assert stat_file.exists() + assert stat_path is not None diff --git a/source/tests/pt/test_train_utils.py b/source/tests/pt_expt/test_train_utils.py similarity index 85% rename from source/tests/pt/test_train_utils.py rename to source/tests/pt_expt/test_train_utils.py index c3944f0972..7e119824db 100644 --- a/source/tests/pt/test_train_utils.py +++ b/source/tests/pt_expt/test_train_utils.py @@ -3,10 +3,9 @@ import torch -from deepmd.pt.train.utils import ( +from deepmd.pt_expt.train.utils import ( NonFiniteGradGuard, clip_grad_norm_, - resolve_keep_ckpt_count, ) @@ -112,22 +111,5 @@ def test_resets_after_check(self) -> None: guard.raise_if_nonfinite(self._named(1.0)) -class TestResolveKeepCkptCount(unittest.TestCase): - def test_none_ratio_leaves_count_unchanged(self) -> None: - self.assertIsNone(resolve_keep_ckpt_count(None, 1000, 10)) - - def test_ratio_maps_to_recent_window_count(self) -> None: - # 1000 / 10 = 100 periodic checkpoints; 40% keeps the most recent 40. - self.assertEqual(resolve_keep_ckpt_count(0.4, 1000, 10), 40) - - def test_ratio_rounds_up(self) -> None: - # 4 periodic checkpoints; ceil(0.4 * 4) = ceil(1.6) = 2. - self.assertEqual(resolve_keep_ckpt_count(0.4, 4, 1), 2) - - def test_keeps_at_least_one(self) -> None: - # save_freq larger than num_steps yields a single (final) checkpoint. - self.assertEqual(resolve_keep_ckpt_count(0.4, 5, 100), 1) - - if __name__ == "__main__": unittest.main() diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index f32f860825..d5e5b19414 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -9,12 +9,14 @@ """ import copy -import datetime import math import os import shutil import tempfile import unittest +from collections.abc import ( + Callable, +) from pathlib import ( Path, ) @@ -27,9 +29,6 @@ import pytest import torch -from deepmd.loggers.training import ( - format_training_message, -) from deepmd.pt.optimizer import ( HybridMuonOptimizer, ) @@ -448,7 +447,7 @@ def test_zero_start_warmup_schedulers_construct(self) -> None: os.chdir(old_cwd) shutil.rmtree(tmpdir, ignore_errors=True) - @patch("deepmd.pt.train.validation.FullValidator.evaluate_all_systems") + @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems") def test_full_validation_loop(self, mocked_eval) -> None: """Run pt_expt full validation and verify best-checkpoint outputs.""" mocked_eval.side_effect = [ @@ -494,6 +493,86 @@ def test_full_validation_loop(self, mocked_eval) -> None: finally: shutil.rmtree(tmpdir, ignore_errors=True) + @patch("deepmd.pt_expt.train.validation.FullValidator.evaluate_all_systems") + def test_ema_full_validation_selects_its_own_best(self, mocked_eval) -> None: + """The EMA flow keeps a separate log, best prefix and best record.""" + # Both flows evaluate at every step, the live one first. + mocked_eval.side_effect = [ + {"mae_e_per_atom": 1.0}, + {"mae_e_per_atom": 2.0}, + {"mae_e_per_atom": 3.0}, + {"mae_e_per_atom": 0.5}, + ] + config = _make_config(self.data_dir, numb_steps=2) + config["training"]["save_freq"] = 100 + config["training"]["enable_ema"] = True + config["validating"] = { + "full_validation": True, + "ema_full_validation": True, + "validation_freq": 1, + "validation_metric": "E:MAE", + "full_val_start": 0.0, + } + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ema_full_validation_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + trainer = get_trainer(config) + self.assertIsNotNone(trainer.full_validator) + self.assertIsNotNone(trainer.ema_full_validator) + trainer.run() + + # The live flow improves at step 1, the EMA flow at step 2. + self.assertTrue(os.path.exists("best.ckpt-1.t-1.pt")) + self.assertTrue(os.path.exists("best_ema.ckpt-2.t-1.pt")) + self.assertTrue(os.path.exists("val.log")) + self.assertTrue(os.path.exists("val_ema.log")) + self.assertEqual( + trainer.model_ema.validation_state["full_validation_topk_records"], + [{"metric": 0.5, "step": 2}], + ) + + # The EMA best checkpoint carries the smoothed weights. + best_ema = torch.load("best_ema.ckpt-2.t-1.pt", weights_only=True) + live = torch.load("best.ckpt-1.t-1.pt", weights_only=True) + self.assertTrue( + any( + not torch.equal(value, live["model"][key]) + for key, value in best_ema["model"].items() + if isinstance(value, torch.Tensor) + and torch.is_floating_point(value) + ) + ) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_ema_full_validation_is_ignored_without_ema(self) -> None: + """The flow stays inactive when EMA itself is disabled.""" + config = _make_config(self.data_dir, numb_steps=2) + config["validating"] = { + "ema_full_validation": True, + "validation_freq": 1, + "full_val_start": 0.0, + } + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ema_full_validation_off_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + trainer = get_trainer(config) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + + self.assertIsNone(trainer.full_validator) + self.assertIsNone(trainer.ema_full_validator) + def test_training_loop_dpa4(self) -> None: """Run a few DPA4/SeZM training steps (model type "dpa4" dispatch).""" config = _make_config(self.data_dir, numb_steps=5) @@ -2021,59 +2100,167 @@ def test_compiled_matches_eager_per_task(self) -> None: shutil.rmtree(tmpdir, ignore_errors=True) -class TestFormatTrainingMessageStepTime(unittest.TestCase): - """The pt_expt trainer reports the average wall time per step over each - display interval by passing ``step_time`` to ``format_training_message`` - (replacing the former standalone ``step=... step_time=...`` debug line). - These tests cover both branches of the optional ``step_time``/``eta`` - arguments so the "avg = ... s/step" segment is rendered only when requested. - """ +class TestCheckpointRetention(unittest.TestCase): + """Test where periodic checkpoints land and how many are kept.""" - def test_without_step_time(self) -> None: - """``step_time=None`` (default) omits the step-time segment.""" - msg = format_training_message(batch=100, wall_time=18.41) - self.assertEqual(msg, "Batch 100: total wall time = 18.41 s") - self.assertNotIn("s/step", msg) - - def test_with_step_time(self) -> None: - """``step_time`` is rendered with 4 decimals after the wall time.""" - msg = format_training_message(batch=100, wall_time=18.41, step_time=0.1841) - self.assertEqual( - msg, - "Batch 100: total wall time = 18.41 s, avg = 0.1841 s/step", - ) + @classmethod + def setUpClass(cls) -> None: + data_dir = os.path.join(EXAMPLE_DIR, "data") + if not os.path.isdir(data_dir): + raise unittest.SkipTest(f"Example data not found: {data_dir}") + cls.data_dir = data_dir + + def _run_and_collect_steps( + self, + config: dict, + before_run: Callable[[str], None] | None = None, + ) -> tuple[list[int], str]: + """Train in a scratch directory and report the surviving checkpoints. - def test_step_time_zero_is_shown(self) -> None: - """A literal ``0.0`` step time is still shown (not treated as absent).""" - msg = format_training_message(batch=1, wall_time=0.5, step_time=0.0) - self.assertIn("avg = 0.0000 s/step", msg) + ``before_run`` receives the checkpoint directory and may seed it, which + is how a rerun over an earlier run's output is set up. + """ + tmpdir = tempfile.mkdtemp(prefix="pt_expt_save_dir_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + if before_run is not None: + before_run(os.path.join(tmpdir, "ckpts")) + get_trainer(config).run() + + ckpt_dir = os.path.join(tmpdir, "ckpts") + saved = sorted( + int(name[len("model.ckpt-") : -len(".pt")]) + for name in os.listdir(ckpt_dir) + if name.startswith("model.ckpt-") + ) + return saved, os.path.realpath(os.path.join(tmpdir, "model.ckpt.pt")) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) - def test_with_step_time_and_eta(self) -> None: - """Step time appears before the eta segment.""" - current_time = datetime.datetime( - 2026, 6, 7, 5, 21, 29, tzinfo=datetime.timezone.utc - ) - msg = format_training_message( - batch=100, - wall_time=18.41, - eta=100, - current_time=current_time, - step_time=0.1841, - ) - self.assertIn("total wall time = 18.41 s, avg = 0.1841 s/step, eta = ", msg) - # ordering: wall time -> step time -> eta - self.assertLess(msg.index("s/step"), msg.index("eta =")) - - def test_eta_without_step_time(self) -> None: - """Eta still works when no step time is supplied.""" - current_time = datetime.datetime( - 2026, 6, 7, 5, 21, 29, tzinfo=datetime.timezone.utc - ) - msg = format_training_message( - batch=100, wall_time=18.41, eta=100, current_time=current_time + def test_save_dir_holds_a_sliding_window_of_checkpoints(self) -> None: + config = _make_config(self.data_dir, numb_steps=6) + config["training"]["save_freq"] = 2 + config["training"]["save_dir"] = "ckpts" + config["training"]["max_ckpt_keep"] = 2 + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + saved, latest = self._run_and_collect_steps(config) + + self.assertEqual(saved, [4, 6]) + self.assertTrue(latest.endswith(os.path.join("ckpts", "model.ckpt-6.pt"))) + + def test_rerun_in_a_finished_directory_keeps_its_own_checkpoints(self) -> None: + """A short rerun is not pruned in favour of a longer run's leftovers.""" + config = _make_config(self.data_dir, numb_steps=2) + config["training"]["save_freq"] = 1 + config["training"]["save_dir"] = "ckpts" + config["training"]["max_ckpt_keep"] = 2 + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + def leave_stale_checkpoints(ckpt_dir: str) -> None: + os.makedirs(ckpt_dir, exist_ok=True) + for step in (900, 1000): + open(os.path.join(ckpt_dir, f"model.ckpt-{step}.pt"), "w").close() + + saved, latest = self._run_and_collect_steps( + config, before_run=leave_stale_checkpoints ) - self.assertNotIn("s/step", msg) - self.assertIn("eta = ", msg) + + self.assertEqual(saved, [1, 2]) + self.assertTrue(latest.endswith(os.path.join("ckpts", "model.ckpt-2.pt"))) + + def test_keep_ratio_retains_the_tail_of_the_run(self) -> None: + config = _make_config(self.data_dir, numb_steps=6) + config["training"]["save_freq"] = 2 + config["training"]["save_dir"] = "ckpts" + config["training"]["ckpt_keep_ratio"] = 0.5 + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + saved, _ = self._run_and_collect_steps(config) + + # 3 periodic checkpoints; ceil(0.5 * 3) = 2 most recent are kept. + self.assertEqual(saved, [4, 6]) + + +class TestEmaCheckpoints(unittest.TestCase): + """Test the EMA-smoothed weights and the checkpoints carrying them.""" + + @classmethod + def setUpClass(cls) -> None: + data_dir = os.path.join(EXAMPLE_DIR, "data") + if not os.path.isdir(data_dir): + raise unittest.SkipTest(f"Example data not found: {data_dir}") + cls.data_dir = data_dir + + def _make_ema_config(self, numb_steps: int = 4) -> dict: + config = _make_config(self.data_dir, numb_steps=numb_steps) + config["training"]["enable_ema"] = True + config["training"]["ema_decay"] = 0.9 + config["training"]["save_freq"] = numb_steps + config = update_deepmd_input(config, warning=False) + return normalize(config) + + def test_ema_checkpoint_holds_smoothed_weights(self) -> None: + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ema_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + trainer = get_trainer(self._make_ema_config()) + trainer.run() + + ema_ckpt = os.path.join(tmpdir, "model_ema.ckpt-4.pt") + self.assertTrue(os.path.exists(ema_ckpt)) + self.assertTrue(os.path.islink(os.path.join(tmpdir, "model_ema.ckpt.pt"))) + + ema_state = torch.load(ema_ckpt, weights_only=True) + live_state = torch.load( + os.path.join(tmpdir, "model.ckpt-4.pt"), weights_only=True + ) + # A deployment snapshot carries neither optimizer nor EMA state. + self.assertNotIn("optimizer", ema_state) + self.assertNotIn("ema", ema_state) + self.assertIn("ema", live_state) + + # The smoothed weights lag the live ones after a few updates. + differing = [ + key + for key, value in ema_state["model"].items() + if isinstance(value, torch.Tensor) + and torch.is_floating_point(value) + and not torch.equal(value, live_state["model"][key]) + ] + self.assertTrue(differing, "EMA weights should differ from live weights") + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + + def test_restart_restores_the_ema_shadow(self) -> None: + tmpdir = tempfile.mkdtemp(prefix="pt_expt_ema_restart_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + trainer = get_trainer(self._make_ema_config()) + trainer.run() + shadow = { + key: value.clone() + for key, value in trainer.model_ema.shadow_params.items() + } + + resumed = get_trainer( + self._make_ema_config(numb_steps=8), + restart_model=os.path.join(tmpdir, "model.ckpt-4.pt"), + ) + self.assertEqual(resumed.start_step, 4) + for key, value in shadow.items(): + torch.testing.assert_close(resumed.model_ema.shadow_params[key], value) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) if __name__ == "__main__": From 6d121cc611e2691162dcc8d9c9c6dc471287ec92 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 28 Jul 2026 14:14:04 +0800 Subject: [PATCH 2/5] feat(pt_expt): shard training state across ranks and guard the gradient Two gaps separated pt_expt from pt in distributed training: a run was always plain data parallel, and a step was taken without ever inspecting its gradient. Both are closed here, and the part that is not specific to PyTorch is shared with pt rather than duplicated. `training.zero_stage` now selects the same four strategies as in pt: plain DDP, DDP over a redundancy-sharded optimizer, and FSDP2 sharding the gradients or the parameters as well. What a stage implies -- which wrapper holds the model, how the optimizer is built, how a checkpoint is assembled, whether a gradient norm may be reduced locally -- follows from the stage alone, so `ShardingPolicy` in the backend-independent train layer states it once and both backends query it, rather than each comparing the stage against numbers wherever a decision is due; pt sheds twenty such comparisons. A single-process run drops the requested stage instead of failing, so one configuration stays usable whether or not it is launched across ranks. Assembling a checkpoint out of shards is a collective operation, which the shared training loop had no notion of: it called `save_checkpoint` on the chief alone, which would leave the other ranks waiting at the next barrier. The trainer gained `checkpoint_is_collective`, false by default so that tf2 and jax are unaffected; a backend that opts in is called on every rank and gates the write itself. A run is restored before the model is distributed. A checkpoint records whole tensors, and those cannot be copied into parameters that FSDP2 has already cut into shards. The optimizer is still built after distribution, so its state is restored separately, through the distributed-checkpoint API when the stage calls for it. That reorder also removes a second construction of the learning-rate schedule, and with it a defect it was covering: the schedule was built before the resumed step was known and rebuilt with the true value only when optimizer state was present, so restarting from a checkpoint that carried a step but no optimizer state -- a frozen model, or one saved without it -- resumed at the wrong learning rate. Sharding is rejected alongside multi-task training, EMA from stage two, and `change_bias_after_training`, as in pt. pt_expt additionally rejects `enable_compile`, whose graph is traced from the parameters that FSDP2 replaces with DTensors. Display-time validation is skipped once the parameters are sharded, because its forward gathers them while the display runs on the chief alone; the full-validation flow, which every rank enters together, stays available. pt_expt clipped gradients with the stock `torch.nn.utils.clip_grad_norm_` and never inspected the resulting norm, so a run could write a checkpoint of a model that had already diverged, and a gradient that was large but still representable could be misread as infinite when the sum of squares of the naive reduction overflowed. pt has carried safeguards against both for a while; they now serve pt_expt as well. The two safeguards move out of the training utilities, which had accreted four unrelated concerns, into `deepmd/pt_expt/train/gradient.py`. They share one rationale -- keep the reduction overflow-safe, and keep the verdict off the host until it is needed -- which the module can now state once. `deepmd/pt_expt/train/utils.py` keeps the trainer setup helpers. pt_expt feeds the norm to the guard on every step and consults the guard at the checkpoint boundary. The verdict is deliberately not read anywhere else: the check resets the accumulated state, so a second caller between two boundaries would consume a divergence that the checkpoint about to be written should have seen. Reading it once per boundary also keeps the step free of host synchronization, which is why the state is accumulated on device in the first place. The reduction itself is overflow-safe except where the parameters are sharded, since that reduction has to propagate DTensor sharding instead. Verified on two gloo ranks: every stage trains, checkpoints and restarts, recording whole tensors and restoring the optimizer state. One test pins a defect the sharded path invites, in that a redundancy-sharded optimizer turns each constructor keyword into a param-group default and would therefore record a second copy of the model in every checkpoint of a name-routed optimizer. --- deepmd/dpmodel/train/__init__.py | 4 + deepmd/dpmodel/train/sharding.py | 105 +++++ deepmd/dpmodel/train/trainer.py | 26 +- deepmd/pt/train/training.py | 70 ++- deepmd/pt_expt/train/gradient.py | 192 ++++++++ deepmd/pt_expt/train/training.py | 440 +++++++++++++----- deepmd/pt_expt/train/utils.py | 160 +------ deepmd/pt_expt/train/validation.py | 23 +- deepmd/utils/argcheck.py | 6 +- .../dpmodel/test_train_abstract_trainer.py | 26 ++ .../common/dpmodel/test_train_sharding.py | 52 +++ source/tests/pt/test_validation.py | 9 - ..._train_utils.py => test_train_gradient.py} | 2 +- source/tests/pt_expt/test_training.py | 35 ++ source/tests/pt_expt/test_training_ddp.py | 191 ++++++++ 15 files changed, 1013 insertions(+), 328 deletions(-) create mode 100644 deepmd/dpmodel/train/sharding.py create mode 100644 deepmd/pt_expt/train/gradient.py create mode 100644 source/tests/common/dpmodel/test_train_sharding.py rename source/tests/pt_expt/{test_train_utils.py => test_train_gradient.py} (99%) diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index dc2db6cb4c..9e9c7b1f64 100644 --- a/deepmd/dpmodel/train/__init__.py +++ b/deepmd/dpmodel/train/__init__.py @@ -20,6 +20,9 @@ StepSchedule, resolve_step_schedule, ) +from .sharding import ( + ShardingPolicy, +) from .timing import ( DisplayInterval, TrainingTimer, @@ -45,6 +48,7 @@ "DisplayInterval", "LearningCurveWriter", "RankContext", + "ShardingPolicy", "StepSchedule", "TrainEntrypointOptions", "TrainStepResult", diff --git a/deepmd/dpmodel/train/sharding.py b/deepmd/dpmodel/train/sharding.py new file mode 100644 index 0000000000..2b35438f9c --- /dev/null +++ b/deepmd/dpmodel/train/sharding.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Distribution strategy of a training run. + +``training.zero_stage`` selects how much of the replicated training state is +sharded across ranks, following the ZeRO stages: the optimizer state at stage +one, the gradients as well at stage two, and the parameters on top of that at +stage three. Every consequence of that choice -- which wrapper holds the +model, how the optimizer is built, how a checkpoint is collected, whether a +gradient norm may be reduced locally -- follows from the stage alone. The +stage is therefore resolved into a policy object once, which the backends +then query for the individual decisions. +""" + +from __future__ import ( + annotations, +) + +from dataclasses import ( + dataclass, +) + +__all__ = ["ShardingPolicy"] + +_MAX_STAGE = 3 + + +@dataclass(frozen=True) +class ShardingPolicy: + """What a ZeRO stage implies for the mechanics of a training run. + + Attributes + ---------- + stage : int + The ZeRO stage in effect, between zero and three. A run that is not + distributed is always stage zero: there is nothing to shard across. + """ + + stage: int = 0 + + def __post_init__(self) -> None: + if not 0 <= self.stage <= _MAX_STAGE: + raise ValueError( + f"training.zero_stage must be 0, 1, 2, or 3, got {self.stage}" + ) + + @classmethod + def from_training_params( + cls, + training_params: dict, + *, + is_distributed: bool, + ) -> ShardingPolicy: + """Read the policy from a normalized ``training`` section. + + Parameters + ---------- + training_params : dict + The normalized ``training`` section. + is_distributed : bool + Whether the run spans several ranks. A single-process run cannot + shard anything, so any requested stage is dropped rather than + rejected, which keeps one configuration usable in both settings. + + Returns + ------- + ShardingPolicy + The policy in effect for the run. + """ + # Construct first, so that an out-of-range stage is rejected whether or + # not this run is in a position to honour it. + policy = cls(stage=int(training_params.get("zero_stage", 0))) + return policy if is_distributed else cls() + + @property + def enabled(self) -> bool: + """Whether any training state is sharded.""" + return self.stage > 0 + + @property + def shards_optimizer_state(self) -> bool: + """Whether the optimizer state is split across ranks.""" + return self.stage >= 1 + + @property + def shards_parameters(self) -> bool: + """Whether parameters and gradients live as shards of a whole. + + A sharded parameter is a ``DTensor``, which rules out any reduction + that assumes a rank holds the complete tensor. + """ + return self.stage >= 2 + + @property + def reshards_after_forward(self) -> bool: + """Whether parameters are released again once the forward is done.""" + return self.stage >= 3 + + def describe(self) -> str: + """Return a one-line description of the strategy in effect.""" + if not self.enabled: + return "Distributed data parallel without state sharding." + if self.stage == 1: + return "Enabled DDP + ZeRO Stage-1 Optimizer State Sharding." + stage = "FULL_SHARD (Stage 3)" if self.stage >= 3 else "SHARD_GRAD_OP (Stage 2)" + return f"Enabled FSDP2 {stage}." diff --git a/deepmd/dpmodel/train/trainer.py b/deepmd/dpmodel/train/trainer.py index 62a279eb5f..4377863629 100644 --- a/deepmd/dpmodel/train/trainer.py +++ b/deepmd/dpmodel/train/trainer.py @@ -583,7 +583,7 @@ def run(self, tasks: TrainingTaskCollection) -> None: ) if ( - self.rank_context.is_chief + self._participates_in_checkpoint() and self.trainer_config.save_freq > 0 and display_step % self.trainer_config.save_freq == 0 ): @@ -640,6 +640,17 @@ def on_train_end(self, tasks: TrainingTaskCollection) -> None: """Hook called after training resources have been closed.""" return None + @property + def checkpoint_is_collective(self) -> bool: + """Whether every rank must enter :meth:`save_checkpoint`. + + A backend that shards training state assembles a checkpoint from all + the shards, which is a collective operation: confining the call to the + chief would leave the other ranks waiting at the next barrier. Such a + backend gates the write itself. + """ + return False + def run_full_validation( self, *, @@ -678,7 +689,12 @@ def learning_rate(self, step: int) -> float: @abstractmethod def save_checkpoint(self, step: int) -> None: - """Persist a checkpoint for a one-based step.""" + """Persist a checkpoint for a one-based step. + + Called on the chief alone, unless :attr:`checkpoint_is_collective` + says otherwise; a backend that opts into being called on every rank is + responsible for letting only one of them write to the checkpoint path. + """ def _open_learning_curve(self) -> TextIO | None: if ( @@ -700,8 +716,12 @@ def _should_display(self, display_step: int) -> bool: and display_step % self.trainer_config.disp_freq == 0 ) + def _participates_in_checkpoint(self) -> bool: + """Whether this rank takes part in writing a checkpoint.""" + return self.rank_context.is_chief or self.checkpoint_is_collective + def _should_save_final_checkpoint(self) -> bool: - if not self.rank_context.is_chief: + if not self._participates_in_checkpoint(): return False if self.trainer_config.num_steps <= self.trainer_config.start_step: return False diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index e37e228924..98c083a7d3 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -27,6 +27,7 @@ from deepmd.dpmodel.train import ( DEFAULT_TASK_KEY, CheckpointStore, + ShardingPolicy, TrainingTimer, build_checkpoint_stores, change_model_out_bias, @@ -106,9 +107,11 @@ ModelEMA, get_ema_checkpoint_prefix, ) -from deepmd.pt_expt.train.utils import ( +from deepmd.pt_expt.train.gradient import ( NonFiniteGradGuard, clip_grad_norm_, +) +from deepmd.pt_expt.train.utils import ( count_parameters, infer_env_defaults, resolve_best_checkpoint_dir, @@ -229,18 +232,14 @@ def __init__( self.change_bias_after_training = training_params.get( "change_bias_after_training", False ) - self.zero_stage = int(training_params.get("zero_stage", 0)) - if self.zero_stage not in (0, 1, 2, 3): - raise ValueError( - f"training.zero_stage must be 0, 1, 2, or 3, got {self.zero_stage}" - ) - if self.enable_ema and self.zero_stage >= 2: + self.sharding = ShardingPolicy.from_training_params( + training_params, is_distributed=self.is_distributed + ) + if self.enable_ema and self.sharding.shards_parameters: raise ValueError( "training.enable_ema currently only supports training.zero_stage < 2." ) - if self.zero_stage > 0 and not self.is_distributed: - self.zero_stage = 0 - if self.zero_stage > 0 and self.change_bias_after_training: + if self.sharding.enabled and self.change_bias_after_training: raise ValueError( "training.zero_stage does not support change_bias_after_training." ) @@ -441,11 +440,11 @@ def get_lr(lr_params: dict[str, Any]) -> BaseLR: # Optimizer self.opt_type, self.opt_param = get_opt_param(optimizer_params) - if self.zero_stage > 0 and self.multi_task: + if self.sharding.enabled and self.multi_task: raise ValueError( "training.zero_stage is currently only supported in single-task training." ) - if self.zero_stage > 0 and self.opt_type == "LKF": + if self.sharding.enabled and self.opt_type == "LKF": raise ValueError("training.zero_stage does not support LKF optimizer.") # Loss parameters are also used to select SeZM/DeNS execution modes. @@ -1032,7 +1031,7 @@ def update_finetune_bias( if self.is_distributed: torch.cuda.set_device(LOCAL_RANK) - if self.zero_stage >= 2: + if self.sharding.shards_parameters: if fully_shard is None: raise RuntimeError( "training.zero_stage>=2 requires FSDP2, which is only " @@ -1048,7 +1047,7 @@ def update_finetune_bias( dist.broadcast(p.data, src=0) for b in self.wrapper.buffers(): dist.broadcast(b.data, src=0) - reshard = self.zero_stage >= 3 + reshard = self.sharding.reshards_after_forward self.wrapper = fully_shard(self.wrapper, reshard_after_forward=reshard) else: # zero_stage=0 or 1: standard DDP (ZeRO-1 will wrap the optimizer) @@ -1103,7 +1102,7 @@ def update_finetune_bias( # ops lack DTensor sharding propagation on older PyTorch, so # fall back to the per-tensor path under zero_stage >= 2. # DDP / ZeRO-1 keep plain tensors and use the default. - "use_foreach": False if self.zero_stage >= 2 else None, + "use_foreach": False if self.sharding.shards_parameters else None, } else: raise ValueError(f"Not supported optimizer type '{self.opt_type}'") @@ -1116,7 +1115,7 @@ def update_finetune_bias( ) if self.opt_type == "HybridMuon": target_optimizer = ( - self.optimizer.optim if self.zero_stage == 1 else self.optimizer + self.optimizer.optim if self.sharding.stage == 1 else self.optimizer ) target_optimizer.set_param_names(runtime_named_parameters) self._load_optimizer_state(optimizer_state_dict) @@ -1134,16 +1133,8 @@ def update_finetune_bias( state=ema_state_dict, ) - if self.zero_stage > 0 and self.rank == 0: - if self.zero_stage == 1: - log.info("Enabled DDP + ZeRO Stage-1 Optimizer State Sharding.") - else: - stage = ( - "FULL_SHARD (Stage 3)" - if self.zero_stage >= 3 - else "SHARD_GRAD_OP (Stage 2)" - ) - log.info(f"Enabled FSDP2 {stage}.") + if self.sharding.enabled and self.rank == 0: + log.info(self.sharding.describe()) # Tensorboard self.enable_tensorboard = training_params.get("tensorboard", False) @@ -1240,7 +1231,7 @@ def _create_full_validators( validation_data ), model_ema=self.model_ema, - zero_stage=self.zero_stage, + sharding=self.sharding, ) def _raise_if_full_validation_unsupported( @@ -1266,7 +1257,7 @@ def _raise_if_full_validation_unsupported( "to be configured." ) - if self.zero_stage >= 2: + if self.sharding.shards_parameters: raise ValueError( "validating.full_validation only supports single-task energy " "training with training.zero_stage < 2." @@ -1300,7 +1291,7 @@ def _create_optimizer( torch.optim.Optimizer Constructed optimizer instance. """ - if self.zero_stage == 1: + if self.sharding.stage == 1: return ZeroRedundancyOptimizer( self.wrapper.parameters(), optimizer_class=optimizer_class, @@ -1310,7 +1301,7 @@ def _create_optimizer( def _get_inner_module(self) -> ModelWrapper: """Unwrap DDP if needed. FSDP2 is in-place so no unwrapping required.""" - if self.is_distributed and self.zero_stage <= 1: + if self.is_distributed and not self.sharding.shards_parameters: return self.wrapper.module return self.wrapper @@ -1320,7 +1311,7 @@ def _load_optimizer_state( """Load optimizer state for restart training when available.""" if optimizer_state_dict is None or not self.restart_training: return - if self.zero_stage >= 2: + if self.sharding.shards_parameters: set_optimizer_state_dict( self.wrapper, self.optimizer, @@ -1411,7 +1402,7 @@ def step(_step_id: int, task_key: str = "Default") -> None: # norm. Skip per-param collection in this case to avoid misleading values. if ( self.enable_tensorboard - and self.zero_stage < 2 + and not self.sharding.shards_parameters and ( display_step_id % self.tensorboard_freq == 0 or display_step_id == 1 @@ -1425,7 +1416,7 @@ def step(_step_id: int, task_key: str = "Default") -> None: total_norm = clip_grad_norm_( self.wrapper.parameters(), self.gradient_max_norm, - stable=self.zero_stage < 2, + stable=not self.sharding.shards_parameters, ) self.nonfinite_grad_guard.update(total_norm) with torch.device(DEVICE): @@ -1652,7 +1643,8 @@ def log_loss_valid(_task_key: str = "Default") -> dict: is_train=True, task_key=_key ) if input_dict and not ( - self.is_distributed and self.zero_stage >= 2 + self.is_distributed + and self.sharding.shards_parameters ): _, loss, more_loss = self._get_inner_module()( **input_dict, @@ -1756,7 +1748,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict: self.wrapper.named_parameters ) if should_save_checkpoint and ( - self.zero_stage > 0 or self.rank == 0 or dist.get_rank() == 0 + self.sharding.enabled or self.rank == 0 or dist.get_rank() == 0 ): # Handle the case if rank 0 aborted and re-assigned self.latest_model = self.ckpt_store.path_for(display_step_id) @@ -1871,7 +1863,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict: ) self.ema_ckpt_store.publish(self.latest_ema_model) - if self.num_steps == 0 and self.zero_stage > 0: + if self.num_steps == 0 and self.sharding.enabled: # ZeRO-1 / FSDP: all ranks participate in save_model (collective op) self.latest_model = self.ckpt_store.path_for(0) self.save_model(self.latest_model, lr=0, step=0) @@ -1883,7 +1875,7 @@ def log_loss_valid(_task_key: str = "Default") -> dict: self.rank == 0 or dist.get_rank() == 0 ): # Handle the case if rank 0 aborted and re-assigned if self.num_steps == 0: - if self.zero_stage == 0: + if not self.sharding.enabled: # When num_steps is 0, the checkpoint is never saved in the loop self.latest_model = self.ckpt_store.path_for(0) self.save_model(self.latest_model, lr=0, step=0) @@ -1978,7 +1970,7 @@ def _collect_checkpoint_states( else nullcontext() ) with ema_context: - if self.zero_stage >= 2: + if self.sharding.shards_parameters: # FSDP2: collective op, all ranks participate; rank 0 gets full state options = StateDictOptions(full_state_dict=True, cpu_offload=True) model_state = get_model_state_dict(self.wrapper, options=options) @@ -1989,7 +1981,7 @@ def _collect_checkpoint_states( if include_optimizer else None ) - elif self.zero_stage == 1: + elif self.sharding.stage == 1: # ZeRO-1: consolidate sharded optimizer state to rank 0. model_state = module.state_dict() if use_ema_weights: diff --git a/deepmd/pt_expt/train/gradient.py b/deepmd/pt_expt/train/gradient.py new file mode 100644 index 0000000000..cc3b863015 --- /dev/null +++ b/deepmd/pt_expt/train/gradient.py @@ -0,0 +1,192 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Numerical safeguards around the gradient of a training step. + +Two hazards are handled here, both of which have to be addressed without +slowing the step down. + +A gradient that is large but still representable can overflow the sum of +squares of a naive L2 reduction, so a finite gradient would be reported as +infinite and clipped as if training had diverged. Scaling the gradients by +their largest magnitude before the reduction removes that failure mode. + +A gradient that is genuinely non-finite must stop the run before it reaches +a checkpoint, but reading that condition back on every step would serialize +the host against the device. The condition is therefore accumulated on +device and read only at a checkpoint boundary. +""" + +from __future__ import ( + annotations, +) + +from typing import ( + TYPE_CHECKING, +) + +import torch + +if TYPE_CHECKING: + from collections.abc import ( + Callable, + Iterable, + ) + +__all__ = [ + "NonFiniteGradGuard", + "clip_grad_norm_", + "raise_nonfinite_gradient_norm", +] + + +def clip_grad_norm_( + parameters: Iterable[torch.nn.Parameter], + max_norm: float, + stable: bool = True, +) -> torch.Tensor: + """ + Clip gradients in place so their global L2 norm does not exceed ``max_norm``. + + The norm is computed and applied on device and returned without a host + synchronization. Non-finite gradients are not reported here; see + :class:`NonFiniteGradGuard`. + + Parameters + ---------- + parameters : Iterable[torch.nn.Parameter] + Parameters whose gradients are clipped in place. + max_norm : float + Maximum allowed global L2 norm of the gradients. + stable : bool, optional + Norm reduction strategy. ``True`` scales the gradients by their largest + magnitude before the float64 reduction, keeping the norm finite for an + arbitrarily large but finite gradient of any dtype. ``False`` uses the + native reduction, which propagates ``DTensor`` sharding under FSDP2 but is + not overflow-safe. + + Returns + ------- + torch.Tensor + The global gradient norm before clipping, or a CPU zero when no gradient + is present. + """ + params = [p for p in parameters if p.grad is not None] + if not params: + return torch.zeros((), dtype=torch.float64, device="cpu") + grads = [p.grad for p in params] + + # === Step 1. Global L2 norm === + if stable: + # Normalize by the largest magnitude before the float64 reduction; the + # factor cancels in the product and keeps the sum of squares finite for + # any dtype. The ``foreach`` ops fuse the per-parameter passes. + scale = torch.stack(torch._foreach_norm(grads, float("inf"))).max() + scale = torch.where(scale > 0, scale, scale.new_ones(())) + scaled = torch._foreach_norm(torch._foreach_div(grads, scale), 2.0) + total_norm = scale.double() * torch.linalg.vector_norm( + torch.stack(scaled).double() + ) + else: + total_norm = torch.nn.utils.get_total_norm(grads, error_if_nonfinite=False) + + # === Step 2. Rescale gradients by the clamped coefficient === + torch.nn.utils.clip_grads_with_norm_(params, max_norm, total_norm) + return total_norm + + +class NonFiniteGradGuard: + """ + Detect non-finite gradient norms without a per-step host synchronization. + + :meth:`update` accumulates the non-finite condition on device; the result is + read back only by :meth:`raise_if_nonfinite`. Calling the check before each + checkpoint keeps a diverged interval from being written while leaving the + training step free of host reads. + """ + + def __init__(self) -> None: + self._nonfinite: torch.Tensor | None = None + + def update(self, total_norm: torch.Tensor) -> None: + """ + Accumulate whether ``total_norm`` is non-finite. + + Parameters + ---------- + total_norm : torch.Tensor + The gradient norm returned by :func:`clip_grad_norm_`. + """ + nonfinite = ~torch.isfinite(total_norm) + if self._nonfinite is not None: + nonfinite |= self._nonfinite.to(nonfinite.device) + self._nonfinite = nonfinite + + def raise_if_nonfinite( + self, + named_parameters: Callable[[], Iterable[tuple[str, torch.nn.Parameter]]], + ) -> None: + """ + Raise if any norm accumulated since the previous call was non-finite. + + On failure the current gradient state is reported via + :func:`raise_nonfinite_gradient_norm`. + + Parameters + ---------- + named_parameters : Callable[[], Iterable[tuple[str, torch.nn.Parameter]]] + Accessor for the model's named parameters, consulted only when raising. + + Raises + ------ + RuntimeError + If a non-finite gradient norm was recorded. + """ + if self._nonfinite is None: + return + diverged = bool(self._nonfinite) + self._nonfinite = None + if diverged: + raise_nonfinite_gradient_norm(named_parameters()) + + +def raise_nonfinite_gradient_norm( + named_parameters: Iterable[tuple[str, torch.nn.Parameter]], +) -> None: + """ + Raise a ``RuntimeError`` reporting the current non-finite gradients. + + Parameters whose current gradient is non-finite are listed by name and shape. + When every current individual gradient is finite, the accumulated guard flag + came from an earlier step in the checkpoint interval or from the norm + reduction. + + Parameters + ---------- + named_parameters : Iterable[tuple[str, torch.nn.Parameter]] + The model's named parameters. + + Raises + ------ + RuntimeError + Always; this is the divergence-reporting path. + """ + bad_params = [] + for name, param in named_parameters: + if param.grad is None: + continue + grad_norm = param.grad.detach().norm() + if not torch.isfinite(grad_norm): + bad_params.append( + f" {name}: grad_norm={grad_norm}, shape={list(param.shape)}" + ) + detail = ( + "\n".join(bad_params) + if bad_params + else ( + " (all current individual gradients are finite; non-finite norm was " + "recorded earlier in the checkpoint interval or in the norm reduction)" + ) + ) + raise RuntimeError( + "Non-finite gradient norm; training has diverged.\n" + f"Parameters with non-finite gradients:\n{detail}" + ) diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 6625a1e09c..c5a98c34aa 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -26,11 +26,28 @@ import numpy as np import torch import torch.distributed as dist +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + set_optimizer_state_dict, +) +from torch.distributed.optim import ( + ZeroRedundancyOptimizer, +) + +try: + from torch.distributed.fsdp import ( + fully_shard, + ) +except ImportError: + fully_shard = None # type: ignore[assignment] from deepmd.dpmodel.train import ( DEFAULT_TASK_KEY, AbstractTrainer, RankContext, + ShardingPolicy, TrainerConfig, TrainingTask, TrainingTaskCollection, @@ -85,6 +102,10 @@ ModelEMA, get_ema_checkpoint_prefix, ) +from deepmd.pt_expt.train.gradient import ( + NonFiniteGradGuard, + clip_grad_norm_, +) from deepmd.pt_expt.train.utils import ( count_parameters, infer_env_defaults, @@ -1436,6 +1457,9 @@ 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 + self.sharding = ShardingPolicy.from_training_params( + training_params, is_distributed=self.is_distributed + ) # Iteration config self.disp_file = training_params.get("disp_file", "lcurve.out") @@ -1450,6 +1474,8 @@ def __init__( self.change_bias_after_training = bool( training_params.get("change_bias_after_training", False) ) + self.enable_compile = bool(training_params.get("enable_compile", False)) + self._raise_if_sharding_unsupported() # Model --------------------------------------------------------------- self.models: dict[str, torch.nn.Module] = {} @@ -1597,6 +1623,7 @@ def initialize_statistics( # Gradient clipping self.gradient_max_norm = training_params.get("gradient_max_norm", 0.0) + self.nonfinite_grad_guard = NonFiniteGradGuard() # Model wrapper ------------------------------------------------------- self.wrapper = ModelWrapper(self.model, self.loss, model_params=model_params) @@ -1643,88 +1670,12 @@ def initialize_statistics( **share_kwargs, ) - # DDP wrapping -------------------------------------------------------- - if self.is_distributed: - # Multi-task uses only one fitting_net per step, so unused - # parameters exist in the graph. Single-task doesn't need this. - _find_unused = self.multi_task - if DEVICE.type == "cuda": - from deepmd.pt_expt.utils.env import ( - LOCAL_RANK, - ) - - torch.cuda.set_device(LOCAL_RANK) - self.wrapper = torch.nn.parallel.DistributedDataParallel( - self.wrapper, - device_ids=[LOCAL_RANK], - find_unused_parameters=_find_unused, - output_device=LOCAL_RANK, - ) - else: - # CPU (gloo backend) — no device_ids - self.wrapper = torch.nn.parallel.DistributedDataParallel( - self.wrapper, - find_unused_parameters=_find_unused, - ) - - # Optimiser ----------------------------------------------------------- - opt_type = optimizer_params.get("type", "Adam") - if opt_type not in {"Adam", "AdamW", "HybridMuon"}: - raise ValueError(f"Unsupported optimizer type: {opt_type}") - - # LambdaLR multiplies each param group's initial learning rate by the - # lambda value. Warmup schedules legitimately return zero at step 0, - # so use the nonzero schedule base as the denominator and let the - # lambda initialize the optimizer to the requested warmup value. - initial_lr = float(self.lr_schedule.start_lr) - adam_betas = ( - float(optimizer_params["adam_beta1"]), - float(optimizer_params["adam_beta2"]), - ) - weight_decay = float(optimizer_params["weight_decay"]) - - if opt_type == "Adam": - self.optimizer = torch.optim.Adam( - self.wrapper.parameters(), - lr=initial_lr, - betas=adam_betas, - weight_decay=weight_decay, - ) - elif opt_type == "AdamW": - self.optimizer = torch.optim.AdamW( - self.wrapper.parameters(), - lr=initial_lr, - betas=adam_betas, - weight_decay=weight_decay, - ) - else: # HybridMuon - runtime_named_parameters = tuple(self.wrapper.named_parameters()) - self.optimizer = HybridMuonOptimizer( - self.wrapper.parameters(), - lr=initial_lr, - momentum=float(optimizer_params["momentum"]), - weight_decay=weight_decay, - adam_betas=adam_betas, - lr_adjust=float(optimizer_params["lr_adjust"]), - lr_adjust_coeff=float(optimizer_params["lr_adjust_coeff"]), - muon_mode=str(optimizer_params["muon_mode"]), - named_parameters=runtime_named_parameters, - enable_gram=bool(optimizer_params["enable_gram"]), - flash_muon=bool(optimizer_params["flash_muon"]), - magma_muon=bool(optimizer_params["magma_muon"]), - ) - - for param_group in self.optimizer.param_groups: - param_group["initial_lr"] = initial_lr - - self.scheduler = torch.optim.lr_scheduler.LambdaLR( - self.optimizer, - lambda step: self.lr_schedule.value(step) / initial_lr, - last_epoch=self.start_step - 1, - ) - # Resume -------------------------------------------------------------- + # Weights are restored while the wrapper still owns whole tensors, + # because a checkpoint records the model as a whole and its tensors + # cannot be copied into sharded parameters. ema_state_dict = None + optimizer_state_dict = None if resuming: log.info(f"Resuming from {resume_model}.") is_pte = resume_model.endswith((".pte", ".pt2")) @@ -1914,18 +1865,72 @@ def update_finetune_bias( ), ) - if optimizer_state_dict is not None: - self.optimizer.load_state_dict(optimizer_state_dict) - for param_group in self.optimizer.param_groups: - param_group["initial_lr"] = initial_lr - # rebuild scheduler from the resumed step. - # last_epoch handles the step offset; the lambda must NOT - # add self.start_step again (that would double-count). - self.scheduler = torch.optim.lr_scheduler.LambdaLR( - self.optimizer, - lambda step: self.lr_schedule.value(step) / initial_lr, - last_epoch=self.start_step - 1, - ) + # Distribution -------------------------------------------------------- + # The weights are in place, so a sharding strategy may cut them up. + if self.is_distributed: + self._distribute_wrapper() + self._log_sharding_strategy() + + # Optimiser ----------------------------------------------------------- + opt_type = optimizer_params.get("type", "Adam") + # LambdaLR multiplies each param group's initial learning rate by the + # lambda value. Warmup schedules legitimately return zero at step 0, + # so use the nonzero schedule base as the denominator and let the + # lambda initialize the optimizer to the requested warmup value. + initial_lr = float(self.lr_schedule.start_lr) + adam_betas = ( + float(optimizer_params["adam_beta1"]), + float(optimizer_params["adam_beta2"]), + ) + weight_decay = float(optimizer_params["weight_decay"]) + + if opt_type in ("Adam", "AdamW"): + self.optimizer = self._create_optimizer( + torch.optim.Adam if opt_type == "Adam" else torch.optim.AdamW, + lr=initial_lr, + betas=adam_betas, + weight_decay=weight_decay, + ) + elif opt_type == "HybridMuon": + self.optimizer = self._create_optimizer( + HybridMuonOptimizer, + lr=initial_lr, + momentum=float(optimizer_params["momentum"]), + weight_decay=weight_decay, + adam_betas=adam_betas, + lr_adjust=float(optimizer_params["lr_adjust"]), + lr_adjust_coeff=float(optimizer_params["lr_adjust_coeff"]), + muon_mode=str(optimizer_params["muon_mode"]), + enable_gram=bool(optimizer_params["enable_gram"]), + flash_muon=bool(optimizer_params["flash_muon"]), + magma_muon=bool(optimizer_params["magma_muon"]), + # Sharded parameters are DTensors, and several torch._foreach_* + # ops lack sharding propagation, so the per-tensor path applies. + use_foreach=False if self.sharding.shards_parameters else None, + ) + # The parameter names route each tensor to Muon or Adam. They are + # supplied after construction because a redundancy-sharded + # optimizer treats every constructor keyword as a param-group + # default, which would serialize the whole model into each + # checkpoint. + self._local_optimizer.set_param_names( + tuple(self.wrapper.named_parameters()) + ) + else: + raise ValueError(f"Unsupported optimizer type: {opt_type}") + + if optimizer_state_dict is not None: + self._load_optimizer_state(optimizer_state_dict) + for param_group in self.optimizer.param_groups: + param_group["initial_lr"] = initial_lr + + # The resumed step offset is carried by last_epoch; the lambda must not + # add it again, which would advance the schedule twice. + self.scheduler = torch.optim.lr_scheduler.LambdaLR( + self.optimizer, + lambda step: self.lr_schedule.value(step) / initial_lr, + last_epoch=self.start_step - 1, + ) # Exponential moving average ------------------------------------------- # The shadow tracks the raw models, whose parameter tensors the compiled @@ -1941,7 +1946,6 @@ def update_finetune_bias( ) # torch.compile ------------------------------------------------------- - self.enable_compile = training_params.get("enable_compile", False) if self.enable_compile: check_compile_torch_version() compile_opts = training_params.get("compile_options", {}) @@ -1992,6 +1996,7 @@ def _create_full_validators( validation_data ), model_ema=self.model_ema, + sharding=self.sharding, ) def _raise_if_full_validation_unsupported( @@ -2005,6 +2010,12 @@ def _raise_if_full_validation_unsupported( "training; multi-task training is not supported." ) + if self.sharding.shards_parameters: + raise ValueError( + "validating.full_validation only supports single-task energy " + "training with training.zero_stage < 2." + ) + if self.models[DEFAULT_TASK_KEY].has_spin() or isinstance( self.loss, EnergySpinLoss ): @@ -2262,9 +2273,154 @@ def _epoch_length(self, model_key: str) -> int: return int(np.ceil(total / self.world_size)) # ------------------------------------------------------------------ - # DDP helpers + # Distribution helpers # ------------------------------------------------------------------ + def _raise_if_sharding_unsupported(self) -> None: + """Reject the run configurations that state sharding cannot serve. + + Raises + ------ + ValueError + If the requested stage conflicts with another training option. + """ + if not self.sharding.enabled: + return + if self.multi_task: + raise ValueError( + "training.zero_stage is currently only supported in single-task " + "training." + ) + if self.change_bias_after_training: + raise ValueError( + "training.zero_stage does not support change_bias_after_training." + ) + if not self.sharding.shards_parameters: + return + if self.enable_ema: + raise ValueError( + "training.enable_ema currently only supports training.zero_stage < 2." + ) + if self.enable_compile: + raise ValueError( + "training.enable_compile only supports training.zero_stage < 2: " + "the compiled graph is traced from the parameters, which FSDP2 " + "shards as DTensors." + ) + + def _distribute_wrapper(self) -> None: + """Place the wrapper under the parallel strategy of this run. + + Stages below two replicate the model and keep plain tensors, so the + wrapper is held by ``DistributedDataParallel``. From stage two on the + parameters themselves are sharded by FSDP2, which mutates the wrapper + in place and therefore leaves no module to unwrap. + """ + local_rank = None + if DEVICE.type == "cuda": + from deepmd.pt_expt.utils.env import ( + LOCAL_RANK, + ) + + local_rank = LOCAL_RANK + torch.cuda.set_device(local_rank) + + if not self.sharding.shards_parameters: + # Multi-task uses only one fitting_net per step, so unused + # parameters exist in the graph. Single-task doesn't need this. + kwargs: dict[str, Any] = {"find_unused_parameters": self.multi_task} + if local_rank is not None: + kwargs |= {"device_ids": [local_rank], "output_device": local_rank} + self.wrapper = torch.nn.parallel.DistributedDataParallel( + self.wrapper, **kwargs + ) + return + + if fully_shard is None: + raise RuntimeError( + "training.zero_stage>=2 requires FSDP2 " + "(``torch.distributed.fsdp.fully_shard``), which is missing " + f"from PyTorch {torch.__version__}. Set training.zero_stage " + "to 0 or 1 to stay on the DDP / ZeRO-1 path." + ) + # Unlike the DDP constructor, FSDP2 does not broadcast: the ranks have + # to already agree on the weights before they are cut into shards. + for tensor in (*self.wrapper.parameters(), *self.wrapper.buffers()): + dist.broadcast(tensor.data, src=0) + self.wrapper = fully_shard( + self.wrapper, + reshard_after_forward=self.sharding.reshards_after_forward, + ) + + def _log_sharding_strategy(self) -> None: + """Report the distribution strategy once the wrapper is in place.""" + if self.sharding.enabled and self.rank == 0: + log.info(self.sharding.describe()) + + def _load_optimizer_state(self, optimizer_state_dict: dict[str, Any]) -> None: + """Restore optimizer state recorded as one whole. + + Unlike the weights, the optimizer is necessarily built after the model + is distributed, so under parameter sharding its state is already made + of shards and the recorded state has to be cut up to match. + + Parameters + ---------- + optimizer_state_dict : dict[str, Any] + The optimizer state as recorded in a checkpoint. + """ + if not self.sharding.shards_parameters: + self.optimizer.load_state_dict(optimizer_state_dict) + return + set_optimizer_state_dict( + self.wrapper, + self.optimizer, + optim_state_dict=optimizer_state_dict, + options=StateDictOptions(full_state_dict=True, broadcast_from_rank0=True), + ) + + def _create_optimizer( + self, + optimizer_class: type[torch.optim.Optimizer], + **kwargs: Any, + ) -> torch.optim.Optimizer: + """Construct the optimizer, sharding its state when the stage asks. + + Parameters + ---------- + optimizer_class : type[torch.optim.Optimizer] + The optimizer to construct. + **kwargs + Keyword arguments forwarded to the optimizer. + + Returns + ------- + torch.optim.Optimizer + The optimizer, wrapped in a ``ZeroRedundancyOptimizer`` when the + stage shards optimizer state but not the parameters; from stage + two on FSDP2 already shards the state that the optimizer derives + from its sharded parameters. + """ + if self.sharding.stage == 1: + return ZeroRedundancyOptimizer( + self.wrapper.parameters(), + optimizer_class=optimizer_class, + **kwargs, + ) + return optimizer_class(self.wrapper.parameters(), **kwargs) + + @property + def _local_optimizer(self) -> torch.optim.Optimizer: + """Return the optimizer that performs this rank's update. + + A redundancy-sharded optimizer owns no update of its own: it delegates + to a local optimizer over this rank's share of the parameters, and it + is that one which holds the per-parameter state and the name routing. + """ + if self.sharding.stage == 1: + return self.optimizer.optim + return self.optimizer + @property def _unwrapped(self) -> "ModelWrapper": """Return the raw ModelWrapper, unwrapping DDP if active.""" @@ -2325,17 +2481,27 @@ def _broadcast_value_from_rank0(self, value: Any) -> Any: # Checkpointing # ------------------------------------------------------------------ + @property + def checkpoint_is_collective(self) -> bool: + """Whether assembling a checkpoint needs every rank.""" + return self.sharding.enabled + def save_checkpoint(self, step: int) -> None: + # Abort before writing if any gradient norm since the previous + # checkpoint was non-finite, so a diverged interval is not persisted. + self.nonfinite_grad_guard.raise_if_nonfinite(self.wrapper.named_parameters) ckpt_path = self.ckpt_store.path_for(step) self._save_checkpoint_to_path(ckpt_path, step=step) - self.ckpt_store.publish(ckpt_path) - self.ckpt_store.prune(ckpt_path) - log.info(f"Saved model to {ckpt_path}") + if self.rank == 0: + self.ckpt_store.publish(ckpt_path) + self.ckpt_store.prune(ckpt_path) + log.info(f"Saved model to {ckpt_path}") if self.model_ema is not None: ema_path = self.ema_ckpt_store.path_for(step) self._save_checkpoint_to_path(ema_path, step=step, use_ema_weights=True) - self.ema_ckpt_store.publish(ema_path) - self.ema_ckpt_store.prune(ema_path) + if self.rank == 0: + self.ema_ckpt_store.publish(ema_path) + self.ema_ckpt_store.prune(ema_path) def _save_full_validation_checkpoint( self, @@ -2416,17 +2582,65 @@ def _write_checkpoint( compiled_backup[task_key] = m wrapper.model[task_key] = m.original_model try: - state: dict[str, Any] = {"model": wrapper.state_dict()} - if include_optimizer: - state["optimizer"] = self.optimizer.state_dict() + model_state, optim_state = self._collect_checkpoint_states( + wrapper, include_optimizer=include_optimizer + ) finally: for task_key, compiled in compiled_backup.items(): wrapper.model[task_key] = compiled + # 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: + return + state: dict[str, Any] = {"model": model_state} + if optim_state is not None: + state["optimizer"] = optim_state 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) torch.save(state, ckpt_path) + def _collect_checkpoint_states( + self, + wrapper: "ModelWrapper", + *, + include_optimizer: bool, + ) -> tuple[dict[str, Any], dict[str, Any] | None]: + """Gather the model and optimizer state a checkpoint records. + + Parameters + ---------- + wrapper : ModelWrapper + The unwrapped model wrapper, already stripped of compiled models. + Under parameter sharding it is the sharded wrapper itself, because + FSDP2 shards in place and compilation is rejected alongside it. + include_optimizer : bool + Whether the optimizer state belongs in the checkpoint. + + Returns + ------- + tuple[dict[str, Any], dict[str, Any] | None] + The model state and the optimizer state. Under sharding both are + complete only on the chief; the other ranks contribute their shards + and receive placeholders. + """ + if self.sharding.shards_parameters: + # FSDP2 reassembles the shards, so every rank has to take part. + options = StateDictOptions(full_state_dict=True, cpu_offload=True) + return ( + get_model_state_dict(wrapper, options=options), + get_optimizer_state_dict(wrapper, self.optimizer, options=options) + if include_optimizer + else None, + ) + model_state = wrapper.state_dict() + if not include_optimizer: + return model_state, None + if self.sharding.shards_optimizer_state: + self.optimizer.consolidate_state_dict(to=0) + return model_state, self.optimizer.state_dict() if self.rank == 0 else {} + return model_state, self.optimizer.state_dict() + # ------------------------------------------------------------------ # Training loop # ------------------------------------------------------------------ @@ -2573,8 +2787,15 @@ def train_step(self, task: TrainingTask, step: int) -> TrainStepResult: loss.backward() if self.gradient_max_norm > 0.0: - torch.nn.utils.clip_grad_norm_( - self.wrapper.parameters(), self.gradient_max_norm + self.nonfinite_grad_guard.update( + clip_grad_norm_( + self.wrapper.parameters(), + self.gradient_max_norm, + # A sharded gradient is a DTensor: the overflow-safe + # reduction would measure this rank's shard alone, so the + # distributed-native norm applies instead. + stable=not self.sharding.shards_parameters, + ) ) self._optimizer_step() @@ -2616,8 +2837,15 @@ def evaluate_validation( step: int, step_result: TrainStepResult | None, ) -> dict[str, float] | None: - """Evaluate validation loss terms for one task.""" - if task.validation_data is None: + """Evaluate validation loss terms for one task. + + Sharded parameters are gathered by the forward itself, which makes it + a collective operation, while the display runs on the chief alone. + Validation is therefore skipped from stage two on; the metrics remain + available through the independent full validation flow, which every + rank enters together. + """ + if task.validation_data is None or self.sharding.shards_parameters: return None valid_results: dict[str, float] = {} diff --git a/deepmd/pt_expt/train/utils.py b/deepmd/pt_expt/train/utils.py index 4dde1b9579..5786ff10d0 100644 --- a/deepmd/pt_expt/train/utils.py +++ b/deepmd/pt_expt/train/utils.py @@ -17,15 +17,13 @@ Any, ) -import torch - if TYPE_CHECKING: from collections.abc import ( - Callable, Generator, - Iterable, ) + import torch + def count_parameters(module: torch.nn.Module) -> tuple[int, int]: """ @@ -46,160 +44,6 @@ def count_parameters(module: torch.nn.Module) -> tuple[int, int]: return trainable, total -def clip_grad_norm_( - parameters: Iterable[torch.nn.Parameter], - max_norm: float, - stable: bool = True, -) -> torch.Tensor: - """ - Clip gradients in place so their global L2 norm does not exceed ``max_norm``. - - The norm is computed and applied on device and returned without a host - synchronization. Non-finite gradients are not reported here; see - :class:`NonFiniteGradGuard`. - - Parameters - ---------- - parameters : Iterable[torch.nn.Parameter] - Parameters whose gradients are clipped in place. - max_norm : float - Maximum allowed global L2 norm of the gradients. - stable : bool, optional - Norm reduction strategy. ``True`` scales the gradients by their largest - magnitude before the float64 reduction, keeping the norm finite for an - arbitrarily large but finite gradient of any dtype. ``False`` uses the - native reduction, which propagates ``DTensor`` sharding under FSDP2 but is - not overflow-safe. - - Returns - ------- - torch.Tensor - The global gradient norm before clipping, or a CPU zero when no gradient - is present. - """ - params = [p for p in parameters if p.grad is not None] - if not params: - return torch.zeros((), dtype=torch.float64, device="cpu") - grads = [p.grad for p in params] - - # === Step 1. Global L2 norm === - if stable: - # Normalize by the largest magnitude before the float64 reduction; the - # factor cancels in the product and keeps the sum of squares finite for - # any dtype. The ``foreach`` ops fuse the per-parameter passes. - scale = torch.stack(torch._foreach_norm(grads, float("inf"))).max() - scale = torch.where(scale > 0, scale, scale.new_ones(())) - scaled = torch._foreach_norm(torch._foreach_div(grads, scale), 2.0) - total_norm = scale.double() * torch.linalg.vector_norm( - torch.stack(scaled).double() - ) - else: - total_norm = torch.nn.utils.get_total_norm(grads, error_if_nonfinite=False) - - # === Step 2. Rescale gradients by the clamped coefficient === - torch.nn.utils.clip_grads_with_norm_(params, max_norm, total_norm) - return total_norm - - -class NonFiniteGradGuard: - """ - Detect non-finite gradient norms without a per-step host synchronization. - - :meth:`update` accumulates the non-finite condition on device; the result is - read back only by :meth:`raise_if_nonfinite`. Calling the check before each - checkpoint keeps a diverged interval from being written while leaving the - training step free of host reads. - """ - - def __init__(self) -> None: - self._nonfinite: torch.Tensor | None = None - - def update(self, total_norm: torch.Tensor) -> None: - """ - Accumulate whether ``total_norm`` is non-finite. - - Parameters - ---------- - total_norm : torch.Tensor - The gradient norm returned by :func:`clip_grad_norm_`. - """ - nonfinite = ~torch.isfinite(total_norm) - if self._nonfinite is not None: - nonfinite |= self._nonfinite.to(nonfinite.device) - self._nonfinite = nonfinite - - def raise_if_nonfinite( - self, - named_parameters: Callable[[], Iterable[tuple[str, torch.nn.Parameter]]], - ) -> None: - """ - Raise if any norm accumulated since the previous call was non-finite. - - On failure the current gradient state is reported via - :func:`raise_nonfinite_gradient_norm`. - - Parameters - ---------- - named_parameters : Callable[[], Iterable[tuple[str, torch.nn.Parameter]]] - Accessor for the model's named parameters, consulted only when raising. - - Raises - ------ - RuntimeError - If a non-finite gradient norm was recorded. - """ - if self._nonfinite is None: - return - diverged = bool(self._nonfinite) - self._nonfinite = None - if diverged: - raise_nonfinite_gradient_norm(named_parameters()) - - -def raise_nonfinite_gradient_norm( - named_parameters: Iterable[tuple[str, torch.nn.Parameter]], -) -> None: - """ - Raise a ``RuntimeError`` reporting the current non-finite gradients. - - Parameters whose current gradient is non-finite are listed by name and shape. - When every current individual gradient is finite, the accumulated guard flag - came from an earlier step in the checkpoint interval or from the norm - reduction. - - Parameters - ---------- - named_parameters : Iterable[tuple[str, torch.nn.Parameter]] - The model's named parameters. - - Raises - ------ - RuntimeError - Always; this is the divergence-reporting path. - """ - bad_params = [] - for name, param in named_parameters: - if param.grad is None: - continue - grad_norm = param.grad.detach().norm() - if not torch.isfinite(grad_norm): - bad_params.append( - f" {name}: grad_norm={grad_norm}, shape={list(param.shape)}" - ) - detail = ( - "\n".join(bad_params) - if bad_params - else ( - " (all current individual gradients are finite; non-finite norm was " - "recorded earlier in the checkpoint interval or in the norm reduction)" - ) - ) - raise RuntimeError( - "Non-finite gradient norm; training has diverged.\n" - f"Parameters with non-finite gradients:\n{detail}" - ) - - def infer_env_defaults(validating_params: dict[str, Any]) -> dict[str, str]: """ Translate the eval-time policy options into environment defaults. diff --git a/deepmd/pt_expt/train/validation.py b/deepmd/pt_expt/train/validation.py index a3a822838c..edd89e6d04 100644 --- a/deepmd/pt_expt/train/validation.py +++ b/deepmd/pt_expt/train/validation.py @@ -28,6 +28,9 @@ import torch.distributed as dist from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT +from deepmd.dpmodel.train import ( + ShardingPolicy, +) from deepmd.dpmodel.utils.lmdb_data import ( LmdbTestData, LmdbTestDataNlocView, @@ -202,8 +205,8 @@ def __init__( state_store: dict[str, Any], num_steps: int, rank: int, - zero_stage: int, restart_training: bool, + sharding: ShardingPolicy = ShardingPolicy(), checkpoint_dir: Path | None = None, full_val_file: str | Path | None = None, best_checkpoint_prefix: str = BEST_CKPT_PREFIX, @@ -219,7 +222,7 @@ def __init__( self.profile = select_metric_profile(model) self.state_store = state_store self.rank = rank - self.zero_stage = zero_stage + self.sharding = sharding self.checkpoint_dir = ( Path(checkpoint_dir) if checkpoint_dir is not None else Path(".") ) @@ -333,9 +336,9 @@ def run( if save_path[0] is not None: try: - # ZeRO/FSDP checkpoint collection is collective, so all ranks must - # enter `save_checkpoint` whenever `zero_stage > 0`. - if (self.is_distributed and self.zero_stage != 0) or self.rank == 0: + # Assembling a checkpoint from shards is collective, so every + # rank enters it once any training state is sharded. + if (self.is_distributed and self.sharding.enabled) or self.rank == 0: save_checkpoint(Path(save_path[0]), lr=lr, step=step_id) if self.rank == 0: self._reconcile_best_checkpoints() @@ -922,7 +925,7 @@ def build_full_validators( checkpoint_dir: Path, ensure_supported: Callable[[], None], model_ema: Any | None = None, - zero_stage: int = 0, + sharding: ShardingPolicy | None = None, ) -> tuple[FullValidator | None, FullValidator | None]: """Build the full validators of a training run. @@ -961,9 +964,9 @@ def build_full_validators( The EMA state of the run. Without it the EMA flow stays inactive, so that ``ema_full_validation`` is ignored rather than rejected when EMA itself is disabled. - zero_stage : int, optional - The ZeRO stage of the run, which decides whether checkpoint collection - is a collective operation. + sharding : ShardingPolicy, optional + The distribution strategy of the run, which decides whether checkpoint + collection is a collective operation. Defaults to no sharding. Returns ------- @@ -994,7 +997,7 @@ def make(**overrides: Any) -> FullValidator: model=model, num_steps=num_steps, rank=rank, - zero_stage=zero_stage, + sharding=ShardingPolicy() if sharding is None else sharding, restart_training=restart_training, checkpoint_dir=checkpoint_dir, **overrides, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index a9d57435b4..196f6de1d0 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5508,7 +5508,9 @@ def training_args( "50% more communication (3x model size) due to parameter all-gather in " "both forward and backward passes. " "Default is 0. Requires distributed launch via torchrun. " - "Currently supports single-task training; does not support LKF or change_bias_after_training." + "Currently supports single-task training; does not support LKF or change_bias_after_training. " + "In the PyTorch Exportable backend, stages 2 and 3 additionally exclude " + "`enable_compile`, whose traced graph cannot carry sharded parameters." ) doc_neighbor_graph_method = ( "Select the carry-all neighbor-graph builder for graph-eligible PyTorch " @@ -5712,7 +5714,7 @@ def training_args( int, optional=True, default=0, - doc=supported_backends("pt") + doc_zero_stage, + doc=supported_backends("pt", "pt_expt") + doc_zero_stage, ), Argument( "neighbor_graph_method", diff --git a/source/tests/common/dpmodel/test_train_abstract_trainer.py b/source/tests/common/dpmodel/test_train_abstract_trainer.py index bb450d67b3..23eaeb1a22 100644 --- a/source/tests/common/dpmodel/test_train_abstract_trainer.py +++ b/source/tests/common/dpmodel/test_train_abstract_trainer.py @@ -232,6 +232,32 @@ def test_non_chief_rank_skips_user_visible_outputs(tmp_path: Path) -> None: assert not lcurve.exists() +def test_collective_checkpointing_reaches_every_rank(tmp_path: Path) -> None: + class CollectiveTrainer(DummyTrainer): + @property + def checkpoint_is_collective(self) -> bool: + return True + + lcurve = tmp_path / "lcurve.out" + trainer = CollectiveTrainer( + TrainerConfig( + num_steps=2, + disp_file=str(lcurve), + disp_freq=1, + save_freq=1, + timing_in_training=False, + ), + rank_context=RankContext(rank=1, world_size=2), + ) + + trainer.run(TrainingTaskCollection.single(DummyData([1.0, 2.0]), DummyData([10.0]))) + + # Assembling a checkpoint from shards is collective, so a non-chief rank + # takes part in it while still writing none of the chief's own output. + assert trainer.checkpoints == [1, 2] + assert not lcurve.exists() + + def test_abstract_trainer_runs_full_validation_before_checkpoint( tmp_path: Path, ) -> None: diff --git a/source/tests/common/dpmodel/test_train_sharding.py b/source/tests/common/dpmodel/test_train_sharding.py new file mode 100644 index 0000000000..cb1244847b --- /dev/null +++ b/source/tests/common/dpmodel/test_train_sharding.py @@ -0,0 +1,52 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the distribution strategy of a training run.""" + +import unittest + +from deepmd.dpmodel.train import ( + ShardingPolicy, +) + + +class TestShardingPolicy(unittest.TestCase): + def test_each_stage_widens_what_is_sharded(self) -> None: + # What a stage shards is cumulative, and each of the three decisions + # it drives switches over at a different stage. + expected = { + 0: (False, False, False, False), + 1: (True, True, False, False), + 2: (True, True, True, False), + 3: (True, True, True, True), + } + for stage, entry in expected.items(): + with self.subTest(stage=stage): + policy = ShardingPolicy(stage=stage) + self.assertEqual( + ( + policy.enabled, + policy.shards_optimizer_state, + policy.shards_parameters, + policy.reshards_after_forward, + ), + entry, + ) + + def test_a_single_process_run_drops_the_requested_stage(self) -> None: + # One configuration has to remain usable whether or not it is launched + # across ranks, so an unusable stage is dropped rather than rejected. + policy = ShardingPolicy.from_training_params( + {"zero_stage": 3}, is_distributed=False + ) + self.assertEqual(policy.stage, 0) + + def test_an_out_of_range_stage_is_rejected_even_without_ranks(self) -> None: + for is_distributed in (True, False): + with self.subTest(is_distributed=is_distributed): + with self.assertRaisesRegex(ValueError, "must be 0, 1, 2, or 3"): + ShardingPolicy.from_training_params( + {"zero_stage": 4}, is_distributed=is_distributed + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/source/tests/pt/test_validation.py b/source/tests/pt/test_validation.py index e16913bfbe..52b33a9792 100644 --- a/source/tests/pt/test_validation.py +++ b/source/tests/pt/test_validation.py @@ -291,7 +291,6 @@ def test_full_validator_rotates_best_checkpoint(self) -> None: state_store=train_infos, num_steps=10, rank=0, - zero_stage=0, restart_training=False, ) new_best_path = validator._update_best_state( @@ -361,7 +360,6 @@ def test_full_validator_restores_top_k_checkpoints(self) -> None: state_store=train_infos, num_steps=10, rank=0, - zero_stage=0, restart_training=True, ) finally: @@ -394,7 +392,6 @@ def test_full_validator_writes_best_into_custom_checkpoint_dir(self) -> None: state_store=train_infos, num_steps=10, rank=0, - zero_stage=0, restart_training=False, checkpoint_dir=best_dir, ) @@ -430,7 +427,6 @@ def test_full_validator_reconciles_directory_checkpoints(self) -> None: state_store=train_infos, num_steps=10, rank=0, - zero_stage=0, restart_training=False, best_checkpoint_suffix=".jax", ) @@ -486,7 +482,6 @@ def test_full_validator_lmdb_full_validation_iterates_nloc_groups(self) -> None: state_store={}, num_steps=10, rank=0, - zero_stage=0, restart_training=False, ) observed_natoms = [] @@ -547,7 +542,6 @@ def test_full_validator_lmdb_excludes_default_filled_partial_labels(self) -> Non state_store={}, num_steps=10, rank=0, - zero_stage=0, restart_training=False, ) observed_flags = [] @@ -620,7 +614,6 @@ def test_full_validator_lmdb_groups_nloc_and_label_availability(self) -> None: state_store={}, num_steps=10, rank=0, - zero_stage=0, restart_training=False, ) observed_groups = [] @@ -660,7 +653,6 @@ def test_full_validator_lmdb_snapshot_requires_type_map(self) -> None: state_store={}, num_steps=10, rank=0, - zero_stage=0, restart_training=False, ) @@ -825,7 +817,6 @@ def test_predict_outputs_emits_real_and_magnetic_forces(self) -> None: state_store={}, num_steps=10, rank=0, - zero_stage=0, restart_training=False, ) self.assertIs(validator.profile, SPIN_FULL_VALIDATION_PROFILE) diff --git a/source/tests/pt_expt/test_train_utils.py b/source/tests/pt_expt/test_train_gradient.py similarity index 99% rename from source/tests/pt_expt/test_train_utils.py rename to source/tests/pt_expt/test_train_gradient.py index 7e119824db..f20912f57b 100644 --- a/source/tests/pt_expt/test_train_utils.py +++ b/source/tests/pt_expt/test_train_gradient.py @@ -3,7 +3,7 @@ import torch -from deepmd.pt_expt.train.utils import ( +from deepmd.pt_expt.train.gradient import ( NonFiniteGradGuard, clip_grad_norm_, ) diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index d5e5b19414..0c375aa651 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -2173,6 +2173,41 @@ def leave_stale_checkpoints(ckpt_dir: str) -> None: self.assertEqual(saved, [1, 2]) self.assertTrue(latest.endswith(os.path.join("ckpts", "model.ckpt-2.pt"))) + def test_diverged_interval_is_not_checkpointed(self) -> None: + """A non-finite gradient aborts the run before anything is written. + + Display is switched off so the run reaches the checkpoint boundary: + the loss report would otherwise reject the NaN first, which leaves the + gradient guard untested. + """ + config = _make_config(self.data_dir, numb_steps=2) + config["training"]["gradient_max_norm"] = 1.0 + config["training"]["save_freq"] = 1 + config["training"]["disp_training"] = False + config = update_deepmd_input(config, warning=False) + config = normalize(config) + + tmpdir = tempfile.mkdtemp(prefix="pt_expt_diverged_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + trainer = get_trainer(config) + # An infinite weight drives the forward, and therefore the whole + # gradient, out of the finite range on the very first step. + with torch.no_grad(): + next(iter(trainer.model.parameters())).fill_(float("inf")) + + with self.assertRaisesRegex(RuntimeError, "diverged"): + trainer.run() + + self.assertEqual( + [name for name in os.listdir(tmpdir) if name.endswith(".pt")], + [], + ) + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + def test_keep_ratio_retains_the_tail_of_the_run(self) -> None: config = _make_config(self.data_dir, numb_steps=6) config["training"]["save_freq"] = 2 diff --git a/source/tests/pt_expt/test_training_ddp.py b/source/tests/pt_expt/test_training_ddp.py index 3f723dcfa8..2b32e95b1e 100644 --- a/source/tests/pt_expt/test_training_ddp.py +++ b/source/tests/pt_expt/test_training_ddp.py @@ -657,11 +657,202 @@ def _worker_epoch_schedule(rank, world_size, port, data_dir, drifted, result_dic dist.destroy_process_group() +def _unsharded_parameter_shapes(data_dir: str) -> dict[str, list[int]]: + """Return the tensor shapes a single-process run records. + + They are the yardstick for a sharded run: a checkpoint must describe the + whole model, never the shard one rank happens to hold. + """ + config = _make_config(data_dir, numb_steps=1) + config = update_deepmd_input(config, warning=False) + config = normalize(config) + tmpdir = tempfile.mkdtemp(prefix="ddp_zero_reference_") + old_cwd = os.getcwd() + try: + os.chdir(tmpdir) + state = get_trainer(config)._unwrapped.state_dict() + finally: + os.chdir(old_cwd) + shutil.rmtree(tmpdir, ignore_errors=True) + return _tensor_shapes(state) + + +def _tensor_shapes(state: dict) -> dict[str, list[int]]: + return { + key: list(value.shape) + for key, value in state.items() + if isinstance(value, torch.Tensor) + } + + +def _worker_zero_stage(rank, world_size, port, data_dir, run_dir, stage, result_dict): + """Worker: train under a ZeRO stage, checkpoint, then restart from it. + + Every rank shares ``run_dir``, as ranks of a real run share a filesystem; + the checkpoint the chief writes there is what all of them resume from. + """ + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group(backend=_DDP_BACKEND, rank=rank, world_size=world_size) + old_cwd = os.getcwd() + try: + os.chdir(run_dir) + + def make_config(numb_steps: int) -> dict: + config = _make_config(data_dir, numb_steps=numb_steps) + config["training"]["zero_stage"] = stage + config["training"]["save_freq"] = numb_steps + # Exercise the clipping strategy and the divergence guard, whose + # reductions differ once the gradients are sharded. + config["training"]["gradient_max_norm"] = 5.0 + config = update_deepmd_input(config, warning=False) + return normalize(config) + + get_trainer(make_config(2)).run() + dist.barrier() + + ckpt = os.path.join(run_dir, "model.ckpt-2.pt") + written_here = rank == 0 and os.path.exists(ckpt) + shapes = ( + _tensor_shapes(torch.load(ckpt, weights_only=True)["model"]) + if rank == 0 + else {} + ) + + resumed = get_trainer(make_config(4), restart_model=ckpt) + # The momenta the first run accumulated must survive the round trip, + # whichever way the stage shards the optimizer state. The state is read + # off the optimizer that performs this rank's update, since a + # redundancy-sharded one exposes nothing before consolidation. + restored_optimizer_state = len(resumed._local_optimizer.state) + resumed.run() + + result_dict[rank] = { + "written_here": written_here, + "shapes": shapes, + "resumed_step": resumed.start_step, + "restored_optimizer_state": restored_optimizer_state, + } + finally: + os.chdir(old_cwd) + dist.destroy_process_group() + + +def _worker_zero_stage_muon(rank, world_size, port, data_dir, run_dir, result_dict): + """Worker: checkpoint a ZeRO-1 run driven by the name-routed optimizer.""" + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + dist.init_process_group(backend=_DDP_BACKEND, rank=rank, world_size=world_size) + old_cwd = os.getcwd() + try: + os.chdir(run_dir) + config = _make_config(data_dir, numb_steps=1) + config["training"]["zero_stage"] = 1 + config["training"]["save_freq"] = 1 + config["optimizer"] = {"type": "HybridMuon"} + get_trainer(normalize(update_deepmd_input(config, warning=False))).run() + + if rank != 0: + result_dict[rank] = {} + return + state = torch.load(os.path.join(run_dir, "model.ckpt-1.pt"), weights_only=False) + result_dict[rank] = { + "group_keys": sorted(state["optimizer"]["param_groups"][0]), + "tensor_valued_keys": sorted( + key + for group in state["optimizer"]["param_groups"] + for key, value in group.items() + if _holds_tensor(value) + ), + } + finally: + os.chdir(old_cwd) + dist.destroy_process_group() + + +def _holds_tensor(value) -> bool: + if isinstance(value, torch.Tensor): + return True + if isinstance(value, (list, tuple)): + return any(_holds_tensor(item) for item in value) + return False + + # --------------------------------------------------------------------------- # Test classes # --------------------------------------------------------------------------- +class TestDDPZeroStage(unittest.TestCase): + """Train, checkpoint and restart under each ZeRO stage.""" + + @classmethod + def setUpClass(cls) -> None: + data_dir = os.path.join(EXAMPLE_DIR, "data") + if not os.path.isdir(data_dir): + raise unittest.SkipTest(f"Example data not found: {data_dir}") + cls.data_dir = os.path.join(data_dir, "data_0") + cls.reference = _unsharded_parameter_shapes(cls.data_dir) + + def _assert_round_trip(self, zero_stage: int) -> None: + port = _find_free_port() + result_dict = mp.Manager().dict() + run_dir = tempfile.mkdtemp(prefix=f"ddp_zero{zero_stage}_") + try: + mp.spawn( + _worker_zero_stage, + args=(2, port, self.data_dir, run_dir, zero_stage, result_dict), + nprocs=2, + join=True, + ) + results = dict(result_dict) + finally: + shutil.rmtree(run_dir, ignore_errors=True) + + self.assertTrue(results[0]["written_here"], "the chief writes the checkpoint") + # A checkpoint records the whole model, never one rank's shard, so a + # run at any stage can resume a run at any other. + self.assertEqual(results[0]["shapes"], self.reference) + for rank in (0, 1): + self.assertEqual(results[rank]["resumed_step"], 2) + self.assertGreater( + results[rank]["restored_optimizer_state"], + 0, + "the resumed optimizer should carry the recorded state", + ) + + def test_zero_stage_1_shards_optimizer_state(self) -> None: + self._assert_round_trip(1) + + def test_zero_stage_2_shards_gradients(self) -> None: + self._assert_round_trip(2) + + def test_zero_stage_3_shards_parameters(self) -> None: + self._assert_round_trip(3) + + def test_zero_stage_1_keeps_weights_out_of_the_optimizer_state(self) -> None: + # A redundancy-sharded optimizer turns every constructor keyword into a + # param-group default, which the checkpoint then records. Parameter + # names must therefore reach the name-routed optimizer some other way, + # or each checkpoint would carry a second copy of the model. + port = _find_free_port() + result_dict = mp.Manager().dict() + run_dir = tempfile.mkdtemp(prefix="ddp_zero1_muon_") + try: + mp.spawn( + _worker_zero_stage_muon, + args=(2, port, self.data_dir, run_dir, result_dict), + nprocs=2, + join=True, + ) + chief = dict(result_dict)[0] + finally: + shutil.rmtree(run_dir, ignore_errors=True) + + self.assertNotIn("named_parameters", chief["group_keys"]) + self.assertEqual(chief["tensor_valued_keys"], []) + + class TestDDPEpochSchedule(unittest.TestCase): """An epoch spans the dataset once across the whole world, not per rank.""" From 775b68f405d820c8aba632cb705c40a6a96bbf06 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Tue, 4 Aug 2026 23:00:19 +0800 Subject: [PATCH 3/5] fix(train): inherit EMA checkpoint retention by default --- deepmd/dpmodel/train/checkpoint.py | 8 +++- deepmd/utils/argcheck.py | 10 +++-- doc/train/training-advanced.md | 1 + .../common/dpmodel/test_train_checkpoint.py | 39 +++++++++++++++++++ source/tests/common/test_argcheck_training.py | 15 +++++++ 5 files changed, 67 insertions(+), 6 deletions(-) create mode 100644 source/tests/common/test_argcheck_training.py diff --git a/deepmd/dpmodel/train/checkpoint.py b/deepmd/dpmodel/train/checkpoint.py index d9048fd801..4cd062d8e0 100644 --- a/deepmd/dpmodel/train/checkpoint.py +++ b/deepmd/dpmodel/train/checkpoint.py @@ -253,7 +253,8 @@ def build_checkpoint_stores( training_params : Mapping[str, Any] The normalized ``training`` section. ``save_ckpt``, ``save_dir``, ``save_freq``, ``max_ckpt_keep``, ``ckpt_keep_ratio`` and - ``ema_ckpt_keep`` are read from it. + ``ema_ckpt_keep`` are read from it. When ``ema_ckpt_keep`` is unset, + the EMA family inherits ``max_ckpt_keep``. num_steps : int The resolved run length, needed to turn ``ckpt_keep_ratio`` into a keep count. @@ -272,7 +273,10 @@ def build_checkpoint_stores( save_dir = training_params.get("save_dir") save_freq = int(training_params.get("save_freq", 1000)) max_keep = int(training_params.get("max_ckpt_keep", 5)) - ema_max_keep = int(training_params.get("ema_ckpt_keep", 3)) + configured_ema_max_keep = training_params.get("ema_ckpt_keep") + ema_max_keep = ( + max_keep if configured_ema_max_keep is None else int(configured_ema_max_keep) + ) ckpt_keep_ratio = training_params.get("ckpt_keep_ratio") keep_ckpt_count = resolve_keep_ckpt_count(ckpt_keep_ratio, num_steps, save_freq) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 196f6de1d0..bc771cf89b 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5440,7 +5440,9 @@ def training_args( doc_ema_ckpt_keep = ( "The maximum number of periodic EMA checkpoints to keep. " "EMA checkpoints use the same prefix-based cleanup rule as regular " - "training checkpoints, but with an EMA-specific checkpoint prefix." + "training checkpoints, but with an EMA-specific checkpoint prefix. " + "When unset, it inherits `max_ckpt_keep`, so both checkpoint families " + "retain the same number by default." ) doc_change_bias_after_training = ( "Whether to change the output bias after the last training step, " @@ -5626,11 +5628,11 @@ def training_args( ), Argument( "ema_ckpt_keep", - int, + [int, None], optional=True, - default=3, + default=None, doc=supported_backends("pt", "pt_expt") + doc_ema_ckpt_keep, - extra_check=lambda x: x > 0, + extra_check=lambda x: x is None or x > 0, extra_check_errmsg="must be greater than 0", ), Argument( diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index fd34c85c85..14cbd94ade 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -104,6 +104,7 @@ Other keys in the {ref}`training ` section are explained below: - {ref}`disp_freq ` The frequency of printing learning curve. Set in the unit of training steps - {ref}`save_freq ` The frequency of saving checkpoint. - {ref}`save_dir ` The directory where periodic checkpoints are written (PyTorch and PyTorch Exportable backends). It is created recursively if missing, while the `model.ckpt.pt` symlinks and the `checkpoint` pointer file stay in the working directory. Defaults to the working directory. +- {ref}`max_ckpt_keep ` The number of recent periodic checkpoints retained for each checkpoint family. EMA checkpoints inherit this window by default; {ref}`ema_ckpt_keep ` may override it when a different EMA window is required. - {ref}`ckpt_keep_ratio ` An alternative to `max_ckpt_keep` (PyTorch and PyTorch Exportable backends) that keeps a sliding window of `ceil(ckpt_keep_ratio * ceil(numb_steps / save_freq))` most recent checkpoints, i.e. the final `ckpt_keep_ratio` fraction of the run by step. It overrides `max_ckpt_keep` (and `ema_ckpt_keep`) when set, and works the same whether the run length is given by `numb_steps` or `numb_epoch`. ## Options and environment variables diff --git a/source/tests/common/dpmodel/test_train_checkpoint.py b/source/tests/common/dpmodel/test_train_checkpoint.py index 788fa05e51..f3f560f6cb 100644 --- a/source/tests/common/dpmodel/test_train_checkpoint.py +++ b/source/tests/common/dpmodel/test_train_checkpoint.py @@ -177,3 +177,42 @@ def test_built_stores_share_a_directory_and_split_the_pointer( store.publish(_write(store.path_for(2))) assert Path("checkpoint").read_text() == str(store.path_for(2)) + + +def test_ema_store_inherits_regular_retention_by_default( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + + store, ema_store = build_checkpoint_stores( + { + "save_ckpt": "model.ckpt", + "max_ckpt_keep": 7, + "ema_ckpt_keep": None, + }, + num_steps=10, + ema_prefix="model_ema.ckpt", + ) + + assert store.max_keep == ema_store.max_keep == 7 + + +def test_ema_store_accepts_an_explicit_retention_override( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.chdir(tmp_path) + + store, ema_store = build_checkpoint_stores( + { + "save_ckpt": "model.ckpt", + "max_ckpt_keep": 7, + "ema_ckpt_keep": 2, + }, + num_steps=10, + ema_prefix="model_ema.ckpt", + ) + + assert store.max_keep == 7 + assert ema_store.max_keep == 2 diff --git a/source/tests/common/test_argcheck_training.py b/source/tests/common/test_argcheck_training.py new file mode 100644 index 0000000000..078cb614ca --- /dev/null +++ b/source/tests/common/test_argcheck_training.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for normalized training arguments.""" + +from deepmd.utils.argcheck import ( + training_args, +) + + +def test_ema_checkpoint_retention_is_left_for_runtime_inheritance() -> None: + training_argument = training_args() + + normalized = training_argument.normalize_value({"max_ckpt_keep": 7}) + training_argument.check_value(normalized, strict=True) + + assert normalized["ema_ckpt_keep"] is None From c814ddd9739fc176bfa7959f0cd8a9b9bbf6c3ed Mon Sep 17 00:00:00 2001 From: OutisLi Date: Wed, 5 Aug 2026 11:24:16 +0800 Subject: [PATCH 4/5] fix(train): correct shared runtime edge cases Bound checkpoint retention before slicing, measure elapsed time with a monotonic clock, and preserve aliased parameters across EMA swaps. Apply eval defaults to finetune source models and align cross-platform tests and backend documentation. --- deepmd/dpmodel/train/checkpoint.py | 3 +- deepmd/dpmodel/train/timing.py | 12 ++++---- deepmd/pt/train/training.py | 7 +++-- deepmd/pt_expt/train/ema.py | 12 ++++---- deepmd/utils/argcheck.py | 2 +- doc/train/training-advanced.md | 5 ++++ .../common/dpmodel/test_train_checkpoint.py | 30 +++++++++++++++---- .../tests/common/dpmodel/test_train_timing.py | 2 +- source/tests/pt_expt/test_ema.py | 26 ++++++++++++++++ source/tests/pt_expt/test_entrypoint.py | 10 ------- source/tests/pt_expt/test_training.py | 6 +++- 11 files changed, 80 insertions(+), 35 deletions(-) create mode 100644 source/tests/pt_expt/test_ema.py diff --git a/deepmd/dpmodel/train/checkpoint.py b/deepmd/dpmodel/train/checkpoint.py index 4cd062d8e0..01a2cce2e8 100644 --- a/deepmd/dpmodel/train/checkpoint.py +++ b/deepmd/dpmodel/train/checkpoint.py @@ -193,7 +193,8 @@ def prune(self, current: Path) -> None: # The current checkpoint occupies one slot of the window when this # store holds it. occupied = 1 if current_step is not None else 0 - for _, path in retained[: len(retained) + occupied - self.max_keep]: + excess = max(0, len(retained) + occupied - self.max_keep) + for _, path in retained[:excess]: path.unlink(missing_ok=True) diff --git a/deepmd/dpmodel/train/timing.py b/deepmd/dpmodel/train/timing.py index 603e4f228b..44ae3c20d4 100644 --- a/deepmd/dpmodel/train/timing.py +++ b/deepmd/dpmodel/train/timing.py @@ -72,7 +72,7 @@ def __init__(self, *, start_step: int, num_steps: int, disp_freq: int) -> None: self._start_step = int(start_step) self._num_steps = int(num_steps) self._disp_freq = max(1, int(disp_freq)) - self._interval_start = time.time() + self._interval_start = time.monotonic() self._last_display_step = self._start_step self._timed_time = 0.0 self._timed_steps = 0 @@ -91,10 +91,10 @@ def record(self, display_step: int) -> DisplayInterval: DisplayInterval Wall-clock summary of the interval that just ended. """ - now = time.time() - wall_time = now - self._interval_start + interval_end = time.monotonic() + wall_time = interval_end - self._interval_start steps = max(1, display_step - self._last_display_step) - self._interval_start = now + self._interval_start = interval_end self._last_display_step = display_step if self._counts_toward_average(display_step): self._timed_time += wall_time @@ -111,9 +111,7 @@ def record(self, display_step: int) -> DisplayInterval: eta=int((self._num_steps - display_step) * wall_time / steps) if forecasts else None, - timestamp=datetime.datetime.fromtimestamp( - now, tz=datetime.timezone.utc - ).astimezone(), + timestamp=datetime.datetime.now(datetime.timezone.utc).astimezone(), ) def format_average(self) -> str | None: diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 98c083a7d3..53a8b98e2c 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -813,8 +813,11 @@ def epoch_length(model_key: str) -> int: target_state_dict = self.wrapper.state_dict() # pretrained_model pretrained_model_params = state_dict["_extra_state"]["model_params"] - pretrained_model = get_model_for_wrapper(pretrained_model_params) - pretrained_model_wrapper = ModelWrapper(pretrained_model) + with scoped_env_defaults(eval_env_defaults): + pretrained_model = get_model_for_wrapper( + pretrained_model_params + ) + pretrained_model_wrapper = ModelWrapper(pretrained_model) pretrained_model_wrapper.load_state_dict(state_dict) # update type related params for model_key in self.model_keys: diff --git a/deepmd/pt_expt/train/ema.py b/deepmd/pt_expt/train/ema.py index eaf482a271..d49ecc43db 100644 --- a/deepmd/pt_expt/train/ema.py +++ b/deepmd/pt_expt/train/ema.py @@ -181,11 +181,12 @@ def apply_shadow( model: torch.nn.Module | dict[str, torch.nn.Module], ) -> Iterator[None]: """Temporarily replace model parameters with the EMA shadow state.""" - backups: dict[str, torch.Tensor] = {} + named_parameters = self._named_model_parameters(model) + with torch.no_grad(): + backups = {name: param.detach().clone() for name, param in named_parameters} try: with torch.no_grad(): - for name, param in self._named_model_parameters(model): - backups[name] = param.detach().clone() + for name, param in named_parameters: param.copy_( self.shadow_params[name].to( device=param.device, @@ -195,6 +196,5 @@ def apply_shadow( yield finally: with torch.no_grad(): - for name, param in self._named_model_parameters(model): - if name in backups: - param.copy_(backups[name]) + for name, param in named_parameters: + param.copy_(backups[name]) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index bc771cf89b..6118de49f0 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -6030,7 +6030,7 @@ def validating_args() -> Argument: bool, optional=True, default=False, - doc=supported_backends("pt", "pt_expt") + doc_compiled_infer, + doc=supported_backends("pt") + doc_compiled_infer, ), Argument( "tf32_infer", diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 14cbd94ade..e756f63b6c 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -124,6 +124,11 @@ An explanation will be provided **`--init-model model.ckpt`**, initializes the model training with an existing model that is stored in the path prefix of checkpoint files `model.ckpt`, the network architectures should match. **`--restart model.ckpt`**, continues the training from the checkpoint `model.ckpt`. +For the PyTorch and PyTorch Exportable backends, when checkpoint retention is +enabled, writing the next periodic or EMA checkpoint at step `N` removes +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. **`--init-frz-model frozen_model.pb`**, initializes the training with an existing model that is stored in `frozen_model.pb`. diff --git a/source/tests/common/dpmodel/test_train_checkpoint.py b/source/tests/common/dpmodel/test_train_checkpoint.py index f3f560f6cb..c697a2e8c7 100644 --- a/source/tests/common/dpmodel/test_train_checkpoint.py +++ b/source/tests/common/dpmodel/test_train_checkpoint.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later """Tests for the backend-independent checkpoint layout and retention.""" +import platform from pathlib import ( Path, ) @@ -20,6 +21,16 @@ def _write(path: Path) -> Path: return path +def _assert_prefix_alias(alias: Path, target: Path, relative_target: str) -> None: + assert alias.exists() + if platform.system() == "Windows": + assert alias.read_bytes() == target.read_bytes() + return + assert alias.is_symlink() + assert alias.resolve() == target + assert alias.readlink().as_posix() == relative_target + + def test_numbered_paths_follow_save_dir(tmp_path: Path) -> None: store = CheckpointStore(tmp_path / "run" / "model.ckpt") assert store.path_for(7) == tmp_path / "run" / "model.ckpt-7.pt" @@ -39,9 +50,7 @@ def test_publish_links_the_prefix_relative_to_its_directory(tmp_path: Path) -> N store.publish(path) latest = tmp_path / "model.ckpt.pt" - assert latest.is_symlink() - assert latest.resolve() == path - assert latest.readlink().as_posix() == "model.ckpt-3.pt" + _assert_prefix_alias(latest, path, "model.ckpt-3.pt") assert (tmp_path / "checkpoint").read_text() == str(path) @@ -53,8 +62,7 @@ def test_publish_reaches_across_save_dir(tmp_path: Path) -> None: store.publish(path) latest = tmp_path / "model.ckpt.pt" - assert latest.resolve() == path - assert latest.readlink().as_posix() == "ckpts/model.ckpt-3.pt" + _assert_prefix_alias(latest, path, "ckpts/model.ckpt-3.pt") def test_prune_keeps_the_newest_checkpoints(tmp_path: Path) -> None: @@ -71,6 +79,16 @@ def test_prune_keeps_the_newest_checkpoints(tmp_path: Path) -> None: assert (tmp_path / "model.ckpt.pt").exists() +def test_prune_keeps_every_checkpoint_below_the_window(tmp_path: Path) -> None: + store = CheckpointStore(tmp_path / "model.ckpt", max_keep=10) + for step in range(1, 10): + _write(store.path_for(step)) + + store.prune(store.path_for(9)) + + assert all(store.path_for(step).exists() for step in range(1, 10)) + + def test_prune_drops_checkpoints_left_by_a_longer_run(tmp_path: Path) -> None: """A rerun in a finished directory keeps its own checkpoint. @@ -102,7 +120,7 @@ def test_prune_ignores_foreign_names_and_symlinks(tmp_path: Path) -> None: assert other.exists() assert unnumbered.exists() - assert (tmp_path / "model.ckpt.pt").is_symlink() + _assert_prefix_alias(tmp_path / "model.ckpt.pt", current, "model.ckpt-2.pt") def test_prune_without_a_window_keeps_every_checkpoint(tmp_path: Path) -> None: diff --git a/source/tests/common/dpmodel/test_train_timing.py b/source/tests/common/dpmodel/test_train_timing.py index b8b27a982a..bede244467 100644 --- a/source/tests/common/dpmodel/test_train_timing.py +++ b/source/tests/common/dpmodel/test_train_timing.py @@ -28,7 +28,7 @@ def advance(self, seconds: float) -> None: def clock(monkeypatch: pytest.MonkeyPatch) -> FakeClock: """Replace the clock the timer reads, leaving the stdlib one untouched.""" fake = FakeClock() - monkeypatch.setattr(timing_module, "time", types.SimpleNamespace(time=fake)) + monkeypatch.setattr(timing_module, "time", types.SimpleNamespace(monotonic=fake)) return fake diff --git a/source/tests/pt_expt/test_ema.py b/source/tests/pt_expt/test_ema.py new file mode 100644 index 0000000000..3611d37a59 --- /dev/null +++ b/source/tests/pt_expt/test_ema.py @@ -0,0 +1,26 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for exponential moving-average model weights.""" + +import torch + +from deepmd.pt_expt.train.ema import ( + ModelEMA, +) + + +def test_apply_shadow_restores_parameters_shared_across_models() -> None: + left = torch.nn.Linear(1, 1, bias=False) + right = torch.nn.Linear(1, 1, bias=False) + right.weight = left.weight + models = {"left": left, "right": right} + + with torch.no_grad(): + left.weight.fill_(1.0) + ema = ModelEMA(models, decay=0.9) + for shadow in ema.shadow_params.values(): + shadow.fill_(2.0) + + with ema.apply_shadow(models): + torch.testing.assert_close(left.weight, torch.full_like(left.weight, 2.0)) + + torch.testing.assert_close(left.weight, torch.ones_like(left.weight)) diff --git a/source/tests/pt_expt/test_entrypoint.py b/source/tests/pt_expt/test_entrypoint.py index d6c456fa68..24bd98efaf 100644 --- a/source/tests/pt_expt/test_entrypoint.py +++ b/source/tests/pt_expt/test_entrypoint.py @@ -7,7 +7,6 @@ from deepmd.pt_expt.entrypoints.main import ( PTExptTrainEntrypoint, _ensure_pt_expt_model_suffix, - _ensure_stat_file_path, train, ) @@ -149,12 +148,3 @@ def test_pt_expt_entrypoint_rejects_random_model_key( {"model": {"model_dict": {"RANDOM": {}}}}, TrainEntrypointOptions(input_file="input.json"), ) - - -def test_pt_expt_stat_file_path_creates_hdf5_parent(tmp_path) -> None: - stat_file = tmp_path / "stats" / "model_stat.hdf5" - - stat_path = _ensure_stat_file_path(str(stat_file)) - - assert stat_file.exists() - assert stat_path is not None diff --git a/source/tests/pt_expt/test_training.py b/source/tests/pt_expt/test_training.py index 0c375aa651..0569f21851 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -11,6 +11,7 @@ import copy import math import os +import platform import shutil import tempfile import unittest @@ -2250,7 +2251,10 @@ def test_ema_checkpoint_holds_smoothed_weights(self) -> None: ema_ckpt = os.path.join(tmpdir, "model_ema.ckpt-4.pt") self.assertTrue(os.path.exists(ema_ckpt)) - self.assertTrue(os.path.islink(os.path.join(tmpdir, "model_ema.ckpt.pt"))) + ema_alias = os.path.join(tmpdir, "model_ema.ckpt.pt") + self.assertTrue(os.path.exists(ema_alias)) + if platform.system() != "Windows": + self.assertTrue(os.path.islink(ema_alias)) ema_state = torch.load(ema_ckpt, weights_only=True) live_state = torch.load( From a425d1e59ebe110bab254179b19bb792888f8aa6 Mon Sep 17 00:00:00 2001 From: OutisLi Date: Fri, 7 Aug 2026 09:36:01 +0800 Subject: [PATCH 5/5] fix(train): handle optimizer and checkpoint edge cases --- deepmd/pt/train/training.py | 8 +++++--- deepmd/pt_expt/train/training.py | 6 +++--- deepmd/utils/argcheck.py | 6 ++++-- doc/train/parallel-training.md | 20 ++++++++++++++------ 4 files changed, 26 insertions(+), 14 deletions(-) diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 53a8b98e2c..5a190fe5af 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -1741,9 +1741,11 @@ def log_loss_valid(_task_key: str = "Default") -> dict: ), ) - should_save_checkpoint = ( - (display_step_id) % self.save_freq == 0 and _step_id != self.start_step - ) or (display_step_id) == self.num_steps + should_save_checkpoint = display_step_id == self.num_steps or ( + self.save_freq > 0 + and display_step_id % self.save_freq == 0 + and _step_id != self.start_step + ) if should_save_checkpoint: # Abort before writing if any gradient norm since the previous # checkpoint was non-finite. diff --git a/deepmd/pt_expt/train/training.py b/deepmd/pt_expt/train/training.py index 088c912fc9..f277452fcb 100644 --- a/deepmd/pt_expt/train/training.py +++ b/deepmd/pt_expt/train/training.py @@ -1915,6 +1915,8 @@ def update_finetune_bias( # Optimiser ----------------------------------------------------------- opt_type = optimizer_params.get("type", "Adam") + if opt_type not in ("Adam", "AdamW", "HybridMuon"): + raise ValueError(f"Unsupported optimizer type: {opt_type}") # LambdaLR multiplies each param group's initial learning rate by the # lambda value. Warmup schedules legitimately return zero at step 0, # so use the nonzero schedule base as the denominator and let the @@ -1933,7 +1935,7 @@ def update_finetune_bias( betas=adam_betas, weight_decay=weight_decay, ) - elif opt_type == "HybridMuon": + else: self.optimizer = self._create_optimizer( HybridMuonOptimizer, lr=initial_lr, @@ -1958,8 +1960,6 @@ def update_finetune_bias( self._local_optimizer.set_param_names( tuple(self.wrapper.named_parameters()) ) - else: - raise ValueError(f"Unsupported optimizer type: {opt_type}") if optimizer_state_dict is not None: self._load_optimizer_state(optimizer_state_dict) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 6118de49f0..eae723bc3d 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -5417,8 +5417,10 @@ def training_args( "set, checkpoints are written to the working directory." ) doc_max_ckpt_keep = ( - "The maximum number of checkpoints to keep. " - "The oldest checkpoints will be deleted once the number of checkpoints exceeds max_ckpt_keep. " + "The maximum number of recent periodic checkpoints to keep for the " + "regular checkpoint family. The EMA checkpoint family inherits this " + "value by default unless `ema_ckpt_keep` overrides it. The oldest " + "checkpoints are deleted when a family's retention window is exceeded. " "Defaults to 5." ) doc_ckpt_keep_ratio = ( diff --git a/doc/train/parallel-training.md b/doc/train/parallel-training.md index 7c032179c7..ff8fa2c38e 100644 --- a/doc/train/parallel-training.md +++ b/doc/train/parallel-training.md @@ -92,15 +92,18 @@ optional arguments: master) ``` -## PyTorch Implementation {{ pytorch_icon }} +## PyTorch implementations {{ pytorch_icon }} -Currently, parallel training in pytorch version is implemented in the form of PyTorch Distributed Data Parallelism [DDP](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html). +Parallel training in the PyTorch and PyTorch Exportable backends uses PyTorch +distributed primitives. Stages 0 and 1 use +[Distributed Data Parallelism (DDP)](https://pytorch.org/docs/stable/generated/torch.nn.parallel.DistributedDataParallel.html), +while stages 2 and 3 use FSDP2. DeePMD-kit will decide whether to launch the training in parallel (distributed) mode or in serial mode depending on your execution command. ### Optional ZeRO memory optimization -In PyTorch backend, DeePMD-kit supports ZeRO (Zero Redundancy Optimizer) stages -to reduce per-GPU memory usage during distributed training. +In both PyTorch backends, DeePMD-kit supports ZeRO (Zero Redundancy Optimizer) +stages to reduce per-GPU memory usage during distributed training. | `zero_stage` | Strategy | Communication | Memory saving | | ------------ | ----------------------------- | ------------- | --------------------------------------------- | @@ -137,11 +140,16 @@ Enable it in input config: Constraints: -- Works only in PyTorch backend. +- Works in the PyTorch and PyTorch Exportable backends. - Requires distributed launch with `torchrun`. - Currently single-task only. - Not supported with `LKF` optimizer. -- `change_bias_after_training` must be `false`. +- `training.change_bias_after_training` must be `false`. +- Stages 2 and 3 require PyTorch 2.6 or later and are incompatible with + `training.enable_ema`, `validating.full_validation`, and + `validating.ema_full_validation`. +- In the PyTorch Exportable backend, stages 2 and 3 are additionally + incompatible with `training.enable_compile`. ### Dataloader and Dataset