Add gradient checkpointing

modeling_mpt.py

@@ -30,11 +30,18 @@ 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']
@@ -161,6 +168,10 @@ class MPTModel(MPTPreTrainedModel):
                 raise ValueError('sequence_id is a required argument when MPT is configured with attn_uses_sequence_id=True ' + 'and the model is in train mode.')
             elif self.attn_uses_sequence_id is False and sequence_id is not None:
                 warnings.warn('MPT received non-None input for `sequence_id` but is configured with attn_uses_sequence_id=False. ' + 'This input will be ignored. If you want the model to use `sequence_id`, set attn_uses_sequence_id to True.')
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                use_cache = False
+
         S = input_ids.size(1)
         assert S <= self.config.max_seq_len, f'Cannot forward input with seq_len={S}, this model only supports seq_len<={self.config.max_seq_len}'
         tok_emb = self.wte(input_ids)
@@ -197,7 +208,27 @@ 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
-            (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 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, attn_weights, past_key_value) = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    x,
+                    past_key_value,
+                    attn_bias,
+                    attention_mask,
+                    self.is_causal,
+                )
+            else:
+                (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: