PyTorch 的多卡并行训练

6,266 阅读3分钟

DataParallel

并行的方式分为了数据并行。 DataParallel 会将module复制到多个卡上,也会将每个batch均分到每张卡上,每张卡独立forward自己那份data,而在backward时,每个卡上的梯度会汇总到原始的module上,以此来实现并行。

但是,这样的方式会造成原始module在的那张卡的显存压力比其他卡要大,也就是这种方式存在负载不均衡的情况。具体情况可以看pytorch: 一机多卡训练的尝试这篇文章的实验。

Multiprocessing

另外一种方式是利用Python的多进程,每张卡运行一个进程,每个进程有一个自己的model和一份数据,求梯度时则将多张卡的梯度汇总,然后传播到每张卡上来实现并行。

这种方式就避免了第一种方式负载不均衡的情况,而且多进程也避免了Python的GIL的机制。

但是PyTorch官方文档还是推荐使用 DataParallel 的方式,其说法如下:

Use nn.DataParallel instead of multiprocessing

Most use cases involving batched inputs and multiple GPUs should default to using DataParallel to utilize more than one GPU. Even with the GIL, a single Python process can saturate multiple GPUs. As of version 0.1.9, large numbers of GPUs (8+) might not be fully utilized. However, this is a known issue that is under active development. As always, test your use case. There are significant caveats to using CUDA models with multiprocessing; unless care is taken to meet the data handling requirements exactly, it is likely that your program will have incorrect or undefined behavior.

Example

一个利用多进程实现多卡训练的样例如下。

...
import torch.distributed as dist
import torch.multiprocessing as mp
...

def run(gpu_id):
    ...
    model = Model()
    model.cuda(gpu_id)
    for epoch in range(epochs):
        train(epoch, gpu_id, model, optimizer, data_loader)
        
        if gpu_id == 0:
            validata(...)
        
# 汇总多卡的梯度,平均后传播
def average_gradients(model):
    size = float(dist.get_world_size())
    for p in model.parameters():
        if p.grad is not None:
            dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
            p.grad.data /= size
            
def train(epoch, gpu_id, model, optimizer, data_loader):
    model.train()
    
    for i, data in enumerate(data_loader):
        data = {key: value.to(gpu_id) for key, value in data.items()}
        model.zero_grad()
        
        loss = model.forward(data)
        
        loss.backward()
        
        if word_size > 1:
            average_gradients(model)
        optimizer.step()
        
def init_process(host, port, rank, fn, backend="nccl"):
    """
    host: str, 如果一机多卡可以直接填localhost
    port: str
    rank: int
    fn: train的函数
    backend: 实现多卡通信的机制,有多种,PyTorch有汇总
    """
    os.environ["MASTER_ADDR"] = host
    os.environ["MASTER_ADDR"] = port
    dist.init_process_group(backend, rank=rank, world_size=world_size)
    fn(rank)
    
if __name__ == "__main__":
    mp.set_start_method("spawn")
    
    processes = []
    # 有几张卡可以指定ranks列表
    for rank in ranks:
        p = mp.Process(target=init_process, args=(host, port, rank, run)
        p.start()
        processes.append(p)
    
    for p in processes:
        p.join()

或者直接参考官方文档 WRITING DISTRIBUTED APPLICATIONS WITH PYTORCH

Apex

NVIDIA开发的支持并行和混合精度的辅助函数。 Github地址为:apex

主要是对PyTorch多进程方式的多卡训练代码的封装,重要的是支持fp16的训练,以及混合精度的训练。

但是我还没用过apex,这里先不写了。

Accumulating gradients

还有一种方式是累加梯度Accumulating gradients(参考自Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups,作者是huggingface/pytorch-pretrained-BERT的作者。)

方法很简单,就是假设batch_size=32,但是1个GPU的显存放不下,我们可以在forward的时候将整个batch分为8份,每次过size=4的小batch,计算小batch的梯度后进行backward,然后先不要更新参数,而是继续forward第二个小batch,计算loss,然后backward,而且将第二个小batch的梯度累加到第一个小的batch的梯度上,如此,当4个小batch都计算出梯度,再更新参数。

这样的方式在PyTorch里很容易实现,因为只有我们调用 model.zero_grad() 或者optimizer.zero_grad()时,缓存的梯度才会被重置。

Thomas Wolf给出了一个gist的例子,可以供我们参考, Accumulating gradients。或者我们可以直接参考huggingface/pytorch-pretrained-BERT里的写法。

Accumulating gradients 可以和分布式训练结合使用,来达到增大batch_size的作用,因为对于BERT等大模型来说,再加一些下游框架,batch_size只能设到个位数,而paper中建议的batch_size都在16或者32,太小的batch_size也会影响到模型的性能。因此我们需要更大的batch_size来稳定训练过程。


欢迎访问个人博客Alex Chiu的学习空间