Skip to content

feat(data_plane): mooncake cpu rdma backend - #2935

Merged
ZhiyuLi-Nvidia merged 72 commits into
mainfrom
zhiyul/tq_mooncake_cpu_rdma
Aug 23, 2026
Merged

feat(data_plane): mooncake cpu rdma backend#2935
ZhiyuLi-Nvidia merged 72 commits into
mainfrom
zhiyul/tq_mooncake_cpu_rdma

Conversation

@ZhiyuLi-Nvidia

@ZhiyuLi-Nvidia ZhiyuLi-Nvidia commented Jun 25, 2026

Copy link
Copy Markdown
Contributor

What does this PR do ?

Makes data_plane.backend=mooncake_cpu use CPU RDMA, and only CPU RDMA.

RDMA is now the transport, not an option. _mooncake_transport_config() auto-detects a RoCE-capable mlx5 device and passes it to MooncakeStore. There is no TCP mode: MC_MOONCAKE_PROTOCOL is gone, and a host without a suitable device raises rather than degrading. A silent TCP fallback only ever served to let a run report mooncake/RDMA coverage it had not actually exercised.

This is an enabling fix, not just a default flip. The device probe on main used two dirnames where it needed three, so basename $(dirname $(dirname .../mlx5_3/ports/1/link_layer)) resolved to the literal string ports — not a device. Setting MC_MOONCAKE_PROTOCOL=rdma on main therefore produced device_name="ports" and could never have worked. The probe is now glob + Path.parents[2], which removes the shell entirely along with that class of bug.

Sizing defaults corrected: 512 + 64 GiB → 64 + 4 GiB per client process. global_segment_size and local_buffer_size are registered per client process — one per GPU — so a node pays gpus_per_node x (segment + buffer). Under RDMA that memory is pinned via ibv_reg_mr and resident from setup, not lazily faulted in, so the old defaults meant 4.6 TiB/node on an 8-GPU host. That was survivable only while the transport was TCP and the mapping lazy; making RDMA the default converted a virtual reservation into a real allocation. Upstream TransferQueue defaults to 4 + 1 GiB per client.

Deletes the TQ actor runtime_env injection. _patch_tq_actor_runtime_env monkey-patched TransferQueue's actor classes to pip-install TQ from a git ref at every actor spawn, putting a network call on the actor-creation path. Its own TODO scoped it to removal once the nightly image shipped TQ via pyproject.toml, which it now does. Verified at 256 actors: no ModuleNotFoundError.

Tests skip instead of falling back. The mooncake unit fixture and grpo_dp_mooncake.sh share the adapter's roce_device() predicate and skip when no RoCE device is present, so they exercise the real transport where hardware allows and are honestly absent where it does not. InfiniBand is not auto-selected — RoCE is the fabric this is validated on — but MC_MOONCAKE_DEVICE names any device explicitly.

Verification

Both runs below are backend=mooncake_cpu over RDMA (mlx5_3), with falling back to TCP at zero and RDMA bring-up confirmed on driver and workers across nodes.

Run Config Result W&B
llama-3.1-8B 2n8g, Megatron FP8, 10 steps 10/10 steps, reward 0.343 → 0.437 xdsn2kek
DeepSeek-V3 32n8g (256 GPUs), Megatron, 10 steps 10/10 steps, check_metrics.py 3/3 PASS cpqmwppy

DeepSeek-V3 metric checks: min(token_mult_prob_error) 1.021 < 1.05, max(reward) 0.623 > 0.4, mean(total_step_time) 372 s < 1000. Steady-state ~270–400 s/step; grad_norm 0.053–0.084 throughout, so the step count is not a collapsed no-op.

Controls. Both recipes were also run with data_plane.backend=simple to isolate the data plane. At 2 nodes the backends are numerically indistinguishable (reward agrees step-for-step to ~0.015). On DeepSeek-V3, train/token_mult_prob_error shows occasional large excursions — 37.99 on mooncake, and 2039 on simple, which never touches mooncake or RDMA. The excursions are therefore a property of DeepSeek-V3 on this recipe (MoE routing with use_importance_sampling_correction=True), not of the data plane.

Worth noting separately: the recipe asserts min(token_mult_prob_error) < 1.05, which passes on the healthiest step and cannot detect those excursions on either backend.

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

  • The TransferQueue v0.1.9 bump is already on main (chore: bump TQ to v0.19 #3423); this PR carries no dependency change.
  • MC_GID_INDEX is still pinned to 3 on the driver. Logs show workers never receive it and mooncake's own findBestGidIndex selects 3 independently, so it is effectively inert — left alone rather than removed on inference.

@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia requested a review from a team as a code owner June 25, 2026 16:46
@copy-pr-bot

copy-pr-bot Bot commented Jun 25, 2026

Copy link
Copy Markdown

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.

@ZhiyuLi-Nvidia ZhiyuLi-Nvidia changed the title feat(data_plane): bump TransferQueue to v0.1.8 with RDMA-default moon… feat(data_plane): bump TransferQueue to v0.1.8 for mooncake cpu rdma backend Jun 25, 2026
@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test b2b2e2a

Comment thread nemo_rl/data_plane/observability.py Outdated
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/tq_mooncake_cpu_rdma branch from fdbb9ac to 28e8184 Compare August 7, 2026 23:02
@ZhiyuLi-Nvidia ZhiyuLi-Nvidia changed the title feat(data_plane): bump TransferQueue to v0.1.8 for mooncake cpu rdma backend feat(data_plane): mooncake cpu rdma backend Aug 7, 2026
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/tq_mooncake_cpu_rdma branch from 28e8184 to 79d1e34 Compare August 11, 2026 19:02
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia requested a review from a team as a code owner August 11, 2026 19:02
@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test e1cccea

@ZhiyuLi-Nvidia ZhiyuLi-Nvidia added the CI:L1 Run doctests, unit tests, and functional tests label Aug 11, 2026
@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test e1cccea

@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/tq_mooncake_cpu_rdma branch from e1cccea to 20e2171 Compare August 12, 2026 00:28
@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test 20e2171

@github-actions github-actions Bot added the CI Relating to CI label Aug 13, 2026
@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test 43a5ea0

@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/tq_mooncake_cpu_rdma branch from 4fb9e5c to 76c39a4 Compare August 14, 2026 10:52
@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test 76c39a4

@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test 9a54ad0

@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test 9183b4b

@zyzhou5 zyzhou5 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Overall LGTM. Left some comments.

Comment thread nemo_rl/data_plane/adapters/transfer_queue.py Outdated
Comment thread nemo_rl/data_plane/adapters/transfer_queue.py Outdated
@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test 46de7a2

@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/tq_mooncake_cpu_rdma branch from adb6721 to 250c9c4 Compare August 16, 2026 07:26
ZhiyuLi-Nvidia and others added 23 commits August 22, 2026 23:25
…h test's payload

a19f7f2 fixed test_growing_a_slot_unregisters_before_dropping_the_old_buffer
by shrinking its "big" buffer from 8 MiB to 8 KiB, but the real regression
is in the ceiling, not the payload: before b61f58e threaded max_bytes
through _StagingPool's constructor, the module-level _STAGING_MAX_BYTES
default was 256 MiB, comfortably above the 8 MiB this test was written
to exercise. b61f58e replaced that with a local _MAX = 1 << 20 (1 MiB)
without raising it to cover the existing 8 MiB payload, so growth started
silently taking the outlier-bypass branch instead.

Restore the test's original 8 MiB payload and instead bump _MAX to
1 << 24 (16 MiB), restoring the "comfortably above every payload these
tests stage" invariant the constant's own comment already promises. The
other three tests that share _MAX only stage 1024-byte payloads, so this
is a no-op for them.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
run_grpo_single_controller.py now calls maybe_configure_data_plane_env
(config.data_plane) before init_ray(), but the main_context fixture's
data_plane={"enabled": True} has no impl/backend, so cfg["impl"] raises
KeyError the moment the real (unmocked) function runs.

The script itself requires data_plane.enabled=true, so the fixture can't
just disable it — give it the minimal complete shape instead:
impl="transfer_queue", backend="simple" (simple has no mooncake engine
env to configure, so maybe_configure_engine_env stays a no-op).

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
gb200 CI runners have 2 RoCE rails per NUMA node, and
MC_ENABLE_DEST_DEVICE_AFFINITY has no measurable effect in the pinned
TransferQueue/mooncake wheels (see 88181da's measurement) — verified
against the exact pinned commits (TransferQueue@c516143, mooncake
v0.3.11.post1): selectDevice()'s hint is a soft preference that silently
falls back to random peer-rail selection whenever the peer's topology
doesn't have an HCA registered under that exact name for the buffer's
location, and neither MC_ENABLE_HCA_PEER_AFFINITY nor MC_NIC_PEER_AFFINITY
exist in this mooncake version at all.

So a same-NUMA cross-rail draw is a live possibility on any transfer,
independent of what we set. Retries do help — a failing gb200 functional
test run showed the failing key set shrink each round (18 -> 12 -> 6)
rather than get stuck on one bad pairing — but TransferQueue's mooncake
client hardcodes 3 attempts with no backoff growth (MAX_RETRIES,
RETRY_DELAY_SECONDS=1.0), too few to reliably clear the residual failures
this topology produces.

Add MooncakeCpuConfig.transfer_max_retries (default 3, matching upstream,
so this is a no-op for every other config) and a small monkeypatch —
MAX_RETRIES is a bare module constant read fresh from the module
namespace on every call, so reassigning it before any transfer runs
changes every subsequent retry loop process-wide. Bump it to 10 in
grpo_dp_mooncake.sh specifically, where the topology is known.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
…nfigurable"

This reverts commit bb345b7.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
gb200 nodes have 2 RoCE rails per NUMA domain. MC_ENABLE_DEST_DEVICE_AFFINITY's
same-name peer hint (Topology::getHcaIndex in the pinned mooncake wheel)
silently falls back to random peer selection whenever the lookup misses, and
two same-domain rails give that fallback something to be ambiguous about.
Measured on the real fleet: every cross-rail pair among 4 same-node rails
(mlx5_0..3) failed with "transport retry counter exceeded" (each rail is its
own subnet), and no same-rail pair ever did.

Retrying more doesn't fix this — it just re-rolls the same ambiguous draw.
The actual fix is removing the ambiguity at the source: rdma_devices() now
keeps only the first rail per NUMA domain, read from
/sys/class/infiniband/<dev>/device/numa_node. Every domain still gets a
dedicated rail, so a job spanning multiple domains still gets multi-rail
bandwidth — only genuinely redundant extra rails on an already-represented
domain are dropped. Missing/"-1" NUMA info is never treated as a shared
domain across devices (that would collapse every rail on the host to one),
so it falls back to one entry per device, matching prior behavior.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Default stays on (matches the measured-necessary behavior on the pinned
mooncake wheel), but a fabric/wheel where MC_ENABLE_DEST_DEVICE_AFFINITY's
same-name peer hint actually resolves reliably would get more aggregate
bandwidth per domain from every rail, not just one -- so make it an
explicit opt-out rather than a hardcoded always-on.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Per external review: the measured cross-rail-ambiguity failure this dedup
guards against has only ever been observed on the gb200 runners' RoCE
fabric. Applying the same default-on dedup to InfiniBand would silently
undo this PR's own "use all IB rails" intent (see
11cc549/9b5c604bc) without any measured justification -- IB may not
even share the same routing failure mode RoCE has (a rail-per-subnet
model), and no IB fleet has shown the ambiguity actually causing failures.

InfiniBand rails are now always listed in full regardless of
dedupe_rails_per_numa_domain; the config flag and NUMA-domain grouping
only ever apply to the RoCE fallback path.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
…laim

Verified against the pinned wheel: mooncake's own parsing is
if (std::getenv("MC_ENABLE_DEST_DEVICE_AFFINITY")) -- presence, not
value -- so setting it to "0" does not disable it at the engine level.
The existing test only pinned NeMo-RL's own os.environ.setdefault
behavior (correctly not clobbering an operator-set value) but its
docstring claimed this let a fabric "opt back into random peer
selection", which isn't true of the real engine. Renamed and corrected
the test, and added the same note where the env var is actually set.

There is currently no supported way to disable this knob once set --
rdma_devices()'s per-domain dedup exists because this knob cannot be
turned off or relied on, not despite it being configurable.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Per external review's pre-merge checklist. Getting this exact
kept-rail/dropped-rail/domain mapping off a real gb200 runner was the
single hardest part of this whole investigation -- every attempt this
session to retrieve it live from a running Slurm job failed for
infrastructure reasons unrelated to the actual bug. Logging it at setup
time means the next time this needs debugging, it is already in the job
log instead of requiring another live-repro attempt.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
…oce_rails_per_numa_domain

The option only ever dedupes RoCE rails (InfiniBand is never deduped),
so the old name was ambiguous about scope. No behavior change.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
… init_ray comment

Extract the mlx5/uverbs sysfs check duplicated across
run_functional_in_docker.sh, run_unit_in_docker.sh, and
grpo_dp_mooncake.sh into a shared tests/scripts/detect_rdma.sh
helper. Also collapse the 3-line rationale comment repeated at
every maybe_configure_data_plane_env() call site down to a one-line
pointer, since the rationale already lives in that function's
docstring.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
…th_1B exemplar/reference configs

The dedupe_rails_per_numa_domain -> dedupe_roce_rails_per_numa_domain
rename added a new MooncakeCpuConfig field without reflecting it in
the exemplar or its test_config_v2 reference counterpart, so CI's
v1/v2 config parity check flagged it as a Pydantic default missing
from the YAML.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
…rdma.sh

The 4th and last copy of the mlx5/uverbs detection predicate lived
inline in test-template/action.yml. Source the same
tests/scripts/detect_rdma.sh the other three call sites use, via the
checkout path this step already relies on for its Docker volume mount
(github.run_id + steps.uuid.outputs.id), so all four consumers now
share one canonical check.

Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
…l dedup

Mooncake reads every MC_* var once, in the constructor of a function-local
static singleton built the first time engine code runs in a process
(mooncake-common/src/environ.cpp, Environ::Get). Setting
MC_ENABLE_DEST_DEVICE_AFFINITY from Python therefore lands in os.environ but
never reaches the engine, so mooncake kept choosing each transfer's peer rail
independently. Every RoCE rail on the gb200 runners is its own /32 subnet, so a
cross-rail pair has no route and dies with "transport retry counter exceeded" —
the gb200_L1_Functional_Tests_GRPO_1 failure.

Export it from the container and launcher instead: image-level ENV in both
Dockerfiles, an export in ray.sub, and --env on the CI docker run. That matches
how RAY_ENABLE_UV_RUN_RUNTIME_ENV is already handled in those same three places,
and covers processes NeMo RL never configures itself — TransferQueue's storage
actors and the mooncake_master subprocess inherit the image environment but do
not run TQDataPlaneClient.__init__.

With the hint actually reaching the engine, per-NUMA-domain rail dedup is
unnecessary. It was compensating for the missing hint and did not work anyway:
the pair it leaves behind (mlx5_0 + mlx5_2) is itself cross-rail, and CI failed
with dedup enabled. Remove the config field and offer every rail.

Also correct maybe_configure_data_plane_env's docstring, which claimed
init_ray's env snapshot made a setting cluster-wide and that an engine reading
a knob at startup "cannot be configured any later than this". It covers Ray
workers, not a backend's own subprocesses.

Measured on a gb200 CI-fleet node in CI's exact shape (4 GPUs visible, test
takes indices 0,1 = NUMA 0, the placement CI always draws):

  before: 76x mlx5_2 -> mlx5_0, TRANSFER_FAIL, exit 1
  after:  mlx5_0,mlx5_1,mlx5_2,mlx5_3 - 0 TRANSFER_FAIL, 0 retries, exit 0

Env placement was the only variable: 4/4 runs clean with it in the container
env, 0/3 with it set only from Python, across both 2-rail and 4-rail lists.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
…them

Mooncake reads its whole MC_* configuration once, as its extension loads, so
a write to os.environ afterwards lands in Python's environment while the
engine keeps the value it captured. The failure is silent: os.environ reads
back correctly, and on a rail-isolated RoCE fabric every transfer still dies
with "transport retry counter exceeded".

Move the configuration into nemo_rl/data_plane/adapters/transfer_queue_env.py,
a module that deliberately imports neither transfer_queue nor mooncake, so
reaching it can never be the thing that loads the engine. It is called once on
the driver by maybe_configure_data_plane_env, before init_ray, and Ray hands
the resulting environment to every worker. Calling it too late now raises
rather than silently no-op'ing.

Gate MC_ENABLE_DEST_DEVICE_AFFINITY on the host being RoCE-only, where each
rail is its own subnet so a cross-rail pair has no route. InfiniBand routes
cross-rail and is left alone. This replaces the unconditional image-level ENV
in both Dockerfiles and the export in ray.sub.

Also drop the now-unused pathlib import in the adapter, add the new module to
pyrefly.toml, and fix test_virtual_cluster's env-threading test, which asserted
the affinity knob unconditionally and would otherwise depend on the fabric of
whatever host ran it.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
The lazy per-client staging pool was guarded by a module-level
threading.Lock and stored by injecting client._nrl_staging onto an
upstream TransferQueue object, read back through getattr. That is
module-level mutable state plus the side-channel/out-of-band state and
reflection the style guide tells us to avoid.

Replace both with a _StagingPoolRegistry that owns its lock and a
weakref.WeakKeyDictionary, instantiated inside
_patch_mooncake_staging_buffers and captured by the two patched thread
workers. The lock and the table are now scoped to the patch rather than
to the module, and nothing is written onto the upstream client.

Behaviour is unchanged: the pool is still built at most once per client
across threads, and the lock is still taken on every lookup rather than
double-checked.

Verified against the pinned TransferQueue revision that
MooncakeStoreClient derives from a plain ABC with no __slots__,
__hash__, or __eq__, so it is hashable and weakref-able. Confirmed by
experiment that the registry constructs exactly one pool under 8-thread
contention and that dropping a client drains its entry, so pools do not
outlive their client.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Co-authored-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
…t keys

The mooncake monkey-patch shape guards silently returned when the pinned
transfer_queue revision no longer matched, which would drop RDMA buffer
reuse and unregistered-buffer diagnosability with zero symptom besides
lost throughput. Raise instead, sharing one error-message builder across
the three guard sites.

tests/functional/grpo_sc_gym_router_failover.sh still passed the old flat
data_plane.* keys, which fails validation now that sizing lives under a
backend-named block -- the only remaining flat-key call site in the repo.

Also tightens two docstrings (transfer_queue_env.py, interfaces.py) that
overstated what's overridable / optional.

Co-authored-by: Terry Kong <terryk@nvidia.com>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia force-pushed the zhiyul/tq_mooncake_cpu_rdma branch from 3e4e076 to 2def95b Compare August 23, 2026 06:37
@ZhiyuLi-Nvidia

Copy link
Copy Markdown
Contributor Author

/ok to test 2def95b

@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia enabled auto-merge (squash) August 23, 2026 06:56
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia merged commit ffbf33f into main Aug 23, 2026
259 of 264 checks passed
@ZhiyuLi-Nvidia
ZhiyuLi-Nvidia deleted the zhiyul/tq_mooncake_cpu_rdma branch August 23, 2026 20:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI:L1 Run doctests, unit tests, and functional tests CI Relating to CI

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants