@@ -1812,6 +1812,11 @@ def update(self) -> int:
       lr = self._learning_rate_schedule(self.train_step)
       self.record_metrics("learning_rate", lr)
 
+    # Materialize the step's recorded metrics before the optimizer step: a value that cannot be
+    # converted then raises here, with nothing applied, instead of after the update has landed
+    # and before the step counter and accumulator are reset (a retry would apply it twice).
+    self._metrics_recorder.get_step_metrics(self.train_step)
+
     # Wait for previous computations to finish before dispatching the update step to TPU.
     self._throttler.wait_for_next()
 
@@ -2113,14 +2118,18 @@ def record_metrics(
       aggregation_fn: Callable[[jax.Array], Any] | None = None,
       mode: metrics_module.Mode = metrics_module.Mode.TRAIN,
   ) -> None:
-    """Records a metric into the buffer, appending to JAX arrays.
+    """Records a metric into the step's buffer (host-side; concatenated once at flush time).
 
     Args:
       name: The name of the metric.
       metric: The metric to record.
       aggregation_fn: The aggregation function to apply to the metric.
     """
-    if metric is None:
+    # `loss_fn` auxiliary outputs can contain `None` entries (disabled optional losses),
+    # `"intermediate_outputs"` (sow activations/mutable collections), or tuples/lists
+    # (e.g., `(total_loss, aux)`). Skip these so `buffer_metrics` only receives scalar
+    # metrics, arrays, or nested metric dicts.
+    if metric is None or name == "intermediate_outputs" or isinstance(metric, (tuple, list)):
       return
     if isinstance(metric, dict):
       for sub_k, sub_v in metric.items():
 
@@ -17,13 +17,13 @@
 from __future__ import annotations
 
 from collections.abc import Callable
-import dataclasses
 import enum
 import os
 from typing import Any
 
 from absl import logging
 import jax
+import jax.numpy as jnp
 from maxtext.configs import pyconfig
 from maxtext.training_engine import abstract_engine
 from maxtext.utils import max_logging, max_utils, maxtext_utils
@@ -42,11 +42,26 @@
 ]
 
 # Completed step buffers to keep resident. Each holds live device arrays and nothing on the
-# engine's step path removes one, so unbounded history grows HBM, and checkpoint size with
-# it. A window rather than Tunix's single prior step, so batched readers still work.
+# engine's step path removes one, so unbounded history grows HBM. A window rather than Tunix's
+# single prior step, so batched readers still work.
 _DEFAULT_MAX_BUFFERED_STEPS = 128
 
 
+def _concat_1d(items: list[Any]) -> jax.Array:
+  """Joins the items in order, with the layout the eager `jnp.append` path produced.
+
+  A single item keeps its shape (`atleast_1d`, as that path stored a step's first item); two
+  or more are raveled and concatenated. A list of scalars, one per microbatch and the usual
+  case, is a single `jnp.stack`: a ravel per item plus the concatenate would be one eager
+  launch per microbatch per metric.
+  """
+  if len(items) == 1:
+    return jnp.atleast_1d(jnp.asarray(items[0]))
+  if all(jnp.ndim(x) == 0 for x in items):
+    return jnp.stack(items)
+  return jnp.concatenate([jnp.ravel(jnp.asarray(x)) for x in items])
+
+
 class Mode(str, enum.Enum):
   TRAIN = "train"
   EVAL = "eval"
@@ -75,6 +90,13 @@ def __init__(self, max_buffered_steps: int = _DEFAULT_MAX_BUFFERED_STEPS, mode:
         as new steps start. Zero or less retains everything.
     """
     self._metrics_buffer: list[abstract_engine.MetricsBuffer] = []
+    # Items recorded but not yet materialized, per live buffer. Keyed by the buffer object's
+    # identity rather than its step id: `buffer_metrics` only compares against the newest
+    # buffer, so the history may hold two buffers with the same id (the eager `jnp.append`
+    # code kept them apart), and they must not share a list. Every place a buffer leaves
+    # `_metrics_buffer` drops its key, so an identity is never reused while an entry exists.
+    self._pending_weighted: dict[int, dict[str, list[abstract_engine.WeightedMetric]]] = {}
+    self._pending_scalar: dict[int, dict[str, list[Any]]] = {}
     self._mode = mode
     self._max_buffered_steps = max_buffered_steps
     self._dropped_buffer_count = 0
@@ -83,17 +105,21 @@ def buffer_metrics(
       self,
       train_step: int,
       name: str,
-      metric: jax.Array | float | int | None = None,
+      metric: abstract_engine.WeightedMetric | jax.Array | float | int | None = None,
       aggregation_fn: Callable[[jax.Array], Any] | None = None,
-  ):
-    """Buffers metrics for the given train step.
+  ) -> None:
+    """Buffers a metric for the given train step (host-side; no device work until flush).
 
     Args:
       train_step: The train step for which to buffer metrics.
       name: The name of the metric.
-      metric: The metric to buffer.
+      metric: The metric to buffer. `None` is a no-op (nothing is recorded and no
+        buffer is opened for the step).
       aggregation_fn: Optional aggregation function to apply to the metric.
     """
+    if metric is None:
+      return
+
     if not self._metrics_buffer or self._metrics_buffer[-1].id != train_step:
       new_buffer = abstract_engine.MetricsBuffer(id=train_step, mode=self._mode)
       self._metrics_buffer.append(new_buffer)
@@ -111,6 +137,9 @@ def _evict_old_buffers(self) -> None:
       return
     num_dropped = len(self._metrics_buffer) - self._max_buffered_steps
     oldest_dropped_id = self._metrics_buffer[0].id
+    for evicted in self._metrics_buffer[:num_dropped]:
+      self._pending_weighted.pop(id(evicted), None)
+      self._pending_scalar.pop(id(evicted), None)
     del self._metrics_buffer[:num_dropped]
     self._dropped_buffer_count += num_dropped
     logging.log_every_n(
@@ -130,37 +159,64 @@ def _record_metric(
       metric: abstract_engine.WeightedMetric | jax.Array | float | int,
       aggregation_fn: Callable[[jax.Array], Any] | None = None,
   ) -> None:
-    """Records a metric into the buffer, appending to JAX arrays.
+    """Records a metric into the step's pending Python list without launching XLA kernels.
 
     Args:
       name: The name of the metric.
       metric: The metric to record.
       aggregation_fn: The aggregation function to apply to the metric.
     """
     buffer = self._metrics_buffer[-1]
+    key = id(buffer)
 
     if aggregation_fn is not None:
       buffer.aggregation_fns[name] = aggregation_fn
 
     if isinstance(metric, abstract_engine.WeightedMetric):
-      if name not in buffer.weighted_metrics:
-        buffer.weighted_metrics[name] = abstract_engine.WeightedMetric(
-            unreduced_sum=jax.numpy.atleast_1d(metric.unreduced_sum),
-            denominator=jax.numpy.atleast_1d(metric.denominator),
-            eps=metric.eps,
-            min_denom=metric.min_denom,
-        )
-      else:
-        current_val = buffer.weighted_metrics[name]
-        new_sum = jax.numpy.append(current_val.unreduced_sum, metric.unreduced_sum)
-        new_denom = jax.numpy.append(current_val.denominator, metric.denominator)
-        buffer.weighted_metrics[name] = dataclasses.replace(current_val, unreduced_sum=new_sum, denominator=new_denom)
+      self._pending_weighted.setdefault(key, {}).setdefault(name, []).append(metric)
     else:
-      if name not in buffer.scalar_metrics:
-        buffer.scalar_metrics[name] = jax.numpy.atleast_1d(jax.numpy.asarray(metric))
-      else:
-        buffer.scalar_metrics[name] = jax.numpy.append(buffer.scalar_metrics[name], metric)
-    self._metrics_buffer[-1] = buffer
+      self._pending_scalar.setdefault(key, {}).setdefault(name, []).append(metric)
+
+  def _flush_buffer(self, buffer: abstract_engine.MetricsBuffer) -> None:
+    """Materializes pending Python lists into 1D JAX arrays on `buffer`.
+
+    All-or-nothing: every pending item is converted first, and `buffer` and the pending
+    dicts are only updated once all conversions succeed. An item that `jnp.asarray` rejects
+    therefore raises without dropping the step's other pending metrics, and a retry sees the
+    same pending state.
+
+    Args:
+      buffer: A live entry of `_metrics_buffer` whose pending items should be materialized.
+    """
+    key = id(buffer)
+    pending_w = self._pending_weighted.get(key)
+    pending_s = self._pending_scalar.get(key)
+    if not pending_w and not pending_s:
+      return
+
+    new_weighted: dict[str, abstract_engine.WeightedMetric] = {}
+    for name, items in (pending_w or {}).items():
+      existing = buffer.weighted_metrics.get(name)
+      prefix = [existing] if existing is not None else []
+      # The live items carry the metric's eps/min_denom.
+      template = items[0]
+      new_weighted[name] = abstract_engine.WeightedMetric(
+          unreduced_sum=_concat_1d([m.unreduced_sum for m in prefix + items]),
+          denominator=_concat_1d([m.denominator for m in prefix + items]),
+          eps=template.eps,
+          min_denom=template.min_denom,
+      )
+
+    new_scalar: dict[str, jax.Array] = {}
+    for name, items in (pending_s or {}).items():
+      prefix = [buffer.scalar_metrics[name]] if name in buffer.scalar_metrics else []
+      new_scalar[name] = _concat_1d(prefix + items)
+
+    # Commit point: nothing above mutated `buffer` or the pending dicts.
+    buffer.weighted_metrics.update(new_weighted)
+    buffer.scalar_metrics.update(new_scalar)
+    self._pending_weighted.pop(key, None)
+    self._pending_scalar.pop(key, None)
 
   def get_metrics_history(self, clear_cache: bool = True) -> list[abstract_engine.MetricsBuffer]:
     """Returns every cached step buffer and optionally clears the metrics cache.
@@ -175,24 +231,39 @@ def get_metrics_history(self, clear_cache: bool = True) -> list[abstract_engine.
     Returns:
       One on-device MetricsBuffer per retained train step, oldest first.
     """
+    for buf in self._metrics_buffer:
+      self._flush_buffer(buf)
     metrics_to_return = self._metrics_buffer
     if clear_cache:
       # Clear the metrics buffer.
       self._metrics_buffer = []
+      self._pending_weighted.clear()
+      self._pending_scalar.clear()
     return metrics_to_return
 
   def get_step_metrics(self, step: int) -> abstract_engine.MetricsBuffer | None:
-    """Returns the latest metrics from the metrics buffer."""
+    """Returns the latest flushed metrics buffer if its ID matches `step`.
+
+    Args:
+      step: The train step ID to retrieve.
+
+    Returns:
+      The flushed `MetricsBuffer` for `step`, or `None` if the latest buffer
+      does not match `step`.
+    """
     if not self._metrics_buffer:
       return None
     latest = self._metrics_buffer[-1]
     if latest.id == step:
+      self._flush_buffer(latest)
       return latest
     return None
 
   def cleanup(self) -> None:
     """Cleans up the metrics recorder."""
     self._metrics_buffer = []
+    self._pending_weighted.clear()
+    self._pending_scalar.clear()
 
 
 class MetricsLogger: