apply gradient checkpointing to Attention blocks

modeling_mpt.py

@@ -12,7 +12,7 @@ from torch.utils.checkpoint import checkpoint
 from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.utils import logging
-from .attention import attn_bias_shape, build_attn_bias, PastKeyValue
+from .attention import attn_bias_shape, build_attn_bias, PastKeyValue, MultiheadAttention, MultiQueryAttention
 from .blocks import MPTBlock, MPTBlockOutput
 from .norm import NORM_CLASS_REGISTRY
 from .configuration_mpt import MPTConfig
@@ -41,7 +41,7 @@ class MPTPreTrainedModel(PreTrainedModel):
     _no_split_modules = ['MPTBlock']
     supports_gradient_checkpointing = True
     def _set_gradient_checkpointing(self, module: nn.Module, value=False) -> None:
-        if isinstance(module, MPTModel):
+        if isinstance(module, MPTModel) or isinstance(module, MultiheadAttention) or isinstance(module, MultiQueryAttention):
             module.gradient_checkpointing = value
 
 class MPTModel(MPTPreTrainedModel):