diff --git a/src/maxdiffusion/checkpointing/wan_checkpointer_2_2.py b/src/maxdiffusion/checkpointing/wan_checkpointer_2_2.py
index 20b984447..10dbab172 100644
--- a/src/maxdiffusion/checkpointing/wan_checkpointer_2_2.py
+++ b/src/maxdiffusion/checkpointing/wan_checkpointer_2_2.py
@@ -18,7 +18,7 @@
 import jax
 from typing import Optional, Tuple
 from ..pipelines.wan.wan_pipeline_2_2 import WanPipeline2_2
-from .. import max_logging
+from .. import max_logging, max_utils
 import orbax.checkpoint as ocp
 from maxdiffusion.checkpointing.checkpointing_utils import add_sharding_to_struct, get_cpu_mesh_and_sharding
 from maxdiffusion.checkpointing.wan_checkpointer import WanCheckpointer
@@ -27,6 +27,15 @@
 class WanCheckpointer2_2(WanCheckpointer[WanPipeline2_2]):
   pipeline_class = WanPipeline2_2
 
+  def _create_optimizer(self, model, config, learning_rate, scale_factor: float = 1.0):
+    total_steps = max(1, int(config.max_train_steps * scale_factor))
+    schedule_steps = max(1, int(config.learning_rate_schedule_steps * scale_factor))
+    learning_rate_scheduler = max_utils.create_learning_rate_schedule(
+        learning_rate, schedule_steps, config.warmup_steps_fraction, total_steps
+    )
+    tx = max_utils.create_optimizer(config, learning_rate_scheduler)
+    return tx, learning_rate_scheduler
+
   def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dict], Optional[int]]:
     if step is None:
       step = self.checkpoint_manager.latest_step()
@@ -81,11 +90,20 @@ def load_wan_configs_from_orbax(self, step: Optional[int]) -> Tuple[Optional[dic
     return restored_checkpoint, step
 
   def _extract_opt_state(self, restored_checkpoint):
-    if "opt_state" in restored_checkpoint.low_noise_transformer_state.keys():
-      return restored_checkpoint.low_noise_transformer_state["opt_state"]
-    elif "opt_state" in restored_checkpoint.high_noise_transformer_state.keys():
-      return restored_checkpoint.high_noise_transformer_state["opt_state"]
-    return None
+    low_state = getattr(restored_checkpoint, "low_noise_transformer_state", {})
+    high_state = getattr(restored_checkpoint, "high_noise_transformer_state", {})
+    low_opt = low_state.get("opt_state") if isinstance(low_state, dict) else getattr(low_state, "opt_state", None)
+    high_opt = high_state.get("opt_state") if isinstance(high_state, dict) else getattr(high_state, "opt_state", None)
+    low_step = low_state.get("step") if isinstance(low_state, dict) else getattr(low_state, "step", None)
+    high_step = high_state.get("step") if isinstance(high_state, dict) else getattr(high_state, "step", None)
+    if low_opt is None and high_opt is None:
+      return None
+    return {
+        "low_noise_transformer": low_opt,
+        "high_noise_transformer": high_opt,
+        "low_noise_step": low_step,
+        "high_noise_step": high_step,
+    }
 
   def save_checkpoint(self, train_step, pipeline: WanPipeline2_2, train_states: dict):
     """Saves the training state and model configurations."""
diff --git a/src/maxdiffusion/configs/base_wan_27b.yml b/src/maxdiffusion/configs/base_wan_27b.yml
index 35b17f9af..fd98959f5 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/pyconfig.py b/src/maxdiffusion/pyconfig.py
index d1121ca3f..cd824d726 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):
@@ -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"]
diff --git a/src/maxdiffusion/tests/wan/wan_checkpointer_test.py b/src/maxdiffusion/tests/wan/wan_checkpointer_test.py
index b30006edf..c48f9b13e 100644
--- a/src/maxdiffusion/tests/wan/wan_checkpointer_test.py
+++ b/src/maxdiffusion/tests/wan/wan_checkpointer_test.py
@@ -387,7 +387,8 @@ def test_load_checkpoint_with_optimizer_in_low_noise(self, mock_from_checkpoint,
     )
     self.assertEqual(pipeline, mock_pipeline_instance)
     self.assertIsNotNone(opt_state)
-    self.assertEqual(opt_state["learning_rate"], 0.001)
+    self.assertEqual(opt_state["low_noise_transformer"]["learning_rate"], 0.001)
+    self.assertIsNone(opt_state["high_noise_transformer"])
     self.assertEqual(step, 1)
 
   @patch("maxdiffusion.checkpointing.wan_checkpointer.create_orbax_checkpoint_manager")
@@ -429,7 +430,8 @@ def test_load_checkpoint_with_optimizer_in_high_noise(self, mock_from_checkpoint
     )
     self.assertEqual(pipeline, mock_pipeline_instance)
     self.assertIsNotNone(opt_state)
-    self.assertEqual(opt_state["learning_rate"], 0.002)
+    self.assertIsNone(opt_state["low_noise_transformer"])
+    self.assertEqual(opt_state["high_noise_transformer"]["learning_rate"], 0.002)
     self.assertEqual(step, 1)
 
 
@@ -758,9 +760,10 @@ def test_load_checkpoint_both_optimizers_present(self, mock_from_checkpoint, moc
     checkpointer = WanCheckpointer2_2(config=self.config)
     pipeline, opt_state, step = checkpointer.load_checkpoint(step=1)
 
-    # Should prioritize low_noise_transformer's optimizer state
+    # Should preserve both low_noise_transformer and high_noise_transformer optimizer states
     self.assertIsNotNone(opt_state)
-    self.assertEqual(opt_state["learning_rate"], 0.001)
+    self.assertEqual(opt_state["low_noise_transformer"]["learning_rate"], 0.001)
+    self.assertEqual(opt_state["high_noise_transformer"]["learning_rate"], 0.002)
 
 
 if __name__ == "__main__":
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..de4401e21
--- /dev/null
+++ b/src/maxdiffusion/tests/wan/wan_trainer_2_2_test.py
@@ -0,0 +1,411 @@
+"""
+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 os
+import unittest
+from unittest.mock import MagicMock
+from absl.testing import absltest
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import optax
+
+import numpy as np
+from maxdiffusion import pyconfig
+from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2
+from maxdiffusion.schedulers import FlaxFlowMatchScheduler
+from maxdiffusion.trainers.base_wan_trainer import TrainState
+from maxdiffusion.trainers.wan_trainer_2_2 import (
+    WanTrainer2_2,
+    train_step_2_2,
+    eval_step_2_2,
+)
+
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+BASE_CONFIG_PATH = os.path.join(THIS_DIR, "..", "..", "configs", "base_wan_27b.yml")
+
+
+class MiniWanModel(nnx.Module):
+  """Lightweight Flax module implementing WanModel interface for real training step execution."""
+
+  def __init__(self, rngs: nnx.Rngs, in_channels: int = 4):
+    self.conv = nnx.Conv(in_channels, in_channels, kernel_size=(1, 1, 1), rngs=rngs)
+    self.time_proj = nnx.Linear(1, in_channels, rngs=rngs)
+
+  def __call__(
+      self,
+      hidden_states,
+      timestep,
+      encoder_hidden_states=None,
+      deterministic=False,
+      rngs=None,
+  ):
+    # hidden_states: (bsz, C, T, H, W)
+    t_emb = self.time_proj(timestep[:, None].astype(jnp.float32))[:, :, None, None, None]
+    x = jnp.transpose(hidden_states, (0, 2, 3, 4, 1))
+    x = self.conv(x)
+    x = jnp.transpose(x, (0, 4, 1, 2, 3))
+    return x + t_emb
+
+
+class WanTrainer22Test(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    pyconfig.initialize([
+        None,
+        BASE_CONFIG_PATH,
+        "max_train_steps=100",
+        "learning_rate_schedule_steps=100",
+        "warmup_steps_fraction=0.1",
+        "per_device_batch_size=2",
+        "dataset_type=synthetic",
+        "cache_latents_text_encoder_outputs=True",
+        "replicate_vae=True",
+        "boundary_ratio=0.5",
+        "weights_dtype=float32",
+    ], unittest=True)
+
+  def test_wan_trainer_2_2_initialization(self):
+    """Smoke test to ensure WanTrainer2_2 can be initialized with base config."""
+    config = pyconfig.config
+    trainer = WanTrainer2_2(config)
+    self.assertIsNotNone(trainer)
+    checkpointer = trainer._get_checkpointer()
+    self.assertIsInstance(checkpointer, WanCheckpointer2_2)
+
+    mesh = jax.sharding.Mesh(np.array(jax.devices()[:1]).reshape((1, 1, 1, 1)), ("data", "fsdp", "context", "tensor"))
+    data_shardings = trainer.get_data_shardings(mesh)
+    self.assertIn("latents", data_shardings)
+    self.assertIn("encoder_hidden_states", data_shardings)
+
+    eval_shardings = trainer.get_eval_data_shardings(mesh)
+    self.assertIn("timesteps", eval_shardings)
+
+  def test_calculate_tflops_formula(self):
+    """Verify corrected TFLOPs calculation matches expected mathematical values."""
+    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_transformer_config.text_dim = 4096
+
+    mock_pipeline.low_noise_transformer.config = mock_transformer_config
+
+    train_tflops, total_attn_flops, seq_len = trainer.calculate_tflops(mock_pipeline)
+
+    self.assertEqual(seq_len, 256)
+    self.assertGreater(train_tflops, 0.0)
+    self.assertGreater(total_attn_flops, 0)
+    self.assertAlmostEqual(train_tflops, 3 * 5301600256 / 1e12, places=6)
+
+  def test_optimizer_schedule_scaling(self):
+    """Verify Optax schedule lengths scale proportionally with boundary_ratio."""
+    trainer = WanTrainer2_2(pyconfig.config)
+    checkpointer = trainer._get_checkpointer()
+
+    mock_model = MagicMock()
+
+    _, lr_sched_low = checkpointer._create_optimizer(
+        mock_model, pyconfig.config, pyconfig.config.learning_rate, scale_factor=0.5
+    )
+    _, lr_sched_high = checkpointer._create_optimizer(
+        mock_model, pyconfig.config, pyconfig.config.learning_rate, scale_factor=0.5
+    )
+
+    self.assertEqual(float(lr_sched_low(0)), 0.0)
+    self.assertAlmostEqual(float(lr_sched_low(5)), float(pyconfig.config.learning_rate), places=5)
+    self.assertEqual(float(lr_sched_high(0)), 0.0)
+    self.assertAlmostEqual(float(lr_sched_high(5)), float(pyconfig.config.learning_rate), places=5)
+
+  def test_extract_opt_state(self):
+    """Verify WanCheckpointer2_2 extracts both optimizer states."""
+    checkpointer = WanCheckpointer2_2(config=pyconfig.config)
+    mock_checkpoint = MagicMock()
+    mock_checkpoint.low_noise_transformer_state = {"opt_state": "low_opt_state_data"}
+    mock_checkpoint.high_noise_transformer_state = {"opt_state": "high_opt_state_data"}
+
+    extracted = checkpointer._extract_opt_state(mock_checkpoint)
+    self.assertIn("low_noise_transformer", extracted)
+    self.assertIn("high_noise_transformer", extracted)
+    self.assertEqual(extracted["low_noise_transformer"], "low_opt_state_data")
+    self.assertEqual(extracted["high_noise_transformer"], "high_opt_state_data")
+
+  def test_real_train_step_2_2_execution(self):
+    """Execute real JIT-compiled training steps and verify gradient application and stepping."""
+    rng = jax.random.key(42)
+    rng_low, rng_high, step_rng = jax.random.split(rng, 3)
+
+    model_low = MiniWanModel(rngs=nnx.Rngs(rng_low), in_channels=4)
+    model_high = MiniWanModel(rngs=nnx.Rngs(rng_high), in_channels=4)
+
+    graphdef_low, params_low, rest_of_state_low = nnx.split(model_low, nnx.Param, ...)
+    graphdef_high, params_high, rest_of_state_high = nnx.split(model_high, nnx.Param, ...)
+
+    tx_low = optax.adam(1e-3)
+    tx_high = optax.adam(1e-3)
+    state_low = TrainState.create(
+        apply_fn=graphdef_low.apply, params=params_low, tx=tx_low, graphdef=graphdef_low, rest_of_state=rest_of_state_low
+    )
+    state_high = TrainState.create(
+        apply_fn=graphdef_high.apply,
+        params=params_high,
+        tx=tx_high,
+        graphdef=graphdef_high,
+        rest_of_state=rest_of_state_high,
+    )
+
+    # Use real Flow Match scheduler
+    noise_scheduler = FlaxFlowMatchScheduler(dtype=jnp.float32)
+    noise_scheduler_state = noise_scheduler.create_state()
+    noise_scheduler_state = noise_scheduler.set_timesteps(
+        noise_scheduler_state, num_inference_steps=1000, training=True
+    )
+
+    config = pyconfig.config
+
+    data = {
+        "latents": jnp.ones((2, 4, 2, 4, 4), dtype=jnp.float32),
+        "encoder_hidden_states": jnp.ones((2, 16, 4), dtype=jnp.float32),
+    }
+
+    jitted_train_step = jax.jit(
+        functools.partial(train_step_2_2, scheduler=noise_scheduler, config=config)
+    )
+
+    initial_params_low = jax.tree.map(lambda x: jnp.copy(x), state_low.params)
+    initial_params_high = jax.tree.map(lambda x: jnp.copy(x), state_high.params)
+
+    num_steps = 6
+    for _ in range(num_steps):
+      state_low, state_high, noise_scheduler_state, metrics, step_rng = jitted_train_step(
+          state_low, state_high, data, step_rng, noise_scheduler_state
+      )
+      loss_val = float(metrics["scalar"]["learning/loss"])
+      self.assertFalse(jnp.isnan(loss_val), "Training loss returned NaN")
+      self.assertFalse(jnp.isinf(loss_val), "Training loss returned Inf")
+      self.assertGreater(loss_val, 0.0, "Training loss should be positive")
+
+    # Verify total steps across both transformers equal the number of train steps executed
+    total_steps = int(state_low.step) + int(state_high.step)
+    self.assertEqual(total_steps, num_steps)
+
+    # Verify that at least one transformer had parameter updates
+    low_updated = any(
+        not jnp.array_equal(p1, p2)
+        for p1, p2 in zip(jax.tree.leaves(initial_params_low), jax.tree.leaves(state_low.params))
+    )
+    high_updated = any(
+        not jnp.array_equal(p1, p2)
+        for p1, p2 in zip(jax.tree.leaves(initial_params_high), jax.tree.leaves(state_high.params))
+    )
+    self.assertTrue(low_updated or high_updated, "Parameters should be updated after training steps")
+
+  def test_real_eval_step_2_2_execution(self):
+    """Execute real JIT-compiled evaluation step and verify loss output."""
+    rng = jax.random.key(0)
+    rng_low, rng_high, eval_rng = jax.random.split(rng, 3)
+
+    model_low = MiniWanModel(rngs=nnx.Rngs(rng_low), in_channels=4)
+    model_high = MiniWanModel(rngs=nnx.Rngs(rng_high), in_channels=4)
+
+    graphdef_low, params_low, rest_of_state_low = nnx.split(model_low, nnx.Param, ...)
+    graphdef_high, params_high, rest_of_state_high = nnx.split(model_high, nnx.Param, ...)
+
+    tx = optax.adam(1e-3)
+    state_low = TrainState.create(
+        apply_fn=graphdef_low.apply, params=params_low, tx=tx, graphdef=graphdef_low, rest_of_state=rest_of_state_low
+    )
+    state_high = TrainState.create(
+        apply_fn=graphdef_high.apply, params=params_high, tx=tx, graphdef=graphdef_high, rest_of_state=rest_of_state_high
+    )
+
+    noise_scheduler = FlaxFlowMatchScheduler(dtype=jnp.float32)
+    noise_scheduler_state = noise_scheduler.create_state()
+    noise_scheduler_state = noise_scheduler.set_timesteps(
+        noise_scheduler_state, num_inference_steps=1000, training=True
+    )
+
+    config = pyconfig.config
+
+    data = {
+        "latents": jnp.ones((2, 4, 2, 4, 4), dtype=jnp.float32),
+        "encoder_hidden_states": jnp.ones((2, 16, 4), dtype=jnp.float32),
+        "timesteps": jnp.array([100, 700], dtype=jnp.int64),
+    }
+
+    jitted_eval_step = jax.jit(
+        functools.partial(eval_step_2_2, scheduler=noise_scheduler, config=config)
+    )
+
+    metrics, _ = jitted_eval_step(state_low, state_high, data, eval_rng, noise_scheduler_state)
+    losses = metrics["scalar"]["learning/eval_loss"]
+    self.assertEqual(len(losses), 2)
+    self.assertFalse(jnp.isnan(losses).any(), "Eval losses should not be NaN")
+    self.assertTrue((losses >= 0).all(), "Eval losses should be non-negative")
+
+  def test_boundary_ratio_validation(self):
+    """Verify that boundary_ratio <= 0.0 or >= 1.0 raises ValueError."""
+    mock_config = MagicMock()
+    mock_config.train_text_encoder = False
+
+    for invalid_ratio in [0.0, 1.0, -0.2, 1.5]:
+      mock_config.boundary_ratio = invalid_ratio
+      with self.assertRaises(ValueError):
+        WanTrainer2_2(mock_config)
+
+  def test_eval_expert_routing_by_timestep(self):
+    """Verify eval_step_2_2 routes t < boundary to low-noise model and t >= boundary to high-noise model."""
+    class DistinguishableModel(nnx.Module):
+      def __init__(self, val: float):
+        self.val = nnx.Param(jnp.array(val, dtype=jnp.float32))
+
+      def __call__(self, hidden_states, timestep, encoder_hidden_states=None, deterministic=True):
+        return jnp.full_like(hidden_states, self.val.value)
+
+    model_low = DistinguishableModel(1.0)
+    model_high = DistinguishableModel(10.0)
+
+    graphdef_low, params_low, rest_of_state_low = nnx.split(model_low, nnx.Param, ...)
+    graphdef_high, params_high, rest_of_state_high = nnx.split(model_high, nnx.Param, ...)
+
+    tx = optax.adam(1e-3)
+    state_low = TrainState.create(apply_fn=graphdef_low.apply, params=params_low, tx=tx, graphdef=graphdef_low, rest_of_state=rest_of_state_low)
+    state_high = TrainState.create(apply_fn=graphdef_high.apply, params=params_high, tx=tx, graphdef=graphdef_high, rest_of_state=rest_of_state_high)
+
+    noise_scheduler = FlaxFlowMatchScheduler(dtype=jnp.float32)
+    noise_scheduler_state = noise_scheduler.create_state()
+    noise_scheduler_state = noise_scheduler.set_timesteps(noise_scheduler_state, num_inference_steps=1000, training=True)
+
+    config = pyconfig.config
+    # boundary = int(0.5 * 1000) = 500
+    # t=100 -> low expert (output 1.0)
+    # t=700 -> high expert (output 10.0)
+    data = {
+        "latents": jnp.zeros((2, 4, 2, 4, 4), dtype=jnp.float32),
+        "encoder_hidden_states": jnp.zeros((2, 16, 4), dtype=jnp.float32),
+        "timesteps": jnp.array([100, 700], dtype=jnp.int64),
+    }
+
+    # Disable training weights for predictable pure MSE loss: (target - pred)^2
+    eval_fn = jax.jit(functools.partial(eval_step_2_2, scheduler=noise_scheduler, config=config))
+    metrics, _ = eval_fn(state_low, state_high, data, jax.random.key(0), noise_scheduler_state)
+    losses = metrics["scalar"]["learning/eval_loss"]
+
+    # Since low model output is 1.0 and high model output is 10.0, the losses must be distinctly different
+    self.assertEqual(len(losses), 2)
+    self.assertNotEqual(float(losses[0]), float(losses[1]))
+
+  def test_checkpoint_resume_equivalence(self):
+    """Verify that continuous training and train->checkpoint->restore->train produce identical trajectories."""
+    rng = jax.random.key(99)
+    rng_low, rng_high = jax.random.split(rng, 2)
+
+    def _init_states():
+      model_low = MiniWanModel(rngs=nnx.Rngs(rng_low), in_channels=4)
+      model_high = MiniWanModel(rngs=nnx.Rngs(rng_high), in_channels=4)
+      g_low, p_low, r_low = nnx.split(model_low, nnx.Param, ...)
+      g_high, p_high, r_high = nnx.split(model_high, nnx.Param, ...)
+      tx_low = optax.adam(1e-3)
+      tx_high = optax.adam(1e-3)
+      s_low = TrainState.create(apply_fn=g_low.apply, params=p_low, tx=tx_low, graphdef=g_low, rest_of_state=r_low)
+      s_high = TrainState.create(apply_fn=g_high.apply, params=p_high, tx=tx_high, graphdef=g_high, rest_of_state=r_high)
+      return s_low, s_high
+
+    noise_scheduler = FlaxFlowMatchScheduler(dtype=jnp.float32)
+    noise_scheduler_state = noise_scheduler.create_state()
+    noise_scheduler_state = noise_scheduler.set_timesteps(noise_scheduler_state, num_inference_steps=1000, training=True)
+    config = pyconfig.config
+
+    data = {
+        "latents": jnp.ones((2, 4, 2, 4, 4), dtype=jnp.float32),
+        "encoder_hidden_states": jnp.ones((2, 16, 4), dtype=jnp.float32),
+    }
+
+    train_step_fn = jax.jit(functools.partial(train_step_2_2, scheduler=noise_scheduler, config=config))
+
+    # Run A: 6 continuous steps
+    s_low_a, s_high_a = _init_states()
+    step_rng_a = jax.random.key(1001)
+    sched_state_a = noise_scheduler_state
+    for _ in range(6):
+      s_low_a, s_high_a, sched_state_a, _, step_rng_a = train_step_fn(
+          s_low_a, s_high_a, data, step_rng_a, sched_state_a
+      )
+
+    # Run B: 3 steps, simulate save/restore, 3 additional steps
+    s_low_b, s_high_b = _init_states()
+    step_rng_b = jax.random.key(1001)
+    sched_state_b = noise_scheduler_state
+    for _ in range(3):
+      s_low_b, s_high_b, sched_state_b, _, step_rng_b = train_step_fn(
+          s_low_b, s_high_b, data, step_rng_b, sched_state_b
+      )
+
+    # Simulate checkpoint extraction and restore
+    checkpoint = MagicMock()
+    checkpoint.low_noise_transformer_state = {"params": s_low_b.params, "opt_state": s_low_b.opt_state, "step": s_low_b.step}
+    checkpoint.high_noise_transformer_state = {"params": s_high_b.params, "opt_state": s_high_b.opt_state, "step": s_high_b.step}
+    checkpointer = WanCheckpointer2_2(config=config)
+    opt_state_dict = checkpointer._extract_opt_state(checkpoint)
+
+    # Re-initialize clean state and restore
+    s_low_b_resumed, s_high_b_resumed = _init_states()
+    s_low_b_resumed = s_low_b_resumed.replace(
+        params=checkpoint.low_noise_transformer_state["params"],
+        opt_state=opt_state_dict["low_noise_transformer"],
+        step=opt_state_dict["low_noise_step"],
+    )
+    s_high_b_resumed = s_high_b_resumed.replace(
+        params=checkpoint.high_noise_transformer_state["params"],
+        opt_state=opt_state_dict["high_noise_transformer"],
+        step=opt_state_dict["high_noise_step"],
+    )
+
+    # Verify steps preserved before continuing
+    self.assertEqual(int(s_low_b_resumed.step), int(s_low_b.step))
+    self.assertEqual(int(s_high_b_resumed.step), int(s_high_b.step))
+
+    # Run remaining 3 steps
+    for _ in range(3):
+      s_low_b_resumed, s_high_b_resumed, sched_state_b, _, step_rng_b = train_step_fn(
+          s_low_b_resumed, s_high_b_resumed, data, step_rng_b, sched_state_b
+      )
+
+    # Assert exact equivalence between continuous and resumed runs
+    self.assertEqual(int(s_low_a.step), int(s_low_b_resumed.step))
+    self.assertEqual(int(s_high_a.step), int(s_high_b_resumed.step))
+
+    for p_a, p_b in zip(jax.tree.leaves(s_low_a.params), jax.tree.leaves(s_low_b_resumed.params)):
+      self.assertTrue(jnp.allclose(p_a, p_b, atol=1e-6), "Low noise parameters mismatch between continuous and resumed runs")
+
+    for p_a, p_b in zip(jax.tree.leaves(s_high_a.params), jax.tree.leaves(s_high_b_resumed.params)):
+      self.assertTrue(jnp.allclose(p_a, p_b, atol=1e-6), "High noise parameters mismatch between continuous and resumed runs")
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/src/maxdiffusion/train_utils.py b/src/maxdiffusion/train_utils.py
index b3cf98e29..6bfab9da5 100644
--- a/src/maxdiffusion/train_utils.py
+++ b/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)
 
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..07e400880
--- /dev/null
+++ b/src/maxdiffusion/trainers/wan_trainer_2_2.py
@@ -0,0 +1,606 @@
+"""
+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 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
+
+from maxdiffusion.checkpointing.wan_checkpointer_2_2 import WanCheckpointer2_2
+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 __init__(self, config):
+    super().__init__(config)
+    if not (0.0 < self.config.boundary_ratio < 1.0):
+      raise ValueError(
+          f"boundary_ratio must be strictly between 0 and 1, got {self.config.boundary_ratio}"
+      )
+
+  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))
+    timesteps_axis = self.config.data_sharding[0] if self.config.data_sharding else None
+    timesteps_sharding = jax.sharding.NamedSharding(mesh, P(timesteps_axis))
+    return {"latents": data_sharding, "encoder_hidden_states": data_sharding, "timesteps": timesteps_sharding}
+
+  def load_dataset(self, mesh, pipeline=None, is_training=True):
+    import tensorflow as tf
+    from maxdiffusion.input_pipeline.input_pipeline_interface import make_data_iterator
+
+    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" 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"
+      )
+    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 = transformer_config.num_attention_heads
+    head_dim = transformer_config.attention_head_dim
+    hidden_dim = heads * head_dim
+    ffn_dim = transformer_config.ffn_dim
+    text_dim = getattr(transformer_config, "text_dim", 4096)
+    text_seq_len = 512
+    seq_len = int(((height / 8) * (width / 8) * ((num_frames - 1) // pipeline.vae_scale_factor_temporal + 1)) / 4)
+
+    # Self-attention FLOPs
+    self_attn_qkv_proj_flops = 3 * (2 * seq_len * hidden_dim**2)
+    self_attn_qk_v_flops = 2 * (2 * seq_len**2 * hidden_dim)
+    self_attn_output_proj_flops = 1 * (2 * seq_len * hidden_dim**2)
+
+    # Cross-attention FLOPs
+    cross_attn_q_proj_flops = 1 * (2 * seq_len * hidden_dim**2)
+    cross_attn_kv_proj_flops = 2 * (2 * text_seq_len * text_dim * hidden_dim)
+    cross_attention_qk_v_flops = 2 * (2 * seq_len * text_seq_len * hidden_dim)
+    cross_attn_output_proj_flops = 1 * (2 * seq_len * hidden_dim**2)
+
+    total_attn_flops = (
+        self_attn_qkv_proj_flops
+        + self_attn_qk_v_flops
+        + self_attn_output_proj_flops
+        + cross_attn_q_proj_flops
+        + cross_attn_kv_proj_flops
+        + cross_attention_qk_v_flops
+        + cross_attn_output_proj_flops
+    )
+
+    # SwiGLU FFN FLOPs (3 projections: gate, up, down)
+    ffn_flops = 3 * (2 * seq_len * hidden_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_low = self.checkpointer._create_optimizer(
+        pipeline.low_noise_transformer, self.config, self.config.learning_rate, scale_factor=self.config.boundary_ratio
+    )
+    optimizer_high, learning_rate_scheduler_high = self.checkpointer._create_optimizer(
+        pipeline.high_noise_transformer,
+        self.config,
+        self.config.learning_rate,
+        scale_factor=(1.0 - self.config.boundary_ratio),
+    )
+
+    pipeline = self.training_loop_2_2(
+        pipeline,
+        optimizer_low,
+        optimizer_high,
+        learning_rate_scheduler_low,
+        learning_rate_scheduler_high,
+        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_low,
+      learning_rate_scheduler_high,
+      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, ...)
+
+    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 opt_state_dict and isinstance(opt_state_dict, dict):
+          if opt_state_dict.get("low_noise_transformer") is not None:
+            state_low = state_low.replace(opt_state=opt_state_dict["low_noise_transformer"])
+          if opt_state_dict.get("high_noise_transformer") is not None:
+            state_high = state_high.replace(opt_state=opt_state_dict["high_noise_transformer"])
+          if opt_state_dict.get("low_noise_step") is not None:
+            state_low = state_low.replace(step=opt_state_dict["low_noise_step"])
+          if opt_state_dict.get("high_noise_step") is not None:
+            state_high = state_high.replace(step=opt_state_dict["high_noise_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:
+      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()
+
+        lr_low = learning_rate_scheduler_low(state_low.step)
+        lr_high = learning_rate_scheduler_high(state_high.step)
+        train_utils.record_scalar_metrics(
+            train_metric, last_step_completion - start_step_time, per_device_tflops, lr_low
+        )
+
+        if self.config.write_metrics:
+          train_metric["scalar"]["learning/current_learning_rate_low"] = lr_low
+          train_metric["scalar"]["learning/current_learning_rate_high"] = lr_high
+          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 self.config.checkpoint_every > 0 and (step + 1) % self.config.checkpoint_every == 0:
+          save_step = int(step + 1)
+          max_logging.log(f"Saving checkpoint for step {save_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(save_step, pipeline, train_states)
+
+      _metrics_queue.put(None)
+      writer_thread.join()
+      if writer:
+        writer.flush()
+      if self.config.save_final_checkpoint:
+        save_step = int(self.config.max_train_steps)
+        max_logging.log(f"Saving final checkpoint for step {save_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(save_step, 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):
+  """Wan 2.2 joint dual-expert training step.
+
+  Expert Routing:
+    Routing is performed at the batch level (one expert per step) based on `boundary_ratio`.
+    This prevents loading and executing both ~27B models simultaneously into TPU memory,
+    avoiding out-of-memory (OOM) conditions during large distributed runs.
+
+  Timestep Distributions:
+    - High-noise expert (coarse structure): samples timesteps from Beta(5.0, 2.0) scaled to [boundary, 1000],
+      biasing updates toward higher noise levels where coarse generation occurs.
+    - Low-noise expert (fine details): samples timesteps from Beta(2.0, 5.0) scaled to [0, boundary],
+      biasing updates toward lower noise levels where fine details are refined.
+  """
+  _, 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)
+
+  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 getattr(config, "disable_training_weights", False):
+      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 getattr(config, "disable_training_weights", False):
+      training_weight = jnp.expand_dims(training_weight, axis=(1, 2, 3, 4))
+      loss = loss * training_weight
+
+    loss_mean = jnp.mean(loss)
+    loss_per_example = loss.reshape(loss.shape[0], -1).mean(axis=1)
+    fine_mask = timesteps < 200
+    mid_mask = (timesteps >= 200) & (timesteps < 500)
+    coarse_mask = timesteps >= 500
+
+    loss_fine = jnp.where(jnp.any(fine_mask), jnp.sum(jnp.where(fine_mask, loss_per_example, 0.0)) / jnp.maximum(1, jnp.sum(fine_mask)), jnp.nan)
+    loss_mid = jnp.where(jnp.any(mid_mask), jnp.sum(jnp.where(mid_mask, loss_per_example, 0.0)) / jnp.maximum(1, jnp.sum(mid_mask)), jnp.nan)
+    loss_coarse = jnp.where(jnp.any(coarse_mask), jnp.sum(jnp.where(coarse_mask, loss_per_example, 0.0)) / jnp.maximum(1, jnp.sum(coarse_mask)), 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)
+    nan_val = jnp.array(jnp.nan, dtype=loss.dtype)
+
+    max_grad_norm_high = jaxopt.tree_util.tree_l2_norm(high_grads)
+    max_abs_grad_high = jax.tree_util.tree_reduce(lambda max_val, arr: jnp.maximum(max_val, jnp.max(jnp.abs(arr))), high_grads, initializer=-1.0)
+
+    new_state_high = state_high.apply_gradients(grads=high_grads)
+
+    return loss, nan_val, nan_val, nan_val, nan_val, nan_val, max_grad_norm_high, nan_val, max_abs_grad_high, new_state_high, state_low
+
+  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)
+    nan_val = jnp.array(jnp.nan, dtype=loss.dtype)
+
+    max_grad_norm_low = jaxopt.tree_util.tree_l2_norm(low_grads)
+    max_abs_grad_low = jax.tree_util.tree_reduce(lambda max_val, arr: jnp.maximum(max_val, jnp.max(jnp.abs(arr))), low_grads, initializer=-1.0)
+
+    new_state_low = state_low.apply_gradients(grads=low_grads)
+
+    return nan_val, loss, loss_fine, loss_mid, loss_coarse, max_grad_norm_low, nan_val, max_abs_grad_low, nan_val, state_high, new_state_low
+
+  (loss_high, loss_low, loss_fine, loss_mid, loss_coarse,
+   max_grad_norm_low, max_grad_norm_high, max_abs_grad_low, max_abs_grad_high,
+   new_state_high, new_state_low) = jax.lax.cond(
+      is_high_noise,
+      true_fn,
+      false_fn,
+      operand=(state_high.params, state_low.params)
+  )
+
+  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": {},
+  }
+
+  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)
+
+  model_high = nnx.merge(state_high.graphdef, state_high.params, state_high.rest_of_state)
+  model_low = nnx.merge(state_low.graphdef, state_low.params, state_low.rest_of_state)
+
+  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")
+    rng, new_rng = jax.random.split(rng, num=2)
+
+    noise = jax.random.normal(key=new_rng, shape=latents.shape, dtype=latents.dtype)
+    noisy_latents, training_target, training_weight = scheduler.apply_flow_match(noise, latents, timesteps)
+
+    def _eval_single(noisy_l, t, enc):
+      return jax.lax.cond(
+          t >= boundary,
+          lambda: model_high(hidden_states=noisy_l[None], timestep=t[None], encoder_hidden_states=enc[None], deterministic=True)[0],
+          lambda: model_low(hidden_states=noisy_l[None], timestep=t[None], encoder_hidden_states=enc[None], deterministic=True)[0],
+      )
+
+    model_pred = jax.vmap(_eval_single)(noisy_latents, timesteps, encoder_hidden_states)
+    loss = (training_target - model_pred) ** 2
+    if not getattr(config, "disable_training_weights", False):
+      training_weight = jnp.expand_dims(training_weight, axis=(1, 2, 3, 4))
+      loss = loss * training_weight
+    loss = loss.reshape(loss.shape[0], -1).mean(axis=1)
+    losses = losses.at[start:end].set(loss)
+
+  metrics = {"scalar": {"learning/eval_loss": losses}}
+  return metrics, new_rng