hugging-apps/echo-memory

1import torch2 3 4def create_custom_forward(module):5    def custom_forward(*inputs, **kwargs):6        return module(*inputs, **kwargs)7    return custom_forward8 9 10def gradient_checkpoint_forward(11    model,12    use_gradient_checkpointing,13    use_gradient_checkpointing_offload,14    *args,15    **kwargs,16):17    if use_gradient_checkpointing_offload:18        with torch.autograd.graph.save_on_cpu():19            model_output = torch.utils.checkpoint.checkpoint(20                create_custom_forward(model),21                *args,22                **kwargs,23                use_reentrant=False,24            )25    elif use_gradient_checkpointing:26        model_output = torch.utils.checkpoint.checkpoint(27            create_custom_forward(model),28            *args,29            **kwargs,30            use_reentrant=False,31        )32    else:33        model_output = model(*args, **kwargs)34    return model_output35