Zhang PeiyuanandClaude Opus 5.5
c520878fd2
Adds opt-in regional torch.compile to the YAML-driven modular Wan trainer (training.model.enable_torch_compile / torch_compile_kwargs) with FA4 and Wan modulation forward/backward as opaque custom ops, compiled repeated DiT blocks after FSDP setup, and activation checkpointing applied before FSDP for trainable Wan transformers. Merged with main's #1716 (WanModel resolves the checkpointing type in __init__; the PR's pre-FSDP wrap uses it and the post-load wrap is gone) and #1757 (pre_fsdp_model_transform runs before this PR's pre_fsdp_transform). Maintainer review (Swarm order #776): AdaLN skip names are canonicalized through checkpoint-wrapper prefixes, test_inference_regional_compile.py joins the unit lane, the Wan modulation custom-op test moved to the transformer lane, regional_compile is classified in the schema inventory, and a 1/2-GPU training smoke covers fully_shard(CheckpointWrapper) in eager and regional-compile modes. Validated on GB200: the FA4 CuTe forward/backward parity and fullgraph-backward tests and all four smoke cases pass. Also exempts MMAudio recipes, whose student builds its pipeline config at init, from the fine-tuning recipe construction check that #1648 left failing on main. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Training Configs
Single-step training configurations organized by method and model.
configs/
├── fine_tuning/ # Standard finetuning and DFSFT
├── distribution_matching/ # DMD2 and Self-Forcing
├── knowledge_distillation/ # KD from teacher to student
└── example.yaml # Annotated reference config with all fields
Each method directory contains per-model subdirectories (e.g. wan/, hunyuan/).
Launch any config with:
bash examples/train/run.sh examples/train/configs/<method>/<model>/<config>.yaml
For multi-step training pipelines, see examples/train/scenario/.