mirror of https://github.com/hpcaitech/ColossalAI
135 lines
6.6 KiB
Python
135 lines
6.6 KiB
Python
import torch
|
|
from colossalai.tensor.op_wrapper import colo_op_impl
|
|
from colossalai.context import ParallelMode
|
|
from colossalai.nn.layer.parallel_1d._utils import split_forward_gather_backward, reduce_input, \
|
|
gather_forward_split_backward, reduce_grad
|
|
from colossalai.nn.layer.utils import divide
|
|
from colossalai.core import global_context as gpc
|
|
from packaging import version
|
|
from colossalai.tensor import ComputePattern, TensorSpec, ComputePattern, ParallelAction, ColoTensor
|
|
|
|
|
|
@colo_op_impl(torch.nn.functional.linear)
|
|
def colo_linear(types, args, kwargs, pg):
|
|
"""Handles ``__torch_function__`` dispatch for ``torch.nn.functional.linear``.
|
|
This method computes a linear.
|
|
"""
|
|
input_tensor = args[0]
|
|
weight = args[1]
|
|
|
|
if version.parse(torch.__version__) > version.parse("1.11.0"):
|
|
if len(args) == 3:
|
|
bias = args[2]
|
|
else:
|
|
bias = None
|
|
else:
|
|
bias = kwargs.get('bias', None)
|
|
|
|
bias_spec = None
|
|
if isinstance(bias, ColoTensor):
|
|
bias_spec = bias.shard_spec
|
|
bias = bias.torch_tensor()
|
|
|
|
# Add communication logic before and after linear call.
|
|
if isinstance(weight, ColoTensor):
|
|
if weight.shard_spec == None or weight.shard_spec.num_action == 0:
|
|
assert bias_spec == None or bias_spec.num_action == 0, 'Invalid bias spec for native Linear op'
|
|
if isinstance(input_tensor, ColoTensor):
|
|
input_tensor = input_tensor.torch_tensor()
|
|
if isinstance(weight, ColoTensor):
|
|
weight = weight.torch_tensor()
|
|
return ColoTensor.init_from_torch_tensor(torch.nn.functional.linear(input_tensor, weight, bias))
|
|
elif weight.shard_spec.num_action == 1:
|
|
parallel_action = weight.shard_spec.get_action_by_compute_pattern(ComputePattern.TP1DRow)
|
|
compute_patterns = weight.shard_spec.compute_patterns
|
|
if ComputePattern.TP1DRow in compute_patterns:
|
|
# Input:S[1] x Weight:S[0] = Output:P
|
|
# All-Reduce(Output) + bias = res
|
|
# Input:S[1]
|
|
input_spec = None
|
|
if isinstance(input_tensor, ColoTensor):
|
|
input_spec = input_tensor.shard_spec
|
|
input_tensor = input_tensor.torch_tensor()
|
|
|
|
if input_spec == None or input_spec.num_action == 0:
|
|
# Not splited yet.
|
|
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_per_partition = split_forward_gather_backward(input_tensor, parallel_action.parallel_mode, dim=-1)
|
|
elif input_tensor.shard_spec.num_action == 1:
|
|
if ComputePattern.TP1DCol in input_spec.compute_patterns:
|
|
# Splited by 1Dcol
|
|
assert input_tensor.shape[-1] == weight.size(-1), \
|
|
'Invalid shapes in 1Drow forward: input={}, weight={}. Expected last dim of input {}.'.format(
|
|
input_tensor.shape, weight.size, weight.size(-1))
|
|
input_per_partition = input_tensor
|
|
else:
|
|
raise NotImplementedError
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
# Output:P
|
|
weight_ = weight.torch_tensor()
|
|
partial_output = torch.nn.functional.linear(input_per_partition, weight_)
|
|
# Reduce(Output)
|
|
output = reduce_input(partial_output, parallel_action.parallel_mode)
|
|
# Bias
|
|
if bias is not None:
|
|
assert bias_spec == None or bias_spec.num_action == 0, 'Invalid bias spec for 1Drow Linear op'
|
|
output = output + bias
|
|
output = ColoTensor.init_from_torch_tensor(output)
|
|
return output
|
|
elif ComputePattern.TP1DCol in compute_patterns:
|
|
# Input:B x Weight:S[1] + Bias:S[1] = Output:S[1]
|
|
# All-Gather(Output)
|
|
# Input:B
|
|
input_spec = None
|
|
output_spec = None
|
|
parallel_action = weight.shard_spec.get_action_by_compute_pattern(ComputePattern.TP1DCol)
|
|
if isinstance(input_tensor, ColoTensor):
|
|
input_spec = input_tensor.shard_spec
|
|
input_tensor = input_tensor.torch_tensor()
|
|
|
|
if input_spec == None or input_spec.num_action == 0:
|
|
# Not splited yet.
|
|
assert input_tensor.shape[-1] == weight.size(-1), \
|
|
'Invalid shapes in 1Dcol forward: input={}, weight={}. Expected last dim of input {}.'.format(
|
|
input_tensor.shape, weight.size, weight.size(-1))
|
|
input_parallel = reduce_grad(input_tensor, parallel_action.parallel_mode)
|
|
else:
|
|
raise NotImplementedError
|
|
# Bias:S[1]
|
|
if bias is not None:
|
|
assert bias_spec is not None and bias_spec.num_action == 1 and \
|
|
ComputePattern.TP1DCol in bias_spec.compute_patterns, \
|
|
'Invalid bias spec for 1Dcol Linear op'
|
|
|
|
weight_ = weight.torch_tensor()
|
|
output_parallel = torch.nn.functional.linear(input_parallel, weight_, bias)
|
|
|
|
if parallel_action.gather_out:
|
|
# All-Gather(Output)
|
|
output = gather_forward_split_backward(output_parallel, parallel_action.parallel_mode, dim=-1)
|
|
output = ColoTensor.init_from_torch_tensor(output)
|
|
else:
|
|
output = ColoTensor.init_from_torch_tensor(output_parallel)
|
|
out_parallel_action_list = [
|
|
ParallelAction(
|
|
priority=1, compute_pattern=ComputePattern.TP1DCol,
|
|
parallel_mode=parallel_action.parallel_mode
|
|
)
|
|
]
|
|
output_spec = TensorSpec(out_parallel_action_list)
|
|
# set ColoTensor spec
|
|
if output_spec is not None:
|
|
output.set_spec(output_spec)
|
|
return output
|
|
|
|
else:
|
|
raise NotImplementedError
|
|
else:
|
|
raise NotImplementedError
|
|
else:
|
|
return torch.nn.functional.linear(input_tensor, weight, bias)
|