From 3268da2adee51c219c745645dc7b631ab1b10e9a Mon Sep 17 00:00:00 2001 From: Mengtao Yuan Date: Sat, 7 Sep 2024 08:37:55 -0700 Subject: [PATCH] [eval_llama] Add option to save checkpoint after eager transforms. Differential Revision: D62150021 Pull Request resolved: https://github.com/pytorch/executorch/pull/5045 --- examples/models/llama2/eval_llama_lib.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/examples/models/llama2/eval_llama_lib.py b/examples/models/llama2/eval_llama_lib.py index 3ea4e77a1a..bd650fab1a 100644 --- a/examples/models/llama2/eval_llama_lib.py +++ b/examples/models/llama2/eval_llama_lib.py @@ -158,6 +158,15 @@ def gen_eval_wrapper( else manager.model.eval().to(device="cpu") ) + # Save the checkpoint after the eager model preparation is done. + # The reason for this option is that the checkpoint can be used + # to do evaluations in other evaluation platforms, or with data + # 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: + torch.save(model, args.output_eager_checkpoint_file) + return EagerEvalWrapper( model=model, tokenizer=tokenizer, @@ -196,6 +205,12 @@ def build_args_parser() -> argparse.ArgumentParser: default=None, help="[For ExecuTorch] Path to the Tokenizer binary for evaluating ExecuTorch models via runtime", ) + parser.add_argument( + "--output_eager_checkpoint_file", + type=str, + default=None, + help="Save the checkpoint after source transformations, for other evaluation platform to run the same checkpoint.", + ) return parser