feat(data_plane): mooncake cpu rdma backend - #2935
Merged
Merged
Conversation
Contributor
Author
|
/ok to test b2b2e2a |
ZhiyuLi-Nvidia
commented
Jul 11, 2026
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_mooncake_cpu_rdma
branch
from
August 7, 2026 23:02
fdbb9ac to
28e8184
Compare
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_mooncake_cpu_rdma
branch
from
August 11, 2026 19:02
28e8184 to
79d1e34
Compare
Contributor
Author
|
/ok to test e1cccea |
Contributor
Author
|
/ok to test e1cccea |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_mooncake_cpu_rdma
branch
from
August 12, 2026 00:28
e1cccea to
20e2171
Compare
Contributor
Author
|
/ok to test 20e2171 |
Contributor
Author
|
/ok to test 43a5ea0 |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_mooncake_cpu_rdma
branch
from
August 14, 2026 10:52
4fb9e5c to
76c39a4
Compare
Contributor
Author
|
/ok to test 76c39a4 |
Contributor
Author
|
/ok to test 9a54ad0 |
Contributor
Author
|
/ok to test 9183b4b |
zyzhou5
reviewed
Aug 14, 2026
zyzhou5
left a comment
Contributor
There was a problem hiding this comment.
Overall LGTM. Left some comments.
Contributor
Author
|
/ok to test 46de7a2 |
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_mooncake_cpu_rdma
branch
from
August 16, 2026 07:26
adb6721 to
250c9c4
Compare
…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>
Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
ZhiyuLi-Nvidia
force-pushed
the
zhiyul/tq_mooncake_cpu_rdma
branch
from
August 23, 2026 06:37
3e4e076 to
2def95b
Compare
Contributor
Author
|
/ok to test 2def95b |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Makes
data_plane.backend=mooncake_cpuuse CPU RDMA, and only CPU RDMA.RDMA is now the transport, not an option.
_mooncake_transport_config()auto-detects a RoCE-capablemlx5device and passes it to MooncakeStore. There is no TCP mode:MC_MOONCAKE_PROTOCOLis 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
mainused twodirnames where it needed three, sobasename $(dirname $(dirname .../mlx5_3/ports/1/link_layer))resolved to the literal stringports— not a device. SettingMC_MOONCAKE_PROTOCOL=rdmaonmaintherefore produceddevice_name="ports"and could never have worked. The probe is nowglob+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_sizeandlocal_buffer_sizeare registered per client process — one per GPU — so a node paysgpus_per_node x (segment + buffer). Under RDMA that memory is pinned viaibv_reg_mrand 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_envinjection._patch_tq_actor_runtime_envmonkey-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 viapyproject.toml, which it now does. Verified at 256 actors: noModuleNotFoundError.Tests skip instead of falling back. The mooncake unit fixture and
grpo_dp_mooncake.shshare the adapter'sroce_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 — butMC_MOONCAKE_DEVICEnames any device explicitly.Verification
Both runs below are
backend=mooncake_cpuover RDMA (mlx5_3), withfalling back to TCPat zero and RDMA bring-up confirmed on driver and workers across nodes.check_metrics.py3/3 PASSDeepSeek-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_norm0.053–0.084 throughout, so the step count is not a collapsed no-op.Controls. Both recipes were also run with
data_plane.backend=simpleto 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_errorshows occasional large excursions — 37.99 on mooncake, and 2039 onsimple, which never touches mooncake or RDMA. The excursions are therefore a property of DeepSeek-V3 on this recipe (MoE routing withuse_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:
Additional Information
main(chore: bump TQ to v0.19 #3423); this PR carries no dependency change.MC_GID_INDEXis still pinned to3on the driver. Logs show workers never receive it and mooncake's ownfindBestGidIndexselects 3 independently, so it is effectively inert — left alone rather than removed on inference.