增大Batch训练神经网络：单GPU、多GPU及分布式配置的实用技巧
2018年中的大部分时间，我都在尝试利用训练神经网络克服GPUs的局限。无论是在包含1.5亿个参数的语言模型中，比如OpenAI’s huge Generative Pre-trained Transformer (or the recent and similar BERT model)，还是在拥有3000万个输入元素的神经网络中，我都只能利用GPU处理很少的训练样本。
可是若想利用随机梯度下降算法得出不错的结果，大批量的训练样本必不可少。
如果你的GPU只能处理少量样本，该如何训练大批量模型呢？
接下来，我将介绍几类工具和技巧。
本文主要会讨论PyTorch框架，并就以下几个问题进行探讨：
- 当训练批量甚至单个训练样本大于GPU内存时，如何训练模型；
- 如何高效地利用多GPU机器；
- 如何在分布式设备上简单的使用多个机器。
在一个或多个GPU上训练大批量模型
你构建了一个不错的模型，可在尝试处理更多样本时，却得到CUDA RuntimeError：内存不足。
根据网友的回答你明白，加倍批量可以对结果进行优化。
此时，梯度累积（accumulating gradients）可以帮助到你。
PyTorch代码如下所示：
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中。
梯度累积意味着，在调用potimizer.step()实现梯度下降之前，我们会求取parameter.grad张量中的几个反向操作的梯度和。
如下是使用梯度累积训练模型的示例。
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-checkpoingting）节省计算资源。
梯度检查点会将我们连续计算的元前馈和元反向传播切分成片段。但由于需要增加额外的计算以减少内存需求，该方法效率不高。不过，它在某些示例中又有较为明显的优势，比如在长序列上训练RNN模型，点击此处查看详情。
或有兴趣可进入下列文档进行查询：
TensorFlow：https://github.com/openai/gradient-checkpointing
PyTorch doc：https://pytorch.org/docs/stable/checkpoint.html
A “Memory-poor” strategy that needs O(1) memory (but requires O(n²) computation steps) — From Yaroslav Bulatov’s nice post: https://medium.com/tensorflow/fitting-larger-networks-into-memory-583e3c758ff9
多GPU机器
在多GPU服务器上训练PyTorch模型首选torch.nn.DataParallel。该策略能够在多个指定设备上按照batch dimension分割输入，实现并行化模块。
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.backward() # Backward pass optimizer.step() # Optimizer step predictions = parallel_model(inputs) # Forward pass with new parameters
但DataParallel存在GPU使用不均衡的问题，下图给出了相应解释：
Forward and Backward passes with torch.nn.DataParallel
在前向传播的第四个步骤（见右上）中，GPU-1汇集了所有并行计算的结果。
通过下列所示的方式能够计算出语言模型输出的大小：
Number of elements in the output of a language model
现有如下假设：数据集共含4万词汇，序列中包含250 tokens，每个batch 包含32个示例，每个元素4 bytes，模型的输出占用1.2GB。但我们需要2.4GB的内存才能存储相关的梯度张量。
这种存储方式会使得GPU-1被过度使用，从而造成GPU使用不均衡的问题。
多GPU机器上的负载均衡
想要解决GPU使用不均衡的问题需要将每部分输出都保留在原有的GPU上，而不汇集于GPU-1。
张航开源了名为PyTorch-Encoding的包，可用于缓解上述问题。
我对这个开源包做了一些调整，你可以点击此处下载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 parameters
DataParallelModel不同于torch.nn.DataParallel的是，前向传播的输出（predictions）没有汇集在GPU-1中，而是作为n_gup张量的元组分布在相应的GPU上。
DataParallelCriterion容器封装了损失函数，并且将n_gpu张量的元组和目标标签张量作为输入。
下图描述了DataParallelModel/DataParallelCriterion的内部情况：
下面有两个特殊情况，并给出了解决办法：
- 模型输出了一些张量：你可以利用output_1,output_2 = zip(*predictions)分解它们。
- 若你不想并行计算损失函数，则可以利用gathered_prdictions = parallel.gather(predictions)收集张量。
分布式训练
PyTorch中的DistributedDataParallel可以帮助我们在遇到大批量训练问题时，拥有控制多个服务器的运算能力。
但值得注意的是：由于对每个节点都要启动一个独立的Python训练脚本，在设定时需要注意改变工作流程。
每个脚本在训练中都会拥有：
- 它自己的优化器，在每次迭代中都执行一个完整的优化，不需要参数传输。
- 一个独立的Python解释器：能够避免GIL-freeze
在后面我们将通过代码进行讨论：
torch.distributed包能够为同步分布式运算提供低级原语，基于此构建得到DistributedDataParallel。你可以通过阅读文档以及教程对其进行进一步理解。
接下来，我们将使用具有两个4-GPU的服务器。
The main server (server 1) has an accessible IP and an open port for communication.
升级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地址（示例中为：196.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参数之外，上述两个命令相同。
扩展
如果你觉得在计算机集群上运行一组几乎相同的命令有些枯燥，可点击此处了解GNU并行。
以上为译文
本文由阿里云云栖社区组织翻译。
文章原标题《Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups》，作者：
Thomas Wolf，译者：Elaine，审校：袁虎。