From 6bc86305a8ac7de0e89a1d137cd574d5c87cfa4e Mon Sep 17 00:00:00 2001 From: liuyun7345 Date: Sun, 12 Jul 2026 04:25:58 +0000 Subject: [PATCH 1/2] fix: add zero-guards for divisor/scale edge cases (#7838) Prevent ZeroDivisionError and silent non-finite propagation at four reported sites: 1. groups._ensure_divisibility: reject denominator == 0 before modulo 2. ThroughputTimer: reject non-positive steps_per_output at init and report-boundary time 3. inference_utils.ceil_div: reject divisor == 0 4. HPU FPQuantizer.dequantize: reject zero/non-finite scales before 1/scale Adds focused unit tests covering each failure mode. Picks up the inactive draft #7855. Fixes #7838 Signed-off-by: liuyun7345 --- deepspeed/inference/v2/inference_utils.py | 2 + deepspeed/utils/groups.py | 1 + deepspeed/utils/timer.py | 5 ++ op_builder/hpu/fp_quantizer.py | 11 +++ tests/unit/utils/test_zero_guards.py | 97 +++++++++++++++++++++++ 5 files changed, 116 insertions(+) mode change 100755 => 100644 deepspeed/utils/timer.py create mode 100644 tests/unit/utils/test_zero_guards.py diff --git a/deepspeed/inference/v2/inference_utils.py b/deepspeed/inference/v2/inference_utils.py index 7b2dd4237353..8f2d5f051be2 100644 --- a/deepspeed/inference/v2/inference_utils.py +++ b/deepspeed/inference/v2/inference_utils.py @@ -102,4 +102,6 @@ def ceil_div(a: int, b: int) -> int: """ Return ceil(a / b). """ + if b == 0: + raise ValueError(f"ceil_div divisor must be non-zero (got a={a}, b={b})") return -(-a // b) diff --git a/deepspeed/utils/groups.py b/deepspeed/utils/groups.py index d912625c544b..140652ebefd6 100644 --- a/deepspeed/utils/groups.py +++ b/deepspeed/utils/groups.py @@ -63,6 +63,7 @@ def initialize(ep_size=1, mpu=None): def _ensure_divisibility(numerator, denominator): """Ensure that numerator is divisible by the denominator.""" + assert denominator != 0, f'denominator must be non-zero (got numerator={numerator}, denominator={denominator})' assert numerator % denominator == 0, '{} is not divisible by {}'.format(numerator, denominator) diff --git a/deepspeed/utils/timer.py b/deepspeed/utils/timer.py old mode 100755 new mode 100644 index 0aa7be55d829..389d9982b64d --- a/deepspeed/utils/timer.py +++ b/deepspeed/utils/timer.py @@ -211,6 +211,8 @@ def __init__(self, config, batch_size, start_step=2, steps_per_output=None, moni self.global_step_count = 0 self.total_elapsed_time = 0 self.step_elapsed_time = 0 + if steps_per_output is not None and steps_per_output <= 0: + raise ValueError(f"steps_per_output must be a positive integer or None, got {steps_per_output}") self.steps_per_output = steps_per_output self.monitor_memory = monitor_memory self.logging = logging_fn @@ -241,6 +243,9 @@ def start(self): def _is_report_boundary(self): if self.steps_per_output is None: return False + # Guard against mutation to 0 after construction (see #7838). + if self.steps_per_output <= 0: + raise ValueError(f"steps_per_output must be a positive integer, got {self.steps_per_output}") return self.global_step_count % self.steps_per_output == 0 def stop(self, global_step=False, report_speed=True): diff --git a/op_builder/hpu/fp_quantizer.py b/op_builder/hpu/fp_quantizer.py index c74affb55045..ca6016987367 100644 --- a/op_builder/hpu/fp_quantizer.py +++ b/op_builder/hpu/fp_quantizer.py @@ -4,6 +4,8 @@ # DeepSpeed Team +import math + import torch try: # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed @@ -54,6 +56,15 @@ def selective_dequantize(cls, val_q, scales, indexes, group_size, q_mantisa_bits @classmethod def dequantize(cls, fp_out, input_q, scale, group_size, q_mantisa_bits, q_exponent_bits): + # Reject zero / non-finite scales before inverse-scale computation to avoid + # silently propagating inf/nan into dequantized outputs (#7838). + if torch.is_tensor(scale): + if (not torch.isfinite(scale).all()) or (scale == 0).any(): + raise ValueError("FPQuantizer.dequantize requires finite non-zero scale values") + else: + scale_f = float(scale) + if scale_f == 0.0 or not math.isfinite(scale_f): + raise ValueError("FPQuantizer.dequantize requires a finite non-zero scale") orig_shape = fp_out.shape orig_dtype = fp_out.dtype dequant_out = torch.ops.hpu.cast_from_fp8(input_q, (1.0 / scale), orig_dtype).view(orig_shape) diff --git a/tests/unit/utils/test_zero_guards.py b/tests/unit/utils/test_zero_guards.py new file mode 100644 index 000000000000..bb95000d8924 --- /dev/null +++ b/tests/unit/utils/test_zero_guards.py @@ -0,0 +1,97 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team + +"""Regression tests for zero/division safety gaps reported in #7838.""" + +import math + +import pytest +import torch + +from deepspeed.utils.groups import _ensure_divisibility +from deepspeed.utils.timer import ThroughputTimer +from deepspeed.inference.v2.inference_utils import ceil_div + + +class _DummyTimerConfig: + enabled = False + synchronized = False + + +def test_ensure_divisibility_rejects_zero_denominator(): + with pytest.raises(AssertionError, match="non-zero"): + _ensure_divisibility(8, 0) + + +def test_ensure_divisibility_accepts_valid_inputs(): + _ensure_divisibility(8, 2) + _ensure_divisibility(0, 4) + + +def test_ceil_div_rejects_zero_divisor(): + with pytest.raises(ValueError, match="non-zero"): + ceil_div(10, 0) + + +def test_ceil_div_matches_math_ceil(): + assert ceil_div(10, 3) == math.ceil(10 / 3) + assert ceil_div(9, 3) == 3 + assert ceil_div(1, 1) == 1 + + +def test_throughput_timer_rejects_zero_steps_per_output(): + with pytest.raises(ValueError, match="positive"): + ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=0) + + +def test_throughput_timer_rejects_negative_steps_per_output(): + with pytest.raises(ValueError, match="positive"): + ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=-1) + + +def test_throughput_timer_report_boundary_guards_mutated_zero(): + timer = ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=2) + timer.steps_per_output = 0 + with pytest.raises(ValueError, match="positive"): + timer._is_report_boundary() + + +def test_throughput_timer_report_boundary_none_is_safe(): + timer = ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=None) + assert timer._is_report_boundary() is False + + +def _import_hpu_fp_quantizer_builder(): + try: + from op_builder.hpu.fp_quantizer import FPQuantizerBuilder + return FPQuantizerBuilder + except ImportError: + pytest.skip("HPU FPQuantizer builder is not available") + + +def test_hpu_fp_quantizer_dequantize_rejects_zero_scale(): + FPQuantizerBuilder = _import_hpu_fp_quantizer_builder() + scale = torch.tensor([0.0, 1.0]) + fp_out = torch.empty(2, 4) + input_q = torch.empty(2, 4) + with pytest.raises(ValueError, match="finite non-zero"): + FPQuantizerBuilder.dequantize(fp_out, input_q, scale, group_size=4, q_mantisa_bits=3, q_exponent_bits=4) + + +def test_hpu_fp_quantizer_dequantize_rejects_nonfinite_scale(): + FPQuantizerBuilder = _import_hpu_fp_quantizer_builder() + scale = torch.tensor([float("nan"), 1.0]) + fp_out = torch.empty(2, 4) + input_q = torch.empty(2, 4) + with pytest.raises(ValueError, match="finite non-zero"): + FPQuantizerBuilder.dequantize(fp_out, input_q, scale, group_size=4, q_mantisa_bits=3, q_exponent_bits=4) + + +def test_hpu_fp_quantizer_dequantize_rejects_zero_scalar_scale(): + FPQuantizerBuilder = _import_hpu_fp_quantizer_builder() + fp_out = torch.empty(2, 4) + input_q = torch.empty(2, 4) + with pytest.raises(ValueError, match="finite non-zero"): + FPQuantizerBuilder.dequantize(fp_out, input_q, 0.0, group_size=4, q_mantisa_bits=3, q_exponent_bits=4) From 1098c0a9cf41783604029933ceba6b76ab7cd9d7 Mon Sep 17 00:00:00 2001 From: liuyun7345 Date: Fri, 31 Jul 2026 10:58:37 +0800 Subject: [PATCH 2/2] Generalize FPQuantizer zero-scale guard (review feedback on #8133) Move the zero/non-finite scale validation out of the HPU-only `op_builder/hpu/fp_quantizer.py` and into the shared front-end `FP_Quantize.dequantize` / `selective_dequantize` in `deepspeed/ops/fp_quantizer/quantize.py`. The same `1.0 / scale` failure mode applies to every backend (CUDA included): `quantize` computes `scale = q_range / max_vals`, so inf/nan in the input silently yields `scale = 0` and corrupts dequantized output. One guard at the API boundary covers all backends. Relocate the regression tests from `tests/unit/utils/test_zero_guards.py` into `tests/unit/ops/fp_quantizer/test_fp_quant.py` using the existing accelerator-agnostic pattern (`deepspeed.ops.op_builder.FPQuantizerBuilder` + module-level skip via `__compatible_ops__`), exercising the guard through the public `FP_Quantize.dequantize` API. Drop the direct `op_builder.hpu` import and the now-unused `torch` import from `test_zero_guards.py`. Signed-off-by: liuyun7345 Co-authored-by: Cursor --- deepspeed/ops/fp_quantizer/quantize.py | 16 ++++ op_builder/hpu/fp_quantizer.py | 11 --- tests/unit/ops/fp_quantizer/test_fp_quant.py | 92 ++++++++++++++++++++ tests/unit/utils/test_zero_guards.py | 35 -------- 4 files changed, 108 insertions(+), 46 deletions(-) diff --git a/deepspeed/ops/fp_quantizer/quantize.py b/deepspeed/ops/fp_quantizer/quantize.py index 71fe96267f85..87124a945446 100644 --- a/deepspeed/ops/fp_quantizer/quantize.py +++ b/deepspeed/ops/fp_quantizer/quantize.py @@ -3,6 +3,8 @@ # DeepSpeed Team +import math + import torch import abc from abc import ABC @@ -14,6 +16,18 @@ fp_quant_module = None +def _validate_scale(scale) -> None: + """Reject zero / non-finite scales before inversion to avoid silently + propagating inf/nan into dequantized outputs (#7838).""" + if torch.is_tensor(scale): + if not torch.isfinite(scale).all() or (scale == 0).any(): + raise ValueError("FPQuantizer.dequantize requires finite non-zero scale values") + else: + scale_f = float(scale) + if scale_f == 0.0 or not math.isfinite(scale_f): + raise ValueError("FPQuantizer.dequantize requires a finite non-zero scale") + + class Quantizer(ABC): """ Abstract Quantizer class that implements quantize/dequantize methods. @@ -125,6 +139,7 @@ def dequantize(self, input_q, fp_out=None, q_bits=8, q_mantisa_bits=3, scale=Non f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" if scale is not None: + _validate_scale(scale) assert input_q.numel() == fp_out.numel(), \ '[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() @@ -158,6 +173,7 @@ def selective_dequantize(self, f"Missing {q_bits}-dequantization, please add the template arguments for the kernel to support this precision!" if scale is not None: + _validate_scale(scale) assert input_q.numel() == fp_out.numel(), \ '[De-quantization Error]: quantized data should have the same size as original tensor when scale is not None!' input_q = torch.cat([input_q.reshape(-1, self.group_size), scale], dim=-1).contiguous() diff --git a/op_builder/hpu/fp_quantizer.py b/op_builder/hpu/fp_quantizer.py index ca6016987367..c74affb55045 100644 --- a/op_builder/hpu/fp_quantizer.py +++ b/op_builder/hpu/fp_quantizer.py @@ -4,8 +4,6 @@ # DeepSpeed Team -import math - import torch try: # is op_builder from deepspeed or a 3p version? this should only succeed if it's deepspeed @@ -56,15 +54,6 @@ def selective_dequantize(cls, val_q, scales, indexes, group_size, q_mantisa_bits @classmethod def dequantize(cls, fp_out, input_q, scale, group_size, q_mantisa_bits, q_exponent_bits): - # Reject zero / non-finite scales before inverse-scale computation to avoid - # silently propagating inf/nan into dequantized outputs (#7838). - if torch.is_tensor(scale): - if (not torch.isfinite(scale).all()) or (scale == 0).any(): - raise ValueError("FPQuantizer.dequantize requires finite non-zero scale values") - else: - scale_f = float(scale) - if scale_f == 0.0 or not math.isfinite(scale_f): - raise ValueError("FPQuantizer.dequantize requires a finite non-zero scale") orig_shape = fp_out.shape orig_dtype = fp_out.dtype dequant_out = torch.ops.hpu.cast_from_fp8(input_q, (1.0 / scale), orig_dtype).view(orig_shape) diff --git a/tests/unit/ops/fp_quantizer/test_fp_quant.py b/tests/unit/ops/fp_quantizer/test_fp_quant.py index 0655b0ce26a3..f0f9ce52c2e9 100644 --- a/tests/unit/ops/fp_quantizer/test_fp_quant.py +++ b/tests/unit/ops/fp_quantizer/test_fp_quant.py @@ -132,3 +132,95 @@ def test_fp_quant(dtype, q_bits): ds_error = (x_dequantized - ds_x).abs().sum() / x.numel() assert 0.0004 > abs(qtorch_error.item() - ds_error.item()), f"failed on iteration {i}" + + +@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bf16"]) +def test_fp_quant_dequantize_rejects_zero_scale(dtype): + device_name = get_accelerator().device_name() + quant_config = QuantizationConfig() + quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype() + quant_config.group_size = 128 + fpq = FP_Quantize(quantization_config=quant_config) + + x = torch.rand(4, quant_config.group_size, dtype=dtype, device=device_name) + x_quantized = fpq.quantize(x, q_bits=8) + + zero_scale = torch.zeros(fpq.num_groups, 1, dtype=torch.float32, device=device_name) + with pytest.raises(ValueError, match="finite non-zero"): + fpq.dequantize(x_quantized, q_bits=8, scale=zero_scale) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bf16"]) +def test_fp_quant_dequantize_rejects_nonfinite_scale(dtype): + device_name = get_accelerator().device_name() + quant_config = QuantizationConfig() + quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype() + quant_config.group_size = 128 + fpq = FP_Quantize(quantization_config=quant_config) + + x = torch.rand(4, quant_config.group_size, dtype=dtype, device=device_name) + x_quantized = fpq.quantize(x, q_bits=8) + + nan_scale = torch.full((fpq.num_groups, 1), float("nan"), dtype=torch.float32, device=device_name) + with pytest.raises(ValueError, match="finite non-zero"): + fpq.dequantize(x_quantized, q_bits=8, scale=nan_scale) + + +def test_fp_quant_dequantize_rejects_zero_scalar_scale(): + device_name = get_accelerator().device_name() + quant_config = QuantizationConfig() + quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype() + quant_config.group_size = 128 + fpq = FP_Quantize(quantization_config=quant_config) + + x = torch.rand(4, quant_config.group_size, dtype=torch.bfloat16, device=device_name) + x_quantized = fpq.quantize(x, q_bits=8) + + with pytest.raises(ValueError, match="finite non-zero"): + fpq.dequantize(x_quantized, q_bits=8, scale=0.0) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bf16"]) +def test_fp_quant_dequantize_rejects_zero_scale(dtype): + device_name = get_accelerator().device_name() + quant_config = QuantizationConfig() + quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype() + quant_config.group_size = 128 + fpq = FP_Quantize(quantization_config=quant_config) + + x = torch.rand(4, quant_config.group_size, dtype=dtype, device=device_name) + x_quantized = fpq.quantize(x, q_bits=8) + + zero_scale = torch.zeros(fpq.num_groups, 1, dtype=torch.float32, device=device_name) + with pytest.raises(ValueError, match="finite non-zero"): + fpq.dequantize(x_quantized, q_bits=8, scale=zero_scale) + + +@pytest.mark.parametrize("dtype", [torch.bfloat16], ids=["bf16"]) +def test_fp_quant_dequantize_rejects_nonfinite_scale(dtype): + device_name = get_accelerator().device_name() + quant_config = QuantizationConfig() + quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype() + quant_config.group_size = 128 + fpq = FP_Quantize(quantization_config=quant_config) + + x = torch.rand(4, quant_config.group_size, dtype=dtype, device=device_name) + x_quantized = fpq.quantize(x, q_bits=8) + + nan_scale = torch.full((fpq.num_groups, 1), float("nan"), dtype=torch.float32, device=device_name) + with pytest.raises(ValueError, match="finite non-zero"): + fpq.dequantize(x_quantized, q_bits=8, scale=nan_scale) + + +def test_fp_quant_dequantize_rejects_zero_scalar_scale(): + device_name = get_accelerator().device_name() + quant_config = QuantizationConfig() + quant_config.q_dtype = FPQuantizerBuilder.get_default_quant_dtype() + quant_config.group_size = 128 + fpq = FP_Quantize(quantization_config=quant_config) + + x = torch.rand(4, quant_config.group_size, dtype=torch.bfloat16, device=device_name) + x_quantized = fpq.quantize(x, q_bits=8) + + with pytest.raises(ValueError, match="finite non-zero"): + fpq.dequantize(x_quantized, q_bits=8, scale=0.0) diff --git a/tests/unit/utils/test_zero_guards.py b/tests/unit/utils/test_zero_guards.py index bb95000d8924..2911ab45f9e8 100644 --- a/tests/unit/utils/test_zero_guards.py +++ b/tests/unit/utils/test_zero_guards.py @@ -8,7 +8,6 @@ import math import pytest -import torch from deepspeed.utils.groups import _ensure_divisibility from deepspeed.utils.timer import ThroughputTimer @@ -61,37 +60,3 @@ def test_throughput_timer_report_boundary_guards_mutated_zero(): def test_throughput_timer_report_boundary_none_is_safe(): timer = ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=None) assert timer._is_report_boundary() is False - - -def _import_hpu_fp_quantizer_builder(): - try: - from op_builder.hpu.fp_quantizer import FPQuantizerBuilder - return FPQuantizerBuilder - except ImportError: - pytest.skip("HPU FPQuantizer builder is not available") - - -def test_hpu_fp_quantizer_dequantize_rejects_zero_scale(): - FPQuantizerBuilder = _import_hpu_fp_quantizer_builder() - scale = torch.tensor([0.0, 1.0]) - fp_out = torch.empty(2, 4) - input_q = torch.empty(2, 4) - with pytest.raises(ValueError, match="finite non-zero"): - FPQuantizerBuilder.dequantize(fp_out, input_q, scale, group_size=4, q_mantisa_bits=3, q_exponent_bits=4) - - -def test_hpu_fp_quantizer_dequantize_rejects_nonfinite_scale(): - FPQuantizerBuilder = _import_hpu_fp_quantizer_builder() - scale = torch.tensor([float("nan"), 1.0]) - fp_out = torch.empty(2, 4) - input_q = torch.empty(2, 4) - with pytest.raises(ValueError, match="finite non-zero"): - FPQuantizerBuilder.dequantize(fp_out, input_q, scale, group_size=4, q_mantisa_bits=3, q_exponent_bits=4) - - -def test_hpu_fp_quantizer_dequantize_rejects_zero_scalar_scale(): - FPQuantizerBuilder = _import_hpu_fp_quantizer_builder() - fp_out = torch.empty(2, 4) - input_q = torch.empty(2, 4) - with pytest.raises(ValueError, match="finite non-zero"): - FPQuantizerBuilder.dequantize(fp_out, input_q, 0.0, group_size=4, q_mantisa_bits=3, q_exponent_bits=4)