add llama bench (#6119)

Author: linjieccc

examples/language_model/llama/data.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import copy
+import json
 from dataclasses import dataclass
 from typing import Dict, List
 
@@ -22,6 +23,26 @@
 
 IGNORE_INDEX = -100
 
+PROMPT_DICT = {
+    "prompt_input": (
+        "Below is an instruction that describes a task, paired with an input that provides further context. "
+        "Write a response that appropriately completes the request.\n\n"
+        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:"
+    ),
+    "prompt_no_input": (
+        "Below is an instruction that describes a task. "
+        "Write a response that appropriately completes the request.\n\n"
+        "### Instruction:\n{instruction}\n\n### Response:"
+    ),
+}
+
+
+def reader(data_path):
+    with open(data_path, "r", encoding="utf-8") as f:
+        for line in f:
+            json_line = json.loads(line)
+            yield json_line
+
 
 def convert_example(example, tokenizer, data_args, is_test=False):
     """
@@ -81,40 +102,57 @@ def convert_example(example, tokenizer, data_args, is_test=False):
     )
 
 
-def custom_instruction_convert_example(example, tokenizer, data_args, is_test=False):
+def custom_instruction_convert_example(
+    example, tokenizer, data_args, is_test=False, benchmark=False, model_max_length=512
+):
     """
     Convert an example into necessary features.
     """
 
-    instruction = ""
-    input = ""
-    output = ""
-    if "instruction" in example and "output" in example:
-        instruction = example["instruction"]
-        output = example["output"]
-    else:
-        assert False, "instruction and output are not in the input dictionary."
-    if "input" in example["input"]:
-        input = example["input"]
+    if benchmark:
+        prompt_input, prompt_no_input = PROMPT_DICT["prompt_input"], PROMPT_DICT["prompt_no_input"]
 
-    input_seq = instruction + input
-    output_seq = output
+        if example.get("input", "") != "":
+            input_seq = prompt_input.format_map(example)
+        else:
+            input_seq = prompt_no_input.format_map(example)
 
+        output_seq = example["output"]
+    else:
+        instruction = ""
+        input = ""
+        output = ""
+        if "instruction" in example and "output" in example:
+            instruction = example["instruction"]
+            output = example["output"]
+        else:
+            assert False, "instruction and output are not in the input dictionary."
+        if "input" in example["input"]:
+            input = example["input"]
+
+        input_seq = instruction + input
+        output_seq = output
+
+    # To compatible with compile training mode in benchmark, input will be pad to fix length
     source_tokenized = tokenizer(
         input_seq,
         return_tensors="pd",
-        max_length=data_args.src_length,
+        padding="loggest" if not benchmark else "max_length",
+        max_length=data_args.src_length if not benchmark else model_max_length,
         truncation=True,
     )
 
     source_input_ids_len = (
         source_tokenized["input_ids"].not_equal(paddle.to_tensor(tokenizer.pad_token_id)).sum().item()
     )
 
+    total_length = data_args.src_length + data_args.tgt_length
+
     example_tokenized = tokenizer(
         input_seq + output_seq,
         return_tensors="pd",
-        max_length=data_args.src_length + data_args.tgt_length,
+        padding="loggest" if not benchmark else "max_length",
+        max_length=total_length if not benchmark else model_max_length,
         truncation=True,
     )
 
@@ -134,7 +172,7 @@ def custom_instruction_convert_example(example, tokenizer, data_args, is_test=Fa
     )
 
 
-def left_padding(inputs, pad_id, max_length=0):
+def left_padding(inputs, pad_id, max_length=-1):
     for ids in inputs:
         max_length = max(max_length, len(ids))
 
@@ -156,7 +194,7 @@ class DataCollatorForSupervisedDataset(object):
     """Collate examples for supervised fine-tuning."""
 
     tokenizer: PretrainedTokenizerBase
-    max_length: 0
+    max_length: -1
 
     def __call__(self, features: List[Dict]) -> Dict[str, paddle.Tensor]:
 

examples/language_model/llama/finetune_instruction_generation.py

@@ -18,7 +18,11 @@
 
 import numpy as np
 import paddle
-from data import DataCollatorForSupervisedDataset, custom_instruction_convert_example
+from data import (
+    DataCollatorForSupervisedDataset,
+    custom_instruction_convert_example,
+    reader,
+)
 from sklearn.metrics import accuracy_score
 from utils import LlamaTrainer, compute_metrics, save_infer_result
 
@@ -57,6 +61,14 @@ class ModelArgument:
     prefix_projection: bool = field(default=True, metadata={"help": "Whether to project the prefix tokens"})
     use_flash_attention: bool = field(default=False, metadata={"help": "Whether to use flash attention"})
     do_generation: bool = field(default=False, metadata={"help": "Whether to do generation for evaluation"})
+    benchmark: bool = field(
+        default=False,
+        metadata={"help": "Whether or not run benchmark."},
+    )
+    profiler_options: str = field(
+        default=None,
+        metadata={"help": "profiler_options."},
+    )
 
 
 def main():
@@ -65,6 +77,8 @@ def main():
 
     training_args.print_config(model_args, "Model")
     training_args.print_config(data_args, "Data")
+    training_args.benchmark = model_args.benchmark
+    training_args.profiler_options = model_args.profiler_options
     setattr(training_args, "label_smoothing", model_args.label_smoothing)
     setattr(training_args, "lr_decay_ratio", model_args.lr_decay_ratio)
 
@@ -115,6 +129,7 @@ def main():
         model_args.model_name_or_path,
         padding_side="left",  # Allow batch inference
     )
+    tokenizer.pad_token = tokenizer.unk_token
 
     if model_args.lora:
         # TODO: hardcode parameters for now. Change after MergedLoRA is introduced
@@ -149,12 +164,20 @@ def main():
         model.print_trainable_parameters()
 
     # Load the dataset.
-    train_ds, dev_ds = load_dataset(data_args.data_name, data_args.task_name, splits=["train", "dev"])
-
-    trans_func = partial(custom_instruction_convert_example, tokenizer=tokenizer, data_args=data_args)
+    if training_args.benchmark:
+        train_ds = load_dataset(reader, data_path="./data/train.txt", lazy=False)
+        dev_ds = None
+    else:
+        train_ds, dev_ds = load_dataset(data_args.data_name, data_args.task_name, splits=["train", "dev"])
+
+    trans_func = partial(
+        custom_instruction_convert_example, tokenizer=tokenizer, data_args=data_args, benchmark=training_args.benchmark
+    )
     train_ds = train_ds.map(partial(trans_func))
-    dev_ds = dev_ds.map(partial(trans_func))
-    collate_fn = DataCollatorForSupervisedDataset(tokenizer)
+
+    if not training_args.benchmark:
+        dev_ds = dev_ds.map(partial(trans_func))
+    collate_fn = DataCollatorForSupervisedDataset(tokenizer, max_length=-1)
 
     def compute_metrics_trainer(eval_preds, tokenizer):
         all_preds = []

examples/language_model/llama/utils.py

@@ -14,6 +14,7 @@
 
 import json
 import os
+import time
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -24,7 +25,9 @@
 from rouge import Rouge
 
 from paddlenlp.metrics import BLEU
-from paddlenlp.trainer import Trainer
+from paddlenlp.trainer import PrinterCallback, ProgressCallback, Trainer
+from paddlenlp.trainer.integrations import TrainerCallback
+from paddlenlp.utils.log import logger
 
 
 def save_infer_result(trainer, dev_ds, k=100, src_length=256, tgt_length=512):
@@ -61,9 +64,100 @@ def save_infer_result(trainer, dev_ds, k=100, src_length=256, tgt_length=512):
                 f.write(json.dumps(out, ensure_ascii=False) + "\n")
 
 
+class AverageStatistical(object):
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.total_cnt = 0
+        self.time = 0
+
+    def record(self, val, cnt=1):
+        self.time += val
+        self.total_cnt += cnt
+
+    def get_average(self):
+        if self.total_cnt == 0:
+            return 0
+
+        return self.time / self.total_cnt
+
+    def get_average_per_sec(self):
+        if self.time == 0.0:
+            return 0.0
+
+        return float(self.total_cnt) / self.time
+
+    def get_total_cnt(self):
+        return self.total_cnt
+
+    def get_total_time(self):
+        return self.time
+
+
+class BenchmarkCallback(TrainerCallback):
+    def __init__(self, benchmark=True, profiler_options=None):
+        self.benchmark = benchmark
+        self.profiler_options = profiler_options
+
+    def on_train_begin(self, args, state, control, **kwargs):
+        assert args.gradient_accumulation_steps == 1 and not args.do_eval and not args.do_predict
+        if self.benchmark:
+            self.reader_cost_avg = AverageStatistical()
+
+    def on_epoch_begin(self, args, state, control, **kwargs):
+        if self.benchmark:
+            self.epoch_start = time.time()
+            self.batch_start = time.time()
+
+    def on_step_begin(self, args, state, control, **kwargs):
+        if self.benchmark:
+            self.reader_cost_avg.record(time.time() - self.batch_start)
+
+    def on_step_end(self, args, state, control, **kwargs):
+        if self.benchmark:
+            self.batch_start = time.time()
+            if control.should_log:
+                self.maybe_log_save_evaluate_start = time.time()
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if self.benchmark:
+            if logs is not None and "interval_steps_per_second" in logs:
+                self.batch_start = self.batch_start + (time.time() - self.maybe_log_save_evaluate_start)
+                ips = logs["interval_steps_per_second"] * args.train_batch_size
+                avg_batch_cost = 1 / logs["interval_steps_per_second"]
+                logger.info(
+                    "global step %d / %d, loss: %f, avg_reader_cost: %.5f sec, avg_batch_cost: %.5f sec, avg_samples: %.5f, ips: %.5f sample/sec"
+                    % (
+                        state.global_step,
+                        state.max_steps,
+                        logs["loss"],
+                        self.reader_cost_avg.get_average(),
+                        avg_batch_cost,
+                        args.train_batch_size,
+                        ips,
+                    )
+                )
+                self.reader_cost_avg.reset()
+
+    def on_epoch_end(self, args, state, control, **kwargs):
+        if self.benchmark:
+            train_epoch_cost = time.time() - self.epoch_start
+            logger.info("train epoch: %d, epoch_cost: %.5f s" % (state.epoch, train_epoch_cost))
+
+
 class LlamaTrainer(Trainer):
     def __init__(self, do_generation: bool, **kwargs):
         super().__init__(**kwargs)
+        if self.args.benchmark or self.args.profiler_options is not None:
+            self.add_callback(
+                BenchmarkCallback(benchmark=self.args.benchmark, profiler_options=self.args.profiler_options)
+            )
+            if self.args.benchmark:
+                if self.args.disable_tqdm:
+                    self.pop_callback(PrinterCallback)
+                else:
+                    self.pop_callback(ProgressCallback)
         self.do_generation = do_generation
 
     def prediction_step(

tests/test_tipc/configs/llama/train_infer_python.txt

@@ -0,0 +1,59 @@
+===========================train_params===========================
+model_name:llama
+python:python3.7
+gpu_list:0|0,1
+--device:gpu|gpu
+--fp16:null
+--max_steps:null
+null:null
+--per_device_train_batch_size:null
+null:null
+null:null
+null:null
+null:null
+##
+trainer:norm_train
+norm_train:../examples/language_model/llama/finetune_instruction_generation.py --model_name_or_path facebook/llama-7b-2l --do_train --max_steps 500 --recompute False --benchmark --overwrite_output_dir --output_dir ./checkpoints/ --fp16_opt_level O2 --learning_rate 3e-5 --lr_scheduler_type constant --warmup_steps 0 --seed 23 --logging_steps 1 --max_grad_norm -1
+pact_train:null
+fpgm_train:null
+distill_train:null
+null:null
+null:null
+##
+===========================eval_params=========================== 
+eval:null
+null:null
+##
+===========================infer_params===========================
+null:null
+null:null
+norm_export:null
+quant_export:null
+fpgm_export:null
+distill_export:null
+export1:null
+export2:null
+##
+infer_model:null
+infer_export:null
+infer_quant:null
+inference:null
+null:null
+null:null
+null:null
+null:null
+null:null
+null:null
+null:null
+null:null
+null:null
+null:null
+null:null
+===========================to_static_train_benchmark_params===========================
+to_static_train:--to_static
+===========================train_benchmark_params==========================
+batch_size:8
+fp_items:fp32|fp16
+epoch:500
+--profiler_options:batch_range=[10,20];state=GPU;tracer_option=Default;profile_path=model.profile  
+flags:FLAGS_eager_delete_tensor_gb=0.0;FLAGS_fraction_of_gpu_memory_to_use=0.98;FLAGS_conv_workspace_size_limit=4096

tests/test_tipc/prepare.sh

@@ -400,12 +400,18 @@ elif [ ${MODE} = "benchmark_train" ];then
         tar -zxvf laion400m_demo_data.tar.gz
     fi
 
+    if [[ ${model_name} =~ "llama" ]]; then
+        rm -rf llama_sft_demo_data.tar.gz
+        wget https://paddlenlp.bj.bcebos.com/models/community/facebook/llama_sft_demo_data.tar.gz
+        tar -xvf llama_sft_demo_data.tar.gz
+    fi
+
     export PYTHONPATH=$(dirname "$PWD"):$PYTHONPATH
     python -m pip install --upgrade pip -i https://pypi.tuna.tsinghua.edu.cn/simple
     python -m pip install setuptools_scm 
     python -m pip install Cython 
     python -m pip install -r ../requirements.txt  #-i https://pypi.tuna.tsinghua.edu.cn/simple
-    python -m pip install pybind11 regex sentencepiece tqdm visualdl attrdict pyyaml -i https://mirror.baidu.com/pypi/simple
+    python -m pip install pybind11 regex sentencepiece tqdm visualdl attrdict pyyaml rouge -i https://mirror.baidu.com/pypi/simple
 
     python -m pip install -e ../
     # python -m pip install paddlenlp    # PDC 镜像中安装失败