From 6d3a3b42864e6e0bcffa934b47d15182dfbe4340 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Thu, 15 Feb 2024 19:24:07 +1100 Subject: [PATCH] Fix RoPE precision issues --- unsloth/models/llama.py | 80 +++++++++++++++++++++++++++++++++-------- 1 file changed, 66 insertions(+), 14 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 8a37be1b8e..8604298734 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -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