diff --git a/llama.cpp b/llama.cpp index 4cd0f16eb..391c956ec 100644 --- a/llama.cpp +++ b/llama.cpp @@ -8133,6 +8133,11 @@ void llama_sample_softmax(struct llama_context * ctx, llama_token_data_array * c } void llama_sample_top_k(struct llama_context * ctx, llama_token_data_array * candidates, int32_t k, size_t min_keep) { + // TODO: move bucket sort to separate function so that top_p/tail_free/typical/softmax first is equally fast + // if (k >= (int32_t)candidates->size) { + // return; + // } + const int64_t t_start_sample_us = ggml_time_us(); k = std::max(k, (int) min_keep); diff --git a/tests/test-sampling.cpp b/tests/test-sampling.cpp index 32e58941c..c3b3d6629 100644 --- a/tests/test-sampling.cpp +++ b/tests/test-sampling.cpp @@ -5,11 +5,10 @@ #undef NDEBUG #endif -#include -#include -#include -#include #include +#include +#include +#include static void dump(const llama_token_data_array * candidates) { for (size_t i = 0; i < candidates->size; i++) { @@ -20,11 +19,11 @@ static void dump(const llama_token_data_array * candidates) { #define DUMP(__candidates) do { printf("%s:%d (%s)\n", __FILE__, __LINE__, __func__); dump((__candidates)); printf("-\n"); } while(0) static void test_top_k(const std::vector & probs, const std::vector & expected_probs, int k) { - size_t n_vocab = probs.size(); + const size_t n_vocab = probs.size(); std::vector candidates; candidates.reserve(n_vocab); for (llama_token token_id = 0; token_id < (llama_token)n_vocab; token_id++) { - float logit = log(probs[token_id]); + const float logit = logf(probs[token_id]); candidates.emplace_back(llama_token_data{token_id, logit, 0.0f}); } @@ -41,11 +40,11 @@ static void test_top_k(const std::vector & probs, const std::vector & probs, const std::vector & expected_probs, float p) { - size_t n_vocab = probs.size(); + const size_t n_vocab = probs.size(); std::vector candidates; candidates.reserve(n_vocab); for (llama_token token_id = 0; token_id < (llama_token)n_vocab; token_id++) { - float logit = log(probs[token_id]); + const float logit = logf(probs[token_id]); candidates.emplace_back(llama_token_data{token_id, logit, 0.0f}); } @@ -62,11 +61,11 @@ static void test_top_p(const std::vector & probs, const std::vector & probs, const std::vector & expected_probs, float z) { - size_t n_vocab = probs.size(); + const size_t n_vocab = probs.size(); std::vector candidates; candidates.reserve(n_vocab); for (llama_token token_id = 0; token_id < (llama_token)n_vocab; token_id++) { - float logit = log(probs[token_id]); + const float logit = logf(probs[token_id]); candidates.emplace_back(llama_token_data{token_id, logit, 0.0f}); } @@ -81,12 +80,33 @@ static void test_tfs(const std::vector & probs, const std::vector } } -static void test_typical(const std::vector & probs, const std::vector & expected_probs, float p) { - size_t n_vocab = probs.size(); +static void test_min_p(const std::vector & probs, const std::vector & expected_probs, float p) { + const size_t n_vocab = probs.size(); std::vector candidates; candidates.reserve(n_vocab); for (llama_token token_id = 0; token_id < (llama_token)n_vocab; token_id++) { - float logit = log(probs[token_id]); + const float logit = logf(probs[token_id]); + candidates.emplace_back(llama_token_data{token_id, logit, 0.0f}); + } + + llama_token_data_array candidates_p = { candidates.data(), candidates.size(), false }; + DUMP(&candidates_p); + llama_sample_min_p(nullptr, &candidates_p, p, 1); + DUMP(&candidates_p); + llama_sample_softmax(nullptr, &candidates_p); + + GGML_ASSERT(candidates_p.size == expected_probs.size()); + for (size_t i = 0; i < candidates_p.size; i++) { + GGML_ASSERT(fabs(candidates_p.data[i].p - expected_probs[i]) < 1e-3); + } +} + +static void test_typical(const std::vector & probs, const std::vector & expected_probs, float p) { + const size_t n_vocab = probs.size(); + std::vector candidates; + candidates.reserve(n_vocab); + for (llama_token token_id = 0; token_id < (llama_token)n_vocab; token_id++) { + const float logit = logf(probs[token_id]); candidates.emplace_back(llama_token_data{token_id, logit, 0.0f}); } @@ -107,11 +127,11 @@ static void test_repetition_penalties( ) { GGML_ASSERT(probs.size() == expected_probs.size()); - size_t n_vocab = probs.size(); + const size_t n_vocab = probs.size(); std::vector candidates; candidates.reserve(n_vocab); for (llama_token token_id = 0; token_id < (llama_token)n_vocab; token_id++) { - float logit = log(probs[token_id]); + const float logit = logf(probs[token_id]); candidates.emplace_back(llama_token_data{token_id, logit, 0.0f}); } @@ -128,6 +148,88 @@ static void test_repetition_penalties( } } +static void test_sampler_queue( + const size_t n_vocab, const std::string samplers_sequence, const int top_k, const float top_p, const float min_p +) { + std::vector candidates; + candidates.reserve(n_vocab); + for (llama_token token_id = 0; token_id < (llama_token)n_vocab; token_id++) { + const float logit = logf(token_id); + candidates.emplace_back(llama_token_data{token_id, logit, 0.0f}); + } + + llama_token_data_array candidates_p = { candidates.data(), candidates.size(), false }; + + llama_token min_token_id = 0; + const llama_token max_token_id = n_vocab-1; + + for (auto s : samplers_sequence) { + switch (s){ + case 'k': llama_sample_top_k (nullptr, &candidates_p, top_k, 1); break; + case 'f': GGML_ASSERT(false && "tail_free test not implemented"); break; + case 'y': GGML_ASSERT(false && "typical test not implemented"); break; + case 'p': llama_sample_top_p (nullptr, &candidates_p, top_p, 1); break; + case 'm': llama_sample_min_p (nullptr, &candidates_p, min_p, 1); break; + case 't': GGML_ASSERT(false && "temperature test not implemented"); break; + default : GGML_ASSERT(false && "Unknown sampler"); break; + } + + llama_sample_softmax(nullptr, &candidates_p); // make sure tokens are sorted for tests + + const int size = candidates_p.size; + + if (s == 'k') { + const int expected_size = std::min(size, top_k); + min_token_id = std::max(min_token_id, (llama_token)(n_vocab - top_k)); + + GGML_ASSERT(size == expected_size); + GGML_ASSERT(candidates_p.data[0].id == max_token_id); + GGML_ASSERT(candidates_p.data[expected_size-1].id == min_token_id); + } else if (s == 'p') { + const int softmax_divisor = n_vocab * (n_vocab-1) / 2 - min_token_id * (min_token_id-1) / 2; + const int softmax_numerator_target = ceilf(top_p * softmax_divisor); + + min_token_id = n_vocab; + int expected_size = 0; + int cumsum = 0; + do { // do-while because always at least one token is sampled + min_token_id--; + expected_size++; + + cumsum += min_token_id; + } while (cumsum < softmax_numerator_target); + + // token 0 has p == 0, need special consideration for cumsum because top_p immediately returns + if (min_token_id == 1) { + min_token_id--; + expected_size += 1; + } + + GGML_ASSERT(size == expected_size); + GGML_ASSERT(candidates_p.data[0].id == max_token_id); + GGML_ASSERT(candidates_p.data[expected_size-1].id == min_token_id); + } else if (s == 'm') { + int expected_size = ceilf((1.0f-min_p) * n_vocab); + expected_size = std::max(expected_size, 1); + expected_size = std::min(expected_size, size); + + min_token_id = floorf(min_p * n_vocab); + min_token_id = std::max(min_token_id, 1); + min_token_id = std::max(min_token_id, (llama_token)(n_vocab - size)); + min_token_id = std::min(min_token_id, (llama_token)(n_vocab - 1)); + + GGML_ASSERT(size == expected_size); + GGML_ASSERT(candidates_p.data[0].id == max_token_id); + GGML_ASSERT(candidates_p.data[expected_size-1].id == min_token_id); + } else { + GGML_ASSERT(false); + } + } + + printf("Sampler queue %3s OK with n_vocab=%05ld top_k=%05d top_p=%f min_p=%f\n", + samplers_sequence.c_str(), n_vocab, top_k, top_p, min_p); +} + int main(void) { ggml_time_init(); @@ -139,6 +241,15 @@ int main(void) { test_top_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f, 0.3f, 0.2f}, 0.8f); test_top_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f, 0.3f, 0.2f, 0.1f}, 1); + test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/1.0f, 0.3f/1.0f, 0.2f/1.0f, 0.1f/1.0f}, 0.00f); + test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/1.0f, 0.3f/1.0f, 0.2f/1.0f, 0.1f/1.0f}, 0.24f); + test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.9f, 0.3f/0.9f, 0.2f/0.9f}, 0.26f); + test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.9f, 0.3f/0.9f, 0.2f/0.9f}, 0.49f); + test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.7f, 0.3f/0.7f}, 0.51f); + test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.7f, 0.3f/0.7f}, 0.74f); + test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.4f}, 0.76f); + test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.4f}, 1.00f); + test_tfs({0.1f, 0.15f, 0.2f, 0.25f, 0.3f}, {0.3f}, 0.25f); test_tfs({0.1f, 0.15f, 0.2f, 0.25f, 0.3f}, {0.3f, 0.25f}, 0.75f); test_tfs({0.1f, 0.15f, 0.2f, 0.25f, 0.3f}, {0.3f, 0.25f}, 0.99f); @@ -154,6 +265,34 @@ int main(void) { test_repetition_penalties({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2}, {0.499966f, 0.499966f, 0.000023f, 0.000023f, 0.000023f}, 1.0f, 5.0f, 5.0f); test_repetition_penalties({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, {0.499977f, 0.499977f, 0.000023f, 0.000023f, 0.000000f}, 1.0f, 5.0f, 5.0f); + test_sampler_queue(10000, "k", 10000, 1.0f, 1.0f); + test_sampler_queue(10000, "k", 1, 1.0f, 1.0f); + test_sampler_queue(10000, "p", 10000, 1.0f, 1.0f); + test_sampler_queue(10000, "p", 10000, 0.0f, 1.0f); + test_sampler_queue(10000, "m", 10000, 1.0f, 1.0f); + test_sampler_queue(10000, "m", 10000, 1.0f, 1e-12); + + test_sampler_queue(10000, "k", 100, 1.0000f, 1.0f); + test_sampler_queue(10000, "p", 10000, 0.0002f, 1.0f); + test_sampler_queue(10000, "p", 10000, 0.8000f, 1.0f); + test_sampler_queue(10000, "m", 10000, 1.0000f, 9997.9f/9999.0f); + test_sampler_queue(10000, "m", 10000, 1.0000f, 0.1f); + + test_sampler_queue(10000, "kp", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "km", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "pk", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "pm", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "mk", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "mp", 100, 0.8f, 9997.9f/9999.0f); + test_sampler_queue(10000, "mp", 100, 0.8f, 0.1f); + + test_sampler_queue(10000, "kpm", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "kmp", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "pkm", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "pmk", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "mkp", 100, 0.8f, 0.1f); + test_sampler_queue(10000, "mpk", 100, 0.8f, 0.1f); + printf("OK\n"); return 0;