Pytorch分布式训练原理简介

1. 引言

分布式训练就是指将模型放置在很多台机器并且在每台机器上的多个GPU上进行训练，之所以使用分布式训练的原因一般来说有两种：其一是模型在一块GPU上放不下，其二使用多块GPU进行并行计算能够加速训练。但是需要注意的是随着使用的GPU数量增加，各个设备之间的通信会变得复杂，导致训练速度下降。

一般来说，分布式训练主要分为两种类型：数据并行化 (Data Parallel) 以及模型平行化（Model Parallel）。

2. 数据并行化

当训练的数据量非常大时，假设模型的结构能够放置在单个GPU上时，就可以采用数据平行化的方式来进行分工合作。

常用的做法是依照一些规则将数据分配到不同的GPU上，并且每个GPU都有相同的模型构架，也就是会在每个GPU上复制一份相同的模型，各自进行训练后，将计算结果合并，再进行参数更新。如下所示：

此时的参数更新的方式又分为同步及异步两种：

• 同步方式：所有的GPU在训练时会等待其他GPU计算完毕后，才会进行一次参数更新，因此训练速度上会比使用异步的方式来得慢。但因为在更新参数时会合并其他计算结果，相当于增加了batch size的大小，对于训练结果有一定的提升。
• 异步方式：每个GPU各自进行训练和参数更新，不须等待其他GPU计算完毕。能够提升训练速度，但可能会产生Slow and Stale Gradients（梯度失效、梯度过期）问题，收敛过程较不稳定，训练结果会比使用同步的方式差。

3. 模型并行化

当训练的模型构架太大以至于在一个GPU放不下时，可以采用模型平行化的方式来将模型拆解并分配到不同的GPU上。如下所示：

由于模型层与层之间通常有依赖性，也就是指在进行前向传播、反向传播时，前面及后面的层会作为彼此的输入和输出，在训练速度上会有一定的限制，因此若非使用较大的模型较不建议采用模型平行化。

4. 数据并行化+模型并行化

当然也可以采用结合数据平行化与模型平行化的方式，如下图所示：

此时的做法是将模型拆解并分配到单个机器中不同的GPU上，并且在其他机器中的每个GPU都复制一份相同的模型构架（拆解后的部分模型），而数据数据也会依照一些规则分配到不同的GPU上。

5. Pytorch分布式训练方法

Pytorch提供了两种分布式训练方法—DataParallel（DP）、DistributedDataParallel（DDP）。

• DP采用Parameter Server（PS）构架，并且仅支持单台机器多GPU，任务中只会有一个进程
• DDP则是采用Ring-All-Reduce构架，可支持多台机器多GPU，能够使用多进程。

通常来说，DP的code写法较为简单，但由于DDP的执行速度更快，官方更推荐使用DDP。

6. 总结

本文介绍了分布式训练的三种实现方式，即数据并行化，模型并行化以及二者的混合，并给出了相应的图例来讲解其背后的原理。

您学废了吗？

关注公众号《AI算法之道》，获取更多AI算法资讯。