Fix optimizer state restoration and multi-host sharding in WAN trainer

src/maxdiffusion/generate_ltx_video.py

@@ -213,9 +213,7 @@ def run(config):
 
   for prompt_idx, current_prompt in enumerate(prompts):
     prompt_word_count = len(current_prompt.split())
-    enhance_prompt = (
-        prompt_enhancement_words_threshold > 0 and prompt_word_count < prompt_enhancement_words_threshold
-    )
+    enhance_prompt = prompt_enhancement_words_threshold > 0 and prompt_word_count < prompt_enhancement_words_threshold
 
     s0 = time.perf_counter()
     images = pipeline(

src/maxdiffusion/max_utils.py

@@ -503,9 +503,13 @@ def device_put_replicated(x, sharding):
 
   to also shard an array based on sharding.
   """
+  if isinstance(x, jax.Array) and hasattr(x, "sharding") and x.sharding == sharding:
+    return x
+
   arr = getattr(x, "value", x)
+  arr_np = np.asarray(arr)
   shd = getattr(sharding, "value", sharding)
-  res = jax.make_array_from_callback(arr.shape, shd, lambda index: arr[index])
+  res = jax.make_array_from_callback(arr_np.shape, shd, lambda index: arr_np[index])
   if hasattr(x, "set_value"):
     x.set_value(res)
     return x

src/maxdiffusion/trainers/base_wan_trainer.py

@@ -20,6 +20,7 @@
 import os
 import pprint
 import threading
+from typing import Any
 from flax import nnx
 from flax.linen import partitioning as nn_partitioning
 from flax.training import train_state
@@ -49,6 +50,54 @@ def _to_array(x):
   return x
 
 
+def _normalize_path(path: tuple[Any, ...]) -> tuple[Any, ...]:
+  """Normalizes a PyTree path tuple for optimizer state matching."""
+  normalized = []
+  for p in path:
+    if hasattr(p, "key"):
+      val = p.key
+    elif hasattr(p, "idx"):
+      val = p.idx
+    elif hasattr(p, "name"):
+      val = p.name
+    else:
+      val = p
+
+    if isinstance(val, str) and val.isdigit():
+      val = int(val)
+    normalized.append(val)
+
+  if normalized and normalized[-1] == "value":
+    normalized = normalized[:-1]
+
+  return tuple(normalized)
+
+
+def restore_optimizer_state_by_path(
+    target_opt_shardings: Any,
+    restored_opt_state: Any,
+) -> Any:
+  """Restores optimizer state by matching exact parameter paths and applying sharding."""
+  raw_restored_map, _ = jax.tree_util.tree_flatten_with_path(restored_opt_state)
+
+  normalized_restored_map = {}
+  for raw_path, val in raw_restored_map:
+    norm_path = _normalize_path(raw_path)
+    normalized_restored_map[norm_path] = val
+
+  def _restore_leaf(path, target_sharding):
+    norm_path = _normalize_path(path)
+    if norm_path not in normalized_restored_map:
+      raise KeyError(f"Optimizer checkpoint missing path: {norm_path}.")
+    val = normalized_restored_map[norm_path]
+    return max_utils.device_put_replicated(val, target_sharding)
+
+  return jax.tree_util.tree_map_with_path(
+      _restore_leaf,
+      target_opt_shardings,
+  )
+
+
 def generate_sample(config, pipeline, filename_prefix):
   """
   Generates a video to validate training did not corrupt the model
@@ -178,9 +227,11 @@ def start_training(self):
     with nn_partitioning.axis_rules(self.config.logical_axis_rules):
       pipeline, opt_state, step = self.checkpointer.load_checkpoint()
     restore_args = {}
-    if opt_state and step:
-      restore_args = {"opt_state": opt_state, "step": step}
+    if opt_state is not None:
+      restore_args["opt_state"] = opt_state
       del opt_state
+    if step is not None:
+      restore_args["step"] = step
     if self.config.enable_ssim:
       # Generate a sample before training to compare against generated sample after training.
       pretrained_video_path = generate_sample(self.config, pipeline, filename_prefix="pre-training-")
@@ -261,21 +312,27 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
       state = TrainState.create(
           apply_fn=graphdef.apply, params=params, tx=optimizer, graphdef=graphdef, rest_of_state=rest_of_state
       )
-      if restore_args:
-        step = restore_args.get("step", 0)
-        max_logging.log(f"Restoring optimizer and resuming from step {step}")
-        state.replace(opt_state=restore_args.get("opt_state"), step=restore_args.get("step", 0))
-        del restore_args["opt_state"]
-        del optimizer
       state = jax.tree.map(_to_array, state)
       state_spec = nnx.get_partition_spec(state)
       state = jax.lax.with_sharding_constraint(state, state_spec)
       state_shardings = nnx.get_named_sharding(state, mesh)
-      if jax.process_index() == 0 and restore_args:
-        max_logging.log("--- Optimizer State Sharding Spec (opt_state) ---")
-        pretty_string = pprint.pformat(state_spec.opt_state, indent=4, width=60)
-        max_logging.log(pretty_string)
-        max_logging.log("------------------------------------------------")
+      if restore_args.get("step") is not None:
+        step = restore_args.get("step", 0)
+        state = state.replace(step=step)
+      if restore_args.get("opt_state") is not None:
+        max_logging.log("Restoring optimizer")
+        resharded_opt_state = restore_optimizer_state_by_path(
+            target_opt_shardings=state_shardings.opt_state,
+            restored_opt_state=restore_args["opt_state"],
+        )
+        state = state.replace(opt_state=resharded_opt_state)
+        del restore_args["opt_state"]
+        if jax.process_index() == 0:
+          max_logging.log("--- Optimizer State Sharding Spec (opt_state) ---")
+          pretty_string = pprint.pformat(state_spec.opt_state, indent=4, width=60)
+          max_logging.log(pretty_string)
+          max_logging.log("------------------------------------------------")
+
     if self.config.hardware != "gpu":
       max_utils.delete_pytree(params)
     data_shardings = self.get_data_shardings(mesh)
@@ -312,7 +369,7 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
         first_profiling_step + self.config.profiler_steps - 1, first_profiling_step, self.config.max_train_steps - 1
     )
     if restore_args.get("step", 0):
-      max_logging.log(f"Resuming training from step {step}")
+      max_logging.log(f"Resuming training from step {restore_args.get('step', 0)}")
     start_step = restore_args.get("step", 0)
     per_device_tflops, _, _ = BaseWanTrainer.calculate_tflops(pipeline)
     scheduler_state = pipeline.scheduler_state
@@ -367,7 +424,10 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
         writer.flush()
       if self.config.save_final_checkpoint:
         max_logging.log(f"Saving final checkpoint for step {step}")
-        self.checkpointer.save_checkpoint(self.config.max_train_steps - 1, pipeline, state.params)
+        if self.config.save_optimizer:
+          self.checkpointer.save_checkpoint(self.config.max_train_steps - 1, pipeline, state)
+        else:
+          self.checkpointer.save_checkpoint(self.config.max_train_steps - 1, pipeline, state.params)
         self.checkpointer.checkpoint_manager.wait_until_finished()
       # load new state for trained transformer
       pipeline.transformer = nnx.merge(state.graphdef, state.params, state.rest_of_state)