mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-14 19:27:20 +00:00
d6678fa8cb
Use Bootstrap node and NVML fabric-domain discovery for optimized LL, collectively gate cross-host fabric payload and semaphore support, remove obsolete topology env defaults, and refresh the EP documentation and performance guidance. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: efbacae6-f679-430b-bc16-b45ae162fc76
570 lines
24 KiB
Python
570 lines
24 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) Microsoft Corporation.
|
|
# Licensed under the MIT License.
|
|
"""Unified EP low-latency benchmark driver.
|
|
|
|
Runs the *same* low-latency dispatch/combine benchmark -- identical tokens,
|
|
experts, hidden size, top-k, warmup and iteration counts -- against a
|
|
selectable expert-parallel library, then prints one normalized summary so the
|
|
libraries can be compared apples-to-apples.
|
|
|
|
Backends (``--ep-lib``):
|
|
|
|
* ``mscclpp`` -- this repo's :mod:`ep_bench_ll` (MoECommunicator LL) launched
|
|
with ``torchrun``.
|
|
* ``mscclpp-cpp`` -- the pure-C++ ``MoERuntime`` benchmark launched with MPI.
|
|
* ``nccl-ep`` -- NVIDIA NCCL-EP's ``contrib/nccl_ep/ep_bench`` binary launched
|
|
with ``mpirun`` (HPCX).
|
|
* ``both`` -- run mscclpp then nccl-ep and print them side by side.
|
|
|
|
Both backends emit the identical ``=== Summary (Low Latency, across N ranks) ===``
|
|
block (``ep_bench_ll.py`` was written to mirror ``ep_bench``), so a single parser
|
|
reads either one.
|
|
|
|
NCCL-EP dynamically links its shared libraries (``libnccl.so``, ``libnccl_ep.so``).
|
|
Point the driver at the correct build with ``--nccl-lib-path`` (falls back to the
|
|
``NCCL_LIB_PATH`` environment variable, else the ``lib`` directory beside the
|
|
``--nccl-ep-bench`` build tree); that directory is prepended to ``LD_LIBRARY_PATH``
|
|
for the ``ep_bench`` process so the intended NCCL is loaded.
|
|
|
|
Scope: single node (``--nproc-per-node`` GPUs). Multi-node runs use the existing
|
|
per-backend launchers (mscclpp: run_ep_bench_ll_multinode.sh; nccl-ep: mpirun with
|
|
a hostfile); this driver focuses on the common single-node comparison.
|
|
|
|
Examples
|
|
--------
|
|
Compare both libraries, 4 GPUs, e128::
|
|
|
|
python run_ep_bench.py --ep-lib both -e 128 -t 128 -d 7168 -k 8 -w 10 -i 50 \
|
|
--nccl-lib-path /opt/microsoft/mrc/ep/nccl/build/lib
|
|
|
|
Just mscclpp with in-process CUPTI kernel timing::
|
|
|
|
python run_ep_bench.py --ep-lib mscclpp -e 128 --cupti-inproc
|
|
|
|
Print the commands without running them::
|
|
|
|
python run_ep_bench.py --ep-lib both -e 128 --dry-run
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import glob
|
|
import os
|
|
import re
|
|
import shlex
|
|
import subprocess
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
from typing import Optional
|
|
|
|
CUDA_HOME = os.environ.get("CUDA_HOME", "/usr/local/cuda")
|
|
_HERE = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
def _find_cupti_paths() -> tuple[str, str]:
|
|
target_dirs = sorted(glob.glob(os.path.join(CUDA_HOME, "targets", "*")))
|
|
include_candidates = [os.path.join(path, "include") for path in target_dirs]
|
|
include_candidates.append(os.path.join(CUDA_HOME, "extras", "CUPTI", "include"))
|
|
library_candidates = [os.path.join(path, "lib") for path in target_dirs]
|
|
library_candidates.append(os.path.join(CUDA_HOME, "extras", "CUPTI", "lib64"))
|
|
|
|
include_dir = next(
|
|
(path for path in include_candidates if os.path.isfile(os.path.join(path, "cupti.h"))),
|
|
"",
|
|
)
|
|
library_dir = next(
|
|
(path for path in library_candidates if glob.glob(os.path.join(path, "libcupti.so*"))),
|
|
"",
|
|
)
|
|
if not include_dir or not library_dir:
|
|
raise SystemExit(f"CUPTI was not found under CUDA_HOME={CUDA_HOME}")
|
|
return include_dir, library_dir
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
p = argparse.ArgumentParser(
|
|
description="Unified EP low-latency benchmark driver (mscclpp EP vs NCCL-EP)",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
)
|
|
p.add_argument(
|
|
"--ep-lib",
|
|
required=True,
|
|
choices=["mscclpp", "mscclpp-cpp", "nccl-ep", "both", "all"],
|
|
help="which expert-parallel library to benchmark. mscclpp=MoECommunicator (Python), "
|
|
"mscclpp-cpp=MoERuntime (pure C++), nccl-ep=ep_bench. both=mscclpp+nccl-ep; all=the three.",
|
|
)
|
|
p.add_argument(
|
|
"-a",
|
|
"--algorithm",
|
|
default="ll",
|
|
choices=["ll", "low-latency"],
|
|
help="algorithm mode (only low-latency is wired up here)",
|
|
)
|
|
|
|
# Shared problem shape -- passed to whichever backend is selected.
|
|
p.add_argument("-t", "--num-tokens", type=int, default=128, help="tokens per rank")
|
|
p.add_argument(
|
|
"-d",
|
|
"--hidden",
|
|
type=int,
|
|
default=7168,
|
|
choices=(4096, 7168, 8192, 9216),
|
|
help="hidden dimension",
|
|
)
|
|
p.add_argument("-k", "--num-topk", type=int, default=8, choices=range(1, 10), help="top-k experts per token")
|
|
p.add_argument("-e", "--num-experts", type=int, default=256, help="global number of experts")
|
|
p.add_argument("-w", "--num-warmup", type=int, default=10, help="warmup iterations")
|
|
p.add_argument("-i", "--num-iters", type=int, default=50, help="timed iterations")
|
|
p.add_argument(
|
|
"--dispatch-dtype",
|
|
choices=("bf16", "fp8_e4m3"),
|
|
default="bf16",
|
|
help="MSCCL++ dispatch format; NCCL-EP runs keep their own configured format",
|
|
)
|
|
p.add_argument(
|
|
"--combine-mode",
|
|
choices=("rank_local_reduce", "direct_send"),
|
|
default="rank_local_reduce",
|
|
help="MSCCL++ low-latency combine mode",
|
|
)
|
|
p.add_argument("--num-blocks", type=int, default=130, help="MSCCL++ low-latency dispatch blocks")
|
|
|
|
# Launch / fabric.
|
|
p.add_argument("--nproc-per-node", type=int, default=4, help="GPUs (ranks) on this node")
|
|
p.add_argument(
|
|
"--nodes",
|
|
default="",
|
|
help="space-separated node IPs for a multi-node run (first = master). Empty = single "
|
|
"local node. Applies to the mpirun backends (nccl-ep, mscclpp-cpp); the Python "
|
|
"mscclpp backend is single-node only (torchrun --standalone).",
|
|
)
|
|
p.add_argument("--iface", default="", help="optional socket interface name (NCCL/GLOO/UCX)")
|
|
p.add_argument("--hca", default="", help="optional comma-separated mscclpp HCA devices")
|
|
|
|
# mscclpp backend.
|
|
p.add_argument("--mscclpp-bench", default=os.path.join(_HERE, "ep_bench_ll.py"), help="path to ep_bench_ll.py")
|
|
p.add_argument("--python", default=sys.executable, help="Python interpreter used for the MSCCL++ benchmark")
|
|
p.add_argument(
|
|
"--conda-prefix",
|
|
default="",
|
|
help="optional conda installation prefix; used only when --conda-env is set",
|
|
)
|
|
p.add_argument("--conda-env", default="", help="optional conda env name with torch + mscclpp")
|
|
p.add_argument(
|
|
"--cupti-inproc", action="store_true", help="mscclpp: also collect in-process CUPTI kernel-only timing"
|
|
)
|
|
p.add_argument(
|
|
"--torch-profiler",
|
|
action="store_true",
|
|
help="mscclpp: run the torch.profiler kernel pass (default: host-observed only)",
|
|
)
|
|
p.add_argument(
|
|
"--kernel-only",
|
|
action="store_true",
|
|
help="compare KERNEL execution time only, stripping host/Python launch overhead "
|
|
"(what ep_bench's CUPTI reports). mscclpp uses in-process CUPTI; nccl-ep uses "
|
|
"ep_bench's built-in CUPTI KernelTimer. The unified table then leads with the "
|
|
"kernel dispatch/combine times and a kernel D+C ratio.",
|
|
)
|
|
|
|
# nccl-ep backend.
|
|
p.add_argument(
|
|
"--nccl-lib-path",
|
|
default=os.environ.get("NCCL_LIB_PATH", ""),
|
|
help="directory with libnccl.so / libnccl_ep.so; prepended to LD_LIBRARY_PATH "
|
|
"for ep_bench (falls back to $NCCL_LIB_PATH, else derived from --nccl-ep-bench)",
|
|
)
|
|
p.add_argument(
|
|
"--nccl-ep-bench",
|
|
default="/opt/microsoft/mrc/ep/nccl/build/test/nccl_ep/ep_bench",
|
|
help="path to the NCCL-EP ep_bench binary",
|
|
)
|
|
p.add_argument("--hpcx", default="", help="HPCX install dir (for mpirun); autodetected under /opt if empty")
|
|
p.add_argument(
|
|
"--layout",
|
|
default="em",
|
|
choices=["em", "rm", "fl"],
|
|
help="nccl-ep dispatch layout (em=expert-major, matches mscclpp LL)",
|
|
)
|
|
|
|
# mscclpp-cpp backend (pure C++ MoERuntime binary).
|
|
p.add_argument(
|
|
"--mscclpp-cpp-bench",
|
|
default=os.path.join(_HERE, "build", "mscclpp_ep_bench"),
|
|
help="path to the mscclpp_ep_bench C++ binary (built via test/python/ep/CMakeLists.txt)",
|
|
)
|
|
|
|
p.add_argument("--dry-run", action="store_true", help="print the backend command(s) and exit")
|
|
args = p.parse_args()
|
|
|
|
# These free-form values are interpolated into shell command strings that are
|
|
# executed via bash; constrain them to safe characters to prevent injection
|
|
# and to fail fast on values that would break the launch (spaces, quotes, ...).
|
|
if args.nodes and not re.fullmatch(r"[0-9A-Za-z._:-]+( [0-9A-Za-z._:-]+)*", args.nodes):
|
|
raise SystemExit("--nodes must be space-separated hostnames/IPs")
|
|
if args.iface and not re.fullmatch(r"[0-9A-Za-z._:-]+", args.iface):
|
|
raise SystemExit("--iface must be a valid network interface name")
|
|
if args.hca and not re.fullmatch(r"[0-9A-Za-z._,-]+", args.hca):
|
|
raise SystemExit("--hca must be comma-separated HCA device names")
|
|
if args.conda_env and not args.conda_prefix:
|
|
raise SystemExit("--conda-prefix is required when --conda-env is set")
|
|
if args.dispatch_dtype != "bf16" and args.ep_lib in ("nccl-ep", "both", "all"):
|
|
raise SystemExit("FP8 unified comparison is unsupported because the NCCL-EP command is configured for BF16")
|
|
if args.num_tokens <= 0 or args.num_experts <= 0 or args.nproc_per_node <= 0:
|
|
raise SystemExit("tokens, experts, and nproc-per-node must be positive")
|
|
if args.num_topk > args.num_experts:
|
|
raise SystemExit("num-topk must not exceed num-experts")
|
|
if args.num_warmup < 0 or args.num_iters <= 0:
|
|
raise SystemExit("num-warmup must be non-negative and num-iters must be positive")
|
|
num_nodes = max(1, len(args.nodes.split()))
|
|
num_ranks = num_nodes * args.nproc_per_node
|
|
if args.num_experts % num_ranks != 0:
|
|
raise SystemExit("num-experts must be divisible by the total number of ranks")
|
|
if args.ep_lib in ("mscclpp", "mscclpp-cpp", "both", "all"):
|
|
if not num_ranks + 2 <= args.num_blocks <= 130:
|
|
raise SystemExit("num-blocks must be in [total ranks + 2, 130]")
|
|
|
|
return args
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# Parsing the common "=== Summary (Low Latency ...) ===" block.
|
|
# ----------------------------------------------------------------------------
|
|
@dataclass
|
|
class Phase:
|
|
avg: float = float("nan")
|
|
min: float = float("nan")
|
|
max: float = float("nan")
|
|
|
|
|
|
@dataclass
|
|
class LLResult:
|
|
ep_lib: str
|
|
num_ranks: int = 0
|
|
dispatch: Phase = field(default_factory=Phase)
|
|
combine: Phase = field(default_factory=Phase)
|
|
total: Phase = field(default_factory=Phase)
|
|
# Kernel-only dispatch/combine (avg/min/max) from mscclpp --cupti-inproc or
|
|
# ep_bench's CUPTI KernelTimer, if present.
|
|
kdispatch: Optional[Phase] = None
|
|
kcombine: Optional[Phase] = None
|
|
ok: bool = False
|
|
|
|
|
|
_HOST_RE = {
|
|
"dispatch": re.compile(r"^Dispatch \([^)]+\):\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us"),
|
|
"combine": re.compile(r"^Combine \([^)]+\):\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us"),
|
|
"total": re.compile(r"^Total \(D\+C\):\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us"),
|
|
}
|
|
_RANKS_RE = re.compile(r"=== Summary \(Low Latency, across (\d+) ranks\) ===")
|
|
# Kernel-only Dispatch line, two formats (both carry avg/min/max):
|
|
# mscclpp in-process CUPTI: ``Dispatch: min=M us (representative) [avg=A, max=X us -- ...]``
|
|
# ep_bench CUPTI: ``Dispatch: avg=A us, min=M us, max=X us``
|
|
_KDISP_REP_RE = re.compile(r"^Dispatch:\s+min=([\d.]+)\s*us \(representative\)\s*\[avg=([\d.]+),\s*max=([\d.]+)")
|
|
_KDISP_AMM_RE = re.compile(r"^Dispatch:\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us")
|
|
_KCOMB_REP_RE = re.compile(r"^Combine:\s+min=([\d.]+)\s*us \(representative\)\s*\[avg=([\d.]+),\s*max=([\d.]+)")
|
|
# Kernel-only Combine line (both backends): ``Combine: avg=A us, min=M us, max=X us`` (no ``(BF16)``).
|
|
_KCOMB_RE = re.compile(r"^Combine:\s+avg=([\d.]+)\s*us,\s*min=([\d.]+)\s*us,\s*max=([\d.]+)\s*us")
|
|
|
|
|
|
def parse_ll_summary(text: str, ep_lib: str) -> LLResult:
|
|
res = LLResult(ep_lib=ep_lib)
|
|
for raw in text.splitlines():
|
|
line = raw.strip()
|
|
m = _RANKS_RE.search(line)
|
|
if m:
|
|
res.num_ranks = int(m.group(1))
|
|
continue
|
|
for name, rx in _HOST_RE.items():
|
|
m = rx.match(line)
|
|
if m:
|
|
ph = Phase(float(m.group(1)), float(m.group(2)), float(m.group(3)))
|
|
setattr(res, name, ph)
|
|
# Kernel-only dispatch, first occurrence only. The host lines carry
|
|
# ``(BF16)`` so they never match these bare ``Dispatch:``/``Combine:`` forms.
|
|
if res.kdispatch is None:
|
|
m = _KDISP_REP_RE.match(line)
|
|
if m: # mscclpp: printed order is min, avg, max
|
|
res.kdispatch = Phase(avg=float(m.group(2)), min=float(m.group(1)), max=float(m.group(3)))
|
|
continue
|
|
m = _KDISP_AMM_RE.match(line)
|
|
if m: # ep_bench: printed order is avg, min, max
|
|
res.kdispatch = Phase(avg=float(m.group(1)), min=float(m.group(2)), max=float(m.group(3)))
|
|
continue
|
|
if res.kcombine is None and res.kdispatch is not None:
|
|
m = _KCOMB_REP_RE.match(line)
|
|
if m:
|
|
res.kcombine = Phase(avg=float(m.group(2)), min=float(m.group(1)), max=float(m.group(3)))
|
|
continue
|
|
m = _KCOMB_RE.match(line)
|
|
if m:
|
|
res.kcombine = Phase(avg=float(m.group(1)), min=float(m.group(2)), max=float(m.group(3)))
|
|
res.ok = res.dispatch.avg == res.dispatch.avg # not NaN
|
|
return res
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# Backend command construction.
|
|
# ----------------------------------------------------------------------------
|
|
def build_mscclpp_cmd(args: argparse.Namespace) -> str:
|
|
env_vars = {}
|
|
if args.iface:
|
|
env_vars.update(
|
|
{
|
|
"NCCL_SOCKET_IFNAME": args.iface,
|
|
"GLOO_SOCKET_IFNAME": args.iface,
|
|
"MSCCLPP_SOCKET_IFNAME": args.iface,
|
|
}
|
|
)
|
|
if args.hca:
|
|
env_vars["MSCCLPP_HCA_DEVICES"] = args.hca
|
|
|
|
bench = args.mscclpp_bench
|
|
bench_flags = (
|
|
f"-a ll -t {args.num_tokens} -d {args.hidden} -k {args.num_topk} "
|
|
f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters} "
|
|
f"--dispatch-dtype {args.dispatch_dtype} --combine-mode {args.combine_mode} --num-blocks {args.num_blocks}"
|
|
)
|
|
cupti_build = ""
|
|
extra_exports = ""
|
|
if args.cupti_inproc or args.kernel_only:
|
|
# In-process CUPTI kernel-only timing (near-zero perturbation, matches
|
|
# ep_bench's KernelTimer). Build it under the benchmark build directory.
|
|
bench_flags += " --cupti-inproc"
|
|
cupti_include, cupti_lib = _find_cupti_paths()
|
|
build_dir = os.path.join(os.path.dirname(bench), "build")
|
|
so = os.path.join(build_dir, "libcupti_kernel_timer.so")
|
|
src = os.path.join(os.path.dirname(bench), "cupti_kernel_timer.cpp")
|
|
env_vars["MSCCLPP_EP_CUPTI_TIMER_LIB"] = so
|
|
cupti_build = (
|
|
f"mkdir -p {shlex.quote(build_dir)} && "
|
|
f"if [ ! -f {shlex.quote(so)} ] || [ {shlex.quote(src)} -nt {shlex.quote(so)} ]; then "
|
|
f"g++ -O2 -fPIC -shared {shlex.quote(src)} -o {shlex.quote(so)} "
|
|
f"-I{shlex.quote(cupti_include)} -L{shlex.quote(cupti_lib)} -lcupti; fi && "
|
|
)
|
|
extra_exports = f"export LD_LIBRARY_PATH={shlex.quote(cupti_lib)}:${{LD_LIBRARY_PATH:-}} && "
|
|
elif args.torch_profiler:
|
|
# Opt-in torch.profiler kernel pass (perturbs the LL recv-spin; the
|
|
# in-process CUPTI path is preferred for kernel numbers).
|
|
pass
|
|
else:
|
|
# Default: clean host-observed only (skip the torch.profiler pass, which
|
|
# is slow and inflates the LL dispatch recv-spin).
|
|
bench_flags += " --no-kernel-timing"
|
|
|
|
activation = ""
|
|
python = shlex.quote(args.python)
|
|
if args.conda_env:
|
|
activation = (
|
|
f"source {shlex.quote(args.conda_prefix)}/etc/profile.d/conda.sh && "
|
|
f"conda activate {shlex.quote(args.conda_env)} && "
|
|
)
|
|
python = "python"
|
|
exports = " ".join(f"{name}={shlex.quote(value)}" for name, value in env_vars.items())
|
|
env_export = f"export {exports} && " if exports else ""
|
|
return (
|
|
f"{activation}"
|
|
f"{cupti_build}"
|
|
f"{env_export}"
|
|
f"{extra_exports}"
|
|
f"{python} -m torch.distributed.run --standalone --nnodes=1 --nproc_per_node={args.nproc_per_node} "
|
|
f"{shlex.quote(bench)} {bench_flags}"
|
|
)
|
|
|
|
|
|
def _autodetect_hpcx() -> str:
|
|
import glob
|
|
|
|
cands = sorted(glob.glob("/opt/hpcx-*"))
|
|
return cands[0] if cands else ""
|
|
|
|
|
|
def _mpi_launch(args, np_total):
|
|
"""Common mpirun prefix. Multi-node when --nodes lists >1 IP (writes a
|
|
hostfile, adds an SSH launcher); otherwise a plain single-node launch."""
|
|
nodes = args.nodes.split()
|
|
setup = ""
|
|
hostfile = ""
|
|
if len(nodes) > 1:
|
|
slots = args.nproc_per_node
|
|
lines = "\\n".join(f"{ip} slots={slots}" for ip in nodes)
|
|
hf = "/tmp/ep_unified_hostfile"
|
|
setup = f"printf '{lines}\\n' > {hf} && "
|
|
hostfile = (
|
|
f"--hostfile {hf} " f'-mca plm_rsh_args "-o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null" '
|
|
)
|
|
iface_arg = f"-mca btl_tcp_if_include {shlex.quote(args.iface)} " if args.iface else ""
|
|
root_arg = "--allow-run-as-root " if os.geteuid() == 0 else ""
|
|
return setup, (
|
|
f"mpirun {root_arg}-np {np_total} {hostfile}--map-by ppr:{args.nproc_per_node}:node --bind-to none "
|
|
f"-mca pml ob1 -mca btl self,vader,tcp {iface_arg}"
|
|
f"-mca coll_hcoll_enable 0 -mca coll_ucc_enable 0 "
|
|
)
|
|
|
|
|
|
def build_nccl_ep_cmd(args: argparse.Namespace) -> str:
|
|
nccl_lib = args.nccl_lib_path
|
|
if not nccl_lib:
|
|
# Derive the libnccl / libnccl_ep directory from the ep_bench binary
|
|
# instead of hard-coding it: <nccl>/build/test/nccl_ep/ep_bench ->
|
|
# <nccl>/build/lib.
|
|
bench_dir = os.path.dirname(os.path.abspath(args.nccl_ep_bench))
|
|
nccl_lib = os.path.join(os.path.dirname(os.path.dirname(bench_dir)), "lib")
|
|
hpcx = args.hpcx or _autodetect_hpcx()
|
|
nodes = args.nodes.split()
|
|
nnodes = max(1, len(nodes))
|
|
np_total = nnodes * args.nproc_per_node
|
|
mnnvl = 1 if nnodes > 1 else 0
|
|
bench_flags = (
|
|
f"-a ll -L {args.layout} -t {args.num_tokens} -d {args.hidden} -k {args.num_topk} "
|
|
f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters}"
|
|
)
|
|
setup, mpi_prefix = _mpi_launch(args, np_total)
|
|
opal = f"-x OPAL_PREFIX={shlex.quote(hpcx)}/ompi " if hpcx else ""
|
|
iface_env = (
|
|
f"-x UCX_NET_DEVICES={shlex.quote(args.iface)} " f"-x NCCL_SOCKET_IFNAME={shlex.quote(args.iface)} "
|
|
if args.iface
|
|
else ""
|
|
)
|
|
mpi = (
|
|
f"{mpi_prefix}"
|
|
f"-x LD_LIBRARY_PATH -x PATH -x CUDA_HOME={shlex.quote(CUDA_HOME)} {opal}"
|
|
f"{iface_env}-x UCX_TLS=tcp,sm,self,cuda_copy -x UCX_HANDLE_ERRORS=none "
|
|
f"-x NCCL_NET_PLUGIN=none "
|
|
f"-x NCCL_IB_DISABLE=1 -x NCCL_MNNVL_ENABLE={mnnvl} "
|
|
f"{shlex.quote(args.nccl_ep_bench)} {bench_flags}"
|
|
)
|
|
activation = f"source {shlex.quote(hpcx)}/hpcx-init.sh && hpcx_load && " if hpcx else ""
|
|
return f"{activation}" f"export LD_LIBRARY_PATH={shlex.quote(nccl_lib)}:$LD_LIBRARY_PATH && " f"{setup}{mpi}"
|
|
|
|
|
|
def build_mscclpp_cpp_cmd(args: argparse.Namespace) -> str:
|
|
"""Pure-C++ mscclpp_ep_bench (MoERuntime), launched with mpirun -- no Python."""
|
|
hpcx = args.hpcx or _autodetect_hpcx()
|
|
nodes = args.nodes.split()
|
|
nnodes = max(1, len(nodes))
|
|
np_total = nnodes * args.nproc_per_node
|
|
bench_flags = (
|
|
f"-a ll -t {args.num_tokens} -d {args.hidden} -k {args.num_topk} "
|
|
f"-e {args.num_experts} -w {args.num_warmup} -i {args.num_iters} "
|
|
f"--dispatch-dtype {args.dispatch_dtype} --combine-mode {args.combine_mode} --num-blocks {args.num_blocks}"
|
|
)
|
|
if args.kernel_only or args.cupti_inproc:
|
|
bench_flags += " --kernel-timing"
|
|
setup, mpi_prefix = _mpi_launch(args, np_total)
|
|
env_exports = ""
|
|
if args.hca:
|
|
env_exports += f"-x MSCCLPP_HCA_DEVICES={shlex.quote(args.hca)} "
|
|
if args.iface:
|
|
env_exports += (
|
|
f"-x NCCL_SOCKET_IFNAME={shlex.quote(args.iface)} " f"-x MSCCLPP_SOCKET_IFNAME={shlex.quote(args.iface)} "
|
|
)
|
|
mpi = (
|
|
f"{mpi_prefix}"
|
|
f"-x LD_LIBRARY_PATH -x PATH "
|
|
f"{env_exports}"
|
|
f"{shlex.quote(args.mscclpp_cpp_bench)} {bench_flags}"
|
|
)
|
|
_, cupti_lib = _find_cupti_paths()
|
|
activation = f"source {shlex.quote(hpcx)}/hpcx-init.sh && hpcx_load && " if hpcx else ""
|
|
return f"{activation}" f"export LD_LIBRARY_PATH={shlex.quote(cupti_lib)}:$LD_LIBRARY_PATH && " f"{setup}{mpi}"
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# Run + report.
|
|
# ----------------------------------------------------------------------------
|
|
def run_backend(ep_lib: str, cmd: str, dry_run: bool) -> Optional[LLResult]:
|
|
print(f"\n########## ep-lib={ep_lib} ##########", flush=True)
|
|
print(f"$ {cmd}\n", flush=True)
|
|
if dry_run:
|
|
return None
|
|
proc = subprocess.run(["bash", "-lc", cmd], capture_output=True, text=True)
|
|
sys.stdout.write(proc.stdout)
|
|
if proc.returncode != 0:
|
|
sys.stderr.write(proc.stderr[-4000:])
|
|
print(f"[warn] {ep_lib} exited rc={proc.returncode}", flush=True)
|
|
res = parse_ll_summary(proc.stdout, ep_lib)
|
|
if not res.ok:
|
|
print(f"[warn] could not parse a Low-Latency summary from {ep_lib} output", flush=True)
|
|
return None
|
|
return res
|
|
|
|
|
|
def print_unified(results: list, kernel_only: bool = False) -> None:
|
|
results = [r for r in results if r is not None]
|
|
if not results:
|
|
return
|
|
has_kernel = all(r.kdispatch is not None and r.kcombine is not None for r in results)
|
|
title = "kernel-only" if (kernel_only and has_kernel) else "host-observed"
|
|
print(f"\n=== Unified EP Low-Latency Summary ({title}, us) ===")
|
|
hdr = f"{'metric':<24}" + "".join(f"{r.ep_lib:>14}" for r in results)
|
|
print(hdr)
|
|
print("-" * len(hdr))
|
|
|
|
def row(label, fn):
|
|
print(f"{label:<24}" + "".join(f"{fn(r):>14.2f}" for r in results))
|
|
|
|
if not (kernel_only and has_kernel):
|
|
# Host-observed dispatch/combine/total, full avg/min/max.
|
|
row("Host Dispatch avg", lambda r: r.dispatch.avg)
|
|
row("Host Dispatch min", lambda r: r.dispatch.min)
|
|
row("Host Dispatch max", lambda r: r.dispatch.max)
|
|
row("Host Combine avg", lambda r: r.combine.avg)
|
|
row("Host Combine min", lambda r: r.combine.min)
|
|
row("Host Combine max", lambda r: r.combine.max)
|
|
row("Host D+C avg", lambda r: r.total.avg)
|
|
if has_kernel:
|
|
if kernel_only:
|
|
row("Kernel Dispatch repr", lambda r: r.kdispatch.min)
|
|
row("Kernel Combine repr", lambda r: r.kcombine.min)
|
|
else:
|
|
row("Kernel Dispatch avg", lambda r: r.kdispatch.avg)
|
|
row("Kernel Dispatch min", lambda r: r.kdispatch.min)
|
|
row("Kernel Dispatch max", lambda r: r.kdispatch.max)
|
|
row("Kernel Combine avg", lambda r: r.kcombine.avg)
|
|
row("Kernel Combine min", lambda r: r.kcombine.min)
|
|
row("Kernel Combine max", lambda r: r.kcombine.max)
|
|
elif kernel_only:
|
|
print(
|
|
" NOTE: kernel-only requested but kernel timing missing for a backend "
|
|
"(mscclpp needs --cupti-inproc / libcupti; nccl-ep needs CUPTI-enabled ep_bench)."
|
|
)
|
|
if len(results) == 2 and not kernel_only:
|
|
a, b = results
|
|
if a.total.avg == a.total.avg and b.total.avg == b.total.avg and b.total.avg:
|
|
print(f"\nHost D+C ratio {a.ep_lib}/{b.ep_lib} = {a.total.avg / b.total.avg:.2f}x")
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
if args.ep_lib == "both":
|
|
libs = ["mscclpp", "nccl-ep"]
|
|
elif args.ep_lib == "all":
|
|
libs = ["mscclpp", "mscclpp-cpp", "nccl-ep"]
|
|
else:
|
|
libs = [args.ep_lib]
|
|
|
|
builders = {
|
|
"mscclpp": build_mscclpp_cmd,
|
|
"mscclpp-cpp": build_mscclpp_cpp_cmd,
|
|
"nccl-ep": build_nccl_ep_cmd,
|
|
}
|
|
if len(args.nodes.split()) > 1 and "mscclpp" in libs:
|
|
print(
|
|
"[warn] --nodes multi-node ignored for the Python 'mscclpp' backend "
|
|
"(torchrun --standalone is single-node); use mscclpp-cpp for multi-node.",
|
|
flush=True,
|
|
)
|
|
results = []
|
|
for lib in libs:
|
|
cmd = builders[lib](args)
|
|
results.append(run_backend(lib, cmd, args.dry_run))
|
|
if not args.dry_run:
|
|
print_unified(results, kernel_only=args.kernel_only)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|