From e78454de41e819d883952e428c4e29df9cc889c9 Mon Sep 17 00:00:00 2001 From: Ryuichi Leo Takashige Date: Tue, 31 Mar 2026 20:13:12 +0100 Subject: [PATCH] Don't reject uneven placements in placement --- nix/mlx.nix | 2 +- src/exo/master/placement.py | 26 ++----- src/exo/shared/types/worker/shards.py | 11 ++- src/exo/worker/engines/mlx/auto_parallel.py | 83 ++++++++++++++++++++- uv.lock | 12 +-- 5 files changed, 105 insertions(+), 29 deletions(-) diff --git a/nix/mlx.nix b/nix/mlx.nix index 39b0ac9f..6352de9e 100644 --- a/nix/mlx.nix +++ b/nix/mlx.nix @@ -49,7 +49,7 @@ let owner = "rltakashige"; repo = "mlx-jaccl-fix-small-recv"; rev = uvLockMlxRev; - hash = "sha256-GosFIWxIB48Egb1MqJrR3xhsUsQeWdRk5rV93USY6wQ="; + hash = "sha256-WZGQKGcKQR9uyXf5X/a1+79ycPdbcs/spfTykDUjLE4="; }; patches = [ diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py index 07699dae..191de654 100644 --- a/src/exo/master/placement.py +++ b/src/exo/master/placement.py @@ -128,26 +128,12 @@ def place_instance( if len(cycles_with_sufficient_memory) == 0: raise ValueError("No cycles found with sufficient memory") - if command.sharding == Sharding.Tensor: - if not command.model_card.supports_tensor: - raise ValueError( - f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}" - ) - # TODO: the condition here for tensor parallel is not correct, but it works good enough for now. - kv_heads = command.model_card.num_key_value_heads - cycles_with_sufficient_memory = [ - cycle - for cycle in cycles_with_sufficient_memory - if command.model_card.hidden_size % len(cycle) == 0 - and (kv_heads is None or kv_heads % len(cycle) == 0) - ] - if not cycles_with_sufficient_memory: - raise ValueError( - f"No tensor sharding found for model with " - f"hidden_size={command.model_card.hidden_size}" - f"{f', num_key_value_heads={kv_heads}' if kv_heads is not None else ''}" - f" across candidate cycles" - ) + if command.sharding == Sharding.Tensor and not command.model_card.supports_tensor: + raise ValueError( + f"Requested Tensor sharding but this model does not support tensor parallelism: {command.model_card.model_id}" + ) + + # Uneven tensor sharding handles arbitrary world sizes — no divisibility check needed if command.sharding == Sharding.Pipeline and command.model_card.model_id == ModelId( "mlx-community/DeepSeek-V3.1-8bit" ): diff --git a/src/exo/shared/types/worker/shards.py b/src/exo/shared/types/worker/shards.py index f7d83a4a..21936952 100644 --- a/src/exo/shared/types/worker/shards.py +++ b/src/exo/shared/types/worker/shards.py @@ -1,7 +1,7 @@ from enum import Enum from typing import TypeAlias, final -from pydantic import Field +from pydantic import Field, field_validator from exo.shared.models.model_cards import ModelCard from exo.utils.pydantic_ext import TaggedModel @@ -91,6 +91,15 @@ class TensorShardMetadata(BaseShardMetadata): shard_weights: list[float] | None = None shard_mode: TensorShardMode = TensorShardMode.Constant + @field_validator("shard_mode", mode="before") + @classmethod + def _coerce_shard_mode(cls, v: object) -> TensorShardMode: + if isinstance(v, str): + return TensorShardMode(v) + if isinstance(v, TensorShardMode): + return v + raise ValueError(f"expected TensorShardMode or str, got {type(v).__name__}") + ShardMetadata: TypeAlias = ( PipelineShardMetadata | CfgShardMetadata | TensorShardMetadata diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py index 67d5fb35..f9fb2de5 100644 --- a/src/exo/worker/engines/mlx/auto_parallel.py +++ b/src/exo/worker/engines/mlx/auto_parallel.py @@ -536,13 +536,70 @@ def tensor_auto_parallel( return None return -1, segments - sharded_to_all_linear_in_place = partial( + _base_sharded_to_all_in_place = partial( shard_inplace, sharding=_sharded_to_all, # type: ignore group=group, weights=shard_weights, ) + _base_all_to_sharded_in_place = all_to_sharded_linear_in_place + + def _quantized_moe_shard_inplace( + module: nn.Module, + sharding: Literal["all-to-sharded", "sharded-to-all"], + weights: list[float] | None = None, + ) -> None: + N = group.size() + r = group.rank() + gs = module.group_size # pyright: ignore[reportAttributeAccessIssue] + bits = module.bits # pyright: ignore[reportAttributeAccessIssue] + params = module.parameters() + scales = params["scales"] + + if sharding == "all-to-sharded": + dim = params["weight"].shape[max(params["weight"].ndim - 2, 0)] + sizes = compute_shard_sizes(dim, N, gs, weights) + result: dict[str, Any] = {} + for key, param in params.items(): + if not isinstance(param, mx.array): + result[key] = param + continue + axis = max(param.ndim - 2, 0) + indices = [sum(sizes[:i]) for i in range(1, len(sizes))] + result[key] = mx.contiguous(mx.split(param, indices, axis=axis)[r]) + else: + num_groups = scales.shape[-1] + group_counts = compute_shard_sizes(num_groups, N, 1, weights) + weight_ppg = gs * bits // 32 + result = {} + for key, param in params.items(): + if not isinstance(param, mx.array): + result[key] = param + continue + if key == "weight": + s = [gc * weight_ppg for gc in group_counts] + elif key in ("scales", "biases"): + s = list(group_counts) + else: + result[key] = param + continue + indices = [sum(s[:i]) for i in range(1, len(s))] + result[key] = mx.contiguous(mx.split(param, indices, axis=-1)[r]) + module.update(result) + + def all_to_sharded_linear_in_place(module: nn.Module, **kwargs: Any) -> None: + if getattr(module, "group_size", 0) > 0 and getattr(module, "bits", 0) > 0 and "scales" in module.parameters(): + _quantized_moe_shard_inplace(module, "all-to-sharded", weights=kwargs.get("weights")) + else: + _base_all_to_sharded_in_place(module, **kwargs) + + def sharded_to_all_linear_in_place(module: nn.Module, **kwargs: Any) -> None: + if getattr(module, "group_size", 0) > 0 and getattr(module, "bits", 0) > 0 and "scales" in module.parameters(): + _quantized_moe_shard_inplace(module, "sharded-to-all", weights=kwargs.get("weights")) + else: + _base_sharded_to_all_in_place(module, **kwargs) + if isinstance(model, (LlamaModel, Ministral3Model)): tensor_parallel_sharding_strategy = LlamaShardingStrategy( group, @@ -778,16 +835,20 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy): layer.self_attn.k_proj.weight.shape[0] // head_dim ) + mlp_unit = getattr(layer.mlp.gate_proj, "group_size", 1) layer.mlp.gate_proj = self.all_to_sharded_linear( layer.mlp.gate_proj, + unit=mlp_unit, weights=self._greedy_weights_for("gate", intermediate), ) layer.mlp.down_proj = self.sharded_to_all_linear( layer.mlp.down_proj, + unit=mlp_unit, weights=self._greedy_weights_for("down", intermediate), ) layer.mlp.up_proj = self.all_to_sharded_linear( layer.mlp.up_proj, + unit=mlp_unit, weights=self._greedy_weights_for("up", intermediate), ) mx.eval(layer) @@ -890,16 +951,20 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy): # Shard the MLP if isinstance(layer.mlp, (DeepseekV3MLP, DeepseekV32MLP)): intermediate = layer.mlp.gate_proj.weight.shape[0] + mlp_unit = getattr(layer.mlp.gate_proj, "group_size", 1) layer.mlp.gate_proj = self.all_to_sharded_linear( layer.mlp.gate_proj, + unit=mlp_unit, weights=self._greedy_weights_for("gate", intermediate), ) layer.mlp.down_proj = self.sharded_to_all_linear( layer.mlp.down_proj, + unit=mlp_unit, weights=self._greedy_weights_for("down", intermediate), ) layer.mlp.up_proj = self.all_to_sharded_linear( layer.mlp.up_proj, + unit=mlp_unit, weights=self._greedy_weights_for("up", intermediate), ) @@ -1037,16 +1102,20 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy): if isinstance(layer.mlp, Glm4MoeLiteMLP): intermediate = layer.mlp.gate_proj.weight.shape[0] + mlp_unit = getattr(layer.mlp.gate_proj, "group_size", 1) layer.mlp.gate_proj = self.all_to_sharded_linear( layer.mlp.gate_proj, + unit=mlp_unit, weights=self._greedy_weights_for("gate", intermediate), ) layer.mlp.down_proj = self.sharded_to_all_linear( layer.mlp.down_proj, + unit=mlp_unit, weights=self._greedy_weights_for("down", intermediate), ) layer.mlp.up_proj = self.all_to_sharded_linear( layer.mlp.up_proj, + unit=mlp_unit, weights=self._greedy_weights_for("up", intermediate), ) @@ -1516,16 +1585,20 @@ class QwenShardingStrategy(TensorParallelShardingStrategy): # Shard the MLP else: intermediate = layer.mlp.gate_proj.weight.shape[0] + mlp_unit = getattr(layer.mlp.gate_proj, "group_size", 1) layer.mlp.gate_proj = self.all_to_sharded_linear( layer.mlp.gate_proj, + unit=mlp_unit, weights=self._greedy_weights_for("gate", intermediate), ) layer.mlp.down_proj = self.sharded_to_all_linear( layer.mlp.down_proj, + unit=mlp_unit, weights=self._greedy_weights_for("down", intermediate), ) layer.mlp.up_proj = self.all_to_sharded_linear( layer.mlp.up_proj, + unit=mlp_unit, weights=self._greedy_weights_for("up", intermediate), ) @@ -1622,16 +1695,20 @@ class Glm4MoeShardingStrategy(TensorParallelShardingStrategy): else: intermediate = layer.mlp.gate_proj.weight.shape[0] + mlp_unit = getattr(layer.mlp.gate_proj, "group_size", 1) layer.mlp.gate_proj = self.all_to_sharded_linear( layer.mlp.gate_proj, + unit=mlp_unit, weights=self._greedy_weights_for("gate", intermediate), ) layer.mlp.down_proj = self.sharded_to_all_linear( layer.mlp.down_proj, + unit=mlp_unit, weights=self._greedy_weights_for("down", intermediate), ) layer.mlp.up_proj = self.all_to_sharded_linear( layer.mlp.up_proj, + unit=mlp_unit, weights=self._greedy_weights_for("up", intermediate), ) @@ -1792,16 +1869,20 @@ class Step35ShardingStrategy(TensorParallelShardingStrategy): if isinstance(layer.mlp, Step35MLP): intermediate = layer.mlp.gate_proj.weight.shape[0] + mlp_unit = getattr(layer.mlp.gate_proj, "group_size", 1) layer.mlp.gate_proj = self.all_to_sharded_linear( layer.mlp.gate_proj, + unit=mlp_unit, weights=self._greedy_weights_for("gate", intermediate), ) layer.mlp.up_proj = self.all_to_sharded_linear( layer.mlp.up_proj, + unit=mlp_unit, weights=self._greedy_weights_for("up", intermediate), ) layer.mlp.down_proj = self.sharded_to_all_linear( layer.mlp.down_proj, + unit=mlp_unit, weights=self._greedy_weights_for("down", intermediate), ) else: diff --git a/uv.lock b/uv.lock index a063f5ae..025cd7de 100644 --- a/uv.lock +++ b/uv.lock @@ -558,7 +558,7 @@ dependencies = [ { name = "loguru", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mflux", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cpu"], marker = "sys_platform == 'linux'" }, - { name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", version = "0.31.2.dev20260401+fbfe79c7", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }, marker = "sys_platform == 'darwin'" }, { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mlx-vlm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "msgspec", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -1436,7 +1436,7 @@ dependencies = [ { name = "huggingface-hub", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "matplotlib", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, extra = ["cuda13"], marker = "sys_platform == 'linux'" }, - { name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", version = "0.31.2.dev20260401+fbfe79c7", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }, marker = "sys_platform == 'darwin'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "piexif", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -1494,8 +1494,8 @@ cuda13 = [ [[package]] name = "mlx" -version = "0.31.2.dev20260331+71bcd7a2" -source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" } +version = "0.31.2.dev20260401+fbfe79c7" +source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'darwin'", "python_full_version < '3.14' and sys_platform == 'darwin'", @@ -1531,7 +1531,7 @@ version = "0.31.2" source = { git = "https://github.com/rltakashige/mlx-lm?branch=leo%2Ffix-arrayscache-leak#d36e9b661e55a5fc0f77fb6f17ea643aa2dc87aa" } dependencies = [ { name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, - { name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", version = "0.31.2.dev20260401+fbfe79c7", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }, marker = "sys_platform == 'darwin'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "protobuf", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, @@ -1548,7 +1548,7 @@ dependencies = [ { name = "fastapi", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "miniaudio", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "mlx", version = "0.30.6", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'linux'" }, - { name = "mlx", version = "0.31.2.dev20260331+71bcd7a2", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#71bcd7a2c6bfff535d00ed85a2ee78102d8dca03" }, marker = "sys_platform == 'darwin'" }, + { name = "mlx", version = "0.31.2.dev20260401+fbfe79c7", source = { git = "https://github.com/rltakashige/mlx-jaccl-fix-small-recv.git?branch=leo%2Fadd-uneven-sharding#fbfe79c7b1238273969328abc5720ba18f7265d5" }, marker = "sys_platform == 'darwin'" }, { name = "mlx-lm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "numpy", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" }, { name = "opencv-python", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },