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

revert inference

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-31 03:40:55 +11:00
commit c0edaa46db
2 changed files with 15 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

18
unsloth/models/llama.py
View file

@ -426,7 +426,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 False:#(past_key_value is not None and q_len == 1):
if (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)
@ -597,7 +597,7 @@ def LlamaModel_fast_forward(
hidden_states = inputs_embeds
if self.gradient_checkpointing and self.training:
if past_key_values is None and self.gradient_checkpointing and self.training:
if use_cache:
logger.warning_once(
"Unsloth: `use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`"
@ -656,7 +656,7 @@ def LlamaModel_fast_forward(
pass
bsz, q_len, hd = hidden_states.size()
if False:#(past_key_values is not None and q_len == 1):
if (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)
@ -719,12 +719,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