attention.py

import warnings
import torch

from typing import Optional, Tuple
from transformers.cache_utils import Cache, StaticCache, EncoderDecoderCache
try:
    from transformers.models.whisper.modeling_whisper import WhisperFlashAttention2, WHISPER_ATTENTION_CLASSES
except ImportError:
    WhisperFlashAttention2 = object
    WHISPER_ATTENTION_CLASSES = {}
    warnings.warn("WhisperFlashAttention2 could not be imported from transformers.models.whisper.modeling_whisper. "
                  "Make sure you have transformers<4.53.0 installed when using HiggsAudio models. This might cause errors "
                  "if you use HiggsAudio in this environment.", ImportWarning)
from transformers.utils import logging
from transformers.modeling_flash_attention_utils import _flash_attention_forward
from transformers.integrations import is_deepspeed_available


from .utils import support_deepspeed_ulysses, is_deepspeed_ulysses_enabled, deepspeed_ulysses_attention, deepspeed_ulysses_rope, sequence_chunking_per_rank

DistributedAttention = None
vocab_sequence_parallel_cross_entropy = None
if is_deepspeed_available():
    import importlib
    DistributedAttention = importlib.import_module("deepspeed.sequence.layer").DistributedAttention
    vocab_sequence_parallel_cross_entropy = importlib.import_module("deepspeed.sequence.cross_entropy").vocab_sequence_parallel_cross_entropy


logger = logging.get_logger(__name__)



@deepspeed_ulysses_attention(seq_dim=1, head_dim=2)
def _distributed_flash_attention_forward(*args, **kwargs):
    return _flash_attention_forward(*args, **kwargs)


def _patched_flash_attention_forward(
    query, key, value, *args, **kwargs,
) -> torch.Tensor:
    # IMPORTANT! Implementation here is wrong and is only for the purpose of obtaining the correct attn_weight shape
    if query.size(-2) != key.size(-2):
        key = key.repeat_interleave(query.size(-2) // key.size(-2), -2)
        value = value.repeat_interleave(query.size(-2) // value.size(-2), -2)

    return (query + key + value) / 3.0


def _distributed_higgs_flash_attention_forward(*args, **kwargs):
    if args[0].size(0) == 0 or args[1].size(1) == 0:
        return _patched_flash_attention_forward(*args, **kwargs)
    else:
        return _distributed_flash_attention_forward(*args, **kwargs)

@support_deepspeed_ulysses
class HiggsDistributedWhisperFlashAttention2(WhisperFlashAttention2):
    """
    Higgs Whisper flash attention module. This module inherits from `WhisperFlashAttention2` as the weights of the module stays
    untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
    flash attention and deal with padding tokens in case the input contains any of them.
    """
    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        """Remove the redundant transpose and apply the monkey-patch for zero-shape input."""
        if seq_len == -1:
            return tensor.view(bsz, tensor.shape[1], self.num_heads, self.head_dim)
        else:
            return tensor.view(bsz, seq_len, self.num_heads, self.head_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        key_value_states: Optional[torch.Tensor] = None,
        past_key_value: Optional[EncoderDecoderCache] = None,
        attention_mask: Optional[torch.Tensor] = None,
        layer_head_mask: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        if isinstance(past_key_value, StaticCache):
            raise ValueError(
                "The `static` cache implementation is not compatible with `attn_implementation='flash_attention_2'`. "
                "Use `attn_implementation='sdpa'` in the meantime, and open an issue at https://github.com/huggingface/transformers"
            )
        # WhisperFlashAttention2 attention does not support output_attentions
        if output_attentions:
            raise ValueError("HiggsWhisperFlashAttention2 attention does not support output_attentions")

        # if key_value_states are provided this layer is used as a cross-attention layer
        # for the decoder
        is_cross_attention = key_value_states is not None
        bsz, tgt_len, _ = hidden_states.size()

        # get query proj
        query_states = torch.reshape(self.q_proj(hidden_states), (bsz, tgt_len, self.num_heads, self.head_dim))

        if past_key_value is not None:
            is_updated = past_key_value.is_updated.get(self.layer_idx)
            if is_cross_attention:
                # after the first generated id, we can subsequently re-use all key/value_states from cache
                past_key_value.is_updated[self.layer_idx] = True
                past_key_value = past_key_value.cross_attention_cache
            else:
                past_key_value = past_key_value.self_attention_cache

        # use key_value_states if cross attention
        current_states = key_value_states if key_value_states is not None else hidden_states
        if is_cross_attention and past_key_value and is_updated:
            # reuse k,v, cross_attentions
            key_states = past_key_value.key_cache[self.layer_idx]
            value_states = past_key_value.value_cache[self.layer_idx]
        else:
            key_states = self._shape(self.k_proj(current_states), -1, bsz)
            value_states = self._shape(self.v_proj(current_states), -1, bsz)
            if past_key_value is not None:
                # save all key/value_states to cache to be re-used for fast auto-regressive generation
                cache_position = cache_position if not is_cross_attention else None
                key_states, value_states = past_key_value.update(
                    key_states.transpose(1, 2), value_states.transpose(1, 2), self.layer_idx, {"cache_position": cache_position}
                )

        # Cache layout [batch_size, num_heads, sequence_length, head_dim]
        if past_key_value is not None:
            key_states = key_states.transpose(1, 2)
            value_states = value_states.transpose(1, 2)

        causal_mask = attention_mask
        if attention_mask is not None:  # no matter the length, we just slice it
            if is_deepspeed_ulysses_enabled():
                causal_mask = attention_mask[:, : key_states.shape[-3] * getattr(self, "sp_size", 1)]
            else:
                causal_mask = attention_mask[:, : key_states.shape[-3]]

        # In PEFT, usually we cast the layer norms in float32 for training stability reasons
        # therefore the input hidden states gets silently casted in float32. Hence, we need
        # cast them back in the correct dtype just to be sure everything works as expected.
        # This might slowdown training & inference so it is recommended to not cast the LayerNorms
        # in fp32. (LlamaRMSNorm handles it correctly)

        input_dtype = query_states.dtype
        if input_dtype == torch.float32:
            if torch.is_autocast_enabled():
                target_dtype = torch.get_autocast_gpu_dtype()
            # Handle the case where the model is quantized
            elif hasattr(self.config, "_pre_quantization_dtype"):
                target_dtype = self.config._pre_quantization_dtype
            else:
                target_dtype = self.q_proj.weight.dtype

            logger.warning_once(
                f"The input hidden states seems to be silently casted in float32, this might be related to"
                f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
                f" {target_dtype}."
            )

            query_states = query_states.to(target_dtype)
            key_states = key_states.to(target_dtype)
            value_states = value_states.to(target_dtype)

        attn_output = _distributed_higgs_flash_attention_forward(
            query_states,
            key_states,
            value_states,
            causal_mask,
            tgt_len * getattr(self, "sp_size", 1) if is_deepspeed_ulysses_enabled() else tgt_len,
            dropout=self.dropout if self.training else 0.0,
            is_causal=self.is_causal,
            use_top_left_mask=self._flash_attn_uses_top_left_mask,
        )

        attn_output = attn_output.reshape(bsz, tgt_len, hidden_states.size(2))
        attn_output = self.out_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, past_key_value

WHISPER_ATTENTION_CLASSES["flash_attention_2"] = HiggsDistributedWhisperFlashAttention2