Skip to content

Commit

Permalink
Fix export script to use pt2e if not selected (#7854)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
Di Xu (SWE) authored and facebook-github-bot committed Jan 22, 2025
1 parent ef2444f commit 4424925
Showing 1 changed file with 10 additions and 4 deletions.
14 changes: 10 additions & 4 deletions examples/models/llama/export_llama_lib.py
Original file line number Diff line number Diff line change
Expand Up @@ -662,17 +662,23 @@ 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()

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

Expand Down

0 comments on commit 4424925

Please sign in to comment.