electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: replace has_weight_files with state-based has_local_model + fix pipeline peak memory

Browse Source 
- 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 <[email protected]>
This commit is contained in:
Alex Cheema
2026-02-17 10:05:43 -08:00
co-authored by Claude Opus 4.6
parent 3409f19496
commit 48cc931ba4
6 changed files with 107 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/exo/shared/types/tasks.py
+2 -1
View File
@@ -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
src/exo/worker/engines/mlx/model_transfer.py
+6 -21
View File
@@ -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)))
src/exo/worker/engines/mlx/utils_mlx.py
+21 -13
View File
@@ -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)
src/exo/worker/plan.py
+18 -3
View File
@@ -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
src/exo/worker/runner/runner.py
+5 -4
View File
@@ -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)
src/exo/worker/tests/unittests/test_plan/test_download_and_loading.py
+55
View File
@@ -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