WASTE: Running Massive AI Models on Limited Hardware

The team behind SQLite Cloud has just released WASTE, an open-source inference engine designed to run massive machine learning models on hardware with limited memory. This is a significant step forward in making state-of-the-art AI accessible to more users without requiring expensive cloud infrastructure.

WASTE's first fully supported model is Kimi K3, a 2.78 trillion parameter mixture-of-experts (MoE) model. The original 1.42 TB checkpoint was converted into a 982 GB container and successfully ran on a MacBook Pro with just 64 GB of unified memory. This isn't a distilled or pruned version; it's the full model.

Here’s what a typical execution looks like:

$ waste run ~/models/k3.waste "What is the capital of Italy?"

waste: no --budget, using 46.24 GB of 64.00 GB
       (expert cache 17.56 GB)

The capital of Italy is Rome.

[16 tokens, 49.31 s, 0.32 tok/s]

While the current performance of about one-third of a token per second isn't ideal for most interactive applications, it demonstrates that memory constraints don’t have to be a hard limit on model size.

Under the Hood

Kimi K3 is an MoE model, which means it contains 896 routed experts. However, only 16 experts are selected in each layer for a given token. This property allows WASTE to separate the model into two main parts: the model trunk and the expert weights.

• Model Trunk: Remains resident in memory.
• Expert Weights: Stored on NVMe storage and loaded when needed by the router.
• Bounded Expert Cache: Any remaining RAM is used as a cache for frequently accessed experts.

This architecture changes the primary constraint from fitting the entire model into memory to retrieving the active portion of the model within an acceptable time. For K3, each token accesses approximately 17 GB of expert data, making storage bandwidth and cache behavior central to performance.

WASTE stores experts in a layout optimized for this access pattern. The matrices required for one expert are adjacent, allowing them to be loaded with a single positional read. This reduces the overhead of fetching weights from storage.

Key Takeaways

• Memory Efficiency: WASTE demonstrates that massive models can run on consumer hardware by leveraging NVMe storage and a bounded cache.
• Performance Trade-offs: While current performance isn't suitable for real-time applications, it opens up new possibilities for running large models locally.
• Control and Privacy: By enabling local execution, WASTE provides greater control over infrastructure costs, data privacy, availability, and deployment.

WASTE is an initial step in a broader effort to make increasingly capable AI models available on hardware controlled by the people and organizations using them. This could lead to more decentralized and accessible AI, reducing dependency on cloud providers and enhancing user control over their data and resources.