Onboarding Molmo Model

Signed-off-by: Mohit Soni <[email protected]>

Author: Mohit Soni

QEfficient/transformers/models/modeling_auto.py

@@ -479,7 +479,9 @@ def model_name(self) -> str:
 
     @property
     def get_model_config(self) -> dict:
-        return self.model.model.vision_model.config.__dict__
+        if hasattr(self.model.model, "vision_model"):
+            return self.model.model.vision_model.config.__dict__
+        return self.model.model.config.__dict__
 
 
 class QEffCausalLMForTextImageToTextModel(QEFFBaseModel):
@@ -536,7 +538,9 @@ def model_name(self) -> str:
 
     @property
     def get_model_config(self) -> dict:
-        return self.model.language_model.config.__dict__
+        if hasattr(self.model, "language_model"):
+            return self.model.language_model.config.__dict__
+        return self.model.config.__dict__
 
 
 class _QEffAutoModelForImageTextToTextDualQPC:
@@ -652,7 +656,11 @@ def compile(
 
         custom_io_vision = {}
         kv_cache_dtype = "mxint8" if mxint8_kv_cache else "float16"
+        molmo = hasattr(self.model.config, "model_type") and self.model.config.model_type == "molmo"
+        if molmo:
+            custom_io_vision["image_masks"] = "float16"
         custom_io_vision["pixel_values"] = "float16"
+
         for output_name in output_names["vision"]:
             if output_name.startswith("past_"):
                 custom_io_vision[output_name] = kv_cache_dtype
@@ -804,11 +812,18 @@ def kv_offload_generate(
             inputs[k] = np.array(v)
 
         vision_inputs = {
-            k: v for k, v in inputs.items() if k in {"pixel_values", "aspect_ratio_ids", "aspect_ratio_mask"}
+            k: v
+            for k, v in inputs.items()
+            if k
+            in {"pixel_values", "image_masks", "image_input_idx", "valid_idx", "aspect_ratio_ids", "aspect_ratio_mask"}
         }
 
+        molmo = hasattr(self.model.config, "model_type") and self.model.config.model_type == "molmo"
+
         if vision_inputs:
             vision_inputs["pixel_values"] = vision_inputs["pixel_values"].astype("float16")
+            if molmo:
+                vision_inputs["image_masks"] = vision_inputs["image_masks"].astype("float16")
         vision_start = perf_counter()
 
         vision_outputs = {}
@@ -923,7 +938,10 @@ def __init__(
             self.model.config.llm_config._attn_implementation = "eager"
             self.model.config.vision_config.use_flash_attn = "false"
         else:
-            self.model.config.text_config.use_cache = True
+            if hasattr(self.model.config, "text_config"):
+                self.model.config.text_config.use_cache = True
+            else:
+                self.model.config.use_cache = True
         self.hash_params["qeff_auto_class"] = self.__class__.__name__
 
     @classmethod
@@ -1292,7 +1310,10 @@ def from_pretrained(cls, pretrained_model_name_or_path: str, kv_offload: Optiona
         return cls(model, kv_offload=kv_offload, pretrained_model_name_or_path=pretrained_model_name_or_path, **kwargs)
 
 
-MISCLASSIFIED_CAUSAL_LM_TO_QEFF_AUTO_CLASS_MAP = {"InternVLChatModel": QEFFAutoModelForImageTextToText}
+MISCLASSIFIED_CAUSAL_LM_TO_QEFF_AUTO_CLASS_MAP = {
+    "InternVLChatModel": QEFFAutoModelForImageTextToText,
+    "MolmoForCausalLM": QEFFAutoModelForImageTextToText,
+}
 
 
 class QEFFAutoModelForCausalLM(QEFFBaseModel):

QEfficient/transformers/models/molmo/__init__.py

@@ -0,0 +1,6 @@
+# -----------------------------------------------------------------------------
+#
+# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries.
+# SPDX-License-Identifier: BSD-3-Clause
+#
+# -----------------------------------------------------------------------------