[Bug] Training/log-prob forward materializes full-vocab fp32 logits over the entire packed sequence → OOM for large-vocab, long multi-turn RL

Bug Description

In the actor forward used for the policy loss and log-prob recompute, the model is called with labels=None, so GPTModel returns the full [1, T, vocab] logits and Float16Module upcasts them to fp32;
get_log_probs_and_entropy then computes over the whole [T, vocab] tensor and only slices out the response spans afterward. For RL where only loss_mask==1 (model-generated) tokens contribute to the
loss, this means logits are computed and upcast for every prompt and tool-output position too, then discarded. Peak memory therefore scales with the total packed sequence length × vocab, not with the
number of response tokens. --log-probs-chunk-size doesn't help the peak, because it only chunks the loss-side reduction — the full fp32 [T, vocab] tensor is already materialized in the model forward
before the loss runs.

This becomes fatal with a large vocabulary and long multi-turn/agentic episodes. Concretely, with vocab ≈ 248k (Qwen3.5) and context 32k–64k, a single microbatch's fp32 logits alone can be tens of GB
(e.g., a ~11.6k-token microbatch = 11.6k × 248320 × 4B ≈ 10.8 GB; near-context episodes are far larger), causing OOM in the training forward even with optimizer CPU offload. Note the tool-output/prompt
tokens (which dominate long agentic sequences) are exactly the non-trainable positions being wasted.

Steps to Reproduce

1. Megatron backend, large-vocab model (Qwen3.5-9B, vocab_size=248320), TP=4/DP=2, colocate, on H100 80 GB GPUs.
2. Long packed sequences — multi-turn/agentic rollouts where prompt + tool outputs dominate the sequence and only response tokens carry loss_mask=1; context length 32768–65536.
3. Run a normal GRPO step (policy loss forward, or the old-actor/ref log-prob recompute), labels=None.
4. Hit a microbatch whose packed length is large (e.g. ~11k+ tokens); observe the OOM in the training forward at the fp32 logit conversion. It reproduces even with --log-probs-chunk-size set and
   optimizer CPU offload enabled.

Expected Behavior

Peak logit memory scales with the number of loss/response tokens (or is chunked), so large-vocab + long-context RL trains without OOM — since only response positions affect the loss.

Actual Behavior

The full [1, T, vocab] fp32 logits for the entire packed sequence are materialized before any loss-side chunking, so peak ≈ T × vocab × 4 bytes regardless of how few tokens are trained. Example: a
single ~11.6k-token microbatch = 11.6k × 248320 × 4B ≈ 10.8 GB just for that tensor (episodes near the context limit are far larger), causing torch.OutOfMemoryError in the forward (float16_to_fp32)
even with optimizer offload. --log-probs-chunk-size doesn't help because it's downstream of the materialization.

Environment

• slime version: 0.3.0
• Python version: 3.12.7
• PyTorch version: 2.11.0+cu130
• CUDA/ROCm version: CUDA 13.0 (torch); toolkit 13.1; driver 580.82.07
• GPU type and count: 8× NVIDIA H100 (80 GB)
• OS: Alibaba Cloud Linux 3 (kernel 5.10.134-18.al8.x86_64)
• SGLang version (if relevant): 0.5.13
• Megatron-LM version (if relevant): megatron-core 0.16.0rc0

Logs

Additional Context

No response

Pre-submission Checklist

• I have read the CONTRIBUTING.md and understand the collaboration scope.
• I have read the documentation and my issue is not addressed there.
• I have searched for existing issues and this is not a duplicate.
• I have provided a minimal, reproducible example.

Trace the actor forward, GPTModel, Float16Module, and get_log_probs_and_entropy path used by a normal GRPO step with labels=None; first reproduce the OOM with a large packed sequence and inspect where full logits are materialized. Done means peak logit memory depends on loss/response positions or bounded chunks rather than all T positions, while policy-loss and old-actor/ref log-prob recomputation remain correct.