Skip to content

[FSDPv2] Use the global mesh API#6500

Merged
alanwaketan merged 2 commits intomasterfrom
alanwaketan/fsdp_v2_global
Feb 8, 2024
Merged

[FSDPv2] Use the global mesh API#6500
alanwaketan merged 2 commits intomasterfrom
alanwaketan/fsdp_v2_global

Conversation

@alanwaketan
Copy link
Copy Markdown
Collaborator

Summary:
Make FSDPv2 to use the global mesh API such that users can just set the global mesh.

Test Plan:
python test/spmd/test_fsdp_v2.py -v -k test_fsdp_v2_global_mesh

@alanwaketan alanwaketan self-assigned this Feb 8, 2024
@alanwaketan
Copy link
Copy Markdown
Collaborator Author

Thanks, Jon!

Summary:
Make FSDPv2 to use the global mesh API such that users
can just set the global mesh.

Test Plan:
python test/spmd/test_fsdp_v2.py -v -k test_fsdp_v2_global_mesh
@alanwaketan alanwaketan force-pushed the alanwaketan/fsdp_v2_global branch from 9cede36 to b3e2fad Compare February 8, 2024 07:47
@alanwaketan
Copy link
Copy Markdown
Collaborator Author

Let's skip the GPU tests given the pre-rebased one has passed all.

@alanwaketan alanwaketan merged commit ef8f6c3 into master Feb 8, 2024
@alanwaketan alanwaketan deleted the alanwaketan/fsdp_v2_global branch February 8, 2024 09:55
amithrm pushed a commit to amithrm/xla that referenced this pull request Mar 1, 2024
Summary:
Make FSDPv2 to use the global mesh API such that users can just set the global mesh.

Test Plan:
python test/spmd/test_fsdp_v2.py -v -k test_fsdp_v2_global_mesh
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
Summary:
Make FSDPv2 to use the global mesh API such that users can just set the global mesh.

Test Plan:
python test/spmd/test_fsdp_v2.py -v -k test_fsdp_v2_global_mesh
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.

2 participants