Add pipeline() support to Qwen2 and Qwen3-Next

[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:
Pipeline parallelism wasn't supported for Qwen2 and Qwen3-Next model architectures. Attempting to use pipeline mode would result in it hanging after a "not supported" error happened on a background thread.

The solution:
Added support for pipeline parallelism for Qwen2 and Qwen3-Next model architectures. It follows the same pattern used for other models with similar structures.

I needed this change so I could properly use pipeline processing on my 3 node cluster with Qwen2 and Qwen3-Next models. I've been running it successfully for a while with great results.

[azamamirza]

Same offer as #1861: this PR adds the pipeline paths but no tests, so I wrote coverage and verified it on this branch locally (Apple Silicon, M5 Max):

1. Added qwen2 and qwen3_next tiny configs to test_pipeline: real forward-pass equality against the unpipelined model under a size-1 group passes for both.
2. Added both to the uneven layer-assignment check (7 layers on 2 ranks): every layer runs on exactly one rank for qwen3_moe, qwen2 and qwen3_next.
3. The hybrid rescan behaves correctly at the edge I was most worried about: with full_attention_interval=4 and 7 layers on 2 ranks, the last rank's slice (layers 0-2) is all linear layers and fa_idx correctly stays None, while rank 0 gets its full-attention layer at local index 0. Both asserted.
4. tests/model_parallel_tests.py on this branch with the additions: 4 passed + 10 subtests. tests/test_models.py -k "qwen2 or qwen3" - 5 passed.

One honest limitation: the fa_idx=None forward path only executes under real multi-rank communication, which a unit test can't reach. The assertions cover the index rescan, and the size-1 forward covers the combined path.

Diff below applies cleanly to this branch if you want to fold it in like the #1861 one.

regression tests (applies to this branch)

diff --git a/tests/model_parallel_tests.py b/tests/model_parallel_tests.py
index 1d9bb89..cb9c621 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 qwen2, qwen3_moe, qwen3_next
 from mlx_lm.models.pipeline import PipelineMixin
 
 
@@ -170,6 +170,45 @@ class TestModelParallel(unittest.TestCase):
                 "max_position_embeddings": 256,
                 "tie_word_embeddings": False,
             },
+            {
+                "model_type": "qwen2",
+                "vocab_size": 128,
+                "hidden_size": 64,
+                "intermediate_size": 128,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "rms_norm_eps": 1e-5,
+                "rope_theta": 10000.0,
+                "max_position_embeddings": 256,
+                "tie_word_embeddings": False,
+            },
+            {
+                "model_type": "qwen3_next",
+                "vocab_size": 128,
+                "hidden_size": 64,
+                "num_hidden_layers": 4,
+                "intermediate_size": 128,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "head_dim": 16,
+                "linear_num_value_heads": 4,
+                "linear_num_key_heads": 2,
+                "linear_key_head_dim": 16,
+                "linear_value_head_dim": 16,
+                "linear_conv_kernel_dim": 4,
+                "num_experts": 4,
+                "num_experts_per_tok": 2,
+                "decoder_sparse_step": 1,
+                "shared_expert_intermediate_size": 32,
+                "mlp_only_layers": [0],
+                "moe_intermediate_size": 32,
+                "rms_norm_eps": 1e-5,
+                "rope_theta": 10000.0,
+                "partial_rotary_factor": 0.25,
+                "max_position_embeddings": 256,
+                "full_attention_interval": 4,
+            },
         ]
         mx.random.seed(0)
         for config in test_configs:
@@ -247,15 +286,74 @@ 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)))
+        qwen2_config = {
+            "model_type": "qwen2",
+            "vocab_size": 128,
+            "hidden_size": 64,
+            "intermediate_size": 128,
+            "num_hidden_layers": 7,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "rms_norm_eps": 1e-5,
+            "rope_theta": 10000.0,
+            "max_position_embeddings": 256,
+            "tie_word_embeddings": False,
+        }
+        qwen3_next_config = {
+            "model_type": "qwen3_next",
+            "vocab_size": 128,
+            "hidden_size": 64,
+            "num_hidden_layers": 7,
+            "intermediate_size": 128,
+            "num_attention_heads": 4,
+            "num_key_value_heads": 2,
+            "head_dim": 16,
+            "linear_num_value_heads": 4,
+            "linear_num_key_heads": 2,
+            "linear_key_head_dim": 16,
+            "linear_value_head_dim": 16,
+            "linear_conv_kernel_dim": 4,
+            "num_experts": 4,
+            "num_experts_per_tok": 2,
+            "decoder_sparse_step": 1,
+            "shared_expert_intermediate_size": 32,
+            "mlp_only_layers": [0],
+            "moe_intermediate_size": 32,
+            "rms_norm_eps": 1e-5,
+            "rope_theta": 10000.0,
+            "partial_rotary_factor": 0.25,
+            "max_position_embeddings": 256,
+            "full_attention_interval": 4,
+        }
+        for arch, arch_config in (
+            (qwen3_moe, config),
+            (qwen2, qwen2_config),
+            (qwen3_next, qwen3_next_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)))
+
+        # Hybrid edge case: 7 layers on 2 ranks puts layers 0-2 (all linear,
+        # full_attention_interval=4) on the last rank; its rescan must leave
+        # fa_idx as None rather than pointing at a wrong layer. Rank 0 holds
+        # layers 3-6, so its full-attention layer sits at local index 0.
+        args = qwen3_next.ModelArgs.from_dict(qwen3_next_config)
+        model = qwen3_next.Model(args)
+        model.model.pipeline(Group(1, size))
+        self.assertIsNone(model.model.fa_idx)
+        self.assertEqual(model.model.ssm_idx, 0)
+        model = qwen3_next.Model(args)
+        model.model.pipeline(Group(0, size))
+        self.assertEqual(model.model.fa_idx, 0)
+        self.assertEqual(model.model.ssm_idx, 1)
 
 
 if __name__ == "__main__":

[twallgren]

Thanks again @azamamirza and apologies for the repeat issue with missing tests. I know I had tests in my local changes but they must have gotten dropped by accident during some refactoring. I've applied your diff and confirmed the tests pass.

[michalk8]

LGTM, thanks!