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

lm_head

Browse source
This commit is contained in:
Daniel Han-Chen 2024-03-18 19:43:21 +11:00
commit daf267967b
2 changed files with 4 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

3
unsloth/models/llama.py
View file

@ -753,7 +753,8 @@ def CausalLM_fast_forward(fast_forward_inference):
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())
lm_head = self.lm_head.weight
logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
logits = logits.unsqueeze(0).unsqueeze(0)
else:
logits = self.lm_head(hidden_states)

3
unsloth/models/mistral.py
View file

@ -225,7 +225,8 @@ 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())
lm_head = self.lm_head.weight
logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
logits = logits.unsqueeze(0).unsqueeze(0)
else:
logits = self.lm_head(hidden_states)