大语言模型工具

简介

这是一个基于🤗HuggingFace开发的大语言模型训练和测试工具。它支持不同模型的网页界面和终端预测，支持各种模型的低参数量及全参数模型的预训练、奖励模型训练以及RLHF训练（包括PPO和DPO两种方法）。同时还支持deepspeed分布式训练。

作者习惯于将配置和任务都写在一个配置文件中，然后以一个主函数作为入口直接运行，因此才有了这个项目。喜欢使用命令行的朋友们可以将其改回去使用。

更新日志

环境要求

几个重要的环境要求：

• Python：3.10+
• PyTorch：2.0.1+
• bitsandbytes：不同操作系统下需要对应安装不同的包（Linux下0.39.0+，Windows下要专门下载对应的wheel本地安装）

其它环境要求请参见requirements.txt 目前FlashAttention的作者未主动兼容和测试Windows操作环境问题，若在Windows上使用，则不需要安装flash-attn这个包。

特性

支持的模型

大模型经过SFT（然后进行RLHF）之后可用于对话任务Chat。面世的Chat大部分都没有重新训练基座，或者是基于同样的基座结构用数据重新预训练了一个基座。下表列出了经过验证并被此项目支持的基座模型，相应地也支持同样结构的衍生和Chat模型。

• 未列入上表的模型或参数规模暂时未经本项目测试。

模板提示词

由于许多训练者都是基于上述基础模型或对话模型进行进一步训练，但他们采用了不同的模板提示词，因此在下载相关模型后，需要根据这些模型的要求添加与之匹配的模板提示词。除了加载这些模型官方要求的模板提示词外，本项目还提供了一些模板提示词，如ziya、openbuddy等的模板。

训练方法

• Lora和AdaLora都支持QLora训练，但量化方式需要选择基于bitsandbytes的bnb量化方式，支持4位和8位量化训练。以下是启用Qlora训练所需的必要配置参数（在ModelArguments中）：

quantization: Optional[str] = field(
    default='bnb',
    metadata={
        # 如果使用qlora只能选择bnb，两种量化方式区别不大。
        'help': '要使用的具体模型版本（可以是分支名、标签名或提交ID）。',
        'choices': ['cpm', 'bnb'],
    }
)
quantization_bit: Optional[int] = field(
    default=None,
    metadata={
        # 使用8位量化还是4位量化？
        'help': '量化模型的位数。',
        'choices': [4, 8],
    }
)

量化

两种量化方式分别是基于bitsandbytes的bnb和基于cpm_kernels组件的cpm，其中cpm量化脚本来自quantization.py。

评估指标

在不同的训练阶段运行测试集时，会输出以下一些常规的生成模型评估结果。这些结果仅供参考，目前对大模型的事实性评估没有更好的方法，各个模型提供方或评测机构都是在各个维度上制作数据集进行评测，相对比较主观。

开始使用

在开始之前，需要确定要实验的模型，并从Hugging Face上下载整个模型文件，完成以下两步：

1. 在ModelArguments中配置好model_type和model_path两个参数。如果除了model_path的基座模型外还有adapter模型，则需要将adapter模型的地址配置到checkpoint_dir中。

model_type: str = field(
    default='internlm',
    metadata={
        # 模型类型
        'help': '模型类型。',
        'choices': ['chatglm', 'qwen', 'llama', 'falcon', 'baichuan', 'aquila', 'internlm', 'moss', 'bloom', 'rwkv'],
    }
)
model_path: str = field(
    default='/home/XXXXX/llm_models/internLM/intern-chat-7b',
    metadata={
        # 从huggingface.co/models上下载的模型保存到本地的路径。
        'help': '预训练模型的本地路径或huggingface.co/models上的模型标识符。'
    }
)
checkpoint_dir: Optional[str] = field(
    default=None,
    metadata={
        # 保存下载的或者自己训练的adapter增量模型的地方。
        'help': '自动保存（增量）模型检查点以及配置的路径。',
    }
)

2. 在DataTrainingArguments中修改prompt_template，使用与该模型配套的模板。这个模板通常只有经过SFT训练后的模型才会有，并且与训练者有关。如果该项目未提供，则需要自己修改engines/utils/prompt_template.py文件，添加新的模板。

prompt_template: Optional[str] = field(
    default='internlm',
    metadata={
        # 选择对应模型的模板prompt，一般Chat模型的出品方都会有一个固定的prompt。
        'help': '用于在训练和推理中构建提示的模板。'
    }
)

推理

这里提供两种预测方式，分别是基于Gradio的WebUI预测和终端预测。需要在config.py中相应地修改mode，然后运行main.py。

• 在预测时，模型会优先从你在ModelArguments中定义的checkpoint_dir读取参数。如果该文件夹下没有参数文件，则会从TrainingArguments的output_dir文件夹加载。如果都没有，则只加载最初的基座模型。

NTK

目前原生配置就能支持NTK方法的有chatglm2-6b-32k、LLama系列、Falcon系列和Qwen-7B-Chat：

• 其他模型需要自行修改原始模型文件以支持NTK方法，例如可用于Alibi编码的Baichuan、Falcon、Bloom系列的NTK-ALibi。一般来说，NTK主要用于在推理时突破模型的输入token限制，但在训练时启用NTK可能无法获得预期效果。
• Falcon系列模型在HF官方提供了两种编码方式，分别是Rope和Alibi，但tiiuae官方目前只实现了Alibi，原因不明。因此，这里仅支持使用Rope编码方式的NTK方法。

训练

预训练

预训练数据参考datasets/pretrain/example/train下的文件，数据以txt格式存储。制作数据集时最好能与示例类似，每行一句话，但不要超过模型接收的最大token长度。然后将数据路径填写到DataTrainingArguments配置中：

train_file_dir: Optional[str] = field(
    default='datasets/pretrain/example/train',
    metadata={
        # 训练集保存路径。
        'help': '训练json数据文件夹。'
    }
)
validation_file_dir: Optional[str] = field(
    default='datasets/pretrain/example/train',
    metadata={
        # 验证集保存路径。
        'help': '评估json文件夹。'
    }
)

开始训练时，需要在config.py中将mode改为pretrain，然后运行main.py。

SFT训练

指令微调数据参考datasets/finetune/example/train下的文件，数据由instruction、input、output和history四个字段组成。

[
  {
    "instruction": "10乘以10等于多少？",
    "input": "",
    "output": "10乘以10等于100。",
    "history": [
        "你好呀。",
        "你好，请问您有什么需要帮助的吗？",
        "好的，我想问下你是谁？",
        "我是一个AI模型，能够解决你提出的问题。"
      ]
  },
  ...  
]

如上所示，history字段需要按一问一答的格式存储对话历史，用于模型训练。如果没有历史对话，history应为空列表：

[
  {
    "instruction": "你身份是什么？",
    "input": "",
    "output": "我是一个AI智能助手，由XX公司训练，我将力所能及的帮助你解决问题。",
    "history": []
  },
  ...  
]

使用时将数据路径填写到DataTrainingArguments配置中：

train_file_dir: Optional[str] = field(
    default='datasets/finetune/example/train',
    metadata={
        # 训练集保存路径。
        'help': '训练json数据文件夹。'
    }
)
validation_file_dir: Optional[str] = field(
    default='datasets/finetune/example/eval',
    metadata={
        # 验证集保存路径。
        'help': '评估json文件夹。'
    }
)

开始训练时，需要在config.py中将mode改为sft_train，然后在TrainingArguments中配置好各项训练参数，运行main.py。

框架支持测试SFT训练效果。测试前在DataTrainingArguments中配置test_file为测试数据集路径，然后在config.py中将mode改为sft_batch_test，运行main.py。

test_file: Optional[str] = field(
    default='datasets/finetune/test',
    metadata={
        # 测试集保存路径。
        'help': '测试文件。'
    }
)

RM训练

奖励模型训练数据参考datasets/rm/example/train下的文件，数据由instruction、input、output三个字段组成。output是一个两元素列表，第一个元素是采纳的答案，第二个是拒绝的答案。使用时将训练奖励模型的数据同样填写到DataTrainingArguments配置中。然后需要在config.py中将mode改为rm_train，在TrainingArguments中配置好各项训练参数，运行main.py。

train_file_dir: Optional[str] = field(
    default='datasets/rm/example/train',
    metadata={
        # 训练集保存路径。
        'help': '训练json数据文件夹。'
    }
)
validation_file_dir: Optional[str] = field(
    default='datasets/rm/example/eval',
    metadata={
        # 验证集保存路径。
        'help': '评估json文件夹。'
    }
)

框架支持测试奖励模型训练效果。首先需要在DataTrainingArguments中配置test_file为测试数据集路径，然后在config.py中将mode改为rm_batch_test，运行main.py。奖励模型测试只会输出模型的准确率。

• 奖励模型训练不支持第一代ChatGLM6B，因为项目使用的trl的AutoModelForCausalLMWithValueHead组件是基于CausalLM模型的。而ChatGLM6B是基于Prefix LM实现的。

RLHF 训练

PPO

在进行基于PPO模型的RLHF训练之前，需要一个奖励模型和一个需要被RLHF微调的SFT模型，需要把它们配置到ModelArguments中如下：

checkpoint_dir: Optional[str] = field(
    default='checkpoint/sft',
    metadata={
        # 保存下载的或者自己训练的adapter增量模型的地方，在RLHF时候，此处需要填写指令微调后模型所在的文件地址。
        'help': '自动保存(增量)模型检查点和配置的路径。',
    }
)
reward_model_checkpoint: str = field(
    default='checkpoint/rm',
    metadata={
        # 在RLHF时候，此处需要填写奖励模型所在的文件地址
        '帮助': '奖励模型的检查点。'
    }
)

PPO方法对模型进行强化学习训练的数据和SFT阶段训练的数据的格式是一致的，此外使用的时候还需要在TrainingArguments中把PPO的配置填写好，在config.py中将mode修改为ppo_train，然后运行main.py。训练的结果将会通过wandb的格式记录在训练输出的文件夹中。

DPO

在进行基于DPO模型的RLHF训练之前，只需要一个被RLHF微调的SFT模型，如果是基于adapter的模型还需要把adapter配置到ModelArguments中如下：

model_path: str = field(
    default='/home/XXX/ChatGLM/ChatGLM2-6B-32k',
    metadata={
        # 从huggingface.co/models上下载的模型保存到本地的路径或者自己的模型。
        'help': '从huggingface.co/models下载的预训练模型的本地路径或模型标识符。'
    }
)
checkpoint_dir: Optional[str] = field(
    default='checkpoint/sft',
    metadata={
        # 保存下载的或者自己训练的adapter增量模型的地方，在RLHF时候，此处需要填写指令微调后模型所在的文件地址。
        'help': '自动保存(增量)模型检查点和配置的路径。',
    }
)

DPO方法对模型进行强化学习训练的数据和奖励模型的数据是一致的，在config.py中将mode修改为dpo_train，然后运行main.py。训练的结果将会通过wandb的格式记录在训练输出的文件夹中。

• 如果前面使用的是adapter在SFT模型上训练的模型，RLHF的时候项目会融合前面的adapter后创建新的adapter继续训练。

训练参数

常用的一些参数如下：

• Lora和其它adapter训练方式的配置参数也在TrainingArguments中，这里面要注意lora_target的设置要根据自己的模型结构来，配置中给了一些参考。
• Firefly Loss仅作用在SFT训练阶段且不支持ChatGLM6B等Prefix LM模型。

DeepSpeed

使用deepspeed进行训练需要在TrainingArguments指定deepspeed的config文件(项目中提供了stage2的deepspeed配置)：

deepspeed: Optional[str] = field(
    default='deepspeed_configs/zero_stage2_config.json',
    metadata={
        'help': '启用deepspeed并传递deepspeed json配置文件的路径（例如ds_config.json）或已加载的json文件作为字典'
    }
)

配置好后在终端输入(单机多卡)：

deepspeed --num_gpus 3 --master_port=9901 main.py

• 多机多卡需要指定更多的参数，可以参考hugingface的deepspeed文档。

其他

• merge_peft_model和save_quantized_model需要在ModelArguments中设置输出地址。

quantized_or_merged_output_dir: Optional[str] = field(
    default=None,
    metadata={
        # 当你想保存量化后的模型或者融合后的模型时，处理后的模型保存的地址。
        'help': '手动保存量化或融合后的模型检查点以及配置的路径。',
    }
)

• 使用bnb和cpm量化时，将默认对除输出层外的所有线性层进行量化。
• 使用expand_vocab方法进行词表扩充时，需要指定训练词表的语料路径（支持文件或文件夹）。仅支持 .txt 和 .tsv 格式。词表扩充后，通常需要继续进行预训练。

待办事项

• 奖励模型训练
• PPO模型训练
• DPO模型训练
• 支持Deepspeed训练
• 集成NTK-Aware Scaled RoPE
• 集成多轮对话的Firefly损失函数
• 支持LLM增量预训练
• 为LLama和Falcon添加Flash Attention2
• mmlu、cmmlu和C-Eval自动化评估

引用

如果你在研究中使用了本项目，请按以下格式引用：

@misc{LLMs Tool,
  title={LLMs Tool: 大型语言模型工具},
  author={Shouxian Li},
  year={2023},
  howpublished={\url{https://github.com/stanleylsx/llms_tool}},
}

Star历史