KBaba7/llama.cpp
0
1#include "ggml.h"2#include "llama.h"3 4#ifdef NDEBUG5#undef NDEBUG6#endif7 8#include <algorithm>9#include <cmath>10#include <string>11#include <vector>12 13extern struct llama_sampler * llama_sampler_init_dry_testing(int32_t context_size, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const std::vector<std::vector<llama_token>>& seq_breakers);14 15static void dump(const llama_token_data_array * cur_p) {16 for (size_t i = 0; i < cur_p->size; i++) {17 printf("%d: %f (%f)\n", cur_p->data[i].id, cur_p->data[i].p, cur_p->data[i].logit);18 }19}20 21#define DUMP(__cur_p) do { printf("%s:%d (%s)\n", __FILE__, __LINE__, __func__); dump((__cur_p)); printf("-\n"); } while(0)22 23struct sampler_tester {24 sampler_tester(size_t n_vocab) {25 cur.reserve(n_vocab);26 for (llama_token token_id = 0; token_id < (llama_token)n_vocab; token_id++) {27 const float logit = logf(token_id);28 cur.emplace_back(llama_token_data{token_id, logit, 0.0f});29 }30 31 cur_p = llama_token_data_array { cur.data(), cur.size(), -1, false };32 }33 34 sampler_tester(const std::vector<float> & probs, const std::vector<float> & probs_expected) : probs_expected(probs_expected) {35 cur.reserve(probs.size());36 for (llama_token token_id = 0; token_id < (llama_token)probs.size(); token_id++) {37 const float logit = logf(probs[token_id]);38 cur.emplace_back(llama_token_data{token_id, logit, probs[token_id]});39 }40 41 cur_p = llama_token_data_array { cur.data(), cur.size(), -1, false };42 }43 44 void apply(llama_sampler * sampler) {45 llama_sampler_apply(sampler, &cur_p);46 llama_sampler_free(sampler);47 }48 49 void check() {50 GGML_ASSERT(cur_p.size == probs_expected.size());51 for (size_t i = 0; i < cur_p.size; i++) {52 GGML_ASSERT(fabs(cur_p.data[i].p - probs_expected[i]) < 1e-5);53 }54 }55 56 llama_token_data_array cur_p;57 58private:59 const std::vector<float> probs_expected;60 61 std::vector<llama_token_data> cur;62};63 64static void test_temp(const std::vector<float> & probs, const std::vector<float> & probs_expected, float temp) {65 sampler_tester tester(probs, probs_expected);66 67 DUMP(&tester.cur_p);68 tester.apply(llama_sampler_init_temp(temp));69 tester.apply(llama_sampler_init_dist(0));70 DUMP(&tester.cur_p);71 72 tester.check();73}74 75static void test_temp_ext(const std::vector<float> & probs, const std::vector<float> & probs_expected, float temp, float delta, float exponent) {76 sampler_tester tester(probs, probs_expected);77 78 DUMP(&tester.cur_p);79 tester.apply(llama_sampler_init_temp_ext(temp, delta, exponent));80 tester.apply(llama_sampler_init_dist (0));81 DUMP(&tester.cur_p);82 83 tester.check();84}85 86static void test_top_k(const std::vector<float> & probs, const std::vector<float> & probs_expected, int k) {87 sampler_tester tester(probs, probs_expected);88 89 DUMP(&tester.cur_p);90 tester.apply(llama_sampler_init_top_k(k));91 tester.apply(llama_sampler_init_dist (0));92 DUMP(&tester.cur_p);93 94 tester.check();95}96 97static void test_top_p(const std::vector<float> & probs, const std::vector<float> & probs_expected, float p) {98 sampler_tester tester(probs, probs_expected);99 100 DUMP(&tester.cur_p);101 tester.apply(llama_sampler_init_top_p(p, 1));102 tester.apply(llama_sampler_init_dist (0));103 DUMP(&tester.cur_p);104 105 tester.check();106}107 108static void test_min_p(const std::vector<float> & probs, const std::vector<float> & probs_expected, float p) {109 sampler_tester tester(probs, probs_expected);110 111 DUMP(&tester.cur_p);112 tester.apply(llama_sampler_init_min_p(p, 1));113 tester.apply(llama_sampler_init_dist (0));114 DUMP(&tester.cur_p);115 116 tester.check();117}118 119static void test_xtc(const std::vector<float> & probs, const std::vector<float> & probs_expected, float p, float t) {120 sampler_tester tester(probs, probs_expected);121 122 DUMP(&tester.cur_p);123 tester.apply(llama_sampler_init_xtc(p, t, 0, 0));124 DUMP(&tester.cur_p);125 126 tester.check();127}128 129static void test_typical(const std::vector<float> & probs, const std::vector<float> & probs_expected, float p) {130 sampler_tester tester(probs, probs_expected);131 132 DUMP(&tester.cur_p);133 tester.apply(llama_sampler_init_typical(p, 1));134 DUMP(&tester.cur_p);135 136 tester.check();137}138 139static void test_penalties(140 const std::vector<float> & probs, const std::vector<llama_token> & last_tokens,141 const std::vector<float> & probs_expected, float repeat_penalty, float alpha_frequency, float alpha_presence142) {143 GGML_ASSERT(probs.size() == probs_expected.size());144 145 sampler_tester tester(probs, probs_expected);146 147 auto * sampler = llama_sampler_init_penalties(last_tokens.size(), repeat_penalty, alpha_frequency, alpha_presence);148 149 for (size_t i = 0; i < last_tokens.size(); i++) {150 llama_sampler_accept(sampler, last_tokens[i]);151 }152 153 DUMP(&tester.cur_p);154 tester.apply(sampler);155 tester.apply(llama_sampler_init_dist(0));156 DUMP(&tester.cur_p);157 158 tester.check();159}160 161static void test_dry(162 const std::vector<float> & probs, const std::vector<llama_token> & last_tokens,163 const std::vector<float> & expected_probs, float dry_multiplier, float dry_base,164 int dry_allowed_length, int dry_penalty_last_n,165 const std::vector<std::vector<llama_token>> & seq_breakers166) {167 GGML_ASSERT(probs.size() == expected_probs.size());168 169 sampler_tester tester(probs, expected_probs);170 171 auto * sampler = llama_sampler_init_dry_testing(1024, dry_multiplier, dry_base, dry_allowed_length, dry_penalty_last_n, seq_breakers);172 173 for (size_t i = 0; i < last_tokens.size(); i++) {174 llama_sampler_accept(sampler, last_tokens[i]);175 }176 177 DUMP(&tester.cur_p);178 tester.apply(sampler);179 tester.apply(llama_sampler_init_dist(0));180 DUMP(&tester.cur_p);181 tester.check();182}183 184static 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_p185) {186 sampler_tester tester(n_vocab);187 188 llama_token min_token_id = 0;189 const llama_token max_token_id = n_vocab-1;190 191 for (auto s : samplers_sequence) {192 switch (s){193 case 'k': tester.apply(llama_sampler_init_top_k(top_k)); break;194 case 'y': GGML_ABORT("typical test not implemented");195 case 'p': tester.apply(llama_sampler_init_top_p(top_p, 1)); break;196 case 'm': tester.apply(llama_sampler_init_min_p(min_p, 1)); break;197 case 't': GGML_ABORT("temperature test not implemented");198 default : GGML_ABORT("Unknown sampler");199 }200 201 tester.apply(llama_sampler_init_dist(0));202 203 auto & cur_p = tester.cur_p;204 205 const int size = cur_p.size;206 207 if (s == 'k') {208 const int expected_size = std::min(size, top_k);209 min_token_id = std::max(min_token_id, (llama_token)(n_vocab - top_k));210 211 GGML_ASSERT(size == expected_size);212 GGML_ASSERT(cur_p.data[0].id == max_token_id);213 GGML_ASSERT(cur_p.data[expected_size-1].id == min_token_id);214 } else if (s == 'p') {215 const int softmax_divisor = n_vocab * (n_vocab-1) / 2 - min_token_id * (min_token_id-1) / 2;216 const int softmax_numerator_target = ceilf(top_p * softmax_divisor);217 218 min_token_id = n_vocab;219 int expected_size = 0;220 int cumsum = 0;221 do { // do-while because always at least one token is sampled222 min_token_id--;223 expected_size++;224 225 cumsum += min_token_id;226 } while (cumsum < softmax_numerator_target);227 228 // token 0 has p == 0, need special consideration for cumsum because top_p immediately returns229 if (min_token_id == 1) {230 min_token_id--;231 expected_size += 1;232 }233 234 GGML_ASSERT(size == expected_size);235 GGML_ASSERT(cur_p.data[0].id == max_token_id);236 GGML_ASSERT(cur_p.data[expected_size-1].id == min_token_id);237 } else if (s == 'm') {238 int expected_size = ceilf((1.0f-min_p) * n_vocab);239 expected_size = std::max(expected_size, 1);240 expected_size = std::min(expected_size, size);241 242 min_token_id = floorf(min_p * n_vocab);243 min_token_id = std::max(min_token_id, 1);244 min_token_id = std::max(min_token_id, (llama_token)(n_vocab - size));245 min_token_id = std::min(min_token_id, (llama_token)(n_vocab - 1));246 247 GGML_ASSERT(size == expected_size);248 GGML_ASSERT(cur_p.data[0].id == max_token_id);249 GGML_ASSERT(cur_p.data[expected_size-1].id == min_token_id);250 } else {251 GGML_ABORT("fatal error");252 }253 }254 255 printf("Sampler queue %3s OK with n_vocab=%05zu top_k=%05d top_p=%f min_p=%f\n",256 samplers_sequence.c_str(), n_vocab, top_k, top_p, min_p);257}258 259static void bench(llama_sampler * cnstr, const char * cnstr_name, const std::vector<llama_token_data> & data, int n_iter) {260 std::vector<llama_token_data> cur(data.size());261 std::copy(data.begin(), data.end(), cur.begin());262 llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false };263 llama_sampler_apply(cnstr, &cur_p);264 llama_sampler_reset(cnstr);265 const int64_t t_start = ggml_time_us();266 for (int i = 0; i < n_iter; i++) {267 std::copy(data.begin(), data.end(), cur.begin());268 llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false };269 llama_sampler_apply(cnstr, &cur_p);270 llama_sampler_reset(cnstr);271 }272 const int64_t t_end = ggml_time_us();273 llama_sampler_free(cnstr);274 printf("%-43s: %8.3f us/iter\n", cnstr_name, (t_end - t_start) / (float)n_iter);275}276 277#define BENCH(__cnstr, __data, __n_iter) bench((__cnstr), #__cnstr, (__data), (__n_iter))278 279static void test_perf() {280 const int n_vocab = 1 << 17;281 282 std::vector<llama_token_data> data;283 284 data.reserve(n_vocab);285 for (int i = 0; i < n_vocab; i++) {286 const float logit = 2.0f*((double)(rand())/RAND_MAX - 0.5);287 data.emplace_back(llama_token_data{i, logit, 0.0f});288 }289 290 BENCH(llama_sampler_init_top_k (40), data, 32);291 BENCH(llama_sampler_init_top_p (0.8f, 1), data, 32);292 BENCH(llama_sampler_init_min_p (0.2f, 1), data, 32);293 BENCH(llama_sampler_init_typical(0.5f, 1), data, 32);294 BENCH(llama_sampler_init_xtc (1.0f, 0.1f, 1, 1), data, 32);295}296 297int main(void) {298 ggml_time_init();299 300 test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f, 0.3f, 0.2f, 0.1f}, 1.0f);301 test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {1.0f, 0.0f, 0.0f, 0.0f}, 0.0f);302 303 test_temp_ext({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f, 0.3f, 0.2f, 0.1f}, 1.0f, 0.0f, 1.0f);304 test_temp_ext({0.1f, 0.2f, 0.3f, 0.4f}, {1.0f, 0.0f, 0.0f, 0.0f}, 0.0f, 0.0f, 1.0f);305 306 test_top_k({0.1f, 0.2f, 0.3f, 0.4f}, {1.0f}, 1);307 test_top_k({0.1f, 0.2f, 0.3f, 0.4f}, {0.44444f, 0.33333f, 0.22222f}, 3);308 test_top_k({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f, 0.3f, 0.2f, 0.1f}, 4);309 test_top_k({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f, 0.3f, 0.2f, 0.1f}, 0);310 311 test_top_p({0.1f, 0.2f, 0.3f, 0.4f}, {1.0f}, 0);312 test_top_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.571429f, 0.428571f}, 0.7f);313 test_top_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.44444f, 0.33333f, 0.22222f}, 0.8f);314 test_top_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f, 0.3f, 0.2f, 0.1f}, 1.0f);315 316 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);317 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);318 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);319 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);320 test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.7f, 0.3f/0.7f}, 0.51f);321 test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.7f, 0.3f/0.7f}, 0.74f);322 test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.4f}, 0.76f);323 test_min_p({0.1f, 0.2f, 0.3f, 0.4f}, {0.4f/0.4f}, 1.00f);324 325 printf("XTC should:\n");326 test_xtc({0.4f, 0.3f, 0.2f, 0.1f}, {0.1f}, 0.99f, 0.09f);327 test_xtc({0.4f, 0.3f, 0.2f, 0.1f}, {0.2f, 0.1f}, 0.99f, 0.19f);328 test_xtc({0.4f, 0.3f, 0.2f, 0.1f}, {0.3f, 0.2f, 0.1f}, 0.99f, 0.29f);329 330 printf("XTC should not:\n");331 test_xtc({0.4f, 0.3f, 0.2f, 0.1f}, {0.4f, 0.3f, 0.2f, 0.1f}, 0.99f, 0.39f);332 333 test_typical({0.97f, 0.01f, 0.01f, 0.01f}, {0.97f}, 0.5f);334 test_typical({0.4f, 0.2f, 0.2f, 0.2f}, {0.2f, 0.2f, 0.2f}, 0.5f);335 336 test_penalties({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0}, {0.25f, 0.25f, 0.25f, 0.25f, 0}, 50.0f, 0.0f, 0.0f);337 test_penalties({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2}, {0.5f, 0.5f, 0, 0, 0}, 50.0f, 0.0f, 0.0f);338 test_penalties({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 0}, {0.5f, 0.5f, 0, 0, 0}, 50.0f, 0.0f, 0.0f);339 340 test_penalties({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0}, {0.249997f, 0.249997f, 0.249997f, 0.249997f, 0.000011f}, 1.0f, 5.0f, 5.0f);341 test_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);342 test_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);343 344 345 test_dry({0.25f, 0.25f, 0.25f, 0.25f}, {0, 1}, {0.25f, 0.25f, 0.25f, 0.25f}, 1.0f, 1.1f, 2, 4, {});346 test_dry({0.25f, 0.25f, 0.25f, 0.25f}, {0, 1, 2, 0, 1}, {0.296923f, 0.296923f, 0.296923f, 0.109232f}, 1.0f, 1.1f, 2, 5, {});347 test_dry({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 3, 4, 0, 1}, {0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, 1.0f, 1.1f, 2, 6, {{3}});348 test_dry({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 0, 1}, {0.241818f, 0.241818f, 0.241818f, 0.241818f, 0.032727f}, 2.0f, 1.1f, 2, 5, {});349 test_dry({0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, {0, 1, 2, 3, 4, 0, 1}, {0.2f, 0.2f, 0.2f, 0.2f, 0.2f}, 1.0f, 1.1f, 4, 7, {});350 351 test_sampler_queue(10000, "k", 10000, 1.0f, 1.0f);352 test_sampler_queue(10000, "k", 1, 1.0f, 1.0f);353 test_sampler_queue(10000, "p", 10000, 1.0f, 1.0f);354 test_sampler_queue(10000, "p", 10000, 0.0f, 1.0f);355 test_sampler_queue(10000, "m", 10000, 1.0f, 1.0f);356 test_sampler_queue(10000, "m", 10000, 1.0f, 1e-12);357 358 test_sampler_queue(10000, "k", 100, 1.0000f, 1.0f);359 test_sampler_queue(10000, "p", 10000, 0.0002f, 1.0f);360 test_sampler_queue(10000, "p", 10000, 0.8000f, 1.0f);361 test_sampler_queue(10000, "m", 10000, 1.0000f, 9997.9f/9999.0f);362 test_sampler_queue(10000, "m", 10000, 1.0000f, 0.1f);363 364 test_sampler_queue(10000, "kp", 100, 0.8f, 0.1f);365 test_sampler_queue(10000, "km", 100, 0.8f, 0.1f);366 test_sampler_queue(10000, "pk", 100, 0.8f, 0.1f);367 test_sampler_queue(10000, "pm", 100, 0.8f, 0.1f);368 test_sampler_queue(10000, "mk", 100, 0.8f, 0.1f);369 test_sampler_queue(10000, "mp", 100, 0.8f, 9997.9f/9999.0f);370 test_sampler_queue(10000, "mp", 100, 0.8f, 0.1f);371 372 test_sampler_queue(10000, "kpm", 100, 0.8f, 0.1f);373 test_sampler_queue(10000, "kmp", 100, 0.8f, 0.1f);374 test_sampler_queue(10000, "pkm", 100, 0.8f, 0.1f);375 test_sampler_queue(10000, "pmk", 100, 0.8f, 0.1f);376 test_sampler_queue(10000, "mkp", 100, 0.8f, 0.1f);377 test_sampler_queue(10000, "mpk", 100, 0.8f, 0.1f);378 379 printf("OK\n");380 381 test_perf();382 383 return 0;384}385