Wan 2.2 training

[Toshi-31]

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.

[gemini-code-assist]

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.

[prishajain1]

Can you please add the current timestep in the PR description

[prishajain1]

Please add relevant unittests

[Toshi-31]

Can you please add the current timestep in the PR description

Done!

[Toshi-31]

Please add relevant unittests

Done!