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 1 commit
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
47 changes: 32 additions & 15 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,8 +755,9 @@ def wait(self, handle_dependency=True) -> None:

if self.quantization:
instrument_w_nvtx(self.quantization.quant_handle.wait)()
flat_tensor = self.quantization.backend.dequantize(
self.quantization.quantized_param, self.quantization.scale_buffer).to(self.params[0].device)
flat_tensor = self.quantization.backend.dequantize(self.quantization.quantized_param,
self.quantization.scale_buffer,
dtype=self.params[0].dtype).to(self.params[0].device)

@chatgpt-codex-connector chatgpt-codex-connector Bot Aug 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Keep coalesced dequantization independent of the first dtype

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 👍 / 👎.

@adityasingh2400 adityasingh2400 Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.


self.partitions: List[Parameter] = []
for i in range(self.world_size):
Expand Down Expand Up @@ -865,12 +874,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 @@ -2069,7 +2084,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 @@ -2128,7 +2145,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
36 changes: 36 additions & 0 deletions 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,6 +12,7 @@
import deepspeed

from deepspeed.runtime.zero.config import DeepSpeedZeroConfig
from deepspeed.runtime.zero.partition_parameters import CUDAQuantizer

import torch.nn as nn
import torch
Expand Down Expand Up @@ -40,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 _assert_no_secondary_tensor_group(model: Module) -> None:
for _, param in model.named_parameters():
assert param.ds_secondary_tensor is None
Expand Down
Loading