diff --git a/README.md b/README.md index 98f83e09c7..759057f2a4 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ -### Finetune Mistral, Llama 2-5x faster with 70% less memory! +### Finetune Mistral, Gemma, Llama 2-5x faster with 70% less memory! ![](https://i.ibb.co/sJ7RhGG/image-41.png) @@ -22,28 +22,30 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and | Unsloth supports | Free Notebooks | Performance | Memory use | |-----------------|--------------------------------------------------------------------------------------------------------------------------|-------------|----------| +| **Gemma 7b** | [▶️ Start on Colab](https://colab.research.google.com/drive/10NbwlsRChbma1v55m8LAPYG15uQv6HLo?usp=sharing) | 2.4x faster | 58% less | | **Mistral 7b** | [▶️ Start on Colab](https://colab.research.google.com/drive/1Dyauq4kTZoLewQ1cApceUQVNcnnNTzg_?usp=sharing) | 2.2x faster | 62% less | | **Llama-2 7b** | [▶️ Start on Colab](https://colab.research.google.com/drive/1lBzz5KeZJKXjvivbYvmGarix9Ao6Wxe5?usp=sharing) | 2.2x faster | 43% less | -| **DPO - Zephyr** | [▶️ Start on Colab](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing) | 1.9x faster | 19% less | | **TinyLlama** | [▶️ Start on Colab](https://colab.research.google.com/drive/1AZghoNBQaMDgWJpi4RbffGM1h6raLUj9?usp=sharing) | 3.9x faster | 74% less | | **CodeLlama 34b** A100 | [▶️ Start on Colab](https://colab.research.google.com/drive/1y7A0AxE3y8gdj4AVkl2aZX47Xu3P1wJT?usp=sharing) | 1.9x faster | 27% less | | **Mistral 7b** 1xT4 | [▶️ Start on Kaggle](https://www.kaggle.com/code/danielhanchen/kaggle-mistral-7b-unsloth-notebook) | 5x faster\* | 62% less | +| **DPO - Zephyr** | [▶️ Start on Colab](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing) | 1.9x faster | 19% less | - This [conversational notebook](https://colab.research.google.com/drive/1Aau3lgPzeZKQ-98h69CCu1UJcvIBLmy2?usp=sharing) is useful for ShareGPT ChatML / Vicuna templates. - This [text completion notebook](https://colab.research.google.com/drive/1ef-tab5bhkvWmBOObepl1WgJvfvSzn5Q?usp=sharing) is for raw text. This [DPO notebook](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing) replicates Zephyr. -- Colab provides a free GPU sometimes. Kaggle has 30 hrs free per week on a 12 hr running cap. -- \* Kaggle has 2x T4s, but we use 1. Due to overhead, 1x T4 is 5x faster. Use Colab as Kaggle takes 10 mins to install. +- \* Kaggle has 2x T4s, but we use 1. Due to overhead, 1x T4 is 5x faster. ## 🦥 Unsloth.ai News -- 📣 [DPO support](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing) is now included. [More info](#DPO) on DPO. -- 📣 [TinyLlama 1.1b](https://colab.research.google.com/drive/1AZghoNBQaMDgWJpi4RbffGM1h6raLUj9?usp=sharing) on 3T tokens now works. -- 📣 We did a [blog](https://huggingface.co/blog/unsloth-trl) with 🤗Hugging Face, and we're in their official docs! Check out the [SFT docs](https://huggingface.co/docs/trl/main/en/sft_trainer#accelerate-fine-tuning-2x-using-unsloth) and [DPO docs](https://huggingface.co/docs/trl/main/en/dpo_trainer#accelerate-dpo-fine-tuning-using-unsloth). -- 📣 Now supports **Llama, Yi, Mistral, CodeLlama, Qwen (llamafied), Deepseek** and their derived models (**Open Hermes** etc). Llama 7, 13, 70b; CodeLlama 7, 13, 34, 70b; Yi 6, 34b are all supported! -- 📣 **Download models 4x faster** from 🤗Hugging Face! Eg: `unsloth/mistral-7b-bnb-4bit` See our [HF collection](https://huggingface.co/collections/unsloth/load-4bit-models-4x-faster-659042e3a41c3cbad582e734) for more! +- 📣 [Gemma 7b](https://colab.research.google.com/drive/10NbwlsRChbma1v55m8LAPYG15uQv6HLo?usp=sharing) on 6T tokens now works. And [Gemma 2b notebook](https://colab.research.google.com/drive/15gGm7x_jTm017_Ic8e317tdIpDG53Mtu?usp=sharing) +- 📣 Added [conversational notebooks](https://colab.research.google.com/drive/1ef-tab5bhkvWmBOObepl1WgJvfvSzn5Q?usp=sharing) and [raw text notebooks](https://colab.research.google.com/drive/1bMOKOBzxQWUIGZBs_B0zm8pimuEnZdfM?usp=sharing) +- 📣 [2x faster inference](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing) added for all our models +- 📣 [DPO support](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing) is now included. [More info](#DPO) on DPO +- 📣 We did a [blog](https://huggingface.co/blog/unsloth-trl) with 🤗Hugging Face and are in their official docs! Check out the [SFT docs](https://huggingface.co/docs/trl/main/en/sft_trainer#accelerate-fine-tuning-2x-using-unsloth) and [DPO docs](https://huggingface.co/docs/trl/main/en/dpo_trainer#accelerate-dpo-fine-tuning-using-unsloth) +- 📣 [Download models 4x faster](https://huggingface.co/collections/unsloth/) from 🤗Hugging Face. Eg: `unsloth/mistral-7b-bnb-4bit` ## 🔗 Links and Resources | Type | Links | | ------------------------------- | --------------------------------------- | +| 📚 **Wiki & FAQ** | [Read Our Wiki](https://github.com/unslothai/unsloth/wiki) | | 📜 **Documentation** | [Read The Doc](https://github.com/unslothai/unsloth/tree/main#-documentation) | | 💾 **Installation** | [unsloth/README.md](https://github.com/unslothai/unsloth/tree/main#installation-instructions)| |   **Twitter (aka X)** | [Follow us on X](https://twitter.com/unslothai)| @@ -113,8 +115,8 @@ pip install --upgrade --force-reinstall --no-cache-dir torch==2.1.0 triton \ ```bash pip install "unsloth[cu118] @ git+https://github.com/unslothai/unsloth.git" pip install "unsloth[cu121] @ git+https://github.com/unslothai/unsloth.git" -pip install "unsloth[cu118_ampere] @ git+https://github.com/unslothai/unsloth.git" -pip install "unsloth[cu121_ampere] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu118-ampere] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu121-ampere] @ git+https://github.com/unslothai/unsloth.git" ``` 3. For Pytorch 2.1.1: Use the `"ampere"` path for newer RTX 30xx GPUs or higher. ```bash @@ -122,10 +124,10 @@ pip install --upgrade --force-reinstall --no-cache-dir torch==2.1.1 triton \ --index-url https://download.pytorch.org/whl/cu121 ``` ```bash -pip install "unsloth[cu118_torch211] @ git+https://github.com/unslothai/unsloth.git" -pip install "unsloth[cu121_torch211] @ git+https://github.com/unslothai/unsloth.git" -pip install "unsloth[cu118_ampere_torch211] @ git+https://github.com/unslothai/unsloth.git" -pip install "unsloth[cu121_ampere_torch211] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu118-torch211] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu121-torch211] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu118-ampere-torch211] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu121-ampere-torch211] @ git+https://github.com/unslothai/unsloth.git" ``` 4. For Pytorch 2.2.0: Use the `"ampere"` path for newer RTX 30xx GPUs or higher. ```bash @@ -133,10 +135,10 @@ pip install --upgrade --force-reinstall --no-cache-dir torch==2.2.0 triton \ --index-url https://download.pytorch.org/whl/cu121 ``` ```bash -pip install "unsloth[cu118_torch220] @ git+https://github.com/unslothai/unsloth.git" -pip install "unsloth[cu121_torch220] @ git+https://github.com/unslothai/unsloth.git" -pip install "unsloth[cu118_ampere_torch220] @ git+https://github.com/unslothai/unsloth.git" -pip install "unsloth[cu121_ampere_torch220] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu118-torch220] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu121-torch220] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu118-ampere-torch220] @ git+https://github.com/unslothai/unsloth.git" +pip install "unsloth[cu121-ampere-torch220] @ git+https://github.com/unslothai/unsloth.git" ``` 5. If you get errors, try the below first, then go back to step 1: ```bash diff --git a/pyproject.toml b/pyproject.toml index 049711276d..7e8956c712 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,7 @@ exclude = ["images*"] [project.optional-dependencies] huggingface = [ - "transformers>=4.37.0", + "transformers>=4.38.0", "datasets", "sentencepiece", "accelerate>=0.26.1", diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index eb61056320..6f0237c27b 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -217,6 +217,35 @@ alpaca_eos_token = "eos_token" CHAT_TEMPLATES["alpaca"] = (alpaca_template, alpaca_eos_token,) +# https://huggingface.co/google/gemma-7b-it +# Notice we must use |trim for lstrip and rstrip. maps to 106. +# maps to 107. user and model are normal 1 word tokens. +gemma_template = \ + "{% for message in messages %}"\ + "{% if message['role'] == 'user' %}"\ + "{{'user\n' + message['content'] | trim + '\n'}}"\ + "{% elif message['role'] == 'assistant' %}"\ + "{{'model\n' + message['content'] | trim + '\n' }}"\ + "{% else %}"\ + "{{ 'system\n' + message['content'] | trim + '\n' }}"\ + "{% endif %}"\ + "{% endfor %}"\ + "{% if add_generation_prompt %}"\ + "{{ 'model\n' }}"\ + "{% endif %}" +gemma_eos_token = "" +CHAT_TEMPLATES["gemma"] = (gemma_template, gemma_eos_token,) + + +# Gemma with ChatML instead +gemma_chatml_template = chatml_template +gemma_chatml_eos_token = ( + {"" : "<|im_start|>", "" : "<|im_end|>"}, + "<|im_end|>", +) +CHAT_TEMPLATES["gemma_chatml"] = (gemma_chatml_template, gemma_chatml_eos_token,) + + def get_chat_template( tokenizer, chat_template = "chatml", @@ -229,7 +258,7 @@ def get_chat_template( old_padding_side = tokenizer.padding_side - if type(chat_template) in (list, tuple): + if type(chat_template) in (list, tuple,): chat_template, stop_word = chat_template assert(type(chat_template) is str) assert(type(stop_word) is str) @@ -238,7 +267,38 @@ def get_chat_template( chat_template, stop_word = CHAT_TEMPLATES[chat_template] - if stop_word != "eos_token": + if type(stop_word) in (list, tuple,): + token_mapping, stop_word = stop_word + assert(type(token_mapping) is dict) + else: + token_mapping = None + + assert(type(stop_word) is str) + + # token_mapping = {"" : "<|im_start|>", "" : "<|im_end|>"} + # For Gemma :) + if token_mapping is not None: + + string_vocab = tokenizer._tokenizer.to_str() + + for old_token, new_token in token_mapping.items(): + old_count = string_vocab.count(f'"{old_token}"') + new_count = string_vocab.count(f'"{new_token}"') + if new_count != 0: + print(f"{new_token} is already a token. Skipping.") + elif old_count == 0: + raise RuntimeError(f"{old_token} was not part of the tokenizer!") + else: + string_vocab = string_vocab.replace(f'"{old_token}"', f'"{new_token}"') + pass + pass + + logger.warning_once(f"Unsloth: Will map {stop_word} to EOS = {tokenizer.eos_token}.") + string_vocab = string_vocab.replace(tokenizer.eos_token, stop_word) + new_tokenizer = tokenizer._tokenizer.from_str(string_vocab) + tokenizer = tokenizer.__class__(tokenizer_object = new_tokenizer, eos_token = stop_word) + + elif stop_word != "eos_token": logger.warning_once(f"Unsloth: Will map {stop_word} to EOS = {tokenizer.eos_token}.") # Replaces the old EOS token with a new one. @@ -252,6 +312,7 @@ def get_chat_template( new_tokenizer = tokenizer._tokenizer.from_str(string_vocab) tokenizer = tokenizer.__class__(tokenizer_object = new_tokenizer, eos_token = stop_word) pass + else: raise TypeError( f"Unsloth: `chat_template` must be a tuple of (your_template, eos_token,) or one of\n"\ @@ -318,6 +379,7 @@ def test_chat_templates(): {"role": "user", "content": " No it's 100% 5! "}, ] + # Zephyr from transformers import AutoTokenizer template = zephyr_template correct_tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta") @@ -326,6 +388,7 @@ def test_chat_templates(): our_prompt = correct_tokenizer.apply_chat_template(messages, tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) + # Chatml template = chatml_template correct_tokenizer = AutoTokenizer.from_pretrained("teknium/OpenHermes-2.5-Mistral-7B") correct_prompt = correct_tokenizer.apply_chat_template(messages, tokenize = False, add_generation_prompt = True) @@ -333,6 +396,7 @@ def test_chat_templates(): our_prompt = correct_tokenizer.apply_chat_template(messages, tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) + # Mistral template = mistral_template correct_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2") correct_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) @@ -340,6 +404,7 @@ def test_chat_templates(): our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) + # Llama template = llama_template correct_tokenizer = AutoTokenizer.from_pretrained("unsloth/llama-2-7b-chat") correct_prompt = correct_tokenizer.apply_chat_template(messages, tokenize = False, add_generation_prompt = True) @@ -347,6 +412,7 @@ def test_chat_templates(): our_prompt = correct_tokenizer.apply_chat_template(messages, tokenize = False, add_generation_prompt = True) assert(correct_prompt == our_prompt) + # Vicuna try: from fastchat.conversation import get_conv_template except: @@ -381,4 +447,11 @@ def test_chat_templates(): our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) # We add ourselves assert(correct_prompt == our_prompt.replace("", "")) + + # Gemma + correct_tokenizer = AutoTokenizer.from_pretrained("unsloth/gemma-7b-it") + correct_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) + correct_tokenizer.chat_template = gemma_template + our_prompt = correct_tokenizer.apply_chat_template(messages[1:], tokenize = False, add_generation_prompt = True) + assert(our_prompt == correct_prompt) pass 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/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index 0a73a393ec..260577912f 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -20,12 +20,14 @@ from transformers.models.llama.modeling_llama import logger @triton.jit -def _cross_entropy_forward(logits_ptr, logits_row_stride, - loss_ptr, - lse_ptr, - labels_ptr, - n_cols, - BLOCK_SIZE: tl.constexpr,): +def _cross_entropy_forward( + logits_ptr, logits_row_stride, + loss_ptr, + logsumexp_ptr, + labels_ptr, + VOCAB_SIZE : tl.constexpr, + BLOCK_SIZE : tl.constexpr, +): """ Cross Entropy Loss = 1/n sum [ -yi log(Pi) ] Pi = exp(xi) / sum(exp(xi)) @@ -34,40 +36,114 @@ def _cross_entropy_forward(logits_ptr, logits_row_stride, = y * (log[sum(exp(x))] - x) If y == 0: CE_i = 0 If y == 1: CE_i = logsumexp - x + + logsumexp is also stable + Take y = log[sum(exp(x))] + exp(y) = sum(exp(x)) + exp(y) = sum(exp(x - c)*exp(c)) Since e^(x-c)*e^c = e^x + exp(y) = exp(c)*sum(exp(x - c)) + y = log(exp(c)*sum(exp(x - c))) + y = c + log[sum(exp(x - c))] + This means we can set c = max(x) to make sure + exp(x - c) always is exp(x - max(x)). + This ensures exp(x - max(x))'s maximum is 1 as exp(0) = 1. """ row_idx = tl.program_id(0) - logits_ptr += row_idx * logits_row_stride - loss_ptr += row_idx - lse_ptr += row_idx - labels_ptr += row_idx + logits_ptr += row_idx * logits_row_stride.to(tl.int64) + loss_ptr += row_idx + logsumexp_ptr += row_idx + labels_ptr += row_idx col_offsets = tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols + mask = col_offsets < VOCAB_SIZE - # TODO: Fixup int32 locations to int64 label_idx = tl.load(labels_ptr).to(tl.int32) logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32) - max_logits = tl.max(logits, 0) - # Maximum stops overflow - lse = tl.log(tl.sum(tl.exp(logits - max_logits), 0)) + max_logits - tl.store(lse_ptr, lse) + c = tl.max(logits, 0) + logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0)) if label_idx != -100: - logits_label = tl.load(logits_ptr + label_idx).to(tl.float32) - loss = lse - logits_label + x = tl.load(logits_ptr + label_idx).to(tl.float32) + loss = logsumexp - x else: loss = 0.0 + tl.store(logsumexp_ptr, logsumexp) tl.store(loss_ptr, loss) pass @triton.jit -def _cross_entropy_backward(logits_ptr, logits_row_stride, - dloss_ptr, dloss_row_stride, - lse_ptr, - labels_ptr, - n_cols, - BLOCK_SIZE: tl.constexpr,): +def _chunked_cross_entropy_forward( + logits_ptr, logits_row_stride, + loss_ptr, + logsumexp_ptr, + labels_ptr, + VOCAB_SIZE : tl.constexpr, + N_CHUNKS : tl.constexpr, + BLOCK_SIZE : tl.constexpr, +): + """ + 256K vocab divided in 4 chunks + + |-65536-| |-65536-| |-65536-| |-65536-| + |-------| |-------| |-------| |-------| + |-------| |-------| |-------| |-------| + + If y == 0: CE_i = 0 + If y == 1: CE_i = logsumexp - x + + Notice we can do logsumexp for each chunk and then + logsumexp[chunk_sum(logsumexp)] == logsumexp + + chunk_sum = log[chunk_sum(logsumexp)] + = log[exp(logsumexp(a)) + ... + exp(logsumexp(z))] + = log[exp(log[sum(exp(a))]) + ... + exp(log[sum(exp(z))])] + = log[sum(exp(a)) + ... + sum(exp(z))] + = logsumexp(x) + + This means we can perform a logsumexp for each chunk, then do a + final logsumexp reduction! + + Ie do: logsumexp(chunked_logsumexp) - x + """ + row_idx = tl.program_id(0) + chunk_idx = tl.program_id(1) + logits_ptr += row_idx * logits_row_stride.to(tl.int64) + loss_ptr += row_idx + logsumexp_ptr += row_idx * N_CHUNKS + chunk_idx + labels_ptr += row_idx + + col_offsets = chunk_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < VOCAB_SIZE + + label_idx = tl.load(labels_ptr).to(tl.int32) + logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32) + c = tl.max(logits, 0) + logsumexp = c + tl.log(tl.sum(tl.exp(logits - c), 0)) + + if chunk_idx == 0: + # logsumexp(chunked_logsumexp) - x + # Do the -x separately + if label_idx != -100: + x = tl.load(logits_ptr + label_idx).to(tl.float32) + loss = -1.0 * x + else: + loss = 0.0 + tl.store(loss_ptr, loss) + pass + tl.store(logsumexp_ptr, logsumexp) +pass + + +@triton.jit +def _cross_entropy_backward( + logits_ptr, logits_row_stride, + dloss_ptr, dloss_row_stride, + logsumexp_ptr, + labels_ptr, + VOCAB_SIZE : tl.constexpr, + BLOCK_SIZE : tl.constexpr, +): """ CE_i = -y log(P) = y * (log[sum(exp(x))] - x) dC/dx = d/dx (y * log[sum(exp(x))] - x * y) @@ -83,47 +159,80 @@ def _cross_entropy_backward(logits_ptr, logits_row_stride, If y == 1 and x == label: dC/dlabel = exp[x - logsumexp] - 1 If y == 1 and x != label: dC/dx = exp[x - logsumexp] """ - row_idx = tl.program_id(0) - logits_ptr += row_idx * logits_row_stride + row_idx = tl.program_id(0) + block_idx = tl.program_id(1) + + logits_ptr += row_idx * logits_row_stride.to(tl.int64) dloss_ptr += row_idx * dloss_row_stride - col_offsets = tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols - # TODO: Fixup int32 locations to int64 + col_offsets = block_idx*BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = col_offsets < VOCAB_SIZE label_idx = tl.load(labels_ptr + row_idx).to(tl.int32) if label_idx != -100: dloss = tl.load(dloss_ptr) else: dloss = 0.0 - logits = tl.load(logits_ptr + col_offsets, mask = mask, other = 0).to(tl.float32) - lse = tl.load(lse_ptr + row_idx) - probs = tl.exp(logits - lse) - probs = tl.where(col_offsets == label_idx, probs - 1.0, probs) - tl.store(logits_ptr + col_offsets, dloss * probs, mask = mask) + x = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")).to(tl.float32) + logsumexp = tl.load(logsumexp_ptr + row_idx) + y = tl.exp(x - logsumexp) + y = tl.where( + col_offsets == label_idx, + y - 1.0, # exp(x - logsumexp) - 1 + y, # exp(x - logsumexp) + ) + + # If y == 0: dC/dx = 0 ==> we already masked it to be = 0, so dloss = 0. + tl.store(logits_ptr + col_offsets, dloss * y, mask = mask) pass +MAX_FUSED_SIZE = 65536 # 2**16 + class Fast_CrossEntropyLoss(torch.autograd.Function): @staticmethod def forward(ctx, logits, labels): - n_rows, n_cols = logits.shape - BLOCK_SIZE, num_warps = calculate_settings(n_cols) - losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda") - logsumexp = torch.empty(n_rows, dtype = torch.float32, device = "cuda") + n_rows, vocab_size = logits.shape - _cross_entropy_forward[(n_rows,)]( - logits, logits.stride(0), - losses, - logsumexp, - labels, - n_cols, - BLOCK_SIZE = BLOCK_SIZE, - num_warps = num_warps, - ) + div, mod = divmod(vocab_size, MAX_FUSED_SIZE) + n_chunks = div + (mod != 0) + losses = torch.empty(n_rows, dtype = torch.float32, device = "cuda") + + if n_chunks == 1: + # For small vocabs <= 65336 like Llama, Mistral + BLOCK_SIZE, num_warps = calculate_settings(vocab_size) + logsumexp = torch.empty(n_rows, dtype = torch.float32, device = "cuda") + + _cross_entropy_forward[(n_rows,)]( + logits, logits.stride(0), + losses, + logsumexp, + labels, + VOCAB_SIZE = vocab_size, + BLOCK_SIZE = BLOCK_SIZE, + num_warps = num_warps, + ) + else: + # For large vocabs > 65336 like Gemma 256K + logsumexp = torch.empty((n_rows, n_chunks,), dtype = torch.float32, device = "cuda") + + _chunked_cross_entropy_forward[(n_rows, n_chunks,)]( + logits, logits.stride(0), + losses, + logsumexp, + labels, + VOCAB_SIZE = vocab_size, + N_CHUNKS = n_chunks, + BLOCK_SIZE = MAX_FUSED_SIZE, + num_warps = 32, + ) + # logsumexp(chunked_logsumexp) - x + # Do the -x separately + logsumexp = torch.logsumexp(logsumexp, dim = 1) # Row sum + losses += logsumexp + losses.masked_fill_(labels == -100, 0) # Don't forget to mask padding out! + pass - ctx.BLOCK_SIZE = BLOCK_SIZE - ctx.num_warps = num_warps ctx.save_for_backward(logits, logsumexp, labels) return losses pass @@ -131,23 +240,26 @@ class Fast_CrossEntropyLoss(torch.autograd.Function): @staticmethod def backward(ctx, dlosses): logits, logsumexp, labels = ctx.saved_tensors - n_rows, n_cols = logits.shape + n_rows, vocab_size = logits.shape - _cross_entropy_backward[(n_rows,)]( + BLOCK_SIZE = 4096 + div, mod = divmod(vocab_size, BLOCK_SIZE) + n_blocks = div + (mod != 0) + + _cross_entropy_backward[(n_rows, n_blocks,)]( logits, logits.stride(0), dlosses, dlosses.stride(0), logsumexp, labels, - n_cols, - BLOCK_SIZE = ctx.BLOCK_SIZE, - num_warps = ctx.num_warps, + VOCAB_SIZE = vocab_size, + BLOCK_SIZE = BLOCK_SIZE, + num_warps = 8, ) return logits, None, None, pass pass -slow_cross_entropy_loss = torch.nn.functional.cross_entropy def fast_cross_entropy_loss(logits, labels): """ Arguments: @@ -159,25 +271,10 @@ def fast_cross_entropy_loss(logits, labels): batch, seq_len, d = logits.shape assert(labels.shape == (batch, seq_len)) - # Prelim support Qwen, Deepseek other large vocab sizes > 2^16 - if d > MAX_FUSED_SIZE: - logger.warning_once( - f"Unsloth: Vocab size of {d} exceeds the max CUDA blocksize of {MAX_FUSED_SIZE}.\n"\ - "For now, Unsloth will use Pytorch's CrossEntropyLoss, which will entail a\n"\ - "25% increase in memory usage and be slower. Make an issue on \n"\ - "Unsloth's Github page if you want a faster and more memory efficient kernel!" - ) - loss = slow_cross_entropy_loss( - logits.float().view(batch*seq_len, d), # Must cast to float32 for numerical stability - labels.view(-1), - ) - return loss - else: - loss = Fast_CrossEntropyLoss.apply( - logits.view(batch*seq_len, d), - labels.view(-1), - ) - n_items = torch.count_nonzero(labels != -100) - return loss.sum() / n_items - pass + loss = Fast_CrossEntropyLoss.apply( + logits.view(batch*seq_len, d), + labels.view(-1), + ) + n_items = torch.count_nonzero(labels != -100) + return loss.sum() / n_items pass diff --git a/unsloth/kernels/fast_lora.py b/unsloth/kernels/fast_lora.py index b3a1098355..3ed0d3c914 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,20 @@ 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 +108,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 +124,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 +147,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 @@ -172,24 +163,36 @@ class LoRA_MLP(torch.autograd.Function): return dX.view(batch, seq_len, hd), \ None, None, d_gateA.t(), d_gateB.t(), None, \ None, None, d_upA.t(), d_upB.t(), None, \ - None, None, d_downA.t(), d_downB.t(), None, + None, None, d_downA.t(), d_downB.t(), None, \ + None, None, # _backward and _forward pass 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..4aa634a4bd --- /dev/null +++ b/unsloth/models/gemma.py @@ -0,0 +1,282 @@ +# 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, + GemmaRotaryEmbedding, + apply_rotary_pos_emb, + repeat_kv, +) +from transformers.modeling_attn_mask_utils import ( + _prepare_4d_causal_attention_mask_for_sdpa, +) +# 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, +): + if False:#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.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.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 + + +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 = LlamaForCausalLM_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 + + # Gemma has tied weights! This means lm_head == embed_tokens + if model.model.embed_tokens.weight.data_ptr() != model.lm_head.weight.data_ptr(): + lm_head = torch.nn.Linear(1, 1, bias = None) + del lm_head.weight + lm_head.weight = model.model.embed_tokens.weight + lm_head.in_features = lm_head.weight.shape[1] + lm_head.out_features = lm_head.weight.shape[0] + model.lm_head = lm_head + pass + + # 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 + # Downcast RoPE embedding to correct data type + if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ + and (module.cos_cached.dtype != correct_dtype): + + module.cos_cached = module.cos_cached.to(correct_dtype) + module.sin_cached = module.sin_cached.to(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 + # We do this first since += 1 seems to not be liked by requires_grad = True + 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 + + # Patch RMS Layernorm + for name, module in model.named_modules(): + if isinstance(module, GemmaRMSNorm): + 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 + 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..359761c06a 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -119,6 +119,7 @@ def LlamaAttention_fast_forward_inference( n_groups = self.num_key_value_groups n_kv_heads = self.num_key_value_heads head_dim = self.head_dim + attention_size = n_heads*head_dim # assert(n_kv_heads * n_groups == n_heads) seq_len = K1.shape[-2] kv_seq_len = seq_len + 1 @@ -131,7 +132,7 @@ def LlamaAttention_fast_forward_inference( self.paged_attention_V = self.paged_attention[:,1] self.paged_attention_K[:seq_len] = K1.permute(2, 0, 1, 3) self.paged_attention_V[:seq_len] = V1.permute(2, 0, 1, 3) - self.temp_QA = torch.empty((2, bsz, 1, hd), dtype = dtype, device = "cuda") + self.temp_QA = torch.empty((2, bsz, 1, attention_size), dtype = dtype, device = "cuda") self.temp_KV = torch.empty((2, bsz, 1, n_kv_heads*head_dim), dtype = dtype, device = "cuda") self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = "cuda") self.attention = torch.empty((bsz, n_heads, 1, KV_CACHE_INCREMENT+seq_len), dtype = dtype, device = "cuda") @@ -201,13 +202,13 @@ def LlamaAttention_fast_forward_inference( A[:] = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32)#.to(A.dtype) A = torch.matmul(A, Vnn, out = Qn) A = A.transpose(1, 2) - A = A.reshape(bsz, 1, self.hidden_size) - A = fast_linear_forward(self.o_proj, A, out = self.temp_QA[1]) + A = A.reshape(bsz, 1, attention_size) + A = fast_linear_forward(self.o_proj, A, out = self.temp_QA[1][:,:,:self.hidden_size]) return A, (Kn, Vn) 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 @@ -339,7 +340,7 @@ def LlamaAttention_fast_forward( # Go back to (batch_size, seq_len, n_heads, head_dim) A = A.transpose(1, 2).contiguous() pass - attn_output = A.reshape(bsz, q_len, self.hidden_size) + attn_output = A.reshape(bsz, q_len, n_heads*head_dim) attn_output = self.apply_o(self, attn_output) attn_weights = None return attn_output, attn_weights, past_key_value @@ -390,7 +391,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 +508,14 @@ 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_requires_grad = inputs_embeds.requires_grad + if inputs_requires_grad: inputs_embeds.requires_grad_(False) + inputs_embeds *= math_sqrt(self.config.hidden_size) + if inputs_requires_grad: inputs_embeds.requires_grad_(True) + 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 +655,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) @@ -812,7 +821,7 @@ class LlamaRotaryEmbedding(torch.nn.Module): self.register_buffer("sin_cached", emb.sin().to(dtype=dtype, device=device, non_blocking=True), persistent=False) pass - def forward(self, x, seq_len=None): + def forward(self, x, position_ids=None, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) @@ -886,20 +895,22 @@ class FastLlamaModel: device_map = "sequential", rope_scaling = None, fix_tokenizer = True, + model_patcher = None, **kwargs, ): + if model_patcher is None: model_patcher = FastLlamaModel SUPPORTS_BFLOAT16 = torch.cuda.is_bf16_supported() gpu_stats = torch.cuda.get_device_properties(0) 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 +966,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): @@ -1159,6 +1170,14 @@ class FastLlamaModel: quant_state.dtype = correct_dtype pass pass + # Downcast RoPE embedding to correct data type + if (name.endswith("rotary_emb") or hasattr(module, "cos_cached")) \ + and (module.cos_cached.dtype != correct_dtype): + + module.cos_cached = module.cos_cached.to(correct_dtype) + module.sin_cached = module.sin_cached.to(correct_dtype) + pass + pass pass # Clear deleted GPU items @@ -1309,6 +1328,16 @@ class FastLlamaModel: ) pass + # Get activation function + model_type = model.config.model_type + + 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 + else: + raise NotImplementedError(f"Unsloth: {model_type} 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..67a59c850c 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -24,6 +24,9 @@ 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) +if SUPPORTS_GEMMA: + from .gemma import FastGemmaModel del major, minor @@ -99,6 +102,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 +127,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/mapper.py b/unsloth/models/mapper.py index 323358fff1..afcbdb75f5 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -74,6 +74,22 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/solar-10.7b-bnb-4bit" : ( "upstage/SOLAR-10.7B-v1.0", ), + "unsloth/gemma-7b-bnb-4bit" : ( + "unsloth/gemma-7b", + "google/gemma-7b", + ), + "unsloth/gemma-2b-bnb-4bit" : ( + "unsloth/gemma-2b", + "google/gemma-2b", + ), + "unsloth/gemma-7b-it-bnb-4bit" : ( + "unsloth/gemma-7b-it", + "google/gemma-7b-it", + ), + "unsloth/gemma-2b-bnb-4bit" : ( + "unsloth/gemma-2b-it", + "google/gemma-2b-it", + ), } INT_TO_FLOAT_MAPPER = {} diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 0e36023255..6c9d9ecc5c 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -293,8 +293,10 @@ class FastMistralModel(FastLlamaModel): device_map = "sequential", rope_scaling = None, # Mistral does not support RoPE scaling fix_tokenizer = True, + model_patcher = None, **kwargs, ): + if model_patcher is None: model_patcher = FastMistralModel # Mistral does NOT support RoPE Scaling! if rope_scaling is not None: logger.warning_once("Unsloth: Mistral models do not support RoPE scaling.") @@ -305,13 +307,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 +362,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): diff --git a/unsloth/save.py b/unsloth/save.py index 83e13bd51c..51ddeb30e7 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -369,6 +369,7 @@ def unsloth_save_model( # Switch to our fast saving modules if it's a slow PC! n_cpus = psutil.cpu_count(logical = False) + if n_cpus is None: n_cpus = 1 if safe_serialization is None: safe_serialization = True @@ -669,7 +670,9 @@ def save_to_gguf( pass pass - n_cpus = psutil.cpu_count()*2 + n_cpus = psutil.cpu_count() + if n_cpus is None: n_cpus = 1 + n_cpus *= 2 # Concurrency from https://rentry.org/llama-cpp-conversions#merging-loras-into-a-model final_location = f"./{model_directory}-unsloth.{first_conversion.upper()}.gguf"