[nvdm59]

Summary

Adds TRL's top_entropy_quantile (high-entropy token masking, "Beyond the 80/20 Rule") to LigerFusedLinearGRPOLoss. This addresses the remaining open item of #1082 (item 6, asked about again in August): TRL's Liger path had to raise for top_entropy_quantile < 1.0 because the fused loss never materializes logits, so per-token entropy had nowhere to come from. It now comes from inside the chunked path.

Details

• New _selective_entropy_forward streams over sequence x vocab chunks like the selective-logprob forward, carrying a second online accumulator sum exp(z - m) * z next to the running (m, sum exp); H = log_z - sum_exp_z / sum_exp. No full-vocab logits are materialized, and the existing logprob fast path is untouched.
• LigerFusedLinearPPOBase.forward(top_entropy_quantile=1.0): when < 1.0, one extra no-grad pre-pass over the batch chunks collects entropies, the threshold is quantile(active entropies, 1 - q) over the whole batch (all-gathered across ranks when distributed, as in TRL's get_high_entropy_mask), and the (B, T) mask is chunked and passed to ppo_loss_fn(entropy_mask=...). _quantile_large_safe falls back to a sort beyond torch.quantile's 2^24-element cap.
• ppo_loss_fn multiplies the policy term by the mask in TRL's position: after the loss-type branch, before the vLLM IS correction and before the KL penalty, for every loss type (including vespo). Loss normalization and the metrics layout are unchanged.
• top_entropy_quantile is appended at the end of the Function signature, so positional callers and liger_fused_linear_grpo are unaffected; the default (1.0) path is bit-identical to before. Validated to [0, 1].
• LigerFusedLinearPPOBase.backward now returns one gradient slot per forward input (the four vespo_* slots were missing). feat(grpo): add off-policy sequence masking (OPSM) to fused linear GRPO loss #1557 touches the same tuple; whichever lands second needs a trivial rebase.

Cost: one additional no-grad lm_head forward per step when the feature is on; nothing when off.

Follow-up (out of scope): the Triton _selective_log_softmax_kernel could carry the same second accumulator in its vocab loop to bring this to triton_grpo_loss.

Testing Done

• _selective_entropy_forward vs torch.distributions.Categorical(...).entropy() (fp32/bf16, bias, temperature, shapes crossing both the 2048-row and 4096-vocab chunk boundaries).

• _quantile_large_safe vs torch.quantile, with the sort fallback forced.

• _compute_high_entropy_mask vs a single-process port of TRL's get_high_entropy_mask (incl. the empty-mask case).

• Fused loss vs a full-logits reference for q in {0.2, 0.5, 0.8} x grpo/bnpo/dapo/cispo/sapo/vespo/luspo x token/sequence IS x beta in {0, 0.04} x bias: loss, metrics, input/weight/bias gradients. The reference asserts the mask drops some but not all tokens, so the test cannot pass vacuously.

• Default 1.0 is bit-identical to omitting the argument; out-of-range values raise; chunk_size=1 and chunk_size=B agree (the quantile is batch-global, not per chunk; the test also shows a per-chunk quantile would differ).

• pytest test/chunked_loss/test_grpo_loss.py -k "not 1024-4096" on a CPU-only machine with torch.compile disabled: 696 passed, 338 skipped. The two test_triton_grpo_loss_matches_reference cases fail identically on main in this environment (Triton unavailable locally) and are not regressions.

• ruff check and ruff format --check clean on the changed files.

• Not verified locally: the distributed all-gather path (single process only), the compiled=True path (ran eager here), GPU/bf16, and the H=1024/V=4096 parametrizations. Please let GPU CI confirm.

• Hardware Type: CPU only (Apple M-series, no GPU)

• run make test to ensure correctness

• run make checkstyle to ensure code style

• run make test-convergence to ensure convergence

🤖 Generated with Claude Code

[nvdm59]

@vaibhavjindal when you have a moment, could you take a look at this one? It adds top_entropy_quantile ("Beyond the 80/20 Rule", TRL top_entropy_quantile) to LigerFusedLinearGRPOLoss, the other half of item 6 in #1082. Per-token entropy is computed with a second online accumulator in a no-grad streaming pre-pass over the lm_head, so the chunked memory profile is preserved. Also includes the same LigerFusedLinearPPOBase.backward gradient-count fix as #1557, so whichever of the two lands second will need a trivial merge.

Status: rebased on current main after #1556 landed, checkstyle is green, and test/chunked_loss/test_grpo_loss.py passes locally on the pure-PyTorch chunked path. No GPU here, so the compiled path, the distributed all-gather for the quantile threshold, and GPU CI still need to be exercised on your side. Happy to address any feedback.