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

Remove fast path

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-30 17:33:08 +11:00
commit 270df81b60
2 changed files with 14 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

16
unsloth/models/llama.py
View file

@ -345,7 +345,7 @@ def LlamaDecoderLayer_fast_forward(
past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
"""
bsz, q_len, hd = hidden_states.size()
if (past_key_value is not None and q_len == 1):
if False:#(past_key_value is not None and q_len == 1):
# Self Attention
residual = hidden_states
hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states)
@ -575,7 +575,7 @@ def LlamaModel_fast_forward(
pass
bsz, q_len, hd = hidden_states.size()
if (past_key_values is not None and q_len == 1):
if False:#(past_key_values is not None and q_len == 1):
hidden_states = fast_rms_layernorm_inference(self.norm, hidden_states)
else:
hidden_states = fast_rms_layernorm(self.norm, hidden_states)
@ -638,12 +638,12 @@ def LlamaForCausalLM_fast_forward(
)
hidden_states = outputs[0]
bsz, q_len, hd = hidden_states.shape
if bsz == 1 and q_len == 1:
logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
logits = logits.unsqueeze(0).unsqueeze(0)
else:
logits = self.lm_head(hidden_states)
# bsz, q_len, hd = hidden_states.shape
# if bsz == 1 and q_len == 1:
# logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
# logits = logits.unsqueeze(0).unsqueeze(0)
# else:
logits = self.lm_head(hidden_states)
pass
loss = None

12
unsloth/models/mistral.py
View file

@ -210,12 +210,12 @@ def MistralForCausalLM_fast_forward(
)
hidden_states = outputs[0]
bsz, q_len, hd = hidden_states.shape
if bsz == 1 and q_len == 1:
logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
logits = logits.unsqueeze(0).unsqueeze(0)
else:
logits = self.lm_head(hidden_states)
# bsz, q_len, hd = hidden_states.shape
# if bsz == 1 and q_len == 1:
# logits = torch.mv(self.lm_head.weight, hidden_states.ravel())
# logits = logits.unsqueeze(0).unsqueeze(0)
# else:
logits = self.lm_head(hidden_states)
pass
loss = None