Update on "Save foundation weights separately"

This diff:
1. Introduces SerializationConfig to llm_config. Currently, this allows user to save the foundation weights in a separate file; majorly useful for lora case.
2. Adds a pass to tag foundation (non-lora) weights. This is at the top-level (export_llama_lib). The tags are preserved through run_decomps/other passes, and do not affect functionality.
3. Tags are read when placing constants into the named_data_store.
4. Tagged weights are serialized to a separate file.


Notes
1. Adding tags to node.meta['custom']['blah'] means that they will not be discarded by run_decompositions
2. Adding tags to the lifted model (ep.graph_module) requires the EP to check is_param_node for xnnpack constants. Instead, add tags to the unlifted model (ep.module()), so we do not need to go through a re-export to get the EP.
3. Not an issue for this diff as llama doesn't have any higher order ops. Adding tags to models with higher-order ops is problematic due to nested submodules.

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

[ghstack-poisoned]

Author: lucylq

.ci/scripts/test_llama_lora.sh

@@ -108,7 +108,7 @@ $PYTHON_EXECUTABLE -m extension.llm.export.export_llm \
     backend.xnnpack.enabled=true \
     backend.xnnpack.extended_ops=true \
     export.output_name="${MODEL_SEPARATE}.pte" \
-    serialization.foundation_weights_file="${MODEL_SEPARATE}.ptd"
+    export.foundation_weights_file="${MODEL_SEPARATE}.ptd"
 
 # Run llama runner.
 NOW=$(date +"%H:%M:%S")

examples/models/llama/export_llama_lib.py

@@ -1078,9 +1078,9 @@ def _export_llama(llm_config: LlmConfig) -> LLMEdgeManager:  # noqa: C901
         llm_config.backend.xnnpack.enabled = True
 
     if llm_config.backend.xnnpack.enabled:
-        if llm_config.serialization.foundation_weights_file is not None:
+        if llm_config.export.foundation_weights_file is not None:
             gen_tag_fn: Callable[[torch.fx.Node], str] = lambda x: (
-                llm_config.serialization.foundation_weights_file
+                llm_config.export.foundation_weights_file
                 if "lora" not in x.name
                 else None
             )

extension/llm/export/config/llm_config.py

@@ -211,6 +211,9 @@ 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.
     """
 
     max_seq_length: int = 128
@@ -219,6 +222,7 @@ class ExportConfig:
     output_name: Optional[str] = None
     so_library: Optional[str] = None
     export_only: bool = False
+    foundation_weights_file: Optional[str] = None
 
     def __post_init__(self):
         if self.max_context_length < self.max_seq_length:
@@ -227,20 +231,6 @@ def __post_init__(self):
             )
 
 
-@dataclass
-class SerializationConfig:
-    """
-    Configures properties relevant to the serialization process.
-
-    Attributes:
-        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: Optional[str] = None
-
-
 ################################################################################
 ################################# DebugConfig ##################################
 ################################################################################
@@ -480,7 +470,6 @@ class LlmConfig:
     base: BaseConfig = field(default_factory=BaseConfig)
     model: ModelConfig = field(default_factory=ModelConfig)
     export: ExportConfig = field(default_factory=ExportConfig)
-    serialization: SerializationConfig = field(default_factory=SerializationConfig)
     debug: DebugConfig = field(default_factory=DebugConfig)
     quantization: QuantizationConfig = field(default_factory=QuantizationConfig)
     backend: BackendConfig = field(default_factory=BackendConfig)
@@ -560,12 +549,8 @@ def from_args(cls, args: argparse.Namespace) -> "LlmConfig":  # noqa: C901
             llm_config.export.so_library = args.so_library
         if hasattr(args, "export_only"):
             llm_config.export.export_only = args.export_only
-
-        # SerializationConfig
         if hasattr(args, "foundation_weights_file"):
-            llm_config.serialization.foundation_weights_file = (
-                args.foundation_weights_file
-            )
+            llm_config.export.foundation_weights_file = args.foundation_weights_file
 
         # QuantizationConfig
         if hasattr(args, "quantization_mode"):