diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml
index b64629d214..15fc192fbe 100644
--- a/.github/workflows/build_and_push_docker_image.yml
+++ b/.github/workflows/build_and_push_docker_image.yml
@@ -42,6 +42,10 @@ on:
         required: false
         type: boolean
         default: false
+      run_tests:
+        required: false
+        type: boolean
+        default: true
     outputs:
       build_result:
         description: 'The result of build_and_push job.'
@@ -184,7 +188,7 @@ jobs:
   run_ci_tests:
     name: Run Unit and Integration Tests
     needs: [pre_build_check, build_and_push]
-    if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true
+    if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true && inputs.run_tests == true
     uses: ./.github/workflows/run_ci_tests.yml
     with:
       image_name: ${{ inputs.image_name }}
diff --git a/.github/workflows/nightly_images_pipeline.yml b/.github/workflows/nightly_images_pipeline.yml
index 4ab0a932b3..32b6a95abe 100644
--- a/.github/workflows/nightly_images_pipeline.yml
+++ b/.github/workflows/nightly_images_pipeline.yml
@@ -50,7 +50,7 @@ jobs:
       device_name: v4-8
       cloud_runner: linux-x86-n2-16-buildkit
 
-  build_and_push_docker_image:
+  build_and_push_tpu_docker_images:
     name: Build ${{ matrix.name }} Docker Image
     needs: build_maxtext_package
     strategy:
@@ -75,6 +75,26 @@ jobs:
             workflow: post-training
             image_name: maxtext_post_training_nightly
             dockerfile: maxtext_tpu_dependencies.Dockerfile
+    uses: ./.github/workflows/build_and_push_docker_image.yml
+    with:
+      image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
+      device: ${{ matrix.device }}
+      build_mode: ${{ matrix.build_mode }}
+      workflow: ${{ matrix.workflow }}
+      dockerfile: ${{ matrix.dockerfile }}
+      maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
+      include_test_assets: true
+      run_tests: false
+    secrets:
+      HF_TOKEN: ${{ secrets.HF_TOKEN }}
+
+  build_and_push_gpu_docker_images:
+    name: Build ${{ matrix.name }} Docker Image
+    needs: build_maxtext_package
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
           - name: "GPU Pre-Training Stable"
             device: gpu
             build_mode: stable
@@ -94,14 +114,40 @@ jobs:
       build_mode: ${{ matrix.build_mode }}
       workflow: ${{ matrix.workflow }}
       dockerfile: ${{ matrix.dockerfile }}
-      maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }}
+      maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }}
       include_test_assets: true
     secrets:
       HF_TOKEN: ${{ secrets.HF_TOKEN }}
 
+  run_tpu_ci_tests:
+    name: Run TPU ${{ matrix.name }} CI Tests
+    needs: build_and_push_tpu_docker_images
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: "Pre-Training Stable"
+            image_name: maxtext_jax_stable
+            device: tpu
+            workflow: pre-training
+          - name: "Pre-Training Nightly"
+            image_name: maxtext_jax_nightly
+            device: tpu
+            workflow: pre-training
+          - name: "Post-Training Nightly"
+            image_name: maxtext_post_training_nightly
+            device: tpu
+            workflow: post-training
+    uses: ./.github/workflows/run_ci_tests.yml
+    with:
+      image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }}
+      image_tag: ${{ github.run_id }}
+      device: ${{ matrix.device }}
+      workflow: ${{ matrix.workflow }}
+
   run_e2e_tests:
     name: Run E2E tests
-    needs: build_and_push_docker_image
+    needs: build_and_push_tpu_docker_images
     uses: ./.github/workflows/run_e2e_tests.yml
     with:
       mode: ${{ inputs.image_suffix != '' && format('{0}_{1}', 'nightly', inputs.image_suffix) || 'nightly' }}
@@ -111,7 +157,7 @@ jobs:
 
   notify_failure:
     name: Notify failed build
-    needs: [build_and_push_docker_image, run_e2e_tests]
+    needs: [build_and_push_tpu_docker_images, build_and_push_gpu_docker_images, run_tpu_ci_tests, run_e2e_tests]
     if: ${{ failure() && inputs.image_suffix == '' }}
     runs-on: ubuntu-latest
     permissions:
diff --git a/.github/workflows/run_tests_coordinator.yml b/.github/workflows/run_tests_coordinator.yml
index 429fae15d1..a47ae59655 100644
--- a/.github/workflows/run_tests_coordinator.yml
+++ b/.github/workflows/run_tests_coordinator.yml
@@ -209,7 +209,7 @@ jobs:
             "gpu-integration": "--ignore=tests/post_training",
             "cpu-unit": "--ignore=tests/post_training",
             "cpu-post-training-unit": "",
-            "cpu-torch-reference": "-o addopts= -rf --import-mode=importlib --strict-markers tests/unit/gemma4_layers_test.py tests/unit/gemma4_small_layers_test.py tests/unit/qwen3_next_vs_reference_test.py tests/unit/qwen3_5_layers_test.py"
+            "cpu-torch-reference": "-o addopts= -rf --import-mode=importlib --strict-markers tests/unit/gemma4_layers_test.py tests/unit/gemma4_small_layers_test.py tests/unit/qwen3_next_vs_reference_test.py tests/unit/qwen3_5_layers_test.py tests/unit/qwen3_omni_layers_test.py"
           }')[inputs.flavor] }}
         ${{ inputs.additional_pytest_args }}
 
diff --git a/scripts/diloco/run_diloco.sh b/scripts/diloco/run_diloco.sh
new file mode 100644
index 0000000000..763a01ccd4
--- /dev/null
+++ b/scripts/diloco/run_diloco.sh
@@ -0,0 +1,101 @@
+#!/bin/bash
+
+# This script launches a DiLoCo pre-training workload on a GKE cluster using XPK.
+
+set -e
+
+# --- Cluster Parameters ---
+export PROJECT_ID="${PROJECT_ID:-}"
+export CLUSTER_NAME="${CLUSTER_NAME:-}"
+export ZONE="${ZONE:-}"
+export RESERVATION="${RESERVATION:-}"  # optional
+export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # change to your own GCS bucket for logging and checkpointing
+export DATASET_PATH="${DATASET_PATH:-}" # change to your own GSC bucket for datasets. Make sure datasets exists
+export DOCKER_IMAGE="${DOCKER_IMAGE:-}" # Full path to the Docker image you pushed (e.g., gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:2026-06-22)
+export TPU_TYPE="${TPU_TYPE:-}"  # At least v5p-32 is needed to run Qwen3-30b-a3b. v5p-8 for qwen3-8b
+export NUM_SLICES="${NUM_SLICES:-}"  # you need at least two slices to let diloco take effect
+export WORKLOAD_NAME="${WORKLOAD_NAME:-$(whoami)-diloco-${TPU_TYPE}-$(date +%Y%m%d-%H%M%S)}" # this will be the name of run, for logging purposes
+
+# --- Model Parameters ---
+export MODEL_NAME="${MODEL_NAME:-}"
+export PER_DEVICE_BATCH_SIZE="${PER_DEVICE_BATCH_SIZE:-8}"
+export MAX_TARGET_LENGTH="${MAX_TARGET_LENGTH:-2048}"
+export TRAINING_STEPS="${TRAINING_STEPS:-20}"
+
+# --- DiLoCo Parameters ---
+export DILOCO_SYNC_PERIOD="${DILOCO_SYNC_PERIOD:-10}"
+export DILOCO_OUTER_LR="${DILOCO_OUTER_LR:-0.1}"
+export DILOCO_OUTER_MOMENTUM="${DILOCO_OUTER_MOMENTUM:-0.9}"
+
+# --- XLA Flags ---
+export XLA_FLAGS="${XLA_FLAGS:- \
+  --xla_tpu_scoped_vmem_limit_kib=65536 \
+--xla_tpu_bf16_emission_mode=NATIVE_EMISSION \
+--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \
+--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \
+--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \
+--xla_tpu_enable_all_gather_offload_tracing=true \
+--xla_tpu_use_tc_device_shape_on_sc=True \
+--xla_sc_disable_megacore_partitioning=True \
+--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \
+--xla_enable_async_all_gather=true \
+--xla_tpu_prefer_async_allgather_to_allreduce=true \
+--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \
+--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \
+--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \
+--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \
+--xla_tpu_enable_concurrent_sparse_core_offloading=true \
+--xla_tpu_aggressive_opt_barrier_removal=true \
+--xla_tpu_enable_offloading_gather_to_sparsecore=true \
+--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \
+--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \
+--xla_tpu_enable_sparse_core_collective_aggregator=true \
+--xla_tpu_enable_latency_hiding_layer_scheduler=true \
+--xla_tpu_scheduler_percent_shared_memory_limit=150 \
+--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \
+--xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \
+--xla_tpu_pcie_bandwidth_multiplier=0.03 \
+--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \
+--xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \
+  --xla_tpu_enable_3d_reduce_scatter_decomposer=false }"
+
+if [ "$NUM_SLICES" -lt 2 ]; then
+    echo "Warning: NUM_SLICES is less than 2. DiLoCo will not take effect."
+fi
+
+# MaxText command
+MAXTEXT_COMMAND="cd /deps/src/ && \
+LIBTPU_INIT_ARGS='${XLA_FLAGS}' \
+python3 -m maxtext.trainers.pre_train.train \
+maxtext/configs/base.yml \
+run_name=$WORKLOAD_NAME \
+save_config_to_gcs=true \
+base_output_directory=$BASE_OUTPUT_DIRECTORY \
+dataset_path=$DATASET_PATH \
+dataset_name='c4/en:3.0.1' \
+eval_dataset_name='c4/en:3.0.1' \
+model_name=$MODEL_NAME \
+tokenizer_type=huggingface \
+tokenizer_path=maxtext/assets/tokenizers/qwen3-tokenizer \
+per_device_batch_size=$PER_DEVICE_BATCH_SIZE \
+max_target_length=$MAX_TARGET_LENGTH \
+enable_diloco=true \
+dcn_diloco_parallelism=$NUM_SLICES \
+diloco_sync_period=$DILOCO_SYNC_PERIOD \
+diloco_outer_lr=$DILOCO_OUTER_LR \
+diloco_outer_momentum=$DILOCO_OUTER_MOMENTUM \
+steps=$TRAINING_STEPS"
+
+# Workload Creation
+xpk workload create \
+  --cluster="$CLUSTER_NAME" \
+  --project="$PROJECT_ID" \
+  --zone="$ZONE" \
+  --priority=medium \
+  --max-restarts=0 \
+  --tpu-type="$TPU_TYPE" \
+  --num-slices="$NUM_SLICES" \
+  --docker-image="${DOCKER_IMAGE}" \
+  --workload="${WORKLOAD_NAME}" \
+  ${RESERVATION:+--reservation="$RESERVATION"} \
+  --command="${MAXTEXT_COMMAND}"
diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py
index b8a513c1df..12bc48c676 100644
--- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py
+++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py
@@ -117,7 +117,7 @@ def init_state_fn():
       cfg.checkpoint_period,
   )
 
-  state, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
+  state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
   max_logging.log("start")
   max_utils.print_mem_stats("After params initialized")
 
diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index 0bd3b84fad..c775078f17 100644
--- a/src/maxtext/common/checkpointing.py
+++ b/src/maxtext/common/checkpointing.py
@@ -202,19 +202,27 @@ def _populate_pure_dict_from_partial(abstract_pure, partial_concrete, config=Non
   (NNX-only state the Linen checkpoint never had) get `_default_for_sds` or are preserved.
   """
   if isinstance(abstract_pure, dict):
-    return {
-        k: _populate_pure_dict_from_partial(
-            v,
-            partial_concrete.get(k) if isinstance(partial_concrete, dict) else None,
-            config,
-        )
-        for k, v in abstract_pure.items()
-    }
+    res = {}
+    for k, v in abstract_pure.items():
+      val = _populate_pure_dict_from_partial(
+          v,
+          partial_concrete.get(k) if isinstance(partial_concrete, dict) else None,
+          config,
+      )
+      if val is not None:
+        res[k] = val
+    return res
+
+  if isinstance(partial_concrete, jax.ShapeDtypeStruct):
+    partial_concrete = None
+
   if partial_concrete is not None and not isinstance(partial_concrete, dict):
     return partial_concrete
 
-  # Under LoRA, preserve the already-initialized base weights instead of resetting to zeros
+  # Under LoRA, return None for SDS to omit from state dict and preserve base weights.
   if config and config.lora and config.lora.enable_lora:
+    if isinstance(abstract_pure, jax.ShapeDtypeStruct):
+      return None
     return abstract_pure
 
   return _default_for_sds(abstract_pure)
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 9c834fb7d4..aaa7c1d8fe 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -881,7 +881,7 @@ diloco_sync_period: 36
 diloco_outer_lr: 0.3
 diloco_outer_momentum: 0.9
 # DCN bandwidth throttling parameters (used for simulating slow networks).
-# If dcn_bandwidth_limit is empty, no throttling is applied.
+# If dcn_bandwidth_limit is empty, no throttling is applied. **Per-VM** throughput limit. Example: 10gbit.
 dcn_bandwidth_limit: ""
 # The burst size parameter for the traffic control (tc) token bucket filter.
 dcn_bandwidth_burst: "10mb"
diff --git a/src/maxtext/experimental/rl/grpo_trainer.py b/src/maxtext/experimental/rl/grpo_trainer.py
index 9ea0a68f82..e1d80e8917 100644
--- a/src/maxtext/experimental/rl/grpo_trainer.py
+++ b/src/maxtext/experimental/rl/grpo_trainer.py
@@ -864,7 +864,7 @@ def init_state_fn():
 
   with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION):
     data_iterator = grpo_input_pipeline.create_data_iterator(config_inference, inference_mesh)
-    state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state(
+    state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
         data_iterator, config, mesh, checkpoint_manager, init_state_fn
     )
 
diff --git a/src/maxtext/utils/generate_param_only_checkpoint.py b/src/maxtext/utils/generate_param_only_checkpoint.py
index 8a2488abe1..a44a7819f2 100644
--- a/src/maxtext/utils/generate_param_only_checkpoint.py
+++ b/src/maxtext/utils/generate_param_only_checkpoint.py
@@ -49,6 +49,23 @@
 from maxtext.utils import model_creation_utils
 from maxtext.utils import train_utils
 
+_SCAN_AXIS_NAMES = ("layers_per_stage", "layers", "circular_repeats")
+
+
+def get_scan_axis(config, sharding_or_spec, x=None):
+  """Find the scan axis from sharding spec or use default fallback if valid."""
+  if hasattr(sharding_or_spec, "spec"):
+    spec = sharding_or_spec.spec
+  else:
+    spec = sharding_or_spec
+
+  for name in _SCAN_AXIS_NAMES:
+    if name in spec:
+      return spec.index(name)
+  if x is not None and hasattr(x, "ndim") and x.ndim > config.param_scan_axis:
+    return config.param_scan_axis
+  return None
+
 
 def _possibly_unroll_params(config, training_state, training_state_annotations, mesh):
   """Unroll scanned input layers when force_unroll is set."""
@@ -106,8 +123,17 @@ def _possibly_unroll_params_nnx(config, state, state_mesh_shardings, mesh):
   and removes the original collection. Mirrors the same operation on
   `state_mesh_shardings` so downstream sharding stays correct.
   """
-  decoder_state = state.model.decoder
-  decoder_shardings = state_mesh_shardings.model.decoder
+  param_only_state, other_state = nnx.split_state(state, nnx.Param, ...)
+  param_paths = set(param_only_state.flat_state().paths)
+
+  param_only_shardings_items = [(path, val) for path, val in state_mesh_shardings.flat_state() if path in param_paths]
+  other_shardings_items = [(path, val) for path, val in state_mesh_shardings.flat_state() if path not in param_paths]
+
+  param_only_shardings = nnx.State.from_flat_path(param_only_shardings_items)
+  other_shardings = nnx.State.from_flat_path(other_shardings_items)
+
+  decoder_state = param_only_state.model.decoder
+  decoder_shardings = param_only_shardings.model.decoder
 
   def unroll_layer_group(num_layers, layer_name="layers"):
     layers = decoder_state.get(layer_name, None)
@@ -115,19 +141,42 @@ def unroll_layer_group(num_layers, layer_name="layers"):
     if layers is None or layers_shardings is None:
       raise ValueError(f"Missing {layer_name} in NNX state.model.decoder or state_mesh_shardings.")
 
-    def drop_scan_axis(named_sharding):
+    def drop_scan_axis(named_sharding, x):
       ps = named_sharding.spec
-      return jax.sharding.PartitionSpec(*(ps[0 : config.param_scan_axis] + ps[config.param_scan_axis + 1 :]))
+      val = x[...] if isinstance(x, nnx.Variable) else x
+      scan_axis = get_scan_axis(config, named_sharding, val)
+      if scan_axis is not None:
+        return jax.sharding.PartitionSpec(*(ps[0:scan_axis] + ps[scan_axis + 1 :]))
+      return ps
 
     new_layer_pspec = jax.tree_util.tree_map(
-        drop_scan_axis, layers_shardings, is_leaf=lambda x: isinstance(x, jax.sharding.NamedSharding)
+        drop_scan_axis,
+        layers_shardings,
+        layers,
+        is_leaf=lambda x: isinstance(x, (jax.sharding.NamedSharding, nnx.Variable)),
     )
     new_layer_sharding = jax.tree_util.tree_map(lambda ps: jax.sharding.NamedSharding(mesh, ps), new_layer_pspec)
 
     for i in range(num_layers):
 
+      # pylint: disable=cell-var-from-loop
       def slice_ith(input_layers):
-        return jax.tree_util.tree_map(lambda x: jnp.take(x, i, axis=config.param_scan_axis), input_layers)
+        def _slice_leaf(x, sharding):
+          val = x[...] if isinstance(x, nnx.Variable) else x
+          scan_axis = get_scan_axis(config, sharding, val)
+          if scan_axis is not None:
+            sliced_val = jnp.take(val, i, axis=scan_axis)
+            if isinstance(x, nnx.Variable):
+              return type(x)(sliced_val)
+            return sliced_val
+          return x
+
+        return jax.tree_util.tree_map(
+            _slice_leaf,
+            input_layers,
+            layers_shardings,
+            is_leaf=lambda x: isinstance(x, (jax.sharding.NamedSharding, nnx.Variable)),
+        )
 
       # pylint: disable=not-callable
       new_layer = jax.jit(slice_ith, out_shardings=new_layer_sharding)(layers)
@@ -145,6 +194,13 @@ def slice_ith(input_layers):
   else:
     unroll_layer_group(config.num_decoder_layers, layer_name="layers")
 
+  # Merge modified parameter state and other variables back in-place
+  merged_state = nnx.merge_state(param_only_state, other_state)
+  state.update(merged_state)
+
+  merged_shardings = nnx.merge_state(param_only_shardings, other_shardings)
+  state_mesh_shardings.update(merged_shardings)
+
 
 def _read_train_checkpoint(config, checkpoint_manager, mesh):
   """Read training checkpoint at path defined by load_full_state_path."""
@@ -167,13 +223,14 @@ def init_state_fn():
     tx = optimizers.get_optimizer(config, learning_rate_schedule)
     init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, rng)
 
-  state, state_mesh_notations, state_mesh_shardings, _ = maxtext_utils.setup_training_state(
+  state, state_mesh_notations, state_mesh_shardings, _, _ = maxtext_utils.setup_training_state(
       None, config, mesh, checkpoint_manager, init_state_fn
   )
   if config.pure_nnx:
     # On NNX, state is a flat nnx.State; params live under state.model and the
     # legacy notations are unused (callers receive shardings directly).
-    num_params = max_utils.calculate_num_params_from_pytree(state.model)
+    params, _ = nnx.split_state(state.model, nnx.Param, ...)
+    num_params = max_utils.calculate_num_params_from_pytree(params)
     max_logging.log(f"In input checkpoint Number of model params={num_params/1e9:.3f} billion")
     return state, state_mesh_shardings
   num_params = max_utils.calculate_num_params_from_pytree(state.params)
@@ -283,16 +340,36 @@ def unroll_layer_group(num_layers, layer_name="layers"):
     if layers is None or layers_annotations is None:
       return  # No LoRA on this layer group; nothing to unroll.
 
-    def new_pspec(x):
-      return jax.sharding.PartitionSpec(*(x[0 : config.param_scan_axis] + x[config.param_scan_axis + 1 :]))
-
-    new_layer_annotation = jax.tree_util.tree_map(new_pspec, layers_annotations)
+    def drop_scan_axis(spec, x):
+      scan_axis = get_scan_axis(config, spec, x)
+      if scan_axis is not None:
+        return jax.sharding.PartitionSpec(*(spec[0:scan_axis] + spec[scan_axis + 1 :]))
+      return spec
+
+    new_layer_annotation = jax.tree_util.tree_map(
+        drop_scan_axis,
+        layers_annotations,
+        layers,
+        is_leaf=lambda x: isinstance(x, jax.sharding.PartitionSpec),
+    )
     new_layer_sharding = jax.tree_util.tree_map(lambda x: jax.sharding.NamedSharding(mesh, x), new_layer_annotation)
 
     for i in range(num_layers):
 
+      # pylint: disable=cell-var-from-loop
       def slice_ith(input_layers):
-        return jax.tree_util.tree_map(lambda x: jnp.take(x, i, axis=config.param_scan_axis), input_layers)
+        def _slice_leaf(x, spec):
+          scan_axis = get_scan_axis(config, spec, x)
+          if scan_axis is not None:
+            return jnp.take(x, i, axis=scan_axis)
+          return x
+
+        return jax.tree_util.tree_map(
+            _slice_leaf,
+            input_layers,
+            layers_annotations,
+            is_leaf=lambda x: isinstance(x, jax.sharding.PartitionSpec),
+        )
 
       # pylint: disable=not-callable
       new_layer = jax.jit(slice_ith, out_shardings=new_layer_sharding)(layers)
@@ -400,13 +477,9 @@ def generate_decode_checkpoint(config):
   training_state, training_state_annotations = _read_train_checkpoint(config, checkpoint_manager, mesh)
   if config.pure_nnx:
     # NNX state is a flat nnx.State; opt_state lives under the optimizer sub-state.
-    assert (
-        training_state.optimizer.opt_state
-    ), "missing opt_state in training checkpoint"
+    assert training_state.optimizer.opt_state, "missing opt_state in training checkpoint"
   else:
-    assert (
-        training_state.opt_state != {}
-    ), "missing opt_state in training checkpoint"
+    assert training_state.opt_state != {}, "missing opt_state in training checkpoint"
 
   _possibly_unroll_params(config, training_state, training_state_annotations, mesh)
 
diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py
index 930ff90540..8b8a5d1208 100644
--- a/src/maxtext/utils/maxtext_utils.py
+++ b/src/maxtext/utils/maxtext_utils.py
@@ -1392,7 +1392,9 @@ def get_abstract_param(model, config):
       {"params": key, "dropout": key, "aqt": key},
       np.ones(input_shape, dtype=jnp.int32),
       np.ones(input_shape, dtype=jnp.int32),
-      encoder_images=np.ones(image_shape, dtype=jnp.int32) if config.use_multimodal else None,  # pyrefly: ignore[no-matching-overload]
+      encoder_images=np.ones(image_shape, dtype=jnp.int32)
+      if config.use_multimodal
+      else None,  # pyrefly: ignore[no-matching-overload]
       encoder_audios=np.ones(audio_shape, dtype=jnp.float32) if config.use_audio else None,
   )
   return abstract_vars
@@ -1413,7 +1415,7 @@ def setup_decode_state(config, mesh, checkpoint_manager, init_state_fn):
   if not config.load_parameters_path:
     # generate random params
     max_logging.log("No decode checkpoint specified - generating random weights.")
-    state, state_mesh_annotations, _, _ = setup_initial_state(
+    state, state_mesh_annotations, _, _, _ = setup_initial_state(
         None, config, mesh, checkpoint_manager, init_state_fn, False
     )
   else:
@@ -1468,6 +1470,9 @@ def setup_initial_state(
   Returns:
     train_state: the initialized train state. For NNX, this is a TrainStateNNX instance
     state_mesh_annotations: the mesh annotations for the train state
+    state_mesh_shardings: the mesh shardings for the train state
+    data_iterator: the updated data iterator
+    was_restored: True if state or params were restored from checkpoint, False if freshly initialized
   """
 
   unboxed_abstract_state, state_mesh_annotations, state_mesh_shardings = get_abstract_state(
@@ -1493,6 +1498,8 @@ def setup_initial_state(
         expansion_factor_real_data=config.expansion_factor_real_data,
         maxtext_config=config,
     )
+    # Partial or fully restored
+    was_restored = bool(restored is not None or raw_params is not None)
 
     if restored:
       if isinstance(
@@ -1507,10 +1514,19 @@ def setup_initial_state(
         # The update of data_iterator state happens in place, no need to assign explicitly
         state = restored["items"]
 
-      # For NNX, convert the pure dict to nnx.State using the abstract state as template
+      # For NNX, convert pure dict to nnx.State. Under LoRA, overwrite concrete initial state.
       if config.pure_nnx:
-        nnx.replace_by_pure_dict(unboxed_abstract_state, state)
-        state = unboxed_abstract_state
+        if config.lora.enable_lora:
+          concrete_state = jax.jit(
+              lambda: nnx.state(init_state_fn()),
+              in_shardings=None,
+              out_shardings=state_mesh_shardings,
+          )()
+          nnx.replace_by_pure_dict(concrete_state, state)
+          state = concrete_state
+        else:
+          nnx.replace_by_pure_dict(unboxed_abstract_state, state)
+          state = unboxed_abstract_state
     else:
       init_state_partial = init_state_fn
       init_state_partial.__name__ = "initialize_state"
@@ -1548,7 +1564,7 @@ def _merge_params(p_raw, p_init):
             state = state.replace(params=raw_params)
   if not config.pure_nnx:
     state = max_utils.unbox_logicallypartioned(state)
-  return state, state_mesh_annotations, state_mesh_shardings, data_iterator
+  return state, state_mesh_annotations, state_mesh_shardings, data_iterator, was_restored
 
 
 def get_logical_annotations(config, mesh, init_state_fn):
diff --git a/src/maxtext/utils/standalone_checkpointer.py b/src/maxtext/utils/standalone_checkpointer.py
index 6b1aa264c2..243ad1dd3e 100644
--- a/src/maxtext/utils/standalone_checkpointer.py
+++ b/src/maxtext/utils/standalone_checkpointer.py
@@ -25,7 +25,6 @@
 
 from absl import app
 from flax import nnx
-from flax.linen import partitioning as nn_partitioning
 import jax
 from jax import numpy as jnp
 from maxtext.configs import pyconfig
@@ -69,36 +68,25 @@ def init_state_fn():
     mesh = model.mesh
     _, tx = train_utils.create_training_optimizer(config, model)
     init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, init_rng)
+
   checkpoint_manager = train_utils.create_checkpoint_manager(config, mesh, init_state_fn)
 
-  unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, is_training=True)
   # A barrier to sync all hosts before starting to restore checkpoint
   jax.experimental.multihost_utils.sync_global_devices("Barrier before load")
-  checkpoint_load_start = datetime.datetime.now()
-  with nn_partitioning.axis_rules(config.logical_axis_rules):
-    state, _ = checkpointing.load_state_if_possible(
-        checkpoint_manager,
-        None,
-        config.load_parameters_path,
-        config.load_full_state_path,
-        config.checkpoint_storage_concurrent_gb,
-        unboxed_abstract_state,
-        use_ocdbt=config.checkpoint_storage_use_ocdbt,
-        use_zarr3=config.checkpoint_storage_use_zarr3,
-    )
-    if state:
-      state = state["items"]
 
+  checkpoint_load_start = datetime.datetime.now()
+  # Delegate checkpoint restoration or state initialization to setup_training_state
+  state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, config, mesh, checkpoint_manager, init_state_fn)
   jax.block_until_ready(state)
   checkpoint_load_end = datetime.datetime.now()
-  if state is not None:  # Checkpoint was available for restore
+
+  if was_restored:
     if jax.process_index() == 0:
       max_logging.log(
           "STANDALONE CHECKPOINTER : Checkpoint restored in :" f" {checkpoint_load_end - checkpoint_load_start}"
       )
-  else:  # Checkpoint was unavailable, state needs to be initialized
-    state, _, _, _ = maxtext_utils.setup_training_state(None, config, mesh, checkpoint_manager, init_state_fn)
-  state = add_entropy_to_checkpoint(state)
+  else:  # Checkpoint was unavailable, fresh state needs to be perturbed with entropy
+    state = add_entropy_to_checkpoint(state)
 
   start_step = get_first_step(model, state)  # this is the start_step for training
   for step in np.arange(start_step, config.steps):
@@ -106,7 +94,8 @@ def init_state_fn():
       start_time = datetime.datetime.now()
       # A barrier to sync all hosts before starting to save checkpoint
       jax.experimental.multihost_utils.sync_global_devices("Barrier before save")
-      if checkpointing.save_checkpoint(checkpoint_manager, int(step), state):
+      state_to_save = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) if config.pure_nnx else state
+      if checkpointing.save_checkpoint(checkpoint_manager, int(step), state_to_save):
         checkpoint_manager.wait_until_finished()
         end_time = datetime.datetime.now()
         if jax.process_index() == 0:
diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py
index 9aff1e3aeb..8e9f1d84c5 100644
--- a/src/maxtext/utils/train_utils.py
+++ b/src/maxtext/utils/train_utils.py
@@ -304,7 +304,7 @@ def create_train_state_fn():
     # Create data_loader AFTER reordering wrapper is applied
     data_loader = create_dataloader(config, mesh, data_iterator, recorder, rampup_manager)
 
-    state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state(
+    state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state(
         data_iterator, config, mesh, checkpoint_manager, init_state_fn
     )
     if config.pure_nnx:
diff --git a/tests/integration/peft_integration_test.py b/tests/integration/peft_integration_test.py
new file mode 100644
index 0000000000..56bc617a44
--- /dev/null
+++ b/tests/integration/peft_integration_test.py
@@ -0,0 +1,375 @@
+# Copyright 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.
+
+"""Integration tests for Flax NNX PEFT, LoRA, and QLoRA."""
+
+import os
+import sys
+import unittest
+import tempfile
+import json
+import shutil
+import pytest
+import subprocess
+
+from maxtext.utils.globals import MAXTEXT_ASSETS_ROOT
+from tests.utils.test_helpers import get_test_config_path
+
+
+def _get_jax_backend():
+  """Determine JAX backend by running a quick subprocess to avoid locking TPU in the parent process."""
+  try:
+    result = subprocess.run(
+        [sys.executable, "-c", "import jax; print(jax.default_backend())"], capture_output=True, text=True, check=True
+    )
+    return result.stdout.strip()
+  except Exception:  # pylint: disable=broad-exception-caught
+    return "cpu"
+
+
+def _get_jax_device_count():
+  """Determine JAX device count by running a quick subprocess."""
+  try:
+    result = subprocess.run(
+        [sys.executable, "-c", "import jax; print(jax.device_count())"], capture_output=True, text=True, check=True
+    )
+    return int(result.stdout.strip())
+  except Exception:  # pylint: disable=broad-exception-caught
+    return 1
+
+
+def _get_common_args(temp_dir, model_name, steps=1, overrides=None):
+  """Build common arguments for downscaled integration runs."""
+  tokenizer_path = os.path.join(MAXTEXT_ASSETS_ROOT, "tokenizers", "tokenizer.llama2")
+
+  backend = _get_jax_backend()
+  num_devices = _get_jax_device_count()
+
+  hardware = "cpu"
+  if backend == "tpu":
+    hardware = "tpu"
+  elif backend == "gpu":
+    hardware = "gpu"
+
+  args = [
+      None,
+      get_test_config_path(),
+      f"run_name=test_peft_{model_name}",
+      f"model_name={model_name}",
+      f"steps={steps}",
+      f"base_output_directory={temp_dir}",
+      "dataset_type=synthetic",
+      "per_device_batch_size=1",
+      "max_target_length=32",
+      f"tokenizer_path={tokenizer_path}",
+      "enable_goodput_recording=False",
+      "enable_checkpoint_cloud_logger=False",
+      "monitor_goodput=False",
+      "sharding_tolerance=1.0",  # Bypass sharding checks for downscaled test models
+      f"hardware={hardware}",
+      "skip_jax_distributed_system=True",
+      "ici_fsdp_parallelism=1",
+      f"ici_data_parallelism={num_devices}",
+  ]
+  if overrides:
+    args.extend(overrides)
+  return args
+
+
+def _run_train_subprocess(args):
+  """Run the pre-training trainer inside an isolated Python subprocess."""
+  cmd = [sys.executable, "-m", "maxtext.trainers.pre_train.train", args[1]] + args[2:]
+  env = os.environ.copy()
+  env["PYTHONPATH"] = os.path.abspath("src")
+
+  backend = _get_jax_backend()
+
+  if backend not in ("tpu", "gpu"):
+    env["JAX_PLATFORMS"] = "cpu"
+
+  subprocess.run(cmd, env=env, check=True)
+
+
+def _run_sft_subprocess(args):
+  """Run the SFT native trainer inside an isolated Python subprocess."""
+  cmd = [sys.executable, "-m", "maxtext.trainers.post_train.sft.train_sft_native", args[1]] + args[2:]
+  env = os.environ.copy()
+  env["PYTHONPATH"] = os.path.abspath("src")
+
+  backend = _get_jax_backend()
+
+  if backend not in ("tpu", "gpu"):
+    env["JAX_PLATFORMS"] = "cpu"
+
+  subprocess.run(cmd, env=env, check=True)
+
+
+@pytest.mark.integration_test
+class PEFTIntegrationTest(unittest.TestCase):
+  """Integration checks for pre-training and native SFT with LoRA & QLoRA."""
+
+  def setUp(self):
+    self.temp_dirs = []
+
+  def tearDown(self):
+    for temp_dir in self.temp_dirs:
+      if os.path.exists(temp_dir):
+        shutil.rmtree(temp_dir)
+
+  def _create_temp_dir(self):
+    temp_dir = tempfile.mkdtemp()
+    self.temp_dirs.append(temp_dir)
+    return temp_dir
+
+  def test_native_sft_gemma4_lora(self):
+    """Test native SFT with Gemma4-e2b and LoRA enabled."""
+    temp_dir = self._create_temp_dir()
+    overrides = [
+        "override_model_config=True",
+        "scan_layers=False",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=256",
+        "base_mlp_dim=512",
+        "num_kv_shared_layers=1",
+    ]
+    args = _get_common_args(temp_dir, "gemma4-e2b", steps=2, overrides=overrides)
+    _run_sft_subprocess(args)
+
+  def test_native_sft_gemma4_qlora(self):
+    """Test native SFT with Gemma4-e2b and QLoRA enabled."""
+    temp_dir = self._create_temp_dir()
+    overrides = [
+        "override_model_config=True",
+        "scan_layers=False",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_weight_qtype=int8",
+        "lora.lora_tile_size=32",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=256",
+        "base_mlp_dim=512",
+        "num_kv_shared_layers=1",
+    ]
+    args = _get_common_args(temp_dir, "gemma4-e2b", steps=2, overrides=overrides)
+    _run_sft_subprocess(args)
+
+  def test_native_sft_qwen3_lora(self):
+    """Test native SFT with Qwen3-0.6B and LoRA enabled."""
+    temp_dir = self._create_temp_dir()
+    overrides = [
+        "override_model_config=True",
+        "scan_layers=True",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=256",
+        "base_mlp_dim=512",
+    ]
+    args = _get_common_args(temp_dir, "qwen3-0.6b", steps=2, overrides=overrides)
+    _run_sft_subprocess(args)
+
+  def test_native_sft_qwen3_qlora(self):
+    """Test native SFT with Qwen3-0.6B and QLoRA enabled."""
+    temp_dir = self._create_temp_dir()
+    overrides = [
+        "override_model_config=True",
+        "scan_layers=True",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_weight_qtype=int8",
+        "lora.lora_tile_size=32",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=256",
+        "base_mlp_dim=512",
+    ]
+    args = _get_common_args(temp_dir, "qwen3-0.6b", steps=2, overrides=overrides)
+    _run_sft_subprocess(args)
+
+  def test_pretrain_lora_checkpoint_save_and_restore(self):
+    """Test pre-training with LoRA, checkpoint saving, and subsequent restoration."""
+    temp_dir = self._create_temp_dir()
+    metrics_file_saved = os.path.join(temp_dir, "saved_metrics.txt")
+    metrics_file_restored = os.path.join(temp_dir, "restored_metrics.txt")
+
+    # Common model scaling/configs
+    overrides_base = [
+        "scan_layers=True",
+        "attention=dot_product",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=128",
+        "base_mlp_dim=256",
+        "ici_fsdp_parallelism=1",
+    ]
+
+    # 1. First run: train 1 step with checkpointing enabled
+    args_save = _get_common_args(
+        temp_dir,
+        "default",
+        steps=1,
+        overrides=overrides_base
+        + [
+            "enable_checkpointing=True",
+            "checkpoint_period=1",
+            "async_checkpointing=False",
+            f"metrics_file={metrics_file_saved}",
+        ],
+    )
+    _run_train_subprocess(args_save)
+
+    # Verify checkpoint got written (MaxText writes to base_output_directory/run_name/checkpoints/0/items)
+    checkpoint_dir = os.path.join(temp_dir, "test_peft_default", "checkpoints", "0", "items")
+    self.assertTrue(os.path.exists(checkpoint_dir), f"Checkpoint directory {checkpoint_dir} was not created.")
+
+    # 2. Second run: restore from the checkpoint directory and train 1 step (total steps = 2)
+    args_restore = _get_common_args(
+        temp_dir,
+        "default",
+        steps=2,
+        overrides=overrides_base
+        + [
+            "enable_checkpointing=True",
+            "checkpoint_period=1",
+            "async_checkpointing=False",
+            f"load_parameters_path={checkpoint_dir}",
+            f"metrics_file={metrics_file_restored}",
+        ],
+    )
+    _run_train_subprocess(args_restore)
+
+    # 3. Verify losses are logged and correct
+    self.assertTrue(os.path.exists(metrics_file_saved))
+    self.assertTrue(os.path.exists(metrics_file_restored))
+
+    with open(metrics_file_saved, "r", encoding="utf8") as f:
+      saved_metrics = [json.loads(line) for line in f.readlines()]
+    with open(metrics_file_restored, "r", encoding="utf8") as f:
+      restored_metrics = [json.loads(line) for line in f.readlines()]
+
+    print("SAVED METRICS:", saved_metrics)
+    print("RESTORED METRICS:", restored_metrics)
+
+    self.assertEqual(len(saved_metrics), 1, "First run should have exactly 1 step.")
+    self.assertEqual(len(restored_metrics), 1, "Second run should have exactly 1 step.")
+
+    saved_step0_loss = saved_metrics[0]["learning/loss"]
+    restored_step1_loss = restored_metrics[0]["learning/loss"]
+
+    print(f"Saved run Step 0 Loss: {saved_step0_loss}")
+    print(f"Restored run Step 1 Loss: {restored_step1_loss}")
+
+    # Verify that the restored run resumed from step 0 (logged step 1)
+    self.assertEqual(saved_metrics[0]["step"], 0.0)
+    self.assertEqual(restored_metrics[0]["step"], 1.0)
+
+  def test_pretrain_qlora_checkpoint_save_and_restore(self):
+    """Test pre-training with QLoRA, checkpoint saving, and subsequent restoration."""
+    temp_dir = self._create_temp_dir()
+    metrics_file_saved = os.path.join(temp_dir, "saved_metrics_qlora.txt")
+    metrics_file_restored = os.path.join(temp_dir, "restored_metrics_qlora.txt")
+
+    # Common model scaling/configs
+    overrides_base = [
+        "scan_layers=True",
+        "attention=dot_product",
+        "weight_dtype=bfloat16",
+        "dtype=bfloat16",
+        "lora.enable_lora=True",
+        "lora.lora_weight_qtype=int8",
+        "lora.lora_tile_size=32",
+        "lora.lora_rank=4",
+        "lora.lora_alpha=8.0",
+        "base_num_decoder_layers=2",
+        "base_emb_dim=128",
+        "base_mlp_dim=256",
+        "ici_fsdp_parallelism=1",
+    ]
+
+    # 1. First run: train 1 step with checkpointing enabled
+    args_save = _get_common_args(
+        temp_dir,
+        "default",
+        steps=1,
+        overrides=overrides_base
+        + [
+            "enable_checkpointing=True",
+            "checkpoint_period=1",
+            "async_checkpointing=False",
+            f"metrics_file={metrics_file_saved}",
+        ],
+    )
+    _run_train_subprocess(args_save)
+
+    # Verify checkpoint got written (MaxText writes to base_output_directory/run_name/checkpoints/0/items)
+    checkpoint_dir = os.path.join(temp_dir, "test_peft_default", "checkpoints", "0", "items")
+    self.assertTrue(os.path.exists(checkpoint_dir), f"Checkpoint directory {checkpoint_dir} was not created.")
+
+    # 2. Second run: restore from the checkpoint directory and train 1 step (total steps = 2)
+    args_restore = _get_common_args(
+        temp_dir,
+        "default",
+        steps=2,
+        overrides=overrides_base
+        + [
+            "enable_checkpointing=True",
+            "checkpoint_period=1",
+            "async_checkpointing=False",
+            f"load_parameters_path={checkpoint_dir}",
+            f"metrics_file={metrics_file_restored}",
+        ],
+    )
+    _run_train_subprocess(args_restore)
+
+    # 3. Verify losses are logged and correct
+    self.assertTrue(os.path.exists(metrics_file_saved))
+    self.assertTrue(os.path.exists(metrics_file_restored))
+
+    with open(metrics_file_saved, "r", encoding="utf8") as f:
+      saved_metrics = [json.loads(line) for line in f.readlines()]
+    with open(metrics_file_restored, "r", encoding="utf8") as f:
+      restored_metrics = [json.loads(line) for line in f.readlines()]
+
+    print("SAVED METRICS (QLoRA):", saved_metrics)
+    print("RESTORED METRICS (QLoRA):", restored_metrics)
+
+    self.assertEqual(len(saved_metrics), 1, "First run should have exactly 1 step.")
+    self.assertEqual(len(restored_metrics), 1, "Second run should have exactly 1 step.")
+
+    saved_step0_loss = saved_metrics[0]["learning/loss"]
+    restored_step1_loss = restored_metrics[0]["learning/loss"]
+
+    print(f"Saved run Step 0 Loss (QLoRA): {saved_step0_loss}")
+    print(f"Restored run Step 1 Loss (QLoRA): {restored_step1_loss}")
+
+    # Verify that the restored run resumed from step 0 (logged step 1)
+    self.assertEqual(saved_metrics[0]["step"], 0.0)
+    self.assertEqual(restored_metrics[0]["step"], 1.0)
diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py
index 670fcce5e1..cce1da530d 100644
--- a/tests/unit/maxtext_utils_test.py
+++ b/tests/unit/maxtext_utils_test.py
@@ -387,7 +387,8 @@ def create_train_state_fn():
       init_state_fn = create_train_state_fn
     else:
       init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, tx, self.config, True, rng)
-    state, _, _, _ = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn)
+    state, _, _, _, was_restored = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn)
+    self.assertFalse(was_restored)
     if self.config.pure_nnx:
       self.assertIsNotNone(state.optimizer)
     else:
@@ -1413,7 +1414,8 @@ def create_train_state_fn():
           True,
           rng,
       )
-    state, _, _, _ = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn)
+    state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn)
+    self.assertFalse(was_restored)
     if self.config.pure_nnx:
       self.assertIsNotNone(state.optimizer)
     else:
diff --git a/tests/unit/qwen3_omni_layers_test.py b/tests/unit/qwen3_omni_layers_test.py
index d27190dbd1..b391dce5f2 100644
--- a/tests/unit/qwen3_omni_layers_test.py
+++ b/tests/unit/qwen3_omni_layers_test.py
@@ -25,6 +25,7 @@
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
+import pytest
 from maxtext.configs import pyconfig
 from maxtext.common import common_types
 from maxtext.utils.globals import MAXTEXT_REPO_ROOT
@@ -62,7 +63,6 @@
     copy_patch_embed_weights,
     copy_patch_merger_weights,
     copy_vision_encoder_weights,
-    create_block_diagonal_attention_mask,
     create_random_jax_torch,
 )
 import numpy as np
@@ -777,6 +777,7 @@ def test_hidden_states_unchanged_without_visual_tokens(self):
     np.testing.assert_allclose(np.array(result), hidden_np, rtol=1e-6, atol=1e-6)
 
 
+@pytest.mark.skip(reason="Requires decord, which may not be installed in remote CI runners.")
 class TestQwen3OmniPreprocessing(unittest.TestCase):
   """Test MaxText Qwen3 Omni preprocessor against HuggingFace reference."""
 
@@ -1003,17 +1004,12 @@ def _test_encoder_layer_with_batch_size(self, batch_size):
 
     jax_input, torch_input_3d = create_random_jax_torch(batch_size, seq_len, hidden_size)
 
-    # PyTorch forward pass - expects 2D input (total_seq_len, hidden_dim) with cu_seqlens
-    torch_input_2d = torch_input_3d.reshape(-1, hidden_size)
-
-    # Create cu_seqlens for PyTorch (cumulative sequence lengths for each batch)
-    # For batch_size=2, seq_len=12: [0, 12, 24] indicates two sequences of length 12 each
-    cu_seqlens = torch.tensor([i * seq_len for i in range(batch_size + 1)], dtype=torch.int32)
-
-    attention_mask = create_block_diagonal_attention_mask(cu_seqlens, torch_input_2d.dtype)
-
-    torch_output_1d = torch_layer(torch_input_2d, cu_seqlens=cu_seqlens, attention_mask=attention_mask)[0]
-    torch_output = torch_output_1d.reshape(batch_size, seq_len, hidden_size)
+    # PyTorch audio layers take 2D packed input. Run each batch item separately
+    # to match MaxText's batched attention without relying on HF's old mask API.
+    cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32)
+    torch_output = torch.stack(
+        [torch_layer(torch_input_3d[i], cu_seqlens=cu_seqlens)[0] for i in range(batch_size)], dim=0
+    )
 
     jax_output = maxtext_layer(jax_input, deterministic=True)
 
@@ -1144,18 +1140,15 @@ def test_audio_encoder_matches_torch(self):
     torch_after_pos = torch_conv_out + torch_pos_emb
 
     # Run through encoder layers + layernorm (but not projector)
-    # Process all chunks together
+    # Process chunks separately, matching MaxText's (batch * chunks, seq, hidden) attention shape.
     seq_len_per_chunk = torch_after_pos.shape[1]
-    cu_seqlens = torch.tensor([i * seq_len_per_chunk for i in range(num_chunks + 1)], dtype=torch.int32)
-    attention_mask = create_block_diagonal_attention_mask(cu_seqlens, torch_after_pos.dtype)
-
-    # Flatten: (num_chunks, seq_len_per_chunk, hidden) -> (num_chunks*seq_len_per_chunk, hidden)
-    hidden_state = torch_after_pos.reshape(-1, torch_after_pos.shape[-1])
+    cu_seqlens = torch.tensor([0, seq_len_per_chunk], dtype=torch.int32)
+    hidden_state = torch_after_pos
     for layer in torch_model.layers:
-      hidden_state = layer(hidden_state, cu_seqlens=cu_seqlens, attention_mask=attention_mask)[0]
+      hidden_state = torch.stack([layer(chunk, cu_seqlens=cu_seqlens)[0] for chunk in hidden_state], dim=0)
     hidden_state = torch_model.ln_post(hidden_state)
 
-    # Reshape back: (num_chunks*seq_len_per_chunk, hidden) -> (batch=1, num_chunks*seq_len_per_chunk, hidden)
+    # Reshape back: (num_chunks, seq_len_per_chunk, hidden) -> (batch=1, num_chunks*seq_len_per_chunk, hidden)
     torch_output = hidden_state.reshape(1, num_chunks * seq_len_per_chunk, -1)
 
     # MaxText forward