diff --git a/.agents/memory/codebase-map/README.md b/.agents/memory/codebase-map/README.md index c19ced268..ce6917d9d 100644 --- a/.agents/memory/codebase-map/README.md +++ b/.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 ` | +| DFSFT (Wan causal) | `examples/train/dfsft_wan_causal_t2v_1.3B.yaml` | `torchrun -m fastvideo.train.entrypoint.train --config ` | +| DMD2 distillation | `examples/train/distill_wan2.1_t2v_1.3B_dmd2.yaml` | `torchrun -m fastvideo.train.entrypoint.train --config ` | +| Self-Forcing | `examples/train/self_forcing_wan_causal_t2v_1.3B.yaml` | `torchrun -m fastvideo.train.entrypoint.train --config ` | + +### 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` | diff --git a/.agents/memory/evaluation-registry/README.md b/.agents/memory/evaluation-registry/README.md index c54f960eb..10309d04f 100644 --- a/.agents/memory/evaluation-registry/README.md +++ b/.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 diff --git a/.agents/onboarding/README.md b/.agents/onboarding/README.md index 82b11b669..7590b5420 100644 --- a/.agents/onboarding/README.md +++ b/.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 diff --git a/.agents/onboarding/worldmodel-training/README.md b/.agents/onboarding/worldmodel-training/README.md index 339813119..be1d9a61d 100644 --- a/.agents/onboarding/worldmodel-training/README.md +++ b/.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 \ + -m fastvideo.train.entrypoint.train \ + --config examples/train/.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 `/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. diff --git a/.agents/skills/SKILL_TEMPLATE.md b/.agents/skills/SKILL_TEMPLATE.md index 4819b44a4..15cd50e41 100644 --- a/.agents/skills/SKILL_TEMPLATE.md +++ b/.agents/skills/SKILL_TEMPLATE.md @@ -28,6 +28,7 @@ description: ## Example Usage + ``` ``` @@ -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": "", "description": "", "path": "/SKILL.md", "status": "draft", "trust": "low"} ``` diff --git a/.agents/skills/evaluate-video-quality/SKILL.md b/.agents/skills/evaluate-video-quality/SKILL.md index 1a7b19793..ca50ac9cd 100644 --- a/.agents/skills/evaluate-video-quality/SKILL.md +++ b/.agents/skills/evaluate-video-quality/SKILL.md @@ -38,6 +38,7 @@ pytest fastvideo/tests/ssim/ -vs --video-path --reference-path \ ### 5. Log to experiment journal After launching, append an entry to `.agents/memory/experiment-journal/README.md`: + ```markdown ## [YYYY-MM-DD] Experiment: - **Hypothesis**: @@ -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: diff --git a/.agents/skills/log-experiment/SKILL.md b/.agents/skills/log-experiment/SKILL.md index a53ff55c4..0756655ab 100644 --- a/.agents/skills/log-experiment/SKILL.md +++ b/.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: diff --git a/.agents/skills/monitor-experiment/SKILL.md b/.agents/skills/monitor-experiment/SKILL.md index 5a9a0e8ea..947c014d5 100644 --- a/.agents/skills/monitor-experiment/SKILL.md +++ b/.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("") ``` **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": diff --git a/.agents/skills/search-related-work/SKILL.md b/.agents/skills/search-related-work/SKILL.md index 5fa2373ed..49ecaf9dc 100644 --- a/.agents/skills/search-related-work/SKILL.md +++ b/.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: diff --git a/.agents/skills/summarize-run/SKILL.md b/.agents/skills/summarize-run/SKILL.md index 58084c267..a0312172e 100644 --- a/.agents/skills/summarize-run/SKILL.md +++ b/.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": diff --git a/docs/design/training_architecture.md b/docs/design/training_architecture.md new file mode 100644 index 000000000..ff0b3e90c --- /dev/null +++ b/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. diff --git a/docs/training/train_infra.md b/docs/training/train_infra.md new file mode 100644 index 000000000..597bb6b6c --- /dev/null +++ b/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 `/checkpoint-/`. + +### 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. diff --git a/examples/train/configs/dfsft_wan_causal_t2v_1.3B.yaml b/examples/train/configs/dfsft_wan_causal_t2v_1.3B.yaml new file mode 100644 index 000000000..a2e3903be --- /dev/null +++ b/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 diff --git a/examples/train/configs/distill_wan2.1_t2v_1.3B_dmd2.yaml b/examples/train/configs/distill_wan2.1_t2v_1.3B_dmd2.yaml new file mode 100644 index 000000000..576bff26d --- /dev/null +++ b/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 diff --git a/examples/train/configs/example.yaml b/examples/train/configs/example.yaml new file mode 100644 index 000000000..5c9f226d1 --- /dev/null +++ b/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 diff --git a/examples/train/configs/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml b/examples/train/configs/finetune_wan2.1_t2v_1.3B_vsa_phase3.4_0.9sparsity.yaml new file mode 100644 index 000000000..0e418d680 --- /dev/null +++ b/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 diff --git a/examples/train/configs/self_forcing_wan_causal_t2v_1.3B.yaml b/examples/train/configs/self_forcing_wan_causal_t2v_1.3B.yaml new file mode 100644 index 000000000..ff19ad131 --- /dev/null +++ b/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 diff --git a/examples/train/run.sh b/examples/train/run.sh new file mode 100755 index 000000000..09a2a3f63 --- /dev/null +++ b/examples/train/run.sh @@ -0,0 +1,61 @@ +#!/usr/bin/env bash +# Launch training from a YAML config. +# +# Usage: +# bash examples/train/run.sh [--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/_.log (and also printed to stdout). + +set -euo pipefail + +CONFIG="${1:?Usage: $0 [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}" diff --git a/examples/train/run_slurm.sh b/examples/train/run_slurm.sh new file mode 100755 index 000000000..3d9f3439b --- /dev/null +++ b/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 [--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 [extra flags...]}" +NUM_NODES="${2:?Usage: $0 [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[@]}" < 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 diff --git a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py b/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py index eb4e67e0c..bdddd22eb 100644 --- a/fastvideo/pipelines/basic/wan/wan_causal_pipeline.py +++ b/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")), ) diff --git a/fastvideo/pipelines/stages/__init__.py b/fastvideo/pipelines/stages/__init__.py index 679039016..9bd5280b8 100644 --- a/fastvideo/pipelines/stages/__init__.py +++ b/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 ( diff --git a/fastvideo/pipelines/stages/causal_denoising.py b/fastvideo/pipelines/stages/causal_denoising.py index 69a961a0c..08cf96752 100644 --- a/fastvideo/pipelines/stages/causal_denoising.py +++ b/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, ) diff --git a/fastvideo/train/.style.yapf b/fastvideo/train/.style.yapf new file mode 100644 index 000000000..c9a88d5a6 --- /dev/null +++ b/fastvideo/train/.style.yapf @@ -0,0 +1,3 @@ +[style] +based_on_style = pep8 +column_limit = 120 diff --git a/fastvideo/train/__init__.py b/fastvideo/train/__init__.py new file mode 100644 index 000000000..fed6b183c --- /dev/null +++ b/fastvideo/train/__init__.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 + +from fastvideo.train.trainer import Trainer + +__all__ = [ + "Trainer", +] diff --git a/fastvideo/train/callbacks/__init__.py b/fastvideo/train/callbacks/__init__.py new file mode 100644 index 000000000..d57e048a9 --- /dev/null +++ b/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", +] diff --git a/fastvideo/train/callbacks/callback.py b/fastvideo/train/callbacks/callback.py new file mode 100644 index 000000000..024a92860 --- /dev/null +++ b/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 diff --git a/fastvideo/train/callbacks/ema.py b/fastvideo/train/callbacks/ema.py new file mode 100644 index 000000000..2bd9f01f0 --- /dev/null +++ b/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), ) diff --git a/fastvideo/train/callbacks/grad_clip.py b/fastvideo/train/callbacks/grad_clip.py new file mode 100644 index 000000000..4e8aa248a --- /dev/null +++ b/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, + ) diff --git a/fastvideo/train/callbacks/validation.py b/fastvideo/train/callbacks/validation.py new file mode 100644 index 000000000..29225a579 --- /dev/null +++ b/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) diff --git a/fastvideo/train/entrypoint/__init__.py b/fastvideo/train/entrypoint/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/train/entrypoint/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/train/entrypoint/dcp_to_diffusers.py b/fastvideo/train/entrypoint/dcp_to_diffusers.py new file mode 100644 index 000000000..8f4955c89 --- /dev/null +++ b/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.", + ) + 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- 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() diff --git a/fastvideo/train/entrypoint/misc/__init__.py b/fastvideo/train/entrypoint/misc/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastvideo/train/entrypoint/misc/wan_ode_init_conversion.py b/fastvideo/train/entrypoint/misc/wan_ode_init_conversion.py new file mode 100644 index 000000000..d75bd2652 --- /dev/null +++ b/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": {}}``. + +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() diff --git a/fastvideo/train/entrypoint/train.py b/fastvideo/train/entrypoint/train.py new file mode 100644 index 000000000..789a7fc59 --- /dev/null +++ b/fastvideo/train/entrypoint/train.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +"""YAML-only training entrypoint. + +Usage:: + + torchrun --nproc_per_node= -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) diff --git a/fastvideo/train/methods/__init__.py b/fastvideo/train/methods/__init__.py new file mode 100644 index 000000000..61fd6ef2e --- /dev/null +++ b/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) diff --git a/fastvideo/train/methods/base.py b/fastvideo/train/methods/base.py new file mode 100644 index 000000000..0f2898028 --- /dev/null +++ b/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..``, ``optimizers.``, + ``schedulers.``, ``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) diff --git a/fastvideo/train/methods/consistency_model/__init__.py b/fastvideo/train/methods/consistency_model/__init__.py new file mode 100644 index 000000000..324710b84 --- /dev/null +++ b/fastvideo/train/methods/consistency_model/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 + +__all__: list[str] = [] diff --git a/fastvideo/train/methods/distribution_matching/__init__.py b/fastvideo/train/methods/distribution_matching/__init__.py new file mode 100644 index 000000000..4edb43cf7 --- /dev/null +++ b/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", +] diff --git a/fastvideo/train/methods/distribution_matching/dmd2.py b/fastvideo/train/methods/distribution_matching/dmd2.py new file mode 100644 index 000000000..9162bae63 --- /dev/null +++ b/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 diff --git a/fastvideo/train/methods/distribution_matching/self_forcing.py b/fastvideo/train/methods/distribution_matching/self_forcing.py new file mode 100644 index 000000000..4939a85bd --- /dev/null +++ b/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 diff --git a/fastvideo/train/methods/fine_tuning/__init__.py b/fastvideo/train/methods/fine_tuning/__init__.py new file mode 100644 index 000000000..6f862df4e --- /dev/null +++ b/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) diff --git a/fastvideo/train/methods/fine_tuning/dfsft.py b/fastvideo/train/methods/fine_tuning/dfsft.py new file mode 100644 index 000000000..a0f778886 --- /dev/null +++ b/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] diff --git a/fastvideo/train/methods/fine_tuning/finetune.py b/fastvideo/train/methods/fine_tuning/finetune.py new file mode 100644 index 000000000..624a5ecab --- /dev/null +++ b/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, + ) diff --git a/fastvideo/train/methods/knowledge_distillation/__init__.py b/fastvideo/train/methods/knowledge_distillation/__init__.py new file mode 100644 index 000000000..324710b84 --- /dev/null +++ b/fastvideo/train/methods/knowledge_distillation/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 + +__all__: list[str] = [] diff --git a/fastvideo/train/models/__init__.py b/fastvideo/train/models/__init__.py new file mode 100644 index 000000000..56b47b1af --- /dev/null +++ b/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..family``. +""" diff --git a/fastvideo/train/models/base.py b/fastvideo/train/models/base.py new file mode 100644 index 000000000..6b97a9d0b --- /dev/null +++ b/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]) diff --git a/fastvideo/train/models/wan/__init__.py b/fastvideo/train/models/wan/__init__.py new file mode 100644 index 000000000..9a8113ac1 --- /dev/null +++ b/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, ) diff --git a/fastvideo/train/models/wan/wan.py b/fastvideo/train/models/wan/wan.py new file mode 100644 index 000000000..61e157ca8 --- /dev/null +++ b/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}") diff --git a/fastvideo/train/models/wan/wan_causal.py b/fastvideo/train/models/wan/wan_causal.py new file mode 100644 index 000000000..39b0fdbb1 --- /dev/null +++ b/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)] diff --git a/fastvideo/train/trainer.py b/fastvideo/train/trainer.py new file mode 100644 index 000000000..4a033b2c9 --- /dev/null +++ b/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() diff --git a/fastvideo/train/utils/__init__.py b/fastvideo/train/utils/__init__.py new file mode 100644 index 000000000..d7eba4033 --- /dev/null +++ b/fastvideo/train/utils/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Distillation utilities shared across families/methods/entrypoints.""" diff --git a/fastvideo/train/utils/builder.py b/fastvideo/train/utils/builder.py new file mode 100644 index 000000000..7028afb1e --- /dev/null +++ b/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 diff --git a/fastvideo/train/utils/checkpoint.py b/fastvideo/train/utils/checkpoint.py new file mode 100644 index 000000000..eca95d396 --- /dev/null +++ b/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-'") + 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- + - /path/to/output_dir/checkpoint-/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-' (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() diff --git a/fastvideo/train/utils/config.py b/fastvideo/train/utils/config.py new file mode 100644 index 000000000..40da284fc --- /dev/null +++ b/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.") + 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, + ) diff --git a/fastvideo/train/utils/dataloader.py b/fastvideo/train/utils/dataloader.py new file mode 100644 index 000000000..a1b22d3ba --- /dev/null +++ b/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 diff --git a/fastvideo/train/utils/instantiate.py b/fastvideo/train/utils/instantiate.py new file mode 100644 index 000000000..f8d7f8bc0 --- /dev/null +++ b/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) diff --git a/fastvideo/train/utils/module_state.py b/fastvideo/train/utils/module_state.py new file mode 100644 index 000000000..6d28a005f --- /dev/null +++ b/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 diff --git a/fastvideo/train/utils/moduleloader.py b/fastvideo/train/utils/moduleloader.py new file mode 100644 index 000000000..d775db9bd --- /dev/null +++ b/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 diff --git a/fastvideo/train/utils/optimizer.py b/fastvideo/train/utils/optimizer.py new file mode 100644 index 000000000..43a79d98d --- /dev/null +++ b/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) diff --git a/fastvideo/train/utils/tracking.py b/fastvideo/train/utils/tracking.py new file mode 100644 index 000000000..7ad28a2e4 --- /dev/null +++ b/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, + ) diff --git a/fastvideo/train/utils/training_config.py b/fastvideo/train/utils/training_config.py new file mode 100644 index 000000000..0167751db --- /dev/null +++ b/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" diff --git a/fastvideo/training/trackers.py b/fastvideo/training/trackers.py index 281d79325..fce905363 100644 --- a/fastvideo/training/trackers.py +++ b/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() diff --git a/fastvideo/training/training_utils.py b/fastvideo/training/training_utils.py index d4fa4efb9..89d47e43e 100644 --- a/fastvideo/training/training_utils.py +++ b/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 diff --git a/mkdocs.yml b/mkdocs.yml index 6a5c11bdf..42fb195bd 100644 --- a/mkdocs.yml +++ b/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: