From eda23430569667a963cbee85c04914cd0523e00e Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 22 Jul 2024 22:58:02 -0700 Subject: [PATCH] Llama 3.1 --- unsloth/models/_utils.py | 93 +++++++++++++++++++++++++++++++++++++++- unsloth/models/llama.py | 73 +++++++++++++++++++++++++++++++ unsloth/models/mapper.py | 14 ++++++ 3 files changed, 179 insertions(+), 1 deletion(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 466a5fee70..c7c779a231 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__}) @@ -779,6 +786,90 @@ 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: + # Hack + if self.config.max_position_embeddings == 131072 + self.rotary_emb = {rope_function}( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + scaling_type = self.config.rope_scaling["type"] + scaling_factor = self.config.rope_scaling.get("factor") + if scaling_type == "linear": + self.rotary_emb = {scaled_rope_function}( + self.head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + ) + elif scaling_type == "extended": + self.rotary_emb = {extended_rope_function}( + 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/llama.py b/unsloth/models/llama.py index ff51b90b84..2d224b3cad 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1052,6 +1052,68 @@ 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(LlamaRotaryEmbedding): + def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + super().__init__() + 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 + + def apply_scaling(self, freqs: torch.Tensor): + scale_factor = 8 + low_freq_factor = 1 + high_freq_factor = 4 + old_context_len = 8192 + + 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 +pass + + def _wrap_fast_inference(generate, device_type, dtype, model): # Wraps inference with bfloat16 / float16 @torch.inference_mode @@ -1108,6 +1170,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/mapper.py b/unsloth/models/mapper.py index 38cbdbe992..462c85f2a1 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -218,6 +218,20 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/Mistral-Nemo-Base-2407", "mistralai/Mistral-Nemo-Base-2407", ), + "unsloth/llama-3.1-8b-bnb-4bit" : ( + "unsloth/llama-3.1-8b", + "meta-llama/Meta-Llama-3.1-8B", + ), + "unsloth/llama-3.1-8b-Instruct-bnb-4bit" : ( + "unsloth/llama-3.1-8b-Instruct", + "meta-llama/Meta-Llama-3.1-8B-Instruct", + ), + "unsloth/llama-3.1-70b-bnb-4bit" : ( + "meta-llama/Meta-Llama-3.1-70B", + ), + "unsloth/llama-3.1-70b-Instruct-bnb-4bit" : ( + "meta-llama/Meta-Llama-3.1-70B-Instruct", + ), } INT_TO_FLOAT_MAPPER = {}