《GPU显存瓶颈下的模型训练优化策略与实践》

文章主题：GPU, 硬件发展, Transformer模型, 显存需求

666AI工具大全，助力做AI时代先行者！

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

一. 导言

最近，腾讯发布的混元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）的GPU。此外，训练过程中产生的激活值也需要大量的存储空间。因此，在大模型训练的过程中，算力和存储需求都成为了关键的限制因素。

为了以最小的成本训练大模型，太极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的方式对内存与显存进行管理，确保所有模型状态仅被存储一次，这通常是在内存或显存中存储的。ZeRO-Cache提出了异构统一存储的方案，将内存与显存共同纳入存储空间，这一举措打破了异构存储的障碍，从而极大地扩展了模型存储的应用范围。

原内容描述了CPU时原生Tensor的存储效率低下的问题，并提到了ZeRO-Cache解决方案。这个方案通过实现Tensor底层分片存储的方式，极大地提高了单机可用的存储空间，同时避免了不必要的Pin Memory存储浪费，从而提升了单机可承载的模型规模的上限。为了更准确地传达这一信息，可以这样表达：在处理CPU时的原生Tensor时，其底层存储机制导致了实际占用的内存空间利用率的极度不稳定。为了解决这个问题，ZeRO-Cache采用了Tensor底层分片存储的策略。这种方法不仅极大地扩大了单机可用的存储空间，而且避免了无必要的Pin Memory存储浪费，从而显著提升了单机所能承载的模型规模的上限。

5.2 ZeRO-Cache显存管理器

PyTorch的默认显存管理器具备强大的缓存功能，能对大部分显存进行高效的二次快速分配。在显存需求相对较小的情况下，这种方式确实能带来最佳的性能表现。然而，当模型规模扩大，参数数量激增时，显存压力显著上升。由于参数梯度导致的频繁显存分配，还会使显存碎片化问题愈发严重。在这种情况下，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的表现出现，因此暂时将其性能数据标记为0。另外，各个框架的测试性能数据都是基于Batch Size满GPU的标准来衡量。

经过实验数据分析，我们发现ZeRO-Cache单机性能在1.7B小模型上略逊于B框架，这主要是因为ZeRO-Cache在进行前向和反向计算时需要预取参数，而B框架则无需此步骤。然而，在其他规模模型上，ZeRO-Cache无论是单机还是多机性能都明显优于其他框架。以13B模型为例，我们可以看到，相较于A框架，ZeRO-Cache的性能提升了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的前沿研究和系统建设工作中。我们致力于将河图的创新成果在腾讯的实际业务中得到应用，从而推动产学研的合作进程并加速成果的转化。

附录：参考文献

[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

拓展阅读

• VLDB 2022可扩展数据科学最佳论文解读：HET训练框架提升大模型训练效率

• “太极”如何利用混部资源，助力腾讯广告降本增效

• 腾讯自研万亿级NLP大模型，自动生成和衍生广告文案

GPU, 硬件发展, Transformer模型, 显存需求

AI时代，拥有个人微信机器人AI助手！AI时代不落人后！

免费ChatGPT问答，办公、写作、生活好得力助手！

搜索微信号aigc666aigc999或上边扫码，即可拥有个人AI助手！