diff --git a/src/exo/worker/engines/image/distributed_model.py b/src/exo/worker/engines/image/distributed_model.py index c34bceec..b6e097b4 100644 --- a/src/exo/worker/engines/image/distributed_model.py +++ b/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( diff --git a/src/exo/worker/engines/image/models/base.py b/src/exo/worker/engines/image/models/base.py index f5800d09..1a245a7f 100644 --- a/src/exo/worker/engines/image/models/base.py +++ b/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 + """ + ... diff --git a/src/exo/worker/engines/image/models/flux/adapter.py b/src/exo/worker/engines/image/models/flux/adapter.py index fe505429..20e686e2 100644 --- a/src/exo/worker/engines/image/models/flux/adapter.py +++ b/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, diff --git a/src/exo/worker/engines/image/models/qwen/adapter.py b/src/exo/worker/engines/image/models/qwen/adapter.py index 0b1699e7..14e59a8b 100644 --- a/src/exo/worker/engines/image/models/qwen/adapter.py +++ b/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, diff --git a/src/exo/worker/engines/image/pipeline/adapter.py b/src/exo/worker/engines/image/pipeline/adapter.py index 86c304ac..e28e90bb 100644 --- a/src/exo/worker/engines/image/pipeline/adapter.py +++ b/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, diff --git a/src/exo/worker/engines/image/pipeline/runner.py b/src/exo/worker/engines/image/pipeline/runner.py index 17767417..fed95d57 100644 --- a/src/exo/worker/engines/image/pipeline/runner.py +++ b/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