diff --git a/fastgen/configs/experiments/DiT/config_mf_b.py b/fastgen/configs/experiments/DiT/config_mf_b.py index 490913e..6a7a7c8 100644 --- a/fastgen/configs/experiments/DiT/config_mf_b.py +++ b/fastgen/configs/experiments/DiT/config_mf_b.py @@ -26,7 +26,7 @@ def create_config(): config.model.sample_t_cfg.train_p_std = 1.0 config.model.sample_t_cfg.min_t = 0.0 config.model.sample_t_cfg.max_t = 0.999 - config.model.sample_t_cfg.r_sample_ratio = 0.25 + config.model.sample_t_cfg.flow_matching_ratio = 0.75 config.model.loss_config.norm_method = "poly_1.0" config.model.loss_config.norm_const = 1.0 diff --git a/fastgen/configs/experiments/EDM/config_mf_cifar10.py b/fastgen/configs/experiments/EDM/config_mf_cifar10.py index 19b41ca..79c50b9 100644 --- a/fastgen/configs/experiments/EDM/config_mf_cifar10.py +++ b/fastgen/configs/experiments/EDM/config_mf_cifar10.py @@ -19,7 +19,6 @@ def create_config(): config.model.sample_t_cfg.train_p_mean = -0.6 config.model.sample_t_cfg.train_p_std = 1.6 - config.model.sample_t_cfg.r_sample_ratio = 0.75 config.model.sample_t_cfg.time_dist_type = "logitnormal" config.model.sample_t_cfg.min_t = 0.0 config.model.sample_t_cfg.max_t = 0.999 diff --git a/fastgen/configs/experiments/WanT2V/config_anyflow.py b/fastgen/configs/experiments/WanT2V/config_anyflow.py new file mode 100644 index 0000000..852142b --- /dev/null +++ b/fastgen/configs/experiments/WanT2V/config_anyflow.py @@ -0,0 +1,127 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AnyFlow flow-map pretrain config on Wan-1.3B T2V (paper Stage 1). + +AnyFlow's pretrain objective is MeanFlow's with a fixed ``beta08`` per-timestep +weighting, a finite-difference JVP, shifted timestep sampling, and a +``consistency_ratio`` fraction of the batch pinned to ``r = 0`` — so this config +runs ``MeanFlowModel`` directly. The values below mirror the reference recipe +``train_wan1b_student_shift5_81f_480p_lr5e-5_6k_b32.yml``. + +Known deviations from the reference, both of which also apply to the on-policy stage: +full-rank fine-tuning instead of the paper's rank-256 LoRA; and under +``guidance_fuse_scale`` the 1/g rescaling of dF/dt is applied only to the samples that +kept their condition, where the reference's ``compute_central_difference`` rescales the +whole batch. A dropped sample's fused prediction is plain ``u_uncond``, so an ungated +1/g would regress it onto a different fixed point -- ours is gated on ``keep``. + +The on-policy stage (paper Stage 2) lives in ``config_anyflow_onpolicy.py``. +""" + +import copy + + +import fastgen.configs.methods.config_mean_flow as config_mean_flow +from fastgen.configs.data import VideoLoaderConfig +from fastgen.configs.net import Wan_1_3B_Config +from fastgen.methods import AnyFlowModel + + +def create_config(): + config = config_mean_flow.create_config() + + # ------ network: gated dual-timestep Wan (AnyFlow architecture) ------ + config.model.net = copy.deepcopy(Wan_1_3B_Config) + config.model.net.r_timestep = True + 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 + # Noise-schedule bounds, forwarded to RFNoiseSchedule. + config.model.net.min_t = 0.0 + config.model.net.max_t = 1.0 + + config.model.precision = "bfloat16" + # FSDP2 parameter storage and gradient reduction in fp32 while compute stays + # bfloat16 -- the same split as the reference. Takes + # effect only under FSDP (`trainer.ddp=False`); it is ignored under DDP, where + # params, grads and compute are all `precision`. + config.model.precision_fsdp = "float32" + + # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. + config.model.input_shape = [16, 21, 60, 104] + + # ------ 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 + # Reference epsilon=5 in 1000-step units = 5e-3 in FastGen's continuous time. + config.model.loss_config.jvp_finite_diff_eps = 5e-3 + # Rebalance the flow-map / consistency (r < t) sample losses to the global + # flow-matching (r = t) loss mean (reference scale_weight). + config.model.loss_config.rebalance_to_flow_matching = True + config.model.precision_amp_jvp = "float32" + + # Prediction-side guidance fusion with text dropout (reference: + # 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.guidance_fuse_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.0 + config.model.sample_t_cfg.max_t = 1.0 + # diffusion_ratio=0.5 of the batch keeps r = t (pure flow matching). + config.model.sample_t_cfg.flow_matching_ratio = 0.5 + # consistency_ratio=0.25 of the batch is pinned to r = 0 (the reference + # sets r = 0 pre-shift, and the shift maps 0 to 0). + config.model.sample_t_cfg.consistency_ratio = 0.25 + # The reference assigns both buckets by rank-indexed partition of the + # global batch, not by an independent per-sample draw. + config.model.sample_t_cfg.deterministic_buckets = True + + # ------ optimization (reference: AdamW lr=5e-5, wd=0, betas=(0.9, 0.95), + # max_grad_norm=1.0, 1000-step LR warmup, EMA decay 0.999) ------ + 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 + # Reference `ema_warmup_step: 1000` (as `start_iter = warmup_steps - 1`). + config.trainer.callbacks.ema.start_iter = 999 + + # ------ inference / validation ------ + config.model.student_sample_type = "ode" + config.model.student_sample_steps = 4 + # Shifted schedule under the same map the (t, r) sampling applies. The pretrain + # stage has a single scheduler in the reference, so it shares that shift. + config.model.sample_t_cfg.t_list = AnyFlowModel.rollout_t_list( + config.model.student_sample_steps, config.model.sample_t_cfg.shift, config.model.net.max_t + ).tolist() + + # ------ data / trainer ------ + config.dataloader_train = copy.deepcopy(VideoLoaderConfig) + 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/experiments/WanT2V/config_anyflow_onpolicy.py b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py new file mode 100644 index 0000000..225ae82 --- /dev/null +++ b/fastgen/configs/experiments/WanT2V/config_anyflow_onpolicy.py @@ -0,0 +1,177 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""AnyFlow on-policy distillation config on Wan-1.3B T2V (paper Stage 2). + +DMD2-style distribution matching with a flow-map student that generates via a +compressed rollout (jump -> fine step -> jump, gradient through all segments, +NFE sampled per iteration) and co-trains the Stage-1 flow-map loss at every +student update — see ``AnyFlowModel``. The values below mirror the reference +recipe ``train_wan1b_onpolicy_81f_480p_lr2e-6_1k_b32.yml``, which runs 1200 +generator updates. + +Two checkpoints must be supplied. The student is seeded from a Stage 1 trainer +checkpoint, read directly with no conversion; ``pretrained_ckpt_key_map`` below + takes it from the Stage 1 ``ema`` weights, matching the reference's +``pretrained_weight: ema``:: + + trainer.checkpointer.pretrained_ckpt_path=/checkpoints/0006000.pth + +Known deviations from the reference: full-rank fine-tuning instead of the paper's +rank-256 LoRA; the noising times for the DMD gradient and the fake score are drawn +on [0.001, 0.999] rather than the reference's [0, 1] (its ``dmd_cfg`` sets no +``min_timestep`` / ``max_timestep``, so it clamps to the full range) -- FastGen's +convention for the other rectified-flow Wan configs; and, in the co-trained loss, the +1/g rescaling of dF/dt under ``guidance_fuse_scale`` is gated on the samples that kept +their condition, where the reference's ``compute_central_difference`` rescales the whole +batch (see ``config_anyflow.py`` and ``FlowMapLossMixin._compute_mf_loss``). +""" + +import copy + + +import fastgen.configs.methods.config_anyflow as config_anyflow_default +from fastgen.configs.data import VideoLoaderConfig +from fastgen.configs.net import Wan_1_3B_Config +from fastgen.methods import AnyFlowModel + + +def create_config(): + 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 + # Full [0, 1] noise schedule, as in the pretrain stage + config.model.net.min_t = 0.0 + config.model.net.max_t = 1.0 + + # ------ teacher / fake score: PLAIN single-timestep Wan ------ + # Neither consumes r (DMD2 queries both at the instantaneous velocity r = t) + # and the reference builds them without the flow-map pathway at all. + config.model.teacher = copy.deepcopy(Wan_1_3B_Config) + config.model.teacher.r_timestep = False + config.model.teacher.min_t = config.model.net.min_t + config.model.teacher.max_t = config.model.net.max_t + + # Student init comes from the Stage 1 trainer checkpoint via + # trainer.checkpointer.pretrained_ckpt_path (see the module docstring); + # "ema" mirrors the reference's `pretrained_weight: ema`. + config.trainer.checkpointer.pretrained_ckpt_key_map = {"net": "ema", "ema": "ema"} + + config.model.precision = "bfloat16" + # FSDP2 parameter storage and gradient reduction in fp32 while compute stays + # bfloat16 -- the same split as the reference. Takes + # effect only under FSDP (`trainer.ddp=False`); it is ignored under DDP, where + # params, grads and compute are all `precision`. + config.model.precision_fsdp = "float32" + + # VAE compress ratio: (1 + T/4) * H/8 * W/8. 81-frame, 480p clips. + config.model.input_shape = [16, 21, 60, 104] + + # ------ DMD machinery ------ + # The reference Stage 2 has no adversarial loss: its "discriminator" is + # the fake score network, trained with denoising score matching only. + config.model.gan_loss_weight_gen = 0.0 + # The reference updates generator and fake score 1:1 (both every global + # 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 + + # DMD gradient noising time: reference `generator_loss` draws torch.rand + # then applies the shift -> shifted-uniform. The bounds keep the score models off + # the degenerate endpoints (see the deviation note above); `fake_score_sample_t_cfg` + # inherits them through the deepcopy below. + config.model.sample_t_cfg.time_dist_type = "shifted" + config.model.sample_t_cfg.shift = 5.0 + config.model.sample_t_cfg.min_t = 0.001 + config.model.sample_t_cfg.max_t = 0.999 + + # Fake-score noising time: reference `discriminator_loss` draws + # logit_normal(0, 1) then applies the same shift. This is a DIFFERENT + # density from the DMD path above, so it needs its own config. + config.model.fake_score_sample_t_cfg = copy.deepcopy(config.model.sample_t_cfg) + config.model.fake_score_sample_t_cfg.time_dist_type = "shifted_logitnormal" + config.model.fake_score_sample_t_cfg.train_p_mean = 0.0 + config.model.fake_score_sample_t_cfg.train_p_std = 1.0 + + # ------ student rollout (reference rollout_cfg) ------ + # The rollout grid's shift comes from `cotrain_sample_t_cfg` below: the reference + # builds its rollout pipeline from the same `scheduler` it draws the co-trained + # (t, r) from, separately from the DMD noising time above. + config.model.student_sample_type = "ode" + config.model.student_sample_steps_list = [2, 4, 8, 16, 50] + config.model.student_sample_steps = 4 + + # ------ co-trained Stage-1 flow-map loss (reference cotrain_forward_kl) ------ + # FastGen's VSD loss carries a 0.5 factor the reference's DMD loss does + # not; 0.5 here keeps the DMD : flow-map gradient ratio at the + # reference's 1 : 1. + 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" + 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_flow_matching = True + config.model.guidance_fuse_scale = 3.0 + config.model.cond_dropout_prob = 0.1 + config.model.precision_amp_jvp = "float32" + # (t, r) sampling for the co-trained loss, drawn from the reference's `scheduler`. + # Its shift also drives the student's rollout grid, so it is independent of the DMD + # noising-time shift above; the reference recipe sets both to 5.0. + config.model.cotrain_sample_t_cfg.time_dist_type = "shifted" + config.model.cotrain_sample_t_cfg.shift = 5.0 + config.model.cotrain_sample_t_cfg.min_t = 0.0 + config.model.cotrain_sample_t_cfg.max_t = 1.0 + config.model.cotrain_sample_t_cfg.flow_matching_ratio = 0.5 + config.model.cotrain_sample_t_cfg.consistency_ratio = 0.25 + config.model.cotrain_sample_t_cfg.deterministic_buckets = True + + # Validation schedule at `student_sample_steps` -- the same grid the per-NFE rollout + # builds, so it follows the co-trained sampling shift. + config.model.sample_t_cfg.t_list = AnyFlowModel.rollout_t_list( + config.model.student_sample_steps, config.model.cotrain_sample_t_cfg.shift, config.model.net.max_t + ).tolist() + + # ------ optimization (reference: AdamW lr=2e-6, betas=(0.0, 0.999), wd=0, + # grad clip 1.0, EMA 0.99) ------ + 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.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 + # EMA start = 6399 = 6000 + 400 - 1: a 400-iteration warmup (the reference's + # `ema_warmup_step: 200` x `student_update_freq = 2`) offset by the Stage-1 + # checkpoint iteration this stage seeds from -- `EMACallback` adds + # `model.resume_iter`, so a bare 399 is already past and never warms up. + # Seeding from another checkpoint: use + 399. + config.trainer.callbacks.ema.start_iter = 6399 + + # ------ data / trainer ------ + config.dataloader_train = copy.deepcopy(VideoLoaderConfig) + config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) + config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 + config.dataloader_train.batch_size = 1 + + # The reference runs 1200 global steps, each performing BOTH a generator + # and a fake-score update. FastGen alternates them across iterations + # (student_update_freq=2), so 1200 generator updates need 2 * 1200 = 2400 + # iterations here. + config.trainer.max_iter = 2400 + config.trainer.logging_iter = 100 + config.trainer.save_ckpt_iter = 400 + config.trainer.batch_size_global = 32 + + config.log_config.group = "wan_anyflow_onpolicy" + return config diff --git a/fastgen/configs/experiments/WanT2V/config_mf.py b/fastgen/configs/experiments/WanT2V/config_mf.py index be2e26d..710716e 100644 --- a/fastgen/configs/experiments/WanT2V/config_mf.py +++ b/fastgen/configs/experiments/WanT2V/config_mf.py @@ -58,21 +58,24 @@ def create_config(): config.model.net.r_embedder_init = "zero" config.model.net.norm_temb = False + # The consistency-distillation teacher (use_cd=True) is only ever queried for + # the instantaneous velocity, without r, so give it the plain architecture. + # Deriving it from `net` would allocate, zero-init and keep on device an + # r_embedder that is never invoked. + config.model.teacher = copy.deepcopy(Wan_1_3B_Config) + config.model.teacher.r_timestep = False + # we use simple diffusion version: 0.73 = 0.5 * log((21 * 60 * 104)/(64 * 64)) - 0.1 config.model.enable_preprocessors = False config.model.sample_t_cfg.time_dist_type = "logitnormal" - config.model.sample_t_cfg.r_sample_ratio = 1.0 + config.model.sample_t_cfg.flow_matching_ratio = 0.0 config.model.sample_t_cfg.train_p_mean = -0.8 config.model.sample_t_cfg.train_p_std = 1.6 config.model.sample_t_cfg.min_t = 0.001 config.model.sample_t_cfg.max_t = 0.999 - config.dataloader_train = VideoLatentLoaderConfig + config.dataloader_train = copy.deepcopy(VideoLatentLoaderConfig) config.dataloader_train.batch_size = 1 - # 480p (832x480) resolution - config.dataloader_train.img_size = (config.model.input_shape[-1] * 8, config.model.input_shape[-2] * 8) - config.dataloader_train.sequence_length = (config.model.input_shape[1] - 1) * 4 + 1 - config.log_config.group = "wan_mf" return config diff --git a/fastgen/configs/methods/config_anyflow.py b/fastgen/configs/methods/config_anyflow.py new file mode 100644 index 0000000..7e4113d --- /dev/null +++ b/fastgen/configs/methods/config_anyflow.py @@ -0,0 +1,113 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Config schema for the AnyFlow on-policy method (paper Stage 2). + +AnyFlow's on-policy stage is DMD2 with a flow-map student that generates via +a multi-step rollout-with-gradient, so the config is based on the DMD2 model +config. The flow-map pretrain stage (paper Stage 1) is MeanFlow with +AnyFlow's hyperparameters — see ``configs/experiments/WanT2V/config_anyflow.py``. +""" + +from typing import List, Optional + +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.configs.methods.config_mean_flow import ( + LossConfig as MeanFlowLossConfig, + SampleRConfig as MeanFlowSampleRConfig, + 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 — DMD2 plus the rollout / cotrain knobs. + + The MeanFlow loss / sampling configs drive the co-trained Stage-1 + 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. Separate from + # the inherited `sample_t_cfg`, which DMD2 uses for the noising time: the reference + # draws these from its `scheduler` and the noising time from its `dmd_scheduler`. + cotrain_sample_t_cfg: MeanFlowSampleTConfig = attrs.field(factory=MeanFlowSampleTConfig) + cotrain_sample_r_cfg: MeanFlowSampleRConfig = attrs.field(factory=MeanFlowSampleRConfig) + loss_config: MeanFlowLossConfig = attrs.field(factory=MeanFlowLossConfig) + + # Weight of the co-trained Stage-1 flow-map loss in the student update. + # 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 + + # Prediction-side guidance fusion for the co-trained flow-map loss; see + # `config_mean_flow.ModelConfig`. + # + # Deviates from the AnyFlow reference, which divides dF/dt by g for the whole + # batch. We divide only on samples that kept their condition: a dropped sample's + # fused prediction is plain `u_uncond`, so an ungated 1/g would regress it onto + # a different fixed point. + guidance_fuse_scale: Optional[float] = None + + # Text dropout for the co-trained flow-map loss (reference drop_text_ratio). + cond_dropout_prob: Optional[float] = None + cond_keys_no_dropout: List[str] = attrs.field(factory=list) + + # Precision for autocast in the co-trained loss JVP (None disables autocast there). + precision_amp_jvp: str | None = None + + # MeanFlow's target-side guidance knobs. AnyFlow typically guides on the + # PREDICTION side (`guidance_fuse_scale`), so these stay at their + # no-op defaults + guidance_mixture_ratio: Optional[float] = None + guidance_t_start: float = 0.0 + guidance_t_end: float = 1.0 + + +@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, + } + ) + + # 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] + config.model.fake_score_scheduler.warm_up_steps = [0] + config.model.discriminator_scheduler.warm_up_steps = [0] + + return config diff --git a/fastgen/configs/methods/config_dmd2.py b/fastgen/configs/methods/config_dmd2.py index a59af1e..73230bc 100644 --- a/fastgen/configs/methods/config_dmd2.py +++ b/fastgen/configs/methods/config_dmd2.py @@ -2,6 +2,8 @@ # SPDX-License-Identifier: Apache-2.0 import copy +from typing import Optional + import attrs from omegaconf import DictConfig @@ -9,6 +11,7 @@ from fastgen.configs.config import ( BaseModelConfig, BaseConfig, + SampleTConfig, ) from fastgen.configs.opt import BaseOptimizerConfig, BaseSchedulerConfig from fastgen.methods import DMD2Model @@ -33,6 +36,10 @@ class ModelConfig(BaseModelConfig): discriminator_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseOptimizerConfig)) discriminator_scheduler: DictConfig = attrs.field(factory=lambda: copy.deepcopy(BaseSchedulerConfig)) + # Optional noising-time distribution used ONLY on fake-score update + # iterations (`iteration % student_update_freq != 0`) + fake_score_sample_t_cfg: Optional[SampleTConfig] = None + # student update frequency student_update_freq: int = 5 diff --git a/fastgen/configs/methods/config_mean_flow.py b/fastgen/configs/methods/config_mean_flow.py index a6dba34..743537d 100644 --- a/fastgen/configs/methods/config_mean_flow.py +++ b/fastgen/configs/methods/config_mean_flow.py @@ -33,8 +33,18 @@ class SampleTConfig(BaseSampleTConfig): train_p_mean: float = -1.1 train_p_std: float = 2.0 - # ratio for randomly sampling r - r_sample_ratio: float = 0.0 + # fraction of the batch with r = t (pure flow matching); the rest keeps + # the randomly sampled r, minus the consistency fraction below + flow_matching_ratio: float = 0.25 + + # fraction of the batch forced to r = 0 (consistency-to-clean, used by + # AnyFlow; 0.0 keeps the original MeanFlow behavior) + consistency_ratio: float = 0.0 + + # how samples are assigned to the two buckets above: False draws each + # sample's bucket independently, True partitions the global batch by rank + # index so the bucket sizes are exact every iteration (used by AnyFlow) + deterministic_buckets: bool = False @attrs.define(slots=False) @@ -61,8 +71,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 @@ -71,6 +82,13 @@ class LossConfig: tangent_spatial_invariance: bool = False # loss type (choice between l2 and opt_grad) loss_type: str = "opt_grad" + # optional fixed per-timestep loss weighting evaluated as a function of t + # ("beta08", "gaussian", "uniform"). Multiplies the adaptive norm_method + # weight above; None disables it. + weight_type: Optional[str] = None + # rebalance the flow-map / consistency (r < t) sample losses to the global + # flow-matching (r = t) loss mean via a detached per-sample factor + rebalance_to_flow_matching: bool = False @attrs.define(slots=False) @@ -93,11 +111,17 @@ class ModelConfig(BaseModelConfig): # optimizer net_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(RAdamOptimizerConfig)) + # prediction-side guidance fusion scale (AnyFlow guidance distillation): + # the conditional output is trained to be the guided flow directly via + # (u_cond + (g-1) * u_uncond) / g against the raw data velocity. None + # keeps MeanFlow's target-side guidance (guidance_scale). + guidance_fuse_scale: Optional[float] = None + # condition dropout probability cond_dropout_prob: Optional[float] = None # 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 @@ -105,7 +129,7 @@ class ModelConfig(BaseModelConfig): # guidance t end guidance_t_end: float = 1.0 - # precision for autocast in JVP (none defaults to training precision) + # precision for autocast in JVP (none disables autocast in the JVP region) precision_amp_jvp: str | None = None diff --git a/fastgen/configs/methods/config_scm.py b/fastgen/configs/methods/config_scm.py index 6303bd4..b430fb0 100644 --- a/fastgen/configs/methods/config_scm.py +++ b/fastgen/configs/methods/config_scm.py @@ -77,7 +77,7 @@ class ModelConfig(BaseModelConfig): # optimizer net_optimizer: DictConfig = attrs.field(factory=lambda: copy.deepcopy(RAdamOptimizerConfig)) - # precision for autocast in JVP (none defaults to training precision) + # precision for autocast in JVP (none disables autocast in the JVP region) precision_amp_jvp: str | None = None diff --git a/fastgen/methods/README.md b/fastgen/methods/README.md index 3afb26f..f9770e8 100644 --- a/fastgen/methods/README.md +++ b/fastgen/methods/README.md @@ -14,6 +14,7 @@ Training methods for fast single-step or few-step generation from diffusion mode | | | f-distill | [`FdistillModel`](distribution_matching/f_distill.py) | f-divergence weighted DMD2 | [Xu et al., 2025](https://arxiv.org/abs/2502.15681) | | | | LADD | [`LADDModel`](distribution_matching/ladd.py) | Pure adversarial distillation | [Sauer et al., 2024](https://arxiv.org/abs/2403.12015) | | | | CausVid | [`CausVidModel`](distribution_matching/causvid.py) | Causal DMD2 with diffusion forcing | [Yin et al., 2024](https://arxiv.org/abs/2412.07772) | +| | | AnyFlow | [`AnyFlowModel`](distribution_matching/anyflow.py) | Any-step DMD2 with a flow-map student | [Gu et al., 2026](https://arxiv.org/abs/2605.13724) | | | | Self-Forcing | [`SelfForcingModel`](distribution_matching/self_forcing.py) | Causal DMD2 with self-forcing | [Huang et al., 2025](https://arxiv.org/abs/2506.08009) | | **Fine-Tuning** | [README](fine_tuning/README.md) | SFT | [`SFTModel`](fine_tuning/sft.py) | Finetuning with denoising score matching | [Ho et al., 2020](https://arxiv.org/abs/2006.11239), [Song et al., 2020](https://arxiv.org/abs/2011.13456), [Lipman et al., 2022](https://arxiv.org/abs/2210.02747), [Albergo et al., 2023](https://arxiv.org/abs/2303.08797) | | | | CausalSFT | [`CausalSFTModel`](fine_tuning/sft.py) | Causal version of SFT | [Chen et al., 2024](https://arxiv.org/abs/2407.01392) | diff --git a/fastgen/methods/__init__.py b/fastgen/methods/__init__.py index e902b82..b6ef462 100644 --- a/fastgen/methods/__init__.py +++ b/fastgen/methods/__init__.py @@ -15,6 +15,9 @@ from fastgen.methods.consistency_model.sCM import SCMModel as SCMModel from fastgen.methods.consistency_model.mean_flow import MeanFlowModel as MeanFlowModel +# AnyFlow pulls in mean_flow, so it must come after the CM import above +from fastgen.methods.distribution_matching.anyflow import AnyFlowModel as AnyFlowModel + from fastgen.methods.fine_tuning.sft import SFTModel as SFTModel from fastgen.methods.fine_tuning.sft import CausalSFTModel as CausalSFTModel diff --git a/fastgen/methods/consistency_model/README.md b/fastgen/methods/consistency_model/README.md index 867dfa6..ad0e82b 100644 --- a/fastgen/methods/consistency_model/README.md +++ b/fastgen/methods/consistency_model/README.md @@ -73,7 +73,7 @@ Two-stage training: frozen Stage-1 CM for `t < transition_t`, trainable student ## MeanFlow -**File:** [`mean_flow.py`](mean_flow.py) | **Reference:** [Geng et al., 2025](https://arxiv.org/abs/2505.13447) +**File:** [`mean_flow.py`](mean_flow.py) | **References:** [Geng et al., 2025](https://arxiv.org/abs/2505.13447), [Sabour et al., 2025](https://arxiv.org/abs/2506.14603) Learns average velocity between trajectory points: `x_r = x_t - (t-r) · u(x_t, t, r)`. @@ -81,10 +81,11 @@ Learns average velocity between trajectory points: `x_r = x_t - (t-r) · u(x_t, - `loss_config.use_cd`: Use consistency distillation (requires teacher; `guidance_scale` controls CFG for the teacher) - `loss_config.use_jvp_finite_diff`: Use finite difference for JVP (e.g., for compatibility with Flash Attention and FSDP) - `sample_t_cfg`, `sample_r_cfg`: Configs of the distributions for sampling `t` and `r` -- `sample_t_cfg.r_sample_ratio`: Ratio for flow matching loss +- `sample_t_cfg.flow_matching_ratio`: Fraction of the batch with `r = t` (flow matching loss) +- `guidance_scale`, `guidance_fuse_scale`: If `guidance_fuse_scale` is `None`, use target-side guidance (`guidance_scale`, 3-way mixture via `guidance_mixture_ratio`, t-window via `guidance_t_start` / `guidance_t_end`), otherwise use prediction-side guidance (conditional output learns the guided flow directly) -**Configs:** [`EDM/config_mf_cifar10.py`](../../configs/experiments/EDM/config_mf_cifar10.py), [`DiT/config_mf_b.py`](../../configs/experiments/DiT/config_mf_b.py), [`DiT/config_mf_xl.py`](../../configs/experiments/DiT/config_mf_xl.py), [`WanT2V/config_mf.py`](../../configs/experiments/WanT2V/config_mf.py) +**Configs:** [`EDM/config_mf_cifar10.py`](../../configs/experiments/EDM/config_mf_cifar10.py), [`DiT/config_mf_b.py`](../../configs/experiments/DiT/config_mf_b.py), [`DiT/config_mf_xl.py`](../../configs/experiments/DiT/config_mf_xl.py), [`WanT2V/config_mf.py`](../../configs/experiments/WanT2V/config_mf.py), [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) (Stage 1 of [AnyFlow](../distribution_matching/README.md#anyflow)) **Expected results:** diff --git a/fastgen/methods/consistency_model/mean_flow.py b/fastgen/methods/consistency_model/mean_flow.py index 485e752..5b9710d 100644 --- a/fastgen/methods/consistency_model/mean_flow.py +++ b/fastgen/methods/consistency_model/mean_flow.py @@ -9,8 +9,10 @@ import numpy as np import torch from fastgen.methods import CMModel +from fastgen.networks.noise_schedule import time_shift from fastgen.utils import basic_utils, expand_like from fastgen.utils.basic_utils import convert_cfg_to_dict +from fastgen.utils.distributed import get_rank, world_size import fastgen.utils.logging_utils as logger @@ -48,57 +50,99 @@ def temp_disable_efficient_attn(device_type: str = "cuda"): yield -class MeanFlowModel(CMModel): - def __init__(self, config: ModelConfig): - """ +class FlowMapLossMixin: + """The MeanFlow flow-map regression objective, shared across methods. + + Holds everything needed to turn a real batch into the flow-map loss + ``|| u_theta(x_t, t, r) - sg(v - (t - r) du/dt) ||^2``: the (t, r) sampler + with its flow-matching / consistency buckets, the JVP estimate, the loss + weighting, and the reduction. ``MeanFlowModel`` uses it as its whole + training objective; distribution-matching methods can co-train it alongside + their own objective. + + The host class must provide ``net``, ``device``, ``config`` and ``precision_amp`` + (plus ``teacher`` when ``loss_config.use_cd``), and call ``_init_flow_map_loss`` + from its ``__init__``. + """ + + def _init_flow_map_loss(self, sample_t_cfg: Any, sample_r_cfg: Any) -> None: + """Bind the flow-map loss / sampling configs and the JVP precision. + + The (t, r) configs are passed in rather than read off ``self.config``: the host + may draw the flow-map times from a different density than its own + ``sample_t_cfg`` (AnyFlow does -- see ``AnyFlowModel``), and ``CMModel`` already + owns a ``self.sample_t_cfg`` attribute that must not be rebound here. Args: - config (ModelConfig): The configuration for the MeanFlow model + sample_t_cfg: Config for sampling the flow-map ``t``. + sample_r_cfg: Config for sampling the flow-map ``r``. """ - super().__init__(config) - self.config = config - self.sample_t_cfg = self.config.sample_t_cfg - self.sample_r_cfg = self.config.sample_r_cfg + self.flow_map_sample_t_cfg = sample_t_cfg + self.flow_map_sample_r_cfg = sample_r_cfg self.loss_config = self.config.loss_config - # Precision for JVP - if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp: + # The shift `_sample_t_r_buckets` actually applies: `shift` is inert unless the + # density is one of `RFNoiseSchedule`'s shifted ones, so resolve it once here and + # reuse it wherever the flow-map timestep grid is rebuilt. + shifted = sample_t_cfg.time_dist_type in ("shifted", "shifted_logitnormal") + self.flow_map_shift = sample_t_cfg.shift if shifted else 1.0 + + # None runs the JVP with autocast disabled: the region always sets the autocast + # state, so the enclosing training autocast does not carry over into it. + if self.config.precision_amp_jvp is None: self.precision_amp_jvp = None else: self.precision_amp_jvp = basic_utils.PRECISION_MAP[self.config.precision_amp_jvp] logger.critical(f"Using precision {self.precision_amp_jvp} for JVP") - def _mix_condition( - self, - condition: Any, - neg_condition: torch.Tensor, - dxt_dt: torch.Tensor, - guided_dxt_dt: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - if self.config.cond_dropout_prob is None: - return condition, dxt_dt - - batch_size = dxt_dt.shape[0] - # Decide how many to drop first. - num_to_drop = (torch.rand(batch_size, device=dxt_dt.device) < self.config.cond_dropout_prob).sum() - # Create the mask to keep all but the first samples (the order is important to ensures most dropout happens at flow matching loss) - mask = torch.arange(batch_size, device=dxt_dt.device) >= num_to_drop - dxt_dt = torch.where(expand_like(mask, dxt_dt), guided_dxt_dt, dxt_dt) + # The fixed per-timestep loss weight is normalized to mean one over the + # network's discrete training timesteps (t = 0 excluded), mapped through the + # same shift `_sample_t_r_buckets` applies -- the weight has to be normalized + # over the grid the loss's own t values are drawn from, which is why this reads + # the passed-in config. The constant depends only on it, so derive it once here. + self._timestep_weight_scale: Optional[float] = None + if self.loss_config.weight_type is not None: + num_steps = self.net.noise_scheduler.num_steps + grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64)[:-1] + grid = time_shift(grid, self.flow_map_shift) + self._timestep_weight_scale = float(num_steps / self._timestep_weight_raw(grid).sum()) + + def _drop_condition( + self, condition: Any, neg_condition: Any, batch_size: int, device: torch.device + ) -> Tuple[Any, torch.Tensor]: + """Replace the condition with neg_condition for a per-sample subset. + + Returns ``(condition, keep)``; ``keep`` is the ``[B]`` bool mask of the + samples that stayed conditional, so callers can reuse the same subset. + + ``deterministic_buckets`` decides whether an index carries bucket + information: if so the buckets are cut on the GLOBAL index and an + index-based rule would only hit flow matching on rank 0, so draw per + sample; otherwise drop the first ``num_to_drop``. + """ + # Dropout disabled, or no negative condition to swap in: every sample + # stays conditional. + if self.config.cond_dropout_prob is None or neg_condition is None: + return condition, torch.ones(batch_size, dtype=torch.bool, device=device) + + if self.flow_map_sample_t_cfg.deterministic_buckets: + keep = torch.rand(batch_size, device=device) >= self.config.cond_dropout_prob + else: + num_to_drop = (torch.rand(batch_size, device=device) < self.config.cond_dropout_prob).sum() + keep = torch.arange(batch_size, device=device) >= num_to_drop if isinstance(condition, torch.Tensor): - condition = torch.where(expand_like(mask, condition), condition, neg_condition) - elif isinstance(condition, dict): - condition = condition.copy() - keys_no_drop = self.config.cond_keys_no_dropout - assert set(keys_no_drop).issubset( + return torch.where(expand_like(keep, condition), condition, neg_condition), keep + if isinstance(condition, dict): + keys_no_drop = set(self.config.cond_keys_no_dropout) + assert keys_no_drop.issubset( condition.keys() - ), f"keys_no_drop: {keys_no_drop} not in {condition.keys()}" + ), f"cond_keys_no_dropout: {keys_no_drop} not in {condition.keys()}" + condition = condition.copy() for k in condition.keys() - keys_no_drop: - condition[k] = torch.where(expand_like(mask, condition[k]), condition[k], neg_condition[k]) - else: - raise TypeError(f"Unsupported type: {type(condition)}") - - return condition, dxt_dt + condition[k] = torch.where(expand_like(keep, condition[k]), condition[k], neg_condition[k]) + return condition, keep + raise TypeError(f"Unsupported condition type: {type(condition)}") @torch.no_grad() def _get_velocity( @@ -106,14 +150,37 @@ def _get_velocity( x: torch.Tensor, z: torch.Tensor, t: torch.Tensor, - condition: Optional[torch.Tensor] = None, - neg_condition: Optional[torch.Tensor] = None, - ) -> Tuple[torch.Tensor, torch.Tensor]: + condition: Optional[Any] = None, + neg_condition: Optional[Any] = None, + ) -> Tuple[Any, torch.Tensor, torch.Tensor]: + """Regression target for the flow-map loss, the condition it was built from, + and the ``[B]`` mask of samples that stayed conditional. + + Two independent choices: + + * ``loss_config.use_cd`` picks the target SOURCE -- the teacher, or the + conditional data velocity. + * ``guidance_fuse_scale`` guides the *prediction* in + ``_compute_mf_loss``, so the target needs no guidance of its own and we only + drop the condition. Otherwise we guide the target too: through the teacher, or + -- without one -- by the net's own cond/uncond pass. + """ + fuse_scale = self.config.guidance_fuse_scale + if fuse_scale is not None: + assert fuse_scale > 0, f"guidance_fuse_scale must be > 0, got {fuse_scale} (None disables fusion)" + assert ( + neg_condition is not None + ), "guidance_fuse_scale requires neg_condition: the unconditional branch is queried at the same (t, r)" + condition, keep = self._drop_condition(condition, neg_condition, x.shape[0], x.device) + else: + keep = torch.ones(x.shape[0], dtype=torch.bool, device=x.device) + x_t = self.net.noise_scheduler.forward_process(x, z, t) if self.loss_config.use_cd: dxt_dt = self.teacher(x_t, t, condition=condition, fwd_pred_type="flow") - if self.config.guidance_scale is not None: + # Under fusion the target stays unguided, or it would be guided twice. + if self.config.guidance_scale is not None and fuse_scale is None: guidance_scale = torch.where( ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), self.config.guidance_scale, @@ -125,39 +192,45 @@ def _get_velocity( else: dxt_dt = self.net.noise_scheduler.cond_velocity(x=x, eps=z, t=t) - # unconditional score estimation from meanflow eq (19) - if self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None: + # unconditional score estimation from meanflow eq (19). Skipped under + # prediction-side fusion: that is this same guidance, moved onto the + # prediction, and the dropout it needs already ran above. + if fuse_scale is None and ( + self.config.guidance_scale is not None or self.config.guidance_mixture_ratio is not None + ): # Turn off dropout - self.net.eval() - neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow") - guidance_scale = self.config.guidance_scale or 1.0 - guidance_scale = torch.where( - ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - guidance_scale, - 1.0, - ) - guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) - - if self.config.guidance_mixture_ratio is None: - guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt) - else: - guidance_mixture_ratio = torch.where( + with basic_utils.train_mode(self.net, mode=False): + neg_dxt_dt = self.net(x_t, t, r=t, condition=neg_condition, fwd_pred_type="flow") + guidance_scale = self.config.guidance_scale or 1.0 + guidance_scale = torch.where( ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), - self.config.guidance_mixture_ratio, - 0.0, + guidance_scale, + 1.0, ) - guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype) - cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow") - guided_dxt_dt = ( - guidance_scale * dxt_dt - + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt - + guidance_mixture_ratio * cond_dxt_dt - ) - - self.net.train() - condition, dxt_dt = self._mix_condition(condition, neg_condition, dxt_dt, guided_dxt_dt) - - return condition, dxt_dt + guidance_scale = expand_like(guidance_scale, x_t).to(dtype=x_t.dtype) + + if self.config.guidance_mixture_ratio is None: + guided_dxt_dt = neg_dxt_dt + guidance_scale * (dxt_dt - neg_dxt_dt) + else: + guidance_mixture_ratio = torch.where( + ((t >= self.config.guidance_t_start) & (t <= self.config.guidance_t_end)), + self.config.guidance_mixture_ratio, + 0.0, + ) + guidance_mixture_ratio = expand_like(guidance_mixture_ratio, x_t).to(dtype=x_t.dtype) + cond_dxt_dt = self.net(x_t, t, r=t, condition=condition, fwd_pred_type="flow") + guided_dxt_dt = ( + guidance_scale * dxt_dt + + (1.0 - guidance_scale - guidance_mixture_ratio) * neg_dxt_dt + + guidance_mixture_ratio * cond_dxt_dt + ) + + condition, keep = self._drop_condition(condition, neg_condition, x_t.shape[0], x_t.device) + # Same subset: a kept sample is conditional + guided, a dropped one + # unconditional + unguided. + dxt_dt = torch.where(expand_like(keep, dxt_dt), guided_dxt_dt, dxt_dt) + + return condition, dxt_dt, keep def _estimate_jvp_finite_difference( self, @@ -251,20 +324,137 @@ def net_wrapper(x_t, t, r): return u_theta_jvp + def _warn_on_degenerate_buckets(self, n_flow_matching: int, n_consistency: int, global_bsz: int) -> None: + """Warn once if a requested bucket rounds away under the deterministic partition.""" + requested_but_empty = (self.flow_map_sample_t_cfg.flow_matching_ratio > 0 and n_flow_matching == 0) or ( + self.flow_map_sample_t_cfg.consistency_ratio > 0 and n_consistency == 0 + ) + if not requested_but_empty or getattr(self, "_bucket_warned", False): + return + self._bucket_warned = True + logger.warning( + f"The deterministic (t, r) bucket partition is degenerate: with " + f"world_size * batch_size = {global_bsz}, flow_matching_ratio=" + f"{self.flow_map_sample_t_cfg.flow_matching_ratio} and consistency_ratio=" + f"{self.flow_map_sample_t_cfg.consistency_ratio} yield bucket sizes " + f"({n_flow_matching}, {n_consistency}). The partition spans ranks but not " + f"gradient-accumulation rounds, so empty buckets stay empty every iteration." + ) + + def _sample_t_r_buckets(self, batch_size: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Sample (t, r) with t >= r and assign the per-batch buckets. + + Returns ``(t, r, r_eq_t_mask)`` where ``r_eq_t_mask`` marks the + samples with ``r = t`` (pure flow matching). + + The batch splits three ways: a ``flow_matching_ratio`` fraction gets + ``r = t``, a ``consistency_ratio`` fraction gets ``r = 0`` + (consistency to clean data), and the rest keep the sampled random pair. + """ + t_sample_kwargs = convert_cfg_to_dict(self.flow_map_sample_t_cfg) + t = self.net.noise_scheduler.sample_t(batch_size, **t_sample_kwargs, device=self.device) + r_sample_kwargs = ( + convert_cfg_to_dict(self.flow_map_sample_r_cfg) if self.flow_map_sample_r_cfg.enabled else t_sample_kwargs + ) + r = self.net.noise_scheduler.sample_t(batch_size, **r_sample_kwargs, device=self.device) + t, r = torch.maximum(t, r), torch.minimum(t, r) + assert torch.all(t >= r), "r cannot be larger than t" + + flow_matching_ratio = self.flow_map_sample_t_cfg.flow_matching_ratio + consistency_ratio = self.flow_map_sample_t_cfg.consistency_ratio + assert ( + flow_matching_ratio + consistency_ratio <= 1.0 + ), f"flow_matching_ratio + consistency_ratio must be <= 1, got {flow_matching_ratio} + {consistency_ratio}" + + # Both policies produce two bucket sizes and index the batch the same + # way; they differ only in how the sizes are obtained. + if self.flow_map_sample_t_cfg.deterministic_buckets: + global_bsz = world_size() * batch_size + n_flow_matching = round(flow_matching_ratio * global_bsz) + n_consistency = round(consistency_ratio * global_bsz) + self._warn_on_degenerate_buckets(n_flow_matching, n_consistency, global_bsz) + position = get_rank() * batch_size + torch.arange(batch_size, device=self.device) + else: + # Both sizes come from one uniform draw, cut at either end of [0, 1). + # The two events are disjoint by the assert above, so the buckets + # cannot overlap. + uniform = torch.rand(batch_size, device=self.device) + n_flow_matching = (uniform >= 1.0 - flow_matching_ratio).sum() + n_consistency = (uniform < consistency_ratio).sum() + position = torch.arange(batch_size, device=self.device) + + r_eq_t_mask = position < n_flow_matching + is_consistency = (position >= n_flow_matching) & (position < n_flow_matching + n_consistency) + r = torch.where(r_eq_t_mask, t, r) + r = torch.where(is_consistency, torch.zeros_like(r), r) + + return t, r, r_eq_t_mask + + def _reduce_mf_loss(self, mf_loss: torch.Tensor, r_eq_t_mask: torch.Tensor) -> torch.Tensor: + """Reduce the per-sample loss to a scalar, optionally rebalancing. + + With ``loss_config.rebalance_to_flow_matching`` set, every flow-map / + consistency (r < t) sample's loss is multiplied by the detached factor + ``mean(global flow-matching losses) / (own loss + 1e-5)``, so those + gradients are self-normalized and rescaled to the global + flow-matching-loss mean. + """ + # No rank-local condition here: the branch must be taken (or not) by + # every rank so the collective below cannot deadlock. Applying the + # scale to an empty ~r_eq_t_mask selection is a no-op. + if getattr(self.loss_config, "rebalance_to_flow_matching", False): + with torch.no_grad(): + # The global flow-matching-loss mean only needs the global sum + # and count, so reduce two scalars instead of gathering the + # per-sample losses (equivalent, 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: + torch.distributed.all_reduce(fm_loss_sum) + torch.distributed.all_reduce(fm_count) + scale = torch.ones_like(mf_loss) + 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() + + 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 _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._timestep_weight_raw(t) * self._timestep_weight_scale) assert ( weight.shape == tensor.shape @@ -306,8 +496,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) - loss = torch.sum(loss, dim=list(range(1, loss.ndim))) - weight = self._compute_weight(loss) + if self.loss_config.norm_method is None: + # Without adaptive normalization the reduction matters: the + # per-element mean keeps the loss scale independent of the + # sample dimensionality. + loss = torch.mean(loss, dim=list(range(1, loss.ndim))) + else: + loss = torch.sum(loss, dim=list(range(1, loss.ndim))) + weight = self._compute_weight(loss, t) loss = loss * weight # use explicit gradient @@ -321,7 +517,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))) @@ -402,7 +598,7 @@ def _compute_mf_loss( z = torch.randn_like(real_data) x_t = self.net.noise_scheduler.forward_process(real_data, z, t) - condition, dxt_dt = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) + condition, dxt_dt, keep = self._get_velocity(real_data, z, t, condition=condition, neg_condition=neg_condition) # prevent JVP to use cached conversions (which can break the computational graph) that were created in the no_grad context of _get_velocity torch.clear_autocast_cache() u_theta_jvp = self._jvp(x_t, t, r, dxt_dt, condition=condition) @@ -417,6 +613,21 @@ def _compute_mf_loss( condition=condition, fwd_pred_type="flow", ) + + guidance_fuse_scale = self.config.guidance_fuse_scale + if guidance_fuse_scale is not None: + # Guidance distillation on the PREDICTION side (see `_get_velocity`): the + # conditional output learns the guided flow directly, so only the prediction + # changes. The uncond branch is queried at the SAME (t, r) flow-map slice, + # giving (u_cond + (g - 1) * u_uncond) / g; dF/dt is then the finite + # difference over g on conditional samples, with the unconditional + # derivative dropped. + u_theta_jvp = torch.where(expand_like(keep, u_theta_jvp), u_theta_jvp / guidance_fuse_scale, u_theta_jvp) + # Turn off dropout for the unconditional pass, as the target-side path does. + with basic_utils.train_mode(self.net, mode=False), torch.no_grad(): + u_uncond = self.net(x_t, t, r=r, condition=neg_condition, fwd_pred_type="flow") + u_theta = (u_theta + (guidance_fuse_scale - 1.0) * u_uncond) / guidance_fuse_scale + mf_loss, tangent, loss_weight, warmup_weight = self._mf_pred_to_loss( u_theta=u_theta, u_theta_jvp=u_theta_jvp, x_t=x_t, dxt_dt=dxt_dt, t=t, r=r, iteration=iteration ) @@ -436,6 +647,18 @@ def _compute_mf_loss( warmup_weight, ) + +class MeanFlowModel(FlowMapLossMixin, CMModel): + def __init__(self, config: ModelConfig): + """ + + Args: + config (ModelConfig): The configuration for the MeanFlow model + """ + super().__init__(config) + self.config = config + self._init_flow_map_loss(config.sample_t_cfg, config.sample_r_cfg) + def single_train_step( self, data: Dict[str, Any], iteration: int ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor | Callable]]: @@ -454,19 +677,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) + # sample t and r, with per-batch buckets (flow matching / consistency) + t, r, r_eq_t_mask = self._sample_t_r_buckets(batch_size) ( mf_loss, @@ -486,7 +698,7 @@ def single_train_step( neg_condition=neg_condition, ) - loss = mf_loss.mean() + loss = self._reduce_mf_loss(mf_loss, r_eq_t_mask) loss_map = { "total_loss": loss, "mf_loss": loss, diff --git a/fastgen/methods/consistency_model/sCM.py b/fastgen/methods/consistency_model/sCM.py index 3a8259b..85d8c82 100644 --- a/fastgen/methods/consistency_model/sCM.py +++ b/fastgen/methods/consistency_model/sCM.py @@ -93,8 +93,9 @@ def __init__(self, config: ModelConfig): self.loss_config = self.config.loss_config self.sigma_data = self.sample_t_cfg.sigma_data - # Precision for JVP - if self.config.precision_amp_jvp is None or self.config.precision_amp_jvp == self.precision_amp: + # Precision for JVP. None disables autocast in the JVP region: it always sets the + # autocast state, so the enclosing training autocast does not carry over into it. + if self.config.precision_amp_jvp is None: self.precision_amp_jvp = None else: self.precision_amp_jvp = PRECISION_MAP[self.config.precision_amp_jvp] diff --git a/fastgen/methods/distribution_matching/README.md b/fastgen/methods/distribution_matching/README.md index 021e181..d465fd8 100644 --- a/fastgen/methods/distribution_matching/README.md +++ b/fastgen/methods/distribution_matching/README.md @@ -85,6 +85,23 @@ DMD2 extended for causal video generation with autoregressive chunk-by-chunk pro --- +## AnyFlow + +**File:** [`anyflow.py`](anyflow.py) | **Reference:** [Gu et al. (2026)](https://arxiv.org/abs/2605.13724) + +DMD2 with a flow-map student `u(x_t, t, r)` (mean velocity from `t` back to `r`), supporting arbitrary inference NFE. The student generates by rolling out the flow map from pure noise with the NFE sampled per iteration and gradients through all segments, and co-trains the flow-map loss at every update. Requires a Stage-1 flow-map pretrain, which runs directly on [`MeanFlowModel`](../consistency_model/mean_flow.py) using AnyFlow's hyperparameters. + +**Key Parameters (Stage 2):** +- `student_sample_steps_list`: Rollout NFEs sampled per iteration +- `cotrain_pretrain_weight`: Weight of the co-trained Stage-1 flow-map loss +- `cotrain_sample_t_cfg`, `cotrain_sample_r_cfg`: Separate `(t, r)` distributions for the co-trained flow-map loss; its shift also defines the rollout schedule +- Requires a Stage-1 pretrained checkpoint via `trainer.checkpointer.pretrained_ckpt_path` +- See also the key parameters of DMD2 above (and the key parameters of [MeanFlow](../consistency_model/README.md#meanflow) for Stage 1) + +**Configs:** [`WanT2V/config_anyflow.py`](../../configs/experiments/WanT2V/config_anyflow.py) (Stage 1 pretrain), [`WanT2V/config_anyflow_onpolicy.py`](../../configs/experiments/WanT2V/config_anyflow_onpolicy.py) (Stage 2 on-policy distillation) + +--- + ## 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..32bbcd4 --- /dev/null +++ b/fastgen/methods/distribution_matching/anyflow.py @@ -0,0 +1,183 @@ +# 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 flow-map model ``u_theta(x_t, t, r)`` predicting the +average velocity from ``t`` to ``r`` (``r <= t``), so one model serves any +inference NFE: each Euler-like step picks its own integration interval. + +Stage 1 (flow-map pretrain) is MeanFlow with AnyFlow's hyperparameters, run on +``MeanFlowModel`` directly (``configs/experiments/WanT2V/config_anyflow.py``) — +there is no AnyFlow-specific pretrain code. Stage 2 (this module) is DMD2 on the +pretrained flow-map weights, deviating from stock DMD2 twice: the student +generates via a compressed flow-map rollout with gradient through all segments +(``gen_data_from_net``), and every student update co-trains the Stage-1 flow-map +loss on the real batch (``cotrain_pretrain_weight``). +""" + +from __future__ import annotations + +from typing import Any, Optional, TYPE_CHECKING + +import torch + +from fastgen.methods.consistency_model.mean_flow import FlowMapLossMixin +from fastgen.methods.distribution_matching.dmd2 import DMD2Model +from fastgen.networks.noise_schedule import time_shift +from fastgen.utils.distributed import world_size +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(FlowMapLossMixin, DMD2Model): + """AnyFlow on-policy stage: DMD2 with a flow-map rollout student. + + ``FlowMapLossMixin`` supplies the co-trained Stage-1 objective and the + flow-map validation sample loop, which integrates with ``r = t_next``. + """ + + def __init__(self, config: ModelConfig): + super().__init__(config) + self.config = config + # The co-trained flow-map loss draws its (t, r) from its own config, not from + # `sample_t_cfg` -- DMD2 keeps that one for the noising time. + self._init_flow_map_loss(config.cotrain_sample_t_cfg, config.cotrain_sample_r_cfg) + + logger.info( + f"AnyFlow on-policy: student_sample_steps_list={self.config.student_sample_steps_list}, " + 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}" + ) + + # ------------------------------------------------------------------ + # Rollout + # ------------------------------------------------------------------ + + def _generate_noise_and_time( + self, real_data: torch.Tensor, iteration: Optional[int] = None + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """The student rollout always starts from pure noise at ``max_t``.""" + batch_size = real_data.shape[0] + + eps_student = torch.randn(batch_size, *self.input_shape, device=self.device, dtype=real_data.dtype) + t_student = torch.full( + (batch_size,), + self.net.noise_scheduler.max_t, + device=self.device, + dtype=self.net.noise_scheduler.t_precision, + ) + input_student = self.net.noise_scheduler.latents(noise=eps_student) + + t = self._sample_noising_time(batch_size, iteration) + eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) + return input_student, t_student, t, eps + + def _broadcast_choice(self, high: int) -> int: + """Pick an index in [0, high) on rank 0 and broadcast it.""" + idx = torch.randint(0, high, (1,), device=self.device, dtype=torch.long) + if world_size() > 1: + torch.distributed.broadcast(idx, src=0) + return int(idx.item()) + + @staticmethod + def rollout_t_list(num_steps: int, shift: float = 1.0, max_t: float = 1.0) -> torch.Tensor: + """Shifted timestep schedule ``[max_t, ..., 0]`` with ``num_steps + 1`` entries. + + Equivalent to the reference scheduler's ``set_timesteps``, with ``shift`` the + model's ``flow_map_shift``. Static so the configs can derive the matching + validation ``t_list`` from it; float64 on the CPU for the caller to cast. + """ + grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64) + return time_shift(grid, shift).clamp(max=max_t) + + def gen_data_from_net( + self, + input_student: torch.Tensor, + t_student: torch.Tensor, + condition: Optional[Any] = None, + ) -> torch.Tensor: + """Flow-map rollout compressed into at most three network forwards. + + Mirrors the reference ``WanAnyFlowPipeline.training_rollout``: on a + ``num_steps``-step schedule ``t_0 > ... > t_N = 0`` with a sampled fine-step + position ``g``, jump ``t_0 -> t_g``, take the fine step ``t_g -> t_{g+1}``, + then jump ``t_{g+1} -> 0``, each one forward of mean-velocity sampling + ``u_theta(x_t, t, r=t_next)``. Gradient flows through all segments; the + fake-score update wraps this call in ``no_grad`` at the caller. + """ + del t_student # the rollout schedule is built below + + # This iteration's rollout NFE, as the reference draws it + # (``random.choice(num_inference_steps_list)``, rank-0 broadcast in both the + # generator and the fake-score update); no list means the fixed + # ``student_sample_steps``. + steps_list = self.config.student_sample_steps_list + if steps_list: + num_steps = int(steps_list[self._broadcast_choice(len(steps_list))]) + else: + num_steps = int(self.config.student_sample_steps) + assert num_steps >= 1, f"rollout steps must be >= 1, got {num_steps}" + grad_step = self._broadcast_choice(num_steps) + ns = self.net.noise_scheduler + t_list = self.rollout_t_list(num_steps, self.flow_map_shift, float(ns.max_t)).to( + device=self.device, dtype=ns.t_precision + ) + + # The leading jump exists only for grad_step > 0 and the trailing one + # only for grad_step + 1 < num_steps. + seg_t = [t_list[0]] if grad_step > 0 else [] + seg_t += [t_list[grad_step], t_list[grad_step + 1]] + if grad_step + 1 < num_steps: + seg_t.append(t_list[-1]) + + return self._student_sample_loop( + self.net, + input_student, + t_list=torch.stack(seg_t), + condition=condition, + student_sample_type="ode", + ) + + # ------------------------------------------------------------------ + # Training step — DMD2 plus the co-trained Stage-1 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, iteration=iteration) + + 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-1 bidirection (flow-map) loss on the + # real batch. + t_mf, r_mf, r_eq_t_mask = self._sample_t_r_buckets(real_data.shape[0]) + mf_outputs = self._compute_mf_loss( + 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], 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 + return loss_map, outputs + + return self._fake_score_discriminator_update_step( + input_student, t_student, t, eps, real_data, condition=condition + ) diff --git a/fastgen/methods/distribution_matching/causvid.py b/fastgen/methods/distribution_matching/causvid.py index 9f73b2e..c2d1e0d 100644 --- a/fastgen/methods/distribution_matching/causvid.py +++ b/fastgen/methods/distribution_matching/causvid.py @@ -21,12 +21,15 @@ class CausVidModel(DMD2Model): """CausVid implementation""" def _generate_noise_and_time( - self, real_data: torch.Tensor + self, real_data: torch.Tensor, iteration: Optional[int] = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Generate random noises and time step Args: real_data: Real data tensor of shape [B, C, T, H, W] + iteration: Current training iteration, used only to tell a generator + update from a fake-score one so `fake_score_sample_t_cfg` can + apply on the latter. `None` always uses `sample_t_cfg`. Returns: noisy_real_data: Random noise used by the student @@ -55,11 +58,7 @@ def _generate_noise_and_time( t_inhom_expanded = t_inhom[:, None, :, None, None] # shape [B, 1, T, 1, 1] noisy_real_data = self.net.noise_scheduler.forward_process(real_data, eps_inhom, t_inhom_expanded) - t = self.net.noise_scheduler.sample_t( - batch_size, - **basic_utils.convert_cfg_to_dict(self.config.sample_t_cfg), - device=self.device, - ) + t = self._sample_noising_time(batch_size, iteration) eps = torch.randn_like(eps_inhom, device=self.device, dtype=real_data.dtype) return noisy_real_data, t_inhom, t, eps diff --git a/fastgen/methods/distribution_matching/dmd2.py b/fastgen/methods/distribution_matching/dmd2.py index b106e61..ed31a8b 100644 --- a/fastgen/methods/distribution_matching/dmd2.py +++ b/fastgen/methods/distribution_matching/dmd2.py @@ -76,14 +76,29 @@ def _setup_grad_requirements(self, iteration: int) -> None: if self.config.gan_loss_weight_gen > 0: self.discriminator.train().requires_grad_(True) + def _sample_noising_time(self, batch_size: int, iteration: Optional[int] = None) -> torch.Tensor: + """Draw the noising time `t` for whichever branch this iteration runs. + + The generator and fake-score updates alternate and need not noise from + the same density. Unset `fake_score_sample_t_cfg` means both use + `sample_t_cfg`. + """ + t_cfg = self.config.sample_t_cfg + is_fake_score_step = iteration is not None and iteration % self.config.student_update_freq != 0 + if is_fake_score_step and getattr(self.config, "fake_score_sample_t_cfg", None) is not None: + t_cfg = self.config.fake_score_sample_t_cfg + return self.net.noise_scheduler.sample_t(batch_size, **convert_cfg_to_dict(t_cfg), device=self.device) + def _generate_noise_and_time( - self, real_data: torch.Tensor + self, real_data: torch.Tensor, iteration: Optional[int] = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Generate random noises and time step Args: - batch_size: Batch size real_data: Real data tensor for dtype/device reference + iteration: Current training iteration, used only to tell a generator + update from a fake-score one so `fake_score_sample_t_cfg` can + apply on the latter. `None` always uses `sample_t_cfg`. Returns: rand_z_max: Random noise used by the student @@ -115,9 +130,7 @@ def _generate_noise_and_time( ) input_student = self.net.noise_scheduler.forward_process(real_data, eps_student, t_student) - t = self.net.noise_scheduler.sample_t( - batch_size, **convert_cfg_to_dict(self.config.sample_t_cfg), device=self.device - ) + t = self._sample_noising_time(batch_size, iteration) eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) return input_student, t_student, t, eps @@ -442,7 +455,7 @@ def single_train_step( self._setup_grad_requirements(iteration) # Generate noise and time steps - input_student, t_student, t, eps = self._generate_noise_and_time(real_data) + input_student, t_student, t, eps = self._generate_noise_and_time(real_data, iteration=iteration) # Choose between student update or fake_score/discriminator update if iteration % self.config.student_update_freq == 0: diff --git a/fastgen/methods/distribution_matching/self_forcing.py b/fastgen/methods/distribution_matching/self_forcing.py index 5d7cd2e..5d1e79c 100644 --- a/fastgen/methods/distribution_matching/self_forcing.py +++ b/fastgen/methods/distribution_matching/self_forcing.py @@ -12,7 +12,6 @@ import fastgen.utils.logging_utils as logger from fastgen.networks.network import CausalFastGenNetwork -from fastgen.utils.basic_utils import convert_cfg_to_dict from fastgen.utils.distributed import is_rank0, world_size if TYPE_CHECKING: @@ -37,13 +36,15 @@ def __init__(self, config: ModelConfig): self.config = config def _generate_noise_and_time( - self, real_data: torch.Tensor + self, real_data: torch.Tensor, iteration: Optional[int] = None ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Generate random noises and time step Args: - batch_size: Batch size real_data: Real data tensor for dtype/device reference + iteration: Current training iteration, used only to tell a generator + update from a fake-score one so `fake_score_sample_t_cfg` can + apply on the latter. `None` always uses `sample_t_cfg`. Returns: input_student: Random noise used by the student @@ -62,9 +63,7 @@ def _generate_noise_and_time( ) input_student = self.net.noise_scheduler.latents(noise=eps_student) - t = self.net.noise_scheduler.sample_t( - batch_size, **convert_cfg_to_dict(self.config.sample_t_cfg), device=self.device - ) + t = self._sample_noising_time(batch_size, iteration) eps = torch.randn_like(real_data, device=self.device, dtype=real_data.dtype) diff --git a/fastgen/networks/Flux/network.py b/fastgen/networks/Flux/network.py index efc04b2..1582c34 100644 --- a/fastgen/networks/Flux/network.py +++ b/fastgen/networks/Flux/network.py @@ -700,7 +700,6 @@ def _calculate_shift( mu = image_seq_len * m + b return mu - @torch.no_grad() def sample( self, noise: torch.Tensor, diff --git a/fastgen/networks/QwenImage/network.py b/fastgen/networks/QwenImage/network.py index b699509..0998424 100644 --- a/fastgen/networks/QwenImage/network.py +++ b/fastgen/networks/QwenImage/network.py @@ -647,7 +647,6 @@ def forward( return out, logvar return out - @torch.no_grad() def sample( self, noise: torch.Tensor, diff --git a/fastgen/networks/Wan/network.py b/fastgen/networks/Wan/network.py index 5181764..9f1d67b 100644 --- a/fastgen/networks/Wan/network.py +++ b/fastgen/networks/Wan/network.py @@ -31,6 +31,7 @@ from transformers import AutoTokenizer, UMT5EncoderModel from fastgen.networks.network import FastGenNetwork +from fastgen.networks.Wan.utils import convert_wan_official_state_dict, remap_anyflow_keys from fastgen.networks.noise_schedule import NET_PRED_TYPES from fastgen.utils.basic_utils import prompt_clean, str2bool @@ -274,6 +275,46 @@ 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. + + * ``additive``: project each embedding separately and sum — ``temb + remb``, + ``timestep_proj + r_embedder.time_proj(...)``. + * ``gated``: interpolate first, ``(1-g)·temb + g·remb``, then project the + blend through the condition_embedder's own ``time_proj`` — shared with the + t-only path, so ``Wan.__init__`` drops the unused ``r_embedder.time_proj`` + rather than let it drift during training. + + Returns ``(temb, timestep_proj, r_timestep_proj)``: ``temb`` [B, D] modulates + the output head (``scale_shift_table + temb``), ``timestep_proj`` [B, 6, D] is + the AdaLN modulation every block consumes, and ``r_timestep_proj`` is what the + blocks switch to at ``encoder_depth`` — ``None`` when there is no switch. + """ + fusion = getattr(self.r_embedder, "fusion_mode", "additive") + + if fusion == "gated": + gate = self.r_embedder.gate_value + rt_emb = (1 - gate) * temb + gate * remb + rt_ts_proj = self.condition_embedder.time_proj(self.condition_embedder.act_fn(rt_emb)) + rt_ts_proj = unflatten_timestep_proj(rt_ts_proj, rs_seq_len) + if self.encoder_depth is None: + return rt_emb, rt_ts_proj, None + 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)) + 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, None + return remb, timestep_proj, r_ts_proj + + def classify_forward_prepare( self, hidden_states: torch.Tensor, @@ -338,14 +379,8 @@ 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 - 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") @@ -557,6 +592,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. @@ -609,13 +646,23 @@ def __init__( self.transformer.time_cond_type = time_cond_type if r_timestep: logger.info(f"Initializing r embedder with {r_embedder_init}") - self.transformer.r_embedder = self.init_embedder(r_embedder_init) + self.transformer.r_embedder = self.init_embedder(r_embedder_init, r_embedder_fusion, r_embedder_gate_value) else: self.transformer.r_embedder = None # core functionality to override forward function and other methods in the transformer self.override_transformer_forward(inner_dim=inner_dim) + # Only the base Wan forward routes the r-embedding through + # _fuse_r_embedding (installed just above); the causal / VACE variants + # install their own (additive-only) classify_forward_prepare, where + # "gated" would be silently ignored. + if r_timestep and r_embedder_fusion == "gated" and not hasattr(self.transformer, "_fuse_r_embedding"): + raise ValueError( + f"r_embedder_fusion='gated' is not supported by {type(self).__name__} " + "(its transformer forward does not implement gated t/r fusion)." + ) + # Use lazy initialization as this wont work in a meta context when doing FSDP2. self._unipc_scheduler = None @@ -791,17 +838,44 @@ def get_model_id(cls, model_path_or_id: str | os.PathLike) -> str: name = model_path_or_id[idx_start:].split("/")[0] return f"Wan-AI/{name}" - def init_embedder(self, embedder_init: str) -> None: + def init_embedder( + self, + embedder_init: str, + r_embedder_fusion: str = "additive", + r_embedder_gate_value: float = 0.25, + ) -> torch.nn.Module: + """Build the r-embedder and stash its fusion config on the module. + + The fusion settings live on the module (not on `self`) so the + method-bound forward override can branch on them without changing its + signature. They are set before the meta-device early return so both + paths carry them. + + For gated fusion the r_embedder's own `time_proj` / `act_fn` are dropped + here (the fused embedding goes through the condition_embedder's shared + projection instead). Whether the transformer's forward actually supports + gated fusion can only be checked after `override_transformer_forward` + installs `_fuse_r_embedding`, so that check lives in `__init__`. + """ embedder = copy.deepcopy(self.transformer.condition_embedder) del embedder.text_embedder + if r_embedder_fusion not in ("additive", "gated"): + raise ValueError(f"r_embedder_fusion must be 'additive' or 'gated', got {r_embedder_fusion!r}") + embedder.fusion_mode = r_embedder_fusion + # Plain float (not a buffer) so FSDP's reset_parameters does not need to + # re-materialize it after meta-device init. + embedder.gate_value = float(r_embedder_gate_value) + logger.info( + f"r_embedder fusion={r_embedder_fusion}" + + (f" gate_value={r_embedder_gate_value}" if r_embedder_fusion == "gated" else "") + ) + # Skip initialization if using meta device (weights will be broadcast via FSDP) if self._is_in_meta_context(): logger.info("Skipping r_embedder initialization on meta device (will receive weights via FSDP sync)") - return embedder - # zero init the r_embedder - if embedder_init == "zero": + elif embedder_init == "zero": for param in embedder.parameters(): param.data.zero_() elif embedder_init == "random": @@ -826,12 +900,22 @@ def init_embedder(self, embedder_init: str) -> None: pass else: raise ValueError(f"Invalid embedder_init: {embedder_init}") + + if r_embedder_fusion == "gated": + # The gated path projects the fused (t, r) embedding through the + # condition_embedder's own time_proj, shared with the t-only path. + # Drop the r_embedder's copy so it cannot linger as an untrained + # dead parameter. Done after the init branches above because + # "random" initializes time_proj. + del embedder.time_proj + del embedder.act_fn return embedder def override_transformer_forward(self, inner_dim: int) -> None: # 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 @@ -990,83 +1074,12 @@ def load_state_dict(self, state_dict: Mapping[str, Any], **kwargs): i.e., {'generator': state_dict}. """ if self._use_wan_official_sinusoidal and not any(k.startswith("transformer.") for k in state_dict.keys()): - # Handle original Wan checkpoint formats - # Pick the source state dict (adjust these keys to your file) - state = None - for k in ["generator", "state_dict", "model", "module", "net", None]: - if k is None: - # fallback: assume loaded object IS the state_dict - if isinstance(state_dict, dict) and all(isinstance(v, torch.Tensor) for v in state_dict.values()): - state = state_dict - break - if isinstance(state_dict, dict) and k in state_dict and isinstance(state_dict[k], dict): - state = state_dict[k] - break - assert state is not None, "Could not find a state_dict in checkpoint." - logger.info(f"Loading original Wan checkpoint formats from key: {k}") - - # Rename mapping as list of tuples (order matters for the norm swap) - rename_mapping = [ - ("time_embedding.0", "condition_embedder.time_embedder.linear_1"), - ("time_embedding.2", "condition_embedder.time_embedder.linear_2"), - ("text_embedding.0", "condition_embedder.text_embedder.linear_1"), - ("text_embedding.2", "condition_embedder.text_embedder.linear_2"), - ("time_projection.1", "condition_embedder.time_proj"), - ("head.modulation", "scale_shift_table"), - ("head.head", "proj_out"), - ("modulation", "scale_shift_table"), - ("ffn.0", "ffn.net.0.proj"), - ("ffn.2", "ffn.net.2"), - # swap norm names: norm1, norm3, norm2 -> norm1, norm2, norm3 - ("norm2", "norm__placeholder"), - ("norm3", "norm2"), - ("norm__placeholder", "norm3"), - # I2V extras - ("img_emb.proj.0", "condition_embedder.image_embedder.norm1"), - ("img_emb.proj.1", "condition_embedder.image_embedder.ff.net.0.proj"), - ("img_emb.proj.3", "condition_embedder.image_embedder.ff.net.2"), - ("img_emb.proj.4", "condition_embedder.image_embedder.norm2"), - ("img_emb.emb_pos", "condition_embedder.image_embedder.pos_embed"), - # attention parts - ("self_attn.q", "attn1.to_q"), - ("self_attn.k", "attn1.to_k"), - ("self_attn.v", "attn1.to_v"), - ("self_attn.o", "attn1.to_out.0"), - ("self_attn.norm_q", "attn1.norm_q"), - ("self_attn.norm_k", "attn1.norm_k"), - ("cross_attn.q", "attn2.to_q"), - ("cross_attn.k", "attn2.to_k"), - ("cross_attn.v", "attn2.to_v"), - ("cross_attn.o", "attn2.to_out.0"), - ("cross_attn.norm_q", "attn2.norm_q"), - ("cross_attn.norm_k", "attn2.norm_k"), - ("attn2.to_k_img", "attn2.add_k_proj"), - ("attn2.to_v_img", "attn2.add_v_proj"), - ("attn2.norm_k_img", "attn2.norm_added_k"), - ] - - # Convert keys - def rename_key(k: str) -> str: - # strip common prefixes if present - for prefix in ["model.", "module.", "transformer."]: - if k.startswith(prefix): - k = k[len(prefix) :] - # apply replacements in the specified order - for old, new in rename_mapping: - if old in k: - k = k.replace(old, new) - return k - - new_state = {} - for k, v in state.items(): - # optional: skip buffer-like positional/freq params if not needed - # if "freqs" in k: - # continue - new_k = rename_key(k) - # Add 'transformer.' prefix since the model expects it - new_k = f"transformer.{new_k}" - new_state[new_k] = v - state_dict = new_state + state_dict = convert_wan_official_state_dict(state_dict) + + # Some third-party checkpoints store the r-pathway as + # condition_embedder.delta_embedder; remap to FastGen's r_embedder + # layout (no-op for all other checkpoints). + state_dict = remap_anyflow_keys(state_dict) return super().load_state_dict(state_dict, **kwargs) @@ -1117,6 +1130,7 @@ 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 + 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/fastgen/networks/Wan/utils.py b/fastgen/networks/Wan/utils.py new file mode 100644 index 0000000..9b09141 --- /dev/null +++ b/fastgen/networks/Wan/utils.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Checkpoint-format conversions for the Wan networks. + +FastGen's Wan wraps a diffusers ``WanTransformer3DModel`` and expects keys under a +``transformer.`` prefix. Third-party checkpoints come in other layouts, and the +converters here normalise them. They are applied in ``Wan.load_state_dict``: + +* ``convert_wan_official_state_dict`` — the original (non-diffusers) Wan release + layout, e.g. ``self_attn.q`` -> ``attn1.to_q``. +* ``remap_anyflow_keys`` — the AnyFlow HF releases, which store the flow-map + r-pathway inside ``condition_embedder`` as ``delta_embedder``. + +Each converter is a no-op when its marker keys are absent, so they are safe to call +unconditionally and in sequence. +""" + +from typing import Any, Dict, List, Mapping, Tuple + +import torch + +import fastgen.utils.logging_utils as logger + + +# Rename mapping for the original Wan checkpoint layout, as an ordered list of +# (old, new) substring pairs — order matters for the norm swap below. +WAN_OFFICIAL_RENAME_MAPPING: List[Tuple[str, str]] = [ + ("time_embedding.0", "condition_embedder.time_embedder.linear_1"), + ("time_embedding.2", "condition_embedder.time_embedder.linear_2"), + ("text_embedding.0", "condition_embedder.text_embedder.linear_1"), + ("text_embedding.2", "condition_embedder.text_embedder.linear_2"), + ("time_projection.1", "condition_embedder.time_proj"), + ("head.modulation", "scale_shift_table"), + ("head.head", "proj_out"), + ("modulation", "scale_shift_table"), + ("ffn.0", "ffn.net.0.proj"), + ("ffn.2", "ffn.net.2"), + # swap norm names: norm1, norm3, norm2 -> norm1, norm2, norm3 + ("norm2", "norm__placeholder"), + ("norm3", "norm2"), + ("norm__placeholder", "norm3"), + # I2V extras + ("img_emb.proj.0", "condition_embedder.image_embedder.norm1"), + ("img_emb.proj.1", "condition_embedder.image_embedder.ff.net.0.proj"), + ("img_emb.proj.3", "condition_embedder.image_embedder.ff.net.2"), + ("img_emb.proj.4", "condition_embedder.image_embedder.norm2"), + ("img_emb.emb_pos", "condition_embedder.image_embedder.pos_embed"), + # attention parts + ("self_attn.q", "attn1.to_q"), + ("self_attn.k", "attn1.to_k"), + ("self_attn.v", "attn1.to_v"), + ("self_attn.o", "attn1.to_out.0"), + ("self_attn.norm_q", "attn1.norm_q"), + ("self_attn.norm_k", "attn1.norm_k"), + ("cross_attn.q", "attn2.to_q"), + ("cross_attn.k", "attn2.to_k"), + ("cross_attn.v", "attn2.to_v"), + ("cross_attn.o", "attn2.to_out.0"), + ("cross_attn.norm_q", "attn2.norm_q"), + ("cross_attn.norm_k", "attn2.norm_k"), + ("attn2.to_k_img", "attn2.add_k_proj"), + ("attn2.to_v_img", "attn2.add_v_proj"), + ("attn2.norm_k_img", "attn2.norm_added_k"), +] + +# Keys a checkpoint may nest the model state under. +_NESTED_STATE_KEYS = ["generator", "state_dict", "model", "module", "net"] + +# Prefixes stripped before applying the rename mapping. +_STRIPPED_PREFIXES = ["model.", "module.", "transformer."] + + +def _unwrap_state_dict(state_dict: Mapping[str, Any]) -> Tuple[Dict[str, Any], str | None]: + """Return the tensor state dict, unwrapping one level of nesting if present.""" + for key in _NESTED_STATE_KEYS: + if isinstance(state_dict, dict) and key in state_dict and isinstance(state_dict[key], dict): + return state_dict[key], key + # fallback: assume the loaded object IS the state dict + if isinstance(state_dict, dict) and all(isinstance(v, torch.Tensor) for v in state_dict.values()): + return state_dict, None + raise ValueError("Could not find a state_dict in checkpoint.") + + +def rename_wan_official_key(key: str) -> str: + """Map one original-Wan parameter name onto the diffusers layout.""" + for prefix in _STRIPPED_PREFIXES: + if key.startswith(prefix): + key = key[len(prefix) :] + for old, new in WAN_OFFICIAL_RENAME_MAPPING: + if old in key: + key = key.replace(old, new) + return key + + +def convert_wan_official_state_dict(state_dict: Mapping[str, Any]) -> Dict[str, Any]: + """Convert an original (non-diffusers) Wan checkpoint to FastGen's layout. + + Unwraps a nested state dict if needed, renames the parameters via + ``WAN_OFFICIAL_RENAME_MAPPING``, and adds the ``transformer.`` prefix the + model expects. + """ + state, nested_key = _unwrap_state_dict(state_dict) + logger.info(f"Loading original Wan checkpoint format from key: {nested_key}") + return {f"transformer.{rename_wan_official_key(k)}": v for k, v in state.items()} + + +def remap_anyflow_keys(state_dict: Mapping[str, Any]) -> Mapping[str, Any]: + """Remap an AnyFlow HF release state_dict to FastGen's Wan layout. + + AnyFlow's ``FAR_Wan_Transformer3DModel`` stores the r-pathway inside the main + ``condition_embedder`` as ``delta_embedder``, and uses ONE shared ``time_proj`` + for both t and (t, r). FastGen exposes the r-pathway as a top-level + ``r_embedder``; in gated mode it keeps only the ``time_embedder`` and reuses + ``condition_embedder.time_proj`` (see ``_fuse_r_embedding`` in + ``Wan/network.py``), so the layouts differ by a rename only. + + The function is a no-op when no ``condition_embedder.delta_embedder.*`` keys are + present, so it's safe to call unconditionally. Keys with or without the + ``transformer.`` module prefix are both handled. + """ + delta_marker = "condition_embedder.delta_embedder." + delta_keys = [k for k in state_dict if delta_marker in k] + if not delta_keys: + return state_dict + + new_sd = dict(state_dict) + for k in delta_keys: + # [transformer.]condition_embedder.delta_embedder.linear_1.weight + # -> [transformer.]r_embedder.time_embedder.linear_1.weight + prefix, _, suffix = k.partition(delta_marker) + target = f"{prefix}r_embedder.time_embedder.{suffix}" + if target in new_sd: + raise ValueError( + f"remap_anyflow_keys: rewriting {k!r} would overwrite the existing {target!r}. " + "This checkpoint carries both the AnyFlow and the FastGen r-pathway layouts; " + "drop one of them before loading." + ) + new_sd[target] = new_sd.pop(k) + logger.info(f"remap_anyflow_keys: rewrote {len(delta_keys)} delta_embedder tensors into r_embedder.") + return new_sd diff --git a/fastgen/networks/noise_schedule.py b/fastgen/networks/noise_schedule.py index eefb519..d629323 100644 --- a/fastgen/networks/noise_schedule.py +++ b/fastgen/networks/noise_schedule.py @@ -20,6 +20,22 @@ NET_PRED_TYPES = {"x0", "eps", "v", "flow"} +def time_shift(t: torch.Tensor, shift: float = 1.0) -> torch.Tensor: + """Map timesteps through ``t * shift / (t * (shift - 1) + 1)``. + + Args: + t: Timesteps in [0, 1]. + shift: Shift of the map. 1.0 is the identity. + + Returns: + torch.Tensor: The shifted timesteps. + """ + assert shift >= 1, f"shift must be >= 1, got {shift}" + if shift == 1.0: + return t + return t * shift / (t * (shift - 1) + 1) + + class BaseNoiseSchedule(torch.nn.Module): """Abstract base noise schedule class. @@ -1317,8 +1333,8 @@ def __init__( **kwargs, ): super().__init__(min_t, max_t, num_steps, **kwargs) - self._supported_time_dist_types = self._supported_time_dist_types + ("shifted",) - assert 0 <= min_t < max_t <= 0.999, "RF min_t and max_t must be between 0 and 0.999" + self._supported_time_dist_types = self._supported_time_dist_types + ("shifted", "shifted_logitnormal") + assert 0 <= min_t < max_t <= 1.0, "RF min_t and max_t must be between 0 and 1" self._sigmas = torch.linspace(min_t, max_t, num_steps, dtype=self.t_precision) def _rescale_t(self, t: torch.Tensor) -> torch.Tensor: @@ -1326,9 +1342,7 @@ def _rescale_t(self, t: torch.Tensor) -> torch.Tensor: @property def max_sigma(self) -> float: - t_max_scale = int(self.num_steps * self.max_t) - assert 0 <= t_max_scale < len(self._sigmas) - return self._sigmas[t_max_scale].item() + return self._max_t @property def sigmas(self) -> torch.Tensor: @@ -1412,14 +1426,20 @@ def sample_t( elif time_dist_type == "uniform": t = torch.rand(n, device=target_device, dtype=self.t_precision) * (max_t - min_t) + min_t elif time_dist_type == "shifted": - shift = kwargs.get("shift", 5.0) - assert shift >= 1, f"shift must be >= 1, got {shift}" t = torch.rand(n, device=target_device, dtype=self.t_precision) * (max_t - min_t) + min_t - t = t * shift / (t * (shift - 1) + 1) + t = time_shift(t, kwargs.get("shift", 5.0)) + elif time_dist_type == "shifted_logitnormal": + # Logit-normal base density under the same shift map as "shifted" + t = ( + torch.sigmoid(torch.randn(n, device=target_device, dtype=self.t_precision) * train_p_std + train_p_mean) + * (max_t - min_t) + + min_t + ) + t = time_shift(t, kwargs.get("shift", 5.0)) else: raise ValueError( f"Unsupported time distribution type: {time_dist_type} in RFNoiseSchedule." - f"Currently only supports logitnormal, uniform, and shifted." + f"Currently only supports logitnormal, uniform, shifted, and shifted_logitnormal." ) return self.safe_clamp(t, min_t, max_t) diff --git a/fastgen/utils/basic_utils.py b/fastgen/utils/basic_utils.py index bb6175b..afbace2 100644 --- a/fastgen/utils/basic_utils.py +++ b/fastgen/utils/basic_utils.py @@ -87,44 +87,65 @@ def to_str(obj: Any) -> str | Dict[Any, str]: @contextmanager -def inference_mode(*modules: torch.nn.Module, precision_amp: torch.dtype | None = None, device_type: str = "cuda"): +def train_mode(*modules: torch.nn.Module, mode: bool = True): """ - Wraps torch.inference_mode() and temporarily sets the provided modules - to .eval() mode. If precision_amp is not None, it also wraps the context in torch.autocast(). + Temporarily sets the provided modules to train (mode=True) or eval (mode=False) mode. + + Use mode=False to turn off training-only behavior such as dropout for a single forward + pass, e.g. when querying the trained network as its own teacher. Args: - *modules: Modules to set temporarily to eval mode. - precision_amp: If not None, wraps the context in torch.autocast(). - device_type: Device type to use for autocast. + *modules: Modules to set temporarily to the given mode. + mode: Passed to torch.nn.Module.train(); False selects eval mode. Returns: Generator that yields the context manager. Upon exit, it restores the original .training state of each module. """ - # 1. Capture the original training state of each module - # (True if in train mode, False if in eval mode) + # Capture the original training state of each module + # (True if in train mode, False if in eval mode) modules = [mod for mod in modules if isinstance(mod, torch.nn.Module)] previous_states = [mod.training for mod in modules] try: - # 2. Set all specific modules to eval mode - # This is crucial for layers like Dropout and BatchNorm + # The mode gates training-only behavior of layers like Dropout and BatchNorm for mod in modules: - mod.eval() - - # 3. Enter strict inference mode (disables gradients, etc.) and autocast if needed - with torch.inference_mode(), torch.autocast( - dtype=precision_amp, device_type=device_type, enabled=precision_amp is not None - ): - yield + mod.train(mode) + yield finally: - # 4. Restore the original state of each module + # Restore the original state of each module for mod, was_training in zip(modules, previous_states): mod.train(was_training) +@contextmanager +def inference_mode(*modules: torch.nn.Module, precision_amp: torch.dtype | None = None, device_type: str = "cuda"): + """ + Wraps torch.inference_mode() and temporarily sets the provided modules to .eval() mode. + + The context always sets the autocast state: with precision_amp it autocasts to that + dtype, and with precision_amp=None it explicitly *disables* autocast, so an enclosing + autocast region does not carry over into this context. + + Args: + *modules: Modules to set temporarily to eval mode. + precision_amp: Dtype to autocast to. If None, autocast is disabled in the context. + device_type: Device type to use for autocast. + + Returns: + Generator that yields the context manager. + + Upon exit, it restores the original .training state of each module. + """ + # Set eval mode, enter strict inference mode (disables gradients, etc.) and set the autocast state + with train_mode(*modules, mode=False), torch.inference_mode(), torch.autocast( + dtype=precision_amp, device_type=device_type, enabled=precision_amp is not None + ): + yield + + def set_random_seed( seed: int, iteration: int = 0, by_rank: bool = False, devices: List[torch.device | str | int] | None = None ) -> int: diff --git a/tests/test_anyflowmodel.py b/tests/test_anyflowmodel.py new file mode 100644 index 0000000..b3f0241 --- /dev/null +++ b/tests/test_anyflowmodel.py @@ -0,0 +1,705 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""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 ``MeanFlowModel`` directly. The on-policy tests +exercise ``AnyFlowModel`` (DMD2 with a multi-step rollout-with-gradient +student). Both run on the tiny EDM backbone with ``r_timestep=True`` and +``schedule_type=rf`` so they execute on CPU without pretrained weights. +""" + +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 as AnyFlowModelConfig +from fastgen.configs.methods.config_mean_flow import ModelConfig as MeanFlowModelConfig +from fastgen.methods import AnyFlowModel, MeanFlowModel +from fastgen.networks.noise_schedule import time_shift +from fastgen.utils.test_utils import check_grad_zero + + +def _build_pretrain_model( + weight_type="beta08", consistency_ratio=0.25, flow_matching_ratio=0.5, deterministic_buckets=True +): + """MeanFlow configured the AnyFlow way (paper Stage 1).""" + gc.collect() + instance = MeanFlowModelConfig() + + 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 + + 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.flow_matching_ratio = flow_matching_ratio + instance.sample_t_cfg.consistency_ratio = consistency_ratio + instance.sample_t_cfg.deterministic_buckets = deterministic_buckets + + opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True", "+schedule_type=rf"] + instance.net = override_config_with_opts(instance.net, opts) + 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 = MeanFlowModel(instance) + model.on_train_begin() + model.init_optimizers() + return model + + +def _build_onpolicy_model(cotrain_time_dist_type="uniform", cotrain_shift=5.0): + """On-policy fixture mirrors test_dmd2model: img_resolution=8 so the + discriminator's 4x4 conv kernels can operate. + + ``cotrain_time_dist_type`` defaults to an unshifted density, so ``shift`` stays + inert; pass a shifted one to exercise the shifted rollout grid. + """ + gc.collect() + instance = AnyFlowModelConfig() + + base_opts = ["-", "img_resolution=8", "channel_mult=[1]", "channel_mult_noise=1", "+schedule_type=rf"] + # Teacher / fake score are PLAIN single-timestep nets, as in the reference + # and in config_anyflow_onpolicy.py; only the student carries the r pathway. + instance.teacher = override_config_with_opts(AnyFlowModelConfig().net, list(base_opts)) + opts = base_opts[:1] + ["r_timestep=True"] + base_opts[1:] + instance.net = override_config_with_opts(instance.net, opts) + opts_disc = ["-", "feature_indices=[0]", "all_res=[8]", "in_channels=128"] + instance.discriminator = override_config_with_opts(instance.discriminator, opts_disc) + + 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.student_sample_steps = 2 + instance.input_shape = [3, 8, 8] + + # Co-trained flow-map loss settings (MeanFlow machinery on the tiny net). These + # live on their own config: `sample_t_cfg` stays DMD2's noising time. + instance.cotrain_sample_t_cfg.time_dist_type = cotrain_time_dist_type + instance.cotrain_sample_t_cfg.shift = cotrain_shift + instance.cotrain_sample_t_cfg.min_t = 0.001 + instance.cotrain_sample_t_cfg.max_t = 0.999 + instance.cotrain_sample_t_cfg.flow_matching_ratio = 0.5 + instance.cotrain_sample_t_cfg.consistency_ratio = 0.25 + instance.cotrain_sample_t_cfg.deterministic_buckets = True + instance.loss_config.loss_type = "l2" + instance.loss_config.weight_type = "uniform" + instance.loss_config.norm_method = None + 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() + return model + + +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 { + "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 — MeanFlow with AnyFlow options +# --------------------------------------------------------------------------- + + +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 "mf_loss" in loss_map + assert torch.isfinite(loss_map["total_loss"]).all() + assert "gen_rand" in outputs + + +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_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 + + 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) + 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_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, flow_matching_ratio=0.0) + _t, r, r_eq_t_mask = model._sample_t_r_buckets(4) + assert not r_eq_t_mask.any() + assert torch.allclose(r.float(), torch.zeros_like(r.float())) + + +def test_pretrain_bucket_partition_is_deterministic(): + """With deterministic_buckets, both buckets follow the reference's global + partition: flow-matching head, consistency middle, random-pair tail.""" + model = _build_pretrain_model(consistency_ratio=0.25, flow_matching_ratio=0.5, deterministic_buckets=True) + batch_size = 8 + t, r, r_eq_t_mask = model._sample_t_r_buckets(batch_size) + + n_flow_matching = round(0.5 * batch_size) + n_consistency = round(0.25 * batch_size) + 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_flow_matching : n_flow_matching + n_consistency].float(), + torch.zeros(n_consistency, device=r.device), + ) + + +def test_pretrain_bucket_partition_is_stochastic(): + """Without deterministic_buckets the same two ratios drive binomial bucket + SIZES: the layout is still head/middle/tail, but the sizes vary per call + and average to the configured fractions.""" + model = _build_pretrain_model(consistency_ratio=0.25, flow_matching_ratio=0.5, deterministic_buckets=False) + torch.manual_seed(0) + batch_size = 4096 + t, r, r_eq_t_mask = model._sample_t_r_buckets(batch_size) + + # min_t = 0.001, so only the consistency bucket can carry r == 0. + is_consistency = r == 0 + assert not (r_eq_t_mask & is_consistency).any() + assert torch.equal(r[r_eq_t_mask], t[r_eq_t_mask]) + assert abs(r_eq_t_mask.float().mean().item() - 0.5) < 0.05 + assert abs(is_consistency.float().mean().item() - 0.25) < 0.05 + # Prefix layout, same as the deterministic policy. + n_flow_matching = int(r_eq_t_mask.sum()) + assert r_eq_t_mask[:n_flow_matching].all() and not r_eq_t_mask[n_flow_matching:].any() + + # The sizes are what varies, unlike the deterministic policy. + sizes = {int(model._sample_t_r_buckets(64)[2].sum()) for _ in range(20)} + assert len(sizes) > 1, f"stochastic bucket sizes should vary, got {sizes}" + + +def test_pretrain_rebalance_to_flow_matching(): + """With rebalancing on, each r < t loss is rescaled to the flow-matching + loss mean by a detached per-sample factor.""" + model = _build_pretrain_model() + model.loss_config.rebalance_to_flow_matching = True + + mf_loss = torch.tensor([2.0, 4.0, 10.0, 100.0], dtype=torch.float64, requires_grad=True) + r_eq_t_mask = torch.tensor([True, True, False, False]) + loss = model._reduce_mf_loss(mf_loss, r_eq_t_mask) + + # flow-matching mean = 3.0; each r < t sample becomes ~3.0. + expected = (2.0 + 4.0 + 3.0 * (10.0 / 10.00001) + 3.0 * (100.0 / 100.00001)) / 4.0 + assert abs(loss.item() - expected) < 1e-3 + loss.backward() + 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_flow_matching = 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.""" + model = _build_pretrain_model() + model.config.guidance_fuse_scale = 3.0 + model.config.cond_dropout_prob = 0.5 + data = _make_data(model, batch_size=2) + + 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 + + +def test_pretrain_guidance_fusion_requires_neg_condition(): + """Fusion queries the unconditional branch, so a missing neg_condition must + stop with a configuration error rather than reach the network as None.""" + model = _build_pretrain_model() + model.config.guidance_fuse_scale = 3.0 + data = _make_data(model) + data["neg_condition"] = None + + with pytest.raises(AssertionError, match="requires neg_condition"): + model.single_train_step(data, 0) + + +def test_shifted_variants_share_the_shift_map(): + """Both shifted time distributions apply the same shift map when sampling t, + so the grids rebuilt outside the sampler must pick up `shift` for either one. + + Covers the loss-weight normalization (pretrain) and the rollout schedule + (on-policy); a variant missing from the lookup degrades silently to shift=1. + """ + model = _build_pretrain_model() + shifted_scale = model._timestep_weight_scale + cfg_t, cfg_r = model.config.sample_t_cfg, model.config.sample_r_cfg + + cfg_t.time_dist_type = "shifted_logitnormal" + model._init_flow_map_loss(cfg_t, cfg_r) + assert model._timestep_weight_scale == pytest.approx(shifted_scale) + + # and the shift is what makes it differ from an unshifted grid + cfg_t.time_dist_type = "uniform" + model._init_flow_map_loss(cfg_t, cfg_r) + assert model._timestep_weight_scale != pytest.approx(shifted_scale) + + # Same lookup on the on-policy rollout schedule. The shift is resolved once at + # init, so each variant needs a re-init rather than a bare config mutation. + onpolicy = _build_onpolicy_model(cotrain_time_dist_type="shifted", cotrain_shift=5.0) + cot_t, cot_r = onpolicy.config.cotrain_sample_t_cfg, onpolicy.config.cotrain_sample_r_cfg + + # `rollout_t_list` is static, so read `flow_map_shift` back off the model: that is + # what checks the co-train density actually reaches the rollout grid. + def rollout_grid(): + return onpolicy.rollout_t_list(4, onpolicy.flow_map_shift, float(onpolicy.net.noise_scheduler.max_t)) + + shifted_grid = rollout_grid() + + cot_t.time_dist_type = "shifted_logitnormal" + onpolicy._init_flow_map_loss(cot_t, cot_r) + assert torch.allclose(rollout_grid(), shifted_grid) + + cot_t.time_dist_type = "uniform" + onpolicy._init_flow_map_loss(cot_t, cot_r) + assert not torch.allclose(rollout_grid(), shifted_grid) + + +def test_onpolicy_shifted_rollout_grid_and_step(): + """A shifted co-train density must move the rollout grid, and a student update + must run on it. + + Every other on-policy test uses the unshifted fixture, so without this the shift + never reaches `rollout_t_list` or a real rollout -- the shipped + `config_anyflow_onpolicy.py` runs shift=5. + """ + model = _build_onpolicy_model(cotrain_time_dist_type="shifted", cotrain_shift=5.0) + assert model.flow_map_shift == 5.0 + + max_t = float(model.net.noise_scheduler.max_t) + grid = torch.linspace(1.0, 0.0, 5, dtype=torch.float64) + expected = time_shift(grid, 5.0).clamp(max=max_t) + t_list = model.rollout_t_list(4, model.flow_map_shift, max_t).double() + + assert torch.allclose(t_list, expected) + # endpoints are fixed by the map; the interior is pushed towards max_t + assert t_list[0].item() == pytest.approx(min(1.0, max_t)) and t_list[-1].item() == 0.0 + assert (t_list[1:-1] > grid.clamp(max=max_t)[1:-1]).all() + + # and the whole student update runs on that grid + data = _make_data(model, img_resolution=8) + loss_map, _ = model.single_train_step(data, 0) + assert torch.isfinite(loss_map["total_loss"]).all() + assert "vsd_loss" in loss_map and "bidirection_loss" in loss_map + + +def test_onpolicy_dmd_and_cotrain_noising_times_are_separate(): + """DMD2 draws its noising time from `sample_t_cfg`; the co-trained flow-map loss + draws (t, r) from `cotrain_sample_t_cfg`. + + The reference keeps these on two schedulers (`dmd_scheduler` vs `scheduler`), so + give the two configs disjoint ranges and check neither path reads the other's. + """ + model = _build_onpolicy_model() + dmd = model.config.sample_t_cfg + dmd.time_dist_type, dmd.min_t, dmd.max_t = "uniform", 0.70, 0.80 + cot = model.config.cotrain_sample_t_cfg + cot.time_dist_type, cot.min_t, cot.max_t = "uniform", 0.10, 0.20 + model._init_flow_map_loss(cot, model.config.cotrain_sample_r_cfg) + + # the mixin binds the co-train config, never DMD2's + assert model.flow_map_sample_t_cfg is cot + assert model.flow_map_sample_t_cfg is not dmd + + t_dmd = model._sample_noising_time(64, iteration=0) + assert ((t_dmd >= 0.70) & (t_dmd <= 0.80)).all(), t_dmd + + t_mf, r_mf, _ = model._sample_t_r_buckets(64) + assert ((t_mf >= 0.10) & (t_mf <= 0.20)).all(), t_mf + # r shares t's density unless sample_r_cfg is enabled; the consistency bucket pins 0 + assert ((r_mf >= 0.0) & (r_mf <= 0.20)).all(), r_mf + + +@pytest.mark.parametrize("weight_type", ["beta08", "gaussian", "uniform"]) +def test_timestep_weight_function(weight_type): + """The fixed per-timestep weight is a direct function of t: non-negative, + finite, and normalized like the reference scheduler.""" + model = _build_pretrain_model(weight_type=weight_type) + + # The fixture sets norm_method=None, so _compute_weight on a ones tensor is + # exactly the fixed per-timestep weight w(t). + t = torch.linspace(0.0, 1.0, 101, dtype=torch.float64) + w = model._compute_weight(torch.ones_like(t), t) + assert torch.all(w >= 0), f"{weight_type} weights must be non-negative" + assert torch.isfinite(w).all() + + if weight_type == "uniform": + # Uniform normalizes to exactly 1.0 over the reference grid. + assert torch.allclose(w, torch.ones_like(w)) + + # The weight has mean one over the network's discrete training timesteps + # (t=0 excluded), which at the default num_steps=1000 is the reference's + # set_timesteps grid: sum over the shifted grid == num_steps. + num_steps = model.net.noise_scheduler.num_steps + shift = model.flow_map_sample_t_cfg.shift + grid = torch.linspace(1.0, 0.0, num_steps + 1, dtype=torch.float64)[:-1] + grid = shift * grid / (1 + (shift - 1) * grid) + assert abs(model._compute_weight(torch.ones_like(grid), grid).sum().item() - num_steps) < 1e-6 + + +# --------------------------------------------------------------------------- +# On-policy stage — DMD2 with the rollout-with-gradient student +# --------------------------------------------------------------------------- + + +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 + # The co-trained Stage-1 flow-map loss is part of every student update. + assert "bidirection_loss" in loss_map + assert torch.isfinite(loss_map["total_loss"]).all() + assert "gen_rand" in outputs + + +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_student_starts_from_pure_noise(): + """The rollout is on-policy: the student starts from pure noise at max_t. + + DMD2's multi-step branch would hand back real data noised to a random entry + of `t_list`, which `gen_data_from_net` would then roll out as if it sat at + `t_list[0]` — both off-policy and a latent/timestep mismatch. + """ + model = _build_onpolicy_model() + torch.manual_seed(0) + real = torch.randn(8, 3, 8, 8, device=model.device, dtype=model.precision) + input_student, t_student, _, _ = model._generate_noise_and_time(real, iteration=0) + + max_t = float(model.net.noise_scheduler.max_t) + assert torch.allclose(t_student.float(), torch.full_like(t_student.float(), max_t)) + # Pure noise carries no signal from the batch it was drawn alongside; the + # off-policy failure mode correlates at ~0.7 (1/sqrt(192) ~ 0.07 by chance). + cos = torch.nn.functional.cosine_similarity(input_student.flatten(1).float(), real.flatten(1).float()) + assert cos.abs().max() < 0.35, f"student input correlates with real data: {cos.tolist()}" + + +def test_onpolicy_rollout_compresses_to_three_forwards(): + """Regardless of the sampled NFE, the rollout must run at most three + network forwards (jump -> fine step -> jump), with gradient through all.""" + 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 + try: + gen = model.gen_data_from_net(input_student, t_student, condition=cond) + finally: + model.net.forward = orig_forward + + assert len(calls) <= 3, f"rollout must compress to <= 3 forwards, got {len(calls)}" + assert gen.requires_grad + + +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_rollout_propagates_gradient(): + """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) + 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) + 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" + 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) + 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) + + +# --------------------------------------------------------------------------- +# Wan r-embedder fusion + AnyFlow checkpoint remap (pure functions) +# --------------------------------------------------------------------------- + + +def _make_fake_fusion_self(fusion_mode, gate_value=0.25, encoder_depth=None, dim=4, proj_dim=12): + condition_embedder = torch.nn.Module() + condition_embedder.time_proj = torch.nn.Linear(dim, proj_dim) + condition_embedder.act_fn = torch.nn.SiLU() + + r_embedder = torch.nn.Module() + r_embedder.fusion_mode = fusion_mode + r_embedder.gate_value = gate_value + if fusion_mode == "additive": + # Gated fusion reuses condition_embedder.time_proj, so Wan.__init__ + # drops these from the r_embedder in that mode. + r_embedder.time_proj = torch.nn.Linear(dim, proj_dim) + r_embedder.act_fn = torch.nn.SiLU() + return types.SimpleNamespace( + r_embedder=r_embedder, condition_embedder=condition_embedder, encoder_depth=encoder_depth + ) + + +def test_fuse_r_embedding_gated(): + 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) + + 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 + # The gated projection goes through the SHARED condition_embedder.time_proj. + expected = fake.condition_embedder.time_proj(fake.condition_embedder.act_fn(rt_emb)).unflatten(1, (6, -1)) + assert torch.allclose(out_temb, rt_emb) + assert torch.allclose(out_proj, expected) + assert out_r_proj is None + assert not hasattr(fake.r_embedder, "time_proj") + + +def test_validation_t_list_matches_shift(): + """The validation schedule must equal the student's shifted sampling grid. + + `config_anyflow_onpolicy.py` derives `sample_t_cfg.t_list` from + `cotrain_sample_t_cfg.shift`, the same shift `rollout_t_list` applies, so this pins + the two together and catches a stale override. + + Left as None, DMD2 hands `generator_fn` the noise scheduler's UNSHIFTED + `linspace(max_t, 0, N+1)`, which is far off-policy for a shift=5 model -- + asserted below so the schedule cannot silently degrade to that. + """ + import fastgen.configs.experiments.WanT2V.config_anyflow_onpolicy as mod + + cfg = mod.create_config() + shift = float(cfg.model.cotrain_sample_t_cfg.shift) + n = int(cfg.model.student_sample_steps) + max_t = float(cfg.model.net.max_t) # the bound `rollout_t_list` clamps to + + grid = [1.0 - i / n for i in range(n + 1)] + expected = [min(shift * x / (1 + (shift - 1) * x), max_t) for x in grid] + + t_list = list(cfg.model.sample_t_cfg.t_list) + assert len(t_list) == n + 1, t_list + assert all(abs(a - b) < 1e-9 for a, b in zip(t_list, expected, strict=True)), (t_list, expected) + assert t_list[0] == max_t and t_list[-1] == 0.0 + + # and it must not be the unshifted fallback DMD2 would use for None + unshifted = [max_t * (1.0 - i / n) for i in range(n + 1)] + assert not all(abs(a - b) < 1e-6 for a, b in zip(t_list, unshifted, strict=True)) + + +def test_fuse_r_embedding_gated_trains_the_shared_time_proj(): + """The gated projection must flow gradient into condition_embedder.time_proj. + + A private r_embedder.time_proj would leave the condition_embedder's copy + without gradient (a dead parameter under DDP/FSDP) while the two silently + drift apart, breaking round-trips to the AnyFlow checkpoint layout. + """ + from fastgen.networks.Wan.network import _fuse_r_embedding + + torch.manual_seed(0) + fake = _make_fake_fusion_self("gated") + _, out_proj, _ = _fuse_r_embedding(fake, torch.randn(2, 4), torch.randn(2, 6, 2), torch.randn(2, 4), None) + out_proj.sum().backward() + + grad = fake.condition_embedder.time_proj.weight.grad + assert grad is not None and torch.any(grad != 0) + + +def test_fuse_r_embedding_gated_respects_encoder_depth(): + 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) + # r is already folded into the two above, and without encoder_depth the + # blocks have nothing to switch to. + assert out_r_proj is None + + +@pytest.mark.parametrize("prefix", ["", "transformer."]) +def test_remap_anyflow_keys(prefix): + from fastgen.networks.Wan.utils 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 + # time_proj stays shared on the condition_embedder — no r_embedder copy. + assert f"{prefix}r_embedder.time_proj.weight" not in out + assert torch.equal( + out[f"{prefix}condition_embedder.time_proj.weight"], sd[f"{prefix}condition_embedder.time_proj.weight"] + ) + # Unrelated keys untouched. + assert torch.equal(out[f"{prefix}blocks.0.attn1.to_q.weight"], sd[f"{prefix}blocks.0.attn1.to_q.weight"]) + + +def test_remap_anyflow_keys_noop_without_delta_keys(): + from fastgen.networks.Wan.utils import remap_anyflow_keys + + sd = {"transformer.blocks.0.attn1.to_q.weight": torch.randn(2, 2)} + assert remap_anyflow_keys(sd) is sd + + +def test_rollout_uses_the_flow_map_sample_loop(): + """`gen_data_from_net` must resolve to FlowMapLossMixin's loop, not the base one. + + `AnyFlowModel(FlowMapLossMixin, DMD2Model)` picks it purely by MRO order. + `FastGenModel._student_sample_loop` is an x0-prediction loop that never + passes `r`, and it accepts the same arguments -- so swapping the base order + would silently turn the flow-map rollout (jump / fine step / jump) into a + plain diffusion sampler instead of raising. + """ + from fastgen.methods import AnyFlowModel + from fastgen.methods.consistency_model.mean_flow import FlowMapLossMixin + + assert ( + AnyFlowModel._student_sample_loop.__func__ is FlowMapLossMixin._student_sample_loop.__func__ + ), "AnyFlowModel must inherit the flow-map sample loop; check the base-class order" diff --git a/tests/test_meanflowmodel.py b/tests/test_meanflowmodel.py index 644ae8c..39446c0 100644 --- a/tests/test_meanflowmodel.py +++ b/tests/test_meanflowmodel.py @@ -49,7 +49,7 @@ def test_single_train_step_update(get_model_data): # Run the training step; cifar10 default config assert model.config.sample_t_cfg.train_p_mean == -0.6 assert model.config.sample_t_cfg.train_p_std == 1.6 - assert model.config.sample_t_cfg.r_sample_ratio == 0.75 + assert model.config.sample_t_cfg.flow_matching_ratio == 0.25 norm_method, *norm_args = model.config.loss_config.norm_method.split("_") assert norm_method == "poly" @@ -120,3 +120,137 @@ def test_single_train_step_update_fp32_jvp(): assert "mf_loss" in loss_map assert "gen_rand" in outputs assert isinstance(outputs["gen_rand"], Callable) + + +@pytest.mark.parametrize("cond_dropout_prob, expect_guided", [(None, True), (0.0, True), (1.0, False)]) +def test_target_side_guidance_applies_when_no_cond_dropout(cond_dropout_prob, expect_guided): + """`guidance_scale` must reach the target for every conditional sample. + + No dropout means every sample is conditional, so it must guide exactly like + `p=0.0`. It used not to: `_drop_condition` returned `keep=None` when + `cond_dropout_prob is None` and the caller skipped the update entirely -- + computing `guided_dxt_dt` at the cost of an extra forward pass and then + discarding it, so `guidance_scale` was silently inert on that path. + (Pre-existing: `_mix_condition` returned early on `cond_dropout_prob is None` + before the AnyFlow work.) `_drop_condition` now always returns a mask -- all-True + here -- so the caller has no special case left to forget. + + The net's forward is stubbed to depend only on the condition, so the guided + velocity is exact rather than initialization-dependent -- EDM's `SongUNet` + zero-inits `out_conv`, which would otherwise make the cond and uncond passes + bitwise identical and guidance a provable no-op. + """ + config = create_config() + instance = config.model + opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"] + instance.net = override_config_with_opts(instance.net, opts) + instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + instance.precision = "float32" + instance.pretrained_model_path = "" + instance.input_shape = [3, 2, 2] + instance.cond_dropout_prob = cond_dropout_prob + instance.guidance_scale = 2.0 + instance.guidance_fuse_scale = None + model = MeanFlowModel(instance) + model.on_train_begin() + + batch_size = 4 + real = torch.randn(batch_size, 3, 2, 2, device=model.device, dtype=torch.float32) + z = torch.randn_like(real) + t = torch.full((batch_size,), 0.5, device=model.device, dtype=model.net.noise_scheduler.t_precision) + condition = torch.nn.functional.one_hot(torch.arange(batch_size) % 10, num_classes=10) + condition = condition.to(model.device, torch.float32) + neg_condition = torch.zeros(batch_size, 10, device=model.device, dtype=torch.float32) + + # Depends ONLY on the condition: the all-zero neg pass returns exactly 0. + def condition_only_forward(x_t, t, **kwargs): + val = kwargs["condition"].sum(dim=1).reshape(-1, 1, 1, 1).to(x_t.dtype) + return torch.ones_like(x_t) * val + + orig_forward = model.net.forward + model.net.forward = condition_only_forward + try: + cond_out, dxt_dt, _ = model._get_velocity(real, z, t, condition=condition, neg_condition=neg_condition) + finally: + model.net.forward = orig_forward + + # neg_dxt_dt == 0, so guided == neg + scale * (plain - neg) == 2 * plain. + plain = model.net.noise_scheduler.cond_velocity(x=real, eps=z, t=t) + dropped = (cond_out == neg_condition).all(dim=1) + assert bool(dropped.all()) is not expect_guided + + expected = 2.0 * plain if expect_guided else plain + assert torch.allclose(dxt_dt, expected), (dxt_dt - expected).abs().max() + + +def test_fused_jvp_scaling_is_gated_on_kept_samples(): + """Under `guidance_fuse_scale`, dF/dt must be divided by g only for samples that + stayed conditional. + + A dropped sample's fused prediction collapses to plain `u_uncond` (the fusion + self-cancels), so scaling its derivative too would regress it onto + `v - (t - r) * d(u_uncond)/dt / g` instead of the unconditional MeanFlow identity. + The AnyFlow reference scales the whole batch + (`compute_central_difference(..., guidance)` in + `far/trainers/trainer_wan_anyflow_pretrain.py`); we gate on `keep`. + + The net is stubbed to ignore `condition`, so dropping it changes nothing about the + raw derivative -- making the 1/g gate the ONLY difference between the two runs. + """ + g = 3.0 + config = create_config() + instance = config.model + opts = ["-", "img_resolution=2", "channel_mult=[1]", "channel_mult_noise=1", "r_timestep=True"] + instance.net = override_config_with_opts(instance.net, opts) + instance.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + instance.precision = "float32" + instance.pretrained_model_path = "" + instance.input_shape = [3, 2, 2] + instance.guidance_fuse_scale = g + instance.loss_config.use_jvp_finite_diff = True + instance.sample_t_cfg.deterministic_buckets = False + model = MeanFlowModel(instance) + model.on_train_begin() + + batch_size = 4 + real = torch.randn(batch_size, 3, 2, 2, device=model.device, dtype=torch.float32) + t_prec = model.net.noise_scheduler.t_precision + t = torch.full((batch_size,), 0.6, device=model.device, dtype=t_prec) + r = torch.full((batch_size,), 0.2, device=model.device, dtype=t_prec) + condition = torch.nn.functional.one_hot(torch.arange(batch_size) % 10, num_classes=10) + condition = condition.to(model.device, torch.float32) + neg_condition = torch.zeros(batch_size, 10, device=model.device, dtype=torch.float32) + + # `+ 0.0 * param.sum()` leaves the value untouched but ties the output to the + # autograd graph: `_mf_pred_to_loss` asserts + # `u_theta.requires_grad is torch.is_grad_enabled()`. The JVP runs under `_jvp`'s + # `@torch.no_grad()`, so `u_theta_jvp` stays grad-free as that code also asserts. + param = next(model.net.parameters()) + + def condition_independent_forward(x_t, t, **kwargs): + tt = t.reshape(-1, *([1] * (x_t.ndim - 1))).to(x_t.dtype) + return x_t * (1.0 + tt) + 0.0 * param.sum().to(x_t.dtype) + + orig_forward = model.net.forward + + def run(cond_dropout_prob): + model.config.cond_dropout_prob = cond_dropout_prob + model.net.forward = condition_independent_forward + try: + torch.manual_seed(11) + return model._compute_mf_loss( + real_data=real, + t=t, + r=r, + iteration=0, + condition=condition, + neg_condition=neg_condition, + )[2] + finally: + model.net.forward = orig_forward + + jvp_kept = run(0.0) # every sample conditional -> every derivative divided by g + jvp_dropped = run(1.0) # every sample unconditional -> none divided + + assert torch.allclose(jvp_dropped, g * jvp_kept, atol=1e-5, rtol=1e-4), (jvp_dropped - g * jvp_kept).abs().max() + assert not torch.allclose(jvp_dropped, jvp_kept, atol=1e-6) diff --git a/tests/test_network_fsdp.py b/tests/test_network_fsdp.py index 7dc6e86..bd45f99 100644 --- a/tests/test_network_fsdp.py +++ b/tests/test_network_fsdp.py @@ -692,9 +692,7 @@ def _test_qwen_image_impl(rank: int, world_size: int) -> Dict: from fastgen.configs.net import QwenImageConfig set_env_vars() - return _generic_fsdp_test_impl( - rank, world_size, QwenImageConfig, generate_qwen_image_inputs, apply_checkpointing=True - ) + return _generic_fsdp_test_impl(world_size, QwenImageConfig, generate_qwen_image_inputs, apply_checkpointing=True) # =============================================================================