385 lines
12 KiB
Python
385 lines
12 KiB
Python
"""
|
|
Benchmark JIT custom all-reduce (v2) vs NCCL vs AOT custom all-reduce (v1).
|
|
|
|
Usage (torchrun required for multi-GPU):
|
|
torchrun --nproc_per_node=2 bench_custom_all_reduce.py
|
|
torchrun --nproc_per_node=4 bench_custom_all_reduce.py --dtype float16
|
|
torchrun --nproc_per_node=8 bench_custom_all_reduce.py --warmup 10 --iters 100
|
|
|
|
The script initializes all three backends, then benchmarks each over a sweep
|
|
of message sizes. Results are printed as a comparison table on rank 0.
|
|
"""
|
|
|
|
import argparse
|
|
import contextlib
|
|
import gc
|
|
import logging
|
|
import os
|
|
from math import isnan
|
|
from typing import Dict, List, Optional
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from sglang.jit_kernel.benchmark.utils import is_in_ci
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
|
|
register_cuda_ci(
|
|
est_time=120,
|
|
suite="stage-b-kernel-benchmark-1-gpu-large",
|
|
disabled="requires multi-GPU, self-skips in CI",
|
|
)
|
|
|
|
DTYPE_MAP = {
|
|
"float16": torch.float16,
|
|
"bfloat16": torch.bfloat16,
|
|
"float32": torch.float32,
|
|
}
|
|
|
|
MESSAGE_SIZES_BYTES = [
|
|
4 * 1024, # 4K
|
|
16 * 1024, # 16K
|
|
64 * 1024, # 64K
|
|
128 * 1024, # 128K
|
|
3 * 64 * 1024, # 192K
|
|
4 * 64 * 1024, # 256K
|
|
3 * 128 * 1024, # 384K
|
|
4 * 128 * 1024, # 512K
|
|
5 * 128 * 1024, # 640K
|
|
6 * 128 * 1024, # 768K
|
|
7 * 128 * 1024, # 896K
|
|
1 * 1024 * 1024, # 1M
|
|
2 * 1024 * 1024, # 2M
|
|
3 * 1024 * 1024, # 2M
|
|
4 * 1024 * 1024, # 4M
|
|
8 * 1024 * 1024, # 8M
|
|
16 * 1024 * 1024, # 16M
|
|
32 * 1024 * 1024, # 32M
|
|
]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Backend wrappers - each exposes a uniform interface:
|
|
# .name - display name
|
|
# .capture() - context manager for CUDA-graph recording
|
|
# .all_reduce() - perform an all-reduce and return the result tensor
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class NCCLAllReduceBackend:
|
|
name = "NCCL"
|
|
|
|
def __init__(self, group: dist.ProcessGroup):
|
|
self.group = group
|
|
|
|
def capture(self, register_input: bool):
|
|
return contextlib.nullcontext()
|
|
|
|
def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
|
|
dist.all_reduce(tensor, group=self.group)
|
|
return tensor
|
|
|
|
|
|
class AOTAllReduceBackend:
|
|
name = "AOT"
|
|
|
|
def __init__(self, group: dist.ProcessGroup, device: torch.device):
|
|
from sglang.srt.distributed.device_communicators.custom_all_reduce import (
|
|
CustomAllreduce,
|
|
)
|
|
|
|
max_size = max(MESSAGE_SIZES_BYTES)
|
|
self.comm = CustomAllreduce(group, device, max_size=max_size)
|
|
if self.comm.disabled:
|
|
raise RuntimeError("AOT CustomAllreduce is disabled on this system")
|
|
|
|
def capture(self, register_input: bool):
|
|
return self.comm.capture() # ignore register_input since v1 always requires it
|
|
|
|
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
|
assert self.comm.should_custom_ar(tensor), str(tensor.shape)
|
|
return self.comm.custom_all_reduce(tensor)
|
|
|
|
|
|
class JITAllReduceBackend:
|
|
name = "JIT"
|
|
|
|
def __init__(self, group: dist.ProcessGroup, device: torch.device):
|
|
from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
|
|
CustomAllReduceV2,
|
|
)
|
|
|
|
max_size = max(MESSAGE_SIZES_BYTES)
|
|
self.comm = CustomAllReduceV2(group, device, max_pull_size=max_size)
|
|
if self.comm.disabled:
|
|
raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
|
|
|
|
def capture(self, register_input: bool):
|
|
return self.comm.capture() if register_input else contextlib.nullcontext()
|
|
|
|
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
|
assert self.comm.should_custom_ar(tensor), str(tensor.shape)
|
|
return self.comm.custom_all_reduce(tensor)
|
|
|
|
|
|
class FlashInferAllReduceBackend:
|
|
name = "FI"
|
|
|
|
def __init__(self, group: dist.ProcessGroup, dtype: torch.dtype):
|
|
import flashinfer.comm as comm
|
|
|
|
rank = torch.distributed.get_rank(group=group)
|
|
world_size = torch.distributed.get_world_size(group=group)
|
|
max_size = max(MESSAGE_SIZES_BYTES)
|
|
hidden_dim = min(MESSAGE_SIZES_BYTES) // 2
|
|
num_tokens = max_size // hidden_dim
|
|
self.comm = comm
|
|
self.hidden_dim = hidden_dim
|
|
self.workspace = comm.create_allreduce_fusion_workspace(
|
|
backend="trtllm",
|
|
world_size=world_size,
|
|
rank=rank,
|
|
max_token_num=num_tokens,
|
|
hidden_dim=hidden_dim,
|
|
dtype=dtype,
|
|
)
|
|
|
|
def capture(self, *_):
|
|
return contextlib.nullcontext()
|
|
|
|
def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
|
|
return self.comm.allreduce_fusion(
|
|
input=tensor.view(-1, self.hidden_dim),
|
|
workspace=self.workspace,
|
|
pattern=self.comm.AllReduceFusionPattern.kAllReduce,
|
|
launch_with_pdl=True,
|
|
fp32_acc=True,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Benchmarking helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def parse_args():
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument("--dtype", choices=DTYPE_MAP.keys(), default="bfloat16")
|
|
p.add_argument("--warmup", type=int, default=5)
|
|
p.add_argument("--iters", type=int, default=50)
|
|
p.add_argument("--no-inplace", dest="register_input", action="store_false")
|
|
return p.parse_args()
|
|
|
|
|
|
@torch.inference_mode()
|
|
def bench_one(
|
|
backend,
|
|
inp: torch.Tensor,
|
|
warmup: int,
|
|
iters: int,
|
|
group: dist.ProcessGroup,
|
|
register_input: bool,
|
|
) -> float:
|
|
"""
|
|
Run *warmup* iterations of all-reduce first.
|
|
Return the average time for *iters* iterations of all-reduce.
|
|
"""
|
|
dist.barrier(group=group)
|
|
for _ in range(warmup):
|
|
backend.all_reduce(inp)
|
|
torch.cuda.synchronize()
|
|
|
|
# Capture a CUDA graph with *iters* all-reduce calls.
|
|
inp_batch = torch.stack([inp] * 4)
|
|
graph = torch.cuda.CUDAGraph()
|
|
with backend.capture(register_input):
|
|
with torch.cuda.graph(graph):
|
|
for i in range(iters):
|
|
backend.all_reduce(inp_batch[i % 4])
|
|
|
|
torch.cuda.synchronize()
|
|
# Warm up the graph once.
|
|
graph.replay()
|
|
|
|
# Timed replay.
|
|
start = torch.cuda.Event(enable_timing=True)
|
|
end = torch.cuda.Event(enable_timing=True)
|
|
torch.cuda.synchronize()
|
|
dist.barrier(group=group)
|
|
graph.replay() # make the stream busy
|
|
start.record()
|
|
graph.replay()
|
|
end.record()
|
|
torch.cuda.synchronize()
|
|
return start.elapsed_time(end) / iters
|
|
|
|
|
|
def bench_sweep(
|
|
backend,
|
|
sizes_bytes: List[int],
|
|
dtype: torch.dtype,
|
|
device: torch.device,
|
|
warmup: int,
|
|
iters: int,
|
|
group: dist.ProcessGroup,
|
|
register_input: bool,
|
|
) -> Dict[int, float]:
|
|
"""Benchmark one backend over all message sizes."""
|
|
elem_size = torch.tensor([], dtype=dtype).element_size()
|
|
results: Dict[int, float] = {}
|
|
for sz in sizes_bytes:
|
|
numel = sz // elem_size
|
|
inp = torch.zeros(numel, dtype=dtype, device=device)
|
|
try:
|
|
elapsed_ms = bench_one(backend, inp, warmup, iters, group, register_input)
|
|
results[sz] = elapsed_ms * 1000 # convert to us per iter
|
|
except AssertionError:
|
|
results[sz] = float("nan")
|
|
return results
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Result printing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def print_results(
|
|
backends: list,
|
|
all_results: Dict[str, Dict[int, float]],
|
|
sizes_bytes: List[int],
|
|
) -> None:
|
|
"""Print a comparison table on rank 0."""
|
|
|
|
def human_bytes(n: int) -> str:
|
|
for suffix, unit in [("M", 1 << 20), ("K", 1 << 10)]:
|
|
if n >= unit and n % unit == 0:
|
|
return f"{n // unit}{suffix}"
|
|
return f"{n}B"
|
|
|
|
def fmt_us(v: float) -> str:
|
|
return f"{v:13.1f}" if not isnan(v) else " n/a"
|
|
|
|
names = [b.name for b in backends]
|
|
nccl_name = "NCCL"
|
|
|
|
# Header
|
|
header_cols = [f"{n:>13}" for n in names]
|
|
speedup_cols = [f"{n:>13}/NCCL" for n in names if n != nccl_name]
|
|
header = f"{'Size':>8} " + " ".join(header_cols)
|
|
for sc in speedup_cols:
|
|
header += f" {sc}"
|
|
header += " "
|
|
print()
|
|
print(header)
|
|
print("-" * len(header))
|
|
|
|
# Rows
|
|
for sz in sizes_bytes:
|
|
row = f"{human_bytes(sz):>8}"
|
|
nccl_lat = all_results[nccl_name][sz]
|
|
for n in names:
|
|
row += f" {fmt_us(all_results[n][sz])}"
|
|
for n in names:
|
|
if n == nccl_name:
|
|
continue
|
|
lat = all_results[n][sz]
|
|
if not isnan(lat):
|
|
row += f" {nccl_lat / lat:17.2f}x"
|
|
else:
|
|
row += f" {'n/a':>17}"
|
|
print(row)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Distributed setup
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def init_distributed():
|
|
"""Initialize distributed groups using torchrun env vars.
|
|
|
|
Returns (rank, world_size, device, cpu_group, nccl_group).
|
|
"""
|
|
import sglang.srt.distributed.parallel_state as ps
|
|
|
|
local_rank = int(os.environ.get("LOCAL_RANK", "0"))
|
|
world_size = int(os.environ.get("WORLD_SIZE", "1"))
|
|
rank = local_rank
|
|
device = torch.device(f"cuda:{rank}")
|
|
torch.cuda.set_device(device)
|
|
torch.cuda.set_stream(torch.cuda.Stream()) # use a non-default stream
|
|
|
|
torch.distributed.init_process_group(backend="gloo")
|
|
ps._WORLD = coord = ps.init_world_group(
|
|
ranks=list(range(world_size)),
|
|
local_rank=local_rank,
|
|
backend="nccl",
|
|
)
|
|
|
|
cpu_group = coord.cpu_group
|
|
nccl_group = coord.device_group
|
|
assert nccl_group is not None
|
|
return rank, world_size, device, cpu_group, nccl_group
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def main():
|
|
logging.basicConfig(level=logging.WARNING)
|
|
args = parse_args()
|
|
dtype = DTYPE_MAP[args.dtype]
|
|
|
|
rank, world_size, device, cpu_group, nccl_group = init_distributed()
|
|
|
|
# Instantiate backends.
|
|
backends = [
|
|
NCCLAllReduceBackend(nccl_group),
|
|
JITAllReduceBackend(cpu_group, device),
|
|
]
|
|
if world_size in [2, 4, 6, 8]:
|
|
backends.insert(1, AOTAllReduceBackend(cpu_group, device))
|
|
if world_size in [2, 4, 8]:
|
|
backends.append(FlashInferAllReduceBackend(cpu_group, dtype))
|
|
|
|
# Run benchmarks.
|
|
all_results: Dict[str, Dict[int, float]] = {}
|
|
torch.cuda.synchronize()
|
|
for backend in backends:
|
|
if rank == 0:
|
|
print(f"Benchmarking {backend.name} ...")
|
|
all_results[backend.name] = bench_sweep(
|
|
backend,
|
|
MESSAGE_SIZES_BYTES,
|
|
dtype,
|
|
device,
|
|
args.warmup,
|
|
args.iters,
|
|
cpu_group,
|
|
args.register_input,
|
|
)
|
|
|
|
# Aggregate across ranks (use max to reflect the slowest rank).
|
|
for name in list(all_results):
|
|
for sz in MESSAGE_SIZES_BYTES:
|
|
val = all_results[name].get(sz)
|
|
if val is None:
|
|
continue
|
|
t = torch.tensor([val], dtype=torch.float64, device=device)
|
|
dist.all_reduce(t, op=dist.ReduceOp.MAX, group=nccl_group)
|
|
all_results[name][sz] = t.item()
|
|
|
|
# Print results on rank 0.
|
|
if rank == 0:
|
|
print_results(backends, all_results, MESSAGE_SIZES_BYTES)
|
|
|
|
del backends, all_results
|
|
gc.collect()
|
|
dist.destroy_process_group()
|
|
|
|
|
|
if __name__ == "__main__" and not is_in_ci():
|
|
main()
|