bugfix: Update modeling_t5.T5Stack.forward() for Gradient Checkpointing

Update checkpoint() call such that parameters for the layer_module object are passed correctly.

modeling_t5.py

@@ -1204,14 +1204,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 +1215,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 +1239,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: