Skip to content

torch.export()-only export Llama arg #6695

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Nov 12, 2024
Merged

torch.export()-only export Llama arg #6695

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
c8b4d27
Export only llama arg
jackzhxng Nov 6, 2024
2230f91
Lint
jackzhxng Nov 6, 2024
a6a7162
Fix test
jackzhxng Nov 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 15 additions & 6 deletions examples/models/llama/export_llama_lib.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -443,6 +443,13 @@ def build_args_parser() -> argparse.ArgumentParser:
default=None,
help="path to the input pruning token mapping file (token_map.json)",
)

parser.add_argument(
"--export_only",
default=False,
action="store_true",
help="If true, stops right after torch.export() and saves the exported model.",
)
return parser


Expand Down Expand Up @@ -587,12 +594,14 @@ def _export_llama(modelname, args) -> LLMEdgeManager: # noqa: C901
pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params(args)

# export_to_edge
builder_exported_to_edge = (
_prepare_for_llama_export(modelname, args)
.export()
.pt2e_quantize(quantizers)
.export_to_edge()
)
builder_exported = _prepare_for_llama_export(modelname, args).export()

if args.export_only:
exit()

builder_exported_to_edge = builder_exported.pt2e_quantize(
quantizers
).export_to_edge()

modelname = builder_exported_to_edge.modelname

Expand Down
11 changes: 7 additions & 4 deletions extension/llm/export/builder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,22 +186,25 @@ def export(self) -> "LLMEdgeManager":
# functional graph. See issue https://github.com/pytorch/executorch/pull/4627 for more details
# pyre-fixme[8]: Attribute has type `Optional[GraphModule]`; used as
# `Module`.
self.pre_autograd_graph_module = torch.export.export(
exported_module = torch.export.export(
self.model,
self.example_inputs,
self.example_kwarg_inputs,
dynamic_shapes=dynamic_shape,
strict=True,
).module()
)
else:
# pyre-fixme[8]: Attribute has type `Optional[GraphModule]`; used as
# `Module`.
self.pre_autograd_graph_module = export_for_training(
exported_module = export_for_training(
self.model,
self.example_inputs,
kwargs=self.example_kwarg_inputs,
dynamic_shapes=dynamic_shape,
).module()
)
self.pre_autograd_graph_module = exported_module.module()
if hasattr(self.args, "export_only") and self.args.export_only:
torch.export.save(exported_module, self.args.output_name)

return self

Expand Down
Loading