version: 2559fc82 (2024-04-02 10:17:24)

LLM 训练实例

那么，下面我们终于要开始训自己的模型了。这一章节我们来训一个自己的角色扮演模型。

环境配置

Windows

与上一章节的样例一样，需要WSL。

MacOS

特别注意：貌似只有M2芯片的Mac + MacOS Sonoma (14+)才支持Bfloat16，而即使是最新的PyTorch也不支持float16/bfloat16的AMP Autocast(见https://github.com/pytorch/pytorch/issues/104191)，所以下面的流程对于Mac很可能无法使用。还好我们目前已经支持了MLX的训练。对了，如果你是直接阅读的这一章节，记得参照上一章的"安装其他依赖/MacOS"中的内容来升级mlx-lm和mlx。

目前训练框架对MLX功能支持如下：

  • 梯度累积 Gradient accumulation
  • 梯度暂存 Gradient checkpointing
  • 梯度裁剪 Gradient clipping
  • LoRA权重合并
  • 保存权重
  • 加载权重 & 恢复训练
  • 数据Packing
  • 全量训练
  • QLoRA

另外介绍一个工具asitop可以用来监测M芯片的GPU占用，可以用pip install asitop -i https://pypi.tuna.tsinghua.edu.cn/simple.

通用

到这里，不管你是WSL还是Mac，假设我们都是POSIX的环境。基本的步骤都是通用的了。

  1. 安装训练依赖

    pip install sentencepiece accelerate tokenizers typer jsonref pydantic -i https://pypi.tuna.tsinghua.edu.cn/simple

    [Linux/WSL 必选] 安装训练时所使用的依赖

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

    [可选，仅Linux/WSL]Flash Attention是一个加速框架，用来加速Attention部分的计算。

    pip install https://github.com/Dao-AILab/flash-attention/releases/download/v2.5.5/flash_attn-2.5.5+cu118torch2.3cxx11abiFALSE-cp311-cp311-linux_x86_64.whl

  2. 下载训练框架

    git clone https://github.com/peterjc123/functionary

    我们把functionary目录的绝对路径记为FUNCTIONARY_HOME。

    P.S. 如后续需要更新，只需进入functionary目录，然后git pull即可

下载数据

一般基于Huggingface或者Huggingface镜像(国内)来下载数据。

下载数据实例

对于角色扮演任务，我们使用RoleLLM所使用的Rolebench数据。

下面是下载的命令

# 选择一条执行即可
## 国外 Huggingface
git lfs clone https://huggingface.co/datasets/ZenMoore/RoleBench
## 国内 Huggingface镜像
GIT_LFS_SKIP_SMUDGE=1 git lfs clone https://hf-mirror.com/datasets/ZenMoore/RoleBench

数据预处理

简单的检视一下数据，可以看到数据在rolebench-zh下面，其中rolebench-zh/general/train.jsonl是角色对于通用问题的回答。而为了能扮演任意角色，我们也需要一个简单的人物介绍，这个可以在profiles-zh/desc.json里面找到。所以我们只需要把这些数据给拼起来就好了。

下面是处理数据的Python代码：

import os
import json

from transformers import AutoTokenizer

# RoleBench数据集的路径
home_dir = '<PATH_TO_ROLEBENCH>'
profile_path = os.path.join(home_dir, 'profiles-zh', 'desc.json')
general_data_path = os.path.join(home_dir, 'rolebench-zh', 'general', 'train.jsonl')
role_specific_data_path = os.path.join(home_dir, 'rolebench-zh', 'role_specific', 'train.jsonl')

# 模型的路径
model_dir = '<PATH_TO_QWEN_MODEL>'
tokenizer = AutoTokenizer.from_pretrained(model_dir)

# 输出路径
output_dir = '<PATH_TO_OUTPUT>'
output_path = os.path.join(output_dir, 'rolebench-train.jsonl')

with open(profile_path, 'r') as f:
    profile_dict = json.load(f)

def transform_message(role, profile, question, answer):
    messages = [
        {"role": "system", "content": f"你是{role}，{profile}。请扮演{role}与我对话。"},
        {"role": "human", "content": question},
        {"role": "assistant", "content": answer},
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False
    ) + "<|im_end|>"
    return text

prompts = []
with open(general_data_path, 'r') as f:
    for l in f:
        obj = json.loads(l)
        role = obj['role']
        question = obj['question']
        answer = obj['generated'][0]
        profile = profile_dict[role].strip('。')
        prompt = transform_message(role, profile, question, answer)
        prompts.append(prompt)

os.makedirs(output_dir, exist_ok=True)
with open(output_path, 'w') as f:
    for prompt in prompts:
        f.write(json.dumps({"text": prompt}))
        f.write("\n")

P.S. 这里为了快速演示，只整理了中文的数据。为了更好的效果，可以把英语的数据也一起放进去。有时间可以自己研究一下。

开始训练

使用下面的命令即可开启LoRA模型训练。

export PYTHONPATH=${FUNCTIONARY_HOME}

if [[ "$OSTYPE" == "darwin"* ]]; then
    # 对比下面的命令，目前还不支持gradient_checkpointing，还去掉了几个不支持的flag比如tf32,bf16等等
    python3 ${FUNCTIONARY_HOME}/functionary/train/train_lora_mlx.py \
        --model_name_or_path ${MODEL_PATH}/Qwen1.5-1.8B-Chat \
        --train_data_path ${DATA_PATH}/rolebench-train.jsonl \
        --output_dir ${OUTPUT_PATH}/roleplay_model_lora \
        --num_train_epochs 1 \
        --per_device_train_batch_size 2 \
        --gradient_accumulation_steps 6 \
        --save_strategy "steps" \
        --save_steps 100 \
        --save_total_limit 3 \
        --logging_steps 10 \
        --learning_rate 2e-4 \
        --weight_decay 0. \
        --warmup_ratio 0.03 \
        --lr_scheduler_type "cosine" \
        --model_max_length 512 \
        --report_to none \
        --optim adamw_bnb_8bit \
        --do_train \
        --max_grad_norm 1.0 \
        --seed 42 \
        --lora_dropout 0.05 \
        --lora_r 64 \
        --lora_alpha 16 \
        --use_cpu True
else
    python3 ${FUNCTIONARY_HOME}/functionary/train/train_lora.py \
        --model_name_or_path ${MODEL_PATH}/Qwen1.5-1.8B-Chat \
        --train_data_path ${DATA_PATH}/rolebench-train.jsonl \
        --bf16 True \
        --output_dir ${OUTPUT_PATH}/roleplay_model_lora \
        --num_train_epochs 1 \
        --per_device_train_batch_size 2 \
        --gradient_accumulation_steps 6 \
        --save_strategy "steps" \
        --save_steps 100 \
        --save_total_limit 3 \
        --logging_steps 10 \
        --learning_rate 2e-4 \
        --weight_decay 0. \
        --warmup_ratio 0.03 \
        --lr_scheduler_type "cosine" \
        --tf32 True \
        --model_max_length 512 \
        --report_to none \
        --gradient_checkpointing True \
        --optim adamw_bnb_8bit \
        --do_train \
        --max_grad_norm 1.0 \
        --seed 42 \
        --lora_dropout 0.05 \
        --lora_r 64 \
        --lora_alpha 16
fi

训练的时候中我们正好来了解一些概念

  1. 什么是LoRA?

    虽然这些模型已经很小了，但是我们手头的显卡显存还是太少了，因此没法办法正常训练这些模型。好在有LoRA技术来帮我们，什么是LoRA呢？就是对于一个矩阵W，在做矩阵乘法的时候，我们加入两个小的矩阵A和B，大小分别为[input_dim, lora_r]和[lora_r, output_dim]，其中lora_r << dim，这样他们矩阵乘的结果是大小[input_dim, output_dim]的矩阵W'，我们通过学习这两个小矩阵，而把原来的大矩阵给冻结住，矩阵乘法变为(W+W')(x)=Wx + lora_alpha * ABx，从而达到高效学习的效果。

  2. 这些参数是什么意思？

    通用参数

    • model_name_or_path 模型路径
    • train_data_path 数据路径
    • output_dir 输出模型路径

    训练参数

    • num_train_epochs 训练轮次
    • learning_rate 学习率
    • optim 优化器
    • weight_decay AdamW优化器的参数
    • lr_scheduler_type 学习的调节方式
    • warmup_ratio 学习率warmup的比率
    • max_grad_norm 梯度裁剪的最大值
    • seed 随机数种子
    • logging_steps 汇报频次
    • save_steps 模型保存频次
    • save_strategy 模型保存策略
    • save_total_limit 最大保存的模型数

    模型参数

    • model_max_length 模型最大输入长度，越长约占显存
    • bf16 是否开启Bfloat16支持，需要30x0及以上的显卡，可以加速训练
    • tf32 是否开启Tensor-float32支持，需要30x0及以上的显卡，可以加速训练
    • fp16 如果上面两个不支持，可以使用Float16来加速训练

    LoRA参数

    • lora_r LoRA矩阵的大小
    • lora_alpha 平衡LoRA和原模型权重的参数
    • lora_dropout LoRA层的dropout比例

合并模型

前面我们也知道了LoRA会训练两个矩阵A、B，那么在推理时我们可以将他们和原矩阵W融合起来，下面的脚本可以完成这个工作。

export PYTHONPATH=${FUNCTIONARY_HOME}

if [[ "$OSTYPE" == "darwin"* ]]; then
    # 对比下面的命令，需要我们手工输入lora的参数
    lora_rank=64
    lora_alpha=16
    lora_dropout=0.05
    python3 ${FUNCTIONARY_HOME}/functionary/train/merge_lora_weight_mlx.py \
        ${OUTPUT_PATH}/roleplay_model \
        ${MODEL_PATH}/Qwen1.5-1.8B-Chat \
        ${OUTPUT_PATH}/roleplay_model_lora \
        ${lora_rank} \
        ${lora_alpha} \
        ${lora_dropout}
else
    python3 ${FUNCTIONARY_HOME}/functionary/train/merge_lora_weight.py \
        ${OUTPUT_PATH}/roleplay_model \
        ${MODEL_PATH}/Qwen1.5-1.8B-Chat \
        ${OUTPUT_PATH}/roleplay_model_lora
fi

使用训练好的模型

请参考上一节的使用方式，这里直接做了一些小修改。

import readline
import os
import platform

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
try:
    import mlx_lm
    is_mlx_available = True
except ImportError:
    is_mlx_available = False

use_mlx = True

if platform.system() == "Darwin":
    device = "mps"
    dtype = torch.float16
else:
    device = "cuda"
    dtype = "auto"

path = "<SAVE_PATH>/roleplay_model"

if is_mlx_available and use_mlx:
    model, tokenizer = mlx_lm.load(path)
else:
    model = AutoModelForCausalLM.from_pretrained(
        path,
        torch_dtype=dtype,
        device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(path)

default_messages = [
    {"role": "system", "content": "你是李云龙，桀骜不驯，纪律性差，特立独行，脾气暴躁，是一个没文化又对政治缺乏兴趣的“大老粗”军人。请扮演李云龙和我对话。"},
]

messages = list(default_messages)
while True:
    prompt = input("Human:").strip("\n").strip()
    if prompt == "\\quit":
        break
    elif prompt == "\\new":
        os.system("clear")
        messages = list(default_messages)
    else:
        messages.append({"role": "user", "content": prompt})

    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    if is_mlx_available and use_mlx:
        # 之前MLX的推理不支持高级的采样策略，比如Top-P和Top-K，因此做了一些微小的工作
        # 见 https://github.com/ml-explore/mlx-examples/pull/486
        response = mlx_lm.generate(model, tokenizer, text, temp=0.6, top_p=0.8, max_tokens=512, repetition_penalty=1.1)
    else:
        model_inputs = tokenizer([text], return_tensors="pt").to(device)

        generated_ids = model.generate(
            model_inputs.input_ids,
            max_new_tokens=512,
            temperature=0.6,
            top_p=0.8,
            repetition_penalty=1.1,
        )
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]

        response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(f"Assistant:{response}")
    messages.append({"role": "assistant", "content": response})