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

[yc32768]

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.

[liuhao-labs]

Hi, I traced this again on current main (624b824). The model-output peak described here is still present.

Both forward-only log-prob and actor training already pass aligned full_loss_masks with labels=None, but the pinned Megatron GPTModel._postprocess() does not use that mask for the main LM output. It projects every hidden row to vocabulary logits, and the fp16 wrapper converts the complete [T, V/TP] output to fp32. Slime only chunks/selects rows later in get_log_probs_and_entropy(), after the peak allocation.

I checked the related open PRs: #1905 and #2076 select rows only after full fp32 logits return; #2010 removes an extra temperature-scaling allocation; #2304 adds peak-memory reporting. None avoids the projection/upcast peak reported here.

A further compatibility audit changed one implementation detail from my initial outline: a narrow runtime compatibility layer at Slime's Megatron patch loading point is preferable to editing the Docker Megatron source patch. Both normal GPTModel.forward() and combined-1F1B's PostProcessNode dynamically call the same _postprocess() before fp32 conversion, so the runtime layer can cover both paths and is exercised by the existing prebuilt CI images.

My proposed contract is actor-only and fail-closed: use full_loss_masks to keep trainable rows at the final hidden-state/output-layer boundary before vocab projection and fp32 conversion; compute vocab-parallel log-prob/entropy on compact rows; then scatter scalar results back to the existing response layout. Critic/value models, MTP training, config logging, custom model/loss paths, and unsupported signatures would keep the full-output fallback. Tests would cover CP/sequence-parallel alignment, top-p replay, empty masks, forward parity, backward-gradient parity, and both normal and combined schedules.

One semantic boundary needs maintainer confirmation: values outside the selected loss mask would be undefined/zero after scattering. Built-in masked training can preserve its result, but TIS or custom hooks that inspect off-mask log-prob/entropy may require the full-output fallback.

Before I open a PR, which ownership do you prefer?

A) The narrow, opt-in runtime compatibility layer in Slime plus the compact-output consumer described above.
B) An upstream Megatron selective-output primitive first, with Slime only wiring the mask and compact result.

I can keep the implementation limited to this contract and its regression tests.