大模型分布式训练并行技术(四)张量并行

引言

在Transformer架构里,有两块主要的计算量较大的部分,一是Self-Attention,二是MLP。在前面的文章中,我们已经介绍了模型并行和数据并行,本文将介绍张量并行,这是一种更加细粒度的并行方式,可以进一步提高模型的训练效率。

张量并行使用了矩阵乘法可以并行计算的特性,将模型的参数划分为多个部分,每个部分在不同的设备上进行计算,最后将结果进行汇总。下面,我们分别看FFN和Self-Attention的张量并行实现。

MLP

MLP的主要构建块都是完全连接的 nn.Linear,后跟非线性激活 GeLU。

按照 Megatron[2] 的论文符号,我们可以将其点积部分写为 Y = GeLU(XA),其中 X 和 Y 是输入和输出向量,A 是权重矩阵。

如果我们以矩阵形式查看计算,很容易看出矩阵乘法如何在多个 GPU 之间拆分:

如果我们将权重矩阵 A 按列拆分到 N 个 GPU 上,并行执行矩阵乘法 XA_1 到 XA_n,那么我们将得到 N 个输出向量 Y_1、Y_2、...、Y_n,这些向量可以独立输入到 GeLU 中:

利用这一原理,我们可以更新任意深度的 MLP,而无需 GPU 之间进行任何同步,直到最后,我们才需要重建输出向量。

Megatron-LM 论文作者为此提供了一个有用的例子:

Self-Attention

Self-Attention 的张量并行更简单,因为self-attention天然的是多头注意力机制,可以将每个头的计算分配到不同的 GPU 上。

在上图中,我们可以用2个GPU并行的计算self-attention,其中每个GPU计算一个头的注意力机制。那原则上,有几个头就可以用几个GPU并行计算。

特别注意事项:TP 需要非常快的网络,因此不建议跨多个节点进行 TP。实际上,如果一个节点有 4 个 GPU,则最高 TP 度为 4。如果你需要 8 的 TP 度,则需要使用至少有 8 个 GPU 的节点。

下一篇我们看看混合并行。

参考

1\] [Model Parallelism](https://link.juejin.cn?target=https%3A%2F%2Fhuggingface.co%2Fdocs%2Ftransformers%2Fv4.15.0%2Fen%2Fparallelism "https://huggingface.co/docs/transformers/v4.15.0/en/parallelism") \[2\] [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://link.juejin.cn?target=https%3A%2F%2Farxiv.org%2Fabs%2F2104.04473 "https://arxiv.org/abs/2104.04473") ## 欢迎关注我的GitHub和微信公众号,来不及解释了,快上船! [GitHub: LLMForEverybody](https://link.juejin.cn?target=https%3A%2F%2Fgithub.com%2Fluhengshiwo%2FLLMForEverybody "https://github.com/luhengshiwo/LLMForEverybody") 仓库上有原始的Markdown文件,完全开源,欢迎大家Star和Fork!

相关推荐
倔强青铜三6 小时前
苦练Python第46天:文件写入与上下文管理器
人工智能·python·面试
AI大模型7 小时前
GitHub 狂飙 72k Star,这本大模型书凭啥能圈粉无数?
程序员·llm·agent
canonical_entropy8 小时前
AI时代,我们还需要低代码吗?—— 一场关于模型、演化与软件未来的深度问答
后端·低代码·aigc
堆栈future9 小时前
秒级生成4K图!字节豆包Seedream 4.0实测:完爆GPT-4o和Nano Banana
llm·aigc
天天扭码10 小时前
来全面地review一下Flex布局(面试可用)
前端·css·面试
Mor_11 小时前
UE5 网络通信协议学习笔记
面试
林木森ai11 小时前
爆款AI动物运动会视频,用Coze(扣子)一键搞定全流程(附保姆级拆解)
人工智能·aigc
沐怡旸11 小时前
【底层机制】std::unique_ptr 解决的痛点?是什么?如何实现?怎么正确使用?
c++·面试
前端缘梦11 小时前
Vue Keep-Alive 组件详解:优化性能与保留组件状态的终极指南
前端·vue.js·面试
大模型教程11 小时前
小白学大模型:从零搭建LLaMA
程序员·llm·llama