Fix RoPE alpha after refactor in #4d25874

Author: turboderp

exllamav2/device.py

@@ -123,20 +123,13 @@ def prepare_sincos(self):
             self.cos = self.sin
             return
 
-        base = cfg.rotary_embedding_base
-        alpha = cfg.scale_alpha_value or 1.0
-        scale = cfg.scale_pos_emb or 1.0
-
-        # Alpha scaling for any rope_scaling type
-
-        if alpha != 1.0: base *= alpha ** (cfg.head_dim / (cfg.head_dim - 2))
-
         # RoPE params
 
         inv_freq, scaling_factor = rope.get_rope_params(device, cfg)
 
         # Common
 
+        scale = cfg.scale_pos_emb or 1.0
         t = torch.arange(cfg.max_seq_len, device = device, dtype = torch.float32)
         if scale != 1.0: t /= scale
 

exllamav2/rope.py

@@ -15,6 +15,8 @@ def get_rope_params_su(
 ):
     head_dim = cfg.head_dim
     base = cfg.rotary_embedding_base
+    if cfg.scale_alpha_value and cfg.scale_alpha_value != 1.0:
+        base *= cfg.scale_alpha_value ** (cfg.head_dim / (cfg.head_dim - 2))
 
     a = cfg.max_seq_len
     b = cfg.original_max_seq_len
@@ -28,7 +30,6 @@ def get_rope_params_su(
     inv_freq = 1.0 / (ext_factors * base ** (torch.arange(0, head_dim, 2, device = device).float() / head_dim))
     return inv_freq, scaling_factor
 
-
 # Llama 3.1
 
 def get_rope_params_llama3(
@@ -37,6 +38,8 @@ def get_rope_params_llama3(
 ):
     head_dim = cfg.head_dim
     base = cfg.rotary_embedding_base
+    if cfg.scale_alpha_value and cfg.scale_alpha_value != 1.0:
+        base *= cfg.scale_alpha_value ** (cfg.head_dim / (cfg.head_dim - 2))
 
     def apply_scaling(
         freqs: torch.Tensor,
@@ -80,6 +83,9 @@ def get_rope_params_yarn(
 ):
     head_dim = cfg.head_dim
     base = cfg.rotary_embedding_base
+    if cfg.scale_alpha_value and cfg.scale_alpha_value != 1.0:
+        base *= cfg.scale_alpha_value ** (cfg.head_dim / (cfg.head_dim - 2))
+
     yarn_max_position_embeddings = cfg.max_seq_len
 
     # Only activate if longer than original ctx
@@ -146,6 +152,8 @@ def get_rope_params_default(
 ):
     head_dim = cfg.head_dim
     base = cfg.rotary_embedding_base
+    if cfg.scale_alpha_value and cfg.scale_alpha_value != 1.0:
+        base *= cfg.scale_alpha_value ** (cfg.head_dim / (cfg.head_dim - 2))
 
     inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2, device = device).float() / head_dim))
     return inv_freq, 1.0