From 08e8a30fb7ef5b113298bbc9cedcbdb739a1b4d2 Mon Sep 17 00:00:00 2001 From: ciaranbor Date: Mon, 5 Jan 2026 17:31:24 +0000 Subject: [PATCH] Fix joint block patch recv shape for non-zero ranks --- src/exo/worker/engines/image/pipeline/runner.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py index 171dddab..f51d7c9f 100644 --- a/src/exo/worker/engines/image/pipeline/runner.py +++ b/src/exo/worker/engines/image/pipeline/runner.py @@ -734,8 +734,18 @@ class DiffusionRunner: if self.has_joint_blocks: if not self.is_first_stage or t != self.num_sync_steps: + if self.is_first_stage: + # First stage receives latent-space from last stage (scheduler output) + recv_template = patch + else: + # Other stages receive hidden-space from previous stage + patch_len = patch.shape[1] + recv_template = mx.zeros( + (batch_size, patch_len, hidden_dim), + dtype=patch.dtype, + ) patch = mx.distributed.recv_like( - patch, src=self.prev_rank, group=self.group + recv_template, src=self.prev_rank, group=self.group ) mx.eval(patch) patch_latents[patch_idx] = patch