From 0beaf18908a2f275ebd553b5784a984db7b7b094 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Fri, 23 Feb 2024 01:50:51 +1100 Subject: [PATCH] Gemma --- unsloth/kernels/__init__.py | 4 +- unsloth/kernels/fast_lora.py | 46 +++-- unsloth/kernels/geglu.py | 104 ++++++++++ unsloth/kernels/rms_layernorm.py | 2 +- unsloth/models/gemma.py | 340 +++++++++++++++++++++++++++++++ unsloth/models/llama.py | 27 ++- unsloth/models/loader.py | 12 ++ unsloth/models/mistral.py | 9 +- 8 files changed, 511 insertions(+), 33 deletions(-) create mode 100644 unsloth/kernels/geglu.py create mode 100644 unsloth/models/gemma.py diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py index f5db8fa890..9c231e6ce1 100644 --- a/unsloth/kernels/__init__.py +++ b/unsloth/kernels/__init__.py @@ -16,9 +16,11 @@ from .cross_entropy_loss import fast_cross_entropy_loss from .rms_layernorm import fast_rms_layernorm from .rope_embedding import fast_rope_embedding, inplace_rope_embedding from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel +from .geglu import geglu_forward_kernel, geglu_backward_kernel from .fast_lora import ( get_lora_parameters, - apply_lora_mlp, + apply_lora_mlp_swiglu, + apply_lora_mlp_geglu, apply_lora_qkv, apply_lora_o, ) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index b3a1098355..3a71ffc587 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -14,7 +14,6 @@ import torch from .utils import fast_dequantize, QUANT_STATE, get_lora_parameters -from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel def matmul_lora(X, W, W_quant, A, B, s, out = None): @@ -85,20 +84,22 @@ class LoRA_MLP(torch.autograd.Function): def forward(ctx, X : torch.Tensor, gateW, gateW_quant, gateA, gateB, gateS, upW, upW_quant, upA, upB, upS, - downW, downW_quant, downA, downB, downS): + downW, downW_quant, downA, downB, downS, + _forward_function, _backward_function,): dtype = X.dtype e = matmul_lora(X, gateW, gateW_quant, gateA, gateB, gateS) g = matmul_lora(X, upW, upW_quant, upA, upB, upS) # f = torch.nn.functional.silu(e) # h = f * g - h = swiglu_fg_kernel(e, g) + h = _forward_function(e, g) i = matmul_lora(h, downW, downW_quant, downA, downB, downS) ctx.custom_saved_tensors = ( gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, + _backward_function, ) ctx.save_for_backward(gateA, gateB, upA, upB, downA, downB, X, e, g) @@ -109,8 +110,8 @@ class LoRA_MLP(torch.autograd.Function): @staticmethod @torch.cuda.amp.custom_bwd def backward(ctx, dY : torch.Tensor): - gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, = \ - ctx.custom_saved_tensors + gateW, gateW_quant, gateS, upW, upW_quant, upS, downW, downW_quant, downS, + _backward_function = ctx.custom_saved_tensors gateA, gateB, upA, upB, downA, downB, \ X, e, g = ctx.saved_tensors @@ -125,14 +126,7 @@ class LoRA_MLP(torch.autograd.Function): dtype = X.dtype DW = matmul_lora(dY, downW.t(), downW_quant, downB, downA, downS) - # e = e.float() - # se = 1.0 / (1.0 + torch.exp(-e)) - # f = (se * e).to(dtype) - # h = f * g - # df = DW * f - # dg = DW * g - # de = (dg.float() * se * (1.0 + e * (1.0 - se))).to(dtype) - DW, e, g = swiglu_DWf_DW_dfg_kernel(DW, e, g) + DW, e, g = _backward_function(DW, e, g) h, df, de = DW, e, g # Down projection LoRA weights @@ -155,7 +149,6 @@ class LoRA_MLP(torch.autograd.Function): # dX = matmul_lora(df, upW.t(), upW_quant, upB, upA, upS) # dX += matmul_lora(de, gateW.t(), gateW_quant, gateB, gateA, gateS) - upW = fast_dequantize(upW.t(), upW_quant) dX = torch.matmul(df, upW.t(), out = X) del upW @@ -177,19 +170,30 @@ class LoRA_MLP(torch.autograd.Function): pass -def apply_lora_mlp(self, X): - # gate = self.gate_proj(X) - # up = self. up_proj(X) - # h = torch.nn.functional.silu(gate) * up - # down = self.down_proj(h) - # return down +from .swiglu import swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel +def apply_lora_mlp_swiglu(self, X): gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj) upW, upW_quant, upA, upB, upS = get_lora_parameters(self. up_proj) downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj) out = LoRA_MLP.apply(X, gateW, gateW_quant, gateA, gateB, gateS, upW, upW_quant, upA, upB, upS, - downW, downW_quant, downA, downB, downS) + downW, downW_quant, downA, downB, downS, + swiglu_fg_kernel, swiglu_DWf_DW_dfg_kernel,) + return out +pass + + +from .geglu import geglu_forward_kernel, geglu_backward_kernel +def apply_lora_mlp_geglu(self, X): + gateW, gateW_quant, gateA, gateB, gateS = get_lora_parameters(self.gate_proj) + upW, upW_quant, upA, upB, upS = get_lora_parameters(self. up_proj) + downW, downW_quant, downA, downB, downS = get_lora_parameters(self.down_proj) + out = LoRA_MLP.apply(X, + gateW, gateW_quant, gateA, gateB, gateS, + upW, upW_quant, upA, upB, upS, + downW, downW_quant, downA, downB, downS, + geglu_forward_kernel, geglu_backward_kernel,) return out pass diff --git a/unsloth/kernels/geglu.py b/unsloth/kernels/geglu.py new file mode 100644 index 0000000000..7001b8ff0a --- /dev/null +++ b/unsloth/kernels/geglu.py @@ -0,0 +1,104 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import triton +import triton.language as tl +import torch +from .utils import calculate_settings + + +@triton.jit +def _forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,): + block_idx = tl.program_id(0) + offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + # f = 1/2 * e * (1 + erf(1/sqrt(2) * e)) + # h = f * up + e_row = tl.load(e + offsets, mask = mask, other = 0).to(tl.float32) + g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32) + + f_row = 0.5 * e_row * (tl.math.erf(tl.math.rsqrt(2.0) * e_row) + 1.0) + f_row = f_row.to(g_row.dtype) # Exact copy from HF + h_row = f_row * g_row + + # Store h + tl.store(h + offsets, h_row, mask = mask) +pass + + +def geglu_forward_kernel(gate, up): + batch, seq_len, hd = gate.shape + n_elements = gate.numel() + out = torch.empty((batch, seq_len, hd), dtype = gate.dtype, device = "cuda") + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + _forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,) + return out +pass + + +@triton.jit +def _backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,): + """ + f = 1/2 * e * (1 + erf(1/sqrt(2) * e)) + h = f * up + + df/de (with help of Wolfram :) + df/de = 1/2 * (1 + erf(1/sqrt(2) * e)) + 1/sqrt(2*pi) * e * exp(-1/2 * e^2) + + Reuse via + f = 1/2 * (1 + erf(1/sqrt(2) * e)) * e + """ + block_idx = tl.program_id(0) + offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + DW_row = tl.load(DW + offsets, mask = mask, other = 0)#.to(tl.float32) + e_row = tl.load(e + offsets, mask = mask, other = 0).to(tl.float32) + g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32) + + # Break e_row away for re-use + # f = 1/2 * e * (1 + erf(1/sqrt(2) * e)) + f_partial_row = 0.5 * (tl.math.erf(tl.math.rsqrt(2.0) * e_row) + 1.0) + f_row = f_partial_row * e_row + + f_row = f_row.to(DW_row.dtype) + # h = f * g + h_row = f_row * g_row + # df = DW * f + df_row = DW_row * f_row + # dg = DW * g + dg_row = DW_row * g_row + + # df/de = 1/2 * (1 + erf(1/sqrt(2) * e)) + 1/sqrt(2*pi) * e * exp(-1/2 * e^2) + t = 0.3989422804014327 # 1/sqrt(2*pi) + df_de = f_partial_row + t * e_row * tl.exp(-0.5 * e_row * e_row) + + de_row = dg_row.to(tl.float32) * df_de + de_row = de_row.to(DW_row.dtype) + + # Store derivatives in buffers + tl.store(DW + offsets, h_row, mask = mask) # h = f * g + tl.store(e + offsets, df_row, mask = mask) # df = DW * f + tl.store(g + offsets, de_row, mask = mask) # de +pass + + +def geglu_backward_kernel(DW, e, g): + batch_seq_len, hd = e.shape + n_elements = e.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + _backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) + return DW, e, g +pass diff --git a/unsloth/kernels/rms_layernorm.py b/unsloth/kernels/rms_layernorm.py index ec34880a2c..ccd9f89948 100644 --- a/unsloth/kernels/rms_layernorm.py +++ b/unsloth/kernels/rms_layernorm.py @@ -44,7 +44,7 @@ def _rms_layernorm_forward( W_row = tl.load(W + col_offsets, mask = mask, other = 0)#.to(tl.float32) row_var = tl.sum(X_row * X_row, axis = 0) / n_cols - inv_var = 1.0 / tl.sqrt(row_var + eps) + inv_var = tl.math.rsqrt(row_var + eps) tl.store(r, inv_var) normed = X_row * inv_var normed = normed.to(W_row.dtype) # Exact copy from HF diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py new file mode 100644 index 0000000000..439a7aecbb --- /dev/null +++ b/unsloth/models/gemma.py @@ -0,0 +1,340 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .llama import * +from ._utils import __version__ + +from transformers.models.gemma.modeling_gemma import ( + GemmaAttention, + GemmaDecoderLayer, + GemmaModel, + GemmaForCausalLM, +) +# For Pytorch 2.1.1 +try: + from transformers.models.gemma.modeling_gemma import ( + GemmaSdpaAttention, + GemmaFlashAttention2, + ) +except: + GemmaSdpaAttention = GemmaAttention + GemmaFlashAttention2 = GemmaAttention +pass + + +def fast_geglu_inference(self, X): + # gate = self.gate_proj(X) + # up = self.up_proj(X) + bsz, _, hd = X.shape + mlp_size = self.config.intermediate_size + temp = torch.empty((2, bsz, 1, mlp_size), dtype = X.dtype, device = "cuda") + + gate = fast_linear_forward(self.gate_proj, X, out = temp[0]) + up = fast_linear_forward(self. up_proj, X, out = temp[1]) + gate = torch.nn.functional.gelu(gate) + gate *= up + + # X = self.down_proj(gate) + down = fast_linear_forward(self.down_proj, gate, out = up[:,:,:hd]) + return down +pass + + +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L590 +def GemmaDecoderLayer_fast_forward( + self, + hidden_states: torch.Tensor, + causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + output_attentions: Optional[bool] = False, + use_cache: Optional[bool] = False, + padding_mask: Optional[torch.LongTensor] = None, + *args, **kwargs, +) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: + if past_key_value is not None: + do_prefill = not hasattr(self.self_attn, "paged_attention") + + # Self Attention + residual = hidden_states + hidden_states = fast_rms_layernorm_inference(self.input_layernorm, hidden_states) + hidden_states, present_key_value = LlamaAttention_fast_forward_inference( + self.self_attn, + hidden_states, + past_key_value, + position_ids, + do_prefill = do_prefill, + ) + hidden_states += residual + + # Fully Connected + residual = hidden_states + hidden_states = fast_rms_layernorm_inference(self.post_attention_layernorm, hidden_states) + hidden_states = fast_geglu_inference(self.mlp, hidden_states) + hidden_states += residual + else: + residual = hidden_states + hidden_states = fast_rms_layernorm(self.input_layernorm, hidden_states) + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + causal_mask=causal_mask, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + output_attentions=output_attentions, + use_cache=use_cache, + padding_mask=padding_mask, + ) + hidden_states = residual + hidden_states + + # Fully Connected + residual = hidden_states + hidden_states = fast_rms_layernorm(self.post_attention_layernorm, hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + pass + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights,) + + if use_cache: + outputs += (present_key_value,) + + return outputs +pass + + +from math import sqrt as math_sqrt + +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L825 +@torch.inference_mode +def GemmaModel_fast_forward_inference( + self, + input_ids, + past_key_values, +): + # Fix out of bounds tokenization + input_ids = input_ids[:,:self.max_seq_length] + + hidden_states = self.embed_tokens(input_ids) + hidden_states *= math_sqrt(self.config.hidden_size) + + next_decoder_cache = [] + for idx, decoder_layer in enumerate(self.layers): + # Self Attention + residual = hidden_states + hidden_states = fast_rms_layernorm_inference(decoder_layer.input_layernorm, hidden_states) + hidden_states, present_key_value = LlamaAttention_fast_forward_inference( + decoder_layer.self_attn, + hidden_states, + past_key_values[idx], + None, + ) + hidden_states += residual + + # Fully Connected + residual = hidden_states + hidden_states = fast_rms_layernorm_inference(decoder_layer.post_attention_layernorm, hidden_states) + hidden_states = fast_geglu_inference(decoder_layer.mlp, hidden_states) + hidden_states += residual + + next_decoder_cache.append(present_key_value) + pass + hidden_states = fast_rms_layernorm_inference(self.norm, hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state = hidden_states, + past_key_values = next_decoder_cache, + hidden_states = [], + attentions = [], + ) +pass + + +def GemmaForCausalLM_fast_forward( + self, + input_ids: torch.LongTensor = None, + causal_mask: Optional[xformers.attn_bias.BlockDiagonalCausalMask] = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[List[torch.FloatTensor]] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + labels: Optional[torch.LongTensor] = None, + use_cache: Optional[bool] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + *args, **kwargs, +) -> Union[Tuple, CausalLMOutputWithPast]: + + if causal_mask is None and past_key_values is None: + causal_mask = xformers.attn_bias.LowerTriangularMask() + + output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_hidden_states = ( + output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + self.model._has_no_labels = labels is None + + if past_key_values is not None and \ + hasattr(self.model.layers[0].self_attn, "paged_attention"): + outputs = GemmaModel_fast_forward_inference( + self.model, + input_ids, + past_key_values, + ) + else: + outputs = self.model( + input_ids=input_ids, + causal_mask=causal_mask, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + pass + + hidden_states = outputs[0] + bsz, q_len, hd = hidden_states.shape + if bsz == 1 and q_len == 1: + logits = torch.mv(self.lm_head.weight, hidden_states.ravel()) + logits = logits.unsqueeze(0).unsqueeze(0) + else: + logits = self.lm_head(hidden_states) + pass + + loss = None + if labels is not None: + shift_logits = logits + if not hasattr(self, "extra_ignored_labels"): + # Fixes https://github.com/unslothai/unsloth/issues/10 + self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda") + pass + + shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]])) + loss = fast_cross_entropy_loss( + logits = shift_logits, + labels = shift_labels, + ) + pass + + if not return_dict: + output = (logits,) + outputs[1:] + return (loss,) + output if loss is not None else output + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) +pass + + +class FastGemmaModel(FastLlamaModel): + + @staticmethod + def pre_patch(): + GemmaAttention .forward = LlamaAttention_fast_forward + GemmaSdpaAttention .forward = LlamaAttention_fast_forward + GemmaFlashAttention2.forward = LlamaAttention_fast_forward + GemmaDecoderLayer .forward = GemmaDecoderLayer_fast_forward + GemmaModel .forward = LlamaModel_fast_forward + GemmaForCausalLM .forward = GemmaForCausalLM_fast_forward + PeftModelForCausalLM.forward = PeftModelForCausalLM_fast_forward + + # Solves https://github.com/unslothai/unsloth/issues/168 + # Static KV Cache was introduced in 4.38.0, causing training to be much slower. + # Inferene can now be CUDAGraphed, but we shall retain the old rotary embeddings. + # https://github.com/huggingface/transformers/pull/27931 + # https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py + import transformers.models.gemma.modeling_gemma + transformers.models.gemma.modeling_gemma.GemmaRotaryEmbedding = LlamaRotaryEmbedding + return + pass + + + @staticmethod + def post_patch(model): + # Patch model for Gemma + layers = model.model.layers + + # Torch.compile fails on embedding matrix?? + # Workaround randomnly fixes it for torch versions < 2.2 + model.model.embed_tokens = torch.nn.Embedding.from_pretrained(model.model.embed_tokens.weight) + model.config.update({"unsloth_version" : __version__}) + + # We also do this for the lm_head + lm_head = torch.nn.Linear(1, 1, bias = None) + del lm_head.weight + lm_head.weight = model.lm_head.weight + lm_head.in_features = lm_head.weight.shape[1] + lm_head.out_features = lm_head.weight.shape[0] + model.lm_head = lm_head + + # Also patch all dtypes - BnB seems to not allocate the correct type? + # BnB default dtype seems to be float16! + correct_dtype = lm_head.weight.dtype + + for name, module in model.named_modules(): + if isinstance(module, (Bnb_Linear4bit, Peft_Linear4bit)): + weight = module.weight + quant_state = weight.quant_state + + if type(quant_state) is list: + # BnB seems to have float16 as default! + module.weight.quant_state[2] = correct_dtype # Cast to correct dtype + else: + # https://github.com/TimDettmers/bitsandbytes/pull/763/files + quant_state.dtype = correct_dtype + pass + pass + pass + + # Add 1 to weight + # return output * (1 + self.weight) + # https://github.com/huggingface/transformers/blob/main/src/transformers/models/gemma/modeling_gemma.py#L89 + from transformers.models.gemma.modeling_gemma import GemmaRMSNorm + + # Freeze all parameters except LoRA + for name, param in model.named_parameters(): + if ".lora_A." in name or ".lora_B." in name: + param.requires_grad_(True) + else: + param.requires_grad_(False) + pass + for name, module in model.named_modules(): + if isinstance(module, GemmaRMSNorm): + module.weight += 1.0 # return output * (1 + self.weight) + pass + + # Clear deleted GPU items + import gc + for _ in range(3): + gc.collect() + torch.cuda.empty_cache() + return model + pass +pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 3ca6291fd5..0c227321af 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -207,7 +207,7 @@ def LlamaAttention_fast_forward_inference( pass -def fast_mlp_inference(self, X): +def fast_swiglu_inference(self, X): # gate = self.gate_proj(X) # up = self.up_proj(X) bsz, _, hd = X.shape @@ -390,7 +390,7 @@ def LlamaDecoderLayer_fast_forward( # Fully Connected residual = hidden_states hidden_states = fast_rms_layernorm_inference(self.post_attention_layernorm, hidden_states) - hidden_states = fast_mlp_inference(self.mlp, hidden_states) + hidden_states = fast_swiglu_inference(self.mlp, hidden_states) hidden_states += residual else: residual = hidden_states @@ -507,6 +507,11 @@ def LlamaModel_fast_forward( if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) + # Mormalized from Gemma + if self.config.model_type == "gemma": + inputs_embeds *= math_sqrt(self.config.hidden_size) + pass + # Fix up attention mask by setting elements to 0 # Specifically for DPO if self._has_no_labels and (attention_mask is not None) and (past_key_values is None): @@ -646,7 +651,7 @@ def LlamaModel_fast_forward_inference( # Fully Connected residual = hidden_states hidden_states = fast_rms_layernorm_inference(decoder_layer.post_attention_layernorm, hidden_states) - hidden_states = fast_mlp_inference(decoder_layer.mlp, hidden_states) + hidden_states = fast_swiglu_inference(decoder_layer.mlp, hidden_states) hidden_states += residual next_decoder_cache.append(present_key_value) @@ -886,6 +891,7 @@ class FastLlamaModel: device_map = "sequential", rope_scaling = None, fix_tokenizer = True, + model_patcher = FastLlamaModel, **kwargs, ): SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported() @@ -893,13 +899,13 @@ class FastLlamaModel: max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) statistics = \ - f"==((====))== Unsloth: Fast Llama patching release {__version__}\n"\ + f"==((====))== Unsloth: Fast {model_patcher.__name__[4:-5]} patching release {__version__}\n"\ f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform = {platform_system}.\n"\ f"O^O/ \_/ \\ Pytorch: {torch.__version__}. CUDA = {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit = {torch.version.cuda}.\n"\ f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. Xformers = {xformers_version}. FA = {HAS_FLASH_ATTENTION}.\n"\ f' "-____-" Free Apache license: http://github.com/unslothai/unsloth' print(statistics) - FastLlamaModel.pre_patch() + model_patcher.pre_patch() if dtype is None: dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 @@ -955,7 +961,7 @@ class FastLlamaModel: ) model, tokenizer = patch_tokenizer(model, tokenizer) - model = FastLlamaModel.post_patch(model) + model = model_patcher.post_patch(model) # Patch up QKV / O and MLP for idx, layer in enumerate(model.model.layers): @@ -1309,6 +1315,15 @@ class FastLlamaModel: ) pass + # Get activation function + if model.config.mod == "swiglu": + apply_lora_mlp = apply_lora_mlp_swiglu + elif activation_function == "geglu": + apply_lora_mlp = apply_lora_mlp_geglu + else: + raise NotImplementedError(f"Unsloth: {activation_function} is not yet implemented!") + pass + model = prepare_model_for_kbit_training( model, use_gradient_checkpointing = use_gradient_checkpointing, diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index e4b3561deb..add8e9dca6 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -14,6 +14,7 @@ from .llama import FastLlamaModel, logger from .mistral import FastMistralModel +from .gemma import FastGemmaModel from transformers import AutoConfig from transformers import __version__ as transformers_version from peft import PeftConfig, PeftModel @@ -24,6 +25,7 @@ from .mapper import INT_TO_FLOAT_MAPPER, FLOAT_TO_INT_MAPPER major, minor = transformers_version.split(".")[:2] major, minor = int(major), int(minor) SUPPORTS_FOURBIT = (major > 4) or (major == 4 and minor >= 37) +SUPPORTS_GEMMA = (major > 4) or (major == 4 and minor >= 38) del major, minor @@ -99,6 +101,15 @@ class FastLanguageModel(FastLlamaModel): if model_type == "llama": dispatch_model = FastLlamaModel elif model_type == "mistral": dispatch_model = FastMistralModel + elif model_type == "gemma": + if not SUPPORTS_GEMMA: + raise RuntimeError( + f"Unsloth: Your transformers version of {transformers_version} does not support Gemma.\n"\ + f"The minimum required version is 4.38.\n"\ + f'Try `pip install --upgrade "transformers>=4.38"`\n'\ + f"to obtain the latest transformers build, then restart this session."\ + ) + dispatch_model = FastGemmaModel else: raise NotImplementedError( f"Unsloth: {model_name} not supported yet!\n"\ @@ -115,6 +126,7 @@ class FastLanguageModel(FastLlamaModel): device_map = device_map, rope_scaling = rope_scaling, fix_tokenizer = fix_tokenizer, + model_patcher = dispatch_model, *args, **kwargs, ) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 0e36023255..713205befb 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -293,6 +293,7 @@ class FastMistralModel(FastLlamaModel): device_map = "sequential", rope_scaling = None, # Mistral does not support RoPE scaling fix_tokenizer = True, + model_patcher = FastMistralModel, **kwargs, ): # Mistral does NOT support RoPE Scaling! @@ -305,13 +306,13 @@ class FastMistralModel(FastLlamaModel): max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3) statistics = \ - f"==((====))== Unsloth: Fast Mistral patching release {__version__}\n"\ + f"==((====))== Unsloth: Fast {model_patcher.__name__[4:-5]} patching release {__version__}\n"\ f" \\\ /| GPU: {gpu_stats.name}. Max memory: {max_memory} GB. Platform = {platform_system}.\n"\ f"O^O/ \_/ \\ Pytorch: {torch.__version__}. CUDA = {gpu_stats.major}.{gpu_stats.minor}. CUDA Toolkit = {torch.version.cuda}.\n"\ f"\ / Bfloat16 = {str(SUPPORTS_BFLOAT16).upper()}. Xformers = {xformers_version}. FA = {HAS_FLASH_ATTENTION}.\n"\ - f' "-____-" Apache 2 free license: http://github.com/unslothai/unsloth' + f' "-____-" Free Apache license: http://github.com/unslothai/unsloth' print(statistics) - FastMistralModel.pre_patch() + model_patcher.pre_patch() if dtype is None: dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 @@ -360,7 +361,7 @@ class FastMistralModel(FastLlamaModel): ) model, tokenizer = patch_tokenizer(model, tokenizer) - model = FastMistralModel.post_patch(model) + model = model_patcher.post_patch(model) # Patch up QKV / O and MLP for idx, layer in enumerate(model.model.layers):