feat(trainer): fold rl/ref_kl loss types into one pg skeleton, vectorize compute_loss

[hallerite]

Motivation

default_loss_fn (DPPO), ipo_loss_fn, and ref_kl_loss_fn are three copies of the same per-token loss

-(keep_mask · adv_tau·A · importance_ratio) + kl_tau · loss_mask · log_ratio²

differing only in three parameters:

	advantage A	trust region	kl_tau
default	shipped stream (data)	sign-conditioned DPPO	config
ipo	shipped stream (data)	symmetric threshold	config
ref_kl	(ref_lp − trainer_lp).detach(), recomputed every forward	one-sided (−0.2)	hard-coded 1e-3

This is the REINFORCE identity: ∇KL(π_θ‖π_ref) = −E[(log π_ref − log π_θ)·∇log π_θ] — reverse-KL distillation is policy gradient with a live advantage provider. Keeping three drifting copies hides that, and hides that "which trust region / which κ should OPD use" is a parameter choice, not a structural difference.

Separately, compute_loss looped over packed sequences, issuing O(sequences × components) kernel launches plus a bool(mask.any()) device sync per component per sequence. Benchmarked on a 32k-token micro-batch at 256 seqs/bin, that loop costs 120–190 ms of exposed critical-path time — and none of it hides behind the forward, because the driver launch queue only absorbs ~1–2k pending launches.

Changes

• PGParams — frozen dataclass selecting one instance of the skeleton: advantage: "shipped" | "ref_kl", trust_region: "dppo" | "ipo" | "one_sided", adv_tau, kl_tau, thresholds, metric prefix.
• pg_loss_fn — the single skeleton. default_loss_fn, ipo_loss_fn, and ref_kl_loss_fn remain as thin adapters, so setup_rl_loss_fn, trainer.loss config handling, and all existing imports are unchanged. ref_kl's hard-coded 1e-3 is now a visible parameter.
• compute_loss — concatenates the bin once and runs each component exactly once, instead of once per packed sequence. The per-sequence loop, the per-sequence bool(mask.any()) D2H syncs, and the graph anchor are gone: the rl component always runs (an empty mask contributes an exact zero), so the loss is always backward-able and every rank runs backward — FSDP collectives stay in sync by construction.

No config, wire, or orchestrator changes. The rl/ce/ref_kl streams and their semantics are untouched — the fold is trainer-internal.

Verification

• Old-vs-new equivalence harness over 18 archetype × config combinations (hot path / weighted grpo / opd / echo / sdft-style rl+ref_kl overlap / all-empty × 2 DefaultLossConfigs × IPO): gradients bit-identical in every case (per-token gradients are unaffected by summation order); loss values match to fp reorder (~1e-7 relative).
• tests/unit/train/rl/test_loss.py: 10/10 pass unmodified, including exact-value CE assertions, the empty-components backward test, and hot-path/explicit-ones equality.

Performance (RTX PRO 6000, real compute_loss + backward)

measurement	before	after
32k micro-batch, 256 seqs/bin: hot / opd / echo	149 / 123 / ~205 ms	1.6 / 2.0 / 1.8 ms
exposed critical-path cost pipelined behind a 416 ms forward	30–190 ms	1.2–4 ms
2048 × 131k step, dense 256-seq bins	306–403 s	5.2–6.5 s
2048 × 131k step, 1 seq/bin (long context)	~2–3 s	~2–3 s

Cost is now flat in packing density, as a per-token loss should be.

Behavioral notes

• Metrics: components with empty masks now emit zero-valued metrics instead of being skipped, and logged loss metrics are one token-weighted value per micro-batch rather than per-sequence means stacked — slightly different curves on dashboards, same information.
• Custom loss fns (CustomLossConfig): the fn now receives the packed sequences concatenated in a single call per micro-batch instead of one call per sequence. Per-token-sum losses (like all built-ins) are unaffected; a custom fn doing per-sequence normalization would change behavior.
• Residual syncs: _safe_mean and ce's boolean indexing still trigger nonzero syncs, but a fixed handful per micro-batch instead of O(sequences × components); included in the 1–4 ms above. Replacing them with mask-multiply reductions is a follow-up if we want the loss torch.compile-able.
• Trust-region shapes and coefficients are preserved exactly per loss type — no behavioral unification (e.g. giving ref_kl the sign-conditioned DPPO mask) is smuggled in; that would be a separate, A/B-tested change.

🤖 Generated with Claude Code