From 964a28678f9bd1f71cc139e8ab514a74be4cfecf Mon Sep 17 00:00:00 2001 From: "Kai Wang (Victor Kai)" <37533040+kaiwang960112@users.noreply.github.com> Date: Tue, 28 Mar 2023 18:04:53 +0800 Subject: [PATCH] [NFC] polish initializer_3d.py code style (#3279) --- .../context/process_group_initializer/initializer_3d.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/colossalai/context/process_group_initializer/initializer_3d.py b/colossalai/context/process_group_initializer/initializer_3d.py index b752b8f45..1ed8eec86 100644 --- a/colossalai/context/process_group_initializer/initializer_3d.py +++ b/colossalai/context/process_group_initializer/initializer_3d.py @@ -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)