Add barrier before all_gather

src/exo/worker/engines/mflux/pipefusion/distributed_denoising.py

@@ -6,6 +6,7 @@ from mflux.config.runtime_config import RuntimeConfig
 from mflux.models.flux.model.flux_transformer.transformer import Transformer
 
 from exo.shared.types.worker.shards import PipelineShardMetadata
+from exo.worker.engines.mlx.utils_mlx import mx_barrier
 
 
 def calculate_patch_heights(latent_height: int, num_patches: int, patch_size: int):
@@ -199,6 +200,9 @@ class DistributedDenoising:
                     hidden_states, self.rank + 1, group=self.group
                 )
 
+        mx.eval(hidden_states)
+        mx_barrier(group=self.group)
+
         #
         # === PHASE 5: All-gather Final Output ===
         # All stages participate to receive the final output