标题:微软|ZeRO-Infinity和DeepSpeed:解锁用于深度学习训练的前所未有的模型规模

作者: DeepSpeed Team , Rangan Majumder , Andrey Proskurin

简介:自从去年推出DeepSpeed优化库以来,它已经针对训练大型AI模型推出了许多新颖的优化方法,从而提高了规模,速度,成本和可用性。随着大型模型在过去的一年中迅速发展,DeepSpeed也是如此。使研究人员能够以最先进的准确性创建170亿参数的Microsoft Turing自然语言生成(Turing-NLG),实现最快的BERT训练记录,还是使用单个GPU DeepSpeed支持10倍的大型模型训练继续通过大规模模型训练的最新进展来应对大规模AI方面的挑战。现在,DeepSpeed中包含的新型内存优化技术ZeRO(零冗余优化器)正在对其自身进行进一步的改造。改进的ZeRO-Infinity提供了超越GPU内存墙的系统功能,并可以训练具有数十万亿个参数的模型,这比最新系统所能支持的数量级大。它还为训练100万亿参数模型提供了一条有希望的途径。

ZeRO-Infinity概览:ZeRO-Infinity是一种新颖的深度学习(DL)训练技术,用于大尺度模型训练,从单个GPU到具有数千个GPU的大型超级计算机。它通过利用系统的全部内存容量并同时利用所有异构内存(GPU,CPU和非易失性内存Express或简称NVMe)来提供前所未有的模型大小。在我们的论文“ ZeRO-Infinity:突破GPU内存壁以实现超大规模深度学习”中了解更多信息。 ZeRO-Infinity的亮点包括:

  • 提供系统功能,可在512个NVIDIA V100 Tensor Core GPU上训练具有30万亿个参数的模型,比现有技术大50倍。
  • 通过新颖的数据分区和映射功能,可利用合计的CPU / NVMe内存带宽和CPU计算能力,提供出色的训练效率和超线性吞吐量扩展,在512个NVIDIA V100 GPU上提供超过25 petaflops的持续吞吐量。
  • 通过允许具有单个GPU的数据科学家微调比Open AI GPT-3(1750亿个参数)大的模型,DeepSpeed团队进一步履行使命,使大型模型训练民主化。
  • 通过简化和简化,消除了大型模型训练的入门门槛-ZeRO-Infinity可以扩展超过一万亿个参数,而无需组合多种并行技术的复杂性,也无需更改用户代码。据我们所知,这是唯一可以做到这一点的并行技术。

我们也很高兴宣布DeepSpeed与Azure机器学习和开源解决方案的集成。 Azure机器学习中的DeepSpeed精选环境使用户可以更轻松地开始使用Azure。 DeepSpeed现在已集成在Hugging Face v4.2和PyTorch Lightning v1.2中。 HuggingFace和PyTorch Lightning用户可以通过一个简单的“ deepspeed”标志通过DeepSpeed轻松加速其模型!

 

使用ZeRO-Infinity解决当前和未来大型模型训练的需求


在过去的三年中,最大的训练密集模型已增长了1,000倍以上,从BERT前时代的一亿个参数发展到现在的一千亿个参数。但是,在相同的持续时间内,单个GPU内存仅增加了5倍(从16 GB到80 GB)。因此,增加模型尺寸的方法主要是通过用于训练大型DL模型的系统技术的进步而实现的,同时借助并行技术(例如模型并行性,管线并行性和ZeRO),大型模型可以装入聚合GPU内存中,从而为训练更大,更强大的模型。

大型模型训练技术中的最新技术是3D并行性。它以复杂的方式将模型并行性(张量切片)和管道并行性与数据并行性结合在一起,以通过充分利用聚集的GPU内存和集群计算来有效地扩展模型。 3D并行性已在DeepSpeed和NVIDIA Megatron-LM等框架中使用。

尽管3D并行性在大型模型训练中具有令人难以置信的功能,但我们现在到达了GPU内存壁。总的GPU内存根本不足以支持模型大小的增长。即使使用具有80 GB内存的最新NVIDIA A100 GPU,3D并行性也需要320个GPU才能满足训练所需的万亿参数模型。此外,3D并行性要求数据科学家进行大量的代码重构,这为进入提供了很大的障碍。出现三个问题:

展望未来,我们如何支持模型规模的下一个1,000倍增长,从具有1750亿个参数的GPT-3之类的模型到具有数百万亿个参数的模型?
着眼于现在,我们如何才能使更多的数据科学家可以访问当今的大型模型,而他们可能无法访问数百个适合这些模型的GPU?
 通过消除对模型重构的需求,我们是否可以使大型模型训练变得更容易?
今天,我们引入了ZeRO-Infinity,这是一种新颖的系统,能够应对3D并行性的飞跃,该系统能够解决大型模型训练的所有上述挑战。 ZeRO-Infinity通过在数据映射和高性能异构内存访问方面的新创新扩展了ZeRO技术家族,这使得ZeRO-Infinity通过同时利用CPU和NVMe内存而不受带宽有限的限制,可以在有限的GPU资源上支持大规模模型。 。

ZeRO-Infinity还可以训练这些模型,而无需在3D并行性中组合多种形式的并行性。它通过一种新颖的以内存为中心的计算平铺方法来实现这一目的,该方法旨在减少大型单个层的GPU内存需求,否则该需求可能需要模型并行性(张量切片)才能使模型适合GPU内存。此外,ZeRO-Infinity通过识别和自动化训练任意模型架构所需的所有通信,使大型模型训练变得容易,即使在扩展至数万亿参数时,也几乎消除了对任何模型重构的需求。最后但并非最不重要的一点是,ZeRO-Infinity提供了强大的计算和通信重叠引擎,旨在通过隐藏尽可能多的通信等待时间来将训练效率推到极限。

通过所有这些创新,ZeRO-Infinity重新定义了DL系统的功能,提供了前所未有的模型比例,该模型比例可访问且易于使用,同时实现了出色的训练效率。

 

前所未有的模型规模:在512个GPU上训练30万亿参数模型


ZeRO-Infinity在DL训练系统技术上实现了数量级的飞跃,通过有效利用当前和未来硬件上的异构存储系统,开辟了一条路径,可支持模型规模下一个1,000倍的增长。 它在单个NVIDIA DGX-2节点上运行具有超过1万亿个参数的模型,在32个节点(512个GPU)上具有超过30万亿个参数的模型。 通过集群中的一百个DGX-2节点,我们预计ZeRO-Infinity可以训练具有一百万亿个参数的模型。 (有关详细信息,请参见图1)。

图1

为了实现这种规模的模型训练,ZeRO-Infinity通过针对不同内存瓶颈的独特创新扩展了ZeRO技术家族。

1. ZeRO(ZeRO-3)的第3阶段允许通过跨数据并行过程划分模型状态来消除数据并行训练中的所有内存冗余。


2. Infinity Offload Engine是一个新颖的数据卸载库,通过将分区模型状态卸载到比GPU内存大得多的CPU或NVMe设备内存中,可以完全利用现代异构内存体系结构。


3.带有CPU卸载的激活检查点可以减少激活内存占用空间,在ZeRO-3和Infinity Offload Engine解决了模型状态所需的内存之后,这可能成为GPU上的内存瓶颈。

4.以内存为中心的运算符平铺,这是一种新颖的计算重新调度技术,可与ZeRO数据访问和通信调度配合使用,可减少难以置信的庞大单个层的内存占用,这些单个层可能太大而无法容纳GPU内存,即使是一个时间。

更广泛地访问微调的超大型模型:单个GPU上的GPT-3甚至更大的模型

尽管预训练是创建大规模模型的第一步,但对特定任务进行微调对于充分利用模型在不同情况下的潜力至关重要。使数据科学家易于访问海量模型的微调,可以允许创建许多派生模型,以满足各种应用场景的需求。这些任务的范围可能从语法校正到写作帮助,从图像字幕到代码生成-大型AI模型可能完成的所有任务。

与可能需要数百万个GPU计算小时的预训练不同,对具有数千亿个参数的模型进行微调要便宜得多,所需的GPU计算时间要少得多,并且可以在具有少量GPU的单个计算节点上完成。尽管许多企业和用户都可以访问此类计算资源,但不幸的是,这些计算资源受到这些计算节点上可用内存的限制,从而限制了可以微调的模型的大小。它使大多数无法访问大型GPU群集的企业和公司都无法访问大型模型的微调。

通过使数据科学家能够访问单个节点(例如NVIDIA DGX-2),ZeRO-Infinity彻底改变了这一格局,从而微调了超过一万亿个参数的模型(图4)。实际上,由于它具有足够的CPU和NVMe内存,因此即使在这样一个节点的单个GPU上,它也可以运行超过一万亿个参数的模型。这比单GPU训练的最新技术高出近100倍。使用ZeRO-Infinity,内存瓶颈不再是GPU内存,甚至不再是CPU内存。相反,我们现在可以将它们与更大,更便宜的NVMe内存一起使用。

通过ZeRO-Infinity,我们使AI大众化迈出了又一步,这使资源有限的用户和企业能够将海量模型的强大功能用于其特定于业务的应用程序。

图3

图4

 

代码:https://github.com/microsoft/DeepSpeed

论文下载:https://arxiv.org/pdf/2104.07857

新闻出处:https://www.microsoft.com/en-us/research/blog/zero-infinity-and-deepspeed-unlocking-unprecedented-model-scale-for-deep-learning-training/