fix: forward calibrated KV cache scales on colocated ZMQ IPC refit path

nemo_rl/algorithms/grpo.py

@@ -2222,7 +2222,8 @@ def refit_policy_generation(
             else:
                 # Original ZMQ IPC path for vLLM
                 futures_train = policy.stream_weights_via_ipc_zmq(
-                    buffer_size_bytes=buffer_size_bytes
+                    buffer_size_bytes=buffer_size_bytes,
+                    kv_scales=kv_scales,
                 )
                 futures_inference = policy_generation.update_weights_via_ipc_zmq()
                 # wait for all futures to complete

nemo_rl/models/policy/interfaces.py

@@ -187,7 +187,9 @@ def prepare_refit_info(self) -> Optional[dict[str, Any]]:
 
     @abstractmethod
     def stream_weights_via_ipc_zmq(
-        self, *args: Any, **kwargs: Any
+        self,
+        buffer_size_bytes: int,
+        kv_scales: Optional[dict[str, float]] = None,
     ) -> list[ray.ObjectRef]:
         pass
 

nemo_rl/weight_sync/http_weight_synchronizer.py

@@ -66,6 +66,12 @@ def sync_weights(
         timer: Optional[Timer] = None,
         kv_scales: Optional[dict[str, float]] = None,
     ) -> None:
+        assert kv_scales is None, (
+            "HTTP transport (SGLang) does not support FP8 KV cache scale sync. "
+            "`kv_scales` is accepted only for interface uniformity; if this "
+            "assertion fires, calibrated scales were routed through a path that "
+            "would otherwise silently ignore them."
+        )
         self._policy.offload_before_refit()
         self._generation.prepare_for_generation(tags=["weights"])
 

nemo_rl/weight_sync/interfaces.py

@@ -85,9 +85,8 @@ def sync_weights(
         Args:
             timer: Optional Timer for profiling individual phases.
             kv_scales: Optional KV cache scales for FP8 quantization.
-                **Note**: Only honored by the NCCL collective transport,
-                which forwards them to ``policy.broadcast_weights_for_collective()``.
-                IPC and HTTP transports ignore this parameter.
+                Honored by the IPC/ZMQ and NCCL collective transports. The
+                HTTP transport ignores this parameter.
 
         Returns:
             Optional transport-specific scalar metrics for the current sync.

nemo_rl/weight_sync/ipc_weight_synchronizer.py

@@ -83,7 +83,8 @@ def sync_weights(
                 buffer_size_bytes = self._compute_buffer_size()
 
                 futures_train = self._policy.stream_weights_via_ipc_zmq(
-                    buffer_size_bytes=buffer_size_bytes
+                    buffer_size_bytes=buffer_size_bytes,
+                    kv_scales=kv_scales,
                 )
                 futures_inference = self._generation.update_weights_via_ipc_zmq()
 

tests/unit/algorithms/test_grpo.py

@@ -42,6 +42,7 @@
     compute_and_apply_seq_logprob_error_masking,
     dynamic_sampling,
     grpo_train,
+    refit_policy_generation,
     validate,
 )
 from nemo_rl.algorithms.loss import ClippedPGLossConfig, ClippedPGLossFn
@@ -73,6 +74,33 @@ def _mock_policy_generation() -> MagicMock:
     return policy_generation
 
 
+@patch("nemo_rl.algorithms.grpo.ray")
+def test_refit_policy_generation_forwards_kv_scales_on_colocated_ipc(
+    mock_ray: MagicMock,
+) -> None:
+    mock_ray.get.return_value = [True]
+    policy = MagicMock()
+    policy_generation = MagicMock()
+    # Match VllmGeneration's default; a bare MagicMock would auto-create a truthy
+    # weight_synchronizer and refit_policy_generation would delegate to it instead
+    # of taking the colocated IPC path under test.
+    policy_generation.weight_synchronizer = None
+    kv_scales = {"layer.0": 0.5}
+
+    refit_policy_generation(
+        policy,
+        policy_generation,
+        colocated_inference=True,
+        _refit_buffer_size_gb=1.0,
+        kv_scales=kv_scales,
+    )
+
+    policy.stream_weights_via_ipc_zmq.assert_called_once_with(
+        buffer_size_bytes=1024**3,
+        kv_scales=kv_scales,
+    )
+
+
 class TestMaskSampleFilter:
     def test_masks_env_flagged_samples(self):
         repeated_batch = BatchedDataDict(

tests/unit/weight_sync/test_weight_synchronizer.py

@@ -119,6 +119,19 @@ def test_sync_weights_calls_full_lifecycle(self, mock_ray):
         policy.offload_after_refit.assert_called_once()
         gen.prepare_for_generation.assert_any_call(tags=["kv_cache"])
 
+    @patch("nemo_rl.weight_sync.ipc_weight_synchronizer.ray")
+    def test_sync_weights_passes_kv_scales(self, mock_ray: MagicMock) -> None:
+        mock_ray.get.return_value = [True]
+        policy = _mock_policy()
+        gen = _mock_generation()
+        sync = IPCWeightSynchronizer(policy, gen)
+        kv_scales = {"layer.0": 0.5}
+
+        sync.sync_weights(kv_scales=kv_scales)
+
+        call_kwargs = policy.stream_weights_via_ipc_zmq.call_args
+        assert call_kwargs.kwargs["kv_scales"] == kv_scales
+
     @patch("nemo_rl.weight_sync.ipc_weight_synchronizer.ray")
     def test_sync_weights_raises_on_failure(self, mock_ray):
         mock_ray.get.side_effect = [
@@ -223,6 +236,17 @@ def test_zero_env_ratio_raises(self, mock_ray, monkeypatch):
 
 
 class TestHTTPWeightSynchronizer:
+    def test_sync_weights_rejects_kv_scales(self):
+        policy = _mock_policy()
+        gen = _mock_generation()
+        sync = HTTPWeightSynchronizer(policy, gen)
+
+        with pytest.raises(AssertionError, match="does not support"):
+            sync.sync_weights(kv_scales={"layer.0": 0.5})
+
+        policy.offload_before_refit.assert_not_called()
+        gen.prepare_for_generation.assert_not_called()
+
     @patch("nemo_rl.weight_sync.http_weight_synchronizer.ray")
     def test_sync_weights_calls_full_lifecycle(self, mock_ray):
         mock_ray.get.return_value = [True]