DeepSpeed 技术详解DeepSpeed 是微软开源的一个深度学习优化库专门用来解决大模型训练时遇到的两大核心痛点显存不够用和训练太慢。你可以把它理解为给 PyTorch 加了一个涡轮增压器让原本跑不动的大模型变得可以训练。 先理解问题为什么需要 DeepSpeed训练一个大模型比如 GPT-3 有 1750 亿参数GPU 显存需要存储三类东西模型参数模型本身的权重梯度反向传播时计算出的更新方向优化器状态比如 Adam 优化器会额外保存动量m和方差v传统的数据并行方式下每张 GPU 卡都完整保存这三样东西极其浪费显存。比如一个 10 亿参数的模型用 FP16 精度训练每张卡大约需要 16GB 显存——这还只是复制模型本身并没有变大。⭐ 核心技术ZeRO零冗余优化器ZeRO 是 DeepSpeed 的灵魂技术核心思想用一句话概括就是别每张卡都存完整副本把东西拆开分给不同的卡存。ZeRO 分三个阶段逐步瘦身ZeRO-1只切分优化器状态把 Adam 优化器的 m 和 v 分散到不同卡上。参数和梯度还是完整复制但已经能省约40%显存。适合不想增加太多通信开销、只想省点显存的场景。ZeRO-2切分优化器状态 梯度在 ZeRO-1 的基础上梯度也不再每张卡都存完整副本而是分片存储。显存节省约75%适合百亿级模型的训练。ZeRO-3全部切分最强参数、梯度、优化器状态全部打散到不同卡上。每张卡只存 1/N 的模型需要用到某部分参数时再临时从其他卡借过来通过 AllGather 通信。显存节省超过90%单卡就能跑百亿级模型。打个比方传统方式就像每个员工都买了一整套《大英百科全书》放在桌上ZeRO 就像把百科全书拆成若干册每人只放一册需要时互相借阅。ZeRO-Infinity突破物理显存极限在 ZeRO-3 的基础上更进一步支持把不活跃的模型状态卸载到 CPU 内存甚至 NVMe 固态硬盘上。这样 GPU 显存就不再是瓶颈理论上单卡可以跑万亿级参数的模型。 3D 并行多维度加速DeepSpeed 还支持三种并行策略的自由组合称为3D 并行数据并行DP把训练数据切分到多张卡每张卡跑完整模型适合数据量大、模型中等的场景张量并行TP把单层的矩阵运算切分到多张卡适合单层特别大的模型流水线并行PP把模型按层切分比如卡1跑前10层卡2跑后10层适合超深的 Transformer 模型实际训练中这三种可以自由组合。比如用 1024 张卡训练万亿参数模型可以配置为64(DP) × 8(TP) × 2(PP)实现线性扩展。️ 其他重要优化混合精度训练支持 FP16/BF16/INT8显存减半的同时速度提升 2~5 倍梯度检查点选择性保存中间激活值反向传播时重新计算可节省约 65% 显存通信优化1-bit Adam 等算法可将通信量减少高达 26 倍长序列支持稀疏注意力内核支持比标准 Transformer 长一个数量级的输入序列 总结维度DeepSpeed 做了什么显存优化ZeRO 分片存储消除冗余节省 90% 显存训练加速3D 并行 混合精度 通信优化易用性基于 PyTorch 的轻量封装只需改几行代码规模突破支持从单卡百亿到千卡万亿级模型训练简单来说DeepSpeed 就是让大模型训练从不可能变成可落地的关键工具。像 GPT-3、BLOOM-176B 等知名大模型的训练背后都有 DeepSpeed 的支撑。
