Skip to content

improve reshard_after_forward logic #1094

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

improve reshard_after_forward logic #1094

wants to merge 1 commit into from

Conversation

tianyu-l
Copy link
Contributor

@tianyu-l tianyu-l commented Apr 11, 2025

according to discussions in #1091

The CI failure is because FSDPMemTracker is not compatible of fully_shard on a list of modules. @sanketpurandare will help address this soon. Let's land it after the feature is available.

@facebook-github-bot facebook-github-bot added the CLA Signed This label is managed by the Meta Open Source bot. label Apr 11, 2025
@tianyu-l tianyu-l requested a review from wwwjn April 11, 2025 23:10
@tianyu-l tianyu-l linked an issue Apr 11, 2025 that may be closed by this pull request
@awgu
Copy link
Collaborator

awgu commented Apr 12, 2025

@tianyu-l I think it's also acceptable for now to allow the norm to be assigned to the root module. In other words, just wrap tok_embeddings separately and output separately.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CLA Signed This label is managed by the Meta Open Source bot.
Projects
None yet
Development

Successfully merging this pull request may close these issues.

FSDP2 root level parameter management
4 participants