From aefa0f00d86dcb2bcdcc5cf5f401a22a3f176956 Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Tue, 1 Sep 2026 13:06:02 +0530 Subject: [PATCH 1/6] Implement Wan 2.2 training pipeline with joint timestep routing --- docker_build_dependency_image.sh | 3 +- maxdiffusion_dependencies.Dockerfile | 2 +- setup.sh | 4 +- src/maxdiffusion/configs/base_wan_27b.yml | 10 +- src/maxdiffusion/max_utils.py | 23 +- src/maxdiffusion/models/wan/wan_utils.py | 28 +- src/maxdiffusion/pyconfig.py | 8 +- src/maxdiffusion/train_utils.py | 9 +- src/maxdiffusion/train_wan_2_2.py | 51 ++ src/maxdiffusion/trainers/wan_trainer_2_2.py | 558 +++++++++++++++++++ 10 files changed, 659 insertions(+), 37 deletions(-) create mode 100644 src/maxdiffusion/train_wan_2_2.py create mode 100644 src/maxdiffusion/trainers/wan_trainer_2_2.py diff --git a/docker_build_dependency_image.sh b/docker_build_dependency_image.sh index 5c2c0f8d2..aa26ce9bf 100644 --- a/docker_build_dependency_image.sh +++ b/docker_build_dependency_image.sh @@ -27,8 +27,7 @@ set -e export LOCAL_IMAGE_NAME=maxdiffusion_base_image -# Use Docker BuildKit so we can cache pip packages. -export DOCKER_BUILDKIT=1 +export DOCKER_BUILDKIT=0 echo "Starting to build your docker image. This will take a few minutes but the image can be reused as you iterate." diff --git a/maxdiffusion_dependencies.Dockerfile b/maxdiffusion_dependencies.Dockerfile index 9a9598271..cf025876a 100644 --- a/maxdiffusion_dependencies.Dockerfile +++ b/maxdiffusion_dependencies.Dockerfile @@ -48,7 +48,7 @@ COPY . . RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION" -RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} +RUN bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} # Cleanup RUN rm -rf /root/.cache/pip \ No newline at end of file diff --git a/setup.sh b/setup.sh index 3f1141888..501be7926 100644 --- a/setup.sh +++ b/setup.sh @@ -161,8 +161,8 @@ elif [[ $MODE == "nightly" ]]; then python3 -m uv pip install --pre -U jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html # Install jaxlib-nightly python3 -m uv pip install --pre -U jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html - # Install libtpu-nightly - python3 -m uv pip install --pre -U libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html + # Install libtpu nightly (package is now named libtpu, not libtpu-nightly) + python3 -m uv pip install --pre -U libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html fi echo "Installing nightly tensorboard plugin profile" python3 -m uv pip install tbp-nightly --upgrade diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml index 185b01277..40bd935fe 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -22,7 +22,7 @@ write_metrics: True timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written. write_timing_metrics: True -gcs_metrics: False +gcs_metrics: True # If true save config to GCS in {base_output_directory}/{run_name}/ save_config_to_gcs: False log_period: 100 @@ -31,6 +31,7 @@ pretrained_model_name_or_path: 'Wan-AI/Wan2.2-T2V-A14B-Diffusers' model_name: wan2.2 model_type: 'T2V' + # Overrides the transformer from pretrained_model_name_or_path wan_transformer_pretrained_model_name_or_path: '' @@ -225,9 +226,9 @@ vae_logical_axis_rules: [ ['heads', null], ['norm', null], ['conv_batch', 'redundant'], - ['out_channels', 'vae_spatial'], - ['conv_out', 'vae_spatial'], - ['conv_in', 'vae_spatial'], + ['out_channels', null], + ['conv_out', null], + ['conv_in', null], ] data_sharding: [['data', 'fsdp', 'context', 'tensor']] @@ -317,6 +318,7 @@ output_dir: 'sdxl-model-finetuned' per_device_batch_size: 1.0 # If global_batch_size % jax.device_count is not 0, use FSDP sharding. global_batch_size: 0 +disable_training_weights: False # For creating tfrecords from dataset tfrecords_dir: '' diff --git a/src/maxdiffusion/max_utils.py b/src/maxdiffusion/max_utils.py index b7ee9a5d1..ce4950706 100644 --- a/src/maxdiffusion/max_utils.py +++ b/src/maxdiffusion/max_utils.py @@ -279,21 +279,16 @@ def parse_gcs_bucket_and_prefix(destination_gcs_name): def download_blobs(source_gcs_folder, local_destination): - """Downloads a folder to a local location""" + """Downloads a folder to a local location using gcloud storage to avoid SSL segfaults on large models.""" + import subprocess + Path(local_destination).mkdir(parents=True, exist_ok=True) + subprocess.run(["gcloud", "storage", "cp", "-r", source_gcs_folder, local_destination], check=True) + + # gcloud storage cp -r gs://bucket/prefix /dest creates /dest/prefix bucket_name, prefix_name = parse_gcs_bucket_and_prefix(source_gcs_folder) - storage_client = storage.Client() - bucket = storage_client.get_bucket(bucket_name) - blobs = bucket.list_blobs(prefix=prefix_name) - for blob in blobs: - file_split = blob.name.split("/") - directory = os.path.join(local_destination, "/".join(file_split[0:-1])) - Path(directory).mkdir(parents=True, exist_ok=True) - if len(file_split[-1]) <= 0: - continue - download_to_filename = os.path.join(directory, file_split[-1]) - if not os.path.isfile(download_to_filename): - blob.download_to_filename(download_to_filename) - return os.path.join(local_destination, prefix_name) + # The last component of prefix_name is the folder name created + folder_name = prefix_name.rstrip("/").split("/")[-1] + return os.path.join(local_destination, folder_name) def upload_blob(destination_gcs_name, source_file_name): diff --git a/src/maxdiffusion/models/wan/wan_utils.py b/src/maxdiffusion/models/wan/wan_utils.py index 4c9053219..bf5e1440b 100644 --- a/src/maxdiffusion/models/wan/wan_utils.py +++ b/src/maxdiffusion/models/wan/wan_utils.py @@ -498,20 +498,32 @@ def convert_chunk(ckpt_shard_path, chunk_keys): # across the ~12 shard files. norm_added_q is explicitly ignored by the # diffusers implementation. chunk_size = 16 - tasks = [] + max_logging.log( + f"Load and port {pretrained_model_name_or_path} {subfolder}: {len(model_files)} shards sequentially to save disk" + ) for model_file in model_files: + tasks = [] ckpt_shard_path = resolve_shard_path(model_file) with safe_open(ckpt_shard_path, framework="pt") as f: shard_keys = [k for k in f.keys() if "norm_added_q" not in k] for i in range(0, len(shard_keys), chunk_size): tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size])) - max_logging.log( - f"Load and port {pretrained_model_name_or_path} {subfolder}: {len(model_files)} shards, {len(tasks)} chunks" - ) - with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: - futures = [executor.submit(convert_chunk, path, keys) for path, keys in tasks] - for future in concurrent.futures.as_completed(futures): - future.result() # re-raise conversion errors + + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: + futures = [executor.submit(convert_chunk, path, keys) for path, keys in tasks] + for future in concurrent.futures.as_completed(futures): + future.result() # re-raise conversion errors + + if not local_files: + try: + real_path = os.path.realpath(ckpt_shard_path) + if os.path.exists(real_path): + os.remove(real_path) + if os.path.exists(ckpt_shard_path): + os.remove(ckpt_shard_path) + except Exception as e: + max_logging.log(f"Warning: could not delete shard {ckpt_shard_path}: {e}") + validate_flax_state_dict(eval_shapes, flax_state_dict) if converted_cache_dir and not os.path.isdir(converted_cache_dir): diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index d1121ca3f..bc29ba23e 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/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): @@ -284,11 +284,11 @@ 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"] 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"], "/dev/shm") 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"], "/dev/shm") 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"], "/dev/shm") 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"] diff --git a/src/maxdiffusion/train_utils.py b/src/maxdiffusion/train_utils.py index e82174e92..f1df6f13e 100644 --- a/src/maxdiffusion/train_utils.py +++ b/src/maxdiffusion/train_utils.py @@ -101,7 +101,10 @@ def _tensorboard_writer_worker(writer, config): metrics, step = data if jax.process_index() == 0: for metric_name in metrics.get("scalar", []): - writer.add_scalar(metric_name, np.array(metrics["scalar"][metric_name]), step) + val = np.array(metrics["scalar"][metric_name]) + if not np.isnan(val): + metric_step = metrics.get("steps", {}).get(metric_name, step) + writer.add_scalar(metric_name, val, metric_step) for metric_name in metrics.get("scalars", []): writer.add_scalars(metric_name, metrics["scalars"][metric_name], step) @@ -150,7 +153,9 @@ def write_metrics_to_tensorboard(writer, metrics, step, config): ) if jax.process_index() == 0: for metric_name in metrics.get("scalar", []): - writer.add_scalar(metric_name, np.array(metrics["scalar"][metric_name]), step) + val = np.array(metrics["scalar"][metric_name]) + if not np.isnan(val): + writer.add_scalar(metric_name, val, step) for metric_name in metrics.get("scalars", []): writer.add_scalars(metric_name, metrics["scalars"][metric_name], step) diff --git a/src/maxdiffusion/train_wan_2_2.py b/src/maxdiffusion/train_wan_2_2.py new file mode 100644 index 000000000..f7ec17556 --- /dev/null +++ b/src/maxdiffusion/train_wan_2_2.py @@ -0,0 +1,51 @@ +""" +Copyright 2025 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from typing import Sequence + +import jax +from absl import app +from maxdiffusion import max_logging, pyconfig, max_utils +from maxdiffusion.train_utils import ( + validate_train_config, + transformer_engine_context, +) +import flax + + +def train(config): + from maxdiffusion.trainers.wan_trainer_2_2 import WanTrainer2_2 + + trainer = WanTrainer2_2(config) + trainer.start_training() + + +def main(argv: Sequence[str]) -> None: + pyconfig.initialize(argv, validate_training=True) + config = pyconfig.config + max_utils.ensure_machinelearning_job_runs(pyconfig.config) + validate_train_config(config) + max_logging.log(f"Found {jax.device_count()} devices.") + try: + flax.config.update("flax_always_shard_variable", False) + except LookupError: + pass + train(config) + + +if __name__ == "__main__": + with transformer_engine_context(): + app.run(main) diff --git a/src/maxdiffusion/trainers/wan_trainer_2_2.py b/src/maxdiffusion/trainers/wan_trainer_2_2.py new file mode 100644 index 000000000..b28f9c6c9 --- /dev/null +++ b/src/maxdiffusion/trainers/wan_trainer_2_2.py @@ -0,0 +1,558 @@ +""" +Copyright 2025 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import functools +import datetime +import os +import pprint +import threading +from concurrent.futures import ThreadPoolExecutor +import numpy as np + +from flax import nnx +from flax.linen import partitioning as nn_partitioning +import jax.numpy as jnp +import jax +from jax.sharding import PartitionSpec as P +from jax.experimental import multihost_utils +import jaxopt +import tensorflow as tf + +from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2 +from maxdiffusion.input_pipeline.input_pipeline_interface import make_data_iterator +from maxdiffusion.trainers.base_wan_trainer import BaseWanTrainer, TrainState, _to_array, print_ssim +from maxdiffusion import max_logging, max_utils, train_utils +from maxdiffusion.train_utils import (_metrics_queue, _tensorboard_writer_worker, load_next_batch) +from maxdiffusion.generate_wan import inference_generate_video +from maxdiffusion.generate_wan import run as generate_wan +from maxdiffusion.pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2 + + +class WanTrainer2_2(BaseWanTrainer): + + def _get_checkpointer(self): + return WanCheckpointer2_2(config=self.config) + + def get_data_shardings(self, mesh): + data_sharding = jax.sharding.NamedSharding(mesh, P(*self.config.data_sharding)) + data_sharding = {"latents": data_sharding, "encoder_hidden_states": data_sharding} + return data_sharding + + def get_eval_data_shardings(self, mesh): + data_sharding = jax.sharding.NamedSharding(mesh, P(*self.config.data_sharding)) + data_sharding = {"latents": data_sharding, "encoder_hidden_states": data_sharding, "timesteps": data_sharding} + return data_sharding + + def load_dataset(self, mesh, pipeline=None, is_training=True): + config = self.config + if config.dataset_type == "synthetic": + return make_data_iterator( + config, + jax.process_index(), + jax.process_count(), + mesh, + config.global_batch_size_to_load, + pipeline=pipeline, + is_training=is_training, + ) + + if config.dataset_type != "tfrecord" and not config.cache_latents_text_encoder_outputs: + raise ValueError( + "Wan 2.2 training only supports config.dataset_type set to tfrecords and config.cache_latents_text_encoder_outputs set to True" + ) + feature_description = { + "latents": tf.io.FixedLenFeature([], tf.string), + "encoder_hidden_states": tf.io.FixedLenFeature([], tf.string), + } + + if not is_training: + feature_description["timesteps"] = tf.io.FixedLenFeature([], tf.int64) + + def prepare_sample_train(features): + latents = tf.io.parse_tensor(features["latents"], out_type=tf.float32) + encoder_hidden_states = tf.io.parse_tensor(features["encoder_hidden_states"], out_type=tf.float32) + return {"latents": latents, "encoder_hidden_states": encoder_hidden_states} + + def prepare_sample_eval(features): + latents = tf.io.parse_tensor(features["latents"], out_type=tf.float32) + encoder_hidden_states = tf.io.parse_tensor(features["encoder_hidden_states"], out_type=tf.float32) + timesteps = features["timesteps"] + return {"latents": latents, "encoder_hidden_states": encoder_hidden_states, "timesteps": timesteps} + + data_iterator = make_data_iterator( + config, + jax.process_index(), + jax.process_count(), + mesh, + config.global_batch_size_to_load, + feature_description=feature_description, + prepare_sample_fn=prepare_sample_train if is_training else prepare_sample_eval, + is_training=is_training, + ) + return data_iterator + + def calculate_tflops(self, pipeline): + maxdiffusion_config = pipeline.config + height = pipeline.config.height + width = pipeline.config.width + num_frames = pipeline.config.num_frames + + transformer_config = pipeline.low_noise_transformer.config + num_layers = transformer_config.num_layers + heads = pipeline.low_noise_transformer.config.num_attention_heads + head_dim = pipeline.low_noise_transformer.config.attention_head_dim + ffn_dim = transformer_config.ffn_dim + seq_len = int(((height / 8) * (width / 8) * ((num_frames - 1) // pipeline.vae_scale_factor_temporal + 1)) / 4) + text_encoder_dim = 512 + + self_attn_qkv_proj_flops = 3 * (2 * seq_len * (heads * head_dim) ** 2) + self_attn_qk_v_flops = 2 * (2 * seq_len**2 * (heads * head_dim)) + cross_attn_kv_proj_flops = 3 * (2 * text_encoder_dim * (heads * head_dim) ** 2) + cross_attn_q_proj_flops = 1 * (2 * seq_len * (heads * head_dim) ** 2) + cross_attention_qk_v_flops = 2 * (2 * seq_len * text_encoder_dim * (heads * head_dim)) + attn_output_proj_flops = 2 * (2 * seq_len * (heads * head_dim) ** 2) + + total_attn_flops = ( + self_attn_qkv_proj_flops + + self_attn_qk_v_flops + + cross_attn_kv_proj_flops + + cross_attn_q_proj_flops + + cross_attention_qk_v_flops + + attn_output_proj_flops + ) + + ffn_flops = 2 * (2 * seq_len * (heads * head_dim) * ffn_dim) + flops_per_block = total_attn_flops + ffn_flops + total_transformer_flops = flops_per_block * num_layers + tflops = maxdiffusion_config.per_device_batch_size * total_transformer_flops / 1e12 + train_tflops = 3 * tflops + + max_logging.log(f"Calculated TFLOPs per pass: {train_tflops:.4f}") + return train_tflops, total_attn_flops, seq_len + + def get_train_step(self, pipeline, mesh, state_shardings, data_shardings): + return jax.jit( + functools.partial(train_step_2_2, scheduler=pipeline.scheduler, config=self.config), + in_shardings=(state_shardings["low_noise"], state_shardings["high_noise"], data_shardings, None, None), + out_shardings=(state_shardings["low_noise"], state_shardings["high_noise"], None, None, None), + donate_argnums=(0, 1), + ) + + def get_eval_step(self, pipeline, mesh, state_shardings, eval_data_shardings): + return jax.jit( + functools.partial(eval_step_2_2, scheduler=pipeline.scheduler, config=self.config), + in_shardings=(state_shardings["low_noise"], state_shardings["high_noise"], eval_data_shardings, None, None), + out_shardings=(None, None), + ) + + def generate_sample(self, config, pipeline, filename_prefix): + if not hasattr(pipeline, "vae"): + wan_vae, vae_cache = WanPipeline2_2.load_vae( + pipeline.mesh.devices, pipeline.mesh, nnx.Rngs(jax.random.key(config.seed)), config + ) + pipeline.vae = wan_vae + pipeline.vae_cache = vae_cache + return generate_wan(config, pipeline, filename_prefix) + + def start_training(self): + with nn_partitioning.axis_rules(self.config.logical_axis_rules): + pipeline, opt_state_dict, step = self.checkpointer.load_checkpoint() + + restore_args = {} + if opt_state_dict and step: + restore_args = {"opt_state": opt_state_dict, "step": step} + del opt_state_dict + + if self.config.enable_ssim: + pretrained_video_path = self.generate_sample(self.config, pipeline, filename_prefix="pre-training-") + + if self.config.eval_every == -1 or (not self.config.enable_generate_video_for_eval): + if hasattr(pipeline, "vae"): + del pipeline.vae + if hasattr(pipeline, "vae_cache"): + del pipeline.vae_cache + + mesh = pipeline.mesh + train_data_iterator = self.load_dataset(mesh, pipeline=pipeline, is_training=True) + + scheduler, scheduler_state = self.create_scheduler() + pipeline.scheduler = scheduler + pipeline.scheduler_state = scheduler_state + + optimizer_low, learning_rate_scheduler = self.checkpointer._create_optimizer( + pipeline.low_noise_transformer, self.config, self.config.learning_rate + ) + optimizer_high, _ = self.checkpointer._create_optimizer( + pipeline.high_noise_transformer, self.config, self.config.learning_rate + ) + + pipeline = self.training_loop_2_2( + pipeline, optimizer_low, optimizer_high, learning_rate_scheduler, train_data_iterator, restore_args + ) + + if self.config.enable_ssim: + posttrained_video_path = self.generate_sample(self.config, pipeline, filename_prefix="post-training-") + print_ssim(pretrained_video_path, posttrained_video_path) + + def training_loop_2_2(self, pipeline, optimizer_low, optimizer_high, learning_rate_scheduler, train_data_iterator, restore_args: dict = {}): + mesh = pipeline.mesh + graphdef_low, params_low, rest_of_state_low = nnx.split(pipeline.low_noise_transformer, nnx.Param, ...) + graphdef_high, params_high, rest_of_state_high = nnx.split(pipeline.high_noise_transformer, nnx.Param, ...) + + with mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + state_low = TrainState.create( + apply_fn=graphdef_low.apply, params=params_low, tx=optimizer_low, graphdef=graphdef_low, rest_of_state=rest_of_state_low + ) + state_high = TrainState.create( + apply_fn=graphdef_high.apply, params=params_high, tx=optimizer_high, graphdef=graphdef_high, rest_of_state=rest_of_state_high + ) + + if restore_args: + step = restore_args.get("step", 0) + max_logging.log(f"Restoring optimizer and resuming from step {step}") + opt_state_dict = restore_args.get("opt_state", {}) + if "low_noise_transformer" in opt_state_dict: + state_low = state_low.replace(opt_state=opt_state_dict["low_noise_transformer"]) + if "high_noise_transformer" in opt_state_dict: + state_high = state_high.replace(opt_state=opt_state_dict["high_noise_transformer"]) + state_low = state_low.replace(step=step) + state_high = state_high.replace(step=step) + + state_low = jax.tree.map(_to_array, state_low) + state_high = jax.tree.map(_to_array, state_high) + + state_spec_low = nnx.get_partition_spec(state_low) + state_spec_high = nnx.get_partition_spec(state_high) + + state_low = jax.lax.with_sharding_constraint(state_low, state_spec_low) + state_high = jax.lax.with_sharding_constraint(state_high, state_spec_high) + + state_shardings_low = nnx.get_named_sharding(state_low, mesh) + state_shardings_high = nnx.get_named_sharding(state_high, mesh) + + state_shardings = {"low_noise": state_shardings_low, "high_noise": state_shardings_high} + + if self.config.hardware != "gpu": + max_utils.delete_pytree(params_low) + max_utils.delete_pytree(params_high) + + data_shardings = self.get_data_shardings(mesh) + eval_data_shardings = self.get_eval_data_shardings(mesh) + + writer = max_utils.initialize_summary_writer(self.config) + writer_thread = threading.Thread(target=_tensorboard_writer_worker, args=(writer, self.config), daemon=True) + writer_thread.start() + + num_model_parameters = max_utils.calculate_num_params_from_pytree(state_low.params) + max_utils.calculate_num_params_from_pytree(state_high.params) + max_utils.add_text_to_summary_writer("number_model_parameters", str(num_model_parameters), writer) + max_utils.add_config_to_summary_writer(self.config, writer) + + if jax.process_index() == 0: + max_logging.log("***** Running training *****") + max_logging.log(f" Total optimization steps = {self.config.max_train_steps}") + + p_train_step = self.get_train_step(pipeline, mesh, state_shardings, data_shardings) + p_eval_step = self.get_eval_step(pipeline, mesh, state_shardings, eval_data_shardings) + + rng = jax.random.key(self.config.seed) + rng, eval_rng_key = jax.random.split(rng) + start_step = restore_args.get("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 + + per_device_tflops, _, _ = self.calculate_tflops(pipeline) + scheduler_state = pipeline.scheduler_state + example_batch = load_next_batch(train_data_iterator, None, self.config) + + with ThreadPoolExecutor(max_workers=1) as executor: + step_high = 0 + step_low = 0 + for step in np.arange(start_step, self.config.max_train_steps): + start_step_time = datetime.datetime.now() + next_batch_future = executor.submit(load_next_batch, train_data_iterator, example_batch, self.config) + + with jax.profiler.StepTraceAnnotation("train", step_num=step), pipeline.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + state_low, state_high, scheduler_state, train_metric, rng = p_train_step( + state_low, state_high, example_batch, rng, scheduler_state + ) + train_metric["scalar"]["learning/loss"].block_until_ready() + last_step_completion = datetime.datetime.now() + + train_utils.record_scalar_metrics( + train_metric, last_step_completion - start_step_time, per_device_tflops, learning_rate_scheduler(step) + ) + + loss_high_val = np.array(train_metric["scalar"]["learning/loss_high"]) + loss_low_val = np.array(train_metric["scalar"]["learning/loss_low"]) + if "steps" not in train_metric: + train_metric["steps"] = {} + if not np.isnan(loss_high_val): + step_high += 1 + train_metric["steps"]["learning/loss_high"] = step_high + if not np.isnan(loss_low_val): + step_low += 1 + train_metric["steps"]["learning/loss_low"] = step_low + train_metric["steps"]["learning/loss_low_fine"] = step_low + train_metric["steps"]["learning/loss_low_mid"] = step_low + train_metric["steps"]["learning/loss_low_coarse"] = step_low + + if self.config.write_metrics: + train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config) + + if self.config.eval_every > 0 and (step + 1) % self.config.eval_every == 0: + if self.config.enable_generate_video_for_eval: + pipeline.low_noise_transformer = nnx.merge(state_low.graphdef, state_low.params, state_low.rest_of_state) + pipeline.high_noise_transformer = nnx.merge(state_high.graphdef, state_high.params, state_high.rest_of_state) + inference_generate_video(self.config, pipeline, filename_prefix=f"{step+1}-train_steps-") + self.eval_2_2(mesh, eval_rng_key, step, p_eval_step, state_low, state_high, scheduler_state, writer) + + example_batch = next_batch_future.result() + if step != 0 and self.config.checkpoint_every != -1 and step % self.config.checkpoint_every == 0: + max_logging.log(f"Saving checkpoint for step {step}") + train_states = { + "low_noise_transformer": state_low if self.config.save_optimizer else state_low.params, + "high_noise_transformer": state_high if self.config.save_optimizer else state_high.params, + } + self.checkpointer.save_checkpoint(step, pipeline, train_states) + + _metrics_queue.put(None) + writer_thread.join() + if writer: + writer.flush() + if self.config.save_final_checkpoint: + train_states = { + "low_noise_transformer": state_low.params, + "high_noise_transformer": state_high.params, + } + self.checkpointer.save_checkpoint(self.config.max_train_steps - 1, pipeline, train_states) + self.checkpointer.checkpoint_manager.wait_until_finished() + + pipeline.low_noise_transformer = nnx.merge(state_low.graphdef, state_low.params, state_low.rest_of_state) + pipeline.high_noise_transformer = nnx.merge(state_high.graphdef, state_high.params, state_high.rest_of_state) + return pipeline + + def eval_2_2(self, mesh, eval_rng_key, step, p_eval_step, state_low, state_high, scheduler_state, writer): + eval_data_iterator = self.load_dataset(mesh, is_training=False) + eval_rng = eval_rng_key + eval_losses_by_timestep = {} + while True: + try: + eval_start_time = datetime.datetime.now() + eval_batch = load_next_batch(eval_data_iterator, None, self.config) + with mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): + metrics, eval_rng = p_eval_step(state_low, state_high, eval_batch, eval_rng, scheduler_state) + metrics["scalar"]["learning/eval_loss"].block_until_ready() + + losses = metrics["scalar"]["learning/eval_loss"] + timesteps = eval_batch["timesteps"] + gathered_losses = multihost_utils.process_allgather(losses, tiled=True) + gathered_losses = jax.device_get(gathered_losses) + gathered_timesteps = multihost_utils.process_allgather(timesteps, tiled=True) + gathered_timesteps = jax.device_get(gathered_timesteps) + + if jax.process_index() == 0: + for t, l in zip(gathered_timesteps.flatten(), gathered_losses.flatten()): + timestep = int(t) + if timestep not in eval_losses_by_timestep: + eval_losses_by_timestep[timestep] = [] + eval_losses_by_timestep[timestep].append(l) + eval_end_time = datetime.datetime.now() + eval_duration = eval_end_time - eval_start_time + max_logging.log(f"Eval time: {eval_duration.total_seconds():.2f} seconds.") + except StopIteration: + break + + if eval_losses_by_timestep and jax.process_index() == 0: + mean_per_timestep = [] + for timestep, losses in sorted(eval_losses_by_timestep.items()): + losses = jnp.array(losses) + losses = losses[: min(self.config.eval_max_number_of_samples_in_bucket, len(losses))] + mean_loss = jnp.mean(losses) + mean_per_timestep.append(mean_loss) + final_eval_loss = jnp.mean(jnp.array(mean_per_timestep)) + max_logging.log(f"Step {step}, Final Average Eval loss: {final_eval_loss:.4f}") + if writer: + writer.add_scalar("learning/eval_loss", final_eval_loss, step) + + +def train_step_2_2(state_low, state_high, data, rng, scheduler_state, scheduler, config): + _, new_rng, timestep_rng, dropout_rng, cond_rng, noise_rng = jax.random.split(rng, num=6) + + for k, v in data.items(): + data[k] = v[: config.global_batch_size_to_train_on, :] + + bsz = data["latents"].shape[0] + num_train_timesteps = scheduler.config.num_train_timesteps + boundary = int(config.boundary_ratio * num_train_timesteps) + + is_high_noise = jax.random.uniform(cond_rng) > config.boundary_ratio + + def compute_loss_high(high_params): + t_float = jax.random.beta(timestep_rng, 5.0, 2.0, shape=(bsz,)) + timesteps = boundary + (t_float * (num_train_timesteps - boundary)).astype(jnp.int32) + timesteps = jnp.clip(timesteps, boundary, num_train_timesteps - 1) + model = nnx.merge(state_high.graphdef, high_params, state_high.rest_of_state) + latents = data["latents"].astype(config.weights_dtype) + encoder_hidden_states = data["encoder_hidden_states"].astype(config.weights_dtype) + noise = jax.random.normal(key=noise_rng, shape=latents.shape, dtype=latents.dtype) + noisy_latents, training_target, training_weight = scheduler.apply_flow_match(noise, latents, timesteps) + + model_pred = model( + hidden_states=noisy_latents, + timestep=timesteps, + encoder_hidden_states=encoder_hidden_states, + deterministic=False, + rngs=nnx.Rngs(dropout=dropout_rng), + ) + loss = (training_target - model_pred) ** 2 + if not config.disable_training_weights: + training_weight = jnp.expand_dims(training_weight, axis=(1, 2, 3, 4)) + loss = loss * training_weight + return jnp.mean(loss) + + def compute_loss_low(low_params): + t_float = jax.random.beta(timestep_rng, 2.0, 5.0, shape=(bsz,)) + timesteps = (t_float * boundary).astype(jnp.int32) + timesteps = jnp.clip(timesteps, 0, boundary - 1) + model = nnx.merge(state_low.graphdef, low_params, state_low.rest_of_state) + latents = data["latents"].astype(config.weights_dtype) + encoder_hidden_states = data["encoder_hidden_states"].astype(config.weights_dtype) + noise = jax.random.normal(key=noise_rng, shape=latents.shape, dtype=latents.dtype) + noisy_latents, training_target, training_weight = scheduler.apply_flow_match(noise, latents, timesteps) + + model_pred = model( + hidden_states=noisy_latents, + timestep=timesteps, + encoder_hidden_states=encoder_hidden_states, + deterministic=False, + rngs=nnx.Rngs(dropout=dropout_rng), + ) + loss = (training_target - model_pred) ** 2 + if not config.disable_training_weights: + training_weight = jnp.expand_dims(training_weight, axis=(1, 2, 3, 4)) + loss = loss * training_weight + + loss_mean = jnp.mean(loss) + avg_t = jnp.mean(timesteps) + loss_fine = jnp.where(avg_t < 200, loss_mean, jnp.nan) + loss_mid = jnp.where((avg_t >= 200) & (avg_t < 500), loss_mean, jnp.nan) + loss_coarse = jnp.where(avg_t >= 500, loss_mean, jnp.nan) + + return loss_mean, (loss_fine, loss_mid, loss_coarse) + + def true_fn(params_tuple): + high_params, low_params = params_tuple + loss, high_grads = nnx.value_and_grad(compute_loss_high)(high_params) + low_grads = jax.tree.map(jnp.zeros_like, low_params) + nan_val = jnp.array(jnp.nan, dtype=loss.dtype) + return loss, nan_val, nan_val, nan_val, nan_val, low_grads, high_grads + + def false_fn(params_tuple): + high_params, low_params = params_tuple + (loss, (loss_fine, loss_mid, loss_coarse)), low_grads = nnx.value_and_grad(compute_loss_low, has_aux=True)(low_params) + high_grads = jax.tree.map(jnp.zeros_like, high_params) + nan_val = jnp.array(jnp.nan, dtype=loss.dtype) + return nan_val, loss, loss_fine, loss_mid, loss_coarse, low_grads, high_grads + + loss_high, loss_low, loss_fine, loss_mid, loss_coarse, grads_low, grads_high = jax.lax.cond( + is_high_noise, + true_fn, + false_fn, + operand=(state_high.params, state_low.params) + ) + + max_grad_norm_low = jaxopt.tree_util.tree_l2_norm(grads_low) + max_grad_norm_high = jaxopt.tree_util.tree_l2_norm(grads_high) + + max_abs_grad_low = jax.tree_util.tree_reduce(lambda max_val, arr: jnp.maximum(max_val, jnp.max(jnp.abs(arr))), grads_low, initializer=-1.0) + max_abs_grad_high = jax.tree_util.tree_reduce(lambda max_val, arr: jnp.maximum(max_val, jnp.max(jnp.abs(arr))), grads_high, initializer=-1.0) + + metrics = { + "scalar": { + "learning/loss": jnp.nan_to_num(loss_high, nan=0.0) + jnp.nan_to_num(loss_low, nan=0.0), + "learning/loss_low": loss_low, + "learning/loss_high": loss_high, + "learning/loss_low_fine": loss_fine, + "learning/loss_low_mid": loss_mid, + "learning/loss_low_coarse": loss_coarse, + "learning/max_grad_norm_low": max_grad_norm_low, + "learning/max_grad_norm_high": max_grad_norm_high, + "learning/max_abs_grad_low": max_abs_grad_low, + "learning/max_abs_grad_high": max_abs_grad_high, + }, + "scalars": {}, + } + + new_state_low = state_low.apply_gradients(grads=grads_low) + new_state_high = state_high.apply_gradients(grads=grads_high) + return new_state_low, new_state_high, scheduler_state, metrics, new_rng + + +def eval_step_2_2(state_low, state_high, data, rng, scheduler_state, scheduler, config): + num_train_timesteps = scheduler.config.num_train_timesteps + boundary = int(config.boundary_ratio * num_train_timesteps) + + def loss_fn_high(params, latents, encoder_hidden_states, timesteps, rng): + model = nnx.merge(state_high.graphdef, params, state_high.rest_of_state) + noise = jax.random.normal(key=rng, shape=latents.shape, dtype=latents.dtype) + noisy_latents, training_target, training_weight = scheduler.apply_flow_match(noise, latents, timesteps) + model_pred = model( + hidden_states=noisy_latents, timestep=timesteps, encoder_hidden_states=encoder_hidden_states, deterministic=True + ) + loss = (training_target - model_pred) ** 2 + if not config.disable_training_weights: + training_weight = jnp.expand_dims(training_weight, axis=(1, 2, 3, 4)) + loss = loss * training_weight + return loss.reshape(loss.shape[0], -1).mean(axis=1) + + def loss_fn_low(params, latents, encoder_hidden_states, timesteps, rng): + model = nnx.merge(state_low.graphdef, params, state_low.rest_of_state) + noise = jax.random.normal(key=rng, shape=latents.shape, dtype=latents.dtype) + noisy_latents, training_target, training_weight = scheduler.apply_flow_match(noise, latents, timesteps) + model_pred = model( + hidden_states=noisy_latents, timestep=timesteps, encoder_hidden_states=encoder_hidden_states, deterministic=True + ) + loss = (training_target - model_pred) ** 2 + if not config.disable_training_weights: + training_weight = jnp.expand_dims(training_weight, axis=(1, 2, 3, 4)) + loss = loss * training_weight + return loss.reshape(loss.shape[0], -1).mean(axis=1) + + bs = len(data["latents"]) + single_batch_size = config.global_batch_size_to_train_on + losses = jnp.zeros(bs) + + for i in range(0, bs, single_batch_size): + start = i + end = min(i + single_batch_size, bs) + latents = data["latents"][start:end, :].astype(config.weights_dtype) + encoder_hidden_states = data["encoder_hidden_states"][start:end, :].astype(config.weights_dtype) + timesteps = data["timesteps"][start:end].astype("int64") + _, new_rng = jax.random.split(rng, num=2) + + is_high = timesteps[0] >= boundary + + def true_fn(_): + return loss_fn_high(state_high.params, latents, encoder_hidden_states, timesteps, new_rng) + def false_fn(_): + return loss_fn_low(state_low.params, latents, encoder_hidden_states, timesteps, new_rng) + + loss = jax.lax.cond(is_high, true_fn, false_fn, operand=None) + losses = losses.at[start:end].set(loss) + + metrics = {"scalar": {"learning/eval_loss": losses}} + return metrics, new_rng From 2f45b41ab57ef60ad2a2924b92af97489ebda376 Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Tue, 1 Sep 2026 16:10:37 +0530 Subject: [PATCH 2/6] Address code-assist bot review comments --- src/maxdiffusion/models/wan/wan_utils.py | 18 +++++++++--------- src/maxdiffusion/trainers/wan_trainer_2_2.py | 20 ++++++++------------ 2 files changed, 17 insertions(+), 21 deletions(-) diff --git a/src/maxdiffusion/models/wan/wan_utils.py b/src/maxdiffusion/models/wan/wan_utils.py index bf5e1440b..93f0e72d3 100644 --- a/src/maxdiffusion/models/wan/wan_utils.py +++ b/src/maxdiffusion/models/wan/wan_utils.py @@ -501,15 +501,15 @@ def convert_chunk(ckpt_shard_path, chunk_keys): max_logging.log( f"Load and port {pretrained_model_name_or_path} {subfolder}: {len(model_files)} shards sequentially to save disk" ) - for model_file in model_files: - tasks = [] - ckpt_shard_path = resolve_shard_path(model_file) - with safe_open(ckpt_shard_path, framework="pt") as f: - shard_keys = [k for k in f.keys() if "norm_added_q" not in k] - for i in range(0, len(shard_keys), chunk_size): - tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size])) - - with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor: + for model_file in model_files: + tasks = [] + ckpt_shard_path = resolve_shard_path(model_file) + with safe_open(ckpt_shard_path, framework="pt") as f: + shard_keys = [k for k in f.keys() if "norm_added_q" not in k] + for i in range(0, len(shard_keys), chunk_size): + tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size])) + futures = [executor.submit(convert_chunk, path, keys) for path, keys in tasks] for future in concurrent.futures.as_completed(futures): future.result() # re-raise conversion errors diff --git a/src/maxdiffusion/trainers/wan_trainer_2_2.py b/src/maxdiffusion/trainers/wan_trainer_2_2.py index b28f9c6c9..5ab71596a 100644 --- a/src/maxdiffusion/trainers/wan_trainer_2_2.py +++ b/src/maxdiffusion/trainers/wan_trainer_2_2.py @@ -69,7 +69,7 @@ def load_dataset(self, mesh, pipeline=None, is_training=True): is_training=is_training, ) - if config.dataset_type != "tfrecord" and not config.cache_latents_text_encoder_outputs: + if config.dataset_type != "tfrecord" or not config.cache_latents_text_encoder_outputs: raise ValueError( "Wan 2.2 training only supports config.dataset_type set to tfrecords and config.cache_latents_text_encoder_outputs set to True" ) @@ -207,7 +207,9 @@ def start_training(self): posttrained_video_path = self.generate_sample(self.config, pipeline, filename_prefix="post-training-") print_ssim(pretrained_video_path, posttrained_video_path) - def training_loop_2_2(self, pipeline, optimizer_low, optimizer_high, learning_rate_scheduler, train_data_iterator, restore_args: dict = {}): + def training_loop_2_2(self, pipeline, optimizer_low, optimizer_high, learning_rate_scheduler, train_data_iterator, restore_args: dict | None = None): + if restore_args is None: + restore_args = {} mesh = pipeline.mesh graphdef_low, params_low, rest_of_state_low = nnx.split(pipeline.low_noise_transformer, nnx.Param, ...) graphdef_high, params_high, rest_of_state_high = nnx.split(pipeline.high_noise_transformer, nnx.Param, ...) @@ -392,8 +394,7 @@ def eval_2_2(self, mesh, eval_rng_key, step, p_eval_step, state_low, state_high, def train_step_2_2(state_low, state_high, data, rng, scheduler_state, scheduler, config): _, new_rng, timestep_rng, dropout_rng, cond_rng, noise_rng = jax.random.split(rng, num=6) - for k, v in data.items(): - data[k] = v[: config.global_batch_size_to_train_on, :] + data = {k: v[: config.global_batch_size_to_train_on, :] for k, v in data.items()} bsz = data["latents"].shape[0] num_train_timesteps = scheduler.config.num_train_timesteps @@ -544,14 +545,9 @@ def loss_fn_low(params, latents, encoder_hidden_states, timesteps, rng): timesteps = data["timesteps"][start:end].astype("int64") _, new_rng = jax.random.split(rng, num=2) - is_high = timesteps[0] >= boundary - - def true_fn(_): - return loss_fn_high(state_high.params, latents, encoder_hidden_states, timesteps, new_rng) - def false_fn(_): - return loss_fn_low(state_low.params, latents, encoder_hidden_states, timesteps, new_rng) - - loss = jax.lax.cond(is_high, true_fn, false_fn, operand=None) + loss_high = loss_fn_high(state_high.params, latents, encoder_hidden_states, timesteps, new_rng) + loss_low = loss_fn_low(state_low.params, latents, encoder_hidden_states, timesteps, new_rng) + loss = jnp.where(timesteps >= boundary, loss_high, loss_low) losses = losses.at[start:end].set(loss) metrics = {"scalar": {"learning/eval_loss": losses}} From 9c447945e03ba1d42f2f7539e0d211a4a371446d Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Tue, 1 Sep 2026 17:08:39 +0530 Subject: [PATCH 3/6] Fix ruff linting errors (unused imports and trailing whitespaces) --- src/maxdiffusion/models/wan/wan_utils.py | 4 +-- src/maxdiffusion/trainers/wan_trainer_2_2.py | 38 ++++++++++---------- 2 files changed, 20 insertions(+), 22 deletions(-) diff --git a/src/maxdiffusion/models/wan/wan_utils.py b/src/maxdiffusion/models/wan/wan_utils.py index 93f0e72d3..3db43fb74 100644 --- a/src/maxdiffusion/models/wan/wan_utils.py +++ b/src/maxdiffusion/models/wan/wan_utils.py @@ -509,11 +509,11 @@ def convert_chunk(ckpt_shard_path, chunk_keys): shard_keys = [k for k in f.keys() if "norm_added_q" not in k] for i in range(0, len(shard_keys), chunk_size): tasks.append((ckpt_shard_path, shard_keys[i : i + chunk_size])) - + futures = [executor.submit(convert_chunk, path, keys) for path, keys in tasks] for future in concurrent.futures.as_completed(futures): future.result() # re-raise conversion errors - + if not local_files: try: real_path = os.path.realpath(ckpt_shard_path) diff --git a/src/maxdiffusion/trainers/wan_trainer_2_2.py b/src/maxdiffusion/trainers/wan_trainer_2_2.py index 5ab71596a..4a40719d3 100644 --- a/src/maxdiffusion/trainers/wan_trainer_2_2.py +++ b/src/maxdiffusion/trainers/wan_trainer_2_2.py @@ -16,8 +16,6 @@ import functools import datetime -import os -import pprint import threading from concurrent.futures import ThreadPoolExecutor import numpy as np @@ -170,7 +168,7 @@ def generate_sample(self, config, pipeline, filename_prefix): def start_training(self): with nn_partitioning.axis_rules(self.config.logical_axis_rules): pipeline, opt_state_dict, step = self.checkpointer.load_checkpoint() - + restore_args = {} if opt_state_dict and step: restore_args = {"opt_state": opt_state_dict, "step": step} @@ -250,7 +248,7 @@ def training_loop_2_2(self, pipeline, optimizer_low, optimizer_high, learning_ra if self.config.hardware != "gpu": max_utils.delete_pytree(params_low) max_utils.delete_pytree(params_high) - + data_shardings = self.get_data_shardings(mesh) eval_data_shardings = self.get_eval_data_shardings(mesh) @@ -297,7 +295,7 @@ def training_loop_2_2(self, pipeline, optimizer_low, optimizer_high, learning_ra train_utils.record_scalar_metrics( train_metric, last_step_completion - start_step_time, per_device_tflops, learning_rate_scheduler(step) ) - + loss_high_val = np.array(train_metric["scalar"]["learning/loss_high"]) loss_low_val = np.array(train_metric["scalar"]["learning/loss_low"]) if "steps" not in train_metric: @@ -358,14 +356,14 @@ def eval_2_2(self, mesh, eval_rng_key, step, p_eval_step, state_low, state_high, with mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): metrics, eval_rng = p_eval_step(state_low, state_high, eval_batch, eval_rng, scheduler_state) metrics["scalar"]["learning/eval_loss"].block_until_ready() - + losses = metrics["scalar"]["learning/eval_loss"] timesteps = eval_batch["timesteps"] gathered_losses = multihost_utils.process_allgather(losses, tiled=True) gathered_losses = jax.device_get(gathered_losses) gathered_timesteps = multihost_utils.process_allgather(timesteps, tiled=True) gathered_timesteps = jax.device_get(gathered_timesteps) - + if jax.process_index() == 0: for t, l in zip(gathered_timesteps.flatten(), gathered_losses.flatten()): timestep = int(t) @@ -377,7 +375,7 @@ def eval_2_2(self, mesh, eval_rng_key, step, p_eval_step, state_low, state_high, max_logging.log(f"Eval time: {eval_duration.total_seconds():.2f} seconds.") except StopIteration: break - + if eval_losses_by_timestep and jax.process_index() == 0: mean_per_timestep = [] for timestep, losses in sorted(eval_losses_by_timestep.items()): @@ -393,9 +391,9 @@ def eval_2_2(self, mesh, eval_rng_key, step, p_eval_step, state_low, state_high, def train_step_2_2(state_low, state_high, data, rng, scheduler_state, scheduler, config): _, new_rng, timestep_rng, dropout_rng, cond_rng, noise_rng = jax.random.split(rng, num=6) - + data = {k: v[: config.global_batch_size_to_train_on, :] for k, v in data.items()} - + bsz = data["latents"].shape[0] num_train_timesteps = scheduler.config.num_train_timesteps boundary = int(config.boundary_ratio * num_train_timesteps) @@ -411,7 +409,7 @@ def compute_loss_high(high_params): encoder_hidden_states = data["encoder_hidden_states"].astype(config.weights_dtype) noise = jax.random.normal(key=noise_rng, shape=latents.shape, dtype=latents.dtype) noisy_latents, training_target, training_weight = scheduler.apply_flow_match(noise, latents, timesteps) - + model_pred = model( hidden_states=noisy_latents, timestep=timesteps, @@ -434,7 +432,7 @@ def compute_loss_low(low_params): encoder_hidden_states = data["encoder_hidden_states"].astype(config.weights_dtype) noise = jax.random.normal(key=noise_rng, shape=latents.shape, dtype=latents.dtype) noisy_latents, training_target, training_weight = scheduler.apply_flow_match(noise, latents, timesteps) - + model_pred = model( hidden_states=noisy_latents, timestep=timesteps, @@ -446,13 +444,13 @@ def compute_loss_low(low_params): if not config.disable_training_weights: training_weight = jnp.expand_dims(training_weight, axis=(1, 2, 3, 4)) loss = loss * training_weight - + loss_mean = jnp.mean(loss) avg_t = jnp.mean(timesteps) loss_fine = jnp.where(avg_t < 200, loss_mean, jnp.nan) loss_mid = jnp.where((avg_t >= 200) & (avg_t < 500), loss_mean, jnp.nan) loss_coarse = jnp.where(avg_t >= 500, loss_mean, jnp.nan) - + return loss_mean, (loss_fine, loss_mid, loss_coarse) def true_fn(params_tuple): @@ -470,15 +468,15 @@ def false_fn(params_tuple): return nan_val, loss, loss_fine, loss_mid, loss_coarse, low_grads, high_grads loss_high, loss_low, loss_fine, loss_mid, loss_coarse, grads_low, grads_high = jax.lax.cond( - is_high_noise, - true_fn, - false_fn, + is_high_noise, + true_fn, + false_fn, operand=(state_high.params, state_low.params) ) max_grad_norm_low = jaxopt.tree_util.tree_l2_norm(grads_low) max_grad_norm_high = jaxopt.tree_util.tree_l2_norm(grads_high) - + max_abs_grad_low = jax.tree_util.tree_reduce(lambda max_val, arr: jnp.maximum(max_val, jnp.max(jnp.abs(arr))), grads_low, initializer=-1.0) max_abs_grad_high = jax.tree_util.tree_reduce(lambda max_val, arr: jnp.maximum(max_val, jnp.max(jnp.abs(arr))), grads_high, initializer=-1.0) @@ -536,7 +534,7 @@ def loss_fn_low(params, latents, encoder_hidden_states, timesteps, rng): bs = len(data["latents"]) single_batch_size = config.global_batch_size_to_train_on losses = jnp.zeros(bs) - + for i in range(0, bs, single_batch_size): start = i end = min(i + single_batch_size, bs) @@ -544,7 +542,7 @@ def loss_fn_low(params, latents, encoder_hidden_states, timesteps, rng): encoder_hidden_states = data["encoder_hidden_states"][start:end, :].astype(config.weights_dtype) timesteps = data["timesteps"][start:end].astype("int64") _, new_rng = jax.random.split(rng, num=2) - + loss_high = loss_fn_high(state_high.params, latents, encoder_hidden_states, timesteps, new_rng) loss_low = loss_fn_low(state_low.params, latents, encoder_hidden_states, timesteps, new_rng) loss = jnp.where(timesteps >= boundary, loss_high, loss_low) From 0524c2eca39c0fad40d6b44bb84959f74cc97759 Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Wed, 2 Sep 2026 11:46:24 +0530 Subject: [PATCH 4/6] Move VAE sharding overrides to a dedicated training config (training_wan_27b.yml) --- src/maxdiffusion/configs/base_wan_27b.yml | 10 +- src/maxdiffusion/configs/training_wan_27b.yml | 469 ++++++++++++++++++ 2 files changed, 473 insertions(+), 6 deletions(-) create mode 100644 src/maxdiffusion/configs/training_wan_27b.yml diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml index 40bd935fe..185b01277 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -22,7 +22,7 @@ write_metrics: True timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written. write_timing_metrics: True -gcs_metrics: True +gcs_metrics: False # If true save config to GCS in {base_output_directory}/{run_name}/ save_config_to_gcs: False log_period: 100 @@ -31,7 +31,6 @@ pretrained_model_name_or_path: 'Wan-AI/Wan2.2-T2V-A14B-Diffusers' model_name: wan2.2 model_type: 'T2V' - # Overrides the transformer from pretrained_model_name_or_path wan_transformer_pretrained_model_name_or_path: '' @@ -226,9 +225,9 @@ vae_logical_axis_rules: [ ['heads', null], ['norm', null], ['conv_batch', 'redundant'], - ['out_channels', null], - ['conv_out', null], - ['conv_in', null], + ['out_channels', 'vae_spatial'], + ['conv_out', 'vae_spatial'], + ['conv_in', 'vae_spatial'], ] data_sharding: [['data', 'fsdp', 'context', 'tensor']] @@ -318,7 +317,6 @@ output_dir: 'sdxl-model-finetuned' per_device_batch_size: 1.0 # If global_batch_size % jax.device_count is not 0, use FSDP sharding. global_batch_size: 0 -disable_training_weights: False # For creating tfrecords from dataset tfrecords_dir: '' diff --git a/src/maxdiffusion/configs/training_wan_27b.yml b/src/maxdiffusion/configs/training_wan_27b.yml new file mode 100644 index 000000000..40bd935fe --- /dev/null +++ b/src/maxdiffusion/configs/training_wan_27b.yml @@ -0,0 +1,469 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# This sentinel is a reminder to choose a real run name. +run_name: '' + +metrics_file: "" # for testing, local file that stores scalar metrics. If empty, no metrics are written. +# If true save metrics such as loss and TFLOPS to GCS in {base_output_directory}/{run_name}/metrics/ +write_metrics: True + +timing_metrics_file: "" # for testing, local file that stores function timing metrics such as state creation, compilation. If empty, no metrics are written. +write_timing_metrics: True + +gcs_metrics: True +# If true save config to GCS in {base_output_directory}/{run_name}/ +save_config_to_gcs: False +log_period: 100 + +pretrained_model_name_or_path: 'Wan-AI/Wan2.2-T2V-A14B-Diffusers' +model_name: wan2.2 +model_type: 'T2V' + + +# Overrides the transformer from pretrained_model_name_or_path +wan_transformer_pretrained_model_name_or_path: '' + +unet_checkpoint: '' +revision: '' +# This will convert the weights to this dtype. +# When running inference on TPUv5e, use weights_dtype: 'bfloat16' +weights_dtype: 'bfloat16' +# This sets the layer's dtype in the model. Ex: nn.Dense(dtype=activations_dtype) +activations_dtype: 'bfloat16' +# The dtype for text_encoder model during load/compile +text_encoder_dtype: 'float32' + +# Whether to compile the text_encoder with torch.compile +compile_text_encoder: False + +# Maximum sequence length for the text encoder +max_sequence_length: 512 + +vae_weights_dtype: 'float32' +vae_dtype: 'float32' +scheduler_dtype: 'float32' + +# Replicates vae across devices instead of using the model's sharding annotations for sharding. +replicate_vae: False + +# Chunk size for VAE decode scan. Increase to improve decode time at the cost of memory. +vae_decode_chunk: 1 + +# Chunk size for VAE encode scan. (num_input_frames - 1) must be divisible by this value. +# Increase to improve encode time at the cost of memory. +vae_encode_chunk: 4 +vae_spatial: -1 + +# matmul and conv precision from https://jax.readthedocs.io/en/latest/jax.lax.html#jax.lax.Precision +# Options are "DEFAULT", "HIGH", "HIGHEST" +# fp32 activations and fp32 weights with HIGHEST will provide the best precision +# at the cost of time. +precision: "DEFAULT" +# Use jax.lax.scan for transformer layers +scan_layers: True +# Use jax.lax.scan for the diffusion loop (non-cache path only). +# Note: Enabling this will disable per-step profiling. +scan_diffusion_loop: False + +# if False state is not jitted and instead replicate is called. This is good for debugging on single host +# It must be True for multi-host. +jit_initializers: True + +# Set true to load weights from pytorch +from_pt: True +split_head_dim: True +attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, tokamax_ring_custom, ulysses, ulysses_custom, ulysses_ring, ulysses_ring_custom, ulysses_ring_custom_bidir +# +# Best 2D-ring / USP (Ulysses x ring) configs for WAN2.2-T2V-A14B (720x1280, 81 frames) +# Set attention=ulysses_ring_custom and ulysses_shards=U (ring degree R=CP/U): +# CP4 (v7x-8): ulysses_shards=2 (R=2), BQ=9472 +# CP8 (v7x-8): ulysses_shards=4 (R=2), BQ=9472 +# CP16 (v7x-16): ulysses_shards=8 (R=2), BQ=9472 +use_base2_exp: True +use_experimental_scheduler: True +# For attention=ulysses_ring, hidden Ulysses shard count; ring shards are context / this. +ulysses_shards: -1 +# Splits Ulysses all-to-all into head-group chunks. The last chunk carries any remainder. +# For communication-compute overlap to be effective, enable the following XLA flags: +# --xla_tpu_enable_async_all_to_all=true +# --xla_tpu_overlap_compute_collective_tc=true +# (Refer to README.md for the full recommended XLA_FLAGS list) +ulysses_attention_chunks: 1 +flash_min_seq_length: 4096 +dropout: 0.0 + +# If mask_padding_tokens is True, we pass in segment ids to splash attention to avoid attending to padding tokens. +# Else we do not pass in segment ids and on vpu bound hardware like trillium this is faster. +# However, when padding tokens are significant, this will lead to worse quality and should be set to True. +mask_padding_tokens: True +# Maxdiffusion has 2 types of attention sharding strategies: +# 1. attention_sharding_uniform = True : same sequence sharding rules applied for q in both (self and cross attention) +# 2. attention_sharding_uniform = False : Heads are sharded uniformly across devices for self attention while sequence is sharded +# in cross attention q. +attention_sharding_uniform: True + +flash_block_sizes: { + "block_q" : 512, + "block_kv_compute" : 512, + "block_kv" : 512, + "block_q_dkv" : 512, + "block_kv_dkv" : 512, + "block_kv_dkv_compute" : 512, + "block_q_dq" : 512, + "block_kv_dq" : 512, + "use_fused_bwd_kernel": False, +} +# Use on v6e +# flash_block_sizes: { +# "block_q" : 3024, +# "block_kv_compute" : 1024, +# "block_kv" : 2048, +# "block_q_dkv" : 3024, +# "block_kv_dkv" : 2048, +# "block_kv_dkv_compute" : 2048, +# "block_q_dq" : 3024, +# "block_kv_dq" : 2048 +# "use_fused_bwd_kernel": False, +# } +# GroupNorm groups +# Tile-size auto-tuning. When enable_tile_search: True, generate_wan runs a fast one-DiT-block +# grid search (maxdiffusion/utils/tile_size_grid_search.py) before inference and overwrites +# flash_block_sizes' block_q/block_kv/block_kv_compute with the winner. Default off (no-op). +enable_tile_search: False +tile_search_mode: 'smart' # 'smart' (VMEM-capped candidate ladders) | 'full' (2D sweep) +tile_search_iters: 10 +tile_search_out: '' # dir for the results CSV; '' -> print only + +norm_num_groups: 32 + +# train text_encoder - Currently not supported for SDXL +train_text_encoder: False +text_encoder_learning_rate: 4.25e-6 + +# https://arxiv.org/pdf/2305.08891.pdf +snr_gamma: -1.0 + +timestep_bias: { + # a value of later will increase the frequence of the model's final training steps. + # none, earlier, later, range + strategy: "none", + # multiplier for bias, a value of 2.0 will double the weight of the bias, 0.5 will halve it. + multiplier: 1.0, + # when using strategy=range, the beginning (inclusive) timestep to bias. + begin: 0, + # when using strategy=range, the final step (inclusive) to bias. + end: 1000, + # portion of timesteps to bias. + # 0.5 will bias one half of the timesteps. Value of strategy determines + # whether the biased portions are in the earlier or later timesteps. + portion: 0.25 +} + +# Override parameters from checkpoints's scheduler. +diffusion_scheduler_config: { + _class_name: 'FlaxEulerDiscreteScheduler', + prediction_type: 'epsilon', + rescale_zero_terminal_snr: False, + timestep_spacing: 'trailing' +} + +# Output directory +# Create a GCS bucket, e.g. my-maxtext-outputs and set this to "gs://my-maxtext-outputs/" +base_output_directory: "" + +# Hardware +hardware: 'tpu' # Supported hardware types are 'tpu', 'gpu' +skip_jax_distributed_system: False + +# Parallelism +mesh_axes: ['data', 'fsdp', 'context', 'tensor'] + +# batch : batch dimension of data and activations +# hidden : +# embed : attention qkv dense layer hidden dim named as embed +# heads : attention head dim = num_heads * head_dim +# length : attention sequence length +# temb_in : dense.shape[0] of resnet dense before conv +# out_c : dense.shape[1] of resnet dense before conv +# out_channels : conv.shape[-1] activation +# keep_1 : conv.shape[0] weight +# keep_2 : conv.shape[1] weight +# conv_in : conv.shape[2] weight +# conv_out : conv.shape[-1] weight +logical_axis_rules: [ + ['batch', ['data', 'fsdp']], + ['activation_batch', ['data', 'fsdp']], + ['activation_self_attn_heads', ['context', 'tensor']], + ['activation_cross_attn_q_length', ['context', 'tensor']], + ['activation_length', 'context'], + ['activation_heads', 'tensor'], + ['mlp','tensor'], + ['embed', ['context', 'fsdp']], + ['heads', 'tensor'], + ['norm', 'tensor'], + ['conv_batch', ['data', 'context', 'fsdp']], + ['out_channels', 'tensor'], + ['conv_out', 'context'], + ] +vae_logical_axis_rules: [ + ['activation_batch', 'redundant'], + ['activation_length', 'vae_spatial'], + ['activation_heads', null], + ['activation_kv_length', null], + ['embed', null], + ['heads', null], + ['norm', null], + ['conv_batch', 'redundant'], + ['out_channels', null], + ['conv_out', null], + ['conv_in', null], + ] +data_sharding: [['data', 'fsdp', 'context', 'tensor']] + +# One axis for each parallelism type may hold a placeholder (-1) +# value to auto-shard based on available slices and devices. +# By default, product of the DCN axes should equal number of slices +# and product of the ICI axes should equal number of devices per slice. +dcn_data_parallelism: 1 +dcn_fsdp_parallelism: 1 +dcn_context_parallelism: -1 # recommended DCN axis to be auto-sharded +dcn_tensor_parallelism: 1 +ici_data_parallelism: 1 +ici_fsdp_parallelism: 1 +ici_context_parallelism: -1 # recommended ICI axis to be auto-sharded +ici_tensor_parallelism: 1 + +allow_split_physical_axes: False + +# Dataset +# Replace with dataset path or train_data_dir. One has to be set. +dataset_name: 'diffusers/pokemon-gpt4-captions' +train_split: 'train' +dataset_type: 'tfrecord' +cache_latents_text_encoder_outputs: True +# cache_latents_text_encoder_outputs only apply to dataset_type="tf", +# only apply to small dataset that fits in memory +# prepare image latents and text encoder outputs +# Reduce memory consumption and reduce step time during training +# transformed dataset is saved at dataset_save_location +dataset_save_location: '' +load_tfrecord_cached: True +train_data_dir: '' +dataset_config_name: '' +jax_cache_dir: '' +# Directory for per-shape AOT serialized executables ('' = disabled). +aot_cache_dir: '' +# Directory for memoized torch->flax converted weights ('' = disabled). +converted_weights_dir: '' +hf_data_dir: '' +hf_train_files: '' +hf_access_token: '' +image_column: 'image' +caption_column: 'text' +resolution: 1024 +center_crop: False +random_flip: False +# If cache_latents_text_encoder_outputs is True +# the num_proc is set to 1 +tokenize_captions_num_proc: 4 +transform_images_num_proc: 4 +reuse_example_batch: False +enable_data_shuffling: True + +# Defines the type of gradient checkpoint to enable. +# NONE - means no gradient checkpoint +# FULL - means full gradient checkpoint, whenever possible (minimum memory usage) +# MATMUL_WITHOUT_BATCH - means gradient checkpoint for every linear/matmul operation, +# except for ones that involve batch dimension - that means that all attention and projection +# layers will have gradient checkpoint, but not the backward with respect to the parameters. +# OFFLOAD_MATMUL_WITHOUT_BATCH - same as MATMUL_WITHOUT_BATCH but offload instead of recomputing. +# CUSTOM - set names to offload and save. +remat_policy: "NONE" +# For CUSTOM policy set below, current annotations are for: attn_output, query_proj, key_proj, value_proj +# xq_out, xk_out, ffn_activation +names_which_can_be_saved: [] +names_which_can_be_offloaded: [] + +# checkpoint every number of samples, -1 means don't checkpoint. +checkpoint_every: -1 +checkpoint_dir: "" +# Directory to cache pretrained weights as an orbax checkpoint for fast inference loads. +# On first run (slow, diffusers load), weights are saved here automatically. +# On subsequent runs, weights are loaded from here instead (~10x faster). +pretrained_orbax_dir: "" +# enables one replica to read the ckpt then broadcast to the rest +enable_single_replica_ckpt_restoring: False + +# Training loop +learning_rate: 1.e-5 +scale_lr: False +max_train_samples: -1 +# max_train_steps takes priority over num_train_epochs. +max_train_steps: 1500 +num_train_epochs: 1 +seed: 0 +output_dir: 'sdxl-model-finetuned' +per_device_batch_size: 1.0 +# If global_batch_size % jax.device_count is not 0, use FSDP sharding. +global_batch_size: 0 +disable_training_weights: False + +# For creating tfrecords from dataset +tfrecords_dir: '' +no_records_per_shard: 0 +enable_eval_timesteps: False +timesteps_list: [125, 250, 375, 500, 625, 750, 875] +num_eval_samples: 420 + +warmup_steps_fraction: 0.1 +learning_rate_schedule_steps: -1 # By default the length of the schedule is set to the number of steps. +save_optimizer: False + +# However you may choose a longer schedule (learning_rate_schedule_steps > steps), in which case the training will end before +# dropping fully down. Or you may choose a shorter schedule, where the unspecified steps will have a learning rate of 0. + +# AdamW optimizer parameters +adam_b1: 0.9 # Exponential decay rate to track the first moment of past gradients. +adam_b2: 0.999 # Exponential decay rate to track the second moment of past gradients. +adam_eps: 1.e-8 # A small constant applied to denominator outside of the square root. +adam_weight_decay: 0.0 # AdamW Weight decay +opt_enable_grad_clipping: False +max_grad_value: 1.0 +opt_enable_grad_global_norm_clipping: False +max_grad_norm: 1.0 + +enable_profiler: False +# Skip first n steps for profiling, to omit things like compilation and to give +# the iteration time a chance to stabilize. +skip_first_n_steps_for_profiler: 5 +profiler_steps: 10 + +# Enable JAX named scopes for detailed profiling and debugging +# When enabled, adds named scopes around key operations in transformer and attention layers +enable_jax_named_scopes: False + +# Generation parameters +prompt: "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." +prompt_2: "A cat and a dog baking a cake together in a kitchen. The cat is carefully measuring flour, while the dog is stirring the batter with a wooden spoon. The kitchen is cozy, with sunlight streaming through the window." +negative_prompt: "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards" +do_classifier_free_guidance: True +height: 720 +width: 1280 +num_frames: 81 +# Official Wan2.2 T2V A14B sampling shift (wan_t2v_A14B.py: sample_shift=12.0). This sets where the +# high->low transformer boundary falls in the step schedule (~step 26 of 40 at boundary_ratio 0.875), +# which must match the schedule the seeded mag_ratios_base was calibrated at. (A smaller shift moves +# the boundary earlier, which both degrades the base sample and misaligns MagCache's skip schedule.) +flow_shift: 12.0 + +# Reference for below guidance scale and boundary values: https://github.com/Wan-Video/Wan2.2/blob/main/wan/configs/wan_t2v_A14B.py +# guidance scale factor for low noise transformer +guidance_scale_low: 3.0 + +# guidance scale factor for high noise transformer +guidance_scale_high: 4.0 + +# The timestep threshold. If `t` is at or above this value, +# the `high_noise_model` is considered as the required model. +# timestep to switch between low noise and high noise transformer +boundary_ratio: 0.875 + +# Diffusion CFG cache (FasterCache-style) +use_cfg_cache: False + +# Batch positive and negative prompts in text encoder to save compute. +use_batched_text_encoder: False + +use_kv_cache: False +# SenCache: Sensitivity-Aware Caching (arXiv:2602.24208) — skip forward pass +# when predicted output change (based on accumulated latent/timestep drift) is small +use_sen_cache: False + +# MagCache (https://github.com/Zehong-Ma/MagCache) — skip transformer blocks when +# the accumulated magnitude-ratio error stays below `magcache_thresh`, reusing the +# cached block residual. `magcache_K` caps consecutive skips; `retention_ratio` is +# the fraction at the start of each phase that always computes in full. +use_magcache: False +magcache_thresh: 0.04 +magcache_K: 2 +retention_ratio: 0.2 +# Calibrated average magnitude ratios, interleaved [cond, uncond, ...], length +# num_inference_steps*2 (= 80 for 40 steps). Single curve spanning both phases; +# the dip near the middle marks the high->low boundary. Seeded from the official +# WAN 2.2 T2V ratios — recalibrate for this setup to tune the speedup/quality. +mag_ratios_base: [1.0, 1.0, 1.00124, 1.00155, 0.99822, 0.99851, 0.99696, 0.99687, 0.99703, 0.99732, 0.9966, 0.99679, 0.99602, 0.99658, 0.99578, 0.99664, 0.99484, 0.9949, 0.99633, 0.996, 0.99659, 0.99683, 0.99534, 0.99549, 0.99584, 0.99577, 0.99681, 0.99694, 0.99563, 0.99554, 0.9944, 0.99473, 0.99594, 0.9964, 0.99466, 0.99461, 0.99453, 0.99481, 0.99389, 0.99365, 0.99391, 0.99406, 0.99354, 0.99361, 0.99283, 0.99278, 0.99268, 0.99263, 0.99057, 0.99091, 0.99125, 0.99126, 0.65523, 0.65252, 0.98808, 0.98852, 0.98765, 0.98736, 0.9851, 0.98535, 0.98311, 0.98339, 0.9805, 0.9806, 0.97776, 0.97771, 0.97278, 0.97286, 0.96731, 0.96728, 0.95857, 0.95855, 0.94385, 0.94385, 0.92118, 0.921, 0.88108, 0.88076, 0.80263, 0.80181] + +# Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf +guidance_rescale: 0.0 +num_inference_steps: 40 +fps: 16 +save_final_checkpoint: False + +# SDXL Lightning parameters +lightning_from_pt: True +# Empty or "ByteDance/SDXL-Lightning" to enable lightning. +lightning_repo: "" +# Empty or "sdxl_lightning_4step_unet.safetensors" to enable lightning. +lightning_ckpt: "" + +# LoRA parameters +enable_lora: False +# Values are lists to support multiple LoRA loading during inference in the future. +lora_config: { + rank: [64], + lora_model_name_or_path: ["lightx2v/Wan2.2-Distill-Loras"], + high_noise_weight_name: ["wan2.2_t2v_A14b_high_noise_lora_rank64_lightx2v_4step_1217.safetensors"], + low_noise_weight_name: ["wan2.2_t2v_A14b_low_noise_lora_rank64_lightx2v_4step_1217.safetensors"], + adapter_name: ["wan22-distill-lora"], + scale: [1.0], + from_pt: [] +} +# Ex with values: +# lora_config : { +# lora_model_name_or_path: ["ByteDance/Hyper-SD"], +# weight_name: ["Hyper-SDXL-2steps-lora.safetensors"], +# adapter_name: ["hyper-sdxl"], +# scale: [0.7], +# from_pt: [True] +# } + +enable_mllog: False + +#controlnet +controlnet_model_name_or_path: 'diffusers/controlnet-canny-sdxl-1.0' +controlnet_from_pt: True +controlnet_conditioning_scale: 0.5 +controlnet_image: 'https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Google_%22G%22_logo.svg/1024px-Google_%22G%22_logo.svg.png' +quantization: '' +# Shard the range finding operation for quantization. By default this is set to number of slices. +quantization_local_shard_count: -1 +compile_topology_num_slices: -1 # Number of target slices, set to a positive integer. +use_qwix_quantization: False # Whether to use qwix for quantization. If set to True, the transformer of WAN will be quantized using qwix. +# Quantization calibration method used for weights and activations. Supported methods can be found in https://github.com/google/qwix/blob/dc2a0770351c740e5ab3cce7c0efe9f7beacce9e/qwix/qconfig.py#L70-L80 +quantization_calibration_method: "absmax" +qwix_module_path: ".*" + +# Eval model on per eval_every steps. -1 means don't eval. +eval_every: -1 +eval_data_dir: "" +enable_generate_video_for_eval: False # This will increase the used TPU memory. +eval_max_number_of_samples_in_bucket: 60 # The number of samples per bucket for evaluation. This is calculated by num_eval_samples / len(timesteps_list). + +enable_ssim: False + +# ML Diagnostics settings +enable_ml_diagnostics: False +profiler_gcs_path: "" +enable_ondemand_xprof: False \ No newline at end of file From 80b0c3b5c9bb65e571988e1d4f2ab29c0624a8ae Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Wed, 2 Sep 2026 12:02:41 +0530 Subject: [PATCH 5/6] Add initialization and tflops smoke tests for WanTrainer2_2 --- src/maxdiffusion/configs/base_wan_27b.yml | 3 + src/maxdiffusion/configs/training_wan_27b.yml | 3 + src/maxdiffusion/models/wan/wan_utils.py | 18 ++--- src/maxdiffusion/pyconfig.py | 7 +- .../tests/wan/wan_trainer_2_2_test.py | 81 +++++++++++++++++++ src/maxdiffusion/trainers/wan_trainer_2_2.py | 23 ++---- 6 files changed, 106 insertions(+), 29 deletions(-) create mode 100644 src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml index 185b01277..e890014c9 100644 --- a/src/maxdiffusion/configs/base_wan_27b.yml +++ b/src/maxdiffusion/configs/base_wan_27b.yml @@ -409,6 +409,9 @@ num_inference_steps: 40 fps: 16 save_final_checkpoint: False +# Location to download pretrained weights +checkpoint_save_location: "/tmp" + # SDXL Lightning parameters lightning_from_pt: True # Empty or "ByteDance/SDXL-Lightning" to enable lightning. diff --git a/src/maxdiffusion/configs/training_wan_27b.yml b/src/maxdiffusion/configs/training_wan_27b.yml index 40bd935fe..46cbad763 100644 --- a/src/maxdiffusion/configs/training_wan_27b.yml +++ b/src/maxdiffusion/configs/training_wan_27b.yml @@ -411,6 +411,9 @@ num_inference_steps: 40 fps: 16 save_final_checkpoint: False +# Location to download pretrained weights +checkpoint_save_location: "/tmp" + # SDXL Lightning parameters lightning_from_pt: True # Empty or "ByteDance/SDXL-Lightning" to enable lightning. diff --git a/src/maxdiffusion/models/wan/wan_utils.py b/src/maxdiffusion/models/wan/wan_utils.py index 3db43fb74..a4383befa 100644 --- a/src/maxdiffusion/models/wan/wan_utils.py +++ b/src/maxdiffusion/models/wan/wan_utils.py @@ -514,15 +514,15 @@ def convert_chunk(ckpt_shard_path, chunk_keys): for future in concurrent.futures.as_completed(futures): future.result() # re-raise conversion errors - if not local_files: - try: - real_path = os.path.realpath(ckpt_shard_path) - if os.path.exists(real_path): - os.remove(real_path) - if os.path.exists(ckpt_shard_path): - os.remove(ckpt_shard_path) - except Exception as e: - max_logging.log(f"Warning: could not delete shard {ckpt_shard_path}: {e}") + if not local_files: + try: + real_path = os.path.realpath(ckpt_shard_path) + if os.path.exists(real_path): + os.remove(real_path) + if os.path.exists(ckpt_shard_path): + os.remove(ckpt_shard_path) + except Exception as e: + max_logging.log(f"Warning: could not delete shard {ckpt_shard_path}: {e}") validate_flax_state_dict(eval_shapes, flax_state_dict) diff --git a/src/maxdiffusion/pyconfig.py b/src/maxdiffusion/pyconfig.py index bc29ba23e..cd824d726 100644 --- a/src/maxdiffusion/pyconfig.py +++ b/src/maxdiffusion/pyconfig.py @@ -283,12 +283,13 @@ 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"], "/dev/shm") + 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"], "/dev/shm") + 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"], "/dev/shm") + 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"] diff --git a/src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py b/src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py new file mode 100644 index 000000000..da85e3687 --- /dev/null +++ b/src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py @@ -0,0 +1,81 @@ +""" +Copyright 2025 Google LLC + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import os +import unittest +from absl.testing import absltest +from maxdiffusion import pyconfig +from maxdiffusion.trainers.wan_trainer_2_2 import WanTrainer2_2 + +IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" +THIS_DIR = os.path.dirname(os.path.abspath(__file__)) + +class WanTrainer22Test(unittest.TestCase): + + def test_wan_trainer_2_2_initialization(self): + """Smoke test to ensure WanTrainer2_2 can be initialized with config.""" + pyconfig.initialize([ + None, + os.path.join(THIS_DIR, "..", "..", "configs", "base_wan_27b.yml"), + "max_train_steps=10", + "per_device_batch_size=1", + "dataset_type=synthetic", + "cache_latents_text_encoder_outputs=True", + ], unittest=True) + + config = pyconfig.config + trainer = WanTrainer2_2(config) + self.assertIsNotNone(trainer) + # Check that checkpointer is created + checkpointer = trainer._get_checkpointer() + self.assertIsNotNone(checkpointer) + + + def test_calculate_tflops(self): + from unittest.mock import MagicMock + pyconfig.initialize([ + None, + os.path.join(THIS_DIR, "..", "..", "configs", "base_wan_27b.yml"), + "max_train_steps=10", + "per_device_batch_size=1", + ], unittest=True) + + trainer = WanTrainer2_2(pyconfig.config) + + mock_pipeline = MagicMock() + mock_pipeline.config.height = 256 + mock_pipeline.config.width = 256 + mock_pipeline.config.num_frames = 1 + mock_pipeline.vae_scale_factor_temporal = 4 + mock_pipeline.config.per_device_batch_size = 1 + + mock_transformer_config = MagicMock() + mock_transformer_config.num_layers = 2 + mock_transformer_config.num_attention_heads = 4 + mock_transformer_config.attention_head_dim = 64 + mock_transformer_config.ffn_dim = 256 + + mock_pipeline.low_noise_transformer.config = mock_transformer_config + + train_tflops, total_attn_flops, seq_len = trainer.calculate_tflops(mock_pipeline) + + self.assertIsNotNone(train_tflops) + self.assertIsNotNone(total_attn_flops) + self.assertIsNotNone(seq_len) + self.assertGreater(train_tflops, 0.0) + +if __name__ == "__main__": + absltest.main() diff --git a/src/maxdiffusion/trainers/wan_trainer_2_2.py b/src/maxdiffusion/trainers/wan_trainer_2_2.py index 4a40719d3..7f73e8bc9 100644 --- a/src/maxdiffusion/trainers/wan_trainer_2_2.py +++ b/src/maxdiffusion/trainers/wan_trainer_2_2.py @@ -279,8 +279,6 @@ def training_loop_2_2(self, pipeline, optimizer_low, optimizer_high, learning_ra example_batch = load_next_batch(train_data_iterator, None, self.config) with ThreadPoolExecutor(max_workers=1) as executor: - step_high = 0 - step_low = 0 for step in np.arange(start_step, self.config.max_train_steps): start_step_time = datetime.datetime.now() next_batch_future = executor.submit(load_next_batch, train_data_iterator, example_batch, self.config) @@ -296,19 +294,7 @@ def training_loop_2_2(self, pipeline, optimizer_low, optimizer_high, learning_ra train_metric, last_step_completion - start_step_time, per_device_tflops, learning_rate_scheduler(step) ) - loss_high_val = np.array(train_metric["scalar"]["learning/loss_high"]) - loss_low_val = np.array(train_metric["scalar"]["learning/loss_low"]) - if "steps" not in train_metric: - train_metric["steps"] = {} - if not np.isnan(loss_high_val): - step_high += 1 - train_metric["steps"]["learning/loss_high"] = step_high - if not np.isnan(loss_low_val): - step_low += 1 - train_metric["steps"]["learning/loss_low"] = step_low - train_metric["steps"]["learning/loss_low_fine"] = step_low - train_metric["steps"]["learning/loss_low_mid"] = step_low - train_metric["steps"]["learning/loss_low_coarse"] = step_low + if self.config.write_metrics: train_utils.write_metrics(writer, local_metrics_file, running_gcs_metrics, train_metric, step, self.config) @@ -496,8 +482,11 @@ def false_fn(params_tuple): "scalars": {}, } - new_state_low = state_low.apply_gradients(grads=grads_low) - new_state_high = state_high.apply_gradients(grads=grads_high) + new_state_high, new_state_low = jax.lax.cond( + is_high_noise, + lambda: (state_high.apply_gradients(grads=grads_high), state_low), + lambda: (state_high, state_low.apply_gradients(grads=grads_low)) + ) return new_state_low, new_state_high, scheduler_state, metrics, new_rng From 4ffe7dc401fb53b14503cca1021eb88db21a352e Mon Sep 17 00:00:00 2001 From: Toshi Pahadia Date: Wed, 2 Sep 2026 16:13:20 +0530 Subject: [PATCH 6/6] Revert local debugging changes to docker and setup scripts --- docker_build_dependency_image.sh | 3 ++- maxdiffusion_dependencies.Dockerfile | 2 +- setup.sh | 4 ++-- .../tests/wan/wan_trainer_2_2_test.py | 16 ++++++++-------- 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/docker_build_dependency_image.sh b/docker_build_dependency_image.sh index aa26ce9bf..5c2c0f8d2 100644 --- a/docker_build_dependency_image.sh +++ b/docker_build_dependency_image.sh @@ -27,7 +27,8 @@ set -e export LOCAL_IMAGE_NAME=maxdiffusion_base_image -export DOCKER_BUILDKIT=0 +# Use Docker BuildKit so we can cache pip packages. +export DOCKER_BUILDKIT=1 echo "Starting to build your docker image. This will take a few minutes but the image can be reused as you iterate." diff --git a/maxdiffusion_dependencies.Dockerfile b/maxdiffusion_dependencies.Dockerfile index cf025876a..9a9598271 100644 --- a/maxdiffusion_dependencies.Dockerfile +++ b/maxdiffusion_dependencies.Dockerfile @@ -48,7 +48,7 @@ COPY . . RUN echo "Running command: bash setup.sh MODE=$ENV_MODE JAX_VERSION=$ENV_JAX_VERSION" -RUN bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} +RUN --mount=type=cache,target=/root/.cache/pip bash setup.sh MODE=${ENV_MODE} JAX_VERSION=${ENV_JAX_VERSION} # Cleanup RUN rm -rf /root/.cache/pip \ No newline at end of file diff --git a/setup.sh b/setup.sh index 501be7926..3f1141888 100644 --- a/setup.sh +++ b/setup.sh @@ -161,8 +161,8 @@ elif [[ $MODE == "nightly" ]]; then python3 -m uv pip install --pre -U jax -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html # Install jaxlib-nightly python3 -m uv pip install --pre -U jaxlib -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html - # Install libtpu nightly (package is now named libtpu, not libtpu-nightly) - python3 -m uv pip install --pre -U libtpu -f https://storage.googleapis.com/jax-releases/libtpu_releases.html + # Install libtpu-nightly + python3 -m uv pip install --pre -U libtpu-nightly -f https://storage.googleapis.com/jax-releases/libtpu_releases.html fi echo "Installing nightly tensorboard plugin profile" python3 -m uv pip install tbp-nightly --upgrade diff --git a/src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py b/src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py index da85e3687..a3a7e4844 100644 --- a/src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py +++ b/src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py @@ -24,7 +24,7 @@ THIS_DIR = os.path.dirname(os.path.abspath(__file__)) class WanTrainer22Test(unittest.TestCase): - + def test_wan_trainer_2_2_initialization(self): """Smoke test to ensure WanTrainer2_2 can be initialized with config.""" pyconfig.initialize([ @@ -35,7 +35,7 @@ def test_wan_trainer_2_2_initialization(self): "dataset_type=synthetic", "cache_latents_text_encoder_outputs=True", ], unittest=True) - + config = pyconfig.config trainer = WanTrainer2_2(config) self.assertIsNotNone(trainer) @@ -52,26 +52,26 @@ def test_calculate_tflops(self): "max_train_steps=10", "per_device_batch_size=1", ], unittest=True) - + trainer = WanTrainer2_2(pyconfig.config) - + mock_pipeline = MagicMock() mock_pipeline.config.height = 256 mock_pipeline.config.width = 256 mock_pipeline.config.num_frames = 1 mock_pipeline.vae_scale_factor_temporal = 4 mock_pipeline.config.per_device_batch_size = 1 - + mock_transformer_config = MagicMock() mock_transformer_config.num_layers = 2 mock_transformer_config.num_attention_heads = 4 mock_transformer_config.attention_head_dim = 64 mock_transformer_config.ffn_dim = 256 - + mock_pipeline.low_noise_transformer.config = mock_transformer_config - + train_tflops, total_attn_flops, seq_len = trainer.calculate_tflops(mock_pipeline) - + self.assertIsNotNone(train_tflops) self.assertIsNotNone(total_attn_flops) self.assertIsNotNone(seq_len)