From b418e3ab99c6f38cdf42d6e7c44cfbc07b185500 Mon Sep 17 00:00:00 2001 From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> Date: Wed, 26 Aug 2026 14:49:01 -0400 Subject: [PATCH 1/4] fix(algorithms): raise instead of asserting on the skipped-reference-KL pairing Sync PPO and sync GRPO both guard skip_reference_policy_logprobs_calculation against a non-zero reference_policy_kl_penalty with a bare, message-less assert master_config.loss_fn.reference_policy_kl_penalty == 0 python -O removes it entirely, so the run proceeds to train against a KL term whose reference logprobs were never computed. Even with asserts enabled, a message-less one says nothing about which two settings conflict. async_ppo_train already does this as an if + ValueError. This is the same defect class yuki-97 raised on #3262 ("assert backend == 'vllm' gets stripped under python -O"), which was converted there but left in these two siblings. Both now use the async form, verbatim, so the three read alike. Tests run under python -O as well as normally, which is what separates this from a cosmetic change: with the assert restored, both modes fail. Signed-off-by: Tianyi Zhang Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> --- nemo_rl/algorithms/grpo_sync.py | 6 +- nemo_rl/algorithms/ppo.py | 6 +- .../algorithms/test_reference_kl_guard.py | 107 ++++++++++++++++++ 3 files changed, 117 insertions(+), 2 deletions(-) create mode 100644 tests/unit/algorithms/test_reference_kl_guard.py diff --git a/nemo_rl/algorithms/grpo_sync.py b/nemo_rl/algorithms/grpo_sync.py index e3f1fe46a81..f1b8e7b2155 100644 --- a/nemo_rl/algorithms/grpo_sync.py +++ b/nemo_rl/algorithms/grpo_sync.py @@ -418,7 +418,11 @@ def grpo_train_sync( assert policy_generation is not None if master_config.grpo.skip_reference_policy_logprobs_calculation: - assert master_config.loss_fn.reference_policy_kl_penalty == 0 + if master_config.loss_fn.reference_policy_kl_penalty != 0: + raise ValueError( + "Skipping reference logprobs requires " + "loss_fn.reference_policy_kl_penalty=0" + ) print( "Reference policy logprob calculation will be skipped since `grpo.skip_reference_policy_logprobs_calculation` is set to True and `loss_fn.reference_policy_kl_penalty` is 0." ) diff --git a/nemo_rl/algorithms/ppo.py b/nemo_rl/algorithms/ppo.py index 4db2867a7ed..3f54a5d6482 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -1226,7 +1226,11 @@ def ppo_train( assert policy_generation is not None # for mypy type check if master_config.ppo.skip_reference_policy_logprobs_calculation: - assert master_config.loss_fn.reference_policy_kl_penalty == 0 + if master_config.loss_fn.reference_policy_kl_penalty != 0: + raise ValueError( + "Skipping reference logprobs requires " + "loss_fn.reference_policy_kl_penalty=0" + ) print( "Reference policy logprob calculation will be skipped since `ppo.skip_reference_policy_logprobs_calculation` is set to True and `loss_fn.reference_policy_kl_penalty` is 0." ) diff --git a/tests/unit/algorithms/test_reference_kl_guard.py b/tests/unit/algorithms/test_reference_kl_guard.py new file mode 100644 index 00000000000..2045940065b --- /dev/null +++ b/tests/unit/algorithms/test_reference_kl_guard.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Skipping reference logprobs with a non-zero KL penalty must fail loudly. + +The two sync paths guarded this with a bare, message-less ``assert``. Under +``python -O`` an assert is removed entirely, so the run would proceed to train +against a KL term whose reference logprobs were never computed -- the same +defect class raised on PR #3262, where the async twin was converted to a +``ValueError`` and these two were missed. + +CPU-only: the guard sits above everything in these functions except a Timer +and a MemoryTracker, so mock arguments reach it. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from nemo_rl.algorithms.grpo_sync import grpo_train_sync +from nemo_rl.algorithms.ppo import ppo_train + + +def _master_config(block: str, *, kl_penalty: float): + return SimpleNamespace( + **{ + block: SimpleNamespace(skip_reference_policy_logprobs_calculation=True), + "loss_fn": SimpleNamespace(reference_policy_kl_penalty=kl_penalty), + "checkpointing": {"checkpoint_must_save_by": None}, + } + ) + + +def _call_ppo(master_config): + ppo_train( + policy=MagicMock(), + policy_generation=MagicMock(), + value_model=MagicMock(), + dataloader=MagicMock(), + val_dataloader=None, + tokenizer=MagicMock(), + loss_fn=MagicMock(), + value_loss_fn=MagicMock(), + task_to_env={}, + val_task_to_env=None, + logger=MagicMock(), + checkpointer=MagicMock(), + ppo_save_state=MagicMock(), + master_config=master_config, + ) + + +def _call_grpo(master_config): + grpo_train_sync( + policy=MagicMock(), + policy_generation=MagicMock(), + wrapped_dataloader=MagicMock(), + val_dataloader=None, + tokenizer=MagicMock(), + loss_fn=MagicMock(), + task_to_env={}, + val_task_to_env=None, + logger=MagicMock(), + checkpointer=MagicMock(), + grpo_save_state=MagicMock(), + master_config=master_config, + ) + + +@pytest.mark.parametrize( + ("call", "block"), + [(_call_ppo, "ppo"), (_call_grpo, "grpo")], + ids=["sync-ppo", "sync-grpo"], +) +def test_a_nonzero_kl_penalty_with_skipped_reference_logprobs_raises(call, block): + """A ValueError, not an assert: `python -O` strips asserts, and this one + was message-less on top of that.""" + with pytest.raises(ValueError, match="reference_policy_kl_penalty=0"): + call(_master_config(block, kl_penalty=0.1)) + + +@pytest.mark.parametrize( + ("call", "block"), + [(_call_ppo, "ppo"), (_call_grpo, "grpo")], + ids=["sync-ppo", "sync-grpo"], +) +def test_the_supported_pairing_gets_past_the_guard(call, block): + """kl_penalty=0 is the combination the guard exists to allow. It must not + raise ValueError here -- these mocks fail later, which is fine and is what + keeps the test from passing vacuously if the guard were made + unconditional.""" + with pytest.raises(Exception) as excinfo: + call(_master_config(block, kl_penalty=0.0)) + assert "reference_policy_kl_penalty=0" not in str(excinfo.value) From 4bf8b1da9cc310c2552d2c036b458a92830ac07f Mon Sep 17 00:00:00 2001 From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:14:39 -0400 Subject: [PATCH 2/4] fix(grpo): convert the guard that actually fires, not just the shadowed one grpo_train_sync's assert is unreachable in practice: grpo.setup checks the same pairing before the train loop starts, so on GRPO that is the one a user hits. It was an assert too, so under python -O both were stripped and nothing was left -- converting only the train-loop copy would have fixed the copy nobody reaches. PPO is the other way round: there is no setup-time equivalent, so the one in ppo_train is the only guard and converting it was already right. Keeps the grpo_sync conversion as well. It costs nothing and the three sites now read alike, which was the point. Test asserts on grpo.setup's source rather than calling it, since calling it needs a cluster. What it pins is that the reachable guard is not an assert. Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> --- nemo_rl/algorithms/grpo.py | 9 ++++---- .../algorithms/test_reference_kl_guard.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index ed97496ce5d..9ace9baf215 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -766,10 +766,11 @@ def init_train_dataloader(dataset, suffix: str = ""): # Validate skip_reference_policy_logprobs_calculation if grpo_config.skip_reference_policy_logprobs_calculation: - assert loss_config.reference_policy_kl_penalty == 0, ( - "grpo.skip_reference_policy_logprobs_calculation=True requires " - "loss_fn.reference_policy_kl_penalty == 0" - ) + if loss_config.reference_policy_kl_penalty != 0: + raise ValueError( + "Skipping reference logprobs requires " + "loss_fn.reference_policy_kl_penalty=0" + ) print( "Reference policy logprob calculation will be skipped since `grpo.skip_reference_policy_logprobs_calculation` is set to True and `loss_fn.reference_policy_kl_penalty` is 0." ) diff --git a/tests/unit/algorithms/test_reference_kl_guard.py b/tests/unit/algorithms/test_reference_kl_guard.py index 2045940065b..2bd17340f88 100644 --- a/tests/unit/algorithms/test_reference_kl_guard.py +++ b/tests/unit/algorithms/test_reference_kl_guard.py @@ -80,6 +80,29 @@ def _call_grpo(master_config): ) +def test_the_grpo_guard_that_actually_fires_is_in_setup(): + """`grpo_train_sync`'s guard is shadowed and never reached. + + `grpo.setup` checks the same pairing before the train loop starts, so on + GRPO the setup one is what a user hits. It was an assert too, so under + `python -O` both were stripped and nothing was left. Converting only the + train-loop copy would have fixed the unreachable one. + + Asserted on the source rather than by calling `setup`, which would need a + cluster: the point is that the reachable guard is not an assert. + """ + import inspect + + from nemo_rl.algorithms import grpo + + src = inspect.getsource(grpo.setup) + marker = "skip_reference_policy_logprobs_calculation:" + assert marker in src + guard = src[src.index(marker) : src.index(marker) + 400] + assert "raise ValueError" in guard, "the reachable GRPO guard must not be an assert" + assert "reference_policy_kl_penalty=0" in guard + + @pytest.mark.parametrize( ("call", "block"), [(_call_ppo, "ppo"), (_call_grpo, "grpo")], From 6e1cc83c892031387bcd8bba5b0d91ec45784234 Mon Sep 17 00:00:00 2001 From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> Date: Sun, 30 Aug 2026 15:48:08 -0400 Subject: [PATCH 3/4] test(grpo): exercise reference KL guard through setup Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> --- nemo_rl/algorithms/grpo.py | 15 +++++--- .../algorithms/test_reference_kl_guard.py | 35 ++++++++----------- 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 44190645f84..c101becc144 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -554,6 +554,16 @@ def setup( cluster_config = master_config.cluster checkpointing_config = master_config.checkpointing + # Validate this before logger, checkpoint, dataloader, or cluster setup can + # produce side effects for a configuration that cannot train correctly. + if ( + grpo_config.skip_reference_policy_logprobs_calculation + and loss_config.reference_policy_kl_penalty != 0 + ): + raise ValueError( + "Skipping reference logprobs requires loss_fn.reference_policy_kl_penalty=0" + ) + checkpointing_pretrained = checkpointing_config.get("pretrained_checkpoint") if checkpointing_pretrained is not None: policy_config["pretrained_checkpoint"] = checkpointing_pretrained @@ -787,11 +797,6 @@ def init_train_dataloader(dataset, suffix: str = ""): # Validate skip_reference_policy_logprobs_calculation if grpo_config.skip_reference_policy_logprobs_calculation: - if loss_config.reference_policy_kl_penalty != 0: - raise ValueError( - "Skipping reference logprobs requires " - "loss_fn.reference_policy_kl_penalty=0" - ) print( "Reference policy logprob calculation will be skipped since `grpo.skip_reference_policy_logprobs_calculation` is set to True and `loss_fn.reference_policy_kl_penalty` is 0." ) diff --git a/tests/unit/algorithms/test_reference_kl_guard.py b/tests/unit/algorithms/test_reference_kl_guard.py index 2bd17340f88..45e2715a6a8 100644 --- a/tests/unit/algorithms/test_reference_kl_guard.py +++ b/tests/unit/algorithms/test_reference_kl_guard.py @@ -30,6 +30,7 @@ import pytest +from nemo_rl.algorithms.grpo import setup as grpo_setup from nemo_rl.algorithms.grpo_sync import grpo_train_sync from nemo_rl.algorithms.ppo import ppo_train @@ -80,27 +81,21 @@ def _call_grpo(master_config): ) -def test_the_grpo_guard_that_actually_fires_is_in_setup(): - """`grpo_train_sync`'s guard is shadowed and never reached. - - `grpo.setup` checks the same pairing before the train loop starts, so on - GRPO the setup one is what a user hits. It was an assert too, so under - `python -O` both were stripped and nothing was left. Converting only the - train-loop copy would have fixed the unreachable one. - - Asserted on the source rather than by calling `setup`, which would need a - cluster: the point is that the reachable guard is not an assert. - """ - import inspect - - from nemo_rl.algorithms import grpo +def test_grpo_setup_rejects_skipped_reference_logprobs_with_nonzero_kl(): + """The user-facing setup entry point rejects the invalid pairing early.""" + master_config = SimpleNamespace( + policy={"generation": {}}, + loss_fn=SimpleNamespace(reference_policy_kl_penalty=0.1), + env={}, + data={}, + grpo=SimpleNamespace(skip_reference_policy_logprobs_calculation=True), + logger={}, + cluster={}, + checkpointing={}, + ) - src = inspect.getsource(grpo.setup) - marker = "skip_reference_policy_logprobs_calculation:" - assert marker in src - guard = src[src.index(marker) : src.index(marker) + 400] - assert "raise ValueError" in guard, "the reachable GRPO guard must not be an assert" - assert "reference_policy_kl_penalty=0" in guard + with pytest.raises(ValueError, match="reference_policy_kl_penalty=0"): + grpo_setup(master_config, MagicMock(), MagicMock(), None) @pytest.mark.parametrize( From d64af0e6d7dc4c29c24fe35df4ad2ae1a20d28c6 Mon Sep 17 00:00:00 2001 From: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> Date: Sun, 30 Aug 2026 22:40:53 -0400 Subject: [PATCH 4/4] test(algorithms): avoid broad KL guard exception checks Signed-off-by: Tianyi Zhang <123608656+tianyi-zhang-02@users.noreply.github.com> --- .../unit/algorithms/test_reference_kl_guard.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/unit/algorithms/test_reference_kl_guard.py b/tests/unit/algorithms/test_reference_kl_guard.py index 45e2715a6a8..13d3d48ad19 100644 --- a/tests/unit/algorithms/test_reference_kl_guard.py +++ b/tests/unit/algorithms/test_reference_kl_guard.py @@ -26,7 +26,7 @@ from __future__ import annotations from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest @@ -116,10 +116,11 @@ def test_a_nonzero_kl_penalty_with_skipped_reference_logprobs_raises(call, block ids=["sync-ppo", "sync-grpo"], ) def test_the_supported_pairing_gets_past_the_guard(call, block): - """kl_penalty=0 is the combination the guard exists to allow. It must not - raise ValueError here -- these mocks fail later, which is fine and is what - keeps the test from passing vacuously if the guard were made - unconditional.""" - with pytest.raises(Exception) as excinfo: - call(_master_config(block, kl_penalty=0.0)) - assert "reference_policy_kl_penalty=0" not in str(excinfo.value) + """kl_penalty=0 reaches the first supported-path side effect.""" + + class ReachedSupportedPath(Exception): + pass + + with patch("builtins.print", side_effect=ReachedSupportedPath): + with pytest.raises(ReachedSupportedPath): + call(_master_config(block, kl_penalty=0.0))