使用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参数。