分布式训练之数据并行

分布式训练之数据并行

1. 背景

​ 在深度学习的发展历程中，通过改变神经网络的结构，可以取得比较明显的精度。但是，随着神经网络结构设计技术不断成熟，想通过优化神经网络结构来打破模型的精度瓶颈遇到了很大的挑战。

​ 根据一些研究表明，通过增大数据规模和模型规模，可以进一步提升模型精度。但是，这也意味着训练时间会变长，所以可以增加计算资源，通过分布式训练来缩短训练时间，将单卡的负载拆到多卡上。

​ 数据并行（Data Parallelism） 通过修改Sampler切分输入，每张卡只需要处理一部分数据；

​ 模型并行（Model Parallelism） 通过修改层内的计算方式，将单层的计算负载和显存负载切分到多张卡上；

​ 流水并行（Pipeline Parallelism） 通过将不同的层放到不同的卡上，进而将计算负载和显存负载切分到多张卡上。

2. 数据并行

​ 在深度学习中，数据集越来越大，我们很难一次性将所有的数据集加载到内存中，所以我们会使用随机梯度下降法，即每次只用一个batch的数据进行训练，对损失函数求导，得到每个参数的梯度。例如，在一个有51200个样本的数据集中，每次只取512个样本进行梯度的计算（batch_size=512），所以要将数据集迭代一遍需要100次，训练时间非常长。

​ 为了加快模型的训练，可以使用数据并行的分布式训练方法。将一个大批次分割为很多小批次，使用很多节点进行计算，在每个节点上计算一个小批次，对若干个节点的梯度进行汇总后再加权平均，最终求和就得到了最终的大批次的梯度结果。

$$\frac{\partial L}{\partial w}=\frac{\partial [\frac{1}{n}{\sum }_{i=1}^{n}f(x^{(i)},y^{(i)})]}{\partial w}=\frac{m_1}{n}\frac{\partial l_1}{\partial w}+\frac{m_2}{n}\frac{\partial l_2}{\partial w}+\cdots +\frac{m_k}{n}\frac{\partial l_k}{\partial w}$$
​ 在上面这个公式中，$w$是模型的参数，$\frac{\partial L}{\partial w}$是在batch_size为n的情况下计算得到的真实梯度，共有$k$个节点，在第$j$个节点上有$m_j$个数据，$\frac{\partial l_j}{\partial w}$是在该节点上计算得到的精度，$m_1+m_2+\cdots +m_k=n$（$n$个样本数据被拆分到了$k$个姐弟节点上），$x^{(i)}$和$y^{(i)}$是样本数据$i$的特征和标签，对于样本数据$i$，$f(x^{(i)},y^{(i)})$是前向传播的损失函数。

​ 公式表达的含义是在不同的节点上分别对一部分数据进行梯度的计算，将各个GPU的梯度进行汇总后，计算其加权平均值。

​ 如果每个节点上的数据量是平分的，则
$$\frac{\partial L}{\partial w}=\frac{1}{k}[\frac{\partial l_1}{\partial w}+\frac{\partial l_2}{\partial w}+\cdots +\frac{\partial l_k}{\partial w}]$$
​ 对于每个节点，我们使用相同的模型参数进行前向传播；我们将大批次的数据分割成很多个小批次的数据，发送到不同的节点上，每个节点都正常计算其梯度，计算完返回其梯度计算结果到主节点。这一步是异步的，因为每个节点的速度都可能稍有差别，一旦得到了所有的梯度，就可以计算这些梯度的加权平均，并使用这个加权平均的梯度去更新整个模型的参数。

3. PyTorch数据并行实现

​ DistributedDataParallel(DDP)在模型层面实现了数据并行。使用DDP的应用程序应该派生多个进程，并为每个进程创建一个DDP实例。它使用torch.distributed程序包中的通信集合来同步梯度和缓冲区。更具体地说，DDP给model.parameters()的每个参数都注册了一个自动的hook，在反向传播中相应的梯度被计算出来时，hook将会fire。然后，DDP使用一个信号来引发跨进程间的梯度同步。

​ 推荐使用DDP的方法是给每个复制的模型派生一个进程，一个复制的模型可以跨多台设备。DDP进程可以放在相同的机器，或者跨多台机器，但是GPU设备不能跨进程分享。

3.1 DataParallel和DistributedDataParallel之间的比较

​ 在深入探讨前，先澄清一下，为什么增加了复杂度，但还是要使用DistributedDataParallel而不是DataParallel。

• 首先，DataParallel是单进程、多线程的，只能工作在一台机器上。但是，DistributedDataParallel是多进程，可以在单台或多台机器上使用。DataParallel在单台机器上甚至比DistributedDataParallel要慢。
• 如果模型太大而无法放到一张GPU上，就必须使用模型并行来将其分到多张GPU上。DistributedDataParallel支持模型并行，但是DataParallel暂时还不支持。当DDP和模型并行相结合时，每个DDP进行将使用模型并行，并且所有的进程将一起使用数据并行。

3.2 基本用例

​ 为了创建DDP模块，首先要合理地设置进程组，详见Writing Distributed Applications with PyTorch。

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP


def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    
    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

​ 现在，可以创建一个toy module，用DDP进行包装，并喂给它一些假的输入数据。当DDP把model states从rank 0进程传播到DDP构造器的所有其他进程中时，不需要不同的DDP进程的模型初始化参数值不同。

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)

​ DDP提供了底层的分布式通信细节，并提供了一个清晰的API，就好像它是本地模型一样。当反向传播时，梯度同步就会发生。当backward()返回时，param.grad已经包含了同步精度的张量。

4. 参考资料

大规模训练系列之技术挑战

分布式深度学习简介：数据并行和模型并行