社区编辑申请
注册/登录
不堆参数、不靠时长,Meta加速ViT训练流程,吞吐量4倍提升
人工智能 新闻
有时候,好的训练「技巧」比蛮力堆参更有效。

现阶段,视觉 transformer(ViT)模型已经在图像分类、目标检测与分割等各样各样的计算机视觉任务中得到了广泛应用,并可以在视觉表征与识别中实现 SOTA 结果。由于计算机视觉模型的性能往往与参数量和训练时长呈正相关,AI 社区已经实验了越来越大规模的 ViT 模型。

但应看到,随着模型开始超出万亿次浮点运算的规模,该领域已经遇到了一些主要的瓶颈。训练单个模型可能耗费数月,需要数以千块的 GPU,进而增加了加速器需求并导致大规模 ViT 模型将很多从业者「排除在外」。

为了扩展 ViT 模型的使用范围,Meta AI 的研究者已经开发出了更高效的训练方法。非常重要的一点是对训练进行优化以实现最佳的加速器利用。但是,这一过程耗时费力且需要大量的专业知识。为了设置有序的实验,研究者必须从无数可能的优化方案中进行选择:一次训练过程中执行的百万次运算中的任何一个都有可能受到低效率的影响和阻碍。

Meta AI 发现,通过将一系列优化应用到其图像分类代码库 PyCls 中的 ViT 实现,可以提升计算和存储效率。对于使用 PyCIs 训练的 ViT 模型,Meta AI 的方法可以提升训练速度和每加速器吞吐量(TFLOPS)。

下图展示了使用优化代码库 PyCIs 后每芯片(per chip)加速器吞吐量相较于 V100 基准的相对增加,而 A100 优化的加速器吞吐量是 V100 基准的 4.05 倍。

运行原理

Meta AI 首先对 PyCIs 代码库进行分析以确认低训练效率的潜在来源,最终将注意力放在了对数字格式的选择上。在默认情况下,大多数应用使用 32-bit 单精度浮点格式来表征神经网络值。转换至 16-bit 半精度格式(FP16)可以减少模型的内存占用和执行时间,但往往也会降低准确率。

研究者采取了折中方案,即混合精度。利用它,系统通过单精度格式执行计算以加速训练并减少内存使用,同时通过单精度存储结果以保持准确率。他们没有手动地将部分网络转换至半精度,而是实验了不同模式的自动混合精度训练,这样可以在数字格式之间自动切换。更高级模式的自动混合精度主要依赖半精度运算和模型权重。研究者采用的平衡设置既能大幅度加速训练,同时也不牺牲准确率。

为了使流程更加高效,研究者充分利用了 FairScale 库中的完全分片数据并行(Fully Sharder Data Parallel, FSDP)训练算法,它在 GPU 上对参数、梯度和优化器状态进行分片。通过 FSDP 算法,研究者可以使用更少的 GPU 构建更大量级的模型。此外,研究者还使用了 MTA 优化器、一个池化的 ViT 分类器和一个 batch-second 输入张量布局来跳过冗余转置运算。

下图 X 轴为可能的优化,Y 轴为采用 ViT-H/16 训练时加速器吞吐量相较于分布式数据并行(DDP)基准的相对增加。

研究者在总 patch 大小为 560 时实现了 1.51 倍的加速器吞吐量提升,以每个加速器芯片上每秒执行的浮点运算数量衡量。通过将图像大小从 224 像素增加至 256 像素,他们可以将吞吐量提升至 1.86 倍。但是,改变图像大小意味着超参数的变化,这会对模型的准确率造成影响。在完全 FP16 模式下训练时,相对吞吐量增加至 2.18 倍。尽管有时会降低准确率,但在实验中准确率降低少于 10%。

下图 Y 轴为 epoch 时间,在整个 ImageNet-1K 数据集上一次训练的持续时间。这里专注于现有配置的实际训练时间,这些配置通常使用 224 像素的图像大小。

Meta AI 的研究者使用优化方案,将 epoch 时间(在整个 ImageNet-1K 数据集上一次训练的持续时间)从 0.65 小时减少到 0.43 小时。

下图 X 轴表示特定配置中 A100 GPU 加速器芯片的数量,Y 轴表示每芯片 TFLOPS 的绝对吞吐量。

该研究还讨论了不同 GPU 配置的影响。在每种情况下,系统都实现了比分布式数据并行(DDP)基线水平更高的吞吐量。随着芯片数量的增加,由于设备间通信的开销,我们可以观察到吞吐量略有下降。然而,即使用 64 块 GPU,Meta 的系统也比 DDP 基准快 1.83 倍。

新研究的意义

将 ViT 训练中可实现的吞吐量翻倍可以有效让训练集群规模翻倍,提高加速器利用率直接减少了 AI 模型的碳排放。由于最近大模型的发展带来了更大模型和更长训练时间的趋势,这种优化有望帮助研究领域进一步推动最先进的技术,缩短周转时间并提高生产力。

责任编辑:张燕妮 来源: 机器之心
相关推荐

2022-06-20 22:37:25

Linux操作系统命令

2022-06-20 06:27:58

Meta人工智能AI模型

2022-06-27 17:46:53

PythonFlask

2022-06-07 09:59:21

网络安全安全漏洞

2022-06-16 15:42:16

攻击面管理ASM

2022-07-01 14:25:27

机器学习人工智能工业4.0

2022-06-28 09:26:25

Python配置文件

2022-06-23 12:03:00

网络安全网络安全事故

2022-06-03 09:41:03

DockerKubernetes容器

2022-05-05 08:25:22

模型OpenAI代码

2022-06-23 14:03:26

混合ITCIOIT管理工具

2022-07-01 18:24:36

勒索软件病毒网络攻击

2022-07-01 05:47:19

PyCharm插件开发

2022-06-23 09:22:57

Vue技巧前端

2022-06-27 19:01:04

Python应用程序数据

2022-06-29 08:13:36

漏洞网络攻击网络安全

2022-06-29 10:16:25

数据库SQL

2022-06-16 07:32:38

VSCodePython插件

2022-06-22 09:19:55

HDC鸿蒙ADB命令

2022-06-24 11:14:00

美团开源

同话题下的热门内容

李飞飞划重点的「具身智能」,走到哪一步了?新出生的机器狗,打滚1小时后自己掌握走路,吴恩达大弟子成果聊聊七个常见的智能汽车技术2022年Nature年度指数出炉,涨幅最快50机构,中国占31席!利用机器学习发起攻击的九种方式监督学习有哪些常见算法?都是如何应用的深度报告:大模型驱动 AI 全面提速!黄金十年开启2022年人工智能机器人的五大趋势

编辑推荐

转转公司架构算法部孙玄:AI下的微服务架构Facebook开源相似性搜索类库Faiss,超越已知最快算法8.5倍运维:对不起,这锅,我们不背快消品图像识别丨无人店背后的商品识别技术最全面的百度NLP自然语言处理技术解析
我收藏的内容
点赞
收藏

51CTO技术栈公众号