ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

flash-attn softcapping

Browse source
This commit is contained in:
Daniel Han 2024-07-30 22:48:53 -07:00
commit 9617ecbbb8
6 changed files with 113 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

22
pyproject.toml
View file

@ -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]

14
unsloth/models/_utils.py
View file

@ -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"\

68
unsloth/models/gemma2.py
View file

@ -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

37
unsloth/models/llama.py
View file

@ -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!

16
unsloth/models/loader.py
View file

@ -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

1
unsloth/models/mapper.py
View file

@ -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" : (