-
Notifications
You must be signed in to change notification settings - Fork 505
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
[FSDPv2] Add a user guide #6386
Conversation
Summary: This diff adds a user guide for FSDPv2.
Can you also update https://github.com/pytorch/xla/blob/master/docs/source/index.rst so it will show up in https://pytorch.org/xla/release/2.1/index.html? |
# As normal FSDP, but an extra mesh is needed. | ||
model = FSDPv2(my_module, mesh) | ||
optim = torch.optim.Adam(model.parameters(), lr=0.0001) | ||
output = model(x, y) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We'll need to add input sharding here - do you want to add an example with MpDeviceLoader?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Let's manually shard the x here. Good catch!
Thanks, @JackCaoG! |
Summary: This diff adds a user guide for FSDPv2.
Summary: This diff adds a user guide for FSDPv2.
Summary:
This diff adds a user guide for FSDPv2.