diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index dcc887747f..adec68d46b 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -17,10 +17,16 @@ from .utils import fast_dequantize, QUANT_STATE from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel def get_lora_parameters(proj): - active_adapter = proj.active_adapters[0] if \ - hasattr(proj, "active_adapters") else proj.active_adapter + # For DPO or disabled adapters base_layer = (proj.base_layer if hasattr(proj, "base_layer") else proj) W = base_layer.weight + + if proj.disable_adapters or proj.merged: + return W, QUANT_STATE(W), None, None, None + pass + + active_adapter = proj.active_adapters[0] if \ + hasattr(proj, "active_adapters") else proj.active_adapter A = proj.lora_A [active_adapter].weight B = proj.lora_B [active_adapter].weight s = proj.scaling[active_adapter] @@ -31,7 +37,6 @@ pass def matmul_lora(X, W, W_quant, A, B, s, out = None): dtype = X.dtype W = fast_dequantize(W.t(), W_quant) - A, B = A.t(), B.t() if X.dim() == 3: batch, seq_len, d = X.shape @@ -43,7 +48,13 @@ def matmul_lora(X, W, W_quant, A, B, s, out = None): out = torch.matmul(X, W, out = out) if W_quant is not None: del W - out += (X @ A.to(dtype)) @ (s * B.to(dtype)) + + if A is not None: + # LoRA is enabled + A, B = A.t(), B.t() + out += (X @ A.to(dtype)) @ (s * B.to(dtype)) + pass + return out.view(batch, seq_len, -1) if reshape else out pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index f6cc078356..22179d4410 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -369,8 +369,21 @@ def LlamaModel_fast_forward( raise ValueError("Unsloth: You have to specify either decoder_input_ids or decoder_inputs_embeds") seq_length_with_past = seq_length + + # Fix out of bounds tokenization if hasattr(self, "max_seq_length"): - assert(seq_length <= self.max_seq_length) + if seq_length > self.max_seq_length: + logger.warning_once( + f"Unsloth: Input IDs of length {seq_length} > the model's max sequence length of {self.max_seq_length}.\n"\ + "We shall truncate it ourselves. It's imperative if you correct this issue first." + ) + if input_ids is not None: + input_ids = input_ids[:,:self.max_seq_length] + elif inputs_embeds is not None: + inputs_embeds = inputs_embeds[:,:self.max_seq_length,:] + pass + pass + past_key_values_length = 0 if past_key_values is not None: diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index e4bac44694..bb6a68f59e 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -97,21 +97,36 @@ def MistralAttention_fast_forward( Q = Q.transpose(1, 2) K = K.transpose(1, 2) V = V.transpose(1, 2) + M = bsz * q_len + + has_sliding_window = isinstance(causal_mask, xformers.attn_bias.BlockDiagonalCausalMask) # Group query attention - if n_groups != 1: - K = K .view(bsz, q_len, n_kv_heads, 1, head_dim) - V = V .view(bsz, q_len, n_kv_heads, 1, head_dim) - K = K.expand(bsz, q_len, n_kv_heads, n_groups, head_dim) - V = V.expand(bsz, q_len, n_kv_heads, n_groups, head_dim) - if hidden_states.requires_grad: - # Xformers does not support backward, so we have to convert - # GQA to MQA by cloning K and V - K = K.reshape(bsz, q_len, n_heads, head_dim) # A copy will be made - V = V.reshape(bsz, q_len, n_heads, head_dim) # A copy will be made - else: - # Xformers does support the forward pass though - Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim) + # if n_groups != 1: + K = K .view(bsz, q_len, n_kv_heads, 1, head_dim) + V = V .view(bsz, q_len, n_kv_heads, 1, head_dim) + K = K.expand(bsz, q_len, n_kv_heads, n_groups, head_dim) + V = V.expand(bsz, q_len, n_kv_heads, n_groups, head_dim) + if hidden_states.requires_grad: + # Xformers does not support backward, so we have to convert + # GQA to MQA by cloning K and V + K = K.reshape(bsz, q_len, n_heads, head_dim) # A copy will be made + V = V.reshape(bsz, q_len, n_heads, head_dim) # A copy will be made + + if has_sliding_window: + Q = Q.view(1, M, n_heads, head_dim) + K = K.view(1, M, n_heads, head_dim) + V = V.view(1, M, n_heads, head_dim) + pass + else: + # Xformers does support the forward pass though + Q = Q.view(bsz, q_len, n_kv_heads, n_groups, head_dim) + + if has_sliding_window: + Q = Q.view(1, M, n_kv_heads, n_groups, head_dim) + K = K.view(1, M, n_kv_heads, n_groups, head_dim) + V = V.view(1, M, n_kv_heads, n_groups, head_dim) + pass pass A = xformers_attention(Q, K, V, attn_bias = causal_mask) @@ -131,12 +146,12 @@ def MistralAttention_fast_forward( A = flash_attn_func(Q, K, V, causal = True, window_size = window) else: # Grouped query attention - if n_groups != 1: - 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) - pass + # if n_groups != 1: + 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) + # pass # Needs (batch_size, n_heads, seq_len, head_dim) # is_casual and attention_mask must not be both set! A = scaled_dot_product_attention(Q, K, V, attn_mask = attention_mask, is_causal = False)