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

donot upcast lm_head and embeddings to float32 (#1186)

Browse source
This commit is contained in:
Datta Nimmaturi 2024-10-25 13:58:12 +05:30 • committed by GitHub
commit ffe4ba1a43
1 changed files with 4 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

8
unsloth/models/llama.py
View file

@ -1955,7 +1955,7 @@ class FastLlamaModel:
print("Unsloth: Casting embed_tokens to float32")
model.model.model.embed_tokens.modules_to_save.default\
.to(device = "cuda:0", dtype = torch.float32, non_blocking = True)
.to(device = "cuda:0", non_blocking = True)
model.model.model.embed_tokens.modules_to_save.default.requires_grad_(True)
# [TODO] Move old embed_tokens to CPU - should be disk!
@ -1968,7 +1968,7 @@ class FastLlamaModel:
print("Unsloth: Casting lm_head to float32")
model.model.lm_head.modules_to_save.default\
.to(device = "cuda:0", dtype = torch.float32, non_blocking = True)
.to(device = "cuda:0", non_blocking = True)
model.model.lm_head.modules_to_save.default.requires_grad_(True)
# [TODO] Move old lm_head to CPU - should be disk!
@ -2206,7 +2206,7 @@ class FastLlamaModel:
print("Unsloth: Casting embed_tokens to float32")
assert(hasattr(model.model.model.embed_tokens, "modules_to_save"))
model.model.model.embed_tokens.modules_to_save.default\
.to(device = "cuda:0", dtype = torch.float32, non_blocking = True)
.to(device = "cuda:0", non_blocking = True)
model.model.model.embed_tokens.modules_to_save.default.requires_grad_(True)
pass
@ -2214,7 +2214,7 @@ class FastLlamaModel:
print("Unsloth: Casting lm_head to float32")
assert(hasattr(model.model.lm_head, "modules_to_save"))
model.model.lm_head.modules_to_save.default\
.to(device = "cuda:0", dtype = torch.float32, non_blocking = True)
.to(device = "cuda:0", non_blocking = True)
model.model.lm_head.modules_to_save.default.requires_grad_(True)
pass