diff --git a/benchmark/kernels/all_reduce/benchmark_all_reduce.py b/benchmark/kernels/all_reduce/benchmark_all_reduce.py new file mode 100644 index 000000000..30a7733f3 --- /dev/null +++ b/benchmark/kernels/all_reduce/benchmark_all_reduce.py @@ -0,0 +1,351 @@ +""" +Benchmark SGLang custom all-reduce vs Torch symm-mem all-reduce across message sizes. +Usage: + torchrun --nproc_per_node=2 benchmark_all_reduce.py + torchrun --nproc_per_node=4 benchmark_all_reduce.py + torchrun --nproc_per_node=8 benchmark_all_reduce.py +""" + +import argparse +import os +import sys +import time +from typing import List, Optional, Tuple + +import torch +import torch.distributed as dist + +from sglang.srt.distributed.parallel_state import ( + destroy_distributed_environment, + destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel, +) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Benchmark SGLang custom all-reduce vs Torch symm-mem all-reduce across message sizes." + ) + parser.add_argument( + "--backend", + type=str, + default="gloo", + help="Process group backend for the custom-AR control path (must NOT be nccl).", + ) + parser.add_argument( + "--warmup", + type=int, + default=5, + help="Warmup iterations per size per implementation.", + ) + parser.add_argument( + "--iters-small", + type=int, + default=50, + help="Benchmark iterations for sizes <= 1MB.", + ) + parser.add_argument( + "--iters-large", + type=int, + default=20, + help="Benchmark iterations for sizes > 1MB.", + ) + parser.add_argument( + "--verbose", + action="store_true", + help="Print per-iteration timings on rank 0 for debugging.", + ) + return parser.parse_args() + + +def get_env_rank_world() -> Tuple[int, int, int]: + rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + local_rank = int(os.environ.get("LOCAL_RANK", str(rank))) + return rank, world_size, local_rank + + +def init_dist(backend: str): + rank, world_size, _ = get_env_rank_world() + if not dist.is_initialized(): + dist.init_process_group( + backend=backend, + init_method="env://", + rank=rank, + world_size=world_size, + ) + + device = torch.device(f"cuda:{rank}") + torch.cuda.set_device(device) + distributed_init_method = f"tcp://localhost:23456" + init_distributed_environment( + world_size=world_size, + rank=rank, + distributed_init_method=distributed_init_method, + local_rank=rank, + ) + initialize_model_parallel(tensor_model_parallel_size=world_size) + return dist.group.WORLD + + +def get_device(local_rank: int) -> torch.device: + torch.cuda.set_device(local_rank) + return torch.device(f"cuda:{local_rank}") + + +def human_size(num_bytes: int) -> str: + units = [("B", 1), ("K", 1024), ("M", 1024 * 1024), ("G", 1024 * 1024 * 1024)] + for suf, base in reversed(units): + if num_bytes % base == 0 and num_bytes >= base: + val = num_bytes // base + return f"{val}{suf}" + return f"{num_bytes}B" + + +def get_message_sizes() -> List[int]: + return [ + 32 * 1024, + 64 * 1024, + 128 * 1024, + 256 * 1024, + 512 * 1024, + 1 * 1024 * 1024, + 2 * 1024 * 1024, + 4 * 1024 * 1024, + 8 * 1024 * 1024, + 16 * 1024 * 1024, + 32 * 1024 * 1024, + 64 * 1024 * 1024, + ] + + +@torch.inference_mode() +def run_once(comm, inp: torch.Tensor) -> Optional[torch.Tensor]: + if hasattr(comm, "custom_all_reduce"): + return comm.custom_all_reduce(inp) + if hasattr(comm, "all_reduce"): + return comm.all_reduce(inp) + raise RuntimeError("No known all-reduce method found on the communicator.") + + +@torch.inference_mode() +def bench_impl( + name: str, + comm, + sizes: List[int], + device: torch.device, + warmup: int, + iters_small: int, + iters_large: int, + verbose: bool, + pg: Optional[dist.ProcessGroup] = None, +) -> List[Tuple[int, Optional[float]]]: + rank = dist.get_rank() + world_size = dist.get_world_size() + results: List[Tuple[int, Optional[float]]] = [] + + for size_bytes in sizes: + elems = size_bytes // 2 # float16: 2 bytes per element + inp = torch.empty(elems, dtype=torch.float16, device=device) + inp.uniform_(0, 1) + + disabled = False + dist.barrier(group=pg) + for _ in range(warmup): + torch.cuda.synchronize() + out = run_once(comm, inp) + torch.cuda.synchronize() + if out is None: + disabled = True + break + dist.barrier(group=pg) + + if disabled: + if rank == 0: + print( + f"[{name}] {human_size(size_bytes)}: custom AR disabled (skipped)" + ) + results.append((size_bytes, None)) + continue + + num_iters = iters_small if size_bytes <= (1 * 1024 * 1024) else iters_large + + times_ms: List[float] = [] + for it in range(num_iters): + dist.barrier(group=pg) + torch.cuda.synchronize() + t0 = time.perf_counter() + out = run_once(comm, inp) + torch.cuda.synchronize() + t1 = time.perf_counter() + dist.barrier(group=pg) + + if out is None: + disabled = True + break + + dt_ms = (t1 - t0) * 1000.0 + times_ms.append(dt_ms) + + if verbose and rank == 0: + print( + f"[{name}] size={human_size(size_bytes)} iter={it} time={dt_ms:.3f} ms" + ) + + if disabled or not times_ms: + if rank == 0: + print( + f"[{name}] {human_size(size_bytes)}: custom AR disabled (no timings)" + ) + results.append((size_bytes, None)) + continue + + avg_ms_local = sum(times_ms) / len(times_ms) + avg_tensor = torch.tensor([avg_ms_local], dtype=torch.float64, device=device) + gather_list = [torch.zeros_like(avg_tensor) for _ in range(world_size)] + dist.all_gather(gather_list, avg_tensor, group=pg) + if rank == 0: + avg_ms = float(torch.stack(gather_list).mean().item()) + print( + f"[{name}] {human_size(size_bytes)}: {avg_ms:.3f} ms (avg across ranks)" + ) + results.append((size_bytes, avg_ms)) + else: + results.append((size_bytes, None)) + + return results + + +def main(): + args = parse_args() + rank, world_size, local_rank = get_env_rank_world() + + if world_size not in (2, 4, 6, 8): + print( + f"[rank {rank}] WARNING: world_size={world_size} not in supported set (2,4,6,8). " + "Custom AR may disable itself.", + file=sys.stderr, + ) + + group = init_dist(args.backend) + device = get_device(local_rank) + + # Import after dist init; some libs query torch dist state on import + torch_symm_mem_comm = None + HAVE_SGLANG_CUSTOM = False + HAVE_TORCH_SYMM_MEM = False + + try: + from sglang.srt.distributed.device_communicators.custom_all_reduce import ( + CustomAllreduce as SGLCustomAllreduce, + ) + + HAVE_SGLANG_CUSTOM = True + except Exception as e: + if rank == 0: + print(f"SGLang CustomAllreduce import failed: {e}", file=sys.stderr) + + try: + from sglang.srt.distributed.device_communicators.torch_symm_mem import ( + TorchSymmMemCommunicator as TorchSymmMemAllreduce, + ) + + HAVE_TORCH_SYMM_MEM = True + except Exception as e: + if rank == 0: + print(f"TorchSymmMemAllreduce import failed: {e}", file=sys.stderr) + + if rank == 0: + print(f"Initialized PG backend={args.backend} world_size={world_size}") + print(f"Device: {device.type}:{device.index}") + print( + f"SGLang Custom available: {HAVE_SGLANG_CUSTOM}, Torch Symm-Mem available: {HAVE_TORCH_SYMM_MEM}" + ) + + sizes = get_message_sizes() + max_size = max(sizes) if sizes else (128 * 1024 * 1024) + + if HAVE_SGLANG_CUSTOM: + try: + sgl_custom_comm = SGLCustomAllreduce( + group=group, device=device, max_size=max_size + ) + except Exception as e: + if rank == 0: + print( + f"Failed to construct SGLangCustomAllreduce: {e}", file=sys.stderr + ) + sgl_custom_comm = None + + if HAVE_TORCH_SYMM_MEM: + try: + torch_symm_mem_comm = TorchSymmMemAllreduce(group=group, device=device) + except Exception as e: + if rank == 0: + print( + f"Failed to construct TorchSymmMemAllreduce: {e}", file=sys.stderr + ) + torch_symm_mem_comm = None + + sgl_custom_results: List[Tuple[int, Optional[float]]] = [] + symm_mem_results: List[Tuple[int, Optional[float]]] = [] + + if sgl_custom_comm is not None: + sgl_custom_results = bench_impl( + name="SGLangCustom", + comm=sgl_custom_comm, + sizes=sizes, + device=device, + warmup=args.warmup, + iters_small=args.iters_small, + iters_large=args.iters_large, + verbose=args.verbose, + pg=group, + ) + + if torch_symm_mem_comm is not None: + symm_mem_results = bench_impl( + name="TorchSymmMem", + comm=torch_symm_mem_comm, + sizes=sizes, + device=device, + warmup=args.warmup, + iters_small=args.iters_small, + iters_large=args.iters_large, + verbose=args.verbose, + pg=group, + ) + + for comm in (sgl_custom_comm, torch_symm_mem_comm): + if comm is not None and hasattr(comm, "close"): + try: + comm.close() + except Exception: + pass + + if dist.get_rank() == 0: + print( + f"\nResults (avg ms across {world_size} ranks; None = disabled/unavailable):" + ) + header = f"{'Size':>8} {'CustomAR(ms)':>12} {'TorchSymmMem(ms)':>11}" + print(header) + print("-" * len(header)) + + sgl_custom_map = {s: v for s, v in sgl_custom_results if v is not None} + symm_mem_map = {s: v for s, v in symm_mem_results if v is not None} + + for s in sizes: + sgl_ms = sgl_custom_map.get(s, None) + symm_mem_ms = symm_mem_map.get(s, None) + print( + f"{human_size(s):>8} {('%.3f' % sgl_ms) if sgl_ms is not None else 'None':>12} " + f"{('%.3f' % symm_mem_ms) if symm_mem_ms is not None else 'None':>11}" + ) + torch.distributed.barrier(group=group) + destroy_model_parallel() + destroy_distributed_environment() + + +if __name__ == "__main__": + main() diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index 0d5f612ce..e71f93ebc 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -325,12 +325,12 @@ class CustomAllreduce: # little performance improvement over NCCL. if not _is_hip: if self.world_size == 2 or self.full_nvlink: - return inp_size < self.max_size + return inp_size <= self.max_size return False if _is_hip: if self.full_nvlink: - return inp_size < self.max_size + return inp_size <= self.max_size return False return False