softmax_context.py

# Copyright (c) Microsoft Corporation.
# SPDX-License-Identifier: Apache-2.0

# DeepSpeed Team

import torch
from deepspeed import comm as dist
from ..config import DeepSpeedInferenceConfig
from .base import BaseOp
from .softmax import SoftmaxOp
from deepspeed.ops.transformer.inference.op_binding.workspace import InferenceContext


class SoftmaxContextOp(BaseOp):

    def __init__(self, config: DeepSpeedInferenceConfig):
        super(SoftmaxContextOp, self).__init__(config)
        try:
            if self.config.dtype in [torch.float16, torch.int8]:
                self.softmax_context_func = self.inference_module.softmax_context_fp16
            elif self.config.dtype == torch.bfloat16:
                self.softmax_context_func = self.inference_module.softmax_context_bf16
            else:
                self.softmax_context_func = self.inference_module.softmax_context_fp32
        except AttributeError:
            self.softmax_context_func = self.softmax_context_fallback

    @staticmethod
    def transform4d_0213(x, seq_length):
        assert x.dim() == 3, F"Dim {x.dim()} is not supported"
        batch_size, num_heads, seq_length_head_dim = x.shape
        head_dim = seq_length_head_dim // seq_length
        x = x.view(batch_size, num_heads, seq_length, head_dim)
        x = x.permute(0, 2, 1, 3)

        return x

    @staticmethod
    def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
        batch, num_key_value_heads, slen, head_dim = hidden_states.shape
        if n_rep <= 1 or num_key_value_heads == 1:
            return hidden_states

        hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)

        return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

    @staticmethod
    def bias_add_transform_0213(input, bias, num_heads, trans_count, perform_bias=False):
        assert trans_count == 1 or trans_count == 3, F"Trans count {trans_count} is not supported"
        assert input.dim() == 3, F"Dim {input.dim()} is not supported"
        input_biased = torch.add(input, bias) if perform_bias else input
        batch_size, seq_length, value_size = input_biased.shape
        hid_dim = value_size // trans_count
        head_dim = hid_dim // num_heads

        if trans_count == 1:
            query_layer = input.view(batch_size, seq_length, num_heads, head_dim)
            query_layer = query_layer.permute(0, 2, 1, 3)
            key_layer = torch.zeros_like(query_layer)
            value_layer = torch.zeros_like(query_layer)
            return query_layer, key_layer, value_layer

        qkv_layers = input.view(batch_size, seq_length, 3, num_heads, head_dim)
        query_layer, key_layer, value_layer = qkv_layers[..., 0, :, :], qkv_layers[..., 1, :, :], qkv_layers[...,
                                                                                                             2, :, :]
        query_layer = query_layer.transpose(1, 2)
        key_layer = key_layer.transpose(1, 2)
        value_layer = value_layer.transpose(1, 2)

        return query_layer, key_layer, value_layer

    def softmax_context_fallback(self, query_key_value, attn_mask, rotary_dim, rotate_half, rotate_every_two, heads,
                                 num_kv, norm_factor, triangular_masking, local_attention, window_size, no_masking,
                                 layer_id, num_layers, alibi, rope_theta, is_prompt, token_idx, position_ids):
        bat_0213_query, bat_0213_key, bat_0213_value = self.bias_add_transform_0213(
            query_key_value, None, heads, 3, False)

        if rotary_dim > 0 and rotate_half:
            from transformers.models.llama.modeling_llama import apply_rotary_pos_emb

            rotary = InferenceContext.Instance().get_rotary(rotary_dim, rope_theta, bat_0213_value.device)
            # LlamaRotaryEmbedding.forward takes position_ids, not a token count.
            cos, sin = rotary(bat_0213_value, position_ids)
            # apply_rotary_pos_emb takes them only through cos/sin. transformers 5.0
            # dropped the deprecated position_ids parameter, so the fifth positional
            # slot is unsqueeze_dim there:
            #   4.51.3 .. 4.57.0   (q, k, cos, sin, position_ids=None, unsqueeze_dim=1)
            #   5.0.0  .. 5.16.1   (q, k, cos, sin, unsqueeze_dim=1)
            # Passing four arguments is correct on both, and leaves unsqueeze_dim at its
            # default rather than handing it a tensor.
            bat_0213_query, bat_0213_key = apply_rotary_pos_emb(bat_0213_query, bat_0213_key, cos, sin)

        bat_0213_key, bat_0213_value = InferenceContext.Instance().update_cache(layer_id, token_idx, is_prompt,
                                                                                bat_0213_key, bat_0213_value)

        bat_0213_key = self.repeat_kv(bat_0213_key, num_kv)
        bat_0213_value = self.repeat_kv(bat_0213_value, num_kv)

        bsz = query_key_value.shape[0]
        head_dim = query_key_value.shape[2] // (heads * 3)

        bmm_output = torch.bmm(bat_0213_query.reshape(bsz * heads, bat_0213_query.shape[2], head_dim),
                               bat_0213_key.reshape(bsz * heads, bat_0213_key.shape[2], head_dim).transpose(1, 2))

        layer_scale = 1.0
        if alibi is not None and len(alibi.shape) > 1:
            layer_scale = max(1, layer_id).to(float)

        alpha = norm_factor * norm_factor / layer_scale
        bmm_output *= alpha
        bmm_output_reshape = bmm_output.reshape(bsz, heads, bmm_output.shape[1], bmm_output.shape[2])

        recompute = is_prompt
        if attn_mask is not None and len(attn_mask.shape) > 1 and attn_mask.shape[-1] < bmm_output_reshape.shape[3]:
            attn_mask = torch.nn.functional.pad(attn_mask, (0, bmm_output_reshape.shape[3] - attn_mask.shape[-1]),
                                                value=torch.finfo(attn_mask.dtype).min)
        softmax_output = SoftmaxOp.softmax_fallback(bmm_output_reshape, attn_mask, alibi, triangular_masking,
                                                    recompute, local_attention, window_size, None, layer_scale, 0, 1)

        output = torch.bmm(softmax_output.reshape(bsz * heads, softmax_output.shape[2], softmax_output.shape[3]),
                           bat_0213_value.reshape(bsz * heads, bat_0213_value.shape[2], head_dim))

        output = output.reshape(bsz, heads, output.shape[1], head_dim)
        output = output.reshape(bsz, heads, output.shape[2] * head_dim)
        input_seq_len = query_key_value.shape[1]
        t4d_0123_output = self.transform4d_0213(output, input_seq_len)
        t4d_0123_output = t4d_0123_output.reshape(bsz, t4d_0123_output.shape[1], heads * head_dim)

        if layer_id == num_layers - 1:
            InferenceContext.Instance().advance_tokens()

        return t4d_0123_output, bat_0213_key, bat_0213_value

    def forward(self, query_key_value: torch.Tensor, attn_mask: torch.Tensor, heads: int, num_kv: int,
                norm_factor: float, no_masking: bool, layer_id: int, num_layers: int, alibi: torch.Tensor,
                is_prompt: bool, token_idx: torch.Tensor, position_ids: torch.Tensor):

        if alibi is not None:
            batch_heads = query_key_value.shape[0] * heads
            offset = dist.get_rank() * batch_heads if dist.is_initialized() else 0
            alibi = alibi[offset:batch_heads + offset, :, :]
        else:
            alibi = torch.empty(1)

        output = self.softmax_context_func(query_key_value, attn_mask, self.config.rotary_dim, self.config.rotate_half,
                                           self.config.rotate_every_two, heads, num_kv, norm_factor,
                                           self.config.triangular_masking, self.config.local_attention,
                                           self.config.window_size, no_masking, layer_id, num_layers, alibi,
                                           self.config.rope_theta, is_prompt, token_idx, position_ids)

        return output