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

Update gemma.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-02-24 19:33:04 +11:00
commit b174e54507
1 changed files with 16 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

28
unsloth/models/gemma.py
View file

@ -618,22 +618,22 @@ class FastGemmaModel(FastLlamaModel):
# Also patch all dtypes - BnB seems to not allocate the correct type?
# BnB default dtype seems to be float16!
correct_dtype = lm_head.weight.dtype
# correct_dtype = lm_head.weight.dtype
for name, module in model.named_modules():
if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)):
weight = module.weight
quant_state = weight.quant_state
# for name, module in model.named_modules():
# if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)):
# weight = module.weight
# quant_state = weight.quant_state
if type(quant_state) is list:
# BnB seems to have float16 as default!
module.weight.quant_state[2] = correct_dtype # Cast to correct dtype
else:
# https://github.com/TimDettmers/bitsandbytes/pull/763/files
quant_state.dtype = correct_dtype
pass
pass
pass
# if type(quant_state) is list:
# # BnB seems to have float16 as default!
# module.weight.quant_state[2] = correct_dtype # Cast to correct dtype
# else:
# # https://github.com/TimDettmers/bitsandbytes/pull/763/files
# quant_state.dtype = correct_dtype
# pass
# pass
# pass
# Add 1 to weight
# return output * (1 + self.weight)