微软的ZeRO-Infinity库训练出32万亿个参数的AI模型

微软最近宣布了ZeRO-Infinity,作为他们的开源DeepSpeed人工智能训练库的一个补充,优化了训练非常大的深度学习模型的内存使用。微软使用ZeRO-Infinity在一个由32个GPU组成的集群上,训练了一个具有32万亿个参数的模型,并且在单个GPU上对1万亿个参数的模型进行微调。

DeepSpeed团队在最近的一篇博文中 描述了这些新功能。ZeRO-Infinity是零冗余优化器(ZeRO)系列内存优化技术的最新迭代。ZeRO-Infinity为了解决训练大型深度学习模型时的内存和带宽限制,引入了几个新的策略,包括: 一个用于利用CPU和非易失性内存表达(NVMe)内存的新卸载引擎;以内存为中心的平铺,处理没有模型并行的大型运算器;以带宽为中心的分区,减少带宽成本;以及以重叠为中心的设计,用于调度数据通信。根据DeepSpeed团队的说法:

改进后的ZeRO-Infinity提供了超越GPU内存墙的系统能力,可以训练具有几十万亿个参数的模型,比最先进的系统可以支持的数量级要大。它还为训练100万亿个参数的模型提供了一条有希望的道路。

最近深度学习研究的一个趋势是在更多的数据上训练更大的模型,其中一些最大的模型在某些任务上实现了超人的表现。然而,训练这些模型需要大型且昂贵的GPU集群。许多情况下,模型开发者可以使用转移学习,只需要使用预训练所需的一小部分计算资源,就可以微调一个大型的预训练模型;但是,像GPT-3这样的非常大的模型在一台机器上微调起来就显得过于庞大。这两种情况通常都需要重构代码以利用分布式训练框架。

为了解决这些问题,微软在2020年初首次发布了DeepSpeed库和零冗余优化器(ZeRO),作为其AI at Scale计划的一部分。ZeRO分三个阶段进行改进,每个阶段都增加了模型状态的额外分区,以及将数据和计算从GPU "卸载 "到训练机CPU的能力。第三阶段于今年年初发布,能够在单台机器上训练多达400亿个参数的模型,在512个GPU的集群上训练超过2万亿个参数。

ZeRO的最新迭代,ZeRO-Infinity,带来了解决训练大型模型的两个瓶颈的新方案:内存大小和内存带宽。infinity卸载引擎通过使用CPU和NVMe内存增加了可用于存储模型参数和激活的内存量;与前几代ZeRO不同,infinity引擎可以将整个模型卸载到这些位置。以内存为中心的平铺是另一项新技术,它通过将大的模型层分解成较小的 "平铺 "来减少对内存的占用,这些平铺可以按顺序执行;这允许在不需要模型并行的情况下训练大型模型。为了处理带宽问题,ZeRO-Infinity引入了以带宽为中心的分区,将模型参数划分到多个数据并行进程中,还有一个重叠引擎,同时执行NVM到CPU、CPU到GPU以及GPU到GPU的通信。

该团队进行了几个实验来验证ZeRO-Infinity的扩展能力,训练不同规模的 "类似GPT "的变形器模型。与最先进的3D并行框架相比,ZeRO-Infinity使用相同的计算硬件处理的模型大了40倍。与上一版本的ZeRO相比,新版本在64个GPU集群上实现了2倍的速度提升。当训练一个1T参数的模型时,ZeRO-Infinity在从64个GPU到512个GPU的各种集群规模上进行了超线性扩展。

DeepSpeed库,包括ZeRO系列的内存优化,是为PyTorch深度学习框架编写的,并且已经被其他几个基于PyTorch的项目所采用。HuggingFace是一个流行的预训练人工智能模型的来源,已经与新的ZeRO-Infinity版本集成,PyTorch Lightning是PyTorch的分布式训练包装器,也采用DeepSpeed和ZeRo的前三个阶段。Facebook用于训练大型PyTorch模型的FairScale库也包括一些ZeRO技术。

Reddit上的讨论中,一位评论者将DeepSpeed库描述为 "难以置信的价值"。另一位指出。

然而,这些技术(零卸载、零无限)通常对帮助训练更大的模型没有什么帮助。对于从头开始训练真正的大型模型,内存通常不是瓶颈,计算才是。然而,这些技术对于帮助你进行微调是很有帮助的。

DeepSpeed库代码可在GitHub上找到。

原文链接:www.infoq.com/news/2021/0…