wehub-resource-sync
1454 lines
61 KiB
Python
1454 lines
61 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import torch
|
|
import re
|
|
from deepspeed import comm as dist
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
from torch.nn.parameter import Parameter
|
|
from deepspeed.accelerator import get_accelerator
|
|
from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
|
|
from deepspeed.runtime.zero.utils import is_zero_param
|
|
from abc import ABC, abstractmethod
|
|
from typing import Iterable, Any, Optional, List, Tuple, Dict
|
|
from .fusedqkv_utils import shard_value_with_share_qk, shard_chunk_mlp, prepare_tp_fused_qkvw
|
|
from deepspeed.runtime.tensor_parallel import AUTOTP_MODE
|
|
from deepspeed.checkpoint.constants import DS_AUTOTP_UC_META
|
|
from copy import deepcopy
|
|
from typing import Union
|
|
|
|
__all__ = [
|
|
"TensorParallel_Layer", "LinearAllreduce", "LinearLayer", "LmHeadLinearAllreduce", "Yuan_LinearAllreduce",
|
|
"Yuan_LinearLayer", "GateUpPack_LinearLayer", "Conv_LinearALlreduce", "fused_LinearLayer", "conv_LinearLayer",
|
|
"SubParamLinearLayer", "SubParamLinearAllreduce"
|
|
]
|
|
|
|
DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE
|
|
DS_IS_REPLACED_MODULE = 'ds_is_replaced_module'
|
|
DS_TENSOR_MODEL_PARALLEL = 'tensor_model_parallel'
|
|
|
|
|
|
def _normalize_uc_shape(value):
|
|
return tuple(value) if value is not None else None
|
|
|
|
|
|
def _build_param_uc_conversion_meta(*,
|
|
partition_type,
|
|
partition_dim=None,
|
|
sub_param_shape=None,
|
|
original_shape=None,
|
|
is_bias=False,
|
|
replicated=False):
|
|
"""Build the conversion-facing subset of parameter UC metadata.
|
|
|
|
This is the only schema that should flow into model-level
|
|
`UNIVERSAL_CHECKPOINT_INFO` via `collect_autotp_universal_checkpoint_info()`.
|
|
"""
|
|
return {
|
|
'partition_type': partition_type,
|
|
'partition_dim': partition_dim,
|
|
'sub_param_shape': _normalize_uc_shape(sub_param_shape),
|
|
'original_shape': _normalize_uc_shape(original_shape),
|
|
'is_bias': is_bias,
|
|
'replicated': replicated,
|
|
}
|
|
|
|
|
|
def _build_param_uc_restore_meta(*,
|
|
partition_type,
|
|
partition_dim=None,
|
|
logical_shape=None,
|
|
output_shape=None,
|
|
sub_param_shape=None,
|
|
sub_param_sizes=None,
|
|
target_partition_shape=None,
|
|
original_shape=None,
|
|
is_bias=False,
|
|
replicated=False):
|
|
"""Build the restore-facing parameter UC metadata.
|
|
|
|
Restore metadata stays on the parameter object and may include details that
|
|
are intentionally omitted from model-level conversion schema.
|
|
"""
|
|
return {
|
|
'partition_type':
|
|
partition_type,
|
|
'partition_dim':
|
|
partition_dim,
|
|
'logical_shape':
|
|
_normalize_uc_shape(logical_shape),
|
|
'output_shape':
|
|
_normalize_uc_shape(output_shape),
|
|
'sub_param_shape':
|
|
_normalize_uc_shape(sub_param_shape),
|
|
'sub_param_sizes':
|
|
_normalize_uc_shape(sub_param_sizes),
|
|
'target_partition_shape':
|
|
_normalize_uc_shape(target_partition_shape),
|
|
'original_shape':
|
|
_normalize_uc_shape(original_shape),
|
|
'is_bias':
|
|
is_bias,
|
|
'replicated':
|
|
replicated,
|
|
'conversion':
|
|
_build_param_uc_conversion_meta(partition_type=partition_type,
|
|
partition_dim=partition_dim,
|
|
sub_param_shape=sub_param_shape,
|
|
original_shape=original_shape,
|
|
is_bias=is_bias,
|
|
replicated=replicated),
|
|
}
|
|
|
|
|
|
def get_auto_tp_mode():
|
|
global DEEPSPEED_AUTOTP_MODE
|
|
return DEEPSPEED_AUTOTP_MODE
|
|
|
|
|
|
def is_autotp_training_mode():
|
|
global DEEPSPEED_AUTOTP_MODE
|
|
return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING
|
|
|
|
|
|
def set_autotp_mode(training=False):
|
|
"""
|
|
Set the DEEPSPEED_AUTOTP_MODE based on the training flag
|
|
"""
|
|
global DEEPSPEED_AUTOTP_MODE
|
|
if training:
|
|
DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.TRAINING
|
|
else:
|
|
DEEPSPEED_AUTOTP_MODE = AUTOTP_MODE.INFERENCE
|
|
|
|
|
|
def add_bias(input, bias):
|
|
if bias is None:
|
|
return input
|
|
if is_autotp_training_mode():
|
|
# Training mode - avoid inplace to ensure correct autograd
|
|
input = input + bias
|
|
return input
|
|
else:
|
|
input += bias
|
|
return input
|
|
|
|
|
|
class RowParallel(torch.autograd.Function):
|
|
"""
|
|
A custom autograd function for performing row-wise parallelism.
|
|
"""
|
|
|
|
@staticmethod
|
|
def symbolic(graph, input):
|
|
"""Symbolic function for tracing."""
|
|
return input
|
|
|
|
@staticmethod
|
|
def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, is_inference_mode: bool) -> torch.Tensor:
|
|
"""
|
|
Forward pass.
|
|
"""
|
|
ctx.group = group
|
|
if group == None:
|
|
return input
|
|
if is_inference_mode:
|
|
dist.inference_all_reduce(input, group=group)
|
|
else:
|
|
dist.all_reduce(input.contiguous(), group=group)
|
|
return input
|
|
|
|
@staticmethod
|
|
def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor, None]:
|
|
"""
|
|
Backward pass.
|
|
"""
|
|
return None, grad_output, None
|
|
|
|
|
|
class AsyncColumnParallel(torch.autograd.Function):
|
|
|
|
@staticmethod
|
|
def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor, weight, bias) -> torch.Tensor:
|
|
"""
|
|
Forward pass.
|
|
"""
|
|
ctx.use_bias = bias is not None
|
|
ctx.group = group
|
|
output = torch.matmul(input, weight.transpose(-1, -2))
|
|
if bias is not None:
|
|
output = add_bias(output, bias)
|
|
|
|
ctx.save_for_backward(input, weight)
|
|
|
|
return output
|
|
|
|
@staticmethod
|
|
def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:
|
|
|
|
input, weight = ctx.saved_tensors
|
|
grad_input = grad_output.matmul(weight)
|
|
handle = dist.all_reduce(grad_input.contiguous(), group=ctx.group, async_op=True)
|
|
grad_weight = grad_output.view(-1, grad_output.shape[-1]).t().matmul(input.view(-1, input.shape[-1]))
|
|
grad_bias = grad_output.sum(0) if ctx.use_bias else None
|
|
handle.wait()
|
|
return None, grad_input, grad_weight, grad_bias
|
|
|
|
|
|
class ColumnParallel(torch.autograd.Function):
|
|
"""
|
|
Custom autograd function for column-wise parallelism.
|
|
"""
|
|
|
|
@staticmethod
|
|
def symbolic(graph, input):
|
|
"""Symbolic function for tracing."""
|
|
return dist.all_reduce(input.contiguous(), dist.get_tensor_model_parallel_group())
|
|
|
|
@staticmethod
|
|
def forward(ctx: Any, group: dist.ProcessGroup, input: torch.Tensor) -> torch.Tensor:
|
|
"""
|
|
Forward pass.
|
|
"""
|
|
ctx.group = group
|
|
return input
|
|
|
|
@staticmethod
|
|
def backward(ctx: Any, grad_output: torch.Tensor) -> Tuple[None, torch.Tensor]:
|
|
"""
|
|
Backward pass.
|
|
"""
|
|
if ctx.group == None:
|
|
return None, grad_output
|
|
|
|
dist.all_reduce(grad_output.contiguous(), group=ctx.group)
|
|
return None, grad_output
|
|
|
|
|
|
class TensorParallel_Layer(nn.Module, ABC):
|
|
"""
|
|
A base class for model layers with tensor parallelism support.
|
|
This class is designed to be extended by specific layers that require distributed
|
|
operations and parameter gather/partitioning during inference or training.
|
|
|
|
Attributes:
|
|
mode (str): The mode of operation[INFERENCE or TRAINING], default is "INFERENCE".
|
|
mp_group (Optional[dist.ProcessGroup]): The process group used for model parallelism.
|
|
tp_world_size (int): The world size of tensor parallelism, i.e., the number of parallel workers.
|
|
tp_index (int): The rank (ID) of the current worker in tensor parallelism.
|
|
support_training (bool): Flag indicating whether the layer supports training (default: False).
|
|
name (Optional[str]): The name of the layer, if provided.
|
|
"""
|
|
##### Initialize Parameter List #####
|
|
|
|
# keep_module_on_host determines whether to keep the module on the host.
|
|
# Checkpoints are first loaded to the host (sometimes directly from disk to avoid filling host memory),
|
|
# so an additional copy is unnecessary.
|
|
keep_module_on_host: bool = False
|
|
|
|
##### Runtime Parameter List #####
|
|
tp_overlap_comm: bool = False
|
|
""" Whether to overlap communication with computation. Currently, only allreduce supports overlap. """
|
|
|
|
def __init__(self, mp_group: Optional[dist.ProcessGroup], **kwargs: Any):
|
|
"""
|
|
Initializes the TensorParallel_Layer with optional model parallelism group and layer name.
|
|
|
|
Args:
|
|
mp_group (Optional[dist.ProcessGroup]): The process group for model parallelism.
|
|
If None, no model parallelism is set.
|
|
"""
|
|
super().__init__()
|
|
self.support_training: bool = False
|
|
self.mp_group = mp_group
|
|
if mp_group is not None:
|
|
self.tp_world_size: int = dist.get_world_size(self.mp_group)
|
|
self.tp_index: int = dist.get_rank(self.mp_group)
|
|
else:
|
|
self.tp_world_size: int = 1
|
|
self.tp_index: int = 0
|
|
|
|
# backward compatibility
|
|
self.world_size = self.tp_world_size
|
|
self.rank = self.tp_index
|
|
|
|
self.name = getattr(self, 'name', None)
|
|
if kwargs.get('name') is not None:
|
|
self.name = kwargs.get('name') # Set the layer name if provided.
|
|
|
|
@classmethod
|
|
def set_keep_module_on_host(cls, value: bool):
|
|
"""
|
|
Set the static variable keep_module_on_host.
|
|
|
|
Args:
|
|
value (bool): The new value for keep_module_on_host.
|
|
"""
|
|
cls.keep_module_on_host = value
|
|
|
|
@abstractmethod
|
|
def forward(self, input):
|
|
"""
|
|
Forward pass method. Must be implemented by subclasses to define layer-specific operations.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def gather_params(self, params_list):
|
|
"""
|
|
Gathers parameters across devices for distributed training. Must be implemented by subclasses in "TRAINING" mode.
|
|
"""
|
|
pass
|
|
|
|
@abstractmethod
|
|
def _tp_partition(self, params_list: List[torch.Tensor]):
|
|
"""
|
|
Partitions the parameters for tensor parallelism.
|
|
It is necessary to ensure that this function only involves the logic of params partitioning.
|
|
"""
|
|
pass
|
|
|
|
def config_requires_grad(self, weight):
|
|
if weight is not None:
|
|
if self.is_training_mode():
|
|
if weight.requires_grad is None:
|
|
weight.requires_grad = True
|
|
else:
|
|
weight.requires_grad = False
|
|
|
|
def config_tp_params(self, weight):
|
|
"""
|
|
Configures the weight tensor for training with tensor parallelism. This includes enabling gradients
|
|
and associating necessary methods for parameter gathering and partitioning.
|
|
|
|
Args:
|
|
weight (Optional[torch.Tensor]): The weight tensor to configure for tensor parallelism.
|
|
If None, no action is taken.
|
|
"""
|
|
# # The RNG states have already been synchronized in init_inference.
|
|
if self.is_training_mode():
|
|
assert self.support_training, "No implementation of backward."
|
|
if weight is not None:
|
|
self.config_requires_grad(weight)
|
|
weight.gather_params = self.gather_params
|
|
weight._tp_partition = self._tp_partition
|
|
setattr(weight, DS_TENSOR_MODEL_PARALLEL, True)
|
|
setattr(weight, DS_IS_REPLACED_MODULE, True)
|
|
|
|
def _set_param_uc_meta(self,
|
|
param,
|
|
*,
|
|
partition_type,
|
|
partition_dim=None,
|
|
logical_shape=None,
|
|
output_shape=None,
|
|
sub_param_shape=None,
|
|
sub_param_sizes=None,
|
|
target_partition_shape=None,
|
|
original_shape=None,
|
|
is_bias=False,
|
|
replicated=False):
|
|
if param is None:
|
|
return
|
|
setattr(
|
|
param, DS_AUTOTP_UC_META,
|
|
_build_param_uc_restore_meta(partition_type=partition_type,
|
|
partition_dim=partition_dim,
|
|
logical_shape=logical_shape,
|
|
output_shape=output_shape,
|
|
sub_param_shape=sub_param_shape,
|
|
sub_param_sizes=sub_param_sizes,
|
|
target_partition_shape=target_partition_shape,
|
|
original_shape=original_shape,
|
|
is_bias=is_bias,
|
|
replicated=replicated))
|
|
|
|
def _mark_uc_metadata(self):
|
|
return
|
|
|
|
def _should_materialize_tp_partition(self):
|
|
# AutoTP partitioning should only materialize parameters when an actual
|
|
# TP process group is present. Metadata-only construction with
|
|
# mp_group=None should not touch device placement.
|
|
return self.mp_group is not None
|
|
|
|
def is_training_mode(self):
|
|
global DEEPSPEED_AUTOTP_MODE
|
|
return DEEPSPEED_AUTOTP_MODE == AUTOTP_MODE.TRAINING
|
|
|
|
def __deepcopy__(self, memo):
|
|
# This function is designed for
|
|
# 'mp_group' (a 'ProcessGroup') cannot be pickled during deepcopy in some usage.
|
|
cls = self.__class__
|
|
new_obj = cls.__new__(cls)
|
|
|
|
for key, value in vars(self).items():
|
|
if key == 'mp_group':
|
|
new_obj.mp_group = self.mp_group
|
|
else:
|
|
setattr(new_obj, key, deepcopy(value, memo))
|
|
|
|
memo[id(self)] = new_obj
|
|
return new_obj
|
|
|
|
def extra_repr(self):
|
|
out_features, in_features = None, None
|
|
if self.weight is not None:
|
|
out_features, in_features = self.weight.ds_shape[-2:] if is_zero_param(
|
|
self.weight) else self.weight.shape[-2:]
|
|
dtype = self.weight.dtype if self.weight is not None else None
|
|
return "in_features={}, out_features={}, bias={}, dtype={}".format(in_features, out_features, self.bias
|
|
is not None, dtype)
|
|
|
|
def move(self, tensor):
|
|
# TODO: consider the timing of deletion
|
|
# to save host resources when DP > 1。
|
|
|
|
# keep_module_on_host is used to keep the module on the host. Checkpoints are loaded to the host first (in some
|
|
# cases it can be done from the disk even to prevent filling host's memory), thus no need to create a new copy.
|
|
if tensor.is_meta:
|
|
# Keep tensor in meta device if tensor is meta.
|
|
return tensor
|
|
else:
|
|
device = 'cpu' if self.__class__.keep_module_on_host else get_accelerator().current_device_name()
|
|
return_new_copy = not self.__class__.keep_module_on_host
|
|
|
|
# Using new tensors help in freeing memory (after split for example) was done before by calling clone().
|
|
# Using copy=True instead of clone() will help in case of cpu --> cpu.
|
|
# Otherwise to() will not create a new copy for the view of the full tensor, and it will not be de-referenced.
|
|
cloned_tensor = tensor.to(device, copy=return_new_copy)
|
|
|
|
if return_new_copy:
|
|
# free the memory of the original tensor to reduce memory peak
|
|
# Equivalent to directly deleting the tensor reference outside the function.
|
|
# see https://github.com/microsoft/DeepSpeed/pull/4353
|
|
tensor.data = torch.empty(0, device=tensor.device)
|
|
return cloned_tensor
|
|
|
|
|
|
def configure_tensor_parallel_runtime(config):
|
|
runtime_keys = ['tp_overlap_comm']
|
|
for key in runtime_keys:
|
|
if hasattr(config, key):
|
|
setattr(TensorParallel_Layer, key, getattr(config, key))
|
|
|
|
|
|
def _get_param_uc_conversion_meta(param: torch.Tensor) -> Optional[Dict[str, Any]]:
|
|
"""Return the conversion-facing view of AutoTP UC metadata for a parameter.
|
|
|
|
AutoTP keeps a single parameter-level metadata object with two roles:
|
|
- top-level fields: restore-time details consumed by `universal_checkpoint.py`
|
|
- `conversion`: conversion-time details consumed by
|
|
`collect_autotp_universal_checkpoint_info()` and then aggregated into
|
|
model-level `UNIVERSAL_CHECKPOINT_INFO` for `ds_to_universal.py`
|
|
"""
|
|
meta = getattr(param, DS_AUTOTP_UC_META, None)
|
|
if not meta:
|
|
return None
|
|
return meta.get('conversion', None)
|
|
|
|
|
|
def collect_autotp_universal_checkpoint_info(model: nn.Module) -> Dict[str, Any]:
|
|
"""Collect the model-level conversion schema for AutoTP universal checkpoints.
|
|
|
|
The returned `UNIVERSAL_CHECKPOINT_INFO` is intentionally limited to the
|
|
pattern/schema data needed during checkpoint conversion. It does not include
|
|
restore-time per-parameter details such as `sub_param_sizes` or
|
|
`target_partition_shape`, which stay on the parameter metadata object.
|
|
"""
|
|
from deepspeed.checkpoint.constants import (ORIGINAL_VOCAB_SIZE, PARAMETER_WITH_ROW_PARALLELISM_PATTERNS,
|
|
PARAMETER_WITH_SUB_PARAMS, TP_REPLICATED_PARAMETER_PATTERNS,
|
|
UNIVERSAL_CHECKPOINT_VERSION_KEY, UNIVERSAL_CHECKPOINT_VERSION_VALUE,
|
|
VOCABULARY_PARAMETER_PATTERNS)
|
|
|
|
row_parallel_patterns = []
|
|
replicated_patterns = []
|
|
vocabulary_patterns = []
|
|
parameter_with_sub_params = []
|
|
original_vocab_size = None
|
|
|
|
for module_name, module in model.named_modules():
|
|
marker = getattr(module, "_mark_uc_metadata", None)
|
|
if marker is not None:
|
|
marker()
|
|
|
|
for param_name, param in module.named_parameters(recurse=False):
|
|
conversion_meta = _get_param_uc_conversion_meta(param)
|
|
if not conversion_meta:
|
|
continue
|
|
|
|
full_name = f"{module_name}.{param_name}" if module_name else param_name
|
|
pattern = rf"^{re.escape(full_name)}$"
|
|
|
|
if conversion_meta.get('replicated'):
|
|
replicated_patterns.append(pattern)
|
|
|
|
if conversion_meta.get('partition_type') == 'row' and not conversion_meta.get('is_bias', False):
|
|
row_parallel_patterns.append(pattern)
|
|
|
|
original_shape = conversion_meta.get('original_shape')
|
|
if original_shape and len(original_shape) == 2 and ('embed' in full_name or 'lm_head' in full_name):
|
|
vocabulary_patterns.append(pattern)
|
|
if original_vocab_size is None:
|
|
original_vocab_size = original_shape[0]
|
|
|
|
sub_param_shape = conversion_meta.get('sub_param_shape')
|
|
partition_dim = conversion_meta.get('partition_dim')
|
|
if sub_param_shape is not None and partition_dim is not None and not conversion_meta.get('is_bias', False):
|
|
parameter_with_sub_params.append({
|
|
'patterns': [pattern],
|
|
'shape': list(sub_param_shape),
|
|
'partition_dim': partition_dim,
|
|
})
|
|
|
|
uc_info = {
|
|
UNIVERSAL_CHECKPOINT_VERSION_KEY: UNIVERSAL_CHECKPOINT_VERSION_VALUE,
|
|
PARAMETER_WITH_ROW_PARALLELISM_PATTERNS: sorted(set(row_parallel_patterns)),
|
|
TP_REPLICATED_PARAMETER_PATTERNS: sorted(set(replicated_patterns)),
|
|
VOCABULARY_PARAMETER_PATTERNS: sorted(set(vocabulary_patterns)),
|
|
PARAMETER_WITH_SUB_PARAMS: parameter_with_sub_params,
|
|
}
|
|
if original_vocab_size is not None:
|
|
uc_info[ORIGINAL_VOCAB_SIZE] = original_vocab_size
|
|
return uc_info
|
|
|
|
|
|
class GatherReplacedLayerParams:
|
|
"""
|
|
A context manager for gathering parameters of a replaced layer, enabling partitioning and gathering functionality
|
|
based on the configuration of the model.
|
|
"""
|
|
|
|
def __init__(self,
|
|
params: Union[Iterable[torch.Tensor], torch.Tensor],
|
|
module: torch.nn.Module,
|
|
enabled: bool = True):
|
|
"""
|
|
Initialize the context manager to handle parameter gathering and partitioning for a replaced layer.
|
|
|
|
Args:
|
|
params (Iterable or torch.Tensor): A collection or single parameter to manage.
|
|
module (torch.nn.Module): The module that these parameters belong to.
|
|
enabled (bool): Flag indicating whether the parameter management is enabled (default: True).
|
|
"""
|
|
self.enabled = enabled
|
|
self.module = module
|
|
if not enabled:
|
|
return
|
|
|
|
# Ensure params is a list, whether it's a single param or iterable (e.g., model.parameters())
|
|
if isinstance(params, Iterable) and not isinstance(params, torch.Tensor):
|
|
self.params: List[torch.Tensor] = list(params) # Convert generators to a list for multiple iterations
|
|
else:
|
|
self.params: List[torch.Tensor] = [params] # Wrap single parameter in a list for uniform processing
|
|
|
|
# Check if the parameters belong to a replaced layer (indicated by a specific attribute)
|
|
if not any(self._is_replaced_module_weight(p) for p in params):
|
|
self.enabled = False
|
|
return
|
|
|
|
def _is_replaced_module_weight(self, param: torch.Tensor) -> bool:
|
|
"""
|
|
Helper function to determine if a parameter belongs to a replaced module.
|
|
|
|
Args:
|
|
param (torch.Tensor): The parameter to check.
|
|
|
|
Returns:
|
|
bool: True if the parameter belongs to a replaced module, False otherwise.
|
|
"""
|
|
return getattr(param, DS_IS_REPLACED_MODULE, False)
|
|
|
|
def __enter__(self) -> None:
|
|
"""
|
|
Enter the context manager. If enabled, gather parameters for the replaced module.
|
|
"""
|
|
if self.enabled:
|
|
self.params[0].gather_params(self.params)
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> None:
|
|
"""
|
|
Exit the context manager. If enabled, partition the parameters for the replaced module.
|
|
"""
|
|
#TODO : Check whether there are any missing attributes.
|
|
if self.enabled:
|
|
self.params[0]._tp_partition(self.params)
|
|
|
|
|
|
class LinearAllreduce(TensorParallel_Layer):
|
|
|
|
def __init__(self, module, mp_group, **kwargs):
|
|
super(LinearAllreduce, self).__init__(mp_group, **kwargs)
|
|
self.weight = module.weight
|
|
self.bias = module.bias
|
|
|
|
if self._should_materialize_tp_partition():
|
|
self._tp_partition([self.weight, self.bias])
|
|
self.support_training = True
|
|
self.config_tp_params(self.weight)
|
|
if self.bias is not None:
|
|
# bias here is not tp params
|
|
self.config_requires_grad(self.bias)
|
|
self._mark_uc_metadata()
|
|
|
|
def forward(self, input):
|
|
output = torch.matmul(input, self.weight.transpose(-1, -2))
|
|
output = RowParallel.apply(self.mp_group, output, not self.is_training_mode())
|
|
if self.bias is not None:
|
|
output = add_bias(output, self.bias)
|
|
return output
|
|
|
|
@torch.no_grad()
|
|
def gather_params(self, params_list):
|
|
|
|
for idx, param in enumerate(params_list):
|
|
if param is None or idx > 0:
|
|
# don't gather bias
|
|
return
|
|
params_list[idx].data_partition = param.data
|
|
param = param.transpose(0, 1).contiguous()
|
|
|
|
output_param = torch.empty(self.tp_world_size * param.shape[0],
|
|
param.shape[1],
|
|
dtype=param.dtype,
|
|
device=param.device)
|
|
dist.all_gather_into_tensor(output_param, param, group=self.mp_group)
|
|
params_list[idx].data = output_param.transpose(0, 1).contiguous()
|
|
return
|
|
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
|
|
if not self.is_training_mode():
|
|
self.uneven_partition(params_list)
|
|
return
|
|
|
|
else:
|
|
for idx, param in enumerate(params_list):
|
|
if param is None:
|
|
# don't slipt bias
|
|
return
|
|
if idx > 0: # move bias to device at initialization
|
|
_partition = self.move(param).detach()
|
|
params_list[idx].data = _partition
|
|
return
|
|
|
|
_partition = torch.chunk(param, self.tp_world_size, dim=-1)[self.tp_index]
|
|
|
|
_partition = self.move(_partition).detach()
|
|
|
|
params_list[idx].data = _partition
|
|
|
|
def uneven_partition(self, params_list):
|
|
for idx, param in enumerate(params_list):
|
|
if param is None or idx > 0:
|
|
# don't slipt bias
|
|
return
|
|
assert self.name is not None, "The module name must be provided in the initialization."
|
|
_partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[1], self.tp_world_size,
|
|
self.name),
|
|
dim=1)[self.tp_index]
|
|
|
|
_partition = self.move(_partition).detach()
|
|
params_list[idx].data = _partition
|
|
|
|
def _mark_uc_metadata(self):
|
|
original_weight_shape = (self.weight.shape[0], self.weight.shape[1] * self.tp_world_size)
|
|
self._set_param_uc_meta(self.weight,
|
|
partition_type='row',
|
|
partition_dim=1,
|
|
logical_shape=original_weight_shape,
|
|
output_shape=(original_weight_shape[0], ),
|
|
original_shape=original_weight_shape)
|
|
if self.bias is not None:
|
|
self._set_param_uc_meta(self.bias,
|
|
partition_type='row',
|
|
partition_dim=None,
|
|
logical_shape=tuple(self.bias.shape),
|
|
output_shape=tuple(self.bias.shape),
|
|
original_shape=tuple(self.bias.shape),
|
|
is_bias=True,
|
|
replicated=True)
|
|
|
|
|
|
#remove kwargs from partition.
|
|
class LinearLayer(TensorParallel_Layer):
|
|
|
|
def __init__(self, module, mp_group=None, skip_partition=False, **kwargs):
|
|
super(LinearLayer, self).__init__(mp_group, **kwargs)
|
|
self.weight = module.weight
|
|
self.bias = module.bias
|
|
if not skip_partition and self._should_materialize_tp_partition():
|
|
self._tp_partition([self.weight, self.bias])
|
|
self.support_training = True
|
|
self.config_tp_params(self.weight)
|
|
if self.bias is not None:
|
|
self.config_tp_params(self.bias)
|
|
self._mark_uc_metadata()
|
|
|
|
def forward(self, input):
|
|
if not self.__class__.tp_overlap_comm:
|
|
if getattr(self, 'mp_group', None) is not None:
|
|
input = ColumnParallel.apply(self.mp_group, input)
|
|
output = torch.matmul(input, self.weight.transpose(-1, -2))
|
|
if self.bias is not None:
|
|
output = add_bias(output, self.bias)
|
|
else:
|
|
output = AsyncColumnParallel.apply(self.mp_group, input, self.weight, self.bias)
|
|
|
|
return output
|
|
|
|
@torch.no_grad()
|
|
def gather_params(self, params_list):
|
|
# Does not support uneven shard.
|
|
for idx, param in enumerate(params_list):
|
|
|
|
params_list[idx].data_partition = param.data
|
|
output_param = torch.empty((self.tp_world_size * param.shape[0], *param.shape[1:]),
|
|
dtype=param.dtype,
|
|
device=param.device)
|
|
dist.all_gather_into_tensor(output_param, param, group=self.mp_group)
|
|
params_list[idx].data = output_param.contiguous()
|
|
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
|
|
if not self.is_training_mode():
|
|
self.uneven_partition(params_list)
|
|
return
|
|
for idx, param in enumerate(params_list):
|
|
if param is None:
|
|
return
|
|
#split bias if provide
|
|
_partition = torch.chunk(param, self.tp_world_size, dim=0)[self.tp_index]
|
|
|
|
_partition = self.move(_partition).detach()
|
|
|
|
params_list[idx].data = _partition
|
|
|
|
def uneven_partition(self, params_list):
|
|
|
|
for idx, param in enumerate(params_list):
|
|
if param is None:
|
|
#split bias if provide
|
|
return
|
|
assert self.name is not None, "The module name must be provided in the initialization."
|
|
_partition = params_list[idx].split(get_shard_size_list(params_list[idx].shape[0], self.tp_world_size,
|
|
self.name),
|
|
dim=0)[self.tp_index]
|
|
|
|
_partition = self.move(_partition).detach()
|
|
|
|
params_list[idx].data = _partition
|
|
|
|
def _mark_uc_metadata(self):
|
|
original_out_dim = self.weight.shape[0] * self.tp_world_size
|
|
original_weight_shape = (original_out_dim, self.weight.shape[1])
|
|
self._set_param_uc_meta(self.weight,
|
|
partition_type='column',
|
|
partition_dim=0,
|
|
logical_shape=original_weight_shape,
|
|
output_shape=(original_out_dim, ),
|
|
original_shape=original_weight_shape)
|
|
if self.bias is not None:
|
|
original_bias_shape = (self.bias.shape[0] * self.tp_world_size, )
|
|
self._set_param_uc_meta(self.bias,
|
|
partition_type='column',
|
|
partition_dim=0,
|
|
logical_shape=original_bias_shape,
|
|
output_shape=original_bias_shape,
|
|
original_shape=original_bias_shape,
|
|
is_bias=True)
|
|
|
|
# for bwc
|
|
@classmethod
|
|
def from_weights(cls, weight_shape=None, dtype=torch.half, weight=None, bias=None):
|
|
if weight is not None:
|
|
in_features = weight.shape[1]
|
|
out_features = weight.shape[0]
|
|
linear = nn.Linear(in_features, out_features, bias=(bias is not None))
|
|
linear.weight.data = weight
|
|
if bias is not None:
|
|
linear.bias.data = bias
|
|
else:
|
|
in_features = weight_shape[1]
|
|
out_features = weight_shape[0]
|
|
linear = nn.Linear(in_features, out_features, bias=(bias is not None))
|
|
return cls(linear, skip_partition=True)
|
|
|
|
|
|
class FusedModuleWrapper:
|
|
|
|
def __init__(self, fused_module: nn.Module):
|
|
self.fused_module = fused_module
|
|
|
|
def __getattr__(self, module):
|
|
return self.fused_module
|
|
|
|
|
|
class fused_LinearLayer(LinearLayer):
|
|
|
|
def __init__(self, module, mp_group, skip_partition=False, **kwargs):
|
|
assert kwargs.get('fused_module') is not None, "'fused_module' is required but not provided"
|
|
# Use the warp class to avoid module circular references.
|
|
self.fused_module = FusedModuleWrapper(kwargs.get('fused_module'))
|
|
super().__init__(module, mp_group, skip_partition, **kwargs)
|
|
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
for idx, param in enumerate(params_list):
|
|
if param is None:
|
|
return
|
|
|
|
_partition = prepare_tp_fused_qkvw(self.fused_module.module, param, self.tp_world_size, self.tp_index)
|
|
|
|
_partition = self.move(_partition).detach()
|
|
|
|
params_list[idx].data = _partition
|
|
|
|
|
|
class conv_LinearLayer(LinearLayer):
|
|
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
weight = None
|
|
bias = None
|
|
if len(params_list) == 1:
|
|
weight = params_list[0]
|
|
elif len(params_list) == 2:
|
|
weight, bias = params_list[0], params_list[1]
|
|
_partition = weight.data.split(get_shard_size_list(weight.shape[0], self.tp_world_size, self.name),
|
|
dim=1)[self.tp_index]
|
|
_partition = self.move(_partition).detach()
|
|
weight.data = _partition
|
|
|
|
if bias is not None:
|
|
_partition = bias.data.split(get_shard_size_list(weight.shape[1], self.tp_world_size, self.name),
|
|
dim=0)[self.tp_index]
|
|
_partition = self.move(_partition).detach()
|
|
|
|
bias.data = _partition
|
|
|
|
|
|
#override the subclasses related to weight splitting.
|
|
class Yuan_LinearAllreduce(LinearAllreduce):
|
|
|
|
#Yuan2
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index,
|
|
self.tp_world_size, False)
|
|
params_list[0].data = weight
|
|
if bias is not None:
|
|
params_list[1].data = bias
|
|
|
|
|
|
class Yuan_LinearLayer(LinearLayer):
|
|
#Yuan2
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
weight, bias = shard_value_with_share_qk(params_list[0].data, params_list[1], self.tp_index,
|
|
self.tp_world_size, True)
|
|
params_list[0].data = self.move(weight).detach()
|
|
if bias is not None:
|
|
params_list[1].data = self.move(bias).detach()
|
|
|
|
|
|
class GateUpPack_LinearLayer(LinearLayer):
|
|
# chatGLM2, chatGLM2
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
weight, bias = shard_chunk_mlp(params_list[0].data, params_list[1], self.tp_index, self.tp_world_size)
|
|
params_list[0].data = self.move(weight).detach()
|
|
if bias is not None:
|
|
params_list[1].data = self.move(bias).detach()
|
|
|
|
|
|
class Conv_LinearALlreduce(LinearAllreduce):
|
|
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
for idx, param in enumerate(params_list):
|
|
if param is None:
|
|
return
|
|
param.data = param.data.transpose(-1, -2).contiguous()
|
|
|
|
_partition = param.split(get_shard_size_list(param.shape[0], self.tp_world_size, self.name),
|
|
dim=1)[self.tp_index]
|
|
|
|
_partition = self.move(_partition).detach()
|
|
|
|
params_list[idx].data = _partition
|
|
|
|
|
|
#override the subclasses related to fwd/bwd.
|
|
class LmHeadLinearAllreduce(LinearAllreduce):
|
|
|
|
def __init__(self, module, mp_group, **kwargs):
|
|
# set the fixed name before partition
|
|
self.name = "lm_head"
|
|
|
|
# In some tied_embedding cases, only the lm head is sharded, while the word embedding is not.
|
|
# Reinitialization is used to decouple them and prevent the word embedding from being sharded.
|
|
# This should also be effective for cases where both are sharded in tied_embedding scenarios.
|
|
|
|
# TODO: Training scenario-related tests, is it necessary to re-implement the vocab parallel module?
|
|
module.weight = nn.Parameter(module.weight.clone().detach())
|
|
if hasattr(module, 'bias') and module.bias is not None:
|
|
module.bias = nn.Parameter(module.bias.clone().detach())
|
|
super().__init__(module, mp_group, **kwargs)
|
|
|
|
def forward(self, input):
|
|
input_shard_size = get_shard_size(input.shape[-1], self.tp_world_size, "lm_head")
|
|
input_shard_offset = sum(get_shard_size_list(input.shape[-1], self.tp_world_size, "lm_head")[0:self.tp_index])
|
|
output = torch.matmul(input[:, :, input_shard_offset:input_shard_offset + input_shard_size],
|
|
self.weight.transpose(-1, -2))
|
|
if self.mp_group is not None:
|
|
dist.inference_all_reduce(output, group=self.mp_group)
|
|
if self.bias is not None:
|
|
output = add_bias(output, self.bias)
|
|
return output
|
|
|
|
|
|
class TensorParallelConv2d(nn.Module):
|
|
|
|
def __init__(self, conv, rank, world_size, shard_by_oc):
|
|
super().__init__()
|
|
self.rank = rank
|
|
self.world_size = world_size
|
|
self.shard_by_oc = shard_by_oc
|
|
self.shard_weights(conv)
|
|
|
|
# Split along the input/output channel depending on whether it is the last conv layer.
|
|
def shard_weights(self, conv):
|
|
if self.shard_by_oc:
|
|
total_size = conv.weight.shape[0]
|
|
else:
|
|
total_size = conv.weight.shape[1]
|
|
bias_data = None
|
|
cols_per_rank = [0]
|
|
for i in range(self.world_size - 1, -1, -1):
|
|
cols = total_size // self.world_size
|
|
if i < total_size % self.world_size:
|
|
cols += 1
|
|
cols_per_rank.append(cols_per_rank[-1] + cols)
|
|
weight_data = conv.weight.data
|
|
if self.shard_by_oc:
|
|
# not last conv layer, split output channel
|
|
weight_data = weight_data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]]
|
|
if conv.bias is not None:
|
|
bias_data = conv.bias.data[cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]]
|
|
else:
|
|
# last conv layer, split input channel
|
|
weight_data = weight_data[:, cols_per_rank[self.rank]:cols_per_rank[self.rank + 1]]
|
|
if conv.bias is not None:
|
|
bias_data = conv.bias.data / float(self.world_size)
|
|
self.conv = nn.Conv2d(weight_data.shape[1], weight_data.shape[0], conv.kernel_size, conv.stride, conv.padding,
|
|
conv.dilation, conv.groups, conv.bias is not None, conv.padding_mode)
|
|
self.conv.weight = torch.nn.Parameter(weight_data)
|
|
if conv.bias is not None:
|
|
self.conv.bias = torch.nn.Parameter(bias_data)
|
|
del conv
|
|
|
|
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
|
return self.conv(input)
|
|
|
|
|
|
class TensorParallelOcShardConv2d(TensorParallelConv2d):
|
|
|
|
def __init__(self, conv, rank, world_size):
|
|
super().__init__(conv, rank, world_size, True)
|
|
|
|
|
|
class TensorParallelIcShardConv2d(TensorParallelConv2d):
|
|
|
|
def __init__(self, conv, rank, world_size):
|
|
super().__init__(conv, rank, world_size, False)
|
|
|
|
def forward(self, input: torch.Tensor) -> torch.Tensor:
|
|
out = self.conv(input)
|
|
if self.world_size > 1:
|
|
dist.inference_all_reduce(out)
|
|
return out
|
|
|
|
|
|
class Normalize(nn.Module):
|
|
|
|
def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None, bias=None):
|
|
super(Normalize, self).__init__()
|
|
if weight is not None:
|
|
self.weight = weight
|
|
self.bias = bias
|
|
else:
|
|
self.norm = nn.LayerNorm(dim, eps=eps).to(dtype).to(get_accelerator().current_device_name())
|
|
self.weight = self.norm.weight
|
|
self.bias = self.norm.bias
|
|
|
|
self.eps = eps
|
|
|
|
def forward(self, input):
|
|
return nn.functional.layer_norm(input, input.shape[-1:], self.weight, self.bias, eps=self.eps)
|
|
|
|
|
|
class EmbeddingLayer(nn.Module):
|
|
|
|
def __init__(self, weight_shape=None, dtype=torch.half, weight=None, bias=None):
|
|
super(EmbeddingLayer, self).__init__()
|
|
if weight is None:
|
|
self.weight = Parameter(
|
|
torch.empty(weight_shape[0],
|
|
weight_shape[1],
|
|
dtype=dtype,
|
|
device=get_accelerator().current_device_name()))
|
|
else:
|
|
self.weight = weight
|
|
|
|
def forward(self, input):
|
|
return F.embedding(input, self.weight)
|
|
|
|
|
|
class OPTEmbedding(EmbeddingLayer):
|
|
"""
|
|
This module learns positional embeddings up to a fixed maximum size.
|
|
"""
|
|
|
|
def __init__(self, weight_shape=None, weight=None, bias=None):
|
|
# OPT is set up so that if padding_idx is specified then offset the embedding ids by 2
|
|
# and adjust num_embeddings appropriately. Other models don't have this hack
|
|
self.offset = 2
|
|
super().__init__(weight_shape, weight=weight)
|
|
|
|
def forward(self, attention_mask: torch.LongTensor, past_key_values_length: int = 0, position_ids: int = 0):
|
|
"""`input_ids_shape` is expected to be [bsz x seqlen]."""
|
|
attention_mask = attention_mask.long()
|
|
|
|
# create positions depending on attention_mask
|
|
positions = (torch.cumsum(attention_mask, dim=1).type_as(attention_mask) * attention_mask).long() - 1
|
|
|
|
# cut positions if `past_key_values_length` is > 0
|
|
positions = positions[:, past_key_values_length:]
|
|
|
|
return super().forward(positions + self.offset)
|
|
|
|
|
|
def _shape_prod(values):
|
|
result = 1
|
|
for val in values:
|
|
result *= val
|
|
return result
|
|
|
|
|
|
def _normalize_shape_spec(shape):
|
|
if isinstance(shape, list):
|
|
return tuple(_normalize_shape_spec(item) for item in shape)
|
|
if isinstance(shape, tuple):
|
|
return tuple(_normalize_shape_spec(item) if isinstance(item, list) else item for item in shape)
|
|
return shape
|
|
|
|
|
|
def _infer_subparam_logical_shapes(weight_shape, shape, partition_dim, name=None):
|
|
shape = _normalize_shape_spec(shape)
|
|
if not isinstance(shape, tuple):
|
|
raise ValueError("AutoTP shape must be a tuple for sub-parameter partitioning.")
|
|
if partition_dim < 0 or partition_dim >= len(shape):
|
|
raise ValueError(f"AutoTP partition_dim {partition_dim} is out of range for shape length {len(shape)}.")
|
|
|
|
layer_label = f"AutoTP layer '{name}'" if name else "AutoTP layer"
|
|
partition_elem = shape[partition_dim]
|
|
subparam_sizes = None
|
|
num_subparams = None
|
|
|
|
if isinstance(partition_elem, tuple):
|
|
if len(partition_elem) == 0:
|
|
raise ValueError(f"{layer_label} sub-parameter size tuple cannot be empty.")
|
|
if any(isinstance(val, tuple) for val in partition_elem):
|
|
raise ValueError(f"{layer_label} supports only 1-level nesting at partition_dim.")
|
|
if any((not isinstance(val, int)) or val <= 0 for val in partition_elem):
|
|
raise ValueError(f"{layer_label} sub-parameter sizes must be positive integers.")
|
|
subparam_sizes = tuple(int(val) for val in partition_elem)
|
|
partition_dim_size = sum(subparam_sizes)
|
|
elif isinstance(partition_elem, int):
|
|
if partition_elem == -1:
|
|
partition_dim_size = None
|
|
elif partition_elem > 0:
|
|
num_subparams = partition_elem
|
|
partition_dim_size = None
|
|
else:
|
|
raise ValueError(f"{layer_label} partition_dim spec must be positive integer or -1.")
|
|
else:
|
|
raise ValueError(f"{layer_label} partition_dim spec must be int or tuple.")
|
|
|
|
logical_dims = []
|
|
for idx, dim in enumerate(shape):
|
|
if idx == partition_dim:
|
|
logical_dims.append(partition_dim_size)
|
|
continue
|
|
if isinstance(dim, tuple):
|
|
raise ValueError(f"{layer_label} nested tuple only allowed at partition_dim={partition_dim}.")
|
|
if isinstance(dim, int):
|
|
if dim == -1:
|
|
logical_dims.append(None)
|
|
elif dim > 0:
|
|
logical_dims.append(dim)
|
|
else:
|
|
raise ValueError(f"{layer_label} shape dimensions must be positive integers or -1.")
|
|
else:
|
|
raise ValueError(f"{layer_label} shape dimensions must be integers.")
|
|
|
|
total_numel = _shape_prod(weight_shape)
|
|
known_product = _shape_prod([dim for dim in logical_dims if dim is not None])
|
|
unknown_indices = [idx for idx, dim in enumerate(logical_dims) if dim is None]
|
|
|
|
if len(unknown_indices) == 0:
|
|
if known_product != total_numel:
|
|
raise ValueError(f"{layer_label} shape product {known_product} != weight numel {total_numel}.")
|
|
elif len(unknown_indices) == 1:
|
|
inferred = total_numel // known_product
|
|
if inferred * known_product != total_numel:
|
|
raise ValueError(f"{layer_label} cannot infer shape for weight with numel {total_numel}.")
|
|
logical_dims[unknown_indices[0]] = inferred
|
|
else:
|
|
if len(shape) == len(weight_shape):
|
|
for idx in unknown_indices:
|
|
logical_dims[idx] = weight_shape[idx]
|
|
if _shape_prod(logical_dims) != total_numel:
|
|
raise ValueError(
|
|
f"{layer_label} shape product {_shape_prod(logical_dims)} != weight numel {total_numel}.")
|
|
else:
|
|
raise ValueError(f"{layer_label} shape has multiple inferred dims and is ambiguous for weight.")
|
|
|
|
logical_shape = tuple(logical_dims)
|
|
if logical_shape[-1] != weight_shape[-1]:
|
|
raise ValueError(
|
|
f"{layer_label} shape last dim {logical_shape[-1]} must match weight input dim {weight_shape[-1]}.")
|
|
|
|
output_shape = logical_shape[:-1]
|
|
if len(output_shape) == 0:
|
|
raise ValueError(f"{layer_label} shape must include at least one output dimension.")
|
|
if _shape_prod(output_shape) != weight_shape[0]:
|
|
raise ValueError(
|
|
f"{layer_label} output shape product {_shape_prod(output_shape)} != weight output dim {weight_shape[0]}.")
|
|
|
|
partition_dim_size = logical_shape[partition_dim]
|
|
if partition_dim_size is None or partition_dim_size <= 0:
|
|
raise ValueError(f"{layer_label} partition_dim size must be a positive integer.")
|
|
|
|
if num_subparams is not None:
|
|
if partition_dim_size % num_subparams != 0:
|
|
raise ValueError(
|
|
f"{layer_label} partition_dim size {partition_dim_size} not divisible by sub-param count {num_subparams}."
|
|
)
|
|
subparam_sizes = tuple([partition_dim_size // num_subparams] * num_subparams)
|
|
|
|
if subparam_sizes is not None and sum(subparam_sizes) != partition_dim_size:
|
|
raise ValueError(
|
|
f"{layer_label} sub-parameter sizes sum {sum(subparam_sizes)} != partition_dim size {partition_dim_size}.")
|
|
|
|
bias_partition_dim = partition_dim if partition_dim < len(output_shape) else None
|
|
return logical_shape, output_shape, subparam_sizes, bias_partition_dim
|
|
|
|
|
|
def _partition_logical_tensor(tensor, partition_dim, tp_world_size, tp_index, name=None, subparam_sizes=None):
|
|
if tp_world_size == 1:
|
|
return tensor
|
|
layer_label = f"AutoTP layer '{name}'" if name else "AutoTP layer"
|
|
if subparam_sizes:
|
|
for size in subparam_sizes:
|
|
if size % tp_world_size != 0:
|
|
raise ValueError(f"{layer_label} sub-parameter size {size} not divisible by tp_size {tp_world_size}.")
|
|
sub_params = torch.split(tensor, subparam_sizes, dim=partition_dim)
|
|
partitioned_sub_params = [torch.chunk(sp, tp_world_size, dim=partition_dim)[tp_index] for sp in sub_params]
|
|
return torch.cat(partitioned_sub_params, dim=partition_dim)
|
|
if tensor.shape[partition_dim] % tp_world_size != 0:
|
|
raise ValueError(
|
|
f"{layer_label} partition_dim size {tensor.shape[partition_dim]} not divisible by tp_size {tp_world_size}."
|
|
)
|
|
return torch.chunk(tensor, tp_world_size, dim=partition_dim)[tp_index]
|
|
|
|
|
|
def _all_gather_along_dim(tensor, partition_dim, mp_group, tp_world_size):
|
|
if mp_group is None or tp_world_size == 1:
|
|
return tensor
|
|
perm = [partition_dim] + [idx for idx in range(tensor.dim()) if idx != partition_dim]
|
|
inv_perm = [0] * len(perm)
|
|
for idx, dim in enumerate(perm):
|
|
inv_perm[dim] = idx
|
|
tensor_perm = tensor.permute(perm).contiguous()
|
|
output = torch.empty((tp_world_size * tensor_perm.shape[0], *tensor_perm.shape[1:]),
|
|
dtype=tensor.dtype,
|
|
device=tensor.device)
|
|
dist.all_gather_into_tensor(output, tensor_perm, group=mp_group)
|
|
return output.permute(inv_perm).contiguous()
|
|
|
|
|
|
def _gather_logical_tensor(tensor,
|
|
logical_shape,
|
|
partition_dim,
|
|
mp_group,
|
|
tp_world_size,
|
|
name=None,
|
|
subparam_sizes=None):
|
|
if mp_group is None or tp_world_size == 1:
|
|
return tensor.reshape(logical_shape)
|
|
layer_label = f"AutoTP layer '{name}'" if name else "AutoTP layer"
|
|
if logical_shape[partition_dim] % tp_world_size != 0:
|
|
raise ValueError(
|
|
f"{layer_label} partition_dim size {logical_shape[partition_dim]} not divisible by tp_size {tp_world_size}."
|
|
)
|
|
partitioned_shape = list(logical_shape)
|
|
partitioned_shape[partition_dim] = logical_shape[partition_dim] // tp_world_size
|
|
tensor_view = tensor.reshape(partitioned_shape)
|
|
|
|
if subparam_sizes:
|
|
for size in subparam_sizes:
|
|
if size % tp_world_size != 0:
|
|
raise ValueError(f"{layer_label} sub-parameter size {size} not divisible by tp_size {tp_world_size}.")
|
|
partitioned_sizes = [size // tp_world_size for size in subparam_sizes]
|
|
sub_params = torch.split(tensor_view, partitioned_sizes, dim=partition_dim)
|
|
gathered_sub_params = [_all_gather_along_dim(sp, partition_dim, mp_group, tp_world_size) for sp in sub_params]
|
|
return torch.cat(gathered_sub_params, dim=partition_dim)
|
|
return _all_gather_along_dim(tensor_view, partition_dim, mp_group, tp_world_size)
|
|
|
|
|
|
class SubParamLinearLayer(TensorParallel_Layer):
|
|
"""
|
|
Column-parallel linear layer with sub-parameter support.
|
|
|
|
Handles cases where weights contain multiple logical sub-parameters
|
|
that need to be partitioned separately (e.g., fused QKV, chunked MLP, GQA).
|
|
|
|
The `shape` parameter controls how the weight is viewed and partitioned:
|
|
- (3, -1) with partition_dim=0: 3 equal sub-params, partition each at dim 0
|
|
- ((q, k, v), -1) with partition_dim=0: 3 unequal sub-params (1-level nesting)
|
|
"""
|
|
|
|
def __init__(self, module, mp_group, shape, partition_dim=0, **kwargs):
|
|
super(SubParamLinearLayer, self).__init__(mp_group, **kwargs)
|
|
self.weight = module.weight
|
|
self.bias = module.bias
|
|
self.shape = shape
|
|
self.partition_dim = partition_dim
|
|
|
|
self._orig_weight_shape = tuple(module.weight.shape)
|
|
self._orig_bias_shape = tuple(module.bias.shape) if self.bias is not None else None
|
|
(self._logical_shape, self._output_shape, self._subparam_sizes,
|
|
self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
|
|
self.partition_dim, self.name)
|
|
if self.bias is not None and self.bias.numel() != _shape_prod(self._output_shape):
|
|
raise ValueError(f"AutoTP layer '{self.name}' bias size {self.bias.numel()} does not match output shape "
|
|
f"{self._output_shape}.")
|
|
|
|
if self._should_materialize_tp_partition():
|
|
self._tp_partition([self.weight, self.bias])
|
|
self.support_training = True
|
|
self.config_tp_params(self.weight)
|
|
if self.bias is not None:
|
|
self.config_tp_params(self.bias)
|
|
self._mark_uc_metadata()
|
|
|
|
def forward(self, input):
|
|
if getattr(self, 'mp_group', None) is not None:
|
|
input = ColumnParallel.apply(self.mp_group, input)
|
|
output = torch.matmul(input, self.weight.transpose(-1, -2))
|
|
if self.bias is not None:
|
|
output = add_bias(output, self.bias)
|
|
return output
|
|
|
|
@torch.no_grad()
|
|
def gather_params(self, params_list):
|
|
"""Gather partitioned parameters back to full size."""
|
|
for idx, param in enumerate(params_list):
|
|
if param is None:
|
|
continue
|
|
params_list[idx].data_partition = param.data
|
|
if idx == 0:
|
|
full_view = _gather_logical_tensor(param,
|
|
self._logical_shape,
|
|
self.partition_dim,
|
|
self.mp_group,
|
|
self.tp_world_size,
|
|
name=self.name,
|
|
subparam_sizes=self._subparam_sizes)
|
|
params_list[idx].data = full_view.reshape(self._orig_weight_shape)
|
|
else:
|
|
if self._bias_partition_dim is None:
|
|
params_list[idx].data = param.data
|
|
else:
|
|
full_bias_view = _gather_logical_tensor(param,
|
|
self._output_shape,
|
|
self._bias_partition_dim,
|
|
self.mp_group,
|
|
self.tp_world_size,
|
|
name=self.name,
|
|
subparam_sizes=self._subparam_sizes)
|
|
params_list[idx].data = full_bias_view.reshape(self._orig_bias_shape)
|
|
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
weight = params_list[0]
|
|
if weight is None:
|
|
return
|
|
|
|
weight_view = weight.reshape(self._logical_shape)
|
|
partitioned_view = _partition_logical_tensor(weight_view,
|
|
self.partition_dim,
|
|
self.tp_world_size,
|
|
self.tp_index,
|
|
name=self.name,
|
|
subparam_sizes=self._subparam_sizes)
|
|
params_list[0].data = self.move(partitioned_view.reshape(-1, partitioned_view.shape[-1])).detach()
|
|
|
|
if params_list[1] is not None:
|
|
if self._bias_partition_dim is None:
|
|
params_list[1].data = self.move(params_list[1]).detach()
|
|
else:
|
|
bias_view = params_list[1].reshape(self._output_shape)
|
|
bias_partitioned = _partition_logical_tensor(bias_view,
|
|
self._bias_partition_dim,
|
|
self.tp_world_size,
|
|
self.tp_index,
|
|
name=self.name,
|
|
subparam_sizes=self._subparam_sizes)
|
|
params_list[1].data = self.move(bias_partitioned.reshape(-1)).detach()
|
|
|
|
def _mark_uc_metadata(self):
|
|
self._set_param_uc_meta(self.weight,
|
|
partition_type='column',
|
|
partition_dim=self.partition_dim,
|
|
logical_shape=self._logical_shape,
|
|
output_shape=self._output_shape,
|
|
sub_param_shape=self.shape,
|
|
sub_param_sizes=self._subparam_sizes,
|
|
target_partition_shape=self.weight.shape,
|
|
original_shape=self._orig_weight_shape)
|
|
if self.bias is not None:
|
|
self._set_param_uc_meta(
|
|
self.bias,
|
|
partition_type='column',
|
|
partition_dim=self._bias_partition_dim,
|
|
logical_shape=self._output_shape,
|
|
output_shape=self._output_shape,
|
|
sub_param_shape=self.shape if self._bias_partition_dim is not None else None,
|
|
sub_param_sizes=self._subparam_sizes if self._bias_partition_dim is not None else None,
|
|
target_partition_shape=self.bias.shape,
|
|
original_shape=self._orig_bias_shape,
|
|
is_bias=True,
|
|
replicated=self._bias_partition_dim is None)
|
|
|
|
|
|
class SubParamLinearAllreduce(TensorParallel_Layer):
|
|
"""
|
|
Row-parallel linear layer with sub-parameter support (AllReduce after forward).
|
|
|
|
Handles cases where weights contain multiple logical sub-parameters
|
|
that need to be partitioned separately.
|
|
"""
|
|
|
|
def __init__(self, module, mp_group, shape, partition_dim=1, **kwargs):
|
|
super(SubParamLinearAllreduce, self).__init__(mp_group, **kwargs)
|
|
self.weight = module.weight
|
|
self.bias = module.bias
|
|
self.shape = shape
|
|
self.partition_dim = partition_dim
|
|
|
|
self._orig_weight_shape = tuple(module.weight.shape)
|
|
self._orig_bias_shape = tuple(module.bias.shape) if self.bias is not None else None
|
|
(self._logical_shape, self._output_shape, self._subparam_sizes,
|
|
self._bias_partition_dim) = _infer_subparam_logical_shapes(self._orig_weight_shape, self.shape,
|
|
self.partition_dim, self.name)
|
|
|
|
if self._should_materialize_tp_partition():
|
|
self._tp_partition([self.weight, self.bias])
|
|
self.support_training = True
|
|
self.config_tp_params(self.weight)
|
|
if self.bias is not None:
|
|
self.config_requires_grad(self.bias)
|
|
self._mark_uc_metadata()
|
|
|
|
def forward(self, input):
|
|
output = torch.matmul(input, self.weight.transpose(-1, -2))
|
|
output = RowParallel.apply(self.mp_group, output, not self.is_training_mode())
|
|
if self.bias is not None:
|
|
output = add_bias(output, self.bias)
|
|
return output
|
|
|
|
@torch.no_grad()
|
|
def gather_params(self, params_list):
|
|
"""Gather partitioned parameters back to full size."""
|
|
for idx, param in enumerate(params_list):
|
|
if param is None or idx > 0:
|
|
# don't gather bias for row parallel
|
|
return
|
|
params_list[idx].data_partition = param.data
|
|
full_view = _gather_logical_tensor(param,
|
|
self._logical_shape,
|
|
self.partition_dim,
|
|
self.mp_group,
|
|
self.tp_world_size,
|
|
name=self.name,
|
|
subparam_sizes=self._subparam_sizes)
|
|
params_list[idx].data = full_view.reshape(self._orig_weight_shape)
|
|
|
|
@torch.no_grad()
|
|
def _tp_partition(self, params_list):
|
|
weight = params_list[0]
|
|
if weight is None:
|
|
return
|
|
|
|
weight_view = weight.reshape(self._logical_shape)
|
|
partitioned_view = _partition_logical_tensor(weight_view,
|
|
self.partition_dim,
|
|
self.tp_world_size,
|
|
self.tp_index,
|
|
name=self.name,
|
|
subparam_sizes=self._subparam_sizes)
|
|
params_list[0].data = self.move(partitioned_view.reshape(-1, partitioned_view.shape[-1])).detach()
|
|
|
|
# Bias is not partitioned for row parallel (it's applied after all-reduce)
|
|
if params_list[1] is not None:
|
|
params_list[1].data = self.move(params_list[1]).detach()
|
|
|
|
def _mark_uc_metadata(self):
|
|
self._set_param_uc_meta(self.weight,
|
|
partition_type='row',
|
|
partition_dim=self.partition_dim,
|
|
logical_shape=self._logical_shape,
|
|
output_shape=self._output_shape,
|
|
sub_param_shape=self.shape,
|
|
sub_param_sizes=self._subparam_sizes,
|
|
target_partition_shape=self.weight.shape,
|
|
original_shape=self._orig_weight_shape)
|
|
if self.bias is not None:
|
|
self._set_param_uc_meta(self.bias,
|
|
partition_type='row',
|
|
partition_dim=None,
|
|
logical_shape=self._orig_bias_shape,
|
|
output_shape=self._orig_bias_shape,
|
|
original_shape=self._orig_bias_shape,
|
|
target_partition_shape=self.bias.shape,
|
|
is_bias=True,
|
|
replicated=True)
|
|
|
|
|
|
class RMSNormalize(nn.Module):
|
|
|
|
def __init__(self, dim=None, dtype=torch.float, eps=1e-5, weight=None):
|
|
super(RMSNormalize, self).__init__()
|
|
if weight is not None:
|
|
self.weight = weight
|
|
else:
|
|
self.weight = nn.Parameter(torch.ones(dim, dtype=dtype, device=get_accelerator().current_device_name()))
|
|
|
|
self.eps = eps
|
|
|
|
def forward(self, hidden_states):
|
|
variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
|
|
hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
|
|
if self.weight.dtype in [torch.float16, torch.bfloat16]:
|
|
hidden_states = hidden_states.to(self.weight.dtype)
|
|
|
|
return hidden_states * self.weight
|