AI智能摘要
GPT
这里是萌新AI,这篇文章介绍了 PyTorch 框架中的三种分布式训练方法。首先,Tensor 并行是将模型权重和输出神经元等分到多个 GPU 上。其次,模型并行是将模型的不同层拆分到不同 GPU,需注意输入输出数据的对应放置。最后,流水线并行被提及但未展开。博客以理论讲解为主,并提供了相关代码链接供参考。
URL
type
status
date
slug
summary
tags
category
icon
password
此篇博客主要用于Tensor并行、模型并行和流水线并行记录。博客主要进行理论内容记录,代码仅仅提供Gitee链接。如果对代码有不懂的地方,可以在下方留言!
📝 Tensor并行
我们根据标题很容易猜测出这一小节的内容是什么?Tensor并行(张量并行)将模型权重等分为n份(几个GPU分为几份)。同时,模型各个层中的输出神经元相较于原模型等分为n份。
以single = nn.Linear(4, 8)和 second = nn.Linear(8, 4)为例(请结合下图一起阅读,矩阵计算是博主口算的,有的元素可能计算有误):
首先,input的shape=(4, 4)。根据线性层得知,权重w的shape=(8, 4)。
接着,根据y = x * W.T,即W.T的shape=(4, 8)。那么,将w进行拆分为w1和w2。w1和w2的shape=(4, 4)。
最后,将拆分的权重分别赋值到各自的模型上。

以上便是Tensor并行的思想,代码分为单卡上简单的张量并行(仅进行计算分析)和Pytorch提供的张量并行库实现。
- 单卡简单实现(仅进行计算分析)
- Pytorch库实现张量并行(需要注意的是:Pytorch版本不同,关于张量并行API使用方法也不同。博主的Pytorch版本是2.1.1)
📝 模型并行(简单)
模型并行主要是将模型不同层进行拆分,并将不同层放入不同GPU中。模型并行主要注意两个地方:输入层和输出层。假设:输入层放入0号GPU,输出层放入1号GPU。那么,输入的数据必须也放入0号GPU,则标签必须放入1号GPU,用来进行损失计算。
这里我们两个线性层为例,并以图辅助分析(假设有两块GPU):
首先,将input层放入0号GPU,output层放入1号GPU。同时,input_data数据放入0号GPU,output_data数据放入1号GPU。
接着,input层输出到的数据放入1号GPU。
最后,计算损失,进行梯度回传。

以上便是模型并行的主要思想。代码只实现了手动拆分模型的代码。在Pytorch文档中并没有找到关于模型并行库(如果伙伴们知道在哪里可以在评论区告诉博主)。
- 手动拆分模型
📝 流水线并行
在GPU利用率上,简单的模型并行方法致使GPU存在大量的空转时间(即等待时间)。在前向传播中,每一层必须等待前一层的结果。反向传播亦是如此。如下图:

计算GPU空转占比(面积分析):1/2。
为了解决该问题,Google推出了Gpipe。Gpipe是一种同步更新梯度策略,而且Pytorch实现的流水线并行API也是基于Gpipe。所以该篇博客主要讲解Gpipe方法和Pytorch的API使用样例。
Gpipe方法在模型并行的基础上,加入了数据并行。数据并行将mini_batch拆分为micro_batch。例子:我们有两个GPU,如果仅使用模型并行,mini_batch_size=16,那么当加入数据并行拆分为micro_batch_size=8。在前向传播时,cuda:1处理第一个micro_batch时,coda:0处理第二个micro_batch。在反向传播时,cuda:1处理第二个micro_batch时,coda:0处理第一个micro_batch。如下图:

计算GPU空转占比(面积分析):1/3。
以上便是流水线并行-Gpipe的主要思想,代码通过调用Pytorch提供的API实现。
- Pipline-Pytorch的API接口(Pytorch从2.4.0开始提供pipline接口。博主使用的是2.4.0版本。该版本拆分模型后并没有提供自动分配GPU,需要手动分配GPU(HugFace的Transformer实现了自动分配,大家可以尝试一下!)。如果博主这里有问题,请及时指正!)
🤗 总结归纳
以上便是张量并行、简单的模型并行和pipline并行的主要理论思想。在pipline并行中,Pytorch提供的API接口,实现了拆分模型,但需要手动分配GPU。pipline并行与简单的模型并行区别在模型拆分方式上不同。
Pytorch分布式框架还有一个FSDP没有学习记录。FSDP方法将在Pytorch框架分布式训练(三)记录。请大家敬请期待!
📎 参考文章
以上便是这三种并行方法学习记录,欢迎您在底部评论区留言,一起交流~
- 作者:不爱吃香菜的萌新
- 链接:https://hexo.levsongsw.com//deeplearn/pytorchDT2
- 声明:本文采用 CC BY-NC-SA 4.0 许可协议,转载请注明出处。





![[pytorch distributed] 01 nn.DataParallel 数据并行初步_哔哩哔哩_bilibili](https://www.notion.so/image/https%3A%2F%2Fi2.hdslb.com%2Fbfs%2Farchive%2F92ef5011da5f7ff1158fe181af1880af2104d570.jpg%40100w_100h_1c.png?table=block&id=23b02083-7c21-809b-9c33-c409710ed337&t=23b02083-7c21-809b-9c33-c409710ed337)







