feat: SFT from a HF dataset through the RL stack (dataset_sft)

[mikasenghaas]

Summary

POC: SFT from a static HF dataset through the RL training stack, keeping the orchestrator + inference server for online evals.

• New run-level algorithm [orchestrator.algo] type = "dataset_sft": carries an SFTDatasetConfig (name, subsets/splits/probabilities/stopping_strategy, shuffle, seed, loss_mask) — the same example schema and knobs as the native SFT trainer.
• Config validators enforce the contract: train sources take no effect and must be empty ([[orchestrator.train.source]] is a config error); eval sources stay optional; dataset_sft is rejected as a per-source algorithm.
• Orchestrator: DatasetSFTSource renders dataset examples through the policy renderer into ce-routed TrainingSamples (ce_weights on the loss mask, rl_weights zeroed, dummy logprobs/temperatures). A dataset_loop runs beside the main loop: it packs and ships batches through the regular packer/transport path, while the dispatcher serves evals only. Token batching takes the largest whole-sample prefix within token_batch_size (the first overflowing sample carries to the next step). The data position ({epoch, cursor}, per-epoch reshuffle) round-trips through the orchestrator checkpoint.
• Example rendering (messages / prompt+completion resolution, tools, role loss masking) is extracted from SFTDataset._process into a shared render_example, so both SFT paths tokenize identically.
• Trainer: unchanged loss path — the ce loss component of the algorithm-blind RL trainer executes the shipped weights.

Three pieces of shared infrastructure fell out of chasing the RL trainer's step time down to the native SFT trainer's (ablation details under Verification):

• On-demand weight broadcasts. The trainer broadcast to inference after every step whether anything consumed it or not. The orchestrator knows exactly which policy versions inference will consume, so it stamps the request on the work itself: MicroBatch.sync_weights marks the steps whose weights are needed, and the trainer broadcasts exactly those. Rollout training marks every step except the final TARGET_LAG+1 (the trainer's old final-step skip, moved to the producer); dataset_sft marks eval-trigger steps and any step shipped while evals are queued or in flight — in-flight evals sample the live policy, so long-running inference keeps receiving fresh weights at the trainer's cadence, mid-generation. Producer pacing moves from policy versions to consumption acks on the ZMQ control channel (the READY-barrier socket), and shutdown holds until the last requested sync lands so the trainer is never stranded mid-NCCL-rendezvous.
• Cost-balanced packing (new default in packed_samples_into_micro_bs). First-fit-decreasing minimizes bins but co-locates the longest samples; with few bins per DP rank the hottest rank's quadratic attention cost gates every FSDP step. Packing now tries a cost-balanced variant at the bin count the token total forces anyway (each sample placed longest-first into the cheapest compatible bin, scored by the existing additive bin_cost) and keeps it unless it needs more rounded-up bins than first-fit — never more bins, strictly better balance. Rollout batches benefit the same way whenever forced bin rounding leaves few bins per rank.
• Teardown fix (multi_node_rl.sbatch.j2). wait -n tore a node down when its first background process exited, even cleanly — at max_steps the trainer's exit killed the orchestrator mid-drain of final-step evals, and a dataset_sft orchestrator finishing its ship loop early killed the trainer mid-step. Teardown now waits until every principal process (torchrun, orchestrator) has exited cleanly; any non-zero exit still tears down immediately.

Adds configs/basic/reverse-text/sft-rl.toml (dev-sized counterpart of sft.toml, same token budget per step) and examples/advanced/glm-4.5-air/{sft,sft-rl}.toml (GLM-4.5-Air on INTELLECT-3-SFT at 131k context, cp=4 ulysses + muon on 4 H200 nodes, identical 1,048,576 tokens/step on both entrypoints).

Verification

Reverse-text (Qwen3-0.6B, 100 steps × ~131k tokens, lr 2e-5, one training GPU each):

	uv run sft (1 GPU)	uv run rl + dataset_sft (1 GPU + 1 vLLM)
loss @ 25 / 50 / 75 / 100	2.78 / 1.63 / 1.12 / 0.945	2.98 / 1.60 / 1.16 / 0.879 (nll/mean)
real samples trained	11,191	12,919 (zero-padding FFD packing vs ~14% tail padding in fixed rows)
online evals	—	5 epochs, reward 0.066 → 0.117

GLM-4.5-Air at 131k (INTELLECT-3-SFT swe_swiss+am_if+toucan_tool, 4×8×H200 training, cp=4 ulysses, muon, 1,048,576 tokens/step). The step-time ablation that motivated the design (same batches, broadcast-free medians):

variant	median step
native uv run sft	36.1 s
rl dataset_sft, per-step broadcast + first-fit packing	47.2 s (+31%)
− per-step broadcast	40.7 s
− loss-pipeline extras (entropy, per-sequence loop, stat syncs)	41.4 s (no effect — dropped)
+ cost-balanced packing	34.6 s (−4%)

Broadcast cost and packing imbalance were the whole gap: measured per-rank attention cost (max/mean Σ len²) is first-fit 2.68, native stream-order rows 1.95, cost-balanced 1.64 — which is why the balanced RL trainer edges out the native trainer per step. The loss-pipeline extras cost nothing measurable and stay.

End-to-end on-demand broadcast validation (GLM-4.5-Air 131k, 10 steps, AIME'25 eval every 5 steps): broadcasts landed at exactly v0 (startup), v1–v3 and v5 (batches shipped while evals were in flight — mid-generation weight updates at the trainer's cadence), v4 and v9 (eval-trigger steps), v10 (mid-final-eval); none on the idle steps 6–8. The shutdown hold released 0.2 s after the trainer's final broadcast, and all four training nodes tore down cleanly on their own.

🤖 Generated with Claude Code