Lazy loaded image
记录-Pytorch框架分布式训练(二)
字数 1377阅读时长 4 分钟
2025-7-26
2025-9-6
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)。
最后,将拆分的权重分别赋值到各自的模型上。
notion image
以上便是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。
最后,计算损失,进行梯度回传。
notion image
以上便是模型并行的主要思想。代码只实现了手动拆分模型的代码。在Pytorch文档中并没有找到关于模型并行库(如果伙伴们知道在哪里可以在评论区告诉博主)。
  • 手动拆分模型

📝 流水线并行

在GPU利用率上,简单的模型并行方法致使GPU存在大量的空转时间(即等待时间)。在前向传播中,每一层必须等待前一层的结果。反向传播亦是如此。如下图:
notion image
计算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。如下图:
notion image
计算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框架分布式训练(三)记录。请大家敬请期待!

📎 参考文章

[pytorch distributed] 01 nn.DataParallel 数据并行初步_哔哩哔哩_bilibili
本期code:https://github.com/chunhuizhang/pytorch_distribute_tutorials/blob/main/tutorials/01_multi_gpus_data_parallelism.ipynb, 视频播放量 12942、弹幕量 15、点赞数 316、投硬币枚数 211、收藏人数 660、转发人数 53, 视频作者 五道口纳什, 作者简介 数学,计算机科学,现代人工智能。bridge the gap。,相关视频:[pytorch distributed] torch 分布式基础(process group),点对点通信,集合通信,[pytorch distributed] 从 DDP、模型并行、流水线并行到 FSDP(NCCL,deepspeed 与 Accelerate),[pytorch distributed] accelerate 基本用法(config,launch)数据并行,[pytorch distributed] 05 张量并行(tensor parallel),分块矩阵的角度,作用在 FFN 以及 Attention 上,[pytorch distributed] 04 模型并行(model parallel)on ResNet50,[pytorch distributed] 02 DDP 基本概念(Ring AllReduce,node,world,rank,参数服务器),[Pytorch 分布式] ring-allreduce 算法(scatter-reduce、allgather)以及 FSDP,[RL4LLM] PPO workflow 及 OpenRLHF、veRL 初步介绍,ray distributed debugger,什么是分布式?音视频领域的分布式系统解决了什么问题?2分钟带你了解!【ONITER分布式系统】,[pytorch distributed] amp 原理,automatic mixed precision 自动混合精度
[pytorch distributed] 01 nn.DataParallel 数据并行初步_哔哩哔哩_bilibili
💡
以上便是这三种并行方法学习记录,欢迎您在底部评论区留言,一起交流~
上一篇
记录-Pytorch框架分布式训练(一)
下一篇
记录-Pytorch框架分布式训练(三)

评论
Loading...