使用FP8加速PyTorch训练工作负载

使用FP8加速PyTorch训练工作负载的方法

如何充分发挥现代GPU的潜力

Deva Darshan在Unsplash上的照片

过去几年中,人工智能领域取得了革命性的进步,最好的例证或许是最近LLM类应用程序(如ChatGPT)的流行和普及。这些突破得益于在训练人工智能模型所使用的机器上同样令人兴奋的发展。新颖的创新架构、复杂的张量处理核心和专用硬件加速器使得不断增大规模的人工智能模型能够以越来越快的速度相互融合。本文将重点讨论人工智能专用硬件的一个特定进展——专用的8位浮点(FP8)张量处理核心。出现在最先进的人工智能硬件架构中(例如Nvidia Hopper、Nvidia Ada Lovelace和Habana Gaudi2),FP8张量核心使得每秒浮点运算量(FLOPS)大幅增加,并为人工智能训练和推理工作负载提供了内存优化和能源节约的机会。

要充分利用硬件级别的FP8能力,我们需要在构建人工智能训练和推理应用程序时使用合适的软件栈和开发框架。本文将说明如何修改一个PyTorch训练脚本以利用Nvidia H100 GPU的FP8支持。首先我们将提供使用FP8数据类型的动机。然后我们将回顾由Transformer Engine库提供的FP8特定的PyTorch API支持,并展示如何将其整合到一个简单的训练脚本中。虽然我们不会详细讨论FP8在人工智能训练中的理论,但我们会指出使用FP8可能涉及的潜在挑战。最后,我们将展示FP8数据类型的重大优化机会。

免责声明

请不要将我们对任何软件组件、方法或服务的提及解释为对其使用的认可。最佳的机器学习开发设计将因您自己的人工智能工作负载的具体细节而大相径庭。请还注意我们将提到的某些软件包和组件的API和行为可能在您阅读本文时发生变化。强烈推荐您根据最新的硬件和软件评估任何潜在的设计决策。

动机

随着人工智能模型变得越来越复杂,用于训练它们的机器也越来越复杂。据说支持“空前的性能和可伸缩性”的Nvidia H100 GPU是(在撰写本文时)Nvidia最新和最强大的人工智能加速器,旨在推动下一代人工智能的发展。随着当前人工智能热潮的全面展开,对这些显卡的需求非常庞大(例如在这里)。因此,这些显卡的成本非常高昂,对于很多读者来说甚至是禁得起的。幸运的是,AWS、GCP和Microsoft Azure等云服务提供商提供“按小时/按秒计费”的访问方式,可以使用H100供电的机器,从而为更多的人工智能开发者提供使用的机会。

在AWS中,H100 GPU作为最近宣布的AWS EC2 p5实例系列的组成部分提供。据称,这些实例可以“与基于上一代GPU的EC2实例相比,将解决方案的时间加速提高4倍,并将训练ML模型的成本降低40%”。

最近的一篇文章中,我们讨论了选择ML训练实例时应考虑的一些因素。我们强调了最佳实例类型将非常依赖于手头的项目这一事实。尤其是在p5实例系列方面,更大的并不总是更好。这在成本方面尤其如此(在本文撰写时,p5的费用为每小时98.32美元,对于8-GPU的p5.48xlarge实例)。您可能会发现其他实例类型更合适。

在下一部分中,我们将使用p5.48xlarge训练一个相对较大的计算机视觉模型,并将其性能与包含8个Nvidia A100 GPUp4d.24xlarge进行比较。

玩具模型

在下面的代码块中,我们定义了一个使用流行的timm Python包版本0.9.10支持的Vision Transformer(ViT)支持的分类模型,以及一个随机生成的数据集。ViT骨干结构有各种形状和大小。在这里,我们选择了通常被称为ViT-Huge配置的模型,该模型具有6.32亿个参数,以更好地利用H100对大模型的容量。

import torch, timeimport torch.optimimport torch.utils.dataimport torch.distributed as distfrom torch.nn.parallel.distributed import DistributedDataParallel as DDPimport torch.multiprocessing as mp# modify batch size according to GPU memorybatch_size = 64from timm.models.vision_transformer import VisionTransformerfrom torch.utils.data import Dataset# use random dataclass FakeDataset(Dataset):    def __len__(self):        return 1000000    def __getitem__(self, index):        rand_image = torch.randn([3, 224, 224], dtype=torch.float32)        label = torch.tensor(data=[index % 1000], dtype=torch.int64)        return rand_image, labeldef mp_fn(local_rank, *args):    # configure process    dist.init_process_group("nccl",                            rank=local_rank,                            world_size=torch.cuda.device_count())    torch.cuda.set_device(local_rank)    device = torch.cuda.current_device()        # create dataset and dataloader    train_set = FakeDataset()    train_loader = torch.utils.data.DataLoader(        train_set, batch_size=batch_size,        num_workers=12, pin_memory=True)    # define ViT-Huge model    model = VisionTransformer(            embed_dim=1280,            depth=32,            num_heads=16,        ).cuda(device)    model = DDP(model, device_ids=[local_rank])    # define loss and optimizer    criterion = torch.nn.CrossEntropyLoss()    optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)    model.train()    t0 = time.perf_counter()    summ = 0    count = 0    for step, data in enumerate(train_loader):        # copy data to GPU        inputs = data[0].to(device=device, non_blocking=True)        label = data[1].squeeze(-1).to(device=device, non_blocking=True)          # use mixed precision to take advantage of bfloat16 support        with torch.autocast(device_type='cuda', dtype=torch.bfloat16):            outputs = model(inputs)            loss = criterion(outputs, label)        optimizer.zero_grad(set_to_none=True)        loss.backward()        optimizer.step()                # capture step time        batch_time = time.perf_counter() - t0        if step > 10:  # skip first steps            summ += batch_time            count += 1        t0 = time.perf_counter()        if step > 50:            break    print(f'average step time: {summ/count}')if __name__ == '__main__':    mp.spawn(mp_fn,             args=(),             nprocs=torch.cuda.device_count(),             join=True)

我们使用专用的PyTorch 2.1 AWS深度学习容器 (763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.1.0-gpu-py310-cu121-ubuntu20.04-ec2) 在<p5.48xlarge和<p4d.24xlarge实例类型上训练了这个模型。</p4d.24xlarge</p5.48xlarge

毫不意外,p5的步骤性能超过了p4d性能 – 每个步骤的时间为0.199秒,而p4d为0.41秒 – 快了两倍多!这意味着训练大型机器学习模型的时间减少一半。然而,当考虑到成本的差异(目前而言,p4d的每小时费用为32.77美元,而p5的每小时费用为98.32美元)时,就会出现完全不同的情况。p5的性价比比p4d差约30%!这与<p5发表中出现的40%改进相去甚远。</p5发表

此时,您可以得出两种可能的结论之一。第一种可能性是,尽管所有炒作,p5并不是适合您的机器。第二种可能性是,p5仍然可行,但要充分利用其潜力,需要对模型进行调整。在接下来的几节中,我们将采用第二种方法,演示如何使用FP8数据类型 – 仅适用于p5实例类型 – 可以完全改变比较价格性能结果。

将FP8与Transformer引擎集成

首先,我们应该强调的是,截至目前,PyTorch(版本2.1)不包括本地支持的8位浮点数据类型。为了使我们的脚本使用FP8,我们将使用Transformer Engine(TE),这是一个专用库,用于加速NVIDIA GPU上的Transformer模型。 TE(版本0.12)已预装在AWS PyTorch 2.1 DL容器中。

尽管训练中使用FP8的理论超出了此帖子的范围(例如,请参见此处),但重要的是要知道,与16位替代方案(float16和bfloat16)相比,使用FP8的机制要复杂得多。幸运的是,TE的实现将所有杂乱的细节隐藏在用户背后。请参阅官方文档以及此简单的示例,以了解如何使用TE API的说明。要了解背后的情况,请务必查看以下两个视频教程。

使用Transformer引擎进行FP8训练 | NVIDIA按需

本课程将介绍FP8和混合精度,Transformer引擎功能概述以及…

www.nvidia.com

深度学习的FP8 | NVIDIA按需

FP8是推动深度学习(DL)训练超越现代16位格式的自然发展…

www.nvidia.com

要修改我们的模型以使用TE,我们将以符合timm的块层签名的自定义transformer块类进行TE的专门Transformer Layer封装。

import transformer_engine.pytorch as tefrom transformer_engine.common import recipeclass TE_Block(te.transformer.TransformerLayer):    def __init__(            self,            dim,            num_heads,            mlp_ratio=4.,            qkv_bias=False,            qk_norm=False,            proj_drop=0.,            attn_drop=0.,            init_values=None,            drop_path=0.,            act_layer=None,            norm_layer=None,            mlp_layer=None    ):        super().__init__(            hidden_size=dim,            ffn_hidden_size=int(dim * mlp_ratio),            num_attention_heads=num_heads,            hidden_dropout=proj_drop,            attention_dropout=attn_drop            )

接下来,我们修改VisionTransformer初始化来使用我们自定义的块层:

  model = VisionTransformer(      embed_dim=1280,      depth=32,      num_heads=16,      block_fn=TE_Block      ).cuda(device)

到目前为止我们还没有对H100进行任何特定的更改 – 相同的代码可以在我们的A100提供动力的p4d实例类型上运行。最后的修改是用te.fp8_autocast上下文管理器包装模型正向传递。这个改变需要一个支持FP8的GPU:

with torch.autocast(device_type='cuda', dtype=torch.bfloat16):    with te.fp8_autocast(enabled=True):        outputs = model(inputs)    loss = criterion(outputs, label)

关于使用FP8的一些谨慎说明

使用8位浮点表示(而不是16或32位表示)意味着较低的精度和较低的动态范围。这些可能会对模型收敛的可达性和/或速度产生显著影响。虽然底层TE FP8实现旨在解决这一挑战,但不能保证这对您的模型有效。您可能需要调整底层FP8机制(例如使用TE recipe API),调整一些超参数,和/或限制FP8的应用于模型的部分。您可能会发现,尽管您采取了一切努力,但您的模型可能与FP8不兼容。

结果

在下表中,我们总结了我们在p4d.24xlarge和p5.48xlarge EC2实例类型上的实验结果,包括TE库和不包括TE库的情况。对于p5.48xlarge实验,我们增加了批量大小,以提高80 GB GPU显存的利用率。使用FP8可以减少GPU内存消耗,进一步增加批量大小。

实验结果(作者提供)

我们可以看到,使用TE transformer块在p4d(约19%)和p5(约32%)实例类型上增加了性价比。使用FP8可以使p5性能额外增加约20%。遵循TE和FP8的优化基于H100的p5.48large的价格性能超过了基于A100的p4d.24large – 尽管差距不是非常大(约2%)。考虑到训练速度的3倍增加,我们可以安全地得出结论,对于训练我们优化的模型,p5将是更好的实例类型。

请注意,性能-价格的相对小幅增加(远低于p5公告中提到的40%)使我们希望获得更多的H100特定优化…但这将需要等待另一篇文章:)。

总结

在本文中,我们演示了如何编写一个PyTorch训练脚本来使用8位浮点类型。我们进一步演示了如何使用FP8可以成为实现现代GPU(例如Nvidia H100)最佳性能的关键因素。重要的是,FP8的可行性以及它对训练性能的影响可能会因模型的细节而有所不同。

本文是关于优化机器学习工作负载的系列文章之一。务必查看我们有关此重要主题的其他文章