[kritarth-simile]

Summary

This PR adds SFT on pre-packed multiple-choice data with a pluggable loss, plus forward-only choice evals during training. It also fixes a gradient-clipping bug that affects any SFT run with validation.

Our use case: each example is a prompt and a target probability distribution over its answer choices, with a per-example weight. The loss compares the model's softmax over those choice tokens with the target. Stock SFT has a single target text per row, full-vocabulary cross-entropy and no per-example weights, so it cannot express this. The loss itself stays in user code, behind [loss] import_path, the way RL already supports custom losses.

Changes

Fix: gradient clipping after validation.

• Problem: validation runs between backward and clip_grad_norm_. The [lm_head, norm] FSDP2 group (reshard_after_forward=False) stays unsharded after the validation forward, so clipping sees 112 of 114 parameters.
• Symptom: on validation steps the logged grad norm is 13.7% low on GLM-5.3 (66.6 instead of 77.2), and about 1% low on Qwen3.5-9B.
• Fix: no_grad_forwards(model) reshards after no-grad forwards and wraps validation and choice-eval forwards. Grad norms then agree within 0.04%.

Packed-choice SFT (data.type = "packed_choice").

• Reads pre-tokenized bins (Arrow IPC plus manifest.json, schema simile-packed-choice/v2). Each bin is one micro batch. Each example carries its answer position, its candidate token ids, target probabilities, a weight and loss-mix weights.
• An LM-head mode returns logits only for the candidate tokens at the marked positions, so the full-vocabulary logits are never materialized.
• [loss] import_path / kwargs (SFTChoiceLossConfig): called per micro batch with right-padded choice logits, targets and weights, and returns (loss_sum, weight_sum). Gradients are normalized by the step's global weight_sum.
• max_steps defaults to one pass over the packed bins.
• val.data can also be packed_choice.

Choice evals during training ([choice_eval]).

• Every interval steps (and optionally at step 0 and at the end), runs no-grad forwards over pre-packed eval sets. Rank 0 gathers each row's candidate logits into <run_dir>/choice_evals/<name>/step_<step>/predictions.arrow.
• With import_path set, it calls a user scoring function and logs its metrics at that step. Without one, it only writes the predictions, for scoring outside the trainer, and logs eval time and peak memory.
• A metrics.json marks a scored step, so resumed or requeued runs do not score a step twice.
• A failing scorer is logged and never stops training.

New config surface is opt-in. Existing SFT configs and behavior are unchanged.

Test plan

• New unit tests:

  • tests/unit/train/sft/test_packed_choice.py
  • tests/unit/train/sft/test_choice_eval.py, which covers the no-scoring-function path and config validation
  • tests/unit/train/test_no_grad_forwards.py

  56 passed on macOS with GPU-only modules stubbed.

• Existing tests, tests/unit/train/sft, tests/unit/test_configs.py and tests/unit/train/test_model.py: same result as main on the same machine. 330 passed; the one failure, test_model.py::test_moe_custom_impl, also fails on main there.

• ruff check and ruff format --check with the pre-commit pin (0.13.0).

• GPU, on 8×H200, at 7f9de649aa80. This is the same code minus the optional scoring function. The only commit to the touched files since then is the unrelated Muon gate-sharding change, fix: shard single-row gate weights along columns for Muon #3723.

  • Qwen3.5-9B: step-0 validation loss 1.7839, against an independent reference of 1.787–1.793; step-1 training loss 1.28616 against 1.2857.
  • Eval metrics computed from the gathered predictions match our reference trainer's within 0.0015.
  • GLM-5.3-BF16 (real weights, 8 of 78 layers), EP 8, CP 2, SignSGD: trains, and resume matches an uninterrupted run within 0.0013.

• GPU run of this branch on hosted training via the source-ref overlay.