electron-rare/exo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Don't reject uneven placements in placement

Browse Source
This commit is contained in:
Ryuichi Leo Takashige
2026-03-31 20:13:12 +01:00
parent b30ee156c4
commit e78454de41
5 changed files with 105 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
nix/mlx.nix
+1 -1
View File
@@ -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 = [
src/exo/master/placement.py
+6 -20
View File
@@ -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"
):
src/exo/shared/types/worker/shards.py
+10 -1
View File
@@ -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
src/exo/worker/engines/mlx/auto_parallel.py
+82 -1
View File
@@ -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:
uv.lock
Generated
+6 -6
View File
@@ -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'" },