Skip to content

Commit 6bc8630

Browse files
committed
fix: add zero-guards for divisor/scale edge cases (deepspeedai#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 deepspeedai#7855. Fixes deepspeedai#7838 Signed-off-by: liuyun7345 <liuyun7345@sina.com>
1 parent 429e2ad commit 6bc8630

5 files changed

Lines changed: 116 additions & 0 deletions

File tree

deepspeed/inference/v2/inference_utils.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,4 +102,6 @@ def ceil_div(a: int, b: int) -> int:
102102
"""
103103
Return ceil(a / b).
104104
"""
105+
if b == 0:
106+
raise ValueError(f"ceil_div divisor must be non-zero (got a={a}, b={b})")
105107
return -(-a // b)

deepspeed/utils/groups.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ def initialize(ep_size=1, mpu=None):
6363

6464
def _ensure_divisibility(numerator, denominator):
6565
"""Ensure that numerator is divisible by the denominator."""
66+
assert denominator != 0, f'denominator must be non-zero (got numerator={numerator}, denominator={denominator})'
6667
assert numerator % denominator == 0, '{} is not divisible by {}'.format(numerator, denominator)
6768

6869

deepspeed/utils/timer.py

100755100644
Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,8 @@ def __init__(self, config, batch_size, start_step=2, steps_per_output=None, moni
211211
self.global_step_count = 0
212212
self.total_elapsed_time = 0
213213
self.step_elapsed_time = 0
214+
if steps_per_output is not None and steps_per_output <= 0:
215+
raise ValueError(f"steps_per_output must be a positive integer or None, got {steps_per_output}")
214216
self.steps_per_output = steps_per_output
215217
self.monitor_memory = monitor_memory
216218
self.logging = logging_fn
@@ -241,6 +243,9 @@ def start(self):
241243
def _is_report_boundary(self):
242244
if self.steps_per_output is None:
243245
return False
246+
# Guard against mutation to 0 after construction (see #7838).
247+
if self.steps_per_output <= 0:
248+
raise ValueError(f"steps_per_output must be a positive integer, got {self.steps_per_output}")
244249
return self.global_step_count % self.steps_per_output == 0
245250

246251
def stop(self, global_step=False, report_speed=True):

op_builder/hpu/fp_quantizer.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,8 @@
44

55
# DeepSpeed Team
66

7+
import math
8+
79
import torch
810
try:
911
# 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
5456

5557
@classmethod
5658
def dequantize(cls, fp_out, input_q, scale, group_size, q_mantisa_bits, q_exponent_bits):
59+
# Reject zero / non-finite scales before inverse-scale computation to avoid
60+
# silently propagating inf/nan into dequantized outputs (#7838).
61+
if torch.is_tensor(scale):
62+
if (not torch.isfinite(scale).all()) or (scale == 0).any():
63+
raise ValueError("FPQuantizer.dequantize requires finite non-zero scale values")
64+
else:
65+
scale_f = float(scale)
66+
if scale_f == 0.0 or not math.isfinite(scale_f):
67+
raise ValueError("FPQuantizer.dequantize requires a finite non-zero scale")
5768
orig_shape = fp_out.shape
5869
orig_dtype = fp_out.dtype
5970
dequant_out = torch.ops.hpu.cast_from_fp8(input_q, (1.0 / scale), orig_dtype).view(orig_shape)
Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
# Copyright (c) Microsoft Corporation.
2+
# SPDX-License-Identifier: Apache-2.0
3+
4+
# DeepSpeed Team
5+
6+
"""Regression tests for zero/division safety gaps reported in #7838."""
7+
8+
import math
9+
10+
import pytest
11+
import torch
12+
13+
from deepspeed.utils.groups import _ensure_divisibility
14+
from deepspeed.utils.timer import ThroughputTimer
15+
from deepspeed.inference.v2.inference_utils import ceil_div
16+
17+
18+
class _DummyTimerConfig:
19+
enabled = False
20+
synchronized = False
21+
22+
23+
def test_ensure_divisibility_rejects_zero_denominator():
24+
with pytest.raises(AssertionError, match="non-zero"):
25+
_ensure_divisibility(8, 0)
26+
27+
28+
def test_ensure_divisibility_accepts_valid_inputs():
29+
_ensure_divisibility(8, 2)
30+
_ensure_divisibility(0, 4)
31+
32+
33+
def test_ceil_div_rejects_zero_divisor():
34+
with pytest.raises(ValueError, match="non-zero"):
35+
ceil_div(10, 0)
36+
37+
38+
def test_ceil_div_matches_math_ceil():
39+
assert ceil_div(10, 3) == math.ceil(10 / 3)
40+
assert ceil_div(9, 3) == 3
41+
assert ceil_div(1, 1) == 1
42+
43+
44+
def test_throughput_timer_rejects_zero_steps_per_output():
45+
with pytest.raises(ValueError, match="positive"):
46+
ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=0)
47+
48+
49+
def test_throughput_timer_rejects_negative_steps_per_output():
50+
with pytest.raises(ValueError, match="positive"):
51+
ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=-1)
52+
53+
54+
def test_throughput_timer_report_boundary_guards_mutated_zero():
55+
timer = ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=2)
56+
timer.steps_per_output = 0
57+
with pytest.raises(ValueError, match="positive"):
58+
timer._is_report_boundary()
59+
60+
61+
def test_throughput_timer_report_boundary_none_is_safe():
62+
timer = ThroughputTimer(_DummyTimerConfig(), batch_size=1, steps_per_output=None)
63+
assert timer._is_report_boundary() is False
64+
65+
66+
def _import_hpu_fp_quantizer_builder():
67+
try:
68+
from op_builder.hpu.fp_quantizer import FPQuantizerBuilder
69+
return FPQuantizerBuilder
70+
except ImportError:
71+
pytest.skip("HPU FPQuantizer builder is not available")
72+
73+
74+
def test_hpu_fp_quantizer_dequantize_rejects_zero_scale():
75+
FPQuantizerBuilder = _import_hpu_fp_quantizer_builder()
76+
scale = torch.tensor([0.0, 1.0])
77+
fp_out = torch.empty(2, 4)
78+
input_q = torch.empty(2, 4)
79+
with pytest.raises(ValueError, match="finite non-zero"):
80+
FPQuantizerBuilder.dequantize(fp_out, input_q, scale, group_size=4, q_mantisa_bits=3, q_exponent_bits=4)
81+
82+
83+
def test_hpu_fp_quantizer_dequantize_rejects_nonfinite_scale():
84+
FPQuantizerBuilder = _import_hpu_fp_quantizer_builder()
85+
scale = torch.tensor([float("nan"), 1.0])
86+
fp_out = torch.empty(2, 4)
87+
input_q = torch.empty(2, 4)
88+
with pytest.raises(ValueError, match="finite non-zero"):
89+
FPQuantizerBuilder.dequantize(fp_out, input_q, scale, group_size=4, q_mantisa_bits=3, q_exponent_bits=4)
90+
91+
92+
def test_hpu_fp_quantizer_dequantize_rejects_zero_scalar_scale():
93+
FPQuantizerBuilder = _import_hpu_fp_quantizer_builder()
94+
fp_out = torch.empty(2, 4)
95+
input_q = torch.empty(2, 4)
96+
with pytest.raises(ValueError, match="finite non-zero"):
97+
FPQuantizerBuilder.dequantize(fp_out, input_q, 0.0, group_size=4, q_mantisa_bits=3, q_exponent_bits=4)

0 commit comments

Comments
 (0)