diff --git a/common/sampling.cpp b/common/sampling.cpp index ef5cb43a..60318fa7 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -118,7 +118,7 @@ struct llama_sampling_context * common_sampler_init(const struct llama_vocab* vo } case llama_sampler_type::ADAPTIVE_P: { - result->adapt_p_ctx=llama_sampler_init_adaptive_p(params.adaptive_target, params.adaptive_decay, result->rng()); + result->adapt_p_ctx = llama_init_adaptive_p(params.adaptive_target, params.adaptive_decay, result->rng()); break; } default: @@ -423,7 +423,7 @@ static void sampler_queue( } if (use_adaptive_p) { // adaptive p should be put to the last, so we ignore the order in the sampler - llama_sample_adaptive_p(ctx_main, ctx_sampling->adapt_p_ctx, &cur_p); + llama_sample_adaptive_p(ctx_main, &cur_p, ctx_sampling->adapt_p_ctx); } } @@ -471,15 +471,9 @@ static llama_token llama_sampling_sample_impl( id = llama_sample_token_mirostat_v2(ctx_main, &cur_p, mirostat_tau, mirostat_eta, &ctx_sampling->mirostat_mu); } else if (adaptive_target >= 0.0f && ctx_sampling->adapt_p_ctx!=nullptr) { // adaptive p sampling - static thread_local std::vector orig_probs; - orig_probs.resize(cur_p.size); - - // store original probabilities - for (size_t ii = 0; ii < cur_p.size; ++ii) { - orig_probs[ii] = cur_p.data[ii].p; - } + llama_prep_adaptive_p(&cur_p, ctx_sampling->adapt_p_ctx); sampler_queue(ctx_main, params, ctx_sampling, cur_p, std::max(1, params.min_keep)); - id = llama_sample_token_adaptive_p(ctx_main, &cur_p, ctx_sampling->adapt_p_ctx, orig_probs.data()); + id = llama_sample_token_adaptive_p(ctx_main, &cur_p, ctx_sampling->adapt_p_ctx); } else { // temperature sampling size_t min_keep = std::max(1, params.min_keep); diff --git a/include/llama.h b/include/llama.h index 67f46b50..3d17f9b2 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1384,16 +1384,20 @@ LLAMA_API struct llama_grammar* llama_sampler_init_grammar_lazy_patterns( /// @details Adaptive p sampler initializer /// @param target Select tokens near this probability (valid range 0.0 to 1.0; <0 = disabled) /// @param decay Decay rate for target adaptation over time. lower values -> faster but less stable adaptation. (valid range 0.0 to 1.0; ≤0 = no adaptation) - LLAMA_API struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p( + LLAMA_API struct llama_sampler_adaptive_p * llama_init_adaptive_p( const float target, const float decay, const uint32_t seed); + void llama_prep_adaptive_p( + llama_token_data_array * candidates, + struct llama_sampler_adaptive_p * adapt_p_ctx); + /// @details Adaptive p sampler described in https://github.com/MrJackSpade/adaptive-p-docs/blob/main/README.md void llama_sample_adaptive_p( - struct llama_context * ctx, - struct llama_sampler_adaptive_p * adapt_p_ctx, - llama_token_data_array * candidates); + struct llama_context * ctx, + llama_token_data_array * candidates, + struct llama_sampler_adaptive_p * adapt_p_ctx); /// @details Mirostat 1.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words. @@ -1437,8 +1441,7 @@ LLAMA_API struct llama_grammar* llama_sampler_init_grammar_lazy_patterns( llama_token llama_sample_token_adaptive_p( struct llama_context * ctx, llama_token_data_array * candidates, - struct llama_sampler_adaptive_p * adapt_p_ctx, - float * orig_probs); + struct llama_sampler_adaptive_p * adapt_p_ctx); // // Model split diff --git a/src/llama-sampling.cpp b/src/llama-sampling.cpp index 244c18d8..9d3134e4 100644 --- a/src/llama-sampling.cpp +++ b/src/llama-sampling.cpp @@ -1038,41 +1038,47 @@ struct llama_sampler_dry* llama_sampler_init_dry_impl(const struct llama_vocab& llama_token llama_sample_token_adaptive_p_impl( struct llama_sampling * smpl, llama_token_data_array * candidates, - struct llama_sampler_adaptive_p * adapt_p_ctx, - float * orig_probs) + struct llama_sampler_adaptive_p * adapt_p_ctx) { GGML_ASSERT(candidates->size > 0); const int64_t t_start_sample_us = ggml_time_us(); - const size_t count = candidates->size; - adapt_p_ctx->probs.resize(count); + struct llama_sampler_adaptive_p * ctx = adapt_p_ctx; + ctx->cum_probs.resize(candidates->size); - // cumulative distribution - const float max_logit = adapt_p_ctx->max_logit; + // compute cumulative probability distribution + const float max_logit = ctx->max_xform_logit; float cum_prob = 0.0f; - for (size_t i = 0; i < count; ++i) { + for (size_t i = 0; i < candidates->size; ++i) { cum_prob += expf(candidates->data[i].logit - max_logit); - adapt_p_ctx->probs[i] = cum_prob; + ctx->cum_probs[i] = cum_prob; } - adapt_p_ctx->probs.back() += 1.0f; // safety margin in case rng() ~= rng.max() + ctx->cum_probs.back() += 1.0f; // safety margin in case rng() ~= rng.max() - // find token with cum_prob > target_cum_prob - const float target_cum_prob = cum_prob * (float)adapt_p_ctx->rng() / (float)adapt_p_ctx->rng.max(); - auto iter = std::upper_bound(adapt_p_ctx->probs.begin(), adapt_p_ctx->probs.end(), target_cum_prob); - GGML_ASSERT(iter != adapt_p_ctx->probs.end()); - llama_token id = candidates->data[std::distance(adapt_p_ctx->probs.begin(), iter)].id; + // select first token whose cum_prob > target_cum_prob + const float target_cum_prob = cum_prob * (float)ctx->rng() / (float)ctx->rng.max(); + auto iter = std::upper_bound(ctx->cum_probs.begin(), ctx->cum_probs.end(), target_cum_prob); + GGML_ASSERT(iter != ctx->cum_probs.end()); + const size_t idx = std::distance(ctx->cum_probs.begin(), iter); + llama_token id = candidates->data[idx].id; smpl->t_sample_us += ggml_time_us() - t_start_sample_us; smpl->n_sample++; + float update_prob = candidates->data[idx].p; // not ideal + if (ctx->orig_prob_map.contains(id)) { + // selected token id is among tracked ids + update_prob = ctx->orig_prob_map[id] / ctx->cum_orig_prob; + } + // update history with original probability of selected token - adapt_p_ctx->weighted_sum = adapt_p_ctx->decay * adapt_p_ctx->weighted_sum + orig_probs[id]; - adapt_p_ctx->total_weight = adapt_p_ctx->decay * adapt_p_ctx->total_weight + 1.0f; + ctx->weighted_sum = ctx->decay * ctx->weighted_sum + update_prob; + ctx->total_weight = ctx->decay * ctx->total_weight + 1.0f; return id; } -void llama_sampler_adaptive_p_apply(struct llama_sampler_adaptive_p * adapt_p_ctx, llama_token_data_array * candidates) +void llama_sample_adaptive_p_impl(llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx) { if (adapt_p_ctx->target < 0.0f) { // sampler is disabled @@ -1082,14 +1088,16 @@ void llama_sampler_adaptive_p_apply(struct llama_sampler_adaptive_p * adapt_p_ct // incomplete softmax because final division can be fused float max_l = candidates->data[0].logit; - for (size_t i = 1; i < candidates->size; ++i) { - max_l = std::max(max_l, candidates->data[i].logit); + if (!candidates->sorted) { + for (size_t i = 1; i < candidates->size; ++i) { + max_l = std::max(max_l, candidates->data[i].logit); + } } float cum_sum = 0.0f; for (size_t i = 0; i < candidates->size; ++i) { - const float p = expf(candidates->data[i].logit - max_l); - candidates->data[i].p = p; - cum_sum += p; + const float prob = expf(candidates->data[i].logit - max_l); + candidates->data[i].p = prob; + cum_sum += prob; } // compute adapted target probability @@ -1117,10 +1125,45 @@ void llama_sampler_adaptive_p_apply(struct llama_sampler_adaptive_p * adapt_p_ct max_logit = std::max(max_logit, logit); } candidates->sorted = false; - adapt_p_ctx->max_logit = max_logit; + adapt_p_ctx->max_xform_logit = max_logit; } -struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p_impl( +void llama_prep_adaptive_p_impl( + llama_token_data_array * candidates, + struct llama_sampler_adaptive_p * adapt_p_ctx) +{ + if (!candidates->sorted) { + std::sort(candidates->data, candidates->data + candidates->size, + [](const llama_token_data & a, const llama_token_data & b) { + return a.logit > b.logit; + }); + candidates->sorted = true; + } + const float max_logit = candidates->data[0].logit; + + // decide how many tokens to track based on logit delta + // i.e. do not track unlikely tokens + auto iter = std::lower_bound( + candidates->data, + candidates->data + candidates->size, + max_logit - 16.6f, // delta + [](const llama_token_data & data, const float delta) { + return data.logit > delta; + }); + const size_t n_track = std::distance(candidates->data, iter); + + // store orig_prob_map and cum_orig_prob to estimate original probability later + float cum_prob = 0.0f; + adapt_p_ctx->orig_prob_map.clear(); + for (size_t i = 0; i < n_track; ++i) { + const float prob = expf(candidates->data[i].logit - max_logit); + cum_prob += prob; + adapt_p_ctx->orig_prob_map[candidates->data[i].id] = prob; + } + adapt_p_ctx->cum_orig_prob = cum_prob; +} + +struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl( const float target, const float decay, const uint32_t seed) @@ -1132,12 +1175,13 @@ struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p_impl( /* .rng = */ std::mt19937(seed), /* .weighted_sum = */ target / (1.0f - clamped_decay), /* .total_weight = */ 1.0f / (1.0f - clamped_decay), - /* .max_logit = */ 0.0f, - /* .probs = */ {}, + /* .orig_logit_map = */ {}, + /* .cum_orig_prob = */ 0.0f, + /* .max_xform_logit = */ -INFINITY, + /* .cum_probs = */ {}, }; } - // grammar struct llama_sampler_grammar { diff --git a/src/llama-sampling.h b/src/llama-sampling.h index 61228548..55b4371c 100644 --- a/src/llama-sampling.h +++ b/src/llama-sampling.h @@ -61,6 +61,7 @@ struct llama_sampler_dry * llama_sampler_init_dry_impl( void llama_sampler_dry_apply(struct llama_sampler_dry* smpl, llama_token_data_array* cur_p); + // maintains an exponential moving average of the *ORIGINAL* probabilities of selected tokens // used to compute an adapted target at each sampling step. // see llama.h for a full description of the sampler @@ -70,15 +71,30 @@ struct llama_sampler_adaptive_p { std::mt19937 rng; // RNG float weighted_sum; // sum(p_n * decay^N) float total_weight; // sum(decay^i), converges to 1/(1-decay) - float max_logit; // maximum logit found during transform - std::vector probs; // cumulative probabilities + + // first referenced in prep + std::unordered_map orig_prob_map; // probabilities before sampler_queue + float cum_orig_prob; // for normalizing orig_prob in sample_token + + // first referenced in sample + float max_xform_logit; // maximum logit found during transform + + // first referenced in sample_token + std::vector cum_probs; // cumulative probability distribution }; -void llama_sampler_adaptive_p_apply( - struct llama_sampler_adaptive_p * adapt_p_ctx, - llama_token_data_array * candidates); +struct llama_sampler_adaptive_p * llama_init_adaptive_p_impl( + const float target, + const float decay, + const uint32_t seed); -struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p_impl(const float target, const float decay, const uint32_t seed); +void llama_prep_adaptive_p_impl( + llama_token_data_array * candidates, + struct llama_sampler_adaptive_p * adapt_p_ctx); + +void llama_sample_adaptive_p_impl( + llama_token_data_array * candidates, + struct llama_sampler_adaptive_p * adapt_p_ctx); void llama_sample_repetition_penalties_impl( @@ -101,6 +117,6 @@ llama_token llama_sample_token_mirostat_v2_impl(struct llama_sampling * smpl, ll llama_token llama_sample_token_greedy_impl (struct llama_sampling * smpl, llama_token_data_array * candidates); llama_token llama_sample_token_with_rng_impl (struct llama_sampling * smpl, llama_token_data_array * candidates, std::mt19937 & rng); llama_token llama_sample_token_impl (struct llama_sampling * smpl, llama_token_data_array * candidates); -llama_token llama_sample_token_adaptive_p_impl (struct llama_sampling * smpl, llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx, float * orig_probs); +llama_token llama_sample_token_adaptive_p_impl (struct llama_sampling * smpl, llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx); diff --git a/src/llama.cpp b/src/llama.cpp index 347b3b70..7175eca9 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -7689,11 +7689,18 @@ void llama_sample_dry([[maybe_unused]] struct llama_context* ctx, struct llama_s void llama_sample_adaptive_p( [[maybe_unused]] struct llama_context * ctx, - struct llama_sampler_adaptive_p * adapt_p_ctx, - llama_token_data_array * candidates) { - llama_sampler_adaptive_p_apply(adapt_p_ctx, candidates); + llama_token_data_array * candidates, + struct llama_sampler_adaptive_p * adapt_p_ctx) +{ + llama_sample_adaptive_p_impl(candidates, adapt_p_ctx); } +void llama_prep_adaptive_p(llama_token_data_array * candidates, struct llama_sampler_adaptive_p * adapt_p_ctx) +{ + llama_prep_adaptive_p_impl(candidates, adapt_p_ctx); +} + + void llama_sample_repetition_penalties( struct llama_context * ctx, llama_token_data_array * candidates, @@ -7736,10 +7743,9 @@ llama_token llama_sample_token(struct llama_context * ctx, llama_token_data_arra llama_token llama_sample_token_adaptive_p( struct llama_context * ctx, llama_token_data_array * candidates, - struct llama_sampler_adaptive_p * adapt_p_ctx, - float * orig_probs) + struct llama_sampler_adaptive_p * adapt_p_ctx) { - return llama_sample_token_adaptive_p_impl(&ctx->sampling, candidates, adapt_p_ctx, orig_probs); + return llama_sample_token_adaptive_p_impl(&ctx->sampling, candidates, adapt_p_ctx); } int llama_split_path(char * split_path, size_t maxlen, const char * path_prefix, int split_no, int split_count) { @@ -7794,9 +7800,9 @@ void llama_sampler_dry_accept(struct llama_sampler_dry* smpl, llama_token token) } -struct llama_sampler_adaptive_p * llama_sampler_init_adaptive_p(const float target, const float decay, const uint32_t seed) +struct llama_sampler_adaptive_p * llama_init_adaptive_p(const float target, const float decay, const uint32_t seed) { - return llama_sampler_init_adaptive_p_impl(target, decay, seed); + return llama_init_adaptive_p_impl(target, decay, seed); }