2.4x faster Gemma (#197)
* Update save.py * Update save.py * linking * llama.cpp bugs * Update save.py * Update save.py * saving * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update save.py * Update __init__.py * Update save.py * Update save.py * Update save.py * save * trainer * spaces * original * Gemma * Update pyproject.toml * Update mapper.py * Update fast_lora.py * FastGemmaModel * model_type * Update llama.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update llama.py * Update fast_lora.py * Update llama.py * Update llama.py * Update cross_entropy_loss.py * Update llama.py * Update llama.py * gemma * Update llama.py * Update llama.py * Update llama.py * Update llama.py * Update fast_lora.py * Update fast_lora.py * Fast CE Loss * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * CE * Update llama.py * Update llama.py * Update cross_entropy_loss.py * Update geglu.py * Update cross_entropy_loss.py * revert * Update llama.py * Update llama.py * norm * Update gemma.py * Update gemma.py * position_ids * Update gemma.py * Update gemma.py * pos * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * revert * revert * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update llama.py * Update llama.py * Update llama.py * Update llama.py * rope * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * llama * Update llama.py * gemma * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update gemma.py * Update save.py * RoPE * Update llama.py * Update llama.py * Update llama.py * Update gemma.py * correct_dtype * Update gemma.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Chat Templates * Update README.md * Update README.md
This commit is contained in:
Daniel Han
GitHub
parent
1b7bf718cc
commit
e7c53fb370
14 changed files with 767 additions and 141 deletions
40
README.md
40
README.md
|
|
@ -10,7 +10,7 @@
|
|||
<a href="https://discord.gg/u54VK8m8tk"><img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/Discord button.png" height="48"></a>
|
||||
<a href="https://ko-fi.com/unsloth"><img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/buy me a coffee button.png" height="48"></a>
|
||||
|
||||
### Finetune Mistral, Llama 2-5x faster with 70% less memory!
|
||||
### Finetune Mistral, Gemma, Llama 2-5x faster with 70% less memory!
|
||||
|
||||
|
||||
|
||||
|
|
@ -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)|
|
||||
| <img height="14" src="https://upload.wikimedia.org/wikipedia/commons/6/6f/Logo_of_Twitter.svg" /> **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
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ exclude = ["images*"]
|
|||
|
||||
[project.optional-dependencies]
|
||||
huggingface = [
|
||||
"transformers>=4.37.0",
|
||||
"transformers>=4.38.0",
|
||||
"datasets",
|
||||
"sentencepiece",
|
||||
"accelerate>=0.26.1",
|
||||
|
|
|
|||
|
|
@ -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. <start_of_turn> maps to 106.
|
||||
# <end_of_turn> maps to 107. user and model are normal 1 word tokens.
|
||||
gemma_template = \
|
||||
"{% for message in messages %}"\
|
||||
"{% if message['role'] == 'user' %}"\
|
||||
"{{'<start_of_turn>user\n' + message['content'] | trim + '<end_of_turn>\n'}}"\
|
||||
"{% elif message['role'] == 'assistant' %}"\
|
||||
"{{'<start_of_turn>model\n' + message['content'] | trim + '<end_of_turn>\n' }}"\
|
||||
"{% else %}"\
|
||||
"{{ '<start_of_turn>system\n' + message['content'] | trim + '<end_of_turn>\n' }}"\
|
||||
"{% endif %}"\
|
||||
"{% endfor %}"\
|
||||
"{% if add_generation_prompt %}"\
|
||||
"{{ '<start_of_turn>model\n' }}"\
|
||||
"{% endif %}"
|
||||
gemma_eos_token = "<end_of_turn>"
|
||||
CHAT_TEMPLATES["gemma"] = (gemma_template, gemma_eos_token,)
|
||||
|
||||
|
||||
# Gemma with ChatML instead
|
||||
gemma_chatml_template = chatml_template
|
||||
gemma_chatml_eos_token = (
|
||||
{"<start_of_turn>" : "<|im_start|>", "<end_of_turn>" : "<|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 = {"<start_of_turn>" : "<|im_start|>", "<end_of_turn>" : "<|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 </s> ourselves
|
||||
assert(correct_prompt == our_prompt.replace("</s>", ""))
|
||||
|
||||
# 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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
104
unsloth/kernels/geglu.py
Normal file
104
unsloth/kernels/geglu.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
282
unsloth/models/gemma.py
Normal file
282
unsloth/models/gemma.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue