Wan 2.2 training - #470
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces training support for the Wan 2.2 model, including a new trainer (WanTrainer2_2), training script, configuration updates, and import smoke tests. It also switches GCS model downloads to use gcloud storage to prevent SSL segfaults, and optimizes disk usage during shard conversion by deleting shards after processing. The code review identified several critical issues and improvement opportunities: a logical error in the dataset validation check in WanTrainer2_2 where 'and' was used instead of 'or'; a JAX purity violation caused by mutating the input dictionary in-place inside the JIT-compiled train_step_2_2; incorrect evaluation routing for mixed-timestep batches, which should be resolved using jnp.where instead of checking only the first timestep; potential out-of-memory errors from downloading large models to /dev/shm instead of /tmp; performance overhead from recreating a ThreadPoolExecutor inside a loop during shard conversion; and the use of a mutable default argument in training_loop_2_2.
|
Can you please add the current timestep in the PR description |
|
Please add relevant unittests |
Done! |
Done! |
ff2a971 to
4ffe7dc
Compare
8edcb2d to
1a6b32e
Compare
a1f349a to
385c2d9
Compare
d674d8a to
4ddaa08
Compare
40ff613 to
65bbdf5
Compare
…expert routing - Implement WanTrainer2_2 dual-expert joint training pipeline with probabilistic batch routing based on boundary_ratio - Use unbiased uniform integer timestep sampling [boundary, num_train_timesteps) for high-noise expert and [0, boundary) for low-noise expert - Pass (state_high, state_low) functionally as operands to jax.lax.cond without outer closures - Isolate buffer updates to avoid buffer donation hazards on untouched states - Implement batched conditional evaluation in eval_step_2_2 with exact per-sample routing and constant HLO graph complexity - Strip redundant process_allgather on replicated evaluation metrics in eval_2_2 - Save and restore both low_noise_transformer and high_noise_transformer configurations and states in WanCheckpointer2_2 - Track active expert step counts on host to log active learning rates and eliminate device-to-host sync stalls - Harmonize TensorBoard writer metric_step tracking for per-expert steps in train_utils.py - Document checkpoint_save_location as local staging cache with disk capacity considerations in base_wan_27b.yml - Add comprehensive test suite covering training steps, eval steps, checkpointing, and resume equivalence
65bbdf5 to
97ebfc3
Compare
5 participants
Title:
Add Wan 2.2 Training Pipeline with Joint Timestep Routing
Description:
This PR introduces full training support for the Wan 2.2 models into MaxDiffusion.
1. Joint Timestep Routing
This implementation utilizes a unified, joint-pipeline training strategy. During each training step, the pipeline dynamically samples a timestep and routes the forward pass to either the high-noise or low-noise transformer based on the configured boundary ratio:
is_high_noise = jax.random.uniform(cond_rng) > config.boundary_ratio
This ensures a seamless and efficient training execution graph without needing separate high/low training loops.
2. Pipeline Inheritance & Code Reuse
The new wan_trainer_2_2.py relies directly on the WanPipeline2_2 class to bootstrap the models and load checkpoints. Because WanPipeline2_2 cleanly inherits from the base WanPipeline, the training loop seamlessly reuses all existing weight conversion and fast-loading logic that was previously established for inference.
3. Native Training Quantization Support
Because the trainer hooks directly into the unified WanPipeline.from_pretrained() (and from_checkpoint()) initialization sequence, quantization is automatically supported out-of-the-box for training. By passing use_qwix_quantization=True (and e.g., quantization="fp8"), the pipeline's innate quantize_transformer() loop executes before handing the model over to the trainer. This seamlessly delivers a quantized transformer straight to the optimizer with zero additional training code overhead. (Note: LoRA remains inference-only).
The loss graphs were plotted for around 260 steps, and show a clear downward trend.
For graphs and other artefacts: https://docs.google.com/document/d/1svzC8cVZxb2XxyypeFcoJig13_1ptu6wYIe5QwnC_Lo/edit?usp=sharing
Step time: Around 39.7 seconds per device.