fix(sc): gate reference logprobs on KL penalty

nemo_rl/algorithms/loss/loss_functions.py

@@ -15,7 +15,7 @@
 from typing import Any, NotRequired, Optional, TypedDict, TypeVar
 
 import torch
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from nemo_rl.algorithms.loss.interfaces import (
     LossFunction,
@@ -126,7 +126,7 @@ class ClippedPGLossConfig(BaseModel, extra="allow"):
     ratio_clip_c: Optional[float] = None
 
     # --- KL regularization ---
-    reference_policy_kl_penalty: float = 0.01
+    reference_policy_kl_penalty: float = Field(default=0.01, ge=0, allow_inf_nan=False)
     # Can be set to k1, k2, k3
     # For more details, see http://joschu.net/blog/kl-approx.html
     reference_policy_kl_type: str = "k3"

nemo_rl/algorithms/single_controller.py

@@ -101,7 +101,7 @@
 from nemo_rl.data.interfaces import DatumSpec
 from nemo_rl.data_plane import DATA_PLANE_CHECKPOINT_SCHEMA_VERSION, KVBatchMeta
 from nemo_rl.data_plane.async_utils import call_data_plane
-from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS
+from nemo_rl.data_plane.schema import DP_CALIB_INPUT_FIELDS, DP_TRAIN_FIELDS
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.distributed.refit_watchdog import RefitAborted, is_refit_context_lost
 from nemo_rl.environments.nemo_gym import should_use_nemo_gym
@@ -149,6 +149,20 @@ def _pooled_opd_metrics(
     }
 
 
+def _train_fields_for_step(
+    *,
+    policy_logprobs_required: bool,
+    reference_logprobs_required: bool,
+) -> tuple[str, ...]:
+    """Return only the data-plane columns produced for this train step."""
+    return tuple(
+        field
+        for field in DP_TRAIN_FIELDS
+        if (policy_logprobs_required or field != "prev_logprobs")
+        and (reference_logprobs_required or field != "reference_policy_logprobs")
+    )
+
+
 @ray.remote(num_cpus=1, num_gpus=0)  # pragma: no cover
 class SingleControllerActor:
     """CPU-only Ray actor that orchestrates the RL training loop.
@@ -206,11 +220,17 @@ def __init__(
             master_config.loss_fn.force_on_policy_ratio
             and self._algo_cfg.seq_logprob_error_threshold is None
         )
+        # _build_trainer initializes the reference model only for a positive KL
+        # penalty, so the controller must use the same gate before requesting it.
         self._reference_logprobs_required = bool(
             master_config.loss_fn.reference_policy_kl_penalty > 0
             and not self._algo_cfg.skip_reference_policy_logprobs_calculation
         )
         self._teacher_logprobs_required = opd_module.is_opd_enabled(master_config)
+        self._train_fields = _train_fields_for_step(
+            policy_logprobs_required=self._policy_logprobs_required,
+            reference_logprobs_required=self._reference_logprobs_required,
+        )
         self._dp_client = actor_args.dp_client
         self._gen: Generation = actor_args.gen_handle
         self._trainer: TQPolicy = actor_args.trainer_handle
@@ -1874,6 +1894,7 @@ async def _train_pump(self) -> None:
                                     await asyncio.to_thread(
                                         self._trainer.train_microbatches_from_meta,
                                         train_meta,
+                                        train_fields=self._train_fields,
                                     )
                                     # A PPO step is one chunk: nothing to
                                     # accumulate, so close every epoch here.

nemo_rl/algorithms/single_controller_utils/config.py

@@ -1020,6 +1020,21 @@ def validate_single_controller_config(master_config: MasterConfig) -> None:
             "loss_fn.reference_policy_kl_penalty=0."
         )
 
+    if (
+        master_config.loss_fn.use_kl_in_reward
+        and reference_policy_kl_penalty > 0
+        and master_config.loss_fn.force_on_policy_ratio
+        and algo_cfg.seq_logprob_error_threshold is None
+    ):
+        raise ValueError(
+            "loss_fn.use_kl_in_reward=true with a nonzero "
+            "loss_fn.reference_policy_kl_penalty requires policy logprobs, but "
+            "loss_fn.force_on_policy_ratio=true without "
+            "seq_logprob_error_threshold skips them. Set "
+            "loss_fn.force_on_policy_ratio=false or configure "
+            "seq_logprob_error_threshold."
+        )
+
     # ``env`` is required in production configs, but model_construct-based unit
     # configs can omit it. Only apply rollout-path validation when it is present.
     env_config = getattr(master_config, "env", None)

nemo_rl/models/policy/tq_policy.py

@@ -438,6 +438,7 @@ def train_microbatches_from_meta(
         self,
         meta: KVBatchMeta,
         timer: Optional[Timer] = None,
+        train_fields: tuple[str, ...] = DP_TRAIN_FIELDS,
     ) -> None:
         """Dispatch one meta slice (DP-sharded) into an open train step.
 
@@ -452,12 +453,17 @@ def train_microbatches_from_meta(
         ``.grad``. Returns nothing — per-microbatch metrics accumulate in
         the workers' open-step state and surface once via
         :meth:`finish_train_step`.
+
+        Args:
+            meta: Data-plane metadata for the samples in this chunk.
+            timer: Optional timer for nested policy-training measurements.
+            train_fields: Columns produced for this step and fetched by workers.
         """
         spa, dba = self._packing_args("train_mb_tokens")
         train_meta = self._isolated_meta(
             meta,
             fields=fields_with_optional_routed_experts(
-                DP_TRAIN_FIELDS, enabled=self._router_replay_enabled
+                train_fields, enabled=self._router_replay_enabled
             ),
             task_name="train",
         )

tests/unit/algorithms/test_loss_functions.py

@@ -43,6 +43,17 @@
 )
 
 
+@pytest.mark.parametrize(
+    "invalid_penalty",
+    [-0.01, float("nan"), float("inf"), float("-inf")],
+)
+def test_clipped_pg_loss_config_rejects_invalid_reference_kl_penalty(
+    invalid_penalty: float,
+) -> None:
+    with pytest.raises(ValueError):
+        ClippedPGLossConfig(reference_policy_kl_penalty=invalid_penalty)
+
+
 def setup_dpo_loss_test_data(vocab_size=16, batch_size=1):
     seq_len = 4
     data = {

tests/unit/models/policy/test_split_api_wrappers.py

@@ -164,6 +164,27 @@ def test_train_microbatches_from_meta_dispatches_and_returns_none(self):
         # get_all_worker_results (unlike the single-data fan-outs)
         wg.get_all_worker_results.assert_called_once()
 
+    def test_train_microbatches_fetches_only_requested_fields(self):
+        p, _ = _make_tq_policy()
+        meta = _meta()
+        train_fields = tuple(
+            field
+            for field in DP_TRAIN_FIELDS
+            if field not in {"prev_logprobs", "reference_policy_logprobs"}
+        )
+        with (
+            patch.object(TQPolicy, "_stamp_pad_seqlen"),
+            patch.object(TQPolicy, "_packing_args", return_value=(None, None)),
+            patch(
+                "nemo_rl.models.policy.tq_policy.shard_meta_for_dp",
+                return_value=([meta, meta], None),
+            ) as mock_shard,
+        ):
+            p.train_microbatches_from_meta(meta, train_fields=train_fields)
+
+        train_meta = mock_shard.call_args.args[0]
+        assert train_meta.fields == list(train_fields)
+
     def test_train_microbatches_requests_routed_experts_for_router_replay(self):
         p, _ = _make_tq_policy()
         p._router_replay_enabled = True

tests/unit/single_controller/test_checkpointing.py

@@ -126,7 +126,9 @@ def prepare_for_training(self) -> None:
     def begin_train_step(self, loss_fn: Any) -> None:
         pass
 
-    def train_microbatches_from_meta(self, meta: KVBatchMeta) -> None:
+    def train_microbatches_from_meta(
+        self, meta: KVBatchMeta, *, train_fields: tuple[str, ...]
+    ) -> None:
         pass
 
     def finish_train_step(self) -> dict[str, Any]:

tests/unit/single_controller/test_resiliency_config.py

@@ -21,14 +21,14 @@
 """
 
 import warnings
-from types import SimpleNamespace
 from typing import get_args
 
 import pytest
 from pydantic import ValidationError
 
 from nemo_rl.algorithms.async_utils.staleness_sampler import SamplerConfig
 from nemo_rl.algorithms.grpo import GRPOConfig
+from nemo_rl.algorithms.loss import ClippedPGLossConfig
 from nemo_rl.algorithms.single_controller_utils.config import (
     AsyncRLConfig,
     FleetHealthConfig,
@@ -68,9 +68,7 @@ def _master_config(*, num_prompts_per_step: int = 8, **async_kwargs) -> MasterCo
             "train_global_batch_size": num_prompts_per_step * 4,
             "generation": {"colocated": {"enabled": False}},
         },
-        # The last two are read only on the ready_first branch, which rejects a run
-        # without them before it reaches anything under test here.
-        loss_fn=SimpleNamespace(
+        loss_fn=ClippedPGLossConfig(
             reference_policy_kl_penalty=0,
             use_importance_sampling_correction=True,
             force_on_policy_ratio=False,
@@ -409,7 +407,7 @@ def _master_config(*, use_nemo_gym: bool, rollout_failure: dict) -> MasterConfig
                 "train_global_batch_size": 8,
                 "generation": {"colocated": {"enabled": False}},
             },
-            loss_fn=SimpleNamespace(reference_policy_kl_penalty=0),
+            loss_fn=ClippedPGLossConfig(reference_policy_kl_penalty=0),
             env={"should_use_nemo_gym": use_nemo_gym},
             # Read by the metric_name check upstream #3429 added to this same
             # validator, which runs before the wrong-path check under test.

tests/unit/single_controller/test_setup.py

@@ -427,6 +427,36 @@ def test_single_controller_mopd_recipe_resolves_to_runtime_contract():
     )
 
 
+@pytest.mark.parametrize(
+    ("reference_policy_kl_penalty", "expected_init_reference_model"),
+    [(0.0, False), (0.01, True)],
+)
+def test_build_trainer_initializes_reference_model_only_for_nonzero_kl(
+    reference_policy_kl_penalty: float,
+    expected_init_reference_model: bool,
+) -> None:
+    master_config = _make_master_config(
+        loss_cfg=ClippedPGLossConfig(
+            reference_policy_kl_penalty=reference_policy_kl_penalty
+        )
+    )
+
+    with patch.object(sc_setup_mod, "TQPolicy") as mock_policy:
+        sc_setup_mod._build_trainer(
+            MagicMock(name="train_cluster"),
+            master_config,
+            MagicMock(name="tokenizer"),
+            None,
+            weights_path=None,
+            optimizer_path=None,
+        )
+
+    assert (
+        mock_policy.call_args.kwargs["init_reference_model"]
+        is expected_init_reference_model
+    )
+
+
 class TestSetup:
     """setup arg validation + actor_args assembly."""
 
@@ -435,6 +465,35 @@ def test_raises_when_data_plane_disabled(self):
         with pytest.raises(ValueError, match="data_plane.enabled=True"):
             setup_single_controller(mc, MagicMock())
 
+    def test_nonzero_kl_rejects_skipping_reference_logprobs(self, patched_factories):
+        mc = _make_master_config(
+            loss_cfg=ClippedPGLossConfig(reference_policy_kl_penalty=0.01)
+        )
+        mc.grpo.skip_reference_policy_logprobs_calculation = True
+
+        with pytest.raises(ValueError, match="requires reference_policy_logprobs"):
+            setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        patched_factories["setup_response_data"].assert_not_called()
+        patched_factories["_build_clusters"].assert_not_called()
+        patched_factories["_build_trainer"].assert_not_called()
+
+    def test_reward_kl_rejects_skipping_policy_logprobs(self, patched_factories):
+        mc = _make_master_config(
+            loss_cfg=ClippedPGLossConfig(
+                reference_policy_kl_penalty=0.01,
+                use_kl_in_reward=True,
+                force_on_policy_ratio=True,
+            )
+        )
+
+        with pytest.raises(ValueError, match="requires policy logprobs"):
+            setup_single_controller(mc, MagicMock(pad_token_id=0))
+
+        patched_factories["setup_response_data"].assert_not_called()
+        patched_factories["_build_clusters"].assert_not_called()
+        patched_factories["_build_trainer"].assert_not_called()
+
     def test_rejects_mooncake_data_plane_checkpointing(self):
         mc = _make_master_config()
         mc.data_plane["backend"] = "mooncake_cpu"