We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[tensor] fix some unittests

[tensor] fix some unittests

[tensor] fix some unittests


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py`
Content:
```
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 from types import MethodType
5 from typing import Optional, Tuple
6 
7 import torch
8 import torch.nn.functional as F
9 from transformers.models.llama.modeling_llama import (
10     LlamaRMSNorm,
11     LlamaAttention,
12     LlamaModel,
13     LlamaForCausalLM,
14     apply_rotary_pos_emb,
15     repeat_kv,
16 )
17 
18 from colossalai.logging import get_dist_logger
19 from einops import rearrange
20 
21 from flash_attn.bert_padding import pad_input, unpad_input
22 from flash_attn.flash_attn_interface import (
23     flash_attn_func,
24     flash_attn_varlen_kvpacked_func,
25 )
26 from flash_attn.ops.rms_norm import rms_norm
27 
28 
29 logger = get_dist_logger()
30 
31 
32 def _prepare_decoder_attention_mask(
33     self: LlamaModel,
34     attention_mask: torch.BoolTensor,
35     input_shape: torch.Size,
36     inputs_embeds: torch.Tensor,
37     past_key_values_length: int,
38 ) -> Optional[torch.Tensor]:
39     """
40     Decoder attetion mask
41     """
42     if past_key_values_length > 0 and attention_mask is not None:
43         attention_mask = torch.cat(
44             tensors=(
45                 torch.full(
46                     size=(input_shape[0], past_key_values_length),
47                     fill_value=True,
48                     dtype=attention_mask.dtype,
49                     device=attention_mask.device,
50                 ),
51                 attention_mask,
52             ),
53             dim=-1,
54         )  # (bsz, past_key_values_length + q_len)
55     if attention_mask is not None and torch.all(attention_mask):
56         return None  # Faster
57     return attention_mask
58 
59 
60 def attention_forward(
61     self: LlamaAttention,
62     hidden_states: torch.Tensor,
63     attention_mask: Optional[torch.Tensor] = None,
64     position_ids: Optional[torch.LongTensor] = None,
65     past_key_value: Optional[Tuple[torch.Tensor]] = None,
66     output_attentions: bool = False,
67     use_cache: bool = False,
68 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
69     """
70     Re-define LLaMA-2 `LlamaAttention` forward method using flash-attention.
71     """
72     if output_attentions:
73         logger.warning(
74             "Argument `output_attentions` is not supported for flash-attention patched `LlamaAttention`, "
75             "return `None` instead."
76         )
77 
78     bsz, q_len, _ = hidden_states.size()
79 
80     if self.config.pretraining_tp > 1:
81         q_slicing, kv_slicing = (
82             dim // self.config.pretraining_tp
83             for dim in (
84                 self.num_heads * self.head_dim,
85                 self.num_key_value_heads * self.head_dim,
86             )
87         )  # `Tuple[int, int]`
88         q_slices, k_slices, v_slices = (
89             proj.weight.split(slicing, dim=0)
90             for proj, slicing in (
91                 (self.q_proj, q_slicing),
92                 (self.k_proj, kv_slicing),
93                 (self.v_proj, kv_slicing),
94             )
95         )  # Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor], Tuple[torch.Tensor]]
96         q, k, v = (
97             torch.cat(
98                 [F.linear(hidden_states, slices[i]) for i in range(self.config.pretraining_tp)],
99                 dim=-1,
100             )
101             for slices in (q_slices, k_slices, v_slices)
102         )
103         # `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]` of shape:
104         # (bsz, q_len, num_heads * head_dim),
105         # (bsz, q_len, num_key_value_heads * head_dim),
106         # (bsz, q_len, num_key_value_heads * head_dim)
107     else:
108         q, k, v = (proj(hidden_states) for proj in (self.q_proj, self.k_proj, self.v_proj))
109         # `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]` of shape:
110         # (bsz, q_len, num_heads * head_dim),
111         # (bsz, q_len, num_key_value_heads * head_dim),
112         # (bsz, q_len, num_key_value_heads * head_dim)
113 
114     # (bsz, q_len, num_heads * head_dim) -> (bsz, num_heads, q_len, head_dim);
115     # (bsz, q_len, num_key_value_heads * head_dim) -> (bsz, num_key_value_heads, q_len, head_dim);
116     # (bsz, q_len, num_key_value_heads * head_dim) -> (bsz, num_key_value_heads, q_len, head_dim)
117     q, k, v = (
118         states.view(bsz, q_len, num_heads, self.head_dim).transpose(1, 2)
119         for states, num_heads in (
120             (q, self.num_heads),
121             (k, self.num_key_value_heads),
122             (v, self.num_key_value_heads),
123         )
124     )
125     kv_len = k.shape[-2]  # initially, `kv_len` == `q_len`
126     past_kv_len = 0
127     if past_key_value is not None:
128         # if `past_key_value` is not None, `kv_len` > `q_len`.
129         past_kv_len = past_key_value[0].shape[-2]
130         kv_len += past_kv_len
131 
132     # two `torch.Tensor` objs of shape (1, 1, kv_len, head_dim)
133     cos, sin = self.rotary_emb(v, seq_len=kv_len)
134     # (bsz, num_heads, q_len, head_dim), (bsz, num_key_value_heads, q_len, head_dim)
135     q, k = apply_rotary_pos_emb(q=q, k=k, cos=cos, sin=sin, position_ids=position_ids)
136     if past_key_value is not None:
137         # reuse k, v, self_attention
138         k = torch.cat([past_key_value[0], k], dim=2)
139         v = torch.cat([past_key_value[1], v], dim=2)
140 
141     past_key_value = (k, v) if use_cache else None
142 
143     # repeat k/v heads if n_kv_heads < n_heads
144     k = repeat_kv(hidden_states=k, n_rep=self.num_key_value_groups)
145     # (bsz, num_key_value_heads, q_len, head_dim) -> (bsz, num_heads, q_len, head_dim)
146     v = repeat_kv(hidden_states=v, n_rep=self.num_key_value_groups)
147     # (bsz, num_key_value_heads, q_len, head_dim) -> (bsz, num_heads, q_len, head_dim)
148 
149     key_padding_mask = attention_mask
150     # (bsz, num_heads, q_len, head_dim) -> (bsz, q_len, num_heads, head_dim)
151     q, k, v = (states.transpose(1, 2) for states in (q, k, v))
152 
153     if past_kv_len > 0:
154         q = torch.cat(
155             tensors=(
156                 torch.full(
157                     size=(bsz, past_kv_len, self.num_heads, self.head_dim),
158                     fill_value=0.0,
159                     dtype=q.dtype,
160                     device=q.device,
161                 ),
162                 q,
163             ),
164             dim=1,
165         )  # (bsz, past_kv_len + q_len, num_heads, head_dim)
166 
167     if key_padding_mask is None:
168         # (bsz, past_kv_len + q_len, num_heads, head_dim)
169         output = flash_attn_func(q=q, k=k, v=v, dropout_p=0.0, softmax_scale=None, causal=True)  # (bsz, )
170         output = rearrange(output, pattern="... h d -> ... (h d)")  # (bsz, past_kv_len + q_len, num_heads * head_dim)
171     else:
172         q, indices, cu_q_lens, max_q_len = unpad_input(hidden_states=q, attention_mask=key_padding_mask)
173         kv, _, cu_kv_lens, max_kv_len = unpad_input(
174             hidden_states=torch.stack(tensors=(k, v), dim=2),
175             attention_mask=key_padding_mask,
176         )
177         output_unpad = flash_attn_varlen_kvpacked_func(
178             q=q,
179             kv=kv,
180             cu_seqlens_q=cu_q_lens,
181             cu_seqlens_k=cu_kv_lens,
182             max_seqlen_q=max_q_len,
183             max_seqlen_k=max_kv_len,
184             dropout_p=0.0,
185             softmax_scale=None,
186             causal=True,
187         )
188         output = pad_input(
189             hidden_states=rearrange(output_unpad, pattern="nnz h d -> nnz (h d)"),
190             indices=indices,
191             batch=bsz,
192             seqlen=past_kv_len + q_len,
193         )  # (bsz, past_kv_len + q_len, num_heads * head_dim)
194 
195     if past_kv_len > 0:
196         # Strip off the zero query outputs.
197         output = output[:, past_kv_len:, ...]  # (bsz, q_len, num_heads * head_dim)
198     output = self.o_proj(output)  # (bsz, q_len, hidden_size)
199     return output, None, past_key_value
200 
201 
202 def rms_norm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor) -> torch.Tensor:
203     """
204     Formard function for RMS Norm
205     """
206     return rms_norm(x=hidden_states, weight=self.weight, epsilon=self.variance_epsilon)
207 
208 
209 def replace_with_flash_attention(model: LlamaForCausalLM) -> None:
210     for name, module in model.named_modules():
211         if isinstance(module, LlamaAttention):
212             module.forward = MethodType(attention_forward, module)
213         if isinstance(module, LlamaModel):
214             module._prepare_decoder_attention_mask = MethodType(_prepare_decoder_attention_mask, module)
215         if isinstance(module, LlamaRMSNorm):
216             module.forward = MethodType(rms_norm_forward, module)
217
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py b/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py
--- a/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py
+++ b/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py
@@ -65,6 +65,7 @@
     past_key_value: Optional[Tuple[torch.Tensor]] = None,
     output_attentions: bool = False,
     use_cache: bool = False,
+    **kwargs
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
     """
     Re-define LLaMA-2 `LlamaAttention` forward method using flash-attention.

{"golden_diff": "diff --git a/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py b/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py\n--- a/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py\n+++ b/applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py\n@@ -65,6 +65,7 @@\n     past_key_value: Optional[Tuple[torch.Tensor]] = None,\n     output_attentions: bool = False,\n     use_cache: bool = False,\n+    **kwargs\n ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:\n     \"\"\"\n     Re-define LLaMA-2 `LlamaAttention` forward method using flash-attention.\n", "issue": "[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n[tensor] fix some unittests\n\n", "before_files": [{"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\nfrom types import MethodType\nfrom typing import Optional, Tuple\n\nimport torch\nimport torch.nn.functional as F\nfrom transformers.models.llama.modeling_llama import (\n    LlamaRMSNorm,\n    LlamaAttention,\n    LlamaModel,\n    LlamaForCausalLM,\n    apply_rotary_pos_emb,\n    repeat_kv,\n)\n\nfrom colossalai.logging import get_dist_logger\nfrom einops import rearrange\n\nfrom flash_attn.bert_padding import pad_input, unpad_input\nfrom flash_attn.flash_attn_interface import (\n    flash_attn_func,\n    flash_attn_varlen_kvpacked_func,\n)\nfrom flash_attn.ops.rms_norm import rms_norm\n\n\nlogger = get_dist_logger()\n\n\ndef _prepare_decoder_attention_mask(\n    self: LlamaModel,\n    attention_mask: torch.BoolTensor,\n    input_shape: torch.Size,\n    inputs_embeds: torch.Tensor,\n    past_key_values_length: int,\n) -> Optional[torch.Tensor]:\n    \"\"\"\n    Decoder attetion mask\n    \"\"\"\n    if past_key_values_length > 0 and attention_mask is not None:\n        attention_mask = torch.cat(\n            tensors=(\n                torch.full(\n                    size=(input_shape[0], past_key_values_length),\n                    fill_value=True,\n                    dtype=attention_mask.dtype,\n                    device=attention_mask.device,\n                ),\n                attention_mask,\n            ),\n            dim=-1,\n        )  # (bsz, past_key_values_length + q_len)\n    if attention_mask is not None and torch.all(attention_mask):\n        return None  # Faster\n    return attention_mask\n\n\ndef attention_forward(\n    self: LlamaAttention,\n    hidden_states: torch.Tensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    past_key_value: Optional[Tuple[torch.Tensor]] = None,\n    output_attentions: bool = False,\n    use_cache: bool = False,\n) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:\n    \"\"\"\n    Re-define LLaMA-2 `LlamaAttention` forward method using flash-attention.\n    \"\"\"\n    if output_attentions:\n        logger.warning(\n            \"Argument `output_attentions` is not supported for flash-attention patched `LlamaAttention`, \"\n            \"return `None` instead.\"\n        )\n\n    bsz, q_len, _ = hidden_states.size()\n\n    if self.config.pretraining_tp > 1:\n        q_slicing, kv_slicing = (\n            dim // self.config.pretraining_tp\n            for dim in (\n                self.num_heads * self.head_dim,\n                self.num_key_value_heads * self.head_dim,\n            )\n        )  # `Tuple[int, int]`\n        q_slices, k_slices, v_slices = (\n            proj.weight.split(slicing, dim=0)\n            for proj, slicing in (\n                (self.q_proj, q_slicing),\n                (self.k_proj, kv_slicing),\n                (self.v_proj, kv_slicing),\n            )\n        )  # Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor], Tuple[torch.Tensor]]\n        q, k, v = (\n            torch.cat(\n                [F.linear(hidden_states, slices[i]) for i in range(self.config.pretraining_tp)],\n                dim=-1,\n            )\n            for slices in (q_slices, k_slices, v_slices)\n        )\n        # `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]` of shape:\n        # (bsz, q_len, num_heads * head_dim),\n        # (bsz, q_len, num_key_value_heads * head_dim),\n        # (bsz, q_len, num_key_value_heads * head_dim)\n    else:\n        q, k, v = (proj(hidden_states) for proj in (self.q_proj, self.k_proj, self.v_proj))\n        # `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]` of shape:\n        # (bsz, q_len, num_heads * head_dim),\n        # (bsz, q_len, num_key_value_heads * head_dim),\n        # (bsz, q_len, num_key_value_heads * head_dim)\n\n    # (bsz, q_len, num_heads * head_dim) -> (bsz, num_heads, q_len, head_dim);\n    # (bsz, q_len, num_key_value_heads * head_dim) -> (bsz, num_key_value_heads, q_len, head_dim);\n    # (bsz, q_len, num_key_value_heads * head_dim) -> (bsz, num_key_value_heads, q_len, head_dim)\n    q, k, v = (\n        states.view(bsz, q_len, num_heads, self.head_dim).transpose(1, 2)\n        for states, num_heads in (\n            (q, self.num_heads),\n            (k, self.num_key_value_heads),\n            (v, self.num_key_value_heads),\n        )\n    )\n    kv_len = k.shape[-2]  # initially, `kv_len` == `q_len`\n    past_kv_len = 0\n    if past_key_value is not None:\n        # if `past_key_value` is not None, `kv_len` > `q_len`.\n        past_kv_len = past_key_value[0].shape[-2]\n        kv_len += past_kv_len\n\n    # two `torch.Tensor` objs of shape (1, 1, kv_len, head_dim)\n    cos, sin = self.rotary_emb(v, seq_len=kv_len)\n    # (bsz, num_heads, q_len, head_dim), (bsz, num_key_value_heads, q_len, head_dim)\n    q, k = apply_rotary_pos_emb(q=q, k=k, cos=cos, sin=sin, position_ids=position_ids)\n    if past_key_value is not None:\n        # reuse k, v, self_attention\n        k = torch.cat([past_key_value[0], k], dim=2)\n        v = torch.cat([past_key_value[1], v], dim=2)\n\n    past_key_value = (k, v) if use_cache else None\n\n    # repeat k/v heads if n_kv_heads < n_heads\n    k = repeat_kv(hidden_states=k, n_rep=self.num_key_value_groups)\n    # (bsz, num_key_value_heads, q_len, head_dim) -> (bsz, num_heads, q_len, head_dim)\n    v = repeat_kv(hidden_states=v, n_rep=self.num_key_value_groups)\n    # (bsz, num_key_value_heads, q_len, head_dim) -> (bsz, num_heads, q_len, head_dim)\n\n    key_padding_mask = attention_mask\n    # (bsz, num_heads, q_len, head_dim) -> (bsz, q_len, num_heads, head_dim)\n    q, k, v = (states.transpose(1, 2) for states in (q, k, v))\n\n    if past_kv_len > 0:\n        q = torch.cat(\n            tensors=(\n                torch.full(\n                    size=(bsz, past_kv_len, self.num_heads, self.head_dim),\n                    fill_value=0.0,\n                    dtype=q.dtype,\n                    device=q.device,\n                ),\n                q,\n            ),\n            dim=1,\n        )  # (bsz, past_kv_len + q_len, num_heads, head_dim)\n\n    if key_padding_mask is None:\n        # (bsz, past_kv_len + q_len, num_heads, head_dim)\n        output = flash_attn_func(q=q, k=k, v=v, dropout_p=0.0, softmax_scale=None, causal=True)  # (bsz, )\n        output = rearrange(output, pattern=\"... h d -> ... (h d)\")  # (bsz, past_kv_len + q_len, num_heads * head_dim)\n    else:\n        q, indices, cu_q_lens, max_q_len = unpad_input(hidden_states=q, attention_mask=key_padding_mask)\n        kv, _, cu_kv_lens, max_kv_len = unpad_input(\n            hidden_states=torch.stack(tensors=(k, v), dim=2),\n            attention_mask=key_padding_mask,\n        )\n        output_unpad = flash_attn_varlen_kvpacked_func(\n            q=q,\n            kv=kv,\n            cu_seqlens_q=cu_q_lens,\n            cu_seqlens_k=cu_kv_lens,\n            max_seqlen_q=max_q_len,\n            max_seqlen_k=max_kv_len,\n            dropout_p=0.0,\n            softmax_scale=None,\n            causal=True,\n        )\n        output = pad_input(\n            hidden_states=rearrange(output_unpad, pattern=\"nnz h d -> nnz (h d)\"),\n            indices=indices,\n            batch=bsz,\n            seqlen=past_kv_len + q_len,\n        )  # (bsz, past_kv_len + q_len, num_heads * head_dim)\n\n    if past_kv_len > 0:\n        # Strip off the zero query outputs.\n        output = output[:, past_kv_len:, ...]  # (bsz, q_len, num_heads * head_dim)\n    output = self.o_proj(output)  # (bsz, q_len, hidden_size)\n    return output, None, past_key_value\n\n\ndef rms_norm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor) -> torch.Tensor:\n    \"\"\"\n    Formard function for RMS Norm\n    \"\"\"\n    return rms_norm(x=hidden_states, weight=self.weight, epsilon=self.variance_epsilon)\n\n\ndef replace_with_flash_attention(model: LlamaForCausalLM) -> None:\n    for name, module in model.named_modules():\n        if isinstance(module, LlamaAttention):\n            module.forward = MethodType(attention_forward, module)\n        if isinstance(module, LlamaModel):\n            module._prepare_decoder_attention_mask = MethodType(_prepare_decoder_attention_mask, module)\n        if isinstance(module, LlamaRMSNorm):\n            module.forward = MethodType(rms_norm_forward, module)\n", "path": "applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py"}], "after_files": [{"content": "#!/usr/bin/env python3\n# -*- coding: utf-8 -*-\n\nfrom types import MethodType\nfrom typing import Optional, Tuple\n\nimport torch\nimport torch.nn.functional as F\nfrom transformers.models.llama.modeling_llama import (\n    LlamaRMSNorm,\n    LlamaAttention,\n    LlamaModel,\n    LlamaForCausalLM,\n    apply_rotary_pos_emb,\n    repeat_kv,\n)\n\nfrom colossalai.logging import get_dist_logger\nfrom einops import rearrange\n\nfrom flash_attn.bert_padding import pad_input, unpad_input\nfrom flash_attn.flash_attn_interface import (\n    flash_attn_func,\n    flash_attn_varlen_kvpacked_func,\n)\nfrom flash_attn.ops.rms_norm import rms_norm\n\n\nlogger = get_dist_logger()\n\n\ndef _prepare_decoder_attention_mask(\n    self: LlamaModel,\n    attention_mask: torch.BoolTensor,\n    input_shape: torch.Size,\n    inputs_embeds: torch.Tensor,\n    past_key_values_length: int,\n) -> Optional[torch.Tensor]:\n    \"\"\"\n    Decoder attetion mask\n    \"\"\"\n    if past_key_values_length > 0 and attention_mask is not None:\n        attention_mask = torch.cat(\n            tensors=(\n                torch.full(\n                    size=(input_shape[0], past_key_values_length),\n                    fill_value=True,\n                    dtype=attention_mask.dtype,\n                    device=attention_mask.device,\n                ),\n                attention_mask,\n            ),\n            dim=-1,\n        )  # (bsz, past_key_values_length + q_len)\n    if attention_mask is not None and torch.all(attention_mask):\n        return None  # Faster\n    return attention_mask\n\n\ndef attention_forward(\n    self: LlamaAttention,\n    hidden_states: torch.Tensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    past_key_value: Optional[Tuple[torch.Tensor]] = None,\n    output_attentions: bool = False,\n    use_cache: bool = False,\n    **kwargs\n) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:\n    \"\"\"\n    Re-define LLaMA-2 `LlamaAttention` forward method using flash-attention.\n    \"\"\"\n    if output_attentions:\n        logger.warning(\n            \"Argument `output_attentions` is not supported for flash-attention patched `LlamaAttention`, \"\n            \"return `None` instead.\"\n        )\n\n    bsz, q_len, _ = hidden_states.size()\n\n    if self.config.pretraining_tp > 1:\n        q_slicing, kv_slicing = (\n            dim // self.config.pretraining_tp\n            for dim in (\n                self.num_heads * self.head_dim,\n                self.num_key_value_heads * self.head_dim,\n            )\n        )  # `Tuple[int, int]`\n        q_slices, k_slices, v_slices = (\n            proj.weight.split(slicing, dim=0)\n            for proj, slicing in (\n                (self.q_proj, q_slicing),\n                (self.k_proj, kv_slicing),\n                (self.v_proj, kv_slicing),\n            )\n        )  # Tuple[Tuple[torch.Tensor], Tuple[torch.Tensor], Tuple[torch.Tensor]]\n        q, k, v = (\n            torch.cat(\n                [F.linear(hidden_states, slices[i]) for i in range(self.config.pretraining_tp)],\n                dim=-1,\n            )\n            for slices in (q_slices, k_slices, v_slices)\n        )\n        # `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]` of shape:\n        # (bsz, q_len, num_heads * head_dim),\n        # (bsz, q_len, num_key_value_heads * head_dim),\n        # (bsz, q_len, num_key_value_heads * head_dim)\n    else:\n        q, k, v = (proj(hidden_states) for proj in (self.q_proj, self.k_proj, self.v_proj))\n        # `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]` of shape:\n        # (bsz, q_len, num_heads * head_dim),\n        # (bsz, q_len, num_key_value_heads * head_dim),\n        # (bsz, q_len, num_key_value_heads * head_dim)\n\n    # (bsz, q_len, num_heads * head_dim) -> (bsz, num_heads, q_len, head_dim);\n    # (bsz, q_len, num_key_value_heads * head_dim) -> (bsz, num_key_value_heads, q_len, head_dim);\n    # (bsz, q_len, num_key_value_heads * head_dim) -> (bsz, num_key_value_heads, q_len, head_dim)\n    q, k, v = (\n        states.view(bsz, q_len, num_heads, self.head_dim).transpose(1, 2)\n        for states, num_heads in (\n            (q, self.num_heads),\n            (k, self.num_key_value_heads),\n            (v, self.num_key_value_heads),\n        )\n    )\n    kv_len = k.shape[-2]  # initially, `kv_len` == `q_len`\n    past_kv_len = 0\n    if past_key_value is not None:\n        # if `past_key_value` is not None, `kv_len` > `q_len`.\n        past_kv_len = past_key_value[0].shape[-2]\n        kv_len += past_kv_len\n\n    # two `torch.Tensor` objs of shape (1, 1, kv_len, head_dim)\n    cos, sin = self.rotary_emb(v, seq_len=kv_len)\n    # (bsz, num_heads, q_len, head_dim), (bsz, num_key_value_heads, q_len, head_dim)\n    q, k = apply_rotary_pos_emb(q=q, k=k, cos=cos, sin=sin, position_ids=position_ids)\n    if past_key_value is not None:\n        # reuse k, v, self_attention\n        k = torch.cat([past_key_value[0], k], dim=2)\n        v = torch.cat([past_key_value[1], v], dim=2)\n\n    past_key_value = (k, v) if use_cache else None\n\n    # repeat k/v heads if n_kv_heads < n_heads\n    k = repeat_kv(hidden_states=k, n_rep=self.num_key_value_groups)\n    # (bsz, num_key_value_heads, q_len, head_dim) -> (bsz, num_heads, q_len, head_dim)\n    v = repeat_kv(hidden_states=v, n_rep=self.num_key_value_groups)\n    # (bsz, num_key_value_heads, q_len, head_dim) -> (bsz, num_heads, q_len, head_dim)\n\n    key_padding_mask = attention_mask\n    # (bsz, num_heads, q_len, head_dim) -> (bsz, q_len, num_heads, head_dim)\n    q, k, v = (states.transpose(1, 2) for states in (q, k, v))\n\n    if past_kv_len > 0:\n        q = torch.cat(\n            tensors=(\n                torch.full(\n                    size=(bsz, past_kv_len, self.num_heads, self.head_dim),\n                    fill_value=0.0,\n                    dtype=q.dtype,\n                    device=q.device,\n                ),\n                q,\n            ),\n            dim=1,\n        )  # (bsz, past_kv_len + q_len, num_heads, head_dim)\n\n    if key_padding_mask is None:\n        # (bsz, past_kv_len + q_len, num_heads, head_dim)\n        output = flash_attn_func(q=q, k=k, v=v, dropout_p=0.0, softmax_scale=None, causal=True)  # (bsz, )\n        output = rearrange(output, pattern=\"... h d -> ... (h d)\")  # (bsz, past_kv_len + q_len, num_heads * head_dim)\n    else:\n        q, indices, cu_q_lens, max_q_len = unpad_input(hidden_states=q, attention_mask=key_padding_mask)\n        kv, _, cu_kv_lens, max_kv_len = unpad_input(\n            hidden_states=torch.stack(tensors=(k, v), dim=2),\n            attention_mask=key_padding_mask,\n        )\n        output_unpad = flash_attn_varlen_kvpacked_func(\n            q=q,\n            kv=kv,\n            cu_seqlens_q=cu_q_lens,\n            cu_seqlens_k=cu_kv_lens,\n            max_seqlen_q=max_q_len,\n            max_seqlen_k=max_kv_len,\n            dropout_p=0.0,\n            softmax_scale=None,\n            causal=True,\n        )\n        output = pad_input(\n            hidden_states=rearrange(output_unpad, pattern=\"nnz h d -> nnz (h d)\"),\n            indices=indices,\n            batch=bsz,\n            seqlen=past_kv_len + q_len,\n        )  # (bsz, past_kv_len + q_len, num_heads * head_dim)\n\n    if past_kv_len > 0:\n        # Strip off the zero query outputs.\n        output = output[:, past_kv_len:, ...]  # (bsz, q_len, num_heads * head_dim)\n    output = self.o_proj(output)  # (bsz, q_len, hidden_size)\n    return output, None, past_key_value\n\n\ndef rms_norm_forward(self: LlamaRMSNorm, hidden_states: torch.Tensor) -> torch.Tensor:\n    \"\"\"\n    Formard function for RMS Norm\n    \"\"\"\n    return rms_norm(x=hidden_states, weight=self.weight, epsilon=self.variance_epsilon)\n\n\ndef replace_with_flash_attention(model: LlamaForCausalLM) -> None:\n    for name, module in model.named_modules():\n        if isinstance(module, LlamaAttention):\n            module.forward = MethodType(attention_forward, module)\n        if isinstance(module, LlamaModel):\n            module._prepare_decoder_attention_mask = MethodType(_prepare_decoder_attention_mask, module)\n        if isinstance(module, LlamaRMSNorm):\n            module.forward = MethodType(rms_norm_forward, module)\n", "path": "applications/Colossal-LLaMA-2/colossal_llama2/utils/flash_attention_patch.py"}]}