PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解


Vincent
发布于 2024-08-03 / 94 阅读 / 0 评论 /
PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解 Nvidia 在 Volta 架构中引入 Tensor Core 单元,来支持 FP32 和 FP16 混合精度计算。也在 2018 年提出一个 PyTorch 拓展 apex,来支持模型参数自动混合精度训练。自动混合精度

PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解

Nvidia 在 Volta 架构中引入 Tensor Core 单元,来支持 FP32 和 FP16 混合精度计算。也在 2018 年提出一个 PyTorch 拓展 apex,来支持模型参数自动混合精度训练。自动混合精度(Automatic Mixed Precision, AMP)训练,是在训练一个数值精度 FP32 的模型,一部分算子的操作时,数值精度为 FP16,其余算子的操作精度是 FP32,而具体哪些算子用 FP16,哪些用 FP32,不需要用户关心,amp 自动给它们都安排好了。这样在不改变模型、不降低模型训练精度的前提下,可以缩短训练时间,降低存储需求,因而能支持更多的 batch size、更大模型和尺寸更大的输入进行训练。PyTorch 从 1.6 以后(在此之前 OpenMMLab 已经支持混合精度训练,即 Fp16OptimizerHook),开始原生支持 amp,即torch.cuda.amp module。2020 ECCV,英伟达官方做了一个 tutorial 推广 amp。从官方各种文档网页 claim 的结果来看,amp 在分类、检测、图像生成、3D CNNs、LSTM,以及 NLP 中机器翻译、语义识别等应用中,都在没有降低模型训练精度都前提下,加速了模型的训练速度。

本文是对torch.cuda.amp工作机制,和 module 中接口使用方法介绍,以及在算法角度上对 amp 不掉点原因进行分析,最后补充一点对 amp 存储消耗的解释。

1. 混合精度训练机制

torch.cuda.amp 给用户提供了较为方便的混合精度训练机制,“方便”体现在两个方面:

用户不需要手动对模型参数 dtype 转换,amp 会自动为算子选择合适的数值精度

对于反向传播的时候,FP16 的梯度数值溢出的问题,amp 提供了梯度 scaling 操作,而且在优化器更新参数前,会自动对梯度 unscaling,所以,对用于模型优化的超参数不会有任何影响

以上两点,分别是通过使用amp.autocast和amp.GradScaler来实现的。

autocast可以作为 Python 上下文管理器和装饰器来使用,用来指定脚本中某个区域、或者某些函数,按照自动混合精度来运行。混合精度在操作的时候,是先将 FP32 的模型的参数拷贝一份,拷贝的参数转换成 FP16,而 amp 规定了的 FP16 的算子(例如卷积、全连接),对 FP16 的数值进行操作;FP32 的算子(例如涉及 reduction 的算子,BatchNormalize,softmax…),输入和输出是 FP16,计算的精度是 FP32。在反向传播时,依然是混合精度计算,得到数值精度为 FP16 的梯度。最后,由于 GPU 中的 Tensor Core 天然支持 FP16 乘积的结果与 FP32 的累加(Tensor Core math),优化器的操作是利用 FP16 的梯度对 FP32 的参数进行更新。

对于 FP16 不可避免的问题就是:表示的范围较窄,如下图所示,大量非 0 梯度会遇到溢出问题。解决办法是:对梯度乘一个 2**N 的系数,称为 scale factor,把梯度 shift 到 FP16 的表示范围。

GradScaler的工作就是在反向传播前给 loss 乘一个 scale factor,所以之后反向传播得到的梯度都乘了相同的 scale factor。并且为了不影响学习率,在梯度更新前将梯度unscale。总结amp的基本训练流程:

  1. 维护一个 FP32 数值精度模型的副本;

  2. 在每个iteration。

    1. 拷贝并且转换成 FP16 模型;

    2. 前向传播(FP16 的模型参数);

    3. loss 乘 scale factor s;

    4. 反向传播(FP16 的模型参数和参数梯度);

    5. 参数梯度乘 1/s;

    6. 利用 FP16 的梯度更新 FP32 的模型参数。

    7. 但是,这里会有一个问题,scale factor 应该如何选取?选一个常量显然是不合适的,因为 loss 和梯度的数值在变,scale factor 需要跟随 loss 动态变化。健康的 loss 是振荡中下降,因此GradScaler设计的 scale factor 每隔 N 个 iteration 乘一个大于 1 的系数,再 scale loss;并且每次更新前检查溢出问题(检查梯度中有没有inf和nan),如果有,scale factor 乘一个小于 1 的系数并跳过该 iteration 的参数更新环节,如果没有,就正常更新参数。动态更新 scale factor 是 amp 实际操作中的流程。总结 amp 动态 scale factor 的训练流程:

  3. 维护一个 FP32 数值精度模型的副本;

  4. 初始化 s;

  5. 在每个 iteration + a 拷贝并且转换成FP16模型 + b 前向传播(FP16 的模型参数) + c loss 乘 scale factor s + d 反向传播(FP16 的模型参数和参数梯度) + e 检查有没有inf或者nan的参数梯度 + 如果有:降低 s,回到步骤a + f 参数梯度乘 1/s + g 利用 FP16 的梯度更新 FP32 的模型参数。

2. amp模块的API

用户使用混合精度训练基本操作:

#amp依赖Tensor core架构,所以model参数必须是cuda tensor类型

model = Net().cuda()

optimizer = optim.SGD(model.parameters(), ...)

#GradScaler对象用来自动做梯度缩放

scaler = GradScaler()

for epoch in epochs:

    for input, target in data:

        optimizer.zero_grad()

        # 在autocast enable 区域运行forward

        with autocast():

            # model做一个FP16的副本,forward

            output = model(input)

            loss = loss_fn(output, target)

        # 用scaler,scale loss(FP16),backward得到scaled的梯度(FP16)

        scaler.scale(loss).backward()

        # scaler 更新参数,会先自动unscale梯度

        # 如果有nan或inf,自动跳过

        scaler.step(optimizer)

        # scaler factor更新

        scaler.update()

2.1 autocast类

``autocast(enable=True)`` 可以作为上下文管理器和装饰器来使用,给算子自动安排按照 FP16 或者 FP32 的数值精度来操作。

2.1.1 autocast算子

PyTorch中,只有 CUDA 算子有资格被 autocast,而且只有 “out-of-place” 才可以被 autocast,例如:a.addmm(b, c)是可以被 autocast,但是a.addmm_(b, c)和a.addmm(b, c, out=d)不可以 autocast。amp autocast 成 FP16 的算子有:

autocast 成 FP32 的算子:

剩下没有列出的算子,像dot,add,cat…都是按数据中较大的数值精度,进行操作,即有 FP32 参与计算,就按 FP32,全是 FP16 参与计算,就是 FP16。

2.1.2 MisMatch error

作为上下文管理器使用时,混合精度计算 enable 区域得到的 FP16 数值精度的变量在 enable 区域外需要显式的转成 FP32:

# Creates some tensors in default dtype (here assumed to be float32)

a_float32 = torch.rand((8, 8), device="cuda")

b_float32 = torch.rand((8, 8), device="cuda")

c_float32 = torch.rand((8, 8), device="cuda")

d_float32 = torch.rand((8, 8), device="cuda")

with autocast():

    # torch.mm is on autocast's list of ops that should run in float16.

    # eable 区域内 fp32 自动转化成 fp16

    e_float16 = torch.mm(a_float32, b_float32)   

    # Also handles mixed input types

    f_float16 = torch.mm(d_float32, e_float16)

# After exiting autocast, calls f_float16.float() to use with d_float32

g_float32 = torch.mm(d_float32, f_float16.float()) # eable 区域外 fp16 显示转化成 fp32
2.1.3 autocast 嵌套使用
# Creates some tensors in default dtype (here assumed to be float32)

a_float32 = torch.rand((8, 8), device="cuda")

b_float32 = torch.rand((8, 8), device="cuda")

c_float32 = torch.rand((8, 8), device="cuda")

d_float32 = torch.rand((8, 8), device="cuda")

with autocast():

    e_float16 = torch.mm(a_float32, b_float32)

    with autocast(enabled=False): # enabled=False 关闭fp16

        f_float32 = torch.mm(c_float32, e_float16.float()) # 需要使用fp32进行计算

    g_float16 = torch.mm(d_float32, f_float32)
2.1.4 autocast 作为装饰器

这种情况一般用于 data parallel 的模型的,autocast 设计为 “thread local” 的,所以只在 main thread 上设 autocast 区域是不 work 的:

model = MyModel() 

dp_model = nn.DataParallel(model)

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) 

正确姿势是对 forward 装饰:

MyModel(nn.Module):

    ...

    @autocast()

    def forward(self, input):

       ...

另一个正确姿势是在 forward 的里面设 autocast 区域:

MyModel(nn.Module):

    ...

    def forward(self, input):

        with autocast():

            ...

forward 函数处理之后,在 main thread 里 autocast:

model = MyModel()

dp_model = nn.DataParallel(model)

with autocast():

    output = dp_model(input)

    loss = loss_fn(output)

2.1.5 autocast 自定义函数

对于用户自定义的 autograd 函数,需要用amp.custom_fwd装饰 forward 函数,amp.custom_bwd装饰 backward 函数:

class MyMM(torch.autograd.Function):

    @staticmethod

    @custom_fwd

    def forward(ctx, a, b):

        ctx.save_for_backward(a, b)

        return a.mm(b)

    @staticmethod

    @custom_bwd

    def backward(ctx, grad):

        a, b = ctx.saved_tensors

        return grad.mm(b.t()), a.t().mm(grad)

调用时再 autocast

mymm = MyMM.apply

with autocast():

    output = mymm(input1, input2)
2.1.6 源码分析

autocast主要实现接口有:

A. enter

def enter(self):

    self.prev = torch.is_autocast_enabled()

    torch.set_autocast_enabled(self._enabled)

    torch.autocast_increment_nesting()

B. exit

def exit(self, *args):

    if torch.autocast_decrement_nesting() == 0:

        torch.clear_autocast_cache()

    torch.set_autocast_enabled(self.prev)

    return False

C. call

def call(self, func):

    @functools.wraps(func)

    def decorate_autocast(*args, **kwargs):

        with self:

            return func(*args, **kwargs)

    return decorate_autocast

其中torch.autocast函数是在 pytorch/aten/src/ATen/autocast_mode.cpp 里实现。PyTorch ATen 是 A TENsor library for C++11,ATen 部分有大量的代码是来声明和定义 Tensor 运算相关的逻辑的。autocast_mode.cpp 实现策略是 “ cache fp16 casts of fp32 model weights”。

2.2 GradScaler 类torch.cuda.amp.GradScaler(init_scale=65536.0, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True)用于动态 scale 梯度

  1. init_scale: scale factor 的初始值;

  2. growth_factor: 每次 scale factor 的增长系数;

  3. backoff_factor: scale factor 下降系数;

  4. growth_interval: 每隔多个 interval 增长 scale factor;

  5. enabled: 是否做 scale。

2.2.1 scale(output)方法

对outputs乘 scale factor,并返回,如果enabled=False就原样返回。

2.2.3 step(optimizer, args, *kwargs)方法

step 方法在做两件事情:

  1. 对梯度 unscale,如果之前没有手动调用unscale方法的话;

  2. 检查梯度溢出,如果没有nan/inf,就执行 optimizer 的 step,如果有就跳过

注意:GradScaler的step不支持传 closure。

2.2.4 update(new_scale=None)方法

update方法在每个 iteration 结束前都需要调用,如果参数更新跳过,会给 scale factor 乘backoff_factor,或者到了该增长的 iteration,就给 scale factor 乘growth_factor。也可以用new_scale直接更新 scale factor。

2.3 举例

2.3.1 Gradient clipping
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()

        # unscale 梯度,可以不影响clip的threshold

        scaler.unscale_(optimizer)

        # clip梯度

        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)

        # unscale_()已经被显式调用了,scaler正常执行step更新参数,有nan/inf也会跳过

        scaler.step(optimizer)

        scaler.update()
2.3.2 Gradient accumulation
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 = loss / iters_to_accumulate

        # scale 归一的loss 并backward  

        scaler.scale(loss).backward()

        if (i + 1) % iters_to_accumulate == 0:

            # may unscale_ here if desired 

            # (e.g., to allow clipping unscaled gradients)

            # step() and update() proceed as usual.

            scaler.step(optimizer)

            scaler.update()

            optimizer.zero_grad()
2.3.3. Gradient penalty
scaler = GradScaler()

for epoch in epochs:

    for input, target in data:

        optimizer.zero_grad()

        with autocast():

            output = model(input)

            loss = loss_fn(output, target)

        # 防止溢出,在不是autocast 区域,先用scaled loss 得到 scaled 梯度

        scaled_grad_params = torch.autograd.grad(outputs=scaler.scale(loss),

                                                 inputs=model.parameters(),

                                                 create_graph=True)

        # 梯度unscale

        inv_scale = 1./scaler.get_scale()

        grad_params = [p * inv_scale for p in scaled_grad_params]

        # 在autocast 区域,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

        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()
2.3.4. Multiple models
scaler 一个就够,但 scale(loss) 和 step(optimizer) 要分别执行

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)

        # (retain_graph here is unrelated to amp, it's present because in this

        # example, both backward() calls share some sections of graph.)

        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()
2.3.5. Multiple GPUs

torch DDP 和 torch DP model 的处理方式一样