Lazy loaded image
记录-Pytorch框架分布式训练(二)
字数 1377阅读时长 4 分钟
2025-7-26
2025-9-6
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)。
最后,将拆分的权重分别赋值到各自的模型上。
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...