Skip to content

Commit 65ef50a

Browse files
ggerganovggerganov
authored
tests : refactor test-save-load-state to accept token input (#24073)
* tests : refactor test-save-load-state to accept token input - Default prompt is now empty; when not provided, generate n_batch random tokens (useful for models without a tokenizer) - Tokenization happens once upfront; pass token vector to test functions - generate_tokens prints token IDs instead of decoded pieces - Use llama_model_get_vocab / llama_vocab_n_tokens API - Upgrade log level from LOG_TRC to LOG_INF for visibility Assisted-by: llama.cpp:local pi * cont : use llama_tokens alias
1 parent 3d19986 commit 65ef50a

1 file changed

Lines changed: 44 additions & 31 deletions

File tree

‎tests/test-save-load-state.cpp‎

Lines changed: 44 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#include "llama-cpp.h"
55

66
#include <clocale>
7+
#include <random>
78
#include <vector>
89

910
struct llama_batch_ptr {
@@ -23,16 +24,15 @@ struct llama_batch_ptr {
2324
const llama_batch & get() const { return batch; }
2425
};
2526

26-
static std::string generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) {
27-
std::string result;
27+
static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) {
28+
llama_tokens result;
2829
llama_batch_ptr batch(1, 0, 1);
2930

3031
for (int i = 0; i < n_predict; i++) {
31-
auto next_token = llama_sampler_sample(smpl, ctx, -1);
32-
auto next_token_str = common_token_to_piece(ctx, next_token);
32+
auto next_token = llama_sampler_sample(smpl, ctx, -1);
3333

34-
LOG("%s", next_token_str.c_str());
35-
result += next_token_str;
34+
LOG("%d ", next_token);
35+
result.push_back(next_token);
3636

3737
common_batch_clear(batch.get());
3838
common_batch_add(batch.get(), next_token, n_past, {seq_id}, true);
@@ -48,28 +48,24 @@ static std::string generate_tokens(llama_context * ctx, llama_sampler * smpl, in
4848
}
4949

5050
// Test 1: baseline
51-
// - tokenize the prompt
5251
// - decode all but the last token
5352
// - save state to disk
5453
// - decode the last token
5554
// - generate n_predict tokens
56-
static std::string test_baseline(struct llama_model * model, const struct common_params & params) {
55+
static llama_tokens test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
5756
auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
5857

5958
auto sparams = llama_sampler_chain_default_params();
6059
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
6160
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
6261

63-
auto tokens = common_tokenize(ctx.get(), params.prompt, true);
64-
6562
auto n_past = 0;
6663
if (!common_prompt_batch_decode(ctx.get(), tokens, (int)tokens.size(), n_past, params.n_batch, params.out_file, true)) {
6764
LOG_ERR("%s: failed to decode prompt\n", __func__);
6865
return {};
6966
}
7067

7168
LOG("\n=== Test 1: baseline ===\n");
72-
LOG("%s", params.prompt.c_str());
7369

7470
auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
7571
if (result.empty()) {
@@ -87,20 +83,17 @@ static std::string test_baseline(struct llama_model * model, const struct common
8783
// - load state from file
8884
// - replay the last prompt token
8985
// - generate n_predict tokens and compare against expected result
90-
static bool test_state_load(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
86+
static bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) {
9187
auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
9288

9389
auto sparams = llama_sampler_chain_default_params();
9490
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
9591
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
9692

97-
auto tokens = common_tokenize(ctx.get(), params.prompt, true);
98-
9993
LOG("\n=== Test 2: state load ===\n");
100-
LOG("%s", params.prompt.c_str());
10194

10295
// Load state from file
103-
std::vector<llama_token> unused_sts(tokens.size());
96+
llama_tokens unused_sts(tokens.size());
10497
size_t n_token_count_out = 0;
10598

10699
if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
@@ -139,7 +132,7 @@ static bool test_state_load(struct llama_model * model, const struct common_para
139132
// - replay the last prompt token
140133
// - migrate KV cache from seq 0 to seq 1 via the CPU path
141134
// - generate n_predict tokens on seq 1 and compare against expected result
142-
static bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
135+
static bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) {
143136
auto params_ctx = common_context_params_to_llama(params);
144137
params_ctx.n_seq_max = 2;
145138
auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
@@ -148,13 +141,10 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par
148141
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
149142
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
150143

151-
auto tokens = common_tokenize(ctx.get(), params.prompt, true);
152-
153144
LOG("\n=== Test 3: seq copy (host) ===\n");
154-
LOG("%s", params.prompt.c_str());
155145

156146
// Load state from file
157-
std::vector<llama_token> unused_sts(tokens.size());
147+
llama_tokens unused_sts(tokens.size());
158148
size_t n_token_count_out = 0;
159149

160150
if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
@@ -214,7 +204,7 @@ static bool test_seq_cp_host(struct llama_model * model, const struct common_par
214204
// - replay the last prompt token
215205
// - migrate KV cache from seq 0 to seq 1 via the on-device path
216206
// - generate n_predict tokens on seq 1 and compare against expected result
217-
static bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
207+
static bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) {
218208
auto params_ctx = common_context_params_to_llama(params);
219209
params_ctx.n_seq_max = 2;
220210
auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
@@ -223,13 +213,10 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p
223213
auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
224214
llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
225215

226-
auto tokens = common_tokenize(ctx.get(), params.prompt, true);
227-
228216
LOG("\n=== Test 4: seq copy (device) ===\n");
229-
LOG("%s", params.prompt.c_str());
230217

231218
// Load state from file
232-
std::vector<llama_token> unused_sts(tokens.size());
219+
llama_tokens unused_sts(tokens.size());
233220
size_t n_token_count_out = 0;
234221

235222
if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
@@ -287,7 +274,8 @@ int main(int argc, char ** argv) {
287274
std::setlocale(LC_NUMERIC, "C");
288275

289276
common_params params;
290-
params.prompt = "The quick brown fox";
277+
params.prompt = "";
278+
params.n_batch = 100;
291279
params.out_file = "dump_state.bin";
292280
params.sampling.seed = 1234;
293281

@@ -318,24 +306,49 @@ int main(int argc, char ** argv) {
318306

319307
GGML_ASSERT(llama_init->context() == nullptr);
320308

309+
// Tokenize prompt or generate random tokens
310+
llama_tokens tokens;
311+
if (params.prompt.empty()) {
312+
const int n_prompt = params.n_batch;
313+
314+
// this path is useful for model files that do not have a tokenizer
315+
LOG_INF("%s: no prompt provided, generating %d (n_batch) random tokens\n", __func__, n_prompt);
316+
317+
const auto * vocab = llama_model_get_vocab(model);
318+
const auto n_vocab = llama_vocab_n_tokens(vocab);
319+
320+
std::mt19937 rng(params.sampling.seed);
321+
std::uniform_int_distribution<llama_token> dist(0, n_vocab - 1);
322+
for (int i = 0; i < n_prompt; i++) {
323+
tokens.push_back(dist(rng));
324+
}
325+
} else {
326+
LOG_INF("%s: tokenizing prompt '%s'\n", __func__, params.prompt.c_str());
327+
328+
auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
329+
tokens = common_tokenize(ctx.get(), params.prompt, true);
330+
}
331+
332+
LOG_INF("%s: the input prompt is %d tokens\n", __func__, (int)tokens.size());
333+
321334
// Test 1: baseline (saves state to disk)
322-
auto result_baseline = test_baseline(model, params);
335+
auto result_baseline = test_baseline(model, params, tokens);
323336
if (result_baseline.empty()) {
324337
return 1;
325338
}
326339

327340
// Test 2: state load
328-
if (!test_state_load(model, params, result_baseline)) {
341+
if (!test_state_load(model, params, tokens, result_baseline)) {
329342
return 1;
330343
}
331344

332345
// Test 3: seq copy (host)
333-
if (!test_seq_cp_host(model, params, result_baseline)) {
346+
if (!test_seq_cp_host(model, params, tokens, result_baseline)) {
334347
return 1;
335348
}
336349

337350
// Test 4: seq copy (device)
338-
if (!test_seq_cp_device(model, params, result_baseline)) {
351+
if (!test_seq_cp_device(model, params, tokens, result_baseline)) {
339352
return 1;
340353
}
341354

0 commit comments

Comments
 (0)