Mixing Multiple Datasets

SFTTrainer can train on a weighted mixture of several datasets and evaluate on multiple held-out datasets separately -- no manual preprocessing or concatenation required. This page covers the list-based dataset config and how mixing weights behave.

Weighted multi-source training

Pass parallel lists to train_datasets and train_dataset_splits, plus an optional train_dataset_weights:

bash examples/train/sft/run_sft_megatron_multi_dataset.sh
# or directly on top of any SFT script:
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, the trainer switches to weighted per-source sampling via DataMixingSampler.

How weights work

train_dataset_weights is the approximate per-batch ratio of samples drawn 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 need not sum to 1 (relative scale is what matters) and default to equal mixing (1/N each).

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: a mid-epoch resume replays the in-flight plan, and all subsequent epochs match an uninterrupted run.

Constraints

Each training dataset is independently detected as either chat format (messages_key) or Alpaca format (instruction / input / output), so the two text formats can be mixed. Chat-format datasets use the configured messages_key; the optional tools_key and system_key columns need not be present in every source.

All datasets must use the same modality (all-text or all-image), because a training batch cannot mix text-only and image-bearing examples. Image datasets must use chat format.

Multi-dataset evaluation

Evaluation accepts multiple datasets too. Each is evaluated separately and its loss is logged under eval/{name}/loss:

bash examples/train/sft/run_sft_megatron.sh \
    eval_datasets="['allenai/tulu-3-sft-mixture','yahma/alpaca-cleaned']" \
    eval_dataset_splits="['train[-500:]','train[200:700]']" \
    eval_dataset_names="['tulu3','alpaca']" \
    eval_interval=5

• eval_dataset_names is optional and only affects metric names. It defaults to each dataset name with / replaced by _, and must be unique (set it explicitly to disambiguate the same dataset used with two different splits).
• eval_interval runs eval every N steps; eval also runs once at the end of training. eval_before_train=true logs a baseline pass at step 0.

Config reference

See SFTConfig in the API reference for the full documentation of the dataset and evaluation fields: train_datasets, train_dataset_splits, train_dataset_weights, eval_datasets, eval_dataset_splits, eval_dataset_names, eval_interval, eval_before_train, and messages_key / tools_key / system_key.