老板, 咱们就一台 Titan Xp, 训不动 BERT 呀
没钱买机器, 自己想办法 .
委屈 T^T
我听说 混合精度训练 可以从算法上缓解这个问题?
喵喵喵??
其实小夕的内心是拒绝的, 就一台破 Xp, 再优化能快到哪里去呀 T^T
燃鹅
小夕找了一份开源代码, 结果刚开始跑小夕就震惊了! 什么鬼? 训练速度怎么这么快? 出 bug 了吧????
一毛一样的模型, 超参数和硬件环境, 竟然可以获得 2.X 倍的加速.
关键的关键, 这不是一个特例, 在各类网络训练问题上都提速明显, 遍地开花~~~
混合精度训练
一切还要从 2018 年 ICLR 的一篇论文说起...
《MIXED PRECISION TRAINING》
这篇论文是百度 & Nvidia 研究院一起发表的, 结合 N 卡底层计算优化, 提出了一种灰常有效的神经网络训练加速方法, 不仅是预训练, 在全民 finetune BERT 的今天变得异常有用哇. 而且小夕调研了一下, 发现不仅百度的 paddle 框架支持混合精度训练, 在 Tensorflow 和 Pytorch 中也有相应的实现. 下面我们先来讲讲 理论 , 后面再分析混合精度训练在三大深度学习框架中的 打开方式 .
理论原理
训练过神经网络的小伙伴都知道, 神经网络的参数和中间结果绝大部分 都是 单精度浮点数 (即 float32)存储和计算的, 当网络变得超级大时, 降低浮点数精度, 比如使用 半精度浮点数 , 显然是提高计算速度, 降低存储开销的一个很直接的办法. 然而副作 用也很显然, 如果我们直接降低浮点数的精度直观上必然导致模型训练精度的损失. 但是呢, 天外有天, 这篇文章用了 三种机制 有效地防止了模型的精度损失. 待小夕一一说来 o(*~▽~*)ブ
权重备份(master weights)
我们知道半精度浮点数 (float16) 在计算机中的表示分为 1bit 的 符号位 ,5bits 的 指数位 和 10bits 的 尾数位 , 所以它能表示的最小的 正数即 2^-24 (也就是 精度到此为止了). 当神经网络中的梯度灰常小的时候, 网络训练过程中每一步的 迭代 (灰常小的梯度 也黑小的 learning rate)会变得更小, 小到 float16 精度无法表示的时候, 相应的梯度就无法得到更新.
论文统计了一下在 Mandarin 数据集上训练 DeepSpeech 2 模型时产生过的梯度, 发现在未乘以 learning rate 之前, 就有接近 5% 的梯度直接悲剧的变成 0(精度比 2^-24 还要高的梯度会直接变成 0), 造成重大的损失呀 /(ㄒ o ㄒ)/~~
还有更难的, 假设迭代量逃过一劫准备奉献自己的时候. . . 由于网络中的权重往往远大于我们要更新的量, 当迭代量小于 Float16 当前区间内能表示的最小间隔的时候, 更新也会失败(哭瞎┭┮﹏┭┮我怎么这么难鸭)
所以怎么办呢? 作者这里提出了一个非常 simple but effective 的方法, 就是前向传播和梯度计算都用 float16, 但是存储网络参数的梯度时要用 float32! 这样就可以一定程度上的解决上面说的两个问题啦~~~
我们来看一下训练曲线, 蓝色的线是正常的 float32 精度训练曲线, 橙色的线是使用 float32 存储网络参数的 learning curve, 绿色滴是不使用 float32 存储参数的曲线, 两者一比就相形见绌啦.
损失放缩(loss scaling)
有了上面的 master weights 已经可以足够高精度的训练很多网络啦, 但是有点强迫症的小夕来说怎么还是觉得有点不对呀 o((⊙﹏⊙))o.
虽然使用 float32 来存储梯度, 确实不会丢失精度了, 但是计算过程中出现的指数位小于 -24 的梯度不还是会丢失的嘛! 相当于用漏水的筛子从河边往村里运水, 为了多存点水, 村民们把储水的碗换成了大缸, 燃鹅筛子依然是漏的哇, 在路上的时候水就已经漏的木有了..
于是 loss scaling 方法来了. 首先作者统计了一下训练过程中激活函数梯度的分布情况, 由于网络中的梯度往往都非常小, 导致在使用 FP16 的时候右边有大量的范围是没有使用的. 这种情况下, 我们可以通过放大 loss 来把整个梯度右移, 减少因为精度随时变为 0 的梯度.
那么问题来了, 怎么合理的放大 loss 呢? 一个最简单的 方法是 常数缩放 , 把 lo ss 一股脑统一放大 S 倍. float16 能表示的最大正数是 2^15*(1+1-2^-10)=65504, 我们可以统计网络中的梯度, 计算出一个常数 S, 使得最大的梯度不超过 float16 能表示的最大整数即可.
当然啦, 还有更加智 能的 动态调整 (automatic scaling) o(*~▽~*)ブ
我们先初始化一个很大的 S, 如果梯度溢出, 我们就把 S 缩小为原来的二分之一; 如果在很多次迭代中梯度都没有溢出, 我们也可以尝试把 S 放大两倍. 以此类推, 实现动态的 loss scaling.
运算精度(precison of ops)
精益求精再进一步, 神经网络中的运算主要可以分为四大类, 混合精度训练把一些有更高精度要求的运算, 在计算过程中使用 float32, 存储的时候再转换为 float16.
matrix multiplication: linear, matmul, bmm, conv
pointwise: relu, sigmoid, tanh, exp, log
reductions: batch norm, layer norm, sum, softmax
loss functions: cross entropy, l2 loss, weight decay
像矩阵乘法和绝大多数 pointwise 的计算可以直接使用 float16 来计算并存储, 而 reductions,loss function 和一些 pointwise(如 exp,log,pow 等函数值远大于变量的函数)需要更加精细的处理, 所以在计算中使用用 float32, 再将结果转换为 float16 来存储.
总结陈词
混合精度训练做到了在前向和后向计算过程中均使用半精度浮点数, 并且没有像之前的一些工作一样还引入额外超参, 而且重要的是, 实现非常简单却能带来非常显著的收 益, 在 显存 half 以及 速度 double 的情况下保持模型的精度, 简直不能再厉害啦.
三大深度学习框架的打开方式
看完了硬核技术细节之后, 我们赶紧来看看代码实现吧! 如此强大的混合精度训练的代码实现不要太简单了吧: open_mouth:
Pytorch
导入 Automatic Mixed Precision (AMP), 不要 998 不要 288, 只需 3 行无痛使用!
- from apex import amp
- model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是 "欧一", 不是 "零一"
- with amp.scale_loss(loss, optimizer) as scaled_loss:
- scaled_loss.backward()
来看个例子, 将上面三行按照正确的位置插入到自己原来的代码中就可以实现酷炫的半精度训练啦!
- import torch
- from apex import amp
- model = ...
- optimizer = ...
- # 包装 model 和 optimizer
- model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
- for data, label in data_iter:
- out = model(data)
- loss = criterion(out, label)
- optimizer.zero_grad()
- #loss scaling, 代替 loss.backward()
- with amp.scaled_loss(loss, optimizer) as scaled_loss:
- scaled_loss.backward()
- optimizer.step()
- Tensorflow
一句话实现混合精度训练之 修改环境变量, 在 python 脚本中设置环境变量
os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
除此之外, 也可以用类似 pytorch 的方式来包装 optimizer.
Graph-based 示例
- opt = tf.train.AdamOptimizer()
- #add a line
- opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(
- opt,
- loss_scale='dynamic')
- train_op = opt.miminize(loss)
keras-based 示例
- opt = tf.keras.optimizers.Adam()
- #add a line
- opt = tf.train.experimental.enable_mixed_precision_graph_rewrite(
- opt,
- loss_scale='dynamic')
- model.compile(loss=loss, optimizer=opt)
- model.fit(...)
- PaddlePaddle
一句话实现混合精度训练之 添加 config (惊呆 毕竟混合精度训练是百度家提出的, 内部早就熟练应用了叭 )
--use_fp16=true
举个栗子, 基于 BERT finetune XNLI 任务时, 只需在执行时设置 use_fp16 为 true 即可.
- export FLAGS_sync_nccl_allreduce=0
- export FLAGS_eager_delete_tensor_gb=1
- export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
- BERT_BASE_PATH="chinese_L-12_H-768_A-12"
- TASK_NAME='XNLI'
- DATA_PATH=/path/to/xnli/data/
- CKPT_PATH=/path/to/save/checkpoints/
- python -u run_classifier.py --task_name ${TASK_NAME} \
- --use_fp16=true \ #!!!!!!add a line
- --use_cuda true \
- --do_train true \
- --do_val true \
- --do_test true \
- --batch_size 32 \
- --in_tokens false \
- --init_pretraining_params ${BERT_BASE_PATH}/params \
- --data_dir ${DATA_PATH} \
- --vocab_path ${BERT_BASE_PATH}/vocab.txt \
- --checkpoints ${CKPT_PATH} \
- --save_steps 1000 \
- --weight_decay 0.01 \
- --warmup_proportion 0.1 \
- --validation_steps 100 \
- --epoch 3 \
- --max_seq_len 128 \
- --bert_config_path ${BERT_BASE_PATH}/bert_config.JSON \
- --learning_rate 5e-5 \
- --skip_steps 10 \
- --num_iteration_per_drop_scope 10 \
- --verbose true
来源: http://www.tuicool.com/articles/AbYbmqa