Fixes all P1/P2 items from the review of 7c7704da: - E2E test's data_path pointed at a nonexistent combined_parquet_dataset/ subdir; preprocess_kandinsky5_overfit.py writes directly to its output root. - Stage 2 loaded the raw DCP checkpoint-N/ directly, which isn't a diffusers model dir. Added dcp_to_diffusers --verify (strict-reload the export immediately) and wired the conversion into the E2E test, launcher scripts, and README between stage 1 and stage 2. - Kandinsky5DmdDenoisingStage called the transformer without set_forward_context, so LocalAttention's missing-context guard caught the AssertionError and silently fell back to plain SDPA -- Stage 2 validation never actually exercised ATTN_QAT_TRAIN. Wrapped the call to match the parent stage's contract and added a scoped backend assertion (ATTN_QAT_TRAIN/ATTN_QAT_INFER only, since other backends can legitimately fall back to auto-selection). - Exported DMD checkpoints kept the base T2V model_index.json, so VideoGenerator.from_pretrained resolved them to Kandinsky5T2VPipeline instead of the DMD sampler. Added Kandinsky5DMDConfig and documented the required override_pipeline_cls_name + pipeline_config combo. - The FSDP-training dtype-detection special case (scan params, default to bf16) applied unconditionally, so an explicit fp32 pipeline was silently cast to bf16 for plain inference. Extracted _resolve_target_dtype and gated the scan on isinstance(transformer, FSDPModule). - Preprocessing tokenized captions with a 512-token cap that doesn't account for the 129-token system template runtime strips off; now reads text_encoder_max_lengths from pipeline_config to match runtime. Also replaced unconditional item["cap"][0] indexing (silently takes the first character of a string caption) with a get_caption() helper that accepts/validates both the documented string schema and the list-of-variants schema some producers use. - test_kandinsky5_qat_attention_engages.py set FASTVIDEO_ATTENTION_BACKEND and cleared the attention backend selector cache without teardown, leaking ATTN_QAT_TRAIN into later tests in the same pytest process. Now uses monkeypatch + try/finally. Added regression tests: test_kandinsky5_target_dtype.py, test_kandinsky5_dmd_pipeline_resolution.py, test_kandinsky5_overfit_caption.py. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Training
Quick Start
Single-node
bash examples/train/run.sh <config.yaml> [--dotted.key value ...]
The script auto-detects available GPUs. Override with NUM_GPUS:
NUM_GPUS=4 bash examples/train/run.sh \
examples/train/configs/fine_tuning/wan/t2v.yaml
Multi-node (Slurm)
bash examples/train/run_slurm.sh <config.yaml> <num_nodes> [--dotted.key value ...]
bash examples/train/run_slurm.sh \
examples/train/configs/fine_tuning/wan/t2v.yaml 4 \
--training.distributed.hsdp_shard_dim 8 \
--training.distributed.hsdp_replicate_dim 4
Slurm environment variables:
| Variable | Default | Description |
|---|---|---|
PARTITION |
main |
Slurm partition |
NUM_GPUS |
8 |
GPUs per node |
CPUS_PER_TASK |
128 |
CPUs per task |
MEM |
1440G |
Memory per node |
EXCLUDE |
"" |
Nodes to exclude |
Overriding config values
Any config field can be overridden from the command line using dotted keys:
bash examples/train/run.sh examples/train/configs/fine_tuning/wan/t2v.yaml \
--training.optimizer.learning_rate 1e-5 \
--training.loop.max_train_steps 1000 \
--training.checkpoint.output_dir outputs/my_experiment
Resuming from a checkpoint
bash examples/train/run.sh examples/train/configs/fine_tuning/wan/t2v.yaml \
--training.checkpoint.resume_from_checkpoint outputs/my_experiment/checkpoint-500
W&B Logging
Training metrics and validation videos are logged to Weights & Biases. Set your API key before launching:
export WANDB_API_KEY=your_key_here
To disable W&B (e.g. for local debugging):
export WANDB_MODE=offline
The project name and run name are set in the config under training.tracker:
training:
tracker:
project_name: my_project
run_name: my_run
Directory Layout
examples/train/
├── configs/ # Single-step training configs (by method and model)
├── scenario/ # Multi-step end-to-end training pipelines
├── run.sh # Single-node launcher
└── run_slurm.sh # Multi-node Slurm launcher
See configs/README.md and scenario/README.md for details.
The featured QAD Wan2.1 MixKit scenario runs Attn-QAT supervised fine-tuning, checkpoint export, and three-step DMD2. See the Attn-QAT training guide.