使用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%)，所以可能还有更多的优化方案，我们需要继续研究。

总结

在这篇文章中，我们演示了如何编写PyTorch训练脚本来使用8位浮点类型。展示了FP8的使用是如何从Nvidia H100中获得最佳性能的关键因素。FP8的可行性及其对训练性能的影响可以根据模型的细节而变化很大。

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

作者：Chaim Rand