fix: resolve OOM in long-sequence training via conditional entropy gradient tracking #1524
+29
−10
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Collaborator
|
thanks, we are refactoring this part for better menmory use. |
Contributor
Author
|
@lilei199908 could you review this PR and suggest any edits? I believe this should reduce memory usage. |
Collaborator
LGTM, we will merged it soon! thanks |
Contributor
Author
|
@zhuzilin Can you check and merge this pr? |
Contributor
Author
|
@yitianlian Can you check and merge this pr? |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR addresses the CUDA Out of Memory (OOM) issues #1523 encountered during training with long sequences (e.g., >30k tokens) by implementing conditional gradient tracking for the entropy term.
The Problem
In the previous implementation, the entropy calculation was always differentiable, regardless of the
entropy_coefvalue. To support the backward pass, the system was forced to store massive intermediate activation tensors (logits and softmax outputs) with shapes of(seq_len, vocab_size / TP). For long sequences and large vocabularies, these tensors consumed 5–6 GB of VRAM per sample on the last pipeline stage, leading to OOM even when the entropy contribution to the loss was zero.The Solution
I have decoupled entropy for monitoring (logging) from entropy for training (loss).
args.entropy_coef > 0.entropy_coefis 0.0, entropy for logging is computed within atorch.no_grad()context. This prevents PyTorch from allocating memory for backward tensors, effectively reclaiming several gigabytes of VRAM per sequence.Files Changed
slime/utils/ppo_utils.py: Addedrequires_entropy_gradflag andno_gradcontext tocalculate_log_probs_and_entropy.slime/backends/megatron_utils/loss.py: Passed gradient requirement flag based onentropy_coef.slime/backends/fsdp_utils/actor.py: Implemented conditional tracking and addedcontextlibimport.