Prune blocks before model load

src/exo/worker/engines/image/distributed_model.py

@@ -46,6 +46,14 @@ class DistributedImageModel:
         config = get_config_for_model(model_id)
         adapter = create_adapter_for_model(config, model_id, local_path, quantize)
 
+        if group is not None:
+            adapter.slice_transformer_blocks(
+                start_layer=shard_metadata.start_layer,
+                end_layer=shard_metadata.end_layer,
+                total_joint_blocks=config.joint_block_count,
+                total_single_blocks=config.single_block_count,
+            )
+
         # Create diffusion runner (handles both single-node and distributed modes)
         num_sync_steps = config.get_num_sync_steps("medium") if group else 0
         runner = DiffusionRunner(

src/exo/worker/engines/image/models/base.py

@@ -69,3 +69,24 @@ class BaseModelAdapter(ABC):
     def _get_latent_creator(self) -> type:
         """Return the latent creator class for this model."""
         ...
+
+    @abstractmethod
+    def slice_transformer_blocks(
+        self,
+        start_layer: int,
+        end_layer: int,
+        total_joint_blocks: int,
+        total_single_blocks: int,
+    ):
+        """Remove transformer blocks outside the assigned range.
+
+        This should be called BEFORE mx.eval() to avoid loading unused weights
+        in distributed mode.
+
+        Args:
+            start_layer: First layer index (inclusive) assigned to this node
+            end_layer: Last layer index (exclusive) assigned to this node
+            total_joint_blocks: Total number of joint blocks in the model
+            total_single_blocks: Total number of single blocks in the model
+        """
+        ...

src/exo/worker/engines/image/models/flux/adapter.py

@@ -250,6 +250,36 @@ class FluxModelAdapter(BaseModelAdapter):
             list(self._transformer.single_transformer_blocks),
         )
 
+    def slice_transformer_blocks(
+        self,
+        start_layer: int,
+        end_layer: int,
+        total_joint_blocks: int,
+        total_single_blocks: int,
+    ) -> None:
+        if end_layer <= total_joint_blocks:
+            # All assigned are joint blocks
+            joint_start, joint_end = start_layer, end_layer
+            single_start, single_end = 0, 0
+        elif start_layer >= total_joint_blocks:
+            # All assigned are single blocks
+            joint_start, joint_end = 0, 0
+            single_start = start_layer - total_joint_blocks
+            single_end = end_layer - total_joint_blocks
+        else:
+            # Spans both joint and single
+            joint_start, joint_end = start_layer, total_joint_blocks
+            single_start = 0
+            single_end = end_layer - total_joint_blocks
+
+        all_joint = list(self._transformer.transformer_blocks)
+        self._transformer.transformer_blocks = all_joint[joint_start:joint_end]
+
+        all_single = list(self._transformer.single_transformer_blocks)
+        self._transformer.single_transformer_blocks = all_single[
+            single_start:single_end
+        ]
+
     def merge_streams(
         self,
         hidden_states: mx.array,

src/exo/worker/engines/image/models/qwen/adapter.py

@@ -319,6 +319,17 @@ class QwenModelAdapter(BaseModelAdapter):
         """Qwen has no single blocks."""
         return []
 
+    def slice_transformer_blocks(
+        self,
+        start_layer: int,
+        end_layer: int,
+        total_joint_blocks: int,
+        total_single_blocks: int,
+    ) -> None:
+        all_blocks = list(self._transformer.transformer_blocks)
+        assigned_blocks = all_blocks[start_layer:end_layer]
+        self._transformer.transformer_blocks = assigned_blocks
+
     def merge_streams(
         self,
         hidden_states: mx.array,

src/exo/worker/engines/image/pipeline/adapter.py

@@ -262,6 +262,26 @@ class ModelAdapter(Protocol):
         """Get the list of single transformer blocks from the model."""
         ...
 
+    def slice_transformer_blocks(
+        self,
+        start_layer: int,
+        end_layer: int,
+        total_joint_blocks: int,
+        total_single_blocks: int,
+    ):
+        """Remove transformer blocks outside the assigned range.
+
+        This should be called BEFORE mx.eval() to avoid loading unused weights
+        in distributed mode.
+
+        Args:
+            start_layer: First layer index (inclusive) assigned to this node
+            end_layer: Last layer index (exclusive) assigned to this node
+            total_joint_blocks: Total number of joint blocks in the model
+            total_single_blocks: Total number of single blocks in the model
+        """
+        ...
+
     def merge_streams(
         self,
         hidden_states: mx.array,

src/exo/worker/engines/image/pipeline/runner.py

@@ -145,21 +145,17 @@ class DiffusionRunner:
             self.has_single_blocks or self.end_layer == self.total_joint
         )
 
-        # Slice blocks to only those assigned to this stage
-        all_joint_blocks = self.adapter.get_joint_blocks()
-        all_single_blocks = self.adapter.get_single_blocks()
-
-        assigned_joint_blocks = all_joint_blocks[self.joint_start : self.joint_end]
-        assigned_single_blocks = all_single_blocks[self.single_start : self.single_end]
+        joint_blocks = self.adapter.get_joint_blocks()
+        single_blocks = self.adapter.get_single_blocks()
 
         # Wrap blocks at initialization (reused across all calls)
         self.joint_block_wrappers = [
             JointBlockWrapper(block=block, adapter=self.adapter)
-            for block in assigned_joint_blocks
+            for block in joint_blocks
         ]
         self.single_block_wrappers = [
             SingleBlockWrapper(block=block, adapter=self.adapter)
-            for block in assigned_single_blocks
+            for block in single_blocks
         ]
 
     @property