(Alpha) FP16 训练

动机、参考资料、涉及内容

半精度训练 (full fp16)

本节讲的是 Pytorch 1.6.0 之前的纯半精度训练, 注意与 Pytorch 现在原生支持的混合精度训练做区分, 为此需要稍稍将时间倒回去看, 本节主要参考 mmdetection 与 mmcv 旧版本的实现, 版本号如下:

• pytorch 1.6.0 (2020.7.29): 稳定版的自动混合精度训练
• mmdetection 2.4.0 (2020.9.10)
• mmcv 1.3.14 (2021.5.14)

依赖关系参考: https://github.com/open-mmlab/mmdetection/blob/v2.11.0/docs/get_started.md

mmcv/mmdetection 的实现来源于 2017/10 的混合精度训练原始论文: Mixed Precision Training, 其实现本质上如下:

参考:

• mmdet/core/fp16/hooks.py:Fp16OptimizerHook: 主体部分
• mmdet/models/detectors/base.py:BaseDetector.forward_train: forward 函数带着装饰器 auto_fp16
• /mmcv/runner/fp16_utils.py:auto_fp16: auto_fp16 的具体实现

# 这里注意一个细节: 我们通常的写法里, optimizer 中会保存 model 中权重的引用, 但这里不是
optimizer.param_groups = copy.deepcopy(optimizer.param_groups)  # 优化器中的参数全部是 fp32
model.half()  # 模型参数全部转为 fp16
for item in dataloader:
    # loss 为 fp16
    loss = model(item.half())
    # 清除 fp16 权重的梯度
    model.zero_grad()
    optimizer.zero_grad()
    scaled_loss = loss * loss_scale  # loss_scale 一般为 32
    scaled_loss.backward()  # 执行完毕后 model 中的 fp16 参数有着 fp16 的 grad
    fp32_weights = []
    for param_group in optimizer.param_groups:
        fp32_weights += param_group['params']

    # 通过直接赋值的方式将 fp16 的梯度转换为 fp32 的梯度
    for fp32_param, fp16_param in zip(fp32_weights, model.parameters()):
        if fp16_param.grad is not None:
            if fp32_param.grad is None:
                fp32_param.grad = fp32_param.data.new(fp32_param.size())
            fp32_param.grad.copy_(fp16_param.grad)

    # 手动对fp32梯度除以放缩因子
    for param in fp32_weights:
        if param.grad is not None:
            param.grad.div_(self.loss_scale)
    
    # 更新 fp32 参数
    runner.optimizer.step()

    # 模型参数更新为将优化器中的 fp32 参数转为 fp16
    for fp16_param, fp32_param in zip(model.parameters(), fp32_weights):
        fp16_param.data.copy_(fp32_param.data)

半精度训练内存分析

以 Adam 优化器为例, 假设模型参数量为 P, 自动微分需要保留的中间结果的参数量为 N.

使用 fp32 训练, 显存主要用在了这几处:

• 模型参数: 4*P byte
• 优化器状态: 8*P byte
• 梯度: 4*P byte
• 中间结果: 4*N byte

使用 fp16 训练, 显存主要用在了这几处:

• fp16 模型参数: 2*P byte
• fp32 模型参数 (对应于前一节中的 optimizer.param_groups): 4*P byte
• fp32 优化器状态: 8*P byte
• fp16 梯度 (注意更新 fp32 权重时将 fp16 梯度临时转为 fp32, 不需要存储 fp32 梯度): 2*P byte
• fp16 中间结果: 2*N byte

所以, fp32 训练需要的显存是: 16*P+4*N, 而 fp16 训练需要的显存是: 16*P+2*N, 所以显存的节省上并不多 (N 与 P 的通常倍数关系是怎样的?).

中间结果及其他显存消耗

参考 Zero 论文 中 3.2 节的分析, 除去模型和优化器外, 还需要考虑以下三种显存消耗, 以 GPT-2 (1.5B参数) 为例

• 模型和优化器: 24 GB (1.5*16)
• Activations (中间层结果): 假设 (使用fp16计算) 序列长度为 1000, batchsize 为 32, 这一部分需要消耗 60 GB 显存. 使用 activation checkpointing (activation recomputation) 技术能在增加 33% 额外计算开销的情况下将 Activations 显存消耗减少至 8GB.
• Temporary buffers (临时变量): 论文指出 GPT-2 需要大约 6GB 的临时变量的显存消耗
• Memory Fragmentation (内存碎片): 论文指出在训练大模型(不确定多大)的时候可能会出现实际上还有 30% 的剩余内存, 但因为内存碎片的原因报 OOM 的错误

参考 Zero 论文及 LOMO, Transformer 类型的 Activations 占用基本上是 num_layer*hidden_dim*batch*seq_length 的倍数, 对于 GPT-2 来说, 这个倍数约为 12

num_layer * hidden_dim * batch * seq_length * 12

48*1600*32*1000*12 /1024/1024/1024 * 2byte = 54.9 GB

这里有个疑问: 精确值怎么计算, 12是个约数

混合精度训练 (mixed fp16)

混合精度训练只能在 GPU 上进行, 因为底层是使用 Nvidia 为自家 GPU 提供的 Float16 高效数值运算能力. 平时使用一般只需要用 torch.cuda.amp.GradScaler 以及 torch.cuda.amp.autocast 即可, 并且可以设置 enabled 参数, 当它为 True 时, 则启用 amp 训练, 否则等价于通常的训练方式. 实际体验上: amp训练计算速度与内存消耗未必快…

原理

原理与半精度训练基本一致, 但做了几个额外补充:

• 损失的放缩系数自适应调整
• 如果出现因为 fp16 计算造成 NaN 的情况, 则放弃此次更新, 并调整放缩系数
• Pytorch 内置了一个“黑白名单”, 即在 torch.cuda.amp.autocast 的作用域下, 某些算子会用 fp16 计算, 某些会按 fp32 计算

使用

参考: torch.cuda.amp

use_amp = True
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
scaler = torch.cuda.amp.GradScaler(enabled=use_amp)

for epoch in range(epochs):
    for input, target in zip(data, targets):
        with torch.cuda.amp.autocast(enabled=use_amp):
            output = net(input)
            loss = loss_fn(output, target)
        scaler.scale(loss).backward()
        # scaler.unscale_(optimizer)
        # torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

其中第11行及第12行只有在需要对梯度做修改时才需要做, 此时需要对这两行同时取消注释.

• scaler.scale(loss).backward()可以简单地理解为: (loss*scaler.get_scale()).backward(), 所以造成的结果是 model.parameters() 中每个参数的 grad 属性被放大了一个倍数 (所有参数共用一个倍数)
• scaler.step(optimizer)可以简单理解为, 首先缩小每个参数的 grad 属性 (原地修改, 并且实际上就是调用了 unscale_ 方法), 之后调用 optimizer.step()
• scaler.update(): 大概是更新放缩比例, 必要性参照下节

源码分析

可以通过下面的例子证实上一部分的说明

scaler = torch.cuda.amp.GradScaler()
x = torch.tensor([1., 2.], requires_grad=True, device="cuda:0")
z = torch.tensor([2., 3.], requires_grad=True, device="cuda:0")
y = torch.sum(x) + torch.sum(z)
opt = torch.optim.SGD([x, z], lr=0.001)

scaler.scale(y).backward()
print(x.grad, y)  # [65536, 65536], 8
# opt.step() # 此处若直接用opt调用step, 会利用65536倍的梯度更新, 并且x.grad依然为[65536, 65536], 这种操作会引发错误(相当于学习率被放大), 要避免
scaler.unscale_(opt)
print(x.grad, scaler.get_scale()) # [1, 1], 65535
# torch.nn.utils.clip_grad_norm_([x, z], 1.)  # x.grad = [0.5, 0.5]
torch.nn.utils.clip_grad_norm_(x, 1.)  # 如果将x改为[x, z], 则会将x与z合并起来将梯度剪裁
torch.nn.utils.clip_grad_norm_(z, 1.)
print(x.grad, scaler.get_scale())  # [0.7, 0.7]
scaler.step(opt)
print(x)  # [0.9993, 1.9993]
scaler.update()

为了解答这个疑问: 在实现上, 是否取消前一节的两行注释, 代码都能正常工作, 这是怎么做到的呢?

通过阅读源码, 发现 GradScaler 内部用一个参数记录了当前的状态, 如下

class OptState(Enum):
    READY = 0  # __init__以及update函数会将状态更新为READY
    UNSCALED = 1  # 已经调用了unscale_函数
    STEPPED = 2  # 已经调用了step函数

当手动调用unscale_函数后, 状态会被更新为UNSCALED, 而在执行step函数时, 如果发现状态为READY, 则先调用unscale_, 由此做到自动性(疑问解答). 另外, unscale_函数会首先检查当前状态, 如果是UNSCALED或者STEPPED直接报错, 因此每次调用step后必须使用update才能使用unscale_

__init__:    -> READY
update:      READY/UNSCALED/STEPPED -> READY
unscale_:    READY -> UNSCALED
step:        READY/UNSCALED -> STEPPED

总结: update, unscale_, step函数的顺序不能乱

with autocast():
	output = model(input)  # model, input: all float32
    loss = loss_fn(output, target)  # maybe float32 or float16
scaler.scale(loss).backward()  # grad always float32
# unscale_调用torch._amp_non_finite_check_and_unscale_
scaler.unscale_(optimizer)  # 缩小倍数, 并记录是否存在inf/nan的情况
scaler.step(optimizer)  # 若上一步记录发现inf/nan, 则跳过step
scaler.update()  # 根据inf/nan的情况以及迭代次数来更新_scale