FP8 低精度训练：Transformer Engine 简析

原文：https://zhuanlan.zhihu.com/p/700874387

一、背景介绍

业界广泛采用 FP16、BF16 混合精度（AMP）进行模型训练。AMP 能在下游任务不掉点的前提下提升训练效率、减少显存等资源占用，如今也常用于大模型预训练、微调等任务。

• Pytorch 1.6 版本后原生支持 FP16、BF16 精度的 AMP 训练（torch.amp），过往 AMP 功能由 NVIDIA APEX 库实现。

NVIDIA GPU 自 Hopper 架构起支持 FP8 精度的 Tensor Core 计算，相比于 FP16/BF16 精度，FP8 具有如下优势：

• 更强的计算性能

• 对比 A100 BF16 精度训练，H100 FP8 训练速度提升 2-3x。

• 对比 FP16/BF16，FP8 的计算吞吐提升至 2x，与 A100 相比提升的吞吐更多。

• 更低的训练成本 ：FP8 能提供 2x 的计算速度提升，节省 50%-75% 内存占用，以及节省 50%-75% 的数据通信量。

• 更好的模型优化 ：FP8 的使用促使模型在训练和推理过程中进行量化，这有助于模型的优化和压缩，进一步降低部署成本。

我们仍然先从几个问题出发～

什么是 FP8？

相比于 16bit 精度，FP8 使用了更少的指数 bit 位和尾数 bit 位：

在 NV、Arm、Intel 公布的 FP8 白皮书中[arXiv:2209.05433] 介绍了 FP8 数据的两种精度：E4M3和E5M2。两种数据格式的具体二进制表示如下表：

FP8 E4M3的表示范围为[-448, 448]，E5M2为[-57334, 57334]，根据其数据表示范围和精度需求，一般而言， E4M3 格式更适合 weight、activation 数据，E5M2 格式更适合 grad 数据 。

为什么是 FP8，不是其他精度（比如 int8）？

• 与 int8 的数值表示相比较， FP8 在 LLM 的训练更有优势。因为 int8 在数值空间是均匀分布的，而 FP8 有更宽的动态范围， 更能精准捕获 LLM 中参数的数值分布。

• 尽管 FP16/BF16 已成为业界常用的训练精度，但在大模型训练场景下精度损失问题相对不敏感，仍然可以通过降低精度提升效率。FP8 训练能够在控制精度误差的情况下（Per-tensor Scaling），具有比 16bit 精度更快的计算速度和更少的资源占用，从而提升吞吐、降低训练通信量。

FP8 精度训练的效果如何？在下游任务的表现如何？

FP8 在绝大多数训练任务下都能有 FP16 相当的精度，在少部分下游任务（如数学运算）存在一定差距。

各种 CV 模型在 FP8 精度下训练的分类精度【NV测试结果】：

NLP 预训练任务【NV测试结果】：

LLM Benchmark【NV测试结果】：

SFT 微调效果：

FP8 有哪些应用场景/案例？

• Inflection AI 推出的 Inflection2 模型中，采用了 FP8 技术对其模型进行训练优化。 Inflection-2 采用了 FP8 混合精度在 5000 个 NVIDIA Hopper 架构 GPU 上进行了训练，累计浮点运算次数高达约10^{25}FLOPs。与同属训练计算类别的 Google 旗舰模型 PaLM 2 相比，在包括知名的 MMLU、TriviaQA、HellaSwag 以及 GSM8k 等多项标准人工智能性能基准测试中，Inflection-2 展现出了卓越的性能，成功超越了 PaLM 2，彰显了其在模型训练方面的领先性，同时也印证了 FP8 混合精度训练策略能够保证模型正常收敛并取得良好的性能。

• 零一万物基于 NVIDIA 软硬结合的技术栈，在功能开发、调试和性能层面，与 NVIDIA 团队合作优化， 完成了在大模型的 FP8 训练和验证。 其大模型的训练吞吐相对 BF16 得到了 1.3 倍的性能提升。零一万物的训练框架是基于 NVIDIA Megatron-LM 开发的 Y 训练框架, 其 FP8 训练基于 NVIDIA Transformer Engine。

• Google 与 NVIDIA 团队合作，将 TensorRT-LLM 应用于 Gemma 模型，并结合 FP8 技术进行了推理加速。 使用 Hopper GPU 进行推理时，FP8 对比 FP16 在吞吐量上能够带来 3 倍以上的收益。FP8 能够在相同的时间限制下使用更大的 batch size，从而有更好的 GPU 利用率，达到更高的吞吐量。

• 目前，NVIDIA 有专门使用 FP8 的 开源库—— Transformer Engine。

• Transformer Engine 和 FP8 已经集成到 PyTorch/JAX/Paddle Paddle 等基础深度学习框架中。

• 在专用于LLM的框架，比 如 Megatron/NeMo/DeepSpeed/HuggingFace/Colossal-AI 中也已经集成了 Transformer Engine 和 FP8，并有相应的 FP8 示例。

二、FP16/BF16 AMP

回顾 Pytorch AMP 的实现原理：

• 计算流程

• 显存分布

• Loss Scaling

1.计算流程

通常的 FP16 AMP 计算流程为：

• 数据、模型一开始都是 FP32 精度。

• 进入 torch.autocast 后，模型开始前向计算：

• 如果遇到 FP16 算子，则权重和数据都会转化为 FP16（权重一般有 FP16 cache，除非设置了autocast(cache_enabled=False)），然后在 FP16 精度的算子上进行前向计算，输出的结果也是 FP16 精度；

• 如果遇到 FP32 算子，则计算精度为 FP32，输出的精度也为 FP32。

• 反向计算，不需要在 torch.autocast 区域中，torch 会根据前向计算精度，自动确定反向计算精度。

• 优化器更新权重，利用 Tensor Core，可以直接完成 FP16 + FP32 的加法计算，以 FP32 精度更新权重，整个过程不需要精度转化。

FP16 支持算子： https:// pytorch.org/docs/stable /amp.html#cuda-ops-that-can-autocast-to-float16 （https://pytorch.org/docs/stable/amp.html#cuda-ops-that-can-autocast-to-float16）

Pytorch 使用 AMP 的样例代码如下（注意此处精度为 BF16）：

with torch.cuda.amp.autocast(dtype=torch.bfloat16):
    outputs = model(inputs)
    loss = loss_func(outputs, targets)

loss.backward()
optimizer.step()
optimizer.zero_grad()

2.显存分布

FP16 AMP 训练过程中，显存包含如下数据：

• 用于前向计算的模型权重： FP16

• 梯度： FP16

• 优化器：包含 FP32 Master Model Weight + 2*FP32 Adam States，即一阶矩和二阶矩

• 其余中间计算结果

3.Global Loss Scaling

由于 FP16 能够表示的数值范围更小，因此对于 FP16 精度的 AMP，需要进行 loss scaling。

scaler = torch.cuda.amp.GradScaler()

with torch.cuda.amp.autocast(dtype=torch.float16):
    outputs = model(inputs)
    loss = loss_func(outputs, targets)

scaler.scaled(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()

关于数值范围：FP32: 1-8-23 / BF16: 1-8-7 / FP16: 1-5-10
BF16 的数值范围和 FP32 一致，均有 8 个指数位，不需要 scaler 调整数值范围。
但 FP16 仅有 5 个指数位，当原数值过大或过小时，转换到 FP16 就可能出现 overflow/underflow，对训练造成影响。

实际上，我们会维护一个 全局的 scale 值，并采用 Dynamic Loss Scaling 动态调整这个全局的 scale 值。即每当梯度溢出时候减少损失缩放规模，并且间歇性地尝试增加损失规模，从而实现在不引起溢出的情况下使用 最高损失缩放因子 ，更好地恢复精度。

三、FP8 技术分析

1.宏观实现框架

FP16 所采用的 Loss Scaling 与量化的思想非常相似，它可以看成是 对全局的梯度数据做离线量化（PTQ） 。

FP8 的数据范围有更大的限制，单一的全局 Scale 值无法满足众多数据分布的相对精确表示，因此我们可以仿照量化的思路，将量化的基本单位缩小至 tensor（更细致的量化，如 Block-wise quantization，理论上可以用于更低精度的训练上）。

FP8 对每一个 tensor（无论是输入数据、前向计算结果、反向计算结果）都计算一个 Per-tensor Scaling Factor ，以此做更加细致的量化，充分利用 FP8 为数不多的格点数。

具体而言，每一次前向的 GEMM 计算需要对 3 个 tensor 记录 scale 值： input , weight , output ；而相对应的反向计算需要记录 2 个 tensor 的 scale 值： grad_output 和 grad_input 。在 TE 的 Hybrid 模式下，前向 tensor 数据格式为 E4M3，反向 tensor 数据格式为 E5M2。两种 FP8 精度的量化方式基本相同，均采用 对称线性量化 ，我们只需要关心 scale 值。可以参考之前的内容：

然而，FP8 训练最关键的问题是， 如何在训练过程中高效地寻找到这个 scale 值 ？

NV 在 TE 文档中给出了两种方案：

• Just-in-time scaling ，即先通过计算得到高精度的 output tensor，再在其上计算 amax，然后对 output tensor 做量化，输出 FP8 tensor。

  这些步骤在单个 kernel 层面上是不可能实现的，因为我们需要将完整的 output tensor 搬到 HBM 才能完成量化，因此这会将一个完整的计算过程分散成多次 kernel 调用，增加了数据传输量，拖慢了运行速度。 NV 认为这极大削减了 FP8 低精度带来的好处。

• Delayed scaling ，即假设我们提前知道了 scale 值，那么计算过程就可以在一个 kernel 上完成，而 amax 的计算和 scale 值的更新与其独立，不会中断计算进程，因此这种方式能完全发挥 FP8 的性能。但“提前知道的” scale 值需要额外的空间来记录，同时会引入一定的误差。
  

如下图所示，如果我们知道了 scale 值，那么计算的公式和伪代码就比较直接了：

TE 框架采用 Delayed scaling 方案 ，即对每个 GEMM 算子用到的 tensor 记录一个 amax history 数组，当我们需要 scale 值时，就从这个数组中取出最近一段时间窗口内 amax 的最大值，以此近似现在这个 tensor 的 amax，并默认用以下方式计算 scale 值：

FP8_MAX = maximum_representable_value(fp8_format)
new_scaling_factor = (FP8_MAX / amax) / (2 ^ margin)

用户可以自定义 Delayed scaling 的策略（Recipe），例如：

• margin：即上面公式的 margin，用于调整 scaling factor

• interval：经过多少 steps 更新一次 scaling factor

• fp8_format：指定前向反向的计算精度，默认为前向 E4M3，反向 E5M2

• amax_history_len：amax history 窗口长度，超过窗口长度的历史记录会被覆盖

• ……等等

2.TE 及各类框架集成方法

总的来说，任何结合 FP8 能力的框架都只需要做两件事：

1. 使用 TE 模块搭建 model，因为计算要用到 TE 提供的 FP8 算子；

2. 用 fp8_autocast 装饰前向计算过程。

在实际场景下， FP8 训练通常需要结合 BF16 混合精度训练 。

1）TE 官方案例：

import torch
import transformer_engine.pytorch as te
from transformer_engine.common import recipe

# Set dimensions.
in_features = 768
out_features = 3072
hidden_size = 2048

# Initialize model and inputs.
model = te.Linear(in_features, out_features, bias=True)
inp = torch.randn(hidden_size, in_features, device="cuda")

# Create an FP8 recipe. Note: All input args are optional.
fp8_recipe = recipe.DelayedScaling(margin=0, interval=1, fp8_format=recipe.Format.E4M3)

# Enable autocasting for the forward pass
with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):
    out = model(inp)

loss = out.sum()
loss.backward()

2）Accelerate：支持 DDP 和 FSDP 的 FP8 训练

# We prepare fp8 after, allowing for bf16 autocast to happen first
if getattr(self.fp8_recipe_handler, "backend", None) == "TE":
    if not has_transformer_engine_layers(model):
        with torch.no_grad():
            convert_model(model)
        model._converted_to_transformer_engine = True

    kwargs = self.fp8_recipe_handler.to_kwargs() if self.fp8_recipe_handler is not None else {}
    if "fp8_format" in kwargs:
        kwargs["fp8_format"] = getattr(te_recipe.Format, kwargs["fp8_format"])
    fp8_recipe = te_recipe.DelayedScaling(**kwargs)
    # If we are in DDP or FSDP, we delay `autocast` until after FSDP/DDP has been initialized
    # to make use of the process group
    if not self.delayed_fp8_autocast:
        model.forward = fp8_autocast(enabled=True, fp8_recipe=fp8_recipe)(model.forward)

• Line5 ：如果模型不是 TE 结构的，需要转化为 TE 结构的模型

• Line8-11 ：获取 fp8_format、fp8_recipe

• Line15 ：前向计算（model.forward）在fp8_autocast上下文管理器内完成

3）Megatron Core：支持 Tensor、Sequence、Pipeline 并行与 FP8 训练结合

# define TE model
use_te = args.transformer_impl == "transformer_engine"
if use_te:
    transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec()
model = GPTModel(transformer_layer_spec=transformer_layer_spec)

# set autocast context
class TransformerBlock(MegatronModule):
    def forward():
        if self.config.fp8:
            import transformer_engine  # To keep out TE dependency when not training in fp8

            if self.config.fp8 == "e4m3":
                fp8_format = transformer_engine.common.recipe.Format.E4M3
            elif self.config.fp8 == "hybrid":
                fp8_format = transformer_engine.common.recipe.Format.HYBRID
            else:
                raise ValueError("E4M3 and HYBRID are the only supported FP8 formats.")

            fp8_recipe = TEDelayedScaling(
                config=self.config,
                fp8_format=fp8_format,
                override_linear_precision=(False, False, not self.config.fp8_wgrad),
            )
            fp8_group = None
            if parallel_state.model_parallel_is_initialized():
                fp8_group = parallel_state.get_amax_reduction_group(with_context_parallel=True)
            fp8_context = transformer_engine.pytorch.fp8_autocast(
                enabled=True, fp8_recipe=fp8_recipe, fp8_group=fp8_group
            )
        else:
            fp8_context = nullcontext()

        with fp8_context:
            # Forward pass.

3.FP8 框架 TE 计算流程

首先分析 TE 框架入口 fp8_autocast 的源代码：