diff --git a/src/generators.cpp b/src/generators.cpp index 388e63501e..cd7eed5497 100644 --- a/src/generators.cpp +++ b/src/generators.cpp @@ -380,8 +380,20 @@ Generator::Generator(const Model& model, const GeneratorParams& params) : model_ throw std::runtime_error("search max_length is 0"); if (params.search.max_length > model.config_->model.context_length) throw std::runtime_error("max_length (" + std::to_string(params.search.max_length) + ") cannot be greater than model context_length (" + std::to_string(model.config_->model.context_length) + ")"); - if (params.search.batch_size < 1) - throw std::runtime_error("batch_size must be 1 or greater, is " + std::to_string(params.search.batch_size)); + + constexpr int kMaxBatchSize = 32; + constexpr int kMaxNumBeams = 32; + constexpr int kMaxNumBeamsCuda = 32; + + if (params.search.batch_size < 1 || params.search.batch_size > kMaxBatchSize) + throw std::runtime_error("batch_size (" + std::to_string(params.search.batch_size) + ") must be in [1, " + std::to_string(kMaxBatchSize) + "]"); + + const int max_num_beams = (params.search.num_beams > 1 && + (params.p_device->GetType() == DeviceType::CUDA || params.p_device->GetType() == DeviceType::NvTensorRtRtx)) + ? kMaxNumBeamsCuda + : kMaxNumBeams; + if (params.search.num_beams < 1 || params.search.num_beams > max_num_beams) + throw std::runtime_error("num_beams (" + std::to_string(params.search.num_beams) + ") must be in [1, " + std::to_string(max_num_beams) + "]"); if (params.config.model.vocab_size < 1) throw std::runtime_error("vocab_size must be 1 or greater, is " + std::to_string(params.config.model.vocab_size)); diff --git a/test/model_tests.cpp b/test/model_tests.cpp index a96cd7a7d0..42ebbd53eb 100644 --- a/test/model_tests.cpp +++ b/test/model_tests.cpp @@ -482,6 +482,46 @@ Print all primes between 1 and n } #endif +// Validation tests for search parameter bounds +#if !USE_DML +TEST(ModelTests, NumBeamsUpperBoundThrows) { + auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32"); + auto params = OgaGeneratorParams::Create(*model); + params->SetSearchOption("max_length", 20); + params->SetSearchOption("batch_size", 1); + params->SetSearchOption("num_beams", 512); // exceeds upper bound of 256 + + EXPECT_THROW(OgaGenerator::Create(*model, *params), std::runtime_error); +} + +TEST(ModelTests, BatchSizeUpperBoundThrows) { + auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32"); + auto params = OgaGeneratorParams::Create(*model); + params->SetSearchOption("max_length", 20); + params->SetSearchOption("batch_size", 512); // exceeds upper bound of 256 + + EXPECT_THROW(OgaGenerator::Create(*model, *params), std::runtime_error); +} + +TEST(ModelTests, NumBeamsZeroThrows) { + auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32"); + auto params = OgaGeneratorParams::Create(*model); + params->SetSearchOption("max_length", 20); + params->SetSearchOption("batch_size", 1); + params->SetSearchOption("num_beams", 0); // below lower bound of 1 + + EXPECT_THROW(OgaGenerator::Create(*model, *params), std::runtime_error); +} + +TEST(ModelTests, BatchSizeZeroThrows) { + auto model = OgaModel::Create(MODEL_PATH "hf-internal-testing/tiny-random-gpt2-fp32"); + auto params = OgaGeneratorParams::Create(*model); + params->SetSearchOption("max_length", 20); + params->SetSearchOption("batch_size", 0); // below lower bound of 1 + + EXPECT_THROW(OgaGenerator::Create(*model, *params), std::runtime_error); +} +#endif // --- Validation tests (no model files required) --- TEST(ValidationTests, WindowIndexAcceptsValidParams) { @@ -520,4 +560,4 @@ TEST(ValidationTests, WindowIndexRejectsExcessiveDimensions) { TEST(ValidationTests, WindowIndexRejectsTotalSizeOverflow) { // Each dim individually <= kMaxElements, but grid_t * padded_h * padded_w > kMaxElements EXPECT_THROW(Generators::ValidateWindowIndexParams(1000, 2000000, 2000000, 2, 14, 112), std::runtime_error); -} \ No newline at end of file +}