mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-15 11:44:56 +00:00
Binyli/ep revise (#828)
This pull request makes significant improvements to the MoE (Mixture of Experts) Python API and documentation, focusing on clarifying and expanding the Expert Parallel (EP) interface, especially around quantization, dispatch/combine handles, and overlap configuration. The changes introduce new data structures, update function signatures, and improve documentation to better reflect the current and planned capabilities of the system. Additionally, the base development container is updated to CUDA 13.0, and minor corrections are made to extension naming.
This commit is contained in:
Binyang Li
476
test/python/ep/test_internode_multirank.py
Normal file
476
test/python/ep/test_internode_multirank.py
Normal file
@@ -0,0 +1,476 @@
|
||||
# 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()
|
||||
408
test/python/ep/test_intranode_multirank.py
Normal file
408
test/python/ep/test_intranode_multirank.py
Normal file
@@ -0,0 +1,408 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# Licensed under the MIT License.
|
||||
"""Multi-rank intranode 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).
|
||||
|
||||
Set ``MSCCLPP_EP_BENCH=1`` to also run a post-correctness benchmark pass
|
||||
that times dispatch and combine **separately** with CUDA events and
|
||||
reports per-phase latency (max across ranks) plus aggregate effective
|
||||
NVLink 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``.
|
||||
|
||||
This is a minimal adaptation of DeepEP's tests/test_intranode.py stripped
|
||||
to exercise only the code paths we've ported.
|
||||
"""
|
||||
|
||||
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 init_dist():
|
||||
rank = int(os.environ["RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", rank))
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
init_method=f"tcp://{os.environ.get('MASTER_ADDR','127.0.0.1')}:{os.environ.get('MASTER_PORT','29500')}",
|
||||
world_size=world_size,
|
||||
rank=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)
|
||||
|
||||
# Small settings for functional check
|
||||
num_tokens = 128
|
||||
hidden = 1024
|
||||
num_topk = min(4, num_ranks)
|
||||
num_experts = num_ranks * 4
|
||||
|
||||
torch.manual_seed(0xA1B2 + rank)
|
||||
|
||||
# Build topk layout that maps each token to num_topk distinct ranks/experts
|
||||
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)
|
||||
|
||||
# Expert / rank meta
|
||||
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")
|
||||
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")
|
||||
token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int)
|
||||
is_token_in_rank = token_idx_in_rank >= 0
|
||||
|
||||
# Token payload = rank id (cast to bf16) so we can check correctness
|
||||
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_NUM_SMS", "20")),
|
||||
nvl_chunked_send=int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")),
|
||||
nvl_chunked_recv=int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")),
|
||||
)
|
||||
if rank == 0:
|
||||
print(
|
||||
f"[cfg] num_ranks={num_ranks} num_tokens={num_tokens} hidden={hidden} "
|
||||
f"num_experts={num_experts} num_topk={num_topk}",
|
||||
flush=True,
|
||||
)
|
||||
print(f"[rank {rank}] MoECommunicator created is_available={moe.is_available()}", flush=True)
|
||||
assert moe.is_available()
|
||||
|
||||
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)
|
||||
|
||||
combined_x = moe.combine(recv_x, handle)
|
||||
|
||||
# Expected: we dispatched with x = rank * ones, so every destination r
|
||||
# received the value `rank` for our token. On combine the destinations
|
||||
# send that value back and we sum: combined[t] = rank * (#destinations).
|
||||
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()
|
||||
if rank == 0:
|
||||
print(f"[combine] max|got-expected|={diff:.4e} max|expected|={max_exp:.4e}", flush=True)
|
||||
assert diff < 1e-2, f"rank{rank}: combine mismatch max diff {diff}"
|
||||
|
||||
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
|
||||
|
||||
# Rebuild inputs at bench size. The benchmark creates its own communicator
|
||||
# below so its internal buffers are sized for the benchmark shape.
|
||||
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)
|
||||
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")
|
||||
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")
|
||||
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 first
|
||||
# (uncached) dispatch records the routing layout on the returned handle;
|
||||
# subsequent dispatches reuse it via previous_handle, skipping notify's
|
||||
# host-side counter wait. 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_NUM_SMS", "20")),
|
||||
nvl_chunked_send=int(os.environ.get("MSCCLPP_EP_NVL_SEND", "8")),
|
||||
nvl_chunked_recv=int(os.environ.get("MSCCLPP_EP_NVL_RECV", "256")),
|
||||
)
|
||||
assert moe.is_available()
|
||||
|
||||
# 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) using cached dispatch.
|
||||
for _ in range(warmup):
|
||||
_combine(_dispatch_cached())
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier(group=group)
|
||||
|
||||
# Time dispatch alone (cached mode -- skips notify_dispatch host wait).
|
||||
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
|
||||
|
||||
# Time combine alone (reusing the same dispatch output each iter).
|
||||
dist.barrier(group=group)
|
||||
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
|
||||
# (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
|
||||
# `total_send_tokens`. Recv side counts unique (src_rank, token) pairs
|
||||
# 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],
|
||||
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.
|
||||
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} "
|
||||
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. Avoids noisy
|
||||
# "recvValue failed / Connection was likely closed" stack traces from
|
||||
# ProcessGroupNCCL's HeartbeatMonitor.
|
||||
if dist.is_initialized():
|
||||
try:
|
||||
dist.barrier()
|
||||
except Exception:
|
||||
pass
|
||||
dist.destroy_process_group()
|
||||
335
test/python/ep/test_low_latency_multirank.py
Normal file
335
test/python/ep/test_low_latency_multirank.py
Normal file
@@ -0,0 +1,335 @@
|
||||
# Copyright (c) Microsoft Corporation.
|
||||
# Licensed under the MIT License.
|
||||
"""Multi-rank low-latency functional test for mscclpp_ep.
|
||||
|
||||
Launch with (intra-node, 8 GPUs):
|
||||
torchrun --nproc_per_node=8 test/python/ep/test_low_latency_multirank.py \
|
||||
--num-tokens 128 --hidden 7168 --num-topk 8 --num-experts 256
|
||||
|
||||
Launch with (2 nodes, 1 GPU per node -- DeepEP's recommended LL topology):
|
||||
# node 0:
|
||||
MASTER_ADDR=<master> MASTER_PORT=29600 NODE_RANK=0 \
|
||||
torchrun --nnodes=2 --nproc_per_node=1 --rdzv-backend=c10d \
|
||||
--rdzv-endpoint=<master>:29600 test/python/ep/test_low_latency_multirank.py
|
||||
# node 1:
|
||||
MASTER_ADDR=<master> MASTER_PORT=29600 NODE_RANK=1 \
|
||||
torchrun --nnodes=2 --nproc_per_node=1 --rdzv-backend=c10d \
|
||||
--rdzv-endpoint=<master>:29600 test/python/ep/test_low_latency_multirank.py
|
||||
|
||||
Exercises the LL dispatch + combine round-trip on a single node. The
|
||||
minimal correctness check:
|
||||
- dispatch: per-expert received token counts agree with an all-gathered
|
||||
reference computed from topk_idx across all ranks;
|
||||
- combine: the reconstructed x matches the analytical sum
|
||||
``x * sum(topk_weights, masked by topk_idx == -1)``.
|
||||
|
||||
Known limitation (see src/ext/ep/README.md): the LL kernels drive every
|
||||
peer via MSCCL++ PortChannel. Intra-node IB loopback between two HCAs on
|
||||
the same host (what an 8-GPU single-node launch exercises) currently hangs
|
||||
during dispatch; cross-node LL with one GPU per node works as designed.
|
||||
|
||||
Adapted from DeepEP/tests/test_low_latency.py stripped to the bare checks
|
||||
we need for an LL port smoke test. BF16-only (no FP8 check).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import random
|
||||
|
||||
# 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 parse_args():
|
||||
parser = argparse.ArgumentParser(description="MSCCL++ EP low-latency multi-rank correctness/benchmark test")
|
||||
parser.add_argument("--num-tokens", type=int, default=128)
|
||||
parser.add_argument("--hidden", type=int, default=7168, help="LL kernels are compiled for a fixed hidden set")
|
||||
parser.add_argument("--num-topk", type=int, default=8)
|
||||
parser.add_argument("--num-experts", type=int, default=256)
|
||||
parser.add_argument("--bench", action="store_true", help="Run dispatch/combine benchmark after correctness")
|
||||
parser.add_argument("--bench-warmup", type=int, default=5)
|
||||
parser.add_argument("--bench-iters", type=int, default=20)
|
||||
parser.add_argument("--local-rank", "--local_rank", type=int, default=None, help=argparse.SUPPRESS)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def init_dist():
|
||||
rank = int(os.environ["RANK"])
|
||||
world_size = int(os.environ["WORLD_SIZE"])
|
||||
local_rank = int(os.environ.get("LOCAL_RANK", rank))
|
||||
torch.cuda.set_device(local_rank)
|
||||
dist.init_process_group(
|
||||
backend="nccl",
|
||||
init_method=f"tcp://{os.environ.get('MASTER_ADDR','127.0.0.1')}:{os.environ.get('MASTER_PORT','29500')}",
|
||||
world_size=world_size,
|
||||
rank=rank,
|
||||
)
|
||||
return rank, world_size, local_rank, dist.new_group(list(range(world_size)))
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
rank, num_ranks, local_rank, group = init_dist()
|
||||
from mscclpp import CommGroup
|
||||
import mscclpp.ep as ep
|
||||
|
||||
ep_group = CommGroup(torch_group=group)
|
||||
|
||||
# Shrink the "bf16 precision" anchor to keep values small.
|
||||
rank_offset = 128
|
||||
assert num_ranks - rank_offset < 257, "too many ranks for bf16 precision anchor"
|
||||
|
||||
num_tokens = args.num_tokens
|
||||
hidden = args.hidden
|
||||
num_topk = args.num_topk
|
||||
num_experts = args.num_experts
|
||||
assert num_experts % num_ranks == 0
|
||||
num_local_experts = num_experts // num_ranks
|
||||
|
||||
torch.manual_seed(0xB3C4 + rank)
|
||||
random.seed(0xB3C4 + rank)
|
||||
|
||||
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * (rank - rank_offset)
|
||||
# Encode the per-token index into the last 128 elements so the receiver
|
||||
# can verify which source token it is looking at.
|
||||
x[:, -128:] = torch.arange(num_tokens, device="cuda").to(torch.bfloat16).view(-1, 1)
|
||||
scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs() + 1
|
||||
topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=True)[1]
|
||||
topk_weights = torch.randn((num_tokens, num_topk), dtype=torch.float32, device="cuda").abs()
|
||||
|
||||
# Randomly mask some positions
|
||||
for _ in range(min(10, num_tokens)):
|
||||
topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = -1
|
||||
|
||||
moe_comm = ep.MoECommunicator(
|
||||
comm=ep_group,
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_local_experts,
|
||||
hidden_size=hidden,
|
||||
topk=num_topk,
|
||||
max_tokens_per_rank=num_tokens,
|
||||
mode=ep.MoEMode.LOW_LATENCY,
|
||||
num_rdma_qps_per_rank=max(1, num_experts // num_ranks),
|
||||
)
|
||||
if rank == 0:
|
||||
print(
|
||||
f"[cfg] num_ranks={num_ranks} num_tokens={num_tokens} hidden={hidden} "
|
||||
f"num_experts={num_experts} num_topk={num_topk}",
|
||||
flush=True,
|
||||
)
|
||||
print(
|
||||
f"[rank {rank}] MoECommunicator created is_available={moe_comm.is_available()} "
|
||||
f"is_internode={moe_comm.is_internode_available()}",
|
||||
flush=True,
|
||||
)
|
||||
assert moe_comm.is_available()
|
||||
|
||||
dist.barrier(group=group)
|
||||
torch.cuda.synchronize()
|
||||
print(f"[rank {rank}] pre-dispatch", flush=True)
|
||||
|
||||
# --- Dispatch ---
|
||||
dispatch_output_buffer = torch.empty(
|
||||
(num_local_experts, num_ranks * num_tokens, hidden),
|
||||
dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
)
|
||||
dispatch_out, handle = moe_comm.dispatch(
|
||||
x,
|
||||
topk_idx,
|
||||
topk_weights,
|
||||
output_buffer=dispatch_output_buffer,
|
||||
)
|
||||
packed_recv_x = dispatch_out.tokens
|
||||
assert dispatch_out.layout.num_tokens_per_expert is not None
|
||||
packed_recv_count = dispatch_out.layout.num_tokens_per_expert
|
||||
packed_recv_layout_range = handle.combine_context.layout_range
|
||||
torch.cuda.synchronize()
|
||||
print(f"[rank {rank}] post-dispatch", flush=True)
|
||||
# packed_recv_x: [num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden]
|
||||
# packed_recv_count: [num_local_experts] int32
|
||||
|
||||
# Reference: gather all ranks' topk_idx and count expected tokens per expert.
|
||||
all_topk_idx = torch.empty((num_ranks, num_tokens, num_topk), dtype=topk_idx.dtype, device="cuda")
|
||||
dist.all_gather_into_tensor(all_topk_idx, topk_idx, group=group)
|
||||
|
||||
int_mask = (1 << 32) - 1
|
||||
for i in range(num_local_experts):
|
||||
expert_id = rank * num_local_experts + i
|
||||
recv_count = int(packed_recv_count[i].item())
|
||||
expected_count = int((all_topk_idx == expert_id).sum().item())
|
||||
recv_layout_range = packed_recv_layout_range[i]
|
||||
layout_sum = int((recv_layout_range & int_mask).sum().item())
|
||||
assert (
|
||||
recv_count == expected_count
|
||||
), f"rank{rank} expert{expert_id}: recv_count={recv_count} != expected={expected_count}"
|
||||
assert (
|
||||
layout_sum == recv_count
|
||||
), f"rank{rank} expert{expert_id}: layout range sum {layout_sum} != recv_count {recv_count}"
|
||||
|
||||
if recv_count:
|
||||
recv_x = packed_recv_x[i, :recv_count]
|
||||
# All columns except the last 128 should share the value (src_rank - rank_offset)
|
||||
recv_x_lo = recv_x[:, :-128]
|
||||
amin = recv_x_lo.amin(dim=-1)
|
||||
amax = recv_x_lo.amax(dim=-1)
|
||||
assert torch.equal(amin, amax), f"rank{rank} expert{expert_id}: non-uniform recv block"
|
||||
|
||||
if rank == 0:
|
||||
print(f"[dispatch] OK (ranks={num_ranks})", flush=True)
|
||||
|
||||
# --- Combine ---
|
||||
# Simulate the downstream GEMM output = identity (bf16 copy) so combine
|
||||
# returns sum(x * weight) across experts.
|
||||
simulated_gemm_x = packed_recv_x.clone()
|
||||
out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
|
||||
combined_x = moe_comm.combine(simulated_gemm_x, handle, out=out)
|
||||
|
||||
# Analytical expected: each token i, weighted sum over topk entries that
|
||||
# are not -1. Accumulate in the same top-k order as the kernel; multiplying
|
||||
# by the pre-summed weights can differ by one BF16 ULP for large token IDs.
|
||||
expected_f = torch.zeros_like(x, dtype=torch.float32)
|
||||
x_f = x.float()
|
||||
for j in range(num_topk):
|
||||
weight_j = topk_weights[:, j].masked_fill(topk_idx[:, j] == -1, 0.0).view(-1, 1)
|
||||
expected_f += x_f * weight_j
|
||||
expected = expected_f.to(torch.bfloat16)
|
||||
diff = (combined_x.float() - expected.float()).abs().max().item()
|
||||
max_exp = expected.float().abs().max().item()
|
||||
print(
|
||||
f"[combine r{rank}] max|got-expected|={diff:.4e} max|expected|={max_exp:.4e}",
|
||||
flush=True,
|
||||
)
|
||||
assert torch.isnan(combined_x).any().item() is False
|
||||
assert diff < 1e-2, f"rank{rank}: LL combine mismatch diff={diff}"
|
||||
|
||||
dist.barrier(group=group)
|
||||
if rank == 0:
|
||||
print("PASS", flush=True)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Optional benchmark. Times dispatch and combine separately, reporting
|
||||
# per-iter latency (max across ranks) and aggregate effective bandwidth
|
||||
# (sum across ranks).
|
||||
# ------------------------------------------------------------------
|
||||
if not args.bench:
|
||||
return
|
||||
|
||||
warmup = args.bench_warmup
|
||||
iters = args.bench_iters
|
||||
bench_dispatch_output_buffer = torch.empty_like(dispatch_output_buffer)
|
||||
|
||||
def _dispatch():
|
||||
return moe_comm.dispatch(
|
||||
x,
|
||||
topk_idx,
|
||||
topk_weights,
|
||||
output_buffer=bench_dispatch_output_buffer,
|
||||
)
|
||||
|
||||
# Hoist combine's output-tensor allocation out of the timed loop so the
|
||||
# measurement reflects the kernel cost. (The original test also cloned the
|
||||
# ~58 MB dispatch recv buffer on every iter, adding ~20 us of D2D memcpy
|
||||
# to each combine sample and masking kernel-level changes.)
|
||||
bench_out = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
def _combine(dout, out_):
|
||||
dispatch_out_, handle_ = dout
|
||||
moe_comm.combine(dispatch_out_.tokens, handle_, out=out_)
|
||||
|
||||
for _ in range(warmup):
|
||||
_combine(_dispatch(), bench_out)
|
||||
torch.cuda.synchronize()
|
||||
dist.barrier(group=group)
|
||||
|
||||
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()
|
||||
end_ev.record()
|
||||
torch.cuda.synchronize()
|
||||
disp_us = start_ev.elapsed_time(end_ev) * 1e3 / iters
|
||||
assert dout[0].layout.num_tokens_per_expert is not None
|
||||
recv_tokens = int(dout[0].layout.num_tokens_per_expert.sum().item())
|
||||
|
||||
dist.barrier(group=group)
|
||||
start_ev.record()
|
||||
for _ in range(iters):
|
||||
_combine(dout, bench_out)
|
||||
end_ev.record()
|
||||
torch.cuda.synchronize()
|
||||
comb_us = start_ev.elapsed_time(end_ev) * 1e3 / iters
|
||||
|
||||
# Dispatch payload: recv_tokens × hidden × bf16 (received on this rank).
|
||||
# Combine payload: recv_tokens × hidden × bf16 as well -- each local expert
|
||||
# sends one copy per dispatched token back to its owner, so the bytes on
|
||||
# the wire match dispatch. Using num_tokens × hidden here would under-count
|
||||
# the actual send payload by ~num_topk×.
|
||||
disp_bytes = recv_tokens * hidden * 2
|
||||
comb_bytes = recv_tokens * hidden * 2
|
||||
|
||||
# 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 = disp_bytes / (disp_avg_us * 1e-6) / 1e9
|
||||
comb_bw_per_rank = comb_bytes / (comb_avg_us * 1e-6) / 1e9
|
||||
if rank == 0:
|
||||
print(
|
||||
f"[bench LL] num_ranks={num_ranks} tokens={num_tokens} hidden={hidden} "
|
||||
f"num_experts={num_experts} num_topk={num_topk} 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" 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,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
finally:
|
||||
# Ordered shutdown: barrier so every rank reaches teardown before the
|
||||
# TCPStore server (rank 0) exits, then destroy the PG. Avoids noisy
|
||||
# "recvValue failed / Connection was likely closed" stack traces from
|
||||
# ProcessGroupNCCL's HeartbeatMonitor.
|
||||
if dist.is_initialized():
|
||||
try:
|
||||
dist.barrier()
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
dist.destroy_process_group()
|
||||
except Exception:
|
||||
pass
|
||||
Reference in New Issue
Block a user