Make it works on GPU too.

src/transformers/modeling_bert.py

@@ -19,6 +19,7 @@
 import logging
 import math
 import os
+from functools import partial
 
 import torch
 from torch import nn
@@ -153,7 +154,7 @@ class BertEmbeddings(nn.Module):
         # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
         # any TensorFlow checkpoint file
         self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.dropout = partial(nn.functional.dropout, p=config.hidden_dropout_prob)
 
     def forward(self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None):
         if input_ids is not None:
@@ -176,7 +177,7 @@ class BertEmbeddings(nn.Module):
 
         embeddings = inputs_embeds + position_embeddings + token_type_embeddings
         embeddings = self.LayerNorm(embeddings)
-        embeddings = self.dropout(embeddings)
+        embeddings = self.dropout(embeddings, training=self.training)
         return embeddings
 
 
@@ -193,17 +194,26 @@ class BertSelfAttention(nn.Module):
         self.num_attention_heads = config.num_attention_heads
         self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
         self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.scale_factor = 1. / math.sqrt(self.attention_head_size)
 
         self.query = nn.Linear(config.hidden_size, self.all_head_size)
         self.key = nn.Linear(config.hidden_size, self.all_head_size)
         self.value = nn.Linear(config.hidden_size, self.all_head_size)
 
-        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+        self.qkv_weight = nn.Parameter(torch.cat(
+            (self.query.weight, self.key.weight, self.value.weight), dim=0
+        ).t().contiguous())
 
-    def transpose_for_scores(self, x):
+        self.qkv_bias = nn.Parameter(
-        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+            torch.cat((self.query.bias, self.key.bias, self.value.bias), dim=0)
-        x = x.view(*new_x_shape)
+        )
-        return x.permute(0, 2, 1, 3)
+
+        self.dropout = partial(nn.functional.dropout, p=config.attention_probs_dropout_prob)
+
+    # def transpose_for_scores(self, x):
+    #     new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+    #     x = x.view(*new_x_shape)
+    #     return x.permute(0, 2, 1, 3)
 
     def forward(
         self,
@@ -213,49 +223,45 @@ class BertSelfAttention(nn.Module):
         encoder_hidden_states=None,
         encoder_attention_mask=None,
     ):
-        mixed_query_layer = self.query(hidden_states)
-
         # If this is instantiated as a cross-attention module, the keys
         # and values come from an encoder; the attention mask needs to be
         # such that the encoder's padding tokens are not attended to.
         if encoder_hidden_states is not None:
+            mixed_query_layer = self.query(hidden_states)
             mixed_key_layer = self.key(encoder_hidden_states)
             mixed_value_layer = self.value(encoder_hidden_states)
             attention_mask = encoder_attention_mask
         else:
-            mixed_key_layer = self.key(hidden_states)
+            # Compute qkv_bias + (hidden_states @ qkv_weight)
-            mixed_value_layer = self.value(hidden_states)
+            qkv = torch.matmul(hidden_states, self.qkv_weight) + self.qkv_bias
+            mixed_query_layer, mixed_key_layer, mixed_value_layer = qkv.chunk(3, dim=2)
 
-        query_layer = self.transpose_for_scores(mixed_query_layer)
+        attention_scores = torch.bmm(mixed_query_layer, mixed_key_layer.transpose(2, 1))
-        key_layer = self.transpose_for_scores(mixed_key_layer)
+        attention_scores = attention_scores * self.scale_factor
-        value_layer = self.transpose_for_scores(mixed_value_layer)
 
-        # Take the dot product between "query" and "key" to get the raw attention scores.
-        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
-        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
         if attention_mask is not None:
             # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
-            attention_scores = attention_scores + attention_mask
+            attention_scores = attention_scores + attention_mask.squeeze(2)
 
         # Normalize the attention scores to probabilities.
-        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+        attention_probs = torch.softmax(attention_scores, dim=-1)
 
         # This is actually dropping out entire tokens to attend to, which might
         # seem a bit unusual, but is taken from the original Transformer paper.
-        attention_probs = self.dropout(attention_probs)
+        attention_probs = self.dropout(attention_probs, training=self.training)
 
         # Mask heads if we want to
         if head_mask is not None:
+            bs = attention_probs.size(0)
+
+            # Split attention_probs
+            attention_probs = attention_probs.view(bs, -1, self.num_attention_heads, self.attention_head_size)
             attention_probs = attention_probs * head_mask
+            attention_probs = attention_probs.flatten(2)
 
-        context_layer = torch.matmul(attention_probs, value_layer)
+        context_layer = torch.bmm(attention_probs, mixed_value_layer)
 
-        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        return (context_layer, attention_probs) if self.output_attentions else (context_layer,)
-        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
-        context_layer = context_layer.view(*new_context_layer_shape)
-
-        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
-        return outputs
 
 
 class BertSelfOutput(nn.Module):
@@ -263,11 +269,11 @@ class BertSelfOutput(nn.Module):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
         self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.dropout = partial(nn.functional.dropout, p=config.hidden_dropout_prob)
 
     def forward(self, hidden_states, input_tensor):
         hidden_states = self.dense(hidden_states)
-        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.dropout(hidden_states, training=self.training)
         hidden_states = self.LayerNorm(hidden_states + input_tensor)
         return hidden_states
 
@@ -314,8 +320,7 @@ class BertAttention(nn.Module):
             hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask
         )
         attention_output = self.output(self_outputs[0], hidden_states)
-        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return (attention_output,) + self_outputs[1:]  # add attentions if we output them
-        return outputs
 
 
 class BertIntermediate(nn.Module):
@@ -338,11 +343,11 @@ class BertOutput(nn.Module):
         super().__init__()
         self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
         self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
-        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.dropout = partial(nn.functional.dropout, p=config.hidden_dropout_prob)
 
     def forward(self, hidden_states, input_tensor):
         hidden_states = self.dense(hidden_states)
-        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.dropout(hidden_states, training=self.training)
         hidden_states = self.LayerNorm(hidden_states + input_tensor)
         return hidden_states
 
@@ -427,15 +432,13 @@ class BertPooler(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        self.activation = nn.Tanh()
 
     def forward(self, hidden_states):
         # We "pool" the model by simply taking the hidden state corresponding
         # to the first token.
         first_token_tensor = hidden_states[:, 0]
         pooled_output = self.dense(first_token_tensor)
-        pooled_output = self.activation(pooled_output)
+        return torch.tanh(pooled_output)
-        return pooled_output
 
 
 class BertPredictionHeadTransform(nn.Module):