[dyurk-lila]

What does this PR do?

Compute selected-token log-probabilities and policy entropy from one fused LM-head projection for RL.

How it works

Metric-only entropy reuses the forward reductions without adding entropy-gradient work. Entropy-loss mode keeps the entropy derivative in the fused autograd path.

Torch and dense Triton return log-probabilities and entropy through the same packed TP/CP gather/scatter contract. #2375 adds block-sparse forward computation on top of this shared entropy path, and #2376 adds block-sparse backward computation.

Usage

Existing fused LM-head settings also cover entropy metrics and use_entropy_loss; the caller’s gradient requirement selects the entropy backward behavior.

Performance

Recorded measurements on one 8× NVIDIA H100 80 GB node compared the combined dense Triton path with fused selected-token log-probability scoring followed by the previous separate metric-entropy projection. Across 115 points, the combined path was a median 1.44× faster. The sweep covered four LM-head shapes, 16K–256K tokens, and representative TP/CP layouts with TP × CP <= 8.

For differentiable entropy at TP4/CP2, (H=4096,V=131072), and 16K tokens, the recorded speedup versus fused log-probability scoring plus a separate differentiable entropy projection was:

Active rows	Dense combined / previous two-pass speedup
100%	2.15×
95%	2.13×
73%	2.13×
64%	2.11×
23%	2.17×

The differentiable-entropy measurements used 180 CPU requested and limited. The small variation across activity levels is measurement noise: this PR's dense path does not inspect the activity pattern. #2375 measures the additional benefit from skipping inactive forward tiles against this shared entropy baseline.

These are isolated kernel measurements on one GPU generation and four model shapes, not end-to-end training throughput or universal model/GPU claims. Recorded benchmarks predate the current rebase and were not rerun for it.

Test coverage

Tests compare entropy values, selected-token log-probabilities, hidden gradients, and weight gradients with materialized-logit references, including packed TP2/CP2 and metric-only entropy.

Validation

Before the short-tail metrics follow-up, focused CPU suites passed in both mirrors: 415 tests across collation, DP sampling, worker batching, layout, active spans, and Megatron correctness (13 Megatron-only skips); 40 bin-packer tests; and 3 unpacked active-span guard tests with megatron-core installed. Ruff and Black passed. Seven targeted Triton value/gradient tests passed on one H100; the full GPU suite and benchmarks were not rerun.

After the metrics follow-up, 39 focused SFT CPU tests passed in each mirror, including synchronous and asynchronous short-tail logging; Ruff, Black, and secret detection passed on the changed files.

Review order

Part 6 of 9 in a linked series. The stack is based on main; each later PR targets the preceding layer. Review and merge from bottom to top.

Depends on #2373. Review only this layer.

1. perf(sft): use modified first-fit decreasing packing #2369 — use modified first-fit decreasing packing
2. perf(train): make sequence packing alignment-aware #2370 — make sequence packing alignment-aware
3. refactor(megatron): share packed segment layout derivation #2371 — share packed segment layout derivation
4. perf(sft): align packed bin counts to data parallelism #2372 — align packed bin counts to data parallelism
5. perf(megatron): autotune fused LM-head Triton kernels #2373 — autotune fused LM-head Triton kernels
6. feat(megatron): fuse LM-head entropy for RL #2374 — fuse LM-head entropy for RL
7. feat(megatron): add block-sparse fused LM-head forward #2375 — add block-sparse fused LM-head forward
8. perf(megatron): add block-sparse fused LM-head backward #2376 — add block-sparse fused LM-head backward
9. perf(megatron): fuse sample-support candidate projection with Triton #2521 — fuse sample-support candidate projection with Triton

---

Note

Medium Risk
Changes fused autograd and Megatron RL loss wiring for entropy bonuses; incorrect TP/CP or gradient math would affect training, but behavior is covered by new reference tests.

Overview
Extends fused LM-head log-prob paths (PyTorch chunked and Triton) so a single projection can also return per-token entropy, with optional differentiable backward when entropy_requires_grad is set for RL entropy bonuses.

from_parallel_hidden_to_logprobs and the packed-sequence variant gain return_entropy / entropy_requires_grad, share CP gather and scatter-back via new _unpack_prediction_values, and drop the old separate fused-entropy helpers. The Megatron worker now requests entropy from the fused op (metric-only vs loss-grad gated by use_entropy_loss) instead of erroring or running a second entropy pass.

Tests add materialized-logit parity for entropy values and gradients (torch + Triton) and expand Megatron CI for fused torch/triton with entropy loss.

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

[gemini-code-assist]

Code Review

This pull request integrates optional entropy computation and its gradient contribution directly into the fused LM-head projection path (supporting both Triton and PyTorch backends), allowing RL entropy bonuses to be computed without materializing the full logits tensor. It also unifies the codebase by removing redundant, separate fused entropy functions and introducing a helper for unpacking packed sequences. The review feedback identifies a performance bottleneck in the PyTorch backend, where performing all_reduce inside the chunk loop for entropy introduces blocking collective communication calls. The reviewer suggests deferring the all_reduce to a single call on the concatenated entropy tensor after the loop, and saving the forward-computed entropy in save_for_backward to avoid redundant recomputation and communication during the backward pass.

[gemini-code-assist]

Performing all_reduce inside the chunk loop introduces a blocking collective communication call for every single chunk, which significantly degrades performance. Since entropy is not used within the loop, we can defer the all_reduce to a single collective call over the concatenated entropy tensor after the loop, matching how log_probs are optimized.

Suggested change

	if compute_entropy:
	entropy = -(distributed_log_probs.exp() * distributed_log_probs).sum(dim=-1)
	torch.distributed.all_reduce(entropy, group=tp_group)
	all_entropy.append(entropy)
	if compute_entropy:
	entropy = -(distributed_log_probs.exp() * distributed_log_probs).sum(dim=-1)
	all_entropy.append(entropy)

[gemini-code-assist]

We can perform a single all_reduce on the concatenated entropy tensor here. Additionally, saving the computed entropy tensor in save_for_backward allows us to completely avoid recomputing and performing another all_reduce on chunk_entropy in the backward pass.

        if compute_entropy:
            entropy = torch.cat(all_entropy, dim=1)
            torch.distributed.all_reduce(
                entropy,
                op=torch.distributed.ReduceOp.SUM,
                group=tp_group,
            )
        else:
            entropy = None

        if not inference_only:
            ctx.save_for_backward(hidden, weight, target_mask, masked_target, entropy if entropy_requires_grad else None)
            ctx.chunk_size = chunk_size
            ctx.tp_group = tp_group
            ctx.entropy_requires_grad = entropy_requires_grad

        if not compute_entropy:
            return log_probs

        if not entropy_requires_grad:
            ctx.mark_non_differentiable(entropy)
        return log_probs, entropy

[gemini-code-assist]

Unpack the saved entropy tensor from ctx.saved_tensors to use it in the backward pass.

Suggested change

	hidden, weight, target_mask, masked_target = ctx.saved_tensors
	hidden, weight, target_mask, masked_target, *saved_entropy = ctx.saved_tensors
	entropy = saved_entropy[0] if ctx.entropy_requires_grad else None

[gemini-code-assist]

Instead of recomputing the entropy and performing a blocking all_reduce call per chunk in the backward pass, we can directly slice the saved entropy tensor that was computed during the forward pass.

Suggested change

	if ctx.entropy_requires_grad:
	softmax_output = distributed_log_probs.exp()
	chunk_entropy = -(softmax_output * distributed_log_probs).sum(dim=-1)
	torch.distributed.all_reduce(chunk_entropy, group=tp_group)
	else:
	softmax_output = distributed_log_probs.exp_()
	if ctx.entropy_requires_grad:
	softmax_output = distributed_log_probs.exp()
	chunk_entropy = entropy[:, chunk_start:chunk_end]
	else:
	softmax_output = distributed_log_probs.exp_()

[greptile-apps]

Confidence Score: 4/5

[High risk] Adds entropy computation to the fused LM-head projection path.

The PR should not merge until the distributed entropy parity test compares hidden gradients at the same TP scope.

Findings

1.  TP hidden gradients mismatch ▶

_{Reviews (1) · Last reviewed commit: "refactor(megatron): move prediction unpa..."}

[greptile-apps]

TP hidden gradients mismatch

In the TP=2 and TP=4 cases, gh_el contains the fused operator’s hidden gradient from this rank’s weight shard, while gh_er comes from the full-vocabulary projection. The new assertion compares them directly, so the distributed test fails even when the fused backward is correct. Reduce the hidden gradients across TP before comparing, or use a shard-local reference.

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, have a team admin enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 85cd762. Configure here.}

[cursor]

TP test compares unsynced hidden grads

Medium Severity

The new TP entropy check compares fused grad_hidden from a vocab shard against a full-vocabulary materialized reference. The fused backward only forms this rank's local d_logits @ W_local contribution and leaves the TP reduction to the caller, so those hidden gradients are not comparable. Weight grads are correctly sliced to the local shard, but the hidden-grad assertion will fail or force a contract-breaking all-reduce.

 

^{Reviewed by Cursor Bugbot for commit 85cd762. Configure here.}