混合精度训练之APEX(上)

跬步亦是千里~~~

——子棐

前言

为了更快的在Pytorch中加入对Volta GPU的支持,并实现针对混合精度训练的优化,NVIDIA发布了Apex开源工具库。

对于混合精度训练问题,该工具库提供了AMP(自动混合精度)以及FP16_Optimizer两种不同的库。AMP提供的是自动的混合精度训练支持,该库会自动检查执行的操作,对于FP16安全的操作在训练中Cast到FP16,反之则选择FP32,从而实现一键混合精度训练,实现了在训练稳定性和易用性的平衡。而FP16_Optimizer提供的是高阶版混合精度训练,其提供了更多详细的实现细节,对于整个网络完全采用FP16,适合最为前沿的玩家更加灵活的实现FP16混合精度训练。

除了针对混合精度训练问题的解决方案,Apex还提供了另外五个库,分别是Apex_C、RNN、Parallel、Reparameterization和Example,这五个库分别是部分功能函数Cuda实现、提供RNN在混合精度训练的实现、提供高性能的多GPU训练(MPI+NCCL)实现、解决模型权值正则化实现问题以及Apex使用范例。

不积跬步无以至千里,因此本文将先从AMP开始介绍,高端玩家可以直接前往FP16_Optimizer区,也欢迎投稿分享经验~~~

Apex之AMP

AMP提供了最为简单的混合精度训练支持,其是通过黑白名单的方式予以实现的。在AMP目录下有Lists目录,其中包括了以下三个库的黑白名单:

Functional库(torch.nn.functional):其支持FP16的Function为'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 'conv_transpose2d', 'conv_transpose3d', 'conv_tbc', 'linear';

tensor库(torch.Tensor和torch.autograd.Variable):其支持FP16 Fuction为 '__matmul';

torch库:其支持FP16 Function为:'conv1d', 'conv2d', 'conv3d', 'conv_transpose1d', 'conv_transpose2d', 'conv_transpose3d', 'conv_tbc', 'addmm', 'addmv', 'addr', 'matmul', 'mm', 'mv'。

Apex对RNN并不像黑白名单里面直接进行“on-ramp”的cast操作,而是通过rnn_cast函数来支持的,其是通过Wrap torch.nn 后端的顶层的RNN,然后调用CudnnRNN或者AutogradRNN来实现的。

AMP的初始化实现在amp.py中,初始化函数为amp.init(),其内部初始化了一个AmpHandle对象,并对不同的层、函数、RNN进行Wrap配置。在反向迭代时,通过调用AmpHandle函数内部的scale_loss函数进行反向迭代。其中amp.init()函数有三个输入参数,分别是enable、enable_caching和verbose。其中enable是模式打开,即启用amp,如果关闭则在初始化是初始一个NoOpHandle对象,其实不会进行任何的操作。而enable_caching决定是否缓存每次Cast的FP16对象,以解决RNN的问题。而verbose配置打开后可以看到每次的Cast操作。这里以Mnist为例,每次迭代,AMP对整个网络的Cast操作为:

Float->Half (conv2d)

Float->Half (conv2d)

Float->Half (conv2d)

Float->Half (conv2d)

Float->Half (conv2d)

Float->Half (linear)

Float->Half (linear)

Float->Half (linear)

Float->Half (linear)

Half->Float (log_softmax)

Apex之Automatic Loss Scaling

关于Loss Scaling已经在前面几期的混合精度训练中有过详细的解释,如果对这个概念还感到陌生的话,可以再次重温来自子棐同学的大作——Tesla V100之NVCaffe混合精度训练深度解析篇(中)。在Apex中或者说Pytorch中所提的Loss Scalling概念和原理和NVcaffe所述一致的,但是在实现上有着不同的思路(如果完全一样这一章节也就没有存在的必要了)。

在Apex中不仅有Loss Scalling而且还有强化升级版“Automatic”,实现算法的如下框所示:

Start with a large scaling factor s

For each training iteration

Make an fp16 copy of weights

Fwd prop

Scale the loss by s

Bwd prop

Update scaling factor s

If dW contains Inf/NaN then reduce s skip the update

If no Inf/NaN were detected for N updates then increase s

Scale dW by 1/s

Update W

在训练开始前会设置一个较大的Scaling Factor S。对于每一次的迭代,和Nvcaffe一样需要维护一个Master权值,因此首先是将FP32权值复制一份用以留存,并生成一份FP16权值。之后利用这个FP16权值进行正向迭代,这个过程就可以充分利用到Tensor Core 125T的计算性能。正向迭代完成后计算Loss,此时的Loss计算值为FP32格式。之后将Loss乘以S以在横轴上右移(放大)Loss获得Scaled_Loss,这样可以在反向迭代转换至FP16时能尽可能的保留FP32的分布特性。下一步基于Scaled_Loss进行反向迭代以计算梯度。计算完成后如果dW中存在着Inf/NaN的话证明存在Overflow,会缩小S值并放弃这一次的更新。如果不存在Inf/NaN则证明放大没有问题,则在N次成功的更新后放大S值,在更新Master权值前要将dW缩小S倍,然后再更新到Master Weight上。

上图是训练一个NMT网络时对S的分析结果图,如图可知,在前面的迭代中S一直在扩大,随后达到最大值,之后一直在一定的范围内上下浮动。通过分析图可知:

最小的S为2的20次方,意味着最大dW不超过2的-5次方;

这种方法能非常高效的为网络设置合适的Scaling Factor。

当然如果是观察仔细的小伙伴肯定会发现几个小问题:

多了一个参数N,作为Scaling Factor的扩大周期;

如果跳过参数更新,是否需要增加迭代计数;

如果跳过参数更新,是否需要对该Batch重新计算。

通过大量的实验总结得知,如果跳过参数更新,是否增加迭代计算并不会对性能造成影响,而且无需重新计算该Batch,直接进行下一个Batch的计算和更新即可。对于参数N,在实验中选取的是2000,而且对于S的缩放采用了2次方的形式。至于上述更新依然可行的原因在于,通过Automatic Loss Scaling可以最高效的实现梯度更新,几个Batch的更新缺失并不会对训练效果造成影响。

本章的最后,给各位小伙伴留一个家庭作业:在Nvcaffe是否有Automatic Loss Scaling呢?如果有,又是怎么实现的呢?

Apex之快速使用指南

安装

软件使用第一步永远都是安装,Apex所需依赖环境为:

Python 3

CUDA 9

PyTorch 0.4

使用Docker Hub提供的最新Cuda9.2的镜像以及PIP安装最新的Pytorch即可正常编译安装Apex。

更加推荐使用NGC已经优化的容器镜像,使用方法如下:

前往ngc.nvidia.com注册登录并获取API Key

使用API Key登录到容器镜像仓库,命令:docker login

拉取最新容器镜像:docker pull nvcr.io/nvidia/pytorch:18.05-py3

运行容器:docker run --runtime=nvidia --rm -tinvcr.io/nvidia/pytorch:18.05-py3 bash

下载Apex源码:git clone https://github.com/NVIDIA/apex.git

编译安装Apex:cd apex && python setup.py install

*这里需要注意的是NGC的Pytorch建议使用18.04之后版本。

针对单一Optimizer使能AMP

对于单一Optimizer,使能AMP只需要以下修改两部分即可:

1. 代码中Optimizer之后添加:

optimizer = # ... some optimizer

from apex import amp

amp_handle = amp.init()

2. 修改代码中反向迭代部分:

# ... do a bunch of stuff to compute a loss

loss.backward()

optimizer.step()

# ...finish the iteration

修改为

# ... same as before

with amp_handle.scale_loss(loss, optimizer) as scaled_loss:

scaled_loss.backward()

optimizer.step()

# ... same as before

针对多Optimizer使能AMP

多Optimizer使能AMP相较单一Optimizer复杂,需要显性的warp多个Optimizer,并处理多个Loss的反向迭代,例如:

amp_handle = amp.init()

optimizer = amp_handle.wrap_optimizer(optimizer, num_loss=2)

# ...

optimizer.zero_grad()

loss1 = ComputeLoss1(model)

with optimizer.scale_loss(loss1) as scaled_loss:

scaled_loss.backward()

# ...

loss2 = ComputeLoss2(model)

with optimizer.scale_loss(loss2) as scaled_loss:

scaled_loss.backward()

# ...

optimizer.step()

AMP添加自定义函数

在AMP添加自定义函数的支持有两种方法,分别使用Function annotation或者函数显性注册如下所示:

对于自定义函数Fru

from backend import FRUBackend

def fru(input, hidden, weight, bias):

# call to CUDA code

FRUBackend(input, hidden, weight, bias)

使用Function Annotation

@amp.half_function

def fru(input, hidden, weight, bias):

#...

使用显性注册

import backend

amp.register_half_function(backend, 'FRUBackend')

amp.init()

此处的Function Annotation有三种类型分别是:

类似显性注册也有对应三种方式,分别是:

amp.register_half_function(module, function_name)

amp.register_float_function(module, function_name)

amp.register_promote_function(module, function_name)

这里需要注意的是,采用显性注册需要在amp.init()之前。

至此,关于AMP的介绍和使用指南已经完全阐明,关于前文留的小问题,还欢迎小伙伴们留言回答哟~~~

以上

子棐

  • 发表于:
  • 原文链接https://kuaibao.qq.com/s/20180624G0UZD400?refer=cp_1026
  • 腾讯「云+社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。
  • 如有侵权,请联系 yunjia_community@tencent.com 删除。

同媒体快讯

扫码关注云+社区

领取腾讯云代金券