使用FP8加速PyTorch训练

现代的人工智能硬件架构(例如，Nvidia Hopper, Nvidia Ada Lovelace和Habana Gaudi2)中，FP8张量内核能够显著提高每秒浮点运算(FLOPS)，以及为人工智能训练和推理工作负载提供内存优化和节能的机会。

在这篇文章中，我们将介绍如何修改PyTorch训练脚本，利用Nvidia H100 GPU的FP8数据类型的内置支持。这里主要介绍由Transformer Engine库公开的fp8特定的PyTorch API，并展示如何将它们集成到一个简单的训练脚本中。（我们这里只介绍如何使用FP8，不会介绍FP8具体的理论知识）

随着人工智能模型变得越来越复杂，训练它们所需的机器也越来越复杂。Nvidia H100 GPU据称支持“前所未有的性能和可扩展性”。

在AWS中，H100 gpu是作为AWS EC2 p5实例的一个组件提供的。这些实例声称“与上一代基于gpu的EC2实例相比，可将解决方案的时间加快4倍，并将训练ML模型的成本降低高达40%”。

当涉及到机器学习训练实例时，并不总是越大越好。p5实例族尤其如此。p5可能会比其他实例要快很多，因为H100是无可争议的性能野兽。但是一旦考虑到p5的成本(8-GPU p5.48xlarge实例的成本为每小时98.32美元)，你可能会发现其他实例类型更适合。

下面我们将在p5.48xlarge上训练一个相对较大的计算机视觉模型，并将其性能与p4d进行比较。p4d.24xlarge包含8个Nvidia A100 gpu。

https://avoid.overfit.cn/post/541a04c656db474d91ee5eb1fa5bc5f8