diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index 66ec0b5bc5e..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,10 +797,6 @@ 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" - ) 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/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 14de4c6797f..280a83705ce 100644 --- a/nemo_rl/algorithms/ppo.py +++ b/nemo_rl/algorithms/ppo.py @@ -1253,7 +1253,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..13d3d48ad19 --- /dev/null +++ b/tests/unit/algorithms/test_reference_kl_guard.py @@ -0,0 +1,126 @@ +# 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, patch + +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 + + +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, + ) + + +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={}, + ) + + with pytest.raises(ValueError, match="reference_policy_kl_penalty=0"): + grpo_setup(master_config, MagicMock(), MagicMock(), None) + + +@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 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))