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

Set XLA_USE_SPMD for spmd cpp tests. #6273

Merged
merged 1 commit into from
Jan 11, 2024
Merged

Conversation

vanbasten23
Copy link
Collaborator

No description provided.

Copy link
Contributor

@yeounoh yeounoh left a comment

Choose a reason for hiding this comment

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

LGTM thanks

@vanbasten23 vanbasten23 force-pushed the set_SPMD_mode_for_spmd_test branch from ca33544 to b0d9aec Compare January 9, 2024 17:09
@vanbasten23 vanbasten23 marked this pull request as ready for review January 9, 2024 17:09
@yeounoh
Copy link
Contributor

yeounoh commented Jan 9, 2024

There seems to be a build issue? cc @vanbasten23

@vanbasten23
Copy link
Collaborator Author

There seems to be a build issue? cc @vanbasten23

Yeah, I'm looking at it.

class XLAShardingTest : public AtenXlaTensorTestBase {
protected:
static void SetUpTestCase() {
setenv("XLA_USE_SPMD", "1", /*overwrite=*/true);
Copy link
Collaborator

Choose a reason for hiding this comment

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

big picture question: what's our plan to remove XLA_USE_SPMD flag to improve usability?

@yeounoh @vanbasten23

Copy link
Collaborator Author

Choose a reason for hiding this comment

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

iiuc, on the user facing level (python), we want to replace XLA_USE_SPMD with an API use_spmd(). Internally, it seems fine to keep using the flag since it doesn't affect usability. Here we use the flag only internally.

@vanbasten23 vanbasten23 force-pushed the set_SPMD_mode_for_spmd_test branch from b0d9aec to 932bb33 Compare January 10, 2024 22:39
@vanbasten23
Copy link
Collaborator Author

rebased and triggered ci again

@vanbasten23
Copy link
Collaborator Author

Thanks for reviewing!

@vanbasten23 vanbasten23 merged commit a728afe into master Jan 11, 2024
20 checks passed
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants