simplecloud/VidChain-exercise

1from typing import List, Optional, Tuple2import logging3 4import torch5from torch import nn6 7import transformers8from transformers.models.llama.modeling_llama import apply_rotary_pos_emb9 10from einops import rearrange11 12try:13    from flash_attn.flash_attn_interface import flash_attn_unpadded_qkvpacked_func14except ImportError:15    from flash_attn.flash_attn_interface import flash_attn_varlen_qkvpacked_func as flash_attn_unpadded_qkvpacked_func16from flash_attn.bert_padding import unpad_input, pad_input17 18 19def forward(20    self,21    hidden_states: torch.Tensor,22    attention_mask: Optional[torch.Tensor] = None,23    position_ids: Optional[torch.Tensor] = None,24    past_key_value: Optional[Tuple[torch.Tensor]] = None,25    output_attentions: bool = False,26    use_cache: bool = False,27) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:28    """Input shape: Batch x Time x Channel29 30    attention_mask: [bsz, q_len]31    """32    bsz, q_len, _ = hidden_states.size()33 34    query_states = (35        self.q_proj(hidden_states)36        .view(bsz, q_len, self.num_heads, self.head_dim)37        .transpose(1, 2)38    )39    key_states = (40        self.k_proj(hidden_states)41        .view(bsz, q_len, self.num_heads, self.head_dim)42        .transpose(1, 2)43    )44    value_states = (45        self.v_proj(hidden_states)46        .view(bsz, q_len, self.num_heads, self.head_dim)47        .transpose(1, 2)48    )49    # [bsz, q_len, nh, hd]50    # [bsz, nh, q_len, hd]51 52    kv_seq_len = key_states.shape[-2]53    assert past_key_value is None, "past_key_value is not supported"54 55    cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)56    query_states, key_states = apply_rotary_pos_emb(57        query_states, key_states, cos, sin, position_ids58    )59    # [bsz, nh, t, hd]60    assert not output_attentions, "output_attentions is not supported"61    assert not use_cache, "use_cache is not supported"62 63    # Flash attention codes from64    # https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/flash_attention.py65 66    # transform the data into the format required by flash attention67    qkv = torch.stack(68        [query_states, key_states, value_states], dim=269    )  # [bsz, nh, 3, q_len, hd]70    qkv = qkv.transpose(1, 3)  # [bsz, q_len, 3, nh, hd]71    # We have disabled _prepare_decoder_attention_mask in LlamaModel72    # the attention_mask should be the same as the key_padding_mask73    key_padding_mask = attention_mask74 75    if key_padding_mask is None:76        qkv = rearrange(qkv, "b s ... -> (b s) ...")77        max_s = q_len78        cu_q_lens = torch.arange(79            0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=qkv.device80        )81        output = flash_attn_unpadded_qkvpacked_func(82            qkv, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True83        )84        output = rearrange(output, "(b s) ... -> b s ...", b=bsz)85    else:86        nheads = qkv.shape[-2]87        x = rearrange(qkv, "b s three h d -> b s (three h d)")88        x_unpad, indices, cu_q_lens, max_s = unpad_input(x, key_padding_mask)89        x_unpad = rearrange(90            x_unpad, "nnz (three h d) -> nnz three h d", three=3, h=nheads91        )92        output_unpad = flash_attn_unpadded_qkvpacked_func(93            x_unpad, cu_q_lens, max_s, 0.0, softmax_scale=None, causal=True94        )95        output = rearrange(96            pad_input(97                rearrange(output_unpad, "nnz h d -> nnz (h d)"), indices, bsz, q_len98            ),99            "b s (h d) -> b s h d",100            h=nheads,101        )102    return self.o_proj(rearrange(output, "b s h d -> b s (h d)")), None, None103 104 105# Disable the transformation of the attention mask in LlamaModel as the flash attention106# requires the attention mask to be the same as the key_padding_mask107def _prepare_decoder_attention_mask(108    self, attention_mask, input_shape, inputs_embeds, past_key_values_length109):110    # [bsz, seq_len]111    return attention_mask112 113 114def replace_llama_attn_with_flash_attn():115    cuda_major, cuda_minor = torch.cuda.get_device_capability()116    if cuda_major < 8:117        logging.warning(118            "Flash attention is only supported on A100 or H100 GPU during training due to head dim > 64 backward."119            "ref: https://github.com/HazyResearch/flash-attention/issues/190#issuecomment-1523359593"120        )121    transformers.models.llama.modeling_llama.LlamaModel._prepare_decoder_attention_mask = (122        _prepare_decoder_attention_mask123    )124    transformers.models.llama.modeling_llama.LlamaAttention.forward = forward125