Adapt modelling for gradient checkpointing

modeling_t5.py

@@ -977,10 +977,6 @@ class T5PreTrainedModel(PreTrainedModel):
             if module.has_relative_attention_bias:
                 module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5))
 
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, (T5Attention, T5Stack)):
-            module.gradient_checkpointing = value
-
     def _shift_right(self, input_ids):
         decoder_start_token_id = self.config.decoder_start_token_id
         pad_token_id = self.config.pad_token_id
@@ -1204,14 +1200,8 @@ class T5Stack(T5PreTrainedModel):
 
             if self.gradient_checkpointing and self.training:
 
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        return tuple(module(*inputs, use_cache, output_attentions))
-
-                    return custom_forward
-
                 layer_outputs = checkpoint(
-                    create_custom_forward(layer_module),
+                    layer_module,
                     hidden_states,
                     extended_attention_mask,
                     position_bias,
@@ -1221,10 +1211,15 @@ class T5Stack(T5PreTrainedModel):
                     layer_head_mask,
                     cross_attn_layer_head_mask,
                     None,  # past_key_value is always None with gradient checkpointing
-                    relative_position=relative_position,
-                    sparsity_mask=sparsity_mask,
-                    use_additional_bucket=use_additional_bucket,
+                    use_cache,
+                    output_attentions,
+                    True, # return_dict is true at training time
+                    relative_position,
+                    sparsity_mask,
+                    use_additional_bucket,
+                    use_reentrant=False
                 )
+
             else:
                 layer_outputs = layer_module(
                     hidden_states,
@@ -1240,7 +1235,7 @@ class T5Stack(T5PreTrainedModel):
                     output_attentions=output_attentions,
                     relative_position=relative_position,
                     sparsity_mask=sparsity_mask,
-                    use_additional_bucket=use_additional_bucket,
+                    use_additional_bucket=use_additional_bucket
                 )
 
             # layer_outputs is a tuple with: