From 78c4e286d741454b91912bede3336c586d30d785 Mon Sep 17 00:00:00 2001 From: Charlelie Laurent Date: Fri, 16 Oct 2026 13:14:32 -0700 Subject: [PATCH] Add optional diffusion distillation strategy Signed-off-by: Charlelie Laurent --- CHANGELOG.md | 10 + docs/api/diffusion/distillation.rst | 94 ++- .../diffusion/distillation/__init__.py | 9 + .../diffusion/distillation/_alchemi.py | 615 ++++++++++++++++++ .../diffusion/distillation/_fastgen.py | 22 + physicsnemo/diffusion/distillation/losses.py | 4 + .../diffusion/distillation/strategy.py | 41 ++ pyproject.toml | 3 + test/diffusion/test_distillation_strategy.py | 462 +++++++++++++ uv.lock | 95 ++- 10 files changed, 1353 insertions(+), 2 deletions(-) create mode 100644 physicsnemo/diffusion/distillation/_alchemi.py create mode 100644 physicsnemo/diffusion/distillation/strategy.py create mode 100644 test/diffusion/test_distillation_strategy.py diff --git a/CHANGELOG.md b/CHANGELOG.md index bc8d798286..94ab88b30f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Adds an experimental, independently implemented SMART operator with sampled geometry encoding, reusable latent contexts, and chunked field decoding. +- Adds alpha distillation losses and an optional strategy layer in + `physicsnemo.diffusion.distillation` for sCM, MeanFlow, DMD2, and f-distill. + Use `DistillationLoss` or `WeightedDistillationLoss` in your own training + loop, or let `DistillationStrategy` manage the optimizer loop. Define custom + losses through `BaseDistillationLoss` without depending on FastGen. + Use `CheckpointStateBundle` to resume models, optimizers, and loss statistics. + These alpha APIs may change without backward compatibility. Distillation is + opt-in, with optional dependencies available through extras. Existing diffusion + training and inference remain unchanged, and distilled models continue to use + PhysicsNeMo inference APIs. - Adds `DiffusionExtraTimes` and `DiffusionExtraOutputs` to `physicsnemo.diffusion` for utilities that carry named times or auxiliary outputs. Compatible model interfaces and preconditioners can accept diff --git a/docs/api/diffusion/distillation.rst b/docs/api/diffusion/distillation.rst index 21a2c2b081..fd55971428 100644 --- a/docs/api/diffusion/distillation.rst +++ b/docs/api/diffusion/distillation.rst @@ -265,11 +265,103 @@ 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. +Optional training strategy +-------------------------- + +``DistillationStrategy`` runs a distillation loss with named optimizers, +alternating phases, gradient accumulation, and checkpoints. The loss object +computes loss values. The strategy manages backward passes and optimizer updates. + +Use the strategy when you want a training loop for an existing loss. Call the +same loss directly when your recipe needs its own loop. Install +``nvidia-physicsnemo[diffusion-strategy]`` to use the strategy. Add +``diffusion-fastgen`` for FastGen losses, or supply a custom loss without FastGen. + +.. code-block:: python + + from physicsnemo.diffusion.distillation import DistillationStrategy + + strategy = DistillationStrategy( + models={"student": student, "teacher": teacher, "fake_score": fake_score}, + loss_fn=loss_fn, + optimizer_configs={ + role: [{"optimizer_cls": torch.optim.Adam, + "optimizer_kwargs": {"lr": 1e-4}}] + for role in ("student", "fake_score") + }, + num_steps=10_000, + devices=[torch.device("cuda")], + accumulation_steps=2, + ) + strategy.run(batches) # TensorDict batches with x0 and optional loss arguments + strategy.save_checkpoint("checkpoints", epoch=10_000, include_rng=True) + +Call ``train_batch(batch)`` when your recipe controls batch iteration. Each +batch produces one loss. The strategy computes gradients through the mean of +``total_loss``, divided by ``accumulation_steps``. A custom ``training_fn(models, batch)`` can +adapt batch fields to loss arguments. Let the loss call the models. The loss +receives ``phase``, ``step``, and ``update_state=True`` from the strategy. Reserve +these batch keys and ``return_outputs`` for the strategy. + +For explicit MeanFlow times, keep ``t`` and ``r`` as separate batch tensors. +Assemble the time tuple in ``training_fn`` because a TensorDict cannot store +this tuple as a tensor field: + +.. code-block:: python + + from physicsnemo.diffusion import DiffusionExtraTimes + + def meanflow_arguments(models, batch): + return { + "x0": batch["x0"], + "t": (batch["t"], DiffusionExtraTimes(r=batch["r"], batch_size=batch.batch_size)), + } + + strategy = DistillationStrategy( + models={"student": student}, loss_fn=loss_fn, + optimizer_configs=optimizer_configs, num_steps=1000, + training_fn=meanflow_arguments, + ) + +Single-phase training updates the student. With auxiliary optimizer roles, the +default repeats one student update and four auxiliary updates. Supply +``phase_fn(step)`` and ``phase_roles`` to change that policy. Steps count completed +optimizer updates, not batches. Accumulation keeps a fixed phase. Inactive +optimizers and their schedulers do not advance. Custom losses must control +gradient flow through their model calls. Phase selection does not freeze +inactive models. With ``precision=torch.float16``, gradient scaling skips the +whole active phase if any active optimizer has infinite or NaN gradients. +``precision=torch.bfloat16`` enables autocast without scaling. Use update hooks +for policies such as gradient clipping. + +Set ``num_steps`` to the desired count of completed optimizer updates. +Or set ``num_epochs`` with a data loader that has a length. +Its total batch count must divide evenly by ``accumulation_steps``. +Skipped updates can require extra epochs to finish the requested updates. + +Compile and wrap your models before constructing the loss and strategy. +The strategy preserves your choices of training modes and frozen parameters. +If a setup hook wraps a model used by a custom loss, supply +``bind_models_fn(loss_fn, models)`` to update that loss's model references. +Wrap before construction if wrapping replaces parameters. + +To track an EMA, initialize its target model, add it to ``models`` as ``ema``, +and set ``ema_models={"student": "ema"}``. The strategy updates it after each +successful student update. Keep inference and validation in your recipe. + +The strategy uses ``CheckpointStateBundle`` for all roles, optimizers, schedulers, +loss statistics, counters, EMA targets, and AMP scaler state. Hooks with both +``state_dict`` and ``load_state_dict`` contribute their state. Reconstruct matching +hooks, settings, and phase callbacks before calling ``load_checkpoint``. +Save checkpoints after a batch that completes an optimizer update, not during +accumulation. Add your own data-loader cursor checkpoint if you need exact +batch resumption. + Restarting training ------------------- ``CheckpointStateBundle`` groups named models, optimizers, and loss statistics -for a training restart. A single-model checkpoint stores one model; a bundle +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. diff --git a/physicsnemo/diffusion/distillation/__init__.py b/physicsnemo/diffusion/distillation/__init__.py index 67ead049ba..f8f22a5c0b 100644 --- a/physicsnemo/diffusion/distillation/__init__.py +++ b/physicsnemo/diffusion/distillation/__init__.py @@ -16,6 +16,7 @@ r"""Alpha diffusion distillation APIs with optional training backends.""" +import importlib import warnings from physicsnemo.core.warnings import ExperimentalFeatureWarning @@ -36,4 +37,12 @@ "LossMap", "DistillationLoss", "WeightedDistillationLoss", + "DistillationStrategy", ] + + +def __getattr__(name: str): + r"""Load the optional training strategy only when requested.""" + if name == "DistillationStrategy": + return importlib.import_module(f"{__name__}.strategy").DistillationStrategy + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/physicsnemo/diffusion/distillation/_alchemi.py b/physicsnemo/diffusion/distillation/_alchemi.py new file mode 100644 index 0000000000..6550b1a7b4 --- /dev/null +++ b/physicsnemo/diffusion/distillation/_alchemi.py @@ -0,0 +1,615 @@ +# 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"""Optional diffusion orchestration through the general Alchemi training loop.""" + +import importlib +import math +from collections.abc import Callable, Mapping +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import torch +from tensordict import TensorDict + +from .checkpoint import CheckpointStateBundle + +_backend = importlib.import_module("nvalchemi.training.strategy") +_TrainingStrategy = _backend.TrainingStrategy +_UpdateHook = _backend.TrainingUpdateHook +Field, PrivateAttr = _backend.Field, _backend.PrivateAttr +field_validator, model_validator = _backend.field_validator, _backend.model_validator + +_COUNTERS = ( + "step_count", + "global_step_count", + "batch_count", + "epoch_count", + "epoch_step_count", +) + + +def _batch_arguments( + models: Mapping[str, torch.nn.Module], batch: TensorDict +) -> TensorDict: + r"""Pass loss arguments without evaluating any model before the loss.""" + return batch + + +class _Accumulation(_UpdateHook): + r"""Keep one optimization phase across a complete accumulation window.""" + + priority = 10 + + def __call__(self, ctx: Any, stage: Any, will_skip: bool): + r"""Scale the backward loss and veto intermediate updates.""" + strategy = ctx.workflow + if stage == _backend.TrainingStage.BEFORE_BATCH: + return strategy._micro_step == 0, ctx.loss + if stage == _backend.TrainingStage.DO_BACKWARD: + return True, ctx.loss / strategy.accumulation_steps + if stage == _backend.TrainingStage.DO_OPTIMIZER_STEP: + return strategy._micro_step + 1 == strategy.accumulation_steps, ctx.loss + return True, ctx.loss + + +class _JointScaledUpdate(_UpdateHook): + r"""Skip the entire phase if scaling detects a nonfinite active gradient.""" + + priority = 25 + + def __call__(self, ctx: Any, stage: Any, will_skip: bool): + r"""Coordinate fake-score and discriminator updates after unscaling.""" + if ( + stage == _backend.TrainingStage.DO_OPTIMIZER_STEP + and not will_skip + and ctx.grad_scaler is not None + ): + found_inf = [ + value + for optimizer in ctx.optimizers + for value in ctx.grad_scaler._found_inf_per_device(optimizer).values() + ] + if any(bool(value.item()) for value in found_inf): + ctx.grad_scaler.update() + return False, ctx.loss + return True, ctx.loss + + +class DistillationStrategy(_TrainingStrategy): + r"""Train diffusion models with a loss and named optimizers. + + The strategy manages alternating phases, gradient accumulation, and + checkpoints. Use it to run a distillation loss without writing an optimizer + loop, or call the same loss directly when your recipe needs a custom loop. + Install ``nvidia-physicsnemo[diffusion-strategy]``, supply your models and + optimizer settings, and pass batches to ``run``. Custom losses do not + require FastGen. + + Parameters + ---------- + models : mapping + Models keyed by name, including ``student`` and optional + ``teacher``, ``fake_score``, ``discriminator``, or EMA targets. + loss_fn : callable + Loss accepting batch fields, ``phase``, ``step``, and ``update_state=True``. + Return a dictionary containing ``total_loss``. The strategy computes gradients + through its mean, divided by ``accumulation_steps``. + optimizer_configs : mapping + Per-role lists of dictionaries with ``optimizer_cls`` and + ``optimizer_kwargs``, plus optional scheduler settings. + num_steps, num_epochs : int, optional + Supply exactly one training duration. Steps count completed updates. + phase_fn : callable, optional + Map completed update count to a phase name. The default trains only the + student, or one student update followed by four auxiliary updates. + phase_roles : mapping, optional + Optimizer roles to update in each phase, for example + ``{"student": ("student",), "auxiliary": ("fake_score",)}``. + This controls optimizer and scheduler steps, not gradient computation. + accumulation_steps : int, optional, default=1 + Batches per update. Phase selection stays fixed during accumulation. + bind_models_fn : callable, optional + Callback ``fn(loss_fn, models)`` that updates model references in a + custom loss if a setup hook wraps or replaces a model. + Built-in losses handle this automatically. + precision : torch.dtype, optional + Autocast precision. ``torch.float16`` also enables gradient scaling and + skips all active optimizers if any active gradient is infinite or NaN. + ema_models : mapping, optional + Models to average and their target names, such as + ``{"student": "ema"}``. Initialize the targets and include them in ``models``. + ema_decay : float, optional, default=0.999 + Decay applied only after a successful update of the source role. + hooks, devices : sequence, optional + Training hooks and devices following the underlying training strategy. + + Examples + -------- + This strategy trains one student model for a single batch with Adam. + The batch supplies clean samples, and the final check confirms one optimizer step. + + >>> import torch + >>> from tensordict import TensorDict + >>> model = torch.nn.Linear(3, 3) + >>> def loss_fn(x0, **kwargs): + ... return {"total_loss": model(x0).square().mean()} + >>> # Use the same model key for the optimizer configuration. + >>> strategy = DistillationStrategy( + ... models={"student": model}, loss_fn=loss_fn, num_steps=1, + ... optimizer_configs={"student": [{"optimizer_cls": torch.optim.Adam, + ... "optimizer_kwargs": {"lr": 1e-3}}]}, + ... ) + >>> strategy.train_batch(TensorDict({"x0": torch.randn(2, 3)}, [2])) + >>> strategy.step_count + 1 + + Notes + ----- + Put ``x0`` and optional loss arguments in each ``TensorDict`` batch. + For explicit MeanFlow times, store ``t`` and ``r`` as separate batch + fields. Use ``training_fn(models, batch)`` to return a dictionary of loss + arguments with ``t=(batch["t"], DiffusionExtraTimes(r=batch["r"], + batch_size=batch.batch_size))``. Do not store the tuple inside a TensorDict. + Reserve ``phase``, ``step``, ``update_state``, and ``return_outputs`` for the + strategy. Custom losses control gradient flow through their model calls; + the strategy preserves existing frozen parameters and training modes. + Keep inference and validation in your recipe. + + Compile and wrap models before constructing the loss and strategy. + With ``num_epochs``, the total batch budget must divide evenly by + ``accumulation_steps``. Save only at completed accumulation boundaries. + Recreate matching settings, hooks, and phase callbacks before loading. + Checkpoint your data-loader cursor separately for exact batch resumption. + """ + + models: dict[str, torch.nn.Module] + loss_fn: Any + training_fn: Any = _batch_arguments + phase_fn: Callable[[int], str] | None = None + phase_roles: dict[str, tuple[str, ...]] = Field(default_factory=dict) + accumulation_steps: int = Field(default=1, ge=1) + bind_models_fn: Callable[..., None] | None = None + precision: torch.dtype | None = None + ema_models: dict[str, str] = Field(default_factory=dict) + ema_decay: float = Field(default=0.999, ge=0, le=1) + _micro_step: int = PrivateAttr(default=0) + _phase: str = PrivateAttr(default="student") + _bound_models: dict[str, torch.nn.Module] = PrivateAttr(default_factory=dict) + _model_modes: dict[torch.nn.Module, bool] = PrivateAttr(default_factory=dict) + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + + def __exit__(self, exc_type, exc, tb) -> None: + r"""Close autocast contexts while keeping scaler state across updates.""" + scalers = [] + for hook in self.hooks: + children = ( + hook._hooks + if isinstance(hook, _backend.TrainingUpdateOrchestrator) + else [hook] + ) + scalers.extend( + (child, child._scaler) + for child in children + if isinstance(child, _backend.MixedPrecisionHook) + ) + try: + super().__exit__(exc_type, exc, tb) + finally: + for hook, scaler in scalers: + hook._scaler = scaler + + @model_validator(mode="before") + @classmethod + def _normalize_inputs(cls, data: Any) -> Any: + r"""Accept native modules without atomic-model or target requirements.""" + if not isinstance(data, dict): + return data + values = dict(data) + values["single_model_input"] = False + hooks = list(values.get("hooks", ())) + hooks.append(_Accumulation()) + if values.get("precision") is not None: + hooks.append(_backend.MixedPrecisionHook(precision=values["precision"])) + hooks.append(_JointScaledUpdate()) + values["hooks"] = hooks + return values + + @field_validator("loss_fn", mode="before") + @classmethod + def _normalize_loss_fn(cls, value: Any) -> Any: + r"""Accept structural diffusion losses without an atomic loss wrapper.""" + if not callable(value): + raise ValueError("loss_fn must be callable.") + return value + + @model_validator(mode="after") + def _validate_strategy(self): + r"""Validate diffusion roles while retaining backend hook ownership checks.""" + if (self.num_steps is None) == (self.num_epochs is None): + raise ValueError("Supply exactly one of num_steps or num_epochs.") + if "student" not in self.models or not self.optimizer_configs: + raise ValueError("Configure a student model and named optimizers.") + if set(self.optimizer_configs) - set(self.models) or any( + not value for value in self.optimizer_configs.values() + ): + raise ValueError( + "Each optimizer configuration must name a model and contain a nonempty list." + ) + if not self.devices or len(self.devices) not in (1, len(self.models)): + raise ValueError("devices must contain one device or one per model.") + if self.validation_config is not None: + raise ValueError( + "Diffusion validation is caller-owned; omit validation_config." + ) + if not self.phase_roles: + self.phase_roles = {"student": ("student",)} + auxiliary = tuple( + role + for role in ("fake_score", "discriminator") + if role in self.optimizer_configs + ) + if auxiliary: + self.phase_roles["auxiliary"] = auxiliary + for roles in self.phase_roles.values(): + if not roles or set(roles) - set(self.optimizer_configs): + raise ValueError("Every phase must name configured optimizer roles.") + for source, target in self.ema_models.items(): + if ( + source not in self.optimizer_configs + or target not in self.models + or target in self.optimizer_configs + ): + raise ValueError( + "EMA requires an optimized source and a non-optimized target role." + ) + if len({id(hook) for hook in self.hooks}) != len(self.hooks): + raise ValueError("Hook instances must be distinct.") + _backend._validate_single_do_claimants(self.hooks) + _backend._validate_hook_dependencies(self.hooks) + return self + + def model_post_init(self, __context: Any) -> None: + r"""Initialize the common loop without atomic target assembly.""" + self._init_hooks(list(self.hooks)) + self._refresh_hook_claim_flags() + self._last_batch = self._last_losses = self._last_loss = None + self._optimizers, self._lr_schedulers = [], [] + self._runtime_optimizers = [] + self._target_keys = () + self._bound_models = dict(self.models) + + def _run_setup_hooks(self, dataloader: Any = None) -> Any: + r"""Bind final execution wrappers before building optimizers or computing loss.""" + dataloader = super()._run_setup_hooks(dataloader) + changed = any( + self._bound_models.get(role) is not model + for role, model in self.models.items() + ) + bind = self.bind_models_fn or getattr( + self.loss_fn, "_bind_execution_models", None + ) + if self.bind_models_fn is not None: + bind(self.loss_fn, self.models) + elif bind is not None: + bind(self.models) + elif changed: + raise ValueError( + "Setup changed model wrappers; supply bind_models_fn for this custom loss." + ) + if changed and self._runtime_optimizers: + roles = [ + role + for role, configs in self.optimizer_configs.items() + for _ in configs + ] + for role, optimizer in zip(roles, self._optimizers, strict=True): + parameters = set(self.models[role].parameters()) + if any( + value not in parameters + for group in optimizer.param_groups + for value in group["params"] + ): + raise ValueError( + "Setup replaced parameters after optimizer construction; create a fresh strategy." + ) + self._bound_models = dict(self.models) + return dataloader + + @contextmanager + def _preserve_modes(self): + r"""Preserve caller-selected mixed child modes around backend contexts.""" + self._model_modes = { + module: module.training + for model in self.models.values() + for module in model.modules() + } + try: + yield + finally: + for module, training in self._model_modes.items(): + module.training = training + self._model_modes = {} + + def train_batch(self, batch: TensorDict) -> None: + r"""Compute one batch loss and accumulate its gradients. + + Use this method when your recipe controls batch iteration. Pass loss + arguments in a ``TensorDict``. An optimizer update occurs only after + ``accumulation_steps`` batches, with the same phase throughout. + """ + with self._preserve_modes(): + super().train_batch(batch) + + def run(self, dataloader: Any) -> None: + r"""Run batch iteration until the configured update budget completes. + + Pass an iterable of ``TensorDict`` batches. Set ``num_steps`` for a + completed update count, or ``num_epochs`` with a data loader that has + a length. Skipped updates can require extra epochs. + """ + self._resume_optimizer_state = bool(self._runtime_optimizers) + with self._preserve_modes(): + super().run(dataloader) + + def _resolve_target_step_count(self, batches_per_epoch: int | None) -> int: + r"""Convert epoch-sized batch budgets to completed accumulation windows.""" + target = super()._resolve_target_step_count(batches_per_epoch) + if self.num_steps is not None: + return target + if target % self.accumulation_steps: + raise ValueError( + "The epoch batch budget must divide evenly into accumulation windows; use num_steps otherwise." + ) + return math.ceil(target / self.accumulation_steps) + + def _train_batch_with_optimizers( + self, batch: TensorDict, flat_opts: list, flat_scheds: list + ) -> None: + r"""Select only active optimizers without changing the underlying step loop.""" + if self._micro_step == 0: + self._phase = ( + self.phase_fn(self.step_count) + if self.phase_fn is not None + else ( + "auxiliary" + if "auxiliary" in self.phase_roles and self.step_count % 5 + else "student" + ) + ) + if self._phase not in self.phase_roles: + raise ValueError(f"Unknown phase {self._phase!r}.") + roles = [ + role for role, configs in self.optimizer_configs.items() for _ in configs + ] + active = [ + i for i, role in enumerate(roles) if role in self.phase_roles[self._phase] + ] + for module, training in self._model_modes.items(): + module.training = training + try: + super()._train_batch_with_optimizers( + batch, [flat_opts[i] for i in active], [flat_scheds[i] for i in active] + ) + finally: + self._optimizers, self._lr_schedulers = flat_opts, flat_scheds + + def _run_optimizer_step_phase( + self, batch: TensorDict, flat_opts: list, flat_scheds: list + ) -> bool: + r"""Complete EMA and accumulation bookkeeping before after-batch hooks.""" + updated = super()._run_optimizer_step_phase(batch, flat_opts, flat_scheds) + self._micro_step = (self._micro_step + 1) % self.accumulation_steps + if updated: + with torch.no_grad(): + for source, target in self.ema_models.items(): + if source not in self.phase_roles[self._phase]: + continue + source_values = ( + *self.models[source].parameters(), + *self.models[source].buffers(), + ) + target_values = ( + *self.models[target].parameters(), + *self.models[target].buffers(), + ) + for value, average in zip( + source_values, target_values, strict=True + ): + if average.is_floating_point(): + average.lerp_(value, 1 - self.ema_decay) + else: + average.copy_(value) + return updated + + def _compute_losses( + self, + predictions: Mapping[str, Any], + batch: TensorDict, + *, + step: int, + epoch: int, + ) -> dict: + r"""Evaluate the loss once and translate only its reporting structure.""" + arguments = dict(predictions.items()) + if set(arguments) & {"phase", "step", "update_state", "return_outputs"}: + raise ValueError( + "The strategy owns phase, step, update_state, and return_outputs." + ) + losses = self.loss_fn( + **arguments, phase=self._phase, step=step, update_state=True + ) + return { + "total_loss": losses["total_loss"].mean(), + "per_component_unweighted": dict(losses), + "per_component_sample": dict(losses), + "per_component_weight": {key: 1.0 for key in losses}, + "per_component_raw_weight": {key: 1.0 for key in losses}, + } + + def state_dict(self) -> dict[str, Any]: + r"""Return counters, phase settings, and hook and gradient scaler state. + + Call only at a completed accumulation boundary. Use + ``checkpoint_bundle`` or ``save_checkpoint`` to include models, + optimizers, schedulers, and loss state in the same checkpoint. + """ + if self._micro_step: + raise ValueError("Save only at completed accumulation boundaries.") + hooks = [] + for hook in self.hooks: + hooks.extend( + hook._hooks + if isinstance(hook, _backend.TrainingUpdateOrchestrator) + else [hook] + ) + state = { + "schema_version": 1, + "counters": {name: getattr(self, name) for name in _COUNTERS}, + } + state["hook_types"] = [ + f"{type(hook).__module__}.{type(hook).__qualname__}" for hook in hooks + ] + state["phase_roles"] = self.phase_roles + state["accumulation_steps"] = self.accumulation_steps + state["hooks"] = { + str(i): hook.state_dict() + for i, hook in enumerate(hooks) + if callable(getattr(hook, "state_dict", None)) + and callable(getattr(hook, "load_state_dict", None)) + } + state["scalers"] = { + str(i): hook._scaler.state_dict() + for i, hook in enumerate(hooks) + if isinstance(hook, _backend.MixedPrecisionHook) + and hook._scaler is not None + } + return state + + def load_state_dict(self, state: Mapping[str, Any]) -> None: + r"""Restore counters and hook and gradient scaler state. + + Recreate matching hooks, phase roles, and accumulation settings before + passing a dictionary returned by ``state_dict``. Recreate the phase + callback separately because this state does not include it. + """ + if state.get("schema_version") != 1: + raise ValueError("Unsupported strategy state schema.") + hooks = [] + for hook in self.hooks: + hooks.extend( + hook._hooks + if isinstance(hook, _backend.TrainingUpdateOrchestrator) + else [hook] + ) + if ( + state["hook_types"] + != [f"{type(hook).__module__}.{type(hook).__qualname__}" for hook in hooks] + or state["phase_roles"] != self.phase_roles + or state["accumulation_steps"] != self.accumulation_steps + ): + raise ValueError( + "Reconstruct matching hooks, phase roles, and accumulation settings before loading." + ) + for name in _COUNTERS: + setattr(self, name, int(state["counters"][name])) + for index, value in state["hooks"].items(): + hooks[int(index)].load_state_dict(value) + for index, value in state["scalers"].items(): + hook = hooks[int(index)] + hook._scaler = torch.amp.GradScaler(self.devices[0].type) + hook._scaler.load_state_dict(value) + self._micro_step = 0 + self._resume_optimizer_state = True + + def checkpoint_bundle(self) -> CheckpointStateBundle: + r"""Collect models, optimizers, schedulers, and training state. + + Use the returned ``CheckpointStateBundle`` with a custom checkpoint + workflow. It includes this strategy and any loss exposing + ``state_dict``, and names optimizers by role and index. + """ + self._setup_runtime_optimizers() + names = [ + (f"{role}-{i}", role) + for role, configs in self.optimizer_configs.items() + for i, _ in enumerate(configs) + ] + stateful = {"strategy": self} + if callable(getattr(self.loss_fn, "state_dict", None)): + stateful["loss"] = self.loss_fn + return CheckpointStateBundle( + self.models, + optimizers={ + name: record.optimizer + for (name, _), record in zip( + names, self._runtime_optimizers, strict=True + ) + }, + optimizer_models=dict(names), + schedulers={ + name: record.scheduler + for (name, _), record in zip( + names, self._runtime_optimizers, strict=True + ) + if record.scheduler is not None + }, + stateful=stateful, + ) + + def save_checkpoint( + self, path: Path | str, *, epoch: int, include_rng: bool = False + ) -> None: + r"""Save models and training state to a checkpoint. + + Call at a completed accumulation boundary with the checkpoint ``path`` + and ``epoch`` identifier. Set ``include_rng=True`` to save random number + generator state. Data-loader position requires a separate checkpoint. + """ + self.checkpoint_bundle().save_checkpoint( + path, epoch=epoch, include_rng=include_rng + ) + + def load_checkpoint( + self, path: Path | str, *, epoch: int, restore_rng: bool = False + ) -> None: + r"""Restore models and training state from a checkpoint. + + Recreate matching models, loss, settings, hooks, and phase callbacks, + then pass the saved ``path`` and ``epoch`` identifier. Set + ``restore_rng=True`` to restore saved random number generator state. + """ + with self._preserve_modes(), self: + self._prepare_setup_hooks() + self._validate_runtime_devices() + self.models = _backend.move_to_devices(self.models, self.devices) + self._run_setup_hooks() + self.checkpoint_bundle().load_checkpoint( + path, epoch=epoch, device=self.devices[0], restore_rng=restore_rng + ) + + def to_spec_dict(self) -> dict[str, Any]: + r"""Raise ``NotImplementedError`` for configuration serialization. + + Reconstruct diffusion components explicitly and use ``load_checkpoint`` + to restore training state. + """ + raise NotImplementedError( + "Reconstruct diffusion components explicitly, then use load_checkpoint." + ) diff --git a/physicsnemo/diffusion/distillation/_fastgen.py b/physicsnemo/diffusion/distillation/_fastgen.py index 0bd4a46413..b4b2bac289 100644 --- a/physicsnemo/diffusion/distillation/_fastgen.py +++ b/physicsnemo/diffusion/distillation/_fastgen.py @@ -79,6 +79,28 @@ def residual( class _FastGenBinding: r"""Connect student and teacher models and collect batch loss outputs.""" + def _bind_execution_models(self, models: Mapping[str, torch.nn.Module]) -> None: + r"""Rebind wrappers only when they keep the configured role parameters.""" + for role, attribute in ( + ("student", "net"), + ("teacher", "teacher"), + ("fake_score", "fake_score"), + ("discriminator", "discriminator"), + ): + adapter = getattr(self, attribute, None) + if adapter is None: + continue + if role not in models: + raise ValueError(f"Missing configured loss role {role!r}.") + replacement = models[role] + if {id(value) for value in adapter.model.parameters()} != { + id(value) for value in replacement.parameters() + }: + raise ValueError( + f"Execution wrapper for {role!r} must preserve parameter identities; wrap before constructing the loss otherwise." + ) + adapter.model = replacement + def _jvp(self, *args: Any, **kwargs: Any) -> Float[Tensor, " B *dims"]: r"""Use differentiable attention kernels for backend tangent calculations.""" with sdpa_kernel(SDPBackend.MATH): diff --git a/physicsnemo/diffusion/distillation/losses.py b/physicsnemo/diffusion/distillation/losses.py index 966b4d3578..abe5ae495e 100644 --- a/physicsnemo/diffusion/distillation/losses.py +++ b/physicsnemo/diffusion/distillation/losses.py @@ -137,6 +137,10 @@ def __init__( **auxiliary_options, ) + def _bind_execution_models(self, models: Mapping[str, torch.nn.Module]) -> None: + r"""Rebind trainer execution wrappers while preserving loss configuration.""" + self._implementation._bind_execution_models(models) + def __call__( self, x0: Float[Tensor, " B *dims"], diff --git a/physicsnemo/diffusion/distillation/strategy.py b/physicsnemo/diffusion/distillation/strategy.py new file mode 100644 index 0000000000..cea66ca1fc --- /dev/null +++ b/physicsnemo/diffusion/distillation/strategy.py @@ -0,0 +1,41 @@ +# 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"""Optional training loop for diffusion distillation.""" + +import importlib +from typing import Any + +from physicsnemo.core.version_check import check_version_spec + +_AVAILABLE = check_version_spec("nvalchemi-toolkit", "0.2.0", hard_fail=False) +if _AVAILABLE: + DistillationStrategy = importlib.import_module( + f"{__package__}._alchemi" + ).DistillationStrategy +else: + + class DistillationStrategy: + r"""Install ``diffusion-strategy`` to run a distillation training loop.""" + + def __init__(self, **kwargs: Any) -> None: + raise ImportError( + "PhysicsNeMo DistillationStrategy requires nvalchemi-toolkit>=0.2.0. " + "Install nvidia-physicsnemo[diffusion-strategy] for the tested revision." + ) + + +__all__ = ["DistillationStrategy"] diff --git a/pyproject.toml b/pyproject.toml index 2426814490..03bcd77c0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -266,6 +266,9 @@ dev = [ allow-direct-references = true [project.optional-dependencies] +diffusion-strategy = [ + "nvalchemi-toolkit @ git+https://github.com/NVIDIA/nvalchemi-toolkit.git@f59e28dc6701583a13bb294e536076c2a7379e98", +] diffusion-fastgen = [ "fastgen @ git+https://github.com/CharlelieLrt/FastGen.git@cd7de8c386a92da2be4cbc3bf0558f16153e46b9", ] diff --git a/test/diffusion/test_distillation_strategy.py b/test/diffusion/test_distillation_strategy.py new file mode 100644 index 0000000000..5adc94f2b6 --- /dev/null +++ b/test/diffusion/test_distillation_strategy.py @@ -0,0 +1,462 @@ +# 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"""Diffusion updates through the optional general-purpose training loop.""" + +import subprocess +import sys + +import pytest +import torch +from tensordict import TensorDict + +pytest.importorskip("nvalchemi") + +from nvalchemi.training import TrainingStrategy + +from physicsnemo.core import Module +from physicsnemo.diffusion import DiffusionExtraTimes +from physicsnemo.diffusion.distillation.strategy import DistillationStrategy + +from .helpers import StructuredDiffusionModel, extract_teacher_features + + +class NativeLoss: + def __init__(self, models): + self.models = models + self.calls = [] + + def __call__(self, x0, *, phase, step, t=None, **kwargs): + self.calls.append((phase, step)) + role = "student" if phase == "student" else "fake_score" + t = torch.zeros(x0.shape[0]) if t is None else t + prediction = self.models[role](x0, t)[0] + return {"total_loss": prediction.square().mean()} + + +def make_strategy(*, alternating=False, **kwargs): + roles = ("student", "fake_score") if alternating else ("student",) + models = {name: StructuredDiffusionModel() for name in roles} + loss = NativeLoss(models) + configs = { + name: [ + { + "optimizer_cls": torch.optim.SGD, + "optimizer_kwargs": {"lr": 0.01}, + "scheduler_cls": torch.optim.lr_scheduler.StepLR, + "scheduler_kwargs": {"step_size": 1, "gamma": 0.5}, + } + ] + for name in roles + } + if "num_epochs" not in kwargs and "num_steps" not in kwargs: + kwargs["num_steps"] = 6 + strategy = DistillationStrategy( + models=models, loss_fn=loss, optimizer_configs=configs, **kwargs + ) + return strategy, loss + + +@pytest.mark.parametrize("structured", (False, True)) +def test_native_strategy_one_call_and_direct_update(structured): + def batch_arguments(models, batch): + return { + "x0": batch["x0"], + "t": ( + batch["t"], + DiffusionExtraTimes(r=batch["r"], batch_size=batch.batch_size), + ), + } + + strategy, loss = make_strategy( + **({"training_fn": batch_arguments} if structured else {}) + ) + assert isinstance(strategy, TrainingStrategy) + model = strategy.models["student"] + direct = StructuredDiffusionModel() + direct.load_state_dict(model.state_dict()) + optimizer = torch.optim.SGD(direct.parameters(), lr=0.01) + x = torch.randn(2, 3, 5) + times = ( + ( + torch.tensor([0.3, 0.7]), + DiffusionExtraTimes(r=torch.tensor([0.1, 0.2]), batch_size=[2]), + ) + if structured + else None + ) + reference = NativeLoss({"student": direct})(x, t=times, phase="student", step=0)[ + "total_loss" + ] + reference.backward() + optimizer.step() + batch = TensorDict({"x0": x}, batch_size=[2]) + if times is not None: + batch["t"], batch["r"] = times[0], times[1].r + strategy.train_batch(batch) + assert loss.calls == [("student", 0)] + assert strategy.step_count == 1 + torch.testing.assert_close(model.state_dict(), direct.state_dict()) + + +def test_strategy_phase_and_inactive_scheduler(): + strategy, loss = make_strategy(alternating=True) + batch = TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2]) + fake = strategy.models["fake_score"].gain.detach().clone() + strategy.train_batch(batch) + torch.testing.assert_close(strategy.models["fake_score"].gain, fake) + student = strategy.models["student"].gain.detach().clone() + strategy.train_batch(batch) + torch.testing.assert_close(strategy.models["student"].gain, student) + for _ in range(4): + strategy.train_batch(batch) + assert loss.calls == [ + ("student" if i % 5 == 0 else "auxiliary", i) for i in range(6) + ] + assert [scheduler.last_epoch for scheduler in strategy._lr_schedulers] == [2, 4] + + +def test_strategy_custom_phase(): + strategy, loss = make_strategy( + alternating=True, + phase_fn=lambda step: "auxiliary" if step % 2 == 0 else "student", + ) + batch = TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2]) + strategy.run([batch, batch]) + assert loss.calls == [ + ("auxiliary" if i % 2 == 0 else "student", i) for i in range(6) + ] + + +def test_checkpoint_hook_sees_completed_joint_state(tmp_path): + from nvalchemi.training._stages import TrainingStage + + class SaveHook: + stage = TrainingStage.AFTER_BATCH + frequency = 1 + + def __call__(self, ctx, stage): + if ctx.batch_count == 2: + ctx.workflow.save_checkpoint(tmp_path, epoch=1) + + strategy, _ = make_strategy( + alternating=True, accumulation_steps=2, hooks=[SaveHook()] + ) + batch = TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2]) + strategy.train_batch(batch) + strategy.train_batch(batch) + assert (tmp_path / "step-1" / "state").exists() + + +def test_strategy_resume(tmp_path): + strategy, _ = make_strategy(alternating=True) + batch = TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2]) + strategy.train_batch(batch) + strategy.save_checkpoint(tmp_path, epoch=1, include_rng=True) + strategy.train_batch(batch) + fresh, loss = make_strategy(alternating=True) + fresh.load_checkpoint(tmp_path, epoch=1, restore_rng=True) + fresh.train_batch(batch) + assert loss.calls == [("auxiliary", 1)] + for name in strategy.models: + torch.testing.assert_close( + fresh.models[name].state_dict(), strategy.models[name].state_dict() + ) + for expected, actual in zip(strategy._optimizers, fresh._optimizers): + torch.testing.assert_close(expected.state_dict(), actual.state_dict()) + + +def test_strategy_accumulation_phase_and_checkpoint(tmp_path): + strategy, loss = make_strategy(alternating=True, accumulation_steps=2) + batch = TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2]) + before = strategy.models["student"].gain.detach().clone() + strategy.train_batch(batch) + assert strategy.step_count == 0 + torch.testing.assert_close(strategy.models["student"].gain, before) + with pytest.raises(ValueError, match="accumulation"): + strategy.save_checkpoint(tmp_path, epoch=0) + strategy.train_batch(batch) + strategy.train_batch(batch) + strategy.train_batch(batch) + assert strategy.step_count == 2 + assert loss.calls == [ + ("student", 0), + ("student", 0), + ("auxiliary", 1), + ("auxiliary", 1), + ] + + +def test_strategy_scaler_skip_keeps_phase(): + strategy, loss = make_strategy(alternating=True, precision=torch.float16) + before = strategy.models["student"].gain.detach().clone() + strategy.train_batch( + TensorDict({"x0": torch.full((2, 3, 5), float("inf"))}, batch_size=[2]) + ) + assert strategy.step_count == 0 + assert strategy.batch_count == 1 + torch.testing.assert_close(strategy.models["student"].gain, before) + strategy.train_batch(TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2])) + assert loss.calls == [("student", 0), ("student", 0)] + assert strategy.step_count == 1 + + +def test_strategy_scaler_resume(tmp_path): + strategy, _ = make_strategy(precision=torch.float16) + batch = TensorDict({"x0": torch.randn(2, 3, 5)}, [2]) + strategy.train_batch(batch) + assert strategy.state_dict()["scalers"] + strategy.save_checkpoint(tmp_path, epoch=1) + fresh, _ = make_strategy(precision=torch.float16) + fresh.load_checkpoint(tmp_path, epoch=1) + assert fresh.state_dict()["scalers"] == strategy.state_dict()["scalers"] + fresh.train_batch(batch) + strategy.train_batch(batch) + torch.testing.assert_close( + fresh.models["student"].state_dict(), strategy.models["student"].state_dict() + ) + + +def test_native_setup_wrapper_binding(tmp_path): + from nvalchemi.training._stages import TrainingStage + + class WrapHook: + stage = TrainingStage.SETUP + frequency = 1 + + def __call__(self, ctx, stage): + if not isinstance(ctx.models["student"], CountingWrapper): + ctx.models["student"] = CountingWrapper(ctx.models["student"]) + + strategy, loss = make_strategy( + hooks=[WrapHook()], + bind_models_fn=lambda loss, models: setattr(loss, "models", models), + ) + strategy.train_batch(TensorDict({"x0": torch.randn(2, 3, 5)}, [2])) + assert loss.models["student"].calls == 1 + strategy.save_checkpoint(tmp_path, epoch=1) + fresh, _ = make_strategy( + hooks=[WrapHook()], + bind_models_fn=lambda loss, models: setattr(loss, "models", models), + ) + fresh.load_checkpoint(tmp_path, epoch=1) + assert isinstance(fresh.models["student"], CountingWrapper) + torch.testing.assert_close( + fresh.models["student"].state_dict(), strategy.models["student"].state_dict() + ) + + +def test_accumulation_epoch_duration(): + strategy, loss = make_strategy(num_epochs=1, accumulation_steps=2) + batch = TensorDict({"x0": torch.randn(2, 3, 5)}, [2]) + strategy.run([batch] * 4) + assert len(loss.calls) == 4 + assert strategy.epoch_count == 1 + assert strategy.step_count == 2 + + +class CountingWrapper(Module): + def __init__(self, model): + super().__init__() + self.model = model + self.calls = 0 + + def forward(self, *args, **kwargs): + self.calls += 1 + return self.model(*args, **kwargs) + + +@pytest.mark.parametrize("method", ("sCM", "MeanFlow", "DMD2", "f-distill")) +def test_strategy_binds_existing_execution_wrapper(method): + pytest.importorskip("fastgen") + from physicsnemo.diffusion.distillation import DistillationLoss + from physicsnemo.diffusion.noise_schedulers import RectifiedFlowNoiseScheduler + + from .helpers import FeatureDiscriminator + + model = StructuredDiffusionModel() + roles = {"student": model} + options = {} + if method in ("DMD2", "f-distill"): + roles.update( + teacher=StructuredDiffusionModel(), fake_score=StructuredDiffusionModel() + ) + auxiliary = {"fake_score": roles["fake_score"]} + if method == "f-distill": + roles["discriminator"] = FeatureDiscriminator() + auxiliary["discriminator"] = roles["discriminator"] + options = { + "teacher": roles["teacher"], + "auxiliary_models": auxiliary, + "model_options": { + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + } + loss = DistillationLoss( + model, RectifiedFlowNoiseScheduler(), method=method, **options + ) + wrapper = CountingWrapper(model) + roles["student"] = wrapper + strategy = DistillationStrategy( + models=roles, + loss_fn=loss, + optimizer_configs={ + "student": [ + {"optimizer_cls": torch.optim.SGD, "optimizer_kwargs": {"lr": 0.01}} + ] + }, + num_steps=1, + ) + strategy.train_batch(TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2])) + assert wrapper.calls > 0 + assert strategy._optimizers[0].param_groups[0]["params"][0] is model.gain + + +def test_strategy_compiled_and_mixed_modes(): + strategy, loss = make_strategy() + model = strategy.models["student"] + model.eval() + wrapper = torch.compile(model, backend="eager", fullgraph=True) + strategy.models["student"] = wrapper + strategy.bind_models_fn = lambda loss, models: setattr(loss, "models", models) + strategy.train_batch(TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2])) + assert loss.models["student"] is wrapper + assert not model.training + + +def test_strategy_ema_after_successful_updates(): + strategy, _ = make_strategy() + target = StructuredDiffusionModel() + before = target.gain.detach().clone() + strategy.models["ema"] = target + strategy.ema_models = {"student": "ema"} + strategy.ema_decay = 0.5 + strategy.bind_models_fn = lambda loss, models: setattr(loss, "models", models) + strategy.train_batch(TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2])) + torch.testing.assert_close( + target.gain, (before + strategy.models["student"].gain) / 2 + ) + + +def test_joint_phase_scaler_skips_all_roles(): + strategy, _ = make_strategy( + alternating=True, + precision=torch.float16, + phase_fn=lambda step: "joint", + phase_roles={"joint": ("student", "fake_score")}, + ) + student, fake = (strategy.models[name] for name in ("student", "fake_score")) + strategy.loss_fn = lambda **kwargs: { + "total_loss": student.gain.square() + fake.gain * float("inf") + } + before = student.gain.detach().clone() + strategy.train_batch(TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2])) + torch.testing.assert_close(student.gain, before) + assert strategy.step_count == 0 + + +def test_optional_strategy_import(): + program = """ +import importlib.metadata +import sys +original = importlib.metadata.version +def version(name): + if name == 'nvalchemi-toolkit': + raise importlib.metadata.PackageNotFoundError(name) + return original(name) +importlib.metadata.version = version +from physicsnemo.diffusion.distillation import BaseDistillationLoss, DistillationStrategy +assert not any(name.startswith('nvalchemi') for name in sys.modules) +try: + DistillationStrategy() +except ImportError as exc: + assert 'diffusion-strategy' in str(exc) +else: + raise AssertionError('Expected actionable delayed error') +""" + subprocess.run( # noqa: S603 + [sys.executable, "-c", program], check=True, capture_output=True, text=True + ) + + +@pytest.mark.parametrize("steps", (1, 2)) +def test_strategy_f_distill_resume(tmp_path, steps): + pytest.importorskip("fastgen") + from physicsnemo.diffusion.distillation import DistillationLoss + from physicsnemo.diffusion.noise_schedulers import RectifiedFlowNoiseScheduler + + from .helpers import FeatureDiscriminator + + def build(): + models = { + role: StructuredDiffusionModel() + for role in ("student", "teacher", "fake_score", "ema") + } + models["discriminator"] = FeatureDiscriminator() + loss = DistillationLoss( + models["student"], + RectifiedFlowNoiseScheduler(), + method="f-distill", + teacher=models["teacher"], + auxiliary_models={ + name: models[name] for name in ("fake_score", "discriminator") + }, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + ) + return DistillationStrategy( + models=models, + loss_fn=loss, + num_steps=6, + phase_fn=lambda step: "student" if step % 2 else "auxiliary", + ema_models={"student": "ema"}, + optimizer_configs={ + role: [ + { + "optimizer_cls": torch.optim.Adam, + "optimizer_kwargs": {"lr": 0.001}, + } + ] + for role in ("student", "fake_score", "discriminator") + }, + ) + + strategy = build() + batch = TensorDict({"x0": torch.randn(2, 3, 5)}, batch_size=[2]) + for _ in range(steps): + strategy.train_batch(batch) + strategy.save_checkpoint(tmp_path, epoch=steps, include_rng=True) + strategy.train_batch(batch) + fresh = build() + fresh.load_checkpoint(tmp_path, epoch=steps, restore_rng=True) + fresh.train_batch(batch) + for name in strategy.models: + torch.testing.assert_close( + fresh.models[name].state_dict(), strategy.models[name].state_dict() + ) + for left, right in zip(strategy._optimizers, fresh._optimizers): + torch.testing.assert_close(left.state_dict(), right.state_dict()) + torch.testing.assert_close( + strategy.loss_fn.state_dict(), fresh.loss_fn.state_dict() + ) + torch.testing.assert_close(strategy._last_loss, fresh._last_loss) diff --git a/uv.lock b/uv.lock index b5cea3326c..6c1b801c27 100644 --- a/uv.lock +++ b/uv.lock @@ -1362,6 +1362,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/77/f5/21d2de20e8b8b0408f0681956ca2c69f1320a3848ac50e6e7f39c6159675/babel-2.18.0-py3-none-any.whl", hash = "sha256:e2b422b277c2b9a9630c1d7903c2a00d0830c409c59ac8cae9081c92f1aeba35", size = 10196845, upload-time = "2026-02-01T12:30:53.445Z" }, ] +[[package]] +name = "beartype" +version = "0.23.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/79/34/c0cfdaee9c796c8381439fed0fcab62810753b8b0adb1ec88fa7f41eaa5d/beartype-0.23.1.tar.gz", hash = "sha256:8b805f246b32931c74f2b321ea0939fe5bb2ff43c12a3eff19cc4c854b832e1d", size = 1993714, upload-time = "2026-10-11T05:36:55.251Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/1c/cb/2f6aa982e2860c1156830db3da7f823b71784ed7d533ca289ea9b6c6347a/beartype-0.23.1-py3-none-any.whl", hash = "sha256:4461b4dc57e3fdd6c8a8464b22cf118002eea107c33e9fef651c69068ab3cce3", size = 1582101, upload-time = "2026-10-11T05:36:52.431Z" }, +] + [[package]] name = "beautifulsoup4" version = "4.15.0" @@ -7592,6 +7601,39 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/ce/e5ec180bc41812edcd8daeb8639d205622c0e8c02259d8ab25a0201b3c2a/numpy-2.4.6-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:2803abfebfc990042cd494d8ce2d5f82e9d847af6d35ec486923aa19dbad5e73", size = 12504263, upload-time = "2026-05-18T23:37:09.715Z" }, ] +[[package]] +name = "nvalchemi-toolkit" +version = "0.2.0" +source = { git = "https://github.com/NVIDIA/nvalchemi-toolkit.git?rev=f59e28dc6701583a13bb294e536076c2a7379e98#f59e28dc6701583a13bb294e536076c2a7379e98" } +dependencies = [ + { name = "click" }, + { name = "jaxtyping" }, + { name = "loguru" }, + { name = "numpy" }, + { name = "nvalchemi-toolkit-ops" }, + { name = "nvidia-physicsnemo" }, + { name = "periodictable" }, + { name = "plotext" }, + { name = "plum-dispatch" }, + { name = "pydantic" }, + { name = "rich" }, + { name = "tensordict" }, + { name = "torch", version = "2.13.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-natten-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu13' and extra == 'extra-18-nvidia-physicsnemo-natten-cu12') or (extra == 'extra-18-nvidia-physicsnemo-cu13' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu12' and extra == 'extra-18-nvidia-physicsnemo-natten-cu13') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu13' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12') or (extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13') or (extra != 'extra-18-nvidia-physicsnemo-cu12' and extra != 'extra-18-nvidia-physicsnemo-cu13' and extra != 'extra-18-nvidia-physicsnemo-natten-cu12' and extra != 'extra-18-nvidia-physicsnemo-natten-cu13' and extra != 'extra-18-nvidia-physicsnemo-transformer-engine-cu12' and extra != 'extra-18-nvidia-physicsnemo-transformer-engine-cu13')" }, + { name = "torch", version = "2.13.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "extra == 'extra-18-nvidia-physicsnemo-cu12' or extra == 'extra-18-nvidia-physicsnemo-natten-cu12' or extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12'" }, + { name = "torch", version = "2.13.0+cu130", source = { registry = "https://download.pytorch.org/whl/cu130" }, marker = "extra == 'extra-18-nvidia-physicsnemo-cu13' or extra == 'extra-18-nvidia-physicsnemo-natten-cu13' or extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13'" }, + { name = "zarr", version = "3.1.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12' or (extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-natten-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu13' and extra == 'extra-18-nvidia-physicsnemo-natten-cu12') or (extra == 'extra-18-nvidia-physicsnemo-cu13' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu12' and extra == 'extra-18-nvidia-physicsnemo-natten-cu13') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu13' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12') or (extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13')" }, + { name = "zarr", version = "3.4.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12' or (extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-natten-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13') or (extra == 'extra-18-nvidia-physicsnemo-cu13' and extra == 'extra-18-nvidia-physicsnemo-natten-cu12') or (extra == 'extra-18-nvidia-physicsnemo-cu13' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu12' and extra == 'extra-18-nvidia-physicsnemo-natten-cu13') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13') or (extra == 'extra-18-nvidia-physicsnemo-natten-cu13' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12') or (extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu12' and extra == 'extra-18-nvidia-physicsnemo-transformer-engine-cu13')" }, +] + +[[package]] +name = "nvalchemi-toolkit-ops" +version = "0.5.0" +source = { git = "https://github.com/NVIDIA/nvalchemi-toolkit-ops?rev=0.5.0-rc#6702138e404d40ea79dec62f68c2fa67e858580e" } +dependencies = [ + { name = "numpy" }, + { name = "warp-lang" }, +] + [[package]] name = "nvdlfw-inspect" version = "0.2.2" @@ -9537,6 +9579,9 @@ datapipes-extras = [ diffusion-fastgen = [ { name = "fastgen" }, ] +diffusion-strategy = [ + { name = "nvalchemi-toolkit" }, +] gnns = [ { name = "line-profiler" }, { name = "mlflow" }, @@ -9654,6 +9699,7 @@ requires-dist = [ { name = "netcdf4", marker = "extra == 'datapipes-extras'", specifier = ">=1.7.3" }, { name = "numpy", specifier = ">=2.2.6" }, { name = "numpy", marker = "extra == 'cu13'", specifier = "<2.5" }, + { name = "nvalchemi-toolkit", marker = "extra == 'diffusion-strategy'", git = "https://github.com/NVIDIA/nvalchemi-toolkit.git?rev=f59e28dc6701583a13bb294e536076c2a7379e98" }, { name = "nvidia-dali-cuda120", marker = "extra == 'cu12'", specifier = ">=2.3.0", index = "https://pypi.nvidia.com/" }, { name = "nvidia-dali-cuda130", marker = "extra == 'cu13'", specifier = ">=2.3.0", index = "https://pypi.nvidia.com/" }, { name = "nvtx", specifier = ">=0.2.16" }, @@ -9712,7 +9758,7 @@ requires-dist = [ { name = "xarray", marker = "extra == 'datapipes-extras'", specifier = ">=2026.7.0" }, { name = "zarr", marker = "extra == 'datapipes-extras'", specifier = ">=3.1.6" }, ] -provides-extras = ["cu12", "cu13", "datapipes-extras", "diffusion-fastgen", "gnns", "mesh-extras", "model-extras", "natten-cu12", "natten-cu13", "nn-extras", "sym", "transformer-engine-cu12", "transformer-engine-cu13", "uq-extras", "utils-extras"] +provides-extras = ["cu12", "cu13", "datapipes-extras", "diffusion-fastgen", "diffusion-strategy", "gnns", "mesh-extras", "model-extras", "natten-cu12", "natten-cu13", "nn-extras", "sym", "transformer-engine-cu12", "transformer-engine-cu13", "uq-extras", "utils-extras"] [package.metadata.requires-dev] dev = [ @@ -10216,6 +10262,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/71/e7/40fb618334dcdf7c5a316c0e7343c5cd82d3d866edc100d98e29bc945ecd/partd-1.4.2-py3-none-any.whl", hash = "sha256:978e4ac767ec4ba5b86c6eaa52e5a2a3bc748a2ca839e8cc798f1cc6ce6efb0f", size = 18905, upload-time = "2024-05-06T19:51:39.271Z" }, ] +[[package]] +name = "periodictable" +version = "2.0.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "pyparsing" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/e4/f2a338abfb5a6b9e2a48dc4acaa9df0bd94126df336ab403178f6a7a901a/periodictable-2.0.2-py3-none-any.whl", hash = "sha256:7d0680baf758715dd66bdac733de82dcb976aa03981fbd6d75a84213fb3b5f86", size = 811077, upload-time = "2024-12-20T02:26:26.165Z" }, +] + [[package]] name = "pexpect" version = "4.9.0" @@ -10321,6 +10379,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/40/31/1521c38d175108ac395941adbcb9f11e085340bed81ad166e8ad8c21c17b/platformdirs-4.12.1-py3-none-any.whl", hash = "sha256:7f251621befae78b2cb74616f9fc99d3fa3db0afba5035987731f8dcd430337d", size = 33011, upload-time = "2026-09-28T05:05:15.06Z" }, ] +[[package]] +name = "plotext" +version = "5.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/c9/d7/f75f397af966fe252d0d34ffd3cae765317fce2134f925f95e7d6725d1ce/plotext-5.3.2.tar.gz", hash = "sha256:52d1e932e67c177bf357a3f0fe6ce14d1a96f7f7d5679d7b455b929df517068e", size = 61967, upload-time = "2024-09-24T15:13:37.728Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f6/1e/12fe7c40cd2099a1f454518754ed229b01beaf3bbb343127f0cc13ce6c22/plotext-5.3.2-py3-none-any.whl", hash = "sha256:394362349c1ddbf319548cfac17ca65e6d5dfc03200c40dfdc0503b3e95a2283", size = 64047, upload-time = "2024-09-24T15:13:36.296Z" }, +] + [[package]] name = "plotly" version = "7.1.0" @@ -10343,6 +10410,32 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "plum-dispatch" +version = "2.10.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "beartype" }, + { name = "rich" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/dc/75/807b02f339906867de07baaecd3519fe8704df013a27ae0a57250ed461d3/plum_dispatch-2.10.1.tar.gz", hash = "sha256:93aa1ec5c4683c0777b9d7d1ccc0b31e0e82f38399116345b4ed8d0fb67cc66e", size = 250842, upload-time = "2026-09-07T17:38:44.9Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8a/1a/3736f28127e6aa0e87eaee3286c0f93bf8fd25382786d8409073fc1a9010/plum_dispatch-2.10.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:143b44537476fcfa644a7b3b00cdfce91cd94c501eb5d4a342d0b1dae6e00269", size = 279045, upload-time = "2026-09-07T17:38:26.368Z" }, + { url = "https://files.pythonhosted.org/packages/67/a7/dfc192aa694244278d66c5251bc415bb23234a8cb12013e46dfae2747091/plum_dispatch-2.10.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:193fa7d9809e6ccf38d0ec423b92c2ce88b1bb5dfa7c769ce7a3b216aeb61eb4", size = 264041, upload-time = "2026-09-07T17:38:27.743Z" }, + { url = "https://files.pythonhosted.org/packages/77/8d/15617a11d26a7dfa9e5f6f36103c4a3e27e5651d9261ca6bd6a099a69abc/plum_dispatch-2.10.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b8909bc8c3b776f116c50986718befeda1de17a3dc61e7ae93f5d3d71de5f98d", size = 328403, upload-time = "2026-09-07T17:38:29.315Z" }, + { url = "https://files.pythonhosted.org/packages/f0/20/7e5a1535a858ae4716083ac15d935394207145f8249fea8f693dd979a7f6/plum_dispatch-2.10.1-cp311-cp311-win_amd64.whl", hash = "sha256:4adc9e9335bee9fffd24abd3be0aa6810e403129697ab767976d70cc3c88405c", size = 244207, upload-time = "2026-09-07T17:38:30.608Z" }, + { url = "https://files.pythonhosted.org/packages/5d/d1/10c40a76f0fc3959d73a245e3e63ab576062e2c7db1fb41a4e035a5f4ee0/plum_dispatch-2.10.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:476821e99fe4c747725887021ea09028ff5e013d8732c40c6d44df65815a70ea", size = 285017, upload-time = "2026-09-07T17:38:32.092Z" }, + { url = "https://files.pythonhosted.org/packages/90/91/d0a7d52b1a2e02e883b078e2801cfd544abb05b1a28183e46db43d812667/plum_dispatch-2.10.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:e1a73a06483fffa8f8247bbf6057157075be2f63e76845792eb75d86b11f1188", size = 267393, upload-time = "2026-09-07T17:38:33.461Z" }, + { url = "https://files.pythonhosted.org/packages/b9/c3/d8793642ce7d29ea15cc39c10185ef126a27fcb8ea230154978221deb4b2/plum_dispatch-2.10.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a5b6a2474f53b23cd07254c2ad286fc8371a3c4ed17dee72794eac9c3a6873db", size = 334557, upload-time = "2026-09-07T17:38:34.857Z" }, + { url = "https://files.pythonhosted.org/packages/41/c7/3544b2bf05aa9e1a1b45981d21da5e5d5c135fcc2da0ee1040f73512662f/plum_dispatch-2.10.1-cp312-cp312-win_amd64.whl", hash = "sha256:ae2f76ce97bcbf1d7038e7c352180cbb35c90c23ea0848e2924f35133fb0f223", size = 245830, upload-time = "2026-09-07T17:38:36.228Z" }, + { url = "https://files.pythonhosted.org/packages/16/69/c73f5145dd136158c74295c3f1751062beaf744f508cd04dce61bfe9b059/plum_dispatch-2.10.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:a8cebe7c3380e8cd451008d405fbb3a202626a7da4c140deff36d52dfb05e5aa", size = 283228, upload-time = "2026-09-07T17:38:37.664Z" }, + { url = "https://files.pythonhosted.org/packages/c3/1e/252f11d5443e341cf1a1c4986ec0e5a365bcb605190cfa26410acbe68dc6/plum_dispatch-2.10.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:4e6ffcc378838cffc9f13d67f1ea79585ffea7e04aa00e888941ef63400f175b", size = 265603, upload-time = "2026-09-07T17:38:39.254Z" }, + { url = "https://files.pythonhosted.org/packages/65/11/6cdfa574f1682840b644527a1644e0740ddb3ddc6e759dcbef4a49091df6/plum_dispatch-2.10.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:131145651d4fa89ee97a026e49545638380a68f71f530ca57b0efaa04dfad1b6", size = 332929, upload-time = "2026-09-07T17:38:40.715Z" }, + { url = "https://files.pythonhosted.org/packages/52/9a/665b51655fd98196066d042d06af9641d0ac589f9e23ad3e895002c9563b/plum_dispatch-2.10.1-cp313-cp313-win_amd64.whl", hash = "sha256:f75fcfb61c89ca561e3b856d2d9392167d51a283d23c7b9a5d11937a3120cae7", size = 246611, upload-time = "2026-09-07T17:38:42.249Z" }, + { url = "https://files.pythonhosted.org/packages/55/9f/0adfb5a90d65233b764cb1b07244b72ded1db7cfc66955418544d0da2e38/plum_dispatch-2.10.1-py3-none-any.whl", hash = "sha256:dac63fd5ada8ae8045aefed127dacda0ba02db2635876b12f4e1dfaf6db73f39", size = 49430, upload-time = "2026-09-07T17:38:43.589Z" }, +] + [[package]] name = "pooch" version = "1.9.0"