2023-06-20 03:45:16 +00:00
|
|
|
#!/usr/bin/env python
|
|
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
|
|
|
|
import math
|
2023-07-20 02:39:06 +00:00
|
|
|
from typing import Callable, List, Optional, Tuple, Union
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
import torch
|
|
|
|
import torch.distributed as dist
|
|
|
|
import torch.nn as nn
|
|
|
|
from torch import Tensor
|
|
|
|
from torch.distributed import ProcessGroup
|
|
|
|
from torch.nn.parameter import Parameter
|
|
|
|
|
2023-07-10 02:48:53 +00:00
|
|
|
from colossalai.lazy import LazyInitContext
|
2023-06-20 03:45:16 +00:00
|
|
|
from colossalai.nn import init as init
|
|
|
|
from colossalai.nn.layer.utils import divide
|
2023-06-23 08:07:09 +00:00
|
|
|
from colossalai.tensor.d_tensor.api import (
|
2023-07-20 02:39:06 +00:00
|
|
|
customized_distributed_tensor_to_existing_param,
|
2023-06-23 08:07:09 +00:00
|
|
|
distribute_tensor_with_customization,
|
2023-07-20 02:39:06 +00:00
|
|
|
is_customized_distributed_tensor,
|
|
|
|
is_distributed_tensor,
|
2023-06-23 08:07:09 +00:00
|
|
|
shard_rowwise,
|
2023-07-20 02:39:06 +00:00
|
|
|
sharded_tensor_to_existing_param,
|
2023-06-23 08:07:09 +00:00
|
|
|
)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
from ._operation import (
|
|
|
|
gather_forward_split_backward,
|
2023-08-16 07:41:20 +00:00
|
|
|
linear_reducescatter_forward_gather_backward,
|
2023-07-14 07:56:59 +00:00
|
|
|
linear_with_async_comm,
|
2023-08-16 07:41:20 +00:00
|
|
|
matmul_gather_forward_reducescatter_backward,
|
2023-06-23 08:07:09 +00:00
|
|
|
matmul_with_async_comm,
|
|
|
|
reduce_backward,
|
|
|
|
reduce_forward,
|
2023-06-20 03:45:16 +00:00
|
|
|
split_forward_gather_backward,
|
|
|
|
)
|
2023-06-21 06:30:06 +00:00
|
|
|
from .parallel_module import ParallelModule
|
2023-06-20 03:45:16 +00:00
|
|
|
from .utils import create_randomizer_with_offset
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
__all__ = ["FusedLinear1D_Col", "FusedLinear1D_Row", "GPT2FusedLinearConv1D_Col", "GPT2FusedLinearConv1D_Row"]
|
2023-06-20 03:45:16 +00:00
|
|
|
|
2023-06-28 07:04:35 +00:00
|
|
|
# ====================================
|
|
|
|
# For GPT Only
|
|
|
|
# ====================================
|
2023-06-20 03:45:16 +00:00
|
|
|
|
2023-06-28 07:04:35 +00:00
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
def split_fused_qkv_in_gpt2_style(
|
|
|
|
qkv: torch.Tensor, n_fused: int, process_group: ProcessGroup, is_transposed: bool = False
|
|
|
|
):
|
2023-06-23 08:07:09 +00:00
|
|
|
"""
|
|
|
|
The fused qkv tensor looks like [Q1, Q2, K1, K2, V1, V2], this function will split them into [Q1, K1, V1] and [Q2, K2, V2].
|
2023-06-28 07:04:35 +00:00
|
|
|
|
|
|
|
Args:
|
|
|
|
qkv (torch.Tensor): The fused qkv tensor.
|
|
|
|
n_fused (int): The number items fused together, defaults to 3 (query, key and value).
|
|
|
|
process_group (ProcessGroup): The process group for distributed communication.
|
|
|
|
is_transposed (bool): generally the tensor is the shape of (out_features, in_features). Set this to True if the tensor is in the shape (in_features, out_features).
|
2023-06-23 08:07:09 +00:00
|
|
|
"""
|
|
|
|
# get the number of slice for the fused qkv
|
|
|
|
rank = dist.get_rank(group=process_group)
|
|
|
|
world_size = dist.get_world_size(group=process_group)
|
|
|
|
order = torch.arange(world_size * n_fused)
|
|
|
|
|
|
|
|
# split the fused qkv
|
|
|
|
# from
|
|
|
|
# [Q, K, V]
|
|
|
|
# to
|
|
|
|
# [Q1, Q2, K1, K2, V1, V2]
|
2023-06-28 07:04:35 +00:00
|
|
|
if is_transposed:
|
|
|
|
weight_chunks = torch.chunk(qkv, world_size * n_fused, dim=-1)
|
|
|
|
else:
|
|
|
|
weight_chunks = torch.chunk(qkv, world_size * n_fused, dim=0)
|
2023-06-23 08:07:09 +00:00
|
|
|
|
|
|
|
# rearrange the slice into the final order
|
|
|
|
# from
|
|
|
|
# [Q1, Q2, K1, K2, V1, V2]
|
|
|
|
# to
|
|
|
|
# [Q1, K1, V1], [Q2, K2, V2]
|
|
|
|
weight_chunks_of_current_rank = [weight_chunks[i] for i in order[rank::world_size]]
|
2023-06-28 07:04:35 +00:00
|
|
|
|
|
|
|
if is_transposed:
|
|
|
|
weight_of_current_rank = torch.cat(weight_chunks_of_current_rank, dim=-1)
|
|
|
|
else:
|
|
|
|
weight_of_current_rank = torch.cat(weight_chunks_of_current_rank, dim=0)
|
2023-06-23 08:07:09 +00:00
|
|
|
return weight_of_current_rank
|
|
|
|
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
def gather_fused_qkv_in_gpt2_style(
|
|
|
|
qkv: torch.Tensor, n_fused: int, process_group: ProcessGroup, is_transposed: bool = False
|
|
|
|
):
|
2023-06-23 08:07:09 +00:00
|
|
|
"""
|
|
|
|
The splitted qkv tensor looks like [Q1, K1, V1] and [Q2, K2, V2], this function will gather them into [Q1, Q2, K1, K2, V1, V2].
|
2023-06-28 07:04:35 +00:00
|
|
|
|
|
|
|
Args:
|
|
|
|
qkv (torch.Tensor): The fused qkv tensor.
|
|
|
|
n_fused (int): The number items fused together, defaults to 3 (query, key and value).
|
|
|
|
process_group (ProcessGroup): The process group for distributed communication.
|
|
|
|
is_transposed (bool): generally the tensor is the shape of (out_features, in_features). Set this to True if the tensor is in the shape (in_features, out_features).
|
2023-06-23 08:07:09 +00:00
|
|
|
"""
|
|
|
|
world_size = dist.get_world_size(group=process_group)
|
|
|
|
|
|
|
|
# gather the tensors
|
|
|
|
# from
|
|
|
|
# [Q1, K1, V1], [Q2, K2, V2]
|
|
|
|
# to
|
|
|
|
# [Q1, K1, V1, Q2, K2, V2]
|
|
|
|
origin_device = qkv.device
|
|
|
|
qkv = qkv.cuda()
|
|
|
|
gather_list = [torch.zeros_like(qkv) for _ in range(world_size)]
|
|
|
|
dist.all_gather(gather_list, qkv, group=process_group)
|
2023-06-28 07:04:35 +00:00
|
|
|
|
|
|
|
if is_transposed:
|
|
|
|
gather_weight = torch.cat(gather_list, dim=-1)
|
|
|
|
else:
|
|
|
|
gather_weight = torch.cat(gather_list, dim=0)
|
2023-06-23 08:07:09 +00:00
|
|
|
gather_weight = gather_weight.to(origin_device)
|
|
|
|
qkv = qkv.to(origin_device)
|
|
|
|
|
|
|
|
# rearrange the tensor slices
|
|
|
|
# from
|
|
|
|
# [Q1, K1, V1, Q2, K2, V2]
|
|
|
|
# to
|
|
|
|
# [Q1, Q2, K1, K2, V1, V2]
|
2023-06-28 07:04:35 +00:00
|
|
|
if is_transposed:
|
|
|
|
weight_chunks = torch.chunk(gather_weight, world_size * n_fused, dim=-1)
|
|
|
|
else:
|
|
|
|
weight_chunks = torch.chunk(gather_weight, world_size * n_fused, dim=0)
|
|
|
|
|
2023-06-23 08:07:09 +00:00
|
|
|
reordered_chunk_list = []
|
|
|
|
for i in range(n_fused):
|
|
|
|
reordered_chunk_list.extend(weight_chunks[i::n_fused])
|
2023-06-28 07:04:35 +00:00
|
|
|
|
|
|
|
if is_transposed:
|
|
|
|
reordered_gather_weight = torch.cat(reordered_chunk_list, dim=-1)
|
|
|
|
else:
|
|
|
|
reordered_gather_weight = torch.cat(reordered_chunk_list, dim=0)
|
2023-06-23 08:07:09 +00:00
|
|
|
return reordered_gather_weight
|
|
|
|
|
|
|
|
|
2023-06-28 07:04:35 +00:00
|
|
|
class GPT2FusedLinearConv1D_Col(ParallelModule):
|
2023-06-20 03:45:16 +00:00
|
|
|
r"""Linear layer with column parallelism.
|
|
|
|
|
|
|
|
The linear layer is defined as :math:`Y = XA + b`. A is parallelized along
|
2023-06-23 08:07:09 +00:00
|
|
|
its second dimension as :math:`A = [A_1, ..., A_p]`. This layer is used to fit `Conv1D` layer (Fused QKV) in gpt2 of huggingface.
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
Args:
|
|
|
|
in_features (int): size of each input sample.
|
|
|
|
out_features (int): size of each output sample.
|
|
|
|
bias (bool, optional): If set to ``False``, the layer will not learn an additive bias, defaults to ``True``.
|
|
|
|
dtype (`torch.dtype`): The dtype of parameters, defaults to None.
|
|
|
|
device (`torch.device`): The device of parameters, defaults to None.
|
2023-06-23 08:07:09 +00:00
|
|
|
n_fused (int): The number items fused, defaults to 3 (QKV).
|
2023-06-20 03:45:16 +00:00
|
|
|
process_group (`torch.distributed.ProcessGroup`): The process group to be used for weight sharding and communication, defaults to None.
|
2023-08-16 07:41:20 +00:00
|
|
|
seq_parallel (`bool`): If set to ``True``, it will use sequence parallel, defaults to False.
|
2023-06-20 03:45:16 +00:00
|
|
|
gather_output (bool, optional): If true, call all-gather on output and make Y available
|
|
|
|
to all GPUs, otherwise, every GPU will have its output
|
|
|
|
which is :math:`Y_i = XA_i`, defaults to False
|
|
|
|
skip_bias_add (bool): If set to ``True``, it will skip bias add for linear layer,
|
|
|
|
which is preserved for kernel fusion, defaults to False
|
|
|
|
weight_initializer (`typing.Callable`):
|
|
|
|
The initializer of weight, defaults to kaiming uniform initializer.
|
|
|
|
bias_initializer (`typing.Callable`):
|
|
|
|
The initializer of bias, defaults to xavier uniform initializer.
|
|
|
|
|
|
|
|
More details about ``initializer`` please refer to
|
|
|
|
`init <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/nn/init.py>`_.
|
|
|
|
"""
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
in_features: int,
|
|
|
|
out_features: int,
|
|
|
|
bias: bool = True,
|
|
|
|
dtype: torch.dtype = None,
|
|
|
|
device: torch.device = None,
|
|
|
|
process_group: ProcessGroup = None,
|
|
|
|
async_communication: bool = False,
|
|
|
|
gather_output: bool = False,
|
|
|
|
seq_parallel: bool = False,
|
|
|
|
overlap: bool = False,
|
|
|
|
skip_bias_add: bool = False,
|
|
|
|
n_fused: int = 3,
|
|
|
|
weight: Optional[Parameter] = None,
|
|
|
|
bias_: Optional[Parameter] = None,
|
|
|
|
weight_initializer: Callable = init.kaiming_uniform_(a=math.sqrt(5)),
|
|
|
|
bias_initializer: Callable = init.xavier_uniform_(a=1, scale=1),
|
|
|
|
):
|
2023-06-20 03:45:16 +00:00
|
|
|
super().__init__()
|
|
|
|
|
|
|
|
# Keep input parameters
|
|
|
|
self.in_features = in_features
|
|
|
|
self.out_features = out_features
|
|
|
|
self.gather_output = gather_output
|
2023-08-16 07:41:20 +00:00
|
|
|
self.seq_parallel = seq_parallel
|
2023-08-29 10:30:50 +00:00
|
|
|
self.overlap = overlap
|
2023-06-20 03:45:16 +00:00
|
|
|
self.skip_bias_add = skip_bias_add
|
|
|
|
self.device = device
|
2023-06-23 08:07:09 +00:00
|
|
|
self.n_fused = n_fused
|
2023-06-20 03:45:16 +00:00
|
|
|
self.process_group = process_group
|
2023-06-23 08:07:09 +00:00
|
|
|
self.async_communication = async_communication
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
if skip_bias_add and not bias:
|
2023-09-19 06:20:26 +00:00
|
|
|
raise ValueError("cannot skip bias addition if bias is None")
|
2023-06-20 03:45:16 +00:00
|
|
|
|
2023-07-20 02:39:06 +00:00
|
|
|
# offset the seed with randomizer index and rank
|
|
|
|
seed = torch.random.initial_seed()
|
|
|
|
self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group)
|
|
|
|
|
|
|
|
# sanity check
|
|
|
|
if weight is not None:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert not bias or bias_ is not None, "bias_ must be provided if bias is True when weight is not None"
|
2023-07-20 02:39:06 +00:00
|
|
|
else:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert bias_ is None, "bias_ must be None if weight is None"
|
2023-07-20 02:39:06 +00:00
|
|
|
|
2023-06-20 03:45:16 +00:00
|
|
|
# Parameters.
|
2023-07-20 02:39:06 +00:00
|
|
|
if weight is None:
|
|
|
|
# Initialize weight.
|
2023-09-19 06:20:26 +00:00
|
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
2023-07-20 02:39:06 +00:00
|
|
|
self.weight = Parameter(torch.empty(self.in_features, self.out_features, **factory_kwargs))
|
|
|
|
else:
|
|
|
|
weight.data = weight.data.to(device=device, dtype=dtype)
|
|
|
|
self.weight = weight
|
2023-06-23 08:07:09 +00:00
|
|
|
|
|
|
|
def shard_fn(tensor):
|
2023-06-28 07:04:35 +00:00
|
|
|
return split_fused_qkv_in_gpt2_style(tensor, self.n_fused, self.process_group, True)
|
2023-06-23 08:07:09 +00:00
|
|
|
|
|
|
|
def gather_fn(tensor):
|
2023-07-20 02:39:06 +00:00
|
|
|
return gather_fused_qkv_in_gpt2_style(tensor, self.n_fused, self.process_group, True)
|
2023-06-23 08:07:09 +00:00
|
|
|
|
2023-07-20 02:39:06 +00:00
|
|
|
if not is_customized_distributed_tensor(self.weight):
|
|
|
|
with torch.no_grad():
|
|
|
|
sharded_weight = distribute_tensor_with_customization(self.weight.data, shard_fn, gather_fn)
|
|
|
|
customized_distributed_tensor_to_existing_param(sharded_weight, self.weight)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
if bias:
|
2023-07-20 02:39:06 +00:00
|
|
|
if bias_ is None:
|
|
|
|
self.bias = Parameter(torch.empty(self.out_features, **factory_kwargs))
|
|
|
|
else:
|
|
|
|
bias_.data = bias_.data.to(device=device, dtype=dtype)
|
|
|
|
self.bias = bias_
|
|
|
|
if not is_customized_distributed_tensor(self.bias):
|
|
|
|
with torch.no_grad():
|
|
|
|
sharded_bias = distribute_tensor_with_customization(self.bias.data, shard_fn, gather_fn)
|
|
|
|
customized_distributed_tensor_to_existing_param(sharded_bias, self.bias)
|
2023-06-20 03:45:16 +00:00
|
|
|
else:
|
|
|
|
self.bias = None
|
|
|
|
|
2023-07-20 02:39:06 +00:00
|
|
|
if weight is None:
|
|
|
|
# init weights
|
|
|
|
self.reset_parameters(weight_initializer, bias_initializer)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
@staticmethod
|
2023-09-19 06:20:26 +00:00
|
|
|
def from_native_module(
|
|
|
|
module: nn.Module, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, **kwargs
|
|
|
|
) -> ParallelModule:
|
2023-06-20 03:45:16 +00:00
|
|
|
r"""
|
|
|
|
Convert a huggingface layer `Conv1D` in gpt2 to a parallelized linear layer.
|
2023-06-22 06:40:37 +00:00
|
|
|
|
|
|
|
Args:
|
|
|
|
module (`nn.Linear`): The module to be converted.
|
|
|
|
process_group (`Union[ProcessGroup, List[ProcessGroup]]`): The process group to be used for weight sharding and communication.
|
|
|
|
n_fused (int): The number of layers to be fused. In GPT2, Q,K,V are fused in one weight.
|
2023-06-20 03:45:16 +00:00
|
|
|
"""
|
2023-07-10 02:48:53 +00:00
|
|
|
LazyInitContext.materialize(module)
|
2023-06-20 03:45:16 +00:00
|
|
|
# get the attributes
|
|
|
|
in_features = module.weight.shape[0]
|
|
|
|
out_features = module.weight.shape[1]
|
|
|
|
bias = module.bias is not None
|
|
|
|
device = module.weight.device
|
|
|
|
|
|
|
|
# ensure only one process group is passed
|
|
|
|
if isinstance(process_group, (list, tuple)):
|
2023-09-19 06:20:26 +00:00
|
|
|
assert len(process_group) == 1, f"Expected only one process group, got {len(process_group)}."
|
2023-06-20 03:45:16 +00:00
|
|
|
process_group = process_group[0]
|
|
|
|
|
2023-08-11 07:43:23 +00:00
|
|
|
tp_size = dist.get_world_size(process_group)
|
|
|
|
if out_features < tp_size:
|
|
|
|
return module
|
|
|
|
|
|
|
|
if out_features % tp_size != 0:
|
|
|
|
raise ValueError(
|
2023-09-19 06:20:26 +00:00
|
|
|
f"The size of out_features:{out_features} is not integer multiples of tensor parallel size: {tp_size}!"
|
|
|
|
)
|
|
|
|
|
|
|
|
linear_1d = GPT2FusedLinearConv1D_Col(
|
|
|
|
in_features=in_features,
|
|
|
|
out_features=out_features,
|
|
|
|
bias=bias,
|
|
|
|
device=device,
|
|
|
|
process_group=process_group,
|
|
|
|
weight=module.weight,
|
|
|
|
bias_=module.bias,
|
|
|
|
*args,
|
|
|
|
**kwargs,
|
|
|
|
)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
return linear_1d
|
|
|
|
|
|
|
|
def reset_parameters(self, weight_initializer, bias_initializer) -> None:
|
2023-06-21 06:30:06 +00:00
|
|
|
with self.randomizer.fork_rng(enable_cpu=True):
|
|
|
|
fan_in, fan_out = self.in_features, self.out_features
|
|
|
|
weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out)
|
|
|
|
if self.bias is not None:
|
|
|
|
bias_initializer(self.bias, fan_in=fan_in)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
def forward(self, input_: Tensor) -> Tuple[Tensor, Tensor]:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert (
|
|
|
|
input_.shape[-1] == self.weight.shape[0]
|
|
|
|
), "Invalid shapes in Linear1D_Col forward: input={}, weight={}. Expected last dim of input {}.".format(
|
|
|
|
input_.shape, self.weight.shape, self.weight.shape[-1]
|
|
|
|
)
|
2023-06-23 08:07:09 +00:00
|
|
|
|
2023-06-20 03:45:16 +00:00
|
|
|
# Matrix multiply.
|
|
|
|
bias = self.bias if not self.skip_bias_add else None
|
2023-06-23 08:07:09 +00:00
|
|
|
|
2023-08-16 07:41:20 +00:00
|
|
|
if self.seq_parallel:
|
|
|
|
input_parallel = input_
|
2023-09-19 06:20:26 +00:00
|
|
|
output_parallel = matmul_gather_forward_reducescatter_backward(
|
|
|
|
input_parallel, self.weight, bias, self.process_group, True, 1, self.overlap
|
|
|
|
)
|
2023-08-16 07:41:20 +00:00
|
|
|
else:
|
|
|
|
# Set up backprop all-reduce.
|
|
|
|
input_parallel = reduce_backward(input_, self.process_group)
|
2023-09-19 06:20:26 +00:00
|
|
|
output_parallel = matmul_with_async_comm(
|
|
|
|
input_parallel, self.weight, bias, self.process_group, self.async_communication
|
|
|
|
)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
if self.gather_output:
|
|
|
|
# All-gather across the partitions.
|
|
|
|
output = gather_forward_split_backward(output_parallel, dim=-1, process_group=self.process_group)
|
|
|
|
else:
|
|
|
|
output = output_parallel
|
|
|
|
|
|
|
|
if self.skip_bias_add:
|
|
|
|
return output, self.bias
|
|
|
|
else:
|
|
|
|
return output
|
|
|
|
|
|
|
|
|
2023-06-28 07:04:35 +00:00
|
|
|
class GPT2FusedLinearConv1D_Row(ParallelModule):
|
2023-09-19 06:20:26 +00:00
|
|
|
r"""Linear layer with row parallelism.
|
2023-06-23 08:07:09 +00:00
|
|
|
This layer is used to fit `Conv1D` layer (Fused QKV) in gpt2 of huggingface.
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
Args:
|
|
|
|
in_features (int): size of each input sample.
|
|
|
|
out_features (int): size of each output sample.
|
|
|
|
bias (bool, optional): If set to ``False``, the layer will not learn an additive bias, defaults to ``True``.
|
|
|
|
dtype (`torch.dtype`): The dtype of parameters, defaults to None.
|
|
|
|
parallel_input (bool): If set to ``True``, it's assumed that the input is split, defaults to False.
|
|
|
|
skip_bias_add (bool): If set to ``True``, it will skip bias add for linear layer,
|
2023-08-16 07:41:20 +00:00
|
|
|
seq_parallel (`bool`): If set to ``True``, it will use sequence parallel, defaults to False.
|
2023-06-20 03:45:16 +00:00
|
|
|
which is preserved for kernel fusion, defaults to False
|
|
|
|
weight_initializer (:class:`typing.Callable`, optional):
|
|
|
|
The initializer of weight, defaults to kaiming uniform initializer.
|
|
|
|
bias_initializer (:class:`typing.Callable`, optional):
|
|
|
|
The initializer of bias, defaults to xavier uniform initializer.
|
|
|
|
|
|
|
|
More details about ``initializer`` please refer to
|
|
|
|
`init <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/nn/init.py>`_.
|
|
|
|
"""
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
in_features: int,
|
|
|
|
out_features: int,
|
|
|
|
bias: bool = True,
|
|
|
|
dtype: torch.dtype = None,
|
|
|
|
device: torch.device = None,
|
|
|
|
process_group: ProcessGroup = None,
|
|
|
|
seq_parallel: bool = False,
|
|
|
|
parallel_input: bool = True,
|
|
|
|
skip_bias_add: bool = False,
|
|
|
|
weight: Optional[Parameter] = None,
|
|
|
|
bias_: Optional[Parameter] = None,
|
|
|
|
weight_initializer: Callable = init.kaiming_uniform_(a=math.sqrt(5)),
|
|
|
|
bias_initializer: Callable = init.xavier_uniform_(a=1, scale=1),
|
|
|
|
stream_chunk_num: int = 1,
|
|
|
|
):
|
2023-06-20 03:45:16 +00:00
|
|
|
super().__init__()
|
|
|
|
|
|
|
|
self.stream_chunk_num = stream_chunk_num
|
|
|
|
|
|
|
|
# Keep input parameters
|
|
|
|
self.in_features = in_features
|
|
|
|
self.out_features = out_features
|
|
|
|
self.parallel_input = parallel_input
|
|
|
|
self.skip_bias_add = skip_bias_add
|
|
|
|
self.process_group = process_group
|
2023-08-16 07:41:20 +00:00
|
|
|
self.seq_parallel = seq_parallel
|
2023-06-20 03:45:16 +00:00
|
|
|
self.num_partitions = dist.get_world_size(self.process_group)
|
|
|
|
|
|
|
|
if skip_bias_add and not bias:
|
2023-09-19 06:20:26 +00:00
|
|
|
raise ValueError("cannot skip bias addition if bias is None")
|
2023-06-20 03:45:16 +00:00
|
|
|
|
2023-07-20 02:39:06 +00:00
|
|
|
# offset the seed with randomizer index and rank
|
|
|
|
seed = torch.random.initial_seed()
|
|
|
|
self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group)
|
|
|
|
|
2023-06-20 03:45:16 +00:00
|
|
|
# Divide the weight matrix along the last dimension.
|
|
|
|
self.input_size_per_partition = divide(in_features, self.num_partitions)
|
|
|
|
|
2023-07-20 02:39:06 +00:00
|
|
|
# sanity check
|
|
|
|
if weight is not None:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert not bias or bias_ is not None, "bias_ must be provided if bias is True when weight is not None"
|
2023-07-20 02:39:06 +00:00
|
|
|
else:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert bias_ is None, "bias_ must be None if weight is None"
|
2023-07-20 02:39:06 +00:00
|
|
|
|
2023-06-20 03:45:16 +00:00
|
|
|
# Parameters.
|
2023-07-20 02:39:06 +00:00
|
|
|
if weight is None:
|
|
|
|
# Initialize weight.
|
2023-09-19 06:20:26 +00:00
|
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
2023-07-20 02:39:06 +00:00
|
|
|
self.weight = Parameter(torch.empty(self.in_features, self.out_features, **factory_kwargs))
|
|
|
|
else:
|
|
|
|
weight.data = weight.data.to(device=device, dtype=dtype)
|
|
|
|
self.weight = weight
|
|
|
|
if not is_distributed_tensor(self.weight):
|
|
|
|
sharded_weight = shard_rowwise(self.weight.data, self.process_group)
|
|
|
|
sharded_tensor_to_existing_param(sharded_weight, self.weight)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
if self.stream_chunk_num > 1:
|
|
|
|
# TODO() work for inference only
|
|
|
|
self.chunk_weight()
|
|
|
|
if bias:
|
2023-07-20 02:39:06 +00:00
|
|
|
if bias_ is None:
|
|
|
|
self.bias = Parameter(torch.empty(self.out_features, **factory_kwargs))
|
|
|
|
else:
|
|
|
|
bias_.data = bias_.data.to(device=device, dtype=dtype)
|
|
|
|
self.bias = bias_
|
2023-06-20 03:45:16 +00:00
|
|
|
else:
|
|
|
|
self.bias = None
|
|
|
|
|
2023-07-20 02:39:06 +00:00
|
|
|
if weight is None:
|
|
|
|
# init weights
|
|
|
|
self.reset_parameters(weight_initializer, bias_initializer)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
@staticmethod
|
2023-09-19 06:20:26 +00:00
|
|
|
def from_native_module(
|
|
|
|
module: nn.Linear, process_group: Union[ProcessGroup, List[ProcessGroup]], *args, **kwargs
|
|
|
|
) -> ParallelModule:
|
2023-06-20 03:45:16 +00:00
|
|
|
r"""
|
|
|
|
Convert a native PyTorch linear layer to a parallelized linear layer.
|
|
|
|
"""
|
2023-07-10 02:48:53 +00:00
|
|
|
LazyInitContext.materialize(module)
|
2023-06-20 03:45:16 +00:00
|
|
|
# get the attributes
|
|
|
|
in_features = module.weight.shape[0]
|
|
|
|
out_features = module.weight.shape[1]
|
|
|
|
bias = module.bias is not None
|
|
|
|
device = module.weight.device
|
|
|
|
|
|
|
|
# ensure only one process group is passed
|
|
|
|
if isinstance(process_group, (list, tuple)):
|
2023-09-19 06:20:26 +00:00
|
|
|
assert len(process_group) == 1, f"Expected only one process group, got {len(process_group)}."
|
2023-06-20 03:45:16 +00:00
|
|
|
process_group = process_group[0]
|
|
|
|
|
2023-08-11 07:43:23 +00:00
|
|
|
tp_size = dist.get_world_size(process_group)
|
|
|
|
if in_features < tp_size:
|
|
|
|
return module
|
|
|
|
|
|
|
|
if in_features % tp_size != 0:
|
|
|
|
raise ValueError(
|
2023-09-19 06:20:26 +00:00
|
|
|
f"The size of in_features:{in_features} is not integer multiples of tensor parallel size: {tp_size}!"
|
|
|
|
)
|
|
|
|
|
|
|
|
linear_1d = GPT2FusedLinearConv1D_Row(
|
|
|
|
in_features=in_features,
|
|
|
|
out_features=out_features,
|
|
|
|
bias=bias,
|
|
|
|
device=device,
|
|
|
|
process_group=process_group,
|
|
|
|
weight=module.weight,
|
|
|
|
bias_=module.bias,
|
|
|
|
*args,
|
|
|
|
**kwargs,
|
|
|
|
)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
return linear_1d
|
|
|
|
|
|
|
|
def chunk_weight(self):
|
|
|
|
self.weight_list = torch.chunk(self.weight, self.stream_chunk_num, dim=0)
|
|
|
|
|
|
|
|
def reset_parameters(self, weight_initializer, bias_initializer) -> None:
|
2023-06-21 06:30:06 +00:00
|
|
|
with self.randomizer.fork_rng(enable_cpu=True):
|
|
|
|
fan_in, fan_out = self.in_features, self.out_features
|
|
|
|
weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out)
|
|
|
|
|
|
|
|
if self.bias is not None:
|
|
|
|
bias_initializer(self.bias, fan_in=fan_in)
|
|
|
|
if self.process_group is None:
|
|
|
|
src_rank = 0
|
|
|
|
else:
|
|
|
|
src_rank = dist.distributed_c10d._get_global_rank(self.process_group, 0)
|
|
|
|
|
|
|
|
origin_device = self.bias.device
|
2023-07-10 02:48:53 +00:00
|
|
|
self.bias.data = self.bias.cuda()
|
2023-06-21 06:30:06 +00:00
|
|
|
dist.broadcast(self.bias, src=src_rank, group=self.process_group)
|
2023-07-10 02:48:53 +00:00
|
|
|
self.bias.data = self.bias.to(origin_device)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
def forward(self, input_: Tensor) -> Tensor:
|
|
|
|
# Set up backprop all-reduce.
|
|
|
|
if self.parallel_input:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert (
|
|
|
|
input_.shape[-1] == self.weight.shape[0]
|
|
|
|
), "Invalid shapes in Linear1D_Row forward: input={}, weight={}. Expected last dim of input {}.".format(
|
|
|
|
input_.shape, self.weight.shape, self.weight.shape[0]
|
|
|
|
)
|
2023-06-20 03:45:16 +00:00
|
|
|
input_ = input_
|
|
|
|
else:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert (
|
|
|
|
divide(input_.shape[-1], self.num_partitions) == self.weight.shape[0]
|
|
|
|
), "Invalid shapes in Linear1D_Row forward: input={}, weight={}. Expected last dim of input {}.".format(
|
|
|
|
input_.shape, self.weight.shape, self.weight.shape[0] * self.num_partitions
|
|
|
|
)
|
2023-06-20 03:45:16 +00:00
|
|
|
input_ = split_forward_gather_backward(input_, dim=-1, process_group=self.process_group)
|
|
|
|
|
|
|
|
if self.stream_chunk_num > 1:
|
|
|
|
if self.training:
|
|
|
|
raise RuntimeError("use stream_chunk_num=1 in Linear1D_Row for training!")
|
|
|
|
with torch.no_grad():
|
|
|
|
output_parallel_list = [None for i in range(self.stream_chunk_num)]
|
|
|
|
handle_list = []
|
|
|
|
for i in range(self.stream_chunk_num):
|
2023-06-23 08:07:09 +00:00
|
|
|
output_parallel_list[i] = torch.matmul(input_, self.weight_list[i])
|
2023-09-19 06:20:26 +00:00
|
|
|
handle = torch.distributed.all_reduce(
|
|
|
|
output_parallel_list[i], group=self.process_group, async_op=True
|
|
|
|
)
|
2023-06-20 03:45:16 +00:00
|
|
|
handle_list.append(handle)
|
|
|
|
# output_parallel_list[i] = reduce_input(output_parallel_list[i], ParallelMode.PARALLEL_1D)
|
|
|
|
for handle in handle_list:
|
|
|
|
handle.wait()
|
|
|
|
output = torch.cat(output_parallel_list, dim=-1)
|
|
|
|
else:
|
2023-06-23 08:07:09 +00:00
|
|
|
output_parallel = torch.matmul(input_, self.weight)
|
2023-08-16 07:41:20 +00:00
|
|
|
if self.seq_parallel:
|
|
|
|
output = linear_reducescatter_forward_gather_backward(output_parallel, self.process_group, 1)
|
|
|
|
else:
|
|
|
|
output = reduce_forward(output_parallel, self.process_group)
|
2023-06-20 03:45:16 +00:00
|
|
|
|
|
|
|
if not self.skip_bias_add:
|
|
|
|
if self.bias is not None:
|
|
|
|
output = output + self.bias
|
|
|
|
return output
|
|
|
|
else:
|
|
|
|
return output, self.bias
|
2023-07-14 07:56:59 +00:00
|
|
|
|
|
|
|
|
|
|
|
# ====================================
|
|
|
|
# For Fused torch.nn.Linear
|
|
|
|
# ====================================
|
|
|
|
|
|
|
|
|
|
|
|
class FusedLinear1D_Col(ParallelModule):
|
|
|
|
r"""Fused Linear layer with column parallelism.
|
|
|
|
|
|
|
|
The linear layer is defined as :math:`Y = XA + b`. A is parallelized along
|
|
|
|
its second dimension as :math:`A = [A_1, ..., A_p]`. This layer is used to fit `torch.nn.Linear` layer (Fused QKV) in normal torch layer of huggingface, like SAM.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
in_features (int): size of each input sample.
|
|
|
|
out_features (int): size of each output sample.
|
|
|
|
bias (bool, optional): If set to ``False``, the layer will not learn an additive bias, defaults to ``True``.
|
|
|
|
dtype (`torch.dtype`): The dtype of parameters, defaults to None.
|
|
|
|
device (`torch.device`): The device of parameters, defaults to None.
|
|
|
|
n_fused (int): The number items fused, defaults to 3 (QKV).
|
|
|
|
process_group (`torch.distributed.ProcessGroup`): The process group to be used for weight sharding and communication, defaults to None.
|
|
|
|
gather_output (bool, optional): If true, call all-gather on output and make Y available
|
|
|
|
to all GPUs, otherwise, every GPU will have its output
|
|
|
|
which is :math:`Y_i = XA_i`, defaults to False
|
|
|
|
skip_bias_add (bool): If set to ``True``, it will skip bias add for linear layer,
|
|
|
|
which is preserved for kernel fusion, defaults to False
|
|
|
|
weight_initializer (`typing.Callable`):
|
|
|
|
The initializer of weight, defaults to kaiming uniform initializer.
|
|
|
|
bias_initializer (`typing.Callable`):
|
|
|
|
The initializer of bias, defaults to xavier uniform initializer.
|
|
|
|
|
|
|
|
More details about ``initializer`` please refer to
|
|
|
|
`init <https://github.com/hpcaitech/ColossalAI/blob/main/colossalai/nn/init.py>`_.
|
|
|
|
"""
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
in_features: int,
|
|
|
|
out_features: int,
|
|
|
|
bias: bool = True,
|
|
|
|
dtype: torch.dtype = None,
|
|
|
|
device: torch.device = None,
|
|
|
|
process_group: ProcessGroup = None,
|
|
|
|
async_communication: bool = False,
|
|
|
|
gather_output: bool = False,
|
|
|
|
skip_bias_add: bool = False,
|
|
|
|
n_fused: int = 3,
|
|
|
|
weight: Optional[Parameter] = None,
|
|
|
|
bias_: Optional[Parameter] = None,
|
|
|
|
weight_initializer: Callable = init.kaiming_uniform_(a=math.sqrt(5)),
|
|
|
|
bias_initializer: Callable = init.xavier_uniform_(a=1, scale=1),
|
|
|
|
):
|
2023-07-14 07:56:59 +00:00
|
|
|
super().__init__()
|
|
|
|
# Keep input parameters
|
|
|
|
self.in_features = in_features
|
|
|
|
self.out_features = out_features
|
|
|
|
self.gather_output = gather_output
|
|
|
|
self.skip_bias_add = skip_bias_add
|
|
|
|
self.device = device
|
|
|
|
self.n_fused = n_fused
|
|
|
|
self.process_group = process_group
|
|
|
|
self.async_communication = async_communication
|
|
|
|
|
|
|
|
if skip_bias_add and not bias:
|
2023-09-19 06:20:26 +00:00
|
|
|
raise ValueError("cannot skip bias addition if bias is None")
|
2023-07-14 07:56:59 +00:00
|
|
|
|
2023-08-01 10:02:49 +00:00
|
|
|
# offset the seed with randomizer index and rank
|
|
|
|
seed = torch.random.initial_seed()
|
|
|
|
self.randomizer = create_randomizer_with_offset(seed, process_group=self.process_group)
|
|
|
|
|
|
|
|
# sanity check
|
|
|
|
if weight is not None:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert not bias or bias_ is not None, "bias_ must be provided if bias is True when weight is not None"
|
2023-08-01 10:02:49 +00:00
|
|
|
else:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert bias_ is None, "bias_ must be None if weight is None"
|
2023-08-01 10:02:49 +00:00
|
|
|
|
2023-07-14 07:56:59 +00:00
|
|
|
# Parameters.
|
2023-08-01 10:02:49 +00:00
|
|
|
if weight is None:
|
|
|
|
# Initialize weight.
|
2023-09-19 06:20:26 +00:00
|
|
|
factory_kwargs = {"device": device, "dtype": dtype}
|
2023-08-01 10:02:49 +00:00
|
|
|
self.weight = Parameter(torch.empty(self.out_features, self.in_features, **factory_kwargs))
|
|
|
|
else:
|
|
|
|
weight.data = weight.data.to(device=device, dtype=dtype)
|
|
|
|
self.weight = weight
|
2023-07-14 07:56:59 +00:00
|
|
|
|
|
|
|
def shard_fn(tensor):
|
|
|
|
return split_fused_qkv_in_gpt2_style(tensor, self.n_fused, self.process_group, False)
|
|
|
|
|
|
|
|
def gather_fn(tensor):
|
2023-08-01 10:02:49 +00:00
|
|
|
return gather_fused_qkv_in_gpt2_style(tensor, self.n_fused, self.process_group, False)
|
2023-07-14 07:56:59 +00:00
|
|
|
|
2023-08-01 10:02:49 +00:00
|
|
|
if not is_customized_distributed_tensor(self.weight):
|
|
|
|
with torch.no_grad():
|
|
|
|
sharded_weight = distribute_tensor_with_customization(self.weight.data, shard_fn, gather_fn)
|
|
|
|
customized_distributed_tensor_to_existing_param(sharded_weight, self.weight)
|
2023-07-14 07:56:59 +00:00
|
|
|
|
|
|
|
if bias:
|
2023-08-01 10:02:49 +00:00
|
|
|
if bias_ is None:
|
|
|
|
self.bias = Parameter(torch.empty(self.out_features, **factory_kwargs))
|
|
|
|
else:
|
|
|
|
bias_.data = bias_.data.to(device=device, dtype=dtype)
|
|
|
|
self.bias = bias_
|
|
|
|
if not is_customized_distributed_tensor(self.bias):
|
|
|
|
with torch.no_grad():
|
|
|
|
sharded_bias = distribute_tensor_with_customization(self.bias.data, shard_fn, gather_fn)
|
|
|
|
customized_distributed_tensor_to_existing_param(sharded_bias, self.bias)
|
2023-07-14 07:56:59 +00:00
|
|
|
else:
|
|
|
|
self.bias = None
|
|
|
|
|
2023-08-01 10:02:49 +00:00
|
|
|
if weight is None:
|
|
|
|
# init weights
|
|
|
|
self.reset_parameters(weight_initializer, bias_initializer)
|
2023-07-14 07:56:59 +00:00
|
|
|
|
|
|
|
@staticmethod
|
2023-09-19 06:20:26 +00:00
|
|
|
def from_native_module(
|
|
|
|
module: nn.Module, process_group: Union[ProcessGroup, List[ProcessGroup]], n_fused: int, *args, **kwargs
|
|
|
|
) -> ParallelModule:
|
2023-07-14 07:56:59 +00:00
|
|
|
r"""
|
|
|
|
Convert a fused `torch.nn.linear` layer to a parallelized linear layer.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
module (`nn.Linear`): The module to be converted.
|
|
|
|
process_group (`Union[ProcessGroup, List[ProcessGroup]]`): The process group to be used for weight sharding and communication.
|
|
|
|
n_fused (int): The number of layers to be fused. In common, Q,K,V are fused in one weight.
|
|
|
|
"""
|
|
|
|
# get the attributes
|
|
|
|
in_features = module.in_features
|
|
|
|
out_features = module.out_features
|
|
|
|
bias = module.bias is not None
|
|
|
|
device = module.weight.device
|
|
|
|
|
|
|
|
# ensure only one process group is passed
|
|
|
|
if isinstance(process_group, (list, tuple)):
|
2023-09-19 06:20:26 +00:00
|
|
|
assert len(process_group) == 1, f"Expected only one process group, got {len(process_group)}."
|
2023-07-14 07:56:59 +00:00
|
|
|
process_group = process_group[0]
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
linear_1d = FusedLinear1D_Col(
|
|
|
|
in_features=in_features,
|
|
|
|
out_features=out_features,
|
|
|
|
bias=bias,
|
|
|
|
device=device,
|
|
|
|
process_group=process_group,
|
|
|
|
weight=module.weight,
|
|
|
|
bias_=module.bias,
|
|
|
|
*args,
|
|
|
|
**kwargs,
|
|
|
|
)
|
2023-07-14 07:56:59 +00:00
|
|
|
|
2023-08-01 10:02:49 +00:00
|
|
|
# # TODO: copy the sharded weights
|
|
|
|
# with torch.no_grad():
|
|
|
|
# sharded_weight = split_fused_qkv_in_gpt2_style(module.weight.data,
|
|
|
|
# n_fused=n_fused,
|
|
|
|
# process_group=process_group,
|
|
|
|
# is_transposed=False)
|
|
|
|
# linear_1d.weight.data.copy_(sharded_weight.data)
|
|
|
|
|
|
|
|
# if bias:
|
|
|
|
# sharded_bias = split_fused_qkv_in_gpt2_style(module.bias.data,
|
|
|
|
# n_fused=n_fused,
|
|
|
|
# process_group=process_group,
|
|
|
|
# is_transposed=False)
|
|
|
|
# linear_1d.bias.data.copy_(sharded_bias.data)
|
|
|
|
print(linear_1d.weight.shape)
|
2023-07-14 07:56:59 +00:00
|
|
|
return linear_1d
|
|
|
|
|
|
|
|
def reset_parameters(self, weight_initializer, bias_initializer) -> None:
|
|
|
|
with self.randomizer.fork_rng(enable_cpu=True):
|
|
|
|
fan_in, fan_out = self.in_features, self.out_features
|
|
|
|
weight_initializer(self.weight, fan_in=fan_in, fan_out=fan_out)
|
|
|
|
if self.bias is not None:
|
|
|
|
bias_initializer(self.bias, fan_in=fan_in)
|
|
|
|
|
|
|
|
def forward(self, input_: Tensor) -> Tuple[Tensor, Tensor]:
|
2023-09-19 06:20:26 +00:00
|
|
|
assert (
|
|
|
|
input_.shape[-1] == self.weight.shape[-1]
|
|
|
|
), "Invalid shapes in Linear1D_Col forward: input={}, weight={}. Expected last dim of input {}.".format(
|
|
|
|
input_.shape, self.weight.shape, self.weight.shape[-1]
|
|
|
|
)
|
2023-07-14 07:56:59 +00:00
|
|
|
# Set up backprop all-reduce.
|
|
|
|
# input_parallel = reduce_backward(input_, self.process_group)
|
|
|
|
input_parallel = input_
|
|
|
|
|
|
|
|
# Matrix multiply.
|
|
|
|
bias = self.bias if not self.skip_bias_add else None
|
|
|
|
|
|
|
|
output_parallel = linear_with_async_comm(input_parallel, self.weight, bias, self.process_group, True)
|
|
|
|
|
|
|
|
if self.gather_output:
|
|
|
|
# All-gather across the partitions.
|
|
|
|
output = gather_forward_split_backward(output_parallel, dim=-1, process_group=self.process_group)
|
|
|
|
else:
|
|
|
|
output = output_parallel
|
|
|
|
|
|
|
|
if self.skip_bias_add:
|
|
|
|
return output, self.bias
|
|
|
|
else:
|
|
|
|
return output
|