From ea0588429b6aeb049799568f5bc9786b61d4e7bf Mon Sep 17 00:00:00 2001 From: rltakashige Date: Mon, 19 Jan 2026 12:36:25 +0000 Subject: [PATCH] Custom mlx layer composition (#1201) ## Motivation With a single pipeline layer, PipelineFirstLayer gets composed with PipelineLastLayer. ## Changes ## Why It Works ## Test Plan ### Manual Testing ### Automated Testing Made failing tests. Fixed them! --- src/exo/worker/engines/mlx/auto_parallel.py | 8 +- .../tests/unittests/test_mlx/conftest.py | 202 ++++++++++++++++++ .../unittests/test_mlx/test_auto_parallel.py | 137 ++++++++++++ 3 files changed, 344 insertions(+), 3 deletions(-) create mode 100644 src/exo/worker/tests/unittests/test_mlx/conftest.py create mode 100644 src/exo/worker/tests/unittests/test_mlx/test_auto_parallel.py 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)