paddlepaddle--paddlenlp/slm/examples/benchmark/peft/torch/benchmark.py

# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from dataclasses import dataclass, field
from typing import Optional

import torch
import torch.profiler as profiler
from datasets import load_dataset
from peft import LoraConfig, TaskType, get_peft_model
from transformers import (
    AutoModel,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    DataCollatorForSeq2Seq,
    HfArgumentParser,
    LlamaTokenizer,
    TrainingArguments,
)
from utils import CustomTrainer, ProfilerCallback

"""
单卡
python benchmark.py --model_name_or_path bigscience/bloomz-7b1-mt  \
    --num_train_epochs 1 --per_device_train_batch_size 4 \
    --evaluation_strategy no --save_strategy no \
    --fp16 --lora \
    --logging_steps 50 --output_dir outputs

多卡 deepspeed zero3
python -m torch.distributed.run --nproc_per_node=4 benchmark.py --deepspeed ds_config.json \
    --model_name_or_path bigscience/bloomz-7b1-mt  \
    --num_train_epochs 1 --per_device_train_batch_size 2 \
    --evaluation_strategy no --save_strategy no \
    --fp16 \
    --logging_steps 50 --output_dir outputs
"""


@dataclass
class ModelArguments:
    """
    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
    """

    model_name_or_path: str = field(default=None, metadata={"help": "model name or local path"})
    lora: Optional[bool] = field(default=False, metadata={"help": "whether to use LoRA"})
    qlora: Optional[bool] = field(default=False, metadata={"help": "whether to use qLoRA"})
    english: Optional[bool] = field(default=False, metadata={"help": "whether to english benchmark dataset"})
    profiler: Optional[bool] = field(default=False, metadata={"help": "whether to use profiler"})
    double_quant: bool = field(
        default=True, metadata={"help": "Compress the quantization statistics through double quantization."}
    )
    quant_type: str = field(
        default="nf4", metadata={"help": "Quantization data type to use. Should be one of `fp4` or `nf4`."}
    )
    bits: int = field(default=4, metadata={"help": "How many bits to use."})
    max_memory_MB: int = field(default=80000, metadata={"help": "Free memory per gpu."})
    train_data_size: int = field(default=1000, metadata={"help": "Number of dataset for training"})


def main():
    parser = HfArgumentParser((ModelArguments, TrainingArguments))
    model_args, training_args = parser.parse_args_into_dataclasses()

    if "llama" in model_args.model_name_or_path:
        tokenizer = LlamaTokenizer.from_pretrained(model_args.model_name_or_path, use_fast=False)
        tokenizer.pad_token_id = 0
    elif model_args.model_name_or_path in ["cerebras/Cerebras-GPT-13B", "stanford-crfm/levanter-gpt2-7B"]:
        tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, use_fast=False)
        tokenizer.pad_token_id = 0
    else:
        tokenizer = AutoTokenizer.from_pretrained(model_args.model_name_or_path, trust_remote_code=True)

    compute_dtype = torch.float16 if training_args.fp16 else (torch.bfloat16 if training_args.bf16 else torch.float32)
    if "chatglm" in model_args.model_name_or_path:
        # Add empty_init=False for zero3 training, refer to https://github.com/THUDM/ChatGLM-6B/issues/530
        model = AutoModel.from_pretrained(
            model_args.model_name_or_path,
            empty_init=False if training_args.deepspeed is not None else True,
            trust_remote_code=True,
            torch_dtype="auto",
        )

    else:
        if model_args.qlora:
            n_gpus = torch.cuda.device_count()
            max_memory = f"{model_args.max_memory_MB}MB"
            max_memory = {i: max_memory for i in range(n_gpus)}
            device_map = "auto"

            # if we are in a distributed setting, we need to set the device map and max memory per device
            if os.environ.get("LOCAL_RANK") is not None:
                local_rank = int(os.environ.get("LOCAL_RANK", "0"))
                device_map = {"": local_rank}
                max_memory = {"": max_memory[local_rank]}

            model = AutoModelForCausalLM.from_pretrained(
                model_args.model_name_or_path,
                torch_dtype="auto",
                load_in_4bit=model_args.bits == 4,
                load_in_8bit=model_args.bits == 8,
                device_map=device_map,
                max_memory=max_memory,
                quantization_config=BitsAndBytesConfig(
                    load_in_4bit=model_args.bits == 4,
                    load_in_8bit=model_args.bits == 8,
                    llm_int8_threshold=6.0,
                    llm_int8_has_fp16_weight=False,
                    bnb_4bit_compute_dtype=compute_dtype,
                    bnb_4bit_use_double_quant=model_args.double_quant,
                    bnb_4bit_quant_type=model_args.quant_type,
                ),
            )
        elif model_args.model_name_or_path in ["cerebras/Cerebras-GPT-13B", "stanford-crfm/levanter-gpt2-7B"]:
            model = AutoModelForCausalLM.from_pretrained(
                model_args.model_name_or_path,
                torch_dtype=torch.float16,
            )
        else:
            model = AutoModelForCausalLM.from_pretrained(
                model_args.model_name_or_path,
                torch_dtype="auto",
            )
    if model_args.lora:
        if "llama" in model_args.model_name_or_path:
            target_modules = ["q_proj", "k_proj", "v_proj"]
        elif model_args.model_name_or_path in ["cerebras/Cerebras-GPT-13B", "stanford-crfm/levanter-gpt2-7B"]:
            target_modules = [
                ".*c_attn.*",
                ".*q_attn.*",
                ".*c_proj.*",
                ".*c_fc.*",
            ]
        else:
            target_modules = ["query_key_value"]
        peft_config = LoraConfig(
            task_type=TaskType.CAUSAL_LM, target_modules=target_modules, r=8, lora_alpha=32, lora_dropout=0.0
        )
        model = get_peft_model(model, peft_config)
        model.print_trainable_parameters()

    if model_args.lora and training_args.gradient_checkpointing:
        # For backward compatibility
        if hasattr(model, "enable_input_require_grads"):
            model.enable_input_require_grads()
        else:

            def make_inputs_require_grad(module, input, output):
                output.requires_grad_(True)

            model.get_input_embeddings().register_forward_hook(make_inputs_require_grad)

        # enable gradient checkpointing for memory efficiency
        model.gradient_checkpointing_enable()

    def preprocess_function(example, max_src_length=256, max_tgt_length=384):
        inputs = example["instruction"]
        if "input" in example:
            inputs += example["input"]
        targets = example["output"]
        model_inputs = tokenizer(inputs, max_length=max_src_length, truncation=True, return_attention_mask=False)

        labels = tokenizer(targets, max_length=max_tgt_length, truncation=True, return_attention_mask=False)
        labels_input_ids = labels["input_ids"] + [tokenizer.eos_token_id]

        model_inputs["labels"] = [-100] * len(model_inputs["input_ids"]) + labels_input_ids
        model_inputs["input_ids"] = model_inputs["input_ids"] + labels_input_ids
        return model_inputs

    if model_args.english:
        dataset = load_dataset("tatsu-lab/alpaca")
    else:
        dataset = load_dataset("Chinese-Vicuna/guanaco_belle_merge_v1.0")

    # select first 10k examples for benchmarking
    dataset = dataset["train"].select(range(model_args.train_data_size))
    dataset = dataset.map(
        lambda example: preprocess_function(example), remove_columns=["instruction", "input", "output"]
    )
    total_effective_tokens = sum([len(i["input_ids"]) for i in dataset]) * training_args.num_train_epochs

    if model_args.profiler:
        prof = profiler.profile(
            activities=[
                torch.profiler.ProfilerActivity.CPU,
                torch.profiler.ProfilerActivity.CUDA,
            ],
            schedule=torch.profiler.schedule(wait=1, warmup=1, active=2, repeat=1),
            on_trace_ready=torch.profiler.tensorboard_trace_handler("hf-training-trainer"),
            profile_memory=True,
            with_stack=True,
        )

    data_collator = DataCollatorForSeq2Seq(return_tensors="pt", tokenizer=tokenizer)

    trainer = CustomTrainer(
        model=model,
        tokenizer=tokenizer,
        train_dataset=dataset,
        callbacks=[ProfilerCallback(prof=prof)] if model_args.profiler else [],
        args=training_args,
        data_collator=data_collator,
    )
    model.config.use_cache = False  # silence the warnings. Please re-enable for inference!
    train_metrics = trainer.train()
    tokens_per_second = trainer.total_observed_tokens / train_metrics.metrics["train_runtime"]
    effective_tokens_per_second = total_effective_tokens / train_metrics.metrics["train_runtime"]
    print(f"Tokens per second: {tokens_per_second:.2f}")
    print(f"Effective Tokens per second: {effective_tokens_per_second:.2f}")


if __name__ == "__main__":
    main()