腾讯混元AI大模型训练技术揭秘——太极AngelPTM

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

一、导言

最近，腾讯发布的混元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来提高硬件利用率。