Llama 3.1
This commit is contained in:
Daniel Han
parent
0690914c62
commit
eda2343056
3 changed files with 179 additions and 1 deletions
|
|
@ -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
|
||||
"<pad>", # 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,])
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = {}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue