fix isinstance for qwen3Moe (#1595)

we were checking if qwen3 had a transformers Qwen3DecoderLayer rather
than an mlx Qwen3MoeDecoderLayer causing an assertion error on loading
qwen models - this corrects it to the actual layer type
This commit is contained in:
Evan Quiney
2026-02-23 15:39:13 +00:00
committed by GitHub
parent fbb80e1cc9
commit 0ff99a2c40
+2 -3
View File
@@ -32,13 +32,12 @@ from mlx_lm.models.minimax import MiniMaxAttention
from mlx_lm.models.minimax import Model as MiniMaxModel
from mlx_lm.models.ministral3 import Model as Ministral3Model
from mlx_lm.models.qwen3_moe import Model as Qwen3MoeModel
from mlx_lm.models.qwen3_moe import Qwen3MoeSparseMoeBlock
from mlx_lm.models.qwen3_moe import Qwen3MoeDecoderLayer, Qwen3MoeSparseMoeBlock
from mlx_lm.models.qwen3_next import Model as Qwen3NextModel
from mlx_lm.models.qwen3_next import Qwen3NextDecoderLayer, Qwen3NextSparseMoeBlock
from mlx_lm.models.step3p5 import Model as Step35Model
from mlx_lm.models.step3p5 import Step3p5MLP as Step35MLP
from mlx_lm.models.step3p5 import Step3p5Model as Step35InnerModel
from transformers.models.qwen3.modeling_qwen3 import Qwen3DecoderLayer
from exo.shared.logging import logger
from exo.shared.types.worker.shards import PipelineShardMetadata
@@ -840,7 +839,7 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
for i, layer in enumerate(model.layers):
eval_with_timeout(layer.parameters(), timeout_seconds / total, on_timeout)
# Shard the self attention
if isinstance(layer, Qwen3DecoderLayer):
if isinstance(layer, Qwen3MoeDecoderLayer):
layer.self_attn.q_proj = self.all_to_sharded_linear(
layer.self_attn.q_proj
)