AI智能摘要
GPT
这里是萌新AI,这篇文章介绍了 Pytorch 框架下的三种分布式训练方法。Tensor 并行是将模型权重和输出神经元等分到不同 GPU,以单卡简单实现和 Pytorch 库实现为例。模型并行是将不同层拆分到不同 GPU,需注意输入数据和标签的放置位置。
URL
type
Post
status
Published
date
Jul 26, 2025
slug
pytorchDT2
summary
分布式训练
tags
深度学习
Pytorch
大模型
推荐
实用教程
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)







