mirror of
https://github.com/microsoft/mscclpp.git
synced 2026-07-15 19:54:45 +00:00
ad0b63d21b
Refresh the Python and C++ benchmark paths for BF16 and FP8 dispatch, current MoERuntime signatures, active kernel sources, portable CUPTI discovery, realistic routing, and safe unified reporting. Remove the merged change to the inactive legacy implementation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: efbacae6-f679-430b-bc16-b45ae162fc76
575 lines
25 KiB
Python
575 lines
25 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 = {
|
|
"MSCCLPP_EP_LOCAL_WORLD_SIZE": str(args.nproc_per_node),
|
|
"NCCL_IB_DISABLE": "1",
|
|
"NCCL_MNNVL_ENABLE": "0",
|
|
}
|
|
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())
|
|
return (
|
|
f"{activation}"
|
|
f"{cupti_build}"
|
|
f"export {exports} && "
|
|
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 = (
|
|
f"-x MSCCLPP_EP_LOCAL_WORLD_SIZE={args.nproc_per_node} " f"-x NCCL_IB_DISABLE=1 -x NCCL_MNNVL_ENABLE=0 "
|
|
)
|
|
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()
|