Back to Annotated Deep Learning Paper Implementations

ශුන්ය-ඩීපීමතක ප්රශස්තිකරණය

docs/si/scaling/zero3/index.html

latest30.8 KB
Original Source

homescalingzero3

[View code on Github](https://github.com/labmlai/annotated_deep_learning_paper_implementations/tree/master/labml_nn/scaling/zero3/ init.py)

#

ශුන්ය-ඩීපීමතක ප්රශස්තිකරණය

මෙයකඩදාසි වල හඳුන්වා දුන් ශුන්ය-ඩීපී ක්රියාත්මක කිරීමකි ශුන්ය: පුහුණුව සඳහා මතක ප්රශස්තිකරණය ට්රිලියනයක් පරාමිති ආකෘති,

එයප්රශස්තිකරණ තත්වයේ, ශ්රේණියේ සහ පරාමිතීන් බහු උපාංග/නෝඩ් වලට තබා ගනී. එය මුල් ආකෘතියට මතක පරිභෝජනය අඩු කරයි, පරාමිති ගණන Ψ කොතැනද, කැබලි ගණන ද K , Nd​(2+2+K)Ψ​ Nd​ පරාමිතියකට ප්රශස්තිකරණ බයිට් ගණන. 2+2 බිට් 16-නිරවද්යතාව උපකල්පනය කරන පරාමිතිය සහ ශ්රේණියේ මතකය; එනම් පරාමිතියට බයිට් 2 ක් සහ අනුක්රමික ය. K=12 සඳහා ආදම් ප්රශස්තකරණය එය පරාමිතීන්ගේ පිටපතක් සහ fp32 හි පරාමිතියකට අවස්ථා දෙකක් පවත්වා ගෙන යන බැවිනි.

ශුන්ය-ඩීපීහි සන්නිවේදන පරිමාව වේ O(3Ψ). සංසන්දනය සඳහා දත්ත-සමාන්තර පුහුණුව සන්නිවේදන පරිමාවක් O(2Ψ)ඇත.

මෙයනම් කර ඇතත් Zero3 , අප ක්රියාත්මක කර ඇත්තේ එහි ශුන්ය-ඩීපී කොටස පමණක් මිස අවශේෂ මතක පරිභෝජනය ඉලක්ක කරන Zero-R මතක ප්රශස්තිකරණයන් නොවේ. පිටත ක්රියාත්මක කිරීම සඳහා පුහුණුව සඳහා සහාය වන්නේ පරාමිතීන්ගේ උප කුලකයක් පමණි.

මෙමක්රියාත්මක කිරීම Fairscale FSDPවිසින් දේවානුභාවයෙන්.

ZERO-DP මතක ප්රශස්තිකරණය භාවිතා කරමින් GPT Neox මනාව සකස් කිරීම සඳහා ස්ක්රිප්ට් එකක් මෙන්න .

32importfunctools33fromtypingimportList,Optional,Tuple3435importtorch36importtorch.distributedasdist37fromtorchimportnn

#

Zero3ස්ථරය

ආකෘතියේසෑම ස්ථරයක්ම (හෝ අනුයාත ස්ථර කිහිපයක එකතුවක්) මෙම මොඩියුලය තුළ ඔතා තිබිය යුතුය.

40classZero3Layer(nn.Module):

#

සෑමකොටසක්ම පරාමිතීන් chunk ලැයිස්තුවේ තබා ගනී. පුහුණු කළ හැකි පරාමිතීන් සඳහා chunk[0]``chunk[1] වන අතර ස්ථාවර පරාමිතීන් සඳහා වේ.

49chunk:List[nn.Parameter]

#

chunk ලැයිස්තුවේ ඇති කුට්ටි වල ප්රමාණය මෙයයි.

51chunk\_size:List[int]

#

පළමුකුට්ටිය පුහුණු කළ හැකි පරාමිතීන් සඳහා වේ.

53TRAINING\_PARAMS\_IDX=0

#

පුහුණුකළ හැකි සහ ස්ථාවර පරාමිතීන් ලෙස ලැයිස්තු වලට බෙදී ඇති පරාමිති ලැයිස්තුව මෙයයි.

56param\_refs:List[List[nn.Parameter]]

#

පරාමිතීන්පිහාටු කිරීමට CUDA ධාරාව

59fetch\_stream:Optional[torch.cuda.Stream]

#

උපස්ථ/සමුච්චයඵලය අනුක්රමික කිරීමට CDA ඇළ

61backup\_stream:Optional[torch.cuda.Stream]

#

මෙමස්ථරයට පෙර ස්ථර ලැයිස්තුව

63prev\_layer:List['Zero3Layer']

#

මෙමස්ථරයට පසුව ස්ථර ලැයිස්තුව

65next\_layer:List['Zero3Layer']

#

වත්මන්ස්ථරයේ පිහිටීම; ලඝු-සටහන් නිදොස්කරණය සඳහා මෙය භාවිතා කරයි

67layer\_idx:int

#

පරාමිතීන්ලබා ගෙන තිබේද

70is\_fetched:bool

#

ස්ථරයේඋපාංගය

73device:torch.device

#

ස්ථරයේදත්ත වර්ගය

75dtype:torch.dtype

#

එවියයුතු මොඩියුලය

77module:nn.Module

#

නෝඩ්/උපාංගගණන දත්ත හරහා තියුණු වේ

79world\_size:int

#

  • module ඔතා ගත යුතු මොඩියුලය.
  • rank වත්මන් නෝඩයේ නිලය.
  • world_size දත්ත හරහා තියුණු කර ඇති නෝඩ්/උපාංග ගණන.
  • device ස්ථරයේ උපාංගය.
  • dtype ස්ථරයේ දත්ත වර්ගය.
81def\_\_init\_\_(self,module:nn.Module,rank:int,world\_size:int,device:torch.device,dtype:torch.dtype):

#

89super().\_\_init\_\_()

#

ගුණාංගආරම්භ කරන්න

92self.device=device93self.dtype=dtype94self.module=module95self.prev\_layer=[]96self.next\_layer=[]97self.is\_fetched=False98self.world\_size=world\_size99self.layer\_idx=-1100self.fetch\_stream=None101self.backup\_stream=None102103withtorch.no\_grad():

#

ස්ථරයේසියලු පරාමිතීන් එකතු කරන්න

105all\_param\_refs=[pforpinself.parameters()]

#

පරාමිතීන්ගේහැඩය ගබඩා කර තබන්න, මන්ද ඒවා ප්රතිනිර්මාණය කිරීමට අපට පසුව අවශ්ය වේ

108forpinall\_param\_refs:109p.\_orig\_shape=p.shape

#

සියලුමපරාමිතීන් එකම වර්ගයේ තිබිය යුතුය

112forpinall\_param\_refs:113assertp.dtype==dtype,"All parameters should have same dtype"

#

පුහුණුකළ හැකි සහ ස්ථාවර ලෙස වෙනම පරාමිතීන්

116self.param\_refs=[[pforpinall\_param\_refsifp.requires\_grad],117[pforpinall\_param\_refsifnotp.requires\_grad]]118delall\_param\_refs

#

මෙම rank = 0 නෝඩය එක් එක් උපාංගය/නෝඩය ගබඩා කළ යුතු ප්රමාණය ගණනය කරනු ඇත, ඒ අනුව පරාමිතීන් බෙදා හරිනු ඇත.

122ifrank==0:

#

ඒකාබද්ධකිරීම සහ පෑඩ් පුහුණු කළ හැකි (merged_params[0] ) සහ ස්ථාවර (merged_params[1] ) පරාමිතීන්

124merged\_params=[self.\_merge\_and\_pad\_params(ps)forpsinself.param\_refs]

#

පුහුණුකළ හැකි සහ ස්ථාවර පරාම් ප්රමාණ ගණනය කරන්න

126self.chunk\_size=[(len(p)//world\_sizeifpisnotNoneelse0)forpinmerged\_params]

#

ප්රමාණවිකාශනය කරන්න

128dist.broadcast(torch.tensor(self.chunk\_size,device=device),src=0)129else:

#

ප්රමාණලබා ගැනීම සඳහා හිස් ටෙන්සරයක් සාදන්න

131chunk\_size=torch.tensor([0,0],device=device)

#

ප්රමාණලබා ගන්න

133dist.broadcast(chunk\_size,src=0)134self.chunk\_size=chunk\_size.tolist()

#

පුහුණුකළ හැකි (self.chunk[0] ) සහ ස්ථාවර (self.chunk[1] ) පරාමිතීන් සඳහා වත්මන් උපාංගයේ/නෝඩයේ ගබඩා කිරීම සඳහා පරාමිතීන් සාදන්න

138self.chunk=[nn.Parameter(self.\_empty((s,)),requires\_grad=i==self.TRAINING\_PARAMS\_IDX)139fori,sinenumerate(self.chunk\_size)]

#

පුහුණුකළ හැකි සහ ස්ථාවර පරාමිතීන් ඒකාබද්ධ කිරීම සඳහා හිස් ටෙන්සරයක්

142chunk=self.\_empty((sum(self.chunk\_size),))143144ifrank==0:

#

පුහුණුකළ හැකි සහ ස්ථාවර පරාම් දෙකම සංයුක්ත කරන්න

146all\_params=torch.cat([p.view(world\_size,-1)forpinmerged\_params],dim=-1).view(-1)147delmerged\_params

#

සියලුමනෝඩ්/උපාංග වෙත ඒවා විසුරුවා හරින්න

150dist.scatter(chunk,list(all\_params.split(sum(self.chunk\_size))))151delall\_params152else:

#

පරාමිතීන්ලබා ගන්න

154dist.scatter(chunk)

#

කුට්ටියදත්ත එකතු කරන්න

157chunk=chunk.split(self.chunk\_size)158fori,cinenumerate(chunk):159self.chunk[i].data[:]=c160delchunk

#

සාමාන්යපරාමිතීන් පිරිසිදු කරන්න

163self.\_cleanup\_params()

#

පසුගාමීකොක්කක් එක් කරන්න. මොඩියුලයට සාපේක්ෂව අනුක්රමික ගණනය කළ විට මෙය කැඳවනු ලැබේ.

166self.\_backward\_hook\_ref=self.register\_full\_backward\_hook(self.\_backward\_hook)# type: ignore

#

සියලුපරාමිතීන් ඒකාබද්ධ කර එය බෙදිය හැකි වන පරිදි පෑඩ් කරන්න world_size .

168def\_merge\_and\_pad\_params(self,params:List[nn.Parameter])-\>torch.Tensor:

#

මුළුපරාමිතීන් ගණන

173size=sum(p.shape.numel()forpinparams)

#

එයබෙදිය නොහැකි නම් world_size , එය පෑඩ් කරන්න

176ifsize%self.world\_size!=0:177padding\_fixed=self.world\_size-(size%self.world\_size)

#

එසේනොමැති නම්, පෑඩ් කිරීමට අවශ්ය නැත

179else:180padding\_fixed=0

#

හිස්පුරවන tensor සාදන්න

182padding=self.\_empty((padding\_fixed,))

#

සියලුමපරාමිතීන් සංයුක්ත කර එය පෑඩ් කරන්න

184returntorch.cat([p.view(-1)forpinparams]+[padding],dim=0)

#

පරාමිතීන්ගේපුහුණු කළ හැකි කූඩු/මඩු ලබා ගන්න.

වත්මන්නෝඩයේ ප්රශස්තකරණය වෙත අප යන්නේ මෙයයි.

186defget\_trainable\_chunk(self)-\>List[nn.Parameter]:

#

පුහුණුකළ හැකි පරාමිතීන් නොමැති නම් ආපසු සහ හිස් ලැයිස්තුව

193iflen(self.chunk[self.TRAINING\_PARAMS\_IDX])==0:194return[]

#

පුහුණුකළ හැකි කුට්ටිය ලැයිස්තුවක් ලෙස ආපසු එවන්න

197return[self.chunk[self.TRAINING\_PARAMS\_IDX]]

#

දීඇති හැඩයේ හිස් ටෙන්සරයක් සාදන්න.

199def\_empty(self,shape:Tuple[int,...])-\>torch.Tensor:

#

203returntorch.empty(shape,device=self.device,dtype=self.dtype)

#

පරාමිතිදත්ත පිරිසිදු කරන්න

මෙයස්ථර පරාමිතීන් විසින් භාවිතා කරන සියලුම මතකය මුදා හරිනු ඇත.

[email protected]\_grad()206def\_cleanup\_params(self):

#

පරාමිතීන්ලබා ගත නොහැකි බව දැක්වීමට ධජය සකසන්න

214self.is\_fetched=False

#

සියලුපරාමිතීන් හරහා නැවත

217forpsinself.param\_refs:218forpinps:

#

නවමෙහෙයුම් වලට පෙර පරාමිතීන්ගේ මෙහෙයුම් සම්පූර්ණ වන තෙක් රැඳී සිටින්න

220p.data.record\_stream(torch.cuda.current\_stream())

#

පරාමිතියවෙනත් කිසිවක් සමඟ ගබඩා බෙදා නොගන්නා බවට වග බලා ගන්න පරීක්ෂා කරන්න

222assertp.data.storage\_offset()==0,"The tensor is not the sole occupant of the storage."

#

ගබඩාවවෙනස් කරන්න 0. මෙය පරාමිතිය භාවිතා කරන මතකය මුදා හරිනු ඇත.

ස්වයංක්‍රීයප්‍රස්තාරය ඒ පිළිබඳව සඳහනක් තබා ඇති බැවින් සැකසුම මතකය මුදා p.data නොහරිනු ඇත.

226p.data.storage().resize\_(0)# This is what actually clears the memory

#

පරාමිතියටශ්රේණියේ දත්ත නොමැති බවට වග බලා ගන්න

228assertp.gradisNone,'Gradients should be None'

#

සියලුකැබලි වලින් පරාමිතීන් ලබා ගන්න

මෙයසියලු නෝඩ් වලින් සියලු පරාමිති දත්ත ලබා ගන්නා අතර එක් එක් නෝඩයේ පරාමිතීන් නැවත ගොඩනඟනු ඇත.

[email protected]\_grad()231deffetch\_params(self):

#

මඟහැරීම දැනටමත් ලැබී ඇත

239ifself.is\_fetched:240return

#

ධජයසකසන්න

243self.is\_fetched=True

# ලබා

ගැනීමට හෝ බෙදා ගැනීමට කිසිවක් නොමැති නම් මඟ හරින්න.

246ifsum(self.chunk\_size)==0:247return

#

සියලුකැබලි වලින් පරාමිතීන් ලබා fetch_stream ගැනීමට භාවිතා කරන්න

250withtorch.cuda.stream(self.fetch\_stream):

#

පරාමිතීන්ලබා ගැනීම සඳහා හිස් ටෙන්සරයක් සාදන්න

252buffer=self.\_empty((self.world\_size\*sum(self.chunk\_size),))

#

අඛණ්ඩබෆරය නෝඩ් ගණනට බෙදන්න. මෙම බෙදීම් `බෆර්' පිළිබඳ අදහස් වේ.

254buffers=list(buffer.split(sum(self.chunk\_size)))

#

පුහුණුකළ හැකි සහ ස්ථාවර කුට්ටි දෙකම සංයුක්ත කරන්න

257chunk=torch.cat(self.chunk,dim=0)

#

සියලුමනෝඩ්/උපාංග වලින් පරාමිතීන් එක්රැස් කරන්න

260dist.all\_gather(buffers,chunk)

#

රැස්කරන ලද පරාමිතීන් පුහුණු කළ හැකි සහ ස්ථාවර කුට්ටි වලට බෙදන්න

263params=buffer.view(-1,sum(self.chunk\_size)).split(self.chunk\_size,dim=1)

#

එක්රැස්කිරීමේ මෙහෙයුම සම්පූර්ණ වන තෙක් රැඳී සිටින්න, පසුව අවරෝධක වෙත යොමු කිරීම් ඉවත් කරන්න

265buffer.record\_stream(self.fetch\_stream)266forbinbuffers:267b.record\_stream(self.fetch\_stream)268buffer.record\_stream(self.fetch\_stream)269delbuffer270delbuffers

#

පුහුණුකළ හැකි සහ ස්ථාවර පරාමිතීන් අඛණ්ඩ ආතතීන්ට නැවත සකස් කරන්න

273params=[p.reshape(-1)forpinparams]

#

තනිපරාමිති ආතතීන් එකතු කරන්න

276forcont,psinzip(params,self.param\_refs):

#

පරාමිතීන්නොමැති නම්, මඟ හරින්න

278ifnotps:279continue

#

අඛණ්ඩආතතියෙන් ඕෆ්සෙට්

282offset=0

#

ආදර්ශපරාමිතීන් හරහා නැවත නැවතත් සහ අඛණ්ඩ tensor සිට අගයන් පැවරීම

284forpinps:

#

මුල්පරාමිති හැඩය

286shape=p.\_orig\_shape# type: ignore[attr-defined]

#

පරාමිතියේගබඩා ප්රමාණය වෙනස් කරන්න. අපි පරාමිතීන් පිරිසිදු කළ 0 විට මෙය සකස් කරන ලදී.

288p.data.storage().resize\_(shape.numel())

#

අඛණ්ඩආතතියෙන් අගයන් පැවරීම

290p.data[:]=cont[offset:offset+shape.numel()].reshape(shape)

#

වෙනත්මෙහෙයුම් සිදු කිරීමට පෙර මෙහෙයුම් සම්පූර්ණ වන තෙක් රැඳී සිටින්න

292p.data.record\_stream(self.fetch\_stream)

#

ඕෆ්සෙට්යාවත්කාලීන කරන්න

294offset+=shape.numel()

#

වෙනත්මෙහෙයුම් සිදු කිරීමට පෙර මෙහෙයුම සම්පූර්ණ වන තෙක් රැඳී සිටින්න

297cont.record\_stream(self.fetch\_stream)

#

300delparams

#

ඉදිරිසාමාර්ථය

302defforward(self,\*args,\*\*kwargs):

#

වත්මන්නෝඩයේ සියලු පරාමිතීන් ලබා ගන්න. මෙය පෙර ස්තරය මගින් කැඳවනු ලැබේ, එබැවින් මෙම ඇමතුම පරාමිතීන් ලබා ගත හැකි බවට වග බලා ගැනීම සඳහා පමණි.

309self.fetch\_params()

#

පරාමිතියසම්පූර්ණ කිරීම සඳහා රැඳී සිටින්න.

312torch.cuda.current\_stream().wait\_stream(self.fetch\_stream)

#

ඉදිරියටයන ස්ථරවල පරාමිතීන් ලබා ගැනීම ආරම්භ කරන්න, එවිට වත්මන් ස්තරය එහි ගණනය කිරීම් සිදු කරන ඒවා ලබා ගනු ඇත.

316forlayerinself.next\_layer:317layer.fetch\_params()

#

ඔටෝග්රාඩ්සක්රීය කර ඇත්නම් වත්මන් ස්ථරයේ පරාමිතීන් සඳහා පසුගාමී කොකු එකතු කරන්න.

320iftorch.is\_grad\_enabled():321self.\_add\_backward\_hooks()

#

වත්මන්ස්ථරයේ ප්රතිදානයන් ගණනය කරන්න

324res=self.module(\*args,\*\*kwargs)

#

ස්ථරයේපරාමිතීන් පිරිසිදු කරන්න.

ඔටෝග්රාඩ්සක්රීය කර ඇත්නම් පිරිසිදු කිරීම මඟ හරින්න, මෙය ජාලයේ අවසාන ස්තරය වේ, මන්ද පසුගාමී පාස් සඳහා අපට නැවත පරාමිතීන් ලබා ගැනීමට අවශ්ය වනු ඇත.

330ifnottorch.is\_grad\_enabled()orself.next\_layer:331self.\_cleanup\_params()332333returnres

#

වත්මන්ස්ථරයේ පරාමිතීන් සඳහා පසුගාමී කොකු එකතු කරන්න.

335def\_add\_backward\_hooks(self):

#

පසුගාමීකොකු ගණන එකතු කරන ලදි

341self.\_backward\_hook\_handles=0

#

වත්මන්ස්ථරයේ පුහුණු කළ හැකි පරාමිතීන් හරහා ලූප්

344forpinself.param\_refs[self.TRAINING\_PARAMS\_IDX]:

#

කොක්කක්දැනටමත් එකතු කර නොමැති බවට වග බලා ගන්න

346assertnothasattr(p,"\_hook\_handle"),'Parameter has already been hooked'

#

අපටබාධා කළ හැකි ස්වයංක්රීය පියවරක් නිර්මාණය expand_as කිරීමට භාවිතා කරන්න

348p\_tmp=p.expand\_as(p)

#

පසුගාමීකොක්ක එක් කිරීමට හසුරුව ලබා ගන්න. මෙම බ්ලොග් අඩවිය ගැන සාකච්ඡා කරයි grad_acc.

351grad\_acc=p\_tmp.grad\_fn.next\_functions[0][0]

#

පසුගාමීකොක්ක එක් කරන්න

353handle=grad\_acc.register\_hook(354functools.partial(self.\_post\_backward\_hook,p))

#

හසුරුවගැන සඳහනක් තබා ගන්න

356p.\_hook\_handle=handle

#

එකතුකරන ලද කොකු ගණන වැඩි කරන්න

358self.\_backward\_hook\_handles+=1

#

පසුගාමීසිදුවීමක් හසුරුවන්න

පරාමිතියපසුගාමී කොකු සහ මොඩියුලයේ පසුගාමී කොක්ක මගින් මෙය කැඳවනු ලැබේ.

360def\_backward\_event(self):

#

කොකුකවුන්ටරය අඩු කිරීම

368self.\_backward\_hook\_handles-=1

#

සියලුමකොකු (මොඩියුල කොක්ක ඇතුළුව) කැඳවා ඇත්නම්, අපට අනුක්රමික උපස්ථ කර පරාමිතීන් පිරිසිදු කළ හැකිය.

372ifself.\_backward\_hook\_handles==-1:373self.\_backup\_grads()374self.\_cleanup\_params()

#

පෙරස්ථරයේ පරාමිතීන් ලබා ගැනීම ආරම්භ කරන්න, මන්ද ඔටෝග්රාඩ් ඊළඟට එහි අනුක්රමික ක්රියාවලිය කරනු ඇත.

377forlayerinself.prev\_layer:378layer.fetch\_params()

#

පරාමිතියපසුගාමී කොක්ක

380def\_post\_backward\_hook(self,p:nn.Parameter,\*args):

#

පරාමිතියෙන්හසුරුව ඉවත් කරන්න

385p.\_hook\_handle.remove()# type: ignore[attr-defined]386delattr(p,"\_hook\_handle")

#

පසුගාමීසිදුවීමක් හසුරුවන්න

389self.\_backward\_event()

#

මොඩියුලයපසුගාමී කොක්ක

391def\_backward\_hook(self,\*args,\*\*kwargs):

#

පසුගාමීසිදුවීමක් හසුරුවන්න

396self.\_backward\_event()

#

පෙරස්ථරය පරිගණක අනුක්රමික ආරම්භ වනු ඇත. එය පරාම් ලබා ගැනීම අවසන් කර ඇති බවට අප වග බලා ගත යුතුය.

399torch.cuda.current\_stream().wait\_stream(self.fetch\_stream)

#

402returnNone

#

වත්මන්ස්ථරයේ අනුක්රමික උපස්ථ කරන්න

[email protected]\_grad()405def\_backup\_grads(self):

#

පුහුණුකළ හැකි පරාමිතීන් නොමැති නම් මඟ හරින්න

410ifself.chunk\_size[self.TRAINING\_PARAMS\_IDX]==0:411return

#

අනුක්රමිකඋපස්ථ කිරීමට උපස්ථ ධාරාව භාවිතා කරන්න

414withtorch.cuda.stream(self.backup\_stream):

#

අනුක්රමිකගබඩා කිරීමට බෆරය

416buffer=self.\_empty((self.world\_size\*self.chunk\_size[self.TRAINING\_PARAMS\_IDX],))

#

අඛණ්ඩබෆරය නෝඩ් ගණනකට බෙදන්න. මෙම බෙදීම් `බෆර්' පිළිබඳ අදහස් වේ.

418buffers=list(buffer.split(self.chunk\_size[self.TRAINING\_PARAMS\_IDX]))

#

අඛණ්ඩබෆරයේ ඕෆ්සෙට්

421offset=0

#

පුහුණුකළ හැකි පරාමිතීන් හරහා නැවත ක්රියාත්මක කරන්න

423forpinself.param\_refs[self.TRAINING\_PARAMS\_IDX]:

#

අනුක්රමිකඑකතු

425shape=p.\_orig\_shape# type: ignore[attr-defined]426buffer[offset:offset+shape.numel()]=p.grad.view(-1)

#

ඕෆ්සෙට්යාවත්කාලීන කරන්න

428offset+=shape.numel()

#

අනුක්රමිකපිරිසිදු

430p.grad=None

#

වත්මන්ෂාර්ඩ් වල අනුක්රමික සමුච්චය කිරීම සඳහා හිස් ටෙන්සරය

433grad=self.\_empty((self.chunk\_size[self.TRAINING\_PARAMS\_IDX],))

#

එක්එක් ශාකයේ අනුක්රමික සමුච්චය කරන්න. එය නෝඩ් හරහා බෆර විසිරී යන අතර සෑම නෝඩයක්ම එයට ලැබෙන ආතතීන් එකතු වේ (අඩු කරයි).

436dist.reduce\_scatter(grad,buffers)

#

මෙහෙයුමසම්පූර්ණ වන තෙක් රැඳී සිටින්න, පසුව අවරෝධක වෙත යොමු කිරීම් ඉවත් කරන්න

439forbinbuffers:440b.record\_stream(self.fetch\_stream)441buffer.record\_stream(self.fetch\_stream)442delbuffer443delbuffers

#

කුට්ටිඅනුක්රමික සකසන්න. ප්රශස්තකරණය දකින්නේ මෙයයි.

446self.chunk[self.TRAINING\_PARAMS\_IDX].grad=grad447delgrad

#

Zero3Layer ස්ථර සඳහා අනුක්රමික මොඩියුලය

450classZero3Sequential(nn.Module):

#

  • modules``Zero3Layer ස්ථර ලැයිස්තුව
454def\_\_init\_\_(self,modules:List[Zero3Layer]):

#

458super().\_\_init\_\_()

#

පරාමිතීන්ලබා ගැනීමට CUDA ධාරාව

461self.fetch\_stream=torch.cuda.Stream()

#

උපස්ථකිරීමට CUDA ධාරාව (සමුච්චය) අනුක්රමික

463self.backup\_stream=torch.cuda.Stream()

#

එක්එක් Zero3Layer ස්ථරයක් සඳහා ධාරාවන් සහ පෙර සහ ඉදිරියට යන ස්ථර සකසන්න

466foriinrange(len(modules)):

#

ස්ථරදර්ශකය සකසන්න

468modules[i].layer\_idx=i

#

ධාරාවන්සකසන්න

470modules[i].fetch\_stream=self.fetch\_stream471modules[i].backup\_stream=self.backup\_stream

#

ඉදිරියටයන ස්ථර සකසන්න

473ifi+1\<len(modules):474modules[i].next\_layer.append(modules[i+1])

#

පෙරස්ථර සකසන්න

476ifi-1\>=0:477modules[i].prev\_layer.append(modules[i-1])

#

මොඩියුලගබඩා ලැයිස්තුව

480self.module\_list=nn.ModuleList(modules)

#

482defget\_trainable\_chunk(self):

#

එක්එක් ස්ථරයෙන් පුහුණු කළ හැකි කුට්ටි ලැයිස්තුව ආපසු ලබා දෙන්න

484returnsum([m.get\_trainable\_chunk()forminself.module\_list],[])

#

486defforward(self,x:torch.Tensor):

#

ආපසුදක්වා ඵලය අනුක්රමික සම්පූර්ණ බවට වග බලා ගන්න

488torch.cuda.current\_stream().wait\_stream(self.backup\_stream)

#

ඉදිරිසාමාර්ථය

491forminself.module\_list:492x=m(x)

#

495returnx

Trending Research Paperslabml.ai