From 834fac934c4e04de9b3d723e2b9287a2c60cfd4a Mon Sep 17 00:00:00 2001 From: Awni Hannun Date: Tue, 24 Feb 2026 17:04:43 -0800 Subject: [PATCH] fix qwen3.5 sanitize (#928) --- mlx_lm/models/qwen3_5.py | 6 ++---- mlx_lm/models/qwen3_5_moe.py | 7 +------ 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/mlx_lm/models/qwen3_5.py b/mlx_lm/models/qwen3_5.py index 70d38f4..09dac0c 100644 --- a/mlx_lm/models/qwen3_5.py +++ b/mlx_lm/models/qwen3_5.py @@ -5,7 +5,6 @@ from typing import Any, Dict, List, Optional, Union import mlx.core as mx import mlx.nn as nn -from mlx.utils import tree_flatten, tree_unflatten from .base import ( BaseModelArgs, @@ -364,11 +363,10 @@ class Model(nn.Module): ) def sanitize(self, weights): - weights = tree_unflatten(list(weights.items())) - weights = dict(tree_flatten(weights)) - sanitized = {} for key, value in weights.items(): + if key.startswith("vision_tower") or key.startswith("model.visual"): + continue if key.startswith("model.visual"): continue if key.startswith("model.language_model"): diff --git a/mlx_lm/models/qwen3_5_moe.py b/mlx_lm/models/qwen3_5_moe.py index c94b3fb..53ab853 100644 --- a/mlx_lm/models/qwen3_5_moe.py +++ b/mlx_lm/models/qwen3_5_moe.py @@ -2,8 +2,6 @@ from dataclasses import dataclass -from mlx.utils import tree_flatten, tree_unflatten - from .base import BaseModelArgs from .qwen3_5 import Model as Qwen3_5Model @@ -23,12 +21,9 @@ class ModelArgs(BaseModelArgs): class Model(Qwen3_5Model): def sanitize(self, weights): - weights = tree_unflatten(list(weights.items())) - weights = dict(tree_flatten(weights)) - new_weights = {} for key, value in weights.items(): - if key.startswith("model.visual"): + if key.startswith("vision_tower") or key.startswith("model.visual"): continue if key.startswith("model.language_model"): key = key.replace("model.language_model", "language_model.model")