Skip to content

Update model arg name rope_theta to be consistent with those in llama's website #3147

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion examples/models/llama2/llama_transformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,10 @@ 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
rope_theta: Optional[float] = (
None # The official name to override self.rope_freq_base.
)
rope_freq_base: float = 10000.0 # The base frequency for RoPE. Keep it for BC.
# Additional Model Metadata needed at runtime
bos_idx: int = 1
eos_idx: int = 3
Expand All @@ -99,6 +102,10 @@ def __post_init__(self):
if self.n_kv_heads is None:
self.n_kv_heads = self.n_heads

# rope_theta overrides rope_freq_base since it's the official name.
if self.rope_theta is not None:
self.rope_freq_base = self.rope_theta

if self.use_sdpa_with_kv_cache_op:
assert self.use_kv_cache, "use_sdpa_with_kv_cache_op requires use_kv_cache"

Expand Down