diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py
index c778f92bf9..2e9659543e 100644
--- a/src/maxtext/common/checkpointing.py
+++ b/src/maxtext/common/checkpointing.py
@@ -943,9 +943,11 @@ def map_to_pspec(data):
       )
       _assert_no_shaped_dtype_struct(restored_state)
     return {"items": restored_state}, None
+  if checkpoint_manager is None:
+    max_logging.log("Checkpoint manager is None, skipping checkpoint restoration.")
   else:
     max_logging.log("No existing checkpoints found, not restoring checkpoint.")
-    return None, None
+  return None, None
 
 
 def setup_checkpoint_logger(config) -> Any | None:  # pytype: disable=attribute-error
diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py
index 21bd870bc8..e31175ba14 100644
--- a/src/maxtext/common/data_loader.py
+++ b/src/maxtext/common/data_loader.py
@@ -26,6 +26,7 @@
 from maxtext.trainers.diloco import diloco
 from maxtext.utils import exceptions
 from maxtext.utils.sharding import get_input_data_sharding
+from maxtext.utils import elastic_utils
 
 
 class DataLoader:
@@ -63,6 +64,11 @@ def load_next_batch_pre_sharding(self):
         self.last_batch = example_batch
         self.check_example_batch()
       except Exception as e:  # pylint: disable=broad-except
+        if elastic_utils.elastic_enabled(self.config) and (
+            isinstance(e, jax.errors.JaxRuntimeError)
+            or isinstance(e, elastic_utils.manager.ScaleUpSignalError)
+        ):
+          raise
         if isinstance(e, StopIteration):
           raise exceptions.StopTraining(f"You may have run out of training data. Received {type(e)} exception: ({e})")
         else:
diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py
index 2137dd6482..1357bd4e76 100644
--- a/src/maxtext/common/metric_logger.py
+++ b/src/maxtext/common/metric_logger.py
@@ -83,6 +83,15 @@ def record_activation_metrics(output_metrics, intermediate_outputs, config):
       output_metrics["scalar"][f"{label}/layer_{layer_num:03d}"] = per_layer[layer_num]
 
 
+def _pin_metrics_to_host(metrics):
+  """Pins metrics arrays to host memory."""
+  def _to_pinned_host(x):
+    if hasattr(x, "sharding"):
+      return jax.device_put(x, x.sharding.with_memory_kind("pinned_host"))
+    return x
+  return jax.tree.map(_to_pinned_host, metrics)
+
+
 class MetadataKey(enum.Enum):
   PER_DEVICE_TFLOPS = "per_device_tflops"
   PER_DEVICE_TOKENS = "per_device_tokens"
@@ -149,7 +158,15 @@ def _log_training_metrics(self, metrics, step):
     """Handles training-specific metric logging."""
     # Skip logging if in profiler activation/deactivation steps
     # TODO(b/456828037): Switch to subprocess profiling to avoid timing artifacts at boundary steps.
-    scalars = metrics["scalar"]
+    def _safe_get(val):
+      if isinstance(val, jax.Array):
+        try:
+          return jax.device_get(val)
+        except Exception:
+          return 0
+      return val
+
+    scalars = jax.tree.map(_safe_get, metrics["scalar"])
     loss = scalars["learning/loss"]
     is_rampup = step < self.config.rampup_end_step
     is_metric_hidden_step = self.config.hide_profiler_step_metric and self._is_profiler_boundary_step(step)
@@ -164,10 +181,12 @@ def _log_training_metrics(self, metrics, step):
           f"completed profiler activation/deactivation step: {step}",
       )
     else:
+      active_slices = len(elastic_utils.live_slice_indices(self.config)) if elastic_utils.elastic_enabled(self.config) else 1
       log_parts.extend(
           [
               f"completed step: {step}",
               f"seconds: {scalars['perf/step_time_seconds']:.3f}",
+              f"active_slices: {active_slices}",
           ]
       )
       if elastic_utils.elastic_enabled(self.config):
@@ -384,13 +403,31 @@ def buffer_and_write_metrics(self, metrics, step, step_time_delta=None, is_train
       self._flush_one_buffered_entry(self.buffered_metrics.pop(0))
     if is_training:
       self.record_train_metrics(metrics, step, step_time_delta.total_seconds())
-      self.buffered_metrics.append(("train", step, metrics, step_time_delta))
+      # Pinned host memory transfer to allow recovery from host safely without default device bindings
+      metrics_pinned = _pin_metrics_to_host(metrics)
+      self.buffered_metrics.append(("train", step, metrics_pinned, step_time_delta))
       if self._pending_eval_step_count > 0:
         self._finalize_eval_metrics(step)
     else:
       self._pending_eval_step_count += 1
       self.buffered_metrics.append(("eval", step, metrics, step_time_delta))
 
+  def recover_metrics(self, recovered_metrics=None):
+    """Flushes and prints buffered metrics safely during recovery, then clears the buffer."""
+    if recovered_metrics is not None:
+      try:
+        scalars = recovered_metrics.get("scalar", {})
+        loss = float(scalars.get("learning/loss", 0.0))
+        step_time = float(scalars.get("perf/step_time_seconds", 0.0))
+        max_logging.log(
+            f"[METRIC RECOVERY] Successfully recovered metrics via Snapshotter | Loss: {loss:.3f} | Step Time: {step_time:.3f}s"
+        )
+      except Exception as e:
+        max_logging.log(f"[METRIC RECOVERY] Failed to read Snapshotter recovered metrics: {e}")
+
+    # Cleanly clear the buffer to prevent dead device reference errors downstream
+    self.buffered_metrics.clear()
+
   def _flush_one_buffered_entry(self, entry):
     """Dispatches a single buffered entry to the writer."""
     kind = entry[0]
diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 82e448c530..4346ae5734 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1331,6 +1331,9 @@ 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_snapshot_interval: 10
+elastic_new_slice_check_period: 10
 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 b0f407eff3..8787626a03 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1913,6 +1913,9 @@ 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: 'checkpoint' or 'snapshot'.")
+  elastic_snapshot_interval: int = Field(10, description="The interval in steps to save snapshots to host memory.")
+  elastic_new_slice_check_period: int = Field(10, description="The interval in seconds to poll for newly joined active slices.")
   elastic_timeout_seconds: int = Field(
       300,
       description=(
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index e7f5ea84b6..648fedc317 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -22,6 +22,7 @@
 import datetime
 import functools
 import os
+import threading
 
 from absl import app
 
@@ -29,6 +30,7 @@
 import optax
 
 import pathwaysutils  # pylint: disable=unused-import
+from pathwaysutils.debug import watchdog
 
 import tensorflow as tf
 
@@ -43,11 +45,12 @@
 from maxtext.configs import pyconfig
 from maxtext.utils.globals import EPS
 from maxtext.utils import elastic_utils
+
 # Placeholder: internal
 
 # pylint: disable=too-many-positional-arguments
 from maxtext.layers.multi_token_prediction import calculate_mtp_acceptance_rate, calculate_mtp_loss, mtp_acceptance, mtp_losses
-from maxtext.common import checkpointing, profiler
+from maxtext.common import checkpointing, profiler, train_state_nnx
 from maxtext.common.goodput import (
     GoodputEvent,
     RECORD_JOB_END_TIME,
@@ -72,6 +75,12 @@
 from maxtext.utils.gradient_accumulation import gradient_accumulation_loss_and_grad
 from maxtext.utils.vocabulary_tiling import vocab_tiling_linen_loss, vocab_tiling_nnx_loss
 
+import logging
+from maxtext.utils.snapshot import Snapshotter
+from pathwaysutils.elastic import manager as pathways_manager
+from pathwaysutils.elastic import elastic
+
+_logger = logging.getLogger(__name__)
 VertexTensorboardManager, _vertex_tb_is_stub = vertex_tensorboard_modules()
 
 
@@ -662,6 +671,7 @@ def training_loop_iteration(
   eval_data_iterator = python_vars["eval_data_iterator"]
   metric_logger_instance = python_vars["metric_logger_instance"]
   prof = python_vars["prof"]
+  snapshot_mgr = python_vars.get("snapshot")
 
   # Unpack immutable_data
   config = immutable_data["config"]  # for helpers
@@ -688,7 +698,11 @@ def training_loop_iteration(
     # DiLoCo's inner step takes the rng like the Linen step does.
     if isinstance(model, nn.Module) or config.enable_diloco:
       # pylint: disable=not-callable
-      step_rng_args = (jax.jit(jax.random.fold_in)(init_rng, step),)
+      nextrng = jax.jit(
+          jax.random.fold_in,
+          out_shardings=jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+      )(init_rng, step)
+      step_rng_args = (nextrng,)
     else:
       step_rng_args = ()
     with maybe_record_goodput(recorder, GoodputEvent.STEP, step):
@@ -696,6 +710,15 @@ def training_loop_iteration(
         if shard_optimizer_over_data and isinstance(model, nn.Module):
           state = sharding.maybe_shard_with_name(state, state_mesh_shardings, shard_mode)
         state, metrics = p_train_step(state, example_batch, *step_rng_args)
+        replicated_sharding = jax.sharding.NamedSharding(
+            mesh, jax.sharding.PartitionSpec()
+        ).with_memory_kind("pinned_host")
+        metrics = jax.tree.map(
+            lambda x: jax.device_put(x, replicated_sharding)
+            if isinstance(x, jax.Array)
+            else x,
+            metrics,
+        )
 
   step_time_delta = datetime.datetime.now() - last_step_completion
   last_step_completion = datetime.datetime.now()
@@ -712,7 +735,11 @@ def training_loop_iteration(
         all_host_upload=dump_hlo_upload_all,
     )
 
-  if eval_interval > 0 and step > start_step and (step + 1) % eval_interval == 0:
+  if (
+      eval_interval > 0
+      and step > start_step
+      and (step + 1) % eval_interval == 0
+  ):
     assert eval_data_iterator
     # Explicitly reset the eval iterator and counters before starting the eval loop
     eval_data_iterator.reset()
@@ -736,6 +763,7 @@ def training_loop_iteration(
       )
       eval_step_count += 1
 
+
   prof.maybe_deactivate_profiler(step, state)
 
   if step == start_step:
@@ -743,13 +771,153 @@ def training_loop_iteration(
 
   metric_logger_instance.buffer_and_write_metrics(metrics, step, step_time_delta)
 
+
+  # Async Host Backup (Elastic Mode only)
+  if snapshot_mgr is not None and step % config.elastic_snapshot_interval == 0:
+    if isinstance(model, nn.Module):
+      state_dict = {
+          "step": state.step,
+          "params": state.params,
+          "opt_state": state.opt_state,
+          "metrics": metrics,
+      }
+    else:
+      model_state = nnx.state(state.model) if hasattr(state, "model") else state
+      opt_state = nnx.state(state.optimizer) if hasattr(state, "optimizer") else None
+      state_dict = {"model": nnx.to_pure_dict(model_state), "metrics": metrics}
+      if opt_state is not None:
+        state_dict["optimizer"] = nnx.to_pure_dict(opt_state)
+    snapshot_mgr.save_pytree(step, state_dict)
+
   # Pack mutated state back to dicts
   jax_device_state["state"] = state
   python_vars["last_step_completion"] = last_step_completion
 
 
-def train_loop(config, recorder, state=None):
-  """Main Training loop."""
+def recover(
+    jax_device_state: dict[str, Any],
+    python_vars: dict[str, Any],
+    immutable_data: dict[str, Any],
+    active_state: Any = None,
+):
+  """Rebuilds MaxText JAX device state and restores state from host snapshot."""
+  config = immutable_data["config"]
+
+  metric_logger_instance = python_vars.get("metric_logger_instance")
+  if metric_logger_instance is not None:
+    metric_logger_instance.recover_metrics()
+
+  recorder = python_vars.get("recorder")
+  elastic_manager = python_vars.get("elastic_manager")
+  snapshot_mgr = python_vars.get("snapshot") or python_vars.get("snapshot_mgr") or python_vars.get("snapshot_manager")
+  if snapshot_mgr is None and config.elastic_enabled:
+    replica_axis_idx = config.mesh_axes.index("data")
+    snapshot_mgr = Snapshotter(replica_axis_index=replica_axis_idx)
+    python_vars["snapshot"] = snapshot_mgr
+
+  # 1. Find currently active slices (wait if none are active)
+  min_slices = 1 if config.elastic_min_slice_count == -1 else config.elastic_min_slice_count
+
+  _logger.info("Waiting for at least %d slices to be active for recovery...", min_slices)
+  all_active_slices = elastic.wait_for_slices(
+      slice_count=min_slices,
+      slice_to_devices=elastic_manager.slice_to_devices,
+      timeout=config.elastic_timeout_seconds,
+  )
+  elastic_manager.active_slice_indices = all_active_slices
+  _logger.info(
+      "Active slices after recovery: %s", elastic_manager.active_slice_indices
+  )
+  _logger.info(
+      "Active devices after recovery: %d", len(elastic_utils.live_devices(config))
+  )
+
+  elastic_utils.elastic_manager = elastic_manager
+
+  # Dynamically mutate the config to match the degraded slice topology
+  new_slice_count = elastic_manager.active_slice_count
+  _logger.info(
+      "[*] Dynamically mutating config.num_slices and"
+      " config.dcn_data_parallelism to: %d",
+      new_slice_count,
+  )
+  object.__setattr__(config, "num_slices", new_slice_count)
+  object.__setattr__(config, "dcn_data_parallelism", new_slice_count)
+
+  # Update DCN data parallel axis in dcn_parallelism list
+  data_axis_idx = config.mesh_axes.index("data")
+  config.dcn_parallelism[data_axis_idx] = new_slice_count
+
+  # Recalculate num_target_devices and batch sizes for the new topology
+  new_num_devices = len(elastic_utils.live_devices(config))
+  object.__setattr__(config, "num_target_devices", new_num_devices)
+
+  def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_steps):
+    if per_device_batch_size < 1.0:
+      mbs_load = num_devices * (expansion_factor if expansion_factor > 0 else 1)
+    else:
+      mbs_load = int(num_devices * per_device_batch_size * (expansion_factor if expansion_factor > 0 else 1))
+    mbs_train = int(num_devices * per_device_batch_size)
+    gbs_load = int(mbs_load * grad_accum_steps)
+    gbs_train = int(mbs_train * grad_accum_steps)
+    return gbs_load, gbs_train, mbs_train
+
+  # Update train batch sizes
+  gbs_load, gbs_train, mbs_train = calc_gbs(
+      config.per_device_batch_size,
+      config.expansion_factor_real_data,
+      new_num_devices,
+      config.gradient_accumulation_steps,
+  )
+  object.__setattr__(config, "global_batch_size_to_load", gbs_load)
+  object.__setattr__(config, "global_batch_size_to_train_on", gbs_train)
+  object.__setattr__(config, "micro_batch_size_to_train_on", mbs_train)
+
+  # Update eval batch sizes
+  gbs_load_eval, gbs_eval, mbs_eval = calc_gbs(
+      config.eval_per_device_batch_size,
+      config.expansion_factor_real_data,
+      new_num_devices,
+      1,
+  )
+  object.__setattr__(config, "global_batch_size_to_load_eval", gbs_load_eval)
+  object.__setattr__(config, "global_batch_size_to_eval_on", gbs_eval)
+  object.__setattr__(config, "micro_batch_size_to_eval_on", mbs_eval)
+
+  if config.enable_rampup_batch_size:
+    gbs_load_start = calc_gbs(
+        config.per_device_batch_size_start,
+        config.expansion_factor_real_data,
+        new_num_devices,
+        config.gradient_accumulation_steps,
+    )[0]
+    gbs_load_inc, _, _ = calc_gbs(
+        config.per_device_batch_size_increment,
+        config.expansion_factor_real_data,
+        new_num_devices,
+        config.gradient_accumulation_steps,
+    )
+    object.__setattr__(config, "global_batch_size_to_load_start", gbs_load_start)
+    object.__setattr__(config, "global_batch_size_to_load_increment", gbs_load_inc)
+
+    diff_batch_size = gbs_load - gbs_load_start
+    if gbs_load_inc > 0:
+      num_increments = diff_batch_size // gbs_load_inc
+      if num_increments > 0:
+        rampup_samples_per_increment = config.global_rampup_samples / num_increments
+        object.__setattr__(config, "rampup_samples_per_increment_to_load", rampup_samples_per_increment)
+
+        total_rampup_steps = 0
+        current_batch_size = gbs_load_start
+        for _ in range(int(num_increments)):
+          steps_for_this_stage = (
+              int(np.ceil(rampup_samples_per_increment / current_batch_size)) if current_batch_size > 0 else 0
+          )
+          total_rampup_steps += steps_for_this_stage
+          current_batch_size += gbs_load_inc
+        object.__setattr__(config, "rampup_end_step", total_rampup_steps)
+
+  # 2. Re-run setup_train_loop to rebuild Mesh, Model, Optimizers, Dataloader
   (
       init_rng,
       checkpoint_manager,
@@ -761,8 +929,274 @@ def train_loop(config, recorder, state=None):
       data_loader,
       rampup_manager,
       eval_data_iterator,
-      state,
-  ) = train_utils.setup_train_loop(config, recorder)
+      state,  # Newly initialized scratch state
+  ) = train_utils.setup_train_loop(
+      config,
+      recorder,
+      devices=elastic_utils.live_devices(config),
+      restore_checkpoint=False,
+  )
+  init_rng = jax.device_put(
+      init_rng, jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+  )
+
+  params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
+
+  # 3. Re-compile train and eval steps for the NEW mesh
+  if isinstance(model, nn.Module):
+    jit_model = model
+  else:
+    jit_model, _ = nnx.split(state)
+
+  with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+    p_train_step, p_eval_step = train_utils.jit_train_and_eval_step(
+        config,
+        jit_model,
+        mesh,
+        state,
+        state_mesh_shardings,
+        train_step,
+        eval_step,
+        eval_data_iterator,
+        params_shardings,
+    )
+
+  # 4. Restore TrainState from host snapshot or active state
+  if active_state is not None:
+    _logger.info("[*] Resharding active state directly (device-to-device)...")
+    if isinstance(model, nn.Module):
+      abstract_dict = {
+          "step": state.step,
+          "params": state.params,
+          "opt_state": state.opt_state,
+      }
+      active_dict = {
+          "step": active_state.step,
+          "params": active_state.params,
+          "opt_state": active_state.opt_state,
+      }
+      restored_dict = jax.device_put(
+          active_dict, jax.tree.map(lambda x: x.sharding, abstract_dict)
+      )
+      restored_state = state.replace(
+          step=restored_dict["step"],
+          params=restored_dict["params"],
+          opt_state=restored_dict["opt_state"],
+      )
+      restored_step = int(restored_state.step)
+    else:
+      model_state = nnx.state(state.model) if hasattr(state, "model") else state
+      opt_state = nnx.state(state.optimizer) if hasattr(state, "optimizer") else None
+      abstract_dict = {"model": nnx.to_pure_dict(model_state)}
+      if opt_state is not None:
+        abstract_dict["optimizer"] = nnx.to_pure_dict(opt_state)
+
+      act_model_state = nnx.state(active_state.model) if hasattr(active_state, "model") else active_state
+      act_opt_state = nnx.state(active_state.optimizer) if hasattr(active_state, "optimizer") else None
+      active_dict = {"model": nnx.to_pure_dict(act_model_state)}
+      if act_opt_state is not None:
+        active_dict["optimizer"] = nnx.to_pure_dict(act_opt_state)
+
+      restored_dict = jax.device_put(
+          active_dict, jax.tree.map(lambda x: x.sharding, abstract_dict)
+      )
+      if hasattr(state, "model") and "model" in restored_dict:
+        nnx.update(state.model, restored_dict["model"])
+      if hasattr(state, "optimizer") and "optimizer" in restored_dict:
+        nnx.update(state.optimizer, restored_dict["optimizer"])
+      restored_state = state
+      try:
+        if hasattr(state, "optimizer") and hasattr(state.optimizer, "step"):
+          step_val = state.optimizer.step
+          restored_step = int(getattr(step_val, "value", step_val))
+      except Exception:
+        restored_step = 0
+    _logger.info(
+        "Resharding complete. Retrying. Slices used: %s",
+        elastic_manager.active_slice_indices,
+    )
+  else:
+    if snapshot_mgr.latest is None:
+      raise RuntimeError("No snapshots available to restore from. Cannot recover.")
+    restored_step = snapshot_mgr.latest.step
+    if isinstance(model, nn.Module):
+      abstract_dict = {
+          "step": state.step,
+          "params": state.params,
+          "opt_state": state.opt_state,
+      }
+      restored_dict = snapshot_mgr.load_pytree(abstract_dict)
+      restored_state = state.replace(
+          step=restored_dict["step"],
+          params=restored_dict["params"],
+          opt_state=restored_dict["opt_state"],
+      )
+    else:
+      model_state = nnx.state(state.model) if hasattr(state, "model") else state
+      opt_state = nnx.state(state.optimizer) if hasattr(state, "optimizer") else None
+      abstract_dict = {"model": nnx.to_pure_dict(model_state)}
+      if opt_state is not None:
+        abstract_dict["optimizer"] = nnx.to_pure_dict(opt_state)
+      restored_dict = snapshot_mgr.load_pytree(abstract_dict)
+      if hasattr(state, "model") and "model" in restored_dict:
+        nnx.update(state.model, restored_dict["model"])
+      if hasattr(state, "optimizer") and "optimizer" in restored_dict:
+        nnx.update(state.optimizer, restored_dict["optimizer"])
+      restored_state = state
+      try:
+        if hasattr(state, "optimizer") and hasattr(state.optimizer, "step"):
+          step_val = state.optimizer.step
+          restored_step = int(getattr(step_val, "value", step_val))
+      except Exception:
+        pass
+    if metric_logger_instance is not None:
+      metric_logger_instance.recover_metrics(restored_dict.get("metrics"))
+
+  # Update jax_device_state with the newly built JAX objects
+  if not isinstance(model, nn.Module) and isinstance(restored_state, train_state_nnx.TrainStateNNX):
+    _, restored_state = nnx.split(restored_state)
+  jax_device_state["state"] = restored_state
+  jax_device_state["init_rng"] = init_rng
+  jax_device_state["mesh"] = mesh
+  jax_device_state["state_mesh_shardings"] = state_mesh_shardings
+  jax_device_state["p_train_step"] = p_train_step
+  jax_device_state["p_eval_step"] = p_eval_step
+
+  # Update python_vars with new loop state and dataloader
+  python_vars["step"] = restored_step
+  python_vars["data_loader"] = data_loader
+  python_vars["data_iterator"] = data_iterator
+  python_vars["eval_data_iterator"] = eval_data_iterator
+  python_vars["checkpoint_manager"] = checkpoint_manager
+  python_vars["rampup_manager"] = rampup_manager
+  python_vars["last_step_completion"] = datetime.datetime.now()
+
+  _logger.info(
+      "Recovery complete! Resuming safely at step %d...", restored_step
+  )
+
+
+def scale_up(
+    jax_device_state: dict[str, Any],
+    python_vars: dict[str, Any],
+    immutable_data: dict[str, Any],
+    active_state: Any,
+):
+  """Handles GKE scale-up by waiting for newly joined slices and recovering."""
+  elastic_manager = python_vars["elastic_manager"]
+  _logger.info("[*] GKE Scale-up detected! Re-joining slices...")
+
+  # Get the newly active slice indices
+  elastic_manager.active_slice_indices = elastic.get_active_slice_indices(
+      elastic_manager.slice_to_devices
+  )
+  elastic_manager.new_slice_event.clear()
+
+  recover(
+      jax_device_state,
+      python_vars,
+      immutable_data,
+      active_state=active_state,
+  )
+
+
+def train_loop(config, recorder, state=None):
+  """Main Training loop."""
+  elastic_manager = None
+  snapshot_mgr = None
+  devices = None
+  stop_event = None
+  monitor_thread = None
+
+  if config.elastic_enabled:
+    min_slices = config.elastic_min_slice_count
+    if min_slices == -1:
+      min_slices = config.num_slices
+
+    _logger.info(
+        "[*] Waiting for %d slices to be active at startup...", min_slices
+    )
+    all_devices = jax.devices()
+    slice_to_devices = elastic.get_slice_to_devices(all_devices)
+    elastic.wait_for_slices(
+        slice_count=min_slices,
+        slice_to_devices=slice_to_devices,
+        timeout=config.elastic_timeout_seconds,
+    )
+
+    _logger.info(
+        "[*] Pathways Elastic Training enabled. Initializing Pathways Manager..."
+    )
+    elastic_manager = pathways_manager.Manager()
+    # Use currently active slices populated by Manager constructor
+    _logger.info(
+        "[*] Active slices at startup: %s", elastic_manager.active_slice_indices
+    )
+    stop_event = threading.Event()
+    monitor_thread = threading.Thread(
+        target=elastic_manager._monitor_new_slices,  # pylint: disable=protected-access
+        args=(stop_event, config.elastic_new_slice_check_period),
+        daemon=True,
+    )
+    monitor_thread.start()
+    elastic_utils.elastic_manager = elastic_manager
+    devices = elastic_utils.live_devices(config)
+  else:
+    _logger.info("[*] Standard Non-Elastic Training.")
+
+  # Kills the workload if initialization takes longer than 20 minutes
+  with watchdog.watchdog(name="initialization", timeout=20 * 60, repeat=False):
+    setup_results = {}
+    init_complete_event = threading.Event()
+
+    def run_setup():
+      try:
+        results = train_utils.setup_train_loop(config, recorder, devices=devices)
+        setup_results['results'] = results
+      except Exception as e:
+        setup_results['exception'] = e
+      finally:
+        init_complete_event.set()
+
+    setup_thread = threading.Thread(target=run_setup, daemon=True)
+    setup_thread.start()
+
+    while True:
+      init_done = init_complete_event.wait(timeout=1)
+
+      if elastic_manager and elastic_utils.elastic_enabled(config):
+        new_slice = elastic_manager.new_slice_event.is_set()
+
+        if new_slice and not init_done:
+          max_logging.log("New slice detected during initialization. Triggering retry.")
+          raise elastic_utils.manager.ScaleUpSignalError("Scale up during initialization")
+
+        if init_done and new_slice:
+          raise elastic_utils.manager.ScaleUpSignalError("Both events set during initialization")
+
+      if init_done:
+        break
+
+    if 'exception' in setup_results:
+      raise setup_results['exception']
+
+    (
+        init_rng,
+        checkpoint_manager,
+        state_mesh_shardings,
+        model,
+        mesh,
+        learning_rate_schedule,
+        data_iterator,
+        data_loader,
+        rampup_manager,
+        eval_data_iterator,
+        state,
+    ) = setup_results['results']
+
+    init_rng = jax.device_put(
+        init_rng, jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+    )
 
   # Throttling is applied only if configured (dcn_bandwidth_limit is set).
   # The default flag value is empty, meaning no throttling is applied by default.
@@ -837,6 +1271,26 @@ def train_loop(config, recorder, state=None):
 
   elastic_utils.record_elastic_reinit_end()
 
+  # Initialize host snapshot manager only in elastic mode
+  if config.elastic_enabled:
+    replica_axis_idx = config.mesh_axes.index("data")
+    snapshot_mgr = Snapshotter(replica_axis_index=replica_axis_idx)
+    if isinstance(model, nn.Module):
+      state_dict = {
+          "step": state.step,
+          "params": state.params,
+          "opt_state": state.opt_state,
+      }
+    else:
+      model_state = nnx.state(state.model) if hasattr(state, "model") else state
+      opt_state = nnx.state(state.optimizer) if hasattr(state, "optimizer") else None
+      state_dict = {"model": nnx.to_pure_dict(model_state)}
+      if opt_state is not None:
+        state_dict["optimizer"] = nnx.to_pure_dict(opt_state)
+    snapshot_mgr.save_pytree(start_step, state_dict)
+    # Block on the first snapshot at startup to guarantee it is secured before training begins
+    snapshot_mgr.join()
+
   # Initialize dictionaries for refactored iteration
   jax_device_state = {
       "state": state,
@@ -859,6 +1313,8 @@ def train_loop(config, recorder, state=None):
       "eval_data_iterator": eval_data_iterator,
       "metric_logger_instance": metric_logger_instance,
       "prof": prof,
+      "elastic_manager": elastic_manager,
+      "snapshot": snapshot_mgr,
   }
 
   immutable_data = {
@@ -884,10 +1340,59 @@ def train_loop(config, recorder, state=None):
   try:
     python_vars["last_step_completion"] = datetime.datetime.now()
 
+    # Using while loop to allow for potential dynamic 'steps' adjustment in future
     # Using while loop to allow for potential dynamic 'steps' adjustment in future
     while python_vars["step"] < immutable_data["steps"]:
-      training_loop_iteration(jax_device_state, python_vars, immutable_data)
-      python_vars["step"] += 1
+      # Print the stacktrace every 60s and also exit the workload if longer than 600s
+      with (
+          watchdog.watchdog("step-stack-status", timeout=60),
+          watchdog.watchdog("step-timebomb", timeout=10 * 60, repeat=False),
+      ):
+        try:
+          # Scale-up check at the end of the step (only if elastic)
+          if config.elastic_enabled and elastic_manager.new_slice_event.is_set():
+            scale_up(
+                jax_device_state,
+                python_vars,
+                immutable_data,
+                active_state=jax_device_state["state"],
+            )
+            # Start snapshot save immediately on the new mesh
+            snapshot_mgr = python_vars["snapshot"]
+            state = jax_device_state["state"]
+            if isinstance(model, nn.Module):
+              state_dict = {
+                  "step": state.step,
+                  "params": state.params,
+                  "opt_state": state.opt_state,
+              }
+            else:
+              model_state = nnx.state(state.model) if hasattr(state, "model") else state
+              opt_state = nnx.state(state.optimizer) if hasattr(state, "optimizer") else None
+              state_dict = {"model": nnx.to_pure_dict(model_state)}
+              if opt_state is not None:
+                state_dict["optimizer"] = nnx.to_pure_dict(opt_state)
+            snapshot_mgr.save_pytree(
+                python_vars["step"], state_dict
+            )
+
+          training_loop_iteration(jax_device_state, python_vars, immutable_data)
+          python_vars["step"] += 1
+
+        except jax.errors.JaxRuntimeError as e:
+          if config.elastic_enabled and elastic.is_error_due_to_slice_down(e):
+            # Slice Failure Recovery
+            _logger.exception(
+                "[!] Elastic event detected around step %d", python_vars["step"]
+            )
+            recover(jax_device_state, python_vars, immutable_data)
+          else:
+            # Non-elastic or unrelated JAX error: log and re-raise
+            _logger.exception(
+                "[!] JAX Runtime Error detected around step %d. Re-raising.",
+                python_vars["step"]
+            )
+            raise
 
     # Unpack state for post-loop actions
     state = jax_device_state["state"]
@@ -904,6 +1409,11 @@ def train_loop(config, recorder, state=None):
     max_logging.log(f"Training stopped: {str(e)}")
     _job_completed_gracefully = True
   finally:
+    # Terminate monitoring thread (Elastic Mode only)
+    if stop_event is not None:
+      stop_event.set()
+    if monitor_thread is not None:
+      monitor_thread.join()
     if _job_completed_gracefully:
       record_goodput(recorder, RECORD_JOB_END_TIME)
     metric_logger_instance.flush_metrics_and_cleanup()
@@ -948,9 +1458,11 @@ def run(config, recorder):
 
 
 def get_train_func(config, recorder, argv):
-  """Returns the train function, wrapping in elastic_retry if elastic training is enabled."""
+  """Returns the train function, wrapping in elastic_retry if backup_kind is checkpoint."""
   if config.elastic_enabled:
-    max_logging.log("Elastic utils: Elastic training enabled.")
+    max_logging.log(f"Elastic utils: Elastic training enabled with {config.elastic_backup_kind} backup.")
+
+  if config.elastic_enabled and config.elastic_backup_kind == "checkpoint":
 
     def on_elastic_event():
       elastic_utils.record_elastic_event_start(recorder, config)
@@ -972,7 +1484,6 @@ def elastic_train_wrapper(argv: Sequence[str]) -> None:
         pre_callback_fn=on_slices_ready,
     )(functools.partial(elastic_train_wrapper, argv=argv))
   else:
-    # Use the already initialized variables
     def train_func():
       run(config, recorder)
 
diff --git a/src/maxtext/utils/elastic_utils.py b/src/maxtext/utils/elastic_utils.py
index 29e9cfb3fe..347da7d9c2 100644
--- a/src/maxtext/utils/elastic_utils.py
+++ b/src/maxtext/utils/elastic_utils.py
@@ -16,8 +16,12 @@
 
 import functools
 from collections import Counter
+from typing import Any
 
 import jax
+import jax.numpy as jnp
+from flax import nnx
+from maxtext.common import train_state_nnx
 from maxtext.utils import gcs_utils
 from maxtext.utils import max_logging
 import pathwaysutils
@@ -28,6 +32,56 @@
 pending_elastic_event_type = None
 
 
+def maybe_snapshot_state(
+    elastic_mgr: Any,
+    step: int,
+    state: Any,
+    force: bool = False,
+    block: bool = False,
+) -> None:
+  """Takes an elasticity snapshot of TrainStateNNX or Linen TrainState."""
+  if isinstance(state, train_state_nnx.TrainStateNNX):
+    model_state = nnx.state(state.model)
+    opt_state = nnx.state(state.optimizer)
+    snapshot_jax_arrays = {
+        "model": nnx.to_pure_dict(model_state),
+        "optimizer": nnx.to_pure_dict(opt_state),
+    }
+  else:
+    linen_dict = {
+        "params": getattr(state, "params", None),
+        "opt_state": getattr(state, "opt_state", None),
+        "step": getattr(state, "step", None),
+    }
+    snapshot_jax_arrays = train_state_nnx.from_linen_checkpoint_dict(linen_dict)
+
+  elastic_mgr.maybe_snapshot(
+      step=step,
+      snapshot_jax_arrays=snapshot_jax_arrays,
+      force=force,
+      block=block,
+  )
+
+
+def restore_resharded_state(elastic_mgr: Any, mesh: Any, state: Any):
+  """Restores state from an elasticity snapshot on a new mesh."""
+  step, snapshot_jax_arrays, _ = elastic_mgr.get_resharded_snapshot(mesh)
+
+  if isinstance(state, train_state_nnx.TrainStateNNX):
+    if "model" in snapshot_jax_arrays:
+      nnx.update(state.model, snapshot_jax_arrays["model"])
+    if "optimizer" in snapshot_jax_arrays:
+      nnx.update(state.optimizer, snapshot_jax_arrays["optimizer"])
+      state.optimizer.step.value = jnp.asarray(step, dtype=jnp.uint32)
+  else:
+    linen_dict = train_state_nnx.to_linen_checkpoint_dict(snapshot_jax_arrays)
+    state = state.replace(**linen_dict)
+    state = state.replace(step=state.step.at[None].set(step))
+
+  return step, state
+
+
+
 def record_elastic_event_start(recorder, config) -> None:
   """Records start of an elastic scale up event."""
   global pending_elastic_event_type
@@ -121,9 +175,10 @@ def live_devices(config=None):
     ensure_elastic_manager_initialized(config)
     assert elastic_manager is not None
     # Filter devices that are in active slices
-    return [
+    active_devices = [
         d for d in jax.devices() if d is not None and getattr(d, "slice_index", 0) in elastic_manager.active_slice_indices
     ]
+    return sorted(active_devices, key=lambda d: (getattr(d, "slice_index", 0), getattr(d, "coords", ())))
   return jax.devices()
 
 
@@ -210,6 +265,9 @@ def is_scale_up_event(config) -> bool:
   if elastic_enabled(config):
     ensure_elastic_manager_initialized(config)
     assert elastic_manager is not None
+    available_inactive = getattr(elastic_manager, "available_inactive_slices", None)
+    if available_inactive is not None:
+      return bool(available_inactive)
     return elastic_manager.new_slice_event.is_set()
 
   return False
diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py
index 3451ec824d..de84151daa 100644
--- a/src/maxtext/utils/maxtext_utils.py
+++ b/src/maxtext/utils/maxtext_utils.py
@@ -1475,7 +1475,10 @@ def setup_initial_state(
   )
 
   # Initialization
-  with nn_partitioning.axis_rules(config.logical_axis_rules):
+  elastic_manager = getattr(elastic_utils, "elastic_manager", None)
+  if elastic_manager and elastic_manager.new_slice_event.is_set():
+    raise elastic_utils.manager.ScaleUpSignalError("Scale up during setup (before load_state)")
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(config.logical_axis_rules):
     restored, raw_params = checkpointing.load_state_if_possible(
         checkpoint_manager,
         data_iterator,
diff --git a/src/maxtext/utils/snapshot.py b/src/maxtext/utils/snapshot.py
new file mode 100644
index 0000000000..9c15affbe4
--- /dev/null
+++ b/src/maxtext/utils/snapshot.py
@@ -0,0 +1,208 @@
+"""Manages asynchronous backups of JAX array states to pinned host memory."""
+
+import logging
+import queue
+import threading
+from typing import Any
+
+from etils import epath
+import jax
+from orbax.checkpoint.experimental.v1 import training
+from orbax.checkpoint.experimental.v1._src.tree import types as tree_types
+from pathwaysutils.experimental import concatenate_by_mesh_axis
+from pathwaysutils.experimental import split_by_mesh_axis
+
+_logger = logging.getLogger(__name__)
+
+
+class Snapshotter:
+  """Manages asynchronous backups of JAX array states to pinned host memory."""
+
+  def __init__(self, *, replica_axis_index: int = 0):
+    self._latest_snapshot: tuple[tree_types.PyTree, int] | None = None
+    self._lock = threading.Lock()
+    self._queue = queue.Queue(maxsize=1)
+    self.replica_axis_index = replica_axis_index
+    self._worker_thread = threading.Thread(target=self._worker, daemon=True)
+    self._worker_thread.start()
+
+  def _worker(self):
+    while True:
+      pinned_state, step = self._queue.get()
+      try:
+        _logger.info(
+            "[*] [Snapshot Thread] Waiting for snapshot at step %d to be ready...",
+            step,
+        )
+        jax.block_until_ready(pinned_state)
+        _logger.info(
+            "[*] [Snapshot Thread] Snapshot at step %d is ready and secured.",
+            step,
+        )
+        with self._lock:
+          self._latest_snapshot = (pinned_state, step)
+      except Exception as e:  # pylint: disable=broad-except
+        _logger.warning(
+            "[*] [Snapshot Thread] Failed to secure snapshot at step %d: %s.",
+            step,
+            e,
+        )
+      finally:
+        self._queue.task_done()
+
+  def save_pytree(
+      self, step: int, state: tree_types.PyTreeOf[jax.Array]
+  ) -> None:
+    """Move arrays onto CPU worker devices."""
+    if self._queue.full():
+      _logger.warning("Snapshotter busy. Skipping snapshot for step %d", step)
+      return
+
+    def pin_leaf(x):
+      if isinstance(x, jax.Array):
+        return jax.device_put(x, x.sharding.with_memory_kind("pinned_host"))
+      return x
+
+    pinned_state = jax.tree.map(pin_leaf, state)
+
+    self._queue.put((pinned_state, step))
+
+  def load_pytree(
+      self,
+      abstract_state: tree_types.PyTreeOf[jax.Array],
+      *,
+      reset_snapshot_state: bool = True,
+  ) -> tree_types.PyTree:
+    """Move arrays from workers onto TPU devices.
+
+    Uses `abstract_state.sharding` to properly re-partition onto the new mesh.
+
+    Args:
+      abstract_state: An abstract representation of the state, used to provide
+        the target shardings for the restored arrays on the TPU devices.
+      reset_snapshot_state: If True, clears snapshot history and resets it to
+        contain only the returned restored state (in host-pinned memory).
+
+    Returns:
+      The restored array state.
+
+    Raises:
+      RuntimeError: If no snapshots are available to restore from.
+    """
+    with self._lock:
+      if self._latest_snapshot is None:
+        raise RuntimeError("No snapshots available to restore from.")
+      pinned_state, step = self._latest_snapshot
+
+    def is_replica_active(arr):
+      try:
+        jax.block_until_ready(arr)
+        return True
+      except jax.errors.JaxRuntimeError as _:
+        return False
+
+    def get_active_pytree(x):
+      if not isinstance(x, jax.Array) or not hasattr(x.sharding, "mesh"):
+        return x
+      try:
+        dev_x = x
+        if hasattr(x.sharding, "memory_kind") and x.sharding.memory_kind == "pinned_host":
+          dev_sharding = x.sharding.with_memory_kind("device")
+          dev_x = jax.device_put(x, dev_sharding)
+        mesh_axis_name = dev_x.sharding.mesh.axis_names[self.replica_axis_index]
+        all_replicas = split_by_mesh_axis.split_by_mesh_axis(
+            dev_x,
+            mesh_axis_name,
+        )
+
+        active_replicas = [
+            replica for replica in all_replicas if is_replica_active(replica)
+        ]
+
+        if active_replicas:
+          return concatenate_by_mesh_axis.concatenate_by_mesh_axis(
+              active_replicas,
+              mesh_axis_name,
+          )
+      except Exception as e:
+        _logger.warning("split_by_mesh_axis failed (%s). Extracting addressable shards from remaining worker hosts...", e)
+
+      live_shards = []
+      for shard in getattr(x, "addressable_shards", []):
+        try:
+          jax.block_until_ready(shard.data)
+          live_shards.append(shard.data)
+        except jax.errors.JaxRuntimeError:
+          pass
+      if live_shards:
+        return live_shards[0] if len(live_shards) == 1 else jax.device_put(live_shards[0])
+      return x
+
+    _logger.info("Restoring from snapshot at step %d...", step)
+    active_pinned_state = jax.tree.map(get_active_pytree, pinned_state)
+    metrics = active_pinned_state.pop("metrics", None)
+
+    # Re-shard on host to the target device mesh
+    host_target_shardings = jax.tree.map(
+        lambda x: x.sharding.with_memory_kind("pinned_host")
+        if isinstance(x, jax.Array) and hasattr(x.sharding, "with_memory_kind")
+        else None,
+        abstract_state,
+    )
+
+    host_target_state = jax.tree.map(
+        lambda x, s: jax.device_put(x, s)
+        if isinstance(x, jax.Array) and s is not None
+        else x,
+        active_pinned_state,
+        host_target_shardings,
+    )
+
+    # Move from host back to device (TPU) memory.
+    target_device_shardings = jax.tree.map(
+        lambda x: x.sharding if isinstance(x, jax.Array) else None,
+        abstract_state,
+    )
+
+    restored_state = jax.tree.map(
+        lambda x, s: jax.device_put(x, s)
+        if isinstance(x, jax.Array) and s is not None
+        else x,
+        host_target_state,
+        target_device_shardings,
+    )
+    def safe_block(x):
+      if isinstance(x, jax.Array):
+        try:
+          jax.block_until_ready(x)
+        except Exception as e:
+          _logger.warning("Ignoring block_until_ready error on leaf: %s", e)
+      return x
+
+    jax.tree.map(safe_block, restored_state)
+
+    if metrics is not None:
+      restored_state["metrics"] = metrics
+
+    if reset_snapshot_state:
+      with self._lock:
+        self._latest_snapshot = (host_target_state, step)
+
+    return restored_state
+
+  def join(self) -> None:
+    """Blocks until all snapshots in the queue are ready and secured."""
+    self._queue.join()
+
+  @property
+  def latest(self) -> training.CheckpointMetadata[None] | None:
+    """Returns the training step of the most recently pinned backup."""
+    with self._lock:
+      if self._latest_snapshot is None:
+        return None
+      _, step = self._latest_snapshot
+    return training.CheckpointMetadata(
+        step=step,
+        path=epath.Path(),
+        metadata=None,
+    )
diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py
index 73e5d06b05..3f3204e013 100644
--- a/src/maxtext/utils/train_utils.py
+++ b/src/maxtext/utils/train_utils.py
@@ -37,6 +37,7 @@
 from maxtext.utils import model_creation_utils
 from maxtext.utils import sharding
 from maxtext.utils.rampup_batch import create_rampup_manager
+from maxtext.utils import elastic_utils
 
 
 def create_training_optimizer(config, model):
@@ -202,7 +203,7 @@ def jit_train_and_eval_step(
   return p_train_step, p_eval_step
 
 
-def setup_train_loop(config, recorder, devices=None):
+def setup_train_loop(config, recorder, devices=None, restore_checkpoint=True):
   """Set up prerequisites for the training loop -
 
       checkpoint_manager, PRNG keys, Mesh, Model and optimizer.
@@ -227,8 +228,12 @@ def setup_train_loop(config, recorder, devices=None):
 
   with maybe_record_goodput(recorder, GoodputEvent.TPU_INIT):
     is_training = True
-    init_rng = jax.random.PRNGKey(config.init_weights_seed)
     mesh = maxtext_utils.get_mesh_from_config(config, devices)
+    with jax.set_mesh(mesh):
+      init_rng = jax.random.PRNGKey(config.init_weights_seed)
+      init_rng = jax.device_put(
+          init_rng, jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
+      )
     context_parallel_size = mesh.shape.get(config.context_sharding, 1)
     if config.pure_nnx:
       # Create abstract NNX model.
@@ -254,6 +259,9 @@ def create_train_state_fn():
         validate_completed_steps(checkpoint_step + 1, config.steps)
 
   with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION):
+    elastic_manager = getattr(elastic_utils, "elastic_manager", None)
+    if elastic_manager and elastic_manager.new_slice_event.is_set():
+      raise elastic_utils.manager.ScaleUpSignalError("Scale up during setup (before data iterator)")
     data_iterator, eval_data_iterator = create_data_iterator(config, mesh)
     rampup_manager = create_rampup_manager(config, checkpoint_manager)
     # Validate context parallelism with packing configuration
@@ -299,8 +307,16 @@ 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(
-        data_iterator, config, mesh, checkpoint_manager, init_state_fn
+    if elastic_manager and elastic_manager.new_slice_event.is_set():
+      raise elastic_utils.manager.ScaleUpSignalError("Scale up during setup (before state restore)")
+    state, _, state_mesh_shardings, data_iterator = (
+        maxtext_utils.setup_training_state(
+            data_iterator,
+            config,
+            mesh,
+            checkpoint_manager if restore_checkpoint else None,
+            init_state_fn,
+        )
     )
     if config.pure_nnx:
       with nn_partitioning.axis_rules(config.logical_axis_rules):
diff --git a/tests/unit/elastic_utils_test.py b/tests/unit/elastic_utils_test.py
index 5c344feb80..bba2af2845 100644
--- a/tests/unit/elastic_utils_test.py
+++ b/tests/unit/elastic_utils_test.py
@@ -406,6 +406,79 @@ def __setattr__(self, name, value):
     elastic_utils.ensure_elastic_manager_initialized(config)
     self.assertEqual(elastic_utils.elastic_manager, self.fake_manager)
 
+  def test_is_scale_up_event_with_set(self):
+    config = FakeConfig()
+    config.elastic_enabled = True
+    elastic_utils.elastic_manager = self.fake_manager
+
+    # Scenario 1: available_inactive_slices is present and not empty
+    self.fake_manager.available_inactive_slices = {1, 2}
+    self.assertTrue(elastic_utils.is_scale_up_event(config))
+
+    # Scenario 2: available_inactive_slices is present and empty
+    self.fake_manager.available_inactive_slices = set()
+    self.assertFalse(elastic_utils.is_scale_up_event(config))
+
+    # Scenario 3: available_inactive_slices is missing (fallback to event)
+    if hasattr(self.fake_manager, "available_inactive_slices"):
+      delattr(self.fake_manager, "available_inactive_slices")
+    self.fake_manager.new_slice_event.is_set.return_value = True
+    self.assertTrue(elastic_utils.is_scale_up_event(config))
+
+    self.fake_manager.new_slice_event.is_set.return_value = False
+    self.assertFalse(elastic_utils.is_scale_up_event(config))
+
+  def test_maybe_snapshot_state_nnx(self):
+    """Tests maybe_snapshot_state with TrainStateNNX."""
+    from flax import nnx
+    from maxtext.common.train_state_nnx import TrainStateNNX
+
+    class DummyModel(nnx.Module):
+      def __init__(self, rngs):
+        self.linear = nnx.Linear(2, 2, rngs=rngs)
+
+    model = DummyModel(rngs=nnx.Rngs(0))
+    tx = nnx.optimizer.adam(0.01)
+    optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
+    state = TrainStateNNX(model, optimizer)
+
+    elastic_mgr = Mock()
+    elastic_utils.maybe_snapshot_state(elastic_mgr, step=10, state=state, force=True, block=True)
+
+    elastic_mgr.maybe_snapshot.assert_called_once()
+    _, kwargs = elastic_mgr.maybe_snapshot.call_args
+    self.assertEqual(kwargs["step"], 10)
+    self.assertIn("model", kwargs["snapshot_jax_arrays"])
+    self.assertIn("optimizer", kwargs["snapshot_jax_arrays"])
+
+  def test_restore_resharded_state_nnx(self):
+    """Tests restore_resharded_state with TrainStateNNX."""
+    from flax import nnx
+    from maxtext.common.train_state_nnx import TrainStateNNX
+
+    class DummyModel(nnx.Module):
+      def __init__(self, rngs):
+        self.linear = nnx.Linear(2, 2, rngs=rngs)
+
+    model = DummyModel(rngs=nnx.Rngs(0))
+    tx = nnx.optimizer.adam(0.01)
+    optimizer = nnx.Optimizer(model, tx, wrt=nnx.Param)
+    state = TrainStateNNX(model, optimizer)
+
+    snapshot_jax_arrays = {
+        "model": nnx.to_pure_dict(nnx.state(model)),
+        "optimizer": nnx.to_pure_dict(nnx.state(optimizer)),
+    }
+
+    elastic_mgr = Mock()
+    elastic_mgr.get_resharded_snapshot.return_value = (15, snapshot_jax_arrays, None)
+
+    restored_step, restored_state = elastic_utils.restore_resharded_state(elastic_mgr, mesh=None, state=state)
+
+    self.assertEqual(restored_step, 15)
+    self.assertEqual(int(restored_state.optimizer.step.value), 15)
+
 
 if __name__ == "__main__":
   unittest.main()
+