From a946b0aba09080ef1cab31c12519f0eb48016fd4 Mon Sep 17 00:00:00 2001 From: Charlelie Laurent Date: Tue, 13 Oct 2026 22:35:52 -0700 Subject: [PATCH] Add composable distillation loss hooks Signed-off-by: Charlelie Laurent --- TRAINING_INTEGRATION.md | 89 +++++++++++ fastgen/methods/common_loss.py | 103 +++++++++--- .../methods/consistency_model/mean_flow.py | 85 +++++++--- fastgen/methods/consistency_model/sCM.py | 33 +++- fastgen/methods/distribution_matching/dmd2.py | 113 +++++++++---- .../distribution_matching/f_distill.py | 98 +++++------- fastgen/methods/loss_support.py | 145 +++++++++++++++++ tests/test_dmd2model.py | 8 + tests/test_loss_support.py | 150 ++++++++++++++++++ 9 files changed, 687 insertions(+), 137 deletions(-) create mode 100644 TRAINING_INTEGRATION.md create mode 100644 fastgen/methods/loss_support.py create mode 100644 tests/test_loss_support.py diff --git a/TRAINING_INTEGRATION.md b/TRAINING_INTEGRATION.md new file mode 100644 index 0000000..021b0a7 --- /dev/null +++ b/TRAINING_INTEGRATION.md @@ -0,0 +1,89 @@ +# Custom training integration + +Use FastGen losses in your own training loop with your models, optimizers, +and checkpoints. FastGen's trainer uses the default unmasked losses and +iteration-based phase selection. Supply weights, select phases, or manage +ratio statistics explicitly when your loop needs that control. + +## Mask and weight losses + +`variational_score_distillation_loss`, `denoising_score_matching_loss`, and +`gan_loss_generator` accept a keyword-only `loss_weight`. The discriminator loss +accepts `real_weight` and `fake_weight`. Weights must be finite, nonnegative, and +broadcast-compatible with the predictions or discriminator scores. Positive +weights scale squared residuals or score losses linearly. Zero weights exclude +entries from the loss and from VSD target normalization. Means divide by the +full element count, including masked entries. Omit +weights to include every entry with unit weight. + +```python +loss = variational_score_distillation_loss( + generated, teacher_x0, fake_x0, loss_weight=pixel_weight, +) +``` + +The sCM and MeanFlow `_scm_pred_to_loss` and `_mf_pred_to_loss` extension points +accept the same keyword to exclude masked entries from tangent norms and +adaptive weights. Samples with no valid entries contribute zero loss, +including the sCM log-variance term. + +MeanFlow clips `t - r` to the noise scheduler's time bounds by default. Set +`clip_time_interval=False` when `t - r` already expresses the duration required +by your custom time parameterization. For DSM flow prediction, supply +`flow_target` to replace the default velocity `eps - x0` with a conditional +velocity from your custom path. + +For adversarial losses, exclude corrupted inputs during feature extraction +and pooling, then supply weights for the resulting discriminator scores. +A pixel mask alone cannot identify valid scores. Masking final scores cannot +undo corruption that has already spread through the backbone. + +## Phase selection + +```python +losses, outputs = model.single_train_step(batch, iteration, phase="auxiliary") +``` + +DMD2 and f-distill accept `phase="student"` or `phase="auxiliary"` +independently of iteration. Pass the actual training step as `iteration` for +warmup schedules, and step the optimizer matching your selected phase. Omit +`phase` to follow `student_update_freq`, as FastGen's trainer does. + +`single_train_step` prepares inputs, samples noise, and sets model modes and +gradients for the selected phase. Accept the `phase` keyword when overriding +noise generation in a subclass. Use `_dispatch_training_phase` with prepared +inputs when your loop also manages sampling, model modes, and frozen parameters. + +## Ratio state and process groups + +Use `f_divergence_weights` to compute f-distill weights while keeping running +ratio statistics in your own training state: + +```python +h, new_bins = f_divergence_weights( + scores, t, weighting_fn, bins=bins, + min_t=min_t, max_t=max_t, + logit_weight=score_weight, + statistics_group=data_group, + update_state=True, +) +``` + +Save the returned bins for the next call or in a checkpoint. The function does +not change the input bins. Set `update_state=False` to normalize with saved +statistics without updating them. Omit `bins` to skip bin normalization. +Zero `logit_weight` entries exclude scores from ratio statistics. Positive +magnitudes only mark scores as valid. Use `sample_mask` to exclude whole +samples. Samples with no valid scores receive zero weight. + +Call the function on every rank in `statistics_group`, even if a rank has no +valid samples. This group should join data replicas, not spatial partitions +of one sample. Combine spatial features within your model when needed. During +distributed execution, omitting `statistics_group` uses the +default process group. + +Normalization uses the actual global valid count, supporting unequal batch +sizes. Scalar divergence functions broadcast to the batch, including reverse KL. +FastGen's f-distill method uses this function and stores bins on its student +network when you enable ratio normalization. Its checkpoint format stays +unchanged. External loops can instead checkpoint the returned bins. diff --git a/fastgen/methods/common_loss.py b/fastgen/methods/common_loss.py index 51b2bdb..87eb5c3 100644 --- a/fastgen/methods/common_loss.py +++ b/fastgen/methods/common_loss.py @@ -7,6 +7,7 @@ from fastgen.networks.noise_schedule import BaseNoiseSchedule from fastgen.utils import expand_like +from fastgen.methods.loss_support import LossReduction def denoising_score_matching_loss( @@ -16,13 +17,17 @@ def denoising_score_matching_loss( eps: torch.Tensor = None, noise_scheduler: Optional[BaseNoiseSchedule] = None, t: Optional[torch.Tensor] = None, + *, + loss_weight: Optional[torch.Tensor] = None, + flow_target: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Compute the denoising diffusion objective. - Forward process: - x_t = alpha_t * x_0 + sigma_t * eps + """Compute mean squared error against the selected denoising target. - Note: We currently don't add any loss weighting for simplicity. In the future, we may include - time-dependent weighting (e.g., by SNR, variance schedule). + The forward process is ``x_t = alpha_t * x_0 + sigma_t * eps``. + By default, every residual has unit weight. Supply ``loss_weight`` for + masking or time-dependent weighting. The mean divides by the full element + count. For ``pred_type="flow"``, ``flow_target`` + replaces the default velocity ``eps - x0`` for a custom diffusion path. Args: pred_type (str): Prediction type: 'x0', 'eps', 'v', 'flow'. @@ -31,6 +36,10 @@ def denoising_score_matching_loss( x0 (torch.Tensor): The clean data x_0. eps (torch.Tensor): The epsilon used to compute noised data. t (torch.Tensor): The target data t. + loss_weight (Optional[torch.Tensor]): Finite nonnegative weights that broadcast + to net_pred. Zero entries exclude residuals. Positive values scale squared errors. + flow_target (Optional[torch.Tensor]): Conditional velocity target used only + for flow prediction. Defaults to eps - x0. Raises: NotImplementedError: If an unknown pred_type is used. @@ -40,24 +49,28 @@ def denoising_score_matching_loss( """ if pred_type == "x0": assert x0 is not None, "x0 cannot be None" - loss = F.mse_loss(x0, net_pred, reduction="mean") + target = x0 elif pred_type == "eps": assert eps is not None, "eps cannot be None" - loss = F.mse_loss(eps, net_pred, reduction="mean") + target = eps elif pred_type == "v": assert x0 is not None and eps is not None and t is not None, "x0, eps, t should not be None" assert noise_scheduler is not None, "noise_scheduler should not be None" alpha_t = expand_like(noise_scheduler.alpha(t), x0).to(device=x0.device, dtype=x0.dtype) sigma_t = expand_like(noise_scheduler.sigma(t), x0).to(device=x0.device, dtype=x0.dtype) v = alpha_t * eps - sigma_t * x0 - loss = F.mse_loss(v, net_pred, reduction="mean") + target = v elif pred_type == "flow": - assert x0 is not None and eps is not None, "x0 and eps cannot be None" - flow_velocity = eps - x0 - loss = F.mse_loss(flow_velocity, net_pred, reduction="mean") + if flow_target is None: + assert x0 is not None and eps is not None, "x0 and eps cannot be None" + target = eps - x0 + else: + target = flow_target else: raise NotImplementedError(f"Unknown prediction type {pred_type}") - return loss + if loss_weight is None: + return F.mse_loss(target, net_pred, reduction="mean") + return LossReduction(net_pred, loss_weight).residual(net_pred - target, mean=True).mean() def variational_score_distillation_loss( @@ -65,19 +78,31 @@ def variational_score_distillation_loss( teacher_x0: torch.Tensor, fake_score_x0: torch.Tensor, additional_scale: Optional[torch.Tensor] = None, + *, + loss_weight: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Compute the variational score distillation loss. + + By default, all entries contribute to the target normalization and loss. + Supply ``loss_weight`` to mask entries or scale squared residuals. Zero + weights also exclude entries from the normalization statistics. Positive + magnitudes affect only the residual. The mean divides by the full element + count. + Args: gen_data (torch.Tensor): generated data teacher_x0 (torch.Tensor): x0-prediction from the teacher fake_score_x0 (torch.Tensor): x0-prediction from the fake score additional_scale (Optional[torch.Tensor]): Additional scale parameter for the VSD loss. + loss_weight (Optional[torch.Tensor]): Finite nonnegative weights that broadcast + to gen_data. Zero entries exclude residuals and target normalization statistics. Returns: loss (torch.Tensor): The variational score distillation loss. """ dims = tuple(range(1, teacher_x0.ndim)) + reduction = LossReduction(gen_data, loss_weight) if loss_weight is not None else None with torch.no_grad(): # Perform weight calculation in fp32 for numerical stability @@ -86,7 +111,12 @@ def variational_score_distillation_loss( teacher_x0_fp32 = teacher_x0.float() # Compute weight in fp32 to avoid numerical instability - diff_abs_mean = (gen_data_fp32 - teacher_x0_fp32).abs().mean(dim=dims, keepdim=True) + if reduction is None: + diff_abs_mean = (gen_data_fp32 - teacher_x0_fp32).abs().mean(dim=dims, keepdim=True) + else: + difference = reduction.clean(gen_data_fp32 - teacher_x0_fp32) + diff_abs_mean = difference.abs().sum(dims) / reduction.count.clamp_min(1) + diff_abs_mean = expand_like(diff_abs_mean, gen_data) w_fp32 = 1 / (diff_abs_mean + 1e-6) # Apply additional scale if provided @@ -97,17 +127,31 @@ def variational_score_distillation_loss( w = w_fp32.to(dtype=original_dtype) vsd_grad = (fake_score_x0 - teacher_x0) * w + if reduction is not None: + vsd_grad = reduction.clean(vsd_grad) pseudo_target = gen_data - vsd_grad - loss = 0.5 * F.mse_loss(gen_data, pseudo_target, reduction="mean") + loss = 0.5 * ( + F.mse_loss(gen_data, pseudo_target, reduction="mean") + if reduction is None + else reduction.residual(gen_data - pseudo_target, mean=True).mean() + ) return loss -def gan_loss_generator(fake_logits: torch.Tensor) -> torch.Tensor: +def gan_loss_generator(fake_logits: torch.Tensor, *, loss_weight: Optional[torch.Tensor] = None) -> torch.Tensor: """ - Compute the GAN loss for the generator + Compute the generator's mean softplus loss over discriminator scores. + + By default, every score contributes equally. Supply ``loss_weight`` to + exclude or scale score losses while retaining the full score count in the + mean. Determine valid scores during feature extraction and pooling. A + score mask cannot undo corruption that has spread through the backbone. + Args: fake_logits (torch.Tensor): The logits for the fake data. + loss_weight (Optional[torch.Tensor]): Finite nonnegative weights that broadcast + to fake_logits. Zero entries exclude scores. Positive values scale their losses. Returns: gan_loss (torch.Tensor): The GAN loss for the generator. @@ -115,22 +159,41 @@ def gan_loss_generator(fake_logits: torch.Tensor) -> torch.Tensor: """ assert fake_logits.ndim == 2, f"fake_logits has shape {fake_logits.shape}" - gan_loss = F.softplus(-fake_logits).mean() + if loss_weight is None: + gan_loss = F.softplus(-fake_logits).mean() + else: + reduction = LossReduction(fake_logits, loss_weight) + gan_loss = (reduction.clean(F.softplus(-reduction.clean(fake_logits))) * reduction.weight).mean() return gan_loss -def gan_loss_discriminator(real_logits: torch.Tensor, fake_logits: torch.Tensor) -> torch.Tensor: +def gan_loss_discriminator( + real_logits: torch.Tensor, + fake_logits: torch.Tensor, + *, + real_weight: Optional[torch.Tensor] = None, + fake_weight: Optional[torch.Tensor] = None, +) -> torch.Tensor: """ - Compute the GAN loss for the discriminator + Sum the discriminator's mean softplus losses for real and fake scores. + + By default, every score contributes equally. Use ``real_weight`` and + ``fake_weight`` to exclude or scale each set of score losses, with the + same reduction and feature-validity requirements as ``gan_loss_generator``. + Args: real_logits (torch.Tensor): The logits for the real data. fake_logits (torch.Tensor): The logits for the fake data. + real_weight, fake_weight (Optional[torch.Tensor]): Finite nonnegative weights + that broadcast to the corresponding logits. Zero entries exclude scores. Returns: gan_loss (torch.Tensor): The GAN loss for the discriminator. """ assert fake_logits.ndim == 2, f"fake_logits has shape {fake_logits.shape}" assert real_logits.ndim == 2, f"real_logits has shape {real_logits.shape}" - gan_loss = F.softplus(fake_logits).mean() + F.softplus(-real_logits).mean() + gan_loss = gan_loss_generator(-fake_logits, loss_weight=fake_weight) + gan_loss_generator( + real_logits, loss_weight=real_weight + ) return gan_loss diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 5b9710d..8159cfd 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -3,6 +3,8 @@ from __future__ import annotations +from fastgen.methods.loss_support import LossReduction + import contextlib from typing import Dict, Any, Callable, TYPE_CHECKING, Tuple, Optional @@ -135,9 +137,9 @@ def _drop_condition( return torch.where(expand_like(keep, condition), condition, neg_condition), keep if isinstance(condition, dict): keys_no_drop = set(self.config.cond_keys_no_dropout) - assert keys_no_drop.issubset( - condition.keys() - ), f"cond_keys_no_dropout: {keys_no_drop} not in {condition.keys()}" + assert keys_no_drop.issubset(condition.keys()), ( + f"cond_keys_no_dropout: {keys_no_drop} not in {condition.keys()}" + ) condition = condition.copy() for k in condition.keys() - keys_no_drop: condition[k] = torch.where(expand_like(keep, condition[k]), condition[k], neg_condition[k]) @@ -168,9 +170,9 @@ def _get_velocity( fuse_scale = self.config.guidance_fuse_scale if fuse_scale is not None: assert fuse_scale > 0, f"guidance_fuse_scale must be > 0, got {fuse_scale} (None disables fusion)" - assert ( - neg_condition is not None - ), "guidance_fuse_scale requires neg_condition: the unconditional branch is queried at the same (t, r)" + assert neg_condition is not None, ( + "guidance_fuse_scale requires neg_condition: the unconditional branch is queried at the same (t, r)" + ) condition, keep = self._drop_condition(condition, neg_condition, x.shape[0], x.device) else: keep = torch.ones(x.shape[0], dtype=torch.bool, device=x.device) @@ -312,9 +314,10 @@ def net_wrapper(x_t, t, r): if self.loss_config.use_jvp_finite_diff: u_theta_jvp = self._estimate_jvp_finite_difference(net_wrapper, x_t, t, r, dxt_dt) else: - with torch.random.fork_rng( - devices=[self.device] if self.device.type == "cuda" else [] - ), temp_disable_efficient_attn(self.device.type): + with ( + torch.random.fork_rng(devices=[self.device] if self.device.type == "cuda" else []), + temp_disable_efficient_attn(self.device.type), + ): tangents = ( dxt_dt.to(dtype=x_t.dtype), torch.ones_like(t, dtype=x_t.dtype), @@ -362,9 +365,9 @@ def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tens flow_matching_ratio = self.flow_map_sample_t_cfg.flow_matching_ratio consistency_ratio = self.flow_map_sample_t_cfg.consistency_ratio - assert ( - flow_matching_ratio + consistency_ratio <= 1.0 - ), f"flow_matching_ratio + consistency_ratio must be <= 1, got {flow_matching_ratio} + {consistency_ratio}" + assert flow_matching_ratio + consistency_ratio <= 1.0, ( + f"flow_matching_ratio + consistency_ratio must be <= 1, got {flow_matching_ratio} + {consistency_ratio}" + ) # Both policies produce two bucket sizes and index the batch the same # way; they differ only in how the sizes are obtained. @@ -456,9 +459,9 @@ def _compute_weight(self, tensor: torch.Tensor, t: torch.Tensor) -> torch.Tensor if self.loss_config.weight_type is not None: weight = weight * (self._timestep_weight_raw(t) * self._timestep_weight_scale) - assert ( - weight.shape == tensor.shape - ), f"weight and input tensor must have the same shape: {weight.shape} != {tensor.shape}" + assert weight.shape == tensor.shape, ( + f"weight and input tensor must have the same shape: {weight.shape} != {tensor.shape}" + ) return weight def _mf_pred_to_loss( @@ -470,7 +473,22 @@ def _mf_pred_to_loss( t: torch.Tensor, r: torch.Tensor, iteration: int, + *, + loss_weight: torch.Tensor | None = None, + clip_time_interval: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, float]: + """Compute MeanFlow losses with adaptive weighting and tangent normalization. + + Omit ``loss_weight`` to include every entry. Supply finite nonnegative + weights that broadcast to ``u_theta`` to scale squared residuals and + exclude zero-weight entries from tangent norms and adaptive weights. + Mean reductions divide by the full element count. Samples with no valid + entries contribute zero loss. + + By default, clip ``t - r`` to the noise scheduler's time bounds. Set + ``clip_time_interval=False`` when ``t - r`` already expresses the + duration required by a custom time parameterization. + """ # perform loss computations in higher precision dtype = torch.float64 u_theta = u_theta.to(dtype=dtype) @@ -479,12 +497,13 @@ def _mf_pred_to_loss( dxt_dt = dxt_dt.to(dtype=dtype) t = t.to(dtype=dtype) r = r.to(dtype=dtype) - delta_t = expand_like((t - r).clip(self.net.noise_scheduler.min_t, self.net.noise_scheduler.max_t), u_theta) + interval = t - r + if clip_time_interval: + interval = interval.clip(self.net.noise_scheduler.min_t, self.net.noise_scheduler.max_t) + delta_t = expand_like(interval, u_theta) + reduction = LossReduction(u_theta, loss_weight) if loss_weight is not None else None assert u_theta.requires_grad is torch.is_grad_enabled(), "u_theta requires gradients if gradients are enabled" - assert not any( - t.requires_grad for t in [u_theta_jvp, x_t, dxt_dt, delta_t] - ), "u_theta_jvp, x_t, v_t, dxt_dt, delta_t should not require gradients" # warmup weight if self.loss_config.tangent_warmup_steps > 0: @@ -495,25 +514,41 @@ def _mf_pred_to_loss( # use l2 loss as in the original paper if self.loss_config.loss_type == "l2": tangent = dxt_dt - warmup_weight * delta_t * u_theta_jvp + if reduction is not None: + tangent = reduction.clean(tangent) loss = (u_theta - tangent).pow(2) - if self.loss_config.norm_method is None: + if reduction is not None: + residual = reduction.clean(u_theta - tangent) + statistic = residual.square().sum(reduction.dims) + if self.loss_config.norm_method is None: + statistic = statistic / reduction.count.clamp_min(1) + weight = self._compute_weight(statistic, t) + loss = reduction.residual(residual, mean=self.loss_config.norm_method is None) + elif self.loss_config.norm_method is None: # Without adaptive normalization the reduction matters: the # per-element mean keeps the loss scale independent of the # sample dimensionality. loss = torch.mean(loss, dim=list(range(1, loss.ndim))) else: loss = torch.sum(loss, dim=list(range(1, loss.ndim))) - weight = self._compute_weight(loss, t) + if reduction is None: + weight = self._compute_weight(loss, t) loss = loss * weight # use explicit gradient elif self.loss_config.loss_type == "opt_grad": u_theta_ = u_theta.detach() tangent = dxt_dt - u_theta_ - warmup_weight * delta_t * u_theta_jvp + if reduction is not None: + tangent = reduction.clean(tangent) if self.loss_config.tangent_spatial_invariance: # scale tangent to be dimension invariant - sample_dim_inv = np.sqrt(tangent.shape[0] / tangent.numel()) + sample_dim_inv = ( + np.sqrt(tangent.shape[0] / tangent.numel()) + if reduction is None + else expand_like(reduction.count.clamp_min(1).rsqrt(), tangent) + ) tangent = tangent * sample_dim_inv opt_grad_norm = torch.linalg.vector_norm(tangent.flatten(1), dim=-1) @@ -521,12 +556,16 @@ def _mf_pred_to_loss( weight = expand_like(weight, tangent) loss = (u_theta - (u_theta + tangent * weight).detach()).pow(2) loss = torch.sum(loss, dim=list(range(1, loss.ndim))) + if reduction is not None: + residual = u_theta - (u_theta + tangent * weight).detach() + loss = reduction.residual(residual, mean=False) else: raise ValueError(f"Invalid loss type: {self.loss_config.loss_type}") - assert not tangent.requires_grad, "Tangent should not require gradients" assert not weight.requires_grad, "Loss weight should not require gradients" + if reduction is not None: + loss = torch.where(reduction.active, loss, torch.zeros_like(loss)) return loss, tangent, weight, warmup_weight @classmethod diff --git a/fastgen/methods/consistency_model/sCM.py b/fastgen/methods/consistency_model/sCM.py index 85d8c82..c51d617 100644 --- a/fastgen/methods/consistency_model/sCM.py +++ b/fastgen/methods/consistency_model/sCM.py @@ -3,6 +3,8 @@ from __future__ import annotations +from fastgen.methods.loss_support import LossReduction + from typing import Dict, Any, Callable, TYPE_CHECKING, Tuple, Optional import numpy as np @@ -248,7 +250,17 @@ def _scm_pred_to_loss( sigma_t: torch.Tensor, t_hat: torch.Tensor, iteration: int, + *, + loss_weight: torch.Tensor | None = None, ) -> Tuple[torch.Tensor, torch.Tensor, float]: + """Compute sCM losses with adaptive weighting and tangent normalization. + + Omit ``loss_weight`` to include every entry. Supply finite nonnegative + weights that broadcast to ``F_theta`` to scale squared residuals and + exclude zero-weight entries from tangent norms and dimension counts. + Means divide by the full element count. Samples with no valid entries + contribute zero loss, including the log-variance term. + """ # perform loss computations in higher precision dtype = torch.float64 F_theta = F_theta.to(dtype) @@ -268,22 +280,35 @@ def _scm_pred_to_loss( g2_coeff = expand_like(alpha_hat * sigma_hat, x_t_hat).to(F_theta.dtype) g2 = -(g2_coeff * x_t_hat + self.sigma_data * F_theta_jvp) g = g1 + r * g2 + reduction = LossReduction(F_theta, loss_weight) if loss_weight is not None else None + if reduction is not None: + g = reduction.clean(g) # Tangent normalization (or we can perform Tangent clipping: g = torch.clamp(g, min=-1, max=1)) g_norm = torch.linalg.vector_norm(g, dim=tuple(range(1, g.ndim)), keepdim=True) if self.loss_config.g_norm_spatial_invariance: - g_norm = g_norm * np.sqrt(g_norm.numel() / g.numel()) # Make the norm to be invariant to spatial size + if reduction is None: + g_norm = g_norm * np.sqrt(g_norm.numel() / g.numel()) + else: + g_norm = g_norm / expand_like(reduction.count.clamp_min(1).sqrt(), g) g = g / (g_norm + self.loss_config.tangent_warmup_const) # Calculate loss with adaptive weighting weight = 1.0 / sigma_t if self.loss_config.prior_weighting_enabled else 1.0 D = x_t_hat[0, :].numel() if self.loss_config.divide_x_0_spatial_dim else 1.0 scm_loss_unweighted = torch.mean(torch.square(F_theta - F_theta_ - g), dim=tuple(range(1, F_theta.ndim))) + if reduction is not None: + scm_loss_unweighted = reduction.residual(F_theta - F_theta_ - g, mean=True) + D = reduction.count.clamp_min(1) if self.loss_config.divide_x_0_spatial_dim else 1.0 - assert ( - logvar.shape == scm_loss_unweighted.shape == t_hat.shape - ), f"log_var.shape: {logvar.shape}, scm_loss_unweighted.shape: {scm_loss_unweighted.shape}, t_hat.shape: {t_hat.shape}" + assert logvar.shape == scm_loss_unweighted.shape == t_hat.shape, ( + f"log_var.shape: {logvar.shape}, scm_loss_unweighted.shape: {scm_loss_unweighted.shape}, t_hat.shape: {t_hat.shape}" + ) + if reduction is not None: + logvar = torch.where(reduction.active, logvar, torch.zeros_like(logvar)) scm_loss = weight / (torch.exp(logvar) * D) * scm_loss_unweighted + logvar + if reduction is not None: + scm_loss = torch.where(reduction.active, scm_loss, torch.zeros_like(scm_loss)) return scm_loss, scm_loss_unweighted, r diff --git a/fastgen/methods/distribution_matching/dmd2.py b/fastgen/methods/distribution_matching/dmd2.py index ed31a8b..77904b7 100644 --- a/fastgen/methods/distribution_matching/dmd2.py +++ b/fastgen/methods/distribution_matching/dmd2.py @@ -64,8 +64,16 @@ def build_model(self): synchronize() torch.cuda.empty_cache() - def _setup_grad_requirements(self, iteration: int) -> None: - if iteration % self.config.student_update_freq == 0: + def _select_training_phase(self, iteration: int, phase: str | None = None) -> str: + """Select an explicit phase or follow ``student_update_freq``.""" + if phase is None: + return "student" if iteration % self.config.student_update_freq == 0 else "auxiliary" + if phase not in ("student", "auxiliary"): + raise ValueError("phase must be student or auxiliary.") + return phase + + def _setup_grad_requirements(self, iteration: int, phase: str | None = None) -> None: + if self._select_training_phase(iteration, phase) == "student": # update the student self.fake_score.eval().requires_grad_(False) if self.config.gan_loss_weight_gen > 0: @@ -76,29 +84,33 @@ def _setup_grad_requirements(self, iteration: int) -> None: if self.config.gan_loss_weight_gen > 0: self.discriminator.train().requires_grad_(True) - def _sample_noising_time(self, batch_size: int, iteration: Optional[int] = None) -> torch.Tensor: - """Draw the noising time `t` for whichever branch this iteration runs. + def _sample_noising_time( + self, batch_size: int, iteration: Optional[int] = None, phase: str | None = None + ) -> torch.Tensor: + """Draw noising times for the selected training phase. - The generator and fake-score updates alternate and need not noise from - the same density. Unset `fake_score_sample_t_cfg` means both use - `sample_t_cfg`. + Use ``phase`` to select student or auxiliary sampling explicitly. + Otherwise, follow ``student_update_freq`` at ``iteration``. Auxiliary + updates use ``fake_score_sample_t_cfg`` when set. Both phases use + ``sample_t_cfg`` otherwise. """ t_cfg = self.config.sample_t_cfg - is_fake_score_step = iteration is not None and iteration % self.config.student_update_freq != 0 + is_fake_score_step = self._select_training_phase(iteration or 0, phase) == "auxiliary" if is_fake_score_step and getattr(self.config, "fake_score_sample_t_cfg", None) is not None: t_cfg = self.config.fake_score_sample_t_cfg return self.net.noise_scheduler.sample_t(batch_size, **convert_cfg_to_dict(t_cfg), device=self.device) def _generate_noise_and_time( - self, real_data: torch.Tensor, iteration: Optional[int] = None + self, real_data: torch.Tensor, iteration: Optional[int] = None, *, phase: str | None = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Generate random noises and time step Args: real_data: Real data tensor for dtype/device reference - iteration: Current training iteration, used only to tell a generator - update from a fake-score one so `fake_score_sample_t_cfg` can - apply on the latter. `None` always uses `sample_t_cfg`. + iteration: Training iteration used to select the sampling phase when + you omit phase. With both omitted, use student sampling. + phase: Explicit student or auxiliary sampling phase. Auxiliary + sampling uses fake_score_sample_t_cfg when set. Returns: rand_z_max: Random noise used by the student @@ -130,7 +142,11 @@ def _generate_noise_and_time( ) input_student = self.net.noise_scheduler.forward_process(real_data, eps_student, t_student) - t = self._sample_noising_time(batch_size, iteration) + t = ( + self._sample_noising_time(batch_size, iteration) + if phase is None + else self._sample_noising_time(batch_size, iteration, phase) + ) eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) return input_student, t_student, t, eps @@ -230,12 +246,12 @@ def _student_update_step( fake_score_x0 = self.fake_score(perturbed_data, t, condition=condition, fwd_pred_type="x0") # Compute the teacher x0-prediction and gan loss for generator - assert ( - perturbed_data.dtype == data["real"].dtype == input_student.dtype - ), f"perturbed_data.dtype: {perturbed_data.dtype}, data['real'].dtype: {data['real'].dtype}, input_student.dtype: {input_student.dtype}" - assert ( - t.dtype == t_student.dtype == self.net.noise_scheduler.t_precision - ), f"t.dtype: {t.dtype}, t_student.dtype: {t_student.dtype}, self.net.noise_scheduler.t_precision: {self.net.noise_scheduler.t_precision}" + assert perturbed_data.dtype == data["real"].dtype == input_student.dtype, ( + f"perturbed_data.dtype: {perturbed_data.dtype}, data['real'].dtype: {data['real'].dtype}, input_student.dtype: {input_student.dtype}" + ) + assert t.dtype == t_student.dtype == self.net.noise_scheduler.t_precision, ( + f"t.dtype: {t.dtype}, t_student.dtype: {t_student.dtype}, self.net.noise_scheduler.t_precision: {self.net.noise_scheduler.t_precision}" + ) teacher_x0, gan_loss_gen = self._compute_teacher_prediction_gan_loss(perturbed_data, t, condition=condition) # Apply classifier-free guidance if needed @@ -358,12 +374,12 @@ def _fake_score_discriminator_update_step( # The fake score matches the teacher, but we want to do SDS in x0 space fake_score_pred_type = self.config.fake_score_pred_type or self.teacher.net_pred_type - assert ( - x_t_sg.dtype == real_data.dtype == input_student.dtype - ), f"x_t_sg.dtype: {x_t_sg.dtype}, real_data.dtype: {real_data.dtype}, input_student.dtype: {input_student.dtype}" - assert ( - t.dtype == t_student.dtype == self.net.noise_scheduler.t_precision - ), f"t.dtype: {t.dtype}, t_student.dtype: {t_student.dtype}, self.net.noise_scheduler.t_precision: {self.net.noise_scheduler.t_precision}" + assert x_t_sg.dtype == real_data.dtype == input_student.dtype, ( + f"x_t_sg.dtype: {x_t_sg.dtype}, real_data.dtype: {real_data.dtype}, input_student.dtype: {input_student.dtype}" + ) + assert t.dtype == t_student.dtype == self.net.noise_scheduler.t_precision, ( + f"t.dtype: {t.dtype}, t_student.dtype: {t_student.dtype}, self.net.noise_scheduler.t_precision: {self.net.noise_scheduler.t_precision}" + ) fake_score_pred = self.fake_score(x_t_sg, t, condition=condition, fwd_pred_type=fake_score_pred_type) loss_fakescore = denoising_score_matching_loss( fake_score_pred_type, @@ -434,38 +450,69 @@ def _get_outputs( return {"gen_rand": gen_rand_func, "input_rand": noise, "gen_rand_train": gen_data} def single_train_step( - self, data: Dict[str, Any], iteration: int + self, data: Dict[str, Any], iteration: int, *, phase: str | None = None ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: """ - Single training step for distribution matching distillation (DMD) + Compute losses and outputs for one distribution matching training phase. + + By default, ``student_update_freq`` selects the phase at ``iteration``. + Set ``phase`` when your training loop schedules student and auxiliary + updates independently. Pass the actual iteration for warmup schedules, + and step the optimizer matching the selected phase. This method + prepares inputs, samples noise, and sets model modes and gradients. + Use ``_dispatch_training_phase`` when your loop manages that setup. Args: data (Dict[str, Any]): Data dict for the current iteration. - iteration (int): Current training iteration + iteration (int): Actual training iteration, including when phase is explicit. + phase (str, optional): Select "student" or "auxiliary", or omit to + follow student_update_freq. Returns: loss_map (dict[str, torch.Tensor]): Dictionary containing the loss values outputs (dict[str, torch.Tensor]): Dictionary containing the network output """ + selected_phase = self._select_training_phase(iteration, phase) # Prepare training data and conditions real_data, condition, neg_condition = self._prepare_training_data(data) # Set up gradient requirements based on training phase - self._setup_grad_requirements(iteration) + if phase is None: + self._setup_grad_requirements(iteration) + else: + self._setup_grad_requirements(iteration, phase=phase) # Generate noise and time steps - input_student, t_student, t, eps = self._generate_noise_and_time(real_data, iteration=iteration) + if phase is None: + input_student, t_student, t, eps = self._generate_noise_and_time(real_data, iteration=iteration) + else: + input_student, t_student, t, eps = self._generate_noise_and_time( + real_data, iteration=iteration, phase=phase + ) - # Choose between student update or fake_score/discriminator update - if iteration % self.config.student_update_freq == 0: + return self._dispatch_training_phase( + selected_phase, input_student, t_student, t, eps, data, real_data, condition, neg_condition + ) + + def _dispatch_training_phase( + self, phase, input_student, t_student, t, eps, data, real_data, condition=None, neg_condition=None + ): + """Compute student or auxiliary losses from prepared inputs. + + Use this hook when your training loop manages sampling, model modes, + frozen parameters, and optimizer selection. For automatic input and + model setup, use ``single_train_step``. + """ + if phase == "student": return self._student_update_step( input_student, t_student, t, eps, data, condition=condition, neg_condition=neg_condition ) - else: + elif phase == "auxiliary": return self._fake_score_discriminator_update_step( input_student, t_student, t, eps, real_data, condition=condition ) + raise ValueError("phase must be student or auxiliary.") def init_optimizers(self): """Initialize optimizers, lr_schedulers and grad_scalers""" diff --git a/fastgen/methods/distribution_matching/f_distill.py b/fastgen/methods/distribution_matching/f_distill.py index e7b7eb4..177c52b 100644 --- a/fastgen/methods/distribution_matching/f_distill.py +++ b/fastgen/methods/distribution_matching/f_distill.py @@ -11,7 +11,7 @@ gan_loss_generator, ) import fastgen.utils.logging_utils as logger -from fastgen.utils.distributed import world_size +from fastgen.methods.loss_support import f_divergence_weights if TYPE_CHECKING: from fastgen.configs.methods.config_f_distill import ModelConfig @@ -40,9 +40,9 @@ def __init__(self, config: ModelConfig): assert self.config.gan_loss_weight_gen > 0, "f-distill requires gan_loss_weight_gen > 0" - assert ( - self.config.f_distill.f_div in all_f_div_weighting_function - ), f"Unsupported f-divergence {self.config.f_distill.f_div}" + assert self.config.f_distill.f_div in all_f_div_weighting_function, ( + f"Unsupported f-divergence {self.config.f_distill.f_div}" + ) logger.info(f"Using {self.config.f_distill.f_div}-divergence") self.f_div_weighting_function = all_f_div_weighting_function[self.config.f_distill.f_div] @@ -56,59 +56,43 @@ def build_model(self): # IF you make them not persistent, you need to ensure they are initialized in net.reset_parameters self.net.register_buffer("bins", torch.ones(self.bin_num)) - def _get_f_div_weighting_h(self, fake_logits: torch.Tensor, t: torch.Tensor) -> torch.Tensor: - assert fake_logits.ndim == 2, f"fake_logits expected to be of rank 2, got {fake_logits.ndim}" - - # Store original dtype and convert to fp32/64 for numerical stability - original_dtype = fake_logits.dtype - fake_logits_fp32 = fake_logits.float() - - # Clamp fake_logits to prevent exponential overflow (in fp32) - fake_logits_clamped = torch.clamp(fake_logits_fp32.mean(dim=1), min=-10.0, max=10.0) - ratio = torch.exp(fake_logits_clamped).detach() - ratio = torch.clamp(ratio, self.config.f_distill.ratio_lower, self.config.f_distill.ratio_upper) - assert ratio.shape == t.shape, f"ratio.shape {ratio.shape} != t.shape {t.shape}" - - if self.config.f_distill.ratio_normalization: - # get bin idx (using fp32 precision) - max_t = float(self.net.noise_scheduler.max_t) - min_t = float(self.net.noise_scheduler.min_t) - bin_width = (max_t - min_t) / self.bin_num - idx = ((t.double() - min_t) / bin_width).floor().long().clamp(0, self.bin_num - 1) - - # bin statistics (keep in fp32) - cnt = torch.bincount(idx, minlength=self.bin_num).float() - ratio_sum = torch.bincount(idx, weights=ratio, minlength=self.bin_num).float() - - if world_size() > 1: - torch.distributed.all_reduce(cnt) - torch.distributed.all_reduce(ratio_sum) - - # EMA update only where cnt > 0 (in fp32 for stability) - valid = cnt > 0 - new_vals = ratio_sum / (cnt + 1e-6) - ema_rate = self.config.f_distill.ratio_ema_rate - - # Convert bins to fp32 for stable computation, then back - bins_fp32 = self.net.bins.float() - bins_fp32[valid] = bins_fp32[valid] * ema_rate + (1 - ema_rate) * new_vals[valid] - self.net.bins.copy_(bins_fp32.to(self.net.bins.dtype)) - - # normalize ratio based on EMA histogram with stability epsilon (in fp32) - ratio = ratio / (bins_fp32[idx] + 1e-6) - - # normalize the weighting function h (in fp32) - h = self.f_div_weighting_function(ratio) - h_sum = h.sum() - if world_size() > 1: - torch.distributed.all_reduce(h_sum) - h_mean = h_sum / (len(h) * world_size()) - - # Add epsilon to prevent division by zero in h normalization (in fp32) - h = h / (h_mean + 1e-6) - - # Convert back to original precision - return h.to(dtype=original_dtype) + def _get_f_div_weighting_h( + self, + fake_logits: torch.Tensor, + t: torch.Tensor, + *, + logit_weight: torch.Tensor | None = None, + sample_mask: torch.Tensor | None = None, + statistics_group=None, + update_state: bool = True, + ) -> torch.Tensor: + """Compute f-distill weights and maintain the student's ratio bins. + + Use ``logit_weight``, ``sample_mask``, and ``statistics_group`` with + the semantics of ``f_divergence_weights``. Set ``update_state=False`` + to use saved bins without updating them. Call ``f_divergence_weights`` + directly when your training loop owns the statistics. + """ + options = self.config.f_distill + bins = self.net.bins if options.ratio_normalization else None + h, updated = f_divergence_weights( + fake_logits, + t, + self.f_div_weighting_function, + ratio_lower=options.ratio_lower, + ratio_upper=options.ratio_upper, + bins=bins, + min_t=float(self.net.noise_scheduler.min_t), + max_t=float(self.net.noise_scheduler.max_t), + ema_rate=options.ratio_ema_rate, + logit_weight=logit_weight, + sample_mask=sample_mask, + statistics_group=statistics_group, + update_state=update_state, + ) + if bins is not None and update_state: + bins.copy_(updated.to(bins.dtype)) + return h def _student_update_step( self, diff --git a/fastgen/methods/loss_support.py b/fastgen/methods/loss_support.py new file mode 100644 index 0000000..9810d4e --- /dev/null +++ b/fastgen/methods/loss_support.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Loss support and ratio statistics for custom training integrations.""" + +from collections.abc import Callable + +import torch + + +class LossReduction: + r"""Reduce squared residuals and compute norms over valid entries. + + Omit ``weight`` to include every entry. Set weights to zero to exclude + invalid entries before nonlinear operations. Positive weights scale + squared residuals linearly. Norms and valid counts use only their + positive support. Means divide by the full element count, including + masked entries. + + Parameters + ---------- + reference : torch.Tensor + Shape reference for broadcasting the weights. + weight : torch.Tensor or None, optional, default=None + Finite nonnegative weights. Omit to give every entry unit weight. + Check these values before calling compiled loss code. + + Examples + -------- + The second feature is an outlier excluded by a zero weight. + The reduction retains the first feature and returns its weighted squared residual. + + >>> # The zero weight excludes 99 from the residual calculation. + >>> x = torch.tensor([[2.0, 99.0]]) + >>> LossReduction(x, torch.tensor([[0.5, 0.0]])).residual(x, mean=True) + tensor([1.]) + """ + + def __init__(self, reference: torch.Tensor, weight: torch.Tensor | 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: torch.Tensor) -> torch.Tensor: + """Remove excluded entries before nonlinear operations.""" + return torch.where(self.valid, value, torch.zeros_like(value)) + + def norm(self, value: torch.Tensor) -> torch.Tensor: + """Compute each sample's Euclidean norm using only unmasked entries.""" + return torch.linalg.vector_norm(self.clean(value).flatten(1), dim=1) + + def residual(self, value: torch.Tensor, *, mean: bool) -> torch.Tensor: + """Reduce linearly weighted squared residuals over non-batch dimensions.""" + weighted = self.clean(value).square() * self.weight + return weighted.mean(self.dims) if mean else weighted.sum(self.dims) + + +@torch.no_grad() +def f_divergence_weights( + logits: torch.Tensor, + t: torch.Tensor, + weighting_fn: Callable[[torch.Tensor], torch.Tensor], + *, + ratio_lower: float = 0.1, + ratio_upper: float = 20.0, + bins: torch.Tensor | None = None, + min_t: float = 0.0, + max_t: float = 1.0, + ema_rate: float = 0.0, + logit_weight: torch.Tensor | None = None, + sample_mask: torch.Tensor | None = None, + statistics_group: torch.distributed.ProcessGroup | None = None, + update_state: bool = True, +) -> tuple[torch.Tensor, torch.Tensor | None]: + r"""Compute normalized f-distill weights and return ratio statistics. + + Use this function when your training loop owns the running statistics. + Pass previous ``bins`` to normalize ratios by diffusion time, then save + the returned bins for the next call or a checkpoint. Input bins remain + unchanged. With ``update_state=False``, normalization uses the saved + values without updating them. With ``bins=None``, it skips bin normalization. + + By default, every score and sample contributes. Zero ``logit_weight`` + entries exclude scores, and positive magnitudes do not change their + relative contribution. False ``sample_mask`` entries exclude whole samples. + Samples with no valid scores receive zero weight. Weight normalization + uses the global valid sample count and supports unequal batch sizes. + + Call on every rank in ``statistics_group``, even ranks with no valid samples. + Choose a group joining data replicas, with spatial features combined + within each sample. During distributed execution, ``None`` + uses the default process group. The function leaves its inputs unchanged. + + Examples + -------- + With zero logits, both samples have unit density ratio before normalization. + Omitting bins returns no running ratio statistics. + + >>> # No bins means the function returns no running ratio state. + >>> weights, bins = f_divergence_weights( + ... torch.zeros(2, 3), torch.tensor([0.2, 0.8]), torch.ones_like, + ... ) + >>> weights.shape + torch.Size([2]) + >>> bins is None + True + """ + if logits.ndim != 2 or t.shape != logits.shape[:1]: + raise ValueError("Expected logits [B, scores] and times [B].") + if bins is not None and (bins.ndim != 1 or bins.numel() == 0 or max_t <= min_t): + raise ValueError("Ratio bins require a nonempty vector and increasing time bounds.") + if logit_weight is None: + mean = logits.float().mean(1) + active = torch.ones_like(t, dtype=torch.bool) + else: + valid = torch.broadcast_to(logit_weight, logits.shape) > 0 + count = valid.sum(1) + mean = torch.where(valid, logits.float(), 0).sum(1) / count.clamp_min(1) + active = count > 0 + if sample_mask is not None: + active = active & sample_mask + ratio = mean.clamp(-10, 10).exp().clamp(ratio_lower, ratio_upper) + updated = bins + if bins is not None: + bin_num = bins.numel() + idx = ((t.double() - min_t) / ((max_t - min_t) / bin_num)).floor().long().clamp(0, bin_num - 1) + count = torch.bincount(idx, weights=active.float(), minlength=bin_num).float() + total = torch.bincount(idx, weights=torch.where(active, ratio, 0), minlength=bin_num).float() + statistics = torch.stack((count, total)) + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(statistics, group=statistics_group) + count, total = statistics.unbind(0) + updated = bins.to(device=t.device, dtype=torch.float32) + if update_state: + new_bins = updated * ema_rate + (1 - ema_rate) * total / (count + 1e-6) + updated = torch.where(count > 0, new_bins, updated) + ratio = ratio / (updated[idx] + 1e-6) + h = weighting_fn(ratio).expand_as(ratio) + totals = torch.stack((torch.where(active, h, 0).sum(), active.float().sum())) + if torch.distributed.is_initialized(): + torch.distributed.all_reduce(totals, group=statistics_group) + h_mean = totals[0] / totals[1].clamp_min(1) + return torch.where(active, h / (h_mean + 1e-6), 0).to(logits.dtype), updated diff --git a/tests/test_dmd2model.py b/tests/test_dmd2model.py index 74e169e..9e1b86f 100644 --- a/tests/test_dmd2model.py +++ b/tests/test_dmd2model.py @@ -97,3 +97,11 @@ def test_optimizers(get_model_data): model.optimizers_zero_grad(3) check_grad_zero(model.discriminator) check_grad_zero(model.fake_score) + + +def test_explicit_phase_ignores_iteration_routing(get_model_data): + model, data = get_model_data + losses, _ = model.single_train_step(data, 0, phase="auxiliary") + assert "fake_score_loss" in losses + losses, _ = model.single_train_step(data, 1, phase="student") + assert "vsd_loss" in losses diff --git a/tests/test_loss_support.py b/tests/test_loss_support.py new file mode 100644 index 0000000..ab538ba --- /dev/null +++ b/tests/test_loss_support.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Loss extension points preserve defaults and exclude invalid entries.""" + +from types import SimpleNamespace +from datetime import timedelta + +import pytest +import torch +import torch.multiprocessing as mp + +from fastgen.methods.common_loss import ( + denoising_score_matching_loss, + gan_loss_generator, + variational_score_distillation_loss, +) +from fastgen.methods.consistency_model.mean_flow import MeanFlowModel +from fastgen.methods.consistency_model.sCM import SCMModel +from fastgen.methods.loss_support import LossReduction, f_divergence_weights +from fastgen.networks.noise_schedule import TrigNoiseSchedule + + +def _empty_rank_worker(rank, rendezvous): + torch.distributed.init_process_group( + "gloo", + init_method=f"file://{rendezvous}", + rank=rank, + world_size=2, + timeout=timedelta(seconds=15), + ) + try: + logits = torch.zeros(rank, 1) + h, bins = f_divergence_weights( + logits, + torch.full((rank,), 0.25), + torch.ones_like, + bins=torch.ones(2), + ) + torch.testing.assert_close(h, torch.ones(rank)) + torch.testing.assert_close(bins, torch.ones(2)) + finally: + torch.distributed.destroy_process_group() + + +def test_empty_rank_ratio_collectives(tmp_path): + mp.spawn(_empty_rank_worker, args=(str(tmp_path / "rendezvous"),), nprocs=2, join=True) + + +def test_masked_residual_excludes_nan(): + values = torch.tensor([[1.0, float("nan")]], requires_grad=True) + reduction = LossReduction(values, torch.tensor([[2.0, 0.0]])) + loss = reduction.residual(values, mean=True).sum() + torch.testing.assert_close(loss, torch.tensor(1.0)) + torch.testing.assert_close(torch.autograd.grad(loss, values)[0], torch.tensor([[2.0, 0.0]])) + + +@pytest.mark.parametrize("kind", ("scm", "l2", "opt_grad")) +@pytest.mark.parametrize("empty", (False, True)) +@pytest.mark.parametrize("input_grad", (False, True)) +def test_unit_support_matches_default(kind, empty, input_grad): + cls = SCMModel if kind == "scm" else MeanFlowModel + model = cls.__new__(cls) + torch.nn.Module.__init__(model) + model.loss_config = SimpleNamespace( + tangent_warmup_steps=10, + tangent_warmup_const=0.1, + g_norm_spatial_invariance=True, + prior_weighting_enabled=True, + divide_x_0_spatial_dim=True, + norm_method="poly_0.5", + norm_const=0.01, + loss_type=kind, + weight_type=None, + tangent_spatial_invariance=True, + ) + model.sigma_data = 0.5 + model.net_trigflow = SimpleNamespace(noise_scheduler=TrigNoiseSchedule()) + model.net = SimpleNamespace(noise_scheduler=SimpleNamespace(min_t=0.0, max_t=1.0)) + prediction = torch.randn(2, 3, 4, requires_grad=True) + tangent, state, velocity = [torch.randn_like(prediction) for _ in range(3)] + state.requires_grad_(input_grad) + t = torch.tensor([0.3, 0.7]) + logvar = torch.full((2, 1), float("nan") if empty else 0.0, requires_grad=True) + if kind == "scm": + args = (prediction, tangent, state, velocity, logvar, t, t, 5) + operation = model._scm_pred_to_loss + else: + args = (prediction, tangent, state, velocity, t, t / 2, 5) + operation = model._mf_pred_to_loss + if empty: + actual = operation(*args, loss_weight=torch.zeros_like(prediction))[0] + torch.testing.assert_close(actual, torch.zeros_like(actual)) + inputs = (prediction, logvar) if kind == "scm" else (prediction,) + for value, grad in zip(inputs, torch.autograd.grad(actual.sum(), inputs)): + torch.testing.assert_close(grad, torch.zeros_like(value)) + return + expected = operation(*args)[0] + actual = operation(*args, loss_weight=torch.ones_like(prediction))[0] + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + torch.autograd.grad(actual.sum(), prediction, retain_graph=True)[0], + torch.autograd.grad(expected.sum(), prediction)[0], + ) + + +def test_vsd_unit_and_masked_support(): + generated = torch.randn(2, 3, 4, requires_grad=True) + teacher, fake = torch.randn_like(generated), torch.randn_like(generated) + expected = variational_score_distillation_loss(generated, teacher, fake) + actual = variational_score_distillation_loss(generated, teacher, fake, loss_weight=torch.ones_like(generated)) + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + torch.autograd.grad(actual, generated, retain_graph=True)[0], torch.autograd.grad(expected, generated)[0] + ) + mask = torch.ones_like(generated) + mask[..., 0] = 0 + teacher[..., 0] = float("nan") + actual = variational_score_distillation_loss(generated, teacher, fake, loss_weight=mask) + assert torch.isfinite(actual) + assert torch.isfinite(torch.autograd.grad(actual, generated)[0]).all() + + +def test_ratio_support_and_frozen_state(): + logits = torch.tensor([[0.0, float("nan")], [float("nan"), float("nan")]]) + weight = torch.tensor([[1.0, 0.0], [0.0, 0.0]]) + bins = torch.ones(4) + h, updated = f_divergence_weights( + logits, + torch.tensor([0.2, 0.7]), + lambda r: torch.ones_like(r), + bins=bins, + min_t=0, + max_t=1, + logit_weight=weight, + update_state=False, + ) + torch.testing.assert_close(updated, bins) + torch.testing.assert_close(h, torch.tensor([1.0, 0.0])) + + +def test_masked_score_and_adversarial_losses(): + prediction = torch.tensor([[1.0, float("nan")]], requires_grad=True) + mask = torch.tensor([[2.0, 0.0]]) + target = torch.zeros_like(prediction) + score = denoising_score_matching_loss("flow", prediction, flow_target=target, loss_weight=mask) + torch.testing.assert_close(score, torch.tensor(1.0)) + gan = gan_loss_generator(prediction, loss_weight=mask) + assert torch.isfinite(gan) + assert torch.isfinite(torch.autograd.grad(score + gan, prediction)[0]).all()