Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[JAX] Add checkpoint_name for the recompute granularity control #542

Merged
merged 1 commit into from
Dec 4, 2023

Conversation

zlsh80826
Copy link
Collaborator

@zlsh80826 zlsh80826 commented Nov 28, 2023

Add checkpoint_name for TE's tensors. That enables the support for different checkpoint_policy in PAXML.

@zlsh80826
Copy link
Collaborator Author

/te-ci jax

@timmoon10 timmoon10 requested a review from denera November 29, 2023 23:45
Copy link
Collaborator

@denera denera left a comment

Choose a reason for hiding this comment

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

LGTM!

@zlsh80826 zlsh80826 requested a review from nouiz November 30, 2023 01:02
Signed-off-by: Reese Wang <[email protected]>
@zlsh80826 zlsh80826 force-pushed the rewang/add_checkpoint_name branch from 767474c to ed67507 Compare December 3, 2023 09:05
@zlsh80826
Copy link
Collaborator Author

/te-ci jax

@zlsh80826
Copy link
Collaborator Author

@denera, all tests are passed. Could you help merge the PR? Thanks

@denera denera merged commit c898ab1 into NVIDIA:main Dec 4, 2023
16 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants