现在的模型越来越大,动辄几B甚至几百B。但是显卡显存大小根本无法支撑训练推理。例如,一块RTX2090的10G显存,光把模型加载上去,就会OOM
,更别提后面的训练优化。
作为传统pytorch Dataparallel的一种替代,DeepSpeed的目标,就是为了能够让亿万参数量的模型,能够在自己个人的工作服务器上进行训练推理。
本文旨在简要地介绍Deepspeed进行大规模模型训练的核心理念,以及最基本的使用方法。更多内容,笔者强烈建议阅读HuggingFace Transformer官网对于DeepSpeed的教程:
Transformer DeepSpeed Integration
ZeRO是一种针对大规模分布式深度学习的新型内存优化技术。
在DeepSpeed下,ZeRO训练支持了完整的ZeRO Stages1, 2和3,以及支持将优化器状态、梯度和模型参数从GPU显存下沉到CPU内存或者硬盘上,实现不同程度的显存节省,以便训练更大的模型。