背景 - H200 是什么？
H200 是 H100 的 “增强版”，它提供与 H100 完全相同的计算能力，但有两项改进：
-
更大的全局内存，拥有 141GB 的 HBM3e，而标准的 HBM3 只有 80GB。 -
内存带宽快了约 43%，达到 4.8TB/s，而 H100 为 3.35TB/s。更快的内存传输对训练速度有显著影响，特别是对于 PyTorch 的异步张量处理（AsyncTP）。
什么是 PyTorch Float8 逐行训练？
Float8 逐行训练是相较于之前的 “张量级” Float8 更细粒度的 Float8 分辨率。它旨在确保更细粒度的精度，以支持更大规模的工作负载，因为随着训练的进行和规模的扩大，这些工作负载对量化更加敏感。
Float8 逐行训练有两个关键改进：
-
每行现在都有自己的缩放因子，而不是整个张量使用单一缩放因子，从而提高了量化精度。更细粒度的逐行缩放有助于减少异常值（那些迫使量化缩放因子拉伸、降低正态分布值精度的极端值）的影响，进而确保更高的精度。 -
缩放因子现在通过向下舍入到最接近的 2 的幂次方来实现。这已被证明有助于减少在乘以 / 除以缩放因子时的量化误差，同时确保在正向和反向传播中，大的值都能被缩放到相同的值。
其他大规模模型也曾使用 2K 规模的 Float8 进行训练，结合 1x128 分组和 128x128 分块的方式，并使用 2 的幂次方缩放因子。它们的目标同样是提高 Float8 的精度，以支持大规模训练。
因此，Float8 逐行训练也有望实现大规模训练，但我们需要证明其在大规模训练中的稳定性和收敛性，在 Crusoe H200 2K 集群上的训练为此提供了初步验证。
Float8 逐行训练与 BF16 在 1600 和 1920 个 GPU 规模下的损失收敛情况
为了验证损失收敛情况的可比性，我们使用 TorchTitan 和 Lllama3 70B 模型，分别在 1920 和 1600（1.6K）个 GPU 规模下进行了两次独立运行。1.6K GPU 的运行设置为 2500 次迭代，使用 TorchTitans 的 HSDP2 和上下文并行技术来实现 2D 并行。
损失收敛测试是在 Titan 的确定性模式下运行的。这种模式有效地冻结了每次运行中大多数潜在的变化源，因此有助于确保唯一的实质性变化是我们想要测试的内容，即 BF16 和 Float8 逐行训练的损失收敛情况和损失曲线。
需要注意的是，确定性模式也会降低训练速度，因为各种内核不会进行自动调优以最大化吞吐量（否则我们可能会在不同运行中使用不同的内核，从而引入差异）。
我们完成了两次运行，一次使用 BF16，另一次使用 Float8 逐行训练。
两次运行都顺利完成了指定的 2500 次迭代，这展示了 Crusoe 集群的稳定性。使用 Float8 的训练在 24 小时内精确完成，而 BF16 则在 31 小时 19 分钟后完成。
在 24 小时时，Float8 完成了 2500 次迭代，展示了 Float8 训练的相对加速（即使在确定性模式下）。在 24 小时时，与 BF16 相比，Float8 在相同的 24 小时大规模训练时间内，损失相对改善了 9.21%。
31 小时 19 分钟后，BF16 的运行终于完成了 2500 次迭代。
最终损失数据：
-
BF16 = 2.88109 -
Float8 = 2.86386
从损失曲线来看，我们观察到在开始和最后的三分之一部分，两条曲线非常相似，中间有一个波动区域，在这个区域两者都出现了类似的峰值，但峰值出现的相对时间略有偏差。
因此，我们可以看到，PyTorch 的 Float8 逐行训练在收敛性上与 BF16 相似，但在相同的训练时间内速度提升超过 33%。
Float8 逐行训练的长期稳定性
除了展示可比的收敛性，我们还想展示 Float8 的长期训练稳定性，因此我们在 256 规模下进行了一次为期 4 天、15000 次迭代的运行。
如上图所示，Float8 训练持续运行了 100 多个小时没有出现问题，突出了 Float8 逐行训练的长期稳定性。
TorchTitan 中的确定性
为了验证确定性，并查看较长运行中的波动是否由规模引起，我们还进行了一次较小规模的运行，包括两次 BF16 运行和一次 Float8 运行，规模为 256，并且仅使用 HSDP2（即不使用 2D 上下文并行）。
在这种情况下，两次 BF16 运行的曲线和最终损失完全相同，并且我们在所有三次运行中都看到了类似的波动区域。
-
BF16（两次运行） = 3.28538 -
Float8 逐行训练 = 3.28203
为了保证确定性，实验使用的是序列化的 C4 数据集（未打乱），这意味着峰值可能是由于在数据集中遇到了新内容。
Float8 逐行训练在不同规模下的净加速
我们在不同的 GPU 规模下进行了较短时间的运行，以了解随着集群规模的扩大，Float8 逐行训练在训练加速方面的表现。从 960 个 GPU 扩展到 1920 个 GPU 时，Float8 继续实现令人瞩目的训练加速，与 BF16 相比，加速范围超过 34%-43%。我们还注意到，从 1000 个 GPU 扩展到 2000 个 GPU 时，通信开销可能会增加，我们观察到 BF16 的吞吐量下降了 4%。
如上述大规模长时间训练运行所示，Float8 逐行训练实现了显著的加速，在 1920（DeepSeek）规模下实现了 34% 的加速，同时损失终点相等甚至略有改善。
如何在训练中使用 Float8 逐行训练？
Float8 逐行训练现已可供你在大规模训练中使用。它包含在 TorchAO[2] 的最新版本（0.9 及更高版本）中，如果你想快速上手，可以直接集成到 TorchTitan[3] 中。
在 TorchTitan 中激活 Float8 逐行训练：
-
首先，启用模型转换器，在你的 models.toml 文件中，将 nn.linear 层热替换为 float8 线性层，见第 29 行： 
-
其次，指定 “逐行” float8 配置，见第 72 行： 
请注意，“recipe_name” 有三个选项：
-
rowwise，这是推荐的默认选项； -
tensorwise（旧版的 float8 方式）； -
rowwise_with_gw_hp。
gw_hp 逐行选项在反向传播时，将权重的梯度保持在 BF16 精度，这可以进一步提高对极端敏感工作负载的 float8 精度。但具有讽刺意味的是，如果模型中的大多数矩阵乘法尺寸较小（在 H100 上，估计转折点约为 13 - 16K 维），它的性能可能比普通的逐行选项更好。
因此，虽然我们推荐将 rowwise 作为默认选项，但在你的模型上与 gw_hp 进行比较，验证哪个选项能提供最佳性能可能是值得的，因为 gw_hp 还有更高精度的优势。
未来更新
我们将发布更多更新内容，展示在流水线并行和异步分布式检查点方面的多项改进，请持续关注。