大模型并行训练指南:通俗理解Megatron-DeepSpeed之模型并行与数据并行(下)
共 7557字,需浏览 16分钟
·
2024-04-18 20:11
文末《大模型项目开发线上营》秒杀倒计时↓↓↓
接前文:
(上)篇>>>大模型并行训练指南:通俗理解Megatron-DeepSpeed之模型并行与数据并行(上)
06
BF16Optimizer/CUDA 融合核函数/数据集/嵌入LayerNorm/位置编码
6.1 优先使用BF16 Optimizer,而非FP16
BF16Optimizer,
如果您不熟悉这种数据格式,请查看它的 位布局
BF16 格式的关键是它的指数位数与 FP32 相同,因此不会溢出,但 FP16 经常溢出FP16 的最大数值范围为 64k,您只能进行较小数的乘法。例如你可以做 250*250=62500,但如果你尝试 255*255=65025,你就会溢出,这是导致训练出现问题的主要原因。这意味着你的权重必须保持很小。一种称为损失缩放(loss scaling)的技术有助于缓解这个问题,但是当模型变得非常大时,FP16 较小的数值范围仍然是一个问题
BF16 没有这个问题,你可以很容易地做 10_000*10_000=100_000_000有意思是的,我司七月审稿项目组在用一万多条paper-review数据微调llama2 7B时也遇到了这个问题,即
“原本的LongQLoRA源码训练参数指定使用fp16数据类型,但是训练可能会很不稳定 具体而言,可能会出现数值溢出问题,初期就出现loss严重震荡,甚至高达上百,并在数步后返回模型已收敛(但实际并未)的提示而中断训练,故要设置使用bf16数据类型进行训练,loss就能稳定从4.多开始收敛”
更多详见此文《七月论文审稿GPT第2版:用一万多条paper-review数据集微调LLaMA2 7B最终反超GPT4》的5.1节LLaMA2 7b chat + LongQLoRA训练
-
当然,由于 BF16 和 FP16 的大小相同,均为 2 个字节,因此,没有免费的午餐,当使用 BF16 时,代价就是它的精度非常差。然而,你应该还记得我们在训练时采用的随机梯度下降法及其变体,该方法有点像蹒跚而行,如果你这步没有找到完美的方向其实没关系,你会在接下来的步骤中纠正自己 -
无论使用 BF16 还是 FP16,都有一个权重副本始终在 FP32 中 —— 这是由优化器更新的内容。因此 16 位格式仅用于计算,优化器以全精度更新 FP32 权重,然后将它们转换为 16 位格式以用于下一次迭代 所有 PyTorch 组件都已更新,以确保它们在 FP32 中执行任何累加,因此不会发生精度损失 一个关键问题是梯度累积,它是流水线并行的主要特征之一,因为每个 micro batch 处理的梯度都会累积。在 FP32 中实现梯度累积以保证训练的精确性至关重要,这正是
BF16Optimizer
所做的
除了其他改进之外,我们认为使用 BF16 混合精度训练将潜在的噩梦变成了一个相对平稳的过程,这可以从以下 lm 损失图中看出:
6.2 CUDA 融合核函数
GPU 主要做两件事。它可以将数据写到显存或从显存读数据,并对这些数据执行计算。当 GPU 忙于读写数据时, GPU 的计算单元就会空闲。如果我们想有效地利用 GPU,我们希望将空闲时间降至最低
-
核函数是一组实现特定 PyTorch 操作的指令。例如,当你调用 torch.add 时,它会通过一个 PyTorch 调度器,它会根据输入张量及其他变量的取值来决定它应该运行哪些代码,最后运行它。CUDA 核函数使用 CUDA 来实现这些代码,因此只能在 NVIDIA GPU 上运行 比如举个例子,当使用 GPU 计算 c = torch.add (a, b); e = torch.max ([c,d]) 时
一般情况下,PyTorch 将执行的操作是启动两个单独的核函数,一个执行 a 和 b 的加法,另一个执行取 c 和 d 两者的最大值。在这种情况下,GPU 从其显存中获取 a 和 b,执行加法运算,然后将结果写回显存(一次对a b的读取 一次对c的写入)。然后它获取 c 和 d 并执行 max 操作,然后再次将结果写回显存(再一次对c d的读取 一次对max[c d]结果的写入)
但如果我们要融合这两个操作,即将它们放入一个 “融合核函数” 中,然后启动那个内核,我们不会将中间结果 c 写到显存中,而是将其保留在 GPU 寄存器中(仅一次对 a b的读取),并且仅需要获取 d 来完成最后的计算(一次对d的读取 一次对max[c d]结果的写入)。这节省了大量开销并防止 GPU 空闲,因此整个操作会更加高效
融合核函数就是这样。它们主要将多个离散的计算和进出显存的数据移动替换为有很少数据移动的融合计算。此外,一些融合核函数会对操作进行数学变换,以便可以更快地执行某些计算组合July注:这其实有点类似flash attention的思想,详见此文:通透理解FlashAttention与FlashAttention2:全面降低显存读写、加快计算速度 -
为了快速高效地训练 BLOOM,有必要使用 Megatron-LM 提供的几个自定义 CUDA 融合核函数。特别地,有一个 LayerNorm 的融合核函数以及用于融合缩放、掩码和 softmax 这些操作的各种组合的核函数。Bias Add 也通过 PyTorch 的 JIT 功能与 GeLU 融合。这些操作都是瓶颈在内存的,因此将它们融合在一起以达到最大化每次显存读取后的计算量非常重要。因此,例如,在执行瓶颈在内存的 GeLU 操作时同时执行 Bias Add,运行时间并不会增加。这些核函数都可以在 Megatron-LM repository 代码库 中找到
6.3 数据集
Megatron-LM 的另一个重要特性是高效的数据加载器。在首次训练启动前,每个数据集中的每个样本都被分成固定序列长度(BLOOM 为 2048)的样本,并创建索引以对每个样本进行编号。基于训练超参,我们会确定每个数据集所需要参与的 epoch 数,并基于此创建一个有序的样本索引列表,然后打乱它
举个例子,如果一个数据集中有 10 个样本并应参与 2 个 epoch 的训练,则
-
系统首先按 [0, ..., 9, 0, ..., 9] 顺序排好样本索引,然后打乱该顺序为数据集创建最终的全局顺序 请注意,这意味着训练不会简单地遍历整个数据集然后重复,你有可能在看到另一个样本之前看到同一个样本两次,但在训练结束时模型将只看到每个样本两次。这有助于确保整个训练过程中的训练曲线平滑。这些索引,包括每个样本在原始数据集中的偏移量,被保存到一个文件中,以避免每次开始训练时都重新计算它们 -
最后,可以将其中几个数据集以不同的权重混合到训练最终使用的数据中
6.4 嵌入LayerNorm
在我们努力阻止 104B 模型发散的过程中,我们发现在第一个层词嵌入层之后添加一个额外的 LayerNorm 可以使训练更加稳定
StableEmbedding
操作,它是一个带有 LayerNorm 的普通嵌入,其使用均匀 xavier 函数来初始化
6.5 位置编码
07
其他事项:训练中的困难与重要资源
在训练开始之前,有很多问题需要弄清楚。特别是,我们发现了几个问题,这些问题只有在我们开始在 48 个节点上进行训练后才会出现,而不会在小规模时出现
例如,需要设 CUDA_LAUNCH_BLOCKING=1 来防止框架挂起,我们需要将优化器组分成更小的组,否则框架会再次挂起。你可以在训前编年史 中详细了解这些内容
训练期间遇到的主要问题类型是硬件故障。由于这是一个拥有大约 400 个 GPU 的新集群,平均每周我们会遇到 1-2 个 GPU 故障。我们每 3 小时 (100 次迭代) 保存一个检查点。因此,我们每周因硬件崩溃平均损失 1.5 小时的训练成果。Jean Zay 系统管理员随后将更换有故障的 GPU 并恢复节点。与此同时,我们有备用节点可供使用。
我们还遇到过多次导致 5-10 小时停机的各种其他问题,其中一些与 PyTorch 中的死锁错误有关,另一些则是由于磁盘空间不足。如果您对具体细节有兴趣,请参阅 训练编年史
在对训练这个模型进行可行性分析时,所有这些停机时间都被计划在内了,我们也据此选择了合适的模型大小和我们希望模型消耗的数据量。因此,即使存在这些停机问题,我们还是成功地在预计时间内完成了训练。如前所述,它需要大约 100 万个计算时才能完成
另一个问题是 SLURM 并非设计为供一组人使用。SLURM 作业由单个用户拥有,如果他们不在身边,则该组的其他成员无法对正在运行的作业执行任何操作。我们制定了一个终止方案,允许组中的其他用户终止当前进程,而不需要启动该进程的用户在场。这在 90% 的问题上都很有效。如果 SLURM 设计者读到这篇文章,请添加一个 Unix 组的概念,这样一个 SLURM 作业就可以由一个组拥有
由于训练是全天候 24/7 进行的,我们需要有人随叫随到,但由于我们在欧洲和加拿大西海岸都有人,因此不需要有人携带传呼机,我们能很好地互相备份。当然,周末的训练也得有人看着。我们自动化了大部分事情,包括自动从硬件崩溃中恢复,但有时仍需要人工干预
7.2 重要链接
-
主训练文档 -
Tensorboard -
训练用的 slurm 脚本 -
训练记录
7.3 论文与文章
我们不可能在本文中详细解释所有内容,因此如果此处介绍的技术激起你的好奇心,使你想了解更多信息,请阅读以下论文:
-
Efficient Large-Scale Language Model Training on GPU Clusters. -
Reducing Activation Recomputation in Large Transformer Models
-
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models -
ZeRO-Offload: Democratizing Billion-Scale Model Training -
ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning -
DeepSpeed: Extreme-scale model training for everyone
-
Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model.
-
Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation -
What Language Model to Train if You Have One Million GPU Hours? - 你会在那里找到最终选择 ALiBi 的实验
-
8-bit Optimizers via Block-wise Quantization (使用了该论文中的嵌入 LaynerNorm,但是论文的其他部分及其技术也很妙,没用 8 位优化器的唯一原因是我们已经使用 DeepSpeed-ZeRO 节省了优化器内存)
课程进度:已直播前27课(都有回放),每周两到三次课,均在晚8-10点上课。
校长July寄语
“
为何讲大模型商用项目的课程很罕见
一方面,即便在大厂,虽有技术但没法讲其内部项目,而专门搞应用开发的小厂,很难将其赖之生存的项目拿出来讲
二方面,一般职业讲师 背后没有项目团队 只能搞个demo,至于一般教育机构也很难再搞个项目团队,成本大 招人难 做出成果更难
而我司教育团队外,我司去年专门成立了大模型项目开发团队(由我带全职和在大厂的兼职组成)
一者,for b端做一系列大模型解决方案,正在逐
一上线一系列大模型产品(每个项目均由各个项目组迭代好几个月)
二者,很多项目开发/落地经验、数据/代码都放到了这个线上营中,如表现突出邀请加入至对应项目组且给报酬 ”
① 一年GPU,封装了诸如ChatGLM3等各大主流大模型
② 一个VIP年卡「200多个AI小课、4个 大模型小课(即ChatGPT原理、类ChatGPT微调实战、SD及其二次开发、垂直大模型实战)」