点击上方“AI公园”，关注公众号，选择加“星标“或“置顶”
导读作者：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个步骤等价的PyTorch代码也可以用5行代码编写：
predictions = model(inputs) # Forward pass
loss = loss_function(predictions, labels) # Compute loss function
loss.backward() # Backward pass
optimizer.step() # Optimizer step
predictions = 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 tensors
for 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
充分利用多GPU的机器
现在让我们更具体地讨论一下多gpu上训练模型。
在多gpu服务器上训练PyTorch模型的首选策略是使用torch.nn.DataParallel。它是一个容器，通过将输入分割到指定的设备上，沿着批处理维度分块，从而并行化模块的应用程序。
DataParallel 非常容易使用，我们只要添加一行来封装模型：
parallel_model = torch.nn.DataParallel(model) # Encapsulate the model
predictions = parallel_model(inputs) # Forward pass on multi-GPUs
loss = loss_function(predictions, labels) # Compute loss function
loss.mean().backward() # Average GPU-losses + backward pass
optimizer.step() # Optimizer step
predictions = parallel_model(inputs) # Forward pass with new parameters然而，使用DataParallel可能出现一个问题：不平衡的GPU使用。
在某些设置下，GPU-1将比其他GPU使用的更多。
这个问题是怎么来的？我做了一个例子来更好地解释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, DataParallelCriterion
parallel_model = DataParallelModel(model) # Encapsulate the model
parallel_loss = DataParallelCriterion(loss_function) # Encapsulate the loss function
predictions = parallel_model(inputs) # Parallel forward pass
# "predictions" is a tuple of n_gpu tensors
loss = parallel_loss(predictions, labels) # Compute loss function in parallel
loss.backward() # Backward pass
optimizer.step() # Optimizer step
predictions = parallel_model(inputs) # Parallel forward pass with new parametersDataParallelModel和torch.nn.DataParallel之间的区别只是前向传递(“预测”)的输出不是在GPU-1上收集的，因此是 n_gpu张量的元组，每个张量位于各自的GPU上。
DataParallelCriterion容器封装了损失函数，将 n_gpu张量的元组和目标标签张量作为输入。它在每个GPU上并行计算损失函数，以数据并行的方式分割目标标签张量。
我做了一个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服务器(节点)：
修改你的Python训练脚本来进行分布式训练
首先，我们需要调整我们的脚本，以便它可以在每台机器(节点)上单独运行。我们实际上要完全分布式，为每个节点的每个GPU运行一个单独的进程，总共8个进程。
我们的训练脚本有点长，我们需要初始化分布式后端来进行同步，训练的每个进程都是在数据的一个子集上进行封装模型和准备数据训练数据(每个进程是独立的，所以我们必须让他们自己去处理数据集的不同部分)。以下是更新后的代码：
from torch.utils.data.distributed import DistributedSampler
from 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", type=int)
args = parser.parse_args()
# Initializes the distributed backend which will take care of sychronizing nodes/GPUs
torch.distributed.init_process_group(backend='nccl')
# Encapsulate the model on the GPU assigned to the current process
device = torch.device('cuda', arg.local_rank)
model = model.to(device)
distrib_model = torch.nn.parallel.DistributedDataParallel(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
# Restricts data loading to a subset of the dataset exclusive to the current process
sampler = DistributedSampler(dataset)
dataloader = DataLoader(dataset, sampler=sampler)
for inputs, labels in dataloader:
predictions = distrib_model(inputs.to(device)) # Forward pass
loss = loss_function(predictions, labels.to(device)) # Compute loss function
loss.backward() # Backward pass
optimizer.step() # Optimizer step运行Python训练脚步的多个实例
我们已经快要完成了。我们只需在每个服务器上启动一个训练脚本实例。
要运行脚本，我们将使用PyTorch的torch.distributed.launch。它将负责设置环境变量，并使用正确的 local_rank参数调用每个脚本。
第一个机器将是我们的主机，它需要从所有其他机器访问，因此有一个可访问的IP地址(在我们的示例中是192.168.1.1)和一个开放端口(在我们的示例中是1234)。在这台机器上，我们使用torch.distributed.launch运行我们的训练脚本：
python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=0 --master_addr="192.168.1.1" --master_port=1234 OUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script)在第二台机器上，我们同样启动脚本：
python -m torch.distributed.launch --nproc_per_node=4 --nnodes=2 --node_rank=1 --master_addr="192.168.1.1" --master_port=1234 OUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other arguments of our training script)除了 --node_rank参数在第一台机器上设置为' 0 '，在第二台机器上设置为' 1 '(在另一台服务器上设置为' 2 '，等等)以外，这两个命令是相同的。
在一组机器上运行一组几乎相同的命令的过程可能看起来有点乏味。所以现在可能是学习……GNU parallel的好时机。
—
END—
英文原文：https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255
请长按或扫描二维码关注本公众号
喜欢的话，请给我个好看吧！
292
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
