Unverified Commit 29159d9b authored by Jiarui Fang's avatar Jiarui Fang Committed by GitHub
Browse files

hotfix tensor unittest bugs (#862)

parent 1258af71
......@@ -8,6 +8,7 @@ from colossalai.core import global_context as gpc
from packaging import version
from colossalai.utils.cuda import get_current_device
@colo_op_impl(torch.nn.functional.linear)
def colo_linear(types, args, kwargs, pg):
"""Handles ``__torch_function__`` dispatch for ``torch.nn.functional.linear``.
......@@ -34,7 +35,7 @@ def colo_linear(types, args, kwargs, pg):
elif weight.shard_spec == '1Drow':
# Input:S[1] x Weight:S[0] = Output:P
# All-Reduce(Output) + bias = res
assert divide(input_tensor.shape[-1], gpc.tensor_parallel_size) == weight.size[-1], \
assert divide(input_tensor.shape[-1], gpc.tensor_parallel_size) == weight.size(-1), \
'Invalid shapes in 1Drow forward: input={}, weight={}. Expected last dim of input {}.'.format(
input_tensor.shape, weight.size, weight.size[-1] * gpc.tensor_parallel_size)
# Input:S[1]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment