Qualcomm AI Engine Direct - Support Multimodal VLMs (pytorch#16292)

Summary:
 - add support for multimodal AOT and calibration
 - AOT support for VLMs:
   - SmolVLM-500M
   - InternVL3-1B
 - add quantization recipe for VLMs and vision encoders
 - Refactor LLM AOT and integrated it into multimodal AOT

### Test plan
Test lowering of the quantized SmolVLM-500M
``` bash
python -m backends.qualcomm.tests.test_qnn_delegate TestExampleMultimodalityScript.test_smolvlm_500m_instruct -b build-android --executorch_root . -a . -m SM8750
```
Test lowering of the quantized InternVL3-1B
``` bash
python -m backends.qualcomm.tests.test_qnn_delegate TestExampleMultimodalityScript.test_internvl3_1b -b build-android --executorch_root . -a . -m SM8750
```

Author: DannyYuyang-quic

backends/qualcomm/tests/test_qnn_delegate.py

@@ -6516,6 +6516,126 @@ def test_qwen2_5(self):
                     )
 
 
+class TestExampleMultimodalityScript(TestQNN):
+    def test_smolvlm_500m_instruct(self):
+        if not self.required_envs():
+            self.skipTest("missing required envs")
+
+        prompt = "Can you describe this image?"
+        cmds = [
+            "python",
+            f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py",
+            "--artifact",
+            self.artifact_dir,
+            "--build_folder",
+            self.build_folder,
+            "--model",
+            self.model,
+            "--ip",
+            self.ip,
+            "--port",
+            str(self.port),
+            "--prompt",
+            prompt,
+            "--temperature",
+            "0",
+            "--decoder_model",
+            "smolvlm_500m_instruct",
+            "--model_mode",
+            "kv",
+            "--max_seq_len",
+            "128",
+        ]
+        if self.compile_only:
+            cmds.extend(["--compile_only"])
+        elif self.device:
+            cmds.extend(["--device", self.device])
+        if self.host:
+            cmds.extend(["--host", self.host])
+        elif self.enable_x86_64:
+            cmds.extend(["--enable_x86_64"])
+        if self.pre_gen_pte:
+            cmds.extend(["--pre_gen_pte", self.pre_gen_pte])
+
+        p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)
+        with Listener((self.ip, self.port)) as listener:
+            conn = listener.accept()
+            p.communicate()
+            msg = json.loads(conn.recv())
+            if "Error" in msg:
+                self.fail(msg["Error"])
+            else:
+                if not self.enable_x86_64:
+                    encoder_pte_size = msg["encoder_pte_size"]
+                    text_embedding_pte_size = msg["text_embedding_pte_size"]
+                    decoder_pte_size = msg["pte_size"]
+                    self.assertLessEqual(encoder_pte_size, 110_000_000)  # 110MB
+                    self.assertLessEqual(text_embedding_pte_size, 100_000_000)  # 100MB
+                    self.assertLessEqual(decoder_pte_size, 400_000_000)  # 400MB
+                    print(f"Encoder PTE Size: {encoder_pte_size} bytes")
+                    print(f"Text Embedding PTE Size: {text_embedding_pte_size} bytes")
+                    print(f"Decoder PTE Size: {decoder_pte_size} bytes")
+
+    def test_internvl3_1b(self):
+        if not self.required_envs():
+            self.skipTest("missing required envs")
+
+        prompt = "Can you describe this image?"
+        cmds = [
+            "python",
+            f"{self.executorch_root}/examples/qualcomm/oss_scripts/llama/llama.py",
+            "--artifact",
+            self.artifact_dir,
+            "--build_folder",
+            self.build_folder,
+            "--model",
+            self.model,
+            "--ip",
+            self.ip,
+            "--port",
+            str(self.port),
+            "--prompt",
+            prompt,
+            "--temperature",
+            "0",
+            "--decoder_model",
+            "internvl3_1b",
+            "--model_mode",
+            "kv",
+            "--max_seq_len",
+            "320",
+        ]
+        if self.compile_only:
+            cmds.extend(["--compile_only"])
+        elif self.device:
+            cmds.extend(["--device", self.device])
+        if self.host:
+            cmds.extend(["--host", self.host])
+        elif self.enable_x86_64:
+            cmds.extend(["--enable_x86_64"])
+        if self.pre_gen_pte:
+            cmds.extend(["--pre_gen_pte", self.pre_gen_pte])
+
+        p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)
+        with Listener((self.ip, self.port)) as listener:
+            conn = listener.accept()
+            p.communicate()
+            msg = json.loads(conn.recv())
+            if "Error" in msg:
+                self.fail(msg["Error"])
+            else:
+                if not self.enable_x86_64:
+                    encoder_pte_size = msg["encoder_pte_size"]
+                    text_embedding_pte_size = msg["text_embedding_pte_size"]
+                    decoder_pte_size = msg["pte_size"]
+                    self.assertLessEqual(encoder_pte_size, 425_000_000)  # 425MB
+                    self.assertLessEqual(text_embedding_pte_size, 300_000_000)  # 300MB
+                    self.assertLessEqual(decoder_pte_size, 550_000_000)  # 550MB
+                    print(f"Encoder PTE Size: {encoder_pte_size} bytes")
+                    print(f"Text Embedding PTE Size: {text_embedding_pte_size} bytes")
+                    print(f"Decoder PTE Size: {decoder_pte_size} bytes")
+
+
 class TestExampleOssScript(TestQNN):
     def test_albert(self):
         if not self.required_envs([self.sentence_dataset]):

examples/models/internvl3/1b_config.json

@@ -0,0 +1,14 @@
+{
+  "dim": 896,
+  "ffn_dim_multiplier": 1,
+  "hidden_dim": 4864,
+  "n_heads": 14,
+  "n_kv_heads": 2,
+  "n_layers": 24,
+  "norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "use_scaled_rope": false,
+  "vocab_size": 151674,
+  "use_hf_rope": true,
+  "attention_qkv_bias": true
+}

examples/models/internvl3/__init__.py

@@ -0,0 +1,16 @@
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from executorch.examples.models.internvl3.convert_weights import convert_weights
+from executorch.examples.models.llama.model import Llama2Model
+
+
+class InternVL3Model(Llama2Model):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+__all__ = [
+    "InternVL3Model",
+    "convert_weights",
+]

examples/models/internvl3/convert_weights.py

@@ -0,0 +1,82 @@
+import argparse
+from typing import Dict
+
+import torch
+
+from executorch.examples.models.smollm3.convert_weights import load_checkpoint
+from torchtune.models.convert_weights import get_mapped_key
+
+# Standard _FROM_META weight mapping of Meta weights to TorchTune + additional bias weight mappings.
+_INTERNVL_TO_META = {
+    "language_model.model.embed_tokens.weight": "tok_embeddings.weight",
+    "language_model.model.norm.weight": "norm.weight",
+    "language_model.model.layers.{}.self_attn.k_proj.weight": "layers.{}.attention.wk.weight",
+    "language_model.model.layers.{}.self_attn.k_proj.bias": "layers.{}.attention.wk.bias",
+    "language_model.model.layers.{}.self_attn.q_proj.weight": "layers.{}.attention.wq.weight",
+    "language_model.model.layers.{}.self_attn.q_proj.bias": "layers.{}.attention.wq.bias",
+    "language_model.model.layers.{}.self_attn.v_proj.weight": "layers.{}.attention.wv.weight",
+    "language_model.model.layers.{}.self_attn.v_proj.bias": "layers.{}.attention.wv.bias",
+    "language_model.model.layers.{}.self_attn.o_proj.weight": "layers.{}.attention.wo.weight",
+    "language_model.model.layers.{}.input_layernorm.weight": "layers.{}.attention_norm.weight",
+    "language_model.model.layers.{}.post_attention_layernorm.weight": "layers.{}.ffn_norm.weight",
+    "language_model.model.layers.{}.mlp.gate_proj.weight": "layers.{}.feed_forward.w1.weight",
+    "language_model.model.layers.{}.mlp.down_proj.weight": "layers.{}.feed_forward.w2.weight",
+    "language_model.model.layers.{}.mlp.up_proj.weight": "layers.{}.feed_forward.w3.weight",
+}
+
+
+def internvl_tune_to_meta(
+    state_dict: Dict[str, torch.Tensor]
+) -> Dict[str, torch.Tensor]:
+    """
+    Convert a state dict from torchtune's format to Meta's format. This function
+    doesn't handle any sharding or splitting of state dicts. It follows the
+    state_dict IN -> state_dict OUT pattern.
+
+    Args:
+        state_dict (Dict[str, torch.Tensor]): State dict in torchtune's format.
+
+    Returns:
+        Dict[str, torch.Tensor]: State dict in Meta's format.
+    """
+    converted_text_model_state_dict = {}
+    for key, value in state_dict.items():
+        try:
+            new_key = get_mapped_key(key, _INTERNVL_TO_META)
+            converted_text_model_state_dict[new_key] = value
+        except:
+            # only preserve parameters of text decoder
+            pass
+    converted_text_model_state_dict["output.weight"] = state_dict[
+        "language_model.lm_head.weight"
+    ]
+    return converted_text_model_state_dict
+
+
+def convert_weights(input_dir: str, output_file: str) -> None:
+    print("Loading checkpoint...")
+    sd = load_checkpoint(input_dir)
+    print("Converting checkpoint...")
+    converted_sd = internvl_tune_to_meta(sd)
+    print("Saving checkpoint...")
+    torch.save(converted_sd, output_file)
+    print("Done.")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert InternVL3 weights to Meta format."
+    )
+    parser.add_argument(
+        "input_dir",
+        type=str,
+        help="Path to directory containing checkpoint files",
+    )
+    parser.add_argument("output", type=str, help="Path to the output checkpoint")
+
+    args = parser.parse_args()
+    convert_weights(args.input_dir, args.output)
+
+
+if __name__ == "__main__":
+    main()

examples/models/smolvlm/500M_config.json

@@ -0,0 +1,14 @@
+{
+    "dim": 960,
+    "ffn_dim_multiplier": 1,
+    "hidden_dim": 2560,
+    "n_heads": 15,
+    "n_kv_heads": 5,
+    "n_layers": 32,
+    "norm_eps": 1e-05,
+    "rope_theta": 100000,
+    "use_scaled_rope": false,
+    "vocab_size": 49280,
+    "use_hf_rope": false,
+    "attention_qkv_bias": false
+  }

examples/models/smolvlm/__init__.py

@@ -0,0 +1,16 @@
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from executorch.examples.models.llama.model import Llama2Model
+from executorch.examples.models.smolvlm.convert_weights import convert_weights
+
+
+class SmolVLMModel(Llama2Model):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+
+__all__ = [
+    "SmolVLMModel",
+    "convert_weights",
+]

examples/models/smolvlm/convert_weights.py

@@ -0,0 +1,77 @@
+import argparse
+from typing import Dict
+
+import torch
+
+from executorch.examples.models.smollm3.convert_weights import load_checkpoint
+from torchtune.models.convert_weights import get_mapped_key
+
+# Standard _FROM_META weight mapping of Meta weights to TorchTune + additional bias weight mappings.
+_SMOLVLM_TO_META = {
+    "model.text_model.embed_tokens.weight": "tok_embeddings.weight",
+    "model.text_model.norm.weight": "norm.weight",
+    "model.text_model.layers.{}.self_attn.k_proj.weight": "layers.{}.attention.wk.weight",
+    "model.text_model.layers.{}.self_attn.q_proj.weight": "layers.{}.attention.wq.weight",
+    "model.text_model.layers.{}.self_attn.v_proj.weight": "layers.{}.attention.wv.weight",
+    "model.text_model.layers.{}.self_attn.o_proj.weight": "layers.{}.attention.wo.weight",
+    "model.text_model.layers.{}.input_layernorm.weight": "layers.{}.attention_norm.weight",
+    "model.text_model.layers.{}.post_attention_layernorm.weight": "layers.{}.ffn_norm.weight",
+    "model.text_model.layers.{}.mlp.gate_proj.weight": "layers.{}.feed_forward.w1.weight",
+    "model.text_model.layers.{}.mlp.down_proj.weight": "layers.{}.feed_forward.w2.weight",
+    "model.text_model.layers.{}.mlp.up_proj.weight": "layers.{}.feed_forward.w3.weight",
+}
+
+
+def smolvlm_tune_to_meta(
+    state_dict: Dict[str, torch.Tensor]
+) -> Dict[str, torch.Tensor]:
+    """
+    Convert a state dict from torchtune's format to Meta's format. This function
+    doesn't handle any sharding or splitting of state dicts. It follows the
+    state_dict IN -> state_dict OUT pattern.
+
+    Args:
+        state_dict (Dict[str, torch.Tensor]): State dict in torchtune's format.
+
+    Returns:
+        Dict[str, torch.Tensor]: State dict in Meta's format.
+    """
+    converted_text_model_state_dict = {}
+    for key, value in state_dict.items():
+        try:
+            new_key = get_mapped_key(key, _SMOLVLM_TO_META)
+            converted_text_model_state_dict[new_key] = value
+        except:
+            # only preserve parameters of text decoder
+            pass
+    converted_text_model_state_dict["output.weight"] = state_dict["lm_head.weight"]
+    return converted_text_model_state_dict
+
+
+def convert_weights(input_dir: str, output_file: str) -> None:
+    print("Loading checkpoint...")
+    sd = load_checkpoint(input_dir)
+    print("Converting checkpoint...")
+    converted_sd = smolvlm_tune_to_meta(sd)
+    print("Saving checkpoint...")
+    torch.save(converted_sd, output_file)
+    print("Done.")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Convert SmolVLM weights to Meta format."
+    )
+    parser.add_argument(
+        "input_dir",
+        type=str,
+        help="Path to directory containing checkpoint files",
+    )
+    parser.add_argument("output", type=str, help="Path to the output checkpoint")
+
+    args = parser.parse_args()
+    convert_weights(args.input_dir, args.output)
+
+
+if __name__ == "__main__":
+    main()

examples/qualcomm/oss_scripts/llama/README.md

@@ -19,7 +19,7 @@ This file provides you the instructions to run LLM Decoder model with different
  1. QWEN3 0.6B / 1.7B
  1. SmolLM2 135M
  1. SmolLM3 3B
- 
+
 
 We offer the following modes to execute the model: