[TP] Add gather_out arg to Linear (#541)

colossalai/nn/layer/colossalai_layer/linear.py

 import math
+import inspect
 from typing import Callable
 
 from colossalai.utils import get_current_device
@@ -78,7 +79,11 @@ class Linear(nn.Module):
             if self.layer.bias is not None:
                 bias_initializer(self.layer.bias, fan_in=in_features)
         else:
-            self.layer = _parallel_linear[tensor_parallel](
+            linear_cls = _parallel_linear[tensor_parallel]
+            gather_output = kwargs.pop('gather_output', None)
+            if 'gather_output' in inspect.signature(linear_cls.__init__).parameters.keys(): # gather_out arg is available
+                kwargs['gather_output'] = gather_output
+            self.layer = linear_cls(
                         in_features,
                         out_features,
                         bias=bias,