Add QK Normalization (#1100)

* add qk normalization * Update NeoXArgs docs automatically * Update NeoXArgs docs automatically --------- Co-authored-by: github-actions <[email protected]> Co-authored-by: Quentin Anthony <[email protected]>
EleutherAI · Dec 22, 2023 · f161245 · f161245
1 parent 9283eff
commit f161245
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 1 deletion.
diff --git a/configs/neox_arguments.md b/configs/neox_arguments.md
@@ -111,7 +111,7 @@ Logging Arguments
 
 - **git_hash**: str
 
- Default = a279fc8
+ Default = 1fc0521
 
  current git hash of repository
 
@@ -261,6 +261,14 @@ Model Arguments
 
 
 
+- **use_qk_layernorm**: bool
+
+ Default = False
+
+ Use QK Normalization
+
+
+
 - **layernorm_epsilon**: float
 
  Default = 1e-05

diff --git a/megatron/model/transformer.py b/megatron/model/transformer.py
@@ -284,6 +284,16 @@ def __init__(
  neox_args.num_attention_heads, world_size
  )
  self.pos_emb = neox_args.pos_emb
+ self.use_qk_layernorm = neox_args.use_qk_layernorm
+ if self.use_qk_layernorm:
+ norm, eps = get_norm(neox_args)
+ self.qk_layernorm = norm(
+ [
+ self.num_attention_heads_per_partition,
+ self.hidden_size_per_attention_head,
+ ],
+ eps=eps,
+ )
 
  # Strided linear layer.
  self.query_key_value = mpu.ColumnParallelLinear(
@@ -639,6 +649,11 @@ def forward(self, hidden_states, attention_mask, layer_past=None):
  mixed_x_layer, 3
  )
 
+ # QK Normalization https://arxiv.org/abs/2302.05442
+ if self.use_qk_layernorm:
+ query_layer = self.qk_layernorm(query_layer)
+ key_layer = self.qk_layernorm(key_layer)
+
  if exists(self.rotary_emb):
  if exists(self.rotary_ndims):
  # partial rotary

diff --git a/megatron/neox_arguments/neox_args.py b/megatron/neox_arguments/neox_args.py
@@ -125,6 +125,11 @@ class NeoXArgsModel(NeoXArgsTemplate):
  Normalization layer to use. Choose from "layernorm", "rmsnorm", "scalenorm".
  """
 
+ use_qk_layernorm: bool = False
+ """
+ Use QK Normalization
+ """
+
  layernorm_epsilon: float = 1.0e-5
  """
  Layer norm epsilon.