Repository navigation
Conversation
… update runs The step's final micro-batch (the one that applies the optimizer) was the only train_step that was never overlapped with anything: `_train_loop` awaited it, then the checkpoint, the weight-sync snapshot and the commit, and only then pulled step N+1's first groups from `scored_q` and packed them. On at-sph1-s1a (35B, gbs512) the update took ~27 s and the first pack of the next step another ~4 s, during which the trainer was idle, even though step N+1's groups had already arrived. With `prefetch_next_batch=True`, an optimizer-update train_step is run as a task and, while it is in flight, the loop pulls ready groups from `scored_q` and feeds them to the assembler. The packed micro-batches are held in `self._prefetched` and consumed first by whichever step (or mini-batch) comes next, which also attributes their bookkeeping to it. Nothing is dispatched early: the next micro-batch still waits for the update, the checkpoint, the sync snapshot and the commit. Under prompt-batch ordering the first fetch draws from the queue's cursor (the rest of the current batch for a mid-batch update, otherwise the next batch, which the queue already serves one ahead of the committed one); later fetches are pinned to that batch, so the prefetch never crosses a second boundary. When the buffer's head belongs to a later batch than the step's, the step treats that as the end of its batch, as it would a None from the queue. The prefetch is skipped at the end of the last step, and the buffer is cleared with the assembler on reset. The fetch and feed code is factored into `_fetch_group` / `_feed_group` and shared by the loop and the prefetch. The flag is plumbed through `run_deepswe_dist.py` (`--prefetch_next_batch`), `k8s_launcher.sh` (`PREFETCH_NEXT_BATCH`) and the MLPerf recipes, where the gbs512 and gbs1024 submission recipes enable it alongside `PIPELINE_TRAIN_MICROBATCHES`.
Summary of ChangesHello, I'm Gemini Code Assist1! 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 implements a performance optimization for the RL training loop by overlapping data preparation with the optimizer update step. By pre-packing upcoming groups during the relatively idle time of an optimizer update, the system reduces the time spent waiting for the next micro-batch. The change is configurable via a new flag and includes necessary updates to existing launcher scripts and recipes to ensure seamless integration. Highlights
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 AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a prefetching mechanism (prefetch_next_batch) to the RL training program, allowing upcoming trajectory groups to be pulled and packed from the scored queue while an optimizer update is in flight. The feedback identifies a critical missing import (contextlib) that would cause a runtime crash, and points out several style guide violations. Specifically, the reviewer recommends using strict type hints instead of Any, avoiding speculative getattr calls on structured schemas to follow the 'one well-lit path' principle, and using asyncio.create_task instead of asyncio.ensure_future for consistency.
| with contextlib.suppress(asyncio.CancelledError): | ||
| await fetch |
There was a problem hiding this comment.
The contextlib module is not imported in this file, which will cause a NameError and crash the training loop when a prefetch task is cancelled.
Instead of using contextlib.suppress, use a standard try-except block to handle asyncio.CancelledError safely without requiring additional imports.
| with contextlib.suppress(asyncio.CancelledError): | |
| await fetch | |
| try: | |
| await fetch | |
| except asyncio.CancelledError: | |
| pass |
| """ | ||
|
|
||
| batch_idx: int | None | ||
| scored_items: list[Any] |
There was a problem hiding this comment.
According to the Tunix Code Review Style Guide (Section 5, Type hints & Strict Typing), we should use strict, specific type hints rather than Any whenever a specific type can be expressed.
Since scored_items contains datatypes.TrajectoryItem objects, we should annotate it as list[datatypes.TrajectoryItem].
| scored_items: list[Any] | |
| scored_items: list[datatypes.TrajectoryItem] |
References
- Use strict, specific type hints: Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any. (link)
|
|
||
| async def _fetch_group( | ||
| self, batch_idx: int | None | ||
| ) -> tuple[int | None, list[Any]]: |
There was a problem hiding this comment.
According to the Tunix Code Review Style Guide (Section 5, Type hints & Strict Typing), we should use strict, specific type hints rather than Any whenever a specific type can be expressed.
Annotate the return type of _fetch_group as tuple[int | None, list[datatypes.TrajectoryItem]] instead of using list[Any].
| ) -> tuple[int | None, list[Any]]: | |
| ) -> tuple[int | None, list[datatypes.TrajectoryItem]]: |
References
- Use strict, specific type hints: Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any. (link)
| return batch_idx, await self.scored_q.get_group_batch(num_groups=1) | ||
|
|
||
| async def _feed_group( | ||
| self, scored_items: list[Any] |
There was a problem hiding this comment.
According to the Tunix Code Review Style Guide (Section 5, Type hints & Strict Typing), we should use strict, specific type hints rather than Any whenever a specific type can be expressed.
Annotate the scored_items parameter as list[datatypes.TrajectoryItem] instead of list[Any].
| self, scored_items: list[Any] | |
| self, scored_items: list[datatypes.TrajectoryItem] |
References
- Use strict, specific type hints: Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any. (link)
| """ | ||
| payloads = [] | ||
| for item in scored_items: | ||
| payload = getattr(item, "payload", None) |
There was a problem hiding this comment.
According to the Tunix Code Review Style Guide (Section 5, Error Handling & One Well-Lit Path), we should never use speculative getattr(obj, "field", default) on structured schemas. Instead, we should access attributes directly on a single, strictly typed contract.
Since item is a datatypes.TrajectoryItem and its payload is guaranteed to be set during the critique stage, we should access item.payload directly.
| payload = getattr(item, "payload", None) | |
| payload = item.payload |
References
- One well-lit path (No defensive fallbacks): Never use speculative getattr(obj, "field", default) on structured schemas. Access attributes and keys directly on a single, strictly typed contract. (link)
| item.metadata.pop("routed_experts", None) | ||
| if isinstance(item.traj, dict): | ||
| item.traj.pop("routed_experts", None) | ||
| item_payload = getattr(item, "payload", None) |
There was a problem hiding this comment.
According to the Tunix Code Review Style Guide (Section 5, Error Handling & One Well-Lit Path), we should never use speculative getattr(obj, "field", default) on structured schemas. Instead, we should access attributes directly on a single, strictly typed contract.
Since item is a datatypes.TrajectoryItem and its payload is guaranteed to be set during the critique stage, we should access item.payload directly.
| item_payload = getattr(item, "payload", None) | |
| item_payload = item.payload |
References
- One well-lit path (No defensive fallbacks): Never use speculative getattr(obj, "field", default) on structured schemas. Access attributes and keys directly on a single, strictly typed contract. (link)
| return assembled_batches, packing_time_sec | ||
|
|
||
| def _has_groups_after( | ||
| self, groups_consumed: int, scored_items: list[Any] |
There was a problem hiding this comment.
According to the Tunix Code Review Style Guide (Section 5, Type hints & Strict Typing), we should use strict, specific type hints rather than Any whenever a specific type can be expressed.
Annotate the scored_items parameter as list[datatypes.TrajectoryItem] instead of list[Any].
| self, groups_consumed: int, scored_items: list[Any] | |
| self, groups_consumed: int, scored_items: list[datatypes.TrajectoryItem] |
References
- Use strict, specific type hints: Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any. (link)
| while ( | ||
| not train_task.done() and len(self._prefetched) < self.full_batch_size | ||
| ): | ||
| fetch = asyncio.ensure_future(self._fetch_group(pinned_batch_idx)) |
There was a problem hiding this comment.
For consistency with the rest of the codebase (which uses asyncio.create_task exclusively) and to follow modern asyncio best practices, use asyncio.create_task instead of asyncio.ensure_future.
| fetch = asyncio.ensure_future(self._fetch_group(pinned_batch_idx)) | |
| fetch = asyncio.create_task(self._fetch_group(pinned_batch_idx)) |
References
- Consistency: Adherence to a consistent style across all projects. (link)
2 participants
Problem
The step's final micro-batch (the one that applies the optimizer) was the only
train_stepnever overlapped with anything:_train_loopawaited it, then the checkpoint, the weight-sync snapshot and the commit, and only then pulled step N+1's first groups fromscored_qand packed them. Measured onat-sph1-s1a(35B, gbs512), median per step:Step N+1's groups are already in
scored_qduring the update (the[pipeline] ARRIVElines for the next batch land throughout), so the first pack can happen then instead.Change
New
prefetch_next_batchoption (default off). When set, an optimizer-updatetrain_stepruns as a task and, while it is in flight, the loop pulls ready groups fromscored_qand feeds them to the assembler. Packed micro-batches are held inself._prefetchedand consumed first by whichever step (or mini-batch) comes next, which also attributes their rewards/advantages/packing time to it. Nothing is dispatched early: the next micro-batch still waits for the update, the checkpoint, the sync snapshot and the commit.Details:
BatchOrderedQueueManageralready serves one batch ahead of the committed one by design). Later fetches are pinned to that batch, so the prefetch never crosses a second boundary. When the buffer's head belongs to a later batch than the step's, the step treats it as the end of its batch, exactly as it treats aNonefrom the queue.full_batch_sizegroups buffered, batch drained, queue closed. Skipped at the end of the last step. The buffer is cleared with the assembler on reset._pending_trainso teardown cancels it like any other train step; a prefetch failure is surfaced only after the update has returned, so an optimizer step is never left unobserved._fetch_group/_feed_groupfactor the existing fetch and pack code so the loop and the prefetch share it (no behaviour change there).run_deepswe_dist.py(--prefetch_next_batch),k8s_launcher.sh(PREFETCH_NEXT_BATCH, default false) and the MLPerf recipes; the gbs512 and gbs1024 submission recipes enable it alongsidePIPELINE_TRAIN_MICROBATCHES.Tests
test_prefetch_packs_next_batch_while_optimizer_update_runs(prompt-batch order): step 0's update observes batch 1 fully pulled and packed; each step still trains exactly its own batch andnext_batch_idxends at 2.test_prefetch_skips_the_last_steptest_prefetch_feeds_next_group_while_optimizer_update_runs(arrival order)mlperf_submission_launcher_testasserts the flag reaches the orchestrator JobSet.rl_program_test.py,orchestrator_test.py,config_wiring_test.py(229) pass on CPU. Thedeepswe_distlauncher tests fail on my VM with and without this change (JobSet render / env dependent).