Added changes to enable dynamo export and compile of VLMs using Qwen2.5 and 3 VLs as example.

QEfficient/exporter/weight_free/checkpoint_key_resolver.py

@@ -6,8 +6,6 @@
 # ----------------------------------------------------------------------------
 
 from pathlib import Path
-from typing import Dict, List, Optional
-
 import onnx_ir as ir
 from torch import nn
 
@@ -58,7 +56,7 @@ def _collect_tied_weights(model: nn.Module) -> list[TiedWeightAlias]:
     return [TiedWeightAlias(alias=alias, canonical=canonical) for alias, canonical in tied_mapping.items()]
 
 
-def _moe_weight_aliases(name: str) -> List[str]:
+def _moe_weight_aliases(name: str) -> list[str]:
     """Return equivalent checkpoint aliases for shared MoEWeights parameters."""
     aliases = []
     canonical = name
@@ -75,22 +73,51 @@ def _moe_weight_aliases(name: str) -> List[str]:
     return aliases
 
 
-def _find_checkpoint_key(candidates: List[str], checkpoint_index: Dict[str, str], onnx_name: str) -> Optional[str]:
+def _vlm_wrapper_aliases(name: str) -> list[str]:
+    """Return aliases introduced by multimodal wrapper nesting."""
+    aliases = []
+    if name.startswith("model.model."):
+        aliases.append("model." + name[len("model.model.") :])
+    if name.startswith("model.vision_model."):
+        aliases.append("model.visual." + name[len("model.vision_model.") :])
+    if name.startswith("vision_model."):
+        aliases.append("model.visual." + name[len("vision_model.") :])
+    if name.startswith("visual."):
+        aliases.append("model.visual." + name[len("visual.") :])
+    if name.startswith("language_model."):
+        aliases.append("model.language_model." + name[len("language_model.") :])
+    if name.startswith("model.lm_head."):
+        aliases.append("lm_head." + name[len("model.lm_head.") :])
+    if name.startswith("lm_head."):
+        aliases.append("model.lm_head." + name[len("lm_head.") :])
+    if name.endswith("lm_head.weight"):
+        prefix = name[: -len("lm_head.weight")]
+        aliases.extend(
+            [
+                f"{prefix}language_model.embed_tokens.weight",
+                f"{prefix}embed_tokens.weight",
+            ]
+        )
+    return aliases
+
+
+def _find_checkpoint_key(candidates: list[str], checkpoint_index: dict[str, str], onnx_name: str) -> str | None:
     """Return the unique matching checkpoint key, or fail on ambiguous matches."""
     seen = set()
     matches = []
     for candidate in candidates:
-        if candidate in seen:
-            continue
-        seen.add(candidate)
-        if candidate in checkpoint_index:
-            matches.append(candidate)
-        for alias in _moe_weight_aliases(candidate):
+        for alias in [candidate, *_vlm_wrapper_aliases(candidate)]:
             if alias in seen:
                 continue
             seen.add(alias)
             if alias in checkpoint_index:
                 matches.append(alias)
+            for moe_alias in _moe_weight_aliases(alias):
+                if moe_alias in seen:
+                    continue
+                seen.add(moe_alias)
+                if moe_alias in checkpoint_index:
+                    matches.append(moe_alias)
     if len(matches) > 1:
         raise ValueError(
             f"Ambiguous checkpoint key for ONNX initializer '{onnx_name}': matched {matches}. "
@@ -106,9 +133,9 @@ def _is_computed_initializer(name: str) -> bool:
 
 def find_checkpoint_key(
     onnx_name: str,
-    checkpoint_index: Dict[str, str],
+    checkpoint_index: dict[str, str],
     backbone: nn.Module,
-) -> Optional[str]:
+) -> str | None:
     """Resolve an ONNX initializer name to its safetensors checkpoint key.
 
     Most weights match directly. The fallback rules cover wrapper prefixes,
@@ -179,15 +206,14 @@ def promote_initializers_and_build_spec(onnx_program, model_ref: str, model_name
         for checkpoint_file in checkpoint_files
     ]
     backbone = qeff_model.model.base_model if isinstance(qeff_model.model, PooledModel) else qeff_model.model
-    promoted_inputs: List[WeightSpecInput] = []
+    promoted_inputs: list[WeightSpecInput] = []
 
     for name, init_value in list(model_ir.graph.initializers.items()):
-        if name not in model_names:
-            continue
-
         onnx_name = tied_weight_map.get(name, name)
         checkpoint_key = find_checkpoint_key(onnx_name, checkpoint_index, backbone)
         if checkpoint_key is None:
+            if name not in model_names:
+                continue
             if _is_computed_initializer(onnx_name):
                 continue
             raise ValueError(

QEfficient/exporter/weight_free/export.py

@@ -35,6 +35,34 @@ def _to_meta(value: Any) -> Any:
     return value
 
 
+def _iter_weight_free_configs(qeff_model):
+    model = getattr(qeff_model, "model", None)
+    nested_model = getattr(model, "model", None)
+
+    for config in (
+        getattr(model, "config", None),
+        getattr(qeff_model, "config", None),
+        getattr(getattr(model, "vision_model", None), "config", None),
+        getattr(getattr(nested_model, "vision_model", None), "config", None),
+        getattr(nested_model, "config", None),
+    ):
+        if config is not None:
+            yield config
+
+
+def _resolve_weight_free_config(qeff_model):
+    return next(_iter_weight_free_configs(qeff_model), None)
+
+
+def _resolve_weight_free_target_dtype(qeff_model) -> torch.dtype:
+    for config in _iter_weight_free_configs(qeff_model):
+        for attr in ("dtype", "torch_dtype"):
+            dtype = getattr(config, attr, None)
+            if isinstance(dtype, torch.dtype):
+                return dtype
+    return torch.float32
+
+
 def _run_quantizer_for_wf(qeff_model, target_dtype: torch.dtype):
     """Finish preparing a meta-device QEfficient wrapper for weight-free tracing, in place."""
     model_ref = qeff_model.hash_params.get("pretrained_model_name_or_path")
@@ -44,7 +72,8 @@ def _run_quantizer_for_wf(qeff_model, target_dtype: torch.dtype):
             "Pass `pretrained_model_name_or_path=...` when constructing the QEff model manually."
         )
 
-    quant_config = getattr(qeff_model.model.config, "quantization_config", None)
+    config = _resolve_weight_free_config(qeff_model)
+    quant_config = getattr(config, "quantization_config", None)
 
     if quant_config is not None:
         # For quantized models the meta model must use the same quantized layer types as the
@@ -135,7 +164,17 @@ def _prepare_checkpoint_for_weight_free_export(
     dtype_suffix = str(target_dtype).replace("torch.", "")
     # TODO(wf): For different flavours of the model that expect different checkpoint weight layouts,
     # we end up overriding old one. We need to add support of hashing/caching here.
-    prepared_name = source_dir.name + f"-qeff-prepared-{dtype_suffix}"
+    hash_params = dict(getattr(qeff_model, "hash_params", {}) or {})
+    flavour = hash_params.get("moe_prefill_flavour")
+    if hasattr(flavour, "value"):
+        flavour = flavour.value
+    expert_parallel_suffix = ""
+    if flavour == "expert_parallel":
+        num_parallelized_experts = hash_params.get("moe_prefill_num_parallelized_experts")
+        num_pipeline_stages = hash_params.get("moe_prefill_num_pipeline_stages")
+        if num_parallelized_experts is not None and num_pipeline_stages is not None:
+            expert_parallel_suffix = f"-moe-expert-parallel-{num_parallelized_experts}x{num_pipeline_stages}"
+    prepared_name = source_dir.name + f"-qeff-prepared-{dtype_suffix}{expert_parallel_suffix}"
     if QEFF_CHECKPOINT_HOME:
         prepared_out = QEFF_CHECKPOINT_HOME.expanduser() / prepared_name
     else:
@@ -146,6 +185,7 @@ def _prepare_checkpoint_for_weight_free_export(
             src=source_dir,
             out=prepared_out,
             target_dtype=target_dtype,
+            hash_params=hash_params,
         )
     )
 
@@ -186,7 +226,7 @@ def export_weight_free_onnx(
     tuple
         Meta QEfficient model, updated ONNX transform kwargs, and cleanup callback.
     """
-    target_dtype = qeff_model.model.config.dtype
+    target_dtype = _resolve_weight_free_target_dtype(qeff_model)
     meta_qeff_model = _run_quantizer_for_wf(qeff_model, target_dtype)
 
     # export_wrapper (the @export_wrapper decorator on _export) already ran