From 7189e025c428c70d76796ea7d36a8edffa2ddd6d Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 8 May 2026 06:36:14 +0000
Subject: [PATCH 01/31] Replica resize related changes

---
 src/maxtext/common/metric_logger.py | 2 ++
 src/maxtext/utils/maxtext_utils.py  | 3 +++
 src/maxtext/utils/train_utils.py    | 6 ++++++
 3 files changed, 11 insertions(+)

diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py
index 2137dd6482..6c0a9d292e 100644
--- a/src/maxtext/common/metric_logger.py
+++ b/src/maxtext/common/metric_logger.py
@@ -164,10 +164,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):
diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py
index 3451ec824d..ce6b12ac4c 100644
--- a/src/maxtext/utils/maxtext_utils.py
+++ b/src/maxtext/utils/maxtext_utils.py
@@ -1475,6 +1475,9 @@ def setup_initial_state(
   )
 
   # Initialization
+  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 nn_partitioning.axis_rules(config.logical_axis_rules):
     restored, raw_params = checkpointing.load_state_if_possible(
         checkpoint_manager,
diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py
index 73e5d06b05..f58d70c2c9 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):
@@ -254,6 +255,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,6 +303,8 @@ def create_train_state_fn():
     # Create data_loader AFTER reordering wrapper is applied
     data_loader = create_dataloader(config, mesh, data_iterator, recorder, rampup_manager)
 
+    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, init_state_fn
     )

From 350be3d3233709dc2d823972137d376fd9ca4223 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 8 May 2026 06:36:35 +0000
Subject: [PATCH 02/31] Adding debugging watchdogs

---
 src/maxtext/trainers/pre_train/train.py | 77 ++++++++++++++++++++-----
 1 file changed, 63 insertions(+), 14 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index e7f5ea84b6..6c661b2e40 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -29,6 +29,7 @@
 import optax
 
 import pathwaysutils  # pylint: disable=unused-import
+from pathwaysutils.debug import watchdog
 
 import tensorflow as tf
 
@@ -750,25 +751,68 @@ def training_loop_iteration(
 
 def train_loop(config, recorder, state=None):
   """Main Training loop."""
-  (
-      init_rng,
-      checkpoint_manager,
-      state_mesh_shardings,
-      model,
-      mesh,
-      learning_rate_schedule,
-      data_iterator,
-      data_loader,
-      rampup_manager,
-      eval_data_iterator,
-      state,
-  ) = train_utils.setup_train_loop(config, recorder)
+  # Kills the workload if initialization takes longer than 20 minutes
+  with watchdog.watchdog(name="initialization", timeout=20 * 60, repeat=False):
+    (
+        init_rng,
+        checkpoint_manager,
+        state_mesh_shardings,
+        model,
+        mesh,
+        learning_rate_schedule,
+        data_iterator,
+        data_loader,
+        rampup_manager,
+        eval_data_iterator,
+        state,
+    ) = train_utils.setup_train_loop(config, recorder)
+
+    start_step = get_first_step(model, state)  # this is the start_step for training
+    train_utils.validate_completed_steps(start_step, config.steps)
+
+    if isinstance(model, nn.Module):
+      jit_model = model
+    else:
+      jit_model, state = nnx.split(state)
+
+    params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
+
+    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,
+    )
+
+    with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
+      data_sharding = sharding.get_input_data_sharding(config, mesh)
+      shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding)
+      if config.shard_optimizer_over_data and isinstance(model, nn.Module):
+        state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode)
+      elif config.shard_optimizer_over_data:
+        # NNX: reshard state so params match the data-sharded in_shardings (Zero-1 layout)
+        state = jax.device_put(state, state_mesh_shardings)
+      if isinstance(model, nn.Module):
+        lower_args = (state, shaped_batch, init_rng)
+      else:
+        lower_args = (state, shaped_batch)
+      maxtext_utils.maybe_dump_jaxpr(config, p_train_step, lower_args)
+      if config.compiled_trainstep_file == "":  # compile only when there is no pre-compiled file loaded
+        compiled = p_train_step.lower(state, shaped_batch, init_rng).compile()
+        compiled_stats = compiled.memory_analysis()
+        max_utils.print_compiled_memory_stats(compiled_stats)
 
   # Throttling is applied only if configured (dcn_bandwidth_limit is set).
   # The default flag value is empty, meaning no throttling is applied by default.
   train_utils.maybe_apply_dcn_throttling(config)
 
   start_step = get_first_step(model, state)  # this is the start_step for training
+<<<<<<< HEAD
   train_utils.validate_completed_steps(start_step, config.steps)
 
   if isinstance(model, nn.Module):
@@ -886,7 +930,12 @@ def train_loop(config, recorder, state=None):
 
     # 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)
+      # 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),
+      ):
+        training_loop_iteration(jax_device_state, python_vars, immutable_data)
       python_vars["step"] += 1
 
     # Unpack state for post-loop actions

From 600569996a0ed69c4f2170e84fac93c432250297 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 8 May 2026 06:36:50 +0000
Subject: [PATCH 03/31] Optimize elastic resizing workflow by running setup in
 child thread

---
 src/maxtext/trainers/pre_train/train.py | 39 ++++++++++++++++++++++++-
 1 file changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 6c661b2e40..fefdd96197 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
 
@@ -753,6 +754,42 @@ def train_loop(config, recorder, state=None):
   """Main Training loop."""
   # 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)
+        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()
+
+    elastic_manager = getattr(elastic_utils, "elastic_manager", None)
+
+    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,
@@ -765,7 +802,7 @@ def train_loop(config, recorder, state=None):
         rampup_manager,
         eval_data_iterator,
         state,
-    ) = train_utils.setup_train_loop(config, recorder)
+    ) = setup_results['results']
 
     start_step = get_first_step(model, state)  # this is the start_step for training
     train_utils.validate_completed_steps(start_step, config.steps)

From e46959dbbc921626c3ed8b65fa7e280d8e9d1a7c Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Tue, 23 Jun 2026 18:19:46 +0000
Subject: [PATCH 04/31] Update scale-up check to use available_inactive_slices
 set

---
 src/maxtext/utils/elastic_utils.py |  3 +++
 tests/unit/elastic_utils_test.py   | 22 ++++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/src/maxtext/utils/elastic_utils.py b/src/maxtext/utils/elastic_utils.py
index 29e9cfb3fe..2b423b3cf5 100644
--- a/src/maxtext/utils/elastic_utils.py
+++ b/src/maxtext/utils/elastic_utils.py
@@ -210,6 +210,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/tests/unit/elastic_utils_test.py b/tests/unit/elastic_utils_test.py
index 5c344feb80..8b2d3b772e 100644
--- a/tests/unit/elastic_utils_test.py
+++ b/tests/unit/elastic_utils_test.py
@@ -406,6 +406,28 @@ 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))
+
 
 if __name__ == "__main__":
   unittest.main()

From f5dbe2cd4dce871ff2ae51c546b416026d5a7df1 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 14 May 2026 01:32:12 +0000
Subject: [PATCH 05/31] Add pathways elastic training resiliency (asynchronous
 host snapshotting, dynamic mesh recovery, and GKE scale-up)

Integrates snapshot-based failure recovery directly into the training iteration.
- Implements asynchronous host snapshot manager (snapshot.py) to copy TrainState parameters to CPU pinned host memory in a background daemon thread, minimizing host memory overhead by freeing old snapshot buffers before allocating new transfers.
- Integrates recovery and scale-up pathways inside the training loop in train.py, including dynamic configuration mutation and JAX set_mesh compiler safety isolation.
- Utilizes the standard pathwaysutils.elastic.manager.Manager directly to track slice health, eliminating custom manager duplication and optimizing startup times by avoiding redundant checks.
- Optimizes RNG key sharding during JIT by defining out_shardings directly inside jax.jit(jax.random.fold_in, out_shardings=...).
- Incorporates configurable elastic_snapshot_interval parameter in Pydantic settings.
- Clarifies checkpoint restoration logging.
---
 src/maxtext/common/checkpointing.py     |   4 +-
 src/maxtext/common/metric_logger.py     |  41 ++-
 src/maxtext/configs/base.yml            |   3 +
 src/maxtext/configs/types.py            |   3 +
 src/maxtext/trainers/pre_train/train.py | 337 ++++++++++++++++++++----
 src/maxtext/utils/elastic_utils.py      |   3 +-
 src/maxtext/utils/maxtext_utils.py      |   2 +-
 src/maxtext/utils/snapshot.py           | 164 ++++++++++++
 src/maxtext/utils/train_utils.py        |  18 +-
 9 files changed, 518 insertions(+), 57 deletions(-)
 create mode 100644 src/maxtext/utils/snapshot.py

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/metric_logger.py b/src/maxtext/common/metric_logger.py
index 6c0a9d292e..bde865854d 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"
@@ -386,13 +395,43 @@ 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):
+    """Flushes and prints buffered metrics safely during recovery, then clears the buffer."""
+    train_entry = None
+    for entry in self.buffered_metrics:
+      if entry[0] == "train":
+        train_entry = entry
+        break
+    if train_entry is not None:
+      (_, step_to_write, metrics_to_write, _) = train_entry
+      try:
+        # Pull loss/perplexity to print safely
+        scalars = metrics_to_write["scalar"]
+        loss = float(scalars["learning/loss"])
+        step_time = float(scalars.get("perf/step_time_seconds", 0.0))
+        max_logging.log(
+            f"[METRIC RECOVERY] Successfully recovered metrics for step {step_to_write} | Loss: {loss:.3f} | Step Time: {step_time:.3f}s"
+        )
+        # Try to write them to local/tensorboard if possible, but catch to avoid block
+        try:
+          self.write_metrics(metrics_to_write, step_to_write)
+        except Exception as e:
+          max_logging.log(f"[METRIC RECOVERY] Skipped standard flush: {e}")
+      except Exception as e:
+        max_logging.log(f"[METRIC RECOVERY] Failed to read buffered 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 fefdd96197..c2e113aa70 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -45,6 +45,7 @@
 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
@@ -74,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()
 
 
@@ -664,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
@@ -690,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):
@@ -714,7 +726,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()
@@ -738,6 +754,7 @@ def training_loop_iteration(
       )
       eval_step_count += 1
 
+
   prof.maybe_deactivate_profiler(step, state)
 
   if step == start_step:
@@ -745,13 +762,220 @@ 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:
+    state_dict = {
+        "step": state.step,
+        "params": state.params,
+        "opt_state": state.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 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"]
+  if config.pure_nnx:
+    raise NotImplementedError("Elastic recovery is not supported for NNX.")
+
+  _logger.info("[*] Recovering JAX device state from host snapshot...")
+  elastic_manager = python_vars["elastic_manager"]
+  snapshot_mgr = python_vars["snapshot"]
+  recorder = python_vars["recorder"]
+
+  # Safe Metrics Extraction & Flushing
+  metric_logger = python_vars.get("metric_logger")
+  if metric_logger is not None:
+    metric_logger.recover_metrics()
+
+  # 1. Find currently active slices
+  all_active_slices = elastic.get_active_slice_indices(
+      elastic_manager.slice_to_devices
+  )
+  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
+
+  # 2. Re-run setup_train_loop to rebuild Mesh, Model, Optimizers, Dataloader
+  (
+      init_rng,
+      checkpoint_manager,
+      state_mesh_shardings,
+      model,
+      mesh,
+      learning_rate_schedule,
+      data_iterator,
+      data_loader,
+      rampup_manager,
+      eval_data_iterator,
+      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
+  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,
+        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)...")
+    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:
+    restored_step = snapshot_mgr.latest.step
+    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"],
+    )
+
+  # Update jax_device_state with the newly built JAX objects
+  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:
+    _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 = {}
@@ -759,7 +983,7 @@ def train_loop(config, recorder, state=None):
 
     def run_setup():
       try:
-        results = train_utils.setup_train_loop(config, recorder)
+        results = train_utils.setup_train_loop(config, recorder, devices=devices)
         setup_results['results'] = results
       except Exception as e:
         setup_results['exception'] = e
@@ -769,8 +993,6 @@ def run_setup():
     setup_thread = threading.Thread(target=run_setup, daemon=True)
     setup_thread.start()
 
-    elastic_manager = getattr(elastic_utils, "elastic_manager", None)
-
     while True:
       init_done = init_complete_event.wait(timeout=1)
 
@@ -804,52 +1026,15 @@ def run_setup():
         state,
     ) = setup_results['results']
 
-    start_step = get_first_step(model, state)  # this is the start_step for training
-    train_utils.validate_completed_steps(start_step, config.steps)
-
-    if isinstance(model, nn.Module):
-      jit_model = model
-    else:
-      jit_model, state = nnx.split(state)
-
-    params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings)
-
-    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,
+    init_rng = jax.device_put(
+        init_rng, jax.sharding.NamedSharding(mesh, jax.sharding.PartitionSpec())
     )
 
-    with jax.set_mesh(mesh), mesh, nn_partitioning.axis_rules(config.logical_axis_rules):
-      data_sharding = sharding.get_input_data_sharding(config, mesh)
-      shaped_batch = maxtext_utils.get_shaped_batch(config, batch_sharding=data_sharding)
-      if config.shard_optimizer_over_data and isinstance(model, nn.Module):
-        state = sharding.maybe_shard_with_name(state, state_mesh_shardings, config.shard_mode)
-      elif config.shard_optimizer_over_data:
-        # NNX: reshard state so params match the data-sharded in_shardings (Zero-1 layout)
-        state = jax.device_put(state, state_mesh_shardings)
-      if isinstance(model, nn.Module):
-        lower_args = (state, shaped_batch, init_rng)
-      else:
-        lower_args = (state, shaped_batch)
-      maxtext_utils.maybe_dump_jaxpr(config, p_train_step, lower_args)
-      if config.compiled_trainstep_file == "":  # compile only when there is no pre-compiled file loaded
-        compiled = p_train_step.lower(state, shaped_batch, init_rng).compile()
-        compiled_stats = compiled.memory_analysis()
-        max_utils.print_compiled_memory_stats(compiled_stats)
-
   # Throttling is applied only if configured (dcn_bandwidth_limit is set).
   # The default flag value is empty, meaning no throttling is applied by default.
   train_utils.maybe_apply_dcn_throttling(config)
 
   start_step = get_first_step(model, state)  # this is the start_step for training
-<<<<<<< HEAD
   train_utils.validate_completed_steps(start_step, config.steps)
 
   if isinstance(model, nn.Module):
@@ -918,6 +1103,21 @@ def run_setup():
 
   elastic_utils.record_elastic_reinit_end()
 
+  # Initialize host snapshot manager only in elastic mode
+  if config.elastic_enabled:
+    if not isinstance(model, nn.Module):
+      raise NotImplementedError("Elastic training with snapshots is not supported for NNX yet.")
+    replica_axis_idx = config.mesh_axes.index("data")
+    snapshot_mgr = Snapshotter(replica_axis_index=replica_axis_idx)
+    state_dict = {
+        "step": state.step,
+        "params": state.params,
+        "opt_state": state.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,
@@ -940,6 +1140,8 @@ def run_setup():
       "eval_data_iterator": eval_data_iterator,
       "metric_logger_instance": metric_logger_instance,
       "prof": prof,
+      "elastic_manager": elastic_manager,
+      "snapshot": snapshot_mgr,
   }
 
   immutable_data = {
@@ -965,6 +1167,7 @@ def run_setup():
   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"]:
       # Print the stacktrace every 60s and also exit the workload if longer than 600s
@@ -972,8 +1175,38 @@ def run_setup():
           watchdog.watchdog("step-stack-status", timeout=60),
           watchdog.watchdog("step-timebomb", timeout=10 * 60, repeat=False),
       ):
-        training_loop_iteration(jax_device_state, python_vars, immutable_data)
-      python_vars["step"] += 1
+        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"]
+            snapshot_mgr.save_pytree(
+                python_vars["step"], jax_device_state["state"]
+            )
+
+          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"]
@@ -990,6 +1223,11 @@ def run_setup():
     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()
@@ -1034,9 +1272,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)
@@ -1058,7 +1298,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 2b423b3cf5..ab5f9a9354 100644
--- a/src/maxtext/utils/elastic_utils.py
+++ b/src/maxtext/utils/elastic_utils.py
@@ -121,9 +121,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()
 
 
diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py
index ce6b12ac4c..de84151daa 100644
--- a/src/maxtext/utils/maxtext_utils.py
+++ b/src/maxtext/utils/maxtext_utils.py
@@ -1478,7 +1478,7 @@ def setup_initial_state(
   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 nn_partitioning.axis_rules(config.logical_axis_rules):
+  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..4b97168117
--- /dev/null
+++ b/src/maxtext/utils/snapshot.py
@@ -0,0 +1,164 @@
+"""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
+
+    pinned_shardings = jax.tree.map(
+        lambda x: x.sharding.with_memory_kind("pinned_host"), state
+    )
+
+    pinned_state = jax.device_put(state, pinned_shardings)
+
+    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):
+      mesh_axis_name = x.sharding.mesh.axis_names[self.replica_axis_index]
+      all_replicas = split_by_mesh_axis.split_by_mesh_axis(
+          x,
+          mesh_axis_name,
+      )
+
+      active_replicas = [
+          replica for replica in all_replicas if is_replica_active(replica)
+      ]
+
+      if not active_replicas:
+        raise RuntimeError(
+            "No active replicas found."
+        )
+
+      reconstructed_state = concatenate_by_mesh_axis.concatenate_by_mesh_axis(
+          active_replicas,
+          mesh_axis_name,
+      )
+      return reconstructed_state
+
+    _logger.info("Restoring from snapshot at step %d...", step)
+    pinned_state = jax.tree.map(get_active_pytree, pinned_state)
+
+    # Re-shard on host to the target device mesh
+    host_target_shardings = jax.tree.map(
+        lambda x: x.sharding.with_memory_kind("pinned_host"), abstract_state
+    )
+
+    host_target_state = jax.device_put(
+        pinned_state, host_target_shardings
+    )
+
+    # Move from host back to device (TPU) memory.
+    restored_state = jax.device_put(
+        host_target_state, jax.tree.map(lambda x: x.sharding, abstract_state)
+    )
+    jax.block_until_ready(restored_state)
+
+    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 f58d70c2c9..3f3204e013 100644
--- a/src/maxtext/utils/train_utils.py
+++ b/src/maxtext/utils/train_utils.py
@@ -203,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.
@@ -228,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.
@@ -305,8 +309,14 @@ def create_train_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, init_state_fn
+    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):

From 463e6dc2fe6bf3ae1cfcb4c610e8aaa7734d826a Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 14 May 2026 01:32:29 +0000
Subject: [PATCH 06/31] Add scale test debugging and logging

---
 src/maxtext/trainers/pre_train/train.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index c2e113aa70..c7845bb5c3 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -887,6 +887,10 @@ def recover(
         opt_state=restored_dict["opt_state"],
     )
     restored_step = int(restored_state.step)
+    _logger.info(
+        "Resharding complete. Retrying. Slices used: %s",
+        elastic_manager.active_slice_indices,
+    )
   else:
     restored_step = snapshot_mgr.latest.step
     abstract_dict = {

From 4f0d0b3394587aac54ac477f6464ef87d34cd7b4 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 15 May 2026 06:47:06 +0000
Subject: [PATCH 07/31] Fix elasticity bugs: wait for slices at startup and
 recalculate batch sizes on recovery

---
 src/maxtext/trainers/pre_train/train.py | 85 +++++++++++++++++++++++++
 1 file changed, 85 insertions(+)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index c7845bb5c3..fef1d7979d 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -826,6 +826,76 @@ def recover(
   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)
+
+  # Update rampup batch size parameters if enabled
+  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,
+    )
+    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,
@@ -960,6 +1030,21 @@ def train_loop(config, recorder, state=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..."
     )

From 78304e2ba37febc853c186d80eb878b6543228e1 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 15 May 2026 18:25:07 +0000
Subject: [PATCH 08/31] Fix save_pytree mismatch: save state_dict instead of
 TrainState after scale-up

---
 src/maxtext/trainers/pre_train/train.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index fef1d7979d..8e90a7df17 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -1275,8 +1275,14 @@ def run_setup():
             )
             # Start snapshot save immediately on the new mesh
             snapshot_mgr = python_vars["snapshot"]
+            state = jax_device_state["state"]
+            state_dict = {
+                "step": state.step,
+                "params": state.params,
+                "opt_state": state.opt_state,
+            }
             snapshot_mgr.save_pytree(
-                python_vars["step"], jax_device_state["state"]
+                python_vars["step"], state_dict
             )
 
           training_loop_iteration(jax_device_state, python_vars, immutable_data)

From 1986e9e3e334a659ca2ec700c33437a6a1678682 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 15 May 2026 18:27:28 +0000
Subject: [PATCH 09/31] Fix elasticity recovery: wait for active slices and add
 empty snapshot check

---
 src/maxtext/trainers/pre_train/train.py | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 8e90a7df17..9a44db67e9 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -798,9 +798,16 @@ def recover(
   if metric_logger is not None:
     metric_logger.recover_metrics()
 
-  # 1. Find currently active slices
-  all_active_slices = elastic.get_active_slice_indices(
-      elastic_manager.slice_to_devices
+  # 1. Find currently active slices (wait if none are active)
+  min_slices = config.elastic_min_slice_count
+  if min_slices == -1:
+    min_slices = config.num_slices
+
+  _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(
@@ -962,6 +969,8 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
         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
     abstract_dict = {
         "step": state.step,

From 33a981746e55f3ce20fd6a99e05b85694496982d Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Fri, 5 Jun 2026 17:04:04 +0000
Subject: [PATCH 10/31] Remove comment and avoid unpacking gbs_load_start in
 recover

---
 src/maxtext/trainers/pre_train/train.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 9a44db67e9..915d516c96 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -869,14 +869,13 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
   object.__setattr__(config, "global_batch_size_to_eval_on", gbs_eval)
   object.__setattr__(config, "micro_batch_size_to_eval_on", mbs_eval)
 
-  # Update rampup batch size parameters if enabled
   if config.enable_rampup_batch_size:
-    gbs_load_start, _, _ = calc_gbs(
+    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,

From 5749742977d5ee1ab9a6acdad8979b992ba3874b Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Wed, 10 Jun 2026 23:05:40 +0000
Subject: [PATCH 11/31] Fix recover to use metric_logger_instance

---
 src/maxtext/trainers/pre_train/train.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 915d516c96..e23ac25f44 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -794,9 +794,9 @@ def recover(
   recorder = python_vars["recorder"]
 
   # Safe Metrics Extraction & Flushing
-  metric_logger = python_vars.get("metric_logger")
-  if metric_logger is not None:
-    metric_logger.recover_metrics()
+  metric_logger_instance = python_vars.get("metric_logger_instance")
+  if metric_logger_instance is not None:
+    metric_logger_instance.recover_metrics()
 
   # 1. Find currently active slices (wait if none are active)
   min_slices = config.elastic_min_slice_count

From e834104d6cdcd481e155a0e7cd10a72a2d389f21 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Wed, 1 Jul 2026 07:09:29 +0000
Subject: [PATCH 12/31] Use Snapshotter for metrics recovery and handle
 non-Array leaves

---
 src/maxtext/common/metric_logger.py     | 24 ++++----------
 src/maxtext/trainers/pre_train/train.py | 19 +++++++----
 src/maxtext/utils/snapshot.py           | 43 +++++++++++++++++++------
 3 files changed, 51 insertions(+), 35 deletions(-)

diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py
index bde865854d..c414c6f843 100644
--- a/src/maxtext/common/metric_logger.py
+++ b/src/maxtext/common/metric_logger.py
@@ -404,30 +404,18 @@ def buffer_and_write_metrics(self, metrics, step, step_time_delta=None, is_train
       self._pending_eval_step_count += 1
       self.buffered_metrics.append(("eval", step, metrics, step_time_delta))
 
-  def recover_metrics(self):
+  def recover_metrics(self, recovered_metrics=None):
     """Flushes and prints buffered metrics safely during recovery, then clears the buffer."""
-    train_entry = None
-    for entry in self.buffered_metrics:
-      if entry[0] == "train":
-        train_entry = entry
-        break
-    if train_entry is not None:
-      (_, step_to_write, metrics_to_write, _) = train_entry
+    if recovered_metrics is not None:
       try:
-        # Pull loss/perplexity to print safely
-        scalars = metrics_to_write["scalar"]
-        loss = float(scalars["learning/loss"])
+        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 for step {step_to_write} | Loss: {loss:.3f} | Step Time: {step_time:.3f}s"
+            f"[METRIC RECOVERY] Successfully recovered metrics via Snapshotter | Loss: {loss:.3f} | Step Time: {step_time:.3f}s"
         )
-        # Try to write them to local/tensorboard if possible, but catch to avoid block
-        try:
-          self.write_metrics(metrics_to_write, step_to_write)
-        except Exception as e:
-          max_logging.log(f"[METRIC RECOVERY] Skipped standard flush: {e}")
       except Exception as e:
-        max_logging.log(f"[METRIC RECOVERY] Failed to read buffered metrics: {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()
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index e23ac25f44..c94c4554f1 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -710,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())
+        if "scalar" in metrics:
+          for k, v in metrics["scalar"].items():
+            if isinstance(v, jax.Array):
+              metrics["scalar"][k] = jax.device_put(v, replicated_sharding)
+        if "scalars" in metrics:
+          for k, v in metrics["scalars"].items():
+            if isinstance(v, jax.Array):
+              metrics["scalars"][k] = jax.device_put(v, replicated_sharding)
 
   step_time_delta = datetime.datetime.now() - last_step_completion
   last_step_completion = datetime.datetime.now()
@@ -769,6 +778,7 @@ def training_loop_iteration(
         "step": state.step,
         "params": state.params,
         "opt_state": state.opt_state,
+        "metrics": metrics,
     }
     snapshot_mgr.save_pytree(step, state_dict)
 
@@ -788,13 +798,6 @@ def recover(
   if config.pure_nnx:
     raise NotImplementedError("Elastic recovery is not supported for NNX.")
 
-  _logger.info("[*] Recovering JAX device state from host snapshot...")
-  elastic_manager = python_vars["elastic_manager"]
-  snapshot_mgr = python_vars["snapshot"]
-  recorder = python_vars["recorder"]
-
-  # Safe Metrics Extraction & Flushing
-  metric_logger_instance = python_vars.get("metric_logger_instance")
   if metric_logger_instance is not None:
     metric_logger_instance.recover_metrics()
 
@@ -982,6 +985,8 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
         params=restored_dict["params"],
         opt_state=restored_dict["opt_state"],
     )
+    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
   jax_device_state["state"] = restored_state
diff --git a/src/maxtext/utils/snapshot.py b/src/maxtext/utils/snapshot.py
index 4b97168117..aa3e85ea3c 100644
--- a/src/maxtext/utils/snapshot.py
+++ b/src/maxtext/utils/snapshot.py
@@ -58,11 +58,12 @@ def save_pytree(
       _logger.warning("Snapshotter busy. Skipping snapshot for step %d", step)
       return
 
-    pinned_shardings = jax.tree.map(
-        lambda x: x.sharding.with_memory_kind("pinned_host"), state
-    )
+    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.device_put(state, pinned_shardings)
+    pinned_state = jax.tree.map(pin_leaf, state)
 
     self._queue.put((pinned_state, step))
 
@@ -101,6 +102,8 @@ def is_replica_active(arr):
         return False
 
     def get_active_pytree(x):
+      if not isinstance(x, jax.Array) or not hasattr(x.sharding, "mesh"):
+        return x
       mesh_axis_name = x.sharding.mesh.axis_names[self.replica_axis_index]
       all_replicas = split_by_mesh_axis.split_by_mesh_axis(
           x,
@@ -123,23 +126,43 @@ def get_active_pytree(x):
       return reconstructed_state
 
     _logger.info("Restoring from snapshot at step %d...", step)
-    pinned_state = jax.tree.map(get_active_pytree, pinned_state)
+    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"), abstract_state
+        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.device_put(
-        pinned_state, host_target_shardings
+    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.
-    restored_state = jax.device_put(
-        host_target_state, jax.tree.map(lambda x: x.sharding, abstract_state)
+    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,
     )
     jax.block_until_ready(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)

From e14ab317a455574938aa244a463d39678bd7d9c9 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 05:00:27 +0000
Subject: [PATCH 13/31] Support NNX and Linen in elasticity snapshot functions

---
 src/maxtext/utils/elastic_utils.py | 54 ++++++++++++++++++++++++++++++
 tests/unit/elastic_utils_test.py   | 51 ++++++++++++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/src/maxtext/utils/elastic_utils.py b/src/maxtext/utils/elastic_utils.py
index ab5f9a9354..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
diff --git a/tests/unit/elastic_utils_test.py b/tests/unit/elastic_utils_test.py
index 8b2d3b772e..bba2af2845 100644
--- a/tests/unit/elastic_utils_test.py
+++ b/tests/unit/elastic_utils_test.py
@@ -428,6 +428,57 @@ def test_is_scale_up_event_with_set(self):
     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()
+

From d8783d48abe39aee10da879c66fa1ac6b7731dd2 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:39:27 +0000
Subject: [PATCH 14/31] Support TrainStateNNX in Snapshotter save/load/recover
 for elastic training

---
 src/maxtext/common/metric_logger.py     |  10 +-
 src/maxtext/trainers/pre_train/train.py | 180 ++++++++++++++++--------
 2 files changed, 131 insertions(+), 59 deletions(-)

diff --git a/src/maxtext/common/metric_logger.py b/src/maxtext/common/metric_logger.py
index c414c6f843..1357bd4e76 100644
--- a/src/maxtext/common/metric_logger.py
+++ b/src/maxtext/common/metric_logger.py
@@ -158,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)
diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index c94c4554f1..21838a152b 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -709,16 +709,15 @@ def training_loop_iteration(
       with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules):
         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())
-        if "scalar" in metrics:
-          for k, v in metrics["scalar"].items():
-            if isinstance(v, jax.Array):
-              metrics["scalar"][k] = jax.device_put(v, replicated_sharding)
-        if "scalars" in metrics:
-          for k, v in metrics["scalars"].items():
-            if isinstance(v, jax.Array):
-              metrics["scalars"][k] = jax.device_put(v, replicated_sharding)
+        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()
@@ -774,12 +773,21 @@ def training_loop_iteration(
 
   # Async Host Backup (Elastic Mode only)
   if snapshot_mgr is not None and step % config.elastic_snapshot_interval == 0:
-    state_dict = {
-        "step": state.step,
-        "params": state.params,
-        "opt_state": state.opt_state,
-        "metrics": metrics,
-    }
+    if isinstance(state, train_state_nnx.TrainStateNNX):
+      model_state = nnx.state(state.model)
+      opt_state = nnx.state(state.optimizer)
+      state_dict = {
+          "model": nnx.to_pure_dict(model_state),
+          "optimizer": nnx.to_pure_dict(opt_state),
+          "metrics": metrics,
+      }
+    else:
+      state_dict = {
+          "step": state.step,
+          "params": state.params,
+          "opt_state": state.opt_state,
+          "metrics": metrics,
+      }
     snapshot_mgr.save_pytree(step, state_dict)
 
   # Pack mutated state back to dicts
@@ -798,9 +806,17 @@ def recover(
   if config.pure_nnx:
     raise NotImplementedError("Elastic recovery is not supported for NNX.")
 
+  metric_logger_instance = python_vars.get("metric_logger_instance")
   if metric_logger_instance is not None:
     metric_logger_instance.recover_metrics()
 
+  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 = config.elastic_min_slice_count
   if min_slices == -1:
@@ -947,25 +963,46 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
   # 4. Restore TrainState from host snapshot or active state
   if active_state is not None:
     _logger.info("[*] Resharding active state directly (device-to-device)...")
-    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)
+    if isinstance(state, train_state_nnx.TrainStateNNX):
+      model_state = nnx.state(state.model)
+      opt_state = nnx.state(state.optimizer)
+      abstract_dict = {
+          "model": nnx.to_pure_dict(model_state),
+          "optimizer": nnx.to_pure_dict(opt_state),
+      }
+      act_model_state = nnx.state(active_state.model)
+      act_opt_state = nnx.state(active_state.optimizer)
+      active_dict = {
+          "model": nnx.to_pure_dict(act_model_state),
+          "optimizer": nnx.to_pure_dict(act_opt_state),
+      }
+      restored_dict = jax.device_put(
+          active_dict, jax.tree.map(lambda x: x.sharding, abstract_dict)
+      )
+      nnx.update(state.model, restored_dict["model"])
+      nnx.update(state.optimizer, restored_dict["optimizer"])
+      restored_step = int(state.optimizer.step.value)
+      restored_state = state
+    else:
+      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)
     _logger.info(
         "Resharding complete. Retrying. Slices used: %s",
         elastic_manager.active_slice_indices,
@@ -974,17 +1011,30 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
     if snapshot_mgr.latest is None:
       raise RuntimeError("No snapshots available to restore from. Cannot recover.")
     restored_step = snapshot_mgr.latest.step
-    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"],
-    )
+    if isinstance(state, train_state_nnx.TrainStateNNX):
+      model_state = nnx.state(state.model)
+      opt_state = nnx.state(state.optimizer)
+      abstract_dict = {
+          "model": nnx.to_pure_dict(model_state),
+          "optimizer": nnx.to_pure_dict(opt_state),
+      }
+      restored_dict = snapshot_mgr.load_pytree(abstract_dict)
+      nnx.update(state.model, restored_dict["model"])
+      nnx.update(state.optimizer, restored_dict["optimizer"])
+      restored_step = int(state.optimizer.step.value)
+      restored_state = state
+    else:
+      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"],
+      )
     if metric_logger_instance is not None:
       metric_logger_instance.recover_metrics(restored_dict.get("metrics"))
 
@@ -1207,15 +1257,21 @@ def run_setup():
 
   # Initialize host snapshot manager only in elastic mode
   if config.elastic_enabled:
-    if not isinstance(model, nn.Module):
-      raise NotImplementedError("Elastic training with snapshots is not supported for NNX yet.")
     replica_axis_idx = config.mesh_axes.index("data")
     snapshot_mgr = Snapshotter(replica_axis_index=replica_axis_idx)
-    state_dict = {
-        "step": state.step,
-        "params": state.params,
-        "opt_state": state.opt_state,
-    }
+    if isinstance(state, train_state_nnx.TrainStateNNX):
+      model_state = nnx.state(state.model)
+      opt_state = nnx.state(state.optimizer)
+      state_dict = {
+          "model": nnx.to_pure_dict(model_state),
+          "optimizer": nnx.to_pure_dict(opt_state),
+      }
+    else:
+      state_dict = {
+          "step": state.step,
+          "params": state.params,
+          "opt_state": state.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()
@@ -1289,11 +1345,19 @@ def run_setup():
             # Start snapshot save immediately on the new mesh
             snapshot_mgr = python_vars["snapshot"]
             state = jax_device_state["state"]
-            state_dict = {
-                "step": state.step,
-                "params": state.params,
-                "opt_state": state.opt_state,
-            }
+            if isinstance(state, train_state_nnx.TrainStateNNX):
+              model_state = nnx.state(state.model)
+              opt_state = nnx.state(state.optimizer)
+              state_dict = {
+                  "model": nnx.to_pure_dict(model_state),
+                  "optimizer": nnx.to_pure_dict(opt_state),
+              }
+            else:
+              state_dict = {
+                  "step": state.step,
+                  "params": state.params,
+                  "opt_state": state.opt_state,
+              }
             snapshot_mgr.save_pytree(
                 python_vars["step"], state_dict
             )

From 218564bb7115186e8efdab67580cbb2539f46879 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:41:29 +0000
Subject: [PATCH 15/31] Import train_state_nnx in train.py

---
 src/maxtext/trainers/pre_train/train.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 21838a152b..e4367f5841 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -50,7 +50,7 @@
 
 # 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,

From 88db7236d236606374cfdaca4c7df3c518c5d39d Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:43:47 +0000
Subject: [PATCH 16/31] Support nnx.State in Snapshotter save/load/recover for
 NNX models

---
 src/maxtext/trainers/pre_train/train.py | 86 ++++++++-----------------
 1 file changed, 28 insertions(+), 58 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index e4367f5841..b9864df6a8 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -773,19 +773,16 @@ def training_loop_iteration(
 
   # Async Host Backup (Elastic Mode only)
   if snapshot_mgr is not None and step % config.elastic_snapshot_interval == 0:
-    if isinstance(state, train_state_nnx.TrainStateNNX):
-      model_state = nnx.state(state.model)
-      opt_state = nnx.state(state.optimizer)
+    if isinstance(model, nn.Module):
       state_dict = {
-          "model": nnx.to_pure_dict(model_state),
-          "optimizer": nnx.to_pure_dict(opt_state),
+          "step": state.step,
+          "params": state.params,
+          "opt_state": state.opt_state,
           "metrics": metrics,
       }
     else:
       state_dict = {
-          "step": state.step,
-          "params": state.params,
-          "opt_state": state.opt_state,
+          "nnx_state": nnx.to_pure_dict(state),
           "metrics": metrics,
       }
     snapshot_mgr.save_pytree(step, state_dict)
@@ -963,27 +960,7 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
   # 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(state, train_state_nnx.TrainStateNNX):
-      model_state = nnx.state(state.model)
-      opt_state = nnx.state(state.optimizer)
-      abstract_dict = {
-          "model": nnx.to_pure_dict(model_state),
-          "optimizer": nnx.to_pure_dict(opt_state),
-      }
-      act_model_state = nnx.state(active_state.model)
-      act_opt_state = nnx.state(active_state.optimizer)
-      active_dict = {
-          "model": nnx.to_pure_dict(act_model_state),
-          "optimizer": nnx.to_pure_dict(act_opt_state),
-      }
-      restored_dict = jax.device_put(
-          active_dict, jax.tree.map(lambda x: x.sharding, abstract_dict)
-      )
-      nnx.update(state.model, restored_dict["model"])
-      nnx.update(state.optimizer, restored_dict["optimizer"])
-      restored_step = int(state.optimizer.step.value)
-      restored_state = state
-    else:
+    if isinstance(model, nn.Module):
       abstract_dict = {
           "step": state.step,
           "params": state.params,
@@ -1003,6 +980,15 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
           opt_state=restored_dict["opt_state"],
       )
       restored_step = int(restored_state.step)
+    else:
+      abstract_dict = {"nnx_state": nnx.to_pure_dict(state)}
+      active_dict = {"nnx_state": nnx.to_pure_dict(active_state)}
+      restored_dict = jax.device_put(
+          active_dict, jax.tree.map(lambda x: x.sharding, abstract_dict)
+      )
+      nnx.update(state, restored_dict["nnx_state"])
+      restored_state = state
+      restored_step = int(restored_dict["nnx_state"]["optimizer"]["step"])
     _logger.info(
         "Resharding complete. Retrying. Slices used: %s",
         elastic_manager.active_slice_indices,
@@ -1011,19 +997,7 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
     if snapshot_mgr.latest is None:
       raise RuntimeError("No snapshots available to restore from. Cannot recover.")
     restored_step = snapshot_mgr.latest.step
-    if isinstance(state, train_state_nnx.TrainStateNNX):
-      model_state = nnx.state(state.model)
-      opt_state = nnx.state(state.optimizer)
-      abstract_dict = {
-          "model": nnx.to_pure_dict(model_state),
-          "optimizer": nnx.to_pure_dict(opt_state),
-      }
-      restored_dict = snapshot_mgr.load_pytree(abstract_dict)
-      nnx.update(state.model, restored_dict["model"])
-      nnx.update(state.optimizer, restored_dict["optimizer"])
-      restored_step = int(state.optimizer.step.value)
-      restored_state = state
-    else:
+    if isinstance(model, nn.Module):
       abstract_dict = {
           "step": state.step,
           "params": state.params,
@@ -1035,6 +1009,12 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
           params=restored_dict["params"],
           opt_state=restored_dict["opt_state"],
       )
+    else:
+      abstract_dict = {"nnx_state": nnx.to_pure_dict(state)}
+      restored_dict = snapshot_mgr.load_pytree(abstract_dict)
+      nnx.update(state, restored_dict["nnx_state"])
+      restored_state = state
+      restored_step = int(restored_dict["nnx_state"]["optimizer"]["step"])
     if metric_logger_instance is not None:
       metric_logger_instance.recover_metrics(restored_dict.get("metrics"))
 
@@ -1259,19 +1239,14 @@ def run_setup():
   if config.elastic_enabled:
     replica_axis_idx = config.mesh_axes.index("data")
     snapshot_mgr = Snapshotter(replica_axis_index=replica_axis_idx)
-    if isinstance(state, train_state_nnx.TrainStateNNX):
-      model_state = nnx.state(state.model)
-      opt_state = nnx.state(state.optimizer)
-      state_dict = {
-          "model": nnx.to_pure_dict(model_state),
-          "optimizer": nnx.to_pure_dict(opt_state),
-      }
-    else:
+    if isinstance(model, nn.Module):
       state_dict = {
           "step": state.step,
           "params": state.params,
           "opt_state": state.opt_state,
       }
+    else:
+      state_dict = {"nnx_state": nnx.to_pure_dict(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()
@@ -1345,19 +1320,14 @@ def run_setup():
             # Start snapshot save immediately on the new mesh
             snapshot_mgr = python_vars["snapshot"]
             state = jax_device_state["state"]
-            if isinstance(state, train_state_nnx.TrainStateNNX):
-              model_state = nnx.state(state.model)
-              opt_state = nnx.state(state.optimizer)
-              state_dict = {
-                  "model": nnx.to_pure_dict(model_state),
-                  "optimizer": nnx.to_pure_dict(opt_state),
-              }
-            else:
+            if isinstance(model, nn.Module):
               state_dict = {
                   "step": state.step,
                   "params": state.params,
                   "opt_state": state.opt_state,
               }
+            else:
+              state_dict = {"nnx_state": nnx.to_pure_dict(state)}
             snapshot_mgr.save_pytree(
                 python_vars["step"], state_dict
             )

From 7e8dd8fbffe393b4700cd0a07409fe4388caacc1 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:45:40 +0000
Subject: [PATCH 17/31] Restore p_train_step execution line in
 training_loop_iteration

---
 src/maxtext/trainers/pre_train/train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index b9864df6a8..0e98be5de1 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -709,6 +709,7 @@ def training_loop_iteration(
       with jax.set_mesh(mesh), nn_partitioning.axis_rules(logical_axis_rules):
         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")

From 0823ced463788ee6973bf6a44855c1a34ceb11f4 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:47:56 +0000
Subject: [PATCH 18/31] Enable elastic recovery for pure_nnx in recover()

---
 src/maxtext/trainers/pre_train/train.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 0e98be5de1..bc3b5fe0b4 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -801,8 +801,6 @@ def recover(
 ):
   """Rebuilds MaxText JAX device state and restores state from host snapshot."""
   config = immutable_data["config"]
-  if config.pure_nnx:
-    raise NotImplementedError("Elastic recovery is not supported for NNX.")
 
   metric_logger_instance = python_vars.get("metric_logger_instance")
   if metric_logger_instance is not None:

From 28ec2f99a0035400cffed4fcd7077a83de99f9a8 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:50:38 +0000
Subject: [PATCH 19/31] Set min_slices=1 for elastic slice-down recovery in
 recover()

---
 src/maxtext/trainers/pre_train/train.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index bc3b5fe0b4..09f9633be1 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -814,9 +814,7 @@ def recover(
     python_vars["snapshot"] = snapshot_mgr
 
   # 1. Find currently active slices (wait if none are active)
-  min_slices = config.elastic_min_slice_count
-  if min_slices == -1:
-    min_slices = config.num_slices
+  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(

From 367104a37b91f4d898c8a8a49455918ddd1435fa Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:51:31 +0000
Subject: [PATCH 20/31] Unpack recorder in recover()

---
 src/maxtext/trainers/pre_train/train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 09f9633be1..b5f205296b 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -806,6 +806,7 @@ def recover(
   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:

From abcee2dd64f4c7d9497d9ae2752802cb2b024114 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:54:36 +0000
Subject: [PATCH 21/31] Fix TrainStateNNX state extraction and update in
 recover()

---
 src/maxtext/trainers/pre_train/train.py | 36 +++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index b5f205296b..8be37badfa 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -958,7 +958,27 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
   # 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):
+    if isinstance(state, train_state_nnx.TrainStateNNX):
+      model_state = nnx.state(state.model)
+      opt_state = nnx.state(state.optimizer)
+      abstract_dict = {
+          "model": nnx.to_pure_dict(model_state),
+          "optimizer": nnx.to_pure_dict(opt_state),
+      }
+      act_model_state = nnx.state(active_state.model)
+      act_opt_state = nnx.state(active_state.optimizer)
+      active_dict = {
+          "model": nnx.to_pure_dict(act_model_state),
+          "optimizer": nnx.to_pure_dict(act_opt_state),
+      }
+      restored_dict = jax.device_put(
+          active_dict, jax.tree.map(lambda x: x.sharding, abstract_dict)
+      )
+      nnx.update(state.model, restored_dict["model"])
+      nnx.update(state.optimizer, restored_dict["optimizer"])
+      restored_state = state
+      restored_step = int(state.optimizer.step.value)
+    elif isinstance(model, nn.Module):
       abstract_dict = {
           "step": state.step,
           "params": state.params,
@@ -995,7 +1015,19 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
     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):
+    if isinstance(state, train_state_nnx.TrainStateNNX):
+      model_state = nnx.state(state.model)
+      opt_state = nnx.state(state.optimizer)
+      abstract_dict = {
+          "model": nnx.to_pure_dict(model_state),
+          "optimizer": nnx.to_pure_dict(opt_state),
+      }
+      restored_dict = snapshot_mgr.load_pytree(abstract_dict)
+      nnx.update(state.model, restored_dict["model"])
+      nnx.update(state.optimizer, restored_dict["optimizer"])
+      restored_state = state
+      restored_step = int(state.optimizer.step.value)
+    elif isinstance(model, nn.Module):
       abstract_dict = {
           "step": state.step,
           "params": state.params,

From 738d44d4e1f7e0397a7c80e7a30002bfd091b280 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:57:13 +0000
Subject: [PATCH 22/31] Convert pinned_host memory kind to device before
 split_by_mesh_axis in snapshot.py

---
 src/maxtext/utils/snapshot.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/maxtext/utils/snapshot.py b/src/maxtext/utils/snapshot.py
index aa3e85ea3c..866fff76af 100644
--- a/src/maxtext/utils/snapshot.py
+++ b/src/maxtext/utils/snapshot.py
@@ -104,6 +104,9 @@ def is_replica_active(arr):
     def get_active_pytree(x):
       if not isinstance(x, jax.Array) or not hasattr(x.sharding, "mesh"):
         return x
+      if hasattr(x.sharding, "memory_kind") and x.sharding.memory_kind == "pinned_host":
+        dev_sharding = x.sharding.with_memory_kind("device")
+        x = jax.device_put(x, dev_sharding)
       mesh_axis_name = x.sharding.mesh.axis_names[self.replica_axis_index]
       all_replicas = split_by_mesh_axis.split_by_mesh_axis(
           x,

From 09ea1b60c3fc7e465ccbe8df35ab06ca14d6aa07 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 07:59:45 +0000
Subject: [PATCH 23/31] Add addressable_shards fallback in snapshot.py for
 host-pinned arrays

---
 src/maxtext/utils/snapshot.py | 52 +++++++++++++++++++++--------------
 1 file changed, 31 insertions(+), 21 deletions(-)

diff --git a/src/maxtext/utils/snapshot.py b/src/maxtext/utils/snapshot.py
index 866fff76af..2f082b09d9 100644
--- a/src/maxtext/utils/snapshot.py
+++ b/src/maxtext/utils/snapshot.py
@@ -104,29 +104,39 @@ def is_replica_active(arr):
     def get_active_pytree(x):
       if not isinstance(x, jax.Array) or not hasattr(x.sharding, "mesh"):
         return x
-      if hasattr(x.sharding, "memory_kind") and x.sharding.memory_kind == "pinned_host":
-        dev_sharding = x.sharding.with_memory_kind("device")
-        x = jax.device_put(x, dev_sharding)
-      mesh_axis_name = x.sharding.mesh.axis_names[self.replica_axis_index]
-      all_replicas = split_by_mesh_axis.split_by_mesh_axis(
-          x,
-          mesh_axis_name,
-      )
-
-      active_replicas = [
-          replica for replica in all_replicas if is_replica_active(replica)
-      ]
-
-      if not active_replicas:
-        raise RuntimeError(
-            "No active replicas found."
+      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,
         )
 
-      reconstructed_state = concatenate_by_mesh_axis.concatenate_by_mesh_axis(
-          active_replicas,
-          mesh_axis_name,
-      )
-      return reconstructed_state
+        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)

From a0958f4f9b3003f82960fbde4cb03d5a680d465b Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 08:02:15 +0000
Subject: [PATCH 24/31] Simplify NNX state extraction and restore in recover()

---
 src/maxtext/trainers/pre_train/train.py | 70 ++++++++++---------------
 1 file changed, 29 insertions(+), 41 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 8be37badfa..d217ab0649 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -958,27 +958,7 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
   # 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(state, train_state_nnx.TrainStateNNX):
-      model_state = nnx.state(state.model)
-      opt_state = nnx.state(state.optimizer)
-      abstract_dict = {
-          "model": nnx.to_pure_dict(model_state),
-          "optimizer": nnx.to_pure_dict(opt_state),
-      }
-      act_model_state = nnx.state(active_state.model)
-      act_opt_state = nnx.state(active_state.optimizer)
-      active_dict = {
-          "model": nnx.to_pure_dict(act_model_state),
-          "optimizer": nnx.to_pure_dict(act_opt_state),
-      }
-      restored_dict = jax.device_put(
-          active_dict, jax.tree.map(lambda x: x.sharding, abstract_dict)
-      )
-      nnx.update(state.model, restored_dict["model"])
-      nnx.update(state.optimizer, restored_dict["optimizer"])
-      restored_state = state
-      restored_step = int(state.optimizer.step.value)
-    elif isinstance(model, nn.Module):
+    if isinstance(model, nn.Module):
       abstract_dict = {
           "step": state.step,
           "params": state.params,
@@ -999,14 +979,27 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
       )
       restored_step = int(restored_state.step)
     else:
-      abstract_dict = {"nnx_state": nnx.to_pure_dict(state)}
-      active_dict = {"nnx_state": nnx.to_pure_dict(active_state)}
+      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)
       )
-      nnx.update(state, restored_dict["nnx_state"])
+      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
-      restored_step = int(restored_dict["nnx_state"]["optimizer"]["step"])
+      restored_step = int(state.optimizer.step.value) if hasattr(state, "optimizer") and hasattr(state.optimizer, "step") else 0
     _logger.info(
         "Resharding complete. Retrying. Slices used: %s",
         elastic_manager.active_slice_indices,
@@ -1015,19 +1008,7 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
     if snapshot_mgr.latest is None:
       raise RuntimeError("No snapshots available to restore from. Cannot recover.")
     restored_step = snapshot_mgr.latest.step
-    if isinstance(state, train_state_nnx.TrainStateNNX):
-      model_state = nnx.state(state.model)
-      opt_state = nnx.state(state.optimizer)
-      abstract_dict = {
-          "model": nnx.to_pure_dict(model_state),
-          "optimizer": nnx.to_pure_dict(opt_state),
-      }
-      restored_dict = snapshot_mgr.load_pytree(abstract_dict)
-      nnx.update(state.model, restored_dict["model"])
-      nnx.update(state.optimizer, restored_dict["optimizer"])
-      restored_state = state
-      restored_step = int(state.optimizer.step.value)
-    elif isinstance(model, nn.Module):
+    if isinstance(model, nn.Module):
       abstract_dict = {
           "step": state.step,
           "params": state.params,
@@ -1040,11 +1021,18 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
           opt_state=restored_dict["opt_state"],
       )
     else:
-      abstract_dict = {"nnx_state": nnx.to_pure_dict(state)}
+      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)
-      nnx.update(state, restored_dict["nnx_state"])
+      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
-      restored_step = int(restored_dict["nnx_state"]["optimizer"]["step"])
+      restored_step = int(state.optimizer.step.value) if hasattr(state, "optimizer") and hasattr(state.optimizer, "step") else restored_step
     if metric_logger_instance is not None:
       metric_logger_instance.recover_metrics(restored_dict.get("metrics"))
 

From 4c484d0bf07ebc7318b71e136dd816e268788297 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 08:04:51 +0000
Subject: [PATCH 25/31] Fix all snapshot saving places to use uniform model and
 optimizer PyTree structure for NNX

---
 src/maxtext/trainers/pre_train/train.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index d217ab0649..331b5b31f3 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -782,10 +782,11 @@ def training_loop_iteration(
           "metrics": metrics,
       }
     else:
-      state_dict = {
-          "nnx_state": nnx.to_pure_dict(state),
-          "metrics": metrics,
-      }
+      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
@@ -1264,7 +1265,11 @@ def run_setup():
           "opt_state": state.opt_state,
       }
     else:
-      state_dict = {"nnx_state": nnx.to_pure_dict(state)}
+      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()
@@ -1345,7 +1350,11 @@ def run_setup():
                   "opt_state": state.opt_state,
               }
             else:
-              state_dict = {"nnx_state": nnx.to_pure_dict(state)}
+              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
             )

From ce7bd0d5aaf4b5bbcf2684cd23c08257301c1ae6 Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 08:07:05 +0000
Subject: [PATCH 26/31] Split TrainStateNNX to nnx.State in recover() to match
 p_train_step in_shardings

---
 src/maxtext/trainers/pre_train/train.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 331b5b31f3..89db8dfa5a 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -1038,6 +1038,8 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
       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

From a00e2fd1e029e2e1fde943b3a7ec0391747784ff Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 08:11:40 +0000
Subject: [PATCH 27/31] Use safe leaf-by-leaf block_until_ready in load_pytree
 in snapshot.py

---
 src/maxtext/utils/snapshot.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/src/maxtext/utils/snapshot.py b/src/maxtext/utils/snapshot.py
index 2f082b09d9..9c15affbe4 100644
--- a/src/maxtext/utils/snapshot.py
+++ b/src/maxtext/utils/snapshot.py
@@ -171,7 +171,15 @@ def get_active_pytree(x):
         host_target_state,
         target_device_shardings,
     )
-    jax.block_until_ready(restored_state)
+    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

From c3bc9afd6fcf5a64588a767f0286d8d22672dddd Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 08:18:41 +0000
Subject: [PATCH 28/31] Pass GraphDef jit_model to jit_train_and_eval_step in
 recover() for NNX

---
 src/maxtext/trainers/pre_train/train.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index 89db8dfa5a..e4025b3203 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -943,10 +943,15 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
   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,
-        model,
+        jit_model,
         mesh,
         state,
         state_mesh_shardings,

From 97713adff02be9f3c70e17ed63f7f36c2e16a78d Mon Sep 17 00:00:00 2001
From: Luke Baumann <lukebaumann@google.com>
Date: Thu, 2 Jul 2026 08:21:57 +0000
Subject: [PATCH 29/31] Safely convert step array to int in recover() to
 prevent DATA_LOSS exception

---
 src/maxtext/trainers/pre_train/train.py | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py
index e4025b3203..648fedc317 100644
--- a/src/maxtext/trainers/pre_train/train.py
+++ b/src/maxtext/trainers/pre_train/train.py
@@ -1005,7 +1005,12 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
       if hasattr(state, "optimizer") and "optimizer" in restored_dict:
         nnx.update(state.optimizer, restored_dict["optimizer"])
       restored_state = state
-      restored_step = int(state.optimizer.step.value) if hasattr(state, "optimizer") and hasattr(state.optimizer, "step") else 0
+      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,
@@ -1038,7 +1043,12 @@ def calc_gbs(per_device_batch_size, expansion_factor, num_devices, grad_accum_st
       if hasattr(state, "optimizer") and "optimizer" in restored_dict:
         nnx.update(state.optimizer, restored_dict["optimizer"])
       restored_state = state
-      restored_step = int(state.optimizer.step.value) if hasattr(state, "optimizer") and hasattr(state.optimizer, "step") else restored_step
+      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"))
 

From 0c6ee1659f1574045e029e727176f4313b9eebe2 Mon Sep 17 00:00:00 2001
From: Abhinav Singh <abhinavsing@google.com>
Date: Sat, 11 Jul 2026 18:54:20 +0000
Subject: [PATCH 30/31] Fix Data Loss issue in DataLoader.

---
 src/maxtext/common/data_loader.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py
index 21bd870bc8..0cfa3069e3 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.src.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:

From 551a6882ed5541850acd20c01ba996ddbeef6284 Mon Sep 17 00:00:00 2001
From: Abhinav Singh <abhinavsing@google.com>
Date: Sat, 11 Jul 2026 19:15:05 +0000
Subject: [PATCH 31/31] Fix maxtext import issue.

---
 src/maxtext/common/data_loader.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py
index 0cfa3069e3..e31175ba14 100644
--- a/src/maxtext/common/data_loader.py
+++ b/src/maxtext/common/data_loader.py
@@ -26,7 +26,7 @@
 from maxtext.trainers.diloco import diloco
 from maxtext.utils import exceptions
 from maxtext.utils.sharding import get_input_data_sharding
-from maxtext.src.maxtext.utils import elastic_utils
+from maxtext.utils import elastic_utils
 
 
 class DataLoader: