ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: config.torch_dtype in LlamaModel_fast_forward_inference (#2091)

Browse source 
* fix: config.torch_dtype in LlamaModel_fast_forward_inference

* Update llama.py

* update for consistency

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
This commit is contained in:
lurf21 2025-03-19 17:06:48 +08:00 • committed by GitHub
commit 00a98f17f5
1 changed files with 2 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

10
unsloth/models/llama.py
View file

@ -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