From 008662c874ffc81d59747bc32ad301d62acf84cb Mon Sep 17 00:00:00 2001 From: Roshan Sharma Date: Sat, 22 Aug 2026 18:07:02 -0400 Subject: [PATCH] Validate GGUF tensor dimensions --- mlx/io/gguf.cpp | 37 +++++++++++----- mlx/io/gguf_quants.cpp | 13 ++++-- tests/load_tests.cpp | 98 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 135 insertions(+), 13 deletions(-) diff --git a/mlx/io/gguf.cpp b/mlx/io/gguf.cpp index 6c27c46987..eab225da75 100644 --- a/mlx/io/gguf.cpp +++ b/mlx/io/gguf.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include "mlx/io/gguf.h" @@ -47,11 +48,33 @@ std::optional gguf_type_to_dtype(const uint32_t& gguf_type) { } } +[[noreturn]] void tensor_error( + const gguf_tensor& tensor, + const std::string& what) { + std::ostringstream msg; + msg << "[load_gguf] Tensor '" << std::string(tensor.name, tensor.namelen) + << "' " << what << ". Perhaps an incomplete download or corrupt file?"; + throw std::runtime_error(msg.str()); +} + Shape get_shape(const gguf_tensor& tensor) { Shape shape; + // Only the byte size is checked against the file, so a dimension that does + // not survive the narrowing to ShapeElem, or a product that wraps, would let + // the shape and tensor.num_weights describe different sizes. + uint64_t num_weights = 1; // The dimension order in GGML is the reverse of the order used in MLX. for (int i = tensor.ndim - 1; i >= 0; i--) { - shape.push_back(tensor.dim[i]); + auto dim = tensor.dim[i]; + if (dim > std::numeric_limits::max()) { + tensor_error(tensor, "has a dimension that is too large"); + } + if (dim != 0 && num_weights > std::numeric_limits::max() / dim) { + tensor_error( + tensor, "has a dimension product that does not fit in 64 bits"); + } + num_weights *= dim; + shape.push_back(dim); } return shape; } @@ -303,20 +326,14 @@ std::unordered_map load_metadata(gguf_ctx* ctx) { } void check_tensor_in_file(const gguf_ctx* ctx, const gguf_tensor& tensor) { - auto fail = [&tensor](const std::string& what) { - std::ostringstream msg; - msg << "[load_gguf] Tensor '" << std::string(tensor.name, tensor.namelen) - << "' " << what << ". Perhaps an incomplete download or corrupt file?"; - throw std::runtime_error(msg.str()); - }; if (tensor.offset < ctx->data_off) { - fail("has a data offset that overflows the data section"); + tensor_error(tensor, "has a data offset that overflows the data section"); } if (tensor.offset > ctx->size) { - fail("has a data offset past the end of the file"); + tensor_error(tensor, "has a data offset past the end of the file"); } if (tensor.bsize > ctx->size - tensor.offset) { - fail("extends past the end of the file"); + tensor_error(tensor, "extends past the end of the file"); } } diff --git a/mlx/io/gguf_quants.cpp b/mlx/io/gguf_quants.cpp index 83dfe1f822..1717be5601 100644 --- a/mlx/io/gguf_quants.cpp +++ b/mlx/io/gguf_quants.cpp @@ -112,6 +112,11 @@ void gguf_load_quantized( std::string name(tensor.name, tensor.namelen); auto shape = get_shape(tensor); + if (shape.empty()) { + std::ostringstream msg; + msg << "[load_gguf] quantized tensor " << name << " has no dimensions"; + throw std::runtime_error(msg.str()); + } const uint64_t weights_per_block = 32; if (shape[shape.size() - 1] % weights_per_block != 0) { std::ostringstream msg; @@ -125,15 +130,17 @@ void gguf_load_quantized( auto w_nbytes = uint32.size() * std::accumulate(weights_shape.begin(), weights_shape.end(), - 1, + size_t{1}, std::multiplies()); array weights(allocator::malloc(w_nbytes), std::move(weights_shape), uint32); // For scales and bias shape[shape.size() - 1] = shape[shape.size() - 1] / weights_per_block; - auto sb_nbytes = float16.size() * - std::accumulate(shape.begin(), shape.end(), 1, std::multiplies()); + auto sb_nbytes = + float16.size() * + std::accumulate( + shape.begin(), shape.end(), size_t{1}, std::multiplies()); array scales(allocator::malloc(sb_nbytes), shape, float16); array biases(allocator::malloc(sb_nbytes), std::move(shape), float16); diff --git a/tests/load_tests.cpp b/tests/load_tests.cpp index 6ef7bc276e..b4b54046a3 100644 --- a/tests/load_tests.cpp +++ b/tests/load_tests.cpp @@ -257,6 +257,104 @@ TEST_CASE("test gguf tensor data offset validation") { } } +// Writes a one-tensor GGUF whose tensor header carries `dims` and `type` +// verbatim, so a test can drive the dimension product independently of the +// amount of tensor data in the file. The name ends in ".weight" because the +// quantized loader derives the scales and biases names from it. +void write_raw_gguf_dims( + const std::string& path, + const std::vector& dims, + uint32_t type, + size_t data_bytes) { + std::ofstream out(path, std::ios::binary); + auto u32 = [&out](uint32_t v) { + out.write(reinterpret_cast(&v), 4); + }; + auto u64 = [&out](uint64_t v) { + out.write(reinterpret_cast(&v), 8); + }; + out.write("GGUF", 4); + u32(3); // version + u64(1); // tensor_count + u64(0); // metadata_kv_count + u64(8); // tensor name length + out.write("w.weight", 8); + u32(dims.size()); + for (auto dim : dims) { + u64(dim); + } + u32(type); + u64(0); // tensor data offset + while (out.tellp() % 32 != 0) { // default GGUF alignment + out.put(0); + } + std::vector data(data_bytes, 0); + out.write(data.data(), data.size()); +} + +TEST_CASE("test gguf tensor dimension validation") { + // The element count exists in two widths: gguflib keeps the 64 bit product in + // tensor.num_weights, and get_shape narrows each dimension to a 32 bit + // ShapeElem. Only the byte size is checked against the file, so dimensions + // that make the two disagree size a buffer from one and index it with the + // other. See ml-explore/mlx#4244. + const uint32_t q8_0 = 8; + + SUBCASE("valid quantized tensor loads") { + std::string file_path = get_temp_file("test_gguf_dims_ok.gguf"); + // One block of 32 weights: 2 bytes of scale plus 32 bytes of weights. + write_raw_gguf_dims(file_path, {32}, q8_0, 34); + auto [weights, metadata] = load_gguf(file_path); + CHECK_EQ(weights.at("w.weight").shape(), Shape{8}); + CHECK_EQ(weights.at("w.scales").shape(), Shape{1}); + CHECK_EQ(weights.at("w.biases").shape(), Shape{1}); + } + + SUBCASE("dimension past the ShapeElem range") { + // Narrows to a shape of 32 that no longer describes the file. + std::string file_path = get_temp_file("test_gguf_dim_narrowed.gguf"); + write_raw_gguf_dims(file_path, {(1ull << 32) + 32}, q8_0, 34); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("dimension that narrows to a negative shape") { + std::string file_path = get_temp_file("test_gguf_dim_negative.gguf"); + write_raw_gguf_dims(file_path, {1ull << 31}, q8_0, 34); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("dimension product that wraps to a small byte size") { + // 96 * 384307168202282326 wraps to 64, so the tensor needs only 68 bytes of + // data and passes the byte size check, while the narrowed shape + // {1431655766, 96} still describes 2^37 elements. That drives the scales + // element count past INT32_MAX, which is what truncates the allocation size + // while the extractor still loops over the full count. + std::string file_path = get_temp_file("test_gguf_dim_wrap.gguf"); + write_raw_gguf_dims(file_path, {96, 384307168202282326ull}, q8_0, 68); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("dimension product that wraps to zero") { + // Wraps num_weights to exactly 0, so the tensor claims no data at all. + std::string file_path = get_temp_file("test_gguf_dim_wrap_zero.gguf"); + write_raw_gguf_dims( + file_path, + {(1ull << 32) + (1ull << 25), + (1ull << 32) + (1ull << 15), + (1ull << 32) + (1ull << 24)}, + q8_0, + 0); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } + + SUBCASE("tensor without dimensions") { + // gguf_load_quantized indexes the last shape element unconditionally. + std::string file_path = get_temp_file("test_gguf_no_dims.gguf"); + write_raw_gguf_dims(file_path, {}, q8_0, 0); + CHECK_THROWS_AS(load_gguf(file_path), std::runtime_error); + } +} + // Writes a metadata-only GGUF (no tensors) whose metadata KV section is // `kv_section` verbatim, so a caller can encode values whose lengths exceed the // file to exercise check_metadata_value_in_file(). `kv_count` must match the