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/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/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/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 new file mode 100644 index 000000000000..2911ab45f9e8 --- /dev/null +++ b/tests/unit/utils/test_zero_guards.py @@ -0,0 +1,62 @@ +# 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 + +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