github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-24 19:17:30 +00:00
Code Issues Packages Projects Releases Wiki Activity

[inference] overlap comm and compute in Linear1D_Row when stream_chunk_num > 1 (#1876)

Browse Source
This commit is contained in:
Jiarui Fang
2022-11-10 17:36:42 +08:00
committed by GitHub
parent 1b494ad73c
commit 986f8cbaa7
2 changed files with 20 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
colossalai/nn/layer/parallel_1d/layers.py
View File

@@ -706,13 +706,22 @@ class Linear1D_Row(ParallelLayer):
input_ = split_forward_gather_backward(input_, ParallelMode.PARALLEL_1D, dim=-1)
if self.stream_chunk_num > 1:
output_parallel_list = [None for i in range(self.stream_chunk_num)]
for i in range(self.stream_chunk_num):
output_parallel_list[i] = F.linear(input_, self.weight_list[i])
output_parallel_list[i] = reduce_input(output_parallel_list[i], ParallelMode.PARALLEL_1D)
output = torch.cat(output_parallel_list, dim=-1)
if self.training:
raise RuntimeError("use stream_chunk_num=1 in Linear1D_Row for training!")
with torch.no_grad():
output_parallel_list = [None for i in range(self.stream_chunk_num)]
handle_list = []
for i in range(self.stream_chunk_num):
output_parallel_list[i] = F.linear(input_, self.weight_list[i])
handle = torch.distributed.all_reduce(output_parallel_list[i],
group=gpc.get_group(ParallelMode.PARALLEL_1D),
async_op=True)
handle_list.append(handle)
# output_parallel_list[i] = reduce_input(output_parallel_list[i], ParallelMode.PARALLEL_1D)
for handle in handle_list:
handle.wait()
output = torch.cat(output_parallel_list, dim=-1)
else:
print(input_.shape, self.weight.shape)
output_parallel = F.linear(input_, self.weight)
# output_parallel = linear_with_async_comm(input_, self.weight, None, ParallelMode.PARALLEL_1D, False)
output = reduce_input(output_parallel, ParallelMode.PARALLEL_1D)