Wan 2.1 training: Resolve checkpoint loading issues with larger TPU slices and different topologies

Co-authored-by: martinarroyo <martinarroyo@google.com>

Author: ninatu

src/maxdiffusion/checkpointing/wan_checkpointer_2_1.py

@@ -15,14 +15,15 @@
 """
 
 import json
-import jax
-import numpy as np
 from typing import Optional, Tuple
-from ..pipelines.wan.wan_pipeline_2_1 import WanPipeline2_1
-from .. import max_logging
-import orbax.checkpoint as ocp
 from etils import epath
+import jax
+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
 from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
+import numpy as np
+import orbax.checkpoint as ocp
+from .. import max_logging
+from ..pipelines.wan.wan_pipeline_2_1 import WanPipeline2_1
 
 
 class WanCheckpointer2_1(WanCheckpointer):
@@ -35,13 +36,32 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
         max_logging.log("No WAN checkpoint found.")
         return None, None
     max_logging.log(f"Loading WAN checkpoint from step {step}")
+
+    cpu_devices = np.array(jax.devices(backend="cpu"))
+    mesh = Mesh(cpu_devices, axis_names=("data",))
+    replicated_sharding = NamedSharding(mesh, P())
+
     metadatas = self.checkpoint_manager.item_metadata(step)
-    transformer_metadata = metadatas.wan_state
-    abstract_tree_structure_params = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, transformer_metadata)
+    state = metadatas.wan_state
+
+    def add_sharding_to_struct(leaf_struct, sharding):
+      return jax.ShapeDtypeStruct(
+          shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding
+      )
+
+    target_shardings = jax.tree_util.tree_map(
+        lambda x: replicated_sharding, state
+    )
+
+    with mesh:
+      abstract_train_state_with_sharding = jax.tree_util.tree_map(
+          add_sharding_to_struct, state, target_shardings
+      )
+
     params_restore = ocp.args.PyTreeRestore(
         restore_args=jax.tree.map(
             lambda _: ocp.RestoreArgs(restore_type=np.ndarray),
-            abstract_tree_structure_params,
+            abstract_train_state_with_sharding,
         )
     )
 

src/maxdiffusion/checkpointing/wan_checkpointer_i2v_2p1.py

@@ -15,14 +15,15 @@
 """
 
 import json
-import jax
-import numpy as np
 from typing import Optional, Tuple
-from ..pipelines.wan.wan_pipeline_i2v_2p1 import WanPipelineI2V_2_1
-from .. import max_logging
-import orbax.checkpoint as ocp
 from etils import epath
+import jax
+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
 from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
+import numpy as np
+import orbax.checkpoint as ocp
+from .. import max_logging
+from ..pipelines.wan.wan_pipeline_i2v_2p1 import WanPipelineI2V_2_1
 
 
 class WanCheckpointerI2V_2_1(WanCheckpointer):
@@ -35,13 +36,32 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
         max_logging.log("No WAN checkpoint found.")
         return None, None
     max_logging.log(f"Loading WAN checkpoint from step {step}")
+
+    cpu_devices = np.array(jax.devices(backend="cpu"))
+    mesh = Mesh(cpu_devices, axis_names=("data",))
+    replicated_sharding = NamedSharding(mesh, P())
+
     metadatas = self.checkpoint_manager.item_metadata(step)
-    transformer_metadata = metadatas.wan_state
-    abstract_tree_structure_params = jax.tree_util.tree_map(ocp.utils.to_shape_dtype_struct, transformer_metadata)
+    state = metadatas.wan_state
+
+    def add_sharding_to_struct(leaf_struct, sharding):
+      return jax.ShapeDtypeStruct(
+          shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding
+      )
+
+    target_shardings = jax.tree_util.tree_map(
+        lambda x: replicated_sharding, state
+    )
+
+    with mesh:
+      abstract_train_state_with_sharding = jax.tree_util.tree_map(
+          add_sharding_to_struct, state, target_shardings
+      )
+
     params_restore = ocp.args.PyTreeRestore(
         restore_args=jax.tree.map(
             lambda _: ocp.RestoreArgs(restore_type=np.ndarray),
-            abstract_tree_structure_params,
+            abstract_train_state_with_sharding,
         )
     )
 

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -168,9 +168,26 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
   )
   for path, val in flax.traverse_util.flatten_dict(params).items():
     if restored_checkpoint:
-      path = path[:-1]
+      if path[-1] == "value":
+        path = path[:-1]  # remove 'value'
+
+      try:
+        # Convert block indices to integers, as they might have been loaded as strings from the checkpoint.
+        path = path[:1] + (int(path[1]),) + path[2:]
+      except Exception:
+        pass
+
     sharding = logical_state_sharding[path].value
-    state[path].value = device_put_replicated(val, sharding)
+    try:
+      state[path].value = device_put_replicated(val, sharding)
+    except Exception as e:
+      max_logging.log(f"Failed to device_put_replicated {path}: {e}")
+      max_logging.log(f"Trying to use process_allgather for {path}")
+      val_on_host = jax.experimental.multihost_utils.process_allgather(
+          val, tiled=True
+      )
+      state[path].value = device_put_replicated(val_on_host, sharding)
+      del val_on_host
   state = nnx.from_flat_state(state)
 
   wan_transformer = nnx.merge(graphdef, state, rest_of_state)
@@ -470,7 +487,6 @@ def encode_prompt(
       negative_prompt_embeds: jax.Array = None,
   ):
     prompt = [prompt] if isinstance(prompt, str) else prompt
-    batch_size = len(prompt)
     if prompt_embeds is None:
       prompt_embeds = self._get_t5_prompt_embeds(
           prompt=prompt,
@@ -480,6 +496,7 @@ def encode_prompt(
       prompt_embeds = jnp.array(prompt_embeds.detach().numpy(), dtype=jnp.float32)
 
     if negative_prompt_embeds is None:
+      batch_size = len(prompt_embeds)
       negative_prompt = negative_prompt or ""
       negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
       negative_prompt_embeds = self._get_t5_prompt_embeds(

src/maxdiffusion/pipelines/wan/wan_vace_pipeline_2_1.py

@@ -119,9 +119,26 @@ def create_model(rngs: nnx.Rngs, wan_config: dict):
   )
   for path, val in flax.traverse_util.flatten_dict(params).items():
     if restored_checkpoint:
-      path = path[:-1]
+      if path[-1] == "value":
+        path = path[:-1]  # remove 'value'
+
+      try:
+        # Convert block indices to integers, as they might have been loaded as strings from the checkpoint.
+        path = path[:1] + (int(path[1]),) + path[2:]
+      except Exception:
+        pass
+
     sharding = logical_state_sharding[path].value
-    state[path].value = device_put_replicated(val, sharding)
+    try:
+      state[path].value = device_put_replicated(val, sharding)
+    except Exception as e:
+      max_logging.log(f"Failed to device_put_replicated {path}: {e}")
+      max_logging.log(f"Trying to use process_allgather for {path}")
+      val_on_host = jax.experimental.multihost_utils.process_allgather(
+          val, tiled=True
+      )
+      state[path].value = device_put_replicated(val_on_host, sharding)
+      del val_on_host
   state = nnx.from_flat_state(state)
 
   wan_transformer = nnx.merge(graphdef, state, rest_of_state)