diff --git a/physicsnemo/diffusion/distillation/_fastgen.py b/physicsnemo/diffusion/distillation/_fastgen.py index b4b2bac289..fc9127f651 100644 --- a/physicsnemo/diffusion/distillation/_fastgen.py +++ b/physicsnemo/diffusion/distillation/_fastgen.py @@ -16,6 +16,7 @@ r"""Connect PhysicsNeMo models to FastGen distillation objectives.""" +import importlib from collections.abc import Callable, Mapping from types import SimpleNamespace from typing import Any @@ -26,6 +27,7 @@ from torch import Tensor from torch.nn.attention import SDPBackend, sdpa_kernel +from physicsnemo.core.version_check import check_version_spec from physicsnemo.diffusion._fastgen import _FastGenNetwork from physicsnemo.diffusion.base import ( DiffusionExtraTimes, @@ -36,44 +38,12 @@ ) from physicsnemo.diffusion.noise_schedulers._fastgen import _FastGenNoiseScheduler - -class _ObjectiveReduction: - r"""Compute masked statistics and weighted squared residuals. - - Statistics use :math:`m=1_{w>0}` and count :math:`\sum m`. Residuals - use :math:`w` with the full element count for means. Positive weights - scale residuals linearly and leave the adaptive target unchanged. - """ - - def __init__( - self, - reference: Float[Tensor, " B *dims"], - weight: Float[Tensor, " #B *#dims"] | None, - ) -> None: - self.weight = ( - torch.ones_like(reference) - if weight is None - else torch.broadcast_to(weight, reference.shape) - ) - self.valid = self.weight > 0 - self.dims = tuple(range(1, reference.ndim)) - self.count = self.valid.sum(self.dims) - self.active = self.count > 0 - - def clean(self, value: Float[Tensor, " B *dims"]) -> Float[Tensor, " B *dims"]: - r"""Remove invalid entries before nonlinear operations.""" - return torch.where(self.valid, value, torch.zeros_like(value)) - - def norm(self, value: Float[Tensor, " B *dims"]) -> Float[Tensor, " B"]: - r"""Compute the Euclidean norm on valid entries only.""" - return torch.linalg.vector_norm(self.clean(value).flatten(1), dim=1) - - def residual( - self, value: Float[Tensor, " B *dims"], *, mean: bool - ) -> Float[Tensor, " B"]: - r"""Weight squared residuals linearly, retaining the backend denominator.""" - weighted = self.clean(value).square() * self.weight - return weighted.mean(self.dims) if mean else weighted.sum(self.dims) +if check_version_spec("fastgen", "0.1.0", hard_fail=False): + _ObjectiveReduction = importlib.import_module( + "fastgen.methods.loss_support" + ).LossReduction +else: + _ObjectiveReduction = None class _FastGenBinding: diff --git a/physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py b/physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py index 9d47806847..eed05f11a0 100644 --- a/physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py +++ b/physicsnemo/diffusion/distillation/methods/dmd2/_fastgen.py @@ -22,7 +22,6 @@ 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 @@ -40,7 +39,7 @@ ) from physicsnemo.diffusion.noise_schedulers._fastgen import _FastGenNoiseScheduler -from ..._fastgen import _configure_options, _FastGenBinding, _ObjectiveReduction +from ..._fastgen import _configure_options, _FastGenBinding _AVAILABLE = check_version_spec("fastgen", "0.1.0", hard_fail=False) if _AVAILABLE: @@ -307,15 +306,13 @@ def single_train_step( ) -> 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"] + return self._dispatch_training_phase( + self._phase, + *args, + data, + data["real"], + condition=data["condition"], + neg_condition=data["neg_condition"], ) def _vsd_loss( @@ -325,24 +322,14 @@ def _vsd_loss( 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() + r"""Compute weighted score distillation with normalization over valid entries.""" + return _backend.variational_score_distillation_loss( + generated, + teacher, + fake, + additional_scale=scale, + loss_weight=self._call_weight, + ) def _logits( self, x: Tensor, t: Tensor, condition: Any @@ -421,14 +408,7 @@ def _student_update_step( 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, - ) + gan = _backend.gan_loss_generator(logits, loss_weight=weight) scale = self._ratio_weight(logits, t, weight) vsd = self._vsd_loss(generated, teacher, fake, scale) return { @@ -463,27 +443,18 @@ def _fake_score_discriminator_update_step( 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() - ) + score_loss = _backend.denoising_score_matching_loss( + prediction_type, + prediction, + x0=generated, + eps=eps, + noise_scheduler=source, + t=t, + loss_weight=self._call_weight, + flow_target=source.cond_velocity(generated, eps, t) + if prediction_type == "flow" + else None, + ) gan = score_loss * 0 r1 = score_loss * 0 if self.config.gan_loss_weight_gen > 0: @@ -497,19 +468,12 @@ def _fake_score_discriminator_update_step( 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 = _backend.gan_loss_discriminator( + real_logits, + fake_logits, + real_weight=real_weight, + fake_weight=fake_weight, ) - 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 diff --git a/physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.py b/physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.py index c624555e17..7ddd0d0459 100644 --- a/physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.py +++ b/physicsnemo/diffusion/distillation/methods/f_distill/_fastgen.py @@ -62,6 +62,7 @@ def __init__( backend = importlib.import_module( "fastgen.methods.distribution_matching.f_distill" ) + self._loss_support = importlib.import_module("fastgen.methods.loss_support") 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}.") @@ -87,75 +88,37 @@ def __init__( "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) + r"""Normalize divergence weights over valid samples, optionally updating running bins.""" + options = self._ratio_options + sample_mask = ( + None + if self._call_weight is None + else _ObjectiveReduction(self._call_weight, self._call_weight).active + ) + h, bins = self._loss_support.f_divergence_weights( + logits, + t, + self.f_div_weighting_function, + ratio_lower=options.ratio_lower, + ratio_upper=options.ratio_upper, + bins=self._bins if options.ratio_normalization else None, + min_t=self.net.noise_scheduler.min_t, + max_t=self.net.noise_scheduler.max_t, + ema_rate=options.ratio_ema_rate, + logit_weight=weight, + sample_mask=sample_mask, + statistics_group=options.statistics_group, + update_state=self._update_state, + ) + if options.ratio_normalization and self._update_state: + self._bins = bins.detach() + return h def loss_state_dict(self) -> dict[str, Any]: r"""Return a copy of the running density-ratio bins for checkpointing.""" diff --git a/physicsnemo/diffusion/distillation/methods/meanflow/_fastgen.py b/physicsnemo/diffusion/distillation/methods/meanflow/_fastgen.py index 77560d8940..b07b5b7cfc 100644 --- a/physicsnemo/diffusion/distillation/methods/meanflow/_fastgen.py +++ b/physicsnemo/diffusion/distillation/methods/meanflow/_fastgen.py @@ -34,7 +34,7 @@ LinearGaussianNoiseScheduler, ) -from ..._fastgen import _configure_options, _FastGenBinding, _ObjectiveReduction +from ..._fastgen import _configure_options, _FastGenBinding _AVAILABLE = check_version_spec("fastgen", "0.1.0", hard_fail=False) if _AVAILABLE: @@ -129,56 +129,14 @@ def _get_velocity( module.training = mode def _mf_pred_to_loss( - self, - u_theta: Float[Tensor, " B *dims"], - u_theta_jvp: Float[Tensor, " B *dims"], - x_t: Float[Tensor, " B *dims"], - dxt_dt: Float[Tensor, " B *dims"], - t: Float[Tensor, " B"], - r: Float[Tensor, " B"], - iteration: int, + self, *args: Any, **kwargs: Any ) -> tuple[ Float[Tensor, " B"], Float[Tensor, " B *dims"], Float[Tensor, " *shape"], float ]: - r"""Exclude masked entries from adaptive norms before weighting residuals. - - .. math:: - - h=v-\lambda(t-r)\partial_t u,\qquad - L_{\mathrm{L2}}=\operatorname{reduce}(w(u-\operatorname{sg}(h))^2). - - The explicit-gradient variant also subtracts the stopped prediction - from its tangent, retaining FastGen's adaptive weighting function. - """ - u = u_theta.double() - reduction = _ObjectiveReduction(u, self._call_weight) - shape = (-1,) + (1,) * (u.ndim - 1) - delta = (t.double() - r.double()).reshape(shape) # (B, 1, ...) - steps = self.loss_config.tangent_warmup_steps - warmup = min(1.0, iteration / steps) if steps > 0 else 1.0 - tangent = dxt_dt.double() - warmup * delta * u_theta_jvp.double() - if self.loss_config.loss_type == "l2": - tangent = reduction.clean(tangent) - residual = reduction.clean(u - tangent) - statistic = residual.square().sum(reduction.dims) - if self.loss_config.norm_method is None: - statistic = statistic / reduction.count.clamp_min(1) - adaptive = self._compute_weight(statistic, t) - loss = ( - reduction.residual(residual, mean=self.loss_config.norm_method is None) - * adaptive - ) - else: - tangent = reduction.clean(tangent - u.detach()) - if self.loss_config.tangent_spatial_invariance: - tangent = tangent / reduction.count.clamp_min(1).sqrt().reshape( - shape - ) # (B, 1, ...) - adaptive = self._compute_weight(reduction.norm(tangent), t) - target = (u + tangent * adaptive.reshape(shape)).detach() # (B, 1, ...) - loss = reduction.residual(u - target, mean=False) - loss = torch.where(reduction.active, loss, torch.zeros_like(loss)) - return loss, tangent, adaptive, warmup + r"""Compute the weighted MeanFlow loss without clipping native time intervals.""" + return super()._mf_pred_to_loss( + *args, loss_weight=self._call_weight, clip_time_interval=False, **kwargs + ) def single_train_step( self, data: Mapping[str, Any], iteration: int diff --git a/physicsnemo/diffusion/distillation/methods/scm/_fastgen.py b/physicsnemo/diffusion/distillation/methods/scm/_fastgen.py index 2657479902..c2d376eb0c 100644 --- a/physicsnemo/diffusion/distillation/methods/scm/_fastgen.py +++ b/physicsnemo/diffusion/distillation/methods/scm/_fastgen.py @@ -35,7 +35,7 @@ ) from physicsnemo.diffusion.noise_schedulers._fastgen import _FastGenNoiseScheduler -from ..._fastgen import _configure_options, _FastGenBinding, _ObjectiveReduction +from ..._fastgen import _configure_options, _FastGenBinding _AVAILABLE = check_version_spec("fastgen", "0.1.0", hard_fail=False) if _AVAILABLE: @@ -105,67 +105,10 @@ def __init__( self.teacher_trigflow.noise_scheduler = self.net_trigflow.noise_scheduler def _scm_pred_to_loss( - self, - F_theta: Float[Tensor, " B *dims"], - F_theta_jvp: Float[Tensor, " B *dims"], - x_t_hat: Float[Tensor, " B *dims"], - dxt_dt: Float[Tensor, " B *dims"], - logvar: Float[Tensor, "B 1"], - sigma_t: Float[Tensor, " B"], - t_hat: Float[Tensor, " B"], - iteration: int, + self, *args: Any, **kwargs: Any ) -> tuple[Float[Tensor, " B"], Float[Tensor, " B"], float]: - r"""Use valid-support tangent norms and linearly weighted residuals. - - .. math:: - - g = -\alpha^2(\sigma_d\operatorname{sg}(F)-v) - -r(\alpha\sigma x+\sigma_d\operatorname{JVP}(F)),\qquad - L = w_t e^{-s}\operatorname{mean}(w(F-\operatorname{sg}(F)-\bar g)^2)/D+s. - - Zero-support samples omit both residual and log-variance terms. - """ - if self._call_weight is None: - return super()._scm_pred_to_loss( - F_theta, F_theta_jvp, x_t_hat, dxt_dt, logvar, sigma_t, t_hat, iteration - ) - F_theta = F_theta.double() - reduction = _ObjectiveReduction(F_theta, self._call_weight) - shape = (-1,) + (1,) * (F_theta.ndim - 1) - alpha = self.net_trigflow.noise_scheduler.alpha(t_hat.double()).reshape( - shape - ) # (B, 1, ...) - sigma = self.net_trigflow.noise_scheduler.sigma(t_hat.double()).reshape( - shape - ) # (B, 1, ...) - warmup = min(1.0, iteration / self.loss_config.tangent_warmup_steps) - g = -alpha.square() * ( - self.sigma_data * F_theta.detach() - dxt_dt.double() - ) - warmup * ( - alpha * sigma * x_t_hat.double() + self.sigma_data * F_theta_jvp.double() - ) - g = reduction.clean(g) - norm = reduction.norm(g) - if self.loss_config.g_norm_spatial_invariance: - norm = norm / reduction.count.clamp_min(1).sqrt() - g = g / (norm + self.loss_config.tangent_warmup_const).reshape( - shape - ) # (B, 1, ...) - unweighted = reduction.residual(F_theta - F_theta.detach() - g, mean=True) - dimension = ( - reduction.count.clamp_min(1) - if self.loss_config.divide_x_0_spatial_dim - else 1 - ) - prior = sigma_t.reciprocal() if self.loss_config.prior_weighting_enabled else 1 - s = logvar[:, 0].double() - s = torch.where(reduction.active, s, torch.zeros_like(s)) - loss = prior * s.neg().exp() * unweighted / dimension + s - return ( - torch.where(reduction.active, loss, torch.zeros_like(loss)), - unweighted, - warmup, - ) + r"""Compute the weighted consistency loss with tangent norms over valid entries.""" + return super()._scm_pred_to_loss(*args, loss_weight=self._call_weight, **kwargs) def single_train_step( self, data: Mapping[str, Any], iteration: int diff --git a/pyproject.toml b/pyproject.toml index 03bcd77c0f..7ba67d1f5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -270,7 +270,7 @@ diffusion-strategy = [ "nvalchemi-toolkit @ git+https://github.com/NVIDIA/nvalchemi-toolkit.git@f59e28dc6701583a13bb294e536076c2a7379e98", ] diffusion-fastgen = [ - "fastgen @ git+https://github.com/CharlelieLrt/FastGen.git@cd7de8c386a92da2be4cbc3bf0558f16153e46b9", + "fastgen @ git+https://github.com/CharlelieLrt/FastGen.git@81814e4f6a0863682a08c4823b22c93258dd3216", ] # CUDA backend extras (mutually exclusive via [tool.uv] conflicts). # Combine with feature extras, e.g. nvidia-physicsnemo[nn-extras,cu13]. diff --git a/test/diffusion/test_distillation_losses.py b/test/diffusion/test_distillation_losses.py index 33b7c34be3..8739e63b90 100644 --- a/test/diffusion/test_distillation_losses.py +++ b/test/diffusion/test_distillation_losses.py @@ -337,6 +337,25 @@ def fn(x): class TestGradientFlow: + @pytest.mark.parametrize("loss_type", ("l2", "opt_grad")) + def test_input_gradients(self, loss_type): + student = StructuredDiffusionModel() + objective = DistillationLoss( + student, + RectifiedFlowNoiseScheduler(), + method="MeanFlow", + method_options={"loss_type": loss_type}, + ) + x = torch.randn(2, 3, 5, requires_grad=True) + time = ( + torch.tensor([0.3, 0.7]), + DiffusionExtraTimes(r=torch.tensor([0.1, 0.2]), batch_size=[2]), + ) + result = objective(x, time, step=1000)["total_loss"] + gradient = torch.autograd.grad(result, x)[0] + assert torch.isfinite(gradient).all() + assert torch.count_nonzero(gradient) > 0 + @pytest.mark.parametrize("training", (False, True), ids=("eval", "mixed")) def test_guidance_preserves_modes(self, training): student = StochasticDiffusionModel().train(training) diff --git a/uv.lock b/uv.lock index 6c1b801c27..53ee2eac52 100644 --- a/uv.lock +++ b/uv.lock @@ -3975,7 +3975,7 @@ wheels = [ [[package]] name = "fastgen" version = "0.1.0" -source = { git = "https://github.com/CharlelieLrt/FastGen.git?rev=cd7de8c386a92da2be4cbc3bf0558f16153e46b9#cd7de8c386a92da2be4cbc3bf0558f16153e46b9" } +source = { git = "https://github.com/CharlelieLrt/FastGen.git?rev=81814e4f6a0863682a08c4823b22c93258dd3216#81814e4f6a0863682a08c4823b22c93258dd3216" } dependencies = [ { name = "accelerate" }, { name = "av", version = "18.1.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')" }, @@ -9677,7 +9677,7 @@ requires-dist = [ { name = "cupy-cuda13x", marker = "extra == 'cu13'", specifier = ">=14.2.0" }, { name = "dask", marker = "extra == 'datapipes-extras'", specifier = ">=2026.8.0" }, { name = "einops", specifier = ">=0.8.2" }, - { name = "fastgen", marker = "extra == 'diffusion-fastgen'", git = "https://github.com/CharlelieLrt/FastGen.git?rev=cd7de8c386a92da2be4cbc3bf0558f16153e46b9" }, + { name = "fastgen", marker = "extra == 'diffusion-fastgen'", git = "https://github.com/CharlelieLrt/FastGen.git?rev=81814e4f6a0863682a08c4823b22c93258dd3216" }, { name = "fsspec", specifier = ">=2026.9.0" }, { name = "gitpython", specifier = ">=3.1.62" }, { name = "gpytorch", marker = "extra == 'uq-extras'", specifier = ">=1.15.2" },