Expert Offloading to CPU or NVMe

The immense parameter count of a Mixture of Experts model is its greatest strength and its most significant deployment challenge. A model like Mixtral 8x7B contains weights for eight distinct experts per MoE layer, leading to a total parameter size that far exceeds the capacity of all but the largest and most expensive GPU accelerators. However, during any given forward pass, only one or two of these experts are activated per token. This computational sparsity is the main factor in managing inference. Expert offloading exploits this property by storing the majority of inactive expert parameters in a larger, more economical memory pool, such as system RAM or NVMe storage.

The fundamental principle is to treat GPU VRAM as a high-speed, limited-capacity cache for expert weights. Instead of loading the entire model onto the GPU, we only load the non-expert layers (the model's backbone) and the gating networks. The expert weights themselves reside "off-chip" and are dynamically moved into VRAM on-demand.

The Memory vs. Latency Trade-off

Offloading is not a free lunch. It resolves the VRAM capacity issue at the cost of introducing data transfer latency. Moving gigabytes of expert weights from CPU RAM or an NVMe drive to the GPU over the PCIe bus takes time, an operation that is orders of magnitude slower than accessing data already in the GPU's high-bandwidth memory (HBM).

The performance of an offloaded system is therefore governed by this trade-off. The goal is to design a system that minimizes these data transfers and hides their latency as much as possible.

Data flow for an offloaded expert forward pass. When a required expert is not in the on-GPU cache, its weights must be transferred from system memory over the PCIe bus, introducing latency.

Offloading Strategies

Offloading to CPU RAM

This is the most common and balanced approach. System RAM is significantly larger and cheaper than GPU VRAM, while still offering reasonable transfer speeds.

The workflow is as follows:

1. Initialization: The model's shared components (attention layers, embeddings, normalization layers) and all gating networks are loaded into GPU VRAM. All expert weights are loaded into CPU RAM.
2. Gating: For an incoming batch of tokens, the gating network on the GPU runs and determines which experts are needed for each token.
3. Transfer: The system identifies the set of unique experts required for the entire batch. For each required expert not already present on the GPU, an asynchronous memory copy is initiated to transfer the weights from CPU RAM to an allocated buffer in GPU VRAM.
4. Computation: Once the expert weights are on the GPU, the corresponding tokens are dispatched to them for computation.
5. Eviction: After computation, the expert weights can be kept on the GPU if memory permits or evicted to make space for the next required experts.

Using asynchronous copy operations (e.g., torch.Tensor.to('cuda', non_blocking=True) in PyTorch) is important, as it allows the CPU to prepare the next transfer while the GPU is busy with other work, partially hiding the transfer latency.

Offloading to NVMe Storage

For extremely large models that may not even fit into system RAM, or on hardware with limited RAM, offloading can be extended to high-speed NVMe SSDs. This provides access to terabytes of storage but comes with a severe latency penalty.

The data path becomes longer: NVMe -> CPU RAM -> GPU VRAM. While technologies like GPUDirect Storage can create a more direct path from NVMe to GPU, they add system complexity. This strategy is typically reserved for offline, throughput-oriented batch processing jobs where per-token latency is not the primary concern.

Improving Performance with Caching

The latency penalty of offloading can be substantially reduced by implementing a cache for expert weights in GPU VRAM. A simple Least Recently Used (LRU) caching policy is highly effective. The GPU allocates a portion of its VRAM to hold a small number of experts.

When an expert is required:

• Cache Hit: If the expert is already in the GPU cache, it is used immediately, and no data transfer occurs.
• Cache Miss: If the expert is not in the cache, the system must fetch it from the offload location (CPU RAM). If the cache is full, the least recently used expert is evicted from the GPU to make room for the incoming one.

The size of this cache is a critical hyperparameter. A larger cache increases the probability of a cache hit, reducing average latency, but it also consumes more of the precious GPU VRAM that could be used for larger batches.

Impact of on-GPU expert caching on inference latency. Even a small cache that holds a fraction of the total experts can dramatically reduce the average latency by avoiding frequent data transfers over the PCIe bus.

An Implementation Example

To make this concrete, consider a simplified OffloadedMoE layer in Python. This example shows the core logic of checking a cache and loading an expert on a miss.

import torch

class OffloadedMoE(torch.nn.Module):
    def __init__(self, experts, cache_size=4):
        super().__init__()
        # Experts are initially on CPU
        self.experts_cpu = torch.nn.ModuleList(experts)
        self.num_experts = len(experts)

        # GPU cache management
        self.cache_size = cache_size
        self.expert_cache_gpu = {} # Maps expert_id to GPU tensor
        self.expert_cache_lru = [] # Stores expert_ids in usage order

    def _load_expert_to_gpu(self, expert_id):
        # Evict if cache is full
        if len(self.expert_cache_lru) >= self.cache_size:
            evict_id = self.expert_cache_lru.pop(0)
            del self.expert_cache_gpu[evict_id]

        # Load expert from CPU to GPU
        expert = self.experts_cpu[expert_id]
        self.expert_cache_gpu[expert_id] = expert.to('cuda', non_blocking=True)
        self.expert_cache_lru.append(expert_id)

    def forward(self, x, gating_output):
        # gating_output contains router decisions, e.g., expert indices
        required_ids = torch.unique(gating_output.top_k_indices).tolist()

        # Ensure all required experts are loaded
        for expert_id in required_ids:
            if expert_id not in self.expert_cache_gpu:
                self._load_expert_to_gpu(expert_id)
            else:
                # Move to end of LRU list to mark as recently used
                self.expert_cache_lru.remove(expert_id)
                self.expert_cache_lru.append(expert_id)

        # ... logic to dispatch tokens to the correct experts on GPU ...
        # final_output = perform_expert_computation(x, self.expert_cache_gpu)
        # return final_output

This example omits the complex token dispatch logic but shows the caching mechanism. A production system like deepspeed-mii or Hugging Face's accelerate provides optimized implementations of this pattern. By combining intelligent caching with asynchronous transfers, expert offloading makes it practical to run massive MoE models on commodity hardware, democratizing access to their capabilities.