ColossalAI/colossalai/nn/_ops/linear.py

from copy import deepcopy
from typing import Optional

import torch.nn.functional as F

from colossalai.tensor import ColoTensor, ColoTensorSpec, ComputePattern, ComputeSpec, ReplicaSpec, ShardSpec
from colossalai.tensor.op_wrapper import colo_op_impl
from colossalai.tensor.sharding_spec import ShardingSpec

from ._utils import GeneralTensor, convert_to_colo_tensor, reduce_grad, reduce_input


def colo_linear_1drow(input_tensor: ColoTensor, weight: ColoTensor, bias: Optional[ColoTensor]) -> 'ColoTensor':
    # Input:S[1] x Weight:S[0] = Output:P
    # All-Reduce(Output) + bias = res
    # Input:S[1]
    pg = weight.get_process_group()
    input_tensor = input_tensor.redistribute(ShardSpec([-1], [weight.get_tp_world_size()]), pg)

    # Output:P
    partial_output = F.linear(input_tensor, weight)
    # Reduce(Output)

    output = reduce_input(partial_output, pg)
    # Bias
    if bias is not None:
        assert not bias.has_compute_spec(), 'Invalid bias spec for 1Drow Linear op'
        output = output + bias

    output = ColoTensor.from_torch_tensor(output, spec=ColoTensorSpec(pg, ReplicaSpec()))
    return output


def colo_linear_1dcol(input_tensor: ColoTensor, weight: ColoTensor, bias: Optional[ColoTensor]) -> 'ColoTensor':
    # Input:B x Weight:S[1] + Bias:S[1] = Output:S[1]
    # All-Gather(Output)
    # Input:B
    compute_spec = weight.compute_spec
    input_tensor = input_tensor.redistribute(ReplicaSpec())
    input_parallel = reduce_grad(input_tensor, weight.get_process_group())

    output_parallel = F.linear(input_parallel, weight, bias)
    output = ColoTensor.from_torch_tensor(output_parallel,
                                          spec=ColoTensorSpec(weight.get_process_group(),
                                                              ShardSpec([-1], [weight.get_tp_world_size()]),
                                                              ComputeSpec(ComputePattern.TP1D)))
    if compute_spec.output_replicate:
        return output.to_replicate()
    else:
        return output


def colo_linear_1d(mode: str, input_tensor: ColoTensor, weight: ColoTensor, bias: Optional[ColoTensor]) -> 'ColoTensor':
    assert mode in ('row', 'col')
    funcs = {'row': colo_linear_1drow, 'col': colo_linear_1dcol}
    return funcs[mode](input_tensor, weight, bias)


# @register_colo_graph(input_pos=[1], param_pos=[2, 3])
def colo_linear_imp(input_tensor: GeneralTensor,
                    weight: GeneralTensor,
                    bias: Optional[GeneralTensor] = None) -> 'ColoTensor':
    """Handles ``__torch_function__`` dispatch for ``torch.nn.functional.linear``.
    This method computes a linear.
    """
    assert isinstance(weight, ColoTensor)
    pg = weight.get_process_group()
    assert pg
    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)))

    # Add communication logic before and after linear call.
    ret_tensor = None
    if not weight.has_compute_spec():    # No Model Parallel Applied
        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'
        ret_tensor = ColoTensor.from_torch_tensor(F.linear(input_tensor, weight, bias), spec=ColoTensorSpec(pg))
    elif weight.has_compute_pattern(ComputePattern.TP1D):    # Single Model Parallel Applied
        if weight.is_shard_1dcol() and (bias is None or bias.is_replicate()):
            mode = 'row'
        elif weight.is_shard_1drow() and (bias is None or bias.is_shard_1drow() or bias.is_shard_1dcol()):
            mode = 'col'
        else:
            raise RuntimeError(f"the weight or bias tensor spec is not valid, weight {weight}, bias {bias}")
        ret_tensor = colo_linear_1d(mode, input_tensor, weight, bias)
    else:
        raise NotImplementedError

    return ret_tensor


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):
    """
    A tentative function to check whether the tensor is using the new sharding spec API. We assume that the sharding spec object is
    set as the attribute `sharding_spec` on a tensor.
    """
    return hasattr(tensor, 'sharding_spec')


@colo_op_impl(F.linear)
def colo_linear(input: GeneralTensor, weight: GeneralTensor, bias: Optional[GeneralTensor] = None) -> 'ColoTensor':
    if _has_sharding_spec(weight):
        return _new_colo_linear_imp(input, weight, bias)
    else:
        return colo_linear_imp(input, weight, bias)