Why Parallel

我們?cè)趦煞N情況下進(jìn)行并行化訓(xùn)練[1]：

「模型一張卡放不下」：我們需要將模型不同的結(jié)構(gòu)放置到不同的GPU上運(yùn)行，這種情況叫ModelParallel(MP)

「一張卡的batch size(bs)過小」：有些時(shí)候數(shù)據(jù)的最大長度調(diào)的比較高（e.g., 512），可用的bs就很小，較小的bs會(huì)導(dǎo)致收斂不穩(wěn)定，因而將數(shù)據(jù)分發(fā)到多個(gè)GPU上進(jìn)行并行訓(xùn)練，這種情況叫DataParallel(DP)。當(dāng)然，DP肯定還可以加速訓(xùn)練，常見于大模型的訓(xùn)練中

這里只講一下DP在pytorch中的原理和相關(guān)實(shí)現(xiàn)，即DataParallel和DistributedParallel

Data Parallel

實(shí)現(xiàn)原理

實(shí)現(xiàn)就是循環(huán)往復(fù)一個(gè)過程：數(shù)據(jù)分發(fā)，模型復(fù)制，各自前向傳播，匯聚輸出，計(jì)算損失，梯度回傳，梯度匯聚更新，可以參見下圖[2]：

pytorch中部分關(guān)鍵源碼[3]截取如下：

defdata_parallel(
module,
input,
device_ids,
output_device=None
):
ifnotdevice_ids:
returnmodule(input)

ifoutput_deviceisNone:
output_device=device_ids[0]

#復(fù)制模型
replicas=nn.parallel.replicate(module,device_ids)
#拆分?jǐn)?shù)據(jù)
inputs=nn.parallel.scatter(input,device_ids)
replicas=replicas[:len(inputs)]
#各自前向傳播
outputs=nn.parallel.parallel_apply(replicas,inputs)
#匯聚輸出
returnnn.parallel.gather(outputs,output_device)

代碼使用

因?yàn)檫\(yùn)行時(shí)會(huì)將數(shù)據(jù)平均拆分到GPU上，所以我們準(zhǔn)備數(shù)據(jù)的時(shí)候， batch size = per_gpu_batch_size * n_gpus

同時(shí)，需要注意主GPU需要進(jìn)行匯聚等操作，因而需要比單卡運(yùn)行時(shí)多留出一些空間

importtorch.nnasnn
#device_ids默認(rèn)所有可使用的設(shè)備
#output_device默認(rèn)cuda:0
net=nn.DataParallel(model,device_ids=[0,1,2],
output_device=None,dim=0)
#input_varcanbeonanydevice,includingCPU
output=net(input_var)

接下來看個(gè)更詳細(xì)的例子[4]，需要注意的是被DP包裹之后涉及到模型相關(guān)的，需要調(diào)用DP.module，比如加載模型

classModel(nn.Module):
#Ourmodel
def__init__(self,input_size,output_size):
super(Model,self).__init__()
#forconvenience
self.fc=nn.Linear(input_size,output_size)

defforward(self,input):
output=self.fc(input)
print("	InModel:inputsize",input.size(),
"outputsize",output.size())
returnoutput

bs,input_size,output_size=6,8,10
#defineinputs
inputs=torch.randn((bs,input_size)).cuda()
model=Model(input_size,output_size)
iftorch.cuda.device_count()>1:
print("Let'suse",torch.cuda.device_count(),"GPUs!")
#dim=0[6,xxx]->[2,...],[2,...],[2,...]on3GPUs
model=nn.DataParallel(model)
#先DataParallel，再cuda
model=model.cuda()
outputs=model(inputs)
print("Outside:inputsize",inputs.size(),
"output_size",outputs.size())
#assume2GPUSareavailable
#Let'suse2GPUs!
#InModel:inputsizetorch.Size([3,8])outputsizetorch.Size([3,10])
#InModel:inputsizetorch.Size([3,8])outputsizetorch.Size([3,10])
#Outside:inputsizetorch.Size([6,8])output_sizetorch.Size([6,10])

#savethemodel
torch.save(model.module.state_dict(),PATH)
#loadagain
model.module.load_state_dict(torch.load(PATH))
#doanythingyouwant

如果經(jīng)常使用huggingface，這里有兩個(gè)誤區(qū)需要小心：

#dataparallelobjecthasnosave_pretrained
model=xxx.from_pretrained(PATH)
model=nn.DataParallel(model).cuda()
model.save_pretrained(NEW_PATH)#error
#因?yàn)閙odel被DPwrap了，得先取出模型#
model.module.save_pretrained(NEW_PATH)

#HF實(shí)現(xiàn)貌似是返回N個(gè)loss（N為GPU數(shù)量）
#然后對(duì)N個(gè)loss取mean
outputs=model(**inputs)
loss,logits=outputs.loss,outputs.logits
loss=loss.mean()
loss.backward()

#返回的logits是匯聚后的
#HF實(shí)現(xiàn)和我們手動(dòng)算loss有細(xì)微差異
#手動(dòng)算略好于HF
loss2=loss_fct(logits,labels)
assertloss!=loss2
True

顯存不均勻

了解前面的原理后，就會(huì)明白為什么會(huì)顯存不均勻。因?yàn)镚PU0比其他GPU多了匯聚的工作，得留一些顯存，而其他GPU顯然是不需要的。那么，解決方案就是讓其他GPU的batch size開大點(diǎn)，GPU0維持原狀，即不按照默認(rèn)實(shí)現(xiàn)的平分?jǐn)?shù)據(jù)

首先我們繼承原來的DataParallel（此處參考[5])），這里我們給定第一個(gè)GPU的bs就可以，這個(gè)是實(shí)際的bs而不是乘上梯度后的。假如你想要總的bs為64，梯度累積為2，一共2張GPU，而一張最多只能18，那么保險(xiǎn)一點(diǎn)GPU0設(shè)置為14，GPU1是18，也就是說你DataLoader每個(gè)batch大小是32，gpu0_bsz=14

classBalancedDataParallel(DataParallel):
def__init__(self,gpu0_bsz,*args,**kwargs):
self.gpu0_bsz=gpu0_bsz
super().__init__(*args,**kwargs)

核心代碼就在于我們重新分配chunk_sizes，實(shí)現(xiàn)思路就是將總的減去第一個(gè)GPU的再除以剩下的設(shè)備，源碼的話有些死板，用的時(shí)候不妨參考我的[6]

defscatter(self,inputs,kwargs,device_ids):
#不同于源碼，獲取batchsize更加靈活
#支持只有kwargs的情況，如model(**inputs)
iflen(inputs)>0:
bsz=inputs[0].size(self.dim)
elifkwargs:
bsz=list(kwargs.values())[0].size(self.dim)
else:
raiseValueError("Youmustpassinputstothemodel!")

num_dev=len(self.device_ids)
gpu0_bsz=self.gpu0_bsz
#除第一塊之外每塊GPU的bsz
bsz_unit=(bsz-gpu0_bsz)//(num_dev-1)
ifgpu0_bszbsz_0=16,bsz_1=bsz_2=18
#總的=53=>bsz_0=16,bsz_1=19,bsz_2=18
foriinrange(delta):
chunk_sizes[i+1]+=1
ifgpu0_bsz==0:
chunk_sizes=chunk_sizes[1:]
else:
returnsuper().scatter(inputs,kwargs,device_ids)

returnscatter_kwargs(inputs,kwargs,device_ids,chunk_sizes,dim=self.dim)

優(yōu)缺點(diǎn)

優(yōu)點(diǎn)：便于操作，理解簡單

缺點(diǎn)：GPU分配不均勻；每次更新完都得銷毀「線程」（運(yùn)行程序后會(huì)有一個(gè)進(jìn)程，一個(gè)進(jìn)程可以有很多個(gè)線程）重新復(fù)制模型，因而速度慢

審核編輯：湯梓紅