Single-page summary of AnyFlow's two-stage recipe:
- Algorithm intuition (u_θ(x_t, t, r) → any-step Euler)
- Stage 1 pretrain: central-difference target + (t, r) sampling
- Stage 2 on-policy DMD: multi-step rollout with grad-step broadcast
- Launch commands for both YAMLs
- Note on loading nvidia/AnyFlow-Wan2.1-T2V-* checkpoints (handled by
param_names_mapping, no separate adapter)
- Note on fuse_guidance_scale parameterization
Registered under mkdocs nav so it surfaces in the built doc tree.
Mirrors test_distill_dmd.py's subprocess-torchrun pattern but points at
the new YAML entrypoint (fastvideo.train.entrypoint.train). Auto-skipped
when fewer than 2 CUDA devices are visible — the new TrainingMethod
framework's distributed barriers don't tolerate single-GPU bring-up.
On Buildkite the test fires under /test distillation; on GMI we can
run it directly with NUM_GPUS=2 NUM_NODES=1 pytest -v -k anyflow_smoke.
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.
Subclasses DMD2Method and overrides _student_rollout to run
student_sample_steps Euler-flow steps from pure noise. One step is
gradient-enabled, with the chosen index broadcast from rank 0 in
distributed runs so every worker agrees (matching AnyFlow's
WanAnyFlowPipeline.training_rollout).
Method config knobs:
- student_sample_steps (default 4)
- use_mean_velocity (default True; r = t_next vs r = t at each step)
- t_list_override (optional pinned schedule, e.g. AnyFlow paper's
[999, 937, 833, 624, 0] for 4-step Wan2.1)
- dmd_score_r_value (default 0.0; r used by the inherited DMD loss)
- real_score_guidance_scale (inherited from DMD2; default 1.0)
The inherited _dmd_loss / _critic_flow_matching_loss machinery is reused
verbatim — they call predict_x0 / predict_noise on the student, which
go through the single-timestep WanModel forward path (since DMD2 scores
against r=0 implicitly via the inherited code; if r=t is desired this
can be wired in a follow-up by overriding _dmd_loss to call
predict_velocity_with_r).
CPU-only tests via object.__new__ bypassing the full __init__ wire-up;
GPU integration coverage lives in the smoke test (Task 11).
Implements AnyFlow's central-difference loss path end-to-end:
1. WanModel gains predict_velocity_with_r(noisy, t, r, batch, ...) — mirrors
predict_noise's forward_context + autocast plumbing but injects
r_timestep into the transformer kwargs.
2. anyflow_pretrain._central_difference_dF_dt — symmetric finite difference
over (t ± delta) with the sample also moved along the flow trajectory by
v_pred * (delta / num_train_timesteps), mirroring AnyFlow's reference
trainer_wan_anyflow_pretrain.py::compute_central_difference. Wrapped in
torch.no_grad so the two extra forwards stay out of the backward graph.
3. AnyFlowPretrainMethod.single_train_step:
- sample (t, r) ∈ [0, 1] via _sample_pair_timesteps
- apply scheduler shift + scale to absolute units
- rebuild noisy with the flow-map scheduler's add_noise
- run conditional + (optional) unconditional student forwards
- apply guidance distillation when fuse_guidance_scale != 1
- compute target = (eps - x0) - (t - r) * dF/dt
- per-sample MSE * per-timestep weight
- stop-grad scale balance so non-diffusion losses match the diffusion
branch's magnitude
- emit metrics for diffusion/consistency fractions + scale weight mean
4. AnyFlowPretrainMethod.backward overrides TrainingMethod.backward to
route the loss through self.student.backward inside the correct
forward_context, matching DMD2Method's pattern.
Adds the single-student AnyFlow flow-map pretrain method as a
TrainingMethod subclass. The __init__ parses and validates all method
config knobs (diffusion_ratio, consistency_ratio, epsilon, weight_type,
fuse_guidance_scale, shift), builds a FlowMapEulerDiscreteScheduler, and
wires a single optimizer+scheduler over the student.
The (t, r) sampling helper _sample_pair_timesteps lives at module scope
so it can be unit-tested without instantiating the full method. It
matches the AnyFlow paper formulation:
- two uniform draws u1, u2; t = max, r = min
- first diffusion_ratio * B: r := t (plain flow matching)
- next consistency_ratio * B: r := 0 (consistency to clean data)
- remainder: free reconstruction range
single_train_step raises NotImplementedError for now — filled in by the
next commit which adds the central-difference target and the loss
assembly (Task 7).
Standalone any-step Euler scheduler implementing AnyFlow's flow-map step
formula x_r = x_t - ((t - r) / num_train_timesteps) * u(x_t, t, r). Also
provides the AnyFlow training-time helpers apply_shift (flow-matching
shift transform), get_train_weight (per-timestep loss weight with
beta08 = t * sqrt(1-t) renormalized to num_train_timesteps total mass),
and add_noise (linear flow-matching interpolation).
Subclasses BaseScheduler so it slots into the existing FastVideo
scheduler discovery surface (timesteps/order/num_train_timesteps
attributes plus set_shift/set_timesteps/scale_model_input). Has no
diffusers ConfigMixin/SchedulerMixin dependency.
set_timesteps accepts custom_timesteps so configs can pin AnyFlow's
hand-tuned schedules (e.g. [999, 937, 833, 624, 0] for the 4-step
Wan2.1 setting from the paper).
The forward signature now accepts r_timestep as an explicit kwarg (not
silently swallowed by **kwargs), gets flattened in parallel with timestep
when the input is 2-D (Wan 2.2 ti2v case), and is forwarded to the
condition_embedder.
The constructor now propagates the four arch_config flags
(r_embedder/r_embedder_fusion/r_embedder_gate_value/r_embedder_deltatime_type)
into the WanTimeTextImageEmbedding so a YAML can opt into AnyFlow's
dual-timestep conditioning through training.dit_config overrides.
Defaults preserve byte-identity with the legacy single-timestep path —
when r_timestep is omitted the delta_embedder branch is skipped entirely
(test_embedder_enabled_without_r_timestep_is_bit_identical_to_legacy).
When r_embedder=True, allocates a delta_embedder (deep-copied from
time_embedder for matching initialization, per the AnyFlow reference's
setup_flowmap_model() pattern) and registers a non-persistent gate buffer.
Forward now accepts an optional r_timestep:
- additive (default): temb_t + g * delta_emb
- gated: (1 - g) * temb_t + g * delta_emb
Both branches are gated by r_embedder=True AND r_timestep is not None,
so existing call sites that don't pass r_timestep stay byte-identical
(verified by test_embedder_enabled_without_r_timestep_is_bit_identical_to_legacy).
The gate is a non-persistent buffer — it's a hyperparameter, not a
learned weight, and stays out of state_dict so checkpoints remain
portable across gate values.
deltatime_type controls whether the delta_embedder consumes r directly
(default, matching AnyFlow's deltatime_type='r') or (t - r).
Adds four fields to WanVideoArchConfig for AnyFlow dual-timestep conditioning:
- r_embedder: bool — enables the delta_embedder allocation
- r_embedder_fusion: 'additive' (default) or 'gated'
- r_embedder_gate_value: float — gate g in gated fusion
- r_embedder_deltatime_type: 'r' (default) or 't-r'
Also adds two regex entries to param_names_mapping so HF AnyFlow checkpoints
load delta_embedder weights into the FastVideo-internal mlp.fc_in/fc_out layout.
Both regex are no-ops on plain Wan checkpoints.
The defaults preserve bit-identity with every existing Wan-based method
(DMD2, Self-Forcing, KD, DFSFT) — no delta_embedder is allocated and no
extra computation runs on the embedder forward path.
Enderfga