DeepSpeed加速大模型训练

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下都出现了显存占用增大的情况，这个正常来说是可以节省显存的，不太明白为何会这样，留待以后继续分析。