diff --git a/pyproject.toml b/pyproject.toml index 7e8956c712..385366a391 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,8 +33,8 @@ exclude = ["images*"] [project.optional-dependencies] huggingface = [ - "transformers>=4.38.0", - "datasets", + "transformers>=4.38.2", + "datasets>=2.16.0", "sentencepiece", "accelerate>=0.26.1", "trl>=0.7.9", @@ -64,6 +64,16 @@ cu121onlytorch211 = [ "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'", "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'", ] +cu118onlytorch212 = [ + "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'", + "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'", + "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.23.post1%2Bcu118-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'", +] +cu121onlytorch212 = [ + "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'", + "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'", + "xformers @ https://download.pytorch.org/whl/cu121/xformers-0.0.23.post1-cp311-cp311-manylinux2014_x86_64.whl ; python_version=='3.11'", +] cu118onlytorch220 = [ "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp39-cp39-manylinux2014_x86_64.whl ; python_version=='3.9'", "xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.24%2Bcu118-cp310-cp310-manylinux2014_x86_64.whl ; python_version=='3.10'", diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py index 9c231e6ce1..8ff255e4a4 100644 --- a/unsloth/kernels/__init__.py +++ b/unsloth/kernels/__init__.py @@ -16,11 +16,17 @@ 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 .geglu import ( + geglu_exact_forward_kernel, + geglu_exact_backward_kernel, + geglu_approx_forward_kernel, + geglu_approx_backward_kernel, +) from .fast_lora import ( get_lora_parameters, apply_lora_mlp_swiglu, - apply_lora_mlp_geglu, + apply_lora_mlp_geglu_exact, + apply_lora_mlp_geglu_approx, apply_lora_qkv, apply_lora_o, ) diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index 3ed0d3c914..6568bba681 100644 --- a/unsloth/kernels/fast_lora.py +++ b/unsloth/kernels/fast_lora.py @@ -183,8 +183,8 @@ def apply_lora_mlp_swiglu(self, X): pass -from .geglu import geglu_forward_kernel, geglu_backward_kernel -def apply_lora_mlp_geglu(self, X): +from .geglu import geglu_exact_forward_kernel, geglu_exact_backward_kernel +def apply_lora_mlp_geglu_exact(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) @@ -192,7 +192,21 @@ def apply_lora_mlp_geglu(self, 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,) + geglu_exact_forward_kernel, geglu_exact_backward_kernel,) + return out +pass + + +from .geglu import geglu_approx_forward_kernel, geglu_approx_backward_kernel +def apply_lora_mlp_geglu_approx(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_approx_forward_kernel, geglu_approx_backward_kernel,) return out pass diff --git a/unsloth/kernels/geglu.py b/unsloth/kernels/geglu.py index 7001b8ff0a..df80fcb79b 100644 --- a/unsloth/kernels/geglu.py +++ b/unsloth/kernels/geglu.py @@ -19,7 +19,7 @@ from .utils import calculate_settings @triton.jit -def _forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,): +def _exact_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 @@ -38,18 +38,18 @@ def _forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,): pass -def geglu_forward_kernel(gate, up): +def geglu_exact_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,) + _exact_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,): +def _exact_backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,): """ f = 1/2 * e * (1 + erf(1/sqrt(2) * e)) h = f * up @@ -95,10 +95,109 @@ def _backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,): pass -def geglu_backward_kernel(DW, e, g): +def geglu_exact_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,) + _exact_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) + return DW, e, g +pass + + +@triton.jit +def _approx_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 + tanh( sqrt(2/pi) * (x + 0.044715 * x^3 ) )) + # f = 1/2 * e * (1 + tanh( sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) )) + # h = f * up + s = 0.7978845608028654 # math.sqrt(2 / math.pi) + + 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.tanh(s * e_row * (1.0 + 0.044715 * e_row * 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_approx_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']),) + _approx_forward_kernel[grid](gate, up, out, n_elements, BLOCK_SIZE = 1024,) + return out +pass + + +@triton.jit +def _approx_backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,): + """ + f = 1/2 * e * (1 + tanh( sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) )) + h = f * up + + df/de (with help from https://arxiv.org/pdf/2305.12073.pdf :)) + df/de = 1/2 * [1 + tanh( sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) )] + + 1/2 * sech^2 [ sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) ] * \ + ( sqrt(2/pi) * x * (1 + 0.044715 * x^2 * 3 ) ) + + Notice sech^2(x) = 1 - tanh^2(x) + So reuse tanh( sqrt(2/pi) * x * (1 + 0.044715 * x^2 ) ) + + See https://www.desmos.com/calculator/nqprfoni6x + """ + 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) + + # See https://www.desmos.com/calculator/nqprfoni6x + s = 0.7978845608028654 # math.sqrt(2 / math.pi) + a = s * e_row # a = sqrt(2 / pi) * x + b = a * 0.044715 * e_row * e_row # b = a * 0.044715 * x^2 + T = 1.0 + tl.math.tanh(a + b) + T2 = 0.5 * T + # Q = 0.5 * -T * (T - 2.0) * (a + 3.0 * b) + Q2 = -T2 * (T - 2.0) * (a + 3.0 * b) + df_de = T2 + Q2 # 1/2 * (T + Q) + + # f = 1/2 * e * (1 + tanh( sqrt(2/pi) * (x + 0.044715 * x^3 ) )) + f_row = T2 * 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 + + 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_approx_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']),) + _approx_backward_kernel[grid](DW, e, g, n_elements, BLOCK_SIZE = 1024,) return DW, e, g pass diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index e92be5e993..3e3b8ff22c 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -25,7 +25,7 @@ from platform import system as platform_system platform_system = platform_system() import math -__version__ = "2024.2" +__version__ = "2024.3" # Get Flash Attention v2 if Ampere (RTX 30xx, A100) major_version, minor_version = torch.cuda.get_device_capability() diff --git a/unsloth/models/dpo.py b/unsloth/models/dpo.py index 3ae4d636f4..b7c7305bb3 100644 --- a/unsloth/models/dpo.py +++ b/unsloth/models/dpo.py @@ -12,11 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. -from transformers.utils.notebook import ( - IntervalStrategy, - NotebookTrainingTracker, - NotebookProgressCallback, -) +try: + from transformers.utils.notebook import ( + IntervalStrategy, + NotebookTrainingTracker, + NotebookProgressCallback, + ) + HAS_NOTEBOOK = True +except: + HAS_NOTEBOOK = False +pass DPOTrainer_metrics = [ "rewards/chosen", @@ -101,13 +106,15 @@ pass def PatchDPOTrainer(): - from transformers.trainer import is_in_notebook - if is_in_notebook(): - # Patch DPO notebook printing - NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line - from transformers.trainer import DEFAULT_PROGRESS_CALLBACK - DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin - DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log + if HAS_NOTEBOOK: + from transformers.trainer import is_in_notebook + if is_in_notebook(): + # Patch DPO notebook printing + NotebookTrainingTracker.write_line = NotebookTrainingTracker_write_line + from transformers.trainer import DEFAULT_PROGRESS_CALLBACK + DEFAULT_PROGRESS_CALLBACK.on_train_begin = NotebookProgressCallback_on_train_begin + DEFAULT_PROGRESS_CALLBACK.on_log = NotebookProgressCallback_on_log + pass pass pass diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index 4aa634a4bd..97da833c12 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -48,7 +48,7 @@ def fast_geglu_inference(self, X): 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 = torch.nn.functional.gelu(gate, approximate = "tanh") gate *= up # X = self.down_proj(gate) @@ -70,7 +70,7 @@ def GemmaDecoderLayer_fast_forward( padding_mask: Optional[torch.LongTensor] = None, *args, **kwargs, ): - if False:#past_key_value is not None: + if past_key_value is not None: do_prefill = not hasattr(self.self_attn, "paged_attention") # Self Attention @@ -267,6 +267,9 @@ class FastGemmaModel(FastLlamaModel): # Patch RMS Layernorm for name, module in model.named_modules(): if isinstance(module, GemmaRMSNorm): + # Must be in float32 + # https://github.com/keras-team/keras-nlp/blob/v0.8.2/keras_nlp/models/gemma/rms_normalization.py#L36 + module = module.to(torch.float32) module.weight += 1.0 # return output * (1 + self.weight) if not hasattr(module, "variance_epsilon"): module.variance_epsilon = module.eps # Gemma doesn't use variance_epsilon diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 54e016d1af..20552644d9 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -511,7 +511,12 @@ def LlamaModel_fast_forward( # Mormalized from Gemma if self.config.model_type == "gemma": inputs_requires_grad = inputs_embeds.requires_grad - if inputs_requires_grad: inputs_embeds.requires_grad_(False) + if not inputs_embeds.is_leaf: + inputs_embeds = inputs_embeds.detach() + inputs_requires_grad = True + elif inputs_requires_grad: + inputs_embeds.requires_grad_(False) + pass inputs_embeds *= math_sqrt(self.config.hidden_size) if inputs_requires_grad: inputs_embeds.requires_grad_(True) pass @@ -522,7 +527,12 @@ def LlamaModel_fast_forward( # Careful for inference the attention_mask is size (1, kv_seq_len) # Whilst the input_embeds is size (1, 1, 4096) inputs_requires_grad = inputs_embeds.requires_grad - if inputs_requires_grad: inputs_embeds.requires_grad_(False) + if not inputs_embeds.is_leaf: + inputs_embeds = inputs_embeds.detach() + inputs_requires_grad = True + elif inputs_requires_grad: + inputs_embeds.requires_grad_(False) + pass inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2) if inputs_requires_grad: inputs_embeds.requires_grad_(True) pass @@ -1335,7 +1345,7 @@ class FastLlamaModel: if model_type == "llama": apply_lora_mlp = apply_lora_mlp_swiglu elif model_type == "mistral": apply_lora_mlp = apply_lora_mlp_swiglu - elif model_type == "gemma": apply_lora_mlp = apply_lora_mlp_geglu + elif model_type == "gemma": apply_lora_mlp = apply_lora_mlp_geglu_approx else: raise NotImplementedError(f"Unsloth: {model_type} is not yet implemented!") pass