[rl] Add trainer entropy metric

[acisseJZhong]

Summary

Add a trainer-policy entropy metric to watch for entropy -> reward collapse during RL (motivated by checking whether MoE / Qwen3-30B-A3B runs collapse), and surface the existing train/inference KL signal.

Why

Entropy collapse is the leading indicator of reward collapse in GRPO/DAPO: the policy sharpens to near-deterministic -> within-group reward variance -> 0 -> zero advantage -> no gradient. Entropy + KL on the dashboard catch this early.

What

• compute_logprobs(..., return_entropy=True) returns per-token entropy alongside logprobs from a single vocab gather + single fp32 upcast (no extra copy); default call is unchanged (returns logprobs).
• DAPO logs trainer/entropy/mean (mean H(p) over response tokens); documents bit_wise/logprob_diff/mean as the k1 estimate of -KL(q||p) (-1x verl's ppo_kl).
• Rename train/* -> trainer/* (grad_norm, lr, policy_version); add entropy to the console keys.

Test plan

• Ran locally (rl_grpo_qwen3_1_7b search_r1 + alphabet_sort); trainer/entropy/mean logs to W&B with the expected sharpen-then-plateau.

[tianyu-l]

This way you'd have two .float() conversion, one in logprobs, one for entropy. Have you thought about that?

[acisseJZhong]

good point, we shouldn't keep two copies of float32 logits, we can share one.

[tianyu-l]

Tow much value is it to keep this standalone? There are two other options so that we don't have to introduce many subfunctions:

• remove this, and only use the combined function
• having a flag e.g. return_entropy to if-else

curious what's the tradeoff in your mind?

[acisseJZhong]

compute_logprobs is only used in some test scripts that we have under experimental/rl. I think combining these two functions(second option) looks better.

[tianyu-l]

The other option is to return logprobs, None when return_entropy=False. Not sure which is better.

Update: OK claude prefers o/w, so no need to change

[tianyu-l]

@claude review this PR, including my comments

[claude]

Claude finished @tianyu-l's task in 2m 45s —— View job

---

PR Review: [rl] Add trainer entropy metric

Overall this is a clean, well-scoped change. The entropy math is correct, the single-upcast refactor preserves logprob numerics bit-for-bit, and the train/* -> trainer/* rename is consistent across the trainer, processor, and test comment. Below are responses to @tianyu-l's comments plus a couple of findings.

Responses to the inline comments

1. Two .float() conversions (loss.py) — resolved correctly.
The refactor hoists logits = logits.float() (loss.py:422) above both the cross_entropy and the entropy computation, so there is now a single upcast shared by both paths. Since .float() on an already-fp32 tensor is a no-op (returns the same tensor, no copy), there's no cost for callers that were already fp32. Importantly, the logprob path is bit-identical to before: F.cross_entropy(logits.float().reshape(...)) vs. logits = logits.float(); F.cross_entropy(logits.reshape(...)) are the same op on the same values. 👍

2. Standalone helper vs. combined function — the return_entropy flag is the right call.
Agreed with the resolution to go with option 2 (single function + flag). It keeps the DTensor / spmd_types vocab-gather logic (loss.py:394-419) in one place — duplicating that into a separate entropy function would be the exact kind of divergence the distributed rules warn about.

3. return logprobs, None vs. the union return type (line 377).
I'd keep the current torch.Tensor | tuple[...] union rather than always returning (logprobs, None). Three existing callsites use the single-tensor return directly:

• torchtitan/experiments/rl/batch_invariance.py:124
• torchtitan/experiments/rl/tests/test_bitwise_parity.py:379,413

Switching to an always-tuple return would silently break those (logprobs would become a 2-tuple), forcing a [0] unpack at every callsite for no real benefit. The union return keeps the common path unchanged and the tuple is opt-in. If the concern is purely type-checker ergonomics, that's a weaker reason than backward-compatibility here.

Findings

A. The unit test claimed in the test plan isn't in the diff.
The PR description says: "Unit: entropy matches -sum(p log p) / log(V) for uniform; logprobs bit-identical with/without the flag." But the only test-file change is a comment rename in test_metrics.py:463. There's no test asserting the entropy value or the bit-identical-logprobs invariant. Per CLAUDE.md ("New features need CPU unit tests at minimum"), this is worth adding — it's cheap and both properties are easy to check on CPU (uniform logits -> H = log(V); compute_logprobs(x, y) equals compute_logprobs(x, y, return_entropy=True)[0]). Fix this →

B. (minor, optional) Extra full-size allocation in the entropy path.
torch.softmax(logits, dim=-1) * logits (loss.py:433-435) materializes a second [B, L, V] fp32 tensor on top of the upcast logits. It's under no_grad so it frees immediately, but for large vocab / long sequences this transiently doubles the peak of the largest tensor. If that ever matters you can fuse via log_softmax to avoid the separate softmax*logits temp, e.g.:

with torch.no_grad():
    logp = torch.log_softmax(logits, dim=-1)
    entropy = -(logp.exp() * logp).sum(dim=-1)

Not blocking — the current form is correct and readable.

Things I confirmed are correct

• Entropy formula (loss.py:433): logsumexp(z) - sum(softmax(z) * z) = -sum p_i (z_i - logsumexp) = -sum p_i log p_i. Correct.
• Masking (dapo.py:138): entropy is aggregated over response_mask (the original loss_mask, not the finite-filtered one) and divided by loss_denominator, consistent with the other SUM-reduction metrics. Using response_mask rather than the nan-filtered loss_mask is right here — trainer-policy entropy is always finite.
• no_grad on entropy is correct: it's a metric only and must not build an autograd graph over the [B, L, V] softmax.
• Rename consistency: trainer/lr error string (trainer.py:441), metric keys (trainer.py:469-471), console keys (processor.py:66-68), and the test comment (test_metrics.py:463) all line up.

Nothing blocking beyond adding the unit test (A). Nice change.
• branch: add_metric