added query-key norm to accomodate OLMo2 (#1894)

Co-authored-by: Andrei-Aksionov <aksionau.andrei@gmail.com>
Co-authored-by: Andrei-Aksionov <58434077+Andrei-Aksionov@users.noreply.github.com>

Author: ysjprojects

litgpt/config.py

@@ -25,6 +25,7 @@ class Config:
     # Transformer block (structure, normalizations)
     norm_class_name: Literal["LayerNorm", "RMSNorm"] = "LayerNorm"
     norm_eps: float = 1e-5
+    norm_qk: bool = False
     post_attention_norm: bool = False
     post_mlp_norm: bool = False
     parallel_residual: bool = True

litgpt/model.py

@@ -341,6 +341,13 @@ def __init__(self, config: Config, block_idx: int) -> None:
             config.sliding_window_size is not None and
             block_idx % config.sliding_window_layer_stride == 0
         )
+
+        if config.norm_qk:
+            self.norm_q = config.norm_class(config.head_size * config.n_head, eps=config.norm_eps)
+            self.norm_k = config.norm_class(config.head_size * config.n_query_groups, eps=config.norm_eps)
+        else:
+            self.norm_q = self.norm_k = None
+
         self.config = config
         self.block_idx = block_idx
 
@@ -377,6 +384,10 @@ def forward(
         # Split qkv into query, key and value matrices.
         q, k, v = qkv.split((query_size, key_size, value_size), dim=-1)  # 3x(B, T, C*)
 
+        if self.config.norm_qk:
+            q = self.norm_q(q)
+            k = self.norm_k(k)
+
         # To place the num_heads (nh) dimension right after the batch (B) dimension, the first step is to decouple the
         # embedding size (C) into num_heads (nh) and head_size (hs).
         q = q.view(B, T, n_head, head_size)  # (B, T, nh_q, hs)

tests/test_readme.py

@@ -32,6 +32,7 @@ def run_command(command):
 
 
 @pytest.mark.dependency()
+@pytest.mark.flaky(reruns=5, reruns_delay=2)
 def test_download_model():
     repo_id = str(REPO_ID).replace("\\", "/")  # fix for Windows CI
     command = ["litgpt", "download", str(repo_id)]
@@ -48,6 +49,7 @@ def test_download_model():
 
 
 @pytest.mark.dependency()
+@pytest.mark.flaky(reruns=5, reruns_delay=2)
 def test_download_books():
     CUSTOM_TEXTS_DIR.mkdir(parents=True, exist_ok=True)