2022-12-26 07:03:54 +00:00
|
|
|
from copy import deepcopy
|
2022-05-19 04:44:59 +00:00
|
|
|
from typing import Optional
|
2022-12-26 07:03:54 +00:00
|
|
|
|
|
|
|
import torch.nn.functional as F
|
|
|
|
|
|
|
|
from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, ComputeSpec, ReplicaSpec, ShardSpec
|
2022-04-21 06:15:48 +00:00
|
|
|
from colossalai.tensor.op_wrapper import colo_op_impl
|
2022-08-12 03:33:09 +00:00
|
|
|
from colossalai.tensor.sharding_spec import ShardingSpec
|
2022-12-26 07:03:54 +00:00
|
|
|
|
|
|
|
from ._utils import GeneralTensor, convert_to_colo_tensor, reduce_grad, reduce_input
|
2022-04-26 05:23:59 +00:00
|
|
|
|
2022-04-25 02:06:53 +00:00
|
|
|
|
2022-07-13 15:06:12 +00:00
|
|
|
def colo_linear_1drow(input_tensor: ColoTensor, weight: ColoTensor, bias: Optional[ColoTensor]) -> 'ColoTensor':
|
2022-04-28 02:55:40 +00:00
|
|
|
# Input:S[1] x Weight:S[0] = Output:P
|
|
|
|
# All-Reduce(Output) + bias = res
|
2022-04-28 06:43:22 +00:00
|
|
|
# Input:S[1]
|
2022-07-08 06:18:30 +00:00
|
|
|
pg = weight.get_process_group()
|
2022-07-13 15:06:12 +00:00
|
|
|
input_tensor = input_tensor.redistribute(ShardSpec([-1], [weight.get_tp_world_size()]), pg)
|
2022-04-28 02:55:40 +00:00
|
|
|
|
|
|
|
# Output:P
|
2022-05-19 04:44:59 +00:00
|
|
|
partial_output = F.linear(input_tensor, weight)
|
2022-04-28 02:55:40 +00:00
|
|
|
# Reduce(Output)
|
2022-07-08 06:18:30 +00:00
|
|
|
|
|
|
|
output = reduce_input(partial_output, pg)
|
2022-04-28 02:55:40 +00:00
|
|
|
# Bias
|
|
|
|
if bias is not None:
|
2022-06-24 05:08:54 +00:00
|
|
|
assert not bias.has_compute_spec(), 'Invalid bias spec for 1Drow Linear op'
|
2022-05-19 04:44:59 +00:00
|
|
|
output = output + bias
|
|
|
|
|
2022-07-11 07:51:48 +00:00
|
|
|
output = ColoTensor.from_torch_tensor(output, spec=ColoTensorSpec(pg, ReplicaSpec()))
|
2022-04-28 02:55:40 +00:00
|
|
|
return output
|
|
|
|
|
2022-04-28 06:43:22 +00:00
|
|
|
|
2022-07-13 15:06:12 +00:00
|
|
|
def colo_linear_1dcol(input_tensor: ColoTensor, weight: ColoTensor, bias: Optional[ColoTensor]) -> 'ColoTensor':
|
2022-04-28 02:55:40 +00:00
|
|
|
# Input:B x Weight:S[1] + Bias:S[1] = Output:S[1]
|
|
|
|
# All-Gather(Output)
|
|
|
|
# Input:B
|
2022-07-06 08:15:16 +00:00
|
|
|
compute_spec = weight.compute_spec
|
2022-07-11 07:51:48 +00:00
|
|
|
input_tensor = input_tensor.redistribute(ReplicaSpec())
|
2022-07-06 08:15:16 +00:00
|
|
|
input_parallel = reduce_grad(input_tensor, weight.get_process_group())
|
2022-04-28 06:43:22 +00:00
|
|
|
|
2022-05-19 04:44:59 +00:00
|
|
|
output_parallel = F.linear(input_parallel, weight, bias)
|
2022-05-20 12:19:58 +00:00
|
|
|
output = ColoTensor.from_torch_tensor(output_parallel,
|
2022-07-06 08:15:16 +00:00
|
|
|
spec=ColoTensorSpec(weight.get_process_group(),
|
2022-07-11 07:51:48 +00:00
|
|
|
ShardSpec([-1], [weight.get_tp_world_size()]),
|
2022-07-06 08:15:16 +00:00
|
|
|
ComputeSpec(ComputePattern.TP1D)))
|
2022-06-24 05:08:54 +00:00
|
|
|
if compute_spec.output_replicate:
|
|
|
|
return output.to_replicate()
|
|
|
|
else:
|
|
|
|
return output
|
2022-04-28 02:55:40 +00:00
|
|
|
|
2022-04-28 06:43:22 +00:00
|
|
|
|
2022-06-22 08:47:20 +00:00
|
|
|
def colo_linear_1d(mode: str, input_tensor: ColoTensor, weight: ColoTensor, bias: Optional[ColoTensor]) -> 'ColoTensor':
|
2022-05-19 04:44:59 +00:00
|
|
|
assert mode in ('row', 'col')
|
2022-07-13 15:06:12 +00:00
|
|
|
funcs = {'row': colo_linear_1drow, 'col': colo_linear_1dcol}
|
2022-05-19 04:44:59 +00:00
|
|
|
return funcs[mode](input_tensor, weight, bias)
|
|
|
|
|
|
|
|
|
2022-07-06 08:15:16 +00:00
|
|
|
# @register_colo_graph(input_pos=[1], param_pos=[2, 3])
|
2022-06-22 08:47:20 +00:00
|
|
|
def colo_linear_imp(input_tensor: GeneralTensor,
|
|
|
|
weight: GeneralTensor,
|
|
|
|
bias: Optional[GeneralTensor] = None) -> 'ColoTensor':
|
2022-04-21 03:42:37 +00:00
|
|
|
"""Handles ``__torch_function__`` dispatch for ``torch.nn.functional.linear``.
|
|
|
|
This method computes a linear.
|
|
|
|
"""
|
2022-07-06 08:15:16 +00:00
|
|
|
assert isinstance(weight, ColoTensor)
|
|
|
|
pg = weight.get_process_group()
|
2022-07-07 10:09:18 +00:00
|
|
|
assert pg
|
2022-07-06 08:15:16 +00:00
|
|
|
input_tensor = convert_to_colo_tensor(input_tensor, pg)
|
|
|
|
bias = convert_to_colo_tensor(bias, pg)
|
|
|
|
# input_tensor, weight, bias = tuple(map(convert_to_colo_tensor, (input_tensor, weight, bias)))
|
2022-04-28 06:43:22 +00:00
|
|
|
|
2022-04-28 02:55:40 +00:00
|
|
|
# Add communication logic before and after linear call.
|
2022-05-07 09:10:37 +00:00
|
|
|
ret_tensor = None
|
2022-06-24 05:08:54 +00:00
|
|
|
if not weight.has_compute_spec(): # No Model Parallel Applied
|
2022-07-06 08:15:16 +00:00
|
|
|
assert weight.is_replicate(), 'Invalid weight spec for native Linear op'
|
|
|
|
assert bias is None or bias.is_replicate(), 'Invalid bias spec for native Linear op'
|
2022-07-07 10:09:18 +00:00
|
|
|
ret_tensor = ColoTensor.from_torch_tensor(F.linear(input_tensor, weight, bias), spec=ColoTensorSpec(pg))
|
2022-07-06 08:15:16 +00:00
|
|
|
elif weight.has_compute_pattern(ComputePattern.TP1D): # Single Model Parallel Applied
|
|
|
|
if weight.is_shard_1dcol() and (bias is None or bias.is_replicate()):
|
2022-05-19 04:44:59 +00:00
|
|
|
mode = 'row'
|
2022-07-06 08:15:16 +00:00
|
|
|
elif weight.is_shard_1drow() and (bias is None or bias.is_shard_1drow() or bias.is_shard_1dcol()):
|
2022-05-19 04:44:59 +00:00
|
|
|
mode = 'col'
|
2022-04-24 05:43:12 +00:00
|
|
|
else:
|
2022-07-06 08:15:16 +00:00
|
|
|
raise RuntimeError(f"the weight or bias tensor spec is not valid, weight {weight}, bias {bias}")
|
2022-05-19 04:44:59 +00:00
|
|
|
ret_tensor = colo_linear_1d(mode, input_tensor, weight, bias)
|
2022-04-21 03:42:37 +00:00
|
|
|
else:
|
2022-04-28 02:55:40 +00:00
|
|
|
raise NotImplementedError
|
2022-05-07 09:10:37 +00:00
|
|
|
|
|
|
|
return ret_tensor
|
2022-06-22 08:47:20 +00:00
|
|
|
|
|
|
|
|
2022-08-12 03:33:09 +00:00
|
|
|
def _new_colo_linear_imp(input_tensor: GeneralTensor,
|
|
|
|
weight: GeneralTensor,
|
|
|
|
bias: Optional[GeneralTensor] = None) -> 'ColoTensor':
|
|
|
|
"""
|
|
|
|
A tentative function to compute the distributed linear layer with the latest sharding spec.
|
|
|
|
This function is subject to future change as the current sharding API is not stable.
|
|
|
|
"""
|
|
|
|
# get mesh info
|
|
|
|
input_sharding_seq = input_tensor.sharding_spec.sharding_sequence
|
|
|
|
weight_sharding_seq = weight.sharding_spec.sharding_sequence
|
|
|
|
if bias is not None:
|
|
|
|
bias_sharding_seq = bias.sharding_spec.sharding_sequence
|
|
|
|
device_mesh = weight.sharding_spec.device_mesh
|
|
|
|
pg_axis0 = weight.pg_axis0
|
|
|
|
pg_axis1 = weight.pg_axis1
|
|
|
|
|
|
|
|
# the last dim of input should have the same spec as the first dim of weight
|
|
|
|
# the weight is transposed, so we look at the second dimension
|
|
|
|
assert input_sharding_seq[-1] == weight_sharding_seq[1]
|
|
|
|
|
|
|
|
if bias is not None:
|
|
|
|
assert bias_sharding_seq[0] == weight_sharding_seq[0]
|
|
|
|
|
|
|
|
# compute the output sharding sequence
|
|
|
|
# as weight is transposed, so we look at the first dimension
|
|
|
|
output_shard_seq = input_sharding_seq[:-1] + weight_sharding_seq[:1]
|
|
|
|
output_shard_seq = deepcopy(output_shard_seq)
|
|
|
|
|
|
|
|
# TODO: add reduce grad logic
|
|
|
|
|
|
|
|
# handle column and row parallel linear
|
|
|
|
# by reusing the implementation above
|
|
|
|
out = F.linear(input_tensor, weight)
|
|
|
|
|
|
|
|
# run all reduce if necessary
|
|
|
|
last_dim_spec = input_sharding_seq[-1]
|
|
|
|
if last_dim_spec.is_replica:
|
|
|
|
pass
|
|
|
|
elif last_dim_spec.shard_list is not None:
|
|
|
|
for dim in last_dim_spec.shard_list:
|
|
|
|
if dim == 0:
|
|
|
|
reduce_input(out, pg_axis0)
|
|
|
|
elif dim == 1:
|
|
|
|
reduce_input(out, pg_axis1)
|
|
|
|
else:
|
|
|
|
raise RuntimeError("Found invalid sharding axis {dim}, only 0 or 1 is expected")
|
|
|
|
# add bias
|
|
|
|
if bias is not None:
|
|
|
|
out += bias
|
|
|
|
|
|
|
|
# convert shard seq to partition dict
|
|
|
|
output_partition_dict = {}
|
|
|
|
for index, dim_spec in enumerate(output_shard_seq):
|
|
|
|
if not dim_spec.is_replica:
|
|
|
|
if index not in output_partition_dict:
|
|
|
|
output_partition_dict[index] = []
|
|
|
|
output_partition_dict[index].extend(dim_spec.shard_list)
|
|
|
|
|
|
|
|
entire_shape = out.shape
|
|
|
|
output_sharding_spec = ShardingSpec(device_mesh, entire_shape, output_partition_dict)
|
|
|
|
ret_tensor = ColoTensor.from_torch_tensor(out)
|
|
|
|
setattr(ret_tensor, 'sharding_spec', output_sharding_spec)
|
|
|
|
return ret_tensor
|
|
|
|
|
|
|
|
|
|
|
|
def _has_sharding_spec(tensor):
|
|
|
|
"""
|
2022-12-26 07:03:54 +00:00
|
|
|
A tentative function to check whether the tensor is using the new sharding spec API. We assume that the sharding spec object is
|
2022-08-12 03:33:09 +00:00
|
|
|
set as the attribute `sharding_spec` on a tensor.
|
|
|
|
"""
|
|
|
|
return hasattr(tensor, 'sharding_spec')
|
|
|
|
|
|
|
|
|
2022-06-22 08:47:20 +00:00
|
|
|
@colo_op_impl(F.linear)
|
2022-12-26 07:03:54 +00:00
|
|
|
def colo_linear(input: GeneralTensor, weight: GeneralTensor, bias: Optional[GeneralTensor] = None) -> 'ColoTensor':
|
2022-08-12 03:33:09 +00:00
|
|
|
if _has_sharding_spec(weight):
|
2022-12-26 07:03:54 +00:00
|
|
|
return _new_colo_linear_imp(input, weight, bias)
|
2022-08-12 03:33:09 +00:00
|
|
|
else:
|
2022-12-26 07:03:54 +00:00
|
|
|
return colo_linear_imp(input, weight, bias)
|