使用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。

模型

我们定义了一个Vision Transformer (ViT)支持的分类模型(使用流行的timm Python包版本0.9.10)以及一个随机生成的数据集。ViT主干有多种形状和大小。我们选择了通常被称为ViT-Huge的配置-具有6.32亿个参数-这样能够更好地利用H100对大型模型的容量。

 import torch, time
 import torch.optim
 import torch.utils.data
 import torch.distributed as dist
 from torch.nn.parallel.distributed import DistributedDataParallel as DDP
 import torch.multiprocessing as mp
 
 # modify batch size according to GPU memory
 batch_size = 64
 
 from timm.models.vision_transformer import VisionTransformer
 
 from torch.utils.data import Dataset
 
 
 # use random data
 class 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, label
 
 
 def 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上都训练了这个模型。

p5的性能远远超过了p4d的性能——每步0.199秒比0.41秒——快了两倍多!!这意味着训练大型机器学习模型的时间将减少一半。但是当你考虑到成本的差异(p4d每小时32.77美元，p5每小时98.32美元)，p5的性价比比p4d差30% !!

在这一点上，可能会得出两个可能的结论之一。第一种可能性是，尽管有这么多宣传，但p5根本不适合您。第二个是p5仍然是可行的，但是需要对模型进行调整，充分利用它的潜力。

FP8与Transformer Engine的集成

PyTorch(版本2.1)不包括FP8数据类型。为了将我们的脚本编程为使用FP8，我们将使用Transformer Engine (TE)，这是一个用于在NVIDIA gpu上加速Transformer模型的专用库。TE(版本0.12)预装在AWS PyTorch 2.1 DL容器中。

使用FP8的机制比16位(float16和bfloat16)要复杂得多。TE库实现向用户隐藏了所有杂乱的细节。有关如何使用TE api的说明（请参阅官方文档）。

为了修改我们的模型以使用TE，我们将TE的专用Transformer层，所以需要我们自己写一个包装器：

 import transformer_engine.pytorch as te
 from transformer_engine.common import recipe
 
 
 class 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位表示)意味着较低的精度和较低的动态范围。这些可以对模型收敛的可达性和/或速度产生有意义的影响，但不能保证这将适用于所有的模型。所以可能需要调整底层FP8机制(例如，使用TEapi)，调整一些超参数，和/或将FP8的应用限制在模型的子模型（一部分）。最坏的可能是尽管进行了所有尝试，模型还是无法与FP8兼容。

结果

在下表中总结了在两个p4d上的实验结果。24xlarge和p5.48xlarge EC2实例类型，使用和不使用TE库。对于p5.48xlarge实验，我们将批处理大小加倍，这样提高80 GB GPU内存的利用率。使用FP8可以减少GPU内存消耗，从而进一步增加批处理大小。

可以看到，使用TE提高了p4d(19%)和p5(32%)的性价比。使用FP8可将p5上的性能额外提高约20%。在TE和FP8优化之后，基于h100的p5.48large的性价比优于基于a100的p4d.xlarge——虽然差距不大(2%)。考虑到训练速度提高了3倍，我们可以有把握地得出结论，p5将是训练优化模型的更好的实例类型。

但是我们也看到了，这是相对较小的性价比提升(远低于p5公告中提到的40%)，所以可能还有更多的优化方案，我们需要继续研究。