add act checkpoint at sub layer level

llmfoundry/models/mpt/modeling_mpt.py

@@ -45,7 +45,9 @@
 from transformers.models.llama.modeling_llama import \
     LlamaRotaryEmbedding as HFRotaryEmbedding
 
-from llmfoundry.models.layers.attention import attn_bias_shape, build_attn_bias
+from llmfoundry.models.layers.attention import (ATTN_CLASS_REGISTRY,
+                                                attn_bias_shape,
+                                                build_attn_bias)
 from llmfoundry.models.layers.blocks import MPTBlock
 from llmfoundry.models.layers.custom_embedding import SharedEmbedding
 from llmfoundry.models.layers.fc import FC_CLASS_REGISTRY as FC_CLASS_REGISTRY
@@ -705,7 +707,34 @@ def fsdp_wrap_fn(self, module: nn.Module) -> bool:
 
     # Activation Checkpointing
     def activation_checkpointing_fn(self, module: nn.Module) -> bool:
-        return isinstance(module, MPTBlock)
+        if not hasattr(self.config, 'activation_checkpointing_target'
+                      ) or self.config.activation_checkpointing_target is None:
+            log.info(
+                f'activation checkpointing MPTBlock as activation_checkpointing_target is not set in model_config'
+            )
+            return isinstance(module, MPTBlock)
+        else:
+            act_ckpt_list = self.config.activation_checkpointing_target
+            if 'MPTBlock' in act_ckpt_list:
+                act_ckpt_list = ['MPTBlock']
+                warnings.warn(
+                    f'activation checkpointing MPTBlock, ignoring other sub-block modules if specified'
+                )
+            mod_types = ()
+            for mod_name in act_ckpt_list:
+                if mod_name.lower() == 'mptblock':
+                    mod_types += (MPTBlock,)
+                elif mod_name in ATTN_CLASS_REGISTRY:
+                    mod_types += (ATTN_CLASS_REGISTRY[mod_name],)
+                elif mod_name in FFN_CLASS_REGISTRY:
+                    mod_types += (FFN_CLASS_REGISTRY[mod_name],)
+                elif mod_name in NORM_CLASS_REGISTRY:
+                    mod_types += (NORM_CLASS_REGISTRY[mod_name],)
+                else:
+                    warnings.warn(
+                        f'module name specified in activation_checkpointing_target ({mod_name}) not recognized, available options are names in ATTN_CLASS_REGISTRY, FFN_CLASS_REGISTRY, NORM_CLASS_REGISTRY, or MPTBlock.'
+                    )
+            return isinstance(module, mod_types)
 
     def prepare_inputs_for_generation(
         self,