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/flake.nix b/flake.nix index 14d5a20e..77942d4a 100644 --- a/flake.nix +++ b/flake.nix @@ -51,8 +51,8 @@ }; nixConfig = { - extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI="; - extra-substituters = "https://exo.cachix.org"; + extra-trusted-public-keys = "exo.cachix.org-1:okq7hl624TBeAR3kV+g39dUFSiaZgLRkLsFBCuJ2NZI= cache.nixos-cuda.org:74DUi4Ye579gUqzH4ziL9IyiJBlDpMRn9MBN8oNan9M="; + extra-substituters = "https://exo.cachix.org https://cache.nixos-cuda.org"; }; outputs = @@ -76,6 +76,8 @@ let # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs) pkgsSwift = import inputs.nixpkgs-swift { inherit system; }; + + pkgsCuda = import ./nix/cuda-pkgs.nix { nixpkgs = inputs.nixpkgs; inherit system; }; in { # Allow unfree for metal-toolchain (needed for Darwin Metal packages) @@ -112,65 +114,137 @@ }; }; - packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin ( - let - uvLock = builtins.fromTOML (builtins.readFile ./uv.lock); - mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package); - uvLockMlxVersion = mlxPackage.version; - in - { - metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { }; - mlx = pkgs.callPackage ./nix/mlx.nix { - inherit (self'.packages) metal-toolchain; - inherit uvLockMlxVersion; - }; - default = self'.packages.exo; - } - ); + packages = lib.optionalAttrs pkgs.stdenv.hostPlatform.isDarwin + ( + let + uvLock = builtins.fromTOML (builtins.readFile ./uv.lock); + mlxPackage = builtins.head (builtins.filter (p: p.name == "mlx" && p.source ? git) uvLock.package); + uvLockMlxVersion = mlxPackage.version; + in + { + metal-toolchain = pkgs.callPackage ./nix/metal-toolchain.nix { }; + mlx = pkgs.callPackage ./nix/mlx.nix { + inherit (self'.packages) metal-toolchain; + inherit uvLockMlxVersion; + }; + default = self'.packages.exo; + } + ) // lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux { + # CUDA-compiled PyTorch and vLLM (built from source by nixpkgs) + torch-cuda = pkgsCuda.python313Packages.torch; + vllm-cuda = pkgsCuda.python313Packages.vllm; - devShells.default = with pkgs; pkgs.mkShell { - inputsFrom = [ self'.checks.cargo-build ]; - - packages = - [ - # FORMATTING - config.treefmt.build.wrapper - - # PYTHON - python313 - uv - ruff - basedpyright - - # RUST - config.rust.toolchain - maturin - - # NIX - nixpkgs-fmt - - # SVELTE - nodejs - - # MISC - just - jq - ] - ++ lib.optionals stdenv.isLinux [ - unixtools.ifconfig - ] - ++ lib.optionals stdenv.isDarwin [ - macmon + # Smoke test script for verifying vLLM + CUDA GPU setup + vllm-check = pkgs.writeShellApplication { + name = "vllm-check"; + runtimeInputs = [ + (pkgsCuda.python313.withPackages (ps: [ ps.torch ps.vllm ])) ]; + # On non-NixOS hosts, NVIDIA driver libraries live in /usr/lib and must be + # LD_PRELOAD'd individually (adding the whole dir causes SIGILL from conflicts). + # These are: CUDA driver, NVML, and the PTX JIT compiler (for flash attention). + # libnvJitLink comes from the nix CUDA toolkit via LD_LIBRARY_PATH. + text = '' + for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do + if [ -e "$dir/libcuda.so.1" ]; then + NVIDIA_LIBS="$dir/libcuda.so.1" + for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do + [ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib" + done + export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}" + break + fi + done + export LD_LIBRARY_PATH="${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + exec python ${inputs.self + /tests/test_vllm_smoke.py} + ''; + }; - OPENSSL_NO_VENDOR = "1"; + # exo with CUDA torch + vLLM — wraps the uv2nix-built package with host driver libs + exo-cuda = pkgs.writeShellApplication { + name = "exo-cuda"; + runtimeInputs = [ self'.packages.exo-cuda-unwrapped ]; + text = '' + for dir in /usr/lib/aarch64-linux-gnu /usr/lib/x86_64-linux-gnu /usr/lib; do + if [ -e "$dir/libcuda.so.1" ]; then + NVIDIA_LIBS="$dir/libcuda.so.1" + for lib in libnvidia-ml.so.1 libnvidia-ptxjitcompiler.so.1; do + [ -e "$dir/$lib" ] && NVIDIA_LIBS="$NVIDIA_LIBS:$dir/$lib" + done + export LD_PRELOAD="$NVIDIA_LIBS''${LD_PRELOAD:+:$LD_PRELOAD}" + break + fi + done + export LD_LIBRARY_PATH="${pkgsCuda.cudaPackages.libnvjitlink}/lib''${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + exec exo-cuda "$@" + ''; + }; + }; - shellHook = '' - export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib" - ${lib.optionalString stdenv.isLinux '' - export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH" - ''} - ''; + # CUDA development shell with torch + vLLM (Linux only) + devShells = lib.optionalAttrs pkgs.stdenv.hostPlatform.isLinux + { + cuda = pkgs.mkShell { + packages = [ + (pkgsCuda.python313.withPackages (ps: [ + ps.torch + ps.vllm + ])) + pkgs.uv + pkgs.just + ]; + + shellHook = '' + echo "CUDA dev shell with torch + vLLM" + python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')" 2>/dev/null || true + ''; + }; + } // { + + default = with pkgs; pkgs.mkShell { + inputsFrom = [ self'.checks.cargo-build ]; + + packages = + [ + # FORMATTING + config.treefmt.build.wrapper + + # PYTHON + python313 + uv + ruff + basedpyright + + # RUST + config.rust.toolchain + maturin + + # NIX + nixpkgs-fmt + + # SVELTE + nodejs + + # MISC + just + jq + ] + ++ lib.optionals stdenv.isLinux [ + unixtools.ifconfig + ] + ++ lib.optionals stdenv.isDarwin [ + macmon + ]; + + OPENSSL_NO_VENDOR = "1"; + + shellHook = '' + export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${python313}/lib" + ${lib.optionalString stdenv.isLinux '' + export LD_LIBRARY_PATH="${openssl.out}/lib:$LD_LIBRARY_PATH" + ''} + ''; + }; }; }; }; diff --git a/nix/cuda-pkgs.nix b/nix/cuda-pkgs.nix new file mode 100644 index 00000000..6904b27f --- /dev/null +++ b/nix/cuda-pkgs.nix @@ -0,0 +1,75 @@ +{ nixpkgs, system }: +let + pkgs = import nixpkgs { inherit system; }; +in +if pkgs.stdenv.hostPlatform.isLinux then + import nixpkgs + { + inherit system; + config = { + allowUnfree = true; + allowBroken = true; + allowUnsupportedSystem = true; + cudaSupport = true; + cudaCapabilities = [ "12.1" ]; + }; + overlays = [ + (final: prev: + let + cudaCompatStub = cfinal: cprev: { + cuda_compat = prev.runCommand "cuda13.0-cuda_compat-stub" { } "mkdir -p $out"; + }; + in + { + cudaPackages = prev.cudaPackages_13.overrideScope cudaCompatStub // { + override = args: + (prev.cudaPackages_13.override args).overrideScope cudaCompatStub; + }; + + pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [ + (_pyFinal: pyPrev: { + cupy = pyPrev.cupy.override { + cudaPackages = final.cudaPackages; + }; + + bitsandbytes = pyPrev.bitsandbytes.overrideAttrs (old: { + preConfigure = (old.preConfigure or "") + '' + export CXXFLAGS="''${CXXFLAGS:-} -I${final.cudaPackages.cuda_crt}/include" + export CUDAFLAGS="''${CUDAFLAGS:-} -I${final.cudaPackages.cuda_crt}/include" + export CMAKE_CUDA_FLAGS="''${CMAKE_CUDA_FLAGS:-} -I${final.cudaPackages.cuda_crt}/include" + ''; + buildInputs = (old.buildInputs or [ ]) ++ [ final.cudaPackages.cuda_crt ]; + }); + + vllm = pyPrev.vllm.overrideAttrs (old: { + buildInputs = (old.buildInputs or [ ]) ++ [ + final.cuda_cccl_with_prefix + final.cudaPackages.cuda_crt + ]; + preConfigure = (old.preConfigure or "") + '' + export CXXFLAGS="''${CXXFLAGS:-} -I${final.cuda_cccl_with_prefix}/include -I${final.cudaPackages.cuda_crt}/include" + export CUDAFLAGS="''${CUDAFLAGS:-} -I${final.cuda_cccl_with_prefix}/include -I${final.cudaPackages.cuda_crt}/include" + ''; + }); + }) + ]; + + magma-cuda-static = prev.magma-cuda-static.overrideAttrs (old: { + postPatch = (old.postPatch or "") + '' + sed -i '/err = cudaGetDeviceProperties( &prop, dev );/a\ int clock_khz = 0; cudaDeviceGetAttribute(\&clock_khz, cudaDevAttrClockRate, dev);' interface_cuda/interface.cpp + sed -i 's/prop\.clockRate/clock_khz/g' interface_cuda/interface.cpp + ''; + }); + + cuda_cccl_with_prefix = prev.runCommand "cuda13.0-cuda_cccl-with-cccl-prefix" { } '' + mkdir -p $out/include + ln -s ${final.cudaPackages.cuda_cccl}/include $out/include/cccl + ''; + + opencv = prev.opencv.override { enableCuda = false; }; + opencv4 = prev.opencv4.override { enableCuda = false; }; + }) + ]; + } +else + null diff --git a/python/parts.nix b/python/parts.nix index d1d1b6d6..41c263b4 100644 --- a/python/parts.nix +++ b/python/parts.nix @@ -3,6 +3,7 @@ perSystem = { config, self', pkgs, lib, system, ... }: let + pkgsCuda = import ../nix/cuda-pkgs.nix { nixpkgs = inputs.nixpkgs; inherit system; }; # Load workspace from uv.lock workspace = inputs.uv2nix.lib.workspace.loadWorkspace { workspaceRoot = inputs.self; @@ -99,16 +100,18 @@ } ); + baseOverlays = [ + inputs.pyproject-build-systems.overlays.default + overlay + exoOverlay + buildSystemsOverlay + linuxOverlay + ]; + pythonSet = (pkgs.callPackage inputs.pyproject-nix.build.packages { inherit python; }).overrideScope ( - lib.composeManyExtensions [ - inputs.pyproject-build-systems.overlays.default - overlay - exoOverlay - buildSystemsOverlay - linuxOverlay - ] + lib.composeManyExtensions baseOverlays ); # mlx-cpu and mlx-cuda-13 both ship mlx/ site-packages files; keep first. # mlx-cpu/mlx-cuda-13 and nvidia-cudnn-cu12/cu13 ship overlapping files. @@ -172,6 +175,35 @@ --set EXO_RESOURCES_DIR ${inputs.self + /resources} \ ${lib.optionalString pkgs.stdenv.hostPlatform.isDarwin "--prefix PATH : ${pkgs.macmon}/bin"} ''; + + vllmEnv = pkgsCuda.python313.withPackages (ps: [ ps.vllm ]); + + vllmSite = pkgs.runCommand "vllm-site-filtered" { } '' + mkdir -p $out + for pkg in ${vllmEnv}/${python.sitePackages}/*; do + name=$(basename "$pkg") + case "$name" in + anyio*|pydantic*) ;; + *) ln -s "$pkg" "$out/$name" ;; + esac + done + ''; + + exoCudaVenv = (pythonSet.mkVirtualEnv "exo-cuda-env" exoDeps).overrideAttrs { + venvIgnoreCollisions = venvCollisionPaths; + }; + + exoCudaPackage = pkgs.runCommand "exo-cuda" + { + nativeBuildInputs = [ pkgs.makeWrapper ]; + } + '' + mkdir -p $out/bin + makeWrapper ${exoCudaVenv}/bin/exo $out/bin/exo-cuda \ + --set EXO_DASHBOARD_DIR ${self'.packages.dashboard} \ + --set EXO_RESOURCES_DIR ${inputs.self + /resources} \ + --prefix PYTHONPATH : "${vllmSite}" + ''; in { # Python package only available on macOS (requires MLX/Metal) @@ -180,7 +212,9 @@ exo = exoPackage; # Test environment for running pytest outside of Nix sandbox (needs GPU access) exo-test-env = testVenv; - } // { + } // lib.optionalAttrs (pkgsCuda != null) { + exo-cuda-unwrapped = exoCudaPackage; + } // { exo-bench = mkBenchScript "exo-bench" (inputs.self + /bench/exo_bench.py); 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); diff --git a/src/exo/main.py b/src/exo/main.py index 2ecb62c2..cbeb09c1 100644 --- a/src/exo/main.py +++ b/src/exo/main.py @@ -255,8 +255,83 @@ class Node: self.api.unpause(result.won_clock) +def _run_vllm_server(args: "Args") -> None: + from typing import TYPE_CHECKING + + if not TYPE_CHECKING: + from vllm.entrypoints.openai.api_server import ( + run_server, + ) + from vllm.entrypoints.openai.cli_args import ( + make_arg_parser, + ) + from vllm.utils.argparse_utils import ( + FlexibleArgumentParser, + ) + + vllm_argv = [ + "--model", + args.model or "Qwen/Qwen2.5-0.5B-Instruct", + "--host", + "0.0.0.0", + "--port", + str(args.api_port), + ] + if args.max_model_len is not None: + vllm_argv += ["--max-model-len", str(args.max_model_len)] + + parser = make_arg_parser(FlexibleArgumentParser()) + vllm_args = parser.parse_args(vllm_argv) + logger.info( + f"vLLM detected — starting OpenAI-compatible server on port {args.api_port}" + ) + anyio.run(run_server, vllm_args) + + +def _check_vllm(args: "Args") -> None: + import sys + + logger.info("Checking PyTorch + CUDA...") + try: + import torch + except ImportError: + logger.critical("FAIL: PyTorch not installed") + sys.exit(1) + + if not torch.cuda.is_available(): + logger.critical("FAIL: CUDA not available") + sys.exit(1) + + gpu_name: str = torch.cuda.get_device_name(0) + compute_cap: tuple[int, int] = torch.cuda.get_device_capability(0) + cuda_version = torch.version.cuda or "unknown" + logger.info(f" PyTorch {torch.__version__}") + logger.info(f" CUDA {cuda_version}") + logger.info(f" GPU: {gpu_name} (compute {compute_cap[0]}.{compute_cap[1]})") + + logger.info("Checking vLLM...") + try: + import vllm # type: ignore + except ImportError: + logger.critical("FAIL: vLLM not installed") + sys.exit(1) + + logger.warning(f" vLLM {vllm.__version__}") # type: ignore + logger.warning("All checks passed.") + + try: + _run_vllm_server(args) + return + except ImportError: + pass + + def main(): args = Args.parse() + + if args.check_vllm: + _check_vllm(args) + soft, hard = resource.getrlimit(resource.RLIMIT_NOFILE) target = min(max(soft, 65535), hard) resource.setrlimit(resource.RLIMIT_NOFILE, (target, hard)) @@ -282,6 +357,15 @@ def main(): os.environ["EXO_FAST_SYNCH"] = "off" logger.info("FAST_SYNCH forced OFF") + try: + import vllm # pyright: ignore[reportMissingImports, reportUnusedImport] # noqa: F401 + except ImportError: + vllm = None # pyright: ignore[reportAssignmentType] + + if vllm is not None: + _run_vllm_server(args) + return + node = anyio.run(Node.create, args) try: anyio.run(node.run) @@ -306,6 +390,9 @@ class Args(CamelCaseModel): offline: bool = os.getenv("EXO_OFFLINE", "false").lower() == "true" no_batch: bool = False fast_synch: bool | None = None # None = auto, True = force on, False = force off + check_vllm: bool = False + model: str | None = None # vLLM model name/path + max_model_len: int | None = None # vLLM max sequence length @classmethod def parse(cls) -> Self: @@ -377,6 +464,25 @@ class Args(CamelCaseModel): dest="fast_synch", help="Force MLX FAST_SYNCH off", ) + parser.add_argument( + "--check-vllm", + action="store_true", + dest="check_vllm", + help="Check vLLM + CUDA GPU setup and exit", + ) + parser.add_argument( + "--model", + type=str, + default=None, + help="Model name or path (used in vLLM mode)", + ) + parser.add_argument( + "--max-model-len", + type=int, + default=None, + dest="max_model_len", + help="Maximum model context length (used in vLLM mode)", + ) args = parser.parse_args() return cls(**vars(args)) # pyright: ignore[reportAny] - We are intentionally validating here, we can't do it statically diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 8b22e6b3..859cb37c 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -163,6 +163,94 @@ def initialize_mlx( return mlx_distributed_init(bound_instance) +# Quant methods that mlx_lm can handle natively (no conversion needed) +_MLX_NATIVE_QUANT_METHODS = frozenset( + {"awq", "gptq", "bitnet", "mxfp4", "compressed-tensors"} +) + + +def _needs_mlx_conversion(model_path: Path) -> bool: + """Check if a model uses a quantization format that mlx_lm cannot load natively. + + Returns True for formats like fp8 that need dequantization via mlx_lm.convert. + Returns False for native MLX models, AWQ/GPTQ (without g_idx), etc. + """ + config_file = model_path / "config.json" + if not config_file.exists(): + return False + try: + with open(config_file) as f: + config = json.load(f) # pyright: ignore[reportAny] + except (json.JSONDecodeError, OSError): + return False + + quant_config: dict[str, object] | None = config.get("quantization_config") # pyright: ignore[reportAny] + if not quant_config: + text_config: dict[str, object] = config.get("text_config", {}) # pyright: ignore[reportAny] + quant_config = text_config.get("quantization_config") # pyright: ignore[reportAssignmentType] + if not quant_config: + return False + + quant_method = str(quant_config.get("quant_method", "")) + + # GPTQ with g_idx is explicitly unsupported by mlx_lm + if quant_method == "gptq" and quant_config.get("desc_act", False): + return True + + # Check for any weight files containing g_idx (GPTQ models that mlx_lm will reject) + if quant_method == "gptq": + try: + from safetensors import safe_open + + for st_file in model_path.glob("*.safetensors"): + with safe_open(str(st_file), framework="numpy") as st: # pyright: ignore[reportUnknownVariableType] + keys: list[str] = st.keys() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] + if any("g_idx" in str(k) for k in keys): # pyright: ignore[reportUnknownArgumentType, reportUnknownVariableType] + return True + break # Only need to check one file for key names + except Exception: + pass + + # FP8 and other non-native methods need conversion + return bool(quant_method and quant_method not in _MLX_NATIVE_QUANT_METHODS) + + +def _convert_to_mlx(model_path: Path) -> Path: + """Convert a non-MLX model to MLX format using mlx_lm.convert --dequantize. + + The converted model is saved alongside the original with a '-mlx' suffix. + Returns the path to the converted model directory. + """ + converted_path = model_path.parent / (model_path.name + "-mlx") + if ( + converted_path.exists() + and (converted_path / "config.json").exists() + and list(converted_path.glob("*.safetensors")) + ): + logger.info(f"Using previously converted MLX model at {converted_path}") + return converted_path + + logger.info(f"Converting model at {model_path} to MLX format (dequantizing)...") + from mlx_lm.convert import convert # pyright: ignore[reportUnknownVariableType] + + convert( + hf_path=str(model_path), + mlx_path=str(converted_path), + dequantize=True, + ) + logger.info(f"Conversion complete: {converted_path}") + return converted_path + + +def _maybe_convert_model(model_path: Path) -> Path: + """If the model needs conversion to MLX format, convert it and return the new path. + Otherwise return the original path unchanged. + """ + if _needs_mlx_conversion(model_path): + return _convert_to_mlx(model_path) + return model_path + + def load_mlx_items( bound_instance: BoundInstance, group: Group | None, @@ -171,7 +259,9 @@ def load_mlx_items( ) -> tuple[Model, TokenizerWrapper]: if group is None: logger.info(f"Single device used for {bound_instance.instance}") - model_path = build_model_path(bound_instance.bound_shard.model_card.model_id) + model_path = _maybe_convert_model( + build_model_path(bound_instance.bound_shard.model_card.model_id) + ) start_time = time.perf_counter() model, _ = load_model(model_path, lazy=True, strict=False) # Eval layers one by one for progress reporting @@ -219,7 +309,9 @@ def shard_and_load( on_timeout: TimeoutCallback | None, on_layer_loaded: LayerLoadedCallback | None, ) -> tuple[nn.Module, TokenizerWrapper]: - model_path = build_model_path(shard_metadata.model_card.model_id) + model_path = _maybe_convert_model( + build_model_path(shard_metadata.model_card.model_id) + ) model, _ = load_model(model_path, lazy=True, strict=False) logger.debug(model)