Add kwarg support for no_reentrant checkpoint (#80987)
Supports kwargs input to function when `torch.utils.checkpoint` with use_reentrant=False. This is required to unblock T5 activation checkpointing and MetaSeq use cases.
Closes https://github.com/pytorch/pytorch/issues/79887
Pull Request resolved: https://github.com/pytorch/pytorch/pull/80987
Approved by: https://github.com/zhaojuanmao