From e3fa7aa908e0d6e6c0f92a65c0726bc79ccdeca4 Mon Sep 17 00:00:00 2001 From: moontidef <53668275+relic-yuexi@users.noreply.github.com> Date: Mon, 5 Aug 2024 14:45:34 +0800 Subject: [PATCH] fix: fix config.torch_dtype bug (#874) fix the bug #404 and the bug https://github.com/hiyouga/LLaMA-Factory/issues/4698#issue-2393500878 --- unsloth/models/llama.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index e6c9280bc5..445e5026f0 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -570,7 +570,14 @@ def LlamaModel_fast_forward( # Embed positions if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) - + + if self.config.torch_dtype == "float32": + self.config.torch_dtype = torch.float32 + elif self.config.torch_dtype == "bfloat16": + self.config.torch_dtype = torch.bfloat16 + elif self.config.torch_dtype == "float16": + self.config.torch_dtype = torch.float16 + inputs_embeds = inputs_embeds.to(self.config.torch_dtype) # Normalized from Gemma