-
Notifications
You must be signed in to change notification settings - Fork 417
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
Fix HSDP with ShardDegree < 8 #3313
Conversation
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.
Tested with : mpt-125m-hybrid-resumption-1-node-7QStEE
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.
LGTM! sorry one more request before approval
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.
Big Ning stays WinNing
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.
Nice! But why would anyone use HSDP with only one node?
I guess at least unit test needs it ? lol |
what
@j316chuck found HSDP with 1 node 2X4 device mesh run can't resume. rank 0 failed this checkpoint files broadcast . It's because rank 4 (which is in the same replica pg group as rank 0) doesn't join this collective call because of this if condition.
So the fix is to let all the first rank of shard pg group joins the broadcast , but only local rank 0 writes to file
test
1 node device mesh: 2 X 4
before:
[nccl timeout] mpt-125m-hybrid-resumption-1-node-tYv2bB
after:
mpt-125m-hybrid-resumption-1-node-pzWqX1
mpt-125m-hybrid-resumption-1-node-RV4ovO
2 node device mesh: 2 X 8 still works
mpt-125m-hybrid-resumption-1-node-NlOtP5
2 node device mesh: 4 x 4 works:
mpt-125m-hybrid-resumption-1-node-ajvSHb
checked the logging, rank 8 is not downloading any checkpoint files, but it got all files broadcast from rank 0