Skip to content

Commit 863039d

Browse files
Di Xu (SWE)facebook-github-bot
authored andcommitted
Fix export script to use pt2e if not selected
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 863039d

File tree

1 file changed

+8
-4
lines changed

1 file changed

+8
-4
lines changed

examples/models/llama/export_llama_lib.py

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -662,17 +662,21 @@ 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(args)
674+
builder_exported_to_edge = builder_exported.pt2e_quantize(
675+
quantizers
676+
).export_to_edge()
677+
else:
678+
pt2e_quant_params = None
679+
builder_exported_to_edge = builder_exported.export_to_edge()
676680

677681
modelname = builder_exported_to_edge.modelname
678682

0 commit comments

Comments
 (0)