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/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/src/storage_view.cc b/src/storage_view.cc index 8f5fe33d4..027def049 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/ops_test.cc b/tests/ops_test.cc index 9aa6f42e4..89af2b39d 100644 --- a/tests/ops_test.cc +++ b/tests/ops_test.cc @@ -1452,10 +1452,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{ diff --git a/tests/storage_view_test.cc b/tests/storage_view_test.cc index 80ffe780d..2c3e3317b 100644 --- a/tests/storage_view_test.cc +++ b/tests/storage_view_test.cc @@ -148,3 +148,114 @@ INSTANTIATE_TEST_SUITE_P(CUDA, StorageViewDeviceTest, ::testing::Values(Device:: #ifdef CT2_WITH_MPS INSTANTIATE_TEST_SUITE_P(MPS, StorageViewDeviceTest, ::testing::Values(Device::MPS)); #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); +}