From 0f6b518ee1a50faa77ded37d96bb98212c8fc00d Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 6 Jan 2025 18:56:26 -0800
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 48 ++++++++++++++++++++++++++++++++++++-----
 1 file changed, 43 insertions(+), 5 deletions(-)

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