New kv_cache API insufficient to restore model state

[abetlen]

I may be doing something wrong or misunderstanding the purpose of the kv_cache API but I believe the recent PR #685 by @chrfalch which added the ability to get / set the kv_cache is still insufficient to restore the state of the model even when resetting external model state such as last_n_tokens_data and n_past.

Here is a minimal example

#include "llama.h"
#include <vector>
#include <iostream>

using namespace std;

int main() {
    // init
    auto params = llama_context_default_params();
    auto ctx = llama_init_from_file("../../models/ggml-model.bin", params);
    auto tokens = vector<llama_token>(params.n_ctx);
    auto prompt = "The quick brown fox";
    auto n_tokens = llama_tokenize(ctx, prompt, tokens.data(), tokens.size(), true);

    // evaluate prompt
    llama_eval(ctx, tokens.data(), n_tokens, 0, 12);
    auto last_n_tokens_size = 64;
    auto last_n_tokens_data = vector<llama_token>(last_n_tokens_size, 0);
    last_n_tokens_data.insert(last_n_tokens_data.end(), tokens.data(), tokens.data() + n_tokens);
    auto n_past = n_tokens;

    // save state
    auto kv_cache_size = llama_get_kv_cache_size(ctx);
    auto kv_cache_token_count = llama_get_kv_cache_token_count(ctx);
    auto kv_cache = llama_get_kv_cache(ctx);
    auto kv_cache_copy = vector<uint8_t>(kv_cache, kv_cache + kv_cache_size);
    auto n_past_copy = n_past;
    auto last_n_tokens_data_copy = vector<llama_token>(last_n_tokens_data);
    
    // first run
    cout << prompt;
    for (auto i = 0; i < 6; i++) {
        auto next_token = llama_sample_top_p_top_k(
            ctx,
            last_n_tokens_data.data() + last_n_tokens_data.size() - n_past,
            last_n_tokens_size,
            1,
            1.0,
            0.0,
            1.1
        );
        auto next_token_str = llama_token_to_str(ctx, next_token);
        last_n_tokens_data.push_back(next_token);
        cout << next_token_str;
        llama_eval(ctx, &next_token, 1, n_past, 12);
        n_past += 1;
    }
    cout << endl;
    //

    // restore state
    llama_set_kv_cache(ctx, kv_cache_copy.data(), kv_cache_size, kv_cache_token_count);
    last_n_tokens_data = last_n_tokens_data_copy;
    n_past = n_past_copy;
    //

    // second run
    cout << prompt;
    for (auto i = 0; i < 6; i++) {
        auto next_token = llama_sample_top_p_top_k(
            ctx,
            last_n_tokens_data.data() + last_n_tokens_data.size() - n_past,
            last_n_tokens_size,
            1,
            1.0,
            0.0,
            1.1
        );
        auto next_token_str = llama_token_to_str(ctx, next_token);
        last_n_tokens_data.push_back(next_token);
        cout << next_token_str;
        llama_eval(ctx, &next_token, 1, n_past, 12);
        n_past += 1;
    }
    cout << endl;
    //
    return 0;
}

I'd expect the following output

The quick brown fox jumps over the lazy dog
The quick brown fox jumps over the lazy dog

But instead I get

The quick brown fox jumps over the lazy dog
The quick brown fox.
The quick brown fo

Which implies the model is still generating from the end of the first run.

[abetlen]

Woops, sorry I just realized you obviously still need to eval the prompt tokens again.

Here's is the working version for future reference.

#include "llama.h"
#include <vector>
#include <iostream>

using namespace std;

int main() {
    // init
    auto params = llama_context_default_params();
    auto ctx = llama_init_from_file("../../models/ggml-model.bin", params);
    auto tokens = vector<llama_token>(params.n_ctx);
    auto prompt = "The quick brown fox";
    auto n_tokens = llama_tokenize(ctx, prompt, tokens.data(), tokens.size(), true);

    // evaluate prompt
    llama_eval(ctx, tokens.data(), n_tokens, 0, 12);
    auto last_n_tokens_size = 64;
    auto last_n_tokens_data = vector<llama_token>(last_n_tokens_size, 0);
    last_n_tokens_data.insert(last_n_tokens_data.end(), tokens.data(), tokens.data() + n_tokens);
    auto n_past = n_tokens;

    // save state
    auto kv_cache_size = llama_get_kv_cache_size(ctx);
    auto kv_cache_token_count = llama_get_kv_cache_token_count(ctx);
    auto kv_cache = llama_get_kv_cache(ctx);
    auto kv_cache_copy = vector<uint8_t>(kv_cache, kv_cache + kv_cache_size);
    auto n_past_copy = n_past;
    auto last_n_tokens_data_copy = vector<llama_token>(last_n_tokens_data);
    
    // first run
    cout << prompt;
    for (auto i = 0; i < 6; i++) {
        auto next_token = llama_sample_top_p_top_k(
            ctx,
            last_n_tokens_data.data() + last_n_tokens_data.size() - n_past,
            last_n_tokens_size,
            1,
            1.0,
            0.0,
            1.1
        );
        auto next_token_str = llama_token_to_str(ctx, next_token);
        last_n_tokens_data.push_back(next_token);
        cout << next_token_str;
        llama_eval(ctx, &next_token, 1, n_past, 12);
        n_past += 1;
    }
    cout << endl;
    //

    // restore state
    llama_set_kv_cache(ctx, kv_cache_copy.data(), kv_cache_size, kv_cache_token_count);
    last_n_tokens_data = last_n_tokens_data_copy;
    n_past = n_past_copy;
    // call eval again on prompt tokens
    llama_eval(ctx, tokens.data(), n_tokens, 0, 12);
    //

    // second run
    cout << prompt;
    for (auto i = 0; i < 6; i++) {
        auto next_token = llama_sample_top_p_top_k(
            ctx,
            last_n_tokens_data.data() + last_n_tokens_data.size() - n_past,
            last_n_tokens_size,
            1,
            1.0,
            0.0,
            1.1
        );
        auto next_token_str = llama_token_to_str(ctx, next_token);
        last_n_tokens_data.push_back(next_token);
        cout << next_token_str;
        llama_eval(ctx, &next_token, 1, n_past, 12);
        n_past += 1;
    }
    cout << endl;
    //
    return 0;
}

[chrfalch]

It is correct that the pr does not implement this - but it describes that last tokens etc is is needed to save full state :) just wanted to implement the missing api for implementing the functionality for a prompt saving mechanism.

[abetlen]

@chrfalch sorry to bug you again on this one but I think I'm missing something.

From my understanding based on your response you should be able to save the internal state to disk assuming you also save n_past and last_n_tokens, however I'm still not able to do this correctly / in a way that reduces processing time once the model is reloaded.

1. llama_init_from_file the model.
2. Tokenize the initial prompt (e.g. the quick brown fox jumps)
3. llama_eval the prompt and set n_past to the number of prompt tokens, and start to fill last_n_tokens_data with the prompt tokens.
4. Save kv_cache, kv_cache_size, kv_cache_token_count, n_past and last_n_tokens.
5. Generate some tokens in a llama_sample_top_p_top_k / llama_eval loop (e.g over the lazy dog)
6. llama_free the context
7. Reload via llama_init_from_file
8. Restore the state via llama_set_kv_cache with the saved values from above.
9. Restore n_past and last_n_tokens to the saved value.
10. Generate some tokens in a llama_sample_top_p_top_k / llama_eval loop (e.g over the lazy dog) starting from the saved value of n_past and last_n_tokens

I would now expect to get the same output based on the original prompt e.g. over the lazy dog but it seems that the model is not taking this into account and instead I get back a random response.

Appreciate any help on this one, cheers.

EDIT: And just to clarify, if I call eval after restoring the kv_cache as I did above it doesn't seem to reduce processing time.

[ivanstepanovftw]

I dont think you need to eval initial prompt, because you'd wanted to avoid this

[abetlen]

@ivanstepanovftw do you have an example? I've tried not eval'ing but in that case even with n_past saved the model fails to generate the same output (just random generation).

[chrfalch]

A great test is to give the model a prompt saying its name, and then the test would be to ask it "What's your name" - if it responds with the correct name everything works.

Here is how I have implemented saving the "state" of the model:

1. Save kv_cache size by calling llama_get_kv_cache_size
2. Get memory in kv_cache by calling llama_get_kv_cache
3. Save token count by calling llama_get_kv_cache_token_count
4. Save n_past
5. Save last_n_tokens and its size

This is all you need. After you have eval'ed the prompt you do the above steps and save the results.

Then to restore you do the opposite - and test it all with the AI name prompt trick.

[abetlen]

@ivanstepanovftw if I don't eval anything after restoring the kv_state in a new context I just get a segfault (I assume some internal buffers only get initialised on first eval). See example below:

#include <vector>
#include <iostream>
#include <chrono>

#include "llama.h"

using namespace std;

int main() {
    auto seed = 42;
    auto thread_count = 12;
    auto last_n_tokens_size = 64;
    auto prompt = "The quick brown fox";
    auto model_path = "../models/ggml-alpaca.bin";

    auto n_past = 0;
    auto last_n_tokens_data = vector<llama_token>(last_n_tokens_size, 0);

    // init
    auto params = llama_context_default_params();
    params.seed = seed;
    auto ctx = llama_init_from_file(model_path, params);
    auto tokens = vector<llama_token>(params.n_ctx);
    auto n_prompt_tokens = llama_tokenize(ctx, prompt, tokens.data(), tokens.size(), true);

    if (n_prompt_tokens < 1) {
        cout << "Failed to tokenize prompt" << endl;
        return 1;
    }

    // evaluate prompt

    llama_eval(ctx, tokens.data(), n_prompt_tokens, n_past, thread_count);

    last_n_tokens_data.insert(last_n_tokens_data.end(), tokens.data(), tokens.data() + n_prompt_tokens);
    n_past += n_prompt_tokens;

    // save kv state, last n tokens and n_past
    auto kv_cache_size = llama_get_kv_cache_size(ctx);
    auto kv_cache_token_count = llama_get_kv_cache_token_count(ctx);
    auto kv_cache = llama_get_kv_cache(ctx);
    auto kv_cache_saved = vector<uint8_t>(kv_cache, kv_cache + kv_cache_size);

    auto last_n_tokens_data_saved = vector<llama_token>(last_n_tokens_data);
    auto n_past_saved = n_past;

    // first run
    cout << endl << prompt;
    for (auto i = 0; i < 6; i++) {
        auto next_token = llama_sample_top_p_top_k(
            ctx,
            &last_n_tokens_data.back() - last_n_tokens_size,
            last_n_tokens_size,
            40,
            1.0,
            1.0,
            1.1
        );
        auto next_token_str = llama_token_to_str(ctx, next_token);
        last_n_tokens_data.push_back(next_token);
        cout << next_token_str;
        if (llama_eval(ctx, &next_token, 1, n_past, thread_count)) {
            cout << endl << "Failed to evaluate" << endl;
            return 1;
        }
        n_past += 1;
    }
    cout << endl << endl;

    // free old model
    llama_free(ctx);

    // load new model
    params = llama_context_default_params();
    params.seed = seed;
    ctx = llama_init_from_file(model_path, params);

    // restore state
    llama_set_kv_cache(ctx, kv_cache_saved.data(), kv_cache_size, kv_cache_token_count);
    last_n_tokens_data = last_n_tokens_data_saved;
    n_past = n_past_saved;

    // second run
    cout << endl << prompt;
    for (auto i = 0; i < 6; i++) {
        auto next_token = llama_sample_top_p_top_k(
            ctx,
            &last_n_tokens_data.back() - last_n_tokens_size,
            last_n_tokens_size,
            40,
            1.0,
            1.0,
            1.1
        );
        auto next_token_str = llama_token_to_str(ctx, next_token);
        last_n_tokens_data.push_back(next_token);
        cout << next_token_str;
        if (llama_eval(ctx, &next_token, 1, n_past, thread_count)) {
            cout << endl << "Failed to evaluate" << endl;
            return 1;
        }
        n_past += 1;
    }
    cout << endl << endl;
    return 0;
}

To avoid this I also tried saving the first generated token and just evaling that with n_past = n_prompt_tokens but no luck I just get random output still.

@chrfalch I would imagine this should be the same as what you're suggesting, the model should "know" about the evaluated prompt throught the kv_state.

#include <vector>
#include <iostream>
#include <chrono>

#include "llama.h"

using namespace std;

int main() {
    auto seed = 42;
    auto thread_count = 12;
    auto last_n_tokens_size = 64;
    auto prompt = "The quick brown fox";
    auto model_path = "../models/ggml-alpaca.bin";

    auto n_past = 0;
    auto last_n_tokens_data = vector<llama_token>(last_n_tokens_size, 0);

    // init
    auto params = llama_context_default_params();
    params.seed = seed;
    auto ctx = llama_init_from_file(model_path, params);
    auto tokens = vector<llama_token>(params.n_ctx);
    auto n_prompt_tokens = llama_tokenize(ctx, prompt, tokens.data(), tokens.size(), true);

    if (n_prompt_tokens < 1) {
        cout << "Failed to tokenize prompt" << endl;
        return 1;
    }

    // evaluate prompt

    llama_eval(ctx, tokens.data(), n_prompt_tokens, n_past, thread_count);

    last_n_tokens_data.insert(last_n_tokens_data.end(), tokens.data(), tokens.data() + n_prompt_tokens);
    n_past += n_prompt_tokens;

    // save kv state, last n tokens and n_past
    auto kv_cache_size = llama_get_kv_cache_size(ctx);
    auto kv_cache_token_count = llama_get_kv_cache_token_count(ctx);
    auto kv_cache = llama_get_kv_cache(ctx);
    auto kv_cache_saved = vector<uint8_t>(kv_cache, kv_cache + kv_cache_size);

    auto last_n_tokens_data_saved = vector<llama_token>(last_n_tokens_data);
    auto n_past_saved = n_past;

    // save first generated token
    auto first_generated_token = llama_token(0);

    // first run
    cout << endl << prompt;
    for (auto i = 0; i < 6; i++) {
        auto next_token = llama_sample_top_p_top_k(
            ctx,
            &last_n_tokens_data.back() - last_n_tokens_size,
            last_n_tokens_size,
            40,
            1.0,
            1.0,
            1.1
        );
        if (i == 0) {
            first_generated_token = next_token;
        }
        auto next_token_str = llama_token_to_str(ctx, next_token);
        last_n_tokens_data.push_back(next_token);
        cout << next_token_str;
        if (llama_eval(ctx, &next_token, 1, n_past, thread_count)) {
            cout << endl << "Failed to evaluate" << endl;
            return 1;
        }
        n_past += 1;
    }
    cout << endl << endl;

    // free old model
    llama_free(ctx);

    // load new model
    params = llama_context_default_params();
    params.seed = seed;
    ctx = llama_init_from_file(model_path, params);

    // restore state
    llama_set_kv_cache(ctx, kv_cache_saved.data(), kv_cache_size, kv_cache_token_count);
    last_n_tokens_data = last_n_tokens_data_saved;
    n_past = n_past_saved;

    // restore first generated token so we can safely sample
    llama_eval(
        ctx,
        &first_generated_token,
        1,
        n_past,
        thread_count
    );
    last_n_tokens_data.push_back(first_generated_token);
    n_past += 1;

    // second run
    cout << endl << prompt << llama_token_to_str(ctx, first_generated_token);
    for (auto i = 0; i < 5; i++) {
        auto next_token = llama_sample_top_p_top_k(
            ctx,
            &last_n_tokens_data.back() - last_n_tokens_size,
            last_n_tokens_size,
            40,
            1.0,
            1.0,
            1.1
        );
        auto next_token_str = llama_token_to_str(ctx, next_token);
        last_n_tokens_data.push_back(next_token);
        cout << next_token_str;
        if (llama_eval(ctx, &next_token, 1, n_past, thread_count)) {
            cout << endl << "Failed to evaluate" << endl;
            return 1;
        }
        n_past += 1;
    }
    cout << endl << endl;
    return 0;
}