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

[babyplutokurt]

Fixes #3225

What

Forward the trainer-calibrated per-tensor FP8 KV cache scales (kv_scales)
on the colocated ZMQ-IPC refit path, so they actually reach the generation
worker.

Why

calibrate_qkv_fp8_scales computes per-tensor K/V scales every step, but they
were only forwarded on the non-colocated NCCL path
(broadcast_weights_for_collective(kv_scales=...)). The colocated ZMQ-IPC path
called stream_weights_via_ipc_zmq(buffer_size_bytes=...) without
kv_scales, so _iter_params_with_optional_kv_scales received None and
defaulted every per-tensor K/V scale to 1.0
(megatron_policy_worker.py, scale_value = 1.0).

Net effect: in the common colocated GRPO/PPO setup, FP8 KV cache runs
silently used scale 1.0 instead of the calibrated values. The calibration ran
each step but its output was discarded at the refit call site — no error, just
uncalibrated KV cache.

Change

Two one-line additions, passing kv_scales on both colocated IPC call sites so
they match the collective path:

• nemo_rl/algorithms/grpo.py — refit_policy_generation, the ZMQ IPC branch.
• nemo_rl/weight_sync/ipc_weight_synchronizer.py — IPCWeightSynchronizer.sync_weights
  (the method already accepted kv_scales but dropped it at the call).

No new machinery; precision-independent; the receiving
stream_weights_via_ipc_zmq → _iter_params_with_optional_kv_scales chain
already threads kv_scales end-to-end.

Verification

On a colocated run, the streamed kv_scales now populate the worker's
per-tensor scales each step (calibrated, non-1.0, drifting with activations)
instead of defaulting to 1.0.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[terrykong]

@sharonyu-115 to review

[ZhiyuLi-Nvidia]

Thanks a lot @babyplutokurt.

Fix works — median max_seq_mult_prob_error drops 20× (132 → 6.5), median policy_kl_error drops and median policy_kl_error drops 0.27 → 0.05 (5×).

[ZhiyuLi-Nvidia]

LGTM.

Just 2 nits:

1. Tighten stream_weights_via_ipc_zmq interface signature (structural root cause)

nemo_rl/models/policy/interfaces.py:189 currently uses *args, **kwargs, which is why the type checker couldn't warn the two IPC call sites that dropped kv_scales. All four concrete implementations (lm_policy, MegatronPolicyWorker, DTensorPolicyWorker, DTensorPolicyWorkerV2) already have the concrete typed signature, and both callers (grpo.py::refit_policy_generation, IPCWeightSynchronizer.sync_weights) pass compatible kwargs — so tightening the abstract is a no-op at runtime and safe against every existing call:

--- a/nemo_rl/models/policy/interfaces.py
+++ b/nemo_rl/models/policy/interfaces.py
@@ -186,9 +186,11 @@
     ) -> list[ray.ObjectRef]:
         pass

     @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

2. Loud-guard the HTTP synchronizer's silent kv_scales swallow (defensive)

HTTPWeightSynchronizer is wired exclusively for SGLang (nemo_rl/weight_sync/factory.py:91), and SGLangGeneration never overrides the default requires_kv_scale_sync = False, so refit_policy_generation never sends non-None kv_scales down this path today. Not a live bug. But HTTPWeightSynchronizer.sync_weights still accepts and silently drops the kv_scales kwarg (http_weight_synchronizer.py:67) — exactly the pattern this PR is fixing on the IPC path. A one-line assert keeps it interface-uniform without adding real behavior, and makes any future misconfig loud:

--- a/nemo_rl/weight_sync/http_weight_synchronizer.py
+++ b/nemo_rl/weight_sync/http_weight_synchronizer.py
@@ -63,6 +63,10 @@
     def sync_weights(
         self,
         *,
         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 you're "
+            "seeing this, somethiated scales down "
+            "a path that will silently ignore them."
+        )
         self._policy.offload_before_refit()
         self._generation.prepare_"])

[ZhiyuLi-Nvidia]

/ok to test a6ed634

[ZhiyuLi-Nvidia]

By the way, could you also rebase onto main branch? CICD NeMo RL / Check if PR branch is up to date (push)Failing after 3s

[babyplutokurt]

@ZhiyuLi-Nvidia Cool, added your suggested changes and rebased.

[ZhiyuLi-Nvidia]

/ok to test f9bd628

[ZhiyuLi-Nvidia]

/ok to test 5ee617c

[babyplutokurt]

@ZhiyuLi-Nvidia Weird, one of the presubmit failed, but I don't think it's related to my code changes.

L1_Functional_Tests_Megatron_1.sh skipped grpo_megatron.sh and grpo_megatron_mbridge_restore.sh and ran only:

1. tests/functional/audio_grpo_megatron.sh
2. tests/functional/grpo_megatron_eagle3_online.sh

And the job dies at mid-run, can you take a look?

[ZhiyuLi-Nvidia]

Sure. Let me rerun might be transition error.

[babyplutokurt]

From logs, looks like a huggingface flaky outage issue

[babyplutokurt]

  Retrying in 2s [Retry 2/5]. ... Retrying in 8s [Retry 5/5].
  ...
  File "/opt/nemo-rl/nemo_rl/data/datasets/response_datasets/avqa.py", line 97, in __init__
      self.dataset = load_dataset("gijs/avqa-processed", split=split)
  FileNotFoundError: Couldn't find any data file at /opt/nemo-rl/gijs/avqa-processed.
    Couldn't find 'gijs/avqa-processed' on the Hugging Face Hub either:
    LocalEntryNotFoundError: An error happened while trying to locate the file on the Hub

[babyplutokurt]

1. L0_Unit_Tests_Algorithms failure is on me (Now fiexed):

FAILED test_refit_policy_generation_forwards_kv_scales_on_colocated_ipc
AssertionError: Expected 'stream_weights_via_ipc_zmq' to be called once. Called 0 times.

refit_policy_generation opens with a synchronizer delegation guard:

# nemo_rl/algorithms/grpo.py:2167
synchronizer = getattr(policy_generation, "weight_synchronizer", None)
if synchronizer is not None:
    return synchronizer.sync_weights(timer=timer, kv_scales=kv_scales) or {}

VllmGeneration declares self.weight_synchronizer: WeightSynchronizer | None = None
(vllm_generation.py:106), so in production this is None and control falls through to
the colocated IPC branch. But MagicMock() auto-creates any attribute on access, so
getattr(...) returned a truthy child mock, the call returned early via
synchronizer.sync_weights(...), and the IPC branch under test never executed.

Fix: pin the attribute to the real default:

 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}

2. fast_L1_Functional_Tests_Megatron_1 — unrelated HF Hub outage

Died 1m13s into audio_grpo_megatron.sh, in dataset setup, before any training/refit ran:

HTTP Error 503 ... HEAD https://huggingface.co/datasets/gijs/avqa-processed/resolve/main/README.md
Retrying in 2s [Retry 2/5]. ... Retrying in 8s [Retry 5/5].
...
File "nemo_rl/data/datasets/response_datasets/avqa.py", line 97, in __init__
    self.dataset = load_dataset("gijs/avqa-processed", split=split)
FileNotFoundError: Couldn't find any data file at /opt/nemo-rl/gijs/avqa-processed.

37 HTTP 503s in the log; all five huggingface_hub retries exhausted. Because the shard
runs under set -eou pipefail, this aborted before grpo_megatron_eagle3_online.sh ran
at all. No code path from this PR was reached.

@ZhiyuLi-Nvidia could you re-trigger when you get a chance?

[terrykong]

lgtm. thanks @babyplutokurt for the fix

[terrykong]

/ok to test 36aef4b