llama + spec: MTP Support (#22673)
* spec: support MTP * fix batch size * rename files * cont : simplify (#7) * MTP: clean-up (#9) * MTP: clean-up * review: use llama_context_type instead of llama_graph_type * review: remove llama_model_has_mtp * review: fix convert issues * convert: fix pycheck * review: formatting * use `mtp-` for identifying mtp models * convert: fix mtp conversion * mtp -> draft-mtp * remove unused llama_arch * add need_embd in speculative * llama: allow partial seq_rm for GDN models for speculative decoding Currently speculative checkpoint needs to restart from a checkpoint after some draft tokens are not accepted, this leads to some wastage in running the target again. This PR adds the ability to rollback upto `draft_max` by storing the GDN intermediates. * fix pending state * vulkan: add GDN partial rollback * meta: extend check to axis 1 * metal: add GDN partial rollback Extend the gated delta net kernel to store intermediate states for partial rollback support on the Metal backend. - Add K (snapshot slot count) as a function constant - Read input state from slot 0 of the 3D state tensor - Write intermediate states to different slots during token loop - For K=1, maintain backward-compatible single-slot behavior Ref: https://github.com/ggml-org/llama.cpp/commit/8c05923630110223669f069af2000e9cf10c02bc Assisted-by: llama.cpp:local pi * delta_net_base: use ggml_pad instead of new_tensor * review: add need_rs_seq * review: rename part_bounded to n_rs * review: deslop comments * review: rename, add asserts * server : adjust checkpoint logic (#11) * server : adjust checkpoint logic * cont : rm asserts * server-context: fix early exit * spec : fix compatibility with n-gram and add TODOs (#13) * metal : cleanup * llama : fix faulty bitwise check in recurrent memory * server : disable RS-based MTP in combination with other spec types * spec : add TODOs * cont : fix comment * cont : update comment * common : fix logic for ngram + mtp compat * llama-memory: enable checkpointing with partial rollback * cont: add test-case for loading into a dirty ctx * llama-memory-recurrent: clear rs_idx in clear * download: fix mtp path * llama-arch: fix enorm op * docs: update docs * conversion: fix type annotations --------- Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
Aman Gupta
co-authored by
Georgi Gerganov
Georgi Gerganov
parent
b81c2cdd74
commit
255582687b
@@ -252,6 +252,9 @@ llama_build_and_test(test-backend-sampler.cpp LABEL "model")
|
||||
llama_build_and_test(test-state-restore-fragmented.cpp LABEL "model" ARGS -m "${MODEL_DEST}")
|
||||
set_tests_properties(test-state-restore-fragmented PROPERTIES FIXTURES_REQUIRED test-download-model)
|
||||
|
||||
llama_build_and_test(test-recurrent-state-rollback.cpp LABEL "model" ARGS -m "${MODEL_DEST}")
|
||||
set_tests_properties(test-recurrent-state-rollback PROPERTIES FIXTURES_REQUIRED test-download-model)
|
||||
|
||||
if (NOT GGML_BACKEND_DL)
|
||||
# these tests use the backends directly and cannot be built with dynamic loading
|
||||
llama_build_and_test(test-barrier.cpp)
|
||||
|
||||
@@ -3832,16 +3832,17 @@ struct test_gated_delta_net : public test_case {
|
||||
const int v_repeat;
|
||||
const bool permuted;
|
||||
const bool kda;
|
||||
const int64_t K; // snapshot slot count: 1 = final-only, >1 = last K states
|
||||
|
||||
std::string vars() override {
|
||||
return VARS_TO_STR8(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda);
|
||||
return VARS_TO_STR9(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K);
|
||||
}
|
||||
|
||||
test_gated_delta_net(ggml_type type = GGML_TYPE_F32,
|
||||
int64_t head_count = 4, int64_t head_size = 16, int64_t n_seq_tokens = 1, int64_t n_seqs = 1,
|
||||
int v_repeat = 1, bool permuted = false, bool kda = false)
|
||||
int v_repeat = 1, bool permuted = false, bool kda = false, int64_t K = 1)
|
||||
: type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs),
|
||||
v_repeat(v_repeat), permuted(permuted), kda(kda) {}
|
||||
v_repeat(v_repeat), permuted(permuted), kda(kda), K(K) {}
|
||||
|
||||
ggml_tensor * build_graph(ggml_context * ctx) override {
|
||||
ggml_tensor * q;
|
||||
@@ -3863,7 +3864,7 @@ struct test_gated_delta_net : public test_case {
|
||||
const int64_t g_ne0 = kda ? head_size : 1;
|
||||
ggml_tensor * g = ggml_new_tensor_4d(ctx, type, g_ne0, head_count * v_repeat, n_seq_tokens, n_seqs);
|
||||
ggml_tensor * beta = ggml_new_tensor_4d(ctx, type, 1, head_count * v_repeat, n_seq_tokens, n_seqs);
|
||||
ggml_tensor * state = ggml_new_tensor_2d(ctx, type, head_size * v_repeat * head_size * head_count, n_seqs);
|
||||
ggml_tensor * state = ggml_new_tensor_3d(ctx, type, head_size * v_repeat * head_size * head_count, K, n_seqs);
|
||||
ggml_set_name(g, "g");
|
||||
ggml_set_name(beta, "beta");
|
||||
ggml_set_name(state, "state");
|
||||
@@ -9034,6 +9035,18 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 33, 1, 1, false, true));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 100, 1, 1, false, true));
|
||||
|
||||
// K > 1: output keeps the last min(n_tokens, K) per-token snapshots in the trailing K-token region.
|
||||
// exact-match cases (K == n_seq_tokens):
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 16, 2, 1, 1, false, false, /*K=*/2));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 4, 1, 1, false, false, /*K=*/4));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, false, /*K=*/4));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 4, 1, 1, false, false, /*K=*/4));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, true, /*K=*/4));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 32, 4, 2, 2, false, true, /*K=*/4));
|
||||
// overflow: n_tokens > K — only the last K snapshots kept.
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, /*K=*/3));
|
||||
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 16, 2, 1, false, false, /*K=*/4));
|
||||
|
||||
#if 0
|
||||
// these tests are disabled to save execution time, sbut they can be handy for debugging
|
||||
test_cases.emplace_back(new test_llama(2, true));
|
||||
|
||||
@@ -0,0 +1,185 @@
|
||||
#include "arg.h"
|
||||
#include "common.h"
|
||||
#include "llama.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <clocale>
|
||||
#include <cmath>
|
||||
#include <cstdio>
|
||||
#include <vector>
|
||||
|
||||
static llama_context * make_ctx(const common_params & params, llama_model * model) {
|
||||
auto cparams = common_context_params_to_llama(params);
|
||||
cparams.n_seq_max = 1;
|
||||
cparams.n_rs_seq = 8;
|
||||
cparams.n_batch = std::max(cparams.n_batch, (uint32_t) (cparams.n_rs_seq + 1));
|
||||
cparams.n_ubatch = std::max(cparams.n_ubatch, (uint32_t) (cparams.n_rs_seq + 1));
|
||||
return llama_init_from_model(model, cparams);
|
||||
}
|
||||
|
||||
static bool decode_tokens(llama_context * ctx, const std::vector<llama_token> & tokens, uint32_t count) {
|
||||
llama_batch batch = llama_batch_init(count, 0, 1);
|
||||
for (uint32_t pos = 0; pos < count; ++pos) {
|
||||
common_batch_add(batch, tokens[pos], pos, { 0 }, false);
|
||||
}
|
||||
const bool ok = llama_decode(ctx, batch) == 0;
|
||||
llama_batch_free(batch);
|
||||
return ok;
|
||||
}
|
||||
|
||||
static bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos) {
|
||||
llama_batch batch = llama_batch_init(1, 0, 1);
|
||||
common_batch_add(batch, tok, pos, { 0 }, true);
|
||||
const bool ok = llama_decode(ctx, batch) == 0;
|
||||
llama_batch_free(batch);
|
||||
return ok;
|
||||
}
|
||||
|
||||
int main(int argc, char ** argv) {
|
||||
std::setlocale(LC_NUMERIC, "C");
|
||||
|
||||
common_params params;
|
||||
params.sampling.seed = 1234;
|
||||
params.n_predict = 1;
|
||||
|
||||
common_init();
|
||||
|
||||
if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
ggml_backend_load_all();
|
||||
|
||||
common_init_result_ptr llama_init = common_init_from_params(params);
|
||||
llama_model * model = llama_init->model();
|
||||
if (model == nullptr) {
|
||||
fprintf(stderr, "%s : failed to init model\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (!llama_model_is_recurrent(model) && !llama_model_is_hybrid(model)) {
|
||||
fprintf(stderr, "%s : skipping for non-recurrent model\n", __func__);
|
||||
return 0;
|
||||
}
|
||||
|
||||
const llama_vocab * vocab = llama_model_get_vocab(model);
|
||||
const int n_vocab = llama_vocab_n_tokens(vocab);
|
||||
|
||||
llama_context * ctx_src = make_ctx(params, model);
|
||||
llama_context * ctx_dst = make_ctx(params, model);
|
||||
if (ctx_src == nullptr || ctx_dst == nullptr) {
|
||||
fprintf(stderr, "%s : failed to init contexts\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
if (llama_n_rs_seq(ctx_src) == 0) {
|
||||
fprintf(stderr, "%s : skipping because n_rs_seq is disabled\n", __func__);
|
||||
llama_free(ctx_src);
|
||||
llama_free(ctx_dst);
|
||||
return 0;
|
||||
}
|
||||
|
||||
std::vector<llama_token> tokens = common_tokenize(ctx_src, "The quick brown fox jumps", true);
|
||||
const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src);
|
||||
if (tokens.size() > n_rs_seq + 1) {
|
||||
tokens.resize(n_rs_seq + 1);
|
||||
}
|
||||
if (tokens.size() < 2) {
|
||||
fprintf(stderr, "%s : not enough prompt tokens\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
const uint32_t n_tokens = tokens.size();
|
||||
const llama_token last_tok = tokens.back();
|
||||
const llama_pos last_pos = (llama_pos) n_tokens - 2;
|
||||
|
||||
// Decode the full prompt on the source, then roll back the last position.
|
||||
// Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0).
|
||||
if (!decode_tokens(ctx_src, tokens, n_tokens)) {
|
||||
fprintf(stderr, "%s : failed to decode prompt\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, last_pos, -1)) {
|
||||
fprintf(stderr, "%s : rollback failed\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
// Save the rolled-back state and restore it into a fresh context.
|
||||
common_prompt_checkpoint ckpt;
|
||||
ckpt.update_tgt(ctx_src, 0, 0);
|
||||
ckpt.load_tgt(ctx_dst, 0, 0);
|
||||
|
||||
// Replay the rolled-back token on both contexts and compare logits.
|
||||
if (!decode_one(ctx_src, last_tok, last_pos) ||
|
||||
!decode_one(ctx_dst, last_tok, last_pos)) {
|
||||
fprintf(stderr, "%s : replay failed\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
const float * logits_src = llama_get_logits_ith(ctx_src, 0);
|
||||
const float * logits_dst = llama_get_logits_ith(ctx_dst, 0);
|
||||
if (logits_src == nullptr || logits_dst == nullptr) {
|
||||
fprintf(stderr, "%s : missing logits\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
constexpr float eps = 1e-5f;
|
||||
for (int i = 0; i < n_vocab; ++i) {
|
||||
if (std::fabs(logits_src[i] - logits_dst[i]) > eps) {
|
||||
fprintf(stderr, "%s : logits mismatch at token %d (%g != %g)\n",
|
||||
__func__, i, (double) logits_src[i], (double) logits_dst[i]);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
// Repeat the load into a context that already has its own rollback state:
|
||||
// groups 1..n_rs_seq hold a *different* prompt's history, and rs_idx[0] is
|
||||
// non-zero at load time. The restore must wipe that state and still match.
|
||||
llama_context * ctx_dirty = make_ctx(params, model);
|
||||
if (ctx_dirty == nullptr) {
|
||||
fprintf(stderr, "%s : failed to init dirty ctx\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
std::vector<llama_token> noise = tokens;
|
||||
for (auto & t : noise) {
|
||||
t = (t + 1) % n_vocab;
|
||||
if (t < 0) {
|
||||
t = 0;
|
||||
}
|
||||
}
|
||||
if (!decode_tokens(ctx_dirty, noise, n_tokens)) {
|
||||
fprintf(stderr, "%s : dirty prompt decode failed\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, last_pos, -1)) {
|
||||
fprintf(stderr, "%s : dirty rollback failed\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
ckpt.load_tgt(ctx_dirty, 0, 0);
|
||||
|
||||
if (!decode_one(ctx_dirty, last_tok, last_pos)) {
|
||||
fprintf(stderr, "%s : dirty replay failed\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0);
|
||||
if (logits_dirty == nullptr) {
|
||||
fprintf(stderr, "%s : missing dirty logits\n", __func__);
|
||||
return 1;
|
||||
}
|
||||
|
||||
for (int i = 0; i < n_vocab; ++i) {
|
||||
if (std::fabs(logits_src[i] - logits_dirty[i]) > eps) {
|
||||
fprintf(stderr, "%s : dirty-ctx logits mismatch at token %d (%g != %g)\n",
|
||||
__func__, i, (double) logits_src[i], (double) logits_dirty[i]);
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
|
||||
fprintf(stderr, "%s : recurrent rollback checkpoint restored successfully\n", __func__);
|
||||
llama_free(ctx_src);
|
||||
llama_free(ctx_dst);
|
||||
llama_free(ctx_dirty);
|
||||
return 0;
|
||||
}
|
||||
Reference in New Issue
Block a user