ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update llama.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-02-04 03:50:31 +11:00
commit 3bfb0ebf9d
1 changed files with 10 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

29
unsloth/models/llama.py
View file

@ -316,21 +316,12 @@ def LlamaAttention_fast_forward(
K = K.reshape(bsz, n_heads, kv_seq_len, head_dim)
V = V.reshape(bsz, n_heads, kv_seq_len, head_dim)
pass
attn_weights = torch.matmul(Q, K.transpose(2, 3)) / math_sqrt(self.head_dim)
attn_weights = attn_weights + attention_mask
# upcast attention to fp32
attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(Q.dtype)
attn_weights = torch.nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
A = torch.matmul(attn_weights, V)
# Must be contiguous or else results are False!
# https://github.com/pytorch/pytorch/issues/112577
# Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
# # Needs (batch_size, n_heads, seq_len, head_dim)
# # is_casual and attention_mask must not be both set!
# A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
Q, K, V = Q.contiguous(), K.contiguous(), V.contiguous()
# Needs (batch_size, n_heads, seq_len, head_dim)
# is_casual and attention_mask must not be both set!
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)
# Go back to (batch_size, seq_len, n_heads, head_dim)
A = A.transpose(1, 2).contiguous()
pass
@ -367,7 +358,7 @@ def LlamaDecoderLayer_fast_forward(
(see `past_key_values`).
past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
"""
if False:#past_key_value is not None:
if past_key_value is not None:
do_prefill = not hasattr(self.self_attn, "paged_attention")
# Self Attention
@ -480,7 +471,7 @@ def LlamaModel_fast_forward(
pass
# We already handle KV cache position_ids ourselves.
if (past_key_values_length != 0):
if False:#(past_key_values_length != 0):
position_ids = torch.arange(
past_key_values_length, seq_length + past_key_values_length,
dtype = torch.int32,
@ -673,7 +664,7 @@ def LlamaForCausalLM_fast_forward(
*args, **kwargs,
) -> Union[Tuple, CausalLMOutputWithPast]:
if causal_mask is None:# and past_key_values is None:
if causal_mask is None and past_key_values is not None:
causal_mask = xformers.attn_bias.LowerTriangularMask()
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
@ -685,8 +676,8 @@ def LlamaForCausalLM_fast_forward(
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
self.model._has_no_labels = labels is None
if False:#past_key_values is not None and \
#hasattr(self.model.layers[0].self_attn, "paged_attention"):
if past_key_values is not None and \
hasattr(self.model.layers[0].self_attn, "paged_attention"):
outputs = LlamaModel_fast_forward_inference(
self.model,
input_ids,
@ -709,7 +700,7 @@ def LlamaForCausalLM_fast_forward(
hidden_states = outputs[0]
bsz, q_len, hd = hidden_states.shape
if False:#bsz == 1 and q_len == 1:
if bsz == 1 and q_len == 1:
logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
logits = logits.unsqueeze(0).unsqueeze(0)
else: