From 1d1d36de469613a7542c2fd1fb176a19d7f0d881 Mon Sep 17 00:00:00 2001
From: Edd <68678137+Erland366@users.noreply.github.com>
Date: Mon, 28 Oct 2024 02:06:45 +0400
Subject: [PATCH] Fix/casting continue pretraining (#1200)

* 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
---
 unsloth/models/llama.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

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