[feat][SFT] Union sampling for multi-dataset runs without explicit weights

[avigyabb]

What

Multiple train_datasets (or pretokenized_dataset_paths) with no train_dataset_weights now sample as the union of the datasets: one shuffle over the concatenation per epoch, without replacement -- every sample of every dataset is seen exactly once per epoch, and each dataset's share of the batches follows its size.

Explicit weights behave exactly as before: weighted per-source mixing via DataMixingSampler, with replacement (the only way per-batch ratios can hold independent of dataset sizes).

Why

Previously, unset weights were silently defaulted to equal mixing (1/N per source, with replacement). Two consequences surprised users who mentally modeled "I gave it two datasets, it trains on their union":

• A 10k-row source mixed with a 1M-row source got 50/50 batch representation by default -- the small source oversampled ~100x, with no signal that rebalancing was happening.
• num_epochs=1 was len(union) weighted draws, not a pass: some rows repeated, others were never seen (~37% missed even in the equal-size case), so "1 epoch" did not mean "saw everything once".

With this change the two sampling semantics map onto explicit user intent:

Config	Semantics	Replacement	Epoch means
weights unset (default)	union of the datasets	without	every row of every dataset exactly once
weights set	rebalanced mix per the ratios	with	len(union) weighted draws

How

Small diff -- the machinery already existed:

• _normalize_mixing_weights no longer fills [1/N] * N; unset weights stay None as the union-sampling signal (explicit weights are validated as before).
• build_train_sampler returns None (the dataloader's built-in stateful shuffle=True -- a randperm over the ConcatSFTDataset) for sampler="random" unless weights are set. Checkpoint/resume is unchanged: the built-in shuffle path is the same one single-dataset runs already use.
• Docs (multi_dataset.mdx, overview.mdx, README) updated to describe both modes and the replacement semantics of each.

⚠️ Behavior change

Multi-dataset runs that relied on the implicit equal-weight default will now train on the union instead of a 1/N mix. To keep the old behavior, set the weights explicitly, e.g. train_dataset_weights="[0.5,0.5]".

Tests

• Config: unset weights stay None (no defaulted mixing); explicit-weights validation unchanged.
• Sampler dispatch: multi + no weights -> None (union); multi + weights -> DataMixingSampler.
• Coverage: one epoch over two unequal sources yields every index of both exactly once, shuffled, with a fresh permutation the next epoch.
• Trainer-level: two pretokenized stores without weights -> union; with weights -> mixing.
• Full tests/train suite green: 893 passed.

🤖 Generated with Claude Code