5 ms·
I follow the MLX team on Twitter and they sometimes post about using MLX on two or more joined together Macs to run models that need more than 512GB of RAM. A
by simonw 11mo ago
I follow the MLX team on Twitter and they sometimes post about using MLX on two or more joined together Macs to run models that need more than 512GB of RAM.
A couple of examples:
Kimi K2 Thinking (1 trillion parameters): https://x.com/awnihannun/status/1986601104130646266 https://x.com/awnihannun/status/1986601104130646266
DeepSeek R1 (671B): https://x.com/awnihannun/status/1881915166922863045 https://x.com/awnihannun/status/1881915166922863045 - that one came with setup instructions in a Gist: https://gist.github.com/awni/ec071fd27940698edd14a4191855bba6 https://gist.github.com/awni/ec071fd27940698edd14a4191855bba...
- awnihannun 11mo agoFor a bit more context, those posts are using pipeline parallelism. For N machines put the first L/N layers on machine 1, next L/N layers on machine 2, etc. With pipeline parallelism you don't get a speedup over one machine - it just buys you the ability to use larger models than you can fit on a single machine. The release in Tahoe 26.2 will enable us to do fast tensor parallelism in MLX. Each layer of the model is sharded across all machines. With this type of parallelism you can get close to N-times faster for N machines. The main challenge is latency since you have to do much more frequent communication.
- liuliu 11mo agoBut that's only for prefilling right? Or is it beneficial for decoding too (I guess you can do KV lookup on shards, not sure how much speed-up that will be though).
- zackangelo 11mo agoNo you use tensor parallelism in both cases. The way it typically works in an attention block is: smaller portions of the Q, K and V linear layers are assigned to each node and are processed independently. Attention, rope norm etc is run on the node-specific output of that. Then, when the output linear layer is applied an "all reduce" is computed which combines the output of all the nodes. EDIT: just realized it wasn't clear -- this means that each node ends up holding a portion of the KV cache specific to its KV tensor shards. This can change based on the specific style of attention (e.g., in GQA where there are fewer KV heads than ranks you end up having to do some replication etc)
- liuliu 11mo agoI usually call it "head parallelism" (which is a type of tensor parallelism, but paralllelize for small clusters, and specific to attention). That is what you described: sharding input tensor by number of heads and send to respective Q, K, V shard. They can do Q / K / V projection, rope, qk norm whatever and attention all inside that particular shard. The out projection will be done in that shard too but then need to all reduce sum amongst shard to get the final out projection broadcasted to every participating shard, then carry on to do whatever else themselves. I am asking, however, is whether that will speed up decoding as linearly as it would for prefilling.
- awnihannun 11mo agoRight, my comment was mostly about decoding speed. For prefill you can get a speed up but there you are less latency bound. In our benchmarks with MLX / mlx-lm it's as much as 3.5x for token generation (decoding) at batch size 1 over 4 machines. In that case you are memory bandwidth bound so sharding the model and KV cache 4-ways means each machine only needs to access 1/4th as much memory.
- liuliu 11mo agoOh! That's great to hear. Congrats! Now, I want to get the all-to-all primitives ready in s4nnc...
- monster_truck 11mo agoEven if it wasn't outright beneficial for decoding by itself, it would still allow you to connect a second machine running a smaller, more heavily quantized version of the model for speculative decoding which can net you >4x without quality loss
- dpe82 11mo ago> The main challenge is latency since you have to do much more frequent communication. Earlier this year I experimented with building a cluster to do tensor parallelism across large cache CPUs (AMD EPYC 7773X have 768mb of L3). My thought was to keep an entire model in SRAM and take advantage of the crazy memory bandwidth between CPU cores and their cache, and use Infiniband between nodes for the scatter/gather operations. Turns out the sum of intra-core latency and PCIe latency absolutely dominate. The Infiniband fabric is damn fast once you get data to it, but getting it there quickly is a struggle. CXL would help but I didn't have the budget for newer hardware. Perhaps modern Apple hardware is better for this than x86 stuff.
- wmf 11mo agoThat's how Groq works. A cluster of LPUv2s would probably be faster and cheaper than an Infiniband cluster of Epycs.
- fooblaster 11mo agowhat is an lpuv2
- wmf 11mo agoThe chip that Groq makes.
- dpe82 11mo agoYeah I'm familiar; I was hoping I could do something related on previous generation commodity(ish) hardware. It didn't work but I learned a ton.
- deleted 11mo ago[deleted]
- aimanbenbaha 11mo agoExo-Labs is an open source project that allows this too, pipeline parallelism I mean not the latter, and it's device agnostic meaning you can daisy-chain anything you have that has memory and the implementation will intelligently shard model layers across them, though its slow but scales linearly with concurrent requests. Exo-Labs: https://github.com/exo-explore/exo https://github.com/exo-explore/exo
- andy99 11mo agoI’m hoping this isn’t as attractive as it sounds for non-hobbyists because the performance won’t scale well to parallel workloads or even context processing, where parallelism can be better used. Hopefully this makes it really nice for people that want the experiment with LLMs and have a local model but means well funded companies won’t have any reason to grab them all vs GPUs.
- codazoda 11mo agoI haven’t looked yet but I might be a candidate for something like this, maybe. I’m RAM constrained and, to a lesser extent, CPU constrained. It would be nice to offload some of that. That said, I don’t think I would buy a cluster of Macs for that. I’d probably buy a machine that can take a GPU.
- ChrisMarshallNY 11mo agoI’m not particularly interested in training models, but it would be nice to have eGPUs again. When Apple Silicon came out, support for them dried up. I sold my old BlackMagic eGPU. That said, the need for them also faded. The new chips have performance every bit as good as the eGPU-enhanced Intel chips.
- andy_ppp 11mo agoeGPU with an Apple accelerator with a bunch or RAM and GPU cores could be really interesting honestly. I’m pretty sure they are capable of designing something very competitive especially in terms of performance per watt.
- CamperBob2 11mo agoAlmost the most impressive thing about that is the power consumption. ~50 watts for both of them? Am I reading it wrong?
- wmf 11mo agoYeah, two Mac Studios is going to be ~400 W.
- CamperBob2 11mo agoWhat am I missing? https://i.imgur.com/YpcnlCH.png https://i.imgur.com/YpcnlCH.png (Edit: interesting, thanks. So the underlying OS APIs that supply the power-consumption figures reported by asitop are just outright broken. The discrepancy is far too large to chalk up to static power losses or die-specific calibration factors that the video talks about.)
- wmf 11mo agohttps://www.youtube.com/watch?v=zCkbVLqUedg https://www.youtube.com/watch?v=zCkbVLqUedg
- m-s-y 11mo agoCan confirm. My M3 Ultra tops out at 210W when ComfyUI or ollama is running flat out. Confirmed via smart plug.
- anemll 11mo agoTensor Parallel test with RDMA last week https://x.com/anemll/status/1996349871260107102 https://x.com/anemll/status/1996349871260107102 Note fast sync workaround