From 43f5cc617b50642965c06197b16641ea7b2a7ec8 Mon Sep 17 00:00:00 2001 From: rasbt Date: Mon, 21 Oct 2024 19:51:03 -0500 Subject: [PATCH] RoPE increase --- .../converting-llama2-to-llama3.ipynb | 18 +++++++++--------- ch05/07_gpt_to_llama/standalone-llama32.ipynb | 6 +++--- ch05/07_gpt_to_llama/tests/tests.py | 6 +++--- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/ch05/07_gpt_to_llama/converting-llama2-to-llama3.ipynb b/ch05/07_gpt_to_llama/converting-llama2-to-llama3.ipynb index 1c5fbf9c..4b4459fc 100644 --- a/ch05/07_gpt_to_llama/converting-llama2-to-llama3.ipynb +++ b/ch05/07_gpt_to_llama/converting-llama2-to-llama3.ipynb @@ -254,12 +254,12 @@ "- Llama 3 uses rotary position embeddings (RoPE) similar to Llama 2 (for a detailed explanation, please see the [RoPE paper](https://arxiv.org/abs/2104.09864))\n", "- There are some subtle differences in the RoPE settings, though\n", " - Llama 3 now supports up to 8,192 tokens, twice as many as Llama 2 (4,096)\n", - " - The base value for the so-called RoPE $\\theta$ (see equation below) was increased from 10,000 (Llama 2) to 50,000 (Llama 3) in the following equation (adapted from the [RoPE paper](https://arxiv.org/abs/2104.09864))\n", + " - The base value for the so-called RoPE $\\theta$ (see equation below) was increased from 10,000 (Llama 2) to 500,000 (Llama 3) in the following equation (adapted from the [RoPE paper](https://arxiv.org/abs/2104.09864))\n", "\n", "$$\\Theta = \\left\\{\\theta_i = \\text{base}^{\\frac{-2(i-1)}{d}}, i \\in \\left[1, 2, ..., d/2\\right]\\right\\}$$\n", "\n", "- These $\\theta$ values are a set of predefined parameters that are used to determine the rotational angles in the rotary matrix, where $d$ is the dimensionality of the embedding space\n", - "- Increasing the base from 10,000 to 50,000 makes the frequencies (or rotation angles) decay more slowly across the dimensions, which means that higher dimensions will be associated with larger angles than before (essentially, it's a decompression of the frequencies)\n", + "- Increasing the base from 10,000 to 500,000 makes the frequencies (or rotation angles) decay more slowly across the dimensions, which means that higher dimensions will be associated with larger angles than before (essentially, it's a decompression of the frequencies)\n", "- In addition, we introduce a `freq_config` section in the code below that adjusts the frequency; however, we won't be needing it in Llama 3 (only Llama 3.1 and Llama 3.2), so we will revisit this `freq_config` later (it's set to `None` and ignored by default)" ] }, @@ -274,7 +274,7 @@ "source": [ "import torch\n", "\n", - "def precompute_rope_params(head_dim, theta_base=10000, context_length=4096, freq_config=None):\n", + "def precompute_rope_params(head_dim, theta_base=10_000, context_length=4096, freq_config=None):\n", " assert head_dim % 2 == 0, \"Embedding dimension must be even\"\n", "\n", " # Compute the inverse frequencies\n", @@ -347,7 +347,7 @@ "llama_3_context_len = 8192\n", "\n", "llama_2_theta_base = 10_000\n", - "llama_3_theta_base = 50_000" + "llama_3_theta_base = 500_000" ] }, { @@ -907,7 +907,7 @@ " \"n_layers\": 32, # Number of layers\n", " \"hidden_dim\": 14_336, # NEW: Larger size of the intermediate dimension in FeedForward\n", " \"n_kv_groups\": 8, # NEW: Key-Value groups for grouped-query attention\n", - " \"rope_base\": 50_000, # NEW: The base in RoPE's \"theta\" was increased to 50_000\n", + " \"rope_base\": 500_000, # NEW: The base in RoPE's \"theta\" was increased to 500_000\n", " \"rope_freq\": None, # NEW: Additional configuration for adjusting the RoPE frequencies\n", " \"dtype\": torch.bfloat16 # Lower-precision dtype to save memory\n", "}" @@ -2060,7 +2060,7 @@ " \"n_layers\": 32, # Number of layers\n", " \"hidden_dim\": 14_336, # Size of the intermediate dimension in FeedForward\n", " \"n_kv_groups\": 8, # Key-Value groups for grouped-query attention\n", - " \"rope_base\": 50_000, # The base in RoPE's \"theta\"\n", + " \"rope_base\": 500_000, # The base in RoPE's \"theta\"\n", " \"rope_freq\": None, # Additional configuration for adjusting the RoPE frequencies\n", " \"dtype\": torch.bfloat16 # Lower-precision dtype to save memory\n", "}\n", @@ -2073,7 +2073,7 @@ " \"n_layers\": 32, # Number of layers\n", " \"hidden_dim\": 14_336, # Size of the intermediate dimension in FeedForward\n", " \"n_kv_groups\": 8, # Key-Value groups for grouped-query attention\n", - " \"rope_base\": 50_000, # The base in RoPE's \"theta\"\n", + " \"rope_base\": 500_000, # The base in RoPE's \"theta\"\n", " \"dtype\": torch.bfloat16, # Lower-precision dtype to save memory\n", " \"rope_freq\": { # NEW: RoPE frequency scaling\n", " \"factor\": 8.0,\n", @@ -2421,7 +2421,7 @@ " \"n_layers\": 32, # Number of layers\n", " \"hidden_dim\": 14_336, # Size of the intermediate dimension in FeedForward\n", " \"n_kv_groups\": 8, # Key-Value groups for grouped-query attention\n", - " \"rope_base\": 50_000, # The base in RoPE's \"theta\"\n", + " \"rope_base\": 500_000, # The base in RoPE's \"theta\"\n", " \"dtype\": torch.bfloat16, # Lower-precision dtype to save memory\n", " \"rope_freq\": { # NEW: RoPE frequency scaling\n", " \"factor\": 8.0,\n", @@ -2440,7 +2440,7 @@ " \"n_layers\": 16, # NEW: Half the number of layers\n", " \"hidden_dim\": 8192, # NEW: Almost half the size of the intermediate dimension in FeedForward\n", " \"n_kv_groups\": 8, # Key-Value groups for grouped-query attention\n", - " \"rope_base\": 50_000, # The base in RoPE's \"theta\"\n", + " \"rope_base\": 500_000, # The base in RoPE's \"theta\"\n", " \"dtype\": torch.bfloat16, # Lower-precision dtype to save memory\n", " \"rope_freq\": { # RoPE frequency scaling\n", " \"factor\": 32.0, # NEW: Adjustment of the rescaling factor\n", diff --git a/ch05/07_gpt_to_llama/standalone-llama32.ipynb b/ch05/07_gpt_to_llama/standalone-llama32.ipynb index a9398a25..4201f959 100644 --- a/ch05/07_gpt_to_llama/standalone-llama32.ipynb +++ b/ch05/07_gpt_to_llama/standalone-llama32.ipynb @@ -129,7 +129,7 @@ "metadata": {}, "outputs": [], "source": [ - "def precompute_rope_params(head_dim, theta_base=10000, context_length=4096, freq_config=None):\n", + "def precompute_rope_params(head_dim, theta_base=10_000, context_length=4096, freq_config=None):\n", " assert head_dim % 2 == 0, \"Embedding dimension must be even\"\n", "\n", " # Compute the inverse frequencies\n", @@ -407,7 +407,7 @@ " \"n_layers\": 16, # Number of layers\n", " \"hidden_dim\": 8192, # Size of the intermediate dimension in FeedForward\n", " \"n_kv_groups\": 8, # Key-Value groups for grouped-query attention\n", - " \"rope_base\": 50_000, # The base in RoPE's \"theta\"\n", + " \"rope_base\": 500_000, # The base in RoPE's \"theta\"\n", " \"dtype\": torch.bfloat16, # Lower-precision dtype to save memory\n", " \"rope_freq\": { # RoPE frequency scaling\n", " \"factor\": 32.0,\n", @@ -427,7 +427,7 @@ "# \"n_layers\": 28, # Number of layers\n", "# \"hidden_dim\": 8192, # Size of the intermediate dimension in FeedForward\n", "# \"n_kv_groups\": 8, # Key-Value groups for grouped-query attention\n", - "# \"rope_base\": 50_000, # The base in RoPE's \"theta\"\n", + "# \"rope_base\": 500_000, # The base in RoPE's \"theta\"\n", "# \"dtype\": torch.bfloat16, # Lower-precision dtype to save memory\n", "# \"rope_freq\": { # RoPE frequency scaling\n", "# \"factor\": 32.0,\n", diff --git a/ch05/07_gpt_to_llama/tests/tests.py b/ch05/07_gpt_to_llama/tests/tests.py index d4725861..e753ae74 100644 --- a/ch05/07_gpt_to_llama/tests/tests.py +++ b/ch05/07_gpt_to_llama/tests/tests.py @@ -111,7 +111,7 @@ def test_rope_llama3(notebook): context_len = 8192 num_heads = 4 head_dim = 16 - theta_base = 50_000 + theta_base = 500_000 # Instantiate RoPE parameters cos, sin = nb2.precompute_rope_params( @@ -155,7 +155,7 @@ def test_rope_llama3_12(notebook): context_len = 8192 num_heads = 4 head_dim = 16 - rope_theta = 50_000 + rope_theta = 500_000 rope_config = { "factor": 8.0, @@ -194,7 +194,7 @@ class RoPEConfig: rope_scaling = hf_rope_params factor = 1.0 dim: int = head_dim - rope_theta = 50_000 + rope_theta = 500_000 max_position_embeddings: int = 8192 hidden_size = head_dim * num_heads num_attention_heads = num_heads