Export lora weights to sep file

Differential Revision: [D83777195](https://our.internmc.facebook.com/intern/diff/D83777195/)

ghstack-source-id: 313704499
Pull Request resolved: #14756

Author: lucylq

examples/models/llama/export_llama_lib.py

@@ -1089,11 +1089,18 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager:  # noqa: C901
 
     if llm_config.backend.xnnpack.enabled:
         if llm_config.export.foundation_weights_file is not None:
-            gen_tag_fn: Callable[[torch.fx.Node], Optional[str]] = lambda x: (
-                llm_config.export.foundation_weights_file
-                if "lora" not in x.name
-                else None
-            )
+            if llm_config.export.lora_weights_file is not None:
+                gen_tag_fn: Callable[[torch.fx.Node], Optional[str]] = lambda x: (
+                    llm_config.export.foundation_weights_file
+                    if "lora" not in x.name
+                    else None
+                )
+            else:
+                gen_tag_fn: Callable[[torch.fx.Node], Optional[str]] = lambda x: (
+                    llm_config.export.foundation_weights_file
+                    if "lora" not in x.name
+                    else llm_config.export.lora_weights_file
+                )
 
             from executorch.exir.passes.external_constants_pass import (
                 delegate_external_constants_pass_unlifted,

extension/llm/export/config/llm_config.py

@@ -215,9 +215,10 @@ class ExportConfig:
         so_library: Shared library to specify custom quantized operators.
         export_only: Whether to stop right after torch.export() and
             just save the exported .pt2 graph file.
-        foundation_weights_file: configure the foundation weights of a model
-            to be placed in a separate file, external to the PTE. Pass the
-            intended file name here.
+        foundation_weights_file: place the foundation weights of the model into
+            a separate file, external to the PTE. Pass the file name here.
+        lora_weights_file: place the lora weights of the model into a
+            separate file, external to the PTE. Pass the file name here.
     """
 
     max_seq_length: int = 128
@@ -227,6 +228,7 @@ class ExportConfig:
     so_library: Optional[str] = None
     export_only: bool = False
     foundation_weights_file: Optional[str] = None
+    lora_weights_file: Optional[str] = None
 
     def __post_init__(self):
         if self.max_context_length < self.max_seq_length:
@@ -572,6 +574,8 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig":  # noqa: C901
             llm_config.export.export_only = args.export_only
         if hasattr(args, "foundation_weights_file"):
             llm_config.export.foundation_weights_file = args.foundation_weights_file
+        if hasattr(args, "lora_weights_file"):
+            llm_config.export.lora_weights_file = args.lora_weights_file
 
         # QuantizationConfig
         if hasattr(args, "quantization_mode"):