[feat][SFT] Support multiple datasets for training and evaluation in SFTTrainer

examples/train/sft/README.md

@@ -6,7 +6,7 @@ This example demonstrates supervised fine-tuning using SkyRL, with support for b
 
 By default, the example uses the [Alpaca-Cleaned](https://huggingface.co/datasets/yahma/alpaca-cleaned) dataset (`yahma/alpaca-cleaned`). No manual download is required -- the dataset is loaded automatically via HuggingFace `datasets`.
 
-You can switch to a different dataset by overriding `dataset_name` and `dataset_split` on the command line.
+You can switch to a different dataset by overriding `train_datasets` and `train_dataset_splits` on the command line, or train on a weighted mixture of several datasets (see [Multi-dataset training](#multi-dataset-training-and-evaluation)). The singular `dataset_name`/`dataset_split` (and `eval_dataset_name`/`eval_dataset_split`) fields are deprecated: they still work but emit a `DeprecationWarning` and are translated to the list form internally.
 
 ## Quickstart
 
@@ -85,8 +85,14 @@ All SFT configuration is defined in [`skyrl/train/config/sft_config.py`](../../.
 |-----------|---------|-------------|
 | `strategy` | `megatron` | Backend: `megatron` or `fsdp` |
 | `model.path` | `Qwen/Qwen3-0.6B` | HuggingFace model ID or local path |
-| `dataset_name` | `yahma/alpaca-cleaned` | HuggingFace dataset name |
-| `dataset_split` | `train[:100]` | Dataset split/slice |
+| `train_datasets` | `[yahma/alpaca-cleaned]` | List of HuggingFace dataset names to train on; multiple entries are mixed per `train_dataset_weights` |
+| `train_dataset_splits` | `[train[:100]]` | Split/slice per training dataset (same length as `train_datasets`) |
+| `train_dataset_weights` | equal (`1/N`) | Per-dataset sampling ratios within a batch, independent of dataset sizes; `sampler=random` only |
+| `eval_datasets` | `None` | List of eval dataset names; `None` disables eval. Metrics logged under `eval/{name}/` |
+| `eval_dataset_splits` | `None` | Split per eval dataset (same length as `eval_datasets`) |
+| `eval_dataset_names` | dataset names | Shorthand names used only for logging (`eval/{name}/loss`); must be unique |
+| `dataset_name` / `dataset_split` | deprecated | Use `train_datasets` / `train_dataset_splits` (still accepted with a `DeprecationWarning`) |
+| `eval_dataset_name` / `eval_dataset_split` | deprecated | Use `eval_datasets` / `eval_dataset_splits` |
 | `max_length` | `None` | Maximum sequence length |
 | `num_steps` | `None` | Number of training steps |
 | `batch_size` | `4` | Global batch size |
@@ -158,6 +164,9 @@ Three sampler strategies are built in:
 - `sampler=random` (default) — reshuffles every epoch using `seed`. The
   in-progress epoch resumes bit-exactly; later epochs are re-shuffled into a
   valid (but not byte-identical) order, matching the RL trainer's behavior.
+  With multiple `train_datasets` this strategy switches to the core
+  [`DataMixingSampler`](../../../skyrl/train/dataset/samplers.py) (weighted
+  per-source mixing, fresh plan each epoch, bit-exact resume via RNG state).
 - `sampler=sequential` — iterates the dataset in order
   ([`StatefulSequentialSampler`](../../../skyrl/train/dataset/samplers.py)).
 - `sampler=custom` — loads your own stateful sampler from `sampler_class_path`,
@@ -188,24 +197,75 @@ bash examples/train/sft/run_sft_megatron.sh \
     'sampler_kwargs={lengths: [34, 33, 33], num_samples: 40, seed: 42}'
 ```
 
-### Data mixing example
+## Multi-dataset training and evaluation
 
-[`data_mixing_sampler.py`](data_mixing_sampler.py) is a reference custom sampler
-(`DataMixingSampler`) that mixes a concatenation of sources with per-source
-`weights`, using torch's native `WeightedRandomSampler` for the weighted draw.
-Each source's weight is divided across its examples, so the source-level mixing
-proportion matches `weights` regardless of how many examples each source has.
-Order the dataset by source and give `lengths`/`weights` per source:
+Training on a weighted mixture of datasets is built in: pass lists to
+`train_datasets`/`train_dataset_splits` (note the quoting -- OmegaConf list
+overrides must be a single shell argument):
 
 ```bash
+bash examples/train/sft/run_sft_megatron_multi_dataset.sh
+# or directly:
 bash examples/train/sft/run_sft_megatron.sh \
+    "train_datasets=['allenai/tulu-3-sft-mixture','yahma/alpaca-cleaned']" \
+    "train_dataset_splits=['train[:50000]','train[:10000]']" \
+    "train_dataset_weights=[0.8,0.2]"
+```
+
+Each dataset is tokenized (and cached) independently, then concatenated. With
+multiple datasets and the default `sampler=random`, batches are drawn by the
+core [`DataMixingSampler`](../../../skyrl/train/dataset/samplers.py):
+`train_dataset_weights` gives the approximate per-batch ratio of samples from
+each source, **independent of the dataset sizes** (a 1k-example dataset with
+weight 0.5 contributes half of every batch even when mixed with a 1M-example
+one). Weights default to equal mixing (`1/N`). The sampler draws a fresh
+weighted plan every epoch and stores its RNG state in the checkpoint, so
+`resume_from` reproduces the exact sample stream. All datasets must share the
+same `messages_key`/`tools_key`/`system_key` columns and modality (all-text or
+all-image).
+
+Evaluation similarly accepts multiple datasets, each evaluated separately with
+its loss logged under `eval/{name}/loss`:
+
+```bash
+    "eval_datasets=['allenai/tulu-3-sft-mixture','yahma/alpaca-cleaned']" \
+    "eval_dataset_splits=['train[-500:]','train[200:700]']" \
+    "eval_dataset_names=[tulu3,alpaca]"
+```
+
+`eval_dataset_names` is optional (defaults to the dataset name with `/`
+replaced by `_`) and only affects metric names. **Note (breaking):** eval
+metrics are now always nested per dataset -- the former `eval/eval_loss` key is
+`eval/{name}/loss`, even with a single eval dataset.
+
+### Custom samplers over multiple datasets
+
+A custom sampler still controls the mixture however it likes (e.g. mixing
+weights that change as training progresses). With multiple `train_datasets`,
+the trainer injects the tokenized per-dataset `lengths` into `sampler_kwargs`
+(user-supplied `lengths` win), so the constructor must accept a `lengths`
+kwarg. Explicit `train_dataset_weights` are rejected outside `sampler=random`
+-- pass your ratios through `sampler_kwargs` instead:
+
+```bash
+bash examples/train/sft/run_sft_megatron.sh \
+    "train_datasets=['allenai/tulu-3-sft-mixture','yahma/alpaca-cleaned']" \
+    "train_dataset_splits=['train[:80]','train[:20]']" \
     sampler=custom \
-    sampler_class_path=examples.train.sft.data_mixing_sampler.DataMixingSampler \
-    'sampler_kwargs={lengths: [80, 20], weights: [0.5, 0.5], num_samples: 40, seed: 42}'
+    sampler_class_path=examples.train.sft.curriculum_sampler.CurriculumLearningSampler \
+    'sampler_kwargs={num_samples: 40, seed: 42}'
 ```
 
+(Here the curriculum sampler treats the concatenated datasets as its
+difficulty-ordered stages via the injected `lengths`.)
+
+[`data_mixing_sampler.py`](data_mixing_sampler.py) now simply re-exports the
+core `DataMixingSampler`, so existing `sampler_class_path` configs pointing at
+it keep working. Behavior change vs. the old example: the core sampler
+re-plans every epoch instead of replaying one fixed plan (old position-only
+checkpoints resume best-effort with a warning).
+
 ## Limitations
 
 - **Limited `train_on_what` options**: Supports training on all or the last assistant message.
 - **Two data formats only.** Supports chat-template (`messages` column) and Alpaca (`instruction`/`output` columns). Raw pre-tokenized or plain-text continuation formats are not supported.
-- **Single dataset.** No built-in multi-dataset mixing or weighting. Only one `dataset_name` + `dataset_split` pair can be specified.

examples/train/sft/data_mixing_sampler.py

@@ -1,118 +1,25 @@
-"""Reference custom sampler: weighted data mixing for ``SFTTrainer``.
+"""Deprecated location: ``DataMixingSampler`` is now part of core SkyRL.
 
-Like ``curriculum_sampler.py``, this is an *example* of a user-supplied stateful
-sampler plugged into ``SFTTrainer`` via ``sampler=custom`` -- it is intentionally
-NOT part of the core library. Point ``sampler_class_path`` at this class via a
-dotted path importable from the repo root (run from the repo root; no
-``__init__.py`` needed thanks to namespace packages), passing the sampler
-config as overrides on the base SFT example script::
+Weighted multi-dataset mixing no longer requires a custom sampler -- configure it
+natively via the list-based dataset fields, e.g.::
 
     bash examples/train/sft/run_sft_megatron.sh \
-        sampler=custom \
-        sampler_class_path=examples.train.sft.data_mixing_sampler.DataMixingSampler \
-        'sampler_kwargs={lengths: [80, 20], weights: [0.5, 0.5], num_samples: 40, seed: 42}'
-
-(``lengths`` must sum to the dataset size -- 100 for the script's
-``train[:100]`` split -- and ``num_samples`` should be
-``num_steps * batch_size``.)
-
-Note: the import happens inside a Ray task, which does not inherit the driver's
-``PYTHONPATH`` -- so the path must resolve from the worker's ``sys.path``
-(which includes the repo root when launching from it).
-
-It mixes a concatenation of sources (``ConcatDataset([src_a, src_b, ...])``)
-according to per-source ``weights``, using torch's native
-``WeightedRandomSampler`` for the weighted draw. Because each source's weight is
-divided across its examples, the *source-level* mixing proportion matches
-``weights`` regardless of how many examples each source has.
+        'train_datasets=[allenai/tulu-3-sft-mixture,yahma/alpaca-cleaned]' \
+        'train_dataset_splits=[train[:50000],train[:10000]]' \
+        'train_dataset_weights=[0.8,0.2]'
+
+With multiple ``train_datasets`` and the default ``sampler=random``, the trainer
+uses :class:`skyrl.train.dataset.samplers.DataMixingSampler` automatically,
+configured with the tokenized per-dataset lengths and ``train_dataset_weights``.
+
+This module re-exports the core class so existing configs pointing
+``sampler_class_path`` here keep working. Note a behavior change from the old
+example: the core sampler draws a *fresh* weighted plan every epoch (the example
+replayed one fixed plan), and its ``state_dict`` now includes the generator
+state so mid-epoch checkpoint resume is exact. Old position-only checkpoints
+load with a warning.
 """
 
-from __future__ import annotations
-
-from typing import Iterator, List, Optional, Sequence
-
-import torch
-
-
-class DataMixingSampler(torch.utils.data.Sampler[int]):
-    """Weighted multi-source sampler built on ``WeightedRandomSampler``.
-
-    The dataset is a concatenation of sources with sizes ``lengths`` (in order);
-    ``weights`` gives a sampling weight per source. Per-example weights are set to
-    ``weight_source / size_source`` so each *source* is sampled in proportion to
-    its weight independent of its size. A deterministic plan of ``num_samples``
-    indices is materialized up front from ``seed`` via
-    ``torch.utils.data.WeightedRandomSampler``, so iteration is reproducible and a
-    single ``position`` cursor is sufficient for checkpoint/resume.
-
-    Args:
-        data_source: The (concatenated) training dataset; only its length is used.
-        lengths: Size of each source, in the order they appear in the dataset.
-        weights: Per-source sampling weight (need not sum to 1; relative scale matters).
-        num_samples: Total number of indices to emit. Set to
-            ``num_steps * batch_size`` to cover the whole run in one pass (keeps
-            the sampler state intact across epoch boundaries and makes resume
-            bit-exact across the entire run). Defaults to ``len(data_source)``.
-        seed: Seed for the deterministic weighted draw.
-        replacement: Whether to sample with replacement (passed through to
-            ``WeightedRandomSampler``; mixing across sources generally wants True).
-    """
-
-    def __init__(
-        self,
-        data_source: Sequence,
-        lengths: Sequence[int],
-        weights: Sequence[float],
-        num_samples: Optional[int] = None,
-        seed: int = 0,
-        replacement: bool = True,
-    ):
-        self.data_source = data_source
-        n = len(data_source)
-        if sum(lengths) != n:
-            raise ValueError(f"DataMixingSampler: sum(lengths)={sum(lengths)} must equal len(data_source)={n}.")
-        if len(weights) != len(lengths):
-            raise ValueError(f"DataMixingSampler: weights ({len(weights)}) and lengths ({len(lengths)}) must align.")
-        if any(length <= 0 for length in lengths):
-            raise ValueError(f"DataMixingSampler: all lengths must be > 0, got {list(lengths)}.")
-        if any(w < 0 for w in weights) or sum(weights) <= 0:
-            raise ValueError(
-                f"DataMixingSampler: weights must be non-negative with a positive sum, got {list(weights)}."
-            )
-
-        self.num_samples = num_samples if num_samples is not None else n
-        if self.num_samples <= 0:
-            raise ValueError(f"DataMixingSampler: num_samples must be > 0, got {self.num_samples}.")
-        self.position = 0
-
-        # Spread each source's weight across its examples so the source-level
-        # mixing proportion matches ``weights`` regardless of source size.
-        per_example_weights: List[float] = []
-        for length, weight in zip(lengths, weights):
-            per_example_weights.extend([weight / length] * length)
-
-        generator = torch.Generator()
-        generator.manual_seed(seed)
-        weighted = torch.utils.data.WeightedRandomSampler(
-            per_example_weights,
-            num_samples=self.num_samples,
-            replacement=replacement,
-            generator=generator,
-        )
-        self._plan: List[int] = list(weighted)
-
-    def __iter__(self) -> Iterator[int]:
-        while self.position < len(self._plan):
-            idx = self._plan[self.position]
-            self.position += 1
-            yield idx
-        self.position = 0
-
-    def __len__(self) -> int:
-        return len(self._plan)
-
-    def state_dict(self) -> dict:
-        return {"position": self.position}
+from skyrl.train.dataset.samplers import DataMixingSampler
 
-    def load_state_dict(self, state: dict) -> None:
-        self.position = state["position"]
+__all__ = ["DataMixingSampler"]

examples/train/sft/run_sft_fsdp.sh

@@ -16,8 +16,8 @@ uv run --isolated --extra fsdp \
     python -m skyrl.train.main_sft \
     strategy=fsdp \
     model.path=Qwen/Qwen2.5-0.5B-Instruct \
-    dataset_name=yahma/alpaca-cleaned \
-    dataset_split="train[:100]" \
+    "train_datasets=['yahma/alpaca-cleaned']" \
+    "train_dataset_splits=['train[:100]']" \
     messages_key=messages \
     max_length=512 \
     num_steps=10 \

examples/train/sft/run_sft_lora.sh

@@ -21,8 +21,8 @@ uv run --isolated --extra fsdp \
     model.lora.rank=32 \
     model.lora.alpha=16 \
     model.lora.target_modules=all-linear \
-    dataset_name=yahma/alpaca-cleaned \
-    dataset_split="train[:100]" \
+    "train_datasets=['yahma/alpaca-cleaned']" \
+    "train_dataset_splits=['train[:100]']" \
     messages_key=messages \
     max_length=512 \
     num_steps=10 \

examples/train/sft/run_sft_megatron.sh

@@ -18,8 +18,8 @@ uv run --isolated --extra megatron \
     python -m skyrl.train.main_sft \
     strategy=megatron \
     model.path=Qwen/Qwen2.5-0.5B-Instruct \
-    dataset_name=yahma/alpaca-cleaned \
-    dataset_split="train[:100]" \
+    "train_datasets=['yahma/alpaca-cleaned']" \
+    "train_dataset_splits=['train[:100]']" \
     messages_key=messages \
     max_length=512 \
     num_steps=10 \

examples/train/sft/run_sft_megatron_apigen_mt.sh

@@ -18,10 +18,10 @@ uv run --isolated --extra megatron --python 3.12 \
     python -m skyrl.train.main_sft \
     strategy=megatron \
     model.path=Qwen/Qwen2.5-1.5B-Instruct \
-    dataset_name="$DATA_DIR" \
-    dataset_split="train[:4000]" \
-    eval_dataset_name="$DATA_DIR" \
-    eval_dataset_split="train[4000:]" \
+    "train_datasets=['$DATA_DIR']" \
+    "train_dataset_splits=['train[:4000]']" \
+    "eval_datasets=['$DATA_DIR']" \
+    "eval_dataset_splits=['train[4000:]']" \
     messages_key=messages \
     tools_key=tools \
     system_key=system \

examples/train/sft/run_sft_megatron_multi_dataset.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+set -x
+
+# SFT training on a weighted mixture of multiple datasets (Megatron backend).
+#
+# Trains on Tulu3 + Alpaca-Cleaned mixed 80/20 per batch (independent of the
+# dataset sizes) via the core DataMixingSampler, and evaluates on held-out
+# slices of both datasets separately. Eval metrics are logged per dataset under
+# eval/{name}/loss (names from eval_dataset_names).
+#
+# Note the quoting: OmegaConf list overrides must be a single shell argument.
+#
+# Usage:
+#   bash examples/train/sft/run_sft_megatron_multi_dataset.sh [extra overrides...]
+#
+# Example:
+#   bash examples/train/sft/run_sft_megatron_multi_dataset.sh num_steps=20 'train_dataset_weights=[0.5,0.5]'
+
+uv run --isolated --extra megatron \
+    python -m skyrl.train.main_sft \
+    strategy=megatron \
+    model.path=Qwen/Qwen2.5-0.5B-Instruct \
+    "train_datasets=['allenai/tulu-3-sft-mixture','yahma/alpaca-cleaned']" \
+    "train_dataset_splits=['train[:800]','train[:200]']" \
+    "train_dataset_weights=[0.8,0.2]" \
+    "eval_datasets=['allenai/tulu-3-sft-mixture','yahma/alpaca-cleaned']" \
+    "eval_dataset_splits=['train[-100:]','train[200:300]']" \
+    "eval_dataset_names=[tulu3,alpaca]" \
+    eval_interval=5 \
+    messages_key=messages \
+    max_length=512 \
+    num_steps=10 \
+    batch_size=4 \
+    micro_train_batch_size_per_gpu=1 \
+    remove_microbatch_padding=true \
+    seed=42 \
+    optimizer_config.lr=1e-6 \
+    optimizer_config.weight_decay=1e-2 \
+    optimizer_config.max_grad_norm=1.0 \
+    optimizer_config.num_warmup_steps=0 \
+    optimizer_config.scheduler=constant_with_warmup \
+    placement.num_nodes=1 \
+    placement.num_gpus_per_node=4 \
+    megatron_config.tensor_model_parallel_size=1 \
+    megatron_config.pipeline_model_parallel_size=1 \
+    megatron_config.context_parallel_size=1 \
+    logger=console \
+    project_name=skyrl_sft \
+    run_name=skyrl_sft_megatron_multi_dataset_run \
+    ckpt_path="" \
+    ckpt_interval=0 \
+    resume_from="" \
+    "$@"

examples/train/sft/run_sft_megatron_tulu3_50k.sh

@@ -19,10 +19,10 @@ uv run --isolated --extra megatron \
     python -m skyrl.train.main_sft \
     strategy=megatron \
     model.path=Qwen/Qwen2.5-0.5B-Instruct \
-    dataset_name=allenai/tulu-3-sft-mixture \
-    dataset_split="train[:50000]" \
-    eval_dataset_name=allenai/tulu-3-sft-mixture \
-    eval_dataset_split="train[-500:]" \
+    "train_datasets=['allenai/tulu-3-sft-mixture']" \
+    "train_dataset_splits=['train[:50000]']" \
+    "eval_datasets=['allenai/tulu-3-sft-mixture']" \
+    "eval_dataset_splits=['train[-500:]']" \
     messages_key=messages \
     max_length=4096 \
     num_steps=4166 \

examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh

@@ -28,10 +28,10 @@ uv run --isolated --extra megatron \
     python -m skyrl.train.main_sft \
     strategy=megatron \
     model.path=Qwen/Qwen2.5-0.5B-Instruct \
-    dataset_name=allenai/tulu-3-sft-mixture \
-    dataset_split="train[:50000]" \
-    eval_dataset_name=allenai/tulu-3-sft-mixture \
-    eval_dataset_split="train[-500:]" \
+    "train_datasets=['allenai/tulu-3-sft-mixture']" \
+    "train_dataset_splits=['train[:50000]']" \
+    "eval_datasets=['allenai/tulu-3-sft-mixture']" \
+    "eval_dataset_splits=['train[-500:]']" \
     messages_key=messages \
     max_length=4096 \
     num_steps=4166 \