ChatGLM2-6B Lora 微调训练医疗问答任务-编程知识

一、ChatGLM2-6B Lora 微调

LoRA 微调技术的思想很简单，在原始 PLM (Pre-trained Language Model) 增加一个旁路，一般是在 transformer 层，做一个降维再升维的操作，模型的输入输出维度不变，来模拟 intrinsic rank，如下图的 A 和 B。训练时冻结 PLM 的参数，只训练 A 和 B ，，输出时将旁路输出与 PLM 的参数叠加，进而影响原始模型的效果。该方式，可以大大降低训练的参数量，而性能可以优于其它参数高效微调方法，甚至和全参数微调（Fine-Tuning）持平甚至超过。

对于 A 和 B 参数的初始化，A 使用随机高斯分布，B 使用 0 矩阵，这样在最初时可以保证旁路为一个 0 矩阵，最开始时使用原始模型的能力。

在这里插入图片描述
对于 lora 微调的实现可以使用 HuggingFace 开源的 PEFT 库，地址如下：

https://github.com/huggingface/peft

下载依赖：

pip install peft -i https://pypi.tuna.tsinghua.edu.cn/simple

使用方式也很简单，例如先查看 ChatGLM2-6B 的模型结构：

from transformers import AutoModelmodel_name = "chatglm-6b"
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
print(model)

输出结果：

ChatGLMForConditionalGeneration((transformer): ChatGLMModel((embedding): Embedding((word_embeddings): Embedding(65024, 4096))(rotary_pos_emb): RotaryEmbedding()(encoder): GLMTransformer((layers): ModuleList((0-27): 28 x GLMBlock((input_layernorm): RMSNorm()(self_attention): SelfAttention((query_key_value): Linear(in_features=4096, out_features=4608, bias=True)(core_attention): CoreAttention((attention_dropout): Dropout(p=0.0, inplace=False))(dense): Linear(in_features=4096, out_features=4096, bias=False))(post_attention_layernorm): RMSNorm()(mlp): MLP((dense_h_to_4h): Linear(in_features=4096, out_features=27392, bias=False)(dense_4h_to_h): Linear(in_features=13696, out_features=4096, bias=False))))(final_layernorm): RMSNorm())(output_layer): Linear(in_features=4096, out_features=65024, bias=False))
)

可以看出 ChatGLM 主要由 28 层的 GLMBlock 进行提取和理解语义特征，下面借助 PEFT 库将 Lora 旁路层注入到模型中，主要关注下 query_key_value 层的变化：

from transformers import AutoTokenizer, AutoModel, AutoConfig
from peft import LoraConfig, get_peft_model, TaskTypemodel_name = "chatglm-6b"
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)config = LoraConfig(peft_type="LORA",task_type=TaskType.CAUSAL_LM,inference_mode=False,r=8,lora_alpha=16,lora_dropout=0.1,fan_in_fan_out=False,bias='lora_only',target_modules=["query_key_value"]
)model = get_peft_model(model, config)
print(model)

其中 r 就是 lora 中秩的大小。

输出结果：

PeftModelForCausalLM((base_model): LoraModel((model): ChatGLMForConditionalGeneration((transformer): ChatGLMModel((embedding): Embedding((word_embeddings): Embedding(65024, 4096))(rotary_pos_emb): RotaryEmbedding()(encoder): GLMTransformer((layers): ModuleList((0-27): 28 x GLMBlock((input_layernorm): RMSNorm()(self_attention): SelfAttention((query_key_value): Linear(in_features=4096, out_features=4608, bias=True(lora_dropout): ModuleDict((default): Dropout(p=0.1, inplace=False))(lora_A): ModuleDict((default): Linear(in_features=4096, out_features=8, bias=False))(lora_B): ModuleDict((default): Linear(in_features=8, out_features=4608, bias=False))(lora_embedding_A): ParameterDict()(lora_embedding_B): ParameterDict())(core_attention): CoreAttention((attention_dropout): Dropout(p=0.0, inplace=False))(dense): Linear(in_features=4096, out_features=4096, bias=False))(post_attention_layernorm): RMSNorm()(mlp): MLP((dense_h_to_4h): Linear(in_features=4096, out_features=27392, bias=False)(dense_4h_to_h): Linear(in_features=13696, out_features=4096, bias=False))))(final_layernorm): RMSNorm())(output_layer): Linear(in_features=4096, out_features=65024, bias=False))))
)

可以对比下原始的 ChatGLM 模型结构， query_key_value 层中已经被加入下 lora 的 A、B 层，下面可以通过 model.print_trainable_parameters() 打印可训练的参数量：

trainable params: 2,078,720 || all params: 6,245,533,696 || trainable%: 0.03328330453698988

可以看到可训练的参数量只有 0.03328330453698988 。

下面依然借助前面文章使用的医疗问答数据集，在 ChatGLM2 lora 微调下的效果。

对该数据集不了解的小伙伴可以参考下面这篇文章：

ChatGLM2-6B P-Tuning v2 微调训练医疗问答任务

二、ChatGLM2-6B Lora 微调

解析数据，构建 Dataset 数据集 qa_dataset.py：

# -*- coding: utf-8 -*-
from torch.utils.data import Dataset
import torch
import json
import numpy as npclass QADataset(Dataset):def __init__(self, data_path, tokenizer, max_source_length, max_target_length) -> None:super().__init__()self.tokenizer = tokenizerself.max_source_length = max_source_lengthself.max_target_length = max_target_lengthself.max_seq_length = self.max_source_length + self.max_target_lengthself.data = []with open(data_path, "r", encoding='utf-8') as f:for line in f:if not line or line == "":continuejson_line = json.loads(line)content = json_line["content"]summary = json_line["summary"]self.data.append({"question": content,"answer": summary})print("data load ， size：", len(self.data))def preprocess(self, question, answer):prompt = self.tokenizer.build_prompt(question, None)a_ids = self.tokenizer.encode(text=prompt, add_special_tokens=True, truncation=True,max_length=self.max_source_length)b_ids = self.tokenizer.encode(text=answer, add_special_tokens=False, truncation=True,max_length=self.max_target_length)context_length = len(a_ids)input_ids = a_ids + b_ids + [self.tokenizer.eos_token_id]labels = [self.tokenizer.pad_token_id] * context_length + b_ids + [self.tokenizer.eos_token_id]pad_len = self.max_seq_length - len(input_ids)input_ids = input_ids + [self.tokenizer.pad_token_id] * pad_lenlabels = labels + [self.tokenizer.pad_token_id] * pad_lenlabels = [(l if l != self.tokenizer.pad_token_id else -100) for l in labels]return input_ids, labelsdef __getitem__(self, index):item_data = self.data[index]input_ids, labels = self.preprocess(**item_data)return {"input_ids": torch.LongTensor(np.array(input_ids)),"labels": torch.LongTensor(np.array(labels))}def __len__(self):return len(self.data)

构造 Lora 结构，微调训练 train_lora.py：

# -*- coding: utf-8 -*-
import pandas as pd
from torch.utils.data import DataLoader
from transformers import AutoTokenizer, AutoModel
from qa_dataset import QADataset
from peft import LoraConfig, get_peft_model, TaskType
from tqdm import tqdm
import torch
import os, time, sysdef train(epoch, model, device, loader, optimizer, gradient_accumulation_steps):model.train()time1 = time.time()for index, data in enumerate(tqdm(loader, file=sys.stdout, desc="Train Epoch: " + str(epoch))):input_ids = data['input_ids'].to(device, dtype=torch.long)labels = data['labels'].to(device, dtype=torch.long)outputs = model(input_ids=input_ids,labels=labels,)loss = outputs.loss# 反向传播，计算当前梯度loss.backward()# 梯度累积步数if (index % gradient_accumulation_steps == 0 and index != 0) or index == len(loader) - 1:# 更新网络参数optimizer.step()# 清空过往梯度optimizer.zero_grad()# 100轮打印一次 lossif index % 100 == 0 or index == len(loader) - 1:time2 = time.time()tqdm.write(f"{index}, epoch: {epoch} -loss: {str(loss)} ; each step's time spent: {(str(float(time2 - time1) / float(index + 0.0001)))}")def validate(tokenizer, model, device, loader, max_length):model.eval()predictions = []actuals = []with torch.no_grad():for _, data in enumerate(tqdm(loader, file=sys.stdout, desc="Validation Data")):input_ids = data['input_ids'].to(device, dtype=torch.long)labels = data['labels'].to(device, dtype=torch.long)generated_ids = model.generate(input_ids=input_ids,max_length=max_length,do_sample=False,temperature=0)preds = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g ingenerated_ids]target = [tokenizer.decode(t, skip_special_tokens=True, clean_up_tokenization_spaces=True) for t in labels]predictions.extend(preds)actuals.extend(target)return predictions, actualsdef main():model_name = "chatglm-6b"train_json_path = "./data/train.json"val_json_path = "./data/val.json"max_source_length = 128max_target_length = 512epochs = 5batch_size = 1lr = 1e-4lora_rank = 8lora_alpha = 32gradient_accumulation_steps = 16model_output_dir = "output"# 设备device = torch.device("cuda:0")# 加载分词器和模型tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)model = AutoModel.from_pretrained(model_name, trust_remote_code=True)# setup peftpeft_config = LoraConfig(task_type=TaskType.CAUSAL_LM,inference_mode=False,r=lora_rank,lora_alpha=lora_alpha,lora_dropout=0.1)model = get_peft_model(model, peft_config)model.is_parallelizable = Truemodel.model_parallel = Truemodel.print_trainable_parameters()# 转为半精度model = model.half()model.float()print("Start Load Train Data...")train_params = {"batch_size": batch_size,"shuffle": True,"num_workers": 0,}training_set = QADataset(train_json_path, tokenizer, max_source_length, max_target_length)training_loader = DataLoader(training_set, **train_params)print("Start Load Validation Data...")val_params = {"batch_size": batch_size,"shuffle": False,"num_workers": 0,}val_set = QADataset(val_json_path, tokenizer, max_source_length, max_target_length)val_loader = DataLoader(val_set, **val_params)optimizer = torch.optim.AdamW(params=model.parameters(), lr=lr)model = model.to(device)print("Start Training...")for epoch in range(epochs):train(epoch, model, device, training_loader, optimizer, gradient_accumulation_steps)print("Save Model To ", model_output_dir)model.save_pretrained(model_output_dir)# 验证print("Start Validation...")with torch.no_grad():predictions, actuals = validate(tokenizer, model, device, val_loader, max_target_length)# 验证结果存储final_df = pd.DataFrame({"Generated Text": predictions, "Actual Text": actuals})val_data_path = os.path.join(model_output_dir, "predictions.csv")final_df.to_csv(val_data_path)print("Validation Data To ", val_data_path)if __name__ == '__main__':main()

开始训练：

在这里插入图片描述

等待训练结束后，可以在输出目录看到保存的模型，仅只有 lora 层的参数，所以模型比较小：

在这里插入图片描述

此时可以查看下 predictions.csv 中验证集的效果。

三、模型测试

from transformers import AutoTokenizer, AutoModel, AutoConfig
from peft import PeftConfig, PeftModel, LoraConfig, get_peft_model, TaskType
import torchdef load_lora_config(model):config = LoraConfig(task_type=TaskType.CAUSAL_LM,inference_mode=False,r=8,lora_alpha=32,lora_dropout=0.1,target_modules=["query_key_value"])return get_peft_model(model, config)device = torch.device("cuda:0")model_name = "chatglm-6b"
lora_dir = "output"model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)config = PeftConfig.from_pretrained(lora_dir)
model = PeftModel.from_pretrained(model, lora_dir)model = model.to(device)
model.eval()response, history = model.chat(tokenizer, "5月至今上腹靠右隐痛，右背隐痛带酸，便秘，喜睡，时有腹痛，头痛，腰酸症状？", history=[])
print("回答：", response)

输出：

在这里插入图片描述

回答：你好，根据你的叙述，考虑是胃炎引来的。建议你平时留意饮食规律，不要吃辛辣刺激性食物，多喝热水，可以口服奥美拉唑肠溶胶囊和阿莫西林胶囊实施救治，如果效果不好，建议去医院做胃镜仔细检查。除了及时救治胃痛外，患者朋友理应始终保持愉快的心态去直面疾病，只有这样才能令得患者及时对症救治，同时要多看重自身饮食护理，多观注自身的症状变动，认为这样一定能将胃痛撵走。