From 04093395ed4bd4987dcd28da57bc325db0d9b41b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 26 Dec 2023 04:32:04 +1100
Subject: [PATCH] Fix RoPE Scaling issues (#52)

* Fix RoPE Scaling

* Update llama.py

* Update llama.py
---
 unsloth/models/llama.py   | 10 ++++++++--
 unsloth/models/mistral.py |  6 +++---
 2 files changed, 11 insertions(+), 5 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 5f9b41d4c..590465ae8 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -369,6 +369,7 @@ def LlamaModel_fast_forward(
         raise ValueError("Unsloth: You have to specify either decoder_input_ids or decoder_inputs_embeds")
 
     seq_length_with_past = seq_length
+    assert(seq_length <= self.max_seq_length)
     past_key_values_length = 0
 
     if past_key_values is not None:
@@ -661,6 +662,9 @@ class FastLlamaModel:
                 bnb_4bit_compute_dtype    = dtype,
             )
 
+        # https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/discussions/12
+        # RoPE Scaling's max_position_embeddings must be updated
+        max_position_embeddings = max(max_seq_length, model_max_seq_length)
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             device_map = device_map,
@@ -668,6 +672,7 @@ class FastLlamaModel:
             quantization_config = bnb_config,
             token = token,
             rope_scaling = rope_scaling,
+            max_position_embeddings = max_position_embeddings,
         )
         tokenizer = AutoTokenizer.from_pretrained(
             model_name,
@@ -685,7 +690,7 @@ class FastLlamaModel:
             layer.self_attn.apply_o   = original_apply_o
         pass
 
-        model.max_seq_length = max_seq_length
+        model.max_seq_length = max_position_embeddings
         return model, tokenizer
     pass
 
@@ -746,7 +751,7 @@ class FastLlamaModel:
         layers_to_transform = None,
         use_gradient_checkpointing = True,
         random_state = 3407,
-        max_seq_length = 2048,
+        max_seq_length = 2048, # not used anymore
         **kwargs,
     ):
         assert(max_seq_length <= model.max_seq_length)
@@ -824,6 +829,7 @@ class FastLlamaModel:
 
         # Patch cross entropy loss labels
         # Fixes https://github.com/unslothai/unsloth/issues/10
+        max_seq_length = model.max_seq_length
         extra_ignored_labels = torch.full((max_seq_length, 1), -100, device = "cuda")
         model.model.extra_ignored_labels = extra_ignored_labels
         internal_model = model
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index f3973c9dd..865493edc 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -125,7 +125,7 @@ def MistralAttention_fast_forward(
         V = V.transpose(1, 2)
 
         # Flash Attention v2 auto supports grouped query attention
-        sliding_window = self.config.sliding_window
+        sliding_window = getattr(self.config, "sliding_window")
         sliding_window = q_len if sliding_window is None else sliding_window
         window = (-1, -1) if (q_len <= sliding_window) else (sliding_window, sliding_window)
         A = flash_attn_func(Q, K, V, causal = True, window_size = window)
@@ -169,7 +169,7 @@ def MistralForCausalLM_fast_forward(
 
     if causal_mask is None:
         bsz, q_len = input_ids.shape
-        sliding_window = self.config.sliding_window
+        sliding_window = getattr(self.config, "sliding_window")
         if sliding_window is None or sliding_window <= 0:
             causal_mask = xformers.attn_bias.LowerTriangularMask()
         elif q_len <= sliding_window:
@@ -312,7 +312,7 @@ class FastMistralModel(FastLlamaModel):
             layer.self_attn.apply_o   = original_apply_o
         pass
 
-        model.max_seq_length = max_seq_length
+        model.max_seq_length = max(max_seq_length, model.config.max_position_embeddings)
         return model, tokenizer
     pass
 pass