From bdd3cd55be4965ae870dac4d9e35646070c71344 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Tue, 13 Oct 2026 15:37:51 -0500 Subject: [PATCH 01/43] train: one rule for the recurrent states a graph keeps, so training runs in a context with rollback slots Serving keeps recurrent rollback slots for drafts (n_rs_seq > 0); the exact walk refused any context that had them, which is what forced a second, rollback-free training context beside serving. llm_graph_context::rs_keep() is now the one rule every family's state write and every snapshot-producing op reads (delta-net, mamba, lfm2, bailing; each decided it separately before): serving with rollback keeps n_rs_seq + 1 states, a training graph keeps one, the live state in group 0, which is the same bytes in both layouts. The walk now refuses only a sequence with an undo PENDING (llama_memory_recurrent::rollback_pending), and refuses by name a model whose state write does not report its exit to the walk, where the gradient across chunks was silently dropped before. Measured on the M5 (Metal), base 178224903: - serving unchanged: 19 recurrent/hybrid archs x n_rs_seq {0, 4}, random weights, 100-token prefill + 4 single decodes + (where rollback is live: qwen35, qwen35moe, deepseek4, nemotron_h, nemotron_h_moe, bailingmoe3, lfm2, lfm2moe) a 2-token undo and redo: logits and full sequence state byte-identical to the base, 38 of 38 rows. - training: test-walk-exact on qwen3.5-4b Q4_K_M, the exact walk in a context with n_rs_seq 4 vs 0: cosine 1.0000, norm ratio 1.0000; every other check in the file unchanged (exit 0). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- src/llama-context-opt.cpp | 21 ++++++++++++++------ src/llama-graph.cpp | 5 +++++ src/llama-graph.h | 9 +++++++++ src/llama-memory-recurrent.h | 6 ++++++ src/models/bailingmoe3.cpp | 9 ++++----- src/models/delta-net-base.cpp | 10 +++++----- src/models/lfm2.cpp | 2 +- src/models/mamba-base.cpp | 2 +- tests/test-walk-exact.cpp | 37 +++++++++++------------------------ 9 files changed, 57 insertions(+), 44 deletions(-) diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 3c183f3d255c..6cdd89002c4e 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -317,6 +317,14 @@ bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { } } for (const auto & ws : res->t_walk_state) { + // a state read without its write reporting the exit (build_walk_state_exit) would drop + // the gradient later chunks put on it: a wrong step, silently. Refused by name. + if (ws.tap && ws.gs == nullptr && lctx.cparams.walk_state_surrogate) { + ggml_free(ctx_compute_opt); + lctx.opt_fail(std::string("this model's recurrent state write (") + ws.cache->name + + ") does not report its exit to the walk (llm_graph_context::build_walk_state_exit): its gradient across chunks would be dropped"); + return false; + } if (ws.tap) { ggml_opt_accumulate_grad(opt_ctx, ws.tap, acc.state.at(ws.cache)); } @@ -406,12 +414,13 @@ void llama_context::opt_walk::run_exact() { // Each window starts from the REQUESTED horizon: a halving that fit one window's device budget // must not shrink every later window and epoch, which may be shorter (Fable). lctx.opt_walk_horizon = lctx.opt_walk_horizon_req; - // A recurrent context with rollback slots keeps several states per sequence; the walk carries - // exactly one. Refused by name, never asserted: training runs inside the serving process, and - // an assert there takes serving down (2026-10-07 04:09Z on the 5090). - if (recr != nullptr && lctx.cparams.n_rs_seq > 0) { - lctx.opt_fail("the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has " - + std::to_string(lctx.cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0"); + // A serving context keeps rollback slots for drafts; the walk's graphs keep one state, the live + // one (llm_graph_context::rs_keep), so they train in that context. Only a PENDING undo reads a + // snapshot instead of the live state: that alone is refused, by name, never asserted, since + // training runs inside the serving process (2026-10-07 04:09Z on the 5090). + if (recr != nullptr && recr->rollback_pending(0)) { + lctx.opt_fail("the exact walk cannot train a sequence with a pending recurrent rollback (n_rs_seq = " + + std::to_string(lctx.cparams.n_rs_seq) + "): finish or drop the draft first"); return; } plan.mark_exact(recr != nullptr, lctx.opt_walk_horizon); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9e9071876ab9..6145d5ba3643 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3690,6 +3690,11 @@ ggml_tensor * llm_graph_context::build_rs( get_state_rows); } +int64_t llm_graph_context::rs_keep() const { + const bool training = cparams.training && cparams.walk_exact; + return cparams.n_rs_seq > 0 && !training ? (int64_t) cparams.n_rs_seq + 1 : 1; +} + void llm_graph_context::build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const { if (!(cparams.training && cparams.walk_exact && cparams.walk_state_surrogate)) { return; diff --git a/src/llama-graph.h b/src/llama-graph.h index 0e0adb7d4bda..3b181afed3fb 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1367,6 +1367,15 @@ struct llm_graph_context { int32_t n_seqs, const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const; + // THE ONE RULE for how many recurrent states a graph keeps per sequence, read by every family's + // state write (delta-net, mamba, lfm2, bailing) and by the ops that produce snapshots. Serving with + // rollback slots keeps K = n_rs_seq + 1 (group g = the state g tokens back, so a draft can be undone); + // a context without them, and a TRAINING graph, keep one: the live state, group 0, which is the same + // bytes in both layouts. So training reads and writes the state serving reads, in the serving + // context; the undo groups it leaves stale are rewritten by the next serving decode, and the walk + // refuses a sequence with an undo pending (llama_memory_recurrent::rollback_pending). + int64_t rs_keep() const; + // THE EXACT WALK: the state a chunk leaves (`exit`, what is written back to `cache`) meets the // gradient later chunks put on it, through the surrogate . A no-op outside the walk. void build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5e..8ef7420a987a 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -78,6 +78,12 @@ class llama_memory_recurrent : public llama_memory_i { void set_rs_idx(llama_seq_id seq_id, uint32_t idx); + // an undo of a draft is pending on this sequence: its next read takes a snapshot group, not the + // live state (group 0) a training graph reads and writes (llm_graph_context::rs_keep) + bool rollback_pending(llama_seq_id seq_id) const { + return n_rs_seq > 0 && seq_id >= 0 && (size_t) seq_id < rs_idx.size() && rs_idx[seq_id] != 0; + } + // computed before each graph build uint32_t n = 0; diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp index 0637931cc0c9..e794c0914afc 100644 --- a/src/models/bailingmoe3.cpp +++ b/src/models/bailingmoe3.cpp @@ -183,7 +183,7 @@ static ggml_tensor * bailingmoe3_causal_conv1d( int64_t n_tokens, int64_t cache_head, uint32_t mem_size, - uint32_t n_rs_seq) { + int64_t K) { // states kept per sequence: llm_graph_context::rs_keep const int64_t d_inner = head_dim * n_head; const int64_t conv_state_size = (d_conv - 1) * d_inner; const int64_t total_state_size = 3 * conv_state_size; @@ -197,7 +197,6 @@ static ggml_tensor * bailingmoe3_causal_conv1d( x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0); - const int64_t K = (int64_t) n_rs_seq + 1; const int64_t n_written = std::min(n_seq_tokens, K); for (int64_t slot = 0; slot < n_written; ++slot) { @@ -262,13 +261,13 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph ggml_tensor * q = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * k = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * v = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); gate = ggml_add(ctx0, gate, layer.ssm_dt_b); diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index e0a6368bd2ed..f613db67f5b9 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -476,7 +476,7 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( const size_t row_size = ggml_row_size(conv_states_all->type, row_count); - if (cparams.n_rs_seq == 0) { + if (rs_keep() == 1) { const int64_t s_idx = conv_input->ne[0] - conv_states->ne[0]; const int64_t s_slot = 0; @@ -496,12 +496,12 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update)); build_walk_state_exit(conv_states_all, conv_state_last); } else { - GGML_ASSERT(!cparams.walk_exact && "the exact walk keeps one recurrent state per sequence (n_rs_seq = 0)"); + GGML_ASSERT(!cparams.walk_exact && "a training graph keeps one recurrent state per sequence (rs_keep)"); // [TAG_RECURRENT_ROLLBACK_SPLITS] // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside // the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument - const int64_t K = (int64_t) cparams.n_rs_seq + 1; + const int64_t K = rs_keep(); for (int64_t t = 1; t <= K; ++t) { const int64_t s_idx = std::max(0, conv_input->ne[0] - conv_states->ne[0] - K + t); @@ -545,7 +545,7 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( const int64_t n_seqs = s->ne[3]; const int64_t n_seq_tokens = q->ne[2]; - const bool keep = cparams.n_rs_seq > 0; + const bool keep = rs_keep() > 1; if (!keep) { auto attn_out = build_delta_net(q, k, v, g, b, s, il); @@ -564,7 +564,7 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( } const int64_t D = S_v * S_v * H_v; - const int64_t K = cparams.n_rs_seq + 1; + const int64_t K = rs_keep(); // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output. ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K); diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp index 9a4295557053..d2e954844608 100644 --- a/src/models/lfm2.cpp +++ b/src/models/lfm2.cpp @@ -205,7 +205,7 @@ llama_model_lfm2::graph::graph(const llama_model & model, const llm_graph_ GGML_ASSERT(bx->ne[0] > conv->ne[0]); // write conv states: slot 0 = the final state, slot s = the state s tokens back (partial rollback) - const int64_t K = hparams.causal_attn && cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; + const int64_t K = hparams.causal_attn ? rs_keep() : 1; const int64_t n_written = std::min(n_seq_tokens, K); const auto mem_size = mctx_cur->get_size(); const size_t row_size = ggml_row_size(conv_state->type, (int64_t) d_conv * n_embd); diff --git a/src/models/mamba-base.cpp b/src/models/mamba-base.cpp index 03ee3805bf80..af8fd15198a8 100644 --- a/src/models/mamba-base.cpp +++ b/src/models/mamba-base.cpp @@ -167,7 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, const int64_t n_seqs = ubatch.n_seqs; const int64_t n_seq_tokens = ubatch.n_seq_tokens; - const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; + const int64_t K = rs_keep(); GGML_ASSERT(n_seqs != 0); GGML_ASSERT(ubatch.equal_seqs()); diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 7e4dd1717a6e..b46144a1e0d0 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -463,36 +463,21 @@ int main(int argc, char ** argv) { } } if (llama_model_is_hybrid(model)) { - // 4. regression for the 5090 crash (2026-10-07 04:09Z): a recurrent context with rollback - // slots (serving's, inherited by the training context) made the exact walk ASSERT inside - // the serving process. It must refuse by name, and the adapter stays untouched. - llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); - const std::vector before = adapter_params(adapter); + // 4. TRAINING IN THE SERVING CONTEXT (Joel, 2026-10-13: "training is basically llama inference with + // backprop"). Serving keeps recurrent rollback slots for drafts; a training graph keeps one state, the + // live one (llm_graph_context::rs_keep), so a context WITH rollback slots takes the same step as one + // without. Was a refusal (the 5090 crash, 2026-10-07 04:09Z), which forced a second training + // context beside serving. + std::vector all(WINDOW + 1, 1); + const auto without = step_delta(params, model, init, tokens, all, WINDOW / 4, true); g_n_rs_seq = 4; - llama_context * ctx = make_ctx(params, model, WINDOW / 4); + const auto with = step_delta(params, model, init, tokens, all, WINDOW / 4, true); g_n_rs_seq = 0; - float scale = 1.0f; - GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); - llama_opt_params lopt{}; - lopt.param_filter = llama_opt_param_filter_all; - lopt.get_opt_pars = sgd_pars; - lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; - lopt.adapter = adapter; - lopt.walk_exact = true; - llama_opt_init(ctx, model, lopt); - std::vector> seqs = { std::vector(tokens.begin(), tokens.begin() + WINDOW + 1) }; - std::vector> loss = { std::vector(WINDOW + 1, 1) }; - ggml_opt_dataset_t dataset = common_opt_dataset_init_masked(WINDOW, seqs, loss, tokens[0]); - llama_opt_epoch(ctx, dataset, nullptr, nullptr, /*idata_split =*/ 1, nullptr, nullptr); - const std::string why = llama_opt_failure(ctx); - printf(" rollback slots: failed=%d, \"%s\"\n", (int) llama_opt_failed(ctx), why.c_str()); - if (!llama_opt_failed(ctx) || why.find("rollback") == std::string::npos || adapter_params(adapter) != before) { - fprintf(stderr, "FAILED: a recurrent context with rollback slots was not refused by name\n"); + if (!same_step(compare("rollback slots: exact walk in a context with n_rs_seq 4 vs 0", with, without))) { + fprintf(stderr, "FAILED: rollback slots changed the exact walk's step\n"); + worst_tensors(with, without, 24); ++failures; } - ggml_opt_dataset_free(dataset); - llama_free(ctx); - llama_adapter_lora_free(adapter); } { From e966b47b74d8825ce52189641e5349f472c38340 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Tue, 13 Oct 2026 22:09:54 -0500 Subject: [PATCH 02/43] metal: OUT_PROD runs batched, broadcast and narrow shapes on the GPU, so training stops falling back to the CPU A LoRA training epoch on the M5 spent most of its time on the CPU: Metal's OUT_PROD refused (1) any ne[2]/ne[3] above 1, which is every batched backward with grouped heads (the 4B's delta-net backward), and (2) any src0 width not a multiple of 16, which is every LoRA parameter gradient at rank 8. The 16-multiple rule only matters for the dequantizing copy of a quantized src0; it now applies only there. The 2D path becomes a per-slice helper, and the op loops over dims 2-3 with the CPU backend's broadcast rule (dst per src0 = ne2/ne02, ne3/ne03), reusing one slice of scratch. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 64 +++++++++++++++++++------- tests/test-backend-ops.cpp | 13 ++++++ 2 files changed, 60 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index cb2e266e12c9..2290f2cda254 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -3996,7 +3996,9 @@ bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { if (op->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32) { return false; } - if (src0->ne[2] != 1 || src0->ne[3] != 1 || src1->ne[2] != 1 || src1->ne[3] != 1) { + // batched and broadcast as the CPU backend defines it (dst per src0 = ne2/ne02, ne3/ne03) + if (src1->ne[2] != op->ne[2] || src1->ne[3] != op->ne[3] || + op->ne[2] % src0->ne[2] != 0 || op->ne[3] % src0->ne[3] != 0) { return false; } // src1 as the graph emits it: contiguous rows, or the TRANSPOSED view of the gradient @@ -4009,8 +4011,12 @@ bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { if (src1->nb[0] != sizeof(float) && src1->nb[1] != sizeof(float)) { return false; } - if (src0->ne[0] % 16 != 0 || src0->ne[0] % ggml_blck_size(src0->type) != 0) { - return false; // the dequantizing copy works in 4x4 tiles along whole quant blocks of a row + // only a quantized src0 is dequantized, and that copy works in 4x4 tiles along whole quant + // blocks of a row; an F32/F16/BF16 src0 copies element-wise at any width. Refusing every + // width that is not a multiple of 16 sent each LoRA parameter gradient (rank 8 is ne[0]) + // to the CPU: 2,732 of 6,112 splits of the 4B training graphs (Fable, 2026-10-13) + if (ggml_is_quantized(src0->type) && (src0->ne[0] % 16 != 0 || src0->ne[0] % ggml_blck_size(src0->type) != 0)) { + return false; } switch (src0->type) { case GGML_TYPE_F32: @@ -4046,7 +4052,7 @@ size_t ggml_metal_op_out_prod_extra(const ggml_tensor * op) { res += GGML_PAD(rows*ne00*sizeof(float), 256); // X res += GGML_PAD(rows*ne10*sizeof(float), 256); // Y if (rows < src0->ne[1]) { - res += GGML_PAD(ggml_nbytes(op), 256); // P + res += GGML_PAD(ne00*ne10*sizeof(float), 256); // P, one slice } return res; } @@ -4063,20 +4069,14 @@ static ggml_tensor ggml_metal_out_prod_view(const ggml_tensor * like, ggml_type return t; } -int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { - ggml_tensor * op = ctx->node(idx); - const ggml_tensor * src0 = op->src[0]; - const ggml_tensor * src1 = op->src[1]; - - GGML_ASSERT(ggml_metal_op_out_prod_supported(op)); - +// One 2D slice: dst[:, :] = out_prod(src0, src1) for 2D views src0 [ne00, ne01] and src1 +// [ne10, ne01], written at dst_data (inside op). The scratch after op is reused by every slice: +// each step ends on a memory barrier, so the slices run one after another. +static void ggml_metal_op_out_prod_slice(ggml_metal_op_t ctx, const ggml_tensor * op, + const ggml_tensor * src0, const ggml_tensor * src1, char * dst_data, int64_t rows_blk) { const int64_t ne00 = src0->ne[0]; const int64_t ne01 = src0->ne[1]; const int64_t ne10 = src1->ne[0]; - GGML_ASSERT(src1->ne[1] == ne01); - GGML_ASSERT(op->ne[0] == ne00 && op->ne[1] == ne10); - - const int64_t rows_blk = ggml_metal_op_out_prod_block_rows(op); // scratch after dst, in dst's buffer (ggml_backend_metal_buffer_type_get_alloc_size) char * scratch = (char *) op->data + ggml_nbytes(op); @@ -4132,7 +4132,7 @@ int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { // P = mul_mat(X, Y): [ne00, ne10], contraction over rows const bool first = k0 == 0; ggml_tensor p = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, ne10, - sizeof(float), ne00*sizeof(float), first ? op->data : (void *) p_data); + sizeof(float), ne00*sizeof(float), first ? (void *) dst_data : (void *) p_data); p.op = GGML_OP_MUL_MAT; p.src[0] = &x; p.src[1] = &y; @@ -4146,7 +4146,7 @@ int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { if (!first) { // dst = dst + P ggml_tensor d = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, ne10, - sizeof(float), ne00*sizeof(float), op->data); + sizeof(float), ne00*sizeof(float), dst_data); ggml_tensor acc = d; acc.op = GGML_OP_ADD; acc.src[0] = &d; @@ -4155,6 +4155,36 @@ int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { ggml_metal_op_concurrency_reset(ctx); } } +} + +int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + const ggml_tensor * src0 = op->src[0]; + const ggml_tensor * src1 = op->src[1]; + + GGML_ASSERT(ggml_metal_op_out_prod_supported(op)); + GGML_ASSERT(src1->ne[1] == src0->ne[1]); + GGML_ASSERT(op->ne[0] == src0->ne[0] && op->ne[1] == src1->ne[0]); + + const int64_t rows_blk = ggml_metal_op_out_prod_block_rows(op); + + // batched and broadcast as the CPU backend defines it: dst and src1 share dims 2-3, and + // src0 is broadcast over them (dps = dst per src0: grouped-query heads). Each slice is the + // 2D product; before this, any ne[2] or ne[3] above 1 ran on the CPU (a delta-net + // backward on the 4B: most of the training CPU time, 2026-10-13) + const int64_t dps2 = op->ne[2]/src0->ne[2]; + const int64_t dps3 = op->ne[3]/src0->ne[3]; + for (int64_t i3 = 0; i3 < op->ne[3]; ++i3) { + for (int64_t i2 = 0; i2 < op->ne[2]; ++i2) { + const int64_t i02 = i2/dps2; + const int64_t i03 = i3/dps3; + const ggml_tensor a = ggml_metal_out_prod_view(src0, src0->type, src0->ne[0], src0->ne[1], + src0->nb[0], src0->nb[1], (char *) src0->data + i02*src0->nb[2] + i03*src0->nb[3]); + const ggml_tensor b = ggml_metal_out_prod_view(src1, GGML_TYPE_F32, src1->ne[0], src1->ne[1], + src1->nb[0], src1->nb[1], (char *) src1->data + i2*src1->nb[2] + i3*src1->nb[3]); + ggml_metal_op_out_prod_slice(ctx, op, &a, &b, (char *) op->data + i2*op->nb[2] + i3*op->nb[3], rows_blk); + } + } return 1; } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index b53590be849e..068f90e81e3f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9774,6 +9774,19 @@ static std::vector> make_test_cases_eval() { type_a == GGML_TYPE_F32 ? 1024 : 256, 1024, 8, {1, 1}, {1, 1}, trans_b, true)); } } + // the batched, broadcast shape a delta-net backward hands OUT_PROD on the 4B (16 value heads + // over 4 key heads, head width 256 and 128), both gradient layouts + for (bool trans_b : {false, true}) { + test_cases.emplace_back(new test_out_prod(GGML_TYPE_F32, GGML_TYPE_F32, 256, 128, 256, {4, 1}, {4, 1}, trans_b)); + } + // a LoRA parameter gradient: the rank is the OUTPUT width (dA, dB = out_prod(x, grad) with + // m = rank), not the contraction; ranks below and between multiples of 16, at the 4B walk's + // shapes (8192 tokens, 128 and 1024 wide), both layouts + for (int64_t rank : {4, 8, 24}) { + test_cases.emplace_back(new test_out_prod(GGML_TYPE_F32, GGML_TYPE_F32, rank, 128, 8192, {1, 1}, {1, 1}, true)); + test_cases.emplace_back(new test_out_prod(GGML_TYPE_F32, GGML_TYPE_F32, rank, 8192, 128, {1, 1}, {1, 1}, false)); + test_cases.emplace_back(new test_out_prod(GGML_TYPE_F32, GGML_TYPE_F32, rank, 1024, 128, {1, 1}, {1, 1}, false)); + } // the q8_0 three-block case in both gradient layouts: the training graph hands // out_prod a TRANSPOSED gradient (out_prod(W, ggml_transpose(grad))) for (bool trans_b : {false, true}) { From 746882623eddf3541f0321ee61deb1e25618324c Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Tue, 13 Oct 2026 23:06:37 -0500 Subject: [PATCH 03/43] metal: a native OUT_PROD kernel reads the weight in place; the composed copy-and-transpose path is deleted OUT_PROD was composed from a dequantizing copy, two transposing copies, a mul_mm and an add per row block, over scratch reserved after dst. Every pass wrote device memory and read it back, so the backward matmul ran at about half the forward's rate. kernels/out_prod.metal is mul_mm with the contraction along the rows of both operands. Each thread dequantizes 16 consecutive elements of one weight row in the weight's own block order and stages them transposed in threadgroup memory. The gradient is read with its own strides, so both layouts work, and staging is float. Batching and broadcast come from the grid. Deleted with it: the row blocks, OUT_PROD_BLOCK_BYTES, the scratch after dst and its get_alloc_size hook, and the synthetic-view helper. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-metal/CMakeLists.txt | 1 + ggml/src/ggml-metal/ggml-metal-device.cpp | 21 ++ ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-device.m | 1 + ggml/src/ggml-metal/ggml-metal-impl.h | 16 ++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 225 ++++----------------- ggml/src/ggml-metal/ggml-metal-ops.h | 1 - ggml/src/ggml-metal/ggml-metal.cpp | 4 - ggml/src/ggml-metal/kernels/out_prod.metal | 157 ++++++++++++++ 9 files changed, 236 insertions(+), 191 deletions(-) create mode 100644 ggml/src/ggml-metal/kernels/out_prod.metal diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt index 140c5d809e02..6f34c7b8fa8c 100644 --- a/ggml/src/ggml-metal/CMakeLists.txt +++ b/ggml/src/ggml-metal/CMakeLists.txt @@ -33,6 +33,7 @@ set(METALLIB_KERNEL_SOURCES kernels/fa.metal kernels/mul_mv.metal kernels/mul_mm.metal + kernels/out_prod.metal kernels/quantize.metal kernels/softmax.metal kernels/norm.metal diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 22a37453bf39..47dc506b391e 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -909,6 +909,27 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_out_prod(ggml_metal_library_t lib, const ggml_tensor * op) { + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_out_prod_%s", ggml_type_name(op->src[0]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + // kernels/out_prod.metal: a 64 x 32 tile, float staging of both 32-deep operands + res.nr0 = 64; + res.nr1 = 32; + res.nsg = 4; + res.smem = (64 + 32)*32*sizeof(float); + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_metal_library_t lib, const ggml_tensor * op) { GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 18ae227a5ca9..2e7afa938cf5 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -140,6 +140,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_del struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_out_prod (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_fp32 (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0 (ggml_metal_library_t lib, int ne02, int ne20); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c916ca8a1657..cbdb4e32b8f7 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -109,6 +109,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi X(FA, fa) \ X(MUL_MV, mul_mv) \ X(MUL_MM, mul_mm) \ + X(OUT_PROD, out_prod) \ X(QUANTIZE, quantize) \ X(SOFTMAX, softmax) \ X(NORM, norm) \ diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index adc883e20f1d..109bb6cf9154 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -481,6 +481,22 @@ typedef struct { int16_t r3; } ggml_metal_kargs_mul_mm; +typedef struct { + int32_t ne00; // m: dst rows + int32_t ne01; // k: the contraction, rows of both sources + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; // n: dst columns + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne2; + int32_t dps2; // dst per src0 along dim 2 (grouped heads) + int32_t dps3; +} ggml_metal_kargs_out_prod; + typedef struct { int32_t ne00; int32_t ne01; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 2290f2cda254..dcaa13cbba61 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -3962,67 +3962,27 @@ static void ggml_metal_encode_add(ggml_metal_op_t ctx, const ggml_tensor * op) { } } -// OUT_PROD: dst[ne00, ne10] = sum over k of src0[:, k] (x) src1[:, k] โ€” the activation gradient -// through a FROZEN weight in a training graph (out_prod(W, grad^T), the mul_mat backward). -// This backend had no OUT_PROD at all, so a Mac could not train a LoRA on the weights it -// serves (Continuum card 55f314b8, the Metal half of the doctor lane). Composed from kernels -// the backend already has, over scratch that the graph allocator reserves after dst: -// -// for each block of rows k0..k0+rows of W: -// S = cpy(W[k0..], F32) (the dequantizing copy; skipped when W is F32) -// X = cpy(S^T) contiguous [rows, ne00] โ€” contraction dim first -// Y = cpy(src1[:, k0..]^T) contiguous [rows, ne10] -// P = mul_mm(X, Y) [ne00, ne10]; the first block writes dst directly -// dst += P for every block after the first -// -// Row blocks bound the transient to OUT_PROD_BLOCK_BYTES of F32 for any weight (the output -// head of a 27B is 3 GB in F32 whole), the same shape as the CUDA out_prod. First cut: -// correctness by test-backend-ops parity against the CPU backend; a fused dequantize-GEMM -// is the second cut if the ยง0 speed gate asks for it. -static constexpr int64_t OUT_PROD_BLOCK_BYTES = 64ll*1024*1024; - -static int64_t ggml_metal_op_out_prod_block_rows(const ggml_tensor * op) { - const ggml_tensor * src0 = op->src[0]; - const int64_t ne00 = src0->ne[0]; - const int64_t ne01 = src0->ne[1]; - int64_t rows = OUT_PROD_BLOCK_BYTES/(ne00*(int64_t) sizeof(float)); - rows = std::max(64, (rows/64)*64); - return std::min(rows, ne01); -} - +// OUT_PROD: dst[ne00, ne10] = sum over k of src0[:, k] (x) src1[:, k]. In a training graph it +// is the activation gradient through a frozen weight (out_prod(W, grad^T), the mul_mat +// backward) and the gradient of each LoRA factor. kernels/out_prod.metal computes it in one +// dispatch, reading the weight in place in its own block order. Batched and broadcast +// operands follow the CPU backend: dst and src1 share dims 2-3, and src0 is broadcast over +// them (dst per src0 = ne2/ne02, ne3/ne03). bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { const ggml_tensor * src0 = op->src[0]; const ggml_tensor * src1 = op->src[1]; - if (op->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32) { + if (op->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || !ggml_is_contiguous(op)) { return false; } - // batched and broadcast as the CPU backend defines it (dst per src0 = ne2/ne02, ne3/ne03) if (src1->ne[2] != op->ne[2] || src1->ne[3] != op->ne[3] || op->ne[2] % src0->ne[2] != 0 || op->ne[3] % src0->ne[3] != 0) { return false; } - // src1 as the graph emits it: contiguous rows, or the TRANSPOSED view of the gradient - // (out_prod(W, ggml_transpose(grad)) is what the mul_mat backward builds; Cormac on - // llama.cpp#5). Either way Y is copied contiguous per block, so both layouts are - // exact; anything else is refused. - if (!ggml_is_contiguous(op)) { - return false; - } - if (src1->nb[0] != sizeof(float) && src1->nb[1] != sizeof(float)) { - return false; - } - // only a quantized src0 is dequantized, and that copy works in 4x4 tiles along whole quant - // blocks of a row; an F32/F16/BF16 src0 copies element-wise at any width. Refusing every - // width that is not a multiple of 16 sent each LoRA parameter gradient (rank 8 is ne[0]) - // to the CPU: 2,732 of 6,112 splits of the 4B training graphs (Fable, 2026-10-13) - if (ggml_is_quantized(src0->type) && (src0->ne[0] % 16 != 0 || src0->ne[0] % ggml_blck_size(src0->type) != 0)) { - return false; - } switch (src0->type) { case GGML_TYPE_F32: - return ggml_is_contiguous_rows(src0); case GGML_TYPE_F16: case GGML_TYPE_BF16: + return src0->nb[0] == ggml_type_size(src0->type); // read element by element, any width case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -4033,130 +3993,13 @@ bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: - return ggml_is_contiguous(src0); // kernel_cpy__f32 exists + // dequantized 16 consecutive elements at a time, from whole blocks of a row + return ggml_is_contiguous_rows(src0) && src0->ne[0] % 16 == 0; default: return false; } } -size_t ggml_metal_op_out_prod_extra(const ggml_tensor * op) { - const ggml_tensor * src0 = op->src[0]; - const ggml_tensor * src1 = op->src[1]; - const int64_t ne00 = src0->ne[0]; - const int64_t ne10 = src1->ne[0]; - const int64_t rows = ggml_metal_op_out_prod_block_rows(op); - size_t res = 0; - if (src0->type != GGML_TYPE_F32) { - res += GGML_PAD(rows*ne00*sizeof(float), 256); // S - } - res += GGML_PAD(rows*ne00*sizeof(float), 256); // X - res += GGML_PAD(rows*ne10*sizeof(float), 256); // Y - if (rows < src0->ne[1]) { - res += GGML_PAD(ne00*ne10*sizeof(float), 256); // P, one slice - } - return res; -} - -// A synthetic 2D view over the same Metal buffer as `like`, at `data`. -static ggml_tensor ggml_metal_out_prod_view(const ggml_tensor * like, ggml_type type, - int64_t ne0, int64_t ne1, size_t nb0, size_t nb1, void * data) { - ggml_tensor t = {}; - t.type = type; - t.ne[0] = ne0; t.ne[1] = ne1; t.ne[2] = 1; t.ne[3] = 1; - t.nb[0] = nb0; t.nb[1] = nb1; t.nb[2] = nb1*ne1; t.nb[3] = t.nb[2]; - t.buffer = like->buffer; - t.data = data; - return t; -} - -// One 2D slice: dst[:, :] = out_prod(src0, src1) for 2D views src0 [ne00, ne01] and src1 -// [ne10, ne01], written at dst_data (inside op). The scratch after op is reused by every slice: -// each step ends on a memory barrier, so the slices run one after another. -static void ggml_metal_op_out_prod_slice(ggml_metal_op_t ctx, const ggml_tensor * op, - const ggml_tensor * src0, const ggml_tensor * src1, char * dst_data, int64_t rows_blk) { - const int64_t ne00 = src0->ne[0]; - const int64_t ne01 = src0->ne[1]; - const int64_t ne10 = src1->ne[0]; - - // scratch after dst, in dst's buffer (ggml_backend_metal_buffer_type_get_alloc_size) - char * scratch = (char *) op->data + ggml_nbytes(op); - char * s_data = nullptr; - if (src0->type != GGML_TYPE_F32) { - s_data = scratch; - scratch += GGML_PAD(rows_blk*ne00*sizeof(float), 256); - } - char * x_data = scratch; scratch += GGML_PAD(rows_blk*ne00*sizeof(float), 256); - char * y_data = scratch; scratch += GGML_PAD(rows_blk*ne10*sizeof(float), 256); - char * p_data = rows_blk < ne01 ? scratch : nullptr; - - for (int64_t k0 = 0; k0 < ne01; k0 += rows_blk) { - const int64_t rows = std::min(rows_blk, ne01 - k0); - - // S: this block of W, dequantized to F32, [ne00, rows] contiguous - ggml_tensor s_view; - if (src0->type == GGML_TYPE_F32) { - s_view = ggml_metal_out_prod_view(src0, GGML_TYPE_F32, ne00, rows, - src0->nb[0], src0->nb[1], (char *) src0->data + k0*src0->nb[1]); - } else { - ggml_tensor w_view = ggml_metal_out_prod_view(src0, src0->type, ne00, rows, - src0->nb[0], src0->nb[1], (char *) src0->data + k0*src0->nb[1]); - s_view = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, rows, - sizeof(float), ne00*sizeof(float), s_data); - ggml_metal_op_cpy_impl(ctx, &w_view, &s_view); - // the encoder dispatches concurrently: every kernel below reads what the one - // above wrote, so each step ends on a memory barrier (as cumsum does) - ggml_metal_op_concurrency_reset(ctx); - } - - // X = S^T: a transposed view of S copied contiguous, [rows, ne00] - // based on S, never on op: for an F32 src0, S IS src0's memory, which lives in another - // buffer (a LoRA weight's adapter buffer). Based on op, the view named op's buffer with - // src0's address, the lookup found no buffer ('buffer is nil'), the copy below did not - // run, and the product read stale scratch: every OUT_PROD with a LoRA weight as src0 - // was wrong on Metal, differently per graph layout (Fable, 2026-10-06) - ggml_tensor st_view = ggml_metal_out_prod_view(&s_view, GGML_TYPE_F32, rows, ne00, - s_view.nb[1], s_view.nb[0], s_view.data); - ggml_tensor x = ggml_metal_out_prod_view(op, GGML_TYPE_F32, rows, ne00, - sizeof(float), rows*sizeof(float), x_data); - ggml_metal_op_cpy_impl(ctx, &st_view, &x); - ggml_metal_op_concurrency_reset(ctx); - - // Y = src1[:, k0..k0+rows]^T copied contiguous, [rows, ne10] - ggml_tensor bt_view = ggml_metal_out_prod_view(src1, GGML_TYPE_F32, rows, ne10, - src1->nb[1], src1->nb[0], (char *) src1->data + k0*src1->nb[1]); - ggml_tensor y = ggml_metal_out_prod_view(op, GGML_TYPE_F32, rows, ne10, - sizeof(float), rows*sizeof(float), y_data); - ggml_metal_op_cpy_impl(ctx, &bt_view, &y); - ggml_metal_op_concurrency_reset(ctx); - - // P = mul_mat(X, Y): [ne00, ne10], contraction over rows - const bool first = k0 == 0; - ggml_tensor p = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, ne10, - sizeof(float), ne00*sizeof(float), first ? (void *) dst_data : (void *) p_data); - p.op = GGML_OP_MUL_MAT; - p.src[0] = &x; - p.src[1] = &y; - // FLOAT staging when the simdgroup kernel runs: X and Y carry gradients, which pass - // half's range (65504) where activations do not; half staging turned them into Inf - // and cost every product ~1e-3 relative precision (2026-10-06, Qwen3.5 0.8B backward) - const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(ctx->lib))->has_tensor; - ggml_metal_encode_mul_mm(ctx, &p, !has_tensor); - ggml_metal_op_concurrency_reset(ctx); - - if (!first) { - // dst = dst + P - ggml_tensor d = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, ne10, - sizeof(float), ne00*sizeof(float), dst_data); - ggml_tensor acc = d; - acc.op = GGML_OP_ADD; - acc.src[0] = &d; - acc.src[1] = &p; - ggml_metal_encode_add(ctx, &acc); - ggml_metal_op_concurrency_reset(ctx); - } - } -} - int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); const ggml_tensor * src0 = op->src[0]; @@ -4166,25 +4009,35 @@ int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { GGML_ASSERT(src1->ne[1] == src0->ne[1]); GGML_ASSERT(op->ne[0] == src0->ne[0] && op->ne[1] == src1->ne[0]); - const int64_t rows_blk = ggml_metal_op_out_prod_block_rows(op); - - // batched and broadcast as the CPU backend defines it: dst and src1 share dims 2-3, and - // src0 is broadcast over them (dps = dst per src0: grouped-query heads). Each slice is the - // 2D product; before this, any ne[2] or ne[3] above 1 ran on the CPU (a delta-net - // backward on the 4B: most of the training CPU time, 2026-10-13) - const int64_t dps2 = op->ne[2]/src0->ne[2]; - const int64_t dps3 = op->ne[3]/src0->ne[3]; - for (int64_t i3 = 0; i3 < op->ne[3]; ++i3) { - for (int64_t i2 = 0; i2 < op->ne[2]; ++i2) { - const int64_t i02 = i2/dps2; - const int64_t i03 = i3/dps3; - const ggml_tensor a = ggml_metal_out_prod_view(src0, src0->type, src0->ne[0], src0->ne[1], - src0->nb[0], src0->nb[1], (char *) src0->data + i02*src0->nb[2] + i03*src0->nb[3]); - const ggml_tensor b = ggml_metal_out_prod_view(src1, GGML_TYPE_F32, src1->ne[0], src1->ne[1], - src1->nb[0], src1->nb[1], (char *) src1->data + i2*src1->nb[2] + i3*src1->nb[3]); - ggml_metal_op_out_prod_slice(ctx, op, &a, &b, (char *) op->data + i2*op->nb[2] + i3*op->nb[3], rows_blk); - } - } + ggml_metal_kargs_out_prod args = { + /*.ne00 =*/ (int32_t) src0->ne[0], + /*.ne01 =*/ (int32_t) src0->ne[1], + /*.nb01 =*/ src0->nb[1], + /*.nb02 =*/ src0->nb[2], + /*.nb03 =*/ src0->nb[3], + /*.ne10 =*/ (int32_t) src1->ne[0], + /*.nb10 =*/ src1->nb[0], + /*.nb11 =*/ src1->nb[1], + /*.nb12 =*/ src1->nb[2], + /*.nb13 =*/ src1->nb[3], + /*.ne2 =*/ (int32_t) op->ne[2], + /*.dps2 =*/ (int32_t) (op->ne[2]/src0->ne[2]), + /*.dps3 =*/ (int32_t) (op->ne[3]/src0->ne[3]), + }; + + auto pipeline = ggml_metal_library_get_pipeline_out_prod(ctx->lib, op); + + ggml_metal_encoder_t enc = ctx->enc; + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(src0), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(src1), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); + ggml_metal_encoder_dispatch_threadgroups(enc, + (op->ne[1] + pipeline.nr1 - 1)/pipeline.nr1, + (op->ne[0] + pipeline.nr0 - 1)/pipeline.nr0, + op->ne[2]*op->ne[3], 32, pipeline.nsg, 1); return 1; } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index e76003f71bca..9c4853a9d872 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -76,7 +76,6 @@ int ggml_metal_op_mul_mat_id (ggml_metal_op_t ctx, int idx); int ggml_metal_op_add_id (ggml_metal_op_t ctx, int idx); int ggml_metal_op_flash_attn_ext (ggml_metal_op_t ctx, int idx); int ggml_metal_op_bin (ggml_metal_op_t ctx, int idx); -size_t ggml_metal_op_out_prod_extra(const struct ggml_tensor * op); bool ggml_metal_op_out_prod_supported(const struct ggml_tensor * op); int ggml_metal_op_out_prod (ggml_metal_op_t ctx, int idx); int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 0ae45e3d7bc3..c462c0686371 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -233,10 +233,6 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_ res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor); res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor); } break; - case GGML_OP_OUT_PROD: - { - res += ggml_metal_op_out_prod_extra(tensor); - } break; case GGML_OP_GATED_DELTA_NET_BACK: { res += ggml_metal_op_gated_delta_net_back_extra(tensor); diff --git a/ggml/src/ggml-metal/kernels/out_prod.metal b/ggml/src/ggml-metal/kernels/out_prod.metal new file mode 100644 index 000000000000..5156dd0f0539 --- /dev/null +++ b/ggml/src/ggml-metal/kernels/out_prod.metal @@ -0,0 +1,157 @@ +#include "common.h" +#include "dequantize.h" + +// OUT_PROD: dst[m, n] = sum_k src0[m, k] * src1[n, k]. It is the backward of a matmul through +// a frozen weight (dx = out_prod(W, dy^T)) and the gradient of a LoRA factor (dA, dB). +// +// This is mul_mm with the contraction along the ROWS of both operands. Each thread reads 16 +// consecutive elements of one src0 row, in the weight's own block order, dequantizes them in +// registers, and stages them transposed in threadgroup memory. Nothing is copied or transposed +// in device memory. The tiles and the simdgroup product are mul_mm's: a 64 x 32 output tile, +// 32-deep k steps, 4 simdgroups. Staging is float because gradients exceed half's range. +// +// nl == 0 selects the element path (F32/F16/BF16, any width, a LoRA rank of 8 included). +// nl > 0 dequantizes whole 16-element chunks, which needs ne00 % 16 == 0. + +template +kernel void kernel_out_prod( + constant ggml_metal_kargs_out_prod & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr short OP_NR0 = 64; // output rows (m) per threadgroup + constexpr short OP_NR1 = 32; // output columns (n) per threadgroup + constexpr short OP_NK = 32; // contraction depth per step + + threadgroup float * sa = (threadgroup float *) shmem; // OP_NR0 x OP_NK + threadgroup float * sb = sa + OP_NR0*OP_NK; // OP_NR1 x OP_NK + + const int r0 = tgpig.y*OP_NR0; + const int r1 = tgpig.x*OP_NR1; + const int i2 = tgpig.z % args.ne2; + const int i3 = tgpig.z / args.ne2; + + // src0 is broadcast over dims 2-3 (grouped heads); src1 and dst are not + src0 += (i2/args.dps2)*args.nb02 + (i3/args.dps3)*args.nb03; + src1 += i2*args.nb12 + i3*args.nb13; + + // src0 staging: thread -> one k row, 16 consecutive m + const short ka = tiitg/4; + const short ma = 16*(tiitg%4); + // src1 staging: thread -> one n, 8 consecutive k + const short nb = tiitg/4; + const short kb = 8*(tiitg%4); + + simdgroup_float8x8 ma8[4]; + simdgroup_float8x8 mb8[2]; + simdgroup_float8x8 mc[8]; + for (short i = 0; i < 8; i++) { + mc[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int k0 = 0; k0 < args.ne01; k0 += OP_NK) { + float4x4 a = float4x4(0.f); + const int k = k0 + ka; + const int m = r0 + ma; + if (k < args.ne01 && m < args.ne00) { + device const char * row = src0 + (uint64_t) k*args.nb01; + if (nl == 0) { + device const T0 * x = (device const T0 *) row + m; + for (short i = 0; i < 16; i++) { + a[i/4][i%4] = m + i < args.ne00 ? (float) x[i] : 0.f; + } + } else { + dequantize_func((device const block_q *) row + m/(16*nl), (m/16) % nl, a); + } + } + + float b[8]; + { + const int n = r1 + nb; + for (short i = 0; i < 8; i++) { + const int kk = k0 + kb + i; + b[i] = kk < args.ne01 && n < args.ne10 + ? *(device const float *) (src1 + (uint64_t) kk*args.nb11 + (uint64_t) n*args.nb10) : 0.f; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // the same 8x8-block layout mul_mm stages: sa block (k/8, m/8) holds [k%8][m%8], + // sb block (k/8, n/8) holds [n%8][k%8] + for (short i = 0; i < 16; i++) { + const short mm = ma + i; + sa[64*(8*(ka/8) + mm/8) + 8*(ka%8) + mm%8] = a[i/4][i%4]; + } + for (short i = 0; i < 8; i++) { + const short kk = kb + i; + sb[64*(4*(kk/8) + nb/8) + 8*(nb%8) + kk%8] = b[i]; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const float * lsma = sa + 4*64*(sgitg%2); + threadgroup const float * lsmb = sb + 2*64*(sgitg/2); + for (short ik = 0; ik < OP_NK/8; ik++) { + for (short i = 0; i < 4; i++) { + simdgroup_load(ma8[i], lsma + 64*i, 8, 0, false); + } + for (short i = 0; i < 2; i++) { + simdgroup_load(mb8[i], lsmb + 64*i, 8, 0, false); + } + for (short i = 0; i < 8; i++) { + simdgroup_multiply_accumulate(mc[i], mb8[i/4], ma8[i%4], mc[i]); + } + lsma += 8*64; + lsmb += 4*64; + } + } + + device float * d = (device float *) dst + (uint64_t) tgpig.z*args.ne00*args.ne10; + + if (r0 + OP_NR0 <= args.ne00 && r1 + OP_NR1 <= args.ne10) { + device float * C = d + (r0 + 32*(sgitg & 1)) + (r1 + 16*(sgitg >> 1))*args.ne00; + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne00*(i/4), args.ne00, 0, false); + } + return; + } + + // an edge tile: stage it, then write only the part inside dst + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float * t = (threadgroup float *) shmem + 32*(sgitg & 1) + 16*(sgitg >> 1)*OP_NR0; + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], t + 8*(i%4) + 8*OP_NR0*(i/4), OP_NR0, 0, false); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const int nr0 = min((int) OP_NR0, args.ne00 - r0); + const int nr1 = min((int) OP_NR1, args.ne10 - r1); + for (int j = tiitg; j < nr0*nr1; j += 128) { + const int jm = j % nr0; + const int jn = j / nr0; + d[(r0 + jm) + (r1 + jn)*args.ne00] = ((threadgroup float *) shmem)[jm + jn*OP_NR0]; + } +} + +typedef decltype(kernel_out_prod) out_prod_t; + +template [[host_name("kernel_out_prod_f32")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_f16")]] kernel out_prod_t kernel_out_prod; +#if defined(GGML_METAL_HAS_BF16) +template [[host_name("kernel_out_prod_bf16")]] kernel out_prod_t kernel_out_prod; +#endif +template [[host_name("kernel_out_prod_q4_0")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q4_1")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q5_0")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q5_1")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q8_0")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q2_K")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q3_K")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q4_K")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q5_K")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q6_K")]] kernel out_prod_t kernel_out_prod; From 22f75563e02939113a4d0eb899a62b267800b382 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Tue, 13 Oct 2026 23:12:17 -0500 Subject: [PATCH 04/43] metal: left padding, strided norm/softmax backward and GET_ROWS_BACK, so a LoRA training graph stays on the GPU The ops a Metal training graph still sent to the CPU after OUT_PROD, each refused by supports_op: - PAD with left padding. The kernel now reads src[i - lp] per dim, through the source's element stride (a training graph pads a view). - RMS_NORM_BACK and SOFT_MAX_BACK on strided rows (a per-head norm). Rows are now addressed through each tensor's dim 1-3 strides, and each row stays contiguous. - GET_ROWS_BACK (no kernel). This is the gradient of the output-row gather. One threadgroup per dst row sums, in id order, the source rows whose id is that row: no atomics, and the same summation order as the CPU reference. Tests: left-pad cases per dim on contiguous, non-contiguous and permuted sources, and an RMS_NORM_BACK over row views of a twice-wide tensor. Not in this change: ADD with a transposed src1 (a strided read in the hot binary kernel; it measured about 1% of a 4B epoch). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-metal/ggml-metal-device.cpp | 3 +- ggml/src/ggml-metal/ggml-metal-device.m | 14 ++++--- ggml/src/ggml-metal/ggml-metal-impl.h | 21 ++++++++-- ggml/src/ggml-metal/ggml-metal-ops.cpp | 47 ++++++++++++++++++++--- ggml/src/ggml-metal/ggml-metal-ops.h | 1 + ggml/src/ggml-metal/kernels/misc.metal | 15 +++++--- ggml/src/ggml-metal/kernels/softmax.metal | 40 ++++++++++++++----- tests/test-backend-ops.cpp | 30 +++++++++++---- 8 files changed, 132 insertions(+), 39 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 47dc506b391e..846f651e43fd 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -351,13 +351,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_lo return res; } -// RMS_NORM_BACK, SOFT_MAX_BACK, REPEAT_BACK: one f32 kernel each, named after the op. +// RMS_NORM_BACK, SOFT_MAX_BACK, REPEAT_BACK, GET_ROWS_BACK: one f32 kernel each, named after the op. ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back(ggml_metal_library_t lib, const ggml_tensor * op) { const char * kernel = nullptr; switch (op->op) { case GGML_OP_RMS_NORM_BACK: kernel = "kernel_rms_norm_back_f32"; break; case GGML_OP_SOFT_MAX_BACK: kernel = "kernel_soft_max_back_f32"; break; case GGML_OP_REPEAT_BACK: kernel = "kernel_repeat_back_f32"; break; + case GGML_OP_GET_ROWS_BACK: kernel = "kernel_get_rows_back_f32"; break; default: GGML_ABORT("fatal error"); } diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index cbdb4e32b8f7..fe2d103943f4 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1471,13 +1471,18 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te } case GGML_OP_RMS_NORM_BACK: case GGML_OP_SOFT_MAX_BACK: - // src0 the incoming gradient, src1 the forward input/output; rows contiguous. - // SOFT_MAX_BACK carries no max_bias (the CPU reference asserts the same). + // src0 the incoming gradient, src1 the forward input/output; each row contiguous, the + // rows themselves at any stride. SOFT_MAX_BACK carries no max_bias (the CPU reference + // asserts the same). return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && - ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op) && + ggml_are_same_shape(op->src[0], op->src[1]) && + ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]) && ggml_is_contiguous_rows(op) && (op->op != GGML_OP_SOFT_MAX_BACK || ggml_get_op_params_f32(op, 1) == 0.0f); case GGML_OP_REPEAT_BACK: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op); + case GGML_OP_GET_ROWS_BACK: + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_I32 && + ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op); case GGML_OP_CROSS_ENTROPY_LOSS: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); @@ -1610,8 +1615,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return false; } - return (ggml_get_op_params_i32(op, 0) == 0) && (ggml_get_op_params_i32(op, 2) == 0) && - (ggml_get_op_params_i32(op, 4) == 0) && (ggml_get_op_params_i32(op, 6) == 0); + return true; // left and right zero padding on every dim case GGML_OP_PAD_REFLECT_1D: case GGML_OP_TIMESTEP_EMBEDDING: return op->src[0]->type == GGML_TYPE_F32; diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 109bb6cf9154..6d1d484635e3 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -905,14 +905,25 @@ typedef struct { } ggml_metal_kargs_cross_entropy_loss; // a per-row backward with one scalar parameter (eps for RMS_NORM_BACK, scale for SOFT_MAX_BACK) +// rows are addressed through each source's own strides, so a strided view (a per-head norm) +// is read in place typedef struct { int32_t ne00; - uint64_t nb01; - uint64_t nb11; - uint64_t nb1; + uint64_t nb01, nb02, nb03; + uint64_t nb11, nb12, nb13; + uint64_t nb1, nb2, nb3; float param; } ggml_metal_kargs_row_back; +// GET_ROWS_BACK: dst = 0, then dst[:, ids[i]] += src0[:, i] for each i (the gradient of a row +// gather; in a training graph, the selection of the output rows) +typedef struct { + int32_t ne00; // row width + int32_t nr; // number of ids + uint64_t nb01; + uint64_t nb1; +} ggml_metal_kargs_get_rows_back; + typedef struct { int64_t ne00; int64_t ne01; @@ -1144,6 +1155,10 @@ typedef struct { uint64_t nb1; uint64_t nb2; uint64_t nb3; + int32_t lp0; // left padding per dim: dst[i] = src[i - lp] + int32_t lp1; + int32_t lp2; + int32_t lp3; } ggml_metal_kargs_pad; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index dcaa13cbba61..90392dbd6a65 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -321,6 +321,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_repeat_back(ctx, idx); } break; + case GGML_OP_GET_ROWS_BACK: + { + n_fuse = ggml_metal_op_get_rows_back(ctx, idx); + } break; case GGML_OP_SILU_BACK: { n_fuse = ggml_metal_op_silu_back(ctx, idx); @@ -4131,13 +4135,14 @@ int ggml_metal_op_row_back(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_t enc = ctx->enc; const int32_t ne00 = op->src[0]->ne[0]; - const int32_t nr = ggml_nrows(op->src[0]); + const ggml_tensor * s0 = op->src[0]; + const ggml_tensor * s1 = op->src[1]; ggml_metal_kargs_row_back args = { /*.ne00 =*/ ne00, - /*.nb01 =*/ op->src[0]->nb[1], - /*.nb11 =*/ op->src[1]->nb[1], - /*.nb1 =*/ op->nb[1], + /*.nb01 =*/ s0->nb[1], s0->nb[2], s0->nb[3], + /*.nb11 =*/ s1->nb[1], s1->nb[2], s1->nb[3], + /*.nb1 =*/ op->nb[1], op->nb[2], op->nb[3], /*.param =*/ ggml_get_op_params_f32(op, 0), // eps, or scale }; @@ -4155,7 +4160,7 @@ int ggml_metal_op_row_back(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ggml_metal_encoder_set_threadgroup_memory_size(enc, 32*sizeof(float), 0); - ggml_metal_encoder_dispatch_threadgroups(enc, nr, 1, 1, nth, 1, 1); + ggml_metal_encoder_dispatch_threadgroups(enc, s0->ne[1], s0->ne[2], s0->ne[3], nth, 1, 1); return 1; } @@ -4191,6 +4196,32 @@ int ggml_metal_op_repeat_back(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_get_rows_back(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + ggml_metal_kargs_get_rows_back args = { + /*.ne00 =*/ (int32_t) op->src[0]->ne[0], + /*.nr =*/ (int32_t) ggml_nelements(op->src[1]), + /*.nb01 =*/ op->src[0]->nb[1], + /*.nb1 =*/ op->nb[1], + }; + + auto pipeline = ggml_metal_library_get_pipeline_row_back(lib, op); + + const int nth = std::min(args.ne00, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], 1, 1, nth, 1, 1); + + return 1; +} + int ggml_metal_op_silu_back(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -5225,7 +5256,11 @@ int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) { /*.nb0 =*/ nb0, /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, - /*.nb3 =*/ nb3 + /*.nb3 =*/ nb3, + /*.lp0 =*/ ggml_get_op_params_i32(op, 0), + /*.lp1 =*/ ggml_get_op_params_i32(op, 2), + /*.lp2 =*/ ggml_get_op_params_i32(op, 4), + /*.lp3 =*/ ggml_get_op_params_i32(op, 6), }; auto pipeline = ggml_metal_library_get_pipeline_pad(lib, op); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index 9c4853a9d872..ad24ef22a6ab 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -82,6 +82,7 @@ int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx); int ggml_metal_op_cross_entropy_loss_back(ggml_metal_op_t ctx, int idx); int ggml_metal_op_row_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_repeat_back (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_get_rows_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_silu_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_l2_norm (ggml_metal_op_t ctx, int idx); int ggml_metal_op_group_norm (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/kernels/misc.metal b/ggml/src/ggml-metal/kernels/misc.metal index 11104b4d8d12..6206ef9eae2e 100644 --- a/ggml/src/ggml-metal/kernels/misc.metal +++ b/ggml/src/ggml-metal/kernels/misc.metal @@ -127,11 +127,13 @@ kernel void kernel_pad_impl( const int32_t k0 = tgpig.x/args.ne1; const int32_t i1 = tgpig.x - k0*args.ne1; - const int32_t i03 = i3; - const int32_t i02 = i2; - const int32_t i01 = i1; + // the source element of dst[i] is src[i - lp]; outside the source it is zero + const int32_t i03 = i3 - args.lp3; + const int32_t i02 = i2 - args.lp2; + const int32_t i01 = i1 - args.lp1; + const bool row = i01 >= 0 && i01 < args.ne01 && i02 >= 0 && i02 < args.ne02 && i03 >= 0 && i03 < args.ne03; - device const T * src0_ptr = (device const T *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01); + device const char * src0_row = src0 + (row ? i03*args.nb03 + i02*args.nb02 + i01*args.nb01 : 0); device T * dst_ptr = (device T *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1); for (int32_t l0 = 0; l0 < 1024; l0 += ntg.x) { @@ -140,8 +142,9 @@ kernel void kernel_pad_impl( break; } - if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) { - dst_ptr[i0] = src0_ptr[i0]; + const int32_t i00 = i0 - args.lp0; + if (row && i00 >= 0 && i00 < args.ne00) { + dst_ptr[i0] = *(device const T *) (src0_row + i00*args.nb00); // a view may have any element stride } else { dst_ptr[i0] = 0.0f; } diff --git a/ggml/src/ggml-metal/kernels/softmax.metal b/ggml/src/ggml-metal/kernels/softmax.metal index a6ad5baaa9bd..b5d7ec685f72 100644 --- a/ggml/src/ggml-metal/kernels/softmax.metal +++ b/ggml/src/ggml-metal/kernels/softmax.metal @@ -363,11 +363,9 @@ kernel void kernel_rms_norm_back_f32( uint sgitg[[simdgroup_index_in_threadgroup]], uint tiisg[[thread_index_in_simdgroup]], uint3 tptg[[threads_per_threadgroup]]) { - const int32_t row = tgpig.x; - - device const float * dz = (device const float *) (src0 + row*args.nb01); - device const float * x = (device const float *) (src1 + row*args.nb11); - device float * dx = (device float *) (dst + row*args.nb1); + device const float * dz = (device const float *) (src0 + tgpig.x*args.nb01 + tgpig.y*args.nb02 + tgpig.z*args.nb03); + device const float * x = (device const float *) (src1 + tgpig.x*args.nb11 + tgpig.y*args.nb12 + tgpig.z*args.nb13); + device float * dx = (device float *) (dst + tgpig.x*args.nb1 + tgpig.y*args.nb2 + tgpig.z*args.nb3); float lxx = 0.0f; float lxdz = 0.0f; @@ -400,11 +398,9 @@ kernel void kernel_soft_max_back_f32( uint sgitg[[simdgroup_index_in_threadgroup]], uint tiisg[[thread_index_in_simdgroup]], uint3 tptg[[threads_per_threadgroup]]) { - const int32_t row = tgpig.x; - - device const float * dy = (device const float *) (src0 + row*args.nb01); - device const float * y = (device const float *) (src1 + row*args.nb11); - device float * dx = (device float *) (dst + row*args.nb1); + device const float * dy = (device const float *) (src0 + tgpig.x*args.nb01 + tgpig.y*args.nb02 + tgpig.z*args.nb03); + device const float * y = (device const float *) (src1 + tgpig.x*args.nb11 + tgpig.y*args.nb12 + tgpig.z*args.nb13); + device float * dx = (device float *) (dst + tgpig.x*args.nb1 + tgpig.y*args.nb2 + tgpig.z*args.nb3); float ldot = 0.0f; for (int i = tpitg.x; i < args.ne00; i += tptg.x) { @@ -456,3 +452,27 @@ kernel void kernel_repeat_back_f32( d[k0] = acc; } } + +// GET_ROWS_BACK: one threadgroup per dst row r; each thread owns columns and sums, in id order, +// every src0 row whose id is r. No atomics, and the order of the sum is the CPU reference's. +kernel void kernel_get_rows_back_f32( + constant ggml_metal_kargs_get_rows_back & args, + device const char * src0, + device const int32_t * ids, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t r = tgpig.x; + device float * d = (device float *) (dst + r*args.nb1); + + for (int c = tpitg.x; c < args.ne00; c += tptg.x) { + float acc = 0.0f; + for (int i = 0; i < args.nr; i++) { + if (ids[i] == r) { + acc += ((device const float *) (src0 + i*args.nb01))[c]; + } + } + d[c] = acc; + } +} diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 068f90e81e3f..36559cd310c3 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -3516,22 +3516,28 @@ struct test_rms_norm_back : public test_case { const ggml_type type; const std::array ne; const float eps; + const bool strided; // a and b are row views into tensors twice as wide (a per-head norm) std::string vars() override { - return VARS_TO_STR3(type, ne, eps); + return VARS_TO_STR4(type, ne, eps, strided); } test_rms_norm_back(ggml_type type = GGML_TYPE_F32, std::array ne = {64, 5, 4, 3}, - float eps = 1e-6f) - : type(type), ne(ne), eps(eps) {} + float eps = 1e-6f, bool strided = false) + : type(type), ne(ne), eps(eps), strided(strided) {} - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); + ggml_tensor * input(ggml_context * ctx, const char * name) { + std::array ne_full = ne; + ne_full[0] *= strided ? 2 : 1; + ggml_tensor * t = ggml_new_tensor(ctx, type, 4, ne_full.data()); + ggml_set_name(t, name); + return strided ? ggml_view_4d(ctx, t, ne[0], ne[1], ne[2], ne[3], t->nb[1], t->nb[2], t->nb[3], 0) : t; + } - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(b, "b"); + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a = input(ctx, "a"); + ggml_tensor * b = input(ctx, "b"); ggml_tensor * out = ggml_rms_norm_back(ctx, a, b, eps); ggml_set_name(out, "out"); @@ -9330,6 +9336,7 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, false, eps, true)); test_cases.emplace_back(new test_rms_norm_back(GGML_TYPE_F32, { n, 5, 4, 3 }, eps)); + test_cases.emplace_back(new test_rms_norm_back(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, true)); test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false)); test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, true)); test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false, true)); @@ -10137,6 +10144,13 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_pad()); test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {33, 17, 2, 1}, 4, 3, true)); // circular test_cases.emplace_back(new test_pad_ext()); + // left padding on one dim at a time, and on views (a training graph pads a strided view): + // Metal refused every left pad until it read src[i - lp] + for (int tfrm : {0, 1, 2}) { + test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {128, 64, 3, 2}, 3, 0, 0, 0, 0, 0, 0, 0, tfrm)); + test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {128, 64, 3, 2}, 0, 0, 5, 2, 0, 0, 0, 0, tfrm)); + test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {128, 64, 3, 2}, 0, 0, 0, 0, 2, 0, 1, 1, tfrm)); + } test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 1, 1, 1}, 1, 0, false)); test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 2, 1, 1}, 1, 0, false)); test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 16, 1, 1}, 0, 1, false)); From 32b86575ff9b3fef788bd6df61f816e9e9f7bce4 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Tue, 13 Oct 2026 23:39:19 -0500 Subject: [PATCH 05/43] batch: a per-token label makes a position a training row, carried into every ubatch P1 step 1 (continuum #5032 section 9.6): training is inference plus extras that switch on and off, so the switch is in the batch. llama_batch gains `label`: the target token at that position, -1 for none, and NULL for an inference batch, which is unchanged. llama_batch_init allocates it filled with -1. The allocator carries labels into each ubatch the same way it carries the output flags, counts them (get_n_labels), and refuses by name a labelled position that is not an output, since the loss reads that position's logits row. Next (step 2): decode() runs a ubatch with labels through ggml_opt in the same loop, and llama_context::opt_walk::train_chunk, the second copy of decode's batch loop, is deleted. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- include/llama.h | 6 ++++++ src/llama-batch.cpp | 31 +++++++++++++++++++++++++++++++ src/llama-batch.h | 6 ++++++ src/llama-kv-cache-dsv4.cpp | 1 + 4 files changed, 44 insertions(+) diff --git a/include/llama.h b/include/llama.h index 3135ea509aa2..53d747830d8b 100644 --- a/include/llama.h +++ b/include/llama.h @@ -252,6 +252,11 @@ extern "C" { // - if embeddings: all tokens are output // - if not: only the last token is output // ) + // - label : the token the model should have predicted at this position, which makes the + // position a training row: llama_decode then also runs the loss, the backward and + // the optimizer write of the context's trainable adapter (llama_opt_init), in the + // same graph as the batch's other rows. -1 marks a position with no label. + // (if set to NULL, the batch is inference only, unchanged) // typedef struct llama_batch { int32_t n_tokens; @@ -262,6 +267,7 @@ extern "C" { int32_t * n_seq_id; llama_seq_id ** seq_id; int8_t * logits; // TODO: rename this to "output" + llama_token * label; } llama_batch; enum llama_model_kv_override_type { diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f48f..7305cda7062b 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -157,6 +157,20 @@ bool llama_batch_allocr::init( n_outputs += batch.logits[i] != 0; } + // a labelled position is a training row: its logits row is what the loss reads + if (batch.label) { + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.label[i] < 0) { + continue; + } + if (batch.logits[i] == 0) { + LLAMA_LOG_ERROR("%s: position %d has a label but is not an output: a training row must be output\n", __func__, i); + return false; + } + ++n_labels; + } + } + has_cpl = false; // determine coupled sequences @@ -224,6 +238,7 @@ bool llama_batch_allocr::init( /*.seq_id_unq =*/ this->seq_id_unq.data(), /*.seq_idx =*/ this->seq_idx.data(), /*.output =*/ batch.logits, + /*.label =*/ batch.label, /*.data =*/ {}, }; @@ -430,6 +445,7 @@ llama_ubatch llama_batch_allocr::ubatch_reserve(uint32_t n_seq_tokens, uint32_t /*.seq_id_unq =*/ udata->seq_id_unq.data(), /*.seq_idx =*/ udata->seq_idx.data(), /*.output =*/ udata->output.data(), + /*.label =*/ nullptr, /*.data =*/ std::move(udata), }; @@ -448,6 +464,10 @@ uint32_t llama_batch_allocr::get_n_outputs() const { return n_outputs; } +uint32_t llama_batch_allocr::get_n_labels() const { + return n_labels; +} + uint32_t llama_batch_allocr::get_n_used() const { return n_used; } @@ -722,6 +742,7 @@ llama_ubatch llama_batch_allocr::split_seq(uint32_t n_ubatch) { void llama_batch_allocr::clear() { n_outputs = 0; + n_labels = 0; batch = {}; @@ -764,6 +785,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u udata->seq_id_unq.resize(0); udata->seq_idx .resize(LLAMA_MAX_SEQ, -1); udata->output .resize(n_tokens); + udata->label .resize(batch.label ? n_tokens : 0); udata->seq_id_data.reserve(n_tokens); @@ -789,6 +811,9 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u udata->n_seq_id[i] = batch.n_seq_id[idxs[i]]; udata->output[i] = batch.logits[idxs[i]]; + if (batch.label) { + udata->label[i] = batch.label[idxs[i]]; + } for (int s = 0; s < udata->n_seq_id[i]; ++s) { const llama_seq_id seq_id = batch.seq_id[idxs[i]][s]; @@ -831,6 +856,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u /*.seq_id_unq =*/ udata->seq_id_unq.data(), /*.seq_idx =*/ udata->seq_idx.data(), /*.output =*/ udata->output.data(), + /*.label =*/ batch.label ? udata->label.data() : nullptr, /*.data =*/ std::move(udata), }; @@ -939,6 +965,7 @@ struct llama_batch llama_batch_get_one( /*n_seq_id =*/ nullptr, /*seq_id =*/ nullptr, /*logits =*/ nullptr, + /*label =*/ nullptr, }; } @@ -951,6 +978,7 @@ struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_ /*n_seq_id =*/ nullptr, /*seq_id =*/ nullptr, /*logits =*/ nullptr, + /*label =*/ nullptr, }; if (embd) { @@ -968,6 +996,8 @@ struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_ batch.seq_id[n_tokens_alloc] = nullptr; batch.logits = (int8_t *) malloc(sizeof(int8_t) * n_tokens_alloc); + batch.label = (llama_token *) malloc(sizeof(llama_token) * n_tokens_alloc); + std::fill(batch.label, batch.label + n_tokens_alloc, -1); // no position trains until labelled return batch; } @@ -984,4 +1014,5 @@ void llama_batch_free(struct llama_batch batch) { free(batch.seq_id); } if (batch.logits) free(batch.logits); + if (batch.label) free(batch.label); } diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d4a04..6441fe839eee 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -50,6 +50,7 @@ struct llama_ubatch { llama_seq_id * seq_id_unq; // [n_seqs_unq] | s | seq_id int32_t * seq_idx; // [LLAMA_MAX_SEQ] | - | seq_idx int8_t * output; // [n_tokens] | i | - + llama_token * label; // [n_tokens] | i | training target, -1 for none (nullptr: inference) struct data_t { std::vector token; @@ -60,6 +61,7 @@ struct llama_ubatch { std::vector seq_id_unq; std::vector seq_idx; std::vector output; + std::vector label; std::vector seq_id_data; }; @@ -87,6 +89,9 @@ class llama_batch_allocr { uint32_t get_n_tokens() const; uint32_t get_n_outputs() const; + + // labelled positions (training rows) in the batch; 0 for an inference batch + uint32_t get_n_labels() const; uint32_t get_n_used() const; // the array of output indices in the order they were encountered during the ubatch splitting @@ -136,6 +141,7 @@ class llama_batch_allocr { uint32_t n_embd; uint32_t n_seq_max; uint32_t n_outputs; + uint32_t n_labels = 0; std::array seq_id_0 = {{ 0 }}; // default sequence id diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 948d08146fbf..b073ed303e5d 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -164,6 +164,7 @@ static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) { /*.seq_id_unq =*/ data->seq_id_unq.data(), /*.seq_idx =*/ data->seq_idx.data(), /*.output =*/ data->output.data(), + /*.label =*/ nullptr, /*.data =*/ data, }; From 0b8ddcef7f6accb0eb452e897418152534096cb4 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 01:02:18 -0500 Subject: [PATCH 06/43] llama: a labelled ubatch is a training step inside decode; train_chunk's copy of the batch loop is gone P1 step 2 (continuum #5032 section 9.6). process_ubatch has one switch: a ubatch whose rows carry labels, on a context with a trainable adapter, is built fresh and runs through ggml_opt (llama_context::opt_ubatch: the loss over the labelled rows, the backward, the optimizer step at a period's end, and the active walk's taps). Every other ubatch runs graph_compute as before. opt_walk::train_chunk, which repeated decode's batch loop, now fills a chunk's labels and calls llama_decode. - cparams.training joins the graph reuse key. It was missing, so a serving decode with the same shape as a training ubatch could reuse the training graph. - Every batch view (server, speculative, mtmd, batched-bench, perplexity, parallel) carries its slice of the labels. - test-walk-exact: the device-refusal case is skipped on a CPU-only run, because the budget caps GPU memory only (ggml_opt_alloc). Its old skip test (graph size == 0) never fired: the CPU graph is measured in host buffers. Measured (fork bdd3cd55b + 32b86575f + this, M5, --device none, qwen3.5-4b hybrid Q4_K_M, 2026-10-14T04:46Z-06:01Z): every exact row is 1.0000 through the new path (exact walk vs one graph, checkpointed state, f16 and q8_0 caches, context + reply, recompute, rollback slots). The plain walk's rows (0.8376, 0.6710) are close to the pre-P1 Metal values (0.8378, 0.6714); the difference is backend, not path. There is no pre-P1 CPU run to compare against. Still owed: the Metal run, for the device-refusal case (the 4B, P1 vs bdd3cd55b, in a quiet window). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- common/speculative.cpp | 2 + examples/parallel/parallel.cpp | 1 + src/llama-context-opt.cpp | 234 +++++++++++--------------- src/llama-context.cpp | 25 ++- src/llama-context.h | 13 ++ src/llama-graph.h | 1 + tests/test-walk-exact.cpp | 8 +- tools/batched-bench/batched-bench.cpp | 1 + tools/mtmd/mtmd-helper-common.h | 2 + tools/perplexity/perplexity.cpp | 1 + tools/server/server-context.cpp | 1 + 11 files changed, 146 insertions(+), 143 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 4eef2212e751..321c6de3effb 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -633,6 +633,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, + /*.label =*/ nullptr, }; const int32_t rc = llama_encode(ctx_dft, enc_batch); if (rc != 0) { @@ -1126,6 +1127,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, + /*.label =*/ nullptr, }; int32_t rc = llama_encode(ctx_dft, enc_batch); diff --git a/examples/parallel/parallel.cpp b/examples/parallel/parallel.cpp index a46400c5b943..45fded7f54f5 100644 --- a/examples/parallel/parallel.cpp +++ b/examples/parallel/parallel.cpp @@ -400,6 +400,7 @@ int main(int argc, char ** argv) { batch.n_seq_id + i, batch.seq_id + i, batch.logits + i, + batch.label ? batch.label + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 6cdd89002c4e..1a2058f2e6f0 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -227,168 +227,124 @@ bool llama_context::opt_walk::decode_span(uint32_t p0, uint32_t p1) { // one training chunk [pos_ctx, pos_ctx + n_chunk): forward and backward over it alone, attending // to the cached context before it +// One chunk of a walk is one llama_decode with labels: the batch loop is decode's, and the +// extras switch on in process_ubatch (llama_context::opt_ubatch). This function only fills the +// chunk's tokens and labels and says whose walk it is. bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { - const llama_model & model = lctx.model; - llama_cparams & cparams = lctx.cparams; - ggml_opt_context_t opt_ctx = lctx.opt_ctx; - llama_walk_fill(batch, tokens, pos_ctx, n_chunk); for (uint32_t pos_batch = 0; pos_batch < n_chunk; ++pos_batch) { // an exact walk's context chunk has no label but still trains (its K/V carry later // chunks' gradient): one output row, weighted zero, so the graph has its loss batch.logits[pos_batch] = labels[pos_ctx + pos_batch] >= 0 || (xc != nullptr && xc->n_labels == 0 && pos_batch + 1 == n_chunk); + batch.label[pos_batch] = labels[pos_ctx + pos_batch]; } - // output_all = false: the batch's own logits flags decide which rows exist - if (!lctx.balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, false)) { - LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); + lctx.opt_walk_active = this; + lctx.opt_result_active = xc != nullptr && xc->n_labels == 0 ? nullptr : result; + lctx.opt_train_unlabelled = xc != nullptr && xc->n_labels == 0; + const int ret = lctx.decode(batch); + lctx.opt_walk_active = nullptr; + lctx.opt_result_active = nullptr; + lctx.opt_train_unlabelled = false; + std::fill(batch.label, batch.label + n_chunk, -1); + + if (ret != 0) { + // a refusal has already stopped the epoch by name (opt_fail); anything else is decode's + if (!lctx.opt_alloc_failed.load()) { + lctx.opt_fail("the training chunk failed to decode (llama_decode returned " + std::to_string(ret) + ")"); + } return false; } - - lctx.n_queued_tokens += lctx.balloc->get_n_tokens(); - - lctx.embd_seq.clear(); - - const uint32_t n_outputs_all = lctx.balloc->get_n_outputs(); - // Unreachable: the walk trains only chunks of a labelled run. It RETURNS from the whole - // window, since a later chunk would need this one in the cache. - if (n_outputs_all == 0) { - LLAMA_LOG_ERROR("%s: a training batch with no labelled position (every label masked): nothing to learn from it\n", __func__); - return false; + if (callback && !(xc != nullptr && xc->n_labels == 0)) { + callback(train, lctx.opt_ctx, dataset, result, idata_in_loop + pos_ctx/n_batch + 1, ndata_in_loop, t_loop_start); } + return true; +} - auto mctx = lctx.memory->init_batch(*lctx.balloc, cparams.n_ubatch, true); - if (!mctx || mctx->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { - LLAMA_LOG_ERROR("%s: could not initialize batch\n", __func__); - return false; - } - - // reserve output buffer - if (lctx.output_reserve(n_outputs_all) < n_outputs_all) { - LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); - GGML_ABORT("TODO: handle this error"); - }; - - uint32_t pos_batch = 0; - do { - const auto & ubatch = mctx->get_ubatch(); +bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch) { + opt_walk * walk = opt_walk_active; + const llama_walk_chunk * xc = walk != nullptr ? walk->xc : nullptr; - // the labelled rows of this ubatch, as decode counts its outputs - { - uint32_t n_outputs_ubatch = 0; - for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { - n_outputs_ubatch += (uint32_t) (ubatch.output[i] != 0); + struct ggml_context * ctx_compute_opt; + { + const size_t size_gf = ggml_graph_size(gf); + const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); + struct ggml_init_params params = { + /*.mem_size =*/ size_meta, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx_compute_opt = ggml_init(params); + } + ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); + if (xc != nullptr) { + ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); + for (const auto & io : res->t_walk) { + if (io.tk) { + ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); + ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); } - lctx.n_outputs = n_outputs_ubatch; - } - - if (!mctx->apply()) { - LLAMA_LOG_ERROR("%s: failed to update the memory context\n", __func__); - return false; - } - - auto * res = lctx.gf_res_prev.get(); - - const auto gparams = lctx.graph_params(res, ubatch, mctx.get(), lctx.graph_type()); - - res->reset(); - - auto * gf = model.build_graph(gparams); - - struct ggml_context * ctx_compute_opt; - { - const size_t size_gf = ggml_graph_size(gf); - const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); - struct ggml_init_params params = { - /*.mem_size =*/ size_meta, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ctx_compute_opt = ggml_init(params); } - ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); - if (xc != nullptr) { - ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); - for (const auto & io : res->t_walk) { - if (io.tk) { - ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); - ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); - } + for (const auto & ws : res->t_walk_state) { + if (ws.tap) { + ggml_opt_accumulate_grad(opt_ctx, ws.tap, walk->acc.state.at(ws.cache)); } - for (const auto & ws : res->t_walk_state) { - // a state read without its write reporting the exit (build_walk_state_exit) would drop - // the gradient later chunks put on it: a wrong step, silently. Refused by name. - if (ws.tap && ws.gs == nullptr && lctx.cparams.walk_state_surrogate) { - ggml_free(ctx_compute_opt); - lctx.opt_fail(std::string("this model's recurrent state write (") + ws.cache->name - + ") does not report its exit to the walk (llm_graph_context::build_walk_state_exit): its gradient across chunks would be dropped"); - return false; - } - if (ws.tap) { - ggml_opt_accumulate_grad(opt_ctx, ws.tap, acc.state.at(ws.cache)); - } - } - } - if (!ggml_opt_alloc(opt_ctx, train)) { - ggml_free(ctx_compute_opt); - lctx.opt_fail(ggml_opt_refusal(opt_ctx), - std::string("the ") + (train ? "training" : "evaluation") + " graph was refused (window " + std::to_string(n_ctx) + ")"); - return false; } + } + const bool train = walk == nullptr || walk->train; + if (!ggml_opt_alloc(opt_ctx, train)) { + ggml_free(ctx_compute_opt); + opt_fail(ggml_opt_refusal(opt_ctx), + std::string("the ") + (train ? "training" : "evaluation") + " graph was refused (window " + std::to_string(walk != nullptr ? walk->n_ctx : ubatch.n_tokens) + ")"); + return false; + } - res->set_inputs(&ubatch); - { - // one-hot rows for the labelled positions only, in the order the ubatch emits its - // outputs (the same order the logits rows take) - struct ggml_tensor * labels_t = ggml_opt_labels(opt_ctx); - GGML_ASSERT(labels_t->ne[1] == (int64_t) lctx.n_outputs); - ggml_set_zero(labels_t); - const float onef = 1.0f; - uint32_t row = 0; - for (uint32_t pos_ubatch = 0; pos_ubatch < ubatch.n_tokens; ++pos_ubatch) { - if (ubatch.output[pos_ubatch] == 0) { - continue; - } - const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch; - // an exact walk's context chunk: its one row has no label (weighted zero) - GGML_ASSERT((xc != nullptr && labels[ilabel] < 0) || (labels[ilabel] >= 0 && labels[ilabel] < labels_t->ne[0])); - if (labels[ilabel] >= 0) { - ggml_backend_tensor_set(labels_t, &onef, (row*labels_t->ne[0] + labels[ilabel])*sizeof(float), sizeof(float)); - } - ++row; + res->set_inputs(&ubatch); + { + // one-hot rows for the labelled positions, in the order the ubatch emits its outputs + // (the order the logits rows take); an unlabelled output row stays zero, so it adds no + // LOCAL loss (its sequence's context still carries gradient to later labelled rows) + struct ggml_tensor * labels_t = ggml_opt_labels(opt_ctx); + GGML_ASSERT(labels_t->ne[1] == (int64_t) n_outputs); + ggml_set_zero(labels_t); + const float onef = 1.0f; + uint32_t row = 0; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (ubatch.output[i] == 0) { + continue; } - GGML_ASSERT(row == lctx.n_outputs); - } - if (xc != nullptr) { - // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the - // accumulator takes this chunk's own - for (const auto & ws : res->t_walk_state) { - ggml_tensor * a = acc.state.at(ws.cache); - if (ws.gs) { - GGML_ASSERT(ggml_nbytes(a) == ggml_nbytes(ws.gs)); - ggml_backend_tensor_set(ws.gs, a->data, 0, ggml_nbytes(ws.gs)); - } - memset(a->data, 0, ggml_nbytes(a)); + const llama_token y = ubatch.label != nullptr ? ubatch.label[i] : -1; + GGML_ASSERT(y < labels_t->ne[0]); + if (y >= 0) { + ggml_backend_tensor_set(labels_t, &onef, (row*labels_t->ne[0] + y)*sizeof(float), sizeof(float)); } + ++row; } - ggml_opt_eval(opt_ctx, xc != nullptr && xc->n_labels == 0 ? nullptr : result); - if (const char * why = ggml_opt_refusal(opt_ctx); why[0] != '\0') { - // the backend failed the graph (ggml_opt_eval): the run fails, as a refused graph does - ggml_free(ctx_compute_opt); - lctx.opt_fail(why); - return false; - } - if (xc != nullptr && !res->t_walk_state.empty()) { - acc.state_ready = true; - } - if (callback && !(xc != nullptr && xc->n_labels == 0)) { - callback(train, opt_ctx, dataset, result, idata_in_loop + (pos_ctx + pos_batch)/n_batch + 1, ndata_in_loop, t_loop_start); + GGML_ASSERT(row == n_outputs); + } + if (xc != nullptr) { + // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the + // accumulator takes this chunk's own + for (const auto & ws : res->t_walk_state) { + ggml_tensor * a = walk->acc.state.at(ws.cache); + if (ws.gs) { + GGML_ASSERT(ggml_nbytes(a) == ggml_nbytes(ws.gs)); + ggml_backend_tensor_set(ws.gs, a->data, 0, ggml_nbytes(ws.gs)); + } + memset(a->data, 0, ggml_nbytes(a)); } - ggml_free(ctx_compute_opt); - - pos_batch += ubatch.n_tokens; - } while (mctx->next()); + } + ggml_opt_eval(opt_ctx, opt_result_active); + ggml_free(ctx_compute_opt); + if (const char * why = ggml_opt_refusal(opt_ctx); why[0] != '\0') { + // the backend failed the graph (ggml_opt_eval): the run fails, as a refused graph does + opt_fail(why); + return false; + } + if (xc != nullptr && !res->t_walk_state.empty()) { + walk->acc.state_ready = true; + } return true; } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a296f0c2cbdb..f8a9c5014848 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1334,11 +1334,23 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll auto * res = gf_res_prev.get(); auto * gf = res->get_gf(); + // a ubatch with labelled rows is a training step on this context's trainable adapter + // (llama_opt_init): the same graph with the extras on (opt_ubatch). It is built fresh, and + // cparams.training is part of the reuse key, so a serving graph never reuses it. + bool train = false; + if (opt_ctx != nullptr) { + train = opt_train_unlabelled; + for (uint32_t i = 0; !train && ubatch.label != nullptr && i < ubatch.n_tokens; ++i) { + train = ubatch.label[i] >= 0; + } + cparams.training = train; + } + // the new graph parameters // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters const auto gparams = graph_params(res, ubatch, mctx, gtype); - if (!graph_reuse_disable && res->can_reuse(gparams)) { + if (!train && !graph_reuse_disable && res->can_reuse(gparams)) { //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__); // with pipeline parallelism, the previous graph_compute_async may still be running @@ -1367,6 +1379,15 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return nullptr; } + if (train) { + if (!opt_ubatch(res, gf, ubatch)) { + ret = GGML_STATUS_FAILED; + return nullptr; + } + ret = GGML_STATUS_SUCCESS; + return res; + } + if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) { LLAMA_LOG_ERROR("%s: failed to allocate graph\n", __func__); ret = GGML_STATUS_ALLOC_FAILED; @@ -1862,7 +1883,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } // extract logits - if (logits.data && t_logits && n_outputs > 0 && needs_raw_logits(ubatch, sampling.samplers)) { + if (logits.data && t_logits && n_outputs > 0 && opt_walk_active == nullptr && needs_raw_logits(ubatch, sampling.samplers)) { ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits); GGML_ASSERT(backend_res != nullptr); GGML_ASSERT(logits.data != nullptr); diff --git a/src/llama-context.h b/src/llama-context.h index 2f3fc48b6efd..d2758811f6cb 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -244,6 +244,12 @@ struct llama_context { // the step that failed, for the log void opt_fail(std::string why, const std::string & what = ""); + // THE EXTRAS A LABELLED UBATCH SWITCHES ON (llama-context-opt.cpp). process_ubatch runs a + // ubatch whose rows carry labels through ggml_opt instead of graph_compute: the same graph, + // plus the loss over the labelled rows, the backward and, at a period's end, the optimizer + // step. A walk driving the decode adds its taps to that graph. + bool opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch); + // // output // @@ -388,6 +394,13 @@ struct llama_context { uint32_t opt_walk_horizon_req = 0; // llama_opt_params::walk_horizon, as requested uint32_t opt_walk_horizon = 0; // this window's: starts at the request, shrinks to fit the device + // set by a walk for the decode it drives: whose taps join the graph, where the loss lands, + // whether a ubatch without labels still trains (an exact walk's context chunk), and that the + // logits rows are the loss's alone (no host copy) + opt_walk * opt_walk_active = nullptr; + ggml_opt_result_t opt_result_active = nullptr; + bool opt_train_unlabelled = false; + ggml_threadpool_t threadpool = nullptr; ggml_threadpool_t threadpool_batch = nullptr; diff --git a/src/llama-graph.h b/src/llama-graph.h index 3b181afed3fb..22b821a7564a 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -875,6 +875,7 @@ struct llm_graph_params { cparams.embeddings_nextn == other.cparams.embeddings_nextn && cparams.embeddings_nextn_masked == other.cparams.embeddings_nextn_masked && cparams.causal_attn == other.cparams.causal_attn && + cparams.training == other.cparams.training && arch == other.arch && gtype == other.gtype && cvec == other.cvec && diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index b46144a1e0d0..f0babeba974b 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -496,8 +496,12 @@ int main(int argc, char ** argv) { step_delta(params, model, init, tokens, all, WINDOW / 4, true); const size_t graph_bytes = g_last_graph_bytes; printf(" device graph: %.1f MiB per chunk\n", graph_bytes / 1048576.0); - if (graph_bytes == 0) { - printf(" device refusal: skipped (no device graph measured: a CPU-only run)\n"); + // the budget caps GPU memory only (ggml_opt_alloc), so a CPU-only run cannot refuse; it + // still measures a graph (in host buffers), which is why the graph size alone was the + // wrong test for "a CPU-only run" + const bool cpu_only = !params.devices.empty() && params.devices[0] == nullptr; + if (graph_bytes == 0 || cpu_only) { + printf(" device refusal: skipped (a CPU-only run: the device budget caps GPU memory only)\n"); } else { llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); const std::vector before = adapter_params(adapter); diff --git a/tools/batched-bench/batched-bench.cpp b/tools/batched-bench/batched-bench.cpp index e2dcd0b2e717..04cb25c62ac1 100644 --- a/tools/batched-bench/batched-bench.cpp +++ b/tools/batched-bench/batched-bench.cpp @@ -91,6 +91,7 @@ int llama_batched_bench(int argc, char ** argv) { batch.n_seq_id + i, batch.seq_id + i, batch.logits + i, + batch.label ? batch.label + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); diff --git a/tools/mtmd/mtmd-helper-common.h b/tools/mtmd/mtmd-helper-common.h index f907346c7b58..19c0f9ed0942 100644 --- a/tools/mtmd/mtmd-helper-common.h +++ b/tools/mtmd/mtmd-helper-common.h @@ -96,6 +96,7 @@ struct decode_embd_batch { /*n_seq_id =*/ n_seq_id.data(), /*seq_id =*/ seq_ids.data(), /*logits =*/ logits.data(), + /*label =*/ nullptr, }; } @@ -179,6 +180,7 @@ struct decode_embd_batch { /*n_seq_id =*/ batch.n_seq_id + offset, /*seq_id =*/ batch.seq_id + offset, /*logits =*/ batch.logits + offset, + /*label =*/ batch.label ? batch.label + offset : nullptr, }; } }; diff --git a/tools/perplexity/perplexity.cpp b/tools/perplexity/perplexity.cpp index ba41287d8e3d..80c834bb79b5 100644 --- a/tools/perplexity/perplexity.cpp +++ b/tools/perplexity/perplexity.cpp @@ -674,6 +674,7 @@ static bool decode_helper(llama_context * ctx, llama_batch & batch, std::vector< batch.n_seq_id + i, batch.seq_id + i, batch.logits + i, + batch.label ? batch.label + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7294a8d876e3..8fe17bba52e8 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -205,6 +205,7 @@ struct server_batch { batch.n_seq_id + off, batch.seq_id + off, batch.logits + off, + batch.label ? batch.label + off : nullptr, }; return view; From 725c893669d5f99a7bc9b2475b9d9fd537d83927 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 07:43:49 -0500 Subject: [PATCH 07/43] llama: the trainable layer applies to labelled rows only, and a trained run is restored P1 step 3a (continuum #5032 section 9.6). The objective is each labelled run (a reply) given everything before it as served: only labelled rows see the shadow (the adapter being trained), so the context before a run carries no gradient and is never trained. The exact walk's backward now covers the reply only, not the window. - The shadow is added on the labelled rows by row selection (llm_graph_context::shadow_add): those rows of the base and of the shadow's term are gathered, added, and written into a copy of the base, so every other row is the stable graph's byte for byte. A ubatch with no label never applies the shadow; the label count is in the graph reuse key. (A 0/1 mask multiply was rejected in review: inf or NaN times zero is NaN, and x + 0 turns -0 into +0.) - SET_ROWS gains the destination's gradient in ggml backward (the result's gradient except at the replaced rows), which row selection needs. - The walk goes run by run: the context before a run decodes under her stable set; the run decodes under the shadow, trains (exact: in reverse, with the K/V and recurrent-state adjoints inside the run only; plain: a step per chunk), and is then restored, its cells and recurrent state decoded again under her stable set. One checkpoint helper and one rewind serve both walks; the plain walk's scratch sequence (n_seq_max >= 2) is gone. - decode_span no longer toggles cparams.training: process_ubatch sets it per ubatch (one of the four forward differences in section 9.6 removed). - The batch allocator refuses, before any mutation, a labelled row that is not an output or has more than one sequence, and a sequence whose labelled rows are not one run ending its rows in the batch. A training ubatch whose output rows are not exactly its labelled rows fails by name: the loss is their mean. - Tests: test-walk-plan for the per-run marks; test-walk-exact case 8, the cache after either walk byte-identical to a fresh stable decode; test-backend-ops SET_ROWS with grad=1. Measured: in the follow-up receipts (CPU and Metal, this SHA). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml.c | 13 +- src/llama-batch.cpp | 17 +- src/llama-batch.h | 9 + src/llama-context-opt.cpp | 363 +++++++++++++++++----------------- src/llama-context.cpp | 10 +- src/llama-context.h | 14 +- src/llama-graph.cpp | 51 ++++- src/llama-graph.h | 31 +++ src/llama-walk.cpp | 45 +++-- src/llama-walk.h | 22 ++- tests/test-backend-ops.cpp | 17 +- tests/test-walk-exact.cpp | 63 +++++- tests/test-walk-plan.cpp | 31 +-- tools/server/server-train.cpp | 6 +- 14 files changed, 449 insertions(+), 243 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 5f572229d5f9..ab9c8911059e 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7212,13 +7212,18 @@ static void ggml_compute_backward( } break; case GGML_OP_SET_ROWS: { // result = the destination (src2) with its rows src1 replaced by src0. The gradient - // on src0 is the result's gradient at those rows. The destination's gradient (the - // result's everywhere else) is not built: a written buffer (a cache) is a constant. - GGML_ASSERT(!src2_needs_grads && "set_rows: the destination it writes carries no gradient"); - if (src0_needs_grads) { + // on src0 is the result's gradient at those rows; the destination's is the result's + // everywhere else (a cache it writes is a constant and needs none). + if (src0_needs_grads || src2_needs_grads) { GGML_ASSERT(src1->type == GGML_TYPE_I32 && "set_rows backward gathers with I32 row indices"); + } + if (src0_needs_grads) { ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, ggml_get_rows(ctx, grad, src1), src0)); } + if (src2_needs_grads) { + struct ggml_tensor * replaced = ggml_fill(ctx, ggml_get_rows(ctx, grad, src1), 0.0f); + ggml_add_or_set(ctx, cgraph, isrc2, ggml_set_rows(ctx, ggml_dup(ctx, grad), replaced, src1)); + } } break; case GGML_OP_GET_ROWS: { if (src0_needs_grads) { diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 7305cda7062b..e4b06e21b994 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -157,16 +157,31 @@ bool llama_batch_allocr::init( n_outputs += batch.logits[i] != 0; } - // a labelled position is a training row: its logits row is what the loss reads + // a labelled position is a training row: its logits row is what the loss reads. A run of + // them trains given everything before it as served (the shadow adapter applies to labelled + // rows only) and is restored to her stable set after, so in one batch a sequence's labelled + // rows are one run that ends its rows: anything after it would read the run through the + // shadow instead of as served. if (batch.label) { + bool in_run[LLAMA_MAX_SEQ] = {}; for (int32_t i = 0; i < batch.n_tokens; ++i) { + const llama_seq_id s = batch.seq_id[i][0]; if (batch.label[i] < 0) { + if (in_run[s]) { + LLAMA_LOG_ERROR("%s: position %d follows a labelled run of sequence %d in one batch: a run ends its sequence's rows\n", __func__, i, s); + return false; + } continue; } if (batch.logits[i] == 0) { LLAMA_LOG_ERROR("%s: position %d has a label but is not an output: a training row must be output\n", __func__, i); return false; } + if (batch.n_seq_id[i] != 1) { + LLAMA_LOG_ERROR("%s: position %d has a label and %d sequences: a training row belongs to one sequence\n", __func__, i, batch.n_seq_id[i]); + return false; + } + in_run[s] = true; ++n_labels; } } diff --git a/src/llama-batch.h b/src/llama-batch.h index 6441fe839eee..ef1f9908f24f 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -17,6 +17,15 @@ struct llama_ubatch { return b_equal_seqs != 0; } + // the training rows (a label >= 0) + uint32_t n_labels() const { + uint32_t n = 0; + for (uint32_t i = 0; label != nullptr && i < n_tokens; ++i) { + n += label[i] >= 0; + } + return n; + } + // typical for M-RoPE cases: // 0 - sequential position of the tokens/embeddings in the sequence // 1 - y position in the image diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 1a2058f2e6f0..b0fc285625a1 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -102,6 +102,8 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params n_params += llama_set_param(ab.b, param_filter, param_filter_ud); } GGML_ASSERT(n_params > 0 && "the param filter selected none of the adapter's tensors"); + // the trained adapter is the shadow: labelled rows only, never a serving row + opt_shadow = lopt_params.adapter; LLAMA_LOG_INFO("%s: LoRA-only training: %lld adapter tensors trainable, base frozen\n", __func__, (long long) n_params); return; } @@ -129,13 +131,17 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params } } -// THE WALK (Fable's design, continuum 2026-10-06): one pass over a window in one context. A run -// with no label is context she did not write: a plain decode into the cache, no backward. A -// labelled run is her reply: it trains in chunks of n_batch, each attending to everything before -// it read from the cache as a constant (build_attn, cparams.training), then the reply is decoded -// into the cache under the adapter as it now is, so the next reply sees it. Memory is chunk x -// window, never window x window, so her whole lived window trains. Nothing after the last label -// is computed. One per window; a member of the context, so it runs on the context's own state. +// THE WALK (Fable's design, continuum 2026-10-06; objective 2026-10-14): one pass over a window +// in one context. A run with no label is context: a plain decode into the cache, as served, no +// backward. A labelled run is her reply, trained GIVEN EVERYTHING BEFORE IT AS SERVED: only its +// rows see the shadow adapter (the label mask), so the context before it carries no gradient and +// is never trained. The run trains in chunks of n_batch, each attending to everything before it +// in the cache; its chunks write their own cells under the shadow, so a later chunk of the run +// reads an earlier one as the shadow computes it. When the run is done it is RESTORED: its cells +// and the recurrent state are decoded again under her stable set, so the cache is always her +// stable self's and the next run reads this one as served. Memory is chunk x window, never +// window x window, so her whole lived window trains. Nothing after the last label is computed. +// One per window; a member of the context, so it runs on the context's own state. struct llama_context::opt_walk { llama_context & lctx; ggml_opt_dataset_t dataset; @@ -171,8 +177,11 @@ struct llama_context::opt_walk { n_ctx(n_ctx), n_batch(n_batch), plan(llama_walk_plan::of(labels, n_ctx, n_batch)) {} bool yield_point(); - bool decode_span(uint32_t p0, uint32_t p1); + bool decode_span(uint32_t p0, uint32_t p1, bool shadow); + std::vector checkpoint(); + bool rewind(uint32_t p_pop, const std::vector & snap, uint32_t p_to); bool train_chunk(uint32_t pos_ctx, uint32_t n_chunk); + bool restore(const llama_walk_run & run, const std::vector & snap); void run_exact(); void run_plain(); @@ -198,55 +207,87 @@ bool llama_context::opt_walk::yield_point() { return !lctx.opt_stop_requested.load(std::memory_order_relaxed); } -// a plain inference decode of [p0, p1) into the cache: the training graph's attention never -// writes the cache, so context and finished replies reach it this way -bool llama_context::opt_walk::decode_span(uint32_t p0, uint32_t p1) { - llama_cparams & cparams = lctx.cparams; - const bool training = cparams.training; - cparams.training = false; - lctx.gf_res_prev->reset(); +// an inference decode of [p0, p1) into the cache, no loss and no backward. shadow: the labelled +// rows decode under the shadow adapter (a run's forward pass); otherwise every row decodes under +// her stable set (context, and a run restored). +bool llama_context::opt_walk::decode_span(uint32_t p0, uint32_t p1, bool shadow) { bool ok = true; for (uint32_t c0 = p0; c0 < p1 && ok; c0 += n_batch) { if (!yield_point()) { - cparams.training = training; - lctx.gf_res_prev->reset(); return false; // stopped or cancelled at a chunk boundary: not a decode failure } const uint32_t c1 = std::min(c0 + n_batch, p1); llama_walk_fill(batch, tokens, c0, c1 - c0); + for (uint32_t i = 0; shadow && i < c1 - c0; ++i) { + // a label marks the row the shadow sees; the allocator wants it an output row + batch.label [i] = labels[c0 + i]; + batch.logits[i] = labels[c0 + i] >= 0; + } batch.logits[c1 - c0 - 1] = true; + lctx.opt_walk_active = shadow ? this : nullptr; + lctx.opt_forward_only = shadow; ok = lctx.decode(batch) == 0; + lctx.opt_walk_active = nullptr; + lctx.opt_forward_only = false; + std::fill(batch.label, batch.label + (c1 - c0), -1); } - cparams.training = training; - lctx.gf_res_prev->reset(); if (!ok) { - LLAMA_LOG_ERROR("%s: the context decode of [%u, %u) failed\n", __func__, p0, p1); + LLAMA_LOG_ERROR("%s: the decode of [%u, %u) failed\n", __func__, p0, p1); } return ok; } -// one training chunk [pos_ctx, pos_ctx + n_chunk): forward and backward over it alone, attending -// to the cached context before it -// One chunk of a walk is one llama_decode with labels: the batch loop is decode's, and the -// extras switch on in process_ubatch (llama_context::opt_ubatch). This function only fills the -// chunk's tokens and labels and says whose walk it is. +// the recurrent state of the walk's sequence through the sequence-state API (the recurrent part +// only, PARTIAL_ONLY), after waiting for the queued decode; empty for a model with none. A +// hand-rolled per-layer copy skipped that wait and, on Metal, carried a torn mid-graph state into +// the reverse pass (2026-10-11). +std::vector llama_context::opt_walk::checkpoint() { + std::vector snap; + if (recr != nullptr) { + lctx.synchronize(); + snap.resize(lctx.state_seq_get_size(0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)); + GGML_ASSERT(lctx.state_seq_get_data(0, snap.data(), snap.size(), LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == snap.size()); + } + return snap; +} + +// back to position p_pop: the cache holds [0, p_pop), the recurrent state is `snap` (taken at +// p_pop); then forward to p_to, under the shadow for the run's labelled rows. A restore puts back +// the cell's position with its rows. Also run before every retry of a refused chunk: a refused +// chunk graph is refused AFTER its ubatch was applied, so its cells already sit in the cache and +// the recurrent cell's position has moved to the chunk's end; retrying on top of that fails to +// prepare the ubatch at all (the 5090, 2026-10-10: "init_batch: failed to prepare attention +// ubatches" after "retrying at 512"). +bool llama_context::opt_walk::rewind(uint32_t p_pop, const std::vector & snap, uint32_t p_to) { + const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : lctx.memory->seq_rm(0, p_pop, -1); + if (!popped) { + LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u)\n", __func__, p_pop); + lctx.opt_stop_requested.store(true); + return false; + } + if (recr != nullptr) { + lctx.synchronize(); + GGML_ASSERT(lctx.state_seq_set_data(0, snap.data(), snap.size(), LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == snap.size()); + } + return p_pop >= p_to || decode_span(p_pop, p_to, /*shadow =*/ true); +} + +// one training chunk [pos_ctx, pos_ctx + n_chunk) of a run: one llama_decode with labels. The +// batch loop is decode's, and the extras switch on in process_ubatch (the shadow's masked term, +// then llama_context::opt_ubatch). This function only fills the chunk's tokens and labels and +// says whose walk it is. bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { llama_walk_fill(batch, tokens, pos_ctx, n_chunk); for (uint32_t pos_batch = 0; pos_batch < n_chunk; ++pos_batch) { - // an exact walk's context chunk has no label but still trains (its K/V carry later - // chunks' gradient): one output row, weighted zero, so the graph has its loss - batch.logits[pos_batch] = labels[pos_ctx + pos_batch] >= 0 - || (xc != nullptr && xc->n_labels == 0 && pos_batch + 1 == n_chunk); - batch.label[pos_batch] = labels[pos_ctx + pos_batch]; + batch.logits[pos_batch] = labels[pos_ctx + pos_batch] >= 0; + batch.label [pos_batch] = labels[pos_ctx + pos_batch]; } - lctx.opt_walk_active = this; - lctx.opt_result_active = xc != nullptr && xc->n_labels == 0 ? nullptr : result; - lctx.opt_train_unlabelled = xc != nullptr && xc->n_labels == 0; + lctx.opt_walk_active = this; + lctx.opt_result_active = result; const int ret = lctx.decode(batch); - lctx.opt_walk_active = nullptr; - lctx.opt_result_active = nullptr; - lctx.opt_train_unlabelled = false; + lctx.opt_walk_active = nullptr; + lctx.opt_result_active = nullptr; std::fill(batch.label, batch.label + n_chunk, -1); if (ret != 0) { @@ -256,13 +297,19 @@ bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { } return false; } - if (callback && !(xc != nullptr && xc->n_labels == 0)) { + if (callback) { callback(train, lctx.opt_ctx, dataset, result, idata_in_loop + pos_ctx/n_batch + 1, ndata_in_loop, t_loop_start); } return true; } bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch) { + // the loss is the mean over the logits rows: a serving row among them would dilute it + if (ubatch.n_labels() != n_outputs) { + opt_fail("a training ubatch's output rows are its labelled rows (" + std::to_string(ubatch.n_labels()) + + " labelled, " + std::to_string(n_outputs) + " output): the loss is their mean"); + return false; + } opt_walk * walk = opt_walk_active; const llama_walk_chunk * xc = walk != nullptr ? walk->xc : nullptr; @@ -350,22 +397,23 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l // THE EXACT WALK. The plain walk steps per chunk and stops the gradient at every chunk boundary: a // reply learns only through its own chunk's K/V, never through what an earlier chunk computed for -// it. Here the window is decoded once under the adapter as it is, then trained in REVERSE: each -// chunk's graph attends to the cache, with its last walk_horizon prefix positions behind a -// gradient tap, and adds the surrogate + on its own K/V, where G is what every later -// chunk put there. Its backward then carries that gradient through itself into the adapter and -// into its own prefix's G. ONE optimizer step per window, on the gradient of the window's loss -// (exact within the horizon). Memory stays chunk x window; the cost is one decode plus one forward -// and backward per chunk, context chunks included. -// THE INVARIANT THAT MAKES THE REVERSE PASS VALID (Fable): the adapter does not change between the -// forward decode and the last chunk's backward. One optimizer step per window, taken by the chunk -// the reverse pass trains last, is what keeps each chunk's recomputed K/V (and recurrent state) +// it. Here each run is decoded once under the shadow, then trained in REVERSE: each chunk's graph +// attends to the cache, with the run's positions within walk_horizon behind a gradient tap, and +// adds the surrogate + on its own K/V, where G is what every later chunk of the run +// put there. Its backward then carries that gradient through itself into the adapter and into the +// run's earlier positions' G. ONE optimizer step per window, on the gradient of the window's loss +// (exact within the horizon). Memory stays chunk x window; the cost is one decode of the window, +// and per run one more decode plus one forward and backward per chunk and the restore. +// THE INVARIANT THAT MAKES THE REVERSE PASS VALID (Fable): the adapter does not change between a +// run's forward decode and its last chunk's backward. One optimizer step per window, taken by the +// chunk the walk trains last, is what keeps each chunk's recomputed K/V (and recurrent state) // equal to the cached values its successors attended to; a step inside the window would make // every surrogate a gradient of a different function. // A recurrent model's state is carried the same way, one chunk at a time: the state entering a // chunk is tapped into a host accumulator, the state it leaves meets the next chunk's gradient // (build_rs, build_walk_state_exit). The reverse pass restores each chunk's entry state from a -// checkpoint taken at its boundary in the forward pass. +// checkpoint taken at its boundary in the forward pass. The state entering a run is her stable +// set's: no gradient crosses it. void llama_context::opt_walk::run_exact() { // Each window starts from the REQUESTED horizon: a halving that fit one window's device budget // must not shrink every later window and epoch, which may be shorter (Fable). @@ -379,7 +427,7 @@ void llama_context::opt_walk::run_exact() { + std::to_string(lctx.cparams.n_rs_seq) + "): finish or drop the draft first"); return; } - plan.mark_exact(recr != nullptr, lctx.opt_walk_horizon); + plan.mark_exact(lctx.opt_walk_horizon); const std::vector & chunks = plan.chunks; const uint32_t n_pos = (uint32_t) plan.last_label + 1; @@ -409,151 +457,118 @@ void llama_context::opt_walk::run_exact() { return; } - // the window as the adapter now reads it, chunk by chunk: a recurrent model's state is - // checkpointed at each chunk's start through the sequence-state API (the recurrent part only, - // PARTIAL_ONLY), after waiting for the queued decode; a restore puts back the cell's position - // with its rows, and an empty checkpoint (before the first chunk) restores the zero state. A - // hand-rolled per-layer copy here skipped that wait and, on Metal, carried a torn mid-graph - // state into the reverse pass (2026-10-11). - std::vector> snaps(recr != nullptr ? chunks.size() : 0); - const llama_state_seq_flags state_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; - for (size_t j = 0; j < chunks.size(); ++j) { - if (recr != nullptr && j % stride == 0) { - lctx.synchronize(); - snaps[j].resize(lctx.state_seq_get_size(0, state_flags)); - GGML_ASSERT(lctx.state_seq_get_data(0, snaps[j].data(), snaps[j].size(), state_flags) == snaps[j].size()); - } - if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n)) { + uint32_t p = 0; // the cache holds [0, p) as served + for (const llama_walk_run & run : plan.runs) { + const uint32_t r0 = plan.begin(run); + const uint32_t r1 = plan.end(run); + if (!decode_span(p, r0, /*shadow =*/ false)) { return; } - } - for (int64_t j = (int64_t) chunks.size() - 1; j >= plan.first; --j) { - const llama_walk_chunk & c = chunks[j]; - if (!c.needed) { - continue; - } - if (!yield_point()) { - break; - } - // attention: pop the cache to [0, c0); recurrent: restore the state at the chunk's - // checkpoint and decode forward to its start (that decode rebuilds the attention K/V in - // between too, under the same adapter: the same values). Run again before every retry - // below: a refused chunk graph is refused AFTER its ubatch was applied, so the chunk's - // cells already sit in the cache and the recurrent cell's position has moved to the - // chunk's end; retrying on top of that fails to prepare the ubatch at all (the 5090, - // 2026-10-10: "init_batch: failed to prepare attention ubatches" after "retrying at 512", - // and the epoch stopped instead of shrinking the horizon). Always a plain decode: the - // walk's graph flags are off whenever it runs. - const auto rewind = [&]() -> bool { - const size_t cp = recr != nullptr ? (size_t) j - (size_t) j % stride : (size_t) j; - const uint32_t p_pop = chunks[cp].c0; - const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : lctx.memory->seq_rm(0, p_pop, -1); - if (!popped) { - LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u) for the reverse pass\n", __func__, p_pop); - lctx.opt_stop_requested.store(true); - return false; + // the run under the shadow, checkpointed every `stride` chunks from its start + std::vector> snaps(run.j1 - run.j0); + for (size_t j = run.j0; j < run.j1; ++j) { + if ((j - run.j0) % stride == 0) { + snaps[j - run.j0] = checkpoint(); } - if (recr != nullptr) { - lctx.synchronize(); - GGML_ASSERT(lctx.state_seq_set_data(0, snaps[cp].data(), snaps[cp].size(), state_flags) == snaps[cp].size()); - if (p_pop < c.c0 && !decode_span(p_pop, c.c0)) { - return false; - } + if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n, /*shadow =*/ true)) { + return; } - return true; - }; - if (!rewind()) { - break; } - const bool state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && acc.state_ready; - // THE HORIZON FITS THE DEVICE. Its gradient taps are the chunk graph's largest growth with - // the window (layers x horizon x (k + v) floats, twice with their gradient); the graph - // preflight refuses a chunk that does not fit before allocating anything, so a memory - // refusal halves the horizon (whole chunks) and the chunk is tried again; every later - // chunk keeps the smaller one. The forward still attends to the whole window and the - // recurrent state is carried exactly; only attention gradient past the horizon is dropped. - // A refusal at a horizon of one chunk, or for a node the device cannot run, stands. - llama_walk_chunk cc = c; - bool ok; - for (;;) { - cc.grad_from = llama_walk_plan::grad_from(cc.c0, lctx.opt_walk_horizon); - set_flags(true, cc.grad_from, c.surrogate, state_surrogate); - xc = &cc; - ok = train_chunk(cc.c0, cc.n); - xc = nullptr; - set_flags(false, 0, false, false); - const uint32_t span = cc.c0 - cc.grad_from; - if (ok || !lctx.opt_failure.empty() || !lctx.opt_alloc_failed.load() || span <= n_batch) { - break; + // the reverse pass: back to the chunk's checkpoint, forward to its start, train it + acc.state_ready = false; + for (size_t j = run.j1; j-- > run.j0; ) { + const llama_walk_chunk & c = chunks[j]; + if (!yield_point()) { + return; } - lctx.opt_walk_horizon = std::max(n_batch, (span/2)/n_batch*n_batch); - LLAMA_LOG_WARN("%s: the chunk at %u did not fit with a gradient horizon of %u positions: retrying at %u\n", - __func__, cc.c0, span, lctx.opt_walk_horizon); - lctx.opt_alloc_failed.store(false); - lctx.opt_stop_requested.store(false); - if (!rewind()) { - ok = false; - break; + const size_t cp = recr != nullptr ? (j - run.j0) - (j - run.j0) % stride : j - run.j0; + const auto back = [&]() { return rewind(chunks[run.j0 + cp].c0, snaps[cp], c.c0); }; + if (!back()) { + return; + } + const bool state_surrogate = recr != nullptr && j + 1 < run.j1 && acc.state_ready; + // THE HORIZON FITS THE DEVICE. Its gradient taps are the chunk graph's largest growth + // with the run (layers x horizon x (k + v) floats, twice with their gradient); the graph + // preflight refuses a chunk that does not fit before allocating anything, so a memory + // refusal halves the horizon (whole chunks) and the chunk is tried again; every later + // chunk keeps the smaller one. The forward still attends to the whole window and the + // recurrent state is carried exactly; only attention gradient past the horizon is + // dropped. A refusal at a horizon of one chunk, or for a node the device cannot run, stands. + llama_walk_chunk cc = c; + bool ok; + for (;;) { + cc.grad_from = std::max(r0, llama_walk_plan::grad_from(cc.c0, lctx.opt_walk_horizon)); + set_flags(true, cc.grad_from, c.surrogate, state_surrogate); + xc = &cc; + ok = train_chunk(cc.c0, cc.n); + xc = nullptr; + set_flags(false, 0, false, false); + const uint32_t span = cc.c0 - cc.grad_from; + if (ok || !lctx.opt_failure.empty() || !lctx.opt_alloc_failed.load() || span <= n_batch) { + break; + } + lctx.opt_walk_horizon = std::max(n_batch, (span/2)/n_batch*n_batch); + LLAMA_LOG_WARN("%s: the chunk at %u did not fit with a gradient horizon of %u positions: retrying at %u\n", + __func__, cc.c0, span, lctx.opt_walk_horizon); + lctx.opt_alloc_failed.store(false); + lctx.opt_stop_requested.store(false); + if (!back()) { + ok = false; + break; + } + } + lctx.opt_walk_horizon_used = lctx.opt_walk_horizon; + if (!ok) { + return; } } - lctx.opt_walk_horizon_used = lctx.opt_walk_horizon; - if (!ok) { - break; + acc.state_ready = false; + // the run restored to her stable self: its cells and state as served + if (!restore(run, snaps[0])) { + return; } + p = r1; } - acc.state_ready = false; } -// THE PLAIN WALK: each labelled chunk steps on its own gradient, then joins the context +// a run's cells and recurrent state decoded again under her stable set, from the checkpoint at +// its start (her stable state, since everything before the run is) +bool llama_context::opt_walk::restore(const llama_walk_run & run, const std::vector & snap) { + const uint32_t r0 = plan.begin(run); + const uint32_t r1 = plan.end(run); + return rewind(r0, snap, r0) && decode_span(r0, r1, /*shadow =*/ false); +} + +// THE PLAIN WALK: each labelled chunk steps on its own gradient, then joins the cache under the +// shadow as it now is, so the run's next chunk reads it; the run is restored when it is done void llama_context::opt_walk::run_plain() { - const int64_t last_label = plan.last_label; - for (const llama_walk_chunk & c : plan.chunks) { - if (lctx.opt_stop_requested.load(std::memory_order_relaxed)) { + uint32_t p = 0; + for (const llama_walk_run & run : plan.runs) { + if (!decode_span(p, plan.begin(run), /*shadow =*/ false)) { return; } - const uint32_t c0 = c.c0; - const uint32_t c1 = c.c0 + c.n; - if (c.n_labels == 0) { - if (!decode_span(c0, c1)) { + const std::vector at_run = checkpoint(); + for (size_t j = run.j0; j < run.j1; ++j) { + const llama_walk_chunk & c = plan.chunks[j]; + if (!yield_point()) { return; } - continue; - } - // The training forward advances a recurrent state in place, and a recurrent memory - // cannot remove a partial range: snapshot the state to the scratch sequence first - // (Fable), restore it after, then the decode below advances it from where the context - // left it. The cache's own copy-on-write keeps it on the device; the state is per - // sequence and small. - const bool snapshot = recr != nullptr && c1 <= (uint32_t) last_label; - if (!yield_point()) { - return; - } - if (snapshot) { - recr->seq_cp(0, 1, -1, -1); - } - if (!train_chunk(c0, c.n)) { - return; - } - if (snapshot) { - recr->seq_rm(0, -1, -1); - recr->seq_cp(1, 0, -1, -1); - recr->seq_rm(1, -1, -1); - } - // the chunk joins the context under the adapter as it now is: the training graph never - // writes the cache, so its cells hold nothing until this decode - if (c1 <= (uint32_t) last_label) { - // drop the chunk's empty cells from the attention memory, then decode them for real - const bool removed = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, c0, -1)) : lctx.memory->seq_rm(0, c0, -1); - if (!removed) { - LLAMA_LOG_ERROR("%s: could not remove the trained chunk [%u, %u) from the cache before decoding it\n", __func__, c0, c1); - lctx.opt_stop_requested.store(true); + // the training forward advances a recurrent state in place and may write the chunk's + // cells: back to the chunk's start, then decode it under the stepped shadow + const std::vector at_chunk = j + 1 < run.j1 ? checkpoint() : std::vector(); + if (!train_chunk(c.c0, c.n)) { return; } - if (!decode_span(c0, c1)) { + if (j + 1 < run.j1 && !rewind(c.c0, at_chunk, c.c0 + c.n)) { lctx.opt_stop_requested.store(true); return; } } + if (!restore(run, at_run)) { + lctx.opt_stop_requested.store(true); + return; + } + p = plan.end(run); } } @@ -590,12 +605,6 @@ void llama_context::opt_epoch_iter( walk.run_exact(); return; } - // the plain walk snapshots a recurrent state into a scratch sequence - if (walk.recr != nullptr && cparams.n_seq_max < 2) { - LLAMA_LOG_ERROR("%s: a recurrent model trains with a scratch sequence for its state snapshot: create the training context with n_seq_max >= 2\n", __func__); - opt_stop_requested.store(true); - return; - } walk.run_plain(); } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index f8a9c5014848..0ff2ff008a10 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1335,14 +1335,12 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll auto * gf = res->get_gf(); // a ubatch with labelled rows is a training step on this context's trainable adapter - // (llama_opt_init): the same graph with the extras on (opt_ubatch). It is built fresh, and + // (llama_opt_init): the same graph with the extras on, the shadow's masked term + // (graph_params) and the loss and backward (opt_ubatch). It is built fresh, and // cparams.training is part of the reuse key, so a serving graph never reuses it. bool train = false; if (opt_ctx != nullptr) { - train = opt_train_unlabelled; - for (uint32_t i = 0; !train && ubatch.label != nullptr && i < ubatch.n_tokens; ++i) { - train = ubatch.label[i] >= 0; - } + train = !opt_forward_only && ubatch.n_labels() > 0; cparams.training = train; } @@ -2508,6 +2506,8 @@ llm_graph_params llama_context::graph_params( /*.backend_cpu =*/ backend_cpu, /*.cvec =*/ cvec.get(), /*.loras =*/ loras.get(), + /*.shadow =*/ opt_shadow, + /*.n_shadow_rows =*/ opt_shadow != nullptr ? ubatch.n_labels() : 0, /*.mctx =*/ mctx, /*.cross =*/ &cross, /*.samplers =*/ sampling.samplers, diff --git a/src/llama-context.h b/src/llama-context.h index d2758811f6cb..56dd3d34156f 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -394,12 +394,16 @@ struct llama_context { uint32_t opt_walk_horizon_req = 0; // llama_opt_params::walk_horizon, as requested uint32_t opt_walk_horizon = 0; // this window's: starts at the request, shrinks to fit the device + // the adapter being trained (llama_opt_params::adapter): applied to labelled rows only + // (llm_graph_params::shadow); null when the base itself trains + const llama_adapter_lora * opt_shadow = nullptr; + // set by a walk for the decode it drives: whose taps join the graph, where the loss lands, - // whether a ubatch without labels still trains (an exact walk's context chunk), and that the - // logits rows are the loss's alone (no host copy) - opt_walk * opt_walk_active = nullptr; - ggml_opt_result_t opt_result_active = nullptr; - bool opt_train_unlabelled = false; + // whether its labelled rows only fill the cache (the shadow on, no loss, no backward), and + // that the logits rows are the loss's alone (no host copy) + opt_walk * opt_walk_active = nullptr; + ggml_opt_result_t opt_result_active = nullptr; + bool opt_forward_only = false; ggml_threadpool_t threadpool = nullptr; ggml_threadpool_t threadpool_batch = nullptr; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 6145d5ba3643..f6a4e5819677 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -223,6 +223,19 @@ void llm_graph_input_out_ids::set_input(const llama_ubatch * ubatch) { } } +void llm_graph_input_label_rows::set_input(const llama_ubatch * ubatch) { + GGML_ASSERT(rows && ggml_backend_buffer_is_host(rows->buffer)); + int32_t * data = (int32_t *) rows->data; + int64_t n = 0; + for (uint32_t i = 0; i < ubatch->n_tokens; ++i) { + if (ubatch->label != nullptr && ubatch->label[i] >= 0) { + GGML_ASSERT(n < rows->ne[0]); + data[n++] = (int32_t) i; + } + } + GGML_ASSERT(n == rows->ne[0]); +} + bool llm_graph_input_out_ids::can_reuse(const llm_graph_params & params) { bool res = true; @@ -1501,6 +1514,8 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) : backend_cpu (params.backend_cpu), cvec (params.cvec), loras (params.loras), + shadow (params.shadow), + n_shadow_rows (params.n_shadow_rows), mctx (params.mctx), cross (params.cross), samplers (params.samplers), @@ -1537,7 +1552,7 @@ ggml_tensor * llm_graph_context::build_lora_mm( for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); - if (lw == nullptr) { + if (lw == nullptr || !lora_on(lora.first)) { continue; } @@ -1550,7 +1565,7 @@ ggml_tensor * llm_graph_context::build_lora_mm( ); ab_cur = ggml_scale(ctx0, ab_cur, scale); - res = ggml_add(ctx0, res, ab_cur); + res = lora.first == shadow ? shadow_add(res, ab_cur) : ggml_add(ctx0, res, ab_cur); } return res; @@ -1573,7 +1588,7 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( } for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); - if (lw == nullptr) { + if (lw == nullptr || !lora_on(lora.first)) { continue; } @@ -1588,7 +1603,7 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( ); ab_cur = ggml_scale(ctx0, ab_cur, scale); - res = ggml_add(ctx0, res, ab_cur); + res = lora.first == shadow ? shadow_add(res, ab_cur) : ggml_add(ctx0, res, ab_cur); } return res; @@ -1731,7 +1746,7 @@ ggml_tensor * llm_graph_context::build_ffn( return false; } for (const auto & lora : *loras) { - if (lora.first->get_weight(w) != nullptr) { + if (lora.first->get_weight(w) != nullptr && lora_on(lora.first)) { return true; } } @@ -2352,7 +2367,7 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const { // apply lora for embedding tokens if needed for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(tok_embd); - if (lw == nullptr) { + if (lw == nullptr || !lora_on(lora.first)) { continue; } @@ -2364,7 +2379,7 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const { ggml_get_rows(ctx0, lw->a, inp->tokens) ), scale); - cur = ggml_add(ctx0, cur, inpL_delta); + cur = lora.first == shadow ? shadow_add(cur, inpL_delta) : ggml_add(ctx0, cur, inpL_delta); } if (n_embd_inp != n_embd) { @@ -2429,6 +2444,28 @@ ggml_tensor * llm_graph_context::build_inp_train_rows() const { return inp_train_rows; } +bool llm_graph_context::lora_on(const llama_adapter_lora * adapter) const { + return adapter != shadow || n_shadow_rows > 0; +} + +ggml_tensor * llm_graph_context::shadow_add(ggml_tensor * base, ggml_tensor * term) const { + if (inp_label_rows == nullptr) { + auto inp = std::make_unique(); + inp->rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_shadow_rows); + ggml_set_input(inp->rows); + ggml_set_name(inp->rows, "label_rows"); + inp_label_rows = inp->rows; + res->add_input(std::move(inp)); + } + // one row per token: the tokens are base's outermost dimension + GGML_ASSERT(ggml_is_contiguous(base) && (n_tokens == 1 || base->ne[ggml_n_dims(base) - 1] == n_tokens)); + const int64_t n = ggml_nelements(base) / n_tokens; + ggml_tensor * b2 = ggml_reshape_2d(ctx0, base, n, n_tokens); + ggml_tensor * t2 = ggml_reshape_2d(ctx0, ggml_cont(ctx0, term), n, n_tokens); + ggml_tensor * sel = ggml_add(ctx0, ggml_get_rows(ctx0, b2, inp_label_rows), ggml_get_rows(ctx0, t2, inp_label_rows)); + return ggml_reshape(ctx0, ggml_set_rows(ctx0, ggml_dup(ctx0, b2), sel, inp_label_rows), base); +} + ggml_tensor * llm_graph_context::build_inp_pos() const { auto inp = std::make_unique(hparams.n_pos_per_embd()); diff --git a/src/llama-graph.h b/src/llama-graph.h index 22b821a7564a..7ffae356a7aa 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -233,6 +233,17 @@ class llm_graph_input_out_ids : public llm_graph_input_i { const uint32_t n_outputs; }; +// the ubatch's labelled rows, in order: where the shadow adapter applies (shadow_add) +class llm_graph_input_label_rows : public llm_graph_input_i { +public: + llm_graph_input_label_rows() = default; + virtual ~llm_graph_input_label_rows() = default; + + void set_input(const llama_ubatch * ubatch) override; + + ggml_tensor * rows = nullptr; // I32 [n_shadow_rows] +}; + class llm_graph_input_mean : public llm_graph_input_i { public: llm_graph_input_mean(const llama_cparams & cparams) : cparams(cparams) {} @@ -783,6 +794,11 @@ struct llm_graph_params { const llama_adapter_cvec * cvec; const llama_adapter_loras * loras; + // THE SHADOW: the attached adapter being trained (llama_opt_init), the layer that switches on + // and off. It applies only to the ubatch's labelled rows (n_shadow_rows of them); a ubatch + // with no label never applies it, so serving is the stable set's graph + const llama_adapter_lora * shadow; + uint32_t n_shadow_rows; const llama_memory_context_i * mctx; const llama_cross * cross; @@ -880,6 +896,8 @@ struct llm_graph_params { gtype == other.gtype && cvec == other.cvec && loras == other.loras && + shadow == other.shadow && + n_shadow_rows == other.n_shadow_rows && cross == other.cross; } }; @@ -1055,6 +1073,8 @@ struct llm_graph_context { const llama_adapter_cvec * cvec; const llama_adapter_loras * loras; + const llama_adapter_lora * shadow; + const uint32_t n_shadow_rows; const llama_memory_context_i * mctx; const llama_cross * cross; @@ -1065,6 +1085,7 @@ struct llm_graph_context { llm_graph_result * res; mutable ggml_tensor * inp_train_rows = nullptr; // build_inp_train_rows, once per graph + mutable ggml_tensor * inp_label_rows = nullptr; // shadow_add, once per graph ggml_context * ctx0 = nullptr; ggml_cgraph * gf = nullptr; @@ -1188,6 +1209,16 @@ struct llm_graph_context { // is its cell). One per graph, shared by every attention layer. ggml_tensor * build_inp_train_rows() const; + // whether an attached adapter's term is in this graph: every one but the shadow always, the + // shadow only when the ubatch has labelled rows + bool lora_on(const llama_adapter_lora * adapter) const; + // base plus the shadow's term, on the labelled rows only (Cormac, 2026-10-14): those rows of + // base and of the term are gathered and added, and written back into a copy of base; every + // other row is base, byte for byte. A multiplied 0/1 mask was not: an untrained layer's inf or + // NaN times zero is NaN, and x + 0 turns a -0 into +0, on a serving row co-batched with + // training. + ggml_tensor * shadow_add(ggml_tensor * base, ggml_tensor * term) const; + // Whether this layer's attention runs through FLASH_ATTN_EXT: the one decision // build_attn_mha makes and a training graph reads to attend to the cache in place. bool uses_flash_attn(ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla) const; diff --git a/src/llama-walk.cpp b/src/llama-walk.cpp index 481e8a8e9078..3ea65e6f8691 100644 --- a/src/llama-walk.cpp +++ b/src/llama-walk.cpp @@ -23,11 +23,15 @@ llama_walk_plan llama_walk_plan::of(const std::vector & labels, uin while ((int64_t) end <= plan.last_label && (labels[end] >= 0) == labelled) { ++end; } + const size_t j0 = plan.chunks.size(); for (uint32_t c0 = p; c0 < end; c0 += n_batch) { const uint32_t n = std::min(c0 + n_batch, end) - c0; plan.chunks.push_back({ c0, n, labelled ? n : 0 }); plan.n_labels += labelled ? n : 0; } + if (labelled) { + plan.runs.push_back({ j0, plan.chunks.size() }); + } p = end; } return plan; @@ -37,35 +41,44 @@ uint32_t llama_walk_plan::grad_from(uint32_t c0, uint32_t horizon) { return horizon == 0 || c0 < horizon ? 0 : c0 - horizon; } -void llama_walk_plan::mark_exact(bool recurrent, uint32_t horizon) { - uint32_t reach = UINT32_MAX; // the lowest position a later trained chunk's gradient reaches - first = -1; - for (int64_t j = (int64_t) chunks.size() - 1; j >= 0; --j) { - llama_walk_chunk & c = chunks[j]; - c.surrogate = reach < c.c0 + c.n; - c.needed = c.n_labels > 0 || c.surrogate || (recurrent && first >= 0); +void llama_walk_plan::mark_exact(uint32_t horizon) { + for (llama_walk_chunk & c : chunks) { + c.needed = false; + c.surrogate = false; c.period_end = false; - if (!c.needed) { - continue; + } + for (const llama_walk_run & r : runs) { + uint32_t reach = UINT32_MAX; // the lowest position a later chunk of the run reaches + for (size_t j = r.j1; j-- > r.j0; ) { + llama_walk_chunk & c = chunks[j]; + c.needed = true; + c.surrogate = reach < c.c0 + c.n; + c.grad_from = std::max(begin(r), grad_from(c.c0, horizon)); + c.loss_scale = (float) c.n_labels / (float) n_labels; + reach = std::min(reach, c.grad_from); } - c.grad_from = grad_from(c.c0, horizon); - c.loss_scale = (float) c.n_labels / (float) n_labels; - reach = std::min(reach, c.grad_from); - first = j; } - GGML_ASSERT(first >= 0); + GGML_ASSERT(!runs.empty()); + first = (int64_t) runs.back().j0; chunks[first].period_end = true; } size_t llama_walk_plan::host_bytes(size_t stride, size_t kv_bytes, size_t state_bytes) const { - const size_t n_checkpoints = chunks.empty() ? 0 : (chunks.size() + stride - 1)/stride - 1; + size_t n_checkpoints = 0; + for (const llama_walk_run & r : runs) { + n_checkpoints = std::max(n_checkpoints, (r.j1 - r.j0 + stride - 1)/stride); + } return kv_bytes + n_checkpoints*state_bytes; } size_t llama_walk_plan::stride_for(size_t budget, size_t kv_bytes, size_t state_bytes) const { + size_t longest = 1; + for (const llama_walk_run & r : runs) { + longest = std::max(longest, r.j1 - r.j0); + } size_t stride = 1; if (budget > 0) { - while (stride < chunks.size() && host_bytes(stride, kv_bytes, state_bytes) > budget) { + while (stride < longest && host_bytes(stride, kv_bytes, state_bytes) > budget) { ++stride; } } diff --git a/src/llama-walk.h b/src/llama-walk.h index 67a68bf12db7..057d50fd21b3 100644 --- a/src/llama-walk.h +++ b/src/llama-walk.h @@ -32,26 +32,38 @@ struct llama_walk_chunk { float loss_scale = 0.0f; // its labelled positions over the window's }; +// chunks [j0, j1) of the plan: one labelled run, which trains given everything before it as +// served (the shadow adapter applies to labelled rows only) +struct llama_walk_run { + size_t j0; + size_t j1; +}; + struct llama_walk_plan { // the window's runs, each cut every n_batch positions from its start; both walks take // them in this order. Nothing after the last label is in the plan. std::vector chunks; + std::vector runs; // the labelled ones, in order int64_t last_label = -1; // -1: the window has no label, nothing to learn from it uint32_t n_labels = 0; int64_t first = -1; // the exact walk's period end: the chunk its reverse pass trains last static llama_walk_plan of(const std::vector & labels, uint32_t n_ctx, uint32_t n_batch); - // which chunks the exact walk trains (a labelled one, one a later gradient reaches, and on a - // recurrent model every one before a trained one, since the state reaches every later - // chunk), how far back each gradient reaches, and the period end - void mark_exact(bool recurrent, uint32_t horizon); + // positions [begin, end) of a run + uint32_t begin(const llama_walk_run & r) const { return chunks[r.j0].c0; } + uint32_t end (const llama_walk_run & r) const { return chunks[r.j1 - 1].c0 + chunks[r.j1 - 1].n; } + + // which chunks the exact walk trains (the labelled ones: the context before a run is her + // stable set's and carries no gradient), how far back each gradient reaches (never before + // its run's start), and the period end (the first chunk of the last run, trained last) + void mark_exact(uint32_t horizon); // the first cached position a chunk at c0 carries gradient to (horizon 0: the whole prefix) static uint32_t grad_from(uint32_t c0, uint32_t horizon); // THE HOST BUDGET (Fable): the K/V accumulators, plus a recurrent state checkpoint every - // `stride` chunks (the first chunk's entry is the zero state and costs nothing) + // `stride` chunks of the run being walked (the largest run's) size_t host_bytes(size_t stride, size_t kv_bytes, size_t state_bytes) const; // the smallest stride under the budget (0: no cap); the reverse pass rebuilds a chunk's // entry state by decoding forward from its checkpoint, so a longer stride trades host diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index b53590be849e..06389e69a212 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2390,9 +2390,12 @@ struct test_set_rows : public test_case { const std::array nr23; // broadcast only dims 2 and 3 const int r; // rows to set const bool v; // view (non-contiguous src1) + // gradients for both the rows written and the destination, written into a copy of it (the + // shadow adapter's row selection, llm_graph_context::shadow_add) + const bool grad; std::string vars() override { - return VARS_TO_STR7(type_src, type_dst, type_idx, ne, nr23, r, v); + return VARS_TO_STR8(type_src, type_dst, type_idx, ne, nr23, r, v, grad); } test_set_rows(ggml_type type_src, @@ -2400,8 +2403,8 @@ struct test_set_rows : public test_case { ggml_type type_idx, std::array ne, std::array nr23, - int r, bool v = false) - : type_src(type_src), type_dst(type_dst), type_idx(type_idx), ne(ne), nr23(nr23), r(r), v(v) {} + int r, bool v = false, bool grad = false) + : type_src(type_src), type_dst(type_dst), type_idx(type_idx), ne(ne), nr23(nr23), r(r), v(v), grad(grad) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * dst = ggml_new_tensor_4d(ctx, type_dst, ne[0], ne[1], ne[2]*nr23[0], ne[3]*nr23[1]); @@ -2410,6 +2413,13 @@ struct test_set_rows : public test_case { ggml_tensor * src = ggml_new_tensor_4d(ctx, type_src, ne[0], r, ne[2]*nr23[0], ne[3]*nr23[1]); ggml_set_name(src, "src"); + if (grad) { + ggml_set_param(dst); + ggml_set_param(src); + dst = ggml_dup(ctx, dst); + ggml_set_name(dst, "copy_of_dst"); + } + ggml_tensor * row_idxs = ggml_new_tensor_3d(ctx, type_idx, r, ne[2], ne[3]); ggml_set_name(row_idxs, "row_idxs"); @@ -8734,6 +8744,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I32, { 64, 9, 1, 1 }, { 1, 1 }, 4, false, true)); test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_Q8_0, GGML_TYPE_I32, { 256, 5, 1, 3 }, { 1, 1, }, 1, false)); for (ggml_type src_type : {GGML_TYPE_F16, GGML_TYPE_F32}) { for (ggml_type type : all_types) { diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index f0babeba974b..04cf6aec61d8 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -95,8 +95,7 @@ static llama_context * make_ctx(const common_params & params, llama_model * mode cparams.n_ctx = g_window; cparams.n_batch = chunk; cparams.n_ubatch = chunk; - cparams.n_seq_max = 2; // the plain walk snapshots a recurrent state into a scratch sequence - cparams.kv_unified = true; // ...and the window keeps every cell + cparams.kv_unified = true; // the window keeps every cell // a quantized V cache needs flash attention (it stores V untransposed); training graphs take // the explicit path either way cparams.flash_attn_type = g_cache_type == GGML_TYPE_F32 ? LLAMA_FLASH_ATTN_TYPE_DISABLED : LLAMA_FLASH_ATTN_TYPE_ENABLED; @@ -142,7 +141,8 @@ static std::vector> g_layout; // one SGD step from the fresh adapter: the step itself, -lr * the gradient it took static std::vector step_delta(const common_params & params, llama_model * model, const std::string & init, const std::vector & tokens, const std::vector & labelled, - uint32_t chunk, bool exact, size_t host_budget = 0, size_t * host_bytes = nullptr) { + uint32_t chunk, bool exact, size_t host_budget = 0, size_t * host_bytes = nullptr, + std::vector * state_after = nullptr) { llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); GGML_ASSERT(adapter != nullptr); const std::vector before = adapter_params(adapter); @@ -179,6 +179,10 @@ static std::vector step_delta(const common_params & params, llama_model * } g_last_graph_bytes = llama_opt_graph_bytes(ctx); g_last_horizon = llama_opt_walk_horizon(ctx); + if (state_after) { + state_after->resize(llama_state_seq_get_size(ctx, 0)); + GGML_ASSERT(llama_state_seq_get_data(ctx, state_after->data(), state_after->size(), 0) == state_after->size()); + } ggml_opt_result_free(result); ggml_opt_dataset_free(dataset); llama_free(ctx); @@ -192,6 +196,30 @@ static std::vector step_delta(const common_params & params, llama_model * return d; } +// her stable self's cache after [0, r0) and then [r0, r1): no adapter, decoded in chunks from +// each start as the walk decodes them (a matmul's result depends on its batch size), so a walk +// that restores its run leaves exactly these bytes +static std::vector stable_state(const common_params & params, llama_model * model, + const std::vector & tokens, uint32_t r0, uint32_t r1, uint32_t chunk) { + llama_context * ctx = make_ctx(params, model, chunk); + llama_batch batch = llama_batch_init(chunk, 0, 1); + for (const auto & [p0, p1] : { std::make_pair(0u, r0), std::make_pair(r0, r1) }) { + for (uint32_t c0 = p0; c0 < p1; c0 += chunk) { + const uint32_t n = std::min(c0 + chunk, p1) - c0; + common_batch_clear(batch); + for (uint32_t i = 0; i < n; ++i) { + common_batch_add(batch, tokens[c0 + i], c0 + i, { 0 }, i + 1 == n); + } + GGML_ASSERT(llama_decode(ctx, batch) == 0); + } + } + std::vector state(llama_state_seq_get_size(ctx, 0)); + GGML_ASSERT(llama_state_seq_get_data(ctx, state.data(), state.size(), 0) == state.size()); + llama_batch_free(batch); + llama_free(ctx); + return state; +} + struct agreement { double cosine; double norm_ratio; @@ -368,8 +396,9 @@ int main(int argc, char ** argv) { } } { - // 2. context, then a reply (the walk's real case): the exact walk is invariant to how the - // window is chunked, context chunks included (they train through the surrogate alone) + // 2. context, then a reply (the walk's real case): the reply trains given the context as + // served, so only its own chunks train, and the exact walk is invariant to how the reply + // is chunked std::vector reply(WINDOW + 1, 0); for (uint32_t i = WINDOW / 2 + 37; i <= WINDOW; ++i) { reply[i] = 1; @@ -387,6 +416,30 @@ int main(int argc, char ** argv) { } } + { + // 8. what this catches: a walk that leaves the shadow in her cache. The reply's cells + // are written under the shadow while it trains, then restored, so after either walk the + // sequence (every attention cell, the recurrent state) is byte-identical to her stable + // self's decode of the same window: serving never reads what training wrote. + std::vector reply(WINDOW + 1, 0); + for (uint32_t i = WINDOW / 2 + 37; i <= WINDOW; ++i) { + reply[i] = 1; + } + const uint32_t r0 = WINDOW / 2 + 36; // the first labelled position (it predicts reply[r0 + 1]) + const std::vector want = stable_state(params, model, tokens, r0, WINDOW, WINDOW / 4); + for (bool exact : { true, false }) { + std::vector got; + step_delta(params, model, init, tokens, reply, WINDOW / 4, exact, 0, nullptr, &got); + const bool same = got == want; + printf(" cache after the %s walk: %zu bytes, %s her stable decode\n", exact ? "exact" : "plain", got.size(), + same ? "byte-identical to" : "DIFFERS from"); + if (!same) { + fprintf(stderr, "FAILED: the %s walk left the shadow in the cache\n", exact ? "exact" : "plain"); + ++failures; + } + } + } + { // 7. per-layer recompute under the exact walk. The chunk's activations, every layer's // kept for the backward pass, are what a large model's chunk graph is made of (Kimi's diff --git a/tests/test-walk-plan.cpp b/tests/test-walk-plan.cpp index ed3749131194..04d145cc6b4f 100644 --- a/tests/test-walk-plan.cpp +++ b/tests/test-walk-plan.cpp @@ -38,22 +38,30 @@ int main(int argc, char ** argv) { }); // what this catches: a step taken anywhere but the reverse pass's last chunk, a loss that is - // not the window's mean, a gradient reaching past the horizon + // not the window's mean, a gradient reaching past the horizon or before its own run (the + // context before a run is her stable set's: the shadow never touched it), a context chunk + // trained for nothing t.test("exact_marks", [](testing & t) { - auto plan = llama_walk_plan::of(window({ {1024, false}, {512, true}, {512, false}, {512, true} }), 2560, 512); - plan.mark_exact(/*recurrent =*/ false, /*horizon =*/ 512); + auto plan = llama_walk_plan::of(window({ {1024, false}, {1536, true}, {512, false}, {512, true} }), 3584, 512); + t.assert_equal(2u, plan.runs.size()); + plan.mark_exact(/*horizon =*/ 512); float loss = 0.0f; int ends = 0; for (const auto & c : plan.chunks) { - t.assert_true("every chunk before a reply carries its gradient", c.needed); - t.assert_equal(llama_walk_plan::grad_from(c.c0, 512), c.grad_from); - loss += c.loss_scale; + t.assert_true("exactly the labelled chunks train", c.needed == (c.n_labels > 0)); + loss += c.needed ? c.loss_scale : 0.0f; ends += c.period_end; } - t.assert_equal(0, plan.first); + // run 1 is chunks 2..4 (positions 1024..2560), run 2 is chunk 6 (3072..3584) + t.assert_equal(1024u, plan.chunks[2].grad_from); + t.assert_equal(1024u, plan.chunks[3].grad_from); + t.assert_equal(1536u, plan.chunks[4].grad_from); + t.assert_equal(3072u, plan.chunks[6].grad_from); + t.assert_true("a chunk a later one of its run reaches", plan.chunks[2].surrogate && plan.chunks[3].surrogate); + t.assert_true("the last chunk of a run has no later gradient", !plan.chunks[4].surrogate && !plan.chunks[6].surrogate); + t.assert_equal(6, plan.first); t.assert_equal(1, ends); t.assert_true("loss scales sum to one", loss > 0.9999f && loss < 1.0001f); - t.assert_true("the last chunk has no later gradient", !plan.chunks.back().surrogate); t.assert_equal(0u, llama_walk_plan::grad_from(1024, 0)); t.assert_equal(512u, llama_walk_plan::grad_from(1024, 512)); }); @@ -61,13 +69,14 @@ int main(int argc, char ** argv) { // what this catches: a stride that does not shrink the checkpoints to the budget, or one // that grows past the window when nothing fits t.test("host_budget_stride", [](testing & t) { + // a run's start is a checkpoint too: the walk restores the run to her stable self from it const auto plan = llama_walk_plan::of(window({ {8*512, true} }), 8*512, 512); - t.assert_equal(170u, plan.host_bytes(1, 100, 10)); + t.assert_equal(180u, plan.host_bytes(1, 100, 10)); t.assert_equal(1u, plan.stride_for(0, 100, 10)); t.assert_equal(2u, plan.stride_for(150, 100, 10)); - t.assert_equal(130u, plan.host_bytes(2, 100, 10)); + t.assert_equal(140u, plan.host_bytes(2, 100, 10)); t.assert_equal(8u, plan.stride_for(50, 100, 10)); - t.assert_equal(100u, plan.host_bytes(8, 100, 10)); + t.assert_equal(110u, plan.host_bytes(8, 100, 10)); }); return t.summary(); diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index 8dd7f4d87426..46e82ede20f1 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -907,10 +907,8 @@ void server_trainer::run(json req, examples_data ex) { cparams.n_ctx = window; cparams.n_batch = chunk; cparams.n_ubatch = chunk; - // two sequences: 0 is the walk; 1 is the scratch a recurrent state is snapshotted into - // around each training chunk (llama_context::opt_epoch_iter). Unified, so the attention - // cache stays one stream of the window, not one per sequence. - cparams.n_seq_max = 2; + // one sequence, the walk (it checkpoints a recurrent state through the sequence-state API) + cparams.n_seq_max = 1; cparams.kv_unified = true; // Flash attention ON for the context: the walk's context decodes use it, and the cache then // stores V un-transposed, which a quantized V cache requires. The training graphs build From ee9b17965c661f598e510e7d45f52d0f09048400 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 08:02:05 -0500 Subject: [PATCH 08/43] llama: the exact walk's K/V gradient holds the reply being trained, not the window P1 step 3b (continuum #5032 section 9.6). Under the served-prefix objective only a run's own positions carry K/V gradient, so the host accumulator is sized to the longest run and rebased to each run's start (llama_walk_acc::start_run, cparams.walk_acc_base). Before, it was sized to the window: at her served window on a 27B that alone was gigabytes of host memory. - The in-place (flash) path taps the window's cells but keeps only the run's rows before the chunk: ggml_opt_accumulate_grad_rows adds rows [row0, row0 + n) of a tensor's gradient into a caller's span; the rest stays on the device for its consumers. What later chunks put on a chunk's own rows reaches it through the surrogate, now one helper for both attention paths (build_walk_surrogate). - The host budget's refusal names the reply, which is what the K/V gradient grows with. Measured (this commit, M5 Metal, test-walk-exact, 2026-10-14): qwen2.5-0.5b 12:56:43-12:57:41Z and qwen3.5-4b hybrid 12:57:41-13:01:53Z, both OK; every exact row 1.0000/1.0000; the cache after either walk byte-identical to her stable decode; device refusal 3 retries, 4 refusals, 0 ubatch failures. Owed: the 27B at her served window with a 2k reply, accumulator bytes and memory peak measured. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/include/ggml-opt.h | 5 +++ ggml/src/ggml-opt.cpp | 22 +++++++++++++ src/llama-context-opt.cpp | 18 +++++++---- src/llama-cparams.h | 8 +++-- src/llama-graph.cpp | 66 ++++++++++++++++++++++----------------- src/llama-graph.h | 6 ++++ src/llama-walk.cpp | 17 ++++++++++ src/llama-walk.h | 16 ++++++---- 8 files changed, 115 insertions(+), 43 deletions(-) diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index 148349e99e20..1364ce4fd96f 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h @@ -238,6 +238,11 @@ extern "C" { // zeroes it, except at its own first allocation, which precedes any accumulation. GGML_API void ggml_opt_accumulate_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, struct ggml_tensor * acc); + // The same for rows [row0, row0 + acc->ne[1]) of a 2-D `tensor`'s gradient only: the rest of + // its gradient lives on the device for its consumers and is never kept, so `acc` holds a + // span (a reply's positions) of a window-sized tensor (a cache's cells). + GGML_API void ggml_opt_accumulate_grad_rows(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, int64_t row0, struct ggml_tensor * acc); + // ############################################################################ // ## The high-level functions start here. They do not depend on any private ## // ## functions or structs and can be copied to and adapted for user code. ## diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 87cd73369d30..6bf7d7c23b14 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -89,6 +89,9 @@ struct ggml_opt_context { // ggml_opt_accumulate_grad: GRAD nodes whose gradient adds into a caller's tensor, for the // graph prepared last (ggml_opt_prepare_alloc clears them) std::unordered_map next_grad_accs; + // ggml_opt_accumulate_grad_rows: GRAD nodes whose gradient's rows [row0, row0 + acc rows) add + // into a caller's tensor, for the graph prepared last + std::unordered_map> next_grad_rows; std::vector grad_accs; int64_t iter = 1; @@ -703,6 +706,15 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { opt_ctx->leaf_grads.clear(); for (int i = 0; i < opt_ctx->gf->n_nodes; ++i) { ggml_tensor * node = opt_ctx->gf->nodes[i]; + if (const auto it = opt_ctx->next_grad_rows.find(node); it != opt_ctx->next_grad_rows.end()) { + ggml_tensor * grad = ggml_graph_get_grad(opt_ctx->gb_grad, node); + if (grad) { + const auto & [row0, acc] = it->second; + ggml_tensor * rows = ggml_view_2d(opt_ctx->ctx_compute, grad, acc->ne[0], acc->ne[1], grad->nb[1], row0*grad->nb[1]); + ggml_build_forward_expand(opt_ctx->gb_grad, ggml_add_inplace(opt_ctx->ctx_compute, acc, rows)); + } + continue; + } if ((node->flags & GGML_TENSOR_FLAG_GRAD) && !opt_ctx->next_grad_accs.count(node)) { ggml_tensor * grad = ggml_graph_get_grad(opt_ctx->gb_grad, node); if (grad) { @@ -836,6 +848,15 @@ void ggml_opt_accumulate_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * t opt_ctx->next_grad_accs[tensor] = acc; } +void ggml_opt_accumulate_grad_rows(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, int64_t row0, struct ggml_tensor * acc) { + GGML_ASSERT(!opt_ctx->eval_ready && "register an accumulator before ggml_opt_alloc"); + GGML_ASSERT(acc->type == GGML_TYPE_F32 && ggml_n_dims(tensor) <= 2 && acc->ne[0] == tensor->ne[0]); + GGML_ASSERT(row0 >= 0 && row0 + acc->ne[1] <= tensor->ne[1]); + GGML_ASSERT((acc->view_src ? acc->view_src : acc)->buffer != nullptr && "the accumulator lives in the caller's buffer"); + ggml_set_grad(tensor); + opt_ctx->next_grad_rows[tensor] = { row0, acc }; +} + struct ggml_tensor * ggml_opt_leaf_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * leaf) { const auto it = opt_ctx->leaf_grads.find(leaf); return it == opt_ctx->leaf_grads.end() ? nullptr : it->second; @@ -976,6 +997,7 @@ void ggml_opt_prepare_alloc( opt_ctx->inputs = inputs; opt_ctx->outputs = outputs; opt_ctx->next_grad_accs.clear(); // registrations belong to the graph prepared last + opt_ctx->next_grad_rows.clear(); } void ggml_opt_set_alloc_budget(ggml_opt_context_t opt_ctx, size_t bytes) { diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index b0fc285625a1..5813e004b71f 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -329,8 +329,14 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); for (const auto & io : res->t_walk) { if (io.tk) { - ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); - ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); + // the explicit path taps the span itself; the in-place path taps the window's cells + if (ggml_are_same_shape(io.tk, io.ak)) { + ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); + ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); + } else { + ggml_opt_accumulate_grad_rows(opt_ctx, io.tk, io.g0, io.ak); + ggml_opt_accumulate_grad_rows(opt_ctx, io.tv, io.g0, io.av); + } } } for (const auto & ws : res->t_walk_state) { @@ -429,7 +435,7 @@ void llama_context::opt_walk::run_exact() { } plan.mark_exact(lctx.opt_walk_horizon); const std::vector & chunks = plan.chunks; - const uint32_t n_pos = (uint32_t) plan.last_label + 1; + const uint32_t n_pos = plan.longest_run(); // THE HOST BUDGET (llama_walk_plan::stride_for): the K/V accumulators and the recurrent state // checkpoints are both known before anything runs @@ -442,8 +448,8 @@ void llama_context::opt_walk::run_exact() { if (budget > 0 && bytes > budget) { char why[512]; snprintf(why, sizeof(why), - "the exact walk needs %.1f MiB of host memory for this window (%u positions, the K/V gradient alone), over its budget of %.1f MiB: " - "the K/V gradient grows with the window; neither the horizon nor the chunk size changes it", + "the exact walk needs %.1f MiB of host memory for this window (its longest reply, %u positions, the K/V gradient alone), over its budget of %.1f MiB: " + "the K/V gradient grows with the reply; neither the horizon nor the chunk size changes it", bytes/1048576.0, n_pos, budget/1048576.0); lctx.opt_fail(why); return; @@ -475,7 +481,7 @@ void llama_context::opt_walk::run_exact() { } } // the reverse pass: back to the chunk's checkpoint, forward to its start, train it - acc.state_ready = false; + acc.start_run(r0); for (size_t j = run.j1; j-- > run.j0; ) { const llama_walk_chunk & c = chunks[j]; if (!yield_point()) { diff --git a/src/llama-cparams.h b/src/llama-cparams.h index f32f7f465b6a..a5f4c57cd740 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -18,14 +18,16 @@ struct llama_cparams { // gradient (positions [walk_grad_from, n_past) attend through a gradient tap; [0, // walk_grad_from) stay constants). walk_surrogate: later chunks accumulated a gradient on // this chunk's own K/V, so its graph adds the surrogate term that carries it in. - // walk_acc_k / walk_acc_v: per layer, dL/d(cached K / V) of the window, [n_embd_*_gqa, - // window] F32 in host memory (null for a recurrent layer); a chunk's backward adds its prefix's - // share into rows [walk_grad_from, n_past) and its surrogate reads its own rows. + // walk_acc_k / walk_acc_v: per layer, dL/d(cached K / V) of the run being trained, [n_embd_*_gqa, + // the run's positions] F32 in host memory (null for a recurrent layer), row 0 at position + // walk_acc_base; a chunk's backward adds its prefix's share into rows [walk_grad_from, n_past) + // and its surrogate reads its own rows. bool walk_exact = false; uint32_t walk_grad_from = 0; bool walk_surrogate = false; ggml_tensor * const * walk_acc_k = nullptr; ggml_tensor * const * walk_acc_v = nullptr; + uint32_t walk_acc_base = 0; // a later chunk's gradient sits on the recurrent state this chunk leaves bool walk_state_surrogate = false; uint32_t n_ctx_seq; // context for a single sequence diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index f6a4e5819677..38f52b951e78 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2908,19 +2908,23 @@ ggml_tensor * llm_graph_context::build_attn( if (cparams.training && uses_flash_attn(kq_b, sinks, v_mla)) { // THE TRAINING GRAPH IS THE SERVING GRAPH. This chunk writes its K/V into the cache // exactly as serving does (SET_ROWS, the cache's own type) and attention reads the - // cache in place: no copy of the context, at any window. The walk decoded the window + // cache in place: no copy of the context, at any window. The walk decoded the run // into the cache under the same adapter, so the write is idempotent; the backward // gathers this chunk's gradient from the rows it wrote (SET_ROWS backward). In the - // exact walk the write is TAPPED (cparams.walk_acc_k / _v): the backward adds - // dL/d(cached K/V) of the whole window into the host accumulator, and reads its own - // rows back from it, so what later chunks put on them reaches the adapter through - // this chunk with no surrogate term. A quantized or f16 cache trains on the values - // serving reads; its gradient passes straight through the cast. + // exact walk the write is TAPPED: the backward adds dL/d(cached K/V) of the run's + // positions before this chunk ([walk_grad_from, n_past)) into the run's host + // accumulator (cparams.walk_acc_k / _v, row 0 at walk_acc_base), and what later chunks + // put on this chunk's own rows reaches it through the surrogate. A quantized or f16 + // cache trains on the values serving reads; its gradient passes straight through the + // cast. const int64_t n_tokens = k_cur->ne[2]; const int64_t n_past = ubatch.pos[0]; const int64_t n_win = n_past + n_tokens; ggml_tensor * rows = build_inp_train_rows(); - llm_graph_result::walk_layer io = { il, (uint32_t) n_past }; + // the run's positions before this chunk carry a gradient: [g0, n_past), rows from base + const int64_t base = cparams.walk_acc_base; + const int64_t g0 = std::max(cparams.walk_grad_from, base); + llm_graph_result::walk_layer io = { il, (uint32_t) n_past, (uint32_t) g0 }; // the window's cells [0, n_win) of one layer's cache with this chunk's rows written auto write = [&](ggml_tensor * all, ggml_tensor * cur, ggml_tensor * acc, ggml_tensor *& tap, ggml_tensor *& tap_acc) { const int64_t n_embd = cur->ne[0]*cur->ne[1]; @@ -2928,9 +2932,9 @@ ggml_tensor * llm_graph_context::build_attn( GGML_ASSERT(ggml_row_size(cur->type, cur->ne[0]) == cur->nb[1]); ggml_tensor * win = ggml_view_2d(ctx0, all, n_embd, n_win, all->nb[2], 0); ggml_tensor * w = ggml_set_rows(ctx0, win, ggml_view_2d(ctx0, cur, n_embd, n_tokens, cur->nb[2], 0), rows); - if (acc != nullptr) { + if (acc != nullptr && n_past > g0) { tap = w; - tap_acc = ggml_view_2d(ctx0, acc, n_embd, n_win, acc->nb[1], 0); + tap_acc = ggml_view_2d(ctx0, acc, n_embd, n_past - g0, acc->nb[1], (g0 - base)*acc->nb[1]); } return ggml_reshape_3d(ctx0, w, cur->ne[0], cur->ne[1], n_win); }; @@ -2940,6 +2944,9 @@ ggml_tensor * llm_graph_context::build_attn( k = write(mctx_cur->get_k(ctx0, il), k_cur, tapped ? cparams.walk_acc_k[il] : nullptr, io.tk, io.ak); v = write(v_all, v_cur, tapped ? cparams.walk_acc_v[il] : nullptr, io.tv, io.av); if (tapped) { + if (cparams.walk_surrogate) { + build_walk_surrogate(il, k_cur, v_cur, n_past - base); + } res->t_walk.push_back(io); } if (kq_mask->ne[0] > n_win) { @@ -2994,15 +3001,9 @@ ggml_tensor * llm_graph_context::build_attn( // (test-walk-exact, q8_0 and f16 caches): cosine 0.998 against one graph, the same as // at an F32 cache, so training reads the numbers serving runs on (Joel: align the bit // depth to inference). - const int64_t g0 = std::min(cparams.walk_grad_from, n_past); - llm_graph_result::walk_layer io = { il, (uint32_t) n_past }; - ggml_tensor * acc_k = cparams.walk_acc_k[il]; - ggml_tensor * acc_v = cparams.walk_acc_v[il]; - // rows [p0, p0 + n) of a layer's accumulator, shaped like that span's K / V - auto acc_rows = [&](ggml_tensor * acc, const ggml_tensor * like, int64_t p0) { - return ggml_view_3d(ctx0, acc, like->ne[0], like->ne[1], like->ne[2], - like->ne[0]*sizeof(float), acc->nb[1], p0*acc->nb[1]); - }; + const int64_t base = cparams.walk_acc_base; + const int64_t g0 = std::min(std::max(cparams.walk_grad_from, base), n_past); + llm_graph_result::walk_layer io = { il, (uint32_t) n_past, (uint32_t) g0 }; // [const prefix][gradient prefix][this chunk]: the order the cells hold them in std::vector k_parts, v_parts; if (g0 > 0) { @@ -3015,8 +3016,8 @@ ggml_tensor * llm_graph_context::build_attn( io.tv = v_span(g0, n_past); ggml_format_name(io.tk, "walk_tk-%d", il); ggml_format_name(io.tv, "walk_tv-%d", il); - io.ak = acc_rows(acc_k, io.tk, g0); - io.av = acc_rows(acc_v, io.tv, g0); + io.ak = walk_acc_rows(cparams.walk_acc_k[il], io.tk, g0 - base); + io.av = walk_acc_rows(cparams.walk_acc_v[il], io.tv, g0 - base); k_parts.push_back(io.tk); v_parts.push_back(io.tv); } @@ -3029,14 +3030,7 @@ ggml_tensor * llm_graph_context::build_attn( v = ggml_concat(ctx0, v, v_parts[i], 2); } if (cparams.walk_surrogate) { - ggml_tensor * gk = acc_rows(acc_k, k_cur, n_past); - ggml_tensor * gv = acc_rows(acc_v, v_cur, n_past); - ggml_tensor * k32 = k_cur->type == GGML_TYPE_F32 ? k_cur : ggml_cast(ctx0, k_cur, GGML_TYPE_F32); - ggml_tensor * v32 = v_cur->type == GGML_TYPE_F32 ? v_cur : ggml_cast(ctx0, v_cur, GGML_TYPE_F32); - ggml_tensor * term = ggml_add(ctx0, - ggml_sum(ctx0, ggml_mul(ctx0, k32, gk)), - ggml_sum(ctx0, ggml_mul(ctx0, v32, gv))); - res->t_walk_surrogate = res->t_walk_surrogate ? ggml_add(ctx0, res->t_walk_surrogate, term) : term; + build_walk_surrogate(il, k_cur, v_cur, n_past - base); } res->t_walk.push_back(io); } else if (n_past > 0) { @@ -3732,6 +3726,22 @@ int64_t llm_graph_context::rs_keep() const { return cparams.n_rs_seq > 0 && !training ? (int64_t) cparams.n_rs_seq + 1 : 1; } +ggml_tensor * llm_graph_context::walk_acc_rows(ggml_tensor * acc, const ggml_tensor * like, int64_t row) const { + return ggml_view_3d(ctx0, acc, like->ne[0], like->ne[1], like->ne[2], + like->ne[0]*sizeof(float), acc->nb[1], row*acc->nb[1]); +} + +void llm_graph_context::build_walk_surrogate(int il, ggml_tensor * k_cur, ggml_tensor * v_cur, int64_t row) const { + ggml_tensor * gk = walk_acc_rows(cparams.walk_acc_k[il], k_cur, row); + ggml_tensor * gv = walk_acc_rows(cparams.walk_acc_v[il], v_cur, row); + ggml_tensor * k32 = k_cur->type == GGML_TYPE_F32 ? k_cur : ggml_cast(ctx0, k_cur, GGML_TYPE_F32); + ggml_tensor * v32 = v_cur->type == GGML_TYPE_F32 ? v_cur : ggml_cast(ctx0, v_cur, GGML_TYPE_F32); + ggml_tensor * term = ggml_add(ctx0, + ggml_sum(ctx0, ggml_mul(ctx0, k32, gk)), + ggml_sum(ctx0, ggml_mul(ctx0, v32, gv))); + res->t_walk_surrogate = res->t_walk_surrogate ? ggml_add(ctx0, res->t_walk_surrogate, term) : term; +} + void llm_graph_context::build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const { if (!(cparams.training && cparams.walk_exact && cparams.walk_state_surrogate)) { return; diff --git a/src/llama-graph.h b/src/llama-graph.h index 7ffae356a7aa..d63901ff10a1 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -962,6 +962,7 @@ class llm_graph_result { struct walk_layer { int il; uint32_t n_past; + uint32_t g0; // the first position whose K/V gradient is kept ggml_tensor * tk = nullptr; // the cached K/V of [grad_from, n_past) as F32, tapped: their ggml_tensor * tv = nullptr; // gradient adds into ak / av (ggml_opt_accumulate_grad) ggml_tensor * ak = nullptr; // the accumulator's rows [grad_from, n_past), shaped like tk / tv @@ -1411,6 +1412,11 @@ struct llm_graph_context { // THE EXACT WALK: the state a chunk leaves (`exit`, what is written back to `cache`) meets the // gradient later chunks put on it, through the surrogate . A no-op outside the walk. void build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const; + // rows [row, row + n) of a layer's walk accumulator, shaped like `like`'s K / V span + ggml_tensor * walk_acc_rows(ggml_tensor * acc, const ggml_tensor * like, int64_t row) const; + // + : what later chunks of the run put on this chunk's own K/V + // (accumulator rows from `row`), carried into it through the loss + void build_walk_surrogate(int il, ggml_tensor * k_cur, ggml_tensor * v_cur, int64_t row) const; ggml_tensor * build_rwkv_token_shift_load( llm_graph_input_rs * inp, diff --git a/src/llama-walk.cpp b/src/llama-walk.cpp index 3ea65e6f8691..73e9fa2449f0 100644 --- a/src/llama-walk.cpp +++ b/src/llama-walk.cpp @@ -63,6 +63,14 @@ void llama_walk_plan::mark_exact(uint32_t horizon) { chunks[first].period_end = true; } +uint32_t llama_walk_plan::longest_run() const { + uint32_t n = 0; + for (const llama_walk_run & r : runs) { + n = std::max(n, end(r) - begin(r)); + } + return n; +} + size_t llama_walk_plan::host_bytes(size_t stride, size_t kv_bytes, size_t state_bytes) const { size_t n_checkpoints = 0; for (const llama_walk_run & r : runs) { @@ -133,13 +141,22 @@ bool llama_walk_acc::init(const llama_hparams & hparams, const llama_memory_recu bound = &cparams; cparams.walk_acc_k = k.data(); cparams.walk_acc_v = v.data(); + cparams.walk_acc_base = 0; return true; } +void llama_walk_acc::start_run(uint32_t base) { + GGML_ASSERT(bound != nullptr); + ggml_backend_buffer_clear(buf.get(), 0); + state_ready = false; + bound->walk_acc_base = base; +} + llama_walk_acc::~llama_walk_acc() { if (bound != nullptr) { bound->walk_acc_k = nullptr; bound->walk_acc_v = nullptr; + bound->walk_acc_base = 0; } } diff --git a/src/llama-walk.h b/src/llama-walk.h index 057d50fd21b3..35ceedfeef8d 100644 --- a/src/llama-walk.h +++ b/src/llama-walk.h @@ -50,9 +50,11 @@ struct llama_walk_plan { static llama_walk_plan of(const std::vector & labels, uint32_t n_ctx, uint32_t n_batch); - // positions [begin, end) of a run + // positions [begin, end) of a run, and the longest run's length (what the exact walk's K/V + // gradient holds: a run's, never the window's) uint32_t begin(const llama_walk_run & r) const { return chunks[r.j0].c0; } uint32_t end (const llama_walk_run & r) const { return chunks[r.j1 - 1].c0 + chunks[r.j1 - 1].n; } + uint32_t longest_run() const; // which chunks the exact walk trains (the labelled ones: the context before a run is her // stable set's and carries no gradient), how far back each gradient reaches (never before @@ -72,11 +74,11 @@ struct llama_walk_plan { }; // THE EXACT WALK's gradients, in host memory, each added into in place by the backward: per -// attention layer dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, position] -// F32 (null for a recurrent layer), which the graph reads through cparams while bound (never -// past the accumulator's life); per recurrent state tensor (a layer's conv state r_l and its -// recurrent state s_l) dL/d(the state entering the chunk trained last), which is what the -// chunk before it puts on the state it leaves (state_ready once one is in) +// attention layer dL/d(cached K) and dL/d(cached V) of the run being trained, [n_embd_*_gqa, +// position - its start] F32 (null for a recurrent layer), which the graph reads through cparams +// while bound (never past the accumulator's life); per recurrent state tensor (a layer's conv +// state r_l and its recurrent state s_l) dL/d(the state entering the chunk trained last), which +// is what the chunk before it puts on the state it leaves (state_ready once one is in) struct llama_walk_acc { std::vector k, v; std::unordered_map state; @@ -90,6 +92,8 @@ struct llama_walk_acc { // zeroed accumulators for n_pos positions, bound into cparams; false: the host memory // could not be allocated bool init(const llama_hparams & hparams, const llama_memory_recurrent * recr, uint32_t n_pos, llama_cparams & cparams); + // zeroed again for the run starting at position base (row 0) + void start_run(uint32_t base); // what init allocates for the K/V of n_pos positions, and one sequence's recurrent state static size_t kv_bytes(const llama_hparams & hparams, uint32_t n_pos); From 0e6e0d949a45c8dd15d6bab1a1bc5de957eed45c Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 08:24:41 -0500 Subject: [PATCH 09/43] ggml: cross-entropy takes one target index per row, so training builds no vocabulary-wide one-hot P1 (continuum #5032 section 9.6, the precision and memory rows). A training step's labels were a one-hot F32 tensor the shape of the logits: [vocab x rows], 485 MiB at a 248k vocabulary and a 512-row chunk, beside the logits themselves. - CROSS_ENTROPY_LOSS and its backward accept I32 labels, one target per row, -1 for a row that carries no loss (ggml_cross_entropy_labels_ok names the two forms). CPU, Metal and CUDA kernels; SYCL refuses the I32 form until it has one. - ggml-opt: GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS (appended to the enum) builds an I32 labels input and counts accuracy against it directly. llama trains with it and writes each output row's target token. - GGML_OPT_TOP=N lists a training graph's N largest tensors at allocation: where its memory goes, measured. - test-backend-ops: CROSS_ENTROPY_LOSS with I32 targets (masked rows included), forward and finite-difference gradient. test_out_prod no longer crashes gradient mode (it built its weight in a null context), which had stopped every grad run. Measured (M5 Metal, 2026-10-14): test-backend-ops CROSS_ENTROPY_LOSS test and grad OK on CPU and MTL0 ([10,5,4,3] and [30000,2,1,1]); test-walk-exact OK on qwen2.5-0.5b (13:17-13:18Z) and qwen3.5-4b hybrid (13:18-13:23Z), every exact row 1.0000/1.0000 and the cache after either walk byte-identical to her stable decode. qwen3.5-4b chunk graph at window 8192, chunk 512, recompute (13:23-13:24Z): 2,417 -> 2,108 MiB on the flash path, 2,265 -> 2,037 MiB explicit. CUDA: not run here; owed on the 5090. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/include/ggml-opt.h | 4 ++ ggml/include/ggml.h | 7 +- ggml/src/ggml-cpu/ops.cpp | 52 +++++++++----- ggml/src/ggml-cuda/cross-entropy-loss.cu | 85 ++++++++++++++++------- ggml/src/ggml-metal/ggml-metal-device.cpp | 4 +- ggml/src/ggml-metal/ggml-metal-device.m | 6 +- ggml/src/ggml-metal/kernels/softmax.metal | 79 +++++++++++++++++++++ ggml/src/ggml-opt.cpp | 33 +++++++-- ggml/src/ggml-sycl/ggml-sycl.cpp | 4 +- ggml/src/ggml.c | 11 ++- src/llama-context-opt.cpp | 30 ++++---- tests/test-backend-ops.cpp | 55 ++++++++++++++- 12 files changed, 295 insertions(+), 75 deletions(-) diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index 1364ce4fd96f..c59190330951 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h @@ -32,6 +32,10 @@ extern "C" { GGML_OPT_LOSS_TYPE_SUM, GGML_OPT_LOSS_TYPE_CROSS_ENTROPY, GGML_OPT_LOSS_TYPE_MEAN_SQUARED_ERROR, + // GGML_OPT_LOSS_TYPE_CROSS_ENTROPY with one I32 target per output row (-1: no loss) in place of a one-hot + // row as wide as the outputs: what a language model's training step needs (its outputs are + // vocabulary-wide logits) + GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS, }; // ====== Dataset ====== diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 1c23bbf98c85..7cefe220106b 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2779,7 +2779,12 @@ extern "C" { int n_tasks, void * userdata); - // loss function + // loss function: the mean over rows of -sum(labels * log softmax(logits)). The labels are + // either a distribution per row (F32, the logits' shape) or one target index per row (I32, + // one per row of the logits, -1 for a row that carries no loss): a target index never builds + // a vocabulary-wide one-hot tensor beside the logits. + + GGML_API bool ggml_cross_entropy_labels_ok(const struct ggml_tensor * logits, const struct ggml_tensor * labels); GGML_API struct ggml_tensor * ggml_cross_entropy_loss( struct ggml_context * ctx, diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index e1b6f0e7bc60..51000117c2bb 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11746,11 +11746,12 @@ static void ggml_compute_forward_cross_entropy_loss_f32( const ggml_tensor * src1 = dst->src[1]; GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_I32); GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type)); GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type)); - GGML_ASSERT(ggml_are_same_shape(src0, src1)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(src0, src1)); GGML_ASSERT(ggml_is_scalar(dst)); + const bool index = src1->type == GGML_TYPE_I32; GGML_ASSERT(dst->type == GGML_TYPE_F32); // TODO: handle transposed/permuted matrices @@ -11775,18 +11776,21 @@ static void ggml_compute_forward_cross_entropy_loss_f32( for (int64_t i1 = ir0; i1 < ir1; ++i1) { const float * s0 = (const float *)((const char *) src0->data + i1*src0->nb[1]); - const float * s1 = (const float *)((const char *) src1->data + i1*src1->nb[1]); - -#ifndef NDEBUG - for (int64_t i = 0; i < nc; ++i) { - //printf("p[%d] = %f\n", i, p[i]); - assert(!isnan(s0[i])); - assert(!isnan(s1[i])); - } -#endif // NDEBUG float max = -INFINITY; ggml_vec_max_f32(nc, &max, s0); + if (index) { + // one target: its log softmax, -(max + log sum exp(x - max)) + x[t] + const int32_t t = ((const int32_t *) src1->data)[i1]; + if (t >= 0) { + GGML_ASSERT(t < nc); + const ggml_float lse = ggml_vec_log_soft_max_f32(nc, st, s0, max); + sum_thread += st[t] - lse; + } + continue; + } + const float * s1 = (const float *)((const char *) src1->data + i1*src1->nb[1]); + const ggml_float sum_softmax = ggml_vec_log_soft_max_f32(nc, st, s0, max); assert(sum_softmax >= 0.0); @@ -11846,7 +11850,8 @@ static void ggml_compute_forward_cross_entropy_loss_back_f32( GGML_ASSERT(ggml_is_contiguous(src0f)); GGML_ASSERT(ggml_is_contiguous(src1f)); GGML_ASSERT(ggml_is_contiguous(grad)); - GGML_ASSERT(ggml_are_same_shape(src0f, src1f) && ggml_are_same_shape(src0f, dst)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(src0f, src1f) && ggml_are_same_shape(src0f, dst)); + const bool index = src1f->type == GGML_TYPE_I32; const int64_t ith = params->ith; const int64_t nth = params->nth; @@ -11867,15 +11872,24 @@ static void ggml_compute_forward_cross_entropy_loss_back_f32( for (int64_t i1 = ir0; i1 < ir1; i1++) { float * ds0 = (float *)((char *) dst->data + i1*dst->nb[1]); const float * s0 = (const float *)((const char *) src0f->data + i1*src0f->nb[1]); - const float * s1 = (const float *)((const char *) src1f->data + i1*src1f->nb[1]); -#ifndef NDEBUG - for (int64_t i = 0; i < nc; ++i) { - //printf("p[%d] = %f\n", i, p[i]); - assert(!isnan(s0[i])); - assert(!isnan(s1[i])); + if (index) { + // one target t: (softmax(x) - onehot(t)) * d / nr; a row without one is zero + const int32_t t = ((const int32_t *) src1f->data)[i1]; + if (t < 0) { + memset(ds0, 0, nc*sizeof(float)); + continue; + } + GGML_ASSERT(t < nc); + float max = -INFINITY; + ggml_vec_max_f32(nc, &max, s0); + const ggml_float sum = ggml_vec_soft_max_f32(nc, ds0, s0, max); + ggml_vec_scale_f32(nc, ds0, 1.0/sum); + ds0[t] -= 1.0f; + ggml_vec_scale_f32(nc, ds0, d_by_nr); + continue; } -#endif // NDEBUG + const float * s1 = (const float *)((const char *) src1f->data + i1*src1f->nb[1]); // soft_max float max = -INFINITY; diff --git a/ggml/src/ggml-cuda/cross-entropy-loss.cu b/ggml/src/ggml-cuda/cross-entropy-loss.cu index 9cfbedd93f3b..277f5cfdbb9e 100644 --- a/ggml/src/ggml-cuda/cross-entropy-loss.cu +++ b/ggml/src/ggml-cuda/cross-entropy-loss.cu @@ -5,13 +5,15 @@ #include #include -template +// index: the labels are one I32 target per row (-1: no loss), not a distribution per row +template static __global__ void cross_entropy_loss_f32( - const float * __restrict__ logits, const float * __restrict__ labels, float * __restrict__ dst, const int nclasses, const int k) { + const float * __restrict__ logits, const void * __restrict__ labels_v, float * __restrict__ dst, const int nclasses, const int k) { extern __shared__ float tmp[]; logits += int64_t(blockIdx.x)*nclasses; - labels += int64_t(blockIdx.x)*nclasses; + const float * labels = index ? nullptr : (const float *) labels_v + int64_t(blockIdx.x)*nclasses; + const int32_t target = index ? ((const int32_t *) labels_v)[blockIdx.x] : -1; // Find maximum for softmax: float max_logit = -INFINITY; @@ -36,11 +38,15 @@ static __global__ void cross_entropy_loss_f32( // log(exp(logits - max) / sum) = (logits - max) - log(sum) float loss = 0.0f; - for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { - const float logit_i = use_shared ? tmp[i] : logits[i]; - loss += (logit_i - max_logit - sum) * labels[i]; + if (index) { + loss = target >= 0 ? -(logits[target] - max_logit - sum) / (float)k : 0.0f; + } else { + for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { + const float logit_i = use_shared ? tmp[i] : logits[i]; + loss += (logit_i - max_logit - sum) * labels[i]; + } + loss = -warp_reduce_sum(loss) / (float)k; } - loss = -warp_reduce_sum(loss) / (float)k; if (threadIdx.x != 0) { return; @@ -49,15 +55,22 @@ static __global__ void cross_entropy_loss_f32( dst[blockIdx.x] = loss; } -template +template static __global__ void cross_entropy_loss_back_f32( - const float * __restrict__ grad, const float * __restrict__ logits, const float * __restrict__ labels, + const float * __restrict__ grad, const float * __restrict__ logits, const void * __restrict__ labels_v, float * __restrict__ dst, const int nclasses) { extern __shared__ float tmp[]; logits += int64_t(blockIdx.x)*nclasses; - labels += int64_t(blockIdx.x)*nclasses; dst += int64_t(blockIdx.x)*nclasses; + const float * labels = index ? nullptr : (const float *) labels_v + int64_t(blockIdx.x)*nclasses; + const int32_t target = index ? ((const int32_t *) labels_v)[blockIdx.x] : -1; + if (index && target < 0) { + for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { + dst[i] = 0.0f; + } + return; + } float maxval = -INFINITY; for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { @@ -85,17 +98,21 @@ static __global__ void cross_entropy_loss_back_f32( // d/dx of -sum(y log softmax(x)) is softmax(x)*sum(y) - y: with sum(y) a masked (all-zero) // label row gets a zero gradient; without it the row would still push softmax(x) back. - float lsum = 0.0f; - for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { - lsum += labels[i]; + float lsum = 1.0f; + if (!index) { + lsum = 0.0f; + for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { + lsum += labels[i]; + } + lsum = warp_reduce_sum(lsum); } - lsum = warp_reduce_sum(lsum); const float sm_scale = lsum/sum; const float d_by_nrows = *grad/gridDim.x; for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { const float val = use_shared ? tmp[i] : dst[i]; - dst[i] = (val*sm_scale - labels[i])*d_by_nrows; + const float y = index ? (i == target ? 1.0f : 0.0f) : labels[i]; + dst[i] = (val*sm_scale - y)*d_by_nrows; } } @@ -104,8 +121,9 @@ void ggml_cuda_cross_entropy_loss(ggml_backend_cuda_context & ctx, ggml_tensor * const ggml_tensor * src1 = dst->src[1]; GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_I32); GGML_ASSERT( dst->type == GGML_TYPE_F32); + const bool index = src1->type == GGML_TYPE_I32; GGML_ASSERT(ggml_is_contiguous(src0)); GGML_ASSERT(ggml_is_contiguous(src1)); @@ -115,7 +133,7 @@ void ggml_cuda_cross_entropy_loss(ggml_backend_cuda_context & ctx, ggml_tensor * const int64_t nrows = ggml_nrows(src0); const float * src0_d = (const float *) src0->data; - const float * src1_d = (const float *) src1->data; + const void * src1_d = src1->data; float * dst_d = (float *) dst->data; ggml_cuda_pool & pool = ctx.pool(); @@ -131,10 +149,17 @@ void ggml_cuda_cross_entropy_loss(ggml_backend_cuda_context & ctx, ggml_tensor * ggml_cuda_pool_alloc dst_tmp(pool, blocks_num.x); if (nbytes_shared <= smpbo) { - CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_f32), smpbo); - cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); + if (index) { + CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_f32), smpbo); + cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); + } else { + CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_f32), smpbo); + cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); + } + } else if (index) { + cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); } else { - cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); + cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); } CUDA_CHECK(cudaGetLastError()); @@ -148,15 +173,16 @@ void ggml_cuda_cross_entropy_loss_back(ggml_backend_cuda_context & ctx, ggml_ten const ggml_tensor * src1f = dst->src[2]; GGML_ASSERT(src0f->type == GGML_TYPE_F32); - GGML_ASSERT(src1f->type == GGML_TYPE_F32); + GGML_ASSERT(src1f->type == GGML_TYPE_F32 || src1f->type == GGML_TYPE_I32); GGML_ASSERT( grad->type == GGML_TYPE_F32); + const bool index = src1f->type == GGML_TYPE_I32; GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(ggml_is_scalar(grad)); GGML_ASSERT(ggml_is_contiguous(src0f)); GGML_ASSERT(ggml_is_contiguous(src1f)); GGML_ASSERT(ggml_is_contiguous(dst)); - GGML_ASSERT(ggml_are_same_shape(src0f, src1f)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(src0f, src1f)); GGML_ASSERT(ggml_are_same_shape(src0f, dst)); const int64_t ne00 = src0f->ne[0]; @@ -164,7 +190,7 @@ void ggml_cuda_cross_entropy_loss_back(ggml_backend_cuda_context & ctx, ggml_ten const float * grad_d = (const float *) grad->data; const float * src0f_d = (const float *) src0f->data; - const float * src1f_d = (const float *) src1f->data; + const void * src1f_d = src1f->data; float * dst_d = (float *) dst->data; cudaStream_t stream = ctx.stream(); @@ -177,9 +203,16 @@ void ggml_cuda_cross_entropy_loss_back(ggml_backend_cuda_context & ctx, ggml_ten const size_t smpbo = ggml_cuda_info().devices[id].smpbo; if (nbytes_shared <= smpbo) { - CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_back_f32), smpbo); - cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); + if (index) { + CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_back_f32), smpbo); + cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); + } else { + CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_back_f32), smpbo); + cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); + } + } else if (index) { + cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); } else { - cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); + cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); } } diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 22a37453bf39..08b2ef272896 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -340,7 +340,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_lo char base[256]; char name[256]; - snprintf(base, 256, "kernel_cross_entropy_loss%s_f32", op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK ? "_back" : ""); + const bool back = op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK; + const ggml_tensor * labels = back ? op->src[2] : op->src[1]; + snprintf(base, 256, "kernel_cross_entropy_loss%s_%s", back ? "_back" : "", labels->type == GGML_TYPE_I32 ? "i32" : "f32"); snprintf(name, 256, "%s", base); ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c916ca8a1657..0ff9e84c17ff 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1478,10 +1478,12 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_REPEAT_BACK: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op); case GGML_OP_CROSS_ENTROPY_LOSS: - return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && + (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_I32) && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_F32 && + return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_I32) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) && ggml_is_contiguous(op); case GGML_OP_OUT_PROD: // The backward through a frozen weight: out_prod(W, grad^T). Composed on this diff --git a/ggml/src/ggml-metal/kernels/softmax.metal b/ggml/src/ggml-metal/kernels/softmax.metal index a6ad5baaa9bd..53e30fd18063 100644 --- a/ggml/src/ggml-metal/kernels/softmax.metal +++ b/ggml/src/ggml-metal/kernels/softmax.metal @@ -350,6 +350,85 @@ kernel void kernel_cross_entropy_loss_back_f32( } +// the same with one target index per row (I32 labels, -1: the row carries no loss), so no +// vocabulary-wide one-hot tensor exists beside the logits +kernel void kernel_cross_entropy_loss_i32( + constant ggml_metal_kargs_cross_entropy_loss & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t row = tgpig.x; + const int32_t t = ((device const int32_t *) src1)[row]; + + device const float * s0 = (device const float *) (src0 + row*args.nb01); + + float lmax = -INFINITY; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lmax = MAX(lmax, s0[i]); + } + const float max_val = ce_reduce_max(lmax, buf, sgitg, tiisg, tptg); + + float lsum = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lsum += exp(s0[i] - max_val); + } + const float lse = log(ce_reduce_sum(lsum, buf, sgitg, tiisg, tptg)); + + if (tpitg.x == 0) { + ((device float *) dst)[row] = t >= 0 ? -(s0[t] - max_val - lse)/(float) args.nr : 0.0f; + } +} + +kernel void kernel_cross_entropy_loss_back_i32( + constant ggml_metal_kargs_cross_entropy_loss & args, + device const char * src0, // grad: one float + device const char * src1, // logits + device const char * src2, // labels: one target per row + device char * dst, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t row = tgpig.x; + const int32_t t = ((device const int32_t *) src2)[row]; + + device const float * s0 = (device const float *) (src1 + row*args.nb01); + device float * d = (device float *) (dst + row*args.nb1); + + if (t < 0) { + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + d[i] = 0.0f; + } + return; + } + + const float d_by_nr = ((device const float *) src0)[0]/(float) args.nr; + + float lmax = -INFINITY; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lmax = MAX(lmax, s0[i]); + } + const float max_val = ce_reduce_max(lmax, buf, sgitg, tiisg, tptg); + + float lsum = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lsum += exp(s0[i] - max_val); + } + const float inv_sum = 1.0f/ce_reduce_sum(lsum, buf, sgitg, tiisg, tptg); + + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + d[i] = (exp(s0[i] - max_val)*inv_sum - (i == t ? 1.0f : 0.0f))*d_by_nr; + } +} + // RMS_NORM_BACK: dx = (dz - x * sum(x*dz)/(sum(x*x) + eps*n)) / sqrt(sum(x*x)/n + eps) // (the CPU reference, ggml-cpu ops.cpp rms_norm_back_f32). src0 = dz, src1 = x. kernel void kernel_rms_norm_back_f32( diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 6bf7d7c23b14..aa5fae235b46 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -585,8 +585,11 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { opt_ctx->loss_per_datapoint = false; break; } - case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY: { - opt_ctx->labels = ggml_dup_tensor(ctx_results, opt_ctx->outputs); + case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY: + case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS: { + opt_ctx->labels = opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS + ? ggml_new_tensor_1d(ctx_results, GGML_TYPE_I32, ggml_nrows(opt_ctx->outputs)) + : ggml_dup_tensor(ctx_results, opt_ctx->outputs); ggml_set_input(opt_ctx->labels); ggml_set_name(opt_ctx->labels, "labels"); opt_ctx->loss = ggml_cross_entropy_loss(ctx_results, opt_ctx->outputs, opt_ctx->labels); @@ -620,7 +623,7 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { ggml_set_output(opt_ctx->loss); if (opt_ctx->next_manual) { GGML_ASSERT(!opt_ctx->static_graphs && "a caller-driven period needs graphs built per step"); - GGML_ASSERT(opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY); + GGML_ASSERT(opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY || opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS); // the loss the backward differentiates: the caller's weight of this graph's loss, plus // its extra term; opt_ctx->loss stays the unweighted loss the result reports struct ggml_tensor * total = ggml_scale(ctx_results, opt_ctx->loss, opt_ctx->next_loss_scale); @@ -637,13 +640,15 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { ggml_build_forward_expand(opt_ctx->gf, opt_ctx->loss); } - if (opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY) { + if (opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY || opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS) { opt_ctx->pred = ggml_argmax(ctx_results, opt_ctx->outputs); ggml_set_name(opt_ctx->pred, "pred"); ggml_set_output(opt_ctx->pred); ggml_build_forward_expand(opt_ctx->gf, opt_ctx->pred); - opt_ctx->ncorrect = ggml_count_equal(ctx_results, opt_ctx->pred, ggml_argmax(ctx_results, opt_ctx->labels)); + ggml_tensor * truth = opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS + ? opt_ctx->labels : ggml_argmax(ctx_results, opt_ctx->labels); + opt_ctx->ncorrect = ggml_count_equal(ctx_results, opt_ctx->pred, truth); ggml_set_name(opt_ctx->ncorrect, "ncorrect"); ggml_set_output(opt_ctx->ncorrect); ggml_build_forward_expand(opt_ctx->gf, opt_ctx->ncorrect); @@ -1172,6 +1177,24 @@ bool ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { } opt_ctx->allocated_graph = graph; + // GGML_OPT_TOP=N: the graph's N largest tensors (views excluded), where its memory goes + if (const char * top = getenv("GGML_OPT_TOP")) { + ggml_cgraph * g = opt_ctx->allocated_graph_copy; + std::vector> rows; + for (int i = 0; i < g->n_nodes; ++i) { + if (g->nodes[i]->view_src == nullptr) { + rows.emplace_back(ggml_nbytes(g->nodes[i]), g->nodes[i]); + } + } + std::sort(rows.begin(), rows.end(), [](const auto & x, const auto & y) { return x.first > y.first; }); + const size_t n = std::min(rows.size(), (size_t) atoi(top)); + for (size_t i = 0; i < n; ++i) { + const ggml_tensor * t = rows[i].second; + GGML_LOG_INFO("%s: top %2zu: %8.1f MiB %-12s %-6s [%lld,%lld,%lld,%lld] %s\n", __func__, i, rows[i].first/1048576.0, + ggml_op_desc(t), ggml_type_name(t->type), (long long) t->ne[0], (long long) t->ne[1], (long long) t->ne[2], (long long) t->ne[3], t->name); + } + } + opt_ctx->eval_ready = true; return true; } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d834e..4b00fdace548 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6336,9 +6336,11 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_FILL: case GGML_OP_CUMSUM: case GGML_OP_DIAG: + return true; case GGML_OP_CROSS_ENTROPY_LOSS: + return op->src[1]->type == GGML_TYPE_F32; // no I32 targets on this backend yet case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - return true; + return op->src[2]->type == GGML_TYPE_F32; case GGML_OP_SOLVE_TRI: return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K; case GGML_OP_FLASH_ATTN_EXT: diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index ab9c8911059e..0b58d97094fe 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6198,11 +6198,18 @@ struct ggml_tensor * ggml_custom_inplace( } // ggml_cross_entropy_loss +bool ggml_cross_entropy_labels_ok(const struct ggml_tensor * logits, const struct ggml_tensor * labels) { + if (labels->type == GGML_TYPE_I32) { + return ggml_is_contiguous(labels) && ggml_nelements(labels) == ggml_nrows(logits); + } + return ggml_are_same_shape(logits, labels); +} + struct ggml_tensor * ggml_cross_entropy_loss( struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b) { - GGML_ASSERT(ggml_are_same_shape(a, b)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(a, b)); struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); @@ -6221,7 +6228,7 @@ struct ggml_tensor * ggml_cross_entropy_loss_back( struct ggml_tensor * b, struct ggml_tensor * c) { GGML_ASSERT(ggml_is_scalar(a)); - GGML_ASSERT(ggml_are_same_shape(b, c)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(b, c)); struct ggml_tensor * result = ggml_dup_tensor(ctx, b); diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 5813e004b71f..c4e39fc1c1f3 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -76,7 +76,8 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params max_nodes, cparams.pipeline_parallel, cparams.op_offload)); } - ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY); + // one target token per row: no one-hot tensor as wide as the vocabulary beside the logits + ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS); opt_params.opt_period = n_batch / n_ubatch; opt_params.get_opt_pars = lopt_params.get_opt_pars; // per-layer recompute: the layer outputs are the checkpoints (llm_graph names them l_out-N) @@ -355,26 +356,21 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l res->set_inputs(&ubatch); { - // one-hot rows for the labelled positions, in the order the ubatch emits its outputs - // (the order the logits rows take); an unlabelled output row stays zero, so it adds no - // LOCAL loss (its sequence's context still carries gradient to later labelled rows) + // the target token of each output row, in the order the ubatch emits its outputs (the + // order the logits rows take) struct ggml_tensor * labels_t = ggml_opt_labels(opt_ctx); - GGML_ASSERT(labels_t->ne[1] == (int64_t) n_outputs); - ggml_set_zero(labels_t); - const float onef = 1.0f; - uint32_t row = 0; + GGML_ASSERT(labels_t->type == GGML_TYPE_I32 && labels_t->ne[0] == (int64_t) n_outputs); + std::vector targets; + targets.reserve(n_outputs); for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { - if (ubatch.output[i] == 0) { - continue; + if (ubatch.output[i] != 0) { + const llama_token y = ubatch.label != nullptr ? ubatch.label[i] : -1; + GGML_ASSERT(y < (llama_token) model.vocab.n_tokens()); + targets.push_back(y); } - const llama_token y = ubatch.label != nullptr ? ubatch.label[i] : -1; - GGML_ASSERT(y < labels_t->ne[0]); - if (y >= 0) { - ggml_backend_tensor_set(labels_t, &onef, (row*labels_t->ne[0] + y)*sizeof(float), sizeof(float)); - } - ++row; } - GGML_ASSERT(row == n_outputs); + GGML_ASSERT(targets.size() == n_outputs); + ggml_backend_tensor_set(labels_t, targets.data(), 0, ggml_nbytes(labels_t)); } if (xc != nullptr) { // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 06389e69a212..fa80a8573970 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -5078,7 +5078,7 @@ struct test_out_prod : public test_case { } ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) override { - ggml_tensor * a = ggml_new_tensor_4d(a_own_buffer ? ctx_weights : ctx, type_a, m, k, bs[0], bs[1]); + ggml_tensor * a = ggml_new_tensor_4d(a_own_buffer && ctx_weights ? ctx_weights : ctx, type_a, m, k, bs[0], bs[1]); // grad mode passes no weights context ggml_set_name(a, "a"); ggml_tensor * b; @@ -7491,6 +7491,57 @@ static void ce_test_init_mask(ggml_tensor * t) { ggml_backend_tensor_set(t, m.data(), 0, m.size()*sizeof(float)); } +// GGML_OP_CROSS_ENTROPY_LOSS with one I32 target per row (-1: a row with no loss), the form a +// training step's labels take so no vocabulary-wide one-hot tensor is built +struct test_cross_entropy_loss_index : public test_case { + const std::array ne; + + std::string vars() override { + return VARS_TO_STR1(ne); + } + + test_cross_entropy_loss_index(std::array ne = {10, 5, 4, 3}) : ne(ne) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * logits = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); + ggml_set_param(logits); + ggml_set_name(logits, "logits"); + + ggml_tensor * targets = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[1]*ne[2]*ne[3]); + ggml_set_name(targets, "targets"); + + ggml_tensor * out = ggml_cross_entropy_loss(ctx, logits, targets); + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + std::default_random_engine rng(7); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_I32) { + // every third row masked (-1), the rest a random target + std::uniform_int_distribution pick(0, (int32_t) ne[0] - 1); + std::vector v(ggml_nelements(t)); + for (size_t i = 0; i < v.size(); ++i) { + v[i] = i % 3 == 2 ? -1 : pick(rng); + } + ggml_backend_tensor_set(t, v.data(), 0, ggml_nbytes(t)); + } else { + init_tensor_uniform(t, -3.0f, 3.0f); + } + } + } + + float grad_eps() override { + return 1.0f; + } + + bool grad_precise() override { + return true; + } +}; + struct test_cross_entropy_loss : public test_case { const ggml_type type; const std::array ne; @@ -10340,6 +10391,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false)); + test_cases.emplace_back(new test_cross_entropy_loss_index()); + test_cases.emplace_back(new test_cross_entropy_loss_index({30000, 2, 1, 1})); // a vocabulary's width; few rows keep the finite difference above f32 rounding test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); From e231baa8eb888d7ff301f2c40e0cb194f1e200c0 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 08:53:42 -0500 Subject: [PATCH 10/43] ggml-opt: a training step's head is computed a slice of rows at a time P1 (continuum #5032 section 9.6, the memory rows). The logits of a chunk and their gradient were the two largest tensors of its graph, [vocab x rows] F32 each (485 MiB apiece on the 4B at a 248k vocabulary and 512 rows), and their size does not shrink with the model. - ggml_opt_prepare_alloc_head: given the hidden state and the outputs (a row-wise function of it: MUL_MAT, MUL, ADD, SCALE, as a language model's head with its adapters and scales is), the head leaves the forward graph and is built again for each slice of rows; the loss is the mean over every row, each slice weighted by its share. Under recompute only one slice's logits and their gradient exist at a time. A head with any other node is not sliced. - llama trains with it, OPT_HEAD_SLICE_BYTES (64 MiB) of logits per slice: 67 rows at a 248k vocabulary. - ggml-opt rebuilt a per-step graph's loss into its predecessor's tensors when it chained the slices, and kept pred/ncorrect from the previous graph when a graph had none: both reset per graph now (the second read a freed tensor). Measured (M5 Metal, 2026-10-14): test-walk-exact OK on qwen2.5-0.5b (13:45:52-13:46:41Z) and qwen3.5-4b hybrid (13:46:41-13:52:10Z): every exact row 1.0000/1.0000, the cache after either walk byte-identical to her stable decode, the forward loss identical to the unsliced head (0.5B chunk 512: 4.122972 before and after). qwen3.5-4b chunk graph at window 8192, chunk 512, recompute (13:52:10-13:53:28Z): 2,108 -> 1,623 MiB on the flash path, 2,037 -> 1,552 explicit; the largest tensors are now the flash-attention backward's score tiles. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/include/ggml-opt.h | 16 +++++ ggml/src/ggml-opt.cpp | 130 ++++++++++++++++++++++++++++++++++++-- src/llama-context-opt.cpp | 10 ++- 3 files changed, 150 insertions(+), 6 deletions(-) diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index c59190330951..707a6ee3dbe7 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h @@ -192,6 +192,22 @@ extern "C" { struct ggml_tensor * inputs, struct ggml_tensor * outputs); + // The same, for an `outputs` that is a row-wise function of `hidden` (a language model's + // head: logits = W x, with its adapters and scales; MUL_MAT, MUL, ADD and SCALE nodes): that + // head is taken out of gf and computed for `rows` rows of hidden at a time, and the loss is + // the mean over every row (GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS only). Under recompute + // (checkpoint_prefix) only one slice's outputs and their gradient exist at a time, never the + // [n_outputs x rows] tensor and its gradient. A head with any other node is not sliced. + // Returns the number of slices (1: not sliced; accuracy is reported only then). + GGML_API int ggml_opt_prepare_alloc_head( + ggml_opt_context_t opt_ctx, + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * inputs, + struct ggml_tensor * hidden, + struct ggml_tensor * outputs, + int64_t rows); + // allocate the next graph for evaluation, either forward or forward + backward // must be called exactly once prior to calling ggml_opt_eval // returns false when the graph could not be allocated (device memory): the graph is then NOT diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index aa5fae235b46..0611db11c66c 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -93,6 +93,8 @@ struct ggml_opt_context { // into a caller's tensor, for the graph prepared last std::unordered_map> next_grad_rows; std::vector grad_accs; + // ggml_opt_prepare_alloc_head: the outputs' row slices, in row order (empty: one outputs) + std::vector output_slices; int64_t iter = 1; int32_t opt_period = 1; @@ -505,6 +507,18 @@ static ggml_cgraph * dup_graph(ggml_context * ctx, ggml_cgraph * src) { return dst; } +// the outputs' rows: every slice's when sliced +static int64_t ggml_opt_output_rows(ggml_opt_context_t opt_ctx) { + if (opt_ctx->output_slices.empty()) { + return opt_ctx->outputs->ne[1]; + } + int64_t n = 0; + for (const ggml_tensor * out : opt_ctx->output_slices) { + n += out->ne[1]; + } + return n; +} + static void ggml_opt_build(ggml_opt_context_t opt_ctx) { GGML_ASSERT(opt_ctx->ctx_compute && "no compute context set, either use static graphs or set one with ggml_opt_prepare_alloc"); GGML_ASSERT((!opt_ctx->static_graphs || opt_ctx->inputs->data) && "when using static graphs the inputs must be allocated statically"); @@ -518,7 +532,9 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { opt_ctx->optimizer == GGML_OPT_OPTIMIZER_TYPE_ADAMW; ggml_set_input(opt_ctx->inputs); - ggml_set_output(opt_ctx->outputs); + if (opt_ctx->output_slices.empty()) { + ggml_set_output(opt_ctx->outputs); + } int n_param = 0; for (int i = 0; i < opt_ctx->gf->n_nodes; ++i) { @@ -588,11 +604,26 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY: case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS: { opt_ctx->labels = opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS - ? ggml_new_tensor_1d(ctx_results, GGML_TYPE_I32, ggml_nrows(opt_ctx->outputs)) + ? ggml_new_tensor_1d(ctx_results, GGML_TYPE_I32, ggml_opt_output_rows(opt_ctx)) : ggml_dup_tensor(ctx_results, opt_ctx->outputs); ggml_set_input(opt_ctx->labels); ggml_set_name(opt_ctx->labels, "labels"); - opt_ctx->loss = ggml_cross_entropy_loss(ctx_results, opt_ctx->outputs, opt_ctx->labels); + if (opt_ctx->output_slices.empty()) { + opt_ctx->loss = ggml_cross_entropy_loss(ctx_results, opt_ctx->outputs, opt_ctx->labels); + } else { + // the mean over every row: each slice's mean weighted by its share of the rows + const int64_t n = opt_ctx->labels->ne[0]; + int64_t r0 = 0; + ggml_tensor * sum = nullptr; + for (ggml_tensor * out : opt_ctx->output_slices) { + const int64_t nr = out->ne[1]; + ggml_tensor * lab = ggml_view_1d(ctx_results, opt_ctx->labels, nr, r0*sizeof(int32_t)); + ggml_tensor * l = ggml_scale(ctx_results, ggml_cross_entropy_loss(ctx_results, out, lab), (float) nr / (float) n); + sum = sum ? ggml_add(ctx_results, sum, l) : l; + r0 += nr; + } + opt_ctx->loss = sum; + } ggml_set_name(opt_ctx->loss, "loss_cross_entropy"); if (opt_ctx->next_manual) { // weighted by the caller below, beside its extra term @@ -640,7 +671,11 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { ggml_build_forward_expand(opt_ctx->gf, opt_ctx->loss); } - if (opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY || opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS) { + // per-step graphs rebuild these: one left from the previous graph would be read after it is freed + opt_ctx->pred = nullptr; + opt_ctx->ncorrect = nullptr; + if ((opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY || opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS) + && opt_ctx->output_slices.empty()) { opt_ctx->pred = ggml_argmax(ctx_results, opt_ctx->outputs); ggml_set_name(opt_ctx->pred, "pred"); ggml_set_output(opt_ctx->pred); @@ -1003,6 +1038,91 @@ void ggml_opt_prepare_alloc( opt_ctx->outputs = outputs; opt_ctx->next_grad_accs.clear(); // registrations belong to the graph prepared last opt_ctx->next_grad_rows.clear(); + opt_ctx->output_slices.clear(); +} + +int ggml_opt_prepare_alloc_head( + ggml_opt_context_t opt_ctx, + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * inputs, + struct ggml_tensor * hidden, + struct ggml_tensor * outputs, + int64_t rows) { + ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, outputs); + const int64_t n = hidden->ne[1]; + if (opt_ctx->loss_type != GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS || rows <= 0 || rows >= n || + !ggml_is_contiguous(hidden) || ggml_n_dims(hidden) > 2 || outputs->ne[1] != n) { + return 1; + } + // the head: the nodes between hidden and outputs, in graph order + std::unordered_set on_path = { hidden }; + std::vector head; + for (int i = 0; i < gf->n_nodes; ++i) { + ggml_tensor * node = gf->nodes[i]; + bool reads_path = false; + for (int k = 0; k < GGML_MAX_SRC; ++k) { + reads_path |= node->src[k] != nullptr && on_path.count(node->src[k]) > 0; + } + if (!reads_path) { + continue; + } + const bool replayable = node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL || + node->op == GGML_OP_ADD || node->op == GGML_OP_SCALE; + if (!replayable || node->view_src != nullptr) { + return 1; // a head this does not know how to replay per slice: the whole outputs + } + on_path.insert(node); + head.push_back(node); + } + if (head.empty() || head.back() != outputs) { + return 1; + } + // nothing else in the graph may read the head (it is about to leave the graph) + for (int i = 0; i < gf->n_nodes; ++i) { + const ggml_tensor * node = gf->nodes[i]; + if (on_path.count(node) > 0) { + continue; + } + for (int k = 0; k < GGML_MAX_SRC; ++k) { + if (node->src[k] != nullptr && node->src[k] != hidden && on_path.count(node->src[k]) > 0) { + return 1; + } + } + } + int j = 0; + for (int i = 0; i < gf->n_nodes; ++i) { + if (gf->nodes[i] == hidden || on_path.count(gf->nodes[i]) == 0) { + gf->nodes[j++] = gf->nodes[i]; + } + } + gf->n_nodes = j; + // the head again, one slice of rows at a time + for (int64_t r0 = 0; r0 < n; r0 += rows) { + const int64_t nr = std::min(rows, n - r0); + std::unordered_map as_slice; + as_slice[hidden] = ggml_view_2d(ctx_compute, hidden, hidden->ne[0], nr, hidden->nb[1], r0*hidden->nb[1]); + for (ggml_tensor * node : head) { + auto src = [&](int k) { const auto it = as_slice.find(node->src[k]); return it != as_slice.end() ? it->second : node->src[k]; }; + ggml_tensor * t = nullptr; + switch (node->op) { + case GGML_OP_MUL_MAT: t = ggml_mul_mat(ctx_compute, src(0), src(1)); break; + case GGML_OP_MUL: t = ggml_mul (ctx_compute, src(0), src(1)); break; + case GGML_OP_ADD: t = ggml_add (ctx_compute, src(0), src(1)); break; + case GGML_OP_SCALE: { + float s, b; + memcpy(&s, (const float *) node->op_params + 0, sizeof(float)); + memcpy(&b, (const float *) node->op_params + 1, sizeof(float)); + t = ggml_scale_bias(ctx_compute, src(0), s, b); + } break; + default: GGML_ABORT("not a head op"); + } + ggml_format_name(t, "%s-rows%lld", node->name, (long long) r0); + as_slice[node] = t; + } + opt_ctx->output_slices.push_back(as_slice.at(outputs)); + } + return (int) opt_ctx->output_slices.size(); } void ggml_opt_set_alloc_budget(ggml_opt_context_t opt_ctx, size_t bytes) { @@ -1283,7 +1403,7 @@ void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result) { GGML_ASSERT(result->opt_period == opt_ctx->opt_period); } - const int64_t ndata = opt_ctx->outputs->ne[1]; + const int64_t ndata = ggml_opt_output_rows(opt_ctx); result->ndata += ndata; result->ndata_batch.push_back(ndata); diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index c4e39fc1c1f3..4891ca771e08 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -304,6 +304,9 @@ bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { return true; } +// what one slice of the training head's logits may hold (see opt_ubatch) +static constexpr size_t OPT_HEAD_SLICE_BYTES = 64u*1024*1024; + bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch) { // the loss is the mean over the logits rows: a serving row among them would dilute it if (ubatch.n_labels() != n_outputs) { @@ -325,7 +328,12 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l }; ctx_compute_opt = ggml_init(params); } - ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); + // THE HEAD A SLICE OF ROWS AT A TIME: the logits of a chunk and their gradient were the + // largest tensors of its graph ([vocab x rows] F32 each: 485 MiB apiece at a 248k vocabulary + // and 512 rows, the 4B at window 8192, 2026-10-14), and their size does not shrink with the + // model. A slice holds OPT_HEAD_SLICE_BYTES of logits. + const int64_t head_rows = std::max(1, (int64_t) (OPT_HEAD_SLICE_BYTES / (sizeof(float) * (size_t) model.vocab.n_tokens()))); + ggml_opt_prepare_alloc_head(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_embd(), res->get_logits(), head_rows); if (xc != nullptr) { ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); for (const auto & io : res->t_walk) { From 88c5c1e26c9691debab279ae2ea529c12c758820 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 09:13:07 -0500 Subject: [PATCH 11/43] ggml: the attention backward is one op (ggml_flash_attn_back), on CPU and Metal P1 (continuum #5032 section 9.6, the memory rows). FLASH_ATTN_EXT's backward was composed from generic ops over 4096-position tiles, three passes, keeping [tile x rows x heads] F32 score tensors (128 MiB each on the 4B at chunk 512) and adding up to 96 graph nodes per tile per layer, so it grew with the window. - ggml_flash_attn_back replaces the dead stub of the same name (it aborted: "TODO: adapt to ggml_flash_attn_ext"): q, k, v, mask in flash_attn_ext's layouts, d the output's gradient; one F32 result holding dQ, dK, dV (ggml_flash_attn_back_offset) and each query row's log-sum-exp and D. Scores are recomputed, never kept. - CPU: the reference (a pass over query rows, then over KV positions with the GQA group summed). - Metal: kernel_flash_attn_back_q / _kv for F32, F16 and q8_0 caches, one simdgroup per row; dq = scale (sum p dp k - D sum p k), so two sweeps over the window instead of three. - FLASH_ATTN_EXT's backward builds the op when the forward's back tile is GGML_FA_BACK_NATIVE; llama selects it when every GPU it computes on has the kernel (probed at the model's attention shape and the cache's types in opt_init), else the composed tiles (CUDA until its kernel). - test-backend-ops: FLASH_ATTN_BACK each backend against the CPU (F16, F32, q8_0, GQA, head 256, odd rows); FLASH_ATTN_EXT gradients through either backward, by finite difference. Measured (M5, 2026-10-14): test-backend-ops FLASH_ATTN_BACK 4/4 on MTL0 against CPU. Finite differences: the native and composed backward pass the direct check and sit at the same level on the SUM-wrapped one (1e-4 to 7e-4 against a 1e-4 bound, varying with the random init, for both): that check does not separate them; the walk test below does. test-walk-exact (native selected) OK on qwen2.5-0.5b (14:04:36-14:05:28Z) and qwen3.5-4b hybrid (14:05:28-14:11:14Z): every exact row 1.0000/1.0000, the cache byte-identical after either walk; the flash path against the F32 explicit gradient went from 0.9999/0.9995 to 1.0000/1.0000 on the 0.5B at an f16 cache. qwen3.5-4b chunk graph, window 8192, chunk 512, recompute (14:11:14-14:12:43Z): 1,623 -> 1,306 MiB on the flash path. SLOWER for now: the 4B walk test took 346 s against 329 s with the composed backward (the kernel is dot products per row, not simdgroup matrices); the matrix version is next. CUDA: no kernel yet (the composed path runs there). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/include/ggml.h | 16 +- ggml/src/ggml-cpu/ggml-cpu.c | 23 +- ggml/src/ggml-cpu/ops.cpp | 452 +++++++--------------- ggml/src/ggml-cpu/ops.h | 5 +- ggml/src/ggml-metal/ggml-metal-device.cpp | 14 + ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-device.m | 9 + ggml/src/ggml-metal/ggml-metal-impl.h | 27 ++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 61 +++ ggml/src/ggml-metal/ggml-metal-ops.h | 1 + ggml/src/ggml-metal/kernels/fa.metal | 190 +++++++++ ggml/src/ggml.c | 109 +++--- src/llama-context-opt.cpp | 42 ++ src/llama-context.h | 1 + src/llama-cparams.h | 3 + src/llama-graph.cpp | 2 +- tests/test-backend-ops.cpp | 102 +++++ 17 files changed, 673 insertions(+), 385 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 7cefe220106b..4597ee2cded6 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2480,14 +2480,26 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * sinks); - // TODO: needs to be adapted to ggml_flash_attn_ext + // The gradients of ggml_flash_attn_ext(q, k, v, mask, scale) (no ALiBi, softcap or sinks), given + // d, the gradient of its output in the output's layout [Dv, H, N]. One op, recomputing the + // scores per block instead of keeping them: memory is the gradients alone. The result is F32: + // dQ ([D, N, H]), dK (k's shape), dV (v's shape), each contiguous at + // ggml_flash_attn_back_offset, then each query row's log-sum-exp and D = sum_j p_ij dp_ij, + // which the kernels pass between their passes. GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, struct ggml_tensor * v, + struct ggml_tensor * mask, struct ggml_tensor * d, - bool masked); + float scale); + + // byte offset in a ggml_flash_attn_back result of: 0 dQ, 1 dK, 2 dV, 3 the per-row scratch + GGML_API size_t ggml_flash_attn_back_offset(const struct ggml_tensor * back, int part); + + // the backward of FLASH_ATTN_EXT as the one op above (ggml_flash_attn_ext_set_back_tile) + #define GGML_FA_BACK_NATIVE (-1) GGML_API struct ggml_tensor * ggml_ssm_conv( struct ggml_context * ctx, diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 61e0c9adb136..c2610efc0375 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2003,10 +2003,7 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm } break; case GGML_OP_FLASH_ATTN_BACK: { - int32_t t = ggml_get_op_params_i32(tensor, 0); - GGML_ASSERT(t == 0 || t == 1); - bool masked = t != 0; - ggml_compute_forward_flash_attn_back(params, masked, tensor); + ggml_compute_forward_flash_attn_back(params, tensor); } break; case GGML_OP_SSM_CONV: { @@ -2975,19 +2972,11 @@ struct ggml_cplan ggml_graph_plan( } break; case GGML_OP_FLASH_ATTN_BACK: { - const int64_t D = node->src[0]->ne[0]; - const int64_t ne11 = ggml_up(node->src[1]->ne[1], GGML_SOFT_MAX_UNROLL); - const int64_t mxDn = MAX(D, ne11) * 2; // *2 because of S and SM in ggml_compute_forward_flash_attn_back - if (node->src[1]->type == GGML_TYPE_F32) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } else if (node->src[1]->type == GGML_TYPE_F16) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } else if (node->src[1]->type == GGML_TYPE_BF16) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } + // per thread: a row's scores, then a K, V, query and output-gradient row + const int64_t D = node->src[0]->ne[0]; + const int64_t M = node->src[1]->ne[1]; + const int64_t Dv = node->src[2]->ne[0]; + cur = sizeof(float)*(M + 2*D + 2*Dv)*n_tasks; } break; case GGML_OP_CROSS_ENTROPY_LOSS: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 51000117c2bb..8c7c8eb7137c 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -9224,316 +9224,166 @@ void ggml_compute_forward_flash_attn_ext( } // ggml_compute_forward_flash_attn_back +// +// The gradients of FLASH_ATTN_EXT (ggml_flash_attn_back), the reference every backend's kernel is +// checked against. Pass 1, over query rows (i, h): the row's log-sum-exp, D_i = sum_j p_ij dp_ij +// (from the same recomputed p and dp the gradients use, never from the forward's output) and +// dq_i = scale sum_j ds_ij k_j, with ds_ij = p_ij (dp_ij - D_i). Pass 2, over KV positions +// (j, kv head): dk_j and dv_j summed over the query rows and the KV head's group of query heads. +// Nothing window x rows is kept: a row's scores live in one per-thread buffer. + +static inline float ggml_fa_back_mask(const ggml_tensor * mask, int64_t j, int64_t i, int64_t h) { + if (mask == nullptr) { + return 0.0f; + } + const char * m = (const char *) mask->data + i*mask->nb[1] + (h % mask->ne[2])*mask->nb[2]; + return GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) m)[j]); +} static void ggml_compute_forward_flash_attn_back_f32( const ggml_compute_params * params, - const bool masked, - ggml_tensor * dst) { + ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * d = dst->src[4]; - const ggml_tensor * q = dst->src[0]; - const ggml_tensor * k = dst->src[1]; - const ggml_tensor * v = dst->src[2]; - const ggml_tensor * d = dst->src[3]; + const float scale = ggml_get_op_params_f32(dst, 0); - GGML_TENSOR_LOCALS(int64_t, neq, q, ne) - GGML_TENSOR_LOCALS(size_t, nbq, q, nb) - GGML_TENSOR_LOCALS(int64_t, nek, k, ne) - GGML_TENSOR_LOCALS(size_t, nbk, k, nb) - GGML_TENSOR_LOCALS(int64_t, nev, v, ne) - GGML_TENSOR_LOCALS(size_t, nbv, v, nb) - GGML_TENSOR_LOCALS(int64_t, ned, d, ne) - GGML_TENSOR_LOCALS(size_t, nbd, d, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + const int64_t D = q->ne[0]; + const int64_t N = q->ne[1]; + const int64_t H = q->ne[2]; + const int64_t M = k->ne[1]; + const int64_t Hkv = k->ne[2]; + const int64_t Dv = v->ne[0]; + const int64_t G = H / Hkv; // query head h reads KV head h / G + + float * dq = (float *) ((char *) dst->data + ggml_flash_attn_back_offset(dst, 0)); + float * dk = (float *) ((char *) dst->data + ggml_flash_attn_back_offset(dst, 1)); + float * dv = (float *) ((char *) dst->data + ggml_flash_attn_back_offset(dst, 2)); + float * lse = (float *) ((char *) dst->data + ggml_flash_attn_back_offset(dst, 3)); + float * dvec = lse + N*H; const int ith = params->ith; const int nth = params->nth; - const int64_t D = neq0; - const int64_t N = neq1; - const int64_t P = nek1 - N; - const int64_t M = P + N; - - const int Mup = ggml_up(M, GGML_SOFT_MAX_UNROLL); - const int mxDM = MAX(D, Mup); - - // GGML_ASSERT(ne0 == D); - // GGML_ASSERT(ne1 == N); - GGML_ASSERT(P >= 0); - - GGML_ASSERT(nbq0 == sizeof(float)); - GGML_ASSERT(nbk0 == sizeof(float)); - GGML_ASSERT(nbv0 == sizeof(float)); - - GGML_ASSERT(neq0 == D); - GGML_ASSERT(nek0 == D); - GGML_ASSERT(nev1 == D); - GGML_ASSERT(ned0 == D); - - GGML_ASSERT(neq1 == N); - GGML_ASSERT(nek1 == N + P); - GGML_ASSERT(nev1 == D); - GGML_ASSERT(ned1 == N); - - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); + // per thread: a row's scores, then one K row, one V row and one query row as F32 + float * s = (float *) params->wdata + ith*(M + D + Dv + D + Dv); + float * kr = s + M; + float * vr = kr + D; + float * qr = vr + Dv; + float * dor = qr + D; + + const ggml_to_float_t k_to_float = ggml_get_type_traits(k->type)->to_float; + const ggml_to_float_t v_to_float = ggml_get_type_traits(v->type)->to_float; + auto k_row = [&](int64_t j, int64_t hk) { + const char * r = (const char *) k->data + j*k->nb[1] + hk*k->nb[2]; + if (k->type == GGML_TYPE_F32) { memcpy(kr, r, D*sizeof(float)); } else { k_to_float(r, kr, D); } + }; + auto v_row = [&](int64_t j, int64_t hk) { + const char * r = (const char *) v->data + j*v->nb[1] + hk*v->nb[2]; + if (v->type == GGML_TYPE_F32) { memcpy(vr, r, Dv*sizeof(float)); } else { v_to_float(r, vr, Dv); } + }; + auto q_row = [&](int64_t i, int64_t h) { + for (int64_t c = 0; c < D; ++c) { + qr[c] = *(const float *) ((const char *) q->data + c*q->nb[0] + i*q->nb[1] + h*q->nb[2]); + } + }; + auto d_row = [&](int64_t i, int64_t h) { + for (int64_t c = 0; c < Dv; ++c) { + dor[c] = *(const float *) ((const char *) d->data + c*d->nb[0] + h*d->nb[1] + i*d->nb[2]); + } + }; + auto dot = [](const float * a, const float * b, int64_t n) { + ggml_float r = 0.0; + for (int64_t c = 0; c < n; ++c) { + r += (ggml_float) a[c]*b[c]; + } + return (float) r; + }; - if (ith == 0) { - memset(dst->data, 0, nb0*ne0*ne1*ne2*ne3); + // pass 1: the query rows + for (int64_t r = ith; r < N*H; r += nth) { + const int64_t i = r % N; + const int64_t h = r / N; + const int64_t hk = h / G; + q_row(i, h); + d_row(i, h); + float smax = -INFINITY; + for (int64_t j = 0; j < M; ++j) { + k_row(j, hk); + s[j] = dot(qr, kr, D)*scale + ggml_fa_back_mask(mask, j, i, h); + smax = std::max(smax, s[j]); + } + float l = -INFINITY; + if (smax != -INFINITY) { + ggml_float sum = 0.0; + for (int64_t j = 0; j < M; ++j) { + sum += expf(s[j] - smax); + } + l = smax + (float) log(sum); + } + // s becomes p; D from the recomputed p and dp + ggml_float dd = 0.0; + for (int64_t j = 0; j < M; ++j) { + s[j] = l == -INFINITY ? 0.0f : expf(s[j] - l); + if (s[j] != 0.0f) { + v_row(j, hk); + dd += (ggml_float) s[j]*dot(dor, vr, Dv); + } + } + lse [h*N + i] = l; + dvec[h*N + i] = (float) dd; + float * g = dq + (h*N + i)*D; + memset(g, 0, D*sizeof(float)); + for (int64_t j = 0; j < M; ++j) { + if (s[j] == 0.0f) { + continue; + } + v_row(j, hk); + k_row(j, hk); + const float ds = s[j]*(dot(dor, vr, Dv) - (float) dd)*scale; + for (int64_t c = 0; c < D; ++c) { + g[c] += ds*kr[c]; + } + } } - ggml_barrier(params->threadpool); - - const int64_t elem_q = ggml_nelements(q); - const int64_t elem_k = ggml_nelements(k); - - ggml_type result_type = dst->type; - GGML_ASSERT(ggml_blck_size(result_type) == 1); - const size_t tsize = ggml_type_size(result_type); - - const size_t offs_q = 0; - const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); - const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); - - void * grad_q = (char *) dst->data; - void * grad_k = (char *) dst->data + offs_k; - void * grad_v = (char *) dst->data + offs_v; - - const size_t nbgq1 = nb0*neq0; - const size_t nbgq2 = nb0*neq0*neq1; - const size_t nbgq3 = nb0*neq0*neq1*neq2; - - const size_t nbgk1 = nb0*nek0; - const size_t nbgk2 = nb0*nek0*nek1; - const size_t nbgk3 = nb0*nek0*nek1*neq2; - - const size_t nbgv1 = nb0*nev0; - const size_t nbgv2 = nb0*nev0*nev1; - const size_t nbgv3 = nb0*nev0*nev1*neq2; - - // parallelize by k rows using ggml_vec_dot_f32 - // total rows in k - const int nr = nek2*nek3; - - // rows per thread - const int dr = (nr + nth - 1)/nth; - - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); - - const float scale = 1.0f/sqrtf(D); - - //printf("P=%d N=%d D=%d ir0=%d ir1=%d scale = %f\n", P, N, D, ir0, ir1, scale); - - // how often k2 (and v2) is repeated in q2 - int nrep = neq2/nek2; - - for (int ir = ir0; ir < ir1; ++ir) { - // q indices - const int ik3 = ir/(nek2); - const int ik2 = ir - ik3*nek2; - - const int iq3 = ik3; - const int id3 = ik3; - const int iv3 = ik3; - const int iv2 = ik2; - - for (int irep = 0; irep < nrep; ++irep) { - const int iq2 = ik2 + irep*nek2; - const int id2 = iq2; - - // (ik2 + irep*nek2) % nek2 == ik2 - for (int iq1 = 0; iq1 < neq1; ++iq1) { - const int id1 = iq1; - - // not sure about CACHE_LINE_SIZE_F32.. - // - maybe it must not be multiplied by 2 and excluded from .. in SM 1*(..) offset? - float * S = (float *) params->wdata + ith*2*(mxDM + CACHE_LINE_SIZE_F32) + 0*(mxDM+CACHE_LINE_SIZE_F32); - float * SM = (float *) params->wdata + ith*2*(mxDM + CACHE_LINE_SIZE_F32) + 1*(mxDM+CACHE_LINE_SIZE_F32); + ggml_barrier(params->threadpool); - for (int i = M; i < Mup; ++i) { - S[i] = -INFINITY; + // pass 2: the KV positions, each summed over its group's query heads and every query row + for (int64_t r = ith; r < M*Hkv; r += nth) { + const int64_t j = r % M; + const int64_t hk = r / M; + float * gk = dk + (hk*M + j)*D; + float * gv = dv + (hk*M + j)*Dv; + memset(gk, 0, D*sizeof(float)); + memset(gv, 0, Dv*sizeof(float)); + k_row(j, hk); + v_row(j, hk); + for (int64_t h = hk*G; h < (hk + 1)*G; ++h) { + for (int64_t i = 0; i < N; ++i) { + const float l = lse[h*N + i]; + if (l == -INFINITY) { + continue; } - - const int64_t masked_begin = masked ? (P + iq1 + 1) : M; - for (int64_t ic = 0; ic < masked_begin; ++ic) { - // k indices - const int ik1 = ic; - - // S indices - const int i1 = ik1; - - ggml_vec_dot_f32(neq0, - S + i1, 0, - (float *) ((char *) k->data + (ik1*nbk1 + ik2*nbk2 + ik3*nbk3)), 0, - (float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)), 0, 1); + const float m = ggml_fa_back_mask(mask, j, i, h); + if (m == -INFINITY) { + continue; } - - // scale - ggml_vec_scale_f32(masked_begin, S, scale); - - for (int64_t i = masked_begin; i < M; i++) { - S[i] = -INFINITY; + q_row(i, h); + const float p = expf(dot(qr, kr, D)*scale + m - l); + if (p == 0.0f) { + continue; } - - // softmax - // exclude known -INF S[..] values from max and loop - // dont forget to set their SM values to zero - { - float max = -INFINITY; - ggml_vec_max_f32(masked_begin, &max, S); - - ggml_float sum = 0.0; - { -#ifdef GGML_SOFT_MAX_ACCELERATE - max = -max; - vDSP_vsadd(SM, 1, &max, SM, 1, Mup); - vvexpf(SM, SM, &Mup); - ggml_vec_sum_f32(Mup, &sum, SM); -#else - sum = ggml_vec_soft_max_f32(Mup, SM, S, max); -#endif - } - - assert(sum > 0.0); - - sum = 1.0/sum; - ggml_vec_scale_f32(masked_begin, SM, sum); - + d_row(i, h); + const float ds = p*(dot(dor, vr, Dv) - dvec[h*N + i])*scale; + for (int64_t c = 0; c < D; ++c) { + gk[c] += ds*qr[c]; } - - // step-by-step explanation - { - // forward-process shape grads from backward process - // parallel_for ik2,ik3: - // for irep: - // iq2 = ik2 + irep*nek2 - // k[:D,:M,:,:] [D,M,:,:] grad[k][:D,:M,ik2,ik3] += grad[kcur] - // q[:D,:N,:,:] [D,N,:,:] grad[q][:D,iq1,iq2,iq3] += grad[qcur] - // v[:M,:D,:,:] [M,D,:,:] grad[v][:M,:D,iv2,iv3] += grad[vcur] - // for iq1: - // kcur = k[:D,:M,ik2,ik3] [D,M,1,1] grad[kcur] = grad[S1].T @ qcur - // qcur = q[:D,iq1,iq2,iq3] [D,1,1,1] grad[qcur] = grad[S1] @ kcur - // vcur = v[:M,:D,iv2,iv3] [M,D,1,1] grad[vcur] = grad[S5].T @ S4 - // S0 = -Inf [D,1,1,1] - // ~S1[i] = dot(kcur[:D,i], qcur) - // S1 = qcur @ kcur.T [M,1,1,1] grad[S1] = grad[S2] * scale - // S2 = S1 * scale [M,1,1,1] grad[S2] = diag_mask_zero(grad[S3], P) - // S3 = diag_mask_inf(S2, P) [M,1,1,1] grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // S4 = softmax(S3) [M,1,1,1] grad[S4] = grad[S5] @ vcur - // ~S5[i] = dot(vcur[:,i], S4) - // S5 = S4 @ vcur.T [D,1,1,1] grad[S5] = d[:D,id1,id2,id3] - // ~dst[i,iq1,iq2,iq3] = S5[i] ^ - // dst[:D,iq1,iq2,iq3] = S5 | grad[dst[:D,iq1,iq2,iq3]] = d[:D,id1,id2,id3] - // dst backward-/ grad[dst] = d - // - // output gradients with their dependencies: - // - // grad[kcur] = grad[S1].T @ qcur - // grad[S1] = diag_mask_zero(grad[S3], P) * scale - // grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // grad[S4] = grad[S5] @ vcur - // grad[S4] = d[:D,id1,id2,id3] @ vcur - // grad[qcur] = grad[S1] @ kcur - // grad[vcur] = grad[S5].T @ S4 - // grad[vcur] = d[:D,id1,id2,id3].T @ S4 - // - // in post-order: - // - // S1 = qcur @ kcur.T - // S2 = S1 * scale - // S3 = diag_mask_inf(S2, P) - // S4 = softmax(S3) - // grad[S4] = d[:D,id1,id2,id3] @ vcur - // grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // grad[S1] = diag_mask_zero(grad[S3], P) * scale - // grad[qcur] = grad[S1] @ kcur - // grad[kcur] = grad[S1].T @ qcur - // grad[vcur] = d[:D,id1,id2,id3].T @ S4 - // - // using less variables (SM=S4): - // - // S = diag_mask_inf(qcur @ kcur.T * scale, P) - // SM = softmax(S) - // S = d[:D,iq1,iq2,iq3] @ vcur - // dot_SM_gradSM = dot(SM, S) - // S = SM * (S - dot(SM, S)) - // S = diag_mask_zero(S, P) * scale - // - // grad[q][:D,iq1,iq2,iq3] += S @ kcur - // grad[k][:D,:M,ik2,ik3] += S.T @ qcur - // grad[v][:M,:D,iv2,iv3] += d[:D,id1,id2,id3].T @ SM - } - - // S = gradSM = d[:D,id1,id2,id3] @ vcur[:,:,iv2,iv3] - // S = d[:D,id1,id2,id3] @ vcur[:,:,iv2,iv3] - // for ic: - // S[:M] += vcur[:M,ic,iv2,iv3] * d[ic,id1,id2,id3] - // exclude known future zero S[..] values from operation - ggml_vec_set_f32(masked_begin, S, 0); - for (int64_t ic = 0; ic < D; ++ic) { - ggml_vec_mad_f32(masked_begin, - S, - (float *) ((char *) v->data + ( ic*nbv1 + iv2*nbv2 + iv3*nbv3)), - *(float *) ((char *) d->data + (ic*nbd0 + id1*nbd1 + id2*nbd2 + id3*nbd3))); - } - - // S = SM * (S - dot(SM, S)) - float dot_SM_gradSM = 0; - ggml_vec_dot_f32 (masked_begin, &dot_SM_gradSM, 0, SM, 0, S, 0, 1); - ggml_vec_acc1_f32(M, S, -dot_SM_gradSM); - ggml_vec_mul_f32 (masked_begin, S, S, SM); - - // S = diag_mask_zero(S, P) * scale - // already done by above ggml_vec_set_f32 - - // exclude known zero S[..] values from operation - ggml_vec_scale_f32(masked_begin, S, scale); - - // S shape [M,1] - // SM shape [M,1] - // kcur shape [D,M] - // qcur shape [D,1] - // vcur shape [M,D] - - // grad[q][:D,iq1,iq2,iq3] += S @ kcur - // grad[q][:D,iq1,iq2,iq3] += shape[M,1] @ shape[D,M] - // for ic: - // grad[q][:D,iq1,iq2,iq3] += S[ic] * kcur[:D,ic,ik2,ik3] - // exclude known zero S[..] values from loop - for (int64_t ic = 0; ic < masked_begin; ++ic) { - ggml_vec_mad_f32(D, - (float *) ((char *) grad_q + (iq1*nbgq1 + iq2*nbgq2 + iq3*nbgq3)), - (float *) ((char *) k->data + (ic*nbk1 + ik2*nbk2 + ik3*nbk3)), - S[ic]); - } - - // grad[k][:D,:M,iq2,iq3] += S.T @ qcur - // for ic: - // grad[k][:D,ic,iq2,iq3] += S.T[0,ic] * qcur[:D,0] - // grad[k][:D,ic,iq2,iq3] += S[ic] * qcur[:D,0] - // exclude known zero S[..] values from loop - for (int64_t ic = 0; ic < masked_begin; ++ic) { - ggml_vec_mad_f32(D, - (float *) ((char *) grad_k + (ic*nbgk1 + ik2*nbgk2 + ik3*nbgk3)), - (float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)), - S[ic]); - } - - // grad[v][:M,:D,iv2,iv3] += d[:D,id1,id2,id3].T @ SM - // for ic: - // grad[v][:M,ic,iv2,iv3] += d[:D,id1,id2,id3].T[0,ic] * SM[:M] - // grad[v][:M,ic,iv2,iv3] += d[ic,id1,id2,id3] * SM[:M] - // exclude known zero SM[..] values from mad - for (int64_t ic = 0; ic < D; ++ic) { - ggml_vec_mad_f32(masked_begin, - (float *) ((char *) grad_v + ( ic*nbgv1 + iv2*nbgv2 + iv3*nbgv3)), - SM, - *(float *) ((char *) d->data + (ic*nbd0 + id1*nbd1 + id2*nbd2 + id3*nbd3))); + for (int64_t c = 0; c < Dv; ++c) { + gv[c] += p*dor[c]; } } } @@ -9542,21 +9392,9 @@ static void ggml_compute_forward_flash_attn_back_f32( void ggml_compute_forward_flash_attn_back( const ggml_compute_params * params, - const bool masked, ggml_tensor * dst) { - - const ggml_tensor * q = dst->src[0]; - - switch (q->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_flash_attn_back_f32(params, masked, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } + GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32); + ggml_compute_forward_flash_attn_back_f32(params, dst); } // ggml_compute_forward_ssm_conv diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 1c8b66bfaaa0..07027df5102b 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -88,10 +88,7 @@ void ggml_compute_forward_leaky_relu(const struct ggml_compute_params * params, void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_flash_attn_back( - const struct ggml_compute_params * params, - const bool masked, - struct ggml_tensor * dst); +void ggml_compute_forward_flash_attn_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_ssm_conv(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_ssm_scan(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_win_part(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 08b2ef272896..7059261b2716 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -353,6 +353,20 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_lo return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_back(ggml_metal_library_t lib, const ggml_tensor * op, bool kv_pass) { + assert(op->op == GGML_OP_FLASH_ATTN_BACK); + + char name[256]; + snprintf(name, 256, "kernel_flash_attn_back_%s_%s", kv_pass ? "kv" : "q", ggml_type_name(op->src[1]->type)); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); + } + + return res; +} + // RMS_NORM_BACK, SOFT_MAX_BACK, REPEAT_BACK: one f32 kernel each, named after the op. ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back(ggml_metal_library_t lib, const ggml_tensor * op) { const char * kernel = nullptr; diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 18ae227a5ca9..c82db0dd521e 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -121,6 +121,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_loss(ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_back(ggml_metal_library_t lib, const struct ggml_tensor * op, bool kv_pass); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_blk (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 0ff9e84c17ff..b31d6df4da24 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1477,6 +1477,15 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te (op->op != GGML_OP_SOFT_MAX_BACK || ggml_get_op_params_f32(op, 1) == 0.0f); case GGML_OP_REPEAT_BACK: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op); + case GGML_OP_FLASH_ATTN_BACK: + // one simdgroup per row: head sizes a multiple of 32, at most 512 (FA_BACK_MAX_D) + return op->src[0]->type == GGML_TYPE_F32 && op->src[4]->type == GGML_TYPE_F32 && + op->src[1]->type == op->src[2]->type && + (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_Q8_0) && + op->src[0]->ne[0] % 32 == 0 && op->src[2]->ne[0] % 32 == 0 && + op->src[0]->ne[0] <= 512 && op->src[2]->ne[0] <= 512 && + op->src[0]->nb[0] == sizeof(float) && op->src[4]->nb[0] == sizeof(float) && + (op->src[3] == NULL || op->src[3]->type == GGML_TYPE_F16); case GGML_OP_CROSS_ENTROPY_LOSS: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_I32) && diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index adc883e20f1d..078b3c50f282 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -888,6 +888,33 @@ typedef struct { uint64_t nb1; } ggml_metal_kargs_cross_entropy_loss; +// FLASH_ATTN_BACK (ggml_flash_attn_back): q [D, N, H], k [D, M, Hkv], v [Dv, M, Hkv], d [Dv, H, N], +// mask [M', N', ne32] F16; the result's parts at byte offsets offs_dk / offs_dv / offs_lse +typedef struct { + int32_t D; + int32_t Dv; + int32_t N; + int32_t H; + int32_t M; + int32_t Hkv; + uint64_t nbq1; + uint64_t nbq2; + uint64_t nbk1; + uint64_t nbk2; + uint64_t nbv1; + uint64_t nbv2; + uint64_t nbd1; + uint64_t nbd2; + uint64_t nbm1; + uint64_t nbm2; + int32_t nem2; + int32_t has_mask; + float scale; + uint64_t offs_dk; + uint64_t offs_dv; + uint64_t offs_lse; +} ggml_metal_kargs_flash_attn_back; + // a per-row backward with one scalar parameter (eps for RMS_NORM_BACK, scale for SOFT_MAX_BACK) typedef struct { int32_t ne00; diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index cb2e266e12c9..fbf84f1fe2ea 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -312,6 +312,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_cross_entropy_loss_back(ctx, idx); } break; + case GGML_OP_FLASH_ATTN_BACK: + { + n_fuse = ggml_metal_op_flash_attn_back(ctx, idx); + } break; case GGML_OP_RMS_NORM_BACK: case GGML_OP_SOFT_MAX_BACK: { @@ -4235,6 +4239,63 @@ int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx) { return 1; } +// FLASH_ATTN_BACK: the Q pass (dq, and each row's lse and D into the scratch after dv), then the +// KV pass, which reads them +int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + const ggml_tensor * q = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * v = op->src[2]; + const ggml_tensor * mask = op->src[3]; + const ggml_tensor * d = op->src[4]; + + ggml_metal_kargs_flash_attn_back args = { + /*.D =*/ (int32_t) q->ne[0], + /*.Dv =*/ (int32_t) v->ne[0], + /*.N =*/ (int32_t) q->ne[1], + /*.H =*/ (int32_t) q->ne[2], + /*.M =*/ (int32_t) k->ne[1], + /*.Hkv =*/ (int32_t) k->ne[2], + /*.nbq1 =*/ q->nb[1], + /*.nbq2 =*/ q->nb[2], + /*.nbk1 =*/ k->nb[1], + /*.nbk2 =*/ k->nb[2], + /*.nbv1 =*/ v->nb[1], + /*.nbv2 =*/ v->nb[2], + /*.nbd1 =*/ d->nb[1], + /*.nbd2 =*/ d->nb[2], + /*.nbm1 =*/ mask ? mask->nb[1] : 0, + /*.nbm2 =*/ mask ? mask->nb[2] : 0, + /*.nem2 =*/ mask ? (int32_t) mask->ne[2] : 1, + /*.has_mask =*/ mask ? 1 : 0, + /*.scale =*/ ggml_get_op_params_f32(op, 0), + /*.offs_dk =*/ ggml_flash_attn_back_offset(op, 1), + /*.offs_dv =*/ ggml_flash_attn_back_offset(op, 2), + /*.offs_lse =*/ ggml_flash_attn_back_offset(op, 3), + }; + + for (int pass = 0; pass < 2; ++pass) { + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_back(lib, op, pass == 1); + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(q), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(k), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(v), 3); + ggml_metal_encoder_set_buffer (enc, mask ? ggml_metal_get_buffer_id(mask) : ggml_metal_get_buffer_id(q), 4); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(d), 5); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 6); + const int ngroups = pass == 0 ? args.N*args.H : args.M*args.Hkv; + ggml_metal_encoder_dispatch_threadgroups(enc, ngroups, 1, 1, 32, 1, 1); + ggml_metal_op_concurrency_reset(ctx); + } + + return 1; +} + int ggml_metal_op_cross_entropy_loss_back(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); return ggml_metal_op_cross_entropy_loss_rows(ctx, op, ggml_metal_get_buffer_id(op), op->nb[1]); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index e76003f71bca..e2cfd4f6cc01 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -81,6 +81,7 @@ bool ggml_metal_op_out_prod_supported(const struct ggml_tensor * op); int ggml_metal_op_out_prod (ggml_metal_op_t ctx, int idx); int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx); int ggml_metal_op_cross_entropy_loss_back(ggml_metal_op_t ctx, int idx); +int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx); int ggml_metal_op_row_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_repeat_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_silu_back (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index 6faf923c425e..3cc9ae274113 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -2254,3 +2254,193 @@ template [[host_name("kernel_lightning_indexer_q4_1")]] kernel kernel_lightning_ template [[host_name("kernel_lightning_indexer_q5_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer; template [[host_name("kernel_lightning_indexer_q5_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer; template [[host_name("kernel_lightning_indexer_q8_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer; + +// FLASH_ATTN_BACK: the gradients of FLASH_ATTN_EXT, recomputing the scores instead of keeping +// them, so the memory is the gradients alone (the CPU reference: ggml-cpu ops.cpp +// ggml_compute_forward_flash_attn_back_f32). One simdgroup per row; each lane holds every 32nd +// element of a row. Pass Q, per query row and head: the row's log-sum-exp over the window, then +// in one more sweep D = sum_j p_j dp_j, A = sum_j p_j dp_j k_j and B = sum_j p_j k_j, so +// dq = scale (A - D B). Pass KV, per KV position and KV head: dk and dv summed over the query +// rows of every head in its group, from the lse and D the Q pass wrote. + +#define FA_BACK_MAX_D 512 + +template +static inline float fa_back_elem(device const char * row, int c); + +template <> inline float fa_back_elem(device const char * row, int c) { return ((device const float *) row)[c]; } +template <> inline float fa_back_elem (device const char * row, int c) { return (float) ((device const half *) row)[c]; } +template <> inline float fa_back_elem(device const char * row, int c) { + device const block_q8_0 * b = (device const block_q8_0 *) row + c/QK8_0; + return (float) b->d * (float) b->qs[c%QK8_0]; +} + +static inline float fa_back_mask(constant ggml_metal_kargs_flash_attn_back & args, device const char * mask, int j, int i, int h) { + if (!args.has_mask) { + return 0.0f; + } + return (float) ((device const half *) (mask + i*args.nbm1 + (h % args.nem2)*args.nbm2))[j]; +} + +template +kernel void kernel_flash_attn_back_q( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]]) { + const int i = tgpig % args.N; + const int h = tgpig / args.N; + const int hk = h / (args.H / args.Hkv); + + const int nq = args.D / 32; + const int nv = args.Dv / 32; + float qr[FA_BACK_MAX_D/32], dr[FA_BACK_MAX_D/32], a[FA_BACK_MAX_D/32], b[FA_BACK_MAX_D/32]; + for (int c = 0; c < nq; ++c) { + qr[c] = *(device const float *) (q + (tiisg + 32*c)*sizeof(float) + i*args.nbq1 + h*args.nbq2); + a[c] = 0.0f; + b[c] = 0.0f; + } + for (int c = 0; c < nv; ++c) { + dr[c] = *(device const float *) (d + (tiisg + 32*c)*sizeof(float) + h*args.nbd1 + i*args.nbd2); + } + + // the log-sum-exp, online + float m = -INFINITY; + float l = 0.0f; + for (int j = 0; j < args.M; ++j) { + const float mj = fa_back_mask(args, mask, j, i, h); + if (mj == -INFINITY) { + continue; + } + device const char * kr = k + j*args.nbk1 + hk*args.nbk2; + float s = 0.0f; + for (int c = 0; c < nq; ++c) { + s += qr[c]*fa_back_elem(kr, tiisg + 32*c); + } + s = simd_sum(s)*args.scale + mj; + const float mn = max(m, s); + l = l*exp(m - mn) + exp(s - mn); + m = mn; + } + const float lse = m == -INFINITY ? -INFINITY : m + log(l); + + // D, A and B in one sweep + float dd = 0.0f; + if (lse != -INFINITY) { + for (int j = 0; j < args.M; ++j) { + const float mj = fa_back_mask(args, mask, j, i, h); + if (mj == -INFINITY) { + continue; + } + device const char * kr = k + j*args.nbk1 + hk*args.nbk2; + device const char * vr = v + j*args.nbv1 + hk*args.nbv2; + float s = 0.0f; + for (int c = 0; c < nq; ++c) { + s += qr[c]*fa_back_elem(kr, tiisg + 32*c); + } + float dp = 0.0f; + for (int c = 0; c < nv; ++c) { + dp += dr[c]*fa_back_elem(vr, tiisg + 32*c); + } + const float p = exp(simd_sum(s)*args.scale + mj - lse); + dp = simd_sum(dp); + dd += p*dp; + for (int c = 0; c < nq; ++c) { + const float kc = fa_back_elem(kr, tiisg + 32*c); + a[c] += p*dp*kc; + b[c] += p*kc; + } + } + } + + device float * dq = (device float *) dst + (h*args.N + i)*args.D; + for (int c = 0; c < nq; ++c) { + dq[tiisg + 32*c] = args.scale*(a[c] - dd*b[c]); + } + if (tiisg == 0) { + device float * lsev = (device float *) (dst + args.offs_lse); + lsev[h*args.N + i] = lse; + lsev[args.N*args.H + h*args.N + i] = dd; + } +} + +template +kernel void kernel_flash_attn_back_kv( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]]) { + const int j = tgpig % args.M; + const int hk = tgpig / args.M; + const int G = args.H / args.Hkv; + + const int nq = args.D / 32; + const int nv = args.Dv / 32; + float kr[FA_BACK_MAX_D/32], vr[FA_BACK_MAX_D/32], gk[FA_BACK_MAX_D/32], gv[FA_BACK_MAX_D/32]; + device const char * kp = k + j*args.nbk1 + hk*args.nbk2; + device const char * vp = v + j*args.nbv1 + hk*args.nbv2; + for (int c = 0; c < nq; ++c) { + kr[c] = fa_back_elem(kp, tiisg + 32*c); + gk[c] = 0.0f; + } + for (int c = 0; c < nv; ++c) { + vr[c] = fa_back_elem(vp, tiisg + 32*c); + gv[c] = 0.0f; + } + device const float * lsev = (device const float *) (dst + args.offs_lse); + for (int h = hk*G; h < (hk + 1)*G; ++h) { + for (int i = 0; i < args.N; ++i) { + const float lse = lsev[h*args.N + i]; + const float mj = fa_back_mask(args, mask, j, i, h); + if (lse == -INFINITY || mj == -INFINITY) { + continue; + } + device const char * qp = q + i*args.nbq1 + h*args.nbq2; + device const char * dp = d + h*args.nbd1 + i*args.nbd2; + float s = 0.0f; + float t = 0.0f; + for (int c = 0; c < nq; ++c) { + s += ((device const float *) qp)[tiisg + 32*c]*kr[c]; + } + for (int c = 0; c < nv; ++c) { + t += ((device const float *) dp)[tiisg + 32*c]*vr[c]; + } + const float p = exp(simd_sum(s)*args.scale + mj - lse); + const float ds = p*(simd_sum(t) - lsev[args.N*args.H + h*args.N + i])*args.scale; + for (int c = 0; c < nq; ++c) { + gk[c] += ds*((device const float *) qp)[tiisg + 32*c]; + } + for (int c = 0; c < nv; ++c) { + gv[c] += p*((device const float *) dp)[tiisg + 32*c]; + } + } + } + device float * dk = (device float *) (dst + args.offs_dk) + (hk*args.M + j)*args.D; + device float * dv = (device float *) (dst + args.offs_dv) + (hk*args.M + j)*args.Dv; + for (int c = 0; c < nq; ++c) { + dk[tiisg + 32*c] = gk[c]; + } + for (int c = 0; c < nv; ++c) { + dv[tiisg + 32*c] = gv[c]; + } +} + +typedef decltype(kernel_flash_attn_back_q) kernel_flash_attn_back_q_t; +typedef decltype(kernel_flash_attn_back_kv) kernel_flash_attn_back_kv_t; + +template [[host_name("kernel_flash_attn_back_q_f32")]] kernel kernel_flash_attn_back_q_t kernel_flash_attn_back_q; +template [[host_name("kernel_flash_attn_back_q_f16")]] kernel kernel_flash_attn_back_q_t kernel_flash_attn_back_q; +template [[host_name("kernel_flash_attn_back_q_q8_0")]] kernel kernel_flash_attn_back_q_t kernel_flash_attn_back_q; +template [[host_name("kernel_flash_attn_back_kv_f32")]] kernel kernel_flash_attn_back_kv_t kernel_flash_attn_back_kv; +template [[host_name("kernel_flash_attn_back_kv_f16")]] kernel kernel_flash_attn_back_kv_t kernel_flash_attn_back_kv; +template [[host_name("kernel_flash_attn_back_kv_q8_0")]] kernel kernel_flash_attn_back_kv_t kernel_flash_attn_back_kv; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 0b58d97094fe..210aad5531f4 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5542,71 +5542,52 @@ void ggml_flash_attn_ext_add_sinks( // ggml_flash_attn_back +size_t ggml_flash_attn_back_offset(const struct ggml_tensor * back, int part) { + const struct ggml_tensor * q = back->src[0]; + const struct ggml_tensor * k = back->src[1]; + const struct ggml_tensor * v = back->src[2]; + const size_t sizes[3] = { + GGML_PAD((size_t) q->ne[0]*q->ne[1]*q->ne[2]*sizeof(float), GGML_MEM_ALIGN), + GGML_PAD((size_t) ggml_nelements(k)*sizeof(float), GGML_MEM_ALIGN), + GGML_PAD((size_t) ggml_nelements(v)*sizeof(float), GGML_MEM_ALIGN), + }; + size_t offs = 0; + for (int i = 0; i < part; ++i) { + offs += sizes[i]; + } + return offs; +} + struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, struct ggml_tensor * v, + struct ggml_tensor * mask, struct ggml_tensor * d, - bool masked) { - GGML_ABORT("TODO: adapt to ggml_flash_attn_ext() changes"); - - GGML_ASSERT(ggml_can_mul_mat(k, q)); - // TODO: check if vT can be multiplied by (k*qT) - - // d shape [D,N,ne2,ne3] - // q shape [D,N,ne2,ne3] - // k shape [D,M,kvne2,ne3] - // v shape [M,D,kvne2,ne3] - - const int64_t D = q->ne[0]; - const int64_t N = q->ne[1]; - const int64_t M = k->ne[1]; - const int64_t ne2 = q->ne[2]; - const int64_t ne3 = q->ne[3]; - const int64_t kvne2 = k->ne[2]; - - GGML_ASSERT(k->ne[0] == D); - GGML_ASSERT(v->ne[0] == M); - GGML_ASSERT(v->ne[1] == D); - GGML_ASSERT(d->ne[0] == D); - GGML_ASSERT(d->ne[1] == N); - GGML_ASSERT(k->ne[2] == kvne2); - GGML_ASSERT(k->ne[3] == ne3); - GGML_ASSERT(v->ne[2] == kvne2); - GGML_ASSERT(v->ne[3] == ne3); - GGML_ASSERT(d->ne[2] == ne2); - GGML_ASSERT(d->ne[3] == ne3); - - GGML_ASSERT(ne2 % kvne2 == 0); - - // store gradients of q, k and v as continuous tensors concatenated in result. - // note: v and gradv are actually transposed, i.e. v->ne[0] != D. - const int64_t elem_q = ggml_nelements(q); - const int64_t elem_k = ggml_nelements(k); - const int64_t elem_v = ggml_nelements(v); - - enum ggml_type result_type = GGML_TYPE_F32; - GGML_ASSERT(ggml_blck_size(result_type) == 1); - const size_t tsize = ggml_type_size(result_type); - - const size_t offs_q = 0; - const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); - const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); - const size_t end = offs_v + GGML_PAD(elem_v * tsize, GGML_MEM_ALIGN); - - const size_t nelements = (end + tsize - 1)/tsize; - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nelements); - - int32_t masked_i = masked ? 1 : 0; - ggml_set_op_params(result, &masked_i, sizeof(masked_i)); + float scale) { + // q [D, N, H, 1], k [D, M, Hkv, 1], v [Dv, M, Hkv, 1], d [Dv, H, N, 1], mask [M', N', 1|H, 1] + GGML_ASSERT(q->type == GGML_TYPE_F32 && d->type == GGML_TYPE_F32); + GGML_ASSERT(q->ne[3] == 1 && k->ne[3] == 1 && v->ne[3] == 1 && "one sequence stream"); + GGML_ASSERT(k->ne[0] == q->ne[0] && v->ne[1] == k->ne[1] && v->ne[2] == k->ne[2]); + GGML_ASSERT(q->ne[2] % k->ne[2] == 0); + GGML_ASSERT(d->ne[0] == v->ne[0] && d->ne[1] == q->ne[2] && d->ne[2] == q->ne[1]); + GGML_ASSERT(!mask || (mask->type == GGML_TYPE_F16 && mask->ne[0] >= k->ne[1] && mask->ne[1] >= q->ne[1])); + + const size_t rows = (size_t) q->ne[1]*q->ne[2]; + struct ggml_tensor tmp = { 0 }; + tmp.src[0] = q; tmp.src[1] = k; tmp.src[2] = v; + const size_t bytes = ggml_flash_attn_back_offset(&tmp, 3) + 2*rows*sizeof(float); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) (bytes/sizeof(float))); + ggml_set_op_params_f32(result, 0, scale); result->op = GGML_OP_FLASH_ATTN_BACK; result->src[0] = q; result->src[1] = k; result->src[2] = v; - result->src[3] = d; + result->src[3] = mask; + result->src[4] = d; return result; } @@ -6826,6 +6807,26 @@ static void ggml_flash_attn_ext_backward( GGML_ASSERT(q->ne[3] == 1 && "flash_attn_ext backward: one sequence stream"); int32_t tile = ggml_get_op_params_i32(tensor, 4); + if (tile == GGML_FA_BACK_NATIVE) { + struct ggml_tensor * b = ggml_flash_attn_back(ctx, q, k, v, mask, grad, scale); + // views of the packed result, each a contiguous tensor shaped like its source + struct ggml_tensor * dq = ggml_view_3d(ctx, b, q->ne[0], q->ne[1], q->ne[2], q->ne[0]*sizeof(float), q->ne[0]*q->ne[1]*sizeof(float), + ggml_flash_attn_back_offset(b, 0)); + struct ggml_tensor * dk = ggml_view_4d(ctx, b, k->ne[0], k->ne[1], k->ne[2], k->ne[3], k->ne[0]*sizeof(float), + k->ne[0]*k->ne[1]*sizeof(float), k->ne[0]*k->ne[1]*k->ne[2]*sizeof(float), ggml_flash_attn_back_offset(b, 1)); + struct ggml_tensor * dv = ggml_view_4d(ctx, b, v->ne[0], v->ne[1], v->ne[2], v->ne[3], v->ne[0]*sizeof(float), + v->ne[0]*v->ne[1]*sizeof(float), v->ne[0]*v->ne[1]*v->ne[2]*sizeof(float), ggml_flash_attn_back_offset(b, 2)); + if (dq_needed) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape_4d(ctx, ggml_cont(ctx, dq), q->ne[0], q->ne[1], q->ne[2], q->ne[3])); + } + if (dk_needed) { + ggml_add_or_set(ctx, cgraph, isrc1, dk); + } + if (dv_needed) { + ggml_add_or_set(ctx, cgraph, isrc2, dv); + } + return; + } if (tile <= 0) { tile = GGML_FA_BACK_TILE_DEFAULT; } diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 4891ca771e08..f65abef27e0c 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -40,6 +40,44 @@ static bool llama_set_param(struct ggml_tensor * tensor, llama_opt_param_filter return true; } +// whether every GPU this context computes on has a kernel for ggml_flash_attn_back, at this +// model's attention shape and this context's cache types (the CPU has one) +bool llama_context::opt_fa_back_native() const { + const llama_hparams & hp = model.hparams; + uint32_t il = 0; + while (il < hp.n_layer() && hp.is_recr(il)) { + ++il; + } + if (il == hp.n_layer()) { + return false; // no attention layer + } + // the cache the step's attention reads in place + const llama_kv_cache * kv = dynamic_cast(memory.get()); + if (const auto * hybrid = dynamic_cast(memory.get())) { + kv = hybrid->get_mem_attn(); + } + if (kv == nullptr) { + return false; + } + ggml_init_params ip = { 8*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true }; + ggml_context_ptr ctx(ggml_init(ip)); + const int64_t H = hp.n_head(il), Hkv = hp.n_head_kv(il), D = hp.n_embd_head_k(il), Dv = hp.n_embd_head_v(il); + ggml_tensor * q = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, D, 8, H); + ggml_tensor * k = ggml_new_tensor_3d(ctx.get(), kv->type_k(), D, 64, Hkv); + ggml_tensor * v = ggml_new_tensor_3d(ctx.get(), kv->type_v(), Dv, 64, Hkv); + ggml_tensor * m = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F16, 64, 8); + ggml_tensor * d = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, Dv, H, 8); + ggml_tensor * op = ggml_flash_attn_back(ctx.get(), q, k, v, m, d, 1.0f); + for (const auto & backend : backends) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend.get()); + const auto type = dev ? ggml_backend_dev_type(dev) : GGML_BACKEND_DEVICE_TYPE_CPU; + if ((type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU) && !ggml_backend_dev_supports_op(dev, op)) { + return false; + } + } + return true; +} + void llama_context::opt_init(struct llama_model * model, struct llama_opt_params lopt_params) { GGML_ASSERT(!opt_ctx); // The training window lives on this context, not the shared model (a serving context may @@ -76,6 +114,10 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params max_nodes, cparams.pipeline_parallel, cparams.op_offload)); } + cparams.fa_back_native = opt_fa_back_native(); + LLAMA_LOG_INFO("%s: the attention backward is %s\n", __func__, + cparams.fa_back_native ? "one op (ggml_flash_attn_back): no score is kept" : "composed from tiles (a device has no ggml_flash_attn_back kernel)"); + // one target token per row: no one-hot tensor as wide as the vocabulary beside the logits ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS); opt_params.opt_period = n_batch / n_ubatch; diff --git a/src/llama-context.h b/src/llama-context.h index 56dd3d34156f..88bd187b9145 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -249,6 +249,7 @@ struct llama_context { // plus the loss over the labelled rows, the backward and, at a period's end, the optimizer // step. A walk driving the decode adds its taps to that graph. bool opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch); + bool opt_fa_back_native() const; // // output diff --git a/src/llama-cparams.h b/src/llama-cparams.h index a5f4c57cd740..8e798ef85e32 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -22,6 +22,9 @@ struct llama_cparams { // the run's positions] F32 in host memory (null for a recurrent layer), row 0 at position // walk_acc_base; a chunk's backward adds its prefix's share into rows [walk_grad_from, n_past) // and its surrogate reads its own rows. + // the attention backward is one op (ggml_flash_attn_back) on every device that would run it + // (llama_context::opt_init probes them), else the composed tiles + bool fa_back_native = false; bool walk_exact = false; uint32_t walk_grad_from = 0; bool walk_surrogate = false; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 38f52b951e78..5d47d281ecae 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2666,7 +2666,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_flash_attn_ext_add_sinks(cur, sinks); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); if (cparams.training) { - ggml_flash_attn_ext_set_back_tile(cur, 0); // the backward's default tile + ggml_flash_attn_ext_set_back_tile(cur, cparams.fa_back_native ? GGML_FA_BACK_NATIVE : 0); } if (v_mla) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index fa80a8573970..0c622d6ec557 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7467,6 +7467,102 @@ struct test_flash_attn_ext : public test_case { } }; +// FLASH_ATTN_EXT's gradients through either backward (the native ggml_flash_attn_back or the +// composed tiles), against finite differences: GQA, a causal mask, q, k and v all trained +struct test_flash_attn_ext_grad : public test_case { + const int64_t hs, nh, nhkv, kv, nb; + const bool native; + + std::string vars() override { + return VARS_TO_STR6(hs, nh, nhkv, kv, nb, native); + } + + test_flash_attn_ext_grad(int64_t hs = 16, int64_t nh = 4, int64_t nhkv = 2, int64_t kv = 24, int64_t nb = 8, bool native = true) + : hs(hs), nh(nh), nhkv(nhkv), kv(kv), nb(nb), native(native) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, nb, nh); + ggml_tensor * k = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, kv, nhkv); + ggml_tensor * v = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, kv, nhkv); + ggml_set_param(q); ggml_set_name(q, "q"); + ggml_set_param(k); ggml_set_name(k, "k"); + ggml_set_param(v); ggml_set_name(v, "v"); + ggml_tensor * m = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kv, nb); + ggml_set_name(m, "m"); + ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf((float) hs), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + ggml_flash_attn_ext_set_back_tile(out, native ? GGML_FA_BACK_NATIVE : 0); + // weighted by a fixed random tensor: a plain sum of attention outputs nearly cancels the + // query gradient (each row's weights sum to one), and the finite difference drowns in it + ggml_tensor * w = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, out->ne[0], out->ne[1], out->ne[2]); + ggml_set_name(w, "w"); + out = ggml_mul(ctx, out, w); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "m") == 0) { + // causal: query row i sees KV positions j <= kv - nb + i + std::vector m(ggml_nelements(t)); + for (int64_t i = 0; i < t->ne[1]; ++i) { + for (int64_t j = 0; j < t->ne[0]; ++j) { + m[i*t->ne[0] + j] = ggml_fp32_to_fp16(j <= kv - nb + i ? 0.0f : -INFINITY); + } + } + ggml_backend_tensor_set(t, m.data(), 0, ggml_nbytes(t)); + } else { + init_tensor_uniform(t); + } + } + } + + bool grad_precise() override { + return true; + } +}; + +// GGML_OP_FLASH_ATTN_BACK itself: each backend's kernel against the CPU reference, at the cache +// types training reads (F16, F32, q8_0) and a GQA group +struct test_flash_attn_back : public test_case { + const int64_t hs, nh, nhkv, kv, nb; + const ggml_type type_kv; + + std::string vars() override { + return VARS_TO_STR6(hs, nh, nhkv, kv, nb, type_kv); + } + + test_flash_attn_back(int64_t hs = 128, int64_t nh = 8, int64_t nhkv = 2, int64_t kv = 512, int64_t nb = 64, ggml_type type_kv = GGML_TYPE_F16) + : hs(hs), nh(nh), nhkv(nhkv), kv(kv), nb(nb), type_kv(type_kv) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, nb, nh); + ggml_tensor * k = ggml_new_tensor_3d(ctx, type_kv, hs, kv, nhkv); + ggml_tensor * v = ggml_new_tensor_3d(ctx, type_kv, hs, kv, nhkv); + ggml_tensor * m = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kv, nb); + ggml_tensor * d = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, nh, nb); + ggml_set_name(q, "q"); ggml_set_name(k, "k"); ggml_set_name(v, "v"); ggml_set_name(m, "m"); ggml_set_name(d, "d"); + ggml_tensor * out = ggml_flash_attn_back(ctx, q, k, v, m, d, 1.0f/sqrtf((float) hs)); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "m") == 0) { + init_tensor_kq_mask(t); + } else { + init_tensor_uniform(t); + } + } + } + + double max_nmse_err() override { + return 1e-5; // the kernel and the reference sum in different orders over the window + } +}; + // GGML_OP_CROSS_ENTROPY_LOSS // Completion-only training masks rows: a masked position's label row is all zero. Every other // row sums to 1. "masked" zeroes every other row, so the finite-difference check sees whether the @@ -10391,6 +10487,12 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext_grad(16, 4, 2, 24, 8, true)); + test_cases.emplace_back(new test_flash_attn_ext_grad(16, 4, 2, 24, 8, false)); + for (ggml_type t : { GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_Q8_0 }) { + test_cases.emplace_back(new test_flash_attn_back(128, 8, 2, 512, 64, t)); + } + test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 300, 37, GGML_TYPE_F16)); test_cases.emplace_back(new test_cross_entropy_loss_index()); test_cases.emplace_back(new test_cross_entropy_loss_index({30000, 2, 1, 1})); // a vocabulary's width; few rows keep the finite difference above f32 rounding test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); From 999d682aac013aef843bcaa72fd2d804dffe226f Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 14:41:25 -0500 Subject: [PATCH 12/43] llama: connect labelled decode to physical cache row gradients Assisted-by: OpenAI Codex --- ggml/src/ggml-cpu/ggml-cpu.cpp | 2 +- ggml/src/ggml-cpu/ops.cpp | 24 ++--- ggml/src/ggml-opt.cpp | 3 +- ggml/src/ggml.c | 12 +-- src/llama-batch.cpp | 20 +++- src/llama-context-opt.cpp | 48 ++++++---- src/llama-context.cpp | 29 ++++++ src/llama-cparams.h | 1 + src/llama-graph.cpp | 24 ++--- src/llama-kv-cache.cpp | 53 +++++++---- src/llama-kv-cache.h | 10 +- tests/test-batch-alloc.cpp | 27 ++++++ tests/test-opt-dynamic-accum.cpp | 152 +++++++++++++++++++++++++++++++ tests/test-walk-exact.cpp | 96 +++++++++++++++++++ 14 files changed, 427 insertions(+), 74 deletions(-) diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 6f75f1254d27..96601da9a074 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -475,7 +475,7 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16; case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || - ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && + ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src1->ne[2] % src0->ne[2] == 0 && src1->ne[3] % src0->ne[3] == 0)) && src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_CONV_2D: return ggml_is_contiguous(op->src[0]); diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 8c7c8eb7137c..238c966566c1 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4371,8 +4371,8 @@ static void ggml_compute_forward_out_prod_q_f32( const ggml_type type = src0->type; ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; - GGML_ASSERT(ne02 == ne12); - GGML_ASSERT(ne03 == ne13); + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); GGML_ASSERT(ne2 == ne12); GGML_ASSERT(ne3 == ne13); @@ -4387,8 +4387,8 @@ static void ggml_compute_forward_out_prod_q_f32( GGML_ASSERT(ne0 == ne00); GGML_ASSERT(ne1 == ne10); - GGML_ASSERT(ne2 == ne02); - GGML_ASSERT(ne3 == ne03); + + // nb01 >= nb00 - src0 is not transposed // compute by src0 rows @@ -4425,8 +4425,8 @@ static void ggml_compute_forward_out_prod_q_f32( const int64_t i2 = (ir - i3*ne2*ne1)/ne1; const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); - const int64_t i02 = i2; - const int64_t i03 = i3; + const int64_t i02 = i2 / (ne12 / ne02); + const int64_t i03 = i3 / (ne13 / ne03); //const int64_t i10 = i1; const int64_t i12 = i2; @@ -4461,8 +4461,8 @@ static void ggml_compute_forward_out_prod_f16_f32( GGML_ASSERT(src1->type == GGML_TYPE_F32); GGML_ASSERT(dst->type == GGML_TYPE_F32); - GGML_ASSERT(ne02 == ne12); - GGML_ASSERT(ne03 == ne13); + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); GGML_ASSERT(ne2 == ne12); GGML_ASSERT(ne3 == ne13); @@ -4471,8 +4471,8 @@ static void ggml_compute_forward_out_prod_f16_f32( GGML_ASSERT(ne0 == ne00); GGML_ASSERT(ne1 == ne10); - GGML_ASSERT(ne2 == ne02); - GGML_ASSERT(ne3 == ne03); + + if (ith == 0) { ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); @@ -4491,8 +4491,8 @@ static void ggml_compute_forward_out_prod_f16_f32( const int64_t i2 = (ir - i3*ne2*ne1)/ne1; const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); - const int64_t i02 = i2; - const int64_t i03 = i3; + const int64_t i02 = i2 / (ne12 / ne02); + const int64_t i03 = i3 / (ne13 / ne03); const int64_t i12 = i2; const int64_t i13 = i3; diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 0611db11c66c..ad968b7794cb 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -1208,8 +1208,7 @@ bool ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { if (backward) { ggml_backend_sched_t sched = opt_ctx->backend_sched; ggml_backend_t device = ggml_backend_sched_get_backend(sched, 0); - ggml_backend_dev_t dev = device ? ggml_backend_get_device(device) : nullptr; - if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { + if (device != nullptr) { ggml_cgraph * g = opt_ctx->allocated_graph_copy; for (int i = 0; i < ggml_graph_n_nodes(g); ++i) { ggml_tensor * node = ggml_graph_node(g, i); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 210aad5531f4..17cbf45b8c24 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6667,7 +6667,7 @@ static void ggml_acc_or_set( } cgraph->grads[isrc] = ggml_acc_impl(ctx, acc_into, tensor, nb1, nb2, nb3, offset, cgraph->grad_accs[isrc]); } else { - struct ggml_tensor * a_zero = ggml_scale(ctx, src, 0.0f); // FIXME this is going to produce NaN if a contains inf/NaN + struct ggml_tensor * a_zero = ggml_fill(ctx, ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, src->ne), 0.0f); cgraph->grads[isrc] = ggml_acc_impl(ctx, a_zero, tensor, nb1, nb2, nb3, offset, false); } ggml_format_name(cgraph->grads[isrc], "grad for %s", cgraph->visited_hash_set.keys[isrc]->name); @@ -7181,10 +7181,10 @@ static void ggml_compute_backward( size_t nb2 = tensor->nb[2]; size_t nb3 = tensor->nb[3]; - if (cgraph->grads[isrc0] && src0->type != cgraph->grads[isrc0]->type) { - // gradient is typically F32, but src0 could be other type - size_t ng = ggml_element_size(cgraph->grads[isrc0]); - size_t n0 = ggml_element_size(src0); + if (src0->type != GGML_TYPE_F32) { + // The gradient is F32 even when the viewed source stores quantized blocks. + size_t ng = sizeof(float) * ggml_blck_size(src0->type); + size_t n0 = ggml_type_size(src0->type); GGML_ASSERT(offset % n0 == 0); GGML_ASSERT(nb1 % n0 == 0); GGML_ASSERT(nb2 % n0 == 0); @@ -7701,7 +7701,7 @@ void ggml_build_backward_expand( // A gradient is F32 whatever its tensor's type. A quantized source is admitted where // the backward only re-shapes the gradient or casts its tile to F32 (a training graph // that attends to a q8_0 cache in place, as serving does); every other op needs float. - const bool shapes_only = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + const bool shapes_only = node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_FLASH_ATTN_EXT; GGML_ASSERT(node->src[j]->type == GGML_TYPE_F32 || node->src[j]->type == GGML_TYPE_F16 || (shapes_only && ggml_is_quantized(node->src[j]->type))); diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index e4b06e21b994..7b873d81d6cc 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -56,7 +56,15 @@ bool llama_batch_allocr::init( } if (batch.seq_id) { + if (!batch.n_seq_id) { + LLAMA_LOG_ERROR("%s: sequence IDs require sequence counts\n", __func__); + return false; + } for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] <= 0 || batch.n_seq_id[i] > (int32_t) n_seq_max || !batch.seq_id[i]) { + LLAMA_LOG_ERROR("%s: invalid sequence list at position %d\n", __func__, i); + return false; + } for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { if (batch.seq_id && (batch.seq_id[i][s] < 0 || batch.seq_id[i][s] >= (llama_seq_id) n_seq_max)) { LLAMA_LOG_ERROR("%s: invalid seq_id[%d][%d] = %d >= %d\n", __func__, i, s, batch.seq_id[i][s], (llama_seq_id) n_seq_max); @@ -165,11 +173,17 @@ bool llama_batch_allocr::init( if (batch.label) { bool in_run[LLAMA_MAX_SEQ] = {}; for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.label[i] < -1 || (batch.label[i] >= 0 && (uint32_t) batch.label[i] >= vocab.n_tokens())) { + LLAMA_LOG_ERROR("%s: invalid label[%d] = %d\n", __func__, i, batch.label[i]); + return false; + } const llama_seq_id s = batch.seq_id[i][0]; if (batch.label[i] < 0) { - if (in_run[s]) { - LLAMA_LOG_ERROR("%s: position %d follows a labelled run of sequence %d in one batch: a run ends its sequence's rows\n", __func__, i, s); - return false; + for (int32_t j = 0; j < batch.n_seq_id[i]; ++j) { + if (in_run[batch.seq_id[i][j]]) { + LLAMA_LOG_ERROR("%s: position %d follows a labelled run of sequence %d in one batch: a run ends its sequence's rows\n", __func__, i, batch.seq_id[i][j]); + return false; + } } continue; } diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index f65abef27e0c..41c6ce4ed089 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -350,12 +350,7 @@ bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { static constexpr size_t OPT_HEAD_SLICE_BYTES = 64u*1024*1024; bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch) { - // the loss is the mean over the logits rows: a serving row among them would dilute it - if (ubatch.n_labels() != n_outputs) { - opt_fail("a training ubatch's output rows are its labelled rows (" + std::to_string(ubatch.n_labels()) - + " labelled, " + std::to_string(n_outputs) + " output): the loss is their mean"); - return false; - } + const uint32_t n_labels = ubatch.n_labels(); opt_walk * walk = opt_walk_active; const llama_walk_chunk * xc = walk != nullptr ? walk->xc : nullptr; @@ -375,7 +370,33 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l // and 512 rows, the 4B at window 8192, 2026-10-14), and their size does not shrink with the // model. A slice holds OPT_HEAD_SLICE_BYTES of logits. const int64_t head_rows = std::max(1, (int64_t) (OPT_HEAD_SLICE_BYTES / (sizeof(float) * (size_t) model.vocab.n_tokens()))); - ggml_opt_prepare_alloc_head(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_embd(), res->get_logits(), head_rows); + ggml_tensor * loss_rows = nullptr; + std::vector rows; + std::vector targets; + rows.reserve(n_labels); + targets.reserve(n_labels); + int32_t output_row = 0; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + if (ubatch.output[i]) { + if (ubatch.label[i] >= 0) { + rows.push_back(output_row); + targets.push_back(ubatch.label[i]); + } + ++output_row; + } + } + if (walk != nullptr && n_labels == n_outputs) { + ggml_opt_prepare_alloc_head(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_embd(), res->get_logits(), head_rows); + } else { + // Direct decode returns the original outputs; only the objective selects labelled rows. + ggml_tensor * objective = res->get_logits(); + if (n_labels != n_outputs) { + loss_rows = ggml_new_tensor_1d(ctx_compute_opt, GGML_TYPE_I32, n_labels); + ggml_set_input(loss_rows); + objective = ggml_get_rows(ctx_compute_opt, objective, loss_rows); + } + ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), objective); + } if (xc != nullptr) { ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); for (const auto & io : res->t_walk) { @@ -409,17 +430,10 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l // the target token of each output row, in the order the ubatch emits its outputs (the // order the logits rows take) struct ggml_tensor * labels_t = ggml_opt_labels(opt_ctx); - GGML_ASSERT(labels_t->type == GGML_TYPE_I32 && labels_t->ne[0] == (int64_t) n_outputs); - std::vector targets; - targets.reserve(n_outputs); - for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { - if (ubatch.output[i] != 0) { - const llama_token y = ubatch.label != nullptr ? ubatch.label[i] : -1; - GGML_ASSERT(y < (llama_token) model.vocab.n_tokens()); - targets.push_back(y); - } + GGML_ASSERT(labels_t->type == GGML_TYPE_I32 && labels_t->ne[0] == (int64_t) n_labels); + if (loss_rows) { + ggml_backend_tensor_set(loss_rows, rows.data(), 0, ggml_nbytes(loss_rows)); } - GGML_ASSERT(targets.size() == n_outputs); ggml_backend_tensor_set(labels_t, targets.data(), 0, ggml_nbytes(labels_t)); } if (xc != nullptr) { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0ff2ff008a10..82e8e0c8293d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1342,6 +1342,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll if (opt_ctx != nullptr) { train = !opt_forward_only && ubatch.n_labels() > 0; cparams.training = train; + cparams.walk_active = opt_walk_active != nullptr; } // the new graph parameters @@ -1717,6 +1718,34 @@ int llama_context::decode(const llama_batch & batch_inp) { return -1; } + if (balloc->get_n_labels() > 0) { + if (!opt_ctx) { + LLAMA_LOG_ERROR("%s: labelled rows require an initialized optimizer\n", __func__); + return -1; + } + auto * recurrent = dynamic_cast(memory.get()); + auto * attention = dynamic_cast(memory.get()); + if (auto * hybrid = dynamic_cast(memory.get())) { + recurrent = hybrid->get_mem_recr(); + attention = hybrid->get_mem_attn(); + } + if (attention && !attention->supports_row_gradients()) { + LLAMA_LOG_ERROR("%s: cache row indices exceed the backward gather's index range\n", __func__); + return -1; + } + if (attention && cparams.flash_attn && attention->get_n_stream() != 1) { + LLAMA_LOG_ERROR("%s: flash attention backward requires a unified cache stream\n", __func__); + return -1; + } + const auto & batch = balloc->get_batch(); + for (int32_t i = 0; recurrent && i < batch.n_tokens; ++i) { + if (batch.label[i] >= 0 && recurrent->rollback_pending(batch.seq_id[i][0])) { + LLAMA_LOG_ERROR("%s: sequence %d has pending rollback before training\n", __func__, batch.seq_id[i][0]); + return -1; + } + } + } + const uint32_t n_tokens_all = balloc->get_n_tokens(); const uint32_t n_outputs_all = balloc->get_n_outputs(); diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 8e798ef85e32..9857c99f066d 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -26,6 +26,7 @@ struct llama_cparams { // (llama_context::opt_init probes them), else the composed tiles bool fa_back_native = false; bool walk_exact = false; + bool walk_active = false; uint32_t walk_grad_from = 0; bool walk_surrogate = false; ggml_tensor * const * walk_acc_k = nullptr; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 5d47d281ecae..a579e1bc0ec2 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2647,15 +2647,13 @@ ggml_tensor * llm_graph_context::build_attn_mha( v = ggml_transpose(ctx0, v); } - // this can happen when KV cache is not used (e.g. an embedding model with non-causal attn). - // A training graph keeps its chunk's own K/V at F32 (every flash-attention backend takes - // F32 K/V): an F16 cast there rounded the trained K/V and moved the step (test-walk-exact, - // 2026-10-11: cosine 0.9969 against the one-graph reference at an f16 cache, 1.0000 at F32) - if (!cparams.training && k->type == GGML_TYPE_F32) { + // Direct decode uses serving precision for every row. The walk keeps its chunk's F32 K/V + // for its existing gradient comparison against a single graph. + if ((!cparams.training || !cparams.walk_active) && k->type == GGML_TYPE_F32) { k = ggml_cast(ctx0, k, GGML_TYPE_F16); } - if (!cparams.training && v->type == GGML_TYPE_F32) { + if ((!cparams.training || !cparams.walk_active) && v->type == GGML_TYPE_F32) { v = ggml_cast(ctx0, v, GGML_TYPE_F16); } @@ -2905,7 +2903,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * q = q_cur; ggml_tensor * k; ggml_tensor * v; - if (cparams.training && uses_flash_attn(kq_b, sinks, v_mla)) { + if (cparams.training && cparams.walk_active && uses_flash_attn(kq_b, sinks, v_mla)) { // THE TRAINING GRAPH IS THE SERVING GRAPH. This chunk writes its K/V into the cache // exactly as serving does (SET_ROWS, the cache's own type) and attention reads the // cache in place: no copy of the context, at any window. The walk decoded the run @@ -2953,7 +2951,7 @@ ggml_tensor * llm_graph_context::build_attn( kq_mask = ggml_cont(ctx0, ggml_view_4d(ctx0, kq_mask, n_win, kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3], kq_mask->nb[1], kq_mask->nb[2], kq_mask->nb[3], 0)); } - } else if (cparams.training) { + } else if (cparams.training && cparams.walk_active) { // EXPLICIT ATTENTION (a layer FLASH_ATTN_EXT's backward does not cover yet: ALiBi, a // logit softcap, sinks, a KQ bias, MLA) reads a copy of the context instead. // A backward pass cannot go through the cache: the write is an in-place SET_ROWS @@ -3048,11 +3046,13 @@ ggml_tensor * llm_graph_context::build_attn( 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 * written_k = mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il); + ggml_tensor * written_v = mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il); + ggml_build_forward_expand(gf, written_k); + ggml_build_forward_expand(gf, written_v); - k = mctx_cur->get_k(ctx0, il); - v = mctx_cur->get_v(ctx0, il); + k = mctx_cur->get_k(ctx0, il, cparams.training ? written_k : nullptr); + v = mctx_cur->get_v(ctx0, il, cparams.training ? written_v : nullptr); } ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index f815cd76656a..1d544de4a895 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1249,7 +1249,7 @@ uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { return result; } -ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { +ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo, ggml_tensor * written) const { const int32_t ikv = map_layer_ids.at(il); auto * k = layers[ikv].k; @@ -1261,7 +1261,7 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; - return ggml_view_4d(ctx, k, + return ggml_view_4d(ctx, written ? written : k, hparams.n_embd_head_k(il), hparams.n_head_kv(il), n_kv, ns, ggml_row_size(k->type, hparams.n_embd_head_k(il)), ggml_row_size(k->type, n_embd_k_gqa), @@ -1269,7 +1269,7 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k ggml_row_size(k->type, n_embd_k_gqa*kv_size)*sinfo.s0); } -ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { +ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo, ggml_tensor * written) const { const int32_t ikv = map_layer_ids.at(il); auto * v = layers[ikv].v; @@ -1284,7 +1284,7 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k if (!v_trans) { // note: v->nb[1] <= v->nb[2] - return ggml_view_4d(ctx, v, + return ggml_view_4d(ctx, written ? written : v, hparams.n_embd_head_v(il), hparams.n_head_kv(il), n_kv, ns, ggml_row_size(v->type, hparams.n_embd_head_v(il)), // v->nb[1] ggml_row_size(v->type, n_embd_v_gqa), // v->nb[2] @@ -1293,7 +1293,7 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k } // note: v->nb[1] > v->nb[2] - return ggml_view_4d(ctx, v, + return ggml_view_4d(ctx, written ? written : v, n_kv, hparams.n_head_kv(il), hparams.n_embd_head_v(il), ns, ggml_row_size(v->type, kv_size*hparams.n_embd_head_v(il)), // v->nb[1] ggml_row_size(v->type, kv_size), // v->nb[2] @@ -1392,10 +1392,29 @@ ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggm return ggml_set_rows(ctx, v_view, v_cur, v_idxs); } +static ggml_type cache_index_type(uint64_t extent) { + return extent <= (uint64_t) INT32_MAX ? GGML_TYPE_I32 : GGML_TYPE_I64; +} + +bool llama_kv_cache::supports_row_gradients() const { + const uint64_t extent = (uint64_t) get_size()*n_stream; + return cache_index_type(extent) == GGML_TYPE_I32 && + cache_index_type(v_trans ? extent*hparams.n_embd_v_gqa_max() : extent) == GGML_TYPE_I32; +} + +static void set_cache_index(ggml_tensor * dst, size_t i, int64_t value) { + if (dst->type == GGML_TYPE_I32) { + GGML_ASSERT(value >= 0 && value <= INT32_MAX); + ((int32_t *) dst->data)[i] = (int32_t) value; + } else { + GGML_ASSERT(dst->type == GGML_TYPE_I64); + ((int64_t *) dst->data)[i] = value; + } +} ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { const uint32_t n_tokens = ubatch.n_tokens; - ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); + ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, cache_index_type((uint64_t) get_size()*n_stream), n_tokens); ggml_set_input(k_idxs); @@ -1408,9 +1427,9 @@ ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama ggml_tensor * v_idxs; if (!v_trans) { - v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); + v_idxs = ggml_new_tensor_1d(ctx, cache_index_type((uint64_t) get_size()*n_stream), n_tokens); } else { - v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens*hparams.n_embd_v_gqa_max()); + v_idxs = ggml_new_tensor_1d(ctx, cache_index_type((uint64_t) get_size()*n_stream*hparams.n_embd_v_gqa_max()), n_tokens*hparams.n_embd_v_gqa_max()); } ggml_set_input(v_idxs); @@ -1464,13 +1483,13 @@ void llama_kv_cache::set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ub GGML_ASSERT(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream()); GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); - int64_t * data = (int64_t *) dst->data; + for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { const int64_t offs = sinfo.strm[s]*get_size(); for (uint32_t i = 0; i < sinfo.size(); ++i) { - data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i]; + set_cache_index(dst, s*sinfo.size() + i, offs + sinfo.idxs[s][i]); } } } @@ -1480,14 +1499,14 @@ void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ub GGML_ASSERT(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream()); GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); - int64_t * data = (int64_t *) dst->data; + if (!v_trans) { for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { const int64_t offs = sinfo.strm[s]*get_size(); for (uint32_t i = 0; i < sinfo.size(); ++i) { - data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i]; + set_cache_index(dst, s*sinfo.size() + i, offs + sinfo.idxs[s][i]); } } } else { @@ -1501,7 +1520,7 @@ void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ub for (uint32_t i = 0; i < sinfo.size(); ++i) { for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { - data[s*sinfo.size()*n_embd_v_gqa + i*n_embd_v_gqa + j] = offs + j*kv_size + sinfo.idxs[s][i]; + set_cache_index(dst, s*sinfo.size()*n_embd_v_gqa + i*n_embd_v_gqa + j, offs + j*kv_size + sinfo.idxs[s][i]); } } } @@ -2594,12 +2613,12 @@ ggml_type llama_kv_cache_context::type_v() const { return kv->type_v(); } -ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il) const { - return kv->get_k(ctx, il, n_kv, sinfos[i_cur]); +ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il, ggml_tensor * written) const { + return kv->get_k(ctx, il, n_kv, sinfos[i_cur], written); } -ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) const { - return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); +ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il, ggml_tensor * written) const { + return kv->get_v(ctx, il, n_kv, sinfos[i_cur], written); } ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const { diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index f00ec2137bb3..8e27fdf12e15 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -158,6 +158,8 @@ class llama_kv_cache : public llama_memory_i { uint32_t get_size() const; uint32_t get_n_stream() const; + bool supports_row_gradients() const; + bool get_has_shift() const; ggml_type type_k() const; @@ -175,8 +177,8 @@ class llama_kv_cache : public llama_memory_i { uint32_t get_n_kv(const slot_info & sinfo) const; // get views of the current state of the cache - ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; - ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; + ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo, ggml_tensor * written = nullptr) const; + ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo, ggml_tensor * written = nullptr) const; // store k_cur and v_cur in the cache based on the provided head location ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; @@ -372,8 +374,8 @@ class llama_kv_cache_context : public llama_memory_context_i { ggml_type type_v() const; // get views of the current state of the cache - ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; - ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; + ggml_tensor * get_k(ggml_context * ctx, int32_t il, ggml_tensor * written = nullptr) const; + ggml_tensor * get_v(ggml_context * ctx, int32_t il, ggml_tensor * written = nullptr) const; // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp index 66d29d6f5164..e41870193ae8 100644 --- a/tests/test-batch-alloc.cpp +++ b/tests/test-batch-alloc.cpp @@ -96,6 +96,33 @@ struct batch_builder { static void test_init(testing & t) { llama_vocab vocab; + t.test("rejects_invalid_labels_before_memory_mutation", [&](testing & t) { + batch_builder bb; + bb.add(0, {0}, true); + mock_memory mem; + llama_batch batch = bb.make(); + llama_token label = -2; + batch.label = &label; + llama_batch_allocr ba(1); + t.assert_true("only -1 denotes no label", !ba.init(batch, vocab, &mem, bb.n_embd, 1, false)); + label = 0; + t.assert_true("target must belong to vocabulary", !ba.init(batch, vocab, &mem, bb.n_embd, 1, false)); + label = -1; + t.assert_true("unlabelled inference remains valid", ba.init(batch, vocab, &mem, bb.n_embd, 1, false)); + }); + + t.test("rejects_missing_or_empty_sequence_lists", [&](testing & t) { + batch_builder bb; + bb.add(0, {0}, true); + llama_batch batch = bb.make(); + llama_batch_allocr ba(1); + batch.n_seq_id = nullptr; + t.assert_true("missing counts", !ba.init(batch, vocab, nullptr, bb.n_embd, 1, false)); + batch = bb.make(); + batch.n_seq_id[0] = 0; + t.assert_true("empty sequence list", !ba.init(batch, vocab, nullptr, bb.n_embd, 1, false)); + }); + t.test("rejects_n_seq_max_too_large", [&](testing & t) { batch_builder bb; bb.add(0, {0}, true); diff --git a/tests/test-opt-dynamic-accum.cpp b/tests/test-opt-dynamic-accum.cpp index ff5445fda605..6b392968c8ec 100644 --- a/tests/test-opt-dynamic-accum.cpp +++ b/tests/test-opt-dynamic-accum.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #define CHECK(cond) do { if (!(cond)) { fprintf(stderr, "FAILED %s:%d: %s\n", __FILE__, __LINE__, #cond); exit(1); } } while (0) @@ -181,11 +182,162 @@ static void run_manual() { ggml_backend_free(cpu); } +static void run_out_prod_broadcast(ggml_type type, bool transpose) { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + CHECK(cpu != nullptr); + ggml_init_params cp = { 32 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; + ggml_context * ctx = ggml_init(cp); + ggml_tensor * input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 32, 2, 2, 2); + ggml_tensor * a = ggml_cast(ctx, input, type); + ggml_tensor * decoded = ggml_cast(ctx, a, GGML_TYPE_F32); + ggml_tensor * b_storage = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, transpose ? 2 : 1, transpose ? 1 : 2, 4, 4); + ggml_tensor * b = transpose ? ggml_transpose(ctx, b_storage) : b_storage; + ggml_tensor * out = ggml_out_prod(ctx, a, b); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, out); + ggml_build_forward_expand(graph, decoded); + CHECK(ggml_backend_supports_op(cpu, out)); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, cpu); + float av[32 * 2 * 2 * 2], bv[2 * 4 * 4]; + for (int batch = 0; batch < 4; ++batch) { + for (int row = 0; row < 2; ++row) { + for (int x = 0; x < 32; ++x) { + av[(batch * 2 + row) * 32 + x] = float((batch + 1) * (row + 1)); + } + } + } + for (int batch = 0; batch < 16; ++batch) { + bv[batch * 2] = float(batch + 1); + bv[batch * 2 + 1] = -float(batch + 2); + } + ggml_backend_tensor_set(input, av, 0, sizeof(av)); + ggml_backend_tensor_set(b_storage, bv, 0, sizeof(bv)); + CHECK(ggml_backend_graph_compute(cpu, graph) == GGML_STATUS_SUCCESS); + float actual[32 * 4 * 4]; + ggml_backend_tensor_get(out, actual, 0, sizeof(actual)); + ggml_backend_tensor_get(decoded, av, 0, sizeof(av)); + for (int z = 0; z < 4; ++z) { + for (int y = 0; y < 4; ++y) { + const int a_batch = (z / 2) * 2 + y / 2; + const int b_batch = z * 4 + y; + for (int x = 0; x < 32; ++x) { + const float expected = av[a_batch * 64 + x] * bv[b_batch * 2] + + av[a_batch * 64 + 32 + x] * bv[b_batch * 2 + 1]; + if (std::fabs(actual[b_batch * 32 + x] - expected) >= 1e-4f) { + fprintf(stderr, "out_prod type=%s transpose=%d y=%d z=%d actual=%g expected=%g\n", + ggml_type_name(type), transpose, y, z, actual[b_batch * 32 + x], expected); + } + CHECK(std::fabs(actual[b_batch * 32 + x] - expected) < 1e-4f); + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + ggml_backend_free(cpu); +} + +static void run_cache_view_gradient(ggml_type type) { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 1, GGML_DEFAULT_GRAPH_SIZE, false, true); + ggml_init_params sp = { 8 * ggml_tensor_overhead(), nullptr, true }; + ggml_context * storage = ggml_init(sp); + ggml_tensor * weights = ggml_new_tensor_2d(storage, GGML_TYPE_F32, 32, 1); + ggml_tensor * cache = ggml_new_tensor_2d(storage, type, 32, 4); + ggml_set_param(weights); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(storage, cpu); + float values[32]; + for (float & value : values) { value = 1.0f; } + ggml_backend_tensor_set(weights, values, 0, sizeof(values)); + ggml_backend_tensor_memset(cache, 0, 0, ggml_nbytes(cache)); + ggml_opt_params params = ggml_opt_default_params(sched, GGML_OPT_LOSS_TYPE_SUM); + params.optimizer = GGML_OPT_OPTIMIZER_TYPE_SGD; + params.get_opt_pars = sgd_pars; + ggml_opt_context_t opt = ggml_opt_init(params); + ggml_init_params cp = { GGML_DEFAULT_GRAPH_SIZE * ggml_tensor_overhead() + 4 * ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), nullptr, true }; + ggml_context * ctx = ggml_init(cp); + ggml_tensor * row = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_set_input(row); + ggml_tensor * written = ggml_set_rows(ctx, cache, weights, row); + ggml_tensor * view = ggml_view_2d(ctx, written, 32, 1, written->nb[1], 2 * written->nb[1]); + view = ggml_cast(ctx, view, GGML_TYPE_F32); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, true); + ggml_build_forward_expand(graph, view); + ggml_opt_prepare_alloc(opt, ctx, graph, row, view); + CHECK(ggml_opt_alloc(opt, true)); + const int32_t index = 2; + ggml_backend_tensor_set(row, &index, 0, sizeof(index)); + ggml_opt_eval(opt, nullptr); + CHECK(ggml_opt_refusal(opt)[0] == '\0'); + ggml_backend_tensor_get(weights, values, 0, sizeof(values)); + for (float value : values) { CHECK(std::fabs(value - (1.0f - LR)) < 1e-6f); } + ggml_opt_free(opt); + ggml_free(ctx); + ggml_backend_buffer_free(buf); + ggml_free(storage); + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); +} + +static void run_label_projection() { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + CHECK(cpu != nullptr); + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 1, GGML_DEFAULT_GRAPH_SIZE, false, true); + ggml_init_params sp = { 8 * ggml_tensor_overhead(), nullptr, true }; + ggml_context * storage = ggml_init(sp); + ggml_tensor * weights = ggml_new_tensor_2d(storage, GGML_TYPE_F32, 2, 3); + ggml_set_param(weights); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(storage, cpu); + const float before[] = { 8.0f, -8.0f, 0.0f, 0.0f, -8.0f, 8.0f }; + ggml_backend_tensor_set(weights, before, 0, sizeof(before)); + ggml_opt_params params = ggml_opt_default_params(sched, GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS); + params.optimizer = GGML_OPT_OPTIMIZER_TYPE_SGD; + params.get_opt_pars = sgd_pars; + ggml_opt_context_t opt = ggml_opt_init(params); + ggml_init_params cp = { GGML_DEFAULT_GRAPH_SIZE * ggml_tensor_overhead() + 4 * ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), nullptr, true }; + ggml_context * graph_ctx = ggml_init(cp); + ggml_tensor * logits = ggml_scale(graph_ctx, weights, 1.0f); + ggml_set_output(logits); + ggml_tensor * rows = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 1); + ggml_set_input(rows); + ggml_tensor * objective = ggml_get_rows(graph_ctx, logits, rows); + ggml_cgraph * gf = ggml_new_graph_custom(graph_ctx, GGML_DEFAULT_GRAPH_SIZE, true); + ggml_build_forward_expand(gf, logits); + ggml_opt_prepare_alloc(opt, graph_ctx, gf, rows, objective); + CHECK(ggml_opt_alloc(opt, true)); + const int32_t row = 1, target = 0; + ggml_backend_tensor_set(rows, &row, 0, sizeof(row)); + ggml_backend_tensor_set(ggml_opt_labels(opt), &target, 0, sizeof(target)); + ggml_opt_eval(opt, nullptr); + CHECK(ggml_opt_refusal(opt)[0] == '\0'); + float after[6], served[6]; + ggml_backend_tensor_get(weights, after, 0, sizeof(after)); + ggml_backend_tensor_get(logits, served, 0, sizeof(served)); + for (int i = 0; i < 6; ++i) { + const float delta = i == 2 ? LR * 0.5f : i == 3 ? -LR * 0.5f : 0.0f; + CHECK(std::fabs(after[i] - before[i] - delta) < 1e-6f); + CHECK(served[i] == before[i]); + } + ggml_opt_free(opt); + ggml_free(graph_ctx); + ggml_backend_buffer_free(buf); + ggml_free(storage); + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); +} + int main() { ggml_backend_load_all(); run(1); run(2); run_manual(); + run_label_projection(); + for (ggml_type type : { GGML_TYPE_F16, GGML_TYPE_Q8_0 }) { + run_cache_view_gradient(type); + run_out_prod_broadcast(type, false); + run_out_prod_broadcast(type, true); + } printf("test-opt-dynamic-accum: OK\n"); return 0; } diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 04cf6aec61d8..f4cc172e50eb 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -21,6 +21,8 @@ #include "common.h" #include "llama.h" #include "../src/llama-adapter.h" +#include "../src/llama-memory-hybrid.h" +#include "../src/llama-memory-recurrent.h" #include #include @@ -270,6 +272,86 @@ static bool same_step(const agreement & r) { return r.cosine > 0.995 && std::fabs(r.norm_ratio - 1.0) < 0.02; } +static std::vector sequence_state(llama_context * ctx, llama_seq_id seq) { + std::vector bytes(llama_state_seq_get_size(ctx, seq)); + GGML_ASSERT(llama_state_seq_get_data(ctx, bytes.data(), bytes.size(), seq) == bytes.size()); + return bytes; +} + +static void mixed_decode(const common_params & params, llama_model * model, const std::string & init, + const std::vector & tokens, bool flash, ggml_type cache_type, uint32_t undo = 0) { + auto cp = common_context_params_to_llama(params); + cp.n_ctx = 128; + cp.n_batch = cp.n_ubatch = undo ? 8 : 2; + cp.n_seq_max = 2; + cp.kv_unified = true; + cp.n_rs_seq = undo; + cp.flash_attn_type = flash ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + cp.type_k = cp.type_v = cache_type; + llama_context * ctx = llama_init_from_model(model, cp); + GGML_ASSERT(ctx); + llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); + float scale = 1.0f; + GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); + llama_opt_params op{}; + op.param_filter = llama_opt_param_filter_all; + op.get_opt_pars = sgd_pars; + op.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; + op.adapter = adapter; + llama_batch batch = llama_batch_init(2, 0, 1); + common_batch_add(batch, tokens[0], 0, {0}, true); + common_batch_add(batch, tokens[1], 0, {1}, true); + GGML_ASSERT(llama_decode(ctx, batch) == 0); + const auto reference_neighbor = sequence_state(ctx, 1); + const auto reference_target = sequence_state(ctx, 0); + batch.label[0] = tokens[2]; + batch.pos[0] = batch.pos[1] = 1; + GGML_ASSERT(llama_decode(ctx, batch) != 0); + GGML_ASSERT(sequence_state(ctx, 0) == reference_target && sequence_state(ctx, 1) == reference_neighbor); + llama_memory_clear(llama_get_memory(ctx), true); + batch.pos[0] = batch.pos[1] = 0; + llama_opt_init(ctx, model, op); + const auto initial_weights = adapter_params(adapter); + batch.label[0] = tokens[2]; + GGML_ASSERT(llama_decode(ctx, batch) == 0); + GGML_ASSERT(llama_get_logits_ith(ctx, 0) && llama_get_logits_ith(ctx, 1)); + const auto neighbor = sequence_state(ctx, 1); + GGML_ASSERT(neighbor == reference_neighbor); + const auto trained = sequence_state(ctx, 0); + GGML_ASSERT(neighbor != trained); + const auto weights = adapter_params(adapter); + GGML_ASSERT(weights != initial_weights); + common_batch_clear(batch); + common_batch_add(batch, tokens[2], 1, {0}, true); + batch.label[0] = llama_vocab_n_tokens(llama_model_get_vocab(model)); + GGML_ASSERT(llama_decode(ctx, batch) != 0); + GGML_ASSERT(sequence_state(ctx, 0) == trained && sequence_state(ctx, 1) == neighbor); + GGML_ASSERT(adapter_params(adapter) == weights); + batch.label[0] = tokens[3]; + GGML_ASSERT(llama_decode(ctx, batch) == 0); + GGML_ASSERT(sequence_state(ctx, 1) == neighbor); + GGML_ASSERT(adapter_params(adapter) != weights); + if (undo) { + auto * memory = llama_get_memory(ctx); + auto * recurrent = dynamic_cast(memory); + if (auto * hybrid = dynamic_cast(memory)) { + recurrent = hybrid->get_mem_recr(); + } + GGML_ASSERT(recurrent); + recurrent->set_rs_idx(0, 1); + const auto pending = sequence_state(ctx, 0); + const auto pending_weights = adapter_params(adapter); + batch.pos[0] = 2; + GGML_ASSERT(llama_decode(ctx, batch) != 0); + GGML_ASSERT(sequence_state(ctx, 0) == pending && sequence_state(ctx, 1) == neighbor); + GGML_ASSERT(adapter_params(adapter) == pending_weights); + } + printf("mixed decode flash=%d cache=%s undo=%u: same-position sequences, neighbour preservation, admission refusal OK\n", flash, ggml_type_name(cache_type), undo); + llama_batch_free(batch); + llama_free(ctx); + llama_adapter_lora_free(adapter); +} + int main(int argc, char ** argv) { common_params params; if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { @@ -302,6 +384,20 @@ int main(int argc, char ** argv) { llama_free(tok_ctx); GGML_ASSERT(tokens.size() > REFUSAL_WINDOW + 1); + if (getenv("WALK_MIXED_ONLY")) { + mixed_decode(params, model, init, tokens, false, GGML_TYPE_F32); + for (ggml_type cache : { GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0 }) { + mixed_decode(params, model, init, tokens, true, cache); + } + if (llama_model_is_hybrid(model) || llama_model_is_recurrent(model)) { + mixed_decode(params, model, init, tokens, true, GGML_TYPE_Q8_0, 2); + } + std::remove(init.c_str()); + llama_model_free(model); + llama_backend_free(); + return 0; + } + // MEASUREMENT MODE (WALK_MEASURE="2048,8192,..."): the device graph per chunk of a context + // reply window (the last 512 positions labelled), on the explicit path (F32 cache, no flash // attention) and through FLASH_ATTN_EXT (f16 cache, flash attention), at each window. Prints a From 87a0f17e93acf41468ecce9d99a800992ceb39b0 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 15:03:46 -0500 Subject: [PATCH 13/43] llama: keep serving indices and compose training admission in memory owners Assisted-by: OpenAI Codex --- src/llama-context.cpp | 22 ++---------------- src/llama-graph.cpp | 38 ++++++++++++++++---------------- src/llama-kv-cache-dsa-iswa.cpp | 7 ++++++ src/llama-kv-cache-dsa-iswa.h | 2 ++ src/llama-kv-cache-dsa.cpp | 4 ++++ src/llama-kv-cache-dsa.h | 2 ++ src/llama-kv-cache-dsv4.cpp | 4 ++++ src/llama-kv-cache-dsv4.h | 2 ++ src/llama-kv-cache-iswa.cpp | 7 ++++++ src/llama-kv-cache-iswa.h | 2 ++ src/llama-kv-cache-msa.cpp | 4 ++++ src/llama-kv-cache-msa.h | 2 ++ src/llama-kv-cache.cpp | 28 +++++++++++++++-------- src/llama-kv-cache.h | 10 +++++---- src/llama-memory-hybrid-idx.cpp | 7 ++++++ src/llama-memory-hybrid-idx.h | 2 ++ src/llama-memory-hybrid-iswa.cpp | 7 ++++++ src/llama-memory-hybrid-iswa.h | 2 ++ src/llama-memory-hybrid.cpp | 7 ++++++ src/llama-memory-hybrid.h | 2 ++ src/llama-memory-recurrent.cpp | 9 ++++++++ src/llama-memory-recurrent.h | 2 ++ src/llama-memory.h | 3 +++ tests/test-batch-alloc.cpp | 2 ++ 24 files changed, 125 insertions(+), 52 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 82e8e0c8293d..a79c04585427 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1723,29 +1723,11 @@ int llama_context::decode(const llama_batch & batch_inp) { LLAMA_LOG_ERROR("%s: labelled rows require an initialized optimizer\n", __func__); return -1; } - auto * recurrent = dynamic_cast(memory.get()); - auto * attention = dynamic_cast(memory.get()); - if (auto * hybrid = dynamic_cast(memory.get())) { - recurrent = hybrid->get_mem_recr(); - attention = hybrid->get_mem_attn(); - } - if (attention && !attention->supports_row_gradients()) { - LLAMA_LOG_ERROR("%s: cache row indices exceed the backward gather's index range\n", __func__); + if (const char * why = memory->training_error(balloc->get_batch(), cparams.flash_attn)) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, why); return -1; } - if (attention && cparams.flash_attn && attention->get_n_stream() != 1) { - LLAMA_LOG_ERROR("%s: flash attention backward requires a unified cache stream\n", __func__); - return -1; - } - const auto & batch = balloc->get_batch(); - for (int32_t i = 0; recurrent && i < batch.n_tokens; ++i) { - if (batch.label[i] >= 0 && recurrent->rollback_pending(batch.seq_id[i][0])) { - LLAMA_LOG_ERROR("%s: sequence %d has pending rollback before training\n", __func__, batch.seq_id[i][0]); - return -1; - } - } } - const uint32_t n_tokens_all = balloc->get_n_tokens(); const uint32_t n_outputs_all = balloc->get_n_outputs(); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index a579e1bc0ec2..e6521d2c9c2d 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2844,8 +2844,8 @@ static std::unique_ptr build_attn_inp_kv_impl( { GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA"); - inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch); - inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3094,7 +3094,7 @@ static std::unique_ptr build_attn_inp_k_impl( { GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA"); - inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3468,14 +3468,14 @@ static std::unique_ptr build_attn_inp_k_dsa_impl( auto inp = std::make_unique(hparams, cparams, mctx_cur); { - inp->self_k_idxs_mla = mctx_cur->get_mla()->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs_mla = mctx_cur->get_mla()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask_mla = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla(), ubatch, cparams); inp->self_kq_mask_mla_cnv = inp->self_kq_mask_mla; } { - inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); // ensure that mask type matches fused lightning indexer use (requires f16 mask) auto cparams_copy = cparams; @@ -3506,7 +3506,7 @@ llm_graph_input_attn_k_dsa_iswa * llm_graph_context::build_attn_inp_k_dsa_iswa() // build_attn_inp_k_impl rejects SWA caches, so construct the input directly auto inp_swa = std::make_unique(hparams, cparams, mctx_cur->get_swa()); - inp_swa->self_k_idxs = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); + inp_swa->self_k_idxs = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp_swa->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp_swa->self_kq_mask_cnv = inp_swa->self_kq_mask; @@ -3527,8 +3527,8 @@ llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_ { GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA"); - inp->self_k_idxs = mctx_base->build_input_k_idxs(ctx0, ubatch); - inp->self_v_idxs = mctx_base->build_input_v_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_base->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp->self_v_idxs = mctx_base->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_base, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3538,7 +3538,7 @@ llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_ inp->self_v_rot = mctx_base->build_input_v_rot(ctx0); if (msa_enabled) { - inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); } return (llm_graph_input_attn_kv_msa *) res->add_input(std::move(inp)); @@ -3553,8 +3553,8 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const auto inp = std::make_unique(hparams, cparams, mctx_cur); { - inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch); - inp->self_v_idxs = mctx_cur->get_base()->build_input_v_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp->self_v_idxs = mctx_cur->get_base()->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3563,8 +3563,8 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const { GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA"); - inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); - inp->self_v_idxs_swa = mctx_cur->get_swa()->build_input_v_idxs(ctx0, ubatch); + inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp->self_v_idxs_swa = mctx_cur->get_swa()->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa; @@ -3585,7 +3585,7 @@ llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { auto inp = std::make_unique(hparams, cparams, mctx_cur); { - inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3594,7 +3594,7 @@ llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { { GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA"); - inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa; @@ -3832,16 +3832,16 @@ llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa() auto inp_attn = std::make_unique(hparams, cparams, attn_ctx); { - inp_attn->self_k_idxs = attn_ctx->get_base()->build_input_k_idxs(ctx0, ubatch); - inp_attn->self_v_idxs = attn_ctx->get_base()->build_input_v_idxs(ctx0, ubatch); + inp_attn->self_k_idxs = attn_ctx->get_base()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp_attn->self_v_idxs = attn_ctx->get_base()->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp_attn->self_kq_mask = build_attn_inp_kq_mask(ctx0, attn_ctx->get_base(), ubatch, cparams); inp_attn->self_kq_mask_cnv = inp_attn->self_kq_mask; } { - inp_attn->self_k_idxs_swa = attn_ctx->get_swa()->build_input_k_idxs(ctx0, ubatch); - inp_attn->self_v_idxs_swa = attn_ctx->get_swa()->build_input_v_idxs(ctx0, ubatch); + inp_attn->self_k_idxs_swa = attn_ctx->get_swa()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp_attn->self_v_idxs_swa = attn_ctx->get_swa()->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp_attn->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, attn_ctx->get_swa(), ubatch, cparams); inp_attn->self_kq_mask_swa_cnv = inp_attn->self_kq_mask_swa; diff --git a/src/llama-kv-cache-dsa-iswa.cpp b/src/llama-kv-cache-dsa-iswa.cpp index dc10342a19c0..83d8e6cd3528 100644 --- a/src/llama-kv-cache-dsa-iswa.cpp +++ b/src/llama-kv-cache-dsa-iswa.cpp @@ -339,3 +339,10 @@ const llama_kv_cache_context * llama_kv_cache_dsa_iswa_context::get_swa() const return static_cast(ctx_swa.get()); } + +const char * llama_kv_cache_dsa_iswa::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = kv_swa->training_error(batch, flash_attn)) { + return why; + } + return kv_dsa->training_error(batch, flash_attn); +} diff --git a/src/llama-kv-cache-dsa-iswa.h b/src/llama-kv-cache-dsa-iswa.h index 28cf95bf0518..08e270acf838 100644 --- a/src/llama-kv-cache-dsa-iswa.h +++ b/src/llama-kv-cache-dsa-iswa.h @@ -34,6 +34,8 @@ class llama_kv_cache_dsa_iswa : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp index 96cb045d2e5d..54a9b4b440f7 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -260,3 +260,7 @@ const llama_kv_cache_context * llama_kv_cache_dsa_context::get_lid() const { return static_cast(ctx_lid.get()); } + +const char * llama_kv_cache_dsa::training_error(const llama_batch &, bool) const { + return "DSA cache row gradients are not supported"; +} diff --git a/src/llama-kv-cache-dsa.h b/src/llama-kv-cache-dsa.h index e74fc4d9100f..0fe3b538e8b9 100644 --- a/src/llama-kv-cache-dsa.h +++ b/src/llama-kv-cache-dsa.h @@ -36,6 +36,8 @@ class llama_kv_cache_dsa : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index b073ed303e5d..52e843388bba 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -2252,3 +2252,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ return reserve_plan_lid; } + +const char * llama_kv_cache_dsv4::training_error(const llama_batch &, bool) const { + return "DSV4 cache row gradients are not supported"; +} diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index ce39867c0342..da5af159720f 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -109,6 +109,8 @@ class llama_kv_cache_dsv4 : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp index e91866469adf..42bd02a59782 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -361,3 +361,10 @@ const llama_kv_cache_context * llama_kv_cache_iswa_context::get_swa() const { return static_cast(ctx_swa.get()); } + +const char * llama_kv_cache_iswa::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = kv_base->training_error(batch, flash_attn)) { + return why; + } + return kv_swa->training_error(batch, flash_attn); +} diff --git a/src/llama-kv-cache-iswa.h b/src/llama-kv-cache-iswa.h index 7dab6eaa82c8..7f7364174f40 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -54,6 +54,8 @@ class llama_kv_cache_iswa : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache-msa.cpp b/src/llama-kv-cache-msa.cpp index 55ef286cafa7..32dc7c91116d 100644 --- a/src/llama-kv-cache-msa.cpp +++ b/src/llama-kv-cache-msa.cpp @@ -393,3 +393,7 @@ void llama_kv_cache_msa_context::set_input_pos_mask(ggml_tensor * dst, const lla } } } + +const char * llama_kv_cache_msa::training_error(const llama_batch &, bool) const { + return "MSA cache row gradients are not supported"; +} diff --git a/src/llama-kv-cache-msa.h b/src/llama-kv-cache-msa.h index f09b6d32b044..cb93a00b034f 100644 --- a/src/llama-kv-cache-msa.h +++ b/src/llama-kv-cache-msa.h @@ -33,6 +33,8 @@ class llama_kv_cache_msa : public llama_memory_i { // llama_memory_i + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 1d544de4a895..ce72a7a20998 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1411,25 +1411,25 @@ static void set_cache_index(ggml_tensor * dst, size_t i, int64_t value) { ((int64_t *) dst->data)[i] = value; } } -ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { +ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients) const { const uint32_t n_tokens = ubatch.n_tokens; - ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, cache_index_type((uint64_t) get_size()*n_stream), n_tokens); + ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, row_gradients ? cache_index_type((uint64_t) get_size()*n_stream) : GGML_TYPE_I64, n_tokens); ggml_set_input(k_idxs); return k_idxs; } -ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { +ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients) const { const uint32_t n_tokens = ubatch.n_tokens; ggml_tensor * v_idxs; if (!v_trans) { - v_idxs = ggml_new_tensor_1d(ctx, cache_index_type((uint64_t) get_size()*n_stream), n_tokens); + v_idxs = ggml_new_tensor_1d(ctx, row_gradients ? cache_index_type((uint64_t) get_size()*n_stream) : GGML_TYPE_I64, n_tokens); } else { - v_idxs = ggml_new_tensor_1d(ctx, cache_index_type((uint64_t) get_size()*n_stream*hparams.n_embd_v_gqa_max()), n_tokens*hparams.n_embd_v_gqa_max()); + v_idxs = ggml_new_tensor_1d(ctx, row_gradients ? cache_index_type((uint64_t) get_size()*n_stream*hparams.n_embd_v_gqa_max()) : GGML_TYPE_I64, n_tokens*hparams.n_embd_v_gqa_max()); } ggml_set_input(v_idxs); @@ -2629,12 +2629,12 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_ return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]); } -ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { - return kv->build_input_k_idxs(ctx, ubatch); +ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients) const { + return kv->build_input_k_idxs(ctx, ubatch, row_gradients); } -ggml_tensor * llama_kv_cache_context::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { - return kv->build_input_v_idxs(ctx, ubatch); +ggml_tensor * llama_kv_cache_context::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients) const { + return kv->build_input_v_idxs(ctx, ubatch, row_gradients); } ggml_tensor * llama_kv_cache_context::build_input_k_rot(ggml_context * ctx) const { @@ -2672,3 +2672,13 @@ void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const { void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const { kv->set_input_v_rot(dst); } + +const char * llama_kv_cache::training_error(const llama_batch &, bool flash_attn) const { + if (!supports_row_gradients()) { + return "cache row indices exceed the backward gather's index range"; + } + if (flash_attn && get_n_stream() != 1) { + return "flash attention backward requires a unified cache stream"; + } + return nullptr; +} diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 8e27fdf12e15..696db3546492 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -122,6 +122,8 @@ class llama_kv_cache : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, @@ -206,8 +208,8 @@ class llama_kv_cache : public llama_memory_i { // input API // - ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; - ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; + ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients = false) const; + ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients = false) const; ggml_tensor * build_input_k_rot(ggml_context * ctx) const; ggml_tensor * build_input_v_rot(ggml_context * ctx) const; @@ -389,8 +391,8 @@ class llama_kv_cache_context : public llama_memory_context_i { // create destination indices for each head of the current batch for where it would be written in the KV cache // the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but // helps understand the implementation logic of cpy_k and cpy_v - ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; - ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; + ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients = false) const; + ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients = false) const; ggml_tensor * build_input_k_rot(ggml_context * ctx) const; ggml_tensor * build_input_v_rot(ggml_context * ctx) const; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index e7f6e565dca1..bfba14c56454 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -674,3 +674,10 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } } + +const char * llama_memory_hybrid_idx::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = llama_memory_hybrid::training_error(batch, flash_attn)) { + return why; + } + return get_mem_idx() ? "indexed hybrid cache row gradients are not supported" : nullptr; +} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index a17e184209d1..f8a3f2dad905 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -48,6 +48,8 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index 06f7fd5428c4..21c39149a358 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -283,3 +283,10 @@ const llama_kv_cache_iswa_context * llama_memory_hybrid_iswa_context::get_attn() const llama_memory_recurrent_context * llama_memory_hybrid_iswa_context::get_recr() const { return static_cast(ctx_recr.get()); } + +const char * llama_memory_hybrid_iswa::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = mem_attn->training_error(batch, flash_attn)) { + return why; + } + return mem_recr->training_error(batch, flash_attn); +} diff --git a/src/llama-memory-hybrid-iswa.h b/src/llama-memory-hybrid-iswa.h index c9d3f9f57c50..693782a1c2ce 100644 --- a/src/llama-memory-hybrid-iswa.h +++ b/src/llama-memory-hybrid-iswa.h @@ -47,6 +47,8 @@ class llama_memory_hybrid_iswa : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6f..dfb4e9b947f5 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -277,3 +277,10 @@ const llama_kv_cache_context * llama_memory_hybrid_context::get_attn() const { const llama_memory_recurrent_context * llama_memory_hybrid_context::get_recr() const { return static_cast(ctx_recr.get()); } + +const char * llama_memory_hybrid::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = mem_attn->training_error(batch, flash_attn)) { + return why; + } + return mem_recr->training_error(batch, flash_attn); +} diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74991..2a92bdbaa81a 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -47,6 +47,8 @@ class llama_memory_hybrid : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7e..93c3ca873c12 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -1270,3 +1270,12 @@ int32_t llama_memory_recurrent_context::s_copy(int i) const { } return (int32_t)(idx * mem->size) + src0; } + +const char * llama_memory_recurrent::training_error(const llama_batch & batch, bool) const { + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.label[i] >= 0 && rollback_pending(batch.seq_id[i][0])) { + return "a labelled sequence has pending recurrent rollback"; + } + } + return nullptr; +} diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 8ef7420a987a..46b65f5dfa67 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -32,6 +32,8 @@ class llama_memory_recurrent : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645e..bc4935ede0a5 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -82,6 +82,9 @@ struct llama_memory_i { virtual ~llama_memory_i() = default; + // Read-only admission for a validated batch; nullptr means admitted. + virtual const char * training_error(const llama_batch & batch, bool flash_attn) const = 0; + // split the input batch into a set of ubatches and verify that they can fit into the cache // return a context object containing the ubatches and memory state required to process them // check the llama_memory_context_i::get_status() for the result diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp index e41870193ae8..963fb1a06b36 100644 --- a/tests/test-batch-alloc.cpp +++ b/tests/test-batch-alloc.cpp @@ -21,6 +21,8 @@ struct mock_memory : public llama_memory_i { llama_memory_context_ptr init_full() override { GGML_ASSERT(false && "not implemented"); } llama_memory_context_ptr init_update(llama_context *, bool) override { GGML_ASSERT(false && "not implemented"); } + const char * training_error(const llama_batch &, bool) const override { return nullptr; } + bool get_can_shift() const override { GGML_ASSERT(false && "not implemented"); } void clear(bool) override { GGML_ASSERT(false && "not implemented"); } From 44abe563f1e4d0453ce5d09bcf3c7fc8966d108e Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 15:17:13 -0500 Subject: [PATCH 14/43] test: compare multi-sequence label steps and isolate recompute Assisted-by: OpenAI Codex --- tests/test-walk-exact.cpp | 59 ++++++++++++++++++++++++++++++++++++++- 1 file changed, 58 insertions(+), 1 deletion(-) diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index f4cc172e50eb..468a3088e817 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -352,6 +352,52 @@ static void mixed_decode(const common_params & params, llama_model * model, cons llama_adapter_lora_free(adapter); } +static void multi_label_decode(const common_params & params, llama_model * model, const std::string & init, + const std::vector & tokens, bool flash, ggml_type cache_type) { + auto step = [&](const std::vector & sequences) { + auto cp = common_context_params_to_llama(params); + cp.n_ctx = 128; + cp.n_batch = cp.n_ubatch = 2; + cp.n_seq_max = 2; + cp.kv_unified = true; + cp.flash_attn_type = flash ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + cp.type_k = cp.type_v = cache_type; + llama_context * ctx = llama_init_from_model(model, cp); + GGML_ASSERT(ctx); + llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); + float scale = 1.0f; + GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); + llama_opt_params op{}; + op.param_filter = llama_opt_param_filter_all; + op.get_opt_pars = sgd_pars; + op.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; + op.adapter = adapter; + llama_opt_init(ctx, model, op); + const auto before = adapter_params(adapter); + llama_batch batch = llama_batch_init(2, 0, 1); + for (int seq : sequences) { + common_batch_add(batch, tokens[seq], 0, {seq}, true); + batch.label[batch.n_tokens - 1] = tokens[seq + 2]; + } + GGML_ASSERT(llama_decode(ctx, batch) == 0); + auto delta = adapter_params(adapter); + for (size_t i = 0; i < delta.size(); ++i) { + delta[i] -= before[i]; + } + llama_batch_free(batch); + llama_free(ctx); + llama_adapter_lora_free(adapter); + return delta; + }; + auto mean = step({0}); + const auto second = step({1}); + for (size_t i = 0; i < mean.size(); ++i) { + mean[i] = (mean[i] + second[i]) / 2; + } + const auto joint = step({0, 1}); + GGML_ASSERT(same_step(compare("two labelled sequences vs mean independent steps", joint, mean))); + GGML_ASSERT(same_step(compare("two labelled sequences with reversed row order", step({1, 0}), joint))); +} int main(int argc, char ** argv) { common_params params; if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { @@ -385,6 +431,8 @@ int main(int argc, char ** argv) { GGML_ASSERT(tokens.size() > REFUSAL_WINDOW + 1); if (getenv("WALK_MIXED_ONLY")) { + multi_label_decode(params, model, init, tokens, false, GGML_TYPE_F32); + multi_label_decode(params, model, init, tokens, true, GGML_TYPE_Q8_0); mixed_decode(params, model, init, tokens, false, GGML_TYPE_F32); for (ggml_type cache : { GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0 }) { mixed_decode(params, model, init, tokens, true, cache); @@ -439,7 +487,9 @@ int main(int argc, char ** argv) { return 0; } + const bool recompute_only = getenv("WALK_RECOMPUTE_ONLY") != nullptr; int failures = 0; + if (!recompute_only) { { // 1. every position labelled: chunk = window is one graph, the true gradient std::vector all(WINDOW + 1, 1); @@ -536,6 +586,8 @@ int main(int argc, char ** argv) { } } + } + { // 7. per-layer recompute under the exact walk. The chunk's activations, every layer's // kept for the backward pass, are what a large model's chunk graph is made of (Kimi's @@ -558,12 +610,15 @@ int main(int argc, char ** argv) { fprintf(stderr, "FAILED: recompute changed the exact walk's step\n"); ++failures; } - if (!(rebuilt_bytes < kept_bytes)) { + if (kept_bytes == 0 && rebuilt_bytes == 0) { + printf(" chunk graph reduction: unavailable (device-only allocation counter)\n"); + } else if (!(rebuilt_bytes < kept_bytes)) { fprintf(stderr, "FAILED: recompute did not shrink the exact walk's chunk graph\n"); ++failures; } } + if (!recompute_only) { { // 3. the host budget: a window whose accumulators and snapshots exceed it refuses by // name before anything runs, and the adapter is untouched @@ -692,6 +747,8 @@ int main(int argc, char ** argv) { g_window = WINDOW; } + } + std::remove(init.c_str()); llama_model_free(model); llama_backend_free(); From 07af30f988407376790f9c1617a89588b0087ffd Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 15:49:56 -0500 Subject: [PATCH 15/43] ggml: recount graph uses after checkpoint rewriting Assisted-by: OpenAI Codex --- ggml/src/ggml-impl.h | 1 + ggml/src/ggml-opt.cpp | 1 + ggml/src/ggml.c | 13 +++++++++++++ 3 files changed, 15 insertions(+) diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index 62b76abbcec9..e0577226b484 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -350,6 +350,7 @@ struct ggml_cgraph { // the slice does not have leafs or gradients // if you need the gradients, get them from the original graph struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph, int i0, int i1); +void ggml_graph_update_use_counts(struct ggml_cgraph * cgraph); // ggml-alloc.c: true if the operation can reuse memory from its sources GGML_API bool ggml_op_can_inplace(enum ggml_op op); diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index ad968b7794cb..9ad5c97d48cd 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -445,6 +445,7 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd, GGML_ASSERT((int) order.size() <= g->size && "recompute clones overflow the graph: raise the training node budget"); std::copy(order.begin(), order.end(), g->nodes); g->n_nodes = (int) order.size(); + ggml_graph_update_use_counts(g); } static ggml_tensor * map_tensor(std::map & tensor_map, ggml_context * ctx, ggml_tensor * tensor) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 17cbf45b8c24..e759bf741acf 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7942,6 +7942,19 @@ void ggml_graph_reset(struct ggml_cgraph * cgraph) { } } +void ggml_graph_update_use_counts(struct ggml_cgraph * cgraph) { + memset(cgraph->use_counts, 0, cgraph->visited_hash_set.size * sizeof(cgraph->use_counts[0])); + for (int i = 0; i < cgraph->n_nodes; ++i) { + for (int j = 0; j < GGML_MAX_SRC; ++j) { + struct ggml_tensor * src = cgraph->nodes[i]->src[j]; + if (src) { + const size_t pos = ggml_hash_find(&cgraph->visited_hash_set, src); + GGML_ASSERT(pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, pos)); + cgraph->use_counts[pos]++; + } + } + } +} void ggml_graph_clear(struct ggml_cgraph * cgraph) { cgraph->n_leafs = 0; cgraph->n_nodes = 0; From 343f369229e899bdb568689a479ca00a5c9b0327 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 16:47:06 -0500 Subject: [PATCH 16/43] llama: separate labelled head rows from public outputs Share row projection and bounded head replay between direct decode and the walk. Keep shadow scatter coordinates explicit and preserve replayed operator precision. Cover sparse outputs and head adapters against the independent optimizer objective. Assisted-by: OpenAI Codex --- common/common.cpp | 13 ++++ ggml/include/ggml-opt.h | 26 +++++++ ggml/src/ggml-opt.cpp | 118 +++++++++++++++++++++++------ src/llama-batch.cpp | 4 - src/llama-batch.h | 18 +++++ src/llama-context-opt.cpp | 45 +++-------- src/llama-context.cpp | 10 ++- src/llama-graph.cpp | 125 ++++++++++++++++++++++++++----- src/llama-graph.h | 18 +++-- src/llama-model.cpp | 5 ++ tests/test-opt-dynamic-accum.cpp | 86 ++++++++++++++++----- tests/test-walk-exact.cpp | 90 ++++++++++++++++++++++ 12 files changed, 450 insertions(+), 108 deletions(-) diff --git a/common/common.cpp b/common/common.cpp index 8147060448d7..bda497246329 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2401,8 +2401,21 @@ bool common_lora_write_fresh(const llama_model * model, const std::string & path std::vector shapes; const int32_t n_layer = llama_model_n_layer(model); const int32_t first = top_layers > 0 && top_layers < n_layer ? n_layer - top_layers : 0; + for (const auto & target : targets) { + if (target == "token_embd.weight") { + LOG_ERR("%s: embedding adapters require the embedding-specific layout\n", __func__); + return false; + } + const ggml_tensor * w = llama_model_get_tensor(model, target.c_str()); + if (w != nullptr && ggml_n_dims(w) == 2) { + shapes.push_back({ target, w->ne[0], w->ne[1] }); + } + } for (int32_t il = first; il < n_layer; ++il) { for (const auto & target : targets) { + if (llama_model_get_tensor(model, target.c_str()) != nullptr) { + continue; + } const std::string name = "blk." + std::to_string(il) + "." + target + ".weight"; const ggml_tensor * w = llama_model_get_tensor(model, name.c_str()); if (w != nullptr && ggml_n_dims(w) == 2) { diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index 707a6ee3dbe7..8de657ed299c 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h @@ -208,6 +208,32 @@ extern "C" { struct ggml_tensor * outputs, int64_t rows); + // Project the existing row-wise head before consumers such as sampling are built. + // Row indices address hidden's token rows. A null public_rows requests no public logits. + // Returns zero without changing gf when the head cannot be replayed. + GGML_API int ggml_opt_project_head( + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * hidden, + struct ggml_tensor * outputs, + struct ggml_tensor * shadow_rows, + struct ggml_tensor * public_rows, + struct ggml_tensor * label_rows, + struct ggml_tensor * public_shadow_rows, + struct ggml_tensor ** slice_shadow_rows, + int64_t rows, + struct ggml_tensor ** public_outputs, + struct ggml_tensor ** slices, + size_t slices_capacity); + + GGML_API void ggml_opt_prepare_alloc_slices( + ggml_opt_context_t opt_ctx, + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * inputs, + struct ggml_tensor ** slices, + size_t n_slices); + // allocate the next graph for evaluation, either forward or forward + backward // must be called exactly once prior to calling ggml_opt_eval // returns false when the graph could not be allocated (device memory): the graph is then NOT diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 9ad5c97d48cd..bc963c6b882d 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -1042,19 +1042,24 @@ void ggml_opt_prepare_alloc( opt_ctx->output_slices.clear(); } -int ggml_opt_prepare_alloc_head( - ggml_opt_context_t opt_ctx, +int ggml_opt_project_head( struct ggml_context * ctx_compute, struct ggml_cgraph * gf, - struct ggml_tensor * inputs, struct ggml_tensor * hidden, struct ggml_tensor * outputs, - int64_t rows) { - ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, outputs); - const int64_t n = hidden->ne[1]; - if (opt_ctx->loss_type != GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS || rows <= 0 || rows >= n || - !ggml_is_contiguous(hidden) || ggml_n_dims(hidden) > 2 || outputs->ne[1] != n) { - return 1; + struct ggml_tensor * shadow_rows, + struct ggml_tensor * public_rows, + struct ggml_tensor * label_rows, + struct ggml_tensor * public_shadow_rows, + struct ggml_tensor ** slice_shadow_rows, + int64_t rows, + struct ggml_tensor ** public_outputs, + struct ggml_tensor ** slices, + size_t slices_capacity) { + const int64_t n = label_rows ? label_rows->ne[0] : hidden->ne[1]; + if (rows <= 0 || n <= 0 || !ggml_is_contiguous(hidden) || ggml_n_dims(hidden) > 2 || + outputs->ne[1] != hidden->ne[1] || slices_capacity < (size_t) ((n + rows - 1) / rows)) { + return 0; } // the head: the nodes between hidden and outputs, in graph order std::unordered_set on_path = { hidden }; @@ -1069,15 +1074,27 @@ int ggml_opt_prepare_alloc_head( continue; } const bool replayable = node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL || - node->op == GGML_OP_ADD || node->op == GGML_OP_SCALE; - if (!replayable || node->view_src != nullptr) { - return 1; // a head this does not know how to replay per slice: the whole outputs + node->op == GGML_OP_ADD || node->op == GGML_OP_SCALE || + (shadow_rows && (node->op == GGML_OP_DUP || node->op == GGML_OP_GET_ROWS || node->op == GGML_OP_SET_ROWS)); + if (!replayable || (node->view_src != nullptr && node->op != GGML_OP_SET_ROWS)) { + return 0; // unsupported head + } + if ((node->op == GGML_OP_GET_ROWS && node->src[1] != shadow_rows) || + (node->op == GGML_OP_SET_ROWS && node->src[1] != shadow_rows)) { + return 0; + } + if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { + for (int k = 0; k < 2; ++k) { + if (!on_path.count(node->src[k]) && node->src[k]->ne[1] != 1) { + return 0; + } + } } on_path.insert(node); head.push_back(node); } if (head.empty() || head.back() != outputs) { - return 1; + return 0; } // nothing else in the graph may read the head (it is about to leave the graph) for (int i = 0; i < gf->n_nodes; ++i) { @@ -1087,7 +1104,7 @@ int ggml_opt_prepare_alloc_head( } for (int k = 0; k < GGML_MAX_SRC; ++k) { if (node->src[k] != nullptr && node->src[k] != hidden && on_path.count(node->src[k]) > 0) { - return 1; + return 0; } } } @@ -1098,15 +1115,31 @@ int ggml_opt_prepare_alloc_head( } } gf->n_nodes = j; - // the head again, one slice of rows at a time - for (int64_t r0 = 0; r0 < n; r0 += rows) { - const int64_t nr = std::min(rows, n - r0); + auto replay = [&](ggml_tensor * selected, ggml_tensor * local_shadow_rows, int64_t row) { std::unordered_map as_slice; - as_slice[hidden] = ggml_view_2d(ctx_compute, hidden, hidden->ne[0], nr, hidden->nb[1], r0*hidden->nb[1]); + as_slice[hidden] = selected; + if (shadow_rows) { + as_slice[shadow_rows] = local_shadow_rows; + } for (ggml_tensor * node : head) { auto src = [&](int k) { const auto it = as_slice.find(node->src[k]); return it != as_slice.end() ? it->second : node->src[k]; }; ggml_tensor * t = nullptr; + if (node->op == GGML_OP_SET_ROWS && shadow_rows && !local_shadow_rows) { + as_slice[node] = src(2); + continue; + } + bool missing = false; + for (int k = 0; k < GGML_MAX_SRC; ++k) { + missing |= node->src[k] != nullptr && src(k) == nullptr; + } + if (missing) { + as_slice[node] = nullptr; + continue; + } switch (node->op) { + case GGML_OP_DUP: t = ggml_dup(ctx_compute, src(0)); break; + case GGML_OP_GET_ROWS: t = ggml_get_rows(ctx_compute, src(0), src(1)); break; + case GGML_OP_SET_ROWS: t = ggml_set_rows(ctx_compute, src(2), src(0), src(1)); break; case GGML_OP_MUL_MAT: t = ggml_mul_mat(ctx_compute, src(0), src(1)); break; case GGML_OP_MUL: t = ggml_mul (ctx_compute, src(0), src(1)); break; case GGML_OP_ADD: t = ggml_add (ctx_compute, src(0), src(1)); break; @@ -1118,14 +1151,57 @@ int ggml_opt_prepare_alloc_head( } break; default: GGML_ABORT("not a head op"); } - ggml_format_name(t, "%s-rows%lld", node->name, (long long) r0); + ggml_set_op_params(t, node->op_params, sizeof(node->op_params)); + ggml_format_name(t, "%s-rows%lld", node->name, (long long) row); as_slice[node] = t; } - opt_ctx->output_slices.push_back(as_slice.at(outputs)); + return as_slice.at(outputs); + }; + *public_outputs = nullptr; + if (public_rows) { + *public_outputs = replay(ggml_get_rows(ctx_compute, hidden, public_rows), + public_shadow_rows, -1); + ggml_build_forward_expand(gf, *public_outputs); + } + ggml_tensor * selected = label_rows ? ggml_get_rows(ctx_compute, hidden, label_rows) : hidden; + size_t count = 0; + for (int64_t r0 = 0; r0 < n; r0 += rows) { + const int64_t nr = std::min(rows, n - r0); + slices[count] = replay( + ggml_view_2d(ctx_compute, selected, selected->ne[0], nr, selected->nb[1], r0 * selected->nb[1]), + slice_shadow_rows ? slice_shadow_rows[count] : nullptr, r0); + ++count; } - return (int) opt_ctx->output_slices.size(); + ggml_graph_update_use_counts(gf); + return (int) count; +} + +void ggml_opt_prepare_alloc_slices( + ggml_opt_context_t opt_ctx, ggml_context * ctx_compute, ggml_cgraph * gf, + ggml_tensor * inputs, ggml_tensor ** slices, size_t n_slices) { + GGML_ASSERT(n_slices > 0); + ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, slices[0]); + opt_ctx->output_slices.assign(slices, slices + n_slices); } +int ggml_opt_prepare_alloc_head( + ggml_opt_context_t opt_ctx, ggml_context * ctx_compute, ggml_cgraph * gf, + ggml_tensor * inputs, ggml_tensor * hidden, ggml_tensor * outputs, int64_t rows) { + ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, outputs); + const int64_t n = hidden->ne[1]; + if (opt_ctx->loss_type != GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS || rows <= 0 || rows >= n) { + return 1; + } + std::vector slices((n + rows - 1) / rows); + ggml_tensor * public_outputs = nullptr; + const int count = ggml_opt_project_head(ctx_compute, gf, hidden, outputs, nullptr, nullptr, nullptr, nullptr, nullptr, + rows, &public_outputs, slices.data(), slices.size()); + if (count == 0) { + return 1; + } + ggml_opt_prepare_alloc_slices(opt_ctx, ctx_compute, gf, inputs, slices.data(), count); + return count; +} void ggml_opt_set_alloc_budget(ggml_opt_context_t opt_ctx, size_t bytes) { opt_ctx->alloc_budget = bytes; } diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 7b873d81d6cc..5ec36436c7e8 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -187,10 +187,6 @@ bool llama_batch_allocr::init( } continue; } - if (batch.logits[i] == 0) { - LLAMA_LOG_ERROR("%s: position %d has a label but is not an output: a training row must be output\n", __func__, i); - return false; - } if (batch.n_seq_id[i] != 1) { LLAMA_LOG_ERROR("%s: position %d has a label and %d sequences: a training row belongs to one sequence\n", __func__, i, batch.n_seq_id[i]); return false; diff --git a/src/llama-batch.h b/src/llama-batch.h index ef1f9908f24f..6430ac1d457a 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -13,6 +13,24 @@ // keep this struct lightweight struct llama_ubatch { + struct row_map { + std::vector public_rows; + std::vector label_rows; + }; + + row_map rows() const { + row_map result; + for (uint32_t i = 0; i < n_tokens; ++i) { + if (output && output[i]) { + result.public_rows.push_back((int32_t) i); + } + if (label && label[i] >= 0) { + result.label_rows.push_back((int32_t) i); + } + } + return result; + } + bool equal_seqs() const { return b_equal_seqs != 0; } diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 41c6ce4ed089..877ae4143692 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -322,7 +322,7 @@ bool llama_context::opt_walk::rewind(uint32_t p_pop, const std::vector bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { llama_walk_fill(batch, tokens, pos_ctx, n_chunk); for (uint32_t pos_batch = 0; pos_batch < n_chunk; ++pos_batch) { - batch.logits[pos_batch] = labels[pos_ctx + pos_batch] >= 0; + batch.logits[pos_batch] = false; batch.label [pos_batch] = labels[pos_ctx + pos_batch]; } @@ -346,9 +346,6 @@ bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { return true; } -// what one slice of the training head's logits may hold (see opt_ubatch) -static constexpr size_t OPT_HEAD_SLICE_BYTES = 64u*1024*1024; - bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch) { const uint32_t n_labels = ubatch.n_labels(); opt_walk * walk = opt_walk_active; @@ -365,38 +362,18 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l }; ctx_compute_opt = ggml_init(params); } - // THE HEAD A SLICE OF ROWS AT A TIME: the logits of a chunk and their gradient were the - // largest tensors of its graph ([vocab x rows] F32 each: 485 MiB apiece at a 248k vocabulary - // and 512 rows, the 4B at window 8192, 2026-10-14), and their size does not shrink with the - // model. A slice holds OPT_HEAD_SLICE_BYTES of logits. - const int64_t head_rows = std::max(1, (int64_t) (OPT_HEAD_SLICE_BYTES / (sizeof(float) * (size_t) model.vocab.n_tokens()))); - ggml_tensor * loss_rows = nullptr; - std::vector rows; std::vector targets; - rows.reserve(n_labels); targets.reserve(n_labels); - int32_t output_row = 0; - for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { - if (ubatch.output[i]) { - if (ubatch.label[i] >= 0) { - rows.push_back(output_row); - targets.push_back(ubatch.label[i]); - } - ++output_row; - } + for (int32_t row : ubatch.rows().label_rows) { + targets.push_back(ubatch.label[row]); } - if (walk != nullptr && n_labels == n_outputs) { - ggml_opt_prepare_alloc_head(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_embd(), res->get_logits(), head_rows); - } else { - // Direct decode returns the original outputs; only the objective selects labelled rows. - ggml_tensor * objective = res->get_logits(); - if (n_labels != n_outputs) { - loss_rows = ggml_new_tensor_1d(ctx_compute_opt, GGML_TYPE_I32, n_labels); - ggml_set_input(loss_rows); - objective = ggml_get_rows(ctx_compute_opt, objective, loss_rows); - } - ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), objective); + if (res->t_loss_slices.empty()) { + ggml_free(ctx_compute_opt); + opt_fail("training head was not projected into bounded label slices"); + return false; } + ggml_opt_prepare_alloc_slices(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), + res->t_loss_slices.data(), res->t_loss_slices.size()); if (xc != nullptr) { ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); for (const auto & io : res->t_walk) { @@ -431,9 +408,7 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l // order the logits rows take) struct ggml_tensor * labels_t = ggml_opt_labels(opt_ctx); GGML_ASSERT(labels_t->type == GGML_TYPE_I32 && labels_t->ne[0] == (int64_t) n_labels); - if (loss_rows) { - ggml_backend_tensor_set(loss_rows, rows.data(), 0, ggml_nbytes(loss_rows)); - } + ggml_backend_tensor_set(labels_t, targets.data(), 0, ggml_nbytes(labels_t)); } if (xc != nullptr) { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a79c04585427..cb95a0333630 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1719,6 +1719,14 @@ int llama_context::decode(const llama_batch & batch_inp) { } if (balloc->get_n_labels() > 0) { + if (!model.output) { + LLAMA_LOG_ERROR("%s: labelled training requires a language-model output weight\n", __func__); + return -1; + } + if (cparams.pooling_type != LLAMA_POOLING_TYPE_NONE) { + LLAMA_LOG_ERROR("%s: labelled training requires a token-wise output head, not pooled outputs\n", __func__); + return -1; + } if (!opt_ctx) { LLAMA_LOG_ERROR("%s: labelled rows require an initialized optimizer\n", __func__); return -1; @@ -2522,7 +2530,7 @@ llm_graph_params llama_context::graph_params( /*.mctx =*/ mctx, /*.cross =*/ &cross, /*.samplers =*/ sampling.samplers, - /*.n_outputs =*/ n_outputs, + /*.n_outputs =*/ (cparams.training || (opt_shadow && ubatch.n_labels() > 0)) ? ubatch.n_tokens : n_outputs, /*.cb =*/ graph_get_cb(), /*.res =*/ res, }; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index e6521d2c9c2d..66f45effa4eb 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -5,6 +5,7 @@ #include "llama-batch.h" #include "llama-cparams.h" #include "llama-sampler.h" +#include "ggml-opt.h" #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" @@ -223,17 +224,35 @@ void llm_graph_input_out_ids::set_input(const llama_ubatch * ubatch) { } } -void llm_graph_input_label_rows::set_input(const llama_ubatch * ubatch) { - GGML_ASSERT(rows && ggml_backend_buffer_is_host(rows->buffer)); - int32_t * data = (int32_t *) rows->data; - int64_t n = 0; - for (uint32_t i = 0; i < ubatch->n_tokens; ++i) { - if (ubatch->label != nullptr && ubatch->label[i] >= 0) { - GGML_ASSERT(n < rows->ne[0]); - data[n++] = (int32_t) i; +void llm_graph_input_training_rows::set_input(const llama_ubatch * ubatch) { + const auto map = ubatch->rows(); + if (public_rows) { + GGML_ASSERT(public_rows->ne[0] == (int64_t) map.public_rows.size()); + ggml_backend_tensor_set(public_rows, map.public_rows.data(), 0, ggml_nbytes(public_rows)); + } + if (label_rows) { + GGML_ASSERT(label_rows->ne[0] == (int64_t) map.label_rows.size()); + ggml_backend_tensor_set(label_rows, map.label_rows.data(), 0, ggml_nbytes(label_rows)); + } + if (shadow_rows && shadow_rows->buffer) { + ggml_backend_tensor_set(shadow_rows, map.label_rows.data(), 0, ggml_nbytes(shadow_rows)); + } + if (public_shadow_rows && public_shadow_rows->buffer) { + std::vector local; + for (size_t i = 0; i < map.public_rows.size(); ++i) { + if (ubatch->label[map.public_rows[i]] >= 0) { + local.push_back((int32_t) i); + } + } + ggml_backend_tensor_set(public_shadow_rows, local.data(), 0, ggml_nbytes(public_shadow_rows)); + } + for (auto * rows : slice_shadow_rows) { + if (rows->buffer) { + std::vector local(rows->ne[0]); + std::iota(local.begin(), local.end(), 0); + ggml_backend_tensor_set(rows, local.data(), 0, ggml_nbytes(rows)); } } - GGML_ASSERT(n == rows->ne[0]); } bool llm_graph_input_out_ids::can_reuse(const llm_graph_params & params) { @@ -1346,6 +1365,7 @@ void llm_graph_result::reset() { t_inp_tokens = nullptr; t_inp_embd = nullptr; t_logits = nullptr; + t_loss_slices.clear(); t_embd = nullptr; t_embd_pooled = nullptr; t_h_nextn = nullptr; @@ -2448,24 +2468,89 @@ bool llm_graph_context::lora_on(const llama_adapter_lora * adapter) const { return adapter != shadow || n_shadow_rows > 0; } -ggml_tensor * llm_graph_context::shadow_add(ggml_tensor * base, ggml_tensor * term) const { - if (inp_label_rows == nullptr) { - auto inp = std::make_unique(); - inp->rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_shadow_rows); - ggml_set_input(inp->rows); - ggml_set_name(inp->rows, "label_rows"); - inp_label_rows = inp->rows; +llm_graph_input_training_rows * llm_graph_context::training_rows() const { + if (!inp_training_rows) { + auto inp = std::make_unique(); + inp_training_rows = inp.get(); res->add_input(std::move(inp)); } + return inp_training_rows; +} + +ggml_tensor * llm_graph_context::shadow_add(ggml_tensor * base, ggml_tensor * term) const { + if (training_rows()->shadow_rows == nullptr) { + auto * inp = training_rows(); + inp->shadow_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_shadow_rows); + ggml_set_input(inp->shadow_rows); + ggml_set_name(inp->shadow_rows, "label_rows"); + } // one row per token: the tokens are base's outermost dimension GGML_ASSERT(ggml_is_contiguous(base) && (n_tokens == 1 || base->ne[ggml_n_dims(base) - 1] == n_tokens)); const int64_t n = ggml_nelements(base) / n_tokens; - ggml_tensor * b2 = ggml_reshape_2d(ctx0, base, n, n_tokens); - ggml_tensor * t2 = ggml_reshape_2d(ctx0, ggml_cont(ctx0, term), n, n_tokens); - ggml_tensor * sel = ggml_add(ctx0, ggml_get_rows(ctx0, b2, inp_label_rows), ggml_get_rows(ctx0, t2, inp_label_rows)); - return ggml_reshape(ctx0, ggml_set_rows(ctx0, ggml_dup(ctx0, b2), sel, inp_label_rows), base); + ggml_tensor * b2 = ggml_n_dims(base) <= 2 ? base : ggml_reshape_2d(ctx0, base, n, n_tokens); + ggml_tensor * t2 = ggml_n_dims(term) <= 2 ? term : ggml_reshape_2d(ctx0, ggml_cont(ctx0, term), n, n_tokens); + ggml_tensor * sel = ggml_add(ctx0, ggml_get_rows(ctx0, b2, training_rows()->shadow_rows), ggml_get_rows(ctx0, t2, training_rows()->shadow_rows)); + ggml_tensor * updated = ggml_set_rows(ctx0, ggml_dup(ctx0, b2), sel, training_rows()->shadow_rows); + return ggml_n_dims(base) <= 2 ? updated : ggml_reshape(ctx0, updated, base); } +bool llm_graph_context::build_training_head() const { + if (!cparams.training && n_shadow_rows == 0) { + return true; + } + if (!res->t_embd || !res->t_logits || pooling_type != LLAMA_POOLING_TYPE_NONE) { + return false; + } + const auto map = ubatch.rows(); + auto * input = training_rows(); + if (cparams.training) { + input->label_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, map.label_rows.size()); + ggml_set_input(input->label_rows); + } + if (!map.public_rows.empty()) { + input->public_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, map.public_rows.size()); + ggml_set_input(input->public_rows); + } + // Preserve the existing training head's allocation budget. + constexpr size_t head_slice_bytes = 64u*1024*1024; + const int64_t rows = std::max(1, head_slice_bytes / (sizeof(float) * res->t_logits->ne[0])); + const int64_t loss_rows = cparams.training ? map.label_rows.size() : n_tokens; + res->t_loss_slices.resize((loss_rows + rows - 1) / rows); + if (input->shadow_rows) { + size_t public_labels = 0; + for (int32_t row : map.public_rows) { + public_labels += ubatch.label[row] >= 0; + } + if (public_labels) { + input->public_shadow_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, public_labels); + ggml_set_input(input->public_shadow_rows); + } + for (int64_t r0 = 0; r0 < loss_rows; r0 += rows) { + auto * local = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, std::min(rows, loss_rows - r0)); + ggml_set_input(local); + input->slice_shadow_rows.push_back(local); + } + } + ggml_tensor * public_logits = nullptr; + const int count = ggml_opt_project_head(ctx0, gf, res->t_embd, res->t_logits, input->shadow_rows, + input->public_rows, input->label_rows, input->public_shadow_rows, + input->slice_shadow_rows.empty() ? nullptr : input->slice_shadow_rows.data(), rows, &public_logits, + res->t_loss_slices.data(), res->t_loss_slices.size()); + if (!count) { + res->t_loss_slices.clear(); + return false; + } + res->t_logits = public_logits; + if (cparams.embeddings && input->public_rows) { + res->t_embd = ggml_get_rows(ctx0, res->t_embd, input->public_rows); + ggml_build_forward_expand(gf, res->t_embd); + } + if (cparams.embeddings_nextn && cparams.embeddings_nextn_masked && input->public_rows) { + res->t_h_nextn = ggml_get_rows(ctx0, res->t_h_nextn, input->public_rows); + ggml_build_forward_expand(gf, res->t_h_nextn); + } + return true; +} ggml_tensor * llm_graph_context::build_inp_pos() const { auto inp = std::make_unique(hparams.n_pos_per_embd()); diff --git a/src/llama-graph.h b/src/llama-graph.h index d63901ff10a1..c018ffd0e6b5 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -233,15 +233,14 @@ class llm_graph_input_out_ids : public llm_graph_input_i { const uint32_t n_outputs; }; -// the ubatch's labelled rows, in order: where the shadow adapter applies (shadow_add) -class llm_graph_input_label_rows : public llm_graph_input_i { +class llm_graph_input_training_rows : public llm_graph_input_i { public: - llm_graph_input_label_rows() = default; - virtual ~llm_graph_input_label_rows() = default; - void set_input(const llama_ubatch * ubatch) override; - - ggml_tensor * rows = nullptr; // I32 [n_shadow_rows] + ggml_tensor * public_rows = nullptr; + ggml_tensor * label_rows = nullptr; + ggml_tensor * shadow_rows = nullptr; + ggml_tensor * public_shadow_rows = nullptr; + std::vector slice_shadow_rows; }; class llm_graph_input_mean : public llm_graph_input_i { @@ -951,6 +950,7 @@ class llm_graph_result { ggml_tensor * t_inp_tokens = nullptr; ggml_tensor * t_inp_embd = nullptr; // [n_embd_inp, n_tokens] ggml_tensor * t_logits = nullptr; + std::vector t_loss_slices; ggml_tensor * t_embd = nullptr; ggml_tensor * t_embd_pooled = nullptr; ggml_tensor * t_h_nextn = nullptr; // [n_embd, n_outputs] hidden state before final output norm @@ -1086,7 +1086,7 @@ struct llm_graph_context { llm_graph_result * res; mutable ggml_tensor * inp_train_rows = nullptr; // build_inp_train_rows, once per graph - mutable ggml_tensor * inp_label_rows = nullptr; // shadow_add, once per graph + mutable llm_graph_input_training_rows * inp_training_rows = nullptr; ggml_context * ctx0 = nullptr; ggml_cgraph * gf = nullptr; @@ -1452,6 +1452,8 @@ struct llm_graph_context { // void build_sampling() const; + bool build_training_head() const; + llm_graph_input_training_rows * training_rows() const; // // dense (out) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 9d34a2d28484..63fa565e7843 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2646,6 +2646,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { std::unique_ptr llm = build_arch_graph(params); + if (!llm->build_training_head()) { + LLAMA_LOG_ERROR("%s: training head is not a supported row-wise projection\n", __func__); + return nullptr; + } + // add on pooling layer llm->build_pooling(cls, cls_b, cls_out, cls_out_b, cls_norm); diff --git a/tests/test-opt-dynamic-accum.cpp b/tests/test-opt-dynamic-accum.cpp index 6b392968c8ec..f72dce6ea54d 100644 --- a/tests/test-opt-dynamic-accum.cpp +++ b/tests/test-opt-dynamic-accum.cpp @@ -279,45 +279,93 @@ static void run_cache_view_gradient(ggml_type type) { ggml_backend_free(cpu); } -static void run_label_projection() { +static void run_label_projection(bool projected) { ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); CHECK(cpu != nullptr); ggml_backend_t backends[] = { cpu }; ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 1, GGML_DEFAULT_GRAPH_SIZE, false, true); ggml_init_params sp = { 8 * ggml_tensor_overhead(), nullptr, true }; ggml_context * storage = ggml_init(sp); - ggml_tensor * weights = ggml_new_tensor_2d(storage, GGML_TYPE_F32, 2, 3); + ggml_tensor * weights = ggml_new_tensor_2d(storage, GGML_TYPE_F32, 2, 4); + ggml_tensor * identity = ggml_new_tensor_2d(storage, GGML_TYPE_F32, 2, 2); ggml_set_param(weights); ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(storage, cpu); - const float before[] = { 8.0f, -8.0f, 0.0f, 0.0f, -8.0f, 8.0f }; + const float before[] = { 8.0f, -8.0f, 0.5f, -0.5f, -8.0f, 8.0f, 0.25f, 0.75f }; + const float eye[] = { 1, 0, 0, 1 }; ggml_backend_tensor_set(weights, before, 0, sizeof(before)); + ggml_backend_tensor_set(identity, eye, 0, sizeof(eye)); ggml_opt_params params = ggml_opt_default_params(sched, GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS); params.optimizer = GGML_OPT_OPTIMIZER_TYPE_SGD; params.get_opt_pars = sgd_pars; ggml_opt_context_t opt = ggml_opt_init(params); ggml_init_params cp = { GGML_DEFAULT_GRAPH_SIZE * ggml_tensor_overhead() + 4 * ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), nullptr, true }; ggml_context * graph_ctx = ggml_init(cp); - ggml_tensor * logits = ggml_scale(graph_ctx, weights, 1.0f); - ggml_set_output(logits); - ggml_tensor * rows = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 1); + ggml_tensor * base = ggml_mul_mat(graph_ctx, identity, weights); + ggml_mul_mat_set_prec(base, GGML_PREC_F32); + ggml_tensor * rows = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 2); + ggml_tensor * public_rows = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 3); ggml_set_input(rows); - ggml_tensor * objective = ggml_get_rows(graph_ctx, logits, rows); + ggml_set_input(public_rows); + ggml_tensor * selected = ggml_add(graph_ctx, ggml_get_rows(graph_ctx, base, rows), + ggml_get_rows(graph_ctx, ggml_scale(graph_ctx, base, 0.25f), rows)); + ggml_tensor * logits = ggml_set_rows(graph_ctx, ggml_dup(graph_ctx, base), selected, rows); ggml_cgraph * gf = ggml_new_graph_custom(graph_ctx, GGML_DEFAULT_GRAPH_SIZE, true); ggml_build_forward_expand(gf, logits); - ggml_opt_prepare_alloc(opt, graph_ctx, gf, rows, objective); + ggml_tensor * public_logits = nullptr; + ggml_tensor * public_shadow = nullptr; + ggml_tensor * local_shadow[2] = {}; + if (projected) { + public_shadow = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 1); + ggml_set_input(public_shadow); + for (auto & local : local_shadow) { + local = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 1); + ggml_set_input(local); + } + ggml_tensor * slices[2]; + CHECK(ggml_opt_project_head(graph_ctx, gf, weights, logits, rows, public_rows, rows, + public_shadow, local_shadow, 1, &public_logits, slices, 2) == 2); + ggml_opt_prepare_alloc_slices(opt, graph_ctx, gf, rows, slices, 2); + } else { + public_logits = ggml_get_rows(graph_ctx, logits, public_rows); + ggml_build_forward_expand(gf, public_logits); + ggml_opt_prepare_alloc(opt, graph_ctx, gf, rows, ggml_get_rows(graph_ctx, logits, rows)); + } + ggml_set_output(public_logits); + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + auto * node = ggml_graph_node(gf, i); + if (node->op == GGML_OP_MUL_MAT) { + CHECK(node->op_params[0] == GGML_PREC_F32); + } + } CHECK(ggml_opt_alloc(opt, true)); - const int32_t row = 1, target = 0; - ggml_backend_tensor_set(rows, &row, 0, sizeof(row)); - ggml_backend_tensor_set(ggml_opt_labels(opt), &target, 0, sizeof(target)); + const int32_t labels[] = { 1, 3 }, public_ids[] = { 2, 1, 0 }, targets[] = { 0, 1 }; + ggml_backend_tensor_set(rows, labels, 0, sizeof(labels)); + ggml_backend_tensor_set(public_rows, public_ids, 0, sizeof(public_ids)); + if (projected) { + const int32_t public_label = 1, local_label = 0; + ggml_backend_tensor_set(public_shadow, &public_label, 0, sizeof(public_label)); + for (auto * local : local_shadow) { + ggml_backend_tensor_set(local, &local_label, 0, sizeof(local_label)); + } + } + ggml_backend_tensor_set(ggml_opt_labels(opt), targets, 0, sizeof(targets)); ggml_opt_eval(opt, nullptr); CHECK(ggml_opt_refusal(opt)[0] == '\0'); - float after[6], served[6]; + float after[8], served[6]; ggml_backend_tensor_get(weights, after, 0, sizeof(after)); - ggml_backend_tensor_get(logits, served, 0, sizeof(served)); - for (int i = 0; i < 6; ++i) { - const float delta = i == 2 ? LR * 0.5f : i == 3 ? -LR * 0.5f : 0.0f; - CHECK(std::fabs(after[i] - before[i] - delta) < 1e-6f); - CHECK(served[i] == before[i]); + ggml_backend_tensor_get(public_logits, served, 0, sizeof(served)); + for (int row = 0; row < 4; ++row) { + const bool labelled = row == 1 || row == 3; + const float p0 = 1.0f / (1.0f + std::exp(1.25f * (before[2*row + 1] - before[2*row]))); + const float delta = labelled ? LR * 1.25f * ((row == 1 ? 1.0f : 0.0f) - p0) / 2 : 0; + CHECK(std::fabs(after[2*row] - before[2*row] - delta) < 1e-6f); + CHECK(std::fabs(after[2*row + 1] - before[2*row + 1] + delta) < 1e-6f); + } + for (int row = 0; row < 3; ++row) { + const float scale = public_ids[row] == 1 ? 1.25f : 1; + for (int col = 0; col < 2; ++col) { + CHECK(served[2*row + col] == before[2*public_ids[row] + col] * scale); + } } ggml_opt_free(opt); ggml_free(graph_ctx); @@ -326,13 +374,13 @@ static void run_label_projection() { ggml_backend_sched_free(sched); ggml_backend_free(cpu); } - int main() { ggml_backend_load_all(); run(1); run(2); run_manual(); - run_label_projection(); + run_label_projection(false); + run_label_projection(true); for (ggml_type type : { GGML_TYPE_F16, GGML_TYPE_Q8_0 }) { run_cache_view_gradient(type); run_out_prod_broadcast(type, false); diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 468a3088e817..affa93f8bab5 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -398,6 +398,80 @@ static void multi_label_decode(const common_params & params, llama_model * model GGML_ASSERT(same_step(compare("two labelled sequences vs mean independent steps", joint, mean))); GGML_ASSERT(same_step(compare("two labelled sequences with reversed row order", step({1, 0}), joint))); } +static void sparse_head_decode(const common_params & params, llama_model * model, const std::string & init, + const std::vector & tokens) { + struct result { + std::vector delta, logits; + std::vector neighbour; + }; + auto step = [&](int public_mode, bool reversed, bool train = true) { + auto cp = common_context_params_to_llama(params); + cp.n_ctx = 128; + cp.n_batch = cp.n_ubatch = 4; + cp.n_seq_max = 2; + cp.kv_unified = true; + cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + auto * ctx = llama_init_from_model(model, cp); + auto * adapter = llama_adapter_lora_init(model, init.c_str()); + GGML_ASSERT(ctx && adapter); + float scale = 1.0f; + GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); + llama_opt_params op{}; + op.param_filter = llama_opt_param_filter_all; + op.get_opt_pars = sgd_pars; + op.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; + op.adapter = adapter; + if (train) { + llama_opt_init(ctx, model, op); + } + const auto before = adapter_params(adapter); + llama_batch batch = llama_batch_init(3, 0, 1); + int neighbour_row = -1; + const std::vector order = reversed ? std::vector{2, 0, 1} : std::vector{0, 1, 2}; + for (int row : order) { + common_batch_add(batch, tokens[row], row == 1 ? 1 : 0, {row == 2 ? 1 : 0}, + public_mode == 2 || (public_mode == 1 && row == 2)); + if (train && row == 1) { + batch.label[batch.n_tokens - 1] = tokens[3]; + } + if (row == 2) { + neighbour_row = batch.n_tokens - 1; + } + } + GGML_ASSERT(llama_decode(ctx, batch) == 0); + printf("sparse head allocation: train=%d public=%d reversed=%d cache=%s graph_bytes=%zu\n", + train, public_mode, reversed, ggml_type_name(cp.type_k), llama_opt_graph_bytes(ctx)); + result out; + out.delta = adapter_params(adapter); + for (size_t i = 0; i < before.size(); ++i) { + out.delta[i] -= before[i]; + } + if (public_mode > 0) { + const float * logits = llama_get_logits_ith(ctx, neighbour_row); + GGML_ASSERT(logits); + out.logits.assign(logits, logits + llama_vocab_n_tokens(llama_model_get_vocab(model))); + } + out.neighbour = sequence_state(ctx, 1); + llama_batch_free(batch); + llama_free(ctx); + llama_adapter_lora_free(adapter); + return out; + }; + const auto dense = step(2, false); + const auto serving = step(1, false, false); + GGML_ASSERT(serving.logits == dense.logits); + GGML_ASSERT(serving.neighbour == dense.neighbour); + for (bool reversed : {false, true}) { + for (int public_mode : {0, 1}) { + const auto sparse = step(public_mode, reversed); + GGML_ASSERT(same_step(compare("sparse public rows vs dense objective", sparse.delta, dense.delta))); + GGML_ASSERT(public_mode == 0 || sparse.logits == dense.logits); + GGML_ASSERT(sparse.neighbour == dense.neighbour); + } + } + printf("sparse head: label-only rows, public logits, neighbour bytes and row reordering OK\n"); +} + int main(int argc, char ** argv) { common_params params; if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { @@ -430,6 +504,22 @@ int main(int argc, char ** argv) { llama_free(tok_ctx); GGML_ASSERT(tokens.size() > REFUSAL_WINDOW + 1); + if (getenv("WALK_HEAD_ONLY")) { + sparse_head_decode(params, model, init, tokens); + if (llama_model_get_tensor(model, "output.weight")) { + const auto head_init = init + ".head.gguf"; + GGML_ASSERT(common_lora_write_fresh(model, head_init, 8, 16, { "output.weight" }, 42)); + sparse_head_decode(params, model, head_init, tokens); + std::remove(head_init.c_str()); + } else { + printf("output adapter: unavailable (model has no independent output.weight tensor)\n"); + } + std::remove(init.c_str()); + llama_model_free(model); + llama_backend_free(); + return 0; + } + if (getenv("WALK_MIXED_ONLY")) { multi_label_decode(params, model, init, tokens, false, GGML_TYPE_F32); multi_label_decode(params, model, init, tokens, true, GGML_TYPE_Q8_0); From eda8574b6ad0210655412f780f5058dfb4cdd712 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 17:36:23 -0500 Subject: [PATCH 17/43] ggml: acc_or_set's zero is a node, not a full-size leaf PR64 (999d682aa) built the zero tensor of ggml_acc_or_set with ggml_fill over a new tensor, to remove 0 * inf = NaN. A tensor made by ggml_new_tensor is a graph leaf, and leaves stay allocated for the whole graph, so every such zero (cache-window-sized ones included) stayed resident. The zero is now a one-element filled leaf repeated to src's shape: the same NaN-free value, freed after its use. Measured (M5 Metal, qwen3.5-4b hybrid, WALK_MEASURE window 8192, chunk 512, recompute, 2026-10-14): chunk graph 2,118.6 MiB at PR64's head -> 1,286.6 MiB (the base 88c5c1e26 measures 1,305.6); explicit path 2,324.8 -> 1,551.8. Bisected: 88c5c1e26 + PR64's ggml.c alone 2,118.6; reverting only this hunk 1,305.6. test-walk-exact OK on qwen2.5-0.5b (22:25:13-22:26:04Z) and the 4B (22:26:04-22:36:01Z): every exact row 1.0000/1.0000, cache byte-identical after the walk, 4B chunk graph 666.8 MiB with recompute (PR64 head: 904.1). test-opt-dynamic-accum OK; test-backend-ops grad SET_ROWS and FLASH_ATTN_EXT on CPU: 0 failures. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml.c | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index e759bf741acf..f08bf6aa0c7e 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6667,7 +6667,11 @@ static void ggml_acc_or_set( } cgraph->grads[isrc] = ggml_acc_impl(ctx, acc_into, tensor, nb1, nb2, nb3, offset, cgraph->grad_accs[isrc]); } else { - struct ggml_tensor * a_zero = ggml_fill(ctx, ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, src->ne), 0.0f); + // zeros by fill (0 * inf would be NaN), built from a one-element leaf repeated to src's + // shape: a full-size leaf stays allocated for the whole graph (+813 MiB on the 4B at + // window 8192, 2026-10-14), a node is freed after its use + struct ggml_tensor * a_zero = ggml_repeat_4d(ctx, ggml_fill(ctx, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1), 0.0f), + src->ne[0], src->ne[1], src->ne[2], src->ne[3]); cgraph->grads[isrc] = ggml_acc_impl(ctx, a_zero, tensor, nb1, nb2, nb3, offset, false); } ggml_format_name(cgraph->grads[isrc], "grad for %s", cgraph->visited_hash_set.keys[isrc]->name); From 8ec4f9d84305e3dd2c938eb1ee299ba3d7e1d8cf Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 18:08:17 -0500 Subject: [PATCH 18/43] metal: the attention backward runs on simdgroup matrices, tiled, with dO scaled for half tiles FLASH_ATTN_BACK's Metal kernel was one simdgroup per row with no reuse: 196 ms at a training shape (head 256, 16 heads over 4 KV heads, 256 rows, 2048 window, f16 cache) where the forward takes 2.5 ms. Now three passes (FlashAttention-2 style; the per-row kernel stays for K and V heads of different sizes): - stats (query blocks of 8): each row's log-sum-exp and D in one sweep, online; - dkv (KV blocks of 16): dK += dS^T Q and dV += P^T dO over the KV head's whole group; - dq (query blocks of 8): dQ += dS K. S and dP are 8x8 simdgroup tiles of half inputs accumulated in F32; staging steps rows with threads across columns (no division per element); partial blocks write through threadgroup memory. dO is scaled by the power of two that brings its largest |value| to [4, 8) (ggml_flash_attn_back_scale, recorded in the result after D by every backend): training gradients sit below half's normal range, and unscaled the half tiles cost the gradient its precision (4B against the F32 gradient at an f16 cache: 0.9943/0.9946 unscaled, 1.0000/0.9998 scaled). test-backend-ops: a perf case beside the forward at the same shape, with the backward's FLOPs. Measured (M5 Metal, 2026-10-14): FLASH_ATTN_BACK 4/4 against the CPU (f16, f32, q8_0, GQA, head 256, 300 positions x 37 rows). Perf at the shape above: 196 ms (per row) -> 67 ms per call, 321 GFLOPS (the forward: 2.5 ms, 3.45 TFLOPS): faster, not yet near the forward; the next work. 4B training measure at window 8192 (both resident slots busy, confounded): 77.8 s against 78 s for the composed backward, chunk graph 1,286.6 MiB against its 1,623. test-walk-exact OK on qwen2.5-0.5b and the 4B hybrid (23:02:52-23:08:04Z): every exact row 1.0000/1.0000, against the F32 gradient back to the composed path's values, cache byte-identical after the walk. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/include/ggml.h | 7 +- ggml/src/ggml-cpu/ops.cpp | 14 + ggml/src/ggml-metal/ggml-metal-device.cpp | 8 +- ggml/src/ggml-metal/ggml-metal-device.h | 2 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 24 +- ggml/src/ggml-metal/kernels/fa.metal | 361 ++++++++++++++++++++++ ggml/src/ggml.c | 12 +- tests/test-backend-ops.cpp | 13 + 8 files changed, 432 insertions(+), 9 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 4597ee2cded6..d4d92cacfe42 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2485,7 +2485,8 @@ extern "C" { // scores per block instead of keeping them: memory is the gradients alone. The result is F32: // dQ ([D, N, H]), dK (k's shape), dV (v's shape), each contiguous at // ggml_flash_attn_back_offset, then each query row's log-sum-exp and D = sum_j p_ij dp_ij, - // which the kernels pass between their passes. + // which the kernels pass between their passes, and the power of two dO is scaled by in a + // kernel's reduced-precision tiles (ggml_flash_attn_back_scale). GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, struct ggml_tensor * q, @@ -2496,6 +2497,10 @@ extern "C" { float scale); // byte offset in a ggml_flash_attn_back result of: 0 dQ, 1 dK, 2 dV, 3 the per-row scratch + // the power of two that brings the largest |dO| to [4, 8): training gradients sit far below half + // precision's normal range (6.1e-5) and lose their precision in half tiles; 1 when dO is zero + // (the Metal kernel computes the same rule, kernel_flash_attn_back_scale) + GGML_API float ggml_flash_attn_back_scale(float max_abs_do); GGML_API size_t ggml_flash_attn_back_offset(const struct ggml_tensor * back, int part); // the backward of FLASH_ATTN_EXT as the one op above (ggml_flash_attn_ext_set_back_tile) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 238c966566c1..0cdcad1c5443 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -9303,6 +9303,20 @@ static void ggml_compute_forward_flash_attn_back_f32( return (float) r; }; + // the dO scale a reduced-precision kernel uses (recorded so every backend's result agrees; the + // CPU works in F32 and does not apply it) + if (ith == 0) { + float mx = 0.0f; + for (int64_t i = 0; i < N; ++i) { + for (int64_t h = 0; h < H; ++h) { + for (int64_t c = 0; c < Dv; ++c) { + mx = std::max(mx, fabsf(*(const float *) ((const char *) d->data + c*d->nb[0] + h*d->nb[1] + i*d->nb[2]))); + } + } + } + dvec[N*H] = ggml_flash_attn_back_scale(mx); + } + // pass 1: the query rows for (int64_t r = ith; r < N*H; r += nth) { const int64_t i = r % N; diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 7059261b2716..d17d8689cc4d 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -353,11 +353,15 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_lo return res; } -ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_back(ggml_metal_library_t lib, const ggml_tensor * op, bool kv_pass) { +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_back(ggml_metal_library_t lib, const ggml_tensor * op, const char * pass) { assert(op->op == GGML_OP_FLASH_ATTN_BACK); char name[256]; - snprintf(name, 256, "kernel_flash_attn_back_%s_%s", kv_pass ? "kv" : "q", ggml_type_name(op->src[1]->type)); + if (strcmp(pass, "scale") == 0) { + snprintf(name, 256, "kernel_flash_attn_back_scale"); + } else { + snprintf(name, 256, "kernel_flash_attn_back_%s_%s", pass, ggml_type_name(op->src[1]->type)); + } ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); if (!res.pipeline) { diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index c82db0dd521e..76c29b4cf62b 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -121,7 +121,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_loss(ggml_metal_library_t lib, const struct ggml_tensor * op); -struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_back(ggml_metal_library_t lib, const struct ggml_tensor * op, bool kv_pass); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_back(ggml_metal_library_t lib, const struct ggml_tensor * op, const char * pass); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_blk (ggml_metal_library_t lib, const struct ggml_tensor * op); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index fbf84f1fe2ea..4fa1f75cca8d 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -4278,8 +4278,8 @@ int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx) { /*.offs_lse =*/ ggml_flash_attn_back_offset(op, 3), }; - for (int pass = 0; pass < 2; ++pass) { - auto pipeline = ggml_metal_library_get_pipeline_flash_attn_back(lib, op, pass == 1); + auto run = [&](const char * pass, int ngx, int ngy, int nth, size_t shmem) { + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_back(lib, op, pass); ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(q), 1); @@ -4288,9 +4288,25 @@ int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, mask ? ggml_metal_get_buffer_id(mask) : ggml_metal_get_buffer_id(q), 4); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(d), 5); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 6); - const int ngroups = pass == 0 ? args.N*args.H : args.M*args.Hkv; - ggml_metal_encoder_dispatch_threadgroups(enc, ngroups, 1, 1, 32, 1, 1); + if (shmem > 0) { + ggml_metal_encoder_set_threadgroup_memory_size(enc, shmem, 0); + } + ggml_metal_encoder_dispatch_threadgroups(enc, ngx, ngy, 1, nth, 1, 1); ggml_metal_op_concurrency_reset(ctx); + }; + // the dO scale (one threadgroup; every path records it so backends agree) + run("scale", 1, 1, 1024, 32*sizeof(float)); + if (args.D == args.Dv && args.D % 32 == 0) { + // tiled (fa.metal FAB_*): query blocks of 8 rows, KV blocks of 16, four simdgroups + const int br = 8, bc = 16, nsg = 4; + const size_t shmem = (size_t) (2*br + 2*bc)*args.D*sizeof(uint16_t) + (2*br*bc + 2*br)*sizeof(float) + 2*br*bc*sizeof(uint16_t); + run("stats", (args.N + br - 1)/br, args.H, 32*nsg, shmem); + run("dkv", (args.M + bc - 1)/bc, args.Hkv, 32*nsg, shmem); + run("dq", (args.N + br - 1)/br, args.H, 32*nsg, shmem); + } else { + // one simdgroup per row + run("q", args.N*args.H, 1, 32, 0); + run("kv", args.M*args.Hkv, 1, 32, 0); } return 1; diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index 3cc9ae274113..c210b005e6b8 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -2435,6 +2435,367 @@ kernel void kernel_flash_attn_back_kv( } } + +// THE TILED BACKWARD (FlashAttention-2 style, simdgroup matrices; D == Dv, D % 32 == 0). Query +// blocks of FAB_BR rows and KV blocks of FAB_BC positions; the four simdgroups of a threadgroup +// each compute one 8x8 tile of S = Q K^T or dP = dO V^T per block (inputs staged as half, +// accumulated in F32), and the gradients accumulate in simdgroup matrices split by column. +// pass A (query blocks): each row's log-sum-exp and D = sum_j p_j dp_j, online in one sweep +// pass B (KV blocks): dK += dS^T Q, dV += P^T dO over every query row of the KV head's group +// pass C (query blocks): dQ += dS K +#define FAB_BR 8 +#define FAB_BC 16 +#define FAB_NSG 4 +// dO (and so dP, D, dS and the gradients) is carried through the half tiles scaled by the power of +// two that brings its largest magnitude to [4, 8) (ggml_flash_attn_back_scale; pass "scale" below): +// training gradients sit far below half's normal range (6.1e-5) and lost their precision there +// (test-walk-exact, M5 2026-10-14: 0.9943 against the F32 gradient unscaled, 1.0000 scaled) + +#define FAB_SCALE_NT 1024 // the scale pass's threads (ggml_metal_op_flash_attn_back dispatches this many) + +kernel void kernel_flash_attn_back_scale( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]]) { + const ushort ntg = FAB_SCALE_NT; + threadgroup float * smax = (threadgroup float *) shmem; + float mx = 0.0f; + for (int r = 0; r < args.N*args.H; ++r) { + device const float * row = (device const float *) (d + (r % args.H)*args.nbd1 + (r / args.H)*args.nbd2); + for (int c = tiitg; c < args.Dv; c += ntg) { + mx = max(mx, fabs(row[c])); + } + } + mx = simd_max(mx); + if (tiisg == 0) { + smax[sgitg] = mx; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiitg == 0) { + for (int s = 1; s < ntg/32; ++s) { + mx = max(mx, smax[s]); + } + float scale = 1.0f; + if (mx > 0.0f) { + int e; + frexp(mx, e); + scale = ldexp(1.0f, clamp(3 - e, -24, 40)); + } + ((device float *) (dst + args.offs_lse))[2*args.N*args.H] = scale; + } +} + +// rows [r0, r0 + FAB_BR|FAB_BC) of a head into threadgroup half, zeros past `n` rows +// (rows outer, threads across a row's columns: no division per element) +template +static inline void fab_stage_kv(threadgroup half * dst, device const char * base, uint64_t nb_row, int r0, int rows, int n, int Dx, + ushort tid) { + for (int r = 0; r < rows; ++r) { + device const char * row = base + (r0 + r)*nb_row; + const bool in = r0 + r < n; + for (int c = tid; c < Dx; c += FAB_NSG*32) { + dst[r*Dx + c] = in ? (half) fa_back_elem(row, c) : (half) 0.0f; + } + } +} + +static inline void fab_stage_rows(threadgroup half * dst, device const char * base, uint64_t nb_row, int r0, int n, int Dx, + ushort tid, float scale = 1.0f) { + for (int r = 0; r < FAB_BR; ++r) { + device const float * row = (device const float *) (base + (r0 + r)*nb_row); + const bool in = r0 + r < n; + for (int c = tid; c < Dx; c += FAB_NSG*32) { + dst[r*Dx + c] = in ? (half) (row[c]*scale) : (half) 0.0f; + } + } +} + +// S (sg 0, 1: keys 0-7, 8-15) and dP (sg 2, 3) for one block into ss / sdp [FAB_BR][FAB_BC] +static inline void fab_scores(threadgroup const half * sq, threadgroup const half * sdo, threadgroup const half * sk, + threadgroup const half * sv, threadgroup float * ss, threadgroup float * sdp, int D, ushort sgitg) { + const bool dp = sgitg >= 2; + const int t = sgitg & 1; + threadgroup const half * a = dp ? sdo : sq; + threadgroup const half * b = (dp ? sv : sk) + 8*t*D; + simdgroup_float8x8 acc = make_filled_simdgroup_matrix(0.0f); + simdgroup_half8x8 ma, mb; + for (int d0 = 0; d0 < D; d0 += 8) { + simdgroup_load(ma, a + d0, D); + simdgroup_load(mb, b + d0, D, 0, true); // K^T: [d][key] + simdgroup_multiply_accumulate(acc, ma, mb, acc); + } + simdgroup_store(acc, (dp ? sdp : ss) + 8*t, FAB_BC); +} + +template +kernel void kernel_flash_attn_back_stats( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const float ds = ((device const float *) (dst + args.offs_lse))[2*args.N*args.H]; // the dO scale + const int D = args.D; + const int h = tgpig.y; + const int i0 = tgpig.x*FAB_BR; + const int hk = h / (args.H / args.Hkv); + + threadgroup half * sq = (threadgroup half *) shmem; + threadgroup half * sdo = sq + FAB_BR*D; + threadgroup half * sk = sdo + FAB_BR*D; + threadgroup half * sv = sk + FAB_BC*D; + threadgroup float * ss = (threadgroup float *) (sv + FAB_BC*D); + threadgroup float * sdp = ss + FAB_BR*FAB_BC; + + fab_stage_rows(sq, q + h*args.nbq2, args.nbq1, i0, args.N, D, tiitg); + fab_stage_rows(sdo, d + h*args.nbd1, args.nbd2, i0, args.N, D, tiitg, ds); + + float m = -INFINITY, l = 0.0f, a = 0.0f; // thread r < FAB_BR owns row i0 + r + for (int j0 = 0; j0 < args.M; j0 += FAB_BC) { + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_stage_kv(sk, k + hk*args.nbk2, args.nbk1, j0, FAB_BC, args.M, D, tiitg); + fab_stage_kv(sv, v + hk*args.nbv2, args.nbv1, j0, FAB_BC, args.M, D, tiitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_scores(sq, sdo, sk, sv, ss, sdp, D, sgitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiitg < FAB_BR && i0 + tiitg < args.N) { + const int i = i0 + tiitg; + for (int c = 0; c < FAB_BC && j0 + c < args.M; ++c) { + const float sc = ss[tiitg*FAB_BC + c]*args.scale + fa_back_mask(args, mask, j0 + c, i, h); + if (sc == -INFINITY) { + continue; + } + const float mn = max(m, sc); + const float r = exp(m - mn); + const float e = exp(sc - mn); + l = l*r + e; + a = a*r + e*sdp[tiitg*FAB_BC + c]; + m = mn; + } + } + } + if (tiitg < FAB_BR && i0 + tiitg < args.N) { + device float * lsev = (device float *) (dst + args.offs_lse); + lsev[h*args.N + i0 + tiitg] = m == -INFINITY ? -INFINITY : m + log(l); + lsev[args.N*args.H + h*args.N + i0 + tiitg] = m == -INFINITY ? 0.0f : a/l/ds; + } +} + +// P and dS of one block from ss / sdp into half sp / sds [FAB_BR][FAB_BC] +static inline void fab_probs(constant ggml_metal_kargs_flash_attn_back & args, device const char * mask, + threadgroup const float * ss, threadgroup const float * sdp, threadgroup const float * slse, + threadgroup half * sp, threadgroup half * sds, int i0, int j0, int h, ushort tiitg) { + const int r = tiitg / FAB_BC; + const int c = tiitg % FAB_BC; + const int i = i0 + r; + const int j = j0 + c; + float p = 0.0f; + float ds = 0.0f; + if (i < args.N && j < args.M && slse[r] != -INFINITY) { + const float sc = ss[r*FAB_BC + c]*args.scale + fa_back_mask(args, mask, j, i, h); + p = sc == -INFINITY ? 0.0f : exp(sc - slse[r]); + ds = p*(sdp[r*FAB_BC + c] - slse[FAB_BR + r])*args.scale; + } + sp [r*FAB_BC + c] = (half) p; + sds[r*FAB_BC + c] = (half) ds; +} + +template +kernel void kernel_flash_attn_back_dkv( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const float ds = ((device const float *) (dst + args.offs_lse))[2*args.N*args.H]; // the dO scale + const int D = args.D; + const int hk = tgpig.y; + const int j0 = tgpig.x*FAB_BC; + const int G = args.H / args.Hkv; + const int CW = D / FAB_NSG; // the columns this simdgroup owns + const int c0 = sgitg*CW; + + threadgroup half * sq = (threadgroup half *) shmem; + threadgroup half * sdo = sq + FAB_BR*D; + threadgroup half * sk = sdo + FAB_BR*D; + threadgroup half * sv = sk + FAB_BC*D; + threadgroup float * ss = (threadgroup float *) (sv + FAB_BC*D); + threadgroup float * sdp = ss + FAB_BR*FAB_BC; + threadgroup float * slse = sdp + FAB_BR*FAB_BC; // [lse x FAB_BR][D x FAB_BR] + threadgroup half * sp = (threadgroup half *) (slse + 2*FAB_BR); + threadgroup half * sds = sp + FAB_BR*FAB_BC; + + fab_stage_kv(sk, k + hk*args.nbk2, args.nbk1, j0, FAB_BC, args.M, D, tiitg); + fab_stage_kv(sv, v + hk*args.nbv2, args.nbv1, j0, FAB_BC, args.M, D, tiitg); + + simdgroup_float8x8 gk[2][FA_BACK_MAX_D/8/FAB_NSG]; + simdgroup_float8x8 gv[2][FA_BACK_MAX_D/8/FAB_NSG]; + for (int r = 0; r < 2; ++r) { + for (int t = 0; t < CW/8; ++t) { + gk[r][t] = make_filled_simdgroup_matrix(0.0f); + gv[r][t] = make_filled_simdgroup_matrix(0.0f); + } + } + device const float * lsev = (device const float *) (dst + args.offs_lse); + for (int h = hk*G; h < (hk + 1)*G; ++h) { + for (int i0 = 0; i0 < args.N; i0 += FAB_BR) { + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_stage_rows(sq, q + h*args.nbq2, args.nbq1, i0, args.N, D, tiitg); + fab_stage_rows(sdo, d + h*args.nbd1, args.nbd2, i0, args.N, D, tiitg, ds); + if (tiitg < FAB_BR) { + const int i = i0 + tiitg; + slse[tiitg] = i < args.N ? lsev[h*args.N + i] : -INFINITY; + slse[FAB_BR + tiitg] = i < args.N ? lsev[args.N*args.H + h*args.N + i]*ds : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_scores(sq, sdo, sk, sv, ss, sdp, D, sgitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_probs(args, mask, ss, sdp, slse, sp, sds, i0, j0, h, tiitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + // dV[16 x CW] += P^T[16 x 8] dO[8 x CW]; dK[16 x CW] += dS^T[16 x 8] Q[8 x CW] + simdgroup_half8x8 pt, dst8, mo, mq; + for (int r = 0; r < 2; ++r) { + simdgroup_load(pt, sp + 8*r, FAB_BC, 0, true); + simdgroup_load(dst8, sds + 8*r, FAB_BC, 0, true); + for (int t = 0; t < CW/8; ++t) { + simdgroup_load(mo, sdo + c0 + 8*t, D); + simdgroup_load(mq, sq + c0 + 8*t, D); + simdgroup_multiply_accumulate(gv[r][t], pt, mo, gv[r][t]); + simdgroup_multiply_accumulate(gk[r][t], dst8, mq, gk[r][t]); + } + } + } + } + // through threadgroup memory (the staged K/V are done with), so a partial block writes only its rows + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float * so = (threadgroup float *) sk; // [FAB_BC][D] + for (int part = 0; part < 2; ++part) { + for (int r = 0; r < 2; ++r) { + for (int t = 0; t < CW/8; ++t) { + simdgroup_store(part == 0 ? gk[r][t] : gv[r][t], so + 8*r*D + c0 + 8*t, D); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + device float * out = (device float *) (dst + (part == 0 ? args.offs_dk : args.offs_dv)) + (hk*args.M + j0)*D; + for (int e = tiitg; e < FAB_BC*D; e += FAB_NSG*32) { + if (j0 + e / D < args.M) { + out[e] = so[e]/ds; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +template +kernel void kernel_flash_attn_back_dq( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const float ds = ((device const float *) (dst + args.offs_lse))[2*args.N*args.H]; // the dO scale + const int D = args.D; + const int h = tgpig.y; + const int i0 = tgpig.x*FAB_BR; + const int hk = h / (args.H / args.Hkv); + const int CW = D / FAB_NSG; + const int c0 = sgitg*CW; + + threadgroup half * sq = (threadgroup half *) shmem; + threadgroup half * sdo = sq + FAB_BR*D; + threadgroup half * sk = sdo + FAB_BR*D; + threadgroup half * sv = sk + FAB_BC*D; + threadgroup float * ss = (threadgroup float *) (sv + FAB_BC*D); + threadgroup float * sdp = ss + FAB_BR*FAB_BC; + threadgroup float * slse = sdp + FAB_BR*FAB_BC; + threadgroup half * sp = (threadgroup half *) (slse + 2*FAB_BR); + threadgroup half * sds = sp + FAB_BR*FAB_BC; + + fab_stage_rows(sq, q + h*args.nbq2, args.nbq1, i0, args.N, D, tiitg); + fab_stage_rows(sdo, d + h*args.nbd1, args.nbd2, i0, args.N, D, tiitg, ds); + device const float * lsev = (device const float *) (dst + args.offs_lse); + if (tiitg < FAB_BR) { + const int i = i0 + tiitg; + slse[tiitg] = i < args.N ? lsev[h*args.N + i] : -INFINITY; + slse[FAB_BR + tiitg] = i < args.N ? lsev[args.N*args.H + h*args.N + i]*ds : 0.0f; + } + + simdgroup_float8x8 gq[FA_BACK_MAX_D/8/FAB_NSG]; + for (int t = 0; t < CW/8; ++t) { + gq[t] = make_filled_simdgroup_matrix(0.0f); + } + for (int j0 = 0; j0 < args.M; j0 += FAB_BC) { + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_stage_kv(sk, k + hk*args.nbk2, args.nbk1, j0, FAB_BC, args.M, D, tiitg); + fab_stage_kv(sv, v + hk*args.nbv2, args.nbv1, j0, FAB_BC, args.M, D, tiitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_scores(sq, sdo, sk, sv, ss, sdp, D, sgitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_probs(args, mask, ss, sdp, slse, sp, sds, i0, j0, h, tiitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + // dQ[8 x CW] += dS[8 x 16] K[16 x CW] + simdgroup_half8x8 ms, mk; + for (int kk = 0; kk < 2; ++kk) { + simdgroup_load(ms, sds + 8*kk, FAB_BC); + for (int t = 0; t < CW/8; ++t) { + simdgroup_load(mk, sk + 8*kk*D + c0 + 8*t, D); + simdgroup_multiply_accumulate(gq[t], ms, mk, gq[t]); + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float * so = (threadgroup float *) sk; // [FAB_BR][D] + for (int t = 0; t < CW/8; ++t) { + simdgroup_store(gq[t], so + c0 + 8*t, D); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + device float * out = (device float *) dst + (h*args.N + i0)*D; + for (int e = tiitg; e < FAB_BR*D; e += FAB_NSG*32) { + if (i0 + e / D < args.N) { + out[e] = so[e]/ds; + } + } +} + +typedef decltype(kernel_flash_attn_back_stats) kernel_flash_attn_back_tiled_t; + +template [[host_name("kernel_flash_attn_back_stats_f32")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_stats; +template [[host_name("kernel_flash_attn_back_stats_f16")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_stats; +template [[host_name("kernel_flash_attn_back_stats_q8_0")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_stats; +template [[host_name("kernel_flash_attn_back_dkv_f32")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dkv; +template [[host_name("kernel_flash_attn_back_dkv_f16")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dkv; +template [[host_name("kernel_flash_attn_back_dkv_q8_0")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dkv; +template [[host_name("kernel_flash_attn_back_dq_f32")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dq; +template [[host_name("kernel_flash_attn_back_dq_f16")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dq; +template [[host_name("kernel_flash_attn_back_dq_q8_0")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dq; + typedef decltype(kernel_flash_attn_back_q) kernel_flash_attn_back_q_t; typedef decltype(kernel_flash_attn_back_kv) kernel_flash_attn_back_kv_t; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index f08bf6aa0c7e..e42f71627d91 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5542,6 +5542,16 @@ void ggml_flash_attn_ext_add_sinks( // ggml_flash_attn_back +float ggml_flash_attn_back_scale(float max_abs_do) { + if (!(max_abs_do > 0.0f)) { + return 1.0f; + } + int e; + frexpf(max_abs_do, &e); // max_abs_do in [2^(e-1), 2^e) + const int k = 3 - e; // scaled into [2^2, 2^3) + return ldexpf(1.0f, k < -24 ? -24 : (k > 40 ? 40 : k)); +} + size_t ggml_flash_attn_back_offset(const struct ggml_tensor * back, int part) { const struct ggml_tensor * q = back->src[0]; const struct ggml_tensor * k = back->src[1]; @@ -5577,7 +5587,7 @@ struct ggml_tensor * ggml_flash_attn_back( const size_t rows = (size_t) q->ne[1]*q->ne[2]; struct ggml_tensor tmp = { 0 }; tmp.src[0] = q; tmp.src[1] = k; tmp.src[2] = v; - const size_t bytes = ggml_flash_attn_back_offset(&tmp, 3) + 2*rows*sizeof(float); + const size_t bytes = ggml_flash_attn_back_offset(&tmp, 3) + (2*rows + 1)*sizeof(float); // lse, D, the dO scale struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) (bytes/sizeof(float))); ggml_set_op_params_f32(result, 0, scale); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 0c622d6ec557..6cd6452c15b5 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7561,6 +7561,13 @@ struct test_flash_attn_back : public test_case { double max_nmse_err() override { return 1e-5; // the kernel and the reference sum in different orders over the window } + + // the five products of the ideal backward (S, dP, dV, dQ, dK), so perf reports TFLOPS against + // the forward's two + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + return 2ull*nb*kv*nh*(3*hs + 2*hs); + } }; // GGML_OP_CROSS_ENTROPY_LOSS @@ -10683,6 +10690,12 @@ static std::vector> make_test_cases_eval() { static std::vector> make_test_cases_perf() { std::vector> test_cases; + // the attention backward beside its forward at one training shape (head 256, 16 heads over 4 KV + // heads, 256 rows over a 2048 window, f16 cache): small enough to time on every kernel edit; a + // backward costs about 2.5 forwards of matmul work + test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 2048, 256, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {4, 1}, 2048, 256, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // OUT_PROD at the shapes a LoRA backward emits (dx = out_prod(W, dy^T), trans_b): a 1.5B // (n_embd 1536, ffn 8960) and a 27B (n_embd 5120, ffn 17408) at 1024 tokens, plus the // output head (vocab 151936) โ€” the largest frozen matmul in the step. F32 beside Q4_K From 22c112cc86f0f0d8cf7b38932c804dcfb9658ea3 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 00:28:40 -0500 Subject: [PATCH 19/43] metal: the tiled attention backward is chosen only when its tiles fit threadgroup memory Cormac's review of 8ec4f9d84: the tiled path's threadgroup memory is 26,176 bytes at head 256 but 50,752 at head 512, over Apple's 32 KiB, and the only guard was D == Dv && D % 32 == 0. The route now also requires the computed size to fit the device's maximum; otherwise the per-row path runs. test-backend-ops gains a head-512 case, which takes the per-row path. Measured (M5 Metal, 2026-10-15): FLASH_ATTN_BACK 5/5 against the CPU, the head-512 case included. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 11 +++++++---- tests/test-backend-ops.cpp | 1 + 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 4fa1f75cca8d..ffa2a9ca2f78 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -4296,10 +4296,13 @@ int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx) { }; // the dO scale (one threadgroup; every path records it so backends agree) run("scale", 1, 1, 1024, 32*sizeof(float)); - if (args.D == args.Dv && args.D % 32 == 0) { - // tiled (fa.metal FAB_*): query blocks of 8 rows, KV blocks of 16, four simdgroups - const int br = 8, bc = 16, nsg = 4; - const size_t shmem = (size_t) (2*br + 2*bc)*args.D*sizeof(uint16_t) + (2*br*bc + 2*br)*sizeof(float) + 2*br*bc*sizeof(uint16_t); + // tiled (fa.metal FAB_*): query blocks of 8 rows, KV blocks of 16, four simdgroups; its tiles + // must fit the device's threadgroup memory (head 256: 26,176 bytes; head 512 would need 50,752, + // over Apple's 32 KiB: Cormac's review of 8ec4f9d84) + const int br = 8, bc = 16, nsg = 4; + const size_t shmem = (size_t) (2*br + 2*bc)*args.D*sizeof(uint16_t) + (2*br*bc + 2*br)*sizeof(float) + 2*br*bc*sizeof(uint16_t); + const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); + if (args.D == args.Dv && args.D % 32 == 0 && shmem <= props_dev->max_theadgroup_memory_size) { run("stats", (args.N + br - 1)/br, args.H, 32*nsg, shmem); run("dkv", (args.M + bc - 1)/bc, args.Hkv, 32*nsg, shmem); run("dq", (args.N + br - 1)/br, args.H, 32*nsg, shmem); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 6cd6452c15b5..c0072af3d020 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10500,6 +10500,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_back(128, 8, 2, 512, 64, t)); } test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 300, 37, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_back(512, 4, 2, 128, 9, GGML_TYPE_F16)); // the tiled path's tiles exceed 32 KiB: the per-row path test_cases.emplace_back(new test_cross_entropy_loss_index()); test_cases.emplace_back(new test_cross_entropy_loss_index({30000, 2, 1, 1})); // a vocabulary's width; few rows keep the finite difference above f32 rounding test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); From 6fc7bc66fc6a847d78bd580534281edc9622a840 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 00:45:08 -0500 Subject: [PATCH 20/43] llama: llama_opt_learn trains one reply on a live sequence of a serving context P2 (continuum #5032 section 9.6). The engine half of "the server batches learning with serving": a reply trains on the sequence that served it, in the serving context, with no second context and no second copy of her window. - llama_opt_learn(ctx, seq, n_cached, tokens, labels, n, result): the sequence holds her context exactly as served ([0, n_cached)); attention cells past it are dropped, and a recurrent part must already end there (restore the served-prefix checkpoint first: a recurrent state cannot be cut back), refused by name otherwise, as is a pending rollback. The reply decodes with labels (the shadow on labelled rows only, PR64's physical-row path), then is restored under her stable set from the recurrent snapshot. Her context is never decoded again; other sequences are not touched. One batch is one exact step; a longer reply steps per batch. - The one-ubatch-per-batch rule moves from llama_opt_init into the walk (a named refusal there), so a serving context with -b > -ub can learn. - test-walk-exact case 9: two sequences in one context, her reply learned in place on sequence 0. Measured (M5 Metal, 2026-10-15): test-walk-exact OK. qwen2.5-0.5b: learned in place vs the walk's exact step cosine 1.0000, norm ratio 0.9997; her sequence byte-identical to a stable decode of the reply; the neighbour byte-identical. qwen3.5-4b hybrid (05:35:29-05:44:47Z): cosine 1.0000, norm ratio 1.0000, her sequence and the neighbour byte-identical; every other row unchanged. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- include/llama.h | 20 +++++++ src/llama-context-opt.cpp | 113 +++++++++++++++++++++++++++++++++++++- src/llama-context.h | 3 + tests/test-walk-exact.cpp | 88 +++++++++++++++++++++++++++++ 4 files changed, 223 insertions(+), 1 deletion(-) diff --git a/include/llama.h b/include/llama.h index 53d747830d8b..ee65be13aa27 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1781,6 +1781,26 @@ extern "C" { // window): the requested one, or smaller where a chunk's graph did not fit the device. LLAMA_API uint32_t llama_opt_walk_horizon(struct llama_context * lctx); + // LEARN ON A LIVE SEQUENCE (P2): train one reply on `seq` of a serving context, given the + // context before it exactly as served. The sequence must hold positions [0, n_cached) as served + // and nothing after them for a recurrent part (restore the served-prefix checkpoint first: + // a recurrent state cannot be cut back); attention cells past n_cached are dropped. The reply + // tokens[0..n) decode at positions n_cached.. with labels[i] the target after tokens[i] (-1: + // none): the trainable adapter applies to labelled rows only, the context is never decoded + // again, and afterwards the reply is restored under her stable set, so the sequence (every + // cell, the recurrent state) is what serving left and every other sequence is untouched. + // A reply within one batch is one exact step; a longer one steps per batch, its gradient + // stopped at each batch boundary. false: refused or failed by name (llama_opt_failure), + // with the sequence restored. + LLAMA_API bool llama_opt_learn( + struct llama_context * lctx, + llama_seq_id seq, + llama_pos n_cached, + const llama_token * tokens, + const llama_token * labels, + int32_t n, + ggml_opt_result_t result); + LLAMA_API void llama_opt_epoch( struct llama_context * lctx, ggml_opt_dataset_t dataset, diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 877ae4143692..06ec6f2fab6a 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -95,7 +95,8 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params // and the context before it from the cache as a constant (build_attn); opt_epoch_iter // walks the window decoding context and training each labelled chunk. One chunk is one // batch is one ubatch: a chunk's backward needs all of its own K/V in one graph. - GGML_ASSERT(n_ubatch == n_batch && "training needs one ubatch per batch: set -ub = -b"); + // (the walk alone needs one ubatch per batch: a chunk's backward needs all of its own K/V in one + // graph; opt_epoch refuses otherwise. llama_opt_learn trains in a serving context's ubatches.) // In a flash-attention context the training graph reads attention through FLASH_ATTN_EXT too // (build_attn_mha), whose backward is tiled over the window; flash attention also serves the // walk's context decodes and lets the cache hold V un-transposed. @@ -625,6 +626,10 @@ void llama_context::opt_epoch_iter( int64_t ndata_in_loop, int64_t t_loop_start) { GGML_ASSERT(opt_ctx); + if (this->n_ubatch() != std::min(this->n_batch(), opt_n_ctx_train)) { + opt_fail("the walk needs one ubatch per batch (a chunk's backward needs all of its own K/V in one graph): set -ub = -b"); + return; + } const uint32_t n_ctx = opt_n_ctx_train; const uint32_t n_batch = std::min(this->n_batch(), n_ctx); @@ -778,6 +783,112 @@ const char * llama_opt_failure(struct llama_context * ctx) { return ctx->opt_failure.c_str(); } +bool llama_context::opt_learn(llama_seq_id seq, llama_pos n_cached, const llama_token * tokens, const llama_token * labels, + int32_t n, ggml_opt_result_t result) { + if (opt_ctx == nullptr) { + opt_fail("llama_opt_learn needs llama_opt_init on this context first"); + return false; + } + if (n <= 0 || n_cached < 0) { + opt_fail("llama_opt_learn: an empty reply or a negative served prefix"); + return false; + } + // the parts of the memory: a recurrent state cannot be cut back, attention cells can + llama_memory_recurrent * recr = nullptr; + llama_memory_i * attn = memory.get(); + if (auto * hybrid = dynamic_cast(memory.get())) { + recr = hybrid->get_mem_recr(); + attn = hybrid->get_mem_attn(); + } else if (auto * r = dynamic_cast(memory.get())) { + recr = r; + attn = nullptr; + } + // refusals before anything is touched + const llama_pos held = memory->seq_pos_max(seq) + 1; + if (held < n_cached) { + opt_fail("llama_opt_learn: sequence " + std::to_string(seq) + " holds " + std::to_string(held) + + " positions, fewer than the " + std::to_string(n_cached) + " served before the reply"); + return false; + } + if (recr != nullptr && recr->seq_pos_max(seq) + 1 != n_cached) { + opt_fail("llama_opt_learn: the recurrent state of sequence " + std::to_string(seq) + " ends at " + + std::to_string(recr->seq_pos_max(seq) + 1) + ", not at the served prefix " + std::to_string(n_cached) + + ": restore the served-prefix checkpoint first"); + return false; + } + if (recr != nullptr && recr->rollback_pending(seq)) { + opt_fail("llama_opt_learn: sequence " + std::to_string(seq) + " has a pending recurrent rollback: finish or drop the draft first"); + return false; + } + if (attn != nullptr && !attn->seq_rm(seq, n_cached, -1)) { + opt_fail("llama_opt_learn: could not drop sequence " + std::to_string(seq) + "'s cells past " + std::to_string(n_cached)); + return false; + } + // her stable recurrent state at the served prefix, for the restore + std::vector snap; + if (recr != nullptr) { + synchronize(); + snap.resize(state_seq_get_size(seq, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)); + GGML_ASSERT(state_seq_get_data(seq, snap.data(), snap.size(), LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == snap.size()); + } + const uint32_t nb = n_batch(); + llama_batch batch = llama_batch_init(nb, 0, 1); + // [c0, c0 + m) of the reply into the batch; with labels, the trainable layer sees those rows + auto fill = [&](int32_t c0, int32_t m, bool labelled) { + batch.n_tokens = m; + for (int32_t i = 0; i < m; ++i) { + batch.token [i] = tokens[c0 + i]; + batch.pos [i] = n_cached + c0 + i; + batch.n_seq_id[i] = 1; + batch.seq_id [i][0] = seq; + batch.label [i] = labelled ? labels[c0 + i] : -1; + batch.logits [i] = labelled && labels[c0 + i] >= 0; + } + if (!labelled) { + batch.logits[m - 1] = true; + } + }; + bool ok = true; + opt_result_active = result; + for (int32_t c0 = 0; c0 < n && ok; c0 += (int32_t) nb) { + const int32_t m = std::min((int32_t) nb, n - c0); + fill(c0, m, true); + ok = decode(batch) == 0; + if (!ok && !opt_alloc_failed.load()) { + opt_fail("llama_opt_learn: the reply's training batch at " + std::to_string(n_cached + c0) + " failed to decode"); + } + } + opt_result_active = nullptr; + // RESTORE: the reply's cells and the recurrent state decoded again under her stable set + bool restored = attn == nullptr || attn->seq_rm(seq, n_cached, -1); + if (recr != nullptr) { + synchronize(); + restored = restored && state_seq_set_data(seq, snap.data(), snap.size(), LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == snap.size(); + } + for (int32_t c0 = 0; c0 < n && restored; c0 += (int32_t) nb) { + const int32_t m = std::min((int32_t) nb, n - c0); + fill(c0, m, false); + restored = decode(batch) == 0; + } + llama_batch_free(batch); + if (!restored) { + opt_fail("llama_opt_learn: sequence " + std::to_string(seq) + " could not be restored after training: re-prefill it"); + return false; + } + return ok; +} + +bool llama_opt_learn( + struct llama_context * ctx, + llama_seq_id seq, + llama_pos n_cached, + const llama_token * tokens, + const llama_token * labels, + int32_t n, + ggml_opt_result_t result) { + return ctx->opt_learn(seq, n_cached, tokens, labels, n, result); +} + void llama_opt_epoch( struct llama_context * ctx, ggml_opt_dataset_t dataset, diff --git a/src/llama-context.h b/src/llama-context.h index 88bd187b9145..ffaff8af72cf 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -215,6 +215,9 @@ struct llama_context { // the largest training graph the allocation preflight measured on a GPU device (bytes) size_t opt_graph_bytes() const; + bool opt_learn(llama_seq_id seq, llama_pos n_cached, const llama_token * tokens, const llama_token * labels, int32_t n, + ggml_opt_result_t result); + void opt_epoch( ggml_opt_dataset_t dataset, ggml_opt_result_t result_train, diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index affa93f8bab5..eb1d8cba7ee5 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -652,6 +652,94 @@ int main(int argc, char ** argv) { } } + { + // 9. LEARNING ON A LIVE SEQUENCE (llama_opt_learn, P2). What this catches: training that + // needs its own context, re-decodes her context, disturbs another slot, or leaves the shadow + // in her cache. A serving-shaped context holds two sequences; sequence 0 is her context as + // served, sequence 1 a neighbour. Her reply trains on sequence 0 from the served prefix: + // the step must be the walk's exact step for the same window (one chunk per run), her + // sequence must end as a stable decode of the reply would leave it, and the neighbour's + // state must not change by a byte. + std::vector reply(WINDOW + 1, 0); + for (uint32_t i = WINDOW / 2 + 37; i <= WINDOW; ++i) { + reply[i] = 1; + } + const llama_pos r0 = WINDOW / 2 + 36; // the first labelled position + const auto want = step_delta(params, model, init, tokens, reply, WINDOW, true); + + auto two_seq_ctx = [&]() { + auto cp = common_context_params_to_llama(params); + cp.n_ctx = 2*WINDOW; cp.n_batch = WINDOW; cp.n_ubatch = WINDOW; cp.n_seq_max = 2; cp.kv_unified = true; + cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; cp.type_k = GGML_TYPE_F32; cp.type_v = GGML_TYPE_F32; + return llama_init_from_model(model, cp); + }; + auto decode_span = [&](llama_context * c, llama_seq_id sq, const llama_token * t, llama_pos p0, int32_t n) { + llama_batch b = llama_batch_init(WINDOW, 0, 1); + common_batch_clear(b); + for (int32_t i = 0; i < n; ++i) { + common_batch_add(b, t[i], p0 + i, { sq }, i + 1 == n); + } + GGML_ASSERT(llama_decode(c, b) == 0); + llama_batch_free(b); + }; + auto seq_state = [&](llama_context * c, llama_seq_id sq) { + std::vector st(llama_state_seq_get_size(c, sq)); + GGML_ASSERT(llama_state_seq_get_data(c, st.data(), st.size(), sq) == st.size()); + return st; + }; + const llama_token * neighbour = tokens.data() + WINDOW + 1; // other text, beyond her window + + // the control: her context, then her reply, decoded plainly beside the same neighbour + llama_context * ctl = two_seq_ctx(); + decode_span(ctl, 1, neighbour, 0, 200); + decode_span(ctl, 0, tokens.data(), 0, r0); + decode_span(ctl, 0, tokens.data() + r0, r0, WINDOW - r0); + const std::vector want_her = seq_state(ctl, 0); + llama_free(ctl); + + llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); + const std::vector before = adapter_params(adapter); + llama_context * ctx = two_seq_ctx(); + float scale = 1.0f; + GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); + decode_span(ctx, 1, neighbour, 0, 200); + decode_span(ctx, 0, tokens.data(), 0, r0); // her context, as served + const std::vector neighbour_before = seq_state(ctx, 1); + llama_opt_params lopt{}; + lopt.param_filter = llama_opt_param_filter_all; + lopt.get_opt_pars = sgd_pars; + lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; + lopt.adapter = adapter; + llama_opt_init(ctx, model, lopt); + std::vector labels(WINDOW - r0); + for (int32_t i = 0; i < (int32_t) WINDOW - r0; ++i) { + labels[i] = tokens[r0 + i + 1]; + } + const bool learned = llama_opt_learn(ctx, 0, r0, tokens.data() + r0, labels.data(), WINDOW - r0, nullptr); + const std::vector her_after = seq_state(ctx, 0); + const std::vector neighbour_after = seq_state(ctx, 1); + llama_free(ctx); + const std::vector after = adapter_params(adapter); + llama_adapter_lora_free(adapter); + std::vector got(before.size()); + for (size_t i = 0; i < got.size(); ++i) { + got[i] = after[i] - before[i]; + } + printf(" live sequence: learned=%d \"%s\"\n", (int) learned, learned ? "" : "see the log"); + if (!learned || !same_step(compare("live sequence: her reply learned in place vs the walk's exact step", got, want))) { + fprintf(stderr, "FAILED: learning on a live sequence is not the walk's step\n"); + ++failures; + } + const bool her_ok = her_after == want_her; + const bool nb_ok = neighbour_after == neighbour_before; + printf(" live sequence: her cache %s a stable decode of the reply; the neighbour %s\n", + her_ok ? "byte-identical to" : "DIFFERS from", nb_ok ? "byte-identical" : "CHANGED"); + if (!her_ok || !nb_ok) { + fprintf(stderr, "FAILED: learning on a live sequence disturbed the cache\n"); + ++failures; + } + } + { // 8. what this catches: a walk that leaves the shadow in her cache. The reply's cells // are written under the shadow while it trains, then restored, so after either walk the From 4c0ee17bebdb01f46571e34b030a2f142809c927 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 00:48:09 -0500 Subject: [PATCH 21/43] llama: llama_opt_learn refuses null inputs and an overlong reply before touching the cache Cormac's review of 6fc7bc66f: null tokens or labels crashed in the batch fill, and a reply past the sequence's context failed inside the training decode (and then the restore) after her reply cells were already dropped. Both are refused by name before any change. Each call also clears the failure state a previous refusal left, so a refusal never refuses the next call. test-walk-exact case 9 gains the refusals: null labels, past the context, a prefix the sequence does not hold; each refused, her sequence byte-identical after all three. Measured (M5 Metal, qwen2.5-0.5b, 2026-10-15): the three refused with her sequence byte-identical; the following live learn unchanged (cosine 1.0000, norm ratio 0.9997 against the walk's exact step; her cache and the neighbour byte-identical). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- src/llama-context-opt.cpp | 13 +++++++++++-- tests/test-walk-exact.cpp | 15 +++++++++++++++ 2 files changed, 26 insertions(+), 2 deletions(-) diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 06ec6f2fab6a..0ab6e8df7a03 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -789,8 +789,17 @@ bool llama_context::opt_learn(llama_seq_id seq, llama_pos n_cached, const llama_ opt_fail("llama_opt_learn needs llama_opt_init on this context first"); return false; } - if (n <= 0 || n_cached < 0) { - opt_fail("llama_opt_learn: an empty reply or a negative served prefix"); + // each call answers for itself: a refusal before it does not refuse this one + opt_alloc_failed.store(false); + opt_stop_requested.store(false); + opt_failure.clear(); + if (n <= 0 || n_cached < 0 || tokens == nullptr || labels == nullptr) { + opt_fail("llama_opt_learn: an empty reply, a negative served prefix, or no tokens / labels"); + return false; + } + if ((int64_t) n_cached + n > (int64_t) cparams.n_ctx_seq) { + opt_fail("llama_opt_learn: the reply ends at " + std::to_string((int64_t) n_cached + n) + ", past the sequence's context of " + + std::to_string(cparams.n_ctx_seq) + " positions"); return false; } // the parts of the memory: a recurrent state cannot be cut back, attention cells can diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index eb1d8cba7ee5..01c89c911a69 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -715,6 +715,21 @@ int main(int argc, char ** argv) { for (int32_t i = 0; i < (int32_t) WINDOW - r0; ++i) { labels[i] = tokens[r0 + i + 1]; } + // refusals come before any change: her sequence byte-identical after each + { + const std::vector her_before = seq_state(ctx, 0); + const bool r_null = llama_opt_learn(ctx, 0, r0, tokens.data() + r0, nullptr, WINDOW - r0, nullptr); + const bool r_long = llama_opt_learn(ctx, 0, r0, tokens.data() + r0, labels.data(), 2*WINDOW, nullptr); + const bool r_more = llama_opt_learn(ctx, 0, r0 + 5, tokens.data() + r0, labels.data(), 10, nullptr); + const bool same = seq_state(ctx, 0) == her_before; + printf(" live sequence refusals: null labels %s, past the context %s, prefix not held %s; her sequence %s\n", + r_null ? "ACCEPTED" : "refused", r_long ? "ACCEPTED" : "refused", r_more ? "ACCEPTED" : "refused", + same ? "byte-identical" : "CHANGED"); + if (r_null || r_long || r_more || !same) { + fprintf(stderr, "FAILED: a refused live learn was accepted or changed her sequence\n"); + ++failures; + } + } const bool learned = llama_opt_learn(ctx, 0, r0, tokens.data() + r0, labels.data(), WINDOW - r0, nullptr); const std::vector her_after = seq_state(ctx, 0); const std::vector neighbour_after = seq_state(ctx, 1); From e740f8911922503fe21b3d640fd13d0505e487b1 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 00:50:49 -0500 Subject: [PATCH 22/43] metal: the tiled attention backward needs simdgroup matrices (Apple7+); older GPUs take the per-row path Raised by Cormac's lane H on IntelMac: the tiled route checked head sizes and threadgroup memory but not has_simdgroup_mm. Measured (M5 Metal, 2026-10-15): FLASH_ATTN_BACK 5/5 against the CPU. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 3 ++- tools/server/server-context.cpp | 5 ++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index ffa2a9ca2f78..c8807949d564 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -4302,7 +4302,8 @@ int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx) { const int br = 8, bc = 16, nsg = 4; const size_t shmem = (size_t) (2*br + 2*bc)*args.D*sizeof(uint16_t) + (2*br*bc + 2*br)*sizeof(float) + 2*br*bc*sizeof(uint16_t); const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); - if (args.D == args.Dv && args.D % 32 == 0 && shmem <= props_dev->max_theadgroup_memory_size) { + // and simdgroup matrices exist only on Apple7+ GPUs (M1 and newer): older Macs take the per-row path + if (props_dev->has_simdgroup_mm && args.D == args.Dv && args.D % 32 == 0 && shmem <= props_dev->max_theadgroup_memory_size) { run("stats", (args.N + br - 1)/br, args.H, 32*nsg, shmem); run("dkv", (args.M + bc - 1)/bc, args.Hkv, 32*nsg, shmem); run("dq", (args.N + br - 1)/br, args.H, 32*nsg, shmem); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 8fe17bba52e8..07188105e86e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3786,8 +3786,11 @@ struct server_context_impl { // - 4 + n_ubatch // - 4 // ref: https://github.com/ggml-org/llama.cpp/pull/20288 + // A lane that learns (--train-dir) also checkpoints ONE token before the end: + // learning a reply re-decodes the last prompt token with the reply's first + // token as its label, from her recurrent state exactly there (llama_opt_learn) if (do_checkpoint) { - static const int checkpoint_offsets[] = {4 + n_ubatch, 4}; + const int checkpoint_offsets[] = {4 + n_ubatch, 4, params_base.train_dir.empty() ? 4 : 1}; bool should_break = false; for (int offset : checkpoint_offsets) { From 018ab27e4bebf670a7ec828163f8f38813f512e8 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 01:02:08 -0500 Subject: [PATCH 23/43] server: POST /learn trains a reply on the slot that served it, in the serving context P2, server half (continuum #5032 section 9.6). A learning turn is a task on the slot loop (SERVER_TASK_TYPE_LEARN), run between that slot's turns: no second context, no copy of her window. - POST /learn {id_slot, n_prefix, reply | reply_tokens, lora_id, lr} (only with --train-dir). The served prompt is what the slot holds; her recurrent state comes back from a checkpoint one token before the prompt's end, which a learning lane keeps (the prompt checkpoint offsets gain 1 when --train-dir is set); the last prompt token is re-decoded labelled with the reply's first token; then the slot returns to her whole served prompt (checkpoint + that token decoded again), so her next turn reuses it. The shadow joins her adapter set for the step only, at scale 1; llama_opt_init runs once (AdamW, the request's learning rate), one shadow per server run for now. - llama_opt_learn: trains to the last labelled row (an unlabelled row after the run is refused in one batch), takes exactly ONE optimizer step per call (manual period end at the ubatch that completes its labelled rows, each ubatch weighted by its share), and reserves no host logits for rows that feed only the loss (opt_outputs_internal; the walk too): a 512-row chunk at a 248k vocabulary was 485 MiB of host RAM, and the serving context's output cap tripped on it. - KV training admission: a context with a stream per slot may train with flash attention when the batch touches one stream (was: refused whenever the cache had several streams, which is how a serving lane runs). - ggml-opt: a caller-driven step's loss is recorded unscaled by opt_period, so the result reports it right (was 4x at -b 2048 -ub 512). - test-walk-exact case 9 runs on a unified F32 cache with explicit attention AND a stream per slot with an f16 cache and flash attention. Measured (M5 Metal, 2026-10-15): test-walk-exact case 9 on qwen2.5-0.5b: unified 1.0000/0.9997, streams + flash + f16 1.0000/1.0000 against the walk's exact step, her sequence and the neighbour byte-identical, the three refusals refused with her sequence unchanged. test-opt-dynamic-accum OK. llama-server, qwen3.5-4b hybrid, 2 slots, -fa on, --train-dir: a served turn, then POST /learn of a 17-token correction four times: loss 2.860 -> 2.525 -> 1.769 -> 1.111, 2.3-2.6 s per learn (slow: GET_ROWS_BACK and OUT_PROD are not on Metal in this branch; the Metal training-op branches merge next). Serving after learning unchanged on her slot (the shadow applies to labelled rows only). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-opt.cpp | 5 +- src/llama-context-opt.cpp | 36 +++++-- src/llama-context.cpp | 13 +-- src/llama-context.h | 7 ++ src/llama-kv-cache.cpp | 16 ++- tests/test-walk-exact.cpp | 19 +++- tools/server/server-context.cpp | 179 ++++++++++++++++++++++++++++++++ tools/server/server-context.h | 1 + tools/server/server-task.cpp | 8 ++ tools/server/server-task.h | 20 ++++ tools/server/server.cpp | 2 + 11 files changed, 286 insertions(+), 20 deletions(-) diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index bc963c6b882d..c0db1dfb594c 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -1397,6 +1397,9 @@ bool ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result) { GGML_ASSERT(opt_ctx->eval_ready); + // a caller-driven step's loss is not scaled by 1/opt_period (ggml_opt_build): recorded as if it + // were, so the result's reader (which multiplies a per-datapoint loss by opt_period) reports it right + const bool manual = opt_ctx->next_manual; if (opt_ctx->allocated_graph == opt_ctx->gb_opt) { const ggml_opt_optimizer_params & opt_pars = opt_ctx->get_opt_pars(opt_ctx->get_opt_pars_ud); @@ -1487,7 +1490,7 @@ void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result) { GGML_ASSERT(opt_ctx->loss->type == GGML_TYPE_F32); float loss; ggml_backend_tensor_get(opt_ctx->loss, &loss, 0, ggml_nbytes(opt_ctx->loss)); - result->loss.push_back(loss); + result->loss.push_back(manual && opt_ctx->loss_per_datapoint ? loss / opt_ctx->opt_period : loss); if (opt_ctx->pred) { GGML_ASSERT(opt_ctx->pred->type == GGML_TYPE_I32); diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 0ab6e8df7a03..9cbb8729bd72 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -105,9 +105,12 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params // (no backward pass) and before any adapter attached since (its nodes uncounted). // Size both again for the training graph, from the same number, plus the nodes each // attention layer's tiled backward adds for this window (one count, owned by ggml). + // the attention backward: one op where every GPU has the kernel, else composed tiles whose node + // count grows with the window (decided first: it sizes the graph) + cparams.fa_back_native = opt_fa_back_native(); { size_t max_nodes = graph_max_nodes(n_ubatch); - if (cparams.flash_attn) { + if (cparams.flash_attn && !cparams.fa_back_native) { max_nodes += (size_t) model->hparams.n_layer() * ggml_flash_attn_ext_back_n_nodes(opt_n_ctx_train, 0); } gf_res_prev.reset(new llm_graph_result(max_nodes)); @@ -115,7 +118,6 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params max_nodes, cparams.pipeline_parallel, cparams.op_offload)); } - cparams.fa_back_native = opt_fa_back_native(); LLAMA_LOG_INFO("%s: the attention backward is %s\n", __func__, cparams.fa_back_native ? "one op (ggml_flash_attn_back): no score is kept" : "composed from tiles (a device has no ggml_flash_attn_back kernel)"); @@ -375,6 +377,12 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l } ggml_opt_prepare_alloc_slices(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->t_loss_slices.data(), res->t_loss_slices.size()); + if (xc == nullptr && opt_learn_labels > 0) { + // llama_opt_learn: one step per call, at the ubatch that completes its labelled rows + const int64_t nl = ubatch.n_labels(); + opt_learn_labels_done += nl; + ggml_opt_set_next_step(opt_ctx, opt_learn_labels_done >= opt_learn_labels, (float) nl / (float) opt_learn_labels, nullptr); + } if (xc != nullptr) { ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); for (const auto & io : res->t_walk) { @@ -857,17 +865,33 @@ bool llama_context::opt_learn(llama_seq_id seq, llama_pos n_cached, const llama_ batch.logits[m - 1] = true; } }; + // train up to the last labelled row: rows after it cannot change the loss, and in one batch an + // unlabelled row after the run is refused (a run ends its sequence's rows); the restore below + // still decodes all n + int32_t n_train = n; + while (n_train > 0 && labels[n_train - 1] < 0) { + --n_train; + } bool ok = true; - opt_result_active = result; - for (int32_t c0 = 0; c0 < n && ok; c0 += (int32_t) nb) { - const int32_t m = std::min((int32_t) nb, n - c0); + opt_result_active = result; + opt_outputs_internal = true; + opt_learn_labels = 0; + opt_learn_labels_done = 0; + for (int32_t i = 0; i < n_train; ++i) { + opt_learn_labels += labels[i] >= 0; + } + for (int32_t c0 = 0; c0 < n_train && ok; c0 += (int32_t) nb) { + const int32_t m = std::min((int32_t) nb, n_train - c0); fill(c0, m, true); ok = decode(batch) == 0; if (!ok && !opt_alloc_failed.load()) { opt_fail("llama_opt_learn: the reply's training batch at " + std::to_string(n_cached + c0) + " failed to decode"); } } - opt_result_active = nullptr; + opt_result_active = nullptr; + opt_outputs_internal = false; + opt_learn_labels = 0; + opt_learn_labels_done = 0; // RESTORE: the reply's cells and the recurrent state decoded again under her stable set bool restored = attn == nullptr || attn->seq_rm(seq, n_cached, -1); if (recr != nullptr) { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index cb95a0333630..7f3b1cc2e579 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1818,8 +1818,9 @@ int llama_context::decode(const llama_batch & batch_inp) { break; } - // reserve output buffer - if (output_reserve(n_outputs_all) < n_outputs_all) { + // reserve output buffer (none for outputs that only feed a training loss) + const bool host_outputs = !opt_outputs_internal && opt_walk_active == nullptr; + if (output_reserve(host_outputs ? n_outputs_all : 0) < (host_outputs ? n_outputs_all : 0)) { LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); return -2; }; @@ -1900,7 +1901,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } // extract logits - if (logits.data && t_logits && n_outputs > 0 && opt_walk_active == nullptr && needs_raw_logits(ubatch, sampling.samplers)) { + if (logits.data && t_logits && n_outputs > 0 && host_outputs && needs_raw_logits(ubatch, sampling.samplers)) { ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits); GGML_ASSERT(backend_res != nullptr); GGML_ASSERT(logits.data != nullptr); @@ -1915,7 +1916,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } // extract embeddings - if (embd.data && t_embd && n_outputs > 0) { + if (embd.data && t_embd && n_outputs > 0 && host_outputs) { ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(sched.get(), t_embd); GGML_ASSERT(backend_embd != nullptr); @@ -1983,7 +1984,7 @@ int llama_context::decode(const llama_batch & batch_inp) { const int64_t n_rows = masked ? n_outputs : (int64_t) ubatch.n_tokens; const int64_t offset = masked ? n_outputs_prev : n_tokens_prev; - if (embd_nextn.data && t_h_nextn && n_rows > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { + if (embd_nextn.data && t_h_nextn && n_rows > 0 && host_outputs && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_nextn); GGML_ASSERT(backend_h != nullptr); @@ -2010,7 +2011,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } while (mctx->next()); // set to total number of outputs in the batch, for use in llama_get_logits_ith - n_outputs = n_outputs_all; + n_outputs = host_outputs ? n_outputs_all : 0; // set output mappings if (n_outputs > 0) { diff --git a/src/llama-context.h b/src/llama-context.h index ffaff8af72cf..c605744e29b1 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -408,6 +408,13 @@ struct llama_context { opt_walk * opt_walk_active = nullptr; ggml_opt_result_t opt_result_active = nullptr; bool opt_forward_only = false; + // the batch's outputs feed only the in-graph loss (llama_opt_learn; a walk): no host logits are + // reserved, copied or mapped for them (a 512-row chunk at a 248k vocabulary is 485 MiB of host RAM) + bool opt_outputs_internal = false; + // llama_opt_learn's one step: its labelled rows in all, and those trained so far (the ubatch that + // completes them takes the step, each weighted by its share) + int64_t opt_learn_labels = 0; + int64_t opt_learn_labels_done = 0; ggml_threadpool_t threadpool = nullptr; ggml_threadpool_t threadpool_batch = nullptr; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ce72a7a20998..fc118c44814c 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2673,12 +2673,24 @@ void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const { kv->set_input_v_rot(dst); } -const char * llama_kv_cache::training_error(const llama_batch &, bool flash_attn) const { +const char * llama_kv_cache::training_error(const llama_batch & batch, bool flash_attn) const { if (!supports_row_gradients()) { return "cache row indices exceed the backward gather's index range"; } + // the flash-attention backward reads one stream: a training batch may touch only one (a serving + // context with a stream per slot learns on one slot's stream at a time) if (flash_attn && get_n_stream() != 1) { - return "flash attention backward requires a unified cache stream"; + int64_t stream = -1; + for (int32_t i = 0; i < batch.n_tokens; ++i) { + for (int32_t k = 0; k < (batch.n_seq_id ? batch.n_seq_id[i] : 1); ++k) { + const llama_seq_id sq = batch.seq_id ? batch.seq_id[i][k] : 0; + const int64_t st = seq_to_stream[sq]; + if (stream >= 0 && st != stream) { + return "a training batch with flash attention may touch one cache stream: this one spans several"; + } + stream = st; + } + } } return nullptr; } diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 01c89c911a69..b75d164803ad 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -665,12 +665,20 @@ int main(int argc, char ** argv) { reply[i] = 1; } const llama_pos r0 = WINDOW / 2 + 36; // the first labelled position - const auto want = step_delta(params, model, init, tokens, reply, WINDOW, true); - + const auto want_f32 = step_delta(params, model, init, tokens, reply, WINDOW, true); + g_cache_type = GGML_TYPE_F16; + const auto want_f16 = step_delta(params, model, init, tokens, reply, WINDOW, true); + g_cache_type = GGML_TYPE_F32; + + for (const bool streams : { false, true }) { + // unified F32 cache with explicit attention, then a stream per slot with an f16 cache and + // flash attention: how a serving lane runs (each slot attends over its own stream) + printf(" live sequence, %s:\n", streams ? "a stream per slot, f16 cache, flash attention" : "one unified F32 cache, explicit attention"); auto two_seq_ctx = [&]() { auto cp = common_context_params_to_llama(params); - cp.n_ctx = 2*WINDOW; cp.n_batch = WINDOW; cp.n_ubatch = WINDOW; cp.n_seq_max = 2; cp.kv_unified = true; - cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; cp.type_k = GGML_TYPE_F32; cp.type_v = GGML_TYPE_F32; + cp.n_ctx = 2*WINDOW; cp.n_batch = WINDOW; cp.n_ubatch = WINDOW; cp.n_seq_max = 2; cp.kv_unified = !streams; + cp.flash_attn_type = streams ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + cp.type_k = streams ? GGML_TYPE_F16 : GGML_TYPE_F32; cp.type_v = cp.type_k; return llama_init_from_model(model, cp); }; auto decode_span = [&](llama_context * c, llama_seq_id sq, const llama_token * t, llama_pos p0, int32_t n) { @@ -741,7 +749,7 @@ int main(int argc, char ** argv) { got[i] = after[i] - before[i]; } printf(" live sequence: learned=%d \"%s\"\n", (int) learned, learned ? "" : "see the log"); - if (!learned || !same_step(compare("live sequence: her reply learned in place vs the walk's exact step", got, want))) { + if (!learned || !same_step(compare("live sequence: her reply learned in place vs the walk's exact step", got, streams ? want_f16 : want_f32))) { fprintf(stderr, "FAILED: learning on a live sequence is not the walk's step\n"); ++failures; } @@ -753,6 +761,7 @@ int main(int argc, char ** argv) { fprintf(stderr, "FAILED: learning on a live sequence disturbed the cache\n"); ++failures; } + } } { diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 07188105e86e..706187056127 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -892,6 +892,17 @@ struct server_context_impl { common_speculative_init_result_ptr spec_init; common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; + + // LEARNING IN THE SERVING CONTEXT (SERVER_TASK_TYPE_LEARN): the shadow adapter llama_opt_init was + // given (one per server run for now), and the learning rate its next step reads + int32_t learn_lora_id = -1; + float learn_lr = 1e-4f; + static ggml_opt_optimizer_params learn_opt_pars(void * ud) { + ggml_opt_optimizer_params p = ggml_opt_get_default_optimizer_params(nullptr); + p.adamw.alpha = *(const float *) ud; + p.adamw.wd = 0.0f; + return p; + } common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_speculative_ptr spec; @@ -2855,6 +2866,138 @@ struct server_context_impl { res->path = la.path; queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_LEARN: + { + auto res = std::make_unique(); + res->id = task.id; + const auto & L = task.learn; + const int64_t t0 = ggml_time_us(); + server_slot * slot = get_slot_by_id(L.id_slot); + auto refuse = [&](const std::string & why) { + res->ok = false; + res->failure = why; + SRV_WRN("learn refused: %s\n", why.c_str()); + queue_results.send(std::move(res)); + }; + if (slot == nullptr) { + refuse("no slot " + std::to_string(L.id_slot)); + break; + } + if (slot->is_processing()) { + // her slot is serving: learn between its turns, never during one + queue_tasks.defer(std::move(task)); + break; + } + auto & loras = params_base.lora_adapters; + if (L.lora_id < 0 || (size_t) L.lora_id >= loras.size() || loras[L.lora_id].ptr == nullptr) { + refuse("no loaded adapter " + std::to_string(L.lora_id) + " to train (POST /lora-adapters/load it first)"); + break; + } + if (learn_lora_id >= 0 && learn_lora_id != L.lora_id) { + refuse("this server trains adapter " + std::to_string(learn_lora_id) + "; one shadow per server run"); + break; + } + const int32_t n0 = L.n_prefix - 1; // the last prompt token is re-decoded with the reply's first token as its label + if (n0 < 0 || L.n_prefix > slot->prompt.n_tokens() || L.reply.empty()) { + refuse("slot " + std::to_string(slot->id) + " holds " + std::to_string(slot->prompt.n_tokens()) + + " prompt tokens; the served prefix is " + std::to_string(L.n_prefix) + " and the reply must not be empty"); + break; + } + // the recurrent part cannot be cut back: her state one token before the prompt's end + const common_prompt_checkpoint * cp = nullptr; + const bool needs_cp = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS; + if (needs_cp) { + for (const auto & c : slot->prompt.checkpoints) { + if (c.n_tokens == n0) { + cp = &c; + } + } + if (cp == nullptr) { + refuse("slot " + std::to_string(slot->id) + " has no checkpoint one token before its prompt's end (" + + std::to_string(n0) + "): the lane learns only with --train-dir, which keeps one"); + break; + } + } + // the tokens from the last prompt token on, each labelled with the next + std::vector toks; + toks.push_back(slot->prompt.tokens[n0]); + toks.insert(toks.end(), L.reply.begin(), L.reply.end()); + std::vector labels(toks.size(), -1); + for (size_t i = 0; i + 1 < toks.size(); ++i) { + labels[i] = toks[i + 1]; + } + // her stable set plus the shadow: the shadow applies to labelled rows only + std::vector set = slot->lora.empty() ? construct_lora_list({}) : slot->lora; + if ((size_t) L.lora_id >= set.size()) { + set = construct_lora_list({}); + } + set[L.lora_id].scale = 1.0f; + common_set_adapter_lora(ctx_tgt, set); + learn_lr = L.lr; + if (learn_lora_id < 0) { + llama_opt_params lopt{}; + lopt.param_filter = llama_opt_param_filter_all; + lopt.get_opt_pars = learn_opt_pars; + lopt.get_opt_pars_ud = &learn_lr; + lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_ADAMW; + lopt.adapter = loras[L.lora_id].ptr; + llama_opt_init(ctx_tgt, model_tgt, lopt); + learn_lora_id = L.lora_id; + } + // drop what follows n0 (attention) and put her recurrent state back at n0 + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot->id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + ggml_opt_result_t result = ggml_opt_result_init(); + const bool ok = llama_opt_learn(ctx_tgt, slot->id, n0, toks.data(), labels.data(), (int32_t) toks.size(), result); + double loss = 0.0, unc = 0.0; + if (ok) { + ggml_opt_result_loss(result, &loss, &unc); + } + ggml_opt_result_free(result); + // her slot goes back to her whole served prompt (attention cells, recurrent state, + // tokens and checkpoints agree; her next turn reuses all of it): back to n0 from the + // checkpoint, then the last prompt token decoded again under her stable set. A draft + // model's cache is cut to n0 with its own checkpoint, and the slot's tokens with it. + llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot->id, n0, -1); + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot->id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + int32_t n_keep = n0; + if (ctx_dft != nullptr) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot->id, n0, -1); + if (cp != nullptr) { + cp->load_dft(ctx_dft, slot->id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + } else { + llama_batch one = llama_batch_get_one(&toks[0], 1); + llama_pos pos = n0; + one.pos = &pos; + llama_seq_id sid = slot->id; + int32_t nsid = 1; + llama_seq_id * sids = &sid; + one.n_seq_id = &nsid; + one.seq_id = &sids; + if (llama_decode(ctx_tgt, one) == 0) { + n_keep = L.n_prefix; + } else { + llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot->id, n0, -1); + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot->id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + } + } + slot->prompt.tokens.keep_first(n_keep); + slot->prompt.checkpoints.remove_if([&](const common_prompt_checkpoint & c) { return c.n_tokens > n_keep; }); + res->ok = ok; + res->failure = ok ? "" : llama_opt_failure(ctx_tgt); + res->loss = loss; + res->n_labels = (int32_t) toks.size() - 1; + res->t_ms = (ggml_time_us() - t0) / 1000.0; + SRV_INF("learned on slot %d: %d labels from position %d, loss %.4f, %.1f ms%s\n", slot->id, res->n_labels, n0, loss, + res->t_ms, ok ? "" : (" (failed: " + res->failure + ")").c_str()); + queue_results.send(std::move(res)); + } break; case SERVER_TASK_TYPE_UNLOAD_LORA: { auto & loras = params_base.lora_adapters; @@ -5610,6 +5753,42 @@ void server_routes::init_routes() { return res; }; + // POST /learn: one reply learned on the slot that served its prompt, between that slot's turns, in + // this serving context (SERVER_TASK_TYPE_LEARN). Body: id_slot, n_prefix (the served prompt's + // tokens), reply (text) or reply_tokens, lora_id (a loaded adapter: the shadow), lr. + this->post_learn = [this](const server_http_req & req) { + auto res = create_response(); + if (params.train_dir.empty()) { + res->error(format_error_response("learning is off: start the server with --train-dir", ERROR_TYPE_NOT_SUPPORTED)); + return res; + } + const json body = json::parse(req.body); + server_task task(SERVER_TASK_TYPE_LEARN); + task.learn.id_slot = json_value(body, "id_slot", -1); + task.learn.n_prefix = json_value(body, "n_prefix", 0); + task.learn.lora_id = json_value(body, "lora_id", -1); + task.learn.lr = json_value(body, "lr", 1e-4f); + if (body.contains("reply_tokens")) { + task.learn.reply = body.at("reply_tokens").get>(); + } else { + task.learn.reply = common_tokenize(ctx_server.vocab, json_value(body, "reply", std::string()), false, true); + } + auto & rd = res->rd; + task.id = rd.get_new_id(); + rd.post_task(std::move(task)); + auto result = rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + res->ok(result->to_json()); + return res; + }; + this->post_lora_adapters_unload = [this](const server_http_req & req) { auto res = create_response(); const json body = json::parse(req.body); diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 349904972d8d..09bbc2eda320 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -167,6 +167,7 @@ struct server_routes { server_http_context::handler_t post_lora_adapters; server_http_context::handler_t post_lora_adapters_load; server_http_context::handler_t post_lora_adapters_unload; + server_http_context::handler_t post_learn; server_http_context::handler_t get_train; server_http_context::handler_t post_train; server_http_context::handler_t post_train_cancel; diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index aeb6ea508511..35c75c6d32fc 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1693,6 +1693,14 @@ json server_task_result_lora_change::to_json() { return res; } +json server_task_result_learn::to_json() { + json res = { { "success", ok }, { "loss", loss }, { "n_labels", n_labels }, { "t_ms", t_ms } }; + if (!ok) { + res["failure"] = failure; + } + return res; +} + // // server_prompt_cache // diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 312259e7091e..5d7169276a31 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -29,6 +29,7 @@ enum server_task_type { SERVER_TASK_TYPE_SET_LORA, SERVER_TASK_TYPE_LOAD_LORA, SERVER_TASK_TYPE_UNLOAD_LORA, + SERVER_TASK_TYPE_LEARN, }; // TODO: change this to more generic "response_format" to replace the "format_response_*" in server-common @@ -185,6 +186,16 @@ struct server_task { // used by SERVER_TASK_TYPE_UNLOAD_LORA: the adapter id to retire int32_t unload_lora = -1; + // used by SERVER_TASK_TYPE_LEARN: one reply learned on the slot that served its prompt, in this + // serving context (llama_opt_learn): no second context, her prompt never decoded again + struct { + int32_t id_slot = -1; + int32_t n_prefix = 0; // the served prompt's tokens, before the reply + std::vector reply; // what she should have said (her reply, or the reviewed correction) + int32_t lora_id = -1; // the trainable adapter (the shadow), loaded dormant + float lr = 1e-4f; + } learn; + server_task() = default; server_task(server_task_type type) : type(type) {} @@ -582,6 +593,15 @@ struct server_task_result_lora_change : server_task_result { virtual json to_json() override; }; +struct server_task_result_learn : server_task_result { + bool ok = false; + std::string failure; + double loss = 0.0; + int32_t n_labels = 0; + double t_ms = 0.0; + virtual json to_json() override; +}; + struct server_prompt { server_tokens tokens; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 7691dff34922..2bad96993a1b 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -219,6 +219,7 @@ int llama_server(common_params & params, int argc, char ** argv) { routes.post_lora_adapters = models_routes->proxy_post; routes.post_lora_adapters_load = models_routes->proxy_post; routes.post_lora_adapters_unload = models_routes->proxy_post; + routes.post_learn = models_routes->proxy_post; routes.get_slots = models_routes->proxy_get; routes.post_slots = models_routes->proxy_post; @@ -273,6 +274,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/lora-adapters", ex_wrapper(routes.post_lora_adapters)); ctx_http.post("/lora-adapters/load", ex_wrapper(routes.post_lora_adapters_load)); ctx_http.post("/lora-adapters/unload", ex_wrapper(routes.post_lora_adapters_unload)); + ctx_http.post("/learn", ex_wrapper(routes.post_learn)); ctx_http.get ("/train", ex_wrapper(routes.get_train)); ctx_http.post("/train", ex_wrapper(routes.post_train)); ctx_http.post("/train/cancel", ex_wrapper(routes.post_train_cancel)); From 3ee6d39d91755c71c697789c745432fa82aa7315 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 01:07:44 -0500 Subject: [PATCH 24/43] server: the learn restore runs under her own adapter set Cormac's review of 018ab27e4 (P1-b): the re-decode of the last prompt token ran with the learn set (shadow at scale 1). The shadow was already off for a batch without labels (llama-graph.cpp lora_on: n_shadow_rows is the ubatch's label count, llama-context.cpp graph_params), so the restored state was hers; her slot's own set is now applied before the restore so this does not depend on that gate. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/server-context.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 706187056127..488a0b3363a2 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2955,6 +2955,10 @@ struct server_context_impl { ggml_opt_result_loss(result, &loss, &unc); } ggml_opt_result_free(result); + // her own adapter set again before anything else is decoded for her (the shadow is + // already off for a batch without labels: llama-graph.cpp lora_on, n_shadow_rows = + // the ubatch's label count; this does not rely on it) + common_set_adapter_lora(ctx_tgt, slot->lora); // her slot goes back to her whole served prompt (attention cells, recurrent state, // tokens and checkpoints agree; her next turn reuses all of it): back to n0 from the // checkpoint, then the last prompt token decoded again under her stable set. A draft From e39b2b6cfa61284e7308789bf09c7fb6e5d97233 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 01:14:04 -0500 Subject: [PATCH 25/43] server: learning takes at most a tenth of the loop while another slot works Measured on the M5 (qwen3.5-4b hybrid, 2 slots, the 4B test server; her live lane on the same GPU was busy throughout, so compare the pairs, not the runs): learns run back to back took the other slot from p50 218 / p95 450 ms per token to 686 / 2,567. Learning must barely touch the flow. While any other slot is processing, a learn waits out a rest of t*(1/LEARN_DUTY - 1) after a learn of t (LEARN_DUTY = 0.10); with every other slot idle it runs at once. A deferred learn is bound to her slot (task.id_slot), so it comes back when a slot frees: learning lands in the gaps. Measured after (06:13:32-06:13:53Z): the other slot p50 46.8 -> 46.1, p95 73.3 -> 73.8 ms per token with learning on; the two learns waited for the gap (p50 10.4 s wall) instead of the GPU. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/server-context.cpp | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 488a0b3363a2..ca7e30b9550a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -897,6 +897,12 @@ struct server_context_impl { // given (one per server run for now), and the learning rate its next step reads int32_t learn_lora_id = -1; float learn_lr = 1e-4f; + // LEARNING BARELY TOUCHES THE FLOW: while any other slot is working, learning may take at most + // LEARN_DUTY of the loop's time (the step after a learn of t ms waits t*(1/LEARN_DUTY - 1) ms); + // with every other slot idle it runs at once. Back to back, learns took the 4B's other slot from + // p50 218 / p95 450 ms per token to 686 / 2,567 (M5, 2026-10-15). + static constexpr double LEARN_DUTY = 0.10; + int64_t learn_next_us = 0; static ggml_opt_optimizer_params learn_opt_pars(void * ud) { ggml_opt_optimizer_params p = ggml_opt_get_default_optimizer_params(nullptr); p.adamw.alpha = *(const float *) ud; @@ -2888,6 +2894,16 @@ struct server_context_impl { queue_tasks.defer(std::move(task)); break; } + { + bool others_busy = false; + for (const auto & other : slots) { + others_busy = others_busy || (other.id != slot->id && other.is_processing()); + } + if (others_busy && ggml_time_us() < learn_next_us) { + queue_tasks.defer(std::move(task)); // within the duty cycle's rest + break; + } + } auto & loras = params_base.lora_adapters; if (L.lora_id < 0 || (size_t) L.lora_id >= loras.size() || loras[L.lora_id].ptr == nullptr) { refuse("no loaded adapter " + std::to_string(L.lora_id) + " to train (POST /lora-adapters/load it first)"); @@ -2998,6 +3014,7 @@ struct server_context_impl { res->loss = loss; res->n_labels = (int32_t) toks.size() - 1; res->t_ms = (ggml_time_us() - t0) / 1000.0; + learn_next_us = ggml_time_us() + (int64_t) ((ggml_time_us() - t0) * (1.0/LEARN_DUTY - 1.0)); SRV_INF("learned on slot %d: %d labels from position %d, loss %.4f, %.1f ms%s\n", slot->id, res->n_labels, n0, loss, res->t_ms, ok ? "" : (" (failed: " + res->failure + ")").c_str()); queue_results.send(std::move(res)); @@ -5769,6 +5786,7 @@ void server_routes::init_routes() { const json body = json::parse(req.body); server_task task(SERVER_TASK_TYPE_LEARN); task.learn.id_slot = json_value(body, "id_slot", -1); + task.id_slot = task.learn.id_slot; // a deferred learn comes back when her slot frees task.learn.n_prefix = json_value(body, "n_prefix", 0); task.learn.lora_id = json_value(body, "lora_id", -1); task.learn.lr = json_value(body, "lr", 1e-4f); From b75293dccfddf9c1804ecf33098698e0e8fe85c4 Mon Sep 17 00:00:00 2001 From: joelteply Date: Thu, 15 Oct 2026 01:41:42 -0500 Subject: [PATCH 26/43] ggml-cpu: the attention backward converts each K and V row once per block, on SIMD dot products FLASH_ATTN_BACK's CPU kernel converted every f16 K and V row again for each query row (twice in pass 1), accumulated its dot products as scalar doubles, and in pass 2 regathered the strided query and output-gradient rows for every KV position. It is the reference other backends are checked against and the backward on machines with no usable GPU. Now pass 1 works blocks of 8 query rows of one head and pass 2 blocks of 16 KV positions of one KV head: each K and V row is converted once per block, each query and dO row gathered once per block, the dot products and accumulations run on ggml_vec_dot_f32 / ggml_vec_mad_f32, and D stays a double. Scratch per thread: max(8*(M + D + Dv), 16*(D + Dv)) + D + Dv floats (ops.h), sized by the work planner (ggml-cpu.c). Measured (IntelMac i7-8850H, 12 threads, test-backend-ops perf, 2026-10-15): at head 256, 16 heads over 4 KV heads, 256 rows, 2048 window, f16 cache: 2,195.5 -> 316.2 ms/call (9.78 -> 67.92 GFLOPS); the forward at the same shape is 51.6 ms. Exactness: the old and new kernels on identical seeded inputs (causal mask, GQA, f16 cache, dO at 1e-3) agree to float rounding on every part of the result: NMSE dq <= 1.97e-13, dk <= 1.49e-13, dv <= 1.00e-13, lse <= 6.6e-16, D <= 4.7e-14, scale identical, at (128,8,2,512,64), (64,4,2,300,37) and (256,16,4,2048,256). test-backend-ops grad: the native backward's MUL case OK; its SUM finite-difference case misses 1e-4 on both the old and new kernels (3/5 and 4/5 runs, random inputs), and goes to inf with dV negated, so it reaches this kernel. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE --- ggml/src/ggml-cpu/ggml-cpu.c | 4 +- ggml/src/ggml-cpu/ops.cpp | 225 +++++++++++++++++++++-------------- ggml/src/ggml-cpu/ops.h | 9 ++ 3 files changed, 146 insertions(+), 92 deletions(-) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index c2610efc0375..0ba316726bb0 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2972,11 +2972,11 @@ struct ggml_cplan ggml_graph_plan( } break; case GGML_OP_FLASH_ATTN_BACK: { - // per thread: a row's scores, then a K, V, query and output-gradient row + // per thread: the blocked kernel's scratch (ops.h) const int64_t D = node->src[0]->ne[0]; const int64_t M = node->src[1]->ne[1]; const int64_t Dv = node->src[2]->ne[0]; - cur = sizeof(float)*(M + 2*D + 2*Dv)*n_tasks; + cur = sizeof(float)*ggml_flash_attn_back_cpu_wsize(D, Dv, M)*n_tasks; } break; case GGML_OP_CROSS_ENTROPY_LOSS: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 0cdcad1c5443..d8e373701b0d 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -9268,39 +9268,38 @@ static void ggml_compute_forward_flash_attn_back_f32( const int ith = params->ith; const int nth = params->nth; - // per thread: a row's scores, then one K row, one V row and one query row as F32 - float * s = (float *) params->wdata + ith*(M + D + Dv + D + Dv); - float * kr = s + M; - float * vr = kr + D; - float * qr = vr + Dv; - float * dor = qr + D; + // Rows are worked in blocks so each K and V row is converted to F32 once per block, not once per + // query row, and the dot products run on ggml's SIMD vec_dot (pass 1: BR query rows of one head; + // pass 2: BC KV positions of one KV head). Per thread: BR rows of scores, the block's query and + // output-gradient rows, one K and one V row; pass 2 reuses the same space for its K and V block. + const int64_t BR = GGML_FA_BACK_BR; + const int64_t BC = GGML_FA_BACK_BC; + float * wbase = (float *) params->wdata + ith*ggml_flash_attn_back_cpu_wsize(D, Dv, M); const ggml_to_float_t k_to_float = ggml_get_type_traits(k->type)->to_float; const ggml_to_float_t v_to_float = ggml_get_type_traits(v->type)->to_float; - auto k_row = [&](int64_t j, int64_t hk) { + auto k_row = [&](int64_t j, int64_t hk, float * out) { const char * r = (const char *) k->data + j*k->nb[1] + hk*k->nb[2]; - if (k->type == GGML_TYPE_F32) { memcpy(kr, r, D*sizeof(float)); } else { k_to_float(r, kr, D); } + if (k->type == GGML_TYPE_F32) { memcpy(out, r, D*sizeof(float)); } else { k_to_float(r, out, D); } }; - auto v_row = [&](int64_t j, int64_t hk) { + auto v_row = [&](int64_t j, int64_t hk, float * out) { const char * r = (const char *) v->data + j*v->nb[1] + hk*v->nb[2]; - if (v->type == GGML_TYPE_F32) { memcpy(vr, r, Dv*sizeof(float)); } else { v_to_float(r, vr, Dv); } + if (v->type == GGML_TYPE_F32) { memcpy(out, r, Dv*sizeof(float)); } else { v_to_float(r, out, Dv); } }; - auto q_row = [&](int64_t i, int64_t h) { + auto q_row = [&](int64_t i, int64_t h, float * out) { for (int64_t c = 0; c < D; ++c) { - qr[c] = *(const float *) ((const char *) q->data + c*q->nb[0] + i*q->nb[1] + h*q->nb[2]); + out[c] = *(const float *) ((const char *) q->data + c*q->nb[0] + i*q->nb[1] + h*q->nb[2]); } }; - auto d_row = [&](int64_t i, int64_t h) { + auto d_row = [&](int64_t i, int64_t h, float * out) { for (int64_t c = 0; c < Dv; ++c) { - dor[c] = *(const float *) ((const char *) d->data + c*d->nb[0] + h*d->nb[1] + i*d->nb[2]); + out[c] = *(const float *) ((const char *) d->data + c*d->nb[0] + h*d->nb[1] + i*d->nb[2]); } }; auto dot = [](const float * a, const float * b, int64_t n) { - ggml_float r = 0.0; - for (int64_t c = 0; c < n; ++c) { - r += (ggml_float) a[c]*b[c]; - } - return (float) r; + float r; + ggml_vec_dot_f32((int) n, &r, 0, a, 0, b, 0, 1); + return r; }; // the dO scale a reduced-precision kernel uses (recorded so every backend's result agrees; the @@ -9317,87 +9316,133 @@ static void ggml_compute_forward_flash_attn_back_f32( dvec[N*H] = ggml_flash_attn_back_scale(mx); } - // pass 1: the query rows - for (int64_t r = ith; r < N*H; r += nth) { - const int64_t i = r % N; - const int64_t h = r / N; - const int64_t hk = h / G; - q_row(i, h); - d_row(i, h); - float smax = -INFINITY; - for (int64_t j = 0; j < M; ++j) { - k_row(j, hk); - s[j] = dot(qr, kr, D)*scale + ggml_fa_back_mask(mask, j, i, h); - smax = std::max(smax, s[j]); - } - float l = -INFINITY; - if (smax != -INFINITY) { - ggml_float sum = 0.0; + // pass 1: blocks of BR query rows of one head + { + float * s = wbase; // [BR][M] scores, then p + float * qb = s + BR*M; // [BR][D] + float * db = qb + BR*D; // [BR][Dv] + float * kr = db + BR*Dv; // [D] + float * vr = kr + D; // [Dv] + float l[GGML_FA_BACK_BR]; + ggml_float dd[GGML_FA_BACK_BR]; + + const int64_t nblk = (N + BR - 1)/BR; + for (int64_t t = ith; t < nblk*H; t += nth) { + const int64_t i0 = (t % nblk)*BR; + const int64_t h = t / nblk; + const int64_t hk = h / G; + const int64_t nb = std::min(BR, N - i0); + for (int64_t b = 0; b < nb; ++b) { + q_row(i0 + b, h, qb + b*D); + d_row(i0 + b, h, db + b*Dv); + } for (int64_t j = 0; j < M; ++j) { - sum += expf(s[j] - smax); + k_row(j, hk, kr); + for (int64_t b = 0; b < nb; ++b) { + s[b*M + j] = dot(qb + b*D, kr, D)*scale + ggml_fa_back_mask(mask, j, i0 + b, h); + } + } + for (int64_t b = 0; b < nb; ++b) { + float * sb = s + b*M; + float smax = -INFINITY; + for (int64_t j = 0; j < M; ++j) { + smax = std::max(smax, sb[j]); + } + l[b] = -INFINITY; + if (smax != -INFINITY) { + ggml_float sum = 0.0; + for (int64_t j = 0; j < M; ++j) { + sum += expf(sb[j] - smax); + } + l[b] = smax + (float) log(sum); + } + for (int64_t j = 0; j < M; ++j) { + sb[j] = l[b] == -INFINITY ? 0.0f : expf(sb[j] - l[b]); + } + dd[b] = 0.0; } - l = smax + (float) log(sum); - } - // s becomes p; D from the recomputed p and dp - ggml_float dd = 0.0; - for (int64_t j = 0; j < M; ++j) { - s[j] = l == -INFINITY ? 0.0f : expf(s[j] - l); - if (s[j] != 0.0f) { - v_row(j, hk); - dd += (ggml_float) s[j]*dot(dor, vr, Dv); + // D from the recomputed p and dp + for (int64_t j = 0; j < M; ++j) { + bool any = false; + for (int64_t b = 0; b < nb; ++b) { any = any || s[b*M + j] != 0.0f; } + if (!any) { continue; } + v_row(j, hk, vr); + for (int64_t b = 0; b < nb; ++b) { + const float p = s[b*M + j]; + if (p != 0.0f) { + dd[b] += (ggml_float) p*dot(db + b*Dv, vr, Dv); + } + } } - } - lse [h*N + i] = l; - dvec[h*N + i] = (float) dd; - float * g = dq + (h*N + i)*D; - memset(g, 0, D*sizeof(float)); - for (int64_t j = 0; j < M; ++j) { - if (s[j] == 0.0f) { - continue; + for (int64_t b = 0; b < nb; ++b) { + lse [h*N + i0 + b] = l[b]; + dvec[h*N + i0 + b] = (float) dd[b]; + memset(dq + (h*N + i0 + b)*D, 0, D*sizeof(float)); } - v_row(j, hk); - k_row(j, hk); - const float ds = s[j]*(dot(dor, vr, Dv) - (float) dd)*scale; - for (int64_t c = 0; c < D; ++c) { - g[c] += ds*kr[c]; + for (int64_t j = 0; j < M; ++j) { + bool any = false; + for (int64_t b = 0; b < nb; ++b) { any = any || s[b*M + j] != 0.0f; } + if (!any) { continue; } + k_row(j, hk, kr); + v_row(j, hk, vr); + for (int64_t b = 0; b < nb; ++b) { + const float p = s[b*M + j]; + if (p == 0.0f) { + continue; + } + const float ds = p*(dot(db + b*Dv, vr, Dv) - (float) dd[b])*scale; + ggml_vec_mad_f32((int) D, dq + (h*N + i0 + b)*D, kr, ds); + } } } } ggml_barrier(params->threadpool); - // pass 2: the KV positions, each summed over its group's query heads and every query row - for (int64_t r = ith; r < M*Hkv; r += nth) { - const int64_t j = r % M; - const int64_t hk = r / M; - float * gk = dk + (hk*M + j)*D; - float * gv = dv + (hk*M + j)*Dv; - memset(gk, 0, D*sizeof(float)); - memset(gv, 0, Dv*sizeof(float)); - k_row(j, hk); - v_row(j, hk); - for (int64_t h = hk*G; h < (hk + 1)*G; ++h) { - for (int64_t i = 0; i < N; ++i) { - const float l = lse[h*N + i]; - if (l == -INFINITY) { - continue; - } - const float m = ggml_fa_back_mask(mask, j, i, h); - if (m == -INFINITY) { - continue; - } - q_row(i, h); - const float p = expf(dot(qr, kr, D)*scale + m - l); - if (p == 0.0f) { - continue; - } - d_row(i, h); - const float ds = p*(dot(dor, vr, Dv) - dvec[h*N + i])*scale; - for (int64_t c = 0; c < D; ++c) { - gk[c] += ds*qr[c]; - } - for (int64_t c = 0; c < Dv; ++c) { - gv[c] += p*dor[c]; + // pass 2: blocks of BC KV positions of one KV head, each summed over its group's query heads and + // every query row + { + float * kb = wbase; // [BC][D] + float * vb = kb + BC*D; // [BC][Dv] + float * qr = vb + BC*Dv; // [D] + float * dr = qr + D; // [Dv] + + const int64_t nblk = (M + BC - 1)/BC; + for (int64_t t = ith; t < nblk*Hkv; t += nth) { + const int64_t j0 = (t % nblk)*BC; + const int64_t hk = t / nblk; + const int64_t nb = std::min(BC, M - j0); + for (int64_t b = 0; b < nb; ++b) { + k_row(j0 + b, hk, kb + b*D); + v_row(j0 + b, hk, vb + b*Dv); + memset(dk + (hk*M + j0 + b)*D, 0, D*sizeof(float)); + memset(dv + (hk*M + j0 + b)*Dv, 0, Dv*sizeof(float)); + } + for (int64_t h = hk*G; h < (hk + 1)*G; ++h) { + for (int64_t i = 0; i < N; ++i) { + const float l = lse[h*N + i]; + if (l == -INFINITY) { + continue; + } + bool loaded = false; + for (int64_t b = 0; b < nb; ++b) { + const float m = ggml_fa_back_mask(mask, j0 + b, i, h); + if (m == -INFINITY) { + continue; + } + if (!loaded) { + q_row(i, h, qr); + d_row(i, h, dr); + loaded = true; + } + const float p = expf(dot(qr, kb + b*D, D)*scale + m - l); + if (p == 0.0f) { + continue; + } + const float ds = p*(dot(dr, vb + b*Dv, Dv) - dvec[h*N + i])*scale; + ggml_vec_mad_f32((int) D, dk + (hk*M + j0 + b)*D, qr, ds); + ggml_vec_mad_f32((int) Dv, dv + (hk*M + j0 + b)*Dv, dr, p); + } } } } diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 07027df5102b..4506167b6d3e 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -89,6 +89,15 @@ void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_flash_attn_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +// FLASH_ATTN_BACK's blocks: BR query rows (pass 1) and BC KV positions (pass 2) share one conversion of +// each K and V row. Its per-thread scratch, in floats; the work-size planner reserves this per task. +#define GGML_FA_BACK_BR 8 +#define GGML_FA_BACK_BC 16 +static inline int64_t ggml_flash_attn_back_cpu_wsize(int64_t D, int64_t Dv, int64_t M) { + const int64_t pass1 = GGML_FA_BACK_BR*(M + D + Dv) + D + Dv; + const int64_t pass2 = GGML_FA_BACK_BC*(D + Dv) + D + Dv; + return pass1 > pass2 ? pass1 : pass2; +} void ggml_compute_forward_ssm_conv(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_ssm_scan(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_win_part(const struct ggml_compute_params * params, struct ggml_tensor * dst); From 85cac9d557bc9af68c94c1a9aa02fdc4d4e57c79 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 01:55:15 -0500 Subject: [PATCH 27/43] server: /learn takes the served conversation and checks her slot still holds its prompt Learning usually comes after review, when her slot may have served other turns since: n_prefix alone said only that the slot held that many tokens, not that they were this turn's. /learn now also takes `messages` (the conversation she was served, its last message her reply or the correction): rendered through the same chat template serving framed it with, the served prompt is everything before the reply, and her slot must hold exactly those tokens, else the learn is refused with the first differing position. One renderer: /train's render_example becomes train_render, shared by both routes, and server_trainer exposes its serving renderer. Measured (M5 Metal, qwen3.5-4b hybrid test server, 2026-10-15): a turn served through /v1/chat/completions (30 prompt tokens), then /learn with the same messages plus a corrected assistant reply, three times: loss 1.573 -> 1.085 -> 0.474, 632-687 ms each. A conversation her slot never served: refused, "slot 0 no longer holds this turn's prompt (it differs at token 23 of 30): serve the prompt again, then learn". Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/server-context.cpp | 42 ++++++++++++++++++++++++++++++++- tools/server/server-task.h | 1 + tools/server/server-train.cpp | 8 +++---- tools/server/server-train.h | 10 ++++++++ 4 files changed, 56 insertions(+), 5 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index ca7e30b9550a..1d3037f604a8 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2919,6 +2919,19 @@ struct server_context_impl { + " prompt tokens; the served prefix is " + std::to_string(L.n_prefix) + " and the reply must not be empty"); break; } + // the served prompt, when the request carries it: her slot must still hold exactly it (a + // slot that has served other turns since, or compacted, would teach the wrong context) + if (!L.prompt.empty()) { + int32_t i = 0; + while (i < L.n_prefix && slot->prompt.tokens[i] == L.prompt[i]) { + ++i; + } + if (i < L.n_prefix) { + refuse("slot " + std::to_string(slot->id) + " no longer holds this turn's prompt (it differs at token " + + std::to_string(i) + " of " + std::to_string(L.n_prefix) + "): serve the prompt again, then learn"); + break; + } + } // the recurrent part cannot be cut back: her state one token before the prompt's end const common_prompt_checkpoint * cp = nullptr; const bool needs_cp = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS; @@ -5790,7 +5803,34 @@ void server_routes::init_routes() { task.learn.n_prefix = json_value(body, "n_prefix", 0); task.learn.lora_id = json_value(body, "lora_id", -1); task.learn.lr = json_value(body, "lr", 1e-4f); - if (body.contains("reply_tokens")) { + if (body.contains("messages")) { + // the conversation she was served, its last message her reply (or the correction): rendered + // through the same template serving used; the served prompt is everything before the reply + const json & messages = body.at("messages"); + std::vector trained(messages.size(), false); + if (!messages.empty() && messages.back().value("role", std::string()) == "assistant") { + trained.back() = true; + } + std::vector toks; + std::vector loss; + std::string why; + if (!ctx_server.trainer || !train_render(messages, trained, json_value(body, "tools", json()), ctx_server.trainer->renderer(), + ctx_server.vocab, "the learn request", toks, loss, why)) { + res->error(format_error_response(why.empty() ? "the conversation could not be rendered" : why, ERROR_TYPE_INVALID_REQUEST)); + return res; + } + size_t b = 0; + while (b < loss.size() && !loss[b]) { + ++b; + } + size_t e = loss.size(); + while (e > b && !loss[e - 1]) { + --e; + } + task.learn.prompt.assign(toks.begin(), toks.begin() + b); + task.learn.reply.assign(toks.begin() + b, toks.begin() + e); + task.learn.n_prefix = (int32_t) b; + } else if (body.contains("reply_tokens")) { task.learn.reply = body.at("reply_tokens").get>(); } else { task.learn.reply = common_tokenize(ctx_server.vocab, json_value(body, "reply", std::string()), false, true); diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 5d7169276a31..150555c02ef6 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -192,6 +192,7 @@ struct server_task { int32_t id_slot = -1; int32_t n_prefix = 0; // the served prompt's tokens, before the reply std::vector reply; // what she should have said (her reply, or the reviewed correction) + std::vector prompt; // the served prompt, when known: her slot must hold exactly it int32_t lora_id = -1; // the trainable adapter (the shadow), loaded dormant float lr = 1e-4f; } learn; diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index 46e82ede20f1..845d64a213f9 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -74,7 +74,7 @@ std::string confine_out(const std::string & dir, const std::string & name) { // // Returns false with `why` when the example is malformed (the run is refused); true otherwise, // with `toks`/`loss` filled. -static bool render_example(const json & messages, const std::vector & trained, const json & tools, +bool train_render(const json & messages, const std::vector & trained, const json & tools, const server_trainer::render_fn & render, const llama_vocab * vocab, const std::string & at, std::vector & toks, std::vector & loss, std::string & why) { auto head = [&](size_t n) { // the first n messages @@ -267,7 +267,7 @@ static example_fit fit_example(const json & ex, const std::string & at, const se } toks.clear(); loss.clear(); - if (!render_example(messages, trained, tools, render, vocab, at, toks, loss, why)) { + if (!train_render(messages, trained, tools, render, vocab, at, toks, loss, why)) { return example_fit::invalid; } full = (int64_t) toks.size(); @@ -284,7 +284,7 @@ static example_fit fit_example(const json & ex, const std::string & at, const se if (dropped) { std::vector t; std::vector l; - if (!render_example(m, tr, tools, render, vocab, at, t, l, why)) { + if (!train_render(m, tr, tools, render, vocab, at, t, l, why)) { return example_fit::invalid; } *min = (int64_t) t.size(); @@ -308,7 +308,7 @@ static example_fit fit_example(const json & ex, const std::string & at, const se return example_fit::skipped; // the head plus her last exchange alone overflow the window } cut = true; - if (!render_example(messages, trained, tools, render, vocab, at, toks, loss, why)) { + if (!train_render(messages, trained, tools, render, vocab, at, toks, loss, why)) { return example_fit::invalid; } } diff --git a/tools/server/server-train.h b/tools/server/server-train.h index 69aff3982ce2..f1fb21166bd1 100644 --- a/tools/server/server-train.h +++ b/tools/server/server-train.h @@ -104,6 +104,9 @@ class server_trainer { }; ~server_trainer(); + // the serving chat template (train_render's render) + const render_fn & renderer() const { return render_chat; } + // Starts a run on a worker thread; refuses (ok=false) while one is running or on bad input. common_json start(const common_json & body_in); // Measures "examples" against "window" by the same fit rule a run trains by, training @@ -206,6 +209,13 @@ class server_trainer { // "fit": "left" for one tokenized example (see server-train.cpp): keep the window+1 tokens that // end the last trained turn (the leading BOS kept when `bos`), mask a trained turn cut at the // front, and return false when the last trained turn alone does not fit. +// A conversation through the serving chat template: its tokens, and which carry loss (the assistant +// turns marked trained), tokenized span by span so no token straddles a boundary. /train and /learn +// share it, so what learns is exactly what serving framed. +bool train_render(const common_json & messages, const std::vector & trained, const common_json & tools, + const server_trainer::render_fn & render, const llama_vocab * vocab, const std::string & at, + std::vector & toks, std::vector & loss, std::string & why); + bool train_fit_left(std::vector & toks, std::vector & loss, int64_t window, bool bos); // "fit": "middle" for one conversation (see server-train.cpp): the index of the oldest history From e77b823d9db014043d5075524441266d81e94baa Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 02:04:49 -0500 Subject: [PATCH 28/43] ggml: cross-entropy over target indices takes a weight per row ggml_cross_entropy_loss_weighted(logits, targets, weights): row i contributes w_i * -log softmax(logits_i)[t_i] / rows and its gradient scales by w_i, so how much a row counts is the caller's to say (her relevance on what she perceived), not only whether. A zero weight skips the row's softmax entirely. The weights are F32, one per row: 4 bytes a row beside a vocabulary-wide logits row. CPU and Metal implement it; CUDA refuses the weighted form in supports_op (the scheduler runs it on the CPU) until its kernel takes the weights. Measured (M5, 2026-10-15): test-backend-ops CROSS_ENTROPY_LOSS on MTL0 7/7, gradients (finite difference, weighted rows including zeros) OK on MTL0 and CPU at [10,5,4,3] and [30000,2,1,1]; full grad mode 19815/19815 on both. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/include/ggml.h | 9 ++++++++ ggml/src/ggml-cpu/ops.cpp | 12 ++++++----- ggml/src/ggml-cuda/ggml-cuda.cu | 2 ++ ggml/src/ggml-metal/ggml-metal-device.m | 6 ++++-- ggml/src/ggml-metal/ggml-metal-impl.h | 1 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 16 +++++++++----- ggml/src/ggml-metal/kernels/softmax.metal | 10 ++++++--- ggml/src/ggml.c | 16 +++++++++++++- tests/test-backend-ops.cpp | 26 +++++++++++++++++++---- 9 files changed, 78 insertions(+), 20 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index d4d92cacfe42..7f6353bb719c 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2803,6 +2803,15 @@ extern "C" { GGML_API bool ggml_cross_entropy_labels_ok(const struct ggml_tensor * logits, const struct ggml_tensor * labels); + // The same with one target index per row (I32) and a weight per row (F32, one per row): row i + // contributes w_i * -log softmax(logits_i)[t_i] / rows, and its gradient scales by w_i. How much a + // row counts is the caller's to say (a mind's relevance on what it perceived), not only whether. + GGML_API struct ggml_tensor * ggml_cross_entropy_loss_weighted( + struct ggml_context * ctx, + struct ggml_tensor * logits, + struct ggml_tensor * targets, + struct ggml_tensor * weights); + GGML_API struct ggml_tensor * ggml_cross_entropy_loss( struct ggml_context * ctx, struct ggml_tensor * a, // logits diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index d8e373701b0d..3c6786022301 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -11677,12 +11677,13 @@ static void ggml_compute_forward_cross_entropy_loss_f32( float max = -INFINITY; ggml_vec_max_f32(nc, &max, s0); if (index) { - // one target: its log softmax, -(max + log sum exp(x - max)) + x[t] + // one target: its log softmax, -(max + log sum exp(x - max)) + x[t], times the row's weight const int32_t t = ((const int32_t *) src1->data)[i1]; - if (t >= 0) { + const float w = dst->src[2] ? ((const float *) dst->src[2]->data)[i1] : 1.0f; + if (t >= 0 && w != 0.0f) { GGML_ASSERT(t < nc); const ggml_float lse = ggml_vec_log_soft_max_f32(nc, st, s0, max); - sum_thread += st[t] - lse; + sum_thread += w*(st[t] - lse); } continue; } @@ -11773,7 +11774,8 @@ static void ggml_compute_forward_cross_entropy_loss_back_f32( if (index) { // one target t: (softmax(x) - onehot(t)) * d / nr; a row without one is zero const int32_t t = ((const int32_t *) src1f->data)[i1]; - if (t < 0) { + const float w = dst->src[3] ? ((const float *) dst->src[3]->data)[i1] : 1.0f; + if (t < 0 || w == 0.0f) { memset(ds0, 0, nc*sizeof(float)); continue; } @@ -11783,7 +11785,7 @@ static void ggml_compute_forward_cross_entropy_loss_back_f32( const ggml_float sum = ggml_vec_soft_max_f32(nc, ds0, s0, max); ggml_vec_scale_f32(nc, ds0, 1.0/sum); ds0[t] -= 1.0f; - ggml_vec_scale_f32(nc, ds0, d_by_nr); + ggml_vec_scale_f32(nc, ds0, d_by_nr*w); continue; } const float * s1 = (const float *)((const char *) src1f->data + i1*src1f->nb[1]); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 1bff033486cb..10a8afb21349 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5351,7 +5351,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_FLASH_ATTN_EXT: return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); case GGML_OP_CROSS_ENTROPY_LOSS: + return op->src[2] == nullptr; // per-row weights: not on CUDA yet, they run on the CPU case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + return op->src[3] == nullptr; case GGML_OP_OPT_STEP_ADAMW: case GGML_OP_OPT_STEP_SGD: case GGML_OP_FILL: diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 1cbe2da8db5f..8faa9a66da77 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -1495,11 +1495,13 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te case GGML_OP_CROSS_ENTROPY_LOSS: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_I32) && - ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); + ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && + (op->src[2] == NULL || (op->src[2]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_I32)); case GGML_OP_CROSS_ENTROPY_LOSS_BACK: return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_I32) && - ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) && ggml_is_contiguous(op); + ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) && ggml_is_contiguous(op) && + (op->src[3] == NULL || (op->src[3]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_I32)); case GGML_OP_OUT_PROD: // The backward through a frozen weight: out_prod(W, grad^T). Composed on this // backend from a dequantizing copy, two transposing copies and the F32 simdgroup diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index e1b2b3c6ea35..6ef1f69dafe2 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -902,6 +902,7 @@ typedef struct { uint64_t nb01; uint64_t nb11; uint64_t nb1; + int32_t has_weights; // a per-row F32 weight follows the labels (I32 targets only) } ggml_metal_kargs_cross_entropy_loss; // FLASH_ATTN_BACK (ggml_flash_attn_back): q [D, N, H], k [D, M, Hkv], v [Dv, M, Hkv], d [Dv, H, N], diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 4274ad63aa8b..17ffba056b1e 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -4064,12 +4064,14 @@ static int ggml_metal_op_cross_entropy_loss_rows(ggml_metal_op_t ctx, const ggml const int32_t ne00 = logits->ne[0]; const int32_t nr = ggml_nrows(logits); + const ggml_tensor * weights = back ? op->src[3] : op->src[2]; ggml_metal_kargs_cross_entropy_loss args = { - /*.ne00 =*/ ne00, - /*.nr =*/ nr, - /*.nb01 =*/ logits->nb[1], - /*.nb11 =*/ labels->nb[1], - /*.nb1 =*/ nb1, + /*.ne00 =*/ ne00, + /*.nr =*/ nr, + /*.nb01 =*/ logits->nb[1], + /*.nb11 =*/ labels->nb[1], + /*.nb1 =*/ nb1, + /*.has_weights =*/ weights != nullptr, }; auto pipeline = ggml_metal_library_get_pipeline_cross_entropy_loss(lib, op); @@ -4082,15 +4084,19 @@ static int ggml_metal_op_cross_entropy_loss_rows(ggml_metal_op_t ctx, const ggml ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + // the weights ride in the last slot (the labels' buffer stands in when there are none: never read) + const ggml_metal_buffer_id bid_w = ggml_metal_get_buffer_id(weights ? weights : labels); if (back) { ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(logits), 2); ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(labels), 3); ggml_metal_encoder_set_buffer(enc, bid_dst, 4); + ggml_metal_encoder_set_buffer(enc, bid_w, 5); } else { ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(logits), 1); ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(labels), 2); ggml_metal_encoder_set_buffer(enc, bid_dst, 3); + ggml_metal_encoder_set_buffer(enc, bid_w, 4); } ggml_metal_encoder_set_threadgroup_memory_size(enc, 32*sizeof(float), 0); ggml_metal_encoder_dispatch_threadgroups(enc, nr, 1, 1, nth, 1, 1); diff --git a/ggml/src/ggml-metal/kernels/softmax.metal b/ggml/src/ggml-metal/kernels/softmax.metal index b3bc06f4d737..d7fae3830425 100644 --- a/ggml/src/ggml-metal/kernels/softmax.metal +++ b/ggml/src/ggml-metal/kernels/softmax.metal @@ -357,6 +357,7 @@ kernel void kernel_cross_entropy_loss_i32( device const char * src0, device const char * src1, device char * dst, + device const float * weights, threadgroup float * buf [[threadgroup(0)]], uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], @@ -381,7 +382,8 @@ kernel void kernel_cross_entropy_loss_i32( const float lse = log(ce_reduce_sum(lsum, buf, sgitg, tiisg, tptg)); if (tpitg.x == 0) { - ((device float *) dst)[row] = t >= 0 ? -(s0[t] - max_val - lse)/(float) args.nr : 0.0f; + const float w = args.has_weights ? weights[row] : 1.0f; + ((device float *) dst)[row] = t >= 0 && w != 0.0f ? -w*(s0[t] - max_val - lse)/(float) args.nr : 0.0f; } } @@ -391,6 +393,7 @@ kernel void kernel_cross_entropy_loss_back_i32( device const char * src1, // logits device const char * src2, // labels: one target per row device char * dst, + device const float * weights, threadgroup float * buf [[threadgroup(0)]], uint3 tgpig[[threadgroup_position_in_grid]], uint3 tpitg[[thread_position_in_threadgroup]], @@ -403,14 +406,15 @@ kernel void kernel_cross_entropy_loss_back_i32( device const float * s0 = (device const float *) (src1 + row*args.nb01); device float * d = (device float *) (dst + row*args.nb1); - if (t < 0) { + const float w = args.has_weights ? weights[row] : 1.0f; + if (t < 0 || w == 0.0f) { for (int i = tpitg.x; i < args.ne00; i += tptg.x) { d[i] = 0.0f; } return; } - const float d_by_nr = ((device const float *) src0)[0]/(float) args.nr; + const float d_by_nr = w*((device const float *) src0)[0]/(float) args.nr; float lmax = -INFINITY; for (int i = tpitg.x; i < args.ne00; i += tptg.x) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index e42f71627d91..c4ffc4ddcaf4 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6211,6 +6211,18 @@ struct ggml_tensor * ggml_cross_entropy_loss( return result; } +struct ggml_tensor * ggml_cross_entropy_loss_weighted( + struct ggml_context * ctx, + struct ggml_tensor * logits, + struct ggml_tensor * targets, + struct ggml_tensor * weights) { + GGML_ASSERT(targets->type == GGML_TYPE_I32 && ggml_cross_entropy_labels_ok(logits, targets)); + GGML_ASSERT(weights->type == GGML_TYPE_F32 && ggml_is_contiguous(weights) && ggml_nelements(weights) == ggml_nrows(logits)); + struct ggml_tensor * result = ggml_cross_entropy_loss(ctx, logits, targets); + result->src[2] = weights; + return result; +} + // ggml_cross_entropy_loss_back struct ggml_tensor * ggml_cross_entropy_loss_back( @@ -7501,7 +7513,9 @@ static void ggml_compute_backward( } break; case GGML_OP_CROSS_ENTROPY_LOSS: { if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_cross_entropy_loss_back(ctx, grad, src0, src1)); + struct ggml_tensor * back = ggml_cross_entropy_loss_back(ctx, grad, src0, src1); + back->src[3] = src2; // the per-row weights, when the loss has them + ggml_add_or_set(ctx, cgraph, isrc0, back); } GGML_ASSERT(!src1_needs_grads && "backward pass for labels not implemented"); } break; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 308eea0837e3..87341ca5cc92 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7606,10 +7606,12 @@ struct test_cross_entropy_loss_index : public test_case { const std::array ne; std::string vars() override { - return VARS_TO_STR1(ne); + return VARS_TO_STR2(ne, weighted); } - test_cross_entropy_loss_index(std::array ne = {10, 5, 4, 3}) : ne(ne) {} + const bool weighted; // a weight per row: how much each row counts, not only whether + + test_cross_entropy_loss_index(std::array ne = {10, 5, 4, 3}, bool weighted = false) : ne(ne), weighted(weighted) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * logits = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); @@ -7619,7 +7621,14 @@ struct test_cross_entropy_loss_index : public test_case { ggml_tensor * targets = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[1]*ne[2]*ne[3]); ggml_set_name(targets, "targets"); - ggml_tensor * out = ggml_cross_entropy_loss(ctx, logits, targets); + ggml_tensor * out; + if (weighted) { + ggml_tensor * weights = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[1]*ne[2]*ne[3]); + ggml_set_name(weights, "weights"); + out = ggml_cross_entropy_loss_weighted(ctx, logits, targets, weights); + } else { + out = ggml_cross_entropy_loss(ctx, logits, targets); + } ggml_set_name(out, "out"); return out; @@ -7628,7 +7637,14 @@ struct test_cross_entropy_loss_index : public test_case { void initialize_tensors(ggml_context * ctx) override { std::default_random_engine rng(7); for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { - if (t->type == GGML_TYPE_I32) { + if (strcmp(t->name, "weights") == 0) { + // a spread of relevance, zero included (a row the mind did not care about) + std::vector w(ggml_nelements(t)); + for (size_t i = 0; i < w.size(); ++i) { + w[i] = i % 4 == 1 ? 0.0f : 0.25f + 0.5f*(float) (i % 5); + } + ggml_backend_tensor_set(t, w.data(), 0, ggml_nbytes(t)); + } else if (t->type == GGML_TYPE_I32) { // every third row masked (-1), the rest a random target std::uniform_int_distribution pick(0, (int32_t) ne[0] - 1); std::vector v(ggml_nelements(t)); @@ -10530,6 +10546,8 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_back(512, 4, 2, 128, 9, GGML_TYPE_F16)); // the tiled path's tiles exceed 32 KiB: the per-row path test_cases.emplace_back(new test_cross_entropy_loss_index()); test_cases.emplace_back(new test_cross_entropy_loss_index({30000, 2, 1, 1})); // a vocabulary's width; few rows keep the finite difference above f32 rounding + test_cases.emplace_back(new test_cross_entropy_loss_index({10, 5, 4, 3}, true)); + test_cases.emplace_back(new test_cross_entropy_loss_index({30000, 2, 1, 1}, true)); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); From 01359323e0491c09e0a31ae018168847d1b11009 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 02:38:33 -0500 Subject: [PATCH 29/43] train: the attention backward computes K/V gradients for the run's rows only The walk reads the K/V gradient of the run's positions only ([g0, n_win)); everything before g0 is her served prefix, a constant. The attention backward computed dK/dV for the whole window anyway and the graph discarded all but the run's rows. - ggml_flash_attn_back(..., kv0): dK/dV for the key rows [kv0, M) only, row j at j - kv0 (op_params[1]); pass 1 (dQ, lse, D) still over every key. CPU and Metal (tiled and per-row). - ggml_flash_attn_ext_set_grad_sinks(a, k_tap, k_cur, v_tap, v_cur): an attention whose K/V gradient goes to the run's rows (the tap of the cached rows before the chunk, then the chunk's own K/V) instead of the window it reads; its backward takes kv0 = M - (tap + chunk rows) and adds each part to its sink. The composed backward (devices without the one op) skips the tiles before kv0 for dK/dV the same way. - The in-place training graph uses the sinks, so no window-sized K/V gradient exists: the tap is now the run's cached rows as F32 (as on the explicit path), and ggml_opt_accumulate_grad_rows, which tapped rows of a window-sized gradient, is deleted with its only caller. Measured (M5 Metal, 2026-10-15): - test-walk-exact, qwen2.5-0.5b and qwen3.5-4b hybrid: every exact row cosine 1.0000 (f16 and q8_0 caches, context + reply, recompute, rollback slots, learning on a live sequence); caches after the walk byte-identical to her stable decode. - test-backend-ops MTL0 vs CPU: FLASH_ATTN_BACK 8/8 (kv0 = M-37 unaligned, q8_0 at 203, per-row head 512 at 100); grad mode 19818/19818. - FLASH_ATTN_BACK at head 256, 16/4 heads, 256 rows, window 2048: 66.0 ms per call at kv0 = 0, 45.7 ms at kv0 = 1792 (the reply's keys). Pass 1 over every key is what remains. - Graph peak, 4B, window 8192, chunk 512, recompute: 1,084.8 MiB, unchanged: the per-layer backward results were never alive together, so at this window the head slices set the peak. The long-window peak is measured next. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/include/ggml-opt.h | 5 - ggml/include/ggml.h | 18 +++- ggml/src/ggml-cpu/ops.cpp | 16 ++-- ggml/src/ggml-metal/ggml-metal-impl.h | 1 + ggml/src/ggml-metal/ggml-metal-ops.cpp | 6 +- ggml/src/ggml-metal/kernels/fa.metal | 13 +-- ggml/src/ggml-opt.cpp | 22 ----- ggml/src/ggml.c | 128 ++++++++++++++++++++----- src/llama-context-opt.cpp | 13 +-- src/llama-graph.cpp | 54 +++++++---- src/llama-graph.h | 4 + tests/test-backend-ops.cpp | 17 +++- 12 files changed, 195 insertions(+), 102 deletions(-) diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index 8de657ed299c..6a78d2795290 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h @@ -284,11 +284,6 @@ extern "C" { // zeroes it, except at its own first allocation, which precedes any accumulation. GGML_API void ggml_opt_accumulate_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, struct ggml_tensor * acc); - // The same for rows [row0, row0 + acc->ne[1]) of a 2-D `tensor`'s gradient only: the rest of - // its gradient lives on the device for its consumers and is never kept, so `acc` holds a - // span (a reply's positions) of a window-sized tensor (a cache's cells). - GGML_API void ggml_opt_accumulate_grad_rows(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, int64_t row0, struct ggml_tensor * acc); - // ############################################################################ // ## The high-level functions start here. They do not depend on any private ## // ## functions or structs and can be copied to and adapted for user code. ## diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 7f6353bb719c..701410e34e11 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2487,6 +2487,9 @@ extern "C" { // ggml_flash_attn_back_offset, then each query row's log-sum-exp and D = sum_j p_ij dp_ij, // which the kernels pass between their passes, and the power of two dO is scaled by in a // kernel's reduced-precision tiles (ggml_flash_attn_back_scale). + // dK and dV cover the key rows [kv0, M) only (row j at j - kv0): keys before kv0 are constants + // whose gradient nobody reads (a frozen served prefix), so it is neither computed nor stored. + // dQ always sums over every key. GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, struct ggml_tensor * q, @@ -2494,7 +2497,8 @@ extern "C" { struct ggml_tensor * v, struct ggml_tensor * mask, struct ggml_tensor * d, - float scale); + float scale, + int64_t kv0); // byte offset in a ggml_flash_attn_back result of: 0 dQ, 1 dK, 2 dV, 3 the per-row scratch // the power of two that brings the largest |dO| to [4, 8): training gradients sit far below half @@ -2503,6 +2507,18 @@ extern "C" { GGML_API float ggml_flash_attn_back_scale(float max_abs_do); GGML_API size_t ggml_flash_attn_back_offset(const struct ggml_tensor * back, int part); + // Where FLASH_ATTN_EXT's K/V gradient goes when only the last rows of its K/V can carry one: + // k_tap / v_tap (or NULL) and k_cur / v_cur, each [D, Hkv, rows], hold the values of K/V's last + // rows in order, ending at M (a run's cached rows, then the chunk's own). The backward computes + // dK/dV from the first of those rows only and adds each part to its sink; the K/V the attention + // reads (the whole window) get no gradient. The sinks are not read by the forward. + GGML_API void ggml_flash_attn_ext_set_grad_sinks( + struct ggml_tensor * a, + struct ggml_tensor * k_tap, + struct ggml_tensor * k_cur, + struct ggml_tensor * v_tap, + struct ggml_tensor * v_cur); + // the backward of FLASH_ATTN_EXT as the one op above (ggml_flash_attn_ext_set_back_tile) #define GGML_FA_BACK_NATIVE (-1) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 3c6786022301..dedb124f3147 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -9249,12 +9249,14 @@ static void ggml_compute_forward_flash_attn_back_f32( const ggml_tensor * mask = dst->src[3]; const ggml_tensor * d = dst->src[4]; - const float scale = ggml_get_op_params_f32(dst, 0); + const float scale = ggml_get_op_params_f32(dst, 0); + const int64_t kv0 = ggml_get_op_params_i32(dst, 1); // dK/dV cover the key rows [kv0, M) const int64_t D = q->ne[0]; const int64_t N = q->ne[1]; const int64_t H = q->ne[2]; const int64_t M = k->ne[1]; + const int64_t L = M - kv0; const int64_t Hkv = k->ne[2]; const int64_t Dv = v->ne[0]; const int64_t G = H / Hkv; // query head h reads KV head h / G @@ -9407,16 +9409,16 @@ static void ggml_compute_forward_flash_attn_back_f32( float * qr = vb + BC*Dv; // [D] float * dr = qr + D; // [Dv] - const int64_t nblk = (M + BC - 1)/BC; + const int64_t nblk = (L + BC - 1)/BC; for (int64_t t = ith; t < nblk*Hkv; t += nth) { - const int64_t j0 = (t % nblk)*BC; + const int64_t j0 = kv0 + (t % nblk)*BC; const int64_t hk = t / nblk; const int64_t nb = std::min(BC, M - j0); for (int64_t b = 0; b < nb; ++b) { k_row(j0 + b, hk, kb + b*D); v_row(j0 + b, hk, vb + b*Dv); - memset(dk + (hk*M + j0 + b)*D, 0, D*sizeof(float)); - memset(dv + (hk*M + j0 + b)*Dv, 0, Dv*sizeof(float)); + memset(dk + (hk*L + j0 - kv0 + b)*D, 0, D*sizeof(float)); + memset(dv + (hk*L + j0 - kv0 + b)*Dv, 0, Dv*sizeof(float)); } for (int64_t h = hk*G; h < (hk + 1)*G; ++h) { for (int64_t i = 0; i < N; ++i) { @@ -9440,8 +9442,8 @@ static void ggml_compute_forward_flash_attn_back_f32( continue; } const float ds = p*(dot(dr, vb + b*Dv, Dv) - dvec[h*N + i])*scale; - ggml_vec_mad_f32((int) D, dk + (hk*M + j0 + b)*D, qr, ds); - ggml_vec_mad_f32((int) Dv, dv + (hk*M + j0 + b)*Dv, dr, p); + ggml_vec_mad_f32((int) D, dk + (hk*L + j0 - kv0 + b)*D, qr, ds); + ggml_vec_mad_f32((int) Dv, dv + (hk*L + j0 - kv0 + b)*Dv, dr, p); } } } diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 6ef1f69dafe2..74a4bcbd15a6 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -930,6 +930,7 @@ typedef struct { uint64_t offs_dk; uint64_t offs_dv; uint64_t offs_lse; + int32_t kv0; // dK/dV cover the key rows [kv0, M), row j at j - kv0 } ggml_metal_kargs_flash_attn_back; // a per-row backward with one scalar parameter (eps for RMS_NORM_BACK, scale for SOFT_MAX_BACK) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 17ffba056b1e..27c07e76939b 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -4169,7 +4169,9 @@ int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx) { /*.offs_dk =*/ ggml_flash_attn_back_offset(op, 1), /*.offs_dv =*/ ggml_flash_attn_back_offset(op, 2), /*.offs_lse =*/ ggml_flash_attn_back_offset(op, 3), + /*.kv0 =*/ ggml_get_op_params_i32(op, 1), }; + const int L = args.M - args.kv0; // the key rows a dK/dV is computed for auto run = [&](const char * pass, int ngx, int ngy, int nth, size_t shmem) { auto pipeline = ggml_metal_library_get_pipeline_flash_attn_back(lib, op, pass); @@ -4198,12 +4200,12 @@ int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx) { // and simdgroup matrices exist only on Apple7+ GPUs (M1 and newer): older Macs take the per-row path if (props_dev->has_simdgroup_mm && args.D == args.Dv && args.D % 32 == 0 && shmem <= props_dev->max_theadgroup_memory_size) { run("stats", (args.N + br - 1)/br, args.H, 32*nsg, shmem); - run("dkv", (args.M + bc - 1)/bc, args.Hkv, 32*nsg, shmem); + run("dkv", (L + bc - 1)/bc, args.Hkv, 32*nsg, shmem); run("dq", (args.N + br - 1)/br, args.H, 32*nsg, shmem); } else { // one simdgroup per row run("q", args.N*args.H, 1, 32, 0); - run("kv", args.M*args.Hkv, 1, 32, 0); + run("kv", L*args.Hkv, 1, 32, 0); } return 1; diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index c210b005e6b8..58c745834ee7 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -2380,8 +2380,9 @@ kernel void kernel_flash_attn_back_kv( device char * dst, uint tgpig[[threadgroup_position_in_grid]], ushort tiisg[[thread_index_in_simdgroup]]) { - const int j = tgpig % args.M; - const int hk = tgpig / args.M; + const int L = args.M - args.kv0; + const int j = args.kv0 + tgpig % L; + const int hk = tgpig / L; const int G = args.H / args.Hkv; const int nq = args.D / 32; @@ -2425,8 +2426,8 @@ kernel void kernel_flash_attn_back_kv( } } } - device float * dk = (device float *) (dst + args.offs_dk) + (hk*args.M + j)*args.D; - device float * dv = (device float *) (dst + args.offs_dv) + (hk*args.M + j)*args.Dv; + device float * dk = (device float *) (dst + args.offs_dk) + (hk*L + j - args.kv0)*args.D; + device float * dv = (device float *) (dst + args.offs_dv) + (hk*L + j - args.kv0)*args.Dv; for (int c = 0; c < nq; ++c) { dk[tiisg + 32*c] = gk[c]; } @@ -2631,7 +2632,7 @@ kernel void kernel_flash_attn_back_dkv( const float ds = ((device const float *) (dst + args.offs_lse))[2*args.N*args.H]; // the dO scale const int D = args.D; const int hk = tgpig.y; - const int j0 = tgpig.x*FAB_BC; + const int j0 = args.kv0 + tgpig.x*FAB_BC; const int G = args.H / args.Hkv; const int CW = D / FAB_NSG; // the columns this simdgroup owns const int c0 = sgitg*CW; @@ -2697,7 +2698,7 @@ kernel void kernel_flash_attn_back_dkv( } } threadgroup_barrier(mem_flags::mem_threadgroup); - device float * out = (device float *) (dst + (part == 0 ? args.offs_dk : args.offs_dv)) + (hk*args.M + j0)*D; + device float * out = (device float *) (dst + (part == 0 ? args.offs_dk : args.offs_dv)) + (hk*(args.M - args.kv0) + j0 - args.kv0)*D; for (int e = tiitg; e < FAB_BC*D; e += FAB_NSG*32) { if (j0 + e / D < args.M) { out[e] = so[e]/ds; diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index c0db1dfb594c..a6b528fb3228 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -89,9 +89,6 @@ struct ggml_opt_context { // ggml_opt_accumulate_grad: GRAD nodes whose gradient adds into a caller's tensor, for the // graph prepared last (ggml_opt_prepare_alloc clears them) std::unordered_map next_grad_accs; - // ggml_opt_accumulate_grad_rows: GRAD nodes whose gradient's rows [row0, row0 + acc rows) add - // into a caller's tensor, for the graph prepared last - std::unordered_map> next_grad_rows; std::vector grad_accs; // ggml_opt_prepare_alloc_head: the outputs' row slices, in row order (empty: one outputs) std::vector output_slices; @@ -747,15 +744,6 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { opt_ctx->leaf_grads.clear(); for (int i = 0; i < opt_ctx->gf->n_nodes; ++i) { ggml_tensor * node = opt_ctx->gf->nodes[i]; - if (const auto it = opt_ctx->next_grad_rows.find(node); it != opt_ctx->next_grad_rows.end()) { - ggml_tensor * grad = ggml_graph_get_grad(opt_ctx->gb_grad, node); - if (grad) { - const auto & [row0, acc] = it->second; - ggml_tensor * rows = ggml_view_2d(opt_ctx->ctx_compute, grad, acc->ne[0], acc->ne[1], grad->nb[1], row0*grad->nb[1]); - ggml_build_forward_expand(opt_ctx->gb_grad, ggml_add_inplace(opt_ctx->ctx_compute, acc, rows)); - } - continue; - } if ((node->flags & GGML_TENSOR_FLAG_GRAD) && !opt_ctx->next_grad_accs.count(node)) { ggml_tensor * grad = ggml_graph_get_grad(opt_ctx->gb_grad, node); if (grad) { @@ -889,15 +877,6 @@ void ggml_opt_accumulate_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * t opt_ctx->next_grad_accs[tensor] = acc; } -void ggml_opt_accumulate_grad_rows(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, int64_t row0, struct ggml_tensor * acc) { - GGML_ASSERT(!opt_ctx->eval_ready && "register an accumulator before ggml_opt_alloc"); - GGML_ASSERT(acc->type == GGML_TYPE_F32 && ggml_n_dims(tensor) <= 2 && acc->ne[0] == tensor->ne[0]); - GGML_ASSERT(row0 >= 0 && row0 + acc->ne[1] <= tensor->ne[1]); - GGML_ASSERT((acc->view_src ? acc->view_src : acc)->buffer != nullptr && "the accumulator lives in the caller's buffer"); - ggml_set_grad(tensor); - opt_ctx->next_grad_rows[tensor] = { row0, acc }; -} - struct ggml_tensor * ggml_opt_leaf_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * leaf) { const auto it = opt_ctx->leaf_grads.find(leaf); return it == opt_ctx->leaf_grads.end() ? nullptr : it->second; @@ -1038,7 +1017,6 @@ void ggml_opt_prepare_alloc( opt_ctx->inputs = inputs; opt_ctx->outputs = outputs; opt_ctx->next_grad_accs.clear(); // registrations belong to the graph prepared last - opt_ctx->next_grad_rows.clear(); opt_ctx->output_slices.clear(); } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index c4ffc4ddcaf4..0f990953c35a 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5524,6 +5524,26 @@ void ggml_flash_attn_ext_set_back_tile( ggml_set_op_params_i32(a, 4, tile); // scale, max_bias, softcap, prec come first } +void ggml_flash_attn_ext_set_grad_sinks( + struct ggml_tensor * a, + struct ggml_tensor * k_tap, + struct ggml_tensor * k_cur, + struct ggml_tensor * v_tap, + struct ggml_tensor * v_cur) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT && k_cur && v_cur && (k_tap == NULL) == (v_tap == NULL)); + const struct ggml_tensor * k = a->src[1]; + const struct ggml_tensor * v = a->src[2]; + const int64_t n = k_cur->ne[2] + (k_tap ? k_tap->ne[2] : 0); + GGML_ASSERT(v_cur->ne[2] + (v_tap ? v_tap->ne[2] : 0) == n && n <= k->ne[1]); + GGML_ASSERT(k_cur->ne[0] == k->ne[0] && k_cur->ne[1] == k->ne[2] && v_cur->ne[0] == v->ne[0] && v_cur->ne[1] == v->ne[2]); + GGML_ASSERT(!k_tap || (k_tap->ne[0] == k->ne[0] && k_tap->ne[1] == k->ne[2] && v_tap->ne[0] == v->ne[0] && v_tap->ne[1] == v->ne[2])); + // src[4] is sinks (the attention's own); the gradient sinks follow it + a->src[5] = k_tap; + a->src[6] = k_cur; + a->src[7] = v_tap; + a->src[8] = v_cur; +} + void ggml_flash_attn_ext_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks) { @@ -5556,10 +5576,11 @@ size_t ggml_flash_attn_back_offset(const struct ggml_tensor * back, int part) { const struct ggml_tensor * q = back->src[0]; const struct ggml_tensor * k = back->src[1]; const struct ggml_tensor * v = back->src[2]; + const int64_t kv0 = ggml_get_op_params_i32(back, 1); // dK/dV hold the key rows [kv0, M) const size_t sizes[3] = { GGML_PAD((size_t) q->ne[0]*q->ne[1]*q->ne[2]*sizeof(float), GGML_MEM_ALIGN), - GGML_PAD((size_t) ggml_nelements(k)*sizeof(float), GGML_MEM_ALIGN), - GGML_PAD((size_t) ggml_nelements(v)*sizeof(float), GGML_MEM_ALIGN), + GGML_PAD((size_t) k->ne[0]*(k->ne[1] - kv0)*k->ne[2]*sizeof(float), GGML_MEM_ALIGN), + GGML_PAD((size_t) v->ne[0]*(v->ne[1] - kv0)*v->ne[2]*sizeof(float), GGML_MEM_ALIGN), }; size_t offs = 0; for (int i = 0; i < part; ++i) { @@ -5575,8 +5596,10 @@ struct ggml_tensor * ggml_flash_attn_back( struct ggml_tensor * v, struct ggml_tensor * mask, struct ggml_tensor * d, - float scale) { + float scale, + int64_t kv0) { // q [D, N, H, 1], k [D, M, Hkv, 1], v [Dv, M, Hkv, 1], d [Dv, H, N, 1], mask [M', N', 1|H, 1] + GGML_ASSERT(kv0 >= 0 && kv0 < k->ne[1] && kv0 <= INT32_MAX); GGML_ASSERT(q->type == GGML_TYPE_F32 && d->type == GGML_TYPE_F32); GGML_ASSERT(q->ne[3] == 1 && k->ne[3] == 1 && v->ne[3] == 1 && "one sequence stream"); GGML_ASSERT(k->ne[0] == q->ne[0] && v->ne[1] == k->ne[1] && v->ne[2] == k->ne[2]); @@ -5587,10 +5610,12 @@ struct ggml_tensor * ggml_flash_attn_back( const size_t rows = (size_t) q->ne[1]*q->ne[2]; struct ggml_tensor tmp = { 0 }; tmp.src[0] = q; tmp.src[1] = k; tmp.src[2] = v; + ggml_set_op_params_i32(&tmp, 1, (int32_t) kv0); const size_t bytes = ggml_flash_attn_back_offset(&tmp, 3) + (2*rows + 1)*sizeof(float); // lse, D, the dO scale struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) (bytes/sizeof(float))); ggml_set_op_params_f32(result, 0, scale); + ggml_set_op_params_i32(result, 1, (int32_t) kv0); result->op = GGML_OP_FLASH_ATTN_BACK; result->src[0] = q; @@ -6815,9 +6840,45 @@ static struct ggml_tensor * ggml_fa_back_join(struct ggml_context * ctx, struct return ggml_concat(ctx, ggml_fa_back_join(ctx, parts, h), ggml_fa_back_join(ctx, parts + h, n - h), 1); } +// whether src needs a gradient in this backward pass +static bool ggml_backward_needs_grad(struct ggml_cgraph * cgraph, const bool * grads_needed, const struct ggml_tensor * src, size_t * isrc) { + if (!src) { + return false; + } + *isrc = ggml_hash_find(&cgraph->visited_hash_set, src); + return *isrc != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, *isrc) && grads_needed[*isrc]; +} + +// FLASH_ATTN_EXT's K (part 0) or V (part 1) gradient g [D, M - kv0, Hkv] (key rows [kv0, M)): to the +// grad sinks when the attention has them (ggml_flash_attn_ext_set_grad_sinks), else to K/V itself +static void ggml_fa_back_route_kv(struct ggml_context * ctx, struct ggml_cgraph * cgraph, const bool * grads_needed, + struct ggml_tensor * tensor, int part, struct ggml_tensor * g, size_t isrc, bool needed) { + if (tensor->src[6] == NULL) { + if (needed) { + const struct ggml_tensor * kv = tensor->src[1 + part]; + ggml_add_or_set(ctx, cgraph, isrc, ggml_reshape_4d(ctx, ggml_cont(ctx, g), kv->ne[0], kv->ne[1], kv->ne[2], kv->ne[3])); + } + return; + } + int64_t r = 0; + for (int s = 0; s < 2; ++s) { + struct ggml_tensor * sink = tensor->src[5 + 2*part + s]; // tap, then the chunk's own + size_t is = 0; + if (sink == NULL) { + continue; + } + const int64_t n = sink->ne[2]; + if (ggml_backward_needs_grad(cgraph, grads_needed, sink, &is)) { + struct ggml_tensor * rows = ggml_view_3d(ctx, g, g->ne[0], n, g->ne[2], g->nb[1], g->nb[2], r*g->nb[1]); + ggml_add_or_set(ctx, cgraph, is, ggml_cont(ctx, ggml_permute(ctx, rows, 0, 2, 1, 3))); // [D, Hkv, n] + } + r += n; + } +} + static void ggml_flash_attn_ext_backward( struct ggml_context * ctx, struct ggml_cgraph * cgraph, struct ggml_tensor * tensor, struct ggml_tensor * grad, - size_t isrc0, size_t isrc1, size_t isrc2, bool dq_needed, bool dk_needed, bool dv_needed) { + size_t isrc0, size_t isrc1, size_t isrc2, bool dq_needed, bool dk_needed, bool dv_needed, const bool * grads_needed) { struct ggml_tensor * q = tensor->src[0]; // [D, N, H, 1] struct ggml_tensor * k = tensor->src[1]; // [D, M, Hkv, 1] struct ggml_tensor * v = tensor->src[2]; // [Dv, M, Hkv, 1] @@ -6832,25 +6893,31 @@ static void ggml_flash_attn_ext_backward( GGML_ASSERT(tensor->src[4] == NULL && "flash_attn_ext backward: sinks are not implemented"); GGML_ASSERT(q->ne[3] == 1 && "flash_attn_ext backward: one sequence stream"); + // with grad sinks only K/V's last rows carry a gradient: from kv0 on, and nothing to K/V itself + const bool sunk = tensor->src[6] != NULL; + const int64_t kv0 = sunk ? k->ne[1] - tensor->src[6]->ne[2] - (tensor->src[5] ? tensor->src[5]->ne[2] : 0) : 0; + if (sunk) { + size_t is = 0; + dk_needed = ggml_backward_needs_grad(cgraph, grads_needed, tensor->src[5], &is) || ggml_backward_needs_grad(cgraph, grads_needed, tensor->src[6], &is); + dv_needed = ggml_backward_needs_grad(cgraph, grads_needed, tensor->src[7], &is) || ggml_backward_needs_grad(cgraph, grads_needed, tensor->src[8], &is); + } + int32_t tile = ggml_get_op_params_i32(tensor, 4); if (tile == GGML_FA_BACK_NATIVE) { - struct ggml_tensor * b = ggml_flash_attn_back(ctx, q, k, v, mask, grad, scale); - // views of the packed result, each a contiguous tensor shaped like its source + struct ggml_tensor * b = ggml_flash_attn_back(ctx, q, k, v, mask, grad, scale, kv0); + // views of the packed result: dQ shaped like q, dK/dV like k/v from row kv0 + const int64_t L = k->ne[1] - kv0; struct ggml_tensor * dq = ggml_view_3d(ctx, b, q->ne[0], q->ne[1], q->ne[2], q->ne[0]*sizeof(float), q->ne[0]*q->ne[1]*sizeof(float), ggml_flash_attn_back_offset(b, 0)); - struct ggml_tensor * dk = ggml_view_4d(ctx, b, k->ne[0], k->ne[1], k->ne[2], k->ne[3], k->ne[0]*sizeof(float), - k->ne[0]*k->ne[1]*sizeof(float), k->ne[0]*k->ne[1]*k->ne[2]*sizeof(float), ggml_flash_attn_back_offset(b, 1)); - struct ggml_tensor * dv = ggml_view_4d(ctx, b, v->ne[0], v->ne[1], v->ne[2], v->ne[3], v->ne[0]*sizeof(float), - v->ne[0]*v->ne[1]*sizeof(float), v->ne[0]*v->ne[1]*v->ne[2]*sizeof(float), ggml_flash_attn_back_offset(b, 2)); + struct ggml_tensor * dk = ggml_view_3d(ctx, b, k->ne[0], L, k->ne[2], k->ne[0]*sizeof(float), k->ne[0]*L*sizeof(float), + ggml_flash_attn_back_offset(b, 1)); + struct ggml_tensor * dv = ggml_view_3d(ctx, b, v->ne[0], L, v->ne[2], v->ne[0]*sizeof(float), v->ne[0]*L*sizeof(float), + ggml_flash_attn_back_offset(b, 2)); if (dq_needed) { ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape_4d(ctx, ggml_cont(ctx, dq), q->ne[0], q->ne[1], q->ne[2], q->ne[3])); } - if (dk_needed) { - ggml_add_or_set(ctx, cgraph, isrc1, dk); - } - if (dv_needed) { - ggml_add_or_set(ctx, cgraph, isrc2, dv); - } + ggml_fa_back_route_kv(ctx, cgraph, grads_needed, tensor, 0, dk, isrc1, dk_needed); + ggml_fa_back_route_kv(ctx, cgraph, grads_needed, tensor, 1, dv, isrc2, dv_needed); return; } if (tile <= 0) { @@ -6900,9 +6967,13 @@ static void ggml_flash_attn_ext_backward( struct ggml_tensor * dq = NULL; struct ggml_tensor ** dk = (struct ggml_tensor **) alloca(sizeof(struct ggml_tensor *) * nt); struct ggml_tensor ** dv = (struct ggml_tensor **) alloca(sizeof(struct ggml_tensor *) * nt); + const int t0 = (int) (kv0 / tile); // the first tile holding a key row that carries a gradient for (int t = 0; t < nt; ++t) { const int64_t j0 = (int64_t) t * tile; const int64_t n = MIN(tile, M - j0); + if (t < t0 && !dq_needed) { + continue; + } struct ggml_tensor * kt = ggml_fa_back_f32(ctx, ggml_view_3d(ctx, k, k->ne[0], n, Hkv, k->nb[1], k->nb[2], j0*k->nb[1])); struct ggml_tensor * vt = ggml_fa_back_f32(ctx, ggml_view_3d(ctx, v, v->ne[0], n, Hkv, v->nb[1], v->nb[2], j0*v->nb[1])); struct ggml_tensor * p = ggml_exp(ctx, ggml_sub(ctx, ggml_fa_back_scores(ctx, q, k, mask, scale, j0, n), lse)); // [n, N, H] @@ -6913,11 +6984,11 @@ static void ggml_flash_attn_ext_backward( struct ggml_tensor * g = ggml_fa_back_mm(ctx, ggml_cont(ctx, ggml_transpose(ctx, kt)), ds); // [D, N, H] dq = dq ? ggml_add(ctx, dq, g) : g; } - if (dk_needed) { + if (dk_needed && t >= t0) { // dk[d,j] = sum_i q[d,i] ds[j,i] dk[t] = ggml_fa_back_group_sum(ctx, ggml_fa_back_mm(ctx, qT, ggml_cont(ctx, ggml_transpose(ctx, ds))), Hkv); // [D, n, Hkv] } - if (dv_needed) { + if (dv_needed && t >= t0) { // dv[d,j] = sum_i dO[d,i] p[j,i] dv[t] = ggml_fa_back_group_sum(ctx, ggml_fa_back_mm(ctx, dOT, ggml_cont(ctx, ggml_transpose(ctx, p))), Hkv); // [Dv, n, Hkv] } @@ -6926,11 +6997,14 @@ static void ggml_flash_attn_ext_backward( if (dq_needed) { ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape_4d(ctx, ggml_cont(ctx, dq), q->ne[0], q->ne[1], q->ne[2], q->ne[3])); } - if (dk_needed) { - ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape_4d(ctx, ggml_fa_back_join(ctx, dk, nt), k->ne[0], k->ne[1], k->ne[2], k->ne[3])); - } - if (dv_needed) { - ggml_add_or_set(ctx, cgraph, isrc2, ggml_reshape_4d(ctx, ggml_fa_back_join(ctx, dv, nt), v->ne[0], v->ne[1], v->ne[2], v->ne[3])); + // the tiles from t0 joined, then their rows from kv0 ([D, M - kv0, Hkv]) + for (int part = 0; part < 2; ++part) { + if (!(part == 0 ? dk_needed : dv_needed)) { + continue; + } + struct ggml_tensor * j = ggml_fa_back_join(ctx, (part == 0 ? dk : dv) + t0, nt - t0); + struct ggml_tensor * g = ggml_view_3d(ctx, j, j->ne[0], M - kv0, j->ne[2], j->nb[1], j->nb[2], (kv0 - (int64_t) t0*tile)*j->nb[1]); + ggml_fa_back_route_kv(ctx, cgraph, grads_needed, tensor, part, g, part == 0 ? isrc1 : isrc2, true); } } @@ -7537,7 +7611,7 @@ static void ggml_compute_backward( } break; case GGML_OP_FLASH_ATTN_EXT: { ggml_flash_attn_ext_backward(ctx, cgraph, tensor, grad, isrc0, isrc1, isrc2, - src0_needs_grads, src1_needs_grads, src2_needs_grads); + src0_needs_grads, src1_needs_grads, src2_needs_grads, grads_needed); } break; case GGML_OP_NONE: { // noop @@ -7549,8 +7623,10 @@ static void ggml_compute_backward( } GGML_ASSERT(!src0_needs_grads || ggml_are_same_shape(src0, cgraph->grads[isrc0])); - GGML_ASSERT(!src1_needs_grads || ggml_are_same_shape(src1, cgraph->grads[isrc1])); - GGML_ASSERT(!src2_needs_grads || ggml_are_same_shape(src2, cgraph->grads[isrc2])); + // an attention with grad sinks puts its K/V gradient on the sinks, never on K/V themselves + const bool kv_sunk = tensor->op == GGML_OP_FLASH_ATTN_EXT && tensor->src[6] != NULL; + GGML_ASSERT(!src1_needs_grads || kv_sunk || ggml_are_same_shape(src1, cgraph->grads[isrc1])); + GGML_ASSERT(!src2_needs_grads || kv_sunk || ggml_are_same_shape(src2, cgraph->grads[isrc2])); } static size_t ggml_visit_parents_graph(struct ggml_cgraph * cgraph, struct ggml_tensor * node, bool compute) { diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 9cbb8729bd72..3090bfcd2293 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -67,7 +67,7 @@ bool llama_context::opt_fa_back_native() const { ggml_tensor * v = ggml_new_tensor_3d(ctx.get(), kv->type_v(), Dv, 64, Hkv); ggml_tensor * m = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F16, 64, 8); ggml_tensor * d = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, Dv, H, 8); - ggml_tensor * op = ggml_flash_attn_back(ctx.get(), q, k, v, m, d, 1.0f); + ggml_tensor * op = ggml_flash_attn_back(ctx.get(), q, k, v, m, d, 1.0f, 0); for (const auto & backend : backends) { ggml_backend_dev_t dev = ggml_backend_get_device(backend.get()); const auto type = dev ? ggml_backend_dev_type(dev) : GGML_BACKEND_DEVICE_TYPE_CPU; @@ -387,14 +387,9 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); for (const auto & io : res->t_walk) { if (io.tk) { - // the explicit path taps the span itself; the in-place path taps the window's cells - if (ggml_are_same_shape(io.tk, io.ak)) { - ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); - ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); - } else { - ggml_opt_accumulate_grad_rows(opt_ctx, io.tk, io.g0, io.ak); - ggml_opt_accumulate_grad_rows(opt_ctx, io.tv, io.g0, io.av); - } + // the run's cached rows before this chunk, as F32: their gradient into the accumulator + ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); + ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); } } for (const auto & ws : res->t_walk_state) { diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 66f45effa4eb..1c7bfe10acd4 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2751,6 +2751,11 @@ ggml_tensor * llm_graph_context::build_attn_mha( if (cparams.training) { ggml_flash_attn_ext_set_back_tile(cur, cparams.fa_back_native ? GGML_FA_BACK_NATIVE : 0); } + if (walk_grad_sinks[1] != nullptr) { + // the in-place training graph's K/V gradient goes to the run's rows only + ggml_flash_attn_ext_set_grad_sinks(cur, walk_grad_sinks[0], walk_grad_sinks[1], walk_grad_sinks[2], walk_grad_sinks[3]); + walk_grad_sinks = {}; + } if (v_mla) { #if 0 @@ -2992,41 +2997,52 @@ ggml_tensor * llm_graph_context::build_attn( // THE TRAINING GRAPH IS THE SERVING GRAPH. This chunk writes its K/V into the cache // exactly as serving does (SET_ROWS, the cache's own type) and attention reads the // cache in place: no copy of the context, at any window. The walk decoded the run - // into the cache under the same adapter, so the write is idempotent; the backward - // gathers this chunk's gradient from the rows it wrote (SET_ROWS backward). In the - // exact walk the write is TAPPED: the backward adds dL/d(cached K/V) of the run's - // positions before this chunk ([walk_grad_from, n_past)) into the run's host - // accumulator (cparams.walk_acc_k / _v, row 0 at walk_acc_base), and what later chunks - // put on this chunk's own rows reaches it through the surrogate. A quantized or f16 - // cache trains on the values serving reads; its gradient passes straight through the - // cast. + // into the cache under the same adapter, so the write is idempotent. + // Only the run's rows carry a K/V gradient, and they are the attention's grad sinks + // (ggml_flash_attn_ext_set_grad_sinks): the chunk's own K/V and, in the exact walk, the + // run's positions before it ([g0, n_past), read from the cache as F32 and TAPPED: their + // gradient adds into the run's host accumulator, row 0 at walk_acc_base). The backward + // computes dK/dV from g0 on and nothing for the window before it, a constant (her served + // prefix): no window-sized gradient exists. What later chunks put on this chunk's own rows + // reaches it through the surrogate. A quantized or f16 cache trains on the values serving + // reads; its gradient passes straight through to the F32 K/V. const int64_t n_tokens = k_cur->ne[2]; const int64_t n_past = ubatch.pos[0]; const int64_t n_win = n_past + n_tokens; ggml_tensor * rows = build_inp_train_rows(); - // the run's positions before this chunk carry a gradient: [g0, n_past), rows from base const int64_t base = cparams.walk_acc_base; - const int64_t g0 = std::max(cparams.walk_grad_from, base); + const int64_t g0 = std::min(std::max(cparams.walk_grad_from, base), n_past); llm_graph_result::walk_layer io = { il, (uint32_t) n_past, (uint32_t) g0 }; // the window's cells [0, n_win) of one layer's cache with this chunk's rows written - auto write = [&](ggml_tensor * all, ggml_tensor * cur, ggml_tensor * acc, ggml_tensor *& tap, ggml_tensor *& tap_acc) { + auto write = [&](ggml_tensor * all, ggml_tensor * cur) { const int64_t n_embd = cur->ne[0]*cur->ne[1]; GGML_ASSERT(all->ne[3] == 1 && "a training context reads one stream"); GGML_ASSERT(ggml_row_size(cur->type, cur->ne[0]) == cur->nb[1]); ggml_tensor * win = ggml_view_2d(ctx0, all, n_embd, n_win, all->nb[2], 0); ggml_tensor * w = ggml_set_rows(ctx0, win, ggml_view_2d(ctx0, cur, n_embd, n_tokens, cur->nb[2], 0), rows); - if (acc != nullptr && n_past > g0) { - tap = w; - tap_acc = ggml_view_2d(ctx0, acc, n_embd, n_past - g0, acc->nb[1], (g0 - base)*acc->nb[1]); - } return ggml_reshape_3d(ctx0, w, cur->ne[0], cur->ne[1], n_win); }; + // the cached positions [p0, p1) as F32 [D, Hkv, p1 - p0] + auto span = [&](ggml_tensor * all, int64_t p0, int64_t p1) { + return ggml_cast(ctx0, ggml_view_3d(ctx0, all, all->ne[0], all->ne[1], p1 - p0, all->nb[1], all->nb[2], p0*all->nb[2]), + GGML_TYPE_F32); + }; + ggml_tensor * k_all = mctx_cur->get_k(ctx0, il); ggml_tensor * v_all = mctx_cur->get_v(ctx0, il); GGML_ASSERT(v_all->nb[1] <= v_all->nb[2] && "a flash-attention cache holds V untransposed"); - const bool tapped = cparams.walk_exact; - k = write(mctx_cur->get_k(ctx0, il), k_cur, tapped ? cparams.walk_acc_k[il] : nullptr, io.tk, io.ak); - v = write(v_all, v_cur, tapped ? cparams.walk_acc_v[il] : nullptr, io.tv, io.av); - if (tapped) { + GGML_ASSERT(k_cur->type == GGML_TYPE_F32 && v_cur->type == GGML_TYPE_F32); + k = write(k_all, k_cur); + v = write(v_all, v_cur); + if (cparams.walk_exact && n_past > g0) { + io.tk = span(k_all, g0, n_past); + io.tv = span(v_all, g0, n_past); + ggml_format_name(io.tk, "walk_tk-%d", il); + ggml_format_name(io.tv, "walk_tv-%d", il); + io.ak = walk_acc_rows(cparams.walk_acc_k[il], io.tk, g0 - base); + io.av = walk_acc_rows(cparams.walk_acc_v[il], io.tv, g0 - base); + } + walk_grad_sinks = { io.tk, k_cur, io.tv, v_cur }; + if (cparams.walk_exact) { if (cparams.walk_surrogate) { build_walk_surrogate(il, k_cur, v_cur, n_past - base); } diff --git a/src/llama-graph.h b/src/llama-graph.h index c018ffd0e6b5..4c7f60bd8137 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -6,6 +6,7 @@ #include "llama-adapter.h" #include +#include #include #include #include @@ -1414,6 +1415,9 @@ struct llm_graph_context { void build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const; // rows [row, row + n) of a layer's walk accumulator, shaped like `like`'s K / V span ggml_tensor * walk_acc_rows(ggml_tensor * acc, const ggml_tensor * like, int64_t row) const; + // the in-place training attention's K/V grad sinks (tap k, chunk k, tap v, chunk v), set by + // build_attn before build_attn_mha and taken by the FLASH_ATTN_EXT it builds + mutable std::array walk_grad_sinks = {}; // + : what later chunks of the run put on this chunk's own K/V // (accumulator rows from `row`), carried into it through the loss void build_walk_surrogate(int il, ggml_tensor * k_cur, ggml_tensor * v_cur, int64_t row) const; diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 87341ca5cc92..93458c9bec8b 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7534,13 +7534,15 @@ struct test_flash_attn_ext_grad : public test_case { struct test_flash_attn_back : public test_case { const int64_t hs, nh, nhkv, kv, nb; const ggml_type type_kv; + const int64_t kv0; // the first key row with a dK/dV (a frozen prefix before it) std::string vars() override { - return VARS_TO_STR6(hs, nh, nhkv, kv, nb, type_kv); + return VARS_TO_STR7(hs, nh, nhkv, kv, nb, type_kv, kv0); } - test_flash_attn_back(int64_t hs = 128, int64_t nh = 8, int64_t nhkv = 2, int64_t kv = 512, int64_t nb = 64, ggml_type type_kv = GGML_TYPE_F16) - : hs(hs), nh(nh), nhkv(nhkv), kv(kv), nb(nb), type_kv(type_kv) {} + test_flash_attn_back(int64_t hs = 128, int64_t nh = 8, int64_t nhkv = 2, int64_t kv = 512, int64_t nb = 64, ggml_type type_kv = GGML_TYPE_F16, + int64_t kv0 = 0) + : hs(hs), nh(nh), nhkv(nhkv), kv(kv), nb(nb), type_kv(type_kv), kv0(kv0) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, nb, nh); @@ -7549,7 +7551,7 @@ struct test_flash_attn_back : public test_case { ggml_tensor * m = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kv, nb); ggml_tensor * d = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, nh, nb); ggml_set_name(q, "q"); ggml_set_name(k, "k"); ggml_set_name(v, "v"); ggml_set_name(m, "m"); ggml_set_name(d, "d"); - ggml_tensor * out = ggml_flash_attn_back(ctx, q, k, v, m, d, 1.0f/sqrtf((float) hs)); + ggml_tensor * out = ggml_flash_attn_back(ctx, q, k, v, m, d, 1.0f/sqrtf((float) hs), kv0); ggml_set_name(out, "out"); return out; } @@ -7572,7 +7574,7 @@ struct test_flash_attn_back : public test_case { // the forward's two uint64_t op_flops(ggml_tensor * t) override { GGML_UNUSED(t); - return 2ull*nb*kv*nh*(3*hs + 2*hs); + return 2ull*nb*nh*(3*kv*hs + 2*(kv - kv0)*hs); // pass 1 over every key, pass 2 from kv0 } }; @@ -10544,6 +10546,10 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 300, 37, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_back(512, 4, 2, 128, 9, GGML_TYPE_F16)); // the tiled path's tiles exceed 32 KiB: the per-row path + // dK/dV from a first key row: a frozen prefix, unaligned to any tile, and on the per-row path + test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 300, 37, GGML_TYPE_F16, 300 - 37)); + test_cases.emplace_back(new test_flash_attn_back(128, 8, 2, 512, 64, GGML_TYPE_Q8_0, 203)); + test_cases.emplace_back(new test_flash_attn_back(512, 4, 2, 128, 9, GGML_TYPE_F16, 100)); test_cases.emplace_back(new test_cross_entropy_loss_index()); test_cases.emplace_back(new test_cross_entropy_loss_index({30000, 2, 1, 1})); // a vocabulary's width; few rows keep the finite difference above f32 rounding test_cases.emplace_back(new test_cross_entropy_loss_index({10, 5, 4, 3}, true)); @@ -10740,6 +10746,7 @@ static std::vector> make_test_cases_perf() { // heads, 256 rows over a 2048 window, f16 cache): small enough to time on every kernel edit; a // backward costs about 2.5 forwards of matmul work test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 2048, 256, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 2048, 256, GGML_TYPE_F16, 2048 - 256)); // the reply's keys only test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {4, 1}, 2048, 256, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); // OUT_PROD at the shapes a LoRA backward emits (dx = out_prod(W, dy^T), trans_b): a 1.5B From 8bbdf012e7c6c7c418db8704d945431d3e68df4f Mon Sep 17 00:00:00 2001 From: joelteply Date: Thu, 15 Oct 2026 03:46:19 -0500 Subject: [PATCH 30/43] ggml-cpu: OUT_PROD with a quantized src0 dequantizes each src0 tile once per block of dst rows, on simd_gemm The LoRA backward's dx = out_prod(W, dy^T) through a frozen quantized base weight is the CPU training step's largest term. The kernel dequantized every src0 row again for EVERY dst row (n times over) and accumulated with one scalar-weighted mad per row. Now each task owns a block of 32 dst rows of one (i2,i3) plane, dequantizes src0 in tiles of 32 rows once per block, gathers the matching src1 values, and accumulates dst[BN x ne0] += S1[BN x KT] * W[KT x ne0] on simd_gemm (the forward attention's GEMM). Non-contiguous dst rows keep the mad loop. Scratch per thread: KT*ne0 + BN*KT floats (ops.h), sized by the work planner. Measured (IntelMac i7-8850H, 12 threads, test-backend-ops perf, q4_K, 2026-10-15): the 1.5B LoRA backward shapes at 1024 tokens, trans_b: m=1536 k=1536 410.9 -> 34.7 ms/call (11.8x); m=1536 k=8960 2,463.2 -> 194.3 ms (12.7x); m=8960 k=1536 494.4 ms (its baseline run was cut). Exactness: old and new kernels on identical seeded q4_K inputs give BIT-IDENTICAL results (NMSE 0, max abs 0) at 1536x1024x1536, 1536x1024x8960, 8960x1024x1536 (transposed b) and 1536x300x1536 (contiguous b): the same FMAs in the same k order, from a dequantized tile instead of a fresh row. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE --- ggml/src/ggml-cpu/ggml-cpu.c | 6 ++- ggml/src/ggml-cpu/ops.cpp | 74 ++++++++++++++++++------------------ ggml/src/ggml-cpu/ops.h | 7 ++++ 3 files changed, 49 insertions(+), 38 deletions(-) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 0ba316726bb0..f7c622be3466 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2884,8 +2884,10 @@ struct ggml_cplan ggml_graph_plan( } break; case GGML_OP_OUT_PROD: { - if (ggml_is_quantized(node->src[0]->type) || - node->src[0]->type == GGML_TYPE_F16) { + if (ggml_is_quantized(node->src[0]->type)) { + // the blocked kernel's tile and gathered values (ops.h) + cur = ggml_type_size(GGML_TYPE_F32) * ggml_out_prod_cpu_wsize(node->src[0]->ne[0]) * n_tasks; + } else if (node->src[0]->type == GGML_TYPE_F16) { cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; } } break; diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index dedb124f3147..7f47eb2321d8 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4392,55 +4392,57 @@ static void ggml_compute_forward_out_prod_q_f32( // nb01 >= nb00 - src0 is not transposed // compute by src0 rows + // + // dst[i0,i1] = sum_i01 src0[i0,i01] * src1[i1,i01]: a GEMM dst[ne1 x ne0] = S1[ne1 x ne01] * + // W[ne01 x ne0]. Each task owns a block of BN dst rows of one (i2,i3) plane, dequantizes src0 in + // tiles of KT rows ONCE per block (the row loop dequantized every src0 row again for every dst + // row), gathers the matching src1 values, and accumulates on simd_gemm. + const int64_t BN = GGML_OUT_PROD_BN; + const int64_t KT = GGML_OUT_PROD_KT; if (ith == 0) { ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); } ggml_barrier(params->threadpool); - // parallelize by last three dimensions - - // total rows in dst - const int64_t nr = ne1*ne2*ne3; - - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; + float * wt = (float *) params->wdata + ith*ggml_out_prod_cpu_wsize(ne0); // [KT][ne0] + float * at = wt + KT*ne0; // [BN][KT] - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); - - // dst[:,:,:,:] = 0 - // for i2,i3: - // for i1: - // for i01: - // for i0: - // dst[i0,i1,i2,i3] += src0[i0,i01,i2,i3] * src1[i1,i01,i2,i3] - - float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith; + const int64_t nblk = (ne1 + BN - 1)/BN; + const int64_t ntasks = nblk*ne2*ne3; + const bool dense = nb1 == (size_t) ne0*sizeof(float); // simd_gemm's C rows are contiguous - for (int64_t ir = ir0; ir < ir1; ++ir) { - // dst indices - const int64_t i3 = ir/(ne2*ne1); - const int64_t i2 = (ir - i3*ne2*ne1)/ne1; - const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); + for (int64_t t = ith; t < ntasks; t += nth) { + const int64_t i3 = t/(ne2*nblk); + const int64_t i2 = (t - i3*ne2*nblk)/nblk; + const int64_t b0 = (t - i3*ne2*nblk - i2*nblk)*BN; + const int64_t nb_ = std::min(BN, ne1 - b0); const int64_t i02 = i2 / (ne12 / ne02); const int64_t i03 = i3 / (ne13 / ne03); - //const int64_t i10 = i1; - const int64_t i12 = i2; - const int64_t i13 = i3; - - for (int64_t i01 = 0; i01 < ne01; ++i01) { - const int64_t i11 = i01; + float * d0 = (float *) ((char *) dst->data + (b0*nb1 + i2*nb2 + i3*nb3)); - float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); - float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); - float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); - - dequantize_row_q(s0, wdata, ne0); - ggml_vec_mad_f32(ne0, d, wdata, *s1); + for (int64_t k0 = 0; k0 < ne01; k0 += KT) { + const int64_t kt = std::min(KT, ne01 - k0); + for (int64_t r = 0; r < kt; ++r) { + dequantize_row_q((const char *) src0->data + ((k0 + r)*nb01 + i02*nb02 + i03*nb03), wt + r*ne0, ne0); + } + for (int64_t b = 0; b < nb_; ++b) { + for (int64_t r = 0; r < kt; ++r) { + at[b*kt + r] = *(const float *) ((const char *) src1->data + ((b0 + b)*nb10 + (k0 + r)*nb11 + i2*nb12 + i3*nb13)); + } + } + if (dense) { + simd_gemm(d0, at, wt, (int) nb_, (int) kt, (int) ne0); + } else { + for (int64_t b = 0; b < nb_; ++b) { + float * d = (float *) ((char *) d0 + b*nb1); + for (int64_t r = 0; r < kt; ++r) { + ggml_vec_mad_f32(ne0, d, wt + r*ne0, at[b*kt + r]); + } + } + } } } } diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 4506167b6d3e..9c1e15dc2765 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -89,6 +89,13 @@ void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_flash_attn_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +// OUT_PROD with a quantized src0: BN dst rows per task share one dequantization of each KT-row tile +// of src0. Its per-thread scratch, in floats, for a row of ne0; the work planner reserves this. +#define GGML_OUT_PROD_BN 32 +#define GGML_OUT_PROD_KT 32 +static inline int64_t ggml_out_prod_cpu_wsize(int64_t ne0) { + return GGML_OUT_PROD_KT*ne0 + GGML_OUT_PROD_BN*GGML_OUT_PROD_KT; +} // FLASH_ATTN_BACK's blocks: BR query rows (pass 1) and BC KV positions (pass 2) share one conversion of // each K and V row. Its per-thread scratch, in floats; the work-size planner reserves this per task. #define GGML_FA_BACK_BR 8 From f0045c036220eb60fbde9845bc9a40b3edb0ad93 Mon Sep 17 00:00:00 2001 From: joelteply Date: Thu, 15 Oct 2026 03:59:14 -0500 Subject: [PATCH 31/43] ggml-cpu: OUT_PROD with an f16 src0 takes the same blocked kernel; the per-row f16 copy is deleted ggml_compute_forward_out_prod_f16_f32 was the quantized kernel's old per-row loop with ggml_fp16_to_fp32_row in place of the dequantizer. F16's to_float trait is that function, so f16 now routes through the blocked kernel (renamed out_prod_converted_f32: it converts any src0 row type through its to_float trait), and the duplicate is gone. Measured (IntelMac i7-8850H, 12 threads, same seeded inputs, best of 2-3, 2026-10-15): f16 1536x1024x1536 transposed b 208.9 -> 23.3 ms (9.0x); 1536x1024x8960 1,276.6 -> 145.5 ms (8.8x); 1536x300x1536 contiguous b 64.8 -> 9.8 ms (6.6x). BIT-IDENTICAL to the deleted kernel on all three (cmp), and the q4_K path's output is unchanged by the rename. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE --- ggml/src/ggml-cpu/ggml-cpu.c | 4 +- ggml/src/ggml-cpu/ops.cpp | 71 ++---------------------------------- 2 files changed, 5 insertions(+), 70 deletions(-) diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index f7c622be3466..b14248897227 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2884,11 +2884,9 @@ struct ggml_cplan ggml_graph_plan( } break; case GGML_OP_OUT_PROD: { - if (ggml_is_quantized(node->src[0]->type)) { + if (ggml_is_quantized(node->src[0]->type) || node->src[0]->type == GGML_TYPE_F16) { // the blocked kernel's tile and gathered values (ops.h) cur = ggml_type_size(GGML_TYPE_F32) * ggml_out_prod_cpu_wsize(node->src[0]->ne[0]) * n_tasks; - } else if (node->src[0]->type == GGML_TYPE_F16) { - cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; } } break; case GGML_OP_SET_ROWS: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 7f47eb2321d8..85ecb257e137 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4356,7 +4356,7 @@ static void ggml_compute_forward_out_prod_f32( } } -static void ggml_compute_forward_out_prod_q_f32( +static void ggml_compute_forward_out_prod_converted_f32( const ggml_compute_params * params, ggml_tensor * dst) { @@ -4447,70 +4447,6 @@ static void ggml_compute_forward_out_prod_q_f32( } } -static void ggml_compute_forward_out_prod_f16_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_TENSOR_BINARY_OP_LOCALS; - - const int ith = params->ith; - const int nth = params->nth; - - GGML_ASSERT(src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - GGML_ASSERT(ne12 % ne02 == 0); - GGML_ASSERT(ne13 % ne03 == 0); - GGML_ASSERT(ne2 == ne12); - GGML_ASSERT(ne3 == ne13); - - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb0 == sizeof(float)); - - GGML_ASSERT(ne0 == ne00); - GGML_ASSERT(ne1 == ne10); - - - - if (ith == 0) { - ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); - } - ggml_barrier(params->threadpool); - - const int64_t nr = ne1*ne2*ne3; - const int64_t dr = (nr + nth - 1)/nth; - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); - - float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith; - - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i3 = ir/(ne2*ne1); - const int64_t i2 = (ir - i3*ne2*ne1)/ne1; - const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); - - const int64_t i02 = i2 / (ne12 / ne02); - const int64_t i03 = i3 / (ne13 / ne03); - - const int64_t i12 = i2; - const int64_t i13 = i3; - - float * d = (float *) ((char *) dst->data + (i1*nb1 + i2*nb2 + i3*nb3)); - - for (int64_t i01 = 0; i01 < ne01; ++i01) { - const int64_t i11 = i01; - ggml_fp16_t * s0 = (ggml_fp16_t *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03)); - float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); - ggml_fp16_to_fp32_row(s0, wdata, ne0); - ggml_vec_mad_f32(ne0, d, wdata, *s1); - } - } -} - void ggml_compute_forward_out_prod( const ggml_compute_params * params, ggml_tensor * dst) { @@ -4544,11 +4480,12 @@ void ggml_compute_forward_out_prod( case GGML_TYPE_IQ3_S: case GGML_TYPE_IQ2_S: { - ggml_compute_forward_out_prod_q_f32(params, dst); + ggml_compute_forward_out_prod_converted_f32(params, dst); } break; case GGML_TYPE_F16: { - ggml_compute_forward_out_prod_f16_f32(params, dst); + // f16 converts through its to_float trait exactly as a quantized row dequantizes + ggml_compute_forward_out_prod_converted_f32(params, dst); } break; case GGML_TYPE_F32: { From a0705b09598dcb1be4568af96e85709e84867707 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 04:08:36 -0500 Subject: [PATCH 32/43] train: the training graph no longer grows with the window Three window-sized holds in a training graph, found by measuring the 4B at 8k and 32k: - Metal reserved a window-sized F16 copy of K/V inside every FLASH_ATTN_EXT's output, at any K type, "to avoid graph reallocations". Only a quantized K is ever dequantized (and only at 32+ rows), so a K the kernel never dequantizes now reserves nothing; a quantized K keeps the reservation, so decode and prefill graphs still match. - A training graph keeps each attention's output for the backward, so a quantized cache's F16 copy lived per layer: training now dequantizes through a cast node of its own, and the native backward reads the cast's quantized source, so the copy is not kept for the backward. - Recompute never rebuilt a node that read a buffer the forward writes in place (the cache), so every layer's cast stayed alive from the forward to that layer's backward. A clone runs after the forward, so it reads such a buffer as the forward left it: the same value unless a write comes AFTER the original read. Only reads followed by a write are kept now. GGML_OPT_TOP also lists the graph's input leaves, each tensor's sources, and the graph's total. Measured (M5 Metal, qwen3.5-4b hybrid, chunk 512 at the window's end, recompute; graph peak MiB; loss identical before and after, 2.904847 / 2.904831 / 2.905061 at 32k): window explicit F32 flash f16 flash q8_0 (her lane's cache) 8192 1,546 -> 1,220 1,085 -> 557 1,096 -> 550 32768 4,608 -> 3,540 2,020 -> 581 1,888 -> 615 test-walk-exact on qwen2.5-0.5b and qwen3.5-4b: every exact row cosine 1.0000, recompute vs without 1.0000, caches byte-identical; test-backend-ops MTL0 FLASH_ATTN_EXT 4809/4809, grad 19818/19818. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-metal/ggml-metal-ops.cpp | 36 +++++++++++++--------- ggml/src/ggml-opt.cpp | 42 +++++++++++++++++++++----- ggml/src/ggml.c | 11 +++++++ src/llama-graph.cpp | 13 ++++++++ tests/test-walk-exact.cpp | 11 ++++--- 5 files changed, 87 insertions(+), 26 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 27c07e76939b..a704f4e40e89 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2961,16 +2961,8 @@ bool ggml_metal_op_flash_attn_ext_use_vec(const ggml_tensor * op) { // ref: https://github.com/ggml-org/llama.cpp/pull/27390 // dequantize the quantized KV cache to F16 before running the F16 flash attention kernels -static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { - assert(op->op == GGML_OP_FLASH_ATTN_EXT); - - // depending on compute/bandwidth ratio, dequant to f16 kv is not always beneficial - // ref: https://github.com/ggml-org/llama.cpp/pull/27390#issuecomment-5355152767 - // TODO: tune per device - if (op->src[0]->ne[1] < 32) { - return false; - } - +// whether the kernel may dequantize this K to F16 at some batch size (a quantized K) +static bool ggml_metal_op_flash_attn_ext_kv_dequantizable(const ggml_tensor * op) { switch (op->src[1]->type) { case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: @@ -2983,6 +2975,19 @@ static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { } } +static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + // depending on compute/bandwidth ratio, dequant to f16 kv is not always beneficial + // ref: https://github.com/ggml-org/llama.cpp/pull/27390#issuecomment-5355152767 + // TODO: tune per device + if (op->src[0]->ne[1] < 32) { + return false; + } + + return ggml_metal_op_flash_attn_ext_kv_dequantizable(op); +} + // in some models (e.g. MLA-based), V is a view of K (the first ne20 elements of each K row); // the dequantized V is then a view of the dequantized K and does not need its own dequant or scratch // - ref: https://github.com/ggml-org/llama.cpp/pull/13435 @@ -3133,10 +3138,13 @@ size_t ggml_metal_op_flash_attn_ext_extra_tmp(const ggml_tensor * op) { size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); - // note: always reserve the temp buffer to avoid graph reallocations - //if (!ggml_metal_op_flash_attn_ext_use_kv_f16(op)) { - // return 0; - //} + // reserved at every batch size for a quantized K (its use depends on the batch, and a graph + // must not reallocate between decode and prefill); a K the kernel never dequantizes needs none: + // the reservation lives as long as the attention's output, and a training graph keeps one per + // layer for the backward (the M5, qwen3.5-4b, window 32768: 128 MiB a layer, never read) + if (!ggml_metal_op_flash_attn_ext_kv_dequantizable(op)) { + return 0; + } GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index a6b528fb3228..0beea89f3bad 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -367,7 +367,8 @@ static ggml_tensor * ggml_opt_recompute_node( static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd, const char * prefix) { std::unordered_set forward; std::unordered_set keep; - std::unordered_set mutated; // roots of the buffers the forward pass writes in place + // roots of the buffers the forward pass writes in place, each with its last write's position + std::unordered_map mutated; const size_t plen = strlen(prefix); for (int i = 0; i < n_fwd; ++i) { ggml_tensor * node = g->nodes[i]; @@ -380,18 +381,27 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd, const bool view_op = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE; if ((node->op == GGML_OP_CPY && !ggml_opt_is_cast(node)) || node->op == GGML_OP_SET_ROWS || (node->view_src != nullptr && !view_op)) { - mutated.insert(ggml_opt_view_root(node)); + mutated[ggml_opt_view_root(node)] = i; } } + // A clone runs after the whole forward pass, so it sees a mutated buffer as the forward left + // it: the same value the original read unless the buffer is written AFTER the original. A read + // of a buffer whose writes all precede it (a layer's cache read after its own chunk's write) is + // recomputed like any other node; keeping it held a window-sized cast of the cache per layer + // from the forward to the backward (the M5, qwen3.5-4b, q8_0 cache, window 32768: 16 x 64 MiB). + auto written_after = [&](ggml_tensor * t, int i) { + const auto it = mutated.find(ggml_opt_view_root(t)); + return it != mutated.end() && it->second > i; + }; for (int i = 0; i < n_fwd; ++i) { ggml_tensor * node = g->nodes[i]; bool reads_mutated = false; for (int k = 0; k < GGML_MAX_SRC && !reads_mutated; ++k) { - reads_mutated = node->src[k] != nullptr && mutated.count(ggml_opt_view_root(node->src[k])); + reads_mutated = node->src[k] != nullptr && written_after(node->src[k], i); } if (strncmp(ggml_get_name(node), prefix, plen) == 0 || (node->flags & GGML_TENSOR_FLAG_LOSS) || (node->flags & GGML_TENSOR_FLAG_OUTPUT) || reads_mutated || - (node->view_src != nullptr && mutated.count(ggml_opt_view_root(node)))) { + (node->view_src != nullptr && written_after(node, i))) { keep.insert(node); } } @@ -1351,7 +1361,8 @@ bool ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { } opt_ctx->allocated_graph = graph; - // GGML_OPT_TOP=N: the graph's N largest tensors (views excluded), where its memory goes + // GGML_OPT_TOP=N: the graph's N largest tensors, nodes and the leaves it allocates (inputs, not + // views or weights), where its memory goes if (const char * top = getenv("GGML_OPT_TOP")) { ggml_cgraph * g = opt_ctx->allocated_graph_copy; std::vector> rows; @@ -1360,12 +1371,29 @@ bool ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { rows.emplace_back(ggml_nbytes(g->nodes[i]), g->nodes[i]); } } + for (int i = 0; i < g->n_leafs; ++i) { + const ggml_tensor * l = g->leafs[i]; + if (l->view_src == nullptr && (l->flags & GGML_TENSOR_FLAG_INPUT)) { + rows.emplace_back(ggml_nbytes(l), l); + } + } std::sort(rows.begin(), rows.end(), [](const auto & x, const auto & y) { return x.first > y.first; }); + size_t sum = 0; + for (const auto & r : rows) { + sum += r.first; + } + GGML_LOG_INFO("%s: graph of %d nodes, %d leafs: its tensors sum to %.1f MiB, the peak so far %.1f MiB\n", __func__, + g->n_nodes, g->n_leafs, sum/1048576.0, opt_ctx->peak_graph_bytes/1048576.0); const size_t n = std::min(rows.size(), (size_t) atoi(top)); for (size_t i = 0; i < n; ++i) { const ggml_tensor * t = rows[i].second; - GGML_LOG_INFO("%s: top %2zu: %8.1f MiB %-12s %-6s [%lld,%lld,%lld,%lld] %s\n", __func__, i, rows[i].first/1048576.0, - ggml_op_desc(t), ggml_type_name(t->type), (long long) t->ne[0], (long long) t->ne[1], (long long) t->ne[2], (long long) t->ne[3], t->name); + GGML_LOG_INFO("%s: top %2zu: %8.1f MiB %-12s %-6s [%lld,%lld,%lld,%lld] %s <- %s %s [%lld,%lld] | %s %s [%lld,%lld]\n", __func__, i, + rows[i].first/1048576.0, ggml_op_desc(t), ggml_type_name(t->type), (long long) t->ne[0], (long long) t->ne[1], + (long long) t->ne[2], (long long) t->ne[3], t->name, + t->src[0] ? ggml_op_desc(t->src[0]) : "-", t->src[0] ? t->src[0]->name : "", t->src[0] ? (long long) t->src[0]->ne[0] : 0LL, + t->src[0] ? (long long) t->src[0]->ne[1] : 0LL, + t->src[1] ? ggml_op_desc(t->src[1]) : "-", t->src[1] ? t->src[1]->name : "", t->src[1] ? (long long) t->src[1]->ne[0] : 0LL, + t->src[1] ? (long long) t->src[1]->ne[1] : 0LL); } } diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 0f990953c35a..167f954de238 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -6884,6 +6884,17 @@ static void ggml_flash_attn_ext_backward( struct ggml_tensor * v = tensor->src[2]; // [Dv, M, Hkv, 1] struct ggml_tensor * mask = tensor->src[3]; // [M', N', 1|H, 1] F16, or NULL + // the native backward reads a quantized K/V itself: one the forward read through an F16 cast of + // a quantized cache is read at the cast's source, so the window-sized cast is not kept alive + // for the backward + const bool native = ggml_get_op_params_i32(tensor, 4) == GGML_FA_BACK_NATIVE; + if (native && k->op == GGML_OP_CPY && ggml_is_quantized(k->src[0]->type) && ggml_are_same_shape(k, k->src[0])) { + k = k->src[0]; + } + if (native && v->op == GGML_OP_CPY && ggml_is_quantized(v->src[0]->type) && ggml_are_same_shape(v, v->src[0])) { + v = v->src[0]; + } + float scale, max_bias, softcap; memcpy(&scale, (const float *) tensor->op_params + 0, sizeof(float)); memcpy(&max_bias, (const float *) tensor->op_params + 1, sizeof(float)); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 1c7bfe10acd4..d1571b907ec7 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2742,6 +2742,19 @@ ggml_tensor * llm_graph_context::build_attn_mha( v = ggml_cast(ctx0, v, GGML_TYPE_F16); } + // A training graph reads a quantized cache as F16 through a node of its own, as the kernel + // would dequantize it anyway: inside the attention the F16 copy lives as long as the + // attention's output, which a training graph keeps for the backward, one window-sized copy + // per layer; the backward reads the cast's quantized source (ggml_flash_attn_ext_backward), + // so this copy dies with the forward (the M5, qwen3.5-4b, q8_0, window 32768: 1,888 MiB + // against 853 at an f16 cache) + if (cparams.training && cparams.walk_active && ggml_is_quantized(k->type)) { + k = ggml_cast(ctx0, k, GGML_TYPE_F16); + } + if (cparams.training && cparams.walk_active && ggml_is_quantized(v->type)) { + v = ggml_cast(ctx0, v, GGML_TYPE_F16); + } + cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias, hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f); res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il}); diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index b75d164803ad..947b30c7a0e3 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -553,7 +553,7 @@ int main(int argc, char ** argv) { const uint32_t chunk = getenv("WALK_MEASURE_CHUNK") ? (uint32_t) atoi(getenv("WALK_MEASURE_CHUNK")) : 256; g_recompute = getenv("WALK_MEASURE_RECOMPUTE") != nullptr; // per-layer recompute, as the served runs train g_walk_horizon = getenv("WALK_MEASURE_HORIZON") ? (uint32_t) atoi(getenv("WALK_MEASURE_HORIZON")) : 0; // 0 = the whole window - printf(" window chunk explicit MiB flash MiB ratio\n"); + printf(" window chunk explicit MiB flash MiB flash q8_0 MiB ratio\n"); for (const char * p = env; *p; ) { char * end = nullptr; const uint32_t w = (uint32_t) strtoul(p, &end, 10); @@ -562,15 +562,16 @@ int main(int argc, char ** argv) { g_window = w; std::vector reply(w + 1, 0); for (uint32_t i = w - 511; i <= w; ++i) reply[i] = 1; - double mib[2] = { 0, 0 }; - const ggml_type types[2] = { GGML_TYPE_F32, GGML_TYPE_F16 }; - for (int a = 0; a < 2; ++a) { + double mib[3] = { 0, 0, 0 }; + // explicit attention at an F32 cache; flash at an f16 cache; flash at q8_0 (her lane's cache) + const ggml_type types[3] = { GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0 }; + for (int a = 0; a < 3; ++a) { g_cache_type = types[a]; step_delta(params, model, init, long_tokens, reply, chunk, true); mib[a] = g_last_graph_bytes / 1048576.0; } g_cache_type = GGML_TYPE_F32; - printf(" %6u %5u %12.1f %9.1f %5.2fx\n", w, chunk, mib[0], mib[1], mib[1] > 0 ? mib[0] / mib[1] : 0.0); + printf(" %6u %5u %12.1f %9.1f %14.1f %5.2fx\n", w, chunk, mib[0], mib[1], mib[2], mib[1] > 0 ? mib[0] / mib[1] : 0.0); fflush(stdout); } std::remove(init.c_str()); From 5b8500ee21fa92850e18e211f7e40fd7672893e2 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 04:43:33 -0500 Subject: [PATCH 33/43] ggml-opt: recompute never clones an in-place write into a buffer the forward did not compute The recompute rule now rebuilds a read whose buffer is written only before it. The LAST in-place write (a view-result op, not a view) into a persistent buffer (a cache, a state) was therefore no longer kept, and its clone, whose view source resolves to the live buffer, would apply the write a second time in the backward. Such a write is a side effect, like SET_ROWS and a non-cast CPY: returned unchanged (Cormac's review of continuum#5057). Measured (M5 Metal, 2026-10-15): test-walk-exact on qwen2.5-0.5b (8 exact rows) and qwen3.5-4b (11) all cosine 1.0000, recompute vs without 1.0000; the 4B graph at window 32768 unchanged at 3,540.3 / 581.4 / 615.0 MiB (explicit / f16 / q8_0). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-opt.cpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 0beea89f3bad..11519f0fa2f8 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -329,6 +329,13 @@ static ggml_tensor * ggml_opt_recompute_node( if (node->flags & (GGML_TENSOR_FLAG_PARAM | GGML_TENSOR_FLAG_INPUT)) { return node; } + // an in-place write into a buffer the forward did not compute (a cache, a state) is a side + // effect: a clone would apply it a second time to the live buffer (Cormac's review of #5057) + const bool view_op = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE; + if (node->view_src != nullptr && !view_op && !forward.count(ggml_opt_view_root(node))) { + return node; + } switch (node->op) { case GGML_OP_NONE: case GGML_OP_SET_ROWS: From c01026943128a64e37e7404bf151db710d1771ca Mon Sep 17 00:00:00 2001 From: joelteply Date: Thu, 15 Oct 2026 05:09:20 -0500 Subject: [PATCH 34/43] tests: the CPU grad suite runs further: EXPM1 grad cases in a measurable range, no params on ops with no backward (partial, card 397cefc6) `test-backend-ops grad -b CPU` (full suite) never finished: it aborted at the first case that marks params on an op ggml_compute_backward has no case for, so no summary was ever printed and every later grad case went unchecked on the CPU. - test_unary's EXPM1 grad cases use [-5, 5]: over [-150, 150] exp overflows f32 (both sides inf, the "SUM nonfinite" failures), and over [-10, 10] the summed loss is too coarse in f32 for the finite difference to see the small gradients (MAA ~1). The EXPM1 backward itself is correct (direct check: grad = exp(x) exactly at -2, 0, 0.5, 3). - test_glu_split marks params only for SWIGLU (the only glu with a backward); test_swiglu_oai and test_pool1d mark none (no backward). Still aborts later, at IM2COL_3D: the remaining cases want one structural fix (ggml says which ops have a backward and the grad harness skips the rest), not one guard per test. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE --- tests/test-backend-ops.cpp | 27 ++++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 93458c9bec8b..354f4f334ae1 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -1992,6 +1992,7 @@ struct test_example : public test_case { // GGML_OP_UNARY +static float EXPM1_GRAD_RANGE = 5.0f; struct test_unary : public test_case { const ggml_unary_op op; const ggml_type type; @@ -2052,6 +2053,13 @@ struct test_unary : public test_case { min = -10.f; max = 10.f; } + // EXPM1's gradient is checked: exp past ~88.7 overflows f32, and over a wide range the + // summed loss is too coarse in f32 for the finite difference to see the small gradients + // (exp(-10) against a sum of exp(10)s reads 0). A range where both sides are measurable. + if (op == GGML_UNARY_OP_EXPM1) { + min = -EXPM1_GRAD_RANGE; + max = EXPM1_GRAD_RANGE; + } for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { // test extended range of values to check for NaNs in GELU @@ -2142,30 +2150,33 @@ struct test_glu_split : public test_case { : op(op), type(type), ne_a(ne_a), v(v) {} ggml_tensor * build_graph(ggml_context * ctx) override { + // only SWIGLU has a backward (ggml_compute_backward); marking another glu's inputs as + // params aborts the whole grad suite there and leaves every later case unchecked + const bool grad_supported = op == GGML_GLU_OP_SWIGLU; ggml_tensor * a; ggml_tensor * b; if (v & 1) { auto ne = ne_a; ne[0] *= 3; a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); + if (grad_supported) { ggml_set_param(a); } ggml_set_name(a, "a"); a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); ggml_set_name(a, "view_of_a"); b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(b); + if (grad_supported) { ggml_set_param(b); } ggml_set_name(b, "b"); b = ggml_view_4d(ctx, b, ne_a[0], ne_a[1], ne_a[2], ne_a[3], b->nb[1], b->nb[2], b->nb[3], 0); ggml_set_name(a, "view_of_b"); } else { a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(a); + if (grad_supported) { ggml_set_param(a); } ggml_set_name(a, "a"); b = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(b); + if (grad_supported) { ggml_set_param(b); } ggml_set_name(b, "b"); } @@ -2201,31 +2212,29 @@ struct test_swiglu_oai : public test_case { float limit = 7.0f) : type(type), ne_a(ne_a), v(v), alpha(alpha), limit(limit) {} + // no params: SWIGLU_OAI has no backward (ggml_compute_backward), and a param would abort the + // whole grad suite here, leaving every later case unchecked ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * a; ggml_tensor * b; if (v & 1) { auto ne = ne_a; ne[0] *= 3; a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(a); ggml_set_name(a, "a"); a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); ggml_set_name(a, "view_of_a"); b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_param(b); ggml_set_name(b, "b"); b = ggml_view_4d(ctx, b, ne_a[0], ne_a[1], ne_a[2], ne_a[3], b->nb[1], b->nb[2], b->nb[3], 0); ggml_set_name(a, "view_of_b"); } else { a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(a); ggml_set_name(a, "a"); b = ggml_new_tensor(ctx, type, 4, ne_a.data()); - ggml_set_param(b); ggml_set_name(b, "b"); } @@ -5798,7 +5807,7 @@ struct test_pool1d : public test_case { ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * input = ggml_new_tensor(ctx, type_input, 4, ne_input.data()); - ggml_set_param(input); + // no param: POOL_1D has no backward (ggml_compute_backward aborts the grad suite on it) ggml_set_name(input, "input"); ggml_tensor * out = ggml_pool_1d(ctx, input, pool_type, k0, s0, p0); From de1d97dab6a49bf2e7e920ac297270f0199b35c6 Mon Sep 17 00:00:00 2001 From: joelteply Date: Thu, 15 Oct 2026 06:33:47 -0500 Subject: [PATCH 35/43] ggml: ggml_node_has_backward is the one answer to "can a backward run through this node"; the grad suite reports instead of aborting (card 397cefc6) `test-backend-ops grad -b CPU` never finished: the first case that marked params on an op ggml_compute_backward has no case for (REGLU, then SWIGLU_OAI, POOL_1D, IM2COL_3D, an in-place SET, ...) aborted the whole run, so no summary ever printed and every later case went unchecked. Upstream (ggml-org master) lacks those backwards too. - ggml_node_has_backward(tensor): the op (and unary or glu op) has a backward case and the node is not an in-place write other than the re-shaping views. ggml_build_backward_expand and ggml_compute_backward both assert it, so the list and the switch cannot disagree silently: a missing case or a missing list entry fails the grad suite loudly. - test_case::eval_grad walks the forward graph and reports a param-reachable node without one as "not supported (no backward for OP)", the way supports_op already reports forward ops. The per-test param guards of c010269 are reverted: this owner replaces them. - test_unary's EXPM1 grad follows the smooth ops' convention (test_sin): eps 0.2, grad_precise, inputs in [-2, 2], tolerance 1e-3. Its grad_eps was 15 (meant for ABS/SGN/RELU over +-150), which no central difference of exp survives, and inputs to 150 overflowed f32 (the "SUM nonfinite" failures). The EXPM1 backward is correct: a direct check gives grad = exp(x) exactly, 1024 elements on 12 threads. Measured (IntelMac CPU, 2026-10-15): the full suite now COMPLETES with a summary (exit 1 instead of 134). 19673/19818 pass; 234 cases report "no backward for ..." instead of aborting. The 145 that fail are tolerance misses (MAA 1.2e-4 to 7e-4): 74 reported as SUM, 68 MUL_MAT and 3 MUL_MAT_HADAMARD. 66 of the MUL_MAT are f16 src0, and they fail the same way at 0135932 (67 there), before the CPU OUT_PROD change, so they are not that kernel's. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE --- ggml/include/ggml.h | 4 ++ ggml/src/ggml.c | 90 ++++++++++++++++++++++++++++++++++++-- tests/test-backend-ops.cpp | 65 ++++++++++++++++++++------- 3 files changed, 138 insertions(+), 21 deletions(-) diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 701410e34e11..12d01193c5f1 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2899,6 +2899,10 @@ extern "C" { struct ggml_cgraph * cgraph, struct ggml_tensor * tensor); + // whether a backward can be built through this node (ggml_build_backward_expand aborts on one + // that cannot: an op without a backward, or an in-place write) + GGML_API bool ggml_node_has_backward(const struct ggml_tensor * tensor); + GGML_API void ggml_build_backward_expand( struct ggml_context * ctx, // context for gradient computation struct ggml_cgraph * cgraph, diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 167f954de238..94bbb9d40fbc 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7019,6 +7019,85 @@ static void ggml_flash_attn_ext_backward( } } +// Whether a backward can be built through this node: its op (and unary or glu op) has a case in +// ggml_compute_backward, and it is not an in-place write other than the re-shaping views. The ONE +// answer a caller asks before building a backward graph, instead of meeting an abort there; both +// ggml_build_backward_expand and ggml_compute_backward assert it, so they cannot disagree silently. +bool ggml_node_has_backward(const struct ggml_tensor * tensor) { + if (tensor->view_src != NULL && tensor->op != GGML_OP_CPY && tensor->op != GGML_OP_VIEW && + tensor->op != GGML_OP_RESHAPE && tensor->op != GGML_OP_PERMUTE && tensor->op != GGML_OP_TRANSPOSE && + tensor->op != GGML_OP_SET_ROWS) { + return false; // in-place operations are not supported + } + switch (tensor->op) { + case GGML_OP_DUP: + case GGML_OP_ADD: + case GGML_OP_ADD1: + case GGML_OP_ACC: + case GGML_OP_SUB: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_LOG: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_SUM: + case GGML_OP_SUM_ROWS: + case GGML_OP_MEAN: + case GGML_OP_REPEAT: + case GGML_OP_REPEAT_BACK: + case GGML_OP_RMS_NORM: + case GGML_OP_MUL_MAT: + case GGML_OP_SCALE: + case GGML_OP_SET: + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_SET_ROWS: + case GGML_OP_GET_ROWS: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_DIAG_MASK_ZERO: + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_IM2COL: + case GGML_OP_POOL_2D: + case GGML_OP_WIN_PART: + case GGML_OP_WIN_UNPART: + case GGML_OP_CONCAT: + case GGML_OP_SSM_CONV: + case GGML_OP_L2_NORM: + case GGML_OP_GATED_DELTA_NET: + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_NONE: + return true; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(tensor)) { + case GGML_UNARY_OP_ABS: + case GGML_UNARY_OP_SGN: + case GGML_UNARY_OP_NEG: + case GGML_UNARY_OP_STEP: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_EXP: + case GGML_UNARY_OP_SIGMOID: + case GGML_UNARY_OP_EXPM1: + case GGML_UNARY_OP_SOFTPLUS: + return true; + default: + return false; + } + case GGML_OP_GLU: + return ggml_get_glu_op(tensor) == GGML_GLU_OP_SWIGLU; + default: + return false; + } +} + static void ggml_compute_backward( struct ggml_context * ctx, struct ggml_cgraph * cgraph, int i, const bool * grads_needed) { struct ggml_tensor * tensor = cgraph->nodes[i]; @@ -7039,6 +7118,9 @@ static void ggml_compute_backward( const bool src1_needs_grads = src1 && isrc1 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc1) && grads_needed[isrc1]; const bool src2_needs_grads = src2 && isrc2 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc2) && grads_needed[isrc2]; + if (!ggml_node_has_backward(tensor)) { + GGML_ABORT("%s: no backward implemented for op %s (%s)", __func__, ggml_op_name(tensor->op), ggml_op_desc(tensor)); + } switch (tensor->op) { case GGML_OP_DUP: { if (src0_needs_grads) { @@ -7827,10 +7909,10 @@ void ggml_build_backward_expand( continue; } - // inplace operations are currently not supported - GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || - node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE || - node->op == GGML_OP_SET_ROWS); + if (!ggml_node_has_backward(node)) { + GGML_ABORT("%s: no backward through %s (%s): its op has none, or it is an in-place write", __func__, + ggml_op_name(node->op), ggml_op_desc(node)); + } const size_t ihash = ggml_hash_find(&cgraph->visited_hash_set, node); GGML_ASSERT(ihash != GGML_HASHSET_FULL); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 354f4f334ae1..a90000500b1f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -44,6 +44,7 @@ #include #include #include +#include #ifdef __EMSCRIPTEN__ # define N_THREADS 1 @@ -1770,6 +1771,28 @@ struct test_case { ggml_set_loss(out); ggml_build_forward_expand(gf, out); + // A param-reachable node whose op has no backward is not supported, not an abort: the + // full grad suite must reach its summary (ggml_node_has_backward is the one list). + { + std::unordered_set reach; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + ggml_tensor * node = ggml_graph_node(gf, i); + bool reached = false; + for (int k = 0; k < GGML_MAX_SRC; ++k) { + const ggml_tensor * s = node->src[k]; + reached = reached || (s != nullptr && ((s->flags & GGML_TENSOR_FLAG_PARAM) || reach.count(s))); + } + if (!reached) { + continue; + } + reach.insert(node); + if (!ggml_node_has_backward(node)) { + output_printer->print_operation(test_operation_info(op_desc(out), vars(), ggml_backend_name(backend), + test_status_t::NOT_SUPPORTED, std::string("no backward for ") + ggml_op_desc(node))); + return true; + } + } + } ggml_graph_cpy(gf, gb); ggml_build_backward_expand(ctx.get(), gb, nullptr); if (expect.size() != 1 || expect[0] != 0.0f) { @@ -1992,7 +2015,6 @@ struct test_example : public test_case { // GGML_OP_UNARY -static float EXPM1_GRAD_RANGE = 5.0f; struct test_unary : public test_case { const ggml_unary_op op; const ggml_type type; @@ -2053,12 +2075,11 @@ struct test_unary : public test_case { min = -10.f; max = 10.f; } - // EXPM1's gradient is checked: exp past ~88.7 overflows f32, and over a wide range the - // summed loss is too coarse in f32 for the finite difference to see the small gradients - // (exp(-10) against a sum of exp(10)s reads 0). A range where both sides are measurable. + // EXPM1's gradient is checked against a finite difference: exp past ~88.7 overflows f32, + // and the summed loss must stay small enough in f32 for a small step to register. if (op == GGML_UNARY_OP_EXPM1) { - min = -EXPM1_GRAD_RANGE; - max = EXPM1_GRAD_RANGE; + min = -2.f; + max = 2.f; } for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { @@ -2068,7 +2089,18 @@ struct test_unary : public test_case { } float grad_eps() override { - return 15.0f; + // 15 suits the piecewise-linear ops over [-150, 150]; exp's slope changes by e^15 over such + // a step, so EXPM1 takes the smooth ops' convention (test_sin): a small step, the 4-point + // difference, and an f32-sum tolerance (its range is narrowed above) + return op == GGML_UNARY_OP_EXPM1 ? 0.2f : 15.0f; + } + + bool grad_precise() override { + return op == GGML_UNARY_OP_EXPM1; + } + + double max_maa_err() override { + return op == GGML_UNARY_OP_EXPM1 ? 1e-3 : test_case::max_maa_err(); } std::vector grad_expect() override { @@ -2150,33 +2182,30 @@ struct test_glu_split : public test_case { : op(op), type(type), ne_a(ne_a), v(v) {} ggml_tensor * build_graph(ggml_context * ctx) override { - // only SWIGLU has a backward (ggml_compute_backward); marking another glu's inputs as - // params aborts the whole grad suite there and leaves every later case unchecked - const bool grad_supported = op == GGML_GLU_OP_SWIGLU; ggml_tensor * a; ggml_tensor * b; if (v & 1) { auto ne = ne_a; ne[0] *= 3; a = ggml_new_tensor(ctx, type, 4, ne.data()); - if (grad_supported) { ggml_set_param(a); } + ggml_set_param(a); ggml_set_name(a, "a"); a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); ggml_set_name(a, "view_of_a"); b = ggml_new_tensor(ctx, type, 4, ne.data()); - if (grad_supported) { ggml_set_param(b); } + ggml_set_param(b); ggml_set_name(b, "b"); b = ggml_view_4d(ctx, b, ne_a[0], ne_a[1], ne_a[2], ne_a[3], b->nb[1], b->nb[2], b->nb[3], 0); ggml_set_name(a, "view_of_b"); } else { a = ggml_new_tensor(ctx, type, 4, ne_a.data()); - if (grad_supported) { ggml_set_param(a); } + ggml_set_param(a); ggml_set_name(a, "a"); b = ggml_new_tensor(ctx, type, 4, ne_a.data()); - if (grad_supported) { ggml_set_param(b); } + ggml_set_param(b); ggml_set_name(b, "b"); } @@ -2212,29 +2241,31 @@ struct test_swiglu_oai : public test_case { float limit = 7.0f) : type(type), ne_a(ne_a), v(v), alpha(alpha), limit(limit) {} - // no params: SWIGLU_OAI has no backward (ggml_compute_backward), and a param would abort the - // whole grad suite here, leaving every later case unchecked ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * a; ggml_tensor * b; if (v & 1) { auto ne = ne_a; ne[0] *= 3; a = ggml_new_tensor(ctx, type, 4, ne.data()); + ggml_set_param(a); ggml_set_name(a, "a"); a = ggml_view_4d(ctx, a, ne_a[0], ne_a[1], ne_a[2], ne_a[3], a->nb[1], a->nb[2], a->nb[3], 0); ggml_set_name(a, "view_of_a"); b = ggml_new_tensor(ctx, type, 4, ne.data()); + ggml_set_param(b); ggml_set_name(b, "b"); b = ggml_view_4d(ctx, b, ne_a[0], ne_a[1], ne_a[2], ne_a[3], b->nb[1], b->nb[2], b->nb[3], 0); ggml_set_name(a, "view_of_b"); } else { a = ggml_new_tensor(ctx, type, 4, ne_a.data()); + ggml_set_param(a); ggml_set_name(a, "a"); b = ggml_new_tensor(ctx, type, 4, ne_a.data()); + ggml_set_param(b); ggml_set_name(b, "b"); } @@ -5807,7 +5838,7 @@ struct test_pool1d : public test_case { ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * input = ggml_new_tensor(ctx, type_input, 4, ne_input.data()); - // no param: POOL_1D has no backward (ggml_compute_backward aborts the grad suite on it) + ggml_set_param(input); ggml_set_name(input, "input"); ggml_tensor * out = ggml_pool_1d(ctx, input, pool_type, k0, s0, p0); From 3a1c7e21264c7ae3b2b4ce26770a4e52263ac4ab Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 09:08:52 -0500 Subject: [PATCH 36/43] server: /learn builds the served prompt through the same parser /v1/chat/completions uses Measured on her lane (the M5, Qwen3.8-27B, 2026-10-15): every /learn with messages was refused, "slot 1 holds 27439 prompt tokens; the served prefix is 27475 and the reply must not be empty": /learn rendered the conversation through /train's renderer, which ignores chat_template_kwargs, and Qwen3.8's thinking template renders a different prompt without enable_thinking=false (and drops the reply span). It passed on the 4B only because that template renders the same either way. The served prompt is now every message but the last through oaicompat_chat_params_parse and tokenize_input_prompts, the request's own template arguments included: the exact tokens /v1/chat/completions served. The reply is what the model would generate after it: the last message's text and the end of the turn. /learn no longer uses /train's renderer, so its accessor and the exported train_render are removed (train_render is static again). Measured (M5, qwen2.5-0.5b served with Qwen3.8-27B's own chat template from her lane's /props, enable_thinking false): a 1,708-token served conversation + correction, three learns: loss 2.092 -> 1.093 -> 0.316, 221-242 ms each. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/server-context.cpp | 44 ++++++++++++++++----------------- tools/server/server-train.cpp | 2 +- tools/server/server-train.h | 9 ------- 3 files changed, 22 insertions(+), 33 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 1d3037f604a8..a1dcec02376c 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -5804,32 +5804,30 @@ void server_routes::init_routes() { task.learn.lora_id = json_value(body, "lora_id", -1); task.learn.lr = json_value(body, "lr", 1e-4f); if (body.contains("messages")) { - // the conversation she was served, its last message her reply (or the correction): rendered - // through the same template serving used; the served prompt is everything before the reply + // the conversation she was served, its last message her reply (or the correction). The + // served prompt is every earlier message through the SAME parser and template arguments + // /v1/chat/completions used (chat_template_kwargs included: a thinking template renders a + // different prompt without them), with the generation prompt; the reply is what the model + // would have generated after it: the text's tokens and the end of the turn const json & messages = body.at("messages"); - std::vector trained(messages.size(), false); - if (!messages.empty() && messages.back().value("role", std::string()) == "assistant") { - trained.back() = true; - } - std::vector toks; - std::vector loss; - std::string why; - if (!ctx_server.trainer || !train_render(messages, trained, json_value(body, "tools", json()), ctx_server.trainer->renderer(), - ctx_server.vocab, "the learn request", toks, loss, why)) { - res->error(format_error_response(why.empty() ? "the conversation could not be rendered" : why, ERROR_TYPE_INVALID_REQUEST)); + if (!messages.is_array() || messages.size() < 2 || !messages.back().is_object() + || messages.back().value("role", std::string()) != "assistant") { + res->error(format_error_response("\"messages\" must end with the assistant reply to learn, after the conversation it answers", + ERROR_TYPE_INVALID_REQUEST)); return res; } - size_t b = 0; - while (b < loss.size() && !loss[b]) { - ++b; - } - size_t e = loss.size(); - while (e > b && !loss[e - 1]) { - --e; - } - task.learn.prompt.assign(toks.begin(), toks.begin() + b); - task.learn.reply.assign(toks.begin() + b, toks.begin() + e); - task.learn.n_prefix = (int32_t) b; + json served = body; + served["messages"] = json::array(); + for (size_t i = 0; i + 1 < messages.size(); ++i) { + served["messages"].push_back(messages[i]); + } + std::vector files; + const json parsed = oaicompat_chat_params_parse(served, meta->chat_params, files); + task.learn.prompt = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, parsed.at("prompt"), true, true)[0].get_text_tokens(); + task.learn.n_prefix = (int32_t) task.learn.prompt.size(); + task.learn.reply = common_tokenize(ctx_server.vocab, messages.back().value("content", std::string()), false, true); + const llama_token eot = llama_vocab_eot(ctx_server.vocab) != LLAMA_TOKEN_NULL ? llama_vocab_eot(ctx_server.vocab) : llama_vocab_eos(ctx_server.vocab); + task.learn.reply.push_back(eot); } else if (body.contains("reply_tokens")) { task.learn.reply = body.at("reply_tokens").get>(); } else { diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index 845d64a213f9..df43130d4c69 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -74,7 +74,7 @@ std::string confine_out(const std::string & dir, const std::string & name) { // // Returns false with `why` when the example is malformed (the run is refused); true otherwise, // with `toks`/`loss` filled. -bool train_render(const json & messages, const std::vector & trained, const json & tools, +static bool train_render(const json & messages, const std::vector & trained, const json & tools, const server_trainer::render_fn & render, const llama_vocab * vocab, const std::string & at, std::vector & toks, std::vector & loss, std::string & why) { auto head = [&](size_t n) { // the first n messages diff --git a/tools/server/server-train.h b/tools/server/server-train.h index f1fb21166bd1..34c282363940 100644 --- a/tools/server/server-train.h +++ b/tools/server/server-train.h @@ -104,8 +104,6 @@ class server_trainer { }; ~server_trainer(); - // the serving chat template (train_render's render) - const render_fn & renderer() const { return render_chat; } // Starts a run on a worker thread; refuses (ok=false) while one is running or on bad input. common_json start(const common_json & body_in); @@ -209,13 +207,6 @@ class server_trainer { // "fit": "left" for one tokenized example (see server-train.cpp): keep the window+1 tokens that // end the last trained turn (the leading BOS kept when `bos`), mask a trained turn cut at the // front, and return false when the last trained turn alone does not fit. -// A conversation through the serving chat template: its tokens, and which carry loss (the assistant -// turns marked trained), tokenized span by span so no token straddles a boundary. /train and /learn -// share it, so what learns is exactly what serving framed. -bool train_render(const common_json & messages, const std::vector & trained, const common_json & tools, - const server_trainer::render_fn & render, const llama_vocab * vocab, const std::string & at, - std::vector & toks, std::vector & loss, std::string & why); - bool train_fit_left(std::vector & toks, std::vector & loss, int64_t window, bool bos); // "fit": "middle" for one conversation (see server-train.cpp): the index of the oldest history From 56507cceb277d96900d2e5382c42b184a27e1dd7 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 09:11:56 -0500 Subject: [PATCH 37/43] server: one serving tokenizer for chat and /learn; /learn refuses replies it would mislearn Cormac's review of continuum#5061: - P3: serving_prompt_tokens(vocab, mctx, prompt, files, chat) is now the one place a parsed request becomes the tokens it is served as (mtmd for a chat request with a projector, else tokenize_input_prompts). handle_completions_impl and /learn both call it, so "the slot holds what serving put there" is true by construction. - P2: a reply is learned as its text then the end of the turn. A reply whose substance is elsewhere is refused by name until it renders through the template: tool_calls (a coding citizen's commonest turn; it would have taught ending the turn instead of calling the tool), reasoning_content, or no text. A conversation with media is refused too (text prompts only). Measured (M5, qwen2.5-0.5b with Qwen3.8-27B's chat template): learning unchanged, loss 2.092 -> 1.093 -> 0.316 (226-255 ms); a tool-call reply, a reasoning reply and an empty reply each refused with 400 and the reason. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/server-context.cpp | 51 ++++++++++++++++++++++++++------- 1 file changed, 40 insertions(+), 11 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a1dcec02376c..d88bf25a3ad1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -4698,6 +4698,21 @@ void server_context::set_state_callback(server_state_callback_t callback) { // server_routes // +// The tokens a parsed request is served as: a chat request with a projector loaded goes through +// mtmd (its media markers and files), everything else through tokenize_input_prompts. One place, so +// what /learn checks a slot against is what serving put there by construction. +static std::vector serving_prompt_tokens(const llama_vocab * vocab, mtmd_context * mctx, const json & prompt, + const std::vector & files, bool chat) { + if (chat && mctx != nullptr) { + // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. + std::vector out; + out.push_back(process_mtmd_prompt(mctx, prompt.get(), files)); + return out; + } + // Everything else, including multimodal completions. + return tokenize_input_prompts(vocab, mctx, prompt, true, true); +} + std::unique_ptr server_routes::handle_completions_impl( const server_http_req & req, server_task_type type, @@ -4733,15 +4748,8 @@ std::unique_ptr server_routes::handle_completions_impl( } // process prompt - std::vector inputs; - - if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) { - // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. - inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files)); - } else { - // Everything else, including multimodal completions. - inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true); - } + std::vector inputs = serving_prompt_tokens(ctx_server.vocab, ctx_server.mctx, prompt, files, + res_type != TASK_RESPONSE_TYPE_NONE); // tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks @@ -5816,6 +5824,23 @@ void server_routes::init_routes() { ERROR_TYPE_INVALID_REQUEST)); return res; } + // the reply is learned as its text then the end of the turn; a reply whose substance is + // elsewhere (tool calls, reasoning) would teach her to end the turn instead (Cormac's + // review of #5061), so it is refused until those render through the template + const json & last = messages.back(); + const json no_value; + const json & calls = last.contains("tool_calls") ? last.at("tool_calls") : no_value; + const json & reasoning = last.contains("reasoning_content") ? last.at("reasoning_content") : no_value; + const bool has_calls = calls.is_array() && !calls.empty(); + const bool has_reasoning = reasoning.is_string() && !reasoning.get().empty(); + const bool has_text = last.contains("content") && last.at("content").is_string() && !last.at("content").get().empty(); + if (has_calls || has_reasoning || !has_text) { + res->error(format_error_response(has_calls ? "a reply with tool_calls cannot be learned yet (only its text would be: it would teach ending the turn)" + : has_reasoning ? "a reply with reasoning_content cannot be learned yet (its reasoning would be dropped)" + : "the reply to learn has no text content", + ERROR_TYPE_INVALID_REQUEST)); + return res; + } json served = body; served["messages"] = json::array(); for (size_t i = 0; i + 1 < messages.size(); ++i) { @@ -5823,9 +5848,13 @@ void server_routes::init_routes() { } std::vector files; const json parsed = oaicompat_chat_params_parse(served, meta->chat_params, files); - task.learn.prompt = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, parsed.at("prompt"), true, true)[0].get_text_tokens(); + if (!files.empty()) { + res->error(format_error_response("a conversation with media cannot be learned yet: only text prompts", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + task.learn.prompt = serving_prompt_tokens(ctx_server.vocab, ctx_server.mctx, parsed.at("prompt"), files, true)[0].get_text_tokens(); task.learn.n_prefix = (int32_t) task.learn.prompt.size(); - task.learn.reply = common_tokenize(ctx_server.vocab, messages.back().value("content", std::string()), false, true); + task.learn.reply = common_tokenize(ctx_server.vocab, last.at("content").get(), false, true); const llama_token eot = llama_vocab_eot(ctx_server.vocab) != LLAMA_TOKEN_NULL ? llama_vocab_eot(ctx_server.vocab) : llama_vocab_eos(ctx_server.vocab); task.learn.reply.push_back(eot); } else if (body.contains("reply_tokens")) { From b7c92db19160a803f7d030d7985494982b8dff18 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 15:47:52 -0500 Subject: [PATCH 38/43] llama: a labelled position carries a weight in the loss (llama_batch.label_weight) How much a token counts, not only whether: a mind's relevance on what it perceived scales its surprise. llama_batch gains a trailing `float * label_weight` (NULL: every labelled position counts 1; >= 0 and finite; read for labelled positions only). llama_batch_init allocates it at 1; the ubatch split carries it; every batch view that slices `label` slices it too (parallel, batched-bench, perplexity, mtmd's media batch), so a view never drops a weight. ggml-opt: ggml_opt_set_next_label_weights(opt_ctx, weighted) asks the next graph for a weighted cross-entropy (ggml_cross_entropy_loss_weighted per head slice, over ggml_opt_label_weights); an unweighted batch keeps the unweighted loss, so CUDA, which has no weighted kernel yet, trains an unweighted batch on the GPU as before. llama sets it per ubatch from label_weight. Measured (M5, CPU quick check, qwen2.5-0.5b, test-walk-exact WALK_MIXED_ONLY, F32 and flash q8_0): weights 1,1 vs an unweighted batch 1.0000/1.0000; weights 2,2 vs twice the step 1.0000/1.0000; weights 1,0 vs half of row 0's own step 1.0000/1.0000; the mixed-decode checks still OK. Build has zero warnings. Metal run owed (held while the M5's deploy can relaunch its lane). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- examples/parallel/parallel.cpp | 1 + ggml/include/ggml-opt.h | 8 ++++++- ggml/src/ggml-opt.cpp | 30 +++++++++++++++++++++++++-- include/llama.h | 4 ++++ src/llama-batch.cpp | 18 ++++++++++++++++ src/llama-batch.h | 2 ++ src/llama-context-opt.cpp | 10 +++++++++ src/llama-kv-cache-dsv4.cpp | 1 + tests/test-walk-exact.cpp | 20 ++++++++++++++++-- tools/batched-bench/batched-bench.cpp | 1 + tools/mtmd/mtmd-helper-common.h | 2 ++ tools/perplexity/perplexity.cpp | 1 + 12 files changed, 93 insertions(+), 5 deletions(-) diff --git a/examples/parallel/parallel.cpp b/examples/parallel/parallel.cpp index 45fded7f54f5..643a05530994 100644 --- a/examples/parallel/parallel.cpp +++ b/examples/parallel/parallel.cpp @@ -401,6 +401,7 @@ int main(int argc, char ** argv) { batch.seq_id + i, batch.logits + i, batch.label ? batch.label + i : nullptr, + batch.label && batch.label_weight ? batch.label_weight + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index 6a78d2795290..c41c791aed78 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h @@ -158,7 +158,13 @@ extern "C" { // if not using static graphs these pointers become invalid with the next call to ggml_opt_alloc GGML_API struct ggml_tensor * ggml_opt_inputs( ggml_opt_context_t opt_ctx); // forward graph input tensor GGML_API struct ggml_tensor * ggml_opt_outputs( ggml_opt_context_t opt_ctx); // forward graph output tensor - GGML_API struct ggml_tensor * ggml_opt_labels( ggml_opt_context_t opt_ctx); // labels to compare outputs against + GGML_API struct ggml_tensor * ggml_opt_labels( ggml_opt_context_t opt_ctx); + // the per-row weights of a weighted cross-entropy step (F32, one per label), or NULL when the + // graph prepared last is unweighted (ggml_opt_set_next_label_weights) + GGML_API struct ggml_tensor * ggml_opt_label_weights(ggml_opt_context_t opt_ctx); + // whether the next graph's cross-entropy weighs each row (by ggml_opt_label_weights); an + // unweighted step keeps the unweighted loss, whose kernels every backend has + GGML_API void ggml_opt_set_next_label_weights(ggml_opt_context_t opt_ctx, bool weighted); // labels to compare outputs against GGML_API struct ggml_tensor * ggml_opt_loss( ggml_opt_context_t opt_ctx); // scalar tensor that contains the loss GGML_API struct ggml_tensor * ggml_opt_pred( ggml_opt_context_t opt_ctx); // predictions made by outputs GGML_API struct ggml_tensor * ggml_opt_ncorrect(ggml_opt_context_t opt_ctx); // number of matching predictions between outputs and labels diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 11519f0fa2f8..e2a1d0fcb2dc 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -90,6 +90,9 @@ struct ggml_opt_context { // graph prepared last (ggml_opt_prepare_alloc clears them) std::unordered_map next_grad_accs; std::vector grad_accs; + // a weighted cross-entropy for the next graph, and that graph's weights (one per label row) + bool next_label_weights = false; + struct ggml_tensor * label_weights = nullptr; // ggml_opt_prepare_alloc_head: the outputs' row slices, in row order (empty: one outputs) std::vector output_slices; @@ -623,8 +626,22 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { : ggml_dup_tensor(ctx_results, opt_ctx->outputs); ggml_set_input(opt_ctx->labels); ggml_set_name(opt_ctx->labels, "labels"); + opt_ctx->label_weights = nullptr; + if (opt_ctx->next_label_weights) { + GGML_ASSERT(opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS && "a weight per row needs one target index per row"); + opt_ctx->label_weights = ggml_new_tensor_1d(ctx_results, GGML_TYPE_F32, opt_ctx->labels->ne[0]); + ggml_set_input(opt_ctx->label_weights); + ggml_set_name(opt_ctx->label_weights, "label_weights"); + } + auto row_loss = [&](ggml_tensor * out, ggml_tensor * lab, int64_t r0) { + if (opt_ctx->label_weights == nullptr) { + return ggml_cross_entropy_loss(ctx_results, out, lab); + } + ggml_tensor * w = ggml_view_1d(ctx_results, opt_ctx->label_weights, out->ne[1], r0*sizeof(float)); + return ggml_cross_entropy_loss_weighted(ctx_results, out, lab, w); + }; if (opt_ctx->output_slices.empty()) { - opt_ctx->loss = ggml_cross_entropy_loss(ctx_results, opt_ctx->outputs, opt_ctx->labels); + opt_ctx->loss = row_loss(opt_ctx->outputs, opt_ctx->labels, 0); } else { // the mean over every row: each slice's mean weighted by its share of the rows const int64_t n = opt_ctx->labels->ne[0]; @@ -633,7 +650,7 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { for (ggml_tensor * out : opt_ctx->output_slices) { const int64_t nr = out->ne[1]; ggml_tensor * lab = ggml_view_1d(ctx_results, opt_ctx->labels, nr, r0*sizeof(int32_t)); - ggml_tensor * l = ggml_scale(ctx_results, ggml_cross_entropy_loss(ctx_results, out, lab), (float) nr / (float) n); + ggml_tensor * l = ggml_scale(ctx_results, row_loss(out, lab, r0), (float) nr / (float) n); sum = sum ? ggml_add(ctx_results, sum, l) : l; r0 += nr; } @@ -924,6 +941,15 @@ struct ggml_tensor * ggml_opt_labels(ggml_opt_context_t opt_ctx) { return opt_ctx->labels; } +struct ggml_tensor * ggml_opt_label_weights(ggml_opt_context_t opt_ctx) { + return opt_ctx->label_weights; +} + +void ggml_opt_set_next_label_weights(ggml_opt_context_t opt_ctx, bool weighted) { + GGML_ASSERT(!opt_ctx->eval_ready && "set before ggml_opt_alloc"); + opt_ctx->next_label_weights = weighted; +} + struct ggml_tensor * ggml_opt_loss(ggml_opt_context_t opt_ctx) { return opt_ctx->loss; } diff --git a/include/llama.h b/include/llama.h index ee65be13aa27..c322a486deaf 100644 --- a/include/llama.h +++ b/include/llama.h @@ -257,6 +257,9 @@ extern "C" { // the optimizer write of the context's trainable adapter (llama_opt_init), in the // same graph as the batch's other rows. -1 marks a position with no label. // (if set to NULL, the batch is inference only, unchanged) + // - label_weight : how much a labelled position counts in the loss (>= 0, finite; 0 counts + // nothing): a mind's relevance on what it perceived, not only whether. Read for + // labelled positions only. (if set to NULL, every labelled position counts 1) // typedef struct llama_batch { int32_t n_tokens; @@ -268,6 +271,7 @@ extern "C" { llama_seq_id ** seq_id; int8_t * logits; // TODO: rename this to "output" llama_token * label; + float * label_weight; } llama_batch; enum llama_model_kv_override_type { diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 5ec36436c7e8..b4a1626afba6 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1,5 +1,7 @@ #include "llama-batch.h" +#include + #include "llama-impl.h" #include "llama-vocab.h" #include "llama-memory.h" @@ -177,6 +179,10 @@ bool llama_batch_allocr::init( LLAMA_LOG_ERROR("%s: invalid label[%d] = %d\n", __func__, i, batch.label[i]); return false; } + if (batch.label[i] >= 0 && batch.label_weight && !(std::isfinite(batch.label_weight[i]) && batch.label_weight[i] >= 0.0f)) { + LLAMA_LOG_ERROR("%s: invalid label_weight[%d] = %f: a weight is finite and >= 0\n", __func__, i, batch.label_weight[i]); + return false; + } const llama_seq_id s = batch.seq_id[i][0]; if (batch.label[i] < 0) { for (int32_t j = 0; j < batch.n_seq_id[i]; ++j) { @@ -264,6 +270,7 @@ bool llama_batch_allocr::init( /*.seq_idx =*/ this->seq_idx.data(), /*.output =*/ batch.logits, /*.label =*/ batch.label, + /*.label_weight =*/ batch.label_weight, /*.data =*/ {}, }; @@ -471,6 +478,7 @@ llama_ubatch llama_batch_allocr::ubatch_reserve(uint32_t n_seq_tokens, uint32_t /*.seq_idx =*/ udata->seq_idx.data(), /*.output =*/ udata->output.data(), /*.label =*/ nullptr, + /*.label_weight =*/ nullptr, /*.data =*/ std::move(udata), }; @@ -811,6 +819,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u udata->seq_idx .resize(LLAMA_MAX_SEQ, -1); udata->output .resize(n_tokens); udata->label .resize(batch.label ? n_tokens : 0); + udata->label_weight.resize(batch.label && batch.label_weight ? n_tokens : 0); udata->seq_id_data.reserve(n_tokens); @@ -838,6 +847,9 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u udata->output[i] = batch.logits[idxs[i]]; if (batch.label) { udata->label[i] = batch.label[idxs[i]]; + if (batch.label_weight) { + udata->label_weight[i] = batch.label_weight[idxs[i]]; + } } for (int s = 0; s < udata->n_seq_id[i]; ++s) { @@ -882,6 +894,7 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u /*.seq_idx =*/ udata->seq_idx.data(), /*.output =*/ udata->output.data(), /*.label =*/ batch.label ? udata->label.data() : nullptr, + /*.label_weight =*/ batch.label && batch.label_weight ? udata->label_weight.data() : nullptr, /*.data =*/ std::move(udata), }; @@ -991,6 +1004,7 @@ struct llama_batch llama_batch_get_one( /*seq_id =*/ nullptr, /*logits =*/ nullptr, /*label =*/ nullptr, + /*label_weight =*/ nullptr, }; } @@ -1004,6 +1018,7 @@ struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_ /*seq_id =*/ nullptr, /*logits =*/ nullptr, /*label =*/ nullptr, + /*label_weight =*/ nullptr, }; if (embd) { @@ -1023,6 +1038,8 @@ struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_ batch.logits = (int8_t *) malloc(sizeof(int8_t) * n_tokens_alloc); batch.label = (llama_token *) malloc(sizeof(llama_token) * n_tokens_alloc); std::fill(batch.label, batch.label + n_tokens_alloc, -1); // no position trains until labelled + batch.label_weight = (float *) malloc(sizeof(float) * n_tokens_alloc); + std::fill(batch.label_weight, batch.label_weight + n_tokens_alloc, 1.0f); // every labelled position counts 1 until weighted return batch; } @@ -1040,4 +1057,5 @@ void llama_batch_free(struct llama_batch batch) { } if (batch.logits) free(batch.logits); if (batch.label) free(batch.label); + if (batch.label_weight) free(batch.label_weight); } diff --git a/src/llama-batch.h b/src/llama-batch.h index 6430ac1d457a..fca6f0571f6b 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -78,6 +78,7 @@ struct llama_ubatch { int32_t * seq_idx; // [LLAMA_MAX_SEQ] | - | seq_idx int8_t * output; // [n_tokens] | i | - llama_token * label; // [n_tokens] | i | training target, -1 for none (nullptr: inference) + float * label_weight; // [n_tokens] | i | a labelled row's weight in the loss (nullptr: 1) struct data_t { std::vector token; @@ -89,6 +90,7 @@ struct llama_ubatch { std::vector seq_idx; std::vector output; std::vector label; + std::vector label_weight; std::vector seq_id_data; }; diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 3090bfcd2293..42969c259a86 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -366,10 +366,16 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l ctx_compute_opt = ggml_init(params); } std::vector targets; + std::vector weights; // a labelled row's weight in the loss, when the batch carries them targets.reserve(n_labels); for (int32_t row : ubatch.rows().label_rows) { targets.push_back(ubatch.label[row]); + if (ubatch.label_weight != nullptr) { + weights.push_back(ubatch.label_weight[row]); + } } + // an unweighted batch keeps the unweighted loss (every backend has its kernels) + ggml_opt_set_next_label_weights(opt_ctx, ubatch.label_weight != nullptr); if (res->t_loss_slices.empty()) { ggml_free(ctx_compute_opt); opt_fail("training head was not projected into bounded label slices"); @@ -414,6 +420,10 @@ bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const l GGML_ASSERT(labels_t->type == GGML_TYPE_I32 && labels_t->ne[0] == (int64_t) n_labels); ggml_backend_tensor_set(labels_t, targets.data(), 0, ggml_nbytes(labels_t)); + if (struct ggml_tensor * weights_t = ggml_opt_label_weights(opt_ctx)) { + GGML_ASSERT(weights.size() == (size_t) n_labels); + ggml_backend_tensor_set(weights_t, weights.data(), 0, ggml_nbytes(weights_t)); + } } if (xc != nullptr) { // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 52e843388bba..e6b2465cffad 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -165,6 +165,7 @@ static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) { /*.seq_idx =*/ data->seq_idx.data(), /*.output =*/ data->output.data(), /*.label =*/ nullptr, + /*.label_weight =*/ nullptr, /*.data =*/ data, }; diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 947b30c7a0e3..9a446b7b86c1 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -354,7 +354,8 @@ static void mixed_decode(const common_params & params, llama_model * model, cons static void multi_label_decode(const common_params & params, llama_model * model, const std::string & init, const std::vector & tokens, bool flash, ggml_type cache_type) { - auto step = [&](const std::vector & sequences) { + // weights empty: an unweighted batch (label_weight NULL) + auto step = [&](const std::vector & sequences, const std::vector & weights = {}) { auto cp = common_context_params_to_llama(params); cp.n_ctx = 128; cp.n_batch = cp.n_ubatch = 2; @@ -375,11 +376,20 @@ static void multi_label_decode(const common_params & params, llama_model * model llama_opt_init(ctx, model, op); const auto before = adapter_params(adapter); llama_batch batch = llama_batch_init(2, 0, 1); - for (int seq : sequences) { + for (size_t k = 0; k < sequences.size(); ++k) { + const int seq = sequences[k]; common_batch_add(batch, tokens[seq], 0, {seq}, true); batch.label[batch.n_tokens - 1] = tokens[seq + 2]; + if (!weights.empty()) { + batch.label_weight[batch.n_tokens - 1] = weights[k]; + } + } + float * kept_weights = batch.label_weight; + if (weights.empty()) { + batch.label_weight = nullptr; // unweighted: the loss every backend has } GGML_ASSERT(llama_decode(ctx, batch) == 0); + batch.label_weight = kept_weights; auto delta = adapter_params(adapter); for (size_t i = 0; i < delta.size(); ++i) { delta[i] -= before[i]; @@ -397,6 +407,12 @@ static void multi_label_decode(const common_params & params, llama_model * model const auto joint = step({0, 1}); GGML_ASSERT(same_step(compare("two labelled sequences vs mean independent steps", joint, mean))); GGML_ASSERT(same_step(compare("two labelled sequences with reversed row order", step({1, 0}), joint))); + // what this catches: the per-row weight not reaching the loss, or reaching it wrong (each row + // counts w_i * its loss / rows, its gradient scaled by w_i; a mind's relevance per token) + auto scaled = [](std::vector v, float s) { for (float & x : v) { x *= s; } return v; }; + GGML_ASSERT(same_step(compare("weights 1,1 vs an unweighted batch", step({0, 1}, {1.0f, 1.0f}), joint))); + GGML_ASSERT(same_step(compare("weights 2,2 vs twice the unweighted step", step({0, 1}, {2.0f, 2.0f}), scaled(joint, 2.0f)))); + GGML_ASSERT(same_step(compare("weights 1,0 vs half of row 0's own step (a mean over 2 rows)", step({0, 1}, {1.0f, 0.0f}), scaled(step({0}), 0.5f)))); } static void sparse_head_decode(const common_params & params, llama_model * model, const std::string & init, const std::vector & tokens) { diff --git a/tools/batched-bench/batched-bench.cpp b/tools/batched-bench/batched-bench.cpp index 04cb25c62ac1..4d39af212aa6 100644 --- a/tools/batched-bench/batched-bench.cpp +++ b/tools/batched-bench/batched-bench.cpp @@ -92,6 +92,7 @@ int llama_batched_bench(int argc, char ** argv) { batch.seq_id + i, batch.logits + i, batch.label ? batch.label + i : nullptr, + batch.label && batch.label_weight ? batch.label_weight + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); diff --git a/tools/mtmd/mtmd-helper-common.h b/tools/mtmd/mtmd-helper-common.h index 19c0f9ed0942..dab628d1a0d6 100644 --- a/tools/mtmd/mtmd-helper-common.h +++ b/tools/mtmd/mtmd-helper-common.h @@ -97,6 +97,7 @@ struct decode_embd_batch { /*seq_id =*/ seq_ids.data(), /*logits =*/ logits.data(), /*label =*/ nullptr, + /*label_weight =*/ nullptr, }; } @@ -181,6 +182,7 @@ struct decode_embd_batch { /*seq_id =*/ batch.seq_id + offset, /*logits =*/ batch.logits + offset, /*label =*/ batch.label ? batch.label + offset : nullptr, + /*label_weight =*/ batch.label && batch.label_weight ? batch.label_weight + offset : nullptr, }; } }; diff --git a/tools/perplexity/perplexity.cpp b/tools/perplexity/perplexity.cpp index 80c834bb79b5..a7e12041ea9d 100644 --- a/tools/perplexity/perplexity.cpp +++ b/tools/perplexity/perplexity.cpp @@ -675,6 +675,7 @@ static bool decode_helper(llama_context * ctx, llama_batch & batch, std::vector< batch.seq_id + i, batch.logits + i, batch.label ? batch.label + i : nullptr, + batch.label && batch.label_weight ? batch.label_weight + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); From 9a1099343722a151e0d26c525697a052000188a2 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 18:19:47 -0500 Subject: [PATCH 39/43] server: a completion can carry a learn, run on its own slot before the slot is released On the M5's 27B (one slot, busy almost always) a separate POST /learn lost the race for the slot 12 times in 12, 2026-10-15: another request took it between the turn and the learn. A completion request may now carry "learn": {"reply": text, "lora_id", "lr"}: when the turn finishes, its reply is learned on that same slot before release (no request can take it in between), and the outcome is in the final response's "learn" field. One learner: learn_on_slot(slot, request, t0) holds everything from the adapter checks to the slot's restore to its served prompt; the SERVER_TASK_TYPE_LEARN case keeps only what is a task's (finding the slot, deferring while it serves, the duty cycle) and calls it, and both finish paths (plain and speculative) call learn_if_asked before send_final_response and release. The learn types are named (learn_request, learn_outcome) and shared by the task, the result and the completion params. The server's batch view slices label_weight with label. Measured (M5, Metal, qwen2.5-0.5b with Qwen3.8-27B's chat template, ONE slot, a second client taking it continuously): 4 learn-carrying turns, 4 succeeded, loss 2.178 -> 1.368 -> 0.282 -> 0.092, 889-962 ms each, 5 competing requests served on the same slot meanwhile; a learn with an empty reply refused (400). Full build: zero warnings. Found, NOT caused by this (card owed): training on a CPU-only server (-ngl 0) crashes in ggml_backend_tensor_copy (a null copy function) inside llama_opt_learn's decode, through POST /learn as well; every earlier learn receipt ran on Metal. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/server-context.cpp | 286 ++++++++++++++++++-------------- tools/server/server-schema.cpp | 16 ++ tools/server/server-task.cpp | 11 +- tools/server/server-task.h | 48 ++++-- 4 files changed, 217 insertions(+), 144 deletions(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index d88bf25a3ad1..04478e4b87b6 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -206,6 +206,7 @@ struct server_batch { batch.seq_id + off, batch.logits + off, batch.label ? batch.label + off : nullptr, + batch.label && batch.label_weight ? batch.label_weight + off : nullptr, }; return view; @@ -347,6 +348,7 @@ struct server_slot { int32_t n_gen_last = 0; void reset() { + learned.reset(); SLT_DBG(*this, "%s", "\n"); spec_is_replay = false; @@ -521,6 +523,9 @@ struct server_slot { prompt.tokens.insert(spec_draft); } + // the outcome of the finishing task's learn (task_params::learn), for its final response + std::optional learned; + void release() { if (is_processing()) { GGML_ASSERT(task); @@ -2057,11 +2062,160 @@ struct server_context_impl { queue_results.send(std::move(res)); } + // One reply learned on `slot`, which holds the served prompt (POST /learn between its turns, or a + // completion's own learn before its slot is released): the slot's stable set plus the shadow on + // the reply's labelled rows, then the slot back to its whole served prompt. + learn_outcome learn_on_slot(server_slot & slot, const learn_request & L, int64_t t0) { + auto refused = [&](const std::string & why) { + learn_outcome out; + out.failure = why; + SRV_WRN("learn refused: %s\n", why.c_str()); + return out; + }; + auto & loras = params_base.lora_adapters; + if (L.lora_id < 0 || (size_t) L.lora_id >= loras.size() || loras[L.lora_id].ptr == nullptr) { + return refused("no loaded adapter " + std::to_string(L.lora_id) + " to train (POST /lora-adapters/load it first)"); + } + if (learn_lora_id >= 0 && learn_lora_id != L.lora_id) { + return refused("this server trains adapter " + std::to_string(learn_lora_id) + "; one shadow per server run"); + } + const int32_t n0 = L.n_prefix - 1; // the last prompt token is re-decoded with the reply's first token as its label + if (n0 < 0 || L.n_prefix > slot.prompt.n_tokens() || L.reply.empty()) { + return refused("slot " + std::to_string(slot.id) + " holds " + std::to_string(slot.prompt.n_tokens()) + + " prompt tokens; the served prefix is " + std::to_string(L.n_prefix) + " and the reply must not be empty"); + } + // the served prompt, when the request carries it: her slot must still hold exactly it (a + // slot that has served other turns since, or compacted, would teach the wrong context) + if (!L.prompt.empty()) { + int32_t i = 0; + while (i < L.n_prefix && slot.prompt.tokens[i] == L.prompt[i]) { + ++i; + } + if (i < L.n_prefix) { + return refused("slot " + std::to_string(slot.id) + " no longer holds this turn's prompt (it differs at token " + + std::to_string(i) + " of " + std::to_string(L.n_prefix) + "): serve the prompt again, then learn"); + } + } + // the recurrent part cannot be cut back: her state one token before the prompt's end + const common_prompt_checkpoint * cp = nullptr; + const bool needs_cp = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS; + if (needs_cp) { + for (const auto & c : slot.prompt.checkpoints) { + if (c.n_tokens == n0) { + cp = &c; + } + } + if (cp == nullptr) { + return refused("slot " + std::to_string(slot.id) + " has no checkpoint one token before its prompt's end (" + + std::to_string(n0) + "): the lane learns only with --train-dir, which keeps one"); + } + } + // the tokens from the last prompt token on, each labelled with the next + std::vector toks; + toks.push_back(slot.prompt.tokens[n0]); + toks.insert(toks.end(), L.reply.begin(), L.reply.end()); + std::vector labels(toks.size(), -1); + for (size_t i = 0; i + 1 < toks.size(); ++i) { + labels[i] = toks[i + 1]; + } + // her stable set plus the shadow: the shadow applies to labelled rows only + std::vector set = slot.lora.empty() ? construct_lora_list({}) : slot.lora; + if ((size_t) L.lora_id >= set.size()) { + set = construct_lora_list({}); + } + set[L.lora_id].scale = 1.0f; + common_set_adapter_lora(ctx_tgt, set); + learn_lr = L.lr; + if (learn_lora_id < 0) { + llama_opt_params lopt{}; + lopt.param_filter = llama_opt_param_filter_all; + lopt.get_opt_pars = learn_opt_pars; + lopt.get_opt_pars_ud = &learn_lr; + lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_ADAMW; + lopt.adapter = loras[L.lora_id].ptr; + llama_opt_init(ctx_tgt, model_tgt, lopt); + learn_lora_id = L.lora_id; + } + // drop what follows n0 (attention) and put her recurrent state back at n0 + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + ggml_opt_result_t result = ggml_opt_result_init(); + const bool ok = llama_opt_learn(ctx_tgt, slot.id, n0, toks.data(), labels.data(), (int32_t) toks.size(), result); + double loss = 0.0, unc = 0.0; + if (ok) { + ggml_opt_result_loss(result, &loss, &unc); + } + ggml_opt_result_free(result); + // her own adapter set again before anything else is decoded for her (the shadow is + // already off for a batch without labels: llama-graph.cpp lora_on, n_shadow_rows = + // the ubatch's label count; this does not rely on it) + common_set_adapter_lora(ctx_tgt, slot.lora); + // her slot goes back to her whole served prompt (attention cells, recurrent state, + // tokens and checkpoints agree; her next turn reuses all of it): back to n0 from the + // checkpoint, then the last prompt token decoded again under her stable set. A draft + // model's cache is cut to n0 with its own checkpoint, and the slot's tokens with it. + llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot.id, n0, -1); + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + int32_t n_keep = n0; + if (ctx_dft != nullptr) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, n0, -1); + if (cp != nullptr) { + cp->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + } else { + llama_batch one = llama_batch_get_one(&toks[0], 1); + llama_pos pos = n0; + one.pos = &pos; + llama_seq_id sid = slot.id; + int32_t nsid = 1; + llama_seq_id * sids = &sid; + one.n_seq_id = &nsid; + one.seq_id = &sids; + if (llama_decode(ctx_tgt, one) == 0) { + n_keep = L.n_prefix; + } else { + llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot.id, n0, -1); + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + } + } + slot.prompt.tokens.keep_first(n_keep); + slot.prompt.checkpoints.remove_if([&](const common_prompt_checkpoint & c) { return c.n_tokens > n_keep; }); + learn_outcome out; + out.ok = ok; + out.failure = ok ? "" : llama_opt_failure(ctx_tgt); + out.loss = loss; + out.n_labels = (int32_t) toks.size() - 1; + out.t_ms = (ggml_time_us() - t0) / 1000.0; + SRV_INF("learned on slot %d: %d labels from position %d, loss %.4f, %.1f ms%s\n", slot.id, out.n_labels, n0, loss, + out.t_ms, ok ? "" : (" (failed: " + out.failure + ")").c_str()); + return out; + } + + // A completion that carries a learn learns its reply on its own slot as it finishes, before the + // slot is released: no other request can take the slot between the turn and its learning (on + // a one-slot lane that is busy almost always, a separate POST /learn lost that race 12 times in + // 12 on the M5's 27B, 2026-10-15). + void learn_if_asked(server_slot & slot) { + if (!slot.task->params.learn) { + return; + } + learn_request L = *slot.task->params.learn; + L.id_slot = slot.id; + L.n_prefix = slot.task->n_tokens(); + slot.learned = learn_on_slot(slot, L, ggml_time_us()); + } + void send_final_response(server_slot & slot) { auto res = std::make_unique(); res->id = slot.task->id; res->id_slot = slot.id; + res->learned = slot.learned; res->index = slot.task->index; @@ -2880,8 +3034,8 @@ struct server_context_impl { const int64_t t0 = ggml_time_us(); server_slot * slot = get_slot_by_id(L.id_slot); auto refuse = [&](const std::string & why) { - res->ok = false; - res->failure = why; + res->outcome.ok = false; + res->outcome.failure = why; SRV_WRN("learn refused: %s\n", why.c_str()); queue_results.send(std::move(res)); }; @@ -2904,132 +3058,8 @@ struct server_context_impl { break; } } - auto & loras = params_base.lora_adapters; - if (L.lora_id < 0 || (size_t) L.lora_id >= loras.size() || loras[L.lora_id].ptr == nullptr) { - refuse("no loaded adapter " + std::to_string(L.lora_id) + " to train (POST /lora-adapters/load it first)"); - break; - } - if (learn_lora_id >= 0 && learn_lora_id != L.lora_id) { - refuse("this server trains adapter " + std::to_string(learn_lora_id) + "; one shadow per server run"); - break; - } - const int32_t n0 = L.n_prefix - 1; // the last prompt token is re-decoded with the reply's first token as its label - if (n0 < 0 || L.n_prefix > slot->prompt.n_tokens() || L.reply.empty()) { - refuse("slot " + std::to_string(slot->id) + " holds " + std::to_string(slot->prompt.n_tokens()) - + " prompt tokens; the served prefix is " + std::to_string(L.n_prefix) + " and the reply must not be empty"); - break; - } - // the served prompt, when the request carries it: her slot must still hold exactly it (a - // slot that has served other turns since, or compacted, would teach the wrong context) - if (!L.prompt.empty()) { - int32_t i = 0; - while (i < L.n_prefix && slot->prompt.tokens[i] == L.prompt[i]) { - ++i; - } - if (i < L.n_prefix) { - refuse("slot " + std::to_string(slot->id) + " no longer holds this turn's prompt (it differs at token " - + std::to_string(i) + " of " + std::to_string(L.n_prefix) + "): serve the prompt again, then learn"); - break; - } - } - // the recurrent part cannot be cut back: her state one token before the prompt's end - const common_prompt_checkpoint * cp = nullptr; - const bool needs_cp = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS; - if (needs_cp) { - for (const auto & c : slot->prompt.checkpoints) { - if (c.n_tokens == n0) { - cp = &c; - } - } - if (cp == nullptr) { - refuse("slot " + std::to_string(slot->id) + " has no checkpoint one token before its prompt's end (" - + std::to_string(n0) + "): the lane learns only with --train-dir, which keeps one"); - break; - } - } - // the tokens from the last prompt token on, each labelled with the next - std::vector toks; - toks.push_back(slot->prompt.tokens[n0]); - toks.insert(toks.end(), L.reply.begin(), L.reply.end()); - std::vector labels(toks.size(), -1); - for (size_t i = 0; i + 1 < toks.size(); ++i) { - labels[i] = toks[i + 1]; - } - // her stable set plus the shadow: the shadow applies to labelled rows only - std::vector set = slot->lora.empty() ? construct_lora_list({}) : slot->lora; - if ((size_t) L.lora_id >= set.size()) { - set = construct_lora_list({}); - } - set[L.lora_id].scale = 1.0f; - common_set_adapter_lora(ctx_tgt, set); - learn_lr = L.lr; - if (learn_lora_id < 0) { - llama_opt_params lopt{}; - lopt.param_filter = llama_opt_param_filter_all; - lopt.get_opt_pars = learn_opt_pars; - lopt.get_opt_pars_ud = &learn_lr; - lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_ADAMW; - lopt.adapter = loras[L.lora_id].ptr; - llama_opt_init(ctx_tgt, model_tgt, lopt); - learn_lora_id = L.lora_id; - } - // drop what follows n0 (attention) and put her recurrent state back at n0 - if (cp != nullptr) { - cp->load_tgt(ctx_tgt, slot->id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - } - ggml_opt_result_t result = ggml_opt_result_init(); - const bool ok = llama_opt_learn(ctx_tgt, slot->id, n0, toks.data(), labels.data(), (int32_t) toks.size(), result); - double loss = 0.0, unc = 0.0; - if (ok) { - ggml_opt_result_loss(result, &loss, &unc); - } - ggml_opt_result_free(result); - // her own adapter set again before anything else is decoded for her (the shadow is - // already off for a batch without labels: llama-graph.cpp lora_on, n_shadow_rows = - // the ubatch's label count; this does not rely on it) - common_set_adapter_lora(ctx_tgt, slot->lora); - // her slot goes back to her whole served prompt (attention cells, recurrent state, - // tokens and checkpoints agree; her next turn reuses all of it): back to n0 from the - // checkpoint, then the last prompt token decoded again under her stable set. A draft - // model's cache is cut to n0 with its own checkpoint, and the slot's tokens with it. - llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot->id, n0, -1); - if (cp != nullptr) { - cp->load_tgt(ctx_tgt, slot->id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - } - int32_t n_keep = n0; - if (ctx_dft != nullptr) { - llama_memory_seq_rm(llama_get_memory(ctx_dft), slot->id, n0, -1); - if (cp != nullptr) { - cp->load_dft(ctx_dft, slot->id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - } - } else { - llama_batch one = llama_batch_get_one(&toks[0], 1); - llama_pos pos = n0; - one.pos = &pos; - llama_seq_id sid = slot->id; - int32_t nsid = 1; - llama_seq_id * sids = &sid; - one.n_seq_id = &nsid; - one.seq_id = &sids; - if (llama_decode(ctx_tgt, one) == 0) { - n_keep = L.n_prefix; - } else { - llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot->id, n0, -1); - if (cp != nullptr) { - cp->load_tgt(ctx_tgt, slot->id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); - } - } - } - slot->prompt.tokens.keep_first(n_keep); - slot->prompt.checkpoints.remove_if([&](const common_prompt_checkpoint & c) { return c.n_tokens > n_keep; }); - res->ok = ok; - res->failure = ok ? "" : llama_opt_failure(ctx_tgt); - res->loss = loss; - res->n_labels = (int32_t) toks.size() - 1; - res->t_ms = (ggml_time_us() - t0) / 1000.0; + res->outcome = learn_on_slot(*slot, L, t0); learn_next_us = ggml_time_us() + (int64_t) ((ggml_time_us() - t0) * (1.0/LEARN_DUTY - 1.0)); - SRV_INF("learned on slot %d: %d labels from position %d, loss %.4f, %.1f ms%s\n", slot->id, res->n_labels, n0, loss, - res->t_ms, ok ? "" : (" (failed: " + res->failure + ")").c_str()); queue_results.send(std::move(res)); } break; case SERVER_TASK_TYPE_UNLOAD_LORA: @@ -4291,6 +4321,7 @@ struct server_context_impl { if (!process_token(result, slot)) { // release slot because of stop condition slot.print_timings(); + learn_if_asked(slot); send_final_response(slot); slot.release(); @@ -4412,6 +4443,7 @@ struct server_context_impl { if (!process_token(result, slot)) { slot.print_timings(); + learn_if_asked(slot); send_final_response(slot); slot.release(); diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 64b9251295ce..828eff7285d3 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -546,6 +546,22 @@ task_params eval_llama_cmpl_schema( f->eval(ctx, data); } + // learn a reply on this request's slot before it is released: {"reply": text, "lora_id", "lr"}; + // the reply is learned as the tokens the model would generate after this prompt (its text, then + // the end of the turn) + if (data.contains("learn")) { + const json & l = data.at("learn"); + if (!l.is_object() || !l.contains("reply") || !l.at("reply").is_string() || l.at("reply").get().empty()) { + throw std::invalid_argument("\"learn\" must be an object with a non-empty \"reply\" string"); + } + learn_request L; + L.lora_id = json_value(l, "lora_id", -1); + L.lr = json_value(l, "lr", 1e-4f); + L.reply = common_tokenize(vocab, l.at("reply").get(), false, true); + L.reply.push_back(llama_vocab_eot(vocab) != LLAMA_TOKEN_NULL ? llama_vocab_eot(vocab) : llama_vocab_eos(vocab)); + params.learn = std::move(L); + } + // post-processing { // if "reasoning_format" is not provided, its handler will not be called, we will need to handle it here diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 35c75c6d32fc..a46dd7a0ca24 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -359,6 +359,9 @@ json server_task_result_cmpl_final::to_json_non_oaicompat() { if (!stream && !probs_output.empty()) { res["completion_probabilities"] = completion_token_output::probs_vector_to_json(probs_output, post_sampling_probs); } + if (learned) { + res["learn"] = learned->to_json(); + } return response_fields.empty() ? res : json_get_nested_values(response_fields, res); } @@ -404,6 +407,9 @@ json server_task_result_cmpl_final::to_json_oaicompat() { if (verbose) { res["__verbose"] = to_json_non_oaicompat(); } + if (learned) { + res["learn"] = learned->to_json(); + } if (stats.is_set()) { res["timings"] = stats.to_json(); } @@ -452,6 +458,9 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat() { if (verbose) { res["__verbose"] = to_json_non_oaicompat(); } + if (learned) { + res["learn"] = learned->to_json(); + } if (stats.is_set()) { res["timings"] = stats.to_json(); } @@ -1693,7 +1702,7 @@ json server_task_result_lora_change::to_json() { return res; } -json server_task_result_learn::to_json() { +json learn_outcome::to_json() const { json res = { { "success", ok }, { "loss", loss }, { "n_labels", n_labels }, { "t_ms", t_ms } }; if (!ok) { res["failure"] = failure; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 150555c02ef6..9d64126d0830 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -3,6 +3,7 @@ #include "common.h" #include "llama.h" +#include #include #include #include @@ -50,6 +51,27 @@ enum stop_type { STOP_TYPE_LIMIT, }; +// One reply learned on the slot that served its prompt, in this serving context (llama_opt_learn): +// no second context, her prompt never decoded again. Asked by POST /learn, or carried by a +// completion request (task_params::learn), which learns on its own slot before releasing it. +struct learn_request { + int32_t id_slot = -1; + int32_t n_prefix = 0; // the served prompt's tokens, before the reply + std::vector reply; // what she should have said (her reply, or the reviewed correction) + std::vector prompt; // the served prompt, when known: her slot must hold exactly it + int32_t lora_id = -1; // the trainable adapter (the shadow), loaded dormant + float lr = 1e-4f; +}; + +struct learn_outcome { + bool ok = false; + std::string failure; + double loss = 0.0; + int32_t n_labels = 0; + double t_ms = 0.0; + json to_json() const; +}; + struct task_params { bool stream = false; bool include_usage = false; @@ -102,6 +124,10 @@ struct task_params { json format_logit_bias(const std::vector & logit_bias) const; json to_json(bool only_metrics = false) const; + + // learn a reply on this request's slot when it finishes, before the slot is released (no other + // request can take the slot in between): POST /learn's work, atomic with the turn + std::optional learn; }; // struct for tracking the state of a task (e.g., for streaming) @@ -186,16 +212,8 @@ struct server_task { // used by SERVER_TASK_TYPE_UNLOAD_LORA: the adapter id to retire int32_t unload_lora = -1; - // used by SERVER_TASK_TYPE_LEARN: one reply learned on the slot that served its prompt, in this - // serving context (llama_opt_learn): no second context, her prompt never decoded again - struct { - int32_t id_slot = -1; - int32_t n_prefix = 0; // the served prompt's tokens, before the reply - std::vector reply; // what she should have said (her reply, or the reviewed correction) - std::vector prompt; // the served prompt, when known: her slot must hold exactly it - int32_t lora_id = -1; // the trainable adapter (the shadow), loaded dormant - float lr = 1e-4f; - } learn; + // used by SERVER_TASK_TYPE_LEARN + learn_request learn; server_task() = default; @@ -340,6 +358,8 @@ struct completion_token_output { }; struct server_task_result_cmpl_final : server_task_result { + // the outcome of the request's learn (task_params::learn), run before its slot was released + std::optional learned; std::string content; llama_tokens tokens; @@ -595,12 +615,8 @@ struct server_task_result_lora_change : server_task_result { }; struct server_task_result_learn : server_task_result { - bool ok = false; - std::string failure; - double loss = 0.0; - int32_t n_labels = 0; - double t_ms = 0.0; - virtual json to_json() override; + learn_outcome outcome; + virtual json to_json() override { return outcome.to_json(); } }; struct server_prompt { From 0cb014e6be2ee27af3824fc2920bc16567b5f9e9 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 18:20:55 -0500 Subject: [PATCH 40/43] server: a learn the request could not run (an empty prompt) is reported in the final response, not omitted Cormac's review question on 9a1099343: every way a turn ends either learns or says why. Both finish paths (plain, speculative) learn before release; a cancelled turn has no reader; an errored turn returns its error; the empty-prompt final response now carries learn.success=false with the reason instead of omitting the field. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/server-context.cpp | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 04478e4b87b6..5be1bf182df5 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3572,6 +3572,11 @@ struct server_context_impl { SLT_WRN(slot, "%s", "empty prompt - releasing slot\n"); slot.print_timings(); + if (slot.task->params.learn) { + // a learn that could not run is said, not omitted (Cormac's review) + slot.learned = learn_outcome{}; + slot.learned->failure = "the request had no prompt to learn the reply after"; + } send_final_response(slot); slot.release(); From d5a4c772dc0d3bea3b9fdf3ed37378819347ac80 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 18:24:21 -0500 Subject: [PATCH 41/43] server: every final-response format carries the learn outcome; the Anthropic conversion keeps "learn" Cormac's review of continuum#5082: the learn outcome was serialized only by /completion, /v1/completions and /v1/chat/completions without streaming, while continuum's chat requests stream. It is now on the final chunk or event of every format: chat stream (the last chunk, beside timings), responses and responses stream (response.completed), Anthropic and Anthropic stream (message_delta, beside usage). Testing every format also found the Anthropic request conversion copied a fixed list of fields without "learn", so a learn there never ran; it now carries it. tools/server/tests/learn-every-format.py asserts exactly one successful learn outcome for each of the 10 endpoint x stream combinations. Measured (M5, Metal, qwen2.5-0.5b with Qwen3.8's template, one slot): all 10 OK (before the conversion fix: 8 OK, /v1/messages plain and stream carried none). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/server-chat.cpp | 4 ++- tools/server/server-task.cpp | 20 ++++++++++++ tools/server/tests/learn-every-format.py | 41 ++++++++++++++++++++++++ 3 files changed, 64 insertions(+), 1 deletion(-) create mode 100755 tools/server/tests/learn-every-format.py diff --git a/tools/server/server-chat.cpp b/tools/server/server-chat.cpp index a6fe3c6ba619..c0c8712c46c0 100644 --- a/tools/server/server-chat.cpp +++ b/tools/server/server-chat.cpp @@ -576,7 +576,9 @@ json server_chat_convert_anthropic_to_oai(const json & body) { } // Pass through common params - for (const auto & key : {"temperature", "top_p", "top_k", "stream", "chat_template_kwargs"}) { + // "learn" too: a learn on an Anthropic request was dropped here, so it never ran (a reply learned + // on the turn's own slot, task_params::learn) + for (const auto & key : {"temperature", "top_p", "top_k", "stream", "chat_template_kwargs", "learn"}) { if (body.contains(key)) { oai_body[key] = body.at(key); } diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index a46dd7a0ca24..e239b39f2706 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -526,6 +526,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat_stream() { deltas.back()["timings"] = stats.to_json(); } + if (learned) { + deltas.back()["learn"] = learned->to_json(); // the final chunk + } + // extra fields for debugging purposes if (verbose && !deltas.empty()) { deltas.front()["__verbose"] = to_json_non_oaicompat(); @@ -602,6 +606,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp() { }}, }; + if (learned) { + res["learn"] = learned->to_json(); + } + return res; } @@ -719,6 +727,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() { server_sent_events.back().at("data")["timings"] = stats.to_json(); } + if (learned) { + server_sent_events.back().at("data")["learn"] = learned->to_json(); // response.completed + } + return server_sent_events; } @@ -800,6 +812,10 @@ json server_task_result_cmpl_final::to_json_anthropic() { }} }; + if (learned) { + res["learn"] = learned->to_json(); + } + return res; } @@ -981,6 +997,10 @@ json server_task_result_cmpl_final::to_json_anthropic_stream() { }} }); + if (learned) { + events.back().at("data")["learn"] = learned->to_json(); // message_delta, beside the usage + } + events.push_back({ {"event", "message_stop"}, {"data", { diff --git a/tools/server/tests/learn-every-format.py b/tools/server/tests/learn-every-format.py new file mode 100755 index 000000000000..61d6e09e7d45 --- /dev/null +++ b/tools/server/tests/learn-every-format.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +# Every final-response format x stream mode carries exactly one learn outcome for a learn-carrying +# request. Run against a server started with --train-dir holding shadow-05.gguf (a fresh LoRA for the +# served model), e.g.: llama-server -m MODEL --port 58199 -np 1 --jinja --train-dir DIR +# what this catches: a serializer, or a request conversion (the Anthropic one dropped "learn"), that +# loses the outcome, so a client learns with nothing on the wire +import os +U = os.environ.get("LLAMA_SERVER_URL", "http://localhost:58199") +def raw(path, body): + req = urllib.request.Request(U + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=300) as r: return r.read().decode() +sid = json.loads(raw("/lora-adapters/load", {"name": "shadow-05.gguf"}))["id"] +learn = {"reply": "Topeka is the capital of Kansas.", "lora_id": sid, "lr": 0.001} +msgs = [{"role": "user", "content": "Which city is the capital of Kansas?"}] +cases = [("/completion", {"prompt": "Q: Which city is the capital of Kansas?\nA:"}), + ("/v1/completions", {"prompt": "Q: Which city is the capital of Kansas?\nA:"}), + ("/v1/chat/completions", {"messages": msgs, "chat_template_kwargs": {"enable_thinking": False}}), + ("/v1/responses", {"input": msgs, "chat_template_kwargs": {"enable_thinking": False}}), + ("/v1/messages", {"messages": msgs, "max_tokens": 2, "chat_template_kwargs": {"enable_thinking": False}})] +fails = 0 +for path, body in cases: + for stream in (False, True): + b = dict(body, stream=stream, learn=learn, temperature=0) + b.setdefault("max_tokens", 2); b["n_predict"] = 2 + text = raw(path, b) + found = [] + for line in (text.splitlines() if stream else [text]): + line = line[5:].strip() if line.startswith("data:") else line.strip() + if not line or line == "[DONE]" or not line.startswith("{"): continue + def walk(o): + if isinstance(o, dict): + if "learn" in o and isinstance(o["learn"], dict): found.append(o["learn"]) + for v in o.values(): walk(v) + elif isinstance(o, list): + for v in o: walk(v) + walk(json.loads(line)) + ok = len(found) == 1 and found[0].get("success") is True + fails += not ok + print(f"{'OK ' if ok else 'BAD'} {path:22s} stream={stream!s:5s} learn outcomes={len(found)} {found[0] if found else ''}") +print("ALL FORMATS CARRY THE LEARN OUTCOME" if fails == 0 else f"{fails} FORMAT(S) WITHOUT IT") +raise SystemExit(1 if fails else 0) From cc2e05ff2512b7784a92555f3518806b59945fcd Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Thu, 15 Oct 2026 18:24:51 -0500 Subject: [PATCH 42/43] server tests: learn-every-format.py imports what it uses (the committed copy dropped its import line) Measured: run against the server as documented in its header: ALL FORMATS CARRY THE LEARN OUTCOME, exit 0. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- tools/server/tests/learn-every-format.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/server/tests/learn-every-format.py b/tools/server/tests/learn-every-format.py index 61d6e09e7d45..fc247fd6d4a9 100755 --- a/tools/server/tests/learn-every-format.py +++ b/tools/server/tests/learn-every-format.py @@ -4,7 +4,7 @@ # served model), e.g.: llama-server -m MODEL --port 58199 -np 1 --jinja --train-dir DIR # what this catches: a serializer, or a request conversion (the Anthropic one dropped "learn"), that # loses the outcome, so a client learns with nothing on the wire -import os +import json, os, urllib.request U = os.environ.get("LLAMA_SERVER_URL", "http://localhost:58199") def raw(path, body): req = urllib.request.Request(U + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) From 35b8eed72f8b137d76428781765d85f9229964e3 Mon Sep 17 00:00:00 2001 From: joelteply Date: Thu, 15 Oct 2026 18:41:28 -0500 Subject: [PATCH 43/43] server: every training path refuses repacked CPU weights alike (continuum card 5b497852) /train refused a server that may hold base weights in a repacked CPU buffer; learn_on_slot (POST /learn and a learn on a completion) did not. Measured 2026-10-15 at cc2e05ff2: on x86 at -ngl 0 the learn failed in the backward ("no backend can run" OUT_PROD 'grad for ffn_swiglu-23' on blk.23.ffn_down.weight); on an M5 at -ngl 0 the scheduler copied the repacked weight to Metal and crashed in ggml_backend_tensor_copy (CPU_REPACK get_tensor/cpy_tensor are NULL). With --no-repack the learn succeeds. The guard is lifted to one function, train_repack_refusal, which /train and learn_on_slot both call; the learn returns learn.success=false with the reason and the turn is still served. Its comment no longer claims "no crash". Test (tools/server/tests/learn-refuses-repacked.py), IntelMac i7-8850H, CPU build, Qwen2.5-0.5B Q4_K_M, 23:40:58Z: repacked -> turn answered, learn refused with the --no-repack remedy, server alive; --no-repack -> learn success, loss 1.6083. The pre-change binary fails the test (its failure text names no remedy, 23:35:32Z). Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE --- tools/server/server-context.cpp | 5 ++ tools/server/server-train.cpp | 49 +++++++++++--------- tools/server/server-train.h | 6 +++ tools/server/tests/learn-refuses-repacked.py | 24 ++++++++++ 4 files changed, 63 insertions(+), 21 deletions(-) create mode 100644 tools/server/tests/learn-refuses-repacked.py diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5be1bf182df5..6972228d5c96 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -2072,6 +2072,11 @@ struct server_context_impl { SRV_WRN("learn refused: %s\n", why.c_str()); return out; }; + // the same refusal /train gives: repacked CPU weights fail the backward (x86) or crash + // the copy to Metal (an M5 at -ngl 0), card 5b497852 + if (const std::string why = train_repack_refusal(model_tgt, params_base); !why.empty()) { + return refused(why); + } auto & loras = params_base.lora_adapters; if (L.lora_id < 0 || (size_t) L.lora_id >= loras.size() || loras[L.lora_id].ptr == nullptr) { return refused("no loaded adapter " + std::to_string(L.lora_id) + " to train (POST /lora-adapters/load it first)"); diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index df43130d4c69..69f47d9b7d48 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -397,6 +397,32 @@ json server_trainer::measure(const json & body) const { return json::object({{"ok", true}, {"window", window}, {"examples", rows}}); } +// REPACKED WEIGHTS CANNOT BE TRAINED THROUGH (Cormac on the /train plan; card 5b497852): the +// backward reads quantized blocks in the standard layout, and a weight in an extra buffer type +// (CPU_REPACK) is not in it, and its buffer cannot even be read back (get_tensor/cpy_tensor NULL, +// ggml-cpu/repack.cpp). Measured 2026-10-15: on x86 a learn fails ("no backend can run" an +// OUT_PROD on blk.23.ffn_down.weight); on a Metal host at -ngl 0 the scheduler copies the weight to +// Metal and crashes in ggml_backend_tensor_copy. With --no-repack the same learn succeeds. Extra +// buffer types only apply to weights kept on the CPU, so training is allowed when they are +// disabled, or when every layer is offloaded and no tensor override pins a weight elsewhere. +std::string train_repack_refusal(const llama_model * model, const common_params & params) { + const int32_t n_layer = llama_model_n_layer(model); + const int32_t ngl = params.n_gpu_layers; + const bool all_gpu = ngl <= -2 || ngl > n_layer; + // the list carries a {nullptr, nullptr} terminator when parsed from args: count real entries + bool override = false; + for (const auto & o : params.tensor_buft_overrides) { + override = override || o.pattern != nullptr; + } + if (params.no_extra_bufts || (all_gpu && !override)) { + return ""; + } + return "this server may hold base weights in a repacked CPU buffer (n_gpu_layers " + std::to_string(ngl) + + " of " + std::to_string(n_layer) + " layers" + (override ? ", tensor overrides set" : "") + + "), and training's backward needs the standard layout. " + "Serve with every layer offloaded (-ngl all) or with --no-repack."; +} + json server_trainer::start(const json & body_in) { json body = body_in; examples_data prepared; @@ -466,27 +492,8 @@ json server_trainer::start(const json & body_in) { if (confine_out(params_base.train_dir, body.at("out").get()).empty()) { return json::object({{"ok", false}, {"error", "\"out\" must be a bare file name ending in .gguf (no directories, no \"..\"); it is written inside the server's --train-dir"}}); } - // REPACKED WEIGHTS GIVE WRONG GRADIENTS, SILENTLY (Cormac on the /train plan): the backward - // reads quantized blocks in the standard layout, and a weight in an extra buffer type - // (CPU_REPACK) is not in it โ€” no crash, garbage gradients. Extra buffer types only apply to - // weights kept on the CPU, so training is allowed when they are disabled, or when every layer - // is offloaded and no tensor override pins a weight elsewhere; otherwise refuse and say so. - { - const int32_t n_layer = llama_model_n_layer(model); - const int32_t ngl = params_base.n_gpu_layers; - const bool all_gpu = ngl <= -2 || ngl > n_layer; - // the list carries a {nullptr, nullptr} terminator when parsed from args: count real entries - bool override = false; - for (const auto & o : params_base.tensor_buft_overrides) { - override = override || o.pattern != nullptr; - } - if (!params_base.no_extra_bufts && (!all_gpu || override)) { - return json::object({{"ok", false}, {"error", - "refusing to train: this server may hold base weights in a repacked CPU buffer (n_gpu_layers " + - std::to_string(ngl) + " of " + std::to_string(n_layer) + " layers" + (override ? ", tensor overrides set" : "") + - "), and the backward reads the standard layout, so gradients would be silently wrong. " - "Serve with every layer offloaded (-ngl all) or with --no-repack."}}); - } + if (const std::string why = train_repack_refusal(model, params_base); !why.empty()) { + return json::object({{"ok", false}, {"error", "refusing to train: " + why}}); } if (body.contains("parse_special") && !body.at("parse_special").is_boolean()) { return json::object({{"ok", false}, {"error", "\"parse_special\" must be true or false"}}); diff --git a/tools/server/server-train.h b/tools/server/server-train.h index 34c282363940..58cff7bd3fb0 100644 --- a/tools/server/server-train.h +++ b/tools/server/server-train.h @@ -61,6 +61,12 @@ struct llama_context; // for every path the server writes or reads there: /train's "out" and /lora-adapters/load. std::string confine_out(const std::string & dir, const std::string & name); +// Why this server must not train on its served weights, or "" when it may: a quantized weight kept +// on the CPU sits in a repacked buffer (CPU_REPACK) whose layout only the CPU kernels read, and the +// backward needs the standard one. One check for every training path (/train, /learn, a learn on a +// completion), so they refuse alike instead of failing differently (continuum card 5b497852). +std::string train_repack_refusal(const llama_model * model, const common_params & params); + class server_trainer { public: // serving: read between training windows; the trainer yields while serving is busy, so a diff --git a/tools/server/tests/learn-refuses-repacked.py b/tools/server/tests/learn-refuses-repacked.py new file mode 100644 index 000000000000..6469c98f9bb1 --- /dev/null +++ b/tools/server/tests/learn-refuses-repacked.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python3 +# A learn-carrying completion on a server that may hold repacked CPU weights still answers its turn, +# and its learn outcome says why it did not learn. Run against a CPU server WITHOUT --no-repack and a +# --train-dir holding shadow-05.gguf, e.g.: +# llama-server -m MODEL --port 58199 -np 2 -c 8192 -ngl 0 --jinja --train-dir DIR +# what this catches (continuum card 5b497852): a learn on repacked weights failed in the backward on +# x86 and crashed the engine in ggml_backend_tensor_copy on a Metal host at -ngl 0; it must refuse +# the way /train does, before any training graph is built, and leave the turn served. +import json, os, urllib.request +U = os.environ.get("LLAMA_SERVER_URL", "http://localhost:58199") +def post(path, body): + req = urllib.request.Request(U + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=300) as r: return json.loads(r.read().decode()) +sid = post("/lora-adapters/load", {"name": "shadow-05.gguf"})["id"] +res = post("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Which city is the capital of Kansas?"}], + "max_tokens": 8, "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}, + "learn": {"reply": "Topeka is the capital of Kansas.", "lora_id": sid, "lr": 0.001}}) +answer = res["choices"][0]["message"]["content"] +learn = res.get("learn") or {} +ok = bool(answer) and learn.get("success") is False and "--no-repack" in learn.get("failure", "") +print(f"answer={answer!r} learn={learn}") +print("REFUSED WITH THE REASON, TURN SERVED" if ok else "WRONG") +raise SystemExit(0 if ok else 1)