fix(mlx sharding): bypass deadlocking pipeline prefill + chunk stream_generate + CPU all_gather (#2108) - #2269
Open
adamteale wants to merge 4 commits into
Open
Conversation
…xo-explore#2108) Two bugs deadlocked pipeline_parallel_prefill for 2-node sharded inference: 1. queue_sends=True + async flush never submitted the send. With queue_sends=True, the PipelineLastLayer appends the send to _pending_prefill_sends and flush_prefill_sends submits it via mx.async_eval(sent). But no eval ever forces that async send, so it is never actually submitted and the receiver's recv deadlocks forever. Fix: queue_sends=False submits the send immediately in the PipelineLastLayer (like the warmup path), so send/recv pair correctly. 2. The distributed_prompt_progress_callback's all_sum/all_gather didn't block for the delayed peer. The callback (agree_on_cancellations + agree_on_tasks) runs collective ops (mx_any all_sum, mx_all_gather_tasks all_gather) on the rank doing the leading dummy (no forward, ~0s) while the peer is still doing the real chunk forward (~8s). The leading rank's collective recv returns stale/garbage data instead of blocking, and the forward rank's all_sum deadlocks. For single-request prefill (no cancellations, no new tasks) the callback is a no-op anyway (all counts are 0). Fix: pass None to skip it. Verified: Qwen3.8-27B-OptiQ-4bit sharded across M4 Max (36GB) + M5 Pro (24GB) over Thunderbolt 5 — 8K (4825 tok) 24s, 16K (8125 tok) 25s, 32K (16225 tok) 67s, all correct generation. Previously deadlocked at the first pipeline callback for any prompt >= prefill_step_size.
…_generate + CPU all_gather (exo-explore#2108) Three changes that together make multi-machine sharded inference (27B across 2 Macs over Thunderbolt) work for all prompt sizes: 1. generate.py prefill(): set prefill_step_size=100000 so the pipeline_parallel_prefill path is never taken. That path has a deeper chunked-ring send/recv deadlock that queue_sends=False + skipping the distributed callback do NOT resolve. All prompts use stream_generate. 2. generate.py prefill(): pass prefill_step_size=2048 to stream_generate (NOT the 100000 bypass threshold) so the prefill is chunked. Without this the entire prompt is processed in one forward pass, which OOMs the 24GB peer's Metal memory at ~6K+ tokens (kIOGPUCommandBufferCallbackErrorOutOfMemory). 3. auto_parallel.py PipelineLastLayer: run the decode all_gather on the CPU stream (stream=mx.default_stream(mx.Device(mx.cpu))). mlx's AllGather has NO GPU implementation, so on the GPU generation_stream it silently never executes and the eval hangs. The CPU stream's collectives DO block (matching mx_barrier in utils_mlx.py). Verified: 27B sharded across M4 Max (36GB) + M5 Pro (24GB): - warmup: both ranks ready (50 tokens) - small prompt (~10 tokens): 8s, correct - 5K prompt (6373 tokens): 34s, correct - 8K prompt (9382 tokens): 47s, correct, no OOM
adamteale
changed the title
…#2108) The mlx ring transport prints '[ring] Too many send/recv errors. Aborting...' to stderr but does NOT crash the process — the in-flight collective hangs forever, leaving the runner alive but unresponsive. exo's diagnostics were only acted upon on runner termination, so the hung runner was never recovered. This wires a fatal-diagnostic callback: when the ring abort is detected on the runner's stderr, the supervisor kills the runner process → RunnerFailed → exo re-places the instance with a fresh distributed group (re-initialized ring). Verified-compiled; needs a sustained-load run to confirm the auto-recovery.
…xo-explore#2108) The ErrorChunk's diagnostics field fails pydantic validation for RunnerRingTransportError (the message field is rejected as extra_forbidden on the receiving side — a pre-existing TaggedModel serialization bug). Clearing the diagnostics list before stopping the runner avoids the validation flood that would otherwise wedge the supervisor + take down the exo API.
adamteale
force-pushed
the
fix/pipeline-prefill-deadlock-2108
branch
from
August 19, 2026 01:57
58a18d0 to
375b0e7
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #2108 — multi-machine sharded inference deadlocks/hangs
Four changes that make multi-machine sharded inference (27B across 2 Macs over Thunderbolt 5) work for all prompt sizes + auto-recover from ring transport failures.
1. Bypass the deadlocking pipeline prefill (
generate.py)prefill_step_size = 100000sopipeline_parallel_prefillis never taken — that path has a deeper chunked-ring send/recv deadlock thatqueue_sends=False+ skipping the distributed callback do NOT resolve. All prompts usestream_generate.2. Chunk stream_generate to avoid OOM (
generate.py)Pass
prefill_step_size=2048tostream_generate(...)(NOT the 100000 threshold). Without this the entire prompt is one forward pass → OOMs the 24GB peer Metal memory at ~6K+ tokens (kIOGPUCommandBufferCallbackErrorOutOfMemory).3. Decode all_gather on the CPU stream (
auto_parallel.py)PipelineLastLayerdecode:mx.distributed.all_gather(..., stream=mx.default_stream(mx.Device(mx.cpu))). mlxAllGatherhas no GPU implementation — on the GPUgeneration_streamit silently never executes → eval hangs. CPU stream collectives DO block (matchingmx_barrierinutils_mlx.py).4. Crash-on-ring-abort auto-recovery (
supervisor.py)The mlx ring transport prints
[ring] Too many send/recv errors. Aborting...to stderr under sustained load (anerrno 14EFAULT on recv — a buffer-lifetime bug in the mlx C++ ring code) but does NOT crash the process — the in-flight collective hangs forever, leaving the runner alive but unresponsive. exo diagnostics were only acted upon on runner termination, so the hung runner was never recovered.This wires a fatal-diagnostic callback: when the ring abort is detected on stderr, the supervisor clears the collected diagnostics (avoids a pre-existing pydantic validation flood in
ErrorChunksending —RunnerRingTransportError.messageis rejected asextra_forbiddenby the receiver) + kills the runner process →RunnerFailed→ exo re-places the instance with a fresh distributed group (re-initialized ring) → the 27B recovers.Verification (27B: M4 Max 36GB + M5 Pro 24GB)
What does NOT work (and why)
pipeline_parallel_prefill(bypassed) — the chunked ring send/recv deadlocks even with thequeue_sends/callback fixes.errno 14EFAULT is an underlying C++ buffer-lifetime bug that cannot be fixed in Python — the crash-on-ring-abort fix makes the system resilient (auto-restart) rather than fixing the ring itself.TaggedModelserialization bug (RunnerRingTransportError.messagerejected asextra_forbidden) is worked around by clearing diagnostics; the real fix would be in the schema.