行业专题研究 极致利用存储空间,ZeRO-Cache策略助力优化模型训练。腾讯为了以最小的成本和最快的性能训练大模型,太极机器学习平台对 DeepSpeed 和Megatron-LM 进行了深度定制优化,推出了 AngelPTM 训练框架,腾讯发布的混元 AI 大模型便是基于 AngelPTM 框架训练而来。在具体技术优化策略上,腾讯是基于 ZERO 策略,将模型的参数、梯度、优化器状态以模型并行的方式切分到所有 GPU,并自研 ZeRO-Cache 框架把内存作为二级存储 offload 参数、梯度、优化器状态到 CPU 内存,同时也支持把 SSD作为第三级存储。而为了最大化和最优化的利用内存和显存进行模型状态的缓存,腾讯引入了显存内存统一存储视角,将存储容量的上界由内存扩容到内存+显存总和。同时,将多流异步化做到极致,在 GPU 计算的同时进行数据 IO 和 NCCL 通信,使用异构流水线均衡设备间的负载,最大化提升整个系统的吞吐。ZeRO-Cache 将 GPU 显存、CPU 内存统一视角管理,减少了冗余存储和内存碎片,增加了内存的利用率,将机器的存