wehub-resource-sync
112 lines
4.8 KiB
Python
112 lines
4.8 KiB
Python
# 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]
|