使用FP8加速PyTorch训练的两种方法总结

本文涉及的产品
实时计算 Flink 版,5000CU*H 3个月
检索分析服务 Elasticsearch 版,2核4GB开发者规格 1个月
大数据开发治理平台 DataWorks,不限时长
简介: 在PyTorch中,FP8数据类型用于高效训练和推理,旨在减少内存占用和加快计算速度。虽然官方尚未全面支持,但在2.2版本中引入了`torch.float8_e4m3fn`和`torch.float8_e5m2`。文章通过示例展示了如何利用FP8优化Vision Transformer模型,使用Transformer Engine库提升性能,并探讨了PyTorch原生FP8支持的初步使用方法。实验表明,结合TE和FP8,训练速度可提升3倍,性能有显著增强,特别是在NVIDIA GPU上。然而,PyTorch的FP8支持仍处于试验阶段,可能带来不稳定性。

在PyTorch中,FP8(8-bit 浮点数)是一个较新的数据类型,用于实现高效的神经网络训练和推理。它主要被设计来降低模型运行时的内存占用,并加快计算速度,同时尽量保持训练和推理的准确性。虽然PyTorch官方在标准发布中尚未全面支持FP8,但是在2.2版本中PyTorch已经包含了对FP8的“有限支持”并且出现了2个新的变量类型,

torch.float8_e4m3fn

torch.float8_e5m2

,而H100也支持这种类型,所以这篇文章我们就来介绍如何使用FP8来提高训练效率

模型架构

我们定义了一个Vision Transformer (ViT)支持的分类模型(使用流行的timm Python包版本0.9.10)以及一个随机生成的数据集。我们选择了ViT-Huge的有6.32亿个参数的最大的模型,这样可以演示FP8的效果。

 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)

Transformer Engine

PyTorch(版本2.1)不包括FP8的数据类型。所以我们需要通过第三方的库Transformer Engine (TE),这是一个用于在NVIDIA gpu上加速Transformer模型的专用库。

使用FP8要比16float16和bfloat16复杂得多。这里我们不用关心细节,因为TE都已经帮我们实现了,我们只要拿来用就可以了。

但是需要对我们上面的模型进行一些简单的修改,需要将transformer变为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)

最后一个修改是用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)

下面我们就可以测试结果:

可以看到,使用TE块提高了p4d(~19%)和p5(~32%)的性价比。使用FP8可将p5上的性能额外提高约20%。在TE和FP8优化之后,基于h100的p5.48large的性价比优于基于a100的p4d.24large 。并且训练速度提高了3倍。

Pytorch的原生FP8

在2.2版本后,pytorch原生FP8支持已经是“有限支持”了,所以我们可以先学习一下如何使用了。

 import torch
 from tabulate import tabulate

 f32_type = torch.float32
 bf16_type = torch.bfloat16
 e4m3_type = torch.float8_e4m3fn
 e5m2_type = torch.float8_e5m2

 # collect finfo for each type
 table = []
 for dtype in [f32_type, bf16_type, e4m3_type, e5m2_type]:
     numbits = 32 if dtype == f32_type else 16 if dtype == bf16_type else 8
     info = torch.finfo(dtype)
     table.append([info.dtype, numbits, info.max, 
                   info.min, info.smallest_normal, info.eps])

 headers = ['data type', 'bits', 'max', 'min', 'smallest normal', 'eps']
 print(tabulate(table, headers=headers))

 '''
 Output:

 data type      bits          max           min  smallest normal          eps
 -------------  ----  -----------  ------------  ---------------  -----------
 float32          32  3.40282e+38  -3.40282e+38      1.17549e-38  1.19209e-07
 bfloat16         16  3.38953e+38  -3.38953e+38      1.17549e-38    0.0078125
 float8_e4m3fn     8          448          -448         0.015625        0.125
 float8_e5m2       8        57344        -57344      6.10352e-05         0.25
 '''

我们可以通过在张量初始化函数中指定dtype来创建FP8张量,如下所示:

 device="cuda"
 e4m3 = torch.tensor(1., device=device, dtype=e4m3_type)
 e5m2 = torch.tensor(1., device=device, dtype=e5m2_type)

也可以强制转换为FP8。在下面的代码中,我们生成一个随机的浮点张量,并比较将它们转换为四种不同的浮点类型的结果:

 x = torch.randn(2, 2, device=device, dtype=f32_type)
 x_bf16 = x.to(bf16_type)
 x_e4m3 = x.to(e4m3_type)
 x_e5m2 = x.to(e5m2_type)

 print(tabulate([[‘float32’, *x.cpu().flatten().tolist()],
                 [‘bfloat16’, *x_bf16.cpu().flatten().tolist()],
                 [‘float8_e4m3fn’, *x_e4m3.cpu().flatten().tolist()],
                 [‘float8_e5m2’, *x_e5m2.cpu().flatten().tolist()]],
                headers=[‘data type’, ‘x1’, ‘x2’, ‘x3’, ‘x4’]))

 '''
 The sample output demonstrates the dynamic range of the different types:

 data type                  x1              x2              x3              x4
 -------------  --------------  --------------  --------------  --------------
 float32        2.073093891143  -0.78251332044  -0.47084918620  -1.32557279110
 bfloat16       2.078125        -0.78125        -0.4707031      -1.328125
 float8_e4m3fn  2.0             -0.8125         -0.46875        -1.375
 float8_e5m2    2.0             -0.75           -0.5            -1.25
 -------------  --------------  --------------  --------------  --------------
 '''

虽然创建FP8张量很容易,但FP8张量上执行一些基本的算术运算是不支持的。并且需要特定的函数,比如torch._scaled_mm来进行矩阵乘法。

 output, output_amax = torch._scaled_mm(
         torch.randn(16,16, device=device).to(e4m3_type),
         torch.randn(16,16, device=device).to(e4m3_type).t(),
         bias=torch.randn(16, device=device).to(bf16_type),
         out_dtype=e4m3_type,
         scale_a=torch.tensor(1.0, device=device),
         scale_b=torch.tensor(1.0, device=device)
     )

那么如何进行模型的训练呢,我们来做一个演示

 import torch
 from timm.models.vision_transformer import VisionTransformer
 from torch.utils.data import Dataset, DataLoader
 import os
 import time

 #float8 imports
 from float8_experimental import config
 from float8_experimental.float8_linear import Float8Linear
 from float8_experimental.float8_linear_utils import (
     swap_linear_with_float8_linear,
     sync_float8_amax_and_scale_history
 )

 #float8 configuration (see documentation)
 config.enable_amax_init = False
 config.enable_pre_and_post_forward = False

 # model configuration controls:
 fp8_type = True # toggle to change floating-point precision
 compile_model = True # toggle to enable model compilation
 batch_size = 32 if fp8_type else 16 # control batch size

 device = torch.device('cuda')

 # use random data
 class FakeDataset(Dataset):
     def __len__(self):
         return 1000000
     def __getitem__(self, index):
         rand_image = torch.randn([3, 256, 256], dtype=torch.float32)
         label = torch.tensor(data=[index % 1024], dtype=torch.int64)
         return rand_image, label

 # get data loader
 def get_data(batch_size):
     ds = FakeDataset()
     return DataLoader(
            ds,
            batch_size=batch_size, 
            num_workers=os.cpu_count(),
            pin_memory=True
          )

 # define the timm model
 def get_model():
     model = VisionTransformer(
         class_token=False,
         global_pool="avg",
         img_size=256,
         embed_dim=1280,
         num_classes=1024,
         depth=32,
         num_heads=16
     )
     if fp8_type:
         swap_linear_with_float8_linear(model, Float8Linear)
     return model

 # define the training step
 def train_step(inputs, label, model, optimizer, criterion):
     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()
     if fp8_type:
         sync_float8_amax_and_scale_history(model)
     optimizer.step()


 model = get_model()
 optimizer = torch.optim.Adam(model.parameters())
 criterion = torch.nn.CrossEntropyLoss()
 train_loader = get_data(batch_size)

 # copy the model to the GPU
 model = model.to(device)
 if compile_model:
     # compile model
     model = torch.compile(model)
 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)

     # train step
     train_step(inputs, label, model, optimizer, criterion)

     # 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}')

这里需要特定的转换函数,将一些操作转换为支持FP8的版本,需要说明的是,因为还在试验阶段所以可能不稳定

FP8线性层的使用使我们的模型的性能比我们的基线实验提高了47%(!!)

对比TE

未编译的TE FP8模型的性能明显优于我们以前的FP8模型,但编译后的PyTorch FP8模型提供了最好的结果。因为TE FP8模块不支持模型编译。所以使用torch.compile会导致“部分编译”,即它在每次使用FP8时将计算分拆为多个图。

总结

在这篇文章中,我们演示了如何编写PyTorch训练脚本来使用8位浮点类型。TE是一个非常好的库,因为它可以让我们的代码修改量最小,而PyTorch原生FP8支持虽然需要修改代码,并且还是在试验阶段(最新的2.3还是在试验阶段),可能会产生问题,但是这会让训练速度更快。

不过总的来说FP8的确可以加快我们的训练速度,提高GPU的使用效率。这里要提一句TE是由NVIDIA开发的,并对其gpu进行了大量定制,所以如果是N卡的话可以直接用TE

https://avoid.overfit.cn/post/0dd1fba546674b48b932260fa8742971

相关实践学习
基于阿里云DeepGPU实例,用AI画唯美国风少女
本实验基于阿里云DeepGPU实例,使用aiacctorch加速stable-diffusion-webui,用AI画唯美国风少女,可提升性能至高至原性能的2.6倍。
目录
相关文章
|
1月前
|
存储 人工智能 PyTorch
基于PyTorch/XLA的高效分布式训练框架
基于PyTorch/XLA的高效分布式训练框架
186 2
|
4天前
|
机器学习/深度学习 算法 PyTorch
【从零开始学习深度学习】45. Pytorch迁移学习微调方法实战:使用微调技术进行2分类图片热狗识别模型训练【含源码与数据集】
【从零开始学习深度学习】45. Pytorch迁移学习微调方法实战:使用微调技术进行2分类图片热狗识别模型训练【含源码与数据集】
|
4天前
|
机器学习/深度学习 算法 PyTorch
【从零开始学习深度学习】44. 图像增广的几种常用方式并使用图像增广训练模型【Pytorch】
【从零开始学习深度学习】44. 图像增广的几种常用方式并使用图像增广训练模型【Pytorch】
|
4天前
|
机器学习/深度学习 PyTorch 算法框架/工具
【从零开始学习深度学习】36. 门控循环神经网络之长短期记忆网络(LSTM)介绍、Pytorch实现LSTM并进行训练预测
【从零开始学习深度学习】36. 门控循环神经网络之长短期记忆网络(LSTM)介绍、Pytorch实现LSTM并进行训练预测
|
4天前
|
机器学习/深度学习 PyTorch 算法框架/工具
【从零开始学习深度学习】35. 门控循环神经网络之门控循环单元(gated recurrent unit,GRU)介绍、Pytorch实现GRU并进行训练预测
【从零开始学习深度学习】35. 门控循环神经网络之门控循环单元(gated recurrent unit,GRU)介绍、Pytorch实现GRU并进行训练预测
|
4天前
|
机器学习/深度学习 PyTorch 算法框架/工具
【从零开始学习深度学习】18. Pytorch中自定义层的几种方法:nn.Module、ParameterList和ParameterDict
【从零开始学习深度学习】18. Pytorch中自定义层的几种方法:nn.Module、ParameterList和ParameterDict
|
4天前
|
机器学习/深度学习 PyTorch 算法框架/工具
【从零开始学习深度学习】17. Pytorch中模型参数的访问、初始化和共享方法
【从零开始学习深度学习】17. Pytorch中模型参数的访问、初始化和共享方法
|
5天前
|
机器学习/深度学习 并行计算 PyTorch
使用PyTorch Profiler进行模型性能分析,改善并加速PyTorch训练
加速机器学习模型训练是工程师的关键需求。PyTorch Profiler提供了一种分析工具,用于测量CPU和CUDA时间,以及内存使用情况。通过在训练代码中嵌入分析器并使用tensorboard查看结果,工程师可以识别性能瓶颈。Profiler的`record_function`功能允许为特定操作命名,便于跟踪。优化策略包括使用FlashAttention或FSDP减少内存使用,以及通过torch.compile提升速度。监控CUDA内核执行和内存分配,尤其是避免频繁的cudaMalloc,能有效提升GPU效率。内存历史记录分析有助于检测内存泄漏和优化批处理大小。
8 1
|
1月前
|
机器学习/深度学习 负载均衡 PyTorch
PyTorch分布式训练:加速大规模数据集的处理
【4月更文挑战第18天】PyTorch分布式训练加速大规模数据集处理,通过数据并行和模型并行提升训练效率。`torch.distributed`提供底层IPC与同步,适合定制化需求;`DistributedDataParallel`则简化并行过程。实际应用注意数据划分、通信开销、负载均衡及错误处理。借助PyTorch分布式工具,可高效应对深度学习的计算挑战,未来潜力无限。
|
1月前
|
机器学习/深度学习 并行计算 PyTorch
PyTorch与CUDA:加速深度学习训练
【4月更文挑战第18天】本文介绍了如何使用PyTorch与CUDA加速深度学习训练。CUDA是NVIDIA的并行计算平台,常用于加速深度学习中的矩阵运算。PyTorch与CUDA集成,允许开发者将模型和数据迁移到GPU,利用`.to(device)`方法加速计算。通过批处理、并行化策略及优化技巧,如混合精度训练,可进一步提升训练效率。监控GPU内存和使用调试工具确保训练稳定性。PyTorch与CUDA的结合对深度学习训练的加速作用显著。

热门文章

最新文章