github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-07-17 17:02:04 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix parallel_input flag for Linear1D_Col gather_output

Browse Source
This commit is contained in:
Wesley 2022-03-31 16:38:14 +08:00 committed by アマデウス
parent a9f778f1b1
commit 666cfd094a
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -302,6 +302,9 @@ class Linear1D_Col(ParallelLayer):
with seed(ParallelMode.TENSOR):
self.reset_parameters(weight_initializer, bias_initializer)
self._set_tensor_parallel_attributes()
if self.gather_output:
set_parallel_input(False)
else:
set_parallel_input(True)
def reset_parameters(self, weight_initializer, bias_initializer) -> None: