From 14de072942d876e54cd9befd99de9368961fd47f Mon Sep 17 00:00:00 2001 From: Sung Hyun Cho Date: Thu, 9 Jul 2026 15:15:14 +0900 Subject: [PATCH] Fix ZeRO-3: use per-param dtype for quantize-scale buffer in _allgather_params_coalesced #8073 switched the output buffers to each param's own dtype, but the quantize-scale buffer in the same method still uses param_list[0]. The loop already binds `param` and uses it for tensor_size, so use it for the scale buffer dtype too. This matches _allgather_params_sequential, which already uses param.ds_tensor.ds_quant_scale.dtype. Signed-off-by: Sung Hyun Cho --- deepspeed/runtime/zero/partition_parameters.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py index 9cdff411237c..e36d51e65c0d 100755 --- a/deepspeed/runtime/zero/partition_parameters.py +++ b/deepspeed/runtime/zero/partition_parameters.py @@ -1959,7 +1959,7 @@ def _allgather_params_coalesced(self, param_list, hierarchy=0, quantize=False): param = param_list[len(allgather_quantize_scale)] tensor_size = psize * self._partition_world_size(param) flat_tensor = torch.empty(tensor_size, - dtype=param_list[0].ds_tensor.ds_quant_scale.dtype, + dtype=param.ds_tensor.ds_quant_scale.dtype, device=self.local_device).view(-1) flat_tensor.requires_grad = False allgather_quantize_scale.append(flat_tensor)