不堆参数、不靠时长，Meta加速ViT训练流程，吞吐量4倍提升

来源：网络浏览：71 2022-11-10 10:25:01

机心报道

机心编辑部

有时，好的训练“技术”比用力参加更有效。

目前，视觉transformer(vit )模型已经广泛应用于图像分类、目标检测和分割等各种计算机视觉任务中，能够通过视觉表示和识别实现SOTA结果。由于计算机视觉模型的性能往往与参数量和训练时间呈正相关，AI社区正在实验越来越大规模的ViT模型。

但是随着模型开始超过万亿次浮点运算的规模，我们可以看到这个领域已经遇到了几个主要瓶颈。训练一个模型需要几个月，可能需要几千个GPU。这将导致对加速器的需求增加，在大型ViT模型中“排除”了许多员工。

为了扩大ViT模型的使用范围，元ai的研究者开发了更高效的训练方法。非常重要的是优化培训以实现最佳的加速器使用。但是，这个过程需要时间和精力，需要很多专业知识。为了设置有序的实验，研究者必须在无数可能的优化方案中做出选择。在一次训练中执行的百万次运算中的任何一次都可能受到低效率的影响和阻碍。

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

下图显示了使用基于优化代码的PyCIs后，与V100基准相比，每芯片的加速器吞吐量相对增加。另一方面，A100优化了加速器吞吐量，是V100基准的4.05倍。

运行原理

Meta AI首先分析PyCIs代码库以确认训练效率低的潜在来源，最终关注数字格式的选择。缺省情况下，大多数APP应用程序使用32位单精度浮点格式来表示神经网络值。转换为16位半精度格式( FP16 )会减少模型的内存消耗和运行时间，但通常会降低精度。

研究人员采用了折中方案，即混合精度。利用它，系统以单精度格式执行计算，在加快训练、减少内存使用的同时，以单精度保存结果并保持准确性。他们没有手动将部分网络转换为半精度，而是实验了不同模式的自动混合精度训练，使得数字格式之间可以自动切换。更高级模式的自动混合精度主要取决于半精度运算和模型权重。研究人员采用的平衡设置可以在不牺牲准确性的情况下大大加速训练。

为了使流程更有效率，研究人员对FairScale库中的完整分片数据并行处理( fsdp；利用fullysharderdataparallel )训练算法，在GPU上分片参数、梯度和优化器状态。通过FSDP算法，研究者可以用更少的GPU构建更大量水平的模型。另外，使用MTA优化器、池化的ViT分类器、batch-second输入张量布局，跳过了冗长的转置运算。

下图中的x轴表示可能的优化，y轴表示与分布式数据并行( DDP )基准相比，采用ViT-H/16训练时的加速器吞吐量相对增加。

研究人员在总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，元系统也比DDP标准快1.83倍。

新研究的意义

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

原文链接： https://ai.Facebook.com/blog/significantly-faster-vision-transformer-training /