From 625209e11febd6d91a5edf8dfdfd04906a013c9f Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Sat, 26 Oct 2024 00:47:54 +0530 Subject: [PATCH] Cleanup upcast logs (#1188) --- unsloth/models/llama.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 9c5499dc75..cb0ce2feaf 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1952,7 +1952,7 @@ class FastLlamaModel: # Offload! # [TODO] First offload lm_head and embed_tokens to CPU (should be disk!!) if "embed_tokens" in new_target_modules: - print("Unsloth: Casting embed_tokens to float32") + print("Unsloth: Training embed_tokens in mixed precision to save VRAM") model.model.model.embed_tokens.modules_to_save.default\ .to(device = "cuda:0", non_blocking = True) @@ -1965,7 +1965,7 @@ class FastLlamaModel: pass if "lm_head" in new_target_modules: - print("Unsloth: Casting lm_head to float32") + print("Unsloth: Training lm_head in mixed precision to save VRAM") model.model.lm_head.modules_to_save.default\ .to(device = "cuda:0", non_blocking = True) @@ -2203,7 +2203,7 @@ class FastLlamaModel: # Now patch lm_head and embed_tokens if train_embed_tokens: - print("Unsloth: Casting embed_tokens to float32") + print("Unsloth: Training embed_tokens in mixed precision to save VRAM") assert(hasattr(model.model.model.embed_tokens, "modules_to_save")) model.model.model.embed_tokens.modules_to_save.default\ .to(device = "cuda:0", non_blocking = True) @@ -2211,7 +2211,7 @@ class FastLlamaModel: pass if train_lm_head: - print("Unsloth: Casting lm_head to float32") + print("Unsloth: Training lm_head in mixed precision to save VRAM") assert(hasattr(model.model.lm_head, "modules_to_save")) model.model.lm_head.modules_to_save.default\ .to(device = "cuda:0", non_blocking = True)