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 的使用促使模型在训练和推理过程中进行 ………………………………