From 48cc931ba4d1eb5c2aa7f7bd928d180ffddbcd10 Mon Sep 17 00:00:00 2001 From: Alex Cheema Date: Thu, 12 Feb 2026 14:26:23 -0800 Subject: [PATCH] feat: replace has_weight_files with state-based has_local_model + fix pipeline peak memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add has_local_model field to LoadModel and TransferModelToDisk tasks, computed from download status in plan.py instead of filesystem checks - Remove has_weight_files() function entirely - In pipeline broadcast path, only load weights for layers in this node's [start_layer, end_layer) range — discard out-of-range results to reduce peak memory from ~22GB to ~10GB for 2-node pipeline Co-Authored-By: Claude Opus 4.6 --- src/exo/shared/types/tasks.py | 3 +- src/exo/worker/engines/mlx/model_transfer.py | 27 ++------- src/exo/worker/engines/mlx/utils_mlx.py | 34 +++++++----- src/exo/worker/plan.py | 21 ++++++- src/exo/worker/runner/runner.py | 9 +-- .../test_plan/test_download_and_loading.py | 55 +++++++++++++++++++ 6 files changed, 107 insertions(+), 42 deletions(-) diff --git a/src/exo/shared/types/tasks.py b/src/exo/shared/types/tasks.py index 6d0be674..5c18e980 100644 --- a/src/exo/shared/types/tasks.py +++ b/src/exo/shared/types/tasks.py @@ -42,7 +42,7 @@ class DownloadModel(BaseTask): # emitted by Worker class LoadModel(BaseTask): # emitted by Worker - pass + has_local_model: bool = Field(default=True) class ConnectToGroup(BaseTask): # emitted by Worker @@ -86,6 +86,7 @@ class TransferModelToDisk(BaseTask): # emitted by Worker """Transfer all model files from source to receivers' disk via MLX distributed.""" shard_metadata: ShardMetadata + has_local_model: bool = Field(default=True) class Shutdown(BaseTask): # emitted by Worker diff --git a/src/exo/worker/engines/mlx/model_transfer.py b/src/exo/worker/engines/mlx/model_transfer.py index 0f76a326..4ada7cb8 100644 --- a/src/exo/worker/engines/mlx/model_transfer.py +++ b/src/exo/worker/engines/mlx/model_transfer.py @@ -34,6 +34,7 @@ Group = mx.distributed.Group CHUNK_SIZE: Final[int] = 100 * 1024 * 1024 # 100 MB _LAYER_RE: Final[re.Pattern[str]] = re.compile(r"(?:^|\.)(layers|h)\.(\d+)\.") + def _all_sum_cpu(x: mx.array, group: Group) -> mx.array: """all_sum on CPU stream to avoid GPU memory pressure.""" return mx.distributed.all_sum( @@ -46,17 +47,6 @@ def _is_metadata_file(filename: str) -> bool: return not filename.endswith(".safetensors") -def has_weight_files(model_path: Path) -> bool: - """Check if model directory has both config and weight files.""" - if not model_path.exists(): - return False - has_config = (model_path / "config.json").exists() - has_weights = any(model_path.glob("*.safetensors")) or any( - model_path.glob("**/*.safetensors") - ) - return has_config and has_weights - - def model_path_for_id(model_id: ModelId) -> Path: """Get model path without requiring directory to exist (unlike build_model_path).""" return EXO_MODELS_DIR / model_id.normalize() @@ -259,18 +249,14 @@ def _transfer_files_to_disk( shutil.rmtree(temp_dir, ignore_errors=True) -def transfer_metadata_files( - model_path: Path, group: Group, is_source: bool -) -> None: +def transfer_metadata_files(model_path: Path, group: Group, is_source: bool) -> None: """ Transfer metadata files (config.json, tokenizer files, etc.) to receivers' disk. All ranks must call this function (collective operation). Only the designated source (is_source=True) should send; all others receive. """ - _transfer_files_to_disk( - model_path, group, is_source=is_source, metadata_only=True - ) + _transfer_files_to_disk(model_path, group, is_source=is_source, metadata_only=True) def transfer_all_files(model_path: Path, group: Group, is_source: bool) -> None: @@ -280,9 +266,7 @@ def transfer_all_files(model_path: Path, group: Group, is_source: bool) -> None: All ranks must call this function (collective operation). Only the designated source (is_source=True) should send; all others receive. """ - _transfer_files_to_disk( - model_path, group, is_source=is_source, metadata_only=False - ) + _transfer_files_to_disk(model_path, group, is_source=is_source, metadata_only=False) def _parse_mx_dtype(dtype_str: str) -> mx.Dtype: @@ -408,7 +392,8 @@ def prepare_weight_broadcast( for wf in weight_files: try: loaded = cast( - dict[str, mx.array], mx.load(str(wf), lazy=True) # pyright: ignore[reportCallIssue] + dict[str, mx.array], + mx.load(str(wf), lazy=True), # pyright: ignore[reportCallIssue] ) except TypeError: loaded = cast(dict[str, mx.array], mx.load(str(wf))) diff --git a/src/exo/worker/engines/mlx/utils_mlx.py b/src/exo/worker/engines/mlx/utils_mlx.py index 899f8747..161d8161 100644 --- a/src/exo/worker/engines/mlx/utils_mlx.py +++ b/src/exo/worker/engines/mlx/utils_mlx.py @@ -63,7 +63,6 @@ from exo.worker.engines.mlx.auto_parallel import ( from exo.worker.engines.mlx.model_transfer import ( WeightBroadcastState, coordinate_transfer, - has_weight_files, model_path_for_id, prepare_weight_broadcast, transfer_metadata_files, @@ -180,6 +179,7 @@ def load_mlx_items( bound_instance: BoundInstance, group: Group | None, on_timeout: TimeoutCallback | None = None, + has_local_model: bool = True, ) -> tuple[Model, TokenizerWrapper]: if group is None: logger.info(f"Single device used for {bound_instance.instance}") @@ -194,7 +194,10 @@ def load_mlx_items( logger.info("Starting distributed init") start_time = time.perf_counter() model, tokenizer = shard_and_load( - bound_instance.bound_shard, group=group, on_timeout=on_timeout + bound_instance.bound_shard, + group=group, + on_timeout=on_timeout, + has_local_model=has_local_model, ) end_time = time.perf_counter() logger.info( @@ -210,20 +213,20 @@ def shard_and_load( shard_metadata: ShardMetadata, group: Group, on_timeout: TimeoutCallback | None = None, + has_local_model: bool = True, ) -> tuple[nn.Module, TokenizerWrapper]: model_id = shard_metadata.model_card.model_id model_path = model_path_for_id(model_id) - has_local_weights = has_weight_files(model_path) # Coordinate: does any rank need a transfer? - needs_transfer, source_rank = coordinate_transfer(group, has_local_weights) + needs_transfer, source_rank = coordinate_transfer(group, has_local_model) broadcast_state: WeightBroadcastState | None = None if needs_transfer: is_source = group.rank() == source_rank logger.info( f"Model transfer needed (source_rank={source_rank}, " - f"is_source={is_source}, local_weights={has_local_weights})" + f"is_source={is_source}, local_weights={has_local_model})" ) # Step 1: Transfer metadata files (config.json, tokenizer, etc.) to disk transfer_metadata_files(model_path, group, is_source) @@ -236,7 +239,7 @@ def shard_and_load( # leaving the model un-quantized. lazy=False would then mx.eval() the full # fp16 model (~72GB for a 36B-param model), causing OOM on the receiver. # We handle quantization ourselves below before loading broadcast weights. - use_lazy = has_local_weights or broadcast_state is not None + use_lazy = has_local_model or broadcast_state is not None model, _ = load_model(model_path, lazy=use_lazy, strict=False) logger.debug(model) if hasattr(model, "model") and isinstance(model.model, DeepseekV3Model): # type: ignore @@ -248,7 +251,7 @@ def shard_and_load( if broadcast_state is not None: # When receiver has no weight files, load_model skips quantization. # Apply it explicitly so QuantizedLinear layers match broadcast weight shapes. - if not has_local_weights: + if not has_local_model: config_path = model_path / "config.json" with open(config_path) as f: config = json.load(f) # pyright: ignore[reportAny] @@ -306,15 +309,20 @@ def shard_and_load( ) case PipelineShardMetadata(): logger.info(f"loading model from {model_path} with pipeline parallelism") - # Pipeline doesn't shard per-layer, so broadcast all layer weights - # sequentially before pipeline_auto_parallel runs. + # Broadcast all layers (all_sum is collective — all ranks must + # participate) but only load weights for layers this node will + # keep after pipeline slicing. Out-of-range results are discarded, + # keeping peak memory proportional to this node's layer count. if broadcast_state is not None: for layer_idx in sorted(broadcast_state.layer_names.keys()): layer_weights = broadcast_state.broadcast_layer(layer_idx) - if layer_weights: - model.load_weights( - list(layer_weights.items()), strict=False - ) + if ( + shard_metadata.start_layer + <= layer_idx + < shard_metadata.end_layer + and layer_weights + ): + model.load_weights(list(layer_weights.items()), strict=False) del layer_weights model = pipeline_auto_parallel(model, group, shard_metadata) eval_with_timeout(model.parameters(), timeout_seconds, on_timeout) diff --git a/src/exo/worker/plan.py b/src/exo/worker/plan.py index be1e724e..037ca9b1 100644 --- a/src/exo/worker/plan.py +++ b/src/exo/worker/plan.py @@ -64,7 +64,7 @@ def plan( or _create_runner(node_id, runners, instances, all_runners) or _model_needs_download(node_id, runners, global_download_status) or _init_distributed_backend(runners, all_runners) - or _transfer_model_to_disk(runners, all_runners) + or _transfer_model_to_disk(runners, all_runners, global_download_status) or _load_model(runners, all_runners, global_download_status) or _ready_to_warmup(runners, all_runners) or _pending_tasks(runners, tasks, all_runners, input_chunk_buffer or {}) @@ -227,6 +227,7 @@ def _init_distributed_backend( def _transfer_model_to_disk( runners: Mapping[RunnerId, RunnerSupervisor], all_runners: Mapping[RunnerId, RunnerStatus], + global_download_status: Mapping[NodeId, Sequence[DownloadProgress]], ) -> TransferModelToDisk | None: """For transfer-only instances: after all ranks are connected, emit TransferModelToDisk.""" for runner in runners.values(): @@ -246,9 +247,15 @@ def _transfer_model_to_disk( ) if is_runner_connected and all_connected_or_further: + has_local = _node_has_download( + runner.bound_instance.bound_node_id, + shard_assignments.model_id, + global_download_status, + ) return TransferModelToDisk( instance_id=instance.instance_id, shard_metadata=runner.bound_instance.bound_shard, + has_local_model=has_local, ) return None @@ -274,7 +281,7 @@ def _load_model( if not _all_downloads_complete(shard_assignments, global_download_status): continue if isinstance(runner.status, RunnerIdle): - return LoadModel(instance_id=instance.instance_id) + return LoadModel(instance_id=instance.instance_id, has_local_model=True) else: # Multi-node: require at least one node to have the model downloaded. # Nodes without the model will receive it via MLX distributed transfer @@ -292,7 +299,15 @@ def _load_model( ) if is_runner_waiting and all_ready_for_model: - return LoadModel(instance_id=instance.instance_id) + has_local = _node_has_download( + runner.bound_instance.bound_node_id, + shard_assignments.model_id, + global_download_status, + ) + return LoadModel( + instance_id=instance.instance_id, + has_local_model=has_local, + ) return None diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py index a7168e2c..5d2d4ff0 100644 --- a/src/exo/worker/runner/runner.py +++ b/src/exo/worker/runner/runner.py @@ -84,7 +84,6 @@ from exo.worker.engines.mlx.cache import KVPrefixCache from exo.worker.engines.mlx.generator.generate import mlx_generate, warmup_inference from exo.worker.engines.mlx.model_transfer import ( coordinate_transfer, - has_weight_files, model_path_for_id, transfer_all_files, ) @@ -208,7 +207,10 @@ def main( if ModelTask.TextGeneration in shard_metadata.model_card.tasks: inference_model, tokenizer = load_mlx_items( - bound_instance, group, on_timeout=on_model_load_timeout + bound_instance, + group, + on_timeout=on_model_load_timeout, + has_local_model=task.has_local_model, ) logger.info( f"model has_tool_calling={tokenizer.has_tool_calling} using tokens {tokenizer.tool_call_start}, {tokenizer.tool_call_end}" @@ -551,8 +553,7 @@ def main( model_path = model_path_for_id( task.shard_metadata.model_card.model_id ) - has_local = has_weight_files(model_path) - _, source_rank = coordinate_transfer(group, has_local) + _, source_rank = coordinate_transfer(group, task.has_local_model) is_source = group.rank() == source_rank transfer_all_files(model_path, group, is_source) diff --git a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py index 60245c64..a1e006b7 100644 --- a/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py +++ b/src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py @@ -112,6 +112,7 @@ def test_plan_loads_model_when_all_shards_downloaded_and_waiting(): assert isinstance(result, LoadModel) assert result.instance_id == INSTANCE_1_ID + assert result.has_local_model is True def test_plan_does_not_request_download_when_shard_already_downloaded(): @@ -207,6 +208,7 @@ def test_plan_loads_model_when_any_node_has_download_for_multi_node(): assert isinstance(result, LoadModel) assert result.instance_id == INSTANCE_1_ID + assert result.has_local_model is True def test_plan_does_not_load_model_when_no_node_has_download(): @@ -252,3 +254,56 @@ def test_plan_does_not_load_model_when_no_node_has_download(): ) assert result is None + + +def test_plan_load_model_has_local_model_false_when_node_missing_download(): + """ + For multi-node instances, when the local node does NOT have the model + but a peer does, LoadModel should be emitted with has_local_model=False. + """ + shard1 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=0, world_size=2) + shard2 = get_pipeline_shard_metadata(MODEL_A_ID, device_rank=1, world_size=2) + instance = get_mlx_ring_instance( + instance_id=INSTANCE_1_ID, + model_id=MODEL_A_ID, + node_to_runner={NODE_A: RUNNER_1_ID, NODE_B: RUNNER_2_ID}, + runner_to_shard={RUNNER_1_ID: shard1, RUNNER_2_ID: shard2}, + ) + + # NODE_B is the local node (bound_node_id=NODE_B), it does NOT have the model + bound_instance = BoundInstance( + instance=instance, bound_runner_id=RUNNER_2_ID, bound_node_id=NODE_B + ) + local_runner = FakeRunnerSupervisor( + bound_instance=bound_instance, status=RunnerConnected() + ) + + runners = {RUNNER_2_ID: local_runner} + instances = {INSTANCE_1_ID: instance} + all_runners = { + RUNNER_1_ID: RunnerConnected(), + RUNNER_2_ID: RunnerConnected(), + } + + # Only NODE_A has the model, NODE_B does not + global_download_status: dict[NodeId, list[DownloadProgress]] = { + NODE_A: [ + DownloadCompleted( + shard_metadata=shard1, node_id=NODE_A, total_bytes=Memory() + ) + ], + NODE_B: [], + } + + result = plan_mod.plan( + node_id=NODE_B, + runners=runners, # type: ignore + global_download_status=global_download_status, + instances=instances, + all_runners=all_runners, + tasks={}, + ) + + assert isinstance(result, LoadModel) + assert result.instance_id == INSTANCE_1_ID + assert result.has_local_model is False