[train] Async batch collation (double-buffering) for the SFT trainer#1809
[train] Async batch collation (double-buffering) for the SFT trainer#1809dyurk-lila wants to merge 4 commits into
Conversation
![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 async batch collation (double-buffering) to the SFT training loop, allowing the CPU-side batch slicing and collation for step N+1 to run on a background thread while step N's forward/backward pass runs on the GPU. This change hides collation latency under GPU compute. It adds the AsyncBatchCollator utility class, integrates it into SFTTrainer, adds configuration options, and includes comprehensive tests. Feedback was provided on skyrl/train/utils/async_batch_collator.py to avoid catching BaseException in the clear method, as it can intercept system-exiting exceptions like SystemExit and KeyboardInterrupt, recommending catching Exception instead.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
b9775f6 to
1bd63a8
Compare
The SFT training loop builds each step's batch with a CPU-side slice + collate on the main thread, which blocks the GPU. The per-step slice is fully deterministic from (global_step, batch_size, len(tokenized)) given the current data order, and the order only changes at epoch boundaries (rng.shuffle). So step N+1's batch can be collated on a background thread while step N's forward/backward runs on the GPU, hiding the collate latency under GPU compute. Add BatchPrefetcher (skyrl/train/utils/batch_prefetcher.py), a generic single-slot async double-buffer: a max_workers=1 ThreadPoolExecutor with a strict submit/get contract that asserts the in-flight step matches the step the loop expects, turning any mis-wiring into a loud failure rather than a silent training-order corruption. Wire it into SFTTrainer.train(): * Cross-epoch prefetch is withheld: _can_prefetch_next mirrors the loop's own reshuffle predicate, so step N+1 is never prefetched against the pre-shuffle order. The first step of each new epoch is computed synchronously against the post-shuffle order. * The prefetch slot is drained (clear()) before any reshuffle, so a worker thread can never read `tokenized` while it is being shuffled. * The executor is shut down in the loop's finally block (drains the in-flight batch + joins the worker) so the background thread is never leaked, even on exception. Gated by SFTConfig.prefetch_data (default True); set False to A/B against the serial data-loading path. Scope: this is an SFT-trainer optimization. The RL trainers do not need it -- in synchronous RL the per-step batch depends on the just-generated rollout, so it cannot be prefetched ahead of generation, and the fully async RL trainer already overlaps generation with training via its asyncio queue. BatchPrefetcher is generic, so a future deterministic producer can reuse it. tests/train/test_sft_prefetch.py: runs the real SFTTrainer.train() loop twice over the same dummy dataset/seed (prefetch OFF vs ON) across two epoch boundaries and asserts every batch handed to train_step is tensor-equal step-for-step, for both DefaultCollator and PackedDataCollator. Examples have distinct token ids, so a stale pre-shuffle batch would diverge from the baseline and fail. Plus BatchPrefetcher invariant unit tests (step-mismatch assertion, single-slot guard, worker-exception propagation, clear() drain). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Standardize on async-batch-collation vocabulary: rename the class BatchPrefetcher → AsyncBatchCollator (module batch_prefetcher.py → async_batch_collator.py), the SFTConfig flag prefetch_data → async_batch_collation, and the loop wiring / helper / test symbols accordingly. Behavior is unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1bd63a8 to
4241f4c
Compare
SumanthRH
left a comment
SumanthRH
left a comment
There was a problem hiding this comment.
LGTM!
Tests look sufficient to capture basic correctness and ckpt-resume behaviour. To be sure, I also did a manual run and tested the PR E2E in the following way:
Script used: run_sft_fsdp.sh: https://github.com/NovaSky-AI/SkyRL/blob/fd79ceec9f66b0562fb15b7e70295652026eefd6/examples/train/sft/run_sft_fsdp.sh
- Serial baseline: Run the above script for 6 steps with
async_batch_collationFalse - Async: Run the script with
async_batch_collation=True - Async, Ckpt + Resume: Run training for 4 steps, ckpt and then resume from step 4 ckpt.
Run 3 matches Run 1 and Run 2 in overall loss values and the input_ids used in step 5 in Run 3 matches step 5 in Run 2.
|
Can you merge from main @dyurk-lila ? |
2 participants
Summary
Overlap the next SFT
StatefulDataLoaderbatch with the current GPU step. A single background thread advances the dataloader and performs collation while the current batch runs forward/backward.This rebased version integrates with the stateful dataloader and custom sampler support added in #1842:
AsyncBatchCollatorprovides a one-worker, one-future buffer with strict step matching.SFTTrainer.train()consumes the buffered batch, handles dataloader exhaustion and epoch transitions, and always joins the worker on exit.SFTConfig.async_batch_collationdefaults toTrue; set it toFalsefor the serial path.Validation
AsyncBatchCollatorunit tests cover step mismatch, single-slot enforcement, worker errors, draining, and shutdown.54 passedacrosstest_async_batch_collation.pyandtest_sft_dataloader.py.