Wan 2.2 training

src/maxdiffusion/checkpointing/checkpointing_utils.py

@@ -74,9 +74,16 @@ def create_orbax_checkpoint_manager(
         "text_encoder_state": ocp.StandardCheckpointHandler(),
     }
   elif checkpoint_type == WAN_CHECKPOINT:
-    item_names = ("low_noise_transformer_state", "high_noise_transformer_state", "wan_state", "wan_config")
+    item_names = (
+        "low_noise_transformer_state",
+        "high_noise_transformer_state",
+        "wan_state",
+        "wan_config",
+        "wan_config_high",
+    )
     item_handlers = {
         "wan_config": ocp.JsonCheckpointHandler(),
+        "wan_config_high": ocp.JsonCheckpointHandler(),
         "wan_state": ocp.StandardCheckpointHandler(),
         "low_noise_transformer_state": ocp.StandardCheckpointHandler(),
         "high_noise_transformer_state": ocp.StandardCheckpointHandler(),

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()
@@ -56,13 +65,25 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
       )
 
     max_logging.log("Restoring WAN 2.2 checkpoint")
+    restore_items = {
+        "low_noise_transformer_state": ocp.args.StandardRestore(abstract_tree_structure_low_params),
+        "high_noise_transformer_state": ocp.args.StandardRestore(abstract_tree_structure_high_params),
+        "wan_config": ocp.args.JsonRestore(),
+    }
+    has_high_config = False
+    if hasattr(metadatas, "wan_config_high"):
+      val = getattr(metadatas, "wan_config_high")
+      if not hasattr(val, "_mock_return_value"):
+        has_high_config = True
+    elif isinstance(metadatas, dict) and "wan_config_high" in metadatas:
+      has_high_config = True
+
+    if has_high_config:
+      restore_items["wan_config_high"] = ocp.args.JsonRestore()
+
     restored_checkpoint = self.checkpoint_manager.restore(
         step=step,
-        args=ocp.args.Composite(
-            low_noise_transformer_state=ocp.args.StandardRestore(abstract_tree_structure_low_params),
-            high_noise_transformer_state=ocp.args.StandardRestore(abstract_tree_structure_high_params),
-            wan_config=ocp.args.JsonRestore(),
-        ),
+        args=ocp.args.Composite(**restore_items),
     )
     max_logging.log(f"restored checkpoint {restored_checkpoint.keys()}")
     max_logging.log(
@@ -81,11 +102,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."""
@@ -96,6 +126,7 @@ def config_to_json(model_or_config):
     max_logging.log(f"Saving checkpoint for step {train_step}")
     items = {
         "wan_config": ocp.args.JsonSave(config_to_json(pipeline.low_noise_transformer)),
+        "wan_config_high": ocp.args.JsonSave(config_to_json(pipeline.high_noise_transformer)),
     }
 
     items["low_noise_transformer_state"] = ocp.args.StandardSave(train_states["low_noise_transformer"])

src/maxdiffusion/configs/base_wan_27b.yml

@@ -409,6 +409,12 @@ num_inference_steps: 40
 fps: 16
 save_final_checkpoint: False
 
+# Staging directory for downloading pretrained weights from Google Cloud Storage (gs://)
+# prior to loading into device memory. Note: On Cloud TPU VMs with limited root disk/tmpfs,
+# ensure this directory resides on a persistent disk or large volume to avoid ENOSPC.
+# Differs from pretrained_model_name_or_path which specifies the remote/local model URI.
+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,17 @@ 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__":