From a3ed0dc4f4d52b01b5499a22946175dff00dfd5b Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Wed, 27 Dec 2023 03:51:31 +1100 Subject: [PATCH] Fix inference --- README.md | 2 +- unsloth/models/llama.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 3d8a460b2a..cf1568f124 100644 --- a/README.md +++ b/README.md @@ -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! diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index de9666bca5..cbbb6b7c9e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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