From 9617ecbbb8bef7864961096f925f03e40ffa7f99 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 30 Jul 2024 22:48:53 -0700 Subject: [PATCH] flash-attn softcapping --- pyproject.toml | 22 ++++++------- unsloth/models/_utils.py | 14 +++++++++ unsloth/models/gemma2.py | 68 +++++++++++++++++++++++++++++----------- unsloth/models/llama.py | 37 ++++++++++++---------- unsloth/models/loader.py | 16 ++++++++++ unsloth/models/mapper.py | 1 - 6 files changed, 112 insertions(+), 46 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 6777f7c26e..e711325be9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,7 +171,7 @@ colab-ampere-torch211 = [ "unsloth[cu121onlytorch211]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] colab-torch220 = [ "unsloth[huggingface]", @@ -184,7 +184,7 @@ colab-ampere-torch220 = [ "unsloth[cu121onlytorch220]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] colab-new = [ "packaging", @@ -215,7 +215,7 @@ colab-ampere = [ "unsloth[colab-ampere-torch220]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] cu118-ampere = [ "unsloth[huggingface]", @@ -223,7 +223,7 @@ cu118-ampere = [ "unsloth[cu118only]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] cu121-ampere = [ "unsloth[huggingface]", @@ -231,7 +231,7 @@ cu121-ampere = [ "unsloth[cu121only]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] cu118-ampere-torch211 = [ "unsloth[huggingface]", @@ -239,7 +239,7 @@ cu118-ampere-torch211 = [ "unsloth[cu118onlytorch211]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] cu121-ampere-torch211 = [ "unsloth[huggingface]", @@ -247,7 +247,7 @@ cu121-ampere-torch211 = [ "unsloth[cu121onlytorch211]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] cu118-ampere-torch220 = [ "unsloth[huggingface]", @@ -255,7 +255,7 @@ cu118-ampere-torch220 = [ "unsloth[cu118onlytorch220]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] cu121-ampere-torch220 = [ "unsloth[huggingface]", @@ -263,7 +263,7 @@ cu121-ampere-torch220 = [ "unsloth[cu121onlytorch220]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] cu118-ampere-torch230 = [ "unsloth[huggingface]", @@ -271,7 +271,7 @@ cu118-ampere-torch230 = [ "unsloth[cu118onlytorch230]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] cu121-ampere-torch230 = [ "unsloth[huggingface]", @@ -279,7 +279,7 @@ cu121-ampere-torch230 = [ "unsloth[cu121onlytorch230]", "packaging", "ninja", - "flash-attn", + "flash-attn>=2.6.3", ] [project.urls] diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index f4e4257b2b..c9bc6065f9 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -21,6 +21,7 @@ __all__ = [ "xformers_version", "__version__", "HAS_FLASH_ATTENTION", + "HAS_FLASH_ATTENTION_SOFTCAPPING", "PRE_CHECK", "platform_system", "patch_tokenizer", @@ -140,6 +141,8 @@ from transformers.utils.import_utils import _is_package_available major_version, minor_version = torch.cuda.get_device_capability() SUPPORTS_BFLOAT16 = False +HAS_FLASH_ATTENTION = False +HAS_FLASH_ATTENTION_SOFTCAPPING = False if major_version >= 8: SUPPORTS_BFLOAT16 = True @@ -148,6 +151,17 @@ if major_version >= 8: try: from flash_attn.flash_attn_interface import flash_attn_cuda HAS_FLASH_ATTENTION = True + + # Also check for softcapping + from flash_attn import __version__ as flash_attn_version + HAS_FLASH_ATTENTION_SOFTCAPPING = Version(flash_attn_version) >= Version("2.6.3") + if not HAS_FLASH_ATTENTION_SOFTCAPPING: + print( + "Unsloth: If you want to finetune Gemma 2, upgrade flash-attn to version 2.6.3 or higher!\n"\ + "Newer versions support faster and less memory usage kernels for Gemma 2's attention softcapping!\n"\ + "To update flash-attn, do the below:\n"\ + '\npip install --no-deps --upgrade "flash-attn>=2.6.3"' + ) except: print( "Unsloth: Your Flash Attention 2 installation seems to be broken?\n"\ diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 0d21c47b00..ecd45fbce2 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -56,6 +56,8 @@ except: Gemma2FlashAttention2 = Gemma2Attention pass +if HAS_FLASH_ATTENTION_SOFTCAPPING: + from flash_attn import flash_attn_func # [TODO] We must randomnly use torch.compile? # I checked the gradients and formulas and I'm sure it's correct. @@ -126,8 +128,31 @@ def Gemma2Attention_fast_forward( V = torch.cat([past_key_value[1], V], dim = 2) pass past_key_value = (K, V) if use_cache else None - - A = slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, kv_seq_len) + + # Only enable if the attention_mask is True + has_sliding_window = type(attention_mask) is bool and attention_mask is True + if HAS_FLASH_ATTENTION_SOFTCAPPING and type(attention_mask) is bool: + window = (-1, -1) + if has_sliding_window: + sw = getattr(self.config, "sliding_window", None) + sw = kv_seq_len if (sw is None or sw == "null") else sw + window = (-1, -1) if (kv_seq_len <= sw) else (sw, sw) + pass + + Q = Q.transpose(1, 2) + K = K.transpose(1, 2) + V = V.transpose(1, 2) + A = flash_attn_func( + Q, K, V, + causal = True, + softcap = self.config.attn_logit_softcapping, + softmax_scale = self.config.query_pre_attn_scalar, + window_size = window, + ) + A = A.reshape(bsz, q_len, n_heads*head_dim) + else: + A = slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, kv_seq_len) + pass A = self.apply_o(self, A) return A, None, past_key_value pass @@ -205,6 +230,8 @@ pass from math import sqrt as math_sqrt KV_CACHE_INCREMENT = 256 # KV Cache update size torch_nn_functional_softmax = torch.nn.functional.softmax +torch_matmul = torch.matmul +torch_tanh = torch.tanh def Gemma2Attention_fast_forward_inference( self, @@ -322,13 +349,13 @@ def Gemma2Attention_fast_forward_inference( # if bsz == 1: Qn *= self.scalar # See https://github.com/ggerganov/llama.cpp/issues/7805#issuecomment-2153349963 # It seems like doing (Q * scalar) @ K is better than (Q @ K) * scalar to stop overflows - A = torch.matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len]) + A = torch_matmul(Qn, Knn.transpose(2, 3), out = self.attention[:,:,:,:cached_len]) # if attention_mask is not None: A += attention_mask # Must add attention_mask for batched - A *= self.reciprocal_t; torch.tanh(A, out = A); A *= self.t; # Logit softcapping + A *= self.reciprocal_t; torch_tanh(A, out = A); A *= self.t; # Logit softcapping A[:] = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype) - A = torch.matmul(A, Vnn, out = Qn) + A = torch_matmul(A, Vnn, out = Qn) # else: # A = scaled_dot_product_attention(Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = False) # pass @@ -359,19 +386,24 @@ def Gemma2Model_fast_forward_inference( bsz, q_len, hd = hidden_states.shape seq_len = past_key_values[0][0].shape[-2] if bsz != 1: - SWA = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, - (bsz, q_len), - hidden_states, - seq_len, - sliding_window = self.config.sliding_window, - ) - GA = _prepare_4d_causal_attention_mask_for_sdpa( - attention_mask, - (bsz, q_len), - hidden_states, - seq_len, - ) + if HAS_FLASH_ATTENTION_SOFTCAPPING: + SWA = True + GA = False + else: + SWA = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (bsz, q_len), + hidden_states, + seq_len, + sliding_window = self.config.sliding_window, + ) + GA = _prepare_4d_causal_attention_mask_for_sdpa( + attention_mask, + (bsz, q_len), + hidden_states, + seq_len, + ) + pass else: SWA = attention_mask GA = attention_mask diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 496a37e7a3..b5244ed4ee 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -682,23 +682,28 @@ def LlamaModel_fast_forward( # Gemma2 has alternating SWA and global attn if IS_GEMMA2 and not hasattr(self, "SWA_mask"): - n = self.config.max_position_embeddings - # masked_fill is making stuff slower! - # self. GA_mask = create_boolean_mask(n = n, sliding_window = 0) - # self.SWA_mask = create_boolean_mask(n = n, sliding_window = self.config.sliding_window) - from transformers.modeling_attn_mask_utils import AttentionMaskConverter - self.SWA_mask = AttentionMaskConverter( - is_causal = True, - sliding_window = self.config.sliding_window, - )\ - .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda:0",)\ - .squeeze(0).squeeze(0) + if HAS_FLASH_ATTENTION_SOFTCAPPING: + self.SWA_mask = True + self.GA_mask = False + else: + n = self.config.max_position_embeddings + # masked_fill is making stuff slower! + # self. GA_mask = create_boolean_mask(n = n, sliding_window = 0) + # self.SWA_mask = create_boolean_mask(n = n, sliding_window = self.config.sliding_window) + from transformers.modeling_attn_mask_utils import AttentionMaskConverter + self.SWA_mask = AttentionMaskConverter( + is_causal = True, + sliding_window = self.config.sliding_window, + )\ + .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda:0",)\ + .squeeze(0).squeeze(0) - self.GA_mask = AttentionMaskConverter( - is_causal = True, - )\ - .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda:0",)\ - .squeeze(0).squeeze(0) + self.GA_mask = AttentionMaskConverter( + is_causal = True, + )\ + .to_causal_4d(1, n, n, dtype = inputs_embeds.dtype, device = "cuda:0",)\ + .squeeze(0).squeeze(0) + pass pass # Go through every layer! diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index fb6d5c5018..47152d6764 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from ._utils import is_bfloat16_supported, HAS_FLASH_ATTENTION, HAS_FLASH_ATTENTION_SOFTCAPPING from .llama import FastLlamaModel, logger from .mistral import FastMistralModel from .qwen2 import FastQwen2Model @@ -233,6 +234,21 @@ class FastLanguageModel(FastLlamaModel): f'Try `pip install --upgrade "transformers>=4.42.3"`\n'\ f"to obtain the latest transformers build, then restart this session."\ ) + # Also check for softcapping support in flash-attn which is faster! + if is_bfloat16_supported() and not HAS_FLASH_ATTENTION: + print( + "Unsloth: If you want to finetune Gemma 2, install flash-attn to make it faster!\n"\ + "To install flash-attn, do the below:\n"\ + '\npip install --no-deps --upgrade "flash-attn>=2.6.3"' + ) + elif HAS_FLASH_ATTENTION and not HAS_FLASH_ATTENTION_SOFTCAPPING: + print( + "Unsloth: If you want to finetune Gemma 2, upgrade flash-attn to version 2.6.3 or higher!\n"\ + "Newer versions support faster and less memory usage kernels for Gemma 2's attention softcapping!\n"\ + "To update flash-attn, do the below:\n"\ + '\npip install --no-deps --upgrade "flash-attn>=2.6.3"' + ) + dispatch_model = FastGemma2Model elif model_type == "qwen2": dispatch_model = FastQwen2Model diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index 254a68a420..462555f317 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -67,7 +67,6 @@ __INT_TO_FLOAT_MAPPER = \ "codellama/CodeLlama-7b-hf", ), "unsloth/codellama-13b-bnb-4bit" : ( - "unsloth/codellama-13b", "codellama/CodeLlama-13b-hf", ), "unsloth/yi-6b-bnb-4bit" : (