如何通过打表找到训练大型模型的最佳并行配置？AI Infra论文解读——高效LLM训练分析

这篇paper的干货在这一节展示。

0x5.1 Fused Kernels 和 Flash Attention

paper里对Fasl Attention进行了更大规模模型的测试，此外还测试了一下Meagtron-LM自带的fused softmax attention kernel以及Faslh Attention v2里面的Fused RMSNorm kernel。

0x5.1.1 Attention

Figure1中展示了不同注意力实现的扫描结果，可以看到MFU最高的配置是Flash Attention2+RMS Kernel，然后micro-size=TP=PP=1，这样可以达到最佳MFU。需要特别注意的是Figure1中的测试都没有使用Activation Checkpointing（但Falsh Attention内部自己是有重计算的）。

0x5.1.2 RMSNorm Kernel

在Figure1厘米还展示了Flash Attention库里面优化的RMSNorm kernel在在训练效率上得到了很多的提升，和没用这个RMS Kernel相比可以提升14个百分点的MFU。当然有很大的原因是因为这个kernel更加节省内存，可以让并行配置从pp2变成纯数据并行，减少了通信开销。这里还提了一句，RMSNorm和Activation Checkpointing一起用会出错。

0x5.2 Activation Checkpointing

在Figure2里面报告了不同模型大小最佳配置的MFU，包含对每一层使用以及不使用Activation Checkpointing。结论就是在不使用Activation Checkpointing并且通过较小的micro-batch或者更高程度的模型并行来补偿因此产生的内存成本，可以实现最佳的训练MFU。这里没有使用Flash Attention提供的RMSNorm Kernel。

对于LLAMA 30B模型，序列长度为8k，在训练过程中即使TP Size大小达到4，PPO Size达到16，Activation Checkpointing也是将模型装入内存的必要条件。因为LLAMA 30B模型有52个注意力头，不能被8整除，所以我们无法增加TP Size。在paper的4.1节中，展示了添加RMSNorm kernel进一步降低了所需内存，使得Activation Checkpointing不再必要。在这种情况下，我们再次看到，不使用Activation Checkpointing的设置实现了最佳吞吐量。

另外，这里提到Flash Attention由于本身就自带重计算，所以一个更加合理的方法是对于一个Transformer Block来说，只对FFN这部分用重计算，避免多余的重复计算。

0x5.3 Micro-Batch 大小

这一小节评估了Micro-Batch 大小与所需模型（张量或流水线）并行度和Activation Checkpointing之间的权衡。之前已经有工作基准测试了不同的Micro-Batch大小，并固定了张量和流水线并行度，显示出更大的Micro-Batch大小带来更高的吞吐量。但是，较小的Micro-Batch大小可能使得不同的、更高效的并行化配置成为可能。此外，张量和流水线并行度在实践中也往往不是固定的。

在图3中，paper展示了评估的每种模型类型的最佳表现（Activation Checkpointing、张量并行度大小、流水线并行度大小）配置。为了公平地评估 Activation Checkpointing，没有包括使用 Flash Attention2 的 RMSNorm Kernel，因为该kernel与Activation Checkpointing结合使用时会导致错误。最终发现对于所有模型类型，Micro-Batch大小为1的情况下达到了最佳MFU。一般来说是：Micro-Batch大小越小，MFU越好。使用大于2的Micro-Batch大小时，序列长度为8k的模型在任何并行配置下都无法塞进内存。

因此，这里得出结论，选择Micro-Batch大小为1在大多数情况下是有益的。Micro-Batch大小为1的卓越性能可归因于以下三个因素。

• 最小程度的模型并行化：最高效的训练通常需要最少量的模型（张量或流水线）并行化，这在Micro-Batch大小最小时实现。
• 避免Activation Checkpointing：对于某些模型（例如，LLAMA 65B），微批量大小为1是唯一允许无Activation Checkpointing训练的配置。正如前一节所讨论的，不使用Activation Checkpointing通常允许最高吞吐量的配置。LLAMA 30B 8k模型在不使用RMSNorm kernel的情况下无法装入内存。
• 减少流水线bubble时间：较小的Micro-Batch大小减少了在批处理开始和结束时在流水线bubble中花费的时间。这里使用的是1F1B流水线并行调度算法。Meagtron支持这种实现。

0x5.4 Tensor Parallelism和Pipline Parallelism

这一节的内容主要就是看Figure4，结论就是流水并行比张量并行更容易获得更高的MFU，应该优先流水并行。

0x5.5 Sequence Parallelism

这一节就是看Figure5，对于LLAMA 13B和30B模型，序列长度为2k时，最佳配置是不使用任何张量并行，因此激活 Sequence Parallelism 没有效果。在序列长度为8k的13B模型中，最佳配置采用了2的张量并行大小；然而，没有观察到训练效率的提升。对于最大的模型和序列长度，即序列长度为8k的30B和65B模型，我们可以看到在使用Sequence Parallelism时有2-6个百分点的MFU提升。在这两种情况下，Sequence Parallelism由于减少了内存需求，使得模型并行化程度降低，从而提高了训练效率。

结论是，当与本工作中探索的其他几种优化结合使用时，Sequence Parallelism的使用只有在模型大小超过30B参数或序列长度超过2k时，才能显著提高训练效率。

0x5.6 端到端性能

实际上就是对下面这张表的解读：

这里的AA-Scaling是作者团队内部的基于Megatron-LM开发的仓库，然后就是说基于本文搜索的最佳并行配置以及相关技术组合可以获得相比于公开框架和模型更高的MFU。