DeepSpeed
b079c309 - Fix ZeRO-3: use per-param dtype for quantize-scale buffer in _allgather_params_coalesced (#8128)

Commit
51 days ago
Fix ZeRO-3: use per-param dtype for quantize-scale buffer in _allgather_params_coalesced (#8128) ## Problem #8073 fixed `_allgather_params_coalesced` to allocate each output buffer with its own parameter's dtype, resolving the mixed-dtype crash under ZeRO-3 (bf16 base + fp32 LoRA adapters, #8072). The quantize-scale buffer in the same method was left on the old pattern. It still uses the first parameter's scale dtype: ```python for psize in quantize_scale_sizes: 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, device=self.local_device).view(-1) ``` Each scale buffer is gathered by its own `all_gather_into_tensor` call using that param's own scale tensor, so a persistent set with mixed scale dtypes hits the same "output tensor must have the same type as input tensor" error that #8073 removed for the main buffers. ## Fix The loop already binds `param` and uses it for `tensor_size`, so use it for the scale buffer dtype as well: ```python dtype=param.ds_tensor.ds_quant_scale.dtype, ``` This matches the main output buffer (fixed in #8073) and the sequential path `_allgather_params_sequential`, which already uses `param.ds_tensor.ds_quant_scale.dtype`. No cast is added, so numerics are unchanged, and it is a no-op when the scale dtypes are uniform. The trigger is narrow (quantized weights with mixed scale dtypes), so this is a consistency fix that removes the last shared-dtype assumption in the method rather than a widely hit crash. Signed-off-by: Sung Hyun Cho <hope5487@gmail.com>
Author
Parents
Loading