diff --git a/deepmd/dpmodel/train/__init__.py b/deepmd/dpmodel/train/__init__.py index e6124d8ce9..9e9c7b1f64 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,13 @@ StepSchedule, resolve_step_schedule, ) +from .sharding import ( + ShardingPolicy, +) +from .timing import ( + DisplayInterval, + TrainingTimer, +) from .trainer import ( DEFAULT_TASK_KEY, AbstractTrainer, @@ -32,8 +44,11 @@ "DEFAULT_TASK_KEY", "AbstractTrainEntrypoint", "AbstractTrainer", + "CheckpointStore", + "DisplayInterval", "LearningCurveWriter", "RankContext", + "ShardingPolicy", "StepSchedule", "TrainEntrypointOptions", "TrainStepResult", @@ -41,10 +56,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..01a2cce2e8 --- /dev/null +++ b/deepmd/dpmodel/train/checkpoint.py @@ -0,0 +1,311 @@ +# 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 + excess = max(0, len(retained) + occupied - self.max_keep) + for _, path in retained[:excess]: + 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. 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. + 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)) + 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) + 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/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/timing.py b/deepmd/dpmodel/train/timing.py new file mode 100644 index 0000000000..44ae3c20d4 --- /dev/null +++ b/deepmd/dpmodel/train/timing.py @@ -0,0 +1,140 @@ +# 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.monotonic() + 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. + """ + interval_end = time.monotonic() + wall_time = interval_end - self._interval_start + steps = max(1, display_step - self._last_display_step) + self._interval_start = interval_end + 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.now(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..4377863629 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, @@ -588,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 ): @@ -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() @@ -644,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, *, @@ -682,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 ( @@ -704,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 @@ -713,37 +729,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..5a190fe5af 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,12 @@ import numpy as np import torch -from deepmd.common import ( - symlink_prefix_files, -) from deepmd.dpmodel.train import ( + DEFAULT_TASK_KEY, + CheckpointStore, + ShardingPolicy, + TrainingTimer, + build_checkpoint_stores, change_model_out_bias, resolve_step_schedule, ) @@ -39,6 +39,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 +70,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 +102,25 @@ 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.gradient import ( + NonFiniteGradGuard, + clip_grad_norm_, +) +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.utils.data import ( DataRequirementItem, has_data_requirement, @@ -198,13 +200,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,34 +223,23 @@ 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) 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." ) @@ -455,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. @@ -475,7 +460,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 +701,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) @@ -836,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: @@ -1054,7 +1034,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 " @@ -1070,7 +1050,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) @@ -1125,7 +1105,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}'") @@ -1138,7 +1118,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) @@ -1156,16 +1136,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) @@ -1174,14 +1146,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 +1212,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, + sharding=self.sharding, ) - return start_step is not None and start_step <= self.num_steps def _raise_if_full_validation_unsupported( self, @@ -1356,48 +1260,19 @@ 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." ) - @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, @@ -1419,7 +1294,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, @@ -1429,7 +1304,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 @@ -1439,7 +1314,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, @@ -1530,7 +1405,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 @@ -1544,7 +1419,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): @@ -1771,7 +1646,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, @@ -1823,39 +1699,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: @@ -1888,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. @@ -1898,30 +1753,21 @@ 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 = 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 +1807,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,70 +1854,48 @@ 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: + if self.num_steps == 0 and self.sharding.enabled: # 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 ( 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 = 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 = ( @@ -2149,7 +1975,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) @@ -2160,7 +1986,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: @@ -2180,28 +2006,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 +2022,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 +2030,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 +2050,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 +2065,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 +2077,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 +2117,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 +2132,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 95% rename from deepmd/pt/train/ema.py rename to deepmd/pt_expt/train/ema.py index eaf482a271..d49ecc43db 100644 --- a/deepmd/pt/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/pt/train/utils.py b/deepmd/pt_expt/train/gradient.py similarity index 60% rename from deepmd/pt/train/utils.py rename to deepmd/pt_expt/train/gradient.py index fe6d29b24e..cc3b863015 100644 --- a/deepmd/pt/train/utils.py +++ b/deepmd/pt_expt/train/gradient.py @@ -1,23 +1,26 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -"""Training utility functions.""" +"""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, ) -import os -from contextlib import ( - contextmanager, -) -from math import ( - ceil, -) -from pathlib import ( - Path, -) from typing import ( TYPE_CHECKING, - Any, ) import torch @@ -25,10 +28,15 @@ if TYPE_CHECKING: from collections.abc import ( Callable, - Generator, Iterable, ) +__all__ = [ + "NonFiniteGradGuard", + "clip_grad_norm_", + "raise_nonfinite_gradient_norm", +] + def clip_grad_norm_( parameters: Iterable[torch.nn.Parameter], @@ -182,104 +190,3 @@ def raise_nonfinite_gradient_norm( "Non-finite gradient norm; training has diverged.\n" f"Parameters with non-finite gradients:\n{detail}" ) - - -@contextmanager -def scoped_env_defaults(defaults: dict[str, str]) -> Generator[None, None, None]: - """Temporarily set missing environment variables and restore them afterward.""" - previous = {key: os.environ.get(key) for key in defaults} - try: - for key, value in defaults.items(): - os.environ.setdefault(key, value) - yield - finally: - for key, value in previous.items(): - if value is None: - os.environ.pop(key, None) - else: - 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: - """ - Resolve the directory for full-validation best checkpoints. - - Parameters - ---------- - validating_params : dict - The ``validating`` section of the training configuration. - save_ckpt : str - The regular checkpoint prefix from ``training.save_ckpt``. - - Returns - ------- - Path - ``validating.save_best_dir`` when set, otherwise the directory derived - from ``save_ckpt``. - """ - save_best_dir = validating_params.get("save_best_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_expt/train/training.py b/deepmd/pt_expt/train/training.py index 2a8165c90f..f277452fcb 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, @@ -27,15 +26,33 @@ 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, TrainStepResult, + build_checkpoint_stores, change_model_out_bias, change_model_out_bias_by_task, resolve_step_schedule, @@ -50,16 +67,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 +97,25 @@ 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.gradient import ( + NonFiniteGradGuard, + clip_grad_norm_, +) +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, ) @@ -372,14 +404,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 # --------------------------------------------------------------------------- @@ -930,12 +954,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 @@ -949,6 +975,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 @@ -1085,18 +1150,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, @@ -1107,11 +1163,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 @@ -1273,29 +1326,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 @@ -1458,30 +1499,41 @@ 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") 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( 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] = {} 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 ---------------------------------------------------------------- @@ -1596,6 +1648,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 @@ -1603,6 +1665,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) @@ -1649,87 +1712,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")) @@ -1743,10 +1731,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: @@ -1914,29 +1907,92 @@ 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") + 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 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, + ) + else: + 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()) + ) + + 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 + # 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") ) # 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", {}) - 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() @@ -1949,53 +2005,42 @@ 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, + sharding=self.sharding, ) - 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, @@ -2007,6 +2052,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 ): @@ -2140,9 +2191,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, ) @@ -2240,9 +2293,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.""" @@ -2303,13 +2501,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: - ckpt_path = Path(f"{self.save_ckpt}-{step}.pt") + # 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) - 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}") + 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) + if self.rank == 0: + self.ema_ckpt_store.publish(ema_path) + self.ema_ckpt_store.prune(ema_path) def _save_full_validation_checkpoint( self, @@ -2321,8 +2533,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 @@ -2337,32 +2602,64 @@ 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(), - } + 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 _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) + 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 @@ -2396,7 +2693,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: @@ -2405,7 +2701,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.""" @@ -2442,16 +2739,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.""" @@ -2507,11 +2807,20 @@ 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() + if self.model_ema is not None: + self.model_ema.update(self.model) return TrainStepResult( task_key=task_key, step=step, @@ -2548,8 +2857,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 new file mode 100644 index 0000000000..5786ff10d0 --- /dev/null +++ b/deepmd/pt_expt/train/utils.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Training utility functions.""" + +from __future__ import ( + annotations, +) + +import os +from contextlib import ( + contextmanager, +) +from pathlib import ( + Path, +) +from typing import ( + TYPE_CHECKING, + Any, +) + +if TYPE_CHECKING: + from collections.abc import ( + Generator, + ) + + import torch + + +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 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.""" + previous = {key: os.environ.get(key) for key in defaults} + try: + for key, value in defaults.items(): + os.environ.setdefault(key, value) + yield + finally: + for key, value in previous.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + + +def resolve_best_checkpoint_dir( + validating_params: dict[str, Any], save_ckpt: str +) -> Path: + """ + Resolve the directory for full-validation best checkpoints. + + Parameters + ---------- + validating_params : dict + The ``validating`` section of the training configuration. + save_ckpt : str + The regular checkpoint prefix from ``training.save_ckpt``. + + Returns + ------- + Path + ``validating.save_best_dir`` when set, otherwise the directory derived + from ``save_ckpt``. + """ + save_best_dir = validating_params.get("save_best_dir") + if save_best_dir: + return Path(save_best_dir) + return Path(save_ckpt).parent diff --git a/deepmd/pt/train/validation.py b/deepmd/pt_expt/train/validation.py similarity index 83% rename from deepmd/pt/train/validation.py rename to deepmd/pt_expt/train/validation.py index 784212ecb5..edd89e6d04 100644 --- a/deepmd/pt/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, @@ -35,20 +38,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 +83,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 = "# " @@ -204,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, @@ -221,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(".") ) @@ -335,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() @@ -425,24 +426,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 +455,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 +896,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, + sharding: ShardingPolicy | None = None, +) -> 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. + sharding : ShardingPolicy, optional + The distribution strategy of the run, which decides whether checkpoint + collection is a collective operation. Defaults to no sharding. + + 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, + sharding=ShardingPolicy() if sharding is None else sharding, + 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..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 = ( @@ -5440,7 +5442,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, " @@ -5508,7 +5512,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 " @@ -5591,7 +5597,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 +5608,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,24 +5617,24 @@ 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", ), Argument( "ema_ckpt_keep", - int, + [int, None], optional=True, - default=3, - doc=supported_backends("pt") + doc_ema_ckpt_keep, - extra_check=lambda x: x > 0, + default=None, + doc=supported_backends("pt", "pt_expt") + doc_ema_ckpt_keep, + extra_check=lambda x: x is None or x > 0, extra_check_errmsg="must be greater than 0", ), Argument( @@ -5712,7 +5718,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", @@ -5936,7 +5942,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 +5968,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", @@ -6038,7 +6046,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/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 diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md index 0112462581..e756f63b6c 100644 --- a/doc/train/training-advanced.md +++ b/doc/train/training-advanced.md @@ -103,8 +103,9 @@ 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}`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 @@ -123,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_abstract_trainer.py b/source/tests/common/dpmodel/test_train_abstract_trainer.py index 5ac21b71ee..23eaeb1a22 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( @@ -200,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_checkpoint.py b/source/tests/common/dpmodel/test_train_checkpoint.py new file mode 100644 index 0000000000..c697a2e8c7 --- /dev/null +++ b/source/tests/common/dpmodel/test_train_checkpoint.py @@ -0,0 +1,236 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Tests for the backend-independent checkpoint layout and retention.""" + +import platform +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 _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" + + 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_prefix_alias(latest, path, "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_prefix_alias(latest, path, "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_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. + + 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_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: + """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)) + + +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/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/common/dpmodel/test_train_timing.py b/source/tests/common/dpmodel/test_train_timing.py new file mode 100644 index 0000000000..bede244467 --- /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(monotonic=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_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 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..52b33a9792 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, ) @@ -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_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 9068df58d3..24bd98efaf 100644 --- a/source/tests/pt_expt/test_entrypoint.py +++ b/source/tests/pt_expt/test_entrypoint.py @@ -148,76 +148,3 @@ def test_pt_expt_entrypoint_rejects_random_model_key( {"model": {"model_dict": {"RANDOM": {}}}}, TrainEntrypointOptions(input_file="input.json"), ) - - -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 - - trainer.save_checkpoint(1) - - 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" diff --git a/source/tests/pt/test_train_utils.py b/source/tests/pt_expt/test_train_gradient.py similarity index 85% rename from source/tests/pt/test_train_utils.py rename to source/tests/pt_expt/test_train_gradient.py index c3944f0972..f20912f57b 100644 --- a/source/tests/pt/test_train_utils.py +++ b/source/tests/pt_expt/test_train_gradient.py @@ -3,10 +3,9 @@ import torch -from deepmd.pt.train.utils import ( +from deepmd.pt_expt.train.gradient 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..0569f21851 100644 --- a/source/tests/pt_expt/test_training.py +++ b/source/tests/pt_expt/test_training.py @@ -9,12 +9,15 @@ """ import copy -import datetime import math import os +import platform import shutil import tempfile import unittest +from collections.abc import ( + Callable, +) from pathlib import ( Path, ) @@ -27,9 +30,6 @@ import pytest import torch -from deepmd.loggers.training import ( - format_training_message, -) from deepmd.pt.optimizer import ( HybridMuonOptimizer, ) @@ -448,7 +448,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 +494,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 +2101,205 @@ 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 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) + 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_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 + ``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_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_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 + 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)) + 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( + 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__": 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."""