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 has no Triton fused kernel yet — it is rejected at config time when use_triton=True, and falls 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, duplicate kinds) at validate time.
• None (default) — Triton on GPU only if every child is supported, else the compiled monolithic.
• False — always compiled.

Validation

Validated on an H100 GPU (FAST_LLM_SKIP_TRITON_AUTOTUNE=TRUE):

• test_lm_losses functional — the monolithic kernel against the PyTorch and standalone-Triton references (rms-relative ~1e-7): CE / z / CE+z, plus the GRPO and GSPO objectives through their kernel slots, across masking, logit/grad scaling, fp16, gradient accumulation, and the looped multi-block path.
• A metrics-without-gradient case — a zero-weight policy loss still logging its diagnostics (loss / new-log-probs / entropy) with no gradient written.
• test_lm_head fused_triton configs — single-pass and the GSPO seam, across masking / splits / bf16 / logit scaling / softcap.
• Distributed gloo world_size=2 monolithic subtests for the tensor-parallel path (CE+z, GRPO, GSPO).
• A plain-CPU run leaves the compiled path intact and skips the Triton head configs.

The Triton interpreter (TRITON_INTERPRET=1) can't run this kernel — the grid trips an InterpreterError on a constexpr-shaped dimension (reproduces on main, so it's pre-existing) — so validation of the monolithic Triton path is GPU-only.

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

🤖 Generated with Claude Code