训练大型网络模型的技巧

vlambda
2020-04-02

训练大型网络模型的技巧

本文主要介绍模型训练中速度和内存的优化策略，针对以下几种情况：

我明天就要答辩了，今天必须把这十个实验跑完
我的模型有些大，好不容易放到一张卡上，训完一亿样本之前我就可以领N+1了
我想出了一个绝妙的T6模型，却加载不进12GB的卡里，又拿不到今年的best paper了

（以上纯属虚构，如有雷同请赶紧看下文）

现实总是残酷的，其实限制大模型训练只有两个因素：时间和空间(=GPU=钱)，根据不同情况可以使用的方案大致如下：

一

梯度累加 Gradient Accumulation

如果只有单卡，且可以加载模型，但batch受限的话可以使用梯度累加，进行N次前向后反向更新一次参数，相当于扩大了N倍的batch size。

正常的训练代码是这样的：

 
   
   
 
  
    
    
   
    
    
    
    
    
  
  
    
    
  for i, (inputs, labels) in enumerate(training_set):  loss = model(inputs, labels)              # 计算loss  optimizer.zero_grad()                      # 清空梯度  loss.backward()                           # 反向计算梯度  optimizer.step()                          # 更新参数

加入梯度累加后：

 
   
   
 
  
    
    
   
    
    
    
    
    
    
    
  
  
    
    
  for i, (inputs, labels) in enumerate(training_set):  loss = model(inputs, labels)                    # 计算loss  loss = loss / accumulation_steps                # Normalize our loss (if averaged)  loss.backward()                                 # 反向计算梯度，累加到之前梯度上  if (i+1) % accumulation_steps == 0:      optimizer.step()                            # 更新参数      model.zero_grad()                           # 清空梯度

要注意的是，batch扩大后，如果想保持样本权重相等，学习率也要线性扩大或者适当调整。另外batchnorm也会受到影响，小batch下的均值和方差肯定不如大batch的精准，可以调整BN中的momentum参数解决[2]。

二

梯度检查点 Gradient Checkpointing

如果只有一张卡，又想训大模型，可以尝试压缩模型所占显存。

梯度检查点是一种以时间换空间的方法，通过减少保存的激活值压缩模型占用空间，但是在计算梯度时必须从新计算没有存储的激活值。

细节可以参考陈天奇的Training Deep Nets with Sublinear Memory Cost[3]，最近的Reformer（参考RevNet）也是这种思想[4]。

注：第一行节点是前向，第二行是反向

三

混合精度训练 Mixed Precision Training

混合精度训练在单卡和多卡情况下都可以使用，通过cuda计算中的half2类型提升运算效率。一个half2类型中会存储两个FP16的浮点数，在进行基本运算时可以同时进行，因此FP16的期望速度是FP32的两倍。举个Gelu的FP16优化栗子：

 
   
   
 
  
    
    
   
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
  
  
    
    
  //FP32的gelu运算float gelu(float x){  float cdf = 0.5f * (1.0f + tanhf((0.7978845608028654f * (x + 0.044715f * x * x * x))));  return x * cdf;}//FP16的gelu运算half2 gelu(half2 val){  half2 val_pow3 = __hmul2(val, __hmul2(val, val)); //同时计算两个x*x*x  float2 tmp_pow = __half22float2(val_pow3);  float2 cdf =  __half22float2(val);  //由于tanhf不支持half2类型，只能分开算  cdf.x = 0.5f * (1.0f + tanhf((0.7978845608028654f * (cdf.x + 0.044715f * tmp_pow.x))));  cdf.y = 0.5f * (1.0f + tanhf((0.7978845608028654f * (cdf.y + 0.044715f * tmp_pow.y))));  //同时计算两个x * cdf;return __hmul2(val, __float22half2_rn(cdf));}

混合精度训练[5]不是很难理解，但要注意以下几点：

混合精度训练不是单纯地把FP32转成FP16去计算就可以了，只用FP16会造成80%的精度损失
Loss scaling：由于梯度值都很小，用FP16会下溢，因此先用FP32存储loss并放大，使得梯度也得到放大，可以用FP16存储，更新时变成FP32再缩放
在涉及到累加操作时，比如BatchNorm、Softmax，FP16会上溢，需要用FP32保存，一般使用GPU中TensorCore的FP16*FP16+FP32=FP32运算

整体流程：FP32权重 -> FP16权重 -> FP16计算前向 -> FP32的loss，扩大 -> 转为FP16 -> FP16反向计算梯度 -> 缩放为FP32的梯度更新权重

vlambda博客
学习文章列表