Add minimax tensor sharding (ml-explore#760)

* add minimax sharding

* fix

---------

Co-authored-by: Awni Hannun <awni@apple.com>

Author: Evanev7

mlx_lm/models/minimax.py

@@ -5,6 +5,7 @@
 
 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 .switch_layers import SwitchGLU
@@ -118,8 +119,12 @@ def __init__(self, args: ModelArgs):
             args.hidden_size, args.intermediate_size, args.num_local_experts
         )
         self.e_score_correction_bias = mx.zeros((args.num_local_experts,))
+        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)
+
         gates = self.gate(x.astype(mx.float32))
 
         scores = mx.sigmoid(gates)
@@ -135,6 +140,10 @@ def __call__(self, x: mx.array) -> mx.array:
 
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -267,6 +276,53 @@ def dequant(weight, scale_inv):
 
         return weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        rank = group.rank()
+        for layer in self.model.layers:
+            # Shard the self attention
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            if layer.self_attn.use_qk_norm:
+                layer.self_attn.q_norm.weight = layer.self_attn.q_norm.weight.split(
+                    N, axis=-1
+                )[rank]
+                layer.self_attn.k_norm.weight = layer.self_attn.k_norm.weight.split(
+                    N, axis=-1
+                )[rank]
+
+            layer.self_attn.num_attention_heads //= N
+            layer.self_attn.num_key_value_heads //= N
+
+            # Shard the MLP
+            shard_inplace(
+                layer.block_sparse_moe.switch_mlp.gate_proj,
+                "all-to-sharded",
+                group=group,
+            )
+            shard_inplace(
+                layer.block_sparse_moe.switch_mlp.down_proj,
+                "sharded-to-all",
+                group=group,
+            )
+            shard_inplace(
+                layer.block_sparse_moe.switch_mlp.up_proj,
+                "all-to-sharded",
+                group=group,
+            )
+            layer.block_sparse_moe.sharding_group = group
+
     @property
     def layers(self):
         return self.model.layers