[feat] Refactor training framework into fastvideo/train (#1159)

Co-authored-by: Peiyuan Zhang <a1286225768@gmail.com>

.agents/memory/codebase-map/README.md

@@ -1,6 +1,6 @@
 # FastVideo-WorldModel — Codebase Map
 
-High-level structural index for agent orientation. Updated 2026-03-02.
+High-level structural index for agent orientation. Updated 2026-03-08.
 
 ## Repository Layout
 
@@ -22,7 +22,18 @@ FastVideo-WorldModel/
 │   ├── pipelines/             # End-to-end pipelines
 │   │   ├── basic/             #   Per-model pipelines (wan/, ltx2/, ...)
 │   │   └── stages/            #   Reusable pipeline stages
-│   ├── training/              # Training infrastructure
+│   ├── train/                 # Refactored training framework (YAML-driven, preferred)
+│   │   ├── trainer.py         #   Main training loop coordinator
+│   │   ├── entrypoint/        #   Training entrypoint (train.py) + checkpoint conversion
+│   │   ├── methods/           #   Training algorithms (FineTune, DFSFT, DMD2, SelfForcing)
+│   │   │   ├── base.py        #     TrainingMethod ABC
+│   │   │   ├── fine_tuning/   #     FineTuneMethod, DiffusionForcingSFTMethod
+│   │   │   └── distribution_matching/  # DMD2Method, SelfForcingMethod
+│   │   ├── models/            #   Per-role model wrappers (ModelBase, CausalModelBase)
+│   │   │   └── wan/           #     WanModel, WanCausalModel
+│   │   ├── callbacks/         #   Composable hooks (grad_clip, ema, validation)
+│   │   └── utils/             #   Config, builder, checkpoint, optimizer, tracking
+│   ├── training/              # Legacy training infrastructure (being phased out)
 │   │   ├── trackers.py        #   W&B tracker (BaseTracker → WandbTracker)
 │   │   ├── training_utils.py  #   Checkpointing, grad clipping, state dicts
 │   │   ├── training_pipeline.py        # Base training pipeline
@@ -65,6 +76,17 @@ FastVideo-WorldModel/
 
 ## Key Training Entrypoints
 
+### New framework (`fastvideo/train/`) — preferred
+
+| Method | Config Example | Launch Pattern |
+|--------|---------------|----------------|
+| FineTune (Wan) | `examples/train/finetune_wan2.1_t2v_1.3B_vsa_*.yaml` | `torchrun -m fastvideo.train.entrypoint.train --config <yaml>` |
+| DFSFT (Wan causal) | `examples/train/dfsft_wan_causal_t2v_1.3B.yaml` | `torchrun -m fastvideo.train.entrypoint.train --config <yaml>` |
+| DMD2 distillation | `examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml` | `torchrun -m fastvideo.train.entrypoint.train --config <yaml>` |
+| Self-Forcing | `examples/train/self_forcing_wan_causal_t2v_1.3B.yaml` | `torchrun -m fastvideo.train.entrypoint.train --config <yaml>` |
+
+### Legacy pipelines (`fastvideo/training/`) — being phased out
+
 | Pipeline | Entrypoint | Launch Pattern |
 |----------|-----------|----------------|
 | Wan T2V finetune | `fastvideo/training/wan_training_pipeline.py` | `torchrun --nproc_per_node N` |

.agents/memory/evaluation-registry/README.md

@@ -280,7 +280,7 @@ A comprehensive benchmark examining **four critical capabilities**:
 **Priority**: **Highest** — this is the most important evaluation signal
 **Trust**: Highest — but expensive
 
-#### What It Measures
+### What It Measures
 Human evaluators compare generated videos and rate them on dimensions like:
 - Overall quality and realism
 - Temporal coherence and smoothness

.agents/onboarding/README.md

@@ -22,9 +22,10 @@ Read these files to build your context:
 |----------|------|----------------|
 | 1 | `AGENTS.md` | Coding guidelines, build/test commands, PR conventions |
 | 2 | `docs/design/overview.md` | Architecture: models, pipelines, configs, registry |
-| 3 | `docs/training/overview.md` | Training data flow and preprocessing |
-| 4 | `docs/training/finetune.md` | Training arguments, parallelism, LoRA, validation |
-| 5 | `docs/contributing/coding_agents.md` | How to add model pipelines with agent assistance |
+| 3 | `fastvideo/train/` | Refactored training framework (YAML-driven, modular methods/models/callbacks) |
+| 4 | `docs/training/overview.md` | Training data flow and preprocessing |
+| 5 | `docs/training/finetune.md` | Training arguments, parallelism, LoRA, validation |
+| 6 | `docs/contributing/coding_agents.md` | How to add model pipelines with agent assistance |
 
 ## Step 2: Discover Available Resources
 

.agents/onboarding/worldmodel-training/README.md

@@ -17,8 +17,25 @@ auto-regressive streaming generation.
 - Full finetuning and LoRA on Wan / LTX-2 / MatrixGame models
 - DMD-based distillation (few-step generation)
 - Self-Forcing distillation (causal streaming)
-- Consistency finetuning with causal ODE initialization
-- Action injection modules for interactive control
+- Diffusion-Forcing SFT (DFSFT) for causal models
+- VSA (Variable Sparsity Acceleration) for efficient training
+
+---
+
+## Training Code: Two Generations
+
+### New modular framework: `fastvideo/train/` (preferred)
+
+The refactored training code uses a **YAML-only config-driven** architecture
+with composable methods, per-role models, and a callback system. All new
+training work should use this framework.
+
+### Legacy pipelines: `fastvideo/training/` (deprecated)
+
+The old monolithic pipeline classes (`WanTrainingPipeline`,
+`DistillationPipeline`, etc.) still exist but are being phased out. The new
+framework imports select utilities from `fastvideo/training/` for backward
+compatibility (EMA, gradient clipping, checkpoint wrappers).
 
 ---
 
@@ -29,54 +46,189 @@ Read these **in order** before touching any training code:
 | # | File | What You Learn |
 |---|------|----------------|
 | 1 | `docs/training/overview.md` | Training data flow: raw video → text embeddings + video latents → training |
-| 2 | `docs/training/finetune.md` | All training arguments, parallelism (SP/TP), LoRA, validation settings |
+| 2 | `docs/training/finetune.md` | Training arguments, parallelism (SP/TP), LoRA, validation settings |
 | 3 | `docs/training/data_preprocess.md` | How to preprocess datasets into the expected format |
 | 4 | `docs/design/overview.md` | Architecture: models, pipelines, configs, registry |
 
 ---
 
-## Training Pipelines
+## New Training Framework (`fastvideo/train/`)
 
-Each pipeline is a Python entrypoint launched via `torchrun`:
+### Architecture Overview
 
-| Pipeline | Entrypoint | Use Case |
-|----------|-----------|----------|
-| **Wan T2V finetune** | `fastvideo/training/wan_training_pipeline.py` | Standard text-to-video full finetune / LoRA |
-| **Wan I2V finetune** | `fastvideo/training/wan_i2v_training_pipeline.py` | Image-to-video (condition on first frame) |
-| **MatrixGame finetune** | `fastvideo/training/matrixgame_training_pipeline.py` | Action-conditioned world model training |
-| **LTX-2 finetune** | `fastvideo/training/ltx2_training_pipeline.py` | LTX-2 architecture finetuning |
-| **Wan DMD distillation** | `fastvideo/training/wan_distillation_pipeline.py` | Few-step distillation via DMD |
-| **Self-Forcing distill** | `fastvideo/training/wan_self_forcing_distillation_pipeline.py` | Causal streaming distillation |
-| **Base training** | `fastvideo/training/training_pipeline.py` | Base class — not called directly |
+```
+fastvideo/train/
+├── __init__.py                    → exports Trainer
+├── trainer.py                     → main training loop coordinator
+├── entrypoint/
+│   ├── train.py                   → YAML-only training entrypoint
+│   └── dcp_to_diffusers.py        → checkpoint conversion utility
+├── methods/                       → training algorithms (TrainingMethod ABC)
+│   ├── base.py                    → TrainingMethod base class
+│   ├── fine_tuning/
+│   │   ├── finetune.py            → FineTuneMethod (supervised finetuning)
+│   │   └── dfsft.py               → DiffusionForcingSFTMethod (causal)
+│   ├── distribution_matching/
+│   │   ├── dmd2.py                → DMD2Method (distribution matching distill)
+│   │   └── self_forcing.py        → SelfForcingMethod (causal streaming)
+│   ├── knowledge_distillation/    → (stub, not yet implemented)
+│   └── consistency_model/         → (stub, not yet implemented)
+├── models/                        → per-role model instances
+│   ├── base.py                    → ModelBase & CausalModelBase (ABC)
+│   └── wan/
+│       ├── wan.py                 → WanModel (non-causal)
+│       └── wan_causal.py          → WanCausalModel (causal streaming)
+├── callbacks/                     → training hooks & monitoring
+│   ├── callback.py                → Callback base class + CallbackDict
+│   ├── grad_clip.py               → GradNormClipCallback
+│   ├── ema.py                     → EMACallback (shadow weights)
+│   └── validation.py              → ValidationCallback (sampling + eval)
+└── utils/                         → configuration, building, checkpointing
+    ├── builder.py                 → build_from_config() (config → runtime)
+    ├── checkpoint.py              → CheckpointManager (DCP-based)
+    ├── config.py                  → load_run_config() (YAML → RunConfig)
+    ├── training_config.py         → TypedConfig dataclasses
+    ├── optimizer.py               → build_optimizer_and_scheduler()
+    ├── instantiate.py             → resolve_target() + instantiate()
+    ├── tracking.py                → build_tracker() (W&B, etc.)
+    ├── dataloader.py              → dataloader utilities
+    ├── module_state.py            → apply_trainable()
+    └── moduleloader.py            → load_module_from_path()
+```
+
+### Key Concepts
+
+**TrainingMethod** (`methods/base.py`): Abstract base class for all training
+algorithms. Owns role models (student, teacher, critic), manages checkpoint
+state, and defines the training step interface.
+
+**ModelBase** (`models/base.py`): Per-role model wrapper. Each role (student,
+teacher, critic) gets its own `ModelBase` instance owning a `transformer` and
+`noise_scheduler`. `CausalModelBase` extends this for streaming models.
+
+**Callback system** (`callbacks/`): Composable hooks for gradient clipping,
+EMA, validation, etc. Configured via YAML, dispatched by `CallbackDict`.
+
+**Config system** (`utils/config.py`, `utils/training_config.py`): YAML files
+are parsed into typed `RunConfig` dataclass trees. Models and methods use
+`_target_` fields for instantiation (similar to Hydra).
+
+### Training Flow
+
+```
+run_training_from_config(config_path)
+  → load_run_config()           # YAML → RunConfig
+  → init_distributed()          # TP/SP setup
+  → build_from_config()         # instantiate models, method, dataloader
+  → Trainer.run()               # main loop:
+      ├─ callbacks.on_train_start()
+      ├─ checkpoint_manager.maybe_resume()
+      ├─ for step in range(max_steps):
+      │    ├─ method.single_train_step(batch)
+      │    ├─ method.backward()
+      │    ├─ callbacks.on_before_optimizer_step()
+      │    ├─ method.optimizers_schedulers_step()
+      │    ├─ tracker.log(metrics, step)
+      │    ├─ callbacks.on_training_step_end()
+      │    └─ checkpoint_manager.maybe_save(step)
+      ├─ callbacks.on_train_end()
+      └─ checkpoint_manager.save_final()
+```
+
+### Training Methods
+
+| Method | Class | Use Case |
+|--------|-------|----------|
+| **FineTune** | `FineTuneMethod` | Single-role supervised finetuning |
+| **DFSFT** | `DiffusionForcingSFTMethod` | Diffusion-forcing SFT with inhomogeneous timesteps |
+| **DMD2** | `DMD2Method` | Multi-role distribution matching distillation (student + teacher + critic) |
+| **Self-Forcing** | `SelfForcingMethod` | Extends DMD2 for causal student rollouts |
+
+### Launching Training (New Framework)
+
+Training is launched via `torchrun` with a single YAML config:
+
+```bash
+torchrun --nproc_per_node <N_GPUS> \
+  -m fastvideo.train.entrypoint.train \
+  --config examples/train/<config>.yaml
+```
+
+### Example YAML Configs
+
+| Config | Method | Description |
+|--------|--------|-------------|
+| `examples/train/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml` | FineTune | Wan 1.3B finetuning with VSA sparsity |
+| `examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml` | DMD2 | Wan 1.3B distillation (student + teacher + critic) |
+| `examples/train/dfsft_wan_causal_t2v_1.3B.yaml` | DFSFT | Causal Wan 1.3B diffusion-forcing SFT |
+| `examples/train/self_forcing_wan_causal_t2v_1.3B.yaml` | Self-Forcing | Causal streaming distillation |
+
+### Checkpointing (New Framework)
+
+**CheckpointManager** (`utils/checkpoint.py`) saves via `torch.distributed.checkpoint`:
+
+```
+output_dir/
+└─ checkpoint-{step}/
+   ├─ dcp/                    # DCP state dict
+   ├─ config.json             # resolved training config
+   └─ .fastvideo_metadata.json
+```
+
+Checkpoint state includes: role model weights, per-role optimizers/schedulers,
+CUDA RNG state, and callback state (e.g., EMA shadow weights).
+
+### Config Structure
+
+A YAML config defines the full training pipeline:
+
+```yaml
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel
+    model_path: ...
+    trainable: true
+  teacher:  # optional, for distillation
+    _target_: fastvideo.train.models.wan.WanModel
+    model_path: ...
+    trainable: false
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.FineTuneMethod
+  # method-specific params...
+
+training:
+  distributed: { num_gpus: 8, tp_size: 1, sp_size: 8 }
+  data: { data_path: ..., batch_size: 1 }
+  optimizer: { lr: 1e-5, lr_scheduler: constant_with_warmup }
+  loop: { max_train_steps: 1000 }
+  checkpoint: { output_dir: ./outputs }
+  tracker: { trackers: [wandb], project_name: ... }
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.ValidationCallback
+    validation_steps: 100
+```
 
 ---
 
-## Example Scripts
+## Legacy Training Pipelines (`fastvideo/training/`)
 
-Ready-to-run training launches. Use these as starting templates:
+> **Note:** Use the new `fastvideo/train/` framework for new work. This section
+> is retained for reference on existing pipelines not yet migrated.
 
-### Finetuning
-| Model | Script | Notes |
-|-------|--------|-------|
-| Wan T2V 1.3B | `examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh` | Smallest, fastest for testing |
-| Wan T2V 1.3B LoRA | `examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v_lora.sh` | Lightweight adapter |
-| Wan I2V 14B | `examples/training/finetune/wan_i2v_14B_480p/crush_smol/finetune_i2v.sh` | Large I2V model |
-| LTX-2 | `examples/training/finetune/ltx2/finetune_t2v.sh` | Alternative architecture |
-| MatrixGame 2.0 | `examples/training/finetune/MatrixGame2.0/finetune_i2v.sh` | Action-conditioned |
-
-### Distillation
-| Method | Script | Notes |
-|--------|--------|-------|
-| DMD Wan | `scripts/distill/v1_distill_dmd_wan.sh` | Full distillation launch |
-| DMD Wan + VSA | `scripts/distill/v1_distill_dmd_wan_VSA.sh` | With variable-step acceleration |
-| Consistency (causal ODE init) | `examples/training/consistency_finetune/causal_ode_init/finetune_ode_init.sh` | Consistency tuning |
-
-### Data Preprocessing
-Each model directory includes a `preprocess_*` script. Always preprocess first:
-```bash
-# Example for Wan T2V
-bash examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh
-```
+| Pipeline | Entrypoint | Use Case |
+|----------|-----------|----------|
+| Wan T2V finetune | `fastvideo/training/wan_training_pipeline.py` | Standard text-to-video finetune / LoRA |
+| Wan I2V finetune | `fastvideo/training/wan_i2v_training_pipeline.py` | Image-to-video (first frame conditioned) |
+| MatrixGame finetune | `fastvideo/training/matrixgame_training_pipeline.py` | Action-conditioned world model |
+| LTX-2 finetune | `fastvideo/training/ltx2_training_pipeline.py` | LTX-2 architecture finetuning |
+| Wan DMD distillation | `fastvideo/training/wan_distillation_pipeline.py` | Few-step distillation via DMD |
+| Self-Forcing distill | `fastvideo/training/wan_self_forcing_distillation_pipeline.py` | Causal streaming distillation |
 
 ---
 
@@ -84,19 +236,12 @@ bash examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.
 
 ### W&B Integration
 - **Tracker**: `fastvideo/training/trackers.py` — `WandbTracker` class
+- **New framework tracker**: `fastvideo/train/utils/tracking.py` — `build_tracker()`
 - **Env vars**: `WANDB_API_KEY`, `WANDB_BASE_URL`, `WANDB_MODE`
-- **Summaries**: Saved to `<output_dir>/tracker/wandb/latest-run/files/wandb-summary.json`
-- **Tests**: `fastvideo/tests/training/Vanilla/test_training_loss.py` (compares against reference summaries)
-
-### Checkpointing
-- **Save**: `fastvideo/training/training_utils.py:save_checkpoint()`
-- **Load**: `fastvideo/training/training_utils.py:load_checkpoint()`
-- **Distill save**: `training_utils.py:save_distillation_checkpoint()` (multi-model)
-- **Format**: FSDP distributed checkpoint → converted to HF format
 
 ### Parallelism
-- **SP** (Sequence Parallel): splits video frames across GPUs — `--sp_size N`
-- **TP** (Tensor Parallel): splits model layers across GPUs — `--tp_size N`
+- **SP** (Sequence Parallel): splits video frames across GPUs — `sp_size: N`
+- **TP** (Tensor Parallel): splits model layers across GPUs — `tp_size: N`
 - Typical configs: SP=2–8, TP=1–2
 
 ---
@@ -148,23 +293,10 @@ their own error accumulation during long auto-regressive generation.
 
 ### DMD Distillation (Distribution Matching Distillation)
 Reduces inference steps from ~50 to 3–4 by training a student model to match
-the output distribution of the teacher model. Uses ODE pairs collected from
-the teacher model as training targets.
+the output distribution of the teacher model. Uses a critic network to estimate
+distribution divergence.
 
----
-
-## Quick-Start: Minimal Training Test
-
-To verify the training infrastructure works, run the smallest possible experiment:
-
-```bash
-# 1. Download crush_smol dataset
-bash examples/training/finetune/wan_t2v_1.3B/crush_smol/download_dataset.sh
-
-# 2. Preprocess
-bash examples/training/finetune/wan_t2v_1.3B/crush_smol/preprocess_wan_data_t2v.sh
-
-# 3. Short training run (5 steps)
-# Edit finetune_t2v.sh: set --max_train_steps 5
-bash examples/training/finetune/wan_t2v_1.3B/crush_smol/finetune_t2v.sh
-```
+### Diffusion-Forcing SFT (DFSFT)
+Supervised finetuning with **inhomogeneous timesteps** across chunks — each
+chunk in a causal sequence can have a different noise level, training the model
+to handle mixed-fidelity contexts.

.agents/skills/SKILL_TEMPLATE.md

@@ -28,6 +28,7 @@ description: <one-line description — Codex uses this for implicit invocation m
 - <What this skill produces>
 
 ## Example Usage
+
 ```
 <Example invocation or prompt snippet>
 ```
@@ -50,6 +51,7 @@ Each skill lives in its own directory under `.agents/skills/`:
 ```
 
 After creating a new skill, add an entry to `.agents/skills/index.jsonl`:
+
 ```json
 {"name": "<skill-name>", "description": "<description>", "path": "<skill-name>/SKILL.md", "status": "draft", "trust": "low"}
 ```

.agents/skills/evaluate-video-quality/SKILL.md

@@ -38,6 +38,7 @@ pytest fastvideo/tests/ssim/ -vs --video-path <generated> --reference-path <refe
 ```
 
 Or use the SSIM utility directly:
+
 ```python
 from fastvideo.tests.ssim.ssim_utils import compute_ssim
 score = compute_ssim(generated_video, reference_video)

.agents/skills/index-related-work/SKILL.md

@@ -76,6 +76,7 @@ If not, create it:
 - Updated catalog.
 
 ## Example Usage
+
 ```
 Index the Self-Forcing paper:
 

.agents/skills/launch-experiment/SKILL.md

@@ -88,6 +88,7 @@ torchrun --nnodes 1 --nproc_per_node <num_gpus> \
 ### 5. Log to experiment journal
 
 After launching, append an entry to `.agents/memory/experiment-journal/README.md`:
+
 ```markdown
 ## [YYYY-MM-DD] Experiment: <run_name>
 - **Hypothesis**: <user-provided or auto-generated>
@@ -101,6 +102,7 @@ After launching, append an entry to `.agents/memory/experiment-journal/README.md
 - An experiment journal entry.
 
 ## Example Usage
+
 ```
 Launch a Wan T2V 1.3B finetune on 4 GPUs with lr=5e-5 and max_train_steps=1000:
 

.agents/skills/log-experiment/SKILL.md

@@ -63,6 +63,7 @@ may be a repeat. If it's `status: running`, assume this is an update.
 - Updated `.agents/memory/experiment-journal/README.md`.
 
 ## Example Usage
+
 ```
 Log a completed experiment:
 

.agents/skills/monitor-experiment/SKILL.md

@@ -32,6 +32,7 @@ report before committing more resources.
 ### 1. Connect to the run
 
 **Online mode** (preferred):
+
 ```python
 import wandb
 api = wandb.Api()
@@ -39,6 +40,7 @@ run = api.run("<run_id>")
 ```
 
 **Offline fallback**:
+
 ```python
 import json
 summary_path = f"{output_dir}/tracker/wandb/latest-run/files/wandb-summary.json"
@@ -70,6 +72,7 @@ with open(summary_path) as f:
 ### 4. Emit structured status
 
 Output format (agent-consumable):
+
 ```json
 {
   "run_id": "...",
@@ -111,6 +114,7 @@ After the first 30 minutes of wall-clock time:
 - 30-minute checkpoint quality report.
 
 ## Example Usage
+
 ```
 Monitor W&B run "fastvideo/Wan_distillation/abc123":
 

.agents/skills/search-related-work/SKILL.md

@@ -65,6 +65,7 @@ For each matching file:
 - Ranked list of relevant related work entries with excerpts.
 
 ## Example Usage
+
 ```
 Search for work related to video quality evaluation metrics:
 

.agents/skills/summarize-run/SKILL.md

@@ -31,6 +31,7 @@ online (W&B API) and offline (local `wandb-summary.json`) modes.
 ### 1. Load run summary
 
 **Online**:
+
 ```python
 import wandb
 api = wandb.Api()
@@ -40,6 +41,7 @@ config = dict(run.config)
 ```
 
 **Offline** (existing codebase pattern from `fastvideo/tests/training/`):
+
 ```python
 import json
 summary_path = f"{output_dir}/tracker/wandb/latest-run/files/wandb-summary.json"
@@ -114,6 +116,7 @@ with the final metrics and status.
 - Updated experiment journal entry.
 
 ## Example Usage
+
 ```
 Summarize the run in output directory "outputs/wan_finetune":
 

docs/design/training_architecture.md

@@ -0,0 +1,369 @@
+# Training Architecture
+
+!!! warning "Work in Progress"
+    This training architecture (`fastvideo/train/`) is under active development
+    and is replacing the older `fastvideo/training/` module. APIs, config
+    formats, and supported methods may change. See the
+    [Current Status](#current-status) section for what is implemented so far.
+
+FastVideo's training framework (`fastvideo/train/`) is built around a
+**pluggable, YAML-driven architecture** that cleanly separates **models**,
+**training methods**, and **infrastructure** into independent, composable
+layers. A single YAML config file is all that is needed to train any supported
+model with any supported algorithm — no code changes required to mix and match.
+
+---
+
+## Motivation
+
+Training video diffusion models involves a tangle of concerns: model loading,
+noise scheduling, distillation algorithms, distributed strategies,
+checkpointing, and validation. Existing training scripts tend to hard-wire
+these together, making it painful to:
+
+1. **Try a new distillation algorithm** on an existing model (requires forking
+   the training loop).
+2. **Add a new model** to an existing algorithm (requires re-implementing
+   boilerplate).
+3. **Switch distributed strategies** (FSDP, TP, SP) without touching algorithm
+   code.
+4. **Resume, checkpoint, and validate** uniformly across all combinations.
+
+The training framework solves this by making each axis of variation an
+independent plugin.
+
+---
+
+## Architecture Overview
+
+```
+YAML Config
+    |
+    v
++------------------+     +---------------------+     +------------------+
+|   Models Layer   |     |   Methods Layer     |     | Infrastructure   |
+|  (per-role)      |     |   (algorithm)       |     |   Layer          |
+|                  |     |                     |     |                  |
+| - ModelBase      |<----|  - TrainingMethod   |---->| - Trainer        |
+| - CausalModelBase|     |  - single_train_step|     | - Callbacks      |
+|                  |     |  - backward         |     | - Checkpoint     |
+| Roles:           |     |  - optimizers       |     | - Tracker (W&B)  |
+|  student         |     |                     |     | - Dataloader     |
+|  teacher         |     | Algorithms:         |     |                  |
+|  critic          |     |  DMD2, SelfForcing, |     | Distributed:     |
+|                  |     |  SFT, DFSFT         |     |  HSDP, TP, SP    |
++------------------+     +---------------------+     +------------------+
+```
+
+### Three Layers
+
+| Layer | Responsibility | Extension point |
+|-------|---------------|-----------------|
+| **Models** (`fastvideo/train/models/`) | Load transformer + scheduler, define `predict_noise`, `predict_x0`, `add_noise`, `backward`. Each training role (student/teacher/critic) is an independent instance. | Subclass `ModelBase` (or `CausalModelBase` for streaming). |
+| **Methods** (`fastvideo/train/methods/`) | Implement the training algorithm: own role models, define `single_train_step` + `backward`, manage optimizers/schedulers. | Subclass `TrainingMethod`. |
+| **Infrastructure** (`fastvideo/train/trainer.py`, `utils/`, `callbacks/`) | Training loop, gradient accumulation, distributed setup, checkpointing (DCP), W&B tracking, validation, EMA, grad clipping. | Add callbacks; everything else is shared. |
+
+---
+
+## YAML-Driven Configuration
+
+Everything is configured declaratively. The `_target_` field selects the Python
+class to instantiate:
+
+```yaml
+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
+  critic:
+    _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.dmd2.DMD2Method
+  rollout_mode: simulate
+  dmd_denoising_steps: [1000, 850, 700, 550, 350, 275, 200, 125]
+  generator_update_interval: 5
+  real_score_guidance_scale: 3.5
+  # ...
+
+training:
+  distributed: { num_gpus: 8, sp_size: 1, tp_size: 1 }
+  data: { data_path: ..., num_latent_t: 20, num_frames: 77 }
+  optimizer: { learning_rate: 2.0e-6, betas: [0.0, 0.999] }
+  loop: { max_train_steps: 4000 }
+  checkpoint: { output_dir: outputs/my_run }
+
+callbacks:
+  grad_clip: { max_grad_norm: 1.0 }
+  validation: { pipeline_target: ..., every_steps: 100 }
+```
+
+To switch from DMD2 to SFT, change the `method._target_` and remove the
+teacher/critic — no code changes needed.
+
+---
+
+## Model Abstraction
+
+### `ModelBase` — Standard (Bidirectional) Models
+
+Every role gets its own `ModelBase` instance owning a `transformer` and
+`noise_scheduler`. The base class defines:
+
+- **`prepare_batch()`** — Convert raw dataloader output into forward-ready
+  `TrainingBatch`.
+- **`add_noise()`** — Apply forward-process noise at a given timestep.
+- **`predict_noise()` / `predict_x0()`** — Run the transformer and return
+  predictions.
+- **`backward()`** — Backward pass that restores forward context (attention
+  metadata, timesteps).
+- **`init_preprocessors()`** — Lazy-load VAE, build dataloader (called only on
+  the student).
+
+### `CausalModelBase` — Streaming / Causal Models
+
+Extends `ModelBase` with streaming inference primitives for causal video
+generation:
+
+```python
+class CausalModelBase(ModelBase):
+    def clear_caches(self, *, cache_tag: str = "pos") -> None: ...
+    def predict_noise_streaming(
+        self, ..., cache_tag, store_kv, cur_start_frame
+    ) -> Tensor | None: ...
+    def predict_x0_streaming(
+        self, ..., cache_tag, store_kv, cur_start_frame
+    ) -> Tensor | None: ...
+```
+
+KV caches are **internal** to the model instance, keyed by `cache_tag`. The
+method controls when to store (`store_kv=True`) vs. read-only
+(`store_kv=False`), enabling block-by-block causal rollout during training.
+
+---
+
+## Training Methods
+
+### DMD2 (Distribution Matching Distillation)
+
+**Roles:** student (trainable) + teacher (frozen) + critic (trainable)
+
+The student learns to generate clean video in few steps by matching the
+teacher's score function, with a critic network providing a learned fake-score
+baseline.
+
+- **Rollout modes:**
+  - `simulate` — Student starts from pure noise and iteratively denoises
+    through the full step schedule.
+  - `data_latent` — Student denoises from a single randomly-noised data
+    sample.
+- **Losses:** Generator loss (DMD gradient) + critic flow-matching loss, with
+  alternating updates (`generator_update_interval`).
+
+### Self-Forcing (Causal DMD)
+
+**Roles:** student (causal, trainable) + teacher (frozen) + critic (trainable)
+
+Extends DMD2 for **streaming/causal video generation**. The key idea: during
+training, the student processes video in temporal chunks, using its own
+previously-denoised outputs as context for future chunks — simulating online
+autoregressive rollout.
+
+- Video is split into blocks of `chunk_size` latent frames.
+- Each block is denoised through the student's step schedule; a random
+  early-exit step is sampled per block.
+- After denoising a block, its output is fed back (with optional
+  `context_noise`) as KV cache context for subsequent blocks via
+  `predict_noise_streaming(store_kv=True)`.
+- Supports SDE and ODE sampling during rollout.
+- Selective gradient control: `enable_gradient_in_rollout`,
+  `start_gradient_frame`.
+
+### Supervised Fine-Tuning (SFT)
+
+**Roles:** student only
+
+Standard flow-matching loss between predicted and ground-truth noise/x0.
+
+### Diffusion-Forcing SFT (DFSFT)
+
+**Roles:** student only
+
+SFT with **inhomogeneous (per-chunk) timesteps** — each temporal chunk in a
+video gets a different noise level. This trains the model to handle mixed-noise
+inputs, which is a prerequisite for causal/streaming inference where earlier
+frames are cleaner than later ones.
+
+---
+
+## Training Loop
+
+The `Trainer` runs a standard loop with pluggable method and callbacks:
+
+```
+for step in range(start_step, max_steps):
+    for accum_iter in range(grad_accum_steps):
+        batch <- dataloader
+        loss_map, outputs, metrics <- method.single_train_step(batch, step)
+        method.backward(loss_map, outputs)
+
+    callbacks.on_before_optimizer_step()   # grad clipping
+    method.optimizers_schedulers_step()
+    method.optimizers_zero_grad()
+    callbacks.on_training_step_end()       # logging
+    checkpoint_manager.maybe_save(step)
+    callbacks.on_validation_begin()        # periodic inference
+```
+
+### Callbacks
+
+- **GradNormClipCallback** — Per-module gradient norm logging + global
+  clipping.
+- **ValidationCallback** — Periodic inference sampling with configurable
+  pipeline, sampling steps, and guidance scale.
+- **EMACallback** — Exponential moving average of student weights.
+
+### Checkpointing
+
+- DCP (Distributed Checkpoint) format, compatible with FSDP/HSDP.
+- Saves: model weights, optimizer states, scheduler states, RNG states (per
+  role).
+- Full resume support: auto-restores step counter and all RNG states.
+
+---
+
+## Getting Started
+
+```bash
+# Install
+uv pip install -e .[dev]
+
+# Run DMD2 distillation on Wan 2.1
+torchrun --nproc_per_node=8 -m fastvideo.train.entrypoint.train \
+    --config examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml
+
+# Run SFT fine-tuning
+torchrun --nproc_per_node=8 -m fastvideo.train.entrypoint.train \
+    --config examples/train/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml
+```
+
+Example configs are in `examples/train/`.
+
+---
+
+## File Structure
+
+```
+fastvideo/train/
+  trainer.py                  # Training loop
+  models/
+    base.py                   # ModelBase, CausalModelBase ABCs
+    wan/wan.py                # Wan 2.1 T2V model plugin
+    wangame/wangame.py        # WanGame 2.1 I2V model plugin
+    wangame/wangame_causal.py # WanGame causal (streaming) plugin
+  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/
+    grad_clip.py              # Gradient clipping + norm logging
+    validation.py             # Periodic inference validation
+    ema.py                    # EMA weight averaging
+  entrypoint/
+    train.py                  # CLI entrypoint (torchrun)
+  utils/
+    config.py                 # YAML parser -> RunConfig
+    builder.py                # build_from_config: model/method instantiation
+    training_config.py        # TrainingConfig dataclass
+    dataloader.py             # Dataset/dataloader construction
+    optimizer.py              # Optimizer/scheduler construction
+    checkpoint.py             # DCP save/resume
+    tracking.py               # W&B tracker
+```
+
+---
+
+## Current Status
+
+| Component | Status |
+|-----------|--------|
+| Core framework (trainer, config, callbacks) | Implemented and tested |
+| `WanModel` (Wan 2.1 T2V) | Implemented and tested |
+| `WanGameModel` (WanGame 2.1 I2V) | Implemented and tested |
+| `WanGameCausalModel` (streaming) | Implemented and tested |
+| `WanCausalModel` (Wan T2V causal) | In progress |
+| DMD2 method | Implemented and tested |
+| Self-Forcing method | Implemented and tested |
+| SFT method | Implemented and tested |
+| DFSFT method | Implemented and tested |
+| DCP checkpointing + resume | Implemented and tested |
+| EMA callback | Implemented |
+| Validation callback | Implemented and tested |
+| Causal DMD inference pipeline | Implemented |
+
+---
+
+## Open Questions
+
+We welcome community feedback on the following topics:
+
+### Model Plugin API
+
+The current `ModelBase` interface requires implementing 6 methods. Is this the
+right granularity?
+
+- Should `prepare_batch` be split into separate concerns (noise sampling,
+  timestep sampling, attention metadata)?
+- Should `backward` be lifted out of the model and into the method/trainer?
+
+### Causal Streaming Interface
+
+`CausalModelBase` adds `predict_noise_streaming` / `predict_x0_streaming` with
+cache management. Alternatives considered:
+
+- **(a) Current:** Cache is internal to the model, keyed by `cache_tag`.
+  Simple but couples cache lifecycle to model.
+- **(b) External cache:** Method owns the cache dict, passes it into predict
+  calls. More explicit but verbose.
+- **(c) Context manager:** `with model.streaming_context(tag) as ctx: ...` —
+  cleaner lifecycle but harder to compose.
+
+### Method Composition
+
+Currently, methods are monolithic classes. Should we support composing methods
+(e.g., DFSFT pre-training followed by Self-Forcing distillation) within a
+single config? Or is sequential training with checkpoint handoff sufficient?
+
+### New Models and Methods
+
+What models and training methods should we prioritize next?
+
+- **Models:** HunyuanVideo, CogVideoX, other Wan variants?
+- **Methods:** Consistency models, progressive distillation, reward-based
+  fine-tuning?
+
+### Distributed Strategy
+
+Currently supports HSDP (hybrid sharded data parallel) + TP + SP. Are there
+scenarios where the current distributed setup is insufficient? Should we add
+pipeline parallelism for very large models?
+
+---
+
+## References
+
+- [Self-Forcing paper](https://arxiv.org/abs/2406.05477) — Chen et al., 2024.
+- [DMD2 paper](https://arxiv.org/abs/2405.14867) — Yin et al., 2024.
+- [Diffusion Forcing paper](https://arxiv.org/abs/2407.01392) — Chen et al.,
+  2024.

docs/training/train_infra.md

@@ -0,0 +1,595 @@
+# 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.
+
+!!! note "Relationship to legacy training"
+    This system replaces the older script-based training in `fastvideo/training/`.
+    The legacy scripts still work for basic fine-tuning, but new development
+    should use the config-driven system documented here.
+
+---
+
+## Quick Start
+
+### Launch with the helper script
+
+```bash
+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:
+
+```bash
+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
+
+```bash
+torchrun --nproc_per_node=8 \
+    fastvideo/train/entrypoint/train.py \
+    --config examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml
+```
+
+### CLI flags
+
+| Flag | Description |
+|------|-------------|
+| `--config` | Path to YAML config file (required) |
+| `--resume-from-checkpoint` | Path to a DCP checkpoint directory to resume from |
+| `--override-output-dir` | Override `training.checkpoint.output_dir` |
+| `--dry-run` | Validate config and exit without training |
+
+---
+
+## 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:
+
+```yaml
+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:
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `_target_` | *(required)* | Python class path for the model |
+| `init_from` | *(required)* | HuggingFace repo ID or local checkpoint path |
+| `trainable` | `true` | Whether the model's parameters require gradients |
+| `disable_custom_init_weights` | `false` | Skip custom weight initialization (use for teacher/critic) |
+| `flow_shift` | `3.0` | Timestep shifting factor |
+| `enable_gradient_checkpointing_type` | `null` | Gradient checkpointing (`"full"` or `null`) |
+
+Which roles are needed depends on the training method:
+
+| Method | Required roles |
+|--------|---------------|
+| Fine-tune (SFT) | `student` |
+| Diffusion-Forcing SFT | `student` |
+| DMD2 | `student`, `teacher`, `critic` |
+| Self-Forcing | `student` (causal), `teacher`, `critic` |
+
+### `method` — Training algorithm
+
+Selects and configures the training algorithm:
+
+```yaml
+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](#training-methods) for details on each algorithm.
+
+### `training` — Typed infrastructure config
+
+This section maps to typed dataclasses with defaults and validation:
+
+```yaml
+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
+```
+
+### `callbacks` — Pluggable hooks
+
+Callbacks run at specific points in the training loop (before/after optimizer
+steps, at validation time, etc.):
+
+```yaml
+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](#callbacks) for details on each callback.
+
+### `pipeline` — Inference pipeline overrides
+
+Optional overrides for the inference pipeline used during validation:
+
+```yaml
+pipeline:
+  flow_shift: 8
+```
+
+---
+
+## 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.
+
+```yaml
+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"
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `attn_kind` | `"dense"` | Attention mode: `"dense"` (standard) or `"vsa"` (sparse) |
+
+### 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.
+
+```yaml
+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
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `chunk_size` | `3` | Latent frames per temporal chunk |
+| `min_timestep_ratio` | `0.0` | Lower bound of timestep sampling range |
+| `max_timestep_ratio` | `1.0` | Upper bound of timestep sampling range |
+| `attn_kind` | `"dense"` | `"dense"` or `"vsa"` |
+
+### 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.
+
+```yaml
+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
+```
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `rollout_mode` | *(required)* | `"simulate"` (pure noise) or `"data_latent"` (from data) |
+| `dmd_denoising_steps` | *(required)* | Timestep schedule for student rollout |
+| `generator_update_interval` | `1` | Update student every N critic steps |
+| `real_score_guidance_scale` | `1.0` | CFG scale for teacher predictions |
+| `fake_score_learning_rate` | *(required)* | Critic optimizer learning rate |
+| `fake_score_betas` | *(required)* | Critic optimizer Adam betas |
+| `fake_score_lr_scheduler` | *(required)* | Critic LR scheduler type |
+
+### 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:
+
+```yaml
+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:
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `student_sample_type` | `"sde"` | `"sde"` or `"ode"` for intermediate steps |
+| `same_step_across_blocks` | `false` | Use same exit timestep for all blocks |
+| `last_step_only` | `false` | Always exit at the final denoising step |
+| `context_noise` | `0.0` | Noise added to context frames (0 = clean) |
+| `enable_gradient_in_rollout` | `true` | Enable backprop through rollout |
+| `start_gradient_frame` | `0` | Frame index where gradients begin |
+
+---
+
+## 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.
+
+```yaml
+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.
+
+```yaml
+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).
+
+```yaml
+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
+
+```yaml
+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:
+
+```bash
+# 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:
+
+```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):
+
+```yaml
+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:
+
+```yaml
+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:
+
+```yaml
+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:
+
+```yaml
+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:
+
+```yaml
+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/training_architecture.md) — design
+  rationale, model/method abstractions, and open questions.
+- [Training Overview](overview.md) — data requirements and preprocessing.
+- [Data Preprocessing](data_preprocess.md) — how to prepare datasets.
+- [Config Reference](../../examples/train/example.yaml) — fully-commented
+  YAML config with all fields and defaults.

examples/train/configs/dfsft_wan_causal_t2v_1.3B.yaml

@@ -0,0 +1,71 @@
+# DFSFT (Diffusion-Forcing SFT): Wan 2.1 T2V 1.3B Causal
+#
+# - Student: trainable causal Wan model
+# - Training: inhomogeneous timesteps per chunk (diffusion forcing)
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+    trainable: true
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.dfsft.DiffusionForcingSFTMethod
+  attn_kind: dense
+  chunk_size: 3
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    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
+    lr_warmup_steps: 0
+
+  loop:
+    max_train_steps: 4000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/wan2.1_causal_dfsft
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: distillation_wan_r
+    run_name: wan2.1_causal_dfsft
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_pipeline.WanCausalPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_64.json
+    every_steps: 50
+    sampling_steps: [40]
+    guidance_scale: 6.0
+    num_frames: 69
+
+pipeline:
+  flow_shift: 8

examples/train/configs/distill_wan2.1_t2v_1.3B_dmd2.yaml

@@ -0,0 +1,94 @@
+# DMD2 distillation: Wan 2.1 T2V 1.3B (teacher 50-step -> student 4-step).
+#
+# - Teacher: frozen pretrained Wan 2.1 T2V 1.3B
+# - Student: trainable, initialized from the same pretrained weights
+# - Critic:  trainable, initialized from the same pretrained weights
+# - Validation: 4-step SDE sampling
+
+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
+  generator_update_interval: 5
+  real_score_guidance_scale: 4.5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+
+  # Critic optimizer (required — no fallback to training.optimizer)
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  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_dmd2_4steps
+    training_state_checkpointing_steps: 20
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: distillation_wan
+    run_name: wan2.1_dmd2_4steps
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_dmd_pipeline.WanDMDPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
+    every_steps: 50
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    guidance_scale: 6.0
+  ema:                                                                                                                     
+    _target_: fastvideo.train.callbacks.ema.EMACallback                                                                    
+    decay: 0.98                                                                                                          
+    start_iter: 0        # delay EMA updates until this iteration 
+
+pipeline:
+  flow_shift: 8

examples/train/configs/example.yaml

@@ -0,0 +1,197 @@
+# ==============================================================================
+# Full configuration reference for fastvideo.train
+#
+# Legend:
+#   [TYPED]    — parsed into a typed dataclass; fields are validated with
+#                defaults. Unknown keys are silently ignored.
+#   [FREE]     — free-form dict passed as-is to the target class / method.
+#                Keys depend on the _target_ class constructor / method_config.
+#   [RESOLVED] — parsed by PipelineConfig.from_kwargs(); auto-populated from
+#                the model's config files. Only scalar overrides are useful.
+# ==============================================================================
+
+# ------------------------------------------------------------------------------
+# models: [FREE]
+#
+# Each role is instantiated via _target_(*, training_config=..., **kwargs).
+# Keys here are constructor kwargs of the _target_ class (e.g. WanModel).
+# You can define any role name (student, teacher, critic, etc.).
+# ------------------------------------------------------------------------------
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanModel    # required
+    init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers     # required: HF repo or local path
+    trainable: true                                   # default: true
+    disable_custom_init_weights: false                # default: false
+    flow_shift: 3.0                                   # default: 3.0
+    enable_gradient_checkpointing_type: null           # default: null (falls back to training.model)
+
+  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: [FREE]
+#
+# Instantiated via _target_(*, cfg=RunConfig, role_models=...).
+# All keys besides _target_ are available in self.method_config (a plain dict).
+# Keys depend entirely on the method class.
+# ------------------------------------------------------------------------------
+method:
+  _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method  # required
+
+  # --- DMD2-specific keys (read from self.method_config) ---
+  rollout_mode: simulate              # required: "simulate" or "data_latent"
+  generator_update_interval: 5        # default: 1
+  dmd_denoising_steps: [1000, 750, 500, 250]  # SDE timestep schedule
+
+  # Critic optimizer (all required — no fallback)
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+  # CFG conditioning policy (optional)
+  # cfg_uncond:
+  #   on_missing: error               # "error" or "ignore"
+  #   text: keep                      # "keep", "zero", "drop", "negative_prompt"
+  #   image: keep                     # "keep", "zero", "drop"
+  #   action: keep                    # "keep", "zero", "drop"
+
+  # --- FineTuneMethod keys (if using finetune instead) ---
+  # _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+  # attn_kind: vsa                    # "dense" or "vsa"
+
+# ------------------------------------------------------------------------------
+# training: [TYPED] -> TrainingConfig
+#
+# Every field below has a typed default. Unknown keys are ignored.
+# ------------------------------------------------------------------------------
+training:
+
+  # --- training.distributed [TYPED] -> DistributedConfig ---
+  distributed:
+    num_gpus: 8                       # default: 1
+    tp_size: 1                        # default: 1
+    sp_size: 1                        # default: 1 (defaults to num_gpus in loader)
+    hsdp_replicate_dim: 1             # default: 1
+    hsdp_shard_dim: 8                 # default: -1 (defaults to num_gpus in loader)
+    pin_cpu_memory: false             # default: false
+
+  # --- training.data [TYPED] -> DataConfig ---
+  data:
+    data_path: data/my_dataset        # default: ""
+    train_batch_size: 1               # default: 1
+    dataloader_num_workers: 4         # default: 0
+    training_cfg_rate: 0.1            # default: 0.0
+    seed: 1000                        # default: 0
+    num_height: 448                   # default: 0
+    num_width: 832                    # default: 0
+    num_latent_t: 20                  # default: 0
+    num_frames: 77                    # default: 0
+
+  # --- training.optimizer [TYPED] -> OptimizerConfig ---
+  # Note: only for the student optimizer. Critic optimizer is in method config.
+  optimizer:
+    learning_rate: 2.0e-6             # default: 0.0
+    betas: [0.9, 0.999]              # default: [0.9, 0.999]
+    weight_decay: 0.01                # default: 0.0
+    lr_scheduler: constant            # default: "constant"
+    lr_warmup_steps: 0                # default: 0
+    lr_num_cycles: 0                  # default: 0
+    lr_power: 0.0                     # default: 0.0
+    min_lr_ratio: 0.5                 # default: 0.5
+
+  # --- training.loop [TYPED] -> TrainingLoopConfig ---
+  loop:
+    max_train_steps: 10000            # default: 0
+    gradient_accumulation_steps: 1    # default: 1
+
+  # --- training.checkpoint [TYPED] -> CheckpointConfig ---
+  checkpoint:
+    output_dir: outputs/my_run        # default: ""
+    resume_from_checkpoint: ""        # default: "" (or use --resume-from-checkpoint CLI)
+    training_state_checkpointing_steps: 1000  # default: 0 (disabled)
+    checkpoints_total_limit: 3        # default: 0 (keep all)
+
+  # --- training.tracker [TYPED] -> TrackerConfig ---
+  tracker:
+    trackers: []                      # default: [] (auto-adds "wandb" if project_name is set)
+    project_name: my_project          # default: "fastvideo"
+    run_name: my_run                  # default: ""
+
+  # --- training.vsa [TYPED] -> VSAConfig ---
+  vsa:
+    sparsity: 0.0                     # default: 0.0 (0.0 = disabled)
+    decay_rate: 0.0                   # default: 0.0
+    decay_interval_steps: 0           # default: 0
+
+  # --- training.model [TYPED] -> ModelTrainingConfig ---
+  model:
+    weighting_scheme: uniform         # default: "uniform"
+    logit_mean: 0.0                   # default: 0.0
+    logit_std: 1.0                    # default: 1.0
+    mode_scale: 1.0                   # default: 1.0
+    precondition_outputs: false       # default: false
+    moba_config: {}                   # default: {}
+    enable_gradient_checkpointing_type: full  # default: null ("full" or null)
+
+  # --- training top-level [TYPED] ---
+  dit_precision: fp32                 # default: "fp32" (master weight precision)
+  # model_path: ...                   # default: "" (auto-derived from models.student.init_from)
+
+# ------------------------------------------------------------------------------
+# callbacks: [FREE]
+#
+# Each callback is instantiated via _target_(*, **kwargs).
+# The callback name (e.g. "grad_clip") is arbitrary — only _target_ matters.
+# training_config is injected automatically (not from YAML).
+# ------------------------------------------------------------------------------
+callbacks:
+
+  # --- GradNormClipCallback ---
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback  # optional if using default registry
+    max_grad_norm: 1.0                # default: 0.0 (0.0 = disabled)
+    log_grad_norms: false             # default: false
+
+  # --- EMACallback ---
+  # ema:
+  #   _target_: fastvideo.train.callbacks.ema.EMACallback
+  #   decay: 0.9999                   # default: 0.9999
+  #   start_iter: 0                   # default: 0
+
+  # --- ValidationCallback ---
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback  # optional if using default registry
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline  # required
+    dataset_file: path/to/validation.json     # required
+    every_steps: 100                  # default: 100
+    sampling_steps: [4]               # default: [40]
+    guidance_scale: 5.0               # default: null (uses model default)
+    num_frames: null                  # default: null (derived from training.data)
+    output_dir: null                  # default: null (falls back to training.checkpoint.output_dir)
+    sampling_timesteps: null          # default: null (explicit timestep list for streaming)
+
+# ------------------------------------------------------------------------------
+# pipeline: [RESOLVED] -> PipelineConfig
+#
+# Parsed by PipelineConfig.from_kwargs(). Most fields are auto-populated from
+# the model's config files (vae_config, dit_config, text_encoder_configs, etc.).
+# Only scalar overrides are typically needed here.
+# ------------------------------------------------------------------------------
+pipeline:
+  flow_shift: 3                       # default: null (model-specific)
+  # flow_shift_sr: null               # default: null (super-resolution shift)
+  # embedded_cfg_scale: 6.0           # default: 6.0
+  # is_causal: false                  # default: false
+  # vae_tiling: true                  # default: true
+  # vae_sp: true                      # default: true
+  # disable_autocast: false           # default: false

examples/train/configs/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml

@@ -0,0 +1,80 @@
+# V3 config: Wan 2.1 T2V 1.3B finetune with VSA (phase 3.4, 0.9 sparsity).
+#
+# Uses _target_-based instantiation — each model role is an independent
+# class instance; the method class is resolved directly from the YAML.
+
+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: vsa
+
+training:
+  distributed:
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 4
+    hsdp_shard_dim: 1
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.1
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 1.0e-6
+    betas: [0.9, 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/phase3.4_wan2.1_finetune_vsa_0.9_v3
+    training_state_checkpointing_steps: 100
+    checkpoints_total_limit: 3
+    resume_from_checkpoint: latest
+
+  tracker:
+    project_name: distillation_wan_r
+    run_name: phase3.4_wan_finetune_vsa_0.9_v3
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+  vsa:
+    sparsity: 0.9
+    decay_rate: 0.03
+    decay_interval_steps: 1
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  ema:
+    _target_: fastvideo.train.callbacks.ema.EMACallback
+    decay: 0.9999
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_pipeline.WanPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
+    every_steps: 50
+    sampling_steps: [50]
+    guidance_scale: 5.0
+
+pipeline:
+  flow_shift: 3

examples/train/configs/self_forcing_wan_causal_t2v_1.3B.yaml

@@ -0,0 +1,98 @@
+# Self-Forcing distillation: Wan 2.1 T2V 1.3B Causal
+#
+# - Teacher: frozen pretrained Wan 2.1 T2V 1.3B
+# - Student: trainable causal Wan model
+# - Critic:  trainable, initialized from pretrained weights
+# - Training: streaming rollout with SDE sampling
+
+models:
+  student:
+    _target_: fastvideo.train.models.wan.WanCausalModel
+    # init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers
+    init_from: /mnt/weka/home/hao.zhang/alex/weights/WanOdeInit
+    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.self_forcing.SelfForcingMethod
+  rollout_mode: simulate
+  generator_update_interval: 5
+  real_score_guidance_scale: 4.5
+  dmd_denoising_steps: [1000, 750, 500, 250]
+
+  # Self-forcing specific
+  chunk_size: 3
+  student_sample_type: sde
+  context_noise: 0.0
+  enable_gradient_in_rollout: true
+  start_gradient_frame: 0
+
+  # Critic optimizer
+  fake_score_learning_rate: 8.0e-6
+  fake_score_betas: [0.0, 0.999]
+  fake_score_lr_scheduler: constant
+
+training:
+  distributed:
+    num_gpus: 8
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 8
+
+  data:
+    data_path: data/Wan-Syn_77x448x832_600k
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    training_cfg_rate: 0.0
+    seed: 1000
+    num_latent_t: 20
+    num_height: 448
+    num_width: 832
+    num_frames: 77
+
+  optimizer:
+    learning_rate: 1e-5
+    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_causal_self_forcing
+    training_state_checkpointing_steps: 1000
+    checkpoints_total_limit: 3
+
+  tracker:
+    project_name: distillation_wan_r
+    run_name: wan2.1_causal_self_forcing
+
+  model:
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    max_grad_norm: 1.0
+  validation:
+    pipeline_target: fastvideo.pipelines.basic.wan.wan_causal_dmd_pipeline.WanCausalDMDPipeline
+    dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json
+    every_steps: 50
+    sampling_steps: [4]
+    sampling_timesteps: [1000, 750, 500, 250]
+    guidance_scale: 6.0
+
+pipeline:
+  flow_shift: 5

examples/train/run.sh

@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+# Launch training from a YAML config.
+#
+# Usage:
+#   bash examples/train/run.sh <config.yaml> [--dotted.key value ...]
+#
+# Examples:
+#   bash examples/train/run.sh examples/train/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml
+#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml --dry-run
+#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml \
+#       --training.distributed.num_gpus 4 \
+#       --training.optimizer.learning_rate 1e-5
+#   bash examples/train/run.sh examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml \
+#       --training.checkpoint.resume_from_checkpoint outputs/my_run/checkpoint-1000
+#
+# Logs are written to logs/<config_name>_<timestamp>.log (and also printed to stdout).
+
+set -euo pipefail
+
+CONFIG="${1:?Usage: $0 <config.yaml> [extra flags...]}"
+shift
+
+# ── GPU / node settings ──────────────────────────────────────────
+NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L 2>/dev/null | wc -l)}"
+NUM_GPUS="${NUM_GPUS:-1}"
+NNODES="${NNODES:-1}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-29501}"
+export TOKENIZERS_PARALLELISM=false
+# ── W&B ──────────────────────────────────────────────────────────
+export WANDB_API_KEY="${WANDB_API_KEY:-}"
+export WANDB_MODE="${WANDB_MODE:-online}"
+
+# ── Log file ─────────────────────────────────────────────────────
+CONFIG_NAME="$(basename "${CONFIG}" .yaml)"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+LOG_DIR="${LOG_DIR:-examples/train}"
+mkdir -p "${LOG_DIR}"
+LOG_FILE="${LOG_DIR}/${CONFIG_NAME}_${TIMESTAMP}.log"
+
+echo "=== Train Training ==="
+echo "Config:      ${CONFIG}"
+echo "Num GPUs:    ${NUM_GPUS}"
+echo "Num Nodes:   ${NNODES}"
+echo "Node Rank:   ${NODE_RANK}"
+echo "Master:      ${MASTER_ADDR}:${MASTER_PORT}"
+echo "Extra args:  $*"
+echo "Log file:    ${LOG_FILE}"
+echo "=============================="
+
+torchrun \
+    --nnodes "${NNODES}" \
+    --node_rank "${NODE_RANK}" \
+    --nproc_per_node "${NUM_GPUS}" \
+    --master_addr "${MASTER_ADDR}" \
+    --master_port "${MASTER_PORT}" \
+    fastvideo/train/entrypoint/train.py \
+    --config "${CONFIG}" \
+    "$@" \
+    2>&1 | tee "${LOG_FILE}"

examples/train/run_slurm.sh

@@ -0,0 +1,114 @@
+#!/usr/bin/env bash
+# Submit a multi-node Slurm training job.
+#
+# Usage:
+#   bash examples/train/run_slurm.sh <config.yaml> <num_nodes> [--dotted.key value ...]
+#
+# Examples:
+#   bash examples/train/run_slurm.sh examples/train/configs/example.yaml 2
+#   bash examples/train/run_slurm.sh examples/train/configs/distill_wan2.1_t2v_1.3B_dmd2.yaml 4 \
+#       --training.optimizer.learning_rate 1e-5
+#   bash examples/train/run_slurm.sh examples/train/configs/example.yaml 8 \
+#       --training.checkpoint.resume_from_checkpoint outputs/my_run/checkpoint-1000
+#
+# Environment variables (override defaults):
+#   PARTITION       Slurm partition           (default: main)
+#   NUM_GPUS        GPUs per node             (default: 8)
+#   CPUS_PER_TASK   CPUs per task             (default: 128)
+#   MEM             Memory per node           (default: 1440G)
+#   JOB_NAME        Slurm job name            (default: derived from config)
+#   OUTPUT_DIR      Directory for slurm logs  (default: slurm_logs)
+#   MASTER_PORT     Rendezvous port           (default: 29500)
+#   EXCLUDE         Nodes to exclude          (default: "")
+#   WANDB_API_KEY   W&B API key               (default: "")
+#   WANDB_MODE      W&B mode                  (default: online)
+
+set -euo pipefail
+
+CONFIG="${1:?Usage: $0 <config.yaml> <num_nodes> [extra flags...]}"
+NUM_NODES="${2:?Usage: $0 <config.yaml> <num_nodes> [extra flags...]}"
+shift 2
+
+# ── Defaults ──────────────────────────────────────────────────────
+PARTITION="${PARTITION:-main}"
+NUM_GPUS="${NUM_GPUS:-8}"
+CPUS_PER_TASK="${CPUS_PER_TASK:-128}"
+MEM="${MEM:-1440G}"
+MASTER_PORT="${MASTER_PORT:-29500}"
+EXCLUDE="${EXCLUDE:-}"
+WANDB_API_KEY="${WANDB_API_KEY:-}"
+WANDB_MODE="${WANDB_MODE:-online}"
+
+TOTAL_GPUS=$(( NUM_NODES * NUM_GPUS ))
+
+CONFIG_NAME="$(basename "${CONFIG}" .yaml)"
+JOB_NAME="${JOB_NAME:-${CONFIG_NAME}}"
+OUTPUT_DIR="${OUTPUT_DIR:-logs/slurm}"
+mkdir -p "${OUTPUT_DIR}"
+
+# ── Build sbatch args ─────────────────────────────────────────────
+SBATCH_ARGS=(
+    --job-name="${JOB_NAME}"
+    --partition="${PARTITION}"
+    --nodes="${NUM_NODES}"
+    --ntasks="${NUM_NODES}"
+    --ntasks-per-node=1
+    --gres="gpu:${NUM_GPUS}"
+    --cpus-per-task="${CPUS_PER_TASK}"
+    --mem="${MEM}"
+    --output="${OUTPUT_DIR}/${JOB_NAME}_%j.out"
+    --error="${OUTPUT_DIR}/${JOB_NAME}_%j.err"
+    --exclusive
+)
+
+if [[ -n "${EXCLUDE}" ]]; then
+    SBATCH_ARGS+=(--exclude="${EXCLUDE}")
+fi
+
+# ── Collect extra overrides for the training script ───────────────
+EXTRA_ARGS=("$@")
+
+echo "=== Slurm Training Submission ==="
+echo "Config:      ${CONFIG}"
+echo "Nodes:       ${NUM_NODES}"
+echo "GPUs/node:   ${NUM_GPUS}"
+echo "Total GPUs:  ${TOTAL_GPUS}"
+echo "Partition:   ${PARTITION}"
+echo "Job name:    ${JOB_NAME}"
+echo "Extra args:  ${EXTRA_ARGS[*]:-}"
+echo "================================="
+
+# ── Submit ────────────────────────────────────────────────────────
+sbatch "${SBATCH_ARGS[@]}" <<EOF
+#!/bin/bash
+set -e -x
+
+# ── Environment ───────────────────────────────────────────────────
+export NCCL_P2P_DISABLE=1
+export TORCH_NCCL_ENABLE_MONITORING=0
+export TRITON_CACHE_DIR=/tmp/triton_cache_\${SLURM_PROCID}
+export TOKENIZERS_PARALLELISM=false
+export WANDB_API_KEY="${WANDB_API_KEY}"
+export WANDB_MODE="${WANDB_MODE}"
+
+# ── Rendezvous ────────────────────────────────────────────────────
+export MASTER_PORT=${MASTER_PORT}
+nodes=( \$(scontrol show hostnames \$SLURM_JOB_NODELIST) )
+export MASTER_ADDR=\${nodes[0]}
+export NODE_RANK=\$SLURM_PROCID
+
+echo "MASTER_ADDR: \$MASTER_ADDR"
+echo "NODE_RANK:   \$NODE_RANK"
+
+# ── Launch ────────────────────────────────────────────────────────
+srun torchrun \\
+    --nnodes \$SLURM_JOB_NUM_NODES \\
+    --nproc_per_node ${NUM_GPUS} \\
+    --node_rank \$SLURM_PROCID \\
+    --rdzv_backend=c10d \\
+    --rdzv_endpoint="\$MASTER_ADDR:\$MASTER_PORT" \\
+    fastvideo/train/entrypoint/train.py \\
+    --config ${CONFIG} \\
+    --training.distributed.num_gpus ${TOTAL_GPUS} \\
+    ${EXTRA_ARGS[*]:-}
+EOF

fastvideo/dataset/parquet_dataset_map_style.py

@@ -343,11 +343,16 @@ class LatentsParquetMapStyleDataset(Dataset):
             for idx in indices
         ]
 
+        # Inject sample indices for deterministic CFG dropout
+        # that is reproducible across checkpoint resume.
+        for row, idx in zip(rows, indices):
+            row["_sample_index"] = idx
+
         batch = collate_rows_from_parquet_schema(rows,
                                                  self.parquet_schema,
                                                  self.text_padding_length,
                                                  cfg_rate=self.cfg_rate,
-                                                 rng=self.rng)
+                                                 seed=self.seed)
         return batch
 
     def __len__(self):

fastvideo/dataset/utils.py

@@ -100,7 +100,8 @@ def collate_rows_from_parquet_schema(rows,
                                      parquet_schema,
                                      text_padding_length,
                                      cfg_rate=0.0,
-                                     rng=None) -> dict[str, Any]:
+                                     rng=None,
+                                     seed=0) -> dict[str, Any]:
     """
     Collate rows from parquet files based on the provided schema.
     Dynamically processes tensor fields based on schema and returns batched data.
@@ -149,13 +150,30 @@ def collate_rows_from_parquet_schema(rows,
                 if len(bytes_data) == 0:
                     tensor = torch.zeros(0, dtype=torch.bfloat16)
                 else:
-                    # Convert bytes to tensor using float32 as default
-                    if tensor_name == 'text_embedding' and (rng.random(
-                    ) if rng else random.random()) < cfg_rate:
-                        data = np.zeros((512, 4096), dtype=np.float32)
+                    # Deterministic per-sample CFG dropout
+                    # using sample index (resume-safe).
+                    drop = False
+                    if (tensor_name == 'text_embedding'
+                            and cfg_rate > 0):
+                        sample_idx = row.get(
+                            "_sample_index")
+                        if sample_idx is not None:
+                            drop = (random.Random(
+                                seed ^ sample_idx
+                            ).random() < cfg_rate)
+                        else:
+                            drop = ((rng.random()
+                                     if rng else
+                                     random.random())
+                                    < cfg_rate)
+                    if drop:
+                        data = np.zeros(
+                            (512, 4096), dtype=np.float32)
                     else:
                         data = np.frombuffer(
-                            bytes_data, dtype=np.float32).reshape(shape).copy()
+                            bytes_data,
+                            dtype=np.float32,
+                        ).reshape(shape).copy()
                     tensor = torch.from_numpy(data)
                     # if len(data.shape) == 3:
                     #     B, L, D = tensor.shape

fastvideo/pipelines/basic/wan/wan_causal_pipeline.py

@@ -30,21 +30,23 @@ class WanCausalPipeline(LoRAPipeline, ComposedPipelineBase):
     """Wan causal pipeline with standard multi-step denoising."""
 
     _required_config_modules = [
-        "text_encoder", "tokenizer", "vae",
-        "transformer", "scheduler",
+        "text_encoder",
+        "tokenizer",
+        "vae",
+        "transformer",
+        "scheduler",
     ]
 
     def initialize_pipeline(
-        self, fastvideo_args: FastVideoArgs,
+        self,
+        fastvideo_args: FastVideoArgs,
     ):
-        self.modules["scheduler"] = (
-            FlowUniPCMultistepScheduler(
-                shift=fastvideo_args.pipeline_config.flow_shift,
-            )
-        )
+        self.modules["scheduler"] = (FlowUniPCMultistepScheduler(
+            shift=fastvideo_args.pipeline_config.flow_shift, ))
 
     def create_pipeline_stages(
-        self, fastvideo_args: FastVideoArgs,
+        self,
+        fastvideo_args: FastVideoArgs,
     ) -> None:
         self.add_stage(
             stage_name="input_validation_stage",
@@ -54,10 +56,8 @@ class WanCausalPipeline(LoRAPipeline, ComposedPipelineBase):
         self.add_stage(
             stage_name="prompt_encoding_stage",
             stage=TextEncodingStage(
-                text_encoders=[
-                    self.get_module("text_encoder")],
-                tokenizers=[
-                    self.get_module("tokenizer")],
+                text_encoders=[self.get_module("text_encoder")],
+                tokenizers=[self.get_module("tokenizer")],
             ),
         )
 
@@ -69,16 +69,14 @@ class WanCausalPipeline(LoRAPipeline, ComposedPipelineBase):
         self.add_stage(
             stage_name="timestep_preparation_stage",
             stage=TimestepPreparationStage(
-                scheduler=self.get_module("scheduler"),
-            ),
+                scheduler=self.get_module("scheduler"), ),
         )
 
         self.add_stage(
             stage_name="latent_preparation_stage",
             stage=LatentPreparationStage(
                 scheduler=self.get_module("scheduler"),
-                transformer=self.get_module(
-                    "transformer", None),
+                transformer=self.get_module("transformer", None),
             ),
         )
 
@@ -92,8 +90,7 @@ class WanCausalPipeline(LoRAPipeline, ComposedPipelineBase):
 
         self.add_stage(
             stage_name="decoding_stage",
-            stage=DecodingStage(
-                vae=self.get_module("vae")),
+            stage=DecodingStage(vae=self.get_module("vae")),
         )
 
 

fastvideo/pipelines/stages/__init__.py

@@ -7,8 +7,8 @@ complete diffusion pipelines.
 """
 
 from fastvideo.pipelines.stages.base import PipelineStage
-from fastvideo.pipelines.stages.causal_denoising import (
-    CausalDMDDenosingStage, CausalDenoisingStage)
+from fastvideo.pipelines.stages.causal_denoising import (CausalDMDDenosingStage,
+                                                         CausalDenoisingStage)
 from fastvideo.pipelines.stages.conditioning import ConditioningStage
 from fastvideo.pipelines.stages.decoding import DecodingStage
 from fastvideo.pipelines.stages.denoising import (

fastvideo/pipelines/stages/causal_denoising.py

@@ -500,31 +500,22 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
         fastvideo_args: FastVideoArgs,
     ) -> ForwardBatch:
         target_dtype = torch.bfloat16
-        autocast_enabled = (
-            target_dtype != torch.float32
-            and not fastvideo_args.disable_autocast
-        )
+        autocast_enabled = (target_dtype != torch.float32
+                            and not fastvideo_args.disable_autocast)
 
-        latent_seq_length = (
-            batch.latents.shape[-1] * batch.latents.shape[-2]
-        )
-        patch_ratio = (
-            self.transformer.config.arch_config.patch_size[-1]
-            * self.transformer.config.arch_config.patch_size[-2]
-        )
+        latent_seq_length = (batch.latents.shape[-1] * batch.latents.shape[-2])
+        patch_ratio = (self.transformer.config.arch_config.patch_size[-1] *
+                       self.transformer.config.arch_config.patch_size[-2])
         self.frame_seq_length = latent_seq_length // patch_ratio
 
         pos_cond_kwargs = self.prepare_extra_func_kwargs(
             self.transformer.forward,
             {
-                "encoder_attention_mask":
-                    batch.prompt_attention_mask,
+                "encoder_attention_mask": batch.prompt_attention_mask,
             },
         )
 
-        assert batch.latents is not None, (
-            "latents must be provided"
-        )
+        assert batch.latents is not None, ("latents must be provided")
         latents = batch.latents  # [B, C, T, H, W]
         b, c, t, h, w = latents.shape
         prompt_embeds = batch.prompt_embeds
@@ -540,21 +531,16 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
         )
         crossattn_cache = self._initialize_crossattn_cache(
             batch_size=b,
-            max_text_len=(
-                fastvideo_args.pipeline_config
-                .text_encoder_configs[0]
-                .arch_config.text_len
-            ),
+            max_text_len=(fastvideo_args.pipeline_config.
+                          text_encoder_configs[0].arch_config.text_len),
             dtype=target_dtype,
             device=latents.device,
         )
 
         # Determine block sizes
         if t % self.num_frames_per_block != 0:
-            raise ValueError(
-                "num_frames must be divisible by "
-                "num_frames_per_block"
-            )
+            raise ValueError("num_frames must be divisible by "
+                             "num_frames_per_block")
         num_blocks = t // self.num_frames_per_block
         block_sizes = [self.num_frames_per_block] * num_blocks
         start_index = 0
@@ -562,17 +548,19 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
 
         context_noise = getattr(
             fastvideo_args.pipeline_config,
-            "context_noise", 0,
+            "context_noise",
+            0,
         )
 
-        with self.progress_bar(
-            total=len(block_sizes) * num_inference_steps
-        ) as progress_bar:
+        with self.progress_bar(total=len(block_sizes) *
+                               num_inference_steps) as progress_bar:
             for current_num_frames in block_sizes:
                 current_latents = latents[
-                    :, :,
+                    :,
+                    :,
                     start_index:start_index + current_num_frames,
-                    :, :,
+                    :,
+                    :,
                 ]
 
                 # Reset scheduler per block so its multi-step
@@ -584,9 +572,7 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
                 timesteps = self.scheduler.timesteps
 
                 for i, t_cur in enumerate(timesteps):
-                    latent_model_input = current_latents.to(
-                        target_dtype
-                    )
+                    latent_model_input = current_latents.to(target_dtype)
                     t_expanded = t_cur * torch.ones(
                         (b, 1),
                         device=latents.device,
@@ -594,16 +580,16 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
                     )
 
                     with (
-                        torch.autocast(
-                            device_type="cuda",
-                            dtype=target_dtype,
-                            enabled=autocast_enabled,
-                        ),
-                        set_forward_context(
-                            current_timestep=i,
-                            attn_metadata=None,
-                            forward_batch=batch,
-                        ),
+                            torch.autocast(
+                                device_type="cuda",
+                                dtype=target_dtype,
+                                enabled=autocast_enabled,
+                            ),
+                            set_forward_context(
+                                current_timestep=i,
+                                attn_metadata=None,
+                                forward_batch=batch,
+                            ),
                     ):
                         # Transformer returns [B, C, T, H, W],
                         # permute to [B, T, C, H, W] then flatten
@@ -614,10 +600,8 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
                             t_expanded,
                             kv_cache=kv_cache,
                             crossattn_cache=crossattn_cache,
-                            current_start=(
-                                (pos_start_base + start_index)
-                                * self.frame_seq_length
-                            ),
+                            current_start=((pos_start_base + start_index) *
+                                           self.frame_seq_length),
                             start_frame=start_index,
                             **pos_cond_kwargs,
                         )
@@ -625,36 +609,29 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
                     # Flatten [B,C,T,H,W] -> [B*T,C,H,W] for
                     # scheduler, then unflatten back.
                     nf = current_num_frames
-                    noise_pred_flat = (
-                        noise_pred
-                        .permute(0, 2, 1, 3, 4)
-                        .flatten(0, 1)
-                    )
-                    latents_flat = (
-                        current_latents
-                        .permute(0, 2, 1, 3, 4)
-                        .flatten(0, 1)
-                    )
+                    noise_pred_flat = (noise_pred.permute(0, 2, 1, 3,
+                                                          4).flatten(0, 1))
+                    latents_flat = (current_latents.permute(0, 2, 1, 3,
+                                                            4).flatten(0, 1))
                     updated_flat = self.scheduler.step(
                         noise_pred_flat,
                         t_cur,
                         latents_flat,
                         return_dict=False,
                     )[0]
-                    current_latents = (
-                        updated_flat
-                        .unflatten(0, (b, nf))
-                        .permute(0, 2, 1, 3, 4)
-                    )
+                    current_latents = (updated_flat.unflatten(
+                        0, (b, nf)).permute(0, 2, 1, 3, 4))
 
                     if progress_bar is not None:
                         progress_bar.update()
 
                 # Write denoised block back
                 latents[
-                    :, :,
+                    :,
+                    :,
                     start_index:start_index + current_num_frames,
-                    :, :,
+                    :,
+                    :,
                 ] = current_latents
 
                 # Update KV cache with clean context
@@ -665,16 +642,16 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
                 ) * int(context_noise)
                 context_bcthw = current_latents.to(target_dtype)
                 with (
-                    torch.autocast(
-                        device_type="cuda",
-                        dtype=target_dtype,
-                        enabled=autocast_enabled,
-                    ),
-                    set_forward_context(
-                        current_timestep=0,
-                        attn_metadata=None,
-                        forward_batch=batch,
-                    ),
+                        torch.autocast(
+                            device_type="cuda",
+                            dtype=target_dtype,
+                            enabled=autocast_enabled,
+                        ),
+                        set_forward_context(
+                            current_timestep=0,
+                            attn_metadata=None,
+                            forward_batch=batch,
+                        ),
                 ):
                     self.transformer(
                         context_bcthw,
@@ -682,10 +659,8 @@ class CausalDenoisingStage(CausalDMDDenosingStage):
                         t_context.unsqueeze(1),
                         kv_cache=kv_cache,
                         crossattn_cache=crossattn_cache,
-                        current_start=(
-                            (pos_start_base + start_index)
-                            * self.frame_seq_length
-                        ),
+                        current_start=((pos_start_base + start_index) *
+                                       self.frame_seq_length),
                         start_frame=start_index,
                         **pos_cond_kwargs,
                     )

fastvideo/train/.style.yapf

@@ -0,0 +1,3 @@
+[style]
+based_on_style = pep8
+column_limit = 120

fastvideo/train/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.train.trainer import Trainer
+
+__all__ = [
+    "Trainer",
+]

fastvideo/train/callbacks/__init__.py

@@ -0,0 +1,20 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.train.callbacks.callback import (
+    Callback,
+    CallbackDict,
+)
+from fastvideo.train.callbacks.ema import (
+    EMACallback, )
+from fastvideo.train.callbacks.grad_clip import (
+    GradNormClipCallback, )
+from fastvideo.train.callbacks.validation import (
+    ValidationCallback, )
+
+__all__ = [
+    "Callback",
+    "CallbackDict",
+    "EMACallback",
+    "GradNormClipCallback",
+    "ValidationCallback",
+]

fastvideo/train/callbacks/callback.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Callback base class and CallbackDict manager.
+
+Adapted from FastGen's callback pattern to FastVideo's types.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Any, TYPE_CHECKING
+
+from fastvideo.logger import init_logger
+from fastvideo.train.utils.instantiate import instantiate
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.base import TrainingMethod
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+
+logger = init_logger(__name__)
+
+# Well-known callback names that don't need ``_target_`` in YAML.
+_BUILTIN_CALLBACKS: dict[str, str] = {
+    "grad_clip": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback",
+    "validation": "fastvideo.train.callbacks.validation.ValidationCallback",
+    "ema": "fastvideo.train.callbacks.ema.EMACallback",
+}
+
+
+class Callback:
+    """Base callback with no-op hooks.
+
+    Subclasses override whichever hooks they need.  The
+    ``training_config`` and ``method`` attributes are set by
+    ``CallbackDict`` after instantiation.
+    """
+
+    training_config: TrainingConfig
+    method: TrainingMethod
+    _callback_dict: CallbackDict | None
+
+    def on_train_start(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        pass
+
+    def on_training_step_end(
+        self,
+        method: TrainingMethod,
+        loss_dict: dict[str, Any],
+        iteration: int = 0,
+    ) -> None:
+        pass
+
+    def on_before_optimizer_step(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        pass
+
+    def on_validation_begin(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        pass
+
+    def on_validation_end(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        pass
+
+    def on_train_end(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        pass
+
+    def state_dict(self) -> dict[str, Any]:
+        return {}
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Any],
+    ) -> None:
+        pass
+
+
+class CallbackDict:
+    """Manages a collection of named callbacks.
+
+    Instantiates each callback from its ``_target_`` config and
+    dispatches hook calls to all registered callbacks.
+    """
+
+    def __init__(
+        self,
+        callback_configs: dict[str, dict[str, Any]],
+        training_config: TrainingConfig,
+    ) -> None:
+        self._callbacks: dict[str, Callback] = {}
+        if not callback_configs:
+            return
+        for name, cb_cfg in callback_configs.items():
+            cb_cfg = dict(cb_cfg)
+            if "_target_" not in cb_cfg:
+                if name in _BUILTIN_CALLBACKS:
+                    cb_cfg["_target_"] = (_BUILTIN_CALLBACKS[name])
+                else:
+                    logger.warning(
+                        "Callback %r is missing "
+                        "'_target_', skipping: %s",
+                        name,
+                        cb_cfg,
+                    )
+                    continue
+            logger.info(
+                "Instantiating callback %r: %s",
+                name,
+                cb_cfg,
+            )
+            cb = instantiate(cb_cfg)
+            if not isinstance(cb, Callback):
+                raise TypeError(f"Callback {name!r} resolved to "
+                                f"{type(cb).__name__}, expected a "
+                                f"Callback subclass.")
+            cb.training_config = training_config
+            cb._callback_dict = self
+            self._callbacks[name] = cb
+
+    def __getattr__(
+        self,
+        method_name: str,
+    ) -> Callable[..., Any]:
+        if method_name.startswith("_"):
+            raise AttributeError(method_name)
+
+        if method_name == "state_dict":
+
+            def _state_dict() -> dict[str, Any]:
+                return {n: cb.state_dict() for n, cb in self._callbacks.items()}
+
+            return _state_dict
+
+        if method_name == "load_state_dict":
+
+            def _load_state_dict(state_dict: dict[str, Any], ) -> None:
+                for n, cb in self._callbacks.items():
+                    if n in state_dict:
+                        cb.load_state_dict(state_dict[n])
+                    else:
+                        logger.warning(
+                            "Callback %r not found in "
+                            "checkpoint.",
+                            n,
+                        )
+
+            return _load_state_dict
+
+        def _dispatch(*args: Any, **kwargs: Any) -> None:
+            for cb in self._callbacks.values():
+                fn = getattr(cb, method_name, None)
+                if fn is None:
+                    continue
+                if not callable(fn):
+                    continue
+                fn(*args, **kwargs)
+
+        return _dispatch

fastvideo/train/callbacks/ema.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+"""EMA (Exponential Moving Average) callback.
+
+Owns the full EMA lifecycle: creation, per-step updates, weight
+swapping for validation, and checkpoint state.  All EMA config
+lives under ``callbacks.ema`` in the YAML file.
+"""
+
+from __future__ import annotations
+
+import contextlib
+from collections.abc import Generator
+from typing import Any, TYPE_CHECKING
+
+import torch
+
+from fastvideo.logger import init_logger
+from fastvideo.train.callbacks.callback import Callback
+from fastvideo.training.training_utils import EMA_FSDP
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.base import TrainingMethod
+
+logger = init_logger(__name__)
+
+
+class EMACallback(Callback):
+    """Manage EMA shadow weights for the student transformer.
+
+    All configuration lives in the YAML ``callbacks.ema`` section:
+
+    .. code-block:: yaml
+
+        callbacks:
+          ema:
+            decay: 0.9999
+            start_iter: 0
+
+    The callback creates an ``EMA_FSDP`` instance at train start,
+    updates it after each optimizer step, and exposes an
+    ``ema_context()`` context manager for temporarily swapping
+    EMA weights into the live model (used by validation).
+    """
+
+    def __init__(
+        self,
+        *,
+        decay: float = 0.9999,
+        start_iter: int = 0,
+    ) -> None:
+        self._decay = float(decay)
+        self._start_iter = int(start_iter)
+        self._ema_started = False
+        self.student_ema: EMA_FSDP | None = None
+
+    # ----------------------------------------------------------
+    # Hooks
+    # ----------------------------------------------------------
+
+    def on_train_start(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        student = getattr(method, "student", None)
+        if student is None or student.transformer is None:
+            raise ValueError("No student transformer found on method, cannot initialize EMA")
+
+        logger.info(
+            "Initializing EMA (local_shard) with "
+            "decay=%s from student transformer",
+            self._decay,
+        )
+        self.student_ema = EMA_FSDP(
+            student.transformer,
+            decay=self._decay,
+            mode="local_shard",
+        )
+        logger.info(
+            "EMA callback enabled (decay=%s, "
+            "start_iter=%d).",
+            self._decay,
+            self._start_iter,
+        )
+
+    def on_training_step_end(
+        self,
+        method: TrainingMethod,
+        loss_dict: dict[str, Any],
+        iteration: int = 0,
+    ) -> None:
+        if self.student_ema is None:
+            return
+
+        if iteration < self._start_iter:
+            return
+        if not self._ema_started:
+            logger.info(
+                "Starting EMA updates at iteration %d "
+                "(re-initializing shadow from current "
+                "model).",
+                iteration,
+            )
+            self.student_ema._init_shadow(method.student.transformer, )
+            self._ema_started = True
+
+        self.student_ema.update(method.student.transformer, )
+
+        tracker = getattr(method, "tracker", None)
+        if tracker is not None:
+            tracker.log(
+                {"ema/decay": self.student_ema.decay},
+                iteration,
+            )
+
+    # ----------------------------------------------------------
+    # EMA context manager
+    # ----------------------------------------------------------
+
+    @contextlib.contextmanager
+    def ema_context(
+        self,
+        transformer: torch.nn.Module,
+    ) -> Generator[torch.nn.Module, None, None]:
+        """Temporarily swap EMA weights into *transformer*.
+
+        If EMA is not active, yields the transformer unchanged.
+        """
+        if (self.student_ema is not None and self._ema_started):
+            with self.student_ema.apply_to_model(transformer, ):
+                yield transformer
+        else:
+            yield transformer
+
+    # ----------------------------------------------------------
+    # Checkpoint state
+    # ----------------------------------------------------------
+
+    def state_dict(self) -> dict[str, Any]:
+        if self.student_ema is None:
+            return {}
+        return {
+            "student_ema": self.student_ema.state_dict(),
+            "ema_started": self._ema_started,
+        }
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Any],
+    ) -> None:
+        ema_state = state_dict.get("student_ema")
+        if (ema_state is not None and self.student_ema is not None):
+            self.student_ema.load_state_dict(ema_state)
+        self._ema_started = bool(state_dict.get("ema_started", False), )

fastvideo/train/callbacks/grad_clip.py

@@ -0,0 +1,61 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Gradient norm clipping callback.
+
+Clips gradients on modules returned by
+``method.get_grad_clip_targets()`` before the optimizer step.
+Optionally logs per-module grad norms to the tracker.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from fastvideo.logger import init_logger
+from fastvideo.train.callbacks.callback import Callback
+from fastvideo.train.utils.optimizer import (
+    clip_grad_norm_if_needed, )
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.base import TrainingMethod
+
+logger = init_logger(__name__)
+
+
+class GradNormClipCallback(Callback):
+    """Clip gradient norms before the optimizer step.
+
+    ``max_grad_norm`` must be set explicitly in the callback
+    config (``callbacks.grad_clip.max_grad_norm``).
+    """
+
+    def __init__(
+        self,
+        *,
+        max_grad_norm: float = 0.0,
+        log_grad_norms: bool = False,
+    ) -> None:
+        self._max_grad_norm = float(max_grad_norm)
+        self._log_grad_norms = bool(log_grad_norms)
+
+    def on_before_optimizer_step(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        max_norm = self._max_grad_norm
+        if max_norm <= 0.0:
+            return
+
+        targets = method.get_grad_clip_targets(iteration)
+        tracker = getattr(method, "tracker", None)
+
+        for name, module in targets.items():
+            grad_norm = clip_grad_norm_if_needed(
+                module,
+                max_norm,
+            )
+            if (self._log_grad_norms and tracker is not None and grad_norm > 0.0):
+                tracker.log(
+                    {f"grad_norm/{name}": grad_norm},
+                    iteration,
+                )

fastvideo/train/callbacks/validation.py

@@ -0,0 +1,459 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Validation callback.
+
+All configuration is read from the YAML ``callbacks.validation``
+section.  The pipeline class is resolved from
+``pipeline_target``.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+from dataclasses import dataclass
+from typing import Any, TYPE_CHECKING
+
+import imageio
+import numpy as np
+import torch
+import torchvision
+from einops import rearrange
+from torch.utils.data import DataLoader
+
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.dataset.validation_dataset import (
+    ValidationDataset, )
+from fastvideo.distributed import (
+    get_sp_group,
+    get_world_group,
+)
+from fastvideo.logger import init_logger
+from fastvideo.pipelines import ForwardBatch
+from fastvideo.train.callbacks.callback import Callback
+from fastvideo.train.utils.instantiate import resolve_target
+from fastvideo.train.utils.moduleloader import (
+    make_inference_args, )
+from fastvideo.training.trackers import DummyTracker
+from fastvideo.utils import shallow_asdict
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.base import TrainingMethod
+
+logger = init_logger(__name__)
+
+
+@dataclass(slots=True)
+class _ValidationStepResult:
+    videos: list[list[np.ndarray]]
+    captions: list[str]
+
+
+class ValidationCallback(Callback):
+    """Generic validation callback driven entirely by YAML
+    config.
+
+    Works with any pipeline that follows the
+    ``PipelineCls.from_pretrained(...)`` + ``pipeline.forward()``
+    contract.
+    """
+
+    def __init__(
+        self,
+        *,
+        pipeline_target: str,
+        dataset_file: str,
+        every_steps: int = 100,
+        sampling_steps: list[int] | None = None,
+        guidance_scale: float | None = None,
+        num_frames: int | None = None,
+        output_dir: str | None = None,
+        sampling_timesteps: list[int] | None = None,
+        **pipeline_kwargs: Any,
+    ) -> None:
+        self.pipeline_target = str(pipeline_target)
+        self.dataset_file = str(dataset_file)
+        self.every_steps = int(every_steps)
+        self.sampling_steps = ([int(s) for s in sampling_steps] if sampling_steps else [40])
+        self.guidance_scale = (float(guidance_scale) if guidance_scale is not None else None)
+        self.num_frames = (int(num_frames) if num_frames is not None else None)
+        self.output_dir = (str(output_dir) if output_dir is not None else None)
+        self.sampling_timesteps = ([int(s) for s in sampling_timesteps] if sampling_timesteps is not None else None)
+        self.pipeline_kwargs = dict(pipeline_kwargs)
+
+        # Set after on_train_start.
+        self._pipeline: Any | None = None
+        self._pipeline_key: tuple[Any, ...] | None = None
+        self._sampling_param: SamplingParam | None = None
+        self.tracker: Any = DummyTracker()
+        self.validation_random_generator: (torch.Generator | None) = None
+        self.seed: int = 0
+
+    # ----------------------------------------------------------
+    # Callback hooks
+    # ----------------------------------------------------------
+
+    def on_train_start(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        self.method = method
+        tc = self.training_config
+
+        self.world_group = get_world_group()
+        self.sp_group = get_sp_group()
+        self.global_rank = self.world_group.rank
+        self.rank_in_sp_group = (self.sp_group.rank_in_group)
+        self.sp_world_size = self.sp_group.world_size
+
+        seed = tc.data.seed
+        if seed is None:
+            raise ValueError("training.data.seed must be set "
+                             "for validation")
+        self.seed = int(seed)
+        self.validation_random_generator = (torch.Generator(device="cpu").manual_seed(self.seed))
+
+        tracker = getattr(method, "tracker", None)
+        if tracker is not None:
+            self.tracker = tracker
+
+    def on_validation_begin(
+        self,
+        method: TrainingMethod,
+        iteration: int = 0,
+    ) -> None:
+        if self.every_steps <= 0:
+            return
+        if iteration % self.every_steps != 0:
+            return
+
+        self._run_validation(method, iteration)
+
+    # ----------------------------------------------------------
+    # Core validation logic
+    # ----------------------------------------------------------
+
+    def _run_validation(
+        self,
+        method: TrainingMethod,
+        step: int,
+    ) -> None:
+
+        transformer = method.student.transformer
+        # Look for an EMA callback to temporarily swap
+        # EMA weights during validation.
+        ema_cb = self._find_ema_callback()
+        if ema_cb is not None:
+            ctx = ema_cb.ema_context(transformer)
+        else:
+            ctx = contextlib.nullcontext(transformer)
+        with ctx as t:
+            self._run_validation_inner(
+                method,
+                step,
+                t,
+            )
+
+    def _find_ema_callback(self) -> Any | None:
+        """Find the EMA callback in the callback dict."""
+        from fastvideo.train.callbacks.ema import (
+            EMACallback, )
+
+        cb_dict = getattr(self, "_callback_dict", None)
+        if cb_dict is not None:
+            for cb in cb_dict._callbacks.values():
+                if isinstance(cb, EMACallback):
+                    return cb
+        return None
+
+    def _run_validation_inner(
+        self,
+        method: TrainingMethod,
+        step: int,
+        transformer: torch.nn.Module,
+    ) -> None:
+        tc = self.training_config
+        was_training = bool(getattr(transformer, "training", False))
+
+        output_dir = (self.output_dir or tc.checkpoint.output_dir)
+
+        try:
+            transformer.eval()
+            num_sp_groups = (self.world_group.world_size // self.sp_group.world_size)
+
+            for num_inference_steps in self.sampling_steps:
+                result = self._run_validation_for_steps(
+                    num_inference_steps,
+                    transformer=transformer,
+                )
+
+                if self.rank_in_sp_group != 0:
+                    continue
+
+                if self.global_rank == 0:
+                    all_videos = list(result.videos)
+                    all_captions = list(result.captions)
+                    for sp_idx in range(1, num_sp_groups):
+                        src = (sp_idx * self.sp_world_size)
+                        recv_v = (self.world_group.recv_object(src=src))
+                        recv_c = (self.world_group.recv_object(src=src))
+                        all_videos.extend(recv_v)
+                        all_captions.extend(recv_c)
+
+                    os.makedirs(
+                        output_dir,
+                        exist_ok=True,
+                    )
+                    video_filenames: list[str] = []
+                    sp = self._get_sampling_param()
+                    for i, video in enumerate(all_videos):
+                        fname = os.path.join(
+                            output_dir,
+                            f"validation_step_{step}"
+                            f"_inference_steps_"
+                            f"{num_inference_steps}"
+                            f"_video_{i}.mp4",
+                        )
+                        imageio.mimsave(
+                            fname,
+                            video,
+                            fps=sp.fps,
+                        )
+                        video_filenames.append(fname)
+
+                    video_logs = []
+                    for fname, cap in zip(
+                            video_filenames,
+                            all_captions,
+                            strict=True,
+                    ):
+                        art = self.tracker.video(
+                            fname,
+                            caption=cap,
+                        )
+                        if art is not None:
+                            video_logs.append(art)
+                    if video_logs:
+                        logs = {
+                            f"validation_videos_"
+                            f"{num_inference_steps}"
+                            f"_steps": video_logs
+                        }
+                        self.tracker.log_artifacts(
+                            logs,
+                            step,
+                        )
+                else:
+                    self.world_group.send_object(
+                        result.videos,
+                        dst=0,
+                    )
+                    self.world_group.send_object(
+                        result.captions,
+                        dst=0,
+                    )
+        finally:
+            if was_training:
+                transformer.train()
+
+    # ----------------------------------------------------------
+    # Pipeline management
+    # ----------------------------------------------------------
+
+    def _get_sampling_param(self) -> SamplingParam:
+        if self._sampling_param is None:
+            self._sampling_param = (SamplingParam.from_pretrained(self.training_config.model_path))
+        return self._sampling_param
+
+    def _get_pipeline(
+        self,
+        *,
+        transformer: torch.nn.Module,
+    ) -> Any:
+        key = (id(transformer), )
+        if (self._pipeline is not None and self._pipeline_key == key):
+            return self._pipeline
+
+        tc = self.training_config
+        PipelineCls = resolve_target(self.pipeline_target)
+        flow_shift = getattr(
+            tc.pipeline_config,
+            "flow_shift",
+            None,
+        )
+
+        kwargs: dict[str, Any] = {
+            "inference_mode": True,
+            "loaded_modules": {
+                "transformer": transformer,
+            },
+            "tp_size": tc.distributed.tp_size,
+            "sp_size": tc.distributed.sp_size,
+            "num_gpus": tc.distributed.num_gpus,
+            "pin_cpu_memory": (tc.distributed.pin_cpu_memory),
+            "dit_cpu_offload": True,
+        }
+        if flow_shift is not None:
+            kwargs["flow_shift"] = float(flow_shift)
+
+        self._pipeline = PipelineCls.from_pretrained(
+            tc.model_path,
+            **kwargs,
+        )
+        self._pipeline_key = key
+        return self._pipeline
+
+    # ----------------------------------------------------------
+    # Batch preparation
+    # ----------------------------------------------------------
+
+    def _prepare_validation_batch(
+        self,
+        sampling_param: SamplingParam,
+        validation_batch: dict[str, Any],
+        num_inference_steps: int,
+    ) -> ForwardBatch:
+        tc = self.training_config
+
+        sampling_param.prompt = validation_batch["prompt"]
+        sampling_param.height = tc.data.num_height
+        sampling_param.width = tc.data.num_width
+        sampling_param.num_inference_steps = int(num_inference_steps)
+        sampling_param.data_type = "video"
+        if self.guidance_scale is not None:
+            sampling_param.guidance_scale = float(self.guidance_scale)
+        sampling_param.seed = self.seed
+
+        # image_path for I2V pipelines.
+        img_path = (validation_batch.get("image_path") or validation_batch.get("video_path"))
+        if img_path is not None and (img_path.startswith("http") or os.path.isfile(img_path)):
+            sampling_param.image_path = img_path
+
+        temporal_compression_factor = int(
+            tc.pipeline_config.vae_config.arch_config.temporal_compression_ratio  # type: ignore[union-attr]
+        )
+        default_num_frames = ((tc.data.num_latent_t - 1) * temporal_compression_factor + 1)
+        if self.num_frames is not None:
+            sampling_param.num_frames = int(self.num_frames)
+        else:
+            sampling_param.num_frames = int(default_num_frames)
+
+        latents_size = [
+            (sampling_param.num_frames - 1) // 4 + 1,
+            sampling_param.height // 8,
+            sampling_param.width // 8,
+        ]
+        n_tokens = (latents_size[0] * latents_size[1] * latents_size[2])
+
+        sampling_timesteps_tensor = (torch.tensor(
+            [int(s) for s in self.sampling_timesteps],
+            dtype=torch.long,
+        ) if self.sampling_timesteps is not None else None)
+
+        inference_args = make_inference_args(
+            tc,
+            model_path=tc.model_path,
+        )
+
+        batch = ForwardBatch(
+            **shallow_asdict(sampling_param),
+            latents=None,
+            generator=self.validation_random_generator,
+            n_tokens=n_tokens,
+            eta=0.0,
+            VSA_sparsity=tc.vsa.sparsity,
+            timesteps=sampling_timesteps_tensor,
+        )
+        batch._inference_args = inference_args  # type: ignore[attr-defined]
+
+        # Conditionally set I2V fields.
+        if ("image" in validation_batch and validation_batch["image"] is not None):
+            batch.pil_image = validation_batch["image"]
+
+        return batch
+
+    # ----------------------------------------------------------
+    # Validation loop
+    # ----------------------------------------------------------
+
+    def _run_validation_for_steps(
+        self,
+        num_inference_steps: int,
+        *,
+        transformer: torch.nn.Module,
+    ) -> _ValidationStepResult:
+        tc = self.training_config
+        pipeline = self._get_pipeline(transformer=transformer, )
+        sampling_param = self._get_sampling_param()
+
+        dataset = ValidationDataset(self.dataset_file)
+        dataloader = DataLoader(
+            dataset,
+            batch_size=None,
+            num_workers=0,
+        )
+
+        inference_args = make_inference_args(
+            tc,
+            model_path=tc.model_path,
+        )
+
+        videos: list[list[np.ndarray]] = []
+        captions: list[str] = []
+
+        for validation_batch in dataloader:
+            batch = self._prepare_validation_batch(
+                sampling_param,
+                validation_batch,
+                num_inference_steps,
+            )
+
+            assert (batch.prompt is not None and isinstance(batch.prompt, str))
+            captions.append(batch.prompt)
+
+            with torch.no_grad():
+                output_batch = pipeline.forward(
+                    batch,
+                    inference_args,
+                )
+
+            samples = output_batch.output.cpu()
+            if self.rank_in_sp_group != 0:
+                continue
+
+            video = rearrange(
+                samples,
+                "b c t h w -> t b c h w",
+            )
+            frames: list[np.ndarray] = []
+            for x in video:
+                x = torchvision.utils.make_grid(
+                    x,
+                    nrow=6,
+                )
+                x = (x.transpose(0, 1).transpose(1, 2).squeeze(-1))
+                frames.append((x * 255).numpy().astype(np.uint8))
+            videos.append(frames)
+
+        return _ValidationStepResult(
+            videos=videos,
+            captions=captions,
+        )
+
+    # ----------------------------------------------------------
+    # State management
+    # ----------------------------------------------------------
+
+    def state_dict(self) -> dict[str, Any]:
+        state: dict[str, Any] = {}
+        if self.validation_random_generator is not None:
+            state["validation_rng"] = (self.validation_random_generator.get_state())
+        return state
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Any],
+    ) -> None:
+        rng_state = state_dict.get("validation_rng")
+        if (rng_state is not None and self.validation_random_generator is not None):
+            self.validation_random_generator.set_state(rng_state)

fastvideo/train/entrypoint/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: Apache-2.0

fastvideo/train/entrypoint/dcp_to_diffusers.py

@@ -0,0 +1,397 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Convert a DCP training checkpoint to a diffusers-style model directory.
+
+Works on a single GPU regardless of how many GPUs were used for training
+(DCP handles resharding automatically).
+
+Usage (no torchrun needed)::
+
+    python -m fastvideo.train.entrypoint.dcp_to_diffusers \
+        --checkpoint /path/to/checkpoint-1000 \
+        --output-dir /path/to/diffusers_output
+
+Or with torchrun (also fine)::
+
+    torchrun --nproc_per_node=1 \
+        -m fastvideo.train.entrypoint.dcp_to_diffusers \
+        --checkpoint ... --output-dir ...
+
+The checkpoint must contain ``metadata.json`` (written by
+``CheckpointManager``).  If the checkpoint predates metadata
+support, pass ``--config`` explicitly to provide the training
+YAML.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import Any
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def _ensure_distributed() -> None:
+    """Set up a single-process distributed env if needed.
+
+    When running under ``torchrun`` the env vars are already set.
+    For plain ``python`` we fill in the minimum required vars so
+    that ``init_process_group`` succeeds with world_size=1.
+    """
+    for key, default in [
+        ("RANK", "0"),
+        ("LOCAL_RANK", "0"),
+        ("WORLD_SIZE", "1"),
+        ("MASTER_ADDR", "127.0.0.1"),
+        ("MASTER_PORT", "29500"),
+    ]:
+        os.environ.setdefault(key, default)
+
+
+def _save_role_pretrained(
+    *,
+    role: str,
+    base_model_path: str,
+    output_dir: str,
+    module_names: list[str] | None = None,
+    overwrite: bool = False,
+    model: Any,
+) -> str:
+    """Export a role's modules into a diffusers-style model dir.
+
+    Produces a ``model_path`` loadable by
+    ``PipelineComponentLoader`` (``model_index.json``,
+    ``transformer/``, ``vae/``, etc. copied from
+    ``base_model_path``).
+    """
+    import shutil
+    from pathlib import Path
+
+    import torch
+    import torch.distributed as dist
+    from torch.distributed.checkpoint.state_dict import (
+        StateDictOptions,
+        get_model_state_dict,
+    )
+
+    from fastvideo.utils import maybe_download_model
+
+    def _rank() -> int:
+        if dist.is_available() and dist.is_initialized():
+            return int(dist.get_rank())
+        return 0
+
+    def _barrier() -> None:
+        if dist.is_available() and dist.is_initialized():
+            dist.barrier()
+
+    local_base = Path(maybe_download_model(str(base_model_path))).resolve()
+    dst = Path(os.path.expanduser(str(output_dir))).resolve()
+
+    if _rank() == 0:
+        if dst.exists():
+            if overwrite:
+                shutil.rmtree(dst, ignore_errors=True)
+            else:
+                raise FileExistsError(f"Refusing to overwrite existing "
+                                      f"directory: {dst}. "
+                                      "Pass --overwrite to replace it.")
+
+        def _copy_or_link(src: str, dest: str) -> None:
+            try:
+                os.link(src, dest)
+            except OSError:
+                shutil.copy2(src, dest)
+
+        logger.info(
+            "Creating pretrained export dir at %s "
+            "(base=%s)",
+            dst,
+            local_base,
+        )
+        shutil.copytree(
+            local_base,
+            dst,
+            symlinks=True,
+            copy_function=_copy_or_link,
+        )
+
+    _barrier()
+
+    modules: dict[str, torch.nn.Module] = {}
+    if model.transformer is not None:
+        modules["transformer"] = model.transformer
+
+    if module_names is None:
+        module_names = sorted(modules.keys())
+
+    for module_name in module_names:
+        if module_name not in modules:
+            raise KeyError(f"Role {role!r} does not have module "
+                           f"{module_name!r}. "
+                           f"Available: {sorted(modules.keys())}")
+
+        module_dir = dst / module_name
+        if not module_dir.is_dir():
+            raise FileNotFoundError(f"Export directory missing component "
+                                    f"dir {module_name!r}: {module_dir}")
+
+        options = StateDictOptions(
+            full_state_dict=True,
+            cpu_offload=True,
+        )
+        state_dict = get_model_state_dict(
+            modules[module_name],
+            options=options,
+        )
+
+        if _rank() == 0:
+            for path in module_dir.glob("*.safetensors"):
+                path.unlink(missing_ok=True)
+
+            tensor_state: dict[str, torch.Tensor] = {}
+            for key, value in state_dict.items():
+                if not isinstance(value, torch.Tensor):
+                    raise TypeError(f"Expected tensor in state_dict "
+                                    f"for {module_name}.{key}, "
+                                    f"got {type(value).__name__}")
+                tensor_state[key] = value.detach().cpu()
+
+            from safetensors.torch import save_file
+
+            out_path = module_dir / "model.safetensors"
+            logger.info(
+                "Saving %s weights to %s (%s tensors)",
+                module_name,
+                out_path,
+                len(tensor_state),
+            )
+            save_file(tensor_state, str(out_path))
+
+        _barrier()
+
+    return str(dst)
+
+
+def convert(
+    *,
+    checkpoint_dir: str,
+    output_dir: str,
+    config_path: str | None = None,
+    role: str = "student",
+    overwrite: bool = False,
+) -> str:
+    """Load a DCP checkpoint and export as a diffusers model.
+
+    Returns the path to the exported model directory.
+    """
+    _ensure_distributed()
+
+    from fastvideo.distributed import (
+        maybe_init_distributed_environment_and_model_parallel, )
+    from fastvideo.train.utils.builder import build_from_config
+    from fastvideo.train.utils.checkpoint import (
+        CheckpointManager,
+        _resolve_resume_checkpoint,
+    )
+    from fastvideo.train.utils.config import (
+        RunConfig,
+        load_run_config,
+    )
+
+    import torch.distributed.checkpoint as dcp
+
+    # -- Resolve checkpoint directory --
+    resolved = _resolve_resume_checkpoint(
+        checkpoint_dir,
+        output_dir=checkpoint_dir,
+    )
+    dcp_dir = resolved / "dcp"
+    if not dcp_dir.is_dir():
+        raise FileNotFoundError(f"Missing dcp/ under {resolved}")
+
+    # -- Obtain config --
+    cfg: RunConfig
+    if config_path is not None:
+        cfg = load_run_config(config_path)
+    else:
+        metadata = CheckpointManager.load_metadata(resolved)
+        raw_config = metadata.get("config")
+        if raw_config is None:
+            raise ValueError("Checkpoint metadata.json does not "
+                             "contain 'config'. Pass --config "
+                             "explicitly.")
+        cfg = _run_config_from_raw(raw_config)
+
+    tc = cfg.training
+
+    # -- Init distributed (1 GPU is enough; DCP reshards) --
+    maybe_init_distributed_environment_and_model_parallel(
+        tp_size=1,
+        sp_size=1,
+    )
+
+    # Override distributed config so model loading uses 1 GPU.
+    tc.distributed.tp_size = 1
+    tc.distributed.sp_size = 1
+    tc.distributed.num_gpus = 1
+    tc.distributed.hsdp_replicate_dim = 1
+    tc.distributed.hsdp_shard_dim = 1
+
+    # -- Build model (loads pretrained weights + FSDP) --
+    _, method, _, _ = build_from_config(cfg)
+
+    # -- Load DCP weights into the model --
+    states = method.checkpoint_state()
+    logger.info(
+        "Loading DCP checkpoint from %s",
+        resolved,
+    )
+    dcp.load(states, checkpoint_id=str(dcp_dir))
+
+    # -- Export to diffusers format --
+    model = method._role_models[role]
+    base_model_path = str(tc.model_path)
+    if not base_model_path:
+        raise ValueError("Cannot determine base_model_path from "
+                         "config. Ensure models.student.init_from "
+                         "is set.")
+
+    logger.info(
+        "Exporting role=%s to %s (base=%s)",
+        role,
+        output_dir,
+        base_model_path,
+    )
+    result = _save_role_pretrained(
+        role=role,
+        base_model_path=base_model_path,
+        output_dir=output_dir,
+        overwrite=overwrite,
+        model=model,
+    )
+    logger.info("Export complete: %s", result)
+    return result
+
+
+def _run_config_from_raw(raw: dict[str, Any], ) -> Any:
+    """Reconstruct a RunConfig from a raw config dict.
+
+    This mirrors ``load_run_config`` but operates on an
+    already-parsed dict (from metadata.json) instead of
+    reading from a YAML file.
+    """
+    from fastvideo.train.utils.config import (
+        RunConfig,
+        _build_training_config,
+        _parse_pipeline_config,
+        _require_mapping,
+        _require_str,
+    )
+
+    models_raw = _require_mapping(
+        raw.get("models"),
+        where="models",
+    )
+    models: dict[str, dict[str, Any]] = {}
+    for role_key, model_cfg_raw in models_raw.items():
+        role_str = _require_str(
+            role_key,
+            where="models.<role>",
+        )
+        model_cfg = _require_mapping(
+            model_cfg_raw,
+            where=f"models.{role_str}",
+        )
+        models[role_str] = dict(model_cfg)
+
+    method_raw = _require_mapping(
+        raw.get("method"),
+        where="method",
+    )
+    method = dict(method_raw)
+
+    callbacks_raw = raw.get("callbacks")
+    callbacks: dict[str, dict[str, Any]] = (_require_mapping(
+        callbacks_raw,
+        where="callbacks",
+    ) if callbacks_raw is not None else {})
+
+    pipeline_config = _parse_pipeline_config(
+        raw,
+        models=models,
+    )
+
+    training_raw = _require_mapping(
+        raw.get("training"),
+        where="training",
+    )
+    t = dict(training_raw)
+    training = _build_training_config(
+        t,
+        models=models,
+        pipeline_config=pipeline_config,
+    )
+
+    return RunConfig(
+        models=models,
+        method=method,
+        training=training,
+        callbacks=callbacks,
+        raw=raw,
+    )
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=("Convert a DCP training checkpoint to a "
+                                                  "diffusers-style model directory. "
+                                                  "Only 1 GPU needed (DCP reshards "
+                                                  "automatically)."), )
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        required=True,
+        help=("Path to checkpoint-<step> dir, its dcp/ "
+              "subdir, or an output_dir (auto-picks "
+              "latest)."),
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        required=True,
+        help="Destination for the diffusers model.",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default=None,
+        help=("Training YAML config. If omitted, read "
+              "from checkpoint metadata.json."),
+    )
+    parser.add_argument(
+        "--role",
+        type=str,
+        default="student",
+        help="Role to export (default: student).",
+    )
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Overwrite output-dir if it exists.",
+    )
+    args = parser.parse_args(sys.argv[1:])
+
+    convert(
+        checkpoint_dir=args.checkpoint,
+        output_dir=args.output_dir,
+        config_path=args.config,
+        role=args.role,
+        overwrite=args.overwrite,
+    )
+
+
+if __name__ == "__main__":
+    main()

fastvideo/train/entrypoint/misc/wan_ode_init_conversion.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Convert Self-Forcing ode_init.pt to HuggingFace diffusers format.
+
+The official ode_init.pt from
+https://huggingface.co/gdhe17/Self-Forcing/resolve/main/checkpoints/ode_init.pt
+stores weights under ``{"generator": {<original_wan_keys>}}``.
+
+This script converts those keys to diffusers
+``WanTransformer3DModel`` format, verifies them against a reference
+model, and saves a complete diffusers-compatible model directory
+(transformer + scheduler + vae + text_encoder + tokenizer).
+
+Usage:
+    python -m fastvideo.train.entrypoint.misc.wan_ode_init_conversion \
+        --input /path/to/ode_init.pt \
+        --output /path/to/WanOdeInit \
+        --base-model Wan-AI/Wan2.1-T2V-1.3B-Diffusers
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+
+import torch
+
+# -- Key conversion -----------------------------------------------
+
+
+def _convert_key(key: str) -> str:
+    """Map a single original-Wan key to diffusers key."""
+    k = key
+    if k.startswith("model."):
+        k = k[len("model."):]
+
+    # Top-level modules
+    k = k.replace("head.modulation", "scale_shift_table")
+    k = k.replace("head.head.", "proj_out.")
+    k = k.replace(
+        "time_embedding.0.",
+        "condition_embedder.time_embedder.linear_1.",
+    )
+    k = k.replace(
+        "time_embedding.2.",
+        "condition_embedder.time_embedder.linear_2.",
+    )
+    k = k.replace(
+        "time_projection.1.",
+        "condition_embedder.time_proj.",
+    )
+    k = k.replace(
+        "text_embedding.0.",
+        "condition_embedder.text_embedder.linear_1.",
+    )
+    k = k.replace(
+        "text_embedding.2.",
+        "condition_embedder.text_embedder.linear_2.",
+    )
+
+    # Block-level: order matters — rename attn sub-modules
+    # before renaming q/k/v/o so prefixes don't collide.
+    k = k.replace(".modulation", ".scale_shift_table")
+    k = k.replace(".self_attn.", ".attn1.")
+    k = k.replace(".cross_attn.", ".attn2.")
+    k = k.replace(".norm3.", ".norm2.")
+
+    # Attention projections
+    k = k.replace(".attn1.q.", ".attn1.to_q.")
+    k = k.replace(".attn1.k.", ".attn1.to_k.")
+    k = k.replace(".attn1.v.", ".attn1.to_v.")
+    k = k.replace(".attn1.o.", ".attn1.to_out.0.")
+    k = k.replace(".attn2.q.", ".attn2.to_q.")
+    k = k.replace(".attn2.k.", ".attn2.to_k.")
+    k = k.replace(".attn2.v.", ".attn2.to_v.")
+    k = k.replace(".attn2.o.", ".attn2.to_out.0.")
+
+    # FFN
+    k = k.replace(".ffn.0.", ".ffn.net.0.proj.")
+    k = k.replace(".ffn.2.", ".ffn.net.2.")
+
+    return k
+
+
+def convert_state_dict(orig_sd: dict[str, torch.Tensor], ) -> dict[str, torch.Tensor]:
+    """Convert an entire original-Wan state dict."""
+    return {_convert_key(k): v for k, v in orig_sd.items()}
+
+
+# -- Main ---------------------------------------------------------
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Convert ode_init.pt to diffusers format", )
+    parser.add_argument(
+        "--input",
+        required=True,
+        help="Path to ode_init.pt",
+    )
+    parser.add_argument(
+        "--output",
+        required=True,
+        help="Output directory for diffusers model",
+    )
+    parser.add_argument(
+        "--base-model",
+        default="Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+        help="HF repo or local path for base model "
+        "(provides config, VAE, tokenizer, etc.)",
+    )
+    parser.add_argument(
+        "--skip-verify",
+        action="store_true",
+        help="Skip strict load verification",
+    )
+    args = parser.parse_args()
+
+    # 1. Load checkpoint
+    print(f"Loading {args.input} ...")
+    ckpt = torch.load(args.input, map_location="cpu", weights_only=False)
+    if "generator" in ckpt:
+        orig_sd = ckpt["generator"]
+    elif isinstance(ckpt, dict) and any(k.startswith("model.") for k in ckpt):
+        orig_sd = ckpt
+    else:
+        raise ValueError("Cannot find weights in checkpoint. "
+                         "Expected key 'generator' or keys starting "
+                         "with 'model.'.")
+    print(f"  Found {len(orig_sd)} weight tensors")
+
+    # 2. Convert keys
+    print("Converting keys ...")
+    new_sd = convert_state_dict(orig_sd)
+
+    # 3. Verify against reference model
+    if not args.skip_verify:
+        from diffusers import WanTransformer3DModel
+
+        print(f"Loading reference model from "
+              f"{args.base_model} for verification ...")
+        ref_model = WanTransformer3DModel.from_pretrained(
+            args.base_model,
+            subfolder="transformer",
+            torch_dtype=torch.bfloat16,
+        )
+        ref_keys = set(ref_model.state_dict().keys())
+        new_keys = set(new_sd.keys())
+        missing = sorted(ref_keys - new_keys)
+        extra = sorted(new_keys - ref_keys)
+        if missing:
+            print(f"  WARNING: {len(missing)} missing keys:")
+            for k in missing[:10]:
+                print(f"    {k}")
+            if len(missing) > 10:
+                print(f"    ... and {len(missing) - 10} more")
+        if extra:
+            print(f"  WARNING: {len(extra)} extra keys:")
+            for k in extra[:10]:
+                print(f"    {k}")
+            if len(extra) > 10:
+                print(f"    ... and {len(extra) - 10} more")
+        if missing or extra:
+            raise RuntimeError("Key mismatch — conversion mapping needs "
+                               "updating. Use --skip-verify to bypass.")
+
+        ref_model.load_state_dict(new_sd, strict=True)
+        print("  Strict load OK — all keys match!")
+    else:
+        ref_model = None
+        print("  Skipping verification")
+
+    # 4. Save transformer
+    output_dir = os.path.abspath(args.output)
+    os.makedirs(output_dir, exist_ok=True)
+    transformer_dir = os.path.join(output_dir, "transformer")
+
+    if ref_model is not None:
+        print(f"Saving transformer to {transformer_dir} ...")
+        ref_model.save_pretrained(transformer_dir)
+    else:
+        from safetensors.torch import save_file
+
+        os.makedirs(transformer_dir, exist_ok=True)
+        print(f"Saving transformer weights to "
+              f"{transformer_dir} ...")
+        save_file(new_sd, os.path.join(transformer_dir, "model.safetensors"))
+
+    # 5. Copy non-transformer files from base model
+    from huggingface_hub import snapshot_download
+
+    print(f"Downloading base model files from "
+          f"{args.base_model} ...")
+    base_path = snapshot_download(
+        args.base_model,
+        allow_patterns=[
+            "model_index.json",
+            "scheduler/*",
+            "tokenizer/*",
+            "vae/*",
+            "text_encoder/*",
+        ],
+    )
+    for item in sorted(os.listdir(base_path)):
+        if item == "transformer":
+            continue
+        src = os.path.join(base_path, item)
+        dst = os.path.join(output_dir, item)
+        if os.path.isfile(src):
+            shutil.copy2(src, dst)
+            print(f"  Copied {item}")
+        elif os.path.isdir(src):
+            if os.path.exists(dst):
+                shutil.rmtree(dst)
+            shutil.copytree(src, dst)
+            print(f"  Copied {item}/")
+
+    print(f"\nDone! Diffusers model saved to {output_dir}")
+
+
+if __name__ == "__main__":
+    main()

fastvideo/train/entrypoint/train.py

@@ -0,0 +1,139 @@
+# SPDX-License-Identifier: Apache-2.0
+"""YAML-only training entrypoint.
+
+Usage::
+
+    torchrun --nproc_per_node=<N> -m fastvideo.train.entrypoint.train \
+        --config path/to/run.yaml
+
+Any unknown ``--dotted.key value`` arguments are applied as
+overrides to the YAML config before parsing.  For example::
+
+    torchrun --nproc_per_node=8 -m fastvideo.train.entrypoint.train \
+        --config path/to/run.yaml \
+        --training.distributed.num_gpus 8 \
+        --training.optimizer.learning_rate 1e-5
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+from typing import Any
+
+import torch
+
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def run_training_from_config(
+    config_path: str,
+    *,
+    dry_run: bool = False,
+    overrides: list[str] | None = None,
+) -> None:
+    """YAML-only training entrypoint (schema v2)."""
+
+    from fastvideo.distributed import (
+        maybe_init_distributed_environment_and_model_parallel, )
+    from fastvideo.train import Trainer
+    from fastvideo.train.utils.checkpoint import (
+        CheckpointConfig,
+        CheckpointManager,
+    )
+    from fastvideo.train.utils.builder import build_from_config
+    from fastvideo.train.utils.config import load_run_config
+
+    # Enable deterministic mode for reproducibility.
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+
+    cfg = load_run_config(config_path, overrides=overrides)
+    tc = cfg.training
+
+    maybe_init_distributed_environment_and_model_parallel(
+        tc.distributed.tp_size,
+        tc.distributed.sp_size,
+    )
+
+    _, method, dataloader, start_step = build_from_config(cfg)
+
+    if dry_run:
+        logger.info("Dry-run: config parsed and "
+                    "build_from_config succeeded.")
+        return
+
+    trainer = Trainer(
+        tc,
+        config=cfg.resolved_config(),
+        callback_configs=cfg.callbacks,
+    )
+
+    # Attach the exact YAML used for this run to the
+    # tracker (e.g., W&B Files).
+    trainer.tracker.log_file(
+        os.path.abspath(os.path.expanduser(config_path)),
+        name="run.yaml",
+    )
+
+    ckpt_config = CheckpointConfig(
+        save_steps=int(tc.checkpoint.training_state_checkpointing_steps or 0),
+        keep_last=int(tc.checkpoint.checkpoints_total_limit or 0),
+    )
+
+    checkpoint_manager = CheckpointManager(
+        method=method,
+        dataloader=dataloader,
+        output_dir=tc.checkpoint.output_dir,
+        config=ckpt_config,
+        callbacks=trainer.callbacks,
+        raw_config=cfg.raw,
+    )
+
+    trainer.run(
+        method,
+        dataloader=dataloader,
+        max_steps=tc.loop.max_train_steps,
+        start_step=start_step,
+        checkpoint_manager=checkpoint_manager,
+    )
+
+
+def main(
+    args: Any,
+    overrides: list[str] | None = None,
+) -> None:
+    config_path = str(args.config)
+    dry_run = bool(args.dry_run)
+    logger.info(
+        "Starting training from config=%s",
+        config_path,
+    )
+    run_training_from_config(
+        config_path,
+        dry_run=dry_run,
+        overrides=overrides,
+    )
+    logger.info("Training completed")
+
+
+if __name__ == "__main__":
+    argv = sys.argv
+    parser = argparse.ArgumentParser(description="YAML-only training entrypoint.", )
+    parser.add_argument(
+        "--config",
+        type=str,
+        required=True,
+        help=("Path to training YAML config (schema v2)."),
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help=("Parse config and build runtime, "
+              "but do not start training."),
+    )
+    args, unknown = parser.parse_known_args(argv[1:])
+    main(args, overrides=unknown if unknown else None)

fastvideo/train/methods/__init__.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.train.methods.base import TrainingMethod
+
+__all__ = [
+    "TrainingMethod",
+    "DMD2Method",
+    "FineTuneMethod",
+    "SelfForcingMethod",
+    "DiffusionForcingSFTMethod",
+]
+
+
+def __getattr__(name: str) -> object:
+    if name == "DMD2Method":
+        from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
+        return DMD2Method
+    if name == "FineTuneMethod":
+        from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
+        return FineTuneMethod
+    if name == "SelfForcingMethod":
+        from fastvideo.train.methods.distribution_matching.self_forcing import SelfForcingMethod
+        return SelfForcingMethod
+    if name == "DiffusionForcingSFTMethod":
+        from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod
+        return DiffusionForcingSFTMethod
+    raise AttributeError(name)

fastvideo/train/methods/base.py

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from typing import Any, Literal, TypeAlias, cast
+
+import torch
+
+from fastvideo.logger import init_logger
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.checkpoint import _RoleModuleContainer
+from fastvideo.training.checkpointing_utils import (
+    ModelWrapper,
+    OptimizerWrapper,
+    SchedulerWrapper,
+)
+
+logger = init_logger(__name__)
+
+LogScalar: TypeAlias = float | int | torch.Tensor
+
+
+class TrainingMethod(torch.nn.Module, ABC):
+    """Base training method (algorithm layer).
+
+    Subclasses own their role models (student, teacher, critic, …) as
+    plain attributes and manage optimizers directly — no ``RoleManager``
+    or ``RoleHandle``.
+
+    The constructor receives *role_models* (a ``dict[str, ModelBase]``)
+    and a *cfg* object.  It calls ``init_preprocessors`` on the student
+    and builds ``self.role_modules`` for FSDP wrapping.
+
+    A single shared CUDA RNG generator (``cuda_generator``) is
+    created in :meth:`on_train_start`.  All ``torch.randn`` /
+    ``torch.randint`` calls in methods **and** models must use this
+    generator instead of relying on global RNG state.
+    """
+
+    # Shared CUDA RNG generator (initialized in on_train_start).
+    cuda_generator: torch.Generator | None = None
+
+    def __init__(
+        self,
+        *,
+        cfg: Any,
+        role_models: dict[str, ModelBase],
+    ) -> None:
+        super().__init__()
+        self.tracker: Any | None = None
+        self._role_models: dict[str, ModelBase] = dict(role_models)
+
+        self.student = role_models["student"]
+        self.training_config = cfg.training
+        self.method_config: dict[str, Any] = dict(cfg.method)
+        self.validation_config: dict[str, Any] = dict(getattr(cfg, "validation", {}) or {})
+
+        # Build nn.ModuleDict for FSDP / checkpoint visibility.
+        self.role_modules = torch.nn.ModuleDict()
+        for role, model in role_models.items():
+            mods: dict[str, torch.nn.Module] = {}
+            transformer = getattr(model, "transformer", None)
+            if isinstance(transformer, torch.nn.Module):
+                mods["transformer"] = transformer
+            if mods:
+                self.role_modules[role] = torch.nn.ModuleDict(mods)
+
+    # ------------------------------------------------------------------
+
+    def set_tracker(self, tracker: Any) -> None:
+        self.tracker = tracker
+
+    @abstractmethod
+    def single_train_step(
+        self,
+        batch: dict[str, Any],
+        iteration: int,
+        *,
+        current_vsa_sparsity: float = 0.0,
+    ) -> tuple[
+            dict[str, torch.Tensor],
+            dict[str, Any],
+            dict[str, LogScalar],
+    ]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_optimizers(
+        self,
+        iteration: int,
+    ) -> Sequence[torch.optim.Optimizer]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_lr_schedulers(
+        self,
+        iteration: int,
+    ) -> Sequence[Any]:
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def _optimizer_dict(self) -> dict[str, Any]:
+        ...
+
+    @property
+    @abstractmethod
+    def _lr_scheduler_dict(self) -> dict[str, Any]:
+        ...
+
+    def checkpoint_state(self) -> dict[str, Any]:
+        """Return DCP-ready checkpoint state for all trainable roles.
+
+        Keys follow the convention:
+        ``roles.<role>.<module>``, ``optimizers.<role>``,
+        ``schedulers.<role>``, ``random_state.*``.
+
+        EMA state is managed by the ``EMACallback`` and is
+        checkpointed through the callback state mechanism.
+        """
+        states: dict[str, Any] = {}
+
+        for role, model in self._role_models.items():
+            if not getattr(model, "_trainable", False):
+                continue
+
+            modules: dict[str, torch.nn.Module] = {}
+            if model.transformer is not None:
+                modules["transformer"] = model.transformer
+
+            container = _RoleModuleContainer(modules)
+
+            for module_name, module in modules.items():
+                states[f"roles.{role}.{module_name}"] = ModelWrapper(module)
+
+            opt = self._optimizer_dict.get(role)
+            if opt is not None:
+                states[f"optimizers.{role}"] = OptimizerWrapper(container, opt)
+
+            sched = self._lr_scheduler_dict.get(role)
+            if sched is not None:
+                states[f"schedulers.{role}"] = SchedulerWrapper(sched)
+
+        return states
+
+    def backward(
+        self,
+        loss_map: dict[str, torch.Tensor],
+        outputs: dict[str, Any],
+        *,
+        grad_accum_rounds: int = 1,
+    ) -> None:
+        del outputs
+        grad_accum_rounds = max(1, int(grad_accum_rounds))
+        (loss_map["total_loss"] / grad_accum_rounds).backward()
+
+    def optimizers_schedulers_step(
+        self,
+        iteration: int,
+    ) -> None:
+        for optimizer in self.get_optimizers(iteration):
+            optimizer.step()
+        for scheduler in self.get_lr_schedulers(iteration):
+            scheduler.step()
+
+    def optimizers_zero_grad(
+        self,
+        iteration: int,
+    ) -> None:
+        for optimizer in self.get_optimizers(iteration):
+            try:
+                optimizer.zero_grad(set_to_none=True)
+            except TypeError:
+                optimizer.zero_grad()
+
+    def seed_optimizer_state_for_resume(self) -> None:
+        """Seed optimizer state so DCP can load saved state.
+
+        A fresh optimizer has empty state (exp_avg, exp_avg_sq,
+        step are only created on the first optimizer.step()).
+        DCP needs matching entries to load into; without them
+        the saved optimizer state is silently dropped.
+        """
+        for opt in self.get_optimizers(0):
+            for group in opt.param_groups:
+                for p in group["params"]:
+                    if not p.requires_grad:
+                        continue
+                    if len(opt.state.get(p, {})) > 0:
+                        continue
+                    opt.state[p] = {
+                        "step": torch.tensor(0.0),
+                        "exp_avg": torch.zeros_like(p),
+                        "exp_avg_sq": torch.zeros_like(p),
+                    }
+
+    # -- Shared hooks (override in subclasses as needed) --
+
+    def get_grad_clip_targets(
+        self,
+        iteration: int,
+    ) -> dict[str, torch.nn.Module]:
+        """Return modules whose gradients should be clipped.
+
+        Override in subclasses to add/conditionally include
+        modules (e.g. critic, conditionally student).
+        Default: student transformer.
+        """
+        return {"student": self.student.transformer}
+
+    def on_train_start(self) -> None:
+        from fastvideo.distributed import (
+            get_world_group, )
+        from fastvideo.utils import set_random_seed
+
+        seed = self.training_config.data.seed
+        if seed is None:
+            raise ValueError("training.data.seed must be set")
+        seed = int(seed)
+
+        world_group = get_world_group()
+        global_rank = int(world_group.rank)
+        sp_size = int(self.training_config.distributed.sp_size or 1)
+
+        # Ranks within the same SP group share a seed.
+        if sp_size > 1:
+            sp_group_seed = seed + (global_rank // sp_size)
+        else:
+            sp_group_seed = seed + global_rank
+
+        set_random_seed(seed)
+
+        self.cuda_generator = torch.Generator(device=self.student.device).manual_seed(sp_group_seed)
+
+        self.student.on_train_start()
+
+    def get_rng_generators(self, ) -> dict[str, torch.Generator]:
+        generators: dict[str, torch.Generator] = {}
+        if self.cuda_generator is not None:
+            generators["cuda"] = self.cuda_generator
+        return generators
+
+    @staticmethod
+    def _parse_attn_kind(raw: Any, ) -> Literal["dense", "vsa"]:
+        if raw in (None, ""):
+            return "dense"
+        kind = str(raw).strip().lower()
+        if kind not in {"dense", "vsa"}:
+            raise ValueError("method_config.attn_kind must be one of "
+                             f"{{'dense', 'vsa'}}, got {raw!r}.")
+        return cast(Literal["dense", "vsa"], kind)

fastvideo/train/methods/consistency_model/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__all__: list[str] = []

fastvideo/train/methods/distribution_matching/__init__.py

@@ -0,0 +1,10 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
+from fastvideo.train.methods.distribution_matching.self_forcing import (
+    SelfForcingMethod, )
+
+__all__ = [
+    "DMD2Method",
+    "SelfForcingMethod",
+]

fastvideo/train/methods/distribution_matching/dmd2.py

@@ -0,0 +1,633 @@
+# SPDX-License-Identifier: Apache-2.0
+"""DMD2 distillation method (algorithm layer)."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import torch
+import torch.nn.functional as F
+
+from fastvideo.train.methods.base import TrainingMethod, LogScalar
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.optimizer import (
+    build_optimizer_and_scheduler, )
+from fastvideo.train.utils.config import (
+    get_optional_float,
+    get_optional_int,
+    parse_betas,
+)
+
+
+class DMD2Method(TrainingMethod):
+    """DMD2 distillation algorithm (method layer).
+
+    Owns role model instances directly:
+    - ``self.student`` — trainable student :class:`ModelBase`
+    - ``self.teacher`` — frozen teacher :class:`ModelBase`
+    - ``self.critic`` — trainable critic :class:`ModelBase`
+    """
+
+    def __init__(
+        self,
+        *,
+        cfg: Any,
+        role_models: dict[str, ModelBase],
+    ) -> None:
+        super().__init__(cfg=cfg, role_models=role_models)
+
+        if "student" not in role_models:
+            raise ValueError("DMD2Method requires role 'student'")
+        if "teacher" not in role_models:
+            raise ValueError("DMD2Method requires role 'teacher'")
+        if "critic" not in role_models:
+            raise ValueError("DMD2Method requires role 'critic'")
+
+        self.teacher = role_models["teacher"]
+        self.critic = role_models["critic"]
+
+        if not self.student._trainable:
+            raise ValueError("DMD2Method requires student to be trainable")
+        if self.teacher._trainable:
+            raise ValueError("DMD2Method requires teacher to be "
+                             "non-trainable")
+        if not self.critic._trainable:
+            raise ValueError("DMD2Method requires critic to be trainable")
+        self._cfg_uncond = self._parse_cfg_uncond()
+        self._rollout_mode = self._parse_rollout_mode()
+        self._denoising_step_list: torch.Tensor | None = (None)
+
+        # Initialize preprocessors on student.
+        self.student.init_preprocessors(self.training_config)
+
+        self._init_optimizers_and_schedulers()
+
+    @property
+    def _optimizer_dict(self, ) -> dict[str, torch.optim.Optimizer]:
+        return {
+            "student": self._student_optimizer,
+            "critic": self._critic_optimizer,
+        }
+
+    @property
+    def _lr_scheduler_dict(self) -> dict[str, Any]:
+        return {
+            "student": self._student_lr_scheduler,
+            "critic": self._critic_lr_scheduler,
+        }
+
+    # TrainingMethod override: single_train_step
+    def single_train_step(
+        self,
+        batch: dict[str, Any],
+        iteration: int,
+        *,
+        current_vsa_sparsity: float = 0.0,
+    ) -> tuple[
+            dict[str, torch.Tensor],
+            dict[str, Any],
+            dict[str, LogScalar],
+    ]:
+        latents_source: Literal["data", "zeros"] = "data"
+        if self._rollout_mode == "simulate":
+            latents_source = "zeros"
+
+        training_batch = self.student.prepare_batch(
+            batch,
+            generator=self.cuda_generator,
+            current_vsa_sparsity=current_vsa_sparsity,
+            latents_source=latents_source,
+        )
+
+        update_student = self._should_update_student(iteration)
+
+        generator_loss = torch.zeros(
+            (),
+            device=training_batch.latents.device,
+            dtype=training_batch.latents.dtype,
+        )
+        student_ctx = None
+        if update_student:
+            generator_pred_x0 = self._student_rollout(training_batch, with_grad=True)
+            student_ctx = (
+                training_batch.timesteps,
+                training_batch.attn_metadata_vsa,
+            )
+            generator_loss = self._dmd_loss(generator_pred_x0, training_batch)
+
+        (
+            fake_score_loss,
+            critic_ctx,
+            critic_outputs,
+        ) = self._critic_flow_matching_loss(training_batch)
+
+        total_loss = generator_loss + fake_score_loss
+        loss_map = {
+            "total_loss": total_loss,
+            "generator_loss": generator_loss,
+            "fake_score_loss": fake_score_loss,
+        }
+
+        outputs: dict[str, Any] = dict(critic_outputs)
+        outputs["_fv_backward"] = {
+            "update_student": update_student,
+            "student_ctx": student_ctx,
+            "critic_ctx": critic_ctx,
+        }
+        metrics: dict[str, LogScalar] = {"update_student": float(update_student)}
+        return loss_map, outputs, metrics
+
+    # TrainingMethod override: backward
+    def backward(
+        self,
+        loss_map: dict[str, torch.Tensor],
+        outputs: dict[str, Any],
+        *,
+        grad_accum_rounds: int = 1,
+    ) -> None:
+        grad_accum_rounds = max(1, int(grad_accum_rounds))
+        backward_ctx = outputs.get("_fv_backward")
+        if not isinstance(backward_ctx, dict):
+            super().backward(
+                loss_map,
+                outputs,
+                grad_accum_rounds=grad_accum_rounds,
+            )
+            return
+
+        update_student = bool(backward_ctx.get("update_student", False))
+        if update_student:
+            student_ctx = backward_ctx.get("student_ctx")
+            if student_ctx is None:
+                raise RuntimeError("Missing student backward context")
+            self.student.backward(
+                loss_map["generator_loss"],
+                student_ctx,
+                grad_accum_rounds=grad_accum_rounds,
+            )
+
+        critic_ctx = backward_ctx.get("critic_ctx")
+        if critic_ctx is None:
+            raise RuntimeError("Missing critic backward context")
+        self.critic.backward(
+            loss_map["fake_score_loss"],
+            critic_ctx,
+            grad_accum_rounds=grad_accum_rounds,
+        )
+
+    # TrainingMethod override: get_optimizers
+    def get_optimizers(
+        self,
+        iteration: int,
+    ) -> list[torch.optim.Optimizer]:
+        optimizers: list[torch.optim.Optimizer] = []
+        optimizers.append(self._critic_optimizer)
+        if self._should_update_student(iteration):
+            optimizers.append(self._student_optimizer)
+        return optimizers
+
+    # TrainingMethod override: get_lr_schedulers
+    def get_lr_schedulers(
+        self,
+        iteration: int,
+    ) -> list[Any]:
+        schedulers: list[Any] = []
+        schedulers.append(self._critic_lr_scheduler)
+        if self._should_update_student(iteration):
+            schedulers.append(self._student_lr_scheduler)
+        return schedulers
+
+    # TrainingMethod override: get_grad_clip_targets
+    def get_grad_clip_targets(
+        self,
+        iteration: int,
+    ) -> dict[str, torch.nn.Module]:
+        targets: dict[str, torch.nn.Module] = {}
+        if self._should_update_student(iteration):
+            targets["student"] = (self.student.transformer)
+        targets["critic"] = self.critic.transformer
+        return targets
+
+    def _parse_rollout_mode(self, ) -> Literal["simulate", "data_latent"]:
+        raw = self.method_config.get("rollout_mode", None)
+        if raw is None:
+            raise ValueError("method_config.rollout_mode must be set "
+                             "for DMD2")
+        if not isinstance(raw, str):
+            raise ValueError("method_config.rollout_mode must be a "
+                             "string, "
+                             f"got {type(raw).__name__}")
+        mode = raw.strip().lower()
+        if mode in ("simulate", "sim"):
+            return "simulate"
+        if mode in ("data_latent", "data", "vae_latent"):
+            return "data_latent"
+        raise ValueError("method_config.rollout_mode must be one of "
+                         "{simulate, data_latent}, got "
+                         f"{raw!r}")
+
+    def _parse_cfg_uncond(self, ) -> dict[str, Any] | None:
+        raw = self.method_config.get("cfg_uncond", None)
+        if raw is None:
+            return None
+        if not isinstance(raw, dict):
+            raise ValueError("method_config.cfg_uncond must be a dict "
+                             f"when set, got {type(raw).__name__}")
+
+        cfg: dict[str, Any] = dict(raw)
+
+        on_missing_raw = cfg.get("on_missing", "error")
+        if on_missing_raw is None:
+            on_missing_raw = "error"
+        if not isinstance(on_missing_raw, str):
+            raise ValueError("method_config.cfg_uncond.on_missing must "
+                             "be a string, got "
+                             f"{type(on_missing_raw).__name__}")
+        on_missing = on_missing_raw.strip().lower()
+        if on_missing not in {"error", "ignore"}:
+            raise ValueError("method_config.cfg_uncond.on_missing must "
+                             "be one of {error, ignore}, got "
+                             f"{on_missing_raw!r}")
+        cfg["on_missing"] = on_missing
+
+        for channel, policy_raw in list(cfg.items()):
+            if channel == "on_missing":
+                continue
+            if policy_raw is None:
+                continue
+            if not isinstance(policy_raw, str):
+                raise ValueError("method_config.cfg_uncond values must "
+                                 "be strings, got "
+                                 f"{channel}="
+                                 f"{type(policy_raw).__name__}")
+            policy = policy_raw.strip().lower()
+            allowed = {"keep", "zero", "drop"}
+            if channel == "text":
+                allowed = {*allowed, "negative_prompt"}
+            if policy not in allowed:
+                raise ValueError("method_config.cfg_uncond values must "
+                                 "be one of "
+                                 f"{sorted(allowed)}, got "
+                                 f"{channel}={policy_raw!r}")
+            cfg[channel] = policy
+
+        return cfg
+
+    def _init_optimizers_and_schedulers(self) -> None:
+        tc = self.training_config
+
+        # Student optimizer/scheduler.
+        student_lr = float(tc.optimizer.learning_rate)
+        student_betas = tc.optimizer.betas
+        student_sched = str(tc.optimizer.lr_scheduler)
+        student_params = [p for p in self.student.transformer.parameters() if p.requires_grad]
+        (
+            self._student_optimizer,
+            self._student_lr_scheduler,
+        ) = build_optimizer_and_scheduler(
+            params=student_params,
+            optimizer_config=tc.optimizer,
+            loop_config=tc.loop,
+            learning_rate=student_lr,
+            betas=student_betas,
+            scheduler_name=student_sched,
+        )
+
+        # Critic optimizer/scheduler — must be set in
+        # method config.
+        critic_lr_raw = get_optional_float(
+            self.method_config,
+            "fake_score_learning_rate",
+            where="method.fake_score_learning_rate",
+        )
+        if critic_lr_raw is None or critic_lr_raw == 0.0:
+            raise ValueError("method.fake_score_learning_rate must "
+                             "be set to a positive value")
+        critic_lr = float(critic_lr_raw)
+
+        critic_betas_raw = self.method_config.get("fake_score_betas", None)
+        if critic_betas_raw is None:
+            raise ValueError("method.fake_score_betas must be set "
+                             "(e.g. [0.0, 0.999])")
+        critic_betas = parse_betas(
+            critic_betas_raw,
+            where="method.fake_score_betas",
+        )
+
+        critic_sched_raw = self.method_config.get("fake_score_lr_scheduler", None)
+        if critic_sched_raw is None:
+            raise ValueError("method.fake_score_lr_scheduler must "
+                             "be set (e.g. 'constant')")
+        critic_sched = str(critic_sched_raw)
+        critic_params = [p for p in self.critic.transformer.parameters() if p.requires_grad]
+        (
+            self._critic_optimizer,
+            self._critic_lr_scheduler,
+        ) = build_optimizer_and_scheduler(
+            params=critic_params,
+            optimizer_config=tc.optimizer,
+            loop_config=tc.loop,
+            learning_rate=critic_lr,
+            betas=critic_betas,
+            scheduler_name=critic_sched,
+        )
+
+    def _should_update_student(
+        self,
+        iteration: int,
+    ) -> bool:
+        interval = get_optional_int(
+            self.method_config,
+            "generator_update_interval",
+            where="method.generator_update_interval",
+        )
+        if interval is None:
+            interval = 1
+        if interval <= 0:
+            return True
+        return iteration % interval == 0
+
+    def _get_denoising_step_list(
+        self,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if (self._denoising_step_list is not None and self._denoising_step_list.device == device):
+            return self._denoising_step_list
+
+        raw = self.method_config.get("dmd_denoising_steps", None)
+        if not isinstance(raw, list) or not raw:
+            raise ValueError("method_config.dmd_denoising_steps must "
+                             "be set for DMD2 distillation")
+
+        steps = torch.tensor(
+            [int(s) for s in raw],
+            dtype=torch.long,
+            device=device,
+        )
+
+        warp = self.method_config.get("warp_denoising_step", None)
+        if warp is None:
+            warp = False
+        if bool(warp):
+            timesteps = torch.cat((
+                self.student.noise_scheduler.timesteps.to("cpu"),
+                torch.tensor([0], dtype=torch.float32),
+            )).to(device)
+            steps = timesteps[1000 - steps]
+
+        self._denoising_step_list = steps
+        return steps
+
+    def _sample_rollout_timestep(
+        self,
+        device: torch.device,
+    ) -> torch.Tensor:
+        step_list = self._get_denoising_step_list(device)
+        index = torch.randint(
+            0,
+            len(step_list),
+            [1],
+            device=device,
+            dtype=torch.long,
+            generator=self.cuda_generator,
+        )
+        return step_list[index]
+
+    def _student_rollout(
+        self,
+        batch: Any,
+        *,
+        with_grad: bool,
+    ) -> torch.Tensor:
+        latents = batch.latents
+        device = latents.device
+        dtype = latents.dtype
+        step_list = self._get_denoising_step_list(device)
+
+        if self._rollout_mode != "simulate":
+            timestep = self._sample_rollout_timestep(device)
+            noise = torch.randn(
+                latents.shape,
+                device=device,
+                dtype=dtype,
+                generator=self.cuda_generator,
+            )
+            noisy_latents = self.student.add_noise(latents, noise, timestep)
+            pred_x0 = self.student.predict_x0(
+                noisy_latents,
+                timestep,
+                batch,
+                conditional=True,
+                cfg_uncond=self._cfg_uncond,
+                attn_kind="vsa",
+            )
+            batch.dmd_latent_vis_dict["generator_timestep"] = timestep
+            return pred_x0
+
+        target_timestep_idx = torch.randint(
+            0,
+            len(step_list),
+            [1],
+            device=device,
+            dtype=torch.long,
+            generator=self.cuda_generator,
+        )
+        target_timestep_idx_int = int(target_timestep_idx.item())
+        target_timestep = step_list[target_timestep_idx]
+
+        current_noise_latents = torch.randn(
+            latents.shape,
+            device=device,
+            dtype=dtype,
+            generator=self.cuda_generator,
+        )
+        current_noise_latents_copy = (current_noise_latents.clone())
+
+        max_target_idx = len(step_list) - 1
+        noise_latents: list[torch.Tensor] = []
+        noise_latent_index = target_timestep_idx_int - 1
+
+        if max_target_idx > 0:
+            with torch.no_grad():
+                for step_idx in range(max_target_idx):
+                    current_timestep = step_list[step_idx]
+                    current_timestep_tensor = (current_timestep * torch.ones(
+                        1,
+                        device=device,
+                        dtype=torch.long,
+                    ))
+
+                    pred_clean = self.student.predict_x0(
+                        current_noise_latents,
+                        current_timestep_tensor,
+                        batch,
+                        conditional=True,
+                        cfg_uncond=self._cfg_uncond,
+                        attn_kind="vsa",
+                    )
+
+                    next_timestep = step_list[step_idx + 1]
+                    next_timestep_tensor = (next_timestep * torch.ones(
+                        1,
+                        device=device,
+                        dtype=torch.long,
+                    ))
+                    noise = torch.randn(
+                        latents.shape,
+                        device=device,
+                        dtype=pred_clean.dtype,
+                        generator=self.cuda_generator,
+                    )
+                    current_noise_latents = (self.student.add_noise(
+                        pred_clean,
+                        noise,
+                        next_timestep_tensor,
+                    ))
+                    noise_latents.append(current_noise_latents.clone())
+
+        if noise_latent_index >= 0:
+            if noise_latent_index >= len(noise_latents):
+                raise RuntimeError("noise_latent_index is out of bounds")
+            noisy_input = noise_latents[noise_latent_index]
+        else:
+            noisy_input = current_noise_latents_copy
+
+        if with_grad:
+            pred_x0 = self.student.predict_x0(
+                noisy_input,
+                target_timestep,
+                batch,
+                conditional=True,
+                cfg_uncond=self._cfg_uncond,
+                attn_kind="vsa",
+            )
+        else:
+            with torch.no_grad():
+                pred_x0 = self.student.predict_x0(
+                    noisy_input,
+                    target_timestep,
+                    batch,
+                    conditional=True,
+                    cfg_uncond=self._cfg_uncond,
+                    attn_kind="vsa",
+                )
+
+        batch.dmd_latent_vis_dict["generator_timestep"] = target_timestep.float().detach()
+        return pred_x0
+
+    def _critic_flow_matching_loss(
+        self,
+        batch: Any,
+    ) -> tuple[torch.Tensor, Any, dict[str, Any]]:
+        with torch.no_grad():
+            generator_pred_x0 = self._student_rollout(batch, with_grad=False)
+
+        device = generator_pred_x0.device
+        fake_score_timestep = torch.randint(
+            0,
+            int(self.student.num_train_timesteps),
+            [1],
+            device=device,
+            dtype=torch.long,
+            generator=self.cuda_generator,
+        )
+        fake_score_timestep = (self.student.shift_and_clamp_timestep(fake_score_timestep))
+
+        noise = torch.randn(
+            generator_pred_x0.shape,
+            device=device,
+            dtype=generator_pred_x0.dtype,
+            generator=self.cuda_generator,
+        )
+        noisy_x0 = self.student.add_noise(generator_pred_x0, noise, fake_score_timestep)
+
+        pred_noise = self.critic.predict_noise(
+            noisy_x0,
+            fake_score_timestep,
+            batch,
+            conditional=True,
+            cfg_uncond=self._cfg_uncond,
+            attn_kind="dense",
+        )
+        target = noise - generator_pred_x0
+        flow_matching_loss = torch.mean((pred_noise - target)**2)
+
+        batch.fake_score_latent_vis_dict = {
+            "generator_pred_video": generator_pred_x0,
+            "fake_score_timestep": fake_score_timestep,
+        }
+        outputs = {"fake_score_latent_vis_dict": (batch.fake_score_latent_vis_dict)}
+        return (
+            flow_matching_loss,
+            (batch.timesteps, batch.attn_metadata),
+            outputs,
+        )
+
+    def _dmd_loss(
+        self,
+        generator_pred_x0: torch.Tensor,
+        batch: Any,
+    ) -> torch.Tensor:
+        guidance_scale = get_optional_float(
+            self.method_config,
+            "real_score_guidance_scale",
+            where="method.real_score_guidance_scale",
+        )
+        if guidance_scale is None:
+            guidance_scale = 1.0
+        device = generator_pred_x0.device
+
+        with torch.no_grad():
+            timestep = torch.randint(
+                0,
+                int(self.student.num_train_timesteps),
+                [1],
+                device=device,
+                dtype=torch.long,
+                generator=self.cuda_generator,
+            )
+            timestep = (self.student.shift_and_clamp_timestep(timestep))
+
+            noise = torch.randn(
+                generator_pred_x0.shape,
+                device=device,
+                dtype=generator_pred_x0.dtype,
+                generator=self.cuda_generator,
+            )
+            noisy_latents = self.student.add_noise(generator_pred_x0, noise, timestep)
+
+            faker_x0 = self.critic.predict_x0(
+                noisy_latents,
+                timestep,
+                batch,
+                conditional=True,
+                cfg_uncond=self._cfg_uncond,
+                attn_kind="dense",
+            )
+            real_cond_x0 = self.teacher.predict_x0(
+                noisy_latents,
+                timestep,
+                batch,
+                conditional=True,
+                cfg_uncond=self._cfg_uncond,
+                attn_kind="dense",
+            )
+            real_uncond_x0 = self.teacher.predict_x0(
+                noisy_latents,
+                timestep,
+                batch,
+                conditional=False,
+                cfg_uncond=self._cfg_uncond,
+                attn_kind="dense",
+            )
+            real_cfg_x0 = real_uncond_x0 + (real_cond_x0 - real_uncond_x0) * guidance_scale
+
+            denom = torch.abs(generator_pred_x0 - real_cfg_x0).mean()
+            grad = (faker_x0 - real_cfg_x0) / denom
+            grad = torch.nan_to_num(grad)
+
+        loss = 0.5 * F.mse_loss(
+            generator_pred_x0.float(),
+            (generator_pred_x0.float() - grad.float()).detach(),
+        )
+        return loss

fastvideo/train/methods/distribution_matching/self_forcing.py

@@ -0,0 +1,576 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Self-Forcing distillation method (algorithm layer)."""
+
+from __future__ import annotations
+
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+
+from fastvideo.train.models.base import (
+    CausalModelBase,
+    ModelBase,
+)
+from fastvideo.train.methods.distribution_matching.dmd2 import (
+    DMD2Method, )
+from fastvideo.train.utils.config import (
+    get_optional_float,
+    get_optional_int,
+)
+from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (
+    SelfForcingFlowMatchScheduler, )
+from fastvideo.models.utils import pred_noise_to_pred_video
+
+if TYPE_CHECKING:
+    from fastvideo.pipelines import TrainingBatch
+
+
+def _require_bool(raw: Any, *, where: str) -> bool:
+    if isinstance(raw, bool):
+        return raw
+    raise ValueError(f"Expected bool at {where}, got {type(raw).__name__}")
+
+
+def _require_str(raw: Any, *, where: str) -> str:
+    if not isinstance(raw, str) or not raw.strip():
+        raise ValueError(f"Expected non-empty string at {where}")
+    return raw
+
+
+class SelfForcingMethod(DMD2Method):
+    """Self-Forcing DMD2 (distribution matching) method.
+
+    Requires a causal student implementing ``CausalModelBase``.
+    """
+
+    def __init__(
+        self,
+        *,
+        cfg: Any,
+        role_models: dict[str, ModelBase],
+    ) -> None:
+        super().__init__(
+            cfg=cfg,
+            role_models=role_models,
+        )
+
+        # Validate causal student.
+        if not isinstance(self.student, CausalModelBase):
+            raise ValueError("SelfForcingMethod requires a causal student "
+                             "implementing CausalModelBase.")
+
+        if self._rollout_mode != "simulate":
+            raise ValueError("SelfForcingMethod only supports "
+                             "method_config.rollout_mode='simulate'")
+
+        mcfg = self.method_config
+
+        chunk_size = get_optional_int(
+            mcfg,
+            "chunk_size",
+            where="method_config.chunk_size",
+        )
+        if chunk_size is None:
+            chunk_size = 3
+        if chunk_size <= 0:
+            raise ValueError("method_config.chunk_size must be a positive "
+                             f"integer, got {chunk_size}")
+        self._chunk_size = int(chunk_size)
+
+        sample_type_raw = mcfg.get("student_sample_type", "sde")
+        sample_type = _require_str(
+            sample_type_raw,
+            where="method_config.student_sample_type",
+        )
+        sample_type = sample_type.strip().lower()
+        if sample_type not in {"sde", "ode"}:
+            raise ValueError("method_config.student_sample_type must be one "
+                             f"of {{sde, ode}}, got {sample_type_raw!r}")
+        self._student_sample_type: Literal["sde", "ode"] = (
+            sample_type  # type: ignore[assignment]
+        )
+
+        same_step_raw = mcfg.get("same_step_across_blocks", False)
+        if same_step_raw is None:
+            same_step_raw = False
+        self._same_step_across_blocks = _require_bool(
+            same_step_raw,
+            where="method_config.same_step_across_blocks",
+        )
+
+        last_step_raw = mcfg.get("last_step_only", False)
+        if last_step_raw is None:
+            last_step_raw = False
+        self._last_step_only = _require_bool(
+            last_step_raw,
+            where="method_config.last_step_only",
+        )
+
+        context_noise = get_optional_float(
+            mcfg,
+            "context_noise",
+            where="method_config.context_noise",
+        )
+        if context_noise is None:
+            context_noise = 0.0
+        if context_noise < 0.0:
+            raise ValueError("method_config.context_noise must be >= 0, "
+                             f"got {context_noise}")
+        self._context_noise = float(context_noise)
+
+        enable_grad_raw = mcfg.get("enable_gradient_in_rollout", True)
+        if enable_grad_raw is None:
+            enable_grad_raw = True
+        self._enable_gradient_in_rollout = _require_bool(
+            enable_grad_raw,
+            where="method_config.enable_gradient_in_rollout",
+        )
+
+        start_grad_frame = get_optional_int(
+            mcfg,
+            "start_gradient_frame",
+            where="method_config.start_gradient_frame",
+        )
+        if start_grad_frame is None:
+            start_grad_frame = 0
+        if start_grad_frame < 0:
+            raise ValueError("method_config.start_gradient_frame must be "
+                             f">= 0, got {start_grad_frame}")
+        self._start_gradient_frame = int(start_grad_frame)
+
+        shift = float(getattr(
+            self.training_config.pipeline_config,
+            "flow_shift",
+            0.0,
+        ) or 0.0)
+        self._sf_scheduler = SelfForcingFlowMatchScheduler(
+            num_inference_steps=1000,
+            num_train_timesteps=int(self.student.num_train_timesteps),
+            shift=shift,
+            sigma_min=0.0,
+            extra_one_step=True,
+            training=True,
+        )
+
+        self._sf_denoising_step_list: torch.Tensor | None = None
+
+    def _get_denoising_step_list(self, device: torch.device) -> torch.Tensor:
+        if (self._sf_denoising_step_list is not None and self._sf_denoising_step_list.device == device):
+            return self._sf_denoising_step_list
+
+        raw = self.method_config.get("dmd_denoising_steps", None)
+        if not isinstance(raw, list) or not raw:
+            raise ValueError("method_config.dmd_denoising_steps must be set "
+                             "for self_forcing")
+        steps = torch.tensor(
+            [int(s) for s in raw],
+            dtype=torch.long,
+            device=device,
+        )
+
+        warp = self.method_config.get("warp_denoising_step", None)
+        if warp is None:
+            warp = False
+        if bool(warp):
+            timesteps = torch.cat((
+                self._sf_scheduler.timesteps.to("cpu"),
+                torch.tensor([0], dtype=torch.float32),
+            )).to(device)
+            steps = timesteps[int(self.student.num_train_timesteps) - steps]
+
+        self._sf_denoising_step_list = steps
+        return steps
+
+    def _predict_x0_with_scheduler(
+        self,
+        model: ModelBase,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        attn_kind: Literal["dense", "vsa"],
+    ) -> torch.Tensor:
+        pred_noise = model.predict_noise(
+            noisy_latents,
+            timestep,
+            batch,
+            conditional=conditional,
+            cfg_uncond=self._cfg_uncond,
+            attn_kind=attn_kind,
+        )
+        pred_x0 = pred_noise_to_pred_video(
+            pred_noise=pred_noise.flatten(0, 1),
+            noise_input_latent=noisy_latents.flatten(0, 1),
+            timestep=timestep,
+            scheduler=self._sf_scheduler,
+        ).unflatten(0, pred_noise.shape[:2])
+        return pred_x0
+
+    def _sf_add_noise(
+        self,
+        clean_latents: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        b, t = clean_latents.shape[:2]
+        noisy = self._sf_scheduler.add_noise(
+            clean_latents.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep,
+        ).unflatten(0, (b, t))
+        return noisy
+
+    def _timestep_to_sigma(self, timestep: torch.Tensor) -> torch.Tensor:
+        sigmas = self._sf_scheduler.sigmas.to(device=timestep.device, dtype=torch.float32)
+        timesteps = self._sf_scheduler.timesteps.to(device=timestep.device, dtype=torch.float32)
+        t = timestep.to(device=timestep.device, dtype=torch.float32)
+        if t.ndim == 2:
+            t = t.flatten(0, 1)
+        elif t.ndim == 1 and t.numel() == 1:
+            t = t.expand(1)
+        elif t.ndim != 1:
+            raise ValueError("Invalid timestep shape: "
+                             f"{tuple(timestep.shape)}")
+        idx = torch.argmin(
+            (timesteps.unsqueeze(0) - t.unsqueeze(1)).abs(),
+            dim=1,
+        )
+        return sigmas[idx]
+
+    def _sample_exit_indices(
+        self,
+        *,
+        num_blocks: int,
+        num_steps: int,
+        device: torch.device,
+    ) -> list[int]:
+        if num_blocks <= 0:
+            return []
+        if num_steps <= 0:
+            raise ValueError("num_steps must be positive")
+
+        shape = ((1, ) if self._same_step_across_blocks else (num_blocks, ))
+
+        if not dist.is_initialized() or dist.get_rank() == 0:
+            if self._last_step_only:
+                indices = torch.full(
+                    shape,
+                    num_steps - 1,
+                    dtype=torch.long,
+                    device=device,
+                )
+            else:
+                indices = torch.randint(
+                    low=0,
+                    high=num_steps,
+                    size=shape,
+                    device=device,
+                    generator=self.cuda_generator,
+                )
+        else:
+            indices = torch.empty(shape, dtype=torch.long, device=device)
+
+        if dist.is_initialized():
+            dist.broadcast(indices, src=0)
+
+        if self._same_step_across_blocks:
+            return [int(indices.item()) for _ in range(num_blocks)]
+        return [int(i) for i in indices.tolist()]
+
+    def _student_rollout(self, batch: Any, *, with_grad: bool) -> torch.Tensor:
+        if not isinstance(self.student, CausalModelBase):
+            raise ValueError("SelfForcingMethod requires a causal student "
+                             "implementing CausalModelBase.")
+        return self._student_rollout_streaming(batch, with_grad=with_grad)
+
+    def _student_rollout_streaming(self, batch: Any, *, with_grad: bool) -> torch.Tensor:
+        assert isinstance(self.student, CausalModelBase)
+        latents = batch.latents
+        if latents is None:
+            raise RuntimeError("TrainingBatch.latents is required for "
+                               "self-forcing rollout")
+        if latents.ndim != 5:
+            raise ValueError("TrainingBatch.latents must be [B, T, C, H, W]"
+                             f", got shape={tuple(latents.shape)}")
+
+        device = latents.device
+        dtype = latents.dtype
+        batch_size = int(latents.shape[0])
+        num_frames = int(latents.shape[1])
+
+        denoising_steps = self._get_denoising_step_list(device)
+        num_steps = int(denoising_steps.numel())
+
+        noise_full = torch.randn_like(latents, device=device, dtype=dtype)
+
+        chunk = int(self._chunk_size)
+        if chunk <= 0:
+            raise ValueError("chunk_size must be positive")
+
+        remaining = num_frames % chunk
+        num_blocks = num_frames // chunk
+        if num_blocks == 0:
+            num_blocks = 1
+            remaining = num_frames
+
+        exit_indices = self._sample_exit_indices(
+            num_blocks=num_blocks,
+            num_steps=num_steps,
+            device=device,
+        )
+
+        denoised_blocks: list[torch.Tensor] = []
+
+        cache_tag = "pos"
+        self.student.clear_caches(cache_tag=cache_tag)
+
+        for block_idx in range(num_blocks):
+            if block_idx == 0:
+                start = 0
+                end = remaining + chunk if remaining else chunk
+            else:
+                start = remaining + block_idx * chunk
+                end = remaining + (block_idx + 1) * chunk
+            start = int(start)
+            end = int(min(end, num_frames))
+            if start >= end:
+                break
+
+            noisy_block = noise_full[:, start:end]
+            exit_idx = int(exit_indices[block_idx])
+
+            for step_idx, current_timestep in enumerate(denoising_steps):
+                exit_flag = step_idx == exit_idx
+
+                timestep_block = (current_timestep * torch.ones(
+                    (batch_size, end - start),
+                    device=device,
+                    dtype=torch.float32,
+                ))
+
+                enable_grad = (bool(with_grad) and bool(self._enable_gradient_in_rollout) and torch.is_grad_enabled()
+                               and start >= int(self._start_gradient_frame))
+
+                if not exit_flag:
+                    with torch.no_grad():
+                        pred_noise = (self.student.predict_noise_streaming(
+                            noisy_block,
+                            timestep_block,
+                            batch,
+                            conditional=True,
+                            cache_tag=cache_tag,
+                            store_kv=False,
+                            cur_start_frame=start,
+                            cfg_uncond=self._cfg_uncond,
+                            attn_kind="vsa",
+                        ))
+                        if pred_noise is None:
+                            raise RuntimeError("predict_noise_streaming "
+                                               "returned None "
+                                               "(store_kv=False)")
+                        pred_x0_chunk = pred_noise_to_pred_video(
+                            pred_noise=pred_noise.flatten(0, 1),
+                            noise_input_latent=(noisy_block.flatten(0, 1)),
+                            timestep=timestep_block,
+                            scheduler=self._sf_scheduler,
+                        ).unflatten(0, pred_noise.shape[:2])
+
+                    if step_idx + 1 >= num_steps:
+                        break
+                    next_timestep = denoising_steps[step_idx + 1]
+                    if self._student_sample_type == "sde":
+                        noisy_block = self._sf_add_noise(
+                            pred_x0_chunk,
+                            torch.randn_like(pred_x0_chunk),
+                            next_timestep * torch.ones(
+                                (batch_size, end - start),
+                                device=device,
+                                dtype=torch.float32,
+                            ),
+                        )
+                    else:
+                        sigma_cur = self._timestep_to_sigma(timestep_block).view(batch_size, end - start, 1, 1, 1)
+                        sigma_next = self._timestep_to_sigma(next_timestep * torch.ones(
+                            (batch_size, end - start),
+                            device=device,
+                            dtype=torch.float32,
+                        )).view(batch_size, end - start, 1, 1, 1)
+                        eps = (noisy_block - (1 - sigma_cur) * pred_x0_chunk) / sigma_cur.clamp_min(1e-8)
+                        noisy_block = ((1 - sigma_next) * pred_x0_chunk + sigma_next * eps)
+                    continue
+
+                with torch.set_grad_enabled(enable_grad):
+                    pred_noise = (self.student.predict_noise_streaming(
+                        noisy_block,
+                        timestep_block,
+                        batch,
+                        conditional=True,
+                        cache_tag=cache_tag,
+                        store_kv=False,
+                        cur_start_frame=start,
+                        cfg_uncond=self._cfg_uncond,
+                        attn_kind="vsa",
+                    ))
+                    if pred_noise is None:
+                        raise RuntimeError("predict_noise_streaming returned "
+                                           "None (store_kv=False)")
+                    pred_x0_chunk = pred_noise_to_pred_video(
+                        pred_noise=pred_noise.flatten(0, 1),
+                        noise_input_latent=(noisy_block.flatten(0, 1)),
+                        timestep=timestep_block,
+                        scheduler=self._sf_scheduler,
+                    ).unflatten(0, pred_noise.shape[:2])
+                break
+
+            denoised_blocks.append(pred_x0_chunk)
+
+            with torch.no_grad():
+                if self._context_noise > 0.0:
+                    context_timestep = torch.ones(
+                        (batch_size, end - start),
+                        device=device,
+                        dtype=torch.float32,
+                    ) * float(self._context_noise)
+                    context_latents = self._sf_add_noise(
+                        pred_x0_chunk.detach(),
+                        torch.randn_like(pred_x0_chunk),
+                        context_timestep,
+                    )
+                else:
+                    context_timestep = torch.zeros(
+                        (batch_size, end - start),
+                        device=device,
+                        dtype=torch.float32,
+                    )
+                    context_latents = pred_x0_chunk.detach()
+
+                _ = self.student.predict_noise_streaming(
+                    context_latents,
+                    context_timestep,
+                    batch,
+                    conditional=True,
+                    cache_tag=cache_tag,
+                    store_kv=True,
+                    cur_start_frame=start,
+                    cfg_uncond=self._cfg_uncond,
+                    attn_kind="vsa",
+                )
+
+        if not denoised_blocks:
+            raise RuntimeError("Self-forcing rollout produced no blocks")
+
+        self.student.clear_caches(cache_tag=cache_tag)
+        return torch.cat(denoised_blocks, dim=1)
+
+    def _critic_flow_matching_loss(self, batch: Any) -> tuple[torch.Tensor, Any, dict[str, Any]]:
+        with torch.no_grad():
+            generator_pred_x0 = self._student_rollout(batch, with_grad=False)
+
+        device = generator_pred_x0.device
+        fake_score_timestep = torch.randint(
+            0,
+            int(self.student.num_train_timesteps),
+            [1],
+            device=device,
+            dtype=torch.long,
+            generator=self.cuda_generator,
+        )
+        fake_score_timestep = (self.student.shift_and_clamp_timestep(fake_score_timestep))
+
+        noise = torch.randn(
+            generator_pred_x0.shape,
+            device=device,
+            dtype=generator_pred_x0.dtype,
+            generator=self.cuda_generator,
+        )
+        noisy_x0 = self._sf_add_noise(generator_pred_x0, noise, fake_score_timestep)
+
+        pred_noise = self.critic.predict_noise(
+            noisy_x0,
+            fake_score_timestep,
+            batch,
+            conditional=True,
+            cfg_uncond=self._cfg_uncond,
+            attn_kind="dense",
+        )
+        target = noise - generator_pred_x0
+        flow_matching_loss = torch.mean((pred_noise - target)**2)
+
+        batch.fake_score_latent_vis_dict = {
+            "generator_pred_video": generator_pred_x0,
+            "fake_score_timestep": fake_score_timestep,
+        }
+        outputs = {"fake_score_latent_vis_dict": (batch.fake_score_latent_vis_dict)}
+        return (
+            flow_matching_loss,
+            (batch.timesteps, batch.attn_metadata),
+            outputs,
+        )
+
+    def _dmd_loss(
+        self,
+        generator_pred_x0: torch.Tensor,
+        batch: Any,
+    ) -> torch.Tensor:
+        guidance_scale = get_optional_float(
+            self.method_config,
+            "real_score_guidance_scale",
+            where="method.real_score_guidance_scale",
+        )
+        if guidance_scale is None:
+            guidance_scale = 1.0
+        device = generator_pred_x0.device
+
+        with torch.no_grad():
+            timestep = torch.randint(
+                0,
+                int(self.student.num_train_timesteps),
+                [1],
+                device=device,
+                dtype=torch.long,
+                generator=self.cuda_generator,
+            )
+            timestep = self.student.shift_and_clamp_timestep(timestep)
+
+            noise = torch.randn(
+                generator_pred_x0.shape,
+                device=device,
+                dtype=generator_pred_x0.dtype,
+                generator=self.cuda_generator,
+            )
+            noisy_latents = self._sf_add_noise(generator_pred_x0, noise, timestep)
+
+            faker_x0 = self._predict_x0_with_scheduler(
+                self.critic,
+                noisy_latents,
+                timestep,
+                batch,
+                conditional=True,
+                attn_kind="dense",
+            )
+            real_cond_x0 = self._predict_x0_with_scheduler(
+                self.teacher,
+                noisy_latents,
+                timestep,
+                batch,
+                conditional=True,
+                attn_kind="dense",
+            )
+            real_uncond_x0 = self._predict_x0_with_scheduler(
+                self.teacher,
+                noisy_latents,
+                timestep,
+                batch,
+                conditional=False,
+                attn_kind="dense",
+            )
+            real_cfg_x0 = real_uncond_x0 + (real_cond_x0 - real_uncond_x0) * guidance_scale
+
+            denom = torch.abs(generator_pred_x0 - real_cfg_x0).mean()
+            grad = (faker_x0 - real_cfg_x0) / denom
+            grad = torch.nan_to_num(grad)
+
+        loss = 0.5 * torch.mean((generator_pred_x0.float() - (generator_pred_x0.float() - grad.float()).detach())**2)
+        return loss

fastvideo/train/methods/fine_tuning/__init__.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from fastvideo.train.methods.fine_tuning.dfsft import DiffusionForcingSFTMethod
+    from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
+
+__all__ = [
+    "DiffusionForcingSFTMethod",
+    "FineTuneMethod",
+]
+
+
+def __getattr__(name: str) -> object:
+    # Lazy import to avoid circular imports during registry bring-up.
+    if name == "DiffusionForcingSFTMethod":
+        from fastvideo.train.methods.fine_tuning.dfsft import (
+            DiffusionForcingSFTMethod, )
+
+        return DiffusionForcingSFTMethod
+    if name == "FineTuneMethod":
+        from fastvideo.train.methods.fine_tuning.finetune import FineTuneMethod
+
+        return FineTuneMethod
+    raise AttributeError(name)

fastvideo/train/methods/fine_tuning/dfsft.py

@@ -0,0 +1,325 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Diffusion-forcing SFT method (DFSFT; algorithm layer)."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import torch
+import torch.nn.functional as F
+
+from fastvideo.train.methods.base import TrainingMethod, LogScalar
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.optimizer import (
+    build_optimizer_and_scheduler, )
+
+
+class DiffusionForcingSFTMethod(TrainingMethod):
+    """Diffusion-forcing SFT (DFSFT): train only ``student``
+    with inhomogeneous timesteps.
+    """
+
+    def __init__(
+        self,
+        *,
+        cfg: Any,
+        role_models: dict[str, ModelBase],
+    ) -> None:
+        super().__init__(cfg=cfg, role_models=role_models)
+
+        if "student" not in role_models:
+            raise ValueError("DFSFT requires role 'student'")
+        if not self.student._trainable:
+            raise ValueError("DFSFT requires student to be trainable")
+        self._attn_kind: Literal["dense", "vsa"] = (self._parse_attn_kind(self.method_config.get("attn_kind", None)))
+
+        self._chunk_size = self._parse_chunk_size(self.method_config.get("chunk_size", None))
+        self._timestep_index_range = (self._parse_timestep_index_range())
+
+        # Initialize preprocessors on student.
+        self.student.init_preprocessors(self.training_config)
+
+        self._init_optimizers_and_schedulers()
+
+    @property
+    def _optimizer_dict(self) -> dict[str, Any]:
+        return {"student": self._student_optimizer}
+
+    @property
+    def _lr_scheduler_dict(self) -> dict[str, Any]:
+        return {"student": self._student_lr_scheduler}
+
+    # TrainingMethod override: single_train_step
+    def single_train_step(
+        self,
+        batch: dict[str, Any],
+        iteration: int,
+        *,
+        current_vsa_sparsity: float = 0.0,
+    ) -> tuple[
+            dict[str, torch.Tensor],
+            dict[str, Any],
+            dict[str, LogScalar],
+    ]:
+        del iteration
+        training_batch = self.student.prepare_batch(
+            batch,
+            generator=self.cuda_generator,
+            current_vsa_sparsity=current_vsa_sparsity,
+            latents_source="data",
+        )
+
+        if training_batch.latents is None:
+            raise RuntimeError("prepare_batch() must set TrainingBatch.latents")
+
+        clean_latents = training_batch.latents
+        if not torch.is_tensor(clean_latents):
+            raise TypeError("TrainingBatch.latents must be a torch.Tensor")
+        if clean_latents.ndim != 5:
+            raise ValueError("TrainingBatch.latents must be "
+                             "[B, T, C, H, W], got "
+                             f"shape={tuple(clean_latents.shape)}")
+
+        batch_size, num_latents = (
+            int(clean_latents.shape[0]),
+            int(clean_latents.shape[1]),
+        )
+
+        expected_chunk = getattr(
+            self.student.transformer,
+            "num_frame_per_block",
+            None,
+        )
+        if (expected_chunk is not None and int(expected_chunk) != int(self._chunk_size)):
+            raise ValueError("DFSFT chunk_size must match "
+                             "transformer.num_frame_per_block for "
+                             f"causal training (got {self._chunk_size}, "
+                             f"expected {expected_chunk}).")
+
+        timestep_indices = self._sample_t_inhom_indices(
+            batch_size=batch_size,
+            num_latents=num_latents,
+            device=clean_latents.device,
+        )
+        sp_size = int(self.training_config.distributed.sp_size)
+        sp_group = getattr(self.student, "sp_group", None)
+        if (sp_size > 1 and sp_group is not None and hasattr(sp_group, "broadcast")):
+            sp_group.broadcast(timestep_indices, src=0)
+
+        scheduler = self.student.noise_scheduler
+        if scheduler is None:
+            raise ValueError("DFSFT requires student.noise_scheduler")
+
+        schedule_timesteps = scheduler.timesteps.to(device=clean_latents.device, dtype=torch.float32)
+        schedule_sigmas = scheduler.sigmas.to(
+            device=clean_latents.device,
+            dtype=clean_latents.dtype,
+        )
+        t_inhom = schedule_timesteps[timestep_indices]
+
+        # Override the homogeneous timesteps from prepare_batch
+        # so that set_forward_context (in predict_noise and
+        # backward) receives the correct per-chunk timesteps.
+        training_batch.timesteps = t_inhom
+
+        noise = getattr(training_batch, "noise", None)
+        if noise is None:
+            noise = torch.randn_like(clean_latents)
+        else:
+            if not torch.is_tensor(noise):
+                raise TypeError("TrainingBatch.noise must be a "
+                                "torch.Tensor when set")
+            noise = noise.permute(0, 2, 1, 3, 4).to(dtype=clean_latents.dtype)
+
+        noisy_latents = self.student.add_noise(
+            clean_latents,
+            noise,
+            t_inhom.flatten(),
+        )
+
+        pred = self.student.predict_noise(
+            noisy_latents,
+            t_inhom,
+            training_batch,
+            conditional=True,
+            attn_kind=self._attn_kind,
+        )
+
+        if bool(self.training_config.model.precondition_outputs):
+            sigmas = schedule_sigmas[timestep_indices]
+            sigmas = sigmas.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
+            pred_x0 = noisy_latents - pred * sigmas
+            loss = F.mse_loss(pred_x0.float(), clean_latents.float())
+        else:
+            target = noise - clean_latents
+            loss = F.mse_loss(pred.float(), target.float())
+
+        if self._attn_kind == "vsa":
+            attn_metadata = training_batch.attn_metadata_vsa
+        else:
+            attn_metadata = training_batch.attn_metadata
+
+        loss_map = {"total_loss": loss, "dfsft_loss": loss}
+        outputs: dict[str, Any] = {
+            "_fv_backward": (
+                training_batch.timesteps,
+                attn_metadata,
+            )
+        }
+        metrics: dict[str, LogScalar] = {}
+        return loss_map, outputs, metrics
+
+    # TrainingMethod override: backward
+    def backward(
+        self,
+        loss_map: dict[str, torch.Tensor],
+        outputs: dict[str, Any],
+        *,
+        grad_accum_rounds: int = 1,
+    ) -> None:
+        grad_accum_rounds = max(1, int(grad_accum_rounds))
+        ctx = outputs.get("_fv_backward")
+        if ctx is None:
+            super().backward(
+                loss_map,
+                outputs,
+                grad_accum_rounds=grad_accum_rounds,
+            )
+            return
+        self.student.backward(
+            loss_map["total_loss"],
+            ctx,
+            grad_accum_rounds=grad_accum_rounds,
+        )
+
+    # TrainingMethod override: get_optimizers
+    def get_optimizers(
+        self,
+        iteration: int,
+    ) -> list[torch.optim.Optimizer]:
+        del iteration
+        return [self._student_optimizer]
+
+    # TrainingMethod override: get_lr_schedulers
+    def get_lr_schedulers(
+        self,
+        iteration: int,
+    ) -> list[Any]:
+        del iteration
+        return [self._student_lr_scheduler]
+
+    def _parse_chunk_size(self, raw: Any) -> int:
+        if raw in (None, ""):
+            return 3
+        if isinstance(raw, bool):
+            raise ValueError("method_config.chunk_size must be an int, "
+                             "got bool")
+        if isinstance(raw, float) and not raw.is_integer():
+            raise ValueError("method_config.chunk_size must be an int, "
+                             "got float")
+        if isinstance(raw, str) and not raw.strip():
+            raise ValueError("method_config.chunk_size must be an int, "
+                             "got empty string")
+        try:
+            value = int(raw)
+        except (TypeError, ValueError) as e:
+            raise ValueError("method_config.chunk_size must be an int, "
+                             f"got {type(raw).__name__}") from e
+        if value <= 0:
+            raise ValueError("method_config.chunk_size must be > 0")
+        return value
+
+    def _parse_ratio(
+        self,
+        raw: Any,
+        *,
+        where: str,
+        default: float,
+    ) -> float:
+        if raw in (None, ""):
+            return float(default)
+        if isinstance(raw, bool):
+            raise ValueError(f"{where} must be a number/string, got bool")
+        if isinstance(raw, int | float):
+            return float(raw)
+        if isinstance(raw, str) and raw.strip():
+            return float(raw)
+        raise ValueError(f"{where} must be a number/string, "
+                         f"got {type(raw).__name__}")
+
+    def _parse_timestep_index_range(self, ) -> tuple[int, int]:
+        scheduler = self.student.noise_scheduler
+        if scheduler is None:
+            raise ValueError("DFSFT requires student.noise_scheduler")
+        num_steps = int(getattr(scheduler, "config", scheduler).num_train_timesteps)
+
+        min_ratio = self._parse_ratio(
+            self.method_config.get("min_timestep_ratio", None),
+            where="method.min_timestep_ratio",
+            default=0.0,
+        )
+        max_ratio = self._parse_ratio(
+            self.method_config.get("max_timestep_ratio", None),
+            where="method.max_timestep_ratio",
+            default=1.0,
+        )
+
+        if not (0.0 <= min_ratio <= 1.0 and 0.0 <= max_ratio <= 1.0):
+            raise ValueError("DFSFT timestep ratios must be in [0,1], "
+                             f"got min={min_ratio}, max={max_ratio}")
+        if max_ratio < min_ratio:
+            raise ValueError("method_config.max_timestep_ratio must be "
+                             ">= min_timestep_ratio")
+
+        min_index = int(min_ratio * num_steps)
+        max_index = int(max_ratio * num_steps)
+        min_index = max(0, min(min_index, num_steps - 1))
+        max_index = max(0, min(max_index, num_steps - 1))
+
+        if max_index <= min_index:
+            max_index = min(num_steps - 1, min_index + 1)
+
+        return min_index, max_index + 1
+
+    def _init_optimizers_and_schedulers(self) -> None:
+        tc = self.training_config
+        student_lr = float(tc.optimizer.learning_rate)
+        if student_lr <= 0.0:
+            raise ValueError("training.learning_rate must be > 0 "
+                             "for dfsft")
+
+        student_betas = tc.optimizer.betas
+        student_sched = str(tc.optimizer.lr_scheduler)
+        student_params = [p for p in self.student.transformer.parameters() if p.requires_grad]
+        (
+            self._student_optimizer,
+            self._student_lr_scheduler,
+        ) = build_optimizer_and_scheduler(
+            params=student_params,
+            optimizer_config=tc.optimizer,
+            loop_config=tc.loop,
+            learning_rate=student_lr,
+            betas=student_betas,
+            scheduler_name=student_sched,
+        )
+
+    def _sample_t_inhom_indices(
+        self,
+        *,
+        batch_size: int,
+        num_latents: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        chunk_size = self._chunk_size
+        num_chunks = ((num_latents + chunk_size - 1) // chunk_size)
+        low, high = self._timestep_index_range
+        chunk_indices = torch.randint(
+            low=low,
+            high=high,
+            size=(batch_size, num_chunks),
+            device=device,
+            dtype=torch.long,
+            generator=self.cuda_generator,
+        )
+        expanded = chunk_indices.repeat_interleave(chunk_size, dim=1)
+        return expanded[:, :num_latents]

fastvideo/train/methods/fine_tuning/finetune.py

@@ -0,0 +1,183 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Supervised finetuning method (algorithm layer)."""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+import torch
+import torch.nn.functional as F
+
+from fastvideo.train.methods.base import TrainingMethod, LogScalar
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.optimizer import (
+    build_optimizer_and_scheduler, )
+
+
+class FineTuneMethod(TrainingMethod):
+    """Supervised finetuning: only ``student`` participates."""
+
+    def __init__(
+        self,
+        *,
+        cfg: Any,
+        role_models: dict[str, ModelBase],
+    ) -> None:
+        super().__init__(cfg=cfg, role_models=role_models)
+
+        if "student" not in role_models:
+            raise ValueError("FineTuneMethod requires role 'student'")
+        if not self.student._trainable:
+            raise ValueError("FineTuneMethod requires student to be "
+                             "trainable")
+        self._attn_kind: Literal["dense", "vsa"] = (self._parse_attn_kind(self.method_config.get("attn_kind", None)))
+
+        # Initialize preprocessors on student.
+        self.student.init_preprocessors(self.training_config)
+
+        self._init_optimizers_and_schedulers()
+
+    @property
+    def _optimizer_dict(self) -> dict[str, Any]:
+        return {"student": self._student_optimizer}
+
+    @property
+    def _lr_scheduler_dict(self) -> dict[str, Any]:
+        return {"student": self._student_lr_scheduler}
+
+    # TrainingMethod override: single_train_step
+    def single_train_step(
+        self,
+        batch: dict[str, Any],
+        iteration: int,
+        *,
+        current_vsa_sparsity: float = 0.0,
+    ) -> tuple[
+            dict[str, torch.Tensor],
+            dict[str, Any],
+            dict[str, LogScalar],
+    ]:
+        del iteration
+        training_batch = self.student.prepare_batch(
+            batch,
+            generator=self.cuda_generator,
+            current_vsa_sparsity=current_vsa_sparsity,
+            latents_source="data",
+        )
+
+        if training_batch.latents is None:
+            raise RuntimeError("prepare_batch() must set "
+                               "TrainingBatch.latents")
+        if training_batch.noisy_model_input is None:
+            raise RuntimeError("prepare_batch() must set "
+                               "TrainingBatch.noisy_model_input")
+        if training_batch.noise is None:
+            raise RuntimeError("prepare_batch() must set "
+                               "TrainingBatch.noise")
+        if training_batch.sigmas is None:
+            raise RuntimeError("prepare_batch() must set "
+                               "TrainingBatch.sigmas")
+        if training_batch.timesteps is None:
+            raise RuntimeError("prepare_batch() must set "
+                               "TrainingBatch.timesteps")
+
+        clean_latents = training_batch.latents
+        noisy_latents = (training_batch.noisy_model_input.permute(0, 2, 1, 3, 4))
+        noise = training_batch.noise.permute(0, 2, 1, 3, 4)
+        sigmas = training_batch.sigmas
+        timesteps = training_batch.timesteps
+
+        pred = self.student.predict_noise(
+            noisy_latents,
+            timesteps,
+            training_batch,
+            conditional=True,
+            attn_kind=self._attn_kind,
+        )
+
+        if bool(self.training_config.model.precondition_outputs):
+            pred_x0 = noisy_latents - pred * sigmas
+            loss = F.mse_loss(pred_x0.float(), clean_latents.float())
+        else:
+            target = noise - clean_latents
+            loss = F.mse_loss(pred.float(), target.float())
+
+        if self._attn_kind == "vsa":
+            attn_metadata = training_batch.attn_metadata_vsa
+        else:
+            attn_metadata = training_batch.attn_metadata
+
+        loss_map = {
+            "total_loss": loss,
+            "finetune_loss": loss,
+        }
+        outputs: dict[str, Any] = {
+            "_fv_backward": (
+                training_batch.timesteps,
+                attn_metadata,
+            )
+        }
+        metrics: dict[str, LogScalar] = {}
+        return loss_map, outputs, metrics
+
+    # TrainingMethod override: backward
+    def backward(
+        self,
+        loss_map: dict[str, torch.Tensor],
+        outputs: dict[str, Any],
+        *,
+        grad_accum_rounds: int = 1,
+    ) -> None:
+        grad_accum_rounds = max(1, int(grad_accum_rounds))
+        ctx = outputs.get("_fv_backward")
+        if ctx is None:
+            super().backward(
+                loss_map,
+                outputs,
+                grad_accum_rounds=grad_accum_rounds,
+            )
+            return
+        self.student.backward(
+            loss_map["total_loss"],
+            ctx,
+            grad_accum_rounds=grad_accum_rounds,
+        )
+
+    # TrainingMethod override: get_optimizers
+    def get_optimizers(
+        self,
+        iteration: int,
+    ) -> list[torch.optim.Optimizer]:
+        del iteration
+        return [self._student_optimizer]
+
+    # TrainingMethod override: get_lr_schedulers
+    def get_lr_schedulers(
+        self,
+        iteration: int,
+    ) -> list[Any]:
+        del iteration
+        return [self._student_lr_scheduler]
+
+    def _init_optimizers_and_schedulers(self) -> None:
+        tc = self.training_config
+
+        student_lr = float(tc.optimizer.learning_rate)
+        if student_lr <= 0.0:
+            raise ValueError("training.learning_rate must be > 0 "
+                             "for finetune")
+
+        student_betas = tc.optimizer.betas
+        student_sched = str(tc.optimizer.lr_scheduler)
+        student_params = [p for p in self.student.transformer.parameters() if p.requires_grad]
+        (
+            self._student_optimizer,
+            self._student_lr_scheduler,
+        ) = build_optimizer_and_scheduler(
+            params=student_params,
+            optimizer_config=tc.optimizer,
+            loop_config=tc.loop,
+            learning_rate=student_lr,
+            betas=student_betas,
+            scheduler_name=student_sched,
+        )

fastvideo/train/methods/knowledge_distillation/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+
+__all__: list[str] = []

fastvideo/train/models/__init__.py

@@ -0,0 +1,5 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Model build plugins for Phase 2/2.9 distillation.
+
+These are "model plugins" selected by ``recipe.family`` / ``roles.<role>.family``.
+"""

fastvideo/train/models/base.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+
+from fastvideo.models.utils import pred_noise_to_pred_video
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+    from fastvideo.pipelines import TrainingBatch
+
+
+class ModelBase(ABC):
+    """Per-role model instance.
+
+    Every role (student, teacher, critic, …) gets its own ``ModelBase``
+    instance.  Each instance owns its own ``transformer`` and
+    ``noise_scheduler``.  Heavyweight resources (VAE, dataloader, RNG
+    seeds) are loaded lazily via :meth:`init_preprocessors`, which the
+    method calls **only on the student**.
+    """
+
+    transformer: torch.nn.Module
+    noise_scheduler: Any
+    _trainable: bool
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def init_preprocessors(  # noqa: B027
+            self,
+            training_config: TrainingConfig,
+    ) -> None:
+        """Load VAE, build dataloader, seed RNGs.
+
+        Called only on the student by the method's ``__init__``.
+        Default is a no-op so teacher/critic instances skip this.
+        """
+
+    def on_train_start(self) -> None:  # noqa: B027
+        """Called once before the training loop begins."""
+
+    # ------------------------------------------------------------------
+    # Timestep helpers
+    # ------------------------------------------------------------------
+
+    @property
+    def num_train_timesteps(self) -> int:
+        """Return the scheduler's training timestep horizon."""
+        return int(self.noise_scheduler.num_train_timesteps)
+
+    def shift_and_clamp_timestep(self, timestep: torch.Tensor) -> torch.Tensor:
+        """Apply model/pipeline timestep shifting and clamp."""
+        return timestep
+
+    # ------------------------------------------------------------------
+    # Runtime primitives
+    # ------------------------------------------------------------------
+
+    @abstractmethod
+    def prepare_batch(
+        self,
+        raw_batch: dict[str, Any],
+        *,
+        generator: torch.Generator,
+        current_vsa_sparsity: float = 0.0,
+        latents_source: Literal["data", "zeros"] = "data",
+    ) -> TrainingBatch:
+        """Convert a dataloader batch into forward primitives."""
+
+    @abstractmethod
+    def add_noise(
+        self,
+        clean_latents: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply forward-process noise at *timestep*."""
+
+    @abstractmethod
+    def predict_noise(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor:
+        """Predict noise/flow for the given noisy latents."""
+
+    def predict_x0(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor:
+        """Predict x0 via ``predict_noise`` + conversion."""
+        pred_noise = self.predict_noise(
+            noisy_latents,
+            timestep,
+            batch,
+            conditional=conditional,
+            cfg_uncond=cfg_uncond,
+            attn_kind=attn_kind,
+        )
+        return pred_noise_to_pred_video(
+            pred_noise=pred_noise.flatten(0, 1),
+            noise_input_latent=noisy_latents.flatten(0, 1),
+            timestep=timestep,
+            scheduler=self.noise_scheduler,
+        ).unflatten(0, pred_noise.shape[:2])
+
+    @abstractmethod
+    def backward(
+        self,
+        loss: torch.Tensor,
+        ctx: Any,
+        *,
+        grad_accum_rounds: int,
+    ) -> None:
+        """Backward that may restore forward-context."""
+
+
+class CausalModelBase(ModelBase):
+    """Extension for causal / streaming model plugins.
+
+    Cache state is internal to the model instance and keyed by
+    *cache_tag* (no role handle needed).
+    """
+
+    @abstractmethod
+    def clear_caches(self, *, cache_tag: str = "pos") -> None:
+        """Clear internal caches before starting a new rollout."""
+
+    @abstractmethod
+    def predict_noise_streaming(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cache_tag: str = "pos",
+        store_kv: bool = False,
+        cur_start_frame: int = 0,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor | None:
+        """Streaming predict-noise that may update internal caches."""
+
+    def predict_x0_streaming(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cache_tag: str = "pos",
+        store_kv: bool = False,
+        cur_start_frame: int = 0,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor | None:
+        """Predict x0 streaming via
+        ``predict_noise_streaming`` + conversion."""
+        pred_noise = self.predict_noise_streaming(
+            noisy_latents,
+            timestep,
+            batch,
+            conditional=conditional,
+            cache_tag=cache_tag,
+            store_kv=store_kv,
+            cur_start_frame=cur_start_frame,
+            cfg_uncond=cfg_uncond,
+            attn_kind=attn_kind,
+        )
+        if pred_noise is None:
+            return None
+        return pred_noise_to_pred_video(
+            pred_noise=pred_noise.flatten(0, 1),
+            noise_input_latent=noisy_latents.flatten(0, 1),
+            timestep=timestep,
+            scheduler=self.noise_scheduler,
+        ).unflatten(0, pred_noise.shape[:2])

fastvideo/train/models/wan/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan model plugin package."""
+
+from fastvideo.train.models.wan.wan import (
+    WanModel as WanModel, )
+from fastvideo.train.models.wan.wan_causal import (
+    WanCausalModel as WanCausalModel, )

fastvideo/train/models/wan/wan.py

@@ -0,0 +1,668 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan model plugin (per-role instance)."""
+
+from __future__ import annotations
+
+import copy
+import gc
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+
+import fastvideo.envs as envs
+from fastvideo.configs.sample import SamplingParam
+from fastvideo.distributed import (
+    get_local_torch_device,
+    get_sp_group,
+    get_world_group,
+)
+from fastvideo.forward_context import set_forward_context
+from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler, )
+from fastvideo.pipelines import TrainingBatch
+from fastvideo.pipelines.basic.wan.wan_pipeline import (
+    WanPipeline, )
+from fastvideo.pipelines.pipeline_batch_info import (
+    ForwardBatch, )
+from fastvideo.training.activation_checkpoint import (
+    apply_activation_checkpointing, )
+from fastvideo.training.training_utils import (
+    compute_density_for_timestep_sampling,
+    get_sigmas,
+    normalize_dit_input,
+    shift_timestep,
+)
+from fastvideo.utils import (
+    is_vmoba_available,
+    is_vsa_available,
+)
+
+from fastvideo.train.models.base import ModelBase
+from fastvideo.train.utils.module_state import (
+    apply_trainable, )
+from fastvideo.train.utils.moduleloader import (
+    load_module_from_path, )
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+
+try:
+    from fastvideo.attention.backends.video_sparse_attn import (
+        VideoSparseAttentionMetadataBuilder, )
+    from fastvideo.attention.backends.vmoba import (
+        VideoMobaAttentionMetadataBuilder, )
+except Exception:
+    VideoSparseAttentionMetadataBuilder = None  # type: ignore[assignment]
+    VideoMobaAttentionMetadataBuilder = None  # type: ignore[assignment]
+
+
+class WanModel(ModelBase):
+    """Wan per-role model: owns transformer + noise_scheduler."""
+
+    _transformer_cls_name: str = "WanTransformer3DModel"
+
+    def __init__(
+        self,
+        *,
+        init_from: str,
+        training_config: TrainingConfig,
+        trainable: bool = True,
+        disable_custom_init_weights: bool = False,
+        flow_shift: float = 3.0,
+        enable_gradient_checkpointing_type: str
+        | None = None,
+    ) -> None:
+        self._init_from = str(init_from)
+        self._trainable = bool(trainable)
+
+        self.transformer = self._load_transformer(
+            init_from=self._init_from,
+            trainable=self._trainable,
+            disable_custom_init_weights=(disable_custom_init_weights),
+            enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
+            training_config=training_config,
+        )
+
+        self.noise_scheduler = (FlowMatchEulerDiscreteScheduler(shift=float(flow_shift)))
+
+        # Filled by init_preprocessors (student only).
+        self.vae: Any = None
+        self.training_config: TrainingConfig = training_config
+        self.dataloader: Any = None
+        self.validator: Any = None
+        self.start_step: int = 0
+
+        self.world_group: Any = None
+        self.sp_group: Any = None
+        self.device: Any = get_local_torch_device()
+
+        self.negative_prompt_embeds: (torch.Tensor | None) = None
+        self.negative_prompt_attention_mask: (torch.Tensor | None) = None
+
+        # Timestep mechanics.
+        self.timestep_shift: float = float(flow_shift)
+        self.num_train_timestep: int = int(self.noise_scheduler.num_train_timesteps)
+        self.min_timestep: int = 0
+        self.max_timestep: int = self.num_train_timestep
+
+    def _load_transformer(
+        self,
+        *,
+        init_from: str,
+        trainable: bool,
+        disable_custom_init_weights: bool,
+        enable_gradient_checkpointing_type: str | None,
+        training_config: TrainingConfig,
+    ) -> torch.nn.Module:
+        transformer = load_module_from_path(
+            model_path=init_from,
+            module_type="transformer",
+            training_config=training_config,
+            disable_custom_init_weights=(disable_custom_init_weights),
+            override_transformer_cls_name=(self._transformer_cls_name),
+        )
+        transformer = apply_trainable(transformer, trainable=trainable)
+        # Fall back to training_config.model if not set on the
+        # model YAML section directly.
+        ckpt_type = (enable_gradient_checkpointing_type or getattr(
+            getattr(training_config, "model", None),
+            "enable_gradient_checkpointing_type",
+            None,
+        ))
+        if trainable and ckpt_type:
+            transformer = apply_activation_checkpointing(
+                transformer,
+                checkpointing_type=ckpt_type,
+            )
+        return transformer
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def init_preprocessors(self, training_config: TrainingConfig) -> None:
+        self.vae = load_module_from_path(
+            model_path=str(training_config.model_path),
+            module_type="vae",
+            training_config=training_config,
+        )
+
+        self.world_group = get_world_group()
+        self.sp_group = get_sp_group()
+
+        self._init_timestep_mechanics()
+
+        from fastvideo.dataset.dataloader.schema import (
+            pyarrow_schema_t2v, )
+        from fastvideo.train.utils.dataloader import (
+            build_parquet_t2v_train_dataloader, )
+
+        text_len = (
+            training_config.pipeline_config.text_encoder_configs[  # type: ignore[union-attr]
+                0].arch_config.text_len)
+        self.dataloader = build_parquet_t2v_train_dataloader(
+            training_config.data,
+            text_len=int(text_len),
+            parquet_schema=pyarrow_schema_t2v,
+        )
+        self.start_step = 0
+
+    @property
+    def num_train_timesteps(self) -> int:
+        return int(self.num_train_timestep)
+
+    def shift_and_clamp_timestep(self, timestep: torch.Tensor) -> torch.Tensor:
+        timestep = shift_timestep(
+            timestep,
+            self.timestep_shift,
+            self.num_train_timestep,
+        )
+        return timestep.clamp(self.min_timestep, self.max_timestep)
+
+    def on_train_start(self) -> None:
+        self.ensure_negative_conditioning()
+
+    # ------------------------------------------------------------------
+    # Runtime primitives
+    # ------------------------------------------------------------------
+
+    def prepare_batch(
+        self,
+        raw_batch: dict[str, Any],
+        *,
+        generator: torch.Generator,
+        current_vsa_sparsity: float = 0.0,
+        latents_source: Literal["data", "zeros"] = "data",
+    ) -> TrainingBatch:
+        self.ensure_negative_conditioning()
+        assert self.training_config is not None
+        tc = self.training_config
+
+        dtype = self._get_training_dtype()
+        device = self.device
+
+        training_batch = TrainingBatch(current_vsa_sparsity=current_vsa_sparsity)
+        encoder_hidden_states = raw_batch["text_embedding"]
+        encoder_attention_mask = raw_batch["text_attention_mask"]
+        infos = raw_batch.get("info_list")
+
+        if latents_source == "zeros":
+            batch_size = encoder_hidden_states.shape[0]
+            vae_config = (
+                tc.pipeline_config.vae_config.arch_config  # type: ignore[union-attr]
+            )
+            num_channels = vae_config.z_dim
+            spatial_compression_ratio = (vae_config.spatial_compression_ratio)
+            latent_height = (tc.data.num_height // spatial_compression_ratio)
+            latent_width = (tc.data.num_width // spatial_compression_ratio)
+            latents = torch.zeros(
+                batch_size,
+                num_channels,
+                tc.data.num_latent_t,
+                latent_height,
+                latent_width,
+                device=device,
+                dtype=dtype,
+            )
+        elif latents_source == "data":
+            if "vae_latent" not in raw_batch:
+                raise ValueError("vae_latent not found in batch "
+                                 "and latents_source='data'")
+            latents = raw_batch["vae_latent"]
+            latents = latents[:, :, :tc.data.num_latent_t]
+            latents = latents.to(device, dtype=dtype)
+        else:
+            raise ValueError(f"Unknown latents_source: "
+                             f"{latents_source!r}")
+
+        training_batch.latents = latents
+        training_batch.encoder_hidden_states = (encoder_hidden_states.to(device, dtype=dtype))
+        training_batch.encoder_attention_mask = (encoder_attention_mask.to(device, dtype=dtype))
+        training_batch.infos = infos
+
+        training_batch.latents = normalize_dit_input("wan", training_batch.latents, self.vae)
+        training_batch = self._prepare_dit_inputs(training_batch, generator)
+        training_batch = self._build_attention_metadata(training_batch)
+
+        training_batch.attn_metadata_vsa = copy.deepcopy(training_batch.attn_metadata)
+        if training_batch.attn_metadata is not None:
+            training_batch.attn_metadata.VSA_sparsity = 0.0  # type: ignore[attr-defined]
+
+        return training_batch
+
+    def add_noise(
+        self,
+        clean_latents: torch.Tensor,
+        noise: torch.Tensor,
+        timestep: torch.Tensor,
+    ) -> torch.Tensor:
+        b, t = clean_latents.shape[:2]
+        noisy = self.noise_scheduler.add_noise(
+            clean_latents.flatten(0, 1),
+            noise.flatten(0, 1),
+            timestep,
+        ).unflatten(0, (b, t))
+        return noisy
+
+    def predict_noise(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: TrainingBatch,
+        *,
+        conditional: bool,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor:
+        device_type = self.device.type
+        dtype = noisy_latents.dtype
+        if conditional:
+            text_dict = batch.conditional_dict
+            if text_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+        else:
+            text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond)
+
+        if attn_kind == "dense":
+            attn_metadata = batch.attn_metadata
+        elif attn_kind == "vsa":
+            attn_metadata = batch.attn_metadata_vsa
+        else:
+            raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
+
+        with torch.autocast(device_type, dtype=dtype), set_forward_context(
+                current_timestep=batch.timesteps,
+                attn_metadata=attn_metadata,
+        ):
+            input_kwargs = (self._build_distill_input_kwargs(noisy_latents, timestep, text_dict))
+            transformer = self._get_transformer(timestep)
+            pred_noise = transformer(**input_kwargs).permute(0, 2, 1, 3, 4)
+        return pred_noise
+
+    def backward(
+        self,
+        loss: torch.Tensor,
+        ctx: Any,
+        *,
+        grad_accum_rounds: int,
+    ) -> None:
+        timesteps, attn_metadata = ctx
+        with set_forward_context(
+                current_timestep=timesteps,
+                attn_metadata=attn_metadata,
+        ):
+            (loss / max(1, int(grad_accum_rounds))).backward()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _get_training_dtype(self) -> torch.dtype:
+        return torch.bfloat16
+
+    def _init_timestep_mechanics(self) -> None:
+        assert self.training_config is not None
+        tc = self.training_config
+        self.timestep_shift = float(tc.pipeline_config.flow_shift  # type: ignore[union-attr]
+                                    )
+        self.num_train_timestep = int(self.noise_scheduler.num_train_timesteps)
+        # min/max timestep ratios now come from method_config;
+        # default to full range.
+        self.min_timestep = 0
+        self.max_timestep = self.num_train_timestep
+
+    def ensure_negative_conditioning(self) -> None:
+        if self.negative_prompt_embeds is not None:
+            return
+
+        assert self.training_config is not None
+        tc = self.training_config
+        world_group = self.world_group
+        device = self.device
+        dtype = self._get_training_dtype()
+
+        from fastvideo.train.utils.moduleloader import (
+            make_inference_args, )
+
+        neg_embeds: torch.Tensor | None = None
+        neg_mask: torch.Tensor | None = None
+
+        if world_group.rank_in_group == 0:
+            sampling_param = SamplingParam.from_pretrained(tc.model_path)
+            negative_prompt = sampling_param.negative_prompt
+
+            inference_args = make_inference_args(tc, model_path=tc.model_path)
+
+            prompt_pipeline = WanPipeline.from_pretrained(
+                tc.model_path,
+                args=inference_args,
+                inference_mode=True,
+                loaded_modules={"transformer": self.transformer},
+                tp_size=tc.distributed.tp_size,
+                sp_size=tc.distributed.sp_size,
+                num_gpus=tc.distributed.num_gpus,
+                pin_cpu_memory=(tc.distributed.pin_cpu_memory),
+                dit_cpu_offload=True,
+            )
+
+            batch_negative = ForwardBatch(
+                data_type="video",
+                prompt=negative_prompt,
+                prompt_embeds=[],
+                prompt_attention_mask=[],
+            )
+            result_batch = prompt_pipeline.prompt_encoding_stage(  # type: ignore[attr-defined]
+                batch_negative,
+                inference_args,
+            )
+
+            neg_embeds = result_batch.prompt_embeds[0].to(device=device, dtype=dtype)
+            neg_mask = (result_batch.prompt_attention_mask[0].to(device=device, dtype=dtype))
+
+            del prompt_pipeline
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+        meta = torch.zeros((2, ), device=device, dtype=torch.int64)
+        if world_group.rank_in_group == 0:
+            assert neg_embeds is not None
+            assert neg_mask is not None
+            meta[0] = neg_embeds.ndim
+            meta[1] = neg_mask.ndim
+        world_group.broadcast(meta, src=0)
+        embed_ndim, mask_ndim = (
+            int(meta[0].item()),
+            int(meta[1].item()),
+        )
+
+        max_ndim = 8
+        embed_shape = torch.full((max_ndim, ), -1, device=device, dtype=torch.int64)
+        mask_shape = torch.full((max_ndim, ), -1, device=device, dtype=torch.int64)
+        if world_group.rank_in_group == 0:
+            assert neg_embeds is not None
+            assert neg_mask is not None
+            embed_shape[:embed_ndim] = torch.tensor(
+                list(neg_embeds.shape),
+                device=device,
+                dtype=torch.int64,
+            )
+            mask_shape[:mask_ndim] = torch.tensor(
+                list(neg_mask.shape),
+                device=device,
+                dtype=torch.int64,
+            )
+        world_group.broadcast(embed_shape, src=0)
+        world_group.broadcast(mask_shape, src=0)
+
+        embed_sizes = tuple(int(x) for x in embed_shape[:embed_ndim].tolist())
+        mask_sizes = tuple(int(x) for x in mask_shape[:mask_ndim].tolist())
+
+        if world_group.rank_in_group != 0:
+            neg_embeds = torch.empty(embed_sizes, device=device, dtype=dtype)
+            neg_mask = torch.empty(mask_sizes, device=device, dtype=dtype)
+        assert neg_embeds is not None
+        assert neg_mask is not None
+
+        world_group.broadcast(neg_embeds, src=0)
+        world_group.broadcast(neg_mask, src=0)
+
+        self.negative_prompt_embeds = neg_embeds
+        self.negative_prompt_attention_mask = neg_mask
+
+    def _sample_timesteps(
+        self,
+        batch_size: int,
+        device: torch.device,
+        generator: torch.Generator,
+    ) -> torch.Tensor:
+        assert self.training_config is not None
+        tc = self.training_config
+
+        u = compute_density_for_timestep_sampling(
+            weighting_scheme=tc.model.weighting_scheme,
+            batch_size=batch_size,
+            generator=generator,
+            device=device,
+            logit_mean=tc.model.logit_mean,
+            logit_std=tc.model.logit_std,
+            mode_scale=tc.model.mode_scale,
+        )
+        indices = (u * self.noise_scheduler.config.num_train_timesteps).long()
+        return self.noise_scheduler.timesteps[indices.cpu()].to(device=device)
+
+    def _build_attention_metadata(self, training_batch: TrainingBatch) -> TrainingBatch:
+        assert self.training_config is not None
+        tc = self.training_config
+        latents_shape = training_batch.raw_latent_shape
+        patch_size = (
+            tc.pipeline_config.dit_config.patch_size  # type: ignore[union-attr]
+        )
+        current_vsa_sparsity = (training_batch.current_vsa_sparsity)
+        assert latents_shape is not None
+        assert training_batch.timesteps is not None
+
+        if (envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN"):
+            if (not is_vsa_available() or VideoSparseAttentionMetadataBuilder is None):
+                raise ImportError("FASTVIDEO_ATTENTION_BACKEND is "
+                                  "VIDEO_SPARSE_ATTN, but "
+                                  "fastvideo_kernel is not correctly "
+                                  "installed or detected.")
+            training_batch.attn_metadata = VideoSparseAttentionMetadataBuilder().build(  # type: ignore[misc]
+                raw_latent_shape=latents_shape[2:5],
+                current_timestep=(training_batch.timesteps),
+                patch_size=patch_size,
+                VSA_sparsity=current_vsa_sparsity,
+                device=self.device,
+            )
+        elif (envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN"):
+            if (not is_vmoba_available() or VideoMobaAttentionMetadataBuilder is None):
+                raise ImportError("FASTVIDEO_ATTENTION_BACKEND is "
+                                  "VMOBA_ATTN, but fastvideo_kernel "
+                                  "(or flash_attn>=2.7.4) is not "
+                                  "correctly installed.")
+            moba_params = tc.model.moba_config.copy()
+            moba_params.update({
+                "current_timestep": (training_batch.timesteps),
+                "raw_latent_shape": (training_batch.raw_latent_shape[2:5]),
+                "patch_size": patch_size,
+                "device": self.device,
+            })
+            training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder().build(**
+                                                                                     moba_params)  # type: ignore[misc]
+        else:
+            training_batch.attn_metadata = None
+
+        return training_batch
+
+    def _prepare_dit_inputs(
+        self,
+        training_batch: TrainingBatch,
+        generator: torch.Generator,
+    ) -> TrainingBatch:
+        assert self.training_config is not None
+        tc = self.training_config
+        latents = training_batch.latents
+        assert isinstance(latents, torch.Tensor)
+        batch_size = latents.shape[0]
+
+        noise = torch.randn(
+            latents.shape,
+            generator=generator,
+            device=latents.device,
+            dtype=latents.dtype,
+        )
+        timesteps = self._sample_timesteps(
+            batch_size,
+            latents.device,
+            generator,
+        )
+        if int(tc.distributed.sp_size or 1) > 1:
+            self.sp_group.broadcast(timesteps, src=0)
+
+        sigmas = get_sigmas(
+            self.noise_scheduler,
+            latents.device,
+            timesteps,
+            n_dim=latents.ndim,
+            dtype=latents.dtype,
+        )
+        noisy_model_input = ((1.0 - sigmas) * latents + sigmas * noise)
+
+        training_batch.noisy_model_input = (noisy_model_input)
+        training_batch.timesteps = timesteps
+        training_batch.sigmas = sigmas
+        training_batch.noise = noise
+        training_batch.raw_latent_shape = latents.shape
+
+        training_batch.conditional_dict = {
+            "encoder_hidden_states": (training_batch.encoder_hidden_states),
+            "encoder_attention_mask": (training_batch.encoder_attention_mask),
+        }
+
+        if (self.negative_prompt_embeds is not None and self.negative_prompt_attention_mask is not None):
+            neg_embeds = self.negative_prompt_embeds
+            neg_mask = (self.negative_prompt_attention_mask)
+            if (neg_embeds.shape[0] == 1 and batch_size > 1):
+                neg_embeds = neg_embeds.expand(batch_size, *neg_embeds.shape[1:]).contiguous()
+            if (neg_mask.shape[0] == 1 and batch_size > 1):
+                neg_mask = neg_mask.expand(batch_size, *neg_mask.shape[1:]).contiguous()
+            training_batch.unconditional_dict = {
+                "encoder_hidden_states": neg_embeds,
+                "encoder_attention_mask": neg_mask,
+            }
+
+        training_batch.latents = (training_batch.latents.permute(0, 2, 1, 3, 4))
+        return training_batch
+
+    def _build_distill_input_kwargs(
+        self,
+        noise_input: torch.Tensor,
+        timestep: torch.Tensor,
+        text_dict: dict[str, torch.Tensor] | None,
+    ) -> dict[str, Any]:
+        if text_dict is None:
+            raise ValueError("text_dict cannot be None for "
+                             "Wan distillation")
+        return {
+            "hidden_states": noise_input.permute(0, 2, 1, 3, 4),
+            "encoder_hidden_states": text_dict["encoder_hidden_states"],
+            "encoder_attention_mask": text_dict["encoder_attention_mask"],
+            "timestep": timestep,
+            "return_dict": False,
+        }
+
+    def _get_transformer(self, timestep: torch.Tensor) -> torch.nn.Module:
+        return self.transformer
+
+    def _get_uncond_text_dict(
+        self,
+        batch: TrainingBatch,
+        *,
+        cfg_uncond: dict[str, Any] | None,
+    ) -> dict[str, torch.Tensor]:
+        if cfg_uncond is None:
+            text_dict = getattr(batch, "unconditional_dict", None)
+            if text_dict is None:
+                raise RuntimeError("Missing unconditional_dict; "
+                                   "ensure_negative_conditioning() "
+                                   "may have failed")
+            return text_dict
+
+        on_missing_raw = cfg_uncond.get("on_missing", "error")
+        if not isinstance(on_missing_raw, str):
+            raise ValueError("method_config.cfg_uncond.on_missing "
+                             "must be a string, got "
+                             f"{type(on_missing_raw).__name__}")
+        on_missing = on_missing_raw.strip().lower()
+        if on_missing not in {"error", "ignore"}:
+            raise ValueError("method_config.cfg_uncond.on_missing "
+                             "must be one of {error, ignore}, got "
+                             f"{on_missing_raw!r}")
+
+        for channel, policy_raw in cfg_uncond.items():
+            if channel in {"on_missing", "text"}:
+                continue
+            if policy_raw is None:
+                continue
+            if not isinstance(policy_raw, str):
+                raise ValueError("method_config.cfg_uncond values "
+                                 "must be strings, got "
+                                 f"{channel}="
+                                 f"{type(policy_raw).__name__}")
+            policy = policy_raw.strip().lower()
+            if policy == "keep":
+                continue
+            if on_missing == "ignore":
+                continue
+            raise ValueError("WanModel does not support "
+                             "cfg_uncond channel "
+                             f"{channel!r} (policy={policy!r}). "
+                             "Set cfg_uncond.on_missing=ignore or "
+                             "remove the channel.")
+
+        text_policy_raw = cfg_uncond.get("text", None)
+        if text_policy_raw is None:
+            text_policy = "negative_prompt"
+        elif not isinstance(text_policy_raw, str):
+            raise ValueError("method_config.cfg_uncond.text must be "
+                             "a string, got "
+                             f"{type(text_policy_raw).__name__}")
+        else:
+            text_policy = (text_policy_raw.strip().lower())
+
+        if text_policy in {"negative_prompt"}:
+            text_dict = getattr(batch, "unconditional_dict", None)
+            if text_dict is None:
+                raise RuntimeError("Missing unconditional_dict; "
+                                   "ensure_negative_conditioning() "
+                                   "may have failed")
+            return text_dict
+        if text_policy == "keep":
+            if batch.conditional_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+            return batch.conditional_dict
+        if text_policy == "zero":
+            if batch.conditional_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+            cond = batch.conditional_dict
+            enc = cond["encoder_hidden_states"]
+            mask = cond["encoder_attention_mask"]
+            if not torch.is_tensor(enc) or not torch.is_tensor(mask):
+                raise TypeError("conditional_dict must contain "
+                                "tensor text inputs")
+            return {
+                "encoder_hidden_states": (torch.zeros_like(enc)),
+                "encoder_attention_mask": (torch.zeros_like(mask)),
+            }
+        if text_policy == "drop":
+            raise ValueError("cfg_uncond.text=drop is not supported "
+                             "for Wan. Use "
+                             "{negative_prompt, keep, zero}.")
+        raise ValueError("cfg_uncond.text must be one of "
+                         "{negative_prompt, keep, zero, drop}, got "
+                         f"{text_policy_raw!r}")

fastvideo/train/models/wan/wan_causal.py

@@ -0,0 +1,437 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan causal model plugin (per-role instance, streaming/cache)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal, TYPE_CHECKING
+
+import torch
+
+from fastvideo.forward_context import set_forward_context
+
+from fastvideo.train.models.base import CausalModelBase
+from fastvideo.train.models.wan.wan import WanModel
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+
+
+@dataclass(slots=True)
+class _StreamingCaches:
+    kv_cache: list[dict[str, Any]]
+    crossattn_cache: list[dict[str, Any]] | None
+    frame_seq_length: int
+    local_attn_size: int
+    sliding_window_num_frames: int
+    batch_size: int
+    dtype: torch.dtype
+    device: torch.device
+
+
+class WanCausalModel(WanModel, CausalModelBase):
+    """Wan per-role model with causal/streaming primitives."""
+
+    _transformer_cls_name: str = ("CausalWanTransformer3DModel")
+
+    def __init__(
+        self,
+        *,
+        init_from: str,
+        training_config: TrainingConfig,
+        trainable: bool = True,
+        disable_custom_init_weights: bool = False,
+        flow_shift: float = 3.0,
+        enable_gradient_checkpointing_type: str
+        | None = None,
+    ) -> None:
+        super().__init__(
+            init_from=init_from,
+            training_config=training_config,
+            trainable=trainable,
+            disable_custom_init_weights=(disable_custom_init_weights),
+            flow_shift=flow_shift,
+            enable_gradient_checkpointing_type=(enable_gradient_checkpointing_type),
+        )
+        self._streaming_caches: (dict[tuple[int, str], _StreamingCaches]) = {}
+
+    # --- CausalModelBase override: clear_caches ---
+    def clear_caches(
+        self,
+        *,
+        cache_tag: str = "pos",
+    ) -> None:
+        self._streaming_caches.pop((id(self), str(cache_tag)), None)
+
+    # --- CausalModelBase override: predict_noise_streaming ---
+    def predict_noise_streaming(
+        self,
+        noisy_latents: torch.Tensor,
+        timestep: torch.Tensor,
+        batch: Any,
+        *,
+        conditional: bool,
+        cache_tag: str = "pos",
+        store_kv: bool = False,
+        cur_start_frame: int = 0,
+        cfg_uncond: dict[str, Any] | None = None,
+        attn_kind: Literal["dense", "vsa"] = "dense",
+    ) -> torch.Tensor | None:
+        if attn_kind == "dense":
+            attn_metadata = batch.attn_metadata
+        elif attn_kind == "vsa":
+            attn_metadata = batch.attn_metadata_vsa
+        else:
+            raise ValueError(f"Unknown attn_kind: {attn_kind!r}")
+
+        cache_tag = str(cache_tag)
+        cur_start_frame = int(cur_start_frame)
+        if cur_start_frame < 0:
+            raise ValueError("cur_start_frame must be >= 0")
+
+        batch_size = int(noisy_latents.shape[0])
+        num_frames = int(noisy_latents.shape[1])
+        timestep_full = self._ensure_per_frame_timestep(
+            timestep=timestep,
+            batch_size=batch_size,
+            num_frames=num_frames,
+            device=noisy_latents.device,
+        )
+
+        transformer = self._get_transformer(timestep_full)
+        caches = self._get_or_init_streaming_caches(
+            cache_tag=cache_tag,
+            transformer=transformer,
+            noisy_latents=noisy_latents,
+        )
+
+        frame_seq_length = int(caches.frame_seq_length)
+        kv_cache = caches.kv_cache
+        crossattn_cache = caches.crossattn_cache
+
+        if (self._should_snapshot_streaming_cache() and torch.is_grad_enabled()):
+            kv_cache = self._snapshot_kv_cache_indices(kv_cache)
+
+        model_kwargs: dict[str, Any] = {
+            "kv_cache": kv_cache,
+            "crossattn_cache": crossattn_cache,
+            "current_start": (cur_start_frame * frame_seq_length),
+            "start_frame": cur_start_frame,
+            "is_cache": bool(store_kv),
+        }
+
+        device_type = self.device.type
+        dtype = noisy_latents.dtype
+
+        if conditional:
+            text_dict = batch.conditional_dict
+            if text_dict is None:
+                raise RuntimeError("Missing conditional_dict in "
+                                   "TrainingBatch")
+        else:
+            text_dict = self._get_uncond_text_dict(batch, cfg_uncond=cfg_uncond)
+
+        with (
+                torch.autocast(device_type, dtype=dtype),
+                set_forward_context(
+                    current_timestep=batch.timesteps,
+                    attn_metadata=attn_metadata,
+                ),
+        ):
+            input_kwargs = (self._build_distill_input_kwargs(
+                noisy_latents,
+                timestep_full,
+                text_dict,
+            ))
+            input_kwargs["timestep"] = (timestep_full.to(
+                device=self.device,
+                dtype=torch.long,
+            ))
+            input_kwargs.update(model_kwargs)
+
+            if store_kv:
+                with torch.no_grad():
+                    _ = transformer(**input_kwargs)
+                return None
+
+            pred_noise = transformer(**input_kwargs, ).permute(0, 2, 1, 3, 4)
+        return pred_noise
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _ensure_per_frame_timestep(
+        self,
+        *,
+        timestep: torch.Tensor,
+        batch_size: int,
+        num_frames: int,
+        device: torch.device,
+    ) -> torch.Tensor:
+        if timestep.ndim == 0:
+            return (timestep.view(1, 1).expand(batch_size, num_frames).to(device=device))
+        if timestep.ndim == 1:
+            if int(timestep.shape[0]) == batch_size:
+                return (timestep.view(batch_size, 1).expand(batch_size, num_frames).to(device=device))
+            raise ValueError("streaming timestep must be scalar, "
+                             "[B], or [B, T]; got shape="
+                             f"{tuple(timestep.shape)}")
+        if timestep.ndim == 2:
+            return timestep.to(device=device)
+        raise ValueError("streaming timestep must be scalar, "
+                         "[B], or [B, T]; got ndim="
+                         f"{int(timestep.ndim)}")
+
+    def _get_or_init_streaming_caches(
+        self,
+        *,
+        cache_tag: str,
+        transformer: torch.nn.Module,
+        noisy_latents: torch.Tensor,
+    ) -> _StreamingCaches:
+        key = (id(self), cache_tag)
+        cached = self._streaming_caches.get(key)
+
+        batch_size = int(noisy_latents.shape[0])
+        dtype = noisy_latents.dtype
+        device = noisy_latents.device
+
+        frame_seq_length = (self._compute_frame_seq_length(transformer, noisy_latents))
+        local_attn_size = self._get_local_attn_size(transformer)
+        sliding_window_num_frames = (self._get_sliding_window_num_frames(transformer))
+
+        meta = (
+            frame_seq_length,
+            local_attn_size,
+            sliding_window_num_frames,
+            batch_size,
+            dtype,
+            device,
+        )
+
+        if cached is not None:
+            cached_meta = (
+                cached.frame_seq_length,
+                cached.local_attn_size,
+                cached.sliding_window_num_frames,
+                cached.batch_size,
+                cached.dtype,
+                cached.device,
+            )
+            if cached_meta == meta:
+                return cached
+
+        kv_cache = self._initialize_kv_cache(
+            transformer=transformer,
+            batch_size=batch_size,
+            dtype=dtype,
+            device=device,
+            frame_seq_length=frame_seq_length,
+            local_attn_size=local_attn_size,
+            sliding_window_num_frames=(sliding_window_num_frames),
+            checkpoint_safe=(self._should_use_checkpoint_safe_kv_cache()),
+        )
+        crossattn_cache = (self._initialize_crossattn_cache(
+            transformer=transformer,
+            device=device,
+        ))
+
+        caches = _StreamingCaches(
+            kv_cache=kv_cache,
+            crossattn_cache=crossattn_cache,
+            frame_seq_length=frame_seq_length,
+            local_attn_size=local_attn_size,
+            sliding_window_num_frames=(sliding_window_num_frames),
+            batch_size=batch_size,
+            dtype=dtype,
+            device=device,
+        )
+        self._streaming_caches[key] = caches
+        return caches
+
+    def _compute_frame_seq_length(
+        self,
+        transformer: torch.nn.Module,
+        noisy_latents: torch.Tensor,
+    ) -> int:
+        latent_seq_length = (int(noisy_latents.shape[-1]) * int(noisy_latents.shape[-2]))
+        patch_size = getattr(transformer, "patch_size", None)
+        if patch_size is None:
+            patch_size = getattr(
+                getattr(
+                    getattr(transformer, "config", None),
+                    "arch_config",
+                    None,
+                ),
+                "patch_size",
+                None,
+            )
+        if patch_size is None:
+            raise ValueError("Unable to determine "
+                             "transformer.patch_size "
+                             "for causal streaming")
+        patch_ratio = (int(patch_size[-1]) * int(patch_size[-2]))
+        if patch_ratio <= 0:
+            raise ValueError("Invalid patch_size for causal "
+                             "streaming")
+        return latent_seq_length // patch_ratio
+
+    def _get_sliding_window_num_frames(
+        self,
+        transformer: torch.nn.Module,
+    ) -> int:
+        cfg = getattr(transformer, "config", None)
+        arch_cfg = getattr(cfg, "arch_config", None)
+        value = (getattr(
+            arch_cfg,
+            "sliding_window_num_frames",
+            None,
+        ) if arch_cfg is not None else None)
+        if value is None:
+            return 15
+        return int(value)
+
+    def _get_local_attn_size(
+        self,
+        transformer: torch.nn.Module,
+    ) -> int:
+        try:
+            value = getattr(transformer, "local_attn_size", -1)
+        except Exception:
+            value = -1
+        if value is None:
+            return -1
+        return int(value)
+
+    def _initialize_kv_cache(
+        self,
+        *,
+        transformer: torch.nn.Module,
+        batch_size: int,
+        dtype: torch.dtype,
+        device: torch.device,
+        frame_seq_length: int,
+        local_attn_size: int,
+        sliding_window_num_frames: int,
+        checkpoint_safe: bool,
+    ) -> list[dict[str, Any]]:
+        num_blocks = len(getattr(transformer, "blocks", []))
+        if num_blocks <= 0:
+            raise ValueError("Unexpected transformer.blocks "
+                             "for causal streaming")
+
+        try:
+            num_attention_heads = int(transformer.num_attention_heads)  # type: ignore[attr-defined]
+        except AttributeError as e:
+            raise ValueError("Transformer is missing "
+                             "num_attention_heads") from e
+
+        try:
+            attention_head_dim = int(transformer.attention_head_dim)  # type: ignore[attr-defined]
+        except AttributeError:
+            try:
+                hidden_size = int(transformer.hidden_size)  # type: ignore[attr-defined]
+            except AttributeError as e:
+                raise ValueError("Transformer is missing "
+                                 "attention_head_dim and "
+                                 "hidden_size") from e
+            attention_head_dim = (hidden_size // max(1, num_attention_heads))
+
+        if local_attn_size != -1:
+            kv_cache_size = (int(local_attn_size) * int(frame_seq_length))
+        else:
+            kv_cache_size = (int(frame_seq_length) * int(sliding_window_num_frames))
+
+        if checkpoint_safe:
+            tc = getattr(self, "training_config", None)
+            total_frames = int(tc.data.num_frames if tc is not None else 0)
+            if total_frames <= 0:
+                raise ValueError("training.num_frames must be set "
+                                 "to enable checkpoint-safe "
+                                 "streaming KV cache; got "
+                                 f"{total_frames}")
+            kv_cache_size = max(
+                kv_cache_size,
+                int(frame_seq_length) * total_frames,
+            )
+
+        kv_cache: list[dict[str, Any]] = []
+        for _ in range(num_blocks):
+            kv_cache.append({
+                "k":
+                torch.zeros(
+                    [
+                        batch_size,
+                        kv_cache_size,
+                        num_attention_heads,
+                        attention_head_dim,
+                    ],
+                    dtype=dtype,
+                    device=device,
+                ),
+                "v":
+                torch.zeros(
+                    [
+                        batch_size,
+                        kv_cache_size,
+                        num_attention_heads,
+                        attention_head_dim,
+                    ],
+                    dtype=dtype,
+                    device=device,
+                ),
+                "global_end_index":
+                torch.zeros((), dtype=torch.long, device=device),
+                "local_end_index":
+                torch.zeros((), dtype=torch.long, device=device),
+            })
+
+        return kv_cache
+
+    def _should_use_checkpoint_safe_kv_cache(self, ) -> bool:
+        tc = getattr(self, "training_config", None)
+        if tc is not None:
+            checkpointing_type = (tc.model.enable_gradient_checkpointing_type)
+        else:
+            checkpointing_type = None
+        return (bool(checkpointing_type) and bool(self._trainable))
+
+    def _should_snapshot_streaming_cache(self, ) -> bool:
+        return (self._should_use_checkpoint_safe_kv_cache())
+
+    def _snapshot_kv_cache_indices(
+        self,
+        kv_cache: list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        snapshot: list[dict[str, Any]] = []
+        for block_cache in kv_cache:
+            global_end_index = block_cache.get("global_end_index")
+            local_end_index = block_cache.get("local_end_index")
+            if (not isinstance(global_end_index, torch.Tensor) or not isinstance(local_end_index, torch.Tensor)):
+                raise ValueError("Unexpected kv_cache index "
+                                 "tensors; expected tensors at "
+                                 "kv_cache[*].{global_end_index, "
+                                 "local_end_index}")
+
+            copied = dict(block_cache)
+            copied["global_end_index"] = (global_end_index.detach().clone())
+            copied["local_end_index"] = (local_end_index.detach().clone())
+            snapshot.append(copied)
+        return snapshot
+
+    def _initialize_crossattn_cache(
+        self,
+        *,
+        transformer: torch.nn.Module,
+        device: torch.device,
+    ) -> list[dict[str, Any]] | None:
+        num_blocks = len(getattr(transformer, "blocks", []))
+        if num_blocks <= 0:
+            return None
+        return [{
+            "is_init": False,
+            "k": torch.empty(0, device=device),
+            "v": torch.empty(0, device=device),
+        } for _ in range(num_blocks)]

fastvideo/train/trainer.py

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import time
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import Any, TYPE_CHECKING
+
+import torch
+from tqdm.auto import tqdm
+
+from fastvideo.distributed import get_sp_group, get_world_group
+from fastvideo.train.callbacks.callback import CallbackDict
+from fastvideo.train.methods.base import TrainingMethod
+from fastvideo.train.utils.tracking import build_tracker
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+
+
+def _coerce_log_scalar(value: Any, *, where: str) -> float:
+    if isinstance(value, torch.Tensor):
+        if value.numel() != 1:
+            raise ValueError(f"Expected scalar tensor at {where}, "
+                             f"got shape={tuple(value.shape)}")
+        return float(value.detach().item())
+    if isinstance(value, float | int):
+        return float(value)
+    raise TypeError(f"Expected a scalar (float/int/Tensor) at "
+                    f"{where}, got {type(value).__name__}")
+
+
+@dataclass(slots=True)
+class TrainLoopState:
+    step: int
+    accum_iter: int
+    current_vsa_sparsity: float
+
+
+class Trainer:
+
+    def __init__(
+        self,
+        training_config: TrainingConfig,
+        *,
+        config: dict[str, Any] | None = None,
+        callback_configs: dict[str, dict[str, Any]]
+        | None = None,
+    ) -> None:
+        self.training_config = training_config
+        self.world_group = get_world_group()
+        self.sp_group = get_sp_group()
+        self.global_rank = self.world_group.rank
+        self.local_rank = self.world_group.local_rank
+        self.tracker = build_tracker(
+            training_config.tracker,
+            training_config.checkpoint,
+            config=config,
+        )
+        self.callbacks = CallbackDict(
+            callback_configs or {},
+            training_config,
+        )
+
+    def _iter_dataloader(self, dataloader: Any) -> Iterator[dict[str, Any]]:
+        data_iter = iter(dataloader)
+        while True:
+            batch = next(data_iter, None)
+            if batch is None:
+                data_iter = iter(dataloader)
+                batch = next(data_iter)
+            yield batch
+
+    def _get_current_vsa_sparsity(self, step: int) -> float:
+        tc = self.training_config
+        vsa_sparsity = tc.vsa.sparsity
+        vsa_decay_rate = tc.vsa.decay_rate
+        vsa_decay_interval_steps = (tc.vsa.decay_interval_steps)
+        if vsa_decay_interval_steps > 1:
+            current_decay_times = min(
+                step // vsa_decay_interval_steps,
+                int(vsa_sparsity // vsa_decay_rate),
+            )
+            return current_decay_times * vsa_decay_rate
+        return vsa_sparsity
+
+    def run(
+        self,
+        method: TrainingMethod,
+        *,
+        dataloader: Any,
+        max_steps: int,
+        start_step: int = 0,
+        checkpoint_manager: Any | None = None,
+    ) -> None:
+        tc = self.training_config
+        grad_accum = max(
+            1,
+            int(tc.loop.gradient_accumulation_steps or 1),
+        )
+
+        method.set_tracker(self.tracker)
+        method.on_train_start()
+        self.callbacks.on_train_start(
+            method,
+            iteration=start_step,
+        )
+
+        resume_from_checkpoint = (tc.checkpoint.resume_from_checkpoint or "")
+        if checkpoint_manager is not None:
+            if resume_from_checkpoint:
+                method.seed_optimizer_state_for_resume()
+            resumed_step = (checkpoint_manager.maybe_resume(resume_from_checkpoint=(resume_from_checkpoint)))
+            if resumed_step is not None:
+                start_step = int(resumed_step)
+        self.callbacks.on_validation_begin(
+            method,
+            iteration=start_step,
+        )
+        method.optimizers_zero_grad(start_step)
+
+        data_stream = self._iter_dataloader(dataloader)
+
+        # Restore the RNG snapshot LAST — after dcp.load,
+        # after iter(dataloader), after everything that may
+        # have advanced the RNG as a side-effect.
+        if (checkpoint_manager is not None and resume_from_checkpoint):
+            checkpoint_manager.load_rng_snapshot(resume_from_checkpoint, )
+        progress = tqdm(
+            range(start_step + 1, max_steps + 1),
+            initial=start_step,
+            desc="Steps",
+            disable=self.local_rank > 0,
+        )
+        for step in progress:
+            t0 = time.perf_counter()
+            current_vsa_sparsity = (self._get_current_vsa_sparsity(step))
+
+            loss_sums: dict[str, float] = {}
+            metric_sums: dict[str, float] = {}
+            for accum_iter in range(grad_accum):
+                batch = next(data_stream)
+                loss_map, outputs, step_metrics = method.single_train_step(
+                    batch,
+                    step,
+                    current_vsa_sparsity=(current_vsa_sparsity),
+                )
+
+                method.backward(
+                    loss_map,
+                    outputs,
+                    grad_accum_rounds=grad_accum,
+                )
+
+                for k, v in loss_map.items():
+                    if isinstance(v, torch.Tensor):
+                        loss_sums[k] = loss_sums.get(k, 0.0) + float(v.detach().item())
+                for k, v in step_metrics.items():
+                    if k in loss_sums:
+                        raise ValueError(f"Metric key {k!r} collides "
+                                         "with loss key. Use a "
+                                         "different name (e.g. prefix "
+                                         "with 'train/').")
+                    metric_sums[k] = metric_sums.get(k, 0.0) + _coerce_log_scalar(
+                        v,
+                        where=("method.single_train_step()"
+                               f".metrics[{k!r}]"),
+                    )
+
+            self.callbacks.on_before_optimizer_step(
+                method,
+                iteration=step,
+            )
+            method.optimizers_schedulers_step(step)
+            method.optimizers_zero_grad(step)
+
+            metrics = {k: v / grad_accum for k, v in loss_sums.items()}
+            metrics.update({k: v / grad_accum for k, v in metric_sums.items()})
+            metrics["step_time_sec"] = (time.perf_counter() - t0)
+            metrics["vsa_sparsity"] = float(current_vsa_sparsity)
+            if self.global_rank == 0 and metrics:
+                self.tracker.log(metrics, step)
+
+            self.callbacks.on_training_step_end(
+                method,
+                metrics,
+                iteration=step,
+            )
+
+            if checkpoint_manager is not None:
+                checkpoint_manager.maybe_save(step)
+
+            self.callbacks.on_validation_begin(
+                method,
+                iteration=step,
+            )
+            self.callbacks.on_validation_end(
+                method,
+                iteration=step,
+            )
+
+        self.callbacks.on_train_end(
+            method,
+            iteration=max_steps,
+        )
+
+        if checkpoint_manager is not None:
+            checkpoint_manager.save_final(max_steps)
+
+        self.tracker.finish()

fastvideo/train/utils/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Distillation utilities shared across families/methods/entrypoints."""

fastvideo/train/utils/builder.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Assembly: build method + dataloader from a ``_target_``-based config."""
+
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from fastvideo.train.utils.instantiate import (
+    instantiate,
+    resolve_target,
+)
+from fastvideo.train.utils.config import RunConfig
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+    from fastvideo.train.methods.base import TrainingMethod
+
+
+def build_from_config(cfg: RunConfig, ) -> tuple[TrainingConfig, TrainingMethod, Any, int]:
+    """Build method + dataloader from a v3 run config.
+
+    1. Instantiate each model in ``cfg.models`` via ``_target_``.
+    2. Resolve the method class from ``cfg.method["_target_"]``
+       and construct it with ``(cfg=cfg, role_models=...)``.
+    3. Return ``(training_args, method, dataloader, start_step)``.
+    """
+    from fastvideo.train.models.base import ModelBase
+
+    # --- 1. Build role model instances ---
+    role_models: dict[str, ModelBase] = {}
+    for role, model_cfg in cfg.models.items():
+        model = instantiate(model_cfg, training_config=cfg.training)
+        if not isinstance(model, ModelBase):
+            raise TypeError(f"models.{role}._target_ must resolve to a "
+                            f"ModelBase subclass, got {type(model).__name__}")
+        role_models[role] = model
+
+    # --- 2. Build method ---
+    method_cfg = dict(cfg.method)
+    method_target = str(method_cfg.pop("_target_"))
+    method_cls = resolve_target(method_target)
+
+    # The student model provides the dataloader.
+    student = role_models.get("student")
+
+    method = method_cls(
+        cfg=cfg,
+        role_models=role_models,
+    )
+
+    # --- 3. Gather dataloader and start_step ---
+    dataloader = (getattr(student, "dataloader", None) if student is not None else None)
+    start_step = int(getattr(student, "start_step", 0) if student is not None else 0)
+
+    return cfg.training, method, dataloader, start_step

fastvideo/train/utils/checkpoint.py

@@ -0,0 +1,388 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import json
+import os
+import random
+import re
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+_CHECKPOINT_DIR_RE = re.compile(r"^checkpoint-(\d+)$")
+
+
+def _is_stateful(obj: Any) -> bool:
+    return callable(getattr(obj, "state_dict", None)) and callable(getattr(obj, "load_state_dict", None))
+
+
+def _rank() -> int:
+    if dist.is_available() and dist.is_initialized():
+        return int(dist.get_rank())
+    return 0
+
+
+def _barrier() -> None:
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+
+
+def _parse_step_from_dir(checkpoint_dir: Path) -> int:
+    match = _CHECKPOINT_DIR_RE.match(checkpoint_dir.name)
+    if not match:
+        raise ValueError(f"Invalid checkpoint directory name {checkpoint_dir.name!r}; "
+                         "expected 'checkpoint-<step>'")
+    return int(match.group(1))
+
+
+def _find_latest_checkpoint(output_dir: Path) -> Path | None:
+    if not output_dir.exists():
+        return None
+
+    candidates: list[tuple[int, Path]] = []
+    for child in output_dir.iterdir():
+        if not child.is_dir():
+            continue
+        if not _CHECKPOINT_DIR_RE.match(child.name):
+            continue
+        if not (child / "dcp").is_dir():
+            continue
+        try:
+            step = _parse_step_from_dir(child)
+        except Exception:
+            continue
+        candidates.append((step, child))
+
+    if not candidates:
+        return None
+    candidates.sort(key=lambda x: x[0])
+    return candidates[-1][1]
+
+
+def _resolve_resume_checkpoint(resume_from_checkpoint: str, *, output_dir: str) -> Path | None:
+    """Resolve a user-provided resume path to a concrete checkpoint dir.
+
+    Accepted values:
+    - "latest" (auto-pick latest checkpoint-*/dcp under output_dir,
+      or ``None`` if no checkpoint exists yet — starts from scratch)
+    - /path/to/output_dir/checkpoint-<step>
+    - /path/to/output_dir/checkpoint-<step>/dcp
+    - /path/to/output_dir (auto-pick latest checkpoint-*/dcp)
+    """
+
+    if str(resume_from_checkpoint).strip().lower() == "latest":
+        out = Path(os.path.expanduser(str(output_dir))).resolve()
+        latest = _find_latest_checkpoint(out)
+        if latest is None:
+            logger.info(
+                "resume_from_checkpoint='latest' but no "
+                "checkpoints found under %s; starting from "
+                "scratch.",
+                out,
+            )
+        return latest
+
+    raw = os.path.expanduser(str(resume_from_checkpoint))
+    path = Path(raw).resolve()
+    if not path.exists():
+        raise FileNotFoundError(f"resume_from_checkpoint not found: {path}")
+
+    if path.is_dir() and path.name == "dcp":
+        path = path.parent
+
+    if path.is_dir() and _CHECKPOINT_DIR_RE.match(path.name):
+        if not (path / "dcp").is_dir():
+            raise FileNotFoundError(f"Missing dcp dir under checkpoint: {path / 'dcp'}")
+        return path
+
+    # Treat as output_dir -> pick latest.
+    latest = _find_latest_checkpoint(path)
+    if latest is not None:
+        return latest
+
+    # Give a clearer error message.
+    out = Path(os.path.expanduser(str(output_dir))).resolve()
+    raise ValueError("Could not resolve resume checkpoint. Expected a checkpoint directory "
+                     f"named 'checkpoint-<step>' (with 'dcp/' inside), or an output_dir "
+                     f"containing such checkpoints. Got: {path} (output_dir={out}).")
+
+
+class _RoleModuleContainer(torch.nn.Module):
+    """Ephemeral container to expose multiple role modules as a single
+    ``nn.Module``.
+
+    Used by ``OptimizerWrapper`` which expects a single root module
+    covering all parameters owned by the optimizer.
+    """
+
+    def __init__(self, modules: dict[str, torch.nn.Module]) -> None:
+        super().__init__()
+        for name, module in modules.items():
+            self.add_module(name, module)
+
+
+class _CallbackStateWrapper:
+    """Wraps a CallbackDict for DCP save/load."""
+
+    def __init__(self, callbacks: Any) -> None:
+        self._callbacks = callbacks
+
+    def state_dict(self) -> dict[str, Any]:
+        return self._callbacks.state_dict()
+
+    def load_state_dict(
+        self,
+        state_dict: dict[str, Any],
+    ) -> None:
+        self._callbacks.load_state_dict(state_dict)
+
+
+@dataclass(slots=True)
+class CheckpointConfig:
+    save_steps: int
+    keep_last: int
+
+
+class CheckpointManager:
+    """Role-based checkpoint manager for training runtime.
+
+    - Checkpoint policy lives in YAML (via TrainingArgs fields).
+    - Resume path is typically provided via CLI (``--resume-from-checkpoint``).
+    """
+
+    def __init__(
+        self,
+        *,
+        method: Any,
+        dataloader: Any,
+        output_dir: str,
+        config: CheckpointConfig,
+        callbacks: Any | None = None,
+        raw_config: dict[str, Any] | None = None,
+    ) -> None:
+        self.method = method
+        self.dataloader = dataloader
+        self.output_dir = str(output_dir)
+        self.config = config
+        self._callbacks = callbacks
+        self._raw_config = raw_config
+        self._last_saved_step: int | None = None
+
+    def _build_states(self) -> dict[str, Any]:
+        states: dict[str, Any] = self.method.checkpoint_state()
+
+        # Dataloader (optional but recommended for exact resume).
+        if _is_stateful(self.dataloader):
+            states["dataloader"] = self.dataloader
+
+        # Callback state (e.g. EMA shadow weights, validation RNG).
+        if self._callbacks is not None and _is_stateful(self._callbacks):
+            states["callbacks"] = _CallbackStateWrapper(self._callbacks, )
+
+        return states
+
+    def _checkpoint_dir(self, step: int) -> Path:
+        return Path(self.output_dir) / f"checkpoint-{step}"
+
+    def _dcp_dir(self, step: int) -> Path:
+        return self._checkpoint_dir(step) / "dcp"
+
+    def maybe_save(self, step: int) -> None:
+        save_steps = int(self.config.save_steps or 0)
+        if save_steps <= 0:
+            return
+        if step % save_steps != 0:
+            return
+        if self._last_saved_step == step:
+            return
+        self.save(step)
+
+    def save_final(self, step: int) -> None:
+        save_steps = int(self.config.save_steps or 0)
+        if save_steps <= 0:
+            return
+        self.save(step)
+
+    def save(self, step: int) -> None:
+        checkpoint_dir = self._checkpoint_dir(step)
+        dcp_dir = self._dcp_dir(step)
+        os.makedirs(dcp_dir, exist_ok=True)
+
+        states = self._build_states()
+        if _rank() == 0:
+            logger.info(
+                "Saving checkpoint to %s",
+                checkpoint_dir,
+            )
+            self._write_metadata(checkpoint_dir, step)
+        dcp.save(states, checkpoint_id=str(dcp_dir))
+        _barrier()
+
+        # Save RNG state AFTER dcp.save so it captures the
+        # exact state the continuous run continues with.
+        # dcp.save triggers FSDP all-gather ops that can
+        # advance the RNG between when DCP captures it and
+        # when the save completes.
+        self._save_rng_snapshot(checkpoint_dir)
+        _barrier()
+
+        self._last_saved_step = step
+
+        self._cleanup_old_checkpoints()
+
+    def _write_metadata(
+        self,
+        checkpoint_dir: Path,
+        step: int,
+    ) -> None:
+        metadata: dict[str, Any] = {"step": step}
+        if self._raw_config is not None:
+            metadata["config"] = self._raw_config
+        meta_path = checkpoint_dir / "metadata.json"
+        with open(meta_path, "w", encoding="utf-8") as f:
+            json.dump(metadata, f, indent=2)
+
+    @staticmethod
+    def load_metadata(checkpoint_dir: str | Path, ) -> dict[str, Any]:
+        """Read ``metadata.json`` from a checkpoint dir."""
+        meta_path = Path(checkpoint_dir) / "metadata.json"
+        if not meta_path.is_file():
+            raise FileNotFoundError(f"No metadata.json in {checkpoint_dir}")
+        with open(meta_path, encoding="utf-8") as f:
+            return json.load(f)  # type: ignore[no-any-return]
+
+    def _save_rng_snapshot(self, checkpoint_dir: Path) -> None:
+        """Save per-rank RNG state to a separate file.
+
+        Called AFTER ``dcp.save`` so the snapshot reflects
+        the exact state the continuous run continues with.
+        Each rank saves its own file because CUDA RNG and
+        custom generators differ across ranks.
+        """
+        rank = _rank()
+        rng: dict[str, Any] = {
+            "torch_rng": torch.get_rng_state(),
+            "python_rng": random.getstate(),
+            "numpy_rng": np.random.get_state(),
+        }
+        if torch.cuda.is_available():
+            rng["cuda_rng"] = torch.cuda.get_rng_state()
+        generators = (self.method.get_rng_generators() if hasattr(self.method, "get_rng_generators") else {})
+        for name, gen in (generators or {}).items():
+            if gen is not None:
+                rng[f"gen_{name}"] = gen.get_state()
+        torch.save(
+            rng,
+            checkpoint_dir / f"rng_state_rank{rank}.pt",
+        )
+
+    def load_rng_snapshot(
+        self,
+        checkpoint_path: str,
+    ) -> None:
+        """Restore per-rank RNG state from the snapshot file.
+
+        Must be called AFTER ``dcp.load`` **and** after
+        ``iter(dataloader)`` so no later operation can
+        clobber the restored state.
+        """
+        resolved = _resolve_resume_checkpoint(
+            checkpoint_path,
+            output_dir=self.output_dir,
+        )
+        if resolved is None:
+            return
+        rank = _rank()
+        rng_path = resolved / f"rng_state_rank{rank}.pt"
+        if not rng_path.is_file():
+            # Fall back to legacy single-file snapshot.
+            rng_path = resolved / "rng_state.pt"
+        if not rng_path.is_file():
+            logger.warning(
+                "No rng_state in %s; skipping "
+                "RNG snapshot restore.",
+                resolved,
+            )
+            return
+
+        rng = torch.load(
+            rng_path,
+            map_location="cpu",
+            weights_only=False,
+        )
+        if "torch_rng" in rng:
+            torch.set_rng_state(rng["torch_rng"])
+        if "python_rng" in rng:
+            random.setstate(rng["python_rng"])
+        if "numpy_rng" in rng:
+            np.random.set_state(rng["numpy_rng"])
+        if torch.cuda.is_available() and "cuda_rng" in rng:
+            torch.cuda.set_rng_state(rng["cuda_rng"])
+        generators = (self.method.get_rng_generators() if hasattr(self.method, "get_rng_generators") else {})
+        for name, gen in (generators or {}).items():
+            key = f"gen_{name}"
+            if key in rng and gen is not None:
+                gen.set_state(rng[key])
+        logger.info(
+            "Restored RNG snapshot from %s",
+            rng_path,
+        )
+
+    def maybe_resume(self, *, resume_from_checkpoint: str | None) -> int | None:
+        if not resume_from_checkpoint:
+            return None
+
+        resolved = _resolve_resume_checkpoint(
+            resume_from_checkpoint,
+            output_dir=self.output_dir,
+        )
+        if resolved is None:
+            return None
+        step = _parse_step_from_dir(resolved)
+
+        states = self._build_states()
+        logger.info("Loading Phase 2 checkpoint from %s", resolved)
+        dcp.load(states, checkpoint_id=str(resolved / "dcp"))
+        _barrier()
+        logger.info("Checkpoint loaded; resuming from step=%s", step)
+        return step
+
+    def _cleanup_old_checkpoints(self) -> None:
+        keep_last = int(self.config.keep_last or 0)
+        if keep_last <= 0:
+            return
+
+        if _rank() != 0:
+            _barrier()
+            return
+
+        output_dir = Path(self.output_dir)
+        candidates: list[tuple[int, Path]] = []
+        for child in output_dir.iterdir():
+            if not child.is_dir():
+                continue
+            if not _CHECKPOINT_DIR_RE.match(child.name):
+                continue
+            try:
+                step = _parse_step_from_dir(child)
+            except Exception:
+                continue
+            candidates.append((step, child))
+
+        candidates.sort(key=lambda x: x[0])
+        to_delete = candidates[:-keep_last] if len(candidates) > keep_last else []
+        for step, path in to_delete:
+            logger.info("Removing old checkpoint (keep_last=%s): %s", keep_last, path)
+            shutil.rmtree(path, ignore_errors=True)
+
+        _barrier()

fastvideo/train/utils/config.py

@@ -0,0 +1,511 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Training run config (``_target_`` based YAML)."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from fastvideo.logger import init_logger
+from fastvideo.train.utils.training_config import (
+    CheckpointConfig,
+    DataConfig,
+    TrainingConfig,
+    DistributedConfig,
+    ModelTrainingConfig,
+    OptimizerConfig,
+    TrackerConfig,
+    TrainingLoopConfig,
+    VSAConfig,
+)
+
+logger = init_logger(__name__)
+
+
+@dataclass(slots=True)
+class RunConfig:
+    """Parsed run config loaded from YAML."""
+
+    models: dict[str, dict[str, Any]]
+    method: dict[str, Any]
+    training: TrainingConfig
+    callbacks: dict[str, dict[str, Any]]
+    raw: dict[str, Any]
+
+    def resolved_config(self) -> dict[str, Any]:
+        """Return a fully-resolved config dict with defaults.
+
+        Suitable for logging to W&B so that every parameter
+        (including defaults) is visible.
+        """
+        import dataclasses
+
+        def _safe_asdict(obj: Any) -> Any:
+            if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
+                return {
+                    f.name: _safe_asdict(getattr(obj, f.name))
+                    for f in dataclasses.fields(obj) if not callable(getattr(obj, f.name))
+                }
+            if isinstance(obj, dict):
+                return {k: _safe_asdict(v) for k, v in obj.items()}
+            if isinstance(obj, list | tuple):
+                return type(obj)(_safe_asdict(v) for v in obj)
+            return obj
+
+        resolved: dict[str, Any] = {}
+        resolved["models"] = dict(self.models)
+        resolved["method"] = dict(self.method)
+        resolved["training"] = _safe_asdict(self.training)
+        resolved["callbacks"] = dict(self.callbacks)
+        return resolved
+
+
+# ---- parsing helpers (kept for use by methods) ----
+
+
+def _resolve_existing_file(path: str) -> str:
+    if not path:
+        return path
+    expanded = os.path.expanduser(path)
+    resolved = Path(expanded).resolve()
+    if not resolved.exists():
+        raise FileNotFoundError(f"Config file not found: {resolved}")
+    if not resolved.is_file():
+        raise ValueError(f"Expected a file path, got: {resolved}")
+    return str(resolved)
+
+
+def _require_mapping(raw: Any, *, where: str) -> dict[str, Any]:
+    if not isinstance(raw, dict):
+        raise ValueError(f"Expected mapping at {where}, "
+                         f"got {type(raw).__name__}")
+    return raw
+
+
+def _require_str(raw: Any, *, where: str) -> str:
+    if not isinstance(raw, str) or not raw.strip():
+        raise ValueError(f"Expected non-empty string at {where}")
+    return raw
+
+
+def get_optional_int(mapping: dict[str, Any], key: str, *, where: str) -> int | None:
+    raw = mapping.get(key)
+    if raw is None:
+        return None
+    if isinstance(raw, bool):
+        raise ValueError(f"Expected int at {where}, got bool")
+    if isinstance(raw, int):
+        return int(raw)
+    if isinstance(raw, float) and raw.is_integer():
+        return int(raw)
+    if isinstance(raw, str) and raw.strip():
+        return int(raw)
+    raise ValueError(f"Expected int at {where}, "
+                     f"got {type(raw).__name__}")
+
+
+def get_optional_float(mapping: dict[str, Any], key: str, *, where: str) -> float | None:
+    raw = mapping.get(key)
+    if raw is None:
+        return None
+    if isinstance(raw, bool):
+        raise ValueError(f"Expected float at {where}, got bool")
+    if isinstance(raw, int | float):
+        return float(raw)
+    if isinstance(raw, str) and raw.strip():
+        return float(raw)
+    raise ValueError(f"Expected float at {where}, "
+                     f"got {type(raw).__name__}")
+
+
+def parse_betas(raw: Any, *, where: str) -> tuple[float, float]:
+    if raw is None:
+        raise ValueError(f"Missing betas for {where}")
+    if isinstance(raw, tuple | list) and len(raw) == 2:
+        return float(raw[0]), float(raw[1])
+    if isinstance(raw, str):
+        parts = [p.strip() for p in raw.split(",") if p.strip()]
+        if len(parts) != 2:
+            raise ValueError(f"Expected betas as 'b1,b2' at {where}, "
+                             f"got {raw!r}")
+        return float(parts[0]), float(parts[1])
+    raise ValueError(f"Expected betas as 'b1,b2' at {where}, "
+                     f"got {type(raw).__name__}")
+
+
+# ---- config convenience helpers ----
+
+
+def require_positive_int(
+    mapping: dict[str, Any],
+    key: str,
+    *,
+    default: int | None = None,
+    where: str | None = None,
+) -> int:
+    """Read an int that must be > 0."""
+    loc = where or key
+    raw = mapping.get(key)
+    if raw is None:
+        if default is not None:
+            return default
+        raise ValueError(f"Missing required key {loc!r}")
+    val = get_optional_int(mapping, key, where=loc)
+    if val is None or val <= 0:
+        raise ValueError(f"{loc} must be a positive integer, got {raw!r}")
+    return val
+
+
+def require_non_negative_int(
+    mapping: dict[str, Any],
+    key: str,
+    *,
+    default: int | None = None,
+    where: str | None = None,
+) -> int:
+    """Read an int that must be >= 0."""
+    loc = where or key
+    raw = mapping.get(key)
+    if raw is None:
+        if default is not None:
+            return default
+        raise ValueError(f"Missing required key {loc!r}")
+    val = get_optional_int(mapping, key, where=loc)
+    if val is None or val < 0:
+        raise ValueError(f"{loc} must be a non-negative integer, "
+                         f"got {raw!r}")
+    return val
+
+
+def require_non_negative_float(
+    mapping: dict[str, Any],
+    key: str,
+    *,
+    default: float | None = None,
+    where: str | None = None,
+) -> float:
+    """Read a float that must be >= 0."""
+    loc = where or key
+    raw = mapping.get(key)
+    if raw is None:
+        if default is not None:
+            return default
+        raise ValueError(f"Missing required key {loc!r}")
+    val = get_optional_float(mapping, key, where=loc)
+    if val is None or val < 0.0:
+        raise ValueError(f"{loc} must be a non-negative float, "
+                         f"got {raw!r}")
+    return val
+
+
+def require_choice(
+    mapping: dict[str, Any],
+    key: str,
+    choices: set[str] | frozenset[str],
+    *,
+    default: str | None = None,
+    where: str | None = None,
+) -> str:
+    """Read a string that must be one of *choices*."""
+    loc = where or key
+    raw = mapping.get(key)
+    if raw is None:
+        if default is not None:
+            if default not in choices:
+                raise ValueError(f"Default {default!r} not in {choices}")
+            return default
+        raise ValueError(f"Missing required key {loc!r}")
+    if not isinstance(raw, str) or not raw.strip():
+        raise ValueError(f"{loc} must be a non-empty string, "
+                         f"got {type(raw).__name__}")
+    val = raw.strip().lower()
+    if val not in choices:
+        raise ValueError(f"{loc} must be one of {sorted(choices)}, "
+                         f"got {raw!r}")
+    return val
+
+
+def require_bool(
+    mapping: dict[str, Any],
+    key: str,
+    *,
+    default: bool | None = None,
+    where: str | None = None,
+) -> bool:
+    """Read a bool value."""
+    loc = where or key
+    raw = mapping.get(key)
+    if raw is None:
+        if default is not None:
+            return default
+        raise ValueError(f"Missing required key {loc!r}")
+    if not isinstance(raw, bool):
+        raise ValueError(f"{loc} must be a bool, "
+                         f"got {type(raw).__name__}")
+    return raw
+
+
+def _parse_pipeline_config(
+    cfg: dict[str, Any],
+    *,
+    models: dict[str, dict[str, Any]],
+) -> Any:
+    """Resolve PipelineConfig from the ``pipeline:`` YAML key."""
+    from fastvideo.configs.pipelines.base import PipelineConfig
+
+    pipeline_raw = cfg.get("pipeline")
+    if pipeline_raw is None:
+        return None
+
+    # Derive model_path from models.student.init_from —
+    # needed by PipelineConfig.from_kwargs.
+    model_path: str | None = None
+    student_cfg = models.get("student")
+    if student_cfg is not None:
+        init_from = student_cfg.get("init_from")
+        if init_from is not None:
+            model_path = str(init_from)
+
+    kwargs: dict[str, Any] = {"pipeline_config": pipeline_raw}
+    if model_path is not None:
+        kwargs["model_path"] = model_path
+
+    if isinstance(pipeline_raw, str):
+        kwargs["pipeline_config"] = _resolve_existing_file(pipeline_raw)
+
+    return PipelineConfig.from_kwargs(kwargs)
+
+
+def _build_training_config(
+    t: dict[str, Any],
+    *,
+    models: dict[str, dict[str, Any]],
+    pipeline_config: Any,
+) -> TrainingConfig:
+    """Build TrainingConfig from nested training: YAML."""
+    d = dict(t.get("distributed", {}) or {})
+    da = dict(t.get("data", {}) or {})
+    o = dict(t.get("optimizer", {}) or {})
+    lo = dict(t.get("loop", {}) or {})
+    ck = dict(t.get("checkpoint", {}) or {})
+    tr = dict(t.get("tracker", {}) or {})
+    vs = dict(t.get("vsa", {}) or {})
+    m = dict(t.get("model", {}) or {})
+
+    num_gpus = int(d.get("num_gpus", 1) or 1)
+
+    betas_raw = o.get("betas", "0.9,0.999")
+    betas = parse_betas(betas_raw, where="training.optimizer.betas")
+
+    model_path = str(t.get("model_path", "") or "")
+    if not model_path:
+        student_cfg = models.get("student")
+        if student_cfg is not None:
+            init_from = student_cfg.get("init_from")
+            if init_from is not None:
+                model_path = str(init_from)
+
+    return TrainingConfig(
+        distributed=DistributedConfig(
+            num_gpus=num_gpus,
+            tp_size=int(d.get("tp_size", 1) or 1),
+            sp_size=int(d.get("sp_size", num_gpus) or num_gpus),
+            hsdp_replicate_dim=int(d.get("hsdp_replicate_dim", 1) or 1),
+            hsdp_shard_dim=int(d.get("hsdp_shard_dim", num_gpus) or num_gpus),
+            pin_cpu_memory=bool(d.get("pin_cpu_memory", False)),
+        ),
+        data=DataConfig(
+            data_path=str(da.get("data_path", "") or ""),
+            train_batch_size=int(da.get("train_batch_size", 1) or 1),
+            dataloader_num_workers=int(da.get("dataloader_num_workers", 0) or 0),
+            training_cfg_rate=float(da.get("training_cfg_rate", 0.0) or 0.0),
+            seed=int(da.get("seed", 0) or 0),
+            num_height=int(da.get("num_height", 0) or 0),
+            num_width=int(da.get("num_width", 0) or 0),
+            num_latent_t=int(da.get("num_latent_t", 0) or 0),
+            num_frames=int(da.get("num_frames", 0) or 0),
+        ),
+        optimizer=OptimizerConfig(
+            learning_rate=float(o.get("learning_rate", 0.0) or 0.0),
+            betas=betas,
+            weight_decay=float(o.get("weight_decay", 0.0) or 0.0),
+            lr_scheduler=str(o.get("lr_scheduler", "constant") or "constant"),
+            lr_warmup_steps=int(o.get("lr_warmup_steps", 0) or 0),
+            lr_num_cycles=int(o.get("lr_num_cycles", 0) or 0),
+            lr_power=float(o.get("lr_power", 0.0) or 0.0),
+            min_lr_ratio=float(o.get("min_lr_ratio", 0.5) or 0.5),
+        ),
+        loop=TrainingLoopConfig(
+            max_train_steps=int(lo.get("max_train_steps", 0) or 0),
+            gradient_accumulation_steps=int(lo.get("gradient_accumulation_steps", 1) or 1),
+        ),
+        checkpoint=CheckpointConfig(
+            output_dir=str(ck.get("output_dir", "") or ""),
+            resume_from_checkpoint=str(ck.get("resume_from_checkpoint", "") or ""),
+            training_state_checkpointing_steps=int(ck.get("training_state_checkpointing_steps", 0) or 0),
+            checkpoints_total_limit=int(ck.get("checkpoints_total_limit", 0) or 0),
+        ),
+        tracker=TrackerConfig(
+            trackers=list(tr.get("trackers", []) or []),
+            project_name=str(tr.get("project_name", "fastvideo") or "fastvideo"),
+            run_name=str(tr.get("run_name", "") or ""),
+        ),
+        vsa=VSAConfig(
+            sparsity=float(vs.get("sparsity", 0.0) or 0.0),
+            decay_rate=float(vs.get("decay_rate", 0.0) or 0.0),
+            decay_interval_steps=int(vs.get("decay_interval_steps", 0) or 0),
+        ),
+        model=ModelTrainingConfig(
+            weighting_scheme=str(m.get("weighting_scheme", "uniform") or "uniform"),
+            logit_mean=float(m.get("logit_mean", 0.0) or 0.0),
+            logit_std=float(m.get("logit_std", 1.0) or 1.0),
+            mode_scale=float(m.get("mode_scale", 1.0) or 1.0),
+            precondition_outputs=bool(m.get("precondition_outputs", False)),
+            moba_config=dict(m.get("moba_config", {}) or {}),
+            enable_gradient_checkpointing_type=(m.get("enable_gradient_checkpointing_type")),
+        ),
+        pipeline_config=pipeline_config,
+        model_path=model_path,
+        dit_precision=str(t.get("dit_precision", "fp32") or "fp32"),
+    )
+
+
+def _parse_cli_overrides(overrides: list[str], ) -> dict[str, Any]:
+    """Parse ``--dotted.key value`` CLI overrides.
+
+    Returns a flat dict mapping dotted keys to parsed
+    values.  Values are auto-cast: ``true``/``false`` to
+    bool, integers/floats to numbers, and YAML list
+    literals (``[a, b]``) to lists.
+    """
+    result: dict[str, Any] = {}
+    i = 0
+    while i < len(overrides):
+        arg = overrides[i]
+        if not arg.startswith("--"):
+            raise ValueError(f"Expected --dotted.key, got {arg!r}")
+        key = arg[2:]  # strip leading --
+        if "=" in key:
+            key, raw_val = key.split("=", 1)
+        else:
+            i += 1
+            if i >= len(overrides):
+                raise ValueError(f"Missing value for override {arg!r}")
+            raw_val = overrides[i]
+        result[key] = _cast_value(raw_val)
+        i += 1
+    return result
+
+
+def _cast_value(raw: str) -> Any:
+    """Best-effort cast a CLI string to a Python value."""
+    if raw.lower() == "true":
+        return True
+    if raw.lower() == "false":
+        return False
+    if raw.lower() in ("none", "null"):
+        return None
+    # Try int
+    try:
+        return int(raw)
+    except ValueError:
+        pass
+    # Try float
+    try:
+        return float(raw)
+    except ValueError:
+        pass
+    # Try YAML list literal like [1, 2]
+    if raw.startswith("[") and raw.endswith("]"):
+        try:
+            return yaml.safe_load(raw)
+        except yaml.YAMLError:
+            pass
+    return raw
+
+
+def _apply_overrides(
+    cfg: dict[str, Any],
+    overrides: dict[str, Any],
+) -> None:
+    """Apply dotted-key overrides to a nested dict."""
+    for dotted_key, value in overrides.items():
+        parts = dotted_key.split(".")
+        d = cfg
+        for part in parts[:-1]:
+            if part not in d or not isinstance(d[part], dict):
+                d[part] = {}
+            d = d[part]
+        d[parts[-1]] = value
+
+
+def load_run_config(
+    path: str,
+    overrides: list[str] | None = None,
+) -> RunConfig:
+    """Load a run config from YAML.
+
+    Expected top-level keys: ``models``, ``method``,
+    ``training`` (nested), and optionally ``callbacks``
+    and ``pipeline``.
+
+    Args:
+        path: Path to the YAML config file.
+        overrides: Optional list of CLI override tokens,
+            e.g. ``["--training.distributed.num_gpus", "4"]``.
+            Dotted keys map to nested YAML paths.
+    """
+    path = _resolve_existing_file(path)
+    with open(path, encoding="utf-8") as f:
+        raw = yaml.safe_load(f)
+    cfg = _require_mapping(raw, where=path)
+
+    # Apply CLI overrides before building typed config.
+    if overrides:
+        parsed = _parse_cli_overrides(overrides)
+        _apply_overrides(cfg, parsed)
+        logger.info("Applied CLI overrides: %s", parsed)
+
+    # --- models ---
+    models_raw = _require_mapping(cfg.get("models"), where="models")
+    models: dict[str, dict[str, Any]] = {}
+    for role, model_cfg_raw in models_raw.items():
+        role_str = _require_str(role, where="models.<role>")
+        model_cfg = _require_mapping(model_cfg_raw, where=f"models.{role_str}")
+        if "_target_" not in model_cfg:
+            raise ValueError(f"models.{role_str} must have a "
+                             "'_target_' key")
+        models[role_str] = dict(model_cfg)
+
+    # --- method ---
+    method_raw = _require_mapping(cfg.get("method"), where="method")
+    if "_target_" not in method_raw:
+        raise ValueError("method must have a '_target_' key")
+    method = dict(method_raw)
+
+    # --- callbacks ---
+    callbacks_raw = cfg.get("callbacks", None)
+    if callbacks_raw is None:
+        callbacks: dict[str, dict[str, Any]] = {}
+    else:
+        callbacks = _require_mapping(callbacks_raw, where="callbacks")
+
+    # --- pipeline config ---
+    pipeline_config = _parse_pipeline_config(cfg, models=models)
+
+    # --- training config ---
+    training_raw = _require_mapping(cfg.get("training"), where="training")
+    t = dict(training_raw)
+    training = _build_training_config(t, models=models, pipeline_config=pipeline_config)
+
+    return RunConfig(
+        models=models,
+        method=method,
+        training=training,
+        callbacks=callbacks,
+        raw=cfg,
+    )

fastvideo/train/utils/dataloader.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        DataConfig, )
+
+
+def build_parquet_t2v_train_dataloader(
+    data_config: DataConfig,
+    *,
+    text_len: int,
+    parquet_schema: Any,
+) -> Any:
+    """Build a parquet dataloader for T2V-style datasets."""
+
+    from fastvideo.dataset import (
+        build_parquet_map_style_dataloader, )
+
+    _dataset, dataloader = (build_parquet_map_style_dataloader(
+        data_config.data_path,
+        data_config.train_batch_size,
+        num_data_workers=(data_config.dataloader_num_workers),
+        parquet_schema=parquet_schema,
+        cfg_rate=data_config.training_cfg_rate,
+        drop_last=True,
+        text_padding_length=int(text_len),
+        seed=int(data_config.seed or 0),
+    ))
+    return dataloader

fastvideo/train/utils/instantiate.py

@@ -0,0 +1,89 @@
+# SPDX-License-Identifier: Apache-2.0
+"""``_target_``-based instantiation utilities.
+
+These helpers resolve a dotted Python path to a class and instantiate it,
+filtering constructor kwargs through ``inspect.signature`` so that only
+recognized parameters are forwarded.  Unrecognized keys emit a warning
+rather than raising — this keeps YAML configs forward-compatible when
+a class drops a parameter in a later version.
+"""
+
+from __future__ import annotations
+
+import importlib
+import inspect
+import warnings
+from typing import Any
+
+
+def resolve_target(target: str) -> type:
+    """Import and return the class (or callable) at *target*.
+
+    *target* must be a fully-qualified dotted path, e.g.
+    ``"fastvideo.train.models.wan.wan.WanModel"``.
+    """
+    if not isinstance(target, str) or not target.strip():
+        raise ValueError(f"_target_ must be a non-empty dotted path string, "
+                         f"got {target!r}")
+    target = target.strip()
+    parts = target.rsplit(".", 1)
+    if len(parts) != 2:
+        raise ValueError(f"_target_ must contain at least one dot "
+                         f"(module.ClassName), got {target!r}")
+    module_path, attr_name = parts
+    try:
+        module = importlib.import_module(module_path)
+    except ModuleNotFoundError as exc:
+        raise ImportError(f"Cannot import module {module_path!r} "
+                          f"(from _target_={target!r})") from exc
+    try:
+        cls = getattr(module, attr_name)
+    except AttributeError as exc:
+        raise ImportError(f"Module {module_path!r} has no attribute "
+                          f"{attr_name!r} (from _target_={target!r})") from exc
+    return cls
+
+
+def instantiate(cfg: dict[str, Any], **extra: Any) -> Any:
+    """Instantiate the class specified by ``cfg["_target_"]``.
+
+    All remaining keys in *cfg* (minus ``_target_``) plus any *extra*
+    keyword arguments are forwarded to the constructor.  Keys that do
+    not match an ``__init__`` parameter are silently warned about and
+    dropped, so callers can safely pass a superset.
+    """
+    if not isinstance(cfg, dict):
+        raise TypeError(f"instantiate() expects a dict with '_target_', "
+                        f"got {type(cfg).__name__}")
+    target_str = cfg.get("_target_")
+    if target_str is None:
+        raise KeyError("Config dict is missing '_target_' key")
+
+    cls = resolve_target(str(target_str))
+    kwargs: dict[str, Any] = {k: v for k, v in cfg.items() if k != "_target_"}
+    kwargs.update(extra)
+
+    sig = inspect.signature(cls.__init__)  # type: ignore[misc]
+    params = sig.parameters
+    has_var_keyword = any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
+
+    if not has_var_keyword:
+        valid_names = {
+            name
+            for name, p in params.items() if p.kind in (
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                inspect.Parameter.KEYWORD_ONLY,
+            )
+        }
+        valid_names.discard("self")
+        unrecognized = set(kwargs) - valid_names
+        if unrecognized:
+            warnings.warn(
+                f"instantiate({target_str}): dropping unrecognized "
+                f"kwargs {sorted(unrecognized)}",
+                stacklevel=2,
+            )
+            for key in unrecognized:
+                del kwargs[key]
+
+    return cls(**kwargs)

fastvideo/train/utils/module_state.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import torch
+
+
+def apply_trainable(module: torch.nn.Module, *, trainable: bool) -> torch.nn.Module:
+    """Apply train/eval mode + requires_grad based on a role's trainable flag."""
+
+    module.requires_grad_(bool(trainable))
+    if trainable:
+        module.train()
+    else:
+        module.eval()
+    return module

fastvideo/train/utils/moduleloader.py

@@ -0,0 +1,144 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from typing import Any, TYPE_CHECKING
+
+import torch
+
+from fastvideo.configs.pipelines.base import PipelineConfig
+from fastvideo.fastvideo_args import ExecutionMode, TrainingArgs
+from fastvideo.models.loader.component_loader import (
+    PipelineComponentLoader, )
+from fastvideo.utils import (
+    maybe_download_model,
+    verify_model_config_and_directory,
+)
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        TrainingConfig, )
+
+# ------------------------------------------------------------------
+# TrainingArgs builders (only place that creates FastVideoArgs)
+# ------------------------------------------------------------------
+
+
+def _make_training_args(
+    tc: TrainingConfig,
+    *,
+    model_path: str,
+) -> TrainingArgs:
+    """Build a TrainingArgs for PipelineComponentLoader."""
+    pipeline_config = tc.pipeline_config or PipelineConfig()
+    # Propagate dit_precision from TrainingConfig to PipelineConfig
+    # so that TransformerLoader.load() picks up the correct
+    # default_dtype (e.g. fp32 master weights for training).
+    if tc.dit_precision and tc.dit_precision != pipeline_config.dit_precision:
+        pipeline_config.dit_precision = tc.dit_precision
+    return TrainingArgs(
+        model_path=model_path,
+        mode=ExecutionMode.DISTILLATION,
+        inference_mode=False,
+        pipeline_config=pipeline_config,
+        num_gpus=tc.distributed.num_gpus,
+        tp_size=tc.distributed.tp_size,
+        sp_size=tc.distributed.sp_size,
+        hsdp_replicate_dim=tc.distributed.hsdp_replicate_dim,
+        hsdp_shard_dim=tc.distributed.hsdp_shard_dim,
+        pin_cpu_memory=tc.distributed.pin_cpu_memory,
+        dit_cpu_offload=False,
+        dit_layerwise_offload=False,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=False,
+        image_encoder_cpu_offload=False,
+        use_fsdp_inference=False,
+        enable_torch_compile=False,
+    )
+
+
+def make_inference_args(
+    tc: TrainingConfig,
+    *,
+    model_path: str,
+) -> TrainingArgs:
+    """Build a TrainingArgs for inference (validation / pipelines)."""
+    args = _make_training_args(tc, model_path=model_path)
+    args.inference_mode = True
+    args.mode = ExecutionMode.INFERENCE
+    args.dit_cpu_offload = True
+    args.VSA_sparsity = tc.vsa.sparsity
+    return args
+
+
+# ------------------------------------------------------------------
+# Module loading
+# ------------------------------------------------------------------
+
+
+def load_module_from_path(
+    *,
+    model_path: str,
+    module_type: str,
+    training_config: TrainingConfig,
+    disable_custom_init_weights: bool = False,
+    override_transformer_cls_name: str | None = None,
+) -> torch.nn.Module:
+    """Load a single pipeline component module.
+
+    Accepts a ``TrainingConfig`` and internally builds the
+    ``TrainingArgs`` needed by ``PipelineComponentLoader``.
+    """
+    fastvideo_args: Any = _make_training_args(training_config, model_path=model_path)
+
+    local_model_path = maybe_download_model(model_path)
+    config = verify_model_config_and_directory(local_model_path)
+
+    if module_type not in config:
+        raise ValueError(f"Module {module_type!r} not found in "
+                         f"config at {local_model_path}")
+
+    module_info = config[module_type]
+    if module_info is None:
+        raise ValueError(f"Module {module_type!r} has null value in "
+                         f"config at {local_model_path}")
+
+    transformers_or_diffusers, _architecture = module_info
+    component_path = os.path.join(local_model_path, module_type)
+
+    old_override: str | None = None
+    if override_transformer_cls_name is not None:
+        old_override = getattr(
+            fastvideo_args,
+            "override_transformer_cls_name",
+            None,
+        )
+        fastvideo_args.override_transformer_cls_name = str(override_transformer_cls_name)
+
+    if disable_custom_init_weights:
+        fastvideo_args._loading_teacher_critic_model = True
+    try:
+        module = PipelineComponentLoader.load_module(
+            module_name=module_type,
+            component_model_path=component_path,
+            transformers_or_diffusers=(transformers_or_diffusers),
+            fastvideo_args=fastvideo_args,
+        )
+    finally:
+        if disable_custom_init_weights and hasattr(fastvideo_args, "_loading_teacher_critic_model"):
+            del fastvideo_args._loading_teacher_critic_model
+        if override_transformer_cls_name is not None:
+            if old_override is None:
+                if hasattr(
+                        fastvideo_args,
+                        "override_transformer_cls_name",
+                ):
+                    fastvideo_args.override_transformer_cls_name = (None)
+            else:
+                fastvideo_args.override_transformer_cls_name = (old_override)
+
+    if not isinstance(module, torch.nn.Module):
+        raise TypeError(f"Loaded {module_type!r} is not a "
+                        f"torch.nn.Module: {type(module)}")
+    return module

fastvideo/train/utils/optimizer.py

@@ -0,0 +1,72 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from fastvideo.training.training_utils import (
+    clip_grad_norm_while_handling_failing_dtensor_cases,
+    get_scheduler,
+)
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        OptimizerConfig,
+        TrainingLoopConfig,
+    )
+
+
+def build_optimizer_and_scheduler(
+    *,
+    params: list[torch.nn.Parameter],
+    optimizer_config: OptimizerConfig,
+    loop_config: TrainingLoopConfig,
+    learning_rate: float,
+    betas: tuple[float, float],
+    scheduler_name: str,
+) -> tuple[torch.optim.Optimizer, object]:
+    """Build an AdamW optimizer and LR scheduler.
+
+    Returns ``(optimizer, lr_scheduler)`` so the caller can store them
+    as method-level attributes.
+    """
+    if not params:
+        raise ValueError("No trainable parameters passed to "
+                         "build_optimizer_and_scheduler")
+
+    optimizer = torch.optim.AdamW(
+        params,
+        lr=float(learning_rate),
+        betas=betas,
+        weight_decay=float(optimizer_config.weight_decay),
+        eps=1e-8,
+    )
+
+    scheduler = get_scheduler(
+        str(scheduler_name),
+        optimizer=optimizer,
+        num_warmup_steps=int(optimizer_config.lr_warmup_steps),
+        num_training_steps=int(loop_config.max_train_steps),
+        num_cycles=int(optimizer_config.lr_num_cycles),
+        power=float(optimizer_config.lr_power),
+        min_lr_ratio=float(optimizer_config.min_lr_ratio),
+        last_epoch=-1,
+    )
+
+    return optimizer, scheduler
+
+
+def clip_grad_norm_if_needed(
+    module: torch.nn.Module,
+    max_grad_norm: float,
+) -> float:
+    if max_grad_norm <= 0.0:
+        return 0.0
+    grad_norm = (clip_grad_norm_while_handling_failing_dtensor_cases(
+        [p for p in module.parameters()],
+        max_grad_norm,
+        foreach=None,
+    ))
+    return (float(grad_norm.item()) if grad_norm is not None else 0.0)

fastvideo/train/utils/tracking.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import os
+from typing import Any, TYPE_CHECKING
+
+from fastvideo.distributed import get_world_group
+from fastvideo.training.trackers import (
+    initialize_trackers,
+    Trackers,
+)
+
+if TYPE_CHECKING:
+    from fastvideo.train.utils.training_config import (
+        CheckpointConfig,
+        TrackerConfig,
+    )
+
+
+def build_tracker(
+    tracker_config: TrackerConfig,
+    checkpoint_config: CheckpointConfig,
+    *,
+    config: dict[str, Any] | None,
+) -> Any:
+    """Build a tracker instance for a distillation run."""
+
+    world_group = get_world_group()
+
+    trackers = list(tracker_config.trackers)
+    if not trackers and str(tracker_config.project_name):
+        trackers.append(Trackers.WANDB.value)
+    if world_group.rank != 0:
+        trackers = []
+
+    tracker_log_dir = (checkpoint_config.output_dir or os.getcwd())
+    if trackers:
+        tracker_log_dir = os.path.join(tracker_log_dir, "tracker")
+
+    tracker_config_dict = config if trackers else None
+    tracker_run_name = tracker_config.run_name or None
+    project = (tracker_config.project_name or "fastvideo")
+
+    return initialize_trackers(
+        trackers,
+        experiment_name=project,
+        config=tracker_config_dict,
+        log_dir=tracker_log_dir,
+        run_name=tracker_run_name,
+    )

fastvideo/train/utils/training_config.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Typed training config — replaces TrainingArgs."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from fastvideo.configs.pipelines.base import PipelineConfig
+
+
+@dataclass(slots=True)
+class DistributedConfig:
+    num_gpus: int = 1
+    tp_size: int = 1
+    sp_size: int = 1
+    hsdp_replicate_dim: int = 1
+    hsdp_shard_dim: int = -1
+    pin_cpu_memory: bool = False
+
+
+@dataclass(slots=True)
+class DataConfig:
+    data_path: str = ""
+    train_batch_size: int = 1
+    dataloader_num_workers: int = 0
+    training_cfg_rate: float = 0.0
+    seed: int = 0
+    num_height: int = 0
+    num_width: int = 0
+    num_latent_t: int = 0
+    num_frames: int = 0
+
+
+@dataclass(slots=True)
+class OptimizerConfig:
+    learning_rate: float = 0.0
+    betas: tuple[float, float] = (0.9, 0.999)
+    weight_decay: float = 0.0
+    lr_scheduler: str = "constant"
+    lr_warmup_steps: int = 0
+    lr_num_cycles: int = 0
+    lr_power: float = 0.0
+    min_lr_ratio: float = 0.5
+
+
+@dataclass(slots=True)
+class TrainingLoopConfig:
+    max_train_steps: int = 0
+    gradient_accumulation_steps: int = 1
+
+
+@dataclass(slots=True)
+class CheckpointConfig:
+    output_dir: str = ""
+    resume_from_checkpoint: str = ""
+    training_state_checkpointing_steps: int = 0
+    checkpoints_total_limit: int = 0
+
+
+@dataclass(slots=True)
+class TrackerConfig:
+    trackers: list[str] = field(default_factory=list)
+    project_name: str = "fastvideo"
+    run_name: str = ""
+
+
+@dataclass(slots=True)
+class VSAConfig:
+    sparsity: float = 0.0
+    decay_rate: float = 0.0
+    decay_interval_steps: int = 0
+
+
+@dataclass(slots=True)
+class ModelTrainingConfig:
+    weighting_scheme: str = "uniform"
+    logit_mean: float = 0.0
+    logit_std: float = 1.0
+    mode_scale: float = 1.0
+    precondition_outputs: bool = False
+    moba_config: dict = field(default_factory=dict)
+    enable_gradient_checkpointing_type: str | None = None
+
+
+@dataclass(slots=True)
+class TrainingConfig:
+    distributed: DistributedConfig = field(default_factory=DistributedConfig)
+    data: DataConfig = field(default_factory=DataConfig)
+    optimizer: OptimizerConfig = field(default_factory=OptimizerConfig)
+    loop: TrainingLoopConfig = field(default_factory=TrainingLoopConfig)
+    checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)
+    tracker: TrackerConfig = field(default_factory=TrackerConfig)
+    vsa: VSAConfig = field(default_factory=VSAConfig)
+    model: ModelTrainingConfig = field(default_factory=ModelTrainingConfig)
+    pipeline_config: PipelineConfig | None = None
+    model_path: str = ""
+    dit_precision: str = "fp32"

fastvideo/training/trackers.py

@@ -89,6 +89,13 @@ class BaseTracker:
         if artifacts:
             self.log(artifacts, step)
 
+    def log_file(
+        self,
+        file_path: str,
+        name: str | None = None,
+    ) -> None:
+        """Attach a file to the tracker run (e.g. config YAML)."""
+
     def finish(self) -> None:  # pragma: no cover - interface
         """Finalize the tracker session."""
 
@@ -151,6 +158,16 @@ class WandbTracker(BaseTracker):
             self._run.log(metrics, step=step)
         self._timed_metrics = {}
 
+    def log_file(
+        self,
+        file_path: str,
+        name: str | None = None,
+    ) -> None:
+        self._wandb.save(
+            file_path,
+            base_path=os.path.dirname(file_path),
+        )
+
     def finish(self) -> None:
         self._run.finish()
 
@@ -201,6 +218,14 @@ class SequentialTracker(BaseTracker):
             tracker.log_artifacts(artifacts, step)
         self._timed_metrics = {}
 
+    def log_file(
+        self,
+        file_path: str,
+        name: str | None = None,
+    ) -> None:
+        for tracker in self._trackers:
+            tracker.log_file(file_path, name=name)
+
     def finish(self) -> None:
         for tracker in self._trackers:
             tracker.finish()

fastvideo/training/training_utils.py

@@ -65,7 +65,8 @@ def gather_state_dict_on_cpu_rank0(
 def compute_density_for_timestep_sampling(
     weighting_scheme: str,
     batch_size: int,
-    generator,
+    generator: torch.Generator,
+    device: torch.device | str = "cpu",
     logit_mean: float | None = None,
     logit_std: float | None = None,
     mode_scale: float | None = None,
@@ -83,15 +84,15 @@ def compute_density_for_timestep_sampling(
             mean=logit_mean,
             std=logit_std,
             size=(batch_size, ),
-            device="cpu",
+            device=device,
             generator=generator,
         )
         u = torch.nn.functional.sigmoid(u)
     elif weighting_scheme == "mode":
-        u = torch.rand(size=(batch_size, ), device="cpu", generator=generator)
+        u = torch.rand(size=(batch_size, ), device=device, generator=generator)
         u = 1 - u - mode_scale * (torch.cos(math.pi * u / 2)**2 - 1 + u)
     else:
-        u = torch.rand(size=(batch_size, ), device="cpu", generator=generator)
+        u = torch.rand(size=(batch_size, ), device=device, generator=generator)
     return u
 
 

mkdocs.yml

@@ -137,6 +137,7 @@ nav:
       - STA Mask Search (Archived): inference/examples/sta_mask_search.md
   - Training:
     - Overview: training/overview.md
+    - Training Infrastructure: training/train_infra.md
     - Data Preprocessing: training/data_preprocess.md
     - Fine-tuning: training/finetune.md
     - Examples:
@@ -159,6 +160,7 @@ nav:
     - Debugging: utilities/debugging.md
   - Design:
     - Overview: design/overview.md
+    - Training Architecture: design/training_architecture.md
   - Developer Guide:
     - Overview: contributing/overview.md
     - Developer Environment: