DeepSpeed是微软推出的一个框架，可以对Pytorch的模型进行包装，提供了加快模型的训练速度，降低对GPU显存的占用，以及方便进行分布式训练等等高级特性。在这里我也对DeepSpeed进行了测试，看看是否能提高我的transformer模型的训练性能。
在之前的博客中我介绍了如何对GPT 2模型进行SFT的训练，召唤神龙打造自己的ChatGPT_gzroy的博客-CSDN博客，我将基于之前的模型，用DeepSpeed来进行训练，看看效果如何。
要使用DeepSpeed，我们需要进行配置的定义，以下是我定义的一个JSON配置文件：
{"train_batch_size": 4,"steps_per_print": 100,"fp16": {"enabled": true},"gradient_accumulation_steps": 1,"optimizer": {"type": "AdamW","params": {"lr": 0.00006,"betas": [0.9, 0.95],"weight_decay": 0.01}},"scheduler": {"type": "WarmupDecayLR","params": {"warmup_min_lr": 0.000006,"warmup_max_lr": 0.00006,"warmup_num_steps": 1000,"total_num_steps": 40000}}
}建立一个正常的Pytorch模型，如我之前的博客提到的方式，然后用DeepSpeed进行封装，
model, _, _, lr_scheduler = deepspeed.initialize(model=model, model_parameters=optim_groups, config=args.deepspeedcfg)这里面的args.deepspeedcfg指向我们之前配置的JSON文件。
之后在训练的步骤里面改写一下即可。
logits, loss = model(x, y)
model.backward(loss)
model.step()最后就是在我本地的2080Ti显卡上进行测试，以下是测试结果：
1. 不启用DeepSpeed，batch_size=4，训练4000个batch，显存占用为8050MB，耗时748秒
2. 启用DeepSpeed，batch_size=4，采用DeepSpeed的AdamW优化器，不启用ZeRO，显存占用为6392MB，耗时613秒
3. 启用DeepSpeed，batch_size=4，采用Pytorch的AdamW优化器，不启用ZeRO，显存占用为6392MB，耗时802秒
4. 启用DeepSpeed，batch_size=8，采用DeepSpeed的AdamW优化器，不启用ZeRO，显存占用为8688MB，耗时1031秒
5. 启用DeepSpeed，batch_size=8，采用DeepSpeed的AdamW优化器，启用ZeRO Stage 0，显存占用为8688MB。
6. 启用DeepSpeed，batch_size=8，采用DeepSpeed的AdamW优化器，启用ZeRO Stage 1，显存占用为9382MB。
7. 启用DeepSpeed，batch_size=8，采用DeepSpeed的AdamW优化器，启用ZeRO Stage 2，显存占用为10336MB。
8. 启用DeepSpeed，batch_size=8，采用DeepSpeed的AdamW优化器，启用ZeRO Stage 3，如果启用torch.compile(model)会报deepcopy错误，取消torch.compile后运行，显示在AMD CPU上不支持FP16，另外我的设备也不支持offload到nvme，测试中止。
从以上测试结果，可以看到在采用了DeepSpeed自带的优化器的情况下，可以大大的节约显存以及提高训练速度。但是开启ZeRO之后，在Stage 0,1,2下都出现了显存占用增大的情况，这个正常来说是可以节省显存的，不太明白为何会这样，留待以后继续分析。
