From d2fd589f514b13d2526fd9467f1ae53bff2aca46 Mon Sep 17 00:00:00 2001 From: Xuan Son Nguyen Date: Sun, 16 Aug 2026 21:06:08 +0200 Subject: [PATCH] training: fix no KV cache --- examples/training/README.md | 2 + src/llama-context.cpp | 20 ++++++- src/llama-cparams.h | 2 + src/llama-graph.cpp | 109 +++++++++++++++++++++++++++++++----- src/llama-graph.h | 4 ++ 5 files changed, 122 insertions(+), 15 deletions(-) diff --git a/examples/training/README.md b/examples/training/README.md index df42527926..146e4154eb 100644 --- a/examples/training/README.md +++ b/examples/training/README.md @@ -6,6 +6,8 @@ Finetuning of Stories 260K and LLaMA 3.2 1b seems to work with 24 GB of memory. **For CPU training, compile llama.cpp without any additional backends such as CUDA.** **For CUDA training, use the maximum number of GPU layers.** +Gradients cannot flow through the KV cache, so training attends to the tokens of the current ubatch instead. This requires `-c`, `-b` and `-ub` to be equal. Flash attention is disabled during training because `FLASH_ATTN_EXT` has no backward pass. + Proof of concept: ``` sh diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 792bb6cac8..e62a02d77a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -482,7 +482,7 @@ llama_context::~llama_context() { // wait for any pending asynchronous copies into the output buffers before they are freed synchronize(); - if (!model.hparams.no_alloc) { + if (!model.hparams.no_alloc && !cparams.training) { for (size_t i = 0; i < backend_ptrs.size(); ++i) { ggml_backend_t backend = backend_ptrs[i]; ggml_backend_buffer_type_t buft = backend_buft[i]; @@ -2331,6 +2331,13 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { if (n_sampling_outputs_max > 1) { res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max; } + + if (cparams.training) { + // backward pass for ggml_opt + // TODO: maybe improve this later + res *= 4; + } + return res; } @@ -3315,6 +3322,17 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params GGML_ASSERT(model->hparams.n_ctx_train % n_batch == 0); GGML_ASSERT(n_batch % n_ubatch == 0); + cparams.training = true; + + if (cparams.flash_attn) { + LLAMA_LOG_INFO("%s: disabling flash attention, FLASH_ATTN_EXT has no backward pass\n", __func__); + cparams.flash_attn = false; + } + + // the graph shape and size change when training, need reserve again + sched_need_reserve = true; + sched_reserve(); + ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY); opt_params.opt_period = n_batch / n_ubatch; opt_params.get_opt_pars = lopt_params.get_opt_pars; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 574ce95920..15bc6a1691 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -54,6 +54,8 @@ struct llama_cparams { bool kv_unified; bool pipeline_parallel; + bool training = false; // set by llama_opt_init(), see [TAG_TRAINING_NO_KV_CACHE] + std::vector embeddings_layer_inp; // [n_layer()] extract input embeddings for layer enum llama_context_type ctx_type; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 1896758c5d..da32d2ba2c 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -405,6 +405,54 @@ static void print_mask(const T * data, int64_t n_tokens, int64_t n_kv, int64_t n } } +static void fill_kq_mask_self( + ggml_tensor * mask, + const llama_ubatch * ubatch, + const llama_hparams & hparams, + const llama_cparams & cparams) { + GGML_ASSERT(mask); + GGML_ASSERT(ggml_backend_buffer_is_host(mask->buffer)); + + const int64_t n_tokens = ubatch->n_tokens; + + GGML_ASSERT(mask->ne[0] == n_tokens && mask->ne[1] == n_tokens); + + const auto fill = [&](auto * data) { + using T = std::remove_reference_t; + std::fill(data, data + ggml_nelements(mask), llama_cast(-INFINITY)); + + for (int64_t i1 = 0; i1 < n_tokens; ++i1) { + const llama_seq_id s1 = ubatch->seq_id[i1][0]; + const llama_pos p1 = ubatch->pos[i1]; + + for (int64_t i0 = 0; i0 < n_tokens; ++i0) { + const llama_seq_id s0 = ubatch->seq_id[i0][0]; + const llama_pos p0 = ubatch->pos[i0]; + + if (s0 != s1) { + continue; + } + + if (cparams.causal_attn && p0 > p1) { + continue; + } + + if (llama_hparams::is_masked_swa(hparams.n_swa, hparams.swa_type, p0, p1)) { + continue; + } + + data[i1*n_tokens + i0] = llama_cast(hparams.use_alibi ? -std::abs(p0 - p1) : 0.0f); + } + } + }; + + if (mask->type == GGML_TYPE_F16) { + fill((ggml_fp16_t *) mask->data); + } else { + fill((float *) mask->data); + } +} + void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) { const int64_t n_kv = ubatch->n_tokens; const int64_t n_tokens = ubatch->n_tokens; @@ -467,6 +515,12 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) { } void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) { + // the cache inputs do not exist when training, see [TAG_TRAINING_NO_KV_CACHE] + if (self_kq_mask_train) { + fill_kq_mask_self(self_kq_mask_train, ubatch, hparams, cparams); + return; + } + mctx->set_input_k_idxs(self_k_idxs, ubatch); mctx->set_input_v_idxs(self_v_idxs, ubatch); @@ -490,6 +544,11 @@ bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) { this->mctx = mctx; + // [TAG_TRAINING_NO_KV_CACHE] + if (self_kq_mask_train) { + return self_kq_mask_train->ne[0] == params.ubatch.n_tokens; + } + bool res = true; res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; @@ -2735,9 +2794,19 @@ static std::unique_ptr build_attn_inp_kv_impl( auto inp = std::make_unique(hparams, cparams, mctx_cur); - { - GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA"); + GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA"); + // [TAG_TRAINING_NO_KV_CACHE] the cache is not used, so only build the mask + if (cparams.training) { + const auto type_mask = cparams.flash_attn ? GGML_TYPE_F16 : GGML_TYPE_F32; + + inp->self_kq_mask_train = ggml_new_tensor_4d(ctx0, type_mask, ubatch.n_tokens, ubatch.n_tokens, 1, 1); + ggml_set_input(inp->self_kq_mask_train); + + return inp; + } + + { inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch); inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch); @@ -2792,20 +2861,32 @@ ggml_tensor * llm_graph_context::build_attn( const auto * mctx_cur = inp->mctx; - // store to KV cache - { - const auto & k_idxs = inp->get_k_idxs(); - const auto & v_idxs = inp->get_v_idxs(); - - ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); - ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); - } - - ggml_tensor * kq_mask = inp->get_kq_mask(); + ggml_tensor * kq_mask; ggml_tensor * q = q_cur; - ggml_tensor * k = mctx_cur->get_k(ctx0, il); - ggml_tensor * v = mctx_cur->get_v(ctx0, il); + ggml_tensor * k; + ggml_tensor * v; + + if (cparams.training) { + // [TAG_TRAINING_NO_KV_CACHE] + // when training, we don't use cache + kq_mask = inp->get_kq_mask_train(); + k = k_cur; + v = v_cur; + } else { + // store to KV cache + { + const auto & k_idxs = inp->get_k_idxs(); + const auto & v_idxs = inp->get_v_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + } + + kq_mask = inp->get_kq_mask(); + k = mctx_cur->get_k(ctx0, il); + v = mctx_cur->get_v(ctx0, il); + } ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il); diff --git a/src/llama-graph.h b/src/llama-graph.h index 94324c7457..f2237ae1dd 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -337,12 +337,16 @@ public: ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; } + ggml_tensor * get_kq_mask_train() const { return self_kq_mask_train; } + ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch] ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa] ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask_train = nullptr; // F32/F16 [n_tokens, n_tokens], see [TAG_TRAINING_NO_KV_CACHE] + // note: assumes v_rot^2 == I ggml_tensor * self_k_rot = nullptr; ggml_tensor * self_v_rot = nullptr;