Triton monolithic head-loss kernel (#507)

[jlamypoirier]

Claude Opus 4.8 (1M context) note: authored this PR.

Follow-up to #549, stacked on jlp_monolithic_head_loss. Gives the monolithic loss a single Triton kernel for its shared-softmax path, alongside the existing torch.compile path.

What it does

The monolithic loss can now run its shared softmax on one @triton_jit kernel via use_triton, covering the label-based objective set: cross-entropy (label), z_loss, the GRPO objective, and the GSPO objective. Any subset sharing one softmax works (CE+z, GRPO+z, GSPO+z, …), at most one of each kind.

Distillation and GRPO metrics have no Triton fused kernel yet — they are rejected at config time when use_triton=True, and fall back to the compiled path when use_triton is left at its None (auto) default. DPO stays separate as before.

Design

Every enabled loss stores its own forward scalar, but the gradients superpose into two per-row coefficients:

grad_j = prob_coeff · softmax_j − label_coeff · δ_{j, label}

so the gradient-write loop is loss-agnostic and each loss's math only contributes to the coefficients (z-loss adds to prob_coeff only; CE/GRPO/GSPO share the same label column). The kernel is gated purely by the presence of each loss's input/output pointers — no separate enable flags.

Three dispatch cases mirror the existing per-loss Triton kernels:

• No tensor parallelism, no GSPO — one fused kernel, softmax computed in-kernel.
• Tensor parallelism — a reduced forward pass (reusing the shared cross-entropy forward kernel) feeds a load-mode kernel.
• GSPO present — its eager segment seam (from Fuse the GSPO segment seam into compiled blocks (#507) #554) runs between the forward and the load-mode backward, handing the kernel a per-token coefficient that superposes with the in-kernel losses.

GSPO's seam-argument assembly is factored out of finish so the compiled and Triton paths share it.

Config

use_triton: bool | None on MonolithicLossConfig:

• True — rejects unsupported children (distillation, GRPO-with-metrics, duplicate kinds) at validate time.
• None (default) — Triton on GPU only if every child is supported, else the compiled monolithic.
• False — always compiled.

Validation

Numerically validated on CPU via the Triton interpreter (TRITON_INTERPRET=1):

• Functional CE / z / CE+z / GRPO / GRPO+z against the PyTorch and standalone-Triton references (rms-relative ~1e-7).
• test_lm_head fused_triton configs — 36/36 (single-pass and the GSPO seam, across masking / splits / bf16 / logit scaling / softcap).
• test_monolithic_loss_triton functional — 20/20.
• Distributed gloo world_size=2 monolithic-triton subtest for the tensor-parallel path — passing.
• Plain-CPU run (no interpreter) leaves the compiled path intact and skips the Triton head configs.

Caveat: this is interpreter-only so far; validation on compiled Triton on a GPU is still pending.

Follow-ups: distillation in the Triton kernel, and the GRPO/GSPO metrics rework.

🤖 Generated with Claude Code