perf(sample-support): reuse replay scores for entropy

[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 can't show the intermediate branches here, so the diff is cumulative on top of main — the changes new to this PR sit on top of:

• perf(r3): accelerate routed-expert transport with packed arrays #1909 — perf(r3) packed-array transport
• perf(r3): cast expert indices to int32 for pp local layers only #1910 — perf(r3) expand routes to local layers only
• refactor(inference): extract reusable RemoteInferenceGenerator #1911 — refactor extract RemoteInferenceGenerator
• refactor(generators): assemble routed-expert traces incrementally #1912 — refactor incremental routed-expert traces
• feat(megatron): add bounded sampler-support replay #1913 — feat(megatron) bounded sampler-support replay
• feat(fsdp): add bounded sampler-support replay #1914 — feat(fsdp) sampler-support replay
• perf(sample-support): add packed base64 transport #1915 — perf(sample-support) packed base64 transport

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

Problem

Sample-support replay renormalizes token log-probabilities over the recorded sampling support, but entropy still takes the full-vocabulary path. That loses most of the compute and memory benefit when support is narrow, and it prevents differentiable entropy from composing with the fused LM-head path.

Implementation

• return log-probabilities, optional entropy, and valid-row metadata together as SampleSupportScores;
• reuse the support logits and exponentials already computed for replay, adding the entropy statistic to the same tensor-parallel reduction;
• use support-conditioned entropy in both the FSDP/Hugging Face and Megatron trainers;
• keep the synthetic-EOS full-vocabulary log-probability fallback, while excluding rows without recorded support from the entropy reduction;
• make entropy computation and gradient tracking separate, explicit boolean controls;
• preserve the existing full-vocabulary entropy path when sample-support replay is disabled.

Entropy is computed for the policy conditioned on the recorded support set. A singleton support therefore has exactly zero entropy and zero entropy gradient. The change keeps the existing dense support representation and does not add a CSR path.

Performance

The replay entropy path scales with recorded support width rather than vocabulary size. It also shares the tensor-parallel collective with replay log-probability computation. With a fused LM head, it uses the existing selected hidden-state/weight projections and does not materialize full-vocabulary logits.

Testing

• value and gradient parity against direct support-conditioned references;
• detached metric entropy and differentiable entropy-loss modes;
• singleton, empty-support, and synthetic-EOS behavior;
• sequence-parallel metadata alignment;
• FSDP integration verifies that replay entropy bypasses full-vocabulary entropy;
• focused test suite: 23 passed;
• Black 24.10.0, Ruff 0.11.9, bytecode compilation, and explicit-call validation pass.

[gemini-code-assist]

Code Review

This pull request implements sample support replay, allowing policy logprobs to be renormalized over the sampler's recorded support. It introduces a packed wire format for dense sampler-support vocab IDs, adds a reusable RemoteGenerateClient for handling raw token generation requests, and integrates support-conditioned logprob and entropy computations into the Megatron and FSDP training pipelines. Additionally, Ulysses sequence parallelism utilities are updated to preserve trailing dimensions for token-aligned metadata. The review feedback suggests optimizing the _selected_hidden_projection function in sample_support_replay.py by chunking along the row dimension instead of flat pairs to avoid costly gather operations and memory allocations on the large hidden tensor.

[gemini-code-assist]

In _selected_hidden_projection, using hidden.index_select with repeated row_ids triggers a costly gather operation and memory allocation on the large hidden tensor (shape [num_tokens, hidden_dim]).

Since hidden is already organized by token rows, we can optimize this by chunking along the row dimension instead of flat pairs. Slicing hidden[start:end] is a $O(1)$ view operation in PyTorch that avoids copying the hidden states. We can then use broadcasting to compute the dot products efficiently.

def _selected_hidden_projection(
    hidden: torch.Tensor,
    token_ids: torch.Tensor,
    local_mask: torch.Tensor,
    lm_head_weight: torch.Tensor,
    temperature: float,
    chunk_size: int | None,
    invalid_value: float,
) -> torch.Tensor:
    """Project selected candidate pairs without materializing vocabulary logits."""
    num_rows, width = token_ids.shape
    output = torch.empty(token_ids.shape, dtype=torch.float32, device=hidden.device)
    # Bound the temporary projection by chunking along the token-row dimension.
    row_chunk_size = num_rows if chunk_size is None else max(1, chunk_size // width)
    for start in range(0, num_rows, row_chunk_size):
        end = min(start + row_chunk_size, num_rows)
        chunk_hidden = hidden[start:end].to(lm_head_weight.dtype)
        chunk_token_ids = token_ids[start:end]
        chunk_mask = local_mask[start:end]

        selected_weight = lm_head_weight.index_select(0, chunk_token_ids.reshape(-1)).reshape(
            end - start, width, -1
        )
        projected = (chunk_hidden.unsqueeze(1) * selected_weight).sum(dim=-1) / temperature
        output[start:end] = torch.where(
            chunk_mask,
            projected.to(torch.float32),
            invalid_value,
        )
    return output