From df91f924d63e04f8cb61d8cbfa32ddb279c7cb12 Mon Sep 17 00:00:00 2001 From: Taylor Wallgren <5424153+twallgren@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:28:47 -0600 Subject: [PATCH 1/4] Convert DeepSeek-V3.2 to PipelineMixin; fix qwen3_moe.py sanitize() gate #1816 already fixed PipelineMixin's own layer-dropping bug (start_idx computed from this rank's own size rather than a real prefix sum, silently omitting a layer whenever ranks have different sizes) with a cleaner API than what this branch originally proposed for that part -- dropped here in favor of it. deepseek_v32.py wasn't touched by that fix because it never used PipelineMixin at all -- it carried its own fully independent copy of the same buggy logic, unrelated to and unfixed by #1816. Rather than patching that copy a third time, this converts DeepseekV32Model to actually inherit PipelineMixin (matching qwen3_5.py/ministral3.py), so it gets #1816's tested fix directly instead of maintaining a parallel implementation of the same bug-prone logic. Same behavior-preserving transformation applied elsewhere: __call__ walks pipeline_layers instead of manually indexing layers[start_idx + i] over a hand-tracked num_layers, and the outer Model.layers property returns model.pipeline_layers instead of re-deriving the same slice by hand. Also added the same empty-cache guards used elsewhere in this codebase (cache and cache[-1] is not None, etc.) for a rank that ends up with zero local layers. qwen3_moe.py's sanitize() has an unrelated bug in the same problem space, not touched by #1816: it gates MoE expert-weight-stacking on a fixed model.layers.0... key. Under pipeline sharding a rank only has its own local layers' weight files downloaded, so a rank not owning layer 0 sees that key absent and skips stacking entirely for its own local layers -- silently wrong for any raw (unconverted) checkpoint under pipeline parallelism. Fixed to detect and iterate by which layers are actually present in the local weights dict instead. --- mlx_lm/models/deepseek_v32.py | 37 ++++++++--------------------------- mlx_lm/models/qwen3_moe.py | 15 +++++++++++--- 2 files changed, 20 insertions(+), 32 deletions(-) diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index e8d22d3f0..67907d923 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/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,23 +435,22 @@ 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 + h, cache[0][0] if cache and 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: h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size) - if cache[-1] is not None: + if cache and cache[-1] is not None: cache[-1][0].keys = mx.depends(cache[-1][0].keys, h) # Broadcast h while keeping it in the graph @@ -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): diff --git a/mlx_lm/models/qwen3_moe.py b/mlx_lm/models/qwen3_moe.py index a9274da9a..565683d31 100644 --- a/mlx_lm/models/qwen3_moe.py +++ b/mlx_lm/models/qwen3_moe.py @@ -260,9 +260,18 @@ def __call__( def sanitize(self, weights): if self.args.tie_word_embeddings: weights.pop("lm_head.weight", None) - if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights: - return weights - for l in range(self.args.num_hidden_layers): + # Presence-based, not a fixed layer-0 probe: under pipeline + # sharding, a rank only downloads its own local layers' weight + # files, so a raw (unstacked-experts) checkpoint may have no + # layer-0 key at all even though its own local layers still need + # stacking — a layer-0-only gate would skip them entirely. + local_moe_layers = sorted( + int(k.split(".")[2]) + for k in weights + if k.startswith("model.layers.") + and k.endswith(".mlp.experts.0.up_proj.weight") + ) + for l in local_moe_layers: prefix = f"model.layers.{l}" for n in ["up_proj", "down_proj", "gate_proj"]: if f"{prefix}.mlp.experts.0.{n}.weight" in weights: From b8b3b94d274eed6d1993fb893602f5355b2bbed4 Mon Sep 17 00:00:00 2001 From: Taylor Wallgren <5424153+twallgren@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:05:30 -0600 Subject: [PATCH 2/4] Add deepseek_v32 pipeline layer-assignment test coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Contributed by @azamamirza in PR review: extends the existing qwen3_moe uneven-assignment regression test to also cover deepseek_v32 now that it uses the shared PipelineMixin. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- tests/model_parallel_tests.py | 68 +++++++++++++++++++++++++++++------ 1 file changed, 58 insertions(+), 10 deletions(-) diff --git a/tests/model_parallel_tests.py b/tests/model_parallel_tests.py index 3a325aba5..577511b2e 100644 --- a/tests/model_parallel_tests.py +++ b/tests/model_parallel_tests.py @@ -5,7 +5,7 @@ import mlx.core as mx -from mlx_lm.models import qwen3_moe +from mlx_lm.models import deepseek_v32, qwen3_moe from mlx_lm.models.pipeline import PipelineMixin @@ -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: @@ -246,15 +269,40 @@ def test_pipeline_uneven_model(self): } 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__": From 9d0a1c7c97f97a90322b1c18521e603d6f2ab939 Mon Sep 17 00:00:00 2001 From: Taylor Wallgren <5424153+twallgren@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:44:23 -0600 Subject: [PATCH 3/4] Address review: drop redundant cache guards and duplicate test coverage cache and cache[0]/cache[-1] guards in DeepseekV32Model.__call__ are unnecessary and inconsistent with qwen3_moe.py's PipelineMixin usage. test_pipeline_uneven_model's deepseek_v32 subtest duplicated coverage already provided by the qwen3_moe case now that both models delegate to the same PipelineMixin.pipeline() implementation. --- mlx_lm/models/deepseek_v32.py | 4 ++-- tests/model_parallel_tests.py | 45 ++++++++--------------------------- 2 files changed, 12 insertions(+), 37 deletions(-) diff --git a/mlx_lm/models/deepseek_v32.py b/mlx_lm/models/deepseek_v32.py index 67907d923..0dba68765 100644 --- a/mlx_lm/models/deepseek_v32.py +++ b/mlx_lm/models/deepseek_v32.py @@ -437,7 +437,7 @@ def __call__( if cache is None: cache = [None] * len(self.pipeline_layers) mask = create_attention_mask( - h, cache[0][0] if cache and cache[0] else None, return_array=True + h, cache[0][0] if cache[0] else None, return_array=True ) # Receive from the previous process in the pipeline @@ -450,7 +450,7 @@ def __call__( # 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: + if cache[-1] is not None: cache[-1][0].keys = mx.depends(cache[-1][0].keys, h) # Broadcast h while keeping it in the graph diff --git a/tests/model_parallel_tests.py b/tests/model_parallel_tests.py index 577511b2e..ab022e927 100644 --- a/tests/model_parallel_tests.py +++ b/tests/model_parallel_tests.py @@ -5,7 +5,7 @@ import mlx.core as mx -from mlx_lm.models import deepseek_v32, qwen3_moe +from mlx_lm.models import qwen3_moe from mlx_lm.models.pipeline import PipelineMixin @@ -269,40 +269,15 @@ def test_pipeline_uneven_model(self): } size = 2 - 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))) + 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))) if __name__ == "__main__": From 886c3e5d91cabbb73eff3f4883f55d3c0d22ed53 Mon Sep 17 00:00:00 2001 From: Taylor Wallgren <5424153+twallgren@users.noreply.github.com> Date: Thu, 10 Sep 2026 11:52:24 -0600 Subject: [PATCH 4/4] Split off qwen3_moe.py sanitize() fix into a separate PR Per review feedback on #1861: this PR now only converts DeepseekV32Model to PipelineMixin. The qwen3_moe.py sanitize() gate fix moved to its own branch/PR since it's an independent bug in a different model. --- mlx_lm/models/qwen3_moe.py | 15 +++------------ 1 file changed, 3 insertions(+), 12 deletions(-) diff --git a/mlx_lm/models/qwen3_moe.py b/mlx_lm/models/qwen3_moe.py index 565683d31..a9274da9a 100644 --- a/mlx_lm/models/qwen3_moe.py +++ b/mlx_lm/models/qwen3_moe.py @@ -260,18 +260,9 @@ def __call__( def sanitize(self, weights): if self.args.tie_word_embeddings: weights.pop("lm_head.weight", None) - # Presence-based, not a fixed layer-0 probe: under pipeline - # sharding, a rank only downloads its own local layers' weight - # files, so a raw (unstacked-experts) checkpoint may have no - # layer-0 key at all even though its own local layers still need - # stacking — a layer-0-only gate would skip them entirely. - local_moe_layers = sorted( - int(k.split(".")[2]) - for k in weights - if k.startswith("model.layers.") - and k.endswith(".mlp.experts.0.up_proj.weight") - ) - for l in local_moe_layers: + if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights: + return weights + for l in range(self.args.num_hidden_layers): prefix = f"model.layers.{l}" for n in ["up_proj", "down_proj", "gate_proj"]: if f"{prefix}.mlp.experts.0.{n}.weight" in weights: