electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

try placing in pipeline last layer

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-01-21 11:35:25 +00:00
parent ded7b35a2e
commit f5886c1953
1 changed files with 12 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/worker/engines/mlx/auto_parallel.py
+12 -15
View File
@@ -152,6 +152,18 @@ class PipelineLastLayer(CustomMlxLayer):
if cache is not None:
cache.keys = mx.depends(cache.keys, sent) # type: ignore[reportUnknownMemberType]
if self.r == self.s - 1:
mx.eval(logits)
logger.info(f"[patch_pipeline_model] send logits to rank 0, shape={logits.shape}")
sent = mx.distributed.send(logits, dst=0, group=group)
mx.eval(sent)
logger.info(f"[patch_pipeline_model] send done")
elif rank == 0:
logger.info(f"[patch_pipeline_model] recv_like logits from rank {self.s - 1}, shape={logits.shape}")
logits = mx.distributed.recv_like(logits, src=self.s - 1, group=group)
mx.eval(logits)
logger.info(f"[patch_pipeline_model] recv_like done")
return output
@@ -262,21 +274,6 @@ 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
# 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:
mx.eval(logits)
logger.info(f"[patch_pipeline_model] send logits to rank 0, shape={logits.shape}")
sent = mx.distributed.send(logits, dst=0, group=group)
mx.eval(sent)
logger.info(f"[patch_pipeline_model] send done")
elif rank == 0:
logger.info(f"[patch_pipeline_model] recv_like logits from rank {world_size - 1}, shape={logits.shape}")
logits = mx.distributed.recv_like(logits, src=world_size - 1, group=group)
mx.eval(logits)
logger.info(f"[patch_pipeline_model] recv_like done")
return logits
cls.__call__ = patched_call