|
import warnings |
|
from typing import Optional, Tuple |
|
|
|
import torch |
|
from flash_attn import __version__ as flash_attn_version |
|
from flash_attn.bert_padding import pad_input, unpad_input |
|
from flash_attn.flash_attn_interface import ( |
|
flash_attn_func, |
|
flash_attn_varlen_kvpacked_func, |
|
) |
|
from transformers.models.llama.modeling_llama import ( |
|
LlamaAttention, |
|
LlamaModel, |
|
rotate_half, |
|
) |
|
|
|
|
|
def apply_rotary_pos_emb(q, k, cos_sin, position_ids): |
|
gather_indices = position_ids[:, :, None, None] |
|
gather_indices = gather_indices.repeat( |
|
1, 1, cos_sin[0].shape[1], cos_sin[0].shape[3] |
|
) |
|
bsz = gather_indices.shape[0] |
|
cos, sin = ( |
|
torch.gather(x.transpose(1, 2).repeat(bsz, 1, 1, 1), 1, gather_indices) |
|
for x in cos_sin |
|
) |
|
q, k = ((x * cos) + (rotate_half(x) * sin) for x in (q, k)) |
|
return q, k |
|
|
|
|
|
def forward( |
|
self, |
|
hidden_states: torch.Tensor, |
|
attention_mask: Optional[torch.Tensor] = None, |
|
position_ids: Optional[torch.Tensor] = None, |
|
past_key_value: Optional[Tuple[torch.Tensor]] = None, |
|
output_attentions: bool = False, |
|
use_cache: bool = False, |
|
padding_mask: Optional[torch.Tensor] = None, |
|
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: |
|
if output_attentions: |
|
warnings.warn( |
|
"Output attentions is not supported for patched `LlamaAttention`, returning `None` instead." |
|
) |
|
|
|
bsz, q_len, _ = hidden_states.size() |
|
kv_heads = getattr(self, "num_key_value_heads", self.num_heads) |
|
|
|
q, k, v = ( |
|
op(hidden_states).view(bsz, q_len, nh, self.head_dim) |
|
for op, nh in ( |
|
(self.q_proj, self.num_heads), |
|
(self.k_proj, kv_heads), |
|
(self.v_proj, kv_heads), |
|
) |
|
) |
|
|
|
|
|
kv_seq_len = k.shape[1] |
|
past_kv_len = 0 |
|
if past_key_value is not None: |
|
past_kv_len = past_key_value[0].shape[2] |
|
kv_seq_len += past_kv_len |
|
|
|
cos_sin = self.rotary_emb(v, seq_len=kv_seq_len) |
|
q, k = apply_rotary_pos_emb(q, k, cos_sin, position_ids) |
|
|
|
if past_key_value is not None: |
|
assert ( |
|
flash_attn_version >= "2.1.0" |
|
), "past_key_value support requires flash-attn >= 2.1.0" |
|
|
|
k = torch.cat([past_key_value[0].transpose(1, 2), k], dim=1) |
|
v = torch.cat([past_key_value[1].transpose(1, 2), v], dim=1) |
|
|
|
past_key_value = (k.transpose(1, 2), v.transpose(1, 2)) if use_cache else None |
|
|
|
if attention_mask is None: |
|
output = flash_attn_func(q, k, v, 0.0, softmax_scale=None, causal=True).view( |
|
bsz, q_len, -1 |
|
) |
|
else: |
|
q, indices, cu_q_lens, max_s = unpad_input(q, attention_mask[:, -q_len:]) |
|
|
|
kv, _, cu_k_lens, max_k = unpad_input( |
|
torch.stack((k, v), dim=2), attention_mask |
|
) |
|
output_unpad = flash_attn_varlen_kvpacked_func( |
|
q, |
|
kv, |
|
cu_q_lens, |
|
cu_k_lens, |
|
max_s, |
|
max_k, |
|
0.0, |
|
softmax_scale=None, |
|
causal=True, |
|
) |
|
output_unpad = output_unpad.reshape(-1, self.num_heads * self.head_dim) |
|
output = pad_input(output_unpad, indices, bsz, q_len) |
|
|
|
return self.o_proj(output), None, past_key_value |
|
|
|
|
|
|
|
|
|
def _prepare_decoder_attention_mask( |
|
self, attention_mask, input_shape, inputs_embeds, past_key_values_length |
|
): |
|
|
|
if past_key_values_length > 0 and attention_mask is not None: |
|
attention_mask = torch.cat( |
|
( |
|
torch.full( |
|
(input_shape[0], past_key_values_length), |
|
True, |
|
dtype=attention_mask.dtype, |
|
device=attention_mask.device, |
|
), |
|
attention_mask, |
|
), |
|
dim=-1, |
|
) |
|
|
|
if attention_mask is not None and torch.all(attention_mask): |
|
return None |
|
|
|
return attention_mask |
|
|
|
|
|
def replace_llama_attn_with_flash_attn(): |
|
cuda_major, cuda_minor = torch.cuda.get_device_capability() |
|
if cuda_major < 8: |
|
warnings.warn( |
|
"Flash attention is only supported on A100 or H100 GPU during training due to head dim > 64 backward." |
|
"ref: https://github.com/HazyResearch/flash-attention/issues/190#issuecomment-1523359593" |
|
) |
|
|
|
LlamaModel._prepare_decoder_attention_mask = _prepare_decoder_attention_mask |
|
LlamaAttention.forward = forward |
|
|
|
|
|
def test(): |
|
from fastchat.train.llama_flash_attn_monkey_patch import forward as fastchat_forward |
|
from transformers.models.llama.configuration_llama import LlamaConfig |
|
|
|
config = LlamaConfig( |
|
hidden_size=1024, |
|
intermediate_size=128, |
|
num_hidden_layers=1, |
|
num_attention_heads=8, |
|
max_position_embeddings=16, |
|
) |
|
device = torch.device("cuda") |
|
model = LlamaModel(config) |
|
attn = LlamaAttention(config).to(device).half() |
|
bsz, hs, seqlen = 2, config.hidden_size, config.max_position_embeddings |
|
position_ids = torch.arange(seqlen, dtype=torch.long, device=device).view( |
|
-1, seqlen |
|
) |
|
|
|
mask = torch.full((bsz, seqlen), True, dtype=torch.bool, device=device) |
|
for i in range(4): |
|
hidden = torch.rand((bsz, seqlen, hs), dtype=torch.float16, device=device) |
|
if i: |
|
mask[0, -i:] = False |
|
mask[1, :i] = False |
|
|
|
lmask = model._prepare_decoder_attention_mask(mask, hidden.shape[:2], hidden, 0) |
|
ref, _, _ = attn.forward( |
|
hidden, attention_mask=lmask, position_ids=position_ids |
|
) |
|
|
|
fast, _, _ = fastchat_forward( |
|
attn, hidden, attention_mask=mask, position_ids=position_ids |
|
) |
|
|
|
lmask = _prepare_decoder_attention_mask( |
|
model, mask, hidden.shape[:2], hidden, 0 |
|
) |
|
test, _, _ = forward( |
|
attn, hidden, attention_mask=lmask, position_ids=position_ids |
|
) |
|
|
|
print(f"Mean(abs(ref)) = {torch.mean(torch.abs(ref))}") |
|
print(f"Mean(abs(ref - fast)) = {torch.mean(torch.abs(ref - fast))}") |
|
print(f"Mean(abs(ref - test)) = {torch.mean(torch.abs(ref - test))}") |
|
print(f"Mean(abs(fast - test)) = {torch.mean(torch.abs(fast - test))}") |
|
print(f"allclose(fast, test) = {torch.allclose(fast, test)}") |
|
|
|
with torch.no_grad(): |
|
|
|
hidden = torch.rand((bsz, seqlen, hs), dtype=torch.float16, device=device) |
|
part_len = seqlen // 4 |
|
assert part_len * 4 == seqlen |
|
mask = torch.full((bsz, seqlen), True, dtype=torch.bool, device=device) |
|
mask[0, -2:] = False |
|
lmask = _prepare_decoder_attention_mask( |
|
model, mask, hidden.shape[:2], hidden, 0 |
|
) |
|
oneshot, _, _ = forward( |
|
attn, hidden, attention_mask=lmask, position_ids=position_ids |
|
) |
|
parts = [] |
|
past_kv, past_kv_len = None, 0 |
|
for i in range(4): |
|
start = part_len * i |
|
end = start + part_len |
|
hidden_part = hidden[:, start:end, ...] |
|
lmask = _prepare_decoder_attention_mask( |
|
model, |
|
mask[:, start:end], |
|
hidden_part.shape[:2], |
|
hidden_part, |
|
past_kv_len, |
|
) |
|
part, _, past_kv = forward( |
|
attn, |
|
hidden_part.clone(), |
|
attention_mask=lmask, |
|
position_ids=position_ids[:, start:end], |
|
past_key_value=past_kv, |
|
use_cache=True, |
|
) |
|
parts.append(part) |
|
past_kv_len = past_kv[0].shape[2] |
|
|
|
print( |
|
f"allclose(oneshot[:, 0], parts[0]) = {torch.allclose(oneshot[:, :part_len], parts[0])}" |
|
) |
|
print( |
|
f"allclose(oneshot, parts) = {torch.allclose(oneshot, torch.cat(parts, dim=1))}" |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
test() |
|
|
