Fix bf16 dtype mismatch in ZeRO-3 with zero_quantized_weights#7792
Open
juyterman1000 wants to merge 3 commits intodeepspeedai:masterfrom
Open
Fix bf16 dtype mismatch in ZeRO-3 with zero_quantized_weights#7792juyterman1000 wants to merge 3 commits intodeepspeedai:masterfrom
juyterman1000 wants to merge 3 commits intodeepspeedai:masterfrom
Conversation
8f82004 to
1c95ade
Compare
Collaborator
|
@juyterman1000 Thanks for working on this. Can you check two things here: 1. Double-check |
e25d8f1 to
72da75f
Compare
Contributor
Author
|
Thanks @PKUWZP for the suggestions! I've updated the PR with both items. |
377b76b to
83bd5d1
Compare
Signed-off-by: juyterman1000 <fastrunner10090@gmail.com>
83bd5d1 to
3940f49
Compare
Contributor
Author
|
Hi @tohtana, |
90ae56e to
3940f49
Compare
Signed-off-by: juyterman1000 <fastrunner10090@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
When using ZeRO-3 with zero_quantized_weights=True and bf16 enabled, the dequantized weights were incorrectly cast to fp16 instead of preserving the original bf16 dtype. This caused RuntimeError during training.
The fix adds original_dtype tracking to AllGatherCoalescedHandle, mirroring the existing pattern in AllGatherHandle, to ensure weights are converted back to their original dtype after dequantization.