From f14d37a778c853a1d7484acc0141cfb983c09f41 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 27 Oct 2024 19:08:14 -0700 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 044ea6e244..a53f52fe27 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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