From e04235d13a3724e2d5870c1a88307a6184f1c981 Mon Sep 17 00:00:00 2001 From: Muhammad Osama Date: Sun, 5 Jan 2025 18:18:42 -0600 Subject: [PATCH 1/2] changing model to base_model if peft model is already used --- unsloth/models/llama.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index c94514966f..128e0fd762 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1967,29 +1967,29 @@ 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\ + dtype = model.base_model.model.embed_tokens.modules_to_save.default.weight.dtype + model.base_model.model.embed_tokens.modules_to_save.default\ .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) + model.base_model.model.embed_tokens.modules_to_save.default.requires_grad_(True) # [TODO] Move old embed_tokens to CPU - should be disk! - model.model.model.embed_tokens.original_module\ + model.base_model.model.embed_tokens.original_module\ .to(device = "cpu", non_blocking = True) - model.model.model.embed_tokens.original_module.requires_grad_(False) + model.base_model.model.embed_tokens.original_module.requires_grad_(False) pass 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\ + dtype = model.base_model.model.lm_head.modules_to_save.default.weight.dtype + model.base_model.lm_head.modules_to_save.default\ .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) + model.base_model.lm_head.modules_to_save.default.requires_grad_(True) # [TODO] Move old lm_head to CPU - should be disk! - model.model.lm_head.original_module\ + model.base_model.lm_head.original_module\ .to(device = "cpu", non_blocking = True) - model.model.lm_head.original_module.requires_grad_(False) + model.base_model.lm_head.original_module.requires_grad_(False) pass return model From 706b4ed717fd419ff4a08abf690b9f4a6324ed69 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 6 Jan 2025 22:06:00 -0800 Subject: [PATCH 2/2] Update llama.py --- unsloth/models/llama.py | 69 +++++++++++++++++++++++++++-------------- 1 file changed, 46 insertions(+), 23 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 128e0fd762..0cfa1d04a8 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1967,29 +1967,41 @@ class FastLlamaModel: if "embed_tokens" in new_target_modules: print("Unsloth: Training embed_tokens in mixed precision to save VRAM") - dtype = model.base_model.model.embed_tokens.modules_to_save.default.weight.dtype - model.base_model.model.embed_tokens.modules_to_save.default\ - .to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True) - model.base_model.model.embed_tokens.modules_to_save.default.requires_grad_(True) + new_dtype = model.get_input_embeddings().modules_to_save.default.weight.dtype + if new_dtype == torch.float16: + # See https://github.com/unslothai/unsloth/pull/1200 + # Tesla T4 must use float32 and not float16 + new_dtype = torch.float32 + pass + + model.get_input_embeddings().modules_to_save.default\ + .to(device = "cuda:0", dtype = new_dtype, non_blocking = True) + model.get_input_embeddings().modules_to_save.default.requires_grad_(True) # [TODO] Move old embed_tokens to CPU - should be disk! - model.base_model.model.embed_tokens.original_module\ + model.get_input_embeddings().original_module\ .to(device = "cpu", non_blocking = True) - model.base_model.model.embed_tokens.original_module.requires_grad_(False) + model.get_input_embeddings().original_module.requires_grad_(False) pass if "lm_head" in new_target_modules: print("Unsloth: Training lm_head in mixed precision to save VRAM") - dtype = model.base_model.model.lm_head.modules_to_save.default.weight.dtype - model.base_model.lm_head.modules_to_save.default\ - .to(device = "cuda:0", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True) - model.base_model.lm_head.modules_to_save.default.requires_grad_(True) + new_dtype = model.get_output_embeddings().modules_to_save.default.weight.dtype + if new_dtype == torch.float16: + # See https://github.com/unslothai/unsloth/pull/1200 + # Tesla T4 must use float32 and not float16 + new_dtype = torch.float32 + pass + + model.get_output_embeddings().modules_to_save.default\ + .to(device = "cuda:0", dtype = new_dtype, non_blocking = True) + model.get_output_embeddings().modules_to_save.default.requires_grad_(True) # [TODO] Move old lm_head to CPU - should be disk! - model.base_model.lm_head.original_module\ + model.get_output_embeddings().original_module\ .to(device = "cpu", non_blocking = True) - model.base_model.lm_head.original_module.requires_grad_(False) + model.get_output_embeddings().original_module.requires_grad_(False) pass return model @@ -2216,25 +2228,36 @@ 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")) + assert(hasattr(model.get_input_embeddings(), "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", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True) - model.model.model.embed_tokens.modules_to_save.default.requires_grad_(True) + new_dtype = model.get_input_embeddings().modules_to_save.default.weight.dtype + if new_dtype == torch.float16: + # See https://github.com/unslothai/unsloth/pull/1200 + # Tesla T4 must use float32 and not float16 + new_dtype = torch.float32 + pass + + model.get_input_embeddings().modules_to_save.default\ + .to(device = "cuda:0", dtype = new_dtype, non_blocking = True) + model.get_input_embeddings().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")) + assert(hasattr(model.get_output_embeddings(), "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", dtype=(dtype if (dtype != torch.float16) else torch.float32), non_blocking = True) - model.model.lm_head.modules_to_save.default.requires_grad_(True) + new_dtype = model.get_output_embeddings().modules_to_save.default.weight.dtype + if new_dtype == torch.float16: + # See https://github.com/unslothai/unsloth/pull/1200 + # Tesla T4 must use float32 and not float16 + new_dtype = torch.float32 + pass + + model.get_output_embeddings().modules_to_save.default\ + .to(device = "cuda:0", dtype = new_dtype, non_blocking = True) + model.get_output_embeddings().modules_to_save.default.requires_grad_(True) pass # Patch tokenizer to pad to the right