Support llama3.1

Summary: Add scaled RoPE

Test Plan: Test official checkpoint and gives meaningful result.

Reviewers:

Subscribers:

Tasks:

Tags:

ghstack-source-id: 1e29544
Pull Request resolved: #4376

Author: larryliu0820

examples/models/llama2/llama_transformer.py

@@ -101,6 +101,7 @@ class ModelArgs:
         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.
+    use_scaled_rope: bool = False  # Use scaled RoPE, introduced in llama3.1.
     # Additional Model Metadata needed at runtime
     bos_idx: int = 1
     eos_idx: int = 3
@@ -462,6 +463,7 @@ def __init__(self, params: ModelArgs):
                 else params.max_seq_len * 2  # Sharded checkpoint.
             ),
             params.rope_freq_base,
+            params.use_scaled_rope,
         )
         self.register_buffer("freqs_cos", freqs_cos, persistent=False)
         self.register_buffer("freqs_sin", freqs_sin, persistent=False)

examples/models/llama2/rope.py

@@ -8,17 +8,46 @@
 # Different RoPE implementations
 
 from typing import Tuple
-
+import math
 import torch
 
 # ======================== Stock Implementation ========================
 
 
-def precompute_freqs_cis(dim: int, end: int, theta: float):
+def apply_scaling(freqs: torch.Tensor):
+    # Values obtained from grid search
+    scale_factor = 8
+    low_freq_factor = 1
+    high_freq_factor = 4
+    old_context_len = 8192  # original llama3 length
+
+    low_freq_wavelen = old_context_len / low_freq_factor
+    high_freq_wavelen = old_context_len / high_freq_factor
+    new_freqs = []
+    for freq in freqs:
+        wavelen = 2 * math.pi / freq
+        if wavelen < high_freq_wavelen:
+            new_freqs.append(freq)
+        elif wavelen > low_freq_wavelen:
+            new_freqs.append(freq / scale_factor)
+        else:
+            assert low_freq_wavelen != high_freq_wavelen
+            smooth = (old_context_len / wavelen - low_freq_factor) / (
+                high_freq_factor - low_freq_factor
+            )
+            new_freqs.append((1 - smooth) * freq / scale_factor + smooth * freq)
+    return torch.tensor(new_freqs, dtype=freqs.dtype, device=freqs.device)
+
+
+def precompute_freqs_cis(
+    dim: int, end: int, theta: float = 10000.0, use_scaled: bool = False
+):
     freqs = 1.0 / (
         theta ** (torch.arange(0, dim, 2, device="cpu")[: (dim // 2)].float() / dim)
     )
     t = torch.arange(end, device=freqs.device)  # pyre-ignore
+    if use_scaled:
+        freqs = apply_scaling(freqs)
     freqs = torch.outer(t, freqs).float()  # pyre-ignore
     freqs_cos = torch.cos(freqs)
     freqs_sin = torch.sin(freqs)