diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 3752d46d0a..9d75fda16a 100644 --- a/unsloth/models/_utils.py +++ b/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 diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 0765e42892..4ffb18f681 100644 --- a/unsloth/models/llama.py +++ b/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