Add debug logs

src/exo/worker/engines/mflux/distributed_flux.py

@@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Any, Literal, Optional
 import mlx.core as mx
 from mflux.callbacks.callbacks import Callbacks
 from mflux.config.config import Config
+from exo.worker.runner.bootstrap import logger
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.models.common.latent_creator.latent_creator import Img2Img, LatentCreator
 from mflux.models.flux.latent_creator.flux_latent_creator import FluxLatentCreator
@@ -260,6 +261,7 @@ class DistributedFlux1:
         # Scale model input if needed by the scheduler
         latents = config.scheduler.scale_model_input(latents, t)
 
+        logger.info("running transformer")
         # Predict noise (communication happens in block wrappers)
         noise = model.transformer(
             t=t,

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

@@ -1,5 +1,6 @@
 from typing import Any
 
+from exo.worker.runner.bootstrap import logger
 import mlx.core as mx
 from mflux.config.runtime_config import RuntimeConfig
 from mflux.models.flux.model.flux_transformer.transformer import Transformer
@@ -99,6 +100,7 @@ class DistributedTransformer:
         controlnet_single_block_samples: list[mx.array] | None = None,
         kontext_image_ids: mx.array | None = None,
     ) -> mx.array:
+        logger.info("running distributed transformer")
         """Forward pass with inline distributed communication."""
         transformer = self.transformer
 
@@ -116,6 +118,7 @@ class DistributedTransformer:
         if self.has_joint_blocks:
             # Receive from previous stage (if not first stage)
             if not self.is_first_stage:
+                logger.info("receiving joint block inputs")
                 hidden_states = mx.distributed.recv_like(
                     hidden_states, self.rank - 1, group=self.group
                 )
@@ -125,6 +128,7 @@ class DistributedTransformer:
 
             # Run assigned joint blocks
             for idx in range(self.joint_start, self.joint_end):
+                logger.info(f"running joint block {idx}")
                 block = transformer.transformer_blocks[idx]
                 encoder_hidden_states, hidden_states = block(
                     hidden_states=hidden_states,
@@ -135,6 +139,7 @@ class DistributedTransformer:
 
         # === PHASE 3: Joint→Single Transition ===
         if self.is_concat_stage:
+            logger.info("concatenating")
             # Concatenate encoder and hidden states
             concatenated = mx.concatenate(
                 [encoder_hidden_states, hidden_states], axis=1
@@ -148,6 +153,7 @@ class DistributedTransformer:
                 mx.distributed.send(concatenated, self.rank + 1, group=self.group)
                 # This stage is done with blocks, but will participate in all_gather
         elif self.has_joint_blocks and not self.is_last_stage:
+            logger.info("sending joint block outputs")
             # Send joint block outputs to next stage (which has more joint blocks)
             mx.distributed.send(hidden_states, self.rank + 1, group=self.group)
             mx.distributed.send(encoder_hidden_states, self.rank + 1, group=self.group)
@@ -156,12 +162,14 @@ class DistributedTransformer:
         if self.has_single_blocks:
             # Receive from previous stage if we didn't do concatenation
             if not self.is_concat_stage and not self.is_first_stage:
+                logger.info(f"receiving single block inputs: {hidden_states.shape}")
                 hidden_states = mx.distributed.recv_like(
                     hidden_states, self.rank - 1, group=self.group
                 )
 
             # Run assigned single blocks
             for idx in range(self.single_start, self.single_end):
+                logger.info(f"running single block: {idx}")
                 block = transformer.single_transformer_blocks[idx]
                 hidden_states = block(
                     hidden_states=hidden_states,
@@ -171,6 +179,7 @@ class DistributedTransformer:
 
             # Send to next stage if not last
             if not self.is_last_stage:
+                logger.info(f"sending single block outputs: {hidden_states.shape}")
                 mx.distributed.send(hidden_states, self.rank + 1, group=self.group)
 
         # === PHASE 5: Final Projection (last stage only) ===
@@ -182,6 +191,7 @@ class DistributedTransformer:
 
         # === PHASE 6: All-gather Final Output ===
         # All stages participate to receive the final output
+        logger.info("gathering final output")
         hidden_states = mx.distributed.all_gather(hidden_states, group=self.group)[
             -hidden_states.shape[0] :
         ]