Update DeepSeek-V3.2 to use PipelineMixin

mlx_lm/models/deepseek_v32.py

@@ -12,6 +12,7 @@
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import CacheList, KVCache
 from .mla import MultiLinear
+from .pipeline import PipelineMixin
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
@@ -412,7 +413,7 @@ def __call__(
         return h + r
 
 
-class DeepseekV32Model(nn.Module):
+class DeepseekV32Model(PipelineMixin, nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
         self.vocab_size = config.vocab_size
@@ -421,28 +422,7 @@ def __init__(self, config: ModelArgs):
             DeepseekV32DecoderLayer(config, idx)
             for idx in range(config.num_hidden_layers)
         ]
-        self.start_idx = 0
-        self.end_idx = len(self.layers)
-        self.num_layers = self.end_idx
-
         self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.pipeline_rank = 0
-        self.pipeline_size = 1
-
-    def pipeline(self, group):
-        # Split layers in reverse so rank=0 gets the last layers and
-        # rank=pipeline_size-1 gets the first
-        self.pipeline_rank = group.rank()
-        self.pipeline_size = group.size()
-        layers_per_rank = len(self.layers) // self.pipeline_size
-        extra = len(self.layers) - layers_per_rank * self.pipeline_size
-        if self.pipeline_rank < extra:
-            layers_per_rank += 1
-        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
-        self.end_idx = self.start_idx + layers_per_rank
-        self.layers = self.layers[: self.end_idx]
-        self.layers[: self.start_idx] = [None] * self.start_idx
-        self.num_layers = len(self.layers) - self.start_idx
 
     def __call__(
         self,
@@ -455,18 +435,17 @@ def __call__(
         pipeline_size = self.pipeline_size
 
         if cache is None:
-            cache = [None] * self.num_layers
+            cache = [None] * len(self.pipeline_layers)
         mask = create_attention_mask(
             h, cache[0][0] if cache[0] else None, return_array=True
         )
 
         # Receive from the previous process in the pipeline
-
         if pipeline_rank < pipeline_size - 1:
             h = mx.distributed.recv_like(h, (pipeline_rank + 1))
 
-        for i in range(self.num_layers):
-            h = self.layers[self.start_idx + i](h, mask, cache[i])
+        for layer, c in zip(self.pipeline_layers, cache):
+            h = layer(h, mask, c)
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
@@ -646,7 +625,7 @@ def shard_heads(w):
 
     @property
     def layers(self):
-        return self.model.layers[self.model.start_idx : self.model.end_idx]
+        return self.model.pipeline_layers
 
     @property
     def cast_predicate(self):

tests/model_parallel_tests.py

@@ -169,6 +169,29 @@ def test_pipeline(self):
                 "max_position_embeddings": 256,
                 "tie_word_embeddings": False,
             },
+            {
+                "model_type": "deepseek_v32",
+                "vocab_size": 128,
+                "hidden_size": 64,
+                "index_head_dim": 16,
+                "index_n_heads": 2,
+                "index_topk": 4,
+                "intermediate_size": 128,
+                "moe_intermediate_size": 32,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 4,
+                "n_shared_experts": 1,
+                "n_routed_experts": 4,
+                "kv_lora_rank": 8,
+                "q_lora_rank": 8,
+                "qk_rope_head_dim": 8,
+                "v_head_dim": 16,
+                "qk_nope_head_dim": 16,
+                "num_experts_per_tok": 2,
+                "first_k_dense_replace": 1,
+                "max_position_embeddings": 256,
+            },
         ]
         mx.random.seed(0)
         for config in test_configs: