Skip to content

[FSDPv2] Enable auto-wrapping#6499

Merged
alanwaketan merged 3 commits intomasterfrom
alanwaketan/fsdp2_auto_wrap
Feb 8, 2024
Merged

[FSDPv2] Enable auto-wrapping#6499
alanwaketan merged 3 commits intomasterfrom
alanwaketan/fsdp2_auto_wrap

Conversation

@alanwaketan
Copy link
Copy Markdown
Collaborator

Summary:
This pull request enables auto-wrapping for FSDPv2 following the original XlaFullyShardedDataParallel design and implementation.

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

Summary:
This pull request enables auto-wrapping for FSDPv2 following the original
XlaFullyShardedDataParallel design and implementation.

Test Plan:
python test/spmd/test_fsdp_v2.py -v -k test_fsdp_v2_auto_wrap_
@alanwaketan alanwaketan self-assigned this Feb 8, 2024
@alanwaketan alanwaketan marked this pull request as ready for review February 8, 2024 01:11
Copy link
Copy Markdown
Collaborator

@jonb377 jonb377 left a comment

Choose a reason for hiding this comment

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

LGTM!

@alanwaketan
Copy link
Copy Markdown
Collaborator Author

Thanks Jon for the quick review.

@alanwaketan alanwaketan merged commit cbe0ca8 into master Feb 8, 2024
amithrm pushed a commit to amithrm/xla that referenced this pull request Mar 1, 2024
Summary:
This pull request enables auto-wrapping for FSDPv2 following the original XlaFullyShardedDataParallel design and implementation.

Test Plan:
python test/spmd/test_fsdp_v2.py -v -k test_fsdp_v2_auto_wrap_
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
Summary:
This pull request enables auto-wrapping for FSDPv2 following the original XlaFullyShardedDataParallel design and implementation.

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