← Benchmark results · Markdown source

RFC: stream MoE experts from disk in from_pretrained

Proposal for huggingface/transformers. Implementation: huggingface/transformers#50717 (proof of concept). Numbers refer to the results page.

Feature request

Let from_pretrained keep MoE experts on disk and stream them through a bounded device cache, so large MoE models run on machines that cannot hold their experts. The options depend on the platform:

# Apple silicon (unified memory): 8-bit dense layers, experts streamed from the SSD
model = AutoModelForCausalLM.from_pretrained(
    "qwen38-flash-next/dense",
    dtype=torch.bfloat16,
    device_map="mps",
    quantization_config=MetalConfig(bits=8),  # dense layers only
    experts_offload={
        "cache_bytes": 12 << 30,  # device cache = host RAM, no host tier needed
        "quantized_experts": "qwen38-flash-next/experts-q3",
        "mps_commit": True,  # commit queued MPS work early, so the GPU keeps busy during host syncs
    },
)

# CUDA (discrete GPU): BF16 dense layers, a small device cache backed by a host tier in pinned RAM
model = AutoModelForCausalLM.from_pretrained(
    "qwen38-flash-next/dense",
    dtype=torch.bfloat16,
    device_map="cuda",
    experts_offload={
        "cache_bytes": 8 << 30,  # VRAM
        "quantized_experts": "qwen38-flash-next/experts-q3",
        "host_cache_bytes": 50 << 30,  # pinned host RAM, holds all 49 GB of 3-bit experts
        "host_preload": True,  # fill the host tier at load time
    },
)

Motivation

Qwen3.8-Flash-Next is a 48-layer MoE with 512 experts per layer. Its experts alone are 121 GB in FP8, next to ~10 GB of dense weights and a 51 GB n-gram embedding table, so it cannot be loaded on a 32 GB laptop: not by vanilla Transformers, and not by device_map offloading either. Yet each token only uses 10 of the 512 experts of a layer (2%), and dedicated engines already run this model on a 32 GB Mac by streaming experts from the SSD: Strata reports 9.2 tok/s on an M4 Mac mini with IQ3_XXS experts, the build closest in quality to our 3-bit ones (13.5 with Q2_0, at ~10% worse perplexity).

With a prototype of this feature, Transformers runs it on a 32 GB MacBook Pro M5:

• experts quantized once to 3-bit (49 GB on disk) and streamed through a 12 GiB device cache, the n-gram table read from disk row by row, dense layers in 8-bit with MetalConfig;
• 17.5 GiB of device memory after loading, the cache included;
• 10.0 tok/s decode (median of four different 200-token answers, starting from a cold cache; 10.7 on the story alone; on an M5, which has ~28% more memory bandwidth than the M4), and 67–70 tok/s prefill on a 4K-token prompt (1K chunks, 10 GiB cache, cold);
• perplexity unchanged with 8-bit dense layers (4.901 vs 4.901 with BF16 dense layers, same 3-bit experts).

The same prototype runs on CUDA (one A10G, 24 GB, a Triton kernel for the quantized experts): 4.1 tok/s decode with experts read from a network file system, 5.8 tok/s with every expert preloaded in the host tier (BF16 dense layers, 8 GiB device cache).

Full results, the architecture with every kernel, and raw data: https://huggingface.co/spaces/tarekziade/moe-expert-streaming

Accelerate disk offload is not an alternative here: it moves whole modules per forward call, so every token would read every expert of every layer (in earlier cloud runs, 0.16 tok/s for the much smaller Qwen3.6-35B-A3B), while an MoE token only needs its top-k experts.

How the prototype works

• Loading: the model is built on meta; every module decorated with @use_experts_implementation (gated, concatenated, untransposed, no bias) gets its expert banks removed and switches to a ssd_streaming experts implementation. The loader skips the expert tensors (a new LoadStateDictConfig.skip_keys); everything else loads as usual, and a MetalConfig quantizes the dense layers. About 20 lines in modeling_utils.py; the rest lives in integrations/ssd_moe.py.
• Cache: one preallocated device pool of expert slots shared by all layers (LRU, optional LFU), with a per-layer expert → slot table kept on the host. Quantized experts run through gather_qmm from kernels-community/mlx-quantization-metal-kernels on the pool (BF16 experts through the existing grouped_mm experts kernel). When every expert is resident, layers never synchronize with the host.
• Reads: each layer copies its routing to the host once; misses are read on I/O threads into recycled pinned buffers, one contiguous read and one non-blocking host-to-device copy per expert. The next two layers' experts are guessed by running their routers on the current hidden states and read in the background.
• Host tier (optional): host_cache_bytes adds a pool of expert rows in pinned host memory between storage and the device cache, so a device-cache miss costs one host-to-device copy instead of a storage read; host_preload=True fills it at load time. It is for discrete GPUs with more host than device memory; on unified memory the device cache already is host RAM.
• Prefill: a long prompt routes to nearly every expert; they run through the sorted gather_qmm in groups of 64 experts so the transient bank stays small.
• Quantized experts: an optional one-time quantize_experts(checkpoint, out, bits) writes the experts in MLX's affine format, one contiguous row per expert.
• Disk embeddings: disk_embeddings={module: directory} replaces a huge embedding table (Qwen3.8's n-gram table) with row reads from disk.

Questions for maintainers

1. API shape: a plain experts_offload={...} kwarg, an ExpertsOffloadConfig class (like GenerationConfig), or a device_map value such as "experts:disk"?
2. Quantized experts: keep them a property of this feature (quantized_experts=path), or route them through the quantizer framework (a quantizer that only targets experts, since MetalConfig only converts nn.Linear today)?
3. Scope: start with inference on a single device (MPS / CUDA / CPU) and gated concatenated experts, and leave training, multi-device and save_pretrained unsupported, with clear errors?
4. Where it lives: integrations/ssd_moe.py plus the experts_offload hook in from_pretrained, or in Accelerate? And do disk-backed embeddings belong to this feature or to a separate one?

Happy to open a PR once the API direction is settled.