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 2024-10-27 19:08:14 -07:00
commit f14d37a778
1 changed files with 6 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

7
unsloth/models/llama.py
View file

@ -1869,14 +1869,19 @@ class FastLlamaModel:
# We also do this for the lm_head
if old_output_embedding.numel() != 0:
requires_grad = old_output_embedding.requires_grad
lm_head = torch.nn.Linear(1, 1, bias = None)
del lm_head.weight
lm_head.weight = old_output_embedding if not is_tied else old_input_embedding
lm_head.in_features = lm_head.weight.shape[1]
lm_head.out_features = lm_head.weight.shape[0]
lm_head.weight.requires_grad_(requires_grad)
model.lm_head = lm_head
model.set_output_embeddings(lm_head)
if hasattr(model, "lm_head"): model.lm_head = lm_head
correct_dtype = lm_head.weight.dtype
else:
correct_dtype = old_input_embedding.dtype