Auto change is_bfloat16_supported

unsloth/models/_utils.py

@@ -15,6 +15,10 @@
 __version__ = "2025.1.1"
 
 __all__ = [
+    "SUPPORTS_BFLOAT16",
+    "is_bfloat16_supported",
+    "USE_BFLOAT16",
+
     "prepare_model_for_kbit_training",
     "xformers",
     "xformers_attention",
@@ -30,7 +34,6 @@ __all__ = [
     "offload_to_disk",
     "offload_input_embeddings",
     "offload_output_embeddings",
-    "is_bfloat16_supported",
     "unsloth_offloaded_gradient_checkpoint",
     "torch_compile_options",
     "patch_linear_scaling",
@@ -773,9 +776,13 @@ def offload_output_embeddings(model, temporary_location : str = "_unsloth_tempor
 pass
 
 
+# Log dtype used - sometimes people use float16 on bfloat16 platforms
+global USE_BFLOAT16
+USE_BFLOAT16 = SUPPORTS_BFLOAT16
 # Fixes a weird Torch 2.3 bug which says T4s have bfloat16
 def is_bfloat16_supported():
-    return SUPPORTS_BFLOAT16
+    global USE_BFLOAT16
+    return SUPPORTS_BFLOAT16 and USE_BFLOAT16
 pass
 
 

unsloth/models/llama.py

@@ -68,6 +68,8 @@ pass
 from triton import __version__ as triton_version
 BlockDiagonalCausalMask = xformers.attn_bias.BlockDiagonalCausalMask if xformers is not None else None
 
+from ._utils import SUPPORTS_BFLOAT16, USE_BFLOAT16
+
 
 def original_apply_qkv(self, X):
     Q = self.q_proj(X)
@@ -1387,7 +1389,8 @@ class LongRopeRotaryEmbedding(torch.nn.Module):
         # self._set_cos_sin_cache(seq_len=self.current_rope_size, device=device, dtype=torch.get_default_dtype())
 
         # Short sequences
-        dtype = torch.bfloat16 if is_bfloat16_supported() else torch.float16
+        global USE_BFLOAT16
+        dtype = torch.bfloat16 if USE_BFLOAT16 else torch.float16
         t = torch.arange(original_max_position_embeddings, device=self.short_inv_freq.device, dtype=torch.int64).float()
         freqs = torch.outer(t, self.short_inv_freq)
         emb = torch.cat((freqs, freqs), dim=-1)
@@ -1580,7 +1583,6 @@ class FastLlamaModel:
         pass
         if token is None: token = get_token()
         if model_patcher is None: model_patcher = FastLlamaModel
-        SUPPORTS_BFLOAT16 = is_bfloat16_supported()
         gpu_stats = torch.cuda.get_device_properties(0)
         max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
 
@@ -1612,6 +1614,10 @@ class FastLlamaModel:
 
         assert(dtype == torch.float16 or dtype == torch.bfloat16 or dtype == torch.float32)
 
+        # Log global device type used
+        global USE_BFLOAT16
+        USE_BFLOAT16 = True if dtype == torch.bfloat16 else False
+
         # RoPE Scaling
         model_config = AutoConfig.from_pretrained(model_name, token = token)
         model_max_seq_length = model_config.max_position_embeddings