mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-18 17:57:25 +00:00
Refactor HT EP for direct fabric domains (#837)
Use communicator-backed direct mappings, remove RDMA paths, and flatten the HT source layout. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 15f71a84-4219-4ae9-a87e-e5fab4205de6
This commit is contained in:
Binyang Li
@@ -7,7 +7,6 @@
|
||||
#
|
||||
# Example:
|
||||
# cmake -S test/python/ep -B test/python/ep/build \
|
||||
# -DMSCCLPP_EP_NUM_MAX_NVL_PEERS=8 \
|
||||
# -DCMAKE_CUDA_ARCHITECTURES=90
|
||||
# cmake --build test/python/ep/build -j 64
|
||||
|
||||
@@ -26,8 +25,6 @@ endif()
|
||||
|
||||
get_filename_component(_default_mscclpp_src "${CMAKE_CURRENT_SOURCE_DIR}/../../.." ABSOLUTE)
|
||||
set(MSCCLPP_SRC "${_default_mscclpp_src}" CACHE PATH "mscclpp source tree")
|
||||
set(MSCCLPP_EP_NUM_MAX_NVL_PEERS "8" CACHE STRING
|
||||
"Compile-time NUM_MAX_NVL_PEERS for the EP kernels (8 for HGX, 4 for GB200)")
|
||||
|
||||
find_package(MPI REQUIRED)
|
||||
find_package(CUDAToolkit REQUIRED)
|
||||
@@ -76,10 +73,7 @@ target_include_directories(mscclpp_ep_bench PRIVATE
|
||||
"${MSCCLPP_INSTALL_DIR}/include"
|
||||
"${CUPTI_INCLUDE_DIR}")
|
||||
|
||||
target_compile_definitions(mscclpp_ep_bench PRIVATE
|
||||
MSCCLPP_USE_CUDA
|
||||
EP_DISPATCH_NCCLEP
|
||||
NUM_MAX_NVL_PEERS=${MSCCLPP_EP_NUM_MAX_NVL_PEERS})
|
||||
target_compile_definitions(mscclpp_ep_bench PRIVATE MSCCLPP_USE_CUDA)
|
||||
|
||||
target_compile_options(mscclpp_ep_bench PRIVATE
|
||||
$<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>
|
||||
|
||||
@@ -1,476 +0,0 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# Licensed under the MIT License.
|
||||
"""Multi-rank internode (HT) functional validation for mscclpp_ep.
|
||||
|
||||
Launch on each node with (example: 2 nodes x 8 GPUs = 16 ranks):
|
||||
|
||||
# on master (NODE_RANK=0):
|
||||
MASTER_ADDR=<master_ip> MASTER_PORT=29600 NODE_RANK=0 \
|
||||
torchrun --nnodes=2 --nproc_per_node=8 \
|
||||
--rdzv-backend=c10d --rdzv-endpoint=<master_ip>:29600 \
|
||||
test/python/ep/test_internode_multirank.py
|
||||
|
||||
# on worker (NODE_RANK=1):
|
||||
MASTER_ADDR=<master_ip> MASTER_PORT=29600 NODE_RANK=1 \
|
||||
torchrun --nnodes=2 --nproc_per_node=8 \
|
||||
--rdzv-backend=c10d --rdzv-endpoint=<master_ip>:29600 \
|
||||
test/python/ep/test_internode_multirank.py
|
||||
|
||||
Round-trip dispatch + combine using internode HT kernels across nodes.
|
||||
|
||||
Set ``MSCCLPP_EP_BENCH=1`` to also run a post-correctness benchmark pass
|
||||
that times dispatch and combine **separately** with CUDA events. Reports
|
||||
per-phase latency (max across ranks) plus aggregate effective bandwidth
|
||||
(sum across ranks). Override iteration counts with
|
||||
``MSCCLPP_EP_BENCH_WARMUP`` / ``MSCCLPP_EP_BENCH_ITERS`` and the bench
|
||||
problem size with ``MSCCLPP_EP_BENCH_TOKENS`` / ``_HIDDEN``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
|
||||
# Disable ProcessGroupNCCL's HeartbeatMonitor before importing torch.distributed.
|
||||
# It runs in a background thread polling the TCPStore; under mpirun, rank 0
|
||||
# (the store server) can exit before non-zero ranks finish teardown, producing
|
||||
# noisy 'recvValue failed / Connection was likely closed' stack traces.
|
||||
os.environ.setdefault("TORCH_NCCL_ENABLE_MONITORING", "0")
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
|
||||
def _detect_local_world_size():
|
||||
"""Number of GPUs per node (4 on GB200, 8 on H100/A100, etc.).
|
||||
|
||||
Resolution order:
|
||||
1. `MSCCLPP_EP_LOCAL_WORLD_SIZE` env var (matches the C++ side).
|
||||
2. `LOCAL_WORLD_SIZE` (torchrun) or `OMPI_COMM_WORLD_LOCAL_SIZE` (mpirun).
|
||||
3. `torch.cuda.device_count()` on the current host.
|
||||
"""
|
||||
for var in ("MSCCLPP_EP_LOCAL_WORLD_SIZE", "LOCAL_WORLD_SIZE", "OMPI_COMM_WORLD_LOCAL_SIZE"):
|
||||
v = os.environ.get(var)
|
||||
if v and int(v) > 0:
|
||||
return int(v)
|
||||
return max(1, torch.cuda.device_count())
|
||||
|
||||
|
||||
def init_dist():
|
||||
rank = int(os.environ["RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
local_world_size = _detect_local_world_size()
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", rank % local_world_size))
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(
|
||||
backend="nccl", world_size=world_size, rank=rank, device_id=torch.device(f"cuda:{local_rank}")
|
||||
)
|
||||
return rank, world_size, local_rank, dist.new_group(list(range(world_size)))
|
||||
|
||||
|
||||
def inplace_unique(x: torch.Tensor, num_slots: int):
|
||||
assert x.dim() == 2
|
||||
mask = x < 0
|
||||
x_padded = x.masked_fill(mask, num_slots)
|
||||
bin_count = torch.zeros((x.size(0), num_slots + 1), dtype=x.dtype, device=x.device)
|
||||
bin_count.scatter_add_(1, x_padded, torch.ones_like(x_padded))
|
||||
bin_count = bin_count[:, :num_slots]
|
||||
sorted_bin_count, sorted_bin_idx = torch.sort(bin_count, dim=-1, descending=True)
|
||||
sorted_bin_idx.masked_fill_(sorted_bin_count == 0, -1)
|
||||
sorted_bin_idx = torch.sort(sorted_bin_idx, descending=True, dim=-1).values
|
||||
x[:, :].fill_(-1)
|
||||
valid_len = min(num_slots, x.size(1))
|
||||
x[:, :valid_len] = sorted_bin_idx[:, :valid_len]
|
||||
|
||||
|
||||
def main():
|
||||
rank, num_ranks, local_rank, group = init_dist()
|
||||
from mscclpp import CommGroup
|
||||
import mscclpp.ep as ep
|
||||
|
||||
ep_group = CommGroup(torch_group=group)
|
||||
|
||||
NUM_MAX_NVL_PEERS = _detect_local_world_size()
|
||||
assert (
|
||||
num_ranks % NUM_MAX_NVL_PEERS == 0 and num_ranks > NUM_MAX_NVL_PEERS
|
||||
), f"expected >1 node with {NUM_MAX_NVL_PEERS} GPUs each, got num_ranks={num_ranks}"
|
||||
num_nodes = num_ranks // NUM_MAX_NVL_PEERS
|
||||
num_local_ranks = NUM_MAX_NVL_PEERS
|
||||
|
||||
# Small settings for functional check
|
||||
import os as _os
|
||||
|
||||
num_tokens = int(_os.environ.get("MSCCLPP_EP_BENCH_TOKENS", "128"))
|
||||
hidden = int(_os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "1024"))
|
||||
num_topk = int(_os.environ.get("MSCCLPP_EP_BENCH_TOPK", str(min(4, num_ranks))))
|
||||
_experts_env = _os.environ.get("MSCCLPP_EP_BENCH_EXPERTS", "")
|
||||
num_experts = int(_experts_env) if _experts_env else num_ranks * 4
|
||||
assert num_experts % num_ranks == 0
|
||||
|
||||
torch.manual_seed(0xA1B2 + rank)
|
||||
|
||||
scores = torch.randn((num_tokens, num_experts), device="cuda", dtype=torch.float32).abs() + 1
|
||||
topk_idx = torch.topk(scores, num_topk, dim=-1, sorted=False).indices
|
||||
topk_weights = torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda")
|
||||
|
||||
rank_idx = topk_idx // (num_experts // num_ranks)
|
||||
rank_idx.masked_fill_(topk_idx == -1, -1)
|
||||
inplace_unique(rank_idx, num_ranks)
|
||||
|
||||
rdma_rank_idx = rank_idx // num_local_ranks
|
||||
rdma_rank_idx.masked_fill_(rank_idx == -1, -1)
|
||||
inplace_unique(rdma_rank_idx, num_nodes)
|
||||
|
||||
num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda")
|
||||
for i in range(num_experts):
|
||||
num_tokens_per_expert[i] = (topk_idx == i).sum()
|
||||
|
||||
num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda")
|
||||
num_tokens_per_rdma_rank = torch.empty((num_nodes,), dtype=torch.int, device="cuda")
|
||||
token_idx_in_rank = torch.full((num_ranks, num_tokens), -1, dtype=torch.long, device="cuda")
|
||||
for i in range(num_ranks):
|
||||
num_tokens_per_rank[i] = (rank_idx == i).sum()
|
||||
token_sel = (rank_idx == i).max(dim=-1).values
|
||||
cnt = token_sel.sum().item()
|
||||
tokens = torch.sort(token_sel.to(torch.int), descending=True).indices
|
||||
tokens[:cnt] = torch.sort(tokens[:cnt]).values
|
||||
token_idx_in_rank[i][tokens[:cnt]] = torch.arange(cnt, dtype=torch.long, device="cuda")
|
||||
for i in range(num_nodes):
|
||||
num_tokens_per_rdma_rank[i] = (rdma_rank_idx == i).sum()
|
||||
token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int)
|
||||
is_token_in_rank = token_idx_in_rank >= 0
|
||||
|
||||
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * float(rank)
|
||||
|
||||
moe = ep.MoECommunicator(
|
||||
comm=ep_group,
|
||||
num_experts=num_experts,
|
||||
hidden_size=hidden,
|
||||
topk=num_topk,
|
||||
max_tokens_per_rank=num_tokens,
|
||||
mode=ep.MoEMode.HIGH_THROUGHPUT,
|
||||
num_sms=int(os.environ.get("MSCCLPP_EP_NSM", "152")),
|
||||
nvl_chunked_send=int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")),
|
||||
nvl_chunked_recv=int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")),
|
||||
rdma_chunked_send=int(os.environ.get("MSCCLPP_EP_RDMA_SEND", "16")),
|
||||
rdma_chunked_recv=int(os.environ.get("MSCCLPP_EP_RDMA_RECV", "128")),
|
||||
)
|
||||
if rank == 0:
|
||||
print(
|
||||
f"[cfg] num_nodes={num_nodes} num_ranks={num_ranks} num_tokens={num_tokens} "
|
||||
f"hidden={hidden} num_experts={num_experts} num_topk={num_topk}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
print(
|
||||
f"[rank {rank}] MoECommunicator created is_available={moe.is_available()} "
|
||||
f"is_internode={moe.is_internode_available()}",
|
||||
flush=True,
|
||||
)
|
||||
assert moe.is_available() and moe.is_internode_available()
|
||||
assert moe.is_internode(), "expected the communicator to select the internode HT transport"
|
||||
|
||||
dispatch_out, handle = moe.dispatch(
|
||||
x,
|
||||
topk_idx,
|
||||
topk_weights,
|
||||
)
|
||||
recv_x = dispatch_out.tokens
|
||||
dist.barrier(group=group)
|
||||
|
||||
assert recv_x.dim() == 2 and recv_x.size(1) == hidden
|
||||
local_experts = num_experts // num_ranks
|
||||
all_expert_counts = torch.empty((num_ranks, num_experts), dtype=num_tokens_per_expert.dtype, device="cuda")
|
||||
dist.all_gather_into_tensor(all_expert_counts, num_tokens_per_expert, group=group)
|
||||
expected_counts = all_expert_counts[:, rank * local_experts : (rank + 1) * local_experts].sum(dim=0).cpu().tolist()
|
||||
assert dispatch_out.layout.num_tokens_per_expert is not None
|
||||
actual_counts = [int(count) for count in dispatch_out.layout.num_tokens_per_expert]
|
||||
assert actual_counts == [int(count) for count in expected_counts]
|
||||
if rank == 0:
|
||||
print(f"[dispatch] OK (recv {recv_x.size(0)} tokens)", flush=True)
|
||||
|
||||
# Keep the existing dispatch/combine phase guard for internode HT until the
|
||||
# backend wires a proper stream-dependency hand-off.
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier(group=group)
|
||||
|
||||
combined_x = moe.combine(recv_x, handle)
|
||||
|
||||
num_dst = is_token_in_rank.sum(dim=1).to(torch.float32)
|
||||
expected = num_dst * float(rank)
|
||||
got = combined_x.float().mean(dim=1)
|
||||
diff = (got - expected).abs().max().item()
|
||||
max_exp = expected.abs().max().item()
|
||||
print(f"[combine r{rank}] max|got-expected|={diff:.4e} max|expected|={max_exp:.4e}", flush=True)
|
||||
# bf16 accumulator has 7-bit mantissa; intermediate partial sums can
|
||||
# round at ulp = max_exp * 2**-7. Use a tolerance that scales with magnitude.
|
||||
tol = max(1e-2, max_exp * (1.0 / 64))
|
||||
assert diff <= tol, f"rank{rank}: combine mismatch max diff {diff} > tol {tol} (max_exp={max_exp})"
|
||||
|
||||
dist.barrier(group=group)
|
||||
if rank == 0:
|
||||
print("PASS", flush=True)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Optional benchmark (enable with MSCCLPP_EP_BENCH=1).
|
||||
# ------------------------------------------------------------------
|
||||
if os.environ.get("MSCCLPP_EP_BENCH", "0") != "1":
|
||||
return
|
||||
|
||||
warmup = int(os.environ.get("MSCCLPP_EP_BENCH_WARMUP", "5"))
|
||||
iters = int(os.environ.get("MSCCLPP_EP_BENCH_ITERS", "20"))
|
||||
bench_tokens = int(os.environ.get("MSCCLPP_EP_BENCH_TOKENS", "4096"))
|
||||
bench_hidden = int(os.environ.get("MSCCLPP_EP_BENCH_HIDDEN", "7168"))
|
||||
# Allow overriding num_experts / num_topk for the bench phase to match
|
||||
# NCCL-EP's `ep_bench -a ht` defaults (256 experts, top-8). The functional
|
||||
# check above still uses the smaller (num_experts=num_ranks*4, topk=4)
|
||||
# configuration.
|
||||
bench_num_experts = int(os.environ.get("MSCCLPP_EP_BENCH_EXPERTS", str(num_experts)))
|
||||
bench_num_topk = int(os.environ.get("MSCCLPP_EP_BENCH_TOPK", str(num_topk)))
|
||||
if bench_num_experts % num_ranks != 0:
|
||||
if rank == 0:
|
||||
print(
|
||||
f"[bench] skip: num_experts={bench_num_experts} not divisible " f"by num_ranks={num_ranks}", flush=True
|
||||
)
|
||||
return
|
||||
if bench_num_topk > bench_num_experts:
|
||||
if rank == 0:
|
||||
print(f"[bench] skip: topk={bench_num_topk} > experts={bench_num_experts}", flush=True)
|
||||
return
|
||||
|
||||
scores_b = torch.randn((bench_tokens, bench_num_experts), device="cuda", dtype=torch.float32).abs() + 1
|
||||
topk_idx_b = torch.topk(scores_b, bench_num_topk, dim=-1, sorted=False).indices
|
||||
topk_weights_b = torch.ones((bench_tokens, bench_num_topk), dtype=torch.float32, device="cuda")
|
||||
rank_idx_b = topk_idx_b // (bench_num_experts // num_ranks)
|
||||
rank_idx_b.masked_fill_(topk_idx_b == -1, -1)
|
||||
inplace_unique(rank_idx_b, num_ranks)
|
||||
rdma_rank_idx_b = rank_idx_b // num_local_ranks
|
||||
rdma_rank_idx_b.masked_fill_(rank_idx_b == -1, -1)
|
||||
inplace_unique(rdma_rank_idx_b, num_nodes)
|
||||
|
||||
num_tokens_per_expert_b = torch.zeros((bench_num_experts,), dtype=torch.int, device="cuda")
|
||||
for i in range(bench_num_experts):
|
||||
num_tokens_per_expert_b[i] = (topk_idx_b == i).sum()
|
||||
num_tokens_per_rank_b = torch.empty((num_ranks,), dtype=torch.int, device="cuda")
|
||||
num_tokens_per_rdma_rank_b = torch.empty((num_nodes,), dtype=torch.int, device="cuda")
|
||||
token_idx_in_rank_b = torch.full((num_ranks, bench_tokens), -1, dtype=torch.long, device="cuda")
|
||||
for i in range(num_ranks):
|
||||
num_tokens_per_rank_b[i] = (rank_idx_b == i).sum()
|
||||
token_sel = (rank_idx_b == i).max(dim=-1).values
|
||||
cnt = token_sel.sum().item()
|
||||
tokens = torch.sort(token_sel.to(torch.int), descending=True).indices
|
||||
tokens[:cnt] = torch.sort(tokens[:cnt]).values
|
||||
token_idx_in_rank_b[i][tokens[:cnt]] = torch.arange(cnt, dtype=torch.long, device="cuda")
|
||||
for i in range(num_nodes):
|
||||
num_tokens_per_rdma_rank_b[i] = (rdma_rank_idx_b == i).sum()
|
||||
token_idx_in_rank_b = token_idx_in_rank_b.T.contiguous().to(torch.int)
|
||||
is_token_in_rank_b = token_idx_in_rank_b >= 0
|
||||
x_b = torch.ones((bench_tokens, bench_hidden), dtype=torch.bfloat16, device="cuda") * float(rank)
|
||||
|
||||
# Drive the benchmark through the public high-level API. The communicator
|
||||
# auto-selects internode HT when the RDMA size hint is non-zero. The first
|
||||
# (uncached) dispatch records routing layout on the returned handle;
|
||||
# subsequent dispatches reuse it via previous_handle, skipping host-side
|
||||
# layout computation. This isolates the on-GPU dispatch-kernel cost
|
||||
# (NCCL-EP ep_bench convention).
|
||||
moe = ep.MoECommunicator(
|
||||
comm=ep_group,
|
||||
num_experts=bench_num_experts,
|
||||
hidden_size=bench_hidden,
|
||||
topk=bench_num_topk,
|
||||
max_tokens_per_rank=bench_tokens,
|
||||
mode=ep.MoEMode.HIGH_THROUGHPUT,
|
||||
num_sms=int(os.environ.get("MSCCLPP_EP_NSM", "152")),
|
||||
nvl_chunked_send=int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")),
|
||||
nvl_chunked_recv=int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")),
|
||||
rdma_chunked_send=int(os.environ.get("MSCCLPP_EP_RDMA_SEND", "16")),
|
||||
rdma_chunked_recv=int(os.environ.get("MSCCLPP_EP_RDMA_RECV", "128")),
|
||||
)
|
||||
assert moe.is_available() and moe.is_internode_available()
|
||||
assert moe.is_internode(), "expected the communicator to select the internode HT transport"
|
||||
|
||||
# One uncached dispatch to build the cached routing layout on the handle.
|
||||
_handle0 = moe.dispatch(x_b, topk_idx_b, topk_weights_b)[1]
|
||||
|
||||
def _dispatch_cached():
|
||||
return moe.dispatch(x_b, topk_idx_b, topk_weights_b, previous_handle=_handle0)
|
||||
|
||||
def _combine(dout):
|
||||
dispatch_out_, handle_ = dout
|
||||
moe.combine(dispatch_out_.tokens, handle_)
|
||||
|
||||
# Warmup (full round-trip with the sync/barrier guard between phases,
|
||||
# matching the correctness-path invariant: internode combine must observe
|
||||
# the completed dispatch outputs before it launches).
|
||||
for _ in range(warmup):
|
||||
dout = _dispatch_cached()
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier(group=group)
|
||||
_combine(dout)
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier(group=group)
|
||||
|
||||
# Time dispatch alone (cached mode -- skips the host-side layout computation).
|
||||
start_ev = torch.cuda.Event(enable_timing=True)
|
||||
end_ev = torch.cuda.Event(enable_timing=True)
|
||||
start_ev.record()
|
||||
dout = None
|
||||
for _ in range(iters):
|
||||
dout = _dispatch_cached()
|
||||
end_ev.record()
|
||||
torch.cuda.synchronize()
|
||||
disp_us = start_ev.elapsed_time(end_ev) * 1e3 / iters
|
||||
|
||||
# Required guard before combine sees the dispatch outputs (see correctness
|
||||
# path's XXX note). Not included in either phase's timing.
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier(group=group)
|
||||
|
||||
# Time combine alone (reusing the same dispatch output each iter).
|
||||
start_ev.record()
|
||||
for _ in range(iters):
|
||||
_combine(dout)
|
||||
end_ev.record()
|
||||
torch.cuda.synchronize()
|
||||
comb_us = start_ev.elapsed_time(end_ev) * 1e3 / iters
|
||||
|
||||
# Per-rank "send bytes" matches NCCL-EP's `ep_bench` accounting (`RDMA_send`):
|
||||
# bench_tokens * hidden * sizeof(bf16). Each rank ships its `bench_tokens`
|
||||
# input rows out (some replicated to multiple peers); NCCL-EP normalizes by
|
||||
# the input footprint, not by the recv-side fan-out. We use the same
|
||||
# convention here so `per_rank_bw` is directly comparable across stacks.
|
||||
bytes_one_way = bench_tokens * bench_hidden * x_b.element_size()
|
||||
|
||||
# NCCL-EP `ep_bench` six-metric breakdown.
|
||||
# Send-side accounting follows NCCL-EP: count unique (token, dst_node) pairs.
|
||||
# `num_tokens_per_rdma_rank_b[n]` is exactly that count for node `n`.
|
||||
# Recv-side accounting: each rank reports `num_tokens_per_rank_b[r]`
|
||||
# (tokens it sends to dst rank `r`); an `all_to_all_single` lets every
|
||||
# rank read how many tokens each source rank sent to it.
|
||||
bytes_per_token = bench_hidden * x_b.element_size()
|
||||
local_node = rank // num_local_ranks
|
||||
nodes_unique = num_tokens_per_rdma_rank_b.to(torch.int64)
|
||||
total_send_tokens_local = int(nodes_unique.sum().item())
|
||||
nvl_send_tokens_local = int(nodes_unique[local_node].item())
|
||||
rdma_send_tokens_local = total_send_tokens_local - nvl_send_tokens_local
|
||||
# Replaced dist.all_to_all_single (NCCL socket transport fails with
|
||||
# NCCL_IB_DISABLE=1 internode) with all_gather_into_tensor + transpose,
|
||||
# which works on the same socket-NCCL setup the LL test uses.
|
||||
_send_row = num_tokens_per_rank_b.to(torch.int64).contiguous()
|
||||
_gathered = torch.empty(num_ranks * num_ranks, dtype=torch.int64, device="cuda")
|
||||
dist.all_gather_into_tensor(_gathered, _send_row, group=group)
|
||||
recv_from_src = _gathered.view(num_ranks, num_ranks)[:, rank].contiguous()
|
||||
src_node = torch.arange(num_ranks, device="cuda") // num_local_ranks
|
||||
remote_mask = (src_node != local_node).to(torch.int64)
|
||||
total_recv_tokens_local = int(recv_from_src.sum().item())
|
||||
rdma_recv_tokens_local = int((recv_from_src * remote_mask).sum().item())
|
||||
|
||||
# Average per-rank token counts across ranks (matches NCCL-EP `Byte counts (per rank avg)`).
|
||||
counts_t = torch.tensor(
|
||||
[total_send_tokens_local, rdma_send_tokens_local, total_recv_tokens_local, rdma_recv_tokens_local],
|
||||
dtype=torch.float64,
|
||||
device="cuda",
|
||||
)
|
||||
dist.all_reduce(counts_t, op=dist.ReduceOp.SUM, group=group)
|
||||
counts_avg = (counts_t / num_ranks).tolist()
|
||||
total_send_avg, rdma_send_avg, total_recv_avg, rdma_recv_avg = counts_avg
|
||||
total_send_bytes = total_send_avg * bytes_per_token
|
||||
rdma_send_bytes = rdma_send_avg * bytes_per_token
|
||||
total_recv_bytes = total_recv_avg * bytes_per_token
|
||||
rdma_recv_bytes = rdma_recv_avg * bytes_per_token
|
||||
nvl_send_bytes = total_send_bytes - rdma_send_bytes
|
||||
nvl_recv_bytes = total_recv_bytes - rdma_recv_bytes
|
||||
|
||||
# Reduce timings: report min/avg/max and base BW on AVG to match NCCL-EP's
|
||||
# `ep_bench.cu` convention.
|
||||
disp_min_t = torch.tensor([disp_us], dtype=torch.float64, device="cuda")
|
||||
disp_avg_t = torch.tensor([disp_us], dtype=torch.float64, device="cuda")
|
||||
disp_max_t = torch.tensor([disp_us], dtype=torch.float64, device="cuda")
|
||||
comb_min_t = torch.tensor([comb_us], dtype=torch.float64, device="cuda")
|
||||
comb_avg_t = torch.tensor([comb_us], dtype=torch.float64, device="cuda")
|
||||
comb_max_t = torch.tensor([comb_us], dtype=torch.float64, device="cuda")
|
||||
dist.all_reduce(disp_min_t, op=dist.ReduceOp.MIN, group=group)
|
||||
dist.all_reduce(disp_avg_t, op=dist.ReduceOp.SUM, group=group)
|
||||
dist.all_reduce(disp_max_t, op=dist.ReduceOp.MAX, group=group)
|
||||
dist.all_reduce(comb_min_t, op=dist.ReduceOp.MIN, group=group)
|
||||
dist.all_reduce(comb_avg_t, op=dist.ReduceOp.SUM, group=group)
|
||||
dist.all_reduce(comb_max_t, op=dist.ReduceOp.MAX, group=group)
|
||||
disp_avg_us = disp_avg_t.item() / num_ranks
|
||||
comb_avg_us = comb_avg_t.item() / num_ranks
|
||||
disp_bw_per_rank = bytes_one_way / (disp_avg_us * 1e-6) / 1e9
|
||||
comb_bw_per_rank = bytes_one_way / (comb_avg_us * 1e-6) / 1e9
|
||||
# Six-metric BW (NCCL-EP convention). Combine reverses send<->recv:
|
||||
# in combine, this rank pushes back what it received in dispatch.
|
||||
disp_t_s = disp_avg_us * 1e-6
|
||||
comb_t_s = comb_avg_us * 1e-6
|
||||
d_send_total_bw = total_send_bytes / disp_t_s / 1e9
|
||||
d_send_nvl_bw = nvl_send_bytes / disp_t_s / 1e9
|
||||
d_send_rdma_bw = rdma_send_bytes / disp_t_s / 1e9
|
||||
d_recv_total_bw = total_recv_bytes / disp_t_s / 1e9
|
||||
d_recv_nvl_bw = nvl_recv_bytes / disp_t_s / 1e9
|
||||
d_recv_rdma_bw = rdma_recv_bytes / disp_t_s / 1e9
|
||||
c_send_total_bw = total_recv_bytes / comb_t_s / 1e9
|
||||
c_send_nvl_bw = nvl_recv_bytes / comb_t_s / 1e9
|
||||
c_send_rdma_bw = rdma_recv_bytes / comb_t_s / 1e9
|
||||
c_recv_total_bw = total_send_bytes / comb_t_s / 1e9
|
||||
c_recv_nvl_bw = nvl_send_bytes / comb_t_s / 1e9
|
||||
c_recv_rdma_bw = rdma_send_bytes / comb_t_s / 1e9
|
||||
if rank == 0:
|
||||
print(
|
||||
f"[bench internode HT] nodes={num_nodes} num_ranks={num_ranks} "
|
||||
f"tokens={bench_tokens} hidden={bench_hidden} "
|
||||
f"experts={bench_num_experts} topk={bench_num_topk} "
|
||||
f"warmup={warmup} iters={iters}",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f" dispatch: avg={disp_avg_us:.1f}us min={disp_min_t.item():.1f}us max={disp_max_t.item():.1f}us "
|
||||
f"per_rank_bw={disp_bw_per_rank:.2f} GB/s "
|
||||
f"agg_bw={disp_bw_per_rank * num_ranks:.2f} GB/s (BW @ avg time)",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f" send: total={d_send_total_bw:.2f} nvl={d_send_nvl_bw:.2f} rdma={d_send_rdma_bw:.2f} GB/s "
|
||||
f"recv: total={d_recv_total_bw:.2f} nvl={d_recv_nvl_bw:.2f} rdma={d_recv_rdma_bw:.2f} GB/s",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f" combine : avg={comb_avg_us:.1f}us min={comb_min_t.item():.1f}us max={comb_max_t.item():.1f}us "
|
||||
f"per_rank_bw={comb_bw_per_rank:.2f} GB/s "
|
||||
f"agg_bw={comb_bw_per_rank * num_ranks:.2f} GB/s (BW @ avg time)",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f" send: total={c_send_total_bw:.2f} nvl={c_send_nvl_bw:.2f} rdma={c_send_rdma_bw:.2f} GB/s "
|
||||
f"recv: total={c_recv_total_bw:.2f} nvl={c_recv_nvl_bw:.2f} rdma={c_recv_rdma_bw:.2f} GB/s",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f" byte counts (per rank avg): "
|
||||
f"total_send={total_send_bytes/1e6:.2f} MB ({total_send_avg:.0f} tok) "
|
||||
f"rdma_send={rdma_send_bytes/1e6:.2f} MB ({rdma_send_avg:.0f} tok) "
|
||||
f"total_recv={total_recv_bytes/1e6:.2f} MB ({total_recv_avg:.0f} tok) "
|
||||
f"rdma_recv={rdma_recv_bytes/1e6:.2f} MB ({rdma_recv_avg:.0f} tok)",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except Exception:
|
||||
import traceback
|
||||
|
||||
traceback.print_exc()
|
||||
sys.exit(1)
|
||||
finally:
|
||||
# Ordered shutdown: barrier so every rank reaches teardown before the
|
||||
# TCPStore server (rank 0) exits, then destroy the PG. Without this,
|
||||
# ProcessGroupNCCL's HeartbeatMonitor on non-zero ranks logs noisy
|
||||
# "recvValue failed / Connection was likely closed" stack traces.
|
||||
if dist.is_initialized():
|
||||
try:
|
||||
dist.barrier()
|
||||
except Exception:
|
||||
pass
|
||||
dist.destroy_process_group()
|
||||
@@ -1,13 +1,13 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# Licensed under the MIT License.
|
||||
"""Multi-rank intranode functional validation for mscclpp_ep.
|
||||
"""Multi-rank direct-fabric HT functional validation for mscclpp_ep.
|
||||
|
||||
Launch with:
|
||||
torchrun --nproc_per_node=<N> test/python/ep/test_intranode_multirank.py
|
||||
|
||||
Tests that the high-level ``MoECommunicator`` succeeds across N GPUs on a single
|
||||
node and that a round-trip dispatch + combine preserves data (sum of top-k
|
||||
weighted copies).
|
||||
Tests that the high-level ``MoECommunicator`` succeeds across GPUs in one
|
||||
detected GPU IPC/NVL fabric domain, including domains that span hosts, and that
|
||||
a round-trip dispatch + combine preserves data.
|
||||
|
||||
Set ``MSCCLPP_EP_BENCH=1`` to also run a post-correctness benchmark pass
|
||||
that times dispatch and combine **separately** with CUDA events and
|
||||
@@ -127,6 +127,10 @@ def main():
|
||||
)
|
||||
print(f"[rank {rank}] MoECommunicator created is_available={moe.is_available()}", flush=True)
|
||||
assert moe.is_available()
|
||||
local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", str(num_ranks)))
|
||||
expected_internode = num_ranks > local_world_size
|
||||
assert moe.is_internode_available() == expected_internode
|
||||
assert moe.is_internode() == expected_internode
|
||||
|
||||
dispatch_out, handle = moe.dispatch(
|
||||
x,
|
||||
@@ -271,16 +275,13 @@ def main():
|
||||
torch.cuda.synchronize()
|
||||
comb_us = start_ev.elapsed_time(end_ev) * 1e3 / iters
|
||||
|
||||
# Per-rank "send bytes" matches NCCL-EP's `ep_bench` accounting (`RDMA_send`):
|
||||
# Per-rank "send bytes" matches NCCL-EP's `ep_bench` accounting:
|
||||
# bench_tokens * hidden * sizeof(bf16). Each rank ships its `bench_tokens`
|
||||
# input rows out (some replicated to multiple peers); NCCL-EP normalizes by
|
||||
# the input footprint, not by the recv-side fan-out. We use the same
|
||||
# convention here so `per_rank_bw` is directly comparable across stacks.
|
||||
bytes_one_way = bench_tokens * bench_hidden * x_b.element_size()
|
||||
|
||||
# NCCL-EP `ep_bench` six-metric breakdown
|
||||
# (intranode -> single node, so rdma_*=0; nvl_*=total_*).
|
||||
#
|
||||
# Send side follows NCCL-EP: count unique (token, dst_node) pairs. With a
|
||||
# single node every selected destination collapses to that node, so a
|
||||
# token with at least one valid expert contributes exactly one to
|
||||
@@ -288,32 +289,23 @@ def main():
|
||||
# landing on this rank.
|
||||
bytes_per_token = bench_hidden * x_b.element_size()
|
||||
total_send_tokens_local = int(is_token_in_rank_b.any(dim=1).sum().item())
|
||||
rdma_send_tokens_local = 0 # intranode: no remote nodes
|
||||
# Replaced dist.all_to_all_single (NCCL socket transport fails with
|
||||
# NCCL_IB_DISABLE=1 internode) with all_gather_into_tensor + transpose,
|
||||
# which works on the same socket-NCCL setup the LL test uses.
|
||||
_send_row = num_tokens_per_rank_b.to(torch.int64).contiguous()
|
||||
_gathered = torch.empty(num_ranks * num_ranks, dtype=torch.int64, device="cuda")
|
||||
dist.all_gather_into_tensor(_gathered, _send_row, group=group)
|
||||
recv_from_src = _gathered.view(num_ranks, num_ranks)[:, rank].contiguous()
|
||||
total_recv_tokens_local = int(recv_from_src.sum().item())
|
||||
rdma_recv_tokens_local = 0 # intranode
|
||||
|
||||
# Average per-rank token counts across ranks (matches NCCL-EP `Byte counts (per rank avg)`).
|
||||
counts_t = torch.tensor(
|
||||
[total_send_tokens_local, rdma_send_tokens_local, total_recv_tokens_local, rdma_recv_tokens_local],
|
||||
[total_send_tokens_local, total_recv_tokens_local],
|
||||
dtype=torch.float64,
|
||||
device="cuda",
|
||||
)
|
||||
dist.all_reduce(counts_t, op=dist.ReduceOp.SUM, group=group)
|
||||
counts_avg = (counts_t / num_ranks).tolist()
|
||||
total_send_avg, rdma_send_avg, total_recv_avg, rdma_recv_avg = counts_avg
|
||||
total_send_avg, total_recv_avg = counts_avg
|
||||
total_send_bytes = total_send_avg * bytes_per_token
|
||||
rdma_send_bytes = rdma_send_avg * bytes_per_token
|
||||
total_recv_bytes = total_recv_avg * bytes_per_token
|
||||
rdma_recv_bytes = rdma_recv_avg * bytes_per_token
|
||||
nvl_send_bytes = total_send_bytes - rdma_send_bytes
|
||||
nvl_recv_bytes = total_recv_bytes - rdma_recv_bytes
|
||||
|
||||
# Reduce timings: report min/avg/max and base BW on AVG to match NCCL-EP's
|
||||
# `ep_bench.cu` convention.
|
||||
@@ -333,21 +325,12 @@ def main():
|
||||
comb_avg_us = comb_avg_t.item() / num_ranks
|
||||
disp_bw_per_rank = bytes_one_way / (disp_avg_us * 1e-6) / 1e9
|
||||
comb_bw_per_rank = bytes_one_way / (comb_avg_us * 1e-6) / 1e9
|
||||
# Six-metric BW (NCCL-EP convention). Combine reverses send<->recv.
|
||||
disp_t_s = disp_avg_us * 1e-6
|
||||
comb_t_s = comb_avg_us * 1e-6
|
||||
d_send_total_bw = total_send_bytes / disp_t_s / 1e9
|
||||
d_send_nvl_bw = nvl_send_bytes / disp_t_s / 1e9
|
||||
d_send_rdma_bw = rdma_send_bytes / disp_t_s / 1e9
|
||||
d_recv_total_bw = total_recv_bytes / disp_t_s / 1e9
|
||||
d_recv_nvl_bw = nvl_recv_bytes / disp_t_s / 1e9
|
||||
d_recv_rdma_bw = rdma_recv_bytes / disp_t_s / 1e9
|
||||
c_send_total_bw = total_recv_bytes / comb_t_s / 1e9 # combine sends back what dispatch received
|
||||
c_send_nvl_bw = nvl_recv_bytes / comb_t_s / 1e9
|
||||
c_send_rdma_bw = rdma_recv_bytes / comb_t_s / 1e9
|
||||
c_recv_total_bw = total_send_bytes / comb_t_s / 1e9 # combine receives back what dispatch sent
|
||||
c_recv_nvl_bw = nvl_send_bytes / comb_t_s / 1e9
|
||||
c_recv_rdma_bw = rdma_send_bytes / comb_t_s / 1e9
|
||||
if rank == 0:
|
||||
print(
|
||||
f"[bench intranode HT] tokens={bench_tokens} hidden={bench_hidden} "
|
||||
@@ -362,8 +345,7 @@ def main():
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f" send: total={d_send_total_bw:.2f} nvl={d_send_nvl_bw:.2f} rdma={d_send_rdma_bw:.2f} GB/s "
|
||||
f"recv: total={d_recv_total_bw:.2f} nvl={d_recv_nvl_bw:.2f} rdma={d_recv_rdma_bw:.2f} GB/s",
|
||||
f" send={d_send_total_bw:.2f} GB/s recv={d_recv_total_bw:.2f} GB/s",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
@@ -373,16 +355,13 @@ def main():
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f" send: total={c_send_total_bw:.2f} nvl={c_send_nvl_bw:.2f} rdma={c_send_rdma_bw:.2f} GB/s "
|
||||
f"recv: total={c_recv_total_bw:.2f} nvl={c_recv_nvl_bw:.2f} rdma={c_recv_rdma_bw:.2f} GB/s",
|
||||
f" send={c_send_total_bw:.2f} GB/s recv={c_recv_total_bw:.2f} GB/s",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f" byte counts (per rank avg): "
|
||||
f"total_send={total_send_bytes/1e6:.2f} MB ({total_send_avg:.0f} tok) "
|
||||
f"rdma_send={rdma_send_bytes/1e6:.2f} MB ({rdma_send_avg:.0f} tok) "
|
||||
f"total_recv={total_recv_bytes/1e6:.2f} MB ({total_recv_avg:.0f} tok) "
|
||||
f"rdma_recv={rdma_recv_bytes/1e6:.2f} MB ({rdma_recv_avg:.0f} tok)",
|
||||
f"total_recv={total_recv_bytes/1e6:.2f} MB ({total_recv_avg:.0f} tok)",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user