diff --git a/examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml b/examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml new file mode 100644 index 000000000..2fe7b928b --- /dev/null +++ b/examples/train/configs/distribution_matching/wan/anyflow_onpolicy_t2v.yaml @@ -0,0 +1,105 @@ +# AnyFlow on-policy DMD — Wan 2.1 T2V 1.3B. +# +# Stage 2 of the AnyFlow two-stage recipe. Continues from the pretrain +# checkpoint; refines the student via DMD2 with a multi-step Euler-flow +# rollout from pure noise. Teacher provides the real score, critic +# learns the fake score; both inherited from DMD2Method. +# +# Replace with the output of the pretrain stage, +# or with the NVIDIA-released checkpoint +# nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers to bootstrap directly from +# the paper weights (the delta_embedder rename is handled by the +# param_names_mapping in WanVideoArchConfig). + +models: + student: + _target_: fastvideo.train.models.wan.WanModel + init_from: + trainable: true + teacher: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-14B-Diffusers + trainable: false + disable_custom_init_weights: true + critic: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.distribution_matching.anyflow.AnyFlowMethod + rollout_mode: simulate + generator_update_interval: 5 + real_score_guidance_scale: 3.0 + dmd_denoising_steps: [999, 937, 833, 624] + warp_denoising_step: false + + # AnyFlow rollout knobs. + student_sample_steps: 4 + use_mean_velocity: true + t_list_override: [999.0, 937.0, 833.0, 624.0, 0.0] + dmd_score_r_value: 0.0 # DMD scoring conditioning is at r=0 (consistency target). + + # Critic optimizer (DMD2 inherited). + fake_score_learning_rate: 8.0e-6 + fake_score_betas: [0.0, 0.999] + fake_score_lr_scheduler: constant + + attn_kind: vsa + +training: + distributed: + num_gpus: 8 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 8 + + data: + data_path: data/preprocessed + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 21 + num_height: 480 + num_width: 832 + num_frames: 81 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.0, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 4000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wan2.1_anyflow_onpolicy + training_state_checkpointing_steps: 500 + checkpoints_total_limit: 3 + resume_from_checkpoint: latest + + tracker: + project_name: anyflow-wan + run_name: wan2.1_t2v_anyflow_onpolicy + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + +pipeline: + flow_shift: 5.0 + dit_config: + r_embedder: true + r_embedder_fusion: gated + r_embedder_gate_value: 0.25 + r_embedder_deltatime_type: r diff --git a/examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml b/examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml new file mode 100644 index 000000000..0e870a4a1 --- /dev/null +++ b/examples/train/configs/distribution_matching/wan/anyflow_pretrain_t2v.yaml @@ -0,0 +1,83 @@ +# AnyFlow pretrain (flow-map central-difference) — Wan 2.1 T2V 1.3B. +# +# Stage 1 of the AnyFlow two-stage recipe. Trains the dual-timestep +# u_θ(x_t, t, r) on the central-difference target so the same checkpoint +# can be sampled at arbitrary NFE in the on-policy stage. +# +# Initialize from base Wan 2.1 T2V 1.3B. No teacher or critic at this +# stage; AnyFlowPretrainMethod owns a single student + one optimizer. + +models: + student: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + +method: + _target_: fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod + diffusion_ratio: 0.5 + consistency_ratio: 0.25 + epsilon: 5 # finite-difference step in absolute train-timestep units + weight_type: beta08 # per-timestep loss weight = t * sqrt(1 - t), renormalized + fuse_guidance_scale: 3.0 + # shift is taken from pipeline.flow_shift below. + +training: + distributed: + num_gpus: 8 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 8 + + data: + data_path: data/preprocessed + dataloader_num_workers: 4 + train_batch_size: 4 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 21 + num_height: 480 + num_width: 832 + num_frames: 81 + + optimizer: + learning_rate: 5.0e-5 + betas: [0.9, 0.999] + weight_decay: 0.0 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 6000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wan2.1_anyflow_pretrain + training_state_checkpointing_steps: 500 + checkpoints_total_limit: 3 + resume_from_checkpoint: latest + + tracker: + project_name: anyflow-wan + run_name: wan2.1_t2v_anyflow_pretrain + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback + max_grad_norm: 1.0 + +pipeline: + flow_shift: 5.0 + dit_config: + # Enable AnyFlow dual-timestep conditioning. The student loads from + # base Wan 2.1 — its checkpoint has no delta_embedder weights, so they + # get initialized identically to time_embedder via deep-copy in + # WanTimeTextImageEmbedding.__init__. + r_embedder: true + r_embedder_fusion: gated + r_embedder_gate_value: 0.25 + r_embedder_deltatime_type: r diff --git a/fastvideo/tests/training/distill/test_anyflow_pretrain.py b/fastvideo/tests/training/distill/test_anyflow_pretrain.py index 59259db6b..301d31a1d 100644 --- a/fastvideo/tests/training/distill/test_anyflow_pretrain.py +++ b/fastvideo/tests/training/distill/test_anyflow_pretrain.py @@ -537,3 +537,48 @@ def test_central_difference_dF_dt_rejects_zero_delta() -> None: delta=0.0, num_train_timesteps=1000.0, ) + + +# --------------------------------------------------------------------------- +# Task 9: param_names_mapping handles AnyFlow checkpoints + is a no-op on plain +# Wan checkpoints. We don't ship a separate remap_anyflow_keys helper because +# the existing param_names_mapping regex mechanism does the job (delta_embedder +# rename is a no-op when the source state dict doesn't contain those keys). +# --------------------------------------------------------------------------- + + +def test_param_names_mapping_includes_delta_embedder_rename() -> None: + """The Wan arch config's param_names_mapping must rename HF AnyFlow + delta_embedder weights into FastVideo's internal mlp.fc_in/fc_out layout.""" + arch = WanVideoConfig().arch_config + mapping_keys = list(arch.param_names_mapping.keys()) + assert any("delta_embedder" in k for k in mapping_keys), ( + "WanVideoArchConfig.param_names_mapping must include delta_embedder " + "rename so HF AnyFlow checkpoints load without a separate adapter") + + +def test_param_names_mapping_default_doesnt_break_plain_wan_keys() -> None: + """The new delta_embedder regex must not match any key in a plain + pretrained Wan2.1 checkpoint (those don't have delta_embedder).""" + import re + + plain_wan_keys = [ + "patch_embedding.weight", + "condition_embedder.time_embedder.linear_1.weight", + "condition_embedder.time_embedder.linear_2.weight", + "condition_embedder.time_proj.weight", + "condition_embedder.text_embedder.linear_1.weight", + "blocks.0.attn1.to_q.weight", + "blocks.0.ffn.net.0.proj.weight", + ] + arch = WanVideoConfig().arch_config + delta_regexes = [ + k for k in arch.param_names_mapping if "delta_embedder" in k + ] + assert delta_regexes, "expected at least one delta_embedder regex" + + for plain in plain_wan_keys: + for rx in delta_regexes: + assert re.match(rx, plain) is None, ( + f"delta_embedder regex {rx!r} unexpectedly matched plain " + f"Wan key {plain!r}")