From b2d7cdded11332fb839b4538f27d75bee413ee23 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Sat, 16 May 2026 01:12:13 +0800 Subject: [PATCH 01/19] Add AnyFlow algorithm MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AnyFlow is an any-step video diffusion method that trains a single model u_theta(x_t, t, r) to predict the average velocity from t back to r, so the same checkpoint supports arbitrary inference NFE. Training has two stages, switched via config.loss_config.training_stage: * pretrain — flow-map prediction with a central-difference target target = (eps - x0) - (t - r) * dF/dt with dF/dt estimated by central differences at (t ± delta). Per-batch sampling assigns r=t to a `diffusion_ratio` fraction (pure flow matching) and r=0 to a `consistency_ratio` fraction (consistency to clean data). * onpolicy — distribution-matching distillation with r=0 conditioning on top of the pretrained flow-map weights. Inherits DMD2's alternating fake_score / teacher / discriminator updates. The backbone requirement (a secondary timestep r) is already satisfied by the Wan transformer with r_timestep=True, which MeanFlow also exercises; no Wan-side changes are needed. New files: fastgen/methods/distribution_matching/anyflow.py fastgen/methods/distribution_matching/anyflow_scheduler.py fastgen/configs/methods/config_anyflow.py fastgen/configs/experiments/WanT2V/config_anyflow.py tests/test_anyflowmodel.py Modified: fastgen/methods/__init__.py (+1 import) fastgen/methods/distribution_matching/README.md (+1 algorithm entry) The multi-step rollout-with-gradient training (matching self_forcing.py's rollout_with_gradient) is intentionally left for a follow-up PR — the on-policy stage here uses single-step student generation. Signed-off-by: Enderfga --- .../experiments/WanT2V/config_anyflow.py | 74 +++ fastgen/configs/methods/config_anyflow.py | 92 ++++ fastgen/methods/__init__.py | 1 + .../methods/distribution_matching/README.md | 25 + .../methods/distribution_matching/anyflow.py | 519 ++++++++++++++++++ .../anyflow_scheduler.py | 118 ++++ tests/test_anyflowmodel.py | 203 +++++++ 7 files changed, 1032 insertions(+) create mode 100644 fastgen/configs/experiments/WanT2V/config_anyflow.py create mode 100644 fastgen/configs/methods/config_anyflow.py create mode 100644 fastgen/methods/distribution_matching/anyflow.py create mode 100644 fastgen/methods/distribution_matching/anyflow_scheduler.py create mode 100644 tests/test_anyflowmodel.py diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py new file mode 100644 index 0000000..08badbf --- /dev/null +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reference AnyFlow experiment config on Wan-1.3B T2V. + +Mirrors the AnyFlow paper's pretrain configuration: 1.3B student initialised +from a Wan2.1-T2V checkpoint, flow-matching shift=5, beta08 loss weighting, +6k iterations with batch_size_global=32 and lr=5e-5. + +Switching to the on-policy stage: + + config.model.loss_config.training_stage = "onpolicy" + config.model.pretrained_student_net_path = "" + +and adjust ``student_update_freq`` / ``gan_loss_weight_gen`` to taste. +""" + +import fastgen.configs.methods.config_anyflow as config_anyflow_default +from fastgen.configs.data import VideoLoaderConfig +from fastgen.configs.discriminator import Discriminator_Wan_1_3B_Config +from fastgen.configs.net import Wan_1_3B_Config + + +def create_config(): + config = config_anyflow_default.create_config() + + # Default to the pretrain stage; flip the switch to "onpolicy" once the + # flow-map pretrain checkpoint is available. + config.model.loss_config.training_stage = "pretrain" + config.model.loss_config.jvp_finite_diff_eps = 5e-3 + config.model.loss_config.diffusion_ratio = 0.5 + config.model.loss_config.consistency_ratio = 0.25 + config.model.loss_config.weight_type = "beta08" + config.model.loss_config.shift = 5.0 + + config.model.net = Wan_1_3B_Config + config.model.net.r_timestep = True + + # The on-policy stage uses these too, but they are harmless in pretrain. + config.model.discriminator = Discriminator_Wan_1_3B_Config + config.model.discriminator.disc_type = "multiscale_down_mlp_large" + config.model.discriminator.feature_indices = [15, 22, 29] + config.model.gan_loss_weight_gen = 0.0 # disabled by default in pretrain + config.model.guidance_scale = 5.0 + + config.model.precision = "bfloat16" + # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. + config.model.input_shape = [16, 21, 60, 104] + + config.model.net_optimizer.lr = 5e-5 + config.model.fake_score_optimizer.lr = 5e-5 + config.model.discriminator_optimizer.lr = 5e-5 + + config.model.sample_t_cfg.time_dist_type = "shifted" + config.model.sample_t_cfg.min_t = 0.001 + config.model.sample_t_cfg.max_t = 0.999 + + config.model.student_sample_type = "ode" + # Any-step model — multiple NFEs validated at inference time. + config.model.student_sample_steps = 4 + config.model.sample_t_cfg.t_list = [0.999, 0.937, 0.833, 0.624, 0.0] + + config.dataloader_train = VideoLoaderConfig + config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) + config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 + config.dataloader_train.batch_size = 1 + + config.trainer.max_iter = 6000 + config.trainer.logging_iter = 100 + config.trainer.save_ckpt_iter = 500 + config.trainer.batch_size_global = 32 + + config.log_config.group = "wan_anyflow" + return config diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py new file mode 100644 index 0000000..96700a8 --- /dev/null +++ b/fastgen/configs/methods/config_anyflow.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Config schema for the AnyFlow method. + +AnyFlow inherits the DMD2 model config (so the on-policy stage gets fake_score +/ discriminator / alternating-step machinery for free) and adds a +``LossConfig`` describing the flow-map pretrain hyperparameters. +""" + +import attrs +from omegaconf import DictConfig + +from fastgen.configs.callbacks import ( + EMA_CALLBACK, + GPUStats_CALLBACK, + GradClip_CALLBACK, + ParamCount_CALLBACK, + TrainProfiler_CALLBACK, + WANDB_CALLBACK, +) +from fastgen.configs.config import BaseConfig +from fastgen.configs.methods.config_dmd2 import ModelConfig as DMD2ModelConfig +from fastgen.methods import AnyFlowModel +from fastgen.utils import LazyCall as L + + +@attrs.define(slots=False) +class LossConfig: + """Hyperparameters for the AnyFlow flow-map loss and on-policy switch.""" + + # Which stage to train. "pretrain" runs the central-difference flow-map + # objective; "onpolicy" inherits DMD2's alternating distillation with + # dual-timestep r=0 conditioning. + training_stage: str = "pretrain" + + # Central-difference step size for estimating dF/dt. Lives in the same + # units as the noise scheduler's timesteps. The default (5e-3) matches the + # AnyFlow paper's choice of epsilon=5 with num_train_timesteps=1000. + jvp_finite_diff_eps: float = 5e-3 + + # Per-batch fraction with r = t (recovers pure flow matching). + diffusion_ratio: float = 0.5 + # Per-batch fraction with r = min_t (forces consistency to clean data). + consistency_ratio: float = 0.25 + + # Per-timestep loss weighting scheme — passed through to the flow-map + # scheduler. One of "gaussian", "beta08", "uniform". + weight_type: str = "beta08" + # Flow-matching schedule shift for the weighting / sampling scheduler. + # Wan video defaults use 5.0; image use 1.0. + shift: float = 1.0 + # Resolution of the discrete weighting grid; matches the AnyFlow reference. + num_train_timesteps: int = 1000 + + +@attrs.define(slots=False) +class ModelConfig(DMD2ModelConfig): + """AnyFlow model config — inherits DMD2 fields, adds the flow-map loss config.""" + + loss_config: LossConfig = attrs.field(factory=LossConfig) + + +@attrs.define(slots=False) +class Config(BaseConfig): + model: ModelConfig = attrs.field(factory=ModelConfig) + model_class: DictConfig = L(AnyFlowModel)( + config=None, + ) + + +def create_config(): + config = Config() + config.trainer.callbacks = DictConfig( + { + **GradClip_CALLBACK, + **EMA_CALLBACK, + **GPUStats_CALLBACK, + **TrainProfiler_CALLBACK, + **ParamCount_CALLBACK, + **WANDB_CALLBACK, + } + ) + + # Pretrain stage relies on a flow-matching net_pred_type and dual-timestep input. + config.model.use_ema = True + config.model.net.r_timestep = True + config.model.net_scheduler.warm_up_steps = [0] + config.model.fake_score_scheduler.warm_up_steps = [0] + config.model.discriminator_scheduler.warm_up_steps = [0] + + return config diff --git a/fastgen/methods/__init__.py b/fastgen/methods/__init__.py index e902b82..907b3f3 100644 --- a/fastgen/methods/__init__.py +++ b/fastgen/methods/__init__.py @@ -9,6 +9,7 @@ from fastgen.methods.distribution_matching.causvid import CausVidModel as CausVidModel from fastgen.methods.distribution_matching.self_forcing import SelfForcingModel as SelfForcingModel +from fastgen.methods.distribution_matching.anyflow import AnyFlowModel as AnyFlowModel from fastgen.methods.consistency_model.CM import CMModel as CMModel from fastgen.methods.consistency_model.TCM import TCMModel as TCMModel diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index 021e181..6e9c278 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -85,6 +85,31 @@ DMD2 extended for causal video generation with autoregressive chunk-by-chunk pro --- +## AnyFlow + +**File:** [`anyflow.py`](anyflow.py) | **Reference:** AnyFlow — any-step video diffusion framework on flow maps + +Single model that supports arbitrary inference NFE by learning a flow map `u_θ(x_t, t, r)` (average velocity from `t` back to `r`). Trained in two stages: + +1. **Pretrain** — flow-map prediction with a central-difference target that reuses the network's own forward at `(t ± δ, r)` to estimate `dF/dt`. Per-batch sampling assigns `r = t` to a `diffusion_ratio` fraction (recovering plain flow matching) and `r = 0` to a `consistency_ratio` fraction (forcing consistency to clean data). +2. **On-policy** — distribution-matching distillation with `r = 0` conditioning on top of the pretrained flow-map weights. Inherits DMD2's alternating fake_score / discriminator / VSD machinery. + +**Key Parameters:** +- `loss_config.training_stage`: `"pretrain"` or `"onpolicy"` +- `loss_config.jvp_finite_diff_eps`: central-difference step δ (in noise scheduler t-units) +- `loss_config.diffusion_ratio` / `loss_config.consistency_ratio`: per-batch fraction with `r=t` / `r=0` +- `loss_config.weight_type`: `gaussian` | `beta08` | `uniform` per-timestep loss weight +- `loss_config.shift`: flow-matching schedule shift (5.0 for Wan video) +- See also key parameters of DMD2 above (used by the on-policy stage) + +**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). + +**Note:** the on-policy stage in this PR uses single-step student generation. Multi-step rollout-with-gradient (matching `self_forcing.py`'s `rollout_with_gradient`) is intentionally deferred to a follow-up PR. + +**Configs:** [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) + +--- + ## Self-Forcing **File:** [`self_forcing.py`](self_forcing.py) | **Reference:** [Huang et al., 2025](https://arxiv.org/abs/2506.08009) diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py new file mode 100644 index 0000000..cbae8fe --- /dev/null +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -0,0 +1,519 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AnyFlow — any-step video diffusion with flow maps and on-policy distillation. + +AnyFlow trains a single model :math:`u_\\theta(x_t, t, r)` that predicts the +average velocity from time ``t`` to ``r`` (with ``r \\le t``). Once trained, +the same model supports arbitrary inference step counts: each Euler-like +sampling step picks its own integration interval ``(t \\rightarrow r)``. + +Training has two stages, selected via ``config.loss_config.training_stage``: + +* ``"pretrain"`` — flow-map prediction with a central-difference target + + v(x_t, t) = eps - x_0 # instantaneous flow (flow matching) + dF/dt ~= (u_theta(x_{t+d}, t+d, r) - u_theta(x_{t-d}, t-d, r)) / 2d + target = v - (t - r) * dF/dt + loss = weight(t) * MSE(u_theta(x_t, t, r), target) + + Per-batch sampling assigns ``r = t`` for a ``diffusion_ratio`` fraction + (recovering plain flow matching), ``r = 0`` for a ``consistency_ratio`` + fraction (forcing consistency to clean data), and a uniform random pair + otherwise — matching the AnyFlow paper. + +* ``"onpolicy"`` — distribution-matching distillation on top of pretrained + flow-map weights. Inherits DMD2's fake_score / teacher / discriminator + machinery and alternating-step optimisation, but conditions all forwards + on ``r = 0`` (predicting the full flow from ``t`` to clean). Multi-step + rollout-with-gradient is intentionally deferred to a follow-up PR. + +The network must support a secondary timestep argument ``r`` (Wan with +``r_timestep=True`` does; MeanFlow already exercises this same code path). +""" + +from __future__ import annotations + +from functools import partial +from typing import Any, Callable, Dict, Optional, TYPE_CHECKING + +import torch +import torch.nn.functional as F + +from fastgen.methods.distribution_matching.anyflow_scheduler import FlowMapDiscreteScheduler +from fastgen.methods.distribution_matching.dmd2 import DMD2Model +from fastgen.utils import expand_like +import fastgen.utils.logging_utils as logger + + +if TYPE_CHECKING: + from fastgen.configs.methods.config_anyflow import ModelConfig + + +class AnyFlowModel(DMD2Model): + """AnyFlow training method. + + See module docstring for the algorithm. + """ + + def __init__(self, config: ModelConfig): + super().__init__(config) + self.config = config + self.loss_config = self.config.loss_config + + if self.loss_config.training_stage not in ("pretrain", "onpolicy"): + raise ValueError( + f"training_stage must be 'pretrain' or 'onpolicy', got {self.loss_config.training_stage!r}" + ) + + # Standalone scheduler used for inference and for the per-timestep + # training weight in the pretrain stage. Training noising still goes + # through ``self.net.noise_scheduler`` to stay compatible with DMD2. + self._flowmap_scheduler = FlowMapDiscreteScheduler( + num_train_timesteps=self.loss_config.num_train_timesteps, + shift=self.loss_config.shift, + weight_type=self.loss_config.weight_type, + ) + + if self.loss_config.training_stage == "pretrain": + logger.info( + f"AnyFlow pretrain stage: epsilon={self.loss_config.jvp_finite_diff_eps}, " + f"diffusion_ratio={self.loss_config.diffusion_ratio}, " + f"consistency_ratio={self.loss_config.consistency_ratio}, " + f"weight_type={self.loss_config.weight_type}" + ) + else: + logger.info( + f"AnyFlow on-policy stage: student_update_freq={self.config.student_update_freq}, " + f"gan_loss_weight_gen={self.config.gan_loss_weight_gen}" + ) + + # ------------------------------------------------------------------ + # Build / optimisation overrides — skip DMD2 plumbing in pretrain + # ------------------------------------------------------------------ + + def build_model(self): + """In pretrain mode skip fake_score / discriminator entirely.""" + if self.config.loss_config.training_stage == "pretrain": + # Bypass DMD2Model.build_model — pretrain only needs the student. + # Call grandparent's build_model (FastGenModel) directly. + super(DMD2Model, self).build_model() + self.load_student_weights_and_ema() + return + super().build_model() + + def init_optimizers(self): + """Pretrain skips the DMD2 fake_score / discriminator optimisers.""" + if self.config.loss_config.training_stage == "pretrain": + # Bypass DMD2Model.init_optimizers — only the student optimiser exists. + super(DMD2Model, self).init_optimizers() + return + super().init_optimizers() + + @property + def model_dict(self): + if self.config.loss_config.training_stage == "pretrain": + return super(DMD2Model, self).model_dict + return super().model_dict + + @property + def optimizer_dict(self): + if self.config.loss_config.training_stage == "pretrain": + return super(DMD2Model, self).optimizer_dict + return super().optimizer_dict + + @property + def scheduler_dict(self): + if self.config.loss_config.training_stage == "pretrain": + return super(DMD2Model, self).scheduler_dict + return super().scheduler_dict + + # ------------------------------------------------------------------ + # Pretrain stage + # ------------------------------------------------------------------ + + def _sample_pair_timesteps( + self, batch_size: int, dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Sample ``(t, r)`` with ``t >= r``, plus a per-sample diffusion mask. + + Implements AnyFlow's partitioning of the batch: + * a ``diffusion_ratio`` fraction has ``r = t`` (pure flow matching) + * a ``consistency_ratio`` fraction has ``r = min_t`` + * the rest gets a uniform random pair + """ + ns = self.net.noise_scheduler + t_dtype = ns.t_precision + + t_1 = torch.rand(batch_size, device=self.device, dtype=t_dtype) + t_2 = torch.rand(batch_size, device=self.device, dtype=t_dtype) + t_norm = torch.maximum(t_1, t_2) + r_norm = torch.minimum(t_1, t_2) + + # Shift to match the flow-matching schedule (Wan default uses shift=5). + t_norm = self._flowmap_scheduler.apply_shift(t_norm) + r_norm = self._flowmap_scheduler.apply_shift(r_norm) + + # Rescale unit-interval timesteps into the noise scheduler's [min_t, max_t]. + max_t = float(ns.max_t) + min_t = float(ns.min_t) + scale = max_t - min_t + t = t_norm * scale + min_t + r = r_norm * scale + min_t + + # Per-batch bucket assignment. We shuffle so the buckets are randomly + # distributed within the local batch — this matches the paper's intent + # without requiring global cross-rank coordination. + n_diffusion = int(round(self.loss_config.diffusion_ratio * batch_size)) + n_consistency = int(round(self.loss_config.consistency_ratio * batch_size)) + n_diffusion = min(n_diffusion, batch_size) + n_consistency = min(n_consistency, batch_size - n_diffusion) + + perm = torch.randperm(batch_size, device=self.device) + is_diffusion = torch.zeros(batch_size, dtype=torch.bool, device=self.device) + is_consistency = torch.zeros(batch_size, dtype=torch.bool, device=self.device) + is_diffusion[perm[:n_diffusion]] = True + is_consistency[perm[n_diffusion : n_diffusion + n_consistency]] = True + + r = torch.where(is_diffusion, t, r) + r = torch.where(is_consistency, torch.full_like(r, min_t), r) + + return t.to(dtype=t_dtype), r.to(dtype=t_dtype), is_diffusion + + @torch.no_grad() + def _compute_central_difference_target( + self, + x_t: torch.Tensor, + t: torch.Tensor, + r: torch.Tensor, + v: torch.Tensor, + condition: Optional[Any], + ) -> torch.Tensor: + """Compute the AnyFlow flow-map target ``v - (t - r) * dF/dt``. + + ``dF/dt`` is estimated by central difference at ``(t ± delta, r)``. + Boundary cases near ``min_t`` / ``max_t`` fall back to one-sided + differences so the estimate stays valid for the whole timestep range. + """ + ns = self.net.noise_scheduler + max_t = float(ns.max_t) + min_t = float(ns.min_t) + delta = float(self.loss_config.jvp_finite_diff_eps) + + # Validity masks for each finite-difference direction. + is_fwd_valid = (t + delta) <= max_t + is_bwd_valid = ((t - delta) >= min_t) & ((t - delta) > r) + + use_central = is_fwd_valid & is_bwd_valid + use_fwd_only = is_fwd_valid & ~is_bwd_valid + use_bwd_only = ~is_fwd_valid & is_bwd_valid + + # Build per-sample (t_plus, t_minus, denom) with broadcasting-safe shapes. + t_plus = torch.where(is_fwd_valid, t + delta, t) + t_minus = torch.where(is_bwd_valid, t - delta, t) + denom = torch.where( + use_central, + torch.full_like(t, 2 * delta), + torch.where(use_fwd_only | use_bwd_only, torch.full_like(t, delta), torch.full_like(t, 1.0)), + ) + + # Linear-path extrapolation along the flow direction: + # x_t = t * eps + (1 - t) * x_0 => d x_t / d t = eps - x_0 = v + x_t_plus = x_t + expand_like(t_plus - t, x_t) * v + x_t_minus = x_t + expand_like(t_minus - t, x_t) * v + + F_plus = self.net(x_t_plus, t_plus, r=r, condition=condition, fwd_pred_type="flow") + F_minus = self.net(x_t_minus, t_minus, r=r, condition=condition, fwd_pred_type="flow") + + dF_dt = (F_plus - F_minus) / expand_like(denom, x_t) + + # Where no finite-difference direction was valid (extremely rare with a + # small delta), fall back to dF/dt = 0 so we recover pure flow matching. + no_diff_valid = ~(use_central | use_fwd_only | use_bwd_only) + if no_diff_valid.any(): + dF_dt = torch.where(expand_like(no_diff_valid, dF_dt), torch.zeros_like(dF_dt), dF_dt) + + target = v - expand_like(t - r, x_t) * dF_dt + return target + + def _pretrain_single_train_step( + self, data: Dict[str, Any], iteration: int + ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: + """Single training step for AnyFlow pretrain.""" + real_data, condition, _ = self._prepare_training_data(data) + batch_size = real_data.shape[0] + + t, r, is_diffusion = self._sample_pair_timesteps(batch_size, dtype=real_data.dtype) + + # Forward noising along the linear flow-matching path. + eps = torch.randn_like(real_data) + x_t = self.net.noise_scheduler.forward_process(real_data, eps, t) + + # Ground-truth instantaneous flow direction at time t. + v = eps - real_data + + # Central-difference target (no grad through the finite-difference probes). + target = self._compute_central_difference_target(x_t, t, r, v, condition) + + # Student forward — this is the term whose gradient flows. + u_theta = self.net(x_t, t, r=r, condition=condition, fwd_pred_type="flow") + + # Per-sample MSE in float for numerical stability under bf16/fp16 AMP. + sq_err = (u_theta.float() - target.float()).pow(2) + loss_per_sample = sq_err.flatten(1).mean(dim=-1) + + # Per-timestep weight from the flow-map scheduler (beta08 default). + weight = self._flowmap_scheduler.get_train_weight(t).to(loss_per_sample.device, loss_per_sample.dtype) + loss = (loss_per_sample * weight).mean() + + # x0 approximation (monitoring only). + with torch.no_grad(): + x0_approx = self.net.noise_scheduler.flow_to_x0(x_t, u_theta.detach(), t) + + loss_map = { + "total_loss": loss, + "anyflow_loss": loss, + "flow_matching_loss": loss_per_sample[is_diffusion].mean() if is_diffusion.any() else loss.detach() * 0, + "dF_dt_target_norm": (target - v).flatten(1).norm(dim=-1).mean(), + } + outputs = self._get_outputs(x0_approx, input_student=x_t, condition=condition) + return loss_map, outputs + + # ------------------------------------------------------------------ + # On-policy stage — DMD2 with r=0 conditioning + # ------------------------------------------------------------------ + + def _zeros_like_t(self, t: torch.Tensor) -> torch.Tensor: + ns = self.net.noise_scheduler + return torch.full_like(t, float(ns.min_t)) + + def _onpolicy_student_update_step( + self, + input_student: torch.Tensor, + t_student: torch.Tensor, + t: torch.Tensor, + eps: torch.Tensor, + data: Dict[str, Any], + condition: Optional[Any], + neg_condition: Optional[Any], + ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + r_zero = self._zeros_like_t(t_student) + r_zero_t = self._zeros_like_t(t) + + # Student rollout to a single x0 estimate, conditioned on r=0. + gen_data = self.net(input_student, t_student, r=r_zero, condition=condition, fwd_pred_type="x0") + perturbed_data = self.net.noise_scheduler.forward_process(gen_data, eps, t) + + with torch.no_grad(): + fake_score_x0 = self.fake_score(perturbed_data, t, r=r_zero_t, condition=condition, fwd_pred_type="x0") + + # Teacher prediction + optional GAN loss for the generator. Mirrors + # DMD2._compute_teacher_prediction_gan_loss but with r=0 conditioning. + if self.config.gan_loss_weight_gen > 0: + teacher_x0, fake_feat = self.teacher( + perturbed_data, + t, + r=r_zero_t, + condition=condition, + feature_indices=self.discriminator.feature_indices, + fwd_pred_type="x0", + ) + from fastgen.methods.common_loss import gan_loss_generator + + gan_loss_gen = gan_loss_generator(self.discriminator(fake_feat)) + else: + teacher_x0 = self.teacher(perturbed_data, t, r=r_zero_t, condition=condition, fwd_pred_type="x0") + gan_loss_gen = torch.tensor(0.0, device=self.device, dtype=teacher_x0.dtype) + teacher_x0 = teacher_x0.detach() + + # Optional CFG on the teacher. + if self.config.guidance_scale is not None: + with torch.no_grad(): + teacher_x0_neg = self.teacher( + perturbed_data, t, r=r_zero_t, condition=neg_condition, fwd_pred_type="x0" + ) + teacher_x0 = teacher_x0 + (self.config.guidance_scale - 1) * (teacher_x0 - teacher_x0_neg) + + from fastgen.methods.common_loss import variational_score_distillation_loss + + vsd_loss = variational_score_distillation_loss(gen_data, teacher_x0, fake_score_x0) + loss = vsd_loss + self.config.gan_loss_weight_gen * gan_loss_gen + + loss_map = { + "total_loss": loss, + "vsd_loss": vsd_loss, + "gan_loss_gen": gan_loss_gen, + } + outputs = self._get_outputs(gen_data, input_student, condition=condition) + return loss_map, outputs + + def _onpolicy_fake_score_discriminator_update_step( + self, + input_student: torch.Tensor, + t_student: torch.Tensor, + t: torch.Tensor, + eps: torch.Tensor, + real_data: torch.Tensor, + condition: Optional[Any], + ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: + r_zero = self._zeros_like_t(t_student) + r_zero_t = self._zeros_like_t(t) + + with torch.no_grad(): + gen_data = self.net(input_student, t_student, r=r_zero, condition=condition, fwd_pred_type="x0") + x_t_sg = self.net.noise_scheduler.forward_process(gen_data, eps, t) + + from fastgen.methods.common_loss import ( + denoising_score_matching_loss, + gan_loss_discriminator, + ) + + fake_score_pred_type = self.config.fake_score_pred_type or self.teacher.net_pred_type + fake_score_pred = self.fake_score( + x_t_sg, t, r=r_zero_t, condition=condition, fwd_pred_type=fake_score_pred_type + ) + loss_fakescore = denoising_score_matching_loss( + fake_score_pred_type, + net_pred=fake_score_pred, + noise_scheduler=self.net.noise_scheduler, + x0=gen_data, + eps=eps, + t=t, + ) + + gan_loss_disc = torch.zeros_like(loss_fakescore) + gan_loss_ar1 = torch.zeros_like(loss_fakescore) + if self.config.gan_loss_weight_gen > 0: + with torch.no_grad(): + fake_feat = self.teacher( + x_t_sg, + t, + r=r_zero_t, + condition=condition, + return_features_early=True, + feature_indices=self.discriminator.feature_indices, + ) + # Real data path — mirror DMD2._compute_real_feat but pass r=0. + from fastgen.utils.basic_utils import convert_cfg_to_dict + + if self.config.gan_use_same_t_noise: + t_real, eps_real = t, eps + else: + t_real = self.net.noise_scheduler.sample_t( + real_data.shape[0], + **convert_cfg_to_dict(self.config.sample_t_cfg), + device=self.device, + ) + eps_real = torch.randn_like(real_data) + perturbed_real = self.net.noise_scheduler.forward_process(real_data, eps_real, t_real) + r_zero_real = self._zeros_like_t(t_real) + real_feat = self.teacher( + perturbed_real, + t_real, + r=r_zero_real, + condition=condition, + return_features_early=True, + feature_indices=self.discriminator.feature_indices, + ) + + real_feat_logit = self.discriminator(real_feat) + gan_loss_disc = gan_loss_discriminator(real_feat_logit, self.discriminator(fake_feat)) + + if self.config.gan_r1_reg_weight > 0: + perturbed_real_alpha = real_data.add(self.config.gan_r1_reg_alpha * torch.randn_like(real_data)) + with torch.no_grad(): + real_feat_alpha = self.teacher( + perturbed_real_alpha, + t_real, + r=r_zero_real, + condition=condition, + return_features_early=True, + feature_indices=self.discriminator.feature_indices, + ) + real_feat_alpha_logit = self.discriminator(real_feat_alpha) + gan_loss_ar1 = F.mse_loss(real_feat_logit, real_feat_alpha_logit, reduction="mean") + + loss = loss_fakescore + gan_loss_disc + self.config.gan_r1_reg_weight * gan_loss_ar1 + loss_map = { + "total_loss": loss, + "fake_score_loss": loss_fakescore, + "gan_loss_disc": gan_loss_disc, + } + if self.config.gan_loss_weight_gen > 0 and self.config.gan_r1_reg_weight > 0: + loss_map["gan_loss_ar1"] = gan_loss_ar1 + outputs = self._get_outputs(gen_data, input_student, condition=condition) + return loss_map, outputs + + def _onpolicy_single_train_step( + self, data: Dict[str, Any], iteration: int + ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: + real_data, condition, neg_condition = self._prepare_training_data(data) + self._setup_grad_requirements(iteration) + input_student, t_student, t, eps = self._generate_noise_and_time(real_data) + + if iteration % self.config.student_update_freq == 0: + return self._onpolicy_student_update_step( + input_student, t_student, t, eps, data, condition=condition, neg_condition=neg_condition + ) + return self._onpolicy_fake_score_discriminator_update_step( + input_student, t_student, t, eps, real_data, condition=condition + ) + + # ------------------------------------------------------------------ + # FastGenModel interface + # ------------------------------------------------------------------ + + def _get_outputs( + self, + gen_data: torch.Tensor, + input_student: Optional[torch.Tensor] = None, + condition: Any = None, + ) -> Dict[str, torch.Tensor | Callable]: + # Pretrain stage uses a direct x0 approximation tensor. + if self.loss_config.training_stage == "pretrain": + assert input_student is not None, "input_student must be provided" + ns = self.net.noise_scheduler + noise = input_student / (ns.max_sigma if hasattr(ns, "max_sigma") else 1.0) + return {"gen_rand": gen_data, "input_rand": noise} + + # On-policy stage delegates to DMD2's get_outputs path so multi-step + # generators are produced consistently with the rest of the family. + if self.config.student_sample_steps == 1: + assert input_student is not None, "input_student must be provided" + ns = self.net.noise_scheduler + noise = input_student / (ns.max_sigma if hasattr(ns, "max_sigma") else 1.0) + return {"gen_rand": gen_data, "input_rand": noise} + + noise = torch.randn_like(gen_data, dtype=self.precision) + gen_rand_func = partial( + self.generator_fn, + net=self.net_inference, + noise=noise, + condition=condition, + student_sample_steps=self.config.student_sample_steps, + student_sample_type=self.config.student_sample_type, + t_list=self.config.sample_t_cfg.t_list, + precision_amp=self.precision_amp_infer, + ) + 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 + ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: + if self.loss_config.training_stage == "pretrain": + return self._pretrain_single_train_step(data, iteration) + return self._onpolicy_single_train_step(data, iteration) + + def get_optimizers(self, iteration: int) -> list: + """Pretrain stage uses only the student optimizer. + + On-policy stage inherits DMD2's alternating optimisation. + """ + if self.loss_config.training_stage == "pretrain": + return [self.net_optimizer] + return super().get_optimizers(iteration) + + def get_lr_schedulers(self, iteration: int) -> list: + if self.loss_config.training_stage == "pretrain": + return [self.net_lr_scheduler] + return super().get_lr_schedulers(iteration) diff --git a/fastgen/methods/distribution_matching/anyflow_scheduler.py b/fastgen/methods/distribution_matching/anyflow_scheduler.py new file mode 100644 index 0000000..b17ffb2 --- /dev/null +++ b/fastgen/methods/distribution_matching/anyflow_scheduler.py @@ -0,0 +1,118 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Lightweight flow-map scheduler for any-step inference. + +Ported from the AnyFlow reference implementation +(``far/schedulers/scheduling_flowmap_euler_discrete.py``) with the +``diffusers.ConfigMixin`` dependency removed so it can be used standalone. + +The scheduler operates on timesteps in ``[0, num_train_timesteps]`` (matching +the Wan T2V/I2V conventions). For pair-step sampling, ``step`` takes both the +current timestep ``t`` and the target ``r`` (with ``r < t``) and integrates the +flow map prediction in one shot: + + x_r = x_t - (t - r) * u_theta(x_t, t, r) + +This is the flow-map analogue of an Euler step where the integration interval +``t - r`` is chosen freely at inference time, enabling any-step sampling. +""" + +from __future__ import annotations + +from typing import Union + +import torch + + +class FlowMapDiscreteScheduler: + """Any-step flow-map scheduler. + + Args: + num_train_timesteps: Maximum timestep value used during training. + shift: Flow-matching schedule shift (Wan default: 5.0 for video, 1.0 for image). + weight_type: Per-timestep loss weighting scheme — ``gaussian``, ``beta08``, + or ``uniform``. ``beta08`` matches AnyFlow's default. + """ + + def __init__( + self, + num_train_timesteps: int = 1000, + shift: float = 1.0, + weight_type: str = "beta08", + ): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.weight_type = weight_type + + # Initialise with train-time uniform spacing; overridden by set_timesteps() + self.set_timesteps(num_train_timesteps, device="cpu") + self._build_train_weights() + + def _build_train_weights(self) -> None: + if self.weight_type == "gaussian": + x = self.timesteps + y = torch.exp(-2 * ((x - self.num_train_timesteps / 2) / self.num_train_timesteps) ** 2) + y_shifted = y - y.min() + self.linear_timesteps_weights = y_shifted * (self.num_train_timesteps / y_shifted.sum()) + elif self.weight_type == "beta08": + t = self.timesteps / self.num_train_timesteps + y = (t**1.0) * ((1 - t) ** 0.5) + self.linear_timesteps_weights = y * (self.num_train_timesteps / y.sum()) + elif self.weight_type == "uniform": + self.linear_timesteps_weights = torch.ones_like(self.timesteps) + else: + raise ValueError(f"Invalid weight_type: {self.weight_type!r}") + + @torch.no_grad() + def get_train_weight(self, timesteps: torch.Tensor) -> torch.Tensor: + """Look up the per-timestep training loss weight via nearest neighbour.""" + device_weights = self.linear_timesteps_weights.to(timesteps.device) + device_timesteps = self.timesteps.to(timesteps.device) + diffs = (device_timesteps.unsqueeze(1) - timesteps.flatten().unsqueeze(0)).abs() + timestep_id = torch.argmin(diffs, dim=0).reshape(timesteps.shape) + return device_weights[timestep_id] + + def apply_shift(self, sigmas: torch.Tensor) -> torch.Tensor: + """Apply the flow-matching schedule shift to normalized sigmas in [0, 1].""" + if self.shift == 1.0: + return sigmas + return self.shift * sigmas / (1 + (self.shift - 1) * sigmas) + + def set_timesteps( + self, + num_inference_steps: int, + device: Union[str, torch.device, None] = None, + ) -> None: + timesteps = torch.linspace(1.0, 0.0, num_inference_steps + 1, dtype=torch.float64, device=device) + timesteps = self.apply_shift(timesteps) + self.timesteps = timesteps * self.num_train_timesteps + + def scale_noise( + self, + sample: torch.Tensor, + timestep: Union[float, torch.Tensor], + noise: torch.Tensor, + ) -> torch.Tensor: + """Forward-noise ``sample`` to ``timestep`` along a linear flow-matching path.""" + timestep = torch.as_tensor(timestep, device=sample.device, dtype=sample.dtype) + timestep = timestep / self.num_train_timesteps + timestep = timestep.view(*timestep.shape, *([1] * (noise.ndim - timestep.ndim))) + return timestep * noise + (1.0 - timestep) * sample + + def step( + self, + model_output: torch.Tensor, + sample: torch.Tensor, + timestep: Union[float, torch.Tensor], + r_timestep: Union[float, torch.Tensor], + ) -> torch.Tensor: + """Pair-step Euler integration ``x_r = x_t - (t - r) * model_output``.""" + timestep = torch.as_tensor(timestep, device=sample.device, dtype=sample.dtype) + r_timestep = torch.as_tensor(r_timestep, device=sample.device, dtype=sample.dtype) + timestep = timestep / self.num_train_timesteps + r_timestep = r_timestep / self.num_train_timesteps + timestep = timestep.view(*timestep.shape, *([1] * (model_output.ndim - timestep.ndim))) + r_timestep = r_timestep.view(*r_timestep.shape, *([1] * (model_output.ndim - r_timestep.ndim))) + prev_sample = sample - (timestep - r_timestep) * model_output + return prev_sample.to(model_output.dtype) diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py new file mode 100644 index 0000000..ac732e8 --- /dev/null +++ b/tests/test_anyflowmodel.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for AnyFlowModel. + +The tests run on the tiny EDM backbone with ``r_timestep=True`` (same trick +``test_meanflowmodel.py`` uses) so they execute on CPU without downloading +any pretrained weights. +""" + +import gc + +import pytest +import torch + +from fastgen.configs.config_utils import override_config_with_opts +from fastgen.configs.methods.config_anyflow import ModelConfig +from fastgen.methods import AnyFlowModel +from fastgen.methods.distribution_matching.anyflow_scheduler import FlowMapDiscreteScheduler +from fastgen.utils.test_utils import check_grad_zero + + +def _build_pretrain_model(): + gc.collect() + instance = ModelConfig() + instance.loss_config.training_stage = "pretrain" + # Use a small finite-difference step relative to t in [0, 1]. + instance.loss_config.jvp_finite_diff_eps = 1e-2 + + opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"] + instance.net = override_config_with_opts(instance.net, opts) + instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16" + instance.pretrained_model_path = "" + instance.input_shape = [3, 2, 2] + + model = AnyFlowModel(instance) + model.on_train_begin() + model.init_optimizers() + return model + + +def _build_onpolicy_model(): + """On-policy fixture mirrors test_dmd2model: img_resolution=8 so the + discriminator's 4x4 conv kernels can operate.""" + gc.collect() + instance = ModelConfig() + instance.loss_config.training_stage = "onpolicy" + + opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"] + instance.net = override_config_with_opts(instance.net, opts) + opts_disc = ["-", "feature_indices=[0]", "all_res=[8]", "in_channels=128"] + instance.discriminator = override_config_with_opts(instance.discriminator, opts_disc) + + instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16" + instance.pretrained_model_path = "" + instance.student_update_freq = 2 + instance.input_shape = [3, 8, 8] + + model = AnyFlowModel(instance) + model.on_train_begin() + model.init_optimizers() + return model + + +def _make_data(model, img_resolution: int = 2): + batch_size = 1 + labels = torch.nn.functional.one_hot(torch.randint(0, 10, (batch_size,)), num_classes=10) + neg_labels = torch.zeros(batch_size, 10) + return { + "real": torch.randn(batch_size, 3, img_resolution, img_resolution).to(model.device, model.precision), + "condition": labels.to(model.device, model.precision), + "neg_condition": neg_labels.to(model.device, model.precision), + } + + +# --------------------------------------------------------------------------- +# Pretrain stage +# --------------------------------------------------------------------------- + + +def test_pretrain_single_train_step(): + model = _build_pretrain_model() + data = _make_data(model) + + loss_map, outputs = model.single_train_step(data, 0) + + assert "total_loss" in loss_map + assert "anyflow_loss" in loss_map + assert "dF_dt_target_norm" in loss_map + assert torch.isfinite(loss_map["total_loss"]).all() + assert "gen_rand" in outputs + assert isinstance(outputs["gen_rand"], torch.Tensor) + + +def test_pretrain_no_fake_score_or_discriminator(): + """Pretrain stage must not instantiate DMD2's fake_score / discriminator.""" + model = _build_pretrain_model() + assert not hasattr(model, "fake_score") or model.fake_score is None or "fake_score" not in model.model_dict + assert "fake_score" not in model.model_dict + assert "discriminator" not in model.model_dict + + +def test_pretrain_optimizer_step(): + model = _build_pretrain_model() + data = _make_data(model) + for iteration in range(2): + model.optimizers_zero_grad(iteration) + loss_map, _ = model.single_train_step(data, iteration) + model.grad_scaler.scale(loss_map["total_loss"]).backward() + model.optimizers_schedulers_step(iteration) + # After one zero_grad with no backward in between, gradients should be cleared. + model.optimizers_zero_grad(2) + check_grad_zero(model.net) + + +def test_pretrain_central_difference_falls_back_at_boundaries(): + """When (t ± δ) leaves [min_t, max_t], the one-sided fallback should still + produce a finite target. We synthesise a worst-case t at the boundary.""" + model = _build_pretrain_model() + ns = model.net.noise_scheduler + + real = torch.randn(2, 3, 2, 2, device=model.device, dtype=model.precision) + cond = torch.nn.functional.one_hot(torch.tensor([0, 1]), num_classes=10).to(model.device, model.precision) + + t = torch.tensor([float(ns.min_t), float(ns.max_t)], device=model.device, dtype=ns.t_precision) + r = torch.tensor([float(ns.min_t), float(ns.min_t)], device=model.device, dtype=ns.t_precision) + + eps_noise = torch.randn_like(real) + x_t = ns.forward_process(real, eps_noise, t) + v = eps_noise - real + + target = model._compute_central_difference_target(x_t, t, r, v, cond) + assert torch.isfinite(target).all(), "boundary samples must yield finite targets" + + +# --------------------------------------------------------------------------- +# On-policy stage — inherits DMD2's alternating updates +# --------------------------------------------------------------------------- + + +def test_onpolicy_student_update_step(): + model = _build_onpolicy_model() + data = _make_data(model, img_resolution=8) + loss_map, outputs = model.single_train_step(data, 0) # iteration 0 -> student update + assert "total_loss" in loss_map + assert "vsd_loss" in loss_map + assert "gen_rand" in outputs + + +def test_onpolicy_fake_score_discriminator_update_step(): + model = _build_onpolicy_model() + model.precision = torch.float32 + model.on_train_begin() + data = _make_data(model, img_resolution=8) + for k, v in data.items(): + if isinstance(v, torch.Tensor): + data[k] = v.to(model.precision) + loss_map, outputs = model.single_train_step(data, 1) # iteration 1 -> fake_score/disc update + assert "fake_score_loss" in loss_map + assert "gan_loss_disc" in loss_map + assert "gen_rand" in outputs + + +def test_onpolicy_optimizer_step(): + model = _build_onpolicy_model() + data = _make_data(model, img_resolution=8) + for iteration in range(2): + model.optimizers_zero_grad(iteration) + loss_map, _ = model.single_train_step(data, iteration) + model.grad_scaler.scale(loss_map["total_loss"]).backward() + model.optimizers_schedulers_step(iteration) + + +# --------------------------------------------------------------------------- +# Scheduler +# --------------------------------------------------------------------------- + + +def test_flowmap_scheduler_apply_shift_identity_for_shift1(): + scheduler = FlowMapDiscreteScheduler(num_train_timesteps=1000, shift=1.0, weight_type="beta08") + sigmas = torch.linspace(0.0, 1.0, 11) + assert torch.allclose(scheduler.apply_shift(sigmas), sigmas) + + +def test_flowmap_scheduler_step_zero_interval(): + scheduler = FlowMapDiscreteScheduler(num_train_timesteps=1000, shift=1.0, weight_type="uniform") + sample = torch.randn(2, 4, 8, 8) + model_output = torch.randn_like(sample) + t = torch.tensor([500.0, 500.0]) + # r = t => zero-length integration interval; the sample should be unchanged. + out = scheduler.step(model_output, sample, timestep=t, r_timestep=t.clone()) + assert torch.allclose(out, sample, atol=1e-5) + + +@pytest.mark.parametrize("weight_type", ["gaussian", "beta08", "uniform"]) +def test_flowmap_scheduler_weights_positive(weight_type): + scheduler = FlowMapDiscreteScheduler(num_train_timesteps=1000, shift=1.0, weight_type=weight_type) + t = torch.tensor([100.0, 500.0, 900.0]) + w = scheduler.get_train_weight(t) + assert torch.all(w >= 0), f"{weight_type} weights must be non-negative" + assert torch.isfinite(w).all() From aec727d5fc10081564b895a815b4a575ac883110 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Sat, 16 May 2026 09:37:23 +0800 Subject: [PATCH 02/19] Wan: add gated r-embedder fusion + AnyFlow weight remap AnyFlow's released HF checkpoints store the r-pathway as ``condition_embedder.delta_embedder.*`` inside the shared ``WanTwoTimeTextImageEmbedding`` module and use ONE shared ``time_proj`` for both t and (t, r). Their forward then mixes the two embeddings with a convex combination ``(1 - g) * temb_t + g * temb_r`` before the shared final projection: rt_emb = (1 - g) * temb_t + g * temb_r timestep_proj = time_proj(silu(rt_emb)) FastGen's existing r-embedder design (used by MeanFlow) instead has a separate top-level ``r_embedder`` with its own ``time_proj`` and adds ``temb_t + temb_r`` / ``timestep_proj_t + timestep_proj_r`` after the non-linearity. The two layouts are not functionally equivalent because ``silu`` is non-linear. Two changes: * ``Wan.__init__``: add ``r_embedder_fusion: str = "additive"`` (default preserves MeanFlow's behaviour) and ``r_embedder_gate_value: float = 0.25``. When ``r_embedder_fusion="gated"``, ``classify_forward_prepare`` computes the convex-mix variant and uses ``r_embedder.time_proj`` (which ``init_embedder`` already deep-copies from ``condition_embedder.time_proj``) for the shared final projection. * ``fastgen/methods/distribution_matching/anyflow.py``: add ``remap_anyflow_keys`` helper that rewrites AnyFlow's ``condition_embedder.delta_embedder.linear_{1,2}.*`` to FastGen's ``r_embedder.time_embedder.linear_{1,2}.*`` and duplicates ``condition_embedder.time_proj.*`` into ``r_embedder.time_proj.*`` so the two projections start identical. The function is a no-op when no AnyFlow-format keys are present. Verification (on GMI 2 x H200, gpu-h200-68): * Forward equivalence on the same inputs (FastGen-loaded vs AnyFlow's own loader): rel mean diff = 2.8% in bf16 (forward noise floor). * Training-step loss equivalence (AnyFlow ``train_bidirection`` math reproduced inline on both code paths, same seed): AnyFlow loss 0.381619 vs FastGen loss 0.397162, rel diff = 4.07%. * 4-step Euler-flow inference end-to-end (text encoder + FastGen Wan + VAE decode) produces a finite 81-frame 480x832 video matching the AnyFlow paper's any-step inference pattern. Signed-off-by: Enderfga --- .../methods/distribution_matching/anyflow.py | 39 ++++++++++++++++++ fastgen/networks/Wan/network.py | 41 ++++++++++++++++--- 2 files changed, 74 insertions(+), 6 deletions(-) diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index cbae8fe..38dd4ea 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -46,6 +46,45 @@ import fastgen.utils.logging_utils as logger +def remap_anyflow_keys(state_dict: dict) -> dict: + """Remap an AnyFlow HF release state_dict to FastGen's Wan layout. + + AnyFlow's ``FAR_Wan_Transformer3DModel`` stores the r-pathway inside the + main ``condition_embedder`` as ``delta_embedder``, and uses ONE shared + ``time_proj`` for both t and (t, r). FastGen exposes a separate top-level + ``r_embedder`` with its own ``time_embedder`` + ``time_proj``. The two + layouts are functionally equivalent (FastGen's ``r_embedder.time_proj`` + starts as a deepcopy of ``condition_embedder.time_proj`` per + :meth:`Wan.init_embedder`), so we just rename / duplicate the tensors. + + The function is a no-op when no ``condition_embedder.delta_embedder.*`` + keys are present, so it's safe to call unconditionally. + """ + delta_keys = [k for k in state_dict if k.startswith("condition_embedder.delta_embedder.")] + if not delta_keys: + return state_dict + new_sd = dict(state_dict) + for k in delta_keys: + # condition_embedder.delta_embedder.linear_1.weight + # -> r_embedder.time_embedder.linear_1.weight + new_k = k.replace("condition_embedder.delta_embedder.", "r_embedder.time_embedder.") + new_sd[new_k] = new_sd.pop(k) + # AnyFlow's gated fusion shares the final time_proj. FastGen has a + # separate r_embedder.time_proj that mathematically substitutes for the + # shared one when fusion="gated"; copy the weights across so the two + # projections start identical (and AnyFlow's training never diverges them). + for sub in ("weight", "bias"): + src = f"condition_embedder.time_proj.{sub}" + dst = f"r_embedder.time_proj.{sub}" + if src in new_sd and dst not in new_sd: + new_sd[dst] = new_sd[src].clone() + logger.info( + f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors " + "and duplicated time_proj weights into r_embedder." + ) + return new_sd + + if TYPE_CHECKING: from fastgen.configs.methods.config_anyflow import ModelConfig diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index 5181764..02021e2 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -338,14 +338,27 @@ def classify_forward_prepare( r_timestep = r_timestep.to(time_embedder_dtype) remb = self.r_embedder.time_embedder(r_timestep).type_as(encoder_hidden_states) - r_timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb)) - r_timestep_proj = unflatten_timestep_proj(r_timestep_proj, rs_seq_len) - if self.encoder_depth is None: - timestep_proj = timestep_proj + r_timestep_proj - temb = temb + remb + # AnyFlow-style gated mixing: convex-combine t- and r-embeddings BEFORE + # the shared final projection, matching the WanTwoTimeTextImageEmbedding + # forward in the AnyFlow reference. The released AnyFlow HF checkpoints + # require this fusion to reproduce the published forward pass. + if getattr(self.r_embedder, "fusion_mode", "additive") == "gated": + gate = self.r_embedder.gate_value.to(remb.dtype) + rt_emb = (1 - gate) * temb + gate * remb + timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(rt_emb)) + timestep_proj = unflatten_timestep_proj(timestep_proj, rs_seq_len) + r_timestep_proj = None + temb = rt_emb else: - temb = remb + r_timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb)) + r_timestep_proj = unflatten_timestep_proj(r_timestep_proj, rs_seq_len) + + if self.encoder_depth is None: + timestep_proj = timestep_proj + r_timestep_proj + temb = temb + remb + else: + temb = remb elif r_timestep is not None: # Raise an error here, otherwise we silently ignore the r_timestep raise ValueError("r_timestep provided but no r_embedder is present") @@ -557,6 +570,8 @@ def __init__( load_pretrained: bool = True, use_fsdp_checkpoint: bool = True, use_wan_official_sinusoidal: bool = False, + r_embedder_fusion: str = "additive", + r_embedder_gate_value: float = 0.25, **model_kwargs, ): """Wan2.1/2.2 model constructor. @@ -610,6 +625,20 @@ def __init__( if r_timestep: logger.info(f"Initializing r embedder with {r_embedder_init}") self.transformer.r_embedder = self.init_embedder(r_embedder_init) + # Stash fusion config on the r_embedder module so the (method-bound) + # forward override can branch on it without changing its signature. + if r_embedder_fusion not in ("additive", "gated"): + raise ValueError(f"r_embedder_fusion must be 'additive' or 'gated', got {r_embedder_fusion!r}") + self.transformer.r_embedder.fusion_mode = r_embedder_fusion + self.transformer.r_embedder.register_buffer( + "gate_value", + torch.tensor([float(r_embedder_gate_value)], dtype=torch.float32), + persistent=False, + ) + logger.info( + f"r_embedder fusion={r_embedder_fusion}" + + (f" gate_value={r_embedder_gate_value}" if r_embedder_fusion == "gated" else "") + ) else: self.transformer.r_embedder = None From bcd78d679253938a214fab17e890df6a6d3772dc Mon Sep 17 00:00:00 2001 From: Enderfga Date: Sat, 16 May 2026 13:15:51 +0800 Subject: [PATCH 03/19] anyflow: multi-step rollout-with-gradient on-policy student generation Replace the single-step student forward in AnyFlow's on-policy stage with a multi-step Euler-flow rollout that enables gradients at one randomly-chosen step. This matches AnyFlow's ``WanAnyFlowPipeline.training_rollout`` (the published on-policy training mode in the reference repo) and gives the DMD generator update a usable gradient through a full denoising window instead of a single forward. Changes: * ``AnyFlowModel._rollout_with_gradient(batch_size, dtype, condition)``: start from pure noise at ``ns.max_t``, iterate ``student_sample_steps`` Euler-flow updates with ``r = t_next`` (mean-velocity, matching the reference default), and toggle ``torch.set_grad_enabled`` at the randomly-selected step. ``grad_step`` is broadcast from rank 0 in distributed runs so all ranks share the same gradient window. The step schedule honours ``sample_t_cfg.t_list`` when set, otherwise falls back to ``noise_scheduler.get_t_list``. * ``_onpolicy_student_update_step`` and ``_onpolicy_fake_score_discriminator_update_step``: source ``gen_data`` from the rollout instead of a single ``self.net(input_student, ...)`` forward. ``input_student`` / ``t_student`` from ``_generate_noise_and_time`` become unused for on-policy and are discarded explicitly. * ``_get_outputs``: when on-policy, always take the multi-step generator callable path (no longer special-cases ``student_sample_steps == 1`` for the validation hook, since the rollout output is always usable). * ``tests/test_anyflowmodel.py``: bump ``student_sample_steps`` to 2 in the on-policy fixtures and add ``test_onpolicy_rollout_propagates_gradient`` which asserts the rollout output keeps a usable autograd graph and that ``backward()`` reaches the student weights. All 13 unit tests pass (`make pytest tests/test_anyflowmodel.py`). Signed-off-by: Enderfga --- .../methods/distribution_matching/anyflow.py | 106 +++++++++++++++--- tests/test_anyflowmodel.py | 26 +++++ 2 files changed, 117 insertions(+), 15 deletions(-) diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index 38dd4ea..31cd861 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -326,6 +326,77 @@ def _zeros_like_t(self, t: torch.Tensor) -> torch.Tensor: ns = self.net.noise_scheduler return torch.full_like(t, float(ns.min_t)) + def _sample_grad_step(self, num_steps: int) -> int: + """Pick one step index in ``[0, num_steps - 1]`` to enable gradients on. + + Broadcast from rank 0 in distributed runs so all ranks agree on the + same window — matches AnyFlow's reference (`training_rollout` in + ``pipeline_wan_anyflow.py`` ``broadcast(sample_step, src=0)``). + """ + idx = torch.randint(0, num_steps, (1,), device=self.device, dtype=torch.long) + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.broadcast(idx, src=0) + return int(idx.item()) + + def _rollout_with_gradient( + self, + batch_size: int, + dtype: torch.dtype, + condition: Optional[Any], + ) -> torch.Tensor: + """Multi-step student rollout from pure noise with one gradient-enabled step. + + Mirrors AnyFlow's ``WanAnyFlowPipeline.training_rollout`` (see + ``pipeline_wan_anyflow.py`` lines 370--454). Starts from + :math:`x_T \\sim \\mathcal{N}(0, \\sigma_{\\max}^2)`, runs + :attr:`student_sample_steps` Euler-flow steps with ``r = t_next`` + (mean-velocity sampling, matching AnyFlow's ``use_mean_velocity=True`` + default), and enables gradients at one randomly-chosen step index so + the DMD generator update receives a usable gradient through one full + denoising forward. + """ + num_steps = int(self.config.student_sample_steps) + if num_steps < 1: + raise ValueError(f"student_sample_steps must be >=1, got {num_steps}") + + ns = self.net.noise_scheduler + grad_step = self._sample_grad_step(num_steps) + + # Initial latents at the maximum-noise timestep. + eps_init = torch.randn(batch_size, *self.input_shape, device=self.device, dtype=dtype) + x = ns.latents(noise=eps_init) + + # Timestep schedule. Use config-provided t_list when set (matches + # AnyFlow's hand-tuned step lists, e.g. [0.999, 0.937, 0.833, 0.624, 0.0] + # for 4-step Wan); otherwise fall back to the scheduler's default. + if self.config.sample_t_cfg.t_list is not None: + t_list = torch.tensor(self.config.sample_t_cfg.t_list, device=self.device, dtype=ns.t_precision) + if len(t_list) != num_steps + 1: + raise ValueError( + f"sample_t_cfg.t_list has {len(t_list)} entries, " + f"expected {num_steps + 1} for student_sample_steps={num_steps}" + ) + else: + t_list = ns.get_t_list(sample_steps=num_steps, device=self.device) + + for step in range(num_steps): + t_cur = t_list[step].expand(batch_size).to(ns.t_precision) + t_next = t_list[step + 1].expand(batch_size).to(ns.t_precision) + + enable_grad = (step == grad_step) and torch.is_grad_enabled() + with torch.set_grad_enabled(enable_grad): + # Mean-velocity flow prediction: u_theta(x_t, t, r=t_next) + flow_pred = self.net(x, t_cur, r=t_next, condition=condition, fwd_pred_type="flow") + + # Euler-flow step. Keep ``x`` attached to the autograd graph + # unconditionally: steps run inside ``torch.no_grad()`` do not add + # graph nodes anyway, but detaching would also strip the gradient + # that earlier grad-enabled step(s) installed onto ``x``. + delta_t = (t_cur - t_next).view(batch_size, *([1] * (x.ndim - 1))).to(x.dtype) + x = x - delta_t * flow_pred + + return x + def _onpolicy_student_update_step( self, input_student: torch.Tensor, @@ -336,11 +407,16 @@ def _onpolicy_student_update_step( condition: Optional[Any], neg_condition: Optional[Any], ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: - r_zero = self._zeros_like_t(t_student) r_zero_t = self._zeros_like_t(t) + del input_student, t_student # unused — rollout starts from fresh pure noise - # Student rollout to a single x0 estimate, conditioned on r=0. - gen_data = self.net(input_student, t_student, r=r_zero, condition=condition, fwd_pred_type="x0") + # Multi-step rollout-with-gradient from pure noise; gradient is enabled + # at one randomly-chosen step matching AnyFlow's training_rollout. + gen_data = self._rollout_with_gradient( + batch_size=data["real"].shape[0], + dtype=data["real"].dtype, + condition=condition, + ) perturbed_data = self.net.noise_scheduler.forward_process(gen_data, eps, t) with torch.no_grad(): @@ -383,7 +459,7 @@ def _onpolicy_student_update_step( "vsd_loss": vsd_loss, "gan_loss_gen": gan_loss_gen, } - outputs = self._get_outputs(gen_data, input_student, condition=condition) + outputs = self._get_outputs(gen_data, condition=condition) return loss_map, outputs def _onpolicy_fake_score_discriminator_update_step( @@ -395,11 +471,16 @@ def _onpolicy_fake_score_discriminator_update_step( real_data: torch.Tensor, condition: Optional[Any], ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: - r_zero = self._zeros_like_t(t_student) r_zero_t = self._zeros_like_t(t) + del input_student, t_student # unused in rollout path + # Generate via the same multi-step rollout (no gradient needed on the + # fake-score / discriminator update step — caller wraps the whole + # update in no_grad effectively by detaching everything below). with torch.no_grad(): - gen_data = self.net(input_student, t_student, r=r_zero, condition=condition, fwd_pred_type="x0") + gen_data = self._rollout_with_gradient( + batch_size=real_data.shape[0], dtype=real_data.dtype, condition=condition + ) x_t_sg = self.net.noise_scheduler.forward_process(gen_data, eps, t) from fastgen.methods.common_loss import ( @@ -480,7 +561,7 @@ def _onpolicy_fake_score_discriminator_update_step( } if self.config.gan_loss_weight_gen > 0 and self.config.gan_r1_reg_weight > 0: loss_map["gan_loss_ar1"] = gan_loss_ar1 - outputs = self._get_outputs(gen_data, input_student, condition=condition) + outputs = self._get_outputs(gen_data, condition=condition) return loss_map, outputs def _onpolicy_single_train_step( @@ -515,14 +596,9 @@ def _get_outputs( noise = input_student / (ns.max_sigma if hasattr(ns, "max_sigma") else 1.0) return {"gen_rand": gen_data, "input_rand": noise} - # On-policy stage delegates to DMD2's get_outputs path so multi-step - # generators are produced consistently with the rest of the family. - if self.config.student_sample_steps == 1: - assert input_student is not None, "input_student must be provided" - ns = self.net.noise_scheduler - noise = input_student / (ns.max_sigma if hasattr(ns, "max_sigma") else 1.0) - return {"gen_rand": gen_data, "input_rand": noise} - + # On-policy stage: gen_data already comes from the multi-step + # rollout (`_rollout_with_gradient`), so a fresh noise sample is + # sufficient for the validation generator hook. noise = torch.randn_like(gen_data, dtype=self.precision) gen_rand_func = partial( self.generator_fn, diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index ac732e8..e27a73a 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -56,6 +56,10 @@ def _build_onpolicy_model(): instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16" instance.pretrained_model_path = "" instance.student_update_freq = 2 + # Exercise the multi-step rollout-with-gradient path (matches the + # AnyFlow paper's on-policy training mode). student_sample_steps=2 is + # the smallest value that runs the rollout loop more than once. + instance.student_sample_steps = 2 instance.input_shape = [3, 8, 8] model = AnyFlowModel(instance) @@ -163,6 +167,28 @@ def test_onpolicy_fake_score_discriminator_update_step(): assert "gen_rand" in outputs +def test_onpolicy_rollout_propagates_gradient(): + """The multi-step rollout must allow gradient flow on the chosen step. + + Mirrors AnyFlow's ``training_rollout`` (pipeline_wan_anyflow.py L370): + one randomly-chosen step in the rollout has gradients enabled; the + remaining steps are no_grad. ``gen_data.requires_grad`` should be True + at the rollout output so the DMD generator update has a valid gradient. + """ + model = _build_onpolicy_model() + real = torch.randn(1, 3, 8, 8, device=model.device, dtype=model.precision) + cond = torch.nn.functional.one_hot(torch.tensor([0]), num_classes=10).to(model.device, model.precision) + # Exercise rollout under grad-enabled context (mirrors student update step). + gen = model._rollout_with_gradient(batch_size=real.shape[0], dtype=real.dtype, condition=cond) + assert tuple(gen.shape) == (1, 3, 8, 8), f"rollout output shape mismatch: {gen.shape}" + assert gen.requires_grad, "rollout output must keep autograd graph at the chosen step" + # Trivial scalar loss; backward should succeed without NaN. + loss = gen.float().pow(2).mean() + loss.backward() + grad_seen = any(p.grad is not None and torch.isfinite(p.grad).all() for p in model.net.parameters()) + assert grad_seen, "no gradient reached the student network through the rollout" + + def test_onpolicy_optimizer_step(): model = _build_onpolicy_model() data = _make_data(model, img_resolution=8) From 5437b4b0a4b5aaefb4be760d3d20000bbb9e9364 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Fri, 22 May 2026 10:02:35 +0800 Subject: [PATCH 04/19] Wan: extract _fuse_r_embedding helper; ship AnyFlow on-policy config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses two pieces of PR #25 reviewer feedback: (1) Code sharing with MeanFlow. The previous commit added AnyFlow's gated t/r mixing as an inline branch inside ``classify_forward_prepare``, which made it visually hard to tell which lines were MeanFlow's additive path and which were AnyFlow-specific. This commit factors both fusion modes into a single ``_fuse_r_embedding`` method bound on the transformer (parallel pattern to ``classify_forward_prepare`` and friends). Both paths still share ``r_embedder.time_embedder`` / ``time_proj`` / ``act_fn`` modules — the helper just makes that sharing explicit and shrinks the call site to three lines. Forward semantics are bit-identical to the previous commit for both additive (MeanFlow) and gated (AnyFlow) modes across all three ``encoder_depth`` cases. (2) Ship a paper-aligned on-policy stage config. Previously the only documented way to run Stage 3 was an inline tweak in the pretrain config docstring. New file ``fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py`` inherits the pretrain config and flips the loss into "onpolicy" with the paper's Stage 3 hyperparameters (lr=2e-6, 1200 iter, GAN on at the DMD2-default 0.03, ``student_update_freq=5``). The docstring notes that the AnyFlow paper's rank-256 LoRA variant is not reproduced here because FastGen does not ship a PEFT/LoRA training path; this config is a full-rank fine-tune of a Stage 2 pretrain checkpoint. The AnyFlow method README is updated to (a) document the new ``r_embedder_fusion="gated"`` requirement when loading the released AnyFlow HF checkpoints, (b) replace the stale "multi-step rollout deferred to a follow-up" note (already landed in ab1174d) with an explicit acknowledgement that end-to-end convergence-scale validation on the paper's training corpus is deferred to a follow-up, and (c) cross-reference both pretrain and on-policy configs. Tests: all 13 AnyFlow + 3 MeanFlow unit tests pass. Signed-off-by: Enderfga --- .../WanT2V/config_anyflow_onpolicy.py | 40 +++++++++++ .../methods/distribution_matching/README.md | 8 ++- fastgen/networks/Wan/network.py | 67 +++++++++++++------ 3 files changed, 92 insertions(+), 23 deletions(-) create mode 100644 fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py new file mode 100644 index 0000000..e0086df --- /dev/null +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -0,0 +1,40 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Reference AnyFlow on-policy distillation config on Wan-1.3B T2V (Stage 3). + +Inherits the pretrain config and flips the loss into the on-policy stage, +turning on DMD2's alternating fake_score / discriminator updates with the +``r=0`` dual-timestep conditioning that AnyFlow keeps from MeanFlow. Mirrors +the paper's Stage 3 hyperparameters: 1.2k iterations at lr=2e-6 on top of a +Stage 2 flow-map pretrain checkpoint. + +Note: the AnyFlow paper trains this stage with a rank-256 LoRA adapter, but +FastGen does not ship a PEFT/LoRA training path today, so this config does +full-rank fine-tuning. Set ``config.model.pretrained_student_net_path`` to a +checkpoint produced by :mod:`config_anyflow` before launching. +""" + +import fastgen.configs.experiments.WanT2V.config_anyflow as config_anyflow_pretrain + + +def create_config(): + config = config_anyflow_pretrain.create_config() + + config.model.loss_config.training_stage = "onpolicy" + config.model.pretrained_student_net_path = "" + + # Re-enable the DMD2 alternating-update machinery. + config.model.gan_loss_weight_gen = 0.03 + config.model.student_update_freq = 5 + + # Stage 3 learning rates from the AnyFlow paper. + config.model.net_optimizer.lr = 2e-6 + config.model.fake_score_optimizer.lr = 2e-6 + config.model.discriminator_optimizer.lr = 2e-6 + + config.trainer.max_iter = 1200 + config.trainer.save_ckpt_iter = 200 + + config.log_config.group = "wan_anyflow_onpolicy" + return config diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index 6e9c278..c9c78a8 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -102,11 +102,13 @@ Single model that supports arbitrary inference NFE by learning a flow map `u_θ( - `loss_config.shift`: flow-matching schedule shift (5.0 for Wan video) - See also key parameters of DMD2 above (used by the on-policy stage) -**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). +**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). When loading the published AnyFlow HF checkpoints, set `r_embedder_fusion="gated"` on the Wan constructor — this routes the t/r mix through `Wan/network.py::_fuse_r_embedding`'s gated branch (shared with MeanFlow's additive default) so the released weights reproduce bit-for-bit. -**Note:** the on-policy stage in this PR uses single-step student generation. Multi-step rollout-with-gradient (matching `self_forcing.py`'s `rollout_with_gradient`) is intentionally deferred to a follow-up PR. +**Note:** correctness of the port is established via forward-parity and single-step training-step parity against the AnyFlow reference (see PR #25 discussion). End-to-end convergence-scale validation on the paper's training corpus is deferred to a follow-up. -**Configs:** [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) +**Configs:** +- [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) — Stage 2 pretrain (6k iter, lr=5e-5, shift=5, beta08, paper-aligned) +- [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) — Stage 3 on-policy distillation (1.2k iter, lr=2e-6, GAN on; full-rank fine-tune of a Stage 2 pretrain ckpt — the paper's rank-256 LoRA variant awaits a PEFT path in FastGen) --- diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index 02021e2..8172e39 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -274,6 +274,49 @@ def classify_forward( return out +def _fuse_r_embedding( + self, + temb: torch.Tensor, + timestep_proj: torch.Tensor, + remb: torch.Tensor, + rs_seq_len: Optional[torch.LongTensor], +) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """Combine the t- and r-time embeddings into the final timestep projection. + + Two fusion modes share the same ``r_embedder.time_proj`` / ``act_fn`` modules: + + * ``additive`` (MeanFlow default, ``r_embedder.fusion_mode`` absent or "additive"): + ``r_timestep_proj = time_proj(act_fn(remb))`` is added to ``timestep_proj`` + and ``remb`` is added to ``temb``. T2V dual-stream variants (``encoder_depth`` + set) instead replace ``temb`` with ``remb`` and leave ``timestep_proj`` alone. + + * ``gated`` (AnyFlow, opt-in): convex-combine the two embeddings *before* the + shared projection — ``rt_emb = (1-g)·temb + g·remb`` then + ``timestep_proj = time_proj(act_fn(rt_emb))``. This matches the + ``WanTwoTimeTextImageEmbedding.forward_timestep`` path in the AnyFlow + reference and is required to reproduce the published HF checkpoint + forward bit-for-bit. + + Returns ``(temb, timestep_proj, r_timestep_proj)`` where ``r_timestep_proj`` + is ``None`` in the gated branch (the t/r mix is folded into ``timestep_proj`` + itself). + """ + fusion = getattr(self.r_embedder, "fusion_mode", "additive") + + if fusion == "gated": + gate = self.r_embedder.gate_value.to(remb.dtype) + rt_emb = (1 - gate) * temb + gate * remb + ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(rt_emb)) + return rt_emb, unflatten_timestep_proj(ts_proj, rs_seq_len), None + + # additive — MeanFlow original path, bit-identical + r_ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb)) + r_ts_proj = unflatten_timestep_proj(r_ts_proj, rs_seq_len) + if self.encoder_depth is None: + return temb + remb, timestep_proj + r_ts_proj, r_ts_proj + return remb, timestep_proj, r_ts_proj + + def classify_forward_prepare( self, hidden_states: torch.Tensor, @@ -339,26 +382,9 @@ def classify_forward_prepare( remb = self.r_embedder.time_embedder(r_timestep).type_as(encoder_hidden_states) - # AnyFlow-style gated mixing: convex-combine t- and r-embeddings BEFORE - # the shared final projection, matching the WanTwoTimeTextImageEmbedding - # forward in the AnyFlow reference. The released AnyFlow HF checkpoints - # require this fusion to reproduce the published forward pass. - if getattr(self.r_embedder, "fusion_mode", "additive") == "gated": - gate = self.r_embedder.gate_value.to(remb.dtype) - rt_emb = (1 - gate) * temb + gate * remb - timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(rt_emb)) - timestep_proj = unflatten_timestep_proj(timestep_proj, rs_seq_len) - r_timestep_proj = None - temb = rt_emb - else: - r_timestep_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb)) - r_timestep_proj = unflatten_timestep_proj(r_timestep_proj, rs_seq_len) - - if self.encoder_depth is None: - timestep_proj = timestep_proj + r_timestep_proj - temb = temb + remb - else: - temb = remb + temb, timestep_proj, r_timestep_proj = self._fuse_r_embedding( + temb, timestep_proj, remb, rs_seq_len + ) elif r_timestep is not None: # Raise an error here, otherwise we silently ignore the r_timestep raise ValueError("r_timestep provided but no r_embedder is present") @@ -861,6 +887,7 @@ def override_transformer_forward(self, inner_dim: int) -> None: # Override transformer forward methods with custom implementations for block in self.transformer.blocks: block.forward = types.MethodType(block_forward, block) + self.transformer._fuse_r_embedding = types.MethodType(_fuse_r_embedding, self.transformer) self.transformer.classify_forward_prepare = types.MethodType(classify_forward_prepare, self.transformer) self.transformer.classify_forward_block_forward = types.MethodType( classify_forward_block_forward, self.transformer From 994aad15b79da97af7400c968853572842d7154a Mon Sep 17 00:00:00 2001 From: Enderfga Date: Wed, 10 Jun 2026 22:25:45 +0800 Subject: [PATCH 05/19] anyflow: reuse MeanFlow for pretrain and stock DMD2 for on-policy MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review feedback on PR #25: - Pretrain (Stage 2) now runs MeanFlowModel directly: add a fixed per-timestep loss weighting (loss_config.weight_type, evaluated as a function of t) and a consistency bucket (sample_t_cfg.consistency_ratio) to MeanFlow; both default off. Drop FlowMapDiscreteScheduler — (t, r) pair sampling uses noise_scheduler.sample_t with the shifted distribution, weights need no precomputed table. - On-policy (Stage 3) keeps only two DMD2 overrides: _generate_noise_and_time (start from pure noise at max_t) and gen_data_from_net (multi-step Euler-flow rollout with one gradient-enabled step). Teacher/fake_score are flow-map networks queried at the instantaneous velocity r=t — the reference passes r_timestep=timesteps to both (not r=0 as before) — implemented as the network-level default for dual-timestep nets when r is not passed. - Wan: store r_embedder.gate_value as a plain float (no buffer to re-materialize in reset_parameters for FSDP); gated fusion respects encoder_depth (mirrors additive); move remap_anyflow_keys here and apply it inside Wan.load_state_dict. - Configs: set r_embedder_fusion=gated + time_cond_type=abs on both stages, matching the published checkpoints (deltatime_type 'r', gate 0.25); imports at module top throughout. --- .../experiments/WanT2V/config_anyflow.py | 84 ++- .../WanT2V/config_anyflow_onpolicy.py | 60 +- fastgen/configs/methods/config_anyflow.py | 44 +- fastgen/configs/methods/config_mean_flow.py | 8 + fastgen/methods/__init__.py | 4 +- .../methods/consistency_model/mean_flow.py | 52 +- .../methods/distribution_matching/README.md | 24 +- .../methods/distribution_matching/anyflow.py | 651 +++--------------- .../anyflow_scheduler.py | 118 ---- fastgen/networks/EDM/network.py | 6 + fastgen/networks/Wan/network.py | 82 ++- tests/test_anyflowmodel.py | 257 +++++-- 12 files changed, 535 insertions(+), 855 deletions(-) delete mode 100644 fastgen/methods/distribution_matching/anyflow_scheduler.py diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py index 08badbf..e95027c 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -1,65 +1,83 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Reference AnyFlow experiment config on Wan-1.3B T2V. +"""AnyFlow flow-map pretrain config on Wan-1.3B T2V (paper Stage 2). -Mirrors the AnyFlow paper's pretrain configuration: 1.3B student initialised -from a Wan2.1-T2V checkpoint, flow-matching shift=5, beta08 loss weighting, -6k iterations with batch_size_global=32 and lr=5e-5. +AnyFlow's pretrain objective is the MeanFlow objective with a fixed +``beta08`` per-timestep loss weighting, a finite-difference JVP, shifted +timestep sampling, and a ``consistency_ratio`` fraction of the batch pinned +to ``r = min_t`` — so this config runs :class:`MeanFlowModel` directly. -Switching to the on-policy stage: +Mirrors the AnyFlow paper's pretrain recipe (``shift=5``, ``beta08`` +weighting, ``epsilon=5e-3``, ``diffusion_ratio=0.5``, +``consistency_ratio=0.25``, lr=5e-5, 6k iterations, batch_size_global=32, +text dropout 0.1 with velocity-fused guidance 3.0) on the gated dual-timestep +Wan architecture (``gate=0.25``, absolute-r conditioning, matching +``deltatime_type: r`` in the reference). - config.model.loss_config.training_stage = "onpolicy" - config.model.pretrained_student_net_path = "" - -and adjust ``student_update_freq`` / ``gan_loss_weight_gen`` to taste. +The on-policy stage (paper Stage 3) lives in ``config_anyflow_onpolicy.py``. """ -import fastgen.configs.methods.config_anyflow as config_anyflow_default +import copy + +import fastgen.configs.methods.config_mean_flow as config_mean_flow from fastgen.configs.data import VideoLoaderConfig -from fastgen.configs.discriminator import Discriminator_Wan_1_3B_Config from fastgen.configs.net import Wan_1_3B_Config def create_config(): - config = config_anyflow_default.create_config() - - # Default to the pretrain stage; flip the switch to "onpolicy" once the - # flow-map pretrain checkpoint is available. - config.model.loss_config.training_stage = "pretrain" - config.model.loss_config.jvp_finite_diff_eps = 5e-3 - config.model.loss_config.diffusion_ratio = 0.5 - config.model.loss_config.consistency_ratio = 0.25 - config.model.loss_config.weight_type = "beta08" - config.model.loss_config.shift = 5.0 + config = config_mean_flow.create_config() - config.model.net = Wan_1_3B_Config + # ------ network: gated dual-timestep Wan (AnyFlow architecture) ------ + config.model.net = copy.deepcopy(Wan_1_3B_Config) config.model.net.r_timestep = True - - # The on-policy stage uses these too, but they are harmless in pretrain. - config.model.discriminator = Discriminator_Wan_1_3B_Config - config.model.discriminator.disc_type = "multiscale_down_mlp_large" - config.model.discriminator.feature_indices = [15, 22, 29] - config.model.gan_loss_weight_gen = 0.0 # disabled by default in pretrain - config.model.guidance_scale = 5.0 + config.model.net.encoder_depth = None + # AnyFlow conditions the r-pathway on the absolute r (deltatime_type "r") + # and fuses the two time embeddings with a fixed convex gate of 0.25. + config.model.net.time_cond_type = "abs" + config.model.net.r_embedder_fusion = "gated" + config.model.net.r_embedder_gate_value = 0.25 config.model.precision = "bfloat16" # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. config.model.input_shape = [16, 21, 60, 104] - config.model.net_optimizer.lr = 5e-5 - config.model.fake_score_optimizer.lr = 5e-5 - config.model.discriminator_optimizer.lr = 5e-5 + # ------ AnyFlow loss: MeanFlow l2 with fixed beta08 weighting ------ + config.model.loss_config.use_cd = False + config.model.loss_config.loss_type = "l2" + config.model.loss_config.weight_type = "beta08" + config.model.loss_config.use_jvp_finite_diff = True + config.model.loss_config.jvp_finite_diff_eps = 5e-3 + config.model.precision_amp_jvp = "float32" + # Velocity-level guidance fusion with text dropout (reference: + # drop_text_ratio=0.1, fuse_guidance_scale=3.0). + config.model.guidance_scale = 3.0 + config.model.cond_dropout_prob = 0.1 + + # ------ (t, r) sampling: shifted uniform pairs + AnyFlow buckets ------ config.model.sample_t_cfg.time_dist_type = "shifted" + config.model.sample_t_cfg.shift = 5.0 config.model.sample_t_cfg.min_t = 0.001 config.model.sample_t_cfg.max_t = 0.999 + # diffusion_ratio=0.5 of the batch keeps r = t (pure flow matching); + # r_sample_ratio is the complementary fraction that keeps the sampled r. + config.model.sample_t_cfg.r_sample_ratio = 0.5 + # consistency_ratio=0.25 of the batch is pinned to r = min_t. + config.model.sample_t_cfg.consistency_ratio = 0.25 + + # ------ optimization (reference: AdamW lr=5e-5, wd=0, betas=(0.9, 0.95)) ------ + config.model.net_optimizer.optim_type = "adamw" + config.model.net_optimizer.lr = 5e-5 + config.model.net_optimizer.betas = (0.9, 0.95) + config.model.net_optimizer.weight_decay = 0.0 + # ------ inference / validation ------ config.model.student_sample_type = "ode" - # Any-step model — multiple NFEs validated at inference time. config.model.student_sample_steps = 4 config.model.sample_t_cfg.t_list = [0.999, 0.937, 0.833, 0.624, 0.0] + # ------ data / trainer ------ config.dataloader_train = VideoLoaderConfig config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index e0086df..f5c5b9f 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -1,13 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Reference AnyFlow on-policy distillation config on Wan-1.3B T2V (Stage 3). +"""AnyFlow on-policy distillation config on Wan-1.3B T2V (paper Stage 3). -Inherits the pretrain config and flips the loss into the on-policy stage, -turning on DMD2's alternating fake_score / discriminator updates with the -``r=0`` dual-timestep conditioning that AnyFlow keeps from MeanFlow. Mirrors -the paper's Stage 3 hyperparameters: 1.2k iterations at lr=2e-6 on top of a -Stage 2 flow-map pretrain checkpoint. +DMD2's alternating fake_score / discriminator updates with a flow-map student +that generates via a multi-step rollout-with-gradient (see +:class:`~fastgen.methods.distribution_matching.anyflow.AnyFlowModel`). +Mirrors the paper's Stage 3 hyperparameters: 1.2k iterations at lr=2e-6 on +top of a Stage 2 flow-map pretrain checkpoint (``config_anyflow.py``). Note: the AnyFlow paper trains this stage with a rank-256 LoRA adapter, but FastGen does not ship a PEFT/LoRA training path today, so this config does @@ -15,26 +15,64 @@ checkpoint produced by :mod:`config_anyflow` before launching. """ -import fastgen.configs.experiments.WanT2V.config_anyflow as config_anyflow_pretrain +import copy + +import fastgen.configs.methods.config_anyflow as config_anyflow_default +from fastgen.configs.data import VideoLoaderConfig +from fastgen.configs.discriminator import Discriminator_Wan_1_3B_Config +from fastgen.configs.net import Wan_1_3B_Config def create_config(): - config = config_anyflow_pretrain.create_config() + config = config_anyflow_default.create_config() + + # ------ network: gated dual-timestep Wan, same as the pretrain stage ------ + config.model.net = copy.deepcopy(Wan_1_3B_Config) + config.model.net.r_timestep = True + config.model.net.encoder_depth = None + config.model.net.time_cond_type = "abs" + config.model.net.r_embedder_fusion = "gated" + config.model.net.r_embedder_gate_value = 0.25 - config.model.loss_config.training_stage = "onpolicy" config.model.pretrained_student_net_path = "" - # Re-enable the DMD2 alternating-update machinery. + config.model.precision = "bfloat16" + # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. + config.model.input_shape = [16, 21, 60, 104] + + # ------ DMD2 machinery ------ + config.model.discriminator = Discriminator_Wan_1_3B_Config + config.model.discriminator.disc_type = "multiscale_down_mlp_large" + config.model.discriminator.feature_indices = [15, 22, 29] config.model.gan_loss_weight_gen = 0.03 config.model.student_update_freq = 5 + config.model.guidance_scale = 3.0 - # Stage 3 learning rates from the AnyFlow paper. + config.model.sample_t_cfg.time_dist_type = "shifted" + config.model.sample_t_cfg.shift = 5.0 + config.model.sample_t_cfg.min_t = 0.001 + config.model.sample_t_cfg.max_t = 0.999 + + # ------ student rollout (AnyFlow's hand-tuned 4-step schedule) ------ + config.model.student_sample_type = "ode" + config.model.student_sample_steps = 4 + config.model.sample_t_cfg.t_list = [0.999, 0.937, 0.833, 0.624, 0.0] + + # ------ Stage 3 learning rates from the AnyFlow paper ------ config.model.net_optimizer.lr = 2e-6 config.model.fake_score_optimizer.lr = 2e-6 config.model.discriminator_optimizer.lr = 2e-6 + # ------ data / trainer ------ + config.dataloader_train = VideoLoaderConfig + config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) + config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 + config.dataloader_train.batch_size = 1 + config.trainer.max_iter = 1200 + config.trainer.logging_iter = 100 config.trainer.save_ckpt_iter = 200 + config.trainer.batch_size_global = 32 config.log_config.group = "wan_anyflow_onpolicy" return config diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py index 96700a8..bf6b6d2 100644 --- a/fastgen/configs/methods/config_anyflow.py +++ b/fastgen/configs/methods/config_anyflow.py @@ -1,11 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Config schema for the AnyFlow method. +"""Config schema for the AnyFlow on-policy method (paper Stage 3). -AnyFlow inherits the DMD2 model config (so the on-policy stage gets fake_score -/ discriminator / alternating-step machinery for free) and adds a -``LossConfig`` describing the flow-map pretrain hyperparameters. +AnyFlow's on-policy stage is DMD2 with a flow-map student that generates via +a multi-step rollout-with-gradient, so the config is the DMD2 model config +unchanged. The flow-map pretrain stage (paper Stage 2) is MeanFlow with +AnyFlow's hyperparameters — see ``configs/experiments/WanT2V/config_anyflow.py``. """ import attrs @@ -25,40 +26,9 @@ from fastgen.utils import LazyCall as L -@attrs.define(slots=False) -class LossConfig: - """Hyperparameters for the AnyFlow flow-map loss and on-policy switch.""" - - # Which stage to train. "pretrain" runs the central-difference flow-map - # objective; "onpolicy" inherits DMD2's alternating distillation with - # dual-timestep r=0 conditioning. - training_stage: str = "pretrain" - - # Central-difference step size for estimating dF/dt. Lives in the same - # units as the noise scheduler's timesteps. The default (5e-3) matches the - # AnyFlow paper's choice of epsilon=5 with num_train_timesteps=1000. - jvp_finite_diff_eps: float = 5e-3 - - # Per-batch fraction with r = t (recovers pure flow matching). - diffusion_ratio: float = 0.5 - # Per-batch fraction with r = min_t (forces consistency to clean data). - consistency_ratio: float = 0.25 - - # Per-timestep loss weighting scheme — passed through to the flow-map - # scheduler. One of "gaussian", "beta08", "uniform". - weight_type: str = "beta08" - # Flow-matching schedule shift for the weighting / sampling scheduler. - # Wan video defaults use 5.0; image use 1.0. - shift: float = 1.0 - # Resolution of the discrete weighting grid; matches the AnyFlow reference. - num_train_timesteps: int = 1000 - - @attrs.define(slots=False) class ModelConfig(DMD2ModelConfig): - """AnyFlow model config — inherits DMD2 fields, adds the flow-map loss config.""" - - loss_config: LossConfig = attrs.field(factory=LossConfig) + """AnyFlow on-policy model config — identical to DMD2's.""" @attrs.define(slots=False) @@ -82,7 +52,7 @@ def create_config(): } ) - # Pretrain stage relies on a flow-matching net_pred_type and dual-timestep input. + # The student is a flow-map network with a dual-timestep input. config.model.use_ema = True config.model.net.r_timestep = True config.model.net_scheduler.warm_up_steps = [0] diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index a6dba34..ff8fda4 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -36,6 +36,10 @@ class SampleTConfig(BaseSampleTConfig): # ratio for randomly sampling r r_sample_ratio: float = 0.0 + # fraction of the batch forced to r = min_t (consistency-to-clean, + # used by AnyFlow; 0.0 keeps the original MeanFlow behavior) + consistency_ratio: float = 0.0 + @attrs.define(slots=False) class SampleRConfig(BaseSampleTConfig): @@ -71,6 +75,10 @@ class LossConfig: tangent_spatial_invariance: bool = False # loss type (choice between l2 and opt_grad) loss_type: str = "opt_grad" + # fixed per-timestep loss weighting evaluated as a function of t + # ("beta08", "gaussian", "uniform"; used by AnyFlow). None keeps the + # adaptive norm_method weighting above. Only applies to loss_type="l2". + weight_type: Optional[str] = None @attrs.define(slots=False) diff --git a/fastgen/methods/__init__.py b/fastgen/methods/__init__.py index 907b3f3..0508a27 100644 --- a/fastgen/methods/__init__.py +++ b/fastgen/methods/__init__.py @@ -9,13 +9,15 @@ from fastgen.methods.distribution_matching.causvid import CausVidModel as CausVidModel from fastgen.methods.distribution_matching.self_forcing import SelfForcingModel as SelfForcingModel -from fastgen.methods.distribution_matching.anyflow import AnyFlowModel as AnyFlowModel from fastgen.methods.consistency_model.CM import CMModel as CMModel from fastgen.methods.consistency_model.TCM import TCMModel as TCMModel from fastgen.methods.consistency_model.sCM import SCMModel as SCMModel from fastgen.methods.consistency_model.mean_flow import MeanFlowModel as MeanFlowModel +# AnyFlow reuses MeanFlow's student sample loop, so it must come after. +from fastgen.methods.distribution_matching.anyflow import AnyFlowModel as AnyFlowModel + from fastgen.methods.fine_tuning.sft import SFTModel as SFTModel from fastgen.methods.fine_tuning.sft import CausalSFTModel as CausalSFTModel diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 485e752..20a2ac5 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -68,6 +68,10 @@ def __init__(self, config: ModelConfig): self.precision_amp_jvp = basic_utils.PRECISION_MAP[self.config.precision_amp_jvp] logger.critical(f"Using precision {self.precision_amp_jvp} for JVP") + # Normalization constant for the fixed per-timestep loss weighting + # (computed lazily on first use, see _get_timestep_weight). + self._timestep_weight_scale: Optional[float] = None + def _mix_condition( self, condition: Any, @@ -251,6 +255,34 @@ def net_wrapper(x_t, t, r): return u_theta_jvp + def _timestep_weight_raw(self, t: torch.Tensor) -> torch.Tensor: + """Unnormalized per-timestep weight as a direct function of t in [0, 1].""" + weight_type = self.loss_config.weight_type + if weight_type == "beta08": + return t * (1 - t).clamp(min=0).sqrt() + if weight_type == "gaussian": + # exp(-2 (t - 1/2)^2), shifted so the minimum over [0, 1] is zero. + return (torch.exp(-2 * (t - 0.5) ** 2) - float(torch.exp(torch.tensor(-0.5)))).clamp(min=0) + if weight_type == "uniform": + return torch.ones_like(t) + raise ValueError(f"Invalid weight_type: {weight_type!r}") + + @torch.no_grad() + def _get_timestep_weight(self, t: torch.Tensor) -> torch.Tensor: + """Fixed per-timestep loss weight w(t) (AnyFlow-style). + + The weight is evaluated directly as a function of t. The normalization + constant matches the AnyFlow reference scheduler, which normalizes the + weights over its shifted discrete timestep grid; we compute the same + constant once and cache it. + """ + if self._timestep_weight_scale is None: + shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 + grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64) + grid = shift * grid / (1 + (shift - 1) * grid) + self._timestep_weight_scale = float(1000.0 / self._timestep_weight_raw(grid).sum()) + return self._timestep_weight_raw(t) * self._timestep_weight_scale + @torch.no_grad() def _compute_weight(self, tensor: torch.Tensor) -> torch.Tensor: norm_method, *norm_args = self.loss_config.norm_method.split("_") @@ -306,9 +338,15 @@ def _mf_pred_to_loss( if self.loss_config.loss_type == "l2": tangent = dxt_dt - warmup_weight * delta_t * u_theta_jvp loss = (u_theta - tangent).pow(2) - loss = torch.sum(loss, dim=list(range(1, loss.ndim))) - weight = self._compute_weight(loss) - loss = loss * weight + if self.loss_config.weight_type is not None: + # Fixed per-timestep weighting with mean reduction (AnyFlow). + loss = torch.mean(loss, dim=list(range(1, loss.ndim))) + weight = self._get_timestep_weight(t) + loss = loss * weight + else: + loss = torch.sum(loss, dim=list(range(1, loss.ndim))) + weight = self._compute_weight(loss) + loss = loss * weight # use explicit gradient elif self.loss_config.loss_type == "opt_grad": @@ -468,6 +506,14 @@ def single_train_step( zero_mask = torch.arange(batch_size, device=self.device) < flow_matching_size r = torch.where(zero_mask, t, r) + # set r=min_t (consistency-to-clean, AnyFlow) for a subset of the batch. + # Taken from the tail so it never overlaps the flow-matching head. + if self.sample_t_cfg.consistency_ratio > 0: + num_consistency = (torch.rand(batch_size, device=self.device) < self.sample_t_cfg.consistency_ratio).sum() + num_consistency = torch.minimum(num_consistency, batch_size - flow_matching_size) + consistency_mask = torch.arange(batch_size, device=self.device) >= batch_size - num_consistency + r = torch.where(consistency_mask, torch.full_like(r, self.net.noise_scheduler.min_t), r) + ( mf_loss, v_loss, diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index c9c78a8..4fa411e 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -91,18 +91,20 @@ DMD2 extended for causal video generation with autoregressive chunk-by-chunk pro Single model that supports arbitrary inference NFE by learning a flow map `u_θ(x_t, t, r)` (average velocity from `t` back to `r`). Trained in two stages: -1. **Pretrain** — flow-map prediction with a central-difference target that reuses the network's own forward at `(t ± δ, r)` to estimate `dF/dt`. Per-batch sampling assigns `r = t` to a `diffusion_ratio` fraction (recovering plain flow matching) and `r = 0` to a `consistency_ratio` fraction (forcing consistency to clean data). -2. **On-policy** — distribution-matching distillation with `r = 0` conditioning on top of the pretrained flow-map weights. Inherits DMD2's alternating fake_score / discriminator / VSD machinery. +1. **Pretrain** — the MeanFlow objective with AnyFlow's hyperparameters, run directly on [`MeanFlowModel`](../consistency_model/mean_flow.py): finite-difference JVP, a fixed `beta08` per-timestep loss weight (`loss_config.weight_type`), shifted timestep sampling, and a `sample_t_cfg.consistency_ratio` fraction of the batch pinned to `r = min_t`. There is no AnyFlow-specific pretrain code. +2. **On-policy** — distribution-matching distillation on top of the pretrained flow-map weights ([`anyflow.py`](anyflow.py)). Stock DMD2 with two overrides: the student generates via a multi-step Euler-flow rollout from pure noise with gradient enabled at one randomly-chosen step (`gen_data_from_net`), and always starts from `max_t` (`_generate_noise_and_time`). The teacher / fake_score are flow-map networks queried at the instantaneous velocity `r = t` (the network-level default for gated-fusion Wan when `r` is not passed). -**Key Parameters:** -- `loss_config.training_stage`: `"pretrain"` or `"onpolicy"` -- `loss_config.jvp_finite_diff_eps`: central-difference step δ (in noise scheduler t-units) -- `loss_config.diffusion_ratio` / `loss_config.consistency_ratio`: per-batch fraction with `r=t` / `r=0` -- `loss_config.weight_type`: `gaussian` | `beta08` | `uniform` per-timestep loss weight -- `loss_config.shift`: flow-matching schedule shift (5.0 for Wan video) -- See also key parameters of DMD2 above (used by the on-policy stage) - -**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). When loading the published AnyFlow HF checkpoints, set `r_embedder_fusion="gated"` on the Wan constructor — this routes the t/r mix through `Wan/network.py::_fuse_r_embedding`'s gated branch (shared with MeanFlow's additive default) so the released weights reproduce bit-for-bit. +**Key Parameters (on-policy):** +- `student_sample_steps` / `sample_t_cfg.t_list`: rollout length and hand-tuned step schedule +- See also key parameters of DMD2 above + +**Key Parameters (pretrain, on MeanFlow):** +- `loss_config.weight_type`: `beta08` | `gaussian` | `uniform` fixed per-timestep loss weight +- `loss_config.use_jvp_finite_diff` / `loss_config.jvp_finite_diff_eps`: central-difference step δ +- `sample_t_cfg.r_sample_ratio` / `sample_t_cfg.consistency_ratio`: per-batch fraction with sampled `r` / `r = min_t` +- `sample_t_cfg.time_dist_type="shifted"`, `sample_t_cfg.shift`: shifted timestep sampling (5.0 for Wan video) + +**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). When loading the published AnyFlow HF checkpoints, set `r_embedder_fusion="gated"` and `time_cond_type="abs"` on the Wan constructor — this routes the t/r mix through `Wan/network.py::_fuse_r_embedding`'s gated branch (shared with MeanFlow's additive default) so the released weights reproduce bit-for-bit. `Wan.load_state_dict` remaps the AnyFlow checkpoint layout (`condition_embedder.delta_embedder.*`) automatically. **Note:** correctness of the port is established via forward-parity and single-step training-step parity against the AnyFlow reference (see PR #25 discussion). End-to-end convergence-scale validation on the paper's training corpus is deferred to a follow-up. diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index 31cd861..2d8629f 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -3,369 +3,132 @@ """AnyFlow — any-step video diffusion with flow maps and on-policy distillation. -AnyFlow trains a single model :math:`u_\\theta(x_t, t, r)` that predicts the -average velocity from time ``t`` to ``r`` (with ``r \\le t``). Once trained, -the same model supports arbitrary inference step counts: each Euler-like -sampling step picks its own integration interval ``(t \\rightarrow r)``. - -Training has two stages, selected via ``config.loss_config.training_stage``: - -* ``"pretrain"`` — flow-map prediction with a central-difference target - - v(x_t, t) = eps - x_0 # instantaneous flow (flow matching) - dF/dt ~= (u_theta(x_{t+d}, t+d, r) - u_theta(x_{t-d}, t-d, r)) / 2d - target = v - (t - r) * dF/dt - loss = weight(t) * MSE(u_theta(x_t, t, r), target) - - Per-batch sampling assigns ``r = t`` for a ``diffusion_ratio`` fraction - (recovering plain flow matching), ``r = 0`` for a ``consistency_ratio`` - fraction (forcing consistency to clean data), and a uniform random pair - otherwise — matching the AnyFlow paper. - -* ``"onpolicy"`` — distribution-matching distillation on top of pretrained - flow-map weights. Inherits DMD2's fake_score / teacher / discriminator - machinery and alternating-step optimisation, but conditions all forwards - on ``r = 0`` (predicting the full flow from ``t`` to clean). Multi-step - rollout-with-gradient is intentionally deferred to a follow-up PR. - -The network must support a secondary timestep argument ``r`` (Wan with -``r_timestep=True`` does; MeanFlow already exercises this same code path). +AnyFlow trains a single flow-map model :math:`u_\\theta(x_t, t, r)` that +predicts the average velocity from time ``t`` to ``r`` (with ``r \\le t``). +Once trained, the same model supports arbitrary inference step counts: each +Euler-like sampling step picks its own integration interval +``(t \\rightarrow r)``. + +Training has two stages: + +* **Flow-map pretrain** (paper Stage 2) is MeanFlow with AnyFlow's + hyperparameters — fixed ``beta08`` per-timestep loss weighting, + finite-difference JVP, a ``consistency_ratio`` fraction of the batch pinned + to ``r = min_t``, and shifted timestep sampling. It is configured directly + on :class:`~fastgen.methods.consistency_model.mean_flow.MeanFlowModel` + (see ``configs/experiments/WanT2V/config_anyflow.py``); there is no + AnyFlow-specific pretrain code. + +* **On-policy** (paper Stage 3, this module) is DMD2 on top of the pretrained + flow-map weights. The only differences from stock DMD2 are the two narrow + overrides below: + + - :meth:`AnyFlowModel._generate_noise_and_time` — the student always starts + from pure noise at ``max_t`` (DMD2's multi-step branch would noise real + data instead); + - :meth:`AnyFlowModel.gen_data_from_net` — the student generates via a + multi-step Euler-flow rollout with ``r = t_next`` (mean-velocity sampling) + and gradient enabled at one randomly-chosen step, matching AnyFlow's + ``WanAnyFlowPipeline.training_rollout``. + + The teacher and fake_score are flow-map networks queried at the + instantaneous velocity ``r = t`` — for gated-fusion Wan networks this is the + network-level default when ``r`` is not passed (see + ``fastgen/networks/Wan/network.py``), so DMD2's student / fake-score / + discriminator update steps are reused unchanged. """ from __future__ import annotations -from functools import partial -from typing import Any, Callable, Dict, Optional, TYPE_CHECKING +from typing import Any, Optional, TYPE_CHECKING import torch -import torch.nn.functional as F -from fastgen.methods.distribution_matching.anyflow_scheduler import FlowMapDiscreteScheduler +from fastgen.methods.consistency_model.mean_flow import MeanFlowModel from fastgen.methods.distribution_matching.dmd2 import DMD2Model -from fastgen.utils import expand_like +from fastgen.utils.basic_utils import convert_cfg_to_dict import fastgen.utils.logging_utils as logger -def remap_anyflow_keys(state_dict: dict) -> dict: - """Remap an AnyFlow HF release state_dict to FastGen's Wan layout. - - AnyFlow's ``FAR_Wan_Transformer3DModel`` stores the r-pathway inside the - main ``condition_embedder`` as ``delta_embedder``, and uses ONE shared - ``time_proj`` for both t and (t, r). FastGen exposes a separate top-level - ``r_embedder`` with its own ``time_embedder`` + ``time_proj``. The two - layouts are functionally equivalent (FastGen's ``r_embedder.time_proj`` - starts as a deepcopy of ``condition_embedder.time_proj`` per - :meth:`Wan.init_embedder`), so we just rename / duplicate the tensors. - - The function is a no-op when no ``condition_embedder.delta_embedder.*`` - keys are present, so it's safe to call unconditionally. - """ - delta_keys = [k for k in state_dict if k.startswith("condition_embedder.delta_embedder.")] - if not delta_keys: - return state_dict - new_sd = dict(state_dict) - for k in delta_keys: - # condition_embedder.delta_embedder.linear_1.weight - # -> r_embedder.time_embedder.linear_1.weight - new_k = k.replace("condition_embedder.delta_embedder.", "r_embedder.time_embedder.") - new_sd[new_k] = new_sd.pop(k) - # AnyFlow's gated fusion shares the final time_proj. FastGen has a - # separate r_embedder.time_proj that mathematically substitutes for the - # shared one when fusion="gated"; copy the weights across so the two - # projections start identical (and AnyFlow's training never diverges them). - for sub in ("weight", "bias"): - src = f"condition_embedder.time_proj.{sub}" - dst = f"r_embedder.time_proj.{sub}" - if src in new_sd and dst not in new_sd: - new_sd[dst] = new_sd[src].clone() - logger.info( - f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors " - "and duplicated time_proj weights into r_embedder." - ) - return new_sd - - if TYPE_CHECKING: from fastgen.configs.methods.config_anyflow import ModelConfig class AnyFlowModel(DMD2Model): - """AnyFlow training method. + """AnyFlow on-policy stage: DMD2 with a multi-step rollout-with-gradient student.""" - See module docstring for the algorithm. - """ + # Validation sampling integrates the flow map with r = t_next (ode) just + # like MeanFlow; FastGenModel's default x0-prediction loop never passes r. + _student_sample_loop = MeanFlowModel._student_sample_loop def __init__(self, config: ModelConfig): super().__init__(config) self.config = config - self.loss_config = self.config.loss_config - - if self.loss_config.training_stage not in ("pretrain", "onpolicy"): - raise ValueError( - f"training_stage must be 'pretrain' or 'onpolicy', got {self.loss_config.training_stage!r}" - ) - - # Standalone scheduler used for inference and for the per-timestep - # training weight in the pretrain stage. Training noising still goes - # through ``self.net.noise_scheduler`` to stay compatible with DMD2. - self._flowmap_scheduler = FlowMapDiscreteScheduler( - num_train_timesteps=self.loss_config.num_train_timesteps, - shift=self.loss_config.shift, - weight_type=self.loss_config.weight_type, - ) - - if self.loss_config.training_stage == "pretrain": - logger.info( - f"AnyFlow pretrain stage: epsilon={self.loss_config.jvp_finite_diff_eps}, " - f"diffusion_ratio={self.loss_config.diffusion_ratio}, " - f"consistency_ratio={self.loss_config.consistency_ratio}, " - f"weight_type={self.loss_config.weight_type}" - ) - else: - logger.info( - f"AnyFlow on-policy stage: student_update_freq={self.config.student_update_freq}, " - f"gan_loss_weight_gen={self.config.gan_loss_weight_gen}" - ) - - # ------------------------------------------------------------------ - # Build / optimisation overrides — skip DMD2 plumbing in pretrain - # ------------------------------------------------------------------ - - def build_model(self): - """In pretrain mode skip fake_score / discriminator entirely.""" - if self.config.loss_config.training_stage == "pretrain": - # Bypass DMD2Model.build_model — pretrain only needs the student. - # Call grandparent's build_model (FastGenModel) directly. - super(DMD2Model, self).build_model() - self.load_student_weights_and_ema() - return - super().build_model() - - def init_optimizers(self): - """Pretrain skips the DMD2 fake_score / discriminator optimisers.""" - if self.config.loss_config.training_stage == "pretrain": - # Bypass DMD2Model.init_optimizers — only the student optimiser exists. - super(DMD2Model, self).init_optimizers() - return - super().init_optimizers() - - @property - def model_dict(self): - if self.config.loss_config.training_stage == "pretrain": - return super(DMD2Model, self).model_dict - return super().model_dict - - @property - def optimizer_dict(self): - if self.config.loss_config.training_stage == "pretrain": - return super(DMD2Model, self).optimizer_dict - return super().optimizer_dict - - @property - def scheduler_dict(self): - if self.config.loss_config.training_stage == "pretrain": - return super(DMD2Model, self).scheduler_dict - return super().scheduler_dict - - # ------------------------------------------------------------------ - # Pretrain stage - # ------------------------------------------------------------------ - - def _sample_pair_timesteps( - self, batch_size: int, dtype: torch.dtype - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Sample ``(t, r)`` with ``t >= r``, plus a per-sample diffusion mask. - - Implements AnyFlow's partitioning of the batch: - * a ``diffusion_ratio`` fraction has ``r = t`` (pure flow matching) - * a ``consistency_ratio`` fraction has ``r = min_t`` - * the rest gets a uniform random pair - """ - ns = self.net.noise_scheduler - t_dtype = ns.t_precision - - t_1 = torch.rand(batch_size, device=self.device, dtype=t_dtype) - t_2 = torch.rand(batch_size, device=self.device, dtype=t_dtype) - t_norm = torch.maximum(t_1, t_2) - r_norm = torch.minimum(t_1, t_2) - - # Shift to match the flow-matching schedule (Wan default uses shift=5). - t_norm = self._flowmap_scheduler.apply_shift(t_norm) - r_norm = self._flowmap_scheduler.apply_shift(r_norm) - - # Rescale unit-interval timesteps into the noise scheduler's [min_t, max_t]. - max_t = float(ns.max_t) - min_t = float(ns.min_t) - scale = max_t - min_t - t = t_norm * scale + min_t - r = r_norm * scale + min_t - - # Per-batch bucket assignment. We shuffle so the buckets are randomly - # distributed within the local batch — this matches the paper's intent - # without requiring global cross-rank coordination. - n_diffusion = int(round(self.loss_config.diffusion_ratio * batch_size)) - n_consistency = int(round(self.loss_config.consistency_ratio * batch_size)) - n_diffusion = min(n_diffusion, batch_size) - n_consistency = min(n_consistency, batch_size - n_diffusion) - - perm = torch.randperm(batch_size, device=self.device) - is_diffusion = torch.zeros(batch_size, dtype=torch.bool, device=self.device) - is_consistency = torch.zeros(batch_size, dtype=torch.bool, device=self.device) - is_diffusion[perm[:n_diffusion]] = True - is_consistency[perm[n_diffusion : n_diffusion + n_consistency]] = True - - r = torch.where(is_diffusion, t, r) - r = torch.where(is_consistency, torch.full_like(r, min_t), r) - - return t.to(dtype=t_dtype), r.to(dtype=t_dtype), is_diffusion - - @torch.no_grad() - def _compute_central_difference_target( - self, - x_t: torch.Tensor, - t: torch.Tensor, - r: torch.Tensor, - v: torch.Tensor, - condition: Optional[Any], - ) -> torch.Tensor: - """Compute the AnyFlow flow-map target ``v - (t - r) * dF/dt``. - - ``dF/dt`` is estimated by central difference at ``(t ± delta, r)``. - Boundary cases near ``min_t`` / ``max_t`` fall back to one-sided - differences so the estimate stays valid for the whole timestep range. - """ - ns = self.net.noise_scheduler - max_t = float(ns.max_t) - min_t = float(ns.min_t) - delta = float(self.loss_config.jvp_finite_diff_eps) - - # Validity masks for each finite-difference direction. - is_fwd_valid = (t + delta) <= max_t - is_bwd_valid = ((t - delta) >= min_t) & ((t - delta) > r) - - use_central = is_fwd_valid & is_bwd_valid - use_fwd_only = is_fwd_valid & ~is_bwd_valid - use_bwd_only = ~is_fwd_valid & is_bwd_valid - - # Build per-sample (t_plus, t_minus, denom) with broadcasting-safe shapes. - t_plus = torch.where(is_fwd_valid, t + delta, t) - t_minus = torch.where(is_bwd_valid, t - delta, t) - denom = torch.where( - use_central, - torch.full_like(t, 2 * delta), - torch.where(use_fwd_only | use_bwd_only, torch.full_like(t, delta), torch.full_like(t, 1.0)), + logger.info( + f"AnyFlow on-policy: student_sample_steps={self.config.student_sample_steps}, " + f"student_update_freq={self.config.student_update_freq}, " + f"gan_loss_weight_gen={self.config.gan_loss_weight_gen}" ) - # Linear-path extrapolation along the flow direction: - # x_t = t * eps + (1 - t) * x_0 => d x_t / d t = eps - x_0 = v - x_t_plus = x_t + expand_like(t_plus - t, x_t) * v - x_t_minus = x_t + expand_like(t_minus - t, x_t) * v - - F_plus = self.net(x_t_plus, t_plus, r=r, condition=condition, fwd_pred_type="flow") - F_minus = self.net(x_t_minus, t_minus, r=r, condition=condition, fwd_pred_type="flow") - - dF_dt = (F_plus - F_minus) / expand_like(denom, x_t) - - # Where no finite-difference direction was valid (extremely rare with a - # small delta), fall back to dF/dt = 0 so we recover pure flow matching. - no_diff_valid = ~(use_central | use_fwd_only | use_bwd_only) - if no_diff_valid.any(): - dF_dt = torch.where(expand_like(no_diff_valid, dF_dt), torch.zeros_like(dF_dt), dF_dt) - - target = v - expand_like(t - r, x_t) * dF_dt - return target - - def _pretrain_single_train_step( - self, data: Dict[str, Any], iteration: int - ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: - """Single training step for AnyFlow pretrain.""" - real_data, condition, _ = self._prepare_training_data(data) - batch_size = real_data.shape[0] - - t, r, is_diffusion = self._sample_pair_timesteps(batch_size, dtype=real_data.dtype) - - # Forward noising along the linear flow-matching path. - eps = torch.randn_like(real_data) - x_t = self.net.noise_scheduler.forward_process(real_data, eps, t) - - # Ground-truth instantaneous flow direction at time t. - v = eps - real_data - - # Central-difference target (no grad through the finite-difference probes). - target = self._compute_central_difference_target(x_t, t, r, v, condition) - - # Student forward — this is the term whose gradient flows. - u_theta = self.net(x_t, t, r=r, condition=condition, fwd_pred_type="flow") - - # Per-sample MSE in float for numerical stability under bf16/fp16 AMP. - sq_err = (u_theta.float() - target.float()).pow(2) - loss_per_sample = sq_err.flatten(1).mean(dim=-1) - - # Per-timestep weight from the flow-map scheduler (beta08 default). - weight = self._flowmap_scheduler.get_train_weight(t).to(loss_per_sample.device, loss_per_sample.dtype) - loss = (loss_per_sample * weight).mean() - - # x0 approximation (monitoring only). - with torch.no_grad(): - x0_approx = self.net.noise_scheduler.flow_to_x0(x_t, u_theta.detach(), t) - - loss_map = { - "total_loss": loss, - "anyflow_loss": loss, - "flow_matching_loss": loss_per_sample[is_diffusion].mean() if is_diffusion.any() else loss.detach() * 0, - "dF_dt_target_norm": (target - v).flatten(1).norm(dim=-1).mean(), - } - outputs = self._get_outputs(x0_approx, input_student=x_t, condition=condition) - return loss_map, outputs - - # ------------------------------------------------------------------ - # On-policy stage — DMD2 with r=0 conditioning - # ------------------------------------------------------------------ - - def _zeros_like_t(self, t: torch.Tensor) -> torch.Tensor: - ns = self.net.noise_scheduler - return torch.full_like(t, float(ns.min_t)) - def _sample_grad_step(self, num_steps: int) -> int: - """Pick one step index in ``[0, num_steps - 1]`` to enable gradients on. + """Pick one rollout step index in ``[0, num_steps - 1]`` to enable gradients on. Broadcast from rank 0 in distributed runs so all ranks agree on the - same window — matches AnyFlow's reference (`training_rollout` in - ``pipeline_wan_anyflow.py`` ``broadcast(sample_step, src=0)``). + same window — matches AnyFlow's reference (``training_rollout`` in + ``pipeline_wan_anyflow.py``, ``broadcast(sample_step, src=0)``). """ idx = torch.randint(0, num_steps, (1,), device=self.device, dtype=torch.long) if torch.distributed.is_available() and torch.distributed.is_initialized(): torch.distributed.broadcast(idx, src=0) return int(idx.item()) - def _rollout_with_gradient( + def _generate_noise_and_time( + self, real_data: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The student rollout always starts from pure noise at ``max_t``.""" + batch_size = real_data.shape[0] + + eps_student = torch.randn(batch_size, *self.input_shape, device=self.device, dtype=real_data.dtype) + t_student = torch.full( + (batch_size,), + self.net.noise_scheduler.max_t, + device=self.device, + dtype=self.net.noise_scheduler.t_precision, + ) + input_student = self.net.noise_scheduler.latents(noise=eps_student) + + t = self.net.noise_scheduler.sample_t( + batch_size, **convert_cfg_to_dict(self.config.sample_t_cfg), device=self.device + ) + eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) + return input_student, t_student, t, eps + + def gen_data_from_net( self, - batch_size: int, - dtype: torch.dtype, - condition: Optional[Any], + input_student: torch.Tensor, + t_student: torch.Tensor, + condition: Optional[Any] = None, ) -> torch.Tensor: - """Multi-step student rollout from pure noise with one gradient-enabled step. - - Mirrors AnyFlow's ``WanAnyFlowPipeline.training_rollout`` (see - ``pipeline_wan_anyflow.py`` lines 370--454). Starts from - :math:`x_T \\sim \\mathcal{N}(0, \\sigma_{\\max}^2)`, runs - :attr:`student_sample_steps` Euler-flow steps with ``r = t_next`` - (mean-velocity sampling, matching AnyFlow's ``use_mean_velocity=True`` - default), and enables gradients at one randomly-chosen step index so - the DMD generator update receives a usable gradient through one full - denoising forward. + """Multi-step Euler-flow rollout with one gradient-enabled step. + + Mirrors AnyFlow's ``WanAnyFlowPipeline.training_rollout``: runs + ``student_sample_steps`` Euler-flow updates with ``r = t_next`` + (mean-velocity sampling, AnyFlow's ``use_mean_velocity=True`` default) + starting from ``input_student`` (pure noise at ``max_t``), and enables + gradients at one randomly-chosen step so the DMD generator update + receives a gradient through one full denoising forward. When the + caller wraps this in ``no_grad`` (the fake-score / discriminator + update), all steps run gradient-free. """ + del t_student # the rollout schedule comes from t_list below num_steps = int(self.config.student_sample_steps) if num_steps < 1: - raise ValueError(f"student_sample_steps must be >=1, got {num_steps}") + raise ValueError(f"student_sample_steps must be >= 1, got {num_steps}") ns = self.net.noise_scheduler + batch_size = input_student.shape[0] grad_step = self._sample_grad_step(num_steps) - # Initial latents at the maximum-noise timestep. - eps_init = torch.randn(batch_size, *self.input_shape, device=self.device, dtype=dtype) - x = ns.latents(noise=eps_init) - # Timestep schedule. Use config-provided t_list when set (matches # AnyFlow's hand-tuned step lists, e.g. [0.999, 0.937, 0.833, 0.624, 0.0] # for 4-step Wan); otherwise fall back to the scheduler's default. @@ -379,6 +142,7 @@ def _rollout_with_gradient( else: t_list = ns.get_t_list(sample_steps=num_steps, device=self.device) + x = input_student for step in range(num_steps): t_cur = t_list[step].expand(batch_size).to(ns.t_precision) t_next = t_list[step + 1].expand(batch_size).to(ns.t_precision) @@ -388,247 +152,10 @@ def _rollout_with_gradient( # Mean-velocity flow prediction: u_theta(x_t, t, r=t_next) flow_pred = self.net(x, t_cur, r=t_next, condition=condition, fwd_pred_type="flow") - # Euler-flow step. Keep ``x`` attached to the autograd graph - # unconditionally: steps run inside ``torch.no_grad()`` do not add - # graph nodes anyway, but detaching would also strip the gradient - # that earlier grad-enabled step(s) installed onto ``x``. + # Euler-flow step. Steps run under no_grad add no graph nodes, and + # keeping ``x`` attached preserves the gradient installed by the + # grad-enabled step. delta_t = (t_cur - t_next).view(batch_size, *([1] * (x.ndim - 1))).to(x.dtype) x = x - delta_t * flow_pred return x - - def _onpolicy_student_update_step( - self, - input_student: torch.Tensor, - t_student: torch.Tensor, - t: torch.Tensor, - eps: torch.Tensor, - data: Dict[str, Any], - condition: Optional[Any], - neg_condition: Optional[Any], - ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: - r_zero_t = self._zeros_like_t(t) - del input_student, t_student # unused — rollout starts from fresh pure noise - - # Multi-step rollout-with-gradient from pure noise; gradient is enabled - # at one randomly-chosen step matching AnyFlow's training_rollout. - gen_data = self._rollout_with_gradient( - batch_size=data["real"].shape[0], - dtype=data["real"].dtype, - condition=condition, - ) - perturbed_data = self.net.noise_scheduler.forward_process(gen_data, eps, t) - - with torch.no_grad(): - fake_score_x0 = self.fake_score(perturbed_data, t, r=r_zero_t, condition=condition, fwd_pred_type="x0") - - # Teacher prediction + optional GAN loss for the generator. Mirrors - # DMD2._compute_teacher_prediction_gan_loss but with r=0 conditioning. - if self.config.gan_loss_weight_gen > 0: - teacher_x0, fake_feat = self.teacher( - perturbed_data, - t, - r=r_zero_t, - condition=condition, - feature_indices=self.discriminator.feature_indices, - fwd_pred_type="x0", - ) - from fastgen.methods.common_loss import gan_loss_generator - - gan_loss_gen = gan_loss_generator(self.discriminator(fake_feat)) - else: - teacher_x0 = self.teacher(perturbed_data, t, r=r_zero_t, condition=condition, fwd_pred_type="x0") - gan_loss_gen = torch.tensor(0.0, device=self.device, dtype=teacher_x0.dtype) - teacher_x0 = teacher_x0.detach() - - # Optional CFG on the teacher. - if self.config.guidance_scale is not None: - with torch.no_grad(): - teacher_x0_neg = self.teacher( - perturbed_data, t, r=r_zero_t, condition=neg_condition, fwd_pred_type="x0" - ) - teacher_x0 = teacher_x0 + (self.config.guidance_scale - 1) * (teacher_x0 - teacher_x0_neg) - - from fastgen.methods.common_loss import variational_score_distillation_loss - - vsd_loss = variational_score_distillation_loss(gen_data, teacher_x0, fake_score_x0) - loss = vsd_loss + self.config.gan_loss_weight_gen * gan_loss_gen - - loss_map = { - "total_loss": loss, - "vsd_loss": vsd_loss, - "gan_loss_gen": gan_loss_gen, - } - outputs = self._get_outputs(gen_data, condition=condition) - return loss_map, outputs - - def _onpolicy_fake_score_discriminator_update_step( - self, - input_student: torch.Tensor, - t_student: torch.Tensor, - t: torch.Tensor, - eps: torch.Tensor, - real_data: torch.Tensor, - condition: Optional[Any], - ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]: - r_zero_t = self._zeros_like_t(t) - del input_student, t_student # unused in rollout path - - # Generate via the same multi-step rollout (no gradient needed on the - # fake-score / discriminator update step — caller wraps the whole - # update in no_grad effectively by detaching everything below). - with torch.no_grad(): - gen_data = self._rollout_with_gradient( - batch_size=real_data.shape[0], dtype=real_data.dtype, condition=condition - ) - x_t_sg = self.net.noise_scheduler.forward_process(gen_data, eps, t) - - from fastgen.methods.common_loss import ( - denoising_score_matching_loss, - gan_loss_discriminator, - ) - - fake_score_pred_type = self.config.fake_score_pred_type or self.teacher.net_pred_type - fake_score_pred = self.fake_score( - x_t_sg, t, r=r_zero_t, condition=condition, fwd_pred_type=fake_score_pred_type - ) - loss_fakescore = denoising_score_matching_loss( - fake_score_pred_type, - net_pred=fake_score_pred, - noise_scheduler=self.net.noise_scheduler, - x0=gen_data, - eps=eps, - t=t, - ) - - gan_loss_disc = torch.zeros_like(loss_fakescore) - gan_loss_ar1 = torch.zeros_like(loss_fakescore) - if self.config.gan_loss_weight_gen > 0: - with torch.no_grad(): - fake_feat = self.teacher( - x_t_sg, - t, - r=r_zero_t, - condition=condition, - return_features_early=True, - feature_indices=self.discriminator.feature_indices, - ) - # Real data path — mirror DMD2._compute_real_feat but pass r=0. - from fastgen.utils.basic_utils import convert_cfg_to_dict - - if self.config.gan_use_same_t_noise: - t_real, eps_real = t, eps - else: - t_real = self.net.noise_scheduler.sample_t( - real_data.shape[0], - **convert_cfg_to_dict(self.config.sample_t_cfg), - device=self.device, - ) - eps_real = torch.randn_like(real_data) - perturbed_real = self.net.noise_scheduler.forward_process(real_data, eps_real, t_real) - r_zero_real = self._zeros_like_t(t_real) - real_feat = self.teacher( - perturbed_real, - t_real, - r=r_zero_real, - condition=condition, - return_features_early=True, - feature_indices=self.discriminator.feature_indices, - ) - - real_feat_logit = self.discriminator(real_feat) - gan_loss_disc = gan_loss_discriminator(real_feat_logit, self.discriminator(fake_feat)) - - if self.config.gan_r1_reg_weight > 0: - perturbed_real_alpha = real_data.add(self.config.gan_r1_reg_alpha * torch.randn_like(real_data)) - with torch.no_grad(): - real_feat_alpha = self.teacher( - perturbed_real_alpha, - t_real, - r=r_zero_real, - condition=condition, - return_features_early=True, - feature_indices=self.discriminator.feature_indices, - ) - real_feat_alpha_logit = self.discriminator(real_feat_alpha) - gan_loss_ar1 = F.mse_loss(real_feat_logit, real_feat_alpha_logit, reduction="mean") - - loss = loss_fakescore + gan_loss_disc + self.config.gan_r1_reg_weight * gan_loss_ar1 - loss_map = { - "total_loss": loss, - "fake_score_loss": loss_fakescore, - "gan_loss_disc": gan_loss_disc, - } - if self.config.gan_loss_weight_gen > 0 and self.config.gan_r1_reg_weight > 0: - loss_map["gan_loss_ar1"] = gan_loss_ar1 - outputs = self._get_outputs(gen_data, condition=condition) - return loss_map, outputs - - def _onpolicy_single_train_step( - self, data: Dict[str, Any], iteration: int - ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: - real_data, condition, neg_condition = self._prepare_training_data(data) - self._setup_grad_requirements(iteration) - input_student, t_student, t, eps = self._generate_noise_and_time(real_data) - - if iteration % self.config.student_update_freq == 0: - return self._onpolicy_student_update_step( - input_student, t_student, t, eps, data, condition=condition, neg_condition=neg_condition - ) - return self._onpolicy_fake_score_discriminator_update_step( - input_student, t_student, t, eps, real_data, condition=condition - ) - - # ------------------------------------------------------------------ - # FastGenModel interface - # ------------------------------------------------------------------ - - def _get_outputs( - self, - gen_data: torch.Tensor, - input_student: Optional[torch.Tensor] = None, - condition: Any = None, - ) -> Dict[str, torch.Tensor | Callable]: - # Pretrain stage uses a direct x0 approximation tensor. - if self.loss_config.training_stage == "pretrain": - assert input_student is not None, "input_student must be provided" - ns = self.net.noise_scheduler - noise = input_student / (ns.max_sigma if hasattr(ns, "max_sigma") else 1.0) - return {"gen_rand": gen_data, "input_rand": noise} - - # On-policy stage: gen_data already comes from the multi-step - # rollout (`_rollout_with_gradient`), so a fresh noise sample is - # sufficient for the validation generator hook. - noise = torch.randn_like(gen_data, dtype=self.precision) - gen_rand_func = partial( - self.generator_fn, - net=self.net_inference, - noise=noise, - condition=condition, - student_sample_steps=self.config.student_sample_steps, - student_sample_type=self.config.student_sample_type, - t_list=self.config.sample_t_cfg.t_list, - precision_amp=self.precision_amp_infer, - ) - 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 - ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: - if self.loss_config.training_stage == "pretrain": - return self._pretrain_single_train_step(data, iteration) - return self._onpolicy_single_train_step(data, iteration) - - def get_optimizers(self, iteration: int) -> list: - """Pretrain stage uses only the student optimizer. - - On-policy stage inherits DMD2's alternating optimisation. - """ - if self.loss_config.training_stage == "pretrain": - return [self.net_optimizer] - return super().get_optimizers(iteration) - - def get_lr_schedulers(self, iteration: int) -> list: - if self.loss_config.training_stage == "pretrain": - return [self.net_lr_scheduler] - return super().get_lr_schedulers(iteration) diff --git a/fastgen/methods/distribution_matching/anyflow_scheduler.py b/fastgen/methods/distribution_matching/anyflow_scheduler.py deleted file mode 100644 index b17ffb2..0000000 --- a/fastgen/methods/distribution_matching/anyflow_scheduler.py +++ /dev/null @@ -1,118 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 - -"""Lightweight flow-map scheduler for any-step inference. - -Ported from the AnyFlow reference implementation -(``far/schedulers/scheduling_flowmap_euler_discrete.py``) with the -``diffusers.ConfigMixin`` dependency removed so it can be used standalone. - -The scheduler operates on timesteps in ``[0, num_train_timesteps]`` (matching -the Wan T2V/I2V conventions). For pair-step sampling, ``step`` takes both the -current timestep ``t`` and the target ``r`` (with ``r < t``) and integrates the -flow map prediction in one shot: - - x_r = x_t - (t - r) * u_theta(x_t, t, r) - -This is the flow-map analogue of an Euler step where the integration interval -``t - r`` is chosen freely at inference time, enabling any-step sampling. -""" - -from __future__ import annotations - -from typing import Union - -import torch - - -class FlowMapDiscreteScheduler: - """Any-step flow-map scheduler. - - Args: - num_train_timesteps: Maximum timestep value used during training. - shift: Flow-matching schedule shift (Wan default: 5.0 for video, 1.0 for image). - weight_type: Per-timestep loss weighting scheme — ``gaussian``, ``beta08``, - or ``uniform``. ``beta08`` matches AnyFlow's default. - """ - - def __init__( - self, - num_train_timesteps: int = 1000, - shift: float = 1.0, - weight_type: str = "beta08", - ): - self.num_train_timesteps = num_train_timesteps - self.shift = shift - self.weight_type = weight_type - - # Initialise with train-time uniform spacing; overridden by set_timesteps() - self.set_timesteps(num_train_timesteps, device="cpu") - self._build_train_weights() - - def _build_train_weights(self) -> None: - if self.weight_type == "gaussian": - x = self.timesteps - y = torch.exp(-2 * ((x - self.num_train_timesteps / 2) / self.num_train_timesteps) ** 2) - y_shifted = y - y.min() - self.linear_timesteps_weights = y_shifted * (self.num_train_timesteps / y_shifted.sum()) - elif self.weight_type == "beta08": - t = self.timesteps / self.num_train_timesteps - y = (t**1.0) * ((1 - t) ** 0.5) - self.linear_timesteps_weights = y * (self.num_train_timesteps / y.sum()) - elif self.weight_type == "uniform": - self.linear_timesteps_weights = torch.ones_like(self.timesteps) - else: - raise ValueError(f"Invalid weight_type: {self.weight_type!r}") - - @torch.no_grad() - def get_train_weight(self, timesteps: torch.Tensor) -> torch.Tensor: - """Look up the per-timestep training loss weight via nearest neighbour.""" - device_weights = self.linear_timesteps_weights.to(timesteps.device) - device_timesteps = self.timesteps.to(timesteps.device) - diffs = (device_timesteps.unsqueeze(1) - timesteps.flatten().unsqueeze(0)).abs() - timestep_id = torch.argmin(diffs, dim=0).reshape(timesteps.shape) - return device_weights[timestep_id] - - def apply_shift(self, sigmas: torch.Tensor) -> torch.Tensor: - """Apply the flow-matching schedule shift to normalized sigmas in [0, 1].""" - if self.shift == 1.0: - return sigmas - return self.shift * sigmas / (1 + (self.shift - 1) * sigmas) - - def set_timesteps( - self, - num_inference_steps: int, - device: Union[str, torch.device, None] = None, - ) -> None: - timesteps = torch.linspace(1.0, 0.0, num_inference_steps + 1, dtype=torch.float64, device=device) - timesteps = self.apply_shift(timesteps) - self.timesteps = timesteps * self.num_train_timesteps - - def scale_noise( - self, - sample: torch.Tensor, - timestep: Union[float, torch.Tensor], - noise: torch.Tensor, - ) -> torch.Tensor: - """Forward-noise ``sample`` to ``timestep`` along a linear flow-matching path.""" - timestep = torch.as_tensor(timestep, device=sample.device, dtype=sample.dtype) - timestep = timestep / self.num_train_timesteps - timestep = timestep.view(*timestep.shape, *([1] * (noise.ndim - timestep.ndim))) - return timestep * noise + (1.0 - timestep) * sample - - def step( - self, - model_output: torch.Tensor, - sample: torch.Tensor, - timestep: Union[float, torch.Tensor], - r_timestep: Union[float, torch.Tensor], - ) -> torch.Tensor: - """Pair-step Euler integration ``x_r = x_t - (t - r) * model_output``.""" - timestep = torch.as_tensor(timestep, device=sample.device, dtype=sample.dtype) - r_timestep = torch.as_tensor(r_timestep, device=sample.device, dtype=sample.dtype) - timestep = timestep / self.num_train_timesteps - r_timestep = r_timestep / self.num_train_timesteps - timestep = timestep.view(*timestep.shape, *([1] * (model_output.ndim - timestep.ndim))) - r_timestep = r_timestep.view(*r_timestep.shape, *([1] * (model_output.ndim - r_timestep.ndim))) - prev_sample = sample - (timestep - r_timestep) * model_output - return prev_sample.to(model_output.dtype) diff --git a/fastgen/networks/EDM/network.py b/fastgen/networks/EDM/network.py index 0c9a550..331b481 100644 --- a/fastgen/networks/EDM/network.py +++ b/fastgen/networks/EDM/network.py @@ -924,6 +924,12 @@ def forward( else condition.reshape(-1, self.label_dim) ) + # A dual-timestep network always consumes the r-pathway (the embedding + # widths are sized for it), so r=None is not a valid input for it. + # Default to r=t, i.e. the instantaneous velocity u(x_t, t, t). + if r is None and getattr(self.model, "r_timestep", None) is not None: + r = t + # Preconditioning weights for input x_t_in, t_in = x_t, t if self.drop_precond not in ["input", "both"]: diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index 8172e39..869681e 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -295,19 +295,23 @@ def _fuse_r_embedding( ``timestep_proj = time_proj(act_fn(rt_emb))``. This matches the ``WanTwoTimeTextImageEmbedding.forward_timestep`` path in the AnyFlow reference and is required to reproduce the published HF checkpoint - forward bit-for-bit. + forward bit-for-bit. With ``encoder_depth`` set, the first + ``encoder_depth`` blocks keep the t-only ``timestep_proj`` and the + remaining blocks switch to the gated projection (mirroring additive). Returns ``(temb, timestep_proj, r_timestep_proj)`` where ``r_timestep_proj`` - is ``None`` in the gated branch (the t/r mix is folded into ``timestep_proj`` - itself). + is ``None`` when the r-information is already folded into ``timestep_proj``. """ fusion = getattr(self.r_embedder, "fusion_mode", "additive") if fusion == "gated": - gate = self.r_embedder.gate_value.to(remb.dtype) + gate = self.r_embedder.gate_value rt_emb = (1 - gate) * temb + gate * remb - ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(rt_emb)) - return rt_emb, unflatten_timestep_proj(ts_proj, rs_seq_len), None + rt_ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(rt_emb)) + rt_ts_proj = unflatten_timestep_proj(rt_ts_proj, rs_seq_len) + if self.encoder_depth is None: + return rt_emb, rt_ts_proj, None + return rt_emb, timestep_proj, rt_ts_proj # additive — MeanFlow original path, bit-identical r_ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb)) @@ -462,6 +466,50 @@ def classify_forward_block_forward( return hidden_states, features +def remap_anyflow_keys(state_dict: Mapping[str, Any]) -> Mapping[str, Any]: + """Remap an AnyFlow HF release state_dict to FastGen's Wan layout. + + AnyFlow's ``FAR_Wan_Transformer3DModel`` stores the r-pathway inside the + main ``condition_embedder`` as ``delta_embedder``, and uses ONE shared + ``time_proj`` for both t and (t, r). FastGen exposes a separate top-level + ``r_embedder`` with its own ``time_embedder`` + ``time_proj``. The two + layouts are functionally equivalent (FastGen's ``r_embedder.time_proj`` + starts as a deepcopy of ``condition_embedder.time_proj`` per + :meth:`Wan.init_embedder`), so we just rename / duplicate the tensors. + + The function is a no-op when no ``condition_embedder.delta_embedder.*`` + keys are present, so it's safe to call unconditionally. Keys with or + without the ``transformer.`` module prefix are both handled. + """ + delta_marker = "condition_embedder.delta_embedder." + delta_keys = [k for k in state_dict if delta_marker in k] + if not delta_keys: + return state_dict + + new_sd = dict(state_dict) + prefixes = set() + for k in delta_keys: + # [transformer.]condition_embedder.delta_embedder.linear_1.weight + # -> [transformer.]r_embedder.time_embedder.linear_1.weight + prefix, _, suffix = k.partition(delta_marker) + prefixes.add(prefix) + new_sd[f"{prefix}r_embedder.time_embedder.{suffix}"] = new_sd.pop(k) + # AnyFlow's gated fusion shares the final time_proj. FastGen has a + # separate r_embedder.time_proj that substitutes for the shared one when + # fusion="gated"; copy the weights so the two projections start identical. + for prefix in prefixes: + for sub in ("weight", "bias"): + src = f"{prefix}condition_embedder.time_proj.{sub}" + dst = f"{prefix}r_embedder.time_proj.{sub}" + if src in new_sd and dst not in new_sd: + new_sd[dst] = new_sd[src].clone() + logger.info( + f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors " + "and duplicated time_proj weights into r_embedder." + ) + return new_sd + + class WanTextEncoder: def __init__(self, model_id_or_local_path): self.text_encoder = UMT5EncoderModel.from_pretrained( @@ -656,11 +704,9 @@ def __init__( if r_embedder_fusion not in ("additive", "gated"): raise ValueError(f"r_embedder_fusion must be 'additive' or 'gated', got {r_embedder_fusion!r}") self.transformer.r_embedder.fusion_mode = r_embedder_fusion - self.transformer.r_embedder.register_buffer( - "gate_value", - torch.tensor([float(r_embedder_gate_value)], dtype=torch.float32), - persistent=False, - ) + # Plain float (not a buffer) so FSDP's reset_parameters does not + # need to re-materialize it after meta-device init. + self.transformer.r_embedder.gate_value = float(r_embedder_gate_value) logger.info( f"r_embedder fusion={r_embedder_fusion}" + (f" gate_value={r_embedder_gate_value}" if r_embedder_fusion == "gated" else "") @@ -1124,6 +1170,10 @@ def rename_key(k: str) -> str: new_state[new_k] = v state_dict = new_state + # AnyFlow HF releases store the r-pathway as condition_embedder.delta_embedder; + # remap to FastGen's r_embedder layout (no-op for all other checkpoints). + state_dict = remap_anyflow_keys(state_dict) + return super().load_state_dict(state_dict, **kwargs) def forward( @@ -1173,6 +1223,16 @@ def forward( assert fwd_pred_type in NET_PRED_TYPES, f"{fwd_pred_type} is not supported as fwd_pred_type" condition = torch.stack(condition, dim=0) if isinstance(condition, list) else condition + + # A gated-fusion flow-map network always consumes the r-pathway (the + # fusion happens inside the shared time projection), so r=None has no + # in-distribution meaning for it. Default to r=t, i.e. the + # instantaneous velocity u(x_t, t, t) — this is how the AnyFlow + # reference queries its real/fake score models. Additive-fusion + # (MeanFlow) networks keep the skip-r-pathway behavior. + if r is None and getattr(self.transformer.r_embedder, "fusion_mode", None) == "gated": + r = t + timestep_mask = torch.ones_like(x_t[:, 0]) # shape: [batch_size, num_latent_frames, H, W] timestep = self._compute_timestep_inputs(t, timestep_mask) r_timestep = None if r is None else self._compute_timestep_inputs(r, timestep_mask) diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index e27a73a..57ab5c2 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -1,40 +1,54 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for AnyFlowModel. - -The tests run on the tiny EDM backbone with ``r_timestep=True`` (same trick -``test_meanflowmodel.py`` uses) so they execute on CPU without downloading -any pretrained weights. +"""Unit tests for AnyFlow. + +The pretrain stage is MeanFlow with AnyFlow's hyperparameters (fixed +per-timestep loss weighting, finite-difference JVP, consistency bucket), so +the pretrain tests drive :class:`MeanFlowModel` directly. The on-policy tests +exercise :class:`AnyFlowModel` (DMD2 with a multi-step rollout-with-gradient +student). Both run on the tiny EDM backbone with ``r_timestep=True`` and +``schedule_type=rf`` so they execute on CPU without pretrained weights. """ import gc +import types import pytest import torch from fastgen.configs.config_utils import override_config_with_opts -from fastgen.configs.methods.config_anyflow import ModelConfig -from fastgen.methods import AnyFlowModel -from fastgen.methods.distribution_matching.anyflow_scheduler import FlowMapDiscreteScheduler +from fastgen.configs.methods.config_anyflow import ModelConfig as AnyFlowModelConfig +from fastgen.configs.methods.config_mean_flow import ModelConfig as MeanFlowModelConfig +from fastgen.methods import AnyFlowModel, MeanFlowModel from fastgen.utils.test_utils import check_grad_zero -def _build_pretrain_model(): +def _build_pretrain_model(weight_type="beta08", consistency_ratio=0.25, r_sample_ratio=0.5): + """MeanFlow configured the AnyFlow way (paper Stage 2).""" gc.collect() - instance = ModelConfig() - instance.loss_config.training_stage = "pretrain" - # Use a small finite-difference step relative to t in [0, 1]. + instance = MeanFlowModelConfig() + + instance.loss_config.loss_type = "l2" + instance.loss_config.weight_type = weight_type + instance.loss_config.use_jvp_finite_diff = True instance.loss_config.jvp_finite_diff_eps = 1e-2 - opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"] + instance.sample_t_cfg.time_dist_type = "shifted" + instance.sample_t_cfg.shift = 5.0 + instance.sample_t_cfg.min_t = 0.001 + instance.sample_t_cfg.max_t = 0.999 + instance.sample_t_cfg.r_sample_ratio = r_sample_ratio + instance.sample_t_cfg.consistency_ratio = consistency_ratio + + opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True", "+schedule_type=rf"] instance.net = override_config_with_opts(instance.net, opts) instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16" instance.pretrained_model_path = "" instance.input_shape = [3, 2, 2] - model = AnyFlowModel(instance) + model = MeanFlowModel(instance) model.on_train_begin() model.init_optimizers() return model @@ -44,10 +58,9 @@ def _build_onpolicy_model(): """On-policy fixture mirrors test_dmd2model: img_resolution=8 so the discriminator's 4x4 conv kernels can operate.""" gc.collect() - instance = ModelConfig() - instance.loss_config.training_stage = "onpolicy" + instance = AnyFlowModelConfig() - opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"] + opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True", "+schedule_type=rf"] instance.net = override_config_with_opts(instance.net, opts) opts_disc = ["-", "feature_indices=[0]", "all_res=[8]", "in_channels=128"] instance.discriminator = override_config_with_opts(instance.discriminator, opts_disc) @@ -56,9 +69,8 @@ def _build_onpolicy_model(): instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16" instance.pretrained_model_path = "" instance.student_update_freq = 2 - # Exercise the multi-step rollout-with-gradient path (matches the - # AnyFlow paper's on-policy training mode). student_sample_steps=2 is - # the smallest value that runs the rollout loop more than once. + # student_sample_steps=2 is the smallest value that runs the rollout loop + # more than once. instance.student_sample_steps = 2 instance.input_shape = [3, 8, 8] @@ -68,8 +80,7 @@ def _build_onpolicy_model(): return model -def _make_data(model, img_resolution: int = 2): - batch_size = 1 +def _make_data(model, img_resolution: int = 2, batch_size: int = 1): labels = torch.nn.functional.one_hot(torch.randint(0, 10, (batch_size,)), num_classes=10) neg_labels = torch.zeros(batch_size, 10) return { @@ -80,7 +91,7 @@ def _make_data(model, img_resolution: int = 2): # --------------------------------------------------------------------------- -# Pretrain stage +# Pretrain stage — MeanFlow with AnyFlow options # --------------------------------------------------------------------------- @@ -91,19 +102,9 @@ def test_pretrain_single_train_step(): loss_map, outputs = model.single_train_step(data, 0) assert "total_loss" in loss_map - assert "anyflow_loss" in loss_map - assert "dF_dt_target_norm" in loss_map + assert "mf_loss" in loss_map assert torch.isfinite(loss_map["total_loss"]).all() assert "gen_rand" in outputs - assert isinstance(outputs["gen_rand"], torch.Tensor) - - -def test_pretrain_no_fake_score_or_discriminator(): - """Pretrain stage must not instantiate DMD2's fake_score / discriminator.""" - model = _build_pretrain_model() - assert not hasattr(model, "fake_score") or model.fake_score is None or "fake_score" not in model.model_dict - assert "fake_score" not in model.model_dict - assert "discriminator" not in model.model_dict def test_pretrain_optimizer_step(): @@ -119,9 +120,9 @@ def test_pretrain_optimizer_step(): check_grad_zero(model.net) -def test_pretrain_central_difference_falls_back_at_boundaries(): - """When (t ± δ) leaves [min_t, max_t], the one-sided fallback should still - produce a finite target. We synthesise a worst-case t at the boundary.""" +def test_pretrain_finite_difference_falls_back_at_boundaries(): + """When (t ± eps) leaves [min_t, max_t], the one-sided fallback should + still produce a finite JVP estimate. We synthesise worst-case boundary t.""" model = _build_pretrain_model() ns = model.net.noise_scheduler @@ -133,14 +134,53 @@ def test_pretrain_central_difference_falls_back_at_boundaries(): eps_noise = torch.randn_like(real) x_t = ns.forward_process(real, eps_noise, t) - v = eps_noise - real + dxt_dt = eps_noise - real + + u_theta_jvp = model._jvp(x_t, t, r, dxt_dt, condition=cond) + assert torch.isfinite(u_theta_jvp).all(), "boundary samples must yield finite JVP estimates" + + +def test_pretrain_consistency_bucket_pins_r_to_min_t(): + """With consistency_ratio=1.0 (and no flow-matching head), every sample's + r must be pinned to min_t.""" + model = _build_pretrain_model(consistency_ratio=1.0, r_sample_ratio=1.0) + data = _make_data(model, batch_size=4) + + captured = {} + orig = model._compute_mf_loss + + def spy(real_data, t, r, **kwargs): + captured["r"] = r + return orig(real_data=real_data, t=t, r=r, **kwargs) - target = model._compute_central_difference_target(x_t, t, r, v, cond) - assert torch.isfinite(target).all(), "boundary samples must yield finite targets" + model._compute_mf_loss = spy + model.single_train_step(data, 0) + + min_t = float(model.net.noise_scheduler.min_t) + assert torch.allclose(captured["r"].float(), torch.full_like(captured["r"].float(), min_t)) + + +@pytest.mark.parametrize("weight_type", ["beta08", "gaussian", "uniform"]) +def test_timestep_weight_function(weight_type): + """The fixed per-timestep weight is a direct function of t: non-negative, + finite, and normalized to ~unit mean over the shifted timestep grid.""" + model = _build_pretrain_model(weight_type=weight_type) + model._timestep_weight_scale = None # force re-derivation for this weight_type + + t = torch.linspace(0.0, 1.0, 101, dtype=torch.float64) + w = model._get_timestep_weight(t) + assert torch.all(w >= 0), f"{weight_type} weights must be non-negative" + assert torch.isfinite(w).all() + + # The normalization matches the reference: sum over the shifted grid == 1000. + shift = model.sample_t_cfg.shift + grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64) + grid = shift * grid / (1 + (shift - 1) * grid) + assert abs(model._get_timestep_weight(grid).sum().item() - 1000.0) < 1e-6 # --------------------------------------------------------------------------- -# On-policy stage — inherits DMD2's alternating updates +# On-policy stage — DMD2 with the rollout-with-gradient student # --------------------------------------------------------------------------- @@ -170,25 +210,39 @@ def test_onpolicy_fake_score_discriminator_update_step(): def test_onpolicy_rollout_propagates_gradient(): """The multi-step rollout must allow gradient flow on the chosen step. - Mirrors AnyFlow's ``training_rollout`` (pipeline_wan_anyflow.py L370): - one randomly-chosen step in the rollout has gradients enabled; the - remaining steps are no_grad. ``gen_data.requires_grad`` should be True - at the rollout output so the DMD generator update has a valid gradient. + Mirrors AnyFlow's ``training_rollout`` (pipeline_wan_anyflow.py): one + randomly-chosen step in the rollout has gradients enabled; the remaining + steps are no_grad. The rollout output must keep the autograd graph so the + DMD generator update has a valid gradient. """ model = _build_onpolicy_model() real = torch.randn(1, 3, 8, 8, device=model.device, dtype=model.precision) cond = torch.nn.functional.one_hot(torch.tensor([0]), num_classes=10).to(model.device, model.precision) - # Exercise rollout under grad-enabled context (mirrors student update step). - gen = model._rollout_with_gradient(batch_size=real.shape[0], dtype=real.dtype, condition=cond) + + input_student, t_student, _, _ = model._generate_noise_and_time(real) + gen = model.gen_data_from_net(input_student, t_student, condition=cond) + assert tuple(gen.shape) == (1, 3, 8, 8), f"rollout output shape mismatch: {gen.shape}" assert gen.requires_grad, "rollout output must keep autograd graph at the chosen step" - # Trivial scalar loss; backward should succeed without NaN. loss = gen.float().pow(2).mean() loss.backward() grad_seen = any(p.grad is not None and torch.isfinite(p.grad).all() for p in model.net.parameters()) assert grad_seen, "no gradient reached the student network through the rollout" +def test_onpolicy_rollout_no_grad_under_no_grad(): + """Under torch.no_grad (the fake-score update path), the rollout must run + fully gradient-free.""" + model = _build_onpolicy_model() + real = torch.randn(1, 3, 8, 8, device=model.device, dtype=model.precision) + cond = torch.nn.functional.one_hot(torch.tensor([0]), num_classes=10).to(model.device, model.precision) + + input_student, t_student, _, _ = model._generate_noise_and_time(real) + with torch.no_grad(): + gen = model.gen_data_from_net(input_student, t_student, condition=cond) + assert not gen.requires_grad + + def test_onpolicy_optimizer_step(): model = _build_onpolicy_model() data = _make_data(model, img_resolution=8) @@ -200,30 +254,97 @@ def test_onpolicy_optimizer_step(): # --------------------------------------------------------------------------- -# Scheduler +# Wan r-embedder fusion + AnyFlow checkpoint remap (pure functions) # --------------------------------------------------------------------------- -def test_flowmap_scheduler_apply_shift_identity_for_shift1(): - scheduler = FlowMapDiscreteScheduler(num_train_timesteps=1000, shift=1.0, weight_type="beta08") - sigmas = torch.linspace(0.0, 1.0, 11) - assert torch.allclose(scheduler.apply_shift(sigmas), sigmas) +def _make_fake_fusion_self(fusion_mode, gate_value=0.25, encoder_depth=None, dim=4, proj_dim=12): + r_embedder = torch.nn.Module() + r_embedder.time_proj = torch.nn.Linear(dim, proj_dim) + r_embedder.act_fn = torch.nn.SiLU() + r_embedder.fusion_mode = fusion_mode + r_embedder.gate_value = gate_value + return types.SimpleNamespace(r_embedder=r_embedder, encoder_depth=encoder_depth) -def test_flowmap_scheduler_step_zero_interval(): - scheduler = FlowMapDiscreteScheduler(num_train_timesteps=1000, shift=1.0, weight_type="uniform") - sample = torch.randn(2, 4, 8, 8) - model_output = torch.randn_like(sample) - t = torch.tensor([500.0, 500.0]) - # r = t => zero-length integration interval; the sample should be unchanged. - out = scheduler.step(model_output, sample, timestep=t, r_timestep=t.clone()) - assert torch.allclose(out, sample, atol=1e-5) +def test_fuse_r_embedding_gated(): + from fastgen.networks.Wan.network import _fuse_r_embedding + torch.manual_seed(0) + fake = _make_fake_fusion_self("gated") + temb = torch.randn(2, 4) + remb = torch.randn(2, 4) + timestep_proj = torch.randn(2, 6, 2) -@pytest.mark.parametrize("weight_type", ["gaussian", "beta08", "uniform"]) -def test_flowmap_scheduler_weights_positive(weight_type): - scheduler = FlowMapDiscreteScheduler(num_train_timesteps=1000, shift=1.0, weight_type=weight_type) - t = torch.tensor([100.0, 500.0, 900.0]) - w = scheduler.get_train_weight(t) - assert torch.all(w >= 0), f"{weight_type} weights must be non-negative" - assert torch.isfinite(w).all() + out_temb, out_proj, out_r_proj = _fuse_r_embedding(fake, temb, timestep_proj, remb, None) + + gate = fake.r_embedder.gate_value + rt_emb = (1 - gate) * temb + gate * remb + expected = fake.r_embedder.time_proj(fake.r_embedder.act_fn(rt_emb)).unflatten(1, (6, -1)) + assert torch.allclose(out_temb, rt_emb) + assert torch.allclose(out_proj, expected) + assert out_r_proj is None + + +def test_fuse_r_embedding_gated_respects_encoder_depth(): + from fastgen.networks.Wan.network import _fuse_r_embedding + + torch.manual_seed(0) + fake = _make_fake_fusion_self("gated", encoder_depth=2) + temb = torch.randn(2, 4) + remb = torch.randn(2, 4) + timestep_proj = torch.randn(2, 6, 2) + + out_temb, out_proj, out_r_proj = _fuse_r_embedding(fake, temb, timestep_proj, remb, None) + + # Encoder blocks keep the t-only projection; the gated projection is + # returned separately so the block loop switches at encoder_depth. + assert torch.allclose(out_proj, timestep_proj) + assert out_r_proj is not None and out_r_proj.shape == timestep_proj.shape + gate = fake.r_embedder.gate_value + assert torch.allclose(out_temb, (1 - gate) * temb + gate * remb) + + +def test_fuse_r_embedding_additive_unchanged(): + from fastgen.networks.Wan.network import _fuse_r_embedding + + torch.manual_seed(0) + fake = _make_fake_fusion_self("additive") + temb = torch.randn(2, 4) + remb = torch.randn(2, 4) + timestep_proj = torch.randn(2, 6, 2) + + out_temb, out_proj, out_r_proj = _fuse_r_embedding(fake, temb, timestep_proj, remb, None) + + r_proj = fake.r_embedder.time_proj(fake.r_embedder.act_fn(remb)).unflatten(1, (6, -1)) + assert torch.allclose(out_temb, temb + remb) + assert torch.allclose(out_proj, timestep_proj + r_proj) + assert torch.allclose(out_r_proj, r_proj) + + +@pytest.mark.parametrize("prefix", ["", "transformer."]) +def test_remap_anyflow_keys(prefix): + from fastgen.networks.Wan.network import remap_anyflow_keys + + sd = { + f"{prefix}condition_embedder.delta_embedder.linear_1.weight": torch.randn(2, 2), + f"{prefix}condition_embedder.time_proj.weight": torch.randn(2, 2), + f"{prefix}condition_embedder.time_proj.bias": torch.randn(2), + f"{prefix}blocks.0.attn1.to_q.weight": torch.randn(2, 2), + } + out = remap_anyflow_keys(sd) + + assert f"{prefix}r_embedder.time_embedder.linear_1.weight" in out + assert f"{prefix}condition_embedder.delta_embedder.linear_1.weight" not in out + # The shared time_proj is duplicated into the r_embedder. + assert torch.equal(out[f"{prefix}r_embedder.time_proj.weight"], sd[f"{prefix}condition_embedder.time_proj.weight"]) + assert torch.equal(out[f"{prefix}r_embedder.time_proj.bias"], sd[f"{prefix}condition_embedder.time_proj.bias"]) + # Unrelated keys untouched. + assert torch.equal(out[f"{prefix}blocks.0.attn1.to_q.weight"], sd[f"{prefix}blocks.0.attn1.to_q.weight"]) + + +def test_remap_anyflow_keys_noop_without_delta_keys(): + from fastgen.networks.Wan.network import remap_anyflow_keys + + sd = {"transformer.blocks.0.attn1.to_q.weight": torch.randn(2, 2)} + assert remap_anyflow_keys(sd) is sd From 667c47dbcf5a376e8bf628e4412f1ba24ec34b96 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Wed, 10 Jun 2026 23:36:55 +0800 Subject: [PATCH 06/19] anyflow: align training math with the reference implementation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adversarial review against the reference surfaced several silent deviations in the training objective; all fixed: Pretrain (MeanFlow, all opt-in via config): - rebalance_to_diffusion: non-diffusion (flow-map / consistency) sample losses are rescaled by the detached factor mean(global diffusion losses) / (own loss + 1e-5), all-gathered across ranks, matching the reference's scale_weight. - guidance_fuse_scale: prediction-side guidance distillation — the conditional output learns the guided flow via (u_cond + (g-1) u_uncond) / g against the raw data velocity, with the unconditional branch queried at the SAME (t, r) slice, the finite-difference dF/dt divided by g, plain text dropout, and the probes extrapolated along the raw velocity. MeanFlow's target-side eq. 19 fusion (guidance_scale) is a different mechanism and stays untouched. - consistency bucket pins r = 0 (not min_t) and, together with the flow-matching head, uses the reference's deterministic global-batch partition by rank index instead of independent binomial draws (which biased the effective ratios at small per-GPU batches). - weight_type=uniform is exactly 1 (the reference applies no grid normalization to it). On-policy: - rollout NFE sampled per iteration from student_sample_steps_list ([2, 4, 8, 16, 50]) with rank-0 broadcast; schedule computed from the shifted grid per NFE. - rollout compressed to <= 3 flow-map forwards (jump t0->tg, fine step tg->tg+1, jump to 0) with gradient through ALL segments — the previous N-step loop with gradient on one step did not match training_rollout. - every student update co-trains the Stage-2 flow-map loss on the real batch (cotrain_pretrain_weight, reference cotrain_forward_kl). - no adversarial loss: the reference 'discriminator' is the fake score network; gan_loss_weight_gen=0. - student/fake-score updates alternate 1:1 (student_update_freq=2), teacher CFG strength corrected to the reference's cond + 3*(cond-uncond) (FastGen formula: guidance_scale=4), optimizer betas (0.0, 0.999), wd=0, grad clip 1.0, EMA 0.99. Configs also align the pretrain recipe (grad clip 1.0, 1000-step LR warmup, EMA 0.999, exact shifted 4-step eval schedule). --- .../experiments/WanT2V/config_anyflow.py | 22 +- .../WanT2V/config_anyflow_onpolicy.py | 83 ++++-- fastgen/configs/methods/config_anyflow.py | 36 ++- fastgen/configs/methods/config_mean_flow.py | 9 + .../methods/consistency_model/mean_flow.py | 175 ++++++++++--- .../methods/distribution_matching/README.md | 14 +- .../methods/distribution_matching/anyflow.py | 237 +++++++++++++----- tests/test_anyflowmodel.py | 126 ++++++++-- 8 files changed, 543 insertions(+), 159 deletions(-) diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py index e95027c..e64bfab 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -48,11 +48,17 @@ def create_config(): config.model.loss_config.weight_type = "beta08" config.model.loss_config.use_jvp_finite_diff = True config.model.loss_config.jvp_finite_diff_eps = 5e-3 + # Rebalance the non-diffusion (flow-map / consistency) sample losses to + # the global diffusion-loss mean (reference scale_weight). + config.model.loss_config.rebalance_to_diffusion = True config.model.precision_amp_jvp = "float32" - # Velocity-level guidance fusion with text dropout (reference: - # drop_text_ratio=0.1, fuse_guidance_scale=3.0). - config.model.guidance_scale = 3.0 + # Prediction-side guidance fusion with text dropout (reference: + # drop_text_ratio=0.1, fuse_guidance_scale=3.0): the conditional output + # learns the guided flow directly. guidance_scale stays None — MeanFlow's + # target-side eq. 19 fusion is a different mechanism. + config.model.guidance_scale = None + config.model.loss_config.guidance_fuse_scale = 3.0 config.model.cond_dropout_prob = 0.1 # ------ (t, r) sampling: shifted uniform pairs + AnyFlow buckets ------ @@ -66,16 +72,22 @@ def create_config(): # consistency_ratio=0.25 of the batch is pinned to r = min_t. config.model.sample_t_cfg.consistency_ratio = 0.25 - # ------ optimization (reference: AdamW lr=5e-5, wd=0, betas=(0.9, 0.95)) ------ + # ------ optimization (reference: AdamW lr=5e-5, wd=0, betas=(0.9, 0.95), + # max_grad_norm=1.0, 1000-step LR warmup, EMA decay 0.999) ------ config.model.net_optimizer.optim_type = "adamw" config.model.net_optimizer.lr = 5e-5 config.model.net_optimizer.betas = (0.9, 0.95) config.model.net_optimizer.weight_decay = 0.0 + config.model.net_scheduler.warm_up_steps = [1000] + config.trainer.callbacks.grad_clip.grad_norm = 1.0 + config.trainer.callbacks.ema.beta = 0.999 # ------ inference / validation ------ config.model.student_sample_type = "ode" config.model.student_sample_steps = 4 - config.model.sample_t_cfg.t_list = [0.999, 0.937, 0.833, 0.624, 0.0] + # 4-step shifted schedule shift*s/(1+(shift-1)*s) on s=linspace(1,0,5), + # first entry clamped to max_t (the reference samples from t=1.0). + config.model.sample_t_cfg.t_list = [0.999, 0.9375, 0.83333333, 0.625, 0.0] # ------ data / trainer ------ config.dataloader_train = VideoLoaderConfig diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index f5c5b9f..c16960e 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -3,23 +3,34 @@ """AnyFlow on-policy distillation config on Wan-1.3B T2V (paper Stage 3). -DMD2's alternating fake_score / discriminator updates with a flow-map student -that generates via a multi-step rollout-with-gradient (see -:class:`~fastgen.methods.distribution_matching.anyflow.AnyFlowModel`). -Mirrors the paper's Stage 3 hyperparameters: 1.2k iterations at lr=2e-6 on -top of a Stage 2 flow-map pretrain checkpoint (``config_anyflow.py``). - -Note: the AnyFlow paper trains this stage with a rank-256 LoRA adapter, but -FastGen does not ship a PEFT/LoRA training path today, so this config does -full-rank fine-tuning. Set ``config.model.pretrained_student_net_path`` to a -checkpoint produced by :mod:`config_anyflow` before launching. +DMD2-style distribution matching with a flow-map student that generates via +a compressed rollout (jump -> fine step -> jump, gradient through all +segments, NFE sampled per iteration from [2, 4, 8, 16, 50]) and co-trains the +Stage-2 flow-map loss at every student update — see +:class:`~fastgen.methods.distribution_matching.anyflow.AnyFlowModel`. + +Mirrors the reference recipe +(``train_wan1b_onpolicy_81f_480p_lr2e-6_1k_b32.yml``): 1.2k iterations at +lr=2e-6, AdamW betas=(0.0, 0.999), wd=0, grad clip 1.0, EMA 0.99, +real-score CFG strength 3, no adversarial loss (the reference +"discriminator" is the fake score network). Set +``config.model.pretrained_student_net_path`` to a Stage 2 checkpoint from +``config_anyflow.py`` before launching, and point the teacher +(``config.model.pretrained_model_path``) at a flow-map teacher checkpoint — +the reference initializes both real and fake score from a separately +fine-tuned flow-map teacher, not stock Wan2.1. + +Known deviations from the reference, both documented: +* full-rank fine-tuning instead of the paper's rank-256 LoRA (FastGen has no + PEFT path today); +* the fake-score noising time is shifted-uniform (FastGen's sampler) instead + of shifted-logit-normal(0, 1). """ import copy import fastgen.configs.methods.config_anyflow as config_anyflow_default from fastgen.configs.data import VideoLoaderConfig -from fastgen.configs.discriminator import Discriminator_Wan_1_3B_Config from fastgen.configs.net import Wan_1_3B_Config @@ -40,28 +51,52 @@ def create_config(): # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. config.model.input_shape = [16, 21, 60, 104] - # ------ DMD2 machinery ------ - config.model.discriminator = Discriminator_Wan_1_3B_Config - config.model.discriminator.disc_type = "multiscale_down_mlp_large" - config.model.discriminator.feature_indices = [15, 22, 29] - config.model.gan_loss_weight_gen = 0.03 - config.model.student_update_freq = 5 - config.model.guidance_scale = 3.0 + # ------ DMD machinery ------ + # The reference Stage 3 has no adversarial loss: its "discriminator" is + # the fake score network, trained with denoising score matching only. + config.model.gan_loss_weight_gen = 0.0 + # The reference updates generator and fake score 1:1 (both every global + # step); in DMD2's alternating scheme that is student_update_freq=2. + config.model.student_update_freq = 2 + # Reference real_guidance_scale=3.0 applies cond + 3*(cond - uncond); + # FastGen's CFG formula is cond + (g-1)*(cond - uncond), so g=4. + config.model.guidance_scale = 4.0 config.model.sample_t_cfg.time_dist_type = "shifted" config.model.sample_t_cfg.shift = 5.0 config.model.sample_t_cfg.min_t = 0.001 config.model.sample_t_cfg.max_t = 0.999 - # ------ student rollout (AnyFlow's hand-tuned 4-step schedule) ------ + # ------ student rollout (reference rollout_cfg) ------ config.model.student_sample_type = "ode" - config.model.student_sample_steps = 4 - config.model.sample_t_cfg.t_list = [0.999, 0.937, 0.833, 0.624, 0.0] - - # ------ Stage 3 learning rates from the AnyFlow paper ------ + config.model.student_sample_steps = 4 # used for validation sampling + config.model.student_sample_steps_list = [2, 4, 8, 16, 50] + config.model.sample_t_cfg.t_list = None # rollout schedules are computed per NFE + + # ------ co-trained Stage-2 flow-map loss (reference cotrain_forward_kl) ------ + config.model.cotrain_pretrain_weight = 1.0 + config.model.loss_config.use_cd = False + config.model.loss_config.loss_type = "l2" + config.model.loss_config.weight_type = "beta08" + config.model.loss_config.use_jvp_finite_diff = True + config.model.loss_config.jvp_finite_diff_eps = 5e-3 + config.model.loss_config.rebalance_to_diffusion = True + config.model.loss_config.guidance_fuse_scale = 3.0 + config.model.cond_dropout_prob = 0.1 + config.model.precision_amp_jvp = "float32" + config.model.sample_t_cfg.r_sample_ratio = 0.5 + config.model.sample_t_cfg.consistency_ratio = 0.25 + + # ------ optimization (reference: AdamW lr=2e-6, betas=(0.0, 0.999), wd=0, + # grad clip 1.0, EMA 0.99) ------ config.model.net_optimizer.lr = 2e-6 + config.model.net_optimizer.betas = (0.0, 0.999) + config.model.net_optimizer.weight_decay = 0.0 config.model.fake_score_optimizer.lr = 2e-6 - config.model.discriminator_optimizer.lr = 2e-6 + config.model.fake_score_optimizer.betas = (0.0, 0.999) + config.model.fake_score_optimizer.weight_decay = 0.0 + config.trainer.callbacks.grad_clip.grad_norm = 1.0 + config.trainer.callbacks.ema.beta = 0.99 # ------ data / trainer ------ config.dataloader_train = VideoLoaderConfig diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py index bf6b6d2..ad299d0 100644 --- a/fastgen/configs/methods/config_anyflow.py +++ b/fastgen/configs/methods/config_anyflow.py @@ -9,6 +9,8 @@ AnyFlow's hyperparameters — see ``configs/experiments/WanT2V/config_anyflow.py``. """ +from typing import List, Optional + import attrs from omegaconf import DictConfig @@ -22,13 +24,45 @@ ) from fastgen.configs.config import BaseConfig from fastgen.configs.methods.config_dmd2 import ModelConfig as DMD2ModelConfig +from fastgen.configs.methods.config_mean_flow import ( + LossConfig as MeanFlowLossConfig, + SampleRConfig, + SampleTConfig as MeanFlowSampleTConfig, +) from fastgen.methods import AnyFlowModel from fastgen.utils import LazyCall as L @attrs.define(slots=False) class ModelConfig(DMD2ModelConfig): - """AnyFlow on-policy model config — identical to DMD2's.""" + """AnyFlow on-policy model config — DMD2 plus the rollout / cotrain knobs. + + The MeanFlow loss / sampling configs drive the co-trained Stage-2 + flow-map loss inside the student update (the reference's + ``cotrain_forward_kl``). + """ + + # MeanFlow-style (t, r) sampling for the co-trained flow-map loss; the + # extra fields are ignored by the DMD2 noising-time sampling. + sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) + sample_r_cfg: SampleRConfig = attrs.field(factory=SampleRConfig) + loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig) + + # Weight of the co-trained Stage-2 flow-map loss in the student update. + # The reference runs it at weight 1 (cotrain_forward_kl: True); 0 disables. + cotrain_pretrain_weight: float = 1.0 + + # Rollout NFE list, sampled uniformly per iteration with rank-0 broadcast + # (reference rollout_cfg.num_inference_steps_list). None falls back to + # the fixed student_sample_steps. + student_sample_steps_list: Optional[List[int]] = None + + # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). + cond_dropout_prob: Optional[float] = None + cond_keys_no_dropout: List[str] = [] + + # Precision for autocast in the co-trained loss JVP (None = training precision). + precision_amp_jvp: str | None = None @attrs.define(slots=False) diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index ff8fda4..1cc1fa4 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -79,6 +79,15 @@ class LossConfig: # ("beta08", "gaussian", "uniform"; used by AnyFlow). None keeps the # adaptive norm_method weighting above. Only applies to loss_type="l2". weight_type: Optional[str] = None + # prediction-side guidance fusion scale (AnyFlow guidance distillation): + # the conditional output is trained to be the guided flow directly via + # (u_cond + (g-1) * u_uncond) / g against the raw data velocity, with + # plain text dropout (cond_dropout_prob). None keeps MeanFlow's + # target-side guidance (guidance_scale, eq. 19). + guidance_fuse_scale: Optional[float] = None + # rebalance non-diffusion (r < t) sample losses to the global + # diffusion-loss mean via a detached per-sample factor (AnyFlow) + rebalance_to_diffusion: bool = False @attrs.define(slots=False) diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 20a2ac5..8f96ed1 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -104,6 +104,25 @@ def _mix_condition( return condition, dxt_dt + def _drop_condition(self, condition: Any, neg_condition: Any) -> Any: + """Replace the condition with neg_condition for a random per-sample subset. + + Plain text dropout (used by the prediction-side guidance fusion); the + dropped samples then train the unconditional pathway. + """ + if self.config.cond_dropout_prob is None or neg_condition is None: + return condition + ref = neg_condition if isinstance(neg_condition, torch.Tensor) else next(iter(neg_condition.values())) + keep = torch.rand(ref.shape[0], device=ref.device) >= self.config.cond_dropout_prob + if isinstance(condition, torch.Tensor): + return torch.where(expand_like(keep, condition), condition, neg_condition) + if isinstance(condition, dict): + condition = condition.copy() + for k in condition.keys() - set(self.config.cond_keys_no_dropout): + condition[k] = torch.where(expand_like(keep, condition[k]), condition[k], neg_condition[k]) + return condition + raise TypeError(f"Unsupported condition type: {type(condition)}") + @torch.no_grad() def _get_velocity( self, @@ -255,6 +274,78 @@ def net_wrapper(x_t, t, r): return u_theta_jvp + def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Sample (t, r) with t >= r and assign the per-batch buckets. + + Returns ``(t, r, is_diffusion)`` where ``is_diffusion`` marks the + samples with ``r = t`` (pure flow matching). + + With ``consistency_ratio == 0`` (MeanFlow default) the flow-matching + head size is stochastic, as before. With ``consistency_ratio > 0`` + (AnyFlow) the buckets are assigned deterministically over the GLOBAL + batch by rank index, exactly like the AnyFlow reference + (``sample_timestep`` in ``trainer_wan_anyflow_pretrain.py``): the + first ``round((1 - r_sample_ratio) * global_bsz)`` samples get + ``r = t``, the next ``round(consistency_ratio * global_bsz)`` get + ``r = 0`` (consistency to clean data), and the rest keep the sampled + random pair. + """ + t_sample_kwargs = convert_cfg_to_dict(self.sample_t_cfg) + t = self.net.noise_scheduler.sample_t(batch_size, **t_sample_kwargs, device=self.device) + r_sample_kwargs = convert_cfg_to_dict(self.sample_r_cfg) if self.sample_r_cfg.enabled else t_sample_kwargs + r = self.net.noise_scheduler.sample_t(batch_size, **r_sample_kwargs, device=self.device) + t, r = torch.maximum(t, r), torch.minimum(t, r) + assert torch.all(t >= r), "r cannot be larger than t" + + if self.sample_t_cfg.consistency_ratio > 0: + if torch.distributed.is_available() and torch.distributed.is_initialized(): + world_size, rank = torch.distributed.get_world_size(), torch.distributed.get_rank() + else: + world_size, rank = 1, 0 + global_bsz = world_size * batch_size + n_diffusion = round((1.0 - self.sample_t_cfg.r_sample_ratio) * global_bsz) + n_consistency = round(self.sample_t_cfg.consistency_ratio * global_bsz) + global_idx = rank * batch_size + torch.arange(batch_size, device=self.device) + is_diffusion = global_idx < n_diffusion + is_consistency = (global_idx >= n_diffusion) & (global_idx < n_diffusion + n_consistency) + r = torch.where(is_diffusion, t, r) + r = torch.where(is_consistency, torch.zeros_like(r), r) + else: + # set t=r (flow matching loss) for a subset of the batch + flow_matching_size = (torch.rand(batch_size, device=self.device) >= self.sample_t_cfg.r_sample_ratio).sum() + is_diffusion = torch.arange(batch_size, device=self.device) < flow_matching_size + r = torch.where(is_diffusion, t, r) + + return t, r, is_diffusion + + def _reduce_mf_loss(self, mf_loss: torch.Tensor, is_diffusion: torch.Tensor) -> torch.Tensor: + """Reduce the per-sample loss to a scalar, optionally rebalancing. + + With ``loss_config.rebalance_to_diffusion`` set (AnyFlow), every + non-diffusion sample's loss is multiplied by the detached factor + ``mean(global diffusion losses) / (own loss + 1e-5)``, matching the + AnyFlow reference (``train_bidirection``): flow-map / consistency + gradients are self-normalized and rescaled to the global + diffusion-loss mean. + """ + if getattr(self.loss_config, "rebalance_to_diffusion", False) and (~is_diffusion).any(): + with torch.no_grad(): + if torch.distributed.is_available() and torch.distributed.is_initialized(): + world_size = torch.distributed.get_world_size() + gathered_loss = [torch.zeros_like(mf_loss) for _ in range(world_size)] + gathered_mask = [torch.zeros_like(is_diffusion) for _ in range(world_size)] + torch.distributed.all_gather(gathered_loss, mf_loss.contiguous()) + torch.distributed.all_gather(gathered_mask, is_diffusion.contiguous()) + global_loss = torch.cat(gathered_loss, dim=0) + global_mask = torch.cat(gathered_mask, dim=0) + else: + global_loss, global_mask = mf_loss, is_diffusion + scale = torch.ones_like(mf_loss) + if global_mask.any(): + scale[~is_diffusion] = global_loss[global_mask].mean() / (mf_loss[~is_diffusion] + 1e-5) + mf_loss = mf_loss * scale + return mf_loss.mean() + def _timestep_weight_raw(self, t: torch.Tensor) -> torch.Tensor: """Unnormalized per-timestep weight as a direct function of t in [0, 1].""" weight_type = self.loss_config.weight_type @@ -276,6 +367,9 @@ def _get_timestep_weight(self, t: torch.Tensor) -> torch.Tensor: weights over its shifted discrete timestep grid; we compute the same constant once and cache it. """ + if self.loss_config.weight_type == "uniform": + # The reference uses exactly 1.0 for uniform (no grid normalization). + return torch.ones_like(t) if self._timestep_weight_scale is None: shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64) @@ -440,21 +534,47 @@ def _compute_mf_loss( z = torch.randn_like(real_data) x_t = self.net.noise_scheduler.forward_process(real_data, z, t) - condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) - # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity - torch.clear_autocast_cache() - u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) - assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" - - # additional forward pass to get u_theta with gradient; see also https://github.com/Gsunshine/py-meanflow?tab=readme-ov-file#note-on-jvp - assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" - u_theta = self.net( - x_t, - t, - r=r, - condition=condition, - fwd_pred_type="flow", - ) + guidance_fuse_scale = getattr(self.loss_config, "guidance_fuse_scale", None) + if guidance_fuse_scale is not None: + # AnyFlow-style guidance distillation: the guidance is fused on the + # PREDICTION side — the conditional output is trained to be the + # guided flow directly. Matches the reference ``train_bidirection``: + # the regression target stays the raw data velocity, the + # unconditional branch is queried at the SAME (t, r) flow-map + # slice, and the effective prediction is + # (u_cond + (g - 1) * u_uncond) / g. Text dropout replaces the + # condition with neg_condition for a random subset beforehand. + g = float(guidance_fuse_scale) + condition = self._drop_condition(condition, neg_condition) + dxt_dt = self.net.noise_scheduler.cond_velocity(x=real_data, eps=z, t=t) + torch.clear_autocast_cache() + # dF/dt of the fused prediction: finite difference of the + # conditional output divided by g (the unconditional derivative is + # dropped, as in the reference's compute_central_difference). + u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) / g + assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" + + assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" + u_theta = self.net(x_t, t, r=r, condition=condition, fwd_pred_type="flow") + with torch.no_grad(): + u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow") + u_theta = (u_theta + (g - 1.0) * u_uncond) / g + else: + condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) + # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity + torch.clear_autocast_cache() + u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) + assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" + + # additional forward pass to get u_theta with gradient; see also https://github.com/Gsunshine/py-meanflow?tab=readme-ov-file#note-on-jvp + assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" + u_theta = self.net( + x_t, + t, + r=r, + condition=condition, + fwd_pred_type="flow", + ) mf_loss, tangent, loss_weight, warmup_weight = self._mf_pred_to_loss( u_theta=u_theta, u_theta_jvp=u_theta_jvp, x_t=x_t, dxt_dt=dxt_dt, t=t, r=r, iteration=iteration ) @@ -492,27 +612,8 @@ def single_train_step( real_data, condition, neg_condition = self._prepare_training_data(data) batch_size = real_data.shape[0] - # sample t and r - t_sample_kwargs = convert_cfg_to_dict(self.sample_t_cfg) - t = self.net.noise_scheduler.sample_t(batch_size, **t_sample_kwargs, device=self.device) - r_sample_kwargs = convert_cfg_to_dict(self.sample_r_cfg) if self.sample_r_cfg.enabled else t_sample_kwargs - r = self.net.noise_scheduler.sample_t(batch_size, **r_sample_kwargs, device=self.device) - t, r = torch.maximum(t, r), torch.minimum(t, r) - assert torch.all(t >= r), "r cannot be larger than t" - - # set t=r (flow matching loss) for a subset of the batch - batch_size = real_data.shape[0] - flow_matching_size = (torch.rand(batch_size, device=self.device) >= self.sample_t_cfg.r_sample_ratio).sum() - zero_mask = torch.arange(batch_size, device=self.device) < flow_matching_size - r = torch.where(zero_mask, t, r) - - # set r=min_t (consistency-to-clean, AnyFlow) for a subset of the batch. - # Taken from the tail so it never overlaps the flow-matching head. - if self.sample_t_cfg.consistency_ratio > 0: - num_consistency = (torch.rand(batch_size, device=self.device) < self.sample_t_cfg.consistency_ratio).sum() - num_consistency = torch.minimum(num_consistency, batch_size - flow_matching_size) - consistency_mask = torch.arange(batch_size, device=self.device) >= batch_size - num_consistency - r = torch.where(consistency_mask, torch.full_like(r, self.net.noise_scheduler.min_t), r) + # sample t and r, with per-batch buckets (flow matching / consistency) + t, r, is_diffusion = self._sample_t_r_buckets(batch_size) ( mf_loss, @@ -532,7 +633,7 @@ def single_train_step( neg_condition=neg_condition, ) - loss = mf_loss.mean() + loss = self._reduce_mf_loss(mf_loss, is_diffusion) loss_map = { "total_loss": loss, "mf_loss": loss, diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index 4fa411e..b3db151 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -91,20 +91,24 @@ DMD2 extended for causal video generation with autoregressive chunk-by-chunk pro Single model that supports arbitrary inference NFE by learning a flow map `u_θ(x_t, t, r)` (average velocity from `t` back to `r`). Trained in two stages: -1. **Pretrain** — the MeanFlow objective with AnyFlow's hyperparameters, run directly on [`MeanFlowModel`](../consistency_model/mean_flow.py): finite-difference JVP, a fixed `beta08` per-timestep loss weight (`loss_config.weight_type`), shifted timestep sampling, and a `sample_t_cfg.consistency_ratio` fraction of the batch pinned to `r = min_t`. There is no AnyFlow-specific pretrain code. -2. **On-policy** — distribution-matching distillation on top of the pretrained flow-map weights ([`anyflow.py`](anyflow.py)). Stock DMD2 with two overrides: the student generates via a multi-step Euler-flow rollout from pure noise with gradient enabled at one randomly-chosen step (`gen_data_from_net`), and always starts from `max_t` (`_generate_noise_and_time`). The teacher / fake_score are flow-map networks queried at the instantaneous velocity `r = t` (the network-level default for gated-fusion Wan when `r` is not passed). +1. **Pretrain** — the MeanFlow objective with AnyFlow's hyperparameters, run directly on [`MeanFlowModel`](../consistency_model/mean_flow.py): finite-difference JVP, a fixed `beta08` per-timestep loss weight (`loss_config.weight_type`), prediction-side guidance fusion (`loss_config.guidance_fuse_scale` — the conditional output learns the guided flow directly), global rebalancing of the flow-map / consistency sample losses to the diffusion-loss mean (`loss_config.rebalance_to_diffusion`), shifted timestep sampling, and a `sample_t_cfg.consistency_ratio` fraction of the batch pinned to `r = 0`. There is no AnyFlow-specific pretrain code. +2. **On-policy** — distribution-matching distillation on top of the pretrained flow-map weights ([`anyflow.py`](anyflow.py)). DMD2 with the reference's three deviations: the student generates via a flow-map rollout compressed to at most three forwards — jump `t_0 → t_g`, fine step `t_g → t_{g+1}`, jump to 0 — with the NFE sampled per iteration from `student_sample_steps_list` and gradient through all segments (`gen_data_from_net`); the student always starts from pure noise at `max_t` (`_generate_noise_and_time`); and every student update co-trains the Stage-2 flow-map loss (`cotrain_pretrain_weight`). There is no adversarial loss (the reference "discriminator" is the fake score network). The teacher / fake_score are flow-map networks queried at the instantaneous velocity `r = t` (the network-level default for gated-fusion Wan when `r` is not passed). **Key Parameters (on-policy):** -- `student_sample_steps` / `sample_t_cfg.t_list`: rollout length and hand-tuned step schedule +- `student_sample_steps_list`: rollout NFEs sampled per iteration (reference: `[2, 4, 8, 16, 50]`) +- `cotrain_pretrain_weight`: weight of the co-trained Stage-2 flow-map loss (reference: 1) +- `guidance_scale`: real-score CFG; FastGen applies `cond + (g-1)·(cond - uncond)`, so the reference's strength 3 is `g=4` - See also key parameters of DMD2 above **Key Parameters (pretrain, on MeanFlow):** - `loss_config.weight_type`: `beta08` | `gaussian` | `uniform` fixed per-timestep loss weight +- `loss_config.guidance_fuse_scale` + `cond_dropout_prob`: prediction-side guidance distillation with text dropout +- `loss_config.rebalance_to_diffusion`: rescale non-diffusion losses to the global diffusion-loss mean - `loss_config.use_jvp_finite_diff` / `loss_config.jvp_finite_diff_eps`: central-difference step δ -- `sample_t_cfg.r_sample_ratio` / `sample_t_cfg.consistency_ratio`: per-batch fraction with sampled `r` / `r = min_t` +- `sample_t_cfg.r_sample_ratio` / `sample_t_cfg.consistency_ratio`: per-batch fraction with sampled `r` / `r = 0` - `sample_t_cfg.time_dist_type="shifted"`, `sample_t_cfg.shift`: shifted timestep sampling (5.0 for Wan video) -**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). When loading the published AnyFlow HF checkpoints, set `r_embedder_fusion="gated"` and `time_cond_type="abs"` on the Wan constructor — this routes the t/r mix through `Wan/network.py::_fuse_r_embedding`'s gated branch (shared with MeanFlow's additive default) so the released weights reproduce bit-for-bit. `Wan.load_state_dict` remaps the AnyFlow checkpoint layout (`condition_embedder.delta_embedder.*`) automatically. +**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). When loading the published AnyFlow HF checkpoints, set `r_embedder_fusion="gated"` and `time_cond_type="abs"` on the Wan constructor — this routes the t/r mix through `Wan/network.py::_fuse_r_embedding`'s gated branch (shared with MeanFlow's additive default) so the released weights reproduce bit-for-bit. `Wan.load_state_dict` remaps the AnyFlow checkpoint layout (`condition_embedder.delta_embedder.*`) automatically; note that loading the HF folder via `model_id_or_local_path` goes through diffusers' `from_pretrained`, which silently drops the `delta_embedder` weights — load the transformer state dict through `Wan.load_state_dict` (or call `remap_anyflow_keys` yourself) instead. **Note:** correctness of the port is established via forward-parity and single-step training-step parity against the AnyFlow reference (see PR #25 discussion). End-to-end convergence-scale validation on the paper's training corpus is deferred to a follow-up. diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index 2d8629f..a7cd69b 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -13,29 +13,35 @@ * **Flow-map pretrain** (paper Stage 2) is MeanFlow with AnyFlow's hyperparameters — fixed ``beta08`` per-timestep loss weighting, - finite-difference JVP, a ``consistency_ratio`` fraction of the batch pinned - to ``r = min_t``, and shifted timestep sampling. It is configured directly - on :class:`~fastgen.methods.consistency_model.mean_flow.MeanFlowModel` + finite-difference JVP, prediction-side guidance fusion + (``guidance_fuse_scale``), global rebalancing of non-diffusion losses + (``rebalance_to_diffusion``), and a ``consistency_ratio`` fraction of the + batch pinned to ``r = 0``. It is configured directly on + :class:`~fastgen.methods.consistency_model.mean_flow.MeanFlowModel` (see ``configs/experiments/WanT2V/config_anyflow.py``); there is no AnyFlow-specific pretrain code. * **On-policy** (paper Stage 3, this module) is DMD2 on top of the pretrained - flow-map weights. The only differences from stock DMD2 are the two narrow - overrides below: + flow-map weights, with the reference's three deviations from stock DMD2: - - :meth:`AnyFlowModel._generate_noise_and_time` — the student always starts - from pure noise at ``max_t`` (DMD2's multi-step branch would noise real - data instead); - - :meth:`AnyFlowModel.gen_data_from_net` — the student generates via a - multi-step Euler-flow rollout with ``r = t_next`` (mean-velocity sampling) - and gradient enabled at one randomly-chosen step, matching AnyFlow's - ``WanAnyFlowPipeline.training_rollout``. + - the student generates via a flow-map rollout compressed into at most three + network forwards — one jump from ``t_0`` to ``t_g``, one fine step + ``t_g \\rightarrow t_{g+1}``, one jump to 0 — with the inference step + count sampled per iteration from ``student_sample_steps_list`` and the + fine-step position ``g`` sampled uniformly (both rank-0 broadcast). + Gradient flows through ALL segments + (:meth:`AnyFlowModel.gen_data_from_net`, mirroring + ``WanAnyFlowPipeline.training_rollout``); + - the student always starts from pure noise at ``max_t`` + (:meth:`AnyFlowModel._generate_noise_and_time`); + - every student update co-trains the Stage-2 flow-map loss on the real + batch (``cotrain_pretrain_weight``, mirroring the reference's + ``cotrain_forward_kl``), borrowing MeanFlow's loss machinery. The teacher and fake_score are flow-map networks queried at the - instantaneous velocity ``r = t`` — for gated-fusion Wan networks this is the - network-level default when ``r`` is not passed (see - ``fastgen/networks/Wan/network.py``), so DMD2's student / fake-score / - discriminator update steps are reused unchanged. + instantaneous velocity ``r = t`` — for gated-fusion Wan networks this is + the network-level default when ``r`` is not passed (see + ``fastgen/networks/Wan/network.py``), so DMD2's update steps run unchanged. """ from __future__ import annotations @@ -46,42 +52,113 @@ from fastgen.methods.consistency_model.mean_flow import MeanFlowModel from fastgen.methods.distribution_matching.dmd2 import DMD2Model +from fastgen.utils import basic_utils, expand_like from fastgen.utils.basic_utils import convert_cfg_to_dict import fastgen.utils.logging_utils as logger if TYPE_CHECKING: + from typing import Dict, Callable + from fastgen.configs.methods.config_anyflow import ModelConfig class AnyFlowModel(DMD2Model): - """AnyFlow on-policy stage: DMD2 with a multi-step rollout-with-gradient student.""" + """AnyFlow on-policy stage: DMD2 with a flow-map rollout student.""" # Validation sampling integrates the flow map with r = t_next (ode) just # like MeanFlow; FastGenModel's default x0-prediction loop never passes r. _student_sample_loop = MeanFlowModel._student_sample_loop + # MeanFlow loss machinery borrowed for the co-trained Stage-2 flow-map + # loss (the reference's cotrain_forward_kl runs the full bidirection loss + # at every generator step). + _sample_t_r_buckets = MeanFlowModel._sample_t_r_buckets + _reduce_mf_loss = MeanFlowModel._reduce_mf_loss + _compute_mf_loss = MeanFlowModel._compute_mf_loss + _drop_condition = MeanFlowModel._drop_condition + _jvp = MeanFlowModel._jvp + _estimate_jvp_finite_difference = MeanFlowModel._estimate_jvp_finite_difference + _mf_pred_to_loss = MeanFlowModel._mf_pred_to_loss + _compute_weight = MeanFlowModel._compute_weight + _timestep_weight_raw = MeanFlowModel._timestep_weight_raw + _get_timestep_weight = MeanFlowModel._get_timestep_weight + + @torch.no_grad() + def _get_velocity(self, x, z, t, condition=None, neg_condition=None): + """Raw data velocity for the co-trained flow-map loss. + + The DMD teacher CFG (``config.guidance_scale``) must not leak into the + co-trained loss — AnyFlow's guidance lives in + ``loss_config.guidance_fuse_scale`` (prediction-side fusion, handled + inside ``_compute_mf_loss``), so MeanFlow's target-side eq. 19 path is + deliberately not used here. + """ + del neg_condition + return condition, self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) + def __init__(self, config: ModelConfig): super().__init__(config) self.config = config + + # Attributes required by the borrowed MeanFlow methods. + self.sample_t_cfg = self.config.sample_t_cfg + self.sample_r_cfg = self.config.sample_r_cfg + self.loss_config = self.config.loss_config + if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp: + self.precision_amp_jvp = None + else: + self.precision_amp_jvp = basic_utils.PRECISION_MAP[self.config.precision_amp_jvp] + self._timestep_weight_scale = None + logger.info( - f"AnyFlow on-policy: student_sample_steps={self.config.student_sample_steps}, " + f"AnyFlow on-policy: student_sample_steps_list={self.config.student_sample_steps_list}, " f"student_update_freq={self.config.student_update_freq}, " + f"cotrain_pretrain_weight={self.config.cotrain_pretrain_weight}, " f"gan_loss_weight_gen={self.config.gan_loss_weight_gen}" ) - def _sample_grad_step(self, num_steps: int) -> int: - """Pick one rollout step index in ``[0, num_steps - 1]`` to enable gradients on. + # ------------------------------------------------------------------ + # Rollout + # ------------------------------------------------------------------ - Broadcast from rank 0 in distributed runs so all ranks agree on the - same window — matches AnyFlow's reference (``training_rollout`` in - ``pipeline_wan_anyflow.py``, ``broadcast(sample_step, src=0)``). - """ - idx = torch.randint(0, num_steps, (1,), device=self.device, dtype=torch.long) + def _broadcast_choice(self, high: int) -> int: + """Pick an index in [0, high) on rank 0 and broadcast it.""" + idx = torch.randint(0, high, (1,), device=self.device, dtype=torch.long) if torch.distributed.is_available() and torch.distributed.is_initialized(): torch.distributed.broadcast(idx, src=0) return int(idx.item()) + def _sample_rollout_steps(self) -> int: + """Sample the rollout NFE for this iteration. + + Matches the reference: ``random.choice(num_inference_steps_list)`` + broadcast from rank 0 in both the generator and fake-score updates. + Falls back to the fixed ``student_sample_steps`` when no list is set. + """ + steps_list = self.config.student_sample_steps_list + if not steps_list: + return int(self.config.student_sample_steps) + return int(steps_list[self._broadcast_choice(len(steps_list))]) + + def _rollout_t_list(self, num_steps: int) -> torch.Tensor: + """Shifted timestep schedule for ``num_steps`` rollout steps. + + Equivalent to the reference scheduler's ``set_timesteps``: a uniform + grid mapped through ``shift * s / (1 + (shift - 1) * s)``. The first + entry is clamped to the scheduler's ``max_t``. A config-provided + ``t_list`` takes precedence when its length matches. + """ + ns = self.net.noise_scheduler + cfg_t_list = self.config.sample_t_cfg.t_list + if cfg_t_list is not None and len(cfg_t_list) == num_steps + 1: + return torch.tensor(cfg_t_list, device=self.device, dtype=ns.t_precision) + + shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 + s = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64, device=self.device) + s = shift * s / (1 + (shift - 1) * s) + return s.clamp(max=float(ns.max_t)).to(ns.t_precision) + def _generate_noise_and_time( self, real_data: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: @@ -109,53 +186,79 @@ def gen_data_from_net( t_student: torch.Tensor, condition: Optional[Any] = None, ) -> torch.Tensor: - """Multi-step Euler-flow rollout with one gradient-enabled step. - - Mirrors AnyFlow's ``WanAnyFlowPipeline.training_rollout``: runs - ``student_sample_steps`` Euler-flow updates with ``r = t_next`` - (mean-velocity sampling, AnyFlow's ``use_mean_velocity=True`` default) - starting from ``input_student`` (pure noise at ``max_t``), and enables - gradients at one randomly-chosen step so the DMD generator update - receives a gradient through one full denoising forward. When the - caller wraps this in ``no_grad`` (the fake-score / discriminator - update), all steps run gradient-free. - """ - del t_student # the rollout schedule comes from t_list below - num_steps = int(self.config.student_sample_steps) - if num_steps < 1: - raise ValueError(f"student_sample_steps must be >= 1, got {num_steps}") + """Flow-map rollout compressed into at most three network forwards. + Mirrors the reference ``WanAnyFlowPipeline.training_rollout``: with an + ``num_steps``-step schedule ``t_0 > ... > t_N = 0`` and a sampled fine + step position ``g``, the model jumps ``t_0 -> t_g`` in ONE flow-map + forward, takes the single fine step ``t_g -> t_{g+1}``, then jumps + ``t_{g+1} -> 0`` in one forward. Gradient flows through all segments + (the reference applies no ``no_grad`` inside ``training_rollout``); + the fake-score update wraps this call in ``no_grad`` at the caller. + Each step uses mean-velocity sampling ``u_theta(x_t, t, r=t_next)``. + """ + del t_student # the rollout schedule is built below ns = self.net.noise_scheduler batch_size = input_student.shape[0] - grad_step = self._sample_grad_step(num_steps) - - # Timestep schedule. Use config-provided t_list when set (matches - # AnyFlow's hand-tuned step lists, e.g. [0.999, 0.937, 0.833, 0.624, 0.0] - # for 4-step Wan); otherwise fall back to the scheduler's default. - if self.config.sample_t_cfg.t_list is not None: - t_list = torch.tensor(self.config.sample_t_cfg.t_list, device=self.device, dtype=ns.t_precision) - if len(t_list) != num_steps + 1: - raise ValueError( - f"sample_t_cfg.t_list has {len(t_list)} entries, " - f"expected {num_steps + 1} for student_sample_steps={num_steps}" - ) - else: - t_list = ns.get_t_list(sample_steps=num_steps, device=self.device) - x = input_student - for step in range(num_steps): - t_cur = t_list[step].expand(batch_size).to(ns.t_precision) - t_next = t_list[step + 1].expand(batch_size).to(ns.t_precision) + num_steps = self._sample_rollout_steps() + if num_steps < 1: + raise ValueError(f"rollout steps must be >= 1, got {num_steps}") + grad_step = self._broadcast_choice(num_steps) + t_list = self._rollout_t_list(num_steps) - enable_grad = (step == grad_step) and torch.is_grad_enabled() - with torch.set_grad_enabled(enable_grad): - # Mean-velocity flow prediction: u_theta(x_t, t, r=t_next) - flow_pred = self.net(x, t_cur, r=t_next, condition=condition, fwd_pred_type="flow") + segments = ( + (t_list[0], t_list[grad_step]), + (t_list[grad_step], t_list[grad_step + 1]), + (t_list[grad_step + 1], t_list[-1]), + ) - # Euler-flow step. Steps run under no_grad add no graph nodes, and - # keeping ``x`` attached preserves the gradient installed by the - # grad-enabled step. - delta_t = (t_cur - t_next).view(batch_size, *([1] * (x.ndim - 1))).to(x.dtype) - x = x - delta_t * flow_pred + x = input_student + for t_cur, t_next in segments: + if float(t_cur) == float(t_next): + continue + t_batch = t_cur.expand(batch_size).to(ns.t_precision) + r_batch = t_next.expand(batch_size).to(ns.t_precision) + flow_pred = self.net(x, t_batch, r=r_batch, condition=condition, fwd_pred_type="flow") + x = x - expand_like(t_batch - r_batch, x).to(x.dtype) * flow_pred return x + + # ------------------------------------------------------------------ + # Training step — DMD2 plus the co-trained Stage-2 flow-map loss + # ------------------------------------------------------------------ + + def single_train_step( + self, data: "Dict[str, Any]", iteration: int + ) -> tuple[dict[str, torch.Tensor], dict[str, "torch.Tensor | Callable"]]: + real_data, condition, neg_condition = self._prepare_training_data(data) + self._setup_grad_requirements(iteration) + input_student, t_student, t, eps = self._generate_noise_and_time(real_data) + + if iteration % self.config.student_update_freq == 0: + loss_map, outputs = self._student_update_step( + input_student, t_student, t, eps, data, condition=condition, neg_condition=neg_condition + ) + if self.config.cotrain_pretrain_weight > 0: + # Reference cotrain_forward_kl: every generator update also + # runs the full Stage-2 bidirection (flow-map) loss on the + # real batch. + t_mf, r_mf, is_diffusion = self._sample_t_r_buckets(real_data.shape[0]) + mf_outputs = self._compute_mf_loss( + real_data=real_data, + t=t_mf, + r=r_mf, + iteration=iteration, + condition=condition, + neg_condition=neg_condition, + ) + bidirection_loss = self._reduce_mf_loss(mf_outputs[0], is_diffusion) + loss_map["bidirection_loss"] = bidirection_loss + loss_map["total_loss"] = ( + loss_map["total_loss"] + self.config.cotrain_pretrain_weight * bidirection_loss + ) + return loss_map, outputs + + return self._fake_score_discriminator_update_step( + input_student, t_student, t, eps, real_data, condition=condition + ) diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index 57ab5c2..2098bcc 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -69,11 +69,20 @@ def _build_onpolicy_model(): instance.precision = "float32" if instance.device == torch.device("cpu") else "bfloat16" instance.pretrained_model_path = "" instance.student_update_freq = 2 - # student_sample_steps=2 is the smallest value that runs the rollout loop - # more than once. instance.student_sample_steps = 2 instance.input_shape = [3, 8, 8] + # Co-trained flow-map loss settings (MeanFlow machinery on the tiny net). + instance.sample_t_cfg.time_dist_type = "uniform" + instance.sample_t_cfg.min_t = 0.001 + instance.sample_t_cfg.max_t = 0.999 + instance.sample_t_cfg.r_sample_ratio = 0.5 + instance.sample_t_cfg.consistency_ratio = 0.25 + instance.loss_config.loss_type = "l2" + instance.loss_config.weight_type = "uniform" + instance.loss_config.use_jvp_finite_diff = True + instance.loss_config.jvp_finite_diff_eps = 1e-2 + model = AnyFlowModel(instance) model.on_train_begin() model.init_optimizers() @@ -140,30 +149,68 @@ def test_pretrain_finite_difference_falls_back_at_boundaries(): assert torch.isfinite(u_theta_jvp).all(), "boundary samples must yield finite JVP estimates" -def test_pretrain_consistency_bucket_pins_r_to_min_t(): +def test_pretrain_consistency_bucket_pins_r_to_zero(): """With consistency_ratio=1.0 (and no flow-matching head), every sample's - r must be pinned to min_t.""" + r must be pinned to 0 (consistency to clean data, as in the reference).""" model = _build_pretrain_model(consistency_ratio=1.0, r_sample_ratio=1.0) - data = _make_data(model, batch_size=4) + t, r, is_diffusion = model._sample_t_r_buckets(4) + assert not is_diffusion.any() + assert torch.allclose(r.float(), torch.zeros_like(r.float())) + + +def test_pretrain_bucket_partition_is_deterministic(): + """With consistency_ratio > 0, buckets follow the reference's deterministic + global partition: diffusion head, consistency middle, random-pair tail.""" + model = _build_pretrain_model(consistency_ratio=0.25, r_sample_ratio=0.5) + batch_size = 8 + t, r, is_diffusion = model._sample_t_r_buckets(batch_size) + + n_diffusion = round(0.5 * batch_size) + n_consistency = round(0.25 * batch_size) + assert is_diffusion.tolist() == [True] * n_diffusion + [False] * (batch_size - n_diffusion) + assert torch.equal(r[:n_diffusion], t[:n_diffusion]) + assert torch.allclose( + r[n_diffusion : n_diffusion + n_consistency].float(), + torch.zeros(n_consistency), + ) + + +def test_pretrain_rebalance_to_diffusion(): + """With rebalancing on, each non-diffusion loss is rescaled to the + diffusion-loss mean by a detached per-sample factor.""" + model = _build_pretrain_model() + model.loss_config.rebalance_to_diffusion = True - captured = {} - orig = model._compute_mf_loss + mf_loss = torch.tensor([2.0, 4.0, 10.0, 100.0], dtype=torch.float64, requires_grad=True) + is_diffusion = torch.tensor([True, True, False, False]) + loss = model._reduce_mf_loss(mf_loss, is_diffusion) + + # diffusion mean = 3.0; each non-diffusion sample becomes ~3.0. + expected = (2.0 + 4.0 + 3.0 * (10.0 / 10.00001) + 3.0 * (100.0 / 100.00001)) / 4.0 + assert abs(loss.item() - expected) < 1e-3 + loss.backward() + assert mf_loss.grad is not None and torch.isfinite(mf_loss.grad).all() - def spy(real_data, t, r, **kwargs): - captured["r"] = r - return orig(real_data=real_data, t=t, r=r, **kwargs) - model._compute_mf_loss = spy - model.single_train_step(data, 0) +def test_pretrain_prediction_side_guidance_fusion(): + """The AnyFlow guidance-distillation branch (guidance_fuse_scale) must run + end to end and keep gradients on the fused prediction.""" + model = _build_pretrain_model() + model.loss_config.guidance_fuse_scale = 3.0 + model.config.cond_dropout_prob = 0.5 + data = _make_data(model, batch_size=2) - min_t = float(model.net.noise_scheduler.min_t) - assert torch.allclose(captured["r"].float(), torch.full_like(captured["r"].float(), min_t)) + loss_map, _ = model.single_train_step(data, 0) + assert torch.isfinite(loss_map["total_loss"]).all() + loss_map["total_loss"].backward() + grad_seen = any(p.grad is not None for p in model.net.parameters()) + assert grad_seen @pytest.mark.parametrize("weight_type", ["beta08", "gaussian", "uniform"]) def test_timestep_weight_function(weight_type): """The fixed per-timestep weight is a direct function of t: non-negative, - finite, and normalized to ~unit mean over the shifted timestep grid.""" + finite, and normalized like the reference scheduler.""" model = _build_pretrain_model(weight_type=weight_type) model._timestep_weight_scale = None # force re-derivation for this weight_type @@ -172,11 +219,15 @@ def test_timestep_weight_function(weight_type): assert torch.all(w >= 0), f"{weight_type} weights must be non-negative" assert torch.isfinite(w).all() - # The normalization matches the reference: sum over the shifted grid == 1000. - shift = model.sample_t_cfg.shift - grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64) - grid = shift * grid / (1 + (shift - 1) * grid) - assert abs(model._get_timestep_weight(grid).sum().item() - 1000.0) < 1e-6 + if weight_type == "uniform": + # The reference uses exactly 1.0 for uniform. + assert torch.allclose(w, torch.ones_like(w)) + else: + # The normalization matches the reference: sum over the shifted grid == 1000. + shift = model.sample_t_cfg.shift + grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64) + grid = shift * grid / (1 + (shift - 1) * grid) + assert abs(model._get_timestep_weight(grid).sum().item() - 1000.0) < 1e-6 # --------------------------------------------------------------------------- @@ -190,9 +241,44 @@ def test_onpolicy_student_update_step(): loss_map, outputs = model.single_train_step(data, 0) # iteration 0 -> student update assert "total_loss" in loss_map assert "vsd_loss" in loss_map + # The co-trained Stage-2 flow-map loss is part of every student update. + assert "bidirection_loss" in loss_map + assert torch.isfinite(loss_map["total_loss"]).all() assert "gen_rand" in outputs +def test_onpolicy_cotrain_can_be_disabled(): + model = _build_onpolicy_model() + model.config.cotrain_pretrain_weight = 0.0 + data = _make_data(model, img_resolution=8) + loss_map, _ = model.single_train_step(data, 0) + assert "bidirection_loss" not in loss_map + + +def test_onpolicy_rollout_compresses_to_three_forwards(): + """Regardless of the sampled NFE, the rollout must run at most three + network forwards (jump -> fine step -> jump), with gradient through all.""" + model = _build_onpolicy_model() + model.config.student_sample_steps_list = [16] + real = torch.randn(1, 3, 8, 8, device=model.device, dtype=model.precision) + cond = torch.nn.functional.one_hot(torch.tensor([0]), num_classes=10).to(model.device, model.precision) + input_student, t_student, _, _ = model._generate_noise_and_time(real) + + calls = [] + orig_forward = model.net.forward + + def counting_forward(*args, **kwargs): + calls.append(1) + return orig_forward(*args, **kwargs) + + model.net.forward = counting_forward + gen = model.gen_data_from_net(input_student, t_student, condition=cond) + model.net.forward = orig_forward + + assert len(calls) <= 3, f"rollout must compress to <= 3 forwards, got {len(calls)}" + assert gen.requires_grad + + def test_onpolicy_fake_score_discriminator_update_step(): model = _build_onpolicy_model() model.precision = torch.float32 From 6b097d7723b500e40443c7ec867ac4350785daa7 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Wed, 10 Jun 2026 23:41:38 +0800 Subject: [PATCH 07/19] anyflow: keep DMD-to-flow-map gradient ratio at the reference's 1:1 --- .../configs/experiments/WanT2V/config_anyflow_onpolicy.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index c16960e..19d9574 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -74,7 +74,10 @@ def create_config(): config.model.sample_t_cfg.t_list = None # rollout schedules are computed per NFE # ------ co-trained Stage-2 flow-map loss (reference cotrain_forward_kl) ------ - config.model.cotrain_pretrain_weight = 1.0 + # FastGen's VSD loss carries a 0.5 factor the reference's DMD loss does + # not; 0.5 here keeps the DMD : flow-map gradient ratio at the + # reference's 1 : 1 (the common 0.5 scale is absorbed by Adam). + config.model.cotrain_pretrain_weight = 0.5 config.model.loss_config.use_cd = False config.model.loss_config.loss_type = "l2" config.model.loss_config.weight_type = "beta08" From a3fd37ec87ff6716d5d6ac5c399bb8e2cf291db5 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 22:50:14 +0800 Subject: [PATCH 08/19] anyflow: address second-round review - Use fastgen.utils.distributed world_size/get_rank instead of raw torch.distributed queries (mean_flow buckets/rebalancing, anyflow rollout broadcast). - Fold the fixed per-timestep weighting into _compute_weight(tensor, t): the adaptive norm_method weight (None disables it) multiplies the optional weight_type weight, for both l2 and opt_grad losses. With norm_method=None the l2 loss reduces with a per-element mean, matching the reference loss scale; the default path is unchanged. - Align the weight-normalization grid with the reference set_timesteps (1000 points, t=0 excluded), which makes the uniform special case redundant; drop it. - Rename the is_diffusion mask to r_eq_t_mask. - Use attrs.field(factory=list) for cond_keys_no_dropout (the plain [] default is shared across config instances). - Formatting fixes for ruff==0.6.9 (trailing newline, line join) and test cleanups (device-safe zeros, unused unpack). Signed-off-by: Enderfga --- .../experiments/WanT2V/config_anyflow.py | 3 + .../WanT2V/config_anyflow_onpolicy.py | 1 + fastgen/configs/methods/config_anyflow.py | 2 +- fastgen/configs/methods/config_mean_flow.py | 13 +- .../methods/consistency_model/mean_flow.py | 113 +++++++++--------- .../methods/distribution_matching/anyflow.py | 11 +- fastgen/networks/Wan/network.py | 4 +- tests/test_anyflowmodel.py | 37 +++--- 8 files changed, 96 insertions(+), 88 deletions(-) diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py index e64bfab..65a9cc7 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -45,7 +45,10 @@ def create_config(): # ------ AnyFlow loss: MeanFlow l2 with fixed beta08 weighting ------ config.model.loss_config.use_cd = False config.model.loss_config.loss_type = "l2" + # Fixed beta08 per-timestep weighting on the per-element mean loss, no + # adaptive normalization — matching the reference train_bidirection. config.model.loss_config.weight_type = "beta08" + config.model.loss_config.norm_method = None config.model.loss_config.use_jvp_finite_diff = True config.model.loss_config.jvp_finite_diff_eps = 5e-3 # Rebalance the non-diffusion (flow-map / consistency) sample losses to diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index 19d9574..6595dc6 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -81,6 +81,7 @@ def create_config(): config.model.loss_config.use_cd = False config.model.loss_config.loss_type = "l2" config.model.loss_config.weight_type = "beta08" + config.model.loss_config.norm_method = None config.model.loss_config.use_jvp_finite_diff = True config.model.loss_config.jvp_finite_diff_eps = 5e-3 config.model.loss_config.rebalance_to_diffusion = True diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py index ad299d0..61671d3 100644 --- a/fastgen/configs/methods/config_anyflow.py +++ b/fastgen/configs/methods/config_anyflow.py @@ -59,7 +59,7 @@ class ModelConfig(DMD2ModelConfig): # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). cond_dropout_prob: Optional[float] = None - cond_keys_no_dropout: List[str] = [] + cond_keys_no_dropout: List[str] = attrs.field(factory=list) # Precision for autocast in the co-trained loss JVP (None = training precision). precision_amp_jvp: str | None = None diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index 1cc1fa4..6541f20 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -65,8 +65,9 @@ class LossConfig: use_jvp_finite_diff: bool = False # epsilon for finite difference estimate of JVP jvp_finite_diff_eps: float = 1e-4 - # normalize JVP - norm_method: str = "poly_1.0" + # adaptive loss normalization as a function of the per-sample loss + # (None disables it; the l2 loss then reduces with a per-element mean) + norm_method: Optional[str] = "poly_1.0" # tangent warmup constant norm_const: float = 1e-1 # tangent warmup steps @@ -75,9 +76,9 @@ class LossConfig: tangent_spatial_invariance: bool = False # loss type (choice between l2 and opt_grad) loss_type: str = "opt_grad" - # fixed per-timestep loss weighting evaluated as a function of t - # ("beta08", "gaussian", "uniform"; used by AnyFlow). None keeps the - # adaptive norm_method weighting above. Only applies to loss_type="l2". + # optional fixed per-timestep loss weighting evaluated as a function of t + # ("beta08", "gaussian", "uniform"; used by AnyFlow). Multiplies the + # adaptive norm_method weight above; None disables it. weight_type: Optional[str] = None # prediction-side guidance fusion scale (AnyFlow guidance distillation): # the conditional output is trained to be the guided flow directly via @@ -114,7 +115,7 @@ class ModelConfig(BaseModelConfig): cond_dropout_prob: Optional[float] = None # list of condition keys that do not drop - cond_keys_no_dropout: List[str] = [] + cond_keys_no_dropout: List[str] = attrs.field(factory=list) # guidance t start guidance_t_start: float = 0.0 diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 8f96ed1..9593635 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -11,6 +11,7 @@ from fastgen.methods import CMModel from fastgen.utils import basic_utils, expand_like from fastgen.utils.basic_utils import convert_cfg_to_dict +from fastgen.utils.distributed import get_rank, world_size import fastgen.utils.logging_utils as logger @@ -277,7 +278,7 @@ def net_wrapper(x_t, t, r): def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Sample (t, r) with t >= r and assign the per-batch buckets. - Returns ``(t, r, is_diffusion)`` where ``is_diffusion`` marks the + Returns ``(t, r, r_eq_t_mask)`` where ``r_eq_t_mask`` marks the samples with ``r = t`` (pure flow matching). With ``consistency_ratio == 0`` (MeanFlow default) the flow-matching @@ -298,51 +299,46 @@ def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tens assert torch.all(t >= r), "r cannot be larger than t" if self.sample_t_cfg.consistency_ratio > 0: - if torch.distributed.is_available() and torch.distributed.is_initialized(): - world_size, rank = torch.distributed.get_world_size(), torch.distributed.get_rank() - else: - world_size, rank = 1, 0 - global_bsz = world_size * batch_size - n_diffusion = round((1.0 - self.sample_t_cfg.r_sample_ratio) * global_bsz) + global_bsz = world_size() * batch_size + n_flow_matching = round((1.0 - self.sample_t_cfg.r_sample_ratio) * global_bsz) n_consistency = round(self.sample_t_cfg.consistency_ratio * global_bsz) - global_idx = rank * batch_size + torch.arange(batch_size, device=self.device) - is_diffusion = global_idx < n_diffusion - is_consistency = (global_idx >= n_diffusion) & (global_idx < n_diffusion + n_consistency) - r = torch.where(is_diffusion, t, r) + global_idx = get_rank() * batch_size + torch.arange(batch_size, device=self.device) + r_eq_t_mask = global_idx < n_flow_matching + is_consistency = (global_idx >= n_flow_matching) & (global_idx < n_flow_matching + n_consistency) + r = torch.where(r_eq_t_mask, t, r) r = torch.where(is_consistency, torch.zeros_like(r), r) else: # set t=r (flow matching loss) for a subset of the batch flow_matching_size = (torch.rand(batch_size, device=self.device) >= self.sample_t_cfg.r_sample_ratio).sum() - is_diffusion = torch.arange(batch_size, device=self.device) < flow_matching_size - r = torch.where(is_diffusion, t, r) + r_eq_t_mask = torch.arange(batch_size, device=self.device) < flow_matching_size + r = torch.where(r_eq_t_mask, t, r) - return t, r, is_diffusion + return t, r, r_eq_t_mask - def _reduce_mf_loss(self, mf_loss: torch.Tensor, is_diffusion: torch.Tensor) -> torch.Tensor: + def _reduce_mf_loss(self, mf_loss: torch.Tensor, r_eq_t_mask: torch.Tensor) -> torch.Tensor: """Reduce the per-sample loss to a scalar, optionally rebalancing. With ``loss_config.rebalance_to_diffusion`` set (AnyFlow), every - non-diffusion sample's loss is multiplied by the detached factor - ``mean(global diffusion losses) / (own loss + 1e-5)``, matching the - AnyFlow reference (``train_bidirection``): flow-map / consistency - gradients are self-normalized and rescaled to the global - diffusion-loss mean. + flow-map / consistency (r < t) sample's loss is multiplied by the + detached factor ``mean(global flow-matching losses) / (own loss + 1e-5)``, + matching the AnyFlow reference (``train_bidirection``): flow-map / + consistency gradients are self-normalized and rescaled to the global + flow-matching-loss mean. """ - if getattr(self.loss_config, "rebalance_to_diffusion", False) and (~is_diffusion).any(): + if getattr(self.loss_config, "rebalance_to_diffusion", False) and (~r_eq_t_mask).any(): with torch.no_grad(): - if torch.distributed.is_available() and torch.distributed.is_initialized(): - world_size = torch.distributed.get_world_size() - gathered_loss = [torch.zeros_like(mf_loss) for _ in range(world_size)] - gathered_mask = [torch.zeros_like(is_diffusion) for _ in range(world_size)] + if world_size() > 1: + gathered_loss = [torch.zeros_like(mf_loss) for _ in range(world_size())] + gathered_mask = [torch.zeros_like(r_eq_t_mask) for _ in range(world_size())] torch.distributed.all_gather(gathered_loss, mf_loss.contiguous()) - torch.distributed.all_gather(gathered_mask, is_diffusion.contiguous()) + torch.distributed.all_gather(gathered_mask, r_eq_t_mask.contiguous()) global_loss = torch.cat(gathered_loss, dim=0) global_mask = torch.cat(gathered_mask, dim=0) else: - global_loss, global_mask = mf_loss, is_diffusion + global_loss, global_mask = mf_loss, r_eq_t_mask scale = torch.ones_like(mf_loss) if global_mask.any(): - scale[~is_diffusion] = global_loss[global_mask].mean() / (mf_loss[~is_diffusion] + 1e-5) + scale[~r_eq_t_mask] = global_loss[global_mask].mean() / (mf_loss[~r_eq_t_mask] + 1e-5) mf_loss = mf_loss * scale return mf_loss.mean() @@ -364,33 +360,41 @@ def _get_timestep_weight(self, t: torch.Tensor) -> torch.Tensor: The weight is evaluated directly as a function of t. The normalization constant matches the AnyFlow reference scheduler, which normalizes the - weights over its shifted discrete timestep grid; we compute the same - constant once and cache it. + weights over its shifted discrete timestep grid (``set_timesteps``: + 1000 points excluding t=0); we compute the same constant once and + cache it. """ - if self.loss_config.weight_type == "uniform": - # The reference uses exactly 1.0 for uniform (no grid normalization). - return torch.ones_like(t) if self._timestep_weight_scale is None: shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 - grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64) + grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64)[:-1] grid = shift * grid / (1 + (shift - 1) * grid) self._timestep_weight_scale = float(1000.0 / self._timestep_weight_raw(grid).sum()) return self._timestep_weight_raw(t) * self._timestep_weight_scale @torch.no_grad() - def _compute_weight(self, tensor: torch.Tensor) -> torch.Tensor: - norm_method, *norm_args = self.loss_config.norm_method.split("_") - - if norm_method == "poly": - power = float(norm_args[0]) - assert len(norm_args) == 1, "poly norm method requires 1 argument" - weight = 1 / (tensor + self.loss_config.norm_const).pow(power) - elif norm_method == "exp": - assert len(norm_args) == 2, "exp norm method requires 2 arguments" - const, scale = float(norm_args[0]), float(norm_args[1]) - weight = const * torch.exp(scale * tensor + self.loss_config.norm_const) + def _compute_weight(self, tensor: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """Per-sample loss weight: the adaptive normalization (``norm_method``, + a function of the per-sample loss; ``None`` disables it) times the + optional fixed per-timestep weight (``weight_type``, a function of t). + """ + if self.loss_config.norm_method is None: + weight = torch.ones_like(tensor) else: - raise ValueError(f"Invalid norm method: {self.loss_config.norm_method}") + norm_method, *norm_args = self.loss_config.norm_method.split("_") + + if norm_method == "poly": + power = float(norm_args[0]) + assert len(norm_args) == 1, "poly norm method requires 1 argument" + weight = 1 / (tensor + self.loss_config.norm_const).pow(power) + elif norm_method == "exp": + assert len(norm_args) == 2, "exp norm method requires 2 arguments" + const, scale = float(norm_args[0]), float(norm_args[1]) + weight = const * torch.exp(scale * tensor + self.loss_config.norm_const) + else: + raise ValueError(f"Invalid norm method: {self.loss_config.norm_method}") + + if self.loss_config.weight_type is not None: + weight = weight * self._get_timestep_weight(t) assert ( weight.shape == tensor.shape @@ -432,15 +436,14 @@ def _mf_pred_to_loss( if self.loss_config.loss_type == "l2": tangent = dxt_dt - warmup_weight * delta_t * u_theta_jvp loss = (u_theta - tangent).pow(2) - if self.loss_config.weight_type is not None: - # Fixed per-timestep weighting with mean reduction (AnyFlow). + if self.loss_config.norm_method is None: + # Without adaptive normalization the reduction matters: the + # per-element mean matches the AnyFlow reference loss scale. loss = torch.mean(loss, dim=list(range(1, loss.ndim))) - weight = self._get_timestep_weight(t) - loss = loss * weight else: loss = torch.sum(loss, dim=list(range(1, loss.ndim))) - weight = self._compute_weight(loss) - loss = loss * weight + weight = self._compute_weight(loss, t) + loss = loss * weight # use explicit gradient elif self.loss_config.loss_type == "opt_grad": @@ -453,7 +456,7 @@ def _mf_pred_to_loss( tangent = tangent * sample_dim_inv opt_grad_norm = torch.linalg.vector_norm(tangent.flatten(1), dim=-1) - weight = self._compute_weight(opt_grad_norm) + weight = self._compute_weight(opt_grad_norm, t) 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))) @@ -613,7 +616,7 @@ def single_train_step( batch_size = real_data.shape[0] # sample t and r, with per-batch buckets (flow matching / consistency) - t, r, is_diffusion = self._sample_t_r_buckets(batch_size) + t, r, r_eq_t_mask = self._sample_t_r_buckets(batch_size) ( mf_loss, @@ -633,7 +636,7 @@ def single_train_step( neg_condition=neg_condition, ) - loss = self._reduce_mf_loss(mf_loss, is_diffusion) + loss = self._reduce_mf_loss(mf_loss, r_eq_t_mask) loss_map = { "total_loss": loss, "mf_loss": loss, diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index a7cd69b..271c7e1 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -54,6 +54,7 @@ from fastgen.methods.distribution_matching.dmd2 import DMD2Model from fastgen.utils import basic_utils, expand_like from fastgen.utils.basic_utils import convert_cfg_to_dict +from fastgen.utils.distributed import world_size import fastgen.utils.logging_utils as logger @@ -125,7 +126,7 @@ def __init__(self, config: ModelConfig): def _broadcast_choice(self, high: int) -> int: """Pick an index in [0, high) on rank 0 and broadcast it.""" idx = torch.randint(0, high, (1,), device=self.device, dtype=torch.long) - if torch.distributed.is_available() and torch.distributed.is_initialized(): + if world_size() > 1: torch.distributed.broadcast(idx, src=0) return int(idx.item()) @@ -243,7 +244,7 @@ def single_train_step( # Reference cotrain_forward_kl: every generator update also # runs the full Stage-2 bidirection (flow-map) loss on the # real batch. - t_mf, r_mf, is_diffusion = self._sample_t_r_buckets(real_data.shape[0]) + t_mf, r_mf, r_eq_t_mask = self._sample_t_r_buckets(real_data.shape[0]) mf_outputs = self._compute_mf_loss( real_data=real_data, t=t_mf, @@ -252,11 +253,9 @@ def single_train_step( condition=condition, neg_condition=neg_condition, ) - bidirection_loss = self._reduce_mf_loss(mf_outputs[0], is_diffusion) + bidirection_loss = self._reduce_mf_loss(mf_outputs[0], r_eq_t_mask) loss_map["bidirection_loss"] = bidirection_loss - loss_map["total_loss"] = ( - loss_map["total_loss"] + self.config.cotrain_pretrain_weight * bidirection_loss - ) + loss_map["total_loss"] = loss_map["total_loss"] + self.config.cotrain_pretrain_weight * bidirection_loss return loss_map, outputs return self._fake_score_discriminator_update_step( diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index 869681e..caf2c5b 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -386,9 +386,7 @@ def classify_forward_prepare( remb = self.r_embedder.time_embedder(r_timestep).type_as(encoder_hidden_states) - temb, timestep_proj, r_timestep_proj = self._fuse_r_embedding( - temb, timestep_proj, remb, rs_seq_len - ) + temb, timestep_proj, r_timestep_proj = self._fuse_r_embedding(temb, timestep_proj, remb, rs_seq_len) elif r_timestep is not None: # Raise an error here, otherwise we silently ignore the r_timestep raise ValueError("r_timestep provided but no r_embedder is present") diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index 2098bcc..cd6817e 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -31,6 +31,7 @@ def _build_pretrain_model(weight_type="beta08", consistency_ratio=0.25, r_sample instance.loss_config.loss_type = "l2" instance.loss_config.weight_type = weight_type + instance.loss_config.norm_method = None instance.loss_config.use_jvp_finite_diff = True instance.loss_config.jvp_finite_diff_eps = 1e-2 @@ -80,6 +81,7 @@ def _build_onpolicy_model(): instance.sample_t_cfg.consistency_ratio = 0.25 instance.loss_config.loss_type = "l2" instance.loss_config.weight_type = "uniform" + instance.loss_config.norm_method = None instance.loss_config.use_jvp_finite_diff = True instance.loss_config.jvp_finite_diff_eps = 1e-2 @@ -153,8 +155,8 @@ def test_pretrain_consistency_bucket_pins_r_to_zero(): """With consistency_ratio=1.0 (and no flow-matching head), every sample's r must be pinned to 0 (consistency to clean data, as in the reference).""" model = _build_pretrain_model(consistency_ratio=1.0, r_sample_ratio=1.0) - t, r, is_diffusion = model._sample_t_r_buckets(4) - assert not is_diffusion.any() + _t, r, r_eq_t_mask = model._sample_t_r_buckets(4) + assert not r_eq_t_mask.any() assert torch.allclose(r.float(), torch.zeros_like(r.float())) @@ -163,15 +165,15 @@ def test_pretrain_bucket_partition_is_deterministic(): global partition: diffusion head, consistency middle, random-pair tail.""" model = _build_pretrain_model(consistency_ratio=0.25, r_sample_ratio=0.5) batch_size = 8 - t, r, is_diffusion = model._sample_t_r_buckets(batch_size) + t, r, r_eq_t_mask = model._sample_t_r_buckets(batch_size) - n_diffusion = round(0.5 * batch_size) + n_flow_matching = round(0.5 * batch_size) n_consistency = round(0.25 * batch_size) - assert is_diffusion.tolist() == [True] * n_diffusion + [False] * (batch_size - n_diffusion) - assert torch.equal(r[:n_diffusion], t[:n_diffusion]) + assert r_eq_t_mask.tolist() == [True] * n_flow_matching + [False] * (batch_size - n_flow_matching) + assert torch.equal(r[:n_flow_matching], t[:n_flow_matching]) assert torch.allclose( - r[n_diffusion : n_diffusion + n_consistency].float(), - torch.zeros(n_consistency), + r[n_flow_matching : n_flow_matching + n_consistency].float(), + torch.zeros(n_consistency, device=r.device), ) @@ -182,8 +184,8 @@ def test_pretrain_rebalance_to_diffusion(): model.loss_config.rebalance_to_diffusion = True mf_loss = torch.tensor([2.0, 4.0, 10.0, 100.0], dtype=torch.float64, requires_grad=True) - is_diffusion = torch.tensor([True, True, False, False]) - loss = model._reduce_mf_loss(mf_loss, is_diffusion) + r_eq_t_mask = torch.tensor([True, True, False, False]) + loss = model._reduce_mf_loss(mf_loss, r_eq_t_mask) # diffusion mean = 3.0; each non-diffusion sample becomes ~3.0. expected = (2.0 + 4.0 + 3.0 * (10.0 / 10.00001) + 3.0 * (100.0 / 100.00001)) / 4.0 @@ -220,14 +222,15 @@ def test_timestep_weight_function(weight_type): assert torch.isfinite(w).all() if weight_type == "uniform": - # The reference uses exactly 1.0 for uniform. + # Uniform normalizes to exactly 1.0 over the reference grid. assert torch.allclose(w, torch.ones_like(w)) - else: - # The normalization matches the reference: sum over the shifted grid == 1000. - shift = model.sample_t_cfg.shift - grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64) - grid = shift * grid / (1 + (shift - 1) * grid) - assert abs(model._get_timestep_weight(grid).sum().item() - 1000.0) < 1e-6 + + # The normalization matches the reference set_timesteps grid (1000 points, + # t=0 excluded): sum over the shifted grid == 1000. + shift = model.sample_t_cfg.shift + grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64)[:-1] + grid = shift * grid / (1 + (shift - 1) * grid) + assert abs(model._get_timestep_weight(grid).sum().item() - 1000.0) < 1e-6 # --------------------------------------------------------------------------- From 01c8ff20f315472dcd91b2040b4b71e2d444c586 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 23:02:37 +0800 Subject: [PATCH 09/19] anyflow: validate guidance-fusion inputs Fail fast with a clear message when guidance_fuse_scale is non-positive (the fused prediction divides by it) or when neg_condition is missing (the unconditional branch is queried at the same (t, r)). Signed-off-by: Enderfga --- fastgen/methods/consistency_model/mean_flow.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 9593635..a197436 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -548,6 +548,10 @@ def _compute_mf_loss( # (u_cond + (g - 1) * u_uncond) / g. Text dropout replaces the # condition with neg_condition for a random subset beforehand. g = float(guidance_fuse_scale) + assert g > 0, f"guidance_fuse_scale must be > 0, got {g} (set it to None to disable guidance fusion)" + assert ( + neg_condition is not None + ), "guidance_fuse_scale requires neg_condition: the unconditional branch is queried at the same (t, r)" condition = self._drop_condition(condition, neg_condition) dxt_dt = self.net.noise_scheduler.cond_velocity(x=real_data, eps=z, t=t) torch.clear_autocast_cache() From b6c153eeed192573207693ebf20b2f35401238dd Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 23:09:04 +0800 Subject: [PATCH 10/19] anyflow: reduce rebalance scalars; warn on degenerate buckets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The rebalance factor only needs the global flow-matching-loss mean, so all_reduce a sum and a count instead of all_gathering the per-sample losses — equivalent to the reference's cat(all_gather(loss)) math, cheaper, and independent of per-rank batch sizes. The deterministic (t, r) bucket partition spans ranks but not gradient-accumulation rounds (as in the AnyFlow reference); log a one-time warning when the configured ratios produce empty buckets so a degenerate setup (e.g. single rank at batch size 1) is visible. Signed-off-by: Enderfga --- .../methods/consistency_model/mean_flow.py | 32 +++++++++++++------ 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index a197436..bed725f 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -302,6 +302,17 @@ def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tens global_bsz = world_size() * batch_size n_flow_matching = round((1.0 - self.sample_t_cfg.r_sample_ratio) * global_bsz) n_consistency = round(self.sample_t_cfg.consistency_ratio * global_bsz) + if (n_flow_matching == 0 or n_consistency == 0) and not getattr(self, "_bucket_warned", False): + self._bucket_warned = True + logger.warning( + f"The deterministic (t, r) bucket partition is degenerate: with " + f"world_size * batch_size = {global_bsz}, r_sample_ratio=" + f"{self.sample_t_cfg.r_sample_ratio} and consistency_ratio=" + f"{self.sample_t_cfg.consistency_ratio} yield bucket sizes " + f"({n_flow_matching}, {n_consistency}). The partition spans ranks but not " + f"gradient-accumulation rounds (as in the AnyFlow reference), so empty " + f"buckets stay empty every iteration." + ) global_idx = get_rank() * batch_size + torch.arange(batch_size, device=self.device) r_eq_t_mask = global_idx < n_flow_matching is_consistency = (global_idx >= n_flow_matching) & (global_idx < n_flow_matching + n_consistency) @@ -327,18 +338,19 @@ def _reduce_mf_loss(self, mf_loss: torch.Tensor, r_eq_t_mask: torch.Tensor) -> t """ if getattr(self.loss_config, "rebalance_to_diffusion", False) and (~r_eq_t_mask).any(): with torch.no_grad(): + # The global flow-matching-loss mean only needs the global sum + # and count, so reduce two scalars instead of gathering the + # per-sample losses (equivalent to the reference's + # cat(all_gather(loss))[mask].mean(), and independent of the + # per-rank batch sizes). + fm_loss_sum = torch.where(r_eq_t_mask, mf_loss, torch.zeros_like(mf_loss)).sum() + fm_count = r_eq_t_mask.sum().to(mf_loss.dtype) if world_size() > 1: - gathered_loss = [torch.zeros_like(mf_loss) for _ in range(world_size())] - gathered_mask = [torch.zeros_like(r_eq_t_mask) for _ in range(world_size())] - torch.distributed.all_gather(gathered_loss, mf_loss.contiguous()) - torch.distributed.all_gather(gathered_mask, r_eq_t_mask.contiguous()) - global_loss = torch.cat(gathered_loss, dim=0) - global_mask = torch.cat(gathered_mask, dim=0) - else: - global_loss, global_mask = mf_loss, r_eq_t_mask + torch.distributed.all_reduce(fm_loss_sum) + torch.distributed.all_reduce(fm_count) scale = torch.ones_like(mf_loss) - if global_mask.any(): - scale[~r_eq_t_mask] = global_loss[global_mask].mean() / (mf_loss[~r_eq_t_mask] + 1e-5) + if fm_count > 0: + scale[~r_eq_t_mask] = (fm_loss_sum / fm_count) / (mf_loss[~r_eq_t_mask] + 1e-5) mf_loss = mf_loss * scale return mf_loss.mean() From 3dae9d9c14e0174aed79246e99581544f555090a Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 23:14:18 +0800 Subject: [PATCH 11/19] anyflow: drop the rank-local guard around the rebalance collective MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The all_reduce in _reduce_mf_loss was gated on the rank-local (~r_eq_t_mask).any(), so a rank holding only flow-matching samples (which the deterministic bucket partition produces by design) skipped the collective while other ranks entered it, deadlocking distributed training. Gate on the config flag only — identical on every rank — and let the empty-selection scale assignment be a no-op, as in the reference (which runs its gather unconditionally). Add a regression test for the all-flow-matching batch. Signed-off-by: Enderfga --- fastgen/methods/consistency_model/mean_flow.py | 5 ++++- tests/test_anyflowmodel.py | 13 +++++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index bed725f..98cdd3f 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -336,7 +336,10 @@ def _reduce_mf_loss(self, mf_loss: torch.Tensor, r_eq_t_mask: torch.Tensor) -> t consistency gradients are self-normalized and rescaled to the global flow-matching-loss mean. """ - if getattr(self.loss_config, "rebalance_to_diffusion", False) and (~r_eq_t_mask).any(): + # No rank-local condition here: the branch must be taken (or not) by + # every rank so the collective below cannot deadlock. Applying the + # scale to an empty ~r_eq_t_mask selection is a no-op. + if getattr(self.loss_config, "rebalance_to_diffusion", False): with torch.no_grad(): # The global flow-matching-loss mean only needs the global sum # and count, so reduce two scalars instead of gathering the diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index cd6817e..6447a20 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -194,6 +194,19 @@ def test_pretrain_rebalance_to_diffusion(): assert mf_loss.grad is not None and torch.isfinite(mf_loss.grad).all() +def test_pretrain_rebalance_all_flow_matching_batch(): + """A rank whose batch is entirely flow-matching (r = t) must still take + the rebalance branch (the collective inside must run on every rank) and + reduce to the plain mean.""" + model = _build_pretrain_model() + model.loss_config.rebalance_to_diffusion = True + + mf_loss = torch.tensor([2.0, 4.0], dtype=torch.float64) + r_eq_t_mask = torch.tensor([True, True]) + loss = model._reduce_mf_loss(mf_loss, r_eq_t_mask) + assert abs(loss.item() - 3.0) < 1e-8 + + def test_pretrain_prediction_side_guidance_fusion(): """The AnyFlow guidance-distillation branch (guidance_fuse_scale) must run end to end and keep gradients on the fused prediction.""" From 01be71f21658ac04c9a15ec031c5252bd44aabab Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 23:20:58 +0800 Subject: [PATCH 12/19] anyflow: fix stale docs The on-policy config disables the adversarial loss (gan_loss_weight_gen=0), so the README config line saying 'GAN on' contradicted both the config and the paragraph above it. The rollout gradient test docstring still described the old one-step-with-gradient scheme instead of the compressed all-segments rollout. Signed-off-by: Enderfga --- fastgen/methods/distribution_matching/README.md | 2 +- tests/test_anyflowmodel.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index b3db151..29abe84 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -114,7 +114,7 @@ Single model that supports arbitrary inference NFE by learning a flow map `u_θ( **Configs:** - [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) — Stage 2 pretrain (6k iter, lr=5e-5, shift=5, beta08, paper-aligned) -- [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) — Stage 3 on-policy distillation (1.2k iter, lr=2e-6, GAN on; full-rank fine-tune of a Stage 2 pretrain ckpt — the paper's rank-256 LoRA variant awaits a PEFT path in FastGen) +- [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) — Stage 3 on-policy distillation (1.2k iter, lr=2e-6, no adversarial loss; full-rank fine-tune of a Stage 2 pretrain ckpt — the paper's rank-256 LoRA variant awaits a PEFT path in FastGen) --- diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index 6447a20..fef00e4 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -310,12 +310,12 @@ def test_onpolicy_fake_score_discriminator_update_step(): def test_onpolicy_rollout_propagates_gradient(): - """The multi-step rollout must allow gradient flow on the chosen step. + """The rollout output must keep the autograd graph so the DMD generator + update has a valid gradient. - Mirrors AnyFlow's ``training_rollout`` (pipeline_wan_anyflow.py): one - randomly-chosen step in the rollout has gradients enabled; the remaining - steps are no_grad. The rollout output must keep the autograd graph so the - DMD generator update has a valid gradient. + Mirrors AnyFlow's ``training_rollout`` (pipeline_wan_anyflow.py): the + compressed jump -> fine step -> jump rollout runs with gradient through + all segments. """ model = _build_onpolicy_model() real = torch.randn(1, 3, 8, 8, device=model.device, dtype=model.precision) From e3d283370844ad0ea1f4daa0e41135a10c1e98a7 Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Tue, 11 Aug 2026 06:38:28 +0000 Subject: [PATCH 13/19] anyflow: reference-exact [0,1] timestep range, on-policy noise start, doc trims Signed-off-by: Julius Berner --- .../configs/experiments/DiT/config_mf_b.py | 2 +- .../experiments/EDM/config_mf_cifar10.py | 2 +- .../experiments/WanT2V/config_anyflow.py | 60 +-- .../WanT2V/config_anyflow_onpolicy.py | 108 ++++-- .../configs/experiments/WanT2V/config_mf.py | 9 +- fastgen/configs/methods/config_anyflow.py | 16 +- fastgen/configs/methods/config_dmd2.py | 11 + fastgen/configs/methods/config_mean_flow.py | 20 +- fastgen/methods/README.md | 1 + fastgen/methods/consistency_model/README.md | 2 +- .../methods/consistency_model/mean_flow.py | 361 +++++++++--------- .../methods/distribution_matching/README.md | 39 +- .../methods/distribution_matching/anyflow.py | 209 ++++------ .../methods/distribution_matching/causvid.py | 11 +- fastgen/methods/distribution_matching/dmd2.py | 27 +- .../distribution_matching/self_forcing.py | 10 +- fastgen/networks/EDM/network.py | 6 - fastgen/networks/Wan/network.py | 233 ++++------- fastgen/networks/Wan/utils.py | 134 +++++++ fastgen/networks/noise_schedule.py | 20 +- tests/test_anyflowmodel.py | 201 ++++++++-- tests/test_meanflowmodel.py | 2 +- 22 files changed, 847 insertions(+), 637 deletions(-) create mode 100644 fastgen/networks/Wan/utils.py diff --git a/fastgen/configs/experiments/DiT/config_mf_b.py b/fastgen/configs/experiments/DiT/config_mf_b.py index 490913e..6a7a7c8 100644 --- a/fastgen/configs/experiments/DiT/config_mf_b.py +++ b/fastgen/configs/experiments/DiT/config_mf_b.py @@ -26,7 +26,7 @@ def create_config(): config.model.sample_t_cfg.train_p_std = 1.0 config.model.sample_t_cfg.min_t = 0.0 config.model.sample_t_cfg.max_t = 0.999 - config.model.sample_t_cfg.r_sample_ratio = 0.25 + config.model.sample_t_cfg.flow_matching_ratio = 0.75 config.model.loss_config.norm_method = "poly_1.0" config.model.loss_config.norm_const = 1.0 diff --git a/fastgen/configs/experiments/EDM/config_mf_cifar10.py b/fastgen/configs/experiments/EDM/config_mf_cifar10.py index 19b41ca..5779786 100644 --- a/fastgen/configs/experiments/EDM/config_mf_cifar10.py +++ b/fastgen/configs/experiments/EDM/config_mf_cifar10.py @@ -19,7 +19,7 @@ def create_config(): config.model.sample_t_cfg.train_p_mean = -0.6 config.model.sample_t_cfg.train_p_std = 1.6 - config.model.sample_t_cfg.r_sample_ratio = 0.75 + config.model.sample_t_cfg.flow_matching_ratio = 0.25 config.model.sample_t_cfg.time_dist_type = "logitnormal" config.model.sample_t_cfg.min_t = 0.0 config.model.sample_t_cfg.max_t = 0.999 diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py index 65a9cc7..f2efd88 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -1,21 +1,17 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""AnyFlow flow-map pretrain config on Wan-1.3B T2V (paper Stage 2). +"""AnyFlow flow-map pretrain config on Wan-1.3B T2V (paper Stage 1). -AnyFlow's pretrain objective is the MeanFlow objective with a fixed -``beta08`` per-timestep loss weighting, a finite-difference JVP, shifted -timestep sampling, and a ``consistency_ratio`` fraction of the batch pinned -to ``r = min_t`` — so this config runs :class:`MeanFlowModel` directly. +AnyFlow's pretrain objective is MeanFlow's with a fixed ``beta08`` per-timestep +weighting, a finite-difference JVP, shifted timestep sampling, and a +``consistency_ratio`` fraction of the batch pinned to ``r = 0`` — so this config +runs``MeanFlowModel`` directly. The values below mirror the reference recipe +``train_wan1b_student_shift5_81f_480p_lr5e-5_6k_b32.yml``. +Known deviations from the reference: full-rank fine-tuning instead of the paper's +rank-256 LoRA — the same deviation applies to the on-policy stage; -Mirrors the AnyFlow paper's pretrain recipe (``shift=5``, ``beta08`` -weighting, ``epsilon=5e-3``, ``diffusion_ratio=0.5``, -``consistency_ratio=0.25``, lr=5e-5, 6k iterations, batch_size_global=32, -text dropout 0.1 with velocity-fused guidance 3.0) on the gated dual-timestep -Wan architecture (``gate=0.25``, absolute-r conditioning, matching -``deltatime_type: r`` in the reference). - -The on-policy stage (paper Stage 3) lives in ``config_anyflow_onpolicy.py``. +The on-policy stage (paper Stage 2) lives in ``config_anyflow_onpolicy.py``. """ import copy @@ -37,8 +33,17 @@ def create_config(): config.model.net.time_cond_type = "abs" config.model.net.r_embedder_fusion = "gated" config.model.net.r_embedder_gate_value = 0.25 + # Noise-schedule bounds, forwarded to RFNoiseSchedule. + config.model.net.min_t = 0.0 + config.model.net.max_t = 1.0 config.model.precision = "bfloat16" + # FSDP2 parameter storage and gradient reduction in fp32 while compute stays + # bfloat16 -- the same split as the reference. Takes + # effect only under FSDP (`trainer.ddp=False`); it is ignored under DDP, where + # params, grads and compute are all `precision`. + config.model.precision_fsdp = "float32" + # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. config.model.input_shape = [16, 21, 60, 104] @@ -50,10 +55,11 @@ def create_config(): config.model.loss_config.weight_type = "beta08" config.model.loss_config.norm_method = None config.model.loss_config.use_jvp_finite_diff = True + # Reference epsilon=5 in 1000-step units = 5e-3 in FastGen's continuous time. config.model.loss_config.jvp_finite_diff_eps = 5e-3 - # Rebalance the non-diffusion (flow-map / consistency) sample losses to - # the global diffusion-loss mean (reference scale_weight). - config.model.loss_config.rebalance_to_diffusion = True + # Rebalance the flow-map / consistency (r < t) sample losses to the global + # flow-matching (r = t) loss mean (reference scale_weight). + config.model.loss_config.rebalance_to_flow_matching = True config.model.precision_amp_jvp = "float32" # Prediction-side guidance fusion with text dropout (reference: @@ -67,13 +73,16 @@ def create_config(): # ------ (t, r) sampling: shifted uniform pairs + AnyFlow buckets ------ config.model.sample_t_cfg.time_dist_type = "shifted" config.model.sample_t_cfg.shift = 5.0 - config.model.sample_t_cfg.min_t = 0.001 - config.model.sample_t_cfg.max_t = 0.999 - # diffusion_ratio=0.5 of the batch keeps r = t (pure flow matching); - # r_sample_ratio is the complementary fraction that keeps the sampled r. - config.model.sample_t_cfg.r_sample_ratio = 0.5 - # consistency_ratio=0.25 of the batch is pinned to r = min_t. + config.model.sample_t_cfg.min_t = 0.0 + config.model.sample_t_cfg.max_t = 1.0 + # diffusion_ratio=0.5 of the batch keeps r = t (pure flow matching). + config.model.sample_t_cfg.flow_matching_ratio = 0.5 + # consistency_ratio=0.25 of the batch is pinned to r = 0 (the reference + # sets r = 0 pre-shift, and the shift maps 0 to 0). config.model.sample_t_cfg.consistency_ratio = 0.25 + # The reference assigns both buckets by rank-indexed partition of the + # global batch, not by an independent per-sample draw. + config.model.sample_t_cfg.deterministic_buckets = True # ------ optimization (reference: AdamW lr=5e-5, wd=0, betas=(0.9, 0.95), # max_grad_norm=1.0, 1000-step LR warmup, EMA decay 0.999) ------ @@ -84,13 +93,14 @@ def create_config(): config.model.net_scheduler.warm_up_steps = [1000] config.trainer.callbacks.grad_clip.grad_norm = 1.0 config.trainer.callbacks.ema.beta = 0.999 + # Reference `ema_warmup_step: 1000` (as `start_iter = warmup_steps - 1`). + config.trainer.callbacks.ema.start_iter = 999 # ------ inference / validation ------ config.model.student_sample_type = "ode" config.model.student_sample_steps = 4 - # 4-step shifted schedule shift*s/(1+(shift-1)*s) on s=linspace(1,0,5), - # first entry clamped to max_t (the reference samples from t=1.0). - config.model.sample_t_cfg.t_list = [0.999, 0.9375, 0.83333333, 0.625, 0.0] + # 4-step shifted schedule shift*s/(1+(shift-1)*s) on s=linspace(1,0,5) + config.model.sample_t_cfg.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625, 0.0] # ------ data / trainer ------ config.dataloader_train = VideoLoaderConfig diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index 6595dc6..e520df5 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -1,30 +1,24 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""AnyFlow on-policy distillation config on Wan-1.3B T2V (paper Stage 3). - -DMD2-style distribution matching with a flow-map student that generates via -a compressed rollout (jump -> fine step -> jump, gradient through all -segments, NFE sampled per iteration from [2, 4, 8, 16, 50]) and co-trains the -Stage-2 flow-map loss at every student update — see -:class:`~fastgen.methods.distribution_matching.anyflow.AnyFlowModel`. - -Mirrors the reference recipe -(``train_wan1b_onpolicy_81f_480p_lr2e-6_1k_b32.yml``): 1.2k iterations at -lr=2e-6, AdamW betas=(0.0, 0.999), wd=0, grad clip 1.0, EMA 0.99, -real-score CFG strength 3, no adversarial loss (the reference -"discriminator" is the fake score network). Set -``config.model.pretrained_student_net_path`` to a Stage 2 checkpoint from -``config_anyflow.py`` before launching, and point the teacher -(``config.model.pretrained_model_path``) at a flow-map teacher checkpoint — -the reference initializes both real and fake score from a separately -fine-tuned flow-map teacher, not stock Wan2.1. - -Known deviations from the reference, both documented: -* full-rank fine-tuning instead of the paper's rank-256 LoRA (FastGen has no - PEFT path today); -* the fake-score noising time is shifted-uniform (FastGen's sampler) instead - of shifted-logit-normal(0, 1). +"""AnyFlow on-policy distillation config on Wan-1.3B T2V (paper Stage 2). + +DMD2-style distribution matching with a flow-map student that generates via a +compressed rollout (jump -> fine step -> jump, gradient through all segments, +NFE sampled per iteration) and co-trains the Stage-1 flow-map loss at every +student update — see ``AnyFlowModel``. The values below mirror the reference +recipe ``train_wan1b_onpolicy_81f_480p_lr2e-6_1k_b32.yml``, which runs 1200 +generator updates. + +Two checkpoints must be supplied. The student is seeded from a Stage 1 trainer +checkpoint, read directly with no conversion; ``pretrained_ckpt_key_map`` below + takes it from the Stage 1 ``ema`` weights, matching the reference's +``pretrained_weight: ema``:: + + trainer.checkpointer.pretrained_ckpt_path=/checkpoints/0006000.pth + +Known deviation from the reference: full-rank fine-tuning instead of the paper's +rank-256 LoRA. """ import copy @@ -44,15 +38,33 @@ def create_config(): config.model.net.time_cond_type = "abs" config.model.net.r_embedder_fusion = "gated" config.model.net.r_embedder_gate_value = 0.25 + # Full [0, 1] noise schedule, as in the pretrain stage + config.model.net.min_t = 0.0 + config.model.net.max_t = 1.0 - config.model.pretrained_student_net_path = "" + # ------ teacher / fake score: PLAIN single-timestep Wan ------ + # Neither consumes r (DMD2 queries both at the instantaneous velocity r = t) + # and the reference builds them without the flow-map pathway at all. + config.model.teacher = copy.deepcopy(Wan_1_3B_Config) + config.model.teacher.r_timestep = False + + # Student init comes from the Stage 1 trainer checkpoint via + # trainer.checkpointer.pretrained_ckpt_path (see the module docstring); + # "ema" mirrors the reference's `pretrained_weight: ema`. + config.trainer.checkpointer.pretrained_ckpt_key_map = {"net": "ema", "ema": "ema"} config.model.precision = "bfloat16" + # FSDP2 parameter storage and gradient reduction in fp32 while compute stays + # bfloat16 -- the same split as the reference. Takes + # effect only under FSDP (`trainer.ddp=False`); it is ignored under DDP, where + # params, grads and compute are all `precision`. + config.model.precision_fsdp = "float32" + # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. config.model.input_shape = [16, 21, 60, 104] # ------ DMD machinery ------ - # The reference Stage 3 has no adversarial loss: its "discriminator" is + # The reference Stage 2 has no adversarial loss: its "discriminator" is # the fake score network, trained with denoising score matching only. config.model.gan_loss_weight_gen = 0.0 # The reference updates generator and fake score 1:1 (both every global @@ -62,21 +74,34 @@ def create_config(): # FastGen's CFG formula is cond + (g-1)*(cond - uncond), so g=4. config.model.guidance_scale = 4.0 + # DMD gradient noising time: reference `generator_loss` draws torch.rand + # then applies the shift -> shifted-uniform. config.model.sample_t_cfg.time_dist_type = "shifted" config.model.sample_t_cfg.shift = 5.0 - config.model.sample_t_cfg.min_t = 0.001 - config.model.sample_t_cfg.max_t = 0.999 + config.model.sample_t_cfg.min_t = 0.0 + config.model.sample_t_cfg.max_t = 1.0 + + # Fake-score noising time: reference `discriminator_loss` draws + # logit_normal(0, 1) then applies the same shift. This is a DIFFERENT + # density from the DMD path above, so it needs its own config. + config.model.fake_score_sample_t_cfg = copy.deepcopy(config.model.sample_t_cfg) + config.model.fake_score_sample_t_cfg.time_dist_type = "shifted_logitnormal" + config.model.fake_score_sample_t_cfg.train_p_mean = 0.0 + config.model.fake_score_sample_t_cfg.train_p_std = 1.0 # ------ student rollout (reference rollout_cfg) ------ config.model.student_sample_type = "ode" - config.model.student_sample_steps = 4 # used for validation sampling config.model.student_sample_steps_list = [2, 4, 8, 16, 50] - config.model.sample_t_cfg.t_list = None # rollout schedules are computed per NFE + config.model.student_sample_steps = 4 + # Validation schedule for `student_sample_steps = 4`: `shift * s / (1 + (shift + # - 1) * s)` on `s = linspace(1, 0, 5)` -- the same grid the per-NFE rollout + # builds. Recompute if `shift`, `max_t` or `student_sample_steps` changes. + config.model.sample_t_cfg.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625, 0.0] - # ------ co-trained Stage-2 flow-map loss (reference cotrain_forward_kl) ------ + # ------ co-trained Stage-1 flow-map loss (reference cotrain_forward_kl) ------ # FastGen's VSD loss carries a 0.5 factor the reference's DMD loss does # not; 0.5 here keeps the DMD : flow-map gradient ratio at the - # reference's 1 : 1 (the common 0.5 scale is absorbed by Adam). + # reference's 1 : 1. config.model.cotrain_pretrain_weight = 0.5 config.model.loss_config.use_cd = False config.model.loss_config.loss_type = "l2" @@ -84,12 +109,13 @@ def create_config(): config.model.loss_config.norm_method = None config.model.loss_config.use_jvp_finite_diff = True config.model.loss_config.jvp_finite_diff_eps = 5e-3 - config.model.loss_config.rebalance_to_diffusion = True + config.model.loss_config.rebalance_to_flow_matching = True config.model.loss_config.guidance_fuse_scale = 3.0 config.model.cond_dropout_prob = 0.1 config.model.precision_amp_jvp = "float32" - config.model.sample_t_cfg.r_sample_ratio = 0.5 + config.model.sample_t_cfg.flow_matching_ratio = 0.5 config.model.sample_t_cfg.consistency_ratio = 0.25 + config.model.sample_t_cfg.deterministic_buckets = True # ------ optimization (reference: AdamW lr=2e-6, betas=(0.0, 0.999), wd=0, # grad clip 1.0, EMA 0.99) ------ @@ -101,6 +127,12 @@ def create_config(): config.model.fake_score_optimizer.weight_decay = 0.0 config.trainer.callbacks.grad_clip.grad_norm = 1.0 config.trainer.callbacks.ema.beta = 0.99 + # EMA start = 6399 = 6000 + 400 - 1: a 400-iteration warmup (the reference's + # `ema_warmup_step: 200` x `student_update_freq = 2`) offset by the Stage-1 + # checkpoint iteration this stage seeds from -- `EMACallback` adds + # `model.resume_iter`, so a bare 399 is already past and never warms up. + # Seeding from another checkpoint: use + 399. + config.trainer.callbacks.ema.start_iter = 6399 # ------ data / trainer ------ config.dataloader_train = VideoLoaderConfig @@ -108,9 +140,13 @@ def create_config(): config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 config.dataloader_train.batch_size = 1 - config.trainer.max_iter = 1200 + # The reference runs 1200 global steps, each performing BOTH a generator + # and a fake-score update. FastGen alternates them across iterations + # (student_update_freq=2), so 1200 generator updates need 2 * 1200 = 2400 + # iterations here. + config.trainer.max_iter = 2400 config.trainer.logging_iter = 100 - config.trainer.save_ckpt_iter = 200 + config.trainer.save_ckpt_iter = 400 config.trainer.batch_size_global = 32 config.log_config.group = "wan_anyflow_onpolicy" diff --git a/fastgen/configs/experiments/WanT2V/config_mf.py b/fastgen/configs/experiments/WanT2V/config_mf.py index be2e26d..9d4f1c2 100644 --- a/fastgen/configs/experiments/WanT2V/config_mf.py +++ b/fastgen/configs/experiments/WanT2V/config_mf.py @@ -58,10 +58,17 @@ def create_config(): config.model.net.r_embedder_init = "zero" config.model.net.norm_temb = False + # The consistency-distillation teacher (use_cd=True) is only ever queried for + # the instantaneous velocity, without r, so give it the plain architecture. + # Deriving it from `net` would allocate, zero-init and keep on device an + # r_embedder that is never invoked. + config.model.teacher = copy.deepcopy(Wan_1_3B_Config) + config.model.teacher.r_timestep = False + # we use simple diffusion version: 0.73 = 0.5 * log((21 * 60 * 104)/(64 * 64)) - 0.1 config.model.enable_preprocessors = False config.model.sample_t_cfg.time_dist_type = "logitnormal" - config.model.sample_t_cfg.r_sample_ratio = 1.0 + config.model.sample_t_cfg.flow_matching_ratio = 0.0 config.model.sample_t_cfg.train_p_mean = -0.8 config.model.sample_t_cfg.train_p_std = 1.6 config.model.sample_t_cfg.min_t = 0.001 diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py index 61671d3..10dc755 100644 --- a/fastgen/configs/methods/config_anyflow.py +++ b/fastgen/configs/methods/config_anyflow.py @@ -1,16 +1,18 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Config schema for the AnyFlow on-policy method (paper Stage 3). +"""Config schema for the AnyFlow on-policy method (paper Stage 2). AnyFlow's on-policy stage is DMD2 with a flow-map student that generates via -a multi-step rollout-with-gradient, so the config is the DMD2 model config -unchanged. The flow-map pretrain stage (paper Stage 2) is MeanFlow with +a multi-step rollout-with-gradient, so the config is based on the DMD2 model +config. The flow-map pretrain stage (paper Stage 1) is MeanFlow with AnyFlow's hyperparameters — see ``configs/experiments/WanT2V/config_anyflow.py``. """ from typing import List, Optional +from typing import Optional + import attrs from omegaconf import DictConfig @@ -26,7 +28,7 @@ from fastgen.configs.methods.config_dmd2 import ModelConfig as DMD2ModelConfig from fastgen.configs.methods.config_mean_flow import ( LossConfig as MeanFlowLossConfig, - SampleRConfig, + SampleRConfig as MeanFlowSampleRConfig, SampleTConfig as MeanFlowSampleTConfig, ) from fastgen.methods import AnyFlowModel @@ -37,7 +39,7 @@ class ModelConfig(DMD2ModelConfig): """AnyFlow on-policy model config — DMD2 plus the rollout / cotrain knobs. - The MeanFlow loss / sampling configs drive the co-trained Stage-2 + The MeanFlow loss / sampling configs drive the co-trained Stage-1 flow-map loss inside the student update (the reference's ``cotrain_forward_kl``). """ @@ -45,10 +47,10 @@ class ModelConfig(DMD2ModelConfig): # MeanFlow-style (t, r) sampling for the co-trained flow-map loss; the # extra fields are ignored by the DMD2 noising-time sampling. sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) - sample_r_cfg: SampleRConfig = attrs.field(factory=SampleRConfig) + sample_r_cfg: MeanFlowSampleRConfig = attrs.field(factory=MeanFlowSampleRConfig) loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig) - # Weight of the co-trained Stage-2 flow-map loss in the student update. + # Weight of the co-trained Stage-1 flow-map loss in the student update. # The reference runs it at weight 1 (cotrain_forward_kl: True); 0 disables. cotrain_pretrain_weight: float = 1.0 diff --git a/fastgen/configs/methods/config_dmd2.py b/fastgen/configs/methods/config_dmd2.py index a59af1e..6f5927a 100644 --- a/fastgen/configs/methods/config_dmd2.py +++ b/fastgen/configs/methods/config_dmd2.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 import copy +from typing import Optional + import attrs from omegaconf import DictConfig @@ -9,6 +11,7 @@ from fastgen.configs.config import ( BaseModelConfig, BaseConfig, + SampleTConfig, ) from fastgen.configs.opt import BaseOptimizerConfig, BaseSchedulerConfig from fastgen.methods import DMD2Model @@ -33,6 +36,14 @@ class ModelConfig(BaseModelConfig): discriminator_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseOptimizerConfig)) discriminator_scheduler: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseSchedulerConfig)) + # Optional noising-time distribution used ONLY on fake-score update + # iterations (`iteration % student_update_freq != 0`). DMD2 alternates the + # generator and fake-score updates, and the two do not have to noise from the + # same density -- AnyFlow's Stage 2, for instance, uses shifted-uniform for the + # DMD gradient and shifted-logit-normal for the fake-score target. `None` + # reuses `sample_t_cfg` for both, which is the historical behaviour. + fake_score_sample_t_cfg: Optional[SampleTConfig] = None + # student update frequency student_update_freq: int = 5 diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index 6541f20..9493541 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -33,13 +33,19 @@ class SampleTConfig(BaseSampleTConfig): train_p_mean: float = -1.1 train_p_std: float = 2.0 - # ratio for randomly sampling r - r_sample_ratio: float = 0.0 + # fraction of the batch with r = t (pure flow matching); the rest keeps + # the randomly sampled r, minus the consistency fraction below + flow_matching_ratio: float = 1.0 - # fraction of the batch forced to r = min_t (consistency-to-clean, - # used by AnyFlow; 0.0 keeps the original MeanFlow behavior) + # fraction of the batch forced to r = 0 (consistency-to-clean, used by + # AnyFlow; 0.0 keeps the original MeanFlow behavior) consistency_ratio: float = 0.0 + # how samples are assigned to the two buckets above: False draws each + # sample's bucket independently, True partitions the global batch by rank + # index so the bucket sizes are exact every iteration (used by AnyFlow) + deterministic_buckets: bool = False + @attrs.define(slots=False) class SampleRConfig(BaseSampleTConfig): @@ -86,9 +92,9 @@ class LossConfig: # plain text dropout (cond_dropout_prob). None keeps MeanFlow's # target-side guidance (guidance_scale, eq. 19). guidance_fuse_scale: Optional[float] = None - # rebalance non-diffusion (r < t) sample losses to the global - # diffusion-loss mean via a detached per-sample factor (AnyFlow) - rebalance_to_diffusion: bool = False + # rebalance the flow-map / consistency (r < t) sample losses to the global + # flow-matching (r = t) loss mean via a detached per-sample factor (AnyFlow) + rebalance_to_flow_matching: bool = False @attrs.define(slots=False) diff --git a/fastgen/methods/README.md b/fastgen/methods/README.md index 3afb26f..f9770e8 100644 --- a/fastgen/methods/README.md +++ b/fastgen/methods/README.md @@ -14,6 +14,7 @@ Training methods for fast single-step or few-step generation from diffusion mode | | | f-distill | [`FdistillModel`](distribution_matching/f_distill.py) | f-divergence weighted DMD2 | [Xu et al., 2025](https://arxiv.org/abs/2502.15681) | | | | LADD | [`LADDModel`](distribution_matching/ladd.py) | Pure adversarial distillation | [Sauer et al., 2024](https://arxiv.org/abs/2403.12015) | | | | CausVid | [`CausVidModel`](distribution_matching/causvid.py) | Causal DMD2 with diffusion forcing | [Yin et al., 2024](https://arxiv.org/abs/2412.07772) | +| | | AnyFlow | [`AnyFlowModel`](distribution_matching/anyflow.py) | Any-step DMD2 with a flow-map student | [Gu et al., 2026](https://arxiv.org/abs/2605.13724) | | | | Self-Forcing | [`SelfForcingModel`](distribution_matching/self_forcing.py) | Causal DMD2 with self-forcing | [Huang et al., 2025](https://arxiv.org/abs/2506.08009) | | **Fine-Tuning** | [README](fine_tuning/README.md) | SFT | [`SFTModel`](fine_tuning/sft.py) | Finetuning with denoising score matching | [Ho et al., 2020](https://arxiv.org/abs/2006.11239), [Song et al., 2020](https://arxiv.org/abs/2011.13456), [Lipman et al., 2022](https://arxiv.org/abs/2210.02747), [Albergo et al., 2023](https://arxiv.org/abs/2303.08797) | | | | CausalSFT | [`CausalSFTModel`](fine_tuning/sft.py) | Causal version of SFT | [Chen et al., 2024](https://arxiv.org/abs/2407.01392) | diff --git a/fastgen/methods/consistency_model/README.md b/fastgen/methods/consistency_model/README.md index 867dfa6..f40863f 100644 --- a/fastgen/methods/consistency_model/README.md +++ b/fastgen/methods/consistency_model/README.md @@ -81,7 +81,7 @@ Learns average velocity between trajectory points: `x_r = x_t - (t-r) · u(x_t, - `loss_config.use_cd`: Use consistency distillation (requires teacher; `guidance_scale` controls CFG for the teacher) - `loss_config.use_jvp_finite_diff`: Use finite difference for JVP (e.g., for compatibility with Flash Attention and FSDP) - `sample_t_cfg`, `sample_r_cfg`: Configs of the distributions for sampling `t` and `r` -- `sample_t_cfg.r_sample_ratio`: Ratio for flow matching loss +- `sample_t_cfg.flow_matching_ratio`: Fraction of the batch with `r = t` (flow matching loss) **Configs:** [`EDM/config_mf_cifar10.py`](../../configs/experiments/EDM/config_mf_cifar10.py), [`DiT/config_mf_b.py`](../../configs/experiments/DiT/config_mf_b.py), [`DiT/config_mf_xl.py`](../../configs/experiments/DiT/config_mf_xl.py), [`WanT2V/config_mf.py`](../../configs/experiments/WanT2V/config_mf.py) diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 98cdd3f..0cbd4e6 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -49,15 +49,23 @@ def temp_disable_efficient_attn(device_type: str = "cuda"): yield -class MeanFlowModel(CMModel): - def __init__(self, config: ModelConfig): - """ - - Args: - config (ModelConfig): The configuration for the MeanFlow model - """ - super().__init__(config) - self.config = config +class FlowMapLossMixin: + """The MeanFlow flow-map regression objective, shared across methods. + + Holds everything needed to turn a real batch into the flow-map loss + ``|| u_theta(x_t, t, r) - sg(v - (t - r) du/dt) ||^2``: the (t, r) sampler + with its flow-matching / consistency buckets, the JVP estimate, the loss + weighting, and the reduction. ``MeanFlowModel`` uses it as its whole + training objective; distribution-matching methods can co-train it alongside + their own objective. + + The host class must provide ``net``, ``device``, ``config``, + ``precision_amp``, a ``_get_velocity`` implementation, and call + ``_init_flow_map_loss`` from its ``__init__``. + """ + + def _init_flow_map_loss(self) -> None: + """Bind the flow-map loss / sampling configs and the JVP precision.""" self.sample_t_cfg = self.config.sample_t_cfg self.sample_r_cfg = self.config.sample_r_cfg self.loss_config = self.config.loss_config @@ -69,120 +77,54 @@ def __init__(self, config: ModelConfig): self.precision_amp_jvp = basic_utils.PRECISION_MAP[self.config.precision_amp_jvp] logger.critical(f"Using precision {self.precision_amp_jvp} for JVP") - # Normalization constant for the fixed per-timestep loss weighting - # (computed lazily on first use, see _get_timestep_weight). + # The fixed per-timestep loss weight is normalized to mean one over the + # network's discrete training timesteps (t = 0 excluded), mapped through + # the same shift the samplers use. The constant depends only on the + # config, so derive it once here. self._timestep_weight_scale: Optional[float] = None + if self.loss_config.weight_type is not None: + num_steps = self.net.noise_scheduler.num_steps + shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 + grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64)[:-1] + grid = shift * grid / (1 + (shift - 1) * grid) + self._timestep_weight_scale = float(num_steps / self._timestep_weight_raw(grid).sum()) - def _mix_condition( - self, - condition: Any, - neg_condition: torch.Tensor, - dxt_dt: torch.Tensor, - guided_dxt_dt: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.config.cond_dropout_prob is None: - return condition, dxt_dt - - batch_size = dxt_dt.shape[0] - # Decide how many to drop first. - num_to_drop = (torch.rand(batch_size, device=dxt_dt.device) < self.config.cond_dropout_prob).sum() - # Create the mask to keep all but the first samples (the order is important to ensures most dropout happens at flow matching loss) - mask = torch.arange(batch_size, device=dxt_dt.device) >= num_to_drop - dxt_dt = torch.where(expand_like(mask, dxt_dt), guided_dxt_dt, dxt_dt) - - if isinstance(condition, torch.Tensor): - condition = torch.where(expand_like(mask, condition), condition, neg_condition) - elif isinstance(condition, dict): - condition = condition.copy() - keys_no_drop = self.config.cond_keys_no_dropout - assert set(keys_no_drop).issubset( - condition.keys() - ), f"keys_no_drop: {keys_no_drop} not in {condition.keys()}" - for k in condition.keys() - keys_no_drop: - condition[k] = torch.where(expand_like(mask, condition[k]), condition[k], neg_condition[k]) - else: - raise TypeError(f"Unsupported type: {type(condition)}") - - return condition, dxt_dt + def _drop_condition(self, condition: Any, neg_condition: Any) -> Tuple[Any, Optional[torch.Tensor]]: + """Replace the condition with neg_condition for a per-sample subset. - def _drop_condition(self, condition: Any, neg_condition: Any) -> Any: - """Replace the condition with neg_condition for a random per-sample subset. + Returns ``(condition, keep)``; ``keep`` is the ``[B]`` bool mask (None if + no dropout), so callers can reuse the same subset. Keys in + ``cond_keys_no_dropout`` are never replaced. - Plain text dropout (used by the prediction-side guidance fusion); the - dropped samples then train the unconditional pathway. + ``deterministic_buckets`` decides whether an index carries bucket + information: if so the buckets are cut on the GLOBAL index and an + index-based rule would only hit flow matching on rank 0, so draw per + sample; otherwise drop the first ``num_to_drop``. """ if self.config.cond_dropout_prob is None or neg_condition is None: - return condition + return condition, None + ref = neg_condition if isinstance(neg_condition, torch.Tensor) else next(iter(neg_condition.values())) - keep = torch.rand(ref.shape[0], device=ref.device) >= self.config.cond_dropout_prob + batch_size = ref.shape[0] + if self.sample_t_cfg.deterministic_buckets: + keep = torch.rand(batch_size, device=ref.device) >= self.config.cond_dropout_prob + else: + num_to_drop = (torch.rand(batch_size, device=ref.device) < self.config.cond_dropout_prob).sum() + keep = torch.arange(batch_size, device=ref.device) >= num_to_drop + if isinstance(condition, torch.Tensor): - return torch.where(expand_like(keep, condition), condition, neg_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()}" condition = condition.copy() - for k in condition.keys() - set(self.config.cond_keys_no_dropout): + for k in condition.keys() - keys_no_drop: condition[k] = torch.where(expand_like(keep, condition[k]), condition[k], neg_condition[k]) - return condition + return condition, keep raise TypeError(f"Unsupported condition type: {type(condition)}") - @torch.no_grad() - def _get_velocity( - self, - x: torch.Tensor, - z: torch.Tensor, - t: torch.Tensor, - condition: Optional[torch.Tensor] = None, - neg_condition: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - x_t = self.net.noise_scheduler.forward_process(x, z, t) - - if self.loss_config.use_cd: - dxt_dt = self.teacher(x_t, t, condition=condition, fwd_pred_type="flow") - if self.config.guidance_scale is not None: - guidance_scale = torch.where( - ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - self.config.guidance_scale, - 1.0, - ) - guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) - neg_dxt_dt = self.teacher(x_t, t, condition=neg_condition, fwd_pred_type="flow") - dxt_dt = dxt_dt + (guidance_scale - 1.0) * (dxt_dt - neg_dxt_dt) - else: - dxt_dt = self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) - - # unconditional score estimation from meanflow eq (19) - if self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None: - # Turn off dropout - self.net.eval() - neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow") - guidance_scale = self.config.guidance_scale or 1.0 - guidance_scale = torch.where( - ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - guidance_scale, - 1.0, - ) - guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) - - if self.config.guidance_mixture_ratio is None: - guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt) - else: - guidance_mixture_ratio = torch.where( - ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - self.config.guidance_mixture_ratio, - 0.0, - ) - guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype) - cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow") - guided_dxt_dt = ( - guidance_scale * dxt_dt - + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt - + guidance_mixture_ratio * cond_dxt_dt - ) - - self.net.train() - condition, dxt_dt = self._mix_condition(condition, neg_condition, dxt_dt, guided_dxt_dt) - - return condition, dxt_dt - def _estimate_jvp_finite_difference( self, net_wrapper: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor], @@ -275,21 +217,32 @@ def net_wrapper(x_t, t, r): return u_theta_jvp + def _warn_on_degenerate_buckets(self, n_flow_matching: int, n_consistency: int, global_bsz: int) -> None: + """Warn once if a requested bucket rounds away under the deterministic partition.""" + requested_but_empty = (self.sample_t_cfg.flow_matching_ratio > 0 and n_flow_matching == 0) or ( + self.sample_t_cfg.consistency_ratio > 0 and n_consistency == 0 + ) + if not requested_but_empty or getattr(self, "_bucket_warned", False): + return + self._bucket_warned = True + logger.warning( + f"The deterministic (t, r) bucket partition is degenerate: with " + f"world_size * batch_size = {global_bsz}, flow_matching_ratio=" + f"{self.sample_t_cfg.flow_matching_ratio} and consistency_ratio=" + f"{self.sample_t_cfg.consistency_ratio} yield bucket sizes " + f"({n_flow_matching}, {n_consistency}). The partition spans ranks but not " + f"gradient-accumulation rounds, so empty buckets stay empty every iteration." + ) + def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Sample (t, r) with t >= r and assign the per-batch buckets. Returns ``(t, r, r_eq_t_mask)`` where ``r_eq_t_mask`` marks the samples with ``r = t`` (pure flow matching). - With ``consistency_ratio == 0`` (MeanFlow default) the flow-matching - head size is stochastic, as before. With ``consistency_ratio > 0`` - (AnyFlow) the buckets are assigned deterministically over the GLOBAL - batch by rank index, exactly like the AnyFlow reference - (``sample_timestep`` in ``trainer_wan_anyflow_pretrain.py``): the - first ``round((1 - r_sample_ratio) * global_bsz)`` samples get - ``r = t``, the next ``round(consistency_ratio * global_bsz)`` get - ``r = 0`` (consistency to clean data), and the rest keep the sampled - random pair. + The batch splits three ways: a ``flow_matching_ratio`` fraction gets + ``r = t``, a ``consistency_ratio`` fraction gets ``r = 0`` + (consistency to clean data), and the rest keep the sampled random pair. """ t_sample_kwargs = convert_cfg_to_dict(self.sample_t_cfg) t = self.net.noise_scheduler.sample_t(batch_size, **t_sample_kwargs, device=self.device) @@ -298,53 +251,53 @@ def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tens t, r = torch.maximum(t, r), torch.minimum(t, r) assert torch.all(t >= r), "r cannot be larger than t" - if self.sample_t_cfg.consistency_ratio > 0: + flow_matching_ratio = self.sample_t_cfg.flow_matching_ratio + consistency_ratio = self.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}" + + # Both policies produce two bucket sizes and index the batch the same + # way; they differ only in how the sizes are obtained. + if self.sample_t_cfg.deterministic_buckets: global_bsz = world_size() * batch_size - n_flow_matching = round((1.0 - self.sample_t_cfg.r_sample_ratio) * global_bsz) - n_consistency = round(self.sample_t_cfg.consistency_ratio * global_bsz) - if (n_flow_matching == 0 or n_consistency == 0) and not getattr(self, "_bucket_warned", False): - self._bucket_warned = True - logger.warning( - f"The deterministic (t, r) bucket partition is degenerate: with " - f"world_size * batch_size = {global_bsz}, r_sample_ratio=" - f"{self.sample_t_cfg.r_sample_ratio} and consistency_ratio=" - f"{self.sample_t_cfg.consistency_ratio} yield bucket sizes " - f"({n_flow_matching}, {n_consistency}). The partition spans ranks but not " - f"gradient-accumulation rounds (as in the AnyFlow reference), so empty " - f"buckets stay empty every iteration." - ) - global_idx = get_rank() * batch_size + torch.arange(batch_size, device=self.device) - r_eq_t_mask = global_idx < n_flow_matching - is_consistency = (global_idx >= n_flow_matching) & (global_idx < n_flow_matching + n_consistency) - r = torch.where(r_eq_t_mask, t, r) - r = torch.where(is_consistency, torch.zeros_like(r), r) + n_flow_matching = round(flow_matching_ratio * global_bsz) + n_consistency = round(consistency_ratio * global_bsz) + self._warn_on_degenerate_buckets(n_flow_matching, n_consistency, global_bsz) + position = get_rank() * batch_size + torch.arange(batch_size, device=self.device) else: - # set t=r (flow matching loss) for a subset of the batch - flow_matching_size = (torch.rand(batch_size, device=self.device) >= self.sample_t_cfg.r_sample_ratio).sum() - r_eq_t_mask = torch.arange(batch_size, device=self.device) < flow_matching_size - r = torch.where(r_eq_t_mask, t, r) + # Both sizes come from one uniform draw, cut at either end of [0, 1). + # The two events are disjoint by the assert above, so the buckets + # cannot overlap. + uniform = torch.rand(batch_size, device=self.device) + n_flow_matching = (uniform >= 1.0 - flow_matching_ratio).sum() + n_consistency = (uniform < consistency_ratio).sum() + position = torch.arange(batch_size, device=self.device) + + r_eq_t_mask = position < n_flow_matching + is_consistency = (position >= n_flow_matching) & (position < n_flow_matching + n_consistency) + r = torch.where(r_eq_t_mask, t, r) + r = torch.where(is_consistency, torch.zeros_like(r), r) return t, r, r_eq_t_mask def _reduce_mf_loss(self, mf_loss: torch.Tensor, r_eq_t_mask: torch.Tensor) -> torch.Tensor: """Reduce the per-sample loss to a scalar, optionally rebalancing. - With ``loss_config.rebalance_to_diffusion`` set (AnyFlow), every - flow-map / consistency (r < t) sample's loss is multiplied by the - detached factor ``mean(global flow-matching losses) / (own loss + 1e-5)``, - matching the AnyFlow reference (``train_bidirection``): flow-map / - consistency gradients are self-normalized and rescaled to the global + With ``loss_config.rebalance_to_flow_matching`` set, every flow-map / + consistency (r < t) sample's loss is multiplied by the detached factor + ``mean(global flow-matching losses) / (own loss + 1e-5)``, so those + gradients are self-normalized and rescaled to the global flow-matching-loss mean. """ # No rank-local condition here: the branch must be taken (or not) by # every rank so the collective below cannot deadlock. Applying the # scale to an empty ~r_eq_t_mask selection is a no-op. - if getattr(self.loss_config, "rebalance_to_diffusion", False): + if getattr(self.loss_config, "rebalance_to_flow_matching", False): with torch.no_grad(): # The global flow-matching-loss mean only needs the global sum # and count, so reduce two scalars instead of gathering the - # per-sample losses (equivalent to the reference's - # cat(all_gather(loss))[mask].mean(), and independent of the + # per-sample losses (equivalent, and independent of the # per-rank batch sizes). fm_loss_sum = torch.where(r_eq_t_mask, mf_loss, torch.zeros_like(mf_loss)).sum() fm_count = r_eq_t_mask.sum().to(mf_loss.dtype) @@ -369,23 +322,6 @@ def _timestep_weight_raw(self, t: torch.Tensor) -> torch.Tensor: return torch.ones_like(t) raise ValueError(f"Invalid weight_type: {weight_type!r}") - @torch.no_grad() - def _get_timestep_weight(self, t: torch.Tensor) -> torch.Tensor: - """Fixed per-timestep loss weight w(t) (AnyFlow-style). - - The weight is evaluated directly as a function of t. The normalization - constant matches the AnyFlow reference scheduler, which normalizes the - weights over its shifted discrete timestep grid (``set_timesteps``: - 1000 points excluding t=0); we compute the same constant once and - cache it. - """ - if self._timestep_weight_scale is None: - shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 - grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64)[:-1] - grid = shift * grid / (1 + (shift - 1) * grid) - self._timestep_weight_scale = float(1000.0 / self._timestep_weight_raw(grid).sum()) - return self._timestep_weight_raw(t) * self._timestep_weight_scale - @torch.no_grad() def _compute_weight(self, tensor: torch.Tensor, t: torch.Tensor) -> torch.Tensor: """Per-sample loss weight: the adaptive normalization (``norm_method``, @@ -409,7 +345,7 @@ def _compute_weight(self, tensor: torch.Tensor, t: torch.Tensor) -> torch.Tensor raise ValueError(f"Invalid norm method: {self.loss_config.norm_method}") if self.loss_config.weight_type is not None: - weight = weight * self._get_timestep_weight(t) + weight = weight * (self._timestep_weight_raw(t) * self._timestep_weight_scale) assert ( weight.shape == tensor.shape @@ -453,7 +389,8 @@ def _mf_pred_to_loss( loss = (u_theta - tangent).pow(2) if self.loss_config.norm_method is None: # Without adaptive normalization the reduction matters: the - # per-element mean matches the AnyFlow reference loss scale. + # 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))) @@ -554,10 +491,9 @@ def _compute_mf_loss( guidance_fuse_scale = getattr(self.loss_config, "guidance_fuse_scale", None) if guidance_fuse_scale is not None: - # AnyFlow-style guidance distillation: the guidance is fused on the - # PREDICTION side — the conditional output is trained to be the - # guided flow directly. Matches the reference ``train_bidirection``: - # the regression target stays the raw data velocity, the + # Guidance distillation fused on the PREDICTION side: the + # conditional output is trained to be the guided flow directly. + # The regression target stays the raw data velocity, the # unconditional branch is queried at the SAME (t, r) flow-map # slice, and the effective prediction is # (u_cond + (g - 1) * u_uncond) / g. Text dropout replaces the @@ -567,12 +503,12 @@ def _compute_mf_loss( assert ( neg_condition is not None ), "guidance_fuse_scale requires neg_condition: the unconditional branch is queried at the same (t, r)" - condition = self._drop_condition(condition, neg_condition) + condition, _ = self._drop_condition(condition, neg_condition) dxt_dt = self.net.noise_scheduler.cond_velocity(x=real_data, eps=z, t=t) torch.clear_autocast_cache() # dF/dt of the fused prediction: finite difference of the # conditional output divided by g (the unconditional derivative is - # dropped, as in the reference's compute_central_difference). + # dropped). u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) / g assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" @@ -616,6 +552,81 @@ def _compute_mf_loss( warmup_weight, ) + +class MeanFlowModel(FlowMapLossMixin, CMModel): + def __init__(self, config: ModelConfig): + """ + + Args: + config (ModelConfig): The configuration for the MeanFlow model + """ + super().__init__(config) + self.config = config + self._init_flow_map_loss() + + @torch.no_grad() + def _get_velocity( + self, + x: torch.Tensor, + z: torch.Tensor, + t: torch.Tensor, + condition: Optional[torch.Tensor] = None, + neg_condition: Optional[torch.Tensor] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + x_t = self.net.noise_scheduler.forward_process(x, z, t) + + if self.loss_config.use_cd: + dxt_dt = self.teacher(x_t, t, condition=condition, fwd_pred_type="flow") + if self.config.guidance_scale is not None: + guidance_scale = torch.where( + ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), + self.config.guidance_scale, + 1.0, + ) + guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) + neg_dxt_dt = self.teacher(x_t, t, condition=neg_condition, fwd_pred_type="flow") + dxt_dt = dxt_dt + (guidance_scale - 1.0) * (dxt_dt - neg_dxt_dt) + else: + dxt_dt = self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) + + # unconditional score estimation from meanflow eq (19) + if self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None: + # Turn off dropout + self.net.eval() + neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow") + guidance_scale = self.config.guidance_scale or 1.0 + guidance_scale = torch.where( + ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), + guidance_scale, + 1.0, + ) + guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) + + if self.config.guidance_mixture_ratio is None: + guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt) + else: + guidance_mixture_ratio = torch.where( + ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), + self.config.guidance_mixture_ratio, + 0.0, + ) + guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype) + cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow") + guided_dxt_dt = ( + guidance_scale * dxt_dt + + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt + + guidance_mixture_ratio * cond_dxt_dt + ) + + self.net.train() + condition, keep = self._drop_condition(condition, neg_condition) + if keep is not None: + # Same subset: a kept sample is conditional + guided, a + # dropped one unconditional + unguided. + dxt_dt = torch.where(expand_like(keep, dxt_dt), guided_dxt_dt, dxt_dt) + + return condition, dxt_dt + def single_train_step( self, data: Dict[str, Any], iteration: int ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index 29abe84..41e860a 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -87,34 +87,25 @@ DMD2 extended for causal video generation with autoregressive chunk-by-chunk pro ## AnyFlow -**File:** [`anyflow.py`](anyflow.py) | **Reference:** AnyFlow — any-step video diffusion framework on flow maps +**File:** [`anyflow.py`](anyflow.py) | **Reference:** [Gu et al., 2026](https://arxiv.org/abs/2605.13724) -Single model that supports arbitrary inference NFE by learning a flow map `u_θ(x_t, t, r)` (average velocity from `t` back to `r`). Trained in two stages: +DMD2 with a flow-map student `u_θ(x_t, t, r)` (average velocity from `t` back to `r`), supporting arbitrary inference NFE. The student generates by rolling out the flow map from pure noise with the NFE sampled per iteration and gradients through all segments, and co-trains the flow-map loss at every update. Requires a Stage-1 flow-map pretrain, which uses the MeanFlow objective with AnyFlow's hyperparameters and runs directly on [`MeanFlowModel`](../consistency_model/mean_flow.py). -1. **Pretrain** — the MeanFlow objective with AnyFlow's hyperparameters, run directly on [`MeanFlowModel`](../consistency_model/mean_flow.py): finite-difference JVP, a fixed `beta08` per-timestep loss weight (`loss_config.weight_type`), prediction-side guidance fusion (`loss_config.guidance_fuse_scale` — the conditional output learns the guided flow directly), global rebalancing of the flow-map / consistency sample losses to the diffusion-loss mean (`loss_config.rebalance_to_diffusion`), shifted timestep sampling, and a `sample_t_cfg.consistency_ratio` fraction of the batch pinned to `r = 0`. There is no AnyFlow-specific pretrain code. -2. **On-policy** — distribution-matching distillation on top of the pretrained flow-map weights ([`anyflow.py`](anyflow.py)). DMD2 with the reference's three deviations: the student generates via a flow-map rollout compressed to at most three forwards — jump `t_0 → t_g`, fine step `t_g → t_{g+1}`, jump to 0 — with the NFE sampled per iteration from `student_sample_steps_list` and gradient through all segments (`gen_data_from_net`); the student always starts from pure noise at `max_t` (`_generate_noise_and_time`); and every student update co-trains the Stage-2 flow-map loss (`cotrain_pretrain_weight`). There is no adversarial loss (the reference "discriminator" is the fake score network). The teacher / fake_score are flow-map networks queried at the instantaneous velocity `r = t` (the network-level default for gated-fusion Wan when `r` is not passed). - -**Key Parameters (on-policy):** -- `student_sample_steps_list`: rollout NFEs sampled per iteration (reference: `[2, 4, 8, 16, 50]`) -- `cotrain_pretrain_weight`: weight of the co-trained Stage-2 flow-map loss (reference: 1) -- `guidance_scale`: real-score CFG; FastGen applies `cond + (g-1)·(cond - uncond)`, so the reference's strength 3 is `g=4` -- See also key parameters of DMD2 above - -**Key Parameters (pretrain, on MeanFlow):** -- `loss_config.weight_type`: `beta08` | `gaussian` | `uniform` fixed per-timestep loss weight -- `loss_config.guidance_fuse_scale` + `cond_dropout_prob`: prediction-side guidance distillation with text dropout -- `loss_config.rebalance_to_diffusion`: rescale non-diffusion losses to the global diffusion-loss mean -- `loss_config.use_jvp_finite_diff` / `loss_config.jvp_finite_diff_eps`: central-difference step δ -- `sample_t_cfg.r_sample_ratio` / `sample_t_cfg.consistency_ratio`: per-batch fraction with sampled `r` / `r = 0` -- `sample_t_cfg.time_dist_type="shifted"`, `sample_t_cfg.shift`: shifted timestep sampling (5.0 for Wan video) - -**Backbone requirement:** the student network must accept a secondary timestep `r` (Wan with `r_timestep=True`). When loading the published AnyFlow HF checkpoints, set `r_embedder_fusion="gated"` and `time_cond_type="abs"` on the Wan constructor — this routes the t/r mix through `Wan/network.py::_fuse_r_embedding`'s gated branch (shared with MeanFlow's additive default) so the released weights reproduce bit-for-bit. `Wan.load_state_dict` remaps the AnyFlow checkpoint layout (`condition_embedder.delta_embedder.*`) automatically; note that loading the HF folder via `model_id_or_local_path` goes through diffusers' `from_pretrained`, which silently drops the `delta_embedder` weights — load the transformer state dict through `Wan.load_state_dict` (or call `remap_anyflow_keys` yourself) instead. +**Key Parameters (Stage 2, on-policy):** +- `student_sample_steps_list`: Rollout NFEs sampled per iteration +- `cotrain_pretrain_weight`: Weight of the co-trained Stage-1 flow-map loss +- `guidance_scale`: CFG scale for teacher, applied as `cond + (g-1)·(cond - uncond)` +- Requires a Stage-1 pretrained checkpoint via `trainer.checkpointer.pretrained_ckpt_path` +- See also the key parameters of DMD2 above -**Note:** correctness of the port is established via forward-parity and single-step training-step parity against the AnyFlow reference (see PR #25 discussion). End-to-end convergence-scale validation on the paper's training corpus is deferred to a follow-up. +**Key Parameters (Stage 1, on MeanFlow):** +- `loss_config.weight_type`: Fixed per-timestep loss weight (`beta08`, `gaussian`, `uniform`) +- `loss_config.guidance_fuse_scale`: Prediction-side guidance fusion, i.e., the conditional output learns the guided flow directly (with `cond_dropout_prob` for text dropout) +- `loss_config.rebalance_to_flow_matching`: Rescale the `r < t` losses to the global flow-matching (`r = t`) loss mean +- `sample_t_cfg`: Shifted timestep sampling (`time_dist_type="shifted"`, `shift=5` for Wan video) and the batch fraction pinned to `r = 0` (`consistency_ratio`) +- See also the key parameters of [MeanFlow](../consistency_model/README.md#meanflow) -**Configs:** -- [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) — Stage 2 pretrain (6k iter, lr=5e-5, shift=5, beta08, paper-aligned) -- [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) — Stage 3 on-policy distillation (1.2k iter, lr=2e-6, no adversarial loss; full-rank fine-tune of a Stage 2 pretrain ckpt — the paper's rank-256 LoRA variant awaits a PEFT path in FastGen) +**Configs:** [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) (Stage 1 pretrain), [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) (Stage 2 on-policy distillation, full-rank fine-tune of a Stage 1 checkpoint) --- diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index 271c7e1..ba61c82 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -3,45 +3,17 @@ """AnyFlow — any-step video diffusion with flow maps and on-policy distillation. -AnyFlow trains a single flow-map model :math:`u_\\theta(x_t, t, r)` that -predicts the average velocity from time ``t`` to ``r`` (with ``r \\le t``). -Once trained, the same model supports arbitrary inference step counts: each -Euler-like sampling step picks its own integration interval -``(t \\rightarrow r)``. - -Training has two stages: - -* **Flow-map pretrain** (paper Stage 2) is MeanFlow with AnyFlow's - hyperparameters — fixed ``beta08`` per-timestep loss weighting, - finite-difference JVP, prediction-side guidance fusion - (``guidance_fuse_scale``), global rebalancing of non-diffusion losses - (``rebalance_to_diffusion``), and a ``consistency_ratio`` fraction of the - batch pinned to ``r = 0``. It is configured directly on - :class:`~fastgen.methods.consistency_model.mean_flow.MeanFlowModel` - (see ``configs/experiments/WanT2V/config_anyflow.py``); there is no - AnyFlow-specific pretrain code. - -* **On-policy** (paper Stage 3, this module) is DMD2 on top of the pretrained - flow-map weights, with the reference's three deviations from stock DMD2: - - - the student generates via a flow-map rollout compressed into at most three - network forwards — one jump from ``t_0`` to ``t_g``, one fine step - ``t_g \\rightarrow t_{g+1}``, one jump to 0 — with the inference step - count sampled per iteration from ``student_sample_steps_list`` and the - fine-step position ``g`` sampled uniformly (both rank-0 broadcast). - Gradient flows through ALL segments - (:meth:`AnyFlowModel.gen_data_from_net`, mirroring - ``WanAnyFlowPipeline.training_rollout``); - - the student always starts from pure noise at ``max_t`` - (:meth:`AnyFlowModel._generate_noise_and_time`); - - every student update co-trains the Stage-2 flow-map loss on the real - batch (``cotrain_pretrain_weight``, mirroring the reference's - ``cotrain_forward_kl``), borrowing MeanFlow's loss machinery. - - The teacher and fake_score are flow-map networks queried at the - instantaneous velocity ``r = t`` — for gated-fusion Wan networks this is - the network-level default when ``r`` is not passed (see - ``fastgen/networks/Wan/network.py``), so DMD2's update steps run unchanged. +AnyFlow trains a single flow-map model ``u_theta(x_t, t, r)`` predicting the +average velocity from ``t`` to ``r`` (``r <= t``), so one model serves any +inference NFE: each Euler-like step picks its own integration interval. + +Stage 1 (flow-map pretrain) is MeanFlow with AnyFlow's hyperparameters, run on +``MeanFlowModel`` directly (``configs/experiments/WanT2V/config_anyflow.py``) — +there is no AnyFlow-specific pretrain code. Stage 2 (this module) is DMD2 on the +pretrained flow-map weights, deviating from stock DMD2 twice: the student +generates via a compressed flow-map rollout with gradient through all segments +(``gen_data_from_net``), and every student update co-trains the Stage-1 flow-map +loss on the real batch (``cotrain_pretrain_weight``). """ from __future__ import annotations @@ -50,10 +22,8 @@ import torch -from fastgen.methods.consistency_model.mean_flow import MeanFlowModel +from fastgen.methods.consistency_model.mean_flow import FlowMapLossMixin from fastgen.methods.distribution_matching.dmd2 import DMD2Model -from fastgen.utils import basic_utils, expand_like -from fastgen.utils.basic_utils import convert_cfg_to_dict from fastgen.utils.distributed import world_size import fastgen.utils.logging_utils as logger @@ -64,36 +34,20 @@ from fastgen.configs.methods.config_anyflow import ModelConfig -class AnyFlowModel(DMD2Model): - """AnyFlow on-policy stage: DMD2 with a flow-map rollout student.""" +class AnyFlowModel(FlowMapLossMixin, DMD2Model): + """AnyFlow on-policy stage: DMD2 with a flow-map rollout student. - # Validation sampling integrates the flow map with r = t_next (ode) just - # like MeanFlow; FastGenModel's default x0-prediction loop never passes r. - _student_sample_loop = MeanFlowModel._student_sample_loop - - # MeanFlow loss machinery borrowed for the co-trained Stage-2 flow-map - # loss (the reference's cotrain_forward_kl runs the full bidirection loss - # at every generator step). - _sample_t_r_buckets = MeanFlowModel._sample_t_r_buckets - _reduce_mf_loss = MeanFlowModel._reduce_mf_loss - _compute_mf_loss = MeanFlowModel._compute_mf_loss - _drop_condition = MeanFlowModel._drop_condition - _jvp = MeanFlowModel._jvp - _estimate_jvp_finite_difference = MeanFlowModel._estimate_jvp_finite_difference - _mf_pred_to_loss = MeanFlowModel._mf_pred_to_loss - _compute_weight = MeanFlowModel._compute_weight - _timestep_weight_raw = MeanFlowModel._timestep_weight_raw - _get_timestep_weight = MeanFlowModel._get_timestep_weight + ``FlowMapLossMixin`` supplies the co-trained Stage-1 objective and the + flow-map validation sample loop, which integrates with ``r = t_next`` — + FastGenModel's default x0-prediction loop never passes ``r``. + """ @torch.no_grad() def _get_velocity(self, x, z, t, condition=None, neg_condition=None): """Raw data velocity for the co-trained flow-map loss. - The DMD teacher CFG (``config.guidance_scale``) must not leak into the - co-trained loss — AnyFlow's guidance lives in - ``loss_config.guidance_fuse_scale`` (prediction-side fusion, handled - inside ``_compute_mf_loss``), so MeanFlow's target-side eq. 19 path is - deliberately not used here. + Note that AnyFlow guides on the prediction side using + ``loss_config.guidance_fuse_scale`` applied in ``_compute_mf_loss``. """ del neg_condition return condition, self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) @@ -101,16 +55,7 @@ def _get_velocity(self, x, z, t, condition=None, neg_condition=None): def __init__(self, config: ModelConfig): super().__init__(config) self.config = config - - # Attributes required by the borrowed MeanFlow methods. - self.sample_t_cfg = self.config.sample_t_cfg - self.sample_r_cfg = self.config.sample_r_cfg - self.loss_config = self.config.loss_config - if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp: - self.precision_amp_jvp = None - else: - self.precision_amp_jvp = basic_utils.PRECISION_MAP[self.config.precision_amp_jvp] - self._timestep_weight_scale = None + self._init_flow_map_loss() logger.info( f"AnyFlow on-policy: student_sample_steps_list={self.config.student_sample_steps_list}, " @@ -123,6 +68,25 @@ def __init__(self, config: ModelConfig): # Rollout # ------------------------------------------------------------------ + def _generate_noise_and_time( + self, real_data: torch.Tensor, iteration: Optional[int] = None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The student rollout always starts from pure noise at ``max_t``.""" + batch_size = real_data.shape[0] + + eps_student = torch.randn(batch_size, *self.input_shape, device=self.device, dtype=real_data.dtype) + t_student = torch.full( + (batch_size,), + self.net.noise_scheduler.max_t, + device=self.device, + dtype=self.net.noise_scheduler.t_precision, + ) + input_student = self.net.noise_scheduler.latents(noise=eps_student) + + t = self._sample_noising_time(batch_size, iteration) + eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) + return input_student, t_student, t, eps + def _broadcast_choice(self, high: int) -> int: """Pick an index in [0, high) on rank 0 and broadcast it.""" idx = torch.randint(0, high, (1,), device=self.device, dtype=torch.long) @@ -131,11 +95,10 @@ def _broadcast_choice(self, high: int) -> int: return int(idx.item()) def _sample_rollout_steps(self) -> int: - """Sample the rollout NFE for this iteration. - - Matches the reference: ``random.choice(num_inference_steps_list)`` - broadcast from rank 0 in both the generator and fake-score updates. - Falls back to the fixed ``student_sample_steps`` when no list is set. + """Sample this iteration's rollout NFE, as the reference does + (``random.choice(num_inference_steps_list)``, rank-0 broadcast in both + the generator and fake-score updates). Falls back to the fixed + ``student_sample_steps`` when no list is set. """ steps_list = self.config.student_sample_steps_list if not steps_list: @@ -145,42 +108,18 @@ def _sample_rollout_steps(self) -> int: def _rollout_t_list(self, num_steps: int) -> torch.Tensor: """Shifted timestep schedule for ``num_steps`` rollout steps. - Equivalent to the reference scheduler's ``set_timesteps``: a uniform - grid mapped through ``shift * s / (1 + (shift - 1) * s)``. The first - entry is clamped to the scheduler's ``max_t``. A config-provided - ``t_list`` takes precedence when its length matches. + Equivalent to the reference scheduler's ``set_timesteps``: a uniform grid + mapped through ``shift * s / (1 + (shift - 1) * s)``, first entry clamped + to ``max_t``. Computed here rather than read from ``sample_t_cfg.t_list`` + because the rollout NFE varies per iteration; that list stays the fixed + *validation* schedule used at ``student_sample_steps``. """ ns = self.net.noise_scheduler - cfg_t_list = self.config.sample_t_cfg.t_list - if cfg_t_list is not None and len(cfg_t_list) == num_steps + 1: - return torch.tensor(cfg_t_list, device=self.device, dtype=ns.t_precision) - shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 s = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64, device=self.device) s = shift * s / (1 + (shift - 1) * s) return s.clamp(max=float(ns.max_t)).to(ns.t_precision) - def _generate_noise_and_time( - self, real_data: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - """The student rollout always starts from pure noise at ``max_t``.""" - batch_size = real_data.shape[0] - - eps_student = torch.randn(batch_size, *self.input_shape, device=self.device, dtype=real_data.dtype) - t_student = torch.full( - (batch_size,), - self.net.noise_scheduler.max_t, - device=self.device, - dtype=self.net.noise_scheduler.t_precision, - ) - input_student = self.net.noise_scheduler.latents(noise=eps_student) - - t = self.net.noise_scheduler.sample_t( - batch_size, **convert_cfg_to_dict(self.config.sample_t_cfg), device=self.device - ) - eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) - return input_student, t_student, t, eps - def gen_data_from_net( self, input_student: torch.Tensor, @@ -189,18 +128,14 @@ def gen_data_from_net( ) -> torch.Tensor: """Flow-map rollout compressed into at most three network forwards. - Mirrors the reference ``WanAnyFlowPipeline.training_rollout``: with an - ``num_steps``-step schedule ``t_0 > ... > t_N = 0`` and a sampled fine - step position ``g``, the model jumps ``t_0 -> t_g`` in ONE flow-map - forward, takes the single fine step ``t_g -> t_{g+1}``, then jumps - ``t_{g+1} -> 0`` in one forward. Gradient flows through all segments - (the reference applies no ``no_grad`` inside ``training_rollout``); - the fake-score update wraps this call in ``no_grad`` at the caller. - Each step uses mean-velocity sampling ``u_theta(x_t, t, r=t_next)``. + Mirrors the reference ``WanAnyFlowPipeline.training_rollout``: on a + ``num_steps``-step schedule ``t_0 > ... > t_N = 0`` with a sampled fine-step + position ``g``, jump ``t_0 -> t_g``, take the fine step ``t_g -> t_{g+1}``, + then jump ``t_{g+1} -> 0``, each one forward of mean-velocity sampling + ``u_theta(x_t, t, r=t_next)``. Gradient flows through all segments; the + fake-score update wraps this call in ``no_grad`` at the caller. """ del t_student # the rollout schedule is built below - ns = self.net.noise_scheduler - batch_size = input_student.shape[0] num_steps = self._sample_rollout_steps() if num_steps < 1: @@ -208,25 +143,23 @@ def gen_data_from_net( grad_step = self._broadcast_choice(num_steps) t_list = self._rollout_t_list(num_steps) - segments = ( - (t_list[0], t_list[grad_step]), - (t_list[grad_step], t_list[grad_step + 1]), - (t_list[grad_step + 1], t_list[-1]), + # The leading jump exists only for grad_step > 0 and the trailing one + # only for grad_step + 1 < num_steps. + seg_t = [t_list[0]] if grad_step > 0 else [] + seg_t += [t_list[grad_step], t_list[grad_step + 1]] + if grad_step + 1 < num_steps: + seg_t.append(t_list[-1]) + + return self._student_sample_loop( + self.net, + input_student, + t_list=torch.stack(seg_t), + condition=condition, + student_sample_type="ode", ) - x = input_student - for t_cur, t_next in segments: - if float(t_cur) == float(t_next): - continue - t_batch = t_cur.expand(batch_size).to(ns.t_precision) - r_batch = t_next.expand(batch_size).to(ns.t_precision) - flow_pred = self.net(x, t_batch, r=r_batch, condition=condition, fwd_pred_type="flow") - x = x - expand_like(t_batch - r_batch, x).to(x.dtype) * flow_pred - - return x - # ------------------------------------------------------------------ - # Training step — DMD2 plus the co-trained Stage-2 flow-map loss + # Training step — DMD2 plus the co-trained Stage-1 flow-map loss # ------------------------------------------------------------------ def single_train_step( @@ -234,7 +167,7 @@ def single_train_step( ) -> tuple[dict[str, torch.Tensor], dict[str, "torch.Tensor | Callable"]]: real_data, condition, neg_condition = self._prepare_training_data(data) self._setup_grad_requirements(iteration) - input_student, t_student, t, eps = self._generate_noise_and_time(real_data) + input_student, t_student, t, eps = self._generate_noise_and_time(real_data, iteration=iteration) if iteration % self.config.student_update_freq == 0: loss_map, outputs = self._student_update_step( @@ -242,7 +175,7 @@ def single_train_step( ) if self.config.cotrain_pretrain_weight > 0: # Reference cotrain_forward_kl: every generator update also - # runs the full Stage-2 bidirection (flow-map) loss on the + # runs the full Stage-1 bidirection (flow-map) loss on the # real batch. t_mf, r_mf, r_eq_t_mask = self._sample_t_r_buckets(real_data.shape[0]) mf_outputs = self._compute_mf_loss( diff --git a/fastgen/methods/distribution_matching/causvid.py b/fastgen/methods/distribution_matching/causvid.py index 9f73b2e..c2d1e0d 100644 --- a/fastgen/methods/distribution_matching/causvid.py +++ b/fastgen/methods/distribution_matching/causvid.py @@ -21,12 +21,15 @@ class CausVidModel(DMD2Model): """CausVid implementation""" def _generate_noise_and_time( - self, real_data: torch.Tensor + self, real_data: torch.Tensor, iteration: Optional[int] = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Generate random noises and time step Args: real_data: Real data tensor of shape [B, C, T, H, W] + 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`. Returns: noisy_real_data: Random noise used by the student @@ -55,11 +58,7 @@ def _generate_noise_and_time( t_inhom_expanded = t_inhom[:, None, :, None, None] # shape [B, 1, T, 1, 1] noisy_real_data = self.net.noise_scheduler.forward_process(real_data, eps_inhom, t_inhom_expanded) - t = self.net.noise_scheduler.sample_t( - batch_size, - **basic_utils.convert_cfg_to_dict(self.config.sample_t_cfg), - device=self.device, - ) + t = self._sample_noising_time(batch_size, iteration) eps = torch.randn_like(eps_inhom, device=self.device, dtype=real_data.dtype) return noisy_real_data, t_inhom, t, eps diff --git a/fastgen/methods/distribution_matching/dmd2.py b/fastgen/methods/distribution_matching/dmd2.py index b106e61..c5bad9e 100644 --- a/fastgen/methods/distribution_matching/dmd2.py +++ b/fastgen/methods/distribution_matching/dmd2.py @@ -76,14 +76,31 @@ 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. + + 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`. + """ + t_cfg = self.config.sample_t_cfg + is_fake_score_step = iteration is not None and iteration % self.config.student_update_freq != 0 + 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 + self, real_data: torch.Tensor, iteration: Optional[int] = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Generate random noises and time step Args: - batch_size: Batch size 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`. Returns: rand_z_max: Random noise used by the student @@ -115,9 +132,7 @@ def _generate_noise_and_time( ) input_student = self.net.noise_scheduler.forward_process(real_data, eps_student, t_student) - t = self.net.noise_scheduler.sample_t( - batch_size, **convert_cfg_to_dict(self.config.sample_t_cfg), device=self.device - ) + t = self._sample_noising_time(batch_size, iteration) eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) return input_student, t_student, t, eps @@ -442,7 +457,7 @@ def single_train_step( self._setup_grad_requirements(iteration) # Generate noise and time steps - input_student, t_student, t, eps = self._generate_noise_and_time(real_data) + input_student, t_student, t, eps = self._generate_noise_and_time(real_data, iteration=iteration) # Choose between student update or fake_score/discriminator update if iteration % self.config.student_update_freq == 0: diff --git a/fastgen/methods/distribution_matching/self_forcing.py b/fastgen/methods/distribution_matching/self_forcing.py index 5d7cd2e..ea8b24f 100644 --- a/fastgen/methods/distribution_matching/self_forcing.py +++ b/fastgen/methods/distribution_matching/self_forcing.py @@ -37,13 +37,15 @@ def __init__(self, config: ModelConfig): self.config = config def _generate_noise_and_time( - self, real_data: torch.Tensor + self, real_data: torch.Tensor, iteration: Optional[int] = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Generate random noises and time step Args: - batch_size: Batch size 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`. Returns: input_student: Random noise used by the student @@ -62,9 +64,7 @@ def _generate_noise_and_time( ) input_student = self.net.noise_scheduler.latents(noise=eps_student) - t = self.net.noise_scheduler.sample_t( - batch_size, **convert_cfg_to_dict(self.config.sample_t_cfg), device=self.device - ) + t = self._sample_noising_time(batch_size, iteration) eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) diff --git a/fastgen/networks/EDM/network.py b/fastgen/networks/EDM/network.py index 331b481..0c9a550 100644 --- a/fastgen/networks/EDM/network.py +++ b/fastgen/networks/EDM/network.py @@ -924,12 +924,6 @@ def forward( else condition.reshape(-1, self.label_dim) ) - # A dual-timestep network always consumes the r-pathway (the embedding - # widths are sized for it), so r=None is not a valid input for it. - # Default to r=t, i.e. the instantaneous velocity u(x_t, t, t). - if r is None and getattr(self.model, "r_timestep", None) is not None: - r = t - # Preconditioning weights for input x_t_in, t_in = x_t, t if self.drop_precond not in ["input", "both"]: diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index caf2c5b..6d4754b 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -31,6 +31,7 @@ from transformers import AutoTokenizer, UMT5EncoderModel from fastgen.networks.network import FastGenNetwork +from fastgen.networks.Wan.utils import convert_wan_official_state_dict, remap_anyflow_keys from fastgen.networks.noise_schedule import NET_PRED_TYPES from fastgen.utils.basic_utils import prompt_clean, str2bool @@ -283,21 +284,21 @@ def _fuse_r_embedding( ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Combine the t- and r-time embeddings into the final timestep projection. - Two fusion modes share the same ``r_embedder.time_proj`` / ``act_fn`` modules: - * ``additive`` (MeanFlow default, ``r_embedder.fusion_mode`` absent or "additive"): - ``r_timestep_proj = time_proj(act_fn(remb))`` is added to ``timestep_proj`` - and ``remb`` is added to ``temb``. T2V dual-stream variants (``encoder_depth`` - set) instead replace ``temb`` with ``remb`` and leave ``timestep_proj`` alone. - - * ``gated`` (AnyFlow, opt-in): convex-combine the two embeddings *before* the - shared projection — ``rt_emb = (1-g)·temb + g·remb`` then - ``timestep_proj = time_proj(act_fn(rt_emb))``. This matches the - ``WanTwoTimeTextImageEmbedding.forward_timestep`` path in the AnyFlow - reference and is required to reproduce the published HF checkpoint - forward bit-for-bit. With ``encoder_depth`` set, the first - ``encoder_depth`` blocks keep the t-only ``timestep_proj`` and the - remaining blocks switch to the gated projection (mirroring additive). + ``r_timestep_proj = r_embedder.time_proj(r_embedder.act_fn(remb))`` is added to + ``timestep_proj`` and ``remb`` is added to ``temb``. T2V dual-stream variants + (``encoder_depth`` set) instead replace ``temb`` with ``remb`` and leave + ``timestep_proj`` alone. + + * ``gated`` (opt-in): convex-combine the two embeddings *before* the + projection — ``rt_emb = (1-g)·temb + g·remb`` then + ``timestep_proj = condition_embedder.time_proj(act_fn(rt_emb))``. The + projection is the condition_embedder's own, i.e. SHARED with the t-only + path; ``Wan.__init__`` therefore drops the redundant + ``r_embedder.time_proj`` in gated mode so it cannot silently drift away + from the shared one during training. With ``encoder_depth`` + set, the first ``encoder_depth`` blocks keep the t-only ``timestep_proj`` + and the remaining blocks switch to the gated projection (mirroring additive). Returns ``(temb, timestep_proj, r_timestep_proj)`` where ``r_timestep_proj`` is ``None`` when the r-information is already folded into ``timestep_proj``. @@ -307,7 +308,7 @@ def _fuse_r_embedding( if fusion == "gated": gate = self.r_embedder.gate_value rt_emb = (1 - gate) * temb + gate * remb - rt_ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(rt_emb)) + rt_ts_proj = self.condition_embedder.time_proj(self.condition_embedder.act_fn(rt_emb)) rt_ts_proj = unflatten_timestep_proj(rt_ts_proj, rs_seq_len) if self.encoder_depth is None: return rt_emb, rt_ts_proj, None @@ -464,50 +465,6 @@ def classify_forward_block_forward( return hidden_states, features -def remap_anyflow_keys(state_dict: Mapping[str, Any]) -> Mapping[str, Any]: - """Remap an AnyFlow HF release state_dict to FastGen's Wan layout. - - AnyFlow's ``FAR_Wan_Transformer3DModel`` stores the r-pathway inside the - main ``condition_embedder`` as ``delta_embedder``, and uses ONE shared - ``time_proj`` for both t and (t, r). FastGen exposes a separate top-level - ``r_embedder`` with its own ``time_embedder`` + ``time_proj``. The two - layouts are functionally equivalent (FastGen's ``r_embedder.time_proj`` - starts as a deepcopy of ``condition_embedder.time_proj`` per - :meth:`Wan.init_embedder`), so we just rename / duplicate the tensors. - - The function is a no-op when no ``condition_embedder.delta_embedder.*`` - keys are present, so it's safe to call unconditionally. Keys with or - without the ``transformer.`` module prefix are both handled. - """ - delta_marker = "condition_embedder.delta_embedder." - delta_keys = [k for k in state_dict if delta_marker in k] - if not delta_keys: - return state_dict - - new_sd = dict(state_dict) - prefixes = set() - for k in delta_keys: - # [transformer.]condition_embedder.delta_embedder.linear_1.weight - # -> [transformer.]r_embedder.time_embedder.linear_1.weight - prefix, _, suffix = k.partition(delta_marker) - prefixes.add(prefix) - new_sd[f"{prefix}r_embedder.time_embedder.{suffix}"] = new_sd.pop(k) - # AnyFlow's gated fusion shares the final time_proj. FastGen has a - # separate r_embedder.time_proj that substitutes for the shared one when - # fusion="gated"; copy the weights so the two projections start identical. - for prefix in prefixes: - for sub in ("weight", "bias"): - src = f"{prefix}condition_embedder.time_proj.{sub}" - dst = f"{prefix}r_embedder.time_proj.{sub}" - if src in new_sd and dst not in new_sd: - new_sd[dst] = new_sd[src].clone() - logger.info( - f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors " - "and duplicated time_proj weights into r_embedder." - ) - return new_sd - - class WanTextEncoder: def __init__(self, model_id_or_local_path): self.text_encoder = UMT5EncoderModel.from_pretrained( @@ -696,18 +653,8 @@ def __init__( self.transformer.time_cond_type = time_cond_type if r_timestep: logger.info(f"Initializing r embedder with {r_embedder_init}") - self.transformer.r_embedder = self.init_embedder(r_embedder_init) - # Stash fusion config on the r_embedder module so the (method-bound) - # forward override can branch on it without changing its signature. - if r_embedder_fusion not in ("additive", "gated"): - raise ValueError(f"r_embedder_fusion must be 'additive' or 'gated', got {r_embedder_fusion!r}") - self.transformer.r_embedder.fusion_mode = r_embedder_fusion - # Plain float (not a buffer) so FSDP's reset_parameters does not - # need to re-materialize it after meta-device init. - self.transformer.r_embedder.gate_value = float(r_embedder_gate_value) - logger.info( - f"r_embedder fusion={r_embedder_fusion}" - + (f" gate_value={r_embedder_gate_value}" if r_embedder_fusion == "gated" else "") + self.transformer.r_embedder = self.init_embedder( + r_embedder_init, r_embedder_fusion, r_embedder_gate_value ) else: self.transformer.r_embedder = None @@ -715,6 +662,16 @@ def __init__( # core functionality to override forward function and other methods in the transformer self.override_transformer_forward(inner_dim=inner_dim) + # Only the base Wan forward routes the r-embedding through + # _fuse_r_embedding (installed just above); the causal / VACE variants + # install their own (additive-only) classify_forward_prepare, where + # "gated" would be silently ignored. + if r_timestep and r_embedder_fusion == "gated" and not hasattr(self.transformer, "_fuse_r_embedding"): + raise ValueError( + f"r_embedder_fusion='gated' is not supported by {type(self).__name__} " + "(its transformer forward does not implement gated t/r fusion)." + ) + # Use lazy initialization as this wont work in a meta context when doing FSDP2. self._unipc_scheduler = None @@ -890,17 +847,44 @@ def get_model_id(cls, model_path_or_id: str | os.PathLike) -> str: name = model_path_or_id[idx_start:].split("/")[0] return f"Wan-AI/{name}" - def init_embedder(self, embedder_init: str) -> None: + def init_embedder( + self, + embedder_init: str, + r_embedder_fusion: str = "additive", + r_embedder_gate_value: float = 0.25, + ) -> torch.nn.Module: + """Build the r-embedder and stash its fusion config on the module. + + The fusion settings live on the module (not on `self`) so the + method-bound forward override can branch on them without changing its + signature. They are set before the meta-device early return so both + paths carry them. + + For gated fusion the r_embedder's own `time_proj` / `act_fn` are dropped + here (the fused embedding goes through the condition_embedder's shared + projection instead). Whether the transformer's forward actually supports + gated fusion can only be checked after `override_transformer_forward` + installs `_fuse_r_embedding`, so that check lives in `__init__`. + """ embedder = copy.deepcopy(self.transformer.condition_embedder) del embedder.text_embedder + if r_embedder_fusion not in ("additive", "gated"): + raise ValueError(f"r_embedder_fusion must be 'additive' or 'gated', got {r_embedder_fusion!r}") + embedder.fusion_mode = r_embedder_fusion + # Plain float (not a buffer) so FSDP's reset_parameters does not need to + # re-materialize it after meta-device init. + embedder.gate_value = float(r_embedder_gate_value) + logger.info( + f"r_embedder fusion={r_embedder_fusion}" + + (f" gate_value={r_embedder_gate_value}" if r_embedder_fusion == "gated" else "") + ) + # Skip initialization if using meta device (weights will be broadcast via FSDP) if self._is_in_meta_context(): logger.info("Skipping r_embedder initialization on meta device (will receive weights via FSDP sync)") - return embedder - # zero init the r_embedder - if embedder_init == "zero": + elif embedder_init == "zero": for param in embedder.parameters(): param.data.zero_() elif embedder_init == "random": @@ -925,6 +909,15 @@ def init_embedder(self, embedder_init: str) -> None: pass else: raise ValueError(f"Invalid embedder_init: {embedder_init}") + + if r_embedder_fusion == "gated": + # The gated path projects the fused (t, r) embedding through the + # condition_embedder's own time_proj, shared with the t-only path. + # Drop the r_embedder's copy so it cannot linger as an untrained + # dead parameter. Done after the init branches above because + # "random" initializes time_proj. + del embedder.time_proj + del embedder.act_fn return embedder def override_transformer_forward(self, inner_dim: int) -> None: @@ -1090,86 +1083,11 @@ def load_state_dict(self, state_dict: Mapping[str, Any], **kwargs): i.e., {'generator': state_dict}. """ if self._use_wan_official_sinusoidal and not any(k.startswith("transformer.") for k in state_dict.keys()): - # Handle original Wan checkpoint formats - # Pick the source state dict (adjust these keys to your file) - state = None - for k in ["generator", "state_dict", "model", "module", "net", None]: - if k is None: - # fallback: assume loaded object IS the state_dict - if isinstance(state_dict, dict) and all(isinstance(v, torch.Tensor) for v in state_dict.values()): - state = state_dict - break - if isinstance(state_dict, dict) and k in state_dict and isinstance(state_dict[k], dict): - state = state_dict[k] - break - assert state is not None, "Could not find a state_dict in checkpoint." - logger.info(f"Loading original Wan checkpoint formats from key: {k}") - - # Rename mapping as list of tuples (order matters for the norm swap) - rename_mapping = [ - ("time_embedding.0", "condition_embedder.time_embedder.linear_1"), - ("time_embedding.2", "condition_embedder.time_embedder.linear_2"), - ("text_embedding.0", "condition_embedder.text_embedder.linear_1"), - ("text_embedding.2", "condition_embedder.text_embedder.linear_2"), - ("time_projection.1", "condition_embedder.time_proj"), - ("head.modulation", "scale_shift_table"), - ("head.head", "proj_out"), - ("modulation", "scale_shift_table"), - ("ffn.0", "ffn.net.0.proj"), - ("ffn.2", "ffn.net.2"), - # swap norm names: norm1, norm3, norm2 -> norm1, norm2, norm3 - ("norm2", "norm__placeholder"), - ("norm3", "norm2"), - ("norm__placeholder", "norm3"), - # I2V extras - ("img_emb.proj.0", "condition_embedder.image_embedder.norm1"), - ("img_emb.proj.1", "condition_embedder.image_embedder.ff.net.0.proj"), - ("img_emb.proj.3", "condition_embedder.image_embedder.ff.net.2"), - ("img_emb.proj.4", "condition_embedder.image_embedder.norm2"), - ("img_emb.emb_pos", "condition_embedder.image_embedder.pos_embed"), - # attention parts - ("self_attn.q", "attn1.to_q"), - ("self_attn.k", "attn1.to_k"), - ("self_attn.v", "attn1.to_v"), - ("self_attn.o", "attn1.to_out.0"), - ("self_attn.norm_q", "attn1.norm_q"), - ("self_attn.norm_k", "attn1.norm_k"), - ("cross_attn.q", "attn2.to_q"), - ("cross_attn.k", "attn2.to_k"), - ("cross_attn.v", "attn2.to_v"), - ("cross_attn.o", "attn2.to_out.0"), - ("cross_attn.norm_q", "attn2.norm_q"), - ("cross_attn.norm_k", "attn2.norm_k"), - ("attn2.to_k_img", "attn2.add_k_proj"), - ("attn2.to_v_img", "attn2.add_v_proj"), - ("attn2.norm_k_img", "attn2.norm_added_k"), - ] - - # Convert keys - def rename_key(k: str) -> str: - # strip common prefixes if present - for prefix in ["model.", "module.", "transformer."]: - if k.startswith(prefix): - k = k[len(prefix) :] - # apply replacements in the specified order - for old, new in rename_mapping: - if old in k: - k = k.replace(old, new) - return k - - new_state = {} - for k, v in state.items(): - # optional: skip buffer-like positional/freq params if not needed - # if "freqs" in k: - # continue - new_k = rename_key(k) - # Add 'transformer.' prefix since the model expects it - new_k = f"transformer.{new_k}" - new_state[new_k] = v - state_dict = new_state - - # AnyFlow HF releases store the r-pathway as condition_embedder.delta_embedder; - # remap to FastGen's r_embedder layout (no-op for all other checkpoints). + state_dict = convert_wan_official_state_dict(state_dict) + + # Some third-party checkpoints store the r-pathway as + # condition_embedder.delta_embedder; remap to FastGen's r_embedder + # layout (no-op for all other checkpoints). state_dict = remap_anyflow_keys(state_dict) return super().load_state_dict(state_dict, **kwargs) @@ -1225,9 +1143,8 @@ def forward( # A gated-fusion flow-map network always consumes the r-pathway (the # fusion happens inside the shared time projection), so r=None has no # in-distribution meaning for it. Default to r=t, i.e. the - # instantaneous velocity u(x_t, t, t) — this is how the AnyFlow - # reference queries its real/fake score models. Additive-fusion - # (MeanFlow) networks keep the skip-r-pathway behavior. + # instantaneous velocity u(x_t, t, t). Additive-fusion (MeanFlow) + # networks keep the skip-r-pathway behavior. if r is None and getattr(self.transformer.r_embedder, "fusion_mode", None) == "gated": r = t diff --git a/fastgen/networks/Wan/utils.py b/fastgen/networks/Wan/utils.py new file mode 100644 index 0000000..f354c00 --- /dev/null +++ b/fastgen/networks/Wan/utils.py @@ -0,0 +1,134 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint-format conversions for the Wan networks. + +FastGen's Wan wraps a diffusers ``WanTransformer3DModel`` and expects keys under a +``transformer.`` prefix. Third-party checkpoints come in other layouts, and the +converters here normalise them. They are applied in ``Wan.load_state_dict``: + +* ``convert_wan_official_state_dict`` — the original (non-diffusers) Wan release + layout, e.g. ``self_attn.q`` -> ``attn1.to_q``. +* ``remap_anyflow_keys`` — the AnyFlow HF releases, which store the flow-map + r-pathway inside ``condition_embedder`` as ``delta_embedder``. + +Each converter is a no-op when its marker keys are absent, so they are safe to call +unconditionally and in sequence. +""" + +from typing import Any, Dict, List, Mapping, Tuple + +import torch + +import fastgen.utils.logging_utils as logger + + +# Rename mapping for the original Wan checkpoint layout, as an ordered list of +# (old, new) substring pairs — order matters for the norm swap below. +WAN_OFFICIAL_RENAME_MAPPING: List[Tuple[str, str]] = [ + ("time_embedding.0", "condition_embedder.time_embedder.linear_1"), + ("time_embedding.2", "condition_embedder.time_embedder.linear_2"), + ("text_embedding.0", "condition_embedder.text_embedder.linear_1"), + ("text_embedding.2", "condition_embedder.text_embedder.linear_2"), + ("time_projection.1", "condition_embedder.time_proj"), + ("head.modulation", "scale_shift_table"), + ("head.head", "proj_out"), + ("modulation", "scale_shift_table"), + ("ffn.0", "ffn.net.0.proj"), + ("ffn.2", "ffn.net.2"), + # swap norm names: norm1, norm3, norm2 -> norm1, norm2, norm3 + ("norm2", "norm__placeholder"), + ("norm3", "norm2"), + ("norm__placeholder", "norm3"), + # I2V extras + ("img_emb.proj.0", "condition_embedder.image_embedder.norm1"), + ("img_emb.proj.1", "condition_embedder.image_embedder.ff.net.0.proj"), + ("img_emb.proj.3", "condition_embedder.image_embedder.ff.net.2"), + ("img_emb.proj.4", "condition_embedder.image_embedder.norm2"), + ("img_emb.emb_pos", "condition_embedder.image_embedder.pos_embed"), + # attention parts + ("self_attn.q", "attn1.to_q"), + ("self_attn.k", "attn1.to_k"), + ("self_attn.v", "attn1.to_v"), + ("self_attn.o", "attn1.to_out.0"), + ("self_attn.norm_q", "attn1.norm_q"), + ("self_attn.norm_k", "attn1.norm_k"), + ("cross_attn.q", "attn2.to_q"), + ("cross_attn.k", "attn2.to_k"), + ("cross_attn.v", "attn2.to_v"), + ("cross_attn.o", "attn2.to_out.0"), + ("cross_attn.norm_q", "attn2.norm_q"), + ("cross_attn.norm_k", "attn2.norm_k"), + ("attn2.to_k_img", "attn2.add_k_proj"), + ("attn2.to_v_img", "attn2.add_v_proj"), + ("attn2.norm_k_img", "attn2.norm_added_k"), +] + +# Keys a checkpoint may nest the model state under. +_NESTED_STATE_KEYS = ["generator", "state_dict", "model", "module", "net"] + +# Prefixes stripped before applying the rename mapping. +_STRIPPED_PREFIXES = ["model.", "module.", "transformer."] + + +def _unwrap_state_dict(state_dict: Mapping[str, Any]) -> Tuple[Dict[str, Any], str | None]: + """Return the tensor state dict, unwrapping one level of nesting if present.""" + for key in _NESTED_STATE_KEYS: + if isinstance(state_dict, dict) and key in state_dict and isinstance(state_dict[key], dict): + return state_dict[key], key + # fallback: assume the loaded object IS the state dict + if isinstance(state_dict, dict) and all(isinstance(v, torch.Tensor) for v in state_dict.values()): + return state_dict, None + raise ValueError("Could not find a state_dict in checkpoint.") + + +def rename_wan_official_key(key: str) -> str: + """Map one original-Wan parameter name onto the diffusers layout.""" + for prefix in _STRIPPED_PREFIXES: + if key.startswith(prefix): + key = key[len(prefix) :] + for old, new in WAN_OFFICIAL_RENAME_MAPPING: + if old in key: + key = key.replace(old, new) + return key + + +def convert_wan_official_state_dict(state_dict: Mapping[str, Any]) -> Dict[str, Any]: + """Convert an original (non-diffusers) Wan checkpoint to FastGen's layout. + + Unwraps a nested state dict if needed, renames the parameters via + ``WAN_OFFICIAL_RENAME_MAPPING``, and adds the ``transformer.`` prefix the + model expects. + """ + state, nested_key = _unwrap_state_dict(state_dict) + logger.info(f"Loading original Wan checkpoint format from key: {nested_key}") + return {f"transformer.{rename_wan_official_key(k)}": v for k, v in state.items()} + + +def remap_anyflow_keys(state_dict: Mapping[str, Any]) -> Mapping[str, Any]: + """Remap an AnyFlow HF release state_dict to FastGen's Wan layout. + + AnyFlow's ``FAR_Wan_Transformer3DModel`` stores the r-pathway inside the main + ``condition_embedder`` as ``delta_embedder``, and uses ONE shared ``time_proj`` + for both t and (t, r). FastGen exposes the r-pathway as a top-level + ``r_embedder``; in gated mode it keeps only the ``time_embedder`` and reuses + ``condition_embedder.time_proj`` (see ``_fuse_r_embedding`` in + ``Wan/network.py``), so the layouts differ by a rename only. + + The function is a no-op when no ``condition_embedder.delta_embedder.*`` keys are + present, so it's safe to call unconditionally. Keys with or without the + ``transformer.`` module prefix are both handled. + """ + delta_marker = "condition_embedder.delta_embedder." + delta_keys = [k for k in state_dict if delta_marker in k] + if not delta_keys: + return state_dict + + new_sd = dict(state_dict) + for k in delta_keys: + # [transformer.]condition_embedder.delta_embedder.linear_1.weight + # -> [transformer.]r_embedder.time_embedder.linear_1.weight + prefix, _, suffix = k.partition(delta_marker) + new_sd[f"{prefix}r_embedder.time_embedder.{suffix}"] = new_sd.pop(k) + logger.info(f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors into r_embedder.") + return new_sd diff --git a/fastgen/networks/noise_schedule.py b/fastgen/networks/noise_schedule.py index eefb519..a7463e1 100644 --- a/fastgen/networks/noise_schedule.py +++ b/fastgen/networks/noise_schedule.py @@ -1317,8 +1317,8 @@ def __init__( **kwargs, ): super().__init__(min_t, max_t, num_steps, **kwargs) - self._supported_time_dist_types = self._supported_time_dist_types + ("shifted",) - assert 0 <= min_t < max_t <= 0.999, "RF min_t and max_t must be between 0 and 0.999" + self._supported_time_dist_types = self._supported_time_dist_types + ("shifted", "shifted_logitnormal") + assert 0 <= min_t < max_t <= 1.0, "RF min_t and max_t must be between 0 and 1" self._sigmas = torch.linspace(min_t, max_t, num_steps, dtype=self.t_precision) def _rescale_t(self, t: torch.Tensor) -> torch.Tensor: @@ -1326,9 +1326,7 @@ def _rescale_t(self, t: torch.Tensor) -> torch.Tensor: @property def max_sigma(self) -> float: - t_max_scale = int(self.num_steps * self.max_t) - assert 0 <= t_max_scale < len(self._sigmas) - return self._sigmas[t_max_scale].item() + return self._max_t @property def sigmas(self) -> torch.Tensor: @@ -1416,10 +1414,20 @@ def sample_t( assert shift >= 1, f"shift must be >= 1, got {shift}" t = torch.rand(n, device=target_device, dtype=self.t_precision) * (max_t - min_t) + min_t t = t * shift / (t * (shift - 1) + 1) + elif time_dist_type == "shifted_logitnormal": + # Logit-normal base density under the same shift map as "shifted" + shift = kwargs.get("shift", 5.0) + assert shift >= 1, f"shift must be >= 1, got {shift}" + t = ( + torch.sigmoid(torch.randn(n, device=target_device, dtype=self.t_precision) * train_p_std + train_p_mean) + * (max_t - min_t) + + min_t + ) + t = t * shift / (t * (shift - 1) + 1) else: raise ValueError( f"Unsupported time distribution type: {time_dist_type} in RFNoiseSchedule." - f"Currently only supports logitnormal, uniform, and shifted." + f"Currently only supports logitnormal, uniform, shifted, and shifted_logitnormal." ) return self.safe_clamp(t, min_t, max_t) diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index fef00e4..7e8f6cb 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -5,8 +5,8 @@ The pretrain stage is MeanFlow with AnyFlow's hyperparameters (fixed per-timestep loss weighting, finite-difference JVP, consistency bucket), so -the pretrain tests drive :class:`MeanFlowModel` directly. The on-policy tests -exercise :class:`AnyFlowModel` (DMD2 with a multi-step rollout-with-gradient +the pretrain tests drive ``MeanFlowModel`` directly. The on-policy tests +exercise ``AnyFlowModel`` (DMD2 with a multi-step rollout-with-gradient student). Both run on the tiny EDM backbone with ``r_timestep=True`` and ``schedule_type=rf`` so they execute on CPU without pretrained weights. """ @@ -24,8 +24,10 @@ from fastgen.utils.test_utils import check_grad_zero -def _build_pretrain_model(weight_type="beta08", consistency_ratio=0.25, r_sample_ratio=0.5): - """MeanFlow configured the AnyFlow way (paper Stage 2).""" +def _build_pretrain_model( + weight_type="beta08", consistency_ratio=0.25, flow_matching_ratio=0.5, deterministic_buckets=True +): + """MeanFlow configured the AnyFlow way (paper Stage 1).""" gc.collect() instance = MeanFlowModelConfig() @@ -39,8 +41,9 @@ def _build_pretrain_model(weight_type="beta08", consistency_ratio=0.25, r_sample instance.sample_t_cfg.shift = 5.0 instance.sample_t_cfg.min_t = 0.001 instance.sample_t_cfg.max_t = 0.999 - instance.sample_t_cfg.r_sample_ratio = r_sample_ratio + instance.sample_t_cfg.flow_matching_ratio = flow_matching_ratio instance.sample_t_cfg.consistency_ratio = consistency_ratio + instance.sample_t_cfg.deterministic_buckets = deterministic_buckets opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True", "+schedule_type=rf"] instance.net = override_config_with_opts(instance.net, opts) @@ -61,7 +64,11 @@ def _build_onpolicy_model(): gc.collect() instance = AnyFlowModelConfig() - opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True", "+schedule_type=rf"] + base_opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "+schedule_type=rf"] + # Teacher / fake score are PLAIN single-timestep nets, as in the reference + # and in config_anyflow_onpolicy.py; only the student carries the r pathway. + instance.teacher = override_config_with_opts(AnyFlowModelConfig().net, list(base_opts)) + opts = base_opts[:1] + ["r_timestep=True"] + base_opts[1:] instance.net = override_config_with_opts(instance.net, opts) opts_disc = ["-", "feature_indices=[0]", "all_res=[8]", "in_channels=128"] instance.discriminator = override_config_with_opts(instance.discriminator, opts_disc) @@ -77,8 +84,9 @@ def _build_onpolicy_model(): instance.sample_t_cfg.time_dist_type = "uniform" instance.sample_t_cfg.min_t = 0.001 instance.sample_t_cfg.max_t = 0.999 - instance.sample_t_cfg.r_sample_ratio = 0.5 + instance.sample_t_cfg.flow_matching_ratio = 0.5 instance.sample_t_cfg.consistency_ratio = 0.25 + instance.sample_t_cfg.deterministic_buckets = True instance.loss_config.loss_type = "l2" instance.loss_config.weight_type = "uniform" instance.loss_config.norm_method = None @@ -154,16 +162,16 @@ def test_pretrain_finite_difference_falls_back_at_boundaries(): def test_pretrain_consistency_bucket_pins_r_to_zero(): """With consistency_ratio=1.0 (and no flow-matching head), every sample's r must be pinned to 0 (consistency to clean data, as in the reference).""" - model = _build_pretrain_model(consistency_ratio=1.0, r_sample_ratio=1.0) + model = _build_pretrain_model(consistency_ratio=1.0, flow_matching_ratio=0.0) _t, r, r_eq_t_mask = model._sample_t_r_buckets(4) assert not r_eq_t_mask.any() assert torch.allclose(r.float(), torch.zeros_like(r.float())) def test_pretrain_bucket_partition_is_deterministic(): - """With consistency_ratio > 0, buckets follow the reference's deterministic - global partition: diffusion head, consistency middle, random-pair tail.""" - model = _build_pretrain_model(consistency_ratio=0.25, r_sample_ratio=0.5) + """With deterministic_buckets, both buckets follow the reference's global + partition: flow-matching head, consistency middle, random-pair tail.""" + model = _build_pretrain_model(consistency_ratio=0.25, flow_matching_ratio=0.5, deterministic_buckets=True) batch_size = 8 t, r, r_eq_t_mask = model._sample_t_r_buckets(batch_size) @@ -177,17 +185,41 @@ def test_pretrain_bucket_partition_is_deterministic(): ) -def test_pretrain_rebalance_to_diffusion(): - """With rebalancing on, each non-diffusion loss is rescaled to the - diffusion-loss mean by a detached per-sample factor.""" +def test_pretrain_bucket_partition_is_stochastic(): + """Without deterministic_buckets the same two ratios drive binomial bucket + SIZES: the layout is still head/middle/tail, but the sizes vary per call + and average to the configured fractions.""" + model = _build_pretrain_model(consistency_ratio=0.25, flow_matching_ratio=0.5, deterministic_buckets=False) + torch.manual_seed(0) + batch_size = 4096 + t, r, r_eq_t_mask = model._sample_t_r_buckets(batch_size) + + # min_t = 0.001, so only the consistency bucket can carry r == 0. + is_consistency = r == 0 + assert not (r_eq_t_mask & is_consistency).any() + assert torch.equal(r[r_eq_t_mask], t[r_eq_t_mask]) + assert abs(r_eq_t_mask.float().mean().item() - 0.5) < 0.05 + assert abs(is_consistency.float().mean().item() - 0.25) < 0.05 + # Prefix layout, same as the deterministic policy. + n_flow_matching = int(r_eq_t_mask.sum()) + assert r_eq_t_mask[:n_flow_matching].all() and not r_eq_t_mask[n_flow_matching:].any() + + # The sizes are what varies, unlike the deterministic policy. + sizes = {int(model._sample_t_r_buckets(64)[2].sum()) for _ in range(20)} + assert len(sizes) > 1, f"stochastic bucket sizes should vary, got {sizes}" + + +def test_pretrain_rebalance_to_flow_matching(): + """With rebalancing on, each r < t loss is rescaled to the flow-matching + loss mean by a detached per-sample factor.""" model = _build_pretrain_model() - model.loss_config.rebalance_to_diffusion = True + model.loss_config.rebalance_to_flow_matching = True mf_loss = torch.tensor([2.0, 4.0, 10.0, 100.0], dtype=torch.float64, requires_grad=True) r_eq_t_mask = torch.tensor([True, True, False, False]) loss = model._reduce_mf_loss(mf_loss, r_eq_t_mask) - # diffusion mean = 3.0; each non-diffusion sample becomes ~3.0. + # flow-matching mean = 3.0; each r < t sample becomes ~3.0. expected = (2.0 + 4.0 + 3.0 * (10.0 / 10.00001) + 3.0 * (100.0 / 100.00001)) / 4.0 assert abs(loss.item() - expected) < 1e-3 loss.backward() @@ -199,7 +231,7 @@ def test_pretrain_rebalance_all_flow_matching_batch(): the rebalance branch (the collective inside must run on every rank) and reduce to the plain mean.""" model = _build_pretrain_model() - model.loss_config.rebalance_to_diffusion = True + model.loss_config.rebalance_to_flow_matching = True mf_loss = torch.tensor([2.0, 4.0], dtype=torch.float64) r_eq_t_mask = torch.tensor([True, True]) @@ -227,10 +259,11 @@ def test_timestep_weight_function(weight_type): """The fixed per-timestep weight is a direct function of t: non-negative, finite, and normalized like the reference scheduler.""" model = _build_pretrain_model(weight_type=weight_type) - model._timestep_weight_scale = None # force re-derivation for this weight_type + # The fixture sets norm_method=None, so _compute_weight on a ones tensor is + # exactly the fixed per-timestep weight w(t). t = torch.linspace(0.0, 1.0, 101, dtype=torch.float64) - w = model._get_timestep_weight(t) + w = model._compute_weight(torch.ones_like(t), t) assert torch.all(w >= 0), f"{weight_type} weights must be non-negative" assert torch.isfinite(w).all() @@ -238,12 +271,14 @@ def test_timestep_weight_function(weight_type): # Uniform normalizes to exactly 1.0 over the reference grid. assert torch.allclose(w, torch.ones_like(w)) - # The normalization matches the reference set_timesteps grid (1000 points, - # t=0 excluded): sum over the shifted grid == 1000. + # The weight has mean one over the network's discrete training timesteps + # (t=0 excluded), which at the default num_steps=1000 is the reference's + # set_timesteps grid: sum over the shifted grid == num_steps. + num_steps = model.net.noise_scheduler.num_steps shift = model.sample_t_cfg.shift - grid = torch.linspace(1.0, 0.0, 1001, dtype=torch.float64)[:-1] + grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64)[:-1] grid = shift * grid / (1 + (shift - 1) * grid) - assert abs(model._get_timestep_weight(grid).sum().item() - 1000.0) < 1e-6 + assert abs(model._compute_weight(torch.ones_like(grid), grid).sum().item() - num_steps) < 1e-6 # --------------------------------------------------------------------------- @@ -257,7 +292,7 @@ def test_onpolicy_student_update_step(): loss_map, outputs = model.single_train_step(data, 0) # iteration 0 -> student update assert "total_loss" in loss_map assert "vsd_loss" in loss_map - # The co-trained Stage-2 flow-map loss is part of every student update. + # The co-trained Stage-1 flow-map loss is part of every student update. assert "bidirection_loss" in loss_map assert torch.isfinite(loss_map["total_loss"]).all() assert "gen_rand" in outputs @@ -271,6 +306,26 @@ def test_onpolicy_cotrain_can_be_disabled(): assert "bidirection_loss" not in loss_map +def test_onpolicy_student_starts_from_pure_noise(): + """The rollout is on-policy: the student starts from pure noise at max_t. + + DMD2's multi-step branch would hand back real data noised to a random entry + of `t_list`, which `gen_data_from_net` would then roll out as if it sat at + `t_list[0]` — both off-policy and a latent/timestep mismatch. + """ + model = _build_onpolicy_model() + torch.manual_seed(0) + real = torch.randn(8, 3, 8, 8, device=model.device, dtype=model.precision) + input_student, t_student, _, _ = model._generate_noise_and_time(real, iteration=0) + + max_t = float(model.net.noise_scheduler.max_t) + assert torch.allclose(t_student.float(), torch.full_like(t_student.float(), max_t)) + # Pure noise carries no signal from the batch it was drawn alongside; the + # off-policy failure mode correlates at ~0.7 (1/sqrt(192) ~ 0.07 by chance). + cos = torch.nn.functional.cosine_similarity(input_student.flatten(1).float(), real.flatten(1).float()) + assert cos.abs().max() < 0.35, f"student input correlates with real data: {cos.tolist()}" + + def test_onpolicy_rollout_compresses_to_three_forwards(): """Regardless of the sampled NFE, the rollout must run at most three network forwards (jump -> fine step -> jump), with gradient through all.""" @@ -361,12 +416,21 @@ def test_onpolicy_optimizer_step(): def _make_fake_fusion_self(fusion_mode, gate_value=0.25, encoder_depth=None, dim=4, proj_dim=12): + condition_embedder = torch.nn.Module() + condition_embedder.time_proj = torch.nn.Linear(dim, proj_dim) + condition_embedder.act_fn = torch.nn.SiLU() + r_embedder = torch.nn.Module() - r_embedder.time_proj = torch.nn.Linear(dim, proj_dim) - r_embedder.act_fn = torch.nn.SiLU() r_embedder.fusion_mode = fusion_mode r_embedder.gate_value = gate_value - return types.SimpleNamespace(r_embedder=r_embedder, encoder_depth=encoder_depth) + if fusion_mode == "additive": + # Gated fusion reuses condition_embedder.time_proj, so Wan.__init__ + # drops these from the r_embedder in that mode. + r_embedder.time_proj = torch.nn.Linear(dim, proj_dim) + r_embedder.act_fn = torch.nn.SiLU() + return types.SimpleNamespace( + r_embedder=r_embedder, condition_embedder=condition_embedder, encoder_depth=encoder_depth + ) def test_fuse_r_embedding_gated(): @@ -382,10 +446,62 @@ def test_fuse_r_embedding_gated(): gate = fake.r_embedder.gate_value rt_emb = (1 - gate) * temb + gate * remb - expected = fake.r_embedder.time_proj(fake.r_embedder.act_fn(rt_emb)).unflatten(1, (6, -1)) + # The gated projection goes through the SHARED condition_embedder.time_proj. + expected = fake.condition_embedder.time_proj(fake.condition_embedder.act_fn(rt_emb)).unflatten(1, (6, -1)) assert torch.allclose(out_temb, rt_emb) assert torch.allclose(out_proj, expected) assert out_r_proj is None + assert not hasattr(fake.r_embedder, "time_proj") + + +def test_validation_t_list_matches_shift(): + """The hardcoded validation schedule must equal the shifted grid. + + `config_anyflow_onpolicy.py` writes `sample_t_cfg.t_list` out as a constant + instead of deriving it at runtime, so nothing detects it going stale if + `shift` or `student_sample_steps` changes. Recompute it here from those two + settings and compare. + + Left as None, DMD2 hands `generator_fn` the noise scheduler's UNSHIFTED + `linspace(max_t, 0, N+1)`, which is far off-policy for a shift=5 model -- + asserted below so the constant cannot silently degrade to that. + """ + import fastgen.configs.experiments.WanT2V.config_anyflow_onpolicy as mod + + cfg = mod.create_config() + shift = float(cfg.model.sample_t_cfg.shift) + n = int(cfg.model.student_sample_steps) + max_t = float(cfg.model.net.max_t) # the bound `_rollout_t_list` clamps to + + grid = [1.0 - i / n for i in range(n + 1)] + expected = [min(shift * x / (1 + (shift - 1) * x), max_t) for x in grid] + + t_list = list(cfg.model.sample_t_cfg.t_list) + assert len(t_list) == n + 1, t_list + assert all(abs(a - b) < 1e-9 for a, b in zip(t_list, expected)), (t_list, expected) + assert t_list[0] == max_t and t_list[-1] == 0.0 + + # and it must not be the unshifted fallback DMD2 would use for None + unshifted = [max_t * (1.0 - i / n) for i in range(n + 1)] + assert not all(abs(a - b) < 1e-6 for a, b in zip(t_list, unshifted)) + + +def test_fuse_r_embedding_gated_trains_the_shared_time_proj(): + """The gated projection must flow gradient into condition_embedder.time_proj. + + A private r_embedder.time_proj would leave the condition_embedder's copy + without gradient (a dead parameter under DDP/FSDP) while the two silently + drift apart, breaking round-trips to the AnyFlow checkpoint layout. + """ + from fastgen.networks.Wan.network import _fuse_r_embedding + + torch.manual_seed(0) + fake = _make_fake_fusion_self("gated") + _, out_proj, _ = _fuse_r_embedding(fake, torch.randn(2, 4), torch.randn(2, 6, 2), torch.randn(2, 4), None) + out_proj.sum().backward() + + grad = fake.condition_embedder.time_proj.weight.grad + assert grad is not None and torch.any(grad != 0) def test_fuse_r_embedding_gated_respects_encoder_depth(): @@ -426,7 +542,7 @@ def test_fuse_r_embedding_additive_unchanged(): @pytest.mark.parametrize("prefix", ["", "transformer."]) def test_remap_anyflow_keys(prefix): - from fastgen.networks.Wan.network import remap_anyflow_keys + from fastgen.networks.Wan.utils import remap_anyflow_keys sd = { f"{prefix}condition_embedder.delta_embedder.linear_1.weight": torch.randn(2, 2), @@ -438,15 +554,34 @@ def test_remap_anyflow_keys(prefix): assert f"{prefix}r_embedder.time_embedder.linear_1.weight" in out assert f"{prefix}condition_embedder.delta_embedder.linear_1.weight" not in out - # The shared time_proj is duplicated into the r_embedder. - assert torch.equal(out[f"{prefix}r_embedder.time_proj.weight"], sd[f"{prefix}condition_embedder.time_proj.weight"]) - assert torch.equal(out[f"{prefix}r_embedder.time_proj.bias"], sd[f"{prefix}condition_embedder.time_proj.bias"]) + # time_proj stays shared on the condition_embedder — no r_embedder copy. + assert f"{prefix}r_embedder.time_proj.weight" not in out + assert torch.equal( + out[f"{prefix}condition_embedder.time_proj.weight"], sd[f"{prefix}condition_embedder.time_proj.weight"] + ) # Unrelated keys untouched. assert torch.equal(out[f"{prefix}blocks.0.attn1.to_q.weight"], sd[f"{prefix}blocks.0.attn1.to_q.weight"]) def test_remap_anyflow_keys_noop_without_delta_keys(): - from fastgen.networks.Wan.network import remap_anyflow_keys + from fastgen.networks.Wan.utils import remap_anyflow_keys sd = {"transformer.blocks.0.attn1.to_q.weight": torch.randn(2, 2)} assert remap_anyflow_keys(sd) is sd + + +def test_rollout_uses_the_flow_map_sample_loop(): + """`gen_data_from_net` must resolve to FlowMapLossMixin's loop, not the base one. + + `AnyFlowModel(FlowMapLossMixin, DMD2Model)` picks it purely by MRO order. + `FastGenModel._student_sample_loop` is an x0-prediction loop that never + passes `r`, and it accepts the same arguments -- so swapping the base order + would silently turn the flow-map rollout (jump / fine step / jump) into a + plain diffusion sampler instead of raising. + """ + from fastgen.methods import AnyFlowModel + from fastgen.methods.consistency_model.mean_flow import FlowMapLossMixin + + assert ( + AnyFlowModel._student_sample_loop.__func__ is FlowMapLossMixin._student_sample_loop.__func__ + ), "AnyFlowModel must inherit the flow-map sample loop; check the base-class order" diff --git a/tests/test_meanflowmodel.py b/tests/test_meanflowmodel.py index 644ae8c..08a88db 100644 --- a/tests/test_meanflowmodel.py +++ b/tests/test_meanflowmodel.py @@ -49,7 +49,7 @@ def test_single_train_step_update(get_model_data): # Run the training step; cifar10 default config assert model.config.sample_t_cfg.train_p_mean == -0.6 assert model.config.sample_t_cfg.train_p_std == 1.6 - assert model.config.sample_t_cfg.r_sample_ratio == 0.75 + assert model.config.sample_t_cfg.flow_matching_ratio == 0.25 norm_method, *norm_args = model.config.loss_config.norm_method.split("_") assert norm_method == "poly" From 5d5916df59458ce0549cbe5544e3efd3e29a05a2 Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Tue, 11 Aug 2026 20:34:14 +0000 Subject: [PATCH 14/19] anyflow: minor edits Signed-off-by: Julius Berner --- .../experiments/EDM/config_mf_cifar10.py | 1 - .../WanT2V/config_anyflow_onpolicy.py | 2 + fastgen/configs/methods/config_dmd2.py | 6 +-- fastgen/configs/methods/config_mean_flow.py | 13 +++---- fastgen/methods/__init__.py | 2 +- fastgen/methods/consistency_model/README.md | 2 +- fastgen/networks/Wan/network.py | 39 ++++++------------- tests/test_anyflowmodel.py | 4 +- 8 files changed, 26 insertions(+), 43 deletions(-) diff --git a/fastgen/configs/experiments/EDM/config_mf_cifar10.py b/fastgen/configs/experiments/EDM/config_mf_cifar10.py index 5779786..79c50b9 100644 --- a/fastgen/configs/experiments/EDM/config_mf_cifar10.py +++ b/fastgen/configs/experiments/EDM/config_mf_cifar10.py @@ -19,7 +19,6 @@ def create_config(): config.model.sample_t_cfg.train_p_mean = -0.6 config.model.sample_t_cfg.train_p_std = 1.6 - config.model.sample_t_cfg.flow_matching_ratio = 0.25 config.model.sample_t_cfg.time_dist_type = "logitnormal" config.model.sample_t_cfg.min_t = 0.0 config.model.sample_t_cfg.max_t = 0.999 diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index e520df5..f37e757 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -47,6 +47,8 @@ def create_config(): # and the reference builds them without the flow-map pathway at all. config.model.teacher = copy.deepcopy(Wan_1_3B_Config) config.model.teacher.r_timestep = False + config.model.teacher.min_t = config.model.net.min_t + config.model.teacher.max_t = config.model.net.max_t # Student init comes from the Stage 1 trainer checkpoint via # trainer.checkpointer.pretrained_ckpt_path (see the module docstring); diff --git a/fastgen/configs/methods/config_dmd2.py b/fastgen/configs/methods/config_dmd2.py index 6f5927a..73230bc 100644 --- a/fastgen/configs/methods/config_dmd2.py +++ b/fastgen/configs/methods/config_dmd2.py @@ -37,11 +37,7 @@ class ModelConfig(BaseModelConfig): discriminator_scheduler: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseSchedulerConfig)) # Optional noising-time distribution used ONLY on fake-score update - # iterations (`iteration % student_update_freq != 0`). DMD2 alternates the - # generator and fake-score updates, and the two do not have to noise from the - # same density -- AnyFlow's Stage 2, for instance, uses shifted-uniform for the - # DMD gradient and shifted-logit-normal for the fake-score target. `None` - # reuses `sample_t_cfg` for both, which is the historical behaviour. + # iterations (`iteration % student_update_freq != 0`) fake_score_sample_t_cfg: Optional[SampleTConfig] = None # student update frequency diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index 9493541..4c3a72e 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -35,7 +35,7 @@ class SampleTConfig(BaseSampleTConfig): # fraction of the batch with r = t (pure flow matching); the rest keeps # the randomly sampled r, minus the consistency fraction below - flow_matching_ratio: float = 1.0 + flow_matching_ratio: float = 0.25 # fraction of the batch forced to r = 0 (consistency-to-clean, used by # AnyFlow; 0.0 keeps the original MeanFlow behavior) @@ -83,17 +83,16 @@ class LossConfig: # loss type (choice between l2 and opt_grad) loss_type: str = "opt_grad" # optional fixed per-timestep loss weighting evaluated as a function of t - # ("beta08", "gaussian", "uniform"; used by AnyFlow). Multiplies the - # adaptive norm_method weight above; None disables it. + # ("beta08", "gaussian", "uniform"). Multiplies the adaptive norm_method + # weight above; None disables it. weight_type: Optional[str] = None # prediction-side guidance fusion scale (AnyFlow guidance distillation): # the conditional output is trained to be the guided flow directly via - # (u_cond + (g-1) * u_uncond) / g against the raw data velocity, with - # plain text dropout (cond_dropout_prob). None keeps MeanFlow's - # target-side guidance (guidance_scale, eq. 19). + # (u_cond + (g-1) * u_uncond) / g against the raw data velocity. None + # keeps MeanFlow's target-side guidance (guidance_scale). guidance_fuse_scale: Optional[float] = None # rebalance the flow-map / consistency (r < t) sample losses to the global - # flow-matching (r = t) loss mean via a detached per-sample factor (AnyFlow) + # flow-matching (r = t) loss mean via a detached per-sample factor rebalance_to_flow_matching: bool = False diff --git a/fastgen/methods/__init__.py b/fastgen/methods/__init__.py index 0508a27..b6ef462 100644 --- a/fastgen/methods/__init__.py +++ b/fastgen/methods/__init__.py @@ -15,7 +15,7 @@ from fastgen.methods.consistency_model.sCM import SCMModel as SCMModel from fastgen.methods.consistency_model.mean_flow import MeanFlowModel as MeanFlowModel -# AnyFlow reuses MeanFlow's student sample loop, so it must come after. +# AnyFlow pulls in mean_flow, so it must come after the CM import above from fastgen.methods.distribution_matching.anyflow import AnyFlowModel as AnyFlowModel from fastgen.methods.fine_tuning.sft import SFTModel as SFTModel diff --git a/fastgen/methods/consistency_model/README.md b/fastgen/methods/consistency_model/README.md index f40863f..fdd0cea 100644 --- a/fastgen/methods/consistency_model/README.md +++ b/fastgen/methods/consistency_model/README.md @@ -84,7 +84,7 @@ Learns average velocity between trajectory points: `x_r = x_t - (t-r) · u(x_t, - `sample_t_cfg.flow_matching_ratio`: Fraction of the batch with `r = t` (flow matching loss) -**Configs:** [`EDM/config_mf_cifar10.py`](../../configs/experiments/EDM/config_mf_cifar10.py), [`DiT/config_mf_b.py`](../../configs/experiments/DiT/config_mf_b.py), [`DiT/config_mf_xl.py`](../../configs/experiments/DiT/config_mf_xl.py), [`WanT2V/config_mf.py`](../../configs/experiments/WanT2V/config_mf.py) +**Configs:** [`EDM/config_mf_cifar10.py`](../../configs/experiments/EDM/config_mf_cifar10.py), [`DiT/config_mf_b.py`](../../configs/experiments/DiT/config_mf_b.py), [`DiT/config_mf_xl.py`](../../configs/experiments/DiT/config_mf_xl.py), [`WanT2V/config_mf.py`](../../configs/experiments/WanT2V/config_mf.py), [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) (Stage 1 of [AnyFlow](../distribution_matching/README.md#anyflow)) **Expected results:** diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index 6d4754b..44e99b6 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -284,24 +284,17 @@ def _fuse_r_embedding( ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: """Combine the t- and r-time embeddings into the final timestep projection. - * ``additive`` (MeanFlow default, ``r_embedder.fusion_mode`` absent or "additive"): - ``r_timestep_proj = r_embedder.time_proj(r_embedder.act_fn(remb))`` is added to - ``timestep_proj`` and ``remb`` is added to ``temb``. T2V dual-stream variants - (``encoder_depth`` set) instead replace ``temb`` with ``remb`` and leave - ``timestep_proj`` alone. - - * ``gated`` (opt-in): convex-combine the two embeddings *before* the - projection — ``rt_emb = (1-g)·temb + g·remb`` then - ``timestep_proj = condition_embedder.time_proj(act_fn(rt_emb))``. The - projection is the condition_embedder's own, i.e. SHARED with the t-only - path; ``Wan.__init__`` therefore drops the redundant - ``r_embedder.time_proj`` in gated mode so it cannot silently drift away - from the shared one during training. With ``encoder_depth`` - set, the first ``encoder_depth`` blocks keep the t-only ``timestep_proj`` - and the remaining blocks switch to the gated projection (mirroring additive). - - Returns ``(temb, timestep_proj, r_timestep_proj)`` where ``r_timestep_proj`` - is ``None`` when the r-information is already folded into ``timestep_proj``. + * ``additive``: project each embedding separately and sum — ``temb + remb``, + ``timestep_proj + r_embedder.time_proj(...)``. + * ``gated``: interpolate first, ``(1-g)·temb + g·remb``, then project the + blend through the condition_embedder's own ``time_proj`` — shared with the + t-only path, so ``Wan.__init__`` drops the unused ``r_embedder.time_proj`` + rather than let it drift during training. + + Returns ``(temb, timestep_proj, r_timestep_proj)``: ``temb`` [B, D] modulates + the output head (``scale_shift_table + temb``), ``timestep_proj`` [B, 6, D] is + the AdaLN modulation every block consumes, and ``r_timestep_proj`` is what the + blocks switch to at ``encoder_depth`` — ``None`` when there is no switch. """ fusion = getattr(self.r_embedder, "fusion_mode", "additive") @@ -318,7 +311,7 @@ def _fuse_r_embedding( r_ts_proj = self.r_embedder.time_proj(self.r_embedder.act_fn(remb)) r_ts_proj = unflatten_timestep_proj(r_ts_proj, rs_seq_len) if self.encoder_depth is None: - return temb + remb, timestep_proj + r_ts_proj, r_ts_proj + return temb + remb, timestep_proj + r_ts_proj, None return remb, timestep_proj, r_ts_proj @@ -1140,14 +1133,6 @@ def forward( condition = torch.stack(condition, dim=0) if isinstance(condition, list) else condition - # A gated-fusion flow-map network always consumes the r-pathway (the - # fusion happens inside the shared time projection), so r=None has no - # in-distribution meaning for it. Default to r=t, i.e. the - # instantaneous velocity u(x_t, t, t). Additive-fusion (MeanFlow) - # networks keep the skip-r-pathway behavior. - if r is None and getattr(self.transformer.r_embedder, "fusion_mode", None) == "gated": - r = t - timestep_mask = torch.ones_like(x_t[:, 0]) # shape: [batch_size, num_latent_frames, H, W] timestep = self._compute_timestep_inputs(t, timestep_mask) r_timestep = None if r is None else self._compute_timestep_inputs(r, timestep_mask) diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index 7e8f6cb..e0bd477 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -537,7 +537,9 @@ def test_fuse_r_embedding_additive_unchanged(): r_proj = fake.r_embedder.time_proj(fake.r_embedder.act_fn(remb)).unflatten(1, (6, -1)) assert torch.allclose(out_temb, temb + remb) assert torch.allclose(out_proj, timestep_proj + r_proj) - assert torch.allclose(out_r_proj, r_proj) + # r is already folded into the two above, and without encoder_depth the + # blocks have nothing to switch to. + assert out_r_proj is None @pytest.mark.parametrize("prefix", ["", "transformer."]) From 23c6ca139e32e6f6163c58503a29e65c29b88dc8 Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Thu, 13 Aug 2026 17:49:21 +0000 Subject: [PATCH 15/19] Finalize integration and configs Signed-off-by: Julius Berner --- .../experiments/WanT2V/config_anyflow.py | 2 +- .../WanT2V/config_anyflow_onpolicy.py | 2 +- fastgen/configs/methods/config_anyflow.py | 11 + fastgen/configs/methods/config_mean_flow.py | 11 +- fastgen/methods/consistency_model/README.md | 3 +- .../methods/consistency_model/mean_flow.py | 213 +++++++++--------- .../methods/distribution_matching/README.md | 16 +- .../methods/distribution_matching/anyflow.py | 10 - tests/test_anyflowmodel.py | 2 +- 9 files changed, 135 insertions(+), 135 deletions(-) diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py index f2efd88..46fccf7 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -67,7 +67,7 @@ def create_config(): # learns the guided flow directly. guidance_scale stays None — MeanFlow's # target-side eq. 19 fusion is a different mechanism. config.model.guidance_scale = None - config.model.loss_config.guidance_fuse_scale = 3.0 + config.model.guidance_fuse_scale = 3.0 config.model.cond_dropout_prob = 0.1 # ------ (t, r) sampling: shifted uniform pairs + AnyFlow buckets ------ diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index f37e757..6dc2a65 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -112,7 +112,7 @@ def create_config(): config.model.loss_config.use_jvp_finite_diff = True config.model.loss_config.jvp_finite_diff_eps = 5e-3 config.model.loss_config.rebalance_to_flow_matching = True - config.model.loss_config.guidance_fuse_scale = 3.0 + config.model.guidance_fuse_scale = 3.0 config.model.cond_dropout_prob = 0.1 config.model.precision_amp_jvp = "float32" config.model.sample_t_cfg.flow_matching_ratio = 0.5 diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py index 10dc755..ff5d80c 100644 --- a/fastgen/configs/methods/config_anyflow.py +++ b/fastgen/configs/methods/config_anyflow.py @@ -59,6 +59,10 @@ class ModelConfig(DMD2ModelConfig): # the fixed student_sample_steps. student_sample_steps_list: Optional[List[int]] = None + # Prediction-side guidance fusion for the co-trained flow-map loss; see + # `config_mean_flow.ModelConfig`. + guidance_fuse_scale: Optional[float] = None + # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). cond_dropout_prob: Optional[float] = None cond_keys_no_dropout: List[str] = attrs.field(factory=list) @@ -66,6 +70,13 @@ class ModelConfig(DMD2ModelConfig): # Precision for autocast in the co-trained loss JVP (None = training precision). precision_amp_jvp: str | None = None + # MeanFlow's target-side guidance knobs. AnyFlow typixcally guides on the + # PREDICTION side (`guidance_fuse_scale`), so these stay at their + # no-op defaults + guidance_mixture_ratio: Optional[float] = None + guidance_t_start: float = 0.0 + guidance_t_end: float = 1.0 + @attrs.define(slots=False) class Config(BaseConfig): diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index 4c3a72e..bdb843d 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -86,11 +86,6 @@ class LossConfig: # ("beta08", "gaussian", "uniform"). Multiplies the adaptive norm_method # weight above; None disables it. weight_type: Optional[str] = None - # prediction-side guidance fusion scale (AnyFlow guidance distillation): - # the conditional output is trained to be the guided flow directly via - # (u_cond + (g-1) * u_uncond) / g against the raw data velocity. None - # keeps MeanFlow's target-side guidance (guidance_scale). - guidance_fuse_scale: Optional[float] = None # rebalance the flow-map / consistency (r < t) sample losses to the global # flow-matching (r = t) loss mean via a detached per-sample factor rebalance_to_flow_matching: bool = False @@ -116,6 +111,12 @@ class ModelConfig(BaseModelConfig): # optimizer net_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(RAdamOptimizerConfig)) + # prediction-side guidance fusion scale (AnyFlow guidance distillation): + # the conditional output is trained to be the guided flow directly via + # (u_cond + (g-1) * u_uncond) / g against the raw data velocity. None + # keeps MeanFlow's target-side guidance (guidance_scale). + guidance_fuse_scale: Optional[float] = None + # condition dropout probability cond_dropout_prob: Optional[float] = None diff --git a/fastgen/methods/consistency_model/README.md b/fastgen/methods/consistency_model/README.md index fdd0cea..ad0e82b 100644 --- a/fastgen/methods/consistency_model/README.md +++ b/fastgen/methods/consistency_model/README.md @@ -73,7 +73,7 @@ Two-stage training: frozen Stage-1 CM for `t < transition_t`, trainable student ## MeanFlow -**File:** [`mean_flow.py`](mean_flow.py) | **Reference:** [Geng et al., 2025](https://arxiv.org/abs/2505.13447) +**File:** [`mean_flow.py`](mean_flow.py) | **References:** [Geng et al., 2025](https://arxiv.org/abs/2505.13447), [Sabour et al., 2025](https://arxiv.org/abs/2506.14603) Learns average velocity between trajectory points: `x_r = x_t - (t-r) · u(x_t, t, r)`. @@ -82,6 +82,7 @@ Learns average velocity between trajectory points: `x_r = x_t - (t-r) · u(x_t, - `loss_config.use_jvp_finite_diff`: Use finite difference for JVP (e.g., for compatibility with Flash Attention and FSDP) - `sample_t_cfg`, `sample_r_cfg`: Configs of the distributions for sampling `t` and `r` - `sample_t_cfg.flow_matching_ratio`: Fraction of the batch with `r = t` (flow matching loss) +- `guidance_scale`, `guidance_fuse_scale`: If `guidance_fuse_scale` is `None`, use target-side guidance (`guidance_scale`, 3-way mixture via `guidance_mixture_ratio`, t-window via `guidance_t_start` / `guidance_t_end`), otherwise use prediction-side guidance (conditional output learns the guided flow directly) **Configs:** [`EDM/config_mf_cifar10.py`](../../configs/experiments/EDM/config_mf_cifar10.py), [`DiT/config_mf_b.py`](../../configs/experiments/DiT/config_mf_b.py), [`DiT/config_mf_xl.py`](../../configs/experiments/DiT/config_mf_xl.py), [`WanT2V/config_mf.py`](../../configs/experiments/WanT2V/config_mf.py), [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) (Stage 1 of [AnyFlow](../distribution_matching/README.md#anyflow)) diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 0cbd4e6..f42648d 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -125,6 +125,90 @@ def _drop_condition(self, condition: Any, neg_condition: Any) -> Tuple[Any, Opti return condition, keep raise TypeError(f"Unsupported condition type: {type(condition)}") + @torch.no_grad() + def _get_velocity( + self, + x: torch.Tensor, + z: torch.Tensor, + t: torch.Tensor, + condition: Optional[Any] = None, + neg_condition: Optional[Any] = None, + ) -> Tuple[Any, torch.Tensor]: + """Regression target for the flow-map loss, plus the condition it was built from. + + Two independent choices: + + * ``loss_config.use_cd`` picks the target SOURCE -- the teacher, or the + conditional data velocity. + * ``guidance_fuse_scale`` guides the *prediction* in + ``_compute_mf_loss``, so the target needs no guidance of its own and we only + drop the condition. Otherwise we guide the target too: through the teacher, or + -- without one -- by the net's own cond/uncond pass. + """ + 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)" + condition, _ = self._drop_condition(condition, neg_condition) + + x_t = self.net.noise_scheduler.forward_process(x, z, t) + + if self.loss_config.use_cd: + dxt_dt = self.teacher(x_t, t, condition=condition, fwd_pred_type="flow") + # Under fusion the target stays unguided, or it would be guided twice. + if self.config.guidance_scale is not None and fuse_scale is None: + guidance_scale = torch.where( + ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), + self.config.guidance_scale, + 1.0, + ) + guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) + neg_dxt_dt = self.teacher(x_t, t, condition=neg_condition, fwd_pred_type="flow") + dxt_dt = dxt_dt + (guidance_scale - 1.0) * (dxt_dt - neg_dxt_dt) + else: + dxt_dt = self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) + + # unconditional score estimation from meanflow eq (19). Skipped under + # prediction-side fusion: that is this same guidance, moved onto the + # prediction, and the dropout it needs already ran above. + if fuse_scale is None and ( + self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None + ): + # Turn off dropout + self.net.eval() + neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow") + guidance_scale = self.config.guidance_scale or 1.0 + guidance_scale = torch.where( + ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), + guidance_scale, + 1.0, + ) + guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) + + if self.config.guidance_mixture_ratio is None: + guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt) + else: + guidance_mixture_ratio = torch.where( + ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), + self.config.guidance_mixture_ratio, + 0.0, + ) + guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype) + cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow") + guided_dxt_dt = ( + guidance_scale * dxt_dt + + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt + + guidance_mixture_ratio * cond_dxt_dt + ) + + self.net.train() + condition, keep = self._drop_condition(condition, neg_condition) + if keep is not None: + # Same subset: a kept sample is conditional + guided, a + # dropped one unconditional + unguided. + dxt_dt = torch.where(expand_like(keep, dxt_dt), guided_dxt_dt, dxt_dt) + + return condition, dxt_dt + def _estimate_jvp_finite_difference( self, net_wrapper: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor], @@ -489,50 +573,34 @@ def _compute_mf_loss( z = torch.randn_like(real_data) x_t = self.net.noise_scheduler.forward_process(real_data, z, t) - guidance_fuse_scale = getattr(self.loss_config, "guidance_fuse_scale", None) + condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) + # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity + torch.clear_autocast_cache() + u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) + assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" + + # additional forward pass to get u_theta with gradient; see also https://github.com/Gsunshine/py-meanflow?tab=readme-ov-file#note-on-jvp + assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" + u_theta = self.net( + x_t, + t, + r=r, + condition=condition, + fwd_pred_type="flow", + ) + + guidance_fuse_scale = self.config.guidance_fuse_scale if guidance_fuse_scale is not None: - # Guidance distillation fused on the PREDICTION side: the - # conditional output is trained to be the guided flow directly. - # The regression target stays the raw data velocity, the - # unconditional branch is queried at the SAME (t, r) flow-map - # slice, and the effective prediction is - # (u_cond + (g - 1) * u_uncond) / g. Text dropout replaces the - # condition with neg_condition for a random subset beforehand. - g = float(guidance_fuse_scale) - assert g > 0, f"guidance_fuse_scale must be > 0, got {g} (set it to None to disable guidance fusion)" - assert ( - neg_condition is not None - ), "guidance_fuse_scale requires neg_condition: the unconditional branch is queried at the same (t, r)" - condition, _ = self._drop_condition(condition, neg_condition) - dxt_dt = self.net.noise_scheduler.cond_velocity(x=real_data, eps=z, t=t) - torch.clear_autocast_cache() - # dF/dt of the fused prediction: finite difference of the - # conditional output divided by g (the unconditional derivative is - # dropped). - u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) / g - assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" - - assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" - u_theta = self.net(x_t, t, r=r, condition=condition, fwd_pred_type="flow") + # Guidance distillation on the PREDICTION side (see `_get_velocity`): the + # conditional output learns the guided flow directly, so only the prediction + # changes. The uncond branch is queried at the SAME (t, r) flow-map slice, + # giving (u_cond + (g - 1) * u_uncond) / g; dF/dt is then the conditional + # finite difference over g, with the unconditional derivative dropped. + u_theta_jvp = u_theta_jvp / guidance_fuse_scale with torch.no_grad(): u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow") - u_theta = (u_theta + (g - 1.0) * u_uncond) / g - else: - condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) - # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity - torch.clear_autocast_cache() - u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) - assert not u_theta_jvp.requires_grad, "u_theta_jvp should not require gradients" - - # additional forward pass to get u_theta with gradient; see also https://github.com/Gsunshine/py-meanflow?tab=readme-ov-file#note-on-jvp - assert x_t.dtype == real_data.dtype, f"x_t.dtype: {x_t.dtype}, real_data.dtype: {real_data.dtype}" - u_theta = self.net( - x_t, - t, - r=r, - condition=condition, - fwd_pred_type="flow", - ) + u_theta = (u_theta + (guidance_fuse_scale - 1.0) * u_uncond) / guidance_fuse_scale + mf_loss, tangent, loss_weight, warmup_weight = self._mf_pred_to_loss( u_theta=u_theta, u_theta_jvp=u_theta_jvp, x_t=x_t, dxt_dt=dxt_dt, t=t, r=r, iteration=iteration ) @@ -564,69 +632,6 @@ def __init__(self, config: ModelConfig): self.config = config self._init_flow_map_loss() - @torch.no_grad() - def _get_velocity( - self, - x: torch.Tensor, - z: torch.Tensor, - t: torch.Tensor, - condition: Optional[torch.Tensor] = None, - neg_condition: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: - x_t = self.net.noise_scheduler.forward_process(x, z, t) - - if self.loss_config.use_cd: - dxt_dt = self.teacher(x_t, t, condition=condition, fwd_pred_type="flow") - if self.config.guidance_scale is not None: - guidance_scale = torch.where( - ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - self.config.guidance_scale, - 1.0, - ) - guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) - neg_dxt_dt = self.teacher(x_t, t, condition=neg_condition, fwd_pred_type="flow") - dxt_dt = dxt_dt + (guidance_scale - 1.0) * (dxt_dt - neg_dxt_dt) - else: - dxt_dt = self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) - - # unconditional score estimation from meanflow eq (19) - if self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None: - # Turn off dropout - self.net.eval() - neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow") - guidance_scale = self.config.guidance_scale or 1.0 - guidance_scale = torch.where( - ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - guidance_scale, - 1.0, - ) - guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) - - if self.config.guidance_mixture_ratio is None: - guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt) - else: - guidance_mixture_ratio = torch.where( - ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - self.config.guidance_mixture_ratio, - 0.0, - ) - guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype) - cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow") - guided_dxt_dt = ( - guidance_scale * dxt_dt - + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt - + guidance_mixture_ratio * cond_dxt_dt - ) - - self.net.train() - condition, keep = self._drop_condition(condition, neg_condition) - if keep is not None: - # Same subset: a kept sample is conditional + guided, a - # dropped one unconditional + unguided. - dxt_dt = torch.where(expand_like(keep, dxt_dt), guided_dxt_dt, dxt_dt) - - return condition, dxt_dt - def single_train_step( self, data: Dict[str, Any], iteration: int ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index 41e860a..16561b8 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -89,23 +89,15 @@ DMD2 extended for causal video generation with autoregressive chunk-by-chunk pro **File:** [`anyflow.py`](anyflow.py) | **Reference:** [Gu et al., 2026](https://arxiv.org/abs/2605.13724) -DMD2 with a flow-map student `u_θ(x_t, t, r)` (average velocity from `t` back to `r`), supporting arbitrary inference NFE. The student generates by rolling out the flow map from pure noise with the NFE sampled per iteration and gradients through all segments, and co-trains the flow-map loss at every update. Requires a Stage-1 flow-map pretrain, which uses the MeanFlow objective with AnyFlow's hyperparameters and runs directly on [`MeanFlowModel`](../consistency_model/mean_flow.py). +DMD2 with a flow-map student `u(x_t, t, r)` (mean velocity from `t` back to `r`), supporting arbitrary inference NFE. The student generates by rolling out the flow map from pure noise with the NFE sampled per iteration and gradients through all segments, and co-trains the flow-map loss at every update. Requires a Stage-1 flow-map pretrain, which runs directly on [`MeanFlowModel`](../consistency_model/mean_flow.py) using AnyFlow's hyperparameters. -**Key Parameters (Stage 2, on-policy):** +**Key Parameters (Stage 2):** - `student_sample_steps_list`: Rollout NFEs sampled per iteration - `cotrain_pretrain_weight`: Weight of the co-trained Stage-1 flow-map loss -- `guidance_scale`: CFG scale for teacher, applied as `cond + (g-1)·(cond - uncond)` - Requires a Stage-1 pretrained checkpoint via `trainer.checkpointer.pretrained_ckpt_path` -- See also the key parameters of DMD2 above - -**Key Parameters (Stage 1, on MeanFlow):** -- `loss_config.weight_type`: Fixed per-timestep loss weight (`beta08`, `gaussian`, `uniform`) -- `loss_config.guidance_fuse_scale`: Prediction-side guidance fusion, i.e., the conditional output learns the guided flow directly (with `cond_dropout_prob` for text dropout) -- `loss_config.rebalance_to_flow_matching`: Rescale the `r < t` losses to the global flow-matching (`r = t`) loss mean -- `sample_t_cfg`: Shifted timestep sampling (`time_dist_type="shifted"`, `shift=5` for Wan video) and the batch fraction pinned to `r = 0` (`consistency_ratio`) -- See also the key parameters of [MeanFlow](../consistency_model/README.md#meanflow) +- See also the key parameters of DMD2 above (and the key parameters of [MeanFlow](../consistency_model/README.md#meanflow) for Stage 1) -**Configs:** [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) (Stage 1 pretrain), [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) (Stage 2 on-policy distillation, full-rank fine-tune of a Stage 1 checkpoint) +**Configs:** [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) (Stage 1 pretrain), [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) (Stage 2 on-policy distillation) --- diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index ba61c82..96d345e 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -42,16 +42,6 @@ class AnyFlowModel(FlowMapLossMixin, DMD2Model): FastGenModel's default x0-prediction loop never passes ``r``. """ - @torch.no_grad() - def _get_velocity(self, x, z, t, condition=None, neg_condition=None): - """Raw data velocity for the co-trained flow-map loss. - - Note that AnyFlow guides on the prediction side using - ``loss_config.guidance_fuse_scale`` applied in ``_compute_mf_loss``. - """ - del neg_condition - return condition, self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) - def __init__(self, config: ModelConfig): super().__init__(config) self.config = config diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index e0bd477..65ea86b 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -243,7 +243,7 @@ def test_pretrain_prediction_side_guidance_fusion(): """The AnyFlow guidance-distillation branch (guidance_fuse_scale) must run end to end and keep gradients on the fused prediction.""" model = _build_pretrain_model() - model.loss_config.guidance_fuse_scale = 3.0 + model.config.guidance_fuse_scale = 3.0 model.config.cond_dropout_prob = 0.5 data = _make_data(model, batch_size=2) From 68a83a406861908dcf28a9730cd8a616b4f2f46b Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Thu, 13 Aug 2026 21:49:58 +0000 Subject: [PATCH 16/19] Address comments Signed-off-by: Julius Berner --- .../experiments/WanT2V/config_anyflow.py | 6 +- .../WanT2V/config_anyflow_onpolicy.py | 2 +- fastgen/configs/methods/config_anyflow.py | 9 +- fastgen/configs/methods/config_mean_flow.py | 2 +- .../methods/consistency_model/mean_flow.py | 59 ++++---- .../methods/distribution_matching/anyflow.py | 5 +- fastgen/networks/Wan/utils.py | 9 +- tests/test_anyflowmodel.py | 6 +- tests/test_meanflowmodel.py | 134 ++++++++++++++++++ tests/test_network_fsdp.py | 4 +- 10 files changed, 192 insertions(+), 44 deletions(-) diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py index 46fccf7..5a630c2 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -4,8 +4,8 @@ """AnyFlow flow-map pretrain config on Wan-1.3B T2V (paper Stage 1). AnyFlow's pretrain objective is MeanFlow's with a fixed ``beta08`` per-timestep -weighting, a finite-difference JVP, shifted timestep sampling, and a -``consistency_ratio`` fraction of the batch pinned to ``r = 0`` — so this config +weighting, a finite-difference JVP, shifted timestep sampling, and a +``consistency_ratio`` fraction of the batch pinned to ``r = 0`` — so this config runs``MeanFlowModel`` directly. The values below mirror the reference recipe ``train_wan1b_student_shift5_81f_480p_lr5e-5_6k_b32.yml``. Known deviations from the reference: full-rank fine-tuning instead of the paper's @@ -103,7 +103,7 @@ def create_config(): config.model.sample_t_cfg.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625, 0.0] # ------ data / trainer ------ - config.dataloader_train = VideoLoaderConfig + config.dataloader_train = copy.deepcopy(VideoLoaderConfig) config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 config.dataloader_train.batch_size = 1 diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index 6dc2a65..0c2cef4 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -137,7 +137,7 @@ def create_config(): config.trainer.callbacks.ema.start_iter = 6399 # ------ data / trainer ------ - config.dataloader_train = VideoLoaderConfig + config.dataloader_train = copy.deepcopy(VideoLoaderConfig) config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 config.dataloader_train.batch_size = 1 diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py index ff5d80c..c7084a2 100644 --- a/fastgen/configs/methods/config_anyflow.py +++ b/fastgen/configs/methods/config_anyflow.py @@ -11,8 +11,6 @@ from typing import List, Optional -from typing import Optional - import attrs from omegaconf import DictConfig @@ -61,6 +59,11 @@ class ModelConfig(DMD2ModelConfig): # Prediction-side guidance fusion for the co-trained flow-map loss; see # `config_mean_flow.ModelConfig`. + # + # Deviates from the AnyFlow reference, which divides dF/dt by g for the whole + # batch. We divide only on samples that kept their condition: a dropped sample's + # fused prediction is plain `u_uncond`, so an ungated 1/g would regress it onto + # a different fixed point. guidance_fuse_scale: Optional[float] = None # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). @@ -70,7 +73,7 @@ class ModelConfig(DMD2ModelConfig): # Precision for autocast in the co-trained loss JVP (None = training precision). precision_amp_jvp: str | None = None - # MeanFlow's target-side guidance knobs. AnyFlow typixcally guides on the + # MeanFlow's target-side guidance knobs. AnyFlow typically guides on the # PREDICTION side (`guidance_fuse_scale`), so these stay at their # no-op defaults guidance_mixture_ratio: Optional[float] = None diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index bdb843d..b491db2 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -83,7 +83,7 @@ class LossConfig: # loss type (choice between l2 and opt_grad) loss_type: str = "opt_grad" # optional fixed per-timestep loss weighting evaluated as a function of t - # ("beta08", "gaussian", "uniform"). Multiplies the adaptive norm_method + # ("beta08", "gaussian", "uniform"). Multiplies the adaptive norm_method # weight above; None disables it. weight_type: Optional[str] = None # rebalance the flow-map / consistency (r < t) sample losses to the global diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index f42648d..4d4411b 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -59,9 +59,9 @@ class FlowMapLossMixin: training objective; distribution-matching methods can co-train it alongside their own objective. - The host class must provide ``net``, ``device``, ``config``, - ``precision_amp``, a ``_get_velocity`` implementation, and call - ``_init_flow_map_loss`` from its ``__init__``. + The host class must provide ``net``, ``device``, ``config`` and ``precision_amp`` + (plus ``teacher`` when ``loss_config.use_cd``), and call ``_init_flow_map_loss`` + from its ``__init__``. """ def _init_flow_map_loss(self) -> None: @@ -70,7 +70,8 @@ def _init_flow_map_loss(self) -> None: self.sample_r_cfg = self.config.sample_r_cfg self.loss_config = self.config.loss_config - # Precision for JVP + # Drop the JVP autocast when it matches the outer one, since a nested region in the + # same dtype is a no-op. if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp: self.precision_amp_jvp = None else: @@ -89,28 +90,29 @@ def _init_flow_map_loss(self) -> None: grid = shift * grid / (1 + (shift - 1) * grid) self._timestep_weight_scale = float(num_steps / self._timestep_weight_raw(grid).sum()) - def _drop_condition(self, condition: Any, neg_condition: Any) -> Tuple[Any, Optional[torch.Tensor]]: + def _drop_condition( + self, condition: Any, neg_condition: Any, batch_size: int, device: torch.device + ) -> Tuple[Any, torch.Tensor]: """Replace the condition with neg_condition for a per-sample subset. - Returns ``(condition, keep)``; ``keep`` is the ``[B]`` bool mask (None if - no dropout), so callers can reuse the same subset. Keys in - ``cond_keys_no_dropout`` are never replaced. + Returns ``(condition, keep)``; ``keep`` is the ``[B]`` bool mask of the + samples that stayed conditional, so callers can reuse the same subset. ``deterministic_buckets`` decides whether an index carries bucket information: if so the buckets are cut on the GLOBAL index and an index-based rule would only hit flow matching on rank 0, so draw per sample; otherwise drop the first ``num_to_drop``. """ + # Dropout disabled, or no negative condition to swap in: every sample + # stays conditional. if self.config.cond_dropout_prob is None or neg_condition is None: - return condition, None + return condition, torch.ones(batch_size, dtype=torch.bool, device=device) - ref = neg_condition if isinstance(neg_condition, torch.Tensor) else next(iter(neg_condition.values())) - batch_size = ref.shape[0] if self.sample_t_cfg.deterministic_buckets: - keep = torch.rand(batch_size, device=ref.device) >= self.config.cond_dropout_prob + keep = torch.rand(batch_size, device=device) >= self.config.cond_dropout_prob else: - num_to_drop = (torch.rand(batch_size, device=ref.device) < self.config.cond_dropout_prob).sum() - keep = torch.arange(batch_size, device=ref.device) >= num_to_drop + num_to_drop = (torch.rand(batch_size, device=device) < self.config.cond_dropout_prob).sum() + keep = torch.arange(batch_size, device=device) >= num_to_drop if isinstance(condition, torch.Tensor): return torch.where(expand_like(keep, condition), condition, neg_condition), keep @@ -133,8 +135,9 @@ def _get_velocity( t: torch.Tensor, condition: Optional[Any] = None, neg_condition: Optional[Any] = None, - ) -> Tuple[Any, torch.Tensor]: - """Regression target for the flow-map loss, plus the condition it was built from. + ) -> Tuple[Any, torch.Tensor, torch.Tensor]: + """Regression target for the flow-map loss, the condition it was built from, + and the ``[B]`` mask of samples that stayed conditional. Two independent choices: @@ -148,7 +151,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)" - condition, _ = self._drop_condition(condition, neg_condition) + 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) x_t = self.net.noise_scheduler.forward_process(x, z, t) @@ -201,13 +206,12 @@ def _get_velocity( ) self.net.train() - condition, keep = self._drop_condition(condition, neg_condition) - if keep is not None: - # Same subset: a kept sample is conditional + guided, a - # dropped one unconditional + unguided. - dxt_dt = torch.where(expand_like(keep, dxt_dt), guided_dxt_dt, dxt_dt) + condition, keep = self._drop_condition(condition, neg_condition, x_t.shape[0], x_t.device) + # Same subset: a kept sample is conditional + guided, a dropped one + # unconditional + unguided. + dxt_dt = torch.where(expand_like(keep, dxt_dt), guided_dxt_dt, dxt_dt) - return condition, dxt_dt + return condition, dxt_dt, keep def _estimate_jvp_finite_difference( self, @@ -573,7 +577,7 @@ def _compute_mf_loss( z = torch.randn_like(real_data) x_t = self.net.noise_scheduler.forward_process(real_data, z, t) - condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) + condition, dxt_dt, keep = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity torch.clear_autocast_cache() u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) @@ -594,9 +598,10 @@ def _compute_mf_loss( # Guidance distillation on the PREDICTION side (see `_get_velocity`): the # conditional output learns the guided flow directly, so only the prediction # changes. The uncond branch is queried at the SAME (t, r) flow-map slice, - # giving (u_cond + (g - 1) * u_uncond) / g; dF/dt is then the conditional - # finite difference over g, with the unconditional derivative dropped. - u_theta_jvp = u_theta_jvp / guidance_fuse_scale + # giving (u_cond + (g - 1) * u_uncond) / g; dF/dt is then the finite + # difference over g on conditional samples, with the unconditional + # derivative dropped. + u_theta_jvp = torch.where(expand_like(keep, u_theta_jvp), u_theta_jvp / guidance_fuse_scale, u_theta_jvp) with torch.no_grad(): u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow") u_theta = (u_theta + (guidance_fuse_scale - 1.0) * u_uncond) / guidance_fuse_scale diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index 96d345e..0cb3136 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -38,8 +38,7 @@ class AnyFlowModel(FlowMapLossMixin, DMD2Model): """AnyFlow on-policy stage: DMD2 with a flow-map rollout student. ``FlowMapLossMixin`` supplies the co-trained Stage-1 objective and the - flow-map validation sample loop, which integrates with ``r = t_next`` — - FastGenModel's default x0-prediction loop never passes ``r``. + flow-map validation sample loop, which integrates with ``r = t_next``. """ def __init__(self, config: ModelConfig): @@ -133,7 +132,7 @@ def gen_data_from_net( grad_step = self._broadcast_choice(num_steps) t_list = self._rollout_t_list(num_steps) - # The leading jump exists only for grad_step > 0 and the trailing one + # The leading jump exists only for grad_step > 0 and the trailing one # only for grad_step + 1 < num_steps. seg_t = [t_list[0]] if grad_step > 0 else [] seg_t += [t_list[grad_step], t_list[grad_step + 1]] diff --git a/fastgen/networks/Wan/utils.py b/fastgen/networks/Wan/utils.py index f354c00..9b09141 100644 --- a/fastgen/networks/Wan/utils.py +++ b/fastgen/networks/Wan/utils.py @@ -129,6 +129,13 @@ def remap_anyflow_keys(state_dict: Mapping[str, Any]) -> Mapping[str, Any]: # [transformer.]condition_embedder.delta_embedder.linear_1.weight # -> [transformer.]r_embedder.time_embedder.linear_1.weight prefix, _, suffix = k.partition(delta_marker) - new_sd[f"{prefix}r_embedder.time_embedder.{suffix}"] = new_sd.pop(k) + target = f"{prefix}r_embedder.time_embedder.{suffix}" + if target in new_sd: + raise ValueError( + f"remap_anyflow_keys: rewriting {k!r} would overwrite the existing {target!r}. " + "This checkpoint carries both the AnyFlow and the FastGen r-pathway layouts; " + "drop one of them before loading." + ) + new_sd[target] = new_sd.pop(k) logger.info(f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors into r_embedder.") return new_sd diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index 65ea86b..6e537ef 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -343,8 +343,10 @@ def counting_forward(*args, **kwargs): return orig_forward(*args, **kwargs) model.net.forward = counting_forward - gen = model.gen_data_from_net(input_student, t_student, condition=cond) - model.net.forward = orig_forward + try: + gen = model.gen_data_from_net(input_student, t_student, condition=cond) + finally: + model.net.forward = orig_forward assert len(calls) <= 3, f"rollout must compress to <= 3 forwards, got {len(calls)}" assert gen.requires_grad diff --git a/tests/test_meanflowmodel.py b/tests/test_meanflowmodel.py index 08a88db..39446c0 100644 --- a/tests/test_meanflowmodel.py +++ b/tests/test_meanflowmodel.py @@ -120,3 +120,137 @@ def test_single_train_step_update_fp32_jvp(): assert "mf_loss" in loss_map assert "gen_rand" in outputs assert isinstance(outputs["gen_rand"], Callable) + + +@pytest.mark.parametrize("cond_dropout_prob, expect_guided", [(None, True), (0.0, True), (1.0, False)]) +def test_target_side_guidance_applies_when_no_cond_dropout(cond_dropout_prob, expect_guided): + """`guidance_scale` must reach the target for every conditional sample. + + No dropout means every sample is conditional, so it must guide exactly like + `p=0.0`. It used not to: `_drop_condition` returned `keep=None` when + `cond_dropout_prob is None` and the caller skipped the update entirely -- + computing `guided_dxt_dt` at the cost of an extra forward pass and then + discarding it, so `guidance_scale` was silently inert on that path. + (Pre-existing: `_mix_condition` returned early on `cond_dropout_prob is None` + before the AnyFlow work.) `_drop_condition` now always returns a mask -- all-True + here -- so the caller has no special case left to forget. + + The net's forward is stubbed to depend only on the condition, so the guided + velocity is exact rather than initialization-dependent -- EDM's `SongUNet` + zero-inits `out_conv`, which would otherwise make the cond and uncond passes + bitwise identical and guidance a provable no-op. + """ + config = create_config() + instance = config.model + opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"] + instance.net = override_config_with_opts(instance.net, opts) + instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + instance.precision = "float32" + instance.pretrained_model_path = "" + instance.input_shape = [3, 2, 2] + instance.cond_dropout_prob = cond_dropout_prob + instance.guidance_scale = 2.0 + instance.guidance_fuse_scale = None + model = MeanFlowModel(instance) + model.on_train_begin() + + batch_size = 4 + real = torch.randn(batch_size, 3, 2, 2, device=model.device, dtype=torch.float32) + z = torch.randn_like(real) + t = torch.full((batch_size,), 0.5, device=model.device, dtype=model.net.noise_scheduler.t_precision) + condition = torch.nn.functional.one_hot(torch.arange(batch_size) % 10, num_classes=10) + condition = condition.to(model.device, torch.float32) + neg_condition = torch.zeros(batch_size, 10, device=model.device, dtype=torch.float32) + + # Depends ONLY on the condition: the all-zero neg pass returns exactly 0. + def condition_only_forward(x_t, t, **kwargs): + val = kwargs["condition"].sum(dim=1).reshape(-1, 1, 1, 1).to(x_t.dtype) + return torch.ones_like(x_t) * val + + orig_forward = model.net.forward + model.net.forward = condition_only_forward + try: + cond_out, dxt_dt, _ = model._get_velocity(real, z, t, condition=condition, neg_condition=neg_condition) + finally: + model.net.forward = orig_forward + + # neg_dxt_dt == 0, so guided == neg + scale * (plain - neg) == 2 * plain. + plain = model.net.noise_scheduler.cond_velocity(x=real, eps=z, t=t) + dropped = (cond_out == neg_condition).all(dim=1) + assert bool(dropped.all()) is not expect_guided + + expected = 2.0 * plain if expect_guided else plain + assert torch.allclose(dxt_dt, expected), (dxt_dt - expected).abs().max() + + +def test_fused_jvp_scaling_is_gated_on_kept_samples(): + """Under `guidance_fuse_scale`, dF/dt must be divided by g only for samples that + stayed conditional. + + A dropped sample's fused prediction collapses to plain `u_uncond` (the fusion + self-cancels), so scaling its derivative too would regress it onto + `v - (t - r) * d(u_uncond)/dt / g` instead of the unconditional MeanFlow identity. + The AnyFlow reference scales the whole batch + (`compute_central_difference(..., guidance)` in + `far/trainers/trainer_wan_anyflow_pretrain.py`); we gate on `keep`. + + The net is stubbed to ignore `condition`, so dropping it changes nothing about the + raw derivative -- making the 1/g gate the ONLY difference between the two runs. + """ + g = 3.0 + config = create_config() + instance = config.model + opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"] + instance.net = override_config_with_opts(instance.net, opts) + instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + instance.precision = "float32" + instance.pretrained_model_path = "" + instance.input_shape = [3, 2, 2] + instance.guidance_fuse_scale = g + instance.loss_config.use_jvp_finite_diff = True + instance.sample_t_cfg.deterministic_buckets = False + model = MeanFlowModel(instance) + model.on_train_begin() + + batch_size = 4 + real = torch.randn(batch_size, 3, 2, 2, device=model.device, dtype=torch.float32) + t_prec = model.net.noise_scheduler.t_precision + t = torch.full((batch_size,), 0.6, device=model.device, dtype=t_prec) + r = torch.full((batch_size,), 0.2, device=model.device, dtype=t_prec) + condition = torch.nn.functional.one_hot(torch.arange(batch_size) % 10, num_classes=10) + condition = condition.to(model.device, torch.float32) + neg_condition = torch.zeros(batch_size, 10, device=model.device, dtype=torch.float32) + + # `+ 0.0 * param.sum()` leaves the value untouched but ties the output to the + # autograd graph: `_mf_pred_to_loss` asserts + # `u_theta.requires_grad is torch.is_grad_enabled()`. The JVP runs under `_jvp`'s + # `@torch.no_grad()`, so `u_theta_jvp` stays grad-free as that code also asserts. + param = next(model.net.parameters()) + + def condition_independent_forward(x_t, t, **kwargs): + tt = t.reshape(-1, *([1] * (x_t.ndim - 1))).to(x_t.dtype) + return x_t * (1.0 + tt) + 0.0 * param.sum().to(x_t.dtype) + + orig_forward = model.net.forward + + def run(cond_dropout_prob): + model.config.cond_dropout_prob = cond_dropout_prob + model.net.forward = condition_independent_forward + try: + torch.manual_seed(11) + return model._compute_mf_loss( + real_data=real, + t=t, + r=r, + iteration=0, + condition=condition, + neg_condition=neg_condition, + )[2] + finally: + model.net.forward = orig_forward + + jvp_kept = run(0.0) # every sample conditional -> every derivative divided by g + jvp_dropped = run(1.0) # every sample unconditional -> none divided + + assert torch.allclose(jvp_dropped, g * jvp_kept, atol=1e-5, rtol=1e-4), (jvp_dropped - g * jvp_kept).abs().max() + assert not torch.allclose(jvp_dropped, jvp_kept, atol=1e-6) diff --git a/tests/test_network_fsdp.py b/tests/test_network_fsdp.py index 7dc6e86..bd45f99 100644 --- a/tests/test_network_fsdp.py +++ b/tests/test_network_fsdp.py @@ -692,9 +692,7 @@ def _test_qwen_image_impl(rank: int, world_size: int) -> Dict: from fastgen.configs.net import QwenImageConfig set_env_vars() - return _generic_fsdp_test_impl( - rank, world_size, QwenImageConfig, generate_qwen_image_inputs, apply_checkpointing=True - ) + return _generic_fsdp_test_impl(world_size, QwenImageConfig, generate_qwen_image_inputs, apply_checkpointing=True) # ============================================================================= From a6b2497853d532ed114194db24602923c7d6a8aa Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Wed, 19 Aug 2026 20:42:13 +0000 Subject: [PATCH 17/19] Lint --- fastgen/methods/distribution_matching/dmd2.py | 4 +--- fastgen/methods/distribution_matching/self_forcing.py | 1 - fastgen/networks/Wan/network.py | 4 +--- 3 files changed, 2 insertions(+), 7 deletions(-) diff --git a/fastgen/methods/distribution_matching/dmd2.py b/fastgen/methods/distribution_matching/dmd2.py index c5bad9e..ed31a8b 100644 --- a/fastgen/methods/distribution_matching/dmd2.py +++ b/fastgen/methods/distribution_matching/dmd2.py @@ -87,9 +87,7 @@ def _sample_noising_time(self, batch_size: int, iteration: Optional[int] = None) is_fake_score_step = iteration is not None and iteration % self.config.student_update_freq != 0 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 - ) + 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 diff --git a/fastgen/methods/distribution_matching/self_forcing.py b/fastgen/methods/distribution_matching/self_forcing.py index ea8b24f..5d1e79c 100644 --- a/fastgen/methods/distribution_matching/self_forcing.py +++ b/fastgen/methods/distribution_matching/self_forcing.py @@ -12,7 +12,6 @@ import fastgen.utils.logging_utils as logger from fastgen.networks.network import CausalFastGenNetwork -from fastgen.utils.basic_utils import convert_cfg_to_dict from fastgen.utils.distributed import is_rank0, world_size if TYPE_CHECKING: diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index 44e99b6..9f1d67b 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -646,9 +646,7 @@ def __init__( self.transformer.time_cond_type = time_cond_type if r_timestep: logger.info(f"Initializing r embedder with {r_embedder_init}") - self.transformer.r_embedder = self.init_embedder( - r_embedder_init, r_embedder_fusion, r_embedder_gate_value - ) + self.transformer.r_embedder = self.init_embedder(r_embedder_init, r_embedder_fusion, r_embedder_gate_value) else: self.transformer.r_embedder = None From 2a1a495fe6d9bcfa09e563beb5f2bae28f2f2ada Mon Sep 17 00:00:00 2001 From: Enderfga Date: Thu, 20 Aug 2026 23:15:36 +0800 Subject: [PATCH 18/19] anyflow: restore the guidance-fusion neg_condition guard The assert from 01c8ff2 was lost in the 23c6ca1 rebase, so a batch without neg_condition now reaches the unconditional forward as None instead of stopping with a configuration error. Restored, with a test this time. Three smaller things from the same pass: - Both shifted time distributions apply the same shift map when sampling t, but the two places that rebuild that grid outside the sampler (the flow-map loss weight normalization and the AnyFlow rollout schedule) only recognized "shifted", so "shifted_logitnormal" would silently fall back to shift=1 there. Both now key off one tuple next to the sampler. No config hits this today -- shifted_logitnormal is only set on fake_score_sample_t_cfg -- but the two literals were bound to drift. - The unconditional pass in the prediction-side fusion ran with the net still in train mode; the target-side path switches to eval() around its own unconditional pass. Matched. - zip(strict=True) in the two test assertions newer ruff flags as B905. Signed-off-by: Enderfga --- .../methods/consistency_model/mean_flow.py | 11 ++++- .../methods/distribution_matching/anyflow.py | 3 +- fastgen/networks/noise_schedule.py | 9 +++- tests/test_anyflowmodel.py | 45 ++++++++++++++++++- 4 files changed, 63 insertions(+), 5 deletions(-) diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 4d4411b..4eaf824 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -9,6 +9,7 @@ import numpy as np import torch from fastgen.methods import CMModel +from fastgen.networks.noise_schedule import SHIFTED_TIME_DIST_TYPES from fastgen.utils import basic_utils, expand_like from fastgen.utils.basic_utils import convert_cfg_to_dict from fastgen.utils.distributed import get_rank, world_size @@ -85,7 +86,7 @@ def _init_flow_map_loss(self) -> None: self._timestep_weight_scale: Optional[float] = None if self.loss_config.weight_type is not None: num_steps = self.net.noise_scheduler.num_steps - shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 + shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type in SHIFTED_TIME_DIST_TYPES else 1.0 grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64)[:-1] grid = shift * grid / (1 + (shift - 1) * grid) self._timestep_weight_scale = float(num_steps / self._timestep_weight_raw(grid).sum()) @@ -151,6 +152,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)" 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) @@ -603,7 +607,12 @@ def _compute_mf_loss( # derivative dropped. u_theta_jvp = torch.where(expand_like(keep, u_theta_jvp), u_theta_jvp / guidance_fuse_scale, u_theta_jvp) with torch.no_grad(): + # Turn off dropout for the unconditional pass, as the target-side path does. + was_training = self.net.training + self.net.eval() u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow") + if was_training: + self.net.train() u_theta = (u_theta + (guidance_fuse_scale - 1.0) * u_uncond) / guidance_fuse_scale mf_loss, tangent, loss_weight, warmup_weight = self._mf_pred_to_loss( diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index 0cb3136..fb0c3fa 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -24,6 +24,7 @@ from fastgen.methods.consistency_model.mean_flow import FlowMapLossMixin from fastgen.methods.distribution_matching.dmd2 import DMD2Model +from fastgen.networks.noise_schedule import SHIFTED_TIME_DIST_TYPES from fastgen.utils.distributed import world_size import fastgen.utils.logging_utils as logger @@ -104,7 +105,7 @@ def _rollout_t_list(self, num_steps: int) -> torch.Tensor: *validation* schedule used at ``student_sample_steps``. """ ns = self.net.noise_scheduler - shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type == "shifted" else 1.0 + shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type in SHIFTED_TIME_DIST_TYPES else 1.0 s = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64, device=self.device) s = shift * s / (1 + (shift - 1) * s) return s.clamp(max=float(ns.max_t)).to(ns.t_precision) diff --git a/fastgen/networks/noise_schedule.py b/fastgen/networks/noise_schedule.py index a7463e1..a222363 100644 --- a/fastgen/networks/noise_schedule.py +++ b/fastgen/networks/noise_schedule.py @@ -1303,6 +1303,13 @@ def __init__(self, *args, model_id="THUDM/CogVideoX-5b", **kwargs): del scheduler +# Time distributions whose draws are mapped through shift * t / (1 + (shift - 1) * t). +# Consumers that rebuild that grid outside the sampler -- the flow-map loss weight +# normalization and the AnyFlow rollout schedule -- key the shift off this tuple, so +# adding a shifted variant below stays in sync with them. +SHIFTED_TIME_DIST_TYPES = ("shifted", "shifted_logitnormal") + + class RFNoiseSchedule(BaseNoiseSchedule): """Rectified Flow noise schedule: x_t = (1-t)*x_0 + t*noise. @@ -1317,7 +1324,7 @@ def __init__( **kwargs, ): super().__init__(min_t, max_t, num_steps, **kwargs) - self._supported_time_dist_types = self._supported_time_dist_types + ("shifted", "shifted_logitnormal") + self._supported_time_dist_types = self._supported_time_dist_types + SHIFTED_TIME_DIST_TYPES assert 0 <= min_t < max_t <= 1.0, "RF min_t and max_t must be between 0 and 1" self._sigmas = torch.linspace(min_t, max_t, num_steps, dtype=self.t_precision) diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index 6e537ef..aa577c4 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -254,6 +254,47 @@ def test_pretrain_prediction_side_guidance_fusion(): assert grad_seen +def test_pretrain_guidance_fusion_requires_neg_condition(): + """Fusion queries the unconditional branch, so a missing neg_condition must + stop with a configuration error rather than reach the network as None.""" + model = _build_pretrain_model() + model.config.guidance_fuse_scale = 3.0 + data = _make_data(model) + data["neg_condition"] = None + + with pytest.raises(AssertionError, match="requires neg_condition"): + model.single_train_step(data, 0) + + +def test_shifted_variants_share_the_shift_map(): + """Both shifted time distributions apply the same shift map when sampling t, + so the grids rebuilt outside the sampler must pick up `shift` for either one. + + Covers the loss-weight normalization (pretrain) and the rollout schedule + (on-policy); a variant missing from the lookup degrades silently to shift=1. + """ + model = _build_pretrain_model() + shifted_scale = model._timestep_weight_scale + + model.config.sample_t_cfg.time_dist_type = "shifted_logitnormal" + model._init_flow_map_loss() + assert model._timestep_weight_scale == pytest.approx(shifted_scale) + + # and the shift is what makes it differ from an unshifted grid + model.config.sample_t_cfg.time_dist_type = "uniform" + model._init_flow_map_loss() + assert model._timestep_weight_scale != pytest.approx(shifted_scale) + + onpolicy = _build_onpolicy_model() + onpolicy.sample_t_cfg.shift = 5.0 + onpolicy.sample_t_cfg.time_dist_type = "shifted" + shifted_t_list = onpolicy._rollout_t_list(4) + onpolicy.sample_t_cfg.time_dist_type = "shifted_logitnormal" + assert torch.allclose(onpolicy._rollout_t_list(4), shifted_t_list) + onpolicy.sample_t_cfg.time_dist_type = "uniform" + assert not torch.allclose(onpolicy._rollout_t_list(4), shifted_t_list) + + @pytest.mark.parametrize("weight_type", ["beta08", "gaussian", "uniform"]) def test_timestep_weight_function(weight_type): """The fixed per-timestep weight is a direct function of t: non-negative, @@ -480,12 +521,12 @@ def test_validation_t_list_matches_shift(): t_list = list(cfg.model.sample_t_cfg.t_list) assert len(t_list) == n + 1, t_list - assert all(abs(a - b) < 1e-9 for a, b in zip(t_list, expected)), (t_list, expected) + assert all(abs(a - b) < 1e-9 for a, b in zip(t_list, expected, strict=True)), (t_list, expected) assert t_list[0] == max_t and t_list[-1] == 0.0 # and it must not be the unshifted fallback DMD2 would use for None unshifted = [max_t * (1.0 - i / n) for i in range(n + 1)] - assert not all(abs(a - b) < 1e-6 for a, b in zip(t_list, unshifted)) + assert not all(abs(a - b) < 1e-6 for a, b in zip(t_list, unshifted, strict=True)) def test_fuse_r_embedding_gated_trains_the_shared_time_proj(): From 63ec0821b55daef5370f9bfd72dd35e11cc5c5a4 Mon Sep 17 00:00:00 2001 From: Julius Berner Date: Thu, 20 Aug 2026 22:46:55 +0000 Subject: [PATCH 19/19] anyflow: split co-train sampling cfg, factor time_shift and train_mode --- .../experiments/WanT2V/config_anyflow.py | 20 ++- .../WanT2V/config_anyflow_onpolicy.py | 46 ++++-- .../configs/experiments/WanT2V/config_mf.py | 6 +- fastgen/configs/methods/config_anyflow.py | 11 +- fastgen/configs/methods/config_mean_flow.py | 2 +- fastgen/configs/methods/config_scm.py | 2 +- .../methods/consistency_model/mean_flow.py | 121 +++++++++------- fastgen/methods/consistency_model/sCM.py | 5 +- .../methods/distribution_matching/README.md | 3 +- .../methods/distribution_matching/anyflow.py | 57 ++++---- fastgen/networks/Flux/network.py | 1 - fastgen/networks/QwenImage/network.py | 1 - fastgen/networks/noise_schedule.py | 33 +++-- fastgen/utils/basic_utils.py | 57 +++++--- tests/test_anyflowmodel.py | 133 ++++++++++++++---- 15 files changed, 318 insertions(+), 180 deletions(-) diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py index 5a630c2..852142b 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -6,19 +6,26 @@ AnyFlow's pretrain objective is MeanFlow's with a fixed ``beta08`` per-timestep weighting, a finite-difference JVP, shifted timestep sampling, and a ``consistency_ratio`` fraction of the batch pinned to ``r = 0`` — so this config -runs``MeanFlowModel`` directly. The values below mirror the reference recipe +runs ``MeanFlowModel`` directly. The values below mirror the reference recipe ``train_wan1b_student_shift5_81f_480p_lr5e-5_6k_b32.yml``. -Known deviations from the reference: full-rank fine-tuning instead of the paper's -rank-256 LoRA — the same deviation applies to the on-policy stage; + +Known deviations from the reference, both of which also apply to the on-policy stage: +full-rank fine-tuning instead of the paper's rank-256 LoRA; and under +``guidance_fuse_scale`` the 1/g rescaling of dF/dt is applied only to the samples that +kept their condition, where the reference's ``compute_central_difference`` rescales the +whole batch. A dropped sample's fused prediction is plain ``u_uncond``, so an ungated +1/g would regress it onto a different fixed point -- ours is gated on ``keep``. The on-policy stage (paper Stage 2) lives in ``config_anyflow_onpolicy.py``. """ import copy + import fastgen.configs.methods.config_mean_flow as config_mean_flow from fastgen.configs.data import VideoLoaderConfig from fastgen.configs.net import Wan_1_3B_Config +from fastgen.methods import AnyFlowModel def create_config(): @@ -99,8 +106,11 @@ def create_config(): # ------ inference / validation ------ config.model.student_sample_type = "ode" config.model.student_sample_steps = 4 - # 4-step shifted schedule shift*s/(1+(shift-1)*s) on s=linspace(1,0,5) - config.model.sample_t_cfg.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625, 0.0] + # Shifted schedule under the same map the (t, r) sampling applies. The pretrain + # stage has a single scheduler in the reference, so it shares that shift. + config.model.sample_t_cfg.t_list = AnyFlowModel.rollout_t_list( + config.model.student_sample_steps, config.model.sample_t_cfg.shift, config.model.net.max_t + ).tolist() # ------ data / trainer ------ config.dataloader_train = copy.deepcopy(VideoLoaderConfig) diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py index 0c2cef4..225ae82 100644 --- a/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -17,15 +17,23 @@ trainer.checkpointer.pretrained_ckpt_path=/checkpoints/0006000.pth -Known deviation from the reference: full-rank fine-tuning instead of the paper's -rank-256 LoRA. +Known deviations from the reference: full-rank fine-tuning instead of the paper's +rank-256 LoRA; the noising times for the DMD gradient and the fake score are drawn +on [0.001, 0.999] rather than the reference's [0, 1] (its ``dmd_cfg`` sets no +``min_timestep`` / ``max_timestep``, so it clamps to the full range) -- FastGen's +convention for the other rectified-flow Wan configs; and, in the co-trained loss, the +1/g rescaling of dF/dt under ``guidance_fuse_scale`` is gated on the samples that kept +their condition, where the reference's ``compute_central_difference`` rescales the whole +batch (see ``config_anyflow.py`` and ``FlowMapLossMixin._compute_mf_loss``). """ import copy + import fastgen.configs.methods.config_anyflow as config_anyflow_default from fastgen.configs.data import VideoLoaderConfig from fastgen.configs.net import Wan_1_3B_Config +from fastgen.methods import AnyFlowModel def create_config(): @@ -77,11 +85,13 @@ def create_config(): config.model.guidance_scale = 4.0 # DMD gradient noising time: reference `generator_loss` draws torch.rand - # then applies the shift -> shifted-uniform. + # then applies the shift -> shifted-uniform. The bounds keep the score models off + # the degenerate endpoints (see the deviation note above); `fake_score_sample_t_cfg` + # inherits them through the deepcopy below. config.model.sample_t_cfg.time_dist_type = "shifted" config.model.sample_t_cfg.shift = 5.0 - config.model.sample_t_cfg.min_t = 0.0 - config.model.sample_t_cfg.max_t = 1.0 + config.model.sample_t_cfg.min_t = 0.001 + config.model.sample_t_cfg.max_t = 0.999 # Fake-score noising time: reference `discriminator_loss` draws # logit_normal(0, 1) then applies the same shift. This is a DIFFERENT @@ -92,13 +102,12 @@ def create_config(): config.model.fake_score_sample_t_cfg.train_p_std = 1.0 # ------ student rollout (reference rollout_cfg) ------ + # The rollout grid's shift comes from `cotrain_sample_t_cfg` below: the reference + # builds its rollout pipeline from the same `scheduler` it draws the co-trained + # (t, r) from, separately from the DMD noising time above. config.model.student_sample_type = "ode" config.model.student_sample_steps_list = [2, 4, 8, 16, 50] config.model.student_sample_steps = 4 - # Validation schedule for `student_sample_steps = 4`: `shift * s / (1 + (shift - # - 1) * s)` on `s = linspace(1, 0, 5)` -- the same grid the per-NFE rollout - # builds. Recompute if `shift`, `max_t` or `student_sample_steps` changes. - config.model.sample_t_cfg.t_list = [1.0, 0.9375, 0.8333333333333334, 0.625, 0.0] # ------ co-trained Stage-1 flow-map loss (reference cotrain_forward_kl) ------ # FastGen's VSD loss carries a 0.5 factor the reference's DMD loss does @@ -115,9 +124,22 @@ def create_config(): config.model.guidance_fuse_scale = 3.0 config.model.cond_dropout_prob = 0.1 config.model.precision_amp_jvp = "float32" - config.model.sample_t_cfg.flow_matching_ratio = 0.5 - config.model.sample_t_cfg.consistency_ratio = 0.25 - config.model.sample_t_cfg.deterministic_buckets = True + # (t, r) sampling for the co-trained loss, drawn from the reference's `scheduler`. + # Its shift also drives the student's rollout grid, so it is independent of the DMD + # noising-time shift above; the reference recipe sets both to 5.0. + config.model.cotrain_sample_t_cfg.time_dist_type = "shifted" + config.model.cotrain_sample_t_cfg.shift = 5.0 + config.model.cotrain_sample_t_cfg.min_t = 0.0 + config.model.cotrain_sample_t_cfg.max_t = 1.0 + config.model.cotrain_sample_t_cfg.flow_matching_ratio = 0.5 + config.model.cotrain_sample_t_cfg.consistency_ratio = 0.25 + config.model.cotrain_sample_t_cfg.deterministic_buckets = True + + # Validation schedule at `student_sample_steps` -- the same grid the per-NFE rollout + # builds, so it follows the co-trained sampling shift. + config.model.sample_t_cfg.t_list = AnyFlowModel.rollout_t_list( + config.model.student_sample_steps, config.model.cotrain_sample_t_cfg.shift, config.model.net.max_t + ).tolist() # ------ optimization (reference: AdamW lr=2e-6, betas=(0.0, 0.999), wd=0, # grad clip 1.0, EMA 0.99) ------ diff --git a/fastgen/configs/experiments/WanT2V/config_mf.py b/fastgen/configs/experiments/WanT2V/config_mf.py index 9d4f1c2..710716e 100644 --- a/fastgen/configs/experiments/WanT2V/config_mf.py +++ b/fastgen/configs/experiments/WanT2V/config_mf.py @@ -74,12 +74,8 @@ def create_config(): config.model.sample_t_cfg.min_t = 0.001 config.model.sample_t_cfg.max_t = 0.999 - config.dataloader_train = VideoLatentLoaderConfig + config.dataloader_train = copy.deepcopy(VideoLatentLoaderConfig) config.dataloader_train.batch_size = 1 - # 480p (832x480) resolution - config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) - config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 - config.log_config.group = "wan_mf" return config diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py index c7084a2..7e4113d 100644 --- a/fastgen/configs/methods/config_anyflow.py +++ b/fastgen/configs/methods/config_anyflow.py @@ -42,10 +42,11 @@ class ModelConfig(DMD2ModelConfig): ``cotrain_forward_kl``). """ - # MeanFlow-style (t, r) sampling for the co-trained flow-map loss; the - # extra fields are ignored by the DMD2 noising-time sampling. - sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) - sample_r_cfg: MeanFlowSampleRConfig = attrs.field(factory=MeanFlowSampleRConfig) + # MeanFlow-style (t, r) sampling for the co-trained flow-map loss. Separate from + # the inherited `sample_t_cfg`, which DMD2 uses for the noising time: the reference + # draws these from its `scheduler` and the noising time from its `dmd_scheduler`. + cotrain_sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) + cotrain_sample_r_cfg: MeanFlowSampleRConfig = attrs.field(factory=MeanFlowSampleRConfig) loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig) # Weight of the co-trained Stage-1 flow-map loss in the student update. @@ -70,7 +71,7 @@ class ModelConfig(DMD2ModelConfig): cond_dropout_prob: Optional[float] = None cond_keys_no_dropout: List[str] = attrs.field(factory=list) - # Precision for autocast in the co-trained loss JVP (None = training precision). + # Precision for autocast in the co-trained loss JVP (None disables autocast there). precision_amp_jvp: str | None = None # MeanFlow's target-side guidance knobs. AnyFlow typically guides on the diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index b491db2..743537d 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -129,7 +129,7 @@ class ModelConfig(BaseModelConfig): # guidance t end guidance_t_end: float = 1.0 - # precision for autocast in JVP (none defaults to training precision) + # precision for autocast in JVP (none disables autocast in the JVP region) precision_amp_jvp: str | None = None diff --git a/fastgen/configs/methods/config_scm.py b/fastgen/configs/methods/config_scm.py index 6303bd4..b430fb0 100644 --- a/fastgen/configs/methods/config_scm.py +++ b/fastgen/configs/methods/config_scm.py @@ -77,7 +77,7 @@ class ModelConfig(BaseModelConfig): # optimizer net_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(RAdamOptimizerConfig)) - # precision for autocast in JVP (none defaults to training precision) + # precision for autocast in JVP (none disables autocast in the JVP region) precision_amp_jvp: str | None = None diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 4eaf824..5b9710d 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -9,7 +9,7 @@ import numpy as np import torch from fastgen.methods import CMModel -from fastgen.networks.noise_schedule import SHIFTED_TIME_DIST_TYPES +from fastgen.networks.noise_schedule import time_shift from fastgen.utils import basic_utils, expand_like from fastgen.utils.basic_utils import convert_cfg_to_dict from fastgen.utils.distributed import get_rank, world_size @@ -65,30 +65,46 @@ class FlowMapLossMixin: from its ``__init__``. """ - def _init_flow_map_loss(self) -> None: - """Bind the flow-map loss / sampling configs and the JVP precision.""" - self.sample_t_cfg = self.config.sample_t_cfg - self.sample_r_cfg = self.config.sample_r_cfg + def _init_flow_map_loss(self, sample_t_cfg: Any, sample_r_cfg: Any) -> None: + """Bind the flow-map loss / sampling configs and the JVP precision. + + The (t, r) configs are passed in rather than read off ``self.config``: the host + may draw the flow-map times from a different density than its own + ``sample_t_cfg`` (AnyFlow does -- see ``AnyFlowModel``), and ``CMModel`` already + owns a ``self.sample_t_cfg`` attribute that must not be rebound here. + + Args: + sample_t_cfg: Config for sampling the flow-map ``t``. + sample_r_cfg: Config for sampling the flow-map ``r``. + """ + self.flow_map_sample_t_cfg = sample_t_cfg + self.flow_map_sample_r_cfg = sample_r_cfg self.loss_config = self.config.loss_config - # Drop the JVP autocast when it matches the outer one, since a nested region in the - # same dtype is a no-op. - if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp: + # The shift `_sample_t_r_buckets` actually applies: `shift` is inert unless the + # density is one of `RFNoiseSchedule`'s shifted ones, so resolve it once here and + # reuse it wherever the flow-map timestep grid is rebuilt. + shifted = sample_t_cfg.time_dist_type in ("shifted", "shifted_logitnormal") + self.flow_map_shift = sample_t_cfg.shift if shifted else 1.0 + + # None runs the JVP with autocast disabled: the region always sets the autocast + # state, so the enclosing training autocast does not carry over into it. + if self.config.precision_amp_jvp is None: self.precision_amp_jvp = None else: self.precision_amp_jvp = basic_utils.PRECISION_MAP[self.config.precision_amp_jvp] logger.critical(f"Using precision {self.precision_amp_jvp} for JVP") # The fixed per-timestep loss weight is normalized to mean one over the - # network's discrete training timesteps (t = 0 excluded), mapped through - # the same shift the samplers use. The constant depends only on the - # config, so derive it once here. + # network's discrete training timesteps (t = 0 excluded), mapped through the + # same shift `_sample_t_r_buckets` applies -- the weight has to be normalized + # over the grid the loss's own t values are drawn from, which is why this reads + # the passed-in config. The constant depends only on it, so derive it once here. self._timestep_weight_scale: Optional[float] = None if self.loss_config.weight_type is not None: num_steps = self.net.noise_scheduler.num_steps - shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type in SHIFTED_TIME_DIST_TYPES else 1.0 grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64)[:-1] - grid = shift * grid / (1 + (shift - 1) * grid) + grid = time_shift(grid, self.flow_map_shift) self._timestep_weight_scale = float(num_steps / self._timestep_weight_raw(grid).sum()) def _drop_condition( @@ -109,7 +125,7 @@ def _drop_condition( if self.config.cond_dropout_prob is None or neg_condition is None: return condition, torch.ones(batch_size, dtype=torch.bool, device=device) - if self.sample_t_cfg.deterministic_buckets: + if self.flow_map_sample_t_cfg.deterministic_buckets: keep = torch.rand(batch_size, device=device) >= self.config.cond_dropout_prob else: num_to_drop = (torch.rand(batch_size, device=device) < self.config.cond_dropout_prob).sum() @@ -183,33 +199,32 @@ def _get_velocity( self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None ): # Turn off dropout - self.net.eval() - neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow") - guidance_scale = self.config.guidance_scale or 1.0 - guidance_scale = torch.where( - ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - guidance_scale, - 1.0, - ) - guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) - - if self.config.guidance_mixture_ratio is None: - guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt) - else: - guidance_mixture_ratio = torch.where( + with basic_utils.train_mode(self.net, mode=False): + neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow") + guidance_scale = self.config.guidance_scale or 1.0 + guidance_scale = torch.where( ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - self.config.guidance_mixture_ratio, - 0.0, - ) - guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype) - cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow") - guided_dxt_dt = ( - guidance_scale * dxt_dt - + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt - + guidance_mixture_ratio * cond_dxt_dt + guidance_scale, + 1.0, ) + guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) + + if self.config.guidance_mixture_ratio is None: + guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt) + else: + guidance_mixture_ratio = torch.where( + ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), + self.config.guidance_mixture_ratio, + 0.0, + ) + guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype) + cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow") + guided_dxt_dt = ( + guidance_scale * dxt_dt + + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt + + guidance_mixture_ratio * cond_dxt_dt + ) - self.net.train() condition, keep = self._drop_condition(condition, neg_condition, x_t.shape[0], x_t.device) # Same subset: a kept sample is conditional + guided, a dropped one # unconditional + unguided. @@ -311,8 +326,8 @@ def net_wrapper(x_t, t, r): def _warn_on_degenerate_buckets(self, n_flow_matching: int, n_consistency: int, global_bsz: int) -> None: """Warn once if a requested bucket rounds away under the deterministic partition.""" - requested_but_empty = (self.sample_t_cfg.flow_matching_ratio > 0 and n_flow_matching == 0) or ( - self.sample_t_cfg.consistency_ratio > 0 and n_consistency == 0 + requested_but_empty = (self.flow_map_sample_t_cfg.flow_matching_ratio > 0 and n_flow_matching == 0) or ( + self.flow_map_sample_t_cfg.consistency_ratio > 0 and n_consistency == 0 ) if not requested_but_empty or getattr(self, "_bucket_warned", False): return @@ -320,8 +335,8 @@ def _warn_on_degenerate_buckets(self, n_flow_matching: int, n_consistency: int, logger.warning( f"The deterministic (t, r) bucket partition is degenerate: with " f"world_size * batch_size = {global_bsz}, flow_matching_ratio=" - f"{self.sample_t_cfg.flow_matching_ratio} and consistency_ratio=" - f"{self.sample_t_cfg.consistency_ratio} yield bucket sizes " + f"{self.flow_map_sample_t_cfg.flow_matching_ratio} and consistency_ratio=" + f"{self.flow_map_sample_t_cfg.consistency_ratio} yield bucket sizes " f"({n_flow_matching}, {n_consistency}). The partition spans ranks but not " f"gradient-accumulation rounds, so empty buckets stay empty every iteration." ) @@ -336,22 +351,24 @@ def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tens ``r = t``, a ``consistency_ratio`` fraction gets ``r = 0`` (consistency to clean data), and the rest keep the sampled random pair. """ - t_sample_kwargs = convert_cfg_to_dict(self.sample_t_cfg) + t_sample_kwargs = convert_cfg_to_dict(self.flow_map_sample_t_cfg) t = self.net.noise_scheduler.sample_t(batch_size, **t_sample_kwargs, device=self.device) - r_sample_kwargs = convert_cfg_to_dict(self.sample_r_cfg) if self.sample_r_cfg.enabled else t_sample_kwargs + r_sample_kwargs = ( + convert_cfg_to_dict(self.flow_map_sample_r_cfg) if self.flow_map_sample_r_cfg.enabled else t_sample_kwargs + ) r = self.net.noise_scheduler.sample_t(batch_size, **r_sample_kwargs, device=self.device) t, r = torch.maximum(t, r), torch.minimum(t, r) assert torch.all(t >= r), "r cannot be larger than t" - flow_matching_ratio = self.sample_t_cfg.flow_matching_ratio - consistency_ratio = self.sample_t_cfg.consistency_ratio + 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}" # Both policies produce two bucket sizes and index the batch the same # way; they differ only in how the sizes are obtained. - if self.sample_t_cfg.deterministic_buckets: + if self.flow_map_sample_t_cfg.deterministic_buckets: global_bsz = world_size() * batch_size n_flow_matching = round(flow_matching_ratio * global_bsz) n_consistency = round(consistency_ratio * global_bsz) @@ -606,13 +623,9 @@ def _compute_mf_loss( # difference over g on conditional samples, with the unconditional # derivative dropped. u_theta_jvp = torch.where(expand_like(keep, u_theta_jvp), u_theta_jvp / guidance_fuse_scale, u_theta_jvp) - with torch.no_grad(): - # Turn off dropout for the unconditional pass, as the target-side path does. - was_training = self.net.training - self.net.eval() + # Turn off dropout for the unconditional pass, as the target-side path does. + with basic_utils.train_mode(self.net, mode=False), torch.no_grad(): u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow") - if was_training: - self.net.train() u_theta = (u_theta + (guidance_fuse_scale - 1.0) * u_uncond) / guidance_fuse_scale mf_loss, tangent, loss_weight, warmup_weight = self._mf_pred_to_loss( @@ -644,7 +657,7 @@ def __init__(self, config: ModelConfig): """ super().__init__(config) self.config = config - self._init_flow_map_loss() + self._init_flow_map_loss(config.sample_t_cfg, config.sample_r_cfg) def single_train_step( self, data: Dict[str, Any], iteration: int diff --git a/fastgen/methods/consistency_model/sCM.py b/fastgen/methods/consistency_model/sCM.py index 3a8259b..85d8c82 100644 --- a/fastgen/methods/consistency_model/sCM.py +++ b/fastgen/methods/consistency_model/sCM.py @@ -93,8 +93,9 @@ def __init__(self, config: ModelConfig): self.loss_config = self.config.loss_config self.sigma_data = self.sample_t_cfg.sigma_data - # Precision for JVP - if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp: + # Precision for JVP. None disables autocast in the JVP region: it always sets the + # autocast state, so the enclosing training autocast does not carry over into it. + if self.config.precision_amp_jvp is None: self.precision_amp_jvp = None else: self.precision_amp_jvp = PRECISION_MAP[self.config.precision_amp_jvp] diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index 16561b8..d465fd8 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -87,13 +87,14 @@ DMD2 extended for causal video generation with autoregressive chunk-by-chunk pro ## AnyFlow -**File:** [`anyflow.py`](anyflow.py) | **Reference:** [Gu et al., 2026](https://arxiv.org/abs/2605.13724) +**File:** [`anyflow.py`](anyflow.py) | **Reference:** [Gu et al. (2026)](https://arxiv.org/abs/2605.13724) DMD2 with a flow-map student `u(x_t, t, r)` (mean velocity from `t` back to `r`), supporting arbitrary inference NFE. The student generates by rolling out the flow map from pure noise with the NFE sampled per iteration and gradients through all segments, and co-trains the flow-map loss at every update. Requires a Stage-1 flow-map pretrain, which runs directly on [`MeanFlowModel`](../consistency_model/mean_flow.py) using AnyFlow's hyperparameters. **Key Parameters (Stage 2):** - `student_sample_steps_list`: Rollout NFEs sampled per iteration - `cotrain_pretrain_weight`: Weight of the co-trained Stage-1 flow-map loss +- `cotrain_sample_t_cfg`, `cotrain_sample_r_cfg`: Separate `(t, r)` distributions for the co-trained flow-map loss; its shift also defines the rollout schedule - Requires a Stage-1 pretrained checkpoint via `trainer.checkpointer.pretrained_ckpt_path` - See also the key parameters of DMD2 above (and the key parameters of [MeanFlow](../consistency_model/README.md#meanflow) for Stage 1) diff --git a/fastgen/methods/distribution_matching/anyflow.py b/fastgen/methods/distribution_matching/anyflow.py index fb0c3fa..32bbcd4 100644 --- a/fastgen/methods/distribution_matching/anyflow.py +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -24,7 +24,7 @@ from fastgen.methods.consistency_model.mean_flow import FlowMapLossMixin from fastgen.methods.distribution_matching.dmd2 import DMD2Model -from fastgen.networks.noise_schedule import SHIFTED_TIME_DIST_TYPES +from fastgen.networks.noise_schedule import time_shift from fastgen.utils.distributed import world_size import fastgen.utils.logging_utils as logger @@ -45,7 +45,9 @@ class AnyFlowModel(FlowMapLossMixin, DMD2Model): def __init__(self, config: ModelConfig): super().__init__(config) self.config = config - self._init_flow_map_loss() + # The co-trained flow-map loss draws its (t, r) from its own config, not from + # `sample_t_cfg` -- DMD2 keeps that one for the noising time. + self._init_flow_map_loss(config.cotrain_sample_t_cfg, config.cotrain_sample_r_cfg) logger.info( f"AnyFlow on-policy: student_sample_steps_list={self.config.student_sample_steps_list}, " @@ -84,31 +86,16 @@ def _broadcast_choice(self, high: int) -> int: torch.distributed.broadcast(idx, src=0) return int(idx.item()) - def _sample_rollout_steps(self) -> int: - """Sample this iteration's rollout NFE, as the reference does - (``random.choice(num_inference_steps_list)``, rank-0 broadcast in both - the generator and fake-score updates). Falls back to the fixed - ``student_sample_steps`` when no list is set. - """ - steps_list = self.config.student_sample_steps_list - if not steps_list: - return int(self.config.student_sample_steps) - return int(steps_list[self._broadcast_choice(len(steps_list))]) - - def _rollout_t_list(self, num_steps: int) -> torch.Tensor: - """Shifted timestep schedule for ``num_steps`` rollout steps. - - Equivalent to the reference scheduler's ``set_timesteps``: a uniform grid - mapped through ``shift * s / (1 + (shift - 1) * s)``, first entry clamped - to ``max_t``. Computed here rather than read from ``sample_t_cfg.t_list`` - because the rollout NFE varies per iteration; that list stays the fixed - *validation* schedule used at ``student_sample_steps``. + @staticmethod + def rollout_t_list(num_steps: int, shift: float = 1.0, max_t: float = 1.0) -> torch.Tensor: + """Shifted timestep schedule ``[max_t, ..., 0]`` with ``num_steps + 1`` entries. + + Equivalent to the reference scheduler's ``set_timesteps``, with ``shift`` the + model's ``flow_map_shift``. Static so the configs can derive the matching + validation ``t_list`` from it; float64 on the CPU for the caller to cast. """ - ns = self.net.noise_scheduler - shift = self.sample_t_cfg.shift if self.sample_t_cfg.time_dist_type in SHIFTED_TIME_DIST_TYPES else 1.0 - s = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64, device=self.device) - s = shift * s / (1 + (shift - 1) * s) - return s.clamp(max=float(ns.max_t)).to(ns.t_precision) + grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64) + return time_shift(grid, shift).clamp(max=max_t) def gen_data_from_net( self, @@ -127,11 +114,21 @@ def gen_data_from_net( """ del t_student # the rollout schedule is built below - num_steps = self._sample_rollout_steps() - if num_steps < 1: - raise ValueError(f"rollout steps must be >= 1, got {num_steps}") + # This iteration's rollout NFE, as the reference draws it + # (``random.choice(num_inference_steps_list)``, rank-0 broadcast in both the + # generator and the fake-score update); no list means the fixed + # ``student_sample_steps``. + steps_list = self.config.student_sample_steps_list + if steps_list: + num_steps = int(steps_list[self._broadcast_choice(len(steps_list))]) + else: + num_steps = int(self.config.student_sample_steps) + assert num_steps >= 1, f"rollout steps must be >= 1, got {num_steps}" grad_step = self._broadcast_choice(num_steps) - t_list = self._rollout_t_list(num_steps) + ns = self.net.noise_scheduler + t_list = self.rollout_t_list(num_steps, self.flow_map_shift, float(ns.max_t)).to( + device=self.device, dtype=ns.t_precision + ) # The leading jump exists only for grad_step > 0 and the trailing one # only for grad_step + 1 < num_steps. diff --git a/fastgen/networks/Flux/network.py b/fastgen/networks/Flux/network.py index efc04b2..1582c34 100644 --- a/fastgen/networks/Flux/network.py +++ b/fastgen/networks/Flux/network.py @@ -700,7 +700,6 @@ def _calculate_shift( mu = image_seq_len * m + b return mu - @torch.no_grad() def sample( self, noise: torch.Tensor, diff --git a/fastgen/networks/QwenImage/network.py b/fastgen/networks/QwenImage/network.py index b699509..0998424 100644 --- a/fastgen/networks/QwenImage/network.py +++ b/fastgen/networks/QwenImage/network.py @@ -647,7 +647,6 @@ def forward( return out, logvar return out - @torch.no_grad() def sample( self, noise: torch.Tensor, diff --git a/fastgen/networks/noise_schedule.py b/fastgen/networks/noise_schedule.py index a222363..d629323 100644 --- a/fastgen/networks/noise_schedule.py +++ b/fastgen/networks/noise_schedule.py @@ -20,6 +20,22 @@ NET_PRED_TYPES = {"x0", "eps", "v", "flow"} +def time_shift(t: torch.Tensor, shift: float = 1.0) -> torch.Tensor: + """Map timesteps through ``t * shift / (t * (shift - 1) + 1)``. + + Args: + t: Timesteps in [0, 1]. + shift: Shift of the map. 1.0 is the identity. + + Returns: + torch.Tensor: The shifted timesteps. + """ + assert shift >= 1, f"shift must be >= 1, got {shift}" + if shift == 1.0: + return t + return t * shift / (t * (shift - 1) + 1) + + class BaseNoiseSchedule(torch.nn.Module): """Abstract base noise schedule class. @@ -1303,13 +1319,6 @@ def __init__(self, *args, model_id="THUDM/CogVideoX-5b", **kwargs): del scheduler -# Time distributions whose draws are mapped through shift * t / (1 + (shift - 1) * t). -# Consumers that rebuild that grid outside the sampler -- the flow-map loss weight -# normalization and the AnyFlow rollout schedule -- key the shift off this tuple, so -# adding a shifted variant below stays in sync with them. -SHIFTED_TIME_DIST_TYPES = ("shifted", "shifted_logitnormal") - - class RFNoiseSchedule(BaseNoiseSchedule): """Rectified Flow noise schedule: x_t = (1-t)*x_0 + t*noise. @@ -1324,7 +1333,7 @@ def __init__( **kwargs, ): super().__init__(min_t, max_t, num_steps, **kwargs) - self._supported_time_dist_types = self._supported_time_dist_types + SHIFTED_TIME_DIST_TYPES + self._supported_time_dist_types = self._supported_time_dist_types + ("shifted", "shifted_logitnormal") assert 0 <= min_t < max_t <= 1.0, "RF min_t and max_t must be between 0 and 1" self._sigmas = torch.linspace(min_t, max_t, num_steps, dtype=self.t_precision) @@ -1417,20 +1426,16 @@ def sample_t( elif time_dist_type == "uniform": t = torch.rand(n, device=target_device, dtype=self.t_precision) * (max_t - min_t) + min_t elif time_dist_type == "shifted": - shift = kwargs.get("shift", 5.0) - assert shift >= 1, f"shift must be >= 1, got {shift}" t = torch.rand(n, device=target_device, dtype=self.t_precision) * (max_t - min_t) + min_t - t = t * shift / (t * (shift - 1) + 1) + t = time_shift(t, kwargs.get("shift", 5.0)) elif time_dist_type == "shifted_logitnormal": # Logit-normal base density under the same shift map as "shifted" - shift = kwargs.get("shift", 5.0) - assert shift >= 1, f"shift must be >= 1, got {shift}" t = ( torch.sigmoid(torch.randn(n, device=target_device, dtype=self.t_precision) * train_p_std + train_p_mean) * (max_t - min_t) + min_t ) - t = t * shift / (t * (shift - 1) + 1) + t = time_shift(t, kwargs.get("shift", 5.0)) else: raise ValueError( f"Unsupported time distribution type: {time_dist_type} in RFNoiseSchedule." diff --git a/fastgen/utils/basic_utils.py b/fastgen/utils/basic_utils.py index bb6175b..afbace2 100644 --- a/fastgen/utils/basic_utils.py +++ b/fastgen/utils/basic_utils.py @@ -87,44 +87,65 @@ def to_str(obj: Any) -> str | Dict[Any, str]: @contextmanager -def inference_mode(*modules: torch.nn.Module, precision_amp: torch.dtype | None = None, device_type: str = "cuda"): +def train_mode(*modules: torch.nn.Module, mode: bool = True): """ - Wraps torch.inference_mode() and temporarily sets the provided modules - to .eval() mode. If precision_amp is not None, it also wraps the context in torch.autocast(). + Temporarily sets the provided modules to train (mode=True) or eval (mode=False) mode. + + Use mode=False to turn off training-only behavior such as dropout for a single forward + pass, e.g. when querying the trained network as its own teacher. Args: - *modules: Modules to set temporarily to eval mode. - precision_amp: If not None, wraps the context in torch.autocast(). - device_type: Device type to use for autocast. + *modules: Modules to set temporarily to the given mode. + mode: Passed to torch.nn.Module.train(); False selects eval mode. Returns: Generator that yields the context manager. Upon exit, it restores the original .training state of each module. """ - # 1. Capture the original training state of each module - # (True if in train mode, False if in eval mode) + # Capture the original training state of each module + # (True if in train mode, False if in eval mode) modules = [mod for mod in modules if isinstance(mod, torch.nn.Module)] previous_states = [mod.training for mod in modules] try: - # 2. Set all specific modules to eval mode - # This is crucial for layers like Dropout and BatchNorm + # The mode gates training-only behavior of layers like Dropout and BatchNorm for mod in modules: - mod.eval() - - # 3. Enter strict inference mode (disables gradients, etc.) and autocast if needed - with torch.inference_mode(), torch.autocast( - dtype=precision_amp, device_type=device_type, enabled=precision_amp is not None - ): - yield + mod.train(mode) + yield finally: - # 4. Restore the original state of each module + # Restore the original state of each module for mod, was_training in zip(modules, previous_states): mod.train(was_training) +@contextmanager +def inference_mode(*modules: torch.nn.Module, precision_amp: torch.dtype | None = None, device_type: str = "cuda"): + """ + Wraps torch.inference_mode() and temporarily sets the provided modules to .eval() mode. + + The context always sets the autocast state: with precision_amp it autocasts to that + dtype, and with precision_amp=None it explicitly *disables* autocast, so an enclosing + autocast region does not carry over into this context. + + Args: + *modules: Modules to set temporarily to eval mode. + precision_amp: Dtype to autocast to. If None, autocast is disabled in the context. + device_type: Device type to use for autocast. + + Returns: + Generator that yields the context manager. + + Upon exit, it restores the original .training state of each module. + """ + # Set eval mode, enter strict inference mode (disables gradients, etc.) and set the autocast state + with train_mode(*modules, mode=False), torch.inference_mode(), torch.autocast( + dtype=precision_amp, device_type=device_type, enabled=precision_amp is not None + ): + yield + + def set_random_seed( seed: int, iteration: int = 0, by_rank: bool = False, devices: List[torch.device | str | int] | None = None ) -> int: diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py index aa577c4..b3f0241 100644 --- a/tests/test_anyflowmodel.py +++ b/tests/test_anyflowmodel.py @@ -21,6 +21,7 @@ from fastgen.configs.methods.config_anyflow import ModelConfig as AnyFlowModelConfig from fastgen.configs.methods.config_mean_flow import ModelConfig as MeanFlowModelConfig from fastgen.methods import AnyFlowModel, MeanFlowModel +from fastgen.networks.noise_schedule import time_shift from fastgen.utils.test_utils import check_grad_zero @@ -58,9 +59,13 @@ def _build_pretrain_model( return model -def _build_onpolicy_model(): +def _build_onpolicy_model(cotrain_time_dist_type="uniform", cotrain_shift=5.0): """On-policy fixture mirrors test_dmd2model: img_resolution=8 so the - discriminator's 4x4 conv kernels can operate.""" + discriminator's 4x4 conv kernels can operate. + + ``cotrain_time_dist_type`` defaults to an unshifted density, so ``shift`` stays + inert; pass a shifted one to exercise the shifted rollout grid. + """ gc.collect() instance = AnyFlowModelConfig() @@ -80,13 +85,15 @@ def _build_onpolicy_model(): instance.student_sample_steps = 2 instance.input_shape = [3, 8, 8] - # Co-trained flow-map loss settings (MeanFlow machinery on the tiny net). - instance.sample_t_cfg.time_dist_type = "uniform" - instance.sample_t_cfg.min_t = 0.001 - instance.sample_t_cfg.max_t = 0.999 - instance.sample_t_cfg.flow_matching_ratio = 0.5 - instance.sample_t_cfg.consistency_ratio = 0.25 - instance.sample_t_cfg.deterministic_buckets = True + # Co-trained flow-map loss settings (MeanFlow machinery on the tiny net). These + # live on their own config: `sample_t_cfg` stays DMD2's noising time. + instance.cotrain_sample_t_cfg.time_dist_type = cotrain_time_dist_type + instance.cotrain_sample_t_cfg.shift = cotrain_shift + instance.cotrain_sample_t_cfg.min_t = 0.001 + instance.cotrain_sample_t_cfg.max_t = 0.999 + instance.cotrain_sample_t_cfg.flow_matching_ratio = 0.5 + instance.cotrain_sample_t_cfg.consistency_ratio = 0.25 + instance.cotrain_sample_t_cfg.deterministic_buckets = True instance.loss_config.loss_type = "l2" instance.loss_config.weight_type = "uniform" instance.loss_config.norm_method = None @@ -275,24 +282,91 @@ def test_shifted_variants_share_the_shift_map(): """ model = _build_pretrain_model() shifted_scale = model._timestep_weight_scale + cfg_t, cfg_r = model.config.sample_t_cfg, model.config.sample_r_cfg - model.config.sample_t_cfg.time_dist_type = "shifted_logitnormal" - model._init_flow_map_loss() + cfg_t.time_dist_type = "shifted_logitnormal" + model._init_flow_map_loss(cfg_t, cfg_r) assert model._timestep_weight_scale == pytest.approx(shifted_scale) # and the shift is what makes it differ from an unshifted grid - model.config.sample_t_cfg.time_dist_type = "uniform" - model._init_flow_map_loss() + cfg_t.time_dist_type = "uniform" + model._init_flow_map_loss(cfg_t, cfg_r) assert model._timestep_weight_scale != pytest.approx(shifted_scale) - onpolicy = _build_onpolicy_model() - onpolicy.sample_t_cfg.shift = 5.0 - onpolicy.sample_t_cfg.time_dist_type = "shifted" - shifted_t_list = onpolicy._rollout_t_list(4) - onpolicy.sample_t_cfg.time_dist_type = "shifted_logitnormal" - assert torch.allclose(onpolicy._rollout_t_list(4), shifted_t_list) - onpolicy.sample_t_cfg.time_dist_type = "uniform" - assert not torch.allclose(onpolicy._rollout_t_list(4), shifted_t_list) + # Same lookup on the on-policy rollout schedule. The shift is resolved once at + # init, so each variant needs a re-init rather than a bare config mutation. + onpolicy = _build_onpolicy_model(cotrain_time_dist_type="shifted", cotrain_shift=5.0) + cot_t, cot_r = onpolicy.config.cotrain_sample_t_cfg, onpolicy.config.cotrain_sample_r_cfg + + # `rollout_t_list` is static, so read `flow_map_shift` back off the model: that is + # what checks the co-train density actually reaches the rollout grid. + def rollout_grid(): + return onpolicy.rollout_t_list(4, onpolicy.flow_map_shift, float(onpolicy.net.noise_scheduler.max_t)) + + shifted_grid = rollout_grid() + + cot_t.time_dist_type = "shifted_logitnormal" + onpolicy._init_flow_map_loss(cot_t, cot_r) + assert torch.allclose(rollout_grid(), shifted_grid) + + cot_t.time_dist_type = "uniform" + onpolicy._init_flow_map_loss(cot_t, cot_r) + assert not torch.allclose(rollout_grid(), shifted_grid) + + +def test_onpolicy_shifted_rollout_grid_and_step(): + """A shifted co-train density must move the rollout grid, and a student update + must run on it. + + Every other on-policy test uses the unshifted fixture, so without this the shift + never reaches `rollout_t_list` or a real rollout -- the shipped + `config_anyflow_onpolicy.py` runs shift=5. + """ + model = _build_onpolicy_model(cotrain_time_dist_type="shifted", cotrain_shift=5.0) + assert model.flow_map_shift == 5.0 + + max_t = float(model.net.noise_scheduler.max_t) + grid = torch.linspace(1.0, 0.0, 5, dtype=torch.float64) + expected = time_shift(grid, 5.0).clamp(max=max_t) + t_list = model.rollout_t_list(4, model.flow_map_shift, max_t).double() + + assert torch.allclose(t_list, expected) + # endpoints are fixed by the map; the interior is pushed towards max_t + assert t_list[0].item() == pytest.approx(min(1.0, max_t)) and t_list[-1].item() == 0.0 + assert (t_list[1:-1] > grid.clamp(max=max_t)[1:-1]).all() + + # and the whole student update runs on that grid + data = _make_data(model, img_resolution=8) + loss_map, _ = model.single_train_step(data, 0) + assert torch.isfinite(loss_map["total_loss"]).all() + assert "vsd_loss" in loss_map and "bidirection_loss" in loss_map + + +def test_onpolicy_dmd_and_cotrain_noising_times_are_separate(): + """DMD2 draws its noising time from `sample_t_cfg`; the co-trained flow-map loss + draws (t, r) from `cotrain_sample_t_cfg`. + + The reference keeps these on two schedulers (`dmd_scheduler` vs `scheduler`), so + give the two configs disjoint ranges and check neither path reads the other's. + """ + model = _build_onpolicy_model() + dmd = model.config.sample_t_cfg + dmd.time_dist_type, dmd.min_t, dmd.max_t = "uniform", 0.70, 0.80 + cot = model.config.cotrain_sample_t_cfg + cot.time_dist_type, cot.min_t, cot.max_t = "uniform", 0.10, 0.20 + model._init_flow_map_loss(cot, model.config.cotrain_sample_r_cfg) + + # the mixin binds the co-train config, never DMD2's + assert model.flow_map_sample_t_cfg is cot + assert model.flow_map_sample_t_cfg is not dmd + + t_dmd = model._sample_noising_time(64, iteration=0) + assert ((t_dmd >= 0.70) & (t_dmd <= 0.80)).all(), t_dmd + + t_mf, r_mf, _ = model._sample_t_r_buckets(64) + assert ((t_mf >= 0.10) & (t_mf <= 0.20)).all(), t_mf + # r shares t's density unless sample_r_cfg is enabled; the consistency bucket pins 0 + assert ((r_mf >= 0.0) & (r_mf <= 0.20)).all(), r_mf @pytest.mark.parametrize("weight_type", ["beta08", "gaussian", "uniform"]) @@ -316,7 +390,7 @@ def test_timestep_weight_function(weight_type): # (t=0 excluded), which at the default num_steps=1000 is the reference's # set_timesteps grid: sum over the shifted grid == num_steps. num_steps = model.net.noise_scheduler.num_steps - shift = model.sample_t_cfg.shift + shift = model.flow_map_sample_t_cfg.shift grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64)[:-1] grid = shift * grid / (1 + (shift - 1) * grid) assert abs(model._compute_weight(torch.ones_like(grid), grid).sum().item() - num_steps) < 1e-6 @@ -498,23 +572,22 @@ def test_fuse_r_embedding_gated(): def test_validation_t_list_matches_shift(): - """The hardcoded validation schedule must equal the shifted grid. + """The validation schedule must equal the student's shifted sampling grid. - `config_anyflow_onpolicy.py` writes `sample_t_cfg.t_list` out as a constant - instead of deriving it at runtime, so nothing detects it going stale if - `shift` or `student_sample_steps` changes. Recompute it here from those two - settings and compare. + `config_anyflow_onpolicy.py` derives `sample_t_cfg.t_list` from + `cotrain_sample_t_cfg.shift`, the same shift `rollout_t_list` applies, so this pins + the two together and catches a stale override. Left as None, DMD2 hands `generator_fn` the noise scheduler's UNSHIFTED `linspace(max_t, 0, N+1)`, which is far off-policy for a shift=5 model -- - asserted below so the constant cannot silently degrade to that. + asserted below so the schedule cannot silently degrade to that. """ import fastgen.configs.experiments.WanT2V.config_anyflow_onpolicy as mod cfg = mod.create_config() - shift = float(cfg.model.sample_t_cfg.shift) + shift = float(cfg.model.cotrain_sample_t_cfg.shift) n = int(cfg.model.student_sample_steps) - max_t = float(cfg.model.net.max_t) # the bound `_rollout_t_list` clamps to + max_t = float(cfg.model.net.max_t) # the bound `rollout_t_list` clamps to grid = [1.0 - i / n for i in range(n + 1)] expected = [min(shift * x / (1 + (shift - 1) * x), max_t) for x in grid]