fix(grpo): clamp log-ratio and k3 KL for numerical stability#9
Open
dongseokmotif wants to merge 2 commits into
Open
fix(grpo): clamp log-ratio and k3 KL for numerical stability#9dongseokmotif wants to merge 2 commits into
dongseokmotif wants to merge 2 commits into
Conversation
…bility Unbounded ``exp`` over the policy/old log-ratio and over the k3 KL ``log_p - log_q`` can overflow fp32 to ``inf`` at masked / low- probability positions. Because the subsequent reduction is a plain ``per_token_loss * attention_mask``, ``inf * 0 == nan`` then contami- nates the entire batch, producing NaN losses on the very first step. Add three optional clamps (defaults ``None`` = no behavior change): - ``log_ratio_clamp_value`` clamps ``policy - old`` before ``exp``, protecting both the token- and sequence-level importance-weight paths (and the PG ``coef_1``). - ``kl_input_clamp_value`` clamps ``ref - policy`` inside ``k3_loss_fn`` before ``exp``. - ``kl_output_clamp_value`` clamps the resulting k3 KL value. This mirrors the guard used by NeMo-RL's ``calculate_kl`` (``input_clamp_value`` / ``output_clamp_value``) and the broader RL-framework convention. Callers that don't set the new arguments get the prior behavior unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Set the new clamp arguments to NeMo-RL's ``calculate_kl`` defaults (``log_ratio_clamp_value=20.0``, ``kl_input_clamp_value=20.0``, ``kl_output_clamp_value=10.0``) instead of ``None`` so the NaN guard is active out of the box. Existing callers automatically get the safety; users that want the prior unbounded behavior can pass ``None`` explicitly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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
main기준으로 GRPO 수치 안정화(clamp) 커밋만 담은 PR입니다.Numerical-stability clamping for the chunked GRPO loss, matching NeMo-RL behavior.
Details
exp(log_ratio_clamp_value).(ref - policy)log-ratio beforeexpinside k3, and clamp the resulting k3 KL output (kl_input_clamp_value,kl_output_clamp_value).nan_to_numguards on the relevant exponentials.Commits
fix(grpo): clamp log-ratio and k3 KL inputs/outputs for numerical stabilitydefault clamp values on (NeMo-RL parity)