鉴于网上此类教程有不少模糊不清,对原理不得其法,代码也难跑通,故而花了几天细究了一下相关原理和实现,欢迎批评指正!
关于此部分的代码,可以去
https://github.com/sherlcok314159/dl-tools
查看
「在开始前,我需要特别致谢一下一位挚友,他送了我双显卡的机器来赞助我做个人研究,否则多卡的相关实验就得付费在云平台上跑了,感谢好朋友一路以来的支持,这份恩情值得一辈子铭记!这篇文章作为礼物赠与挚友。」
我们在两种情况下进行并行化训练[1]:
ModelParallel(MP)
DataParallel(DP)
。当然,DP肯定还可以加速训练,常见于大模型的训练中这里只讲一下DP在pytorch中的原理和相关实现,即DataParallel和DistributedParallel
实现就是循环往复一个过程:数据分发,模型复制,各自前向传播,汇聚输出,计算损失,梯度回传,梯度汇聚更新,可以参见下图[2]:
pytorch中部分关键源码[3]截取如下:
def data_parallel(
module,
input,
device_ids,
output_device=None
):
if not device_ids:
return module(input)
if output_device is None:
output_device = device_ids[0]
# 复制模型
replicas = nn.parallel.replicate(module, device_ids)
# 拆分数据
inputs = nn.parallel.scatter(input, device_ids)
replicas = replicas[:len(inputs)]
# 各自前向传播
outputs = nn.parallel.parallel_apply(replicas, inputs)
# 汇聚输出
return nn.parallel.gather(outputs, output_device)
因为运行时会将数据平均拆分到GPU上,所以我们准备数据的时候, batch size = per_gpu_batch_size * n_gpus
同时,需要注意主GPU需要进行汇聚等操作,因而需要比单卡运行时多留出一些空间
import torch.nn as nn
# device_ids默认所有可使用的设备
# output_device默认cuda:0
net = nn.DataParallel(model, device_ids=[0, 1, 2],
output_device=None, dim=0)
# input_var can be on any device, including CPU
output = net(input_var)
接下来看个更详细的例子[4],需要注意的是被DP包裹之后涉及到模型相关的,需要调用DP.module,比如加载模型
class Model(nn.Module):
# Our model
def __init__(self, input_size, output_size):
super(Model, self).__init__()
# for convenience
self.fc = nn.Linear(input_size, output_size)
def forward(self, input):
output = self.fc(input)
print("\tIn Model: input size", input.size(),
"output size", output.size())
return output
bs, input_size, output_size = 6, 8, 10
# define inputs
inputs = torch.randn((bs, input_size)).cuda()
model = Model(input_size, output_size)
if torch.cuda.device_count() > 1:
print("Let's use", torch.cuda.device_count(), "GPUs!")
# dim = 0 [6, xxx] -> [2, ...], [2, ...], [2, ...] on 3 GPUs
model = nn.DataParallel(model)
# 先DataParallel,再cuda
model = model.cuda()
outputs = model(inputs)
print("Outside: input size", inputs.size(),
"output_size", outputs.size())
# assume 2 GPUS are available
# Let's use 2 GPUs!
# In Model: input size torch.Size([3, 8]) output size torch.Size([3, 10])
# In Model: input size torch.Size([3, 8]) output size torch.Size([3, 10])
# Outside: input size torch.Size([6, 8]) output_size torch.Size([6, 10])
# save the model
torch.save(model.module.state_dict(), PATH)
# load again
model.module.load_state_dict(torch.load(PATH))
# do anything you want
如果经常使用huggingface,这里有两个误区需要小心:
# data parallel object has no save_pretrained
model = xxx.from_pretrained(PATH)
model = nn.DataParallel(model).cuda()
model.save_pretrained(NEW_PATH) # error
# 因为model被DP wrap了,得先取出模型 #
model.module.save_pretrained(NEW_PATH)
# HF实现貌似是返回N个loss(N为GPU数量)
# 然后对N个loss取mean
outputs = model(**inputs)
loss, logits = outputs.loss, outputs.logits
loss = loss.mean()
loss.backward()
# 返回的logits是汇聚后的
# HF实现和我们手动算loss有细微差异
# 手动算略好于HF
loss2 = loss_fct(logits, labels)
assert loss != loss2
True
了解前面的原理后,就会明白为什么会显存不均匀。因为GPU0
比其他GPU多了汇聚的工作,得留一些显存,而其他GPU显然是不需要的。那么,解决方案就是让其他GPU的batch size开大点,GPU0
维持原状,即不按照默认实现的平分数据
首先我们继承原来的DataParallel(此处参考[5])),这里我们给定第一个GPU的bs就可以,这个是实际的bs而不是乘上梯度后的。假如你想要总的bs为64,梯度累积为2,一共2张GPU,而一张最多只能18,那么保险一点GPU0
设置为14,GPU1是18,也就是说你DataLoader每个batch大小是32,gpu0_bsz=14
class BalancedDataParallel(DataParallel):
def __init__(self, gpu0_bsz, *args, **kwargs):
self.gpu0_bsz = gpu0_bsz
super().__init__(*args, **kwargs)
核心代码就在于我们重新分配chunk_sizes,实现思路就是将总的减去第一个GPU的再除以剩下的设备,源码的话有些死板,用的时候不妨参考我的[6]
def scatter(self, inputs, kwargs, device_ids):
# 不同于源码,获取batch size更加灵活
# 支持只有kwargs的情况,如model(**inputs)
if len(inputs) > 0:
bsz = inputs[0].size(self.dim)
elif kwargs:
bsz = list(kwargs.values())[0].size(self.dim)
else:
raise ValueError("You must pass inputs to the model!")
num_dev = len(self.device_ids)
gpu0_bsz = self.gpu0_bsz
# 除第一块之外每块GPU的bsz
bsz_unit = (bsz - gpu0_bsz) // (num_dev - 1)
if gpu0_bsz < bsz_unit:
# adapt the chunk sizes
chunk_sizes = [gpu0_bsz] + [bsz_unit] * (num_dev - 1)
delta = bsz - sum(chunk_sizes)
# 补足偏移量
# 会有显存溢出的风险,因而最好给定的bsz是可以整除的
# e.g., 总的=52 => bsz_0=16, bsz_1=bsz_2=18
# 总的=53 => bsz_0=16, bsz_1=19, bsz_2=18
for i in range(delta):
chunk_sizes[i + 1] += 1
if gpu0_bsz == 0:
chunk_sizes = chunk_sizes[1:]
else:
return super().scatter(inputs, kwargs, device_ids)
return scatter_kwargs(inputs, kwargs, device_ids, chunk_sizes, dim=self.dim)
[1]
Pytorch中多GPU并行计算教程: https://blog.csdn.net/qq_37541097/article/details/109736159
[2]
DDP: https://www.cnblogs.com/ljwgis/p/15471530.html
[3]
MULTI-GPU EXAMPLES: https://pytorch.org/tutorials/beginner/former_torchies/parallelism_tutorial.html?highlight=dataparallel
[4]
OPTIONAL: DATA PARALLELISM: https://pytorch.org/tutorials/beginner/blitz/data_parallel_tutorial.html
[5]
transformer-xl: https://github.com/kimiyoung/transformer-xl
[6]
Balanced Data Parallel: https://github.com/sherlcok314159/dl-tools/blob/main/balanced_data_parallel/README.md
📝论文解读投稿,让你的文章被更多不同背景、不同方向的人看到,不被石沉大海,或许还能增加不少引用的呦~ 投稿加下面微信备注“投稿”即可。