Revert "Revert "Qualcomm AI Engine Direct - improve llama3.2 3B TPS"" (#21150)

Reverts #21143

Author: metascroy

backends/qualcomm/debugger/utils.py

@@ -4,7 +4,6 @@
 import shutil
 import subprocess
 import tempfile
-from pathlib import Path
 from typing import Sequence, Tuple
 
 import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager
@@ -211,7 +210,7 @@ def __init__(
         self.adb = adb
         self.sample_input = sample_input
         self.build_folder = build_folder
-        self.root = str(Path(__file__).resolve().parents[3])
+        self.root = os.getcwd()
         self.config = {
             "backend_extension_config": {
                 "backend_extensions": {

backends/qualcomm/tests/models.py

@@ -1532,7 +1532,7 @@ def forward(self, x):
 class LargeTensorLinear(torch.nn.Module):
     def __init__(self):
         super().__init__()
-        hidden_dim = 8192
+        hidden_dim = 16384
         self.linear1_1 = torch.nn.Linear(512, hidden_dim)
         self.linear1_2 = torch.nn.Linear(512, hidden_dim)
         self.linear1_3 = torch.nn.Linear(512, hidden_dim)

backends/qualcomm/tests/test_qnn_delegate.py

@@ -10470,7 +10470,7 @@ def test_cli(self):
             cmds = [
                 "python",
                 "-m",
-                "examples.qualcomm.util_scripts.cli",
+                "executorch.examples.qualcomm.util_scripts.cli",
                 "quantize",
                 "--artifact",
                 f"{tmp_dir}/relu.pt2",
@@ -10489,7 +10489,7 @@ def test_cli(self):
             cmds = [
                 "python",
                 "-m",
-                "examples.qualcomm.util_scripts.cli",
+                "executorch.examples.qualcomm.util_scripts.cli",
                 "compile",
                 "--artifact",
                 f"{tmp_dir}/q_out/relu_quantized.pt2",
@@ -10507,7 +10507,7 @@ def test_cli(self):
             cmds = [
                 "python",
                 "-m",
-                "examples.qualcomm.util_scripts.cli",
+                "executorch.examples.qualcomm.util_scripts.cli",
                 "execute",
                 "--artifact",
                 f"{tmp_dir}/c_out/relu_quantized.pte",
@@ -10551,7 +10551,7 @@ def test_cli_with_input_list_assignment(self):
             cmds = [
                 "python",
                 "-m",
-                "examples.qualcomm.util_scripts.cli",
+                "executorch.examples.qualcomm.util_scripts.cli",
                 "quantize",
                 "--artifact",
                 f"{tmp_dir}/sub.pt2",
@@ -10570,7 +10570,7 @@ def test_cli_with_input_list_assignment(self):
             cmds = [
                 "python",
                 "-m",
-                "examples.qualcomm.util_scripts.cli",
+                "executorch.examples.qualcomm.util_scripts.cli",
                 "compile",
                 "--artifact",
                 f"{tmp_dir}/q_out/sub_quantized.pt2",
@@ -10588,7 +10588,7 @@ def test_cli_with_input_list_assignment(self):
             cmds = [
                 "python",
                 "-m",
-                "examples.qualcomm.util_scripts.cli",
+                "executorch.examples.qualcomm.util_scripts.cli",
                 "execute",
                 "--artifact",
                 f"{tmp_dir}/c_out/sub_quantized.pte",

examples/qualcomm/oss_scripts/llama/__init__.py

@@ -281,7 +281,7 @@ class Llama3_2_3B_Instruct(LLMModelConfig):
     transform_weight = True
     # The Llama3_2 enabled should be instruct, however, Llama's tokenizer does not provide utility to apply chat template.
     instruct_model = False
-    num_sharding = 4
+    num_sharding = 3
     masked_softmax = False
     seq_mse_candidates = 0
     r1 = False

examples/qualcomm/oss_scripts/llama/model/static_llama.py

@@ -730,7 +730,7 @@ def forward(
             self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin
         )
 
-        hidden_states = self.embedding_scale_factor * self.tok_embeddings(tokens)
+        hidden_states = self.tok_embeddings(tokens)
 
         for ind, decoder_layer in enumerate(self.layers):
             k_caches = None
@@ -865,7 +865,11 @@ def forward(
             self.freqs_sin[input_pos][0] if self.use_kv_cache else self.freqs_sin
         )
 
-        hidden_states = self.embedding_scale_factor * hidden_states
+        hidden_states = (
+            self.embedding_scale_factor * hidden_states
+            if self.embedding_scale_factor != 1.0
+            else hidden_states
+        )
 
         for ind, decoder_layer in enumerate(self.layers):
             k_caches = None
@@ -1024,7 +1028,8 @@ def forward(
             else self.local_freqs_sin
         )
 
-        hidden_states = self.embedding_scale_factor * self.tok_embeddings(tokens)
+        hidden_states = self.tok_embeddings(tokens)
+
         for ind, decoder_layer in enumerate(self.layers):
             k_caches = None
             v_caches = None

examples/qualcomm/oss_scripts/llama/wrappers/llm_wrappers.py

@@ -241,6 +241,9 @@ def _prepare_model(self):  # noqa: C901
                     # Llama does x.to(float16) * w whilst Gemma3 is (x * w).to(float16)
                     # See https://github.com/huggingface/transformers/pull/29402
                     state_dict[k] = v.float() + torch.ones(v.shape, dtype=torch.float32)
+            for k, v in state_dict.items():
+                if "tok_embeddings.weight" == k:
+                    state_dict[k] = v.float() * self.model_args.embedding_scale_factor
         else:
             state_dict = torch.load(
                 self.control_args.checkpoint,
@@ -466,9 +469,7 @@ def _tag_ios(self, node, fixed_point_type):
             (self.meta["get_ar_len"], self.meta["get_head_dim"] // 2),
         }
 
-        freq_op = {
-            exir_ops.edge.aten.select.int,
-        }
+        freq_op = {exir_ops.edge.aten.select.int, exir_ops.edge.aten.select_copy.int}
         quant_io_type = None
 
         if node.op == "placeholder":

examples/qualcomm/util_scripts/cli.py

@@ -166,6 +166,8 @@ def quantize(args):
             quant_dtype=quant_dtype,
             per_channel_conv=args.per_channel,
             per_channel_linear=args.per_row,
+            per_channel_embedding=args.per_channel_embedding,
+            act_symmetric=args.act_symmetric,
             act_observer=act_observer,
             backend=get_backend_type(args.backend),
             soc_model=args.soc_model,
@@ -499,6 +501,16 @@ def main():
         action="store_true",
         help="Use per_row encoding for operator linear.",
     )
+    sub_quantize.add_argument(
+        "--per_channel_embedding",
+        action="store_true",
+        help="Use per_channel encoding for operator embedding.",
+    )
+    sub_quantize.add_argument(
+        "--act_symmetric",
+        action="store_true",
+        help="Use symmetric quantization for activations.",
+    )
     sub_quantize.add_argument(
         "--activation_observer",
         type=str,