选自Medium

GPU捉襟见肘还想训练大批量模型?谁说不可以

深度学习模型和数据集的规模增长速度已经让 GPU 算力也开始捉襟见肘,如果你的 GPU 连一个样本都容不下,你要如何训练大批量模型?通过本文介绍的方法,我们可以在训练批量甚至单个训练样本大于 GPU 内存时,在单个或多个 GPU 服务器上训练模型。

深度学习模型和数据集的规模增长速度已经让 GPU 算力也开始捉襟见肘,如果你的 GPU 连一个样本都容不下,你要如何训练大批量模型?通过本文介绍的方法,我们可以在训练批量甚至单个训练样本大于 GPU 内存时,在单个或多个 GPU 服务器上训练模型。

2018 年的大部分时间我都在试图训练神经网络时克服 GPU 极限。无论是在含有 1.5 亿个参数语言模型(如 OpenAI 的大型生成预训练 Transformer 或最近类似的 BERT 模型)还是馈入 3000 万个元素输入的元学习神经网络(如我们在一篇 ICLR 论文《Meta-Learning a Dynamical Language Model》中提到的模型),我都只能在 GPU 上处理很少的训练样本。

但在多数情况下,随机梯度下降算法需要很大批量才能得出不错的结果。

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

有几个工具、技巧可以帮助你解决上述问题。在本文中,我将自己用过、学过的东西整理出来供大家参考。

在这篇文章中,我将主要讨论 PyTorch 框架。有部分工具尚未包括在 PyTorch(1.0 版本)中,因此我也写了自定义代码。

我们将着重探讨以下问题:

  • 在训练批量甚至单个训练样本大于 GPU 内存,要如何在单个或多个 GPU 服务器上训练模型;

  • 如何尽可能高效地利用多 GPU 机器;

  • 在分布式设备上使用多个机器的最简单训练方法。

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

你建的模型不错,在这个简洁的任务中可能成为新的 SOTA,但每次尝试在一个批量处理更多样本时,你都会得到一个 CUDA RuntimeError:内存不足。

这位网友指出了你的问题!

但你很确定将批量加倍可以优化结果。

你要怎么做呢?

这个问题有一个简单的解决方法:梯度累积。

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

与之对等的 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 中。

累积梯度意味着,在调用 optimizer.step() 实施一步梯度下降之前,我们会对 parameter.grad 张量中的几个反向运算的梯度求和。在 PyTorch 中这一点很容易实现,因为梯度张量在不调用 model.zero_grad() 或 optimizer.zero_grad() 的情况下不会重置。如果损失在训练样本上要取平均,我们还需要除以累积步骤的数量。

以下是使用梯度累积训练模型的要点。在这个例子中,我们可以用一个大于 GPU 最大容量的 accumulation_steps 批量进行训练:

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-checkpointing)来节省计算资源。

基本思路是沿着模型将梯度在小组件中进行反向传播,以额外的前馈传递为代价,节约存储完整的反向传播图的内存。这个方法比较慢,因为我们需要添加额外的计算来减少内存要求,但在某些设置中挺有意思,比如在非常长的序列上训练 RNN 模型(示例参见 https://medium.com/huggingface/from-zero-to-research-an-introduction-to-meta-learning-8e16e677f78a)。

这里不再赘述,读者可以查看以下链接:

  • TensorFlow:https://github.com/openai/gradient-checkpointing

  • PyTorch 文档:https://pytorch.org/docs/stable/checkpoint.html

「节约内存」(Memory-poor)策略需要 O(1) 的内存(但是要求 O(n²) 的计算步)。

充分利用多 GPU 机器

现在我们具体来看如何在多 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 使用不均衡。

在一些设置下,GPU-1 会比其他 GPU 使用率高得多。

这个问题从何而来呢?下图很好地解释了 DataParallel 的行为:

使用 torch.nn.DataParallel 的前向和后向传播。

在前向传播的第四步(右上),所有并行计算的结果都聚集在 GPU-1 上。这对很多分类问题来说是件好事,但如果你在大批量上训练语言模型时,这就会成为问题。

我们可以快速计算语言模型输出的大小:

语言模型输出中的元素数量。

假设我们的数据集有 4 万词汇,每一条序列有 250 个 token、每个 batch 中有 32 条序列,那么序列中的每一个元素需要 4 个字节的内存空间,模型的输出大概为 1.2GB。要储存相关的梯度张量,我们就需要把这个内存翻倍,因此我们的模型输出需要 2.4GB 的内存。

这是典型 10GB GPU 内存的主要部分,意味着相对于其它 GPU,GPU - 1 会被过度使用,从而限制了并行化的效果。

如果不调整模型和/或优化方案,我们就无法轻易减少输出中的元素数量。但我们可以确保内存负载在 GPU 中更均匀地分布。

多 GPU 机器上的均衡负载

解决办法是把每部分输出保留在其 GPU 上,而不是将它们聚集到 GPU-1 上。我们也需要分配损失标准计算,计算损失并进行反向传播。

幸而,张航开源了一个名为 PyTorch-Encoding 的 PyTorch 包,它包含了这些定制的并行化功能。

我提取并稍稍改动了这个模块,你可以从以下地址下载 gist(parallel.py)来纳入并调用你的代码。它主要包括两个模块:DataParallelModel 和 DataParallelCriterion,它们的用途如下:

下载地址:https://gist.github.com/thomwolf/7e2407fbd5945f07821adae3d9fd1312

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_gpu 张量的元组,每个张量分布在相应的 GPU 上。

DataParallelCriterion 容器封装了损失函数,并把 n_gpu 张量元组和目标标签张量作为输入。它在每个 GPU 上并行计算损失函数,像 DataParallel 分割模型输入一样分割目标标签张量

下图说明了 DataParallelModel/DataParallelCriterion 的内部情况:

使用 DataParallelModel 和 DataParallelCriterion。

以下是你可能会遇到的两个特定案例的解决办法:

  • 你的模型输出几个张量:你可能想分解它们:output_1, output_2 = zip(*predictions)

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

分布式训练:在多台机器上训练

在更大的批量上训练时,我们要如何控制多个服务器的算力呢?

最简单的选择是使用 PyTorch 的 DistributedDataParallel,它几乎可以说是以上讨论的 DataParallel 的直接替代元件。

但要注意:尽管代码看起来很相似,但在分布式设定中训练模型要改变工作流程,因为你必须在每个节点上启动一个独立的 Python 训练脚本。正如我们将看到的,一旦启动,这些训练脚本可以通过使用 PyTorch 分布式后端一起同步化。

在实践中,这意味着每个训练脚本将拥有:

  • 它自己的优化器,并在每次迭代中执行一个完整的优化步骤,不需要进行参数传播(DataParallel 中的步骤 2);

  • 一个独立的 Python 解释器:这也将避免 GIL-freeze,这是在单个 Python 解释器上驱动多个并行执行线程时会出现的问题。

当多个并行前向调用由单个解释器驱动时,在前向传播中大量使用 Python 循环/调用的模型可能会被 Python 解释器的 GIL 放慢速度。通过这种设置,DistributedDataParallel 甚至在单台机器设置中也能很方便地替代 DataParallel。

现在我们直接讨论代码和用途。

DistributedDataParallel 是建立在 torch.distributed 包之上的,这个包可以为同步分布式运算提供低级原语,并能以不同的性能使用多种后端(tcp、gloo、mpi、nccl)。在这篇文章中,我将选择一种简单的开箱即用的方式来使用它,但你应该阅读文档和 Séb Arnold 写的教程来深入理解这个模块。

  • 文档:https://pytorch.org/docs/stable/distributed.html

  • 教程:https://pytorch.org/tutorials/intermediate/dist_tuto.html

我们将考虑使用具有两个 4 - GPU 服务器(节点)的简单但通用的设置:

主服务器(服务器 1)拥有一个可访问的 IP 地址和一个用于通信的开放端口。

改写 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)                                     # Optimizer step

在第二台机器上,我们类似地启动脚本:

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,以此类推…)。

工程分布式计算并行计算
7
相关数据
深度学习技术

深度学习(deep learning)是机器学习的分支,是一种试图使用包含复杂结构或由多重非线性变换构成的多个处理层对数据进行高层抽象的算法。 深度学习是机器学习中一种基于对数据进行表征学习的算法,至今已有数种深度学习框架,如卷积神经网络和深度置信网络和递归神经网络等已被应用在计算机视觉、语音识别、自然语言处理、音频识别与生物信息学等领域并获取了极好的效果。

参数技术

在数学和统计学裡,参数(英语:parameter)是使用通用变量来建立函数和变量之间关系(当这种关系很难用方程来阐述时)的一个数量。

损失函数技术

在数学优化,统计学,计量经济学,决策理论,机器学习和计算神经科学等领域,损失函数或成本函数是将一或多个变量的一个事件或值映射为可以直观地表示某种与之相关“成本”的实数的函数。

元学习技术

元学习是机器学习的一个子领域,是将自动学习算法应用于机器学习实验的元数据上。现在的 AI 系统可以通过大量时间和经验从头学习一项复杂技能。但是,我们如果想使智能体掌握多种技能、适应多种环境,则不应该从头开始在每一个环境中训练每一项技能,而是需要智能体通过对以往经验的再利用来学习如何学习多项新任务,因此我们不应该独立地训练每一个新任务。这种学习如何学习的方法,又叫元学习(meta-learning),是通往可持续学习多项新任务的多面智能体的必经之路。

TensorFlow技术

TensorFlow是一个开源软件库,用于各种感知和语言理解任务的机器学习。目前被50个团队用于研究和生产许多Google商业产品,如语音识别、Gmail、Google 相册和搜索,其中许多产品曾使用过其前任软件DistBelief。

张量技术

张量是一个可用来表示在一些矢量、标量和其他张量之间的线性关系的多线性函数,这些线性关系的基本例子有内积、外积、线性映射以及笛卡儿积。其坐标在 维空间内,有 个分量的一种量,其中每个分量都是坐标的函数,而在坐标变换时,这些分量也依照某些规则作线性变换。称为该张量的秩或阶(与矩阵的秩和阶均无关系)。 在数学里,张量是一种几何实体,或者说广义上的“数量”。张量概念包括标量、矢量和线性算子。张量可以用坐标系统来表达,记作标量的数组,但它是定义为“不依赖于参照系的选择的”。张量在物理和工程学中很重要。例如在扩散张量成像中,表达器官对于水的在各个方向的微分透性的张量可以用来产生大脑的扫描图。工程上最重要的例子可能就是应力张量和应变张量了,它们都是二阶张量,对于一般线性材料他们之间的关系由一个四阶弹性张量来决定。

神经网络技术

(人工)神经网络是一种起源于 20 世纪 50 年代的监督式机器学习模型,那时候研究者构想了「感知器(perceptron)」的想法。这一领域的研究者通常被称为「联结主义者(Connectionist)」,因为这种模型模拟了人脑的功能。神经网络模型通常是通过反向传播算法应用梯度下降训练的。目前神经网络有两大主要类型,它们都是前馈神经网络:卷积神经网络(CNN)和循环神经网络(RNN),其中 RNN 又包含长短期记忆(LSTM)、门控循环单元(GRU)等等。深度学习是一种主要应用于神经网络帮助其取得更好结果的技术。尽管神经网络主要用于监督学习,但也有一些为无监督学习设计的变体,比如自动编码器和生成对抗网络(GAN)。

梯度下降技术

梯度下降是用于查找函数最小值的一阶迭代优化算法。 要使用梯度下降找到函数的局部最小值,可以采用与当前点的函数梯度(或近似梯度)的负值成比例的步骤。 如果采取的步骤与梯度的正值成比例,则接近该函数的局部最大值,被称为梯度上升。

随机梯度下降技术

梯度下降(Gradient Descent)是遵循成本函数的梯度来最小化一个函数的过程。这个过程涉及到对成本形式以及其衍生形式的认知,使得我们可以从已知的给定点朝既定方向移动。比如向下朝最小值移动。 在机器学习中,我们可以利用随机梯度下降的方法来最小化训练模型中的误差,即每次迭代时完成一次评估和更新。 这种优化算法的工作原理是模型每看到一个训练实例,就对其作出预测,并重复迭代该过程到一定的次数。这个流程可以用于找出能导致训练数据最小误差的模型的系数。

分类问题技术

分类问题是数据挖掘处理的一个重要组成部分,在机器学习领域,分类问题通常被认为属于监督式学习(supervised learning),也就是说,分类问题的目标是根据已知样本的某些特征,判断一个新的样本属于哪种已知的样本类。根据类别的数量还可以进一步将分类问题划分为二元分类(binary classification)和多元分类(multiclass classification)。

优化器技术

优化器基类提供了计算梯度loss的方法,并可以将梯度应用于变量。优化器里包含了实现了经典的优化算法,如梯度下降和Adagrad。 优化器是提供了一个可以使用各种优化算法的接口,可以让用户直接调用一些经典的优化算法,如梯度下降法等等。优化器(optimizers)类的基类。这个类定义了在训练模型的时候添加一个操作的API。用户基本上不会直接使用这个类,但是你会用到他的子类比如GradientDescentOptimizer, AdagradOptimizer, MomentumOptimizer(tensorflow下的优化器包)等等这些算法。

语言模型技术

语言模型经常使用在许多自然语言处理方面的应用,如语音识别,机器翻译,词性标注,句法分析和资讯检索。由于字词与句子都是任意组合的长度,因此在训练过的语言模型中会出现未曾出现的字串(资料稀疏的问题),也使得在语料库中估算字串的机率变得很困难,这也是要使用近似的平滑n元语法(N-gram)模型之原因。

暂无评论
暂无评论~