diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c98feeca1e..cf05d432c7 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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