← back to Exo
Custom mlx layer composition (#1201)
ea0588429b6aeb049799568f5bc9786b61d4e7bf · 2026-01-19 12:36:25 +0000 · rltakashige
## Motivation
With a single pipeline layer, PipelineFirstLayer gets composed with
PipelineLastLayer.
## Changes
<!-- Describe what you changed in detail -->
## Why It Works
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
### Automated Testing
Made failing tests. Fixed them!
Files touched
M src/exo/worker/engines/mlx/auto_parallel.pyA src/exo/worker/tests/unittests/test_mlx/conftest.pyA src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
Diff
commit ea0588429b6aeb049799568f5bc9786b61d4e7bf
Author: rltakashige <rl.takashige@gmail.com>
Date: Mon Jan 19 12:36:25 2026 +0000
Custom mlx layer composition (#1201)
## Motivation
With a single pipeline layer, PipelineFirstLayer gets composed with
PipelineLastLayer.
## Changes
<!-- Describe what you changed in detail -->
## Why It Works
<!-- Explain why your approach solves the problem -->
## Test Plan
### Manual Testing
### Automated Testing
Made failing tests. Fixed them!
---
src/exo/worker/engines/mlx/auto_parallel.py | 8 +-
.../worker/tests/unittests/test_mlx/conftest.py | 202 +++++++++++++++++++++
.../tests/unittests/test_mlx/test_auto_parallel.py | 137 ++++++++++++++
3 files changed, 344 insertions(+), 3 deletions(-)
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 7ae9c527..c2819f64 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -46,9 +46,11 @@ class CustomMlxLayer(nn.Module):
def __init__(self, original_layer: _LayerCallable):
super().__init__()
- # Set twice to avoid __setattr__ recursion
object.__setattr__(self, "_original_layer", original_layer)
- self.original_layer: _LayerCallable = original_layer
+
+ @property
+ def original_layer(self) -> _LayerCallable:
+ return cast(_LayerCallable, object.__getattribute__(self, "_original_layer"))
# Calls __getattr__ for any attributes not found on nn.Module (e.g. use_sliding)
if not TYPE_CHECKING:
@@ -58,7 +60,7 @@ class CustomMlxLayer(nn.Module):
return super().__getattr__(name)
except AttributeError:
original_layer = object.__getattribute__(self, "_original_layer")
- return object.__getattribute__(original_layer, name)
+ return getattr(original_layer, name)
class PipelineFirstLayer(CustomMlxLayer):
diff --git a/src/exo/worker/tests/unittests/test_mlx/conftest.py b/src/exo/worker/tests/unittests/test_mlx/conftest.py
new file mode 100644
index 00000000..ea263b6f
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_mlx/conftest.py
@@ -0,0 +1,202 @@
+# type: ignore
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from exo.shared.constants import EXO_MODELS_DIR
+
+
+class MockLayer(nn.Module):
+ def __init__(self) -> None:
+ super().__init__()
+ self.custom_attr = "test_value"
+ self.use_sliding = True
+
+ def __call__(self, x: mx.array, *args: object, **kwargs: object) -> mx.array:
+ return x * 2
+
+
+@dataclass(frozen=True)
+class PipelineTestConfig:
+ model_path: Path
+ total_layers: int
+ base_port: int
+ max_tokens: int
+
+
+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:
+ json.dump(hosts, f)
+ hostfile_path = f.name
+
+ return hostfile_path, hosts
+
+
+# Use GPT OSS 20b to test as it is a model with a lot of strange behaviour
+
+DEFAULT_GPT_OSS_CONFIG = PipelineTestConfig(
+ model_path=EXO_MODELS_DIR / "mlx-community--gpt-oss-20b-MXFP4-Q8",
+ total_layers=24,
+ base_port=29600,
+ max_tokens=200,
+)
+
+
+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,
+ )
+
+ start_layer, end_layer = layer_splits[rank]
+
+ shard_meta = PipelineShardMetadata(
+ model_meta=ModelMetadata(
+ model_id=ModelId("mlx-community/gpt-oss-20b-MXFP4-Q8"),
+ pretty_name="GPT-OSS 20B",
+ storage_size=Memory.from_gb(12),
+ n_layers=24,
+ hidden_size=2880,
+ supports_tensor=False,
+ ),
+ device_rank=rank,
+ world_size=world_size,
+ start_layer=start_layer,
+ end_layer=end_layer,
+ n_layers=24,
+ )
+
+ model = pipeline_auto_parallel(model, group, shard_meta)
+
+ # Barrier before generation
+ barrier = mlx_core.distributed.all_sum(mlx_core.array([1.0]), group=group)
+ mlx_core.eval(barrier)
+
+ generated_text = ""
+ for response in stream_generate(
+ model=model,
+ tokenizer=tokenizer,
+ prompt=formatted_prompt,
+ max_tokens=max_tokens,
+ prefill_step_size=prefill_step_size,
+ ):
+ generated_text += response.text
+
+ result_queue.put((rank, True, generated_text)) # pyright: ignore[reportAny]
+
+ except Exception as e:
+ result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}")) # pyright: ignore[reportAny]
+
+
+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)
+
+ model, tokenizer = load(str(model_path))
+
+ base_text = "The quick brown fox jumps over the lazy dog. "
+ base_tokens = tokenizer.encode(base_text)
+ base_len = len(base_tokens)
+
+ repeats = (prompt_tokens // base_len) + 2
+ long_text = base_text * repeats
+ tokens = tokenizer.encode(long_text)
+ 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,
+ )
+
+ 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(
+ model=model,
+ tokenizer=tokenizer,
+ prompt=formatted_prompt,
+ max_tokens=max_tokens,
+ prefill_step_size=prefill_step_size,
+ ):
+ generated_text += response.text
+
+ result_queue.put((rank, True, generated_text)) # pyright: ignore[reportAny]
+
+ except Exception as e:
+ result_queue.put((rank, False, f"{e}\n{traceback.format_exc()}")) # 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
new file mode 100644
index 00000000..14fccfcb
--- /dev/null
+++ b/src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py
@@ -0,0 +1,137 @@
+import multiprocessing as mp
+from typing import Any
+
+import mlx.core as mx
+import pytest
+
+from exo.worker.engines.mlx.auto_parallel import (
+ CustomMlxLayer,
+ PipelineFirstLayer,
+ PipelineLastLayer,
+)
+from exo.worker.tests.unittests.test_mlx.conftest import MockLayer
+
+
+def run_pipeline_device(
+ rank: int,
+ world_size: int,
+ hostfile_path: str,
+ result_queue: Any, # pyright: ignore[reportAny]
+) -> None:
+ import os
+
+ 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:
+ return x * 2
+
+ try:
+ group = mlx_core.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)
+
+ success = result.shape == x.shape
+ result_queue.put((rank, success, result)) # pyright: ignore[reportAny]
+ except Exception as e:
+ result_queue.put((rank, False, str(e))) # pyright: ignore[reportAny]
+
+
+def test_single_wrapper_delegates_attributes() -> None:
+ mock = MockLayer()
+ wrapped = CustomMlxLayer(mock)
+
+ assert wrapped.custom_attr == "test_value" # type: ignore[attr-defined]
+ assert wrapped.use_sliding is True # type: ignore[attr-defined]
+
+
+def test_composed_wrappers_delegate_attributes() -> None:
+ mock = MockLayer()
+ group = mx.distributed.init()
+
+ first = PipelineFirstLayer(mock, r=0, group=group)
+ composed = PipelineLastLayer(first, r=0, s=1, group=group)
+
+ assert composed.custom_attr == "test_value" # type: ignore[attr-defined]
+ assert composed.use_sliding is True # type: ignore[attr-defined]
+
+
+def test_missing_attribute_raises() -> None:
+ mock = MockLayer()
+ wrapped = CustomMlxLayer(mock)
+
+ with pytest.raises(AttributeError):
+ _ = wrapped.nonexistent_attr # type: ignore[attr-defined]
+
+
+def test_composed_call_works() -> None:
+ import json
+ import os
+ import tempfile
+
+ ctx = mp.get_context("spawn")
+
+ world_size = 2
+ base_port = 29500
+
+ 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:
+ json.dump(hosts, f)
+ hostfile_path = f.name
+
+ try:
+ result_queue: Any = ctx.Queue()
+
+ processes: list[Any] = []
+ for rank in range(world_size):
+ p = ctx.Process(
+ target=run_pipeline_device,
+ args=(rank, world_size, hostfile_path, result_queue),
+ )
+ p.start()
+ processes.append(p)
+
+ for p in processes: # pyright: ignore[reportAny]
+ p.join(timeout=10) # pyright: ignore[reportAny]
+
+ results: dict[int, Any] = {}
+ errors: dict[int, str] = {}
+ while not result_queue.empty(): # pyright: ignore[reportAny]
+ rank, success, value = result_queue.get() # pyright: ignore[reportAny]
+ if success:
+ results[rank] = value
+ else:
+ errors[rank] = value
+
+ assert len(results) == world_size, (
+ f"Expected {world_size} results, got {len(results)}. Errors: {errors}"
+ )
+
+ for rank in range(world_size):
+ assert rank in results, (
+ f"Device {rank} failed: {errors.get(rank, 'unknown')}"
+ )
+ result_array = results[rank]
+ # Both devices see the final result (4.0) after all_gather
+ assert (result_array == 4.0).all(), (
+ f"Device {rank}: expected 4.0, got {result_array}"
+ )
+ finally:
+ os.unlink(hostfile_path)
← 73b3f87e Set swa_idx and ga_idx for single layer (#1202)
·
back to Exo
·
Enhance LaTeX rendering in dashboard markdown (#1197) 346b13e2 →