fix(args): reject --kl-coef with GRPO-family advantage estimators - #2368
Open
QJSQJS wants to merge 1 commit into
Open
fix(args): reject --kl-coef with GRPO-family advantage estimators#2368QJSQJS wants to merge 1 commit into
QJSQJS wants to merge 1 commit into
Conversation
--kl-coef was silently ignored by the grpo/gspo/cispo estimators: get_grpo_returns only uses the per-token KL for its shape, so the penalty never entered rewards or loss, while the reference-model forward cost and KL metrics were still paid. Fail fast in slime_validate_args and point users to --use-kl-loss instead.
QJSQJS
force-pushed
the
fix/kl-coef-ignored-by-grpo-estimators
branch
from
September 8, 2026 07:54
28aca97 to
b3950b5
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
1 participant
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.
Problem
--kl-coefis silently ignored by thegrpo/gspo/cispoadvantage estimators:get_grpo_returns(slime/utils/ppo_utils.py) only uses the per-token KL tensors for their shape, so the KL penalty never enters rewards or loss. Users who set--kl-coef 0.01with the default estimator (grpo) get zero regularization with no warning — while the reference-model forward cost is still paid and KL metrics are still logged, which makes the configuration look effective.Fix
slime_validate_args: reject--kl-coef != 0whenadvantage_estimatoris one ofgrpo/gspo/cispo, pointing to--use-kl-loss --kl-loss-coefinstead. Skipped when--custom-advantage-function-pathis set (a custom advantage function replaces the estimator branch and may consumekl_coefitself).--kl-coefhelp text with the estimators that actually apply it.Estimators that do consume
--kl-coef(unchanged, covered by tests):ppo,reinforce_plus_plus,reinforce_plus_plus_baseline.Tests
Added to
tests/test_megatron_argument_validation.py:kl_coefwith each ofgrpo/gspo/cispokl_coefwithppo/reinforce_plus_plus/reinforce_plus_plus_baselinekl_coefwith a custom advantage functionAll 27 tests in the file pass.