From 999d682aac013aef843bcaa72fd2d804dffe226f Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Wed, 14 Oct 2026 14:41:25 -0500 Subject: [PATCH 1/5] 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 2/5] 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 3/5] 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 4/5] 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 5/5] 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);