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

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:/a/159656.html

如若内容造成侵权/违法违规/事实不符,请联系我们进行投诉反馈qq邮箱809451989@qq.com,一经查实,立即删除!

相关文章

【Linux】线程控制

文章目录 线程的概念Linux下的进程Linux下的线程进程再理解Linux线程和接口的认识代码验证二级页表 页表线程的优点线程的缺点线程异常 线程的用途进程和线程的关系线程控制线程线程ID和LWP线程等待线程终止线程分离 线程ID及进程地址空间布局 线程的概念 我们知道&#xff0c…

全新小权云黑系统

小权云黑管理系统 V1.0 功能如下: 1.添加骗子,查询骗子 2.可添加团队后台方便审核用 3.在线反馈留言系统 4.前台提交骗子,后台需要审核才能过 5.后台使用光年UI界面 6.新增导航列表,可给网站添加导航友链 7.可添加云黑类型收录 8.…

基于STC12C5A60S2系列1T 8051单片的IIC总线器件模数芯片PCF8591实现模数转换应用

基于STC12C5A60S2系列1T 8051单片的IIC总线器件模数芯片PCF8591实现模数转换应用 STC12C5A60S2系列1T 8051单片机管脚图STC12C5A60S2系列1T 8051单片机I/O口各种不同工作模式及配置STC12C5A60S2系列1T 8051单片机I/O口各种不同工作模式介绍IIC总线器件模数芯片PCF8591介绍通过I…

注解方式优雅的实现 Redisson 分布式锁

1前言 日常开发中,难免遇到一些并发的场景,为了保证接口执行的一致性,通常采用加锁的方式,因为服务是分布式部署模式,本地锁Reentrantlock和Synchnorized这些就先放到一边了,Redis的setnx锁存在无法抱保证…

asp.net健身会所管理系统sqlserver

asp.net健身会所管理系统sqlserver说明文档 运行前附加数据库.mdf(或sql生成数据库) 主要技术: 基于asp.net架构和sql server数据库 功能模块: 首页 会员注册 教练预约 系统公告 健身课程 在线办卡 用户中心[修改个人信息 修…

Modbus RTU 使用教程(modbus教程、modbus协议)

参考文章:这节课带你吃透Modbus通信协议 文章目录 Modbus RTU 使用教程第1部分:Modbus RTU 协议概述什么是Modbus RTU协议Modbus RTU的特点工作原理 第2部分:硬件和软件要求硬件要求软件要求 第3部分:Modbus RTU数据结构数据帧格式…

android适配鸿蒙系统开发

将一个Android应用迁移到鸿蒙系统需要进行细致的工作,因为两者之间存在一些根本性的差异,涉及到代码、架构、界面等多个方面的修改和适配。以下是迁移工作可能涉及的一些主要方面,希望对大家有所帮助。北京木奇移动技术有限公司,专…

各品牌PLC元件在modbus内区域

1台达: 输出在0区, 040961是在 0区 0xA000~0xA0FF 【Y0~Y377】 输入在1区,124577是在 1区 0x6000~0x60FF 【X0~X377】 M寄存器0区,0000001是 0区,0x000~0x1FFF 【M0~M8191】 D寄存器4区,400000…

CTFhub-RCE-过滤运算符

检查网页源代码发现如果用户输入 || (逻辑或操作符) 或者是 & (按位与操作符),就不会执行。直接进行测试,在 URL 后输入本地 IP 进行 ping 命令测试,发现可以正常回显。检查网页源代码发现如…

2018年五一杯数学建模B题商业银行人民币贷款规模分配及盈利问题解题全过程文档及程序

2019年五一杯数学建模 B题 商业银行人民币贷款规模分配及盈利问题 原题再现 商业银行贷款投放的简单模型是:从客户端吸收存款,缴存法定准备金(法定准备金率:大型金融机构15.5%,中小金融机构12%;法定准备金…

Spring Boot Actuator:自定义端点

要在Spring Boot Actuator中实现自定义端点,可以按照以下步骤进行操作: 1.创建一个自定义端点类 该类需要使用Endpoint注解进行标记,并使用Component注解将其作为Spring Bean进行管理。 package com.example.highactuator.point;import lo…

flowable消息事件

一&#xff1a;启动事件 定义消息。 引用消息。 <startEvent id"msgStart" name"消息启动事件" isInterrupting"true"><messageEventDefinition messageRef"myMsgStart"></messageEventDefinition> </startE…

机器学习第6天:线性回归模型正则化

文章目录 机器学习专栏 正则化介绍 岭回归 岭回归成本函数 核心代码 示例 Lasso回归 Lasso回归损失函数 核心代码 弹性网络 弹性网络成本函数 核心代码 结语 机器学习专栏 机器学习_Nowl的博客-CSDN博客 正则化介绍 作用&#xff1a;正则化是为了防止模型过拟合…

es使用客户端,“grunt” 不是内部或外部命令,多种解决方法

”grunt“不是内部或外部命令&#xff0c;也不是可运行的程序 或批处理文件。 4、问题排查 查看node的安装根目录 npm root -g 在运行grunt -version还是不行 网上找了很多&#xff0c;给出正确解决方案的没几个&#xff0c;所以自己摸索&#xff0c;最后确定了加环境变量的解…

Python实验项目7 :tkinter GUI编程

&#xff08;1&#xff09;利用tkinter 制作界面&#xff0c;效果图如下&#xff1a; from tkinter import * # winTk() for i in range(1,20):Button(width5,height10,bg"black" if i%20 else"white").pack(side"left") win.geometry("8…

Android 屏幕适配

目录 一、为什么要适配 二、几个重要的概念 2.1 屏幕尺寸 2.2 屏幕分辨率 2.3 屏幕像素密度 2.4 屏幕尺寸、分辨率、像素密度三者关系 三、常用单位 3.1 密度无关像素(dp) 3.2 独立比例像素&#xff08;sp&#xff09; 3.3 dp与px的转换 四、解决方案 4.1 今日头条…

Docker之DockerFile解析

DockerFile解析 是什么 Dockerfile是用来构建Docker镜像的文本文件&#xff0c;是由一条条构建镜像所需的指令和参数构成的脚本。 概述 官网 https://docs.docker.com/engine/reference/builder/ 构建三步骤 编写Dockerfile文件 docker build命令构建镜像 docker run依镜像运…

李沐的学习Pytorch环境配置

https://github.com/Miraclelucy/dive_into_deep_learning/tree/main 上面是别人的笔记 可以学一下。 如果没有梯子&#xff0c;按照清华源配置 清华源conda配置 最好下载 1.11版本torch那一套 然后装d2l版本可以装 pip install d2l0.17.6然后可以用 http://localhost:8889/…

DolphinScheduler V3.1.8 海豚调度器【概述、安装、功能介绍、项目运用、邮箱预警设置】轻松拿捏!

概述 Apache DolphinScheduler 是一个分布式易扩展的可视化 DAG 工作流任务调度开源系统。适用于企业级场景&#xff0c;提供了一个可视化操作任务、工作流和全生命周期数据处理过程的解决方案。 Apache DolphinScheduler 旨在解决复杂的大数据任务依赖关系&#xff0c;并为应…

Diffusion Models CLIP

Introduction to Diffusion Models 生成模型 主要指的是无监督学习中的生成模型&#xff0c;在无监督学习中的主要任务是让机器学习给定的样本&#xff0c;然后生成一些新的东西出来。比如&#xff1a;给机器看一些图片&#xff0c;能够生成一些新的图片出来&#xff0c;给机器…