Repository navigation
Experimental: routed experts on disk (--experts-on-disk SHARE) - #56
Merged
Merged
Conversation
A MoE model that doesn't fit memory runs with its routed experts left in the checkpoint: each SwitchGLU keeps SHARE of its experts in an LRU cache and reads the others with pread (page cache bypassed) when the router picks them. A decoding step computes each expert with quantized_matmul on its own arrays; a prompt or a batch runs gather_qmm on a stack of just the experts it uses. Stacked checkpoints keep each expert in one slice per tensor; file names are matched to modules by their layers.N tail. Gemma 4 26B-A4B (15.6 GB) with 25%: 7.8 tok/s, 6.1 GB peak; the server answers single and batched requests. Research and measurements: rromenskyi/quant-ternary moe-offload (Ornith 35B-A3B 7.4 tok/s in 5.9 GB). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XhEwWJCoJPXFyyCzp988Bx
- A module maps to a checkpoint tensor only when its layers.N tail and the names before it match exactly (in any order); otherwise it stays in memory. An MTP head's layer 0 no longer reads the model's. - Experts with a linear bias (GPT-OSS) stay in memory. - The cache never holds more than its share, also when one step needs more experts. - A lock per layer around the cache and the prompt path's weight swap. - The reader's pool and file descriptors close with the model. - Page cache bypass only on macOS (F_NOCACHE). - A share of 0 still leaves experts on disk; adapters and distributed serving refuse the option instead of ignoring it. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XhEwWJCoJPXFyyCzp988Bx
The option is read from the load's model_config only, never from a checkpoint's config.json, so the adapter and distributed refusals can't be bypassed. Names before layers.N compared with their counts (model.model is not model). The reader's pool waits for its threads before its files close. 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.
A MoE model that doesn't fit memory runs with its routed experts left in the checkpoint.
How it works
SwitchGLUkeeps SHARE of its experts in an LRU cache. The rest are read withpreadwhen the router picks them (on macOS the reads bypass the page cache).quantized_matmulon its own arrays.gather_qmmon a stack of just the experts it uses.layers.Ntail and the names before it; anything ambiguous stays in memory.Measured on Gemma 4 26B-A4B (15.6 GB) with 25% kept: 7.8 tok/s at a 6.1 GB peak.
moe-offload(Ornith 35B-A3B: 7.4 tok/s in 5.9 GB, against 20.6 tok/s in 15.2 GB).Tests
test_expert_offload.py: the same logits as in memory at 100% and 25% over a prompt and 6 decode steps; the cache stays within its share; tensor-name mapping, including an MTP head's layer 0.Reviews (codex, muse), fixed in the second commit
fcntl.🤖 Generated with Claude Code
https://claude.ai/code/session_01XhEwWJCoJPXFyyCzp988Bx