From f3cd5e192567683545b6f4cfe1cb5f14479c2aaf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jul 2024 10:51:10 -0700 Subject: [PATCH] Mistral Nemo 12b (#777) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Update gemma2.py * Update llama.py * Update llama.py * Update gemma2.py * init * Update gemma2.py * Update gemma2.py * Update _utils.py * Update _utils.py * Update _utils.py * Update _utils.py * Update _utils.py * Update gemma2.py * Update gemma2.py * Update gemma2.py * All RoPE Scaling support * cleanup * Update llama.py * Update llama.py * Update _utils.py * Update _utils.py * exec * exec * Attention_Module * attention_module * imports * exec * Update llama.py * Update llama.py * boolean mask * revert masking * Update llama.py * Update save.py * Update llama.py * Update gemma2.py * Update gemma2.py * Update gemma2.py * Update utils.py * retry * Update gemma2.py * Update gemma2.py * Update gemma2.py * Update _utils.py * Update _utils.py * Update gemma2.py * Update chat_templates.py * Gemma 2 Ollama support * Update llama.py * Update llama.py * error handling * Update _utils.py * Update _utils.py * Stats for debugging * Update _utils.py * Update _utils.py * Debugging * Update tokenizer_utils.py * Update _utils.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update cross_entropy_loss.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Update rms_layernorm.py * Check exec, eval * Update _utils.py * Update _utils.py * Images * Bug fixes * Update pyproject.toml * Bug fixes * Update _utils.py * Update _utils.py * Deprecation fix * Update chat_templates.py * Now permitting use of pre-installed llama.cpp (#763) * Now permitting use of pre-installed llama.cpp * Update save.py --------- Co-authored-by: Giuseppe Strafforello Co-authored-by: Daniel Han * Update save.py * Deprecation & compile * typo * Update chat_templates.py * Update chat_templates.py * train_on_responses_only * Update llama.py * Update llama.py * Update save.py * Update gemma2.py * Flex Attention * typos * Update _utils.py * Update llama.py * Update __init__.py * Update flex_attention.py * Update llama.py * Update llama.py * emulation * Update __init__.py * Update rope_embedding.py * Update flex_attention.py * Update flex_attention.py * Update rope_embedding.py * libdevice * triton_tanh * Update flex_attention.py * Update flex_attention.py * Update flex_attention.py * Update flex_attention.py * Update flex_attention.py * score * Update llama.py * Update flex_attention.py * Update flex_attention.py * Update flex_attention.py * Update flex_attention.py * Update flex_attention.py * Update llama.py * Update flex_attention.py * Update flex_attention.py * Update flex_attention.py * Update flex_attention.py * Flex Attention removal * upload tensorboard training stats to hub if available (#773) * causal_mask * Update llama.py * Update llama.py * Update flex_attention.py * Update _utils.py * Update mapper.py * Update _utils.py --------- Co-authored-by: pepistrafforello Co-authored-by: Giuseppe Strafforello Co-authored-by: Sébastien De Greef --- unsloth/__init__.py | 19 ++++--- unsloth/kernels/__init__.py | 38 +++++++------ unsloth/kernels/cross_entropy_loss.py | 12 ++--- unsloth/kernels/flex_attention.py | 77 +++++++++++++++++++++++++++ unsloth/kernels/geglu.py | 6 +-- unsloth/kernels/rope_embedding.py | 2 +- unsloth/kernels/utils.py | 12 +++++ unsloth/models/_utils.py | 20 ++++++- unsloth/models/gemma2.py | 50 +---------------- unsloth/models/llama.py | 6 +++ unsloth/models/mapper.py | 8 +++ unsloth/save.py | 14 +++++ 12 files changed, 183 insertions(+), 81 deletions(-) create mode 100644 unsloth/kernels/flex_attention.py diff --git a/unsloth/__init__.py b/unsloth/__init__.py index feb550be5f..0490484cb1 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -60,12 +60,6 @@ except: "We have some installation instructions on our Github page.") pass -# Fix up is_bf16_supported https://github.com/unslothai/unsloth/issues/504 -major_version, minor_version = torch.cuda.get_device_capability() -SUPPORTS_BFLOAT16 = (major_version >= 8) -def is_bf16_supported(): return SUPPORTS_BFLOAT16 -torch.cuda.is_bf16_supported = is_bf16_supported - # We support Pytorch 2 # Fixes https://github.com/unslothai/unsloth/issues/38 torch_version = torch.__version__.split(".") @@ -79,6 +73,19 @@ elif (major_torch == 2) and (minor_torch < 2): del os.environ["PYTORCH_CUDA_ALLOC_CONF"] pass +# Torch 2.5 has including_emulation +major_version, minor_version = torch.cuda.get_device_capability() +SUPPORTS_BFLOAT16 = (major_version >= 8) + +if (major_torch == 2) and (minor_torch >= 5): + old_is_bf16_supported = torch.cuda.is_bf16_supported + def is_bf16_supported(including_emulation = False): + return old_is_bf16_supported(including_emulation) + torch.cuda.is_bf16_supported = is_bf16_supported +else: + def is_bf16_supported(): SUPPORTS_BFLOAT16 + torch.cuda.is_bf16_supported = is_bf16_supported +pass # Try loading bitsandbytes and triton import bitsandbytes as bnb diff --git a/unsloth/kernels/__init__.py b/unsloth/kernels/__init__.py index ebea02afd7..c2de979a6f 100644 --- a/unsloth/kernels/__init__.py +++ b/unsloth/kernels/__init__.py @@ -17,24 +17,32 @@ 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_exact_forward_kernel, - geglu_exact_backward_kernel, - geglu_approx_forward_kernel, - geglu_approx_backward_kernel, + geglu_exact_forward_kernel, + geglu_exact_backward_kernel, + geglu_approx_forward_kernel, + geglu_approx_backward_kernel, ) from .fast_lora import ( - get_lora_parameters, - get_lora_parameters_bias, - apply_lora_mlp_swiglu, - apply_lora_mlp_geglu_exact, - apply_lora_mlp_geglu_approx, - apply_lora_qkv, - apply_lora_o, + get_lora_parameters, + get_lora_parameters_bias, + apply_lora_mlp_swiglu, + apply_lora_mlp_geglu_exact, + apply_lora_mlp_geglu_approx, + apply_lora_qkv, + apply_lora_o, ) from .utils import fast_dequantize, fast_gemv, QUANT_STATE, fast_linear_forward, matmul_lora -try: - print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.") -except: - print("Unsloth: Will patch your computer to enable 2x faster free finetuning.") +from .flex_attention import HAS_FLEX_ATTENTION, slow_attention_softcapping + +if HAS_FLEX_ATTENTION: + from .flex_attention import ( + FLEX_ATTENTION_PADDING, + ) +pass + +try: + print("🦥 Unsloth: Will patch your computer to enable 2x faster free finetuning.") +except: + print("Unsloth: Will patch your computer to enable 2x faster free finetuning.") pass diff --git a/unsloth/kernels/cross_entropy_loss.py b/unsloth/kernels/cross_entropy_loss.py index dc1ad269fe..6074a51538 100644 --- a/unsloth/kernels/cross_entropy_loss.py +++ b/unsloth/kernels/cross_entropy_loss.py @@ -15,7 +15,7 @@ import triton import triton.language as tl import torch -from .utils import calculate_settings, MAX_FUSED_SIZE +from .utils import calculate_settings, MAX_FUSED_SIZE, triton_tanh from transformers.models.llama.modeling_llama import logger @@ -63,7 +63,7 @@ def _cross_entropy_forward( label_idx = tl.load(labels_ptr).to(tl.int32) logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")) # Do logit softcapping for Gemma 2: t * tanh(1/t * x) - if DO_SOFTCAPPING: logits = SOFTCAP * tl.math.tanh(logits / SOFTCAP) + if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits / SOFTCAP) logits = logits.to(tl.float32) c = tl.max(logits, 0) @@ -72,7 +72,7 @@ def _cross_entropy_forward( if label_idx != -100: x = tl.load(logits_ptr + label_idx) # Do logit softcapping for Gemma 2: t * tanh(1/t * x) - if DO_SOFTCAPPING: x = SOFTCAP * tl.math.tanh(x / SOFTCAP) + if DO_SOFTCAPPING: x = SOFTCAP * triton_tanh(x / SOFTCAP) loss = logsumexp - x.to(tl.float32) else: loss = 0.0 @@ -131,7 +131,7 @@ def _chunked_cross_entropy_forward( label_idx = tl.load(labels_ptr).to(tl.int32) logits = tl.load(logits_ptr + col_offsets, mask = mask, other = -float("inf")) # Do logit softcapping for Gemma 2: t * tanh(1/t * x) - if DO_SOFTCAPPING: logits = SOFTCAP * tl.math.tanh(logits / SOFTCAP) + if DO_SOFTCAPPING: logits = SOFTCAP * triton_tanh(logits / SOFTCAP) logits = logits.to(tl.float32) c = tl.max(logits, 0) @@ -143,7 +143,7 @@ def _chunked_cross_entropy_forward( if label_idx != -100: x = tl.load(logits_ptr + label_idx).to(tl.float32) # Do logit softcapping for Gemma 2: t * tanh(1/t * x) - if DO_SOFTCAPPING: x = SOFTCAP * tl.math.tanh(x / SOFTCAP) + if DO_SOFTCAPPING: x = SOFTCAP * triton_tanh(x / SOFTCAP) loss = -1.0 * x.to(tl.float32) else: loss = 0.0 @@ -198,7 +198,7 @@ def _cross_entropy_backward( # Do logit softcapping for Gemma 2: t * tanh(1/t * x) if DO_SOFTCAPPING: # d/dx [t * tanh(1/t * x)] = 1 - tanh^2(1/t * x) - partial = tl.math.tanh(x / SOFTCAP) + partial = triton_tanh(x / SOFTCAP) x = SOFTCAP * partial pass diff --git a/unsloth/kernels/flex_attention.py b/unsloth/kernels/flex_attention.py new file mode 100644 index 0000000000..1eb2486998 --- /dev/null +++ b/unsloth/kernels/flex_attention.py @@ -0,0 +1,77 @@ +# 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 torch +from functools import lru_cache +from transformers.models.llama.modeling_llama import logger + +torch_compile_options = { + "epilogue_fusion" : True, + "max_autotune" : True, + "shape_padding" : True, + "trace.enabled" : False, # Output Triton kernel outputs! + "triton.cudagraphs" : False, +} + +# Flex Attention supported from torch 2.5 onwards only +import torch.nn.attention +if hasattr(torch.nn.attention, "flex_attention"): + import torch.nn.attention.flex_attention + from torch.nn.attention.flex_attention import flex_attention + from torch.nn.attention.flex_attention import create_block_mask + FLEX_ATTENTION_PADDING = getattr( + torch.nn.attention.flex_attention, + "_DEFAULT_SPARSE_BLOCK_SIZE", + 1, + ) + flex_attention = torch.compile(flex_attention, dynamic = False) + HAS_FLEX_ATTENTION = True +else: + HAS_FLEX_ATTENTION = False +pass + +# Logit softcapping +@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) +def slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, q_len): + n_heads = self.num_heads + head_dim = self.head_dim + n_kv_heads = self.num_key_value_heads + n_groups = self.num_key_value_groups + + # Grouped query attention + K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim) + V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim) + K = K.reshape(bsz, n_heads, q_len, head_dim) + V = V.reshape(bsz, n_heads, q_len, head_dim) + + # See https://github.com/google/gemma_pytorch/commit/03e657582d17cb5a8617ebf333c1c16f3694670e + # Gemma 9b should use 256 and not 224 (hs / nah). 27b uses the below + # We default to using the config file itself + # s = self.config.hidden_size // self.config.num_attention_heads + s = self.config.query_pre_attn_scalar + t = self.config.attn_logit_softcapping + + Q = Q * torch.tensor(s**-0.5, dtype = Q.dtype) # Follow Keras exactly + A = torch.matmul(Q, K.transpose(2, 3)) + A = t * torch.tanh(A / t) # Logit softcapping + A += causal_mask[:q_len, :q_len] + # Much slower in torch compile! + # A.masked_fill_(causal_mask[:q_len, :q_len], -float("inf")) + A = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32).to(Q.dtype) + A = torch.matmul(A, V) + A = A.transpose(1, 2).contiguous() + A = A.reshape(bsz, q_len, n_heads*head_dim) + return A +pass + diff --git a/unsloth/kernels/geglu.py b/unsloth/kernels/geglu.py index 006e8c0f34..9fedae769e 100644 --- a/unsloth/kernels/geglu.py +++ b/unsloth/kernels/geglu.py @@ -15,7 +15,7 @@ import triton import triton.language as tl import torch -from .utils import calculate_settings +from .utils import calculate_settings, triton_tanh @triton.jit @@ -119,7 +119,7 @@ def _approx_forward_kernel(e, g, h, n_elements, BLOCK_SIZE : tl.constexpr,): g_row = tl.load(g + offsets, mask = mask, other = 0)#.to(tl.float32) f_row = 0.5 * e_row * ( - tl.math.tanh(s * e_row * (1.0 + 0.044715 * e_row * e_row)) \ + triton_tanh(s * e_row * (1.0 + 0.044715 * e_row * e_row)) \ + 1.0 ) f_row = f_row.to(g_row.dtype) # Exact copy from HF @@ -168,7 +168,7 @@ def _approx_backward_kernel(DW, e, g, n_elements, BLOCK_SIZE : tl.constexpr,): s = 0.7978845608028654 # math.sqrt(2 / math.pi) a = s * e_row # a = sqrt(2 / pi) * x b = a * 0.044715 * e_row * e_row # b = a * 0.044715 * x^2 - T = 1.0 + tl.math.tanh(a + b) + T = 1.0 + triton_tanh(a + b) T2 = 0.5 * T # Q = 0.5 * -T * (T - 2.0) * (a + 3.0 * b) Q2 = -T2 * (T - 2.0) * (a + 3.0 * b) diff --git a/unsloth/kernels/rope_embedding.py b/unsloth/kernels/rope_embedding.py index b32d75ebf3..2934ac41c9 100644 --- a/unsloth/kernels/rope_embedding.py +++ b/unsloth/kernels/rope_embedding.py @@ -16,7 +16,6 @@ import triton import triton.language as tl import torch from .utils import calculate_settings - ROPE_GROUP_SIZE = 4 @triton.heuristics({"BACKWARD_PASS": lambda args: args["BACKWARD_PASS"],}) @@ -36,6 +35,7 @@ def _rope_embedding( RoPE is Q * cos + rotate_half(Q) * sin See our blog post for more info """ + ROPE_GROUP_SIZE = 4 row_position = tl.program_id(0) group_head_position = tl.program_id(1) col_offsets = tl.arange(0, BLOCK_SIZE) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 4b789001f5..23be372217 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -28,6 +28,18 @@ else: pass +# tl.math.tanh now is libdevice.tanh +from packaging.version import Version +import triton +if Version(triton.__version__) >= Version("3.0.0"): + from triton.language.extra import libdevice + triton_tanh = libdevice.tanh +else: + import triton.language as tl + triton_tanh = tl.math.tanh +pass + + def calculate_settings(n): BLOCK_SIZE = next_power_of_2(n) if BLOCK_SIZE > MAX_FUSED_SIZE: diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b224e85a89..025daec136 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -145,7 +145,15 @@ from xformers import __version__ as xformers_version # Temporarily disable 0.0.27 and higher - inference issues if Version(xformers_version) >= Version("0.0.27"): raise ImportError( - f"Unsloth: Your xformers version of {xformers_version} is too new.\n"\ + "Unsloth: If you are in Colab, we updated the top cell install instructions - please change it to below "\ + "then press Disconnect Runtime and then Restart it.\n"\ + "\n"\ + "%%capture\n" + "# Installs Unsloth, Xformers (Flash Attention) and all other packages!\n" + '!pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"\n' + '!pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes\n'\ + '\n'\ + f"Otherwise in local machines, your xformers version of {xformers_version} is too new.\n"\ 'Please downgrade xformers via `pip install --force-reinstall "xformers<0.0.27"' ) pass @@ -154,7 +162,15 @@ pass from trl import __version__ as trl_version if Version(xformers_version) >= Version("0.9.0"): raise ImportError( - f"Unsloth: Your TRL version of {trl_version} is too new.\n"\ + "Unsloth: If you are in Colab, we updated the top cell install instructions - please change it to below "\ + "then press Disconnect Runtime and then Restart it.\n"\ + "\n"\ + "%%capture\n" + "# Installs Unsloth, Xformers (Flash Attention) and all other packages!\n" + '!pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"\n' + '!pip install --no-deps "xformers<0.0.27" "trl<0.9.0" peft accelerate bitsandbytes\n'\ + '\n'\ + f"Otherwise in local machines, your TRL version of {trl_version} is too new.\n"\ 'Please downgrade TRL via `pip install --force-reinstall "trl<0.9.0"' ) pass diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py index 9c055ff09e..0d21c47b00 100644 --- a/unsloth/models/gemma2.py +++ b/unsloth/models/gemma2.py @@ -70,52 +70,6 @@ def fast_rms_layernorm_gemma2_compiled(layernorm, X, gemma = True): pass -# Flex Attention in torch 2.5 and higher -# try: -# from torch.nn.attention._flex_attention import _flex_attention -# from functools import lru_cache -# @lru_cache -# def create_block_mask_from_score_mod(score_mod, B, H, M, N): -# SPARSE_BLOCK = 128 -# block_mask = _create_block_mask(score_mod, B, H, M, N, device = "cuda:0") -# return block_mask - - -# Logit softcapping -@torch.compile(fullgraph = True, dynamic = True, options = torch_compile_options) -def gemma2_attention(Q, K, V, causal_mask, self, bsz, q_len): - n_heads = self.num_heads - head_dim = self.head_dim - n_kv_heads = self.num_key_value_heads - n_groups = self.num_key_value_groups - - # Grouped query attention - K = K[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim) - V = V[:, :, None, :, :].expand(bsz, n_kv_heads, n_groups, q_len, head_dim) - K = K.reshape(bsz, n_heads, q_len, head_dim) - V = V.reshape(bsz, n_heads, q_len, head_dim) - - # See https://github.com/google/gemma_pytorch/commit/03e657582d17cb5a8617ebf333c1c16f3694670e - # Gemma 9b should use 256 and not 224 (hs / nah). 27b uses the below - # We default to using the config file itself - # s = self.config.hidden_size // self.config.num_attention_heads - s = self.config.query_pre_attn_scalar - t = self.config.attn_logit_softcapping - - Q = Q * torch.tensor(s**-0.5, dtype = Q.dtype) # Follow Keras exactly - A = torch.matmul(Q, K.transpose(2, 3)) - A = t * torch.tanh(A / t) # Logit softcapping - A += causal_mask[:q_len, :q_len] - # Much slower in torch compile! - # A.masked_fill_(causal_mask[:q_len, :q_len], -float("inf")) - A = torch.nn.functional.softmax(A, dim = -1, dtype = torch.float32).to(Q.dtype) - A = torch.matmul(A, V) - A = A.transpose(1, 2).contiguous() - A = A.reshape(bsz, q_len, n_heads*head_dim) - return A -pass - - # Logit softcapping def Gemma2Attention_fast_forward( self, @@ -172,8 +126,8 @@ def Gemma2Attention_fast_forward( V = torch.cat([past_key_value[1], V], dim = 2) pass past_key_value = (K, V) if use_cache else None - - A = gemma2_attention(Q, K, V, causal_mask, self, bsz, kv_seq_len) + + A = slow_attention_softcapping(Q, K, V, causal_mask, self, bsz, kv_seq_len) A = self.apply_o(self, A) return A, None, past_key_value pass diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 2d888b8c0e..9bea364ca4 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -662,6 +662,12 @@ def LlamaModel_fast_forward( offloaded_gradient_checkpointing = True pass + # Check for Flex Attention + # if IS_GEMMA2 and HAS_FLEX_ATTENTION: + # if not (seq_length % FLEX_ATTENTION_PADDING == 0): + # USE_FLEX_ATTENTION = True + + # Gemma2 has alternating SWA and global attn if IS_GEMMA2 and not hasattr(self, "SWA_mask"): n = self.config.max_position_embeddings diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index 31b3ab6df6..38cbdbe992 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -210,6 +210,14 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/Phi-3-mini-4k-instruct-v0-bnb-4bit" : ( # Old Phi pre July "unsloth/Phi-3-mini-4k-instruct-v0", ), + "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit" : ( # New 12b Mistral models + "unsloth/Mistral-Nemo-Instruct-2407", + "mistralai/Mistral-Nemo-Instruct-2407", + ), + "unsloth/Mistral-Nemo-Base-2407-bnb-4bit" : ( # New 12b Mistral models + "unsloth/Mistral-Nemo-Base-2407", + "mistralai/Mistral-Nemo-Base-2407", + ), } INT_TO_FLOAT_MAPPER = {} diff --git a/unsloth/save.py b/unsloth/save.py index c8e791c703..a5904efc1b 100644 --- a/unsloth/save.py +++ b/unsloth/save.py @@ -1359,6 +1359,20 @@ def upload_to_huggingface( uploaded_location = file_location pass + # find ftevent file from tensorboard and upload it + import glob + ftevent_files = glob.glob("*out.tfevents*", recursive = True) + if len(ftevent_files) > 0: + print("Unsloth: Uploading tensorboard files... Please wait...", file_location + "*out.tfevents*") + for ftevent_file in ftevent_files: + hf_api.upload_file( + path_or_fileobj = ftevent_file, + path_in_repo = ftevent_file.replace(file_location, ""), + repo_id = save_directory, + repo_type = "model", + commit_message = "(Trained with Unsloth)", + ) + hf_api.upload_file( path_or_fileobj = file_location, path_in_repo = uploaded_location,