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

Fix/casting continue pretraining (#1200)

Browse source 
* Bring back float32 if float16 instead of bfloat16

* Refactor mixed precision handling for lm_head and embed_tokens to ensure correct dtype usage

* Fix dtype retrieval for embed_tokens and lm_head in mixed precision training

* Fix dtype retrieval for embed_tokens and lm_head to use weight dtype in mixed precision training

* Fix dtype handling for embed_tokens and lm_head to ensure correct float32 usage in mixed precision training

* Fix dtype assignment for lm_head modules to ensure correct weight dtype usage in mixed precision training
This commit is contained in:
Edd 2024-10-28 02:06:45 +04:00 • committed by GitHub
commit 539fcea071
1 changed files with 10 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

14
unsloth/models/llama.py
View file

@ -1958,8 +1958,9 @@ class FastLlamaModel:
if "embed_tokens" in new_target_modules:
print("Unsloth: Training embed_tokens in mixed precision to save VRAM")
dtype = model.model.model.embed_tokens.modules_to_save.default.weight.dtype
model.model.model.embed_tokens.modules_to_save.default\
.to(device = "cuda:0", non_blocking = True)
.to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), 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!
@ -1971,8 +1972,9 @@ class FastLlamaModel:
if "lm_head" in new_target_modules:
print("Unsloth: Training lm_head in mixed precision to save VRAM")
dtype = model.model.model.lm_head.modules_to_save.default.weight.dtype
model.model.lm_head.modules_to_save.default\
.to(device = "cuda:0", non_blocking = True)
.to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True)
model.model.lm_head.modules_to_save.default.requires_grad_(True)
# [TODO] Move old lm_head to CPU - should be disk!
@ -2209,16 +2211,20 @@ class FastLlamaModel:
if train_embed_tokens:
print("Unsloth: Training embed_tokens in mixed precision to save VRAM")
assert(hasattr(model.model.model.embed_tokens, "modules_to_save"))
dtype = model.model.model.embed_tokens.modules_to_save.default.weight.dtype
model.model.model.embed_tokens.modules_to_save.default\
.to(device = "cuda:0", non_blocking = True)
.to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True)
model.model.model.embed_tokens.modules_to_save.default.requires_grad_(True)
pass
if train_lm_head:
print("Unsloth: Training lm_head in mixed precision to save VRAM")
assert(hasattr(model.model.lm_head, "modules_to_save"))
dtype = model.model.lm_head.modules_to_save.default.weight.dtype
model.model.lm_head.modules_to_save.default\
.to(device = "cuda:0", non_blocking = True)
.to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True)
model.model.lm_head.modules_to_save.default.requires_grad_(True)
pass