From bdb43e1dbb331a3fdf83ff65c0ea2765dd54065a Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Tue, 13 Jan 2026 19:17:14 +0000 Subject: [PATCH 1/4] nix: drop noisy echos from devshell Drop all the printing when entering a devshell. It's annoying, and not a super accurate description of how to develop exo anyway. --- flake.nix | 3 --- 1 file changed, 3 deletions(-) diff --git a/flake.nix b/flake.nix index 5c7a56cf..15df8eef 100644 --- a/flake.nix +++ b/flake.nix @@ -131,9 +131,6 @@ export PKG_CONFIG_PATH="${pkgs.openssl.dev}/lib/pkgconfig:$PKG_CONFIG_PATH" export LD_LIBRARY_PATH="${pkgs.openssl.out}/lib:$LD_LIBRARY_PATH" ''} - echo - echo "🍎🍎 Run 'just ' to get started" - just --list ''; }; }; From e6434ec446a72af870f074d5934ceadd273a5c1b Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Mon, 12 Jan 2026 17:27:47 +0000 Subject: [PATCH 2/4] nix: add Rust builds with crane and fenix The Rust workspace lacked Nix build support, making it difficult to build packages reproducibly or run checks in CI. Added a flake-parts module at rust/parts.nix that uses crane for Rust builds and fenix for the nightly toolchain. The source filter isolates rust/ and root Cargo files to prevent Python/docs changes from triggering Rust rebuilds. Exports packages (system_custodian, exo_pyo3_bindings wheel, exo-rust-workspace) and checks (cargo-nextest, cargo-doc) for all three target platforms. The devShell now uses inputsFrom to inherit build dependencies from the workspace package, removing the need for manual pkg-config/openssl setup. Test plan: - Ran `nix flake check` successfully - Built `nix build ".#checks.x86_64-linux.cargo-nextest"` and tests pass - Built `nix build ".#exo_pyo3_bindings"` and wheel is produced --- flake.lock | 16 +++++ flake.nix | 44 ++++++-------- rust/parts.nix | 159 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 193 insertions(+), 26 deletions(-) create mode 100644 rust/parts.nix diff --git a/flake.lock b/flake.lock index d45f6f50..41de94a6 100644 --- a/flake.lock +++ b/flake.lock @@ -1,5 +1,20 @@ { "nodes": { + "crane": { + "locked": { + "lastModified": 1767744144, + "narHash": "sha256-9/9ntI0D+HbN4G0TrK3KmHbTvwgswz7p8IEJsWyef8Q=", + "owner": "ipetkov", + "repo": "crane", + "rev": "2fb033290bf6b23f226d4c8b32f7f7a16b043d7e", + "type": "github" + }, + "original": { + "owner": "ipetkov", + "repo": "crane", + "type": "github" + } + }, "fenix": { "inputs": { "nixpkgs": [ @@ -75,6 +90,7 @@ }, "root": { "inputs": { + "crane": "crane", "fenix": "fenix", "flake-parts": "flake-parts", "nixpkgs": "nixpkgs", diff --git a/flake.nix b/flake.nix index 15df8eef..a9ae9ca7 100644 --- a/flake.nix +++ b/flake.nix @@ -9,6 +9,8 @@ inputs.nixpkgs-lib.follows = "nixpkgs"; }; + crane.url = "github:ipetkov/crane"; + fenix = { url = "github:nix-community/fenix"; inputs.nixpkgs.follows = "nixpkgs"; @@ -39,10 +41,11 @@ imports = [ inputs.treefmt-nix.flakeModule + ./rust/parts.nix ]; perSystem = - { config, inputs', pkgs, lib, system, ... }: + { config, self', inputs', pkgs, lib, system, ... }: let fenixToolchain = inputs'.fenix.packages.complete; # Use pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs) @@ -59,7 +62,7 @@ }; rustfmt = { enable = true; - package = fenixToolchain.rustfmt; + package = config.rust.toolchain; }; prettier = { enable = true; @@ -79,6 +82,8 @@ ''; devShells.default = with pkgs; pkgs.mkShell { + inputsFrom = [ self'.checks.cargo-build ]; + packages = [ # FORMATTING @@ -91,14 +96,8 @@ basedpyright # RUST - (fenixToolchain.withComponents [ - "cargo" - "rustc" - "clippy" - "rustfmt" - "rust-src" - ]) - rustup # Just here to make RustRover happy + config.rust.toolchain + maturin # NIX nixpkgs-fmt @@ -110,26 +109,19 @@ just jq ] - ++ (pkgs.lib.optionals pkgs.stdenv.isLinux [ - # IFCONFIG + ++ lib.optionals stdenv.isLinux [ unixtools.ifconfig - - # Build dependencies for Linux - pkg-config - openssl - ]) - ++ (pkgs.lib.optionals pkgs.stdenv.isDarwin [ - # MACMON + ] + ++ lib.optionals stdenv.isDarwin [ macmon - ]); + ]; + + OPENSSL_NO_VENDOR = "1"; shellHook = '' - # PYTHON - export LD_LIBRARY_PATH="$LD_LIBRARY_PATH:${pkgs.python313}/lib" - ${lib.optionalString pkgs.stdenv.isLinux '' - # Build environment for Linux - export PKG_CONFIG_PATH="${pkgs.openssl.dev}/lib/pkgconfig:$PKG_CONFIG_PATH" - export LD_LIBRARY_PATH="${pkgs.openssl.out}/lib:$LD_LIBRARY_PATH" + 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/rust/parts.nix b/rust/parts.nix new file mode 100644 index 00000000..e00b27c6 --- /dev/null +++ b/rust/parts.nix @@ -0,0 +1,159 @@ +{ inputs, ... }: +{ + perSystem = + { config, self', inputs', pkgs, lib, ... }: + let + # Fenix nightly toolchain with all components + fenixPkgs = inputs'.fenix.packages; + rustToolchain = fenixPkgs.complete.withComponents [ + "cargo" + "rustc" + "clippy" + "rustfmt" + "rust-src" + "rust-analyzer" + ]; + + # Crane with fenix toolchain + craneLib = (inputs.crane.mkLib pkgs).overrideToolchain rustToolchain; + + # Source filtering - only include rust/ directory and root Cargo files + # This ensures changes to Python/docs/etc don't trigger Rust rebuilds + src = lib.cleanSourceWith { + src = inputs.self; + filter = + path: type: + let + baseName = builtins.baseNameOf path; + parentDir = builtins.dirOf path; + inRustDir = + (lib.hasInfix "/rust/" path) + || (lib.hasSuffix "/rust" parentDir) + || (baseName == "rust" && type == "directory"); + isRootCargoFile = + (baseName == "Cargo.toml" || baseName == "Cargo.lock") + && (builtins.dirOf path == toString inputs.self); + in + isRootCargoFile + || (inRustDir && (craneLib.filterCargoSources path type || lib.hasSuffix ".toml" path || lib.hasSuffix ".md" path)); + }; + + # Common arguments for all Rust builds + commonArgs = { + inherit src; + pname = "exo-rust"; + version = "0.0.1"; + strictDeps = true; + + nativeBuildInputs = [ + pkgs.pkg-config + pkgs.python313 # Required for pyo3-build-config + ]; + + buildInputs = [ + pkgs.openssl + pkgs.python313 # Required for pyo3 tests + ]; + + OPENSSL_NO_VENDOR = "1"; + + # Required for pyo3 tests to find libpython + LD_LIBRARY_PATH = lib.makeLibraryPath [ pkgs.python313 ]; + }; + + # Build dependencies once for caching + cargoArtifacts = craneLib.buildDepsOnly ( + commonArgs + // { + cargoExtraArgs = "--workspace"; + } + ); + in + { + # Export toolchain for use in treefmt and devShell + options.rust = { + toolchain = lib.mkOption { + type = lib.types.package; + default = rustToolchain; + description = "The Rust toolchain to use"; + }; + }; + + config = { + packages = { + # The system_custodian binary + system_custodian = craneLib.buildPackage ( + commonArgs + // { + inherit cargoArtifacts; + cargoExtraArgs = "-p system_custodian"; + + meta = { + description = "System custodian daemon for exo"; + mainProgram = "system_custodian"; + }; + } + ); + + # Python bindings wheel via maturin + exo_pyo3_bindings = craneLib.buildPackage ( + commonArgs + // { + inherit cargoArtifacts; + pname = "exo_pyo3_bindings"; + + nativeBuildInputs = commonArgs.nativeBuildInputs ++ [ + pkgs.maturin + ]; + + buildPhaseCargoCommand = '' + maturin build \ + --release \ + --manylinux off \ + --manifest-path rust/exo_pyo3_bindings/Cargo.toml \ + --features "pyo3/extension-module,pyo3/experimental-async" \ + --interpreter ${pkgs.python313}/bin/python \ + --out dist + ''; + + # Don't use crane's default install behavior + doNotPostBuildInstallCargoBinaries = true; + + installPhaseCommand = '' + mkdir -p $out + cp dist/*.whl $out/ + ''; + } + ); + }; + + checks = { + # Full workspace build (all crates) + cargo-build = craneLib.buildPackage ( + commonArgs + // { + inherit cargoArtifacts; + cargoExtraArgs = "--workspace"; + } + ); + # Run tests with nextest + cargo-nextest = craneLib.cargoNextest ( + commonArgs + // { + inherit cargoArtifacts; + cargoExtraArgs = "--workspace"; + } + ); + + # Build documentation + cargo-doc = craneLib.cargoDoc ( + commonArgs + // { + inherit cargoArtifacts; + cargoExtraArgs = "--workspace"; + } + ); + }; + }; + }; +} From 3671528fa46a0633fc7853a89e68e863273d0c3e Mon Sep 17 00:00:00 2001 From: Jake Hillion Date: Tue, 13 Jan 2026 22:19:45 +0000 Subject: [PATCH 3/4] nix: add dashboard build with dream2nix Continue working towards a fully Nix based build by building the dashboard with Nix. Continuing the theme of using the existing lock files, use dream2nix to parse the lock file and build the tree of dependency derivations. dream2nix doesn't like the bundleDependencies, so we apply a small patch to the lock file that drops all dependencies that are bundled. This should ideally be contributed upstream but that can be done later. Use this new dashboard build in the build-app CI workflow, meaning future macOS apps will include this reproducible dashboard. Test plan: - Built a DMG, shipped to a cluster, loaded in a browser with no cache and the dashboard looks good. - Directory layout is as expected: ``` $ nix build .#dashboard $ find result/ ... result/_app/immutable/entry result/_app/immutable/entry/app.CTPAnMjf.js result/_app/immutable/entry/start.fUSEa-2O.js result/_app/immutable/nodes result/_app/immutable/nodes/3.DqQr1Obm.js result/_app/immutable/nodes/0.DgEY44RO.js result/_app/immutable/nodes/2.BjZg_lJh.js result/_app/immutable/nodes/1.D6vGUYYT.js result/_app/env.js result/_app/version.json result/exo-logo.png result/favicon.ico result/index.html ``` --- .github/workflows/build-app.yml | 17 +++++- dashboard/dashboard.nix | 60 ++++++++++++++++++ dashboard/parts.nix | 44 +++++++++++++ flake.lock | 105 ++++++++++++++++++++++++++++++++ flake.nix | 6 ++ 5 files changed, 229 insertions(+), 3 deletions(-) create mode 100644 dashboard/dashboard.nix create mode 100644 dashboard/parts.nix diff --git a/.github/workflows/build-app.yml b/.github/workflows/build-app.yml index 44e9c77e..542dc800 100644 --- a/.github/workflows/build-app.yml +++ b/.github/workflows/build-app.yml @@ -113,11 +113,22 @@ jobs: uv python install uv sync --locked + - name: Install Nix + uses: cachix/install-nix-action@v31 + with: + nix_path: nixpkgs=channel:nixos-unstable + + - name: Configure Cachix + uses: cachix/cachix-action@v14 + with: + name: exo + authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}" + - name: Build dashboard run: | - cd dashboard - npm ci - npm run build + DASHBOARD_OUT=$(nix build .#dashboard --print-build-logs --no-link --print-out-paths) + mkdir -p dashboard/build + cp -r "$DASHBOARD_OUT"/* dashboard/build/ - name: Install Sparkle CLI run: | diff --git a/dashboard/dashboard.nix b/dashboard/dashboard.nix new file mode 100644 index 00000000..6ecd1490 --- /dev/null +++ b/dashboard/dashboard.nix @@ -0,0 +1,60 @@ +{ lib +, config +, dream2nix +, ... +}: +let + # Read and parse the lock file + rawLockFile = builtins.fromJSON (builtins.readFile "${config.deps.dashboardSrc}/package-lock.json"); + + # For packages with bundleDependencies, filter out deps that are bundled + # (bundled deps are inside the tarball, not separate lockfile entries) + fixedPackages = lib.mapAttrs + (path: entry: + if entry ? bundleDependencies && entry.bundleDependencies != [ ] + then entry // { + dependencies = lib.filterAttrs + (name: _: !(lib.elem name entry.bundleDependencies)) + (entry.dependencies or { }); + } + else entry + ) + (rawLockFile.packages or { }); + + fixedLockFile = rawLockFile // { packages = fixedPackages; }; +in +{ + imports = [ + dream2nix.modules.dream2nix.nodejs-package-lock-v3 + dream2nix.modules.dream2nix.nodejs-granular-v3 + ]; + + name = "exo-dashboard"; + version = "1.0.0"; + + mkDerivation = { + src = config.deps.dashboardSrc; + + buildPhase = '' + runHook preBuild + npm run build + runHook postBuild + ''; + + installPhase = '' + runHook preInstall + cp -r build $out/build + runHook postInstall + ''; + }; + + deps = { nixpkgs, ... }: { + inherit (nixpkgs) stdenv; + dashboardSrc = null; # Injected by parts.nix + }; + + nodejs-package-lock-v3 = { + # Don't use packageLockFile - provide the fixed lock content directly + packageLock = fixedLockFile; + }; +} diff --git a/dashboard/parts.nix b/dashboard/parts.nix new file mode 100644 index 00000000..80fbb156 --- /dev/null +++ b/dashboard/parts.nix @@ -0,0 +1,44 @@ +{ inputs, ... }: +{ + perSystem = + { pkgs, lib, ... }: + let + # Filter source to only include dashboard directory + src = lib.cleanSourceWith { + src = inputs.self; + filter = + path: type: + let + baseName = builtins.baseNameOf path; + inDashboardDir = + (lib.hasInfix "/dashboard/" path) + || (lib.hasSuffix "/dashboard" (builtins.dirOf path)) + || (baseName == "dashboard" && type == "directory"); + in + inDashboardDir; + }; + + # Build the dashboard with dream2nix (includes node_modules in output) + dashboardFull = inputs.dream2nix.lib.evalModules { + packageSets.nixpkgs = pkgs; + modules = [ + ./dashboard.nix + { + paths.projectRoot = inputs.self; + paths.projectRootFile = "flake.nix"; + paths.package = inputs.self + "/dashboard"; + } + # Inject the filtered source + { + deps.dashboardSrc = lib.mkForce "${src}/dashboard"; + } + ]; + }; + in + { + # Extract just the static site from the full build + packages.dashboard = pkgs.runCommand "exo-dashboard" { } '' + cp -r ${dashboardFull}/build $out + ''; + }; +} diff --git a/flake.lock b/flake.lock index 41de94a6..3c7b28e1 100644 --- a/flake.lock +++ b/flake.lock @@ -15,6 +15,28 @@ "type": "github" } }, + "dream2nix": { + "inputs": { + "nixpkgs": [ + "nixpkgs" + ], + "purescript-overlay": "purescript-overlay", + "pyproject-nix": "pyproject-nix" + }, + "locked": { + "lastModified": 1765953015, + "narHash": "sha256-5FBZbbWR1Csp3Y2icfRkxMJw/a/5FGg8hCXej2//bbI=", + "owner": "nix-community", + "repo": "dream2nix", + "rev": "69eb01fa0995e1e90add49d8ca5bcba213b0416f", + "type": "github" + }, + "original": { + "owner": "nix-community", + "repo": "dream2nix", + "type": "github" + } + }, "fenix": { "inputs": { "nixpkgs": [ @@ -36,6 +58,22 @@ "type": "github" } }, + "flake-compat": { + "flake": false, + "locked": { + "lastModified": 1696426674, + "narHash": "sha256-kvjfFW7WAETZlt09AgDn1MrtKzP7t90Vf7vypd3OL1U=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "0f9255e01c2351cc7d116c072cb317785dd33b33", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, "flake-parts": { "inputs": { "nixpkgs-lib": [ @@ -88,9 +126,54 @@ "type": "github" } }, + "purescript-overlay": { + "inputs": { + "flake-compat": "flake-compat", + "nixpkgs": [ + "dream2nix", + "nixpkgs" + ], + "slimlock": "slimlock" + }, + "locked": { + "lastModified": 1728546539, + "narHash": "sha256-Sws7w0tlnjD+Bjck1nv29NjC5DbL6nH5auL9Ex9Iz2A=", + "owner": "thomashoneyman", + "repo": "purescript-overlay", + "rev": "4ad4c15d07bd899d7346b331f377606631eb0ee4", + "type": "github" + }, + "original": { + "owner": "thomashoneyman", + "repo": "purescript-overlay", + "type": "github" + } + }, + "pyproject-nix": { + "inputs": { + "nixpkgs": [ + "dream2nix", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1763017646, + "narHash": "sha256-Z+R2lveIp6Skn1VPH3taQIuMhABg1IizJd8oVdmdHsQ=", + "owner": "pyproject-nix", + "repo": "pyproject.nix", + "rev": "47bd6f296502842643078d66128f7b5e5370790c", + "type": "github" + }, + "original": { + "owner": "pyproject-nix", + "repo": "pyproject.nix", + "type": "github" + } + }, "root": { "inputs": { "crane": "crane", + "dream2nix": "dream2nix", "fenix": "fenix", "flake-parts": "flake-parts", "nixpkgs": "nixpkgs", @@ -115,6 +198,28 @@ "type": "github" } }, + "slimlock": { + "inputs": { + "nixpkgs": [ + "dream2nix", + "purescript-overlay", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1688756706, + "narHash": "sha256-xzkkMv3neJJJ89zo3o2ojp7nFeaZc2G0fYwNXNJRFlo=", + "owner": "thomashoneyman", + "repo": "slimlock", + "rev": "cf72723f59e2340d24881fd7bf61cb113b4c407c", + "type": "github" + }, + "original": { + "owner": "thomashoneyman", + "repo": "slimlock", + "type": "github" + } + }, "treefmt-nix": { "inputs": { "nixpkgs": [ diff --git a/flake.nix b/flake.nix index a9ae9ca7..1c049537 100644 --- a/flake.nix +++ b/flake.nix @@ -21,6 +21,11 @@ inputs.nixpkgs.follows = "nixpkgs"; }; + dream2nix = { + url = "github:nix-community/dream2nix"; + inputs.nixpkgs.follows = "nixpkgs"; + }; + # Pinned nixpkgs for swift-format (swift is broken on x86_64-linux in newer nixpkgs) nixpkgs-swift.url = "github:NixOS/nixpkgs/08dacfca559e1d7da38f3cf05f1f45ee9bfd213c"; }; @@ -41,6 +46,7 @@ imports = [ inputs.treefmt-nix.flakeModule + ./dashboard/parts.nix ./rust/parts.nix ]; From 82ba42bae9fc111e25478a0a8c8f45fd95163b65 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Wed, 14 Jan 2026 16:33:17 +0000 Subject: [PATCH 4/4] add glm-47, minimax-m21 (#1147) Adds support glm 4.7 and MiniMax M2.1 Manual testing: Tensor + Pipeline execution of both models. Closes #1141 and #1142 --- src/exo/shared/models/model_cards.py | 77 +++++++++++++++++++++ src/exo/worker/engines/mlx/auto_parallel.py | 40 ++++++++++- 2 files changed, 116 insertions(+), 1 deletion(-) diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py index fd4f9003..d46a8115 100644 --- a/src/exo/shared/models/model_cards.py +++ b/src/exo/shared/models/model_cards.py @@ -82,6 +82,7 @@ MODEL_CARDS: dict[str, ModelCard] = { # storage_size=Memory.from_kb(754706307), # n_layers=61, # hidden_size=7168, + # supports_tensor=True, # ), # ), # "deepseek-v3.2-4bit": ModelCard( @@ -96,6 +97,7 @@ MODEL_CARDS: dict[str, ModelCard] = { # storage_size=Memory.from_kb(754706307 // 2), # TODO !!!!! # n_layers=61, # hidden_size=7168, + # supports_tensor=True, # ), # ), # deepseek r1 @@ -554,6 +556,81 @@ MODEL_CARDS: dict[str, ModelCard] = { supports_tensor=True, ), ), + "glm-4.7-4bit": ModelCard( + short_id="glm-4.7-4bit", + model_id=ModelId("mlx-community/GLM-4.7-4bit"), + name="GLM 4.7 4bit", + description="GLM 4.7 4bit", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/GLM-4.7-4bit"), + pretty_name="GLM 4.7 4bit", + storage_size=Memory.from_bytes(198556925568), + n_layers=91, + hidden_size=5120, + supports_tensor=True, + ), + ), + "glm-4.7-6bit": ModelCard( + short_id="glm-4.7-6bit", + model_id=ModelId("mlx-community/GLM-4.7-6bit"), + name="GLM 4.7 6bit", + description="GLM 4.7 6bit", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/GLM-4.7-6bit"), + pretty_name="GLM 4.7 6bit", + storage_size=Memory.from_bytes(286737579648), + n_layers=91, + hidden_size=5120, + supports_tensor=True, + ), + ), + "glm-4.7-8bit-gs32": ModelCard( + short_id="glm-4.7-8bit-gs32", + model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"), + name="GLM 4.7 8bit (gs32)", + description="GLM 4.7 8bit (gs32)", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/GLM-4.7-8bit-gs32"), + pretty_name="GLM 4.7 8bit (gs32)", + storage_size=Memory.from_bytes(396963397248), + n_layers=91, + hidden_size=5120, + supports_tensor=True, + ), + ), + "minimax-m2.1-8bit": ModelCard( + short_id="minimax-m2.1-8bit", + model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"), + name="MiniMax M2.1 8bit", + description="MiniMax M2.1 8bit", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/MiniMax-M2.1-8bit"), + pretty_name="MiniMax M2.1 8bit", + storage_size=Memory.from_bytes(242986745856), + n_layers=61, + hidden_size=3072, + supports_tensor=True, + ), + ), + "minimax-m2.1-3bit": ModelCard( + short_id="minimax-m2.1-3bit", + model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"), + name="MiniMax M2.1 3bit", + description="MiniMax M2.1 3bit", + tags=[], + metadata=ModelMetadata( + model_id=ModelId("mlx-community/MiniMax-M2.1-3bit"), + pretty_name="MiniMax M2.1 3bit", + storage_size=Memory.from_bytes(100086644736), + n_layers=61, + hidden_size=3072, + supports_tensor=True, + ), + ), # "devstral-2-123b-instruct-2512-8bit": ModelCard( # short_id="devstral-2-123b-instruct-2512-8bit", # model_id=ModelId("mlx-community/Devstral-2-123B-Instruct-2512-8bit"), diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index 55a549fd..46d812e3 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -19,6 +19,7 @@ from mlx_lm.models.glm4_moe import MoE from mlx_lm.models.gpt_oss import GptOssMoeModel from mlx_lm.models.gpt_oss import Model as GptOssModel from mlx_lm.models.llama import Model as LlamaModel +from mlx_lm.models.minimax import Model as MiniMaxModel from mlx_lm.models.ministral3 import Model as Ministral3Model from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel from mlx_lm.models.qwen3_moe import Qwen3MoeSparseMoeBlock @@ -252,6 +253,14 @@ def tensor_auto_parallel( all_to_sharded_linear_in_place, sharded_to_all_linear_in_place, ) + elif isinstance(model, MiniMaxModel): + tensor_parallel_sharding_strategy = MiniMaxShardingStrategy( + group, + all_to_sharded_linear, + sharded_to_all_linear, + all_to_sharded_linear_in_place, + sharded_to_all_linear_in_place, + ) elif isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)): tensor_parallel_sharding_strategy = QwenShardingStrategy( group, @@ -394,6 +403,35 @@ class ShardedDeepseekV3MoE(CustomMlxLayer): return y +class MiniMaxShardingStrategy(TensorParallelShardingStrategy): + def shard_model(self, model: nn.Module) -> nn.Module: + model = cast(MiniMaxModel, model) + for layer in model.layers: + # Shard the self attention + layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj) + layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj) + layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj) + layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj) + layer.self_attn.num_attention_heads //= self.N + layer.self_attn.num_key_value_heads //= self.N + + # Shard the MoE. Shard in place since the MoE should be responsible + # for aggregating the results. + self.all_to_sharded_linear_in_place( + layer.block_sparse_moe.switch_mlp.gate_proj + ) + self.sharded_to_all_linear_in_place( + layer.block_sparse_moe.switch_mlp.down_proj + ) + self.all_to_sharded_linear_in_place( + layer.block_sparse_moe.switch_mlp.up_proj + ) + layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] + layer.block_sparse_moe.sharding_group = self.group + + return model + + class QwenShardingStrategy(TensorParallelShardingStrategy): def shard_model(self, model: nn.Module) -> nn.Module: model = cast(Qwen3MoeModel, model) @@ -414,7 +452,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy): self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj) self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj) self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj) - layer.mlp = ShardedQwenMoE(layer.mlp) # type: ignore + layer.mlp = ShardedQwenMoE(layer.mlp) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType] layer.mlp.sharding_group = self.group # Shard the MLP