pytorch混合精度训练

爱被打了一巴掌 2022-12-01 05:28 307阅读 0赞

# 1 混合精度 #

计算机中的浮点数表示，按照IEEE754可以分为三种，分别是半精度浮点数、单精度浮点数和双精度浮点数。三种格式的浮点数因占用的存储位数不同，能够表示的数据精度也不同。

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center]Signed bit用于控制浮点数的正负，0表示正数，1表示负数；  
Exponent部分用于控制浮点数的大小，以2为底进行指数运算；  
Significand部分用于控制浮点数的精度，存储浮点数的有效数字。

从上图中可以看出，fp16、fp32和fp64具有不同的表示范围和数值精度，fp16数值范围和精度最低，存储空间也最小。fp64数值范围和精度最高，存储空间也最大。fp32居两者之间。

默认深度学习模型训练过程中都是使用fp32.

# 2 fp16训练的优劣 #

该部分内容引用自：[https://zhuanlan.zhihu.com/p/79887894][https_zhuanlan.zhihu.com_p_79887894]  
使用fp16存储网络的权重值、激活值和梯度值进行网络训练，好处是：

1.  减少显存占用：上面的图已经很明显的可以看出，fp16的存储空间为fp32的一半，如果使用fp16进行训练，那么可以减少一半的显存占用，因此也就可以使用更大的batchsize进行大模型的训练；
2.  加快训练和推理速度：fp16可以提高模型的训练和推理的速度。

**完全使用fp16进行训练的缺点**：

1.  数值溢出：fp16的动态范围比fp32的动态范围小很多，因此在计算过程中很容易出现数值上溢出和下溢出的问题，溢出之后就会出现“Nan”的问题。在深度学习中，由于激活值的梯度往往比权重的梯度小，更容易出现数值下溢的情况。一旦出现了数值下溢，就会造成发生数值下溢的层及其前面层的权重无法更新，训练失败；
2.  舍入误差：舍入误差是指梯度过小，小于当前区间内的最小间隔时，该次梯度更新可能会失败，如下图所示：  
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 1] 某次梯度更新值小于了当前权重区间的最小间隔时，更新无效。

![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 2]

# 3 混合精度训练的实现思路 #

混合精度训练的思路来自于nvidia和百度合作的论文：[https://arxiv.org/pdf/1710.03740.pdf][https_arxiv.org_pdf_1710.03740.pdf]

下面部分的内容引用自：[https://zhuanlan.zhihu.com/p/103685761][https_zhuanlan.zhihu.com_p_103685761]

1.  FP32权重备份  
    实现思路为，网络前向计算过程中，权重、激活值和梯度都使用fp16进行存储，同时保留一份fp32格式的权重的副本用于进行参数更新。  
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 3] 进行权重更新时，需要对权重的梯度乘以学习率，而学习率一般都是一个远小于1的数字，因此权重的更新值就更小。如果使用fp16进行权重更新，那么更容易引入舍入误差的问题造成参数无效更新。而这里将参数的更新量应用于fp32精度的权重上，因为fp32的数值精度远远大于fp16，也就避免了更新过程中的舍入误差问题。
    
    看到这里，可能有人提出这种 fp32 拷贝weight的方式，那岂不是使得内存占用反而更高了呢？是的， fp32 额外拷贝一份 weight 的确新增加了训练时候存储的占用。 但是实际上，在训练过程中，内存中占据大部分的基本都是 activations 的值。特别是在batchsize 很大的情况下， activations 更是特别占据空间。 保存 activiations 主要是为了在 back-propogation 的时候进行计算。因此，只要 activation 的值基本都是使用 fp16 来进行存储的话，则最终模型与 fp32 相比起来， 内存占用也基本能够减半。
2.  Loss Scale  
    Loss Scale 主要是为了解决 fp16 underflow 的问题。刚才提到，训练到了后期，梯度（特别是激活函数平滑段的梯度）会特别小，fp16 表示容易产生 underflow 现象。 下图展示了 SSD 模型在训练过程中，激活函数梯度的分布情况：可以看到，有67%的梯度小于 2 − 24 2^\{-24\} 2−24，如果用 fp16 来表示，则这些梯度都会变成0。  
    ![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 4] 为了解决梯度过小的问题，论文中对计算出来的loss值进行scale，由于链式法则的存在，loss上的scale会作用也会作用在梯度上。这样比起对每个梯度进行scale更加划算。 scaled 过后的梯度，就会平移到 fp16 有效的展示范围内。
    
    这样，scaled-gradient 就可以一直使用 fp16 进行存储了。只有在进行更新的时候，才会将 scaled-gradient 转化为 fp32，同时将scale抹去。论文指出， scale 并非对于所有网络而言都是必须的。而scale的取值为也会特别大，论文给出在 8 - 32k 之间皆可。
3.  提高计算精度
    
    在论文中还提到一个『计算精度』的问题：在某些模型中，fp16矩阵乘法的过程中，需要利用 fp32 来进行矩阵乘法中间的累加(accumulated)，然后再将 fp32 的值转化为 fp16 进行存储。 换句不太严谨的话来说，也就是利用 **利用fp16进行乘法和存储，利用fp32来进行加法计算**。 这么做的原因主要是**为了减少加法过程中的舍入误差，保证精度不损失**。
    
    在这里也就引出了，为什么网上大家都说，只有 Nvidia Volta 结构的 拥有 TensorCore 的CPU(例如V100)，才能利用 fp16 混合精度来进行加速。 那是因为 TensorCore 能够保证 fp16 的矩阵相乘，利用 fp16 or fp32 来进行累加。在累加阶段能够使用 FP32 大幅减少混合精度训练的精度损失。而其他的GPU 只能支持 fp16 的 multiply-add operation。这里直接贴出原文句子：
    
    Whereas previous GPUs supported only FP16 multiply-add operation, NVIDIA Volta GPUs introduce Tensor Cores that multiply FP16 input matrices and accumulate products into either FP16 or FP32 outputs。

# 4 实现代码 #

## 4.1 nvidia版本 ##

nvidia官方github：[https://github.com/NVIDIA/apex][https_github.com_NVIDIA_apex]

在pytorch中应用：

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()

参数 opt\_level：

O0：纯FP32训练，可以作为accuracy的baseline；
    O1：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（GEMM, 卷积）还是FP32（Softmax）进行计算；
    O2：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算；
    O3：纯FP16训练，很不稳定，但是可以作为speed的baseline。

## 4.2 pytorch官方版本 ##

pytorch1.6开始支持amp，位于torch.cuda.amp模块下，具体例子见[https://pytorch.org/docs/stable/notes/amp\_examples.html][https_pytorch.org_docs_stable_notes_amp_examples.html]

包括 torch.cuda.amp.autocast和torch.cuda.amp.GradScaler两个模块，autocast针对选定的代码块自动选取适合的计算精度，以便在保持模型准确率的情况下最大化改善训练效率；GradScaler即进行梯度缩放，以最大程度避免使用fp16进行运算时的梯度下溢。

### 4.2.1 **示例代码** ###

# Creates model and optimizer in default precision
    model = Net().cuda()
    optimizer = optim.SGD(model.parameters(), ...)
    
    # Creates a GradScaler once at the beginning of training.
    scaler = GradScaler()
    
    for epoch in epochs:
        for input, target in data:
            optimizer.zero_grad()
    
            # Runs the forward pass with autocasting.
            with autocast():
                output = model(input)
                loss = loss_fn(output, target)
    
            # Scales loss.  Calls backward() on scaled loss to create scaled gradients.
            # Backward passes under autocast are not recommended.
            # Backward ops run in the same dtype autocast chose for corresponding forward ops.
            scaler.scale(loss).backward()
    
            # scaler.step() first unscales the gradients of the optimizer's assigned params.
            # If these gradients do not contain infs or NaNs, optimizer.step() is then called,
            # otherwise, optimizer.step() is skipped.
            scaler.step(optimizer)
    
            # Updates the scale for next iteration.
            scaler.update()

backward不要放到autocast作用域下，反向传播的时候梯度和前向时对应的op采用相同的计算精度。

### 4.2.2 梯度反缩放 ###

由scaler.scale(loss).backward()计算的梯度都是缩放后的梯度，如果要使用梯度值，首先应该进行反缩放操作。如进行梯度裁剪时，首先应该对计算的梯度进行反缩放再和给定的阈值进行比较，当然也可以对阈值按照同样的系数进行缩放而不对计算的梯度反缩放，效果是一致的。scaler.unscale\_(optimizer)对optimizer指定的参数的梯度进行反缩放操作。如果训练过程中使用了多个optimizer，则可以对每个optimizer进行单独反缩放，如scaler.unscale\_(optimizer2) 。

**梯度裁剪**：

scaler = GradScaler()
    
    for epoch in epochs:
        for input, target in data:
            optimizer.zero_grad()
            with autocast():
                output = model(input)
                loss = loss_fn(output, target)
            scaler.scale(loss).backward()
    
            # Unscales the gradients of optimizer's assigned params in-place
            scaler.unscale_(optimizer)
    
            # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
    
            # optimizer's gradients are already unscaled, so scaler.step does not unscale them,
            # although it still skips optimizer.step() if the gradients contain infs or NaNs.
            scaler.step(optimizer)
    
            # Updates the scale for next iteration.
            scaler.update()

在梯度裁剪之前进行反缩放。每次step中，只能在梯度计算完成后调用一次unscale\_，在一次step中对一个optimizer调用多次unscale\_会引发异常。

### 4.2.3 使用缩放后的梯度 ###

#### 4.2.3.1 梯度累加 ####

梯度累加是模拟按照batch\_per\_iter \* iters\_to\_accumulate大的batchsize进行梯度计算，当然如果是分布式训练的话，还需要乘以num\_procs。如果计算得到的梯度中出现了inf/NaN，针对这个梯度的step过程会被跳过。进行梯度累加时，仍然应该使用缩放后的梯度，并且是按照有效batchsize的粒度进行梯度累加。如果在累加计算过程中，使用了未缩放的梯度或者是梯度的缩放系数发生了改变，将会把缩放后的梯度累加到未缩放的梯度上(或者是对缩放系数不同的梯度进行了累加)，那么之后将很难对未缩放的梯度应用step操作。所以说，如果想要使用未缩放的梯度，就需要在step函数之前调用unscale\_操作。如下面的代码所示：

scaler = GradScaler()
    
    for epoch in epochs:
        for i, (input, target) in enumerate(data):
            with autocast():
                output = model(input)
                loss = loss_fn(output, target)
                loss = loss / iters_to_accumulate
    
            # Accumulates scaled gradients.
            scaler.scale(loss).backward()
    
            if (i + 1) % iters_to_accumulate == 0:
                # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)
    
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad()

#### 4.2.3.2 梯度正则化 ####

使用torch.autograd.grad()实现梯度正则化，将其加到损失值上。

下面是使用未缩放的梯度和autocasting进行L2梯度正则化的示例代码：

for epoch in epochs:
        for input, target in data:
            optimizer.zero_grad()
            output = model(input)
            loss = loss_fn(output, target)
    
            # Creates gradients
            grad_params = torch.autograd.grad(loss, model.parameters(), create_graph=True)
    
            # Computes the penalty term and adds it to the loss
            grad_norm = 0
            for grad in grad_params:
                grad_norm += grad.pow(2).sum()
            grad_norm = grad_norm.sqrt()
            loss = loss + grad_norm
    
            loss.backward()
    
            # clip gradients here, if desired
    
            optimizer.step()

对于缩放后的梯度，torch.autograd.grad()计算的也是缩放后的梯度，那么在计算梯度正则化项之前首先应该进行反缩放处理。正则化项的计算也属于前向过程，应该放到autocast的作用域中。因此使用缩放后的梯度和autocasting进行L2梯度正则化的示例代码为：

scaler = GradScaler()
    
    for epoch in epochs:
        for input, target in data:
            optimizer.zero_grad()
            with autocast():
                output = model(input)
                loss = loss_fn(output, target)
    
            # Scales the loss for autograd.grad's backward pass, resulting in scaled grad_params
            scaled_grad_params = torch.autograd.grad(scaler.scale(loss), model.parameters(), create_graph=True)
    
            # Creates unscaled grad_params before computing the penalty. scaled_grad_params are
            # not owned by any optimizer, so ordinary division is used instead of scaler.unscale_:
            inv_scale = 1./scaler.get_scale()
            grad_params = [p * inv_scale for p in scaled_grad_params]
    
            # Computes the penalty term and adds it to the loss
            with autocast():
                grad_norm = 0
                for grad in grad_params:
                    grad_norm += grad.pow(2).sum()
                grad_norm = grad_norm.sqrt()
                loss = loss + grad_norm
    
            # Applies scaling to the backward call as usual.
            # Accumulates leaf gradients that are correctly scaled.
            scaler.scale(loss).backward()
    
            # may unscale_ here if desired (e.g., to allow clipping unscaled gradients)
    
            # step() and update() proceed as usual.
            scaler.step(optimizer)
            scaler.update()

### 4.2.4 使用多个模型、损失和优化器 ###

如果网络有多个损失，对每一个损失单独应用scaler.scale；如果网络有多个优化器，对每一个单独应用scaler.unscale\_，同时对每一个优化器单独应用scaler.step函数。但scaler.update()函数只使用一次。

scaler = torch.cuda.amp.GradScaler()
    
    for epoch in epochs:
        for input, target in data:
            optimizer0.zero_grad()
            optimizer1.zero_grad()
            with autocast():
                output0 = model0(input)
                output1 = model1(input)
                loss0 = loss_fn(2 * output0 + 3 * output1, target)
                loss1 = loss_fn(3 * output0 - 5 * output1, target)
    
            scaler.scale(loss0).backward(retain_graph=True)
            scaler.scale(loss1).backward()
    
            # You can choose which optimizers receive explicit unscaling, if you
            # want to inspect or modify the gradients of the params they own.
            scaler.unscale_(optimizer0)
    
            scaler.step(optimizer0)
            scaler.step(optimizer1)
    
            scaler.update()

每一个优化器单独进行自身的梯度检查以决定是否跳过step过程，有可能出现在一次迭代中，某个优化器跳过了step操作而另一个优化器却没有。但由于跳过step过程很少发生，所以这样的设计基本不会阻碍训练的收敛过程。

### 4.2.5 使用多个GPU ###

使用多个GPU时不影响GradScaler，只影响autocast。

#### 4.2.5.1 单个进程中的数据并行 ####

torch.nn.DataParallel是在每个设备上新建线程进行前向计算，autocast只作用于线程内部，所以下面这种写法不能正常的工作(只有主线程进行了autocast，其余线程未应用autocast)：

model = MyModel()
    dp_model = nn.DataParallel(model)
    
    # Sets autocast in the main thread
    with autocast():
        # dp_model's internal threads won't autocast.  The main thread's autocast state has no effect.
        output = dp_model(input)
        # loss_fn still autocasts, but it's too late...
        loss = loss_fn(output)

修改代码的思路是在model的前向过程中应用autocast，如下所示：

MyModel(nn.Module):
        ...
        @autocast()
        def forward(self, input):
           ...
    
    # Alternatively
    MyModel(nn.Module):
        ...
        def forward(self, input):
            with autocast():
                ...

下面的代码在各dp\_model的线程中应用autocast，在主线程中执行loss函数：

model = MyModel()
    dp_model = nn.DataParallel(model)
    
    with autocast():
        output = dp_model(input)
        loss = loss_fn(output)

#### 4.2.5.2 分布式数据并行，一个进程占用一个GPU ####

分布式数据并行中建议一个进程占用一个GPU已取得最佳性能，这种情况下，分布式数据并行没有在内部新建线程，所以对autocast和GradScaler的应用没有影响。

#### 4.2.5.3 分布式数据并行，一个进程占用多个GPU ####

分布式数据并行中有可能会在每一个设备上创建一个副线程进行前向过程，这种情况下的代码修正如4.2.5.1所示，在模型的forward函数中使用autocast。

# 5 他人经验 #

引用自：[https://zhuanlan.zhihu.com/p/79887894][https_zhuanlan.zhihu.com_p_79887894]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 5]  
引用自：[https://zhuanlan.zhihu.com/p/103685761][https_zhuanlan.zhihu.com_p_103685761]  
![在这里插入图片描述][watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 6]  
参考：[https://pytorch.org/docs/stable/notes/amp\_examples.html][https_pytorch.org_docs_stable_notes_amp_examples.html]  
[https://zhuanlan.zhihu.com/p/103685761][https_zhuanlan.zhihu.com_p_103685761]  
[https://zhuanlan.zhihu.com/p/79887894][https_zhuanlan.zhihu.com_p_79887894]  
[http://nvidia.zhidx.com/content-6-1651-1.html][http_nvidia.zhidx.com_content-6-1651-1.html]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center]: /images/20221123/9f989fc3a952458c811ff9d7079d4065.png
[https_zhuanlan.zhihu.com_p_79887894]: https://zhuanlan.zhihu.com/p/79887894
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 1]: /images/20221123/6c71a50e8ecd49d6a4573fea3792db23.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 2]: /images/20221123/cd0e731d8137482d8ae15867c95a55a0.png
[https_arxiv.org_pdf_1710.03740.pdf]: https://arxiv.org/pdf/1710.03740.pdf
[https_zhuanlan.zhihu.com_p_103685761]: https://zhuanlan.zhihu.com/p/103685761
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 3]: /images/20221123/966f3c2f0f4145e4bfec54891e18236b.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 4]: /images/20221123/f40684a67e114d428051ddacc4800579.png
[https_github.com_NVIDIA_apex]: https://github.com/NVIDIA/apex
[https_pytorch.org_docs_stable_notes_amp_examples.html]: https://pytorch.org/docs/stable/notes/amp_examples.html
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 5]: /images/20221123/899a42dfbc3b444b9f97529ef52d890b.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2Nka25pZ2h0X2hhcHB5_size_16_color_FFFFFF_t_70_pic_center 6]: /images/20221123/462caa2a43a64ed997a195a5b97420c3.png
[http_nvidia.zhidx.com_content-6-1651-1.html]: http://nvidia.zhidx.com/content-6-1651-1.html