LLM - ChatGLM-6B Lora 微调与推理

目录

一.引言

二.环境准备

三.ChatGLM-6B Lora 微调

1.样本准备 By Json

2.样本生成 By Tokenizer

3.模型生成 By Trainer

四.ChatGLM-6B Lora 文本生成

1.文本生成 By Chat

2.输出测试

五.总结

---

一.引言

ChatGLM 是一个初具问答和对话功能的千亿中英语言模型， 并针对中文进行了优化，本文基于 ChatGLM-6B 实现 Lora 微调，整体步骤与前面介绍的 LLM - Baichuan7B Lora 训练详解 有一些相似之处。

Tips：

本文涉及样本处理、模型训练和推理可在一张 Tesla V100 x 32G 实现。

二.环境准备

◆ 主要依赖

python 3.9.11
numpy==1.23.5
torch==2.0.1
transformers==4.29.1

可以将上述依赖版本放入 requirements.txt 中，上面是博主自己的配置，也可以根据 ChatGLM 官方给出的 requirements 准备环境 https://github.com/THUDM/ChatGLM-6B。

◆ 激活并配置环境

conda create -n chatglm python=3.9
conda activate chatglm
 
pip install -r requirements.txt

◆ ChatGLM 6B 模型下载

下载地址: https://huggingface.co/THUDM/chatglm-6b

三.ChatGLM-6B Lora 微调

1.样本准备 By Json

{"q": "请计算：39 * 0 = 什么？", "a": "这是简单的乘法运算，39乘以0得到的是0"}
{"q": "题目：51/186的答案是什么?", "a": "这是简单的除法运算，51除以186大概为0.274"}
{"q": "鹿妈妈买了24个苹果，她想平均分给她的3只小鹿吃，每只小鹿可以分到几个苹果？", "a":"鹿妈妈买了24个苹果，平均分给3只小鹿吃，那么每只
小鹿可以分到的苹果数就是总苹果数除以小鹿的只数。\n24÷3=8\n每只小鹿可以分到8个苹果。所以，答案是每只小鹿可以分到8个苹果。"}
{"q": "请计算：39 * 0 = 什么？", "a": "这是简单的乘法运算，39乘以0得到的是0"}
{"q": "题目：51/186的答案是什么?", "a": "这是简单的除法运算，51除以186大概为0.274"}
{"q": "鹿妈妈买了24个苹果，她想平均分给她的3只小鹿吃，每只小鹿可以分到几个苹果？", "a": "鹿妈妈买了24个苹果，平均分给3只小鹿吃，那么每>只小鹿可以分到的苹果数就是总苹果数除以小鹿的只数。\n24÷3=8\n每只小鹿可以分到8个苹果。所以，答案是每只小鹿可以分到8个苹果。"}
{"q": "请计算：39 * 0 = 什么？", "a": "这是简单的乘法运算，39乘以0得到的是0"}
{"q": "题目：51/186的答案是什么?", "a": "这是简单的除法运算，51除以186大概为0.274"}
{"q": "鹿妈妈买了24个苹果，她想平均分给她的3只小鹿吃，每只小鹿可以分到几个苹果？", "a": "鹿妈妈买了24个苹果，平均分给3只小鹿吃，那么每>只小鹿可以分到的苹果数就是总苹果数除以小鹿的只数。\n24÷3=8\n每只小鹿可以分到8个苹果。所以，答案是每只小鹿可以分到8个苹果。"}

与前面 Baichuan-7B 类似，这里我们准备了 10 条测试样本，每一条样本以 QA 的形式构建上下文对话，其中文件格式需保存为 json。

2.样本生成 By Tokenizer

◆ python 脚本

import argparse
import json
from tqdm import tqdm
import datasets
import transformers

# 1.参数准备
parser = argparse.ArgumentParser()
parser.add_argument("--model_checkpoint", type=str, help="checkpoint, like `THUDM/chatglm-6b`") # 必填
parser.add_argument("--input_file", type=str, help="Instruction 数据文件地址，文件中每一行都是json格式，包含一个输出和一个输出") # 必填
parser.add_argument("--prompt_key", type=str, default=f"prompt", help="你的jsonl文件里，Instruction 的输入字段是什么") # 选填
parser.add_argument("--target_key", type=str, default=f"target", help="你的jsonl文件里，Instruction 的输出字段是什么") # 必填
parser.add_argument("--save_name", type=str, default=f"temp", help="经过tokenize之后的数据集的存放位置") # 选填
parser.add_argument("--max_seq_length", type=int, default=2040) # 选填
parser.add_argument("--skip_overlength", type=bool, default=False) # 选填
args = parser.parse_args()
model_checkpoint = args.model_checkpoint

#. 2.处理逻辑
def preprocess(tokenizer, config, example, max_seq_length, prompt_key, target_key):
    print(config.pad_token_id, config.eos_token_id, config.pad_token_id, tokenizer.unk_token, tokenizer.pad_token, tokenizer.eos_token)
    prompt = example[prompt_key]
    target = example[target_key]
    prompt_ids = tokenizer.encode(prompt, max_length=max_seq_length, truncation=True)
    target_ids = tokenizer.encode(target, max_length=max_seq_length, truncation=True, add_special_tokens=False)
    # 最终还是将 instruction 的输入输出都拼在一起，使用经典的 causal-LM 的 next word prediction 方式来训练
    input_ids = prompt_ids + target_ids + [config.eos_token_id] # EOS 用于标识句子结束
    print("p:", prompt_ids)
    print("t:", target_ids)
    print(input_ids, len(prompt_ids))
    return {"input_ids": input_ids, "seq_len": len(prompt_ids)}

# 3.读取训练 JSON
def read_jsonl(path, max_seq_length, prompt_key,target_key,skip_overlength=False):
    # 基于预训练模型加载获取 tokenizer 和 config
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        model_checkpoint, trust_remote_code=True)
    config = transformers.AutoConfig.from_pretrained(
        model_checkpoint, trust_remote_code=True, device_map='auto')
    with open(path, "r") as f:
        for line in tqdm(f.readlines()):
            example = json.loads(line)
            feature = preprocess(tokenizer, config, example, max_seq_length,prompt_key,target_key)
            if skip_overlength and len(feature["input_ids"]) > max_seq_length:
                continue
            # 截取最大长度
            feature["input_ids"] = feature["input_ids"][:max_seq_length]
            yield feature


# 输入文件统一放在 data 文件夹下
# 输出文件统一放在 data/tokenized_data 文件夹下
input_file_path = f'data/{args.input_file}'
save_path = f"data/tokenized_data/{args.save_name}"
dataset = datasets.Dataset.from_generator(
    lambda: read_jsonl(input_file_path, args.max_seq_length, args.prompt_key,args.target_key,args.skip_overlength)
)

dataset.save_to_disk(save_path)

python 脚本的批量处理逻辑主要在 read_jsonl 方法中，方法内会调用处理单条样本的 preprocess 方法，最后使用 datasets.Dataset.from_generator 生成数据 dataset 并调用 save_to_disk 方法保存到磁盘上。这里根据大家习惯可以调整输入和输出的文件夹位置，本例中输入的 input_file 位于 ./data 目录下，tokenizer 后生成的 dataset 位于 ./data/tokenized_data 文件夹下。文件夹下文件内容如下：

◆ shell 脚本

chatGLM="/models/ChatGLM-6B/chatglm-6b/"

input=simple_test.json
outpyt=simple_token_by_chatGLM

CUDA_VISIBLE_DEVICES=0 python tokenize_dataset_rows.py \
    --model_checkpoint $chatGLM \
    --input_file $input \
    --prompt_key q \
    --target_key a \
    --save_name $output \
    --max_seq_length 2000 \
    --skip_overlength False

这里把下载好的 ChatGLM 地址传到脚本和 input、output 传入即可，如果你按照上面 qa 的形式构造样本，则 prompt_key 和 target_key 也无需修改，否则需要根据自己 json 里的 QA 的 key 修改。max_seq_length 用于 token 的截取，skip_overlength 用于样本的取舍。

◆ 运行结果

执行脚本后得到 tokenizer 后的样本：

Generating train split: 0 examples [00:00, ? examples/s]                                                                              3 130005 3                                                                                         | 0/9 [00:00  
p: [5, 68247, 12, 15, 9, 26, 9, 23, 21, 77612, 65267, 31, 130001, 130004]
t: [65356, 67061, 64607, 63947, 79222, 6, 15, 9, 103875, 9, 23, 21, 66359, 63834, 8, 7, 10, 25, 16]
[5, 68247, 12, 15, 9, 26, 9, 23, 21, 77612, 65267, 31, 130001, 130004, 65356, 67061, 64607, 63947, 79222, 6, 15, 9, 103875, 9, 23, 21, 66359, 63834, 8, 7, 10, 25, 16, 130005] 14

其中几个特殊 token 和 id 如下：

eos_token_id、pad_token_id、unk_token、pad_token、eos_token
130005        3                        

这里还有一个坑等下我们训练代码时进行分析。

3.模型生成 By Trainer

◆ python 脚本

from transformers.integrations import TensorBoardCallback
from torch.utils.tensorboard import SummaryWriter
from transformers import TrainingArguments
from transformers import Trainer, HfArgumentParser
from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn as nn
from peft import get_peft_model, LoraConfig, TaskType
from dataclasses import dataclass, field
import datasets
import os
from pprint import pprint as print

model_path = "/models/ChatGLM-6B/chatglm-6b/"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, add_special_tokens=False)
print(tokenizer.mask_token_id)
print(tokenizer.bos_token_id)

@dataclass
class FinetuneArguments:
    tokenized_dataset: str = field(default=" ") # tokenized之后的数据集文件夹
    model_path: str = field(default=" ")
    lora_rank: int = field(default=8)


class CastOutputToFloat(nn.Sequential):
    def forward(self, x):
        return super().forward(x).to(torch.float32)


def data_collator(features: list) -> dict:
    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)
        _ids = torch.LongTensor(ids)
        labels_list.append(torch.LongTensor(labels))
        input_ids.append(_ids)
    input_ids = torch.stack(input_ids)
    labels = torch.stack(labels_list)
    return {
        "input_ids": input_ids,
        "labels": labels,
    }


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):
        from transformers.trainer import TRAINING_ARGS_NAME

        os.makedirs(output_dir, exist_ok=True)
        torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME))
        saved_params = {
            k: v.to("cpu") for k, v in self.model.named_parameters() if v.requires_grad
        }
        torch.save(saved_params, os.path.join(output_dir, "adapter_model.bin"))
        

def main():
    writer = SummaryWriter()
    finetune_args, training_args = HfArgumentParser(
        (FinetuneArguments, TrainingArguments)
    ).parse_args_into_dataclasses()

    # load dataset
    dataset = datasets.load_from_disk('data/tokenized_data/'+finetune_args.tokenized_dataset)
    print(f"\n{len(dataset)=}\n")
    
    # init model
    model = AutoModel.from_pretrained(
        model_path, load_in_8bit=False, trust_remote_code=True, 
        device_map="auto" # 模型不同层会被自动分配到不同GPU上进行计算
        ,empty_init=False
    )
    print(model.hf_device_map)
    
    model.gradient_checkpointing_enable() 
    model.enable_input_require_grads()
    model.lm_head = CastOutputToFloat(model.lm_head)

    # setup peft
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        inference_mode=False,
        r=finetune_args.lora_rank,
        lora_alpha=32,
        lora_dropout=0.1,
    )
    
    model = get_peft_model(model, peft_config)


    # start train
    model.save_pretrained(training_args.output_dir) # 因为adapter_config.json只能通过这个save_pretrained来生成，先这里生成一份，好在训练完之前就可以尝试中间的checkpoint
    trainer = ModifiedTrainer(
        model=model,
        train_dataset=dataset,
        args=training_args,
        callbacks=[TensorBoardCallback(writer)],
        data_collator=data_collator,
    )
    trainer.train()
    writer.close()
    # save model
    model.save_pretrained(training_args.output_dir)


if __name__ == "__main__":
    main()

python 脚本逻辑也不复杂，主要是 Transformer API 和 peft API 的组合使用:

➤  datasets.load_from_disk - 加载 tokenized 生成的训练集

➤  AutoModel.from_pretrained - 加载原生的 ChatGLM-6B 模型

➤  LoraConfig - 配置 Lora 微调参数，主要参数为 r 即 lora_rank

➤  get_peft_model - 在 Base 模型基础上获取 Lora 微调模型

➤  Trainer.train - 继承 Trainer 实现 Loss 计算与 model save 的方法并训练

➤  model.save_pretrained - Lora 微调模型参数保存

◆ shell 脚本

CUDA_VISIBLE_DEVICES=0 python chatglm_lora_tuning.py \
    --tokenized_dataset simple_token_by_chatGLM \
    --lora_rank 4 \
    --per_device_train_batch_size 8 \
    --gradient_accumulation_steps 1 \
    --num_train_epochs 10 \
    --save_steps 200 \
    --save_total_limit 2 \
    --learning_rate 1e-4 \
    --fp16 \
    --remove_unused_columns false \
    --logging_steps 50 \
    --output_dir weights/simple_test_by_chatglm

上述 python 脚本命名为 chatglm_lora_tuning.py，token 后数据保存在 tokenized_data 下的 simple_token_by_chatGLM 文件夹 下，其余 lora_rank、batch_size、epoch、lr 等都是训练的超参，可以根据自己场景设定。最后定义 output_dir 定义 Lora 微调后模型参数的存储地址即可。

◆ 运行结果

执行上述 shell 脚本后模型开始训练，我们一共 9 个样本，batch_size 采用默认的 8，所以全部样本得到 2 个 batch，又 epoch = 10，所以共 20 个 iter 即 20 个 step。训练数据很少，主打一个跑通：

 运行期间显存会提高至 20G 左右，GPU-Util 最高 100%，这个和 batch_size 也有一定关系：

◆ 异常分析 ValueError: 130004 is not in list

这个错误一开始比较懵逼，但是后来分析了 tokenizer 生成的样本和 ChatGLM 样本规则后问题就迎刃而解了。这里 130004 是 ChatGLM Token 里 bos_token 的 token_id：

print(tokenizer.bos_token_id)

再回看 tokenizer 生成的 q_ids 和 a_ids：

可以看到每个 q 的结尾都有一个 bos_token_id，而样本生成逻辑是将 QA 结合在一起再做掩码生成 label，由于 bos_token 标识句子的开始，所以 label 生成时代码为 (seq_len - 1)，因为这里把 130004 对应的 bos_token 算作了 A 的开头：

这里顺便介绍下几个常用的 token_id：

➤ bos_token - Beginning of Sentence Token 表示句子的开头

➤ eop_token - End of Paragraph Token 标识段落的结束

➤ eos_token - End of Sentence Token 表示句子的结束

之前还介绍了 pad-token、unk_token 分别标识填充元素和未知元素。这些特殊标记用于标识文章段落结构，在训练和推断时以便让模型学到文章的不同成分。除此之外，可以发现 q_ids 的结尾除了有 130004 之外，还有一个 130001，我们使用 tokenizer 反推：

tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
token_id = 130001
token = tokenizer.convert_ids_to_tokens(token_id)
print("Id: %d Token: %s" %(token_id, token))

得到 1300001 对应的 token 是：

Id: 130001 Token: [gMASK]

◆ 异常分析结论 ValueError: 130004 is not in list

因为 ChatGLM 对 Input_ids 格式有要求，需要必须有 130001 和 130004，而如果 Tokenizer 读取了之前版本的模型文件或者读取了并非 ChatGLM-6B 对应的模型文件获取 tokenizer 并处理预训练数据，就会导致 tokenizer 后的 dataset 数据中不包含 130001 和 130004。所以出现该问题可以优先检查模型文件读取是否正确，博主出这个问题是因为我读取了 Baichuan-7B 的模型文件生成 tokenizer 处理 ChatGLM-6B 的训练数据，所以报错。当然也有偷懒的方法，因为我们知道了 ChatGLM 的 input_ids 生成风格，我们也可以手动在 q_ids 后面追加一个 130001 和 130004，这样也可以继续后面的训练。

四.ChatGLM-6B Lora 文本生成

1.文本生成 By Chat

Baichuan-7B 和 ChatGLM-6B 都提供 generate 的文本生成方法，使用方法类似。下面使用 ChatGLM 的 chat 方法进行文本生成：

from peft import PeftModel
from transformers import AutoTokenizer, AutoModel
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import time

def cost(st, end):
    # 转换为 ms
    return (end - st) * 1000

# 加载原始 LLM
model_path = "/models/ChatGLM-6B/chatglm-6b/"

device = torch.device(0)
load_st = time.time()
model = AutoModel.from_pretrained(model_path, load_in_8bit=False, trust_remote_code=True).half().to(device)
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
load_end = time.time()

# 原始 LLM 安装上 Lora 模型
lora_model = PeftModel.from_pretrained(model, "weights/simple_test_by_chatglm").half()
load_lora_end = time.time()
print("Load Ori Model: %s Load Lora Model: %s" % (cost(load_st, load_end), cost(load_end, load_lora_end)))

while True:
    inputText = input("请输入信息 [输入'q'退出]\n")
    if inputText == 'q':
        print("Exit!")
        break
    else:
        time_1 = time.time()
        ori_out = model.chat(tokenizer, inputText, history=[])
        time_2 = time.time()
        lora_out = lora_model.chat(tokenizer, inputText, history=[])
        time_3 = time.time()
        print("原始输出:", ori_out)
        print("Lora输出:", lora_out)
        print("Total Cost: %s Ori Cost: %s Lora Cost: %s" % (cost(time_1, time_3), cost(time_1, time_2), cost(time_2, time_3)))

2.输出测试

这里使用 AutoModel 和 PeftModel 分别加载原始 ChatGLM-6B 和 Lora ChatGLM-6B 进行 chat 文本输出。

◆ 自然语言

这里先用唐诗测试了模型的古文常识功能，不论是 Ori 还是 Lora 都能够正常输出。

◆ 逻辑推理

{"q": "请计算：39 * 0 = 什么？", "a": "这是简单的乘法运算，39乘以0得到的是0"}

首先使用训练样本的数据，可以看到 Lora 后其口吻更接近 QA 的训练数据，因为它说的是 "得到" 而原始模型说的是 "结果就是":

简单修改一下问题，换一个数字，39 * 0 改成 39 * 1，原始模型能力未受影响，Lora 模型还在 "得到"，但是计算出现了明显的错误 39 * 1 = 49:

五.总结

ChatGLM-6B 微调的 Demo 大概就这么多，在 Transformer、Hugging Face 以及 Peft 逐渐丰富的 API 包装下，对模型进行微调变得越来越简单。但可以看到，虽然只使用了小批量的样本，也可能在样本 Token 比较相近的情况下，造成原始模型能力的丧失；而与样本 Token 无关时，模型基础能力受到影响相对较小甚至没有。

6月底最新的 ChatGLM2-6B 也已经发布，可以看到新模型的推出和更新速度非常快，这边刚刚搞定 ChatGLM，马上出 ChatGLM2；刚刚搞定 6B 的模型，还有 13B、130B ... 在大模型快速更新迭代的过程中，一方面要调试新的模型，也明白了数据的重要性，想要获得更好的结果可能从最简单的源头 "干净好用的训练数据" 开始会更好。后面有机会也会继续介绍新的模型例如 ChatGLM2、LLama2。

◆ 参考链接

Baichuan-7B - https://github.com/baichuan-inc/baichuan-7B

ChatGLM-6B - THUDM/chatglm-6b · Hugging Face

ChatGLM2-6B - https://github.com/thudm/chatglm2-6b

LLM-Tuning - GitHub - beyondguo/LLM-Tuning

Lora 流程代码详解 -  LLM - Baichuan7B Lora 训练详解_BIT_666的博客-CSDN博客

样本生成和 Lora 代码的一些细节和参数放在上文中 ▲