← back to Exo
pipeline parallel fix
9058b117c0ab39efce6f099636d05b60476c106f · 2025-11-07 18:19:19 -0800 · Alex Cheema
Files touched
M .github/benchmark-dashboard/index.htmlM .github/configs/bench_simple.yamlM .github/scripts/bench.pyM .github/workflows/bench.ymlM .mlx_typings/mlx/utils.pyiM TODO.mdM src/exo/engines/mlx/auto_parallel.pyM src/exo/engines/mlx/utils_mlx.pyM src/exo/worker/runner/generate.pyM src/exo/worker/runner/runner.py
Diff
commit 9058b117c0ab39efce6f099636d05b60476c106f
Author: Alex Cheema <41707476+AlexCheema@users.noreply.github.com>
Date: Fri Nov 7 18:19:19 2025 -0800
pipeline parallel fix
---
.github/benchmark-dashboard/index.html | 23 ++++++---
.github/configs/bench_simple.yaml | 89 +++++++++++++++++++++++++++++---
.github/scripts/bench.py | 20 +++++---
.github/workflows/bench.yml | 58 ++++++++++++---------
.mlx_typings/mlx/utils.pyi | 9 +++-
TODO.md | 28 ++++++++++
src/exo/engines/mlx/auto_parallel.py | 93 ++++++++++++++++++++++------------
src/exo/engines/mlx/utils_mlx.py | 45 ++++++++--------
src/exo/worker/runner/generate.py | 8 +--
src/exo/worker/runner/runner.py | 28 +++++-----
10 files changed, 286 insertions(+), 115 deletions(-)
diff --git a/.github/benchmark-dashboard/index.html b/.github/benchmark-dashboard/index.html
index 341604bf..5b64af48 100644
--- a/.github/benchmark-dashboard/index.html
+++ b/.github/benchmark-dashboard/index.html
@@ -595,8 +595,9 @@
if (!resultStage) return;
- // Format: model [prompt_len/generation_len] iterations every time_between_requests secs
- const name = `${modelName} [${stageConfig.prompt_length}/${stageConfig.generation_length}] ${stageConfig.iterations} iterations every ${stageConfig.time_between_requests}s`;
+ // Format: stage_name: model [prompt_len/generation_len] iterations every time_between_requests secs
+ const stageName = stageConfig.name || `Stage ${stageIdx + 1}`;
+ const name = `${stageName}: ${modelName} [${stageConfig.prompt_length}/${stageConfig.generation_length}] ${stageConfig.iterations}× @ ${stageConfig.time_between_requests}s`;
// Success Rate
let successRate = 'N/A';
@@ -619,15 +620,25 @@
let prefillTime = 'N/A';
if (resultStage.avg_time_to_first_token !== null && resultStage.avg_time_to_first_token !== undefined) {
const ttftMs = resultStage.avg_time_to_first_token * 1000;
- prefillTime = `${ttftMs.toFixed(0)} ms`;
+ if (resultStage.std_time_to_first_token !== null && resultStage.std_time_to_first_token !== undefined) {
+ const stdMs = resultStage.std_time_to_first_token * 1000;
+ prefillTime = `${ttftMs.toFixed(0)} ± ${stdMs.toFixed(0)} ms`;
+ } else {
+ prefillTime = `${ttftMs.toFixed(0)} ms`;
+ }
}
// ms per token (1000 / decode_tps)
let msPerToken = 'N/A';
let msPerTokenClass = '';
- if (resultStage.avg_decode_tps !== null && resultStage.avg_decode_tps !== undefined && resultStage.avg_decode_tps > 0) {
- const ms = 1000 / resultStage.avg_decode_tps;
- msPerToken = `${ms.toFixed(1)} ms`;
+ if (resultStage.avg_ms_per_token !== null && resultStage.avg_ms_per_token !== undefined) {
+ const ms = resultStage.avg_ms_per_token;
+ if (resultStage.std_ms_per_token !== null && resultStage.std_ms_per_token !== undefined) {
+ const stdMs = resultStage.std_ms_per_token;
+ msPerToken = `${ms.toFixed(1)} ± ${stdMs.toFixed(1)} ms`;
+ } else {
+ msPerToken = `${ms.toFixed(1)} ms`;
+ }
// Color code based on performance
if (ms < 50) {
diff --git a/.github/configs/bench_simple.yaml b/.github/configs/bench_simple.yaml
index 26837edd..7ba45e44 100644
--- a/.github/configs/bench_simple.yaml
+++ b/.github/configs/bench_simple.yaml
@@ -4,33 +4,106 @@
# Hardware configuration - maps runner labels to instance counts
hardware_plan:
puffin4: 1
- puffin8: 1
+ # puffin8: 1
# Environment variables to set on each node
environment:
PLACEHOLDER: "placeholder"
- # OVERRIDE_MEMORY_MB: 30000
+ # OVERRIDE_MEMORY_MB: 50000
# MLX_METAL_FAST_SYNCH: 1
# Timeout for instance and runner readiness (seconds)
-timeout_seconds: 900
+timeout_seconds: 1800
# Model instances to run concurrently
model_ids:
- - "mlx-community/DeepSeek-V3.1-8bit"
+ # - "mlx-community/DeepSeek-V3.1-8bit"
# - "mlx-community/Qwen3-235B-A22B-4bit"
# - "mlx-community/Llama-3.3-70B-Instruct-4bit"
+ - "mlx-community/Llama-3.3-70B-Instruct-8bit"
# Placement strategy: "tensor", "pipeline", or "auto"
-strategy: "tensor_rdma"
+strategy: "pipeline"
# If true, run requests sequentially (no overlap); if false, fire-and-forget (default: false)
no_overlap: true
# Benchmark stages
+# pp: 64, 256, 1024, 2048, 4096, 8192, 16384
+# g: 64, 512
stages:
- - name: "simple"
- prompt_length: 512
- generation_length: 10
+ # - name: "simple"
+ # prompt_length: 512
+ # generation_length: 10
+ # time_between_requests: 2.0
+ # iterations: 5
+ - name: "pp64_g64"
+ prompt_length: 64
+ generation_length: 64
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp64_g512"
+ prompt_length: 64
+ generation_length: 512
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp256_g64"
+ prompt_length: 256
+ generation_length: 64
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp256_g512"
+ prompt_length: 256
+ generation_length: 512
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp1024_g64"
+ prompt_length: 1024
+ generation_length: 64
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp1024_g512"
+ prompt_length: 1024
+ generation_length: 512
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp2048_g64"
+ prompt_length: 2048
+ generation_length: 64
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp2048_g512"
+ prompt_length: 2048
+ generation_length: 512
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp4096_g64"
+ prompt_length: 4096
+ generation_length: 64
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp4096_g512"
+ prompt_length: 4096
+ generation_length: 512
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp8192_g64"
+ prompt_length: 8192
+ generation_length: 64
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp8192_g512"
+ prompt_length: 8192
+ generation_length: 512
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp16384_g64"
+ prompt_length: 16384
+ generation_length: 64
+ time_between_requests: 2.0
+ iterations: 10
+ - name: "pp16384_g512"
+ prompt_length: 16384
+ generation_length: 512
time_between_requests: 2.0
iterations: 10
diff --git a/.github/scripts/bench.py b/.github/scripts/bench.py
index 44733da1..6a841fbb 100644
--- a/.github/scripts/bench.py
+++ b/.github/scripts/bench.py
@@ -1,5 +1,6 @@
#!/usr/bin/env python3
+# type: ignore
"""
Unified benchmark script for EXO.
Runs single or multi-stage benchmarks with configurable load patterns.
@@ -54,7 +55,7 @@ def _http_request(url: str, *, method: str = "GET", data: Mapping[str, Any] | No
payload = json.dumps(data).encode("utf-8")
req = urllib.request.Request(url, data=payload, headers=headers, method=method)
try:
- with urllib.request.urlopen(req, timeout=30) as resp: # nosec - runner-local API
+ with urllib.request.urlopen(req, timeout=300) as resp: # nosec - runner-local API
body = resp.read().decode("utf-8")
try:
return json.loads(body)
@@ -72,7 +73,7 @@ async def _http_request_async(url: str, *, method: str = "GET", data: Mapping[st
return await loop.run_in_executor(None, lambda: _http_request(url, method=method, data=data))
-async def _http_stream_async(url: str, *, method: str = "POST", data: Mapping[str, Any], timeout: int = 120) -> list[tuple[str, float]]:
+async def _http_stream_async(url: str, *, method: str = "POST", data: Mapping[str, Any], timeout: int = 300) -> list[tuple[str, float]]:
"""Async streaming request. Returns list of (line, timestamp) tuples."""
def _stream() -> list[tuple[str, float]]:
headers = {"Content-Type": "application/json"}
@@ -400,7 +401,7 @@ async def wait_for_all_instances_deleted(api_base: str, model_id: str) -> None:
await asyncio.sleep(2)
-async def wait_for_tasks_drained(api_base: str, timeout_s: int = 300) -> None:
+async def wait_for_tasks_drained(api_base: str, timeout_s: int = 600) -> None:
"""Wait for all tasks in the cluster to be drained (completed or failed).
Tasks are deleted from state when complete, so we wait until there are no
@@ -550,7 +551,7 @@ async def run_single_request(
prompt: str,
max_tokens: int,
request_id: int,
- timeout: int = 60,
+ timeout: int = 300,
) -> RequestResult:
"""Run a single chat completion request and return its result."""
started_at = time.time()
@@ -770,7 +771,7 @@ async def run_stage(
# Wait for all tasks in the cluster to be drained
print(f"\nHTTP requests completed. Now waiting for cluster tasks to drain...")
- await wait_for_tasks_drained(api_base, timeout_s=300)
+ await wait_for_tasks_drained(api_base, timeout_s=600)
stage_completed_at = time.time()
@@ -786,8 +787,11 @@ async def run_stage(
# Calculate average TTFT and decode TPS for successful requests only
successful_results = [r for r in results if r.success]
+ # Skip first iteration if there are more than 1 iterations (warmup)
+ results_for_stats = successful_results[1:] if len(successful_results) > 1 else successful_results
+
# TTFT statistics
- ttft_values = [r.time_to_first_token_s for r in successful_results if r.time_to_first_token_s is not None]
+ ttft_values = [r.time_to_first_token_s for r in results_for_stats if r.time_to_first_token_s is not None]
avg_ttft = sum(ttft_values) / len(ttft_values) if ttft_values else None
if avg_ttft is not None and len(ttft_values) > 1:
@@ -797,7 +801,7 @@ async def run_stage(
std_ttft = None
# Decode TPS and ms per token statistics
- decode_tps_values = [r.decode_tps for r in successful_results if r.decode_tps is not None]
+ decode_tps_values = [r.decode_tps for r in results_for_stats if r.decode_tps is not None]
avg_decode_tps = sum(decode_tps_values) / len(decode_tps_values) if decode_tps_values else None
# Convert to ms per token
@@ -1162,7 +1166,7 @@ def main() -> int:
parser.add_argument("--config", type=Path, required=True, help="Path to YAML config file")
parser.add_argument("--expected-nodes", type=int, required=True, help="Total number of nodes expected in the cluster")
parser.add_argument("--is-primary", type=str, choices=["true", "false"], required=True)
- parser.add_argument("--timeout-seconds", type=int, default=600)
+ parser.add_argument("--timeout-seconds", type=int, default=1800)
parser.add_argument("--output", type=Path, help="Path to save detailed results JSON")
parser.add_argument("--git-commit", type=str, help="Git commit hash for metadata")
parser.add_argument("--hardware-labels", type=str, help="Comma-separated hardware labels")
diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml
index 746c0704..baa0d20d 100644
--- a/.github/workflows/bench.yml
+++ b/.github/workflows/bench.yml
@@ -1,8 +1,6 @@
name: bench
-on:
- pull_request:
- types: [review_requested]
+on: [push]
jobs:
plan:
@@ -139,29 +137,41 @@ jobs:
- name: Configure local MLX if available
run: |
- RUNNER_LABELS='${{ toJSON(runner.labels) }}'
- if echo "$RUNNER_LABELS" | grep -q "local_mlx"; then
- echo "Runner has 'local_mlx' tag, configuring local MLX paths..."
- MODIFIED=false
- if [ -d "/Users/Shared/mlx" ]; then
- echo "Found /Users/Shared/mlx, enabling local mlx path in pyproject.toml"
- sed -i.bak 's|^# mlx = { path = "/Users/Shared/mlx", editable=true }$|mlx = { path = "/Users/Shared/mlx", editable=true }|' pyproject.toml
- MODIFIED=true
- fi
- if [ -d "/Users/Shared/mlx-lm" ]; then
- echo "Found /Users/Shared/mlx-lm, enabling local mlx-lm path in pyproject.toml"
- sed -i.bak 's|^# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }$|mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }|' pyproject.toml
- MODIFIED=true
- fi
- if [ "$MODIFIED" = true ]; then
- echo "Modified pyproject.toml [tool.uv.sources] section:"
- sed -n '/\[tool\.uv\.sources\]/,/^\[/p' pyproject.toml | head -n -1
- echo "Regenerating uv.lock with local MLX paths..."
- nix --extra-experimental-features nix-command --extra-experimental-features flakes develop --command uv lock --upgrade-package mlx --upgrade-package mlx-lm
- fi
+ echo "=== DEBUG: Checking for local MLX configuration ==="
+ MODIFIED=false
+
+ echo "Checking for /Users/Shared/mlx directory..."
+ if [ -d "/Users/Shared/mlx" ]; then
+ echo "✓ Found /Users/Shared/mlx"
+ ls -la /Users/Shared/mlx | head -5
+ echo "Enabling local mlx path in pyproject.toml"
+ sed -i.bak 's|^# mlx = { path = "/Users/Shared/mlx", editable=true }$|mlx = { path = "/Users/Shared/mlx", editable=true }|' pyproject.toml
+ MODIFIED=true
+ else
+ echo "✗ /Users/Shared/mlx not found, will use PyPI version"
+ fi
+
+ echo "Checking for /Users/Shared/mlx-lm directory..."
+ if [ -d "/Users/Shared/mlx-lm" ]; then
+ echo "✓ Found /Users/Shared/mlx-lm"
+ ls -la /Users/Shared/mlx-lm | head -5
+ echo "Enabling local mlx-lm path in pyproject.toml"
+ sed -i.bak 's|^# mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }$|mlx-lm = { path = "/Users/Shared/mlx-lm", editable=true }|' pyproject.toml
+ MODIFIED=true
+ else
+ echo "✗ /Users/Shared/mlx-lm not found, will use PyPI version"
+ fi
+
+ if [ "$MODIFIED" = true ]; then
+ echo "=== Modified pyproject.toml [tool.uv.sources] section: ==="
+ sed -n '/\[tool\.uv\.sources\]/,/^\[/{/^\[tool\.uv\.sources\]/p; /^\[/!p;}' pyproject.toml
+ echo "=== Regenerating uv.lock with local MLX paths... ==="
+ nix --extra-experimental-features nix-command --extra-experimental-features flakes develop --command uv lock --upgrade-package mlx --upgrade-package mlx-lm
+ echo "✓ Lock file regenerated"
else
- echo "Runner does not have 'local_mlx' tag, using default PyPI packages"
+ echo "⚠ No local MLX directories found, using PyPI packages"
fi
+ echo "=== DEBUG: Local MLX configuration complete ==="
shell: bash
- name: Sync dependencies
diff --git a/.mlx_typings/mlx/utils.pyi b/.mlx_typings/mlx/utils.pyi
index d005e8cd..43738ca7 100644
--- a/.mlx_typings/mlx/utils.pyi
+++ b/.mlx_typings/mlx/utils.pyi
@@ -4,6 +4,8 @@ This type stub file was generated by pyright.
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from mlx.core import MX_ARRAY_TREE
+
def tree_map(
fn: Callable, tree: Any, *rest: Any, is_leaf: Optional[Callable] = ...
) -> Any:
@@ -139,7 +141,12 @@ def tree_unflatten(tree: Union[List[Tuple[str, Any]], Dict[str, Any]]) -> Any:
A Python tree.
"""
-def tree_reduce(fn, tree, initializer=..., is_leaf=...): # -> None:
+def tree_reduce(
+ fn: Callable[[Any, Any], Any],
+ tree: list[MX_ARRAY_TREE] | tuple[MX_ARRAY_TREE, ...] | dict[str, MX_ARRAY_TREE],
+ initializer=...,
+ is_leaf=...,
+) -> None:
"""Applies a reduction to the leaves of a Python tree.
This function reduces Python trees into an accumulated result by applying
diff --git a/TODO.md b/TODO.md
index 889ab52d..c07c2220 100644
--- a/TODO.md
+++ b/TODO.md
@@ -17,9 +17,37 @@
19. Fix mx.distributed.Group typing.
20. Add chat completion cancellations (e.g OpenWebUI has something for cancelling an ongoing request).
21. Make two separate things: tensor or pipeline, and ring or ibv.
+22. When downloading for the first time, stuff times out and I think the model never ends up actually loading into memory, or something.
+23. Do we need cache_limit? We went back and forth on that a lot because we thought it might be causing issues. One problem is it sets it relative to model size. So if you have multiple models loaded in it will take the most recent model size for the cache_limit. This is problematic if you launch DeepSeek -> Llama for example.
Potential refactors:
1. Make ForwarderEvent typed
2. Topology can be simplified
3. Get rid of InstanceReplacedAtomically
+
+Random errors we've run into:
+
+1. exo.shared.types.worker.common.RunnerError: RuntimeError: [ibv] Couldn't connect (error: 60). Traceback: Traceback (most recent call last):
+ File "/Users/puffin4/actions-runner/_work/exo/exo/src/exo/worker/runner/runner.py", line 54, in main
+ model, tokenizer, sampler, group = await loop.run_in_executor(
+ ^^^^^^^^^^^^^^^^^^^^^^^^^^^
+ ...<8 lines>...
+ )
+ ^
+ File "/nix/store/s7ik6dazn4nd2jdg9l36qf5q0z18sjyk-python3-3.13.8/lib/python3.13/concurrent/futures/thread.py", line 59, in run
+ result = self.fn(*self.args, **self.kwargs)
+ File "/Users/puffin4/actions-runner/_work/exo/exo/src/exo/engines/mlx/utils_mlx.py", line 149, in initialize_mlx
+ group = mlx_distributed_init(
+ model_shard_meta.device_rank,
+ ...<4 lines>...
+ or (mlx_ibv_devices is not None and len(mlx_ibv_devices) > 1),
+ )
+ File "/Users/puffin4/actions-runner/_work/exo/exo/src/exo/engines/mlx/utils_mlx.py", line 124, in mlx_distributed_init
+ group = mx.distributed.init(
+ backend="ring" if hosts is not None else "ibv",
+ strict=strict,
+ )
+RuntimeError: [ibv] Couldn't connect (error: 60)
+
+2.
\ No newline at end of file
diff --git a/src/exo/engines/mlx/auto_parallel.py b/src/exo/engines/mlx/auto_parallel.py
index 78109325..d1026779 100644
--- a/src/exo/engines/mlx/auto_parallel.py
+++ b/src/exo/engines/mlx/auto_parallel.py
@@ -22,16 +22,6 @@ from mlx.nn.layers.distributed import (
)
-class IdentityLayer(nn.Module):
- def __init__(self) -> None:
- super().__init__()
- self.use_sliding = False
-
- @override
- def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
- return x
-
-
class _LayerCallable(Protocol):
"""Structural type that any compatible layer must satisfy.
@@ -64,30 +54,55 @@ class CustomMlxLayer(nn.Module):
class PipelineFirstLayer(CustomMlxLayer):
- def __init__(self, original_layer: _LayerCallable, r: int, s: int):
+ def __init__(
+ self,
+ original_layer: _LayerCallable,
+ r: int,
+ s: int,
+ group: mx.distributed.Group | None = None,
+ ):
super().__init__(original_layer)
self.r: int = r
self.s: int = s
+ self.group = group
@override
def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
if self.r != 0:
- x = mx.distributed.recv_like(x, (self.r - 1))
+ x = mx.distributed.recv_like(x, (self.r - 1), group=self.group)
return self.original_layer(x, *args, **kwargs)
class PipelineLastLayer(CustomMlxLayer):
- def __init__(self, original_layer: _LayerCallable, r: int, s: int):
+ def __init__(
+ self,
+ original_layer: _LayerCallable,
+ r: int,
+ s: int,
+ group: mx.distributed.Group | None = None,
+ ):
super().__init__(original_layer)
self.r: int = r
self.s: int = s
+ self.group = group
@override
- def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+ def __call__(
+ self, x: mx.array, *args: object, cache: object = None, **kwargs: object
+ ) -> mx.array:
output: mx.array = self.original_layer(x, *args, **kwargs)
if self.r != self.s - 1:
- output = mx.distributed.send(output, (self.r + 1) % self.s)
- output = mx.distributed.all_gather(output)[-output.shape[0] :]
+ output = mx.distributed.send(
+ output, (self.r + 1) % self.s, group=self.group
+ )
+ if (
+ cache is not None
+ and hasattr(cache, "keys")
+ and getattr(cache, "keys", None) is not None
+ ):
+ cache.keys = mx.depends(cache.keys, output) # type: ignore[reportUnknownMemberType]
+
+ output = mx.distributed.all_gather(output, group=self.group)[-output.shape[0] :]
return output
@@ -98,6 +113,9 @@ class ParallelisationShardStrategy(Protocol):
class PipelineParallelisationStrategy(ParallelisationShardStrategy):
+ def __init__(self, group: mx.distributed.Group):
+ self.group = group
+
def auto_parallel(
self, model: nn.Module, model_shard_meta: ShardMetadata
) -> nn.Module:
@@ -124,27 +142,21 @@ class PipelineParallelisationStrategy(ParallelisationShardStrategy):
else:
raise ValueError("Model must have either a 'layers' or 'h' attribute")
- layers[: model_shard_meta.start_layer] = [
- IdentityLayer() for _ in range(model_shard_meta.start_layer)
- ]
- layers[model_shard_meta.end_layer :] = [
- IdentityLayer() for _ in range(len(layers) - model_shard_meta.end_layer)
- ]
- layers[model_shard_meta.start_layer] = PipelineFirstLayer(
- layers[model_shard_meta.start_layer],
+ layers = layers[model_shard_meta.start_layer : model_shard_meta.end_layer]
+ layers[0] = PipelineFirstLayer(
+ layers[0],
model_shard_meta.device_rank,
model_shard_meta.world_size,
+ group=self.group,
)
- layers[model_shard_meta.end_layer - 1] = PipelineLastLayer(
- layers[model_shard_meta.end_layer - 1],
+ layers[-1] = PipelineLastLayer(
+ layers[-1],
model_shard_meta.device_rank,
model_shard_meta.world_size,
+ group=self.group,
)
- # At this point `layers` *must* be a concrete list.
- assert isinstance(layers, list), (
- "Expected a list of layers after auto-parallel initialisation"
- )
+ PipelineParallelisationStrategy._set_layers(model, layers)
return model
@@ -160,11 +172,28 @@ class PipelineParallelisationStrategy(ParallelisationShardStrategy):
raise ValueError("Model must either have a 'model' or 'transformer' attribute")
+ @staticmethod
+ def _set_layers(model: nn.Module, layers: list[_LayerCallable]) -> None:
+ inner_model_instance = PipelineParallelisationStrategy._inner_model(model)
+ if hasattr(inner_model_instance, "layers"):
+ inner_model_instance.layers = layers
+
+ # Update DeepSeek V3 specific parameters when layers are shrunk
+ if isinstance(model, DeepseekV3Model) and hasattr(
+ inner_model_instance, "num_layers"
+ ):
+ inner_model_instance.start_idx = 0
+ inner_model_instance.end_idx = len(layers)
+ inner_model_instance.num_layers = len(layers)
+ elif hasattr(inner_model_instance, "h"):
+ inner_model_instance.h = layers
+ else:
+ raise ValueError("Model must have either a 'layers' or 'h' attribute")
+
class TensorParallelisationStrategy(ParallelisationShardStrategy):
def __init__(self, group: mx.distributed.Group):
self.group = group
- self.N = self.group.size
def auto_parallel(
self, model: nn.Module, model_shard_meta: ShardMetadata
@@ -236,7 +265,7 @@ class TensorParallelShardingStrategy(ABC):
self.sharded_to_all_linear = sharded_to_all_linear
self.all_to_sharded_linear_in_place = all_to_sharded_linear_in_place
self.sharded_to_all_linear_in_place = sharded_to_all_linear_in_place
- self.group = group or mx.distributed.init()
+ self.group = group
self.N = group.size()
@abstractmethod
diff --git a/src/exo/engines/mlx/utils_mlx.py b/src/exo/engines/mlx/utils_mlx.py
index 8d7bde05..9dfc5df5 100644
--- a/src/exo/engines/mlx/utils_mlx.py
+++ b/src/exo/engines/mlx/utils_mlx.py
@@ -8,23 +8,25 @@ from typing import Any, Callable, cast
from mlx_lm.models.cache import KVCache
from mlx_lm.sample_utils import make_sampler
+from exo.worker.runner.utils import get_weights_size
+
try:
from mlx_lm.tokenizer_utils import load_tokenizer
except ImportError:
from mlx_lm.tokenizer_utils import load as load_tokenizer # type: ignore
from mlx_lm.utils import load_model
+from mlx.utils import tree_reduce
from pydantic import RootModel
import mlx.core as mx
import mlx.nn as nn
from exo.engines.mlx import Model, TokenizerWrapper
from exo.engines.mlx.auto_parallel import (
- IdentityLayer,
PipelineParallelisationStrategy,
TensorParallelisationStrategy,
)
-from exo.shared.types.common import Host
from exo.shared.types.memory import Memory
+from exo.shared.types.common import Host
from exo.shared.types.tasks import ChatCompletionTaskParams
from exo.shared.types.worker.communication import runner_print
from exo.shared.types.worker.shards import ShardMetadata
@@ -72,6 +74,7 @@ def mlx_distributed_init(
hosts: list[Host] | None = None,
mlx_ibv_devices: list[list[str | None]] | None = None,
mlx_ibv_coordinator: str | None = None,
+ strict: bool = True,
) -> mx.distributed.Group:
"""
Initialize the MLX distributed (runs in thread pool).
@@ -80,10 +83,11 @@ def mlx_distributed_init(
- hosts: traditional host-based connectivity using MLX_HOSTFILE
- mlx_ibv_devices: RDMA connectivity matrix using MLX_IBV_DEVICES
- mlx_ibv_coordinator: coordinator address (IP:PORT) for RDMA setup
+ - strict: if True, raise an error if the distributed backend is not available
"""
- runner_print(f"Starting initialization for rank {rank}")
+ runner_print(f"Starting initialization for rank {rank}. Strict: {strict}")
- if mlx_ibv_devices is not None:
+ if mlx_ibv_devices is not None and mlx_ibv_devices != []:
assert mlx_ibv_coordinator is not None, (
"To use ibv backend must set ibv coordinator"
)
@@ -101,8 +105,7 @@ def mlx_distributed_init(
os.environ["MLX_IBV_DEVICES"] = devices_file
os.environ["MLX_RANK"] = str(rank)
os.environ["MLX_IBV_COORDINATOR"] = mlx_ibv_coordinator
-
- elif hosts is not None:
+ elif hosts is not None and hosts != []:
# Traditional host-based connectivity
hostfile = f"./hosts_{rank}.json"
hosts_json = HostList.from_hosts(hosts).model_dump_json()
@@ -116,10 +119,11 @@ def mlx_distributed_init(
os.environ["MLX_RANK"] = str(rank)
os.environ["MLX_RING_VERBOSE"] = "1"
else:
- raise ValueError("Either hosts or mlx_ibv_devices must be provided")
+ runner_print("No distributed setup, using single device mode")
group = mx.distributed.init(
- backend="ring" if hosts is not None else "ibv", strict=True
+ backend="ring" if hosts is not None else "ibv",
+ strict=strict,
)
runner_print(f"Rank {rank} mlx distributed initialization complete")
@@ -147,8 +151,12 @@ def initialize_mlx(
hosts=hosts,
mlx_ibv_devices=mlx_ibv_devices,
mlx_ibv_coordinator=mlx_ibv_coordinator,
+ strict=(hosts is not None and len(hosts) > 1)
+ or (mlx_ibv_devices is not None and len(mlx_ibv_devices) > 1),
)
+ set_wired_limit_for_model(get_weights_size(model_shard_meta))
+
sampler: Callable[[mx.array], mx.array] = make_sampler(temp=0.7)
model, tokenizer = shard_and_load(model_shard_meta, group=group)
@@ -177,11 +185,11 @@ def shard_and_load(
match model_shard_meta.strategy:
case "auto":
- strategy = PipelineParallelisationStrategy()
+ strategy = PipelineParallelisationStrategy(group)
case "pipeline":
- strategy = PipelineParallelisationStrategy()
+ strategy = PipelineParallelisationStrategy(group)
case "pipeline_rdma":
- strategy = PipelineParallelisationStrategy()
+ strategy = PipelineParallelisationStrategy(group)
case "tensor":
strategy = TensorParallelisationStrategy(group)
case "tensor_rdma":
@@ -189,8 +197,6 @@ def shard_and_load(
model = strategy.auto_parallel(model, model_shard_meta)
- runner_print(f"Model after auto_parallel: {str(model)}")
-
mx.eval(model.parameters())
mx.eval(model)
@@ -271,11 +277,7 @@ async def make_kv_cache(
max_kv_size: int | None = None,
) -> list[KVCache]:
assert hasattr(model, "layers")
-
- return [
- NullKVCache() if isinstance(layer, IdentityLayer) else KVCache()
- for layer in model.layers
- ]
+ return [KVCache() for _ in model.layers]
def mlx_force_oom(size: int = 40000) -> None:
@@ -315,6 +317,9 @@ def set_wired_limit_for_model(model_size: Memory):
"MB. This can be slow. See the documentation for possible work-arounds: "
"https://github.com/ml-explore/mlx-lm/tree/main#large-models"
)
- runner_print(f"Setting wired limit to {max_rec_size}")
+ kv_bytes = int(0.02 * model_bytes)
+ target_cache = int(1.10 * (model_bytes + kv_bytes))
+ target_cache = min(target_cache, max_rec_size)
+ mx.set_cache_limit(target_cache)
mx.set_wired_limit(max_rec_size)
- runner_print(f"Wired limit set to {max_rec_size}")
+ runner_print(f"Wired limit set to {max_rec_size}. Cache limit set to {target_cache}.")
diff --git a/src/exo/worker/runner/generate.py b/src/exo/worker/runner/generate.py
index e6821fd0..4f91e9e8 100644
--- a/src/exo/worker/runner/generate.py
+++ b/src/exo/worker/runner/generate.py
@@ -58,7 +58,7 @@ def generate_step(
logits_processors: list[Callable[[mx.array, mx.array], mx.array]] | None = None,
max_kv_size: int | None = None,
prompt_cache: list[KVCache] | None = None,
- prefill_step_size: int = 2048,
+ prefill_step_size: int = 16384,
kv_bits: int | None = None,
kv_group_size: int = 64,
quantized_kv_start: int = 0,
@@ -203,8 +203,8 @@ def generate_step(
)
mx.clear_cache()
- if eval_time > 7.0:
- prefill_step_size = prefill_step_size // 2
+ # if eval_time > 7.0:
+ # prefill_step_size = prefill_step_size // 2
if group is not None:
prefill_step_size = broadcast_from_zero(prefill_step_size)
prefill_step_size = max(1, prefill_step_size)
@@ -351,7 +351,7 @@ async def warmup_inference(
await loop.run_in_executor(mlx_executor, _generate_warmup)
runner_print("Generated ALL warmup tokens")
- mx_barrier()
+ await loop.run_in_executor(mlx_executor, lambda: mx_barrier(group))
return tokens_generated
diff --git a/src/exo/worker/runner/runner.py b/src/exo/worker/runner/runner.py
index 78b782da..79b9b521 100644
--- a/src/exo/worker/runner/runner.py
+++ b/src/exo/worker/runner/runner.py
@@ -5,6 +5,7 @@ from functools import partial
from multiprocessing.connection import Connection
from exo.engines.mlx.utils_mlx import (
+ mx_barrier,
initialize_mlx,
mlx_force_oom,
)
@@ -25,7 +26,7 @@ from exo.shared.types.worker.communication import (
)
from exo.shared.types.worker.shards import ShardMetadata
from exo.utils import ensure_type
-from exo.worker.runner.generate import mlx_generate, warmup_inference
+from exo.worker.runner.generate import mlx_generate
async def main(raw_conn: Connection):
@@ -62,17 +63,20 @@ async def main(raw_conn: Connection):
),
)
- runner_print(
- f"Warming up inference for model_shard_meta: {model_shard_meta} hosts: {hosts}"
- )
- toks = await warmup_inference(
- mlx_executor=mlx_executor,
- model=model,
- tokenizer=tokenizer,
- sampler=sampler,
- group=group,
- )
- runner_print(f"Warmed up by generating {toks} tokens")
+ # runner_print(
+ # f"Warming up inference for model_shard_meta: {model_shard_meta} hosts: {hosts}"
+ # )
+ # toks = await warmup_inference(
+ # mlx_executor=mlx_executor,
+ # model=model,
+ # tokenizer=tokenizer,
+ # sampler=sampler,
+ # group=group,
+ # )
+ # runner_print(f"Warmed up by generating {toks} tokens")
+ runner_print("Synchronizing processes before generation")
+ await loop.run_in_executor(mlx_executor, lambda: mx_barrier(group))
+ runner_print("Synchronized processes before generation")
await conn.send(InitializedResponse(time_taken=time.time() - setup_start_time))
while True:
← 612f58c7 Revert dumb merge mistake
·
back to Exo
·
Worker refactor aa519b8c →