batch & print pro_使用大batch训练神经网络：多GPU的分布式训练，基于PyTorch！

作者：Thomas Wolf编译：ronghuaiyang

导读

基于PyTorch的分布式多GPU训练实践，帮你搞定大Batch训练！

2018年的大部分时间，我都在训练神经网络来解决gpu的限制。无论是有150M参数的语言模型，还是30M参数的元学习神经网络，我只能在一块GPU上一次训练几个样本。

但是大多数情况下，随机梯度下降算法需要一个batch给与更多的样本，才能得到满意的结果。

当你的GPU一个batch只能容纳几个样本的时候，你怎么才能批量训练你的模型？

这里有一些工具，技巧和技巧，你可以用来做这一点，我认为这是对我过去使用和学习的东西的一个很好的总结，都在这个文章里了。

本文中，我将主要讨论PyTorch框架。其中一些工具还没有包含在PyTorch中(截至1.0)，所以我还包含了一些定制代码。

特别的，我们将讨论：

• 当batch量大于GPU内存，甚至一个训练样本都放不进去的时候，如何在单个或多GPU服务器上训练模型
• 如何最有效地使用一个多GPU的机器
• 在分布式设置中使用多台机器训练模型的最简单方法。

让我们从最简单的技巧开始：梯度累积。

在一个或多个GPU上使用大的batch

你已经构建了一个很好的模型，它可能是这个任务上的新SOTA，但是每次你尝试在一个batch中使用多个样本时，你都会得到一个CUDA RuntimeError: out of memory。

但是你非常肯定对batch的数量加倍会改善结果。

你怎么才能做到呢？

对于这个问题有一个简单的解决方案：累积梯度。这里有一个关于随机梯度下降如何工作的快速提示，来自我之前关于元学习的文章：https://medium.com/huggingface/from-zero-to-research-an-introduction-to-meta-learning-8e16e677f78a

梯度下降优化算法的5个步骤

与这5个步骤等价的PyTorch代码也可以用5行代码编写：

predictions = model(inputs) # Forward passloss = loss_function(predictions, labels) # Compute loss functionloss.backward() # Backward passoptimizer.step() # Optimizer steppredictions = model(inputs) # Forward pass with new parameters

在 loss.backward()操作期间，为每个参数计算梯度(在动画中用绿色表示)，并存储在与每个参数相关的张量中：parameter.grad。(我们的动画中间的那张图)。

积累梯度意味着，在调用 optimizer.step()执行梯度下降步骤之前，我们将在 parameter.grad中对几个向后操作的梯度求和。这在PyTorch中很简单，因为除非我们调用 model.zero_grad()或 optimizer.zero_grad()，否则梯度张量不会重置。如果我们的损失在训练样本上取平均值，我们还需要除以累积步骤的数量。

下面是使用梯度积累训练模型的一个简单要点。在这个例子中，我们可以使用批处理大小 accumulation_steps进行训练，它比GPU的最大可容纳大小更大：

model.zero_grad() # Reset gradients tensorsfor i, (inputs, labels) in enumerate(training_set): predictions = model(inputs) # Forward pass loss = loss_function(predictions, labels) # Compute loss function loss = loss / accumulation_steps # Normalize our loss (if averaged) loss.backward() # Backward pass if (i+1) % accumulation_steps == 0: # Wait for several backward steps optimizer.step() # Now we can do an optimizer step model.zero_grad() # Reset gradients tensors if (i+1) % evaluation_steps == 0: # Evaluate the model when we... evaluate_model() # ...have no gradients accumulated

推广到极端情况

如果一个样本都无法放到GPU中，还可以训练模型吗？

如果你的架构没有太多的跳跃连接，是的，这是可能的！解决方案是使用*gradient-checkpoint *，以计算量换取内存。

基本上，这个想法是沿着模型以小块的形式反向传播梯度，用存储完整的反向传播图所需的内存交换与每个块关联的部分正向传递的额外计算。这是一种相当慢的方法，因为我们添加了额外的计算来减少内存需求，但是在某些设置中，它可能会很有趣，例如，在非常长的序列上训练RNN模型。

我不会在这里介绍更多的细节，只会给你提供相关的链接:

• TensorFlow：https://github.com/openai/gradient-checkpointing
• PyTorch doc：https://pytorch.org/docs/stable/checkpoint.html

“Memory-poor”策略，需要O(1)内存，但需要O (n²)计算步骤

充分利用多GPU的机器

现在让我们更具体地讨论一下多gpu上训练模型。

在多gpu服务器上训练PyTorch模型的首选策略是使用torch.nn.DataParallel。它是一个容器，通过将输入分割到指定的设备上，沿着批处理维度分块，从而并行化模块的应用程序。

DataParallel 非常容易使用，我们只要添加一行来封装模型：

parallel_model = torch.nn.DataParallel(model) # Encapsulate the modelpredictions = parallel_model(inputs) # Forward pass on multi-GPUsloss = loss_function(predictions, labels) # Compute loss functionloss.mean().backward() # Average GPU-losses + backward passoptimizer.step() # Optimizer steppredictions = parallel_model(inputs) # Forward pass with new parameters

然而，使用DataParallel可能出现一个问题：不平衡的GPU使用。

在某些设置下，GPU-1将比其他GPU使用的更多。

这个问题是怎么来的？我做了一个例子来更好地解释DataParallel的作用：

使用torch.nn.DataParallel来进行前向和后向传递

在前向传递的第4步(右上角)中，在GPU-1上收集所有并行计算的结果。这对于许多分类问题来说是可以的，但是当你在大的batch上训练语言模型时，就会出现问题。

让我们快速的计算出语言模型的输出大小：

如果我们假设有40k词汇表、序列中有250个tokens，每个批处理32个样本，每个元素存储在内存中需要4个字节，那么模型的输出大约需要1.2 GB。我们需要将其加倍来存储相关的梯度张量，因此我们的模型输出需要2.4 GB的内存！

这样就占了一个典型的10GB GPU内存非常大的一部分，这意味着GPU-1相对于其他GPU将被过度使用，从而限制了并行化的效果。

在不调整模型或优化方案的情况下，我们不能轻易地减少这个输出中的元素数量。但我们可以确保内存负载更均匀地分布在gpu之间。

在多GPU机器上负载均衡

GPU使用不平衡的问题主要有两种解决方案：

• 在前向传递中计算你的模型的的损失
• 以并行方式计算损失

第一个选项是最简单的，但有时由于各种原因，你不能使用它，所以让我们来谈谈第二个解决方案。沿着这条路，我们将学习关于PyTorch多GPU模块是如何工作的有趣的事情。

在这种情况下，解决方案是将每个部分的输出保存在各自的GPU上，而不是将所有输出都收集到GPU-1中。我们需要分配我们的损失计算以及能够计算和反向传播我们的损失。

值得庆幸的是，Hang Zhang开源了一个PyTorch包叫做PyTorch-Encoding，提供了这些定制的并行化功能。

我已经提取并稍微修改了这个模块，你可以从这里下载一个gist (parallel.py)来包含和调用代码。它主要由DataParallelModel和DataParallelCriterion两个模块组成，使用如下：

from parallel import DataParallelModel, DataParallelCriterionparallel_model = DataParallelModel(model) # Encapsulate the modelparallel_loss = DataParallelCriterion(loss_function) # Encapsulate the loss functionpredictions = parallel_model(inputs) # Parallel forward pass # "predictions" is a tuple of n_gpu tensorsloss = parallel_loss(predictions, labels) # Compute loss function in parallelloss.backward() # Backward passoptimizer.step() # Optimizer steppredictions = parallel_model(inputs) # Parallel forward pass with new parameters

DataParallelModel和torch.nn.DataParallel之间的区别只是前向传递(“预测”)的输出不是在GPU-1上收集的，因此是 n_gpu张量的元组，每个张量位于各自的GPU上。

DataParallelCriterion容器封装了损失函数，将 n_gpu张量的元组和目标标签张量作为输入。它在每个GPU上并行计算损失函数，以数据并行的方式分割目标标签张量。

我做了一个DataParallelModel/DataParallelCriterion内部的演示图：

使用DataParallelModel和DataParallelCriterion

以下是如何处理你可能遇到的两种特殊情况：

你的模型输出了几个张量：你希望将它们分开：output_1,output_2=zip(*predictions)

有时候你不想使用并行的损失函数：收集cpu上的所有张量：gathered_predictions=parallel.gather(predictions)

分布式训练：在几个机器上训练

现在，我们如何利用多个服务器的能力进行更大规模的训练呢?

最简单的选择是使用PyTorch的DistributedDataParalle，这几乎是上面讨论的DataParallel 的一个简易替代。

但是要注意：虽然代码看起来很相似，但是在分布式设置中训练模型将会改变你的工作流，因为你实际上必须在每个节点上启动独立的python训练脚本(这些脚本都是相同的)。我们将看到，一旦启动，这些训练脚本将由PyTorch分布式后端同步到一起。

在实践中，这意味着每个训练脚本将有：

• 它自己的优化器，并在每次迭代执行一个完整的优化步骤，不需要参数广播(DataParallel中的第2步)。
• 一个独立的Python解释器：这也将避免GIL-freeze，这可能来自于在一个Python解释器中驱动多个并行执行线程。

当多个并行的前向调用由一个解释器驱动时，在向前传递中大量使用Python循环/调用的模型会被Python解释器的GIL减慢速度。在这些设置中，DistributedDataParallel甚至可以在单机设置中有利地替换DataParallel。

现在让我们直接进入代码和用法。

DistributedDataParallel构建于torch.distributed包之上，该包提供用于同步分布式操作的底层原语，可以使用具有不同功能的多个后端(tcp、gloo、mpi、nccl)。

在这篇文章中，我将选择一种简单的方法开箱即用，但是你应该和Seb Arnold的指南和文档来深入了解这个模块。

我们将考虑一个简单但通用的设置，两个4-GPU服务器(节点)：

主服务器(服务器1)具有可访问的IP和用于通信的开放端口。

修改你的Python训练脚本来进行分布式训练

首先，我们需要调整我们的脚本，以便它可以在每台机器(节点)上单独运行。我们实际上要完全分布式，为每个节点的每个GPU运行一个单独的进程，总共8个进程。

我们的训练脚本有点长，我们需要初始化分布式后端来进行同步，训练的每个进程都是在数据的一个子集上进行封装模型和准备数据训练数据(每个进程是独立的，所以我们必须让他们自己去处理数据集的不同部分)。以下是更新后的代码：

from torch.utils.data.distributed import DistributedSamplerfrom torch.utils.data import DataLoader# Each process runs on 1 GPU device specified by the local_rank argument.parser = argparse.ArgumentParser()parser.add_argument("--local_rank