deepspeedai--deepspeed/deepspeed/inference/v2/model_implementations/sharding/qkv.py

# 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