Tighten async batch collation docs

Author: dyurk-lila

skyrl/train/config/sft_config.py

@@ -203,13 +203,10 @@ def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig":
     num_workers: int = 8
     """Number of worker processes for parallel tokenization during dataset loading. Set to 0 for single-threaded."""
     async_batch_collation: bool = True
-    """Async double-buffer the per-step collate. When True, the CPU-side collate
-    for step N+1 runs on a single background thread while step N's
-    forward/backward runs on the GPU, hiding the collate latency. The collate is
-    deterministic within an epoch, so the collated-ahead batch is byte-identical
-    to the synchronous path; at an epoch boundary (data reshuffle) the
-    collate-ahead is skipped and the next batch is collated synchronously on the
-    post-shuffle order. Set to False to A/B against the serial data-loading path."""
+    """Overlap the next step's CPU collate with the current GPU step.
+
+    Collate-ahead is skipped across epoch reshuffles. Set to False for the
+    serial data-loading path."""
 
     # ---- Tokenized dataset caching ----
     cache_dir: str = os.path.join(

skyrl/train/sft_trainer.py

@@ -1682,15 +1682,8 @@ def train(self):
         self._fire("on_epoch_start")
         epoch_in_progress = True
 
-        # ------------------------------------------------------------------
-        # Async batch collation (double-buffering) setup
-        # ------------------------------------------------------------------
-        # Two consecutive steps in the same epoch see the SAME ``tokenized``
-        # order (it only changes at an epoch boundary reshuffle), so step N+1's
-        # slice is knowable while step N runs. ``_slice_examples`` reproduces the
-        # loop's deterministic wrap-around slice against the *current* order, and
-        # ``_collate_batch`` is the producer run on the background thread (the
-        # heavy collate releases the GIL, so it overlaps the GPU step).
+        # Collate step N+1 on a background thread while step N runs on GPU.
+        # Do not collate ahead across reshuffles; the tokenized order changes.
         n_examples = len(tokenized)
 
         def _epoch_of(step: int) -> int:
@@ -1716,13 +1709,6 @@ def _collate_batch(step: int):
             f"SFT async batch collation (double-buffering): {'ENABLED' if collate_ahead_enabled else 'disabled'}"
         )
 
-        # Whether the step about to run can collate its successor ahead. The loop
-        # reshuffles ``tokenized`` after step N iff ``_epoch_of(N) > cur_epoch``
-        # (the same predicate the epoch-boundary block below uses). When a
-        # reshuffle would occur, step N+1 reads a DIFFERENT order than step N, so
-        # it must not be collated ahead against the pre-shuffle order. Mirroring the
-        # loop's reshuffle decision via ``cur_epoch`` (the authoritative loop
-        # state) keeps the predicate exact regardless of wrap-around alignment.
         def _can_collate_ahead(step: int, cur_epoch: int) -> bool:
             if async_collator is None or step + 1 > num_steps:
                 return False
@@ -1737,24 +1723,14 @@ def _can_collate_ahead(step: int, cur_epoch: int) -> bool:
 
                 with Timer("step", all_timings):
 
-                    # Data loading with wrap-around. With async batch collation
-                    # enabled this measures only the (ideally ~0) wait for the
-                    # already-running background collate; otherwise the full serial collate.
+                    # With async enabled, this is usually just the wait for an
+                    # already-running collate; otherwise it is the full collate.
                     with Timer("data_loading", all_timings):
                         if async_collator is not None and async_collator.pending_step() == self.global_step:
-                            # Consume the batch collated ahead during the previous step.
-                            # ``get`` asserts the in-flight step matches, so a
-                            # stale/mismatched batch fails loudly.
                             batch = async_collator.get(self.global_step)
                         else:
-                            # No valid in-flight batch (first step, or the first
-                            # step after an epoch reshuffle): collate synchronously
-                            # against the live order.
                             batch = _collate_batch(self.global_step)
 
-                        # Kick off the NEXT step's collate on the background thread so
-                        # it overlaps this step's GPU work — only when the successor
-                        # is in the same epoch (no reshuffle between them).
                         if _can_collate_ahead(self.global_step, current_epoch):
                             async_collator.submit(self.global_step + 1)
 
@@ -1847,12 +1823,7 @@ def _can_collate_ahead(step: int, cur_epoch: int) -> bool:
                 if epoch > current_epoch:
                     self._fire("on_epoch_end")
                     epoch_in_progress = False
-                    # Drain any in-flight batch BEFORE reshuffling so a background
-                    # collate can never read ``tokenized`` while it is being
-                    # shuffled, and so the next epoch's first step is collated
-                    # synchronously against the post-shuffle order. ``_can_collate_ahead``
-                    # already withholds cross-epoch submits, so this is normally a
-                    # no-op — it's defense in depth against the reshuffle/collate-ahead race.
+                    # Drain before mutating tokenized order.
                     if async_collator is not None:
                         async_collator.clear()
                     for _ in range(epoch - current_epoch):
@@ -1865,10 +1836,6 @@ def _can_collate_ahead(step: int, cur_epoch: int) -> bool:
 
                 self.global_step += 1
         finally:
-            # Always tear down the async batch collation thread (drains any
-            # in-flight batch and joins the worker) so neither the background
-            # thread nor the dataset reference is leaked, even on exception.
-            # No-op when async batch collation is disabled.
             if async_collator is not None:
                 async_collator.shutdown()
             if self._torch_profiler_enabled:

skyrl/train/utils/async_batch_collator.py

@@ -1,18 +1,4 @@
-"""Single-slot async double-buffer for a deterministic per-step producer.
-
-When a training loop's per-step input is built by a CPU-heavy producer that is
-deterministic given the current loop state (e.g. a slice + collate over an
-in-memory dataset), the producer for step ``N+1`` can run on a background
-thread while step ``N``'s forward/backward runs on the GPU, collating the next
-batch ahead of time and hiding the producer latency under GPU compute.
-
-:class:`AsyncBatchCollator` is the generic mechanism: it owns a single-worker
-executor and a one-slot future, and exposes a strict submit/get contract that
-turns any mis-wiring into a loud failure rather than a silent corruption of the
-input order. The caller supplies a ``compute(step) -> batch`` producer and is
-responsible for only ever collating ahead a step whose input is knowable in
-advance (i.e. no state change happens between submit and consume).
-"""
+"""Single-slot async double-buffer for deterministic per-step collation."""
 
 from __future__ import annotations
 
@@ -21,26 +7,10 @@
 
 
 class AsyncBatchCollator:
-    """Single-slot async double-buffer for a deterministic per-step producer.
-
-    Correctness contract:
-
-    * The producer is a single function ``compute(step) -> batch`` supplied by
-      the caller; it must read the loop's *current* state. The caller only ever
-      submits a step whose input is knowable while the current step runs — i.e.
-      a step for which no state change happens between submit and consume. If
-      the loop is about to mutate the producer's inputs (e.g. a dataset
-      reshuffle at an epoch boundary), it must :meth:`clear` the in-flight batch
-      first and collate the next step synchronously.
-    * At most one batch is in flight (``max_workers=1``). The submitted step is
-      recorded so :meth:`get` can assert the retrieved batch matches the step
-      the caller expects.
-    * The producer must construct fresh outputs per call and mutate no shared
-      state, so the worker thread genuinely overlaps with the GPU stream
-      (NumPy/torch release the GIL during the heavy array ops).
+    """Run ``compute(step)`` in a one-worker, one-future buffer.
 
-    This component owns only its executor and a one-slot future; it holds no
-    reference to caller state, which keeps the threading surface tiny.
+    The caller may only submit steps whose inputs will not change before
+    consumption. ``get`` checks the expected step so stale batches fail loudly.
     """
 
     def __init__(self, compute: Callable[[int], Any], thread_name_prefix: str = "batch-collate"):
@@ -50,11 +20,7 @@ def __init__(self, compute: Callable[[int], Any], thread_name_prefix: str = "bat
         self._pending_step: Optional[int] = None
 
     def submit(self, step: int) -> None:
-        """Schedule the producer for ``step`` on the background thread.
-
-        Must not be called while a batch is already in flight (the single-slot
-        invariant); drain via :meth:`get` before resubmitting.
-        """
+        """Schedule ``compute(step)``; call ``get`` before submitting again."""
         assert self._future is None, (
             f"collate-ahead slot already occupied (pending step {self._pending_step}); "
             f"call get() before submitting step {step}"
@@ -69,12 +35,7 @@ def pending_step(self) -> Optional[int]:
         return self._pending_step
 
     def get(self, expected_step: int) -> Any:
-        """Block until the in-flight batch is ready and return it.
-
-        Asserts the in-flight batch was produced for ``expected_step`` so a
-        mis-wired submit/get pairing fails loudly instead of silently feeding
-        the wrong batch into training.
-        """
+        """Return the in-flight batch for ``expected_step``."""
         assert self._future is not None, "get() called with no in-flight batch"
         assert self._pending_step == expected_step, (
             f"collated-ahead step {self._pending_step} != expected step {expected_step}; "
@@ -87,13 +48,7 @@ def get(self, expected_step: int) -> Any:
         return future.result()
 
     def clear(self) -> None:
-        """Drop any in-flight batch.
-
-        Waits for the worker to finish (so it cannot still be reading caller
-        state) and discards the result. Used before the caller mutates the
-        producer's inputs, so a batch collated against the old state can never
-        be consumed.
-        """
+        """Drain and discard the in-flight batch, propagating worker errors."""
         if self._future is not None:
             self._future.result()
         self._future = None