fix(refit): stream a parameter larger than the IPC staging buffer

nemo_rl/models/policy/utils.py

@@ -443,9 +443,48 @@ def release_staging_buffers() -> None:
                 current_buffer = buffer_a
 
             aligned_size = calculate_aligned_size(tensor.nbytes)
-            assert aligned_size <= buffer_size_bytes, (
-                f"Parameter {name} too large for buffer: {aligned_size} > {buffer_size_bytes}"
-            )
+
+            # A parameter larger than a single staging buffer cannot be packed
+            # at all. Ship it on its own in a buffer sized to fit rather than
+            # failing the refit. The staging buffers are sized from *free
+            # memory* (NRL_REFIT_BUFFER_MEMORY_RATIO, default 0.3, halved again
+            # for ping-pong) with no floor at the largest parameter, so a big
+            # embedding can exceed one: DeepSeek-V3's model.embed_tokens.weight
+            # is 1.73 GiB against a 1.65 GiB buffer. This mirrors the HTTP
+            # streaming path, which already gives an oversized parameter a
+            # bucket of its own instead of raising.
+            if aligned_size > buffer_size_bytes:
+                if param_names:
+                    await_recv = send_buffer_group_overlap(
+                        current_buffer, param_names, used_bytes, await_recv
+                    )
+                    count_of_groups += 1
+                    current_buffer = (
+                        buffer_b if current_buffer is buffer_a else buffer_a
+                    )
+                    used_bytes, param_names = 0, []
+
+                oversized_buffer = torch.empty(
+                    aligned_size,
+                    device=current_buffer.device,
+                    dtype=torch.uint8,
+                    requires_grad=False,
+                )
+                try:
+                    packed_bytes = pack_tensor(oversized_buffer, tensor, 0)
+                    send_buffer_group_overlap(
+                        oversized_buffer, [name], packed_bytes, await_recv
+                    )
+                    count_of_groups += 1
+                    # Unlike the ping-pong pair, this buffer is not kept alive
+                    # across the next send, so its ACK must be consumed here
+                    # before it is freed.
+                    zmq_socket.recv()
+                    await_recv = False
+                finally:
+                    del oversized_buffer
+                    torch.cuda.empty_cache()
+                continue
 
             # Check if we need to send current buffer and switch to the other one
             if used_bytes + aligned_size > buffer_size_bytes:

tests/unit/models/policy/test_refit_oversized_param.py

@@ -0,0 +1,169 @@
+# Copyright (c) 2025, 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.
+
+"""IPC weight streaming when one parameter exceeds the staging buffer.
+
+The staging buffers are sized from *free memory*
+(``NRL_REFIT_BUFFER_MEMORY_RATIO``, default 0.3, halved again for ping-pong)
+with no floor at the largest parameter, so a large embedding can be bigger than
+a single buffer and cannot be packed at all. That used to abort the refit with
+``AssertionError: Parameter ... too large for buffer``; DeepSeek-V3 hit it with
+``model.embed_tokens.weight`` at 1.73 GiB against a 1.65 GiB buffer.
+
+These tests run on CPU tensors with a stub socket -- no GPU, no Ray -- because
+the packing/hand-off logic is plain Python around a byte buffer.
+"""
+
+import torch
+
+from nemo_rl.models.policy import utils
+
+# Small enough to keep the test instant; the arithmetic is scale-free.
+BUFFER_BYTES = 4096
+ALIGNMENT = 512
+
+
+class FakeSocket:
+    """Records the payloads a policy worker streams, and ACKs each one."""
+
+    def __init__(self):
+        self.payloads = []
+        self.completed = False
+        self.pending_acks = 0
+
+    def send_pyobj(self, payload):
+        # The receiver ACKs the end-of-stream marker too, not just data groups.
+        if payload is utils.IPCProtocol.COMPLETE:
+            self.completed = True
+            self.pending_acks += 1
+            return
+        _handle, param_names, used_bytes = payload
+        self.payloads.append((list(param_names), used_bytes))
+        self.pending_acks += 1
+
+    def recv(self):
+        # A receiver only ACKs a group it was actually sent. If the streamer
+        # waits for an ACK it is not owed, that is a protocol bug worth failing
+        # on rather than deadlocking.
+        assert self.pending_acks > 0, "streamer waited for an ACK it was not owed"
+        self.pending_acks -= 1
+        return b""
+
+    def getsockopt(self, _opt):
+        return 0
+
+
+def _stream(monkeypatch, params, buffer_size_bytes=BUFFER_BYTES):
+    """Run the streamer over ``params`` with CUDA calls stubbed out."""
+    monkeypatch.setattr(utils, "get_handle_from_tensor", lambda tensor: ("handle",))
+    # send_buffer_group_overlap synchronizes the current CUDA stream; on a
+    # CPU-only box there is none.
+    monkeypatch.setattr(
+        torch.cuda, "current_stream", lambda *a, **k: _NullStream(), raising=False
+    )
+    monkeypatch.setattr(torch.cuda, "empty_cache", lambda *a, **k: None, raising=False)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+
+    socket = FakeSocket()
+    utils.stream_weights_via_ipc_zmq_impl(
+        params_generator=iter(params),
+        buffer_size_bytes=buffer_size_bytes,
+        zmq_socket=socket,
+        rank=0,
+        worker_name="test-worker",
+    )
+    return socket
+
+
+class _NullStream:
+    def synchronize(self):
+        pass
+
+
+def _param(name, nbytes):
+    return name, torch.zeros(nbytes, dtype=torch.uint8)
+
+
+def test_oversized_parameter_is_streamed_instead_of_aborting(monkeypatch):
+    """The parameter that used to trip the assertion is delivered on its own."""
+    # buffer_size_bytes is halved for ping-pong, so the usable buffer is 2048.
+    # 3072 exceeds it, exactly like embed_tokens.weight vs its staging buffer.
+    params = [
+        _param("model.layers.0.weight", 512),
+        _param("model.embed_tokens.weight", 3072),
+        _param("model.layers.1.weight", 512),
+    ]
+
+    socket = _stream(monkeypatch, params)
+
+    streamed = [name for names, _ in socket.payloads for name in names]
+    assert streamed == [p[0] for p in params], (
+        "every parameter must be delivered exactly once, in order"
+    )
+    assert socket.completed, "the stream must still be terminated with COMPLETE"
+    assert socket.pending_acks == 0, "every sent group must be ACKed"
+
+    # The oversized parameter cannot share a group with anything else.
+    oversized_group = [
+        names for names, _ in socket.payloads if "model.embed_tokens.weight" in names
+    ]
+    assert oversized_group == [["model.embed_tokens.weight"]]
+
+
+def test_oversized_parameter_alone(monkeypatch):
+    """A single oversized parameter needs no preceding group to flush."""
+    socket = _stream(monkeypatch, [_param("model.embed_tokens.weight", 3072)])
+
+    assert socket.payloads == [(["model.embed_tokens.weight"], 3072)]
+    assert socket.completed
+    assert socket.pending_acks == 0
+
+
+def test_consecutive_oversized_parameters(monkeypatch):
+    """Back-to-back oversized parameters each get their own buffer and ACK."""
+    params = [_param("a.weight", 3072), _param("b.weight", 4096)]
+
+    socket = _stream(monkeypatch, params)
+
+    assert [names for names, _ in socket.payloads] == [["a.weight"], ["b.weight"]]
+    assert socket.pending_acks == 0
+
+
+def test_parameters_that_fit_are_still_batched(monkeypatch):
+    """Regression guard: the common path must keep packing many params per group.
+
+    This is the behaviour every currently-passing refit relies on, so it must
+    not change -- the oversized branch is only reachable where the old code
+    raised.
+    """
+    params = [_param(f"layer.{i}.weight", 512) for i in range(8)]
+
+    socket = _stream(monkeypatch, params)
+
+    assert len(socket.payloads) < len(params), (
+        "parameters that fit must be batched, not sent one per group"
+    )
+    streamed = [name for names, _ in socket.payloads for name in names]
+    assert streamed == [p[0] for p in params]
+    assert socket.pending_acks == 0
+
+
+def test_alignment_is_what_decides_oversized(monkeypatch):
+    """A parameter is oversized by its *aligned* size, not its raw size."""
+    # 2048 raw fits the 2048-byte usable buffer exactly; 2049 aligns up to 2560.
+    assert utils.calculate_aligned_size(2048, ALIGNMENT) == 2048
+    assert utils.calculate_aligned_size(2049, ALIGNMENT) == 2560
+
+    socket = _stream(monkeypatch, [_param("exact.weight", 2048)])
+    assert socket.payloads == [(["exact.weight"], 2048)]