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 2025-01-07 00:50:56 -08:00
commit 67e3251577
1 changed files with 6 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

6
unsloth/models/llama.py
View file

@ -498,7 +498,9 @@ def LlamaDecoderLayer_fast_forward(
hidden_states += residual
else:
residual = hidden_states
print(501, hidden_states.dtype)
hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states)
print(503, hidden_states.dtype)
hidden_states, self_attn_weights, present_key_value = self.self_attn(
hidden_states = hidden_states,
causal_mask = causal_mask,
@ -510,12 +512,16 @@ def LlamaDecoderLayer_fast_forward(
padding_mask = padding_mask,
position_embeddings = position_embeddings,
)
print(515, hidden_states.dtype)
hidden_states = residual + hidden_states
# Fully Connected
residual = hidden_states
print(520, hidden_states.dtype)
hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states)
print(522, hidden_states.dtype)
hidden_states = self.mlp(hidden_states)
print(524, hidden_states.dtype)
hidden_states = residual + hidden_states
pass