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下都出現了顯存占用增大的情況，這個正常來說是可以節省顯存的，不太明白為何會這樣，留待以后繼續分析。
