From 159b741427337a2e9a58b08121001545d66b5825 Mon Sep 17 00:00:00 2001
From: Xuan-Son Nguyen <son@huggingface.co>
Date: Wed, 2 Sep 2026 23:53:32 +0200
Subject: [PATCH] finetune: fix no KV cache (#27199)

* training: fix no KV cache

* apply @ ggerganov
 suggestion
---
 examples/training/README.md |  2 ++
 ggml/src/ggml.c             |  2 +-
 src/llama-context.cpp       | 12 +++++++++++-
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/examples/training/README.md b/examples/training/README.md
index df42527926..526ac258fc 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.**
 
+Flash attention is disabled during training because `FLASH_ATTN_EXT` has no backward pass.
+
 Proof of concept:
 
 ``` sh
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
index 8dc0945084..2d5fdb7c10 100644
--- a/ggml/src/ggml.c
+++ b/ggml/src/ggml.c
@@ -7335,7 +7335,7 @@ void ggml_build_backward_expand(
         }
 
         // inplace operations are currently not supported
-        GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW ||
+        GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_SET_ROWS || node->op == GGML_OP_VIEW ||
             node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE);
 
         const size_t ihash = ggml_hash_find(&cgraph->visited_hash_set, node);
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index f286bd1da2..3cc27717ec 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -482,7 +482,8 @@ llama_context::~llama_context() {
     // wait for any pending asynchronous copies into the output buffers before they are freed
     synchronize();
 
-    if (!model.hparams.no_alloc) {
+    // when training, ggml_opt allocates extra buffers through the scheduler, so the sizes no longer match the expectation
+    if (!model.hparams.no_alloc && !opt_ctx) {
         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];
@@ -3408,6 +3409,15 @@ 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);
 
+    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 changes without flash attention, need to 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;