Feed forward chunking in BertLayer class.

src/transformers/configuration_utils.py

@@ -160,7 +160,7 @@ def __init__(self, **kwargs):
         self.no_repeat_ngram_size = kwargs.pop("no_repeat_ngram_size", 0)
         self.bad_words_ids = kwargs.pop("bad_words_ids", None)
         self.num_return_sequences = kwargs.pop("num_return_sequences", 1)
-        self.chunk_size_feed_forward = kwargs.pop("chunk_size_feed_forward", 1)
+        self.chunk_size_feed_forward = kwargs.pop("chunk_size_feed_forward", 0)
 
         # Fine-tuning task arguments
         self.architectures = kwargs.pop("architectures", None)

src/transformers/modeling_bert.py

@@ -378,34 +378,17 @@ def forward(self, hidden_states, input_tensor):
         return hidden_states
 
 
-class ChunkFeedForward(nn.Module):
+class BertLayer(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.chunk_size_feed_forward = config.chunk_size_feed_forward
         self.seq_len_dim = 1
-
-        self.dense = BertIntermediate(config)
-        self.output = BertOutput(config)
-
-    def forward(self, attention_output):
-        return apply_chunking_to_forward(
-            self.chunk_size_feed_forward, self.seq_len_dim, self.forward_chunk, attention_output,
-        )
-
-    def forward_chunk(self, attention_output):
-        intermediate_output = self.dense(attention_output)
-        return self.output(intermediate_output, attention_output)
-
-
-class BertLayer(nn.Module):
-    def __init__(self, config):
-        super().__init__()
         self.attention = BertAttention(config)
         self.is_decoder = config.is_decoder
         if self.is_decoder:
             self.crossattention = BertAttention(config)
-
-        self.feed_forward = ChunkFeedForward(config)
+        self.intermediate = BertIntermediate(config)
+        self.output = BertOutput(config)
 
     def forward(
         self,
@@ -434,10 +417,17 @@ def forward(
             attention_output = cross_attention_outputs[0]
             outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights
 
-        layer_output = self.feed_forward(attention_output)
+        layer_output = apply_chunking_to_forward(
+            self.chunk_size_feed_forward, self.seq_len_dim, self.feed_forward_chunk, attention_output
+        )
         outputs = (layer_output,) + outputs
         return outputs
 
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
 
 class BertEncoder(nn.Module):
     def __init__(self, config):