Add support for NTK-by-Part Rotary Embedding & set correct rotary bas…

…e for Llama-3.1 series
TransformerLensOrg · Oct 25, 2024 · a1feb3f · a1feb3f
1 parent 172cc2a
commit a1feb3f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/transformer_lens/loading_from_pretrained.py b/transformer_lens/loading_from_pretrained.py
@@ -951,7 +951,8 @@ def convert_hf_model_config(model_name: str, **kwargs):
             "rotary_dim": 128,
             "final_rms": True,
             "gated_mlp": True,
-            "use_NTK_by_parts_rope": False,
+            "rotary_base": 500000.0,
+            "use_NTK_by_parts_rope": True,
             "NTK_by_parts_low_freq_factor": 1.0,
             "NTK_by_parts_high_freq_factor": 4.0,
             "NTK_by_parts_factor": 8.0,