diff --git a/docs/api/diffusion/distillation.rst b/docs/api/diffusion/distillation.rst index 91fbcac74d..21a2c2b081 100644 --- a/docs/api/diffusion/distillation.rst +++ b/docs/api/diffusion/distillation.rst @@ -265,6 +265,46 @@ Keep phase selection and f-distill statistics updates outside compilation. Check combined domain parallelism and FSDP on your target hardware before a long training run. +Restarting training +------------------- + +``CheckpointStateBundle`` groups named models, optimizers, and loss statistics +for a training restart. A single-model checkpoint stores one model; a bundle +keeps the state of the participating models and optimizers together. + +Use the model's checkpoint API when you only need its weights for inference. +Use a bundle to resume training with optimizer state and loss statistics. +Give each model a name, including fake-score, discriminator, and EMA models +when present. Use these names for their optimizers: + +.. code-block:: python + + from physicsnemo.diffusion.distillation import CheckpointStateBundle + + bundle = CheckpointStateBundle( + {"student": student, "teacher": teacher, "fake_score": fake_score}, + optimizers={"student": student_optimizer, "fake_score": fake_optimizer}, + stateful={"loss": loss_fn}, + ) + bundle.save_checkpoint("checkpoints", epoch=20, progress={"step": 20}, include_rng=True) + progress = bundle.load_checkpoint("checkpoints", epoch=20, restore_rng=True) + +Recreate the same models, noise schedules, optimizers, and loss settings before +loading. Save after completing an optimizer update and choose a new ``epoch`` +index each time. To restore a data-loader cursor, add it to ``stateful`` with +``state_dict`` and ``load_state_dict`` methods. The bundle does not save pending +gradients or reconstruct your training configuration. + +Call save and load on every rank, and restart with the same world size and +parallel layout. The bundle uses the existing PhysicsNeMo checkpoint utilities. +Use ``optimizer_models`` if optimizer and model names differ. +Each optimizer saves a copy of its model weights, even when optimizers share +a model. + +If loading fails, recreate your objects before retrying: some objects may +already contain restored state. Existing checkpoint indices cannot be +overwritten, including incomplete saves. + Custom implementations ---------------------- diff --git a/physicsnemo/diffusion/distillation/__init__.py b/physicsnemo/diffusion/distillation/__init__.py index b30b34697d..67ead049ba 100644 --- a/physicsnemo/diffusion/distillation/__init__.py +++ b/physicsnemo/diffusion/distillation/__init__.py @@ -21,6 +21,7 @@ from physicsnemo.core.warnings import ExperimentalFeatureWarning from .base import BaseDistillationLoss, LossMap +from .checkpoint import CheckpointStateBundle from .losses import DistillationLoss, WeightedDistillationLoss warnings.warn( @@ -31,6 +32,7 @@ __all__ = [ "BaseDistillationLoss", + "CheckpointStateBundle", "LossMap", "DistillationLoss", "WeightedDistillationLoss", diff --git a/physicsnemo/diffusion/distillation/checkpoint.py b/physicsnemo/diffusion/distillation/checkpoint.py new file mode 100644 index 0000000000..61e3192712 --- /dev/null +++ b/physicsnemo/diffusion/distillation/checkpoint.py @@ -0,0 +1,354 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Named training-state bundles built from existing PhysicsNeMo checkpoints.""" + +import random +import re +from collections.abc import Mapping +from pathlib import Path +from typing import Any, Protocol + +import fsspec +import numpy as np +import torch + +from physicsnemo.distributed import DistributedManager +from physicsnemo.utils import checkpoint as _checkpoint + + +class _Stateful(Protocol): + r"""Objects that expose restart state independently of their configuration.""" + + def state_dict(self) -> dict[str, Any]: ... + + def load_state_dict(self, state: Mapping[str, Any]) -> None: ... + + +class CheckpointStateBundle: + r"""Save and restore named models and their training state together. + A bundle groups your student, teacher, optimizers, and loss statistics. + + Use a bundle to resume distillation with the associated training state. + For model weights alone, use the model's checkpoint API. Save after an + optimizer update, then recreate the same models, noise schedules, and + loss settings before loading. Pending gradients are not part of the + checkpoint. + + Parameters + ---------- + models : mapping + Models to save, keyed by names such as ``student`` and ``fake_score``. + Use the same names when resuming training. + optimizers : mapping, optional + Optimizers to save, keyed by the name of the model they update. + optimizer_models : mapping, optional + Map optimizer names to model names when they differ, for example + ``{"student_head": "student"}``. + schedulers, scalers : mapping, optional + Learning-rate schedulers and gradient scalers keyed by optimizer name. + stateful : mapping, optional + Other objects to save, such as the loss or a data-loader cursor. + Each needs ``state_dict`` and ``load_state_dict`` methods. + + Examples + -------- + Register the student under a stable name and save its state at a chosen epoch. + Loading that epoch restores the registered module and returns its progress data. + + >>> import tempfile + >>> model = torch.nn.Linear(3, 2) + >>> # The key must match when reconstructing a bundle for loading. + >>> bundle = CheckpointStateBundle({"student": model}) + >>> with tempfile.TemporaryDirectory() as directory: + ... bundle.save_checkpoint(directory, epoch=0, progress={"step": 0}) + ... progress = bundle.load_checkpoint(directory, epoch=0) + >>> progress + {'step': 0} + + Notes + ----- + Call save and load on every rank when training across devices. Restart + with the same world size and parallel layout. The bundle uses PhysicsNeMo's + checkpoint utilities and saves each rank's buffers and optional random state. + Each optimizer saves a copy of its model's weights, even when optimizers + share a model. + """ + + def __init__( + self, + models: Mapping[str, torch.nn.Module], + *, + optimizers: Mapping[str, torch.optim.Optimizer] | None = None, + optimizer_models: Mapping[str, str] | None = None, + schedulers: Mapping[str, Any] | None = None, + scalers: Mapping[str, Any] | None = None, + stateful: Mapping[str, _Stateful] | None = None, + ) -> None: + self.models = dict(models) + self.optimizers = dict(optimizers or {}) + self.optimizer_models = dict(optimizer_models or {}) + self.schedulers = dict(schedulers or {}) + self.scalers = dict(scalers or {}) + self.stateful = dict(stateful or {}) + if not self.models: + raise ValueError("A checkpoint bundle requires named models.") + for collection in ( + self.models, + self.optimizers, + self.schedulers, + self.scalers, + self.stateful, + ): + if any( + not isinstance(name, str) + or re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_-]*", name) is None + for name in collection + ): + raise ValueError( + "Checkpoint names must contain letters, digits, hyphens, or underscores." + ) + if set(self.optimizer_models) - set(self.optimizers): + raise ValueError("optimizer_models contains unknown optimizer names.") + if (set(self.schedulers) | set(self.scalers)) - set(self.optimizers): + raise ValueError("Schedulers and scalers must name configured optimizers.") + for name, optimizer in self.optimizers.items(): + role = self.optimizer_models.setdefault(name, name) + if role not in self.models: + raise ValueError(f"Optimizer {name!r} needs a known model role.") + parameters = set(self.models[role].parameters()) + if any( + parameter not in parameters + for group in optimizer.param_groups + for parameter in group["params"] + ): + raise ValueError( + f"Optimizer {name!r} contains parameters outside model role {role!r}." + ) + + def _schema(self) -> dict[str, Any]: + r"""Describe role associations independently of mapping insertion order.""" + return { + "schema_version": 1, + "models": sorted(self.models), + "optimizer_models": dict(self.optimizer_models), + "schedulers": sorted(self.schedulers), + "scalers": sorted(self.scalers), + "stateful": sorted(self.stateful), + "world_size": torch.distributed.get_world_size() + if torch.distributed.is_initialized() + else 1, + } + + @staticmethod + def _rank() -> int: + r"""Initialize the same manager used by the underlying checkpoint utilities.""" + if not DistributedManager.is_initialized(): + DistributedManager.initialize() + return DistributedManager().rank + + @staticmethod + def _barrier() -> None: + r"""Finish collective snapshots before publishing per-rank bundle metadata.""" + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + def _entries(self, root: str, rank: int): + r"""Yield standard checkpoint calls with explicit optimizer associations.""" + entries = [ + (f"models/{role}", model, None) + for role, model in sorted(self.models.items()) + ] + entries += [ + (f"optimizers/{name}", self.models[self.optimizer_models[name]], name) + for name in sorted(self.optimizers) + ] + for name, model, optimizer_name in entries: + named = _checkpoint._unique_model_names([model]) + distributed = any( + _checkpoint._is_distributed_model(value) for value in named.values() + ) + suffix = "distributed" if distributed else f"rank-{rank}" + kwargs = {"models": model} + if optimizer_name is not None: + kwargs.update( + optimizer=self.optimizers[optimizer_name], + optimizer_model=model, + scheduler=self.schedulers.get(optimizer_name), + scaler=self.scalers.get(optimizer_name), + ) + yield f"{root}/{name}/{suffix}", kwargs, named, distributed + + def save_checkpoint( + self, + path: Path | str, + *, + epoch: int, + progress: Mapping[str, Any] | None = None, + include_rng: bool = False, + ) -> None: + r"""Save the bundle's models and training state at one checkpoint index. + + Call after a completed optimizer update on every rank. Choose a new + index for each save, including retries after an incomplete save. + + Parameters + ---------- + path : Path or str + Root directory or filesystem URI accepted by PhysicsNeMo checkpoints. + epoch : int + Nonnegative checkpoint index. Choose a new index for each save. + progress : mapping, optional + Counters to restore, such as ``{"step": 1000}``. + include_rng : bool, optional, default=False + Capture Python, NumPy, CPU, and available CUDA random states per rank. + """ + if epoch < 0: + raise ValueError("epoch must be nonnegative.") + rank = self._rank() + root = f"{str(path).rstrip('/')}/step-{epoch}" + fs = fsspec.filesystem(fsspec.utils.get_protocol(root)) + exists = fs.exists(root) + occupied = [exists] + if torch.distributed.is_initialized(): + occupied = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(occupied, exists) + if any(occupied): + raise FileExistsError(f"Checkpoint index already exists: {root}.") + self._barrier() + metadata = { + "bundle": self._schema(), + "progress": dict(progress or {}), + "stateful": { + name: value.state_dict() for name, value in self.stateful.items() + }, + } + if include_rng: + numpy_state = np.random.get_state() + metadata["rng"] = { + "python": random.getstate(), + "torch": torch.get_rng_state(), + "numpy": (numpy_state[0], numpy_state[1].tolist(), *numpy_state[2:]), + "cuda": torch.cuda.get_rng_state_all() + if torch.cuda.is_available() + else [], + } + for directory, kwargs, _, _ in self._entries(root, rank): + _checkpoint.save_checkpoint( + directory, epoch=epoch, metadata={"bundle_epoch": epoch}, **kwargs + ) + self._barrier() + _checkpoint.save_checkpoint( + f"{root}/state/rank-{rank}", epoch=epoch, metadata=metadata + ) + self._barrier() + + def load_checkpoint( + self, + path: Path | str, + *, + epoch: int, + device: str | torch.device = "cpu", + restore_rng: bool = False, + ) -> dict[str, Any]: + r"""Load saved weights and training state into the objects in this bundle. + + To resume training, recreate the same objects and call on every rank + with the saved world size and parallel layout. If loading fails, + recreate the objects before retrying because some state may have loaded. + + Parameters + ---------- + path : Path or str + Root supplied to :meth:`save_checkpoint`. + epoch : int + Checkpoint index used when saving. All models load from this index. + device : str or torch.device, optional, default="cpu" + Device for checkpoint tensor loading. + restore_rng : bool, optional, default=False + Restore captured random states after all model and state loading. + + Returns + ------- + dict + Progress dictionary passed to :meth:`save_checkpoint`. + """ + if epoch < 0: + raise ValueError("epoch must be nonnegative.") + rank = self._rank() + root = f"{str(path).rstrip('/')}/step-{epoch}" + state_path = f"{root}/state/rank-{rank}" + fs = fsspec.filesystem(fsspec.utils.get_protocol(root)) + entries = list(self._entries(root, rank)) + error = None + metadata = {} + try: + state_file, _ = _checkpoint._resolve_checkpoint_files( + state_path, fs, {}, epoch + ) + if state_file is None: + raise FileNotFoundError( + f"Missing complete bundle state at {state_path}." + ) + _checkpoint.load_checkpoint( + state_path, epoch=epoch, metadata_dict=metadata, device=device + ) + if metadata.get("bundle", {}).get("schema_version") != 1: + raise ValueError("Unsupported checkpoint bundle schema version.") + if metadata["bundle"] != self._schema(): + raise ValueError( + "Checkpoint bundle roles, associations, or world size do not match." + ) + if restore_rng and "rng" not in metadata: + raise ValueError("This checkpoint does not contain random states.") + for directory, _, named, distributed in entries: + state_file, model_files = _checkpoint._resolve_checkpoint_files( + directory, fs, named, epoch, distributed=distributed + ) + if state_file is None: + raise FileNotFoundError(f"Missing role checkpoint at {directory}.") + _checkpoint._validate_checkpoint_files(state_file, model_files) + except Exception as exc: + # Deserialization failures must reach every rank before collectives. + error = (type(exc).__name__, str(exc)) + errors = [error] + if torch.distributed.is_initialized(): + errors = [None] * torch.distributed.get_world_size() + torch.distributed.all_gather_object(errors, error) + if any(errors): + exception = ( + FileNotFoundError + if all( + value is None or value[0] == "FileNotFoundError" for value in errors + ) + else ValueError + ) + raise exception(f"Invalid checkpoint bundle: {errors}") + for directory, kwargs, _, _ in entries: + _checkpoint.load_checkpoint(directory, epoch=epoch, device=device, **kwargs) + for name, value in self.stateful.items(): + value.load_state_dict(metadata["stateful"][name]) + if restore_rng: + rng = metadata["rng"] + random.setstate(rng["python"]) + state = rng["numpy"] + np.random.set_state( + (state[0], np.asarray(state[1], dtype=np.uint32), *state[2:]) + ) + torch.set_rng_state(rng["torch"].cpu()) + if rng["cuda"]: + torch.cuda.set_rng_state_all([value.cpu() for value in rng["cuda"]]) + return metadata["progress"] diff --git a/test/diffusion/test_distillation_checkpoint.py b/test/diffusion/test_distillation_checkpoint.py new file mode 100644 index 0000000000..e4b49fa433 --- /dev/null +++ b/test/diffusion/test_distillation_checkpoint.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +r"""Roundtrip named training roles through ordinary PhysicsNeMo checkpoints.""" + +import random + +import numpy as np +import pytest +import torch + +from physicsnemo.diffusion.distillation.checkpoint import CheckpointStateBundle + +from .helpers import StructuredDiffusionModel + + +def test_bundle_roundtrip(tmp_path): + models = { + name: StructuredDiffusionModel() for name in ("student", "fake_score", "ema") + } + optimizers = { + name: torch.optim.Adam(model.parameters(), lr=0.01) + for name, model in models.items() + if name != "ema" + } + schedulers = { + name: torch.optim.lr_scheduler.StepLR(optimizer, step_size=1) + for name, optimizer in optimizers.items() + } + for name, optimizer in optimizers.items(): + models[name].gain.square().backward() + optimizer.step() + schedulers[name].step() + bundle = CheckpointStateBundle(models, optimizers=optimizers, schedulers=schedulers) + bundle.save_checkpoint(tmp_path, epoch=4, progress={"step": 4}, include_rng=True) + expected_rng = (torch.rand(3), np.random.rand(3), random.random()) + restored = {name: StructuredDiffusionModel() for name in models} + restored_opts = { + name: torch.optim.Adam(restored[name].parameters(), lr=0.1) + for name in optimizers + } + restored_sched = { + name: torch.optim.lr_scheduler.StepLR(optimizer, step_size=1) + for name, optimizer in restored_opts.items() + } + fresh = CheckpointStateBundle( + restored, optimizers=restored_opts, schedulers=restored_sched + ) + progress = fresh.load_checkpoint(tmp_path, epoch=4, restore_rng=True) + assert progress == {"step": 4} + torch.testing.assert_close(torch.rand(3), expected_rng[0]) + np.testing.assert_array_equal(np.random.rand(3), expected_rng[1]) + assert random.random() == expected_rng[2] + for name in models: + torch.testing.assert_close(restored[name].gain, models[name].gain) + for name in optimizers: + torch.testing.assert_close( + restored_opts[name].state_dict(), optimizers[name].state_dict() + ) + assert restored_sched[name].state_dict() == schedulers[name].state_dict() + + +def test_bundle_role_mismatch_before_loading(tmp_path): + original = StructuredDiffusionModel() + CheckpointStateBundle({"student": original}).save_checkpoint(tmp_path, epoch=1) + other = StructuredDiffusionModel() + with torch.no_grad(): + other.gain.fill_(1.5) + with pytest.raises(ValueError, match="roles"): + CheckpointStateBundle({"fake_score": other}).load_checkpoint(tmp_path, epoch=1) + assert other.gain == 1.5 + + +def test_bundle_missing_checkpoint(tmp_path): + bundle = CheckpointStateBundle({"student": StructuredDiffusionModel()}) + with pytest.raises(FileNotFoundError): + bundle.load_checkpoint(tmp_path, epoch=3) + + +def test_bundle_validates_optimizer_association(): + student, other = StructuredDiffusionModel(), StructuredDiffusionModel() + with pytest.raises(ValueError, match="parameters"): + CheckpointStateBundle( + {"student": student}, + optimizers={"student": torch.optim.Adam(other.parameters())}, + ) + + +def test_bundle_rejects_overwrite(tmp_path): + bundle = CheckpointStateBundle({"student": StructuredDiffusionModel()}) + bundle.save_checkpoint(tmp_path, epoch=0) + with pytest.raises(FileExistsError): + bundle.save_checkpoint(tmp_path, epoch=0) + + +def test_bundle_partial_snapshot_does_not_load_models(tmp_path): + model = StructuredDiffusionModel() + bundle = CheckpointStateBundle({"student": model}) + bundle.save_checkpoint(tmp_path, epoch=0) + checkpoint = next((tmp_path / "step-0" / "models").rglob("*.pt")) + checkpoint.unlink() + with torch.no_grad(): + model.gain.fill_(3.0) + with pytest.raises(FileNotFoundError): + bundle.load_checkpoint(tmp_path, epoch=0) + assert model.gain == 3.0 + + +def test_bundle_requires_captured_rng(tmp_path): + bundle = CheckpointStateBundle({"student": StructuredDiffusionModel()}) + bundle.save_checkpoint(tmp_path, epoch=0) + with pytest.raises(ValueError, match="random states"): + bundle.load_checkpoint(tmp_path, epoch=0, restore_rng=True) + + +def test_bundle_corrupt_metadata_reports_validation_error(tmp_path): + bundle = CheckpointStateBundle({"student": StructuredDiffusionModel()}) + bundle.save_checkpoint(tmp_path, epoch=0) + checkpoint = next((tmp_path / "step-0" / "state").rglob("*.pt")) + checkpoint.write_bytes(checkpoint.read_bytes()[:32]) + with pytest.raises(ValueError, match="Invalid checkpoint bundle"): + bundle.load_checkpoint(tmp_path, epoch=0) + + +@pytest.mark.parametrize("phase", ("student", "auxiliary")) +def test_alternating_resume(tmp_path, phase): + pytest.importorskip("fastgen") + from physicsnemo.diffusion.distillation import DistillationLoss + from physicsnemo.diffusion.noise_schedulers import RectifiedFlowNoiseScheduler + + from .helpers import FeatureDiscriminator, extract_teacher_features + + def setup(): + roles = { + name: StructuredDiffusionModel() + for name in ("student", "teacher", "fake_score") + } + roles["discriminator"] = FeatureDiscriminator() + loss_fn = DistillationLoss( + roles["student"], + RectifiedFlowNoiseScheduler(), + method="f-distill", + teacher=roles["teacher"], + auxiliary_models={ + name: roles[name] for name in ("fake_score", "discriminator") + }, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + ) + optimizers = { + name: torch.optim.Adam(model.parameters(), lr=0.01) + for name, model in roles.items() + if name != "teacher" + } + bundle = CheckpointStateBundle( + roles, optimizers=optimizers, stateful={"loss": loss_fn} + ) + return roles, optimizers, loss_fn, bundle + + def update(roles, optimizers, loss_fn, phase, x): + active = ("student",) if phase == "student" else ("fake_score", "discriminator") + for optimizer in optimizers.values(): + optimizer.zero_grad(set_to_none=True) + loss = loss_fn(x, phase=phase)["total_loss"] + loss.backward() + for name in active: + optimizers[name].step() + return loss.detach() + + roles, optimizers, loss_fn, bundle = setup() + x = torch.randn(2, 3, 5) + for initial in ("student", "auxiliary"): + update(roles, optimizers, loss_fn, initial, x) + bundle.save_checkpoint(tmp_path, epoch=2, progress={"step": 2}, include_rng=True) + expected = update(roles, optimizers, loss_fn, phase, x) + fresh, fresh_opts, fresh_loss, fresh_bundle = setup() + assert fresh_bundle.load_checkpoint(tmp_path, epoch=2, restore_rng=True) == { + "step": 2 + } + actual = update(fresh, fresh_opts, fresh_loss, phase, x) + torch.testing.assert_close(actual, expected) + for name, model in roles.items(): + torch.testing.assert_close(fresh[name].state_dict(), model.state_dict()) + for name, optimizer in optimizers.items(): + torch.testing.assert_close( + fresh_opts[name].state_dict(), optimizer.state_dict() + ) + torch.testing.assert_close(fresh_loss.state_dict(), loss_fn.state_dict())