From 9032334c302f7c6860860e165bbb9a3cb53d447c Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 18 Mar 2024 03:22:49 +1100 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 73eaa33652..52ab61b4b8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1336,7 +1336,6 @@ class FastLlamaModel: "We shall do it for you!" ) train_lm_head = True - model.model.embed_tokens.to(torch.float32, non_blocking = True) elif module == "embed_tokens": logger.warning_once( @@ -1344,7 +1343,6 @@ class FastLlamaModel: "We shall do it for you!" ) train_embed_tokens = True - model.lm_head.to(torch.float32, non_blocking = True) else: assert(module in accepted_modules) @@ -1386,8 +1384,12 @@ class FastLlamaModel: # Now patch lm_head and embed_tokens if train_embed_tokens: + model.model.model.embed_tokens.to(torch.float32) model.model.model.embed_tokens.requires_grad_(True) + pass + if train_lm_head: + model.model.lm_head.to(torch.float32) model.model.lm_head.requires_grad_(True) pass