使用FP8加速PyTorch训练

发布时间 2023-11-17 09:46:07作者: deephub

现代的人工智能硬件架构(例如,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