← back to Exo
Wrap pipeline models for explicit mx.depends between cache and logits (#1206)
5fd55594c9e433b6a2444b0b43681e78dafcd6b5 · 2026-01-19 17:49:42 +0000 · rltakashige
## Motivation
GPU timeouts often when prompt size > profile_step_size. It also happens
for seemingly random models.
## Changes
Add mx.depends for cache on the logits.
All gather at the model level rather than the layer level, reducing the
amount of data sent.
## Why It Works
mlx_lm's prefill loop only evaluates cache state, not logits.
When prompt > prefill_step_size, the all_gather is never evaluated,
causing GPU timeout.
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->
### Automated Testing
Added failing test cases and then resolved them.
Files touched
M bench/exo_bench.pyM src/exo/worker/engines/mlx/auto_parallel.pyM src/exo/worker/tests/unittests/test_mlx/conftest.pyM src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.pyA src/exo/worker/tests/unittests/test_mlx/test_distributed_fix.py
Diff
commit 5fd55594c9e433b6a2444b0b43681e78dafcd6b5
Author: rltakashige <rl.takashige@gmail.com>
Date: Mon Jan 19 17:49:42 2026 +0000
Wrap pipeline models for explicit mx.depends between cache and logits (#1206)
## Motivation
GPU timeouts often when prompt size > profile_step_size. It also happens
for seemingly random models.
## Changes
Add mx.depends for cache on the logits.
All gather at the model level rather than the layer level, reducing the
amount of data sent.
## Why It Works
mlx_lm's prefill loop only evaluates cache state, not logits.
When prompt > prefill_step_size, the all_gather is never evaluated,
causing GPU timeout.
## Test Plan
### Manual Testing
<!-- Hardware: (e.g., MacBook Pro M1 Max 32GB, Mac Mini M2 16GB,
connected via Thunderbolt 4) -->
<!-- What you did: -->
<!-- - -->
### Automated Testing
Added failing test cases and then resolved them.
---
bench/exo_bench.py | 25 +--
src/exo/worker/engines/mlx/auto_parallel.py | 31 ++-
.../worker/tests/unittests/test_mlx/conftest.py | 141 +++++++------
.../tests/unittests/test_mlx/test_auto_parallel.py | 37 ++--
.../unittests/test_mlx/test_distributed_fix.py | 230 +++++++++++++++++++++
5 files changed, 364 insertions(+), 100 deletions(-)
diff --git a/bench/exo_bench.py b/bench/exo_bench.py
index 7366dbc0..478ebf9e 100644
--- a/bench/exo_bench.py
+++ b/bench/exo_bench.py
@@ -16,9 +16,6 @@ from urllib.parse import urlencode
from loguru import logger
from transformers import AutoTokenizer
-from exo.shared.models.model_cards import MODEL_CARDS
-from exo.shared.types.memory import Memory
-
class ExoHttpError(RuntimeError):
def __init__(self, status: int, reason: str, body_preview: str):
@@ -490,17 +487,17 @@ def main() -> int:
logger.debug(f" warmup {i + 1}/{args.warmup} done")
for pp in pp_list:
- if (
- pp * n_nodes > 2048
- and "ring" in instance_meta.lower()
- and "tensor" in sharding.lower()
- ):
- model_card = MODEL_CARDS[short_id]
- if model_card.metadata.storage_size > Memory.from_gb(10):
- logger.info(
- f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
- )
- continue
+ # if (
+ # pp * n_nodes > 2048
+ # and "ring" in instance_meta.lower()
+ # and "tensor" in sharding.lower()
+ # ):
+ # model_card = MODEL_CARDS[short_id]
+ # if model_card.metadata.storage_size > Memory.from_gb(10):
+ # logger.info(
+ # f"Skipping tensor ring as this is too slow for model of size {model_card.metadata.storage_size} on {n_nodes=}"
+ # )
+ # continue
for tg in tg_list:
runs: list[dict[str, Any]] = []
for r in range(args.repeat):
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index c2819f64..b81b75be 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -108,7 +108,6 @@ class PipelineLastLayer(CustomMlxLayer):
if cache 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
@@ -193,6 +192,36 @@ def pipeline_auto_parallel(
"Expected a list of layers after auto-parallel initialisation"
)
+ return patch_pipeline_model(model, group)
+
+
+def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
+ # Patch __call__ on the model's class
+ cls = model.__class__
+ original_call = cls.__call__ # type :ignore
+ call_signature = signature(original_call) # type :ignore
+
+ def patched_call(
+ self: T,
+ *args: object,
+ **kwargs: object,
+ ) -> mx.array:
+ logits: mx.array = original_call(self, *args, **kwargs) # type: ignore
+ cache = call_signature.bind_partial(self, *args, **kwargs).arguments.get(
+ "cache", None
+ )
+
+ # Add dependency to last cache entry to ensure distributed ops are evaluated
+ if cache is not None:
+ cache[-1].state = mx.depends(cache[-1].state, logits) # type: ignore
+
+ logits = mx.distributed.all_gather(logits, group=group)[
+ -logits.shape[0] :
+ ] # type :ignore
+
+ return logits
+
+ cls.__call__ = patched_call
return model
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
index ea263b6f..daafc3e0 100644
--- a/src/exo/worker/tests/unittests/test_mlx/conftest.py
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -1,12 +1,24 @@
# type: ignore
+import json
+import os
+import tempfile
+import traceback
from dataclasses import dataclass
from pathlib import Path
-from typing import Any
+from typing import Any, cast
import mlx.core as mx
import mlx.nn as nn
from exo.shared.constants import EXO_MODELS_DIR
+from exo.shared.types.api import ChatCompletionMessage
+from exo.shared.types.memory import Memory
+from exo.shared.types.models import ModelId, ModelMetadata
+from exo.shared.types.tasks import ChatCompletionTaskParams
+from exo.shared.types.worker.shards import PipelineShardMetadata, TensorShardMetadata
+from exo.worker.engines.mlx import Model
+from exo.worker.engines.mlx.generator.generate import mlx_generate
+from exo.worker.engines.mlx.utils_mlx import shard_and_load
class MockLayer(nn.Module):
@@ -28,9 +40,6 @@ class PipelineTestConfig:
def create_hostfile(world_size: int, base_port: int) -> tuple[str, list[str]]:
- import json
- import tempfile
-
hosts = [f"127.0.0.1:{base_port + i}" for i in range(world_size)]
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
@@ -50,60 +59,30 @@ DEFAULT_GPT_OSS_CONFIG = PipelineTestConfig(
)
+DEFAULT_GPT_OSS_MODEL_ID = "mlx-community/gpt-oss-20b-MXFP4-Q8"
+
+
def run_gpt_oss_pipeline_device(
rank: int,
world_size: int,
hostfile_path: str,
- model_path: Path,
layer_splits: list[tuple[int, int]],
prompt_tokens: int,
prefill_step_size: int,
result_queue: Any, # pyright: ignore[reportAny]
max_tokens: int = 200,
) -> None:
- import os
- import traceback
-
os.environ["MLX_HOSTFILE"] = hostfile_path
os.environ["MLX_RANK"] = str(rank)
- import mlx.core as mlx_core
- from mlx_lm import load, stream_generate
-
- from exo.shared.types.memory import Memory
- from exo.shared.types.models import ModelId, ModelMetadata
- from exo.shared.types.worker.shards import PipelineShardMetadata
- from exo.worker.engines.mlx.auto_parallel import pipeline_auto_parallel
-
try:
- group = mlx_core.distributed.init(backend="ring", strict=True)
-
- model, tokenizer = load(str(model_path))
-
- # Generate a prompt of exact token length
- base_text = "The quick brown fox jumps over the lazy dog. "
- base_tokens = tokenizer.encode(base_text)
- base_len = len(base_tokens)
-
- # Build prompt with approximate target length
- repeats = (prompt_tokens // base_len) + 2
- long_text = base_text * repeats
- tokens = tokenizer.encode(long_text)
- # Truncate to exact target length
- tokens = tokens[:prompt_tokens]
- prompt_text = tokenizer.decode(tokens)
-
- formatted_prompt = tokenizer.apply_chat_template(
- [{"role": "user", "content": prompt_text}],
- tokenize=False,
- add_generation_prompt=True,
- )
+ group = mx.distributed.init(backend="ring", strict=True)
start_layer, end_layer = layer_splits[rank]
shard_meta = PipelineShardMetadata(
model_meta=ModelMetadata(
- model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
+ model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID),
pretty_name="GPT-OSS 20B",
storage_size=Memory.from_gb(12),
n_layers=24,
@@ -117,21 +96,37 @@ def run_gpt_oss_pipeline_device(
n_layers=24,
)
- model = pipeline_auto_parallel(model, group, shard_meta)
+ model, tokenizer = shard_and_load(shard_meta, group)
+ model = cast(Model, model)
+
+ # Generate a prompt of exact token length
+ base_text = "The quick brown fox jumps over the lazy dog. "
+ base_tokens = tokenizer.encode(base_text)
+ base_len = len(base_tokens)
+
+ # Build prompt with approximate target length
+ repeats = (prompt_tokens // base_len) + 2
+ long_text = base_text * repeats
+ tokens = tokenizer.encode(long_text)
+ # Truncate to exact target length
+ tokens = tokens[:prompt_tokens]
+ prompt_text = tokenizer.decode(tokens)
- # Barrier before generation
- barrier = mlx_core.distributed.all_sum(mlx_core.array([1.0]), group=group)
- mlx_core.eval(barrier)
+ task = ChatCompletionTaskParams(
+ model=DEFAULT_GPT_OSS_MODEL_ID,
+ messages=[ChatCompletionMessage(role="user", content=prompt_text)],
+ max_tokens=max_tokens,
+ )
generated_text = ""
- for response in stream_generate(
+ for response in mlx_generate(
model=model,
tokenizer=tokenizer,
- prompt=formatted_prompt,
- max_tokens=max_tokens,
- prefill_step_size=prefill_step_size,
+ task=task,
):
generated_text += response.text
+ if response.finish_reason is not None:
+ break
result_queue.put((rank, True, generated_text)) # pyright: ignore[reportAny]
@@ -143,27 +138,36 @@ def run_gpt_oss_tensor_parallel_device(
rank: int,
world_size: int,
hostfile_path: str,
- model_path: Path,
prompt_tokens: int,
prefill_step_size: int,
result_queue: Any, # pyright: ignore[reportAny]
max_tokens: int = 10,
) -> None:
- import os
- import traceback
-
os.environ["MLX_HOSTFILE"] = hostfile_path
os.environ["MLX_RANK"] = str(rank)
- import mlx.core as mlx_core
- from mlx_lm import load, stream_generate
-
- from exo.worker.engines.mlx.auto_parallel import tensor_auto_parallel
-
try:
- group = mlx_core.distributed.init(backend="ring", strict=True)
+ group = mx.distributed.init(backend="ring", strict=True)
+
+ # For tensor parallelism, all devices run all layers
+ shard_meta = TensorShardMetadata(
+ model_meta=ModelMetadata(
+ model_id=ModelId(DEFAULT_GPT_OSS_MODEL_ID),
+ pretty_name="GPT-OSS 20B",
+ storage_size=Memory.from_gb(12),
+ n_layers=24,
+ hidden_size=2880,
+ supports_tensor=True,
+ ),
+ device_rank=rank,
+ world_size=world_size,
+ start_layer=0,
+ end_layer=24,
+ n_layers=24,
+ )
- model, tokenizer = load(str(model_path))
+ model, tokenizer = shard_and_load(shard_meta, group)
+ model = cast(Model, model)
base_text = "The quick brown fox jumps over the lazy dog. "
base_tokens = tokenizer.encode(base_text)
@@ -175,26 +179,21 @@ def run_gpt_oss_tensor_parallel_device(
tokens = tokens[:prompt_tokens]
prompt_text = tokenizer.decode(tokens)
- formatted_prompt = tokenizer.apply_chat_template(
- [{"role": "user", "content": prompt_text}],
- tokenize=False,
- add_generation_prompt=True,
+ task = ChatCompletionTaskParams(
+ model=DEFAULT_GPT_OSS_MODEL_ID,
+ messages=[ChatCompletionMessage(role="user", content=prompt_text)],
+ max_tokens=max_tokens,
)
- model = tensor_auto_parallel(model, group)
-
- barrier = mlx_core.distributed.all_sum(mlx_core.array([1.0]), group=group)
- mlx_core.eval(barrier)
-
generated_text = ""
- for response in stream_generate(
+ for response in mlx_generate(
model=model,
tokenizer=tokenizer,
- prompt=formatted_prompt,
- max_tokens=max_tokens,
- prefill_step_size=prefill_step_size,
+ task=task,
):
generated_text += response.text
+ if response.finish_reason is not None:
+ break
result_queue.put((rank, True, generated_text)) # pyright: ignore[reportAny]
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
index 14fccfcb..850db320 100644
--- a/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
+++ b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
@@ -1,13 +1,18 @@
+import json
import multiprocessing as mp
+import os
+import tempfile
from typing import Any
import mlx.core as mx
+import mlx.nn as mlx_nn
import pytest
from exo.worker.engines.mlx.auto_parallel import (
CustomMlxLayer,
PipelineFirstLayer,
PipelineLastLayer,
+ patch_pipeline_model,
)
from exo.worker.tests.unittests.test_mlx.conftest import MockLayer
@@ -23,30 +28,38 @@ def run_pipeline_device(
os.environ["MLX_HOSTFILE"] = hostfile_path
os.environ["MLX_RANK"] = str(rank)
- import mlx.core as mlx_core
- import mlx.nn as mlx_nn
-
class MockLayerInner(mlx_nn.Module):
def __init__(self) -> None:
super().__init__()
self.custom_attr = "test_value"
- def __call__(
- self, x: mlx_core.array, *args: object, **kwargs: object
- ) -> mlx_core.array:
+ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
return x * 2
+ class MockModel(mlx_nn.Module):
+ def __init__(self, layers: list[mlx_nn.Module]) -> None:
+ super().__init__()
+ self.layers = layers
+
+ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+ for layer in self.layers:
+ x = layer(x, *args, **kwargs) # pyright: ignore[reportUnknownVariableType]
+ return x # pyright: ignore[reportUnknownVariableType]
+
try:
- group = mlx_core.distributed.init(backend="ring", strict=True)
+ group = mx.distributed.init(backend="ring", strict=True)
mock = MockLayerInner()
first = PipelineFirstLayer(mock, r=rank, group=group)
composed = PipelineLastLayer(first, r=rank, s=world_size, group=group)
- x = mlx_core.ones((1, 4))
- result = composed(x)
- mlx_core.eval(result)
+ # Wrap in a mock model, then wrap in PipelineParallelModel for all_gather
+ inner_model = MockModel([composed])
+ model = patch_pipeline_model(inner_model, group)
+ x = mx.ones((1, 4))
+ result = model(x)
+ mx.eval(result)
success = result.shape == x.shape
result_queue.put((rank, success, result)) # pyright: ignore[reportAny]
except Exception as e:
@@ -81,10 +94,6 @@ def test_missing_attribute_raises() -> None:
def test_composed_call_works() -> None:
- import json
- import os
- import tempfile
-
ctx = mp.get_context("spawn")
world_size = 2
diff --git a/src/exo/worker/tests/unittests/test_mlx/test_distributed_fix.py b/src/exo/worker/tests/unittests/test_mlx/test_distributed_fix.py
new file mode 100644
index 00000000..45d5bdbb
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_mlx/test_distributed_fix.py
@@ -0,0 +1,230 @@
+import multiprocessing as mp
+import os
+from dataclasses import dataclass
+from typing import Any, Callable
+
+import pytest
+
+from exo.worker.tests.unittests.test_mlx.conftest import (
+ DEFAULT_GPT_OSS_CONFIG,
+ create_hostfile,
+ run_gpt_oss_pipeline_device,
+ run_gpt_oss_tensor_parallel_device,
+)
+
+
+def _check_model_exists() -> bool:
+ return DEFAULT_GPT_OSS_CONFIG.model_path.exists()
+
+
+pytestmark = [
+ pytest.mark.skipif(
+ not _check_model_exists(),
+ reason=f"GPT-OSS model not found at {DEFAULT_GPT_OSS_CONFIG.model_path}",
+ ),
+]
+
+
+@dataclass
+class DistributedTestResult:
+ timed_out: bool
+ world_size: int
+ results: dict[int, tuple[bool, str]]
+
+ @property
+ def all_success(self) -> bool:
+ if len(self.results) != self.world_size:
+ return False
+ return all(r[0] for r in self.results.values())
+
+
+def run_distributed_test(
+ world_size: int,
+ port_offset: int,
+ process_timeout: int,
+ target: Callable[..., None],
+ make_args: Callable[[int], tuple[Any, ...]],
+) -> DistributedTestResult:
+ ctx = mp.get_context("spawn")
+ hostfile_path, _ = create_hostfile(
+ world_size, DEFAULT_GPT_OSS_CONFIG.base_port + port_offset
+ )
+
+ try:
+ result_queue: Any = ctx.Queue()
+ processes: list[Any] = []
+
+ for rank in range(world_size):
+ args = make_args(rank)
+ p = ctx.Process(
+ target=target,
+ args=(rank, world_size, hostfile_path, *args, result_queue),
+ )
+ p.start()
+ processes.append(p)
+
+ for p in processes: # pyright: ignore[reportAny]
+ p.join(timeout=process_timeout) # pyright: ignore[reportAny]
+
+ timed_out = any(p.is_alive() for p in processes) # pyright: ignore[reportAny]
+
+ for p in processes: # pyright: ignore[reportAny]
+ if p.is_alive(): # pyright: ignore[reportAny]
+ p.terminate() # pyright: ignore[reportAny]
+ p.join(timeout=5) # pyright: ignore[reportAny]
+
+ results: dict[int, tuple[bool, str]] = {}
+ while not result_queue.empty(): # pyright: ignore[reportAny]
+ rank, success, value = result_queue.get() # pyright: ignore[reportAny]
+ results[rank] = (success, value)
+
+ return DistributedTestResult(
+ timed_out=timed_out, world_size=world_size, results=results
+ )
+
+ finally:
+ os.unlink(hostfile_path)
+
+
+def run_pipeline_test(
+ layer_splits: list[tuple[int, int]],
+ prompt_tokens: int,
+ prefill_step_size: int,
+ port_offset: int = 0,
+ process_timeout: int = 60,
+) -> DistributedTestResult:
+ def make_args(rank: int) -> tuple[Any, ...]:
+ return (
+ layer_splits,
+ prompt_tokens,
+ prefill_step_size,
+ )
+
+ return run_distributed_test(
+ world_size=len(layer_splits),
+ port_offset=port_offset,
+ process_timeout=process_timeout,
+ target=run_gpt_oss_pipeline_device,
+ make_args=make_args,
+ )
+
+
+def run_tensor_test(
+ prompt_tokens: int,
+ prefill_step_size: int,
+ port_offset: int = 0,
+ process_timeout: int = 60,
+) -> DistributedTestResult:
+ def make_args(rank: int) -> tuple[Any, ...]:
+ return (
+ prompt_tokens,
+ prefill_step_size,
+ )
+
+ return run_distributed_test(
+ world_size=2,
+ port_offset=port_offset,
+ process_timeout=process_timeout,
+ target=run_gpt_oss_tensor_parallel_device,
+ make_args=make_args,
+ )
+
+
+class TestPipelineParallelFix:
+ BUG_TRIGGER_SPLITS: list[tuple[int, int]] = [(0, 1), (1, 24)]
+
+ def test_pipeline_single_layer_first_device(self) -> None:
+ result = run_pipeline_test(
+ layer_splits=self.BUG_TRIGGER_SPLITS,
+ prompt_tokens=100,
+ prefill_step_size=64,
+ process_timeout=60,
+ )
+ assert not result.timed_out, "Unexpected timeout - fix may not be working"
+ assert result.all_success, f"Failures: {result.results}"
+
+
+class TestPipelineSplitConfigurations:
+ @pytest.mark.parametrize(
+ "layer_splits",
+ [
+ [(0, 1), (1, 24)],
+ [(0, 6), (6, 24)],
+ [(0, 12), (12, 24)],
+ ],
+ ids=["1_23", "6_18", "12_12"],
+ )
+ def test_pipeline_splits(
+ self,
+ layer_splits: list[tuple[int, int]],
+ ) -> None:
+ result = run_pipeline_test(
+ layer_splits=layer_splits,
+ prompt_tokens=600,
+ prefill_step_size=512,
+ port_offset=100,
+ )
+ assert not result.timed_out, f"Timeout with {layer_splits}"
+ assert result.all_success, f"Failures with {layer_splits}: {result.results}"
+
+
+class TestPrefillStepSizeBoundaries:
+ @pytest.mark.parametrize(
+ "prefill_step_size,prompt_tokens",
+ [
+ (512, 511),
+ (512, 512),
+ (512, 513),
+ (512, 1024),
+ ],
+ ids=["under", "exact", "over", "double"],
+ )
+ def test_boundary_conditions(
+ self,
+ prefill_step_size: int,
+ prompt_tokens: int,
+ ) -> None:
+ result = run_pipeline_test(
+ layer_splits=[(0, 12), (12, 24)],
+ prompt_tokens=prompt_tokens,
+ prefill_step_size=prefill_step_size,
+ port_offset=200,
+ )
+ assert not result.timed_out, f"Timeout: {prompt_tokens=}, {prefill_step_size=}"
+ assert result.all_success, f"Failures: {result.results}"
+
+
+class TestTensorParallelFix:
+ def test_tensor_parallel(self) -> None:
+ result = run_tensor_test(
+ prompt_tokens=100,
+ prefill_step_size=64,
+ port_offset=400,
+ )
+ assert not result.timed_out, "Unexpected timeout"
+ assert result.all_success, f"Failures: {result.results}"
+
+
+class TestTensorParallelBoundaries:
+ @pytest.mark.parametrize(
+ "prefill_step_size,prompt_tokens",
+ [
+ (512, 511),
+ (512, 512),
+ (512, 513),
+ (512, 1024),
+ ],
+ ids=["under", "exact", "over", "double"],
+ )
+ def test_tensor_parallel_boundaries(
+ self,
+ prefill_step_size: int,
+ prompt_tokens: int,
+ ) -> None:
+ result = run_tensor_test(
+ prompt_tokens=prompt_tokens,
+ prefill_step_size=prefill_step_size,
+ port_offset=500,
+ )
+ assert not result.timed_out, f"Timeout: {prompt_tokens=}, {prefill_step_size=}"
+ assert result.all_success, f"Failures: {result.results}"
← 5ab1f8b3 NetworkSetupHelper: detect stale startup script content
·
back to Exo
·
Split NodePerformanceProfile into granular state mappings (# ee43b598 →