fix(algorithms): raise instead of asserting on the skipped-reference-KL pairing

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."
         )

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."
         )

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."
         )

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))