rollback

modeling_mpt.py

@@ -30,18 +30,11 @@ class MPTPreTrainedModel(PreTrainedModel):
     base_model_prefix = 'model'
     _no_split_modules = ['MPTBlock']
 
-    supports_gradient_checkpointing = True
-
-    def _set_gradient_checkpointing(self, module, value=False):
-        if isinstance(module, MPTModel):
-            module.gradient_checkpointing = value
-
 class MPTModel(MPTPreTrainedModel):
 
     def __init__(self, config: MPTConfig):
         config._validate_config()
         super().__init__(config)
-        self.gradient_checkpointing = False
         self.attn_impl = config.attn_config['attn_impl']
         self.prefix_lm = config.attn_config['prefix_lm']
         self.attn_uses_sequence_id = config.attn_config['attn_uses_sequence_id']
@@ -151,9 +144,6 @@ class MPTModel(MPTPreTrainedModel):
     def forward(self, input_ids: torch.LongTensor, past_key_values: Optional[List[Tuple[torch.FloatTensor]]]=None, attention_mask: Optional[torch.ByteTensor]=None, prefix_mask: Optional[torch.ByteTensor]=None, sequence_id: Optional[torch.LongTensor]=None, return_dict: Optional[bool]=None, output_attentions: Optional[bool]=None, output_hidden_states: Optional[bool]=None, use_cache: Optional[bool]=None, inputs_embeds: Optional[torch.FloatTensor] = None):
         return_dict = return_dict if return_dict is not None else self.config.return_dict
         use_cache = use_cache if use_cache is not None else self.config.use_cache
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                use_cache = False
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
         elif input_ids is not None:
@@ -228,23 +218,7 @@ class MPTModel(MPTPreTrainedModel):
                 assert all_hidden_states is not None
                 all_hidden_states = all_hidden_states + (x,)
             past_key_value = past_key_values[b_idx] if past_key_values is not None else None
-            if self.gradient_checkpointing and self.training:
-                def create_custom_forward(module):
-                    def custom_forward(*inputs):
-                        # None for past_key_value
-                        return module(*inputs)
-                    return custom_forward
-                (x, past_key_value) = torch.utils.checkpoint.checkpoint(
-                    create_custom_forward(block),
-                    x,
-                    past_key_value,
-                    attn_bias,
-                    attention_mask,
-                    self.is_causal,
-                )
-            else:
-                (x, _, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
-            # (x, attn_weights, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
+            (x, attn_weights, past_key_value) = block(x, past_key_value=past_key_value, attn_bias=attn_bias, attention_mask=attention_mask, is_causal=self.is_causal)
             if past_key_values is not None:
                 past_key_values[b_idx] = past_key_value
             if output_attentions: