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

# What does this PR do?

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.

Note this overlaps the in-memory **slice + collation** that assembles each training batch from the already-tokenized dataset; it is not disk/network data prefetching (that is a separate, orthogonal optimization).

This is **default-safe and byte-identical to the serial path**: the collated-ahead batch is exactly what the synchronous loop would have produced for the same step.

## Changes

- **`AsyncBatchCollator`, `skyrl/train/utils/async_batch_collator.py`** — a generic single-slot async double-buffer: a `max_workers=1` `ThreadPoolExecutor` with a one-slot future and 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. Teardown is exception-safe: `clear()` discards the in-flight batch (and any exception the orphaned worker raised) on a reset/teardown path so it cannot mask the caller's own unwind, and `shutdown()` always joins the executor even if draining the in-flight batch raises.
- **`SFTTrainer.train()`, `skyrl/train/sft_trainer.py`** — wire the async collator into the data-loading block.
- **`SFTConfig.async_batch_collation` (default `True`), `skyrl/train/config/sft_config.py`** — set `False` to A/B against the serial batch-building path. Defaulting to `True` flips the default SFT batch-building path to threaded; it is behavior-preserving (byte-identical batches; see tests) and only changes timing.

### Correctness guarantees

- **Cross-epoch collate-ahead is withheld**: `_can_collate_ahead` mirrors the loop's own reshuffle predicate, so step `N+1` is never collated against the pre-shuffle order. The first step of each new epoch is collated synchronously against the post-shuffle order.
- **Drain before reshuffle**: the in-flight slot is drained (`clear()`) before any reshuffle, so a worker thread can never read `tokenized` while it is being shuffled.
- **No leaked thread, even on exception**: the executor is shut down in the loop's `finally` block (drains the in-flight batch + joins the worker). Teardown is double-fault-safe — if the loop body raises *and* an in-flight collation worker also raised, the original loop exception is what propagates (the worker's exception is swallowed during the reset), and the worker thread is still joined.

## Scope

This is a controller-side **SFT-trainer** optimization. The other trainers do not need it:

- In **synchronous RL** (`trainer.py`), the per-step training batch is built from the *just-generated* rollout (`convert_to_training_input`), so it cannot be collated ahead of generation.
- The **fully async RL** trainer (`fully_async_trainer.py`) already overlaps generation with training via its `asyncio` queue.
- `SFTTrainer.run_eval()` has the same collate-then-dispatch shape that `train()` optimizes here, and (since the eval slice has no wrap-around/reshuffle) would be an even cleaner `AsyncBatchCollator` fit, but it is **intentionally left serial**: eval is off by default (`eval_interval=0`) and infrequent, and the per-batch model forward dominates the small eval collate, so the payoff is marginal. It can adopt the async collator later without API changes.

`AsyncBatchCollator` is written as a generic, reusable component (it holds no trainer state), so a future deterministic per-step producer can reuse it.

## Test plan

- [x] `tests/train/test_async_batch_collation.py` — **7 integration cases + 5 `AsyncBatchCollator` unit tests**:
  - The integration cases run the real `SFTTrainer.train()` loop twice over the same dummy dataset/seed (async collation OFF vs ON) across **two epoch boundaries** and assert every `batch` handed to `train_step` is tensor-equal step-for-step — for **both** `DefaultCollator` (FSDP left-pad) and `PackedDataCollator` (Megatron FFD packing), plus an uneven wrap-around case, a mid-epoch resume case, an eval-enabled case, and a `num_epochs`-derived `num_steps` case. Examples have distinct token ids, so a stale pre-shuffle batch would diverge from the baseline and fail. (Mutation-checked: disabling either the cross-epoch withhold or the epoch-boundary drain makes the integration tests fail.)
  - A teardown double-fault case: a loop-body `train_step` failure coincides with an in-flight collation worker that *also* raises, asserting the original loop exception propagates (not the worker's re-raised error) and that no `sft-batch-collate` worker thread survives after `train()` returns.
  - `AsyncBatchCollator` invariant unit tests: step-mismatch assertion, single-slot guard, worker-exception propagation through `get()`, and `clear()` drain.
  - The collate path is pure Python + torch (no numpy), so the packed-collator test needs no Megatron runtime.
- [x] `ruff==0.11.9` + `black==24.10.0` clean.

```bash
uv run --isolated --extra dev --extra fsdp pytest tests/train/test_async_batch_collation.py -v
```

Author: dyurk-lila

skyrl/train/config/sft_config.py

@@ -199,6 +199,15 @@ 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
+    """Async double-buffer the per-step batch slice + collate. When True, the
+    CPU-side slice+collate that builds step N+1's training batch runs on a single
+    background thread while step N's forward/backward runs on the GPU, hiding the
+    collation latency under GPU compute. The slice+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 batch-building path."""
 
     # ---- Tokenized dataset caching ----
     cache_dir: str = os.path.join(

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,
@@ -1553,118 +1554,194 @@ def train(self):
         self._fire("on_epoch_start")
         epoch_in_progress = True
 
-        while self.global_step <= num_steps:
-            all_timings: dict[str, float] = {}
-
-            with Timer("step", all_timings):
-
-                # Data loading with wrap-around
-                with Timer("data_loading", all_timings):
-                    start_idx = (self.global_step * batch_size) % len(tokenized)
-                    end_idx = start_idx + batch_size
-                    if end_idx > len(tokenized):
-                        batch_examples = tokenized[start_idx:] + tokenized[: end_idx - len(tokenized)]
-                    else:
-                        batch_examples = tokenized[start_idx:end_idx]
-                    batch = self.collator(batch_examples, batch_size=batch_size)
-
-                self._fire("on_step_start", batch=batch)
-
-                # Training step
-                step_result = self.train_step(batch, self.global_step)
-                all_timings.update(step_result["timings"])
-
-            # Compute throughput using actual (non-padding) tokens
-            batch_padded_seq_len = batch["sequences"].shape[1]
-            actual_num_tokens = batch["attention_mask"].sum().item()
-            self._total_tokens_processed += actual_num_tokens
-            tokens_per_second = actual_num_tokens / all_timings["step"]
-
-            # Build log dict
-            log_dict = {
-                "train/loss": step_result["loss"],
-                "train/grad_norm": step_result["grad_norm"],
-                "train/tokens_per_second": tokens_per_second,
-                "train/tokens_per_second_per_gpu": tokens_per_second / self._num_training_gpus,
-                "train/actual_num_tokens": actual_num_tokens,
-                "train/batch_padded_seq_len": batch_padded_seq_len,
-                "train/total_tokens_processed": self._total_tokens_processed,
-            }
-            log_dict.update({f"timing/{k}": v for k, v in all_timings.items()})
-            if self._ray_gpu_monitor is not None:
-                log_dict.update(self._ray_gpu_monitor.flush())
-
-            self._fire("on_step_end", batch=batch, metrics=step_result)
-
-            # Capture callback-driven triggers, then reset so they only fire once.
-            force_save = self._training_control.should_save
-            force_eval = self._training_control.should_evaluate
-            self._training_control.should_save = False
-            self._training_control.should_evaluate = False
-
-            # Checkpoint: interval-driven or callback-requested.
-            interval_save = (
-                self.sft_cfg.ckpt_interval > 0
-                and self.global_step > 0
-                and self.global_step % self.sft_cfg.ckpt_interval == 0
-            )
-            did_save_last_step = force_save or interval_save
-            if did_save_last_step:
-                with Timer("save_checkpoint", all_timings):
-                    ckpt_path = self.save_checkpoint()
-                log_dict["timing/save_checkpoint"] = all_timings["save_checkpoint"]
-                self._fire("on_save", ckpt_path=ckpt_path)
-
-            # HF export at regular intervals
-            if self.sft_cfg.hf_save_interval > 0 and self.global_step % self.sft_cfg.hf_save_interval == 0:
-                with Timer("save_hf_model", all_timings):
-                    self.save_hf_model()
-                log_dict["timing/save_hf_model"] = all_timings["save_hf_model"]
-
-            eval_metrics = None
-            num_eval_batches: int | None = None
-            # Eval fires at step N where N % eval_interval == 0 and N > 0, OR
-            # whenever a callback set ``control.should_evaluate``.
-            interval_eval = self.sft_cfg.eval_interval > 0 and self.global_step % self.sft_cfg.eval_interval == 0
-            if eval_tokenized is not None and (force_eval or interval_eval):
-                self._fire("on_eval_start")
-                with Timer("eval", all_timings):
-                    eval_metrics, num_eval_batches = self.run_eval(eval_tokenized)
-                self._fire("on_eval_end", metrics=eval_metrics)
-                if eval_metrics:
-                    log_dict.update({f"eval/{k}": v for k, v in eval_metrics.items()})
-                    log_dict["timing/eval"] = all_timings["eval"]
-
-            log_dict.update({"train/epoch": current_epoch, "train/global_step": self.global_step})
-            # Callbacks may mutate log_dict in place via on_log.
-            self._fire("on_log", logs=log_dict)
-            self.tracker.log(log_dict, step=self.global_step, commit=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).
+        n_examples = len(tokenized)
+
+        def _epoch_of(step: int) -> int:
+            return (step * batch_size) // n_examples
+
+        def _slice_examples(step: int) -> list:
+            start_idx = (step * batch_size) % n_examples
+            end_idx = start_idx + batch_size
+            if end_idx > n_examples:
+                return tokenized[start_idx:] + tokenized[: end_idx - n_examples]
+            return tokenized[start_idx:end_idx]
+
+        def _collate_batch(step: int) -> TrainingInputBatch:
+            return self.collator(_slice_examples(step), batch_size=batch_size)
+
+        collate_ahead_enabled = self.sft_cfg.async_batch_collation
+        async_collator: Optional[AsyncBatchCollator] = (
+            AsyncBatchCollator(_collate_batch, thread_name_prefix="sft-batch-collate")
+            if collate_ahead_enabled
+            else None
+        )
+        collation_state = "ENABLED" if collate_ahead_enabled else "disabled"
+        logger.info(f"SFT async batch collation (double-buffering): {collation_state}")
+
+        # 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
+            reshuffle_after_step = _epoch_of(step) > cur_epoch
+            return not reshuffle_after_step
 
-            if self.global_step % 5 == 0:
-                logger.info(
-                    f"Step {self.global_step}: loss={step_result['loss']:.4f}, " f"grad_norm={step_result['grad_norm']}"
+        try:
+            while self.global_step <= num_steps:
+                all_timings: dict[str, float] = {}
+
+                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 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)
+
+                    self._fire("on_step_start", batch=batch)
+
+                    # Training step
+                    step_result = self.train_step(batch, self.global_step)
+                    all_timings.update(step_result["timings"])
+
+                # Compute throughput using actual (non-padding) tokens
+                batch_padded_seq_len = batch["sequences"].shape[1]
+                actual_num_tokens = batch["attention_mask"].sum().item()
+                self._total_tokens_processed += actual_num_tokens
+                tokens_per_second = actual_num_tokens / all_timings["step"]
+
+                # Build log dict
+                log_dict = {
+                    "train/loss": step_result["loss"],
+                    "train/grad_norm": step_result["grad_norm"],
+                    "train/tokens_per_second": tokens_per_second,
+                    "train/tokens_per_second_per_gpu": tokens_per_second / self._num_training_gpus,
+                    "train/actual_num_tokens": actual_num_tokens,
+                    "train/batch_padded_seq_len": batch_padded_seq_len,
+                    "train/total_tokens_processed": self._total_tokens_processed,
+                }
+                log_dict.update({f"timing/{k}": v for k, v in all_timings.items()})
+                if self._ray_gpu_monitor is not None:
+                    log_dict.update(self._ray_gpu_monitor.flush())
+
+                self._fire("on_step_end", batch=batch, metrics=step_result)
+
+                # Capture callback-driven triggers, then reset so they only fire once.
+                force_save = self._training_control.should_save
+                force_eval = self._training_control.should_evaluate
+                self._training_control.should_save = False
+                self._training_control.should_evaluate = False
+
+                # Checkpoint: interval-driven or callback-requested.
+                interval_save = (
+                    self.sft_cfg.ckpt_interval > 0
+                    and self.global_step > 0
+                    and self.global_step % self.sft_cfg.ckpt_interval == 0
                 )
+                did_save_last_step = force_save or interval_save
+                if did_save_last_step:
+                    with Timer("save_checkpoint", all_timings):
+                        ckpt_path = self.save_checkpoint()
+                    log_dict["timing/save_checkpoint"] = all_timings["save_checkpoint"]
+                    self._fire("on_save", ckpt_path=ckpt_path)
+
+                # HF export at regular intervals
+                if self.sft_cfg.hf_save_interval > 0 and self.global_step % self.sft_cfg.hf_save_interval == 0:
+                    with Timer("save_hf_model", all_timings):
+                        self.save_hf_model()
+                    log_dict["timing/save_hf_model"] = all_timings["save_hf_model"]
+
+                eval_metrics = None
+                num_eval_batches: int | None = None
+                # Eval fires at step N where N % eval_interval == 0 and N > 0, OR
+                # whenever a callback set ``control.should_evaluate``.
+                interval_eval = self.sft_cfg.eval_interval > 0 and self.global_step % self.sft_cfg.eval_interval == 0
+                if eval_tokenized is not None and (force_eval or interval_eval):
+                    self._fire("on_eval_start")
+                    with Timer("eval", all_timings):
+                        eval_metrics, num_eval_batches = self.run_eval(eval_tokenized)
+                    self._fire("on_eval_end", metrics=eval_metrics)
+                    if eval_metrics:
+                        log_dict.update({f"eval/{k}": v for k, v in eval_metrics.items()})
+                        log_dict["timing/eval"] = all_timings["eval"]
+
+                log_dict.update({"train/epoch": current_epoch, "train/global_step": self.global_step})
+                # Callbacks may mutate log_dict in place via on_log.
+                self._fire("on_log", logs=log_dict)
+                self.tracker.log(log_dict, step=self.global_step, commit=True)
+
+                if self.global_step % 5 == 0:
+                    logger.info(
+                        f"Step {self.global_step}: loss={step_result['loss']:.4f}, "
+                        f"grad_norm={step_result['grad_norm']}"
+                    )
 
-            if eval_metrics:
-                logger.info(
-                    f"Step {self.global_step}: eval_loss={eval_metrics.get('eval_loss', float('nan')):.4f} "
-                    f"over {num_eval_batches} batches"
-                )
+                if eval_metrics:
+                    logger.info(
+                        f"Step {self.global_step}: eval_loss={eval_metrics.get('eval_loss', float('nan')):.4f} "
+                        f"over {num_eval_batches} batches"
+                    )
 
-            # Check for epoch boundary and reshuffle
-            epoch = (self.global_step * batch_size) // len(tokenized)
-            if epoch > current_epoch:
-                self._fire("on_epoch_end")
-                epoch_in_progress = False
-                for _ in range(epoch - current_epoch):
-                    rng.shuffle(tokenized)
-                current_epoch = epoch
-                self._current_epoch = epoch
-                if self.global_step + 1 <= num_steps:
-                    self._fire("on_epoch_start")
-                    epoch_in_progress = True
-
-            self.global_step += 1
+                # Check for epoch boundary and reshuffle. Uses the single-sourced
+                # ``_epoch_of`` predicate (``tokenized`` is shuffled in place so
+                # ``len(tokenized) == n_examples`` is invariant).
+                epoch = _epoch_of(self.global_step)
+                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.
+                    if async_collator is not None:
+                        async_collator.clear()
+                    for _ in range(epoch - current_epoch):
+                        rng.shuffle(tokenized)
+                    current_epoch = epoch
+                    self._current_epoch = epoch
+                    if self.global_step + 1 <= num_steps:
+                        self._fire("on_epoch_start")
+                        epoch_in_progress = True
+
+                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()
         self.global_step = min(self.global_step, num_steps)
 
         # Pair the leading on_epoch_start: fire on_epoch_end if we exited the