From 0ad852870a954f8ebc11f04448cb4b0dd313f3a3 Mon Sep 17 00:00:00 2001 From: Ryuichi Leo Takashige Date: Wed, 21 Jan 2026 11:12:48 +0000 Subject: [PATCH] Use recv/send over all gather --- src/exo/worker/engines/mlx/auto_parallel.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index 159c2a45..d0b8b6b4 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -238,6 +238,9 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T: original_call = cls.__call__ # type :ignore call_signature = signature(original_call) # type :ignore + world_size = group.size() + rank = group.rank() + def patched_call( self: T, *args: object, @@ -252,9 +255,13 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T: if cache is not None: cache[-1].state = mx.depends(cache[-1].state, logits) # type: ignore - logits = mx.distributed.all_gather(logits, group=group)[ - -logits.shape[0] : - ] # type :ignore + # Use send/recv instead of all_gather to get logits to rank 0 + # (JACCL recv picks up all_gather data instead of send data) + if world_size > 1: + if rank == world_size - 1: + sent = mx.distributed.send(logits, dst=0, group=group) + elif rank == 0: + logits = mx.distributed.recv_like(logits, src=world_size - 1, group=group) return logits