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.

Sign up for GitHub

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

Jump to bottom

[DO NOT REVIEW] --experimental.fsdp_sharding_on_largest_dim #607

Open

weifengpy wants to merge 1 commit into pytorch:main from weifengpy:dim-i-sharding

Contributor

weifengpy commented Oct 9, 2024 •

edited

Loading

FSDP2 can shard on largest dim now

pytorch PR with extra copy [FSDP2] Added shard_placement_fn arg pytorch#137496
pytorch PR without extra copy but big cpu overhead using out_view tricks [benchmark] mimic new copy-out kernels with out_views pytorch#137683

becuase of dim-1, there are 2 extra copies (chunk-cat)

1st in all_gather_copy_out
2nd in post_backward_reduce

Screenshot 2024-10-09 at 4 27 33 PM


          [DO NOT REVIEW] --experimental.fsdp_sharding_on_largest_dim

c0b0d42

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

facebook-github-bot added the CLA Signed label

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels