More debug

2026-01-06 10:51:20 +00:00
parent 44e68e4498
commit eecc1da596
1 changed files with 5 additions and 1 deletions
@@ -196,7 +196,11 @@ class DistributedTransformer:
                mx.eval(hidden_states)
                mx_barrier(self.group)
                logger.info(f"sending single block outputs: {hidden_states.shape}")
-                mx.distributed.send(hidden_states, self.rank + 1, group=self.group)
+                hidden_states = mx.distributed.send(
+                    hidden_states, self.rank + 1, group=self.group
+                )
+                logger.info("sent single block outputs")
+                mx.eval(hidden_states)

        # === PHASE 5: Final Projection (last stage only) ===
        if self.is_last_stage: