当前位置: 首页 > 科技观察

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

时间:2023-03-13 17:35:20 科技观察

Meta在不叠加参数和时间的情况下加速了ViT训练过程,吞吐量提高了4倍,已经得到广泛应用,在视觉表示和识别方面可以达到SOTA的效果。由于计算机视觉模型的性能往往与参数数量和训练时间呈正相关,因此AI社区已经尝试了越来越大的ViT模型。然而,应该指出的是,随着模型开始扩展到teraflops以上,该领域遇到了一些主要瓶颈。训练单个模型可能需要数月时间并需要数千个GPU,增加了加速器要求并导致大规模ViT模型“超出”许多从业者的范围。为了扩大ViT模型的使用范围,MetaAI研究人员开发了更高效的训练方法。优化训练以获得最佳加速器利用率非常重要。但是,此过程非常耗时并且需要大量专业知识。为了建立有序的实验,研究人员必须从无数可能的优化中进行选择:在训练期间执行的数百万次操作中的任何一项都可能因效率低下而受到阻碍。MetaAI发现,通过在其图像分类代码库PyCls中对ViT实现应用一系列优化,可以提高计算和存储效率。对于使用PyCI训练的ViT模型,MetaAI的方法提高了训练速度和每个加速器的吞吐量(TFLOPS)。下图显示了使用优化代码库PyCIs后每芯片(perchip)加速器吞吐量相对于V100baseline的相对提升,而A100优化后的加速器吞吐量是V100baseline的4.05倍。它是如何工作的MetaAI首先分析了PyCIs代码库,以确定训练效率低下的潜在来源,最后重点关注数字格式的选择。默认情况下,大多数应用程序使用32位单精度浮点格式来表示神经网络值。转换为16位半精度格式(FP16)可减少模型的内存占用和执行时间,但通常也会降低准确性。研究人员采取了折衷方案,即混合精度。有了它,系统以单精度格式执行计算以加速训练并减少内存使用,同时以单精度格式存储结果以保持准确性。他们没有手动将部分网络转换为半精度,而是尝试了不同的自动混合精度训练模式,可以在数字格式之间自动切换。更高级模式的自动混合精度主要依赖于半精度运算和模型权重。研究人员使用的平衡设置可以在不牺牲准确性的情况下大大加快训练速度。为了提高流程效率,研究人员利用了FairScale库中的FullySharderDataParallel(FSDP)训练算法,该算法在GPU上对参数、梯度和优化器状态进行分片。通过FSDP算法,研究人员可以使用更少的GPU来构建更大规模的模型。此外,研究人员还使用了MTA优化器、池化ViT分类器和batch-second输入张量布局来跳过冗余的转置操作。下图显示了与Y轴上的分布式数据并行(DDP)基准相比,使用ViT-H/16进行训练时X轴上可能的优化以及加速器吞吐量的相对增加。总补丁大小为560,研究人员将加速器吞吐量提高了1.51倍,以每个加速器芯片上每秒执行的浮点运算数来衡量。通过将图像大小从224像素增加到256像素,他们能够将吞吐量提高1.86倍。然而,改变图像大小意味着改变超参数,这将影响模型的准确性。在全FP16模式下训练时,相对吞吐量增加到2.18倍。虽然偶尔会出现精度下降,但在实验中精度下降不到10%。下图的Y轴是纪元时间,即整个ImageNet-1K数据集上训练会话的持续时间。在这里,我们关注现有配置的实际训练时间,这些配置通常使用224像素的图像大小。MetaAI的研究人员使用优化将纪元时间(整个ImageNet-1K数据集上一次训练的持续时间)从0.65小时减少到0.43小时。下图在X轴上显示了特定配置中A100GPU加速器芯片的数量,在Y轴上显示了每个芯片的TFLOPS绝对吞吐量。该研究还讨论了不同GPU配置的影响。在每种情况下,系统都实现了比分布式数据并行(DDP)基线更高的吞吐量。随着芯片数量的增加,我们可以观察到由于设备间通信的开销导致吞吐量略有下降。然而,即使使用64个GPU,Meta的系统也比DDP基准测试快1.83倍。新研究的意义将ViT训练中可实现的吞吐量提高一倍,可以有效地将训练集群的规模扩大一倍,提高加速器利用率直接减少AI模型的碳足迹。由于最近的趋势是更大的模型和大型模型的开发带来的更长的训练时间,这种优化有望帮助研究界进一步推动最先进的技术,减少周转时间并提高生产力。