[train] Async batch collation (double-buffering) for the SFT trainer

skyrl/train/config/sft_config.py

@@ -202,6 +202,11 @@ def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig":
     # ---- Data loading ----
     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
+    """Overlap the next stateful-dataloader batch with the current GPU step.
+
+    Checkpoint state remains pinned after the current batch. Set to False for
+    serial data loading."""
 
     # ---- Dataloader / sampler ----
     dataloader_num_workers: int = 0

skyrl/train/sft_trainer.py

@@ -55,6 +55,7 @@
     get_response_ids_and_loss_mask_from_messages,
 )
 from skyrl.train.utils import get_ray_pg_ready_with_timeout
+from skyrl.train.utils.async_batch_collator import AsyncBatchCollator
 from skyrl.train.utils.callbacks import (
     CallbackHandler,
     CallbackInput,
@@ -812,6 +813,7 @@ def __init__(
         # Stateful dataloaders, built in train() once data is tokenized.
         self.train_dataloader: StatefulDataLoader | None = None
         self.eval_dataloader: StatefulDataLoader | None = None
+        self._checkpoint_dataloader_state: dict | None = None
         self.global_step = 0
         # running count of total non-padding tokens trained on
         self._total_tokens_processed = 0
@@ -1850,6 +1852,16 @@ def train(self):
         # their state across the (conceptual) epoch boundaries.
         data_iter = iter(self.train_dataloader)
 
+        collate_ahead_enabled = self.sft_cfg.async_batch_collation
+        async_collator: Optional[AsyncBatchCollator] = (
+            AsyncBatchCollator(lambda _step: next(data_iter, None), thread_name_prefix="sft-batch-collate")
+            if collate_ahead_enabled
+            else None
+        )
+        logger.info(
+            f"SFT async batch collation (double-buffering): {'ENABLED' if collate_ahead_enabled else 'disabled'}"
+        )
+
         if self._torch_profiler_enabled:
             self.dispatch.start_profile("policy")
         try:
@@ -1858,10 +1870,14 @@ def train(self):
 
                 with Timer("step", all_timings):
 
-                    # Fetch the next batch; on epoch exhaustion, close the epoch and
-                    # restart the iterator (reshuffles the random/sequential samplers).
+                    # With async enabled, this is usually just the wait for an
+                    # already-running collate. ``None`` marks epoch exhaustion.
                     with Timer("data_loading", all_timings):
-                        batch = next(data_iter, None)
+                        if async_collator is not None and async_collator.pending_step() == self.global_step:
+                            batch = async_collator.get(self.global_step)
+                            self._checkpoint_dataloader_state = None
+                        else:
+                            batch = next(data_iter, None)
                     if batch is None:
                         self._fire("on_epoch_end")
                         current_epoch += 1
@@ -1871,6 +1887,13 @@ def train(self):
                         with Timer("data_loading", all_timings):
                             batch = next(data_iter)
 
+                    if async_collator is not None and self.global_step < num_steps:
+                        # Advancing the iterator in the worker moves the live
+                        # dataloader state one batch ahead. Preserve the state after
+                        # the current batch so checkpoints still resume exactly.
+                        self._checkpoint_dataloader_state = self.train_dataloader.state_dict()
+                        async_collator.submit(self.global_step + 1)
+
                     self._fire("on_step_start", batch=batch)
 
                     # Training step
@@ -1965,6 +1988,13 @@ def train(self):
 
                 self.global_step += 1
         finally:
+            # Always tear down the async 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 collation is disabled.
+            if async_collator is not None:
+                async_collator.shutdown()
+            self._checkpoint_dataloader_state = None
             if self._torch_profiler_enabled:
                 self.dispatch.stop_profile("policy")
         self.global_step = min(self.global_step, num_steps)
@@ -2040,7 +2070,12 @@ def save_checkpoint(self) -> str:
             dataloader_save_path = os.path.join(global_step_folder, "data.pt")
             try:
                 with io.open_file(dataloader_save_path, "wb") as f:
-                    torch.save(self.train_dataloader.state_dict(), f)
+                    dataloader_state = (
+                        self._checkpoint_dataloader_state
+                        if self._checkpoint_dataloader_state is not None
+                        else self.train_dataloader.state_dict()
+                    )
+                    torch.save(dataloader_state, f)
                 logger.info(f"Saved dataloader state to {dataloader_save_path}")
             except Exception as e:
                 logger.warning(f"Failed to save dataloader state: {e}")

skyrl/train/utils/async_batch_collator.py

@@ -0,0 +1,58 @@
+"""Single-slot async double-buffer for deterministic per-step collation."""
+
+from concurrent.futures import Future, ThreadPoolExecutor
+from typing import Any, Callable
+
+
+class AsyncBatchCollator:
+    """Run ``compute(step)`` in a one-worker, one-future buffer.
+
+    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"):
+        self._compute = compute
+        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix=thread_name_prefix)
+        self._future: Future | None = None
+        self._pending_step: int | None = None
+
+    def submit(self, step: int) -> None:
+        """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}"
+        )
+        self._pending_step = step
+        self._future = self._executor.submit(self._compute, step)
+
+    def has_pending(self) -> bool:
+        return self._future is not None
+
+    def pending_step(self) -> int | None:
+        return self._pending_step
+
+    def get(self, expected_step: int) -> Any:
+        """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}; "
+            f"refusing to serve a mismatched batch"
+        )
+        future = self._future
+        self._future = None
+        self._pending_step = None
+        # Propagates any exception raised inside the worker thread.
+        return future.result()
+
+    def clear(self) -> None:
+        """Drain and discard the in-flight batch, propagating worker errors."""
+        if self._future is not None:
+            self._future.result()
+        self._future = None
+        self._pending_step = None
+
+    def shutdown(self) -> None:
+        """Drain any in-flight batch and join the worker thread."""
+        self.clear()
+        self._executor.shutdown(wait=True)