def _checkpointed_attention_forward(
        self, query, key, value, attention_mask, rotary_pos_emb=None
    ):
        """Forward method with selective activation checkpointing."""

        def custom_forward(*inputs):
            query = inputs[0]
            key = inputs[1]
            value = inputs[2]
            attention_mask = inputs[3]
            output_ = self.dot_product_attention(query, key, value, attention_mask)
            return output_

        hidden_states = tensor_parallel.checkpoint(
            custom_forward, False, query, key, value, attention_mask, rotary_pos_emb
        )

        return hidden_states
Forward method with selective activation checkpointing.
training/DeepSpeed-Domino/megatron/core/transformer/attention.py