这段时间一直在研究大模型的微调，从ChatGPT到ChatGLM，再到这篇文章的Baichuan，感触颇深，不外乎就是大模型的训练时间很长，成本很高，效果并没有想象中的那么好，但我相信在不考虑成本的情况下针对特点场景下的微调可以达到相应的目标。

Baichuan-7B是百川智能推出的70亿参数的大模型，是一个很好的基座模型，具有非常棒的中文理解能力，但其还不具备聊天的能力；相比于使用现成的通用大模型去聊天，使用一个基座大模型去微调一个具备聊天能力的模型让人更满足。本篇使用QLoRA去微调这个模型，使用一张3090消费级显卡训练3个小时就可以满足训练需求。关于QLoRA的原理，这里就不过多介绍，其实就是LoRA的量化变体，本篇就是使用INT4的模型量化去训练LoRA模型。在训练的过程中发现我INT4量化竟然还会爆显存，这小小的7B我一个3090卡24G显存还不够用？7B的参数量在INT4下占用显存4G显存，加上AdamW训练占用16G，那么总共就20G显存就足够训练了，它竟然会爆显存？可能是我训练的token长度过长导致的显存需求过高，将batch train size调小，并多个小批次累积更新即可解决这个问题。废话不多说，下面贴出源码：

安装必要环境，transformers版本一定要对应起来，不然量化不了，这是一个坑。

#安装环境
!pip install -q transformers==4.30.2
#finetune需要
!pip install -q 'bitsandbytes==0.39.1' #提供4bit量化支持，版本限制非常重要，
!pip install datasets==2.13.1
!pip install -q git+https://github.com/huggingface/accelerate
!pip install  -q git+https://github.com/huggingface/peft  #使用最新版本非常重要，否则可能报错

如果包安装不了，去github使用源码安装。

接着去huggingface下载模型，国内用户会下载很慢，想要更快的获取请使用以下方法：

!pip install modelscope
from modelscope.hub.snapshot_download import snapshot_download
model_dir = snapshot_download('baichuan-inc/baichuan-7B', cache_dir='./baichuan-inc')

请前往baichuan-7B的仓库拉取百川的代码，解压后将下载好的baichuan-7B模型文件拉进代码目录中。下载完之后直接加载模型会报错，你就说坑不吭，它会报缺少configuration_baichuan这个模块，直接将baichuan的模型代码拉到你项目的目录中即可解决这个问题。如下图：

image.png

在代码目录新建qlora-tuning.ipynb，加载模型：

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig, BitsAndBytesConfig


bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True, #QLoRA 设计的 Double Quantization
    bnb_4bit_quant_type="nf4", #QLoRA 设计的 Normal Float 4 量化数据类型
    llm_int8_threshold=6.0,
    llm_int8_has_fp16_weight=False,
)
model_name_or_path = "baichuan-inc/baichuan-7B"
config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, 
                                             device_map="auto", 
                                             quantization_config=bnb_config, trust_remote_code=True)

这里需要注意的是device_map要设置为“auto”，不设置为“auto”后续可能会爆显存。

接着测试下原始模型的能力，使用Markdown来显示后面的输出：

from IPython.display import display, Markdown
device = torch.device('cuda:0')

def display_answer(text):
    inputs = tokenizer(text, return_tensors="pt")
    inputs = inputs.to(device)
    pred = model.generate(**inputs, max_new_tokens=256, repetition_penalty=1.1)
    res = tokenizer.decode(pred.cpu()[0], skip_special_tokens=True).replace(text, "")
    display(Markdown(res))

image.png

回答的效果并不好，接下来处理数据，为微调做好准备。数据集采用Hello-SimpleAI的中文数据集，下载地址请详见我的百度网盘分享：

链接：https://pan.baidu.com/s/1Lvvt9u9UbIE9QhYUViHpGw?pwd=nxmb 
提取码：nxmb 

接着处理处理数据集，如下：

import json, datasets
from tqdm import tqdm

def preprocess(tokenizer, config, file_path, max_seq_length, prompt_key, target_key, skip_overlength=False):
    # 数据预处理
    with open(file_path, "r", encoding="utf8") as f:
        for line in tqdm(f.readlines()):
            example = json.loads(line)
            prompt_ids = tokenizer.encode(example[prompt_key], max_length=max_seq_length, truncation=True)
            target_ids = tokenizer.encode(example[target_key], max_length=max_seq_length, truncation=True)
            input_ids = prompt_ids + target_ids + [config.eos_token_id]
            if skip_overlength and len(input_ids) > max_seq_length:
                continue
            input_ids = input_ids[:max_seq_length]
            yield {
                "input_ids": input_ids,
                "seq_len": len(prompt_ids)
            }


dataset = datasets.Dataset.from_generator(lambda: preprocess(tokenizer, 
                                            config, 
                                            "./hc3_chatgpt_zh_specific_qa.json", 
                                            max_seq_length=2000, 
                                            prompt_key="q",
                                            target_key="a",))

dataset.save_to_disk("h3c-chinese")  # 保存数据集

加载datasets数据集

train_set = datasets.load_from_disk("h3c-chinese")
print(len(train_set))

现在就可以使用peft库去LoRA微调了，导入相应的包

from transformers import TrainingArguments, Trainer
from peft import get_peft_model, LoraConfig

peft预处理INT4模型

from peft import prepare_model_for_kbit_training
model = prepare_model_for_kbit_training(model)

封装函数，找到可微调的参数，peft库会将这些参数转为低秩矩阵相乘

import bitsandbytes as bnb
def find_all_linear_nams(model):
    cls = bnb.nn.Linear4bit
    lora_module_names = set()
    for name, module in model.named_modules():
        if isinstance(module, cls):
            names = name.split('.')
            lora_module_names.add(names[0] if len(names) == 1 else names[-1])
    if "lm_head" in lora_module_names:
        lora_module_names.remove("lm_head")
    return list(lora_module_names)

lora_modules = find_all_linear_nams(model)
print(lora_modules)
# ['up_proj', 'gate_proj', 'o_proj', 'down_proj', 'W_pack']

初始化LoRA配置

peft_config = LoraConfig(
    task_type="CAUSAL_LM",
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=lora_modules,
)

model = get_peft_model(model, peft_config)
# 以下参数为了减少显存消耗，相应的训练时间也会变长，这也是没有显卡资源的无奈
model.supports_gradient_checkpointing = True
model.gradient_checkpointing_enable()
model.enable_input_require_grads()
model.config.use_cache = False
model.isparallelizable = True
model.model_parallel = True
model.print_trainable_parameters()

封装每一批数据forward前预处理的函数

tokenizer.pad_token_id = config.pad_token_id 
def data_collator(features):
    len_ids = [len(feature["input_ids"]) for feature in features]
    longest = max(len_ids)
    input_ids = []
    labels_list = []
    
    for ids_l, feature in sorted(zip(len_ids, features), key=lambda x: -x[0]):
        ids = feature["input_ids"]
        seq_len = feature["seq_len"]
        labels = (
            [-100] * (seq_len) + ids[seq_len:] + [-100] * (longest - ids_l)
        )

        ids = ids + [tokenizer.pad_token_id] * (longest - ids_l)
        input_ids.append(torch.LongTensor(ids))
        labels_list.append(torch.LongTensor(labels))
    return {
        "input_ids": torch.stack(input_ids),
        "labels": torch.stack(labels_list),
    }

创建训练器

class ModifiedTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        return model(
            input_ids=inputs["input_ids"],
            labels=inputs["labels"],
        ).loss
    
    def save_model(self, output_dir=None, _internal_call=False):
        self.model.save_pretrained(output_dir)

开始训练，这里max_steps设置为600，相应的训练会训练2轮，如果batch_size设置过大，3090会爆掉。

batch_size = 6
train_args = TrainingArguments(learning_rate=1e-4, 
                               per_device_train_batch_size=batch_size, 
                               gradient_accumulation_steps=10,
                               max_steps=600,
                               save_steps=100,
                               logging_steps=10,
                               output_dir="baichuan-7b-lora",
                               remove_unused_columns=False,
                              )

trainer = ModifiedTrainer(
    model=model,
    train_dataset=train_set,
    args=train_args,
    data_collator=data_collator,
)
trainer.train()
model.save_pretrained("./output")

模型的推理

以上训练完的模型会保存到output中，加载QLoRA微调后的模型

from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig

model_name_or_path = "baichuan-inc/baichuan-7B"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_name_or_path, trust_remote_code=True).half().cuda()  # 以半精度加载原始模型
model = PeftModel.from_pretrained(model, "output")  # 加载LoRA模型

然后直接调用display_answer，效果如下图所示：

display_answer("如何学习英语，使我顺利通过考试？")

display_answer("小明的老爸有四个儿子，大儿子叫老大，二儿子叫老二，三儿子叫老三，四儿子叫什么？")

image.png

display_answer("中国最高的山叫什么？")

image.png

image.png

本篇的ipynb不会公布出去，按照本篇肯定可以跑通代码的，想要源码的请留言，请先加个关注，后续有关其它模型的微调也会进行更新。