diff --git a/README.md b/README.md index 05977bad73..c666f2d9c8 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and | Unsloth supports | Free Notebooks | Performance | Memory use | |-----------|---------|--------|----------| -| **Llama 3 (8B)** | [▶️ Start for free](https://colab.research.google.com/drive/135ced7oHytdxu3N2DNe1Z0kqjyYIkDXp?usp=sharing) | 2x faster | 60% less | +| **Llama 3.1 (8B)** | [▶️ Start for free](https://colab.research.google.com/drive/135ced7oHytdxu3N2DNe1Z0kqjyYIkDXp?usp=sharing) | 2x faster | 60% less | | **Mistral Nemo (12B)** | [▶️ Start for free](https://colab.research.google.com/drive/17d3U-CAIwzmbDRqbZ9NnpHxCkmXB6LZ0?usp=sharing) | 2x faster | 60% less | | **Gemma 2 (9B)** | [▶️ Start for free](https://colab.research.google.com/drive/1vIrqH5uYDQwsJ4-OO3DErvuv4pBgVwk4?usp=sharing) | 2x faster | 63% less | | **Phi-3 (mini)** | [▶️ Start for free](https://colab.research.google.com/drive/1lN6hPQveB_mHSnTOYifygFcrO8C1bxq4?usp=sharing) | 2x faster | 50% less | @@ -32,13 +32,14 @@ All notebooks are **beginner friendly**! Add your dataset, click "Run All", and | **DPO Zephyr** | [▶️ Start for free](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing) | 1.9x faster | 43% less | | **TinyLlama** | [▶️ Start for free](https://colab.research.google.com/drive/1AZghoNBQaMDgWJpi4RbffGM1h6raLUj9?usp=sharing) | 3.9x faster | 74% less | -- **Kaggle Notebooks** for [Llama 3 (8B)](https://www.kaggle.com/code/danielhanchen/kaggle-llama-3-8b-unsloth-notebook), [Gemma 2 (9B)](https://www.kaggle.com/code/danielhanchen/kaggle-gemma-7b-unsloth-notebook/), [Mistral (7B)](https://www.kaggle.com/code/danielhanchen/kaggle-mistral-7b-unsloth-notebook) +- **Kaggle Notebooks** for [Llama 3.1 (8B)](https://www.kaggle.com/code/danielhanchen/kaggle-llama-3-8b-unsloth-notebook), [Gemma 2 (9B)](https://www.kaggle.com/code/danielhanchen/kaggle-gemma-7b-unsloth-notebook/), [Mistral (7B)](https://www.kaggle.com/code/danielhanchen/kaggle-mistral-7b-unsloth-notebook) - Run [Llama 3 conversational notebook](https://colab.research.google.com/drive/1XamvWYinY6FOSX9GLvnqSjjsNflxdhNc?usp=sharing) and [Mistral v0.3 ChatML](https://colab.research.google.com/drive/15F1xyn8497_dUbxZP4zWmPZ3PJx1Oymv?usp=sharing) - This [text completion notebook](https://colab.research.google.com/drive/1ef-tab5bhkvWmBOObepl1WgJvfvSzn5Q?usp=sharing) is for continued pretraining / raw text - This [continued pretraining notebook](https://colab.research.google.com/drive/1tEd1FrOXWMnCU9UIvdYhs61tkxdMuKZu?usp=sharing) is for learning another language - Click [here](https://github.com/unslothai/unsloth/wiki) for detailed documentation for Unsloth. ## 🦥 Unsloth.ai News +- 📣 NEW! [Llama 3.1 8b, 70b](https://colab.research.google.com/drive/135ced7oHytdxu3N2DNe1Z0kqjyYIkDXp?usp=sharing) both Base and Instruct now supported - 📣 NEW! [Mistral Nemo-12b](https://colab.research.google.com/drive/17d3U-CAIwzmbDRqbZ9NnpHxCkmXB6LZ0?usp=sharing) both Base and Instruct now supported - 📣 NEW! [Gemma-2-9b](https://colab.research.google.com/drive/1vIrqH5uYDQwsJ4-OO3DErvuv4pBgVwk4?usp=sharing) and Gemma-2-27b now supported - 📣 UPDATE! [Phi-3 mini](https://colab.research.google.com/drive/1hhdhBa1j_hsymiW9m-WzxQtgqTH_NHqi?usp=sharing) model updated. [Phi-3 Medium](https://colab.research.google.com/drive/1hhdhBa1j_hsymiW9m-WzxQtgqTH_NHqi?usp=sharing) 2x faster finetuning. diff --git a/pyproject.toml b/pyproject.toml index 29b35577e9..829b35ad3b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,7 +35,7 @@ exclude = ["images*"] huggingface = [ "packaging", "tyro", - "transformers>=4.42.3", + "transformers>=4.43.1", "datasets>=2.16.0", "sentencepiece>=0.2.0", "tqdm", @@ -188,7 +188,7 @@ colab-ampere-torch220 = [ colab-new = [ "packaging", "tyro", - "transformers>=4.42.3", + "transformers>=4.43.1", "datasets>=2.16.0", "sentencepiece>=0.2.0", "tqdm", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 466a5fee70..b021e89e9d 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -33,6 +33,7 @@ __all__ = [ "unsloth_offloaded_gradient_checkpoint", "torch_compile_options", "patch_linear_scaling", + "patch_llama_rope_scaling", "check_nvidia", "create_boolean_mask", "torch_amp_custom_fwd", @@ -332,7 +333,13 @@ def patch_tokenizer(model, tokenizer): Check if pad_token is not the same as eos_token otherwise the loss will ignore it!! Fixes https://github.com/unslothai/unsloth/issues/5 """ - possible_reserved_tokens = ("<|reserved", "<|placeholder", "[control") + possible_reserved_tokens = ( + "<|reserved", # Llama-3 + "<|placeholder", # Phi-3 + "[control", # Forgot where lol + "", # Mistral Nemo + "<|finetune_right_pad_id|>", # Llama-3.1 + ) if model is not None: model.config.update({"unsloth_version" : __version__}) @@ -745,7 +752,7 @@ def patch_linear_scaling( fix_rope_function = """ if getattr(self.config, "rope_scaling", None) is None: self.rotary_emb = {rope_function}( - self.head_dim, + dim = self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, ) @@ -754,7 +761,7 @@ def patch_linear_scaling( scaling_factor = self.config.rope_scaling["factor"] if scaling_type == "linear": self.rotary_emb = {scaled_rope_function}( - self.head_dim, + dim = self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope_theta, @@ -779,6 +786,91 @@ def patch_linear_scaling( pass +# Patches for Llama-3 LlamaExtendedRotaryEmbedding +def patch_llama_rope_scaling( + model_name = "llama", + rope_module = None, + scaled_rope_module = None, + extended_rope_module = None, + attention_module = None, +): + assert(\ + rope_module is not None and \ + scaled_rope_module is not None and \ + extended_rope_module is not None + ) + assert(attention_module is not None) + + rope_name = rope_module.__name__ + scaled_rope_name = scaled_rope_module.__name__ + model_filepath = f"transformers.models.{model_name}.modeling_{model_name}" + exec_code = \ + f"import torch.nn as nn\n"\ + f"from typing import Union, Optional, List, Any, Callable, Tuple\n"\ + f"from {model_filepath} import logger, "\ + f"{model_name.title()}Attention, {model_name.title()}Config" + + try: + function = inspect.getsource(attention_module.__init__) + except: + # Most likely already patched! + return None, None + where = function.find("def") + function = function.split("\n") + function = "\n".join(x[where:] for x in function) + init_name = f"{model_name.title()}Attention__init__" + function = function.replace("def __init__", f"def {init_name}") + function = function.replace( + "super().__init__()", + f"super({model_name.title()}Attention, self).__init__()", + ) + fix_rope_function = """ + if getattr(self.config, "rope_scaling", None) is None: + self.rotary_emb = {rope_function}( + dim = self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + scaling_type1 = self.config.rope_scaling.get("type", None) + scaling_type2 = self.config.rope_scaling.get("rope_type", None) + scaling_type = scaling_type1 if scaling_type1 is not None else scaling_type2 + scaling_factor = self.config.rope_scaling.get("factor") + + if scaling_type == "linear": + self.rotary_emb = {scaled_rope_function}( + dim = self.head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + ) + elif scaling_type == "llama3": + self.rotary_emb = {extended_rope_function}( + dim = self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + raise ValueError(f"Unknown RoPE scaling type {{scaling_type}}") + pass + """ + fix_rope_function = fix_rope_function.format( + rope_function = rope_module.__name__, + scaled_rope_function = scaled_rope_module.__name__, + extended_rope_function = extended_rope_module.__name__, + ) + rotary_emb = re.findall( + "self.rotary_emb = .+?\)", function, + flags = re.DOTALL | re.MULTILINE, + ) + if len(rotary_emb) == 0: return None, function + rotary_emb = rotary_emb[0] + function = function.replace(rotary_emb, fix_rope_function, 1) + function = exec_code + "\n\n" + function + return init_name, function +pass + + def check_nvidia(): # Unsloth doesn't work yet on AMD devices - we're working on it! output = np.array([0,]) diff --git a/unsloth/models/gemma.py b/unsloth/models/gemma.py index ce89ad3be6..e3f1e615db 100644 --- a/unsloth/models/gemma.py +++ b/unsloth/models/gemma.py @@ -205,8 +205,11 @@ class GemmaFixedRotaryEmbedding(torch.nn.Module): # Fixes https://github.com/huggingface/transformers/pull/28837 # https://github.com/microsoft/DeepSpeed/issues/4932 # The precision of RoPE buffers is not correct, so we cast to int64. - def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, + config = None, # [TODO] Hack to pass in config - need to remove later + ): super().__init__() + if config is not None: return # [TODO] Hack to pass in config - need to remove later self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base @@ -264,9 +267,11 @@ class GemmaFixedLinearScalingRotaryEmbedding(GemmaFixedRotaryEmbedding): # Fixes https://github.com/huggingface/transformers/pull/28837 # https://github.com/microsoft/DeepSpeed/issues/4932 # The precision of RoPE buffers is not correct, so we cast to int64. - def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): + def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0, + config = None, # [TODO] Hack to pass in config - need to remove later + ): self.scaling_factor = scaling_factor - super().__init__(dim, max_position_embeddings, base, device) + super().__init__(dim = dim, max_position_embeddings = max_position_embeddings, base = base, device = device, config = config) pass def _set_cos_sin_cache(self, seq_len, device, dtype): diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index ff51b90b84..338ae0a7ce 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -977,8 +977,19 @@ class LlamaRotaryEmbedding(torch.nn.Module): # Fixes https://github.com/huggingface/transformers/pull/28837 # https://github.com/microsoft/DeepSpeed/issues/4932 # The precision of RoPE buffers is not correct, so we cast to int64. - def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, + config = None, # [TODO] Hack to pass in config - need to remove later + ): super().__init__() + if config is not None: + # [TODO] Hack to pass in config - need to remove later + base = config.rope_theta + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + dim = int((config.hidden_size // config.num_attention_heads)) + device = "cuda" + max_position_embeddings = config.max_position_embeddings + pass + self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base @@ -1030,9 +1041,11 @@ class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): # Fixes https://github.com/huggingface/transformers/pull/28837 # https://github.com/microsoft/DeepSpeed/issues/4932 # The precision of RoPE buffers is not correct, so we cast to int64. - def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): + def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0, + config = None, # [TODO] Hack to pass in config - need to remove later + ): self.scaling_factor = scaling_factor - super().__init__(dim, max_position_embeddings, base, device) + super().__init__(dim = dim, max_position_embeddings = max_position_embeddings, base = base, device = device, config = config) pass def _set_cos_sin_cache(self, seq_len, device, dtype): @@ -1052,6 +1065,99 @@ class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): pass +# See https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/rotary_embedding.py#L736 +# For Llama 3.1 +class LlamaExtendedRotaryEmbedding(torch.nn.Module): + def __init__(self, dim = None, max_position_embeddings=2048, base=10000, device=None, + config = None, # [TODO] Hack to pass in config - need to remove later + ): + super().__init__() + if config is not None: + # [TODO] Hack to pass in config - need to remove later + base = config.rope_theta + partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 + dim = int((config.hidden_size // config.num_attention_heads)) + device = "cuda" + max_position_embeddings = config.max_position_embeddings + pass + + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + # Dynamic RoPE we first set it to a max of 4 * 8192 tokens then we iteratively grow this + self.current_rope_size = min(4 * 8192, self.max_position_embeddings) + + # Normal Llama-3 RoPE + inv_freq = 1.0 / ( + self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device="cpu").float() / self.dim) + ) + inv_freq = self.apply_scaling(inv_freq) + self.register_buffer("inv_freq", inv_freq, persistent = False) + + # Build here to make `torch.jit.trace` work. + self._set_cos_sin_cache(seq_len=self.current_rope_size, device=device, dtype=torch.get_default_dtype()) + pass + + def _set_cos_sin_cache(self, seq_len, device, dtype): + # Note: on the original Llama codebase, these tensors are created on the target device (and not on CPU) and + # in FP32. They are applied (multiplied) in FP32 as well. + self.current_rope_size = seq_len + + t = torch.arange(self.current_rope_size, device="cpu", dtype=torch.int64).float() + + freqs = torch.outer(t, self.inv_freq) + # Different from paper, but it uses a different permutation in order to obtain the same calculation + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("cos_cached", emb.cos().to(dtype=dtype, device=device, non_blocking=True), persistent=False) + self.register_buffer("sin_cached", emb.sin().to(dtype=dtype, device=device, non_blocking=True), persistent=False) + pass + + # From https://github.com/meta-llama/llama-models/blob/main/models/llama3_1/api/model.py#L41 + def apply_scaling(self, freqs: torch.Tensor): + # Values obtained from grid search + scale_factor = 8 + low_freq_factor = 1 + high_freq_factor = 4 + old_context_len = 8192 # original llama3 length + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + new_freqs = [] + for freq in freqs: + wavelen = 2 * math.pi / freq + if wavelen < high_freq_wavelen: + new_freqs.append(freq) + elif wavelen > low_freq_wavelen: + new_freqs.append(freq / scale_factor) + else: + assert low_freq_wavelen != high_freq_wavelen + smooth = (old_context_len / wavelen - low_freq_factor) / ( + high_freq_factor - low_freq_factor + ) + new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq) + return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device) + pass + + def forward(self, x, position_ids=None, seq_len=None): + # x: [bs, num_attention_heads, seq_len, head_size] + if seq_len > self.current_rope_size: + self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) + + return ( + self.cos_cached[:seq_len].to(dtype = x.dtype), + self.sin_cached[:seq_len].to(dtype = x.dtype), + ) + pass + + def extend_rope_embedding(self, x, seq_len): + if seq_len <= self.current_rope_size: return + # Iteratively grow by increments of 8192 + self.current_rope_size = int(round(seq_len / 8192)) * 8192 + self._set_cos_sin_cache(self.current_rope_size, device = "cuda:0", dtype = x.dtype) + pass +pass + + def _wrap_fast_inference(generate, device_type, dtype, model): # Wraps inference with bfloat16 / float16 @torch.inference_mode @@ -1108,6 +1214,17 @@ class FastLlamaModel: @staticmethod def pre_patch(): + init_name, function = patch_llama_rope_scaling( + model_name = "llama", + rope_module = LlamaRotaryEmbedding, + scaled_rope_module = LlamaLinearScalingRotaryEmbedding, + extended_rope_module = LlamaExtendedRotaryEmbedding, + attention_module = LlamaAttention, + ) + if init_name is not None: + exec(function, globals()) + LlamaAttention.__init__ = eval(init_name) + pass LlamaAttention .forward = LlamaAttention_fast_forward LlamaSdpaAttention .forward = LlamaAttention_fast_forward LlamaFlashAttention2.forward = LlamaAttention_fast_forward diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 0f170597b1..ece8af2821 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -27,6 +27,7 @@ transformers_version = Version(transformers_version) SUPPORTS_FOURBIT = transformers_version >= Version("4.37") SUPPORTS_GEMMA = transformers_version >= Version("4.38") SUPPORTS_GEMMA2 = transformers_version >= Version("4.42") +SUPPORTS_LLAMA31 = transformers_version >= Version("4.43.1") if SUPPORTS_GEMMA: from .gemma import FastGemmaModel if SUPPORTS_GEMMA2: @@ -130,7 +131,19 @@ class FastLanguageModel(FastLlamaModel): model_type = model_config.model_type - if model_type == "llama": dispatch_model = FastLlamaModel + if model_type == "llama": + scaling_type1 = model_config.rope_scaling.get("type", None) + scaling_type2 = model_config.rope_scaling.get("rope_type", None) + scaling_type = scaling_type1 if scaling_type1 is not None else scaling_type2 + + if scaling_type == "llama3" and not SUPPORTS_LLAMA31: + raise ImportError( + f"Unsloth: Your transformers version of {transformers_version} does not support Llama 3.1.\n"\ + f"The minimum required version is 4.43.1\n"\ + f'Try `pip install --upgrade "transformers>=4.43.1"`\n'\ + f"to obtain the latest transformers build, then restart this session."\ + ) + dispatch_model = FastLlamaModel elif model_type == "mistral": dispatch_model = FastMistralModel elif model_type == "gemma": if not SUPPORTS_GEMMA: diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index 38cbdbe992..fc13c94e84 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -218,6 +218,22 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/Mistral-Nemo-Base-2407", "mistralai/Mistral-Nemo-Base-2407", ), + "unsloth/Meta-Llama-3.1-8B-bnb-4bit" : ( + "unsloth/Meta-Llama-3.1-8B", + "meta-llama/Meta-Llama-3.1-8B", + ), + "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit" : ( + "unsloth/Meta-Llama-3.1-8B-Instruct", + "meta-llama/Meta-Llama-3.1-8B-Instruct", + ), + "unsloth/Meta-Llama-3.1-70B-bnb-4bit" : ( + "unsloth/Meta-Llama-3.1-70B", + "meta-llama/Meta-Llama-3.1-70B", + ), + "unsloth/Meta-Llama-3.1-70B-Instruct-bnb-4bit" : ( + "unsloth/Meta-Llama-3.1-70B-Instruct", + "meta-llama/Meta-Llama-3.1-70B-Instruct", + ), } INT_TO_FLOAT_MAPPER = {}