Skip to content

[PyTorch] Support torch.amp.autocast in TE checkpoint - #791

Merged
denera merged 2 commits into
NVIDIA:mainfrom
denera:te-checkpoint-torch-autocast-support
May 22, 2024
Merged

[PyTorch] Support torch.amp.autocast in TE checkpoint#791
denera merged 2 commits into
NVIDIA:mainfrom
denera:te-checkpoint-torch-autocast-support

Conversation

@denera

Copy link
Copy Markdown
Collaborator

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

Reported in #787.

@deneradenera added the bug Something isn't working label Apr 18, 2024
@denera
denera requested a review from ksivamanApril 18, 2024 13:56
@deneradenera self-assigned this Apr 18, 2024
…rd pass during the recompute phase
Signed-off-by: Alp Dener <adener@nvidia.com>
@denera
deneraforce-pushed the te-checkpoint-torch-autocast-support branch from 60860bb to ab58966CompareMay 3, 2024 18:56
@ptrendx
ptrendx self-requested a review May 3, 2024 18:56

@ptrendxptrendx left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@ptrendx

Copy link
Copy Markdown
Member

/te-ci pytorch

1 similar comment
@denera

Copy link
Copy Markdown
CollaboratorAuthor

/te-ci pytorch

@denera
denera merged commit 7c4887b into NVIDIA:mainMay 22, 2024
ptrendx pushed a commit that referenced this pull request May 23, 2024
TE checkpoint now preserves the torch autocast context from the forward pass during the recompute phase
Signed-off-by: Alp Dener <adener@nvidia.com>
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

1.7.0bugSomething isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@denera@ptrendx