2018年中的大部分时间,我都在尝试利用训练神经网络克服GPUs的局限。无论是在包含1.5亿个参数的语言模型中,比如OpenAI’s huge Generative Pre-trained Transformer (or the recent and similar BERT model),还是在拥有3000万个输入元素的神经网络中,我都只能利用GPU处理很少的训练样本。

可是若想利用随机梯度下降算法得出不错的结果,大批量的训练样本必不可少。

如果你的GPU只能处理少量样本,该如何训练大批量模型呢?

接下来,我将介绍几类工具和技巧。

本文主要会讨论PyTorch框架,并就以下几个问题进行探讨:

  1. 当训练批量甚至单个训练样本大于GPU内存时,如何训练模型;
  2. 如何高效地利用多GPU机器;
  3. 如何在分布式设备上简单的使用多个机器。

在一个或多个GPU上训练大批量模型

你构建了一个不错的模型,可在尝试处理更多样本时,却得到CUDA RuntimeError:内存不足。

怎么看batchsize神经网络(增大Batch训练神经网络)(1)

根据网友的回答你明白,加倍批量可以对结果进行优化。

此时,梯度累积(accumulating gradients)可以帮助到你。

怎么看batchsize神经网络(增大Batch训练神经网络)(2)

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

怎么看batchsize神经网络(增大Batch训练神经网络)(3)

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使用不均衡的问题,下图给出了相应解释:

怎么看batchsize神经网络(增大Batch训练神经网络)(4)

Forward and Backward passes with torch.nn.DataParallel

在前向传播的第四个步骤(见右上)中,GPU-1汇集了所有并行计算的结果。

通过下列所示的方式能够计算出语言模型输出的大小:

怎么看batchsize神经网络(增大Batch训练神经网络)(5)

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的内部情况:

怎么看batchsize神经网络(增大Batch训练神经网络)(6)

下面有两个特殊情况,并给出了解决办法:

  1. 模型输出了一些张量:你可以利用output_1,output_2 = zip(*predictions)分解它们。
  2. 若你不想并行计算损失函数,则可以利用gathered_prdictions = parallel.gather(predictions)收集张量。

分布式训练

PyTorch中的DistributedDataParallel可以帮助我们在遇到大批量训练问题时,拥有控制多个服务器的运算能力。

但值得注意的是:由于对每个节点都要启动一个独立的Python训练脚本,在设定时需要注意改变工作流程。

每个脚本在训练中都会拥有:

  1. 它自己的优化器,在每次迭代中都执行一个完整的优化,不需要参数传输。
  2. 一个独立的Python解释器:能够避免GIL-freeze

在后面我们将通过代码进行讨论:

torch.distributed包能够为同步分布式运算提供低级原语,基于此构建得到DistributedDataParallel。你可以通过阅读文档以及教程对其进行进一步理解。

接下来,我们将使用具有两个4-GPU的服务器。

怎么看batchsize神经网络(增大Batch训练神经网络)(7)

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并行。

怎么看batchsize神经网络(增大Batch训练神经网络)(8)


以上为译文

本文由阿里云云栖社区组织翻译。

文章原标题《Training Neural Nets on Larger Batches: Practical Tips for 1-GPU, Multi-GPU & Distributed setups》,作者:

Thomas Wolf,译者:Elaine,审校:袁虎。

,