wehub-resource-sync
105 lines
4.9 KiB
Python
105 lines
4.9 KiB
Python
# Copyright (c) Microsoft Corporation.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
from typing import Optional, Tuple
|
|
|
|
import torch
|
|
|
|
from .types import ShardingType, DEFAULT_SHARD_GRANULARITY
|
|
|
|
|
|
def get_shard_endpoints(dim_size: int,
|
|
shard_rank: int,
|
|
num_shards: int,
|
|
granularity: int = DEFAULT_SHARD_GRANULARITY) -> Tuple[int, int]:
|
|
"""
|
|
Given a dimension to shard with size dim_size, return the start and end indices of the slice
|
|
that belong to the given rank.
|
|
|
|
The typical use of this is as an internal helper function, so see if there is a higher level
|
|
API that better suits the application.
|
|
|
|
Args:
|
|
dim_size (int): The size of the dimension to shard.
|
|
shard_rank (int): The rank of the shard to return.
|
|
num_shards (int): Total number of shards the dimension will be distributed across.
|
|
granularity (int): The minimum alignment of the shard endpoints. This is used to support
|
|
non-even head counts as well as align dimensions to cleaner GEMM boundaries.
|
|
"""
|
|
assert dim_size % granularity == 0, "Dimension size must be divisible by granularity"
|
|
|
|
total_chunks = dim_size // granularity
|
|
base_chunks_per_rank = total_chunks // num_shards
|
|
remainder_chunks = total_chunks % num_shards
|
|
|
|
start_chunk_id = shard_rank * base_chunks_per_rank + min(shard_rank, remainder_chunks)
|
|
end_chunk_id = start_chunk_id + base_chunks_per_rank + (1 if shard_rank < remainder_chunks else 0)
|
|
|
|
return start_chunk_id * granularity, end_chunk_id * granularity
|
|
|
|
|
|
def shard_param(param: Optional[torch.Tensor],
|
|
shard_mode: ShardingType,
|
|
shard_rank: int,
|
|
num_shards: int,
|
|
num_concatenated_matrices: int = 1,
|
|
granularity: int = 32,
|
|
bias_dims: int = 1) -> torch.Tensor:
|
|
"""
|
|
Utility for sharding a parameter. This will return the slice of the parameter that should
|
|
exist on the given shard_rank given the sharding configuration. The workflow here is
|
|
to find the minimum bounded Tensor to shard, get the slicing endpoints, and then concatenate
|
|
as needed.
|
|
|
|
The typical use of this is as an internal helper function, so see if there is a higher level
|
|
API that better suits the application.
|
|
|
|
Args:
|
|
param (torch.Tensor): The parameter to shard.
|
|
shard_mode (ShardingType): The type of sharding to apply. See ShardingType for more context.
|
|
shard_rank (int): The rank of the shard to return.
|
|
num_shards (int): Total number of shards the parameter will be distrbuted across.
|
|
num_concatenated_matrices (int): The number of matrices that have been concatenated together in the original
|
|
parameter. An example of this is a fused QKV projection matrix, where the `num_concatenated_matrices`
|
|
argument would be 3.
|
|
granularity (int): The minimum alignment of the shard endpoints. For attention projection matrices, this
|
|
should be set to the head size to support non-even sharding.
|
|
bias_dims (int): The number of dimensions that are considered bias dimensions. This is used to support
|
|
sharding of MoE and non-MoE biases on the same codepath.
|
|
"""
|
|
assert shard_rank < num_shards, "Shard rank must be less than num_shards"
|
|
|
|
# Easier to hide this inside of the sharding logic than to add checks in every model
|
|
# implementation.
|
|
if param is None:
|
|
return None
|
|
|
|
if num_shards == 1:
|
|
# Trivial case of no sharding.
|
|
return param
|
|
|
|
if shard_mode == ShardingType.OUTER_DIMENSION:
|
|
|
|
def get_matrices(dim_idx: int) -> torch.Tensor:
|
|
dim_size = param.size(dim_idx) // num_concatenated_matrices
|
|
start_channel_id, end_channel_id = get_shard_endpoints(dim_size, shard_rank, num_shards, granularity)
|
|
return torch.chunk(param, num_concatenated_matrices, dim=dim_idx), start_channel_id, end_channel_id
|
|
|
|
if param.ndim == bias_dims:
|
|
# Special case for bias parameters.
|
|
matrices, start_channel_id, end_channel_id = get_matrices(dim_idx=-1)
|
|
return torch.cat([mat[..., start_channel_id:end_channel_id] for mat in matrices], dim=-1)
|
|
else:
|
|
# General case for weight parameters. This assumes MoE parameters are stored in the format of
|
|
# [num_experts, out_features, in_features]
|
|
matrices, start_channel_id, end_channel_id = get_matrices(dim_idx=-2)
|
|
return torch.cat([mat[..., start_channel_id:end_channel_id, :] for mat in matrices], dim=-2)
|
|
|
|
elif shard_mode == ShardingType.INNER_DIMENSION:
|
|
dim_size = param.size(-1) // num_concatenated_matrices
|
|
start_channel_id, end_channel_id = get_shard_endpoints(dim_size, shard_rank, num_shards, granularity)
|
|
matrices = torch.chunk(param, num_concatenated_matrices, dim=-1)
|
|
return torch.cat([mat[..., start_channel_id:end_channel_id] for mat in matrices], dim=-1)
|