vrr/unsloth
2
0
Fork 0
mirror of https://github.com/unslothai/unsloth.git synced 2026-07-09 15:58:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 29

Fix Mistral

Browse source 
BlockDiagonalCausalMask fix courtesy of https://github.com/Rypo
This commit is contained in:
Daniel Han-Chen 2023-12-16 13:14:33 +11:00
parent 44bf41e113
commit cf5204a52d
2 changed files with 16 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -7,7 +7,7 @@
## 2-5x faster 60% less memory local QLoRA finetuning
* Supports Llama 7b, 13b, 70b, CodeLlama 34b, Mistral 7b, TinyLlama and all Llama archs!
* Llama 7b [Colab T4 example](https://colab.research.google.com/drive/1n-fgduZhRUsSjgpqNtVkXA3rSfE7iBdg?usp=sharing) on 1 T4 2x faster, uses 43% less VRAM (8.4GB) LAION dataset. [Alpaca T4 example](https://colab.research.google.com/drive/1oW55fBmwzCOrBVX66RcpptL3a99qWBxb?usp=sharing) 2x faster on 1 T4, using 6.4GB VRAM.
* Mistral 7b [Colab A100 example](https://colab.research.google.com/drive/1SKrKGV-BZoU4kv5q3g0jtE_OhRgPtrrQ?usp=sharing) on 1 A100 2.2x faster, uses 62% less VRAM (12.4GB).
* Mistral 7b [Colab A100 example](https://colab.research.google.com/drive/1SKrKGV-BZoU4kv5q3g0jtE_OhRgPtrrQ?usp=sharing) on 1 A100 2.2x faster, uses 62% less VRAM (12.4GB). [Colab T4 example](https://colab.research.google.com/drive/15pyLgRN97B_jA56HS0esx56knA9I5tuv?usp=sharing)
* CodeLlama 34b [Colab example](https://colab.research.google.com/drive/1gdHyAx8XJsz2yNV-DHvbHjR1iCef5Qmh?usp=sharing) does not OOM is 1.9x faster, uses 32% less VRAM (27GB).
* Kaggle 2 Tesla T4s 5.28x faster on Alpaca. [Kaggle example](https://www.kaggle.com/danielhanchen/unsloth-laion-t4-ddp)
* All kernels written in [OpenAI's Triton](https://openai.com/research/triton) language.

17
unsloth/models/mistral.py
View file

@ -37,6 +37,17 @@ def MistralAttention_fast_forward(
bsz, q_len, _ = hidden_states.size()
Q, K, V = self.apply_qkv(self, hidden_states)
# Check for inference
if use_cache and past_key_value is not None and q_len == 1:
A, past_key_value = LlamaAttention_fast_forward_inference(
self,
hidden_states,
past_key_value,
position_ids,
)
return A, None, past_key_value
pass
n_heads = self.num_heads
n_groups = self.num_key_value_groups
n_kv_heads = self.num_key_value_heads
@ -152,8 +163,10 @@ def MistralForCausalLM_fast_forward(
elif q_len <= sliding_window:
causal_mask = xformers.attn_bias.LowerTriangularMask()
else:
causal_mask = xformers.attn_bias.BlockDiagonalCausalLocalAttentionMask.\
make_local_attention(window_size = sliding_window)
# Fix from https://github.com/Rypo
causal_mask = xformers.attn_bias.BlockDiagonalCausalMask\
.from_seqlens([qlen]*bsz)\
.make_local_attention(window_size = sliding_window)
pass
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions