在本博文中，我们将分解主要的 FP8 缩放策略 (按张量缩放、延迟和电流缩放以及按块缩放 (包括 Blackwell 支持的 MXFP8 格式)) ，并解释为什么每个策略都对在低精度训练期间保持数值稳定性和准确性至关重要。了解这些方法有助于为您自己的 FP8 工作流程选择合适的方法。
本文将探讨 FP8 训练的实际情况，重点介绍 NVIDIA Nemotron 实验。我们将探讨为什么需要按张量和延迟扩展，以及它们在哪些方面有不足，以及高级方法 (例如按块子通道扩展) 如何为大型模型实现稳定、高效的 FP8 训练。
在上一篇博文中，我们介绍了 FP8 精度，并将其两种主要格式 E4M3 和 E5M2 与更成熟的类型（如 BF16）进行了比较。
按张量缩放
每张量缩放是一种基本的 FP8 策略，可为每个张量分配唯一的缩放系数（例如权重、激活函数或梯度），而不是使用单一的全局缩放。这一点至关重要，因为 FP8 的窄动态范围无法处理跨不同张量的不同值分布。通过根据每个张量的统计数据调整缩放系数，每张量缩放有助于防止数值不稳定，并确保更准确、更稳定的训练。
延迟缩放和当前缩放是每个张量缩放中的两种常见方法。每种方法都提供了不同的缩放因子设置策略，但都在单个张量级别运行，以满足 FP8 的独特要求。
按张量延迟扩展
延迟缩放是一种实用的 FP8 缩放策略，该策略基于在先前迭代窗口中观察到的最大绝对值 (amax) 来计算每个张量的缩放系数，而不仅仅是当前批量。这种基于历史记录的方法可以消除异常值，并降低可能会破坏训练稳定的突然缩放变化的风险。
延迟扩展并非基于当前的数据批量来确定扩展系数，而是着眼于最近几批中的最大值 (称为“amax”) ，并保留这些最大值的运行历史。
当选择新的缩放系数时，系统会查看此历史记录并选择反映近期趋势的值，而不仅仅是最近的批量。这种方法有助于消除数据值的突然峰值或下降，使训练过程更加稳定，并降低可能损害模型性能的突然变化的风险。
虽然延迟扩展适用于 FP8 训练，但其有效性取决于一个关键假设：张量值的统计分布在训练迭代中保持相对稳定。但是，异常值可能会在 amax 历史记录中占据主导地位，导致次优扩展，导致不足或溢出，并可能破坏训练稳定，特别是在大规模运行中。
按张量电流缩放
虽然延迟扩展使用历史数据，但每张量当前扩展提供了一种优先考虑实时适应性的替代方案。此方法根据当前向前或向后传递中各tensor的统计属性确定其缩放系数。
按张量电流缩放可根据当前数据范围动态调整缩放。这种即时响应能力有助于优化 FP8 表征，并可在训练期间改善模型收敛。
其工作原理是在每个训练批次中测量每个张量中的最大绝对值 (amax)，无论是激活值、权重还是梯度。此测量是实时进行的，无需依赖前一批的任何历史记录。
然后，将缩放系数设置为 amax 适合 FP8 范围，其中每个张量使用 FP8 的完整可用精度来执行特定步骤。
借助此缩放系数，张量的浮点值可量化为 FP8，从而在每次迭代中实现高效且准确的训练。当前缩放的优势如下：
- 瞬时适应：当前扩展可适应数据分布、学习率时间表或模型架构中的实时变化。这种实时响应能力可确保在扩展过程中不会出现延迟或不匹配。
- 简单高效：缺少复杂的历史缓冲区或跟踪机制可简化实施，并减少计算和内存用度。资源仍然直接集中在模型训练上。
- 鲁棒性：实时调整可针对当前批量中的异常值提高电流扩展的鲁棒性。由于规模不受过去、过时或错误的 amax 值的影响，因此对于要求苛刻的深度学习工作流而言，它将提供更稳定的训练体验。
图 2 突出显示了在平均推理等具有挑战性的基准测试中，与 BF16 基准相比，低精度格式的效率。
虽然 BF16 在整个 token 范围内的总精度方面保持略微边缘，但当前 scaling 紧随其后，差距仍保持在较小的百分比余量内。
逐块扩展
虽然每张量扩展方法提供了 FP8 训练的基础方法，但在处理单个张量内的块级可变性时，它们通常会面临局限性。
当张量内的不同区域或切片表现出不同的数值范围和统计行为时，就会出现这种挑战。例如，在大型 Transformer 模型中，单个权重矩阵中不同的注意力头可能具有不同的数量级。同样，高级层 (如 SwiGLU 或 mixture-of-experts (MoE) 路由器) 的激活函数因生成由小值和罕见但极端的异常值组成的密集集群的异构混合而声名狼。
当对不同张量应用单个缩放系数时，要么牺牲密集、低数量级区域的精度，要么极端异常值可能会被削减。两者都会降低训练稳定性和模型的最终质量。
每块缩放是一种功能强大的解决方案，可将每个 tensor 划分为更小、更易于管理的连续块，并为每个 张量分配专用的缩放因素。
缩放机制可以适应每个块的局部统计数据，而不是由单个最极端值决定。这意味着高数量级区域得到准确表示，而不会影响同一张量内更小、更典型的值的保真度。
随着深度学习模型的规模和架构复杂性不断增加，细粒度、自适应扩展对于保持数字稳定性以及释放 FP8 在现代硬件上提供的效率和速度至关重要。
接下来，我们将探索如何通过硬件原生 MXFP8 和其他可配置块扩展方法等领先方案，在实践中实施这一先进的扩展概念。
什么是 Micro-Scaling FP8？
Micro-Scaling FP8 (MXFP8) 基于逐块扩展概念构建，代表了 NVIDIA Blackwell 硬件级解决方案，可实现高效稳定的 FP8 训练。
MXFP8 符合更广泛的 MX 数据格式标准，该标准定义了一个通用框架，用于跨各种低精度格式 (FP8、FP6、FP4 和 Int8) 进行共享的细粒度块扩展。该标准具有三个关键组成部分：
- 缩放 (X) 数据类型/编码：定义缩放系数本身的表示方式。
- 元素 (P_i) 数据类型/编码：指定每个块内实际数值元素的数据类型。
- 缩放块大小 (k) ：确定共享通用缩放系数的标量元素的数量。
MXFP8 的工作原理是什么？
MXFP8 实现块级扩展，具有针对 Blackwell 架构优化的特定特征：
- 块级除法：张量被系统划分为连续块，每个块包含 32 个连续值，与 NVIDIA Blackwell Tensor Cores 的设计相匹配。
- 仅指数级缩放系数：32 值块中的每个块都分配了一个以 E8M0 格式存储的专用缩放系数（8 位指数，0 位尾数），并表示一个 2 的幂乘数。这可优化硬件实现，并保持有利于深度学习的数值属性。
- 转置的硬件重构：鉴于缩放发生在特定方向 (例如，逐行或逐列) ，MXFP8 张量及其转置在数字上不等效。当运算需要张量及其转置 (常见于向前和向后传递) 时，Blackwell 单元会自动从高精度输入进行重量化。这可避免因转置量化 FP8 张量而复合量化错误，从而保持准确性。
这种创新的块级扩展使 MXFP8 能够克服按张量扩展的限制，适应单个张量内的数量级变化。通过保留高和低数量级区域，MXFP8 可最大限度地利用 FP8 的动态范围并最小化量化误差。这转化为更快的训练速度和更高的准确性，对于大型模型来说尤其如此。
图 5 比较了 Nemotron 2b 和 8b 模型的 BF16 和 MXFP8 损失。在 2B 参数 Nemotron 模型中，MXFP8 显示的验证损失曲线与 BF16 相同，即使训练令牌数量增加也是如此。这意味着，对于此模型，使用 BF16 和 MXFP8 数据类型之间没有显著的准确性差异，并且模型会按预期收敛。
借助 8B 参数模型，即使训练令牌数量增加，MXFP8 的验证损失曲线也与 BF16 相同。这再次表明，MXFP8 的收敛与 BF16 相似，即使模型大小增加，精度也没有任何重大变化。
块缩放
除了按张量缩放策略和 MXFP8 等特定于硬件的块缩放之外，通用 FP8 块缩放是一种通用且可配置的细粒度精度控制方法。
这项技术因其适应广泛的模型架构和硬件要求而脱颖而出。
用户可以指定各种形状 (例如 1×128 或 128×128) ，以优化不同张量结构的缩放。这些特定于块的扩展因子通常以 FP32 格式存储，从而提供可靠的数值属性。
这种可配置粒度可更精确地适应不同模型组件的独特统计行为，成为跨不同工作负载优化 FP8 训练的强大工具。
块扩展的工作原理是什么？
FP8 块缩放背后的核心原则是通过调整缩放系数以适应局部数据特征来提高精度。与在整个张量上应用单一比例的每张量方法不同，块缩放会将每个张量分为更小、不同的片段。在每个定义的块内，所有值都共享一个单独存储在 FP32 中的通用缩放系数，以确保准确性。
量化按元素块的系数来扩展元素，并将其转换为 FP8；去量化通过将 FP8 值乘以相同的系数来重建原始值。块大小的选择至关重要：较小的块通过更精细的粒度减少量化误差，而较大的块则以增加噪音为代价降低存储开销。
NVIDIA NeMo 框架中的 Recipes
Transformer 引擎 FP8 配置通过高级配置公开，例如在 NVIDIA NeMo 框架中。用户可以选择不同的 fp8_recipe 标志来进行缩放策略，通常结合 BF16 进行混合精度训练。
- 延迟扩展：
fp8_recipe = "delayed" - 每张量电流缩放：
fp8_recipe = "tensorwise"MXFP8：fp8_recipe = "mxfp8"通用块缩放：fp8_recipe = "blockwise"
总结
FP8 训练是适用于大规模深度学习的实用高效解决方案。充分发挥 FP8 潜力的关键在于选择和实施扩展策略。通过访问 FP8 recipes 探索这些技术，并开始使用实用的 FP8 训练配置和代码。
