From 0f7b614354ff9cb1ad87357eb621d5d764f69495 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Wed, 26 Aug 2026 10:21:57 -0700 Subject: [PATCH 1/8] feat(data-plane): add Mooncake GDR support Signed-off-by: Zeyu Zhou --- examples/configs/grpo_math_1B.yaml | 2 + nemo_rl/data_plane/README.md | 24 +++- nemo_rl/data_plane/adapters/transfer_queue.py | 2 + nemo_rl/data_plane/interfaces.py | 30 ++++- nemo_rl/data_plane/worker_mixin.py | 14 ++- tests/unit/data_plane/README.md | 6 + tests/unit/data_plane/test_backend_config.py | 42 ++++++- tests/unit/data_plane/test_mooncake_gdr.py | 113 ++++++++++++++++++ .../unit/reference_configs/grpo_math_1B.yaml | 2 + 9 files changed, 226 insertions(+), 9 deletions(-) create mode 100644 tests/unit/data_plane/test_mooncake_gdr.py diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 99d9ee369cc..6cb0a2eeb4e 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -565,6 +565,8 @@ data_plane: local_buffer_size: 4294967296 # 4 GiB/process reuse_registered_buffers: true # reuse RDMA-registered buffers staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool + use_gdr: false # GPU-memory RDMA staging in CUDA clients + gdr_staging_buffer_mb: 1024 # persistent MiB per active GDR client # observability: # NotRequired # enabled: false diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 44d20e97d8e..e665f928b62 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -437,6 +437,8 @@ data_plane: local_buffer_size: 4294967296 # 4 GiB/process reuse_registered_buffers: true # reuse RDMA-registered buffers staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool + use_gdr: false # GPU-memory RDMA staging in CUDA clients + gdr_staging_buffer_mb: 1024 # persistent MiB per active GDR client # observability: # NotRequired # enabled: false ``` @@ -445,13 +447,29 @@ These keys used to sit directly under `data_plane:`. That spelling is not rejected — it is simply never read. A config still using it silently gets this backend's defaults instead of its own values: an inherited config supplies the nested block, so a surviving flat key always loses the merge, -with no warning either way. +with no warning either way. Earlier revisions of PR #3501 also used flat +`use_gdr` and `gdr_staging_buffer_mb` keys; those two spellings are rejected +with a migration error because silently disabling GDR would invalidate a run. Backend choice: - **`simple`** — ZMQ-backed; lowest setup overhead. Default for tests and small runs. -- **`mooncake_cpu`** — Mooncake transfer engine; higher throughput at - scale. Required for multi-node clusters with large bulk volume. +- **`mooncake_cpu`** — Mooncake's RDMA-only transfer engine. By default, + tensors transfer through registered CPU staging. Set + `mooncake_cpu.use_gdr: true` to let CUDA-initialized clients use + TransferQueue's GDR staging path. CPU-only clients, such as a + SingleController producer, continue to use CPU RDMA. GDR changes the + client-side tensor transfer and staging path; queued objects still reside in + Mooncake-managed host-memory segments. + +The CPU host staging pool's `staging_buffer_size` is independent of the GDR +buffer. `gdr_staging_buffer_mb` is the persistent GPU staging capacity per +active CUDA client and defaults to 1024 MiB. Transfers through that per-client +buffer are serialized. Aggregate fetches may exceed the buffer and are split +into groups. In the mixed CPU-producer/GDR-receiver flow used by +SingleController, however, each individual tensor must currently fit because +the CPU PUT path does not create the chunk metadata required by an oversized +GDR GET. Capacity rule of thumb (any backend): diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 77283fbb9bb..45c864493c9 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -684,6 +684,8 @@ def _init_tq(cfg: DataPlaneConfig) -> None: "metadata_server": f"{local_ip}:50050", "master_server_address": f"{local_ip}:50051", **_mooncake_transport_config(), + "use_gdr": bool(mooncake_cfg.use_gdr), + "gdr_staging_buffer_mb": int(mooncake_cfg.gdr_staging_buffer_mb), }, }, } diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 9f527c4b2c8..7bbff650f01 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -40,7 +40,7 @@ from pathlib import Path from typing import Annotated, Any, Callable, Literal, NotRequired, Sequence, TypedDict -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, PositiveInt from tensordict import TensorDict DATA_PLANE_CHECKPOINT_SCHEMA_VERSION = 2 @@ -80,6 +80,11 @@ class MooncakeCpuConfig(BaseModel, extra="allow"): admitted and never shrink — so raise it only when a per-key payload (one sample of one field) genuinely exceeds it, not for headroom. + ``use_gdr`` lets CUDA-initialized clients transfer through TransferQueue's + persistent GPU staging buffer. ``gdr_staging_buffer_mb`` is the positive + HBM capacity of that buffer per active GDR client. CPU-only clients keep + using the registered host-buffer path. + Every RDMA rail on the host is offered to mooncake (see ``rdma_devices``). That is only safe with ``MC_ENABLE_DEST_DEVICE_AFFINITY=1``, which pins each transfer's peer rail to the local one by name; on a rail-isolated RoCE @@ -91,6 +96,8 @@ class MooncakeCpuConfig(BaseModel, extra="allow"): local_buffer_size: int = 4294967296 # 4 GiB per client process reuse_registered_buffers: bool = True staging_buffer_size: int = 268435456 # 256 MiB per pool slot + use_gdr: bool = False + gdr_staging_buffer_mb: PositiveInt = 1024 class DataPlaneConfig(TypedDict): @@ -117,7 +124,9 @@ class DataPlaneConfig(TypedDict): ``local_buffer_size`` used to sit at this level. A config still using that spelling is not rejected — the flat key is simply never read, and :func:`backend_config` resolves the nested block (or its defaults) as if it - were absent. See there. + were absent. Earlier revisions of PR #3501 also put ``use_gdr`` and + ``gdr_staging_buffer_mb`` here; those two spellings are rejected explicitly + so a stale GDR recipe cannot silently run with GDR disabled. See there. """ enabled: bool @@ -130,6 +139,14 @@ class DataPlaneConfig(TypedDict): ack_timeout_ms: NotRequired[int] observability: NotRequired["ObservabilityConfig"] + # Declared only so Pydantic preserves the first GDR integration's flat + # spellings until backend_config can raise the migration error below. + use_gdr: NotRequired[bool] + gdr_staging_buffer_mb: NotRequired[int] + + +_FIRST_GDR_FLAT_KEYS = ("use_gdr", "gdr_staging_buffer_mb") + _CHECKPOINTABLE_BACKENDS: frozenset[str] = frozenset({"simple"}) @@ -158,9 +175,16 @@ def backend_config(cfg: DataPlaneConfig) -> Any: (block already coerced to a model) or as a plain dict from a test. Sizing is read only from the nested block. A config still using the - pre-nesting flat spelling gets this backend's defaults, not its own values. + pre-nesting flat sizing spelling gets this backend's defaults, not its own + values. The first GDR integration's flat spelling raises instead because + silently disabling a requested transport would invalidate a benchmark. """ backend = cfg["backend"] + stale_gdr = [key for key in _FIRST_GDR_FLAT_KEYS if key in cfg] + if stale_gdr: + keys = ",".join(stale_gdr) + raise ValueError(f"data_plane.{{{keys}}} moved under data_plane.mooncake_cpu") + nested = cfg.get(backend) or {} if isinstance(nested, BaseModel): nested = nested.model_dump(exclude_unset=True) diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index a5309d453f8..f128f24170d 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -36,10 +36,9 @@ import numpy as np import torch -FetchPolicy = Literal["auto", "independent", "leader_broadcast"] - from nemo_rl.data.llm_message_utils import attach_message_log_view from nemo_rl.data.multimodal_utils import PackedTensor +from nemo_rl.data_plane.interfaces import backend_config from nemo_rl.data_plane.schema import ( ELEM_COUNTS_PER_GB, GLOBAL_FORWARD_PAD_SEQLEN, @@ -64,6 +63,8 @@ DataPlaneRuntimeConfig, ) +FetchPolicy = Literal["auto", "independent", "leader_broadcast"] + def _broadcast_batched_data_dict( data: Optional[BatchedDataDict[Any]], @@ -327,6 +328,15 @@ def setup_data_plane(self, cfg: DataPlaneRuntimeConfig) -> None: self._route_fallback_counts = Counter() from nemo_rl.data_plane import build_data_plane_client + if ( + cfg["backend"] == "mooncake_cpu" + and backend_config(cfg).use_gdr + and not torch.cuda.is_initialized() + ): + raise RuntimeError( + "CUDA must be initialized before attaching TransferQueue with GDR" + ) + # bootstrap=False — the driver already created the named # controller actor; this process attaches as a client. self._dp_client = build_data_plane_client(cfg, bootstrap=False) diff --git a/tests/unit/data_plane/README.md b/tests/unit/data_plane/README.md index 14c6cdc55f8..393abd1433d 100644 --- a/tests/unit/data_plane/README.md +++ b/tests/unit/data_plane/README.md @@ -107,6 +107,12 @@ Generated audit of every test function under `tests/unit/data_plane/` with a one - `test_local_node_ip_returns_empty_on_exception` — DNS exception → returns empty string (no crash). - `test_mc_tcp_bind_address_overwrites_existing` — TQDataPlaneClient `__init__` uses direct assignment (not `setdefault`). +## `test_mooncake_gdr.py` (3 tests) + +- `test_init_tq_forwards_nested_gdr_config_and_keeps_rdma` — Nested GDR settings reach TransferQueue without changing #2935's RDMA/all-rail transport. +- `test_cpu_only_client_may_attach_with_gdr_config` — A CPU-only producer may attach because GDR eligibility is client-local. +- `test_gdr_receiver_requires_cuda_initialized` — A policy receiver must initialize CUDA before attaching its GDR client. + ## `test_message_log_decompose.py` (11 tests) - `test_decompose_message_log_basic_shapes` — Basic shapes of decompose output. diff --git a/tests/unit/data_plane/test_backend_config.py b/tests/unit/data_plane/test_backend_config.py index b7bea76c7c8..999ecb2cfcf 100644 --- a/tests/unit/data_plane/test_backend_config.py +++ b/tests/unit/data_plane/test_backend_config.py @@ -61,12 +61,19 @@ def test_checkpointing_capability_defaults_to_unsupported( def test_nested_block_is_used() -> None: cfg = _cfg( "mooncake_cpu", - mooncake_cpu={"global_segment_size": 111, "reuse_registered_buffers": False}, + mooncake_cpu={ + "global_segment_size": 111, + "reuse_registered_buffers": False, + "use_gdr": True, + "gdr_staging_buffer_mb": 512, + }, ) resolved = backend_config(cfg) assert isinstance(resolved, MooncakeCpuConfig) assert resolved.global_segment_size == 111 assert resolved.reuse_registered_buffers is False + assert resolved.use_gdr is True + assert resolved.gdr_staging_buffer_mb == 512 def test_absent_block_falls_back_to_model_defaults() -> None: @@ -81,6 +88,23 @@ def test_absent_block_falls_back_to_model_defaults() -> None: assert resolved.local_buffer_size == 4294967296 # 4 GiB per client process # The opt-out flag defaults on, so omitting it must not disable the pool. assert resolved.reuse_registered_buffers is True + assert resolved.use_gdr is False + assert resolved.gdr_staging_buffer_mb == 1024 + + +@pytest.mark.parametrize("value", [0, -1]) +def test_gdr_staging_size_rejects_non_positive_values(value: int) -> None: + with pytest.raises(pydantic.ValidationError): + backend_config( + _cfg("mooncake_cpu", mooncake_cpu={"gdr_staging_buffer_mb": value}) + ) + + +@pytest.mark.parametrize("backend", ["simple", "mooncake_cpu"]) +@pytest.mark.parametrize("key", ["use_gdr", "gdr_staging_buffer_mb"]) +def test_flat_gdr_key_from_first_integration_raises(backend: str, key: str) -> None: + with pytest.raises(ValueError, match="moved under data_plane.mooncake_cpu"): + backend_config(_cfg(backend, **{key: True})) def test_accepts_an_already_coerced_model() -> None: @@ -133,3 +157,19 @@ def test_schema_validates_without_any_backend_block() -> None: validate, otherwise MasterConfig fails before training starts. """ TypeAdapter(DataPlaneConfig).validate_python(_cfg("simple")) + + +@pytest.mark.parametrize( + ("key", "value"), + [("use_gdr", True), ("gdr_staging_buffer_mb", 512)], +) +def test_pydantic_preserves_flat_gdr_keys_for_migration_error( + key: str, value: object +) -> None: + validated = TypeAdapter(DataPlaneConfig).validate_python( + _cfg("mooncake_cpu", **{key: value}) + ) + + assert key in validated + with pytest.raises(ValueError, match="moved under data_plane.mooncake_cpu"): + backend_config(validated) diff --git a/tests/unit/data_plane/test_mooncake_gdr.py b/tests/unit/data_plane/test_mooncake_gdr.py new file mode 100644 index 00000000000..0babd4bb688 --- /dev/null +++ b/tests/unit/data_plane/test_mooncake_gdr.py @@ -0,0 +1,113 @@ +# 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. +"""GDR configuration plumbing on top of the RDMA-only Mooncake backend.""" + +from __future__ import annotations + +import sys +from types import ModuleType + +import pytest +import torch + +import nemo_rl.data_plane as data_plane_module +from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter +from nemo_rl.data_plane.worker_mixin import TQWorkerMixin + + +def _gdr_cfg(*, staging_mb: int = 384) -> dict: + return { + "enabled": True, + "impl": "transfer_queue", + "backend": "mooncake_cpu", + "claim_meta_poll_interval_s": 0.5, + "mooncake_cpu": { + "use_gdr": True, + "gdr_staging_buffer_mb": staging_mb, + }, + } + + +def test_init_tq_forwards_nested_gdr_config_and_keeps_rdma( + tmp_path, monkeypatch +) -> None: + """GDR changes destination memory, not #2935's transport selection.""" + mooncake_dir = tmp_path / "mooncake" + mooncake_dir.mkdir() + mooncake_init = mooncake_dir / "__init__.py" + mooncake_init.touch() + (mooncake_dir / "mooncake_master").touch() + + mooncake_module = ModuleType("mooncake") + mooncake_module.__file__ = str(mooncake_init) + mooncake_module.__path__ = [str(mooncake_dir)] # type: ignore[attr-defined] + monkeypatch.setitem(sys.modules, "mooncake", mooncake_module) + monkeypatch.setitem(sys.modules, "mooncake.store", ModuleType("mooncake.store")) + + monkeypatch.setattr(tq_adapter.os, "environ", dict(tq_adapter.os.environ)) + monkeypatch.setattr(tq_adapter, "_get_local_node_ip", lambda: "10.0.0.1") + monkeypatch.setattr( + tq_adapter, + "rdma_devices", + lambda: "mlx5_0,mlx5_1,mlx5_2,mlx5_4", + ) + captured: dict = {} + monkeypatch.setattr( + tq_adapter.tq, + "init", + lambda *, conf: captured.update(conf=conf), + ) + + tq_adapter._init_tq(_gdr_cfg(staging_mb=384)) + + store_cfg = captured["conf"]["backend"]["MooncakeStore"] + assert store_cfg["protocol"] == "rdma" + assert store_cfg["device_name"] == "mlx5_0,mlx5_1,mlx5_2,mlx5_4" + assert store_cfg["use_gdr"] is True + assert store_cfg["gdr_staging_buffer_mb"] == 384 + + +class _ReceiverWorker(TQWorkerMixin): + pass + + +def test_cpu_only_client_may_attach_with_gdr_config(monkeypatch) -> None: + """The CUDA guard is receiver-specific; a CPU-only producer remains valid.""" + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: False) + monkeypatch.setattr(tq_adapter, "_connect_existing", lambda: None) + monkeypatch.setattr(tq_adapter, "_get_local_node_ip", lambda: "10.0.0.1") + monkeypatch.setattr(tq_adapter, "_patch_mooncake_register_check", lambda: None) + monkeypatch.setattr( + tq_adapter, "_patch_mooncake_staging_buffers", lambda *_, **__: None + ) + monkeypatch.setattr(tq_adapter.os, "environ", dict(tq_adapter.os.environ)) + + client = tq_adapter.TQDataPlaneClient(_gdr_cfg(), bootstrap=False) + + assert client._closed is False + + +def test_gdr_receiver_requires_cuda_initialized(monkeypatch) -> None: + """A policy receiver must establish its CUDA context before TQ attaches.""" + worker = _ReceiverWorker() + worker._dp_client = None + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: False) + monkeypatch.setattr( + data_plane_module, + "build_data_plane_client", + lambda *args, **kwargs: pytest.fail("client built before CUDA guard"), + ) + + with pytest.raises(RuntimeError, match="CUDA must be initialized"): + worker.setup_data_plane(_gdr_cfg()) diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index 373ad2bffb2..a5a0d134942 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -539,6 +539,8 @@ data_plane: local_buffer_size: 4294967296 # 4 GiB/process reuse_registered_buffers: true # reuse RDMA-registered buffers staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool + use_gdr: false # GPU-memory RDMA staging in CUDA clients + gdr_staging_buffer_mb: 1024 # persistent MiB per active GDR client # observability: # NotRequired # enabled: false From 13764175031d94e1c01583b264041514945ef46d Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Fri, 4 Sep 2026 22:32:19 +0000 Subject: [PATCH 2/8] fix(data-plane): fail closed on GDR PUT fallback Signed-off-by: Zeyu Zhou --- nemo_rl/data_plane/adapters/transfer_queue.py | 42 ++++++++++++++++ tests/unit/data_plane/test_mooncake_gdr.py | 49 ++++++++++++++++++- 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 45c864493c9..029c0711f98 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -27,6 +27,7 @@ import importlib import ipaddress import json +import logging import os import resource import socket @@ -55,6 +56,8 @@ data_plane_supports_checkpointing, ) +LOGGER = logging.getLogger(__name__) + # ────────────────────────────────────────────────────────────────────────── # Backend init — lifted from rl-arena/arena/backends.py. # ────────────────────────────────────────────────────────────────────────── @@ -819,6 +822,13 @@ def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: self._backend = cfg["backend"] self._supports_checkpointing = data_plane_supports_checkpointing(cfg) + # GDR is a mooncake_cpu-only transport knob, so key it off the backend + # directly rather than off any incidental per-backend flag. + self._gdr_requested = self._backend == "mooncake_cpu" and bool( + backend_config(cfg).use_gdr + ) + self._gdr_put_confirmed = False + if bootstrap: _init_tq(cfg) else: @@ -1033,6 +1043,28 @@ def put_samples( wire_fields = detached_fields field_names = [str(key) for key in detached_fields.keys()] + confirm_gdr_put = bool( + self._gdr_requested + and not self._gdr_put_confirmed + and torch.cuda.is_initialized() + and wire_fields is not None + and any( + isinstance(wire_fields.get(key), torch.Tensor) + for key in wire_fields.keys() + ) + ) + gdr_staging = None + if confirm_gdr_put: + tq_client = tq.get_client() + storage_manager = getattr(tq_client, "storage_manager", None) + storage_client = getattr(storage_manager, "storage_client", None) + gdr_staging = getattr(storage_client, "_gdr_staging", None) + if not getattr(storage_client, "use_gdr", False) or gdr_staging is None: + raise RuntimeError( + "GDR was requested for a CUDA-initialized TransferQueue " + "client, but TransferQueue selected CPU RDMA for tensor PUTs" + ) + self._mark_data_operation_started() # TQ's wire vocabulary is `keys=` — translation point. tq.kv_batch_put( @@ -1041,6 +1073,16 @@ def put_samples( fields=wire_fields, tags=user_tags, ) + if confirm_gdr_put: + if not getattr(gdr_staging, "_initialized", False): + raise RuntimeError( + "TransferQueue tensor PUT completed without initializing " + "the requested GDR staging buffer" + ) + LOGGER.info( + "TransferQueue GDR tensor PUT active (partition=%s)", partition_id + ) + self._gdr_put_confirmed = True return KVBatchMeta( partition_id=partition_id, diff --git a/tests/unit/data_plane/test_mooncake_gdr.py b/tests/unit/data_plane/test_mooncake_gdr.py index 0babd4bb688..632d92f1da1 100644 --- a/tests/unit/data_plane/test_mooncake_gdr.py +++ b/tests/unit/data_plane/test_mooncake_gdr.py @@ -15,11 +15,13 @@ from __future__ import annotations +import logging import sys -from types import ModuleType +from types import ModuleType, SimpleNamespace import pytest import torch +from tensordict import TensorDict import nemo_rl.data_plane as data_plane_module from nemo_rl.data_plane.adapters import transfer_queue as tq_adapter @@ -98,6 +100,51 @@ def test_cpu_only_client_may_attach_with_gdr_config(monkeypatch) -> None: assert client._closed is False +def test_gdr_tensor_put_is_confirmed_once_and_never_falls_back( + monkeypatch, caplog +) -> None: + """A CUDA client's first tensor PUT must initialize the GDR path.""" + staging = SimpleNamespace(_initialized=False) + storage_client = SimpleNamespace(use_gdr=True, _gdr_staging=None) + tq_client = SimpleNamespace( + storage_manager=SimpleNamespace(storage_client=storage_client) + ) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: True) + monkeypatch.setattr(tq_adapter, "_connect_existing", lambda: None) + monkeypatch.setattr(tq_adapter, "_get_local_node_ip", lambda: "10.0.0.1") + monkeypatch.setattr(tq_adapter, "_patch_mooncake_register_check", lambda: None) + monkeypatch.setattr( + tq_adapter, "_patch_mooncake_staging_buffers", lambda *_, **__: None + ) + monkeypatch.setattr(tq_adapter.tq, "get_client", lambda: tq_client) + client = tq_adapter.TQDataPlaneClient(_gdr_cfg(), bootstrap=False) + fields = TensorDict( + {"token_ids": torch.tensor([[1]], dtype=torch.int64)}, batch_size=[1] + ) + + monkeypatch.setattr( + tq_adapter.tq, + "kv_batch_put", + lambda **kwargs: pytest.fail("CPU RDMA fallback was allowed"), + ) + with pytest.raises(RuntimeError, match="selected CPU RDMA"): + client.put_samples(["sample-0"], "rollout_staging", fields) + + storage_client._gdr_staging = staging + + def put_with_gdr(**kwargs) -> None: + staging._initialized = True + + monkeypatch.setattr(tq_adapter.tq, "kv_batch_put", put_with_gdr) + with caplog.at_level(logging.INFO, logger=tq_adapter.LOGGER.name): + client.put_samples(["sample-0"], "rollout_staging", fields) + client.put_samples(["sample-1"], "rollout_staging", fields) + + assert caplog.messages.count( + "TransferQueue GDR tensor PUT active (partition=rollout_staging)" + ) == 1 + + def test_gdr_receiver_requires_cuda_initialized(monkeypatch) -> None: """A policy receiver must establish its CUDA context before TQ attaches.""" worker = _ReceiverWorker() From 6955c510288712d664188c67e033efaa750ebc80 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 8 Sep 2026 18:50:45 +0000 Subject: [PATCH 3/8] fix(data-plane): drop the post-PUT GDR check that cannot fail closed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The check ran after `tq.kv_batch_put`, by which point TQ has already written the rows into Mooncake and called `notify_data_update`, so raising could not prevent the fallback it was meant to catch. It only denied the caller its `KVBatchMeta` — and at the one call site where it was reachable, `column_io.write_columns` via `TQWorkerMixin._write_back`, there is no rollback, so the committed rows were left orphaned. It was also unreachable as a genuine failure. The pre-PUT check already tests TQ's own routing condition (`use_gdr and _gdr_staging is not None`), and `_put_tensors_gdr` then calls `lazy_init()` unconditionally, which sets `_initialized` or raises. The only way to trip the post-PUT check was an upstream rename of a TransferQueue private attribute, which would have aborted a healthy GDR run with a message describing a fallback that never happened. The pre-PUT check is kept: it runs before anything is written, so it still fails closed on a real CPU-RDMA fallback, which is what this commit's predecessor set out to do. Signed-off-by: Zeyu Zhou --- nemo_rl/data_plane/adapters/transfer_queue.py | 10 ++++----- tests/unit/data_plane/test_mooncake_gdr.py | 22 +++++++++++-------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index 029c0711f98..c4a9f644ab3 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -1053,8 +1053,11 @@ def put_samples( for key in wire_fields.keys() ) ) - gdr_staging = None if confirm_gdr_put: + # Checked before the put, not after: TQ fixes GDR eligibility when + # the client attaches, so this is decidable up front — and once + # `kv_batch_put` returns, the rows are already durable and the + # controller has been notified, so raising then would strand them. tq_client = tq.get_client() storage_manager = getattr(tq_client, "storage_manager", None) storage_client = getattr(storage_manager, "storage_client", None) @@ -1074,11 +1077,6 @@ def put_samples( tags=user_tags, ) if confirm_gdr_put: - if not getattr(gdr_staging, "_initialized", False): - raise RuntimeError( - "TransferQueue tensor PUT completed without initializing " - "the requested GDR staging buffer" - ) LOGGER.info( "TransferQueue GDR tensor PUT active (partition=%s)", partition_id ) diff --git a/tests/unit/data_plane/test_mooncake_gdr.py b/tests/unit/data_plane/test_mooncake_gdr.py index 632d92f1da1..f8673f3bf55 100644 --- a/tests/unit/data_plane/test_mooncake_gdr.py +++ b/tests/unit/data_plane/test_mooncake_gdr.py @@ -103,8 +103,8 @@ def test_cpu_only_client_may_attach_with_gdr_config(monkeypatch) -> None: def test_gdr_tensor_put_is_confirmed_once_and_never_falls_back( monkeypatch, caplog ) -> None: - """A CUDA client's first tensor PUT must initialize the GDR path.""" - staging = SimpleNamespace(_initialized=False) + """A CUDA client's tensor PUT must take GDR, and say so exactly once.""" + staging = SimpleNamespace() storage_client = SimpleNamespace(use_gdr=True, _gdr_staging=None) tq_client = SimpleNamespace( storage_manager=SimpleNamespace(storage_client=storage_client) @@ -132,17 +132,21 @@ def test_gdr_tensor_put_is_confirmed_once_and_never_falls_back( storage_client._gdr_staging = staging - def put_with_gdr(**kwargs) -> None: - staging._initialized = True - - monkeypatch.setattr(tq_adapter.tq, "kv_batch_put", put_with_gdr) + puts: list[dict] = [] + monkeypatch.setattr( + tq_adapter.tq, "kv_batch_put", lambda **kwargs: puts.append(kwargs) + ) with caplog.at_level(logging.INFO, logger=tq_adapter.LOGGER.name): client.put_samples(["sample-0"], "rollout_staging", fields) client.put_samples(["sample-1"], "rollout_staging", fields) - assert caplog.messages.count( - "TransferQueue GDR tensor PUT active (partition=rollout_staging)" - ) == 1 + assert len(puts) == 2 + assert ( + caplog.messages.count( + "TransferQueue GDR tensor PUT active (partition=rollout_staging)" + ) + == 1 + ) def test_gdr_receiver_requires_cuda_initialized(monkeypatch) -> None: From d3b727a9c090973faca392c2c7263ec0aea71060 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 8 Sep 2026 20:05:07 +0000 Subject: [PATCH 4/8] docs(data-plane): list the fourth GDR test in the unit-test audit `test_gdr_tensor_put_is_confirmed_once_and_never_falls_back` was added without updating the section, which claims to audit every test function under `tests/unit/data_plane/`. Signed-off-by: Zeyu Zhou --- tests/unit/data_plane/README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/data_plane/README.md b/tests/unit/data_plane/README.md index 393abd1433d..719104307b4 100644 --- a/tests/unit/data_plane/README.md +++ b/tests/unit/data_plane/README.md @@ -107,10 +107,11 @@ Generated audit of every test function under `tests/unit/data_plane/` with a one - `test_local_node_ip_returns_empty_on_exception` — DNS exception → returns empty string (no crash). - `test_mc_tcp_bind_address_overwrites_existing` — TQDataPlaneClient `__init__` uses direct assignment (not `setdefault`). -## `test_mooncake_gdr.py` (3 tests) +## `test_mooncake_gdr.py` (4 tests) - `test_init_tq_forwards_nested_gdr_config_and_keeps_rdma` — Nested GDR settings reach TransferQueue without changing #2935's RDMA/all-rail transport. - `test_cpu_only_client_may_attach_with_gdr_config` — A CPU-only producer may attach because GDR eligibility is client-local. +- `test_gdr_tensor_put_is_confirmed_once_and_never_falls_back` — A CUDA client's tensor PUT raises rather than downgrading to CPU RDMA, and logs the GDR confirmation exactly once. - `test_gdr_receiver_requires_cuda_initialized` — A policy receiver must initialize CUDA before attaching its GDR client. ## `test_message_log_decompose.py` (11 tests) From 774308190328c6509bc17bc249e71ae6cb7cda76 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 8 Sep 2026 20:49:53 +0000 Subject: [PATCH 5/8] docs(data-plane): document the HBM cost and sizing of gdr_staging_buffer_mb The knob reserves its full size in HBM per active GDR client for the life of the run, and it has a hard lower bound: below the largest per-sample, per-field payload TransferQueue splits the value into chunk sub-keys, and a GDR read of a value written by a CPU-RDMA client then fails on an assertion inside TransferQueue that carries no message. Sizing also has a knee. An aggregate fetch larger than the buffer is split into groups that transfer one at a time, so lowering it multiplies the group count and costs throughput on both the CPU and GDR paths, while raising it above the default buys little and costs HBM linearly. GDR's gain shows up only once the buffer is large enough that group count is not the bottleneck. None of this was stated anywhere a user would look. The exemplar config is where the value actually gets set, so the guidance goes there; the transfer mechanics stay in nemo_rl/data_plane/README.md. Also aligns the two GDR comments to the block's existing column. Signed-off-by: Zeyu Zhou --- examples/configs/grpo_math_1B.yaml | 12 ++++++++++-- tests/unit/reference_configs/grpo_math_1B.yaml | 12 ++++++++++-- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 6cb0a2eeb4e..1af1e40d768 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -565,8 +565,16 @@ data_plane: local_buffer_size: 4294967296 # 4 GiB/process reuse_registered_buffers: true # reuse RDMA-registered buffers staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool - use_gdr: false # GPU-memory RDMA staging in CUDA clients - gdr_staging_buffer_mb: 1024 # persistent MiB per active GDR client + use_gdr: false # GPU-memory RDMA staging in CUDA clients + gdr_staging_buffer_mb: 1024 # HBM reserved per active GDR client; + # 8 clients/node is ~8 GiB off the training + # budget. Must exceed the largest per-sample, + # per-field payload: below that, GDR reads of + # CPU-written values fail inside TransferQueue. + # Lower values split fetches into more + # serialized groups and cost throughput; + # higher buys little at linear HBM cost. + # GDR needs that headroom to pay off. # observability: # NotRequired # enabled: false diff --git a/tests/unit/reference_configs/grpo_math_1B.yaml b/tests/unit/reference_configs/grpo_math_1B.yaml index a5a0d134942..d2670e97f93 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -539,8 +539,16 @@ data_plane: local_buffer_size: 4294967296 # 4 GiB/process reuse_registered_buffers: true # reuse RDMA-registered buffers staging_buffer_size: 268435456 # 256 MiB/pool slot; bigger transfers bypass the pool - use_gdr: false # GPU-memory RDMA staging in CUDA clients - gdr_staging_buffer_mb: 1024 # persistent MiB per active GDR client + use_gdr: false # GPU-memory RDMA staging in CUDA clients + gdr_staging_buffer_mb: 1024 # HBM reserved per active GDR client; + # 8 clients/node is ~8 GiB off the training + # budget. Must exceed the largest per-sample, + # per-field payload: below that, GDR reads of + # CPU-written values fail inside TransferQueue. + # Lower values split fetches into more + # serialized groups and cost throughput; + # higher buys little at linear HBM cost. + # GDR needs that headroom to pay off. # observability: # NotRequired # enabled: false From 2b2eec292da382a9b5ab5920861fb8561a25d295 Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 8 Sep 2026 21:11:25 +0000 Subject: [PATCH 6/8] refactor(data-plane): drop the flat use_gdr migration error MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The flat `data_plane.use_gdr` spelling never shipped — `git log -S"use_gdr"` at this branch's base returns no commits, and it appears nowhere on main — so no merged config could ever trip this check. Making it fire also required declaring `use_gdr` and `gdr_staging_buffer_mb` as `NotRequired` on `DataPlaneConfig`, which advertises them to any reader as legal top-level keys: the opposite of the intent. It also left two policies for one class of mistake three lines apart, since the pre-nesting flat sizing keys are silently ignored. `backend_config` now has one rule: only the nested block is read. Removes the two declarations, `_FIRST_GDR_FLAT_KEYS`, the raise, the two docstring passages, the README sentence, and the two tests that pinned it. Signed-off-by: Zeyu Zhou --- nemo_rl/data_plane/README.md | 4 +--- nemo_rl/data_plane/interfaces.py | 21 ++---------------- tests/unit/data_plane/test_backend_config.py | 23 -------------------- 3 files changed, 3 insertions(+), 45 deletions(-) diff --git a/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index e665f928b62..fb70eb79c5d 100644 --- a/nemo_rl/data_plane/README.md +++ b/nemo_rl/data_plane/README.md @@ -447,9 +447,7 @@ These keys used to sit directly under `data_plane:`. That spelling is not rejected — it is simply never read. A config still using it silently gets this backend's defaults instead of its own values: an inherited config supplies the nested block, so a surviving flat key always loses the merge, -with no warning either way. Earlier revisions of PR #3501 also used flat -`use_gdr` and `gdr_staging_buffer_mb` keys; those two spellings are rejected -with a migration error because silently disabling GDR would invalidate a run. +with no warning either way. Backend choice: - **`simple`** — ZMQ-backed; lowest setup overhead. Default for tests diff --git a/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 7bbff650f01..a50478295b7 100644 --- a/nemo_rl/data_plane/interfaces.py +++ b/nemo_rl/data_plane/interfaces.py @@ -124,9 +124,7 @@ class DataPlaneConfig(TypedDict): ``local_buffer_size`` used to sit at this level. A config still using that spelling is not rejected — the flat key is simply never read, and :func:`backend_config` resolves the nested block (or its defaults) as if it - were absent. Earlier revisions of PR #3501 also put ``use_gdr`` and - ``gdr_staging_buffer_mb`` here; those two spellings are rejected explicitly - so a stale GDR recipe cannot silently run with GDR disabled. See there. + were absent. See there. """ enabled: bool @@ -139,14 +137,6 @@ class DataPlaneConfig(TypedDict): ack_timeout_ms: NotRequired[int] observability: NotRequired["ObservabilityConfig"] - # Declared only so Pydantic preserves the first GDR integration's flat - # spellings until backend_config can raise the migration error below. - use_gdr: NotRequired[bool] - gdr_staging_buffer_mb: NotRequired[int] - - -_FIRST_GDR_FLAT_KEYS = ("use_gdr", "gdr_staging_buffer_mb") - _CHECKPOINTABLE_BACKENDS: frozenset[str] = frozenset({"simple"}) @@ -175,16 +165,9 @@ def backend_config(cfg: DataPlaneConfig) -> Any: (block already coerced to a model) or as a plain dict from a test. Sizing is read only from the nested block. A config still using the - pre-nesting flat sizing spelling gets this backend's defaults, not its own - values. The first GDR integration's flat spelling raises instead because - silently disabling a requested transport would invalidate a benchmark. + pre-nesting flat spelling gets this backend's defaults, not its own values. """ backend = cfg["backend"] - stale_gdr = [key for key in _FIRST_GDR_FLAT_KEYS if key in cfg] - if stale_gdr: - keys = ",".join(stale_gdr) - raise ValueError(f"data_plane.{{{keys}}} moved under data_plane.mooncake_cpu") - nested = cfg.get(backend) or {} if isinstance(nested, BaseModel): nested = nested.model_dump(exclude_unset=True) diff --git a/tests/unit/data_plane/test_backend_config.py b/tests/unit/data_plane/test_backend_config.py index 999ecb2cfcf..0e6e20f5eda 100644 --- a/tests/unit/data_plane/test_backend_config.py +++ b/tests/unit/data_plane/test_backend_config.py @@ -100,13 +100,6 @@ def test_gdr_staging_size_rejects_non_positive_values(value: int) -> None: ) -@pytest.mark.parametrize("backend", ["simple", "mooncake_cpu"]) -@pytest.mark.parametrize("key", ["use_gdr", "gdr_staging_buffer_mb"]) -def test_flat_gdr_key_from_first_integration_raises(backend: str, key: str) -> None: - with pytest.raises(ValueError, match="moved under data_plane.mooncake_cpu"): - backend_config(_cfg(backend, **{key: True})) - - def test_accepts_an_already_coerced_model() -> None: """Configs arriving via pydantic have the block coerced to a model already.""" cfg = _cfg("mooncake_cpu", mooncake_cpu=MooncakeCpuConfig(global_segment_size=555)) @@ -157,19 +150,3 @@ def test_schema_validates_without_any_backend_block() -> None: validate, otherwise MasterConfig fails before training starts. """ TypeAdapter(DataPlaneConfig).validate_python(_cfg("simple")) - - -@pytest.mark.parametrize( - ("key", "value"), - [("use_gdr", True), ("gdr_staging_buffer_mb", 512)], -) -def test_pydantic_preserves_flat_gdr_keys_for_migration_error( - key: str, value: object -) -> None: - validated = TypeAdapter(DataPlaneConfig).validate_python( - _cfg("mooncake_cpu", **{key: value}) - ) - - assert key in validated - with pytest.raises(ValueError, match="moved under data_plane.mooncake_cpu"): - backend_config(validated) From 9896a67bbc6a65f52ed71bcad93e783680ba3a3b Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Tue, 8 Sep 2026 21:15:23 +0000 Subject: [PATCH 7/8] fix(data-plane): skip the GDR CUDA guard for the process-local config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `setup_data_plane` takes `DataPlaneRuntimeConfig`, which is `DataPlaneConfig | LocalDataPlaneConfig`. The guard subscripted `cfg` directly, so a `LocalDataPlaneConfig` — a BaseModel, not a mapping — would raise `TypeError` before ever reaching `build_data_plane_client`. The guard was written when `setup_data_plane` still took the narrower `DataPlaneConfig`; main widened the parameter, and the two changes never touched the same lines, so the rebase merged cleanly and broke the semantics. `pyrefly` caught it: Argument `LocalDataPlaneConfig | TypedDict[DataPlaneConfig]` is not assignable to parameter `cfg` with type `TypedDict[DataPlaneConfig]` Narrow with `isinstance` first, matching how `factory.py` handles the same union. The local plane has no TQ and no mooncake, so it has no GDR to order against a CUDA context. Signed-off-by: Zeyu Zhou --- nemo_rl/data_plane/worker_mixin.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index f128f24170d..eae8393be6b 100644 --- a/nemo_rl/data_plane/worker_mixin.py +++ b/nemo_rl/data_plane/worker_mixin.py @@ -38,7 +38,7 @@ from nemo_rl.data.llm_message_utils import attach_message_log_view from nemo_rl.data.multimodal_utils import PackedTensor -from nemo_rl.data_plane.interfaces import backend_config +from nemo_rl.data_plane.interfaces import LocalDataPlaneConfig, backend_config from nemo_rl.data_plane.schema import ( ELEM_COUNTS_PER_GB, GLOBAL_FORWARD_PAD_SEQLEN, @@ -328,8 +328,11 @@ def setup_data_plane(self, cfg: DataPlaneRuntimeConfig) -> None: self._route_fallback_counts = Counter() from nemo_rl.data_plane import build_data_plane_client + # ``LocalDataPlaneConfig`` is the process-local plane: no TQ, no + # mooncake, so no GDR to order against a CUDA context. if ( - cfg["backend"] == "mooncake_cpu" + not isinstance(cfg, LocalDataPlaneConfig) + and cfg["backend"] == "mooncake_cpu" and backend_config(cfg).use_gdr and not torch.cuda.is_initialized() ): From 2b15799829e6aeea3950a977747e847dcfc72f4d Mon Sep 17 00:00:00 2001 From: Zeyu Zhou Date: Wed, 9 Sep 2026 06:15:23 +0000 Subject: [PATCH 8/8] fix(data-plane): default the GDR flags at class level `put_samples` reads `self._gdr_requested` / `self._gdr_put_confirmed`, which are the first instance-attribute reads on that path. Any client built without running `__init__` therefore raises: AttributeError: 'TQDataPlaneClient' object has no attribute '_gdr_requested' `test_codec_mooncake.py::test_put_samples_passes_fields_and_tags_through_unchanged` constructs exactly that with `object.__new__`, and the module already notes that a process can unpickle a client without running the constructor. Declaring both as class-level defaults keeps the read safe on those paths without touching a test that is about codec passthrough, not GDR. Signed-off-by: Zeyu Zhou --- nemo_rl/data_plane/adapters/transfer_queue.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/nemo_rl/data_plane/adapters/transfer_queue.py b/nemo_rl/data_plane/adapters/transfer_queue.py index c4a9f644ab3..692f6d6daaf 100644 --- a/nemo_rl/data_plane/adapters/transfer_queue.py +++ b/nemo_rl/data_plane/adapters/transfer_queue.py @@ -776,6 +776,12 @@ def _from_wire(td: TensorDict) -> TensorDict: class TQDataPlaneClient(DataPlaneClient): """Adapter façade — maps NeMo-RL calls onto TransferQueue's public API.""" + # Class-level so ``put_samples`` stays readable on an instance built + # without ``__init__`` — ``object.__new__`` in tests, or a process that + # unpickles a client without running the constructor. + _gdr_requested: bool = False + _gdr_put_confirmed: bool = False + def __init__(self, cfg: DataPlaneConfig, *, bootstrap: bool = True) -> None: """Construct a TQ-backed client.