training_engine: add stage_batch with host-buffered multi-batch queue (TUNIX_STAGE_MAX_DEVICE_BATCHES)

[A9isha]

Description

Adds MaxTextTrainingEngine.stage_batch(payload, key=...), TUNIX_STAGE_BATCH (off | stage | serialize), and a host-buffered multi-batch queue (TUNIX_STAGE_MAX_DEVICE_BATCHES, default 1, 0 = unlimited) so google/tunix#2906 can pre-stage up to 2 micro-batches ahead of fwd_bwd while capping device-resident staged micro-batches at +1 sharded input set on TPU HBM.

Key Changes

• Environment Configuration (TUNIX_STAGE_BATCH, TUNIX_STAGE_MAX_DEVICE_BATCHES):
  • Configured directly via stage_batch_mode() (TUNIX_STAGE_BATCH=off|stage|serialize) and stage_max_device_batches() (TUNIX_STAGE_MAX_DEVICE_BATCHES, default 1).
• Split FIFO vs. Transfer Locks (_stage_lock vs. _transfer_lock):
  • _stage_lock (paired with _stage_cond) guards only microsecond FIFO deque and counter mutations and is never held across jax.device_put or jax.block_until_ready.
  • _transfer_lock serializes _wait_for_input_transfer() + _issue_input_transfer() so H2D input transfers never overlap, while fwd_bwd can pop an already-on-device head immediately without blocking behind an in-progress transfer for the next micro-batch.
• Host-Buffered 2nd Staged Batch & Instrumented Promotion (TUNIX_STAGE_MAX_DEVICE_BATCHES=1):
  • When TUNIX_STAGE_MAX_DEVICE_BATCHES is reached (counting both queued on_device=True entries and the currently executing _in_flight_staged_on_device batch), additional staged batches are kept in trainer host RAM as np.ndarray leaves (on_device=False, zero extra TPU HBM).
  • As soon as fwd_bwd finishes dispatching a staged batch (wrapped in try ... finally), _promote_staged_head_to_device() promotes the next queued host-buffered batch to TPU HBM under _transfer_lock and logs STAGE_PROMOTE key=%s host_bytes=%d wait_prev_s=%.4f issue_s=%.4f queued=%d on_device=%s.
• Key Matching & Idempotency:
  • stage_batch is idempotent by key (replaces an existing entry with the same key in-place rather than appending a duplicate on RPC retry).
  • Keyless fwd_bwd(payload) calls (stage_key=None) never consume or drop keyed FIFO entries.
  • Added reset_staged() to clear queued staged entries and reset in-flight counters on error/reset paths.
• Interaction with AI-Hypercomputer/maxtext#5718:
  • Self-contained with respect to #5718. Both PRs touch the fwd_bwd signature and input-preparation preamble; whichever lands second rebases fwd_bwd so a PreparedFwdBwdInput payload bypasses _take_staged_batch and the explicit transfer, while None or a raw dict payload goes through _take_staged_batch.

Matched with google/tunix#2906.

Tests

Ran the unit test suite (13 passed):

pytest tests/post_training/unit/maxtext_engine_stage_batch_test.py

Covers mode/env parsing, bit-identical losses & gradients against off mode, fwd_bwd(None, stage_key=...), out-of-order key lookup, signature mismatch fallback, serialized transfers, host-buffered np.ndarray leaves and STAGE_PROMOTE device promotion, keyless fwd_bwd(payload) coexistence with keyed staged batches, stage_batch key idempotency, exception recovery / reset_staged(), and multi-threaded ThreadPoolExecutor concurrency. (End-to-end 1-deep stage_batch was measured on siv-sk8; full 2-deep e2e on atwigg/mlperf is pending paired deployment with google/tunix#2906.)

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have added unit tests covering the multi-batch stage_batch queue and host-to-device promotion (2-deep e2e run on atwigg/mlperf pending paired deployment with google/tunix#2906).
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request introduces micro-batch staging (TUNIX_STAGE_BATCH) in MaxTextTrainingEngine to transfer micro-batch inputs onto devices ahead of their forward-backward passes, along with corresponding unit tests. Feedback on the changes highlights a critical runtime error where self._throttler.block_until_ready() is called but does not exist, suggesting wait_for_all() instead. Additionally, to prevent head-of-line blocking when a staged batch is dropped due to a signature mismatch, it is recommended to defer the promotion of the next staged batch to the end of fwd_bwd rather than executing it immediately.

[gemini-code-assist]

The InflightThrottler class does not define a block_until_ready method. Calling self._throttler.block_until_ready() will raise an AttributeError at runtime, crashing the training process. Use wait_for_all() instead to block until all inflight computations have finished.

Suggested change

	self._throttler.block_until_ready()
	self._throttler.wait_for_all()

[gemini-code-assist]

Calling _promote_staged_head_to_device() immediately when a staged batch is dropped due to a signature mismatch causes head-of-line blocking. The promotion issues a transfer for the next staged batch, and the subsequent explicit transfer of the current payload must block and wait for this future batch's transfer to complete. Deferring the promotion to the end of fwd_bwd avoids blocking the critical path of the current step.

Suggested change

	if signature is not None and self._needs_recompile(signature, staged.signature):
	self._staged_batches.popleft()
	self._stage_warn_once(
	"signature",
	"Staged batch %r does not match the fwd_bwd payload's structure; dropped, the payload is transferred "
	"explicitly. Logged once.",
	staged.key,
	)
	self._promote_staged_head_to_device()
	return None
	if signature is not None and self._needs_recompile(signature, staged.signature):
	self._staged_batches.popleft()
	self._stage_warn_once(
	"signature",
	"Staged batch %r does not match the fwd_bwd payload's structure; dropped, the payload is transferred "
	"explicitly. Logged once.",
	staged.key,
	)
	return None

[gemini-code-assist]

To support deferring promotion of the next staged batch (preventing head-of-line blocking when a batch is dropped or when staging is bypassed), always call _promote_staged_head_to_device() at the end of fwd_bwd regardless of whether staged is None.

    with self._stage_lock:
      if staged is not None and staged.on_device:
        self._in_flight_staged_on_device = max(0, self._in_flight_staged_on_device - 1)
      self._promote_staged_head_to_device()