Skip to content

Commit def8b27

Browse files
cccclaifacebook-github-bot
authored andcommitted
add qualcomm option
Summary: Add qualcomm option, including both qnn_quantizer and qnn_partitioner Differential Revision: D55218789
1 parent 14e31f0 commit def8b27

File tree

1 file changed

+66
-0
lines changed

1 file changed

+66
-0
lines changed

examples/models/llama2/export_llama_lib.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -420,6 +420,11 @@ def build_args_parser() -> argparse.ArgumentParser:
420420
parser.add_argument("-V", "--vulkan", action="store_true")
421421
parser.add_argument("--mps", action="store_true")
422422
parser.add_argument("--coreml", action="store_true")
423+
parser.add_argument(
424+
"--qualcomm",
425+
action="store_true",
426+
help="Delegate llama2 to Qualcomm backend, please use it with ----quantization_mode pt2e and --kv_cahce=True",
427+
)
423428

424429
parser.add_argument(
425430
"--expand_rope_table",
@@ -553,6 +558,23 @@ def _export_llama(modelname, args) -> str: # noqa: C901
553558
# export_to_edge
554559
pt2e_quant_params = _get_pt2e_quantization_params(args)
555560
quantizers = get_pt2e_quantizers(pt2e_quant_params, args)
561+
if args.qualcomm:
562+
# reset quantizers and pt2e_quant_params from xnnpack backend
563+
pt2e_quant_params = None
564+
quantizers = []
565+
try:
566+
# pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.quantizer.quantizer`
567+
from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer
568+
except ImportError:
569+
raise ImportError(
570+
"Please install the Qualcomm backend follwing https://pytorch.org/executorch/main/build-run-qualcomm.html"
571+
)
572+
573+
# pyre-ignore: Undefined attribute [16]: Module `executorch.backends` has no attribute `qualcomm`.
574+
qnn_quantizer = QnnQuantizer()
575+
custom_annotations = ()
576+
qnn_quantizer.add_custom_quant_annotations(custom_annotations)
577+
quantizers.append(qnn_quantizer)
556578

557579
builder_exported_to_edge = _prepare_for_llama_export(
558580
modelname, args
@@ -636,6 +658,50 @@ def _export_llama(modelname, args) -> str: # noqa: C901
636658
)
637659
modelname = f"coreml_{modelname}"
638660

661+
if args.qualcomm:
662+
assert (
663+
args.use_kv_cache is True
664+
), "Qualcomm backend currently only supports static shape and use_kv_cache=True is the only way to support it at the moment"
665+
try:
666+
# pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.partition.qnn_partitioner`
667+
from executorch.backends.qualcomm.partition.qnn_partitioner import (
668+
QnnPartitioner,
669+
)
670+
671+
# pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.serialization.qnn_compile_spec_schema`
672+
from executorch.backends.qualcomm.serialization.qnn_compile_spec_schema import (
673+
QcomChipset,
674+
)
675+
676+
# pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `executorch.backends.qualcomm.utils.utils`
677+
from executorch.backends.qualcomm.utils.utils import (
678+
_transform,
679+
generate_htp_compiler_spec,
680+
generate_qnn_executorch_compiler_spec,
681+
)
682+
except ImportError:
683+
raise ImportError(
684+
"Please install the Qualcomm backend follwing https://pytorch.org/executorch/main/build-run-qualcomm.html"
685+
)
686+
687+
# pyre-ignore: Undefined attribute [16]: Module `executorch.backends` has no attribute `qualcomm`
688+
backend_options = generate_htp_compiler_spec(use_fp16=False)
689+
# pyre-ignore: Undefined attribute [16]: Module `executorch.backends` has no attribute `qualcomm`
690+
partitioner = QnnPartitioner(
691+
# pyre-ignore: Undefined attribute [16]: Module `executorch.backends` has no attribute `qualcomm`
692+
generate_qnn_executorch_compiler_spec(
693+
# pyre-ignore: Undefined attribute [16]: Module `executorch.backends` has no attribute `qualcomm`.
694+
soc_model=QcomChipset.SM8650,
695+
backend_options=backend_options,
696+
debug=False,
697+
saver=False,
698+
),
699+
skip_node_id_set={},
700+
skip_node_op_set={},
701+
)
702+
# pyre-ignore: Undefined attribute [16]: Module `executorch.backends` has no attribute `qualcomm`
703+
_transform(builder_exported_to_edge.export_program())
704+
639705
if args.generate_etrecord:
640706
if not builder_exported_to_edge.edge_manager:
641707
raise ValueError("Unable to generate etrecord due to missing edge manager.")

0 commit comments

Comments
 (0)