← back to Exo
Leo/fix basic model shard (#1291)
9968abe81605d5e11531377e8de4ff86bcacf7c3 · 2026-01-26 17:49:09 +0000 · rltakashige
## Motivation
Some models, on some configurations, would have several issues that
caused the model to be stuck on loading.
## Changes
Several loading issues were with upstream mlx lm shard loading for
tensor parallel.
GLM 4.7 Flash now uses GLM 4.7 Lite.
A final portion of the issues were from mlx memory not being properly
released before calling mx.eval(model), causing the system to run out of
memory.
## Test Plan
### Manual Testing
Done a bunch (thanks @AlexCheema), hopefully exhaustive.
### Automated Testing
A bunch of automated testing is imminent but not landed yet.
---------
Co-authored-by: Alex Cheema <alexcheema123@gmail.com>
Files touched
M pyproject.tomlM src/exo/worker/engines/mlx/auto_parallel.pyM uv.lock
Diff
commit 9968abe81605d5e11531377e8de4ff86bcacf7c3
Author: rltakashige <rl.takashige@gmail.com>
Date: Mon Jan 26 17:49:09 2026 +0000
Leo/fix basic model shard (#1291)
## Motivation
Some models, on some configurations, would have several issues that
caused the model to be stuck on loading.
## Changes
Several loading issues were with upstream mlx lm shard loading for
tensor parallel.
GLM 4.7 Flash now uses GLM 4.7 Lite.
A final portion of the issues were from mlx memory not being properly
released before calling mx.eval(model), causing the system to run out of
memory.
## Test Plan
### Manual Testing
Done a bunch (thanks @AlexCheema), hopefully exhaustive.
### Automated Testing
A bunch of automated testing is imminent but not landed yet.
---------
Co-authored-by: Alex Cheema <alexcheema123@gmail.com>
---
pyproject.toml | 2 +-
src/exo/worker/engines/mlx/auto_parallel.py | 114 ++++++++++++++++++++++++----
uv.lock | 16 ++--
3 files changed, 109 insertions(+), 23 deletions(-)
diff --git a/pyproject.toml b/pyproject.toml
index 702e198d..319b3155 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,7 +19,7 @@ dependencies = [
"anyio==4.11.0",
"mlx==0.30.3; sys_platform == 'darwin'",
"mlx[cpu]==0.30.3; sys_platform == 'linux'",
- "mlx-lm @ git+https://github.com/AlexCheema/mlx-lm.git@fix-transformers-5.0.0rc2",
+ "mlx-lm==0.30.5",
"tiktoken>=0.12.0", # required for kimi k2 tokenizer
"hypercorn>=0.18.0",
"openai-harmony>=0.0.8",
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 159c2a45..f2ecb698 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -19,6 +19,8 @@ from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
from mlx_lm.models.glm4_moe import Model as Glm4MoeModel
from mlx_lm.models.glm4_moe import MoE
+from mlx_lm.models.glm4_moe_lite import Glm4MoeLiteDecoderLayer, Glm4MoeLiteMLP
+from mlx_lm.models.glm4_moe_lite import Model as GLM4MoeLiteModel
from mlx_lm.models.gpt_oss import GptOssMoeModel
from mlx_lm.models.gpt_oss import Model as GptOssModel
from mlx_lm.models.llama import Model as LlamaModel
@@ -145,6 +147,10 @@ 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] :
+ ] # type :ignore
+
return output
@@ -252,10 +258,6 @@ def patch_pipeline_model[T](model: T, group: mx.distributed.Group) -> T:
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
@@ -334,15 +336,7 @@ def tensor_auto_parallel(
group=group,
)
- if hasattr(model, "shard") and not isinstance(model, GptOssModel):
- try:
- model.shard(group) # type: ignore
- return patch_tensor_model(model)
- except (AttributeError, TypeError, NameError):
- pass
-
if isinstance(model, (LlamaModel, Ministral3Model)):
- logger.warning("shouldn't be hit - upstream sharding exists")
tensor_parallel_sharding_strategy = LlamaShardingStrategy(
group,
all_to_sharded_linear,
@@ -351,7 +345,6 @@ def tensor_auto_parallel(
sharded_to_all_linear_in_place,
)
elif isinstance(model, (DeepseekV3Model, DeepseekV32Model)):
- logger.warning("shouldn't be hit - upstream sharding exists")
tensor_parallel_sharding_strategy = DeepSeekShardingStrategy(
group,
all_to_sharded_linear,
@@ -367,6 +360,14 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
)
+ elif isinstance(model, GLM4MoeLiteModel):
+ tensor_parallel_sharding_strategy = GLM4MoeLiteShardingStrategy(
+ group,
+ all_to_sharded_linear,
+ sharded_to_all_linear,
+ all_to_sharded_linear_in_place,
+ sharded_to_all_linear_in_place,
+ )
elif isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)):
tensor_parallel_sharding_strategy = QwenShardingStrategy(
group,
@@ -441,7 +442,7 @@ class LlamaShardingStrategy(TensorParallelShardingStrategy):
layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
-
+ mx.eval(layer)
return model
@@ -516,6 +517,8 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
layer.mlp = ShardedDeepseekV3MoE(layer.mlp) # type: ignore
layer.mlp.sharding_group = self.group
+ mx.eval(layer)
+
return model
@@ -533,6 +536,84 @@ class ShardedDeepseekV3MoE(CustomMlxLayer):
return y
+class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
+ def shard_model(
+ self,
+ model: nn.Module,
+ timeout_seconds: float,
+ on_timeout: TimeoutCallback | None,
+ ) -> nn.Module:
+ model = cast(GLM4MoeLiteModel, model)
+ for layer in model.layers: # type: ignore
+ layer = cast(Glm4MoeLiteDecoderLayer, layer)
+ eval_with_timeout(
+ layer.parameters(),
+ timeout_seconds / len(model.layers), # type: ignore
+ on_timeout,
+ )
+ if layer.self_attn.q_lora_rank is None: # type: ignore
+ layer.self_attn.q_proj = self.all_to_sharded_linear(
+ layer.self_attn.q_proj
+ )
+ else:
+ layer.self_attn.q_b_proj = self.all_to_sharded_linear(
+ layer.self_attn.q_b_proj
+ )
+
+ layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+ layer.self_attn.num_heads //= self.N
+
+ # Logic from upstream mlx
+ num_heads = layer.self_attn.num_heads
+ sh = self.group.rank() * num_heads
+ eh = sh + num_heads
+
+ def shard_heads(w: mx.array, sh: int = sh, eh: int = eh) -> mx.array:
+ return w[sh:eh]
+
+ layer.self_attn.embed_q.apply(shard_heads)
+ layer.self_attn.unembed_out.apply(shard_heads)
+
+ if isinstance(layer.mlp, Glm4MoeLiteMLP):
+ layer.mlp.gate_proj = self.all_to_sharded_linear(layer.mlp.gate_proj)
+ layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
+ layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
+
+ else:
+ if getattr(layer.mlp, "shared_experts", None) is not None:
+ self.all_to_sharded_linear_in_place(
+ layer.mlp.shared_experts.gate_proj
+ )
+ self.sharded_to_all_linear_in_place(
+ layer.mlp.shared_experts.down_proj
+ )
+ self.all_to_sharded_linear_in_place(
+ layer.mlp.shared_experts.up_proj
+ )
+ self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.gate_proj)
+ self.sharded_to_all_linear_in_place(layer.mlp.switch_mlp.down_proj)
+ self.all_to_sharded_linear_in_place(layer.mlp.switch_mlp.up_proj)
+ layer.mlp = ShardedGLM4MoeLiteMoE(layer.mlp) # type: ignore
+ layer.mlp.sharding_group = self.group # type: ignore
+ mx.eval(layer)
+
+ return model
+
+
+class ShardedGLM4MoeLiteMoE(CustomMlxLayer):
+ def __init__(self, layer: _LayerCallable):
+ super().__init__(layer)
+ self.sharding_group: mx.distributed.Group | None = None
+
+ def __call__(self, x: mx.array) -> mx.array:
+ if self.sharding_group is not None:
+ x = sum_gradients(self.sharding_group)(x)
+ y = self.original_layer.__call__(x)
+ if self.sharding_group is not None:
+ y = mx.distributed.all_sum(y, group=self.sharding_group)
+ return y
+
+
class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
def shard_model(
self,
@@ -566,7 +647,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
)
layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
layer.block_sparse_moe.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue]
-
+ mx.eval(layer)
return model
@@ -607,6 +688,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
layer.mlp.down_proj = self.sharded_to_all_linear(layer.mlp.down_proj)
layer.mlp.up_proj = self.all_to_sharded_linear(layer.mlp.up_proj)
+ mx.eval(layer)
return model
@@ -661,7 +743,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
layer.mlp = ShardedGptOssMoE(layer.mlp) # type: ignore
layer.mlp.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue]
-
+ mx.eval(layer)
return model
diff --git a/uv.lock b/uv.lock
index c81ef825..79e8b5f1 100644
--- a/uv.lock
+++ b/uv.lock
@@ -415,7 +415,7 @@ requires-dist = [
{ name = "mflux", specifier = "==0.15.4" },
{ name = "mlx", marker = "sys_platform == 'darwin'", specifier = "==0.30.3" },
{ name = "mlx", extras = ["cpu"], marker = "sys_platform == 'linux'", specifier = "==0.30.3" },
- { name = "mlx-lm", git = "https://github.com/AlexCheema/mlx-lm.git?rev=fix-transformers-5.0.0rc2" },
+ { name = "mlx-lm", specifier = "==0.30.5" },
{ name = "openai-harmony", specifier = ">=0.0.8" },
{ name = "pillow", specifier = ">=11.0,<12.0" },
{ name = "psutil", specifier = ">=7.0.0" },
@@ -1072,8 +1072,8 @@ wheels = [
[[package]]
name = "mlx-lm"
-version = "0.30.4"
-source = { git = "https://github.com/AlexCheema/mlx-lm.git?rev=fix-transformers-5.0.0rc2#a5daf2b894f31793dfaef0fdf9bc3ed683176ad6" }
+version = "0.30.5"
+source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "jinja2", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "mlx", marker = "sys_platform == 'darwin'" },
@@ -1083,6 +1083,10 @@ dependencies = [
{ name = "sentencepiece", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "transformers", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
+sdist = { url = "https://files.pythonhosted.org/packages/0b/90/4469d9f75f196e6255f59a89441abe0079925d30a001462e1c1c4bc4e6a1/mlx_lm-0.30.5.tar.gz", hash = "sha256:9e6cb258c65b766c6af25cb90958aef40acab67139f05839eef19864cb3154f6", size = 262367, upload-time = "2026-01-25T15:29:30.125Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/89/ba/66db6e1e5f1ef506655b562932f6bd8f72600116d5f31f92d71c1f200b3f/mlx_lm-0.30.5-py3-none-any.whl", hash = "sha256:a80bc8e3efdebe81813b0f6eb403fb66a7a15071e256f4e7102ada986acb75bb", size = 366716, upload-time = "2026-01-25T15:29:28.29Z" },
+]
[[package]]
name = "mlx-metal"
@@ -2281,7 +2285,7 @@ wheels = [
[[package]]
name = "transformers"
-version = "5.0.0rc2"
+version = "5.0.0rc3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "filelock", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
@@ -2296,9 +2300,9 @@ dependencies = [
{ name = "tqdm", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
{ name = "typer-slim", marker = "sys_platform == 'darwin' or sys_platform == 'linux'" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/94/e2/86b1bd5264272953370a5e50a91da38d7a53a87c5faf3fd3ff62d7353879/transformers-5.0.0rc2.tar.gz", hash = "sha256:9f2fa5e132433dd7eb910dc224b32de0baf758f3b6ffc918dbb632e0af85c07a", size = 8362532, upload-time = "2026-01-07T16:58:02.603Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/3f/a3/7c116a8d85f69ea7749cf4c2df79e64c35d028e5fc7ea0168f299d03b8c7/transformers-5.0.0rc3.tar.gz", hash = "sha256:a0315b92b7e087617ade42ec9e6e92ee7620541cc5d6a3331886c52cbe306f5c", size = 8388520, upload-time = "2026-01-14T16:49:02.952Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/b4/eb/9526a77354a2126f5b220f4792dc8494d573773c098dac6a5ad1fc7a5f17/transformers-5.0.0rc2-py3-none-any.whl", hash = "sha256:f8f2a14060ab11f20a0eec39d827af54c1589c327c5799d82808ae3f4167418a", size = 10067329, upload-time = "2026-01-07T16:57:59.617Z" },
+ { url = "https://files.pythonhosted.org/packages/1e/f2/ae2b8968764253bdf38a48dee3c299b8d0bedf7c8ffbe3449fca9bd95338/transformers-5.0.0rc3-py3-none-any.whl", hash = "sha256:383fad27f4f73092d330e45fae384681e5c8521e1dc1cf6cb1a297780e68bf2d", size = 10107087, upload-time = "2026-01-14T16:48:59.393Z" },
]
[[package]]
← 0e30b083 Fix download system for upstream file changes (#1290)
·
back to Exo
·
Ciaran/image quantization (#1272) ffba340e →