Wan 2.2 training

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.

src/maxdiffusion/max_utils.py

@@ -301,21 +301,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):

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 = []
-  for model_file in model_files:
-    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"
+      f"Load and port {pretrained_model_name_or_path} {subfolder}: {len(model_files)} shards sequentially to save disk"
   )
   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
+    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
+
+      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):

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):
@@ -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"], "/tmp")
+      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"], "/tmp")
+      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"], "/tmp")
+      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"]

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()

src/maxdiffusion/train_utils.py

@@ -130,7 +130,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)
 
@@ -198,7 +201,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)
 

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)