From 1576396cd0d124ee055b5c8bfbf0d40cfdad498f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 16 Jan 2025 03:09:02 -0800 Subject: [PATCH] Bug fixes --- unsloth/models/_utils.py | 8 ++++++-- unsloth/models/mistral.py | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index ff2c8726e4..2c16bf6e72 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -285,7 +285,11 @@ if major_version >= 8: if _is_package_available("flash_attn"): # Check for CUDA linking errors "undefined symbol: _ZNK3c106SymIntltEl" try: - from flash_attn.flash_attn_interface import flash_attn_cuda + try: + # See https://github.com/unslothai/unsloth/issues/1437 + from flash_attn.flash_attn_interface import flash_attn_gpu + except: + from flash_attn.flash_attn_interface import flash_attn_cuda HAS_FLASH_ATTENTION = True # Also check for softcapping @@ -799,7 +803,7 @@ def patch_linear_scaling( f"from typing import Union, Optional, List, Any, Callable, Tuple\n"\ f"from {model_filepath} import logger, "\ f"{model_name.title()}Attention, {model_name.title()}Config" - + try: function = inspect.getsource(attention_module.__init__) except: diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 9a97015f9b..784ca9cb41 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -304,7 +304,7 @@ class FastMistralModel(FastLlamaModel): attention_module = MistralAttention, ) # Just for Mistral Nemo models! - if function is not None: + if function is not None and init_name is not None: function = patch_mistral_nemo_attention(function) # if True:#init_name is not None: exec(function, globals())