Fix: Overhaul WAN checkpointers for robust multi-host restoration

This commit resolves several interrelated checkpointing issues by updating
how Orbax handles metadata, sharding, and PyTree restoration.

Key changes:
* Add explicit `item_handlers`: Defined specific handlers (`JsonCheckpointHandler`
  for configs, `StandardCheckpointHandler` for states) in `CheckpointManager`.
  This ensures metadata is restored correctly, resolving known Orbax limitations
  (reference: google/orbax#986).

* Bypass mesh validation during restore: Replaced `ocp.utils.to_shape_dtype_struct`
  with manual `jax.ShapeDtypeStruct` construction in `add_sharding_to_struct`.
  This makes restoration topology-agnostic, preventing `ValueError` when the
  current device mesh has fewer devices than the saved checkpoint's topology
  (e.g., restoring 32-device metadata on 4 devices).

* Migrate to Standard API: Upgraded all WAN checkpointers from
  the `PyTreeSave`/`PyTreeRestore` APIs to `StandardSave`/`StandardRestore`
  to align with `item_handlers` defined in CheckpointManager.

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

Author: ninatu

src/maxdiffusion/checkpointing/checkpointing_utils.py

@@ -58,10 +58,17 @@ def create_orbax_checkpoint_manager(
   max_logging.log(f"checkpoint dir: {checkpoint_dir}")
   p = epath.Path(checkpoint_dir)
 
+  item_handlers = None
   if checkpoint_type == FLUX_CHECKPOINT:
     item_names = ("flux_state", "flux_config", "vae_state", "vae_config", "scheduler", "scheduler_config")
   elif checkpoint_type == WAN_CHECKPOINT:
     item_names = ("low_noise_transformer_state", "high_noise_transformer_state", "wan_state", "wan_config")
+    item_handlers = {
+        "wan_config": ocp.JsonCheckpointHandler(),
+        "wan_state": ocp.StandardCheckpointHandler(),
+        "low_noise_transformer_state": ocp.StandardCheckpointHandler(),
+        "high_noise_transformer_state": ocp.StandardCheckpointHandler(),
+    }
   else:
     item_names = (
         "unet_config",
@@ -89,6 +96,7 @@ def create_orbax_checkpoint_manager(
       options=CheckpointManagerOptions(
           create=True, save_interval_steps=save_interval_steps, enable_async_checkpointing=use_async
       ),
+      item_handlers=item_handlers,
       logger=orbax_logger,
   )
 

src/maxdiffusion/checkpointing/wan_checkpointer_2_1.py

@@ -16,7 +16,6 @@
 
 import json
 from typing import Optional, Tuple
-from etils import epath
 import jax
 from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
 from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
@@ -44,31 +43,26 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
     metadatas = self.checkpoint_manager.item_metadata(step)
     state = metadatas.wan_state
 
+    # Manually constructs jax.ShapeDtypeStruct with a specific sharding.
+    # This avoids device mesh validation (as in ocp.utils.to_shape_dtype_struct)
+    # allowing for sharding with a different mesh than the one used during
+    # saving.
     def add_sharding_to_struct(leaf_struct, sharding):
-      struct = ocp.utils.to_shape_dtype_struct(leaf_struct)
-      if hasattr(struct, "shape") and hasattr(struct, "dtype"):
-        return jax.ShapeDtypeStruct(shape=struct.shape, dtype=struct.dtype, sharding=sharding)
-      return struct
+      if hasattr(leaf_struct, "shape") and hasattr(leaf_struct, "dtype"):
+        return jax.ShapeDtypeStruct(shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding)
+      return leaf_struct
 
     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=jax.Array),
-            abstract_train_state_with_sharding,
-        )
-    )
-
     max_logging.log("Restoring WAN checkpoint")
     restored_checkpoint = self.checkpoint_manager.restore(
-        directory=epath.Path(self.config.checkpoint_dir),
         step=step,
         args=ocp.args.Composite(
-            wan_state=params_restore,
             wan_config=ocp.args.JsonRestore(),
+            wan_state=ocp.args.StandardRestore(abstract_train_state_with_sharding),
         ),
     )
     max_logging.log(f"restored checkpoint {restored_checkpoint.keys()}")
@@ -106,7 +100,7 @@ def config_to_json(model_or_config):
         "wan_config": ocp.args.JsonSave(config_to_json(pipeline.transformer)),
     }
 
-    items["wan_state"] = ocp.args.PyTreeSave(train_states)
+    items["wan_state"] = ocp.args.StandardSave(train_states)
 
     # Save the checkpoint
     self.checkpoint_manager.save(train_step, args=ocp.args.Composite(**items))

src/maxdiffusion/checkpointing/wan_checkpointer_2_2.py

@@ -17,11 +17,11 @@
 import json
 import jax
 import numpy as np
+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
 from typing import Optional, Tuple
 from ..pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2
 from .. import max_logging
 import orbax.checkpoint as ocp
-from etils import epath
 from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
 
 
@@ -35,39 +35,45 @@ 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)
 
     # Handle low_noise_transformer
     low_noise_transformer_metadata = metadatas.low_noise_transformer_state
-    abstract_tree_structure_low_params = jax.tree_util.tree_map(
-        ocp.utils.to_shape_dtype_struct, low_noise_transformer_metadata
-    )
-    low_params_restore = ocp.args.PyTreeRestore(
-        restore_args=jax.tree.map(
-            lambda _: ocp.RestoreArgs(restore_type=np.ndarray),
-            abstract_tree_structure_low_params,
-        )
-    )
+
+    # Manually constructs jax.ShapeDtypeStruct with a specific sharding.
+    # This avoids device mesh validation (as in ocp.utils.to_shape_dtype_struct)
+    # allowing for sharding with a different mesh than the one used during
+    # saving.
+    def add_sharding_to_struct(leaf_struct, sharding):
+      if hasattr(leaf_struct, "shape") and hasattr(leaf_struct, "dtype"):
+        return jax.ShapeDtypeStruct(shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding)
+      return leaf_struct
+
+    target_shardings = jax.tree_util.tree_map(lambda x: replicated_sharding, low_noise_transformer_metadata)
+    with mesh:
+      abstract_tree_structure_low_params = jax.tree_util.tree_map(
+          add_sharding_to_struct, low_noise_transformer_metadata, target_shardings
+      )
 
     # Handle high_noise_transformer
     high_noise_transformer_metadata = metadatas.high_noise_transformer_state
-    abstract_tree_structure_high_params = jax.tree_util.tree_map(
-        ocp.utils.to_shape_dtype_struct, high_noise_transformer_metadata
-    )
-    high_params_restore = ocp.args.PyTreeRestore(
-        restore_args=jax.tree.map(
-            lambda _: ocp.RestoreArgs(restore_type=np.ndarray),
-            abstract_tree_structure_high_params,
-        )
-    )
+    target_shardings = jax.tree_util.tree_map(lambda x: replicated_sharding, high_noise_transformer_metadata)
+    with mesh:
+      abstract_tree_structure_high_params = jax.tree_util.tree_map(
+          add_sharding_to_struct, high_noise_transformer_metadata, target_shardings
+      )
 
     max_logging.log("Restoring WAN 2.2 checkpoint")
     restored_checkpoint = self.checkpoint_manager.restore(
-        directory=epath.Path(self.config.checkpoint_dir),
         step=step,
         args=ocp.args.Composite(
-            low_noise_transformer_state=low_params_restore,
-            high_noise_transformer_state=high_params_restore,
+            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(),
         ),
     )
@@ -119,8 +125,8 @@ def config_to_json(model_or_config):
         "wan_config": ocp.args.JsonSave(config_to_json(pipeline.low_noise_transformer)),
     }
 
-    items["low_noise_transformer_state"] = ocp.args.PyTreeSave(train_states["low_noise_transformer"])
-    items["high_noise_transformer_state"] = ocp.args.PyTreeSave(train_states["high_noise_transformer"])
+    items["low_noise_transformer_state"] = ocp.args.StandardSave(train_states["low_noise_transformer"])
+    items["high_noise_transformer_state"] = ocp.args.StandardSave(train_states["high_noise_transformer"])
 
     # Save the checkpoint
     self.checkpoint_manager.save(train_step, args=ocp.args.Composite(**items))

src/maxdiffusion/checkpointing/wan_checkpointer_i2v_2p1.py

@@ -44,31 +44,26 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
     metadatas = self.checkpoint_manager.item_metadata(step)
     state = metadatas.wan_state
 
+    # Manually constructs jax.ShapeDtypeStruct with a specific sharding.
+    # This avoids device mesh validation (as in ocp.utils.to_shape_dtype_struct)
+    # allowing for sharding with a different mesh than the one used during
+    # saving.
     def add_sharding_to_struct(leaf_struct, sharding):
-      struct = ocp.utils.to_shape_dtype_struct(leaf_struct)
-      if hasattr(struct, "shape") and hasattr(struct, "dtype"):
-        return jax.ShapeDtypeStruct(shape=struct.shape, dtype=struct.dtype, sharding=sharding)
-      return struct
+      if hasattr(leaf_struct, "shape") and hasattr(leaf_struct, "dtype"):
+        return jax.ShapeDtypeStruct(shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding)
+      return leaf_struct
 
     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=jax.Array),
-            abstract_train_state_with_sharding,
-        )
-    )
-
     max_logging.log("Restoring WAN checkpoint")
     restored_checkpoint = self.checkpoint_manager.restore(
-        directory=epath.Path(self.config.checkpoint_dir),
         step=step,
         args=ocp.args.Composite(
-            wan_state=params_restore,
             wan_config=ocp.args.JsonRestore(),
+            wan_state=ocp.args.StandardRestore(abstract_train_state_with_sharding),
         ),
     )
     max_logging.log(f"restored checkpoint {restored_checkpoint.keys()}")
@@ -106,7 +101,7 @@ def config_to_json(model_or_config):
         "wan_config": ocp.args.JsonSave(config_to_json(pipeline.transformer)),
     }
 
-    items["wan_state"] = ocp.args.PyTreeSave(train_states)
+    items["wan_state"] = ocp.args.StandardSave(train_states)
 
     # Save the checkpoint
     self.checkpoint_manager.save(train_step, args=ocp.args.Composite(**items))

src/maxdiffusion/checkpointing/wan_checkpointer_i2v_2p2.py

@@ -17,6 +17,7 @@
 import json
 import jax
 import numpy as np
+from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
 from typing import Optional, Tuple
 from ..pipelines.wan.wan_pipeline_i2v_2p2 import WanPipelineI2V_2_2
 from .. import max_logging
@@ -35,39 +36,45 @@ 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)
 
     # Handle low_noise_transformer
     low_noise_transformer_metadata = metadatas.low_noise_transformer_state
-    abstract_tree_structure_low_params = jax.tree_util.tree_map(
-        ocp.utils.to_shape_dtype_struct, low_noise_transformer_metadata
-    )
-    low_params_restore = ocp.args.PyTreeRestore(
-        restore_args=jax.tree.map(
-            lambda _: ocp.RestoreArgs(restore_type=np.ndarray),
-            abstract_tree_structure_low_params,
-        )
-    )
+
+    # Manually constructs jax.ShapeDtypeStruct with a specific sharding.
+    # This avoids device mesh validation (as in ocp.utils.to_shape_dtype_struct)
+    # allowing for sharding with a different mesh than the one used during
+    # saving.
+    def add_sharding_to_struct(leaf_struct, sharding):
+      if hasattr(leaf_struct, "shape") and hasattr(leaf_struct, "dtype"):
+        return jax.ShapeDtypeStruct(shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding)
+      return leaf_struct
+
+    target_shardings = jax.tree_util.tree_map(lambda x: replicated_sharding, low_noise_transformer_metadata)
+    with mesh:
+      abstract_tree_structure_low_params = jax.tree_util.tree_map(
+          add_sharding_to_struct, low_noise_transformer_metadata, target_shardings
+      )
 
     # Handle high_noise_transformer
     high_noise_transformer_metadata = metadatas.high_noise_transformer_state
-    abstract_tree_structure_high_params = jax.tree_util.tree_map(
-        ocp.utils.to_shape_dtype_struct, high_noise_transformer_metadata
-    )
-    high_params_restore = ocp.args.PyTreeRestore(
-        restore_args=jax.tree.map(
-            lambda _: ocp.RestoreArgs(restore_type=np.ndarray),
-            abstract_tree_structure_high_params,
-        )
-    )
+    target_shardings = jax.tree_util.tree_map(lambda x: replicated_sharding, high_noise_transformer_metadata)
+    with mesh:
+      abstract_tree_structure_high_params = jax.tree_util.tree_map(
+          add_sharding_to_struct, high_noise_transformer_metadata, target_shardings
+      )
 
     max_logging.log("Restoring WAN 2.2 checkpoint")
     restored_checkpoint = self.checkpoint_manager.restore(
-        directory=epath.Path(self.config.checkpoint_dir),
         step=step,
         args=ocp.args.Composite(
-            low_noise_transformer_state=low_params_restore,
-            high_noise_transformer_state=high_params_restore,
+            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(),
         ),
     )
@@ -119,8 +126,8 @@ def config_to_json(model_or_config):
         "wan_config": ocp.args.JsonSave(config_to_json(pipeline.low_noise_transformer)),
     }
 
-    items["low_noise_transformer_state"] = ocp.args.PyTreeSave(train_states["low_noise_transformer"])
-    items["high_noise_transformer_state"] = ocp.args.PyTreeSave(train_states["high_noise_transformer"])
+    items["low_noise_transformer_state"] = ocp.args.StandardSave(train_states["low_noise_transformer"])
+    items["high_noise_transformer_state"] = ocp.args.StandardSave(train_states["high_noise_transformer"])
 
     # Save the checkpoint
     self.checkpoint_manager.save(train_step, args=ocp.args.Composite(**items))

src/maxdiffusion/checkpointing/wan_vace_checkpointer_2_1.py

@@ -42,11 +42,14 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
     metadatas = self.checkpoint_manager.item_metadata(step)
     state = metadatas.wan_state
 
+    # Manually constructs jax.ShapeDtypeStruct with a specific sharding.
+    # This avoids device mesh validation (as in ocp.utils.to_shape_dtype_struct)
+    # allowing for sharding with a different mesh than the one used during
+    # saving.
     def add_sharding_to_struct(leaf_struct, sharding):
-      struct = ocp.utils.to_shape_dtype_struct(leaf_struct)
-      if hasattr(struct, "shape") and hasattr(struct, "dtype"):
-        return jax.ShapeDtypeStruct(shape=struct.shape, dtype=struct.dtype, sharding=sharding)
-      return struct
+      if hasattr(leaf_struct, "shape") and hasattr(leaf_struct, "dtype"):
+        return jax.ShapeDtypeStruct(shape=leaf_struct.shape, dtype=leaf_struct.dtype, sharding=sharding)
+      return leaf_struct
 
     target_shardings = jax.tree_util.tree_map(lambda x: replicated_sharding, state)