From d6ac9b56c125a7e0a80b7ab006ca9929aa778b3a Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Tue, 12 Mar 2024 20:12:55 +1100 Subject: [PATCH] upcasting --- unsloth/models/gemma.py | 1 + unsloth/models/llama.py | 8 +++++++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index bcd0e1abd9..7bfec43e51 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -154,6 +154,7 @@ def GemmaModel_fast_forward_inference( out_weight = torch.empty_like(self.layers[0].input_layernorm.weight, dtype = torch.float32, device = "cuda") hidden_states = self.embed_tokens(input_ids) + hidden_states = hidden_states.to(self.config.torch_dtype) # 3072**0.5 = 55.5000 in bfloat16, whilst 55.4256 in float32 # 2048**0.5 = 45.2500 in bfloat16, whilst 45.2548 in float32 hidden_states *= torch.tensor(math_sqrt(self.config.hidden_size), dtype = hidden_states.dtype) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index eb66b17d4a..d83d9b76f2 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -509,7 +509,10 @@ def LlamaModel_fast_forward( if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) - # Mormalized from Gemma + # Downcast to the correct dtype ie float32 to float16 + inputs_embeds = inputs_embeds.to(self.config.torch_dtype) + + # Normalized from Gemma IS_GEMMA = self.config.model_type == "gemma" train_embed_tokens = self.embed_tokens.weight.requires_grad @@ -665,6 +668,7 @@ def LlamaModel_fast_forward_inference( input_ids = input_ids[:,:self.max_seq_length] hidden_states = self.embed_tokens(input_ids) + hidden_states = hidden_states.to(self.config.torch_dtype) next_decoder_cache = [] for idx, decoder_layer in enumerate(self.layers): @@ -1334,6 +1338,7 @@ 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( @@ -1341,6 +1346,7 @@ 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)