✨ gradient checkpointing

modeling_mpt.py

@@ -33,7 +33,10 @@ class MPTPreTrainedModel(PreTrainedModel):
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _no_split_modules = ["MPTBlock"]
-
+    
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, MPTModel):
+            module.gradient_checkpointing = value
 
 class MPTModel(MPTPreTrainedModel):
     def __init__(self, config: MPTConfig):