[object Object]

← back to Exo

Fix Qwen3-VL and autodetect vision config (#1893)

2ecefa0cfe1671db51d1e6c7aa8cdb1b9fa2fe7a · 2026-04-14 23:05:55 +0100 · rltakashige

## Motivation

Qwen3 VL TP doesn't work atm, and vision is not behaving.

## Test Plan

### Manual Testing
Works now.

Files touched

Diff

commit 2ecefa0cfe1671db51d1e6c7aa8cdb1b9fa2fe7a
Author: rltakashige <rl.takashige@gmail.com>
Date:   Tue Apr 14 23:05:55 2026 +0100

    Fix Qwen3-VL and autodetect vision config (#1893)
    
    ## Motivation
    
    Qwen3 VL TP doesn't work atm, and vision is not behaving.
    
    ## Test Plan
    
    ### Manual Testing
    Works now.
---
 src/exo/shared/models/model_cards.py        | 15 ++++++++++-----
 src/exo/worker/engines/mlx/auto_parallel.py | 23 ++++++++++++++++++++---
 2 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/src/exo/shared/models/model_cards.py b/src/exo/shared/models/model_cards.py
index bdbe3f63..83ac9d49 100644
--- a/src/exo/shared/models/model_cards.py
+++ b/src/exo/shared/models/model_cards.py
@@ -41,7 +41,7 @@ _BUILTIN_CARD_DIRS = [
 _card_cache: dict[ModelId, "ModelCard"] = {}
 
 
-def _detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
+def detect_vision_from_config(model_id: ModelId) -> "VisionCardConfig | None":
     normalized = model_id.normalize()
     for model_dir in [d / normalized for d in EXO_MODELS_DIRS]:
         config_path = model_dir / "config.json"
@@ -65,10 +65,6 @@ async def _load_cards_from_dir(directory: Path, *, is_custom: bool) -> None:
             card = await ModelCard.load_from_path(toml_file)
             if is_custom:
                 card = card.model_copy(update={"is_custom": True})
-            if card.vision is None:
-                vision = _detect_vision_from_config(card.model_id)
-                if vision is not None:
-                    card = card.model_copy(update={"vision": vision})
             if card.model_id not in _card_cache:
                 _card_cache[card.model_id] = card
         except (ValidationError, TOMLKitError):
@@ -140,6 +136,14 @@ class ModelCard(CamelCaseModel):
     is_custom: bool = False
     vision: VisionCardConfig | None = None
 
+    @model_validator(mode="after")
+    def _autodetect_vision(self) -> "ModelCard":
+        if self.vision is None:
+            detected = detect_vision_from_config(self.model_id)
+            if detected is not None:
+                object.__setattr__(self, "vision", detected)
+        return self
+
     @model_validator(mode="after")
     def _fill_vision_weights_repo(self) -> "ModelCard":
         if self.vision is not None and not self.vision.weights_repo:
@@ -256,6 +260,7 @@ class ConfigData(BaseModel):
             ["Qwen3MoeForCausalLM"],
             ["Qwen3_5MoeForConditionalGeneration"],
             ["Qwen3_5ForConditionalGeneration"],
+            ["Qwen3VLForConditionalGeneration"],
             ["MiniMaxM2ForCausalLM"],
             ["LlamaForCausalLM"],
             ["GptOssForCausalLM"],
diff --git a/src/exo/worker/engines/mlx/auto_parallel.py b/src/exo/worker/engines/mlx/auto_parallel.py
index 9566c838..64df97f2 100644
--- a/src/exo/worker/engines/mlx/auto_parallel.py
+++ b/src/exo/worker/engines/mlx/auto_parallel.py
@@ -38,6 +38,8 @@ from mlx_lm.models.nemotron_h import (
     NemotronHMoE,
 )
 from mlx_lm.models.nemotron_h import NemotronHModel as NemotronHInnerModel
+from mlx_lm.models.qwen3 import Model as Qwen3Model
+from mlx_lm.models.qwen3 import TransformerBlock as Qwen3TransformerBlock
 from mlx_lm.models.qwen3_5 import DecoderLayer as Qwen3_5DecoderLayer
 from mlx_lm.models.qwen3_5 import Model as Qwen3_5TextModel
 from mlx_lm.models.qwen3_5 import Qwen3_5TextModel as Qwen3_5TextModelInner
@@ -52,6 +54,7 @@ from mlx_lm.models.qwen3_next import (
     Qwen3NextSparseMoeBlock,
 )
 from mlx_lm.models.qwen3_next import Qwen3NextModel as Qwen3NextInnerModel
+from mlx_lm.models.qwen3_vl import Model as Qwen3VLModel
 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
@@ -540,7 +543,15 @@ def tensor_auto_parallel(
             sharded_to_all_linear_in_place,
         )
     elif isinstance(
-        model, (Qwen3MoeModel, Qwen3NextModel, Qwen3_5TextModel, Qwen3_5MoeModel)
+        model,
+        (
+            Qwen3Model,
+            Qwen3MoeModel,
+            Qwen3NextModel,
+            Qwen3_5TextModel,
+            Qwen3_5MoeModel,
+            Qwen3VLModel,
+        ),
     ):
         tensor_parallel_sharding_strategy = QwenShardingStrategy(
             group,
@@ -935,13 +946,19 @@ class QwenShardingStrategy(TensorParallelShardingStrategy):
         on_layer_loaded: LayerLoadedCallback | None,
     ) -> nn.Module:
         model = cast(
-            Qwen3MoeModel | Qwen3NextModel | Qwen3_5TextModel | Qwen3_5MoeModel, model
+            Qwen3Model
+            | Qwen3MoeModel
+            | Qwen3NextModel
+            | Qwen3_5TextModel
+            | Qwen3_5MoeModel
+            | Qwen3VLModel,
+            model,
         )
         total = len(model.layers)
         for i, layer in enumerate(model.layers):
             mx.eval(layer.parameters())
             # Shard the self attention
-            if isinstance(layer, Qwen3MoeDecoderLayer):
+            if isinstance(layer, (Qwen3MoeDecoderLayer, Qwen3TransformerBlock)):
                 layer.self_attn.q_proj = self.all_to_sharded_linear(
                     layer.self_attn.q_proj
                 )

← b8eaf707 Add gemma 4 tensor parallelism (#1891)  ·  back to Exo  ·  Fix out of order event idx causing fatal crashes (#1894) 2cd66ae4 →