Auto-select the monolithic head-loss implementation (#507)

[jlamypoirier]

Claude Opus 4.8 note: authored with Claude Code on behalf of @jlamypoirier.

Stacked on #549 (jlp_monolithic_head_loss); this PR only adds the config-ergonomics layer on top of the monolithic kernels.

Motivation

Using the fused kernels currently requires restructuring the head config into a nested monolithic wrapper and hand-setting use_triton:

head:
  losses:
    fused:
      type: monolithic
      use_triton: true
      losses: {cross_entropy: {type: label}, z_loss: {type: z_loss}}

That nested type appears in no example/production config. This PR lets losses stay flat and picks the implementation automatically:

head:
  loss_implementation: auto     # default — nothing to write in the common case
  losses: {cross_entropy: {type: label}, z_loss: {type: z_loss}}

What changed

• loss_implementation head knob {auto, compiled, triton, per_loss}, default auto.
  • auto: fuse combinable losses; use triton when a group is triton-eligible and triton is available, else the compiled path.
  • compiled / triton: force a fused backend (triton errors at config time on an ineligible set).
  • per_loss: unfused per-loss behavior (the pre-change path).
• LanguageModelHeadConfig.get_effective_losses() applies the default-CE fallback and groups combinable losses by effective logits scale (one softmax serves one scale) into an internally-synthesized MonolithicLossConfig. Non-combinable losses (e.g. DPO) stay standalone. Called from _validate, so --validate catches an ineligible triton set early.
• Removed the user-facing type: monolithic from the dynamic-type registry; the class stays as the internal grouping vehicle. The head build loop, has_main_loss/registration logic, and metric names are unchanged.

Behavior note

auto is the default, so it changes the realized kernel for existing configs (including plain single-CE), moving them onto the fused path. Equivalence to per_loss is covered by the head/loss parity tests (rms-close; the fused path accumulates in fp32 and casts once, so it is at least as accurate). per_loss remains the escape hatch.

Tests

• test_get_effective_losses (new): grouping by scale, singleton wrapping, DPO standalone, backend mapping, and config-time rejection of an ineligible triton set.
• tests/layers/test_lm_head.py builds losses flat and drives the backend through loss_implementation (drops the manual wrapping); adds auto distillation variants (interpreter-safe distribution kernel).
• tests/layers/test_lm_losses.py unchanged (it exercises the kernels directly, not via the registry).

CPU validation green: test_lm_head.py 144 passed / 54 skipped (skips = triton, unavailable on CPU), test_lm_losses.py 586 passed / 21 skipped, and fast-llm train gpt --config examples/mistral.yaml --validate accepts the new default. Not yet run on GPU: the triton label-family paths, auto→triton, and the full tests/models suite (the default-kernel change).

🤖 Generated with Claude Code