From 43759984a1066ac09252a47dbec539cc680bbcc5 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Sat, 16 May 2026 01:12:13 +0800 Subject: [PATCH 01/12] 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 99c0415b3f5c3e2dee5c44338b6e27f2bbfabfaa Mon Sep 17 00:00:00 2001 From: Enderfga Date: Sat, 16 May 2026 09:37:23 +0800 Subject: [PATCH 02/12] 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 f6c3668..ede813a 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 ab1174d293a5dd74b83ca75c6d9f1322213640bf Mon Sep 17 00:00:00 2001 From: Enderfga Date: Sat, 16 May 2026 13:15:51 +0800 Subject: [PATCH 03/12] 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 1671bb2be5f5ddd42a567660e595a234b1687cbb Mon Sep 17 00:00:00 2001 From: Enderfga Date: Fri, 22 May 2026 10:02:35 +0800 Subject: [PATCH 04/12] 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 ede813a..0a47a4e 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 3c98dd118e1d7d1c0540589a80d35bdfe1a57ea3 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Wed, 10 Jun 2026 22:25:45 +0800 Subject: [PATCH 05/12] 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 0a47a4e..40d2338 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 "") @@ -1123,6 +1169,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( @@ -1172,6 +1222,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 6bfb8dc954a3f6b687d5ab4ff12047d853356a04 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Wed, 10 Jun 2026 23:36:55 +0800 Subject: [PATCH 06/12] 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 2be625acd7dd8b9fb518f1a528024a856c36bfaf Mon Sep 17 00:00:00 2001 From: Enderfga Date: Wed, 10 Jun 2026 23:41:38 +0800 Subject: [PATCH 07/12] 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 f93a980db8d0f188da98e7226f8f9d37b9809458 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 22:50:14 +0800 Subject: [PATCH 08/12] 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 40d2338..2dadc87 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 de4addebfa98f147dec5c5e41952e1ecba43f1aa Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 23:02:37 +0800 Subject: [PATCH 09/12] 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 552f0776670e8e7a870afbe6c6bb92440a83a302 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 23:09:04 +0800 Subject: [PATCH 10/12] 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 9aa47ff81d51421928ed7eca9184292dfb988c1d Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 23:14:18 +0800 Subject: [PATCH 11/12] 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 e8c8c101d2048a37e30cc8d9a97b0594ecd71a53 Mon Sep 17 00:00:00 2001 From: Enderfga Date: Tue, 14 Jul 2026 23:20:58 +0800 Subject: [PATCH 12/12] 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)