From 421848f717a3f845d48c76a2cc8f128d14a1e8d0 Mon Sep 17 00:00:00 2001 From: "Di Xu (SWE)" Date: Wed, 22 Jan 2025 14:58:49 -0800 Subject: [PATCH] Fix export script to use pt2e if not selected (#7854) Summary: Fix export script to use pt2e if not selected - This is useful for coreml export since it's not using pt2e quantizer Differential Revision: D68522589 --- examples/models/llama/export_llama_lib.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 6d4e1de0c0..73e4cc865f 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -662,7 +662,6 @@ def _validate_args(args): def _export_llama(args) -> LLMEdgeManager: # noqa: C901 _validate_args(args) - pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params(args) # export_to_edge builder_exported = _prepare_for_llama_export(args).export() @@ -670,9 +669,16 @@ def _export_llama(args) -> LLMEdgeManager: # noqa: C901 if args.export_only: exit() - builder_exported_to_edge = builder_exported.pt2e_quantize( - quantizers - ).export_to_edge() + if args.pt2e_quantize: + pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params( + args + ) + builder_exported_to_edge = builder_exported.pt2e_quantize( + quantizers + ).export_to_edge() + else: + pt2e_quant_params = None + builder_exported_to_edge = builder_exported.export_to_edge() modelname = builder_exported_to_edge.modelname