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

[avigyabb]

Implements RFC #1875: native SFTTrainer support for training on a weighted mixture of multiple datasets and evaluating on multiple named datasets, as a thin layer on top of the StatefulDataLoader/sampler infrastructure from #1842.

What's in this PR

Config (skyrl/train/config/sft_config.py)

• New list-based fields: train_datasets / train_dataset_splits / train_dataset_weights and eval_datasets / eval_dataset_splits / eval_dataset_names.
• dataset_name/dataset_split and eval_dataset_name/eval_dataset_split are deprecated: still accepted, translated to the list form with a DeprecationWarning (both CLI and programmatic paths), and rejected when combined with the new fields.
• Validation: list lengths must match, weights strictly positive (default 1/N equal mixing), eval names unique; explicit train_dataset_weights only valid with sampler=random (custom samplers take ratios via sampler_kwargs).

DataMixingSampler promoted to core (skyrl/train/dataset/samplers.py)

• Moved from examples/train/sft/data_mixing_sampler.py (which now re-exports it, so existing sampler_class_path configs keep working).
• Fixes vs the example, per the RFC: draws a fresh weighted plan every epoch (the example materialized _plan once and replayed it), and state_dict now 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)

• Each training dataset is tokenized independently through the existing _load_and_tokenize path (per-(name, split) cache keys preserved), then concatenated with per-dataset lengths tracked.
• Single dataset: behavior completely unchanged (sampler=random still uses the dataloader's built-in shuffle). Multiple datasets: sampler=random selects DataMixingSampler configured with the tokenized lengths and train_dataset_weights — per-batch source ratios, independent of dataset sizes.
• Custom samplers receive the per-dataset lengths via an injected kwarg (user-supplied sampler_kwargs win).
• Eval: one dataloader per eval dataset, evaluated separately, logged under eval/{name}/loss — nested even with a single eval dataset so runs with and without mixing chart the same keys.
• Checkpoint/resume reuses the existing data.pt dataloader state; no new mechanism.

⚠️ Breaking change

The eval metric key eval/eval_loss is now eval/{name}/loss (intentional per the RFC). Dashboards/callbacks keying on the old name need updating.

Tests (CPU-only, per the RFC test plan)

• Batch representation tracks user-specified weights ([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).
• Core sampler: fresh plan per epoch, deterministic multi-epoch stream per seed, exact resume mid-epoch and across epoch boundaries (sampler-level and through StatefulDataLoader), old-format state compatibility.
• Config normalization/deprecation/validation coverage in 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.
• New examples/train/sft/run_sft_megatron_multi_dataset.sh; existing SFT run scripts (and the GPU e2e script) migrated to the list fields so layering train_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 missing vllm/torchvision optional deps in the local env, unrelated to this change.

Closes #1875

🤖 Generated with Claude Code

[avigyabb]

Ran the GPU e2e test (tests/train/gpu_e2e_test/sft_tulu3_megatron.sh) on 4x L40S against this branch — all CI assertions passed ✅

• 100/100 steps completed; train/loss trended down (2.80 @ step 20 → 2.19 @ step 100), no NaN/inf.
• Final eval logged under the new nested key: eval/allenai_tulu-3-sft-mixture/loss = 1.9843 (62 batches) — exercising the default sanitized eval_dataset_names and the eval/{name}/loss scheme end-to-end.
• The run also exercises the migrated script path: the e2e layers "train_dataset_splits=['train[:2000]']" over the base script's list-based fields.
• wandb run: https://wandb.ai/sky-posttraining-uc-berkeley/skyrl_sft_ci/runs/rii88p8v

The first e2e attempt caught a real bug, fixed in fe51254: OmegaConf/YAML can't parse HF slice syntax unquoted inside a list override (train_dataset_splits=[train[:2000]] → ParserError on the nested brackets). All scripts/README examples now quote each element, with a CPU regression test encoding the syntax.

🤖 Generated with Claude Code