← back to Exo
Fix GLM4Moe Tensor Sharding (#1411)
5455a97a8c4186390a72f68fdc1477f18bf00226 · 2026-02-06 16:53:15 +0000 · rltakashige
## Motivation
Recent commit broke glm (non lite) sharding
## Why It Works
Assert is no longer hit, as isinstance check includes
GLM4MoeDecoderLayer.
Added type stubs to keep the type checker happy.
## Test Plan
### Manual Testing
Runs as expected without gibberish.
Files touched
A .mlx_typings/mlx_lm/models/glm4_moe.pyiM src/exo/worker/engines/mlx/auto_parallel.pyM src/exo/worker/engines/mlx/cache.py
Diff
commit 5455a97a8c4186390a72f68fdc1477f18bf00226
Author: rltakashige <rl.takashige@gmail.com>
Date: Fri Feb 6 16:53:15 2026 +0000
Fix GLM4Moe Tensor Sharding (#1411)
## Motivation
Recent commit broke glm (non lite) sharding
## Why It Works
Assert is no longer hit, as isinstance check includes
GLM4MoeDecoderLayer.
Added type stubs to keep the type checker happy.
## Test Plan
### Manual Testing
Runs as expected without gibberish.
---
.mlx_typings/mlx_lm/models/glm4_moe.pyi | 153 ++++++++++++++++++++++++++++
src/exo/worker/engines/mlx/auto_parallel.py | 113 ++++++++++----------
src/exo/worker/engines/mlx/cache.py | 1 -
3 files changed, 214 insertions(+), 53 deletions(-)
diff --git a/.mlx_typings/mlx_lm/models/glm4_moe.pyi b/.mlx_typings/mlx_lm/models/glm4_moe.pyi
new file mode 100644
index 00000000..9e824922
--- /dev/null
+++ b/.mlx_typings/mlx_lm/models/glm4_moe.pyi
@@ -0,0 +1,153 @@
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .switch_layers import SwitchGLU
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+ model_type: str
+ vocab_size: int
+ hidden_size: int
+ intermediate_size: int
+ max_position_embeddings: int
+ moe_intermediate_size: int
+ norm_topk_prob: bool
+ num_attention_heads: int
+ n_group: int
+ head_dim: int
+ topk_group: int
+ n_shared_experts: int
+ n_routed_experts: int
+ routed_scaling_factor: float
+ num_experts_per_tok: int
+ first_k_dense_replace: int
+ num_hidden_layers: int
+ num_key_value_heads: int
+ rms_norm_eps: float
+ rope_theta: float
+ rope_scaling: Optional[Dict[str, Any]]
+ use_qk_norm: bool
+ tie_word_embeddings: bool
+ attention_bias: bool
+ partial_rotary_factor: float
+ scoring_func: str
+ topk_method: str
+
+class Attention(nn.Module):
+ n_heads: int
+ n_kv_heads: int
+ scale: float
+ q_proj: nn.Linear
+ k_proj: nn.Linear
+ v_proj: nn.Linear
+ o_proj: nn.Linear
+ use_qk_norm: bool
+ q_norm: nn.RMSNorm
+ k_norm: nn.RMSNorm
+ rope: nn.RoPE
+
+ def __init__(self, args: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ x: mx.array,
+ mask: Optional[mx.array] = None,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+
+class MLP(nn.Module):
+ config: ModelArgs
+ hidden_size: int
+ intermediate_size: int
+ gate_proj: nn.Linear
+ up_proj: nn.Linear
+ down_proj: nn.Linear
+
+ def __init__(
+ self,
+ config: ModelArgs,
+ hidden_size: Optional[int] = None,
+ intermediate_size: Optional[int] = None,
+ ) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class MoEGate(nn.Module):
+ config: ModelArgs
+ top_k: int
+ norm_topk_prob: bool
+ n_routed_experts: int
+ routed_scaling_factor: float
+ n_group: int
+ topk_group: int
+ weight: mx.array
+ e_score_correction_bias: mx.array
+
+ def __init__(self, config: ModelArgs) -> None: ...
+ def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: ...
+
+class MoE(nn.Module):
+ config: ModelArgs
+ num_experts_per_tok: int
+ switch_mlp: SwitchGLU
+ gate: MoEGate
+ shared_experts: MLP
+ sharding_group: Optional[mx.distributed.Group]
+
+ def __init__(self, config: ModelArgs) -> None: ...
+ def __call__(self, x: mx.array) -> mx.array: ...
+
+class DecoderLayer(nn.Module):
+ self_attn: Attention
+ mlp: MLP | MoE
+ input_layernorm: nn.RMSNorm
+ post_attention_layernorm: nn.RMSNorm
+
+ def __init__(self, config: ModelArgs, layer_idx: int) -> None: ...
+ def __call__(
+ self,
+ x: mx.array,
+ mask: Optional[mx.array] = None,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+
+class LanguageModel(nn.Module):
+ vocab_size: int
+ embed_tokens: nn.Embedding
+ layers: list[DecoderLayer]
+ norm: nn.RMSNorm
+ pipeline_rank: int
+ pipeline_size: int
+ start_idx: int
+ end_idx: Optional[int]
+ num_layers: int
+
+ def __init__(self, config: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ x: mx.array,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+ @property
+ def pipeline_layers(self) -> list[DecoderLayer]: ...
+
+class Model(nn.Module):
+ args: ModelArgs
+ model_type: str
+ model: LanguageModel
+ lm_head: nn.Linear
+
+ def __init__(self, config: ModelArgs) -> None: ...
+ def __call__(
+ self,
+ inputs: mx.array,
+ cache: Optional[Any] = None,
+ ) -> mx.array: ...
+ def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: ...
+ def shard(self, group: Optional[mx.distributed.Group] = None) -> None: ...
+ @property
+ def layers(self) -> list[DecoderLayer]: ...
+ @property
+ def cast_predicate(self) -> Any: ...
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index ae85945e..04e47b2c 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -386,7 +386,15 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
)
- elif isinstance(model, (Qwen3MoeModel, Glm4MoeModel, Qwen3NextModel)):
+ elif isinstance(model, Glm4MoeModel):
+ tensor_parallel_sharding_strategy = Glm4MoeShardingStrategy(
+ 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, Qwen3NextModel)):
tensor_parallel_sharding_strategy = QwenShardingStrategy(
group,
all_to_sharded_linear,
@@ -512,9 +520,6 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
layer.self_attn.q_b_proj
)
- # layer.self_attn.kv_b_proj = self.all_to_sharded_linear(
- # layer.self_attn.kv_b_proj
- # )
layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
layer.self_attn.num_heads //= self.N
@@ -544,7 +549,7 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
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 = ShardedDeepseekV3MoE(layer.mlp) # type: ignore
+ layer.mlp = ShardedMoE(layer.mlp) # type: ignore
layer.mlp.sharding_group = self.group
mx.eval(layer)
@@ -552,7 +557,9 @@ class DeepSeekShardingStrategy(TensorParallelShardingStrategy):
return model
-class ShardedDeepseekV3MoE(CustomMlxLayer):
+class ShardedMoE(CustomMlxLayer):
+ """Wraps any MoE layer with distributed sum_gradients / all_sum."""
+
def __init__(self, layer: _LayerCallable):
super().__init__(layer)
self.sharding_group: mx.distributed.Group | None = None
@@ -623,27 +630,13 @@ class GLM4MoeLiteShardingStrategy(TensorParallelShardingStrategy):
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 = ShardedMoE(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 WrappedMiniMaxAttention(CustomMlxLayer):
def __init__(self, layer: _LayerCallable, group: mx.distributed.Group):
super().__init__(layer)
@@ -756,7 +749,7 @@ class MiniMaxShardingStrategy(TensorParallelShardingStrategy):
self.all_to_sharded_linear_in_place(
layer.block_sparse_moe.switch_mlp.up_proj
)
- layer.block_sparse_moe = ShardedQwenMoE(layer.block_sparse_moe) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
+ layer.block_sparse_moe = ShardedMoE(layer.block_sparse_moe) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
layer.block_sparse_moe.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue]
mx.eval(layer)
return model
@@ -861,9 +854,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
# Shard the MoE. Shard in place since the MoE should be responsible
# for aggregating the results.
- if isinstance(
- layer.mlp, (Qwen3MoeSparseMoeBlock, MoE, Qwen3NextSparseMoeBlock)
- ):
+ if isinstance(layer.mlp, (Qwen3MoeSparseMoeBlock, Qwen3NextSparseMoeBlock)):
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)
@@ -875,7 +866,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
layer.mlp.shared_expert.down_proj
)
self.all_to_sharded_linear_in_place(layer.mlp.shared_expert.up_proj)
- layer.mlp = ShardedQwenMoE(layer.mlp) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
+ layer.mlp = ShardedMoE(layer.mlp) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
layer.mlp.sharding_group = self.group
# Shard the MLP
@@ -888,18 +879,50 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
return model
-class ShardedQwenMoE(CustomMlxLayer):
- def __init__(self, layer: _LayerCallable):
- super().__init__(layer)
- self.sharding_group: mx.distributed.Group | None = None
+class Glm4MoeShardingStrategy(TensorParallelShardingStrategy):
+ def shard_model(
+ self,
+ model: nn.Module,
+ timeout_seconds: float,
+ on_timeout: TimeoutCallback | None,
+ ) -> nn.Module:
+ model = cast(Glm4MoeModel, model)
+ for layer in model.layers:
+ eval_with_timeout(
+ layer.parameters(), timeout_seconds / len(model.layers), on_timeout
+ )
- 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
+ layer.self_attn.q_proj = self.all_to_sharded_linear(layer.self_attn.q_proj)
+ layer.self_attn.k_proj = self.all_to_sharded_linear(layer.self_attn.k_proj)
+ layer.self_attn.v_proj = self.all_to_sharded_linear(layer.self_attn.v_proj)
+ layer.self_attn.o_proj = self.sharded_to_all_linear(layer.self_attn.o_proj)
+ layer.self_attn.n_heads //= self.N
+ layer.self_attn.n_kv_heads //= self.N
+
+ if isinstance(layer.mlp, MoE):
+ 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)
+ 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
+ )
+ layer.mlp = ShardedMoE(layer.mlp) # pyright: ignore[reportAttributeAccessIssue, reportArgumentType]
+ layer.mlp.sharding_group = self.group
+
+ else:
+ 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
class GptOssShardingStrategy(TensorParallelShardingStrategy):
@@ -937,21 +960,7 @@ class GptOssShardingStrategy(TensorParallelShardingStrategy):
self.sharded_to_all_linear_in_place(layer.mlp.experts.down_proj)
self.all_to_sharded_linear_in_place(layer.mlp.experts.up_proj)
- layer.mlp = ShardedGptOssMoE(layer.mlp) # type: ignore
+ layer.mlp = ShardedMoE(layer.mlp) # type: ignore
layer.mlp.sharding_group = self.group # pyright: ignore[reportAttributeAccessIssue]
mx.eval(layer)
return model
-
-
-class ShardedGptOssMoE(CustomMlxLayer):
- def __init__(self, layer: nn.Module):
- 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(x)
- if self.sharding_group is not None:
- y = mx.distributed.all_sum(y, group=self.sharding_group)
- return y
diff --git a/src/exo/worker/engines/mlx/cache.py b/src/exo/worker/engines/mlx/cache.py
index 7f2fb03d..c747ba4e 100644
--- a/src/exo/worker/engines/mlx/cache.py
+++ b/src/exo/worker/engines/mlx/cache.py
@@ -290,7 +290,6 @@ def make_kv_cache(
) -> KVCacheType:
assert hasattr(model, "layers")
- # TODO: Do this for all models
if hasattr(model, "make_cache"):
logger.info("Using MLX LM's make cache")
return model.make_cache() # type: ignore
← 6f0cb999 Ciaran/flux1 kontext (#1394)
·
back to Exo
·
Fix kimi tool calling id (#1413) 3b2f553a →