ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix RoPE precision issues

Browse source
This commit is contained in:
Daniel Han-Chen 2024-02-15 19:24:07 +11:00
commit 6d3a3b4286
1 changed files with 66 additions and 14 deletions
Download patch file Download diff file Expand all files Collapse all files

80
unsloth/models/llama.py
View file

@ -782,30 +782,59 @@ pass
# https://github.com/huggingface/transformers/pull/27931
# https://github.com/huggingface/transformers/blob/v4.37.2/src/transformers/models/llama/modeling_llama.py
class LlamaRotaryEmbedding(torch.nn.Module):
# 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
# inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
# self.register_buffer("inv_freq", inv_freq, persistent=False)
# # Build here to make `torch.jit.trace` work.
# self._set_cos_sin_cache(
# seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
# )
# pass
# def _set_cos_sin_cache(self, seq_len, device, dtype):
# self.max_seq_len_cached = seq_len
# t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
# 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), persistent=False)
# self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
# pass
# 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):
super().__init__()
self.dim = dim
self.max_position_embeddings = max_position_embeddings
self.base = base
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
self.register_buffer("inv_freq", inv_freq, persistent=False)
# Build here to make `torch.jit.trace` work.
self._set_cos_sin_cache(
seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
)
self._set_cos_sin_cache(seq_len=max_position_embeddings, 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.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
inv_freq = 1.0 / (
self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device="cpu").float() / self.dim)
)
t = torch.arange(self.max_seq_len_cached, device="cpu", dtype=torch.int64).float()
freqs = torch.outer(t, self.inv_freq)
freqs = torch.outer(t, 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), persistent=False)
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
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 forward(self, x, seq_len=None):
@ -824,6 +853,26 @@ pass
class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
"""LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
# def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
# self.scaling_factor = scaling_factor
# super().__init__(dim, max_position_embeddings, base, device)
# pass
# def _set_cos_sin_cache(self, seq_len, device, dtype):
# self.max_seq_len_cached = seq_len
# t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
# t = t / self.scaling_factor
# 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), persistent=False)
# self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
# pass
# 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):
self.scaling_factor = scaling_factor
super().__init__(dim, max_position_embeddings, base, device)
@ -831,14 +880,17 @@ class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
def _set_cos_sin_cache(self, seq_len, device, dtype):
self.max_seq_len_cached = seq_len
t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
inv_freq = 1.0 / (
self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64, device="cpu").float() / self.dim)
)
t = torch.arange(self.max_seq_len_cached, device="cpu", dtype=torch.int64).float()
t = t / self.scaling_factor
freqs = torch.outer(t, self.inv_freq)
freqs = torch.outer(t, 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), persistent=False)
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
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
pass