TransformerEngine icon indicating copy to clipboard operation
TransformerEngine copied to clipboard
verified publisher icon NVIDIA

[PyTorch] Support `torch.amp.autocast` in TE checkpoint

Open denera opened this issue 1 year ago • 1 comments

This PR modifies te.distributed.checkpoint(...) to preserve the torch.amp.autocast(...) context from the forward pass during the recompute phase.

Reported in #787.

denera avatar Apr 18 '24 13:04 denera

/te-ci pytorch

ptrendx avatar May 17 '24 23:05 ptrendx

/te-ci pytorch

denera avatar May 21 '24 22:05 denera