feat(wan2.2): Implement Wan 2.2 joint training pipeline with dynamic expert routing

- Add WanTrainer2_2 joint training pipeline with probabilistic batch routing and Beta timestep sampling (Beta(5,2) for high-noise expert, Beta(2,5) for low-noise expert)
- Implement batched conditional evaluation (eval_step_2_2) with exact per-sample timestep routing and zero unrolled loop overhead
- Implement dual transformer checkpoint saving & restoration in WanCheckpointer2_2 with support for weights-only resume
- Prevent TPU VRAM overhead via branch-isolated gradient computation and optimizer application
- Support independent Optax learning rate schedule lengths scaled by boundary_ratio
- Add step-dependent PRNG key folding during evaluation passes
- Add comprehensive unit test suite with real JIT-compiled train_step and eval_step execution

Author: Toshi Pahadia

src/maxdiffusion/checkpointing/wan_checkpointer_2_2.py

@@ -18,7 +18,7 @@
 import jax
 from typing import Optional, Tuple
 from ..pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2
-from .. import max_logging
+from .. import max_logging, max_utils
 import orbax.checkpoint as ocp
 from maxdiffusion.checkpointing.checkpointing_utils import add_sharding_to_struct, get_cpu_mesh_and_sharding
 from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
@@ -27,6 +27,15 @@
 class WanCheckpointer2_2(WanCheckpointer[WanPipeline2_2]):
   pipeline_class = WanPipeline2_2
 
+  def _create_optimizer(self, model, config, learning_rate, scale_factor: float = 1.0):
+    total_steps = max(1, int(config.max_train_steps * scale_factor))
+    schedule_steps = max(1, int(config.learning_rate_schedule_steps * scale_factor))
+    learning_rate_scheduler = max_utils.create_learning_rate_schedule(
+        learning_rate, schedule_steps, config.warmup_steps_fraction, total_steps
+    )
+    tx = max_utils.create_optimizer(config, learning_rate_scheduler)
+    return tx, learning_rate_scheduler
+
   def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dict], Optional[int]]:
     if step is None:
       step = self.checkpoint_manager.latest_step()
@@ -81,11 +90,20 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
     return restored_checkpoint, step
 
   def _extract_opt_state(self, restored_checkpoint):
-    if "opt_state" in restored_checkpoint.low_noise_transformer_state.keys():
-      return restored_checkpoint.low_noise_transformer_state["opt_state"]
-    elif "opt_state" in restored_checkpoint.high_noise_transformer_state.keys():
-      return restored_checkpoint.high_noise_transformer_state["opt_state"]
-    return None
+    low_state = getattr(restored_checkpoint, "low_noise_transformer_state", {})
+    high_state = getattr(restored_checkpoint, "high_noise_transformer_state", {})
+    low_opt = low_state.get("opt_state") if isinstance(low_state, dict) else getattr(low_state, "opt_state", None)
+    high_opt = high_state.get("opt_state") if isinstance(high_state, dict) else getattr(high_state, "opt_state", None)
+    low_step = low_state.get("step") if isinstance(low_state, dict) else getattr(low_state, "step", None)
+    high_step = high_state.get("step") if isinstance(high_state, dict) else getattr(high_state, "step", None)
+    if low_opt is None and high_opt is None:
+      return None
+    return {
+        "low_noise_transformer": low_opt,
+        "high_noise_transformer": high_opt,
+        "low_noise_step": low_step,
+        "high_noise_step": high_step,
+    }
 
   def save_checkpoint(self, train_step, pipeline: WanPipeline2_2, train_states: dict):
     """Saves the training state and model configurations."""

src/maxdiffusion/configs/base_wan_27b.yml

@@ -409,6 +409,9 @@ num_inference_steps: 40
 fps: 16
 save_final_checkpoint: False
 
+# Location to download pretrained weights
+checkpoint_save_location: "/tmp"
+
 # SDXL Lightning parameters
 lightning_from_pt: True
 # Empty or "ByteDance/SDXL-Lightning" to enable lightning.

src/maxdiffusion/pyconfig.py

@@ -43,7 +43,7 @@
 )
 
 _ALLOWED_MODEL_NAMES = {WAN2_1, WAN2_2, LTX2_VIDEO, LTX2_3, Z_IMAGE}
-_ALLOWED_TRAINING_MODEL_NAMES = {WAN2_1}
+_ALLOWED_TRAINING_MODEL_NAMES = {WAN2_1, WAN2_2}
 
 
 def _validate_model_name(model_name: str | None):
@@ -283,12 +283,13 @@ def user_init(raw_keys):
 
     # Orbax doesn't save the tokenizer params, instead it loads them from the pretrained_model_name_or_path
     raw_keys["tokenizer_model_name_or_path"] = raw_keys["pretrained_model_name_or_path"]
+    ckpt_save_loc = raw_keys.get("checkpoint_save_location", "/tmp")
     if "gs://" in raw_keys["pretrained_model_name_or_path"]:
-      raw_keys["pretrained_model_name_or_path"] = max_utils.download_blobs(raw_keys["pretrained_model_name_or_path"], "/tmp")
+      raw_keys["pretrained_model_name_or_path"] = max_utils.download_blobs(raw_keys["pretrained_model_name_or_path"], ckpt_save_loc)
     if "gs://" in raw_keys["unet_checkpoint"]:
-      raw_keys["unet_checkpoint"] = max_utils.download_blobs(raw_keys["unet_checkpoint"], "/tmp")
+      raw_keys["unet_checkpoint"] = max_utils.download_blobs(raw_keys["unet_checkpoint"], ckpt_save_loc)
     if "gs://" in raw_keys["tokenizer_model_name_or_path"]:
-      raw_keys["tokenizer_model_name_or_path"] = max_utils.download_blobs(raw_keys["tokenizer_model_name_or_path"], "/tmp")
+      raw_keys["tokenizer_model_name_or_path"] = max_utils.download_blobs(raw_keys["tokenizer_model_name_or_path"], ckpt_save_loc)
     if "gs://" in raw_keys["dataset_name"]:
       raw_keys["dataset_name"] = max_utils.download_blobs(raw_keys["dataset_name"], raw_keys["dataset_save_location"])
       raw_keys["dataset_save_location"] = raw_keys["dataset_name"]

src/maxdiffusion/tests/wan/wan_checkpointer_test.py

@@ -387,7 +387,8 @@ def test_load_checkpoint_with_optimizer_in_low_noise(self, mock_from_checkpoint,
     )
     self.assertEqual(pipeline, mock_pipeline_instance)
     self.assertIsNotNone(opt_state)
-    self.assertEqual(opt_state["learning_rate"], 0.001)
+    self.assertEqual(opt_state["low_noise_transformer"]["learning_rate"], 0.001)
+    self.assertIsNone(opt_state["high_noise_transformer"])
     self.assertEqual(step, 1)
 
   @patch("maxdiffusion.checkpointing.wan_checkpointer.create_orbax_checkpoint_manager")
@@ -429,7 +430,8 @@ def test_load_checkpoint_with_optimizer_in_high_noise(self, mock_from_checkpoint
     )
     self.assertEqual(pipeline, mock_pipeline_instance)
     self.assertIsNotNone(opt_state)
-    self.assertEqual(opt_state["learning_rate"], 0.002)
+    self.assertIsNone(opt_state["low_noise_transformer"])
+    self.assertEqual(opt_state["high_noise_transformer"]["learning_rate"], 0.002)
     self.assertEqual(step, 1)
 
 
@@ -758,9 +760,10 @@ def test_load_checkpoint_both_optimizers_present(self, mock_from_checkpoint, moc
     checkpointer = WanCheckpointer2_2(config=self.config)
     pipeline, opt_state, step = checkpointer.load_checkpoint(step=1)
 
-    # Should prioritize low_noise_transformer's optimizer state
+    # Should preserve both low_noise_transformer and high_noise_transformer optimizer states
     self.assertIsNotNone(opt_state)
-    self.assertEqual(opt_state["learning_rate"], 0.001)
+    self.assertEqual(opt_state["low_noise_transformer"]["learning_rate"], 0.001)
+    self.assertEqual(opt_state["high_noise_transformer"]["learning_rate"], 0.002)
 
 
 if __name__ == "__main__":