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

Fixes deepspeedai#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>

Author: adityasingh2400

deepspeed/runtime/zero/partition_parameters.py

@@ -64,8 +64,11 @@ def __init__(self, param: Parameter) -> None:
             raise RuntimeError(f"expected param {param.ds_summary()} to be available")
 
         if hasattr(param.ds_tensor, "ds_quant_scale"):
-            param.data = Init.quantizer_module.dequantize(param.ds_tensor.data, param.ds_tensor.ds_quant_scale).to(
-                device=get_accelerator().current_device_name(), non_blocking=True).view(param.ds_shape)
+            param.data = Init.quantizer_module.dequantize(param.ds_tensor.data,
+                                                          param.ds_tensor.ds_quant_scale,
+                                                          dtype=param.dtype).to(
+                                                              device=get_accelerator().current_device_name(),
+                                                              non_blocking=True).view(param.ds_shape)
         else:
             param.data = param.ds_tensor.data.to(device=get_accelerator().current_device_name(),
                                                  non_blocking=True).view(param.ds_shape)
@@ -87,8 +90,11 @@ def __init__(self, params: List[Parameter]) -> None:
             if param.ds_status != ZeroParamStatus.INFLIGHT:
                 raise RuntimeError(f"expected param {param.ds_summary()} to not be available")
             if hasattr(param.ds_tensor, "ds_quant_scale"):
-                param.data = Init.quantizer_module.dequantize(param.ds_tensor.data, param.ds_tensor.ds_quant_scale).to(
-                    device=get_accelerator().current_device_name(), non_blocking=True).view(param.ds_shape)
+                param.data = Init.quantizer_module.dequantize(param.ds_tensor.data,
+                                                              param.ds_tensor.ds_quant_scale,
+                                                              dtype=param.dtype).to(
+                                                                  device=get_accelerator().current_device_name(),
+                                                                  non_blocking=True).view(param.ds_shape)
             else:
                 param.data = param.ds_tensor.data.to(device=get_accelerator().current_device_name(),
                                                      non_blocking=True).view(param.ds_shape)
@@ -708,8 +714,10 @@ def wait(self, handle_dependency=True) -> None:
                 self.__original_dtype).to(self.__param.device)
         elif self.__quantization:
             instrument_w_nvtx(self.__quantization.quant_handle.wait)()
-            self.__param.data = self.__quantization.backend.dequantize(
-                self.__quantization.quantized_param, self.__quantization.scale_buffer).to(self.__param.device)
+            self.__param.data = self.__quantization.backend.dequantize(self.__quantization.quantized_param,
+                                                                       self.__quantization.scale_buffer,
+                                                                       dtype=self.__param.dtype).to(
+                                                                           self.__param.device)
         self.__param.ds_status = ZeroParamStatus.AVAILABLE
 
 
@@ -747,6 +755,9 @@ def wait(self, handle_dependency=True) -> None:
 
         if self.quantization:
             instrument_w_nvtx(self.quantization.quant_handle.wait)()
+            # No dtype here on purpose. A quantized coalesced bucket is not grouped by dtype the
+            # way the non-quantized path is, so params[0].dtype is not necessarily the dtype of
+            # the rest of the bucket. Each slice is cast to its own parameter's dtype below.
             flat_tensor = self.quantization.backend.dequantize(
                 self.quantization.quantized_param, self.quantization.scale_buffer).to(self.params[0].device)
 
@@ -865,12 +876,18 @@ def quantize(self, param, groups=None):
                 assert param.numel(
                 ) > groups, f"Adaptive grouping algorithm cannot find a group size for input tensor of size {param.numel()}"
                 self.group_size_cache[param.numel()] = groups
-        return self.quantizer_cuda_module.quantize(param.to(get_accelerator().device_name()), groups, 8,
-                                                   self.quantizer_cuda_module.Symmetric)
+        # The CUDA kernel reads its input through a __half* and always writes fp16 back out, so a bf16
+        # parameter would be reinterpreted bit-for-bit and silently corrupted. Convert on the way in and
+        # let the caller ask for its own dtype back on the way out.
+        param = param.to(get_accelerator().device_name(), dtype=torch.half)
+        return self.quantizer_cuda_module.quantize(param, groups, 8, self.quantizer_cuda_module.Symmetric)
 
-    def dequantize(self, quantized_param, scale):
-        return self.quantizer_cuda_module.dequantize(quantized_param, scale, scale.numel(), 8,
-                                                     self.quantizer_cuda_module.Symmetric)
+    def dequantize(self, quantized_param, scale, dtype=None):
+        dequantized = self.quantizer_cuda_module.dequantize(quantized_param, scale, scale.numel(), 8,
+                                                            self.quantizer_cuda_module.Symmetric)
+        if dtype is not None and dequantized.dtype != dtype:
+            dequantized = dequantized.to(dtype)
+        return dequantized
 
 
 def _no_gather_coalesced(params: Iterable[Parameter]) -> AllGatherCoalescedHandle:
@@ -2071,7 +2088,9 @@ def _allgather_params_coalesced(self, param_list, hierarchy=0, quantize=False):
         for i, param in enumerate(param_list):
             gathered_tensor = allgather_params[i]
             if quantize:
-                gathered_tensor = self.quantizer_module.dequantize(gathered_tensor, allgather_quantize_scale[i])
+                gathered_tensor = self.quantizer_module.dequantize(gathered_tensor,
+                                                                   allgather_quantize_scale[i],
+                                                                   dtype=param.dtype)
             param.data = gathered_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape).data
 
         # guarantee the communication to be completed
@@ -2130,7 +2149,7 @@ def _allgather_params_sequential(self, param_list, hierarchy=0):
                                     scale_partitions[partition_rank],
                                     group=self.get_partition_dp_group(param),
                                     async_op=False)
-                flat_tensor = self.quantizer_module.dequantize(flat_tensor, flat_scale_tensor)
+                flat_tensor = self.quantizer_module.dequantize(flat_tensor, flat_scale_tensor, dtype=param.dtype)
 
             param.data = flat_tensor.narrow(0, 0, param.ds_numel).view(param.ds_shape)
 

tests/unit/runtime/zero/test_zeropp.py

@@ -12,7 +12,7 @@
 import deepspeed
 
 from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
-from deepspeed.runtime.zero.partition_parameters import Init, ZeroParamStatus
+from deepspeed.runtime.zero.partition_parameters import CUDAQuantizer, Init, ZeroParamStatus
 
 import torch.nn as nn
 import torch
@@ -41,6 +41,41 @@ def test_zero_hpz_partition_size_config():
     assert config.zero_hpz_partition_size == 4
 
 
+class Fp16OnlyQuantizerModule:
+    """Stand-in for the compiled QuantizerBuilder op.
+
+    It mirrors the two properties of the real kernel that matter here: quantize() reads its input
+    through a __half*, and dequantize() always allocates an fp16 output tensor.
+    """
+
+    Symmetric = 0
+
+    def quantize(self, param, groups, num_bits, quant_type):
+        assert param.dtype == torch.half, f"the quantize kernel reads fp16, got {param.dtype}"
+        return param.to(torch.int8), torch.ones(groups, dtype=torch.float32, device=param.device)
+
+    def dequantize(self, quantized_param, scale, num_groups, num_bits, quant_type):
+        return quantized_param.to(torch.half)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.half])
+def test_cuda_quantizer_round_trips_parameter_dtype(monkeypatch, dtype):
+    """zero_quantized_weights must hand a parameter back in its own dtype.
+
+    The quantizer op is fp16-only, so under a bf16 config the weights used to come back as fp16 and
+    break the forward pass with a dtype mismatch. See #7775.
+    """
+    monkeypatch.setattr(CUDAQuantizer, "quantizer_cuda_module", Fp16OnlyQuantizerModule())
+    quantizer = CUDAQuantizer()
+
+    param = torch.randn(4096, dtype=dtype)
+    quantized_param, scale = quantizer.quantize(param)
+    assert quantized_param.dtype == torch.int8
+
+    dequantized = quantizer.dequantize(quantized_param, scale, dtype=param.dtype)
+    assert dequantized.dtype == dtype
+
+
 def test_zero_hpz_small_param_secondary_shard_without_overlap(monkeypatch):
 
     class _FakeAccelerator: