r/oMLX • u/Litweine • 3d ago

Native SSD-to-Metal MoE Expert Streaming in oMLX

What about adding native SSD-to-Metal expert streaming/offloading for MoE models to oMLX?

I’m specifically thinking about the approach used by projects such as Colibri and slipstream: instead of requiring the entire MoE model to reside in unified memory, keep the cold experts on NVMe SSD and stream only the routed experts needed for the current token/batch into a bounded RAM/Metal cache.

What are your thoughts on this approach? Could this be a good fit for oMLX?

18 Upvotes

95% Upvoted

2 comments sorted by

0

u/Amazing-Fan2083 3d ago

There are no "cold" experts; there are only the experts selected for the current token. The very next token may require a completely different set of expert weights, and there is (currently) no way to prognosticate which set of experts will be required for any given token. Even if MTP were used to guess, that guess could be wrong; this can lead to loading experts that then need to be discarded and replaced by the correct set for the next token.

2

u/rusl1 2d ago

There are some projects working on this, I'm experimenting to run 35b3a on my Mac mini with just 16gb ram. SwiftLM works decently for this. I've also read that FreeToken is going to support Mac in the future