From 6622a2e162d9df7f6f48c44cedffa730cfde9a2e Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Mon, 2 Mar 2026 16:07:56 -0800 Subject: [PATCH 1/4] Detect and test mismatch between raw data size and declared data type and shape --- onnxruntime/lora/adapter_format_utils.cc | 14 ++- onnxruntime/test/lora/lora_test.cc | 143 +++++++++++++++++++++++ 2 files changed, 154 insertions(+), 3 deletions(-) diff --git a/onnxruntime/lora/adapter_format_utils.cc b/onnxruntime/lora/adapter_format_utils.cc index 2d061b6066a8a..f4b0988a9a505 100644 --- a/onnxruntime/lora/adapter_format_utils.cc +++ b/onnxruntime/lora/adapter_format_utils.cc @@ -7,8 +7,9 @@ #include "core/framework/allocator.h" #include "core/common/common.h" #include "core/common/endian.h" -#include "core/framework/endian_utils.h" +#include "core/common/safeint.h" #include "core/common/span_utils.h" +#include "core/framework/endian_utils.h" #include "core/framework/ortdevice.h" #include "core/framework/ortmemoryinfo.h" #include "core/framework/ort_value.h" @@ -155,7 +156,14 @@ std::pair CreateOrtValueOverLoraParameter(const Parameter const auto data_type = param.data_type(); // Copying shape takes care of endianess using flatbuffers accessors TensorShapeVector shape(param.dims()->begin(), param.dims()->end()); + TensorShape tensor_shape(shape); const auto elem_type = DataTypeImpl::TensorTypeFromONNXEnum(static_cast(data_type))->GetElementType(); + const size_t expected_raw_data_size = SafeInt(tensor_shape.Size()) * elem_type->Size(); + if (param.raw_data()->size() != expected_raw_data_size) { + ORT_THROW("Lora Param:", param.name(), + "Raw data size does not match the expected size calculated from tensor shape and element type"); + } + static const OrtMemoryInfo cpu_meminfo(CPU, OrtAllocatorType::OrtDeviceAllocator); if constexpr (endian::native == endian::big) { @@ -166,7 +174,7 @@ std::pair CreateOrtValueOverLoraParameter(const Parameter // of raw data // const_cast is necessary due to Tensor class API Tensor::InitOrtValue(elem_type, - TensorShape(shape), + tensor_shape, const_cast(param.raw_data()->data()), cpu_meminfo, result); @@ -174,7 +182,7 @@ std::pair CreateOrtValueOverLoraParameter(const Parameter } else { // const_cast is necessary due to Tensor class API Tensor::InitOrtValue(elem_type, - TensorShape(shape), + tensor_shape, const_cast(param.raw_data()->data()), cpu_meminfo, result); diff --git a/onnxruntime/test/lora/lora_test.cc b/onnxruntime/test/lora/lora_test.cc index 0c55cf45abcdf..e5cdc583981e4 100644 --- a/onnxruntime/test/lora/lora_test.cc +++ b/onnxruntime/test/lora/lora_test.cc @@ -199,6 +199,149 @@ TEST(LoraAdapterTest, Load) { } } +TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_ValidParam) { + // Build a valid adapter with a single float parameter, then call + // CreateOrtValueOverLoraParameter on the deserialized Parameter. + constexpr std::array shape = {8, 4}; + InlinedVector data(32); + std::iota(data.begin(), data.end(), 0.f); + + adapters::utils::AdapterFormatBuilder adapter_builder; + adapter_builder.AddParameter("valid_param", adapters::TensorDataType::FLOAT, + shape, ReinterpretAsSpan(gsl::make_span(data))); + + auto buffer = adapter_builder.Finish(kAdapterVersion, kModelVersion); + + const auto* adapter = adapters::utils::ValidateAndGetAdapterFromBytes(buffer); + ASSERT_NE(adapter, nullptr); + ASSERT_NE(adapter->parameters(), nullptr); + ASSERT_EQ(adapter->parameters()->size(), 1u); + + const auto* param = adapter->parameters()->Get(0); + auto [name, ort_value] = adapters::utils::CreateOrtValueOverLoraParameter(*param); + + ASSERT_EQ(name, "valid_param"); + ASSERT_TRUE(ort_value.IsTensor()); + + const auto& tensor = ort_value.Get(); + ASSERT_EQ(tensor.GetElementType(), ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + + auto dims = tensor.Shape().GetDims(); + ASSERT_EQ(dims.size(), 2u); + ASSERT_EQ(dims[0], 8); + ASSERT_EQ(dims[1], 4); + + auto result_span = tensor.DataAsSpan(); + ASSERT_EQ(result_span.size(), 32u); + for (size_t i = 0; i < result_span.size(); ++i) { + ASSERT_EQ(static_cast(i), result_span[i]); + } +} + +TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_RawDataSizeMismatch) { + // Craft a flatbuffer Parameter where raw_data has fewer bytes than + // shape (8 x 4) * sizeof(float) = 128 bytes. + // We supply only 64 bytes (half the expected amount) so the validation + // inside CreateOrtValueOverLoraParameter must throw. + flatbuffers::FlatBufferBuilder fbb; + + auto name_offset = fbb.CreateString("bad_param"); + std::vector dims = {8, 4}; + auto dims_offset = fbb.CreateVector(dims); + + // 8 * 4 floats = 32 elements = 128 bytes expected. + // Provide only 64 bytes (16 floats worth) to trigger the mismatch. + std::vector short_data(64, 0); + fbb.ForceVectorAlignment(short_data.size(), sizeof(uint8_t), 8); + auto data_offset = fbb.CreateVector(short_data); + + auto param_offset = adapters::CreateParameter( + fbb, name_offset, dims_offset, adapters::TensorDataType::FLOAT, data_offset); + + // Wrap the single parameter inside an Adapter so the buffer is valid flatbuffers. + auto params_offset = fbb.CreateVector(¶m_offset, 1); + auto adapter_offset = adapters::CreateAdapter( + fbb, adapters::kAdapterFormatVersion, kAdapterVersion, kModelVersion, params_offset); + adapters::FinishAdapterBuffer(fbb, adapter_offset); + + auto* buf = fbb.GetBufferPointer(); + auto size = fbb.GetSize(); + + // Retrieve the Parameter from the Adapter + const auto* adapter = adapters::GetAdapter(buf); + ASSERT_NE(adapter, nullptr); + ASSERT_NE(adapter->parameters(), nullptr); + ASSERT_EQ(adapter->parameters()->size(), 1u); + + const auto* param = adapter->parameters()->Get(0); + ASSERT_NE(param, nullptr); + + // The raw_data is 64 bytes but shape says 8x4 floats = 128 bytes. + // CreateOrtValueOverLoraParameter must throw. + ASSERT_THROW(adapters::utils::CreateOrtValueOverLoraParameter(*param), OnnxRuntimeException); +} + +TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_ExcessRawData) { + // Craft a flatbuffer Parameter where raw_data has MORE bytes than expected. + // Shape (2, 2) with float => 4 elements => 16 bytes expected, but we supply 32. + flatbuffers::FlatBufferBuilder fbb; + + auto name_offset = fbb.CreateString("excess_param"); + std::vector dims = {2, 2}; + auto dims_offset = fbb.CreateVector(dims); + + // 2 * 2 floats = 4 elements = 16 bytes expected. Supply 32. + std::vector excess_data(32, 0); + fbb.ForceVectorAlignment(excess_data.size(), sizeof(uint8_t), 8); + auto data_offset = fbb.CreateVector(excess_data); + + auto param_offset = adapters::CreateParameter( + fbb, name_offset, dims_offset, adapters::TensorDataType::FLOAT, data_offset); + + auto params_offset = fbb.CreateVector(¶m_offset, 1); + auto adapter_offset = adapters::CreateAdapter( + fbb, adapters::kAdapterFormatVersion, kAdapterVersion, kModelVersion, params_offset); + adapters::FinishAdapterBuffer(fbb, adapter_offset); + + const auto* adapter = adapters::GetAdapter(fbb.GetBufferPointer()); + ASSERT_NE(adapter, nullptr); + + const auto* param = adapter->parameters()->Get(0); + ASSERT_NE(param, nullptr); + + // Excess data should also trigger the mismatch throw. + ASSERT_THROW(adapters::utils::CreateOrtValueOverLoraParameter(*param), OnnxRuntimeException); +} + +TEST(LoraAdapterTest, Load_RawDataSizeMismatch) { + // End-to-end: loading an adapter whose parameter has mismatched raw data + // should fail during LoraAdapter::Load. + flatbuffers::FlatBufferBuilder fbb; + + auto name_offset = fbb.CreateString("bad_param"); + std::vector dims = {8, 4}; + auto dims_offset = fbb.CreateVector(dims); + + // Provide 64 bytes instead of the expected 128 for float [8, 4]. + std::vector short_data(64, 0); + fbb.ForceVectorAlignment(short_data.size(), sizeof(uint8_t), 8); + auto data_offset = fbb.CreateVector(short_data); + + auto param_offset = adapters::CreateParameter( + fbb, name_offset, dims_offset, adapters::TensorDataType::FLOAT, data_offset); + + auto params_offset = fbb.CreateVector(¶m_offset, 1); + auto adapter_offset = adapters::CreateAdapter( + fbb, adapters::kAdapterFormatVersion, kAdapterVersion, kModelVersion, params_offset); + adapters::FinishAdapterBuffer(fbb, adapter_offset); + + std::vector buffer(fbb.GetBufferPointer(), + fbb.GetBufferPointer() + fbb.GetSize()); + + lora::LoraAdapter adapter; + ASSERT_THROW(adapter.Load(std::move(buffer)), OnnxRuntimeException); +} + #ifdef USE_CUDA TEST(LoraAdapterTest, VerifyDeviceCopy) { auto cpu_ep = DefaultCpuExecutionProvider(); From aeee8d549f3bafe7e0efb560f6534999496b10d8 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Mon, 2 Mar 2026 17:04:27 -0800 Subject: [PATCH 2/4] Update onnxruntime/lora/adapter_format_utils.cc Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxruntime/lora/adapter_format_utils.cc | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/onnxruntime/lora/adapter_format_utils.cc b/onnxruntime/lora/adapter_format_utils.cc index f4b0988a9a505..c3349bde19429 100644 --- a/onnxruntime/lora/adapter_format_utils.cc +++ b/onnxruntime/lora/adapter_format_utils.cc @@ -160,8 +160,10 @@ std::pair CreateOrtValueOverLoraParameter(const Parameter const auto elem_type = DataTypeImpl::TensorTypeFromONNXEnum(static_cast(data_type))->GetElementType(); const size_t expected_raw_data_size = SafeInt(tensor_shape.Size()) * elem_type->Size(); if (param.raw_data()->size() != expected_raw_data_size) { - ORT_THROW("Lora Param:", param.name(), - "Raw data size does not match the expected size calculated from tensor shape and element type"); + ORT_THROW("Lora Param '", name, + "': raw_data size (", param.raw_data()->size(), + ") does not match expected size (", expected_raw_data_size, + ") calculated from tensor shape and element type"); } static const OrtMemoryInfo cpu_meminfo(CPU, OrtAllocatorType::OrtDeviceAllocator); From ff301036172be18855f0b2a999411823654e59ce Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Mon, 2 Mar 2026 17:04:46 -0800 Subject: [PATCH 3/4] Update onnxruntime/test/lora/lora_test.cc Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- onnxruntime/test/lora/lora_test.cc | 1 - 1 file changed, 1 deletion(-) diff --git a/onnxruntime/test/lora/lora_test.cc b/onnxruntime/test/lora/lora_test.cc index e5cdc583981e4..cbb54bf91a13a 100644 --- a/onnxruntime/test/lora/lora_test.cc +++ b/onnxruntime/test/lora/lora_test.cc @@ -265,7 +265,6 @@ TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_RawDataSizeMismatch) { adapters::FinishAdapterBuffer(fbb, adapter_offset); auto* buf = fbb.GetBufferPointer(); - auto size = fbb.GetSize(); // Retrieve the Parameter from the Adapter const auto* adapter = adapters::GetAdapter(buf); From 4a4d308bda4487aaab23d994efcf2be016677e42 Mon Sep 17 00:00:00 2001 From: Dmitri Smirnov Date: Mon, 2 Mar 2026 18:05:37 -0800 Subject: [PATCH 4/4] Add validation for dims and name --- onnxruntime/lora/adapter_format_utils.cc | 26 ++++-- onnxruntime/test/lora/lora_test.cc | 101 +++++++++++++++++++++++ 2 files changed, 121 insertions(+), 6 deletions(-) diff --git a/onnxruntime/lora/adapter_format_utils.cc b/onnxruntime/lora/adapter_format_utils.cc index c3349bde19429..6e2d204b04cf2 100644 --- a/onnxruntime/lora/adapter_format_utils.cc +++ b/onnxruntime/lora/adapter_format_utils.cc @@ -150,18 +150,32 @@ struct ReadDataForBigEndian { std::pair CreateOrtValueOverLoraParameter(const Parameter& param) { OrtValue result; + const auto* param_name = param.name(); + ORT_ENFORCE(param_name != nullptr, "Lora Parameter: name is missing"); + std::string name; - LoadStringFromLoraFormat(name, param.name()); + LoadStringFromLoraFormat(name, param_name); const auto data_type = param.data_type(); + ORT_ENFORCE(data_type != TensorDataType::UNDEFINED, + "Lora Param '", name, "': data_type is UNDEFINED"); + + const auto* dims = param.dims(); + ORT_ENFORCE(dims != nullptr && dims->size() > 0, + "Lora Param '", name, "': dims is missing or empty"); + + const auto* raw_data = param.raw_data(); + ORT_ENFORCE(raw_data != nullptr, + "Lora Param '", name, "': raw_data is missing"); + // Copying shape takes care of endianess using flatbuffers accessors - TensorShapeVector shape(param.dims()->begin(), param.dims()->end()); + TensorShapeVector shape(dims->begin(), dims->end()); TensorShape tensor_shape(shape); const auto elem_type = DataTypeImpl::TensorTypeFromONNXEnum(static_cast(data_type))->GetElementType(); const size_t expected_raw_data_size = SafeInt(tensor_shape.Size()) * elem_type->Size(); - if (param.raw_data()->size() != expected_raw_data_size) { + if (raw_data->size() != expected_raw_data_size) { ORT_THROW("Lora Param '", name, - "': raw_data size (", param.raw_data()->size(), + "': raw_data size (", raw_data->size(), ") does not match expected size (", expected_raw_data_size, ") calculated from tensor shape and element type"); } @@ -177,7 +191,7 @@ std::pair CreateOrtValueOverLoraParameter(const Parameter // const_cast is necessary due to Tensor class API Tensor::InitOrtValue(elem_type, tensor_shape, - const_cast(param.raw_data()->data()), + const_cast(raw_data->data()), cpu_meminfo, result); } @@ -185,7 +199,7 @@ std::pair CreateOrtValueOverLoraParameter(const Parameter // const_cast is necessary due to Tensor class API Tensor::InitOrtValue(elem_type, tensor_shape, - const_cast(param.raw_data()->data()), + const_cast(raw_data->data()), cpu_meminfo, result); } diff --git a/onnxruntime/test/lora/lora_test.cc b/onnxruntime/test/lora/lora_test.cc index cbb54bf91a13a..7455484be28ec 100644 --- a/onnxruntime/test/lora/lora_test.cc +++ b/onnxruntime/test/lora/lora_test.cc @@ -174,6 +174,20 @@ struct TestDataType { } }; +// Helper that wraps a single Parameter offset into a finished Adapter flatbuffer +// and returns a pointer to the deserialized Parameter. +// The FlatBufferBuilder must outlive the returned pointer. +const adapters::Parameter* BuildAdapterAndGetParam(flatbuffers::FlatBufferBuilder& fbb, + flatbuffers::Offset param_offset) { + auto params_offset = fbb.CreateVector(¶m_offset, 1); + auto adapter_offset = adapters::CreateAdapter( + fbb, adapters::kAdapterFormatVersion, kAdapterVersion, kModelVersion, params_offset); + adapters::FinishAdapterBuffer(fbb, adapter_offset); + + const auto* adapter = adapters::GetAdapter(fbb.GetBufferPointer()); + return adapter->parameters()->Get(0); +} + } // namespace TEST(LoraAdapterTest, Load) { @@ -341,6 +355,93 @@ TEST(LoraAdapterTest, Load_RawDataSizeMismatch) { ASSERT_THROW(adapter.Load(std::move(buffer)), OnnxRuntimeException); } +TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_MissingName) { + // Parameter with null name should throw gracefully. + flatbuffers::FlatBufferBuilder fbb; + + std::vector dims = {2, 2}; + std::vector raw_data(16, 0); // 2*2 floats = 16 bytes + + // name is nullptr, all other fields are valid + auto param_offset = adapters::CreateParameterDirect( + fbb, /*name=*/nullptr, &dims, adapters::TensorDataType::FLOAT, &raw_data); + + const auto* param = BuildAdapterAndGetParam(fbb, param_offset); + ASSERT_NE(param, nullptr); + ASSERT_EQ(param->name(), nullptr); + + ASSERT_THROW(adapters::utils::CreateOrtValueOverLoraParameter(*param), OnnxRuntimeException); +} + +TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_MissingDims) { + // Parameter with null dims should throw gracefully. + flatbuffers::FlatBufferBuilder fbb; + + std::vector raw_data(16, 0); + + // dims is nullptr + auto param_offset = adapters::CreateParameterDirect( + fbb, "no_dims_param", /*dims=*/nullptr, adapters::TensorDataType::FLOAT, &raw_data); + + const auto* param = BuildAdapterAndGetParam(fbb, param_offset); + ASSERT_NE(param, nullptr); + ASSERT_EQ(param->dims(), nullptr); + + ASSERT_THROW(adapters::utils::CreateOrtValueOverLoraParameter(*param), OnnxRuntimeException); +} + +TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_EmptyDims) { + // Parameter with an empty dims vector should throw gracefully. + flatbuffers::FlatBufferBuilder fbb; + + std::vector empty_dims; + std::vector raw_data(16, 0); + + auto param_offset = adapters::CreateParameterDirect( + fbb, "empty_dims_param", &empty_dims, adapters::TensorDataType::FLOAT, &raw_data); + + const auto* param = BuildAdapterAndGetParam(fbb, param_offset); + ASSERT_NE(param, nullptr); + ASSERT_EQ(param->dims()->size(), 0u); + + ASSERT_THROW(adapters::utils::CreateOrtValueOverLoraParameter(*param), OnnxRuntimeException); +} + +TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_MissingRawData) { + // Parameter with null raw_data should throw gracefully. + flatbuffers::FlatBufferBuilder fbb; + + std::vector dims = {2, 2}; + + // raw_data is nullptr + auto param_offset = adapters::CreateParameterDirect( + fbb, "no_data_param", &dims, adapters::TensorDataType::FLOAT, /*raw_data=*/nullptr); + + const auto* param = BuildAdapterAndGetParam(fbb, param_offset); + ASSERT_NE(param, nullptr); + ASSERT_EQ(param->raw_data(), nullptr); + + ASSERT_THROW(adapters::utils::CreateOrtValueOverLoraParameter(*param), OnnxRuntimeException); +} + +TEST(LoraAdapterTest, CreateOrtValueOverLoraParameter_UndefinedDataType) { + // Parameter with UNDEFINED data_type should throw gracefully. + flatbuffers::FlatBufferBuilder fbb; + + std::vector dims = {2, 2}; + std::vector raw_data(16, 0); + + // data_type defaults to UNDEFINED when not set + auto param_offset = adapters::CreateParameterDirect( + fbb, "undef_type_param", &dims, adapters::TensorDataType::UNDEFINED, &raw_data); + + const auto* param = BuildAdapterAndGetParam(fbb, param_offset); + ASSERT_NE(param, nullptr); + ASSERT_EQ(param->data_type(), adapters::TensorDataType::UNDEFINED); + + ASSERT_THROW(adapters::utils::CreateOrtValueOverLoraParameter(*param), OnnxRuntimeException); +} + #ifdef USE_CUDA TEST(LoraAdapterTest, VerifyDeviceCopy) { auto cpu_ep = DefaultCpuExecutionProvider();