[feat] AnyFlow YAML configs (pretrain + on-policy) for Wan 2.1 T2V 1.3B
Two reference configs slot into examples/train/configs/distribution_matching/wan/: - anyflow_pretrain_t2v.yaml: single-student pretrain stage. Enables r_embedder under pipeline.dit_config, sets diffusion/consistency ratios, epsilon=5, weight_type=beta08, fuse_guidance_scale=3.0, lr 5e-5, 6k steps, Wan2.1 T2V 1.3B init. - anyflow_onpolicy_t2v.yaml: full DMD2 trio (student + Wan2.1 14B teacher + Wan2.1 1.3B critic). Pins the 4-step rollout schedule [999, 937, 833, 624, 0] from the AnyFlow paper, use_mean_velocity=true, generator_update_interval=5, lr 2e-6, 4k steps. The student loads from a path placeholder which can be either the local pretrain output or the published nvidia/AnyFlow-Wan2.1-T2V-1.3B-Diffusers checkpoint — the latter case relies on the param_names_mapping regex in WanVideoArchConfig to rename delta_embedder weights.
This commit is contained in:
Enderfga
committed by
SolitaryThinker
SolitaryThinker
parent
314a018002
commit
e8516c2fa2
@@ -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 <PATH_TO_PRETRAIN_CKPT> 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: <PATH_TO_PRETRAIN_CKPT>
|
||||
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
|
||||
@@ -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
|
||||
@@ -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}")
|
||||
|
||||
Reference in New Issue
Block a user