github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-28 16:28:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Ziyue Jiang 2022-03-30 09:35:46 +08:00 committed by GitHub
parent 8c90d4df54
commit 763dc325f1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/nn/layer/colossalai_layer/linear.py
View File

@ -1,4 +1,5 @@
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,