Export only llama arg

Author: jackzhxng

examples/models/llama/export_llama_lib.py

@@ -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
 
 
@@ -587,9 +594,16 @@ 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 = (
+    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()
     )

extension/llm/export/builder.py

@@ -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 self.args.export_only:
+                torch.export.save(exported_module, self.args.output_name)
 
         return self