现在的模型越来越大,动辄几B甚至几百B。但是显卡显存大小根本无法支撑训练推理。例如,一块RTX2090的10G显存,光把模型加载上去,就会OOM
,更别提后面的训练优化。
作为传统pytorch Dataparallel的一种替代,DeepSpeed的目标,就是为了能够让亿万参数量的模型,能够在自己个人的工作服务器上进行训练推理。
本文旨在简要地介绍Deepspeed进行大规模模型训练的核心理念,以及最基本的使用方法。更多内容,笔者强烈建议阅读HuggingFace Transformer官网对于DeepSpeed的教程:
Transformer DeepSpeed Integration
1. 核心思想 (TLDR)
DeepSpeed的核心就在于,GPU显存不够,CPU内存来凑。
比方说,我们只有一张10GB的GPU,那么我们很可能需要借助80GB的CPU,才能够训练一个大模型。
看一下官网对于这个理念的描述:
Why would you want to use DeepSpeed with just one GPU?
- It has a ZeRO-offload feature which can delegate some computatio