← back to Exo
Add gemma 4 tensor parallelism (#1891)
b8eaf707a8270cb28614127e95ec867a52eae640 · 2026-04-14 20:31:59 +0100 · rltakashige
Files touched
M resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-4bit.tomlM resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-6bit.tomlM resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-8bit.tomlM src/exo/master/placement.pyM src/exo/worker/engines/mlx/auto_parallel.py
Diff
commit b8eaf707a8270cb28614127e95ec867a52eae640
Author: rltakashige <rl.takashige@gmail.com>
Date: Tue Apr 14 20:31:59 2026 +0100
Add gemma 4 tensor parallelism (#1891)
---
.../mlx-community--gemma-4-26b-a4b-it-4bit.toml | 2 +-
.../mlx-community--gemma-4-26b-a4b-it-6bit.toml | 2 +-
.../mlx-community--gemma-4-26b-a4b-it-8bit.toml | 2 +-
src/exo/master/placement.py | 11 ++++
src/exo/worker/engines/mlx/auto_parallel.py | 63 ++++++++++++++++++++++
5 files changed, 77 insertions(+), 3 deletions(-)
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-4bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-4bit.toml
index a36497e8..95c0be7b 100644
--- a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-4bit.toml
+++ b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-4bit.toml
@@ -2,7 +2,7 @@ model_id = "mlx-community/gemma-4-26b-a4b-it-4bit"
n_layers = 30
hidden_size = 2816
num_key_value_heads = 8
-supports_tensor = true
+supports_tensor = false
tasks = ["TextGeneration"]
family = "gemma"
quantization = "4bit"
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-6bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-6bit.toml
index bd582311..66c45506 100644
--- a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-6bit.toml
+++ b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-6bit.toml
@@ -2,7 +2,7 @@ model_id = "mlx-community/gemma-4-26b-a4b-it-6bit"
n_layers = 30
hidden_size = 2816
num_key_value_heads = 8
-supports_tensor = true
+supports_tensor = false
tasks = ["TextGeneration"]
family = "gemma"
quantization = "6bit"
diff --git a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-8bit.toml b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-8bit.toml
index 9dda6aa4..f0cc1514 100644
--- a/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-8bit.toml
+++ b/resources/inference_model_cards/mlx-community--gemma-4-26b-a4b-it-8bit.toml
@@ -2,7 +2,7 @@ model_id = "mlx-community/gemma-4-26b-a4b-it-8bit"
n_layers = 30
hidden_size = 2816
num_key_value_heads = 8
-supports_tensor = true
+supports_tensor = false
tasks = ["TextGeneration"]
family = "gemma"
quantization = "8bit"
diff --git a/src/exo/master/placement.py b/src/exo/master/placement.py
index f9380693..a8eae188 100644
--- a/src/exo/master/placement.py
+++ b/src/exo/master/placement.py
@@ -153,6 +153,17 @@ def place_instance(
raise ValueError(
"Pipeline parallelism is not supported for DeepSeek V3.1 (8-bit)"
)
+ if (
+ command.sharding == Sharding.Pipeline
+ and command.model_card.base_model.startswith("Gemma 4")
+ ):
+ cycles_with_sufficient_memory = [
+ cycle for cycle in cycles_with_sufficient_memory if len(cycle) == 1
+ ]
+ if not cycles_with_sufficient_memory:
+ raise ValueError(
+ "Pipeline parallelism is not supported for Gemma 4; use tensor parallelism instead."
+ )
smallest_cycles = get_smallest_cycles(cycles_with_sufficient_memory)
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 29c5320e..9566c838 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -19,6 +19,7 @@ from mlx_lm.models.deepseek_v3 import DeepseekV3MLP
from mlx_lm.models.deepseek_v3 import Model as DeepseekV3Model
from mlx_lm.models.deepseek_v32 import DeepseekV32MLP
from mlx_lm.models.deepseek_v32 import Model as DeepseekV32Model
+from mlx_lm.models.gemma4 import Model as Gemma4Model
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
@@ -572,6 +573,14 @@ def tensor_auto_parallel(
all_to_sharded_linear_in_place,
sharded_to_all_linear_in_place,
)
+ elif isinstance(model, Gemma4Model):
+ tensor_parallel_sharding_strategy = Gemma4ShardingStrategy(
+ group,
+ all_to_sharded_linear,
+ sharded_to_all_linear,
+ all_to_sharded_linear_in_place,
+ sharded_to_all_linear_in_place,
+ )
else:
raise ValueError(f"Unsupported model type: {type(model)}")
@@ -1332,3 +1341,57 @@ class NemotronHShardingStrategy(TensorParallelShardingStrategy):
mixer.intermediate_size = is_per_rank
mixer.conv_dim = new_conv_dim
mixer.heads_per_group = heads_per_rank // groups_per_rank
+
+
+class WrappedGemma4Experts(CustomMlxLayer):
+ def __init__(self, layer: _LayerCallable):
+ super().__init__(layer)
+ self.sharding_group: mx.distributed.Group | None = None
+
+ def __call__(
+ self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
+ ) -> mx.array:
+ if self.sharding_group is not None:
+ x = sum_gradients(self.sharding_group)(x)
+ y: mx.array = self.original_layer(x, top_k_indices, top_k_weights)
+ if self.sharding_group is not None:
+ y = mx.distributed.all_sum(y, group=self.sharding_group)
+ return y
+
+
+class Gemma4ShardingStrategy(TensorParallelShardingStrategy):
+ def shard_model(
+ self,
+ model: nn.Module,
+ on_layer_loaded: LayerLoadedCallback | None,
+ ) -> nn.Module:
+ model = cast(Gemma4Model, model)
+ layers = model.language_model.model.layers
+ total = len(layers)
+ for i, layer in enumerate(layers):
+ mx.eval(layer.parameters())
+
+ attn = layer.self_attn
+ attn.q_proj = self.all_to_sharded_linear(attn.q_proj)
+ attn.k_proj = self.all_to_sharded_linear(attn.k_proj)
+ if not attn.use_k_eq_v:
+ attn.v_proj = self.all_to_sharded_linear(attn.v_proj)
+ attn.o_proj = self.sharded_to_all_linear(attn.o_proj)
+ attn.n_heads //= self.N
+ attn.n_kv_heads //= self.N
+
+ 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)
+
+ if layer.enable_moe:
+ self.all_to_sharded_linear_in_place(layer.experts.switch_glu.gate_proj)
+ self.sharded_to_all_linear_in_place(layer.experts.switch_glu.down_proj)
+ self.all_to_sharded_linear_in_place(layer.experts.switch_glu.up_proj)
+ layer.experts = WrappedGemma4Experts(layer.experts) # pyright: ignore[reportAttributeAccessIssue,reportArgumentType]
+ layer.experts.sharding_group = self.group
+
+ mx.eval(layer)
+ if on_layer_loaded is not None:
+ on_layer_loaded(i, total)
+ return model
← 8d81811b Try harder to clean up processes nicely (#1889)
·
back to Exo
·
Fix Qwen3-VL and autodetect vision config (#1893) 2ecefa0c →