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

Add dropout to granite to match HF's implementation (#1557)

Browse source 
Signed-off-by: datta0 <venkatadattasainimmaturi@gmail.com>
This commit is contained in:
Datta Nimmaturi 2025-01-19 17:24:12 +05:30 • committed by GitHub
commit cc33bd05dc
2 changed files with 9 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

7
unsloth/models/granite.py
View file

@ -89,6 +89,7 @@ def GraniteAttention_fast_forward(
n_groups = self.num_key_value_groups
n_kv_heads = self.config.num_key_value_heads
head_dim = self.head_dim
dropout_p = self.config.attention_dropout if self.training else 0
assert(n_kv_heads * n_groups == n_heads)
Q, K, V = self.apply_qkv(self, hidden_states)
@ -135,7 +136,7 @@ def GraniteAttention_fast_forward(
Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim)
pass
A = xformers_attention(Q, K, V, attn_bias = causal_mask, scale=self.scaling)
A = xformers_attention(Q, K, V, attn_bias = causal_mask, scale=self.scaling, p=dropout_p)
A = A.view(bsz, q_len, n_heads, head_dim)
elif HAS_FLASH_ATTENTION and attention_mask is None:
@ -143,7 +144,7 @@ def GraniteAttention_fast_forward(
K = K.transpose(1, 2)
V = V.transpose(1, 2)
window = (kv_seq_len, kv_seq_len)
A = flash_attn_func(Q, K, V, causal = True, window_size = window, softmax_scale=self.scaling)
A = flash_attn_func(Q, K, V, causal = True, window_size = window, softmax_scale=self.scaling, dropout_p=dropout_p)
else:
# Grouped query attention
# if n_groups != 1:
@ -157,7 +158,7 @@ def GraniteAttention_fast_forward(
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, scale = self.scaling, is_causal = False)
A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, scale = self.scaling, is_causal = False, dropout_p=dropout_p)
# Go back to (batch_size, seq_len, n_heads, head_dim)
A = A.transpose(1, 2).contiguous()
pass

6
unsloth/models/llama.py
View file

@ -636,6 +636,7 @@ def LlamaModel_fast_forward(
IS_GEMMA2 = self.config.model_type.startswith("gemma2")
IS_COHERE = self.config.model_type.startswith("cohere")
IS_GRANITE = self.config.model_type.startswith("granite")
train_embed_tokens = self.embed_tokens.weight.requires_grad
if IS_GEMMA:
@ -792,9 +793,12 @@ def LlamaModel_fast_forward(
pass
pass
if IS_ATTENTION_REFACTOR and not hasattr(self.layers[0].self_attn, "rotary_emb"):
if (IS_ATTENTION_REFACTOR and (hasattr(self, "rotary_emb") or not hasattr(self.layers[0].self_attn, "rotary_emb"))) or IS_GRANITE:
# Transformers main has made it mandatory to pass position_embeddings
# https://github.com/huggingface/transformers/pull/34858
# Also, transformers 4.45.0 supports granite but with the attention refactor (it always had the refactor)
# unsloth's check for granite too has "version >= 4.45.0 (rightly so)".
# so let granite always use the attention refactor implementation.
position_embeddings = self.rotary_emb(hidden_states, position_ids, self.config.max_position_embeddings)
else:
position_embeddings = None