Introduce hydra framework with backwards compatibility

Pull Request resolved: #11029




@imported-using-ghimport

Differential Revision: [D75263989](https://our.internmc.facebook.com/intern/diff/D75263989/)
ghstack-source-id: 288798835

Author: jackzhxng

examples/models/llama/TARGETS

@@ -82,6 +82,8 @@ runtime.python_binary(
     ],
     deps = [
         ":export_library",
+        ":export_llama_args",
+        ":export_llama_hydra",
         "//caffe2:torch",
         "//executorch/extension/pybindings:aten_lib",
     ],
@@ -148,6 +150,8 @@ runtime.python_library(
         ":source_transformation",
         "//ai_codesign/gen_ai/fast_hadamard_transform:fast_hadamard_transform",
         "//caffe2:torch",
+        "//executorch/examples/models/llama/config:llm_config",
+        "//executorch/backends/vulkan/_passes:vulkan_passes",
         "//executorch/exir/passes:init_mutable_pass",
         "//executorch/examples/models:model_base",
         "//executorch/examples/models:models",
@@ -230,6 +234,40 @@ runtime.python_library(
     ],
 )
 
+runtime.python_library(
+    name = "export_llama_args",
+    srcs = [
+        "export_llama_args.py",
+    ],
+    _is_external_target = True,
+    base_module = "executorch.examples.models.llama",
+    visibility = [
+        "//executorch/examples/...",
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        ":export_library",
+    ],
+)
+
+runtime.python_library(
+    name = "export_llama_hydra",
+    srcs = [
+        "export_llama_hydra.py",
+    ],
+    _is_external_target = True,
+    base_module = "executorch.examples.models.llama",
+    visibility = [
+        "//executorch/examples/...",
+        "@EXECUTORCH_CLIENTS",
+    ],
+    deps = [
+        ":export_library",
+        "//executorch/examples/models/llama/config:llm_config",
+        "fbsource//third-party/pypi/hydra-core:hydra-core",
+    ],
+)
+
 runtime.python_test(
     name = "quantized_kv_cache_test",
     srcs = [

examples/models/llama/config/llm_config.py

@@ -12,11 +12,12 @@
 Uses dataclasses, which integrate with OmegaConf and Hydra.
 """
 
+import argparse
 import ast
 import re
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import ClassVar, List, Optional
+from typing import ClassVar, List, Optional, Self
 
 
 ################################################################################
@@ -468,6 +469,18 @@ class LlmConfig:
     quantization: QuantizationConfig = field(default_factory=QuantizationConfig)
     backend: BackendConfig = field(default_factory=BackendConfig)
 
+    @staticmethod
+    def from_args(args: argparse.Namespace) -> Self:
+        """
+        To support legacy purposes, this function converts CLI args from
+        argparse to an LlmConfig, which is used by the LLM export process.
+        """
+        llm_config = LlmConfig()
+
+        # TODO: conversion code.
+
+        return llm_config
+
     def __post_init__(self):
         self._validate_low_bit()
 

examples/models/llama/export_llama.py

@@ -4,30 +4,50 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# Example script for exporting Llama2 to flatbuffer
-
-import logging
-
 # force=True to ensure logging while in debugger. Set up logger before any
 # other imports.
+import logging
+
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=FORMAT, force=True)
 
+import argparse
+import runpy
 import sys
 
 import torch
 
-from .export_llama_lib import build_args_parser, export_llama
-
 sys.setrecursionlimit(4096)
 
 
+def parse_hydra_arg():
+    """First parse out the arg for whether to use Hydra or the old CLI."""
+    parser = argparse.ArgumentParser(add_help=True)
+    parser.add_argument("--hydra", action="store_true")
+    args, remaining = parser.parse_known_args()
+    return args.hydra, remaining
+
+
 def main() -> None:
     seed = 42
     torch.manual_seed(seed)
-    parser = build_args_parser()
-    args = parser.parse_args()
-    export_llama(args)
+
+    use_hydra, remaining_args = parse_hydra_arg()
+    if use_hydra:
+        # The import runs the main function of export_llama_hydra with the remaining args
+        # under the Hydra framework.
+        sys.argv = [arg for arg in sys.argv if arg != "--hydra"]
+        print(f"running with {sys.argv}")
+        runpy.run_module(
+            "executorch.examples.models.llama.export_llama_hydra", run_name="__main__"
+        )
+    else:
+        # Use the legacy version of the export_llama script which uses argsparse.
+        from executorch.examples.models.llama.export_llama_args import (
+            main as export_llama_args_main,
+        )
+
+        export_llama_args_main(remaining_args)
 
 
 if __name__ == "__main__":

examples/models/llama/export_llama_args.py

@@ -0,0 +1,21 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Run export_llama with the legacy argparse setup.
+"""
+
+from .export_llama_lib import build_args_parser, export_llama
+
+
+def main(args) -> None:
+    parser = build_args_parser()
+    args = parser.parse_args(args)
+    export_llama(args)
+
+
+if __name__ == "__main__":
+    main()

examples/models/llama/export_llama_hydra.py

@@ -0,0 +1,27 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Run export_llama using the new Hydra CLI.
+"""
+
+import hydra
+
+from executorch.examples.models.llama.config.llm_config import LlmConfig
+from executorch.examples.models.llama.export_llama_lib import export_llama
+from hydra.core.config_store import ConfigStore
+
+cs = ConfigStore.instance()
+cs.store(name="llm_config", node=LlmConfig)
+
+
+@hydra.main(version_base=None, config_name="llm_config")
+def main(llm_config: LlmConfig) -> None:
+    export_llama(llm_config)
+
+
+if __name__ == "__main__":
+    main()

examples/models/llama/export_llama_lib.py

@@ -27,6 +27,10 @@
 from executorch.devtools.backend_debug import print_delegation_info
 
 from executorch.devtools.etrecord import generate_etrecord as generate_etrecord_func
+
+from executorch.examples.models.llama.config.llm_config import (
+    LlmConfig,
+)
 from executorch.examples.models.llama.hf_download import (
     download_and_convert_hf_checkpoint,
 )
@@ -50,6 +54,7 @@
     get_vulkan_quantizer,
 )
 from executorch.util.activation_memory_profiler import generate_memory_trace
+from omegaconf.dictconfig import DictConfig
 
 from ..model_factory import EagerModelFactory
 from .source_transformation.apply_spin_quant_r1_r2 import (
@@ -567,7 +572,23 @@ def canonical_path(path: Union[str, Path], *, dir: bool = False) -> str:
         return return_val
 
 
-def export_llama(args) -> str:
+def export_llama(
+    export_options: Union[argparse.Namespace, DictConfig],
+) -> str:
+    if isinstance(export_options, argparse.Namespace):
+        # Legacy CLI.
+        args = export_options
+        llm_config = LlmConfig.from_args(export_options)  # noqa: F841
+    elif isinstance(export_options, DictConfig):
+        # Hydra CLI.
+        llm_config = export_options  # noqa: F841
+    else:
+        raise ValueError(
+            "Input to export_llama must be either of type argparse.Namespace or LlmConfig"
+        )
+
+    # TODO: refactor rest of export_llama to use llm_config instead of args.
+
     # If a checkpoint isn't provided for an HF OSS model, download and convert the
     # weights first.
     if not args.checkpoint and args.model in HUGGING_FACE_REPO_IDS:

examples/models/llama/install_requirements.sh

@@ -10,7 +10,7 @@
 # Install tokenizers for hf .json tokenizer.
 # Install snakeviz for cProfile flamegraph
 # Install lm-eval for Model Evaluation with lm-evalution-harness.
-pip install huggingface_hub tiktoken torchtune sentencepiece tokenizers snakeviz lm_eval==0.4.5 blobfile
+pip install hydra-core huggingface_hub tiktoken torchtune sentencepiece tokenizers snakeviz lm_eval==0.4.5 blobfile
 
 # Call the install helper for further setup
 python examples/models/llama/install_requirement_helper.py