LLaMA-Factory原理与底层实现

LLaMA-Factory的简要介绍

简化大模型训练的工具

从 Pre-train → SFT(监督微调,supervised fine-tuning)→ RLHF(基于人类反馈的强化学习,reinforcement learning human feedback),一步完成。

一张4090完成70B参数的模型微调。

效率:

算子上提高两倍提速和60%的缓存节省。

一键完成。

通过分布式框架,甚至可以实现最多180B的框架模型微调。

目前已经有数十篇ArXiv论文使用了LLaMA-Factory工具作为微调框架。

为什么要使用LLaMA-Factory

  1. 如果从头训练一个llama2,模型,从LLaMA论文来看,仅7B参数就需要184320GPU小时,是非常遥不可及的时长。
  2. 报错难追踪,因为训练的报错常常是CUDA层面的bug,非常不方便定位。

GaLora微调技术,关于一种全参数的高效训练方法

QLora微调技术,一种Lora训练的量化版本。举例来说,在7B模型下使用2bite的量化微调QLora方法,可以有效的降低到仅仅需要4GB的显存占用。

Lora是在网络的旁边增加一个旁路,将d维度的高秩表征映射到r维的低秩表征,最后再从r维重新映射回d维。由于完整的微调工作是在缩小后的r维上进行的,这样就使得完整的训练中需要修改的参数量大幅度的降低(只用到0.01%的参数,显存占用量减少3倍)。Lora用多了还是有一个明显的缺陷,即对很多任务的拟合依然不是很好。

但是GaLora会是更好的微调方法,因为它具备了全参数微调的能力。

LLaMA-Factory的结构分层

从最底层开始介绍:

底层加速层:已适配的硬件有英伟达、华为、海光

软件层和计算架构层:以PyTorch和Hugging Face 为主。DeepSpeed作为分布式多卡加速训练的方案(使用了一种数据并行的方案,平均分配所有参数到多卡环境中),在工程架构实现以后,未来将支持更多的框架。

大模型和数据层:更高层的抽象,比如Flash Attention RoPE Scaling,以及数据上的集成,比如动态集成数据

高效训练层:预训练/继续训练 + 监督微调 + PPO算法

如何优化每一步的加速过程

简单回顾一下大模型的训练方式,首先是InstructGPT论文,它使用到的训练方法和GPT是一样,称之为三阶段训练方法

三个阶段可以分为问答阶段、模型奖励阶段、强化学习训练

  1. 问答阶段:利用人类标记好的数据,进行指令的微调工作。在这一阶段,模型就初具了能够听懂人类指令完成对话的基本功能。它存在的弊端是,当下的模型依然含有有害内容,而且相对输出的内容过短。
  2. 设计模型奖励阶段:这一步主要任务就是收集数据,构造一个已标注过的文本对比数据集。建立一个奖励模型(RM,reward model),该奖励模型可以对每一个给定的问题和这个问题的答案输出一个标量的分数。
  3. 强化学习:通过这个得到的标量分数,通过PPO算法去微调问答阶段得到的微调模型(偏好对齐)。

通过以上三步,加上LLaMA-Factory的一些技术,使得这一套训练工作流比原始训练更加高效。

预训练

在开始将解SFT之前,首先讲解一下预训练。

预训练是模型训练过程中必要的一个阶段,在有大量无标注的语料情况下,我们希望可以先不进行标注。通过大量的无标注数据的预训练方法,使得模型在我们相关领域的数据上表现的更好(思考为什么?)。

这里我们使用类似GPT3的自回归模型(auto-regressive language model),根据前i-1个tokens来预测第i个token的能力。

在持续训练的过程中,使用传统的因果建模的方式,通过计算交叉熵损失的方式来最小化下面这个对数似然的损失来实现模型的优化。
$$
\mathcal{L}_ {\mathrm{CLM}}=-\frac{1}{n} \sum_ {i=1}^{n} \log P\left(x_ {i} \mid x_ {1}, x_ {2}, \ldots, x_ {i-1} ; \theta\right)
$$
它的优势是:通过无监督的方式来完成一系列优化。

预训练阶段的语料处理

这里,处理这种预训练阶段的数据最高效的方法有两种,分别是序列打包方法(sequence packing)和数据的流式加载的(data streaming),下面对二者分别进行介绍。

序列打包方法

将不同长度的序列打包在一起,省去padding的过程。如图中的过程,左侧的序列是没有打包过的序列,右边序列是打包过的序列。如果我们假设 batch_size = 1,则处理左边的序列需要4次;处理右边的序列只需要3次。

问:打包在一起的语料之间有干扰怎么办?

答:会有干扰,但是影响极小,带来的训练效率提升却是显著的。

通常情况下,默认会使用 sequence packing 的方法处理序列数据。

数据的流式加载

下面的这一种数据流式加载技术,主要是为了极大的洁身CPU的内存。通常,我们预期的阶段需要用到几百GB的内存,在使用了数据的流式加载技术后可以节省巨大的内存占用,同时可以避免在分布式训练上造成超时的问题。

监督微调

在经历了一系列优化工作后,监督微调(SFT,Supervised Fine-Tuning)通过问答对的形式进行一个 sequence to sequence 的学习。

这一阶段的方法和预训练阶段的损失是一样的,采用交叉熵损失。
$$
\mathcal{L}_ {\mathrm{SFT}}=-\frac{1}{n} \sum_ {i=1}^{n} \log P\left(y_{i} \mid x_ {1}, x_ {2}, \ldots, x_ {m} ,y_ {1}, y_ {2}, \ldots, y_ {i-1} ; \theta\right)
$$
在公式中,x代表问答对中的问题,y代表问答对中的回复。在计算损失的过程中,只计算y上面的损失,对于x一般是不计算损失的。

在LLaMA-Factory中,主要就是针对这个多轮对话的数据集去优化sft的loss计算方式。

在多轮对话数据集中,只要计算一次前向传播。

在模型输出的损失函数计算中,将终止标签eos标注添加到标准输出中,这样可以帮助模型学习到在什么位置进行停止。提高5%到10%的运算效率。

奖励模型

奖励模型使用 pairwise ranking loss 来计算损失:
$$
\mathcal{L}_ {\mathrm{RM}}=-\sum_ {(x,y^c,y^r)}\log\left(\sigma(r_ \theta(x,y^c)-r_ \theta(x,y^r))\right)
$$
其中,yc代表了更好的回复,yr代表了更差的回复,r就是对答案的打分。这个函数希望对好回答的打分高于差回答的打分。

在训练过程中,会将所有的r1、r2、r3的分数一起作为参考求平均值。

在推理过程中,只会参考r3的分数最为推理时候的分数。

思考一个问题:为什么训练的过程中需要将句子中每一个token进行采取然后做分数计算,而不能只计算最后一个token的分数?

答案是:这样会使得模型过拟合。

以上是对奖励模型学习过程中的优化过程:针对最后一个输出token进行打分,用于评估奖励模型效果。

强化学习(PPO学习)

PPO学习参考一个多臂老虎机问题(bandit problem)。在大模型的过程中,会对一个Prompt产生一条随机的回复,然后奖励这个模型,并对生成的回复进行打分。PPO模型就是为了希望最大化的奖励这个模型给出的分数。

除了奖励以外,还会对当前这个模型进行KL惩罚以防止PPO对模型的过度优化现象。
$$
\mathcal{J}(\phi)=\mathbb{E}_ {(\mathrm{x},\mathrm{y})\sim\mathcal{D}_ {\pi_\phi^{\mathrm{PPO}}}}\left[r_ \theta(x,y)-\beta\log\frac{\pi_\phi^{\mathrm{PPO}}(y|x)}{\pi^{\mathrm{SFT}}(y|x)}\right]
$$
在上面损失函数中,log后分子上的数学建模是PPO模型,分母上的建模是SFT模型。

因此,可以看到,如果想要进行PPO模型的训练,就需要加载三个模型。如果只有一张显卡,就很难把三个模型都运行起来。

模型权重共享的方法

在三个模型中,参数的大头在backbone部分,这部分的参数和空间占了整个模型空间的98%左右。如果多次读取模型,最冗余的状态就是backbone的重复加载和卸载。

这里我们引入LoRA适配器的方法,共享backbone部分参数,在训练过程中动态加载LoRA和adapter模块。

PPO算法拥有一些强化学习的通病,就是难收敛。这里可以推荐一个替代的强化学习算法DPO,公式上的差别较小:
$$
\mathcal{L}_ {\mathrm{DPO}}=-\sum_ {(x,y^c,y^{\prime})\sim\mathcal{D}}\left[\log\sigma\left(\beta\log\frac{\pi^{\mathrm{DPO}}(y^c|x)}{\pi^{\mathrm{SFT}}(y^c|x)}-\beta\log\frac{\pi^{\mathrm{DPO}}(y^r|x)}{\pi^{\mathrm{SFT}}(y^r|x)}\right)\right]
$$

LLaMA-Factory在应用上的优化

在强化学习的两个算法中,DPO收敛更快,但是更容易过拟合;PPO效果好,但是不容易收敛/收敛慢。

在LLaMA-Factory上提出了混合微调梯度的方法。

将sft损失乘一个损失因子γ后,将整体值加到DPO中。由于它是一个常数项,所以并不会增加梯度计算的负担

实现工具调用训练

使得大模型具有调用外部工具的能力(LangChain)。

如果咨询大模型:当前时刻某个酒店有多少客房

如果直接让大模型回答相关问题,大模型既不知道最新的知识,也不清楚没见过的专有名词(酒店的名称)。

而通过LangChain等方法,就可以将整个工具链上附带工具化调用的能力,其实就是格式化大模型输出的能力。将模型的输出格式,变成可以随着工具进行传参的格式。

在原有论文的基础上,将传入参数的格式变成json格式传入,这样就可以学习到一些嵌套结构例如列表等多参数的传递。

工具链的学习过程没有特别多的技巧,就是一个多轮对话的sft学习,重要的是对标准的传递格式的学习。

使用Flash Attention加速优化

一种计算上无损的加速方式,推理速度提升2~3倍,主要在硬件层实现,其原理是提升了I/O的执行效率。

它的替身效率随上下文长度的提升而显著的提升。

模型量化

GPTQ的方式可以将模型的权重量化到4bit的程度

位置编码

困惑度是指将模型的输入扩充超过最长tokens的限制,那么模型对于输入几乎会丧失所有的推理能力。

位置编码可以将长序列的困惑度大幅度降低。

这里介绍两种插值方法:线性差值(Linear Scaling)和 NTK-aware 差值。

损失函数的优化

在LoRA自身的反向传播过程中,PyTorch的反向传播实现得不够好。

例如,对下面A计算梯度,它是一个m×d维的矩阵,乘m×h维的矩阵,乘h×r维的矩阵。

在PyTorch中,它会先将大矩阵相乘,然后再去乘小矩阵。

而如果先将两个较小的矩阵进行相乘,最后再乘较大的矩阵,计算量是大幅度缩小的。(矩阵的乘法交换律)

那么,通过替换原有的PyTorch中的的梯度计算方法,可以大幅度减少需要计算的浮点运算数量。

计算稳定性上的优化

fp16精度下的训练有不稳定的情况,bf16相比前者更加稳定。但是由于精度上的差异,导致二者在训练和运行推理时的兼容性上有一定的问题。

  • bf16精度训练的模型,可以运行在fp16精度的环境下
  • fp16精度训练的模型,不可以运行在bf16精度的环境下

如果在LLaMA-Factory的环境上强制运行后者,将会得到一个warning警告。

更好的工具调用和总结能力

实例展示了通过相应的数据集在LLaMA-Factory上进行训练后,模型不仅可以正确的理解输入的工具调用指令,还能生成正确对应格式的结果。

Trick

继续预训练(continue pre-training)相关

  • 继续预训练最好在base模型上进行(不要用chat版)
  • Scaling law 是非常必要的
  • 参考前人最常用的超参数设置可以事半功倍
  • 需要竟可能大数量级的tokens去继续训练过程
  • 多阶段训练可能会提高模型的能力
  • 词表扩充有可能害了训练结果(除非你需要换小语种)

监督微调(SFT)相关

  • 数据质量很重要(合理使用GPT-4去标注知识)
  • 数据集非常小(只有几百条),最好直接使用chat模型做微调并且保持和数据集模板的对齐
  • 为了保证微调过程中那个,模型原有的能力不会损失太明显,还是需要再微调的过程中加上一些通用领域的知识。
  • Qlora有时币Lora效果更好(更大幅度的降低参数,带来了更稳定的模型)
  • (一般植入新知识的过程都在继续预训练的过程中进行)如果一定要在监督微调的过程中植入新知识,那么就需要加入“同意不同文(或不同格式)”的几十条数据。例如:在监督微调部分修改模型的自我认知,就需要几百条数据来完成一项知识的植入。
  • 超参数没有最好的,只有实践才知道(实践出真知)

答疑

问:Flash Attention没法在V100上用,有什么好的办法?

答:在PyTorch2.0之后,已经实现了替代了FlashAttention的轮子。

问:超参数设置怎么来?

答:参考别的论文。

问:针对不同的response做优化,答案长度不同的情况下,如何做归一化呢?

答:不同长度的情况下,会使用tokens填充的方法,将较少tokens的向量做padding,最终补齐长度和较长的tokens一致。