Default to only applying QK norm over the head dimension

configs/neox_arguments.md

@@ -490,6 +490,14 @@ Model Arguments
 
 
 
+- **qk_layernorm_over_heads**: bool
+
+    Default = False
+
+    Apply QK normalization over [\*, N, H] instead of [\*, H].
+
+
+
 - **layernorm_epsilon**: float
 
     Default = 1e-05

megatron/model/transformer.py

@@ -309,15 +309,18 @@ def __init__(
         self.pos_emb = neox_args.pos_emb
 
         self.use_qk_layernorm = neox_args.use_qk_layernorm
+        self.qk_layernorm_over_heads = neox_args.qk_layernorm_over_heads
         if self.use_qk_layernorm:
             norm, eps = get_norm(neox_args)
-            self.qk_layernorm = norm(
+            norm_dims = (
                 [
                     self.num_attention_heads_per_partition,
                     self.hidden_size_per_attention_head,
-                ],
-                eps=eps,
+                ]
+                if self.qk_layernorm_over_heads
+                else [self.hidden_size_per_attention_head]
             )
+            self.qk_layernorm = norm(norm_dims, eps=eps)
 
         self.sliding_window_width = neox_args.sliding_window_width
 

megatron/neox_arguments/neox_args.py

@@ -189,6 +189,11 @@ class NeoXArgsModel(NeoXArgsTemplate):
     Use QK Normalization
     """
 
+    qk_layernorm_over_heads: bool = False
+    """
+    Apply QK normalization over [*, N, H] instead of [*, H].
+    """
+
     layernorm_epsilon: float = 1.0e-5
     """
     Layer norm epsilon.