diff --git a/.github/workflows/collectivex-sweep.yml b/.github/workflows/collectivex-sweep.yml index 9a9d9b33fb..0b754c0fbb 100644 --- a/.github/workflows/collectivex-sweep.yml +++ b/.github/workflows/collectivex-sweep.yml @@ -10,7 +10,7 @@ on: description: "EP library to sweep — 'all' runs every EP backend in one matrix" type: choice default: all - options: [all, deepep-v2, mori, uccl-ep] + options: [all, deepep-v2, mori, uccl-ep, nccl-ep] only_sku: description: Restrict to one GHA runner pool; blank = all type: string diff --git a/experimental/CollectiveX/README.md b/experimental/CollectiveX/README.md index ea87b89e89..1e1f0ce521 100644 --- a/experimental/CollectiveX/README.md +++ b/experimental/CollectiveX/README.md @@ -13,9 +13,11 @@ responsibility. The full measurement methodology is in [docs/methodology.md](doc The workload uses packed placement and one pinned `fixed-profile` resource configuration per backend/topology; there is no tuning sweep. Combine is always BF16; dispatch precision is a swept -dimension — a BF16 control plus an FP8 dispatch on every backend, caller-prequantized in `normal` -mode (in `low-latency` the DeepEP and UCCL-EP kernels quantize internally from BF16; MoRI stays -caller-prequantized). Coverage is uniform routing only. Cases run in one of two modes: +dimension — a BF16 control plus, on every backend whose FP8 dispatch is supported upstream +(DeepEP V2, MoRI, UCCL-EP), an FP8 dispatch, caller-prequantized in `normal` mode (in `low-latency` +the DeepEP and UCCL-EP kernels quantize internally from BF16; MoRI stays caller-prequantized). NCCL +EP is BF16-only this release, so it emits the control alone. Coverage is uniform routing only. Cases +run in one of two modes: - `normal` uses `layout-and-dispatch-v1`, rank-deduplicated token payloads, and activation-only, unweighted rank-sum combine. It runs the full decode and prefill ladders. @@ -27,7 +29,9 @@ caller-prequantized). Coverage is uniform routing only. Cases run in one of two decode-phase-only, per-SKU-capability-gated addition whose runnable set differs from `normal`'s, so it is enabled from each SKU's `ll_backends` registry entry (currently DeepEP V2 EP8 on H100/H200/B200, MoRI EP8 on MI300X/MI325X/MI355X, and UCCL-EP EP8 on H100/H200/B200 only — UCCL's low-latency kernel - trips a warp-group assertion on AMD's CU count, so the AMD SKUs keep UCCL-EP normal mode without LL). + trips a warp-group assertion on AMD's CU count, so the AMD SKUs keep UCCL-EP normal mode without LL; + NCCL EP has no low-latency row on any SKU while its decode kernels carry + [NVIDIA/nccl#2303](https://github.com/NVIDIA/nccl/issues/2303)). Scoped single-node EP8 runs over the intra-node NVLink/XGMI low-latency path (no `/dev/gdrdrv` needed — validated on H200 with it absent); NVSHMEM/IBGDA on the wire is only a multi-node scale-out (EP16) concern. @@ -68,6 +72,7 @@ scale-up domain. | DeepEP V2 | `normal` mode is PR #605 `ElasticBuffer` plus exact upstream #630 and #640 fixes: LSA for scale-up and GIN for x86 EP16 scale-out. FP8 dispatch via `use_fp8_dispatch` (blockwise e4m3fn) alongside BF16. `low-latency` mode is the legacy `deep_ep.Buffer` IBGDA decode kernels (per-expert padded layout, weighted combine, `use_fp8` e4m3fn), decode/EP8 only | | MoRI | `normal` mode uses the direct `IntraNode` kernel for scale-up EP8 on every CDNA SKU and pins `InterNodeV1` for EP16 over 2x8 XGMI + RDMA. `low-latency` mode selects the `IntraNodeLL` decode kernel (single-call, pure-intranode, same compact layout and unweighted combine as `IntraNode`), decode/EP8 only. FP8 dispatch is caller-prequantized (per-SKU e4m3fnuz on gfx942, e4m3fn on gfx950); combine stays BF16 (`quant_type=none`) alongside BF16 dispatch | | UCCL-EP | [UCCL](https://github.com/uccl-project/uccl) EP: a drop-in, API-identical DeepEP replacement whose CPU proxies issue GPUDirect RDMA over plain `libibverbs` (no NVSHMEM/IBGDA), with software message ordering, atomics, and flow control; scale-up is single-node `cudaIpc` over NVLink/XGMI (never MNNVL). `normal` mode is the legacy `Buffer` `dispatch`/`combine` (unweighted rank-sum); `low-latency` reuses the legacy `low_latency_dispatch`/`low_latency_combine` decode kernels (weighted combine), decode/EP8 only. FP8 dispatch is caller-prequantized in `normal` mode (blockwise e4m3fn, per-SKU e4m3fnuz on gfx942); in `low-latency` mode the caller sends BF16 and the decode kernel quantizes to e4m3 internally (`use_fp8`). Combine is BF16. Runs on NVIDIA and AMD (H100/H200/B200 + MI300X/MI325X/MI355X), EP8 scale-up. Cross-node EP16 is functional (the internode RDMA path connects and the light case passes correctness) but its CPU-proxy throughput overruns the standardized per-case wall-clock budget on heavy token counts, so EP16 is an unsupported coverage row for now | +| NCCL EP | [NCCL EP](https://github.com/NVIDIA/nccl/tree/master/contrib/nccl_ep): NVIDIA's native MoE dispatch/combine on the NCCL Device API — LSA (NVLink load/store) intra-node, GIN (GPU-Initiated Networking) inter-node — driven through the `nccl4py` bindings. `normal` mode selects the `HIGH_THROUGHPUT` algorithm (FLAT `[N, hidden]` receive, unweighted rank-sum combine); the `LOW_LATENCY` algorithm is implemented in the adapter but has no enabled cell (see the `ll_backends` note above). BF16 only: NCCL EP's FP8 machinery exists upstream but its RELEASE.md lists it unsupported/untested, so no FP8 case is emitted. NVIDIA-only and CUDA 13 only. EP8 scale-up on H100/H200/B200/B300 plus EP8 and EP16 on GB200/GB300, where EP16 stays inside the MNNVL scale-up domain. x86 EP16 scale-out is an unsupported coverage row: the cross-node GIN path faults inside `nccl_ep.cc` identically on RoCE and IB across four SKUs, a GDAKI limit rather than a fabric-selection one | DeepEP V2 means the `ElasticBuffer` implementation introduced by [DeepEP PR #605](https://github.com/deepseek-ai/DeepEP/pull/605), not a newer legacy `Buffer` build. diff --git a/experimental/CollectiveX/bench/ep_nccl.py b/experimental/CollectiveX/bench/ep_nccl.py new file mode 100644 index 0000000000..c088cb469e --- /dev/null +++ b/experimental/CollectiveX/bench/ep_nccl.py @@ -0,0 +1,470 @@ +#!/usr/bin/env python3 +"""NCCL EP adapter: NVIDIA's native MoE dispatch/combine on the NCCL Device API. + +NCCL EP (github.com/NVIDIA/nccl/contrib/nccl_ep, arXiv 2603.13606) is a ground-up MoE +communication library built on NCCL's Device API — LSA (NVLink load/store) intra-node and +GIN (GPU-Initiated Networking) inter-node — with two algorithms selected per case: + normal -> HIGH_THROUGHPUT (HT), the Hybrid-EP-derived prefill/train path. FLAT recv + layout ``[N, hidden]`` (one row per received token, no expert structure) with + an unweighted rank-sum combine — identical semantics to deepep-v2 normal. + low-latency -> LOW_LATENCY (LL), the DeepEP-derived decode path. EXPERT_MAJOR recv layout + ``[num_local_experts, max_dispatch*num_ranks, hidden]`` with a source-side + weighted-kernel-sum combine — identical semantics to deepep-v2 low-latency. +Because both map exactly onto the two combine contracts CollectiveX already models, the +scale_up_domain two-level combine oracle in ep_harness applies unchanged. + +BF16 only. NCCL EP's FP8 machinery exists but RELEASE.md lists it unsupported/untested this +release, so this adapter does not override the FP8 encode hooks (SUPPORTED_PRECISIONS=("bf16",)). + +Communicator bootstrap: NCCL EP forms its OWN NCCL communicator (separate from PyTorch's +process group) via ``Communicator.init(nranks, rank, unique_id)``. Upstream broadcasts the +unique id with MPI; CollectiveX has no MPI, so rank 0 generates the id and we broadcast its +bytes over the already-initialized torch process group (see ``_bootstrap_comm``). + +Python bindings are ``nccl4py`` (``import nccl.core`` + ``nccl.ep``); the API surface used +here is verified against upstream ``bindings/nccl4py/nccl/ep`` and driven exactly as +``contrib/nccl_ep/ep_test.py`` drives it. +""" +from __future__ import annotations + +import os +import sys +import types + +import torch +import torch.distributed as dist + +from ep_backend import EPBackend + +try: + import nccl.core as nccl_core + import nccl.ep as nccl_ep + from nccl.ep import ( + Algorithm, + CombineConfig, + CombineInputs, + CombineOutputs, + DispatchConfig, + DispatchInputs, + DispatchOutputs, + GroupConfig, + HandleConfig, + Layout, + LayoutInfo, + Tensor, + ) +except Exception as exc: # pragma: no cover - requires the benchmark image + print(f"ERROR: NCCL EP import failed: {exc!r}", file=sys.stderr) + raise + + +# ncclUniqueId is a fixed 128-byte blob; we still broadcast the length first so a future size +# change can't silently truncate the id on the non-root ranks. +_UNIQUE_ID_MAX_BYTES = 256 + + +class NCCLEPBackend(EPBackend): + name = "nccl-ep" + # One library, two algorithms selected by args.mode. kernel_generation and the combine + # semantics are switched to their LL values in __init__ (mirrors ep_deepep_v2). + # normal -> HT / FLAT layout / unweighted-rank-sum combine. + # low-latency -> LL / EXPERT_MAJOR layout / source-side weighted-kernel-sum combine. + kernel_generation = "nccl-ep-ht" + SUPPORTED_MODES = ("normal", "low-latency") + SUPPORTED_PRECISIONS = ("bf16",) + stage_device_work = False + combine_needs_redispatch = False + dispatch_needs_combine_cleanup = False + combine_weight_semantics = "unweighted-rank-sum" + + def __init__(self, args, rank, world_size, local_rank, device): + super().__init__(args, rank, world_size, local_rank, device) + # NCCL EP group creation requires the NCCL Device API (LSA symmetric memory), which NCCL + # only advertises (comm.device_api_support) when cuMem allocation is enabled — otherwise + # ncclEpCreateGroup fails its deviceApiSupport gate with ncclInvalidUsage. The launcher + # exports this process-wide for every backend; setdefault before the EP comm is formed + # also covers manual/torchrun invocations. Verified on h100 EP8: absent -> group.create + # fails (error 5); present -> device_api_support=True and HT+LL groups create. + os.environ.setdefault("NCCL_CUMEM_ENABLE", "1") + self.group = dist.group.WORLD + self.experts_per_rank = args.experts // world_size + self.num_local_experts = self.experts_per_rank + self._internode = world_size > int(args.scale_up_domain) + self._ll = self.mode == "low-latency" + if self._ll: + # LL decode kernels apply the top-k gate at the source (weighted), not an + # unweighted rank sum — the benchmark stages the UNWEIGHTED per-expert transform + # and the kernel multiplies by the gate. Same contract as deepep-v2 low-latency. + self.kernel_generation = "nccl-ep-ll" + self.combine_weight_semantics = "weighted-kernel-sum" + # NCCL EP's handle is explicitly reusable across dispatch/combine cycles (ep_test.py + # cached mode redispatches and recombines on one handle), so — unlike DeepEP's legacy + # low-latency Buffer — no timed component needs a fresh dispatch or a draining combine; + # both modes keep combine_needs_redispatch / dispatch_needs_combine_cleanup False. If + # on-metal bring-up ever shows LL result-tensor aliasing, flip both to True for LL. + self._algorithm = Algorithm.LOW_LATENCY if self._ll else Algorithm.HIGH_THROUGHPUT + self._layout = Layout.EXPERT_MAJOR if self._ll else Layout.FLAT + # send_only=0 on every dispatch/combine (no staged execution). Handle.complete() is + # still required after each call — the HT combine only stages the local send and + # finishes the cross-rank gather in complete() (see _finish, verified on h100 EP8). + # FWD pass carries top-k weights on dispatch (HT) and forbids them on the HT combine + # input (the combine is a plain rank sum). + self._dispatch_cfg = DispatchConfig(send_only=0, round_scales=0) + self._combine_cfg = CombineConfig(send_only=0) + self._comm = None + self._ep_group = None + + def buffer_cap(self, args): + if self._ll: + # LL pre-allocates the fixed [num_local_experts, cap*num_ranks, hidden] receive + # buffer, so cap is a hard per-rank dispatch-slot bound (same 256 as ep_deepep_v2 / + # ep_uccl low-latency; the harness clamps the decode ladder and reports drops). + return 256 + return None + + # ---- helpers ----------------------------------------------------------------------------- + + @staticmethod + def _t(x): + """Wrap a torch tensor as an ``nccl.ep.Tensor`` (torch passthrough; the library reads + the device pointer/shape/dtype internally and anchors the torch buffer's lifetime).""" + return Tensor(x) + + def _stream(self): + """Raw handle of torch's current CUDA stream — NCCL EP runs on the same stream torch + times, so its work is captured by the harness's CUDA-event timing.""" + return torch.cuda.current_stream().cuda_stream + + def _finish(self, h, stream): + """Complete a dispatch/combine on the handle. Upstream ep_test.py calls this after + every dispatch and combine unconditionally: with send_only=0 the LL path self-completes + (combine runs SEND|RECV in one call), but the HT path only stages the local send in the + combine call and finishes the cross-rank gather in complete() — without it HT combine + returns just the source rank's own local-expert contribution (verified on h100 EP8).""" + h.handle.complete(stream=stream) + + def _bootstrap_comm(self): + """Form NCCL EP's own communicator, broadcasting the unique id over the torch PG. + + Rank 0 generates the id; we ship its raw bytes through the existing torch process + group (no MPI in CollectiveX) and every rank rebuilds it, then all ranks collectively + ``Communicator.init``. This communicator is independent of PyTorch's — the torch PG is + used only for this broadcast and for the harness's timing all-reduces/barriers. + """ + # Internode EP16 rides GIN over the RDMA fabric. NCCL EP's own ncclEpCreateGroup builds + # the internode DevComm with ginForceEnable + a RAIL connection type, so we do NOT force + # NCCL_GIN_TYPE here — forcing GDAKI(3) can prevent NCCL from selecting a fabric-supported + # GIN transport and trips an illegal-memory-access in the GIN dispatch kernel (on-metal + # checkpoint; still under cross-node bring-up). Set NCCL_GIN_TYPE in the launcher env if a + # specific transport must be pinned. + length = torch.zeros(1, dtype=torch.int64, device=self.device) + payload = torch.zeros(_UNIQUE_ID_MAX_BYTES, dtype=torch.uint8, device=self.device) + if self.rank == 0: + uid_bytes = nccl_core.get_unique_id().as_bytes + assert 0 < len(uid_bytes) <= _UNIQUE_ID_MAX_BYTES, ( + f"unexpected ncclUniqueId size {len(uid_bytes)}" + ) + length[0] = len(uid_bytes) + payload[: len(uid_bytes)] = torch.frombuffer( + bytearray(uid_bytes), dtype=torch.uint8 + ).to(self.device) + dist.broadcast(length, src=0) + dist.broadcast(payload, src=0) + n = int(length.item()) + uid = nccl_core.UniqueId.from_bytes(bytes(payload[:n].cpu().numpy().tobytes())) + self._comm = nccl_core.Communicator.init( + nranks=self.world_size, rank=self.rank, unique_id=uid + ) + + # ---- buffer construction ----------------------------------------------------------------- + + def create_buffer(self, spec): + """Bootstrap the communicator, create the EP group sized from the ladder maximum, and + allocate the persistent receive/combine buffers reused across every ladder shape.""" + self.max_dispatch = spec.max_tokens_per_rank + hidden = self.args.hidden + self._bootstrap_comm() + # max_recv_tokens_per_rank: HT requires >0 and >= max_dispatch; LL auto-derives when 0. + # world*max_dispatch is the recv-slot budget (every peer sends all its tokens here). + max_recv = self.max_dispatch * self.world_size + config = GroupConfig( + algorithm=self._algorithm, + num_experts=self.args.experts, + max_dispatch_tokens_per_rank=self.max_dispatch, + max_recv_tokens_per_rank=max_recv, + max_token_bytes=hidden * 2, # bfloat16 payload + ) + self._ep_group = nccl_ep.Group.create(self._comm, config) + + dev = self.device + if self._ll: + # EXPERT_MAJOR recv: [num_local_experts, max_dispatch*num_ranks, hidden]. + slots = self.max_dispatch * self.world_size + self._recv_x = torch.empty( + (self.num_local_experts, slots, hidden), dtype=torch.bfloat16, device=dev + ) + # Per-local-expert received-token counts, written by NCCL EP during dispatch. + self._recv_count = torch.empty( + (self.num_local_experts,), dtype=torch.int32, device=dev + ) + # Zeroed scratch the combine oracle scatters the transformed rows into. + self._combine_scratch = torch.empty_like(self._recv_x) + self._recv_x_t = self._t(self._recv_x) + self._recv_count_t = self._t(self._recv_count) + else: + # FLAT recv sized to the group's recv-slot budget max_recv_tokens_per_rank = + # world_size * max_dispatch — the max unique tokens this rank can receive (every peer + # sends up to max_dispatch and FLAT delivers one row per received token). HT combine + # copies this whole buffer into the group's expert_input_token IPC staging, which the + # group sized to exactly max_recv; oversizing it (e.g. ep_test's + # num_local_experts*max_dispatch, which only equals max_recv there because + # num_local_experts==n_ranks) overflows that buffer with cudaErrorInvalidValue. + rows = self.max_dispatch * self.world_size + self._recv_rows = rows + self._recv_x = torch.empty((rows, hidden), dtype=torch.bfloat16, device=dev) + self._recv_w = torch.empty((rows, self.args.topk), dtype=torch.float32, device=dev) + self._recv_idx = torch.empty((rows, self.args.topk), dtype=torch.int64, device=dev) + self._recv_x_t = self._t(self._recv_x) + self._recv_w_t = self._t(self._recv_w) + self._recv_idx_t = self._t(self._recv_idx) + # HT FLAT dispatch writes per-local-expert received counts (unpadded int32) here via + # the metadata path. Passing it in the dispatch LayoutInfo is REQUIRED for the + # internode GIN dispatch kernel — a null expert_counters is tolerated intranode (EP8) + # but faults the cross-node kernel (illegal memory access). ep_bench passes this too. + self._ht_disp_counts = torch.empty( + (self.num_local_experts,), dtype=torch.int32, device=dev + ) + self._ht_disp_counts_t = self._t(self._ht_disp_counts) + + def _ensure_handle(self, p): + """Create (once, cached on the problem) the reusable per-step handle for p's routing. + + create_handle is collective and, in HT, performs the metadata exchange that fixes the + received-token count — so it must run in the same order on every rank. It first runs + inside Pass 1's untimed warm (ladder order, identical across ranks); the timed passes + then reuse the cached handle and never enter a collective here. + """ + cached = getattr(p, "_nccl", None) + if cached is not None: + return cached + stream = self._stream() + topk_idx_t = self._t(p.topk_idx) + h = types.SimpleNamespace( + in_tokens_t=self._t(p.dispatch_x), + topk_idx_t=topk_idx_t, + ) + if not self._ll: + h.in_weights_t = self._t(p.topk_weights) + # combined output is restored to original token order: [num_tokens, hidden]. + h.out = torch.empty((p.T, self.args.hidden), dtype=torch.bfloat16, device=self.device) + h.out_t = self._t(h.out) + # HT: request the received-token counters at handle creation (populated by the HT + # metadata exchange). recv_total_counter is the authoritative FLAT row count. + ht_layout_info = None + if not self._ll: + h.recv_total = torch.zeros(1, dtype=torch.int32, device=self.device) + h.recv_experts = torch.zeros( + self.num_local_experts, dtype=torch.int32, device=self.device + ) + ht_layout_info = LayoutInfo( + expert_counters=self._t(h.recv_experts), + recv_total_counter=self._t(h.recv_total), + ) + h.handle = self._ep_group.create_handle( + self._layout, + topk_idx_t, + layout_info=ht_layout_info, + config=HandleConfig(), + stream=stream, + ) + torch.cuda.synchronize() + if not self._ll: + h.count = int(h.recv_total.item()) + p._nccl = h + return h + + # ---- transport contract ------------------------------------------------------------------ + + def dispatch(self, p): + h = self._ensure_handle(p) + stream = self._stream() + if self._ll: + # LL EXPERT_MAJOR: tokens in, 3D per-expert padded tokens out, per-expert recv + # counts written into expert_counters. No weights on the dispatch (the gate is + # applied by the combine kernel at the source). + h.handle.dispatch( + DispatchInputs(tokens=h.in_tokens_t), + DispatchOutputs(tokens=self._recv_x_t), + layout_info=LayoutInfo(expert_counters=self._recv_count_t), + config=self._dispatch_cfg, + stream=stream, + ) + h.recv_x = self._recv_x + h.recv_count = self._recv_count + else: + # HT FLAT: tokens + top-k weights in (FWD requires weights); received tokens, + # received top-k weights and GLOBAL top-k expert ids out. + h.handle.dispatch( + DispatchInputs(tokens=h.in_tokens_t, topk_weights=h.in_weights_t), + DispatchOutputs( + tokens=self._recv_x_t, + topk_weights=self._recv_w_t, + topk_idx=self._recv_idx_t, + ), + layout_info=LayoutInfo(expert_counters=self._ht_disp_counts_t), + config=self._dispatch_cfg, + stream=stream, + ) + h.recv_x = self._recv_x + h.recv_w = self._recv_w + h.recv_idx = self._recv_idx + self._finish(h, stream) + return h + + def stage(self, p, h): + # BF16 combine input is the received buffer itself; no device work (value correctness + # is exercised only through the oracle's combine_transformed path). + h.combine_input_t = self._recv_x_t + + def combine(self, p, h): + stream = self._stream() + if self._ll: + # Weighted LL combine: the kernel multiplies each expert contribution by the + # source token's gate (CombineOutputs.topk_weights) before the FP32 accumulation. + h.handle.combine( + CombineInputs(tokens=h.combine_input_t), + CombineOutputs(tokens=h.out_t, topk_weights=self._t(p.topk_weights)), + config=self._combine_cfg, + stream=stream, + ) + else: + # Unweighted HT combine (FWD forbids input weights): sums the per-token expert + # aggregates back to each token's home rank, restored to original order. + h.handle.combine( + CombineInputs(tokens=h.combine_input_t), + CombineOutputs(tokens=h.out_t), + config=self._combine_cfg, + stream=stream, + ) + self._finish(h, stream) + return h.out + + def recv_tokens(self, h): + if self._ll: + return int(h.recv_count.sum().item()) + return int(h.count) + + # ---- correctness-oracle views ------------------------------------------------------------ + + def _ll_inspect_dispatch(self, p, h): + """Flat per-slot view over the EXPERT_MAJOR padded receive (mirror of + ep_deepep_v2._ll_inspect_dispatch): each local expert's valid tokens are packed at the + front [0:recv_count[e]] of its slot dimension. Flatten to the oracle's compact + (expert, slot) row-major contract and keep the coordinates for the combine scatter.""" + recv_bf16 = h.recv_x # [E, S, hidden] BF16 + num_slots = recv_bf16.shape[1] + counts = h.recv_count.to(torch.int64) # [E] + slot_valid = ( + torch.arange(num_slots, device=recv_bf16.device).unsqueeze(0) < counts.unsqueeze(1) + ) + slot_expert, slot_j = slot_valid.nonzero(as_tuple=True) + h.slot_expert = slot_expert + h.slot_j = slot_j + local_lo = self.rank * self.num_local_experts + return types.SimpleNamespace( + payload=recv_bf16[slot_expert, slot_j], + expert_ids=local_lo + slot_expert.to(torch.int64), + local_expert_counts=counts, + ) + + def inspect_dispatch(self, p, h): + if self._ll: + return self._ll_inspect_dispatch(p, h) + # HT FLAT normal recv: front-packed to recv_total_counter, one row per received token. + # recv_idx holds this rank's LOCAL expert indices [0, experts_per_rank) front-packed per + # row (valid entries first, non-local padded to -1) with recv_w aligned to them — NOT the + # global top-k. (Verified on h100 EP8: rank-1 token recv_idx=[2,17,-1..] for global experts + # 34,49; rank 0 looks global only because its local range starts at 0.) So rebase the valid + # locals to the GLOBAL ids the oracle compares by + rank*experts_per_rank, exactly as + # ep_uccl/ep_deepep_v2 normal do. The oracle sorts each row over the top-k axis and sums + # the per-expert transforms, so token order is free and no per-(token,expert) expansion is + # needed. + count = int(h.count) + local_idx = h.recv_idx[:count].to(torch.int64) # [count, topk] LOCAL ids, -1 non-local + valid = local_idx >= 0 + expert_ids = torch.where( + valid, local_idx + self.rank * self.experts_per_rank, local_idx + ) + weights = h.recv_w[:count].to(torch.float32).masked_fill(~valid, 0) + return types.SimpleNamespace( + payload=h.recv_x[:count], # [count, hidden] BF16 + expert_ids=expert_ids, + weights=weights, + local_expert_counts=torch.bincount( + local_idx[valid], minlength=self.experts_per_rank + ), + ) + + def _ll_combine_transformed(self, p, h, transformed): + """Scatter the oracle-transformed rows back into a zeroed EXPERT_MAJOR combine buffer + at the exact (expert, slot) coordinates inspect read them from, then run the weighted + LL combine (mirror of ep_deepep_v2._ll_combine_transformed). `transformed` is + [N, hidden] in that slot order; the kernel applies p.topk_weights, so the staged + transform is unweighted.""" + combine_buf = self._combine_scratch + combine_buf.zero_() + combine_buf[h.slot_expert, h.slot_j] = transformed.to(combine_buf.dtype) + stream = self._stream() + h.handle.combine( + CombineInputs(tokens=self._t(combine_buf)), + CombineOutputs(tokens=h.out_t, topk_weights=self._t(p.topk_weights)), + config=self._combine_cfg, + stream=stream, + ) + self._finish(h, stream) + return h.out[: p.T] + + def combine_transformed(self, p, h, transformed): + if self._ll: + return self._ll_combine_transformed(p, h, transformed) + # `transformed` is the oracle's per-received-token combine input [count, hidden] + # (already summed over the top-k axis, gate folded in). Write it in place over the + # dispatch-output buffer (recv_x) — the same registered buffer the real expert MLP + # overwrites — so the combine's cross-rank gather reads it from the slots its routing + # map expects. Zero the padding tail first; combine sums across the token's unique + # destination ranks back to each token's home rank. + self._recv_x.zero_() + self._recv_x[: transformed.shape[0]].copy_(transformed.to(self._recv_x.dtype)) + stream = self._stream() + h.handle.combine( + CombineInputs(tokens=self._recv_x_t), + CombineOutputs(tokens=h.out_t), + config=self._combine_cfg, + stream=stream, + ) + self._finish(h, stream) + return h.out + + def finalize(self, rc): + """Clean teardown: NCCL EP's Device-API objects tear down without MoRI's post- + shmem_finalize assertion, so we destroy the handles/group/comm and the torch PG in + order rather than hard-exiting.""" + try: + dist.barrier() + self._destroy_handles() + if self._ep_group is not None: + self._ep_group.destroy() + if self._comm is not None: + self._comm.destroy() + dist.barrier() + dist.destroy_process_group() + except Exception: + return 1 + return rc + + def _destroy_handles(self): + # Per-problem handles are cached on the problem namespaces; the group keeps no registry, + # so there is nothing to walk here. Handles are released when their problems are GC'd + # (Handle.destroy runs in the binding's __del__); the group/comm destroy below reclaims + # the device buffers. Kept as a seam in case bring-up needs explicit handle teardown. + return diff --git a/experimental/CollectiveX/bench/run_ep.py b/experimental/CollectiveX/bench/run_ep.py index caaf9475d4..9e90daf7ea 100644 --- a/experimental/CollectiveX/bench/run_ep.py +++ b/experimental/CollectiveX/bench/run_ep.py @@ -50,7 +50,8 @@ def _runtime_info(torch, *, vendor: str) -> dict: def main() -> int: ap = argparse.ArgumentParser(description="CollectiveX EP dispatch/combine sweep") - ap.add_argument("--backend", required=True, choices=["deepep-v2", "mori", "uccl-ep"]) + ap.add_argument("--backend", required=True, + choices=["deepep-v2", "mori", "uccl-ep", "nccl-ep"]) ep_harness.add_common_args(ap) args = ap.parse_args() @@ -92,18 +93,21 @@ def main() -> int: from ep_mori import MoRIBackend as Backend elif args.backend == "uccl-ep": from ep_uccl import UCCLEPBackend as Backend + elif args.backend == "nccl-ep": + from ep_nccl import NCCLEPBackend as Backend else: from ep_deepep_v2 import DeepEPV2Backend as Backend # MoRI registers the default GPU process group with its SHMEM runtime. Keep that # group device-only so scale-out does not also depend on a host Gloo fabric. if not dist.is_initialized(): - if args.backend in ("mori", "uccl-ep"): + if args.backend in ("mori", "uccl-ep", "nccl-ep"): # MoRI registers this group with its SHMEM runtime; UCCL-EP is portable across # NVIDIA (NCCL) and AMD (RCCL) and bootstraps its Buffer + CPU-proxy ranks from - # it. Both take the explicit rank/world_size form and keep the group device-only - # so scale-out does not also depend on a host Gloo fabric. (uccl-ep's own default- - # PG registration, if any, happens in the adapter's create_buffer, mirroring MoRI.) + # it. NCCL EP forms its OWN NCCL communicator and uses this group only to broadcast + # that communicator's unique id (no MPI in CollectiveX) plus the harness's timing + # collectives. All take the explicit rank/world_size form and keep the group + # device-only so scale-out does not also depend on a host Gloo fabric. dist.init_process_group( backend="nccl", rank=rank, diff --git a/experimental/CollectiveX/configs/platform_config.json b/experimental/CollectiveX/configs/platform_config.json index 45925f9991..bc17013fb4 100644 --- a/experimental/CollectiveX/configs/platform_config.json +++ b/experimental/CollectiveX/configs/platform_config.json @@ -9,7 +9,7 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 2x200GbE", "switch": "Arista 7060DX5-64S (Tomahawk4, 25.6T)"}, "operator": { @@ -33,7 +33,7 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 400G", "switch": "NVIDIA Quantum-2 QM9790 (25.6T, InfiniBand)"}, "operator": { @@ -53,7 +53,7 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16], "uccl-ep": [8]}, + "backends": {"deepep-v2": [8, 16], "uccl-ep": [8], "nccl-ep": [8]}, "ll_backends": {"deepep-v2": [8], "uccl-ep": [8]}, "fabric": {"nic": "ConnectX-7 400GbE", "switch": "Whitebox Tomahawk3 leaf + Tomahawk4 (RoCE)"}, "operator": { @@ -76,7 +76,7 @@ "scale_up_domain": 8, "scale_up_transport": "nvlink", "launcher": "single-slurm", - "backends": {"deepep-v2": [8, 16]}, + "backends": {"deepep-v2": [8, 16], "nccl-ep": [8]}, "fabric": {"nic": "ConnectX-8 2x400GbE", "switch": "NVIDIA Spectrum-X SN5600 (51.2T)"}, "operator": { "partition": "batch_1", @@ -100,7 +100,7 @@ "scale_up_domain": 72, "scale_up_transport": "mnnvl", "launcher": "gb-nv", - "backends": {"deepep-v2": [8, 16]}, + "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]}, "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, "operator": { "partition": "batch", @@ -119,7 +119,7 @@ "scale_up_domain": 72, "scale_up_transport": "mnnvl", "launcher": "gb-nv", - "backends": {"deepep-v2": [8, 16]}, + "backends": {"deepep-v2": [8, 16], "nccl-ep": [8, 16]}, "fabric": {"nic": "MNNVL (scale-out not used)", "switch": "NVLink NVL72"}, "operator": { "partition": "batch_1", diff --git a/experimental/CollectiveX/docs/methodology.md b/experimental/CollectiveX/docs/methodology.md index 332b5cef03..4ec82c7258 100644 --- a/experimental/CollectiveX/docs/methodology.md +++ b/experimental/CollectiveX/docs/methodology.md @@ -21,9 +21,13 @@ It does not predict serving throughput without a separate correlation study. The implemented workload is `deepseek-v3`: hidden 7168, top-k 8, 256 routed experts, packed placement, and one pinned fixed resource profile per backend/topology. Combine is always BF16; -dispatch precision is a swept dimension — a BF16 control and an FP8 dispatch (`bf16`, `fp8`), +dispatch precision is a swept dimension — a BF16 control and, on the backends whose FP8 dispatch is +supported upstream (DeepEP V2, MoRI, UCCL-EP), an FP8 dispatch (`bf16`, `fp8`), caller-prequantized in `normal` mode (the `low-latency` kernels quantize FP8 internally from BF16 on -DeepEP and UCCL-EP, and stay caller-prequantized on MoRI). `normal`-mode cases use the +DeepEP and UCCL-EP, and stay caller-prequantized on MoRI). NCCL EP is BF16-only this release, so its +cells carry the control alone; the per-backend precision set lives in `sweep_matrix.py`'s +`BACKEND_PRECISIONS` and a backend never emits a case for a precision it does not support. +`normal`-mode cases use the `layout-and-dispatch-v1` semantics; `low-latency` cases use each backend's decode-kernel semantics (detailed below). @@ -58,7 +62,14 @@ input. UCCL-EP is a drop-in, API-identical DeepEP replacement that keeps the leg `dispatch`/`combine` (unweighted rank-sum) but routes it over CPU-proxy GPUDirect RDMA on plain `libibverbs` — no NVSHMEM/IBGDA — with software message ordering, atomics, and flow control; its scale-up is single-node `cudaIpc` over NVLink/XGMI (so the scale-up domain is one physical node, -never MNNVL) and its EP16 scale-out uses the same per-SKU RDMA rails as the other backends. Those throughput kernels run across the full token ladder in the `normal` mode. +never MNNVL) and its EP16 scale-out uses the same per-SKU RDMA rails as the other backends. NCCL EP +is NVIDIA's native MoE dispatch/combine on the NCCL Device API, driven through the `nccl4py` +bindings; `normal` mode selects its `HIGH_THROUGHPUT` algorithm, whose FLAT `[N, hidden]` receive and +unweighted rank-sum combine match `layout-and-dispatch-v1` exactly, so the same oracle applies. It is +NVIDIA-only and CUDA 13 only, and runs EP8 scale-up on H100/H200/B200/B300 plus EP8 and EP16 on +GB200/GB300, where EP16 stays inside the MNNVL scale-up domain; x86 EP16 scale-out is an unsupported +coverage row, its cross-node GIN path faulting inside `nccl_ep.cc` identically on RoCE and IB across +four SKUs — a GDAKI limit, not a fabric-selection one. Those throughput kernels run across the full token ladder in the `normal` mode. A second `low-latency` mode adds each backend's decode-optimized kernel family. On DeepEP it drives the legacy `deep_ep.Buffer` low-latency decode kernels (`low_latency_dispatch`/`low_latency_combine`), @@ -76,7 +87,12 @@ cell-by-cell from the registry's `ll_backends` map rather than assumed wherever currently enabled for DeepEP V2 EP8 on H100/H200/B200, MoRI EP8 on MI300X/MI325X/MI355X, and UCCL-EP EP8 on H100/H200/B200 only (the legacy `Buffer` low-latency kernels over UCCL's CPU-proxy transport; the AMD SKUs keep UCCL-EP normal mode but drop LL, whose -kernel trips a warp-group assertion on AMD's CU count). Whether a given SKU/backend/EP/mode cell is attempted is a capability +kernel trips a warp-group assertion on AMD's CU count). NCCL EP implements the mode — its +`LOW_LATENCY` algorithm is the DeepEP-derived decode path, EXPERT_MAJOR receive with a source-side +weighted-kernel-sum combine — but carries no `ll_backends` row on any SKU: the shipped decode kernels +consume stale peer signals under a fixed workload and wedge +([NVIDIA/nccl#2303](https://github.com/NVIDIA/nccl/issues/2303)), so the cells stay out of the matrix +until a fixed wheel ships. Whether a given SKU/backend/EP/mode cell is attempted is a capability fact; whether it succeeded is decided only by the emitted artifact. ## Workload Identity diff --git a/experimental/CollectiveX/launchers/launch_gb-nv.sh b/experimental/CollectiveX/launchers/launch_gb-nv.sh index a2d4b7b5ac..8488c5c739 100644 --- a/experimental/CollectiveX/launchers/launch_gb-nv.sh +++ b/experimental/CollectiveX/launchers/launch_gb-nv.sh @@ -36,7 +36,7 @@ export COLLX_TRANSPORT=mnnvl export COLLX_NODES="$NODES" COLLX_GPUS_PER_NODE="$GPN" COLLX_SCALE_UP_DOMAIN="$SCALE_UP_DOMAIN" export COLLX_NGPUS="$NGPUS" case "$COLLX_BENCH" in - deepep-v2) ;; + deepep-v2 | nccl-ep) ;; *) collx_die "unsupported $PRODUCT EP backend: $COLLX_BENCH" ;; esac collx_require_vars COLLX_IMAGE COLLX_IMAGE_PLATFORM COLLX_PARTITION COLLX_ACCOUNT COLLX_SQUASH_DIR COLLX_STAGE_DIR @@ -53,10 +53,12 @@ collx_select_image "$IMAGE" MOUNT_SRC="$(collx_stage_path "$REPO_ROOT" "$COLLX_STAGE_DIR")" collx_stage_repo "$REPO_ROOT" "$MOUNT_SRC" CONTAINER_MOUNTS="$MOUNT_SRC:/ix" -# ---- backend-setup: pinned DeepEP source + isolated build cache ------------- -# The backend case above admits only deepep-v2, so its staging is unconditional. -collx_prepare_deepep_source "$MOUNT_SRC" \ - || collx_die "cannot stage the pinned backend source" +# ---- backend-setup: pinned source (deepep-v2 only) + isolated build cache ---- +# nccl-ep is pip-only (nccl4py wheel; no source stage); deepep-v2 needs its pinned tree. +if [ "$COLLX_BENCH" = deepep-v2 ]; then + collx_prepare_deepep_source "$MOUNT_SRC" \ + || collx_die "cannot stage the pinned backend source" +fi export COLLX_BACKEND_SOURCE_ROOT=/ix/experimental/CollectiveX/.collx_sources collx_prepare_backend_cache "$COLLX_SQUASH_DIR" \ || collx_die "cannot prepare the isolated backend cache" diff --git a/experimental/CollectiveX/launchers/launch_mi-amds.sh b/experimental/CollectiveX/launchers/launch_mi-amds.sh index 9f0f2a137b..eee35d1fa1 100644 --- a/experimental/CollectiveX/launchers/launch_mi-amds.sh +++ b/experimental/CollectiveX/launchers/launch_mi-amds.sh @@ -126,8 +126,15 @@ CONTAINER_MOUNTS="$MOUNT_SRC:$MOUNT_DIR$DEVICE_MOUNTS" # uccl-ep builds from source, so give it the same cross-allocation backend cache the single-slurm # launcher provides (built once per arch/image/commit under /cx-cache, reused each allocation). # mori ships in the image and needs no cache, so its mounts are left untouched. +# +# The cache parent is the runner-shared stage BASE, not SQUASH_DIR. single-slurm can use its +# squash dir because that sits on shared storage, but this SKU's squash dir is node-local and +# root-owned (/var/lib/squash), which fails twice over: the submit-side mkdir is denied to the +# runner account, and even as root the directory it creates is not the one the compute node +# would bind-mount. COLLX_STAGE_DIR is runner-owned and compute-visible, and the cache lands +# beside job_ rather than inside it, so it still survives stage cleanup between allocations. if [ "$COLLX_BENCH" = uccl-ep ]; then - collx_prepare_backend_cache "$SQUASH_DIR" \ + collx_prepare_backend_cache "$COLLX_STAGE_DIR" \ || collx_die "cannot prepare the isolated backend cache" CONTAINER_MOUNTS="$CONTAINER_MOUNTS,$COLLX_PREPARED_BACKEND_CACHE:/cx-cache" export COLLX_BACKEND_CACHE_ROOT=/cx-cache diff --git a/experimental/CollectiveX/launchers/launch_mi-tw.sh b/experimental/CollectiveX/launchers/launch_mi-tw.sh index 089ebb850e..ec3d5265ab 100644 --- a/experimental/CollectiveX/launchers/launch_mi-tw.sh +++ b/experimental/CollectiveX/launchers/launch_mi-tw.sh @@ -57,6 +57,28 @@ if ! docker ps >/dev/null 2>&1; then fi fi +# Reap containers left behind by an earlier leg before touching the GPUs. +# +# A case container can outlive its GHA job: `docker run --rm` only removes the container when it +# exits, so when Actions kills the runner's process tree (job cancellation, or the 350-minute +# timeout firing on a wedged run) the daemon keeps it alive, and the workflow's non-root cleanup +# step cannot remove a root-owned container. It then pins every GPU on the node, and the failure +# it causes is nothing like a hang: hipIpcGetMemHandle starts returning "invalid argument" for any +# job with >= 3 ranks while 1- and 2-rank jobs still pass. Observed twice - tw032 carried a wedged +# MoRI low-latency leg for 5 days (4 red legs in one sweep, cleared by hand) and tw018 stranded one +# the same way when a MoRI run hit the job timeout. A wedged low-latency leg is a known class, so +# without this the next sweep re-poisons the node. +# +# The -tw fleets run one runner per node, so any container from the pinned image that predates this +# launcher belongs to a finished leg. Should a pool ever put two runners on one node, this needs a +# narrower predicate (match COLLX_CX_LABEL below and exclude the live execution id) instead. +COLLX_CX_LABEL="collectivex.leg" +for stray_id in $("${DOCKER[@]}" ps -q --filter "ancestor=$IMAGE" 2>/dev/null); do + stray_started="$("${DOCKER[@]}" inspect -f '{{.State.StartedAt}}' "$stray_id" 2>/dev/null)" || continue + collx_log "reaping stray container ${stray_id:0:12} from an earlier leg (started $stray_started)" + "${DOCKER[@]}" rm -f "$stray_id" >/dev/null 2>&1 || true +done + # The image is imported once per node and reused; pull only when absent. "${DOCKER[@]}" image inspect "$IMAGE" >/dev/null 2>&1 \ || "${DOCKER[@]}" pull "$IMAGE" >&2 \ @@ -187,6 +209,7 @@ for ((ci = 0; ci < ncases; ci++)); do for attempt in 1 2; do collx_log "case $ci/$ncases attempt $attempt: docker torchrun --nproc-per-node=$NGPUS" if "${DOCKER[@]}" run --rm \ + --label "$COLLX_CX_LABEL=${COLLECTIVEX_EXECUTION_ID:-manual}" \ --device /dev/kfd --device /dev/dri \ --group-add video --group-add render \ --ipc host --shm-size 32g \ diff --git a/experimental/CollectiveX/launchers/launch_single-slurm.sh b/experimental/CollectiveX/launchers/launch_single-slurm.sh index f1f81ab986..f30eef42e4 100644 --- a/experimental/CollectiveX/launchers/launch_single-slurm.sh +++ b/experimental/CollectiveX/launchers/launch_single-slurm.sh @@ -47,7 +47,7 @@ TIME_MIN="${COLLX_TIME:-$DEFAULT_TIME}" IMAGE="$COLLX_IMAGE" TS="$(date -u +%Y-%m-%dT%H-%M-%SZ)" case "$COLLX_BENCH" in - deepep-v2 | uccl-ep) ;; + deepep-v2 | uccl-ep | nccl-ep) ;; *) collx_die "unsupported $RUNNER EP backend: $COLLX_BENCH" ;; esac diff --git a/experimental/CollectiveX/runtime/common.sh b/experimental/CollectiveX/runtime/common.sh index c65bee18fe..24072604af 100644 --- a/experimental/CollectiveX/runtime/common.sh +++ b/experimental/CollectiveX/runtime/common.sh @@ -17,6 +17,16 @@ COLLX_DEEPEP_V2_COMMIT="fa8a9b16898204afd347c663b89e65ef87dc6ce6" COLLX_UCCL_REPO="https://github.com/uccl-project/uccl" COLLX_UCCL_COMMIT="fc1b582031221645ea9fce58aeb57187713145e3" +# NCCL EP (NVIDIA's native MoE dispatch/combine on the NCCL Device API). Primary path is the +# published nccl4py wheel — it bundles libnccl_ep.so's JIT runtime and pulls the matching +# nvidia-nccl-cu13 (>= 2.30, carrying the Device API + GIN nccl.ep needs). The from-source pins +# below are the fallback, deferred until on-metal bring-up shows the wheel is insufficient: +# contrib/nccl_ep is absent from the v2.29.x / v2.30.4 release tags, so any such build must use +# this post-merge master commit (which contains contrib/nccl_ep), NOT a release tag. +COLLX_NCCL4PY_SPEC="nccl4py[cu13]==0.3.1" +COLLX_NCCL_EP_REPO="https://github.com/NVIDIA/nccl" +COLLX_NCCL_EP_COMMIT="9d22d5dfec8391ee65b56df139d471f8e08e921e" + # Print bounded command output without maintaining a parallel failure taxonomy. collx_log_tail() { local log_path="$1" diff --git a/experimental/CollectiveX/runtime/prepare_backend.sh b/experimental/CollectiveX/runtime/prepare_backend.sh index 5ebf6b22a9..651129e0ec 100644 --- a/experimental/CollectiveX/runtime/prepare_backend.sh +++ b/experimental/CollectiveX/runtime/prepare_backend.sh @@ -404,6 +404,122 @@ uccl_prepare() { collx_log "UCCL-EP ready ($COLLX_UCCL_COMMIT, deep_ep wrapper over uccl.ep CPU-proxy runtime)" } +# ---- NCCL EP lifecycle ------------------------------------------------------ + +# Slug of the pinned pip spec, safe as a cache-dir path component. +nccl_ep_spec_slug() { + printf '%s' "$COLLX_NCCL4PY_SPEC" | tr -cs 'A-Za-z0-9_.-' '-' +} + +# Cache root keyed by cpu + build arch + image + pinned wheel spec, under the shared /cx-cache +# mount ($COLLX_BACKEND_CACHE_ROOT). Returns non-zero when no shared cache is mounted (manual +# runs), so the caller falls back to a node-local install. Mirrors uccl_cache_root. +nccl_ep_cache_root() { + local arch="$1" cpu base image slug + cpu="$(uname -m)" + [[ "$cpu" =~ ^[A-Za-z0-9._-]+$ ]] || return 1 + base="${COLLX_BACKEND_CACHE_ROOT:-}" + [[ "$base" = /* ]] || return 1 + image="$(printf '%s' "${COLLECTIVEX_IMAGE:-manual}" | tr -cs 'A-Za-z0-9_.-' '-')" + arch="$(printf '%s' "$arch" | tr -cs 'A-Za-z0-9_.-' '-')" + slug="$(nccl_ep_spec_slug)" + printf '%s/nccl-ep-%s-%s-%s-%s' \ + "$base" "$cpu" "${arch#-}" "${image#-}" "${slug#-}" +} + +# Put the installed wheel ($root/site) on PYTHONPATH for the probe and rank tasks, and the +# wheel-bundled NCCL runtime lib dir ahead of the image torch's older NCCL on the loader path +# (nccl.ep needs NCCL >= 2.29.3's Device API + GIN; the image torch bundles an older NCCL). Both +# PYTHONPATH and LD_LIBRARY_PATH are already carried to the ranks by write_rank_env. +nccl_ep_activate() { + local root="$1" site="$1/site" nccl_lib + [ -d "$site" ] || { collx_log "ERROR: NCCL EP cache site is unavailable"; return 1; } + export PYTHONPATH="$site${PYTHONPATH:+:$PYTHONPATH}" + for nccl_lib in "$site"/nvidia/nccl*/lib; do + if [ -d "$nccl_lib" ]; then + export LD_LIBRARY_PATH="$nccl_lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + break + fi + done + # NCCL EP group creation gates on the NCCL Device API (LSA symmetric memory), which NCCL only + # advertises when cuMem allocation is enabled; without it ncclEpCreateGroup returns + # ncclInvalidUsage. Persisted here (already in RANK_ENV_VARS) so every rank has it, mirroring + # the launcher's process-wide export. Verified on h100 EP8 (2026-07-21). + export NCCL_CUMEM_ENABLE=1 +} + +nccl_ep_probe() { + # import torch FIRST so libc10/libnccl are resident before the nccl.ep extension dlopens; then + # nccl.core (libnccl.so) and nccl.ep (libnccl_ep.so JIT runtime). nccl.ep.__init__ runs its own + # libnccl/libnccl_ep CUDA-major consistency check on import and raises ImportError on mismatch. + python3 - <<'PY' +import torch # noqa: F401 +import nccl.core # noqa: F401 +import nccl.ep # noqa: F401 +PY +} + +# Primary install: the published nccl4py[cu13] wheel + deps into $root/site via pip --target +# (self-contained; the runtime imports it through PYTHONPATH, so cache-hit and cache-miss paths +# import identically — mirrors uccl_install's copy-to-cache scheme). The from-source fallback +# (OpenMPI + build NCCL + contrib/nccl_ep from COLLX_NCCL_EP_COMMIT, with a matching launcher +# source-staging arm) is deferred until bring-up shows the wheel does not ship libnccl_ep.so. +nccl_ep_install() { + local root="$1" site="$1/site" + if [ -e "$root" ] || [ -L "$root" ]; then + rm -rf "$root" || { collx_log "ERROR: incomplete NCCL EP cache-reset failed"; return 1; } + fi + mkdir -m 700 "$root" || { collx_log "ERROR: NCCL EP cache-create failed"; return 1; } + mkdir -p "$site" || { collx_log "ERROR: NCCL EP cache-site-create failed"; return 1; } + collx_log "NCCL EP: installing $COLLX_NCCL4PY_SPEC (pip --target)" + # --target installs into an isolated tree and does not touch the system env, so PEP 668 does + # not apply; torch is imported from the image at runtime (nccl.ep's torch interop resolver). + python3 -m pip install -q --disable-pip-version-check --no-input \ + --target "$site" "$COLLX_NCCL4PY_SPEC" >&2 2>&1 \ + || { collx_log "ERROR: NCCL EP nccl4py install failed"; return 1; } + nccl_ep_activate "$root" \ + || { collx_log "ERROR: NCCL EP environment activation failed"; return 1; } + nccl_ep_probe || { collx_log "ERROR: NCCL EP import probe failed"; return 1; } + : > "$root/.ready" +} + +# NCCL EP lifecycle: install once per (arch, image, wheel-spec) into the shared /cx-cache behind +# an flock + .ready marker, reused on every later allocation (mirrors uccl_prepare); fall back to +# a node-local install when no shared cache is mounted (e.g. a manual run). +nccl_ep_prepare() { + local arch root ready lock_path + command -v python3 >/dev/null || { collx_log "ERROR: python3 unavailable for NCCL EP"; return 1; } + arch="$(cuda_arch)" || return 1 + if root="$(nccl_ep_cache_root "$arch")"; then + ready="$root/.ready"; lock_path="${root}.lock" + command -v flock >/dev/null \ + || { collx_log "ERROR: flock is required for NCCL EP caching"; return 1; } + mkdir -p "${root%/*}" || return 1 + collx_log "NCCL EP: preparing $COLLX_NCCL4PY_SPEC (shared cache $root)" + if ! ( + [ ! -L "$lock_path" ] || { collx_log "ERROR: NCCL EP cache lock is unsafe"; exit 1; } + (umask 077; : >> "$lock_path") && chmod 600 "$lock_path" \ + || { collx_log "ERROR: NCCL EP cache-lock-create failed"; exit 1; } + exec 9<>"$lock_path" || { collx_log "ERROR: NCCL EP cache-lock-open failed"; exit 1; } + flock 9 || { collx_log "ERROR: NCCL EP cache-lock-acquire failed"; exit 1; } + if [ ! -f "$ready" ] || [ ! -d "$root/site" ]; then + nccl_ep_install "$root" || exit 1 + fi + ); then + collx_log "ERROR: shared NCCL EP environment is incomplete"; return 1 + fi + else + root="/tmp/collectivex-nccl-ep-cache-$(nccl_ep_spec_slug)" + collx_log "NCCL EP: preparing $COLLX_NCCL4PY_SPEC (node-local $root; no shared cache mounted)" + if [ ! -f "$root/.ready" ] || [ ! -d "$root/site" ]; then + nccl_ep_install "$root" || return 1 + fi + fi + nccl_ep_activate "$root" || return 1 + nccl_ep_probe || { collx_log "ERROR: NCCL EP import probe failed"; return 1; } + collx_log "NCCL EP ready ($COLLX_NCCL4PY_SPEC; libnccl_ep.so JIT runtime, NCCL Device API LSA/GIN)" +} + # ---- container boundary ---------------------------------------------------- write_rank_env() { @@ -462,6 +578,7 @@ main() { || { collx_log "ERROR: MoRI backend import failed"; return 1; } ;; uccl-ep) uccl_prepare || return 1 ;; + nccl-ep) nccl_ep_prepare || return 1 ;; *) collx_log "ERROR: unknown backend preparation request" return 1 diff --git a/experimental/CollectiveX/sweep_matrix.py b/experimental/CollectiveX/sweep_matrix.py index 2dff3f5497..9a38f7f9ee 100644 --- a/experimental/CollectiveX/sweep_matrix.py +++ b/experimental/CollectiveX/sweep_matrix.py @@ -36,6 +36,9 @@ def _load_config(name: str) -> dict[str, Any]: "deepep-v2": ("bf16", "fp8"), "mori": ("bf16", "fp8"), "uccl-ep": ("bf16", "fp8"), + # NCCL EP is BF16-only this release: its FP8 machinery exists upstream but RELEASE.md + # lists it unsupported/untested, so no FP8 case is emitted (see bench/ep_nccl.py). + "nccl-ep": ("bf16",), } # Short shard-ID slug per non-normal mode. Normal-mode shard IDs carry no mode # segment so existing references stay valid; a low-latency shard adds "-ll". diff --git a/experimental/CollectiveX/tests/test_matrix.py b/experimental/CollectiveX/tests/test_matrix.py index f73c771327..07e3749f66 100644 --- a/experimental/CollectiveX/tests/test_matrix.py +++ b/experimental/CollectiveX/tests/test_matrix.py @@ -100,7 +100,8 @@ def test_runnable_cases_fan_out_over_backend_precisions(self): if precision in sweep_matrix.BACKEND_PRECISIONS[cell[1]] } self.assertEqual(precisions, expected, cell) - # Every backend (deepep-v2, mori, uccl-ep) realizes BF16 and FP8. + # Every backend that lists FP8 (deepep-v2, mori, uccl-ep) realizes BF16 and FP8; + # nccl-ep is BF16-only, so the cross-cell union of realized precisions stays {bf16, fp8}. self.assertEqual( {precision for precisions in by_cell.values() for precision in precisions}, {"bf16", "fp8"}, @@ -202,6 +203,67 @@ def test_uccl_ep_rollout_shape(self): } self.assertEqual(ll_skus, {"h100-dgxc", "h200-dgxc", "b200-dgxc"}) + def test_nccl_ep_rollout_shape(self): + # NCCL-EP's rollout, locked to the on-metal verdict (2026-07-22, all via the real launcher): + # * RDMA scale-out SKUs (h100/h200/b200/b300): EP8 runnable, EP16 an UNSUPPORTED coverage + # row. EP16 cross-node rides NCCL's kernel-initiated GDAKI GIN, which faults identically + # on RoCE (h100/b200/b300) and InfiniBand (h200) — a reproducible NCCL-EP v0.1.0 internode + # limitation, so EP16 is scoped out like uccl-ep's. + # * GB NVL72 SKUs (gb200/gb300, MNNVL, gb-nv launcher, 4 GPU/node): EP8 AND EP16 runnable — + # both stay inside the 72-GPU scale-up domain (world <= scale_up_domain => LSA, no GIN), + # so EP16 works over MNNVL where the RDMA-GIN path walls. + # * AMD SKUs: no rows (NCCL EP is NVIDIA-only; AMD runs mori). + # * LOW-LATENCY: no rows on ANY SKU. The wheel's LL count/flag protocol consumes stale + # double-buffer signals (values carry no generation, so a signal from two calls earlier is + # bit-identical at a repeating workload); a rank that slips one parity cycle gets lapped and + # the pipeline wedges on dispatch/combine receive timeouts, ending in cudaErrorLaunchFailure. + # Reported as NVIDIA/nccl#2303, fixed by NVIDIA/nccl#2306, unfixable here (we install the + # published wheel). Observed first on GB, then reproduced on every x86 SKU — 5/5 over SSH and + # 4/4 in sweep 30155842613 — so ll_backends carries no nccl-ep row anywhere until a fixed + # wheel ships. Restore per-SKU rows only alongside a spec bump that contains the fix. + # BF16 only — no FP8 case (NCCL EP FP8 unsupported this release). + document = matrix(backend="all") + runnable = { + (item["sku"], item["case"]["ep"]) + for item in document["requested_cases"] + if item["case"]["backend"] == "nccl-ep" and item["disposition"] == "runnable" + } + unsupported = { + (item["sku"], item["case"]["ep"]) + for item in document["requested_cases"] + if item["case"]["backend"] == "nccl-ep" and item["disposition"] == "unsupported" + } + rdma_skus = {"h100-dgxc", "h200-dgxc", "b200-dgxc", "b300"} + gb_skus = {"gb200", "gb300"} + # RDMA SKUs: EP8 runnable + EP16 unsupported. GB SKUs: EP8 and EP16 both runnable. + self.assertEqual( + runnable, + {(sku, 8) for sku in rdma_skus} | {(sku, ep) for sku in gb_skus for ep in (8, 16)}, + ) + self.assertEqual(unsupported, {(sku, 16) for sku in rdma_skus}) + # Offered on the six NVIDIA SKUs; never on AMD. + offered = {sku for sku, _ in runnable | unsupported} + self.assertEqual(offered, rdma_skus | gb_skus) + for absent in ("mi355x", "mi325x-tw", "mi300x-tw"): + self.assertNotIn(absent, offered) + # Every nccl-ep case is BF16 (FP8 unsupported this release). + self.assertEqual( + { + item["case"]["precision"] + for item in document["requested_cases"] + if item["case"]["backend"] == "nccl-ep" + }, + {"bf16"}, + ) + # Low-latency: no nccl-ep case on any SKU while the wheel carries the stale-signal wedge. + ll = { + (item["sku"], item["case"]["ep"]) + for item in document["requested_cases"] + if item["case"]["backend"] == "nccl-ep" + and item["case"]["mode"] == "low-latency" + } + self.assertEqual(ll, set()) + def test_invalid_filters_fail_closed(self): for options in ( {"exclude_skus": "unknown"}, diff --git a/experimental/CollectiveX/tests/test_runtime.py b/experimental/CollectiveX/tests/test_runtime.py index a3795ecae1..160e6c562f 100644 --- a/experimental/CollectiveX/tests/test_runtime.py +++ b/experimental/CollectiveX/tests/test_runtime.py @@ -340,7 +340,7 @@ def _run_ep_parser() -> argparse.ArgumentParser: # Mirror of the parser bench/run_ep.py builds in main(). parser = argparse.ArgumentParser() parser.add_argument( - "--backend", required=True, choices=["deepep-v2", "mori", "uccl-ep"] + "--backend", required=True, choices=["deepep-v2", "mori", "uccl-ep", "nccl-ep"] ) ep_harness.add_common_args(parser) return parser @@ -415,6 +415,21 @@ def test_uccl_ep_case_round_trips_through_the_run_ep_parser(self) -> None: self.assertEqual(args.case_id, uccl_case["case_id"]) self.assertEqual(args.out, "results/h200-dgxc_uccl-ep_bf16_decode_TS-c000.json") + def test_nccl_ep_case_round_trips_through_the_run_ep_parser(self) -> None: + # A nccl-ep case flows through the same generic codec; run_ep's --backend choices must + # accept "nccl-ep" and the result filename must carry the backend token so a nccl-ep leg + # never collides with the deepep-v2/uccl-ep legs of the same cell. BF16 only. + nccl_case = { + **self.CASE, + "backend": "nccl-ep", + "case_id": "h200-dgxc-nccl-ep-deepseek-v3-normal-decode-ep16-uniform-bf16", + } + argv = self._case_argv(["16", "2", "8", "8"], case=nccl_case) + args = self._run_ep_parser().parse_args(argv) + self.assertEqual(args.backend, "nccl-ep") + self.assertEqual(args.case_id, nccl_case["case_id"]) + self.assertEqual(args.out, "results/h200-dgxc_nccl-ep_bf16_decode_TS-c000.json") + # logical_byte_provenance is where FP8 changes MEASUREMENT semantics (asymmetric # per-direction byte counts), so its arithmetic and guards are pinned here on CPU.