[NeuralChat] enable lm_eval during training. (intel#1363)

* enable lm_eval during training.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: lkk12014402

intel_extension_for_transformers/llm/evaluation/lm_eval/evaluator.py

@@ -34,7 +34,7 @@
 MODEL_REGISTRY = {
     "hf-causal": huggingface.AutoCausalLM,
     "hf-seq2seq": huggingface.AutoSeq2SeqLM,
-
+    "simple-hf-causal": huggingface.HFModelAdapter,
 }
 
 def itrex_bootstrap_stderr(f, xs, iters):
@@ -69,6 +69,8 @@ def evaluate(model,
              output_base_path=None,
              seed=1234,
              user_model=None,
+             user_tokenizer=None,
+             warmup=False,
              model_format='torch'
             ):
     """Instantiate and evaluate a model on a list of tasks.
@@ -125,6 +127,18 @@ def evaluate(model,
             }
         if user_model:
             kwargs["init_empty_weights"] = True
+
+        if device == "hpu":
+            # if hpu, set user_model
+            kwargs["user_model"] = user_model
+            if model == "hf-causal":
+                model = "simple-hf-causal"
+        if model == "simple-hf-causal":
+            kwargs["warmup"] = warmup
+
+        if user_tokenizer:
+            kwargs["user_tokenizer"] = user_tokenizer
+
         lm = get_model(model).create_from_arg_string(
             model_args, kwargs
         )

intel_extension_for_transformers/llm/evaluation/lm_eval/models/huggingface.py

@@ -29,7 +29,7 @@
 from transformers import BatchEncoding
 
 from lm_eval import utils
-from lm_eval.base import BaseLM
+from lm_eval.base import BaseLM, CacheHook
 import re
 
 TokenSequence = Union[List[int], torch.LongTensor, torch.Tensor, BatchEncoding]
@@ -1078,3 +1078,80 @@ def stop_sequences_criteria(
             ],
         ]
     )
+
+
+class HFModelAdapter(HuggingFaceAutoLM):
+    AUTO_MODEL_CLASS = transformers.AutoModelForCausalLM
+    AUTO_PEFT_CLASS = peft.PeftModel
+
+    def __init__(self, *args, user_model=None, user_tokenizer=None, **kwargs):
+        self.cache_hook = CacheHook(None)
+        self.model = user_model
+        if user_tokenizer is None:
+            self.tokenizer = self._create_auto_tokenizer(
+                    pretrained=kwargs["pretrained"],
+                    revision="main",
+                    subfolder=None,)
+        else:
+            self.tokenizer = user_tokenizer
+        self._batch_size = kwargs["batch_size"]
+        self._add_special_tokens = None
+        self.model_format = "torch"
+        self.buckets = [16, 32, 64, 128, 189, 284]
+        self._device = kwargs["device"]
+        if self._device == "hpu":
+            from optimum.habana.checkpoint_utils import model_is_optimized # pylint: disable=E0611, E0401
+            self.static_shapes = model_is_optimized(self.model.config)
+        else:
+            self.static_shapes = False
+        if kwargs["warmup"]:
+            print("lm-eval warmup for Gaudi.")
+            self.warm_up()
+
+    def warm_up(self):
+        for bucket_size in reversed(self.buckets):
+            inps = torch.ones((self._batch_size, bucket_size), dtype=torch.int64)
+            self._model_call(inps)
+            pass
+
+    @property
+    def eot_token_id(self):
+        return self.model.config.eos_token_id
+
+    @property
+    def max_length(self):
+        return self.buckets[-1]
+
+    @property
+    def max_gen_toks(self):
+        raise NotImplementedError()
+
+    @property
+    def batch_size(self):
+        return self._batch_size
+
+    @property
+    def device(self):
+        # We need to do padding ourselves, otherwise we'll end up with recompilations
+        # Returning 'cpu' to keep tensors on CPU in lm_eval code
+        return "cpu"
+
+    def _model_generate(self, context, max_length, eos_token_id):
+        raise NotImplementedError()
+
+    def find_bucket(self, length):
+        return [b for b in self.buckets if b >= length][0]
+
+    def _model_call(self, inps):
+        bs, seq_length = inps.shape
+        padding_length = 0
+        if self.static_shapes:
+            bucket_length = self.find_bucket(seq_length)
+            padding_length = bucket_length - seq_length
+            inps = F.pad(inps, (0, padding_length), value=self.model.config.pad_token_id)
+        logits = self.model(inps.to(self._device))["logits"].cpu()
+
+        if self.static_shapes and padding_length > 0:
+            logits = logits[:, :-padding_length, :]
+        logits = logits.to(torch.float32)
+        return logits

intel_extension_for_transformers/llm/finetuning/eval_utils.py

@@ -25,6 +25,7 @@
 from transformers.trainer_utils import speed_metrics
 from transformers.debug_utils import DebugOption
 import math
+from transformers import TrainerCallback
 
 @torch.no_grad()
 def compute_rouge_metric(model, tokenizer, eval_dataset, training_args, gen_kwargs):
@@ -155,10 +156,39 @@ def evaluate_plus_ppl(
 
     output.metrics[f"{metric_key_prefix}_ppl"] = math.exp(output.metrics[f"{metric_key_prefix}_loss"])
 
-    self.log(output.metrics)
-
     self.control = self.callback_handler.on_evaluate(self.args, self.state, self.control, output.metrics)
 
+    self.log(output.metrics)
+
     self._memory_tracker.stop_and_update_metrics(output.metrics)
 
     return output.metrics
+
+
+class LMEvalCallback(TrainerCallback):
+    def __init__(self, lm_eval_func, device=None):
+        self.lm_eval = lm_eval_func
+        self.device = device
+        self.warmup = True
+
+    def on_evaluate(self, args, state, control, **kwargs):
+        if not state.is_local_process_zero:
+            return
+        if self.device == "hpu":
+            results = self.lm_eval(user_model=kwargs["model"],
+                    user_tokenizer=kwargs["tokenizer"],
+                    warmup=self.warmup)
+            self.warmup = False
+        else:
+            results = self.lm_eval(model="simple-hf-causal",
+                    user_model=kwargs["model"],
+                    user_tokenizer=kwargs["tokenizer"],
+                    warmup=False)
+        task_metrics = {}
+        for task_name in results["results"]:
+            for metric in results["results"][task_name]:
+                if "stderr" in metric:
+                    continue
+                metric_name = task_name + "_" + metric
+                task_metrics[metric_name] = results["results"][task_name][metric]
+        kwargs["metrics"].update(task_metrics)

intel_extension_for_transformers/llm/finetuning/finetuning.py

@@ -565,6 +565,21 @@ def concatenate_data(dataset, max_seq_length):
             if model_dtype == torch.bfloat16:
                 model = model.to(model_dtype)
 
+            lm_eval_callback = None
+            if training_args.do_eval and finetune_args.do_lm_eval:
+                from .eval_utils import LMEvalCallback
+                from functools import partial
+                from intel_extension_for_transformers.llm.evaluation.lm_eval import evaluate
+                lm_eval_func = partial(evaluate,
+                        model="hf-causal",
+                        model_args='pretrained='+model_args.model_name_or_path+\
+                                ',tokenizer='+model_args.model_name_or_path+',dtype=float16',
+                        device=finetune_args.device,
+                        batch_size=training_args.per_device_eval_batch_size,
+                        tasks=finetune_args.lm_eval_tasks,
+                        limit=data_args.max_eval_samples)
+                lm_eval_callback = LMEvalCallback(lm_eval_func, device=finetune_args.device)
+
             if finetune_args.device != 'hpu':
                 # Initialize our Trainer
                 trainer = Trainer(
@@ -574,6 +589,7 @@ def concatenate_data(dataset, max_seq_length):
                     eval_dataset=eval_dataset if training_args.do_eval else None,
                     tokenizer=tokenizer,
                     data_collator=data_collator,
+                    callbacks=[lm_eval_callback] if lm_eval_callback is not None else None
                 )
             else:
                 from optimum.habana import GaudiConfig, GaudiTrainer # pylint: disable=E0611 E0401
@@ -590,6 +606,7 @@ def concatenate_data(dataset, max_seq_length):
                     eval_dataset=eval_dataset if training_args.do_eval else None,
                     tokenizer=tokenizer,
                     data_collator=data_collator,
+                    callbacks=[lm_eval_callback] if lm_eval_callback is not None else None
                 )
 
             trainer.train(resume_from_checkpoint=training_args.resume_from_checkpoint)

intel_extension_for_transformers/neural_chat/config.py

@@ -349,7 +349,7 @@ class FinetuningArguments:
         metadata={"help": "whether to run the LM evaluation with EleutherAI/lm-evaluation-harness"},
     )
     lm_eval_tasks: Optional[List[str]] = field(
-        default_factory=lambda: ["truthfulqa_mc"],
+        default_factory=lambda: ["truthfulqa_mc", "lambada_openai"],
         metadata={"help": "tasks list for accuracy validation with EleutherAI/lm-evaluation-harness."},
     )
     qlora: bool = field(