ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix inference

Browse source
This commit is contained in:
Daniel Han-Chen 2023-12-27 03:51:31 +11:00
commit a3ed0dc4f4
2 changed files with 4 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -8,7 +8,7 @@
| Llama 7b | Mistral 7b | CodeLlama 34b | Llama 7b Kaggle 2x T4 |
|-----------------------------|-----------------------------|-------------------------|------------------------|
| **2.2x faster, -43% VRAM** | **2.2x faster, -62% VRAM** | **1.9x faster, -27% VRAM** | **5.5x faster, -44% VRAM** |
| [Colab Alpaca example + inference](https://colab.research.google.com/drive/1oW55fBmwzCOrBVX66RcpptL3a99qWBxb?usp=sharing) | [Colab T4 example](https://colab.research.google.com/drive/15pyLgRN97B_jA56HS0esx56knA9I5tuv?usp=sharing) | [A100 example](https://colab.research.google.com/drive/1gdHyAx8XJsz2yNV-DHvbHjR1iCef5Qmh?usp=sharing) | [Kaggle Alpaca example](https://www.kaggle.com/danielhanchen/unsloth-alpaca-t4-ddp) |
| [Colab Alpaca example + inference, saving](https://colab.research.google.com/drive/1lBzz5KeZJKXjvivbYvmGarix9Ao6Wxe5?usp=sharing) | [Colab T4 example + inference, saving](https://colab.research.google.com/drive/1Dyauq4kTZoLewQ1cApceUQVNcnnNTzg_?usp=sharing) | [A100 example](https://colab.research.google.com/drive/1gdHyAx8XJsz2yNV-DHvbHjR1iCef5Qmh?usp=sharing) | [Kaggle Alpaca example](https://www.kaggle.com/danielhanchen/unsloth-alpaca-t4-ddp) |
| [Colab A100 example](https://colab.research.google.com/drive/1YIPY_18xm-K0iJDgvNkRoJsgkPMPAO3G?usp=sharing) | [Colab A100 example](https://colab.research.google.com/drive/1SKrKGV-BZoU4kv5q3g0jtE_OhRgPtrrQ?usp=sharing) | (59 more examples if you scroll down) | [Kaggle Slim Orca](https://www.kaggle.com/danielhanchen/unsloth-slimorca-t4-ddp) |
* Supports Llama (7, 13, 70b), Yi (6, 34b), Mistral (7b), Tinyllama, CodeLlama (7, 13, 34b), and all Llama / Mistral derived architectures!

6
unsloth/models/llama.py
View file

@ -141,8 +141,8 @@ def LlamaAttention_fast_forward_inference(
_, _, cached_len, _ = Kn.shape
Knn = Kn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
Vnn = Vn[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, cached_len, head_dim)
Knn = Knn.view(bsz, n_heads, cached_len, head_dim)
Vnn = Vnn.view(bsz, n_heads, cached_len, head_dim)
Knn = Knn.reshape(bsz, n_heads, cached_len, head_dim)
Vnn = Vnn.reshape(bsz, n_heads, cached_len, head_dim)
else:
Knn, Vnn = Kn, Vn
@ -152,7 +152,7 @@ def LlamaAttention_fast_forward_inference(
A = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32).to(A.dtype)
A = torch.matmul(A, Vnn)
A = A.transpose(1, 2)
A = A.view(bsz, 1, self.hidden_size)
A = A.reshape(bsz, 1, self.hidden_size)
A = original_apply_o(self, A)
return A, (Kn, Vn)
pass