diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py index c2de979a6f..26f632ee12 100644 --- a/unsloth/kernels/__init__.py +++ b/unsloth/kernels/__init__.py @@ -33,7 +33,11 @@ from .fast_lora import ( ) from .utils import fast_dequantize, fast_gemv, QUANT_STATE, fast_linear_forward, matmul_lora -from .flex_attention import HAS_FLEX_ATTENTION, slow_attention_softcapping +from .flex_attention import ( + HAS_FLEX_ATTENTION, + slow_attention_softcapping, + slow_inference_attention_softcapping, +) if HAS_FLEX_ATTENTION: from .flex_attention import ( diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py index a992a02382..9cf999e2b7 100644 --- a/unsloth/kernels/flex_attention.py +++ b/unsloth/kernels/flex_attention.py @@ -80,3 +80,40 @@ def slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len): return A pass + +torch_matmul = torch.matmul +torch_tanh = torch.tanh +torch_nn_functional_softmax = torch.nn.functional.softmax +def slow_inference_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len): + n_heads = self.num_heads + head_dim = self.head_dim + n_kv_heads = self.num_key_value_heads + n_groups = self.num_key_value_groups + + # Grouped query attention + K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim) + V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim) + K = K.reshape(bsz, n_heads, q_len, head_dim) + V = V.reshape(bsz, n_heads, q_len, head_dim) + + # See https://github.com/google/gemma_pytorch/commit/03e657582d17cb5a8617ebf333c1c16f3694670e + # Gemma 9b should use 256 and not 224 (hs / nah). 27b uses the below + # We default to using the config file itself + # s = self.config.hidden_size // self.config.num_attention_heads + s = self.config.query_pre_attn_scalar + t = self.config.attn_logit_softcapping + + Q = Q * torch.tensor(s**-0.5, dtype = Q.dtype) # Follow Keras exactly + A = torch_matmul(Q, K.transpose(2, 3)) + + # Logit softcapping + A /= t; torch_tanh(A, out = A); A *= t; + A += causal_mask[:q_len, :q_len] + # Much slower in torch compile! + # A.masked_fill_(causal_mask[:q_len, :q_len], -float("inf")) + A = torch_nn_functional_softmax(A, dim = -1, dtype = torch.float32).to(Q.dtype) + A = torch_matmul(A, V) + A = A.transpose(1, 2).contiguous() + A = A.reshape(bsz, q_len, n_heads*head_dim) + return A +pass diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index ea9a0c53db..242d234db8 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -39,6 +39,8 @@ __all__ = [ "create_boolean_mask", "torch_amp_custom_fwd", "torch_amp_custom_bwd", + "accelerate_old_send_to_device", + "accelerate_new_send_to_device", ] import torch @@ -287,6 +289,7 @@ if Version(xformers_version) >= Version("0.0.27"): import accelerate.utils.operations if hasattr(accelerate.utils.operations, "send_to_device") and \ accelerate.utils.operations.send_to_device.__name__ != "_fixed_send_to_device": + accelerate_old_send_to_device = accelerate.utils.operations.send_to_device from accelerate.utils.operations import * send_to_device = inspect.getsource(accelerate.utils.operations.send_to_device) send_to_device = re.sub( @@ -296,6 +299,7 @@ if Version(xformers_version) >= Version("0.0.27"): ).replace("def send_to_device", "def _fixed_send_to_device") exec(send_to_device) # accelerate.utils.operations.send_to_device = _fixed_send_to_device + accelerate_new_send_to_device = _fixed_send_to_device pass pass # ============================================= diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 6858f52573..218849ef2f 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -157,7 +157,10 @@ def Gemma2Attention_fast_forward( A = A.reshape(bsz, q_len, n_heads*head_dim) else: mask = causal_mask if attention_mask is None else attention_mask - A = slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, kv_seq_len) + fx = slow_inference_attention_softcapping \ + if "_flag_for_generation" in kwargs else \ + slow_attention_softcapping + A = fx(Q, K, V, causal_mask, self, bsz, kv_seq_len) pass A = self.apply_o(self, A) return A, None, past_key_value @@ -192,6 +195,7 @@ def Gemma2DecoderLayer_fast_forward( output_attentions=output_attentions, use_cache=use_cache, padding_mask=padding_mask, + _flag_for_generation=True, ) hidden_states = fast_rms_layernorm_inference_gemma(self.post_attention_layernorm, hidden_states, out_weight) hidden_states += residual diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 5ccf906acb..3fcb8a76d2 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -953,6 +953,8 @@ def CausalLM_fast_forward(fast_forward_inference): if bsz == 1 and q_len == 1: logits = torch.mv(lm_head, hidden_states.ravel().to(lm_head.dtype)) logits = logits.unsqueeze(0).unsqueeze(0) + elif num_logits_to_keep != 0: + logits = self.lm_head(hidden_states[:, -num_logits_to_keep:, :].to(lm_head.dtype)) else: logits = self.lm_head(hidden_states.to(lm_head.dtype)) pass @@ -1368,8 +1370,14 @@ def _wrap_fast_inference(generate, device_type, dtype, model): pass internal_model._flag_for_generation = True + # Must patch accelerate for Xformers + import accelerate.utils.operations + accelerate.utils.operations.send_to_device = accelerate_new_send_to_device + # For newer HF kwargs["cache_implementation"] = "dynamic" + # For num_logits_to_keep + kwargs["num_logits_to_keep"] = 1 # Remove token_type_ids kwargs.pop("token_type_ids", None) @@ -1402,6 +1410,9 @@ def _wrap_fast_inference(generate, device_type, dtype, model): pass if hasattr(internal_model, "_flag_for_generation"): del internal_model._flag_for_generation + # Return accelerate back + accelerate.utils.operations.send_to_device = accelerate_old_send_to_device + return output pass return _fast_generate