diff --git a/examples/configs/grpo_math_1B.yaml b/examples/configs/grpo_math_1B.yaml index 99d9ee369cc..1af1e40d768 100644 --- a/examples/configs/grpo_math_1B.yaml +++ b/examples/configs/grpo_math_1B.yaml @@ -565,6 +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 # 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/nemo_rl/data_plane/README.md b/nemo_rl/data_plane/README.md index 44d20e97d8e..fb70eb79c5d 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 ``` @@ -450,8 +452,22 @@ with no warning either way. 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..692f6d6daaf 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. # ────────────────────────────────────────────────────────────────────────── @@ -684,6 +687,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), }, }, } @@ -771,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. @@ -817,6 +828,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: @@ -1031,6 +1049,31 @@ 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() + ) + ) + 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) + 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( @@ -1039,6 +1082,11 @@ def put_samples( fields=wire_fields, tags=user_tags, ) + if confirm_gdr_put: + 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/nemo_rl/data_plane/interfaces.py b/nemo_rl/data_plane/interfaces.py index 9f527c4b2c8..a50478295b7 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): diff --git a/nemo_rl/data_plane/worker_mixin.py b/nemo_rl/data_plane/worker_mixin.py index a5309d453f8..eae8393be6b 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 LocalDataPlaneConfig, 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,18 @@ 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 ( + not isinstance(cfg, LocalDataPlaneConfig) + and 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..719104307b4 100644 --- a/tests/unit/data_plane/README.md +++ b/tests/unit/data_plane/README.md @@ -107,6 +107,13 @@ 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` (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) - `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..0e6e20f5eda 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,16 @@ 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}) + ) def test_accepts_an_already_coerced_model() -> None: 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..f8673f3bf55 --- /dev/null +++ b/tests/unit/data_plane/test_mooncake_gdr.py @@ -0,0 +1,164 @@ +# 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 logging +import sys +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 +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_tensor_put_is_confirmed_once_and_never_falls_back( + monkeypatch, caplog +) -> None: + """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) + ) + 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 + + 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 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: + """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..d2670e97f93 100644 --- a/tests/unit/reference_configs/grpo_math_1B.yaml +++ b/tests/unit/reference_configs/grpo_math_1B.yaml @@ -539,6 +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 # 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