diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index 04fb733b3076..072da6823963 100755 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/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) diff --git a/tests/unit/runtime/zero/test_zeropp.py b/tests/unit/runtime/zero/test_zeropp.py index a0155fba452e..b57510235986 100644 --- a/tests/unit/runtime/zero/test_zeropp.py +++ b/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: