diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 763154972..36703127a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -240,6 +240,26 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) FIXTURES_SETUP generate-models ) + # Target model for test-speculative-rejection-sampling + set(SPEC_TARGET_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-speculative-models/target") + file(MAKE_DIRECTORY "${SPEC_TARGET_DIR}") + llama_test( + test-llama-archs + NAME test-generate-speculative-target + LABEL main + ARGS -a llama -s 999 -o "${SPEC_TARGET_DIR}" + ) + set_tests_properties( + test-generate-speculative-target + PROPERTIES FIXTURES_SETUP speculative-model + ) + + # To avoid recreating the model use the fixture exclude any option: + # -FA "^speculative-model$" + llama_build_and_test(test-speculative-rejection-sampling.cpp ARGS + -m "${SPEC_TARGET_DIR}/llama-dense.gguf") + set_tests_properties(test-speculative-rejection-sampling PROPERTIES FIXTURES_REQUIRED speculative-model) + # Test context state handling across all architectures, using the generated dummy models llama_test( test-llama-context diff --git a/tests/test-speculative-rejection-sampling.cpp b/tests/test-speculative-rejection-sampling.cpp new file mode 100644 index 000000000..523436028 --- /dev/null +++ b/tests/test-speculative-rejection-sampling.cpp @@ -0,0 +1,229 @@ +#include "sampling.h" +#include "llama.h" +#include "common.h" +#include "arg.h" +#include "log.h" + +#include +#include +#include +#include +#include +#include + +// Tests that common_sampler_sample_and_accept_n_rejection preserves the target +// distribution regardless of what the draft distribution Q looks like. + +static int test_speculative_rejection_sampling(common_params & params, int n_trials) { + params.sampling.backend_sampling = false; + auto llama_init = common_init_from_params(params); + + auto * model_tgt = llama_init->model(); + auto * ctx_tgt = llama_init->context(); + if (!model_tgt || !ctx_tgt) { + LOG_ERR("failed to initialize target model and context\n"); + return 1; + } + + const llama_tokens prompt = { 1, 2, 3, 1, 2 }; + + // Process the prompt with the target model to get its true distribution for + // the position right after the prompt. + common_batch batch = common_batch_get_one(ctx_tgt, prompt); + if (llama_process(ctx_tgt, LLAMA_PROCESS_TYPE_DECODE, batch.get()) != 0) { + LOG_ERR("failed to decode the test prompt\n"); + return 1; + } + + common_params_sampling sampling; + sampling.seed = 18; + sampling.samplers = { COMMON_SAMPLER_TYPE_TEMPERATURE }; + sampling.temp = 1.0f; + + // Initialize the sampler for the target model. + common_sampler_ptr smpl(common_sampler_init(model_tgt, sampling)); + if (!smpl) { + LOG_ERR("failed to initialize target sampler\n"); + return 1; + } + + // Sample once so the sampler chain runs and populates cur_p. We don't need + // the returned token id itself so the return value is ignored. + common_sampler_sample(smpl.get(), ctx_tgt, -1); + + // Get the target candidates distribution in sorted order. + const llama_token_data_array * dist_tgt = common_sampler_get_candidates(smpl.get(), true); + + // Clone smpl before the trial loop starts. We need the dist_tgt pointer to stay valid. + common_sampler_ptr smpl_verify(common_sampler_clone(smpl.get())); + if (!smpl_verify) { + LOG_ERR("failed to clone target sampler\n"); + return 1; + } + + const llama_vocab * vocab = llama_model_get_vocab(model_tgt); + const int n_vocab = llama_vocab_n_tokens(vocab); + + std::mt19937 rng(18); + std::uniform_real_distribution uni(0.0f, 1.0f); + + std::unordered_map token_counts; + + std::vector q(n_vocab); + for (int i = 0; i < n_vocab; ++i) { + q[i].id = i; + q[i].logit = 0.0f; + } + + for (int trial = 0; trial < n_trials; ++trial) { + float q_sum = 0.0f; + + for (auto & td : q) { + td.p = uni(rng); + q_sum += td.p; + } + + const float inv_sum = 1.0f / q_sum; + for (auto & td : q) { + td.p *= inv_sum; + } + + float u = uni(rng); + + // Fallback to the last bucket to absorb residual probability from FP precision loss + llama_token draft_tok = q.back().id; + for (const auto & e : q) { + u -= e.p; + if (u <= 0.0f) { + draft_tok = e.id; + break; + } + } + + // Both indices reuse the prompt's last output logits. The second entry + // is the bonus draw slot; it is always stale and its result is ignored. + const llama_tokens draft = { draft_tok }; + const std::vector> draft_q = { q }; + const std::vector idxs = { -1, -1 }; + + const std::vector result = + common_sampler_sample_and_accept_n_rejection(smpl_verify.get(), ctx_tgt, idxs, draft, draft_q, false); + + token_counts[result[0]]++; + } + + LOG("\n"); + LOG("Empirical vs target distribution after %d trials:\n", n_trials); + + bool ok = true; + size_t covered = 0; + + auto get_count_for = [&](llama_token id) -> int { + const auto it = token_counts.find(id); + return it != token_counts.end() ? it->second : 0; + }; + + // Iterate over the target distribution + for (size_t i = 0; i < dist_tgt->size; ++i) { + const llama_token id = dist_tgt->data[i].id; + const float p_t = dist_tgt->data[i].p; + + const int count = get_count_for(id); + const float p_e = (float) count / n_trials; + const float diff = std::fabs(p_t - p_e); + + // We only deal with ids in dist_tgt, and if each trial result was one + // of dist_tgt's candidates then covered ends up equal to the number of + // trials. If not this inconsistency tells us how many trials produced + // a token outside of the target's distribution. + covered += count; + + const float standard_error = std::sqrt(p_t * (1.0f - p_t) / n_trials); + const float tolerance = 6.0f * standard_error; + + LOG("token %3d : target = %f, empirical = %f, diff = %f, tol = %f\n", id, p_t, p_e, diff, tolerance); + + if (diff > tolerance) { + LOG_ERR("token %d: empirical probability %f too far from target %f (diff %f > tol %f)\n", + id, p_e, p_t, diff, tolerance); + ok = false; + } + } + + if (covered != (size_t) n_trials) { + LOG_ERR("%zu of %d trials produced a token outside dist_tgt's candidates\n", + (size_t) n_trials - covered, n_trials); + ok = false; + } + + if (!ok) { + return 1; + } + + return 0; +} + +struct parsed_args { + int n_trials; + std::vector argv; +}; + +static std::optional parse_n_trials_arg(int argc, char ** argv, int default_n_trials) { + parsed_args result; + result.n_trials = default_n_trials; + result.argv.push_back(argv[0]); + + for (int i = 1; i < argc; ++i) { + if (strcmp(argv[i], "--n-trials") == 0) { + if (i + 1 >= argc) { + LOG_ERR("%s: --n-trials requires a number argument\n", __func__); + return std::nullopt; + } + int value = std::atoi(argv[i + 1]); + if (value <= 0) { + LOG_ERR("%s: --n-trials requires a positive number argument\n", __func__); + return std::nullopt; + } + result.n_trials = value; + i++; + } else { + result.argv.push_back(argv[i]); + } + } + result.argv.push_back(nullptr); + + return result; +} + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + + common_params params; + params.n_ctx = 256; + + common_init(); + + auto parsed = parse_n_trials_arg(argc, argv, 5000); + if (!parsed) { + return 1; + } + const int fargc = (int) parsed->argv.size() - 1; + + if (!common_params_parse(fargc, parsed->argv.data(), params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + LOG("Using target model '%s'\n", params.model.path.c_str()); + + llama_backend_init(); + + const int result = test_speculative_rejection_sampling(params, parsed->n_trials); + if (result == 0) { + LOG("\nAll tests passed.\n"); + } else { + LOG("\nSome tests failed.\n"); + } + + llama_backend_free(); + return result; +}