Be able to set RopE.Freq_Base explicitly (#2064)

Summary:
Pull Request resolved: #2064

Currently it is set to 10K all the time. Let's make it possible so that one can set it explicitly.

This should be a no-op for existing models and CI jobs.

bypass-github-export-checks
bypass-github-pytorch-ci-checks
bypass-github-executorch-ci-checks

Reviewed By: larryliu0820

Differential Revision: D54131552

fbshipit-source-id: 46e9eb98bc7bb99c7225d3d7b02ab4bd2ecaed1c

Author: mergennachin

examples/models/llama2/llama_transformer.py

@@ -82,6 +82,7 @@ class ModelArgs:
     use_sdpa_with_kv_cache_op: bool = (
         False  # Use custom sdpa op that updates kv cache in-place
     )
+    rope_freq_base: float = 10000.0  # The base frequency for RoPE
     # Additional Model Metadata needed at runtime
     bos_idx: int = 1
     eos_idx: int = 3
@@ -108,7 +109,7 @@ def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
     )
 
 
-def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+def precompute_freqs_cis(dim: int, end: int, theta: float):
     freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
     t = torch.arange(end, device=freqs.device)  # pyre-ignore
     freqs = torch.outer(t, freqs).float()  # pyre-ignore
@@ -411,7 +412,9 @@ def __init__(self, params: ModelArgs):
         self.use_kv_cache = params.use_kv_cache
 
         freqs_cos, freqs_sin = precompute_freqs_cis(
-            self.params.dim // self.params.n_heads, self.params.max_seq_len
+            params.dim // params.n_heads,
+            params.max_seq_len,
+            params.rope_freq_base,
         )
         self.register_buffer("freqs_cos", freqs_cos, persistent=False)
         self.register_buffer("freqs_sin", freqs_sin, persistent=False)