编者按：秉承“技术提效”理念，腾讯广告不断探索技术能力边界，全面升级广告系统，基于“一大平台、两大模型”持续精进创新研发，提升投放效率与投放效果，助力广告主实现高效的全域经营与生意增长。本篇文章是技术团队在大模型训练技术方向的新探索。

---

一、导言

最近，腾讯发布的混元AI万亿大模型登顶权威中文测评基准CLUE榜[1]并超越人类水平。混元AI大模型采用腾讯太极机器学习平台自研的训练框架AngelPTM，相比业届主流的解决方案，太极AngelPTM单机可容纳的模型可达55B，20个节点（A100-40Gx8）可容纳万亿规模模型，节省45%训练资源，并在此基础上训练速度提升1倍！

二、背景

Transformer模型凭借其出色的表达能力在多个人工智能领域均取得了巨大成功，如自然语言处理、计算机视觉和语音处理等。与此同时，随着训练数据量和模型容量的增加可以持续提高模型的泛化能力和通用能力，研究大模型成为了近两年的趋势。如下图所示，近几年NLP预训练模型规模的发展，模型已经从亿级发展到了万亿级参数规模。具体来说，2018年BERT[2]模型最大参数量为340M，2019年GPT-2[3]为十亿级参数的模型。2020年发布的百亿级规模有T5[4]和T-NLG[5]，以及千亿参数规模的GPT-3[6]。在2021年末，Google发布了Switch Transformer[7]，首次将模型规模提升至万亿。

然而GPU硬件发展的速度难以满足Transformer模型规模发展的需求。近四年中，模型参数量增长了十万倍，但GPU的显存仅增长了4倍。举例来说，万亿模型的模型训练仅参数和优化器状态便需要1.7TB以上的存储空间，至少需要425张A100(40G)，这还不包括训练过程中产生的激活值所需的存储。在这样的背景下，大模型训练不仅受限于海量的算力, 更受限于巨大的存储需求。

为了以最小的成本训练大模型，太极AngelPTM基于ZERO策略，将模型的参数、梯度、优化器状态以模型并行的方式切分到所有GPU，并自研ZeRO-Cache框架把内存作为二级存储offload参数、梯度、优化器状态到CPU内存，同时也支持把SSD作为第三级存储。ZeRO-Cache为了最大化最优化的利用内存和显存进行模型状态的缓存，引入了显存内存统一存储视角，将存储容量的上界由内存扩容到内存+显存总和。同时将多流异步化做到了极致，在GPU计算的同时进行数据IO和NCCL通信，使用异构流水线均衡设备间的负载，最大化提升整个系统的吞吐。ZeRO-Cache将GPU显存、CPU内存统一视角管理，减少了冗余存储和内存碎片，增加了内存的利用率，将机器的存储空间压榨到了极致。

三、大模型训练优化遇到的挑战

3.1 多级存储访存带宽不一致

在大模型训练中，激活值、梯度位于GPU中，模型的FP16/FP32参数、优化器状态位于CPU中甚至位于SSD中，模型的前向和反向在GPU上进行运算，而参数更新在CPU做运算，这就需要频繁的进行内存显存以及SSD之间的访问，而GPU访问显存的带宽为1555GB/s，显存与内存数据互传的带宽为32GB/s，CPU访问内存、显存和SSD的带宽分别为200GB/s、32GB/s、3.5GB/s，多级存储访问带宽的不一致很容易导致硬件资源闲置，如何减少硬件资源的闲置时间是大模型训练优化的一大挑战。ZeRO-Cache通过多流异步以及PipelineOptimizer来提高硬件利用率。

3.2 模型状态冗余存储

大模型训练时的模型状态存储于CPU中，在模型训练过程中会不断拷贝到GPU，这就导致模型状态同时存储于CPU和GPU中，这种冗余存储是对本就捉肘见襟的单机存储空间一种严重浪费，如何彻底的去处这种冗余，对低成本训练大模型至关重要。ZeRO-Cache秉承彻底去冗余的思想引入统一视角存储管理从根本上解决模型状态的冗余存储。

3.3 内存碎片过多

大模型拥有巨量的模型状态，单张GPU卡不能完全放置所有模型状态，在训练过程中模型状态被顺序在CPU和GPU之间交换，这种交换导致GPU显存的频繁分配和释放，此外大模型训练过程中海量的Activation也需要频繁分配和释放显存，显存频繁分配和释会产生大量的显存碎片，过多的显存碎片会导致Memory Allocator分配效率低下以及显存浪费，ZeRO-Cache设计开发了Contiguous Memory显存管理器来更细致的管理显存来减少显存碎片率。

3.4 带宽利用率低

在大模型分布式训练中多机之间会进行参数all_gather、梯度reduce_scatter以及MOE层AlltoAll三类通信，单机内部存在模型状态H2D、D2H以及SSD和Host之间的数据传输。通信以及数据传输带宽利用率低是训练框架在分布式训练中最常见的问题，合并通信以及合并数据传输作为解决带宽问题的首选手段在大模型训练中同样适用，合并通信和合并数据传输预示着需要更多的内存和显存缓存，这更加剧了大模型训练的存储压力。为了平衡高带宽利用率以及低缓存，ZeRO-Cache在多流异步的基础上，引入chunk机制管理模型状态通信以及数据传输，在梯度后处理时引入多buffer机制重复利用以分配缓存。

四、太极AngelPTM概览

太极AngelPTM的设计目标是依托太极机器学习平台，为NLP、CV和多模态、AICG等多类预训练任务提供一站式服务。其主要由高性能训练框架、通用加速组件和基础模型仓库组成：

高性能训练框架：包含大模型训练框架ZeRO-Cache，高性能MOE组件，以及3D并行和自动流水并行策略，方便用户结合多种并行策略进行大模型训练。

通用加速组件：包含可减少显存并提高精度的异构Adafactor优化器，可稳定MOE半精度训练loss的Z_loss组件，选择性重计算组件和降低通信代价的PowerSGD组件。

基础模型仓库：包含T5、BERT、GPT以及Transformer等基础模型，供用户直接使用。

五、ZeRO-Cache优化策略

ZeRO-Cache是一款超大规模模型训练的利器，其通过统一视角去管理内存和显存，在去除模型状态冗余的同时扩大单个机器的可用存储空间上限，通过Contiguous Memory显存管理器管理模型参数的显存分配/释放进而减少显存碎片，通过多流均衡各个硬件资源的负载，通过引入SSD进一步扩展单机模型容量。

5.1 统一视角存储管理

大模型训练时模型状态都位于CPU内存中，在训练时会拷贝到GPU显存，这就导致模型状态的冗余存储(CPU和GPU同时存在一份)，此外大模型的训练中会用到大量的pin memory，pin memory的使用会提升性能同时会导致物理内存的大量浪费，如何科学合理的使用pin memory是ZeRO-Cache着重要解决的问题。

ZeRO-Cache本着极致化去冗余的理念引入了chunk对内存和显存进行管理，保证所有模型状态只存储一份通常模型会存储在内存or显存上，ZeRO-Cache提出异构统一存储，采用内存和显存共同作为存储空间，击破了异构存储的壁垒，极大扩充了模型存储可用空间。

在CPU时原生Tensor的底层存储机制对于实际占用的内存空间利用率极不稳定，对此问题ZeRO-Cache实现了Tensor底层分片存储的机制，在极大的扩展了单机可用的存储空间的同时，避免了不必要的pin memory存储浪费，使得单机可负载的模型上限获得了极大提升。

5.2 ZeRO-Cache显存管理器

PyTorch自带的显存管理器可以Cache几乎所有显存进行二次快速分配，在显存压力不大的情况下这种显存分配方式可以达到性能最优，但是对于超大规模参数的模型显存压力剧增且由于参数梯度频繁的显存分配导致显存碎片明显增多，PyTorch Allocator尝试分配显存失败的次数增加，导致训练性能急剧下降。

ZeRO-Cache引入了一个Contiguous Memory显存管理器，其在PyTorch Allocator之上进行二次显存分配管理，模型训练过程中参数需要的显存的分配和释放都由Contiguous Memory统一管理，在实际的大模型的训练中其相比不使用Contiguous Memory显存分配效率以及碎片有显著的提升，模型训练速度有质的飞越。

5.3 PipelineOptimizer

ZeRO-Infinity利用GPU或者CPU更新模型的参数，特别是针对大模型，只可以通过CPU来更新参数，由于CPU更新参数的速度相比GPU更新参数有数十倍的差距且参数更新几乎可以占到整个模型训练时间的一半，在CPU更新参数的时候GPU空闲且显存闲置，造成资源的极大浪费。

ZeRO-Cache会在模型参数算出梯度之后开始Cache模型的优化器状态到GPU显存，并在参数更新的时候异步Host和Device之间的模型状态数据传输，同时支持CPU和GPU同时更新参数。ZeRO-Cache pipeline了模型状态H2D、参数更新、模型状态D2H，最大化的利用硬件资源，避免硬件资源闲置。

此外，ZeRO-Cache自研了异构Adafactor优化器，支持CPU和GPU同时进行参数的更新，其可以减少33%的模型状态存储空间，同时可提高模型训练精度。

5.4 多流异步化

大模型训练过程中有大量的计算和通信，包括GPU计算、H2D和D2H单机通信、NCCL多机通信等，涉及的硬件有GPU、CPU、PCIE等。ZeRO-Cache为了最大化的利用硬件，多流异步化GPU计算、H2D和D2H单机通信、NCCL多机通信，参数预取采用用时同步机制，梯度后处理采用多buffer机制，优化器状态拷贝采用多流机制。

5.5 ZeRO-Cache SSD框架

为了更加低成本的扩展模型参数，ZeRO-Cache进一步引入了SSD作为三级存储，针对GPU高计算吞吐、高通信带宽和SSD低PCIE带宽之间的GAP，ZeRO-Cache放置所有fp16参数和梯度到内存中，让foward和backward的计算不受SSD低带宽影响，同时通过对优化器状态做半精度压缩来缓解SSD读写对性能的影响。

六、实验效果

本小节基于A100 GPU(40G)集群和1.6T ETH RDMA网络环境对比评测目前业界比较主流的几个大模型框架在不同规模模型上的性能表现，包括首次提出ZeRO理念的框架A、可高效扩展并训练Transformer模型的框架B以及集成高效并行化技术的框架C。秉承降低成本、提高训练效率的理念，我们以最少的资源去呈现ZeRO-Cache的高性能和大模型能力，环境配置为4台1T内存，96核CPU机器，每台机器有8张A100 GPU(40G)，模型具体配置如下表，测试均采用真实训练数据，序列长度为512，主要对比测试TFLOPs/s。

七、实验数据

图表说明：图中数值0表示该规模模型不能在当前资源下运行，而C框架因官方缺乏多机训练文档，多次测试失败，暂且以0表示其性能数据。各个框架的测试性能数据均以Batch Size打满GPU为准。

实验数据表明虽然在1.7B小模型上ZeRO-Cache单机略差于B框架（ZeRO-Cache前向和反向计算中需要进行参数的预取，而B框架无需此操作），但是在其他规模模型上无论单机多机性能均优于其他框架，我们以13B模型为例来说明各个框架对比A框架的性能提升，以Tflops/s进行计算，具体见下表。此外，在GPT3-175B大模型上，ZeRO-Cache仅仅只需要32张卡便可进行模型的训练，而其他三个框架均无法在这么少资源下进行训练。

八、业务应用

混元AI大模型先后支持了包括微信、QQ、游戏、腾讯广告、腾讯云等众多产品和业务，通过NLP、CV、跨模态等AI大模型，不仅为业务创造了增量价值而且降低了使用成本。特别是广告场景的落地应用带来了大幅的GMV提升，验证了大模型的商业化能力。

九、展望未来

ZeRO-Cache将继续秉承降低成本、提高训练效率的理念，用更少的资源以更快的速度训练更大的模型。我们后续会在ZeRO-Cache引入模型并行、流水并行等并行策略进一步提高框架的训练速度，同时通过模型状态无损压缩存储对SSD框架做进一步性能优化。

十、致谢：北京大学-腾讯协同创新实验室

北京大学-腾讯协同创新实验室成立于2017年，主要在人工智能、大数据等领域展开前沿探索和人才培养，打造国际领先的校企合作科研平台和产学研用基地。实验室通过合作研究，在理论和技术创新、系统研发和产业应用方面取得重要成果和进展，如分布式机器学习平台Angel，已在国际顶级学术会议和期刊发表学术论文20余篇。

**AngelPTM由腾讯和北京大学河图团队共同设计开发完成。**北大河图团队来自北京大学-腾讯协同创新实验室主任崔斌教授课题组， 自主研发了分布式深度学习系统河图（Hetu[8]，https://github.com/PKU-DAIR/Hetu），面向拥有海量训练数据、超大模型参数的深度学习训练场景，并具备通用、高效、敏捷、灵活和可扩展等特性。AngelPTM集成了多项北大河图团队在大规模预训练模型的系统优化方面的工作，兼顾学术界的创新性和工业界的可用性。

此前，河图团队还参与了腾讯稀疏大模型训练框架HEAP及新一代大模型参数服务器 AngelPS的前瞻研究[9]和系统建设工作，推动河图创新成果在腾讯实际业务中的应用落地，促进产学研合作和成果转化。

附录：参考文献

[1] 腾讯发布万亿大模型训练方法：最快256卡1天训完万亿NLP大模型

https://mp.weixin.qq.com/s/TG2vlT0UrlXSdGgNOlLn5w

[2] BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding https://arxiv.org/abs/1810.04805

[3] Language Models are Unsupervised Multitask Learners https://d4mucfpksywv.cloudfront.net/better-language-models/language-models.pdf

[4] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer [1910.10683] Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer

[5] T-NLG Turing-NLG: A 17-billion-parameter language model by Microsoft - Microsoft Research

[6] Language Models are Few-Shot Learners [2005.14165] Language Models are Few-Shot Learners

[7] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity https://arxiv.org/abs/2101.03961

[8] 北大腾讯联合共建Angel4.0，自研深度学习框架“河图”融入生态

https://mp.weixin.qq.com/s/03GabpMzwCi7MLvDbCGloQ

[9] 腾讯北大合作的稀疏大模型训练加速方案HET入选国际顶会VLDB

https://mp.weixin.qq.com/s/t6y05tv4CHo5dQaQ0A4tbA