fix parallel_input flag for Linear1D_Col gather_output

666cfd09 · Wesley · アマデウス · a9f778f1 · 666cfd09
Commit 666cfd09 authored Mar 31, 2022 by Wesley Committed by アマデウス Mar 31, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 4 additions and 1 deletion

colossalai/nn/layer/parallel_1d/layers.py colossalai/nn/layer/parallel_1d/layers.py +4 -1

No files found.
--- a/colossalai/nn/layer/parallel_1d/layers.py
+++ b/colossalai/nn/layer/parallel_1d/layers.py
@@ -302,7 +302,10 @@ class Linear1D_Col(ParallelLayer):
        with seed(ParallelMode.TENSOR):
            self.reset_parameters(weight_initializer, bias_initializer)
        self._set_tensor_parallel_attributes()
-        set_parallel_input(True)
+        if self.gather_output:
+            set_parallel_input(False)
+        else:
+            set_parallel_input(True)

    def reset_parameters(self, weight_initializer, bias_initializer) -> None:
        fan_in, fan_out = self.in_features, self.out_features