FP16和FP32以及基于Apex的混合精度加速

1. FP16

FP16也称为半精度浮点数是一种计算机使用的二进制浮点数数据类型，使用 2 字节（16 位）存储

图1. FP16和FP32表示范围对比

FP16的表示范围($6\ast 10^{-8}\to 65504$)，FP32表示范围($1.4\ast 10^{-45}\to 1.7\ast 10^{38}$)

FP16的作用

• 减少显存占用，FP16 的内存占用只有 FP32 的一半，自然地就可以帮助训练过程节省一半的显存空间
• 加快训练和推断的计算，在大部分的测试中，基于 FP16 的加速方法能够给模型训练带来多一倍的加速体验

FP16的问题
FP16 带来的问题主要有两个：溢出错误和舍入误差

• 溢出错误（Grad Overflow / Underflow）是由于 FP16 的动态范围、比 FP32 的动态范围要狭窄很多，因此在计算过程中很容易出现上溢出（Overflow，g>65504）和下溢出（Underflow）的错误，溢出之后就会出现“Nan”的问题
• 舍入误差（Rounding Error）舍入误差指的是当梯度过小，小于当前区间内的最小间隔时，该次梯度更新可能会失败
  

2. 混合精度

为了解决FP16的问题，提出了混合精度训练+动态损失放大的方法

• 混合精度训练（Mixed Precision）混合精度训练的精髓在于“在内存中用 FP16 做储存和乘法从而加速计算，用 FP32 做累加避免舍入误差”。混合精度训练的策略有效地缓解了舍入误差的问题
• 损失放大（Loss Scaling）即使用了混合精度训练，还是会存在无法收敛的情况，原因是激活梯度的值太小，造成了下溢出（Underflow）。损失放大的思路是：（1）反向传播前，将损失变化手动增大$2^k$倍，因此反向传播时得到的中间变量（激活函数梯度）则不会溢出（2）反向传播后，将权重梯度缩小$2^k$倍，恢复正常值

3. 混合精度api使用

# 1. 导入包
from apex import amp
# 2. 模型和优化器的初始化
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”
# 3. 损失函数的反向传播
with amp.scale_loss(loss, optimizer) as scaled_loss:
	scaled_loss.backward()
# 4. 梯度裁剪（如果有）
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)

模型和优化器初始化中 opt_level 参数说明：总共可选[O0, O1,O2, O3]
O0 ：纯 FP32 训练，可以作为 accuracy 的 baseline
O1 ：混合精度训练（推荐使用），根据黑白名单自动决定使用 FP16（GEMM, 卷积）还是 FP32（Softmax）进行计算
O2 ：“几乎 FP16”混合精度训练，不存在黑白名单，除了 Batch Norm，几乎都是用 FP16 计算
O3 ：纯 FP16 训练，很不稳定，但是可以作为 speed 的 baseline

使用例子

def train(trainset, evalset, model, tokenizer, model_dir, lr, epochs, device):
    optimizer = AdamW(model.parameters(), lr=lr)
    model, optimizer = amp.initialize(model, optimizer, opt_level="O3")
    for epoch in tqdm(range(epochs), desc="epoch"):
        train_loss, steps = 0, 0
        for batch in tqdm(trainset, desc="train"):
            batch = tuple(input_tensor.to(device) for input_tensor in batch if isinstance(input_tensor, torch.Tensor))
            input_ids, label, mc_ids = batch
            steps += 1
            model.train()
            loss, logits = model(input_ids=input_ids, mc_token_ids=mc_ids, labels=label)
            # loss.backward()
            with amp.scale_loss(loss, optimizer) as scaled_loss:
                scaled_loss.backward()
            train_loss += loss.item()
            torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), 5)
            # torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
            optimizer.step()
            optimizer.zero_grad()
            if steps % 500 == 0:
                print("step:%d  avg_loss:%.3f"%(steps, train_loss/steps))
        eval_res = evaluate(evalset, model, device)
        os.makedirs(model_dir, exist_ok=True)
        model_path = os.path.join(model_dir, "gpt2clsnews.model%d.ckpt"%epoch)
        model.save_pretrained(model_path)
        tokenizer.save_pretrained(os.path.join(model_dir,"gpt2clsnews.tokinizer"))
        logging.info("checkpoint saved in %s"%model_dir)

结果对比

	不使用apm	O0	O1	O3
训练耗时	3m48s	3m45s	3m18s	1m56s
预测耗时	1m12s	1m12s	50s	34s
accuracy	0.859	0.849	0.875	0.004
显存占用	13.76G	13.76G	12.05G	7.79G

上面O3情况下loss 为nan， 应该是出错了，从上表中可以看出使用混合精度O1可以节省部分时间和显存，其中accuracy并不相关，最开始使用了1000个样本测时不使用apm的accuracy比O1高，后来表中换成3000个样本时，O1比不使用apm高

4. 参考

PyTorch必备神器 | 唯快不破：基于Apex的混合精度加速