230 lines
7.0 KiB
C++
230 lines
7.0 KiB
C++
#include "sampling.h"
|
|
#include "llama.h"
|
|
#include "common.h"
|
|
#include "arg.h"
|
|
#include "log.h"
|
|
|
|
#include <cmath>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <optional>
|
|
#include <random>
|
|
#include <unordered_map>
|
|
|
|
// 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<float> uni(0.0f, 1.0f);
|
|
|
|
std::unordered_map<llama_token, int> token_counts;
|
|
|
|
std::vector<llama_token_data> 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<std::vector<llama_token_data>> draft_q = { q };
|
|
const std::vector<int> idxs = { -1, -1 };
|
|
|
|
const std::vector<llama_token> 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<char *> argv;
|
|
};
|
|
|
|
static std::optional<parsed_args> 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;
|
|
}
|