Improve rope

litgpt/config.py

@@ -62,6 +62,7 @@ class Config:
  intermediate_size: Optional[int] = None
  rope_condense_ratio: int = 1
  rope_base: int = 10000
+ rope_adjustments: Optional[dict] = None
  n_expert: int = 0
  n_expert_per_token: int = 0
  attention_logit_softcapping: Optional[float] = None
@@ -893,6 +894,7 @@ def norm_class(self) -> Type:
  mlp_class_name="LLaMAMLP",
  intermediate_size=14336,
  rope_base=500000,
+ rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192)
  ),
  # https://huggingface.co/meta-llama/Meta-Llama-3-70B/blob/main/config.json
  dict(
@@ -931,6 +933,7 @@ def norm_class(self) -> Type:
  mlp_class_name="LLaMAMLP",
  intermediate_size=28672,
  rope_base=500000,
+ rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192)
  ),
  # https://huggingface.co/meta-llama/Meta-Llama-3.1-405B/blob/main/config.json
  dict(
@@ -950,8 +953,9 @@ def norm_class(self) -> Type:
  mlp_class_name="LLaMAMLP",
  intermediate_size=53248,
  rope_base=500000,
+ rope_adjustments=dict(factor=8.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192)
  ),
- # https://huggingface.co/meta-llama/Llama-3.1-8B/blob/main/config.json
+ # https://huggingface.co/meta-llama/Llama-3.2-1B/blob/main/config.json
  dict(
  name="Llama-3.2-1B{}",
  hf_config=dict(org="meta-llama", name="Llama-3.2-1B{}"),
@@ -969,6 +973,7 @@ def norm_class(self) -> Type:
  mlp_class_name="LLaMAMLP",
  intermediate_size=8192,
  rope_base=500000,
+ rope_adjustments=dict(factor=32.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192)
  ),
  # https://huggingface.co/meta-llama/Llama-3.2-3B/blob/main/config.json
  dict(
@@ -988,6 +993,7 @@ def norm_class(self) -> Type:
  mlp_class_name="LLaMAMLP",
  intermediate_size=8192,
  rope_base=500000,
+ rope_adjustments=dict(factor=32.0, low_freq_factor=1.0, high_freq_factor=4.0, original_max_seq_len=8192)
  ),
 ]
 for c in llama_3:

litgpt/model.py

@@ -108,12 +108,39 @@ def from_name(cls, name: str, **kwargs: Any) -> Self:
  return cls(Config.from_name(name, **kwargs))
 
  def rope_cache(self, device: Optional[torch.device] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+
+ if self.config.rope_adjustments is None:
+ extra_config = None
+
+ else:
+ adjusted_params_required = ["factor", "low_freq_factor", "high_freq_factor", "original_max_seq_len"]
+ params_present = [param in self.config.rope_adjustments for param in adjusted_params_required]
+ num_params_present = sum(params_present)
+
+ if num_params_present == 0:
+ extra_config = None # uses standard RoPE
+ elif num_params_present == 4:
+ # These parameters should always be used together so that we don't interfere with standard rope
+ extra_config = {
+ "original_max_seq_len": self.config.rope_adjustments["original_max_seq_len"],
+ "factor": self.config.rope_adjustments["factor"],
+ "low_freq_factor": self.config.rope_adjustments["low_freq_factor"],
+ "high_freq_factor": self.config.rope_adjustments["high_freq_factor"],
+ }
+ else:
+ # Some but not all parameters are specified; raise an error
+ raise ValueError(
+ "The following adjusted RoPE parameters are missing in rope_adjustments."
+ "All adjusted RoPE parameters must be specified together."
+ )
+
  return build_rope_cache(
  seq_len=self.max_seq_length,
  n_elem=self.config.rope_n_elem,
  device=device,
  condense_ratio=self.config.rope_condense_ratio,
  base=self.config.rope_base,
+ extra_config=extra_config,
  )
 
  def set_kv_cache(
@@ -410,17 +437,67 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 
 def build_rope_cache(
- seq_len: int, n_elem: int, device: Optional[torch.device] = None, base: int = 10000, condense_ratio: int = 1
+ seq_len: int,
+ n_elem: int,
+ device: Optional[torch.device] = None,
+ base: int = 10000,
+ condense_ratio: int = 1,
+ extra_config: Optional[dict] = None,
 ) -> Tuple[torch.Tensor, torch.Tensor]:
  """Enhanced Transformer with Rotary Position Embedding.
 
  Derived from: https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/
  transformers/rope/__init__.py. MIT License:
  https://github.com/labmlai/annotated_deep_learning_paper_implementations/blob/master/license.
+
+ Args:
+ seq_len (int): Sequence length.
+ n_elem (int): Number of elements (head dimension).
+ device (torch.device, optional): Device for tensor allocations.
+ base (int, optional): Base for computing inverse frequencies.
+ condense_ratio (int, optional): Ratio to condense the position indices.
+ extra_config (dict, optional): Configuration parameters for frequency adjustments (used by Llama 3.1 and 3.2)
+
+ Returns:
+ Tuple[torch.Tensor, torch.Tensor]: Cosine and sine caches for RoPE.
  """
- # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$
+
+ # $\Theta = {\theta_i = 10000^{\frac{2(i-1)}{d}}, i \in [1, 2, ..., \frac{d}{2}]}$ assert n_elem % 2 == 0, "n_elem (head dimension) must be even"
  theta = 1.0 / (base ** (torch.arange(0, n_elem, 2, device=device).float() / n_elem))
 
+ if extra_config is not None:
+ # Extract configuration parameters
+ orig_context_len = extra_config["original_max_seq_len"]
+ factor = extra_config["factor"]
+ low_freq_factor = extra_config["low_freq_factor"]
+ high_freq_factor = extra_config["high_freq_factor"]
+
+ # Compute wavelength thresholds
+ low_freq_wavelen = orig_context_len / low_freq_factor
+ high_freq_wavelen = orig_context_len / high_freq_factor
+
+ # Compute wavelengths corresponding to the inverse frequencies
+ wavelen = 2 * torch.pi / theta
+
+ # Initialize adjusted inverse frequencies
+ adjusted_theta = theta.clone()
+
+ # Low Frequency Region: wavelen > low_freq_wavelen
+ mask_low_freq = wavelen > low_freq_wavelen
+ adjusted_theta[mask_low_freq] = theta[mask_low_freq] / factor
+
+ # Medium Frequency Region: high_freq_wavelen ≤ wavelen ≤ low_freq_wavelen
+ mask_medium_freq = (wavelen >= high_freq_wavelen) & (wavelen <= low_freq_wavelen)
+ # Compute smooth factor for medium frequencies
+ ratio = orig_context_len / wavelen[mask_medium_freq]
+ smooth_factor = (ratio - low_freq_factor) / (high_freq_factor - low_freq_factor)
+ # Interpolate inverse frequencies
+ adjusted_theta[mask_medium_freq] = (
+ (1 - smooth_factor) * (theta[mask_medium_freq] / factor)
+ + smooth_factor * theta[mask_medium_freq]
+ )
+ theta = adjusted_theta
+
  # Create position indexes `[0, 1, ..., seq_len - 1]`
  seq_idx = torch.arange(seq_len, device=device) / condense_ratio
 

tests/test_rope.py

@@ -1,13 +1,17 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
 import torch
-from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXRotaryEmbedding, apply_rotary_pos_emb
+from transformers.models.gpt_neox.modeling_gpt_neox import GPTNeoXRotaryEmbedding
+from transformers.models.gpt_neox.modeling_gpt_neox import apply_rotary_pos_emb as apply_rotary_pos_emb_gptneo
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb as apply_rotary_pos_emb_llama
+from transformers.models.llama.configuration_llama import LlamaConfig
 
 from litgpt.model import apply_rope, build_rope_cache
 
 
 @torch.inference_mode()
-def test_rope():
+def test_rope_gptneox():
  bs, seq_len, n_head, n_embed = 1, 6, 2, 8
  head_size = n_embed // n_head
  x = torch.randint(0, 10000, size=(bs, n_head, seq_len, head_size)).float()
@@ -22,5 +26,200 @@ def test_rope():
  torch.testing.assert_close(ours_sin_cached, theirs_sin.squeeze())
 
  ours_x_rope = apply_rope(x, ours_cos_cached, ours_sin_cached)
- theirs_x_rope, _ = apply_rotary_pos_emb(x, x, theirs_cos, theirs_sin, position_ids)
+ theirs_x_rope, _ = apply_rotary_pos_emb_gptneo(x, x, theirs_cos, theirs_sin, position_ids)
  torch.testing.assert_close(ours_x_rope, theirs_x_rope)
+
+
+@torch.inference_mode()
+def test_rope_llama_2():
+ head_dim = 64
+ rope_theta = 10_000
+
+ ##################################
+ # Compare cos and sin
+ ##################################
+ # transformer rope
+ rot_emb = LlamaRotaryEmbedding(head_dim, scaling_factor=None, base=rope_theta)
+ batch_size, seq_len = 1, 10
+ qk_tensor = torch.randn(batch_size, seq_len, head_dim)
+ position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)
+ theirs_cos, theirs_sin = rot_emb(qk_tensor, position_ids)
+
+ # our rope
+ ours_cos, ours_sin = build_rope_cache(seq_len, n_elem=head_dim, base=rope_theta)
+ torch.testing.assert_close(theirs_cos.squeeze(0), ours_cos)
+ torch.testing.assert_close(theirs_sin.squeeze(0), ours_sin)
+
+ ##################################
+ # Compare rotated tensors
+ ##################################
+ # Settings
+ num_heads = 4
+
+ # Dummy query and key tensors
+ torch.manual_seed(123)
+ queries = torch.randn(batch_size, num_heads, seq_len, head_dim)
+ keys = torch.randn(batch_size, num_heads, seq_len, head_dim)
+
+ ours_q_rot = apply_rope(queries, ours_cos, ours_sin)
+ ours_k_rot = apply_rope(keys, ours_cos, ours_sin)
+ theirs_q_rot, theirs_k_rot = apply_rotary_pos_emb_llama(queries, keys, theirs_cos, theirs_sin)
+ torch.testing.assert_close(theirs_q_rot, ours_q_rot)
+ torch.testing.assert_close(theirs_k_rot, ours_k_rot)
+
+
+# See https://huggingface.co/meta-llama/Meta-Llama-3-8B/blob/main/config.json for settings
+@torch.inference_mode()
+def test_rope_llama_3():
+ head_dim = 64
+ rope_theta = 50_000
+
+ ##################################
+ # Compare cos and sin
+ ##################################
+ # transformer rope
+ rot_emb = LlamaRotaryEmbedding(head_dim, scaling_factor=None, base=rope_theta)
+ batch_size, seq_len = 1, 10
+ qk_tensor = torch.randn(batch_size, seq_len, head_dim)
+ position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)
+ theirs_cos, theirs_sin = rot_emb(qk_tensor, position_ids)
+
+ # our rope
+ ours_cos, ours_sin = build_rope_cache(seq_len, n_elem=head_dim, base=rope_theta)
+ torch.testing.assert_close(theirs_cos.squeeze(0), ours_cos)
+ torch.testing.assert_close(theirs_sin.squeeze(0), ours_sin)
+
+ ##################################
+ # Compare rotated tensors
+ ##################################
+ # Settings
+ num_heads = 4
+
+ # Dummy query and key tensors
+ torch.manual_seed(123)
+ queries = torch.randn(batch_size, num_heads, seq_len, head_dim)
+ keys = torch.randn(batch_size, num_heads, seq_len, head_dim)
+
+ ours_q_rot = apply_rope(queries, ours_cos, ours_sin)
+ ours_k_rot = apply_rope(keys, ours_cos, ours_sin)
+ theirs_q_rot, theirs_k_rot = apply_rotary_pos_emb_llama(queries, keys, theirs_cos, theirs_sin)
+ torch.testing.assert_close(theirs_q_rot, ours_q_rot)
+ torch.testing.assert_close(theirs_k_rot, ours_k_rot)
+
+
+# See https://huggingface.co/meta-llama/Llama-3.1-8B/blob/main/config.json for settings
+@torch.inference_mode()
+def test_rope_llama_3_1():
+ head_dim = 128
+ rope_theta = 50_000
+
+ their_rope_config = {
+ "factor": 8.0,
+ "low_freq_factor": 1.0,
+ "high_freq_factor": 4.0,
+ "original_max_position_embeddings": 8192,
+ "rope_type": "llama3"
+ }
+
+ our_rope_config = {
+ "factor": 8.0,
+ "low_freq_factor": 1.0,
+ "high_freq_factor": 4.0,
+ "original_max_seq_len": 8192
+ }
+
+ config = LlamaConfig(
+ rope_theta=rope_theta,
+ rope_scaling=their_rope_config
+ )
+
+ ##################################
+ # Compare cos and sin
+ ##################################
+ # transformer rope
+ rot_emb = LlamaRotaryEmbedding(head_dim, base=rope_theta, config=config, rope_type="llama3")
+ batch_size, seq_len = 1, 10
+ qk_tensor = torch.randn(batch_size, seq_len, head_dim)
+ position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)
+ theirs_cos, theirs_sin = rot_emb(qk_tensor, position_ids)
+
+ # our rope
+ ours_cos, ours_sin = build_rope_cache(seq_len, n_elem=head_dim, base=rope_theta, extra_config=our_rope_config)
+ torch.testing.assert_close(theirs_cos.squeeze(0), ours_cos)
+ torch.testing.assert_close(theirs_sin.squeeze(0), ours_sin)
+
+ ##################################
+ # Compare rotated tensors
+ ##################################
+ # Settings
+ num_heads = 4
+
+ # Dummy query and key tensors
+ torch.manual_seed(123)
+ queries = torch.randn(batch_size, num_heads, seq_len, head_dim)
+ keys = torch.randn(batch_size, num_heads, seq_len, head_dim)
+
+ ours_q_rot = apply_rope(queries, ours_cos, ours_sin)
+ ours_k_rot = apply_rope(keys, ours_cos, ours_sin)
+ theirs_q_rot, theirs_k_rot = apply_rotary_pos_emb_llama(queries, keys, theirs_cos, theirs_sin)
+ torch.testing.assert_close(theirs_q_rot, ours_q_rot)
+ torch.testing.assert_close(theirs_k_rot, ours_k_rot)
+
+
+# See https://huggingface.co/meta-llama/Llama-3.2-3B/blob/main/config.json for settings
+@torch.inference_mode()
+def test_rope_llama_3_2():
+ head_dim = 128
+ rope_theta = 50_000
+
+ their_rope_config = {
+ "factor": 32.0,
+ "low_freq_factor": 1.0,
+ "high_freq_factor": 4.0,
+ "original_max_position_embeddings": 8192,
+ "rope_type": "llama3"
+ }
+
+ our_rope_config = {
+ "factor": 32.0,
+ "low_freq_factor": 1.0,
+ "high_freq_factor": 4.0,
+ "original_max_seq_len": 8192
+ }
+
+ config = LlamaConfig(
+ rope_theta=rope_theta,
+ rope_scaling=their_rope_config
+ )
+
+ ##################################
+ # Compare cos and sin
+ ##################################
+ # transformer rope
+ rot_emb = LlamaRotaryEmbedding(head_dim, base=rope_theta, config=config, rope_type="llama3")
+ batch_size, seq_len = 1, 10
+ qk_tensor = torch.randn(batch_size, seq_len, head_dim)
+ position_ids = torch.arange(seq_len, dtype=torch.long).unsqueeze(0)
+ theirs_cos, theirs_sin = rot_emb(qk_tensor, position_ids)
+
+ # our rope
+ ours_cos, ours_sin = build_rope_cache(seq_len, n_elem=head_dim, base=rope_theta, extra_config=our_rope_config)
+ torch.testing.assert_close(theirs_cos.squeeze(0), ours_cos)
+ torch.testing.assert_close(theirs_sin.squeeze(0), ours_sin)
+
+ ##################################
+ # Compare rotated tensors
+ ##################################
+ # Settings
+ num_heads = 4
+
+ # Dummy query and key tensors
+ torch.manual_seed(123)
+ queries = torch.randn(batch_size, num_heads, seq_len, head_dim)
+ keys = torch.randn(batch_size, num_heads, seq_len, head_dim)
+
+ ours_q_rot = apply_rope(queries, ours_cos, ours_sin)
+ ours_k_rot = apply_rope(keys, ours_cos, ours_sin)
+ theirs_q_rot, theirs_k_rot = apply_rotary_pos_emb_llama(queries, keys, theirs_cos, theirs_sin)
+ torch.testing.assert_close(theirs_q_rot, ours_q_rot)
+ torch.testing.assert_close(theirs_k_rot, ours_k_rot)