From 458278a5a7e02cb0644f427ddab478fa92dd5e24 Mon Sep 17 00:00:00 2001 From: Evan Date: Mon, 16 Mar 2026 11:09:05 +0000 Subject: [PATCH] add leo's bandwidth test --- bench/test_mlx_bandwidth.py | 377 ++++++++++++++++++++++++++++++++++++ python/parts.nix | 1 + 2 files changed, 378 insertions(+) create mode 100644 bench/test_mlx_bandwidth.py diff --git a/bench/test_mlx_bandwidth.py b/bench/test_mlx_bandwidth.py new file mode 100644 index 00000000..6b165b8b --- /dev/null +++ b/bench/test_mlx_bandwidth.py @@ -0,0 +1,377 @@ +# type: ignore +import argparse +import json +import os +import statistics +import sys +import tempfile +import time + +import mlx.core as mx + +DTYPE_MAP = { + "float32": (mx.float32, 4), + "float16": (mx.float16, 2), + "bfloat16": (mx.bfloat16, 2), +} + +SIZES = [ + 1 * 1024, + 4 * 1024, + 16 * 1024, + 64 * 1024, + 256 * 1024, + 1 * 1024 * 1024, + 4 * 1024 * 1024, + 16 * 1024 * 1024, + 64 * 1024 * 1024, + 256 * 1024 * 1024, + 1 * 1024 * 1024 * 1024, + 2 * 1024 * 1024 * 1024, + 4 * 1024 * 1024 * 1024, + 8 * 1024 * 1024 * 1024, +] + + +def format_bytes(n: int) -> str: + if n >= 1024 * 1024 * 1024: + return f"{n / (1024 * 1024 * 1024):.0f} GB" + if n >= 1024 * 1024: + return f"{n / (1024 * 1024):.0f} MB" + if n >= 1024: + return f"{n / 1024:.0f} KB" + return f"{n} B" + + +def format_time(seconds: float) -> str: + if seconds >= 1.0: + return f"{seconds:.3f} s" + if seconds >= 0.001: + return f"{seconds * 1000:.2f} ms" + return f"{seconds * 1_000_000:.1f} us" + + +def format_bandwidth(bytes_per_sec: float) -> str: + if bytes_per_sec >= 1024 * 1024 * 1024: + return f"{bytes_per_sec / (1024 * 1024 * 1024):.2f} GB/s" + if bytes_per_sec >= 1024 * 1024: + return f"{bytes_per_sec / (1024 * 1024):.1f} MB/s" + return f"{bytes_per_sec / 1024:.1f} KB/s" + + +def barrier(group: mx.distributed.Group) -> None: + mx.eval(mx.distributed.all_sum(mx.array(1.0), group=group)) + + +def init_ring( + rank: int, self_ip: str, peer_ip: str, port: int, tmpdir: str +) -> mx.distributed.Group: + if rank == 0: + hosts = [f"{self_ip}:{port}", f"{peer_ip}:{port}"] + else: + hosts = [f"{peer_ip}:{port}", f"{self_ip}:{port}"] + + hostfile = os.path.join(tmpdir, "hosts.json") + with open(hostfile, "w") as f: + json.dump(hosts, f) + + for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"): + os.environ.pop(var, None) + + os.environ["MLX_HOSTFILE"] = hostfile + os.environ["MLX_RANK"] = str(rank) + return mx.distributed.init(backend="ring", strict=True) + + +def init_jaccl( + rank: int, interface: str, coordinator: str, port: int, tmpdir: str +) -> mx.distributed.Group: + devices = [[None, interface], [interface, None]] + devfile = os.path.join(tmpdir, "devices.json") + with open(devfile, "w") as f: + json.dump(devices, f) + + for var in ("MLX_HOSTFILE", "MLX_RANK", "MLX_IBV_DEVICES", "MLX_JACCL_COORDINATOR"): + os.environ.pop(var, None) + + os.environ["MLX_IBV_DEVICES"] = devfile + os.environ["MLX_RANK"] = str(rank) + if rank == 0: + os.environ["MLX_JACCL_COORDINATOR"] = f"0.0.0.0:{port}" + else: + os.environ["MLX_JACCL_COORDINATOR"] = coordinator + + return mx.distributed.init(backend="jaccl", strict=True) + + +def bench_unidirectional( + group: mx.distributed.Group, + rank: int, + size_bytes: int, + dtype: mx.Dtype, + element_size: int, + warmup: int, + iterations: int, +) -> list[float]: + n_elements = size_bytes // element_size + tensor = mx.random.normal(shape=(n_elements,)).astype(dtype) + mx.eval(tensor) + + for _ in range(warmup): + if rank == 0: + sent = mx.distributed.send(tensor, dst=1, group=group) + mx.eval(sent) + else: + received = mx.distributed.recv_like(tensor, src=0, group=group) + mx.eval(received) + barrier(group) + + times: list[float] = [] + for _ in range(iterations): + barrier(group) + t0 = time.perf_counter() + if rank == 0: + sent = mx.distributed.send(tensor, dst=1, group=group) + mx.eval(sent) + else: + received = mx.distributed.recv_like(tensor, src=0, group=group) + mx.eval(received) + barrier(group) + t1 = time.perf_counter() + times.append(t1 - t0) + + return times + + +def bench_rtt( + group: mx.distributed.Group, + rank: int, + size_bytes: int, + dtype: mx.Dtype, + element_size: int, + warmup: int, + iterations: int, +) -> list[float]: + n_elements = size_bytes // element_size + tensor = mx.random.normal(shape=(n_elements,)).astype(dtype) + mx.eval(tensor) + + for _ in range(warmup): + if rank == 0: + sent = mx.distributed.send(tensor, dst=1, group=group) + mx.eval(sent) + received = mx.distributed.recv_like(tensor, src=1, group=group) + mx.eval(received) + else: + received = mx.distributed.recv_like(tensor, src=0, group=group) + mx.eval(received) + sent = mx.distributed.send(received, dst=0, group=group) + mx.eval(sent) + barrier(group) + + times: list[float] = [] + for _ in range(iterations): + barrier(group) + t0 = time.perf_counter() + if rank == 0: + sent = mx.distributed.send(tensor, dst=1, group=group) + mx.eval(sent) + received = mx.distributed.recv_like(tensor, src=1, group=group) + mx.eval(received) + else: + received = mx.distributed.recv_like(tensor, src=0, group=group) + mx.eval(received) + sent = mx.distributed.send(received, dst=0, group=group) + mx.eval(sent) + barrier(group) + t1 = time.perf_counter() + times.append(t1 - t0) + + return times + + +def bench_all_gather( + group: mx.distributed.Group, + rank: int, + size_bytes: int, + dtype: mx.Dtype, + element_size: int, + warmup: int, + iterations: int, +) -> list[float]: + n_elements = (size_bytes // 2) // element_size + tensor = mx.random.normal(shape=(n_elements,)).astype(dtype) + mx.eval(tensor) + + for _ in range(warmup): + gathered = mx.distributed.all_gather(tensor, group=group) + mx.eval(gathered) + barrier(group) + + times: list[float] = [] + for _ in range(iterations): + barrier(group) + t0 = time.perf_counter() + gathered = mx.distributed.all_gather(tensor, group=group) + mx.eval(gathered) + t1 = time.perf_counter() + times.append(t1 - t0) + + return times + + +def print_table(title: str, rows: list[dict[str, str]]) -> None: + print(f"\n=== {title} ===") + headers = ["Size", "Median", "Min", "Max", "Bandwidth"] + widths = [ + max(len(h), max((len(r[h]) for r in rows), default=0)) + 2 for h in headers + ] + header_line = "".join(h.ljust(w) for h, w in zip(headers, widths, strict=True)) + print(header_line) + print("-" * len(header_line)) + for row in rows: + print("".join(row[h].ljust(w) for h, w in zip(headers, widths, strict=True))) + + +def run_bench( + name: str, + bench_fn, + group: mx.distributed.Group, + rank: int, + dtype: mx.Dtype, + element_size: int, + warmup: int, + iterations: int, + bw_multiplier: int = 1, +) -> None: + rows: list[dict[str, str]] = [] + for size in SIZES: + if rank == 0: + print(f" {name}: {format_bytes(size)}...", end="", flush=True) + times = bench_fn(group, rank, size, dtype, element_size, warmup, iterations) + if rank == 0: + med = statistics.median(times) + mn = min(times) + mx_ = max(times) + bw = (size * bw_multiplier) / med + rows.append( + { + "Size": format_bytes(size), + "Median": format_time(med), + "Min": format_time(mn), + "Max": format_time(mx_), + "Bandwidth": format_bandwidth(bw), + } + ) + print(f" {format_bandwidth(bw)}") + if rank == 0: + print_table(name, rows) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="MLX Distributed Communication Benchmark" + ) + subparsers = parser.add_subparsers(dest="backend", required=True) + + ring_parser = subparsers.add_parser("ring") + ring_parser.add_argument("--rank", type=int, required=True, choices=[0, 1]) + ring_parser.add_argument("--self-ip", required=True) + ring_parser.add_argument("--peer-ip", required=True) + ring_parser.add_argument("--port", type=int, default=5555) + + jaccl_parser = subparsers.add_parser("jaccl") + jaccl_parser.add_argument("--rank", type=int, required=True, choices=[0, 1]) + jaccl_parser.add_argument("--interface", required=True) + jaccl_parser.add_argument( + "--coordinator", + type=str, + default=None, + help="IP:PORT of rank 0 (required for rank 1)", + ) + jaccl_parser.add_argument( + "--port", type=int, default=9999, help="Coordinator port (rank 0 only)" + ) + + for p in [ring_parser, jaccl_parser]: + p.add_argument("--warmup", type=int, default=3) + p.add_argument("--iterations", type=int, default=10) + p.add_argument("--dtype", choices=list(DTYPE_MAP.keys()), default="float32") + + args = parser.parse_args() + + if args.backend == "jaccl" and args.rank == 1 and args.coordinator is None: + jaccl_parser.error("--coordinator is required for rank 1") + + return args + + +def main() -> int: + args = parse_args() + dtype, element_size = DTYPE_MAP[args.dtype] + + with tempfile.TemporaryDirectory() as tmpdir: + if args.backend == "ring": + print(f"Initializing ring backend (rank {args.rank})...") + group = init_ring(args.rank, args.self_ip, args.peer_ip, args.port, tmpdir) + else: + print(f"Initializing jaccl backend (rank {args.rank})...") + group = init_jaccl( + args.rank, args.interface, args.coordinator or "", args.port, tmpdir + ) + + print(f"Rank {group.rank()} of {group.size()} initialized") + barrier(group) + + if args.rank == 0: + print("\nMLX Distributed Communication Benchmark") + print( + f"Backend: {args.backend} | Dtype: {args.dtype} | Warmup: {args.warmup} | Iterations: {args.iterations}" + ) + + run_bench( + "Unidirectional (rank 0 -> rank 1)", + bench_unidirectional, + group, + args.rank, + dtype, + element_size, + args.warmup, + args.iterations, + ) + run_bench( + "Round-Trip (ping-pong)", + bench_rtt, + group, + args.rank, + dtype, + element_size, + args.warmup, + args.iterations, + bw_multiplier=2, + ) + run_bench( + "All-Gather", + bench_all_gather, + group, + args.rank, + dtype, + element_size, + args.warmup, + args.iterations, + ) + + if args.rank == 0: + print("\nDone.") + else: + print("Rank 1 complete.") + + return 0 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except KeyboardInterrupt: + print("\nInterrupted.") + sys.exit(1) diff --git a/python/parts.nix b/python/parts.nix index d1d1b6d6..ce3fe4b2 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -185,6 +185,7 @@ exo-eval = mkBenchScript "exo-eval" (inputs.self + /bench/exo_eval.py); exo-eval-tool-calls = mkBenchScript "exo-eval-tool-calls" (inputs.self + /bench/eval_tool_calls.py); exo-get-all-models-on-cluster = mkSimplePythonScript "exo-get-all-models-on-cluster" (inputs.self + /tests/get_all_models_on_cluster.py); + exo-mlx-bandwidth-test = mkPythonScript "exo-mlx-bandwidth-test" (inputs.self + /bench/test_mlx_bandwidth.py); }; checks = {