From 8c8cab2a954810216a45d65c1b8d56839c846e5b Mon Sep 17 00:00:00 2001 From: Kyozzz <1147385157@qq.com> Date: Sun, 18 Oct 2026 13:33:38 +0800 Subject: [PATCH 1/2] Guard dim_t overflow in compute_size, reshape and grow (#2097) * Guard dim_t overflow in compute_size, reshape and grow The model loader validates shapes with a per-multiplication overflow check (since #2091), but the StorageView API paths still multiply dims unchecked: compute_size() and StorageView::reshape() can wrap int64_t (e.g. a shape of {2^62+1, 8} wraps to 8), yielding a view whose shape claims far more elements than _size, and grow() can wrap its addition. Apply the same per-multiplication guard used by the loader to compute_size() and reshape(), add an overflow check to grow(), and reject the overflowing cases with std::invalid_argument in two new unit tests, mirroring the #2094 test style. * Accept zero dimensions in the overflow guards A zero dimension yields an empty tensor, which cannot overflow, but the guards divided by it to compare with the running product: x86_64 raises SIGFPE, while AArch64 division by zero returns 0 and made every later check reject valid zero-dim shapes. Skip the division when the dimension is zero. * Check empty dimensions before multiplying the shape size A zero dimension makes the size 0 whatever its position, so the other dimensions are never multiplied and an overflowing prefix next to a zero dimension is not an error. Skipping the division is not enough: the prefix can overflow before the zero is reached, which made the result depend on where the zero appeared in the shape. The same ordering in reshape() also left the known size at 0 when the new shape combined a zero dimension with an unknown one, and that 0 is the divisor of the compatibility check below (crash on master today). The unknown dimension is now rejected in that case. * Cover the grow() guard with unit tests Split the negative increment out of the overflow condition so it gets its own message, and pin the three cases: a negative size (which used to silently shrink the dimension), an overflowing addition (tested on {INT64_MAX, 0}, where the zero dimension keeps the allocation empty), and normal growth. --------- Co-authored-by: Jordi Mas --- include/ctranslate2/storage_view.h | 19 ++++- src/storage_view.cc | 28 +++++++- tests/storage_view_test.cc | 111 +++++++++++++++++++++++++++++ 3 files changed, 156 insertions(+), 2 deletions(-) diff --git a/include/ctranslate2/storage_view.h b/include/ctranslate2/storage_view.h index 1d4ef27d5..da34ec9cb 100644 --- a/include/ctranslate2/storage_view.h +++ b/include/ctranslate2/storage_view.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -29,9 +30,25 @@ namespace ctranslate2 { using Shape = std::vector; inline dim_t compute_size(const Shape& shape) { + bool has_zero_dim = false; + for (const dim_t dim : shape) { + if (dim < 0) + THROW_INVALID_ARGUMENT("shape contains a negative dimension"); + if (dim == 0) + has_zero_dim = true; + } + + // A zero dimension makes the size 0 whatever its position and whatever the + // other dimensions, so nothing is multiplied and no overflow is possible. + if (has_zero_dim) + return 0; + dim_t size = 1; - for (const dim_t dim : shape) + for (const dim_t dim : shape) { + if (size > std::numeric_limits::max() / dim) + THROW_INVALID_ARGUMENT("shape leads to an overflowing tensor size"); size *= dim; + } return size; } diff --git a/src/storage_view.cc b/src/storage_view.cc index 76b143bfd..07179256b 100644 --- a/src/storage_view.cc +++ b/src/storage_view.cc @@ -190,12 +190,15 @@ namespace ctranslate2 { StorageView& StorageView::reshape(Shape new_shape) { dim_t unknown_dim = -1; dim_t known_size = 1; + bool has_zero_dim = false; for (size_t i = 0; i < new_shape.size(); ++i) { const dim_t dim = new_shape[i]; if (dim >= 0) { - known_size *= dim; + // the empty dimensions are only recorded here: the size is computed once + // the whole shape is known to contain no negative dimension + has_zero_dim = has_zero_dim || dim == 0; } else if (dim == -1) { if (unknown_dim >= 0) THROW_INVALID_ARGUMENT("only one dimension can be set to -1, got -1 for dimensions " @@ -207,6 +210,25 @@ namespace ctranslate2 { } } + // A zero dimension makes the new size 0 whatever its position and whatever the + // other dimensions, so nothing is multiplied and no overflow is possible. The + // unknown dimension can't be recovered from an empty size, and the compatibility + // check below would take the modulo by 0. + if (has_zero_dim) { + if (unknown_dim >= 0) + THROW_INVALID_ARGUMENT("can't infer dimension " + std::to_string(unknown_dim) + + " of a new shape with a zero dimension"); + known_size = 0; + } else { + for (const dim_t dim : new_shape) { + if (dim <= 0) + continue; + if (known_size > std::numeric_limits::max() / dim) + THROW_INVALID_ARGUMENT("new shape leads to an overflowing tensor size"); + known_size *= dim; + } + } + if (unknown_dim >= 0) { if (_size % known_size != 0) THROW_INVALID_ARGUMENT("current size (" + std::to_string(_size) @@ -261,6 +283,10 @@ namespace ctranslate2 { StorageView& StorageView::grow(dim_t dim, dim_t size) { GUARD_DIM(dim, rank()); + if (size < 0) + THROW_INVALID_ARGUMENT("can't grow a dimension by a negative size"); + if (_shape[dim] > std::numeric_limits::max() - size) + THROW_INVALID_ARGUMENT("new dimension size leads to an overflowing tensor size"); return resize(dim, _shape[dim] + size); } diff --git a/tests/storage_view_test.cc b/tests/storage_view_test.cc index e670e9dc9..6229cbf44 100644 --- a/tests/storage_view_test.cc +++ b/tests/storage_view_test.cc @@ -145,3 +145,114 @@ INSTANTIATE_TEST_SUITE_P(CPU, StorageViewDeviceTest, ::testing::Values(Device::C #ifdef CT2_WITH_CUDA INSTANTIATE_TEST_SUITE_P(CUDA, StorageViewDeviceTest, ::testing::Values(Device::CUDA)); #endif + +TEST(StorageViewTest, OverflowingShapeThrows) { + // (2^62 + 1) * 8 overflows dim_t: the shape must be rejected, not wrapped + // to a small size that passes allocation. + EXPECT_THROW(StorageView(Shape{(int64_t{1} << 62) + 1, 8}, DataType::INT32), + std::invalid_argument); +} + +TEST(StorageViewTest, OverflowingReshapeThrows) { + StorageView view(Shape{8}, DataType::INT32); + EXPECT_THROW(view.reshape(Shape{(int64_t{1} << 62) + 1, 8}), std::invalid_argument); +} + +TEST(StorageViewTest, ZeroDimIsValidWhateverItsPosition) { + // The size of a shape with a zero dimension is 0, so the other dimensions are + // never multiplied and the overflowing prefix below is not an error. Multiplying + // before looking for the zero dimension made the result depend on its position. + EXPECT_THROW(compute_size(Shape{(int64_t{1} << 62) + 1, 8}), std::invalid_argument); + EXPECT_EQ(compute_size(Shape{(int64_t{1} << 62) + 1, 8, 0}), 0); + EXPECT_EQ(compute_size(Shape{0, (int64_t{1} << 62) + 1, 8}), 0); + + StorageView view(Shape{(int64_t{1} << 62) + 1, 8, 0}, DataType::INT32); + EXPECT_EQ(view.size(), 0); + EXPECT_EQ(view.dim(2), 0); + view.reshape(Shape{0, (int64_t{1} << 62) + 1, 8}); + EXPECT_EQ(view.size(), 0); + EXPECT_EQ(view.dim(0), 0); +} + +TEST(StorageViewTest, ZeroDimDoesNotSkipNegativeCheck) { + EXPECT_THROW(compute_size(Shape{0, -2}), std::invalid_argument); + EXPECT_THROW(compute_size(Shape{(int64_t{1} << 62) + 1, 0, -1}), std::invalid_argument); + + StorageView view(Shape{0}, DataType::INT32); + EXPECT_THROW(view.reshape(Shape{0, 8, -3}), std::invalid_argument); +} + +TEST(StorageViewTest, UnknownDimWithZeroDimThrows) { + // An unknown dimension can't be inferred from an empty shape, and the known size + // of 0 used to be the divisor of the compatibility check below. + StorageView view(Shape{8}, DataType::INT32); + EXPECT_THROW(view.reshape(Shape{0, -1}), std::invalid_argument); + EXPECT_THROW(view.reshape(Shape{-1, 0}), std::invalid_argument); +} + +TEST(StorageViewTest, GrowRejectsNegativeIncrement) { + // grow() with a negative size used to reach resize() and silently shrink the + // dimension, and INT64_MIN made the addition wrap. + StorageView view(Shape{4}, DataType::INT32); + for (const dim_t size : {dim_t(-1), std::numeric_limits::min()}) { + try { + view.grow(0, size); + FAIL() << "Expected a negative growth to be rejected"; + } catch (const std::invalid_argument& e) { + EXPECT_NE(std::string(e.what()).find("can't grow a dimension by a negative size"), + std::string::npos); + } + } + + // A rejected grow leaves the storage untouched. + EXPECT_EQ(view.dim(0), 4); + EXPECT_EQ(view.size(), 4); +} + +TEST(StorageViewTest, GrowRejectsAdditionOverflow) { + // The zero dimension makes the size 0, so this overflowing first dimension costs + // no memory and still exercises the guard on the addition itself. + StorageView view(Shape{std::numeric_limits::max(), 0}, DataType::INT32); + EXPECT_EQ(view.size(), 0); + + EXPECT_THROW(view.grow(0, 1), std::invalid_argument); + EXPECT_THROW(view.grow(0, std::numeric_limits::max()), std::invalid_argument); + + // Growing the zero dimension by INT64_MAX doesn't overflow the addition (0 + + // INT64_MAX), so grow() lets it through and compute_size() rejects the resulting + // shape instead of wrapping it. + EXPECT_THROW(view.grow(1, std::numeric_limits::max()), std::invalid_argument); + + EXPECT_EQ(view.dim(0), std::numeric_limits::max()); + EXPECT_EQ(view.dim(1), 0); + EXPECT_EQ(view.size(), 0); +} + +TEST(StorageViewTest, GrowingToAnUnallocatableSizeIsRejected) { + // INT64_MAX is a valid dimension and INT64_MAX elements still fit in a dim_t, so + // neither grow() nor compute_size() rejects it: reserve() catches the byte size + // before allocating. The three guards leave no path through to the allocator. + StorageView view(Shape{0}, DataType::INT32); + EXPECT_THROW(view.grow(0, std::numeric_limits::max()), std::runtime_error); + EXPECT_EQ(view.size(), 0); +} + +TEST(StorageViewTest, GrowIncreasesShape) { + StorageView view(Shape{2, 3}, DataType::INT32); + EXPECT_EQ(view.size(), 6); + + view.grow(0, 1); + EXPECT_EQ(view.dim(0), 3); + EXPECT_EQ(view.dim(1), 3); + EXPECT_EQ(view.size(), 9); + + view.grow(1, 2); + EXPECT_EQ(view.dim(1), 5); + EXPECT_EQ(view.size(), 15); + + // Growing an empty dimension back to a positive size stays valid. + StorageView empty_view(Shape{0}, DataType::INT32); + empty_view.grow(0, 4); + EXPECT_EQ(empty_view.dim(0), 4); + EXPECT_EQ(empty_view.size(), 4); +} From 0495cfa02ef28e069c5b3c867703a83828aafa05 Mon Sep 17 00:00:00 2001 From: Jordi Mas Date: Mon, 19 Oct 2026 07:23:34 +0200 Subject: [PATCH 2/2] Fix CPU Conv1D dilation and INT8 scale overflow (#2117) * Fix CPU Conv1D dilation in oneDNN and GEMM paths * Keep INT8 quantization scales finite for tiny weight rows --- python/ctranslate2/specs/model_spec.py | 2 ++ python/tests/test_spec.py | 15 +++++++++++++++ src/ops/conv1d_cpu.cc | 10 ++++------ tests/ops_test.cc | 18 ++++++++++++++++-- 4 files changed, 37 insertions(+), 8 deletions(-) diff --git a/python/ctranslate2/specs/model_spec.py b/python/ctranslate2/specs/model_spec.py index 41710ff41..2bf5e96c4 100644 --- a/python/ctranslate2/specs/model_spec.py +++ b/python/ctranslate2/specs/model_spec.py @@ -233,6 +233,8 @@ def _quantize(spec, name, value): value = value.reshape(value.shape[0], -1) amax = np.amax(np.absolute(value), axis=1) amax[amax == 0] = 127.0 + # Keep the scale representable in float32 for tiny nonzero rows. + amax = np.maximum(amax, np.float32(127) / np.finfo(np.float32).max) scale = 127.0 / amax value *= np.expand_dims(scale, 1) value = np.rint(value) diff --git a/python/tests/test_spec.py b/python/tests/test_spec.py index f142bcf59..1790cc0d8 100644 --- a/python/tests/test_spec.py +++ b/python/tests/test_spec.py @@ -116,6 +116,21 @@ def __init__(self): ) +def test_int8_quantization_tiny_weights(): + class Spec(ctranslate2.specs.LayerSpec): + def __init__(self): + self.weight = np.array([[1e-38, 0, 0, 0]], dtype=np.float32) + self.weight_scale = OPTIONAL + + spec = Spec() + spec.validate() + with np.errstate(over="raise", invalid="raise"): + spec.optimize(quantization="int8") + assert np.isfinite(spec.weight.numpy()).all() + assert np.isfinite(spec.weight_scale.numpy()).all() + assert (spec.weight_scale.numpy() > 0).all() + + @pytest.mark.parametrize( "quantization,expected_weight,expected_weight_scale,expected_bias", [ diff --git a/src/ops/conv1d_cpu.cc b/src/ops/conv1d_cpu.cc index a33abb602..0cf7891c6 100644 --- a/src/ops/conv1d_cpu.cc +++ b/src/ops/conv1d_cpu.cc @@ -36,7 +36,7 @@ namespace ctranslate2 { const_cast(weight.buffer())); dnnl::memory::dims stride{_stride}; - dnnl::memory::dims dilation{_dilation > 1 ? _dilation : 0}; + dnnl::memory::dims dilation{_dilation - 1}; dnnl::memory::dims padding{_padding}; std::unique_ptr conv_pd; @@ -130,9 +130,6 @@ namespace ctranslate2 { void Conv1D::compute(const StorageView &input, const StorageView &weight, const StorageView *bias, StorageView &output, const StorageView *qscale) const { - if (_dilation != 1) - throw std::runtime_error("Dilation is not supported in this Conv1D implementation"); - compute_with_gemm(input, weight, output, qscale); // Add bias apply_bias_and_activation(output, bias, _activation_type, /*residual=*/nullptr, /*axis=*/-2); @@ -218,6 +215,7 @@ namespace ctranslate2 { const dim_t batch_size = input.dim(0); const dim_t in_channels = input.dim(1); const dim_t input_length = input.dim(2); + const dim_t effective_kernel_size = _dilation * (kernel_size - 1) + 1; auto* out = output.data (); const auto* in = input.data (); dim_t out_offset = 0; @@ -226,11 +224,11 @@ namespace ctranslate2 { const dim_t in_group_stride = in_channels_per_group * input_length; for (dim_t batch_offset = 0; batch_offset < batch_size * in_batch_stride; batch_offset += in_batch_stride) { for (dim_t group_offset = batch_offset; group_offset < (batch_offset + _groups * in_group_stride); group_offset += in_group_stride) { - for (dim_t ti = -_padding; ti <= (input_length - kernel_size + _padding); ti += _stride) { + for (dim_t ti = -_padding; ti <= (input_length - effective_kernel_size + _padding); ti += _stride) { for (dim_t c = group_offset; c < (group_offset + in_channels_per_group * input_length); c += input_length) { for (int k = 0; k < kernel_size; k++) { // Fill items in [0, input_length) range - auto window_i = k + ti; + auto window_i = k * _dilation + ti; if (0 <= window_i && window_i < input_length) { out[out_offset] = in[window_i + c]; } diff --git a/tests/ops_test.cc b/tests/ops_test.cc index 656778a73..5b5541c0c 100644 --- a/tests/ops_test.cc +++ b/tests/ops_test.cc @@ -1339,10 +1339,24 @@ TEST_P(OpDeviceFPTest, Conv1DPaddingAndStride) { expect_storage_eq(output.to_float32(), expected, error); } +TEST(OpTest, Conv1DDilationCPU) { + const StorageView input({1, 1, 7}, std::vector{1, 2, 3, 4, 5, 6, 7}); + const StorageView weight({1, 1, 3}, std::vector{1, 2, 3}); + const StorageView bias({1}, std::vector{1}); + StorageView output; + + ops::Conv1D(1, 0, 2)(input, weight, output); + expect_storage_eq(output, StorageView({1, 1, 3}, std::vector{22, 28, 34})); + + ops::Conv1D(1, 0, 3)(input, weight, output); + expect_storage_eq(output, StorageView({1, 1, 1}, std::vector{30})); + + ops::Conv1D(2, 2, 2)(input, weight, bias, output); + expect_storage_eq(output, StorageView({1, 1, 4}, std::vector{12, 23, 35, 20})); +} + TEST_P(OpDeviceFPTest, Conv1DDilation) { const Device device = GetParam().device; - if (device == Device::CPU) - GTEST_SKIP() << "Dilated convolution is not implemented for CPU."; const DataType dtype = GetParam().dtype; const float error = GetParam().error; const StorageView expected({2, 4, 1}, std::vector{