[feat][SFT] Support multiple datasets for training and evaluation in SFTTrainer#1883
[feat][SFT] Support multiple datasets for training and evaluation in SFTTrainer#1883avigyabb wants to merge 2 commits into
Conversation
…SFTTrainer Implements RFC NovaSky-AI#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 NovaSky-AI#1875 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
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 <noreply@anthropic.com> Signed-off-by: Avi Basnet <avigyabb@stanford.edu>
|
Ran the GPU e2e test (
The first e2e attempt caught a real bug, fixed in fe51254: OmegaConf/YAML can't parse HF slice syntax unquoted inside a list override ( 🤖 Generated with Claude Code |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces native support for multi-dataset training and evaluation in SFT by transitioning from single-dataset fields to list-based configurations (such as train_datasets and eval_datasets). It also integrates the DataMixingSampler into the core library to handle weighted multi-source mixing natively. The review feedback highlights a few important robustness improvements: first, validating dataset weights defensively to prevent silent passes when NaN values are present; and second, guarding against None values in sampler_kwargs during unpacking to avoid potential TypeError crashes.
| 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)}." | ||
| ) |
There was a problem hiding this comment.
The check any(w < 0 for w in weights) or sum(weights) <= 0 is not robust against NaN values. If weights contains a NaN value, w < 0 evaluates to False and sum(weights) <= 0 also evaluates to False, allowing the check to pass silently. Using any(not (w >= 0) for w in weights) or not (sum(weights) > 0) is a more robust, defensive programming approach that correctly catches both negative values and NaNs.
| 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)}." | |
| ) | |
| if any(not (w >= 0) for w in weights) or not (sum(weights) > 0): | |
| raise ValueError( | |
| f"DataMixingSampler: weights must be non-negative with a positive sum, got {list(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}.") |
There was a problem hiding this comment.
The check any(w <= 0 for w in cfg.train_dataset_weights) is not robust against NaN values. If train_dataset_weights contains a NaN value, w <= 0 evaluates to False, allowing the check to pass silently. Using any(not (w > 0) for w in cfg.train_dataset_weights) is a more robust, defensive programming approach that correctly catches both non-positive values and NaNs.
| 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}.") | |
| if any(not (w > 0) for w in cfg.train_dataset_weights): | |
| raise ValueError(f"train_dataset_weights must all be > 0, got {cfg.train_dataset_weights}.") |
| 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) |
There was a problem hiding this comment.
If self.sft_cfg.sampler_kwargs is explicitly set to None (or parsed as null from a YAML configuration), unpacking it with **sampler_kwargs will raise a TypeError. Guarding against None by defaulting to an empty dictionary (self.sft_cfg.sampler_kwargs or {}) is a safer, more defensive programming practice.
| 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) | |
| sampler_kwargs = self.sft_cfg.sampler_kwargs or {} | |
| if multi_dataset: | |
| # User-provided kwargs win over the injected lengths. | |
| sampler_kwargs = {"lengths": dataset_lengths, **sampler_kwargs} | |
| return sampler_cls(tokenized, **sampler_kwargs) |
SumanthRH
left a comment
SumanthRH
left a comment
There was a problem hiding this comment.
Looking good! Left a few minor comments
| "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]" \ |
There was a problem hiding this comment.
For my understanding: OmegaConf will fail to parse the below override?
train_datasets="['allenai/tulu-3-sft-mixture','yahma/alpaca-cleaned']" \| 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). |
There was a problem hiding this comment.
Can you please make sure that Claude doesn't add such "in the process" comments?
Comments should describe what the code is doing, not what the goal was or what learnings-on-the-journey the agent stumbled into.
See examples here: https://github.com/NovaSky-AI/SkyRL/blob/main/.claude/docs/contributing.md#comments
2 participants
Implements RFC #1875: native
SFTTrainersupport for training on a weighted mixture of multiple datasets and evaluating on multiple named datasets, as a thin layer on top of theStatefulDataLoader/sampler infrastructure from #1842.What's in this PR
Config (
skyrl/train/config/sft_config.py)train_datasets/train_dataset_splits/train_dataset_weightsandeval_datasets/eval_dataset_splits/eval_dataset_names.dataset_name/dataset_splitandeval_dataset_name/eval_dataset_splitare deprecated: still accepted, translated to the list form with aDeprecationWarning(both CLI and programmatic paths), and rejected when combined with the new fields.1/Nequal mixing), eval names unique; explicittrain_dataset_weightsonly valid withsampler=random(custom samplers take ratios viasampler_kwargs).DataMixingSampler promoted to core (
skyrl/train/dataset/samplers.py)examples/train/sft/data_mixing_sampler.py(which now re-exports it, so existingsampler_class_pathconfigs keep working)._planonce and replayed it), andstate_dictnow includes the generator state, so mid-epoch and epoch-boundary checkpoint resumes reproduce the uninterrupted sample stream exactly. Old position-only states load best-effort with a warning.Trainer (
skyrl/train/sft_trainer.py)_load_and_tokenizepath (per-(name, split)cache keys preserved), then concatenated with per-dataset lengths tracked.sampler=randomstill uses the dataloader's built-in shuffle). Multiple datasets:sampler=randomselectsDataMixingSamplerconfigured with the tokenized lengths andtrain_dataset_weights— per-batch source ratios, independent of dataset sizes.lengthsvia an injected kwarg (user-suppliedsampler_kwargswin).eval/{name}/loss— nested even with a single eval dataset so runs with and without mixing chart the same keys.data.ptdataloader state; no new mechanism.The eval metric key
eval/eval_lossis noweval/{name}/loss(intentional per the RFC). Dashboards/callbacks keying on the old name need updating.Tests (CPU-only, per the RFC test plan)
[0.8, 0.2]→ ~80/20), default weights give equal mixing, and ratios are independent of dataset sizes (20-vs-180 sources still ~50/50).StatefulDataLoader), old-format state compatibility.test_sft_config.py; callback test updated to the nested eval keys.Docs & examples
examples/train/sft/README.md: new Multi-dataset training and evaluation section, config reference updates, breaking-change note; removed the single-dataset limitation.examples/train/sft/run_sft_megatron_multi_dataset.sh; existing SFT run scripts (and the GPU e2e script) migrated to the list fields so layeringtrain_datasets=[...]overrides on them doesn't conflict with the deprecated fields they used to set.Testing
tests/train/test_sft_dataloader.py,test_sft_config.py,test_sft_callbacks.py: 106 passed. Broader CPU suite (tests/train/+tests/backends/skyrl_train/, GPU dirs excluded): 965 passed; the 3 remaining failures are missingvllm/torchvision optional deps in the local env, unrelated to this change.Closes #1875
🤖 Generated with Claude Code