[transformer] add qk norm

wenet/transformer/attention.py

@@ -608,9 +608,17 @@ def __init__(self,
                  use_sdpa: bool = False,
                  n_kv_head: Optional[int] = None,
                  head_dim: Optional[int] = None,
+                 qk_norm: bool = False,
                  style='google'):
+
         super().__init__(n_head, n_feat, dropout_rate, query_bias, key_bias,
                          value_bias, use_sdpa, n_kv_head, head_dim)
+        self.qk_norm = qk_norm
+        # https://arxiv.org/pdf/2302.05442
+        if self.qk_norm:
+            from wenet.utils.class_utils import WENET_NORM_CLASSES
+            self.q_norm = WENET_NORM_CLASSES['rms_norm'](self.d_k, eps=1e-6)
+            self.k_norm = WENET_NORM_CLASSES['rms_norm'](self.d_k, eps=1e-6)
         self.style = style
 
     def forward(
@@ -655,11 +663,15 @@ def forward(
         q = self._forward_linearx('query', query, head_first=False)
         k = self._forward_linearx('key', key, head_first=False)
         v = self._forward_linearx('value', value, head_first=False)
+
         # NOTE(Mddct): In order to make the code easier to read,
         #    these two lines are not placed in MultiHeadedAttention.
         q = WENET_APPLY_ROTARY_EMB[self.style](q, pos_emb)
         k = WENET_APPLY_ROTARY_EMB[self.style](k, pos_emb)
 
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
         k, v, new_cache = self._update_kv_and_cache(k,
                                                     v,
                                                     cache,