Fix Llama 3 Tests (#782)

HuggingFace seems to have changed a few things around in what info they
expect to be stored in the config leading the Llama 3 roundtrip tests to
hit errors. AFAICT, the Torch tests aren't running in CI so this just
fixes the regression!


![image](https://github.com/user-attachments/assets/eec1f2a3-ceb9-443a-911e-ea6476fa91bf)

src/levanter/models/rotary.py

@@ -157,6 +157,7 @@ def to_hf_config(self) -> tuple[float, dict]:
             "low_freq_factor": self.low_freq_factor,
             "high_freq_factor": self.high_freq_factor,
             "original_max_position_embeddings": self.original_max_position_embeddings,
+            "rope_type": "llama3",
         }
 
 

tests/test_llama3.py

@@ -26,9 +26,10 @@ def get_config(vocab_size=1000):
             "eos_token_id": 128001,
             "hidden_act": "silu",
             "hidden_size": 4096,
+            "head_dim": 64,
             "initializer_range": 0.02,
             "intermediate_size": 14336,
-            "max_position_embeddings": 8192,
+            "max_position_embeddings": 131072,
             "model_type": "llama",
             "num_attention_heads": 32,
             "num_hidden_layers": 32,
@@ -55,6 +56,7 @@ def get_config(vocab_size=1000):
     llama3_8b_config.hidden_size = 16
     llama3_8b_config.intermediate_size = 64
     llama3_8b_config.num_attention_heads = 4
+    llama3_8b_config.head_dim = 4
     llama3_8b_config.num_hidden_layers = 4
     llama3_8b_config.num_key_value_heads = 2
     llama3_8b_config.max_position_embeddings = 128