feat(training_engine): expose step-timing readiness tokens for Tunix TRAINER_STEP_TIMING

src/maxtext/training_engine/maxtext_engine.py

@@ -981,6 +981,10 @@ def __init__(
     # resumed into completes and its finished state has been checkpointed.
     self._resumed_mid_step = False
     self._cached_losses: list[abstract_engine.WeightedMetric | jax.Array] = []
+    # Readiness tokens for Tunix's TrainerWorker step timer (TRAINER_STEP_TIMING): arrays that
+    # become ready when the last microbatch's / optimizer step's device work has finished.
+    self.last_fwd_bwd_token: jax.Array | None = None
+    self.last_update_token: jax.Array | None = None
     # `create_training_optimizer` returns a raw optax GradientTransformation. `TrainStateNNX.apply_gradients`
     # calls `optimizer.update(model, grads)`, which is the nnx.Optimizer signature, and
     # `checkpointing.CheckpointState` expects an nnx.Optimizer too, so wrap it here. `wrt=nnx.Param`
@@ -2687,6 +2691,9 @@ def fwd_bwd(self, payload: abstract_engine.TrainerPayload | PreparedFwdBwdInput,
     # the update step.
     self._throttler.add_computation(computation=loss, metrics=None)
 
+    # Readiness token for Tunix's step timer (TRAINER_STEP_TIMING): the loss rather than the running sums, which the
+    # next micro-batch's accumulate donates. A dropless rerun dispatched later replaces it (`_resolve_overflow_check`).
+    self.last_fwd_bwd_token = loss.unreduced_sum if isinstance(loss, abstract_engine.WeightedMetric) else loss
     self._accumulated_grads = acc_grads
     self._accumulated_denominator = acc_denom
     if self._dropless_replay:
@@ -2821,6 +2828,8 @@ def _resolve_overflow_check(self) -> None:
             self._accumulated_grads, self._accumulated_denominator, grads, denominator
         )
       self._throttler.add_computation(computation=loss, metrics=None)
+      # The rerun is the last device work dispatched so far, so it carries the readiness token from here on.
+      self.last_fwd_bwd_token = loss.unreduced_sum if isinstance(loss, abstract_engine.WeightedMetric) else loss
       if bool(rerun_overflow):
         raise RuntimeError(
             f"The dropless rerun of micro step {check.micro_step} still overflowed; check apply_dropless_overrides."
@@ -2983,6 +2992,8 @@ def update(self, **kwargs: Any) -> int:
       # Back to pinned host memory, where the compiled update's `out_shardings` would put it.
       self._offload_optimizer_state()
 
+    # Same readiness signal the throttler uses below: the norm comes out of the update executable.
+    self.last_update_token = grad_norm
     if grad_norm is not None:
       self.record_metrics("gradient_norm", grad_norm)
     if self._log_moe_drops:

tests/post_training/unit/maxtext_engine_test.py

@@ -2323,6 +2323,29 @@ def _loss_fn(model, *args, **kwargs):
     # This form carries no aux, so nothing beyond the loss is recorded.
     self.assertEmpty(metrics.scalar_metrics)
 
+  def test_readiness_tokens_track_last_microbatch_and_update(self):
+    """`last_fwd_bwd_token` / `last_update_token` are the arrays Tunix's step timer waits on."""
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    self.assertIsNone(t.last_fwd_bwd_token)
+    self.assertIsNone(t.last_update_token)
+    t.with_loss_fn(
+        lambda model, *args, **kwargs: (
+            abstract_engine.WeightedMetric(unreduced_sum=jnp.sum(model.weights[...]), denominator=jnp.array(1.0)),
+            {},
+        )
+    )
+    payload = DummyPayload()
+    t.fwd_bwd(payload)
+    first = t.last_fwd_bwd_token
+    t.fwd_bwd(payload)
+    # A plain array (the WeightedMetric's sum), replaced by each microbatch.
+    self.assertIsInstance(t.last_fwd_bwd_token, jax.Array)
+    self.assertIsNot(t.last_fwd_bwd_token, first)
+    self.assertIs(t.last_fwd_bwd_token, t._cached_losses[-1].unreduced_sum)
+    t.update()
+    self.assertIsInstance(t.last_update_token, jax.Array)
+    jax.block_until_ready((t.last_fwd_bwd_token, t.last_update_token))
+
   def test_fwd_bwd_with_tunix_spelled_loss_output(self):
     """A loss written against Tunix's API behaves identically to the MaxText spelling.