From 59463ce86416fbf701d9b0ff1d999e204f605462 Mon Sep 17 00:00:00 2001 From: lurf21 <93976703+lurf21@users.noreply.github.com> Date: Wed, 19 Mar 2025 17:06:48 +0800 Subject: [PATCH] fix: config.torch_dtype in LlamaModel_fast_forward_inference (#2091) * fix: config.torch_dtype in LlamaModel_fast_forward_inference * Update llama.py * update for consistency --------- Co-authored-by: Daniel Han --- unsloth/models/llama.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 4bf1357169..61cf05e110 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -652,13 +652,7 @@ def LlamaModel_fast_forward( if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) - # inputs_embeds = inputs_embeds.to(self.config.torch_dtype) - torch_dtype = __DTYPE_MAP.get(self.config.torch_dtype, None) - if torch_dtype is not None: - inputs_embeds = inputs_embeds.to(torch_dtype) - else: - raise TypeError("Unsloth: torch_dtype for models is not bfloat16, float16 or float32!") - pass + inputs_embeds = inputs_embeds.to(_get_dtype(self.config.torch_dtype)) # Normalized from Gemma IS_GEMMA = self.config.model_type.startswith("gemma") @@ -924,7 +918,7 @@ def LlamaModel_fast_forward_inference( mlp_size = self.config.intermediate_size X = self.model.embed_tokens(input_ids) - X = X.to(self.config.torch_dtype) + X = X.to(_get_dtype(self.config.torch_dtype)) bsz, q_len, hd = X.shape assert(q_len == 1) # Get saved buffers to reduce memory movement