mirrors/deepspeedai--deepspeed
1
0
Fork 0
Code Issues Pull Requests Actions 12 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
deepspeed/inference/v2/model_implementations/sharding/__init__.py
+12
View File
@@ -0,0 +1,12 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
from .attn import *
from .attn_out import *
from .embedding import *
from .mlp import *
from .qkv import *
from .types import *
from .unembed import *
deepspeed/inference/v2/model_implementations/sharding/attn.py
+63
View File
@@ -0,0 +1,63 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
from typing import Optional, Tuple
def get_local_heads(shard_rank: int,
num_shards: int,
n_heads_q: int,
n_heads_kv: Optional[int] = None) -> Tuple[int, int]:
"""
Helper to determine the number of local heads of a given shard.
Args:
shard_rank (int): The rank of the shard.
num_shards (int): The total number of shards that attention is distributed over.
n_heads_q (int): The number of query heads.
n_heads_kv (int): The number of key/value heads. If not passed, it is assumed that
the number of query and key/value heads are the same.
"""
if n_heads_q < num_shards:
raise ValueError("There must be at least as many attention heads as there are shards.")
if n_heads_kv is None or n_heads_kv == n_heads_q:
# MHA attention
base_heads = n_heads_q // num_shards
extra_heads = n_heads_q % num_shards
if shard_rank < extra_heads:
return (base_heads + 1), (base_heads + 1)
else:
return base_heads, base_heads
else:
# GQA attention
if n_heads_q % n_heads_kv != 0:
raise ValueError("Must be an even ratio between query and key/value heads.")
if n_heads_kv < num_shards and num_shards % n_heads_kv != 0:
raise ValueError(
"If splitting a group across multiple shards, we must be able to distribute the groups evenly.")
if n_heads_kv >= num_shards and n_heads_kv % num_shards != 0:
raise ValueError("If parallelizing groups, must be able to evenly distribute them.")
q_ratio = n_heads_q // n_heads_kv
if n_heads_kv >= num_shards:
local_kv_heads = n_heads_kv // num_shards
local_q_heads = local_kv_heads * q_ratio
return local_q_heads, local_kv_heads
else:
group_sharding_size = num_shards // n_heads_kv
group_rank_idx = shard_rank % group_sharding_size
base_heads = q_ratio // group_sharding_size
extra_heads = q_ratio % group_sharding_size
if group_rank_idx < extra_heads:
return (base_heads + 1), 1
else:
return base_heads, 1
deepspeed/inference/v2/model_implementations/sharding/attn_out.py
+111
View File
@@ -0,0 +1,111 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
from typing import Optional
import torch
from .types import ShardingType
from .utils import shard_param, get_shard_endpoints
def shard_attn_out_param(param: torch.Tensor,
shard_rank: int,
num_shards: int,
head_size: int,
n_heads_q: Optional[int] = None,
n_heads_kv: Optional[int] = None) -> Optional[torch.Tensor]:
"""
Utility method for sharding an attention output parameter.
"""
if len(param.shape) == 1:
# We will do the bias addition on the 0th rank only rather than scale the parameter and
# implicitly reconstruct this in the distributed reduce.
return param if shard_rank == 0 else None
assert n_heads_kv is None or (n_heads_q is not None
and n_heads_kv is not None), "n_heads_kv should not be passed without n_heads_q"
mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv
if mha_sharding:
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards, granularity=head_size)
else:
assert param.shape[0] == head_size * n_heads_q, "GQA param shape is not correct"
# 32 KV heads, 16 shards for example
even_kv_sharding = n_heads_kv % num_shards == 0
# 8 KV heads, 16 shards for example
even_kv_distribution = num_shards % n_heads_kv == 0
assert even_kv_sharding or even_kv_distribution, "No partitioning algorithm for this yet."
if even_kv_sharding:
# Same as original sharding scenario
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards, granularity=head_size)
else:
# We will first do a sharding on the KV and Q to map to the one KV shard per group of Q.
q_sharding_degree = num_shards // n_heads_kv
kv_head = shard_rank // q_sharding_degree
q_sharding_rank = shard_rank % q_sharding_degree
q_factor = n_heads_q // n_heads_kv
q_chunk = param[..., q_factor * kv_head * head_size:q_factor * (kv_head + 1) * head_size]
return shard_param(q_chunk,
ShardingType.INNER_DIMENSION,
q_sharding_rank,
q_sharding_degree,
granularity=head_size)
def attn_out_in_features(out_features: int,
shard_rank: int,
num_shards: int,
head_size: int,
n_heads_q: Optional[int] = None,
n_heads_kv: Optional[int] = None) -> int:
"""
Helper to calculate the expected output projection dimension of a QKV projection matrix.
Args:
in_features (int): The model dimension.
shard_rank (int): Which rank to return the corresponding size for.
num_shards (int): The total number of shards the parameter is distributed across.
head_size (int): The size of each attention head.
n_heads_q (int): The number of query heads on the model. This only needs to be passed if the number
of query and key/value heads are different. If passed without n_heads_kv, default
MHA partitioning will be used.
n_heads_kv (int): The number of key and value heads on the model. This only needs to be passed
if the number of query and key/value heads are different. This argument cannot be passed without
also passing n_heads_q (we want to explicitly opt into GQA sharding).
"""
assert n_heads_kv is None or (n_heads_q is not None
and n_heads_kv is not None), "n_heads_kv should not be passed without n_heads_q"
mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv
if mha_sharding:
endpoints = get_shard_endpoints(out_features, shard_rank, num_shards, granularity=head_size)
return endpoints[1] - endpoints[0]
else:
if n_heads_kv >= num_shards:
assert n_heads_kv % num_shards == 0, "No partitioning algorithm for this yet."
n_local_groups = n_heads_kv // num_shards
group_size = n_heads_q // n_heads_kv
return n_local_groups * head_size * group_size
else:
assert num_shards % n_heads_kv == 0, "No partitioning algorithm for this yet."
q_split_degree = num_shards // n_heads_kv
q_split_rank = shard_rank % q_split_degree
split_granularity = (n_heads_q // n_heads_kv) * head_size
q_endpoints = get_shard_endpoints(split_granularity, q_split_rank, q_split_degree, granularity=head_size)
return q_endpoints[1] - q_endpoints[0]
deepspeed/inference/v2/model_implementations/sharding/embedding.py
+34
View File
@@ -0,0 +1,34 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import torch
from .types import ShardingType
from .utils import shard_param, get_shard_endpoints
def shard_embedding_param(param: torch.Tensor, shard_rank: int, num_shards: int) -> torch.Tensor:
"""
Utility method for sharding an embedding parameter.
Args:
param (torch.Tensor): The parameter to shard. Should be of shape [vocab_size, model_dim]
shard_rank (int): Which shard of the partitioned tensor to return.
num_shards (int): The total number of shards the parameter is distributed across.
"""
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards)
def sharded_embedding_dim(embedding_size: int, shard_rank: int, num_shards: int) -> int:
"""
Utility method for getting the size of the embedding dimension of a sharded embedding.
Args:
embedding_size (int): The size of the embedding.
shard_rank (int): Which shard of the partitioned tensor to return.
num_shards (int): The total number of shards the parameter is distributed across.
"""
start_idx, end_idx = get_shard_endpoints(embedding_size, shard_rank, num_shards)
return end_idx - start_idx
deepspeed/inference/v2/model_implementations/sharding/mlp.py
+75
View File
@@ -0,0 +1,75 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
from typing import Optional
import torch
from .types import ShardingType, DEFAULT_SHARD_GRANULARITY
from .utils import shard_param, get_shard_endpoints
def shard_mlp_1_param(param: torch.Tensor,
shard_rank: int,
num_shards: int,
gated: bool = False,
is_moe: bool = False) -> torch.Tensor:
"""
Utility method for sharding an MLP 1 parameter. Both biases and weights are supported, as well
as for fused weights for MoE.
Args:
param (torch.Tensor): The parameter to shard.
shard_rank (int): Which shard of the partitioned tensor to return.
num_shards (int): The total number of shards the parameter is distributed across.
gated (bool): Whether or not the parameter is from a gated MLP.
"""
bias_dims = 2 if is_moe else 1
if gated:
return shard_param(param,
ShardingType.OUTER_DIMENSION,
shard_rank,
num_shards,
granularity=DEFAULT_SHARD_GRANULARITY * 2,
bias_dims=bias_dims)
else:
return shard_param(param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, bias_dims=bias_dims)
def shard_mlp_2_param(param: torch.Tensor,
shard_rank: int,
num_shards: int,
is_moe: bool = False) -> Optional[torch.Tensor]:
"""
Utility method for sharding an MLP 2 parameter.
Args:
param (torch.Tensor): The parameter to shard.
shard_rank (int): Which shard of the partitioned tensor to return.
num_shards (int): The total number of shards the parameter is distributed across.
is_moe (bool): Whether or not the parameter is from a MoE model.
"""
bias_dim_size = 2 if is_moe else 1
if len(param.shape) == bias_dim_size:
# We will do the bias addition on the 0th rank only rather than scale the parameter and
# implicitly reconstruct this in the distributed reduce.
return param if shard_rank == 0 else None
return shard_param(param, ShardingType.INNER_DIMENSION, shard_rank, num_shards)
def sharded_intermediate_dim(intermediate_size: int, num_shards: int, shard_rank: int) -> int:
"""
Utility method for getting the size of the intermediate dimension of a sharded MLP.
Args:
intermediate_size (int): The size of the intermediate dimension.
num_shards (int): The total number of shards the parameter is distributed across.
shard_rank (int): Which shard of the partitioned tensor to return.
"""
endpoints = get_shard_endpoints(intermediate_size, shard_rank, num_shards)
return endpoints[1] - endpoints[0]
deepspeed/inference/v2/model_implementations/sharding/qkv.py
+167
View File
@@ -0,0 +1,167 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
from typing import Optional
import torch
from .types import ShardingType
from .utils import shard_param, get_shard_endpoints
def shard_qkv_param(param: torch.Tensor,
shard_rank: int,
num_shards: int,
head_size: int,
n_heads_q: Optional[int] = None,
n_heads_kv: Optional[int] = None) -> Optional[torch.Tensor]:
"""
Utility method for sharding a QKV parameter. Both biases and weights are supported. It is assumed
that the layout of the parameter is such that all Q heads, all K heads, and all V heads
are contiguous with respect to each other.
Args:
param (torch.Tensor): The parameter to shard.
shard_rank (int): Which shard of the partitioned tensor to return.
num_shards (int): The total number of shards the parameter is distributed across.
head_size (int): The size of each head.
n_heads_q (int): The number of query heads. This only needs to be passed if the number
of query and key/value heads are different. If passed without n_heads_kv, default
MHA partitioning will be used.
n_heads_kv (int): The number of key/value heads. This only needs to be passed if the number
of query and key/value heads are different. This argument should not be passed without
n_heads_q (we want to explicitly opt into GQA sharding).
"""
if n_heads_kv is not None and n_heads_q is None:
raise ValueError("n_heads_kv should not be passed without n_heads_q")
if param is None:
raise ValueError("param should not be None")
if n_heads_q is None:
# Guaranteed to be in MHA
if param.shape[0] // 3 % head_size != 0:
raise ValueError("MHA param shape is not correct")
n_heads_q = param.shape[0] // head_size // 3
mha_sharding = True
elif n_heads_kv is None:
mha_sharding = True
else:
mha_sharding = n_heads_q == n_heads_kv
if n_heads_q < num_shards:
raise ValueError("There must be at least as many query heads as there are shards.")
if mha_sharding:
return shard_param(param,
ShardingType.OUTER_DIMENSION,
shard_rank,
num_shards,
num_concatenated_matrices=3,
granularity=head_size)
else:
if n_heads_q % n_heads_kv != 0:
raise ValueError("Must be an even ratio between query and key/value heads.")
if param.shape[0] != head_size * (n_heads_q + 2 * n_heads_kv):
raise ValueError("GQA param shape is not correct")
# 32 KV heads, 16 shards for example
if n_heads_kv >= num_shards and n_heads_kv % num_shards != 0:
raise ValueError("Currently do not support uneven partitioning of KV heads for GQA.")
# 8 KV heads, 16 shards for example
if n_heads_kv < num_shards and num_shards % n_heads_kv != 0:
raise ValueError("Currently do not support distributing KV heads across different numbers of shards.")
else:
even_kv_sharding = n_heads_kv >= num_shards
q_param = param[:head_size * n_heads_q]
kv_param = param[head_size * n_heads_q:]
if even_kv_sharding:
# This is equivalent to the original sharding algorithm since n_heads_q = C * n_heads_kv.
# If n_heads_kv % num_shards == 0, then n_heads_q % num_shards == 0.
q_param = shard_param(q_param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, granularity=head_size)
kv_param = shard_param(kv_param,
ShardingType.OUTER_DIMENSION,
shard_rank,
num_shards,
num_concatenated_matrices=2,
granularity=head_size)
return torch.cat([q_param, kv_param], dim=0)
else:
# We will first do a sharding on the KV and Q to map to the one KV shard per group of Q.
q_sharding_degree = num_shards // n_heads_kv
kv_head = shard_rank // q_sharding_degree
k_param = kv_param[kv_head * head_size:(kv_head + 1) * head_size]
v_param = kv_param[(n_heads_kv + kv_head) * head_size:(n_heads_kv + kv_head + 1) * head_size]
q_sharding_rank = shard_rank % q_sharding_degree
q_factor = n_heads_q // n_heads_kv
q_chunk = q_param[q_factor * kv_head * head_size:q_factor * (kv_head + 1) * head_size]
q_param = shard_param(q_chunk,
ShardingType.OUTER_DIMENSION,
q_sharding_rank,
q_sharding_degree,
granularity=head_size)
return torch.cat([q_param, k_param, v_param], dim=0)
def qkv_out_features(in_features: int,
shard_rank: int,
num_shards: int,
head_size: int,
n_heads_q: Optional[int] = None,
n_heads_kv: Optional[int] = None) -> int:
"""
Helper to calculate the expected output projection dimension of a QKV projection matrix.
Args:
in_features (int): The model dimension.
shard_rank (int): Which rank to return the corresponding size for.
num_shards (int): The total number of shards the parameter is distributed across.
head_size (int): The size of each head.
n_heads_q (int): The number of query heads. This only needs to be passed if the number
of query and key/value heads are different. If passed without n_heads_kv, default
MHA partitioning will be used.
n_heads_kv (int): The number of key/value heads. This only needs to be passed if the number
of query and key/value heads are different. This argument cannot be passed without also
passing n_heads_q (we want to explicitly opt into GQA sharding).
"""
if n_heads_kv is not None and n_heads_q is None:
raise ValueError("n_heads_kv should not be passed without n_heads_q")
mha_sharding = n_heads_kv is None or n_heads_q == n_heads_kv
if n_heads_q is not None and in_features != head_size * n_heads_q:
raise ValueError("in_features is not consistent with n_heads_q and head_size")
if mha_sharding:
endpoints = get_shard_endpoints(in_features, shard_rank, num_shards, granularity=head_size)
return (endpoints[1] - endpoints[0]) * 3
else:
if n_heads_kv >= num_shards:
if n_heads_kv % num_shards != 0:
raise ValueError("The KV heads must be evenly distributed across the shards.")
n_local_groups = n_heads_kv // num_shards
group_size = n_heads_q // n_heads_kv
return n_local_groups * head_size * (2 + group_size)
else:
if num_shards % n_heads_kv != 0:
raise ValueError("A shared KV head must always partition across the same number of shards.")
q_split_degree = num_shards // n_heads_kv
q_split_rank = shard_rank % q_split_degree
split_granularity = (n_heads_q // n_heads_kv) * head_size
q_endpoints = get_shard_endpoints(split_granularity, q_split_rank, q_split_degree, granularity=head_size)
return (q_endpoints[1] - q_endpoints[0]) + 2 * head_size
deepspeed/inference/v2/model_implementations/sharding/types.py
+18
View File
@@ -0,0 +1,18 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
from enum import Enum
DEFAULT_SHARD_GRANULARITY = 32
class ShardingType(Enum):
# Inner dimension sharding corresponds to splitting the Tensor along the K-dimension
# of a matrix multiplication. This would be used for attention_output or MLP2.
INNER_DIMENSION = 1
# Outer dimension sharding corresponds to splitting the Tensor along the N-dimension
# of a matrix multiplication. This would be used for the QKV and MLP1 projections.
OUTER_DIMENSION = 0
deepspeed/inference/v2/model_implementations/sharding/unembed.py
+41
View File
@@ -0,0 +1,41 @@
# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0
# DeepSpeed Team
import torch
from .types import ShardingType
from .utils import shard_param, get_shard_endpoints
def shard_unembed_param(param: torch.Tensor, shard_rank: int, num_shards: int) -> torch.Tensor:
"""
Utility method for sharding an unembed parameter. We shard unembeddings on the vocab dimension
with the expectation of an all-gather to produce the full results.
TODO(cmikeh2): Really ideal would be if MII could have access to the comm and we would do
an A2A and sharded sampling.
Args:
param (torch.Tensor): The parameter to shard. Should be of shape [vocab_size, model_dim]
shard_rank (int): Which shard of the partitioned tensor to return.
num_shards (int): The total number of shards the parameter is distributed across.
Returns:
torch.Tensor: The sharded parameter of shape [sharded_vocab_size, model_dim]
"""
return shard_param(param, ShardingType.OUTER_DIMENSION, shard_rank, num_shards, granularity=1)
def sharded_unembed_dim(vocab_size: int, shard_rank: int, num_shards: int) -> int:
"""
Utility method for determining the sharded vocab size of a sharded unembed parameter.
Args:
vocab_size (int): The size of the vocabulary.
shard_rank (int): Which shard of the partitioned tensor to return.
num_shards (int): The total number of shards the parameter is distributed across.
"""
start_idx, end_idx = get_shard_endpoints(vocab_size, shard_rank, num_shards, granularity=1)
return end_idx - start_idx
deepspeed/inference/v2/model_implementations/sharding/utils.py
+104
View File
@@ -0,0 +1,104 @@
# 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)