From 6651d2e0bfcbd8e6044957bc38fc73b8df486bac Mon Sep 17 00:00:00 2001
From: Evan Quiney <evanev7@gmail.com>
Date: Sat, 17 Jan 2026 15:42:11 +0000
Subject: [PATCH] Add gpt-oss sharding (#761)

* add gpt-oss sharding

* fixes

---------

Co-authored-by: Awni Hannun <awni@apple.com>
---
 mlx_lm/models/gpt_oss.py | 54 +++++++++++++++++++++++++++++++++++++++-
 1 file changed, 53 insertions(+), 1 deletion(-)

diff --git a/mlx_lm/models/gpt_oss.py b/mlx_lm/models/gpt_oss.py
index 6bbdaf6..e0fda5e 100644
--- a/mlx_lm/models/gpt_oss.py
+++ b/mlx_lm/models/gpt_oss.py
@@ -7,6 +7,7 @@ from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
@@ -142,8 +143,12 @@ class MLPBlock(nn.Module):
             bias=True,
         )
         self.router = nn.Linear(config.hidden_size, config.num_local_experts, bias=True)
+        self.sharding_group = None
 
     def __call__(self, x: mx.array) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         g = self.router(x)
         experts, indices = mlx_topk(g, k=self.num_experts_per_tok, axis=-1)
         expert_weights = mx.softmax(experts, axis=-1, precise=True)
@@ -152,7 +157,13 @@ class MLPBlock(nn.Module):
         x = self.experts(x, indices)
 
         x = x * mx.expand_dims(expert_weights, axis=-1)
-        return x.sum(axis=-2)
+
+        y = x.sum(axis=-2)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
 
 
 class TransformerBlock(nn.Module):
@@ -268,6 +279,47 @@ class Model(nn.Module):
 
         return new_weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        R = group.rank()
+
+        for layer in self.model.layers:
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, sharding="all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, sharding="all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, sharding="all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, sharding="sharded-to-all", group=group
+            )
+
+            layer.self_attn.num_attention_heads //= N
+            layer.self_attn.num_key_value_heads //= N
+            layer.self_attn.num_key_value_groups = (
+                layer.self_attn.num_attention_heads
+                // layer.self_attn.num_key_value_heads
+            )
+
+            layer.self_attn.sinks = layer.self_attn.sinks[
+                layer.self_attn.num_attention_heads
+                * R : layer.self_attn.num_attention_heads
+                * (R + 1)
+            ]
+
+            shard_inplace(layer.mlp.experts.gate_proj, "all-to-sharded", group=group)
+            shard_inplace(layer.mlp.experts.down_proj, "sharded-to-all", group=group)
+            layer.mlp.experts.down_proj.bias /= N
+            shard_inplace(
+                layer.mlp.experts.up_proj, sharding="all-to-sharded", group=group
+            )
+
+            layer.mlp.sharding_group = group
+
     @property
     def layers(self):
         return self.model.layers