CPU QDQ SQNR support

Author: winskuo-quic

examples/qualcomm/oss_scripts/llama/decoder_constants.py

@@ -4,6 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# decode qdq file name
+DECODE_QDQ_FILENAME = "decode_qdq.pt2"
+
 # eval methods
 PROMPT_EVAL = "prompt_eval"
 TASKS_EVAL = "tasks_eval"

examples/qualcomm/oss_scripts/llama/decoder_runtime_evaluator.py

@@ -389,7 +389,7 @@ def run(self, prompt):
         self.inference_speed = output_performance_holder[0]
 
         sqnr = compute_error(golden_logits, output_logits_holder[0])
-        return sqnr.item()
+        return sqnr.item(), golden_logits, output_logits_holder[0]
 
 
 class TaskEval(EvalBase):

examples/qualcomm/oss_scripts/llama/llama.py

@@ -7,8 +7,6 @@
 # TODO: reenable pyre after fixing the issues
 # pyre-ignore-all-errors
 
-import copy
-import getpass
 import json
 import logging
 import os
@@ -17,26 +15,6 @@
 from typing import Dict
 
 import torch
-from executorch.backends.qualcomm._passes import FoldQDQ, TagQuantIO
-from executorch.backends.qualcomm._passes.i64_to_i32 import I64toI32
-from executorch.backends.qualcomm._passes.qnn_pass_manager import (
-    get_capture_program_passes,
-)
-from executorch.backends.qualcomm._passes.utils import (
-    get_passes_dependency_for_capture_program,
-)
-
-from executorch.backends.qualcomm.builders.utils import is_graph_output
-from executorch.backends.qualcomm.quantizer.custom_annotation import (
-    annotate_prefill_kv_output,
-)
-
-from executorch.backends.qualcomm.serialization.qc_schema import QcomChipset
-
-from executorch.backends.qualcomm.utils.constants import (
-    QCOM_PASS_ACTIVATE_KEY,
-    QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY,
-)
 from executorch.backends.qualcomm.utils.utils import (
     generate_htp_compiler_spec,
     generate_qnn_executorch_compiler_spec,
@@ -49,6 +27,7 @@
 from executorch.examples.qualcomm.oss_scripts.llama.dataset import DatasetBuilder
 from executorch.examples.qualcomm.oss_scripts.llama.decoder_constants import (
     AUDIO_ENCODER,
+    DECODE_QDQ_FILENAME,
     DECODER_GRAPH_NAMES,
     EVAL_MODE,
     PROMPT_EVAL,
@@ -65,33 +44,13 @@
     SqnrEval,
     TaskEval,
 )
-from executorch.examples.qualcomm.oss_scripts.llama.decoder_utils import (
-    graph_module_inference,
-)
-from executorch.examples.qualcomm.oss_scripts.llama.model.static_llama import (
-    LlamaModel,
-    ModelArgs,
-)
-from executorch.examples.qualcomm.oss_scripts.llama.range_setting_pt2e import (
-    make_custom_quantizer,
-)
-from executorch.examples.qualcomm.oss_scripts.llama.static_llm_quant_recipe import (
-    StaticLLMQuantRecipe,
-)
 from executorch.examples.qualcomm.oss_scripts.llama.tokenizer import TokenizerWrapper
 from executorch.examples.qualcomm.oss_scripts.llama.wrappers import (
     MultiModalManager,
     next_power_of_two,
 )
-
 from executorch.examples.qualcomm.utils import setup_common_args_and_variables
-from executorch.exir.capture._config import ExecutorchBackendConfig
-from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass
-from executorch.extension.llm.custom_ops import model_sharding
-from executorch.extension.llm.export.builder import DType
-from pytorch_tokenizers import get_tokenizer, TiktokenTokenizer
-from pytorch_tokenizers.llama2c import Llama2cTokenizer as SentencePieceTokenizer
+from torchao.quantization.utils import compute_error
 
 
 sys.setrecursionlimit(4096)
@@ -267,15 +226,41 @@ def inference(
             tokenizer=tokenizer,
             runtime_tokenizer_path=runtime_tokenizer_path,
         )
-        sqnr = sqnr_evaluator.run(prompt=prompt)
-        logging.info(f"SQNR Eval Score: {sqnr}")
+        sqnr, golden_logits, _ = sqnr_evaluator.run(prompt=prompt)
+        logging.info(f"SQNR Eval Score between FP32 nn.Module and QNN: {sqnr}")
         eval_results.update(
             {
                 "sqnr": sqnr,
                 "inference_speed": sqnr_evaluator.inference_speed,
             }
         )
 
+        qdq_ep_path = (
+            f"{args.pre_gen_pte}/{DECODE_QDQ_FILENAME}"
+            if args.pre_gen_pte
+            else f"{args.artifact}/{DECODE_QDQ_FILENAME}"
+        )
+        if os.path.exists(qdq_ep_path):
+            qdq_ep = torch.export.load(qdq_ep_path)
+            qdq_sqnr_evaluator = SqnrEval(
+                source_model=qdq_ep.module(),
+                example_input=source_model.get_example_inputs(),
+                args=args,
+                pte_path=pte_path,
+                tokenizer=tokenizer,
+                runtime_tokenizer_path=runtime_tokenizer_path,
+            )
+            qdq_sqnr, cpu_qdq_logits, _ = qdq_sqnr_evaluator.run(prompt=prompt)
+            eval_results["qdq_sqnr"] = qdq_sqnr
+            logging.info(f"SQNR Eval Score between CPU QDQ and QNN: {qdq_sqnr}")
+            logging.info(
+                f"SQNR Eval Score between FP32 nn.Module and CPU QDQ: {compute_error(golden_logits, cpu_qdq_logits).item()}"
+            )
+        else:
+            logging.info(
+                f"Couldn't find saved qdq_ep under {qdq_ep_path}, skip eval sqnr for CPU QDQ."
+            )
+
     if TASKS_EVAL in args.eval_methods:
         # Generate the eval wrapper
         ppl_evaluator = TaskEval(

examples/qualcomm/oss_scripts/llama/wrappers.py

@@ -51,6 +51,7 @@
 )
 from executorch.examples.qualcomm.oss_scripts.llama.decoder_constants import (
     AUDIO_ENCODER,
+    DECODE_QDQ_FILENAME,
     DECODER_GRAPH_NAMES,
     TEXT_DECODER,
     TEXT_EMBEDDING,
@@ -717,8 +718,17 @@ def quantize(self, request: Request):  # noqa: C901
                 tok_embedding=self.tok_embedding,
                 intermediate_outputs=image_embedding,
             )
-
             self.decoder = convert_pt2e(self.decoder)
+
+            # Saving Decode QDQ Model EP for SQNR evaluation
+            if self.mode == self.Mode.DECODE:
+                qdq_ep = torch.export.export(
+                    self.decoder, self.export_input, strict=True
+                )
+                qdq_ep_path = f"{self.control_args.artifact}/{DECODE_QDQ_FILENAME}"
+                torch.export.save(qdq_ep, qdq_ep_path)
+                logging.info(f"QDQ EP saved to {qdq_ep_path}")
+
             if self.apply_embedding:
                 self.tok_embedding = convert_pt2e(self.tok_embedding)