[Bugfix][Training] Avoid entropy grad graph when entropy_coef is zero

[Henry-Jessie]

Summary

This PR reduces actor-update memory usage when entropy is logged but does not contribute to the loss.

Long-context actor updates can become memory-bound because entropy is evaluated over the full vocabulary. When entropy_coef == 0, VIME still computes entropy metrics, but the entropy tensor is only used for reporting and has no gradient contribution. The current implementation always builds the entropy autograd graph and clones the full logits tensor before entropy computation, which increases peak memory without changing the loss.

The change keeps the existing VIME log-probability and entropy implementation, but adds an explicit with_entropy_grad gate:

• get_log_probs_and_entropy() computes
  with_entropy_grad = with_entropy and getattr(args, "entropy_coef", 0.0) != 0.
• calculate_log_probs_and_entropy() computes entropy under torch.no_grad() when with_entropy_grad is false.
• In the no-grad entropy path, logits are reused directly instead of cloned.
• When entropy gradients are needed, the existing clone behavior is preserved.

The PR also includes two allgather-CP graph fixes from slime #2123:

• zero-filled response tensors created during allgather-CP redistribution now use requires_grad=ref_value.requires_grad instead of always requiring gradients;
• empty allgather-CP per-sample slices now use log_prob_full[:0] and entropy_full[:0] instead of newly allocated detached zero tensors.

This keeps entropy values available for metrics while avoiding unnecessary entropy-backward activations when the entropy loss coefficient is zero.

Upstream references

This PR follows recent slime changes, but only applies the pieces that fit VIME's separate log-probability and entropy paths.

• slime #2130 (open, not merged): proposed computing entropy under torch.no_grad() and skipping the defensive logits.clone() when entropy_coef == 0 on the separate _VocabParallelEntropy path. This is the closest conceptual match to VIME's current implementation.

• slime #2144 / slime #2152: slime later moved to a fused _VocabParallelLogProbEntropy function and added an internal with_entropy_grad gate. This PR borrows the same gate derivation, with_entropy_grad = with_entropy and getattr(args, "entropy_coef", 0.0) != 0, but does not port the fused function because VIME still uses separate compute_log_probs and _VocabParallelEntropy paths.

• slime #2123: fixed allgather-CP graph participation issues by using requires_grad=ref_value.requires_grad for redistributed zero fills and log_prob_full[:0] / entropy_full[:0] for empty slices. This PR includes those small fixes to keep allgather-CP tensors' autograd state consistent when entropy may be computed without gradients.

Not included: the fused _VocabParallelLogProbEntropy refactor. Porting it would replace Megatron's fused_vocab_parallel_cross_entropy with a custom vocab-parallel softmax, which is a larger behavioral change than this bug fix needs.

Correctness

The no-grad branch is safe because _VocabParallelEntropy.forward() does not modify the input logits in place. Its in-place exp_() and div_() operations are applied to tensors derived from vocab_parallel_logits - logits_max, not to the original logits tensor.

The clone is still kept when entropy gradients are required, preserving the original behavior for nonzero entropy coefficients.

Using getattr(args, "entropy_coef", 0.0) != 0 keeps the gradient path for any nonzero coefficient, including negative values.

The allgather-CP changes keep empty or missing response pieces aligned with the reference tensor's autograd state. This avoids creating artificial gradient paths for metric-only tensors, while still preserving graph connectivity for tensors that need to participate in distributed backward.

Validation

The profiling run followed the official GRPO training script structure in scripts/run-qwen3.5-27B.sh: Ray job submission to train.py, colocated actor/rollout, DAPO-style math data, GRPO with entropy_coef=0, dynamic batching, recompute, and Megatron context parallelism. The run was adapted to Qwen3.5-2B.

Run setup:

• Hardware: GPU, 4×A100-80GB
• CUDA: 12.9
• PyTorch: 2.11.0+cu129
• vLLM: 0.23.0
• Model: Qwen3.5-2B
• Training mode: GRPO, entropy_coef=0, TP=2, CP=2, rollout max response length 32k
• Main relevant changes from the official script: smaller Qwen3.5-2B checkpoint, 4 local GPUs, rollout-num-gpus-per-engine=4, global-batch-size=128, max-tokens-per-gpu=16384, and log-probs-chunk-size=1024

For matched ~32k-token actor microbatches, the internal loss-segment torch.cuda.max_memory_allocated peak was:

• patched: ~25.6 GiB
• unpatched: ~40.9 GiB

a reduction of ~15 GiB (~37%). This matches expectations: at ~32k tokens with vocab 248320 and TP=2 (V_local = 124160, fp32), one full [T, V_local] tensor is ~15 GiB, and the patch removes the entropy autograd graph and defensive logits clone that contribute that order of magnitude to the loss-segment peak.

AI Assistance

This PR was prepared with AI assistance from OpenAI Codex. I reviewed the changed code, compared it against the relevant slime changes, and validated the patch locally.

[read-the-docs-community]

Documentation build overview

📚 vime | 🛠️ Build #33509275 | 📁 Comparing 4714153 against latest (2702eef)

  🔍 Preview build  

42 files changed · ± 42 modified

± Modified

• genindex.html
• index.html
• search.html
• advanced/arch-support-beyond-megatron.html
• advanced/delta-weight-sync.html
• advanced/external-rollout-engines.html
• advanced/fault-tolerance.html
• advanced/low-precision.html
• advanced/megatron-config.html
• advanced/observability.html
• and 32 more...

[gemini-code-assist]

Code Review

This pull request optimizes memory and computation during entropy calculation by introducing a with_entropy_grad flag. This flag conditionally disables gradient tracking (using torch.no_grad()) and avoids cloning logits when the entropy term does not affect the loss. Additionally, it optimizes empty tensor creation in _extract_per_sample by slicing existing tensors instead of calling torch.zeros, and ensures requires_grad is correctly propagated in _allgather_cp_redistribute. There are no review comments, so I have no feedback to provide.