Support uneven tensor-parallel sharding in shard() (GQA/MLA-aware)

mlx_lm/models/deepseek_v2.py

@@ -10,7 +10,7 @@
 
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .pipeline import PipelineMixin
+from .pipeline import PipelineMixin, _rank_sizes
 from .switch_layers import SwitchGLU
 
 
@@ -448,23 +448,43 @@ def sanitize(self, 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:
+            attn = layer.self_attn
+            # MLA has no separate KV-head count — every head reconstructs
+            # its own K/V from a shared low-rank latent via kv_b_proj, so
+            # (unlike GQA) heads can be distributed individually rather
+            # than in fixed-size groups.
+            head_sizes = _rank_sizes(attn.num_heads, N)
+            if any(s == 0 for s in head_sizes):
+                zero_ranks = [i for i, s in enumerate(head_sizes) if s == 0]
+                raise ValueError(
+                    f"Cannot shard {self.args.model_type}'s {attn.num_heads} attention "
+                    f"head(s) across {N} ranks: rank(s) {zero_ranks} would get zero "
+                    "heads."
+                )
+            q_sizes = [s * attn.q_head_dim for s in head_sizes]
+            kv_out_head_dim = attn.q_head_dim - attn.qk_rope_head_dim + attn.v_head_dim
+            kv_sizes = [s * kv_out_head_dim for s in head_sizes]
+            o_sizes = [s * attn.v_head_dim for s in head_sizes]
+
             # Shard the self attention
-            if layer.self_attn.q_lora_rank is None:
-                layer.self_attn.q_proj = shard_linear(
-                    layer.self_attn.q_proj, "all-to-sharded", group=group
+            if attn.q_lora_rank is None:
+                attn.q_proj = shard_linear(
+                    attn.q_proj, "all-to-sharded", group=group, sizes=q_sizes
                 )
             else:
-                layer.self_attn.q_b_proj = shard_linear(
-                    layer.self_attn.q_b_proj, "all-to-sharded", group=group
+                attn.q_b_proj = shard_linear(
+                    attn.q_b_proj, "all-to-sharded", group=group, sizes=q_sizes
                 )
-            layer.self_attn.kv_b_proj = shard_linear(
-                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
+            attn.kv_b_proj = shard_linear(
+                attn.kv_b_proj, "all-to-sharded", group=group, sizes=kv_sizes
             )
-            layer.self_attn.o_proj = shard_linear(
-                layer.self_attn.o_proj, "sharded-to-all", group=group
+            attn.o_proj = shard_linear(
+                attn.o_proj, "sharded-to-all", group=group, sizes=o_sizes
             )
-            layer.self_attn.num_heads //= N
+            attn.num_heads = head_sizes[rank]
 
             # Shard the MLP
             if isinstance(layer.mlp, DeepseekV2MLP):

mlx_lm/models/deepseek_v3.py

@@ -11,7 +11,7 @@
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .mla import MultiLinear
-from .pipeline import PipelineMixin
+from .pipeline import PipelineMixin, _rank_sizes
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
@@ -483,30 +483,51 @@ 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:
+            attn = layer.self_attn
+            # MLA has no separate KV-head count — every head reconstructs
+            # its own K/V via the batched embed_q/unembed_out projections,
+            # so (unlike GQA) heads can be distributed individually rather
+            # than in fixed-size groups.
+            head_sizes = _rank_sizes(attn.num_heads, N)
+            if any(s == 0 for s in head_sizes):
+                zero_ranks = [i for i, s in enumerate(head_sizes) if s == 0]
+                raise ValueError(
+                    f"Cannot shard {self.args.model_type}'s {attn.num_heads} attention "
+                    f"head(s) across {N} ranks: rank(s) {zero_ranks} would get zero "
+                    "heads."
+                )
+            q_sizes = [s * attn.q_head_dim for s in head_sizes]
+
             # Shard the self attention
-            if layer.self_attn.q_lora_rank is None:
-                layer.self_attn.q_proj = shard_linear(
-                    layer.self_attn.q_proj, "all-to-sharded", group=group
+            if attn.q_lora_rank is None:
+                attn.q_proj = shard_linear(
+                    attn.q_proj, "all-to-sharded", group=group, sizes=q_sizes
                 )
             else:
-                layer.self_attn.q_b_proj = shard_linear(
-                    layer.self_attn.q_b_proj, "all-to-sharded", group=group
+                attn.q_b_proj = shard_linear(
+                    attn.q_b_proj, "all-to-sharded", group=group, sizes=q_sizes
                 )
-            layer.self_attn.num_heads //= N
-            num_heads = layer.self_attn.num_heads
-            sh = rank * num_heads
-            eh = sh + num_heads
+
+            # embed_q/unembed_out are batched per-head (MultiLinear), sliced
+            # directly along their head axis rather than through
+            # shard_linear — boundaries follow the same cumulative,
+            # possibly-uneven head_sizes as q_proj/o_proj above.
+            sh = sum(head_sizes[:rank])
+            eh = sh + head_sizes[rank]
 
             def shard_heads(w):
                 return w[sh:eh]
 
-            layer.self_attn.embed_q.apply(shard_heads)
-            layer.self_attn.unembed_out.apply(shard_heads)
+            attn.embed_q.apply(shard_heads)
+            attn.unembed_out.apply(shard_heads)
 
-            layer.self_attn.o_proj = shard_linear(
-                layer.self_attn.o_proj, "sharded-to-all", group=group
+            o_sizes = [s * attn.v_head_dim for s in head_sizes]
+            attn.o_proj = shard_linear(
+                attn.o_proj, "sharded-to-all", group=group, sizes=o_sizes
             )
+            attn.num_heads = head_sizes[rank]
 
             # Shard the MLP
             if isinstance(layer.mlp, DeepseekV3MLP):

mlx_lm/models/glm4_moe.py

@@ -9,7 +9,7 @@
 
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .pipeline import PipelineMixin
+from .pipeline import PipelineMixin, _rank_sizes
 from .switch_layers import SwitchGLU
 
 
@@ -341,22 +341,54 @@ def sanitize(self, 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
+            attn = layer.self_attn
+            n_heads = attn.n_heads
+            n_kv_heads = attn.n_kv_heads
+            ratio = n_heads // n_kv_heads
+            # q_proj's pre-shard output rows are always n_heads * head_dim,
+            # for both plain and quantized linears (quantization only packs
+            # the input/column axis, not the output/row axis).
+            head_dim = attn.q_proj.weight.shape[0] // n_heads
+
+            # Distribute whole KV-head *groups* (each covering `ratio`
+            # query heads) across ranks, rather than distributing q_proj's
+            # and k_proj's/v_proj's output features independently. This
+            # keeps n_heads_local an exact multiple of n_kv_heads_local on
+            # every rank even when n_heads/n_kv_heads/N don't divide evenly
+            # among each other, so grouped-query-attention's local repeat
+            # of KV heads stays correct after uneven sharding.
+            kv_head_sizes = _rank_sizes(n_kv_heads, N)
+            if any(s == 0 for s in kv_head_sizes):
+                zero_ranks = [i for i, s in enumerate(kv_head_sizes) if s == 0]
+                raise ValueError(
+                    f"Cannot shard {self.args.model_type}'s {n_kv_heads} KV head(s) "
+                    f"across {N} ranks: rank(s) {zero_ranks} would get zero heads."
+                )
+            q_head_sizes = [s * ratio for s in kv_head_sizes]
+            q_sizes = [s * head_dim for s in q_head_sizes]
+            kv_sizes = [s * head_dim for s in kv_head_sizes]
+
+            # Shard the self attention. o_proj's input sizes must match
+            # q_proj's output sizes exactly (same per-rank partition of the
+            # n_heads * head_dim dimension) so the local contraction before
+            # all_sum is consistent.
+            attn.q_proj = shard_linear(
+                attn.q_proj, "all-to-sharded", group=group, sizes=q_sizes
             )
-            layer.self_attn.k_proj = shard_linear(
-                layer.self_attn.k_proj, "all-to-sharded", group=group
+            attn.k_proj = shard_linear(
+                attn.k_proj, "all-to-sharded", group=group, sizes=kv_sizes
             )
-            layer.self_attn.v_proj = shard_linear(
-                layer.self_attn.v_proj, "all-to-sharded", group=group
+            attn.v_proj = shard_linear(
+                attn.v_proj, "all-to-sharded", group=group, sizes=kv_sizes
             )
-            layer.self_attn.o_proj = shard_linear(
-                layer.self_attn.o_proj, "sharded-to-all", group=group
+            attn.o_proj = shard_linear(
+                attn.o_proj, "sharded-to-all", group=group, sizes=q_sizes
             )
-            layer.self_attn.n_heads //= N
-            layer.self_attn.n_kv_heads //= N
+            attn.n_heads = q_head_sizes[rank]
+            attn.n_kv_heads = kv_head_sizes[rank]
 
             # Shard the MLP
             if isinstance(layer.mlp, MLP):

mlx_lm/models/llama.py

@@ -10,6 +10,7 @@
 from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
+from .pipeline import PipelineMixin, _rank_sizes
 from .rope_utils import initialize_rope
 
 
@@ -148,7 +149,7 @@ def __call__(
         return out
 
 
-class LlamaModel(nn.Module):
+class LlamaModel(PipelineMixin, nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
         self.args = args
@@ -181,18 +182,44 @@ def __call__(
         else:
             h = self.embed_tokens(inputs)
 
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
         if cache is None:
-            cache = [None] * len(self.layers)
+            cache = [None] * len(self.pipeline_layers)
+
+        # Build the full-attention / sliding-window masks from whichever
+        # layers this rank actually owns (fa_idx/swa_idx are computed over
+        # the full, unsharded layer list at __init__ time and are no longer
+        # valid once pipeline() has truncated self.layers).
+        fa_mask = None
+        swa_mask = None
+        for c, layer in zip(cache, self.pipeline_layers):
+            if layer.use_sliding:
+                if swa_mask is None:
+                    swa_mask = create_attention_mask(
+                        h, c, window_size=self.sliding_window
+                    )
+            elif fa_mask is None:
+                fa_mask = create_attention_mask(h, c)
+
+        # Receive from the previous process in the pipeline
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+        for layer, c in zip(self.pipeline_layers, cache):
+            mask = swa_mask if layer.use_sliding else fa_mask
+            h = layer(h, mask, cache=c)
 
-        fa_mask = create_attention_mask(h, cache[self.fa_idx])
-        if self.swa_idx is not None:
-            swa_mask = create_attention_mask(
-                h, cache[self.swa_idx], window_size=self.sliding_window
-            )
+        # Send to the next process in the pipeline
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache and cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
 
-        for layer, cache in zip(self.layers, cache):
-            mask = swa_mask if layer.use_sliding else fa_mask
-            h = layer(h, mask, cache=cache)
+        # Broadcast h while keeping it in the graph
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -231,24 +258,59 @@ def sanitize(self, 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
+            attn = layer.self_attn
+            n_heads = attn.n_heads
+            n_kv_heads = attn.n_kv_heads
+            ratio = n_heads // n_kv_heads
+            # q_proj's pre-shard output rows are always n_heads * head_dim,
+            # for both plain and quantized linears (quantization only packs
+            # the input/column axis, not the output/row axis).
+            head_dim = attn.q_proj.weight.shape[0] // n_heads
+
+            # Distribute whole KV-head *groups* (each covering `ratio`
+            # query heads) across ranks, rather than distributing q_proj's
+            # and k_proj's/v_proj's output features independently. This
+            # keeps n_heads_local an exact multiple of n_kv_heads_local on
+            # every rank even when n_heads/n_kv_heads/N don't divide evenly
+            # among each other, so grouped-query-attention's local repeat
+            # of KV heads stays correct after uneven sharding.
+            kv_head_sizes = _rank_sizes(n_kv_heads, N)
+            if any(s == 0 for s in kv_head_sizes):
+                zero_ranks = [i for i, s in enumerate(kv_head_sizes) if s == 0]
+                raise ValueError(
+                    f"Cannot shard {self.args.model_type}'s {n_kv_heads} KV head(s) "
+                    f"across {N} ranks: rank(s) {zero_ranks} would get zero heads."
+                )
+            q_head_sizes = [s * ratio for s in kv_head_sizes]
+            q_sizes = [s * head_dim for s in q_head_sizes]
+            kv_sizes = [s * head_dim for s in kv_head_sizes]
+
+            # Shard the self attention. o_proj's input sizes must match
+            # q_proj's output sizes exactly (same per-rank partition of the
+            # n_heads * head_dim dimension) so the local contraction before
+            # all_sum is consistent.
+            attn.q_proj = shard_linear(
+                attn.q_proj, "all-to-sharded", group=group, sizes=q_sizes
             )
-            layer.self_attn.k_proj = shard_linear(
-                layer.self_attn.k_proj, "all-to-sharded", group=group
+            attn.k_proj = shard_linear(
+                attn.k_proj, "all-to-sharded", group=group, sizes=kv_sizes
             )
-            layer.self_attn.v_proj = shard_linear(
-                layer.self_attn.v_proj, "all-to-sharded", group=group
+            attn.v_proj = shard_linear(
+                attn.v_proj, "all-to-sharded", group=group, sizes=kv_sizes
             )
-            layer.self_attn.o_proj = shard_linear(
-                layer.self_attn.o_proj, "sharded-to-all", group=group
+            attn.o_proj = shard_linear(
+                attn.o_proj, "sharded-to-all", group=group, sizes=q_sizes
             )
-            layer.self_attn.n_heads //= N
-            layer.self_attn.n_kv_heads //= N
+            attn.n_heads = q_head_sizes[rank]
+            attn.n_kv_heads = kv_head_sizes[rank]
 
-            # Shard the MLP
+            # Shard the MLP. No head-alignment constraint here, so the
+            # generic remainder-aware (group_size-blocked for quantized
+            # layers) automatic split in shard_linear/distributed.py is
+            # used directly.
             layer.mlp.gate_proj = shard_linear(
                 layer.mlp.gate_proj, "all-to-sharded", group=group
             )
@@ -261,7 +323,7 @@ def shard(self, group: Optional[mx.distributed.Group] = None):
 
     @property
     def layers(self):
-        return self.model.layers
+        return self.model.pipeline_layers
 
     def make_cache(self):
         return [

mlx_lm/models/pipeline.py

@@ -1,6 +1,17 @@
 # Copyright © 2025 Apple Inc.
 
 
+def _rank_sizes(dim, N, block=1):
+    # Split dim into N ranks as evenly as possible: the first `extra`
+    # ranks get one extra unit (or block, for group_size-aware quantized
+    # splits), the rest get the base amount. Reduces to an exactly even
+    # split whenever dim % (N * block) == 0.
+    n_blocks = dim // block
+    base = n_blocks // N
+    extra = n_blocks - base * N
+    return [(base + (1 if i < extra else 0)) * block for i in range(N)]
+
+
 class PipelineMixin:
     def __init__(self):
         super().__init__()