Training Infrastructure¶

FastVideo's training infrastructure (fastvideo/train/) is a YAML-driven framework for training and distilling video diffusion models. A single config file controls everything — models, algorithms, distributed strategy, checkpointing, and validation — with no code changes needed to mix and match.

Quick Start¶

Launch with the helper script¶

bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml

The script auto-detects available GPUs and sets up torchrun. Override with environment variables:

NUM_GPUS=4 NNODES=2 NODE_RANK=0 \
    MASTER_ADDR=10.0.0.1 MASTER_PORT=29501 \
    bash examples/train/run.sh my_config.yaml

Launch directly with torchrun¶

torchrun --nproc_per_node=8 \
    fastvideo/train/entrypoint/train.py \
    --config examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml

CLI flags¶

Config Format¶

Every run is defined by a single YAML file with five top-level sections. See examples/train/example.yaml for a fully-commented reference.

models — Role-based model instances¶

Each entry defines a model role. The _target_ field specifies the Python class to instantiate:

models:
  student:
    _target_: fastvideo.train.models.wan.WanModel
    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
    trainable: true
  teacher:
    _target_: fastvideo.train.models.wan.WanModel
    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
    trainable: false
    disable_custom_init_weights: true

Common model parameters:

Which roles are needed depends on the training method:

method — Training algorithm¶

Selects and configures the training algorithm:

method:
  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method
  rollout_mode: simulate
  dmd_denoising_steps: [1000, 750, 500, 250]
  generator_update_interval: 5

To switch algorithms, change _target_ and adjust the method-specific keys. See Training Methods for details on each algorithm.

training — Typed infrastructure config¶

This section maps to typed dataclasses with defaults and validation:

training:
  distributed:
    num_gpus: 8
    sp_size: 1            # sequence parallelism
    tp_size: 1            # tensor parallelism
    hsdp_replicate_dim: 1 # HSDP replication dimension
    hsdp_shard_dim: 8     # HSDP sharding dimension

  data:
    data_path: data/my_dataset
    train_batch_size: 1
    dataloader_num_workers: 4
    training_cfg_rate: 0.1  # classifier-free guidance dropout rate
    seed: 1000
    num_latent_t: 20
    num_height: 448
    num_width: 832
    num_frames: 77

  optimizer:
    learning_rate: 2.0e-6
    betas: [0.9, 0.999]
    weight_decay: 0.01
    lr_scheduler: constant  # constant, linear, cosine, polynomial
    lr_warmup_steps: 0

  loop:
    max_train_steps: 4000
    gradient_accumulation_steps: 1

  checkpoint:
    output_dir: outputs/my_run
    training_state_checkpointing_steps: 1000  # 0 = disabled
    checkpoints_total_limit: 3                # 0 = keep all

  tracker:
    project_name: my_project
    run_name: my_run

  model:
    weighting_scheme: uniform   # uniform, logit_normal, mode
    precondition_outputs: false
    enable_gradient_checkpointing_type: full

  vsa:
    sparsity: 0.0         # 0.0 = disabled
    decay_rate: 0.0
    decay_interval_steps: 0

training.data.data_path can also mix multiple preprocessed datasets by using a mapping from dataset path to repeat count:

training:
  data:
    data_path:
      data/zeldam2-clean: 1
      data/multi3d_games: 2

The repeat count duplicates that dataset's parquet file list before shuffling/sampling, so the example above trains with roughly twice as much multi3d_games exposure as zeldam2-clean. Paths are just suggested locations; use any local path that contains a FastVideo preprocessed parquet dataset.

See Training Trackers to configure Weights & Biases or SwanLab, including SwanLab installation and authentication.

callbacks — Pluggable hooks¶

Callbacks run at specific points in the training loop (before/after optimizer steps, at validation time, etc.):

callbacks:
  grad_clip:
    max_grad_norm: 1.0

  ema:
    _target_: fastvideo.train.callbacks.ema.EMACallback
    decay: 0.9999
    start_iter: 0

  validation:
    _target_: fastvideo.train.callbacks.validation.ValidationCallback
    pipeline_target: fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline
    dataset_file: path/to/validation.json
    every_steps: 100
    sampling_steps: [4]
    guidance_scale: 5.0

See Callbacks for details on each callback.

pipeline — Inference pipeline overrides¶

Optional overrides for the inference pipeline used during validation:

pipeline:
  flow_shift: 8

Registered transformer linear-quantization configs can also be selected by name. For example, the LTX-2 NVFP4-QAT recipe applies real FP4 forward GEMMs with a straight-through-estimator backward to its deployment-targeted attention/FFN projections:

pipeline:
  dit_config:
    quant_config: nvfp4_qat_train

The LTX-2 recipe in examples/train/configs/overfit_ltx2_t2v_nvfp4_qat.yaml combines that linear configuration with models.student.attention_backend: ATTN_QAT_TRAIN for video-attention forward/backward. On sm120, its validation callback temporarily switches those layers to ATTN_QAT_INFER. On GB200, set callbacks.validation.attn_qat_infer: false to keep validation on the train-time QAT backend; the inference kernel is sm120-only.

User-adaptable LTX-2 fine-tuning recipes (full, LoRA, and NVFP4 QAT) live in examples/train/configs/fine_tuning/ltx2/, alongside the other model families under examples/train/configs/fine_tuning/.

Training Methods¶

Supervised Fine-Tuning (SFT)¶

Standard flow-matching loss. The simplest method — train the student to predict noise (or clean x0) from noised data samples.

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.fine_tuning.finetune.FineTuneMethod
  attn_kind: dense   # "dense" or "vsa"

Diffusion-Forcing SFT (DFSFT)¶

SFT with per-chunk inhomogeneous timesteps — each temporal chunk of the video gets a different noise level. This is a prerequisite for training causal / streaming models that must handle mixed-noise inputs.

method:
  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
  chunk_size: 3
  min_timestep_ratio: 0.0
  max_timestep_ratio: 1.0
  attn_kind: dense

DMD2 (Distribution Matching Distillation)¶

Distill a many-step teacher into a few-step student. The student learns to match the teacher's score function, guided by a trainable critic network.

models:
  student:
    _target_: fastvideo.train.models.wan.WanModel
    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
    trainable: true
  teacher:
    _target_: fastvideo.train.models.wan.WanModel
    init_from: Wan-AI/Wan2.1-T2V-1.3B-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.dmd2.DMD2Method
  rollout_mode: simulate
  dmd_denoising_steps: [1000, 750, 500, 250]
  generator_update_interval: 5
  real_score_guidance_scale: 4.5

  fake_score_learning_rate: 8.0e-6
  fake_score_betas: [0.0, 0.999]
  fake_score_lr_scheduler: constant

Self-Forcing (Causal DMD)¶

Extends DMD2 for streaming / causal video generation. The student processes video in temporal chunks, feeding its own denoised outputs as context for future chunks — simulating autoregressive rollout during training.

Requires a causal model class (e.g., WanCausalModel) for the student:

models:
  student:
    _target_: fastvideo.train.models.wan.wan_causal.WanCausalModel
    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
    trainable: true

method:
  _target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod
  rollout_mode: simulate
  dmd_denoising_steps: [1000, 750, 500, 250]
  student_sample_type: sde
  context_noise: 0.0
  enable_gradient_in_rollout: true
  start_gradient_frame: 0

Self-Forcing inherits all DMD2 parameters, plus:

Streaming Long Tuning¶

StreamingLongTuningMethod extends Self-Forcing for LongLive-style rollouts. It keeps a streaming state, generates overlapping chunks, and trains only the new frames while preserving context from earlier chunks.

For the MatrixGame2/Zelda world-model example, self-forcing and long tuning are separate runs: first train or load the 1k-step self-forcing checkpoint using examples/train/scenario/worldmodel/zelda/self_forcing_causal_i2v.yaml, then run examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml from that checkpoint for the 3k-step streaming long-tuning stage.

method:
  _target_: fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
  streaming_chunk_size: 9
  streaming_max_length: 39
  streaming_fixed_overlap_latents: 3
  streaming_reencode_overlap_anchor: true
  streaming_anchor_inject_k: 1
  streaming_require_full_blocks: true
  multi_phased_distill_schedule:
    - stage: streaming_long
      start_step: 0
      end_step: 3000
      num_latent_t: 39
      streaming_training: true

See examples/train/scenario/worldmodel/zelda/streaming_long_tuning_causal_i2v.yaml for a complete MatrixGame2/Zelda configuration.

Callbacks¶

Callbacks are pluggable hooks that run at specific points in the training loop. Configure them under the callbacks section.

GradNormClipCallback¶

Clips gradient norms before the optimizer step. Optionally logs per-module gradient norms to the tracker.

callbacks:
  grad_clip:
    max_grad_norm: 1.0      # 0.0 = disabled
    log_grad_norms: false

EMACallback¶

Maintains an exponential moving average of the student's weights. The EMA weights are automatically swapped in during validation.

callbacks:
  ema:
    _target_: fastvideo.train.callbacks.ema.EMACallback
    decay: 0.9999
    start_iter: 0   # delay EMA updates until this iteration

The EMA callback owns its own state and checkpoints independently — EMA weights are saved and restored automatically on resume.

ValidationCallback¶

Runs inference with the trained model at regular intervals, saving generated videos and logging them to the tracker (W&B).

callbacks:
  validation:
    _target_: fastvideo.train.callbacks.validation.ValidationCallback
    pipeline_target: fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline
    dataset_file: path/to/validation.json
    every_steps: 100
    sampling_steps: [4]
    sampling_timesteps: [1000, 750, 500, 250]  # explicit timestep list
    guidance_scale: 5.0
    rollout_mode: parallel  # "parallel" or "streaming"

The validation dataset is a JSON file containing a list of prompt strings. If EMA is enabled, validation automatically uses the EMA weights.

Checkpointing and Resume¶

Checkpoint format¶

Checkpoints use PyTorch Distributed Checkpoint (DCP) format, compatible with FSDP/HSDP sharding. Each checkpoint saves:

• Model weights (all roles)
• Optimizer states (all roles)
• LR scheduler states
• RNG states (for exact reproducibility)
• EMA shadow weights (if enabled)
• Training step counter

Checkpoints are saved to <output_dir>/checkpoint-<step>/.

Saving checkpoints¶

training:
  checkpoint:
    output_dir: outputs/my_run
    training_state_checkpointing_steps: 1000  # save every N steps (0 = off)
    checkpoints_total_limit: 3                # rolling window (0 = keep all)

Resuming training¶

Use --resume-from-checkpoint to resume from a specific checkpoint:

# Via the helper script
bash examples/train/run.sh my_config.yaml --resume outputs/my_run/checkpoint-2000

# Via torchrun directly
torchrun --nproc_per_node=8 \
    fastvideo/train/entrypoint/train.py \
    --config my_config.yaml \
    --resume-from-checkpoint outputs/my_run/checkpoint-2000

Or set it in the YAML:

training:
  checkpoint:
    resume_from_checkpoint: outputs/my_run/checkpoint-2000

Reproducibility¶

The training entrypoint enables deterministic mode automatically:

• torch.backends.cudnn.benchmark = False
• torch.backends.cudnn.deterministic = True
• torch.use_deterministic_algorithms(True)

A shared CUDA RNG generator is seeded from training.data.seed and threaded through all random operations (noise sampling, timestep sampling, etc.). Ranks within the same sequence-parallel group share a seed, ensuring identical noise across SP shards.

Distributed Training¶

The framework supports HSDP (Hybrid Sharded Data Parallel), Tensor Parallelism (TP), and Sequence Parallelism (SP):

training:
  distributed:
    num_gpus: 8
    sp_size: 1            # sequence parallelism group size
    tp_size: 1            # tensor parallelism group size
    hsdp_replicate_dim: 1 # number of HSDP replicas
    hsdp_shard_dim: 8     # number of HSDP shards

HSDP shards model parameters across hsdp_shard_dim GPUs and replicates across hsdp_replicate_dim groups. The product hsdp_replicate_dim * hsdp_shard_dim should equal num_gpus.

Sequence parallelism splits the sequence (video frames) across sp_size GPUs within each data-parallel group. Useful for long videos that don't fit on a single GPU.

VSA (Variable Sparse Attention)¶

VSA progressively increases attention sparsity during training, reducing compute while maintaining quality:

training:
  vsa:
    sparsity: 0.9             # target sparsity level
    decay_rate: 0.03          # sparsity increment per decay interval
    decay_interval_steps: 1   # steps between sparsity increases

The effective sparsity at step t is min(sparsity, decay_rate * (t // decay_interval_steps)).

Extending the Framework¶

Adding a new model¶

1. Create a new module under fastvideo/train/models/ (e.g., fastvideo/train/models/mymodel/mymodel.py).
2. Subclass ModelBase (or CausalModelBase for streaming models).
3. Implement the required methods:
   • prepare_batch() — convert raw dataloader output to TrainingBatch
   • add_noise() — forward-process noise addition
   • predict_noise() — run the transformer forward pass
   • backward() — backward pass with forward context restoration
4. Reference it in your YAML config:

models:
  student:
    _target_: fastvideo.train.models.mymodel.mymodel.MyModel
    init_from: my-org/my-model
    trainable: true

Adding a new training method¶

1. Create a new module under fastvideo/train/methods/.
2. Subclass TrainingMethod.
3. Implement the required methods:
   • single_train_step() — one forward pass returning losses, outputs, metrics
   • get_optimizers() — return optimizer list
   • get_lr_schedulers() — return scheduler list
4. Reference it in your config:

method:
  _target_: fastvideo.train.methods.my_method.MyMethod
  my_param: 42

Method-specific parameters are accessible via self.method_config (a plain dict).

Adding a new callback¶

1. Create a new module under fastvideo/train/callbacks/.
2. Subclass Callback.
3. Override the hooks you need: on_train_start, on_training_step_end, on_before_optimizer_step, etc.
4. Optionally implement state_dict() / load_state_dict() for checkpoint persistence.
5. Add it to your config:

callbacks:
  my_callback:
    _target_: fastvideo.train.callbacks.my_callback.MyCallback
    my_param: 42

File Structure¶

fastvideo/train/
  entrypoint/
    train.py                  # CLI entrypoint (torchrun)
  trainer.py                  # Training loop orchestrator
  models/
    base.py                   # ModelBase, CausalModelBase ABCs
    wan/
      wan.py                  # Wan 2.1 T2V model
      wan_causal.py           # Wan causal (streaming) model
  methods/
    base.py                   # TrainingMethod ABC
    distribution_matching/
      dmd2.py                 # DMD2 distillation
      self_forcing.py         # Self-Forcing (causal DMD)
    fine_tuning/
      finetune.py             # Supervised fine-tuning
      dfsft.py                # Diffusion-forcing SFT
  callbacks/
    callback.py               # Callback ABC and CallbackDict
    grad_clip.py              # Gradient clipping + norm logging
    ema.py                    # EMA weight averaging
    validation.py             # Periodic inference validation
  utils/
    config.py                 # YAML parser -> RunConfig
    training_config.py        # Typed config dataclasses
    builder.py                # Model/method instantiation
    optimizer.py              # Optimizer/scheduler construction
    checkpoint.py             # DCP save/resume
    dataloader.py             # Dataset/dataloader construction
    tracking.py               # W&B tracker

Related Docs¶

• Training Architecture — design rationale, model/method abstractions, and open questions.
• Training Overview — data requirements and preprocessing.
• Data Preprocessing — how to prepare datasets.
• Config Reference — fully-commented YAML config with all fields and defaults.