Fix bf16 dtype mismatch in ZeRO-3 with zero_quantized_weights

[juyterman1000]

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.

[PKUWZP]

@juyterman1000 Thanks for working on this. Can you check two things here: 1. Double-check AllGatherHandle (the non-coalesced version) doesn't have the same bug in its quantization path (lines 697-700 in partition_parameters.py file); 2. (Optional) Add a regression test for bf16 + zero_quantized_weights. I think 1) might be an interesting verification, and if true, we can have a follow up PR.

[juyterman1000]

Thanks @PKUWZP for the suggestions! I've updated the PR with both items.

[juyterman1000]

Hi @tohtana,
please have a look when you have a moment. I’m happy to make any necessary adjustments if you have concerns.