mlx-moe-offload

KO | EN

gitlite — search

#csharp #java #artificial-intelligence #json #electron #bash #tailwindcss #android #tensorflow #list

mlx-moe-offload

(★ 9)

No description available.

File Explorer

ZIP Download

- demo.py
- inkling_omni.py
- __init__.py
- offload.py
- repack.py
- run.py
- expert_offload.ipynb
- offload-diagram.svg
- test_offload.py
.gitignore
LICENSE
pyproject.toml
README.md

# Use via CDN

jsDelivr

jsDelivr serves any public GitHub repository as a CDN with zero setup. Pick a version and a file to get a ready-to-paste link and snippet.

Version

File

Link

Example

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

Run Mixture-of-Experts LLMs that are **bigger than your RAM** on Apple Silicon, by
paging routed experts from SSD on demand. We plan to have a small, self-contained layer on top of
[`mlx-lm`](https://github.com/ml-explore/mlx-lm) and having a PR on it.

An MoE only fires a few experts per token (e.g. **6 of 256**), yet the experts are
~95% of the weights. So keep the always-needed parts resident in unified memory and
**stream the selected experts from SSD**, letting the OS page cache hold the hot ones.
This is the "LLM in a flash" idea, expressed for MLX's `SwitchGLU`.

![How mlx-moe-offload works](https://github.com/huckiyang/mlx-moe-offload/blob/main/tests/offload-diagram.svg)

## Why

| | Resident (RAM) | Paged (SSD) |
|---|---|---|
| what | attention, shared experts, embeddings, norms, router | the 256 routed experts / layer |
| size (Inkling-2bit) | ~35 GB | ~270 GB |
| touched per token | all | only top-6 → the rest stay on disk |

Result: a 315 GB 2-bit model becomes runnable on a **192 GB Mac** (RAM = resident +
hot-expert cache; the cold experts live on SSD). Reads don't wear SSDs, and the mmap'd
expert pages are clean/file-backed so they're reclaimed without swap writes.

## Install

```bash
pip install mlx mlx-lm            # deps
pip install -e .                  # this package
```

## Use (two steps)

```bash
# 1) one-time: repack a build's stacked experts into a per-expert SSD store
python -m mlx_moe_offload.repack --build <mlx_moe_dir> --out <offload_dir>

# 2) run with experts paged from disk
python examples/demo.py                              # tiny, runnable, explains the idea
python -m mlx_moe_offload.run --offload <offload_dir> --prompt "..."   # real models
```

```python
import mlx_lm
from mlx_moe_offload import patch_model

model, tok = mlx_lm.load("<offload_dir>")             # loads only the resident weights
store = patch_model(model, "<offload_dir>")           # experts now page from SSD
# keep prefill_step_size small: a lazy full-prompt prefill fires ~all experts per layer at
# once and pins them until eval -> OOM. Small blocks (mlx-lm evals the cache between them)
# bound the live expert set to ~step*top_k per layer.
print(mlx_lm.generate(model, tok, "The capital of France is",
                      max_tokens=64, prefill_step_size=8, verbose=True))
print(store.stats())                                  # {'hit_rate': ...}
```

Works with any `mlx-lm` MoE that uses `SwitchGLU` (Mixtral, Qwen3-MoE, DeepSeek-V3, GLM-MoE, …).

## Omni Inkling (text + image + audio)

The flagship is **Inkling-2bit** ([mlx-community/Inkling-mlx-2bit](https://huggingface.co/mlx-community/Inkling-mlx-2bit)):
975B total / 41B active, 256 routed experts/layer, ~315 GB. It has a **custom architecture**
(not `mlx_lm.load`-able) plus vision + audio towers, so it runs through its bundled `inkling_mlx`
loader in [`examples/inkling_omni.py`](examples/inkling_omni.py). `repack` copies that loader into
the offload dir, so the dir stays self-contained. Extra deps: `scipy` (audio), `pillow` (image).

```bash
python -m mlx_moe_offload.repack --build ./Inkling-mlx-2bit --out ./Inkling-mlx-2bit-offload
OFF=./Inkling-mlx-2bit-offload

python examples/inkling_omni.py --offload-dir $OFF --prompt "The capital of France is"
python examples/inkling_omni.py --offload-dir $OFF --image cat.png --prompt "What is in this image?"
python examples/inkling_omni.py --offload-dir $OFF --audio q.wav   --prompt "Transcribe in English:"
```

`--prefill-chunk` is the multimodal analog of `prefill_step_size`; the media towers run once, then
the backbone is fed in blocks. Observed on an M4 Max (64 GB): text → *"Tokyo"* ✓; image (Pikachu)
→ coherent but misID'd; audio (*"when does Cafe Macs close"*) → phonetically right but German-worded
— that quality gap is **2-bit quantization**, not the offload path (a 4/6-bit build fixes it).

## How it works

- `repack.py` splits each layer's stacked expert tensor `[E, out, in]` into per-expert
  entries in a memory-mappable per-layer file.
- `ExpertStore` `mx.load`s those files lazily (mmap). Materializing one expert pages in
  only that expert's bytes; the OS page cache keeps hot experts (optional explicit LRU).
- `OffloadedSwitchGLU` replaces `mlx-lm`'s `SwitchGLU`: given the router's `indices`, it
  computes **only** the selected experts, so cold experts never enter RAM.

## Caveats

- **Speed is gated by the routing cache hit-rate** and SSD bandwidth. Uniformly-routed
  models (Inkling) cache poorly; skewed routing caches well. Measure with `store.stats()`.
  Decode is ~0.03–0.4 tok/s single-stream when the store doesn't fit the page cache.
- **Prefill must be chunked**, or a lazy full-prompt forward pins ~all experts across all
  layers at once and OOMs — use small `prefill_step_size` (mlx-lm) / `--prefill-chunk` (Inkling).
- Decode-oriented (per-token expert grouping is python-side); high-throughput prefill
  wants a batched gather.
- Keep resident + cache under your RAM ceiling. On a Mac that's *very* tight for the model
  (e.g. 64 GB for a 315 GB build), long multimodal runs can be **jetsam-killed** even at
  moderate pressure — keep prompts/outputs short, or use the intended 192 GB Mac Studio.
  mmap pages are clean, so eviction is free: no swap writes → no SSD wear.

## Prior art

"[LLM in a flash](https://arxiv.org/abs/2312.11514)"; MoE expert-offload for Mixtral et al.
This packages the idea as a drop-in for MLX `SwitchGLU`, and answers the mlx-lm feature request
[ml-explore/mlx-lm#1438](https://github.com/ml-explore/mlx-lm/issues/1438) (stream MoE experts
from SSD to run models bigger than unified memory).

## License

MIT. See `LICENSE`.

mlx-moe-offload

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

mlx-moe-offload

Owner Repos

File Explorer

# Use via CDN

Was this content helpful?

Recent Feedback

Download README