前言
SFT 可以说是 LLM 的基本操作了,如果只是想把 SFT 跑起来是非常简单的,只需要构造 input_ids 和 labels,然后就可以把训练跑起来。然而,这样的训练效率实际上非常低。
所以在训练时,通常有两个加速方法:
- 多轮合并
- packing
无论是哪种方法,加速后都需要保证 loss 和原来是等价的。本文主要介绍这两种加速方法,以及 loss 计算时遇到的问题。
1、多轮合并
假设我们有一个对话,其中 user 和 bot 交互了 3 轮,我们可以构建三个样本:
input_ids 就是对应的 token id,labels 输入部分(白色)使用 -100,输出部分(绿色)使用 input_ids。
这样计算的 loss 可以表示为:
其中 l_i 表示第 i 个样本的 loss,n_i 表示第 i 个样本输出的 token 数量 (对应绿色部分)。
这样除了训练比较慢,没有什么别的问题。因为不同样本之间有很多重复计算的前缀,实际上这部分计算一次就行。
2、加速计算
如果将三轮三个样本合并成一个样本,可以尝试这种构造形式。
因为存在 causal attention mask,所以每个 token 只能看到前面的 token,计算上和之前是等价的。
但是这样有一个坑:如果还是按照刚才的方式构建 input_ids 和 labels (白色用-100,绿色用input_ids)loss 计算是有问题的。
pytorch CrossEntropyLoss 计算 loss 按照下面的方法,默认是"mean"。
所以我们会得到这样的 loss:
当不同轮次的输出长度不同时,这种 loss 和刚才的不等价。多轮对话中输出较短的权重被降低了,输出较长的被提高了。所以结果就是短输出的数据训练不够充分。
3、Packing
假设我们有两个对话,第一个是单轮对话,第二个是三轮对话。
正确的 loss:
其中 l_ij 表示第 i 个样本第 j 轮对话的 loss,n_ij 同理。
问题:真实场景中的训练集文本长度长短不一,Padding 后矩阵非常稀疏,只有不到一半是有效计算。
加速计算:
将所有样本拼接成一条,并且加入 attention mask 保证后面的样本看不见前面的 token。
比如在 flash attention 中,可以调用 flash_attn_varlen_qkvpacked_func,并传入 cu_seqlens 参数。
和之前一样,如果不修改 loss 计算方法,packing 的样本之间会存在因为长度不同,导致训练不充分的问题。
4、正确方法
一般情况下,loss 计算会经历三次平均:
- micro batch 维度,分母是这个 micro batch 中的所有 label 不是 -100 的 token 数
- DP 维度,分母是 DP size (和GPU数量相关)
- 梯度累加维度,分母是梯度累加数
我们这里要做的就是禁用这三个平均,统一用这个 global batch 的对话轮数作为分母。
在新版 megatron 框架中,开启开关 --calculate-per-token-loss 即可禁用 DP 和梯度累加的平均,然后修改 loss_func。
每个 micro batch 都需要返回这个 micro batch 的轮数,最后框架会自动将所有轮数求和,作为分母。对于分子,需要除以这个轮次的token 数。
正确实现代码如下(loss_token_num, turn_num 是在构建 data 的时候构建的):
def loss_func(output_tensor, loss_mask, loss_token_num, turn_num):
losses = output_tensor.view(-1).float()
loss_mask = loss_mask.view(-1).float()
loss_token_num = loss_token_num.view(-1).float()
# label: [-100, -100, a, a, a, -100, b, b, -100, -100, c, c, c, -100, -100]
# loss_mask: [0, 0, 1, 1, 1, 0, 1, 1, 0, 0, 1, 1, 1, 0, 0]
# losses: [a0, a1, a2, a3, a4, b0, b1, b2, c0, c1, c2, c3, c4, d0, d1]
# losses * loss_mask = [0, 0, a2, a3, a4, 0, b1, b2, 0, 0, c2, c3, c4, 0, 0]
# loss_token_num: [3, 3, 3, 3, 3, 2, 2, 2, 3, 3, 3, 3, 3, 1, 1]
# losses * loss_mask / loss_token_num = [0, 0, a2/3, a3/3, a4/3, 0, b1/2, b2/2, 0, 0, c2/3, c3/3, c4/3, 0, 0]
# sum = 1/3 (a2 + a3 + a4) + 1/2 (b1 + b2) + 1/3 (c2 + c3 + c4)
loss = torch.sum(losses * loss_mask / loss_token_num)
loss_and_turn_num = torch.cat([loss.view(1), turn_num.view(1)])
# Reduce loss for logging.
loss_and_turn_num = loss_and_turn_num.clone().detach()
torch.distributed.all_reduce(loss_and_turn_num, group=mpu.get_data_parallel_group())
# 新版返回结构,开启 calculate_per_token_loss 开关后,返回三个值
# 第一个是反向传播实际使用的 loss, 所有 packing 的 loss 求和
# 第二个是 turn_num, 优化器状态更新时会使用对这个值求和然后缩放梯度
# 第三个是用于日志打印的 loss, 包含两个值,第一个是所有 loss 求和作为分子,第二个是所有 turn_num 求和作为分母
return loss, turn_num, {"lm loss": (loss_and_turn_num[0], loss_and_turn_num[1])}
5、总结
在 SFT 时,如果要加速,需要注意:
- 不同样本之间是等价的;
- 不同轮次之间也是等价的。
在合并多轮 / packing 时,需要修改 loss 计算方法,为每个 token 设置正确的权重,并且关闭 DP / 梯度累加的平均。
最后
为了助力朋友们跳槽面试、升职加薪、职业困境,提高自己的技术,本文给大家整了一套涵盖AI大模型所有技术栈的快速学习方法和笔记。目前已经收到了七八个网友的反馈,说是面试问到了很多这里面的知识点。
面试题展示
1、请解释一下BERT模型的原理和应用场景。
答案:BERT(Bidirectional Encoder Representations from Transformers)是一种预训练的语言模型,通过双向Transformer编码器来学习文本的表示。它在自然语言处理任务中取得了很好的效果,如文本分类、命名实体识别等。
2、什么是序列到序列模型(Seq2Seq),并举例说明其在自然语言处理中的应用。
答案:Seq2Seq模型是一种将一个序列映射到另一个序列的模型,常用于机器翻译、对话生成等任务。例如,将英文句子翻译成法文句子。
3、请解释一下Transformer模型的原理和优势。
答案:Transformer是一种基于自注意力机制的模型,用于处理序列数据。它的优势在于能够并行计算,减少了训练时间,并且在很多自然语言处理任务中表现出色。
4、什么是注意力机制(Attention Mechanism),并举例说明其在深度学习中的应用。
答案:注意力机制是一种机制,用于给予模型对不同部分输入的不同权重。在深度学习中,注意力机制常用于提升模型在处理长序列数据时的性能,如机器翻译、文本摘要等任务。
5、请解释一下卷积神经网络(CNN)在计算机视觉中的应用,并说明其优势。
答案:CNN是一种专门用于处理图像数据的神经网络结构,通过卷积层和池化层提取图像特征。它在计算机视觉任务中广泛应用,如图像分类、目标检测等,并且具有参数共享和平移不变性等优势。
6、请解释一下生成对抗网络(GAN)的原理和应用。
答案:GAN是一种由生成器和判别器组成的对抗性网络结构,用于生成逼真的数据样本。它在图像生成、图像修复等任务中取得了很好的效果。
7、请解释一下强化学习(Reinforcement Learning)的原理和应用。
答案:强化学习是一种通过与环境交互学习最优策略的机器学习方法。它在游戏领域、机器人控制等领域有广泛的应用。
8、请解释一下自监督学习(Self-Supervised Learning)的原理和优势。
答案:自监督学习是一种无需人工标注标签的学习方法,通过模型自动生成标签进行训练。它在数据标注困难的情况下有很大的优势。
9、解释一下迁移学习(Transfer Learning)的原理和应用。
答案:迁移学习是一种将在一个任务上学到的知识迁移到另一个任务上的学习方法。它在数据稀缺或新任务数据量较小时有很好的效果。
10、请解释一下模型蒸馏(Model Distillation)的原理和应用。
答案:模型蒸馏是一种通过训练一个小模型来近似一个大模型的方法。它可以减少模型的计算和存储开销,并在移动端部署时有很大的优势。
11、请解释一下LSTM(Long Short-Term Memory)模型的原理和应用场景。
答案:LSTM是一种特殊的循环神经网络结构,用于处理序列数据。它通过门控单元来学习长期依赖关系,常用于语言建模、时间序列预测等任务。
12、请解释一下BERT模型的原理和应用场景。
答案:BERT(Bidirectional Encoder Representations from Transformers)是一种预训练的语言模型,通过双向Transformer编码器来学习文本的表示。它在自然语言处理任务中取得了很好的效果,如文本分类、命名实体识别等。
13、什么是注意力机制(Attention Mechanism),并举例说明其在深度学习中的应用。
答案:注意力机制是一种机制,用于给予模型对不同部分输入的不同权重。在深度学习中,注意力机制常用于提升模型在处理长序列数据时的性能,如机器翻译、文本摘要等任务。
14、请解释一下生成对抗网络(GAN)的原理和应用。
答案:GAN是一种由生成器和判别器组成的对抗性网络结构,用于生成逼真的数据样本。它在图像生成、图像修复等任务中取得了很好的效果。
15、请解释一下卷积神经网络(CNN)在计算机视觉中的应用,并说明其优势。
答案:CNN是一种专门用于处理图像数据的神经网络结构,通过卷积层和池化层提取图像特征。它在计算机视觉任务中广泛应用,如图像分类、目标检测等,并且具有参数共享和平移不变性等优势。
16、请解释一下强化学习(Reinforcement Learning)的原理和应用。
答案:强化学习是一种通过与环境交互学习最优策略的机器学习方法。它在游戏领域、机器人控制等领域有广泛的应用。
17、请解释一下自监督学习(Self-Supervised Learning)的原理和优势。
答案:自监督学习是一种无需人工标注标签的学习方法,通过模型自动生成标签进行训练。它在数据标注困难的情况下有很大的优势。
18、请解释一下迁移学习(Transfer Learning)的原理和应用。
答案:迁移学习是一种将在一个任务上学到的知识迁移到另一个任务上的学习方法。它在数据稀缺或新任务数据量较小时有很好的效果。
19、请解释一下模型蒸馏(Model Distillation)的原理和应用。
答案:模型蒸馏是一种通过训练一个小模型来近似一个大模型的方法。它可以减少模型的计算和存储开销,并在移动端部署时有很大的优势。
20、请解释一下BERT中的Masked Language Model(MLM)任务及其作用。
答案:MLM是BERT预训练任务之一,通过在输入文本中随机mask掉一部分词汇,让模型预测这些被mask掉的词汇。