fix: add zero-guards for divisor/scale edge cases (#7838)

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)

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()

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)
 
 

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):

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)

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