本文将详细解读微软与OpenAI合作推出的ZeRO-2框架、新的ONNX Runtime和超级计算机。

使用海量数据训练模型，已成为现代深度学习应用的一种常态。以Transformers及其它预训练模型为代表的一些最新的创新性应用，需要特大型训练框架的支持。此类预训练模型通常被视为人工智能领域的最新突破，由此激发了全球范围引领该领域的一些顶级技术企业间的竞争。在最近的Build大会上，微软展示了自身在支持大规模构建和使用大规模预训练模型方面的一些最新举措。

深度学习领域的最新进展，主要围绕着预训练模型和Transformers。这场运动源自于Google提出的著名BERT模型，OpenAI GPT-2、微软Turin-NLG等特大型架构随后迅速跟进。建立此类模型，不仅需要占用大量计算资源，而且通常还需要在架构上针对大规模训练操作做出优化。由于其新颖性，主流深度学习框架尚未包含此类架构。微软的最新举措，全面代表了深度学习社区在简化创建和访问大规模模型上的进展。下面择其重点进行概述。

Zero-2和DeepSpeed

微软于今年年初开源了DeepSpeed的首个版本。DeepSpeed是一种优化大规模训练的框架，是支持微软Turing-NLG模型打破记录的后台基础框架。在发布版中，包括了一个称为“ZeRO”（零冗余优化器，Zero Redundacy Optimizer）的优化模块。从概念上讲，ZeRO寻求模型并行和数据并行的中间点，致力于最大化模型的可扩展性。ZeRO的最大优点，是无需重构代码即可支持在内存中运行大型深度学习模型。这是通过优化模型内存的使用方式而实现的。就此而言，ZeRO可视为一种深度学习模型的内存优化方法。

作为ZeRO的第二版，ZeRO-2相比第一版做了多处创新。其中涉及深度学习模型内存使用的多个方面，包括激活内存、碎片内存和模型状态内存。

模型状态内存（Model State Memory）： ZeRO-2优化了保存模型状态的内存，达到相比传统的数据并行方法降低内存使用8倍。所有深度学习模型的状态，可归为分区优化器状态、梯度和参数这三个基本过程。在ZeRO第一版中，优化了分区优化器状态的内存使用，达到比传统方法降低内存使用4倍。ZeRO-2在该方法上做了扩展，优化了梯度计算过程的内存占用，进一步降低内存使用2倍。
激活内存（Activation Memory）：即便优化了模型状态内存，激活函数依然可导致瓶颈。激活函数计算位于前向过程中，用于支持后向过程。ZeRO-2引入了一种将计算过程加载到主机CPU的激活分区技术，消除了内存瓶颈隐患。
碎片内存（Fragmented Memory）：在一些情况下，深度学习模型的低效是由于碎片所导致的，这些碎片的生成，源自于模型中各个张量的不同生命周期。ZeRO-2通过管理不同张量生命周期的内存使用全过程，消除了低效问题。

不同于其它方法中大多通过将训练分布在各个GPU上进行优化，ZeRO-2的一个最显著创新就是支持单个GPU粒度的内存优化。ZeRO-2团队展示了使用单个GPU达到在44分钟内训练Google BERT版本，明确给出该优化的价值所在。下图是ZeRO-2对比以往方法的改进情况。

此外，如下图所示，在支持具有数十亿参数模型的能力、实现更高通量和高速度、每个GPU和跨GPU集群上的扩展性、单个GPU上训练超大模型等方面，ZeRO-2同样展现了其优势所在。

ZeRO-2 + ONNX Runtime

ONNX Runtime是一种支持不同机器学习框架间互操作的开源框架。机器学习互操作性的一个关注点，在于其是否适用于超大模型。ONNX可直接用于决策树，但其是否同样可用于具有数十亿参数的模型？

为优化ONMX模型的训练和内存使用，微软决定在ONNX Runtime中添加DeepSpeed和ZeRO-2。下图展示了在一些当前全球最大的预训练模型上的测试结果，包括在Visual Studio中的GPT-2、在Office 365中的Turing-NLG，以及Bing中具有5亿参数的大型transformer。

AI优先的超级计算机

基础架构是实现大规模预训练模型的关键。微软和OpenAI的科研人员通过合作，开发了专用于大型深度学习工作负载的超级计算机。该超级计算机在单个系统的基础架构内，不可思议地集成了超过28.5万个CPU内核、1万个GPU，每台GPU服务器的网络互连能力达每秒400千兆位。

这台微软-OpenAI合作超级计算机所使用的基础设施，已跻身全球超级计算机排名的前五位。更值得关注的是，这台超级计算机针对AI工作负载做了优化。从这个意义上讲，该超级计算机设计用于执行Turing-NLG、BERT等大型预训练模型。微软将新的超级计算机以Azure云的服务形式提供给开发人员。

总结

未来几年中，人工智能市场最值得关注的一个趋势，就是大规模人工智能主导地位之争。现在DeepSpeed-ZeRO-2、ONNX Runtime和新型OpenAI超级计算机等举措，表明微软在该竞赛中领先迈出了坚实的一步。

原文链接：

https://medium.com/swlh/inside-microsofts-new-frameworks-to-enable-large-scale-ai-953e9a977912

创作场景

探秘微软大规模人工智能的新框架 ZeRO-2