HF/FSDP backend: pass `logits_to_keep` to avoid materializing full logits tensor

[vsimkus]

Problem

The HF backend's model_wrapper.py computes logits for all sequence positions, then discards observation-position logits immediately after.

action_log_probs = log_probs[:, -num_actions - 1 : -1]

The same is done for entropies in PolicyWorkerBase.

For models with large vocabularies (e.g., Qwen 3.5/3.6 with vocab size 248K), this materializes a massive tensor at the lm_head. At bf16, 262K tokens × 248K vocab = ~123 GiB.

Often, there are significantly fewer action tokens than observation tokens, so not materializing the observation token logits can produce significant savings.

Proposed fix

HuggingFace transformers models already support a logits_to_keep parameter that slices hidden_states before the lm_head projection:

slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
logits = self.lm_head(hidden_states[:, slice_indices, :])

Since num_actions is already available in forward(), passing logits_to_keep=num_actions + 1 would avoid computing logits for observation positions entirely. No custom kernels required.

Standard path (no sample packing): one-line change — pass logits_to_keep=num_actions + 1 as an int.

With remove_microbatch_padding / sample packing: action positions are interleaved across packed sequences, so logits_to_keep needs the tensor form (indices of action positions).

Context

The Megatron backend already solves this via fused linear cross-entropy (#1841, #1765). This issue is about the HF/FSDP backend, which has no equivalent optimization.