fix llama3 rotary in pytorch engine (#2444)

InternLM · Sep 11, 2024 · edcdd8e · edcdd8e
1 parent e681077
commit edcdd8e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/lmdeploy/pytorch/backends/default/rotary_embedding.py b/lmdeploy/pytorch/backends/default/rotary_embedding.py
@@ -134,6 +134,7 @@ def __init__(
 
         inv_freq = self.inv_freq
         factor = self.scaling_factor
+
         wavelen = 2 * math.pi / inv_freq
         # wavelen < high_freq_wavelen: do nothing
         # wavelen > low_freq_wavelen: divide by factor
@@ -149,7 +150,8 @@ def __init__(
                                                             low_freq_wavelen)
         inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq,
                                      inv_freq_llama)
-        self.register_buffer('inv_freq_llama', inv_freq_llama)
+        self.scaling_factor = 1.0
+        self.register_buffer('inv_freq', inv_freq_llama)
 
 
 def yarn_find_correction_dim(num_rotations,