Skip to content

Commit 4424925

Browse files
Di Xu (SWE)facebook-github-bot
authored andcommitted
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
1 parent ef2444f commit 4424925

File tree

1 file changed

+10
-4
lines changed

1 file changed

+10
-4
lines changed

examples/models/llama/export_llama_lib.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -662,17 +662,23 @@ def _validate_args(args):
662662

663663
def _export_llama(args) -> LLMEdgeManager: # noqa: C901
664664
_validate_args(args)
665-
pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params(args)
666665

667666
# export_to_edge
668667
builder_exported = _prepare_for_llama_export(args).export()
669668

670669
if args.export_only:
671670
exit()
672671

673-
builder_exported_to_edge = builder_exported.pt2e_quantize(
674-
quantizers
675-
).export_to_edge()
672+
if args.pt2e_quantize:
673+
pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params(
674+
args
675+
)
676+
builder_exported_to_edge = builder_exported.pt2e_quantize(
677+
quantizers
678+
).export_to_edge()
679+
else:
680+
pt2e_quant_params = None
681+
builder_exported_to_edge = builder_exported.export_to_edge()
676682

677683
modelname = builder_exported_to_edge.modelname
678684

0 commit comments

Comments
 (0)