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

Update llama.py

Browse source
This commit is contained in:
Daniel Han 2025-01-06 18:56:26 -08:00
commit a229db5a85
1 changed files with 43 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

48
unsloth/models/llama.py
View file

@ -1968,8 +1968,18 @@ class FastLlamaModel:
print("Unsloth: Training embed_tokens in mixed precision to save VRAM")
dtype = model.model.model.embed_tokens.modules_to_save.default.weight.dtype
# Now patch lm_head and embed_tokens
if dtype == torch.float16:
# See https://github.com/unslothai/unsloth/pull/1200
# Tesla T4 must use float32 and not float16
modules_to_save_dtype = torch.float32
else:
# Can be bfloat16
modules_to_save_dtype = dtype
pass
model.model.model.embed_tokens.modules_to_save.default\
.to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True)
.to(device = "cuda:0", dtype = modules_to_save_dtype, 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!
@ -1982,8 +1992,17 @@ class FastLlamaModel:
print("Unsloth: Training lm_head in mixed precision to save VRAM")
dtype = model.model.model.lm_head.modules_to_save.default.weight.dtype
# Now patch lm_head and embed_tokens
if dtype == torch.float16:
# See https://github.com/unslothai/unsloth/pull/1200
# Tesla T4 must use float32 and not float16
modules_to_save_dtype = torch.float32
else:
# Can be bfloat16
modules_to_save_dtype = dtype
pass
model.model.lm_head.modules_to_save.default\
.to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True)
.to(device = "cuda:0", dtype = modules_to_save_dtype, non_blocking = True)
model.model.lm_head.modules_to_save.default.requires_grad_(True)
# [TODO] Move old lm_head to CPU - should be disk!
@ -2216,14 +2235,23 @@ class FastLlamaModel:
model = FastLlamaModel.patch_peft_model(model, use_gradient_checkpointing)
# Now patch lm_head and embed_tokens
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
# Now patch lm_head and embed_tokens
if dtype == torch.float16:
# See https://github.com/unslothai/unsloth/pull/1200
# Tesla T4 must use float32 and not float16
modules_to_save_dtype = torch.float32
else:
# Can be bfloat16
modules_to_save_dtype = dtype
pass
model.model.model.embed_tokens.modules_to_save.default\
.to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True)
.to(device = "cuda:0", dtype = modules_to_save_dtype, non_blocking = True)
model.model.model.embed_tokens.modules_to_save.default.requires_grad_(True)
pass
@ -2232,8 +2260,18 @@ class FastLlamaModel:
assert(hasattr(model.model.lm_head, "modules_to_save"))
dtype = model.model.lm_head.modules_to_save.default.weight.dtype
# Now patch lm_head and embed_tokens
if dtype == torch.float16:
# See https://github.com/unslothai/unsloth/pull/1200
# Tesla T4 must use float32 and not float16
modules_to_save_dtype = torch.float32
else:
# Can be bfloat16
modules_to_save_dtype = dtype
pass
model.model.lm_head.modules_to_save.default\
.to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True)
.to(device = "cuda:0", dtype = modules_to_save_dtype, non_blocking = True)
model.model.lm_head.modules_to_save.default.requires_grad_(True)
pass