From 306b412cebb83a35a2897e1792eb6450e0ed47f2 Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Fri, 10 Jul 2026 22:06:35 +0000 Subject: [PATCH 1/6] [feat][SFT] Support multiple datasets for training and evaluation in SFTTrainer Implements RFC #1875: native SFTTrainer support for training on a weighted mixture of datasets and evaluating on multiple named datasets. - Config: add train_datasets/train_dataset_splits/train_dataset_weights and eval_datasets/eval_dataset_splits/eval_dataset_names. The singular dataset_name/dataset_split and eval_dataset_name/eval_dataset_split fields are deprecated: still accepted, translated to the list form with a DeprecationWarning, and rejected when combined with the new fields. - Sampler: promote DataMixingSampler from examples into skyrl.train.dataset.samplers. The core version draws a fresh weighted plan every epoch (the example replayed one fixed plan) and checkpoints its RNG state, so mid-epoch and epoch-boundary resumes reproduce the uninterrupted sample stream exactly. The example file re-exports the core class for backward compatibility with existing sampler_class_path configs. - Trainer: tokenize each training dataset independently (per-dataset cache keys preserved) and concatenate; with multiple datasets, sampler=random selects DataMixingSampler configured with the tokenized lengths and train_dataset_weights (per-batch source ratios, independent of dataset sizes, default 1/N). Custom samplers receive the per-dataset lengths via an injected `lengths` kwarg (user-supplied kwargs win). - Eval: one dataloader per eval dataset, each evaluated separately. Metrics are logged under eval/{name}/loss, nested even for a single eval dataset. BREAKING: the former eval/eval_loss key is now eval/{name}/loss. - Tests: CPU coverage for config normalization/validation, the core sampler (fresh plan per epoch, exact resume, old-format state), sampler dispatch, and the RFC mixing-ratio test plan (user weights, default equal mixing, size-independence). - Docs/examples: README multi-dataset section, new run_sft_megatron_multi_dataset.sh, and example/e2e scripts migrated to the list-based fields. Closes #1875 Co-Authored-By: Claude Fable 5 Signed-off-by: Avi Basnet --- examples/train/sft/README.md | 86 +++++- examples/train/sft/data_mixing_sampler.py | 131 ++------- examples/train/sft/run_sft_fsdp.sh | 4 +- examples/train/sft/run_sft_lora.sh | 4 +- examples/train/sft/run_sft_megatron.sh | 4 +- .../train/sft/run_sft_megatron_apigen_mt.sh | 8 +- .../sft/run_sft_megatron_multi_dataset.sh | 53 ++++ .../train/sft/run_sft_megatron_tulu3_50k.sh | 8 +- .../sft/run_sft_megatron_tulu3_50k_packed.sh | 8 +- examples/train/sft/run_sft_megatron_vlm.sh | 4 +- skyrl/train/config/sft_config.py | 167 ++++++++++- skyrl/train/dataset/samplers.py | 130 ++++++++- skyrl/train/sft_trainer.py | 213 ++++++++++---- .../megatron/test_sft_packing_parity.py | 3 +- .../train/gpu_e2e_test/sft_tulu3_megatron.sh | 2 +- tests/train/test_sft_callbacks.py | 24 +- tests/train/test_sft_config.py | 134 ++++++++- tests/train/test_sft_dataloader.py | 270 ++++++++++++++---- 18 files changed, 977 insertions(+), 276 deletions(-) create mode 100755 examples/train/sft/run_sft_megatron_multi_dataset.sh diff --git a/examples/train/sft/README.md b/examples/train/sft/README.md index d787d18347..0a506b48d1 100644 --- a/examples/train/sft/README.md +++ b/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. diff --git a/examples/train/sft/data_mixing_sampler.py b/examples/train/sft/data_mixing_sampler.py index c6bd1e313a..4c90dc278d 100644 --- a/examples/train/sft/data_mixing_sampler.py +++ b/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"] diff --git a/examples/train/sft/run_sft_fsdp.sh b/examples/train/sft/run_sft_fsdp.sh index e94530539a..3727a36e48 100755 --- a/examples/train/sft/run_sft_fsdp.sh +++ b/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 \ diff --git a/examples/train/sft/run_sft_lora.sh b/examples/train/sft/run_sft_lora.sh index dc0de1373b..220b87da32 100644 --- a/examples/train/sft/run_sft_lora.sh +++ b/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 \ diff --git a/examples/train/sft/run_sft_megatron.sh b/examples/train/sft/run_sft_megatron.sh index cb6a98caf3..83aef32513 100755 --- a/examples/train/sft/run_sft_megatron.sh +++ b/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 \ diff --git a/examples/train/sft/run_sft_megatron_apigen_mt.sh b/examples/train/sft/run_sft_megatron_apigen_mt.sh index bfc9659a44..e6f8ec5efe 100755 --- a/examples/train/sft/run_sft_megatron_apigen_mt.sh +++ b/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 \ diff --git a/examples/train/sft/run_sft_megatron_multi_dataset.sh b/examples/train/sft/run_sft_megatron_multi_dataset.sh new file mode 100755 index 0000000000..6cff607edb --- /dev/null +++ b/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="" \ + "$@" diff --git a/examples/train/sft/run_sft_megatron_tulu3_50k.sh b/examples/train/sft/run_sft_megatron_tulu3_50k.sh index 35be6deff2..9813d6c6d9 100755 --- a/examples/train/sft/run_sft_megatron_tulu3_50k.sh +++ b/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 \ diff --git a/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh b/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh index 342080d660..41635ef6e5 100755 --- a/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh +++ b/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 \ diff --git a/examples/train/sft/run_sft_megatron_vlm.sh b/examples/train/sft/run_sft_megatron_vlm.sh index 84f98ba5fb..2563d6d0c9 100755 --- a/examples/train/sft/run_sft_megatron_vlm.sh +++ b/examples/train/sft/run_sft_megatron_vlm.sh @@ -33,8 +33,8 @@ uv run --isolated --extra megatron \ python -m skyrl.train.main_sft \ strategy=megatron \ model.path=Qwen/Qwen3-VL-2B-Instruct \ - dataset_name="$DATA_DIR" \ - dataset_split="train" \ + "train_datasets=[$DATA_DIR]" \ + "train_dataset_splits=[train]" \ messages_key=messages \ max_length=4096 \ num_steps=30 \ diff --git a/skyrl/train/config/sft_config.py b/skyrl/train/config/sft_config.py index 118c75c4c8..c9a44be7b1 100644 --- a/skyrl/train/config/sft_config.py +++ b/skyrl/train/config/sft_config.py @@ -147,8 +147,23 @@ def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig": # ---- SFT-specific flat fields ---- strategy: str = "megatron" # "megatron" or "fsdp" - dataset_name: str = "yahma/alpaca-cleaned" - dataset_split: str = "train[:100]" + dataset_name: Optional[str] = None + """Deprecated: use ``train_datasets`` instead. Translated to ``train_datasets=[dataset_name]`` + with a DeprecationWarning. Cannot be combined with ``train_datasets``.""" + dataset_split: Optional[str] = None + """Deprecated: use ``train_dataset_splits`` instead.""" + train_datasets: Optional[List[str]] = None + """HuggingFace dataset names (or paths) to train on. With multiple datasets, batches are + mixed per-source by :class:`~skyrl.train.dataset.samplers.DataMixingSampler` according to + ``train_dataset_weights``. Defaults to ``["yahma/alpaca-cleaned"]``. All datasets must share + the same ``messages_key``/``tools_key``/``system_key`` columns and modality.""" + train_dataset_splits: Optional[List[str]] = None + """Split to load for each entry of ``train_datasets`` (e.g. ``"train[:50000]"``). Must match + ``train_datasets`` in length. Defaults to ``["train[:100]"]``.""" + train_dataset_weights: Optional[List[float]] = None + """Per-dataset sampling weights: the approximate per-batch ratio of samples drawn from each + dataset, independent of dataset sizes. Only supported with ``sampler="random"`` (custom + samplers receive ratios via ``sampler_kwargs``). Defaults to equal mixing (``1/N`` each).""" messages_key: str = "messages" # column name for chat-format datasets tools_key: str = "tools" """Column name holding per-row tool/function schemas for tool-calling datasets @@ -158,12 +173,22 @@ def from_cli_overrides(cls, args: Union[List[str], dict]) -> "SFTConfig": """Column name holding a per-row system prompt to prepend when ``messages`` does not already start with a system turn. Ignored if absent.""" - # ---- Evaluation dataset ---- + # ---- Evaluation datasets ---- eval_dataset_name: Optional[str] = None - """HuggingFace dataset name (or path) used to compute eval loss during training. - When ``None`` (default), eval is disabled.""" - eval_dataset_split: str = "validation" - """Split of the eval dataset to load (e.g. ``"validation"``, ``"test[:500]"``).""" + """Deprecated: use ``eval_datasets`` instead. Translated to ``eval_datasets=[eval_dataset_name]`` + with a DeprecationWarning. Cannot be combined with ``eval_datasets``.""" + eval_dataset_split: Optional[str] = None + """Deprecated: use ``eval_dataset_splits`` instead.""" + eval_datasets: Optional[List[str]] = None + """HuggingFace dataset names (or paths) used to compute eval loss during training. + When ``None`` (default), eval is disabled. Metrics are logged per dataset under + ``eval/{name}/`` (nested even with a single eval dataset).""" + eval_dataset_splits: Optional[List[str]] = None + """Split to load for each entry of ``eval_datasets`` (e.g. ``"validation"``, ``"test[:500]"``). + Must match ``eval_datasets`` in length. Defaults to ``["validation"]`` on the deprecated path.""" + eval_dataset_names: Optional[List[str]] = None + """Optional shorthand names used only for logging (``eval/{name}/loss``). Must be unique and + match ``eval_datasets`` in length. Defaults to each dataset name with ``/`` replaced by ``_``.""" eval_interval: int = 0 """Run eval every N training steps. Eval also runs once at the end of training when an eval dataset is configured. ``0`` disables periodic eval.""" @@ -286,6 +311,125 @@ def resolved_bin_capacity(self) -> int: _VALID_STRATEGIES = ("megatron", "fsdp") _VALID_SAMPLERS = ("random", "sequential", "custom") +_DEFAULT_TRAIN_DATASET = "yahma/alpaca-cleaned" +_DEFAULT_TRAIN_SPLIT = "train[:100]" +_DEFAULT_EVAL_SPLIT = "validation" + + +def _normalize_dataset_cfg(cfg: SFTConfig) -> None: + """Translate the deprecated single-dataset fields into the list-based fields + and validate the dataset configuration. + + Idempotent, so it is safe to call from both ``validate_sft_cfg`` and + ``SFTTrainer.__init__`` (the latter covers trainers constructed with a + pre-built ``skyrl_cfg``, which bypass ``build_skyrl_config_for_sft``). + + Post-conditions: + + - ``train_datasets``/``train_dataset_splits`` are equal-length non-empty lists; + ``train_dataset_weights`` too when ``sampler="random"`` (``None`` otherwise). + - ``eval_datasets`` is ``None`` (eval disabled) or an equal-length triple with + ``eval_dataset_splits``/``eval_dataset_names`` (names unique). + - The deprecated ``dataset_name``/``dataset_split``/``eval_dataset_name``/ + ``eval_dataset_split`` fields are ``None``. + """ + import warnings + + # ---- Train datasets ---- + if cfg.train_datasets is not None and (cfg.dataset_name is not None or cfg.dataset_split is not None): + raise ValueError( + "Specify only one of train_datasets/train_dataset_splits and the deprecated " + "dataset_name/dataset_split, not both." + ) + if cfg.train_datasets is None: + if cfg.dataset_name is not None or cfg.dataset_split is not None: + warnings.warn( + "dataset_name/dataset_split are deprecated; use train_datasets/train_dataset_splits instead.", + DeprecationWarning, + stacklevel=2, + ) + cfg.train_datasets = [cfg.dataset_name if cfg.dataset_name is not None else _DEFAULT_TRAIN_DATASET] + if cfg.train_dataset_splits is None: + cfg.train_dataset_splits = [cfg.dataset_split if cfg.dataset_split is not None else _DEFAULT_TRAIN_SPLIT] + cfg.dataset_name = None + cfg.dataset_split = None + + if len(cfg.train_datasets) == 0: + raise ValueError("train_datasets must be a non-empty list.") + if cfg.train_dataset_splits is None or len(cfg.train_dataset_splits) != len(cfg.train_datasets): + raise ValueError( + f"train_dataset_splits must specify one split per entry of train_datasets " + f"({len(cfg.train_datasets)} datasets), got {cfg.train_dataset_splits}." + ) + if cfg.train_dataset_weights is not None: + if cfg.sampler != "random": + raise ValueError( + f"train_dataset_weights is only supported with sampler='random' (weighted mixing via " + f"DataMixingSampler), got sampler='{cfg.sampler}'. Pass mixing ratios to a custom " + f"sampler via sampler_kwargs instead." + ) + if len(cfg.train_dataset_weights) != len(cfg.train_datasets): + raise ValueError( + f"train_dataset_weights must specify one weight per entry of train_datasets " + f"({len(cfg.train_datasets)} datasets), got {cfg.train_dataset_weights}." + ) + if any(w <= 0 for w in cfg.train_dataset_weights): + raise ValueError(f"train_dataset_weights must all be > 0, got {cfg.train_dataset_weights}.") + elif cfg.sampler == "random": + # Default: equal mixing. Left as None for other samplers (sequential + # ignores mixing; custom samplers take ratios via sampler_kwargs). + cfg.train_dataset_weights = [1.0 / len(cfg.train_datasets)] * len(cfg.train_datasets) + + # ---- Eval datasets ---- + if cfg.eval_datasets is not None and (cfg.eval_dataset_name is not None or cfg.eval_dataset_split is not None): + raise ValueError( + "Specify only one of eval_datasets/eval_dataset_splits and the deprecated " + "eval_dataset_name/eval_dataset_split, not both." + ) + if cfg.eval_datasets is None and cfg.eval_dataset_name is not None: + warnings.warn( + "eval_dataset_name/eval_dataset_split are deprecated; use eval_datasets/eval_dataset_splits instead.", + DeprecationWarning, + stacklevel=2, + ) + cfg.eval_datasets = [cfg.eval_dataset_name] + if cfg.eval_dataset_splits is None: + cfg.eval_dataset_splits = [ + cfg.eval_dataset_split if cfg.eval_dataset_split is not None else _DEFAULT_EVAL_SPLIT + ] + cfg.eval_dataset_name = None + cfg.eval_dataset_split = None + + if cfg.eval_datasets is None: + if cfg.eval_dataset_splits is not None or cfg.eval_dataset_names is not None: + raise ValueError("eval_dataset_splits/eval_dataset_names require eval_datasets to be set.") + return + if len(cfg.eval_datasets) == 0: + raise ValueError("eval_datasets must be a non-empty list when set.") + if cfg.eval_dataset_splits is None or len(cfg.eval_dataset_splits) != len(cfg.eval_datasets): + raise ValueError( + f"eval_dataset_splits must specify one split per entry of eval_datasets " + f"({len(cfg.eval_datasets)} datasets), got {cfg.eval_dataset_splits}." + ) + if cfg.eval_dataset_names is None: + names = [name.replace("/", "_") for name in cfg.eval_datasets] + if len(set(names)) != len(names): + raise ValueError( + f"Default eval dataset names collide ({names}), e.g. the same dataset with two " + f"different splits. Set eval_dataset_names explicitly to disambiguate." + ) + cfg.eval_dataset_names = names + else: + if len(cfg.eval_dataset_names) != len(cfg.eval_datasets): + raise ValueError( + f"eval_dataset_names must specify one name per entry of eval_datasets " + f"({len(cfg.eval_datasets)} datasets), got {cfg.eval_dataset_names}." + ) + if any(not name for name in cfg.eval_dataset_names): + raise ValueError(f"eval_dataset_names must be non-empty strings, got {cfg.eval_dataset_names}.") + if len(set(cfg.eval_dataset_names)) != len(cfg.eval_dataset_names): + raise ValueError(f"eval_dataset_names must be unique, got {cfg.eval_dataset_names}.") + def validate_sft_cfg(cfg: SFTConfig) -> None: """Validate SFT-specific configuration. @@ -293,6 +437,7 @@ def validate_sft_cfg(cfg: SFTConfig) -> None: Only checks fields that are relevant to SFT training, unlike ``validate_cfg`` which includes RL-specific validations. """ + _normalize_dataset_cfg(cfg) if cfg.strategy == "fsdp2": import warnings @@ -335,10 +480,10 @@ def validate_sft_cfg(cfg: SFTConfig) -> None: # Eval config if cfg.eval_interval < 0: raise ValueError(f"eval_interval must be >= 0, got {cfg.eval_interval}") - if cfg.eval_interval > 0 and not cfg.eval_dataset_name: - raise ValueError("eval_interval > 0 requires eval_dataset_name to be set") - if cfg.eval_before_train and cfg.eval_dataset_name is None: - raise ValueError("eval_before_train=True requires eval_dataset_name to be set") + if cfg.eval_interval > 0 and not cfg.eval_datasets: + raise ValueError("eval_interval > 0 requires eval_datasets to be set") + if cfg.eval_before_train and cfg.eval_datasets is None: + raise ValueError("eval_before_train=True requires eval_datasets to be set") # checks for megatron if cfg.strategy == "megatron": diff --git a/skyrl/train/dataset/samplers.py b/skyrl/train/dataset/samplers.py index ad8ec16767..6939f9242b 100644 --- a/skyrl/train/dataset/samplers.py +++ b/skyrl/train/dataset/samplers.py @@ -7,7 +7,9 @@ resumes after a checkpoint load. Core ships :class:`StatefulSequentialSampler` (backing the ``sampler="sequential"`` -config option). The ``sampler="custom"`` path loads a user-supplied sampler from +config option) and :class:`DataMixingSampler` (weighted multi-source mixing, used +by default when multiple ``train_datasets`` are configured with ``sampler="random"``). +The ``sampler="custom"`` path loads a user-supplied sampler from ``SFTConfig.sampler_class_path`` via :func:`import_sampler_class`, instantiating it as ``ClassName(tokenized, **sampler_kwargs)``. See ``examples/train/sft/curriculum_sampler.py`` for a reference custom sampler. @@ -16,11 +18,13 @@ from __future__ import annotations import importlib -from typing import Iterator, Sequence +from typing import Iterator, List, Optional, Sequence import torch +from loguru import logger __all__ = [ + "DataMixingSampler", "StatefulSequentialSampler", "import_sampler_class", ] @@ -82,3 +86,125 @@ def state_dict(self) -> dict: def load_state_dict(self, state: dict) -> None: self.position = state["position"] + + +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. + + Each epoch draws a fresh plan of ``num_samples`` indices from a persistent + generator seeded with ``seed``: exhausting the plan clears it, and the next + ``iter()`` (the trainer re-creates the dataloader iterator at epoch + boundaries) re-draws with the advanced generator state. ``state_dict`` + captures the cursor plus the generator state *as of the current plan's + draw*, so a mid-epoch checkpoint resume reproduces the in-flight plan and + all subsequent epochs match the uninterrupted run. + + 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: Number of indices to emit per epoch. Defaults to + ``len(data_source)``, so ``steps_per_epoch`` matches a single + concatenated dataset of the same size. + 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 + self.replacement = replacement + + # 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) + self._per_example_weights = per_example_weights + + self._generator = torch.Generator() + self._generator.manual_seed(seed) + # The current epoch's plan, drawn lazily. ``_plan_gen_state`` snapshots + # the generator *before* the draw so a resumed sampler re-draws the + # identical plan. + self._plan: Optional[List[int]] = None + self._plan_gen_state: Optional[torch.Tensor] = None + + def _ensure_plan(self) -> None: + if self._plan is not None: + return + self._plan_gen_state = self._generator.get_state() + weighted = torch.utils.data.WeightedRandomSampler( + self._per_example_weights, + num_samples=self.num_samples, + replacement=self.replacement, + generator=self._generator, + ) + self._plan = list(weighted) + + def __iter__(self) -> Iterator[int]: + self._ensure_plan() + while self.position < len(self._plan): + idx = self._plan[self.position] + self.position += 1 + yield idx + # Reset for the next epoch: clear the plan so the next ``iter()`` + # draws fresh indices with the advanced generator state. + self.position = 0 + self._plan = None + + def __len__(self) -> int: + return self.num_samples + + def state_dict(self) -> dict: + # Mid-epoch: persist the pre-draw state so resume re-draws the current + # plan. Between epochs: persist the advanced state so the next epoch's + # draw matches the uninterrupted run. + if self._plan is not None: + generator_state = self._plan_gen_state + else: + generator_state = self._generator.get_state() + return {"position": self.position, "generator_state": generator_state} + + def load_state_dict(self, state: dict) -> None: + self.position = state["position"] + if "generator_state" in state: + self._generator.set_state(state["generator_state"]) + else: + # Checkpoint from the old example sampler (position-only): keep the + # freshly-seeded generator and resume the cursor best-effort. + logger.warning( + "DataMixingSampler: checkpoint has no generator_state (old format); " + "resuming position only -- the restored plan may differ from the original run." + ) + self._plan = None + self._plan_gen_state = None diff --git a/skyrl/train/sft_trainer.py b/skyrl/train/sft_trainer.py index 08b955e54e..5ee5f5e726 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -49,6 +49,7 @@ from skyrl.train.config.sft_config import ( SFTConfig, TrainOnWhat, + _normalize_dataset_cfg, build_skyrl_config_for_sft, ) from skyrl.train.generators.utils import ( @@ -779,6 +780,11 @@ def collate_sft_examples( # --------------------------------------------------------------------------- +def _format_eval_metrics(eval_metrics: dict) -> str: + """Render per-dataset eval metrics (``{name}/loss``) for stdout logging.""" + return ", ".join(f"{k}={v:.4f}" for k, v in eval_metrics.items()) + + class SFTTrainer: """SFT trainer supporting FSDP and Megatron backends. @@ -801,6 +807,11 @@ def __init__( callbacks: Optional[list[TrainingCallback]] = None, ): self.sft_cfg = cfg + # Translate deprecated single-dataset fields into the list-based ones. + # Also runs inside validate_sft_cfg (idempotent); calling it here covers + # trainers constructed with a pre-built ``skyrl_cfg``, which skip + # ``build_skyrl_config_for_sft``. + _normalize_dataset_cfg(cfg) # Accept a pre-built bridge config to avoid redundant rebuilds. # When not provided (e.g. standalone usage), build it here. self.cfg = skyrl_cfg if skyrl_cfg is not None else build_skyrl_config_for_sft(cfg) @@ -811,7 +822,10 @@ def __init__( self.tracker: Tracking | None = None # Stateful dataloaders, built in train() once data is tokenized. self.train_dataloader: StatefulDataLoader | None = None - self.eval_dataloader: StatefulDataLoader | None = None + # One ``(name, dataloader)`` pair per configured eval dataset; ``None`` + # when eval is disabled. Names are unique (enforced in config validation) + # and namespace the eval metrics as ``eval/{name}/...``. + self.eval_dataloaders: list[tuple[str, StatefulDataLoader]] | None = None self.global_step = 0 # running count of total non-padding tokens trained on self._total_tokens_processed = 0 @@ -1216,15 +1230,55 @@ def _load_and_tokenize(self, dataset_name: str, dataset_split: str) -> list: shutil.rmtree(tokenizer_cache_dir, ignore_errors=True) - def load_dataset(self) -> list: - """Load and tokenize the training dataset.""" - return self._load_and_tokenize(self.sft_cfg.dataset_name, self.sft_cfg.dataset_split) + def load_dataset(self) -> tuple[list, list[int]]: + """Load and tokenize the training dataset(s). + + Each ``(name, split)`` pair from ``train_datasets``/``train_dataset_splits`` + is tokenized independently through :meth:`_load_and_tokenize` (preserving + per-dataset cache keys), then concatenated in config order. + + Returns: + ``(tokenized, dataset_lengths)`` where ``dataset_lengths`` holds the + tokenized size of each source, used to configure weighted mixing in + :meth:`build_train_sampler`. + """ + tokenized: list = [] + dataset_lengths: list[int] = [] + for name, split in zip(self.sft_cfg.train_datasets, self.sft_cfg.train_dataset_splits): + source = self._load_and_tokenize(name, split) + if len(source) == 0: + raise ValueError(f"Training dataset '{name}' (split '{split}') tokenized to 0 examples.") + tokenized.extend(source) + dataset_lengths.append(len(source)) + if len(dataset_lengths) > 1: + per_dataset = ", ".join( + f"{name}={length}" for name, length in zip(self.sft_cfg.train_datasets, dataset_lengths) + ) + logger.info(f"Concatenated {len(dataset_lengths)} training datasets: {per_dataset}") + return tokenized, dataset_lengths - def load_eval_dataset(self) -> Optional[list]: - """Load and tokenize the eval dataset, or return ``None`` if not configured.""" - if not self.sft_cfg.eval_dataset_name: + def load_eval_datasets(self) -> Optional[list[tuple[str, list]]]: + """Load and tokenize the eval dataset(s), or return ``None`` if not configured. + + Returns: + One ``(name, tokenized)`` pair per entry of ``eval_datasets``, where + ``name`` comes from ``eval_dataset_names`` and namespaces the eval + metrics (``eval/{name}/...``). + """ + if not self.sft_cfg.eval_datasets: return None - return self._load_and_tokenize(self.sft_cfg.eval_dataset_name, self.sft_cfg.eval_dataset_split) + eval_sets: list[tuple[str, list]] = [] + for name, dataset, split in zip( + self.sft_cfg.eval_dataset_names, self.sft_cfg.eval_datasets, self.sft_cfg.eval_dataset_splits + ): + eval_tokenized = self._load_and_tokenize(dataset, split) + if len(eval_tokenized) == 0: + raise ValueError( + f"Eval dataset '{dataset}' (split '{split}') tokenized to 0 examples. " + f"Provide a non-empty eval split or remove it from eval_datasets." + ) + eval_sets.append((name, eval_tokenized)) + return eval_sets def _log_dataset_stats(self, tokenized: list) -> None: """Log tokenized sequence length statistics over the training set. @@ -1275,36 +1329,69 @@ def collate_batch(self, examples: list, batch_size: int) -> TrainingInputBatch: # Dataloaders & samplers # ------------------------------------------------------------------ # - def build_train_sampler(self, tokenized: list) -> Optional[torch.utils.data.Sampler]: + def build_train_sampler( + self, tokenized: list, dataset_lengths: Optional[list[int]] = None + ) -> Optional[torch.utils.data.Sampler]: """Build the training sampler from ``sft_cfg.sampler``. - Returns ``None`` for the default ``"random"`` strategy, signalling - :meth:`build_train_dataloader` to use the dataloader's built-in - ``shuffle=True`` path (which is statefully checkpointed by - ``StatefulDataLoader``). For ``"sequential"`` and ``"custom"`` it - returns an explicit stateful sampler. + Returns ``None`` for the default ``"random"`` strategy over a single + dataset, signalling :meth:`build_train_dataloader` to use the + dataloader's built-in ``shuffle=True`` path (which is statefully + checkpointed by ``StatefulDataLoader``). With multiple training + datasets, ``"random"`` instead returns a :class:`DataMixingSampler` + configured with the per-dataset lengths and ``train_dataset_weights``. + For ``"sequential"`` and ``"custom"`` it returns an explicit stateful + sampler. Custom samplers are imported from ``sft_cfg.sampler_class_path`` and - instantiated as ``ClassName(tokenized, **sft_cfg.sampler_kwargs)``. + instantiated as ``ClassName(tokenized, **sft_cfg.sampler_kwargs)``. With + multiple datasets, the per-dataset ``lengths`` are injected into the + kwargs (unless the user already supplied ``lengths``), so the sampler + constructor must accept them. + + Args: + tokenized: The (concatenated) tokenized training dataset. + dataset_lengths: Tokenized size of each source dataset, in order. + ``None`` is treated as a single source spanning ``tokenized``. """ from skyrl.train.dataset.samplers import ( + DataMixingSampler, StatefulSequentialSampler, import_sampler_class, ) + multi_dataset = dataset_lengths is not None and len(dataset_lengths) > 1 sampler_type = self.sft_cfg.sampler if sampler_type == "random": - return None + if not multi_dataset: + return None + weights = self.sft_cfg.train_dataset_weights + if weights is None: + # Config normalization fills this on the standard path; default + # to equal mixing for directly-constructed trainers. + weights = [1.0 / len(dataset_lengths)] * len(dataset_lengths) + return DataMixingSampler( + tokenized, + lengths=dataset_lengths, + weights=weights, + seed=self.sft_cfg.seed, + ) if sampler_type == "sequential": return StatefulSequentialSampler(tokenized) if sampler_type == "custom": if not self.sft_cfg.sampler_class_path: raise ValueError("sampler='custom' requires sampler_class_path to be set.") sampler_cls = import_sampler_class(self.sft_cfg.sampler_class_path) - return sampler_cls(tokenized, **self.sft_cfg.sampler_kwargs) + sampler_kwargs = self.sft_cfg.sampler_kwargs + if multi_dataset: + # User-provided kwargs win over the injected lengths. + sampler_kwargs = {"lengths": dataset_lengths, **sampler_kwargs} + return sampler_cls(tokenized, **sampler_kwargs) raise ValueError(f"Unknown sampler '{sampler_type}'. Must be one of 'random', 'sequential', 'custom'.") - def build_train_dataloader(self, tokenized: list) -> StatefulDataLoader: + def build_train_dataloader( + self, tokenized: list, dataset_lengths: Optional[list[int]] = None + ) -> StatefulDataLoader: """Build the training ``StatefulDataLoader``. Sampling order is seeded for reproducibility and captured in the dataloader's @@ -1336,7 +1423,7 @@ def build_train_dataloader(self, tokenized: list) -> StatefulDataLoader: seeded_generator = torch.Generator() seeded_generator.manual_seed(self.sft_cfg.seed) - sampler = self.build_train_sampler(tokenized) + sampler = self.build_train_sampler(tokenized, dataset_lengths) num_workers = self.sft_cfg.dataloader_num_workers return StatefulDataLoader( @@ -1485,9 +1572,38 @@ def load_checkpoint(self) -> int: # ------------------------------------------------------------------ # def run_eval(self) -> tuple[dict, int]: - """Compute eval loss over the full eval dataset. + """Compute eval loss over every configured eval dataset. + + Runs :meth:`_run_eval_one` per ``(name, dataloader)`` pair in + :attr:`eval_dataloaders`, namespacing each dataset's metrics by its + name. The keys are later prefixed with ``eval/`` at the logging sites, + yielding ``eval/{name}/loss`` — nested even with a single eval dataset, + so runs with and without dataset mixing chart the same metric keys. - Iterates :attr:`eval_dataloader` (chunks of ``micro_train_batch_size_per_gpu * dp_size``, + Returns: + ``(metrics, num_eval_batches)`` where ``metrics`` maps + ``{name}/loss`` to that dataset's token-weighted mean loss and + ``num_eval_batches`` is the total batch count across datasets + (stdout bookkeeping, not a wandb metric). + """ + if not self.eval_dataloaders: + raise ValueError( + "run_eval called without eval dataloaders. Provide non-empty eval splits or " + "disable eval by setting eval_datasets=None." + ) + metrics: dict[str, float] = {} + total_eval_batches = 0 + for name, eval_dataloader in self.eval_dataloaders: + eval_loss, num_eval_batches = self._run_eval_one(eval_dataloader) + metrics[f"{name}/loss"] = eval_loss + total_eval_batches += num_eval_batches + logger.info(f"Eval dataset '{name}': loss={eval_loss:.4f} over {num_eval_batches} batches") + return metrics, total_eval_batches + + def _run_eval_one(self, eval_dataloader: StatefulDataLoader) -> tuple[float, int]: + """Compute eval loss over one eval dataset. + + Iterates the dataloader (chunks of ``micro_train_batch_size_per_gpu * dp_size``, i.e. exactly one micro-batch per DP rank per dispatch call), calls :meth:`WorkerDispatch.forward` with ``loss_fn="cross_entropy"`` (which runs the model in ``eval()`` mode under ``no_grad``), and aggregates the @@ -1498,25 +1614,11 @@ def run_eval(self) -> tuple[dict, int]: yields the true per-non-pad-token mean across the eval dataset. Returns: - ``(metrics, num_eval_batches)`` where ``metrics`` contains - ``eval_loss`` and ``num_eval_batches`` is bookkeeping for - stdout logging (not a wandb metric). + ``(eval_loss, num_eval_batches)``. """ - if self.eval_dataloader is None: - raise ValueError( - "run_eval called without an eval dataloader. Provide a non-empty eval split or " - "disable eval by setting eval_dataset_name=None." - ) - num_eval = len(self.eval_dataloader.dataset) - if num_eval == 0: - raise ValueError( - "Eval dataset is empty. Provide a non-empty eval split or disable eval " - "by setting eval_dataset_name=None." - ) - # The dataloader yields one chunk per DP rank's micro-batch; the final # (possibly short) chunk is padded below up to the full chunk size. - eval_chunk_size = self.eval_dataloader.batch_size + eval_chunk_size = eval_dataloader.batch_size # Pad a trailing partial batch up to ``eval_chunk_size`` via # ``pad_training_input_batch`` (which zeros ``loss_mask`` on padded rows). @@ -1527,7 +1629,7 @@ def run_eval(self) -> tuple[dict, int]: total_loss_weighted = 0.0 total_tokens = 0 num_eval_batches = 0 - for batch in self.eval_dataloader: + for batch in eval_dataloader: num_eval_batches += 1 # Pad the last (possibly-short) chunk so every dispatch sees exactly # ``eval_chunk_size`` rows. ``pad_training_input_batch`` zeros the @@ -1557,7 +1659,7 @@ def run_eval(self) -> tuple[dict, int]: total_tokens += nonpad_tokens eval_loss = total_loss_weighted / max(total_tokens, 1) - return {"eval_loss": eval_loss}, num_eval_batches + return eval_loss, num_eval_batches def train_step(self, batch: TrainingInputBatch, step: int) -> dict: """Execute a single training step: forward_backward + optim_step. @@ -1734,16 +1836,17 @@ def train(self): logger.warning("resume_from is ignored in dummy run mode") return self._train_dummy() - tokenized = self.load_dataset() + tokenized, dataset_lengths = self.load_dataset() # Log tokenized sequence length statistics (once, before training loop) self._log_dataset_stats(tokenized) - # Load eval dataset (if configured). We load once up-front so the + # Load eval datasets (if configured). We load once up-front so the # tokenization cost is amortized across all eval invocations. - eval_tokenized = self.load_eval_dataset() - if eval_tokenized is not None: - logger.info(f"Eval dataset loaded: {len(eval_tokenized)} examples") + eval_datasets = self.load_eval_datasets() + if eval_datasets is not None: + for eval_name, eval_tokenized in eval_datasets: + logger.info(f"Eval dataset '{eval_name}' loaded: {len(eval_tokenized)} examples") batch_size = self.sft_cfg.batch_size @@ -1752,9 +1855,11 @@ def train(self): # Build stateful dataloaders (replaces manual list shuffling/slicing). # The training sampler is selected by ``sft_cfg.sampler`` and its # position is captured in the checkpoint for resume. - self.train_dataloader = self.build_train_dataloader(tokenized) - if eval_tokenized is not None: - self.eval_dataloader = self.build_eval_dataloader(eval_tokenized) + self.train_dataloader = self.build_train_dataloader(tokenized, dataset_lengths) + if eval_datasets is not None: + self.eval_dataloaders = [ + (eval_name, self.build_eval_dataloader(eval_tokenized)) for eval_name, eval_tokenized in eval_datasets + ] # Validate the invariant the training loop relies on: the dataloader must # yield at least one batch. With drop_last=False (the final short batch is @@ -1826,7 +1931,7 @@ def train(self): # Baseline eval before training begins (logged at step 0). # Wandb's step counter starts at 0; the training loop's first commit # advances it to >=1, so step=0 here does not conflict with later steps. - if self.sft_cfg.eval_before_train and eval_tokenized is not None: + if self.sft_cfg.eval_before_train and self.eval_dataloaders is not None: self._fire("on_eval_start") eval_metrics, num_eval_batches = self.run_eval() self._fire("on_eval_end", metrics=eval_metrics) @@ -1834,8 +1939,7 @@ def train(self): self._fire("on_log", logs=baseline_log) self.tracker.log(baseline_log, step=self.global_step, commit=True) logger.info( - f"Baseline eval before training: " - f"eval_loss={eval_metrics.get('eval_loss', float('nan')):.4f} " + f"Baseline eval before training: {_format_eval_metrics(eval_metrics)} " f"over {num_eval_batches} batches" ) @@ -1934,7 +2038,7 @@ def train(self): # 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): + if self.eval_dataloaders 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() @@ -1956,7 +2060,7 @@ def train(self): if eval_metrics: logger.info( - f"Step {self.global_step}: eval_loss={eval_metrics.get('eval_loss', float('nan')):.4f} " + f"Step {self.global_step}: {_format_eval_metrics(eval_metrics)} " f"over {num_eval_batches} batches" ) @@ -2003,7 +2107,7 @@ def train(self): # ``final_eval_step`` rather than mutating ``self.global_step``: the # bump is purely a wandb-step accounting concern, not real trainer # state. - if eval_tokenized is not None: + if self.eval_dataloaders is not None: already_ran = self.sft_cfg.eval_interval > 0 and num_steps % self.sft_cfg.eval_interval == 0 if not already_ran: final_eval_step = num_steps + 1 @@ -2018,8 +2122,7 @@ def train(self): self._fire("on_log", logs=eval_log) self.tracker.log(eval_log, step=final_eval_step, commit=True) logger.info( - f"Final eval at step {final_eval_step}: " - f"eval_loss={eval_metrics.get('eval_loss', float('nan')):.4f} " + f"Final eval at step {final_eval_step}: {_format_eval_metrics(eval_metrics)} " f"over {num_eval_batches} batches" ) diff --git a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_sft_packing_parity.py b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_sft_packing_parity.py index eea94e3dfa..5bff36356d 100644 --- a/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_sft_packing_parity.py +++ b/tests/backends/skyrl_train/gpu/gpu_ci/megatron/test_sft_packing_parity.py @@ -260,7 +260,8 @@ def _make_sft_cfg(use_sequence_packing: bool, cp: int, gpus: int) -> SFTConfig: max_tokens_per_microbatch=MAX_LENGTH if use_sequence_packing else None, seed=SEED, train_on_what=TrainOnWhat.ALL_ASSISTANT_MESSAGES, - dataset_name="allenai/tulu-3-sft-mixture", + train_datasets=["allenai/tulu-3-sft-mixture"], + train_dataset_splits=["train[:100]"], placement=SFTPlacementConfig(num_nodes=1, num_gpus_per_node=gpus), megatron_config=MegatronConfig( tensor_model_parallel_size=1, diff --git a/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh b/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh index c120427a68..dfd97ef457 100755 --- a/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh +++ b/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh @@ -34,7 +34,7 @@ LOG_FILE="${LOG_FILE:-/tmp/${RUN_NAME}.log}" # * batch_size=8, micro_train_batch_size_per_gpu=2 are sized for L4_ci (4 GPUs). bash examples/train/sft/run_sft_megatron_tulu3_50k.sh \ num_steps=$NUM_STEPS \ - dataset_split="train[:2000]" \ + "train_dataset_splits=[train[:2000]]" \ batch_size=8 \ micro_train_batch_size_per_gpu=2 \ max_length=1024 \ diff --git a/tests/train/test_sft_callbacks.py b/tests/train/test_sft_callbacks.py index 1a7d1c2d6d..a793b2680b 100644 --- a/tests/train/test_sft_callbacks.py +++ b/tests/train/test_sft_callbacks.py @@ -115,15 +115,17 @@ def on_step_end(self, trainer, ci, control): def _build_test_sft_config() -> SFTConfig: cfg = SFTConfig() cfg.strategy = "fsdp" - # model.path / dataset_name are unused — we never load the model and - # monkeypatch _load_and_tokenize. eval_dataset_name must be non-empty so - # load_eval_dataset actually invokes _load_and_tokenize. + # model.path / train_datasets are unused — we never load the model and + # monkeypatch _load_and_tokenize. eval_datasets must be non-empty so + # load_eval_datasets actually invokes _load_and_tokenize. cfg.model.path = "unused" cfg.placement = SFTPlacementConfig(num_nodes=1, num_gpus_per_node=1) - cfg.dataset_name = "unused-monkeypatched" - cfg.dataset_split = "train" - cfg.eval_dataset_name = "unused-monkeypatched" - cfg.eval_dataset_split = "train" + cfg.train_datasets = ["unused-monkeypatched"] + cfg.train_dataset_splits = ["train"] + cfg.eval_datasets = ["unused-monkeypatched"] + cfg.eval_dataset_splits = ["train"] + # Shorthand logging name: eval metrics land under eval/evalset/... + cfg.eval_dataset_names = ["evalset"] # eval_interval=2 means step 1 has no interval-driven eval; only the force-evaluate # callback can trigger eval at step 1. Step 2 still gets an interval-driven eval. cfg.eval_interval = 2 @@ -252,16 +254,16 @@ def _record_load_checkpoint(): assert snap["has_metrics"], "on_step_end should see step metrics" assert "loss" in snap["metrics_keys"], snap["metrics_keys"] - # Both eval ends carry eval metrics + # Both eval ends carry eval metrics, namespaced by the eval dataset name for snap in snaps_by_event["on_eval_end"]: assert snap["has_metrics"], "on_eval_end should see eval metrics" - assert "eval_loss" in snap["metrics_keys"], snap["metrics_keys"] + assert "evalset/loss" in snap["metrics_keys"], snap["metrics_keys"] - # Both on_log calls carry train/loss + eval/eval_loss + # Both on_log calls carry train/loss + eval/{name}/loss for snap in snaps_by_event["on_log"]: log_keys = snap["logs_keys"] assert "train/loss" in log_keys, log_keys - assert "eval/eval_loss" in log_keys, log_keys + assert "eval/evalset/loss" in log_keys, log_keys # on_save fired exactly once at step 2, with the fake ckpt path assert len(snaps_by_event["on_save"]) == 1, snaps_by_event["on_save"] diff --git a/tests/train/test_sft_config.py b/tests/train/test_sft_config.py index 65bb80ef2c..00e6b585bd 100644 --- a/tests/train/test_sft_config.py +++ b/tests/train/test_sft_config.py @@ -13,7 +13,7 @@ SFTConfig, build_skyrl_config_for_sft, ) -from skyrl.train.config.sft_config import validate_sft_cfg +from skyrl.train.config.sft_config import _normalize_dataset_cfg, validate_sft_cfg def _sft_cfg_from_overrides(overrides: list[str]) -> SFTConfig: @@ -259,3 +259,135 @@ def test_bridge_sets_worker_micro_batch_size_to_one(self): cfg = self._packed_cfg(max_tokens_per_microbatch=256, micro_train_batch_size_per_gpu=4) skyrl_cfg = build_skyrl_config_for_sft(cfg) assert skyrl_cfg.trainer.micro_train_batch_size_per_gpu == 1 + + +class TestDatasetConfigNormalization: + """Deprecated single-dataset fields translate to the list-based fields (RFC #1875).""" + + def test_zero_config_defaults(self): + cfg = SFTConfig() + _normalize_dataset_cfg(cfg) + assert cfg.train_datasets == ["yahma/alpaca-cleaned"] + assert cfg.train_dataset_splits == ["train[:100]"] + assert cfg.train_dataset_weights == [1.0] + assert cfg.dataset_name is None and cfg.dataset_split is None + assert cfg.eval_datasets is None # eval stays disabled + + def test_deprecated_train_fields_translate_with_warning(self): + with pytest.warns(DeprecationWarning, match="dataset_name/dataset_split are deprecated"): + cfg = _sft_cfg_from_overrides(["dataset_name=foo/bar", "dataset_split=train[:50]"]) + _normalize_dataset_cfg(cfg) + assert cfg.train_datasets == ["foo/bar"] + assert cfg.train_dataset_splits == ["train[:50]"] + assert cfg.dataset_name is None and cfg.dataset_split is None + + def test_deprecated_translation_works_programmatically(self): + cfg = SFTConfig() + cfg.dataset_name = "foo/bar" + with pytest.warns(DeprecationWarning, match="dataset_name/dataset_split are deprecated"): + validate_sft_cfg(cfg) + assert cfg.train_datasets == ["foo/bar"] + assert cfg.train_dataset_splits == ["train[:100]"] # split default preserved + + def test_deprecated_eval_fields_translate_with_warning(self): + with pytest.warns(DeprecationWarning, match="eval_dataset_name/eval_dataset_split are deprecated"): + cfg = _sft_cfg_from_overrides(["eval_dataset_name=baz/qux"]) + _normalize_dataset_cfg(cfg) + assert cfg.eval_datasets == ["baz/qux"] + assert cfg.eval_dataset_splits == ["validation"] # old default preserved + assert cfg.eval_dataset_names == ["baz_qux"] # "/" sanitized for logging + assert cfg.eval_dataset_name is None and cfg.eval_dataset_split is None + + def test_old_and_new_train_fields_conflict(self): + cfg = _sft_cfg_from_overrides(["dataset_name=x", "train_datasets=[a]", "train_dataset_splits=[s]"]) + with pytest.raises(ValueError, match="only one of train_datasets"): + _normalize_dataset_cfg(cfg) + + def test_old_and_new_eval_fields_conflict(self): + cfg = _sft_cfg_from_overrides(["eval_dataset_name=x", "eval_datasets=[a]", "eval_dataset_splits=[s]"]) + with pytest.raises(ValueError, match="only one of eval_datasets"): + _normalize_dataset_cfg(cfg) + + def test_idempotent(self): + cfg = _sft_cfg_from_overrides(["train_datasets=[a,b]", "train_dataset_splits=[s1,s2]", "sampler=sequential"]) + _normalize_dataset_cfg(cfg) + first = (cfg.train_datasets, cfg.train_dataset_splits, cfg.train_dataset_weights) + _normalize_dataset_cfg(cfg) + assert (cfg.train_datasets, cfg.train_dataset_splits, cfg.train_dataset_weights) == first + + +class TestMultiDatasetValidation: + """List-shaped dataset fields are validated for lengths, weights and names.""" + + def test_default_weights_are_equal_mixing(self): + cfg = _sft_cfg_from_overrides(["train_datasets=[a,b,c,d]", "train_dataset_splits=[s,s,s,s]"]) + _normalize_dataset_cfg(cfg) + assert cfg.train_dataset_weights == pytest.approx([0.25] * 4) + + def test_splits_length_mismatch_rejected(self): + cfg = _sft_cfg_from_overrides(["train_datasets=[a,b]", "train_dataset_splits=[s1]"]) + with pytest.raises(ValueError, match="one split per entry"): + _normalize_dataset_cfg(cfg) + + def test_missing_splits_rejected(self): + cfg = _sft_cfg_from_overrides(["train_datasets=[a,b]"]) + with pytest.raises(ValueError, match="one split per entry"): + _normalize_dataset_cfg(cfg) + + def test_weights_length_mismatch_rejected(self): + cfg = _sft_cfg_from_overrides( + ["train_datasets=[a,b]", "train_dataset_splits=[s,s]", "train_dataset_weights=[1.0]"] + ) + with pytest.raises(ValueError, match="one weight per entry"): + _normalize_dataset_cfg(cfg) + + def test_nonpositive_weights_rejected(self): + cfg = _sft_cfg_from_overrides( + ["train_datasets=[a,b]", "train_dataset_splits=[s,s]", "train_dataset_weights=[1.0,0.0]"] + ) + with pytest.raises(ValueError, match="must all be > 0"): + _normalize_dataset_cfg(cfg) + + def test_weights_require_random_sampler(self): + cfg = _sft_cfg_from_overrides( + [ + "train_datasets=[a,b]", + "train_dataset_splits=[s,s]", + "train_dataset_weights=[0.5,0.5]", + "sampler=sequential", + ] + ) + with pytest.raises(ValueError, match="only supported with sampler='random'"): + _normalize_dataset_cfg(cfg) + + def test_eval_names_default_collision_rejected(self): + cfg = _sft_cfg_from_overrides(["eval_datasets=[a/b,a/b]", "eval_dataset_splits=[s1,s2]"]) + with pytest.raises(ValueError, match="names collide"): + _normalize_dataset_cfg(cfg) + + def test_eval_names_explicit_disambiguate(self): + cfg = _sft_cfg_from_overrides( + ["eval_datasets=[a/b,a/b]", "eval_dataset_splits=[s1,s2]", "eval_dataset_names=[early,late]"] + ) + _normalize_dataset_cfg(cfg) + assert cfg.eval_dataset_names == ["early", "late"] + + def test_eval_names_duplicates_rejected(self): + cfg = _sft_cfg_from_overrides(["eval_datasets=[a,b]", "eval_dataset_splits=[s,s]", "eval_dataset_names=[x,x]"]) + with pytest.raises(ValueError, match="must be unique"): + _normalize_dataset_cfg(cfg) + + def test_eval_names_length_mismatch_rejected(self): + cfg = _sft_cfg_from_overrides(["eval_datasets=[a,b]", "eval_dataset_splits=[s,s]", "eval_dataset_names=[x]"]) + with pytest.raises(ValueError, match="one name per entry"): + _normalize_dataset_cfg(cfg) + + def test_eval_lists_without_eval_datasets_rejected(self): + cfg = _sft_cfg_from_overrides(["eval_dataset_splits=[s]"]) + with pytest.raises(ValueError, match="require eval_datasets"): + _normalize_dataset_cfg(cfg) + + def test_eval_interval_requires_eval_datasets(self): + cfg = _sft_cfg_from_overrides(["eval_interval=5"]) + with pytest.raises(ValueError, match="requires eval_datasets"): + validate_sft_cfg(cfg) diff --git a/tests/train/test_sft_dataloader.py b/tests/train/test_sft_dataloader.py index 3ad325d211..c6502f0f5f 100644 --- a/tests/train/test_sft_dataloader.py +++ b/tests/train/test_sft_dataloader.py @@ -1,6 +1,6 @@ -"""CPU tests for the SFT stateful dataloader and custom samplers. +"""CPU tests for the SFT stateful dataloader and samplers. -Covers the RFC test plan: +Covers the stateful-dataloader RFC test plan: - default random dataloader order is deterministic for identical seeds - random dataloader order differs across different seeds - sequential sampler yields examples in order @@ -8,10 +8,18 @@ - custom sampler state is included in the dataloader checkpoint - data-mixing / curriculum samplers resume correctly -The custom-sampler path is exercised with a small test-local sampler (the core -library ships only ``StatefulSequentialSampler``; curriculum learning is a -user-supplied example under ``examples/train/sft/curriculum_sampler.py``, which -is loaded by file path here to keep it covered). +and the multi-dataset RFC (#1875) test plan: + - user-specified train_dataset_weights: per-batch source representation + tracks the weights + - default weights: equal mixing across datasets + - mixing ratios are independent of the individual dataset sizes + - DataMixingSampler draws a fresh plan every epoch and resumes exactly + (mid-epoch and across epoch boundaries) via its RNG state + +The custom-sampler path is exercised with small test-local samplers (curriculum +learning remains a user-supplied example under +``examples/train/sft/curriculum_sampler.py``, loaded by file path here to keep +it covered). Run:: @@ -27,7 +35,11 @@ from torchdata.stateful_dataloader import StatefulDataLoader from skyrl.train.config import SFTConfig -from skyrl.train.dataset.samplers import StatefulSequentialSampler, import_sampler_class +from skyrl.train.dataset.samplers import ( + DataMixingSampler, + StatefulSequentialSampler, + import_sampler_class, +) from skyrl.train.sft_trainer import SFTTrainer # --------------------------------------------------------------------------- @@ -74,6 +86,18 @@ def load_state_dict(self, state): _CUSTOM_SAMPLER_PATH = f"{__name__}._ExampleKwargSampler" +class _LengthsAwareSampler(_ExampleKwargSampler): + """Test-local custom sampler that accepts the ``lengths`` kwarg the trainer + injects when multiple training datasets are configured.""" + + def __init__(self, data_source, lengths=None, num_samples=None, seed=0): + super().__init__(data_source, num_samples=num_samples, seed=seed) + self.lengths = lengths + + +_LENGTHS_SAMPLER_PATH = f"{__name__}._LengthsAwareSampler" + + def _load_example_cls(filename: str, class_name: str): """Import a sampler class from an examples/train/sft/*.py file by path.""" path = Path(__file__).resolve().parents[2] / "examples" / "train" / "sft" / filename @@ -203,6 +227,55 @@ def test_unknown_sampler(self): with pytest.raises(ValueError, match="Unknown sampler"): trainer.build_train_sampler(list(range(10))) + def test_random_single_source_lengths_returns_none(self): + trainer = _make_trainer(sampler="random") + assert trainer.build_train_sampler(list(range(10)), dataset_lengths=[10]) is None + + def test_random_multi_dataset_returns_mixing_sampler(self): + trainer = _make_trainer(sampler="random", train_dataset_weights=[0.8, 0.2], seed=11) + sampler = trainer.build_train_sampler(list(range(20)), dataset_lengths=[10, 10]) + assert isinstance(sampler, DataMixingSampler) + assert len(sampler) == 20 + # Seeded from the config so runs are reproducible. + other = _make_trainer(sampler="random", train_dataset_weights=[0.8, 0.2], seed=11).build_train_sampler( + list(range(20)), dataset_lengths=[10, 10] + ) + assert list(sampler) == list(other) + + def test_random_multi_dataset_defaults_to_equal_weights(self): + # Bare trainers skip config normalization; equal mixing is the fallback. + trainer = _make_trainer(sampler="random") + sampler = trainer.build_train_sampler(list(range(20)), dataset_lengths=[10, 10]) + assert isinstance(sampler, DataMixingSampler) + + def test_custom_multi_dataset_injects_lengths(self): + trainer = _make_trainer( + sampler="custom", + sampler_class_path=_LENGTHS_SAMPLER_PATH, + sampler_kwargs={"num_samples": 8}, + ) + sampler = trainer.build_train_sampler(list(range(20)), dataset_lengths=[12, 8]) + assert isinstance(sampler, _LengthsAwareSampler) + assert sampler.lengths == [12, 8] + + def test_custom_multi_dataset_user_lengths_win(self): + trainer = _make_trainer( + sampler="custom", + sampler_class_path=_LENGTHS_SAMPLER_PATH, + sampler_kwargs={"num_samples": 8, "lengths": [1, 19]}, + ) + sampler = trainer.build_train_sampler(list(range(20)), dataset_lengths=[12, 8]) + assert sampler.lengths == [1, 19] + + def test_custom_single_dataset_does_not_inject_lengths(self): + trainer = _make_trainer( + sampler="custom", + sampler_class_path=_LENGTHS_SAMPLER_PATH, + sampler_kwargs={"num_samples": 8}, + ) + sampler = trainer.build_train_sampler(list(range(20)), dataset_lengths=[20]) + assert sampler.lengths is None + # --------------------------------------------------------------------------- # build_train_dataloader: determinism + checkpoint/resume @@ -451,99 +524,143 @@ def test_state_dict_resume(self): # --------------------------------------------------------------------------- -# Example: DataMixingSampler (WeightedRandomSampler-based, loaded by path) +# DataMixingSampler (core) # --------------------------------------------------------------------------- -class TestDataMixingExample: +class TestDataMixingSampler: # data = two sources of 10 each: indices [0,10) = source 0, [10,20) = source 1. DATA = list(range(20)) LENGTHS = [10, 10] + def _make(self, weights=(0.7, 0.3), num_samples=40, seed=3, data=None, lengths=None): + return DataMixingSampler( + data if data is not None else self.DATA, + lengths=lengths if lengths is not None else self.LENGTHS, + weights=list(weights), + num_samples=num_samples, + seed=seed, + ) + def test_len_matches_num_samples(self): - cls = _load_data_mixing_cls() - sampler = cls(self.DATA, lengths=self.LENGTHS, weights=[0.5, 0.5], num_samples=64) + sampler = self._make(weights=[0.5, 0.5], num_samples=64) assert len(sampler) == 64 + def test_num_samples_defaults_to_dataset_length(self): + sampler = DataMixingSampler(self.DATA, lengths=self.LENGTHS, weights=[0.5, 0.5]) + assert len(sampler) == 20 + + def test_example_file_reexports_core_class(self): + assert _load_data_mixing_cls() is DataMixingSampler + def test_validates_lengths_sum(self): - cls = _load_data_mixing_cls() with pytest.raises(ValueError, match="must equal len"): - cls(self.DATA, lengths=[10, 5], weights=[0.5, 0.5]) + DataMixingSampler(self.DATA, lengths=[10, 5], weights=[0.5, 0.5]) def test_validates_weights_align(self): - cls = _load_data_mixing_cls() with pytest.raises(ValueError, match="must align"): - cls(self.DATA, lengths=self.LENGTHS, weights=[1.0]) + DataMixingSampler(self.DATA, lengths=self.LENGTHS, weights=[1.0]) def test_rejects_degenerate_weights(self): - cls = _load_data_mixing_cls() with pytest.raises(ValueError, match="non-negative with a positive sum"): - cls(self.DATA, lengths=self.LENGTHS, weights=[0.0, 0.0]) + DataMixingSampler(self.DATA, lengths=self.LENGTHS, weights=[0.0, 0.0]) def test_deterministic_for_same_seed(self): - cls = _load_data_mixing_cls() - a = list(cls(self.DATA, lengths=self.LENGTHS, weights=[0.8, 0.2], num_samples=200, seed=1)) - b = list(cls(self.DATA, lengths=self.LENGTHS, weights=[0.8, 0.2], num_samples=200, seed=1)) + a = list(self._make(weights=[0.8, 0.2], num_samples=200, seed=1)) + b = list(self._make(weights=[0.8, 0.2], num_samples=200, seed=1)) assert a == b def test_differs_for_different_seed(self): - cls = _load_data_mixing_cls() - a = list(cls(self.DATA, lengths=self.LENGTHS, weights=[0.8, 0.2], num_samples=200, seed=1)) - b = list(cls(self.DATA, lengths=self.LENGTHS, weights=[0.8, 0.2], num_samples=200, seed=2)) + a = list(self._make(weights=[0.8, 0.2], num_samples=200, seed=1)) + b = list(self._make(weights=[0.8, 0.2], num_samples=200, seed=2)) assert a != b def test_weighting_biases_the_mix(self): - cls = _load_data_mixing_cls() # Source 0 weighted 4x source 1 -> ~80% of draws from indices [0,10). - plan = list(cls(self.DATA, lengths=self.LENGTHS, weights=[0.8, 0.2], num_samples=4000, seed=7)) + plan = list(self._make(weights=[0.8, 0.2], num_samples=4000, seed=7)) frac_source0 = sum(1 for idx in plan if idx < 10) / len(plan) assert 0.75 < frac_source0 < 0.85, frac_source0 def test_equal_weights_balanced(self): - cls = _load_data_mixing_cls() - plan = list(cls(self.DATA, lengths=self.LENGTHS, weights=[0.5, 0.5], num_samples=4000, seed=7)) + plan = list(self._make(weights=[0.5, 0.5], num_samples=4000, seed=7)) frac_source0 = sum(1 for idx in plan if idx < 10) / len(plan) assert 0.45 < frac_source0 < 0.55, frac_source0 def test_size_imbalance_still_matches_weights(self): - cls = _load_data_mixing_cls() # Source 0 is tiny (2 examples) but weighted equally; source-level mix # should still be ~50/50 because weight is divided across examples. data = list(range(12)) # source 0: [0,2), source 1: [2,12) - plan = list(cls(data, lengths=[2, 10], weights=[0.5, 0.5], num_samples=4000, seed=9)) + plan = list(self._make(data=data, lengths=[2, 10], weights=[0.5, 0.5], num_samples=4000, seed=9)) frac_source0 = sum(1 for idx in plan if idx < 2) / len(plan) assert 0.45 < frac_source0 < 0.55, frac_source0 + def test_fresh_plan_each_epoch(self): + sampler = self._make() + epoch1, epoch2 = list(sampler), list(sampler) + assert epoch1 != epoch2, "each epoch must draw a fresh plan" + # ...but the whole multi-epoch stream is reproducible from the seed. + other = self._make() + assert [list(other), list(other)] == [epoch1, epoch2] + def test_state_dict_resume(self): - cls = _load_data_mixing_cls() - sampler = cls(self.DATA, lengths=self.LENGTHS, weights=[0.7, 0.3], num_samples=40, seed=3) + sampler = self._make() it = iter(sampler) consumed = [next(it) for _ in range(11)] state = sampler.state_dict() - assert state == {"position": 11} + assert state["position"] == 11 + assert "generator_state" in state rest = list(it) - resumed = cls(self.DATA, lengths=self.LENGTHS, weights=[0.7, 0.3], num_samples=40, seed=3) + # A different construction seed: the checkpointed generator state alone + # must determine the restored plan. + resumed = self._make(seed=999) resumed.load_state_dict(state) assert list(resumed) == rest - assert consumed + rest == list(cls(self.DATA, lengths=self.LENGTHS, weights=[0.7, 0.3], num_samples=40, seed=3)) + assert consumed + rest == list(self._make()) + + def test_mid_epoch_resume_matches_across_epoch_boundary(self): + sampler = self._make() + it = iter(sampler) + [next(it) for _ in range(11)] + state = sampler.state_dict() + # Uninterrupted: rest of the current epoch, then the full next epoch. + expected = list(it) + list(sampler) + + resumed = self._make(seed=999) + resumed.load_state_dict(state) + assert list(resumed) + list(resumed) == expected + + def test_between_epoch_resume_draws_same_next_epoch(self): + sampler = self._make() + list(sampler) # exhaust epoch 1 + state = sampler.state_dict() + epoch2 = list(sampler) + + resumed = self._make(seed=999) + resumed.load_state_dict(state) + assert list(resumed) == epoch2 + + def test_old_format_position_only_state_loads(self): + # Checkpoints from the old example sampler carry only the position. + sampler = self._make() + sampler.load_state_dict({"position": 5}) + assert sampler.position == 5 + assert len(list(sampler)) == 35 # resumes the cursor best-effort + + def _make_dl(self): + sampler = self._make(num_samples=40, seed=5) + return StatefulDataLoader( + self.DATA, + batch_size=4, + sampler=sampler, + shuffle=False, + drop_last=True, + collate_fn=lambda x: list(x), + ) def test_resume_through_stateful_dataloader(self): """The sampler's state is captured in the dataloader checkpoint and resumes.""" - cls = _load_data_mixing_cls() - - def make_dl(): - sampler = cls(self.DATA, lengths=self.LENGTHS, weights=[0.7, 0.3], num_samples=40, seed=5) - return StatefulDataLoader( - self.DATA, - batch_size=4, - sampler=sampler, - shuffle=False, - drop_last=True, - collate_fn=lambda x: list(x), - ) - - dl = make_dl() + dl = self._make_dl() it = iter(dl) next(it) next(it) @@ -551,7 +668,62 @@ def make_dl(): assert "_sampler_iter_state" in state rest_full = [b for b in it] - dl2 = make_dl() + dl2 = self._make_dl() dl2.load_state_dict(state) rest_resumed = [b for b in dl2] assert rest_full == rest_resumed + + def test_dataloader_epochs_differ_and_resume_after_boundary(self): + """Fresh plans survive the trainer's epoch-boundary iter() re-creation, + and a checkpoint taken in a later epoch resumes exactly.""" + dl = self._make_dl() + epoch1 = list(dl) + it = iter(dl) # epoch 2: the sampler re-plans on the fresh iterator + first_batch = next(it) + state = dl.state_dict() + rest_full = [b for b in it] + assert [first_batch] + rest_full != epoch1, "epoch 2 must be a fresh plan" + + dl2 = self._make_dl() + dl2.load_state_dict(state) + rest_resumed = [b for b in dl2] + assert rest_full == rest_resumed + + +# --------------------------------------------------------------------------- +# Multi-dataset mixing through build_train_dataloader (RFC #1875 test plan) +# --------------------------------------------------------------------------- + + +class TestMultiDatasetMixing: + """Per-batch source representation tracks ``train_dataset_weights``, + independent of the individual dataset sizes.""" + + @staticmethod + def _source0_fraction(trainer, data, dataset_lengths, boundary, epochs=10): + dl = trainer.build_train_dataloader(data, dataset_lengths=dataset_lengths) + drawn = [] + for _ in range(epochs): + drawn.extend(_flatten(dl)) + return sum(1 for idx in drawn if idx < boundary) / len(drawn) + + def test_user_weights_reflected_in_batches(self): + # Two sources of 100 each, weighted 80/20. + data = list(range(200)) + trainer = _make_trainer(sampler="random", batch_size=20, seed=3, train_dataset_weights=[0.8, 0.2]) + frac = self._source0_fraction(trainer, data, [100, 100], boundary=100) + assert 0.75 < frac < 0.85, frac + + def test_default_weights_mix_equally(self): + data = list(range(200)) + trainer = _make_trainer(sampler="random", batch_size=20, seed=3) + frac = self._source0_fraction(trainer, data, [100, 100], boundary=100) + assert 0.45 < frac < 0.55, frac + + def test_mixing_independent_of_dataset_sizes(self): + # Source 0 has 20 examples, source 1 has 180; equal weights should + # still yield ~50/50 representation in batches. + data = list(range(200)) + trainer = _make_trainer(sampler="random", batch_size=20, seed=3, train_dataset_weights=[0.5, 0.5]) + frac = self._source0_fraction(trainer, data, [20, 180], boundary=20) + assert 0.45 < frac < 0.55, frac From fe512549f06efec224ebef76b3e0d2faac8d123c Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Fri, 10 Jul 2026 23:46:44 +0000 Subject: [PATCH 2/6] [fix][SFT] Quote list elements in dataset override examples and scripts Found by the GPU e2e run: OmegaConf/YAML cannot parse HF slice syntax unquoted inside a flow sequence -- train_dataset_splits=[train[:2000]] fails with a ParserError on the nested brackets. Each element must be quoted: "train_dataset_splits=['train[:2000]']". Fix all migrated run scripts, the e2e script, and the README examples, and add a CPU regression test encoding the supported CLI syntax. Verified end-to-end on 4xL40S via tests/train/gpu_e2e_test/sft_tulu3_megatron.sh: 100 steps, downward loss trend, all wandb CI assertions passed (run: sky-posttraining-uc-berkeley/skyrl_sft_ci/rii88p8v). Co-Authored-By: Claude Fable 5 Signed-off-by: Avi Basnet --- examples/train/sft/README.md | 16 ++++++++-------- examples/train/sft/run_sft_fsdp.sh | 4 ++-- examples/train/sft/run_sft_lora.sh | 4 ++-- examples/train/sft/run_sft_megatron.sh | 4 ++-- examples/train/sft/run_sft_megatron_apigen_mt.sh | 8 ++++---- .../train/sft/run_sft_megatron_multi_dataset.sh | 8 ++++---- examples/train/sft/run_sft_megatron_tulu3_50k.sh | 8 ++++---- .../sft/run_sft_megatron_tulu3_50k_packed.sh | 8 ++++---- examples/train/sft/run_sft_megatron_vlm.sh | 4 ++-- tests/train/gpu_e2e_test/sft_tulu3_megatron.sh | 2 +- tests/train/test_sft_config.py | 14 ++++++++++++++ 11 files changed, 47 insertions(+), 33 deletions(-) diff --git a/examples/train/sft/README.md b/examples/train/sft/README.md index 0a506b48d1..3b0fc054c2 100644 --- a/examples/train/sft/README.md +++ b/examples/train/sft/README.md @@ -207,9 +207,9 @@ overrides must be a single shell argument): 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]' + "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 @@ -228,9 +228,9 @@ 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_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 `/` @@ -249,8 +249,8 @@ kwarg. Explicit `train_dataset_weights` are rejected outside `sampler=random` ```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]]' \ + "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.curriculum_sampler.CurriculumLearningSampler \ 'sampler_kwargs={num_samples: 40, seed: 42}' diff --git a/examples/train/sft/run_sft_fsdp.sh b/examples/train/sft/run_sft_fsdp.sh index 3727a36e48..df9ca7e309 100755 --- a/examples/train/sft/run_sft_fsdp.sh +++ b/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 \ - "train_datasets=[yahma/alpaca-cleaned]" \ - "train_dataset_splits=[train[:100]]" \ + "train_datasets=['yahma/alpaca-cleaned']" \ + "train_dataset_splits=['train[:100]']" \ messages_key=messages \ max_length=512 \ num_steps=10 \ diff --git a/examples/train/sft/run_sft_lora.sh b/examples/train/sft/run_sft_lora.sh index 220b87da32..8d9ca604e8 100644 --- a/examples/train/sft/run_sft_lora.sh +++ b/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 \ - "train_datasets=[yahma/alpaca-cleaned]" \ - "train_dataset_splits=[train[:100]]" \ + "train_datasets=['yahma/alpaca-cleaned']" \ + "train_dataset_splits=['train[:100]']" \ messages_key=messages \ max_length=512 \ num_steps=10 \ diff --git a/examples/train/sft/run_sft_megatron.sh b/examples/train/sft/run_sft_megatron.sh index 83aef32513..82101ded50 100755 --- a/examples/train/sft/run_sft_megatron.sh +++ b/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 \ - "train_datasets=[yahma/alpaca-cleaned]" \ - "train_dataset_splits=[train[:100]]" \ + "train_datasets=['yahma/alpaca-cleaned']" \ + "train_dataset_splits=['train[:100]']" \ messages_key=messages \ max_length=512 \ num_steps=10 \ diff --git a/examples/train/sft/run_sft_megatron_apigen_mt.sh b/examples/train/sft/run_sft_megatron_apigen_mt.sh index e6f8ec5efe..44c0db12b2 100755 --- a/examples/train/sft/run_sft_megatron_apigen_mt.sh +++ b/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 \ - "train_datasets=[$DATA_DIR]" \ - "train_dataset_splits=[train[:4000]]" \ - "eval_datasets=[$DATA_DIR]" \ - "eval_dataset_splits=[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 \ diff --git a/examples/train/sft/run_sft_megatron_multi_dataset.sh b/examples/train/sft/run_sft_megatron_multi_dataset.sh index 6cff607edb..167d827d46 100755 --- a/examples/train/sft/run_sft_megatron_multi_dataset.sh +++ b/examples/train/sft/run_sft_megatron_multi_dataset.sh @@ -20,11 +20,11 @@ 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_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_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 \ diff --git a/examples/train/sft/run_sft_megatron_tulu3_50k.sh b/examples/train/sft/run_sft_megatron_tulu3_50k.sh index 9813d6c6d9..cda6cbde3c 100755 --- a/examples/train/sft/run_sft_megatron_tulu3_50k.sh +++ b/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 \ - "train_datasets=[allenai/tulu-3-sft-mixture]" \ - "train_dataset_splits=[train[:50000]]" \ - "eval_datasets=[allenai/tulu-3-sft-mixture]" \ - "eval_dataset_splits=[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 \ diff --git a/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh b/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh index 41635ef6e5..85ff3698f0 100755 --- a/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh +++ b/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 \ - "train_datasets=[allenai/tulu-3-sft-mixture]" \ - "train_dataset_splits=[train[:50000]]" \ - "eval_datasets=[allenai/tulu-3-sft-mixture]" \ - "eval_dataset_splits=[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 \ diff --git a/examples/train/sft/run_sft_megatron_vlm.sh b/examples/train/sft/run_sft_megatron_vlm.sh index 2563d6d0c9..eeb3f9ed6d 100755 --- a/examples/train/sft/run_sft_megatron_vlm.sh +++ b/examples/train/sft/run_sft_megatron_vlm.sh @@ -33,8 +33,8 @@ uv run --isolated --extra megatron \ python -m skyrl.train.main_sft \ strategy=megatron \ model.path=Qwen/Qwen3-VL-2B-Instruct \ - "train_datasets=[$DATA_DIR]" \ - "train_dataset_splits=[train]" \ + "train_datasets=['$DATA_DIR']" \ + "train_dataset_splits=['train']" \ messages_key=messages \ max_length=4096 \ num_steps=30 \ diff --git a/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh b/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh index dfd97ef457..b8675b4cda 100755 --- a/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh +++ b/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh @@ -34,7 +34,7 @@ LOG_FILE="${LOG_FILE:-/tmp/${RUN_NAME}.log}" # * batch_size=8, micro_train_batch_size_per_gpu=2 are sized for L4_ci (4 GPUs). bash examples/train/sft/run_sft_megatron_tulu3_50k.sh \ num_steps=$NUM_STEPS \ - "train_dataset_splits=[train[:2000]]" \ + "train_dataset_splits=['train[:2000]']" \ batch_size=8 \ micro_train_batch_size_per_gpu=2 \ max_length=1024 \ diff --git a/tests/train/test_sft_config.py b/tests/train/test_sft_config.py index 00e6b585bd..9bf6ec2e44 100644 --- a/tests/train/test_sft_config.py +++ b/tests/train/test_sft_config.py @@ -319,6 +319,20 @@ def test_idempotent(self): class TestMultiDatasetValidation: """List-shaped dataset fields are validated for lengths, weights and names.""" + def test_bracketed_splits_parse_with_inner_quotes(self): + # HF slice syntax nests brackets inside the OmegaConf list, so each + # element must be quoted on the CLI: "train_dataset_splits=['train[:2000]']". + # (Unquoted elements with brackets are a YAML parse error.) + cfg = _sft_cfg_from_overrides( + [ + "train_datasets=['allenai/tulu-3-sft-mixture','yahma/alpaca-cleaned']", + "train_dataset_splits=['train[:50000]','train[:10000]']", + ] + ) + _normalize_dataset_cfg(cfg) + assert list(cfg.train_datasets) == ["allenai/tulu-3-sft-mixture", "yahma/alpaca-cleaned"] + assert list(cfg.train_dataset_splits) == ["train[:50000]", "train[:10000]"] + def test_default_weights_are_equal_mixing(self): cfg = _sft_cfg_from_overrides(["train_datasets=[a,b,c,d]", "train_dataset_splits=[s,s,s,s]"]) _normalize_dataset_cfg(cfg) From a45ff5fe76116892dfb78fcf7a0f631048285917 Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Fri, 17 Jul 2026 18:09:34 +0000 Subject: [PATCH 3/6] [refactor][SFT] Remove data_mixing_sampler.py example shim Address review feedback (#1883): drop the backward-compat re-export of the now-core DataMixingSampler instead of keeping the example file around. Weighted mixing is configured natively via train_datasets + train_dataset_weights; users pointing sampler_class_path at the old example should switch to the native fields. Also remove the shim's test/loader and the README migration note. Co-Authored-By: Claude Fable 5 Signed-off-by: Avi Basnet --- examples/train/sft/README.md | 16 +++++---------- examples/train/sft/data_mixing_sampler.py | 25 ----------------------- tests/train/test_sft_dataloader.py | 8 -------- 3 files changed, 5 insertions(+), 44 deletions(-) delete mode 100644 examples/train/sft/data_mixing_sampler.py diff --git a/examples/train/sft/README.md b/examples/train/sft/README.md index 3b0fc054c2..ac6dbc22d3 100644 --- a/examples/train/sft/README.md +++ b/examples/train/sft/README.md @@ -164,9 +164,8 @@ 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). + With multiple `train_datasets` this strategy switches to weighted per-source + mixing (see [Multi-dataset training](#multi-dataset-training-and-evaluation)). - `sampler=sequential` — iterates the dataset in order ([`StatefulSequentialSampler`](../../../skyrl/train/dataset/samplers.py)). - `sampler=custom` — loads your own stateful sampler from `sampler_class_path`, @@ -213,8 +212,9 @@ bash examples/train/sft/run_sft_megatron.sh \ ``` 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): +multiple datasets and the default `sampler=random`, batches are drawn with +weighted per-source sampling (see +[`samplers.py`](../../../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 @@ -259,12 +259,6 @@ bash examples/train/sft/run_sft_megatron.sh \ (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. diff --git a/examples/train/sft/data_mixing_sampler.py b/examples/train/sft/data_mixing_sampler.py deleted file mode 100644 index 4c90dc278d..0000000000 --- a/examples/train/sft/data_mixing_sampler.py +++ /dev/null @@ -1,25 +0,0 @@ -"""Deprecated location: ``DataMixingSampler`` is now part of core SkyRL. - -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 \ - '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 skyrl.train.dataset.samplers import DataMixingSampler - -__all__ = ["DataMixingSampler"] diff --git a/tests/train/test_sft_dataloader.py b/tests/train/test_sft_dataloader.py index c6502f0f5f..826cc1fd66 100644 --- a/tests/train/test_sft_dataloader.py +++ b/tests/train/test_sft_dataloader.py @@ -112,11 +112,6 @@ def _load_curriculum_cls(): return _load_example_cls("curriculum_sampler.py", "CurriculumLearningSampler") -def _load_data_mixing_cls(): - """Import the example DataMixingSampler by file path.""" - return _load_example_cls("data_mixing_sampler.py", "DataMixingSampler") - - def _make_trainer(**overrides) -> SFTTrainer: """Build a bare SFTTrainer (no setup/Ray/GPU) for dataloader-building tests. @@ -550,9 +545,6 @@ def test_num_samples_defaults_to_dataset_length(self): sampler = DataMixingSampler(self.DATA, lengths=self.LENGTHS, weights=[0.5, 0.5]) assert len(sampler) == 20 - def test_example_file_reexports_core_class(self): - assert _load_data_mixing_cls() is DataMixingSampler - def test_validates_lengths_sum(self): with pytest.raises(ValueError, match="must equal len"): DataMixingSampler(self.DATA, lengths=[10, 5], weights=[0.5, 0.5]) From a4f45dd267c39819157039a01c9b503d7e78f9ca Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Fri, 17 Jul 2026 18:22:45 +0000 Subject: [PATCH 4/6] [refactor][SFT] Rewrite journey-style comments per contributing guide Address review feedback (#1883): comments and docstrings should describe what the code does, not the goal or how the code came to be. Reword the DataMixingSampler position-only-state warning/comment, the test-module docstring, and the position-only-state test to describe behavior directly. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Avi Basnet --- skyrl/train/dataset/samplers.py | 8 ++++---- tests/train/test_sft_dataloader.py | 12 ++++++------ 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/skyrl/train/dataset/samplers.py b/skyrl/train/dataset/samplers.py index 6939f9242b..5afec0cdbd 100644 --- a/skyrl/train/dataset/samplers.py +++ b/skyrl/train/dataset/samplers.py @@ -200,11 +200,11 @@ def load_state_dict(self, state: dict) -> None: if "generator_state" in state: self._generator.set_state(state["generator_state"]) else: - # Checkpoint from the old example sampler (position-only): keep the - # freshly-seeded generator and resume the cursor best-effort. + # Position-only state (no generator state): keep the freshly-seeded + # generator and resume the cursor best-effort. logger.warning( - "DataMixingSampler: checkpoint has no generator_state (old format); " - "resuming position only -- the restored plan may differ from the original run." + "DataMixingSampler: checkpoint has no generator_state; " + "resuming position only -- the restored plan may differ from the run that saved it." ) self._plan = None self._plan_gen_state = None diff --git a/tests/train/test_sft_dataloader.py b/tests/train/test_sft_dataloader.py index 826cc1fd66..36eefd7427 100644 --- a/tests/train/test_sft_dataloader.py +++ b/tests/train/test_sft_dataloader.py @@ -16,10 +16,9 @@ - DataMixingSampler draws a fresh plan every epoch and resumes exactly (mid-epoch and across epoch boundaries) via its RNG state -The custom-sampler path is exercised with small test-local samplers (curriculum -learning remains a user-supplied example under -``examples/train/sft/curriculum_sampler.py``, loaded by file path here to keep -it covered). +The custom-sampler path is exercised with small test-local samplers. The +``CurriculumLearningSampler`` example under +``examples/train/sft/curriculum_sampler.py`` is loaded by file path. Run:: @@ -632,8 +631,9 @@ def test_between_epoch_resume_draws_same_next_epoch(self): resumed.load_state_dict(state) assert list(resumed) == epoch2 - def test_old_format_position_only_state_loads(self): - # Checkpoints from the old example sampler carry only the position. + def test_position_only_state_loads(self): + # A position-only state (no generator_state key) loads with a warning + # and resumes the cursor against the freshly-seeded generator. sampler = self._make() sampler.load_state_dict({"position": 5}) assert sampler.position == 5 From 209ff2577f2a135e8f504df0031c3b2a4c59511f Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Fri, 17 Jul 2026 18:34:26 +0000 Subject: [PATCH 5/6] removed claude comments Signed-off-by: Avi Basnet --- skyrl/train/config/sft_config.py | 4 ---- skyrl/train/sft_trainer.py | 4 ---- tests/train/test_sft_dataloader.py | 16 ---------------- 3 files changed, 24 deletions(-) diff --git a/skyrl/train/config/sft_config.py b/skyrl/train/config/sft_config.py index c9a44be7b1..7ba53abb3c 100644 --- a/skyrl/train/config/sft_config.py +++ b/skyrl/train/config/sft_config.py @@ -320,10 +320,6 @@ def _normalize_dataset_cfg(cfg: SFTConfig) -> None: """Translate the deprecated single-dataset fields into the list-based fields and validate the dataset configuration. - Idempotent, so it is safe to call from both ``validate_sft_cfg`` and - ``SFTTrainer.__init__`` (the latter covers trainers constructed with a - pre-built ``skyrl_cfg``, which bypass ``build_skyrl_config_for_sft``). - Post-conditions: - ``train_datasets``/``train_dataset_splits`` are equal-length non-empty lists; diff --git a/skyrl/train/sft_trainer.py b/skyrl/train/sft_trainer.py index 5ee5f5e726..46530ac2f8 100644 --- a/skyrl/train/sft_trainer.py +++ b/skyrl/train/sft_trainer.py @@ -807,10 +807,6 @@ def __init__( callbacks: Optional[list[TrainingCallback]] = None, ): self.sft_cfg = cfg - # Translate deprecated single-dataset fields into the list-based ones. - # Also runs inside validate_sft_cfg (idempotent); calling it here covers - # trainers constructed with a pre-built ``skyrl_cfg``, which skip - # ``build_skyrl_config_for_sft``. _normalize_dataset_cfg(cfg) # Accept a pre-built bridge config to avoid redundant rebuilds. # When not provided (e.g. standalone usage), build it here. diff --git a/tests/train/test_sft_dataloader.py b/tests/train/test_sft_dataloader.py index 36eefd7427..f8328035b9 100644 --- a/tests/train/test_sft_dataloader.py +++ b/tests/train/test_sft_dataloader.py @@ -1,21 +1,5 @@ """CPU tests for the SFT stateful dataloader and samplers. -Covers the stateful-dataloader RFC test plan: - - default random dataloader order is deterministic for identical seeds - - random dataloader order differs across different seeds - - sequential sampler yields examples in order - - sequential sampler state_dict/load_state_dict resumes at the next sample - - custom sampler state is included in the dataloader checkpoint - - data-mixing / curriculum samplers resume correctly - -and the multi-dataset RFC (#1875) test plan: - - user-specified train_dataset_weights: per-batch source representation - tracks the weights - - default weights: equal mixing across datasets - - mixing ratios are independent of the individual dataset sizes - - DataMixingSampler draws a fresh plan every epoch and resumes exactly - (mid-epoch and across epoch boundaries) via its RNG state - The custom-sampler path is exercised with small test-local samplers. The ``CurriculumLearningSampler`` example under ``examples/train/sft/curriculum_sampler.py`` is loaded by file path. From 62b53360a5c2f8a325a5742ae944101955ce0405 Mon Sep 17 00:00:00 2001 From: Avi Basnet Date: Fri, 17 Jul 2026 20:44:31 +0000 Subject: [PATCH 6/6] [refactor][SFT] Use key="[...]" quoting for list overrides in examples Address review feedback (#1883): quote the value rather than the whole key=value token in the dataset list overrides across the SFT example scripts, the e2e script, and the README. This is shell-equivalent to the previous "key=[...]" form (both strip to the same argv string) and reads cleaner. Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Avi Basnet --- examples/train/sft/README.md | 19 +++++++++---------- examples/train/sft/run_sft_fsdp.sh | 4 ++-- examples/train/sft/run_sft_lora.sh | 4 ++-- examples/train/sft/run_sft_megatron.sh | 4 ++-- .../train/sft/run_sft_megatron_apigen_mt.sh | 8 ++++---- .../sft/run_sft_megatron_multi_dataset.sh | 14 ++++++-------- .../train/sft/run_sft_megatron_tulu3_50k.sh | 8 ++++---- .../sft/run_sft_megatron_tulu3_50k_packed.sh | 8 ++++---- examples/train/sft/run_sft_megatron_vlm.sh | 4 ++-- .../train/gpu_e2e_test/sft_tulu3_megatron.sh | 2 +- 10 files changed, 36 insertions(+), 39 deletions(-) diff --git a/examples/train/sft/README.md b/examples/train/sft/README.md index ac6dbc22d3..d0135db8fb 100644 --- a/examples/train/sft/README.md +++ b/examples/train/sft/README.md @@ -199,16 +199,15 @@ bash examples/train/sft/run_sft_megatron.sh \ ## Multi-dataset training and evaluation 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): +`train_datasets`/`train_dataset_splits`: ```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]" + 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 @@ -228,9 +227,9 @@ 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_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 `/` @@ -249,8 +248,8 @@ kwarg. Explicit `train_dataset_weights` are rejected outside `sampler=random` ```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]']" \ + 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.curriculum_sampler.CurriculumLearningSampler \ 'sampler_kwargs={num_samples: 40, seed: 42}' diff --git a/examples/train/sft/run_sft_fsdp.sh b/examples/train/sft/run_sft_fsdp.sh index df9ca7e309..8bcd1add2d 100755 --- a/examples/train/sft/run_sft_fsdp.sh +++ b/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 \ - "train_datasets=['yahma/alpaca-cleaned']" \ - "train_dataset_splits=['train[:100]']" \ + train_datasets="['yahma/alpaca-cleaned']" \ + train_dataset_splits="['train[:100]']" \ messages_key=messages \ max_length=512 \ num_steps=10 \ diff --git a/examples/train/sft/run_sft_lora.sh b/examples/train/sft/run_sft_lora.sh index 8d9ca604e8..f17658f3c4 100644 --- a/examples/train/sft/run_sft_lora.sh +++ b/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 \ - "train_datasets=['yahma/alpaca-cleaned']" \ - "train_dataset_splits=['train[:100]']" \ + train_datasets="['yahma/alpaca-cleaned']" \ + train_dataset_splits="['train[:100]']" \ messages_key=messages \ max_length=512 \ num_steps=10 \ diff --git a/examples/train/sft/run_sft_megatron.sh b/examples/train/sft/run_sft_megatron.sh index 82101ded50..dc24397a05 100755 --- a/examples/train/sft/run_sft_megatron.sh +++ b/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 \ - "train_datasets=['yahma/alpaca-cleaned']" \ - "train_dataset_splits=['train[:100]']" \ + train_datasets="['yahma/alpaca-cleaned']" \ + train_dataset_splits="['train[:100]']" \ messages_key=messages \ max_length=512 \ num_steps=10 \ diff --git a/examples/train/sft/run_sft_megatron_apigen_mt.sh b/examples/train/sft/run_sft_megatron_apigen_mt.sh index 44c0db12b2..a022f8c7fc 100755 --- a/examples/train/sft/run_sft_megatron_apigen_mt.sh +++ b/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 \ - "train_datasets=['$DATA_DIR']" \ - "train_dataset_splits=['train[:4000]']" \ - "eval_datasets=['$DATA_DIR']" \ - "eval_dataset_splits=['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 \ diff --git a/examples/train/sft/run_sft_megatron_multi_dataset.sh b/examples/train/sft/run_sft_megatron_multi_dataset.sh index 167d827d46..ed24d5c852 100755 --- a/examples/train/sft/run_sft_megatron_multi_dataset.sh +++ b/examples/train/sft/run_sft_megatron_multi_dataset.sh @@ -8,8 +8,6 @@ set -x # 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...] # @@ -20,12 +18,12 @@ 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]" \ + 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 \ diff --git a/examples/train/sft/run_sft_megatron_tulu3_50k.sh b/examples/train/sft/run_sft_megatron_tulu3_50k.sh index cda6cbde3c..d3e5eef55c 100755 --- a/examples/train/sft/run_sft_megatron_tulu3_50k.sh +++ b/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 \ - "train_datasets=['allenai/tulu-3-sft-mixture']" \ - "train_dataset_splits=['train[:50000]']" \ - "eval_datasets=['allenai/tulu-3-sft-mixture']" \ - "eval_dataset_splits=['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 \ diff --git a/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh b/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh index 85ff3698f0..b11e5f80f7 100755 --- a/examples/train/sft/run_sft_megatron_tulu3_50k_packed.sh +++ b/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 \ - "train_datasets=['allenai/tulu-3-sft-mixture']" \ - "train_dataset_splits=['train[:50000]']" \ - "eval_datasets=['allenai/tulu-3-sft-mixture']" \ - "eval_dataset_splits=['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 \ diff --git a/examples/train/sft/run_sft_megatron_vlm.sh b/examples/train/sft/run_sft_megatron_vlm.sh index eeb3f9ed6d..f7eb10d897 100755 --- a/examples/train/sft/run_sft_megatron_vlm.sh +++ b/examples/train/sft/run_sft_megatron_vlm.sh @@ -33,8 +33,8 @@ uv run --isolated --extra megatron \ python -m skyrl.train.main_sft \ strategy=megatron \ model.path=Qwen/Qwen3-VL-2B-Instruct \ - "train_datasets=['$DATA_DIR']" \ - "train_dataset_splits=['train']" \ + train_datasets="['$DATA_DIR']" \ + train_dataset_splits="['train']" \ messages_key=messages \ max_length=4096 \ num_steps=30 \ diff --git a/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh b/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh index b8675b4cda..9a2a1fa713 100755 --- a/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh +++ b/tests/train/gpu_e2e_test/sft_tulu3_megatron.sh @@ -34,7 +34,7 @@ LOG_FILE="${LOG_FILE:-/tmp/${RUN_NAME}.log}" # * batch_size=8, micro_train_batch_size_per_gpu=2 are sized for L4_ci (4 GPUs). bash examples/train/sft/run_sft_megatron_tulu3_50k.sh \ num_steps=$NUM_STEPS \ - "train_dataset_splits=['train[:2000]']" \ + train_dataset_splits="['train[:2000]']" \ batch_size=8 \ micro_train_batch_size_per_gpu=2 \ max_length=1024 \