fix parallel_input flag for Linear1D_Col gather_output

2022-03-31 16:38:14 +08:00 · 2022-03-31 16:38:14 +08:00 · 666cfd094a
parent a9f778f1b1
commit 666cfd094a
1 changed files with 4 additions and 1 deletions
--- a/colossalai/nn/layer/parallel_1d/layers.py
+++ b/colossalai/nn/layer/parallel_1d/layers.py
@ -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: