feat(sample-support): report the policy entropy over the recorded support

[dyurk-lila]

Note on the diff: This PR is part of a routed-expert-replay / sampler-support series and builds on the PRs below. GitHub cannot show the intermediate branches here, so the diff is cumulative on top of main — the changes new to this PR sit on top of:

• refactor(inference): extract reusable RemoteInferenceGenerator #1911 — refactor extract RemoteInferenceGenerator
• refactor(generators): assemble routed-expert traces incrementally #1912 — refactor assemble routed-expert traces incrementally
• perf(r3): collate routes through a container-aware pool, and refuse replay under VPP #2078 — perf(r3) collate routes through a container-aware pool
• refactor(wire): generic packed-ndarray codec and an N-field response-body splice #2079 — refactor(wire) generic packed-ndarray codec
• perf(r3): route rollout experts to the trainer packed with cu_seqlens #2080 — perf(r3) packed routed-expert trainer transport
• perf(batch): ship large per-token side channels to the trainer zero-copy #2081 — perf(batch) zero-copy per-token side channels
• feat(sample-support): capture the sampler's bounded support and trace it across turns #2082 — feat(sample-support) capture and trace sampler support
• perf(sample-support): carry the captured support to the trainer as one packed field #2083 — perf(sample-support) packed trainer transport
• feat(sample-support): score policy logprobs over the recorded support end to end #2084 — feat(sample-support) Megatron scoring end to end
• feat(fsdp): score policy logprobs over the recorded support under Ulysses SP #2085 — feat(fsdp) support-conditioned scoring under Ulysses

Reviewing in PR order (lowest number first) shows each incremental change cleanly.

Problem

Support replay makes policy log-probabilities conditional on the recorded sampler support, but policy entropy still uses the full vocabulary. Besides reporting a different distribution, the full-vocabulary entropy path retains most of the computation and memory that bounded replay avoids—especially when the fused LM-head path otherwise projects only selected candidates.

This PR computes entropy from the same support scores used for replay on both Megatron and FSDP.

Entropy computation

For support logits l_i, shifted by their row maximum, the scorer uses:

H = log(Z) - sum_i p_i * (l_i - max(l))

The exponentials and denominator Z already exist for support-conditioned log-probabilities. Under tensor parallelism, the entropy numerator is added as a third row to the existing combined SUM reduction, so entropy requires no additional collective.

A singleton support has exactly zero entropy and zero entropy gradient. Trailing -1 members do not contribute.

Metric and loss modes

Entropy calculation and gradient tracking are separate controls:

• metric-only entropy detaches support values and probabilities;
• entropy-loss mode preserves the gradient through the support-conditioned distribution; and
• requesting entropy gradients without entropy calculation is rejected.

The fused LM-head scorer reuses the selected candidate projections from replay. It therefore supports differentiable entropy without materializing full-vocabulary logits.

Masking and layout

SampleSupportScores.valid_mask identifies positions backed by a recorded support row. Megatron and FSDP carry that mask back to canonical batch coordinates alongside the entropy values, including packed Megatron layouts and Ulysses gather/unpadding.

The final entropy reduction intersects this mask with the response loss mask. A synthetic EOS still receives its full-vocabulary log-probability fallback, but it is excluded from support-conditioned entropy because no rollout support exists for that token.

When sample-support replay is disabled, both backends retain their existing full-vocabulary entropy paths.

Testing

• Numerical tests compare entropy values with direct softmax references across support widths, padding patterns, temperatures, and tensor-parallel shards.
• Gradient tests cover logits and fused selected-projection weights in both metric-only and entropy-loss modes.
• Edge cases include singleton support, entirely missing support, synthetic EOS, and empty local vocabulary shards.
• Megatron tests cover packed/unpacked layouts, context-parallel scattering, and the fused LM-head path.
• FSDP tests cover padding removal, Ulysses slicing/gather, optimizer updates, and canonical entropy-mask reconstruction.

---

Note

High Risk
Changes how training logprobs and entropy (including entropy-loss gradients) are computed, plus inference capture and packed per-token batch transport. Incorrect support alignment would silently bias RL updates.

Overview
When enable_sample_support_replay is on, policy entropy now matches the support-conditioned logprobs instead of the full vocabulary. Metric-only entropy detaches the support distribution; entropy-loss keeps gradients through it. Tensor-parallel entropy piggybacks on the existing SUM reduction. Synthetic EOS still uses a full-vocab logprob fallback but is excluded from support entropy via valid_mask.

Megatron and FSDP both reduce entropy with that mask intersected with the response loss mask. The fused LM-head path can now take a differentiable entropy term because it reuses selected-candidate projections rather than materializing vocab logits. Full-vocab entropy is unchanged when replay is off.

The same diff also lands the rest of the sample-support stack: vLLM capture of bounded top-k IDs, packed HTTP/trainer transport (PackedTensor), R3 route packing, and support-conditioned scoring. R3 remains refused with step-wise / misaligned generators.

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

[gemini-code-assist]

Code Review

This pull request introduces sample support replay to renormalize policy logprobs over the sampler's recorded bounded support, adding a new PackedTensor utility to manage ragged token-aligned batch fields efficiently. It integrates this capability across both Megatron and FSDP training pipelines, optimizes CPU thread pool sizing using cgroup quotas, and implements zero-copy serialization for side-channel arrays. Feedback on the changes highlights a potential AttributeError in the vLLM server actor if logprobs are not returned during sample-support capture, suggesting a guard to handle None values gracefully.

[gemini-code-assist]

If resp.logprobs is None (which can happen in some edge cases or error states in vLLM), calling _sample_support_from_flat_logprobs will raise an AttributeError when trying to access logprobs.token_ids. Consider adding a guard to check if resp.logprobs is None and raise a descriptive error or handle it gracefully.

                    if resp.logprobs is None:
                        raise HTTPException(
                            status_code=500,
                            detail="vLLM failed to return logprobs for sample-support capture."
                        )
                    content, support_ids = _sample_support_from_flat_logprobs(
                        resp.logprobs,
                        sampling_params_dict["top_k"],
                    )