DeepSpeed

6098f799 - Keep parameter dtype through ZeRO-3 weight quantization (#8215)

Keep parameter dtype through ZeRO-3 weight quantization (#8215) Fixes #7775 The quantizer op is fp16-only in both directions. quantize_kernel in csrc/quantization/pt_binding.cpp casts input_vals.data_ptr() to __half* whatever the tensor dtype actually is, and dequantize is bound as dequantize<__half> so it always allocates an fp16 output. CUDAQuantizer passed parameters straight through, so with bf16 enabled and zero_quantized_weights set, ZeRO-3 quantized bf16 bits reinterpreted as fp16 and then restored param.data as fp16. Training fails on the resulting dtype mismatch, which is the BERT failure in the issue. The bit reinterpretation on the way in is the quieter half of the bug: values are wrong before the dtype mismatch is ever noticed. CUDAQuantizer.quantize now converts to fp16 on the way into the kernel so the values are read correctly, and dequantize takes an optional dtype so each caller can ask for the dtype its parameter actually has. The five call sites in the gather paths pass the parameter dtype. Omitting the argument keeps the previous fp16 return, so no other caller changes behavior. Precision is not a concern here, since the values are being quantized to int8 regardless. Verification: added a parametrized test to tests/unit/runtime/zero/test_zeropp.py that stands in for the compiled op with a stub asserting the fp16 contract, and checks a bf16 and an fp16 parameter both round trip in their own dtype. It passes and it fails against the unmodified code on both halves of the fix. Run on CPU, since the test does not need the compiled op. The root cause is verified by reading csrc/quantization/pt_binding.cpp, not by running on a GPU. yapf and flake8 are clean on the changed files. --------- Signed-off-by: Aditya Singh <adisin650@gmail.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>