From 4789d89b17e034b8d2abe9a9f7ed878d75bcb4cf Mon Sep 17 00:00:00 2001 From: lucylq Date: Thu, 2 Oct 2025 13:21:23 -0700 Subject: [PATCH] Export lora weights to sep file Differential Revision: [D83777195](https://our.internmc.facebook.com/intern/diff/D83777195/) [ghstack-poisoned] --- examples/models/llama/export_llama_lib.py | 17 ++++++++++++----- extension/llm/export/config/llm_config.py | 10 +++++++--- 2 files changed, 19 insertions(+), 8 deletions(-) diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index aa3b157c8da..80a479c827e 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/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, diff --git a/extension/llm/export/config/llm_config.py b/extension/llm/export/config/llm_config.py index b13001c005b..f15aad9e000 100644 --- a/extension/llm/export/config/llm_config.py +++ b/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"):