Add missing Block size + Update Configs to not hardcode rope_scaling (#1128)

* Update Configs to not hardcode rope_scaling fields

* Adding explcit error

Author: Jack-Khuu

torchchat/model.py

@@ -10,7 +10,8 @@
 from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
-from typing import Callable, Dict, Optional, Union
+
+from typing import Any, Callable, Dict, Optional, Union
 from abc import ABC, abstractmethod
 
 import torch
@@ -132,7 +133,7 @@ class TransformerArgs:
     ffn_dim_multiplier: Optional[int] = None
     use_tiktoken: bool = False
     max_seq_length: int = 8192
-    use_scaled_rope: bool = False
+    rope_scaling: Optional[Dict[str, Any]] = None
     # For pipeline parallel
     n_stages: int = 1
     stage_idx: int = 0
@@ -418,8 +419,6 @@ def __init__(self, config: TransformerArgs) -> None:
             self.norm = None
             self.output = None
 
-        # self.freqs_cis: Optional[Tensor] = None
-        # self.mask_cache: Optional[Tensor] = None
         self.max_batch_size = -1
         self.max_seq_length = -1
         # For supporting sequence parallel (default is off, thus value of 1)
@@ -444,7 +443,7 @@ def setup_caches(self, max_batch_size, max_seq_length):
             self.config.dim // self.config.n_heads,
             self.config.block_size * 2,
             self.config.rope_base,
-            use_scaled=self.config.use_scaled_rope,
+            rope_scaling=self.config.rope_scaling,
         )
         self.register_buffer("freqs_cis", freqs_cis, persistent=True)
         causal_mask = torch.tril(
@@ -681,12 +680,16 @@ def forward(self, x: Tensor) -> Tensor:
         return output * self.weight
 
 
-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
+def apply_scaling(freqs: torch.Tensor, rope_scaling: Dict[str, Any]):
+    # Check for the presence of the required keys
+    required_keys = {"factor", "low_freq_factor", "high_freq_factor", "original_max_position_embeddings"}
+    if not required_keys.issubset(rope_scaling.keys()):
+        raise ValueError(f"Missing required keys in apply_scaling. Expected: {required_keys}")
+
+    scale_factor = rope_scaling["factor"]
+    low_freq_factor = rope_scaling["low_freq_factor"]
+    high_freq_factor = rope_scaling["high_freq_factor"]
+    old_context_len = rope_scaling["original_max_position_embeddings"]
 
     low_freq_wavelen = old_context_len / low_freq_factor
     high_freq_wavelen = old_context_len / high_freq_factor
@@ -707,16 +710,16 @@ def apply_scaling(freqs: torch.Tensor):
 
 
 def precompute_freqs_cis(
-    n_elem: int, seq_len: int, base: int = 10000, dtype=None, use_scaled: bool = False
+    n_elem: int, seq_len: int, base: int = 10000, dtype=None, rope_scaling: Optional[Dict[str, Any]] = None
 ) -> Tensor:
     if not dtype:
         dtype = get_precision()
     freqs = 1.0 / (
         base ** (torch.arange(0, n_elem, 2)[: (n_elem // 2)].float() / n_elem)
     )
     t = torch.arange(seq_len, device=freqs.device)
-    if use_scaled:
-        freqs = apply_scaling(freqs)
+    if rope_scaling is not None:
+        freqs = apply_scaling(freqs, rope_scaling)
     freqs = torch.outer(t, freqs)
     freqs_cis = torch.polar(torch.ones_like(freqs), freqs)
     cache = torch.stack([freqs_cis.real, freqs_cis.imag], dim=-1)

torchchat/model_params/Meta-Llama-3-70B.json

@@ -1 +1 @@
-{"dim": 8192, "ffn_dim_multiplier": 1.3, "multiple_of": 4096, "n_heads": 64, "n_local_heads": 8, "n_layers": 80, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true}
+{"block_size": 8192, "dim": 8192, "ffn_dim_multiplier": 1.3, "multiple_of": 4096, "n_heads": 64, "n_local_heads": 8, "n_layers": 80, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true}

torchchat/model_params/Meta-Llama-3-8B.json

@@ -1 +1 @@
-{"dim": 4096, "ffn_dim_multiplier": 1.3, "multiple_of": 1024, "n_heads": 32, "n_local_heads": 8, "n_layers": 32, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true}
+{"block_size": 8192, "dim": 4096, "ffn_dim_multiplier": 1.3, "multiple_of": 1024, "n_heads": 32, "n_local_heads": 8, "n_layers": 32, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true}

torchchat/model_params/Meta-Llama-3.1-70B.json

@@ -1 +1 @@
-{"dim": 8192, "ffn_dim_multiplier": 1.3, "multiple_of": 4096, "n_heads": 64, "n_local_heads": 8, "n_layers": 80, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true, "norm_eps": 1e-05, "use_scaled_rope": true}
+{"block_size": 131072, "dim": 8192, "ffn_dim_multiplier": 1.3, "multiple_of": 4096, "n_heads": 64, "n_local_heads": 8, "n_layers": 80, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true, "norm_eps": 1e-05, "rope_scaling": {"factor": 8.0, "low_freq_factor": 1.0, "high_freq_factor": 4.0, "original_max_position_embeddings": 8192}}

torchchat/model_params/Meta-Llama-3.1-8B.json

@@ -1 +1 @@
-{"dim": 4096, "ffn_dim_multiplier": 1.3, "multiple_of": 1024, "n_heads": 32, "n_local_heads": 8, "n_layers": 32, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true, "norm_eps": 1e-05, "use_scaled_rope": true}
+{"block_size": 131072, "dim": 4096, "ffn_dim_multiplier": 1.3, "multiple_of": 1024, "n_heads": 32, "n_local_heads": 8, "n_layers": 32, "rope_base": 500000.0, "vocab_size": 128256, "use_tiktoken": true, "norm_eps": 1e-05, "rope_scaling": {"factor": 8.0, "low_freq_factor": 1.0, "high_freq_factor": 4.0, "original_max_position_embeddings": 8192}}