如何在多个GPU中训练非常大的模型?
目录
一、并行训练策略
1.数据并行
2.模型并行
3.混合并行:
4.上下文并行
二、内存优化技术
三、总结
在多个GPU上训练超大规模模型(如千亿参数级语言模型)需要结合并行策略、内存优化技术、分布式框架。
一、并行训练策略
1.数据并行
原理:每个GPU保存完整的模型副本,处理不同的数据批次,通过同步梯度更新参数。在每个小批处理结束时,需要同步梯度或权重以免陈旧。
适用场景:模型可单卡容纳,需加速训练速度
2.模型并行
1.张量并行:将模型层内权重拆分到多个GPU(如按行或者按列切分),各GPU计算后合并结果,适用于单层参数过大(如Transformer的注意力头)。
2.流水线并行:将模型按层切分到不同GPU,通过微批次实现并行计算,减少设备空闲时间
3.混合并行:
结合数据并行和模型并行
- DeepSpeed-ZeRO:分片优化器状态、梯度和参数,支持千亿级模型训练。
- FSDP(Fully Sharded Data Parallel):将模型参数、梯度、优化器状态分片到多 GPU,显存占用降低至单卡的 1/N。
4.上下文并行
新兴技术,将长序列上下文分块处理到不同 GPU,提升长文本生成效率(如处理 32k Token 序列)
二、内存优化技术
1.激活重计算:仅保存部分中间激活值,反向传播时重新计算其他部分,显存减少30%-50%
2.混合精度计算:使用FP16/BF16计算前向和反向传播,保留FP32主权重更新参数,显存降低50%
3.梯度累计:小批量训练多次后累加梯度在更新,模型大批量效果,避免显存溢出
4.参数卸载:将暂时不用的参数或激活值转移到CPU内存,需要时在加载回GPU,适合超大模型
三、总结
中小模型:优先使用数据并行+混合精度
超大模型:采用混合并行结合激活冲计算