Update DeepSeek-V3.2 to use PipelineMixin

[twallgren]

• ☑️ I understand it is strictly prohibited to use AI to write PR description
• AI usage disclosure: AI was used to write the fix and the tests with some guidance from me. It went through many rounds of code reviews with a wide variety of agent frameworks and models to check for any issues with the core logic, styling, edge cases, etc.

The issue this is addressing:
deepseek_v32.py had its own independent pipeline implementation, duplicating logic that already existed in the shared PipelineMixin base class but not benefiting from bug fixes.

The solution:
Converted DeepseekV32Model to inherit PipelineMixin like every other pipeline-capable model so it benefits from the shared, tested pipeline logic (including uneven-split support) instead of maintaining another copy of similar code.

Note: an earlier version of this PR also fixed a layer-drop bug directly in pipeline.py's core split calculation, but that was independently fixed and merged upstream in #1816 first so that part was dropped from this PR to avoid duplicating it.

I needed these fixes to get DeepSeek-V3.2-family models running correctly in pipeline mode on my 3-node cluster.

[azamamirza]

Ran into the deepseek_v32 side of this yesterday while working on #1816: it has the same silently-dropped-layers bug that PR fixed in the mixin (7 layers on 2 ranks ran layer 3 on no rank), and @michalk8 suggested exactly this subclassing approach as a follow-up in the #1816 discussion. You got there first, and this version is cleaner than what I had staged: dropping the override entirely also picks up the new split= argument for free.

Verified locally on Apple Silicon (M5 Max):

1. Applied my regression tests from the Fix PipelineMixin dropping layers on uneven splits, add explicit per-rank splits #1816 follow-up work to this branch unmodified: all pass, including a 7-layer / 2-rank uneven assignment check on the real DeepseekV32Model (fails on current main, passes here) and a tiny-config forward-equality test through test_pipeline with a size-1 group.
2. tests/model_parallel_tests.py: 4 passed + 8 subtests. tests/test_models.py -k "deepseek or qwen3": 6 passed.
3. The sanitize() change reads correct to me but I did not execute a pipeline-sharded checkpoint load to verify that half.
4. We hit the original issue testing on a 2 Mac ring of uneven hardware (M4 Max and M5 Max machines with differing specs).

This PR has no committed tests; the diff below adds the deepseek_v32 coverage and applies cleanly to this branch if you want to fold it in.

regression tests (applies to this branch)

diff --git a/tests/model_parallel_tests.py b/tests/model_parallel_tests.py
index 1d9bb89..af73c83 100644
--- a/tests/model_parallel_tests.py
+++ b/tests/model_parallel_tests.py
@@ -6,7 +6,7 @@ import unittest
 import mlx.core as mx
 
 import mlx_lm
-from mlx_lm.models import qwen3_moe
+from mlx_lm.models import deepseek_v32, qwen3_moe
 from mlx_lm.models.pipeline import PipelineMixin
 
 
@@ -170,6 +170,29 @@ class TestModelParallel(unittest.TestCase):
                 "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:
@@ -247,15 +270,40 @@ class TestModelParallel(unittest.TestCase):
         }
 
         size = 2
-        args = qwen3_moe.ModelArgs.from_dict(config)
-        assigned = []
-        for rank in range(size):
-            model = qwen3_moe.Model(args)
-            model.model.pipeline(Group(rank, size))
-            assigned.extend(
-                i for i, l in enumerate(model.model.layers) if l is not None
-            )
-        self.assertEqual(sorted(assigned), list(range(7)))
+        dsv32_config = {
+            "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": 7,
+            "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,
+        }
+        for arch, arch_config in ((qwen3_moe, config), (deepseek_v32, dsv32_config)):
+            with self.subTest(model_type=arch_config["model_type"]):
+                args = arch.ModelArgs.from_dict(arch_config)
+                assigned = []
+                for rank in range(size):
+                    model = arch.Model(args)
+                    model.model.pipeline(Group(rank, size))
+                    assigned.extend(
+                        i for i, l in enumerate(model.model.layers) if l is not None
+                    )
+                self.assertEqual(sorted(assigned), list(range(7)))
 
 
 if __name__ == "__main__":

[twallgren]

Thanks @azamamirza for the review, feedback, testing, and regression tests! I applied your diff and confirmed that the tests pass. I'll remove the draft status.
Since you're working on the uneven split topic, you might be interested in #1862 and #1863 as well.

[michalk8]

Can you please split this PR into two, one for DeepSeek to use PipelineMixin (happy to merge) and the Qwen3 MoE in a separate one (I'm not sure about this one yet, if it's indeed a bug, it might be present in other models).

[twallgren]

Thanks for the review and feedback, @michalk8! I've split the PR into two as you requested. The new PR is #1875.

[michalk8]

LGTM, thanks!