Keep parameter dtype through ZeRO-3 weight quantization - #8215
Keep parameter dtype through ZeRO-3 weight quantization#8215adityasingh2400 wants to merge 3 commits into
Conversation
The quantizer op is fp16-only in both directions: quantize_kernel 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 what deepspeedai#7775 reports for BERT. Convert to fp16 on the way into the kernel so the values are read correctly, and give dequantize an optional dtype so each caller can ask for the dtype its parameter actually has. The five call sites in the gather paths now pass the parameter dtype. Omitting the argument keeps the previous fp16 return, so no other caller changes behavior. Fixes deepspeedai#7775 Signed-off-by: Aditya Singh <adisin650@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6d691bc514
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| flat_tensor = self.quantization.backend.dequantize(self.quantization.quantized_param, | ||
| self.quantization.scale_buffer, | ||
| dtype=self.params[0].dtype).to(self.params[0].device) |
There was a problem hiding this comment.
When a coalesced quantized fetch contains parameters with mixed dtypes, this casts the entire dequantized flat buffer to only self.params[0].dtype before it is split; the per-parameter cast that happens later can restore the dtype, but it cannot undo rounding already applied to later fp16/fp32 parameters when the first parameter is bf16, and it can also inflate the whole buffer when the first parameter is fp32. The non-quantized coalesced path already buckets by dtype, so this path should either keep the kernel's fp16 output until each slice is assigned or bucket quantized coalesces by dtype.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Valid finding, fixed in b7fafc7.
I checked the premise before changing anything and it holds. In all_gather_coalesced, the not quantize branch groups parameters into dtype_params and issues one _all_gather_dtype per dtype, and that helper asserts every parameter in the bucket shares a communication dtype. The quantize branch does none of that, it concatenates every ds_tensor into a single int8 buffer regardless of dtype. So a quantized coalesced bucket really can be mixed, and params[0].dtype is not a safe stand-in for the rest of it.
Rather than bucket the quantized path by dtype, I dropped the argument at that one call site. The flat buffer keeps the kernel's fp16 output exactly as it did before this PR, and the existing per-slice line a few lines down already assigns each parameter its own dtype:
param.data = instrument_w_nvtx(torch.cat)(partitions).view(param.ds_shape).to(param.ds_tensor.dtype)That keeps the change minimal and avoids both failure modes you described, the rounding when params[0] is the narrower dtype and the buffer inflating when it is the wider one.
The other four call sites still pass dtype, and I checked each one: they assign param.data directly from the dequantized tensor with only a .view or a device move, so there is no later cast to restore the dtype and the argument is doing real work there.
yapf is clean on the file.
The quantized coalesced path concatenates every parameter into one int8 buffer without grouping by dtype, unlike the non-quantized path which buckets into one all-gather per dtype. Asking for params[0].dtype therefore casts the whole flat buffer to whatever the first parameter happens to be, and the later per-slice cast cannot undo rounding already applied to the rest. Each slice is already cast to its own parameter's dtype when it is assigned, so the flat buffer can keep the kernel's fp16 output as it did before. Signed-off-by: Aditya Singh <adisin650@gmail.com>
Resolves a conflict in tests/unit/runtime/zero/test_zeropp.py where both sides added tests. Keeps both, and unions the partition_parameters imports so CUDAQuantizer, Init and ZeroParamStatus are all available.
1 participant
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.