Fuse distillation + z-loss in the triton monolithic path (#507)

[jlamypoirier]

Claude Opus 4.8 note: authored with Claude Code on the user's behalf.

Part of the monolithic head-loss work (#507), stacked on #549 ← #558.

What it does

Distillation was the one combinable loss with no slot in the triton monolithic kernel, so any monolithic loss containing it fell back to the compiled path — even though distillation already has a perfectly good standalone triton kernel. This adds a second fused triton kernel so distillation configs run on triton.

The label-family kernel (ce/z/grpo/gspo) and the distribution-family loss don't share a gradient shape — the label losses superpose a one-hot target (grad_j = prob_coeff·softmax_j − label_coeff·δ), while distillation's target is a full teacher distribution. Rather than force both into one kernel, MonolithicLoss now picks between two fused kernels by child kind:

• label-family {ce, z, grpo, gspo} → the existing monolithic kernel (unchanged);
• distribution-family {distillation} or {distillation, z_loss} → the existing from-distribution forward/backward kernels, extended to also emit z-loss over the same shared student softmax.

z-loss folds in as a pure student-probability coefficient — grad_j += 2·grad_z·log_sum_exp · student_prob_j, forward scalar log_sum_exp² — reusing values the distillation kernel already computes. It's direction-agnostic (cross-entropy / forward-KL / reverse-KL) and shares distillation's loss mask and divisor (both come from the same _get_loss_mask / label-count on the head). The term is gated on its output pointer, so it's dead-code-eliminated when no z-loss is present, leaving the standalone distillation path bit-for-bit unchanged.

Distillation mixed with a label or policy loss under use_triton is rejected at config validation (it can't share either kernel) and falls to the compiled path, which already fuses those combinations.

Changes

• functional/triton/entropy_loss.py — optional z slot in both from-distribution f/b kernels; triton_entropy_loss_forward_backward gains a z argument and returns (loss, z_loss, grad_logits).
• loss/monolithic.py — kernel selection by child kind; new distribution dispatch method; _TritonContext distillation slot.
• loss/entropy_loss.py, loss/config.py — distillation triton_kind, its _TritonContext pack/read, and relaxed monolithic validation.
• Tests — a functional parity test (all three KL directions, vs the sum of the standalone kernels), a tensor-parallel distributed subtest, and test_lm_head fused_triton distillation configs.

Validation

On an H100 (FAST_LLM_SKIP_TRITON_AUTOTUNE=TRUE):

• test_lm_losses.py — 586 passed / 21 skipped, including the new distillation-triton parity test across CE / forward-KL / reverse-KL, masking, scaling, fp16, and gradient accumulation, and the world_size=2 tensor-parallel subtests.
• test_lm_head.py — 189 passed, including the new fused_triton_distillation_loss / _and_z_loss / _loss_temperature configs (split and masked variants), exercising the MonolithicLoss dispatch end to end.

The 21 skips are the pre-existing reverse_kl-with-labels combinations. As with the rest of the monolithic triton path, the kernel can't run under TRITON_INTERPRET=1 (a constexpr-shaped grid trips an InterpreterError, reproducing on the base branch), so validation is GPU-only.

🤖 Generated with Claude Code