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

[ultrons]

Exposes the two readiness tokens that Tunix's trainer step timer (TRAINER_STEP_TIMING, google/tunix#2713) waits on, so the timeline reports when each microbatch and the optimizer step finished on device.

• last_fwd_bwd_token: the microbatch's loss array (the WeightedMetric's unreduced_sum when the loss is one), set in fwd_bwd. The loss rather than the running gradient sums, because the next microbatch's accumulate donates those. When an overflowed microbatch is rerun dropless (_resolve_overflow_check), the rerun's loss replaces the token, since it is the last device work dispatched.
• last_update_token: the gradient norm, set in update. It is the same signal the in-flight throttler already waits on, since it comes out of the update executable.

Both start as None. The timer blocks on them from a daemon thread, so the serving path never waits.

Why: with the MaxText backend (what the MLPerf recipes run) neither attribute existed, so the worker's step timer had no device-readiness signal and every ready field in TRAINER_STEP_TIMING stayed None. google/tunix#2713 now warns once when a trainer lacks them; this makes the warning go away for MaxText.

Rebased onto atwigg/mlperf e526ea0 (after #5718); also applies on #5721

We rebased the commit onto atwigg/mlperf at e526ea0, which includes #5718's prepare_fwd_bwd_input / wait_for_device. The fwd_bwd token now sits after the throttler hand-off in the restructured method, and the dropless rerun path sets it as well (2 files, +34 lines). The same commit cherry-picks cleanly onto the #5721 head (14ed039) and onto e938007, the #5721 state that google/tunix#2884 pins.

Test receipt (CPU, trellis-experimental@sha256:d6cc06ff... image, with the checkout's src/ placed ahead of the image's installed maxtext; a probe test in the same session confirmed maxtext.training_engine.maxtext_engine.__file__ resolved to the mounted src/):

• tests/post_training/unit/maxtext_engine_test.py: 112 passed, 1 failed with this change; 111 passed, 1 failed on plain e526ea0. The extra pass is the new test_readiness_tokens_track_last_microbatch_and_update. The one failure, test_maybe_register_pathways_persistence, fails identically on both (Orbax Pathways handler registration on a CPU-only host).
• feat(training_engine): add prepare_fwd_bwd_input and wait_for_device to overlap the next step's input copy (-7.2% s/step, 397B RL trainer bench, 128 v7x) #5718's maxtext_engine_prepared_input_test.py and maxtext_engine_stage_before_update_test.py: 14 passed with this change, 14 passed on plain e526ea0.

Testing (original revision)

• New test_readiness_tokens_track_last_microbatch_and_update.
• tests/post_training/unit/maxtext_engine_test.py on CPU in the atwigg/trellis-experimental:1005 image: 91 passed, 1 failed. The failure, test_maybe_register_pathways_persistence, fails identically without this change (Orbax Pathways handler registration on a CPU-only host).

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable (unit test receipts above; the tokens were exercised end to end by the TRAINER_STEP_TIMING runs described in Log the trainer's device-side step timeline (TRAINER_STEP_TIMING) google/tunix#2713).
• I have made or will make corresponding changes to the doc if needed (no doc page changes: the two attributes are documented in the engine docstrings).

[gemini-code-assist]

Code Review

This pull request introduces readiness tokens (last_fwd_bwd_token and last_update_token) to the MaxTextTrainingEngine to support Tunix's TrainerWorker step timer. These tokens are updated during the forward-backward and optimizer update steps respectively. A corresponding unit test has been added to verify their behavior. There are no review comments, so I have no feedback to provide.