diff --git a/.github/workflows/gemini_invoke.yml b/.github/workflows/gemini_invoke.yml
index c3c6b9bebd..16e50eebd5 100644
--- a/.github/workflows/gemini_invoke.yml
+++ b/.github/workflows/gemini_invoke.yml
@@ -86,7 +86,7 @@ jobs:
           settings: |-
             {
               "model": {
-                "maxSessionTurns": 25
+                "maxSessionTurns": 100
               },
               "telemetry": {
                 "enabled": false,
diff --git a/.github/workflows/gemini_review.yml b/.github/workflows/gemini_review.yml
index 1babf822ba..0e518de050 100644
--- a/.github/workflows/gemini_review.yml
+++ b/.github/workflows/gemini_review.yml
@@ -36,7 +36,7 @@ defaults:
 jobs:
   review:
     runs-on: 'ubuntu-latest'
-    timeout-minutes: 10
+    timeout-minutes: 30
     permissions:
       contents: 'read'
       id-token: 'write'
@@ -103,7 +103,7 @@ jobs:
           settings: |-
             {
               "model": {
-                "maxSessionTurns": 25
+                "maxSessionTurns": 100
               },
               "telemetry": {
                 "enabled": false,
diff --git a/.github/workflows/promote_docker_image.yml b/.github/workflows/promote_docker_image.yml
index f9499c6742..4bbfd42b57 100644
--- a/.github/workflows/promote_docker_image.yml
+++ b/.github/workflows/promote_docker_image.yml
@@ -40,12 +40,14 @@ jobs:
           DAG_RUN_ID: ${{ github.event.client_payload.dag_run_id }}
           SHA: ${{ github.event.client_payload.sha }}
           GITHUB_RUN_ID: ${{ github.event.client_payload.github_run_id }}
+          TEST_TYPE: ${{ github.event.client_payload.test_type }}
         run: |
           echo "================================"
           echo "Github Run ID: ${GITHUB_RUN_ID}"
           echo "DAG ID:        ${DAG_ID}"
           echo "DAG Run ID:    ${DAG_RUN_ID}"
           echo "Commit SHA:    ${SHA}"
+          echo "Test Type:     ${TEST_TYPE}"
           echo "State:         ${STATE}"
           echo "================================"
 
@@ -66,16 +68,17 @@ jobs:
   tag_docker_image:
     name: Promote ${{ matrix.image_name }} Docker Image
     needs: handle_result
-    if: needs.handle_result.result == 'success'
+    if: needs.handle_result.result == 'success' && (github.event.client_payload.test_type == '' || github.event.client_payload.test_type == matrix.test_type)
     runs-on: linux-x86-n2-16-buildkit
     container: google/cloud-sdk:524.0.0
     strategy:
       fail-fast: false
       matrix:
-        image_name:
-          - maxtext_jax_nightly
-          - maxtext_gpu_jax_nightly
-          - maxtext_post_training_nightly
+        include:
+          - test_type: pre_training
+            image_name: maxtext_jax_nightly
+          - test_type: post_training
+            image_name: maxtext_post_training_nightly
     steps:
       - name: Configure Docker
         run: gcloud auth configure-docker us-docker.pkg.dev,gcr.io -q
@@ -83,8 +86,10 @@ jobs:
         shell: bash
         env:
           GITHUB_RUN_ID: ${{ github.event.client_payload.github_run_id }}
+          PROJECT_NAME: ${{ vars.PROJECT_NAME }}
+          IMAGE_NAME: ${{ matrix.image_name }}
         run: |
-          SOURCE_IMAGE="gcr.io/${{ vars.PROJECT_NAME }}/${{ matrix.image_name }}"
+          SOURCE_IMAGE="gcr.io/${PROJECT_NAME}/${IMAGE_NAME}"
 
           # Add the traceability tag to confirm it passed validation suite
           gcloud container images add-tag "${SOURCE_IMAGE}:${GITHUB_RUN_ID}" \
diff --git a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py
index 7e9784e516..93294454dc 100644
--- a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py
+++ b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py
@@ -48,7 +48,7 @@
   
   [Mode 2: MaxText Architecture]
     python -m maxtext.checkpoint_conversion.inspect_checkpoint maxtext \
-        model_name=<maxtext_model_name> scan_layers=<True | False> 
+        model_name=<maxtext_model_name> scan_layers=<True | False> enable_nnx=<True | False>
         (Optional: other maxtext config)
   
   [Mode 3: Orbax]
@@ -68,6 +68,8 @@
 import absl
 from maxtext.inference.inference_utils import str2bool
 from maxtext.checkpoint_conversion.utils.utils import print_peak_memory
+from maxtext.utils.model_creation_utils import create_nnx_abstract_model
+from flax import nnx
 
 
 def natural_sort_key(s: str):
@@ -238,15 +240,19 @@ def inspect_maxtext(args, remaining_args):
   print(argv)
   config = pyconfig.initialize(argv)
 
-  print(f"\n--- Inspecting MaxText Architecture: {config.model_name} (Scan: {config.scan_layers}) ---")
+  print(
+      f"\n--- Inspecting MaxText Architecture: {config.model_name} "
+      f"(scan_layers: {config.scan_layers}, enable_nnx: {config.enable_nnx}) ---"
+  )
   devices_array = maxtext_utils.create_device_mesh(config)
   mesh = jax.sharding.Mesh(devices_array, config.mesh_axes)
-  quant = quantizations.configure_quantization(config)
-  model = Transformer(config, mesh=mesh, quant=quant)
-
-  # Extract abstract parameters. This returns a PyTree of `ShapeDtypeStruct`
-  # objects, without materializing weights.
-  abstract_param = maxtext_utils.get_abstract_param(model, config)
+  if config.enable_nnx:
+    _, abstract_model = create_nnx_abstract_model(config, mesh=mesh)
+    _, abstract_param, _ = nnx.split(abstract_model, nnx.Param, ...)
+  else:
+    quant = quantizations.configure_quantization(config)
+    model = Transformer(config, mesh=mesh, quant=quant)
+    abstract_param = maxtext_utils.get_abstract_param(model, config)
 
   # Calculate and display the total parameter count based purely on abstract shapes.
   num_params = max_utils.calculate_num_params_from_pytree(abstract_param)
@@ -261,8 +267,14 @@ def inspect_maxtext(args, remaining_args):
   for path_tuple, abstract_leaf_value in abstract_params_flat:
     key_parts = param_key_parts_from_path(path_tuple)
 
-    # Construct a MaxText-style parameter key (e.g., "params.params.layer.weight").
-    param_key = "params." + ".".join(key_parts)
+    # Construct a MaxText-style parameter key. Examples:
+    # "params.params.decoder.decoder_norm.scale" (for standard model weights)
+    # "params.Tid2EidVar.decoder.layers_0.mlp.MoeBlock_0.tid2eid" (for legacy custom collections)
+    key_str = ".".join(key_parts)
+    if config.enable_nnx and not key_str.startswith(("params", "Tid2EidVar")):
+      param_key = "params.params." + key_str
+    else:
+      param_key = "params." + key_str
 
     shape = abstract_leaf_value.shape
     param_dict[param_key] = f"shape: {shape}"
@@ -294,6 +306,7 @@ def inspect_orbax(args):
   # Defer imports to avoid overhead when running in other modes.
   import orbax.checkpoint as ocp
   from etils import epath
+  from maxtext.checkpoint_conversion.utils.utils import param_key_parts_from_path
 
   path = epath.Path(args.path)
 
@@ -309,9 +322,13 @@ def inspect_orbax(args):
   # Filter strictly for parameter keys and format them.
   param_dict = {}
   for k, v in dictionary.items():
-    # `k` is a tuple representing the path hierarchy; join it into a single string.
-    # `v` is a metadata object containing `.shape` and `.dtype`.
-    param_key = ".".join(k)
+    # `k` is a tuple representing the path hierarchy.
+    # `v` is a shape-dtype struct.
+    # Pass it through param_key_parts_from_path to normalize NNX structures
+    # (folding list indices like "0" into "layers_0" and dropping "value").
+    key_parts = param_key_parts_from_path(k)
+    param_key = ".".join(key_parts)
+
     if not param_key.startswith("params"):
       continue
 
@@ -372,7 +389,8 @@ def main():
   if args.mode == "hf":
     inspect_hf(args)
   elif args.mode == "maxtext":
-    # remaining_args accepts maxtext config, like `model_name=<maxtext_model_name> scan_layers=<True | False>`
+    # remaining_args accepts maxtext config, like `model_name=<maxtext_model_name>
+    # scan_layers=<True | False> enable_nnx=<True | False>`
     inspect_maxtext(args, remaining_args)
   elif args.mode == "orbax":
     inspect_orbax(args)
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 2f3e26db7b..138d5be4a5 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
@@ -41,6 +41,7 @@
 import sys
 
 from flax import nnx
+from flax import serialization
 import jax
 from jax import random
 from jax.sharding import Mesh
@@ -118,6 +119,12 @@ def init_state_fn():
   )
 
   state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn)
+  if cfg.pure_nnx:
+    state = train_state_nnx.to_checkpoint_dict(state)
+    state.pop("nnx_aux", None)
+  else:
+    state = serialization.to_state_dict(state)
+
   max_logging.log("start")
   max_utils.print_mem_stats("After params initialized")
 
@@ -201,24 +208,10 @@ def init_state_fn():
       "['decoder']['decoder_norm']['bias']": (".params.lm.final_ln.bias", None),
   }
 
-  if cfg.pure_nnx:
-    # NNX state-tree paths after `nnx.split(TrainStateNNX)`. The state is a
-    # nested `nnx.State` (dict-like Mapping) with `nnx.Variable` leaves, so
-    # `jax.tree_util.keystr` produces dict-style entries (`['key']`) plus
-    # `.value` for the Variable leaf, plus `[idx]` for the optax tuple:
-    #   model params   -> ['model']<rest>.value
-    #   adam mu / nu   -> ['optimizer']['opt_state'][0]['mu' | 'nu']<rest>.value
-    #   step           -> ['optimizer']['step'].value
-    #   opt count      -> ['optimizer']['opt_state'][0]['count'].value
-    state_map = {
-        "['optimizer']['step'].value": ("step", None),
-        "['optimizer']['opt_state'][0]['count'].value": ("opt_states_0.no_prefix_0.count", None),
-    }
-  else:
-    state_map = {
-        ".step": ("step", None),
-        ".opt_state.count": ("opt_states_0.no_prefix_0.count", None),
-    }
+  state_map = {
+      "['step']": ("step", None),
+      "['opt_state']['count']": ("opt_states_0.no_prefix_0.count", None),
+  }
 
   def get_layer_prefix(keystr_pax):
     # different path format between decoder_layer variable
@@ -231,26 +224,18 @@ def get_layer_prefix(keystr_pax):
 
   for keystr_maxtext, (keystr_pax, transform_fn) in keystr_map.items():
     prefix_pax_opt_state = get_layer_prefix(keystr_pax)
-    if cfg.pure_nnx:
-      state_map[f"['model']{keystr_maxtext}.value"] = (f"mdl_vars{keystr_pax}", transform_fn)
-      state_map[f"['optimizer']['opt_state'][0]['mu']{keystr_maxtext}.value"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
-          transform_fn,
-      )
-      state_map[f"['optimizer']['opt_state'][0]['nu']{keystr_maxtext}.value"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
-          transform_fn,
-      )
-    else:
-      state_map[f".params['params']{keystr_maxtext}"] = (f"mdl_vars{keystr_pax}", transform_fn)
-      state_map[f".opt_state.mu['params']{keystr_maxtext}"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
-          transform_fn,
-      )
-      state_map[f".opt_state.nu['params']{keystr_maxtext}"] = (
-          f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
-          transform_fn,
-      )
+    state_map[f"['params']['params']{keystr_maxtext}"] = (
+        f"mdl_vars{keystr_pax}",
+        transform_fn,
+    )
+    state_map[f"['opt_state']['mu']['params']{keystr_maxtext}"] = (
+        f"opt_states_0.{prefix_pax_opt_state}.m{keystr_pax}",
+        transform_fn,
+    )
+    state_map[f"['opt_state']['nu']['params']{keystr_maxtext}"] = (
+        f"opt_states_0.{prefix_pax_opt_state}.v{keystr_pax}",
+        transform_fn,
+    )
 
   def verify_fn(key_path, _):
     keystr = jax.tree_util.keystr(key_path)
@@ -302,9 +287,9 @@ def map_fn(key_path, value):
   max_logging.log("converted state finished")
   max_utils.print_mem_stats("converted state finished")
 
-  step_value = int(converted_state.optimizer.step.value) if cfg.pure_nnx else converted_state.step
-  if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state):
-    max_logging.log(f"saved a checkpoint at step {step_value}")
+  step_value = int(converted_state["step"]) if isinstance(converted_state, dict) else int(converted_state.step)
+  if checkpointing.save_checkpoint(checkpoint_manager, step_value, converted_state, config=cfg):
+    max_logging.log(f"Saved a checkpoint at step {step_value}.")
   # Upon preemption, exit when and only when all ongoing saves are complete.
   if checkpointing.reached_preemption(checkpoint_manager, step_value):
     checkpointing.wait_until_finished(checkpoint_manager)
@@ -322,7 +307,12 @@ def map_fn(key_path, value):
       default="gs://mlperf-llm-public2/gpt3_spmd1x64x24_tpuv4-3072_v84_20221101/checkpoints/checkpoint_00004000",
       required=True,
   )
-  parser.add_argument("--maxtext-model-name", choices=["gpt3-175b", "gpt3-52k"], type=str, required=True)
+  parser.add_argument(
+      "--maxtext-model-name",
+      choices=["gpt3-175b", "gpt3-52k"],
+      type=str,
+      required=True,
+  )
   parser.add_argument("--base-output-directory", type=str, required=True)
   parser.add_argument("--run-name", type=str, required=True)
 
@@ -330,4 +320,9 @@ def map_fn(key_path, value):
   if not args.paxml_ckpt_path.startswith("gs://"):
     raise ValueError("--paxml-ckpt-path should be a gcs path starting with gs://")
 
-  convert(args.paxml_ckpt_path, args.maxtext_model_name, args.base_output_directory, args.run_name)
+  convert(
+      args.paxml_ckpt_path,
+      args.maxtext_model_name,
+      args.base_output_directory,
+      args.run_name,
+  )
diff --git a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py
index 098bd403f1..079c35c00d 100644
--- a/src/maxtext/checkpoint_conversion/utils/load_dynamic.py
+++ b/src/maxtext/checkpoint_conversion/utils/load_dynamic.py
@@ -73,12 +73,12 @@
 
 from flax import nnx
 import flax.traverse_util
-from google.cloud import storage
 import huggingface_hub
 import jax
 from maxtext.checkpoint_conversion.utils import hf_model_configs
 from maxtext.checkpoint_conversion.utils import param_mapping
 from maxtext.checkpoint_conversion.utils import tensor_handling
+from maxtext.common.gcloud_stub import gcs_storage
 from maxtext.utils import gcs_utils
 from maxtext.utils import globals as maxtext_globals
 from maxtext.utils import max_logging
@@ -86,6 +86,10 @@
 from orbax.checkpoint._src.arrays import sharding as sharding_utils
 
 
+# Route GCS through the decoupling helper so this module imports cleanly in
+# decoupled environments where google-cloud-storage is intentionally absent.
+storage = gcs_storage()
+
 HF_MODEL_CONFIGS = hf_model_configs.HF_MODEL_CONFIGS
 get_hf_loading_function = tensor_handling.get_hf_loading_function
 
diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py
index b83304f489..94ed14935e 100644
--- a/src/maxtext/checkpoint_conversion/utils/utils.py
+++ b/src/maxtext/checkpoint_conversion/utils/utils.py
@@ -867,12 +867,26 @@ def create_restore_args(tree_metadata):
   for i, path in enumerate(paths):
     checkpoint_path = epath.Path(path)
     metadata = ckptr.metadata(checkpoint_path)
+    checkpoint_tree = metadata.item_metadata.tree
+    if isinstance(checkpoint_tree, dict):
+      if "params" in checkpoint_tree:
+        checkpoint_tree = {"params": checkpoint_tree["params"]}
+        max_logging.log(f"Filtering checkpoint to only load 'params' from {path}")
+      else:
+        filtered_tree = {k: v for k, v in checkpoint_tree.items() if k not in ("opt_state", "optimizer")}
+        if len(filtered_tree) < len(checkpoint_tree):
+          checkpoint_tree = filtered_tree
+          max_logging.log(f"Filtering checkpoint to exclude optimizer keys from {path}")
+
     restore_args = jax.tree_util.tree_map(
         lambda x: create_restore_args(x) if hasattr(x, "shape") else None,
-        metadata.item_metadata.tree,
+        checkpoint_tree,
         is_leaf=lambda x: hasattr(x, "shape"),
     )
-    restored = ckptr.restore(checkpoint_path, restore_args=restore_args)
+    restored = ckptr.restore(
+        checkpoint_path,
+        args=ocp.args.PyTreeRestore(item=checkpoint_tree, restore_args=restore_args, partial_restore=True),
+    )
 
     if i == 0:
       merged_dict = restored
diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index 4827cbc1a4..d0d3bd7028 100644
--- a/src/maxtext/common/checkpointing.py
+++ b/src/maxtext/common/checkpointing.py
@@ -15,6 +15,7 @@
 
 """Create an Orbax CheckpointManager with specified (Async or not) Checkpointer."""
 
+import contextlib
 import datetime
 import importlib
 import time
@@ -716,7 +717,35 @@ def _handle_post_checkpoint_preemption(checkpoint_manager, step, force_ckpt_save
   if force_ckpt_save or is_preempted:
     wait_until_finished(checkpoint_manager)
   if is_preempted:
-    raise exceptions.StopTraining("Job is preempted.")
+    raise exceptions.StopTraining("Job received termination signal (SIGTERM).")
+
+
+@contextlib.contextmanager
+def checkpoint_exception_guard(config, checkpoint_manager, handler_fn=None):
+  """Context manager that wraps checkpointing save Exception handling.
+
+  On block success (checkpoint written without errors): runs the scale-up check
+  if elastic training is active.
+  On block failure: bubbles up JAX/ScaleUp errors if elastic training is active;
+  otherwise delegates to `handler_fn`.
+
+  Args:
+    config: maxtext configuration object.
+    checkpoint_manager: The CheckpointManager instance.
+    handler_fn: Optional callback function(Exception) that handles/wraps
+      non-elastic exceptions. If this handler raises a new exception, that new
+      exception is propagated. If it returns normally (returns None) or is not
+      provided, the original exception is re-raised (preserving its traceback).
+  """
+  try:
+    yield
+    elastic_utils.maybe_elastic_scale_up(config, checkpoint_manager)
+  except Exception as e:  # pylint: disable=broad-except
+    elastic_utils.maybe_bubble_elastic_exception(config, e)
+    if handler_fn:
+      handler_fn(e)
+    else:
+      raise
 
 
 def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step=None):
@@ -753,38 +782,21 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.")
     return
 
-  if config.pure_nnx:
-    # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
-    if config.enable_diloco:
-      # DiLoCoTrainState: persist the synchronized global model (outer params).
-      # The per-replica inner optimizer / outer-momentum state is not checkpointed.
-      step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step
-      state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}})
-    else:
-      # rngs/dropout/batch-stats are packed under items/nnx_aux so the RNG/dropout
-      # stream continues across resumes instead of resetting to a base key.
-      state = train_state_nnx.to_checkpoint_dict(state)
-
-  try:
-    checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save)
-    if checkpoint_saved:
-      print_save_message(actual_step, config.async_checkpointing)
-    if config.elastic_enabled:
-      elastic_utils.maybe_elastic_scale_up(config, checkpoint_manager)
-  except elastic_utils.manager.ScaleUpSignalError as e:
-    if config.elastic_enabled:
-      max_logging.log(f"Elastic event detected, letting exception bubble up: {e}")
-      raise
-    else:
-      raise exceptions.StopTraining("Job is preempted.") from e
-  except jax.errors.JaxRuntimeError as e:
-    if config.elastic_enabled:
-      max_logging.log(f"Elastic event detected, letting exception bubble up: {e}")
-      raise
-    else:
-      raise exceptions.StopTraining("Job is preempted.") from e
-  except Exception as e:
-    raise exceptions.StopTraining(f"Checkpointing failed. {str(e)}") from e
+  def _checkpoint_error_handler(err):
+    """Handles checkpointing errors, when not in an elastic context."""
+    raise exceptions.StopTraining(f"Checkpointing failed. {str(err)}") from err
+
+  with checkpoint_exception_guard(config, checkpoint_manager, _checkpoint_error_handler):
+    checkpoint_saved = save_checkpoint(
+        checkpoint_manager,
+        actual_step,
+        state,
+        config,
+        data_iterator,
+        force_ckpt_save,
+    )
+  if checkpoint_saved:
+    print_save_message(actual_step, config.async_checkpointing)
 
   # Wait for any pending checkpoint save to finish during preemption or final
   # step save, then raise upon preemption.
@@ -793,12 +805,26 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
 
 def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False):
   """Wrapper for saving checkpoint."""
-  if config and config.enable_checkpointing:
+  if not isinstance(state, (dict, nnx.State, train_state.TrainState)):
+    if isinstance(state, train_state_nnx.TrainStateNNX):
+      state = nnx.state(state)
+    elif not isinstance(state, (dict, nnx.State)):
+      state = {}
+
+  if config and getattr(config, "pure_nnx", False) and isinstance(state, nnx.State):
+    # Save in the Linen on-disk layout so pure_nnx and Linen checkpoints are interchangeable.
+    if getattr(config, "enable_diloco", False):
+      step_value = state.step.get_value() if hasattr(state.step, "get_value") else state.step
+      state = train_state_nnx.to_linen_checkpoint_dict({"model": state.params, "optimizer": {"step": step_value}})
+    else:
+      state = train_state_nnx.to_checkpoint_dict(state)
+
+  if config and getattr(config, "enable_checkpointing", False):
     if (
         force
-        or (step % config.checkpoint_period == 0 and not config.enable_continuous_checkpointing)
+        or (step % config.checkpoint_period == 0 and not getattr(config, "enable_continuous_checkpointing", False))
         or (_uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0)
-        or (config.enable_autocheckpoint and reached_preemption(checkpoint_manager, step))
+        or (getattr(config, "enable_autocheckpoint", False) and reached_preemption(checkpoint_manager, step))
     ):
       blocking_until_ready_start = time.time()
       max_logging.log(f"Waiting for step {step} to finish before checkpoint...")
@@ -812,7 +838,13 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
 
   # specify chunk_byte_size to force orbax to control maximum file size in checkpoint
   chunk_byte_size = (
-      config.checkpoint_storage_target_data_file_size_bytes if config else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
+      getattr(
+          config,
+          "checkpoint_storage_target_data_file_size_bytes",
+          DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE,
+      )
+      if config
+      else DEFAULT_OCDBT_TARGET_DATA_FILE_SIZE
   )
 
   checkpoint_args = ocp.args.PyTreeSave(
@@ -822,7 +854,11 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=
   )
   save_args_composite = {"items": checkpoint_args}
 
-  if config and config.dataset_type == "grain" and not isinstance(data_iterator, PlaceHolderDataIterator):
+  if (
+      config
+      and getattr(config, "dataset_type", None) == "grain"
+      and not isinstance(data_iterator, PlaceHolderDataIterator)
+  ):
     if isinstance(data_iterator, RemoteIteratorWrapper):
       # Pass the wrapper directly; GrainCheckpointHandler will call save_state with the step
       save_args_composite["iter"] = grain_utility.GrainCheckpointSave(
diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py
index 21bd870bc8..2bd04c527b 100644
--- a/src/maxtext/common/data_loader.py
+++ b/src/maxtext/common/data_loader.py
@@ -15,19 +15,40 @@
 # pytype: disable=unsupported-operands
 """Module to load data for training."""
 
+import contextlib
 import jax
-import jax.numpy as jnp
 from jax.experimental import checkify
-
+import jax.numpy as jnp
 from maxtext.common.goodput import (
     GoodputEvent,
     maybe_record_goodput,
 )
 from maxtext.trainers.diloco import diloco
+from maxtext.utils import elastic_utils
 from maxtext.utils import exceptions
 from maxtext.utils.sharding import get_input_data_sharding
 
 
+@contextlib.contextmanager
+def loader_exception_guard(config):
+  """Context manager that wraps data loading Exception handling.
+
+  On block failure: bubbles up JAX/ScaleUp errors if elastic training is active;
+  otherwise raises a StopTraining exception.
+
+  Args:
+    config: maxtext configuration object.
+  """
+  try:
+    yield
+  except Exception as e:  # pylint: disable=broad-except
+    elastic_utils.maybe_bubble_elastic_exception(config, e)
+    if isinstance(e, StopIteration):
+      raise exceptions.StopTraining(f"You may have run out of training data. Received {type(e)}" f" exception: ({e})")
+    else:
+      raise exceptions.StopTraining(f"`next(self.data_iterator)` failed with {type(e)} exception: ({e}).")
+
+
 class DataLoader:
   """
   Loads preprocessed data for training.
@@ -54,23 +75,18 @@ def update_data_iterator(self):
   def load_next_batch_pre_sharding(self):
     """Loads the next batch w/o sharding. Can keep reusing the same batch for performance reasons."""
     with maybe_record_goodput(self.goodput_recorder, GoodputEvent.DATA_LOADING):
-      try:
-        if self.config.reuse_example_batch and self.last_batch:
-          example_batch = self.last_batch
-        else:
+      if self.config.reuse_example_batch and self.last_batch:
+        example_batch = self.last_batch
+      else:
+        with loader_exception_guard(self.config):
           example_batch = next(self.data_iterator)
-          self.update_data_iterator()
-        self.last_batch = example_batch
-        self.check_example_batch()
-      except Exception as e:  # pylint: disable=broad-except
-        if isinstance(e, StopIteration):
-          raise exceptions.StopTraining(f"You may have run out of training data. Received {type(e)} exception: ({e})")
-        else:
-          raise exceptions.StopTraining(f"`load_next_batch()` failed with {type(e)} exception: ({e}).")
+        self.update_data_iterator()
+      self.last_batch = example_batch
+      self.check_example_batch()
     return self.last_batch
 
   def load_next_batch(self, *args, **kwargs):
-    """Loads the next batch with sharding hint"""
+    """Loads the next batch with sharding hint."""
     example_batch = self.load_next_batch_pre_sharding()
     if self.config.enable_diloco:
       example_batch = diloco.reshape_first_axis_with_diloco(self.config.num_diloco_replicas, example_batch)
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index b0843a9939..73f6ddc4a6 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -988,7 +988,7 @@ decode_sampling_nucleus_p: -1 # set if you're doing nucleus / top-p
 decode_sampling_top_k: 0 # set if you're doing top-k
 decode_sampling_temperature: 1.
 
-eval_start_step: 0 # start eval after train step is >= eval_start_step
+eval_start_step: 0 # start eval when train step is >= eval_start_step
 eval_interval: -1  # the specific number of train step between eval_step
 eval_steps: -1  # run this number of steps for eval, recommend setting this to prevent error due to running out of evel data
 target_eval_loss: 0.  # early stop once reaching target eval_loss
@@ -1339,6 +1339,7 @@ distill_beta_schedule: "constant"
 ##### Elastic training parameters
 # Elastic training is Pathways-specific and does not work on McJAX.
 elastic_enabled: false
+elastic_backup_kind: "snapshot"
 elastic_timeout_seconds: 300
 elastic_max_retries: 10
 elastic_min_slice_count: -1
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index edc7d3112e..ff4172bac3 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -864,6 +864,10 @@ class MoEGeneral(BaseModel):
       False,
       description="Whether to use ragged kernel for sorting, improve performance when EP is enabled.",
   )
+  ragged_sort_use_single_sparsecore: bool = Field(
+      False,
+      description="Whether to run ragged sort kernels on 1 SparseCore instead of all SparseCores.",
+  )
   use_gather_mosaic_kernel: bool = Field(
       False,
       description="Whether to use a custom mosaic kernel for token gather ops.",
@@ -1600,7 +1604,8 @@ class TrainingLoop(BaseModel):
   log_period: int = Field(100, description="Frequency (in steps) to log metrics and flush Tensorboard.")
   eval_start_step: int = Field(
       0,
-      description="Start evaluation after training step is >= eval_start_step.",
+      ge=0,
+      description="Start evaluation when training step is >= eval_start_step.",
   )
   eval_interval: int = Field(
       -1,
@@ -2035,6 +2040,10 @@ class ElasticTraining(BaseModel):
   """
 
   elastic_enabled: bool = Field(False, description="Whether to enable elastic training.")
+  elastic_backup_kind: str = Field(
+      "snapshot",
+      description=("The kind of backup to use for elastic training: 'snapshot' or 'checkpoint'."),
+  )
   elastic_timeout_seconds: int = Field(
       300,
       description=(
@@ -3227,6 +3236,10 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       )
     if self.elastic_enabled and not self.enable_single_controller:
       raise ValueError("Elastic training is only supported with Pathways (`enable_single_controller=True`).")
+    if self.elastic_backup_kind not in ("snapshot", "checkpoint"):
+      raise ValueError(
+          "elastic_backup_kind must be one of 'snapshot' or 'checkpoint', got" f" '{self.elastic_backup_kind}'."
+      )
     if self.colocated_python_data_input and not self.enable_single_controller:
       raise ValueError(
           "Colocated python data input is only supported with Pathways (single"
diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py
index b13d72d89e..5de06a828b 100644
--- a/src/maxtext/inference/maxengine/maxengine.py
+++ b/src/maxtext/inference/maxengine/maxengine.py
@@ -145,11 +145,16 @@ def __init__(self, config: Any, devices: Any | None = None):
       # nnx.merge. `rest` (RNG state etc.) is materialized in load_params.
       graphdef, _, _, _ = nnx.split(abstract_model, nnx.Param, nnx.Cache, ...)
       self.graphdef = graphdef
+      # Layers may bake their construction-time model_mode into static attributes,
+      # so a call must be merged with the graphdef built for that same mode.
+      graphdef_ar, _, _, _ = nnx.split(abstract_model_ar, nnx.Param, nnx.Cache, ...)
+      self.graphdef_ar = graphdef_ar
       self._create_model_fn = _create_model
       self._nnx_rest_state = None
     else:
       self.model = models.transformer_as_linen(config, mesh=self._mesh, quant=quant, model_mode=MODEL_MODE_PREFILL)
       self.graphdef = None
+      self.graphdef_ar = None
       self._create_model_fn = None
     self.replicated_sharding = jax.sharding.NamedSharding(self._mesh, P(None))
 
@@ -218,10 +223,14 @@ def _nnx_run_model(
     """NNX equivalent of `model.apply(..., mutable=["cache"])`. Returns (logits, new_cache_dict)."""
     cache_state = self._nnx_cache_state_template(mode=model_mode)
     nnx.replace_by_pure_dict(cache_state, cache_dict)
+    # Merge with the graphdef built for this mode. Layers that captured their
+    # model_mode at construction (e.g. the DeepSeek layers) would otherwise run
+    # the prefill attention path during autoregressive decode.
+    graphdef = self.graphdef_ar if model_mode == MODEL_MODE_AUTOREGRESSIVE else self.graphdef
     # copy=True avoids reusing Variable objects across traces (TraceContextError),
     # mirroring the workaround in train.py's diff_wrapper.
     model = nnx.merge(
-        self.graphdef, params, cache_state, self._nnx_rest_state, copy=True
+        graphdef, params, cache_state, self._nnx_rest_state, copy=True
     )  # pyrefly: ignore[no-matching-overload]
     logits = model(
         decoder_input_tokens,
diff --git a/src/maxtext/kernels/ragged/ragged_gather.py b/src/maxtext/kernels/ragged/ragged_gather.py
index e1d31b24c4..0c2aabb1c1 100644
--- a/src/maxtext/kernels/ragged/ragged_gather.py
+++ b/src/maxtext/kernels/ragged/ragged_gather.py
@@ -275,6 +275,8 @@ def get_cost_estimate(
       auto-computing.  -1 (default) means auto-compute.
     bytes_accessed_override: If > 0, use this value as bytes_accessed instead
       of auto-computing.  -1 (default) means auto-compute.
+    use_single_sparsecore: Static bool flag. When True, launch the kernel
+      on 1 SparseCore instead of all SparseCores.
 
   Returns:
     A ``pl.CostEstimate`` suitable for XLA scheduling.
@@ -343,7 +345,14 @@ def calculate_col_size(hidden_size: int) -> int:
 
 
 @functools.partial(
-    jax.jit, static_argnames=("has_weights", "enforce_fallback", "flops_override", "bytes_accessed_override")
+    jax.jit,
+    static_argnames=(
+        "has_weights",
+        "enforce_fallback",
+        "flops_override",
+        "bytes_accessed_override",
+        "use_single_sparsecore",
+    ),
 )
 def ragged_gather(
     x: jax.Array,
@@ -355,6 +364,7 @@ def ragged_gather(
     enforce_fallback: bool = False,
     flops_override: int = -1,
     bytes_accessed_override: int = -1,
+    use_single_sparsecore: bool = False,
 ) -> jax.Array:
   """Perform gather on indices within dynamic array start and end.
 
@@ -375,6 +385,8 @@ def ragged_gather(
       auto-computing.  -1 (default) means auto-compute.
     bytes_accessed_override: If > 0, use this value as bytes_accessed instead
       of auto-computing.  -1 (default) means auto-compute.
+    use_single_sparsecore: Static bool flag. When True, launch the kernel
+      on 1 SparseCore instead of all SparseCores.
 
   Returns:
     Gathered output of shape ``(indices_size, hidden_size)``.
@@ -410,7 +422,8 @@ def ragged_gather(
   out_size = indices.size
 
   num_simd_lanes = sc_info.num_lanes
-  num_cores = sc_info.num_cores * sc_info.num_subcores
+  num_sc_cores = 1 if use_single_sparsecore else sc_info.num_cores
+  num_cores = num_sc_cores * sc_info.num_subcores
   block_size = num_simd_lanes * num_cores
   col_size = calculate_col_size(hidden_size)
 
@@ -427,7 +440,7 @@ def ragged_gather(
   aligned_hidden_size = pl.cdiv(hidden_size, col_size) * col_size
 
   vector_mesh = plsc.VectorSubcoreMesh(
-      num_cores=sc_info.num_cores,
+      num_cores=num_sc_cores,
       num_subcores=sc_info.num_subcores,
       core_axis_name="core",
       subcore_axis_name="subcore",
diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py
index 3ef49a2896..e2116a4630 100644
--- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py
+++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py
@@ -410,7 +410,14 @@ def _preprocess(
 
 
 @functools.partial(
-    jax.jit, static_argnames=("reduce_group_size", "enforce_fallback", "flops_override", "bytes_accessed_override")
+    jax.jit,
+    static_argnames=(
+        "reduce_group_size",
+        "enforce_fallback",
+        "flops_override",
+        "bytes_accessed_override",
+        "use_single_sparsecore",
+    ),
 )
 def ragged_gather_reduce(
     x: jax.Array,
@@ -421,6 +428,7 @@ def ragged_gather_reduce(
     enforce_fallback: bool = False,
     flops_override: int = -1,
     bytes_accessed_override: int = -1,
+    use_single_sparsecore: bool = False,
 ) -> jax.Array:
   """Gathers `x` according to `indices`, applies weights and masks, and reduces.
 
@@ -474,7 +482,8 @@ def ragged_gather_reduce(
   hidden_size = x.shape[-1]
   input_size = indices.size
   num_simd_lanes = sc_info.num_lanes
-  num_cores = sc_info.num_cores * sc_info.num_subcores
+  num_sc_cores = 1 if use_single_sparsecore else sc_info.num_cores
+  num_cores = num_sc_cores * sc_info.num_subcores
 
   # This kernel partitions the output's columns into `num_column_partitions`
   # and partition the output's rows into `num_row_partitions` and run each
@@ -528,7 +537,7 @@ def ragged_gather_reduce(
   )
 
   vector_mesh = plsc.VectorSubcoreMesh(
-      num_cores=sc_info.num_cores,
+      num_cores=num_sc_cores,
       num_subcores=sc_info.num_subcores,
       core_axis_name="core",
       subcore_axis_name="subcore",
diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
index b6ccc1e281..48f3ac413c 100644
--- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
+++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
@@ -625,7 +625,14 @@ def col_loop(col_compute_offset):
 
 
 @functools.partial(
-    jax.jit, static_argnames=("reduce_group_size", "enforce_fallback", "flops_override", "bytes_accessed_override")
+    jax.jit,
+    static_argnames=(
+        "reduce_group_size",
+        "enforce_fallback",
+        "flops_override",
+        "bytes_accessed_override",
+        "use_single_sparsecore",
+    ),
 )
 def ragged_gather_reduce(
     x: jax.Array,
@@ -636,6 +643,7 @@ def ragged_gather_reduce(
     enforce_fallback: bool = False,
     flops_override: int = -1,
     bytes_accessed_override: int = -1,
+    use_single_sparsecore: bool = False,
 ) -> jax.Array:
   """Gathers ``x`` by ``indices``, weights and masks, then reduces by group.
 
@@ -673,7 +681,8 @@ def ragged_gather_reduce(
   input_size = indices.size
   num_simd_lanes = sc_info.num_lanes
   num_lanes = pltpu.get_tpu_info().num_lanes
-  num_cores = sc_info.num_cores * sc_info.num_subcores
+  num_sc_cores = 1 if use_single_sparsecore else sc_info.num_cores
+  num_cores = num_sc_cores * sc_info.num_subcores
 
   num_column_partitions = _calculate_num_column_partitions(hidden_size, input_size, num_cores, num_lanes, num_simd_lanes)
   num_row_partitions = num_cores // num_column_partitions
@@ -711,7 +720,7 @@ def ragged_gather_reduce(
 
   # Step 4: Launch the SparseCore kernel.
   vector_mesh = plsc.VectorSubcoreMesh(
-      num_cores=sc_info.num_cores,
+      num_cores=num_sc_cores,
       num_subcores=sc_info.num_subcores,
       core_axis_name="core",
       subcore_axis_name="subcore",
diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py
index 639d3172b7..f3def646b9 100644
--- a/src/maxtext/kernels/ragged/ragged_sort.py
+++ b/src/maxtext/kernels/ragged/ragged_sort.py
@@ -34,6 +34,7 @@ def ring_ragged_sort(
     gather_reduce_flops_override=-1,
     gather_bytes_accessed_override=-1,
     gather_reduce_bytes_accessed_override=-1,
+    use_single_sparsecore=False,
 ):
   """Ragged-gather variant for AG-RS Expert Parallelism token routing.
 
@@ -111,6 +112,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local):
           enforce_fallback=enforce_gather_fallback,
           flops_override=gather_flops_override,
           bytes_accessed_override=gather_bytes_accessed_override,
+          use_single_sparsecore=use_single_sparsecore,
       )
     else:
       local_buffer_size = buffer_size
@@ -134,6 +136,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local):
           enforce_fallback=enforce_gather_fallback,
           flops_override=gather_flops_override,
           bytes_accessed_override=gather_bytes_accessed_override,
+          use_single_sparsecore=use_single_sparsecore,
       )
 
     out = (x, group_sizes_local, topk_argsort_revert_indices)
@@ -188,6 +191,7 @@ def _ring_ragged_sort_bwd(res, g_out):
           enforce_fallback=enforce_gather_reduce_fallback,
           flops_override=gather_reduce_flops_override,
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
+          use_single_sparsecore=use_single_sparsecore,
       )
     else:
       # Buffering: g_x has size `local_buffer_size` (packed).
@@ -213,6 +217,7 @@ def _ring_ragged_sort_bwd(res, g_out):
           enforce_fallback=enforce_gather_reduce_fallback,
           flops_override=gather_reduce_flops_override,
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
+          use_single_sparsecore=use_single_sparsecore,
       )
     return grad_hidden_states, None
 
@@ -235,6 +240,7 @@ def ring_ragged_unsort(
     gather_reduce_flops_override=-1,
     gather_bytes_accessed_override=-1,
     gather_reduce_bytes_accessed_override=-1,
+    use_single_sparsecore=False,
 ):
   """Dual of :func:`ring_ragged_sort`.
 
@@ -266,14 +272,27 @@ def ring_ragged_unsort(
   """
 
   @jax.custom_vjp
-  def _ring_ragged_unsort(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat):
+  def _ring_ragged_unsort(
+      sorted_tokens_local,
+      group_sizes_local,
+      topk_argsort_revert_indices,
+      topk_weights_flat,
+  ):
     """Unsort and scatter activations."""
     return _ring_ragged_unsort_fwd(
-        sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat
+        sorted_tokens_local,
+        group_sizes_local,
+        topk_argsort_revert_indices,
+        topk_weights_flat,
     )[0]
 
   @jax.named_scope("ragged-unsort-fwd")
-  def _ring_ragged_unsort_fwd(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat):
+  def _ring_ragged_unsort_fwd(
+      sorted_tokens_local,
+      group_sizes_local,
+      topk_argsort_revert_indices,
+      topk_weights_flat,
+  ):
     """Executes unsorting sending tokens back."""
     group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True)
 
@@ -309,6 +328,7 @@ def _ring_ragged_unsort_fwd(sorted_tokens_local, group_sizes_local, topk_argsort
           enforce_fallback=enforce_gather_reduce_fallback,
           flops_override=gather_reduce_flops_override,
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
+          use_single_sparsecore=use_single_sparsecore,
       )
     else:
       # Shift indices so they map to the packed local buffer [0, local_num_tokens).
@@ -327,6 +347,7 @@ def _ring_ragged_unsort_fwd(sorted_tokens_local, group_sizes_local, topk_argsort
           enforce_fallback=enforce_gather_reduce_fallback,
           flops_override=gather_reduce_flops_override,
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
+          use_single_sparsecore=use_single_sparsecore,
       )
 
     res = (
@@ -385,6 +406,7 @@ def _ring_ragged_unsort_bwd(res, g_out):
           enforce_fallback=enforce_gather_fallback,
           flops_override=gather_flops_override,
           bytes_accessed_override=gather_bytes_accessed_override,
+          use_single_sparsecore=use_single_sparsecore,
       )
       # Mask out gradients that correspond to elements outside the valid shard
       # output range.
@@ -408,6 +430,7 @@ def _ring_ragged_unsort_bwd(res, g_out):
           enforce_fallback=enforce_gather_fallback,
           flops_override=gather_flops_override,
           bytes_accessed_override=gather_bytes_accessed_override,
+          use_single_sparsecore=use_single_sparsecore,
       )
       # Mask out gradients for elements beyond the valid limit of the local buffer.
       limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size)
@@ -420,10 +443,22 @@ def _ring_ragged_unsort_bwd(res, g_out):
   # Build the flat weights array from the routing weights.
   topk_weights_flat = topk_weights.astype(jnp.float32)
 
-  return _ring_ragged_unsort(sorted_tokens_local, group_sizes_local, topk_argsort_revert_indices, topk_weights_flat)
+  return _ring_ragged_unsort(
+      sorted_tokens_local,
+      group_sizes_local,
+      topk_argsort_revert_indices,
+      topk_weights_flat,
+  )
 
 
-def a2a_ragged_sort(inputs, sort_indices, valid_end, enforce_gather_fallback=False, enforce_gather_reduce_fallback=False):
+def a2a_ragged_sort(
+    inputs,
+    sort_indices,
+    valid_end,
+    enforce_gather_fallback=False,
+    enforce_gather_reduce_fallback=False,
+    use_single_sparsecore=False,
+):
   """Ragged-gather variant for ``local_permute``.
 
   Unlike :func:`ring_ragged_sort`, the rows valid for this shard live in
@@ -463,7 +498,13 @@ def _a2a_ragged_sort(inputs, sort_indices, valid_end):
   def _a2a_ragged_sort_fwd(inputs, sort_indices, valid_end):
     start = jnp.int32(0)
     end = valid_end.astype(jnp.int32) if hasattr(valid_end, "astype") else jnp.int32(valid_end)
-    out = ragged_gather(inputs, sort_indices, start[None], end[None])
+    out = ragged_gather(
+        inputs,
+        sort_indices,
+        start[None],
+        end[None],
+        use_single_sparsecore=use_single_sparsecore,
+    )
     n = sort_indices.shape[0]
     valid_mask = jnp.arange(n) < end
     out = jnp.where(valid_mask[:, None], out, 0.0)
@@ -487,6 +528,7 @@ def _a2a_ragged_sort_bwd(res, g_out):
         valid_rows_mask=valid_rows_mask[idx_inv],
         reduce_group_size=1,
         enforce_fallback=enforce_gather_reduce_fallback,
+        use_single_sparsecore=use_single_sparsecore,
     )
     # custom_vjp must return one gradient per primal arg; valid_end is integer
     # and non-differentiable, so we return None for it.
@@ -497,7 +539,12 @@ def _a2a_ragged_sort_bwd(res, g_out):
 
 
 def a2a_ragged_unsort(
-    sorted_tokens, revert_indices, valid_end, enforce_gather_fallback=False, enforce_gather_reduce_fallback=False
+    sorted_tokens,
+    revert_indices,
+    valid_end,
+    enforce_gather_fallback=False,
+    enforce_gather_reduce_fallback=False,
+    use_single_sparsecore=False,
 ):
   """Dual of :func:`a2a_ragged_sort`.
 
@@ -540,6 +587,7 @@ def _a2a_ragged_unsort_fwd(sorted_tokens, revert_indices, valid_end):
         valid_rows_mask=valid_rows_mask,
         reduce_group_size=1,
         enforce_fallback=enforce_gather_reduce_fallback,
+        use_single_sparsecore=use_single_sparsecore,
     )
     res = (revert_indices, end, sorted_tokens.shape, start)
     return out, res
@@ -551,7 +599,13 @@ def _a2a_ragged_unsort_bwd(res, g_out):
     # Because revert_indices is a permutation, build the inverse and use
     # ragged_gather to pull the per-row gradients to the right positions.
     idx_inv = jnp.argsort(revert_indices)
-    grad_sorted = ragged_gather(g_out, idx_inv, start[None], end[None])
+    grad_sorted = ragged_gather(
+        g_out,
+        idx_inv,
+        start[None],
+        end[None],
+        use_single_sparsecore=use_single_sparsecore,
+    )
     num_rows = sorted_tokens_shape[0]
     pos = jnp.arange(num_rows)
     valid = pos < end
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index c7c632a29c..915193eb0f 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -917,6 +917,7 @@ def permute(
           gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops,
           gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed,
           gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed,
+          use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
       )
     else:
       flatten_selected_experts = jnp.ravel(selected_experts)
@@ -1005,6 +1006,7 @@ def unpermute(
           gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops,
           gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed,
           gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed,
+          use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
       )
     else:
       unsort_intermediate = _sort_activations(
@@ -1069,6 +1071,7 @@ def local_permute(
       use_custom_sort_vjp=True,
       use_ragged_sort=False,
       ragged_buffer_factor=-1.0,
+      use_single_sparsecore=False,
   ):
     """Permutes tokens locally within an expert shard.
 
@@ -1152,7 +1155,12 @@ def local_permute(
       # the worst-case ragged buffer. Restricting the gather to that prefix
       # makes both forward and backward proportional to the routed token count.
       valid_end = jnp.sum(local_group_size).astype(jnp.int32)
-      sorted_inputs = a2a_ragged_sort(inputs, sorted_indices, valid_end)
+      sorted_inputs = a2a_ragged_sort(
+          inputs,
+          sorted_indices,
+          valid_end,
+          use_single_sparsecore=use_single_sparsecore,
+      )
     else:
       sorted_inputs = _sort_activations(inputs, sorted_indices, use_custom_sort_vjp)
     sorted_experts_ids = expert_indices[sorted_indices]
@@ -1762,6 +1770,7 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in
               use_custom_sort_vjp=self.config.use_custom_sort_vjp,
               use_ragged_sort=self.config.use_ragged_sort,
               ragged_buffer_factor=self.config.ragged_buffer_factor,
+              use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           )
         else:
           x, local_sorted_indices, group_sizes, selected_experts = RoutedMoE.local_permute(
@@ -1774,6 +1783,7 @@ def ra2a_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, in
               use_custom_sort_vjp=self.config.use_custom_sort_vjp,
               use_ragged_sort=self.config.use_ragged_sort,
               ragged_buffer_factor=self.config.ragged_buffer_factor,
+              use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           )
 
       return (
@@ -1959,6 +1969,7 @@ def unsort_output_and_ra2a(
               intermediate_output,
               jnp.argsort(route_metadata.local_sorted_indices),  # pylint: disable=undefined-variable
               valid_end,
+              use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           )
         else:
           local_output = _sort_activations(
diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py
index 53166f00ee..e0e09f4175 100644
--- a/src/maxtext/models/deepseek.py
+++ b/src/maxtext/models/deepseek.py
@@ -211,6 +211,7 @@ def attention_op(
       decoder_segment_ids,
       decoder_positions,
       deterministic,
+      model_mode,
       previous_chunk=None,
       slot: None | int = None,
   ):
@@ -221,7 +222,7 @@ def attention_op(
         decoder_positions,
         decoder_segment_ids=decoder_segment_ids,
         deterministic=deterministic,
-        model_mode=self.model_mode,
+        model_mode=model_mode,
         out_sharding=self.out_sharding,
         previous_chunk=previous_chunk,
         slot=slot,
@@ -270,6 +271,7 @@ def self_attention_with_norm_op(
       decoder_segment_ids,
       decoder_positions,
       deterministic,
+      model_mode,
       previous_chunk=None,
       slot: None | int = None,
   ):
@@ -283,7 +285,7 @@ def self_attention_with_norm_op(
           decoder_segment_ids=decoder_segment_ids,
           inputs_positions=decoder_positions,
           deterministic=deterministic,
-          model_mode=self.model_mode,
+          model_mode=model_mode,
           out_sharding=self.out_sharding,
           previous_chunk=previous_chunk,
           slot=slot,
@@ -295,6 +297,7 @@ def self_attention_with_norm_op(
           decoder_segment_ids,
           decoder_positions,
           deterministic,
+          model_mode,
           previous_chunk,
           slot,
       )
@@ -368,6 +371,7 @@ def __call__(
         decoder_segment_ids,
         decoder_positions,
         deterministic,
+        model_mode,
         previous_chunk,
         slot,
     )
@@ -581,6 +585,7 @@ def extract_fn(x):
         decoder_segment_ids,
         decoder_positions,
         deterministic,
+        model_mode,
         previous_chunk,
         slot,
     )
diff --git a/src/maxtext/models/deepseek4.py b/src/maxtext/models/deepseek4.py
index c259db8c66..6181d684e0 100644
--- a/src/maxtext/models/deepseek4.py
+++ b/src/maxtext/models/deepseek4.py
@@ -151,6 +151,7 @@ def __call__(
         decoder_segment_ids,
         decoder_positions,
         deterministic,
+        model_mode,
         previous_chunk,
         slot,
     )
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index e28316ef15..367695211f 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -718,7 +718,12 @@ def training_loop_iteration(
         all_host_upload=dump_hlo_upload_all,
     )
 
-  if eval_interval > 0 and step >= start_step and step >= eval_start_step and (step + 1) % eval_interval == 0:
+  if (
+      eval_interval > 0
+      and step >= start_step
+      and step >= eval_start_step
+      and (step - eval_start_step) % eval_interval == 0
+  ):
     assert eval_data_iterator
     # Explicitly reset the eval iterator and counters before starting the eval loop
     eval_data_iterator.reset()
diff --git a/src/maxtext/utils/elastic_utils.py b/src/maxtext/utils/elastic_utils.py
index 1cba750774..4ebfb97571 100644
--- a/src/maxtext/utils/elastic_utils.py
+++ b/src/maxtext/utils/elastic_utils.py
@@ -14,8 +14,8 @@
 
 """Utility functions for Elastic Training."""
 
-import functools
 from collections import Counter
+import functools
 from types import SimpleNamespace
 
 import jax
@@ -64,6 +64,22 @@ def elastic_enabled(config) -> bool:
   return pathwaysutils.is_pathways_backend_used() and config.elastic_enabled
 
 
+def elastic_snapshot(config) -> bool:
+  """Returns whether elastic snapshot mode is enabled."""
+  return elastic_enabled(config) and config.elastic_backup_kind == "snapshot"
+
+
+def maybe_bubble_elastic_exception(config, e: Exception) -> None:
+  """Checks JAX/ScaleUp elastic errors and re-raises them if elasticity is enabled.
+
+  Args:
+    config: Maxtext configuration object.
+    e: The exception currently being evaluated.
+  """
+  if elastic_enabled(config) and isinstance(e, (jax.errors.JaxRuntimeError, manager.ScaleUpSignalError)):
+    raise e
+
+
 def should_use_elastic(config) -> bool:
   """Returns whether elastic training should be used."""
   return config is not None and elastic_enabled(config)
@@ -218,6 +234,9 @@ def is_scale_up_event(config) -> bool:
 
 def maybe_elastic_scale_up(config, checkpoint_manager):
   """Waits for a checkpoint to finish before interrupting for scale up."""
+  if not should_use_elastic(config):
+    max_logging.log("maybe_elastic_scale_up: Elastic training is not enabled.")
+    return
   if is_scale_up_event(config):
     max_logging.log(
         "Started a checkpoint and a new slice is available. Waiting for current"
diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py
index 421a6d150f..b4b67a6eaf 100644
--- a/tests/unit/checkpointing_test.py
+++ b/tests/unit/checkpointing_test.py
@@ -149,7 +149,7 @@ class LoadDynamicTest(parameterized.TestCase):
   """Tests for cache downloads and dynamic loading of safetensors."""
 
   @mock.patch("huggingface_hub.HfFileSystem")
-  @mock.patch("google.cloud.storage.Client")
+  @mock.patch.object(load_dynamic.storage, "Client")
   def test_build_gcs_cache_worker_cache_hit(self, mock_storage_client, mock_hf_fs):
     mock_client_instance = mock_storage_client.return_value
     mock_bucket = mock_client_instance.bucket.return_value
@@ -161,7 +161,7 @@ def test_build_gcs_cache_worker_cache_hit(self, mock_storage_client, mock_hf_fs)
     mock_blob.upload_from_file.assert_not_called()
 
   @mock.patch("huggingface_hub.HfFileSystem")
-  @mock.patch("google.cloud.storage.Client")
+  @mock.patch.object(load_dynamic.storage, "Client")
   def test_build_gcs_cache_worker_cache_miss_success(self, mock_storage_client, mock_hf_fs):
     mock_fs_instance = mock_hf_fs.return_value
     mock_remote_file = mock.MagicMock()
@@ -177,7 +177,7 @@ def test_build_gcs_cache_worker_cache_miss_success(self, mock_storage_client, mo
     mock_blob.upload_from_file.assert_called_once_with(mock_remote_file, client=mock_client_instance)
 
   @mock.patch("huggingface_hub.HfFileSystem")
-  @mock.patch("google.cloud.storage.Client")
+  @mock.patch.object(load_dynamic.storage, "Client")
   def test_build_gcs_cache_worker_retry_and_fail(self, mock_storage_client, mock_hf_fs):
     mock_fs_instance = mock_hf_fs.return_value
     mock_fs_instance.open.side_effect = Exception("Download failed")
diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py
index 465adbc307..78dfab1da9 100644
--- a/tests/unit/configs_value_test.py
+++ b/tests/unit/configs_value_test.py
@@ -312,6 +312,17 @@ def test_safetensors_dynamic_disallows_single_controller(self):
     with self.assertRaises(pydantic.ValidationError):
       pyconfig.initialize(argv)
 
+  def test_elastic_backup_kind_validation(self):
+    """Tests that elastic_backup_kind must be either 'snapshot' or 'checkpoint'."""
+    argv = [
+        "",
+        _BASE_CONFIG_PATH,
+        "run_name=test",
+        "elastic_backup_kind=invalid_backup_kind",
+    ]
+    with self.assertRaises(pydantic.ValidationError):
+      pyconfig.initialize(argv)
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/unit/decoder_layer_model_mode_test.py b/tests/unit/decoder_layer_model_mode_test.py
new file mode 100644
index 0000000000..1d60a3b69d
--- /dev/null
+++ b/tests/unit/decoder_layer_model_mode_test.py
@@ -0,0 +1,180 @@
+# 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.
+
+"""A decoder layer must attend in the mode it is called with, not the one it was built in.
+
+Under NNX a model object is built once and reused across modes — MaxEngine builds it
+in PREFILL and then passes MODEL_MODE_AUTOREGRESSIVE per call while decoding. A layer
+that forwards its construction-time `self.model_mode` to attention instead of the
+`model_mode` argument therefore runs the prefill KV-cache path during decode and
+produces wrong tokens from the first generated one, with no error anywhere.
+
+The check stops at the attention boundary and asserts on the mode that arrives there,
+so it stays cheap and does not depend on any particular attention implementation.
+"""
+
+import sys
+import unittest
+from unittest import mock
+
+import jax
+import jax.numpy as jnp
+from absl.testing import parameterized
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+
+from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL
+from maxtext.configs import pyconfig
+from maxtext.layers import attention_mla, attentions
+from maxtext.models import deepseek, gemma, gemma2, llama2, mistral, qwen2, qwen3
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+
+# Config knobs shared by every family; small enough that construction is the only cost.
+_COMMON = {
+    "base_emb_dim": 64,
+    "base_mlp_dim": 64,
+    "base_num_query_heads": 4,
+    "base_num_kv_heads": 4,
+    "base_num_decoder_layers": 3,
+    "vocab_size": 64,
+    "max_prefill_predict_length": 8,
+    "max_target_length": 16,
+    "per_device_batch_size": 1,
+    "scan_layers": False,
+    "attention": "dot_product",
+    "sparse_matmul": False,
+    "dtype": "float32",
+    "weight_dtype": "float32",
+    "enable_checkpointing": False,
+    "skip_jax_distributed_system": True,
+    "pure_nnx": True,
+}
+
+_DEEPSEEK = {
+    "model_name": "deepseek2-16b",
+    "override_model_config": True,
+    "mla_naive_kvcache": False,
+    "first_num_dense_layers": 1,
+    "num_experts": 4,
+    "num_experts_per_tok": 2,
+    "shared_experts": 1,
+}
+
+# Each entry is (case name, layer class, extra config, extra constructor kwargs).
+_LAYERS = [
+    ("deepseek_dense", deepseek.DeepSeekDenseLayer, _DEEPSEEK, {}),
+    ("deepseek_moe", deepseek.DeepSeekMoELayer, _DEEPSEEK, {}),
+    ("llama2", llama2.LlamaDecoderLayer, {}, {}),
+    ("mistral", mistral.MistralDecoderLayer, {"decoder_block": "mistral"}, {}),
+    ("gemma", gemma.GemmaDecoderLayer, {"decoder_block": "gemma"}, {}),
+    ("gemma2", gemma2.Gemma2DecoderLayer, {"decoder_block": "gemma2"}, {}),
+    ("qwen2", qwen2.Qwen2DecoderLayer, {"decoder_block": "qwen2"}, {"quant": None}),
+    ("qwen3", qwen3.Qwen3DecoderLayer, {"decoder_block": "qwen3"}, {"quant": None}),
+]
+
+
+def _mesh(cfg):
+  """Builds the mesh over every available device.
+
+  ici_fsdp_parallelism keeps its base.yml default of -1, so the fsdp axis absorbs
+  however many devices the host has. Pinning the ICI axes instead would fix the mesh
+  at one device and fail everywhere except a single-device runner.
+
+  Args:
+    cfg: Model config.
+
+  Returns:
+    The device mesh.
+  """
+  return jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+
+
+class _StopAtAttention(Exception):
+  """Raised by the spy so the forward pass ends once the mode has been observed."""
+
+
+class DecoderLayerModelModeTest(parameterized.TestCase):
+  """Every decoder layer must forward the call-time model_mode to its attention."""
+
+  def _mode_seen_by_attention(self, layer_cls, extra_config, ctor_kwargs, build_mode, call_mode):
+    """Builds a layer in one mode and calls it in another.
+
+    Args:
+      layer_cls: Decoder layer class to instantiate.
+      extra_config: Config overrides this family needs, on top of _COMMON.
+      ctor_kwargs: Extra constructor arguments this family requires.
+      build_mode: Model mode the layer is constructed with.
+      call_mode: Model mode the layer is called with.
+
+    Returns:
+      The model mode that reached the attention module.
+    """
+    cfg = pyconfig.initialize([sys.argv[0], get_test_config_path()], **(_COMMON | extra_config))
+    mesh = _mesh(cfg)
+    seen = []
+
+    def spy(_self, *args, **kwargs):
+      # model_mode is the 5th positional parameter of Attention.__call__, but every
+      # in-tree caller passes it by keyword. Accept both so the spy is not brittle.
+      seen.append(kwargs.get("model_mode", args[4] if len(args) > 4 else None))
+      raise _StopAtAttention()
+
+    with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh:
+      layer = layer_cls(
+          config=cfg,
+          model_mode=build_mode,
+          mesh=mesh,
+          rngs=nnx.Rngs(params=0, dropout=0),
+          **ctor_kwargs,
+      )
+      batch = cfg.micro_batch_size_to_train_on
+      inputs = jnp.zeros((batch, 1, cfg.emb_dim), dtype=jnp.float32)
+      positions = jnp.zeros((batch, 1), dtype=jnp.int32)
+      segment_ids = jnp.ones((batch, 1), dtype=jnp.int32)
+      with (
+          mock.patch.object(attentions.Attention, "__call__", spy),
+          mock.patch.object(attention_mla.MLA, "__call__", spy),
+      ):
+        try:
+          layer(inputs, segment_ids, positions, True, call_mode)
+        except _StopAtAttention:
+          pass
+
+    self.assertEqual(len(seen), 1, f"expected exactly one attention call, saw {len(seen)}")
+    return seen[0]
+
+  @parameterized.named_parameters(*_LAYERS)
+  def test_decode_call_overrides_prefill_construction(self, layer_cls, extra_config, ctor_kwargs):
+    """Checks a layer built for prefill and called for decode, as MaxEngine does."""
+    seen = self._mode_seen_by_attention(
+        layer_cls, extra_config, ctor_kwargs, MODEL_MODE_PREFILL, MODEL_MODE_AUTOREGRESSIVE
+    )
+    self.assertEqual(
+        seen,
+        MODEL_MODE_AUTOREGRESSIVE,
+        "attention ran in the layer's construction-time mode; decode would use the prefill KV-cache path",
+    )
+
+  @parameterized.named_parameters(*_LAYERS)
+  def test_prefill_call_overrides_decode_construction(self, layer_cls, extra_config, ctor_kwargs):
+    """Checks the reverse direction, so a layer cannot pass by hardcoding either mode."""
+    seen = self._mode_seen_by_attention(
+        layer_cls, extra_config, ctor_kwargs, MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL
+    )
+    self.assertEqual(seen, MODEL_MODE_PREFILL, "attention ran in the layer's construction-time mode")
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/tests/unit/deepseek_decode_consistency_test.py b/tests/unit/deepseek_decode_consistency_test.py
new file mode 100644
index 0000000000..181636d9a8
--- /dev/null
+++ b/tests/unit/deepseek_decode_consistency_test.py
@@ -0,0 +1,268 @@
+# 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.
+
+"""Greedy decode through MaxEngine must match a teacher-forced forward pass.
+
+Golden-logit tests only cover the forward pass, so nothing catches a broken
+prefill -> autoregressive handoff. Here a tiny model is decoded greedily through
+prefill/insert/generate and compared against the argmax rollout of a plain forward
+pass over the same growing prefix, which is the same sequence by definition. That
+holds for any model, so new families can be added to the parameter list.
+
+DeepSeek is the case that motivated this: MaxEngine's NNX path builds the model once
+in PREFILL mode and passes the mode per call, so a layer that reads its
+construction-time model_mode instead of the argument runs the prefill attention path
+during decode and silently produces wrong tokens.
+"""
+
+import sys
+import unittest
+from unittest import mock
+
+import jax
+import jax.numpy as jnp
+import pytest
+from absl.testing import parameterized
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+
+from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_PREFILL, MODEL_MODE_TRAIN
+from maxtext.configs import pyconfig
+from maxtext.models import deepseek
+from maxtext.utils import maxtext_utils, model_creation_utils
+
+pytest.importorskip("jetstream", reason="jetstream not installed")
+from maxtext.inference.maxengine import maxengine
+from tests.utils.test_helpers import get_test_config_path
+
+PROMPT = [3, 17, 42, 5, 9]
+STEPS = 4
+
+_COMMON = {
+    "base_emb_dim": 64,
+    "base_mlp_dim": 64,
+    "base_num_query_heads": 4,
+    "base_num_kv_heads": 4,
+    "base_num_decoder_layers": 3,
+    "vocab_size": 64,
+    "max_prefill_predict_length": 8,
+    "max_target_length": 16,
+    "per_device_batch_size": 1,
+    "scan_layers": False,
+    "attention": "dot_product",
+    "sparse_matmul": False,
+    "dtype": "float32",
+    "weight_dtype": "float32",
+    "matmul_precision": "highest",
+    "decode_sampling_strategy": "greedy",
+    "enable_checkpointing": False,
+    "skip_jax_distributed_system": True,
+    "pure_nnx": True,
+}
+
+_DEEPSEEK = {
+    "model_name": "deepseek2-16b",
+    "override_model_config": True,
+    "num_experts": 4,
+    "num_experts_per_tok": 2,
+    "shared_experts": 1,
+    # The naive MLA kv cache sizes its key cache with the value head dim, so prefill
+    # writes a wider key than the cache holds. Unrelated to decode consistency and it
+    # fails the same way on the Linen path.
+    "mla_naive_kvcache": False,
+}
+
+# MoE and dense DeepSeek layers take different paths to attention, so cover both.
+_DEEPSEEK_MOE = _DEEPSEEK | {"first_num_dense_layers": 1}
+_DEEPSEEK_DENSE = _DEEPSEEK | {"first_num_dense_layers": 3}
+
+
+def _make_config(**overrides):
+  return pyconfig.initialize([sys.argv[0], get_test_config_path()], **(_COMMON | overrides))
+
+
+def _mesh(cfg):
+  """Builds the mesh over every available device.
+
+  ici_fsdp_parallelism keeps its base.yml default of -1, so the fsdp axis absorbs
+  however many devices the host has. Pinning the ICI axes instead would fix the mesh
+  at one device and fail everywhere except a single-device runner.
+
+  Args:
+    cfg: Model config.
+
+  Returns:
+    The device mesh.
+  """
+  return jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+
+
+class DecodeConsistencyTest(parameterized.TestCase):
+  """Decode must agree with the forward pass it is supposed to be replaying."""
+
+  def _forward_rollout(self, cfg, mesh, params_state):
+    """Rolls out greedily, recomputing a full forward pass over the prefix each step.
+
+    Args:
+      cfg: Model config.
+      mesh: Device mesh the model is built on.
+      params_state: nnx.Param state to load into the model.
+
+    Returns:
+      The generated token ids.
+    """
+    with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh:
+      model = model_creation_utils.create_model(
+          cfg, mesh, model_mode=MODEL_MODE_TRAIN, rngs=nnx.Rngs(params=0, dropout=0)
+      )
+    nnx.update(model, params_state)
+
+    tokens = list(PROMPT)
+    generated = []
+    pad = cfg.max_target_length
+    batch = cfg.micro_batch_size_to_train_on
+    for _ in range(STEPS + 1):
+      ids = jnp.tile(jnp.asarray([tokens + [0] * (pad - len(tokens))], dtype=jnp.int32), (batch, 1))
+      positions = jnp.tile(jnp.asarray([list(range(pad))], dtype=jnp.int32), (batch, 1))
+      segment_ids = jnp.tile(jnp.asarray([[1] * len(tokens) + [0] * (pad - len(tokens))], dtype=jnp.int32), (batch, 1))
+      with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh:
+        logits = model(
+            ids,
+            positions,
+            decoder_segment_ids=segment_ids,
+            enable_dropout=False,
+            model_mode=MODEL_MODE_TRAIN,
+        )
+      next_token = int(jnp.argmax(logits[0, len(tokens) - 1]))
+      generated.append(next_token)
+      tokens.append(next_token)
+    return generated
+
+  def _engine_rollout(self, engine, params):
+    """Rolls out greedily through prefill, insert and generate.
+
+    Args:
+      engine: MaxEngine with params already loaded.
+      params: Params returned by load_params.
+
+    Returns:
+      The generated token ids.
+    """
+    padded = jnp.asarray(
+        PROMPT + [0] * (engine.config.max_prefill_predict_length - len(PROMPT)),
+        dtype=jnp.int32,
+    )
+    prefix, first_token = engine.prefill(params=params, padded_tokens=padded, true_length=len(PROMPT))
+    generated = [int(first_token.data[0, 0])]
+
+    decode_state = engine.init_decode_state()
+    decode_state = engine.insert(prefix, decode_state, slot=0)
+    for _ in range(STEPS):
+      decode_state, result = engine.generate(params, decode_state)
+      generated.append(int(result.data[0, 0]))
+    return generated
+
+  def _build(self, cfg):
+    """Builds a freshly initialized model.
+
+    Args:
+      cfg: Model config.
+
+    Returns:
+      A tuple of (mesh, params_state).
+    """
+    mesh = _mesh(cfg)
+    with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh:
+      model = model_creation_utils.create_model(
+          cfg, mesh, model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(params=0, dropout=0)
+      )
+    _, params_state, _ = nnx.split(model, nnx.Param, ...)
+    return mesh, params_state
+
+  @parameterized.named_parameters(
+      ("deepseek_moe", _DEEPSEEK_MOE),
+      ("deepseek_dense", _DEEPSEEK_DENSE),
+      ("generic", {}),
+  )
+  def test_greedy_decode_matches_forward_pass(self, overrides):
+    cfg = _make_config(**overrides)
+    mesh, params_state = self._build(cfg)
+
+    expected = self._forward_rollout(cfg, mesh, params_state)
+
+    engine = maxengine.MaxEngine(cfg, jax.devices())
+    params = engine.load_params(params=params_state)
+    actual = self._engine_rollout(engine, params)
+
+    self.assertEqual(
+        expected,
+        actual,
+        f"decode diverged from the forward pass: expected {expected}, got {actual}",
+    )
+
+
+class EngineGraphdefModeTest(unittest.TestCase):
+  """MaxEngine must merge the graphdef built for the mode it is running.
+
+  Layers are supposed to honor the call-time model_mode, but the engine should not
+  depend on that: merging a prefill graphdef for an autoregressive step leaves every
+  layer's `self.model_mode` reading "prefill". This guards the engine side on its own,
+  so the two defenses cannot regress together silently.
+  """
+
+  def test_generate_merges_autoregressive_graphdef(self):
+    cfg = _make_config(**_DEEPSEEK_MOE)
+    mesh = _mesh(cfg)
+    with nn_partitioning.axis_rules(cfg.logical_axis_rules), mesh:
+      model = model_creation_utils.create_model(
+          cfg, mesh, model_mode=MODEL_MODE_PREFILL, rngs=nnx.Rngs(params=0, dropout=0)
+      )
+    _, params_state, _ = nnx.split(model, nnx.Param, ...)
+
+    engine = maxengine.MaxEngine(cfg, jax.devices())
+    params = engine.load_params(params=params_state)
+
+    # The mode each merged layer was constructed with, recorded from inside the
+    # jitted prefill / generate bodies.
+    built_modes = []
+    original = deepseek.DeepSeekGenericLayer.attention_op
+
+    def recording_attention_op(self, *args, **kwargs):
+      built_modes.append(self.model_mode)
+      return original(self, *args, **kwargs)
+
+    padded = jnp.asarray(
+        PROMPT + [0] * (cfg.max_prefill_predict_length - len(PROMPT)),
+        dtype=jnp.int32,
+    )
+    with mock.patch.object(deepseek.DeepSeekGenericLayer, "attention_op", recording_attention_op):
+      prefix, _ = engine.prefill(params=params, padded_tokens=padded, true_length=len(PROMPT))
+      self.assertTrue(built_modes, "prefill did not reach the DeepSeek attention")
+      self.assertEqual(set(built_modes), {MODEL_MODE_PREFILL})
+
+      decode_state = engine.init_decode_state()
+      decode_state = engine.insert(prefix, decode_state, slot=0)
+      built_modes.clear()
+      engine.generate(params, decode_state)
+
+    self.assertTrue(built_modes, "generate did not reach the DeepSeek attention")
+    self.assertEqual(
+        set(built_modes),
+        {MODEL_MODE_AUTOREGRESSIVE},
+        "generate merged a graphdef built for another mode",
+    )
+
+
+if __name__ == "__main__":
+  unittest.main()
diff --git a/tests/unit/elastic_utils_test.py b/tests/unit/elastic_utils_test.py
index 634ec2f992..9307c7dd30 100644
--- a/tests/unit/elastic_utils_test.py
+++ b/tests/unit/elastic_utils_test.py
@@ -15,16 +15,19 @@
 """Unit tests for Elastic Training utility functions."""
 
 import unittest
-from unittest.mock import create_autospec, Mock
+from unittest.mock import Mock, create_autospec
 from absl.testing import parameterized
-
-
+from maxtext.common import checkpointing
 from maxtext.utils import elastic_utils
 from maxtext.utils import gcs_utils
 import pathwaysutils
 from pathwaysutils.elastic.manager import ScaleUpSignalError
 
 
+class MockJaxRuntimeError(Exception):
+  """Fake JAX Runtime Error class for unit tests."""
+
+
 class FakeDevice:
   """Fake Device object."""
 
@@ -77,6 +80,7 @@ def setUp(self):
     # Inject fakes into elastic_utils namespace
     elastic_utils.pathwaysutils = self.fake_pathwaysutils
     elastic_utils.jax = self.fake_jax
+    self.fake_jax.errors.JaxRuntimeError = MockJaxRuntimeError
     elastic_utils.gcs_utils = self.fake_gcs_utils
     elastic_utils.max_logging = self.fake_logging
 
@@ -467,6 +471,80 @@ def test_is_scale_up_event_with_set(self, available_inactive_slices, expected):
     self.fake_manager.available_inactive_slices = available_inactive_slices
     self.assertEqual(elastic_utils.is_scale_up_event(config), expected)
 
+  def test_maybe_bubble_elastic_exception_bubbles_on_elastic_errors(self):
+    """Tests that elastic exceptions are bubbled up, while other exceptions are returned normally."""
+    config = FakeConfig()
+    config.elastic_enabled = True
+    self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
+
+    # Scenario 1: Elastic JAX error propagates
+    with self.assertRaises(MockJaxRuntimeError):
+      elastic_utils.maybe_bubble_elastic_exception(config, MockJaxRuntimeError("TPU offline"))
+
+    # Scenario 2: ScaleUpSignalError propagates
+    with self.assertRaises(ScaleUpSignalError):
+      elastic_utils.maybe_bubble_elastic_exception(config, ScaleUpSignalError())
+
+    # Scenario 3: Non-elastic error is returned/ignored
+    elastic_utils.maybe_bubble_elastic_exception(config, ValueError("Disk full"))
+
+  def test_maybe_bubble_elastic_exception_disabled_does_not_bubble(self):
+    """If elasticity is disabled, JaxRuntimeError should not bubble."""
+    config = FakeConfig()
+    config.elastic_enabled = False  # Disabled
+    self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
+
+    # Executes normally (no exception raised)
+    elastic_utils.maybe_bubble_elastic_exception(config, MockJaxRuntimeError("JAX error but elasticity disabled"))
+
+  def test_checkpoint_exception_guard_checks_scale_up_on_success(self):
+    """Signals ScaleUpSignalError if scale-up is active when save completes."""
+    config = FakeConfig()
+    config.elastic_enabled = True
+    elastic_utils.elastic_manager = self.fake_manager
+    self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
+    self.fake_manager.available_inactive_slices = {1}  # Trigger scale-up
+
+    mock_checkpoint_manager = Mock(spec=["wait_until_finished"])
+
+    # Successful checkpoint save block raises ScaleUpSignalError to trigger restart
+    with self.assertRaises(ScaleUpSignalError):
+      with checkpointing.checkpoint_exception_guard(config, mock_checkpoint_manager):
+        pass
+
+    mock_checkpoint_manager.wait_until_finished.assert_called_once()
+
+  def test_checkpoint_exception_guard_none_manager(self):
+    """Checks that checkpoint_manager=None doesn't raise AttributeError on scale-up."""
+    config = FakeConfig()
+    config.elastic_enabled = True
+    elastic_utils.elastic_manager = self.fake_manager
+    self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
+    self.fake_manager.available_inactive_slices = {1}  # Trigger scale-up
+
+    with self.assertRaises(ScaleUpSignalError):
+      with checkpointing.checkpoint_exception_guard(config, checkpoint_manager=None):
+        pass
+
+  def test_checkpoint_exception_guard_skips_scale_up_on_failure(self):
+    """If checkpoint save fails, scale-up check should be skipped, and exception handled."""
+    config = FakeConfig()
+    config.elastic_enabled = True
+    elastic_utils.elastic_manager = self.fake_manager
+    self.fake_pathwaysutils.is_pathways_backend_used.return_value = True
+    self.fake_manager.available_inactive_slices = {1}
+
+    handler_called = False
+
+    def handler(_err):
+      nonlocal handler_called
+      handler_called = True
+
+    with checkpointing.checkpoint_exception_guard(config, self.fake_manager, handler):
+      raise ValueError("Save failed")
+
+    self.assertTrue(handler_called)
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py
index 8eb3bd70c9..96f76106ed 100644
--- a/tests/unit/moe_test.py
+++ b/tests/unit/moe_test.py
@@ -756,6 +756,7 @@ def _run_ragged_sort_loss_and_grad(
       ragged_buffer_factor: float = -1.0,
       ragged_gather_fallback: bool = False,
       ragged_gather_reduce_fallback: bool = False,
+      ragged_sort_use_single_sparsecore: bool = False,
   ):
     """Loss and gradient correctness for the use_ragged_sort flag.
 
@@ -789,6 +790,7 @@ def _build_cfg(use_ragged_sort: bool):
           ragged_buffer_factor=effective_buffer_factor,
           ragged_gather_fallback=ragged_gather_fallback,
           ragged_gather_reduce_fallback=ragged_gather_reduce_fallback,
+          ragged_sort_use_single_sparsecore=ragged_sort_use_single_sparsecore,
       )
 
     def _build_model(cfg, mesh):
@@ -900,6 +902,14 @@ def test_ragged_sort_loss_and_grad_no_ring_of_experts_fallback(self):
         use_ring_of_experts=False, ragged_gather_fallback=True, ragged_gather_reduce_fallback=True
     )
 
+  @pytest.mark.tpu_only
+  def test_ragged_sort_single_sparsecore_ring_of_experts(self):
+    self._run_ragged_sort_loss_and_grad(use_ring_of_experts=True, ragged_sort_use_single_sparsecore=True)
+
+  @pytest.mark.tpu_only
+  def test_ragged_sort_single_sparsecore_no_ring_of_experts(self):
+    self._run_ragged_sort_loss_and_grad(use_ring_of_experts=False, ragged_sort_use_single_sparsecore=True)
+
   @pytest.mark.tpu_only
   def test_moe_fsdp_two_stage_parallelism_tpu_only(self):
     # Use an imperative skip inside the test method instead of a static decorator.
diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py
index d3afcb612e..9b7862bc5a 100644
--- a/tests/unit/pyconfig_test.py
+++ b/tests/unit/pyconfig_test.py
@@ -332,6 +332,15 @@ def test_eval_start_step_config(self):
     )
     self.assertEqual(config_override.eval_start_step, 50)
 
+  def test_eval_start_step_negative_raises_error(self):
+    """Verifies that eval_start_step < 0 raises a validation error."""
+    with self.assertRaises((ValueError, Exception)):
+      pyconfig.initialize(
+          [os.path.join(MAXTEXT_PKG_DIR, "train.py"), get_test_config_path()],
+          skip_jax_distributed_system=True,
+          eval_start_step=-1,
+      )
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py
index ef19ac427e..46ee0b73e9 100644
--- a/tests/unit/train_state_nnx_checkpoint_test.py
+++ b/tests/unit/train_state_nnx_checkpoint_test.py
@@ -434,13 +434,31 @@ def test_nnx_and_linen_agree_on_actual_step(self):
     )
 
   def test_nnx_state_is_saved_in_linen_layout(self):
-    """For pure_nnx=True, maybe_save_checkpoint reshapes the NNX state to the Linen on-disk layout."""
+    """For pure_nnx=True, save_checkpoint reshapes the NNX state to the Linen on-disk layout."""
     state = self._build_nnx_state(self.N_STEPS)
     self.assertIsInstance(state, nnx.State)  # precondition: NNX train_step returns an nnx.State
 
-    captured = self._invoke_maybe_save(state, pure_nnx=True)
+    config = self._config(pure_nnx=True, enable_checkpointing=True, checkpoint_period=1)
+    mgr = mock.MagicMock()
+    mgr.reached_preemption.return_value = False
+
+    captured = {}
 
-    # save_checkpoint should receive a plain dict in Linen layout, not the nnx.State.
+    def fake_save(_step, *args, **kwargs):
+      composite = kwargs.get("args")
+      if composite:
+        for key in ["items", "state"]:
+          if hasattr(composite, "_items") and key in composite._items:  # pylint: disable=protected-access
+            val = composite[key]
+            if val is not None and hasattr(val, "item"):
+              captured["state"] = val.item
+              break
+      return True
+
+    mgr.save.side_effect = fake_save
+    checkpointing.save_checkpoint(mgr, self.N_STEPS - 1, state, config=config, force=True)
+
+    # save_checkpoint should pass a plain dict in Linen layout to Orbax, not the nnx.State.
     self.assertIsInstance(captured["state"], dict)
     self.assertNotIsInstance(captured["state"], nnx.State)
     # Linen layout: {params: {params: ...}, step, opt_state}; not the NNX {model, optimizer}.
@@ -558,16 +576,12 @@ def test_maybe_save_checkpoint_allows_local_checkpoint_period(self):
         mgr.reached_preemption.return_value = False
         save_checkpoint_mock = mock.MagicMock(return_value=False)
 
-        with (
-            mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
-            mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock,
-        ):
+        with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock):
           checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
 
         mgr.latest_step.assert_called_once_with()
         mgr.reached_preemption.assert_called_once_with(5)
         mgr.wait_until_finished.assert_not_called()
-        to_checkpoint_dict_mock.assert_called_once_with(state)
         save_checkpoint_mock.assert_called_once()
 
   def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy(
@@ -587,16 +601,12 @@ def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy(
     mgr.reached_preemption.return_value = False
     save_checkpoint_mock = mock.MagicMock(return_value=False)
 
-    with (
-        mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
-        mock.patch.object(train_state_nnx, "to_checkpoint_dict") as to_checkpoint_dict_mock,
-    ):
+    with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock):
       checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
 
     mgr.should_save.assert_called_once_with(5)
     mgr.latest_step.assert_called_once_with()
     mgr.reached_preemption.assert_called_once_with(5)
-    to_checkpoint_dict_mock.assert_called_once_with(state)
     save_checkpoint_mock.assert_called_once()
 
   def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self):
@@ -610,12 +620,10 @@ def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self):
 
     with (
         mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock),
-        mock.patch.object(train_state_nnx, "to_checkpoint_dict", return_value={}) as to_checkpoint_dict_mock,
         mock.patch.object(checkpointing.elastic_utils, "maybe_elastic_scale_up") as mock_maybe_scale_up,
     ):
       checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
 
-    to_checkpoint_dict_mock.assert_called_once_with(state)
     save_checkpoint_mock.assert_called_once()
     mock_maybe_scale_up.assert_called_once_with(config, mgr)