[A9isha]

Summary

Ports intra-step stage_batch onto atwigg/mlperf and extends it to continuous 2-deep intra-step stage_batch pipelining (_stage_ahead_while_pending, max_staged_batches=2) across the orchestrator, RPC server, TrainerWorker, and K8s launch recipes (matched with AI-Hypercomputer/maxtext#5751 and google/tunix#2905).

Value Breakdown (Measured & Estimated from siv-sk8)

1. Porting intra-step stage_batch onto atwigg/mlperf (~20–35 s / step estimated savings vs. inline train_step(payload)):
   • atwigg/mlperf currently has no intra-step staging, putting each micro-batch's wire transfer (8–10 s) plus H2D copy (~3.9 s for 5.03 GB) on the critical path (estimated ~2–4 s x 9 ≈ 20–35 s/step vs. inline train_step(payload)).
   • Measured on siv-sk8 Step 1: when every stage completed in 7.8–10.0 s < 10.35 s, stage_batch allowed the TPU to execute MB 1..8 back-to-back with 0.00 s mid-step idle (ready(k) - ready(k-1) = 10.35 s for every k).
2. 2-Deep over 1-Deep (_stage_ahead_while_pending, max_staged_batches=2):
   • In siv-sk8 Step 2, the raw 5.03 GB stages for MB 1..8 summed to 94.9 s against 82.8 s (8 x 10.35 s) of compute windows — a 12.1 s bandwidth deficit that requires google/tunix#2905 (wire compression) to remove. Standalone on Step 2's front-loaded stall pattern, 2-deep over 1-deep is estimated to recover up to ~5.6 s.
   • Combined with google/tunix#2905 (which shrinks stage transfers to ~4–5 s, leaving ~5–6 s of spare wire time per 10.35 s compute window) and step-boundary pipelining (#2878 + #5718, which dispatches MB 0 ahead so _stage_ahead_while_pending builds a 2-deep cushion from the start of the step), 2-deep staging absorbs 2x transient gRPC contention spikes without stalling TPU dispatch.

Solution (atwigg/mlperf base)

1. RemoteExecutionServer Background Method Executor (tunix/experimental/worker/remote_execution.py):
   • Adds a lazy ThreadPoolExecutor (_method_executor, max_workers=max(2, len(remote_executor_methods)), cleaned up via close_executor()) for methods listed in an instance's remote_executor_methods.
   • Under TUNIX_STAGE_BATCH=stage, TrainerWorker.remote_executor_methods includes {"stage_batch", "fwd_bwd"} so stage_batch(MB k+1) can run its host preparation and H2D transfer on one pool thread while fwd_bwd(MB k) blocks in InflightThrottler.wait_for_next() on the other pool thread.
2. TrainerWorker.stage_batch & reset_staged (tunix/experimental/worker/trainer_worker.py):
   • Materializes any CompactRoutedExperts leaf via _with_dense_routing before forwarding (payload, key=request.request_id) to self._trainer.stage_batch (ensuring seamless composition with google/tunix#2905).
   • Checks hasattr(self._trainer, "stage_batch") up front and degrades gracefully on trainer exceptions (staged=False, reason=...) while keeping the worker in WorkerState.READY so the orchestrator falls back to inline train_step(payload).
3. DistributedRLEngine.stage_train_batch & reset_staged_batches (tunix/experimental/orchestrator/distributed_rl_engine.py):
   • Sends stage_batch to the actor worker and returns stage_key when staged=True, catching (AttributeError, RuntimeError) to fall back cleanly to inline train_step(payload) if staging is unsupported or fails.
4. Continuous 2-Deep Intra-Step Staging & Fallback Metric in StandardRLProgram._train_loop (tunix/experimental/orchestrator/rl_program.py, run_deepswe_dist.py):
   • While self._pending_train (MB k-1) is in flight on the TPU, _stage_ahead_while_pending() keeps up to max_staged_batches=2 micro-batches staged ahead (one active_stage_task RPC in flight at a time from the orchestrator).
   • Once MB k is in staged_queue[0], both background group pulls (active_pull_task with count_exposed=False) and active_stage_task (MB k+1) race against self._pending_train via asyncio.wait(..., return_when=asyncio.FIRST_COMPLETED) so MB k dispatches immediately as soon as MB k-1 finishes.
   • Tracks per-step staging fallbacks via self._stage_fallbacks and logs stage/fallbacks in _collect_and_log_step_metrics whenever stage_train_batch raises or returns False.
   • stage_train_batches defaults to TUNIX_STAGE_BATCH=stage when pipeline_train_microbatches=True and is plumbed through run_deepswe_dist.py (--stage_train_batches, --max_staged_batches).
5. Recipe & K8s Launcher Plumbing (mlperf_base.sh, k8s_launcher.sh):
   • Defines TUNIX_STAGE_BATCH="${TUNIX_STAGE_BATCH:-off}", TUNIX_STAGE_MAX_DEVICE_BATCHES="${TUNIX_STAGE_MAX_DEVICE_BATCHES:-1}", and TUNIX_MAX_STAGED_BATCHES="${TUNIX_MAX_STAGED_BATCHES:-2}" in mlperf_base.sh and k8s_launcher.sh, forwarding TUNIX_STAGE_BATCH and TUNIX_MAX_STAGED_BATCHES to the orchestrator container and TUNIX_STAGE_BATCH and TUNIX_STAGE_MAX_DEVICE_BATCHES to the trainer container.

Sequencing

• Recommended merge order: google/tunix#2905 -> google/tunix#2906 (with AI-Hypercomputer/maxtext#5751 on the MaxText side).

Testing

• pytest tests/experimental/worker/trainer_worker_stage_batch_test.py (14 passed)

Update (commit d8342a169): stage the next step's second micro-batch across the step boundary

With the pipelined boundary (#2878) the next step's first micro-batch is already staged on the trainer's slot before the update. The intra-step queue above only starts filling once a step's first fwd_bwd is pending, so every step began with an empty queue and its early, slow transfers were never hidden. This commit:

• Adds _prestage_across_boundary (rl_program.py): right after the next step's first micro-batch is staged, its second micro-batch is sent through the intra-step queue (stage_batch, keyed step_<k+1>_stage_…), so that transfer overlaps the current step's update, metrics, checkpoint and weight sync. The next step adopts it (_StepState.prestaged / prestage_task) as the first queue entry, or as its in-flight active_stage_task, and keeps filling from there.
• Same refusals as the first micro-batch (not when it would close its step; not at scored checkpoint boundaries); a declined stage falls back to the inline call. Memory bound unchanged: the first micro-batch's payload is dropped once staged, so at most one feed's output is held.
• Trainer: apply_deferred_update joins stage_batch/fwd_bwd on the worker's executor pool under TUNIX_STAGE_BATCH=stage, so its device wait at the boundary does not block the stage request on the server loop thread. No MaxText change beyond AI-Hypercomputer/maxtext#5751.
• Active only with pipeline_step_boundary + stage_next_microbatch + stage_train_batches + max_staged_batches >= 2; no new flag. Launch: TUNIX_STAGE_BATCH=stage (forwarded to both containers by mlperf_base.sh / k8s_launcher.sh).
• Also fixes rl_program_pipelined_boundary_robustness_test's _PrepackProbe._take_group for the count_exposed keyword this PR added (72 of its 100 cases raised TypeError at the previous head).

Modelled on siv-sk8's measured stage times against the merged boundary (image trellis-experimental:1014c): ~46 s saved over six steps with a ~10 s boundary window, vs ~13 s for the queue alone; to be confirmed by an A/B on hardware.

Testing (jax 0.10.1, py3.12): tests/experimental/orchestrator/rl_program_prestage_across_boundary_test.py (18 tests + 10 subtests pass), rl_program_stage_before_update_test (165) and rl_program_pipelined_boundary_test (199) green, rl_program_pipelined_boundary_robustness_test 100/100, tests/experimental/worker/*stage_batch* 75/75.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a 2-deep background staging pipeline and concurrent method execution on the TrainerWorker to significantly reduce training stalls. By pre-staging micro-batches while the trainer is busy with weight updates and utilizing a multi-threaded executor for staging tasks, the system effectively hides latency between micro-batch processing and eliminates significant idle gaps on the TPU.

Highlights

• 2-Deep Background Staging: Implemented background staging for up to 2 micro-batches in the orchestrator, allowing pre-staging across step boundaries to eliminate inter-step idle time.
• Concurrent Trainer Execution: Upgraded the TrainerWorker's method executor to a 2-worker pool, enabling parallel execution of stage_batch and train_step to resolve mid-step TPU idle gaps.
• Testing: Added 9 new unit tests in tests/experimental/worker/trainer_worker_stage_batch_test.py to verify the staging logic and concurrent execution behavior.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request implements background staging of training micro-batches to overlap data transfer with computation, updating the orchestrator, distributed RL engine, and trainer worker. It also introduces a thread pool in the remote execution server to run synchronous executor methods off the main event loop. The review feedback highlights several violations of the style guide's "One well-lit path" and "Fail loud and fast" principles, specifically advising against speculative getattr calls and redundant checks on strictly typed contracts in favor of direct attribute access and standard exception handling.