Don't compute the logprobs and their gradients for tokens no loss com…

[MarioSieg]

The loss math for itself it uncganed, default loss fn still runs elementwise of the full packed sequence and mult with a mask:

pg_loss = keep_mask * advantages * importance
kl_loss = loss_mask*log_importanance^2
loss = sum(per_token_loss)
^^^^^ those are elementwise ops on a vector with len seq with no vocab dim, skipping them doesn't seem to save much
what this PR skips is the matmul
when we do
hidden_i = XV -> logits -> logprob_i -> loss
the matmul XV is of dims [N, 4096] @ [4096, 151936] which is big
so we don't compute the logprobs and their grads for tokens no loss component reads since the loss can't run a token without its logprob - it's basically same as <<don't compute the loss on masked tokens>> but we save where the FLOPS actually are spent which is the matmul not the elementwise ops

bench results on qwen 30b:
┌────────────────────┬───────┬───────┬───────┐
│ │ on │ off │ Δ │
├────────────────────┼───────┼───────┼───────┤
│ throughput (tok/s) │ 5,462 │ 5,178 │ +5.5% │
├────────────────────┼───────┼───────┼───────┤
│ MFU (%) │ 0.7 │ 0.7 │ +5.5% │
├────────────────────┼───────┼───────┼───────┤
│ peak memory (GiB) │ 58.5 │ 58.5 │ +0.0% │
└────────────────────┴───────┴───────┴───────┘
Step-1 loss/grad_norm match within noise (verified against an on-vs-on and off-vs-off control: the on-vs-off gap sits inside the natural run-to-run spread, not outside it — EP all-to-all reduction makes grad_norm inherently noisier
than on dense models, ~5e-3 vs ~3e-5 relative).

---

Note

Medium Risk
Changes default training behavior and LM-head numerics path (mask compaction + zero fill on skipped tokens); correctness is covered by unit tests but distributed RL/SFT parity should be watched on mixed-mask batches.

Overview
Adds skip masked LM head tokens (default on) so the fused LM head runs the expensive vocab matmul only on positions any loss component actually uses—RL loss_mask, CE weights, and ref-KL weights—after the usual next-token shift. The backbone still processes the full sequence; skipped positions get zero logprobs/entropy so importance-ratio math stays stable.

Implements this via a keep_mask on forward and fused/Gemma heads: compact → chunked logprob/entropy → expand back. Vanilla LM head ignores the mask; RL disables skipping when the fused head is off or enable_token_export needs every token. Default fused_lm_head_token_chunk_size is raised to 8192; docs describe the new knob and perf/lm_head_token_fraction logging in RL and SFT.

^{Reviewed by Cursor Bugbot for commit d1b2e56. Bugbot is set up for automated code reviews on this repo. Configure here.}

[garrett361]

Very good. Mostly small comments, which you can address or not per your judgement. But before merging, please:

• Fix CI
• Remove bench.sh