這段時間一直在研究大模型的微調，從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不會公布出去，按照本篇肯定可以跑通代碼的，想要源碼的請留言，請先加個關注，后續有關其它模型的微調也會進行更新。
