Convert Args to DictConfig (#9717)

Summary:
Pull Request resolved: #9717

1. Add structure to the complicated args
2. Convert args to DictConfig, to decouple the cli args
3. Pass needed sub configs to functions (instead of args)

Pull Request resolved: #9450

Test Plan:
```
python3 -m examples.models.llama.export_llama -c stories110M.pt -p params.json -d fp32 -n tinyllama_xnnpack+custom_fp32_main.pte -kv -X --xnnpack-extended-ops -qmode 8da4w -G 128 --use_sdpa_with_kv_cache --output-dir tmp -E 8,64
```

Reviewed By: larryliu0820

Differential Revision: D71557301

Pulled By: jackzhxng

Author: Martin Yuan

examples/models/llama/eval_llama.py

@@ -11,6 +11,7 @@
 import torch
 
 from .eval_llama_lib import (
+    _convert_cli_to_config_format,
     build_args_parser,
     eval_llama,
     eval_llama_with_attention_sink,
@@ -28,10 +29,11 @@ def main() -> None:
     args = parser.parse_args()
     # Overrides this arg, because evaluation requires full logits.
     args.generate_full_logits = True
+    config = _convert_cli_to_config_format(args)
     if args.use_attention_sink:
-        eval_llama_with_attention_sink(modelname, args)  # pyre-ignore
+        eval_llama_with_attention_sink(modelname, config)  # pyre-ignore
     else:
-        eval_llama(modelname, args)  # pyre-ignore
+        eval_llama(modelname, config)  # pyre-ignore
 
 
 if __name__ == "__main__":

examples/models/llama/eval_llama_lib.py

@@ -4,35 +4,27 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-
 import argparse
-
 from typing import Optional, Union
 
 import torch
-
-from datasets import load_dataset
 from executorch.examples.models.llama.export_llama_lib import (
+    _convert_args_to_config,
+    _prepare_for_llama_export,
+    build_args_parser as _build_args_parser,
     get_quantizer_and_quant_params,
 )
 from executorch.examples.models.llama.tokenizer.tiktoken import Tokenizer as Tiktoken
-
 from executorch.extension.llm.export.builder import LLMEdgeManager
 from executorch.extension.llm.tokenizer.tokenizer import (
     Tokenizer as SentencePieceTokenizer,
 )
 from executorch.extension.llm.tokenizer.utils import get_tokenizer
 from lm_eval.evaluator import simple_evaluate
-from torch.nn import CrossEntropyLoss
-from tqdm import tqdm
+from omegaconf import DictConfig, OmegaConf
 
 from .evaluate.eager_eval import EagerEvalWrapper
 
-from .export_llama_lib import (
-    _prepare_for_llama_export,
-    build_args_parser as _build_args_parser,
-)
-
 
 class GraphModuleEvalWrapper(EagerEvalWrapper):
     """
@@ -165,7 +157,7 @@ def _model_call(self, inps):
 
 def gen_eval_wrapper(
     model_name: str,
-    args: argparse.ArgumentParser,
+    config: DictConfig,
 ):
     """
     Generates a wrapper interface around the provided model and tokenizer for
@@ -174,17 +166,17 @@ def gen_eval_wrapper(
     Returns:
         eval_wrapper (LM): A wrapper interface for the lm-evaluation-harness library.
     """
-    tokenizer = get_tokenizer(args.tokenizer_path)  # pyre-ignore
+    tokenizer = get_tokenizer(config.model.tokenizer_path)
 
     # ExecuTorch Binary Evaluation
-    if (model := args.pte) is not None:  # pyre-ignore
-        if (tokenizer_bin := args.tokenizer_bin) is not None:  # pyre-ignore
+    if (model := config.eval.pte) is not None:
+        if (tokenizer_bin := config.eval.tokenizer_bin) is not None:
             # ETRunnerEvalWrapper: Create a wrapper around an ExecuTorch model, evaluated at runtime
             return ETRunnerEvalWrapper(
                 model=model,
                 tokenizer=tokenizer,
                 tokenizer_bin=tokenizer_bin,
-                max_seq_length=args.max_seq_length,  # pyre-ignore
+                max_seq_length=config.sequence.max_seq_length,
             )
 
         # ETPybindEvalWrapper: Create a wrapper around an ExecuTorch model, evaluated with pybindings
@@ -193,12 +185,12 @@ def gen_eval_wrapper(
             tokenizer=tokenizer,
             # Exported model takes at most (max_seq_length - 1) tokens.
             # Note that the eager model takes at most max_seq_length tokens.
-            max_seq_length=args.max_seq_length - 1,
+            max_seq_length=config.sequence.max_seq_length - 1,
         )
 
-    pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params(args)
+    pt2e_quant_params, quantizers, quant_dtype = get_quantizer_and_quant_params(config)
     # GPTFastEvalWrapper: Create a wrapper around a pre-exported model
-    manager: LLMEdgeManager = _prepare_for_llama_export(args)
+    manager: LLMEdgeManager = _prepare_for_llama_export(config)
 
     if len(quantizers) != 0:
         manager = manager.export().pt2e_quantize(quantizers)
@@ -210,9 +202,9 @@ def gen_eval_wrapper(
         return GraphModuleEvalWrapper(
             model=model,
             tokenizer=tokenizer,
-            max_seq_length=args.max_seq_length,
-            use_kv_cache=args.use_kv_cache,  # pyre-ignore
-            enable_dynamic_shape=args.enable_dynamic_shape,  # pyre-ignore
+            max_seq_length=config.sequence.max_seq_length,
+            use_kv_cache=config.kv_cache.use_kv_cache,  # pyre-ignore
+            enable_dynamic_shape=config.misc.enable_dynamic_shape,  # pyre-ignore
         )
     else:
         # TODO: use manager.pre_autograd_graph_module for the eval to remove the if-else branch
@@ -230,18 +222,94 @@ def gen_eval_wrapper(
         # that is not available in this eval_llama. We save the checkpoint
         # here for consistency with eval_llama. The accuracy results we
         # get from eval_llama can be used as a reference to other evaluations.
-        if args.output_eager_checkpoint_file is not None:  # pyre-ignore
-            torch.save(model, args.output_eager_checkpoint_file)
+        if config.eval.output_eager_checkpoint_file is not None:  # pyre-ignore
+            torch.save(model, config.eval.output_eager_checkpoint_file)
 
         return EagerEvalWrapper(
             model=model,
             tokenizer=tokenizer,
-            max_seq_length=args.max_seq_length,
-            use_kv_cache=args.use_kv_cache,
+            max_seq_length=config.sequence.max_seq_length,
+            use_kv_cache=config.kv_cache.use_kv_cache,
+        )
+
+
+def eval_llama(
+    model_name: str,
+    config: DictConfig,
+) -> None:
+    # Generate the eval wrapper
+    eval_wrapper = gen_eval_wrapper(model_name, config)
+
+    # Needed for loading mmlu dataset.
+    # See https://github.com/EleutherAI/lm-evaluation-harness/pull/1998/files
+    if config.eval.tasks and "mmlu" in config.eval.tasks:
+        import datasets
+
+        datasets.config.HF_DATASETS_TRUST_REMOTE_CODE = True
+
+    # Evaluate the model
+    tasks = (
+        None if config.eval.tasks is None else OmegaConf.to_container(config.eval.tasks)
+    )
+    with torch.no_grad():
+        eval_results = simple_evaluate(
+            model=eval_wrapper,
+            tasks=tasks,
+            num_fewshot=config.eval.num_fewshot,
+            limit=config.eval.limit,
+        )
+
+    for task, res in eval_results["results"].items():
+        print(f"{task}: {res}")
+
+
+def eval_llama_with_attention_sink(
+    model_name: str,
+    config: DictConfig,
+) -> None:
+    # Generate the eval wrapper
+    eval_wrapper = gen_eval_wrapper(model_name, config)
+
+    # Needed for loading mmlu dataset.
+    # See https://github.com/EleutherAI/lm-evaluation-harness/pull/1998/files
+    if config.eval.tasks and "mmlu" in config.eval.tasks:
+        import datasets
+
+        datasets.config.HF_DATASETS_TRUST_REMOTE_CODE = True
+
+    # Evaluate the model
+    with torch.no_grad():
+        eval_results = simple_evaluate(
+            model=eval_wrapper,
+            tasks=OmegaConf.to_container(config.eval.tasks),
+            num_fewshot=config.eval.num_fewshot,
+            limit=config.eval.limit,
         )
 
+    for task, res in eval_results["results"].items():
+        print(f"{task}: {res}")
+
+
+def _convert_cli_to_config_format(args) -> DictConfig:
+    """Convert CLI arguments to config format."""
+    # First convert common args using the shared function
+    config = _convert_args_to_config(args)
+
+    # Add evaluation-specific settings
+    config.eval = OmegaConf.create()
+    config.eval.tasks = args.tasks
+    config.eval.limit = args.limit
+    config.eval.num_fewshot = args.num_fewshot
+    config.eval.pte = args.pte
+    config.eval.tokenizer_bin = args.tokenizer_bin
+    config.eval.output_eager_checkpoint_file = args.output_eager_checkpoint_file
+    config.eval.attention_sink_eval_tokens = args.attention_sink_eval_tokens
+
+    return config
+
 
 def build_args_parser() -> argparse.ArgumentParser:
+    """Build argument parser for evaluation, extending the export parser with eval-specific args."""
     # Start with arg parser from export_llama_lib
     parser = _build_args_parser()
 
@@ -288,92 +356,7 @@ def build_args_parser() -> argparse.ArgumentParser:
         help="Save the checkpoint after source transformations, for other evaluation platform to run the same checkpoint.",
     )
 
-    # Set of parameters secpific to AttentionSink.
+    # Set of parameters specific to AttentionSink.
     parser.add_argument("--attention_sink_eval_tokens", type=int, default=0)
 
     return parser
-
-
-def eval_llama(
-    model_name: str,
-    args: argparse.ArgumentParser,
-) -> None:
-    # Generate the eval wrapper
-    eval_wrapper = gen_eval_wrapper(model_name, args)
-
-    # Needed for loading mmlu dataset.
-    # See https://github.com/EleutherAI/lm-evaluation-harness/pull/1998/files
-    # pyre-ignore: Undefined attribute [16]: `argparse.ArgumentParser` has no attribute `tasks`
-    if args.tasks and "mmlu" in args.tasks:
-        import datasets
-
-        datasets.config.HF_DATASETS_TRUST_REMOTE_CODE = True
-
-    # Evaluate the model
-    with torch.no_grad():
-        eval_results = simple_evaluate(
-            model=eval_wrapper,
-            tasks=args.tasks,
-            num_fewshot=args.num_fewshot,  # pyre-ignore: Undefined attribute [16]: `argparse.ArgumentParser` has no attribute `num_fewshot`
-            limit=args.limit,  # pyre-ignore: Undefined attribute [16]: `argparse.ArgumentParser` has no attribute `limit`
-        )
-
-    for task, res in eval_results["results"].items():
-        print(f"{task}: {res}")
-
-
-def eval_llama_with_attention_sink(model_name: str, args: argparse.ArgumentParser):
-    """
-    Evaluate the model's perplexity when AttentionSink is enabled.
-
-    This is mostly copied from https://github.com/mit-han-lab/streaming-llm/blob/main/examples/eval_long_ppl.py
-    """
-    assert args.use_attention_sink is not None  # pyre-ignore [16]
-    assert args.attention_sink_eval_tokens > 0  # pyre-ignore [16]
-    attention_sink_params = args.use_attention_sink.split(",")
-    assert len(attention_sink_params) == 3
-    sink_size = int(attention_sink_params[0])
-    window_size = int(attention_sink_params[1])
-
-    assert args.max_seq_length == sink_size + window_size  # pyre-ignore [16]
-
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    manager: LLMEdgeManager = _prepare_for_llama_export(args)
-    model = manager.model.eval().to(device=device)
-    tokenizer = get_tokenizer(args.tokenizer_path)  # pyre-ignore [16]
-
-    eval_data = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
-
-    nlls = []
-    loss_fn = CrossEntropyLoss(reduction="none")
-    progress_bar = tqdm(total=args.attention_sink_eval_tokens)
-    input_pos = 0
-    while input_pos < args.attention_sink_eval_tokens:
-        for text in eval_data["text"]:  # pyre-ignore [16]
-            tokens = tokenizer.encode(text, bos=False, eos=False)
-            if len(tokens) <= 0:
-                continue
-            with torch.no_grad():
-                num_tokens = min(
-                    len(tokens) - 1, args.attention_sink_eval_tokens - input_pos
-                )
-                logits = model(
-                    torch.tensor(
-                        [tokens[:num_tokens]], dtype=torch.int64, device=device
-                    ),
-                    torch.tensor([input_pos], dtype=torch.int64, device=device),
-                ).squeeze(dim=0)
-                neg_log_likelihood = loss_fn(
-                    logits,
-                    torch.tensor(
-                        [tokens[1 : num_tokens + 1]], dtype=torch.int64, device=device
-                    ).view(-1),
-                )
-                nlls.append(neg_log_likelihood)
-                input_pos += num_tokens
-                progress_bar.update(num_tokens)
-            if input_pos >= args.attention_sink_eval_tokens:
-                break
-    ppl = torch.exp(torch.cat(nlls).mean())
-    print(f"Perplexity: {ppl.item()}")
-    return ppl.item()