Repository navigation
Experimental: run MoE models that don't fit with experts on disk - #394
Merged
Merged
Conversation
Settings › Server › Experimental › "Run big MoE models with experts on disk", off by default. With it on, a MoE model that doesn't fit the GPU limit (less than GPUFit's 2.5 GB left beside its weights) starts with the runtime's --experts-on-disk: the share of its routed experts kept in memory is what the limit leaves after the rest of the weights, the margin and 2.5 GB for its work, in 5% steps, 10-90%. The rest is read from the SSD as needed (ipsupport-llc/mlx-lm#56, pinned). - The experts' size comes from the safetensors headers (stacked switch_mlp/switch_glu tensors with the experts first); a dense model or one without stacked experts is left as it is. - The prompt cache and prefill are sized by what stays resident. - Only when the runtime has the flag. Gemma 4 26B-A4B with a quarter of its experts: 7.8 tok/s at a 6.1 GB peak (it fits on its own, so the switch leaves it alone). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XhEwWJCoJPXFyyCzp988Bx
Review (codex, muse): - Only quantized gate/up/down projections of a stacked SwitchGLU count (Nemotron's fc1/fc2, unquantized or biased experts load whole), and a module's tensors are gathered across shards before that check: Gemma 4 26B-A4B counted 11.7 of its 12.85 GB, its modules split over files. - No share is offered when even 10% of the experts doesn't fit beside the rest: it used to be clamped up to 10% and launched anyway. Checked on real folders: Gemma 4 26B-A4B 12.85 GB, Ornith 35B-A3B 11.86 GB, Nemotron 30B-A3B and the dense Muse Glimmer 0. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XhEwWJCoJPXFyyCzp988Bx
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Settings › Server › Experimental › "Run big MoE models with experts on disk", off by default.
What it does
--experts-on-disk SHARE(Experimental: routed experts on disk (--experts-on-disk SHARE) mlx-lm#56, pinned).Which experts count: only those the runtime offloads, read from the safetensors headers.
Measured on Gemma 4 26B-A4B with a quarter of its experts kept, through the runtime and its server: 7.8 tok/s at a 6.1 GB peak. That model fits this Mac on its own, so the switch leaves it alone here; a model over the limit gets the share.
Settings section rendered in dark mode.
Tests:
ExpertsOnDiskTests(header counting, the share rules, the launch argument). 1389 pass.Reviews (codex, muse), fixed in the second commit:
🤖 Generated with Claude Code
https://claude.ai/code/session_01XhEwWJCoJPXFyyCzp988Bx