Distributed cluster serving across Macs (tensor + pipeline parallel) - #2423
Distributed cluster serving across Macs (tensor + pipeline parallel)#2423ashhart wants to merge 13 commits into
Conversation
|
Highly interested in this! |
|
Thank you for this impressive PR. Distributed cluster serving has been on my long-term roadmap, and I did not expect such a substantial implementation to arrive this quickly. Since this is a large change that will need an extended testing period, I plan to manage and validate it through a separate development release. I’ll revisit the PR after the 0.5.4 release and share my review then. Thanks again for the excellent contribution. |
|
No problem at all, if you need me to make any changes, let me know. What you've done with oMLX is amazing and I'd love to contribute more to the development 🙌🏻 |
|
I have a feeling that many people (myself included) are looking for something like this given the start/stop nature of exo's development and communication. |
|
Exo is phenomenal and I did take inspiration from some of their work. However, every time I tried it with Exo, I was thwart with overcoming issues in terminal and in the end I just gave up as even Claude/Codex could not get it to work. My goal was to make clustering as easy as clicking a button. |
ce7116d to
d8230c9
Compare
|
This looks excellent, thanks for doing it! I've been working on my own version of this for a while, with some different design decisions. Mine's nowhere near PR-ready, but I do have inference basically working across at least two machines for Qwen3.6-27B working using an mlx ring in raw TCP mode. Potentially interesting bits of my implementation:
My experimental branch has a lot more work left to do before it would actually be usable, and a fair bit of the code was authored by Qwen-3.6-27B, so there's likely a lot of slop in there I haven't noticed. I've pushed the squashed branch to my fork as a single commit, if anything in there looks useful feel free to pick ideas out of it and incorporate them. :) I'm going to pull your PR branch and give it a go. Thanks again for your work on this! :) |
|
Impressive! I have 2x Mac Studio M2 ultra 64gb and a Mac Studio M2 max 64gb. This I’m going to test! |
|
Would it help if I ran this on two 512 gb M3 max studios for testing or should I was for a release candidate? |
|
@smithcoin I'd love to see what you can run across both your Studios with this, post feedback! |
|
@ashhart Thank you for this, and sorry it took me a while to come back with a proper review. The scope, and the fact that you verified it on real hardware, both stand out, and distributed serving is a feature I have wanted in oMLX for a long time. I read through the whole branch. I agree with the direction and I want to land it, so here is how I would like to manage it rather than asking you to polish everything in one pass. I am splitting the feedback into three groups. Before I merge (A). These are the only hard blockers, and they are small:
Merge to dev, then follow-up PRs (B). The happy path works, and I would rather iterate on these in the open than block the merge on them:
Ongoing, with tests (C). The cluster tab strings are not localized yet, the idle polling re-runs peer probes and drops the in-progress plan, and there are staging cleanup and lint items. I am happy to take some of these on alongside you. For the dev builds I would also like to put distributed inference behind an explicit toggle in Settings > Advanced, so the cluster tab, the cluster routes, and the Bonjour advertising only appear once someone turns it on. That keeps the surface off by default while the feature matures, and it also stops a normal install from advertising itself over mDNS. The toggle is about exposure, not safety, so the A items above still need fixing regardless. I will send exact repro steps for the other A items so you are fixing rather than hunting. This is a strong first contribution and I would be glad to have you keep building on it. |
|
Repro and fix pointers for the A items, as promised. Line numbers are against the PR head. 1. RCE (SSH argument injection). The three GET handlers split the
Fix is to run each host through the same
2. Forgeable pairing and key-exchange tokens. 3. Keepalive on the wrong channel. The keepalive options are on
Not an A item and not blocking the merge, but the KV-zero reservation reproduces in three lines, so here it is: The cause is the |
|
Thank you for the detailed review and repros. I’ve confirmed all three A items and will address them in focused commits with regression tests. I’ll leave the B/C items and the Advanced toggle for follow-up work against the development branch so this update stays narrowly scoped |
09de67f to
f2749c0
Compare
|
Testing this on two Mac Studio M3 Ultra 256GB. I self-participated the "MLX Distributed" WWDC session earlier this summer, so these machines have already done this successfully in limited capacity. The most immediate issue I'm finding (from the server output) is that interactive ssh authentication messages are disrupting the setup process. Messages mostly about "known hosts from different IP addresses". Will continue tweaking the experiment here and report my results. Attempting to run Deepseek-V4-Flash-0731 distributed, which I'm already running successfully with oMLX on a single node. |
|
@tannerdsilva Thanks, this is very useful. Some setup probes can still inherit OpenSSH's interactive host-key policy when the same Mac is discovered through a different address. I will improve that flow around the paired node identity and verified host key. If possible, please retain the exact SSH stderr and note whether the target shown is That will help verify the fix against your topology. |
Registers mlx_lm.models.minimax_m3_vl over the vendored mlx-vlm implementation and adds the pipeline() contract: assigned-stage pinning, stage-local make_cache, and send/recv at stage boundaries. A rank now materialises only the layers it holds.
Plans check whether an architecture can actually be split before anything is staged, the catalogue reports max context per cluster, and node roles (headless/workstation) decide how much of a Mac a rank may take. Link bandwidth is measured, not assumed from the cable type.
A rank is refused before loading a stage its Mac cannot admit, watched while the weights arrive, and guarded per-prompt once serving. Budgets are derived from one definition shared by the planner and the guard.
Ranks heartbeat while idle, exclude themselves from peer watching, and a vanished peer becomes a stated failure instead of a hung collective. SSH control channels carry keepalives; teardown stops every rank.
Activation plans, preflights peers, launches, eager-loads every rank and declares ready only after a one-token canary. Cluster deployments register through the engine pool and roll back cleanly on failure.
Discovery, model picker backed by the catalogue, pipeline split slider with per-node KV and context, node roles with reasons, and activation that names its blocker instead of going quiet.
f2749c0 to
8d49fd5
Compare
|
I got this error and it wasn't clear the command to run:
I'll keep playing with it and try to post a comparison. Here was my benchmark running that model on exo:
Update - I am having trouble figuring out how to pair them in the GUI. Here are the diagnostic logs from each machine: |
|
@smithcoin Thanks for the report, the screenshot was very helpful. I've pushed an update that makes this setup clearer. The error means the peer Mac's SSH host key has not been trusted yet. The dashboard now provides an Open SSH setup in oMLX action, a copyable Terminal command, and key-generation instructions if the managed oMLX SSH key does not exist yet. After pulling the latest PR branch, please retry the pairing flow. If it still fails, sharing fresh cluster diagnostics from both Macs would be very useful. l've also improved the remote-code check so models requiring trust |
|
Quick status update following your review: All three A merge blockers are now resolved with regression tests:
I originally planned to leave the B items for follow-ups, but after reproducing them, each had a contained fix with test coverage, so these are now resolved too:
The Advanced distributed-inference toggle is also implemented and defaults to off. Subsequent tester reports prompted improvements to prompt-free SSH setup, first-time pairing guidance, and fail-fast handling for models requiring remote code. The relevant test suites are passing. Unless you would like further changes, I propose holding the branch stable now for hardware retesting and your re-review. I’m happy to take ownership of follow-up issues for orphaned ranks, engine-pool locking, cross-machine heartbeat timing, cluster localization, dashboard polling state, and staging cleanup after this lands in the development branch. Thanks. |
|
@ashhart thanks for your hard work on this! Its funny you mention topology - that seems to be the core of my issue for sure. I am not a python dev but I have extensive experience and am happy to contribute any amount "test effort" or tokens to help the cause alongside you! I will test anything you need, just tell me what you want! And apologies for the delayed response..... I had trouble with this stage during the WWDC session. I got it working at the time with some basic hacks....that aren't getting me by here. So I am currently stuck with oMLX in a way I didn't get stuck with Topology:I have two hosts:
During WWDC session, I never had any luck with the thunderbolt bridge before configuring it with static networking. With a static private network (10.200.215.x/24 and fdXX:XXXX:XXXX:XXXX/64) on both hosts, I can get them talking over Thunderbolt. At the time with OBSERVATION:When I am driving from machine OBSERVATION:As of I attempted to correct the IP addresses of both of my hosts, to use the STATIC IPv4 assigned to Thunderbolt I thought this would be the winning formula since this got me working on BETA2As of
TL:DR
|
To answer this directly: the target causing the error was an mDNS name and NOT an IP address. I was riding in the back of an Uber at the time and didn't have sufficient capacity to capture the output - apologies. |
|
Perhaps unrelated but I did want to bring this to attention (I could open another PR). I pulled HEAD from this PR and attempted to install using |
|
The first thing that I noticed in the GUI is that each machine thinks the other machine doesn't have enough memory: Another thing that is slightly 'frustrating' is that when I am trying to work on the cluster this element keeps toggling between this state: which causes a layout shift in the whole GUI making it hard to edit anything. In the GUI it states that RDMA isn't enabled for both machines, but it is: Here are my diagnostic logs from the two machines again: |
|
I'd be willing to download another model and give an apples to apples comparison to exo (if that is the issue), but that was the largest model I got working with that. If there is anything else I should do let me know. Happy to get this working! |
|
Thanks for the detailed testing and diagnostics. I’ve pushed another update addressing the reported memory, topology, RDMA, and dashboard issues:
The complete cluster test suite passes: 743 tests. JavaScript syntax, Python compilation, lint, diff, and privacy checks are also clean. The separate Please pull the latest PR branch and retry both automatic discovery and your static IPv4 Thunderbolt configuration. Fresh diagnostics would be very useful if either route still fails. |
|
Okay I installed Xcode and the Metal toolchain and pulled the latest version. I am still getting the layout shift from this component which hurts the usability of the whole page: It doesn't seem like I way to have the short lived key between the two devices. I am able to ssh from each machine to the other FYI. The machines are still not reading the correct size of the other (or detecting RDMA) I still see this front end error: Here are the logs: Let me know if there in another manner in which you'd like me to test. |
|
@smithcoin Thanks, I’ve pushed a fix for this. SSH errors now remain stable during polling, and the pairing UI explains the dedicated oMLX key with clear steps for both Macs. Please update and restart oMLX on both machines, hard-refresh the browser, then try pairing again. Your diagnostics show ~495 GiB and RDMA enabled on both Macs, so the earlier values appeared to be stale UI state. |














Distributed cluster serving across Macs
Serve one model across two or more Macs from the admin dashboard: pick the
cluster, pick the model, choose tensor or pipeline parallelism, activate.
Planning, staging, memory safety and rollback are handled by the server; the
dashboard exposes one
/v1endpoint when the cluster is ready.Verified on hardware
(jaccl): 28.6 tok/s vs 16.1 single-node (1.78x), byte-identical output.
usable context alone) split across a 128 GB MacBook Pro (workstation role,
18-layer stage pinned exactly as planned) and a 256 GB Mac Studio: one-click
activation from the dashboard reached ready on both ranks and passed the
one-token canary over jaccl. An earlier run served 212 tokens before losing
its remote rank to an idle ssh control-channel drop; keepalives were added
since, and a long sustained-serving run is the remaining verification.
What's in here
pipeline()/shard()) is refused at plan time, before staging — not aftera hundred gigabytes have moved. The catalogue answers "which of my models
fit this cluster, and at how much context".
workstationMac (someone works on it) admits less than aheadlessone, enforced by the memory guard at load, not just planning.watchdog sampling during the load, and a rank-0 prefill guard per request.
mlx-vlm implementation with an assigned-stage
pipeline()so a rankmaterialises only the layers the plan gave it.
disk checks and size verification.
vanished Mac becomes a stated failure instead of a hung collective.
Honest caveats (things a reviewer would find anyway)
memory_guard_tieris not yet carried per-rank (roleis); ranks defaultto
balanced.uneven plan whose even split would not fit, even for architectures that
honour the assigned-stage pin. Fails closed; costs some availability on
workstation splits.
remote rank supervision is future work.
tests/test_glm_mtp_patch.pyhas 2 tolerance failures that pre-exist thisbranch (file unchanged from origin/HEAD).
importorskip("mlx_vlm")on runners without it.Tests
7725 passed, 52 skipped locally (Apple Silicon). Commits are sequenced for
review: models → planning → guards → lifecycle → API → GUI → docs. Happy to
split into stacked PRs if that's easier to review.