Repository navigation
Conversation
…'s first micro-batch is sent earlier After a step commits, `_train_loop` used to run the deferred trainer `get_metrics`, the step metric reduction and logging, `on_step_end`, and `gc.collect()` + `malloc_trim` inline on the event loop before pulling the next step's groups. The trainer sat idle for all of it. On at-sph1-s1a (35B, gbs512) this was ~2.5 s per step boundary on top of packing, and the sync reduction and GC also stalled the dispatch/polling stages. This moves that work into a background task created right after the commit: * The deferred `get_metrics` and the filtered-group drain run immediately. * The GIL-bound work (metric reduction, trajectory logging, `on_step_end`, memory release) waits until the next step has issued its first `train_step`, so it overlaps trainer compute rather than the next step's first packing. It runs in a worker thread. * `_train_loop` awaits the task before the next step's optimizer update (the trainer's `get_metrics` slot is overwritten there, so the previous fetch must have landed) and when the loop ends, so a failure in it still fails the run. If a step's first micro-batch is also its final one, the gate is released before that await so the two cannot wait on each other. `step_time_sec` is now measured at commit rather than after the deferred metrics fetch; `metrics_fetch_time` is still reported separately.
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 optimizes the training loop by decoupling end-of-step bookkeeping from the main execution path. By offloading metrics fetching, reduction, and memory management to background tasks, the system can initiate the next training step's micro-batch earlier, significantly reducing idle time for the trainer. The changes ensure that these background operations are synchronized with the trainer's progress and that any failures are handled gracefully to maintain run integrity. 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 refactors the end-of-step bookkeeping in StandardRLProgram to run as a background task, allowing the next step's first micro-batch to be issued immediately without waiting for metrics fetching, reduction, logging, or host memory release. It also offloads metric reduction and host memory release to background threads. Feedback on these changes suggests improving type safety in the newly introduced _finish_step_bookkeeping method by replacing Any annotations with concrete, specific types in accordance with the repository's strict typing guidelines.
| async def _finish_step_bookkeeping( | ||
| self, | ||
| *, | ||
| current_step: int, | ||
| step_result: Any, | ||
| trainer_metrics: Any, | ||
| pending_metrics_fetch: bool, | ||
| all_step_items: list[Any], | ||
| uncommitted_groups: list[Any], | ||
| scored_items: list[Any], | ||
| step_rewards: list[float], | ||
| step_advantages: list[float], | ||
| generation_metrics: dict[str, Any], | ||
| num_rollouts: int, | ||
| num_microbatches: int, | ||
| step_padding_stats: list[batch_assembly.PaddingStats], | ||
| step_packing_time_sec: float, | ||
| step_time_sec: float, | ||
| consumed_policy_version: int, | ||
| step_sampler_agreement: dict[str, tuple[Any, list[float]]], | ||
| policy_training_time: float, | ||
| exposed_generation_time: float, | ||
| weight_sync_time: float, | ||
| policy_version: int, | ||
| ) -> None: |
There was a problem hiding this comment.
To adhere to the repository style guide on strict typing, we should avoid using Any in the signature of _finish_step_bookkeeping where more specific types can be expressed. Specifically:
all_step_itemscan be typed aslist[datatypes.TrajectoryItem]uncommitted_groupscan be typed aslist[list[datatypes.TrajectoryItem]]scored_itemscan be typed aslist[datatypes.TrajectoryItem]generation_metricscan be typed asdict[str, float]step_sampler_agreementcan be typed asdict[str, tuple[Callable[[list[float]], float], list[float]]]
async def _finish_step_bookkeeping(
self,
*,
current_step: int,
step_result: Any,
trainer_metrics: Any,
pending_metrics_fetch: bool,
all_step_items: list[datatypes.TrajectoryItem],
uncommitted_groups: list[list[datatypes.TrajectoryItem]],
scored_items: list[datatypes.TrajectoryItem],
step_rewards: list[float],
step_advantages: list[float],
generation_metrics: dict[str, float],
num_rollouts: int,
num_microbatches: int,
step_padding_stats: list[batch_assembly.PaddingStats],
step_packing_time_sec: float,
step_time_sec: float,
consumed_policy_version: int,
step_sampler_agreement: dict[str, tuple[Callable[[list[float]], float], list[float]]],
policy_training_time: float,
exposed_generation_time: float,
weight_sync_time: float,
policy_version: int,
) -> None:References
- Always annotate function signatures, attributes, and data structures with concrete, specific types rather than Any or untyped dict/object. Avoid Any whenever a specific type can be expressed. (link)
2 participants
Problem
After a step commits,
_train_loopran the deferred trainerget_metrics, the step metric reduction and logging,on_step_end, andgc.collect()+malloc_triminline on the event loop before pulling step N+1's groups. The trainer sat idle for all of it, and the synchronous reduction/GC also stalled the dispatch, polling and critique stages.Measured on
at-sph1-s1a(35B, gbs512,ASYNC_WEIGHT_SYNC,PIPELINE_TRAIN_MICROBATCHES), median per step boundary:get_metricsRPCChange
End-of-step bookkeeping now runs as a background task created right after the commit:
get_metricsand the filtered-group drain run immediately.on_step_end, memory release) waits until the next step has issued its firsttrain_step, so it overlaps trainer compute instead of the next step's first packing. It runs in a worker thread (asyncio.to_thread)._train_loopawaits the task before the next step's optimizer update — the trainer'sget_metricsslot (_written_metrics) is overwritten there, so the previous fetch must have landed first — and when the loop ends, so a failure in the task still fails the run. If a step's first micro-batch is also its final one, the gate is released before that await so the two cannot wait on each other.run_async/train_stagecancel the task on teardown, like_pending_train/_pending_sync.step_time_secis now measured at commit rather than after the deferred metrics fetch;metrics_fetch_timeis still reported separately.Behavioural note: the
>>> step N+1 startinglog line (fromon_step_begin) can now appear before<<< step N finished/Train step N - loss ..., since those are emitted by the background task.Tests
test_next_step_first_train_step_not_held_by_deferred_get_metrics: step N+1's first micro-batch is issued while step N'sget_metricsis still outstanding, and the fetch lands before step N+1's optimizer update.test_on_step_end_runs_after_next_step_issued_its_first_train_steptest_on_step_end_failure_fails_the_runtest_step_metrics_reduce_off_the_event_looptests/experimental/orchestrator/rl_program_test.py(186) andorchestrator_test.pypass on CPU. Thedeepswe_distlauncher/reaper tests (mlperf_submission_launcher_test,prepull_sandbox_images_test,cluster_reaper_test) fail on my VM with and without this change (JobSet render / env dependent); they do not importrl_program.