介绍

随着AI模型变得越来越复杂，训练它们所需的机器也是如此。Nvidia H100 GPU，据说支持“前所未有的性能和可扩展性”，是Nvidia最新也是最强大的AI加速器，专门设计用于支持下一代AI开发。随着当前对AI的热烈追捧，这些GPU的需求量巨大（例如，见此）。因此，不出所料，这些GPU的成本极为高昂——对许多读者来说甚至是禁止性的。幸运的是，云服务提供商如AWS、GCP和Microsoft Azure，提供了按需（按小时/按秒）使用H100支持机器的方式，从而为更广泛的AI开发者社群打开了使用这些机器的大门。

在AWS中，H100 GPU作为新近宣布的AWS EC2 p5实例系列的组成部分被提供。这些实例据称“与上一代基于GPU的EC2实例相比加速你的解决方案时间高达4倍，并降低训练机器学习模型的成本多达40%”。

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

模型

在下面的代码块中，我们定义了一个以Vision Transformer（ViT）为支撑的分类模型（使用流行的Python包timm版本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)

我们在p5.48xlarge和p4d.24xlarge实例类型上训练了这个模型，使用的是专门的PyTorch 2.1 AWS深度学习容器。

不出所料，p5的每步运行时间性能远超过p4d——0.199秒每步相比之下是0.41秒——快了两倍以上！！这意味着训练你的大型机器学习模型时间可以减半。

此时你可能得出两个可能的结论。第一种可能性是，尽管所有的炒作，p5简单来说并不适合你。第二种可能是p5仍然可能是可行的，但需要对你的模型进行调整，以充分利用其潜在能力。在接下来，我们将采纳第二种方法，并展示如何使用FP8数据类型——p5实例类型独有的——可以完全改变相比性价比的结果。

将FP8与Transformer Engine整合使用

我们应该强调的第一件事是，我们将使用Transformer Engine（TE），一个专门用于在NVIDIA GPU上加速Transformer模型的库。TE（版本0.12）预先安装在AWS PyTorch 2.1深度学习容器中。

尽管训练时使用FP8的理论超出了本帖子的范围，但重要的是要意识到使用FP8的机制远比16位替代品（float16和bfloat16）要复杂得多。幸运的是，TE实现隐藏了所有繁琐的细节，让用户不必处理。

为了修改我们的模型以使用TE，我们用一个自定义的变换器块类包裹TE的专用变换器层，这个类符合timm的块层签名。

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(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配方API），调整一些超参数，和/或限制FP8到模型的子部分应用。你可能会发现，尽管你尝试了各种方法，你的模型简单来说就是不兼容FP8。

结果

在下面的表中，我们总结了我们在p4d.24xlarge和p5.48xlarge EC2实例类型上，使用和未使用TE库的实验结果。对于p5.48xlarge实验，我们加倍了批量大小，以增加80 GB GPU内存的利用率。使用FP8减少了GPU内存消耗，使得批量大小进一步增加。

我们可以看到，使用TE变压器模块在p4d（约19%）和p5（约32%）实例类型上提高了性价比。在p5上使用FP8可进一步将性能提高约20%。在进行了TE和FP8优化后，基于H100的p5.48large的性价比超过了基于A100的p4d.24large——尽管优势不是非常大（约2%）。考虑到训练速度提高了3倍，我们可以有把握地得出结论，p5将是训练我们优化模型的更佳实例类型。

总结

在本文中，我们展示了如何编程PyTorch训练脚本以使用8位浮点类型。我们进一步展示了使用FP8如何成为从现代GPU，例如Nvidia H100中获得最佳性能的关键因素。重要的是，FP8的可行性以及其对训练性能的影响可以根据模型的具体细节而有很大差异。