diff --git a/docs/api/diffusion/distillation.rst b/docs/api/diffusion/distillation.rst index 950a9458fe..91fbcac74d 100644 --- a/docs/api/diffusion/distillation.rst +++ b/docs/api/diffusion/distillation.rst @@ -6,9 +6,9 @@ Distillation losses This API is alpha. Imports emit ``ExperimentalFeatureWarning`` and future releases may introduce changes that break backward compatibility. -``DistillationLoss`` computes sCM or MeanFlow objectives for PhysicsNeMo -models. Each call constructs training targets and returns a differentiable -``total_loss`` with detached values for logging. +``DistillationLoss`` computes sCM, MeanFlow, DMD2, or f-distill losses for +PhysicsNeMo models. Each call constructs training targets and returns a +differentiable ``total_loss`` with detached values for logging. Use it in your training loop with initialized models and a noise scheduler. Your loop controls backpropagation, optimizer updates, and checkpoints. @@ -143,6 +143,128 @@ Save your PhysicsNeMo models and optimizers as usual. Recreate the loss with the same noise schedules and settings when resuming. sCM and MeanFlow have no extra running statistics, so their ``state_dict`` results are empty. +Alternate distribution-matching objectives +----------------------------------------- + +DMD2 and f-distill match a student's generated distribution to a fixed +diffusion teacher. Unlike an ordinary denoising loss that trains a model on +clean data and noise, these objectives also train a fake-score model on +student-generated samples. + +Initialize the models and pass the fake-score model in ``auxiliary_models``. +Choose ``phase="student"`` or ``phase="auxiliary"`` on each call and step +the matching optimizer. Pass the training step separately as ``step``. + +.. code-block:: python + + loss_fn = DistillationLoss( + student, scheduler, method="DMD2", teacher=teacher, + auxiliary_models={"fake_score": fake_score}, + ) + for phase, optimizer in (("auxiliary", fake_optimizer), + ("student", student_optimizer)): + optimizer.zero_grad() + loss_fn(x0, phase=phase, step=step)["total_loss"].backward() + optimizer.step() + +The student phase trains only the student. The auxiliary phase trains the +fake-score model and any discriminator. The teacher stays fixed in both phases. +Use separate parameters for each model. The loss preserves your choices of +training mode and frozen parameters. + +DMD2 trains the student using teacher and fake-score predictions, with an +optional adversarial loss. f-distill also weights this score +matching by a selected f-divergence and requires a discriminator. + +Supply ``discriminator`` in ``auxiliary_models`` to enable adversarial training. +The generator GAN coefficient defaults to ``0.001``. Set +``gan_loss_weight_gen`` to change it. f-distill requires a positive coefficient. +The discriminator accepts a TensorDict of named teacher features and returns +scores with shape :math:`(B, L)`, where :math:`L` counts scores per sample. +Supply a ``feature_fn(prediction, extras)`` callback to select those features +from your teacher output. The teacher must still return its prediction. +Feature fields are model-specific, not reserved diffusion outputs. The callback +interprets them. For example, if your teacher adds a +``bottleneck`` field to a subclass of ``DiffusionExtraOutputs``: + +.. code-block:: python + + from tensordict import TensorDict + + def teacher_features(prediction, extras): + return TensorDict({"bottleneck": extras.bottleneck}, batch_size=extras.batch_size) + + loss_fn = DistillationLoss( + student, scheduler, method="f-distill", teacher=teacher, + auxiliary_models={"fake_score": fake_score, + "discriminator": discriminator}, + model_options={"teacher": {"feature_names": ("bottleneck",), + "feature_fn": teacher_features}}, + method_options={"f_distill": {"f_div": "js"}}, + ) + +Supported divergence names are ``rkl``, ``kl``, ``js``, ``sf``, ``neyman``, +``sh``, and ``jf``. f-distill tracks running ratio statistics in the loss object. +Save and restore ``loss_fn.state_dict()`` alongside model and optimizer state. +Pass ``update_state=False`` to compute a loss without changing these statistics. + +``model_options["fake_score"]`` configures its prediction type or scheduler. +The loss converts predictions when these models use different noise schedules. +Use ``fake_score_pred_type`` to select the fake-score training target and +``fake_score_sampling`` to choose its training-time distribution. + +One-step training draws student inputs at the scheduler's upper training time. +Custom schedulers without that bound need explicit ``input_student`` and +``t_student`` on the loss call, or configured training bounds. Multi-step training +uses ``student_sample_steps`` and an explicit ``training_times`` grid under the +student's ``scheduler_options``. Specify training times independently of your +inference grid. +Pass ``noise`` to control the score perturbation independently. + +Adversarial masking +~~~~~~~~~~~~~~~~~~~ + +Adversarial masks exclude corrupted inputs during feature extraction and +discriminator pooling. The loss needs a callback that maps the input mask to +valid discriminator scores. + +For masked adversarial training, supply a callback with this signature: + +.. code-block:: python + + def adversarial_fn(x_t, t, condition, weight): + ... + return logits, logit_weight + +Return ``logits`` with shape :math:`(B, L)` and finite, +nonnegative ``logit_weight`` values that broadcast to that shape. + +Use the same teacher and discriminator you passed to the loss, and keep +gradients through their inputs for student training. You may need a masked +backbone or feature extractor. Masking only the final scores cannot remove +corruption that has already spread through features or pooled statistics. + +The loss applies the returned weights to generator, discriminator, and approximate +R1 residuals. Zero weights also exclude scores and empty samples from f-distill +ratio statistics. Positive residual weights remain linear. The callback receives +both real and generated perturbations, including the R1 query when enabled. + +Distributed statistics and compilation +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +f-distill reduces ratio statistics over the default process group for ordinary +data parallelism. For a domain-parallel scheduler, explicitly supply the data +group through ``method_options["f_distill"]["statistics_group"]``. It must +group corresponding spatial partitions across independent data replicas, not +different partitions of the same sample. Each rank participates even when its +local batch has no valid entries. The callback must produce consistent global scores +across spatial partitions when the architecture needs spatial aggregation. + +For DMD2, you can compile the individual models with ``torch.compile``. +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. + Custom implementations ---------------------- diff --git a/physicsnemo/diffusion/_fastgen.py b/physicsnemo/diffusion/_fastgen.py index 87a7f2f064..82e2bd068f 100644 --- a/physicsnemo/diffusion/_fastgen.py +++ b/physicsnemo/diffusion/_fastgen.py @@ -17,7 +17,7 @@ r"""Adapt PhysicsNeMo model calls and predictions for FastGen losses.""" import importlib -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import Any, Literal import torch @@ -59,6 +59,11 @@ def __init__( prediction_type: Literal["x0", "epsilon", "score", "flow"], prediction_semantics: Literal["instantaneous", "average"] = "instantaneous", model_kwargs: Mapping[str, Any] | None = None, + feature_names: tuple[str, ...] = (), + feature_fn: Callable[ + [Float[Tensor, " B *dims"], DiffusionExtraOutputs | None], TensorDict + ] + | None = None, ) -> None: if not _AVAILABLE: raise ImportError( @@ -90,6 +95,8 @@ def __init__( self._prediction_semantics = prediction_semantics self._model_kwargs = dict(model_kwargs or {}) self._input_noise_scheduler = None + self._feature_names = tuple(feature_names) + self._feature_fn = feature_fn def set_noise_schedule( self, schedule_type: str | None = None, **noise_schedule_kwargs: Any @@ -131,10 +138,6 @@ def forward( raise ValueError( "An average flow cannot use instantaneous prediction conversions." ) - if return_features_early or feature_indices: - raise NotImplementedError( - "This adapter does not define feature extraction." - ) input_x, input_t = x_t, t source = self._input_noise_scheduler if source is not None: @@ -168,6 +171,21 @@ def forward( prediction, extras = output if isinstance(output, tuple) else (output, None) if not isinstance(prediction, Tensor): raise TypeError("A diffusion model must return a prediction tensor.") + features = None + if return_features_early or feature_indices: + if self._feature_fn is None: + raise ValueError( + "Adversarial feature requests require an explicit feature_fn." + ) + named_features = self._feature_fn(prediction, extras) + indices = ( + sorted(feature_indices) + if feature_indices + else range(len(self._feature_names)) + ) + features = [named_features[self._feature_names[index]] for index in indices] + if return_features_early: + return features if self._prediction_type == "score": prediction = self.noise_scheduler._scheduler.score_to_epsilon(prediction, t) prediction = self.noise_scheduler.convert_model_output( @@ -194,4 +212,4 @@ def forward( f"Expected logvar shape {(x_t.shape[0],)}, got {logvar.shape}." ) return prediction, logvar[:, None] # (B, 1) - return prediction + return (prediction, features) if features is not None else prediction diff --git a/physicsnemo/diffusion/distillation/losses.py b/physicsnemo/diffusion/distillation/losses.py index 9d1938f259..966b4d3578 100644 --- a/physicsnemo/diffusion/distillation/losses.py +++ b/physicsnemo/diffusion/distillation/losses.py @@ -36,20 +36,22 @@ _METHODS = { "sCM": ("scm._fastgen", "_SCMLoss"), "MeanFlow": ("meanflow._fastgen", "_MeanFlowLoss"), + "DMD2": ("dmd2._fastgen", "_DMD2Loss"), + "f-distill": ("f_distill._fastgen", "_FDistillLoss"), } class DistillationLoss: - r"""Compute sCM or MeanFlow objectives for a diffusion student. + r"""Compute sCM, MeanFlow, DMD2, or f-distill losses for a diffusion student. - The selected method defines the training loss and constructs its - targets, with or without a teacher. The loss returns ``total_loss`` for - backpropagation and detached values for logging. Your training loop - controls optimization. + The selected method defines the training loss and constructs its targets. + Each call returns ``total_loss`` for backpropagation and detached values + for logging. Your training loop controls optimization. Initialize and wrap your models before constructing the loss. Call it on - each batch of clean data, backpropagate through ``losses["total_loss"]``, - and step your optimizer. See the distillation guide for compilation options. + each batch, backpropagate through ``losses["total_loss"]``, and step the + optimizer for the selected phase. See the distillation guide for + compilation options. Parameters ---------- @@ -59,18 +61,27 @@ class DistillationLoss: noise_scheduler : LinearGaussianNoiseScheduler or DomainParallelNoiseScheduler Noise schedule for the student. Its ``sample_time`` method supplies training times unless you pass times to the loss call. - method : {"sCM", "MeanFlow"} + method : {"sCM", "MeanFlow", "DMD2", "f-distill"} Training loss. MeanFlow students receive ``(t, DiffusionExtraTimes)`` and predict average velocity over the interval from ``t`` to ``r``. teacher : torch.nn.Module, optional - Model providing training targets. Omit it to train from clean data + Model providing training targets. DMD2 and f-distill require a + teacher. For sCM and MeanFlow, omit it to train from clean data and noise without a pretrained teacher. + auxiliary_models : mapping, optional + Models trained during the auxiliary phase. Supply ``fake_score`` for + DMD2 and f-distill, and ``discriminator`` for adversarial training. + Initialize these models before passing them. method_options : mapping, optional Method-specific loss settings. Unknown settings raise an error. model_options : mapping, optional - Settings for ``student`` or ``teacher``. Set ``prediction_type`` to + Settings for ``student``, ``teacher``, or ``fake_score``. Set ``prediction_type`` to describe what each model predicts. ``model_kwargs`` forwards model-specific options, and ``noise_scheduler`` selects a different schedule for that model. + Adversarial teachers also need ``feature_names`` and a + ``feature_fn(prediction, extras)`` callback that returns named features + as a ``TensorDict``. The teacher always returns its prediction; the + callback selects model-specific extras for the discriminator. Examples -------- @@ -100,6 +111,7 @@ def __init__( *, method: str, teacher: torch.nn.Module | None = None, + auxiliary_models: Mapping[str, torch.nn.Module] | None = None, method_options: Mapping[str, Any] | None = None, model_options: Mapping[str, Mapping[str, Any]] | None = None, ) -> None: @@ -111,12 +123,18 @@ def __init__( implementation = getattr( importlib.import_module(f"{__package__}.methods.{module}"), name ) + auxiliary_options = {} + if method in ("DMD2", "f-distill"): + auxiliary_options["auxiliary_models"] = dict(auxiliary_models or {}) + elif auxiliary_models: + raise ValueError(f"{method} does not use auxiliary models.") self._implementation = implementation( student, noise_scheduler, teacher=teacher, method_options=dict(method_options or {}), model_options=dict(model_options or {}), + **auxiliary_options, ) def __call__( @@ -152,7 +170,9 @@ def __call__( Conditioning forwarded to the supplied models. phase : str, optional Training phase. sCM and MeanFlow accept only ``"student"`` - or ``None``. + or ``None``. For DMD2 and f-distill, choose ``"student"`` + to train the student or ``"auxiliary"`` to train the fake-score + model and optional discriminator. Step the matching optimizers. step : int, optional, default=0 Training step, used to adjust loss warmup. return_outputs : bool, optional, default=False diff --git a/physicsnemo/diffusion/distillation/methods/dmd2/__init__.py b/physicsnemo/diffusion/distillation/methods/dmd2/__init__.py new file mode 100644 index 0000000000..af85283aa4 --- /dev/null +++ b/physicsnemo/diffusion/distillation/methods/dmd2/__init__.py @@ -0,0 +1,15 @@ +# 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. diff --git a/physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py b/physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py new file mode 100644 index 0000000000..9d47806847 --- /dev/null +++ b/physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py @@ -0,0 +1,540 @@ +# 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"""DMD2 distribution matching with separate student and auxiliary objectives.""" + +import importlib +from collections.abc import Callable, Mapping +from contextlib import contextmanager +from typing import Any + +import torch +import torch.nn.functional as F +from jaxtyping import Float, Shaped +from omegaconf import OmegaConf +from tensordict import TensorDict +from torch import Tensor + +from physicsnemo.core import Module +from physicsnemo.core.version_check import check_version_spec +from physicsnemo.diffusion._fastgen import _FastGenNetwork +from physicsnemo.diffusion.base import ( + DiffusionExtraTimes, +) +from physicsnemo.diffusion.noise_schedulers import ( + DomainParallelNoiseScheduler, + LinearGaussianNoiseScheduler, +) +from physicsnemo.diffusion.noise_schedulers._fastgen import _FastGenNoiseScheduler + +from ..._fastgen import _configure_options, _FastGenBinding, _ObjectiveReduction + +_AVAILABLE = check_version_spec("fastgen", "0.1.0", hard_fail=False) +if _AVAILABLE: + _backend = importlib.import_module("fastgen.methods.distribution_matching.dmd2") + _DMD2Base = _backend.DMD2Model +else: + _DMD2Base = Module + + +@contextmanager +def _fixed_parameters(*models: torch.nn.Module | None): + r"""Freeze parameters temporarily while preserving gradients through inputs.""" + parameters = { + p: p.requires_grad + for model in models + if model is not None + for p in model.parameters() + } + try: + for parameter in parameters: + parameter.requires_grad_(False) + yield + finally: + for parameter, requires_grad in parameters.items(): + parameter.requires_grad_(requires_grad) + + +class _Discriminator(torch.nn.Module): + r"""Pass teacher features to a discriminator as a named TensorDict.""" + + def __init__(self, model: torch.nn.Module, names: tuple[str, ...]) -> None: + super().__init__() + self.model = model + self.names = names + self.feature_indices = set(range(len(names))) + + def forward(self, features: list[Tensor]) -> Float[Tensor, " B logits"]: + r"""Return discriminator logits from the named teacher features.""" + values = TensorDict( + dict(zip(self.names, features)), batch_size=[features[0].shape[0]] + ) + logits = self.model(values) + if not isinstance(logits, Tensor) or logits.ndim != 2: + raise ValueError("Discriminator output must have shape [B, logits].") + return logits + + +class _DMD2Loss(_FastGenBinding, _DMD2Base): + r"""Match a student's generated distribution to a fixed diffusion teacher. + + Use separate student and auxiliary phases to train the student and fake-score + model. An optional discriminator adds adversarial training; each phase + preserves the supplied models' training modes and frozen parameters. + """ + + def __init__( + self, + student: torch.nn.Module, + noise_scheduler: LinearGaussianNoiseScheduler | DomainParallelNoiseScheduler, + *, + teacher: torch.nn.Module | None, + auxiliary_models: Mapping[str, torch.nn.Module], + method_options: Mapping[str, Any], + model_options: Mapping[str, Mapping[str, Any]], + ) -> None: + torch.nn.Module.__init__(self) + if teacher is None or "fake_score" not in auxiliary_models: + raise ValueError( + "DMD2 requires teacher and auxiliary_models['fake_score']." + ) + unknown = set(auxiliary_models) - {"fake_score", "discriminator"} + if unknown: + raise ValueError(f"Unknown auxiliary roles: {sorted(unknown)}.") + unknown = set(model_options) - {"student", "teacher", "fake_score"} + if unknown: + raise ValueError(f"Unknown model roles: {sorted(unknown)}.") + self._bind_models( + student, + noise_scheduler, + teacher, + {key: value for key, value in model_options.items() if key != "fake_score"}, + ) + options = dict(model_options.get("fake_score", {})) + schedule = _FastGenNoiseScheduler( + options.pop("noise_scheduler", noise_scheduler) + ) + schedule._configure_training(**options.pop("scheduler_options", {})) + options.setdefault("prediction_type", "x0") + self.fake_score = _FastGenNetwork( + auxiliary_models["fake_score"], schedule, **options + ) + for role in (self.teacher, self.fake_score): + if ( + role.noise_scheduler._scheduler + is not self.net.noise_scheduler._scheduler + ): + role._input_noise_scheduler = self.net.noise_scheduler + discriminator = auxiliary_models.get("discriminator") + self.discriminator = None + if discriminator is not None: + names = tuple(model_options.get("teacher", {}).get("feature_names", ())) + if not names or model_options.get("teacher", {}).get("feature_fn") is None: + raise ValueError( + "Configure teacher feature_names and feature_fn for the discriminator." + ) + self.discriminator = _Discriminator(discriminator, names) + roles = [student, teacher, auxiliary_models["fake_score"]] + if discriminator is not None: + roles.append(discriminator) + owned = [set(model.parameters()) for model in roles] + if any( + left & right for i, left in enumerate(owned) for right in owned[i + 1 :] + ): + raise ValueError("Training roles must not share parameters.") + self.config = _configure_options( + { + "gan_loss_weight_gen": 0.001 if discriminator is not None else 0.0, + "gan_use_same_t_noise": False, + "gan_r1_reg_weight": 0.0, + "gan_r1_reg_alpha": 0.1, + "guidance_scale": None, + "skip_layers": None, + "fake_score_pred_type": None, + "student_sample_steps": 1, + "sampling": {}, + "fake_score_sampling": None, + }, + method_options, + ) + if self.config.gan_loss_weight_gen < 0 or self.config.gan_r1_reg_weight < 0: + raise ValueError("GAN loss weights must be nonnegative.") + if self.config.gan_loss_weight_gen > 0 and discriminator is None: + raise ValueError("GAN objectives require an explicit discriminator.") + if self.config.student_sample_steps < 1: + raise ValueError("student_sample_steps must be positive.") + self.config.sample_t_cfg = OmegaConf.create( + {"time_dist_type": "physicsnemo", **self.config.sampling} + ) + self.config.fake_score_pred_type = ( + self.config.fake_score_pred_type or self.fake_score.net_pred_type + ) + self._phase = None + self._adversarial_fn = None + self._update_state = True + + def evaluate( + self, + x0: Float[Tensor, " B *dims"], + t: Float[Tensor, " B"] + | tuple[Float[Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]] + | None, + condition: Tensor | TensorDict | None, + *, + phase: str | None, + step: int, + return_outputs: bool, + update_state: bool, + weight: Float[Tensor, " #B *#dims"] | None = None, + neg_condition: Tensor | TensorDict | None = None, + sample_t_fn: Callable[ + ..., + Float[Tensor, " B"] + | tuple[Float[Tensor, " B"], Shaped[DiffusionExtraTimes, " B"]], + ] + | None = None, + input_student: Float[Tensor, " B *dims"] | None = None, + t_student: Float[Tensor, " B"] | None = None, + noise: Float[Tensor, " B *dims"] | None = None, + adversarial_fn: Callable[..., tuple[Tensor, Tensor]] | None = None, + ) -> Any: + r"""Compute the selected phase's losses for the supplied training step.""" + if phase not in ("student", "auxiliary"): + raise ValueError("Select phase='student' or phase='auxiliary' explicitly.") + if step < 0: + raise ValueError("step must be nonnegative.") + if t is not None and sample_t_fn is not None: + raise ValueError("Supply explicit times or sample_t_fn, not both.") + if (input_student is None) != (t_student is None): + raise ValueError("Supply input_student and t_student together.") + if x0.ndim < 2 or x0.shape[0] == 0: + raise ValueError("Expected nonempty states with shape [B, ...].") + if weight is not None: + weight = torch.broadcast_to(weight.to(x0.device), x0.shape) + if not torch.isfinite(weight).all() or torch.any(weight < 0): + raise ValueError("Weights must be finite and nonnegative.") + if self.config.gan_loss_weight_gen > 0 and adversarial_fn is None: + raise ValueError( + "Weighted GAN objectives require adversarial_fn to map valid support to logits." + ) + self.device = x0.device + batch = x0.shape[0] + source = self.net.noise_scheduler + if sample_t_fn is not None: + t = sample_t_fn(batch, device=x0.device, dtype=source.t_precision) + if isinstance(t, tuple): + raise ValueError( + "Distribution matching requires a single diffusion time tensor." + ) + if t is None: + sampling = self.config.fake_score_sampling if phase == "auxiliary" else None + t = source.sample_t( + batch, device=x0.device, **(sampling or self.config.sampling) + ) + if input_student is None: + initial_noise = torch.randn_like(x0) + if self.config.student_sample_steps == 1: + t_student = torch.full_like(t, source.max_t) + input_student = source.latents(noise=initial_noise) + else: + t_student = source.sample_from_t_list( + batch, + sample_steps=self.config.student_sample_steps, + device=x0.device, + ) + input_student = source.forward_process(x0, initial_noise, t_student) + noise = torch.randn_like(x0) if noise is None else noise + previous = ( + self._phase, + self._call_weight, + self._adversarial_fn, + self._update_state, + ) + self._phase, self._call_weight, self._adversarial_fn, self._update_state = ( + phase, + weight, + adversarial_fn, + update_state, + ) + try: + fixed = ( + (self.teacher, self.fake_score, self.discriminator) + if phase == "student" + else (self.teacher, self.net) + ) + with _fixed_parameters(*fixed): + losses, outputs = self.single_train_step( + { + "real": x0, + "t": t, + "input_student": input_student, + "t_student": t_student, + "eps": noise, + "condition": condition, + "neg_condition": neg_condition, + }, + step, + ) + finally: + self._phase, self._call_weight, self._adversarial_fn, self._update_state = ( + previous + ) + losses = { + key: value if key == "total_loss" else value.detach() + for key, value in losses.items() + } + return ( + (losses, {key: value.detach() for key, value in outputs.items()}) + if return_outputs + else losses + ) + + def single_train_step( + self, data: Mapping[str, Any], iteration: int + ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + r"""Evaluate the selected student or auxiliary loss.""" + args = (data["input_student"], data["t_student"], data["t"], data["eps"]) + if self._phase == "student": + return self._student_update_step( + *args, + data, + condition=data["condition"], + neg_condition=data["neg_condition"], + ) + return self._fake_score_discriminator_update_step( + *args, data["real"], condition=data["condition"] + ) + + def _vsd_loss( + self, + generated: Tensor, + teacher: Tensor, + fake: Tensor, + scale: Tensor | None = None, + ) -> Tensor: + r"""Exclude invalid entries from the target norm before weighting residuals.""" + if self._call_weight is None: + return _backend.variational_score_distillation_loss( + generated, teacher, fake, additional_scale=scale + ) + reduction = _ObjectiveReduction(generated, self._call_weight) + with torch.no_grad(): + difference = reduction.clean(generated.float() - teacher.float()) + mean = difference.abs().sum(reduction.dims) / reduction.count.clamp_min(1) + factor = (mean + 1e-6).reciprocal() + if scale is not None: + factor = factor * scale.float() + factor = factor.to(generated.dtype).reshape( + -1, *([1] * (generated.ndim - 1)) + ) + gradient = reduction.clean(fake - teacher) * factor + target = generated - gradient + return 0.5 * reduction.residual(generated - target, mean=True).mean() + + def _logits( + self, x: Tensor, t: Tensor, condition: Any + ) -> tuple[Tensor, Tensor | None]: + r"""Return discriminator logits and optional weights from the mask callback.""" + if self._adversarial_fn is not None: + logits, weight = self._adversarial_fn(x, t, condition, self._call_weight) + weight = torch.broadcast_to(weight.to(logits.device), logits.shape) + if ( + logits.ndim != 2 + or not torch.isfinite(weight).all() + or torch.any(weight < 0) + ): + raise ValueError( + "Adversarial logits require rank two and finite nonnegative weights." + ) + return logits, weight + features = self.teacher( + x, + t, + condition=condition, + return_features_early=True, + feature_indices=self.discriminator.feature_indices, + ) + return self.discriminator(features), None + + @staticmethod + def _logit_mean(value: Tensor, weight: Tensor | None) -> Tensor: + r"""Average weighted logits over all entries, with zero weights excluded.""" + return ( + value.mean() + if weight is None + else (torch.where(weight > 0, value, 0) * weight).mean() + ) + + def _student_update_step( + self, + input_student: Tensor, + t_student: Tensor, + t: Tensor, + eps: Tensor, + data: Mapping[str, Any], + condition: Any = None, + neg_condition: Any = None, + ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + r"""Compute student score-matching and optional adversarial losses.""" + generated = self.gen_data_from_net( + input_student, t_student, condition=condition + ) + perturbed = self.net.noise_scheduler.forward_process(generated, eps, t) + with torch.no_grad(): + fake = self.fake_score( + perturbed, t, condition=condition, fwd_pred_type="x0" + ) + gan = generated.new_zeros(()) + scale = None + if self.config.gan_loss_weight_gen > 0 and self._adversarial_fn is None: + teacher, features = self.teacher( + perturbed, + t, + condition=condition, + fwd_pred_type="x0", + feature_indices=self.discriminator.feature_indices, + ) + logits, weight = self.discriminator(features), None + teacher = teacher.detach() + else: + with torch.no_grad(): + teacher = self.teacher( + perturbed, t, condition=condition, fwd_pred_type="x0" + ) + if self.config.gan_loss_weight_gen > 0: + logits, weight = self._logits(perturbed, t, condition) + if self.config.guidance_scale is not None: + teacher = self._apply_classifier_free_guidance( + perturbed, t, teacher, neg_condition=neg_condition + ) + if self.config.gan_loss_weight_gen > 0: + gan = self._logit_mean( + F.softplus( + -torch.where(weight > 0, logits, 0) + if weight is not None + else -logits + ), + weight, + ) + scale = self._ratio_weight(logits, t, weight) + vsd = self._vsd_loss(generated, teacher, fake, scale) + return { + "total_loss": vsd + self.config.gan_loss_weight_gen * gan, + "vsd_loss": vsd, + "gan_loss_gen": gan, + }, self._get_outputs(generated, input_student) + + def _ratio_weight( + self, logits: Tensor, t: Tensor, weight: Tensor | None + ) -> Tensor | None: + r"""DMD2 does not reweight scores with a density ratio.""" + return None + + def _fake_score_discriminator_update_step( + self, + input_student: Tensor, + t_student: Tensor, + t: Tensor, + eps: Tensor, + real_data: Tensor, + condition: Any = None, + ) -> tuple[dict[str, Tensor], dict[str, Tensor]]: + r"""Compute auxiliary losses with student samples held fixed.""" + source = self.net.noise_scheduler + with torch.no_grad(): + generated = self.gen_data_from_net( + input_student, t_student, condition=condition + ) + perturbed = source.forward_process(generated, eps, t) + prediction_type = self.config.fake_score_pred_type + prediction = self.fake_score( + perturbed, t, condition=condition, fwd_pred_type=prediction_type + ) + if self._call_weight is None and prediction_type != "flow": + score_loss = _backend.denoising_score_matching_loss( + prediction_type, + prediction, + x0=generated, + eps=eps, + noise_scheduler=source, + t=t, + ) + else: + targets = {"x0": generated, "eps": eps} + target = ( + source.cond_velocity(generated, eps, t) + if prediction_type == "flow" + else targets[prediction_type] + ) + score_loss = ( + _ObjectiveReduction(prediction, self._call_weight) + .residual(prediction - target, mean=True) + .mean() + ) + gan = score_loss * 0 + r1 = score_loss * 0 + if self.config.gan_loss_weight_gen > 0: + fake_logits, fake_weight = self._logits(perturbed, t, condition) + if self.config.gan_use_same_t_noise: + t_real, eps_real = t, eps + else: + t_real = source.sample_t( + real_data.shape[0], device=real_data.device, **self.config.sampling + ) + eps_real = torch.randn_like(real_data) + real_perturbed = source.forward_process(real_data, eps_real, t_real) + real_logits, real_weight = self._logits(real_perturbed, t_real, condition) + clean_fake = ( + fake_logits + if fake_weight is None + else torch.where(fake_weight > 0, fake_logits, 0) + ) + clean_real = ( + real_logits + if real_weight is None + else torch.where(real_weight > 0, real_logits, 0) + ) + gan = self._logit_mean( + F.softplus(clean_fake), fake_weight + ) + self._logit_mean(F.softplus(-clean_real), real_weight) + if self.config.gan_r1_reg_weight > 0: + jittered = real_data + self.config.gan_r1_reg_alpha * torch.randn_like( + real_data + ) + jitter_logits, jitter_weight = self._logits(jittered, t_real, condition) + support = real_weight + if jitter_weight is not None: + support = ( + jitter_weight + if support is None + else torch.where(jitter_weight > 0, support, 0) + ) + delta = real_logits - jitter_logits + if support is not None: + delta = torch.where(support > 0, delta, 0) + r1 = self._logit_mean(delta.square(), support) + return { + "total_loss": score_loss + gan + self.config.gan_r1_reg_weight * r1, + "fake_score_loss": score_loss, + "gan_loss_disc": gan, + "gan_loss_ar1": r1, + }, self._get_outputs(generated, input_student) + + def _get_outputs( + self, gen_data: Tensor, input_student: Tensor, condition: Any = None + ) -> dict[str, Tensor]: + r"""Return generated samples and their student inputs.""" + return {"prediction": gen_data, "input_student": input_student} diff --git a/physicsnemo/diffusion/distillation/methods/f_distill/__init__.py b/physicsnemo/diffusion/distillation/methods/f_distill/__init__.py new file mode 100644 index 0000000000..af85283aa4 --- /dev/null +++ b/physicsnemo/diffusion/distillation/methods/f_distill/__init__.py @@ -0,0 +1,15 @@ +# 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. diff --git a/physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.py b/physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.py new file mode 100644 index 0000000000..c624555e17 --- /dev/null +++ b/physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.py @@ -0,0 +1,178 @@ +# 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"""f-distill distribution matching with running density-ratio statistics.""" + +import importlib +from collections.abc import Mapping +from typing import Any + +import torch +from jaxtyping import Float +from torch import Tensor + +from physicsnemo.diffusion.noise_schedulers import DomainParallelNoiseScheduler + +from ..._fastgen import _configure_options, _ObjectiveReduction +from ..dmd2._fastgen import _DMD2Loss + + +class _FDistillLoss(_DMD2Loss): + r"""Weight student score matching by a selected f-divergence. + + Use a discriminator to estimate density ratios and running time-bin + statistics to normalize them. Save these statistics with the training + checkpoint, and select a data-replica group for domain-parallel training. + """ + + def __init__( + self, *args: Any, method_options: Mapping[str, Any], **kwargs: Any + ) -> None: + options = dict(method_options) + self._ratio_options = _configure_options( + { + "f_div": "js", + "ratio_lower": 0.1, + "ratio_upper": 20.0, + "ratio_normalization": True, + "bin_num": 10, + "ratio_ema_rate": 0.0, + "statistics_group": None, + }, + options.pop("f_distill", {}), + ) + super().__init__(*args, method_options=options, **kwargs) + if self.config.gan_loss_weight_gen <= 0: + raise ValueError( + "f-distill requires a discriminator and positive gan_loss_weight_gen." + ) + backend = importlib.import_module( + "fastgen.methods.distribution_matching.f_distill" + ) + ratio = self._ratio_options + if ratio.f_div not in backend.all_f_div_weighting_function: + raise ValueError(f"Unsupported f-divergence {ratio.f_div!r}.") + if ( + not 0 < ratio.ratio_lower <= ratio.ratio_upper + or ratio.bin_num < 1 + or not 0 <= ratio.ratio_ema_rate <= 1 + ): + raise ValueError( + "Require positive ratio bounds/bin count and an EMA rate in [0, 1]." + ) + self.f_div_weighting_function = backend.all_f_div_weighting_function[ + ratio.f_div + ] + self._bins = torch.ones(ratio.bin_num) + if ( + isinstance( + self.net.noise_scheduler._sampling_source, DomainParallelNoiseScheduler + ) + and ratio.statistics_group is None + ): + raise ValueError( + "Domain-parallel f-distill requires an explicit data-parallel statistics_group." + ) + + def _sum_statistics(self, value: Tensor) -> Tensor: + r"""Reduce over data replicas without mixing spatial partitions into counts.""" + if torch.distributed.is_initialized(): + torch.distributed.all_reduce( + value, group=self._ratio_options.statistics_group + ) + return value + + def _ratio_weight( + self, + logits: Float[Tensor, " B logits"], + t: Float[Tensor, " B"], + weight: Tensor | None, + ) -> Float[Tensor, " B"]: + r"""Normalize ratios and divergence weights over valid samples only.""" + with torch.no_grad(): + if weight is None: + mean = logits.float().mean(1) + active = torch.ones_like(t, dtype=torch.bool) + else: + valid = weight > 0 + count = valid.sum(1) + mean = torch.where(valid, logits.float(), 0).sum(1) / count.clamp_min(1) + active = count > 0 + if self._call_weight is not None: + active = ( + active + & _ObjectiveReduction(self._call_weight, self._call_weight).active + ) + options = self._ratio_options + ratio = ( + mean.clamp(-10, 10) + .exp() + .clamp(options.ratio_lower, options.ratio_upper) + ) + if options.ratio_normalization: + source = self.net.noise_scheduler + lo, hi = source.min_t, source.max_t + width = (hi - lo) / options.bin_num + idx = ( + ((t.double() - lo) / width) + .floor() + .long() + .clamp(0, options.bin_num - 1) + ) + count = torch.bincount( + idx, weights=active.float(), minlength=options.bin_num + ) + total = torch.bincount( + idx, + weights=torch.where(active, ratio, 0), + minlength=options.bin_num, + ) + statistics = self._sum_statistics(torch.stack((count, total))) + count, total = statistics.unbind(0) + bins = self._bins.to(device=t.device, dtype=torch.float32) + if self._update_state: + new_bins = bins * options.ratio_ema_rate + ( + 1 - options.ratio_ema_rate + ) * total / (count + 1e-6) + bins = torch.where(count > 0, new_bins, bins) + self._bins = bins.detach() + ratio = ratio / (bins[idx] + 1e-6) + h = self.f_div_weighting_function(ratio).expand_as(ratio) + totals = self._sum_statistics( + torch.stack((torch.where(active, h, 0).sum(), active.float().sum())) + ) + h_mean = totals[0] / totals[1].clamp_min(1) + return torch.where(active, h / (h_mean + 1e-6), 0).to(logits.dtype) + + def loss_state_dict(self) -> dict[str, Any]: + r"""Return a copy of the running density-ratio bins for checkpointing.""" + return {"schema_version": 1, "bins": self._bins.detach().clone()} + + def load_loss_state_dict(self, state: Mapping[str, Any]) -> None: + r"""Restore density-ratio bins using the configured bin count.""" + if set(state) != {"schema_version", "bins"} or state["schema_version"] != 1: + raise ValueError("Expected f-distill state schema version 1 with bins.") + bins = state["bins"] + if ( + not isinstance(bins, Tensor) + or bins.shape != self._bins.shape + or not torch.isfinite(bins).all() + or torch.any(bins <= 0) + ): + raise ValueError( + "Expected finite positive ratio bins with the configured shape." + ) + self._bins = bins.detach().clone() diff --git a/test/diffusion/helpers.py b/test/diffusion/helpers.py index a9f61f5c62..076c5b5bdd 100644 --- a/test/diffusion/helpers.py +++ b/test/diffusion/helpers.py @@ -90,6 +90,19 @@ class CustomExtraTimes(DiffusionExtraTimes): r2: Float[torch.Tensor, " B"] | None = None +class FeatureOutputs(DiffusionExtraOutputs): + r"""Model-specific fields selected explicitly by adversarial tests.""" + + spatial: Float[torch.Tensor, " B *dims"] | None = None + + +def extract_teacher_features( + prediction: Float[torch.Tensor, " B *dims"], extras: FeatureOutputs +) -> TensorDict: + r"""Select the teacher field used by the test discriminator.""" + return TensorDict({"spatial": extras.spatial}, batch_size=extras.batch_size) + + class StructuredDiffusionModel(Module): r"""Linear model with optional destination times and learned loss weights.""" @@ -126,7 +139,8 @@ def forward( return prediction if self.output == "empty": return prediction, DiffusionExtraOutputs(batch_size=[x.shape[0]]) - return prediction, DiffusionExtraOutputs( + return prediction, FeatureOutputs( + spatial=x, logvar=( self.gain * t if self.logvar is None else self.logvar.expand(x.shape[0]) ), @@ -154,6 +168,18 @@ def forward( return self.gain * alpha * x / (alpha.square() + sigma.square()) +class FeatureDiscriminator(Module): + r"""Map named spatial features to one logit per sample.""" + + def __init__(self) -> None: + super().__init__() + self.gain = torch.nn.Parameter(torch.tensor(0.4)) + + def forward(self, features: TensorDict) -> Float[torch.Tensor, " B logits"]: + r"""Keep a differentiable input path for generator-gradient checks.""" + return self.gain * features["spatial"].flatten(1).mean(1, keepdim=True) + + class AttentionDiffusionModel(StructuredDiffusionModel): r"""Exercise attention kernels in state and time derivative calculations.""" diff --git a/test/diffusion/test_distribution_distillation.py b/test/diffusion/test_distribution_distillation.py new file mode 100644 index 0000000000..771b39ba7c --- /dev/null +++ b/test/diffusion/test_distribution_distillation.py @@ -0,0 +1,595 @@ +# 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"""Check explicitly selected distribution-matching objectives and role ownership.""" + +import pytest +import torch + +pytest.importorskip("fastgen") + +from physicsnemo.diffusion.distillation import ( + DistillationLoss, + WeightedDistillationLoss, +) +from physicsnemo.diffusion.noise_schedulers import ( + EDMNoiseScheduler, + RectifiedFlowNoiseScheduler, +) + +from .helpers import ( + FeatureDiscriminator, + GaussianDenoiser, + StochasticDiffusionModel, + StructuredDiffusionModel, + extract_teacher_features, +) + + +class TestDMD2: + @pytest.mark.parametrize("empty", (False, True)) + def test_masked_nan_prediction(self, empty): + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + with torch.no_grad(): + fake.gain.fill_(0.3) + loss_fn = WeightedDistillationLoss( + student, + RectifiedFlowNoiseScheduler(), + method="DMD2", + teacher=teacher, + auxiliary_models={"fake_score": fake}, + ) + x = torch.randn(2, 3, 5) + weight = torch.zeros_like(x) if empty else torch.ones_like(x) + weight[:, 1] = 0 + + def corrupt(module, inputs, output): + return torch.where(weight > 0, output[0], float("nan")), output[1] + + hook = student.register_forward_hook(corrupt) + try: + losses = loss_fn( + x, torch.tensor([0.3, 0.7]), phase="student", weight=weight + ) + assert torch.isfinite(losses["total_loss"]) + assert losses["gan_loss_gen"] == 0 + gradient = torch.autograd.grad(losses["total_loss"], student.gain)[0] + assert torch.isfinite(gradient) + if empty: + assert losses["total_loss"] == 0 + assert gradient == 0 + finally: + hook.remove() + + @pytest.mark.parametrize("phase", ("student", "auxiliary")) + def test_masked_statistics(self, phase): + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + with torch.no_grad(): + fake.gain.fill_(0.3) + objective = WeightedDistillationLoss( + student, + RectifiedFlowNoiseScheduler(), + method="DMD2", + teacher=teacher, + auxiliary_models={"fake_score": fake}, + ) + x, t = torch.randn(2, 3, 5), torch.tensor([0.3, 0.7]) + z, noise = torch.randn_like(x), torch.randn_like(x) + weight = torch.ones_like(x) + weight[:, 1] = 0 + corrupted = z.clone() + corrupted[:, 1] = 1e15 + options = dict( + t_student=torch.full_like(t, 0.9), noise=noise, weight=weight, phase=phase + ) + expected = objective(x, t, input_student=z, **options)["total_loss"] + actual = objective(x, t, input_student=corrupted, **options)["total_loss"] + torch.testing.assert_close(actual, expected) + parameter = student.gain if phase == "student" else fake.gain + torch.testing.assert_close( + torch.autograd.grad(actual, parameter)[0], + torch.autograd.grad(expected, parameter)[0], + ) + + def test_explicit_phase(self): + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + loss_fn = DistillationLoss( + student, + RectifiedFlowNoiseScheduler(), + method="DMD2", + teacher=teacher, + auxiliary_models={"fake_score": fake}, + ) + x = torch.randn(2, 3, 5) + with pytest.raises(ValueError, match="phase"): + loss_fn(x) + for phase, target in (("student", student), ("auxiliary", fake)): + for model in (student, teacher, fake): + model.zero_grad(set_to_none=True) + loss = loss_fn(x, torch.tensor([0.3, 0.7]), phase=phase, step=3) + loss["total_loss"].backward() + assert target.gain.grad is not None + for model in (student, teacher, fake): + assert model.gain.requires_grad + if model is not target: + assert model.gain.grad is None + + @pytest.mark.parametrize("phase", ("student", "auxiliary")) + def test_unit_and_zero_weights(self, phase): + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + kwargs = dict( + method="DMD2", teacher=teacher, auxiliary_models={"fake_score": fake} + ) + source = RectifiedFlowNoiseScheduler() + plain = DistillationLoss(student, source, **kwargs) + weighted = WeightedDistillationLoss(student, source, **kwargs) + x, t = torch.randn(2, 3, 5), torch.tensor([0.3, 0.7]) + target = student if phase == "student" else fake + torch.manual_seed(52) + reference = plain(x, t, phase=phase)["total_loss"] + reference_grad = torch.autograd.grad(reference, target.gain)[0] + torch.manual_seed(52) + actual = weighted(x, t, phase=phase, weight=torch.ones_like(x))["total_loss"] + torch.testing.assert_close(actual, reference) + torch.testing.assert_close( + torch.autograd.grad(actual, target.gain)[0], reference_grad + ) + empty = weighted(x, t, phase=phase, weight=torch.zeros_like(x))["total_loss"] + torch.testing.assert_close(empty, torch.zeros_like(empty)) + torch.testing.assert_close( + torch.autograd.grad(empty, target.gain)[0], torch.zeros_like(target.gain) + ) + + @pytest.mark.parametrize("phase", ("student", "auxiliary")) + def test_backend_parity(self, phase): + from fastgen.methods.distribution_matching.dmd2 import DMD2Model + + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + with torch.no_grad(): + fake.gain.fill_(0.25) + loss_fn = DistillationLoss( + student, + RectifiedFlowNoiseScheduler(), + method="DMD2", + teacher=teacher, + auxiliary_models={"fake_score": fake}, + model_options={ + "student": {"scheduler_options": {"time_precision": torch.float32}} + }, + ) + impl = loss_fn._implementation + x = torch.randn(2, 3, 5) + t = torch.tensor([0.3, 0.7], dtype=impl.net.noise_scheduler.t_precision) + z, eps = torch.randn_like(x), torch.randn_like(x) + t_student = torch.full_like(t, 0.9) + actual = loss_fn( + x, t, phase=phase, input_student=z, t_student=t_student, noise=eps + )["total_loss"] + impl.device = x.device + if phase == "student": + expected, _ = DMD2Model._student_update_step( + impl, z, t_student, t, eps, {"real": x} + ) + else: + expected, _ = DMD2Model._fake_score_discriminator_update_step( + impl, z, t_student, t, eps, x + ) + target = student if phase == "student" else fake + torch.testing.assert_close(actual, expected["total_loss"]) + torch.testing.assert_close( + torch.autograd.grad(actual, target.gain)[0], + torch.autograd.grad(expected["total_loss"], target.gain)[0], + ) + + def test_diagnostics_rng(self): + loss_fn = DistillationLoss( + StructuredDiffusionModel(), + RectifiedFlowNoiseScheduler(), + method="DMD2", + teacher=StructuredDiffusionModel(), + auxiliary_models={"fake_score": StructuredDiffusionModel()}, + ) + x = torch.randn(2, 3, 5) + torch.manual_seed(61) + expected = loss_fn(x, phase="student") + state = torch.get_rng_state() + torch.manual_seed(61) + actual, outputs = loss_fn(x, phase="student", return_outputs=True) + torch.testing.assert_close(actual["total_loss"], expected["total_loss"]) + assert torch.equal(state, torch.get_rng_state()) + assert all(not value.requires_grad for value in outputs.values()) + + +@pytest.mark.parametrize("method", ("DMD2", "f-distill")) +class TestAdversarial: + def test_feature_callback_required(self, method): + with pytest.raises(ValueError, match="feature_fn"): + DistillationLoss( + StructuredDiffusionModel(), + RectifiedFlowNoiseScheduler(), + method=method, + teacher=StructuredDiffusionModel(), + auxiliary_models={ + "fake_score": StructuredDiffusionModel(), + "discriminator": FeatureDiscriminator(), + }, + model_options={"teacher": {"feature_names": ("spatial",)}}, + ) + + def test_prediction_required(self, method): + teacher = StructuredDiffusionModel() + hook = teacher.register_forward_hook( + lambda module, args, output: (None, output[1]) + ) + loss_fn = DistillationLoss( + StructuredDiffusionModel(), + RectifiedFlowNoiseScheduler(), + method=method, + teacher=teacher, + auxiliary_models={ + "fake_score": StructuredDiffusionModel(), + "discriminator": FeatureDiscriminator(), + }, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + ) + with pytest.raises(TypeError, match="prediction tensor"): + loss_fn(torch.randn(2, 3, 5), phase="student") + hook.remove() + + @pytest.mark.parametrize("phase", ("student", "auxiliary")) + def test_weighted_logits(self, method, phase): + from tensordict import TensorDict + + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + discriminator = FeatureDiscriminator() + kwargs = dict( + method=method, + teacher=teacher, + auxiliary_models={"fake_score": fake, "discriminator": discriminator}, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + method_options={"gan_use_same_t_noise": True, "gan_r1_reg_weight": 0.5}, + ) + source = RectifiedFlowNoiseScheduler() + plain = DistillationLoss(student, source, **kwargs) + weighted = WeightedDistillationLoss(student, source, **kwargs) + + def adversarial(x, time, condition, weight): + valid = weight > 0 + features = torch.where(valid, x, 0) + logits = discriminator(TensorDict({"spatial": features}, [x.shape[0]])) + return logits, valid.flatten(1).any(1, keepdim=True).to(x.dtype) + + x, t = torch.randn(2, 3, 5), torch.tensor([0.3, 0.7]) + torch.manual_seed(88) + expected = plain(x, t, phase=phase)["total_loss"] + torch.manual_seed(88) + actual = weighted( + x, t, phase=phase, weight=torch.ones_like(x), adversarial_fn=adversarial + )["total_loss"] + torch.testing.assert_close(actual, expected) + targets = ( + (student.gain,) if phase == "student" else (fake.gain, discriminator.gain) + ) + for got, reference in zip( + torch.autograd.grad(actual, targets), torch.autograd.grad(expected, targets) + ): + torch.testing.assert_close(got, reference) + state = weighted.state_dict() + empty = weighted( + x, t, phase=phase, weight=torch.zeros_like(x), adversarial_fn=adversarial + )["total_loss"] + torch.testing.assert_close(empty, torch.zeros_like(empty)) + for gradient in torch.autograd.grad(empty, targets): + torch.testing.assert_close(gradient, torch.zeros_like(gradient)) + if state: + torch.testing.assert_close(weighted.state_dict()["bins"], state["bins"]) + + @pytest.mark.parametrize("phase", ("student", "auxiliary")) + def test_backend_gan_parity(self, method, phase): + from types import SimpleNamespace + + from fastgen.methods.distribution_matching.dmd2 import DMD2Model + from fastgen.methods.distribution_matching.f_distill import FdistillModel + + student, teacher, fake = (StochasticDiffusionModel() for _ in range(3)) + discriminator = FeatureDiscriminator() + objective = DistillationLoss( + student, + RectifiedFlowNoiseScheduler(), + method=method, + teacher=teacher, + auxiliary_models={"fake_score": fake, "discriminator": discriminator}, + model_options={ + "student": {"scheduler_options": {"time_precision": torch.float32}}, + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + }, + }, + method_options={"gan_use_same_t_noise": True, "gan_r1_reg_weight": 0.5}, + ) + impl = objective._implementation + x, t = torch.randn(2, 3, 5), torch.tensor([0.3, 0.7]) + z, eps, ts = torch.randn_like(x), torch.randn_like(x), torch.full_like(t, 0.9) + torch.manual_seed(89) + actual = objective(x, t, phase=phase, input_student=z, t_student=ts, noise=eps)[ + "total_loss" + ] + state = torch.get_rng_state() + if method == "f-distill": + impl.config.f_distill = SimpleNamespace(**vars(impl._ratio_options)) + impl.bin_num = impl._ratio_options.bin_num + impl.net.bins = torch.ones(impl.bin_num) + impl._get_f_div_weighting_h = lambda logits, time: ( + FdistillModel._get_f_div_weighting_h(impl, logits, time) + ) + torch.manual_seed(89) + if phase == "student": + backend = FdistillModel if method == "f-distill" else DMD2Model + expected, _ = backend._student_update_step(impl, z, ts, t, eps, {"real": x}) + else: + expected, _ = DMD2Model._fake_score_discriminator_update_step( + impl, z, ts, t, eps, x + ) + torch.testing.assert_close(actual, expected["total_loss"]) + assert torch.equal(torch.get_rng_state(), state) + targets = ( + (student.gain,) if phase == "student" else (fake.gain, discriminator.gain) + ) + for got, reference in zip( + torch.autograd.grad(actual, targets), + torch.autograd.grad(expected["total_loss"], targets), + ): + torch.testing.assert_close(got, reference) + + def test_teacher_called_once(self, method): + teacher = StructuredDiffusionModel() + calls = [] + hook = teacher.register_forward_hook(lambda *args: calls.append(1)) + loss_fn = DistillationLoss( + StructuredDiffusionModel(), + RectifiedFlowNoiseScheduler(), + method=method, + teacher=teacher, + auxiliary_models={ + "fake_score": StructuredDiffusionModel(), + "discriminator": FeatureDiscriminator(), + }, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + ) + loss_fn(torch.randn(2, 3, 5), phase="student") + hook.remove() + assert len(calls) == 1 + + @pytest.mark.parametrize("phase", ("student", "auxiliary")) + def test_role_gradients(self, method, phase): + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + discriminator = FeatureDiscriminator() + loss_fn = DistillationLoss( + student, + RectifiedFlowNoiseScheduler(), + method=method, + teacher=teacher, + auxiliary_models={"fake_score": fake, "discriminator": discriminator}, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + method_options={"gan_r1_reg_weight": 0.5}, + ) + x = torch.randn(2, 3, 5) + loss_fn(x, torch.tensor([0.3, 0.7]), phase=phase)["total_loss"].backward() + assert (student.gain.grad is not None) == (phase == "student") + assert teacher.gain.grad is None + assert (fake.gain.grad is not None) == (phase == "auxiliary") + assert (discriminator.gain.grad is not None) == (phase == "auxiliary") + assert all( + model.gain.requires_grad + for model in (student, teacher, fake, discriminator) + ) + + def test_mask_capability_required(self, method): + loss_fn = WeightedDistillationLoss( + StructuredDiffusionModel(), + RectifiedFlowNoiseScheduler(), + method=method, + teacher=StructuredDiffusionModel(), + auxiliary_models={ + "fake_score": StructuredDiffusionModel(), + "discriminator": FeatureDiscriminator(), + }, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + ) + x = torch.randn(2, 3, 5) + rng = torch.get_rng_state() + with pytest.raises(ValueError, match="adversarial_fn"): + loss_fn(x, phase="student", weight=torch.ones_like(x)) + assert torch.equal(torch.get_rng_state(), rng) + + +class TestRatioState: + @pytest.mark.parametrize( + "divergence", ("rkl", "kl", "js", "sf", "neyman", "sh", "jf") + ) + def test_divergences(self, divergence): + loss_fn = DistillationLoss( + StructuredDiffusionModel(), + RectifiedFlowNoiseScheduler(), + method="f-distill", + teacher=StructuredDiffusionModel(), + auxiliary_models={ + "fake_score": StructuredDiffusionModel(), + "discriminator": FeatureDiscriminator(), + }, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + method_options={"f_distill": {"f_div": divergence}}, + ) + loss = loss_fn(torch.randn(2, 3, 5), phase="student")["total_loss"] + assert torch.isfinite(loss) + loss.backward() + + def test_update_and_restore(self): + options = dict( + method="f-distill", + teacher=StructuredDiffusionModel(), + auxiliary_models={ + "fake_score": StructuredDiffusionModel(), + "discriminator": FeatureDiscriminator(), + }, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + ) + loss_fn = DistillationLoss( + StructuredDiffusionModel(), RectifiedFlowNoiseScheduler(), **options + ) + x, t = torch.randn(2, 3, 5), torch.tensor([0.3, 0.7]) + before = loss_fn.state_dict() + loss_fn(x, t, phase="student", update_state=False) + torch.testing.assert_close(loss_fn.state_dict()["bins"], before["bins"]) + loss_fn(x, t, phase="student") + state = loss_fn.state_dict() + assert not torch.equal(state["bins"], before["bins"]) + fresh = DistillationLoss( + StructuredDiffusionModel(), RectifiedFlowNoiseScheduler(), **options + ) + fresh.load_state_dict(state) + torch.manual_seed(74) + expected = loss_fn(x, t, phase="student")["total_loss"] + torch.manual_seed(74) + actual = fresh(x, t, phase="student")["total_loss"] + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + fresh.state_dict()["bins"], loss_fn.state_dict()["bins"] + ) + + +class TestDistributionIntegration: + def test_custom_role_coordinates(self): + source, teacher_source = RectifiedFlowNoiseScheduler(), EDMNoiseScheduler() + student, fake = StructuredDiffusionModel(), StructuredDiffusionModel() + common = dict(method="DMD2", auxiliary_models={"fake_score": fake}) + native = DistillationLoss( + student, source, teacher=GaussianDenoiser(source), **common + ) + converted = DistillationLoss( + student, + source, + teacher=GaussianDenoiser(teacher_source), + **common, + model_options={"teacher": {"noise_scheduler": teacher_source}}, + ) + x, t = torch.randn(2, 3, 5), torch.tensor([0.3, 0.7]) + torch.manual_seed(94) + expected = native(x, t, phase="student")["total_loss"] + torch.manual_seed(94) + actual = converted(x, t, phase="student")["total_loss"] + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + torch.autograd.grad(actual, student.gain)[0], + torch.autograd.grad(expected, student.gain)[0], + ) + + @pytest.mark.parametrize("phase", ("student", "auxiliary")) + def test_user_compiled_roles(self, phase): + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + source = RectifiedFlowNoiseScheduler() + eager = DistillationLoss( + student, + source, + method="DMD2", + teacher=teacher, + auxiliary_models={"fake_score": fake}, + ) + compiled = DistillationLoss( + torch.compile(student, backend="eager", fullgraph=True), + source, + method="DMD2", + teacher=torch.compile(teacher, backend="eager", fullgraph=True), + auxiliary_models={ + "fake_score": torch.compile(fake, backend="eager", fullgraph=True) + }, + ) + for _ in range(2): + x, t = torch.randn(2, 3, 5), torch.tensor([0.3, 0.7]) + torch.manual_seed(95) + expected = eager(x, t, phase=phase)["total_loss"] + torch.manual_seed(95) + actual = compiled(x, t, phase=phase)["total_loss"] + torch.testing.assert_close(actual, expected) + target = student.gain if phase == "student" else fake.gain + torch.testing.assert_close( + torch.autograd.grad(actual, target)[0], + torch.autograd.grad(expected, target)[0], + ) + + def test_exception_restores_freeze_mask(self): + student, teacher, fake = (StructuredDiffusionModel() for _ in range(3)) + discriminator = FeatureDiscriminator() + teacher.requires_grad_(False) + loss_fn = WeightedDistillationLoss( + student, + RectifiedFlowNoiseScheduler(), + method="DMD2", + teacher=teacher, + auxiliary_models={"fake_score": fake, "discriminator": discriminator}, + model_options={ + "teacher": { + "feature_names": ("spatial",), + "feature_fn": extract_teacher_features, + } + }, + ) + + def fail(*args): + raise RuntimeError("test callback failure") + + x = torch.randn(2, 3, 5) + with pytest.raises(RuntimeError, match="test callback failure"): + loss_fn(x, phase="student", weight=torch.ones_like(x), adversarial_fn=fail) + assert not teacher.gain.requires_grad + assert all(model.gain.requires_grad for model in (student, fake, discriminator)) + assert loss_fn._implementation._call_weight is None