Add multimodal evaluation and component-aware quantization support

.gitignore

@@ -10,6 +10,8 @@ augmented_model.onnx
 test/**/data
 nc_workspace/
 run_cache/
+/submissions/
+/Delwin_testing/
 
 # Onnx dynamo export artifacts
 *.sarif

docs/source/features/quantization.md

@@ -59,6 +59,36 @@ This pass only supports HuggingFace transformer PyTorch models.
 }
 ```
 
+## MBQ
+`Mbq` implements the decoder reparameterization from [Modality-Balanced Quantization](https://arxiv.org/abs/2412.19509). It uses vision- and answer-token gradients to balance an AWQ-style reconstruction search, then folds the selected channel scales into decoder norms and linear weights. The output remains a regular Hugging Face checkpoint and must be followed by `Rtn`, `Gptq`, or another compatible weight quantizer using the same `bits`, `group_size`, and `sym` values.
+
+Calibration batches must contain processor-ready model inputs, `labels`, and a boolean `vision_mask`. An optional `answer_mask` overrides the default `labels != -100` mask. Both masks must match the expanded decoder sequence, not the pre-expansion prompt. The initial implementation supports `Qwen2DecoderLayer`, `Qwen2VLDecoderLayer`, and the current Transformers `Qwen2_5_VLDecoderLayer`. Qwen3-VL is not yet supported because its text-model loop injects DeepStack visual features between decoder layers, outside the individual layer calls used by MBQ replay.
+
+### Example Configuration
+```json
+{
+    "type": "Mbq",
+    "bits": 4,
+    "group_size": 128,
+    "sym": false,
+    "data_config": "multimodal_calibration"
+}
+```
+
+## Multimodal activation range calibration
+`MultimodalActivationRangeCalibration` collects separate static activation ranges for vision, audio, text, and projector components. It is an MQuant-inspired diagnostic and planning pass: it does **not** insert MQuant MSQ runtime selection, AIFS token reordering, or activation quantization kernels. The resulting `multimodal_activation_ranges` model attribute records the observer configuration, component rules, sample counts, and derived quantization parameters for future runtime integration.
+
+### Example Configuration
+```json
+{
+    "type": "MultimodalActivationRangeCalibration",
+    "data_config": "multimodal_calibration",
+    "components": ["vision", "text"],
+    "bits": 8,
+    "sym": true
+}
+```
+
 ## RTN
 `RTN (Round To Nearest)` is a fast, calibration-free weight quantization method that enables low-bit quantization of large models without relying on gradient-based optimization or calibration datasets. RTN quantization uses simple rounding to the nearest quantization level, making it extremely fast while maintaining reasonable accuracy.
 
@@ -365,4 +395,3 @@ Configurations:
 ```
 
 Please refer to [AimetQuantization](aimet_quantization) for more details about the pass and its config parameters.
-

docs/source/reference/pass.rst

@@ -235,6 +235,24 @@ SelectiveMixedPrecision
 -----------------------
 .. autoconfigclass:: olive.passes.SelectiveMixedPrecision
 
+.. _multimodal_mixed_precision:
+
+MultiModalMixedPrecision
+------------------------
+.. autoconfigclass:: olive.passes.MultiModalMixedPrecision
+
+.. _multimodal_activation_range_calibration:
+
+MultimodalActivationRangeCalibration
+------------------------------------
+.. autoconfigclass:: olive.passes.MultimodalActivationRangeCalibration
+
+.. _mbq:
+
+Mbq
+---
+.. autoconfigclass:: olive.passes.Mbq
+
 .. _lora:
 
 LoRA

olive/cache.py

@@ -532,6 +532,16 @@
             model_path_resource = model_json["config"]["model_path"]
             source_path = Path(model_path_resource.get_path())
             onnx_file_name = model_json["config"].get("onnx_file_name")
+            model_attributes = model_json["config"].get("model_attributes") or {}
+
+            if model_attributes.get("ort_genai_package") and source_path.is_dir():
+                if output_dir.suffix == ".onnx":
+                    raise ValueError("ORT-GenAI packages must be saved to a directory, not an ONNX file path.")
+                output_dir.mkdir(parents=True, exist_ok=True)
+                shutil.copytree(source_path, output_dir, dirs_exist_ok=True)
+                model_json["config"]["model_path"] = str(output_dir)
+                local_resource_names = [rn for rn in local_resource_names if rn != "model_path"]
+                return self._save_additional_files(model_json, output_dir)
 
             # Determine if source has external data or additional files
             has_additional_files = bool(onnx_file_name)

olive/common/hf/wrapper.py

@@ -99,6 +99,7 @@ class LayerWrapper:
     SECOND_LAYER_NORM = {
         "default": "post_attention_layernorm",
         "gemma2": "pre_feedforward_layernorm",
+        "gemma4": "pre_feedforward_layernorm",
         "gpt2": "ln_2",
         "lfm2": "ffn_norm",
         "opt": "final_layer_norm",
@@ -161,9 +162,17 @@ def get_attention_inputs(self, return_name: bool = True):
         if self.attn is None:
             return ([], []) if return_name else []
         attention_inputs, names = get_submodules(
-            self.attn, self.ATTENTION_INPUTS, self.model_type, return_name=True, return_name_prefix=f"{self.attn_name}."
+            self.attn,
+            self.ATTENTION_INPUTS,
+            self.model_type,
+            return_name=True,
+            return_name_prefix=f"{self.attn_name}.",
+            fail_on_not_found=False,
         )
-        if isinstance(attention_inputs[0], UnpackedQKV):
+        present = [(module, name) for module, name in zip(attention_inputs, names) if module is not None]
+        attention_inputs = [module for module, _ in present]
+        names = [name for _, name in present]
+        if attention_inputs and isinstance(attention_inputs[0], UnpackedQKV):
             names = [f"{names[0]}.{part}" for part in ["q_proj", "k_proj", "v_proj"]]
             attention_inputs = [attention_inputs[0].q_proj, attention_inputs[0].k_proj, attention_inputs[0].v_proj]
         return attention_inputs if not return_name else (attention_inputs, names)
@@ -204,49 +213,79 @@ class ModelWrapper:
         "default": ["model.embed_tokens"],
         "bloom": ["transformer.word_embeddings", "transformer.word_embeddings_layernorm"],
         "falcon": ["transformer.word_embeddings"],
+        "gemma4": [
+            "model.language_model.embed_tokens",
+            "model.language_model.embed_tokens_per_layer",
+            "model.language_model.per_layer_model_projection",
+            "model.language_model.per_layer_projection_norm",
+        ],
         "gpt2": ["transformer.wte", "transformer.wpe"],
         "gpt_neox": ["gpt_neox.embed_in"],
         "gptj": ["transformer.wte"],
         "opt": ["model.decoder.embed_tokens", "model.decoder.embed_positions"],
         "qwen": ["transformer.wte"],
+        "qwen2_5_vl": ["model.language_model.embed_tokens"],
+        "qwen3_vl": ["model.language_model.embed_tokens"],
     }
     # in newer transformers versions, there is one rotary embedding per model
     ROTARY_EMBEDDING = {
         "default": "model.rotary_emb",
         "falcon": "transformer.rotary_emb",
+        "gemma4": "model.language_model.rotary_emb",
         "gpt_neox": "gpt_neox.rotary_emb",
         "qwen": "transformer.rotary_emb",
+        "qwen2_5_vl": "model.language_model.rotary_emb",
+        "qwen3_vl": "model.language_model.rotary_emb",
     }
     LM_HEAD = {"default": "lm_head"}
     PRE_HEAD_LAYERNORM = {
         "default": "model.norm",
+        "gemma4": "model.language_model.norm",
         "gpt2": "transformer.ln_f",
         "lfm2": "model.embedding_norm",
         "qwen": "transformer.ln_f",
+        "qwen2_5_vl": "model.language_model.norm",
+        "qwen3_vl": "model.language_model.norm",
     }
     LAYERS = {
         "default": "model.layers",
         "bloom": "transformer.h",
         "falcon": "transformer.h",
+        "gemma4": "model.language_model.layers",
         "gpt2": "transformer.h",
         "gpt_neox": "gpt_neox.layers",
         "gptj": "transformer.h",
         "opt": "model.decoder.layers",
         "qwen": "transformer.h",
+        "qwen2_5_vl": "model.language_model.layers",
+        "qwen3_vl": "model.language_model.layers",
     }
 
     def __init__(self, config: Union[PretrainedConfig, dict]):
         self.config = config if isinstance(config, PretrainedConfig) else PretrainedConfig.from_dict(config)
         self.model_type = getattr(self.config, "model_type", None)
 
+        # Multimodal Hugging Face configs commonly nest decoder metadata under
+        # text_config while keeping the module layout keyed by the composite model_type.
+        text_config = getattr(self.config, "text_config", None)
+
+        def resolve_model_alias(name: str):
+            value = resolve_alias(self.config, name)
+            if value is None and text_config is not None:
+                value = resolve_alias(text_config, name)
+            return value
+
         # model attributes (using unified aliases from defaults.yaml)
-        self.hidden_size = resolve_alias(self.config, "hidden_size")
-        self.num_attention_heads = resolve_alias(self.config, "num_attention_heads")
-        self.num_key_value_heads = resolve_alias(self.config, "num_kv_heads") or self.num_attention_heads
-        self.head_dim = resolve_alias(self.config, "head_dim") or self.hidden_size // self.num_attention_heads
-        self.num_hidden_layers = resolve_alias(self.config, "num_layers")
+        self.hidden_size = resolve_model_alias("hidden_size")
+        self.num_attention_heads = resolve_model_alias("num_attention_heads")
+        self.num_key_value_heads = resolve_model_alias("num_kv_heads") or self.num_attention_heads
+        self.head_dim = resolve_model_alias("head_dim") or self.hidden_size // self.num_attention_heads
+        self.num_hidden_layers = resolve_model_alias("num_layers")
         # MAX_LENGTH uses model_type-based mapping, not aliases
-        self.max_length = find_first_matched_value(self.config, self.MAX_LENGTH)
+        max_length_name = self.MAX_LENGTH.get(self.model_type, self.MAX_LENGTH["default"])
+        self.max_length = find_first_matched_value(self.config, max_length_name)
+        if self.max_length is None and text_config is not None:
+            self.max_length = find_first_matched_value(text_config, max_length_name)
 
         self._model = None
         self._layer_wrappers = None
@@ -263,7 +302,18 @@ def set_model(self, model: "PreTrainedModel"):
         self._layer_wrappers = [LayerWrapper(layer, self.model_type) for layer in self.get_layers(False)]
 
     def get_embeds(self, return_name: bool = True):
-        return get_submodules(self.model, self.EMBEDDINGS, self.model_type, return_name=return_name)
+        embeds, names = get_submodules(
+            self.model,
+            self.EMBEDDINGS,
+            self.model_type,
+            return_name=True,
+            fail_on_not_found=self.model_type != "gemma4",
+        )
+        if isinstance(embeds, list):
+            present = [(module, name) for module, name in zip(embeds, names) if module is not None]
+            embeds = [module for module, _ in present]
+            names = [name for _, name in present]
+        return embeds if not return_name else (embeds, names)
 
     def get_rotary_embed(self, return_name: bool = True):
         return get_submodules(