gradient_checkpointing以计算时间换取内存

gradient_checkpointing的方法能够以计算时间去换取内存，这里我阅读longformer的源代码的时候，读取到了以下的内容，可以作为以计算时间换取内存的经典代码方式

if self.gradient_checkpointing and self.training:
    #gradient_checkpoint:以计算时间换内存的方式，显著减小模型训练对gpu的占用
    def create_custom_forward(module):
        def custom_forward(*inputs):
            return module(*inputs, is_global_attn, output_attentions)

        return custom_forward

    layer_outputs = torch.utils.checkpoint.checkpoint(
        create_custom_forward(layer_module),
        hidden_states,
        attention_mask,
        head_mask[idx] if head_mask is not None else None,
        is_index_masked,
        is_index_global_attn,
    )
else:
    layer_outputs = layer_module(
        hidden_states,
        attention_mask=attention_mask,
        layer_head_mask=head_mask[idx] if head_mask is not None else None,
        is_index_masked=is_index_masked,
        is_index_global_attn=is_index_global_attn,
        is_global_attn=is_global_attn,
        output_attentions=output_attentions,
    )

两者调用的方式本质上是一样的，只不过使用self.gradient_checkpointing的时候，调用了torch.utils.checkpoint模块，能够用时间换取显存的调用