ChatGLM多任务微调

# 创建虚拟环境
conda create -n llm_env python=3.10
conda activate llm_env
# 安装依赖
pip install -r requirements.txt

protobuf>=3.19.5,<3.20.1
transformers==4.33
icetk
cpm_kernels
streamlit==1.18.0
matplotlib
datasets
accelerate>=0.20.3
packaging>=20.0
psutil
pyyaml
peft==0.3.0

mkdir -p /root/autodl-tmp/llm_tuning/THUDM/chatglm-6b
cd /root/autodl-tmp/llm_tuning/THUDM/chatglm-6b
pip install modelscope
modelscope download --model ZhipuAI/ChatGLM-6B --local_dir ./

{
    "context": "Instruction: 你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。\nInput: 找到句子中的三元组信息并输出成json给我:\n\n九玄珠是在纵横中文网连载的一部小说，作者是龙马。\nAnswer: ",
    "target": "```json\n[{\"predicate\": \"连载网站\", \"object_type\": \"网站\", \"subject_type\": \"网络小说\", \"object\": \"纵横中文网\", \"subject\": \"九玄珠\"}, {\"predicate\": \"作者\", \"object_type\": \"人物\", \"subject_type\": \"图书作品\", \"object\": \"龙马\", \"subject\": \"九玄珠\"}]\n```"
}

{
    "context": "Instruction: 你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。\nInput: 下面句子可能是一条关于什么的评论，用列表形式回答：\n\n很不错，很新鲜，快递小哥服务很好，水果也挺甜挺脆的\nAnswer: ",
    "target": "[\"水果\"]"
}

import os.path
import torch

base_dir = os.path.dirname(os.path.abspath(__file__))

class ProjectConfig(object):
    def __init__(self):
        self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
        # ChatGLM-6B 模型路径
        self.pre_model = os.path.join(base_dir, '../THUDM/model/chatglm-6b-int4')
        # 训练 / 验证数据路径
        self.train_path = os.path.join(base_dir, 'data/mixed_train_dataset.jsonl')
        self.dev_path = os.path.join(base_dir, 'data/mixed_dev_dataset.jsonl')
        # 微调方法开关
        self.use_lora = True
        self.use_ptuning = False
        self.lora_rank = 8  # LoRA 秩
        # 训练超参数
        self.batch_size = 4
        self.epochs = 2
        self.learning_rate = 3e-5
        self.weight_decay = 0
        self.warmup_ratio = 0.06
        # 序列长度限制
        self.max_source_seq_len = 100
        self.max_target_seq_len = 100
        # 日志与保存
        self.logging_steps = 10
        self.save_freq = 200
        # P-Tuning 伪 token 长度（如使用）
        self.pre_seq_len = 200
        self.prefix_projection = False  # False 为 p-tuning v1，True 为 p-tuning v2
        # 模型保存路径
        self.save_dir = os.path.join(base_dir, 'save_model')

import json
import numpy as np
from tqdm import tqdm
from transformers import AutoTokenizer
from ptune_chatglm.glm_config import ProjectConfig

pc = ProjectConfig()

def convert_example_chatglm(examples: dict, tokenizer, max_source_seq_len: int, max_target_seq_len: int):
    '''
    将样本数据转换为 ChatGLM 模型接收的输入数据。
    :param examples: 训练数据样本，e.g. -> {"text": ['{"context": "...", "target": "..."}', ...]}
    :param tokenizer: 分词器
    :param max_source_seq_len: prompt 最大长度
    :param max_target_seq_len: 答案最大长度
    :return: dict (str: np.array) -> {'input_ids': [...], 'labels': [...]}
    '''
    tokenized_output = {'input_ids': [], 'labels': []}
    max_seq_len = max_source_seq_len + max_target_seq_len

    for example in tqdm(examples['text']):
        try:
            example = json.loads(example)
            context = example["context"]
            target = example["target"]

            # 不添加特殊标记，因为需要手动添加
            context_ids = tokenizer.encode(context, add_special_tokens=False)
            target_ids = tokenizer.encode(target, add_special_tokens=False)

            # 截断：context 留 1 位给 [gMASK]，target 留 2 位给 <sop> 和 <eop>
            if len(context_ids) > max_source_seq_len - 1:
                context_ids = context_ids[:max_source_seq_len - 1]
            if len(target_ids) > max_target_seq_len - 2:
                target_ids = target_ids[:max_target_seq_len - 2]

            # source_ids + [gMASK] + <sop> + target_ids + <eop>
            input_ids = tokenizer.build_inputs_with_special_tokens(context_ids, target_ids)

            # bos_token_id 的索引位置即 context 长度
            context_length = input_ids.index(tokenizer.bos_token_id)
            # context 部分 label 设为 -100，target 部分保留原 id
            labels = [-100] * context_length + input_ids[context_length:]

            # 填充到固定长度
            pad_len = max_seq_len - len(input_ids)
            input_ids += [tokenizer.pad_token_id] * pad_len
            labels += [-100] * pad_len

            tokenized_output['input_ids'].append(input_ids)
            tokenized_output['labels'].append(labels)
        except Exception as e:
            print(f'Exception-->{e}')
            continue

    for k, v in tokenized_output.items():
        tokenized_output[k] = np.array(v)
    return tokenized_output

from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import default_data_collator, AutoTokenizer
from functools import partial
from ptune_chatglm.data_handle.data_preprocess import convert_example_chatglm
from ptune_chatglm.glm_config import ProjectConfig

pc = ProjectConfig()
tokenizer = AutoTokenizer.from_pretrained(pc.pre_model, trust_remote_code=True)

def get_data():
    # 'text' 加载器：按行读取纯文本，每行作为一条样本放到 text 列表里
    dataset = load_dataset('text', data_files={'train': pc.train_path, 'dev': pc.dev_path})

    # 使用 partial 固定 tokenizer 和长度参数
    new_func = partial(convert_example_chatglm,
                       tokenizer=tokenizer,
                       max_source_seq_len=pc.max_source_seq_len,
                       max_target_seq_len=pc.max_target_seq_len)
    dataset = dataset.map(new_func, batched=True)

    train_dataloader = DataLoader(dataset["train"], shuffle=True,
                                  collate_fn=default_data_collator,
                                  batch_size=pc.batch_size, drop_last=True)
    dev_dataloader = DataLoader(dataset["dev"],
                                collate_fn=default_data_collator,
                                batch_size=pc.batch_size, drop_last=True)
    return train_dataloader, dev_dataloader

import torch
import torch.nn as nn
import copy
from ptune_chatglm.glm_config import ProjectConfig

pc = ProjectConfig()

class CastOutputToFloat(nn.Sequential):
    """将模型输出转换为 float32 类型，作为模型最后一层确保输出类型符合预期。"""
    def forward(self, x):
        return super().forward(x).to(torch.float32)

def second2time(seconds: int):
    """将秒转换成时分秒。"""
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    return "%02d:%02d:%02d" % (h, m, s)

def save_model(model, cur_save_dir: str):
    """保存模型。若使用 LoRA，需先将 LoRA 参数与原始模型合并后再保存。"""
    if pc.use_lora:
        # 直接保存只会保存 adapter 参数，因此需要合并后再保存
        merged_model = copy.deepcopy(model)
        merged_model = merged_model.merge_and_unload()
        merged_model.save_pretrained(cur_save_dir)
    else:
        model.save_pretrained(cur_save_dir)

import sys, os, time
import torch
import peft
from torch.cuda.amp import autocast as autocast
from transformers import AutoConfig, AutoModel, get_scheduler, AutoTokenizer
from ptune_chatglm.data_handle.data_loader import get_data
from ptune_chatglm.glm_config import ProjectConfig
from ptune_chatglm.utils.common_utils import CastOutputToFloat, second2time, save_model

pc = ProjectConfig()

def model2train():
    """训练模型：从配置、数据加载到训练和评估的全过程。"""
    # 1. 获取数据加载器
    train_dataloader, dev_dataloader = get_data()

    # 2. 加载预训练模型
    tokenizer = AutoTokenizer.from_pretrained(pc.pre_model, trust_remote_code=True)
    config = AutoConfig.from_pretrained(pc.pre_model, trust_remote_code=True)

    if pc.use_ptuning:
        config.pre_seq_len = pc.pre_seq_len
        config.prefix_projection = pc.prefix_projection

    model = AutoModel.from_pretrained(pc.pre_model, config=config, trust_remote_code=True)
    # 转为 float16 减少显存
    model = model.half()
    model.config.use_cache = False
    # 梯度检查点：降低内存使用
    model.gradient_checkpointing_enable()
    # 输入层 require_grads
    model.enable_input_require_grads()

    if pc.use_ptuning:
        model.transformer.prefix_encoder.float()

    if pc.use_lora:
        # 输出头转换为 float32
        model.lm_head = CastOutputToFloat(model.lm_head)
        # LoRA 配置
        peft_config = peft.LoraConfig(
            task_type=peft.TaskType.CAUSAL_LM,
            inference_mode=False,
            r=pc.lora_rank,
            lora_alpha=32,
            lora_dropout=0.1
        )
        model = peft.get_peft_model(model, peft_config)

    model = model.to(pc.device)
    model.print_trainable_parameters()

    # 3. 优化器（bias 和 LayerNorm.weight 不做权重衰减）
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
         "weight_decay": pc.weight_decay},
        {"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
         "weight_decay": 0.0},
    ]
    optimizer = torch.optim.AdamW(optimizer_grouped_parameters, lr=pc.learning_rate)

    # 4. 学习率调度器
    num_update_steps_per_epoch = len(train_dataloader)
    max_train_steps = pc.epochs * num_update_steps_per_epoch
    warm_steps = int(pc.warmup_ratio * max_train_steps)
    lr_scheduler = get_scheduler(name='linear', optimizer=optimizer,
                                 num_warmup_steps=warm_steps,
                                 num_training_steps=max_train_steps)

    # 5. 训练循环
    loss_list = []
    tic_train = time.time()
    global_step, best_eval_loss = 0, float('inf')
    for epoch in range(1, pc.epochs + 1):
        print("开始训练")
        model.train()
        for batch in train_dataloader:
            if pc.use_lora:
                # autocast 混合精度训练（仅 GPU 可用）
                # FP16：速度快、显存少；FP32：保证数值稳定性，尤其在权重更新阶段
                with autocast():
                    loss = model(
                        input_ids=batch['input_ids'].to(dtype=torch.long, device=pc.device),
                        labels=batch['labels'].to(dtype=torch.long, device=pc.device)
                    ).loss
            else:
                loss = model(
                    input_ids=batch['input_ids'].to(dtype=torch.long, device=pc.device),
                    labels=batch['labels'].to(dtype=torch.long, device=pc.device)
                ).loss

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            lr_scheduler.step()

            loss_list.append(float(loss.cpu().detach()))
            global_step += 1
            if global_step % pc.logging_steps == 0:
                time_diff = time.time() - tic_train
                loss_avg = sum(loss_list) / len(loss_list)
                print("全局步骤 %d ( %02.2f%% ) , 轮次: %d, 损失: %.5f, 速度: %.2f step/s, ETA: %s"
                      % (global_step, global_step / max_train_steps * 100, epoch, loss_avg,
                         pc.logging_steps / time_diff,
                         second2time(int(max_train_steps - global_step) / (pc.logging_steps / time_diff))))
                tic_train = time.time()

        # 6. 每个 epoch 评估并保存最佳模型
        eval_loss = evaluate_model(model, dev_dataloader)
        print("评估集的损失为: %.5f" % (eval_loss))
        if eval_loss < best_eval_loss:
            print(f"最小的损失已经更新：{best_eval_loss:.5f} --> {eval_loss:.5f}")
            best_eval_loss = eval_loss
            cur_save_dir = os.path.join(pc.save_dir, "model_best")
            save_model(model, cur_save_dir)
            tokenizer.save_pretrained(cur_save_dir)
            print(f'最好的模型已经保存在：{cur_save_dir}.')
        tic_train = time.time()

def evaluate_model(model, dev_dataloader):
    """计算验证集平均 loss。"""
    model.eval()
    loss_list = []
    with torch.no_grad():
        for batch in dev_dataloader:
            if pc.use_lora:
                with autocast():
                    loss = model(
                        input_ids=batch['input_ids'].to(dtype=torch.long, device=pc.device),
                        labels=batch['labels'].to(dtype=torch.long, device=pc.device)
                    ).loss
            else:
                loss = model(
                    input_ids=batch['input_ids'].to(dtype=torch.long, device=pc.device),
                    labels=batch['labels'].to(dtype=torch.long, device=pc.device)
                ).loss
            loss_list.append(float(loss.cpu().detach()))
    return sum(loss_list) / len(loss_list)

if __name__ == '__main__':
    model2train()

import sys, os, time
import torch
from transformers import AutoTokenizer, AutoModel
from ptune_chatglm.glm_config import ProjectConfig

pc = ProjectConfig()

def inference(model, tokenizer, instuction: str, sentence: str):
    with torch.no_grad():
        input_text = f"Instruction: {instuction}\n"
        if sentence:
            input_text += f"Input: {sentence}\n"
        input_text += f"Answer: "
        batch = tokenizer(input_text, return_tensors="pt")
        out = model.generate(
            input_ids=batch["input_ids"].to(pc.device),
            max_new_tokens=max_new_tokens,
            temperature=0
        )
        out_text = tokenizer.decode(out[0])
        answer = out_text.split('Answer: ')[-1]
        return answer

if __name__ == '__main__':
    max_new_tokens = 300
    model_path = os.path.join(pc.save_dir, 'model_best')
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_path, trust_remote_code=True).half().to(pc.device)

    samples = [
        {
            'instruction': "现在你是一个非常厉害的SPO抽取器。",
            "input": "下面这句中包含了哪些三元组，用json列表的形式回答，不要输出除json外的其他答案。\n\n73获奖记录人物评价：黄磊是一个特别幸运的演员，拍第一部戏就碰到了导演陈凯歌，而且在他的下一部电影《夜半歌声》中演对手戏的张国荣、吴倩莲、黎明等都是著名的港台演员。",
        },
        {
            'instruction': "你现在是一个很厉害的阅读理解器，严格按照人类指令进行回答。",
            "input": "下面子中的主语是什么类别，输出成列表形式。\n\n第N次入住了，就是方便去客户那里哈哈。还有啥说的"
        }
    ]

    start = time.time()
    for i, sample in enumerate(samples):
        res = inference(model, tokenizer, sample['instruction'], sample['input'])
        print(f'res {i}: ')
        print(res)
    print(f'Used {round(time.time() - start, 2)}s.')