使用deepspeed并行训练模型的简单入门例程

deepspeed简介

deepspeed不仅可以并行训练模型，而且可以解决大模型训练过程中，显卡内存不够用的问题。和DDP这种并行训练模型的方式的主要区别在于，deepspeed有zerosoffload技术，可以将参数部署到cpu上，和gpu一起训练，从而解决显卡内存不够的问题。本文章基于deepspeed官方公布的代码进行了总结改进，主要是介绍deepspeed简单使用（代码公布在github上:link）。
zerosoffload技术和如何使用vscode调试，可以参考链接: link。

deepspeed安装

pip install deepspeed

deepspeed使用步骤

1.设置配置文件deepspeed_config.json

deepspeed自带很多默认参数，都储存在了外部文件deepspeed_config.json中。本文只是简单给出了一个模板，更多参数的使用需要自己配置。
deepspeed_config.json内容：

{
    "train_batch_size": 4, 
    "optimizer": {
        "type": "SGD",
        "params": {
            "lr": 0.001,
            "momentum": 0.9
        }
    }
}

2.程序中的超参数

parser.add_argument('--local_rank', type=int, default=-1,
                     help='local rank passed from distributed launcher')

local_rank涉及到了多个显卡和服务器中的概念，local_rank是相对于服务器的进程序号，在这里可以理解为使用的gpu编号，训练中，需要将数据放置到不同的gpu中，那就涉及到了gpu编号，而不是简单的.to(‘cuda’)来实现了，如下所示。（本文只考虑了单个服务器上的多个GPU并行训练，如果涉及到多个服务器，本文并不通用）

tensor.cuda(args.local_rank)

同时应该将args的超参数处理成deepspeed需要的形式。

        args, _ = parser.parse_known_args()
        if args.isDeepspeed:
            parser = deepspeed.add_config_arguments(parser)
        args = parser.parse_args()

其中的

        args, _ = parser.parse_known_args()

不要写成了

args = parser.parser_args()

要不然会报错。

3.处理数据集

使用多个gpu并行训练的目的就是希望使用更大的batch_size进行模型训练，因此如果不对数据集进行简单处理，那么每张显卡上面使用dataloader得到的sample_batch都是一样的，并不会起到不同显卡处理不同数据的目的。

from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
if isdeepspeed:
    dataloader_training = DataLoader(dataset, batch_size,sampler=DistributedSampler(dataset, shuffle=True),
                                    num_workers=16, pin_memory=True,drop_last=True)
else:
    dataloader_training = DataLoader(dataset, batch_size,
                                    shuffle=True, num_workers=16, pin_memory=True,drop_last=True)

需要注意的地方有两点：

• 使用deepspeed，应该在DataLoader中加入DistributedSampler(dataset, shuffle=True)，这样起到了不同显卡使用不同数据的作用。
• 如果使用了DistributedSampler，就不要在DataLoader中使用shuffle参数了，要不然会报错。

4.处理模型

这里列出了两种方式，一种是参数需要训练的模型，另一种是参数冻结的模型

• 需要训练的模型：

self.optimizer = Adam(self.model.parameters(), lr=1e-2)
self.model,*_ = deepspeed.initialize(args=self.args, model=self.model,
									model_parameters=self.model.parameters(), optimizer=self.optimizer)

• 不需要训练的模型：

self.model,*_ = deepspeed.initialize(args=self.args,model=self.model)
self.model.eval()

总结

基于上述流程和deepspeed官方公布的使用代码，笔者提供了一个简单的deepspeed代码供参考。后面会写一下如何使用vscode进行deepspeed调试和使用zerosoffload参数。