Optimize MultiHeadAttention forward method (#3550)

facebookresearch · Mar 26, 2021 · bbd6e26 · bbd6e26
1 parent 68850f8
commit bbd6e26
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/parlai/agents/transformer/modules/attention.py b/parlai/agents/transformer/modules/attention.py
@@ -178,8 +178,6 @@ def prepare_head(tensor):
         _, _key_len, dim = key.size()
 
         q = prepare_head(self.q_lin(query))
-        k = prepare_head(self.k_lin(key))
-        v = prepare_head(self.v_lin(value))
 
         # Prepend incremental states. For each of the key, value, and mask, see if
         # a previous incremental state exists, and if so, reshape it to match the shape
@@ -195,15 +193,19 @@ def prepare_head(tensor):
             if static_kv:
                 k = prev_key
             else:
-                k = torch.cat([prev_key, k], dim=1)
+                k = torch.cat([prev_key, prepare_head(self.k_lin(key))], dim=1)
+        else:
+            k = prepare_head(self.k_lin(key))
         if 'prev_value' in incr_state:
             prev_value = incr_state['prev_value'].view(
                 batch_size * n_heads, -1, dim_per_head
             )
             if static_kv:
                 v = prev_value
             else:
-                v = torch.cat([prev_value, v], dim=1)
+                v = torch.cat([prev_value, prepare_head(self.v_lin(value))], dim=1)
+        else:
+            v = prepare_head(self.v_lin(value))
         if 'prev_mask' in incr_state:
             if static_kv:
                 mask = incr_state['prev_mask']