Monolithic fused head-loss kernel (torch.compile) (#507)

[jlamypoirier]

Authored by Claude Opus 4.8 (Claude Code).

Summary

Implements the torch.compile portion of #507: a monolithic head loss that runs the vocabulary softmax once and shares it across a set of combinable losses, emitting each loss's scalar / metrics and the combined logits gradient in a single fused pass — instead of the per-loss loop where each loss re-softmaxes the full vocab and re-issues its own tensor-parallel all-reduces.

It is exposed as a loss type, not a head mode: configure a monolithic loss whose nested losses are the combinable set.

head:
  losses:
    combined:
      type: monolithic          # an ordinary head loss; the head loops over it unchanged
      losses:
        ce:      {type: label}
        z:       {type: z_loss, weight: 0.1}
        distill: {type: distillation, reference_model: teacher}
    dpo: {type: dpo, reference_model: ref}   # non-combinable → plain sibling, runs its own pass

Approach

• A loss, not a mode. MonolithicLoss / MonolithicLossConfig is a dynamic_type loss. The head is unchanged — it loops over head.losses and threads one gradient buffer, so a non-combinable loss (e.g. DPO) is just a sibling entry that runs its own forward_backward. There is no head-level toggle, no fallback dispatch, and no central spec object.
• Per-loss cores, one source of truth. Each combinable loss (label cross-entropy, z-loss, distillation, GRPO) owns a plain post-softmax function (e.g. z_loss_combinable) plus combinable_extract (eager kwargs → tensors) and combinable_core (the math). Its standalone fused_* kernel and the monolithic loop both call the same function, so the per-loss math lives once.
• One fused boundary via object dispatch. _monolithic_core computes the student softmax once, then loops the child losses calling child.combinable_core(...). The child set is fixed per config, so the loop unrolls and each combinable_core inlines inside one @torch.compile graph — confirmed to fuse with no graph breaks. Gradient contributions accumulate in fp32 and cast to the logits dtype once.
• Validation stays local. MonolithicLossConfig._validate requires its entries to be combinable and to agree on logits_scale_factor (one softmax = one scale). Duplicates are allowed and simply run twice over the shared student softmax.

Loss coverage

In the monolithic loss: label cross-entropy, z-loss, from-distribution cross-entropy / forward-KL / reverse-KL (distillation, with temperature), and the GRPO objective + new_logprobs + the GRPO metric family + entropy — all from the shared softmax, which removes the second softmax pass introduced in #494.

Out of scope / follow-ups. DPO (span-based + autograd + reference model) stays a standalone sibling loss. GSPO is a standalone loss for now (its per-segment eager seam doesn't fit the uniform core -> (loss, grad) shape); folding it into MonolithicLoss is a planned follow-up. The Triton port is also a follow-up.

Correctness

The per-loss kernels remain the equivalence oracle: the standalone fused_* kernels (now sharing the same combinable cores) validate against independent references, and the monolithic loss is validated end-to-end at the head level against the same references — across masking × cross-entropy-splits × dtypes, multi-token-prediction, tied embeddings, logit scaling, and softcap. GSPO's standalone three-phase form (compiled forward → eager segment seam → compiled backward) is unchanged and keeps its tensor-parallel coverage.

Validated on CPU (torch.compile) plus the distributed gloo world_size=2 chain for the tensor-parallel vocab-sharded path: head suite 129 passed, functional suite 260 passed, distributed 262 passed.

Performance

Opt-in, so this is upside rather than a default change. The fusion mechanism is a single shared softmax + one fused gradient pass across the enabled losses, so combined-loss configs avoid the redundant softmax/all-reduce passes of the per-loss loop. (Indicative single-GPU functional numbers from the earlier kernel structure showed ~1.5–1.7× on combined-loss configs with identical peak memory; a re-benchmark of the loss-type path is a follow-up. Tensor-parallel configs additionally save the redundant all-reduce set.)

Bundled fix

Per-loss distillation now honors its temperature config — the standalone path previously dropped it and defaulted to 1.0. Existing distillation runs with temperature != 1.0 will see corrected behavior.

Files

• loss/monolithic.py — MonolithicLoss + the compiled _monolithic_core entry loop.
• loss/config.py — MonolithicLossConfig + a combinable flag on each loss config.
• loss/z_loss.py, loss/entropy_loss.py, loss/policy_gradient.py — per-loss combinable_* methods + the shared post-softmax functions; standalone fused_* kernels routed through them. GSPO's segment seam / backward moved back here.
• functional/entropy_loss.py — the shared math cores (softmax_base, predicted_logits_from_labels, cross_entropy_from_labels_core, z_loss_core, the from-distribution cores).
• loss/loss.py — register_combinable_extras hook; _get_reference_model_logits resolves the head logits robustly for nested losses.
• language_model/head.py — reverted to main (no monolithic-specific code).
• tests/layers/test_lm_head.py, tests/layers/test_lm_losses.py — head-level parity for the monolithic loss type; standalone kernel + reference coverage retained.

🤖 Generated with Claude Code