[BUG]: activation checkpoint function can't pass kwargs

[Gy-Lu]

🐛 Describe the bug

See the checkpoint function and how it used in CheckpointModule. Now the only keyword arg can be passed in checkpoint function is use_reentrant, and can't be passed into _forward.

the def of checkpoint function

def checkpoint(function, activation_offload, *args, use_reentrant: bool = True):

its usage in CheckpointModule

    def forward(self, *args, **kwargs):
        if self._use_checkpoint:
            return checkpoint(self._forward, self._offload, *args, **kwargs)
        else:
            return self._forward(*args, **kwargs)

Environment

No response

[Gy-Lu]

Hmm, It seems that keyword args in _forward of CheckpointModule is not supported.

[super-dainiu]

cc @Cypher30

[binmakeswell]

We have updated a lot. This issue was closed due to inactivity. Thanks.