Qualcomm AI Engine Direct - Gemma4 Text model enablement (Unified Architecture)

backends/qualcomm/tests/test_passes.py

@@ -162,9 +162,7 @@ def test_mha_to_sha(self):
         from executorch.examples.qualcomm.oss_scripts.llama.masking_utils import (
             CausalAttentionMask,
         )
-        from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import (
-            LlamaAttention,
-        )
+        from executorch.examples.qualcomm.oss_scripts.llama.model import LlamaAttention
 
         # Initailize model config
         args = ModelArgs()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -8144,6 +8144,14 @@ def setUp(self):
                 hellaswag_acc_norm=None,
                 sqnr=10,
             ),
+            "gemma4-e2b": TestExampleLLMScript.LlmSpecs(
+                SM8650=20,
+                SM8750=30,
+                pte_size=4_500_000_000,  # 4.5 GB
+                wikitext_ppl=120,
+                hellaswag_acc_norm=None,
+                sqnr=10,
+            ),
             "glm-1_5b": TestExampleLLMScript.LlmSpecs(
                 SM8650=42,
                 SM8750=52,
@@ -8345,6 +8353,14 @@ def test_static_llm_model(self):  # noqa: C901
                 ]
             )
 
+        if self.model_name == "gemma4-e2b":
+            cmds.extend(
+                [
+                    "--embedding-quantize",
+                    "4,32",
+                ]
+            )
+
         self.add_default_cmds(cmds)
 
         p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL)

examples/models/gemma4/__init__.py

@@ -6,8 +6,9 @@
 
 from executorch.examples.models.gemma4.text_decoder import (
     convert_hf_to_custom,
+    convert_weights,
     Gemma4Config,
     Gemma4Model,
 )
 
-__all__ = ["Gemma4Config", "Gemma4Model", "convert_hf_to_custom"]
+__all__ = ["Gemma4Config", "Gemma4Model", "convert_hf_to_custom", "convert_weights"]

examples/models/gemma4/config/e2b_config.json

@@ -22,6 +22,9 @@
     "hidden_size_per_layer_input": 256,
     "num_kv_shared_layers": 20,
     "tie_word_embeddings": true,
+    "post_attention_norm": true,
+    "post_ffn_norm": true,
+    "use_ffn_norm": true,
     "layer_types": [
       "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention",
       "sliding_attention","sliding_attention","sliding_attention","sliding_attention","full_attention",

examples/models/gemma4/text_decoder/__init__.py

@@ -5,7 +5,7 @@
 # pyre-unsafe
 # LICENSE file in the root directory of this source tree.
 
-from .convert_weights import convert_hf_to_custom  # noqa: F401
+from .convert_weights import convert_hf_to_custom, convert_weights  # noqa: F401
 from .gemma4_attention import (  # noqa: F401
     apply_rotary_emb,
     apply_rotary_emb_single,

examples/models/gemma4/text_decoder/convert_weights.py

@@ -21,11 +21,41 @@
 
 import torch
 
+from executorch.examples.models.checkpoint import get_mapped_key
+
 from .gemma4_config import Gemma4Config
 
 logger = logging.getLogger(__name__)
 
 
+# Weight mappings from Gemma 4's checkpoint to ExecuTorch's transformer parameters.
+_GEMMA4_TO_EXECUTORCH = {
+    "model.language_model.embed_tokens.weight": "tok_embeddings.weight",
+    "model.language_model.embed_tokens_per_layer.weight": "embed_tokens_per_layer.weight",
+    "model.language_model.per_layer_model_projection.weight": "per_layer_model_projection.weight",
+    "model.language_model.per_layer_projection_norm.weight": "per_layer_projection_norm.weight",
+    "model.language_model.norm.weight": "norm.weight",
+    "model.language_model.lm_head.weight": "output.weight",
+    "model.language_model.layers.{}.self_attn.q_proj.weight": "layers.{}.attention.wq.weight",
+    "model.language_model.layers.{}.self_attn.k_proj.weight": "layers.{}.attention.wk.weight",
+    "model.language_model.layers.{}.self_attn.v_proj.weight": "layers.{}.attention.wv.weight",
+    "model.language_model.layers.{}.self_attn.o_proj.weight": "layers.{}.attention.wo.weight",
+    "model.language_model.layers.{}.self_attn.q_norm.weight": "layers.{}.attention.q_norm_fn.weight",
+    "model.language_model.layers.{}.self_attn.k_norm.weight": "layers.{}.attention.k_norm_fn.weight",
+    "model.language_model.layers.{}.mlp.gate_proj.weight": "layers.{}.feed_forward.w1.weight",
+    "model.language_model.layers.{}.mlp.up_proj.weight": "layers.{}.feed_forward.w3.weight",
+    "model.language_model.layers.{}.mlp.down_proj.weight": "layers.{}.feed_forward.w2.weight",
+    "model.language_model.layers.{}.layer_scalar": "layers.{}.layer_scalar",
+    "model.language_model.layers.{}.per_layer_input_gate.weight": "layers.{}.per_layer_input_gate.weight",
+    "model.language_model.layers.{}.per_layer_projection.weight": "layers.{}.per_layer_projection.weight",
+    "model.language_model.layers.{}.input_layernorm.weight": "layers.{}.attention_norm.weight",
+    "model.language_model.layers.{}.post_attention_layernorm.weight": "layers.{}.post_attention_norm.weight",
+    "model.language_model.layers.{}.pre_feedforward_layernorm.weight": "layers.{}.ffn_norm.weight",
+    "model.language_model.layers.{}.post_feedforward_layernorm.weight": "layers.{}.post_ffn_norm.weight",
+    "model.language_model.layers.{}.post_per_layer_input_norm.weight": "layers.{}.post_per_layer_input_norm.weight",
+}
+
+
 def _download_manifold_file(manifold_path: str, local_path: Path) -> None:
     """Download a file from Manifold to a local path."""
     import io
@@ -279,6 +309,43 @@ def convert_hf_to_custom(
     return converted_state_dict
 
 
+def convert_weights(input_dir: str, output_file: str) -> None:
+    logger.info(f"Loading weights from {input_dir}")
+
+    hf_state_dict = _load_safetensors_weights(input_dir)
+
+    converted_state_dict = {}
+    unmapped_keys = []
+    for hf_key, tensor in hf_state_dict.items():
+        try:
+            qnn_key = get_mapped_key(hf_key, _GEMMA4_TO_EXECUTORCH)
+        except Exception:
+            unmapped_keys.append(hf_key)
+            logger.warning(f"Unmapped key: {hf_key}")
+            continue
+        converted_state_dict[qnn_key] = tensor.to(torch.float32)
+        logger.debug(f"Mapped {hf_key} -> {qnn_key}")
+
+    lm_head_key = "output.weight"
+    embed_tokens_key = "tok_embeddings.weight"
+    if (
+        lm_head_key not in converted_state_dict
+        and embed_tokens_key in converted_state_dict
+    ):
+        logger.info(
+            f"Using tied embeddings: copying {embed_tokens_key} to {lm_head_key}"
+        )
+        converted_state_dict[lm_head_key] = converted_state_dict[
+            embed_tokens_key
+        ].clone()
+
+    logger.info(f"Converted {len(converted_state_dict)} weights")
+    if unmapped_keys:
+        logger.warning(f"Unmapped keys: {len(unmapped_keys)}")
+
+    torch.save(converted_state_dict, output_file)
+
+
 def verify_conversion(
     hf_checkpoint_path: str,
     custom_state_dict: Dict[str, torch.Tensor],

examples/models/llama/evaluate/eager_eval.py

@@ -50,6 +50,8 @@ def eot_token_id(self):
 
     @property
     def prefix_token_id(self):
+        if hasattr(self._tokenizer, "bos_id"):
+            return self._tokenizer.bos_id
         return self.eot_token_id
 
     @property

examples/models/llama/model_args.py

@@ -11,6 +11,7 @@ class ActFn(Enum):
     SILU = "silu"
     GELU = "gelu"
     GELU_APPROX = "gelu_approx"
+    GELU_APPROX_TANH = "gelu_pytorch_tanh"
 
     @classmethod
     def from_string(cls, value: str) -> "ActFn":
@@ -29,7 +30,7 @@ def get_function(self):
             return F.silu
         elif self == ActFn.GELU:
             return F.gelu
-        elif self == ActFn.GELU_APPROX:
+        elif self == ActFn.GELU_APPROX or self == ActFn.GELU_APPROX_TANH:
             return partial(F.gelu, approximate="tanh")
         else:
             raise ValueError(f"Unsupported activation function: {self}")
@@ -172,6 +173,13 @@ class ModelArgs:
     # gemma2 attn and output soft capping
     final_logit_softcapping: Optional[float] = None
     attn_logit_softcapping: Optional[float] = None
+    # gemma4: head_dim override for full-attention layers (sliding layers keep head_dim).
+    global_head_dim: Optional[int] = None
+    # gemma4: per layer embedding
+    hidden_size_per_layer_input: Optional[int] = None
+    vocab_size_per_layer_input: Optional[int] = None
+    # gemma4: cross-decoder (KV-shared) layers use 2x hidden_dim in FeedForward.
+    use_double_wide_mlp: bool = False
 
     # rlformers forward-pass features for on-device model parity
     normalize_tok_embeddings: bool = False
@@ -225,3 +233,49 @@ def find_multiple(n: int, k: int) -> int:
         # Convert string act_fn to enum if needed
         if isinstance(self.act_fn, str):
             self.act_fn = ActFn.from_string(self.act_fn)
+
+    def get_layer_type(self, layer_idx: int) -> str:
+        return self.layer_types[layer_idx]
+
+    def is_sliding_attention(self, layer_idx: int) -> bool:
+        return (
+            self.layer_types and self.get_layer_type(layer_idx) == "sliding_attention"
+        )
+
+    def get_head_dim(self, layer_idx: int) -> int:
+        if not self.global_head_dim:
+            return self.head_dim
+        if self.is_sliding_attention(layer_idx):
+            return self.head_dim
+        return self.global_head_dim
+
+    def is_kv_shared_layer(self, layer_idx: int) -> bool:
+        first_kv_shared_layer_idx = self.n_layers - self.num_kv_shared_layers
+        return layer_idx >= first_kv_shared_layer_idx and first_kv_shared_layer_idx > 0
+
+    def get_intermediate_size(self, layer_idx: int) -> int:
+        if self.use_double_wide_mlp and self.is_kv_shared_layer(layer_idx):
+            return self.hidden_dim * 2
+        return self.hidden_dim
+
+    def get_kv_shared_layer_index(self, layer_idx: int) -> Optional[int]:
+        if not self.is_kv_shared_layer(layer_idx):
+            return None
+        first_kv_shared_layer_idx = self.n_layers - self.num_kv_shared_layers
+        current_layer_type = self.get_layer_type(layer_idx)
+        for i in range(first_kv_shared_layer_idx - 1, -1, -1):
+            if self.get_layer_type(i) == current_layer_type:
+                return i
+        return None
+
+    def is_kv_donor_layer(self, layer_idx: int) -> bool:
+        if self.is_kv_shared_layer(layer_idx):
+            return False
+        first_kv_shared_layer_idx = self.n_layers - self.num_kv_shared_layers
+        if first_kv_shared_layer_idx <= 0:
+            return False
+        current_layer_type = self.get_layer_type(layer_idx)
+        for i in range(first_kv_shared_layer_idx - 1, -1, -1):
+            if self.get_layer_type(i) == current_layer_type:
+                return layer_idx == i
+        return False

examples/qualcomm/oss_scripts/llama/README.md

@@ -14,6 +14,7 @@ This file provides you the instructions to run LLM Decoder model, VLM model, and
  1. Gemma 2B
  1. Gemma2 2B
  1. Gemma3 1B
+ 1. Gemma4 E2B
  1. GLM 1.5B
  1. Granite3.3 2B
  1. Phi4-mini-instruct
@@ -157,6 +158,12 @@ Default example using hybrid mode
 python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model gemma3-1b --prompt "I would like to learn python, could you teach me with a simple example?" --calib_tasks wikitext --calib_limit 1
 ```
 
+#### Gemma4 E2B
+Default example using hybrid mode
+```bash
+python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} --temperature 0 --model_mode hybrid --max_seq_len 1024 --prefill_ar_len 128 --decoder_model gemma4-e2b --prompt "I would like to learn python, could you teach me with a simple example?" --calib_tasks wikitext --calib_limit 1 --embedding-quantize '4,32'
+```
+
 #### GLM 1.5B
 Default example using hybrid mode
 ```bash

examples/qualcomm/oss_scripts/llama/__init__.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import json
 import os
 from abc import ABC
 from dataclasses import dataclass
@@ -17,6 +18,7 @@
 from executorch.examples.models.gemma import convert_weights as convert_gemma_weights
 from executorch.examples.models.gemma2 import convert_weights as convert_gemma2_weights
 from executorch.examples.models.gemma3 import convert_weights as convert_gemma3_weights
+from executorch.examples.models.gemma4 import convert_weights as convert_gemma4_weights
 
 from executorch.examples.models.glm import convert_weights as convert_glm_weights
 from executorch.examples.models.granite import (
@@ -59,16 +61,18 @@
     SmolVLMEncoder,
     VisionModalityConfig,
 )
-from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import (
+from executorch.examples.qualcomm.oss_scripts.llama.model import (
     LlamaModel,
     LlamaModelWithoutEmbedding,
+    ModelArgs,
     MultiScopeAwareLlamaModel,
 )
 
 from executorch.examples.qualcomm.oss_scripts.llama.static_llm_quant_recipe import (
     CodegenQuantRecipe,
     Gemma2QuantRecipe,
     Gemma3QuantRecipe,
+    Gemma4QuantRecipe,
     Gemma_2BQuantRecipe,
     GLM_1_5B_InstructQuantRecipe,
     Granite_3_3_2B_InstructQuantRecipe,
@@ -100,6 +104,7 @@
     "smolvlm_500m_instruct": LlamaModelWithoutEmbedding,
     "internvl3_1b": LlamaModelWithoutEmbedding,
     "gemma2-2b": MultiScopeAwareLlamaModel,
+    "gemma4-e2b": MultiScopeAwareLlamaModel,
 }
 
 
@@ -602,3 +607,56 @@ class SmolVLM_500M(LLMModelConfig):
     r2 = False
     r3 = False
     quant_recipe = SmolVLMQuantRecipe
+
+
+def load_gemma4_model_args(params_path: str) -> ModelArgs:
+    with open(params_path) as f:
+        config_dict = json.load(f)
+    text_config = config_dict.get("text_config", config_dict)
+
+    return ModelArgs(
+        dim=text_config["hidden_size"],
+        n_layers=text_config["num_hidden_layers"],
+        n_heads=text_config["num_attention_heads"],
+        n_kv_heads=text_config["num_key_value_heads"],
+        head_dim=text_config["head_dim"],
+        global_head_dim=text_config["global_head_dim"],
+        vocab_size=text_config["vocab_size"],
+        hidden_dim=text_config["intermediate_size"],
+        norm_eps=text_config["rms_norm_eps"],
+        max_position_embeddings=text_config["max_position_embeddings"],
+        rope_theta=text_config["rope_theta"],
+        local_rope_theta=text_config["rope_local_base_freq"],
+        partial_rotary_factor=text_config["partial_rotary_factor"],
+        sliding_window=text_config["sliding_window"],
+        layer_types=text_config["layer_types"],
+        final_logit_softcapping=text_config["final_logit_softcapping"],
+        hidden_size_per_layer_input=text_config["hidden_size_per_layer_input"],
+        vocab_size_per_layer_input=text_config["vocab_size_per_layer_input"],
+        num_kv_shared_layers=text_config["num_kv_shared_layers"],
+        use_double_wide_mlp=text_config["use_double_wide_mlp"],
+        post_attention_norm=text_config["post_attention_norm"],
+        post_ffn_norm=text_config["post_ffn_norm"],
+        use_ffn_norm=text_config["use_ffn_norm"],
+        act_fn=text_config["hidden_activation"],
+        model_architecture=config_dict["model_type"],
+    )
+
+
+@register_llm_model("gemma4-e2b")
+@dataclass(init=False, frozen=True)
+class Gemma4_E2B(LLMModelConfig):
+    repo_id: str = "google/gemma-4-e2b-it"
+    params_path: str = os.path.join(
+        BASE_DIR, "../../../models/gemma4/config/e2b_config.json"
+    )
+    convert_weights = convert_gemma4_weights
+    transform_weight = False
+    instruct_model = True
+    num_sharding = 4
+    masked_softmax = False
+    seq_mse_candidates = 0
+    r1 = False
+    r2 = False
+    r3 = False
+    quant_recipe = Gemma4QuantRecipe

examples/qualcomm/oss_scripts/llama/decoder_constants.py

@@ -63,4 +63,5 @@
     "glm-1_5b": "glm",
     "smolvlm_500m_instruct": "smolvlm",
     "internvl3_1b": "internvl3",
+    "gemma4-e2b": "gemma4",
 }

examples/qualcomm/oss_scripts/llama/encoder/encoder_config.py

@@ -13,10 +13,10 @@
     InternVL3EncoderQuantRecipe,
     SmolVLMEncoderQuantRecipe,
 )
-from executorch.examples.qualcomm.oss_scripts.llama.model.audio_encoder import (
+from executorch.examples.qualcomm.oss_scripts.llama.model.encoders.audio import (
     GraniteSpeechCTCEncoderWrapper,
 )
-from executorch.examples.qualcomm.oss_scripts.llama.model.vision_encoder import (
+from executorch.examples.qualcomm.oss_scripts.llama.model.encoders.vision import (
     Idefics3VisionEncoder,
     InternVL3VisionEncoder,
 )

examples/qualcomm/oss_scripts/llama/llama.py

@@ -666,6 +666,11 @@ def export_llama(args) -> None:
                 "Pass --calib_tasks to match the previous --tasks behavior."
             )
         raise RuntimeError("Please provide --eval_tasks to eval perplexity")
+    if args.decoder_model == "gemma4-e2b" and not args.embedding_quantize:
+        raise RuntimeError(
+            "gemma4-e2b requires --embedding-quantize: "
+            "per-layer embedding is too large for HTP quantization."
+        )
     assert (
         args.decoder_model in SUPPORTED_LLM_MODELS
     ), f"Unknown decoder_model: {args.decoder_model}."