feat: add top-k entropy approximation for memory-efficient GRPO training#555
Open
saurabhbikram wants to merge 2 commits intoOpenPipe:mainfrom
Open
feat: add top-k entropy approximation for memory-efficient GRPO training#555saurabhbikram wants to merge 2 commits intoOpenPipe:mainfrom
saurabhbikram wants to merge 2 commits intoOpenPipe:mainfrom
Conversation
When training models with large vocabularies (128k+ tokens), computing entropy over the full vocabulary is a major memory bottleneck. This adds a `top_k_entropy` config parameter (default 0 = disabled) that computes entropy over only the top-k logits, dramatically reducing memory usage. Also skips entropy computation entirely for reference model logprobs since entropy is unused in the KL divergence calculation. https://claude.ai/code/session_017Y9KNNQX2RyVWnqpj3A4hh
saurabhbikram
force-pushed
the
top-k-entropy-grpo
branch
from
524672e to
480300f
Compare
Preserve full-vocab entropy as default (top_k_entropy=0), only use top-k approximation when explicitly configured. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
saurabhbikram
force-pushed
the
top-k-entropy-grpo
branch
from
480300f to
14b8711
Compare
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
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.
Summary
When training models with large vocabularies (128k+ tokens, e.g. Qwen3), computing entropy over the full vocabulary during GRPO is a major memory bottleneck. This PR adds a
top_k_entropyconfig parameter:top_k_entropy=0(default): computes full-vocabulary entropy — identical to existing behavior, no regression.top_k_entropy=N(e.g.256): computes entropy over only the top-k logits instead of materializing a[B, chunk_size, V]log-probs tensor, dramatically reducing peak GPU memory.top_k_entropy=0since entropy is unused in KL divergence.Files changed
src/art/unsloth/train.py— addstop_k_entropyparameter threading throughcalculate_logprobs→_calculate_logprobs, adds top-k branch alongside preserved full-vocab entropy default