Skip to content

Keep parameter dtype through ZeRO-3 weight quantization #8215

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
adityasingh2400 wants to merge 3 commits into
base: master
Choose a base branch
from
+68 −14
Open

Keep parameter dtype through ZeRO-3 weight quantization #8215

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6d691bc
Keep parameter dtype through ZeRO-3 weight quantization
adityasingh2400 Aug 5, 2026
b7fafc7
Do not force the coalesced dequantize to params[0] dtype
adityasingh2400 Aug 5, 2026
8d14024
Merge branch 'master' into fix-zero-quantized-weights-bf16
adityasingh2400 Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 32 additions & 13 deletions deepspeed/runtime/zero/partition_parameters.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down
37 changes: 36 additions & 1 deletion tests/unit/runtime/zero/test_zeropp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
Loading