AI Infra论文阅读之通过打表得到训练大模型的最佳并行配置

最新推荐文章于 2026-01-07 10:31:23 发布

原创

最新推荐文章于 2026-01-07 10:31:23 发布 · 2.1k 阅读

15 ·

CC 4.0 BY-SA版权

文章标签：

#人工智能 #论文阅读

0x0. 前言

这次阅读一篇mlsys的一篇新论文，《Efficient Parallelization Layouts
for Large-Scale Distributed Model Training》，虽然该论文还处于open review阶段，但作者在 Megatron-LM 基础上通过对各种训练调优方法进行组合打表获得了一些比较有趣的结论和发现，地址：https://openreview.net/pdf?id=Y0AHNkVDeu 。作者的详细实验表格在：https://github.com/Aleph-Alpha/NeurIPS-WANT-submission-efficient-parallelization-layouts

paper中基于大量搜索实验，最后得出来以下几个训练大模型的建议：

使用Micro Batch大小为1，以实现最小程度的模型并行化，避免Activation Checkpointing，并减少流水线bubble。
优先使用TP和PP，而不是Activation Checkpointing。
只有在无法进一步降低模型并行化程度时，才扩大Micro Batch大小。
对于参数超过30B且序列长度超过2k的模型，使用Sequence Parallelism。

然后在Paper的第4节还有几个有趣的发现，比如应该优先流水并行而不是模型并行，以及Flash Attention自带Activation Checkpointing，只对FFN进行Activation Checkpointing能获得更好的MFU。

这篇文章的解读力求精简，基本没有翻译任何原文，所以节奏非常快只写了核心要点，感兴趣也可以阅读原文，但干货应该都在这里了。

0x1. 摘要

摘要指出了现在训练大语言模型的一些高效的训练方法。包括：1) 需要跨多个硬件加速器进行并行处理和采用计算及内存优化；2) 之前的研究未涉及最新的优化技术，如Flash Attention v1/v2或Sequece Parallel；3) 作者进行了全面的训练配置打表研究，并提出了一些关键的建议以提高训练效率；4) 发现使用 micro-batch=1 能提高训练效率，但较大的micro-batch需要Activation Checkpointing或更高程度的模型并行大小，同时也会导致更大的pipeline bubble；5) 作者得到的最高效的配置在不同模型大小上实现了先进的训练效率，特别是在训练13B LLAMA模型时，达到了70.5%的MFU。

0x2. 介绍

这里对当前训练LLM的一些流行的技术做了简介，比如Zero，Activation Checkpointing，3D Parallel，Flash Attention，Sequence Parallel等等。然后paper的贡献就是把这些技术进行组合打表，最终得到了几个结论。

0x3. 背景

背景里面对Data Parallel，Tensor Parallel，Pipline Paralle，3D Parallel，Sequence Parallel，Activation Checkpointing，Fused Kernels，Flash Attention做了简要介绍，这些都是AI Infra的常见技术大家可以在 https://github.com/BBuf/how-to-optim-algorithm-in-cuda/tree/master/large-language-model-note 这里找到每一种技术的更多介绍资料。

这里需要关注的应该只有下面这张表：