diff --git a/.dockerignore b/.dockerignore
index e567ea2ff6..03564ad8dc 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,2 +1,5 @@
 .git
 maxtext_venv
+.venv
+local_logs
+scratch
diff --git a/run_e2e.sh b/run_e2e.sh
new file mode 100755
index 0000000000..71d1ae28ed
--- /dev/null
+++ b/run_e2e.sh
@@ -0,0 +1,29 @@
+#!/bin/bash
+set -e
+
+echo "=========================================================================="
+echo "Starting MaxText End-to-End TPU Validation Matrix"
+echo "=========================================================================="
+echo "This script will orchestrate 224 training jobs across 8 models and 2 scan modes."
+echo "Models and Checkpoints will be saved to GCS: gs://mesa-maxtext/validation_runs/post_train_layout_v1/"
+echo "Execution Logs will be saved locally to: ./local_logs/"
+echo ""
+
+# Ensure we are in the correct MaxText directory
+if [ ! -f "src/maxtext/trainers/pre_train/train.py" ]; then
+    echo "[ERROR] Please run this script from the root of the maxtext repository."
+    exit 1
+fi
+
+echo "Cleaning up any lingering MaxText python processes to free the TPU..."
+pkill -f "python.*maxtext" || true
+
+# Run the python orchestrator
+python3 run_e2e_matrix.py
+
+echo "=========================================================================="
+echo "Validation Matrix execution completed."
+echo "Results written to validation_summary.csv"
+echo "To view summary:"
+echo "cat validation_summary.csv | column -t -s,"
+echo "=========================================================================="
diff --git a/run_e2e_matrix.py b/run_e2e_matrix.py
new file mode 100644
index 0000000000..420747baca
--- /dev/null
+++ b/run_e2e_matrix.py
@@ -0,0 +1,323 @@
+"""Drives the end-to-end checkpoint validation matrix across models and scan modes.
+
+Each combination trains a base checkpoint and then reloads it from every trainer, so a
+checkpoint written by one is proven readable by the others. Results land in a CSV.
+"""
+
+import os
+import subprocess
+import csv
+
+from maxtext.utils.globals import HF_IDS
+
+MODELS = [
+    "llama3.1-8b",
+    "gemma3-4b",
+    # "gemma2-2b",
+    # "gemma4-e2b",
+    # "qwen2.5-1.5b",
+    # "qwen3-0.6b",
+    # "olmo3-7b",
+    # "gpt-oss-20b",
+]
+SCAN_MODES = ["scanned", "unscanned"]
+
+GCS_BASE = "gs://mesa-maxtext/validation_runs/post_train_layout_v14"
+HF_BASE = "gs://mesa-maxtext/huggingface_transformers"
+LOCAL_LOGS = "local_logs"
+CSV_REPORT = "validation_summary.csv"
+
+
+def get_tokenizer_flags(model_name):
+  """Returns the tokenizer config flags a model needs, so no job falls back to a missing default."""
+  flags = []
+  if "gemma4" in model_name:
+    flags.extend(["tokenizer_path=src/maxtext/assets/tokenizers/tokenizer_gemma4.model", "tokenizer_type=sentencepiece"])
+  elif "gemma3" in model_name:
+    flags.extend(["tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma3", "tokenizer_type=sentencepiece"])
+  elif "gemma" in model_name:
+    flags.extend(["tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.gemma", "tokenizer_type=sentencepiece"])
+  elif "llama" in model_name:
+    flags.extend(["tokenizer_path=src/maxtext/assets/tokenizers/tokenizer_llama3.tiktoken", "tokenizer_type=tiktoken"])
+  elif "mistral" in model_name:
+    flags.extend(["tokenizer_path=src/maxtext/assets/tokenizers/tokenizer.mistral-v3", "tokenizer_type=sentencepiece"])
+  elif "qwen" in model_name or "olmo" in model_name or "gpt-oss" in model_name:
+    # These have no tokenizer asset of their own -- the path built from the model name does not
+    # exist, and a huggingface tokenizer_type then reads it as a repo id and fails. Name the repo.
+    flags.extend([f"tokenizer_path={HF_IDS[model_name]}", "tokenizer_type=huggingface"])
+  return flags
+
+
+def execute_command(cmd, log_path):
+  """Runs one job to completion, tee-ing its output to `log_path`.
+
+  Args:
+    cmd: Argument list to run.
+    log_path: File to write the command and its combined output to.
+
+  Returns:
+    Whether the job exited zero.
+  """
+  os.makedirs(os.path.dirname(log_path), exist_ok=True)
+  env = os.environ.copy()
+
+  cmd_str = " ".join(cmd)
+  print(f"\n[EXECUTING]: {cmd_str}")
+  print(f"[LOG PATH]: {log_path}")
+
+  with open(log_path, "w", encoding="utf-8") as f:
+    f.write(f"Command: {cmd_str}\n\n")
+    f.flush()
+    with subprocess.Popen(cmd, stdout=f, stderr=subprocess.STDOUT, env=env) as process:
+      process.wait()
+
+  if process.returncode != 0:
+    print(f"[ERROR] Job failed. Check logs at: {log_path}")
+  return process.returncode == 0
+
+
+def run_matrix():
+  """Runs every model and scan mode through the matrix, recording each job's result."""
+  with open(CSV_REPORT, "w", newline="", encoding="utf-8") as f:
+    writer = csv.writer(f)
+    writer.writerow(["Model", "Scan Mode", "Phase", "Action", "Run Name", "Status"])
+
+  # The helpers below close over the loop variables, but each is only called within the
+  # iteration that defines it, so the late-binding pylint warns about cannot happen.
+  # pylint: disable=cell-var-from-loop
+  for model in MODELS:
+    for scan_mode in SCAN_MODES:
+      print(f"\n{'='*80}\nStarting matrix for Model: {model} | Scan Mode: {scan_mode}\n{'='*80}")
+      scan_bool = "True" if scan_mode == "scanned" else "False"
+      hf_ckpt = f"{HF_BASE}/{model}/to_maxtext/{scan_mode}/0/items"
+
+      def record_result(phase, action, run_name, success):
+        status = "PASS" if success else "FAIL"
+        print(f"[{status}] {model} | {scan_mode} | {action} | {run_name}")
+        with open(CSV_REPORT, "a", newline="", encoding="utf-8") as f:
+          writer = csv.writer(f)
+          writer.writerow([model, scan_mode, phase, action, run_name, status])
+
+      def get_ckpt_path(action, run_name):
+        step = "0" if action in ("sft", "distill") else "1"
+        return f"{GCS_BASE}/{scan_mode}/{action}/{model}/{run_name}/checkpoints/{step}/items"
+
+      def build_cmd(script, action, run_name, load_path, extra_flags=None):
+        config = "src/maxtext/configs/post_train/rl.yml" if action == "rl" else "src/maxtext/configs/base.yml"
+        cmd = [
+            "python",
+            script,
+            config,
+            f"run_name={run_name}",
+            f"model_name={model}",
+            f"scan_layers={scan_bool}",
+            f"base_output_directory={GCS_BASE}/{scan_mode}/{action}/{model}",
+        ]
+        # Handle steps vs num_batches for RL
+        if action == "rl":
+          cmd.extend(["num_batches=2", "rl.num_generations=2"])
+          if "gemma3" in model:
+            cmd.append("chat_template_path=src/maxtext/examples/chat_templates/gemma-3-27b-chat_template.json")
+            hf_model_name = f"google/gemma-3-{model.rsplit('-', maxsplit=1)[-1]}-it"
+            cmd.append(f"vllm_hf_config_path={hf_model_name}")
+          else:
+            # RL needs a chat template, and the base checkpoints are not all instruction tuned,
+            # so their tokenizers do not ship one. The round trip under test does not care which
+            # template it is.
+            cmd.append("chat_template_path=src/maxtext/examples/chat_templates/gsm8k_rl.json")
+        else:
+          cmd.extend(["steps=5", "per_device_batch_size=1"])
+
+        # Handle datasets: DPO needs HF dataset, others can use synthetic to bypass tokenization
+        if action == "dpo":
+          cmd.extend(
+              [
+                  "dataset_type=hf",
+                  "hf_path=json",
+                  "hf_train_files=tests/assets/local_datasets/dpo/dpo_3_column_dataset.json",
+                  "train_data_columns=\"['prompt', 'chosen', 'rejected']\"",
+                  "tokenize_train_data=True",
+                  "use_dpo=True",
+                  "packing=False",
+              ]
+          )
+        elif action != "rl":
+          cmd.extend(["dataset_type=synthetic"])
+
+        # Inject tokenizers for ALL jobs to prevent missing tokenizer config errors
+        cmd.extend(get_tokenizer_flags(model))
+
+        if action == "rl":
+          # RL samples through vLLM, which wants the HuggingFace model, and applies a chat
+          # template, which MaxText's own tokenizer assets do not carry. Name the repo for both.
+          if "llama3" in model:
+            hf_model_name = f"meta-llama/Meta-Llama-3.1-{model.rsplit('-', maxsplit=1)[-1].upper()}-Instruct"
+          else:
+            hf_model_name = HF_IDS[model]
+          cmd = [c for c in cmd if not c.startswith("tokenizer_path=") and not c.startswith("tokenizer_type=")]
+          cmd.extend([f"tokenizer_path={hf_model_name}", "tokenizer_type=huggingface"])
+          if not any(c.startswith("vllm_hf_config_path=") for c in cmd):
+            cmd.append(f"vllm_hf_config_path={hf_model_name}")
+
+        if action == "dpo":
+          # DPO reads dataset_type=hf, whose pipeline tokenizes through AutoTokenizer. That cannot
+          # open MaxText's own tokenizer assets, so name the HF repo for the model instead.
+          cmd = [c for c in cmd if not c.startswith("tokenizer_path=") and not c.startswith("tokenizer_type=")]
+          cmd.extend([f"tokenizer_path={HF_IDS[model]}", "tokenizer_type=huggingface"])
+
+        if load_path:
+          cmd.append(f"load_parameters_path={load_path}")
+        if ("gemma" in model and scan_bool == "True") and (load_path and HF_BASE in load_path):
+          cmd.append("use_standalone_converter=True")
+
+        if extra_flags:
+          cmd.extend(extra_flags)
+
+        return cmd
+
+      # Script paths
+      pre_train_script = "src/maxtext/trainers/pre_train/train.py"
+      sft_script = "src/maxtext/trainers/post_train/sft/train_sft.py"
+      dpo_script = "src/maxtext/trainers/post_train/dpo/train_dpo.py"
+      rl_script = "src/maxtext/trainers/post_train/rl/train_rl.py"
+      distill_script = "src/maxtext/trainers/post_train/distillation/train_distill.py"
+
+      # ---------------------------------------------------------
+      # Phase A: Pre-train -> Post-train Reloading
+      # ---------------------------------------------------------
+
+      # 1. Generate Pre-train Base
+      run_name = "ckpt_pretrain_base"
+      cmd = build_cmd(pre_train_script, "pre_train", run_name, None)
+      log_path = f"{LOCAL_LOGS}/{scan_mode}/pre_train/{model}/{run_name}.log"
+      success = execute_command(cmd, log_path)
+      record_result("Phase A", "pre_train", run_name, success)
+      pt_ckpt = get_ckpt_path("pre_train", run_name)
+
+      if success:
+        # 2. Pre-train -> SFT Reload
+        run_name = "sft_reload_pt"
+        cmd = build_cmd(sft_script, "sft", run_name, pt_ckpt)
+        log_path = f"{LOCAL_LOGS}/{scan_mode}/sft/{model}/{run_name}.log"
+        success = execute_command(cmd, log_path)
+        record_result("Phase A", "sft", run_name, success)
+
+        # 3. Pre-train -> DPO Reload
+        run_name = "dpo_reload_pt"
+        cmd = build_cmd(dpo_script, "dpo", run_name, pt_ckpt)
+        log_path = f"{LOCAL_LOGS}/{scan_mode}/dpo/{model}/{run_name}.log"
+        success = execute_command(cmd, log_path)
+        record_result("Phase A", "dpo", run_name, success)
+
+        # 4. Pre-train -> RL Reload
+        run_name = "rl_reload_pt"
+        cmd = build_cmd(rl_script, "rl", run_name, pt_ckpt)
+        log_path = f"{LOCAL_LOGS}/{scan_mode}/rl/{model}/{run_name}.log"
+        success = execute_command(cmd, log_path)
+        record_result("Phase A", "rl", run_name, success)
+
+      # ---------------------------------------------------------
+      # Phase B: Post-train -> Pre-train / Self Reloading
+      # ---------------------------------------------------------
+
+      # SFT Validations
+      run_name = "ckpt_sft_base"
+      cmd = build_cmd(sft_script, "sft", run_name, hf_ckpt)
+      log_path = f"{LOCAL_LOGS}/{scan_mode}/sft/{model}/{run_name}.log"
+      success = execute_command(cmd, log_path)
+      record_result("Phase B", "sft", run_name, success)
+      sft_ckpt = get_ckpt_path("sft", run_name)
+
+      if success:
+        rn = "sft_intra_reload"
+        success2 = execute_command(
+            build_cmd(sft_script, "sft", rn, sft_ckpt), f"{LOCAL_LOGS}/{scan_mode}/sft/{model}/{rn}.log"
+        )
+        record_result("Phase B", "sft", rn, success2)
+
+        rn = "sft_cross_reload"
+        success2 = execute_command(
+            build_cmd(pre_train_script, "pre_train", rn, sft_ckpt), f"{LOCAL_LOGS}/{scan_mode}/pre_train/{model}/{rn}.log"
+        )
+        record_result("Phase B", "pre_train", rn, success2)
+
+      # DPO Validations
+      run_name = "ckpt_dpo_base"
+      cmd = build_cmd(dpo_script, "dpo", run_name, hf_ckpt)
+      log_path = f"{LOCAL_LOGS}/{scan_mode}/dpo/{model}/{run_name}.log"
+      success = execute_command(cmd, log_path)
+      record_result("Phase B", "dpo", run_name, success)
+      dpo_ckpt = get_ckpt_path("dpo", run_name)
+
+      if success:
+        rn = "dpo_intra_reload"
+        success2 = execute_command(
+            build_cmd(dpo_script, "dpo", rn, dpo_ckpt), f"{LOCAL_LOGS}/{scan_mode}/dpo/{model}/{rn}.log"
+        )
+        record_result("Phase B", "dpo", rn, success2)
+
+        rn = "dpo_cross_reload"
+        success2 = execute_command(
+            build_cmd(pre_train_script, "pre_train", rn, dpo_ckpt), f"{LOCAL_LOGS}/{scan_mode}/pre_train/{model}/{rn}.log"
+        )
+        record_result("Phase B", "pre_train", rn, success2)
+
+      # RL Validations
+      run_name = "ckpt_rl_base"
+      cmd = build_cmd(rl_script, "rl", run_name, hf_ckpt)
+      log_path = f"{LOCAL_LOGS}/{scan_mode}/rl/{model}/{run_name}.log"
+      success = execute_command(cmd, log_path)
+      record_result("Phase B", "rl", run_name, success)
+      rl_ckpt = get_ckpt_path("rl", run_name)
+
+      if success:
+        rn = "rl_intra_reload"
+        success2 = execute_command(
+            build_cmd(rl_script, "rl", rn, rl_ckpt), f"{LOCAL_LOGS}/{scan_mode}/rl/{model}/{rn}.log"
+        )
+        record_result("Phase B", "rl", rn, success2)
+
+        rn = "rl_cross_reload"
+        success2 = execute_command(
+            build_cmd(pre_train_script, "pre_train", rn, rl_ckpt), f"{LOCAL_LOGS}/{scan_mode}/pre_train/{model}/{rn}.log"
+        )
+        record_result("Phase B", "pre_train", rn, success2)
+
+      # Distill Validations
+      run_name = "ckpt_distill_base"
+      distill_flags = [
+          f"teacher_overrides.load_parameters_path={hf_ckpt}",
+          "teacher_overrides.model_name=" + model,
+          f"teacher_overrides.scan_layers={scan_bool}",
+          "teacher_overrides.per_device_batch_size=1",
+      ]
+      cmd = build_cmd(distill_script, "distill", run_name, hf_ckpt, distill_flags)
+      log_path = f"{LOCAL_LOGS}/{scan_mode}/distill/{model}/{run_name}.log"
+      success = execute_command(cmd, log_path)
+      record_result("Phase B", "distill", run_name, success)
+      distill_ckpt = get_ckpt_path("distill", run_name)
+
+      if success:
+        rn = "distill_intra_reload"
+        distill_flags2 = [
+            f"teacher_overrides.load_parameters_path={hf_ckpt}",  # teacher loads base again
+            "teacher_overrides.model_name=" + model,
+            f"teacher_overrides.scan_layers={scan_bool}",
+            "teacher_overrides.per_device_batch_size=1",
+        ]
+        success2 = execute_command(
+            build_cmd(distill_script, "distill", rn, distill_ckpt, distill_flags2),
+            f"{LOCAL_LOGS}/{scan_mode}/distill/{model}/{rn}.log",
+        )
+        record_result("Phase B", "distill", rn, success2)
+
+        rn = "distill_cross_reload"
+        success2 = execute_command(
+            build_cmd(pre_train_script, "pre_train", rn, distill_ckpt),
+            f"{LOCAL_LOGS}/{scan_mode}/pre_train/{model}/{rn}.log",
+        )
+        record_result("Phase B", "pre_train", rn, success2)
+
+
+if __name__ == "__main__":
+  run_matrix()
diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index 854c1b3968..261fad6976 100644
--- a/src/maxtext/common/checkpointing.py
+++ b/src/maxtext/common/checkpointing.py
@@ -30,6 +30,7 @@
 from flax.training import train_state
 from grain.experimental import ElasticIterator
 import jax
+import jax.numpy as jnp
 from maxtext.checkpoint_conversion.utils.load_dynamic import load_safetensors_dynamic_state
 from maxtext.common import emergency_checkpointing
 from maxtext.common import grain_utility
@@ -150,7 +151,7 @@ def _raise_on_weight_mismatch(want, have, config=None):
   )
 
 
-def _linen_items_to_nnx(restored_linen, abstract_nnx_state):
+def linen_items_to_nnx(restored_linen, abstract_nnx_state):
   """Reshapes a restored Linen-layout `items` dict into an NNX state.
 
   The inverse of `to_checkpoint_dict`, over the same `split_for_checkpoint` partition. The Linen
@@ -184,7 +185,7 @@ def _load_linen_checkpoint_into_nnx(
   """Restores a Linen-layout checkpoint into an NNX state (pure_nnx resume).
 
   Restores a Linen-shape target that includes `nnx_aux`, then reshapes back via
-  `_linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when
+  `linen_items_to_nnx`. rngs/dropout/batch stats come from `items/nnx_aux` when
   present, else keep their fresh init value. A genuinely-missing weight raises.
   """
   max_logging.log(f"Restoring Linen-layout checkpoint into NNX state at {path}")
@@ -222,7 +223,7 @@ def _restored_linen_to_nnx(restored_linen, abstract_nnx_state, config=None):
   itself is the Linen one, since pure_nnx reads and writes the Linen on-disk layout.
   """
   _raise_on_weight_mismatch(*_expected_and_restored_params(abstract_nnx_state, restored_linen), config=config)
-  return _linen_items_to_nnx(restored_linen, abstract_nnx_state)
+  return linen_items_to_nnx(restored_linen, abstract_nnx_state)
 
 
 def _abstract_params(abstract_unboxed_pre_state):
@@ -266,6 +267,190 @@ def _resolve_conversion_fn(checkpoint_conversion_fn):
   return fn
 
 
+def _build_tunix_target_tree(disk_meta, want_tree):
+  """Builds a target tree matching Tunix's on-disk metadata, injecting target shapes and shardings."""
+  if isinstance(disk_meta, dict):
+    if "value" in disk_meta and len(disk_meta) == 1:
+      if want_tree is not None:
+        if isinstance(want_tree, dict) and "value" in want_tree:
+          return want_tree
+        return {"value": want_tree}
+      else:
+        leaf = disk_meta["value"]
+        if hasattr(leaf, "shape") and hasattr(leaf, "dtype"):
+          return {"value": jax.ShapeDtypeStruct(shape=leaf.shape, dtype=leaf.dtype)}
+        return {"value": leaf}
+    res = {}
+    for k, v in disk_meta.items():
+      want_sub = want_tree.get(k) if (want_tree is not None and isinstance(want_tree, dict)) else None
+      res[k] = _build_tunix_target_tree(v, want_sub)
+    return res
+  elif hasattr(disk_meta, "shape") and hasattr(disk_meta, "dtype"):
+    if want_tree is not None:
+      return want_tree
+    return jax.ShapeDtypeStruct(shape=disk_meta.shape, dtype=disk_meta.dtype)
+  return disk_meta
+
+
+def _drop_adapter_level(tree):
+  """Strip the single-key `base` wrapper Tunix adds around adapter-free checkpoints."""
+  if isinstance(tree, dict):
+    if set(tree) == {"base"}:
+      return _drop_adapter_level(tree["base"])
+    return {k: _drop_adapter_level(v) for k, v in tree.items()}
+  if isinstance(tree, list):
+    return [_drop_adapter_level(v) for v in tree]
+  return tree
+
+
+def _drop_inject_hyperparams(opt_state):
+  """Unwrap the `inject_hyperparams` layer Tunix wraps around the optimizer state."""
+  if isinstance(opt_state, dict) and {"count", "hyperparams", "hyperparams_states", "inner_state"}.issubset(
+      opt_state.keys()
+  ):
+    return opt_state["inner_state"]
+  return opt_state
+
+
+def _load_tunix_full_state_from_path(
+    path,
+    abstract_unboxed_pre_state,
+    checkpoint_storage_concurrent_gb,
+    use_ocdbt,
+    use_zarr3,
+    maxtext_config=None,
+):
+  """Restore a Tunix-layout full training state (params + optimizer) into MaxText's structure."""
+  is_nnx = isinstance(abstract_unboxed_pre_state, nnx.State)
+  if is_nnx:
+    want_params = nnx.split_state(abstract_unboxed_pre_state.model, nnx.Param, ...)[0].to_pure_dict()
+    want_opt = abstract_unboxed_pre_state.optimizer
+  else:
+    want_params = abstract_unboxed_pre_state.params
+    want_opt = abstract_unboxed_pre_state.opt_state
+
+  ckptr = ocp.Checkpointer(
+      ocp.PyTreeCheckpointHandler(
+          restore_concurrent_gb=checkpoint_storage_concurrent_gb,
+          use_ocdbt=use_ocdbt,
+          use_zarr3=use_zarr3,
+      )
+  )
+
+  step_path = epath.Path(path)
+  model_params_path = step_path / "model_params" if (step_path / "model_params").exists() else step_path
+  opt_state_path = step_path / "optimizer_state"
+
+  has_base = False
+  has_inject = False
+  has_value_wrapper = True
+  disk_tree = None
+  try:
+    metadata = ckptr.metadata(step_path)
+    item_meta = getattr(metadata, "item_metadata", metadata)
+    if isinstance(item_meta, dict):
+      mp_meta = item_meta.get("model_params", item_meta)
+      disk_tree = getattr(mp_meta, "tree", mp_meta)
+      if isinstance(disk_tree, dict) and "base" in disk_tree:
+        has_base = True
+      opt_meta = item_meta.get("optimizer_state", None)
+      if opt_meta is not None:
+        tree_opt = getattr(opt_meta, "tree", opt_meta)
+        if isinstance(tree_opt, dict) and "inner_state" in tree_opt and "count" in tree_opt:
+          has_inject = True
+    else:
+      disk_tree = item_meta
+      if isinstance(disk_tree, dict) and "base" in disk_tree:
+        has_base = True
+  except Exception:  # pylint: disable=broad-exception-caught
+    pass
+
+  if disk_tree is None:
+    try:
+      mp_metadata = ckptr.metadata(model_params_path)
+      disk_tree = getattr(mp_metadata, "item_metadata", mp_metadata)
+      disk_tree = getattr(disk_tree, "tree", disk_tree)
+      if isinstance(disk_tree, dict) and "base" in disk_tree:
+        has_base = True
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass
+
+  is_linen_collection = (
+      (not is_nnx) and isinstance(want_params, dict) and "params" in want_params and len(want_params) == 1
+  )
+  inner_want_params = want_params["params"] if is_linen_collection else want_params
+
+  if disk_tree is not None and isinstance(disk_tree, dict):
+    target_disk_tree = disk_tree["base"] if has_base and "base" in disk_tree else disk_tree
+    target_want_params = _build_tunix_target_tree(target_disk_tree, inner_want_params)
+  else:
+    target_want_params = (
+        jax.tree.map(lambda v: {"value": v}, inner_want_params) if has_value_wrapper else inner_want_params
+    )
+
+  target_params = {"base": target_want_params} if has_base else target_want_params
+
+  if has_inject:
+    target_opt = {
+        "count": jnp.zeros((), dtype=jnp.int32),
+        "hyperparams": {},
+        "hyperparams_states": {},
+        "inner_state": want_opt,
+    }
+  else:
+    target_opt = want_opt
+
+  restore_args_params = ocp.checkpoint_utils.construct_restore_args(target_params)
+  restored_params = ckptr.restore(
+      model_params_path,
+      item=target_params,
+      restore_args=restore_args_params,
+  )
+
+  if opt_state_path.exists():
+    restore_args_opt = ocp.checkpoint_utils.construct_restore_args(target_opt)
+    restored_opt = ckptr.restore(
+        opt_state_path,
+        item=target_opt,
+        restore_args=restore_args_opt,
+    )
+  else:
+    restored_opt = want_opt
+
+  if has_base:
+    restored_params = _drop_adapter_level(restored_params)
+    if opt_state_path.exists():
+      restored_opt = _drop_adapter_level(restored_opt)
+
+  if has_inject:
+    restored_opt = _drop_inject_hyperparams(restored_opt)
+
+  if has_value_wrapper:
+    restored_params = train_state_nnx._strip_rng_state(restored_params)
+    restored_params = jax.tree.map(
+        lambda v: v["value"] if isinstance(v, dict) and "value" in v else v,
+        restored_params,
+        is_leaf=lambda x: isinstance(x, dict) and "value" in x and not isinstance(x.get("value"), dict),
+    )
+
+  if is_linen_collection:
+    restored_params_collection = {"params": restored_params}
+  else:
+    restored_params_collection = restored_params
+
+  _raise_on_weight_mismatch(want_params, restored_params_collection, config=maxtext_config)
+
+  if is_nnx:
+    nnx.replace_by_pure_dict(abstract_unboxed_pre_state, restored_params)
+    nnx.replace_by_pure_dict(abstract_unboxed_pre_state, {"optimizer": restored_opt})
+    return abstract_unboxed_pre_state
+  else:
+    return abstract_unboxed_pre_state.replace(
+        params=restored_params_collection,
+        opt_state=restored_opt,
+    )
+
+
 def _load_full_state_from_path(
     path,
     abstract_unboxed_pre_state,
@@ -297,6 +482,21 @@ def _load_full_state_from_path(
     The loaded state.
   """
 
+  if source_checkpoint_layout == "orbax":
+    if (epath.Path(path) / "model_params").exists() and (epath.Path(path) / "optimizer_state").exists():
+      max_logging.log(f"Auto-detected Tunix checkpoint layout at {path}")
+      source_checkpoint_layout = "tunix"
+
+  if source_checkpoint_layout == "tunix":
+    return _load_tunix_full_state_from_path(
+        path,
+        abstract_unboxed_pre_state,
+        checkpoint_storage_concurrent_gb,
+        use_ocdbt,
+        use_zarr3,
+        maxtext_config,
+    )
+
   if enable_orbax_v1:
     if source_checkpoint_layout == "orbax":
       # pure_nnx saves in the Linen on-disk layout; reshape it back into the NNX state.
@@ -661,6 +861,7 @@ def map_to_pspec(data):
         checkpoint_storage_concurrent_gb,
         use_ocdbt=use_ocdbt,
         use_zarr3=use_zarr3,
+        maxtext_config=maxtext_config,
     )
     return None, restored_params
   elif load_full_state_from_path != "":
@@ -701,22 +902,134 @@ def setup_checkpoint_logger(config) -> Any | None:  # pytype: disable=attribute-
   return orbax_cloud_logger
 
 
+def _nnx_native_wrapper_key(ckptr, path):
+  """Returns the wrapper key an NNX-native checkpoint nests its weights under, if any.
+
+  Post-training wrote DPO and RL checkpoints through `TunixMaxTextAdapter`, whose only child
+  module is `base`, so their weights sit one level deeper. SFT and the training engine save the
+  model directly and do not.
+
+  Args:
+    ckptr: Checkpointer used to read the checkpoint metadata.
+    path: Path to the checkpoint.
+
+  Returns:
+    "base" if the checkpoint nests its weights under it, else None.
+  """
+  tree = ckptr.metadata(epath.Path(path)).item_metadata
+  tree = getattr(tree, "tree", tree)
+  return "base" if tree is not None and "base" in tree else None
+
+
 def load_params_from_path(
     load_parameters_from_path,
     abstract_unboxed_params,
     checkpoint_storage_concurrent_gb,
     use_ocdbt=True,
     use_zarr3=True,
+    maxtext_config=None,
 ):
   """Load decode params from checkpoint at specified path."""
   assert load_parameters_from_path, "load_parameters_from_path is not defined."
   max_logging.log(f"restoring params from {load_parameters_from_path}")
 
+  # Check if Tunix Layout (either pointing to step dir or model_params dir)
+  is_tunix = False
+  target_path = None
+  path_obj = epath.Path(load_parameters_from_path)
+  if path_obj.name in ("model_params", "model"):
+    is_tunix = True
+    target_path = path_obj.parent
+  elif (path_obj / "model_params").exists():
+    is_tunix = True
+    target_path = path_obj
+
+  is_nnx = isinstance(abstract_unboxed_params, nnx.State)
+  want = abstract_unboxed_params.to_pure_dict() if is_nnx else abstract_unboxed_params
+
+  if is_tunix:
+    max_logging.log(f"Detected Tunix layout for parameters at {target_path}")
+    model_params_path = target_path / "model_params" if (target_path / "model_params").exists() else target_path
+
+    ckptr = ocp.Checkpointer(
+        ocp.PyTreeCheckpointHandler(
+            restore_concurrent_gb=checkpoint_storage_concurrent_gb,
+            use_ocdbt=use_ocdbt,
+            use_zarr3=use_zarr3,
+        )
+    )
+
+    has_base = False
+    has_value_wrapper = True  # Tunix checkpoints use NNX layout with {"value": Array} per leaf
+    disk_tree = None
+    try:
+      metadata = ckptr.metadata(target_path)
+      item_meta = getattr(metadata, "item_metadata", metadata)
+      if isinstance(item_meta, dict):
+        mp_meta = item_meta.get("model_params", item_meta)
+        disk_tree = getattr(mp_meta, "tree", mp_meta)
+      else:
+        disk_tree = item_meta
+      if isinstance(disk_tree, dict) and "base" in disk_tree:
+        has_base = True
+    except Exception:  # pylint: disable=broad-exception-caught
+      pass
+
+    if disk_tree is None:
+      try:
+        mp_metadata = ckptr.metadata(model_params_path)
+        disk_tree = getattr(mp_metadata, "item_metadata", mp_metadata)
+        disk_tree = getattr(disk_tree, "tree", disk_tree)
+        if isinstance(disk_tree, dict) and "base" in disk_tree:
+          has_base = True
+      except Exception:  # pylint: disable=broad-exception-caught
+        pass
+
+    is_linen_collection = (not is_nnx) and isinstance(want, dict) and "params" in want and len(want) == 1
+    inner_want = want["params"] if is_linen_collection else want
+
+    if disk_tree is not None and isinstance(disk_tree, dict):
+      target_disk_tree = disk_tree["base"] if has_base and "base" in disk_tree else disk_tree
+      target_want = _build_tunix_target_tree(target_disk_tree, inner_want)
+    else:
+      target_want = (
+          jax.tree.map(lambda v: {"value": v}, inner_want) if has_value_wrapper else inner_want
+      )
+
+    target_params = {"base": target_want} if has_base else target_want
+    restore_args = ocp.checkpoint_utils.construct_restore_args(target_params)
+
+    restored_weights = ckptr.restore(
+        model_params_path,
+        item=target_params,
+        restore_args=restore_args,
+    )
+
+    if has_base:
+      restored_weights = _drop_adapter_level(restored_weights)
+
+    if has_value_wrapper:
+      restored_weights = train_state_nnx._strip_rng_state(restored_weights)
+      restored_weights = jax.tree.map(
+          lambda v: v["value"] if isinstance(v, dict) and "value" in v else v,
+          restored_weights,
+          is_leaf=lambda x: isinstance(x, dict) and "value" in x and not isinstance(x.get("value"), dict),
+      )
+
+    if is_linen_collection:
+      restored_collection = {"params": restored_weights}
+    else:
+      restored_collection = restored_weights
+
+    _raise_on_weight_mismatch(want, restored_collection, config=maxtext_config)
+    if is_nnx:
+      nnx.replace_by_pure_dict(abstract_unboxed_params, restored_weights)
+      return abstract_unboxed_params
+    return restored_collection
+
   # On disk the weights live at `params/params/...`: an outer key naming the item, and Flax's
   # `params` collection inside it. A Linen TrainState.params is that collection; an NNX params
   # state sits one level below it (bare weights), so wrap it going in and unwrap it coming out.
-  is_nnx = isinstance(abstract_unboxed_params, nnx.State)
-  want = abstract_unboxed_params.to_pure_dict() if is_nnx else abstract_unboxed_params
 
   # Determine the restore key based on the leaf directory name to support native and custom SFT
   restore_key = os.path.basename(load_parameters_from_path)
@@ -743,25 +1056,37 @@ def load_params_from_path(
   # Rather than pass the entire abstract state, which could unnecessarily restore opt_state and such and waste
   # memory, we instead specify here that we are just restoring the params field of the checkpoint
   # (which itself may be a dictionary containing a key named 'params' or 'model').
-  restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection)
-  restored = ckptr.restore(
-      epath.Path(load_parameters_from_path),
-      item={restore_key: params_collection},
-      transforms={},
-      restore_args={restore_key: restore_args},
-  )
-  restored_collection = restored[restore_key]
-
   if restore_key in ("model_params", "model"):
-    restored_weights = restored_collection
+    wrapper_key = _nnx_native_wrapper_key(ckptr, load_parameters_from_path)
+    # Restore into the NNX state itself rather than a pure dict. Flax registers a Variable as a
+    # pytree holding its array under `value`, matching what `nnx.state(model)` wrote, so save and
+    # restore agree without reshaping either side by hand.
+    item = {wrapper_key: abstract_unboxed_params} if wrapper_key else abstract_unboxed_params
+    restored = ckptr.restore(
+        epath.Path(load_parameters_from_path),
+        item=item,
+        transforms={},
+        restore_args=ocp.checkpoint_utils.construct_restore_args(item),
+    )
+    # No `params` collection in this layout, so the weights are the whole collection.
+    restored_weights = nnx.to_pure_dict(restored[wrapper_key] if wrapper_key else restored)
+    restored_collection = restored_weights
   else:
+    restore_args = ocp.checkpoint_utils.construct_restore_args(params_collection)
+    restored = ckptr.restore(
+        epath.Path(load_parameters_from_path),
+        item={"params": params_collection},
+        transforms={},
+        restore_args={"params": restore_args},
+    )
+    restored_collection = restored["params"]
     restored_weights = restored_collection["params"] if is_nnx else restored_collection
 
   # `transforms={}` lets Orbax return an unmaterialized leaf for a weight the checkpoint lacks,
   # and a stored array at its own shape rather than the target's. Either reaches the model and
   # fails much later without naming the weight, so check here -- the params-only load
   # (load_parameters_path, e.g. SFT) has no init state to fall back on.
-  _raise_on_weight_mismatch(want, restored_weights)
+  _raise_on_weight_mismatch(want, restored_weights, config=maxtext_config)
   if is_nnx:
     nnx.replace_by_pure_dict(abstract_unboxed_params, restored_weights)
     return abstract_unboxed_params
@@ -777,26 +1102,66 @@ def save_params_to_path(checkpoint_dir, params, use_ocdbt=True, use_zarr3=True):
   print(f"Quantized params checkpoint saved at: {checkpoint_dir}")
 
 
-def load_checkpoint_metadata(checkpoint_dir_path: str) -> dict[str, Any]:
-  """Loads custom metadata from an Orbax checkpoint.
+def checkpoint_custom_metadata(config) -> dict[str, Any]:
+  """Returns the metadata a checkpoint stores alongside its state.
+
+  `verify_and_sync_scan_layers` and `lora_utils.sync_lora_metadata` read these back on load.
+  Post-training saves through its own manager, so it calls this too.
 
   Args:
-    checkpoint_dir_path: Path to the checkpoint directory.
+    config: The run's config, or None.
 
   Returns:
-    A dictionary containing custom metadata, or an empty dictionary if none is
-    present or loading fails.
+    The metadata dict, empty if there is no config to read it from.
+  """
+  custom_metadata = {}
+  if config:
+    if hasattr(config, "scan_layers"):
+      custom_metadata["scan_layers"] = config.scan_layers
+    if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0:
+      custom_metadata["lora"] = config.lora.model_dump()
+  return custom_metadata
+
+
+def _custom_metadata_at(checkpoint_dir: epath.Path) -> dict[str, Any]:
+  """Reads the custom metadata stored at exactly this directory.
+
+  Args:
+    checkpoint_dir: Directory to read.
+
+  Returns:
+    The metadata dict, empty if there is none or the read fails.
   """
-  checkpoint_dir = epath.Path(checkpoint_dir_path)
   try:
-    ckptr = ocp.StandardCheckpointer()
-    metadata = ckptr.metadata(checkpoint_dir)
+    metadata = ocp.StandardCheckpointer().metadata(checkpoint_dir)
     return metadata.custom_metadata or {}
   except Exception as e:  # pylint: disable=broad-except
     max_logging.log(f"Warning: Failed to load checkpoint metadata: {e}")
     return {}
 
 
+def load_checkpoint_metadata(checkpoint_dir_path: str) -> dict[str, Any]:
+  """Loads custom metadata from an Orbax checkpoint.
+
+  The metadata belongs to the step, so it sits at `<step>/` and not at `<step>/items/`. Callers
+  pass `load_parameters_path`, which points at the item, so fall back to the parent directory.
+
+  Args:
+    checkpoint_dir_path: Path to the checkpoint directory, item level or step level.
+
+  Returns:
+    A dictionary containing custom metadata, or an empty dictionary if none is
+    present or loading fails.
+  """
+  checkpoint_dir = epath.Path(checkpoint_dir_path)
+  metadata = _custom_metadata_at(checkpoint_dir)
+  if not metadata and checkpoint_dir.parent != checkpoint_dir:
+    metadata = _custom_metadata_at(checkpoint_dir.parent)
+  if not metadata and (checkpoint_dir / "model_params").exists():
+    metadata = _custom_metadata_at(checkpoint_dir / "model_params")
+  return metadata
+
+
 def _uses_local_checkpoint_period(config):
   return config.enable_emergency_checkpoint or config.enable_multi_tier_checkpointing
 
@@ -1030,12 +1395,7 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
           item=grain_iters_to_save
       )  # pyrefly: ignore[bad-assignment]
 
-  custom_metadata = {}
-  if config:
-    if hasattr(config, "scan_layers"):
-      custom_metadata["scan_layers"] = config.scan_layers
-    if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0:
-      custom_metadata["lora"] = config.lora.model_dump()
+  custom_metadata = checkpoint_custom_metadata(config)
 
   match (checkpoint_manager, config, data_iterator):
     case (checkpoint_manager, _, _) if isinstance(
diff --git a/src/maxtext/common/train_state_nnx.py b/src/maxtext/common/train_state_nnx.py
index 45dc386576..86b6c68654 100644
--- a/src/maxtext/common/train_state_nnx.py
+++ b/src/maxtext/common/train_state_nnx.py
@@ -120,6 +120,22 @@ def _as_chain_index(key):
   return None
 
 
+def opt_state_to_linen(opt_state):
+  """Reshapes an optimizer state into the Linen on-disk layout.
+
+  `to_checkpoint_dict` already does this for the state it is given. This is for a caller that
+  has unwrapped something the conversion could not see through -- `optax.inject_hyperparams`
+  hides mu and nu behind its own keys -- and needs the inner state converted on its own.
+
+  Args:
+    opt_state: The optimizer state to reshape.
+
+  Returns:
+    The same state in the Linen layout.
+  """
+  return _opt_state_to_linen(opt_state)
+
+
 def _opt_state_to_linen(opt_state):
   """Reshapes the NNX opt_state to Linen's on-disk layout.
 
diff --git a/src/maxtext/eval/README.md b/src/maxtext/eval/README.md
index 7970ce63b2..5eb0a9d889 100644
--- a/src/maxtext/eval/README.md
+++ b/src/maxtext/eval/README.md
@@ -186,7 +186,7 @@ Example (Qwen3-30B-A3B, v6e-8):
 STEP=244
 MODEL=qwen3-30b-a3b
 HF_PATH=Qwen/Qwen3-30B-A3B
-CHECKPOINT=gs://<bucket>/run/checkpoints/actor/${STEP}/model_params
+CHECKPOINT=gs://<bucket>/run/checkpoints/actor/${STEP}/items
 OUTPUT=gs://<bucket>/eval/
 
 python -m maxtext.eval.runner.run \
diff --git a/src/maxtext/examples/rl_llama3_demo.ipynb b/src/maxtext/examples/rl_llama3_demo.ipynb
index 68a9ff95eb..4615c1207b 100644
--- a/src/maxtext/examples/rl_llama3_demo.ipynb
+++ b/src/maxtext/examples/rl_llama3_demo.ipynb
@@ -300,13 +300,13 @@
     "# Define the output directory for the Hugging Face checkpoint\n",
     "hf_output_directory = epath.Path(BASE_OUTPUT_DIRECTORY) / \"hf_checkpoint\"\n",
     "\n",
-    "# Find the latest MaxText checkpoint\n",
-    "checkpoint_dir = epath.Path(config.checkpoint_dir) / 'actor'\n",
+    "# Find the latest MaxText checkpoint. RL checkpoints the actor under its own subdirectory.\n",
+    "checkpoint_dir = epath.Path(config.checkpoint_dir) / \"actor\"\n",
     "step_dirs = [d.name for d in checkpoint_dir.iterdir() if d.name.isdigit() and d.is_dir()]\n",
     "if not step_dirs:\n",
     "    raise ValueError(f\"No checkpoint found in {checkpoint_dir}\")\n",
     "latest_step = max(step_dirs, key=int)\n",
-    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"model_params\"\n",
+    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"items\"\n",
     "\n",
     "print(f\"Converting MaxText checkpoint from: {maxtext_checkpoint_path}\")\n",
     "print(f\"Saving Hugging Face checkpoint to: {hf_output_directory}\")\n",
diff --git a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb
index 216769ac48..2001622fe7 100644
--- a/src/maxtext/examples/sft_llama3_demo_tpu.ipynb
+++ b/src/maxtext/examples/sft_llama3_demo_tpu.ipynb
@@ -341,7 +341,7 @@
     "if not step_dirs:\n",
     "    raise ValueError(f\"No checkpoint found in {checkpoint_dir}\")\n",
     "latest_step = max(step_dirs, key=int)\n",
-    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"model_params\"\n",
+    "maxtext_checkpoint_path = checkpoint_dir / latest_step / \"items\"\n",
     "\n",
     "print(f\"Converting MaxText checkpoint from: {maxtext_checkpoint_path}\")\n",
     "print(f\"Saving Hugging Face checkpoint to: {hf_output_directory}\")\n",
diff --git a/src/maxtext/input_pipeline/hf_data_processing.py b/src/maxtext/input_pipeline/hf_data_processing.py
index 56d9ef0498..89a4954202 100644
--- a/src/maxtext/input_pipeline/hf_data_processing.py
+++ b/src/maxtext/input_pipeline/hf_data_processing.py
@@ -33,10 +33,15 @@
 
 
 def _get_pad_id(tokenizer):
-  if tokenizer.pad_token_id is not None:
+  """Returns the pad token id from the tokenizer, or -1 if not found."""
+  if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is not None:
     pad_id = tokenizer.pad_token_id
-  elif tokenizer.unk_token_id is not None:
+  elif hasattr(tokenizer, "unk_token_id") and tokenizer.unk_token_id is not None:
     pad_id = tokenizer.unk_token_id
+  elif hasattr(tokenizer, "pad_id") and getattr(tokenizer, "pad_id", None) is not None:
+    pad_id = tokenizer.pad_id
+  elif hasattr(tokenizer, "unk_id") and getattr(tokenizer, "unk_id", None) is not None:
+    pad_id = tokenizer.unk_id
   else:
     pad_id = -1
   return pad_id
diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py
index 8517fc4253..08b7327290 100644
--- a/src/maxtext/input_pipeline/input_pipeline_utils.py
+++ b/src/maxtext/input_pipeline/input_pipeline_utils.py
@@ -316,13 +316,25 @@ def apply_chat_template(example, tokenizer_model, data_column_name):
 
 def tokenization(example, hf_tokenizer, truncation, max_length, column_names):
   """Tokenize a HuggingFace dataset"""
+  is_hf = callable(hf_tokenizer)
   for column_name in column_names:
     if isinstance(example[column_name], list):
-      example[column_name] = [
-          hf_tokenizer(x, truncation=truncation, max_length=max_length)["input_ids"] for x in example[column_name]
-      ]
+      if is_hf:
+        example[column_name] = [
+            hf_tokenizer(x, truncation=truncation, max_length=max_length)["input_ids"] for x in example[column_name]
+        ]
+      else:
+        example[column_name] = [
+            hf_tokenizer.encode(x)[:max_length] if truncation else hf_tokenizer.encode(x) for x in example[column_name]
+        ]
     elif isinstance(example[column_name], str):
-      example[column_name] = hf_tokenizer(example[column_name], truncation=truncation, max_length=max_length)["input_ids"]
+      if is_hf:
+        example[column_name] = hf_tokenizer(example[column_name], truncation=truncation, max_length=max_length)[
+            "input_ids"
+        ]
+      else:
+        ids = hf_tokenizer.encode(example[column_name])
+        example[column_name] = ids[:max_length] if truncation else ids
   return example
 
 
diff --git a/src/maxtext/input_pipeline/synthetic_data_processing.py b/src/maxtext/input_pipeline/synthetic_data_processing.py
index 7c79eebfcb..9cbfbb6ab4 100644
--- a/src/maxtext/input_pipeline/synthetic_data_processing.py
+++ b/src/maxtext/input_pipeline/synthetic_data_processing.py
@@ -73,14 +73,14 @@ def _make_packed_segment_ids(batch_size: int, seq_len: int, max_segments_per_seq
 class SyntheticDataIterator:
   """Creates a synthetic data iterator for performance testing work"""
 
-  data_generator: Callable[[pyconfig.HyperParameters, tuple[Any, ...]], dict]
+  data_generator: Callable[[tuple[Any, ...]], dict]
 
   def __init__(self, config, mesh):
     self.mesh = mesh
     self.config = config
     data_pspec_shardings = sharding.get_input_data_sharding(config, mesh)
     self.data_generator = jax.jit(
-        SyntheticDataIterator.raw_generate_synthetic_data, out_shardings=data_pspec_shardings, static_argnums=0
+        SyntheticDataIterator.raw_generate_synthetic_data, out_shardings=data_pspec_shardings
     )
 
     tokens = jax.random.randint(
@@ -113,10 +113,10 @@ def __iter__(self):
 
   def __next__(self):
     with self.mesh:
-      return self.data_generator(self.config, self.data)  # pylint: disable=not-callable
+      return self.data_generator(self.data)  # pylint: disable=not-callable
 
   @staticmethod
-  def raw_generate_synthetic_data(config: pyconfig.HyperParameters, data):
+  def raw_generate_synthetic_data(data):
     """Generates a single batch of synthetic data"""
     tokens, positions, segmentation = data
 
diff --git a/src/maxtext/input_pipeline/tokenizer.py b/src/maxtext/input_pipeline/tokenizer.py
index 10a528205d..502aab6d83 100644
--- a/src/maxtext/input_pipeline/tokenizer.py
+++ b/src/maxtext/input_pipeline/tokenizer.py
@@ -23,7 +23,53 @@
 from sentencepiece import SentencePieceProcessor
 
 
-class TikTokenTokenizer:
+class ChatTemplateMixin:
+  """Mixin to provide Jinja2 chat template rendering for native tokenizers."""
+
+  def apply_chat_template(self, conversation, chat_template=None, add_generation_prompt=False, tokenize=True, **kwargs):
+    """Applies a Jinja2 chat template to a conversation."""
+    if chat_template is None:
+      chat_template = getattr(self, "chat_template", None)
+    if chat_template is None:
+      raise ValueError("Cannot apply chat template because no chat template was provided or set.")
+
+    import jinja2  # pylint: disable=import-outside-toplevel
+
+    env = jinja2.Environment(autoescape=False)
+
+    def raise_exception(message):
+      raise jinja2.exceptions.TemplateError(message)
+
+    env.globals["raise_exception"] = raise_exception
+
+    template = env.from_string(chat_template)
+
+    bos_token = ""
+    eos_token = ""
+    if hasattr(self, "_tokenizer_model"):  # SentencePiece
+      if self.bos_id is not None and self.bos_id >= 0:
+        bos_token = self._tokenizer_model.IdToPiece(self.bos_id)
+      if self.eos_id is not None and self.eos_id >= 0:
+        eos_token = self._tokenizer_model.IdToPiece(self.eos_id)
+    elif hasattr(self, "model"):  # TikToken
+      if self.bos_id is not None and self.bos_id >= 0:
+        bos_token = self.decode([self.bos_id])
+      if self.eos_id is not None and self.eos_id >= 0:
+        eos_token = self.decode([self.eos_id])
+
+    rendered = template.render(
+        messages=conversation,
+        add_generation_prompt=add_generation_prompt,
+        bos_token=bos_token,
+        eos_token=eos_token,
+        **kwargs,
+    )
+    if tokenize:
+      return self.encode(rendered)
+    return rendered
+
+
+class TikTokenTokenizer(ChatTemplateMixin):
   """
   Tokenizing and encoding/decoding text using the Tiktoken tokenizer.
   """
@@ -180,7 +226,7 @@ def _split_whitespaces_or_nonwhitespaces(s: str, max_consecutive_slice_len: int)
     yield s[slice_start:]
 
 
-class SentencePieceTokenizer:
+class SentencePieceTokenizer(ChatTemplateMixin):
   """
   Tokenizing and encoding/decoding text using the native sentencepiece library.
   Supports both local and GCS (gs://) model paths.
@@ -245,6 +291,9 @@ def encode(self, s: str) -> list[int]:
   def decode(self, t: Sequence[int]) -> str:
     return self.tokenizer.decode(t)
 
+  def apply_chat_template(self, *args, **kwargs):
+    return self.tokenizer.apply_chat_template(*args, **kwargs)
+
 
 def build_tokenizer(tokenizer_path, tokenizer_type, add_bos, add_eos, hf_access_token):
   """Loads the tokenizer at `tokenizer_path`"""
diff --git a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
index b272bf79eb..8ab9709358 100644
--- a/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
+++ b/src/maxtext/integration/vllm/maxtext_vllm_rollout.py
@@ -52,6 +52,56 @@
 # entry whose value is None", which means direct-sync-only.
 _NO_RULE_TABLE = object()
 
+import vllm.config.vllm
+import vllm.config.utils as vllm_config_utils
+
+# Monkey-patch VLLM's is_init_field to gracefully handle dynamically added
+# fields (like sharding_config injected by tpu-inference) which would otherwise
+# cause ValueError/AssertionError during vllm_config.with_hf_config replacement.
+_orig_is_init_field = vllm_config_utils.is_init_field
+
+
+def _patched_is_init_field(cls, name):
+  try:
+    return _orig_is_init_field(cls, name)
+  except ValueError:
+    return False
+
+
+vllm_config_utils.is_init_field = _patched_is_init_field
+
+
+_orig_with_hf_config = vllm.config.vllm.VllmConfig.with_hf_config
+
+
+def _patched_with_hf_config(self, *args, **kwargs):
+  """
+  Restore the original data_parallel_size which tpu_platform mutated,
+  so that the new VllmConfig passes the device_indexes length assertion.
+  tpu_inference deletes sharding_config before calling with_hf_config,
+  so we must reverse-engineer the data_parallel_size from device_indexes.
+  """
+  if self.additional_config and "sharding" in self.additional_config:
+    sharding_strategy = self.additional_config["sharding"].get("sharding_strategy", {})
+    device_indexes = sharding_strategy.get("device_indexes")
+    if device_indexes is not None:
+      pc = self.parallel_config
+      tp = sharding_strategy.get("tensor_parallelism") or pc.tensor_parallel_size
+      ep = sharding_strategy.get("expert_parallelism", 1)
+      sp = sharding_strategy.get("sequence_parallelism", 1)
+      attn_dp = sharding_strategy.get("attention_data_parallelism", 1)
+      attn_dp_ep = sharding_strategy.get("attention_data_expert_parallelism", 1)
+      dcp = pc.decode_context_parallel_size
+
+      other_parallelism = tp * ep * sp * attn_dp * attn_dp_ep * dcp
+      if other_parallelism > 0:
+        self.parallel_config.data_parallel_size = len(device_indexes) // other_parallelism
+
+  return _orig_with_hf_config(self, *args, **kwargs)
+
+
+vllm.config.vllm.VllmConfig.with_hf_config = _patched_with_hf_config
+
 
 def _rule_table_for(model_name: str):
   """Maps a MaxText model name to its torchax rule table.
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 1a9fdd48b0..bfb942d85d 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -438,7 +438,7 @@ def __init__(
     self.is_gemma4 = self.config.decoder_block == DecoderBlockType.GEMMA4
     self.is_gemma4_small = self.config.decoder_block == DecoderBlockType.GEMMA4_SMALL
 
-    if config.mhc_expansion_rate > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4:
+    if getattr(config, "mhc_expansion_rate", 1) > 1 and config.decoder_block == DecoderBlockType.DEEPSEEK4:
       self.hc_head = mhc.DeepSeek4HyperHead(
           config=config,
           mesh=self.mesh,
diff --git a/src/maxtext/optimizers/optimizers.py b/src/maxtext/optimizers/optimizers.py
index 67e1f589ca..23d1bfa98f 100644
--- a/src/maxtext/optimizers/optimizers.py
+++ b/src/maxtext/optimizers/optimizers.py
@@ -166,6 +166,33 @@ def skip_update():
   return optax.GradientTransformationExtraArgs(init_fn, update_fn)
 
 
+def add_gradient_clipping(tx, clipping_threshold):
+  """Clips gradients by global norm ahead of `tx`, keeping `tx`'s optimizer state shape.
+
+  `optax.chain(clip_by_global_norm(...), tx)` would nest tx's state under an extra chain level
+  even though the clip is stateless. Pre-training clips raw gradients in its train step instead,
+  so its checkpointed state has no such level and cannot restore one that does.
+
+  Args:
+    tx: The optimizer to clip gradients for.
+    clipping_threshold: Global norm to clip to.
+
+  Returns:
+    A transformation applying the same updates as the chained form, with tx's state tree.
+  """
+  clip = optax.clip_by_global_norm(clipping_threshold)
+  inner = optax.with_extra_args_support(tx)
+
+  def init_fn(params):
+    return inner.init(params)
+
+  def update_fn(updates, state, params=None, **extra_args):
+    updates, _ = clip.update(updates, optax.EmptyState(), params)
+    return inner.update(updates, state, params, **extra_args)
+
+  return optax.GradientTransformationExtraArgs(init_fn, update_fn)
+
+
 def get_optimizer(config, learning_rate_schedule, model=None):
   """Create optimizer."""
   if config.opt_type == "adamw":
diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py
index 725d5cd48a..1b64b4b2ad 100644
--- a/src/maxtext/trainers/post_train/rl/train_rl.py
+++ b/src/maxtext/trainers/post_train/rl/train_rl.py
@@ -139,7 +139,7 @@ def _compat_unstack(src_val, tgt_val, key_path, scan_axis=None):
 from maxtext.common.common_types import DecoderBlockType
 from maxtext.configs import pyconfig, types
 from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR
-from maxtext.integration.vllm.maxtext_vllm_rollout import MaxTextVllmRollout
+from maxtext.integration.vllm.maxtext_vllm_rollout import MaxTextVllmRollout, requires_maxtext_scanned_weight_unroll
 from maxtext.trainers.post_train.rl.evaluate_rl import evaluate
 from maxtext.trainers.post_train.rl import utils_rl
 from maxtext.input_pipeline.instruction_data_processing import load_data_template_from_file
@@ -539,6 +539,7 @@ def create_rl_components(  # pylint: disable=too-many-positional-arguments
               "hf_overrides": trainer_config.vllm_hf_overrides,
               "enable_expert_parallel": sampler_config.enable_expert_parallel,
               "enable_prefix_caching": rollout_prefix_caching_enabled(trainer_config),
+              "trust_remote_code": True,
               # Ensures vLLM model initializes with correct dtype (not float32 default)
               "dtype": trainer_config.weight_dtype.value,
           },
diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py
index c99b5f48b6..ab76f45b39 100644
--- a/src/maxtext/trainers/post_train/sft/train_sft.py
+++ b/src/maxtext/trainers/post_train/sft/train_sft.py
@@ -109,6 +109,7 @@ def train_step(
         grad_accumulator: Any = None,
         inputs: Any = None,
         is_update_step: Any = True,
+        **kwargs: Any,
     ):
       if inputs is None and grad_accumulator is not None:
         # In Tunix versions where train_step only receives (model, optimizer, inputs)
diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py
index a2a1403125..2fb7e4161b 100644
--- a/src/maxtext/utils/model_creation_utils.py
+++ b/src/maxtext/utils/model_creation_utils.py
@@ -49,6 +49,7 @@
 from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_TRAIN
 from maxtext.configs import pyconfig
 from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter
+from maxtext.layers import moe
 from maxtext.layers import quantizations
 from maxtext.models import models
 from maxtext.utils import max_logging
@@ -583,7 +584,8 @@ def create_nnx_abstract_model(
     # serve-mode AQT variables (NamedSharding with `spec=None` rejected under
     # AbstractMesh). Sharding is resolved afterwards via the helper, so the
     # wrap is unnecessary here.
-    abs_model = nnx.eval_shape(_create_model)
+    with jax.set_mesh(None):
+      abs_model = nnx.eval_shape(_create_model)
     if mesh is None:
       mesh = abs_model.mesh
     graphdef, abs_var_state = nnx.split(abs_model)
@@ -941,6 +943,10 @@ def from_pretrained(
 
   with mesh:
     if config.load_parameters_path:
+      load_path = epath.Path(config.load_parameters_path)
+      if (load_path / "model_params").exists():
+        load_path = load_path / "model_params"
+
       ckptr = ocp.Checkpointer(
           ocp.PyTreeCheckpointHandler(
               restore_concurrent_gb=config.checkpoint_storage_concurrent_gb,
@@ -955,8 +961,8 @@ def from_pretrained(
       # waste memory, we instead restore the params field of the checkpoint (which itself may be a dictionary
       #  containing a key named 'params').
 
-      # Get the structure of checkpoint in `config.load_parameters_path`
-      metadata = ckptr.metadata(config.load_parameters_path)
+      # Get the structure of checkpoint in `load_path`
+      metadata = ckptr.metadata(load_path)
       if metadata is None or metadata.item_metadata is None:
         max_logging.log(
             f"ERROR: No valid Orbax checkpoint found at '{config.load_parameters_path}'. "
@@ -1010,7 +1016,9 @@ def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx):
       # types (e.g. `qrhs.frozen`), which are NOT subclasses of `nnx.Param`. Negative filtering with
       # `not isinstance(...)` safely retains all weight-like leaves while excluding transient runtime state.
       param_state = sharded_state.filter(
-          lambda path, var: not isinstance(var, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat))
+          lambda path, var: not isinstance(
+              var, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat, moe.Tid2EidVar, moe.MoEBiasVar)
+          )
       )
       is_nnx_checkpoint = True
       if (
@@ -1102,7 +1110,9 @@ def _free_device_memory(path, node):
         is_custom = any("custom_linear" in str(getattr(p, "key", p)) for p in path)
         if (
             isinstance(node, nnx.Variable)
-            and not isinstance(node, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat))
+            and not isinstance(
+                node, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat, moe.Tid2EidVar, moe.MoEBiasVar)
+            )
             and not is_custom
         ):
           inner = node.get_value() if hasattr(node, "get_value") else node[...]
@@ -1120,7 +1130,7 @@ def _free_device_memory(path, node):
       jax.tree_util.tree_map_with_path(_free_device_memory, sharded_state, is_leaf=lambda n: isinstance(n, nnx.Variable))
 
       restored = ckptr.restore(
-          epath.Path(config.load_parameters_path),
+          load_path,
           item=item_to_restore,
           transforms={},
           restore_args=restore_args,
diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh
index eed23bad03..c7f54c515e 100644
--- a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh
+++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_rl.sh
@@ -60,7 +60,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
 # Step 3: Run inference on the checkpoint generated from the previous run
 python3 -m maxtext.inference.vllm_decode \
     model_name=${MODEL_NAME} \
-    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \
+    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/items \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     hbm_utilization_vllm=0.5 \
     prompt='Suggest some famous landmarks in London.' \
diff --git a/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh b/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh
index ba2e294e57..560ea2ce43 100644
--- a/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh
+++ b/tests/end_to_end/tpu/gemma4/26b/test_gemma4_rl.sh
@@ -65,7 +65,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
 # Step 3: Run inference on the checkpoint generated from the previous run
 python3 -m maxtext.inference.vllm_decode \
     model_name=${MODEL_NAME} \
-    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \
+    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/items \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     hbm_utilization_vllm=0.85 \
     prompt='Suggest some famous landmarks in London.' \
diff --git a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh
index e50fab88f0..1c1dad96eb 100644
--- a/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh
+++ b/tests/end_to_end/tpu/llama3.1/70b/test_llama3.1_70b_rl.sh
@@ -54,7 +54,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
 # Step 3: Run inference on the checkpoint generated from the previous run
 python3 -m maxtext.inference.vllm_decode \
     model_name=${MODEL_NAME} \
-    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \
+    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/items \
     tokenizer_path='meta-llama/Llama-3.1-70B-Instruct' \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     hbm_utilization_vllm=0.6 \
diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh
index 3bd90eb519..a4916d8c24 100755
--- a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh
+++ b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh
@@ -95,7 +95,7 @@ python3 -m maxtext.trainers.post_train.rl.train_rl \
 # Step 3: Run inference on the checkpoint produced by the RL run
 python3 -m maxtext.inference.vllm_decode \
     model_name=${MODEL_NAME} \
-    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/model_params \
+    load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/2/items \
     vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \
     hbm_utilization_vllm=0.85 \
     prompt='Suggest some famous landmarks in London.' \
diff --git a/tests/post_training/unit/distillation_checkpointing_test.py b/tests/post_training/unit/distillation_checkpointing_test.py
index 940f3bbc37..b34f0fffde 100644
--- a/tests/post_training/unit/distillation_checkpointing_test.py
+++ b/tests/post_training/unit/distillation_checkpointing_test.py
@@ -21,6 +21,10 @@
 
 import json
 import os
+from types import SimpleNamespace
+from etils import epath
+import jax.numpy as jnp
+import optax
 import shutil
 import tempfile
 from unittest import mock
@@ -30,6 +34,7 @@
 import jax
 from flax import nnx
 import orbax.checkpoint as ocp
+from maxtext.common import checkpointing
 from maxtext.trainers.post_train.distillation import distillation_utils
 
 
@@ -88,6 +93,8 @@ def test_save_and_restore_iterator(self):
     # 2. Save Checkpoint
     mock_student_config = mock.Mock()
     mock_student_config.learn_to_init_mode = False
+    mock_student_config.scan_layers = True
+    mock_student_config.lora = None
     manager = distillation_utils.MaxTextCheckpointManager(
         raw_iterator=iterator, root_directory=self.test_dir, student_config=mock_student_config, options=self.options
     )
@@ -119,6 +126,8 @@ def test_save_and_restore_iterator(self):
 
     mock_student_config_restore = mock.Mock()
     mock_student_config_restore.learn_to_init_mode = False
+    mock_student_config_restore.scan_layers = True
+    mock_student_config_restore.lora = None
     restore_manager = distillation_utils.MaxTextCheckpointManager(
         raw_iterator=new_iterator,
         root_directory=self.test_dir,
@@ -137,6 +146,8 @@ def test_restore_returns_none_if_no_checkpoint(self):
     iterator = FakeGrainIterator()
     mock_student_config_restore = mock.Mock()
     mock_student_config_restore.learn_to_init_mode = False
+    mock_student_config_restore.scan_layers = True
+    mock_student_config_restore.lora = None
     manager = distillation_utils.MaxTextCheckpointManager(
         raw_iterator=iterator,
         root_directory=self.test_dir,
@@ -149,5 +160,107 @@ def test_restore_returns_none_if_no_checkpoint(self):
     self.assertIsNone(result)
 
 
+class MaxTextCheckpointManagerLayoutTest(absltest.TestCase):
+  """Distillation checkpoints only the student, in MaxText's on-disk layout."""
+
+  class Bundle(nnx.Module):
+    """Stand-in for the teacher/student ModelBundle the trainer holds."""
+
+    def __init__(self, student, teacher):
+      self.student_model = student
+      self.teacher_model = teacher
+
+  def setUp(self):
+    super().setUp()
+    self.test_dir = tempfile.mkdtemp()
+    self.options = ocp.CheckpointManagerOptions(max_to_keep=2, create=True)
+
+  def tearDown(self):
+    if os.path.exists(self.test_dir):
+      shutil.rmtree(self.test_dir)
+    super().tearDown()
+
+  def _save(self, learn_to_init_mode=False):
+    """Saves a checkpoint and returns its on-disk leaf paths."""
+    student, teacher = DummyModel(nnx.Rngs(0)), DummyModel(nnx.Rngs(1))
+    bundle = self.Bundle(student, teacher)
+    optimizer = nnx.Optimizer(student, optax.adamw(1e-3), wrt=nnx.Param)
+    manager = distillation_utils.MaxTextCheckpointManager(
+        raw_iterator=None,
+        root_directory=self.test_dir,
+        student_config=SimpleNamespace(learn_to_init_mode=learn_to_init_mode, scan_layers=True, lora=None),
+        options=self.options,
+    )
+    self.assertTrue(manager.save(1, bundle, optimizer, force=True))
+    manager.wait_until_finished()
+    manager.close()
+
+    metadata = ocp.Checkpointer(ocp.PyTreeCheckpointHandler()).metadata(epath.Path(self.test_dir) / "1" / "items")
+    tree = getattr(metadata.item_metadata, "tree", metadata.item_metadata)
+    return list(ocp.tree.to_flat_dict(tree, sep="/"))
+
+  def test_saves_the_student_in_maxtext_layout(self):
+    keys = self._save()
+
+    self.assertTrue(any(k.startswith("params/params/layer/") for k in keys), keys)
+    self.assertTrue(any(k.startswith("opt_state/") for k in keys), keys)
+    # The bundle's wrapper level and the teacher stay out of the checkpoint.
+    self.assertEqual([k for k in keys if "student_model" in k.split("/") or "teacher_model" in k.split("/")], [])
+
+  def test_the_students_scan_setting_is_recorded(self):
+    """Distillation checkpoints the student, so the metadata has to describe the student."""
+    student, teacher = DummyModel(nnx.Rngs(0)), DummyModel(nnx.Rngs(1))
+    bundle = self.Bundle(student, teacher)
+    optimizer = nnx.Optimizer(student, optax.adamw(1e-3), wrt=nnx.Param)
+    manager = distillation_utils.MaxTextCheckpointManager(
+        raw_iterator=None,
+        root_directory=self.test_dir,
+        student_config=SimpleNamespace(learn_to_init_mode=False, scan_layers=False, lora=None),
+        options=self.options,
+    )
+    manager.save(1, bundle, optimizer, force=True)
+    manager.wait_until_finished()
+    manager.close()
+
+    metadata = checkpointing.load_checkpoint_metadata(os.path.join(self.test_dir, "1", "items"))
+    self.assertIs(metadata.get("scan_layers"), False)
+
+  def test_learn_to_init_mode_leaves_the_optimizer_out(self):
+    keys = self._save(learn_to_init_mode=True)
+
+    self.assertTrue(any(k.startswith("params/params/layer/") for k in keys), keys)
+    self.assertEqual([k for k in keys if k.startswith("opt_state")], [])
+
+  def test_restores_the_student(self):
+    student, teacher = DummyModel(nnx.Rngs(0)), DummyModel(nnx.Rngs(1))
+    bundle = self.Bundle(student, teacher)
+    optimizer = nnx.Optimizer(student, optax.adamw(1e-3), wrt=nnx.Param)
+    optimizer.update(student, jax.tree.map(jnp.ones_like, nnx.state(student, nnx.Param)))
+    trained = jnp.asarray(student.layer.kernel[...])
+
+    def manager():
+      return distillation_utils.MaxTextCheckpointManager(
+          raw_iterator=None,
+          root_directory=self.test_dir,
+          student_config=SimpleNamespace(learn_to_init_mode=False),
+          options=self.options,
+      )
+
+    saver = manager()
+    saver.save(1, bundle, optimizer, force=True)
+    saver.wait_until_finished()
+    saver.close()
+
+    fresh_student = DummyModel(nnx.Rngs(0))
+    fresh_bundle = self.Bundle(fresh_student, teacher)
+    fresh_optimizer = nnx.Optimizer(fresh_student, optax.adamw(1e-3), wrt=nnx.Param)
+    restorer = manager()
+    step, _ = restorer.maybe_restore(fresh_bundle, fresh_optimizer)
+    restorer.close()
+
+    self.assertEqual(step, 1)
+    self.assertTrue(jnp.array_equal(trained, fresh_student.layer.kernel[...]))
+
+
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/post_training/unit/post_train_checkpointing_test.py b/tests/post_training/unit/post_train_checkpointing_test.py
new file mode 100644
index 0000000000..d5d0eaa45f
--- /dev/null
+++ b/tests/post_training/unit/post_train_checkpointing_test.py
@@ -0,0 +1,576 @@
+# Copyright 2023-2026 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.
+
+"""Unit tests for post-training checkpointing in MaxText's on-disk layout."""
+
+import os
+import tempfile
+from types import SimpleNamespace
+import unittest
+from unittest import mock
+
+from etils import epath
+from flax import nnx
+import jax
+import jax.numpy as jnp
+import optax
+import orbax.checkpoint as ocp
+import pytest
+from tunix.sft import checkpoint_manager as tunix_checkpoint_manager
+
+from maxtext.common import checkpointing
+from maxtext.trainers.post_train import checkpointing as post_train_checkpointing
+
+pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only]
+
+
+class _Model(nnx.Module):
+  """Tiny stand-in for a MaxText Transformer."""
+
+  def __init__(self, rngs: nnx.Rngs):
+    self.linear = nnx.Linear(2, 3, rngs=rngs)
+    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
+
+
+class _Adapter(nnx.Module):
+  """Stand-in for TunixMaxTextAdapter: holds the model as its only child."""
+
+  def __init__(self, base):
+    self.base = base
+
+
+class _ScannedModel(nnx.Module):
+  """scan_layers=True: every decoder layer stacked under one `layers` key."""
+
+  def __init__(self, rngs: nnx.Rngs, num_layers=3):
+    self.layers = nnx.Param(jnp.zeros((num_layers, 2, 3)))
+    self.decoder_norm = nnx.Param(jnp.ones((3,)))
+    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
+
+
+class _UnscannedModel(nnx.Module):
+  """scan_layers=False: one key per decoder layer, each without the stacking axis."""
+
+  def __init__(self, rngs: nnx.Rngs, num_layers=3):
+    for i in range(num_layers):
+      setattr(self, f"layers_{i}", nnx.Linear(2, 3, rngs=rngs))
+    self.decoder_norm = nnx.Param(jnp.ones((3,)))
+    self.dropout = nnx.Dropout(rate=0.1, rngs=rngs)
+
+
+def _build(wrapped):
+  model = _Model(nnx.Rngs(0))
+  outer = _Adapter(model) if wrapped else model
+  optimizer = nnx.Optimizer(outer, optax.adamw(1e-3), wrt=nnx.Param)
+  return outer, optimizer
+
+
+def _on_disk_keys(directory, step=1):
+  """Returns the checkpoint's leaf paths, slash-separated.
+
+  Args:
+    directory: Checkpoint root directory.
+    step: Step to read.
+
+  Returns:
+    A list of leaf paths.
+  """
+  metadata = ocp.Checkpointer(ocp.PyTreeCheckpointHandler()).metadata(epath.Path(directory) / str(step) / "items")
+  tree = getattr(metadata.item_metadata, "tree", metadata.item_metadata)
+  return list(ocp.tree.to_flat_dict(tree, sep="/"))
+
+
+def _train_a_step(model, optimizer):
+  """Moves the weights and fills opt_state, so a restore has something to prove."""
+  target = model.base if isinstance(model, _Adapter) else model
+  grads = jax.tree.map(lambda p: jnp.full_like(p, 0.1), nnx.state(model, nnx.Param))
+  optimizer.update(model, grads)
+  return jnp.asarray(target.linear.kernel[...])
+
+
+class PostTrainCheckpointInjectHyperparamsTest(unittest.TestCase):
+  """RL and distillation wrap the optimizer in inject_hyperparams; the layout must not notice."""
+
+  def test_opt_state_matches_an_unwrapped_optimizer(self):
+    """mu and nu belong under the Linen `params` collection either way.
+
+    The conversion finds them by name at the top of the optimizer state. Behind the
+    inject_hyperparams shell they are not there, so stripping it afterwards leaves them
+    unwrapped and pre-training cannot line the optimizer up.
+    """
+
+    def mu_keys(directory, optimizer):
+      model = _Model(nnx.Rngs(0))
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=directory, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      self.assertTrue(manager.save(1, model, optimizer, force=True))
+      manager.close()
+      keys = _on_disk_keys(directory)
+      return sorted({k.split("/")[3] for k in keys if k.startswith("opt_state/0/mu/")})
+
+    with tempfile.TemporaryDirectory() as plain_dir:  # pylint: disable=consider-using-with
+      plain_model = _Model(nnx.Rngs(0))
+      plain = mu_keys(plain_dir, nnx.Optimizer(plain_model, optax.adamw(1e-3), wrt=nnx.Param))
+
+    with tempfile.TemporaryDirectory() as injected_dir:  # pylint: disable=consider-using-with
+      injected_model = _Model(nnx.Rngs(0))
+      # A schedule, as RL and distillation use. A plain float produces a different shell --
+      # optax only adds hyperparams_states for callables -- and would not reproduce this.
+      schedule = optax.constant_schedule(1e-3)
+      tx = optax.inject_hyperparams(optax.adamw)(learning_rate=schedule)
+      injected = mu_keys(injected_dir, nnx.Optimizer(injected_model, tx, wrt=nnx.Param))
+
+    self.assertEqual(plain, ["params"], "an unwrapped optimizer should already be under params")
+    self.assertEqual(injected, plain, "inject_hyperparams changed where mu landed on disk")
+
+
+class PostTrainCheckpointLayoutTest(unittest.TestCase):
+  """The on-disk layout has to be MaxText's, so pre-training can read what post-training wrote."""
+
+  def _save(self, directory, wrapped):
+    model, optimizer = _build(wrapped)
+    trained = _train_a_step(model, optimizer)
+    manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+        root_directory=directory, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+    )
+    self.assertTrue(manager.save(1, model, optimizer, force=True))
+    manager.close()
+    return trained
+
+  def test_saves_maxtext_layout_not_the_tunix_one(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      self._save(d, wrapped=False)
+      self.assertEqual(sorted(os.listdir(os.path.join(d, "1"))), ["_CHECKPOINT_METADATA", "items"])
+      keys = _on_disk_keys(d)
+
+    self.assertTrue(any(k.startswith("params/params/linear/") for k in keys), keys)
+    self.assertTrue(any(k.startswith("opt_state/") for k in keys), keys)
+    self.assertIn("step", keys)
+    # rngs are NNX-only, so they belong in nnx_aux rather than the Linen collections.
+    self.assertTrue(any(k.startswith("nnx_aux/") for k in keys), keys)
+
+  def test_adapter_level_is_stripped_from_weights_and_optimizer(self):
+    """DPO and RL train through the adapter; its `base` level must not reach the checkpoint.
+
+    Pre-training builds its params and opt_state from the bare model, so a stray `base` level
+    puts every weight at a path it will not look for.
+    """
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      self._save(d, wrapped=True)
+      keys = _on_disk_keys(d)
+
+    self.assertTrue(keys)
+    self.assertEqual([k for k in keys if "base" in k.split("/")], [])
+
+  def test_restores_weights_and_optimizer_it_saved(self):
+    for wrapped in (False, True):
+      with self.subTest(wrapped=wrapped):
+        with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+          trained = self._save(d, wrapped=wrapped)
+
+          model, optimizer = _build(wrapped)
+          target = model.base if wrapped else model
+          self.assertFalse(jnp.array_equal(trained, target.linear.kernel[...]), "fresh model should differ")
+
+          manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+              root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+          )
+          step, _ = manager.maybe_restore(model, optimizer)
+          manager.close()
+
+        self.assertEqual(step, 1)
+        self.assertTrue(jnp.array_equal(trained, target.linear.kernel[...]), "weights were not restored")
+        # A resume that dropped opt_state would silently restart the optimizer's moments.
+        opt_leaves = jax.tree.leaves(nnx.state(optimizer, nnx.optimizer.OptState))
+        self.assertTrue(any(jnp.any(jnp.asarray(leaf) != 0) for leaf in opt_leaves), "opt_state came back empty")
+
+  def test_no_checkpoint_yet_reports_step_zero(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      step, metadata = manager.maybe_restore(model, optimizer)
+      manager.close()
+
+    self.assertEqual(step, 0)
+    self.assertEqual(metadata, {})
+
+  def test_custom_metadata_survives_the_round_trip(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      manager.save(1, model, optimizer, force=True, custom_metadata={"run": "abc"})
+      manager.wait_until_finished()
+      _, metadata = manager.maybe_restore(*_build(wrapped=False))
+      manager.close()
+
+    self.assertEqual(metadata.get("run"), "abc")
+
+  def test_config_metadata_is_stamped_like_pre_training_does(self):
+    """`scan_layers` and the LoRA settings have to ride along for the loaders to read them back."""
+    config = SimpleNamespace(scan_layers=False, lora=SimpleNamespace(lora_rank=8, model_dump=lambda: {"lora_rank": 8}))
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1), config=config
+      )
+      manager.save(1, model, optimizer, force=True)
+      manager.wait_until_finished()
+      _, metadata = manager.maybe_restore(*_build(wrapped=False))
+      manager.close()
+
+    self.assertIs(metadata.get("scan_layers"), False)
+    self.assertEqual(metadata.get("lora"), {"lora_rank": 8})
+
+  def test_a_caller_key_wins_over_the_config_derived_one(self):
+    config = SimpleNamespace(scan_layers=True, lora=None)
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1), config=config
+      )
+      manager.save(1, model, optimizer, force=True, custom_metadata={"scan_layers": False, "run": "abc"})
+      manager.wait_until_finished()
+      _, metadata = manager.maybe_restore(*_build(wrapped=False))
+      manager.close()
+
+    self.assertIs(metadata.get("scan_layers"), False)
+    self.assertEqual(metadata.get("run"), "abc")
+
+  def test_no_config_still_saves(self):
+    """The config is optional; a manager built without one just writes no sidecar metadata."""
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      self.assertTrue(manager.save(1, model, optimizer, force=True))
+      manager.wait_until_finished()
+      _, metadata = manager.maybe_restore(*_build(wrapped=False))
+      manager.close()
+
+    self.assertEqual(metadata, {})
+
+  def test_weights_only_when_there_is_no_optimizer(self):
+    """Some callers checkpoint the model alone; opt_state must simply be absent."""
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, _ = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      self.assertTrue(manager.save(1, model, optimizer=None, force=True))
+      manager.close()
+      keys = _on_disk_keys(d)
+
+    self.assertTrue(any(k.startswith("params/params/linear/") for k in keys), keys)
+    self.assertEqual([k for k in keys if k.startswith("opt_state")], [])
+
+
+class PostTrainCheckpointScanLayoutTest(unittest.TestCase):
+  """Both scan settings have to survive the layout conversion.
+
+  Scanned stacks the decoder layers under one `layers` key; unscanned splits them into
+  `layers_0 … layers_N`. The RL scripts ship both and vLLM requires unscanned.
+  """
+
+  def _round_trip(self, build):
+    """Saves a trained model and restores it into a fresh one.
+
+    Args:
+      build: Callable taking rngs and returning the model to checkpoint.
+
+    Returns:
+      A tuple of the on-disk leaf paths, the restored step, and the params before and after.
+    """
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model = build(nnx.Rngs(0))
+      optimizer = nnx.Optimizer(model, optax.adamw(1e-3), wrt=nnx.Param)
+      grads = jax.tree.map(lambda p: jnp.full_like(p, 0.1), nnx.state(model, nnx.Param))
+      optimizer.update(model, grads)
+      trained = jax.tree.leaves(nnx.state(model, nnx.Param))
+
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      self.assertTrue(manager.save(1, model, optimizer, force=True))
+      manager.wait_until_finished()
+      keys = _on_disk_keys(d)
+
+      restored_model = build(nnx.Rngs(0))
+      restored_optimizer = nnx.Optimizer(restored_model, optax.adamw(1e-3), wrt=nnx.Param)
+      step, _ = manager.maybe_restore(restored_model, restored_optimizer)
+      manager.close()
+
+      restored = jax.tree.leaves(nnx.state(restored_model, nnx.Param))
+    return keys, step, trained, restored
+
+  def test_scanned_layers_round_trip(self):
+    keys, step, trained, restored = self._round_trip(_ScannedModel)
+    self.assertEqual(step, 1)
+    self.assertIn("params/params/layers", keys)
+    self.assertEqual([k for k in keys if k.startswith("params/params/layers_")], [])
+    for want, got in zip(trained, restored):
+      self.assertTrue(jnp.array_equal(want, got))
+
+  def test_unscanned_layers_round_trip(self):
+    keys, step, trained, restored = self._round_trip(_UnscannedModel)
+    self.assertEqual(step, 1)
+    layer_keys = {k.split("/")[2] for k in keys if k.startswith("params/params/layers_")}
+    self.assertEqual(layer_keys, {"layers_0", "layers_1", "layers_2"})
+    for want, got in zip(trained, restored):
+      self.assertTrue(jnp.array_equal(want, got))
+
+  def test_the_two_layouts_are_actually_different_on_disk(self):
+    """Guards the test itself: if both models wrote the same keys, neither case would prove much."""
+    scanned, _, _, _ = self._round_trip(_ScannedModel)
+    unscanned, _, _, _ = self._round_trip(_UnscannedModel)
+    self.assertNotEqual(sorted(scanned), sorted(unscanned))
+
+
+class PostTrainCheckpointMetadataReaderTest(unittest.TestCase):
+  """The metadata has to come back through the function the loaders actually call."""
+
+  def test_load_checkpoint_metadata_reads_what_the_manager_wrote(self):
+    config = SimpleNamespace(
+        scan_layers=False,
+        lora=SimpleNamespace(lora_rank=8, model_dump=lambda: {"lora_rank": 8, "lora_alpha": 16.0}),
+    )
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1), config=config
+      )
+      manager.save(1, model, optimizer, force=True)
+      manager.wait_until_finished()
+      manager.close()
+
+      # `verify_and_sync_scan_layers` and `sync_lora_metadata` both read a checkpoint this way.
+      metadata = checkpointing.load_checkpoint_metadata(os.path.join(d, "1", "items"))
+
+    self.assertIs(metadata.get("scan_layers"), False)
+    self.assertEqual(metadata.get("lora"), {"lora_rank": 8, "lora_alpha": 16.0})
+
+
+class PostTrainCheckpointSaveDecisionTest(unittest.TestCase):
+  """Saving has to honour the same enable/interval rules the Tunix manager applied."""
+
+  def test_disabled_when_there_is_no_root_directory(self):
+    model, optimizer = _build(wrapped=False)
+    manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(root_directory=None)
+
+    self.assertFalse(manager.save(1, model, optimizer, force=True))
+    self.assertEqual(manager.maybe_restore(model, optimizer), (0, {}))
+
+  def test_declines_when_the_policy_says_not_to_save(self):
+    """`force` bypasses the policy; without it the manager's decision is honoured."""
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      with mock.patch.object(manager._checkpoint_manager, "should_save", return_value=False):  # pylint: disable=protected-access
+        declined = manager.save(3, model, optimizer)
+      forced = manager.save(3, model, optimizer, force=True)
+      manager.close()
+
+      self.assertFalse(declined)
+      self.assertTrue(forced)
+      self.assertEqual(sorted(x for x in os.listdir(d) if x.isdigit()), ["3"])
+
+
+class PostTrainCheckpointLegacyLayoutTest(unittest.TestCase):
+  """Checkpoints written before the layout change are still in Tunix's, and must still restore."""
+
+  def test_falls_back_to_the_tunix_layout(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      trained = _train_a_step(model, optimizer)
+
+      legacy = tunix_checkpoint_manager.CheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      legacy.save(1, model, optimizer, force=True)
+      legacy.close()
+
+      fresh_model, fresh_optimizer = _build(wrapped=False)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      step, _ = manager.maybe_restore(fresh_model, fresh_optimizer)
+      manager.close()
+
+    self.assertEqual(step, 1)
+    self.assertTrue(jnp.array_equal(trained, fresh_model.linear.kernel[...]))
+
+
+class PostTrainCheckpointSubclassHookTest(unittest.TestCase):
+  """Distillation checkpoints a sub-module and an extra item, through these hooks."""
+
+  class _Bundle(nnx.Module):
+
+    def __init__(self, student):
+      self.student_model = student
+
+  class _Manager(post_train_checkpointing.MaxTextLayoutCheckpointManager):
+    """Stand-in for the distillation manager: checkpoints a sub-module plus an extra item."""
+
+    def __init__(self, root_directory, options):
+      super().__init__(
+          root_directory=root_directory,
+          options=options,
+          extra_item_handlers={"note": ocp.JsonCheckpointHandler()},
+      )
+
+    def model_to_checkpoint(self, model):
+      return getattr(model, "student_model", model)
+
+    def _extra_save_args(self, step):
+      del step
+      return {"note": ocp.args.JsonSave({"hello": "world"})}
+
+  def test_hooks_pick_the_submodule_and_add_the_extra_item(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      student = _Model(nnx.Rngs(0))
+      bundle = self._Bundle(student)
+      optimizer = nnx.Optimizer(student, optax.adamw(1e-3), wrt=nnx.Param)
+      manager = self._Manager(d, ocp.CheckpointManagerOptions(save_interval_steps=1))
+      self.assertTrue(manager.save(1, bundle, optimizer, force=True))
+      manager.close()
+
+      self.assertIn("note", os.listdir(os.path.join(d, "1")))
+      metadata = ocp.Checkpointer(ocp.PyTreeCheckpointHandler()).metadata(epath.Path(d) / "1" / "items")
+      tree = getattr(metadata.item_metadata, "tree", metadata.item_metadata)
+      keys = list(ocp.tree.to_flat_dict(tree, sep="/"))
+
+    # The student's weights, not the bundle's wrapper level.
+    self.assertTrue(any(k.startswith("params/params/linear/") for k in keys), keys)
+    self.assertEqual([k for k in keys if "student_model" in k.split("/")], [])
+
+
+class PostTrainCheckpointBaseManagerTest(unittest.TestCase):
+  """The base class builds a manager over Tunix's item names before we replace it."""
+
+  def test_closes_the_base_class_manager_it_replaces(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      mock_base_cm = mock.MagicMock()
+
+      def fake_base_init(self, root_directory=None, options=None):
+        del root_directory, options
+        self._checkpoint_manager = mock_base_cm
+
+      with mock.patch.object(tunix_checkpoint_manager.CheckpointManager, "__init__", fake_base_init):
+        manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+            root_directory=d,
+            options=ocp.CheckpointManagerOptions(save_interval_steps=1),
+        )
+      mock_base_cm.close.assert_called_once()
+      # pylint: disable=protected-access
+      self.assertIsNotNone(manager._checkpoint_manager)
+      self.assertIsNot(manager._checkpoint_manager, mock_base_cm)
+      # pylint: enable=protected-access
+      manager.close()
+
+
+class InstallTest(unittest.TestCase):
+  """`install` swaps in the MaxText-layout manager and restores what it finds."""
+
+  class _FakeConfig:
+
+    def __init__(self):
+      self.checkpointing_options = ocp.CheckpointManagerOptions(save_interval_steps=1)
+
+    def get_with_default(self, key, default):
+      del key
+      return default
+
+  class _FakeTrainer:
+
+    def __init__(self, model, optimizer, checkpoint_manager):
+      self.model = model
+      self.optimizer = optimizer
+      self.checkpoint_manager = checkpoint_manager
+      self.config = InstallTest._FakeConfig()
+      self._train_steps = 0
+      self._iter_steps = 0
+
+  def test_replaces_the_manager_and_restores_the_step(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      trained = _train_a_step(model, optimizer)
+      manager = post_train_checkpointing.MaxTextLayoutCheckpointManager(
+          root_directory=d, options=ocp.CheckpointManagerOptions(save_interval_steps=1)
+      )
+      manager.save(4, model, optimizer, force=True)
+      manager.close()
+
+      fresh_model, fresh_optimizer = _build(wrapped=False)
+      trainer = self._FakeTrainer(fresh_model, fresh_optimizer, checkpoint_manager=None)
+      post_train_checkpointing.install(trainer, d)
+      trainer.checkpoint_manager.close()
+
+    self.assertIsInstance(trainer.checkpoint_manager, post_train_checkpointing.MaxTextLayoutCheckpointManager)
+    self.assertEqual(trainer._train_steps, 4)  # pylint: disable=protected-access
+    self.assertEqual(trainer._iter_steps, 4)  # pylint: disable=protected-access
+    self.assertTrue(jnp.array_equal(trained, fresh_model.linear.kernel[...]))
+
+  def test_forwards_the_run_config_to_the_manager(self):
+    """The manager reads the config for the metadata it stamps, so `install` has to pass it on."""
+    config = SimpleNamespace(scan_layers=False, lora=None)
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      trainer = self._FakeTrainer(*_build(wrapped=False), checkpoint_manager=None)
+      post_train_checkpointing.install(trainer, d, config)
+      trainer.checkpoint_manager.save(1, trainer.model, trainer.optimizer, force=True)
+      trainer.checkpoint_manager.wait_until_finished()
+      _, metadata = trainer.checkpoint_manager.maybe_restore(*_build(wrapped=False))
+      trainer.checkpoint_manager.close()
+
+    self.assertIs(metadata.get("scan_layers"), False)
+
+  def test_closes_the_manager_it_replaces(self):
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      model, optimizer = _build(wrapped=False)
+      replaced = tunix_checkpoint_manager.CheckpointManager(root_directory=None)
+      closed = []
+      replaced.close = lambda: closed.append(True)
+
+      trainer = self._FakeTrainer(model, optimizer, checkpoint_manager=replaced)
+      post_train_checkpointing.install(trainer, d)
+      trainer.checkpoint_manager.close()
+
+    self.assertEqual(closed, [True])
+
+
+class UnwrapModelTest(unittest.TestCase):
+  """The adapter is matched on its child module, not on its class."""
+
+  def test_unwraps_a_wrapper(self):
+    model = _Model(nnx.Rngs(0))
+    self.assertIs(post_train_checkpointing.unwrap_model(_Adapter(model)), model)
+
+  def test_leaves_a_bare_model_alone(self):
+    model = _Model(nnx.Rngs(0))
+    self.assertIs(post_train_checkpointing.unwrap_model(model), model)
+
+  def test_ignores_a_base_attribute_that_is_not_a_module(self):
+    model = _Model(nnx.Rngs(0))
+    setattr(model, "base", "not a module")
+    self.assertIs(post_train_checkpointing.unwrap_model(model), model)
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py
index 845014db32..52643b8923 100644
--- a/tests/post_training/unit/train_distill_test.py
+++ b/tests/post_training/unit/train_distill_test.py
@@ -95,6 +95,49 @@ def test_maxtext_to_tunix_iterator(self):
     expected_mask = dummy_batch["inputs_segmentation"] != 0
     np.testing.assert_array_equal(tunix_input.input_mask, expected_mask)
 
+  def test_maxtext_to_tunix_iterator_stops_at_the_batch_budget(self):
+    """The trainer is managed externally, so this bound is what ends an endless dataset."""
+
+    def endless():
+      while True:
+        yield {
+            "inputs": np.array([[10, 11]]),
+            "inputs_position": np.array([[0, 1]]),
+            "inputs_segmentation": np.array([[1, 1]]),
+            "targets": np.array([[11, 12]]),
+        }
+
+    adapter = distillation_utils.MaxTextToTunixIterator(endless(), max_batches=3)
+    self.assertEqual(len(list(adapter)), 3)
+
+  def test_maxtext_to_tunix_iterator_is_unbounded_without_a_budget(self):
+    """A finite upstream iterator still ends on its own."""
+    batches = [
+        {
+            "inputs": np.array([[10, 11]]),
+            "inputs_position": np.array([[0, 1]]),
+            "inputs_segmentation": np.array([[1, 1]]),
+            "targets": np.array([[11, 12]]),
+        }
+    ] * 4
+    adapter = distillation_utils.MaxTextToTunixIterator(iter(batches))
+    self.assertEqual(len(list(adapter)), 4)
+
+  def test_maxtext_to_tunix_iterator_spends_nothing_on_a_zero_budget(self):
+    """A resume that is already at its step count must not train further."""
+
+    def endless():
+      while True:
+        yield {
+            "inputs": np.array([[10, 11]]),
+            "inputs_position": np.array([[0, 1]]),
+            "inputs_segmentation": np.array([[1, 1]]),
+            "targets": np.array([[11, 12]]),
+        }
+
+    adapter = distillation_utils.MaxTextToTunixIterator(endless(), max_batches=0)
+    self.assertEqual(len(list(adapter)), 0)
+
   def test_maxtext_to_tunix_iterator_sft(self):
     """Verifies SFT-related fields are handled correctly."""
     # 1. Create a dummy batch with SFT fields
@@ -960,6 +1003,10 @@ def __call__(self, input_tokens, **kwargs):
     config.checkpoint_dir = self.test_dir
     config.dataset_type = "synthetic"
     config.lora_enabled = False
+    # The checkpoint manager reads these for the metadata it stamps, and a bare Mock puts a Mock
+    # where a bool belongs.
+    config.scan_layers = True
+    config.lora = None
 
     # pylint: disable=import-outside-toplevel
     from tunix.sft import peft_trainer
@@ -967,7 +1014,12 @@ def __call__(self, input_tokens, **kwargs):
     train_config = peft_trainer.TrainingConfig(
         max_steps=2,
         eval_every_n_steps=0,
-        checkpointing_options=ocp.CheckpointManagerOptions(save_interval_steps=1, max_to_keep=2, create=True),
+        checkpointing_options=ocp.CheckpointManagerOptions(
+            save_interval_steps=1,
+            max_to_keep=2,
+            create=True,
+            enable_async_checkpointing=False,
+        ),
         gradient_accumulation_steps=1,
     )
 
diff --git a/tests/post_training/unit/train_dpo_test.py b/tests/post_training/unit/train_dpo_test.py
index 489d3ff9ed..9bf51467ee 100644
--- a/tests/post_training/unit/train_dpo_test.py
+++ b/tests/post_training/unit/train_dpo_test.py
@@ -51,5 +51,44 @@ def test_validate_config_invalid_vocab_tiling(self):
       train_dpo.validate_config(config)
 
 
+class TrainDPOTunixConfigTest(unittest.TestCase):
+  """The Tunix config decides who checkpoints and whether the optimizer gets wrapped."""
+
+  def _mt_config(self, grad_accum=1):
+    return SimpleNamespace(
+        checkpoint_period=5,
+        async_checkpointing=False,
+        tensorboard_dir="/tmp/tb",
+        profiler="",
+        eval_interval=1,
+        steps=10,
+        checkpoint_dir="/tmp/ckpt",
+        data_sharding=["data"],
+        gradient_accumulation_steps=grad_accum,
+        max_target_length=128,
+        dpo=SimpleNamespace(
+            algo="dpo",
+            orpo_lambda=1.0,
+            dpo_beta=0.1,
+            dpo_label_smoothing=0.0,
+            max_prompt_length=32,
+        ),
+    )
+
+  @pytest.mark.cpu_only
+  def test_tunix_checkpointing_is_disabled(self):
+    """post_train.checkpointing owns checkpointing, so Tunix's own manager must stay off."""
+    self.assertIsNone(train_dpo.get_tunix_config(self._mt_config()).checkpoint_root_directory)
+
+  @pytest.mark.cpu_only
+  def test_single_step_accumulation_is_not_passed_through(self):
+    """Tunix wraps the optimizer in MultiSteps whenever this is set, changing the state shape."""
+    self.assertIsNone(train_dpo.get_tunix_config(self._mt_config(grad_accum=1)).gradient_accumulation_steps)
+
+  @pytest.mark.cpu_only
+  def test_real_accumulation_is_passed_through(self):
+    self.assertEqual(train_dpo.get_tunix_config(self._mt_config(grad_accum=4)).gradient_accumulation_steps, 4)
+
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py
index 41daec72e9..d48ebbd40b 100644
--- a/tests/post_training/unit/train_rl_test.py
+++ b/tests/post_training/unit/train_rl_test.py
@@ -20,8 +20,11 @@
 import pytest
 from types import SimpleNamespace
 import jax
+import jax.numpy as jnp
+import optax
 
 from maxtext.trainers.post_train.rl import train_rl
+from maxtext.trainers.post_train.rl import utils_rl
 
 pytestmark = [pytest.mark.post_training]
 from maxtext.configs import types
@@ -662,5 +665,43 @@ def apply_chat_template(self, conversation, tokenize=False):
     self.assertEqual(rendered, "Hello!")
 
 
+class RLOptimizerClippingTest(unittest.TestCase):
+  """RL clips gradients without giving the optimizer state an extra chain level."""
+
+  def _config(self, threshold):
+    return SimpleNamespace(
+        learning_rate=1e-3,
+        learning_rate_schedule_steps=-1,
+        steps=-1,
+        train_steps=10,
+        # No warmup, so the learning rate at step 0 is non-zero and updates are comparable.
+        warmup_steps_fraction=0.0,
+        adam_b1=0.9,
+        adam_b2=0.95,
+        adam_weight_decay=0.1,
+        gradient_clipping_threshold=threshold,
+    )
+
+  def _opt_state_structure(self, threshold):
+    params = {"w": jnp.array([3.0, 4.0])}
+    return jax.tree_util.tree_structure(utils_rl.get_optimizer(self._config(threshold)).init(params))
+
+  @pytest.mark.cpu_only
+  def test_clipping_does_not_change_the_optimizer_state_shape(self):
+    self.assertEqual(self._opt_state_structure(1.0), self._opt_state_structure(0.0))
+
+  @pytest.mark.cpu_only
+  def test_gradients_are_clipped(self):
+    params = {"w": jnp.array([3.0, 4.0])}  # global norm 5.0
+    grads = {"w": jnp.array([3.0, 4.0])}
+
+    clipped = utils_rl.get_optimizer(self._config(0.1))
+    unclipped = utils_rl.get_optimizer(self._config(0.0))
+    clipped_updates, _ = clipped.update(grads, clipped.init(params), params)
+    unclipped_updates, _ = unclipped.update(grads, unclipped.init(params), params)
+
+    self.assertLess(float(optax.tree.norm(clipped_updates)), float(optax.tree.norm(unclipped_updates)))
+
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/post_training/unit/train_sft_test.py b/tests/post_training/unit/train_sft_test.py
index 3e71ba6273..72d4d07e4f 100644
--- a/tests/post_training/unit/train_sft_test.py
+++ b/tests/post_training/unit/train_sft_test.py
@@ -104,5 +104,35 @@ def test_maxtext_peft_trainer_train_step_signature(self):
       self.assertEqual(params, ["model", "optimizer", "grad_accumulator", "inputs", "is_update_step"])
 
 
+class TrainSFTTunixConfigTest(unittest.TestCase):
+  """The Tunix config decides who checkpoints and whether the optimizer gets wrapped."""
+
+  def _mt_config(self, grad_accum=1):
+    return SimpleNamespace(
+        checkpoint_period=5,
+        async_checkpointing=False,
+        tensorboard_dir="/tmp/tb",
+        profiler="",
+        eval_interval=1,
+        steps=10,
+        checkpoint_dir="/tmp/ckpt",
+        data_sharding=["data"],
+        gradient_accumulation_steps=grad_accum,
+    )
+
+  @pytest.mark.cpu_only
+  def test_tunix_checkpointing_is_disabled(self):
+    """post_train.checkpointing owns checkpointing, so Tunix's own manager must stay off."""
+    self.assertIsNone(train_sft.get_tunix_config(self._mt_config()).checkpoint_root_directory)
+
+  @pytest.mark.cpu_only
+  def test_single_step_accumulation_is_not_passed_through(self):
+    self.assertIsNone(train_sft.get_tunix_config(self._mt_config(grad_accum=1)).gradient_accumulation_steps)
+
+  @pytest.mark.cpu_only
+  def test_real_accumulation_is_passed_through(self):
+    self.assertEqual(train_sft.get_tunix_config(self._mt_config(grad_accum=4)).gradient_accumulation_steps, 4)
+
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py
index 0c2dae51ff..976d82fc5a 100644
--- a/tests/unit/checkpointing_nnx_load_test.py
+++ b/tests/unit/checkpointing_nnx_load_test.py
@@ -525,7 +525,7 @@ def __init__(self, rngs):
     self.assertIn("table", ckpt["nnx_aux"]["model"])  # the custom variable persists here instead
     self.assertNotIn("cache", ckpt["nnx_aux"]["model"])  # caches are recomputed, never stored
 
-    restored = checkpointing._linen_items_to_nnx(ckpt, nnx.eval_shape(lambda: state))  # pylint: disable=protected-access
+    restored = checkpointing.linen_items_to_nnx(ckpt, nnx.eval_shape(lambda: state))
     pure = restored.to_pure_dict()
     self.assertTrue(jnp.array_equal(pure["model"]["table"], model.table.value))
     self.assertIsInstance(pure["model"]["cache"], jax.ShapeDtypeStruct)  # left for the init to fill
@@ -595,11 +595,11 @@ def test_no_nnx_aux_when_state_has_none(self):
 
 
 class TestLinenItemsToNnx(unittest.TestCase):
-  """checkpointing._linen_items_to_nnx reshapes restored items into the NNX-layout overlay."""
+  """checkpointing.linen_items_to_nnx reshapes restored items into the NNX-layout overlay."""
 
   def _to_nnx(self, restored):
     """Reshape `restored` against the `_ModelDropout` abstract, as the restore paths do."""
-    state = checkpointing._linen_items_to_nnx(restored, _abstract_dropout_state())  # pylint: disable=protected-access
+    state = checkpointing.linen_items_to_nnx(restored, _abstract_dropout_state())
     return state.to_pure_dict()
 
   def test_materialized_aux_is_kept(self):
@@ -723,5 +723,79 @@ def test_expected_and_restored_params_splits_by_param_type(self):
     checkpointing._raise_on_weight_mismatch(want, have)  # pylint: disable=protected-access
 
 
+class TestLoadNnxNativeParams(unittest.TestCase):
+  """Weight-only load of an NNX-native checkpoint -- the layout post-training writes.
+
+  Tunix and the training engine save `nnx.state(model)` under a `model_params` item, so each
+  leaf lands in a `value` box and the tree is the whole checkpoint. DPO and RL train through
+  `TunixMaxTextAdapter`, adding a `base` level on top.
+  """
+
+  def _save(self, directory, weights, wrapper_key=None):
+    """Writes `weights` the way `nnx.state(model)` serializes, under a `model_params` dir."""
+    boxed = jax.tree.map(lambda leaf: {"value": leaf}, weights)
+    path = os.path.join(directory, "model_params")
+    ocp.PyTreeCheckpointer(use_ocdbt=True, use_zarr3=True).save(
+        epath.Path(path),
+        {wrapper_key: boxed} if wrapper_key else boxed,
+        force=True,
+    )
+    return path
+
+  def _weights(self):
+    return {
+        "linear": {
+            "kernel": jnp.arange(2, dtype=jnp.float32).reshape(2, 1),
+            "bias": jnp.array([5.0]),
+        }
+    }
+
+  def _restore(self, path):
+    # The real caller hands over an abstract params state, as load_state_if_possible does.
+    params_abstract = nnx.eval_shape(lambda: nnx.split(_Model(nnx.Rngs(0)), nnx.Param, ...)[1])
+    return checkpointing.load_params_from_path(path, params_abstract, 8)
+
+  def test_restores_weights_saved_by_a_bare_model(self):
+    """SFT, distillation and the training engine save the model directly, with no wrapper level."""
+    weights = self._weights()
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      restored = self._restore(self._save(d, weights))
+
+    pure = restored.to_pure_dict()
+    self.assertTrue(jnp.array_equal(pure["linear"]["kernel"], weights["linear"]["kernel"]))
+    self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"]))
+
+  def test_restores_weights_saved_through_the_tunix_adapter(self):
+    """DPO and RL nest the whole tree under `base`; the weights still have to come back."""
+    weights = self._weights()
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      restored = self._restore(self._save(d, weights, wrapper_key="base"))
+
+    pure = restored.to_pure_dict()
+    self.assertTrue(jnp.array_equal(pure["linear"]["kernel"], weights["linear"]["kernel"]))
+    self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"]))
+
+  def test_linen_target_against_an_nnx_checkpoint_raises(self):
+    """This layout only restores into an NNX state, so say so instead of failing inside Orbax."""
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      path = self._save(d, self._weights())
+      linen_params = jax.tree.map(jnp.zeros_like, {"params": self._weights()})
+      with self.assertRaises(ValueError) as ctx:
+        checkpointing.load_params_from_path(path, linen_params, 8)
+
+    self.assertIn("NNX", str(ctx.exception))
+
+  def test_weight_missing_from_the_checkpoint_raises(self):
+    """A params-only load has no init state to fall back on, so a gap must not pass silently."""
+    weights = self._weights()
+    del weights["linear"]["bias"]
+    with tempfile.TemporaryDirectory() as d:  # pylint: disable=consider-using-with
+      path = self._save(d, weights)
+      with self.assertRaises(ValueError) as ctx:
+        self._restore(path)
+
+    self.assertIn("linear/bias", str(ctx.exception))
+
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py
index 4ef2ab30b7..3f1a213914 100644
--- a/tests/unit/checkpointing_test.py
+++ b/tests/unit/checkpointing_test.py
@@ -17,6 +17,7 @@
 import asyncio
 import json
 import os
+from types import SimpleNamespace
 from unittest import mock
 
 from absl.testing import absltest
@@ -319,6 +320,32 @@ def test_load_checkpoint_metadata(self, mock_checkpointer_cls):
     self.assertEqual(loaded_metadata.get("lora"), {"lora_rank": 8, "lora_alpha": 16.0})
     mock_ckptr.metadata.assert_called_once()
 
+  @mock.patch.object(checkpointing.ocp, "StandardCheckpointer")
+  def test_load_checkpoint_metadata_falls_back_to_the_step_directory(self, mock_checkpointer_cls):
+    """`load_parameters_path` points at the item, but the metadata belongs to the step above it."""
+    mock_ckptr = mock_checkpointer_cls.return_value
+
+    def metadata_for(path):
+      result = mock.MagicMock()
+      result.custom_metadata = {"scan_layers": False} if str(path).endswith("/0") else None
+      return result
+
+    mock_ckptr.metadata.side_effect = metadata_for
+
+    self.assertEqual(checkpointing.load_checkpoint_metadata("/ckpt/0/items"), {"scan_layers": False})
+    self.assertEqual([str(c.args[0]) for c in mock_ckptr.metadata.call_args_list], ["/ckpt/0/items", "/ckpt/0"])
+
+  @mock.patch.object(checkpointing.ocp, "StandardCheckpointer")
+  def test_load_checkpoint_metadata_stops_at_the_path_it_was_given(self, mock_checkpointer_cls):
+    """A hit at the given path is used as is; no reason to look at the parent."""
+    mock_ckptr = mock_checkpointer_cls.return_value
+    metadata = mock.MagicMock()
+    metadata.custom_metadata = {"scan_layers": True}
+    mock_ckptr.metadata.return_value = metadata
+
+    self.assertEqual(checkpointing.load_checkpoint_metadata("/ckpt/0"), {"scan_layers": True})
+    mock_ckptr.metadata.assert_called_once()
+
   @mock.patch.object(checkpointing.ocp, "StandardCheckpointer")
   def test_load_checkpoint_metadata_handles_exceptions(self, mock_checkpointer_cls):
     mock_ckptr = mock_checkpointer_cls.return_value
@@ -326,7 +353,44 @@ def test_load_checkpoint_metadata_handles_exceptions(self, mock_checkpointer_cls
 
     loaded_metadata = checkpointing.load_checkpoint_metadata("corrupt/path")
     self.assertEqual(loaded_metadata, {})
-    mock_ckptr.metadata.assert_called_once()
+    # A failed read falls through to the step directory above, which fails the same way.
+    self.assertEqual(mock_ckptr.metadata.call_count, 2)
+
+
+class CheckpointCustomMetadataTest(parameterized.TestCase):
+  """What a checkpoint records about the run that wrote it.
+
+  Loaders read this back to fill in a value the run left at its default, or to reject one that
+  contradicts the checkpoint. Post-training saves through its own manager and calls the same
+  builder, so the two paths cannot drift apart.
+  """
+
+  def _config(self, scan_layers=True, lora_rank=0):
+    lora = SimpleNamespace(
+        lora_rank=lora_rank,
+        lora_alpha=16.0,
+        model_dump=lambda: {"lora_rank": lora_rank, "lora_alpha": 16.0},
+    )
+    return SimpleNamespace(scan_layers=scan_layers, lora=lora)
+
+  @parameterized.parameters(True, False)
+  def test_scan_layers_is_recorded_either_way(self, scan_layers):
+    metadata = checkpointing.checkpoint_custom_metadata(self._config(scan_layers=scan_layers))
+    self.assertIs(metadata["scan_layers"], scan_layers)
+
+  def test_lora_is_recorded_once_there_is_a_rank(self):
+    metadata = checkpointing.checkpoint_custom_metadata(self._config(lora_rank=8))
+    self.assertEqual(metadata["lora"], {"lora_rank": 8, "lora_alpha": 16.0})
+
+  def test_no_lora_key_without_a_rank(self):
+    """A rank of 0 means LoRA is off; recording it would make `sync_lora_metadata` sync a zero."""
+    self.assertNotIn("lora", checkpointing.checkpoint_custom_metadata(self._config(lora_rank=0)))
+
+  def test_no_lora_key_when_the_config_has_no_lora_section(self):
+    self.assertNotIn("lora", checkpointing.checkpoint_custom_metadata(SimpleNamespace(scan_layers=True)))
+
+  def test_no_config_records_nothing(self):
+    self.assertEqual(checkpointing.checkpoint_custom_metadata(None), {})
 
 
 class GrainCheckpointableEquivalenceTest(parameterized.TestCase):
diff --git a/tests/unit/optimizers_test.py b/tests/unit/optimizers_test.py
index 6f43c420cd..bf45130f56 100644
--- a/tests/unit/optimizers_test.py
+++ b/tests/unit/optimizers_test.py
@@ -825,5 +825,85 @@ def test_muon_newton_schulz_config(self):
       self.assertEqual(kwargs["ns_coeffs"], (3.4445, -4.7750, 2.0315))
 
 
+class AddGradientClippingTest(parameterized.TestCase):
+  """Clipping must stay identical to the chained form while leaving the state tree alone."""
+
+  THRESHOLD = 1.0
+
+  def _params(self):
+    return {"w": jnp.array([3.0, 4.0]), "b": jnp.array([12.0])}
+
+  def _grads(self, scale):
+    # Global norm of [3, 4, 12] is 13, so scale=1.0 is well over the threshold.
+    return jax.tree.map(lambda x: x * scale, self._params())
+
+  @parameterized.named_parameters(("over_threshold", 1.0), ("under_threshold", 0.01))
+  def test_updates_match_the_chained_form(self, scale):
+    params, grads = self._params(), self._grads(scale)
+
+    chained = optax.chain(optax.clip_by_global_norm(self.THRESHOLD), optax.adamw(1e-2))
+    chained_updates, _ = chained.update(grads, chained.init(params), params)
+
+    inline = optimizers.add_gradient_clipping(optax.adamw(1e-2), self.THRESHOLD)
+    inline_updates, _ = inline.update(grads, inline.init(params), params)
+
+    jax.tree.map(lambda a, b: self.assertTrue(jnp.allclose(a, b), f"{a} != {b}"), chained_updates, inline_updates)
+
+  def test_gradients_over_the_threshold_are_clipped(self):
+    """Guards against the wrapper quietly becoming a passthrough."""
+    params, grads = self._params(), self._grads(1.0)
+
+    inline = optimizers.add_gradient_clipping(optax.sgd(1.0), self.THRESHOLD)
+    inline_updates, _ = inline.update(grads, inline.init(params), params)
+    unclipped = optax.sgd(1.0)
+    unclipped_updates, _ = unclipped.update(grads, unclipped.init(params), params)
+
+    inline_norm = optax.tree.norm(inline_updates)
+    self.assertAlmostEqual(float(inline_norm), self.THRESHOLD, places=5)
+    self.assertGreater(float(optax.tree.norm(unclipped_updates)), float(inline_norm))
+
+  def test_state_tree_matches_the_unclipped_optimizer(self):
+    """A checkpointed opt_state must not gain a level from clipping.
+
+    Pre-training clips in its train step, so its optimizer state is the bare tx state. A
+    post-training checkpoint one chain level deeper cannot be resumed by it.
+    """
+    params = self._params()
+    inner = optax.adamw(1e-2)
+
+    bare = jax.tree_util.tree_structure(inner.init(params))
+    inline = jax.tree_util.tree_structure(optimizers.add_gradient_clipping(inner, self.THRESHOLD).init(params))
+    chained = jax.tree_util.tree_structure(optax.chain(optax.clip_by_global_norm(self.THRESHOLD), inner).init(params))
+
+    self.assertEqual(bare, inline)
+    self.assertNotEqual(bare, chained)
+
+  def test_extra_args_reach_the_inner_optimizer(self):
+    """`get_optimizer` can return a transform taking extra args, as skip_step_on_spikes does."""
+    params, grads = self._params(), self._grads(1.0)
+    seen = {}
+
+    def update_fn(updates, state, params=None, **extra_args):
+      del params
+      seen.update(extra_args)
+      return updates, state
+
+    inner = optax.GradientTransformationExtraArgs(lambda p: optax.EmptyState(), update_fn)
+    wrapped = optimizers.add_gradient_clipping(inner, self.THRESHOLD)
+    wrapped.update(grads, wrapped.init(params), params, loss=jnp.array(1.0))
+
+    self.assertEqual(list(seen), ["loss"])
+
+  def test_clipping_composes_with_skip_step_on_spikes(self):
+    """get_optimizer can return a spike-skipping transform, which takes loss/grad_norm."""
+    params, grads = self._params(), self._grads(1.0)
+
+    inner = optimizers.skip_step_on_spikes(optax.adamw(1e-2), interval=4, scaling_factor=2.0)
+    wrapped = optimizers.add_gradient_clipping(inner, self.THRESHOLD)
+    updates, _ = wrapped.update(grads, wrapped.init(params), params, loss=jnp.array(1.0), grad_norm=jnp.array(1.0))
+
+    self.assertEqual(jax.tree_util.tree_structure(updates), jax.tree_util.tree_structure(params))
+
+
 if __name__ == "__main__":
   unittest.main()