LLM - ChatGLM-6B Lora 微调与推理_chatglm2-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 <unk> <pad> <eop>                                                                                      | 0/9 [00:00<?, ?it/s]
p: [71492, 65416, 12, 13, 18, 44, 5, 8, 29, 5, 64061, 31, 130001, 130004]
t: [65356, 67061, 114063, 79222, 6, 13, 18, 104040, 8, 72656, 63829, 8]
[71492, 65416, 12, 13, 18, 44, 5, 8, 29, 5, 64061, 31, 130001, 130004, 65356, 67061, 114063, 79222, 6, 13, 18, 104040, 8, 72656, 63829, 8, 130005] 14
Generating train split: 1 examples [00:01,  1.49s/ examples]3 130005 3 <unk> <pad> <eop>
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            <unk>      <pad>      <eop>

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

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 文件夹