tqcq/llama.cpp
0
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2025-06-28 12:25:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

llama : add llama_sampling API + move grammar in libllama

Browse Source 
ggml-ci
This commit is contained in:
Georgi Gerganov
2024-08-05 10:08:25 +03:00
parent b69a480af4
commit f648ca2cee
48 changed files with 2481 additions and 2590 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
examples/speculative/speculative.cpp
View File

@ -21,7 +21,7 @@ struct seq_draft {
std::vector<llama_token> tokens;
std::vector<std::vector<llama_token_data>> dists;
struct llama_sampling_context * ctx_sampling;
struct llama_sampling * smpl;
};
int main(int argc, char ** argv) {
@ -37,16 +37,16 @@ int main(int argc, char ** argv) {
return 1;
}
// for probabilities to be computed even with temp = 0
params.sparams.n_probs = 16;
// max number of parallel drafting sequences (i.e. tree branches)
const int n_seq_dft = params.n_parallel;
// probability threshold for splitting a draft branch (only for n_seq_dft > 1)
const float p_split = params.p_split;
if (params.seed == LLAMA_DEFAULT_SEED) {
params.seed = time(NULL);
}
std::default_random_engine rng(params.seed);
std::default_random_engine rng(params.sparams.seed);
std::uniform_real_distribution<> u_dist;
#ifndef LOG_DISABLE_LOGS
@ -179,19 +179,15 @@ int main(int argc, char ** argv) {
// used to determine end of generation
bool has_eos = false;
// target model sampling context
struct llama_sampling_context * ctx_sampling = llama_sampling_init(params.sparams);
// target model sampling context (reuse the llama_context's sampling instance)
struct llama_sampling * smpl = llama_sampling_init(model_tgt, params.sparams);
// draft sequence data
std::vector<seq_draft> drafts(n_seq_dft);
params.sparams.grammar.clear(); // the draft samplers will copy the target sampler's grammar
if (params.sparams.temp == 0) {
params.sparams.temp = -1.0f; // force greedy sampling with probs for the draft model
}
for (int s = 0; s < n_seq_dft; ++s) {
drafts[s].ctx_sampling = llama_sampling_init(params.sparams);
// allocate llama_sampling for each draft sequence
drafts[s].smpl = llama_sampling_init(model_dft, params.sparams);
}
llama_batch batch_dft = llama_batch_init(params.n_ctx, 0, 1);
@ -234,9 +230,15 @@ int main(int argc, char ** argv) {
if (params.sparams.temp > 0) {
// stochastic verification
llama_token_data_array dist_tgt = llama_sampling_prepare(ctx_sampling, ctx_tgt, NULL, drafts[s_keep].i_batch_tgt[i_dft], true, NULL);
llama_sample_softmax(ctx_tgt, &dist_tgt);
float p_tgt = 0, p_dft = 0;
llama_sampling_set_logits(smpl, llama_get_logits_ith(ctx_tgt, drafts[s_keep].i_batch_tgt[i_dft]));
auto & dist_tgt = *llama_sampling_get_candidates(smpl);
llama_sampling_grammar(smpl, &dist_tgt);
llama_sampling_softmax(smpl, &dist_tgt);
float p_tgt = 0.0f;
float p_dft = 0.0f;
// GGML_ASSERT(dist_tgt.size() == dist_dft.size());
@ -278,7 +280,7 @@ int main(int argc, char ** argv) {
accept = true;
token_id = drafts[s].tokens[i_dft];
token_str = llama_token_to_piece(ctx_tgt, token_id);
llama_sampling_accept(ctx_sampling, ctx_tgt, token_id, true);
llama_sampling_accept(smpl, token_id, true);
LOG("draft token %d of sequence %d (%d, '%s') accepted\n", i_dft, s, token_id, token_str.c_str());
break;
@ -332,8 +334,8 @@ int main(int argc, char ** argv) {
// all drafted tokens were rejected
// sample from the target model
LOG("all drafted tokens were rejected, sampling from residual distribution\n");
token_id = llama_sample_token(ctx_tgt, &dist_tgt);
llama_sampling_accept(ctx_sampling, ctx_tgt, token_id, true);
token_id = llama_sampling_sample_dist(smpl, &dist_tgt);
llama_sampling_accept(smpl, token_id, true);
token_str = llama_token_to_piece(ctx_tgt, token_id);
}
@ -342,11 +344,11 @@ int main(int argc, char ** argv) {
// sample from the target model
LOG("sampling target: s_keep = %3d, i_dft = %3d, i_batch_tgt = %3d\n", s_keep, i_dft, drafts[s_keep].i_batch_tgt[i_dft]);
token_id = llama_sampling_sample(ctx_sampling, ctx_tgt, NULL, drafts[s_keep].i_batch_tgt[i_dft]);
token_id = llama_sampling_sample(smpl, ctx_tgt, drafts[s_keep].i_batch_tgt[i_dft]);
llama_sampling_accept(ctx_sampling, ctx_tgt, token_id, true);
llama_sampling_accept(smpl, token_id, true);
//LOG("last: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx_tgt, ctx_sampling->prev).c_str());
//LOG("last: %s\n", LOG_TOKENS_TOSTR_PRETTY(ctx_tgt, smpl->prev).c_str());
token_str = llama_token_to_piece(ctx_tgt, token_id);
@ -434,7 +436,7 @@ int main(int argc, char ** argv) {
break;
}
llama_sampling_cp(ctx_sampling, drafts[0].ctx_sampling);
llama_sampling_cp(smpl, drafts[0].smpl);
int n_seq_cur = 1;
int n_past_cur = n_past_dft;
@ -463,20 +465,20 @@ int main(int argc, char ** argv) {
continue;
}
llama_sampling_sample(drafts[s].ctx_sampling, ctx_dft, NULL, drafts[s].i_batch_dft);
llama_sampling_sample(drafts[s].smpl, ctx_dft, drafts[s].i_batch_dft);
const auto & cur_p = drafts[s].ctx_sampling->cur;
const auto * cur_p = llama_sampling_get_candidates(drafts[s].smpl);
for (int k = 0; k < std::min(n_seq_dft + 3, (int) cur_p.size()); ++k) {
for (int k = 0; k < std::min(n_seq_dft + 3, (int) cur_p->size); ++k) {
LOG(" - draft candidate %3d for seq %3d, pos %3d: %6d (%8.3f) '%s'\n",
k, s, i, cur_p[k].id, cur_p[k].p, llama_token_to_piece(ctx_dft, cur_p[k].id).c_str());
k, s, i, cur_p->data[k].id, cur_p->data[k].p, llama_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
}
std::vector<int> sa(1, s);
// attempt to split the branch if the probability is high enough
for (int f = 1; f < 8; ++f) {
if (n_seq_cur < n_seq_dft && cur_p[f].p > p_split) {
if (n_seq_cur < n_seq_dft && cur_p->data[f].p > p_split) {
LOG("splitting seq %3d into %3d\n", s, n_seq_cur);
llama_kv_cache_seq_rm(ctx_dft, n_seq_cur, -1, -1);
@ -503,7 +505,7 @@ int main(int argc, char ** argv) {
drafts[n_seq_cur].i_batch_dft = drafts[s].i_batch_dft;
drafts[n_seq_cur].i_batch_tgt = drafts[s].i_batch_tgt;
llama_sampling_cp(drafts[s].ctx_sampling, drafts[n_seq_cur].ctx_sampling);
llama_sampling_cp(drafts[s].smpl, drafts[n_seq_cur].smpl);
sa.push_back(n_seq_cur);
@ -515,15 +517,15 @@ int main(int argc, char ** argv) {
// add drafted token for each sequence
for (int is = 0; is < (int) sa.size(); ++is) {
const llama_token id = cur_p[is].id;
const llama_token id = cur_p->data[is].id;
const int s = sa[is];
llama_sampling_accept(drafts[s].ctx_sampling, ctx_dft, id, true);
llama_sampling_accept(drafts[s].smpl, id, true);
drafts[s].tokens.push_back(id);
// save cur_p.data into drafts[s].dists
drafts[s].dists.push_back(cur_p);
drafts[s].dists.push_back({cur_p->data, cur_p->data + cur_p->size});
// add unique drafted tokens to the target batch
drafts[s].i_batch_tgt.push_back(batch_tgt.n_tokens);
@ -594,14 +596,15 @@ int main(int argc, char ** argv) {
LOG_TEE("accept = %.3f%%\n", 100.0f * n_accept / n_drafted);
LOG_TEE("\ndraft:\n");
llama_print_timings(ctx_dft);
// TODO: print sampling/grammar timings for all drafts
llama_print_timings(ctx_dft, nullptr);
LOG_TEE("\ntarget:\n");
llama_print_timings(ctx_tgt);
llama_print_timings(ctx_tgt, smpl);
llama_sampling_free(ctx_sampling);
llama_sampling_free(smpl);
for (int s = 0; s < n_seq_dft; ++s) {
llama_sampling_free(drafts[s].ctx_sampling);
llama_sampling_free(drafts[s].smpl);
}
llama_batch_free(batch_dft);