added query-key norm to accomodate OLMo2

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

@@ -287,6 +287,12 @@ def __init__(self, config: Config, block_idx: int) -> None:
             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
 
     def forward(
@@ -324,6 +330,15 @@ def forward(
         k = k.view(B, T, self.config.n_query_groups, self.config.head_size)  # (B, T, nh_k, hs)
         v = v.view(B, T, self.config.n_query_groups, self.config.head_size)  # (B, T, nh_v, hs)
 
+        if self.config.norm_qk:
+            q = q.reshape(B, T, -1)  # (B, T, nh_q * hs)
+            q = self.norm_q(q)
+            q = q.view(B, T, self.config.n_head, self.config.head_size)
+
+            k = k.reshape(B, T, -1)  # (B, T, nh_k * hs)
+            k = self.norm_k(k)
+            k = k.view(B, T, self.config.n_query_groups, self.config.head_size)
+
         # The tensors `query`, `key`, and `value` are now accurately structured: within each batch element (B), there are
         # multiple heads (nh), and within each head, there is a sequence of elements (T), each represented by a vector
         # of size `hs`.