[coreyhu]

Summary

With temperature != 1, the torch fused LM-head path (FusedLinearChunkedDistributedLogprob and the fused entropy metric) divided the BF16 output weight by the temperature before the projection. This PR divides each FP32 chunk of logits instead, and applies 1/T to the FP32 logit gradient before casting it, the same order as the unfused path (logits.float() / T). The Triton backend still receives weight / temperature; its kernel already takes a temperature, but switching it needs GPU validation.

Why

W / T is rounded to BF16, so fused log-probs, gradients and entropy drift from the unfused path. The docstring said the two were numerically identical.

Testing

• New CPU test tests/backends/skyrl_train/distributed/test_fused_lm_head_temperature.py (BF16, T=0.7): fused log-probs and hidden grads now match the unfused path exactly. All 3 cases fail on main (log-prob error 4.2e-2, entropy 2.0e-2).
• pytest tests/backends/skyrl_train/distributed/: 196 passed, 15 skipped (main: 193, 15). CPU suite (tests/train/ tests/backends/skyrl_train/, -m "not vllm"): 2225 passed, 41 skipped (main: 2222); the same 2 Megatron import failures occur on main. pre-commit clean.
• Not run: GPU or Triton tests. This box's CUDA 12.9 driver can't run SkyRL's CUDA 13 stack.

This conflicts textually with #2374, which edits the same op.

AI assistance: Claude Code ported this from our patch set, wrote the tests, and drafted this description.

---

Note

Medium Risk
Changes numerics and autograd for fused log-prob/entropy used in Megatron RL training when temperature ≠ 1; Triton path is untouched.

Overview
Aligns torch fused LM-head temperature scaling with the unfused path by scaling FP32 logits after projection instead of dividing BF16 weights before the matmul.

FusedLinearChunkedDistributedLogprob now takes a temperature argument: forward divides each chunked FP32 logits tensor by T; backward applies 1/T to the FP32 logit gradient before projecting to hidden/weight grads. Callers (from_parallel_hidden_to_logprobs, packed variant, _fused_vocab_parallel_entropy_from_hidden) no longer mutate lm_head_weight; they pass temperature through _fused_lm_head_logprob_apply. The Triton backend is unchanged and still receives weight / temperature.

Adds CPU tests (test_fused_lm_head_temperature.py) that fused log-probs, gradients, and entropy match the reference matmul → float32 → /T behavior at T=0.7 with BF16 weights (chunked and single-chunk).

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

[gemini-code-assist]

Code Review

This pull request updates the temperature scaling logic in the fused LM-head logprob calculation. Instead of scaling the weights before projection, the torch backend now scales the FP32 logits directly after projection, matching the order of operations in the unfused path. The backward pass is updated accordingly to apply the temperature factor to the logit gradients. Additionally, a comprehensive CPU test suite has been added to verify that the fused log-probs, gradients, and entropy metrics match the unfused reference implementation. I have no feedback to provide as there are no review comments.

[greptile-apps]

Confidence Score: 5/5

[Medium risk] Adds temperature scaling to the fused language model head computation.

The PR appears safe to merge; no actionable issues were found.

What we checked:

• Temperature reaches both wrappers: Both wrappers pass temperature to the shared helper, which passes it to the torch operation.
• Entropy loss stays blocked: The caller raises an error when the fused path is selected with use_entropy_loss=True.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    H[Hidden states and weights] --> P[Project each chunk]
    P --> F[Cast logits to FP32]
    F --> T[Divide logits by temperature]
    T --> L[Compute log-probs or entropy]
    L --> G[Compute FP32 logit gradients]
    G --> S[Divide gradients by temperature]
    S --> C[Cast to weight dtype]
    C --> B[Project hidden and weight gradients]

Loading

_{Reviews (1) · Last reviewed commit: "Apply temperature to FP32 logits in the ..." · Reviewed by Greptile}