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:30:32 -08:00
commit 2d66722781
1 changed files with 6 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/models/llama.py
View file

@ -996,18 +996,21 @@ def CausalLM_fast_forward(fast_forward_inference):
lm_head = self.lm_head.weight
logit_softcapping = getattr(self.config, "final_logit_softcapping", 0)
logit_scaling = getattr(self.config, "logit_scale", 0)
dtype = lm_head.dtype
if bsz == 1 and q_len == 1:
logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype))
logits = torch.mv(lm_head, hidden_states.ravel().to(dtype))
logits = logits.unsqueeze(0).unsqueeze(0)
elif num_logits_to_keep != 0:
logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(lm_head.dtype))
logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(dtype))
else:
RETURN_LOGITS = os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1"
# < 1024 Normal Unsloth uses less VRAM!
if bsz*q_len <= 1024: RETURN_LOGITS = True
if not RETURN_LOGITS and HAS_CUT_CROSS_ENTROPY and labels is not None:
print(hidden_states, lm_head)
n_items = kwargs.get("num_items_in_batch", None) or kwargs.get("n_items", None)
loss = fused_linear_cross_entropy(
hidden_states = hidden_states,
@ -1029,7 +1032,7 @@ def CausalLM_fast_forward(fast_forward_inference):
)
return output
pass
logits = self.lm_head(hidden_states.to(lm_head.dtype))
logits = self.lm_head(hidden_states.to(dtype))
pass
torch_dtype = __DTYPE_MAP.get(self.config.torch_dtype, None)