github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish initializer_3d.py code style (#3279)

pull/3313/head
Kai Wang (Victor Kai) 2023-03-28 18:04:53 +08:00 committed by binmakeswell
parent 94eec1c5ad
commit 964a28678f
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/context/process_group_initializer/initializer_3d.py
View File

@ -4,6 +4,7 @@
import math
import torch.distributed as dist
from colossalai.global_variables import tensor_parallel_env as env
from colossalai.registry import DIST_GROUP_INITIALIZER
@ -213,7 +214,8 @@ class Initializer_3D_InputxWeight(ProcessGroupInitializer):
for h in range(self.num_group):
for k in range(self.depth):
ranks = [
h * self.depth**3 + i + self.depth * (j + self.depth * k) for j in range(self.depth)
h * self.depth**3 + i + self.depth * (j + self.depth * k)
for j in range(self.depth)
for i in range(self.depth)
]
group = dist.new_group(ranks)
@ -266,7 +268,8 @@ class Initializer_3D_OutputxWeight(ProcessGroupInitializer):
for h in range(self.num_group):
for j in range(self.depth):
ranks = [
h * self.depth**3 + i + self.depth * (j + self.depth * k) for k in range(self.depth)
h * self.depth**3 + i + self.depth * (j + self.depth * k)
for k in range(self.depth)
for i in range(self.depth)
]
group = dist.new_group(ranks)