Skip to content

Conversation

@alanwaketan
Copy link
Collaborator

Summary:
This diff adds a user guide for FSDPv2.

Summary:
This diff adds a user guide for FSDPv2.
@alanwaketan alanwaketan self-assigned this Jan 26, 2024
@JackCaoG
Copy link
Collaborator

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)
Copy link
Collaborator

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?

Copy link
Collaborator Author

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!

@alanwaketan
Copy link
Collaborator Author

Thanks, @JackCaoG!

@alanwaketan alanwaketan merged commit de5d764 into master Jan 29, 2024
@alanwaketan alanwaketan deleted the alanwaketan/fsdp_v2_tutorial branch January 29, 2024 20:00
alanwaketan added a commit that referenced this pull request Jan 29, 2024
Summary:
This diff adds a user guide for FSDPv2.
bhavya01 pushed a commit that referenced this pull request Apr 22, 2024
Summary:
This diff adds a user guide for FSDPv2.
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.

5 participants