ℹ️ Skipped - page is already crawled
| Filter | Status | Condition | Details |
|---|---|---|---|
| HTTP status | PASS | download_http_code = 200 | HTTP 200 |
| Age cutoff | PASS | download_stamp > now() - 6 MONTH | 1 months ago |
| History drop | PASS | isNull(history_drop_reason) | No drop reason |
| Spam/ban | PASS | fh_dont_index != 1 AND ml_spam_score = 0 | ml_spam_score=0 |
| Canonical | PASS | meta_canonical IS NULL OR = '' OR = src_unparsed | Not set |
| Property | Value |
|---|---|
| URL | https://www.high-flyer.cn/blog/pytorch/ |
| Last Crawled | 2026-03-10 00:03:01 (29 days ago) |
| First Indexed | 2022-03-23 00:09:27 (4 years ago) |
| HTTP Status Code | 200 |
| Meta Title | PyTorch 分布式训练方法 |
| Meta Description | 分布式训练技术逐渐成为 AI 从业者必备技能之一，这是从“小模型”走向“大模型”的必由之路。 |
| Meta Canonical | null |
| Boilerpipe Text | Vachel June 18, 2021
2018年，将近3亿参数的Bert模型横空出世，将NLP领域推向了新的高度。近年来人工智能领域的发展愈来愈趋向于对大模型的研究，各大AI巨头都纷纷发布了其拥有数千亿参数的大模型，诞生出了很多新的AI应用场景。
另一方面，多种因素继续推动大模型的长足发展：1) 社会正经历着深度的数字化转型，大量的数据逐渐融合，催生出许多人工智能的应用场景和需求；2) 硬件技术不断进步：英伟达 A100 GPU，Google的 TPU，阿里的含光800等，推动着AI算力的不断提升。可以这么说，“大数据+大模型+大算力”的组合，是迈入 AI 2.0 的基石。
然而，对于大多数AI从业者们来说，受限于一些因素，平常很难接触到大数据或者大算力，很多高校研究生们往往都
习惯于用一张显卡（个人PC）来加速训练
，这不仅训练缓慢，甚至会因此限制住了想象，禁锢了创造力。
幻方AI根据自身业务需求，构建了“萤火二号”，其旨在提供澎湃如海的深度学习算力服务。团队自研了智能分时调度系统、高效存储系统与网络通信系统，可以将集群当成一台普通计算机来使用，根据任务需求弹性扩展GPU算力；自研了hfai数据仓库、模型仓库，优化AI框架与算子，集成落地了许多前沿的应用场景；通过Client接口或者Jupyter就可以轻松接入，加速训练起AI大模型。
本期文章分享的，是
如何使用起多张显卡，来加速你的AI模型
。分布式训练技术逐渐成为AI从业者必备技能之一，这是从“小模型”走向“大模型”的必由之路。我们以 PyTorch 编写的ResNet训练为例，为大家展示不同的分布式训练方法及其效果。
训练任务
：ResNet + ImageNet
训练框架
：PyTorch 1.8.0
训练平台
：幻方萤火二号
训练代码
：
https://github.com/HFAiLab/pytorch_distributed
训练准备
萤火二号当前完整支持了PyTorch的并行训练环境。笔者先使用1个节点8张卡，测试不同的分布式训练方法所需要的时间、显卡利用效率；之后采用多个节点，来验证并行加速的效果。测试的方法如下：
nn.DataParallel
torch.distributed + torch.multiprocessing
apex 半精度
为充分利用显存，这里
batch_size
设为400，记录每Epoch的耗时来进行对比。我们先采用单机单卡测试作为baseline，每Epoch耗时1786秒上下。
测试结果发现，
有半精度加持的并行加速效果最好
；
DataParallel较慢，且显卡资源利用度不高，不推荐使用
；
多机多卡带来的加速效果非常显著
。整体结果如下：
nn.DataParallel
DataParallel 是PyTorch早期提出的数据并行训练方法，其使用单进程控制，将模型和数据加载到多个GPU中，控制数据在 GPU 之间的流动，协同不同 GPU 上的模型进行并行训练。
使用起来非常方便，我们只需要用
nn.DataParallel
包装模型，再设置一些参数即可。需要定义的参数包括：
参与训练的 GPU 有哪些，device_ids=gpus；
用于汇总梯度的 GPU 是哪个，output_device=gpus[0] 。
DataParallel 会自动帮我们将数据切分 load 到相应 GPU，将模型复制到相应 GPU，进行正向传播计算梯度并汇总：
值得注意的是，
这里因为有8张显卡，
**batch_size**
要设为3200
，由DataParallel均分到不同的显卡里。模型和数据都要先 load 进 GPU中，DataParallel 的 module 才能对其进行处理，否则会报错。
发起训练。从下图我们可以看到，在显存几乎占满的情况下，8张显卡平均GPU利用率不高（36.5%）。
单独查看每一张显卡的利用情况，可以发现资源使用不平均。其中0号显卡作为梯度汇总，资源利用度相比其他显卡要高一点，整体利用率不到50%。
最终，
nn.DataParrallel
方法下，ResNet平均每Epoch
耗时984s
左右，相比单机单卡加速
大约两倍效果
。
torch.distributed
在 pytorch 1.0 之后，PyTorch官方对分布式的常用方法进行了封装，支持 all-reduce，broadcast，send 和 receive 等等。通过 MPI 实现 CPU 通信，通过 NCCL 实现 GPU 通信，解决了 DataParallel 速度慢，GPU 负载不均衡的问题。
与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，我们只需要编写一份代码，torch 就会自动将其分配给 n 个进程，分别在 n 个 GPU 上运行。
具体的，首先使用
init_process_group
设置GPU之间通信使用的后端和端口：
然后，使用
DistributedDataParallel
包装模型，它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值：
model
=
torch
.
nn
.
parallel
.
DistributedDataParallel
(
model
.
cuda
(
)
,
device_ids
=
[
local_rank
]
)
最后，因为是采用了多进程管理，需要使用
DistributedSampler
对数据集进行划分，它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练：
值得注意的是**，与**
**nn.DataParrallel**
不同，这里
**batch_size**
应该为400
，因为其只负责当前rank下对应的partition，8张卡一共组成3200个样本的batch_size。
至于 API 层面，PyTorch 为我们提供了
torch.distributed.launch
启动器，用于在命令行分布式地执行 python 文件。在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 通过参数传递给 python，我们可以这样获得当前进程的 index：
parser
=
argparse
.
ArgumentParser
(
)
parser
.
add_argument
(
'--local_rank'
,
default
=
-
1
,
type
=
int
,
help
=
'node rank for distributed training'
)
args
=
parser
.
parse_args
(
)
print
(
args
.
local_rank
)
启动8个命令行，执行如下命令：
CUDA_VISIBLE_DEVICES
=
0,1
,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node
=
8
train.py
torch.multiprocessing
手动启动命令行比较麻烦，这里有一种更简单的方法：
torch.multiprocessing
会帮助自动创建进程，绕开
torch.distributed.launch
自动控制开启和退出进程的一些小毛病。
如下面的代码所示，将原本需要
torch.distributed.launch
管理的执行内容封装进
main
函数中，通过
torch.multiprocessing.spawn
开启了 nprocs=8 个进程，每个进程执行
main
并向其中传入 local_rank（当前进程 index）。
发起训练。从下图我们可以看到，显存几乎占满，8张显卡平均利用率高（95.8%）。
同时，每张显卡利用都比较充分，利用率都在80%以上：
最终，
torch.distributed
方法下，ResNet平均每Epoch
耗时239s
左右，相比单机单卡加速
大约八倍效果
。
apex
Apex 是 NVIDIA 开源的用于混合精度训练和分布式训练库。Apex 对混合精度训练的过程进行了封装，改两三行配置就可以进行混合精度的训练，从而大幅度降低显存占用，节约运算时间。此外，Apex 也提供了对分布式训练的封装，针对 NVIDIA 的 NCCL 通信库进行了优化。PyTorch 1.6版本之后，开始原生支持，即
torch.cuda.amp
。
直接使用
amp.initialize
包装模型和优化器，apex 就会自动帮助我们管理模型参数和优化器的精度了，根据精度需求不同可以传入其他配置参数。
在分布式训练的封装上，Apex 的改动并不大，主要是优化了 NCCL 的通信。因此，大部分代码仍与
torch.distributed
保持一致。使用的时候只需要将
torch.nn.parallel.DistributedDataParallel
替换为
apex.parallel.DistributedDataParallel
用于包装模型。
在正向传播计算 loss 时，Apex 需要使用
amp.scale_loss
包装，用于根据 loss 值自动对精度进行缩放：
发起训练。从下图我们可以看到，在同等
batch_szie
条件下，显存只占用了60%，而8张显卡平均利用率高（95.8%）。
同时，每张显卡利用都比较充分，利用率都在80%以上：
最终，
apex
方法下，ResNet平均每Epoch
耗时230s
左右，相比
torch.distributed
加速了一点。同时，
apex也节约了GPU的算力资源，对应的可以设置更大的batch_size，获取更快的速度
。我们极限测试下，能在
torch.distributed
基础上再提升50%的性能左右。
多机多卡训练加速
从上面的试验可以看到，8张卡带来的加速效果近似8倍。如果我们使用更多的机器，是否会有更快的提速效果呢？这里，我们对
apex
的代码，配置4个节点，一共32张A100 GPU来进行试验，获得了惊人的效果。
可以看到，
每Epoch总耗时52秒左右
，相比单机单卡速度
提升30多倍
，几乎和显卡数量保持一致。
分布式评价
以上展示了分布式训练的过程。然而，如何对训练的结果进行快速的推理与测试，也是非常重要的问题，例如：
训练样本被切分成了若干个部分，被若干个进程分别控制运行在若干个 GPU 上，如何在进程间进行通信汇总这些（GPU 上的）信息？
使用一张卡进行推理、测试太慢了，如何使用 Distributed 进行分布式地推理和测试，并将结果汇总在一起？
…
要解决这些问题，我们需要一个更为基础的 API，
汇总记录不同 GPU 上生成的准确率、损失函数等指标信息
。这个 API 就是 torch.distributed.reduce。
reduce等 API 是
torch.distributed
中更为基础和底层的 API。这些 API 可以帮助我们控制进程之间的交互，控制 GPU 数据的传输。在自定义 GPU 协作逻辑，汇总 GPU 间少量的统计信息时，大有用处。熟练掌握这些 API 也可以帮助我们自己设计、优化分布式训练、测试流程。
如上图所示，它的工作过程包含以下两步：
在调用 reduce(tensor, op=…)后，当前进程接受其他进程发来的 tensor。
在全部接收完成后，当前进程（例如rank 0）会对当前进程的和接收到的 tensor 进行 op 操作。
通过上述步骤，我们就能够对不同 GPU 上的训练数据的损失函数进行求和了：
除了
reduce
，PyTorch官网提供了诸如
scatter
,
gather
,
all-reduce
等6种集合通信方案，具体的参考官方文档：
https://ptorch.com/docs/1/distributed
总结
本文介绍了不同的PyTorch并行训练方法，并在幻方萤火二号上进行了测试。从测试结果中我们能够得出，多机多卡并行训练能有效提升我们效率，apex 方法结合torch.distributed 能最大限度的发挥显卡的算力，让加速效果与显卡数量一致。同时，我们也发现，随着显卡数量的增多，受限于梯度汇合中的reduce计算，单个显卡的利用率会逐渐降低。如何高效优化reduce算法，尽可能提高gpu算力，从而加速整体训练效率，是接下来值得大家继续研究的课题。 |
| Markdown | [返回博客首页](https://www.high-flyer.cn/blog)
- [训练准备](https://www.high-flyer.cn/blog/pytorch/#%E8%AE%AD%E7%BB%83%E5%87%86%E5%A4%87)
- [nn.DataParallel](https://www.high-flyer.cn/blog/pytorch/#nndataparallel)
- [torch.distributed](https://www.high-flyer.cn/blog/pytorch/#torchdistributed)
- [apex](https://www.high-flyer.cn/blog/pytorch/#apex)
- [多机多卡训练加速](https://www.high-flyer.cn/blog/pytorch/#%E5%A4%9A%E6%9C%BA%E5%A4%9A%E5%8D%A1%E8%AE%AD%E7%BB%83%E5%8A%A0%E9%80%9F)
- [分布式评价](https://www.high-flyer.cn/blog/pytorch/#%E5%88%86%E5%B8%83%E5%BC%8F%E8%AF%84%E4%BB%B7)
- [总结](https://www.high-flyer.cn/blog/pytorch/#%E6%80%BB%E7%BB%93)
# PyTorch 分布式训练方法
Vachel June 18, 2021
[](https://hfai-static.high-flyer.cn/static/9ac28a620ade131f0826acbd7d76181e/99f37/p0.png)
2018年，将近3亿参数的Bert模型横空出世，将NLP领域推向了新的高度。近年来人工智能领域的发展愈来愈趋向于对大模型的研究，各大AI巨头都纷纷发布了其拥有数千亿参数的大模型，诞生出了很多新的AI应用场景。 另一方面，多种因素继续推动大模型的长足发展：1) 社会正经历着深度的数字化转型，大量的数据逐渐融合，催生出许多人工智能的应用场景和需求；2) 硬件技术不断进步：英伟达 A100 GPU，Google的 TPU，阿里的含光800等，推动着AI算力的不断提升。可以这么说，“大数据+大模型+大算力”的组合，是迈入 AI 2.0 的基石。 然而，对于大多数AI从业者们来说，受限于一些因素，平常很难接触到大数据或者大算力，很多高校研究生们往往都**习惯于用一张显卡（个人PC）来加速训练**，这不仅训练缓慢，甚至会因此限制住了想象，禁锢了创造力。 幻方AI根据自身业务需求，构建了“萤火二号”，其旨在提供澎湃如海的深度学习算力服务。团队自研了智能分时调度系统、高效存储系统与网络通信系统，可以将集群当成一台普通计算机来使用，根据任务需求弹性扩展GPU算力；自研了hfai数据仓库、模型仓库，优化AI框架与算子，集成落地了许多前沿的应用场景；通过Client接口或者Jupyter就可以轻松接入，加速训练起AI大模型。
本期文章分享的，是**如何使用起多张显卡，来加速你的AI模型**。分布式训练技术逐渐成为AI从业者必备技能之一，这是从“小模型”走向“大模型”的必由之路。我们以 PyTorch 编写的ResNet训练为例，为大家展示不同的分布式训练方法及其效果。
**训练任务**：ResNet + ImageNet
**训练框架**：PyTorch 1.8.0
**训练平台**：幻方萤火二号
**训练代码**：<https://github.com/HFAiLab/pytorch_distributed>
### 训练准备
萤火二号当前完整支持了PyTorch的并行训练环境。笔者先使用1个节点8张卡，测试不同的分布式训练方法所需要的时间、显卡利用效率；之后采用多个节点，来验证并行加速的效果。测试的方法如下：
1. nn.DataParallel
2. torch.distributed + torch.multiprocessing
3. apex 半精度
为充分利用显存，这里`batch_size`设为400，记录每Epoch的耗时来进行对比。我们先采用单机单卡测试作为baseline，每Epoch耗时1786秒上下。
测试结果发现，**有半精度加持的并行加速效果最好**；**DataParallel较慢，且显卡资源利用度不高，不推荐使用**；**多机多卡带来的加速效果非常显著**。整体结果如下：
[](https://hfai-static.high-flyer.cn/static/2b8b7b8b386810d429647d93028deece/d9ed0/p1.png)
### nn.DataParallel
DataParallel 是PyTorch早期提出的数据并行训练方法，其使用单进程控制，将模型和数据加载到多个GPU中，控制数据在 GPU 之间的流动，协同不同 GPU 上的模型进行并行训练。
使用起来非常方便，我们只需要用 `nn.DataParallel` 包装模型，再设置一些参数即可。需要定义的参数包括：
- 参与训练的 GPU 有哪些，device\_ids=gpus；
- 用于汇总梯度的 GPU 是哪个，output\_device=gpus\[0\] 。
DataParallel 会自动帮我们将数据切分 load 到相应 GPU，将模型复制到相应 GPU，进行正向传播计算梯度并汇总：
[](https://hfai-static.high-flyer.cn/static/93576e68226a3bcf0b52a87376ab932d/f058b/p2.png)
值得注意的是，**这里因为有8张显卡，**`**batch_size**`**要设为3200**，由DataParallel均分到不同的显卡里。模型和数据都要先 load 进 GPU中，DataParallel 的 module 才能对其进行处理，否则会报错。
[](https://hfai-static.high-flyer.cn/static/0e4685b5f245ff537f3b322b8724df2e/63ec5/p3.png)
发起训练。从下图我们可以看到，在显存几乎占满的情况下，8张显卡平均GPU利用率不高（36.5%）。
[](https://hfai-static.high-flyer.cn/static/606dcf947bf33e727f91a0c771c451a2/67fe0/p4.png) 单独查看每一张显卡的利用情况，可以发现资源使用不平均。其中0号显卡作为梯度汇总，资源利用度相比其他显卡要高一点，整体利用率不到50%。
[](https://hfai-static.high-flyer.cn/static/7f6c45154108bd9094fff7ba70ec8c6f/7798c/p5.png)
最终，`nn.DataParrallel`方法下，ResNet平均每Epoch**耗时984s**左右，相比单机单卡加速**大约两倍效果**。
### torch.distributed
在 pytorch 1.0 之后，PyTorch官方对分布式的常用方法进行了封装，支持 all-reduce，broadcast，send 和 receive 等等。通过 MPI 实现 CPU 通信，通过 NCCL 实现 GPU 通信，解决了 DataParallel 速度慢，GPU 负载不均衡的问题。
与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，我们只需要编写一份代码，torch 就会自动将其分配给 n 个进程，分别在 n 个 GPU 上运行。
具体的，首先使用 `init_process_group` 设置GPU之间通信使用的后端和端口：
[](https://hfai-static.high-flyer.cn/static/38fe459308837d43f8e467cc7c32b7fe/024d6/p6.png)
然后，使用 `DistributedDataParallel` 包装模型，它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值：
```
model = torch.nn.parallel.DistributedDataParallel(model.cuda(), device_ids=[local_rank])
```
最后，因为是采用了多进程管理，需要使用 `DistributedSampler` 对数据集进行划分，它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练： [](https://hfai-static.high-flyer.cn/static/f31c5d3b7c1a02c10fb68e7500c64dfa/1f09d/p7.png)
值得注意的是\*\*，与\*\*`**nn.DataParrallel**`**不同，这里**`**batch_size**`**应该为400**，因为其只负责当前rank下对应的partition，8张卡一共组成3200个样本的batch\_size。
至于 API 层面，PyTorch 为我们提供了 `torch.distributed.launch` 启动器，用于在命令行分布式地执行 python 文件。在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 通过参数传递给 python，我们可以这样获得当前进程的 index：
```
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int, help='node rank for distributed training')
args = parser.parse_args()
print(args.local_rank)
```
启动8个命令行，执行如下命令：
```
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8 train.py
```
**torch.multiprocessing**
手动启动命令行比较麻烦，这里有一种更简单的方法：`torch.multiprocessing` 会帮助自动创建进程，绕开 `torch.distributed.launch` 自动控制开启和退出进程的一些小毛病。
如下面的代码所示，将原本需要 `torch.distributed.launch` 管理的执行内容封装进 `main` 函数中，通过 `torch.multiprocessing.spawn` 开启了 nprocs=8 个进程，每个进程执行 `main` 并向其中传入 local\_rank（当前进程 index）。
[](https://hfai-static.high-flyer.cn/static/886058f74182b0d6828dd04c7ddf4f04/55811/p8.png)
发起训练。从下图我们可以看到，显存几乎占满，8张显卡平均利用率高（95.8%）。
[](https://hfai-static.high-flyer.cn/static/034771139e2213367410e5b97647fb52/84ee5/p9.png)
同时，每张显卡利用都比较充分，利用率都在80%以上：
[](https://hfai-static.high-flyer.cn/static/882794889232274d41529dd36abe3f95/d073d/p10.png)
最终，`torch.distributed`方法下，ResNet平均每Epoch**耗时239s**左右，相比单机单卡加速**大约八倍效果**。
### apex
Apex 是 NVIDIA 开源的用于混合精度训练和分布式训练库。Apex 对混合精度训练的过程进行了封装，改两三行配置就可以进行混合精度的训练，从而大幅度降低显存占用，节约运算时间。此外，Apex 也提供了对分布式训练的封装，针对 NVIDIA 的 NCCL 通信库进行了优化。PyTorch 1.6版本之后，开始原生支持，即`torch.cuda.amp`。
直接使用 `amp.initialize` 包装模型和优化器，apex 就会自动帮助我们管理模型参数和优化器的精度了，根据精度需求不同可以传入其他配置参数。
[](https://hfai-static.high-flyer.cn/static/b7a7aa514937fe1a8035b4cf71f9ed59/d44c9/p11.png)
在分布式训练的封装上，Apex 的改动并不大，主要是优化了 NCCL 的通信。因此，大部分代码仍与 `torch.distributed` 保持一致。使用的时候只需要将 `torch.nn.parallel.DistributedDataParallel` 替换为 `apex.parallel.DistributedDataParallel` 用于包装模型。
在正向传播计算 loss 时，Apex 需要使用 `amp.scale_loss` 包装，用于根据 loss 值自动对精度进行缩放：
[](https://hfai-static.high-flyer.cn/static/e13ffa82cd72bc92c84b7cc54408adc5/ff42b/p12.png)
发起训练。从下图我们可以看到，在同等`batch_szie`条件下，显存只占用了60%，而8张显卡平均利用率高（95.8%）。
[](https://hfai-static.high-flyer.cn/static/828d9008cfdd642b61515577bd904dbf/87488/p13.png)
同时，每张显卡利用都比较充分，利用率都在80%以上：
[](https://hfai-static.high-flyer.cn/static/eb8903f9aa4abecd5b1d2c54b3253fbc/21b4d/p14.png)
最终，`apex`方法下，ResNet平均每Epoch**耗时230s**左右，相比`torch.distributed`加速了一点。同时，**apex也节约了GPU的算力资源，对应的可以设置更大的batch\_size，获取更快的速度**。我们极限测试下，能在`torch.distributed`基础上再提升50%的性能左右。
### 多机多卡训练加速
从上面的试验可以看到，8张卡带来的加速效果近似8倍。如果我们使用更多的机器，是否会有更快的提速效果呢？这里，我们对 `apex`的代码，配置4个节点，一共32张A100 GPU来进行试验，获得了惊人的效果。 [](https://hfai-static.high-flyer.cn/static/ea59d00d6bf86822a0d2abceef8462ed/218a4/p15.png) 可以看到，**每Epoch总耗时52秒左右**，相比单机单卡速度**提升30多倍**，几乎和显卡数量保持一致。
### 分布式评价
以上展示了分布式训练的过程。然而，如何对训练的结果进行快速的推理与测试，也是非常重要的问题，例如：
1. 训练样本被切分成了若干个部分，被若干个进程分别控制运行在若干个 GPU 上，如何在进程间进行通信汇总这些（GPU 上的）信息？
2. 使用一张卡进行推理、测试太慢了，如何使用 Distributed 进行分布式地推理和测试，并将结果汇总在一起？
3. …
要解决这些问题，我们需要一个更为基础的 API，**汇总记录不同 GPU 上生成的准确率、损失函数等指标信息**。这个 API 就是 torch.distributed.reduce。
reduce等 API 是 `torch.distributed` 中更为基础和底层的 API。这些 API 可以帮助我们控制进程之间的交互，控制 GPU 数据的传输。在自定义 GPU 协作逻辑，汇总 GPU 间少量的统计信息时，大有用处。熟练掌握这些 API 也可以帮助我们自己设计、优化分布式训练、测试流程。 [](https://hfai-static.high-flyer.cn/static/296becc7cfe080c8de524b9cf43e4ba7/854dc/p16.png)
如上图所示，它的工作过程包含以下两步：
1. 在调用 reduce(tensor, op=…)后，当前进程接受其他进程发来的 tensor。
2. 在全部接收完成后，当前进程（例如rank 0）会对当前进程的和接收到的 tensor 进行 op 操作。
通过上述步骤，我们就能够对不同 GPU 上的训练数据的损失函数进行求和了：
[](https://hfai-static.high-flyer.cn/static/755f130d5fa29aeee931ae3051fe5921/97655/p17.png)
除了`reduce`，PyTorch官网提供了诸如 `scatter`, `gather`, `all-reduce`等6种集合通信方案，具体的参考官方文档： <https://ptorch.com/docs/1/distributed>
### 总结
本文介绍了不同的PyTorch并行训练方法，并在幻方萤火二号上进行了测试。从测试结果中我们能够得出，多机多卡并行训练能有效提升我们效率，apex 方法结合torch.distributed 能最大限度的发挥显卡的算力，让加速效果与显卡数量一致。同时，我们也发现，随着显卡数量的增多，受限于梯度汇合中的reduce计算，单个显卡的利用率会逐渐降低。如何高效优化reduce算法，尽可能提高gpu算力，从而加速整体训练效率，是接下来值得大家继续研究的课题。
***
本文作者： **Vachel**
[返回博客首页](https://www.high-flyer.cn/blog)
- [← 萤火跑模型 \| Informer 上手实践](https://www.high-flyer.cn/blog/informer/)
- [萤火跑模型 \| Autoformer 长时序预测 →](https://www.high-flyer.cn/blog/auto-former/)
***
您可以转载、不违背作品原意地摘录及引用本技术博客的内容，但必须遵守以下条款： 署名 — 您应当署名原作者，但不得以任何方式暗示幻方为您背书，亦不会对幻方的权利造成任何负面影响。 非商业性使用 — 您不得将本技术博客内容用于商业目的。 禁止演绎 — 如果基于该内容改编、转换、或者再创作，您不得公开或分发被修改内容，该内容仅可供个人使用。
联系幻方
地址
杭州市拱墅区环城北路 169 号 汇金国际大厦 A 座 14 层 上海市浦东新区花园石桥路 66 号 东亚银行金融大厦 45 层
电话
\+86-0571-86656960
邮箱
info@high-flyer.cn
[浙ICP备15029780号-1](https://beian.miit.gov.cn/)
[浙公网安备 33010302003447](http://www.beian.gov.cn/portal/registerSystemInfo?recordcode=33010302003447)
©幻方量化版权所有 |
| Readable Markdown | null |
| Shard | 132 (laksa) |
| Root Hash | 5581124212395307732 |
| Unparsed URL | cn,high-flyer!www,/blog/pytorch/ s443 |