From b8a3fba7e3ccac4293c355fcb1168c84bbc04d8a Mon Sep 17 00:00:00 2001 From: Nina Shvetsova Date: Wed, 16 Sep 2026 16:15:10 +0000 Subject: [PATCH] Fix optimizer state restoration and multi-host sharding in WAN trainer - Fix no-op state.replace(...) when restoring opt_state and step by reassigning state. - Add restore_optimizer_state_by_path to normalize and align Orbax checkpoint paths with live Optax PyTree structure. - Update max_utils.device_put_replicated to convert to NumPy host arrays before callback slicing to avoid multi-host dispatch failures, and fast-path matching shardings. --- src/maxdiffusion/generate_ltx_video.py | 4 +- src/maxdiffusion/max_utils.py | 6 +- .../tests/input_pipeline_interface_test.py | 5 +- src/maxdiffusion/trainers/base_wan_trainer.py | 137 +++++++++++++++--- 4 files changed, 128 insertions(+), 24 deletions(-) diff --git a/src/maxdiffusion/generate_ltx_video.py b/src/maxdiffusion/generate_ltx_video.py index bb771424f..a3fddcb59 100644 --- a/src/maxdiffusion/generate_ltx_video.py +++ b/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( diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index 4471e3bd1..171fb91ea 100644 --- a/src/maxdiffusion/max_utils.py +++ b/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 diff --git a/src/maxdiffusion/tests/input_pipeline_interface_test.py b/src/maxdiffusion/tests/input_pipeline_interface_test.py index 40a9a72c6..5e7de3417 100644 --- a/src/maxdiffusion/tests/input_pipeline_interface_test.py +++ b/src/maxdiffusion/tests/input_pipeline_interface_test.py @@ -21,6 +21,7 @@ import pathlib import shutil import subprocess +import tempfile import unittest import pytest from absl.testing import absltest @@ -75,7 +76,8 @@ def setUp(self): def test_make_dreambooth_train_iterator(self): instance_class_gcs_dir = "gs://maxdiffusion-github-runner-test-assets/datasets/dreambooth/instance_class" class_class_gcs_dir = "gs://maxdiffusion-github-runner-test-assets/datasets/dreambooth/class_class" - local_dir = "/tmp/" + self.temp_dir = tempfile.TemporaryDirectory() + local_dir = self.temp_dir.name instance_class_local_dir = max_utils.download_blobs(instance_class_gcs_dir, local_dir) class_class_local_dir = max_utils.download_blobs(class_class_gcs_dir, local_dir) @@ -136,6 +138,7 @@ def test_make_dreambooth_train_iterator(self): cleanup(instance_class_local_dir) cleanup(class_class_local_dir) + self.temp_dir.cleanup() @pytest.mark.skip( "This test is deprecated and will be removed in a future version. Reason: stable diffusion 2 base is no longer in HuggingFace" diff --git a/src/maxdiffusion/trainers/base_wan_trainer.py b/src/maxdiffusion/trainers/base_wan_trainer.py index 1d2c543dc..90157bc6d 100644 --- a/src/maxdiffusion/trainers/base_wan_trainer.py +++ b/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-") @@ -256,26 +307,76 @@ def eval(self, mesh, eval_rng_key, step, p_eval_step, state, scheduler_state, wr def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data_iterator, restore_args: dict = {}): mesh = pipeline.mesh graphdef, params, rest_of_state = nnx.split(pipeline.transformer, nnx.Param, ...) + start_step = 0 with mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): 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("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 restore_args.get("step") is not None: + step = restore_args["step"] + # A checkpoint labelled N is written *after* iteration N completes, + # and the loop starts at step 0, so it already reflects N+1 applied + # updates. Resume at N+1 so that iteration N is not re-run, and so + # that TrainState.step stays aligned with the optimizer's internal + # step count, which actually drives the LR schedule. + state = state.replace(step=step + 1) + start_step = step + 1 + max_logging.log(f"Resuming training from step {start_step}") + else: + # Weights-only: no optimizer state means no schedule position. + # Resetting to 0 keeps the applied LR and the logged LR consistent. + if "step" in restore_args: + max_logging.log( + "Weights-only checkpoint detected (no opt_state); resetting step" + " to 0. The LR schedule will restart from the beginning. If you" + " intended to resume training, re-save with optimizer state." + ) + weights_step = restore_args.pop("step") + + # The loop restarts at 0 but checkpoint_dir still holds steps up to + # weights_step, and Orbax will skip saving the checkpoints at those + # steps. Therefore, we redirect the three output paths that + # pyconfig derives from run_name, then rebuild the CheckpointManager. + previous_run_name = self.config.run_name + previous_checkpoint_dir = self.config.checkpoint_dir + new_run_name = f"{previous_run_name}_from_weights_{weights_step}" + keys = self.config.get_keys() + keys["run_name"] = new_run_name + for key, subdir in ( + ("tensorboard_dir", "tensorboard"), + ("checkpoint_dir", "checkpoints"), + ): + keys[key] = os.path.join(self.config.output_dir, new_run_name, subdir, "") + max_logging.log( + f"Redirecting outputs from run_name {previous_run_name} to" + f" {new_run_name}: {previous_checkpoint_dir} already holds" + f" checkpoints up to step {weights_step}, which this run would" + " otherwise fail to save over." + ) + max_logging.log(f" checkpoint_dir: {self.config.checkpoint_dir}") + max_logging.log(f" tensorboard_dir: {self.config.tensorboard_dir}") + self.checkpointer.checkpoint_manager.close() + self.checkpointer = self._get_checkpointer() + if self.config.hardware != "gpu": max_utils.delete_pytree(params) data_shardings = self.get_data_shardings(mesh) @@ -301,8 +402,6 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data rng = jax.random.key(self.config.seed) rng, eval_rng_key = jax.random.split(rng) - start_step = 0 - last_step_completion = datetime.datetime.now() local_metrics_file = open(self.config.metrics_file, "a", encoding="utf8") if self.config.metrics_file else None running_gcs_metrics = [] if self.config.gcs_metrics else None first_profiling_step = self.config.skip_first_n_steps_for_profiler @@ -311,9 +410,6 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data last_profiling_step = np.clip( 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}") - start_step = restore_args.get("step", 0) per_device_tflops, _, _ = BaseWanTrainer.calculate_tflops(pipeline) scheduler_state = pipeline.scheduler_state example_batch = load_next_batch(train_data_iterator, None, self.config) @@ -367,7 +463,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)