From b93892157b9d6e51cda2d23576da9293102511cb Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Sun, 11 Oct 2026 11:13:23 -0500 Subject: [PATCH 1/6] train: tiled flash-attention backward, and the walk accumulates its K/V gradient in host memory - ggml: FLASH_ATTN_EXT gets a backward built from existing ops, tiled over the KV window (LSE, D, then grads per tile; GQA group-sum; mask floor -1e9). test-flash-attn-back checks it on every backend against an explicit F32 reference (max rel < 1e-4). - Metal: MUL_MAT honors GGML_PREC_F32 for F32 x F32 (fp32 staging), and the F32 FA kernel keeps its query in float; without both the walk's q8 cosine fell to 0.987. - ggml-opt: ggml_opt_accumulate_grad(tensor, acc): the backward adds a node's gradient into a caller-owned tensor; ggml_set_grad accepts any node. - The exact walk taps its cached prefix K/V and accumulates dL/d(cached K/V) straight into a host-memory accumulator. Deleted: the zero GRAD leaves, their device-resident gradients, the host copy-in of the surrogate inputs and the readback-and-add loop. Exactness (test-walk-exact, qwen2.5-0.5b, M5 Metal): exact walk vs one graph cosine 1.0000, f16 1.0000, q8_0 0.9905, recompute vs not 1.0000. Memory receipts go in the PR. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/include/ggml-opt.h | 8 + ggml/include/ggml.h | 14 +- ggml/src/ggml-metal/ggml-metal-ops.cpp | 11 +- ggml/src/ggml-metal/kernels/fa.metal | 12 +- ggml/src/ggml-opt.cpp | 16 +- ggml/src/ggml.c | 195 ++++++++++++++++++++++++- src/llama-context.cpp | 101 ++++++------- src/llama-cparams.h | 7 +- src/llama-graph.cpp | 81 +++++----- src/llama-graph.h | 9 +- tests/CMakeLists.txt | 2 + tests/test-flash-attn-back.cpp | 156 ++++++++++++++++++++ tests/test-walk-exact.cpp | 40 +++++ 13 files changed, 549 insertions(+), 103 deletions(-) create mode 100644 tests/test-flash-attn-back.cpp diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index c04ffba47f67..148349e99e20 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h @@ -230,6 +230,14 @@ extern "C" { // for a GRAD leaf (ggml_set_grad) of the evaluated graph, or NULL when the graph has none. GGML_API struct ggml_tensor * ggml_opt_leaf_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * leaf); + // For the graph last given to ggml_opt_prepare_alloc: its backward ADDS the gradient on `tensor` + // (a node of that graph, flagged GRAD by this call) into `acc`, an F32 tensor of the same shape + // in a buffer the caller owns and zeroes, instead of keeping it in the compute buffer for + // ggml_opt_leaf_grad. Nothing gradient-sized outlives its consumer on the device: the + // accumulator may live in host memory while the graph runs on a GPU. The optimizer never + // zeroes it, except at its own first allocation, which precedes any accumulation. + GGML_API void ggml_opt_accumulate_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, struct ggml_tensor * acc); + // ############################################################################ // ## The high-level functions start here. They do not depend on any private ## // ## functions or structs and can be copied to and adapted for user code. ## diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 619b82b88e62..1c23bbf98c85 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -654,7 +654,7 @@ extern "C" { GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed - GGML_TENSOR_FLAG_GRAD = 32, // ...needs its gradient computed, but no optimizer updates it (a leaf whose gradient is read back) + GGML_TENSOR_FLAG_GRAD = 32, // ...needs its gradient computed, but no optimizer updates it (read back, or accumulated into a caller's tensor) }; enum ggml_tri_type { @@ -2464,6 +2464,18 @@ extern "C" { GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( const struct ggml_tensor * a); + // the KV positions per tile of FLASH_ATTN_EXT's backward pass (0 = the default): the backward + // recomputes each tile's scores, so its memory is queries x tile, never queries x window + GGML_API void ggml_flash_attn_ext_set_back_tile( + struct ggml_tensor * a, + int32_t tile); + + // an upper bound on the graph nodes FLASH_ATTN_EXT's backward adds for n_kv KV positions at + // `tile` (0 = the default): a caller sizing a training graph counts them with its own + GGML_API size_t ggml_flash_attn_ext_back_n_nodes( + int64_t n_kv, + int32_t tile); + GGML_API void ggml_flash_attn_ext_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 2a71a2904d2b..cb2e266e12c9 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -2622,7 +2622,12 @@ int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { // default: break; //} - ggml_metal_encode_mul_mm(ctx, op); + // GGML_PREC_F32 on an F32 x F32 product is honored with the float-staged kernel (CUDA honors + // the same request): the default stages through half, ~1e-3 relative per product, which the + // training backward cannot afford (flash attention's tiled backward, continuum 2026-10-11) + const bool fp32 = !props_dev->has_tensor && ggml_get_op_params_i32(op, 0) == GGML_PREC_F32 && + op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32; + ggml_metal_encode_mul_mm(ctx, op, fp32); } else { auto pipeline = ggml_metal_library_get_pipeline_mul_mv(lib, op); @@ -3412,7 +3417,9 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { // the shared memory needed for the simdgroups to load the KV cache // each thread loads (dequantizes) 16 head elements, there are 32 threads in th SG // -#define FATTN_SMEM(nsg) (GGML_PAD((nqptg*(ne00 + 2*GGML_PAD(ne20, 64) + 2*(2*ncpsg)) + is_q*(16*32*(nsg)))*(sizeof(float)/2), 16)) + // the F32 kernels hold the query as float (kernels/fa.metal, QH): twice the halves per element + const int qh = !use_kv_f16 && op->src[1]->type == GGML_TYPE_F32 ? 2 : 1; +#define FATTN_SMEM(nsg) (GGML_PAD((nqptg*(qh*ne00 + 2*GGML_PAD(ne20, 64) + 2*(2*ncpsg)) + is_q*(16*32*(nsg)))*(sizeof(float)/2), 16)) //int64_t nsgmax = 4; // diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index e95dec258a37..6faf923c425e 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -273,12 +273,16 @@ void kernel_flash_attn_ext_impl( constexpr short SH = 2*C; // shared memory per simdgroup (s_t == float) constexpr short TS = 2*SH; - constexpr short T = DK + 2*PV; // shared memory size per query in (half) + // halves per query element: 1 for a half query, 2 for the F32 kernels' float query (Continuum: + // a training graph reads attention through the F32 kernels, and a query rounded to half moved + // its step, test-walk-exact on the M5 2026-10-11; the host sizes the memory the same way) + constexpr short QH = sizeof(q_t)/sizeof(half); + constexpr short T = QH*DK + 2*PV; // shared memory size per query in (half) threadgroup q_t * sq = (threadgroup q_t *) (shmem_f16 + 0*T); // holds the query data threadgroup q4_t * sq4 = (threadgroup q4_t *) (shmem_f16 + 0*T); // same as above but in q4_t - threadgroup o_t * so = (threadgroup o_t *) (shmem_f16 + 0*T + Q*DK); // the result for all queries in 8x8 matrices (the O matrix from the paper) - threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 0*T + Q*DK); + threadgroup o_t * so = (threadgroup o_t *) (shmem_f16 + 0*T + Q*QH*DK); // the result for all queries in 8x8 matrices (the O matrix from the paper) + threadgroup o4_t * so4 = (threadgroup o4_t *) (shmem_f16 + 0*T + Q*QH*DK); threadgroup s_t * ss = (threadgroup s_t *) (shmem_f16 + Q*T); // scratch buffer for attention, mask and diagonal matrix threadgroup s2_t * ss2 = (threadgroup s2_t *) (shmem_f16 + Q*T); // same as above but in s2_t @@ -913,7 +917,7 @@ kernel void kernel_flash_attn_ext( //float, float4, simdgroup_float8x8 #define FA_TYPES_F32 \ - half, half4, simdgroup_half8x8, \ + float, float4, simdgroup_float8x8, \ float, float4x4, simdgroup_float8x8, \ float, float4x4, simdgroup_float8x8, \ float, simdgroup_float8x8, \ diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 6379f996f4f3..5a604762ef83 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -86,6 +86,9 @@ struct ggml_opt_context { bool period_fresh = true; // the next backward starts a period (zeroed gradients) // the gradients of the evaluated graph's GRAD leaves, readable until the next alloc std::unordered_map leaf_grads; + // ggml_opt_accumulate_grad: GRAD nodes whose gradient adds into a caller's tensor, for the + // graph prepared last (ggml_opt_prepare_alloc clears them) + std::unordered_map next_grad_accs; std::vector grad_accs; int64_t iter = 1; @@ -677,6 +680,8 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { if (accumulate) { opt_ctx->grad_accs[i] = opt_ctx->param_grad_acc.at(node); } + } else if (const auto it = opt_ctx->next_grad_accs.find(node); it != opt_ctx->next_grad_accs.end()) { + opt_ctx->grad_accs[i] = it->second; } } @@ -686,7 +691,7 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { opt_ctx->leaf_grads.clear(); for (int i = 0; i < opt_ctx->gf->n_nodes; ++i) { ggml_tensor * node = opt_ctx->gf->nodes[i]; - if (node->flags & GGML_TENSOR_FLAG_GRAD) { + if ((node->flags & GGML_TENSOR_FLAG_GRAD) && !opt_ctx->next_grad_accs.count(node)) { ggml_tensor * grad = ggml_graph_get_grad(opt_ctx->gb_grad, node); if (grad) { ggml_set_output(grad); // read back after the step: the allocator must not reuse it @@ -811,6 +816,14 @@ void ggml_opt_set_next_step(ggml_opt_context_t opt_ctx, bool period_end, float l opt_ctx->next_extra_loss = extra_loss; } +void ggml_opt_accumulate_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * tensor, struct ggml_tensor * acc) { + GGML_ASSERT(!opt_ctx->eval_ready && "register an accumulator before ggml_opt_alloc"); + GGML_ASSERT(acc->type == GGML_TYPE_F32 && ggml_are_same_shape(tensor, acc)); + GGML_ASSERT((acc->view_src ? acc->view_src : acc)->buffer != nullptr && "the accumulator lives in the caller's buffer"); + ggml_set_grad(tensor); + opt_ctx->next_grad_accs[tensor] = acc; +} + struct ggml_tensor * ggml_opt_leaf_grad(ggml_opt_context_t opt_ctx, struct ggml_tensor * leaf) { const auto it = opt_ctx->leaf_grads.find(leaf); return it == opt_ctx->leaf_grads.end() ? nullptr : it->second; @@ -950,6 +963,7 @@ void ggml_opt_prepare_alloc( opt_ctx->gf = gf; opt_ctx->inputs = inputs; opt_ctx->outputs = outputs; + opt_ctx->next_grad_accs.clear(); // registrations belong to the graph prepared last } void ggml_opt_set_alloc_budget(ggml_opt_context_t opt_ctx, size_t bytes) { diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 8f64ff6399bb..b00ebff70d61 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5517,6 +5517,13 @@ enum ggml_prec ggml_flash_attn_ext_get_prec( return (enum ggml_prec) prec_i32; } +void ggml_flash_attn_ext_set_back_tile( + struct ggml_tensor * a, + int32_t tile) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + ggml_set_op_params_i32(a, 4, tile); // scale, max_bias, softcap, prec come first +} + void ggml_flash_attn_ext_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks) { @@ -6711,6 +6718,188 @@ static void ggml_sub_or_set( ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); } +// The backward of FLASH_ATTN_EXT, built from ops every backend already has, tiled over the KV +// positions so its memory is chunk x tile, never chunk x window (continuum 2026-10-11: training's +// explicit attention held ~4 F32 copies of window x chunk x heads, 75% of a training window's +// bytes, so training fit a quarter of the window serving runs). Per tile it RECOMPUTES the scores +// instead of keeping them, in three passes over the tiles: each row's log-sum-exp; then +// D_i = sum_j p_ij dp_ij; then the gradients. No row-max op is needed: for p = softmax(s), +// log sum exp(s) = sum p * (s - log p), which soft_max computes stably, and the tiles' values +// combine by the same identity. Masked (-inf) scores are clamped to a finite floor first, so a +// row a tile masks entirely contributes exp(floor) = 0 instead of 0/0. +#define GGML_FA_BACK_TILE_DEFAULT 4096 +#define GGML_FA_BACK_FLOOR -1.0e9f +#define GGML_FA_BACK_EPS 1.0e-30f + +// per tile: the scores (8 ops) recomputed in each of three passes, the log-sum-exp (6), D (5) and the +// gradients (~25 with the transposes and the group sums), padded; plus the passes' fixed ops +#define GGML_FA_BACK_NODES_PER_TILE 96 +#define GGML_FA_BACK_NODES_FIXED 32 + +size_t ggml_flash_attn_ext_back_n_nodes(int64_t n_kv, int32_t tile) { + if (tile <= 0) { + tile = GGML_FA_BACK_TILE_DEFAULT; + } + const int64_t nt = (n_kv + tile - 1) / tile; + return (size_t) (nt * GGML_FA_BACK_NODES_PER_TILE + GGML_FA_BACK_NODES_FIXED); +} + +// a K or V tile at F32: a mixed product (F16 x F32) converts the F32 side to F16 before the dot, +// so an F16 cache would round the queries and gradients the backward multiplies them with. Per +// tile, so the copy is tile-sized, never window-sized. +static struct ggml_tensor * ggml_fa_back_f32(struct ggml_context * ctx, struct ggml_tensor * t) { + return t->type == GGML_TYPE_F32 ? t : ggml_cast(ctx, t, GGML_TYPE_F32); +} + +// every product in the backward asks for F32 precision: a backend that stages through half would +// round the scores and the gradients ~1e-3 per product (Metal's default mul_mm) +static struct ggml_tensor * ggml_fa_back_mm(struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b) { + struct ggml_tensor * r = ggml_mul_mat(ctx, a, b); + ggml_mul_mat_set_prec(r, GGML_PREC_F32); + return r; +} + +// the scores of KV positions [j0, j0+n) for every query row, as [n, N, H] F32: scaled, masked, +// and clamped to a finite floor so a fully masked row stays finite +static struct ggml_tensor * ggml_fa_back_scores(struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, + struct ggml_tensor * mask, float scale, int64_t j0, int64_t n) { + struct ggml_tensor * kt = ggml_fa_back_f32(ctx, ggml_view_3d(ctx, k, k->ne[0], n, k->ne[2], k->nb[1], k->nb[2], j0*k->nb[1])); + struct ggml_tensor * s = ggml_fa_back_mm(ctx, kt, q); // [n, N, H] + s = ggml_scale(ctx, s, scale); + if (mask) { + struct ggml_tensor * m = ggml_view_4d(ctx, mask, n, q->ne[1], mask->ne[2], mask->ne[3], + mask->nb[1], mask->nb[2], mask->nb[3], j0*mask->nb[0]); + s = ggml_add(ctx, s, ggml_cast(ctx, m, GGML_TYPE_F32)); + } + return ggml_clamp(ctx, s, GGML_FA_BACK_FLOOR, INFINITY); +} + +// log-sum-exp over rows (ne0) of s, by the softmax identity above +static struct ggml_tensor * ggml_fa_back_lse(struct ggml_context * ctx, struct ggml_tensor * s) { + struct ggml_tensor * p = ggml_soft_max(ctx, s); + struct ggml_tensor * lp = ggml_log(ctx, ggml_scale_bias(ctx, p, 1.0f, GGML_FA_BACK_EPS)); + return ggml_sum_rows(ctx, ggml_mul(ctx, p, ggml_sub(ctx, s, lp))); +} + +// heads [D, T, H] -> their KV heads [D, T, Hkv], summing each KV head's group of query heads +// (the broadcast mul_mat uses: query head h reads KV head h / (H / Hkv)) +static struct ggml_tensor * ggml_fa_back_group_sum(struct ggml_context * ctx, struct ggml_tensor * t, int64_t n_head_kv) { + if (t->ne[2] == n_head_kv) { + return t; + } + const int64_t g = t->ne[2] / n_head_kv; + struct ggml_tensor * v = ggml_view_4d(ctx, t, t->ne[0], t->ne[1], n_head_kv, g, t->nb[1], t->nb[2]*g, t->nb[2], 0); + return ggml_repeat_back(ctx, v, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, t->ne[0], t->ne[1], n_head_kv, 1)); +} + +// tiles joined along ne1 by halves, so a long window copies O(n log n), not O(n^2) +static struct ggml_tensor * ggml_fa_back_join(struct ggml_context * ctx, struct ggml_tensor ** parts, int n) { + if (n == 1) { + return parts[0]; + } + const int h = n / 2; + return ggml_concat(ctx, ggml_fa_back_join(ctx, parts, h), ggml_fa_back_join(ctx, parts + h, n - h), 1); +} + +static void ggml_flash_attn_ext_backward( + struct ggml_context * ctx, struct ggml_cgraph * cgraph, struct ggml_tensor * tensor, struct ggml_tensor * grad, + size_t isrc0, size_t isrc1, size_t isrc2, bool dq_needed, bool dk_needed, bool dv_needed) { + struct ggml_tensor * q = tensor->src[0]; // [D, N, H, 1] + struct ggml_tensor * k = tensor->src[1]; // [D, M, Hkv, 1] + struct ggml_tensor * v = tensor->src[2]; // [Dv, M, Hkv, 1] + struct ggml_tensor * mask = tensor->src[3]; // [M', N', 1|H, 1] F16, or NULL + + float scale, max_bias, softcap; + memcpy(&scale, (const float *) tensor->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) tensor->op_params + 1, sizeof(float)); + memcpy(&softcap, (const float *) tensor->op_params + 2, sizeof(float)); + GGML_ASSERT(max_bias == 0.0f && "flash_attn_ext backward: ALiBi is not implemented"); + GGML_ASSERT(softcap == 0.0f && "flash_attn_ext backward: logit softcap is not implemented"); + GGML_ASSERT(tensor->src[4] == NULL && "flash_attn_ext backward: sinks are not implemented"); + GGML_ASSERT(q->ne[3] == 1 && "flash_attn_ext backward: one sequence stream"); + + int32_t tile = ggml_get_op_params_i32(tensor, 4); + if (tile <= 0) { + tile = GGML_FA_BACK_TILE_DEFAULT; + } + const int64_t M = k->ne[1]; + const int64_t Hkv = k->ne[2]; + const int nt = (int) ((M + tile - 1) / tile); + + // the output's gradient as [Dv, N, H] + struct ggml_tensor * dO = ggml_cont(ctx, ggml_permute(ctx, grad, 0, 2, 1, 3)); + + // pass 1 of 3: each row's log-sum-exp over the whole window + struct ggml_tensor * lse = NULL; + { + struct ggml_tensor * lz = NULL; + for (int t = 0; t < nt; ++t) { + const int64_t j0 = (int64_t) t * tile; + const int64_t n = MIN(tile, M - j0); + struct ggml_tensor * l = ggml_fa_back_lse(ctx, ggml_fa_back_scores(ctx, q, k, mask, scale, j0, n)); // [1, N, H] + lz = lz ? ggml_concat(ctx, lz, l, 0) : l; + } + lse = nt == 1 ? lz : ggml_fa_back_lse(ctx, lz); + } + + // pass 2 of 3: D_i = sum_j p_ij dp_ij, [1, N, H], from the SAME F32 recomputation the gradients use, never + // from the forward's output (D_i = dO_i . O_i only when O is exactly sum_j p_ij v_j): a + // flash-attention forward that accumulates at reduced precision (Metal, CUDA's tensor-core + // kernels) would otherwise put its rounding into every ds_ij, inconsistently with the p_ij + // recomputed here (test-walk-exact on the M5, 2026-10-11). + struct ggml_tensor * dvec = NULL; + for (int t = 0; t < nt; ++t) { + const int64_t j0 = (int64_t) t * tile; + const int64_t n = MIN(tile, M - j0); + struct ggml_tensor * vt = ggml_fa_back_f32(ctx, ggml_view_3d(ctx, v, v->ne[0], n, Hkv, v->nb[1], v->nb[2], j0*v->nb[1])); + struct ggml_tensor * p = ggml_exp(ctx, ggml_sub(ctx, ggml_fa_back_scores(ctx, q, k, mask, scale, j0, n), lse)); + struct ggml_tensor * d = ggml_sum_rows(ctx, ggml_mul(ctx, p, ggml_fa_back_mm(ctx, vt, dO))); + dvec = dvec ? ggml_add(ctx, dvec, d) : d; + } + + // transposed once, for the matmuls below ([N, D, H]): every backend has a batched, + // head-broadcasting mul_mat, while a batched out_prod is not universal + struct ggml_tensor * qT = dk_needed ? ggml_cont(ctx, ggml_transpose(ctx, q)) : NULL; + struct ggml_tensor * dOT = dv_needed ? ggml_cont(ctx, ggml_transpose(ctx, dO)) : NULL; + + // pass 3 of 3: the gradients, tile by tile + struct ggml_tensor * dq = NULL; + struct ggml_tensor ** dk = (struct ggml_tensor **) alloca(sizeof(struct ggml_tensor *) * nt); + struct ggml_tensor ** dv = (struct ggml_tensor **) alloca(sizeof(struct ggml_tensor *) * nt); + for (int t = 0; t < nt; ++t) { + const int64_t j0 = (int64_t) t * tile; + const int64_t n = MIN(tile, M - j0); + struct ggml_tensor * kt = ggml_fa_back_f32(ctx, ggml_view_3d(ctx, k, k->ne[0], n, Hkv, k->nb[1], k->nb[2], j0*k->nb[1])); + struct ggml_tensor * vt = ggml_fa_back_f32(ctx, ggml_view_3d(ctx, v, v->ne[0], n, Hkv, v->nb[1], v->nb[2], j0*v->nb[1])); + struct ggml_tensor * p = ggml_exp(ctx, ggml_sub(ctx, ggml_fa_back_scores(ctx, q, k, mask, scale, j0, n), lse)); // [n, N, H] + struct ggml_tensor * dp = ggml_fa_back_mm(ctx, vt, dO); // [n, N, H] + struct ggml_tensor * ds = ggml_scale(ctx, ggml_mul(ctx, p, ggml_sub(ctx, dp, dvec)), scale); // [n, N, H] + if (dq_needed) { + // dq[d,i] = sum_j k[d,j] ds[j,i] + struct ggml_tensor * g = ggml_fa_back_mm(ctx, ggml_cont(ctx, ggml_transpose(ctx, kt)), ds); // [D, N, H] + dq = dq ? ggml_add(ctx, dq, g) : g; + } + if (dk_needed) { + // dk[d,j] = sum_i q[d,i] ds[j,i] + dk[t] = ggml_fa_back_group_sum(ctx, ggml_fa_back_mm(ctx, qT, ggml_cont(ctx, ggml_transpose(ctx, ds))), Hkv); // [D, n, Hkv] + } + if (dv_needed) { + // dv[d,j] = sum_i dO[d,i] p[j,i] + dv[t] = ggml_fa_back_group_sum(ctx, ggml_fa_back_mm(ctx, dOT, ggml_cont(ctx, ggml_transpose(ctx, p))), Hkv); // [Dv, n, Hkv] + } + } + + if (dq_needed) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape_4d(ctx, ggml_cont(ctx, dq), q->ne[0], q->ne[1], q->ne[2], q->ne[3])); + } + if (dk_needed) { + ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape_4d(ctx, ggml_fa_back_join(ctx, dk, nt), k->ne[0], k->ne[1], k->ne[2], k->ne[3])); + } + if (dv_needed) { + ggml_add_or_set(ctx, cgraph, isrc2, ggml_reshape_4d(ctx, ggml_fa_back_join(ctx, dv, nt), v->ne[0], v->ne[1], v->ne[2], v->ne[3])); + } +} + static void ggml_compute_backward( struct ggml_context * ctx, struct ggml_cgraph * cgraph, int i, const bool * grads_needed) { struct ggml_tensor * tensor = cgraph->nodes[i]; @@ -7295,6 +7484,10 @@ static void ggml_compute_backward( } //break; } } break; + case GGML_OP_FLASH_ATTN_EXT: { + ggml_flash_attn_ext_backward(ctx, cgraph, tensor, grad, isrc0, isrc1, isrc2, + src0_needs_grads, src1_needs_grads, src2_needs_grads); + } break; case GGML_OP_NONE: { // noop } break; @@ -8083,7 +8276,7 @@ void ggml_set_param(struct ggml_tensor * tensor) { } void ggml_set_grad(struct ggml_tensor * tensor) { - GGML_ASSERT(tensor->op == GGML_OP_NONE); + // any tensor: a leaf whose gradient is read back, or a node tapped for the gradient on its value tensor->flags |= GGML_TENSOR_FLAG_GRAD; } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 61ae0cb40ec6..18d094873c17 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3348,15 +3348,19 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params // walks the window decoding context and training each labelled chunk. One chunk is one // batch is one ubatch: a chunk's backward needs all of its own K/V in one graph. GGML_ASSERT(n_ubatch == n_batch && "training needs one ubatch per batch: set -ub = -b"); - // FLASH_ATTN_EXT has no backward: training graphs build explicit attention whatever the - // context's flag (build_attn_mha checks cparams.training); flash attention, when the context - // has it, serves the walk's context decodes and lets the cache hold V un-transposed. + // In a flash-attention context the training graph reads attention through FLASH_ATTN_EXT too + // (build_attn_mha), whose backward is tiled over the window; flash attention also serves the + // walk's context decodes and lets the cache hold V un-transposed. cparams.training = true; // The graph result and the scheduler were sized at context creation: before training // (no backward pass) and before any adapter attached since (its nodes uncounted). - // Size both again for the training graph, from the same number. + // Size both again for the training graph, from the same number, plus the nodes each + // attention layer's tiled backward adds for this window (one count, owned by ggml). { - const size_t max_nodes = graph_max_nodes(n_ubatch); + size_t max_nodes = graph_max_nodes(n_ubatch); + if (cparams.flash_attn) { + max_nodes += (size_t) model->hparams.n_layer() * ggml_flash_attn_ext_back_n_nodes(opt_n_ctx_train, 0); + } gf_res_prev.reset(new llm_graph_result(max_nodes)); sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), max_nodes, cparams.pipeline_parallel, cparams.op_offload)); @@ -3504,15 +3508,23 @@ void llama_context::opt_epoch_iter( uint32_t c0; uint32_t n; uint32_t n_labels; - uint32_t grad_from; // its GRAD leaves cover cached positions [grad_from, c0) + uint32_t grad_from; // its gradient taps cover cached positions [grad_from, c0) bool surrogate; // a later chunk's gradient sits on its own K/V bool needed; bool period_end; // the window's optimizer step float loss_scale; // its labelled positions over the window's }; const exact_chunk * xc = nullptr; - // per layer, per position: dL/d(cached K) and dL/d(cached V), [position][n_embd_*_gqa] - std::vector> walk_gk, walk_gv; + // per layer: dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, position] F32 + // in host memory (null for a recurrent layer); every chunk's backward adds into it in place. + // The graph reads it through cparams, which never outlives it: whatever way the walk returns. + struct walk_acc_t { + llama_cparams & cp; + std::vector k, v; + ggml_context_ptr ctx; + ggml_backend_buffer_ptr buf; + ~walk_acc_t() { cp.walk_acc_k = nullptr; cp.walk_acc_v = nullptr; } + } walk_acc { cparams, {}, {}, nullptr, nullptr }; // per recurrent state tensor: dL/d(the state the chunk just before leaves), from the chunk // trained last (each chunk's exit state feeds exactly one chunk, the next) std::unordered_map> walk_gstate; @@ -3604,6 +3616,12 @@ void llama_context::opt_epoch_iter( ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); if (xc != nullptr) { ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); + for (const auto & io : res->t_walk) { + if (io.tk) { + ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); + ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); + } + } } if (!ggml_opt_alloc(opt_ctx, train)) { const char * why = ggml_opt_refusal(opt_ctx); @@ -3640,22 +3658,6 @@ void llama_context::opt_epoch_iter( } GGML_ASSERT(row == n_outputs); } - if (xc != nullptr) { - // the GRAD leaves start at zero; the surrogate inputs are what later chunks - // accumulated on this chunk's own K/V - for (const auto & io : res->t_walk) { - if (io.dk) { - ggml_backend_tensor_memset(io.dk, 0, 0, ggml_nbytes(io.dk)); - ggml_backend_tensor_memset(io.dv, 0, 0, ggml_nbytes(io.dv)); - } - if (io.gk) { - const size_t kd = io.gk->ne[0]*io.gk->ne[1]; - const size_t vd = io.gv->ne[0]*io.gv->ne[1]; - ggml_backend_tensor_set(io.gk, walk_gk[io.il].data() + (size_t) pos_ctx*kd, 0, ggml_nbytes(io.gk)); - ggml_backend_tensor_set(io.gv, walk_gv[io.il].data() + (size_t) pos_ctx*vd, 0, ggml_nbytes(io.gv)); - } - } - } if (xc != nullptr) { for (const auto & ws : res->t_walk_state) { if (ws.ds) { @@ -3678,26 +3680,6 @@ void llama_context::opt_epoch_iter( opt_stop_requested.store(true); return false; } - if (xc != nullptr) { - // dL/d(cached K/V) of this chunk's prefix joins what later chunks put there - std::vector g; - for (const auto & io : res->t_walk) { - if (!io.dk) { - continue; - } - for (int kv = 0; kv < 2; ++kv) { - ggml_tensor * leaf = kv == 0 ? io.dk : io.dv; - ggml_tensor * grad = ggml_opt_leaf_grad(opt_ctx, leaf); - GGML_ASSERT(grad != nullptr && grad->type == GGML_TYPE_F32 && ggml_is_contiguous(grad)); - g.resize(ggml_nelements(grad)); - ggml_backend_tensor_get(grad, g.data(), 0, ggml_nbytes(grad)); - float * acc = (kv == 0 ? walk_gk : walk_gv)[io.il].data() + (size_t) io.grad_from*leaf->ne[0]*leaf->ne[1]; - for (size_t e = 0; e < g.size(); ++e) { - acc[e] += g[e]; - } - } - } - } if (xc != nullptr) { // dL/d(entry state): what the previous chunk's exit state takes in its surrogate for (const auto & ws : res->t_walk_state) { @@ -3853,14 +3835,29 @@ void llama_context::opt_epoch_iter( __func__, stride, opt_walk_host_budget/1048576.0, bytes/1048576.0); } } - walk_gk.assign(model.hparams.n_layer(), {}); - walk_gv.assign(model.hparams.n_layer(), {}); - for (uint32_t il = 0; il < model.hparams.n_layer(); ++il) { - if (model.hparams.is_recr(il)) { - continue; // a recurrent layer has no K/V: its state carries the chain + { + const uint32_t n_layer = model.hparams.n_layer(); + walk_acc.k.assign(n_layer, nullptr); + walk_acc.v.assign(n_layer, nullptr); + walk_acc.ctx.reset(ggml_init({ 2*n_layer*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true })); + for (uint32_t il = 0; il < n_layer; ++il) { + if (model.hparams.is_recr(il)) { + continue; // a recurrent layer has no K/V: its state carries the chain + } + walk_acc.k[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_k_gqa(il), last_label + 1); + walk_acc.v[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_v_gqa(il), last_label + 1); + } + walk_acc.buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(walk_acc.ctx.get(), ggml_backend_cpu_buffer_type())); + if (!walk_acc.buf) { + opt_failure = "the exact walk's gradient accumulator (host memory) could not be allocated"; + LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); + opt_alloc_failed.store(true); + opt_stop_requested.store(true); + return; } - walk_gk[il].assign((size_t) (last_label + 1)*model.hparams.n_embd_k_gqa(il), 0.0f); - walk_gv[il].assign((size_t) (last_label + 1)*model.hparams.n_embd_v_gqa(il), 0.0f); + ggml_backend_buffer_clear(walk_acc.buf.get(), 0); + cparams.walk_acc_k = walk_acc.k.data(); + cparams.walk_acc_v = walk_acc.v.data(); } // the window as the adapter now reads it, chunk by chunk: a recurrent model's state is @@ -4008,8 +4005,6 @@ void llama_context::opt_epoch_iter( cparams.walk_surrogate = false; cparams.walk_state_surrogate = false; } - walk_gk.clear(); - walk_gv.clear(); walk_gstate.clear(); return; } diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5df493dec5da..f32f7f465b6a 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -15,12 +15,17 @@ struct llama_cparams { // THE EXACT WALK (llama_context::opt_epoch_iter, walk_exact): a chunk's training graph in // the reverse pass. walk_grad_from: the first cached prefix position whose K/V carry a - // gradient (positions [walk_grad_from, n_past) attend through a GRAD leaf; [0, + // gradient (positions [walk_grad_from, n_past) attend through a gradient tap; [0, // walk_grad_from) stay constants). walk_surrogate: later chunks accumulated a gradient on // this chunk's own K/V, so its graph adds the surrogate term that carries it in. + // walk_acc_k / walk_acc_v: per layer, dL/d(cached K / V) of the window, [n_embd_*_gqa, + // window] F32 in host memory (null for a recurrent layer); a chunk's backward adds its prefix's + // share into rows [walk_grad_from, n_past) and its surrogate reads its own rows. bool walk_exact = false; uint32_t walk_grad_from = 0; bool walk_surrogate = false; + ggml_tensor * const * walk_acc_k = nullptr; + ggml_tensor * const * walk_acc_v = nullptr; // a later chunk's gradient sits on the recurrent state this chunk leaves bool walk_state_surrogate = false; uint32_t n_ctx_seq; // context for a single sequence diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 698ec24bdaa7..30bb874bdaf5 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2578,10 +2578,15 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * cur; - // A training graph takes the explicit path, which has a backward (FLASH_ATTN_EXT has none), - // even in a context created with flash attention: there the walk's context decodes use - // flash attention and the cache stores V un-transposed, so it can be quantized. - const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr && !cparams.training; + // A training graph reads attention through FLASH_ATTN_EXT too, whose backward is tiled over + // the KV positions and recomputes each tile's scores (ggml_flash_attn_ext_backward): memory + // queries x tile instead of the explicit path's ~4 F32 copies of queries x window x heads, + // which were ~75% of a 27B's training graph (continuum, the 5090, 2026-10-11). The backward + // covers plain scaled attention; a model with ALiBi, a logit softcap, sinks or MLA trains on + // the explicit path, which has a backward for all of them. + const bool fa_back_ok = !cparams.training || + (hparams.f_max_alibi_bias == 0.0f && !hparams.attn_soft_cap && sinks == nullptr && v_mla == nullptr); + const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr && fa_back_ok; if (use_flash_attn) { GGML_ASSERT(kq_b == nullptr && "Flash attention does not support KQ bias yet"); @@ -2589,12 +2594,15 @@ 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) - if (k->type == GGML_TYPE_F32) { + // 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) { k = ggml_cast(ctx0, k, GGML_TYPE_F16); } - if (v->type == GGML_TYPE_F32) { + if (!cparams.training && v->type == GGML_TYPE_F32) { v = ggml_cast(ctx0, v, GGML_TYPE_F16); } @@ -2604,6 +2612,9 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_flash_attn_ext_add_sinks(cur, sinks); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); + if (cparams.training) { + ggml_flash_attn_ext_set_back_tile(cur, 0); // the backward's default tile + } if (v_mla) { #if 0 @@ -2875,18 +2886,27 @@ ggml_tensor * llm_graph_context::build_attn( }; if (cparams.walk_exact) { // THE EXACT WALK's chunk graph. The prefix positions [grad_from, n_past) are the - // cached values PLUS a zero GRAD leaf, so the backward yields dL/d(cached K/V) for - // them without copying the cache; earlier positions stay constants. The chunk's own - // K/V meet the gradient later chunks put on them through the surrogate - // + : its gradient on k_cur / v_cur IS gk / gv, which the - // backward carries through this chunk into the adapter and into its own prefix. + // cached values as this chunk's (F32) K/V, TAPPED: the backward adds dL/d(cached K/V) for them straight + // into the window's accumulator (cparams.walk_acc_k / _v, host memory), so nothing + // window-sized outlives its layer on the device; earlier positions stay constants. + // The chunk's own K/V meet the gradient later chunks put on them through the surrogate + // + , gk / gv being the accumulator's rows for this chunk: its + // gradient on k_cur / v_cur IS gk / gv, which the backward carries through this chunk + // into the adapter and into its own prefix. // At a quantized or f16 cache the prefix's gradient is taken at the cache's values and // applied to the chunk's own full-precision K/V: straight-through. Measured // (test-walk-exact, q8_0 and f16 caches): cosine 0.998 against one graph, the same as // at an F32 cache, so training reads the numbers serving runs on (Joel: align the bit // depth to inference). const int64_t g0 = std::min(cparams.walk_grad_from, n_past); - llm_graph_result::walk_layer io = { il, (uint32_t) g0, (uint32_t) n_past }; + llm_graph_result::walk_layer io = { il, (uint32_t) n_past }; + ggml_tensor * acc_k = cparams.walk_acc_k[il]; + ggml_tensor * acc_v = cparams.walk_acc_v[il]; + // rows [p0, p0 + n) of a layer's accumulator, shaped like that span's K / V + auto acc_rows = [&](ggml_tensor * acc, const ggml_tensor * like, int64_t p0) { + return ggml_view_3d(ctx0, acc, like->ne[0], like->ne[1], like->ne[2], + like->ne[0]*sizeof(float), acc->nb[1], p0*acc->nb[1]); + }; // [const prefix][gradient prefix][this chunk]: the order the cells hold them in std::vector k_parts, v_parts; if (g0 > 0) { @@ -2894,20 +2914,15 @@ ggml_tensor * llm_graph_context::build_attn( v_parts.push_back(v_span(0, g0)); } if (n_past > g0) { - ggml_tensor * kc = ggml_cont(ctx0, ggml_cast(ctx0, k_span(g0, n_past), GGML_TYPE_F32)); - ggml_tensor * vc = ggml_cont(ctx0, ggml_cast(ctx0, v_span(g0, n_past), GGML_TYPE_F32)); - io.dk = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, kc->ne[0], kc->ne[1], kc->ne[2], 1); - io.dv = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, vc->ne[0], vc->ne[1], vc->ne[2], 1); - ggml_format_name(io.dk, "walk_dk-%d", il); - ggml_format_name(io.dv, "walk_dv-%d", il); - ggml_set_input(io.dk); - ggml_set_input(io.dv); - ggml_set_grad(io.dk); - ggml_set_grad(io.dv); - ggml_tensor * kg = ggml_add(ctx0, kc, io.dk); - ggml_tensor * vg = ggml_add(ctx0, vc, io.dv); - k_parts.push_back(kg->type != k_cur->type ? ggml_cast(ctx0, kg, k_cur->type) : kg); - v_parts.push_back(vg->type != v_cur->type ? ggml_cast(ctx0, vg, v_cur->type) : vg); + GGML_ASSERT(k_cur->type == GGML_TYPE_F32 && v_cur->type == GGML_TYPE_F32); + io.tk = k_span(g0, n_past); + io.tv = v_span(g0, n_past); + ggml_format_name(io.tk, "walk_tk-%d", il); + ggml_format_name(io.tv, "walk_tv-%d", il); + io.ak = acc_rows(acc_k, io.tk, g0); + io.av = acc_rows(acc_v, io.tv, g0); + k_parts.push_back(io.tk); + v_parts.push_back(io.tv); } k_parts.push_back(k_cur); v_parts.push_back(v_cur); @@ -2918,17 +2933,13 @@ ggml_tensor * llm_graph_context::build_attn( v = ggml_concat(ctx0, v, v_parts[i], 2); } if (cparams.walk_surrogate) { - io.gk = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, k_cur->ne[0], k_cur->ne[1], k_cur->ne[2]); - io.gv = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, v_cur->ne[0], v_cur->ne[1], v_cur->ne[2]); - ggml_format_name(io.gk, "walk_gk-%d", il); - ggml_format_name(io.gv, "walk_gv-%d", il); - ggml_set_input(io.gk); - ggml_set_input(io.gv); + ggml_tensor * gk = acc_rows(acc_k, k_cur, n_past); + ggml_tensor * gv = acc_rows(acc_v, v_cur, n_past); ggml_tensor * k32 = k_cur->type == GGML_TYPE_F32 ? k_cur : ggml_cast(ctx0, k_cur, GGML_TYPE_F32); ggml_tensor * v32 = v_cur->type == GGML_TYPE_F32 ? v_cur : ggml_cast(ctx0, v_cur, GGML_TYPE_F32); ggml_tensor * term = ggml_add(ctx0, - ggml_sum(ctx0, ggml_mul(ctx0, k32, io.gk)), - ggml_sum(ctx0, ggml_mul(ctx0, v32, io.gv))); + ggml_sum(ctx0, ggml_mul(ctx0, k32, gk)), + ggml_sum(ctx0, ggml_mul(ctx0, v32, gv))); res->t_walk_surrogate = res->t_walk_surrogate ? ggml_add(ctx0, res->t_walk_surrogate, term) : term; } res->t_walk.push_back(io); diff --git a/src/llama-graph.h b/src/llama-graph.h index 5f924306f7f6..9b9d00758df2 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -942,12 +942,11 @@ class llm_graph_result { // fills before the step and reads after it (llama_context::opt_epoch_iter) struct walk_layer { int il; - uint32_t grad_from; // the GRAD leaves cover cached positions [grad_from, n_past) uint32_t n_past; - ggml_tensor * dk = nullptr; // GRAD leaves, zero-filled: their gradient is dL/d(cached K/V) - ggml_tensor * dv = nullptr; // [n_embd_head, n_head_kv, n_past - grad_from], V untransposed - ggml_tensor * gk = nullptr; // inputs: the gradient later chunks accumulated on this chunk's - ggml_tensor * gv = nullptr; // own K/V, shaped like k_cur / v_cur (null without a surrogate) + ggml_tensor * tk = nullptr; // the cached K/V of [grad_from, n_past) as F32, tapped: their + ggml_tensor * tv = nullptr; // gradient adds into ak / av (ggml_opt_accumulate_grad) + ggml_tensor * ak = nullptr; // the accumulator's rows [grad_from, n_past), shaped like tk / tv + ggml_tensor * av = nullptr; }; std::vector t_walk; // the recurrent states of a reverse-pass chunk graph, one per state tensor of the memory diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 96484141528c..344d9bebab43 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -298,6 +298,8 @@ if (NOT LLAMA_SANITIZE_ADDRESS AND NOT GGML_SCHED_NO_REALLOC) llama_build_and_test(test-opt.cpp) endif() llama_build_and_test(test-backend-ops.cpp) +# the tiled FLASH_ATTN_EXT backward gives plain attention's gradients on every backend (no model needed) +llama_build_and_test(test-flash-attn-back.cpp) llama_build_and_test(test-moe-gather.cpp) llama_build_and_test(test-model-load-cancel.cpp LABEL "model") diff --git a/tests/test-flash-attn-back.cpp b/tests/test-flash-attn-back.cpp new file mode 100644 index 000000000000..c3441edee957 --- /dev/null +++ b/tests/test-flash-attn-back.cpp @@ -0,0 +1,156 @@ +// what this catches: the tiled FLASH_ATTN_EXT backward (ggml_flash_attn_ext_backward) giving a +// gradient other than attention's. Training reads attention through flash attention so its memory is +// queries x tile, never queries x window x heads (continuum, 2026-10-11), and every backend that runs +// it must give the gradients of plain attention: across several tiles and a ragged last one, grouped +// query heads, a causal mask that leaves whole tiles masked for some rows, and F32 or F16 K/V. +// +// The reference is the explicit path (mul_mat, soft_max_ext, mul_mat) on the CPU at F32, the path +// training took before. A finite difference is not used: through a kernel that rounds K/V or stages +// products at reduced precision its noise is larger than the defects this has to find (measured: +// Metal's half-staged mul_mm moved the tiled gradients 3.7e-4 relative, which this catches at 1e-4). +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include + +struct shape { int64_t d, h, hkv, m, n; int tile; }; + +// d(sum(out * w)) / d{q, k, v}, read in logical order +static bool grads(ggml_backend_t be, bool flash, const shape & s, ggml_type kv_type, std::vector (&g)[3]) { + ggml_init_params ip = { ggml_tensor_overhead()*16384 + ggml_graph_overhead_custom(1 << 16, true), nullptr, true }; + ggml_context * ctx = ggml_init(ip); + ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, s.d, s.n, s.h); + ggml_tensor * k = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, s.d, s.m, s.hkv); + ggml_tensor * v = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, s.d, s.m, s.hkv); + ggml_tensor * w = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, s.d, s.h, s.n); + ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, s.m, GGML_PAD(s.n, 64), 1, 1); + ggml_set_param(q); + ggml_set_param(k); + ggml_set_param(v); + // F16 K/V: the flash path casts (as a cache feeds it); the reference reads the same F16-exact + // values at F32 (its explicit backward cannot take a gradient through a transposed cast) + ggml_tensor * kk = kv_type == GGML_TYPE_F32 || !flash ? k : ggml_cast(ctx, k, kv_type); + ggml_tensor * vv = kv_type == GGML_TYPE_F32 || !flash ? v : ggml_cast(ctx, v, kv_type); + const float scale = 1.0f / sqrtf((float) s.d); + ggml_tensor * out; + if (flash) { + out = ggml_flash_attn_ext(ctx, q, kk, vv, m, scale, 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + ggml_flash_attn_ext_set_back_tile(out, s.tile); + } else { + ggml_tensor * mv = ggml_cont(ctx, ggml_cast(ctx, ggml_view_2d(ctx, m, s.m, s.n, m->nb[1], 0), GGML_TYPE_F32)); + ggml_tensor * kq = ggml_mul_mat(ctx, kk, q); + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + kq = ggml_soft_max_ext(ctx, kq, mv, scale, 0.0f); + ggml_tensor * kqv = ggml_mul_mat(ctx, ggml_cont(ctx, ggml_transpose(ctx, vv)), kq); + out = ggml_cont(ctx, ggml_permute(ctx, kqv, 0, 2, 1, 3)); + } + ggml_tensor * loss = ggml_sum(ctx, ggml_mul(ctx, out, w)); + ggml_set_loss(loss); + ggml_cgraph * gf = ggml_new_graph_custom(ctx, 1 << 16, true); + ggml_build_forward_expand(gf, loss); + ggml_build_backward_expand(ctx, gf, nullptr); + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + if (!ggml_backend_supports_op(be, ggml_graph_node(gf, i))) { + ggml_free(ctx); + return false; + } + } + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, be); + std::mt19937 rng(42); + std::uniform_real_distribution uni(-1.0f, 1.0f); + for (ggml_tensor * t : { q, k, v, w }) { + std::vector x(ggml_nelements(t)); + for (auto & e : x) { + e = uni(rng); + if (kv_type == GGML_TYPE_F16 && (t == k || t == v)) { + e = ggml_fp16_to_fp32(ggml_fp32_to_fp16(e)); // F16-exact, so both paths see one value + } + } + ggml_backend_tensor_set(t, x.data(), 0, x.size()*sizeof(float)); + } + // causal over the LAST n positions, the way a chunk sits after its context + std::vector mm(ggml_nelements(m)); + for (int64_t i = 0; i < m->ne[1]; ++i) { + for (int64_t j = 0; j < s.m; ++j) { + mm[i*s.m + j] = ggml_fp32_to_fp16(i < s.n && j <= s.m - s.n + i ? 0.0f : -INFINITY); + } + } + ggml_backend_tensor_set(m, mm.data(), 0, mm.size()*sizeof(ggml_fp16_t)); + ggml_graph_reset(gf); + GGML_ASSERT(ggml_backend_graph_compute(be, gf) == GGML_STATUS_SUCCESS); + ggml_tensor * ps[3] = { q, k, v }; + for (int i = 0; i < 3; ++i) { + ggml_tensor * gr = ggml_graph_get_grad(gf, ps[i]); + std::vector raw(ggml_nbytes(gr)); + ggml_backend_tensor_get(gr, raw.data(), 0, raw.size()); + g[i].clear(); + for (int64_t i2 = 0; i2 < gr->ne[2]; ++i2) { + for (int64_t i1 = 0; i1 < gr->ne[1]; ++i1) { + for (int64_t i0 = 0; i0 < gr->ne[0]; ++i0) { + g[i].push_back(*(const float *) (raw.data() + i0*gr->nb[0] + i1*gr->nb[1] + i2*gr->nb[2])); + } + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + return true; +} + +int main() { + ggml_backend_load_all(); + const shape shapes[] = { + { 64, 6, 2, 300, 37, 64 }, // several tiles, a ragged last one, grouped heads + { 256, 24, 4, 1000, 64, 256 }, // Qwen3.8-27B's attention layout + { 64, 3, 1, 50, 50, 16 }, // early rows whose later tiles are wholly masked + { 64, 4, 4, 200, 32, 512 }, // one tile, no grouping + }; + const double max_rel = 1e-4; + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + int failures = 0; + for (ggml_type kv : { GGML_TYPE_F32, GGML_TYPE_F16 }) { + for (const shape & s : shapes) { + std::vector ref[3]; + GGML_ASSERT(grads(cpu, false, s, kv, ref)); + for (size_t di = 0; di < ggml_backend_dev_count(); ++di) { + ggml_backend_dev_t dev = ggml_backend_dev_get(di); + ggml_backend_t be = ggml_backend_dev_init(dev, nullptr); + if (!be) { + continue; + } + std::vector got[3]; + if (!grads(be, true, s, kv, got)) { + printf(" %-8s d=%lld h=%lld/%lld m=%lld n=%lld tile=%d %s: not supported\n", ggml_backend_name(be), + (long long) s.d, (long long) s.h, (long long) s.hkv, (long long) s.m, (long long) s.n, s.tile, ggml_type_name(kv)); + ggml_backend_free(be); + continue; + } + double worst = 0.0; + for (int i = 0; i < 3; ++i) { + double md = 0.0, mr = 0.0; + for (size_t j = 0; j < ref[i].size(); ++j) { + md = std::fmax(md, std::fabs(got[i][j] - ref[i][j])); + mr = std::fmax(mr, std::fabs(ref[i][j])); + } + worst = std::fmax(worst, md / mr); + } + const bool ok = worst < max_rel; + printf(" %-8s d=%lld h=%lld/%lld m=%lld n=%lld tile=%d %s: max rel %.2e %s\n", ggml_backend_name(be), + (long long) s.d, (long long) s.h, (long long) s.hkv, (long long) s.m, (long long) s.n, s.tile, + ggml_type_name(kv), worst, ok ? "OK" : "FAILED"); + failures += ok ? 0 : 1; + ggml_backend_free(be); + } + } + } + ggml_backend_free(cpu); + printf("%s\n", failures == 0 ? "OK" : "FAILED"); + return failures == 0 ? 0 : 1; +} diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 342a327a7829..f6a69276c9a6 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -235,6 +235,46 @@ int main(int argc, char ** argv) { llama_free(tok_ctx); GGML_ASSERT(tokens.size() > REFUSAL_WINDOW + 1); + // 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 + // table and exits; the bytes are the run's own allocation preflight (llama_opt_graph_bytes). + if (const char * env = getenv("WALK_MEASURE")) { + std::string more = text; + while (more.size() < 6 * 65536) { + lcg = lcg * 1664525u + 1013904223u; + more += words[(lcg >> 16) % 16]; + more += (lcg >> 8) % 7 == 0 ? ".\n" : " "; + } + llama_context * tc = make_ctx(params, model, WINDOW); + const std::vector long_tokens = common_tokenize(tc, more, true); + llama_free(tc); + const uint32_t chunk = getenv("WALK_MEASURE_CHUNK") ? (uint32_t) atoi(getenv("WALK_MEASURE_CHUNK")) : 256; + g_recompute = getenv("WALK_MEASURE_RECOMPUTE") != nullptr; // per-layer recompute, as the served runs train + g_walk_horizon = getenv("WALK_MEASURE_HORIZON") ? (uint32_t) atoi(getenv("WALK_MEASURE_HORIZON")) : 0; // 0 = the whole window + printf(" window chunk explicit MiB flash MiB ratio\n"); + for (const char * p = env; *p; ) { + char * end = nullptr; + const uint32_t w = (uint32_t) strtoul(p, &end, 10); + p = *end == ',' ? end + 1 : end; + if (w == 0 || long_tokens.size() <= w + 1) { break; } + g_window = w; + std::vector reply(w + 1, 0); + for (uint32_t i = w - 511; i <= w; ++i) reply[i] = 1; + double mib[2] = { 0, 0 }; + const ggml_type types[2] = { GGML_TYPE_F32, GGML_TYPE_F16 }; + for (int a = 0; a < 2; ++a) { + g_cache_type = types[a]; + step_delta(params, model, init, long_tokens, reply, chunk, true); + mib[a] = g_last_graph_bytes / 1048576.0; + } + g_cache_type = GGML_TYPE_F32; + printf(" %6u %5u %12.1f %9.1f %5.2fx\n", w, chunk, mib[0], mib[1], mib[1] > 0 ? mib[0] / mib[1] : 0.0); + fflush(stdout); + } + return 0; + } + int failures = 0; { // 1. every position labelled: chunk = window is one graph, the true gradient From 73bdef9a7f40bf8eb6a44a5a306d19b61b708497 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Sun, 11 Oct 2026 12:15:05 -0500 Subject: [PATCH 2/6] ggml-opt: recompute casts; a CPY into its own result is not a side effect The checkpoint pass never recomputed a GGML_OP_CPY (a write into another tensor, a cache) and counted every CPY's buffer as written in place, pinning whatever read it. ggml_cast is a CPY whose destination is itself: a pure conversion. So the exact walk's f32 cast of each layer's cached K/V, and the concat reading it, were computed once in the forward and held until that layer's backward, for every layer at once: window x 2 x n_embd_kv_gqa x 4 bytes x layers. On the 5090 at Kimi's 67,584 window that is ~8.9 GB by arithmetic (16 layers x 1024 x 2 x 4 B), the term that left chunk 64 at 10.4 GB after every smaller chunk. A cast is now recomputed like any op (and does not mark a mutated buffer); a CPY into another tensor still is not. place() skips a cast's self-reference. Exactness (test-walk-exact, qwen2.5-0.5b, M5): exact 1.0000, f16 1.0000, q8_0 0.9905, recompute vs not 1.0000; test-grad-leaf, test-opt-dynamic-accum OK. test-walk-exact gains WALK_MEASURE_HORIZON for measuring the gradient-horizon term. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml-opt.cpp | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 5a604762ef83..87cd73369d30 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -308,6 +308,12 @@ static ggml_tensor * ggml_opt_view_root(ggml_tensor * t) { return t; } +// A CPY whose destination is its own result (ggml_cast) writes no other tensor: it is a pure +// conversion, recomputed like any op. A CPY into another tensor (a cache write) is a side effect. +static bool ggml_opt_is_cast(const ggml_tensor * node) { + return node->op == GGML_OP_CPY && node->src[1] == node; +} + static ggml_tensor * ggml_opt_recompute_node( ggml_context * ctx, const std::unordered_set & forward, const std::unordered_set & keep, std::unordered_map & clones, @@ -323,9 +329,13 @@ static ggml_tensor * ggml_opt_recompute_node( } switch (node->op) { case GGML_OP_NONE: - case GGML_OP_CPY: // a write into another tensor (a cache): side effect, never repeated case GGML_OP_SET_ROWS: return node; + case GGML_OP_CPY: // a write into another tensor (a cache): side effect, never repeated + if (!ggml_opt_is_cast(node)) { + return node; + } + break; default: break; } @@ -367,7 +377,7 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd, // 6.8e-4 with recompute on until in-place ops counted here). const bool view_op = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE; - if (node->op == GGML_OP_CPY || node->op == GGML_OP_SET_ROWS || (node->view_src != nullptr && !view_op)) { + if ((node->op == GGML_OP_CPY && !ggml_opt_is_cast(node)) || node->op == GGML_OP_SET_ROWS || (node->view_src != nullptr && !view_op)) { mutated.insert(ggml_opt_view_root(node)); } } @@ -411,7 +421,9 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd, return; } for (int k = 0; k < GGML_MAX_SRC; ++k) { - place(t->src[k]); + if (t->src[k] != t) { // a cast is its own destination + place(t->src[k]); + } } place(t->view_src); placed.insert(t); From 6d3b7f79721e87494ca0ba32cf99b8584100b50b Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Sun, 11 Oct 2026 12:26:25 -0500 Subject: [PATCH 3/6] train: the training graph is the serving graph; attention reads the cache in place The walk built its own attention input: every layer of every chunk copied the cached prefix out of the cache (a cast to F32 and a concat with the chunk), a window-sized copy per layer, held to its backward. On the 5090 at Kimi's 67,584 window every chunk size was refused (10.4-16 GB against 4.55 GB free) and all 16 examples were parked. Now a layer that trains through FLASH_ATTN_EXT does what serving does: the chunk writes its K/V into the cache (SET_ROWS, the cache's own type) and attention reads the cache in place. The forward decode already put those values there, so the write is idempotent. - ggml: SET_ROWS gets a backward (the written source's gradient is the result's gradient at the written rows; the destination, a cache, carries none); the backward admits a quantized source where it only re-shapes a gradient or casts a tile to F32 (VIEW/RESHAPE/PERMUTE/ TRANSPOSE/FLASH_ATTN_EXT), so a q8_0 cache trains in place. - The exact walk taps the write into the host accumulator over the WHOLE window; the chunk's own rows are read back from it, so later chunks' gradient reaches the adapter with no surrogate term, and no gradient horizon bounds it. - llm_graph_context::uses_flash_attn is the one decision build_attn_mha and training read; build_inp_train_rows reuses llm_graph_input_pos for the chunk's rows. - Explicit attention (ALiBi, softcap, sinks, KQ bias, MLA: no flash backward yet) keeps the copy path, named as the fallback. test-walk-exact (qwen2.5-0.5b, M5), the cache-type check split into its two claims: exact walk vs one graph at the same cache: f16 1.0000, q8_0 1.0000 one graph at the cache vs the F32 gradient: f16 0.9999, q8_0 0.9624 (head dim 64; printed) F32 exact 1.0000, recompute vs not 1.0000 test-flash-attn-back OK, test-grad-leaf OK, test-opt-dynamic-accum OK, SET_ROWS 171/171 (Metal). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- ggml/src/ggml.c | 21 ++++++++++-- src/llama-graph.cpp | 67 ++++++++++++++++++++++++++++++++++++--- src/llama-graph.h | 10 ++++++ tests/test-walk-exact.cpp | 22 ++++++------- 4 files changed, 103 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index b00ebff70d61..5f572229d5f9 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -7210,6 +7210,16 @@ static void ggml_compute_backward( ggml_add_or_set(ctx, cgraph, isrc0, ggml_transpose(ctx, grad)); } } break; + case GGML_OP_SET_ROWS: { + // result = the destination (src2) with its rows src1 replaced by src0. The gradient + // on src0 is the result's gradient at those rows. The destination's gradient (the + // result's everywhere else) is not built: a written buffer (a cache) is a constant. + GGML_ASSERT(!src2_needs_grads && "set_rows: the destination it writes carries no gradient"); + if (src0_needs_grads) { + GGML_ASSERT(src1->type == GGML_TYPE_I32 && "set_rows backward gathers with I32 row indices"); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, ggml_get_rows(ctx, grad, src1), src0)); + } + } break; case GGML_OP_GET_ROWS: { if (src0_needs_grads) { ggml_add_or_set(ctx, cgraph, isrc0, ggml_get_rows_back(ctx, grad, src1, src0)); @@ -7675,7 +7685,13 @@ void ggml_build_backward_expand( if (!node->src[j] || ignore_src[j] || !grads_needed[ggml_hash_find(&cgraph->visited_hash_set, node->src[j])]) { continue; } - GGML_ASSERT(node->src[j]->type == GGML_TYPE_F32 || node->src[j]->type == GGML_TYPE_F16); + // 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 || + 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))); node_needs_grad = true; break; } @@ -7685,7 +7701,8 @@ void ggml_build_backward_expand( // inplace operations are currently not supported GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || - node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE); + node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE || + node->op == GGML_OP_SET_ROWS); const size_t ihash = ggml_hash_find(&cgraph->visited_hash_set, node); GGML_ASSERT(ihash != GGML_HASHSET_FULL); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 30bb874bdaf5..1c98fc717b8f 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2411,6 +2411,24 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const { return cur; } +bool llm_graph_context::uses_flash_attn(ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla) const { + const bool fa_back_ok = !cparams.training || + (hparams.f_max_alibi_bias == 0.0f && !hparams.attn_soft_cap && sinks == nullptr && v_mla == nullptr); + return cparams.flash_attn && kq_b == nullptr && fa_back_ok; +} + +ggml_tensor * llm_graph_context::build_inp_train_rows() const { + if (inp_train_rows == nullptr) { + auto inp = std::make_unique(1); + inp->pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->pos); + ggml_set_name(inp->pos, "train_rows"); + inp_train_rows = inp->pos; + res->add_input(std::move(inp)); + } + return inp_train_rows; +} + ggml_tensor * llm_graph_context::build_inp_pos() const { auto inp = std::make_unique(hparams.n_pos_per_embd()); @@ -2584,9 +2602,7 @@ ggml_tensor * llm_graph_context::build_attn_mha( // which were ~75% of a 27B's training graph (continuum, the 5090, 2026-10-11). The backward // covers plain scaled attention; a model with ALiBi, a logit softcap, sinks or MLA trains on // the explicit path, which has a backward for all of them. - const bool fa_back_ok = !cparams.training || - (hparams.f_max_alibi_bias == 0.0f && !hparams.attn_soft_cap && sinks == nullptr && v_mla == nullptr); - const bool use_flash_attn = cparams.flash_attn && kq_b == nullptr && fa_back_ok; + const bool use_flash_attn = uses_flash_attn(kq_b, sinks, v_mla); if (use_flash_attn) { GGML_ASSERT(kq_b == nullptr && "Flash attention does not support KQ bias yet"); @@ -2852,7 +2868,50 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * q = q_cur; ggml_tensor * k; ggml_tensor * v; - if (cparams.training) { + if (cparams.training && uses_flash_attn(kq_b, sinks, v_mla)) { + // THE TRAINING GRAPH IS THE SERVING GRAPH. This chunk writes its K/V into the cache + // exactly as serving does (SET_ROWS, the cache's own type) and attention reads the + // cache in place: no copy of the context, at any window. The walk decoded the window + // into the cache under the same adapter, so the write is idempotent; the backward + // gathers this chunk's gradient from the rows it wrote (SET_ROWS backward). In the + // exact walk the write is TAPPED (cparams.walk_acc_k / _v): the backward adds + // dL/d(cached K/V) of the whole window into the host accumulator, and reads its own + // rows back from it, so what later chunks put on them reaches the adapter through + // this chunk with no surrogate term. A quantized or f16 cache trains on the values + // serving reads; its gradient passes straight through the cast. + const int64_t n_tokens = k_cur->ne[2]; + const int64_t n_past = ubatch.pos[0]; + const int64_t n_win = n_past + n_tokens; + ggml_tensor * rows = build_inp_train_rows(); + llm_graph_result::walk_layer io = { il, (uint32_t) n_past }; + // the window's cells [0, n_win) of one layer's cache with this chunk's rows written + auto write = [&](ggml_tensor * all, ggml_tensor * cur, ggml_tensor * acc, ggml_tensor *& tap, ggml_tensor *& tap_acc) { + const int64_t n_embd = cur->ne[0]*cur->ne[1]; + GGML_ASSERT(all->ne[3] == 1 && "a training context reads one stream"); + GGML_ASSERT(ggml_row_size(cur->type, cur->ne[0]) == cur->nb[1]); + ggml_tensor * win = ggml_view_2d(ctx0, all, n_embd, n_win, all->nb[2], 0); + ggml_tensor * w = ggml_set_rows(ctx0, win, ggml_view_2d(ctx0, cur, n_embd, n_tokens, cur->nb[2], 0), rows); + if (acc != nullptr) { + tap = w; + tap_acc = ggml_view_2d(ctx0, acc, n_embd, n_win, acc->nb[1], 0); + } + return ggml_reshape_3d(ctx0, w, cur->ne[0], cur->ne[1], n_win); + }; + ggml_tensor * v_all = mctx_cur->get_v(ctx0, il); + GGML_ASSERT(v_all->nb[1] <= v_all->nb[2] && "a flash-attention cache holds V untransposed"); + const bool tapped = cparams.walk_exact; + k = write(mctx_cur->get_k(ctx0, il), k_cur, tapped ? cparams.walk_acc_k[il] : nullptr, io.tk, io.ak); + v = write(v_all, v_cur, tapped ? cparams.walk_acc_v[il] : nullptr, io.tv, io.av); + if (tapped) { + res->t_walk.push_back(io); + } + if (kq_mask->ne[0] > n_win) { + 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) { + // 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 // and the read is a view of a leaf. So this ubatch's K/V are attended to directly, // and carry the gradient. The context before it (positions [0, n_past), decoded into diff --git a/src/llama-graph.h b/src/llama-graph.h index 9b9d00758df2..47df0126241b 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1062,6 +1062,8 @@ struct llm_graph_context { llm_graph_result * res; + mutable ggml_tensor * inp_train_rows = nullptr; // build_inp_train_rows, once per graph + ggml_context * ctx0 = nullptr; ggml_cgraph * gf = nullptr; @@ -1179,6 +1181,14 @@ struct llm_graph_context { ggml_tensor * build_inp_embd(ggml_tensor * tok_embd) const; ggml_tensor * build_inp_pos() const; + + // A training chunk's cache rows: I32 [n_tokens], its positions (in the walk a position + // is its cell). One per graph, shared by every attention layer. + ggml_tensor * build_inp_train_rows() const; + + // Whether this layer's attention runs through FLASH_ATTN_EXT: the one decision + // build_attn_mha makes and a training graph reads to attend to the cache in place. + bool uses_flash_attn(ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla) const; ggml_tensor * build_inp_attn_scale() const; ggml_tensor * build_inp_out_ids() const; ggml_tensor * build_inp_mean() const; diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index f6a69276c9a6..1724b8d18912 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -196,10 +196,6 @@ static agreement compare(const char * what, const std::vector & a, const // Measured on the 1.5B Q4_K_M (5090): the exact walk agrees with one graph at cosine 0.998 and // norm 0.995 (what is left is the quantized matmuls' batch-size variance: a 128-row chunk and a // 512-row graph take different kernels); the plain walk sits at cosine 0.81, norm 2.1. -// the served cache types' bar (see 1b). Measured on the 5090: f16 0.9986 (1.5B) / 0.9981 (hybrid), -// q8_0 0.9978 / 0.9979, against F32's 0.998: the straight-through approximation costs nothing -// measurable, so the bar sits just under F32's own agreement. -static const double STRAIGHT_THROUGH_COSINE = 0.99; static bool same_step(const agreement & r) { return r.cosine > 0.995 && std::fabs(r.norm_ratio - 1.0) < 0.02; @@ -305,21 +301,25 @@ int main(int argc, char ** argv) { } } { - // 1b. the served cache types: the walk reads the prefix at the cache's precision and its - // gradient is taken there, then applied to the chunk's own full-precision K/V (straight- - // through). Measured against the same one-graph reference, which never reads the cache. + // 1b. the served cache types. Training attends to the cache in place, as serving does: + // every K/V it reads, the chunk's own included, is at the cache's precision, and the + // gradient passes straight through the cast. Two claims, measured apart: the walk is + // EXACT at that precision (against one graph reading the same cache), and what the + // precision itself costs against the F32 gradient (printed, the served model's own). std::vector all(WINDOW + 1, 1); all[0] = 0; - const auto ref = step_delta(params, model, init, tokens, all, WINDOW, false); + const auto truth = step_delta(params, model, init, tokens, all, WINDOW, false); for (ggml_type t : { GGML_TYPE_F16, GGML_TYPE_Q8_0 }) { g_cache_type = t; + const auto one = step_delta(params, model, init, tokens, all, WINDOW, false); const auto exact = step_delta(params, model, init, tokens, all, WINDOW / 4, true); g_cache_type = GGML_TYPE_F32; - const agreement r = compare((std::string("all labelled, ") + ggml_type_name(t) + " cache: exact walk in 4 chunks vs one graph").c_str(), exact, ref); - if (!(r.cosine > STRAIGHT_THROUGH_COSINE)) { - fprintf(stderr, "FAILED: at a %s cache the exact walk's step strays from the window's gradient (cosine %.4f)\n", ggml_type_name(t), r.cosine); + const std::string name = std::string("all labelled, ") + ggml_type_name(t) + " cache: "; + if (!same_step(compare((name + "exact walk in 4 chunks vs one graph at the same cache").c_str(), exact, one))) { + fprintf(stderr, "FAILED: at a %s cache the exact walk is not the window's gradient\n", ggml_type_name(t)); ++failures; } + compare((name + "one graph vs the F32 gradient (the cache's own precision)").c_str(), one, truth); } } { From ea1a4d1bb1d29d22018925d9180670833a697c58 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Sun, 11 Oct 2026 19:00:11 -0500 Subject: [PATCH 4/6] train: the exact walk is exact on a GPU hybrid; recurrent state gradients go to host accumulators A hybrid (attention + gated delta-net) exact walk trained on a gradient ~0.93-0.94 aligned with the true one on Metal, varying run to run, while the CPU gave 1.0000. Cause: decode() only queues its graph, and the walk's hand-rolled per-layer state snapshot read and restored the recurrent rows with ggml_backend_tensor_get/set, a bare memcpy on a host-visible buffer that does not wait. The reverse pass started from a torn, mid-graph state: the exact walk's own forward loss was 3.283623 against one graph's 3.252877. The fix deletes the parallel snapshot path and calls the owner: the sequence-state API (state_seq_get/set_data, PARTIAL_ONLY = the recurrent part), after synchronize(). A restore puts the cell's position back with its rows, and an empty checkpoint is the zero state, so the rewind's special case for the first chunk goes too. The state entering a chunk is now tapped (ggml_opt_accumulate_grad into a host accumulator, as the K/V taps are) instead of a GRAD leaf whose gradient was read back after the graph: the 24 delta-net backward outputs were pinned to the graph's end for that readback, and the allocator's high-water was 3512 MiB against 562 MiB live at a 4096 window. test-walk-exact names the worst LoRA tensors when the step strays (this localized the bug to the attention layers' adapters, worsening toward shallow layers), and test-backend-ops gains GATED_DELTA_NET_BACK at the 4B hybrid's own layer (16 key heads x 2, 128, 512 tokens). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- src/llama-context.cpp | 125 +++++++++++++++---------------------- src/llama-graph.cpp | 13 ++-- src/llama-graph.h | 3 +- tests/CMakeLists.txt | 3 +- tests/test-backend-ops.cpp | 4 +- tests/test-walk-exact.cpp | 47 +++++++++++++- 6 files changed, 108 insertions(+), 87 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 18d094873c17..6ed7cb46b2ff 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3515,19 +3515,21 @@ void llama_context::opt_epoch_iter( float loss_scale; // its labelled positions over the window's }; const exact_chunk * xc = nullptr; - // per layer: dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, position] F32 - // in host memory (null for a recurrent layer); every chunk's backward adds into it in place. - // The graph reads it through cparams, which never outlives it: whatever way the walk returns. + // THE EXACT WALK's gradients, in host memory, each added into in place by the backward: + // per attention layer dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, + // position] F32 (null for a recurrent layer), which the graph reads through cparams (never + // outliving it, whatever way the walk returns); per recurrent state tensor (a layer's conv + // state r_l and its recurrent state s_l) dL/d(the state entering the chunk trained last), + // which is what the chunk before it puts on the state it leaves (state_ready once one is in) struct walk_acc_t { llama_cparams & cp; std::vector k, v; + std::unordered_map state; + bool state_ready = false; ggml_context_ptr ctx; ggml_backend_buffer_ptr buf; ~walk_acc_t() { cp.walk_acc_k = nullptr; cp.walk_acc_v = nullptr; } - } walk_acc { cparams, {}, {}, nullptr, nullptr }; - // per recurrent state tensor: dL/d(the state the chunk just before leaves), from the chunk - // trained last (each chunk's exit state feeds exactly one chunk, the next) - std::unordered_map> walk_gstate; + } walk_acc { cparams, {}, {}, {}, false, nullptr, nullptr }; // one training chunk [pos_ctx, pos_ctx + n_tokens): forward and backward over it alone, // attending to the cached context before it @@ -3622,6 +3624,11 @@ void llama_context::opt_epoch_iter( ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); } } + for (const auto & ws : res->t_walk_state) { + if (ws.tap) { + ggml_opt_accumulate_grad(opt_ctx, ws.tap, walk_acc.state.at(ws.cache)); + } + } } if (!ggml_opt_alloc(opt_ctx, train)) { const char * why = ggml_opt_refusal(opt_ctx); @@ -3659,15 +3666,15 @@ void llama_context::opt_epoch_iter( GGML_ASSERT(row == n_outputs); } if (xc != nullptr) { + // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the + // accumulator takes this chunk's own for (const auto & ws : res->t_walk_state) { - if (ws.ds) { - ggml_backend_tensor_memset(ws.ds, 0, 0, ggml_nbytes(ws.ds)); - } + ggml_tensor * acc = walk_acc.state.at(ws.cache); if (ws.gs) { - const auto it = walk_gstate.find(ws.cache); - GGML_ASSERT(it != walk_gstate.end() && it->second.size() == (size_t) ggml_nelements(ws.gs)); - ggml_backend_tensor_set(ws.gs, it->second.data(), 0, ggml_nbytes(ws.gs)); + GGML_ASSERT(ggml_nbytes(acc) == ggml_nbytes(ws.gs)); + ggml_backend_tensor_set(ws.gs, acc->data, 0, ggml_nbytes(ws.gs)); } + memset(acc->data, 0, ggml_nbytes(acc)); } } ggml_opt_eval(opt_ctx, xc != nullptr && xc->n_labels == 0 ? nullptr : result); @@ -3680,18 +3687,8 @@ void llama_context::opt_epoch_iter( opt_stop_requested.store(true); return false; } - if (xc != nullptr) { - // dL/d(entry state): what the previous chunk's exit state takes in its surrogate - for (const auto & ws : res->t_walk_state) { - if (!ws.ds) { - continue; - } - ggml_tensor * grad = ggml_opt_leaf_grad(opt_ctx, ws.ds); - GGML_ASSERT(grad != nullptr && grad->type == GGML_TYPE_F32 && ggml_is_contiguous(grad)); - auto & g = walk_gstate[ws.cache]; - g.resize(ggml_nelements(grad)); - ggml_backend_tensor_get(grad, g.data(), 0, ggml_nbytes(grad)); - } + if (xc != nullptr && !res->t_walk_state.empty()) { + walk_acc.state_ready = true; } if (callback && !(xc != nullptr && xc->n_labels == 0)) { callback(train, opt_ctx, dataset, result, idata_in_loop + (pos_ctx + pos_batch)/n_batch + 1, ndata_in_loop, t_loop_start); @@ -3752,7 +3749,7 @@ void llama_context::opt_epoch_iter( return; } // A recurrent model's state is carried the same way, one chunk at a time: the state - // entering a chunk is a GRAD leaf, the state it leaves meets the next chunk's gradient + // entering a chunk is tapped into a host accumulator, the state it leaves meets the next chunk's gradient // (build_rs, build_walk_state_exit). The reverse pass restores each chunk's entry state // from a host snapshot taken at its boundary in the forward pass. std::vector chunks; @@ -3839,10 +3836,17 @@ void llama_context::opt_epoch_iter( const uint32_t n_layer = model.hparams.n_layer(); walk_acc.k.assign(n_layer, nullptr); walk_acc.v.assign(n_layer, nullptr); - walk_acc.ctx.reset(ggml_init({ 2*n_layer*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true })); + walk_acc.ctx.reset(ggml_init({ 4*n_layer*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true })); + walk_acc.state.clear(); + walk_acc.state_ready = false; for (uint32_t il = 0; il < n_layer; ++il) { if (model.hparams.is_recr(il)) { - continue; // a recurrent layer has no K/V: its state carries the chain + // a recurrent layer has no K/V: its state carries the chain, one sequence's row + if (recr != nullptr) { + walk_acc.state[recr->r_l[il]] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_r(), 1); + walk_acc.state[recr->s_l[il]] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_s(), 1); + } + continue; } walk_acc.k[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_k_gqa(il), last_label + 1); walk_acc.v[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_v_gqa(il), last_label + 1); @@ -3861,50 +3865,25 @@ void llama_context::opt_epoch_iter( } // the window as the adapter now reads it, chunk by chunk: a recurrent model's state is - // snapshotted at each chunk's start (empty before the first: it starts from zero) - struct state_snapshot { - bool empty = true; - llama_pos pos = -1; // the cell's last position at the snapshot - std::vector> r, s; // per layer + // checkpointed at each chunk's start through the sequence-state API (the recurrent part + // only, PARTIAL_ONLY), after waiting for the queued decode; a restore puts back the + // cell's position with its rows, and an empty checkpoint (before the first chunk) + // restores the zero state. A hand-rolled per-layer copy here skipped that wait and, on + // Metal, carried a torn mid-graph state into the reverse pass (2026-10-11). + std::vector> snaps(recr != nullptr ? chunks.size() : 0); + const llama_state_seq_flags state_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; + auto save_state = [&](std::vector & snap) { + synchronize(); + snap.resize(state_seq_get_size(0, state_flags)); + GGML_ASSERT(state_seq_get_data(0, snap.data(), snap.size(), state_flags) == snap.size()); }; - std::vector snaps(recr != nullptr ? chunks.size() : 0); - auto state_rows = [&](bool write, state_snapshot & snap) { - const int32_t cell = recr->cells[0].tail; // the cell holding sequence 0's state - if (!write) { - snap.empty = cell < 0; - snap.pos = cell < 0 ? -1 : recr->cells[cell].pos; - snap.r.assign(recr->r_l.size(), {}); - snap.s.assign(recr->s_l.size(), {}); - } - if (cell < 0) { - return; - } - // the restored state is the state AFTER snap.pos, so the cell says so: left at the - // window's end, the next ubatch read as non-consecutive ("position 65640 after 66345") - if (write) { - recr->cells[cell].pos = snap.pos; - } - for (size_t il = 0; il < recr->r_l.size(); ++il) { - for (int kind = 0; kind < 2; ++kind) { - ggml_tensor * t = kind == 0 ? recr->r_l[il] : recr->s_l[il]; - if (t == nullptr) { - continue; - } - GGML_ASSERT(t->type == GGML_TYPE_F32); - auto & row = kind == 0 ? snap.r[il] : snap.s[il]; - const size_t bytes = t->nb[1]; - if (write) { - ggml_backend_tensor_set(t, row.data(), (size_t) cell*bytes, bytes); - } else { - row.resize(t->ne[0]); - ggml_backend_tensor_get(t, row.data(), (size_t) cell*bytes, bytes); - } - } - } + auto load_state = [&](const std::vector & snap) { + synchronize(); + GGML_ASSERT(state_seq_set_data(0, snap.data(), snap.size(), state_flags) == snap.size()); }; for (size_t j = 0; j < chunks.size(); ++j) { if (recr != nullptr && j % stride == 0) { - state_rows(false, snaps[j]); + save_state(snaps[j]); } if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n)) { return; @@ -3938,11 +3917,7 @@ void llama_context::opt_epoch_iter( return false; } if (recr != nullptr) { - if (snaps[cp].empty) { - recr->seq_rm(0, -1, -1); // the first chunk starts from a zero state - } else { - state_rows(true, snaps[cp]); - } + load_state(snaps[cp]); if (p_pop < c.c0 && !decode_span(p_pop, c.c0)) { return false; } @@ -3955,7 +3930,7 @@ void llama_context::opt_epoch_iter( } cparams.walk_exact = true; cparams.walk_surrogate = c.surrogate; - cparams.walk_state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && !walk_gstate.empty(); + cparams.walk_state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && walk_acc.state_ready; // THE HORIZON FITS THE DEVICE. Its GRAD leaves are the chunk graph's largest growth // with the window (layers x horizon x (k + v) floats, twice with their gradient); the // graph preflight refuses a chunk that does not fit before allocating anything, so a @@ -4005,7 +3980,7 @@ void llama_context::opt_epoch_iter( cparams.walk_surrogate = false; cparams.walk_state_surrogate = false; } - walk_gstate.clear(); + walk_acc.state_ready = false; return; } diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 1c98fc717b8f..9e9071876ab9 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3628,15 +3628,12 @@ ggml_tensor * llm_graph_context::build_rs( // {state_size, rs_size} -> {state_size, n_seqs} ggml_tensor * output_states = get_state_rows(ctx0, states, state_copy_main); if (cparams.training && cparams.walk_exact) { - // THE EXACT WALK: the state entering the chunk is the cached one PLUS a zero GRAD leaf, - // so the backward yields dL/d(entry state): the gradient the previous chunk's exit takes + // THE EXACT WALK: the state entering the chunk is tapped, so the backward adds + // dL/d(entry state), the gradient the previous chunk's exit takes, into a host + // accumulator the moment it is computed: no GRAD leaf, and no backward output held + // to the end of the graph for a readback GGML_ASSERT(output_states->type == GGML_TYPE_F32 && "the exact walk carries an F32 recurrent state"); - ggml_tensor * ds = ggml_new_tensor(ctx0, GGML_TYPE_F32, GGML_MAX_DIMS, output_states->ne); - ggml_format_name(ds, "walk_ds-%s", s->name); - ggml_set_input(ds); - ggml_set_grad(ds); - output_states = ggml_add(ctx0, output_states, ds); - res->walk_state_of(s).ds = ds; + res->walk_state_of(s).tap = output_states; } ggml_build_forward_expand(gf, output_states); diff --git a/src/llama-graph.h b/src/llama-graph.h index 47df0126241b..0e0adb7d4bda 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -953,7 +953,8 @@ class llm_graph_result { // (a layer's conv state r_l and its recurrent state s_l): the chain runs through them too struct walk_state { ggml_tensor * cache = nullptr; // the memory's state tensor this entry belongs to - ggml_tensor * ds = nullptr; // GRAD leaf, zero-filled, on the state ENTERING the chunk + ggml_tensor * tap = nullptr; // the state ENTERING the chunk, tapped: its gradient adds into the + // walk's host accumulator for this state (ggml_opt_accumulate_grad) ggml_tensor * gs = nullptr; // input: the gradient later chunks put on the state it LEAVES }; std::vector t_walk_state; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 344d9bebab43..b2bab4505e5d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -150,7 +150,8 @@ endif () llama_build(test-recurrent-state-rollback.cpp) -# the exact walk trains on the window's whole gradient whatever the chunking (needs -m ) +# the exact walk trains on the window's whole gradient whatever the chunking: run on a pure-attention AND a +# hybrid model (-m), on a GPU backend as well as the CPU (a hybrid passed on the CPU while Metal tore its state) llama_build(test-walk-exact.cpp) # GGML_TENSOR_FLAG_GRAD: a leaf whose gradient the backward computes, never optimized (the walk reads the cached prefix K/V gradient) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 68140927d3a8..b53590be849e 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -10431,7 +10431,9 @@ static std::vector> make_test_cases_eval() { // packed qkv projection, several sequences (the atomic dq/dk accumulation across heads) for (auto [hc, hs, nt, ns, kd, kk, vr, st] : std::vector>{ {2, 16, 10, 1, true, 1, 3, false}, {2, 32, 7, 2, false, 2, 3, true}, {2, 64, 40, 2, true, 3, 3, true}, - {2, 128, 33, 1, false, 1, 3, true}, {3, 32, 17, 2, true, 1, 2, true}}) { + {2, 128, 33, 1, false, 1, 3, true}, {3, 32, 17, 2, true, 1, 2, true}, + // the 4B hybrid's own layer at one walk chunk: 16 key heads shared 2 ways, 512 tokens + {16, 128, 512, 1, false, 1, 2, true}}) { test_cases.emplace_back(new test_gated_delta_net_back(hc, hs, nt, ns, kd, kk, vr, st)); } test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 16, 1, 1)); diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 1724b8d18912..7e4dd1717a6e 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -22,12 +22,14 @@ #include "llama.h" #include "../src/llama-adapter.h" +#include #include #include #include #include #include #include +#include #include static const float LR = 1e-3f; // an SGD step is linear in the gradient at any lr @@ -122,6 +124,21 @@ static std::vector adapter_params(const llama_adapter_lora * adapter) { return out; } +// the name and length of each tensor adapter_params lays out, in the same order +static std::vector> adapter_layout(const llama_adapter_lora * adapter) { + std::map sorted; + for (const auto & [name, w] : adapter->ab_map) { + sorted[name] = &w; + } + std::vector> out; + for (const auto & [name, w] : sorted) { + out.emplace_back(name + ".a", (size_t) ggml_nelements(w->a)); + out.emplace_back(name + ".b", (size_t) ggml_nelements(w->b)); + } + return out; +} +static std::vector> g_layout; + // one SGD step from the fresh adapter: the step itself, -lr * the gradient it took static std::vector step_delta(const common_params & params, llama_model * model, const std::string & init, const std::vector & tokens, const std::vector & labelled, @@ -129,6 +146,7 @@ static std::vector step_delta(const common_params & params, llama_model * llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); GGML_ASSERT(adapter != nullptr); const std::vector before = adapter_params(adapter); + g_layout = adapter_layout(adapter); llama_context * ctx = make_ctx(params, model, chunk); float scale = 1.0f; @@ -193,6 +211,29 @@ static agreement compare(const char * what, const std::vector & a, const return r; } +// a step that strays names where: the tensors whose own cosine is lowest, so a wrong op +// shows as the layers (recurrent or attention, which depth) that carry it +static void worst_tensors(const std::vector & a, const std::vector & ref, int n) { + std::vector> rows; + size_t off = 0; + for (const auto & [name, len] : g_layout) { + double dot = 0.0, na = 0.0, nr = 0.0; + for (size_t i = off; i < off + len; ++i) { + dot += (double) a[i] * ref[i]; + na += (double) a[i] * a[i]; + nr += (double) ref[i] * ref[i]; + } + off += len; + if (na > 0.0 && nr > 0.0) { + rows.emplace_back(dot / std::sqrt(na * nr), name + " norm ratio " + std::to_string(std::sqrt(na / nr))); + } + } + std::sort(rows.begin(), rows.end()); + for (int i = 0; i < n && i < (int) rows.size(); ++i) { + printf(" worst: cosine %.4f %s\n", rows[i].first, rows[i].second.c_str()); + } +} + // Measured on the 1.5B Q4_K_M (5090): the exact walk agrees with one graph at cosine 0.998 and // norm 0.995 (what is left is the quantized matmuls' batch-size variance: a 128-row chunk and a // 512-row graph take different kernels); the plain walk sits at cosine 0.81, norm 2.1. @@ -212,7 +253,9 @@ int main(int argc, char ** argv) { llama_model * model = llama_model_load_from_file(params.model.path.c_str(), mparams); GGML_ASSERT(model != nullptr); - const std::string init = "test-walk-exact.init.gguf"; + // one file per run: two runs in one directory (a measurement beside an exactness check) + // overwrote each other's adapter and failed to load it (the M5, 2026-10-11) + const std::string init = "test-walk-exact." + std::to_string(std::chrono::steady_clock::now().time_since_epoch().count()) + ".init.gguf"; GGML_ASSERT(common_lora_write_fresh(model, init, 8, 16, { "attn_q", "attn_v" }, 42)); // words in a fixed pseudo-random order: a window the model cannot already predict, so its @@ -268,6 +311,7 @@ int main(int argc, char ** argv) { printf(" %6u %5u %12.1f %9.1f %5.2fx\n", w, chunk, mib[0], mib[1], mib[1] > 0 ? mib[0] / mib[1] : 0.0); fflush(stdout); } + std::remove(init.c_str()); return 0; } @@ -282,6 +326,7 @@ int main(int argc, char ** argv) { const auto plain = step_delta(params, model, init, tokens, all, WINDOW / 4, false); if (!same_step(compare("all labelled: exact walk in 4 chunks vs one graph", exact, ref))) { fprintf(stderr, "FAILED: the exact walk's step is not the window's gradient\n"); + worst_tensors(exact, ref, 24); ++failures; } if (same_step(compare("all labelled: plain walk in 4 chunks vs one graph", plain, ref))) { From 1c429b137e1ef9aac4aee95020cbb26458d834d0 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Sun, 11 Oct 2026 20:26:05 -0500 Subject: [PATCH 5/6] llama: training gets its own file; the graph type is one member Code motion, no behaviour change. llama-context.cpp held inference and training in 4,899 lines; the training section (opt_init, opt_epoch, the walk in opt_epoch_iter) and the llama_opt_* interface move verbatim to llama-context-opt.cpp (891 lines), leaving 3,992. The file-static ctx_type_to_graph_type, which both halves call, becomes the member llama_context::graph_type() instead of a second copy. M5 Metal, test-walk-exact built from this tree (diff f8f0195b29cf on 4da182d9a): - qwen3.5-4b hybrid Q4_K_M: exit 0, every row cosine 1.0000, host memory 182.8 to 82.2 MiB checkpointed, device graph 1851.1 MiB per chunk, 2026-10-12T01:25:16Z; the same numbers as 4da182d9a. - qwen2.5-0.5b pure attention: exit 0, 2026-10-12T01:25:51Z. Moved text checked line by line: every removed line is in the new file except the three ctx_type_to_graph_type call sites, now graph_type(). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- src/CMakeLists.txt | 1 + src/llama-context-opt.cpp | 904 +++++++++++++++++++++++++++++++++++++ src/llama-context.cpp | 906 +------------------------------------- src/llama-context.h | 3 + 4 files changed, 918 insertions(+), 896 deletions(-) create mode 100644 src/llama-context-opt.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 8922dc12adc7..4c1c95b856a9 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -16,6 +16,7 @@ add_library(llama llama-batch.cpp llama-chat.cpp llama-context.cpp + llama-context-opt.cpp llama-cparams.cpp llama-grammar.cpp llama-graph.cpp diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp new file mode 100644 index 000000000000..1e20995686ad --- /dev/null +++ b/src/llama-context-opt.cpp @@ -0,0 +1,904 @@ +// Training: llama_context's optimizer members (opt_init, opt_epoch, the walk in +// opt_epoch_iter) and the llama_opt_* interface, apart from inference so each reads as one concern. +#include "llama-context.h" + +#include "ggml.h" +#include "llama-batch.h" +#include "llama-impl.h" +#include "llama-memory-hybrid.h" +#include "llama-memory-recurrent.h" +#include "llama-kv-cache.h" +#include "llama-model.h" +#include "llama.h" + +#include +#include +#include +#include +#include + +// +// training +// + +// Returns whether the tensor was made a trainable parameter. +static bool llama_set_param(struct ggml_tensor * tensor, llama_opt_param_filter param_filter, void * userdata) { + if (!tensor || tensor->type != GGML_TYPE_F32) { + return false; + } + if (!param_filter(tensor, userdata)) { + return false; + } + if (strcmp(tensor->name, "token_embd.weight") == 0) { + return false; // FIXME + } + if (strcmp(tensor->name, "rope_freqs.weight") == 0) { + return false; // FIXME + } + ggml_set_param(tensor); + return true; +} + +void llama_context::opt_init(struct llama_model * model, struct llama_opt_params lopt_params) { + GGML_ASSERT(!opt_ctx); + // The training window lives on this context, not the shared model (a serving context may + // be running on the same weights). + opt_n_ctx_train = lopt_params.n_ctx_train > 0 ? lopt_params.n_ctx_train : n_ctx(); + opt_walk_exact = lopt_params.walk_exact; + opt_walk_horizon_req = lopt_params.walk_horizon; + opt_walk_horizon = lopt_params.walk_horizon; + const uint32_t n_batch = std::min(this->n_batch(), opt_n_ctx_train); + const uint32_t n_ubatch = std::min(this->n_ubatch(), n_batch); + GGML_ASSERT(opt_n_ctx_train % n_batch == 0); + GGML_ASSERT(n_batch % n_ubatch == 0); + // A training graph cannot backprop through the KV cache (SET_ROWS writes in place and + // attention reads the cache as a leaf), so attention consumes this chunk's K/V directly + // and the context before it from the cache as a constant (build_attn); opt_epoch_iter + // walks the window decoding context and training each labelled chunk. One chunk is one + // batch is one ubatch: a chunk's backward needs all of its own K/V in one graph. + GGML_ASSERT(n_ubatch == n_batch && "training needs one ubatch per batch: set -ub = -b"); + // In a flash-attention context the training graph reads attention through FLASH_ATTN_EXT too + // (build_attn_mha), whose backward is tiled over the window; flash attention also serves the + // walk's context decodes and lets the cache hold V un-transposed. + cparams.training = true; + // The graph result and the scheduler were sized at context creation: before training + // (no backward pass) and before any adapter attached since (its nodes uncounted). + // Size both again for the training graph, from the same number, plus the nodes each + // attention layer's tiled backward adds for this window (one count, owned by ggml). + { + size_t max_nodes = graph_max_nodes(n_ubatch); + if (cparams.flash_attn) { + max_nodes += (size_t) model->hparams.n_layer() * ggml_flash_attn_ext_back_n_nodes(opt_n_ctx_train, 0); + } + gf_res_prev.reset(new llm_graph_result(max_nodes)); + sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), + max_nodes, cparams.pipeline_parallel, cparams.op_offload)); + } + + ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY); + opt_params.opt_period = n_batch / n_ubatch; + opt_params.get_opt_pars = lopt_params.get_opt_pars; + // per-layer recompute: the layer outputs are the checkpoints (llm_graph names them l_out-N) + opt_params.checkpoint_prefix = lopt_params.recompute ? "l_out" : nullptr; + opt_params.get_opt_pars_ud = lopt_params.get_opt_pars_ud; + opt_params.optimizer = lopt_params.optimizer_type; + opt_ctx = ggml_opt_init(opt_params); + ggml_opt_set_alloc_budget(opt_ctx, opt_memory_budget); + + llama_opt_param_filter param_filter = lopt_params.param_filter; + void * param_filter_ud = lopt_params.param_filter_ud; + + if (lopt_params.adapter) { + // LoRA-only: the adapter's A/B are the parameters, the base is frozen. An adapter + // that is not attached is not in the graph, so it would train nothing while every + // epoch reported a loss: refuse it here instead. + const auto it = loras->find(lopt_params.adapter); + GGML_ASSERT(it != loras->end() && it->second != 0.0f && "the LoRA to train must be attached with a non-zero scale"); + int64_t n_params = 0; + for (auto & [name, ab] : lopt_params.adapter->ab_map) { + GGML_ASSERT(ab.a->type == GGML_TYPE_F32 && ab.b->type == GGML_TYPE_F32 && "a trainable LoRA holds F32 A/B tensors"); + n_params += llama_set_param(ab.a, param_filter, param_filter_ud); + n_params += llama_set_param(ab.b, param_filter, param_filter_ud); + } + GGML_ASSERT(n_params > 0 && "the param filter selected none of the adapter's tensors"); + LLAMA_LOG_INFO("%s: LoRA-only training: %lld adapter tensors trainable, base frozen\n", __func__, (long long) n_params); + return; + } + + //llama_set_param(model->tok_embd, param_filter, param_filter_ud); // FIXME + llama_set_param(model->type_embd, param_filter, param_filter_ud); + llama_set_param(model->pos_embd, param_filter, param_filter_ud); + llama_set_param(model->tok_norm, param_filter, param_filter_ud); + llama_set_param(model->tok_norm_b, param_filter, param_filter_ud); + llama_set_param(model->output_norm, param_filter, param_filter_ud); + llama_set_param(model->output_norm_b, param_filter, param_filter_ud); + llama_set_param(model->output, param_filter, param_filter_ud); + llama_set_param(model->output_b, param_filter, param_filter_ud); + llama_set_param(model->output_norm_enc, param_filter, param_filter_ud); + llama_set_param(model->cls, param_filter, param_filter_ud); + llama_set_param(model->cls_b, param_filter, param_filter_ud); + llama_set_param(model->cls_out, param_filter, param_filter_ud); + llama_set_param(model->cls_out_b, param_filter, param_filter_ud); + llama_set_param(model->cls_norm, param_filter, param_filter_ud); + + for (struct llama_layer & layer : model->layers) { + for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) { + llama_set_param(reinterpret_cast(&layer)[i], param_filter, param_filter_ud); + } + } +} + +void llama_context::opt_epoch_iter( + ggml_opt_dataset_t dataset, + ggml_opt_result_t result, + const std::vector & tokens, + const std::vector & labels_sparse, + llama_batch & batch, + ggml_opt_epoch_callback callback, + bool train, + int64_t idata_in_loop, + int64_t ndata_in_loop, + int64_t t_loop_start) { + GGML_ASSERT(opt_ctx); + const uint32_t n_ctx = opt_n_ctx_train; + const uint32_t n_batch = std::min(this->n_batch(), n_ctx); + + memory->clear(true); + + // THE WALK (Fable's design, continuum 2026-10-06): one pass over the window in one + // context. A run with no label is context she did not write: a plain decode into the + // cache, no backward. A labelled run is her reply: it trains in chunks of n_batch, each + // attending to everything before it read from the cache as a constant (build_attn, + // cparams.training), then the reply is decoded into the cache under the adapter as it + // now is, so the next reply sees it. Memory is chunk x window, never window x window, so + // her whole lived window trains. Nothing after the last label is computed. + int64_t last_label = -1; + for (int64_t i = (int64_t) n_ctx - 1; i >= 0; --i) { + if (labels_sparse[i] >= 0) { + last_label = i; + break; + } + } + if (last_label < 0) { + LLAMA_LOG_ERROR("%s: a training window with no labelled position: nothing to learn from it\n", __func__); + return; + } + + // THE YIELD POINT IS THE CHUNK, not the window: a turn that arrives mid-walk waits for one + // chunk, never for the whole example (measured on the 5090: one callback per example made a + // turn wait for its entire walk, 1.2-1.8 s on a 1.5B at 15k and 13.5 s on Fable's run). The + // first unit is the window's own boundary, which opt_epoch has already offered. + bool first_unit = true; + auto yield_point = [&]() -> bool { + if (first_unit) { + first_unit = false; + } else if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { + opt_stop_requested.store(true); + } + return !opt_stop_requested.load(std::memory_order_relaxed); + }; + + // a plain inference decode of [p0, p1) into the cache: the training graph's attention + // never writes the cache, so context and finished replies reach it this way + auto decode_span = [&](uint32_t p0, uint32_t p1) -> bool { + const bool training = cparams.training; + cparams.training = false; + gf_res_prev->reset(); + bool ok = true; + for (uint32_t c0 = p0; c0 < p1 && ok; c0 += n_batch) { + if (!yield_point()) { + cparams.training = training; + gf_res_prev->reset(); + return false; // stopped or cancelled at a chunk boundary: not a decode failure + } + const uint32_t c1 = std::min(c0 + n_batch, p1); + batch.n_tokens = c1 - c0; + for (uint32_t i = 0; i < c1 - c0; ++i) { + batch.token [i] = tokens[c0 + i]; + batch.pos [i] = c0 + i; + batch.n_seq_id[i] = 1; + batch.seq_id [i][0] = 0; + batch.logits [i] = i + 1 == c1 - c0; + } + const int rc = decode(batch); + ok = rc == 0; + } + cparams.training = training; + gf_res_prev->reset(); + if (!ok) { + LLAMA_LOG_ERROR("%s: the context decode of [%u, %u) failed\n", __func__, p0, p1); + } + return ok; + }; + + // THE EXACT WALK's state for the chunk being trained (null in the plain walk): its place in + // the window's one optimizer period, and the gradient accumulated on the window's cached K/V + struct exact_chunk { + uint32_t c0; + uint32_t n; + uint32_t n_labels; + uint32_t grad_from; // its gradient taps cover cached positions [grad_from, c0) + bool surrogate; // a later chunk's gradient sits on its own K/V + bool needed; + bool period_end; // the window's optimizer step + float loss_scale; // its labelled positions over the window's + }; + const exact_chunk * xc = nullptr; + // THE EXACT WALK's gradients, in host memory, each added into in place by the backward: + // per attention layer dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, + // position] F32 (null for a recurrent layer), which the graph reads through cparams (never + // outliving it, whatever way the walk returns); per recurrent state tensor (a layer's conv + // state r_l and its recurrent state s_l) dL/d(the state entering the chunk trained last), + // which is what the chunk before it puts on the state it leaves (state_ready once one is in) + struct walk_acc_t { + llama_cparams & cp; + std::vector k, v; + std::unordered_map state; + bool state_ready = false; + ggml_context_ptr ctx; + ggml_backend_buffer_ptr buf; + ~walk_acc_t() { cp.walk_acc_k = nullptr; cp.walk_acc_v = nullptr; } + } walk_acc { cparams, {}, {}, {}, false, nullptr, nullptr }; + + // one training chunk [pos_ctx, pos_ctx + n_tokens): forward and backward over it alone, + // attending to the cached context before it + auto train_chunk = [&](uint32_t pos_ctx, uint32_t n_chunk) -> bool { + batch.n_tokens = n_chunk; + for (uint32_t pos_batch = 0; pos_batch < n_chunk; ++pos_batch) { + batch.token [pos_batch] = tokens[pos_ctx + pos_batch]; + batch.pos [pos_batch] = pos_ctx + pos_batch; + batch.n_seq_id[pos_batch] = 1; + batch.seq_id [pos_batch][0] = 0; + // an exact walk's context chunk has no label but still trains (its K/V carry later + // chunks' gradient): one output row, weighted zero, so the graph has its loss + batch.logits [pos_batch] = labels_sparse[pos_ctx + pos_batch] >= 0 + || (xc != nullptr && xc->n_labels == 0 && pos_batch + 1 == n_chunk); + } + + // output_all = false: the batch's own logits flags decide which rows exist + if (!balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, false)) { + LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); + return false; + } + + const uint32_t n_tokens_all = balloc->get_n_tokens(); + + n_queued_tokens += n_tokens_all; + + embd_seq.clear(); + + const uint32_t n_outputs_all = balloc->get_n_outputs(); + // Unreachable: the walk below trains only chunks of a labelled run. It RETURNS from the + // whole window, since a later chunk would need this one in the cache. + if (n_outputs_all == 0) { + LLAMA_LOG_ERROR("%s: a training batch with no labelled position (every label masked): nothing to learn from it\n", __func__); + return false; + } + + auto mctx = memory->init_batch(*balloc, cparams.n_ubatch, true); + if (!mctx || mctx->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { + LLAMA_LOG_ERROR("%s: could not initialize batch\n", __func__); + return false; + } + + // reserve output buffer + if (output_reserve(n_outputs_all) < n_outputs_all) { + LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); + GGML_ABORT("TODO: handle this error"); + }; + + uint32_t pos_batch = 0; + do { + const auto & ubatch = mctx->get_ubatch(); + + // the labelled rows of this ubatch, as decode counts its outputs + { + uint32_t n_outputs_ubatch = 0; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + n_outputs_ubatch += (uint32_t) (ubatch.output[i] != 0); + } + n_outputs = n_outputs_ubatch; + } + + if (!mctx->apply()) { + LLAMA_LOG_ERROR("%s: failed to update the memory context\n", __func__); + return false; + } + + auto * res = gf_res_prev.get(); + + const auto gparams = graph_params(res, ubatch, mctx.get(), graph_type()); + + res->reset(); + + auto * gf = model.build_graph(gparams); + + struct ggml_context * ctx_compute_opt; + { + const size_t size_gf = ggml_graph_size(gf); + const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); + struct ggml_init_params params = { + /*.mem_size =*/ size_meta, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx_compute_opt = ggml_init(params); + } + ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); + if (xc != nullptr) { + ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); + for (const auto & io : res->t_walk) { + if (io.tk) { + ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); + ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); + } + } + for (const auto & ws : res->t_walk_state) { + if (ws.tap) { + ggml_opt_accumulate_grad(opt_ctx, ws.tap, walk_acc.state.at(ws.cache)); + } + } + } + if (!ggml_opt_alloc(opt_ctx, train)) { + const char * why = ggml_opt_refusal(opt_ctx); + opt_failure = why; + LLAMA_LOG_ERROR("%s: the %s graph was refused (window %u): %s: stopping the epoch\n", + __func__, train ? "training" : "evaluation", n_ctx, + why[0] ? why : "it does not fit in device memory"); + ggml_free(ctx_compute_opt); + opt_alloc_failed.store(true); + opt_stop_requested.store(true); + return false; + } + + res->set_inputs(&ubatch); + { + // one-hot rows for the labelled positions only, in the order the ubatch + // emits its outputs (the same order the logits rows take) + struct ggml_tensor * labels = ggml_opt_labels(opt_ctx); + GGML_ASSERT(labels->ne[1] == (int64_t) n_outputs); + ggml_set_zero(labels); + const float onef = 1.0f; + uint32_t row = 0; + for (uint32_t pos_ubatch = 0; pos_ubatch < ubatch.n_tokens; ++pos_ubatch) { + if (ubatch.output[pos_ubatch] == 0) { + continue; + } + const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch; + // an exact walk's context chunk: its one row has no label (weighted zero) + GGML_ASSERT((xc != nullptr && labels_sparse[ilabel] < 0) || (labels_sparse[ilabel] >= 0 && labels_sparse[ilabel] < labels->ne[0])); + if (labels_sparse[ilabel] >= 0) { + ggml_backend_tensor_set(labels, &onef, (row*labels->ne[0] + labels_sparse[ilabel])*sizeof(float), sizeof(float)); + } + ++row; + } + GGML_ASSERT(row == n_outputs); + } + if (xc != nullptr) { + // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the + // accumulator takes this chunk's own + for (const auto & ws : res->t_walk_state) { + ggml_tensor * acc = walk_acc.state.at(ws.cache); + if (ws.gs) { + GGML_ASSERT(ggml_nbytes(acc) == ggml_nbytes(ws.gs)); + ggml_backend_tensor_set(ws.gs, acc->data, 0, ggml_nbytes(ws.gs)); + } + memset(acc->data, 0, ggml_nbytes(acc)); + } + } + ggml_opt_eval(opt_ctx, xc != nullptr && xc->n_labels == 0 ? nullptr : result); + if (const char * why = ggml_opt_refusal(opt_ctx); why[0] != '\0') { + // the backend failed the graph (ggml_opt_eval): the run fails, as a refused graph does + opt_failure = why; + LLAMA_LOG_ERROR("%s: %s: stopping the epoch\n", __func__, why); + ggml_free(ctx_compute_opt); + opt_alloc_failed.store(true); + opt_stop_requested.store(true); + return false; + } + if (xc != nullptr && !res->t_walk_state.empty()) { + walk_acc.state_ready = true; + } + if (callback && !(xc != nullptr && xc->n_labels == 0)) { + callback(train, opt_ctx, dataset, result, idata_in_loop + (pos_ctx + pos_batch)/n_batch + 1, ndata_in_loop, t_loop_start); + } + ggml_free(ctx_compute_opt); + + pos_batch += ubatch.n_tokens; + } while (mctx->next()); + return true; + }; + + // a recurrent part (pure recurrent or hybrid) is snapshotted around each training chunk; + // a hybrid's attention part removes the chunk's cells like a plain cache + llama_memory_recurrent * recr = nullptr; + llama_kv_cache * attn = nullptr; + if (auto * hybrid = dynamic_cast(memory.get())) { + recr = hybrid->get_mem_recr(); + attn = hybrid->get_mem_attn(); + } else { + recr = dynamic_cast(memory.get()); + } + // the plain walk snapshots a recurrent state into a scratch sequence; the exact walk keeps host snapshots + if (recr != nullptr && cparams.n_seq_max < 2 && !(train && opt_walk_exact)) { + LLAMA_LOG_ERROR("%s: a recurrent model trains with a scratch sequence for its state snapshot: create the training context with n_seq_max >= 2\n", __func__); + opt_stop_requested.store(true); + return; + } + + if (train && opt_walk_exact) { + // THE EXACT WALK. The plain walk below steps per chunk and stops the gradient at every + // chunk boundary: a reply learns only through its own chunk's K/V, never through what + // an earlier chunk computed for it. Here the window is decoded once under the adapter + // as it is, then trained in REVERSE: each chunk's graph attends to the cache, with its + // last walk_horizon prefix positions behind a GRAD leaf, and adds the surrogate + // + on its own K/V, where G is what every later chunk put there. Its + // backward then carries that gradient through itself into the adapter and into its own + // prefix's G. ONE optimizer step per window, on the gradient of the window's loss + // (exact within the horizon). Memory stays chunk x window; the cost is one decode plus + // one forward and backward per chunk, context chunks included. + // THE INVARIANT THAT MAKES THE REVERSE PASS VALID (Fable): the adapter does not change + // between the forward decode and the last chunk's backward. One optimizer step per window, + // taken by the chunk the reverse pass trains last, is what keeps each chunk's recomputed + // K/V (and recurrent state) equal to the cached values its successors attended to; a step + // inside the window would make every surrogate a gradient of a different function. + // + // Each window starts from the REQUESTED horizon: a halving that fit one window's device + // budget must not shrink every later window and epoch, which may be shorter (Fable). + opt_walk_horizon = opt_walk_horizon_req; + // A recurrent context with rollback slots keeps several states per sequence; the walk + // carries exactly one. Refused by name, never asserted: training runs inside the serving + // process, and an assert there takes serving down (2026-10-07 04:09Z on the 5090). + if (recr != nullptr && cparams.n_rs_seq > 0) { + opt_failure = "the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has " + + std::to_string(cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0"; + LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); + opt_alloc_failed.store(true); + opt_stop_requested.store(true); + return; + } + // A recurrent model's state is carried the same way, one chunk at a time: the state + // entering a chunk is tapped into a host accumulator, the state it leaves meets the next chunk's gradient + // (build_rs, build_walk_state_exit). The reverse pass restores each chunk's entry state + // from a host snapshot taken at its boundary in the forward pass. + std::vector chunks; + uint32_t n_labels_window = 0; + for (uint32_t p = 0; p <= (uint32_t) last_label; ) { + const bool labelled = labels_sparse[p] >= 0; + uint32_t end = p; + while (end <= (uint32_t) last_label && (labels_sparse[end] >= 0) == labelled) { + ++end; + } + for (uint32_t c0 = p; c0 < end; c0 += n_batch) { + const uint32_t c1 = std::min(c0 + n_batch, end); + chunks.push_back({ c0, c1 - c0, labelled ? c1 - c0 : 0, 0, false, false, false, 0.0f }); + n_labels_window += labelled ? c1 - c0 : 0; + } + p = end; + } + // which chunks train, and how far back each one's gradient reaches + uint32_t reach = UINT32_MAX; // the lowest position a later trained chunk's gradient reaches + int64_t first = -1; // the last chunk the reverse pass trains: the step + for (int64_t j = (int64_t) chunks.size() - 1; j >= 0; --j) { + exact_chunk & c = chunks[j]; + c.surrogate = reach < c.c0 + c.n; + // a recurrent state reaches every later chunk: a chunk before a trained one trains + c.needed = c.n_labels > 0 || c.surrogate || (recr != nullptr && first >= 0); + if (!c.needed) { + continue; + } + c.grad_from = opt_walk_horizon == 0 || c.c0 < opt_walk_horizon ? 0 : c.c0 - opt_walk_horizon; + c.loss_scale = (float) c.n_labels / (float) n_labels_window; + reach = std::min(reach, c.grad_from); + first = j; + } + GGML_ASSERT(first >= 0); + chunks[first].period_end = true; + + // THE HOST BUDGET (Fable): the accumulators are layers x positions x (k + v) floats, and a + // recurrent state is checkpointed every `stride` chunks (one sequence's state rows each; + // the first chunk's entry is the zero state and costs nothing). Both are known before + // anything runs. Under a budget the stride is the smallest that fits: the reverse pass + // rebuilds a chunk's entry state by decoding forward from its checkpoint, so a longer + // stride trades host memory for decode, never for exactness. + size_t stride = 1; + { + size_t kv_bytes = 0; + for (uint32_t il = 0; il < model.hparams.n_layer(); ++il) { + if (!model.hparams.is_recr(il)) { + kv_bytes += (size_t) (last_label + 1)*(model.hparams.n_embd_k_gqa(il) + model.hparams.n_embd_v_gqa(il))*sizeof(float); + } + } + size_t state = 0; + if (recr != nullptr) { + for (size_t il = 0; il < recr->r_l.size(); ++il) { + state += recr->r_l[il] ? recr->r_l[il]->nb[1] : 0; + state += recr->s_l[il] ? recr->s_l[il]->nb[1] : 0; + } + } + auto bytes_at = [&](size_t s) { return kv_bytes + ((chunks.size() + s - 1)/s - 1)*state; }; + if (opt_walk_host_budget > 0) { + while (stride < chunks.size() && bytes_at(stride) > opt_walk_host_budget) { + ++stride; + } + } + const size_t bytes = bytes_at(stride); + opt_walk_host_bytes = std::max(opt_walk_host_bytes, bytes); + if (opt_walk_host_budget > 0 && bytes > opt_walk_host_budget) { + char why[512]; + snprintf(why, sizeof(why), + "the exact walk needs %.1f MiB of host memory for this window (%lld positions, the K/V gradient alone), over its budget of %.1f MiB: " + "the K/V gradient grows with the window; neither the horizon nor the chunk size changes it", + bytes/1048576.0, (long long) (last_label + 1), opt_walk_host_budget/1048576.0); + opt_failure = why; + LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); + opt_alloc_failed.store(true); + opt_stop_requested.store(true); + return; + } + if (stride > 1) { + LLAMA_LOG_INFO("%s: the exact walk checkpoints the recurrent state every %zu chunks to fit %.1f MiB of host memory (%.1f MiB)\n", + __func__, stride, opt_walk_host_budget/1048576.0, bytes/1048576.0); + } + } + { + const uint32_t n_layer = model.hparams.n_layer(); + walk_acc.k.assign(n_layer, nullptr); + walk_acc.v.assign(n_layer, nullptr); + walk_acc.ctx.reset(ggml_init({ 4*n_layer*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true })); + walk_acc.state.clear(); + walk_acc.state_ready = false; + for (uint32_t il = 0; il < n_layer; ++il) { + if (model.hparams.is_recr(il)) { + // a recurrent layer has no K/V: its state carries the chain, one sequence's row + if (recr != nullptr) { + walk_acc.state[recr->r_l[il]] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_r(), 1); + walk_acc.state[recr->s_l[il]] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_s(), 1); + } + continue; + } + walk_acc.k[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_k_gqa(il), last_label + 1); + walk_acc.v[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_v_gqa(il), last_label + 1); + } + walk_acc.buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(walk_acc.ctx.get(), ggml_backend_cpu_buffer_type())); + if (!walk_acc.buf) { + opt_failure = "the exact walk's gradient accumulator (host memory) could not be allocated"; + LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); + opt_alloc_failed.store(true); + opt_stop_requested.store(true); + return; + } + ggml_backend_buffer_clear(walk_acc.buf.get(), 0); + cparams.walk_acc_k = walk_acc.k.data(); + cparams.walk_acc_v = walk_acc.v.data(); + } + + // the window as the adapter now reads it, chunk by chunk: a recurrent model's state is + // checkpointed at each chunk's start through the sequence-state API (the recurrent part + // only, PARTIAL_ONLY), after waiting for the queued decode; a restore puts back the + // cell's position with its rows, and an empty checkpoint (before the first chunk) + // restores the zero state. A hand-rolled per-layer copy here skipped that wait and, on + // Metal, carried a torn mid-graph state into the reverse pass (2026-10-11). + std::vector> snaps(recr != nullptr ? chunks.size() : 0); + const llama_state_seq_flags state_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; + auto save_state = [&](std::vector & snap) { + synchronize(); + snap.resize(state_seq_get_size(0, state_flags)); + GGML_ASSERT(state_seq_get_data(0, snap.data(), snap.size(), state_flags) == snap.size()); + }; + auto load_state = [&](const std::vector & snap) { + synchronize(); + GGML_ASSERT(state_seq_set_data(0, snap.data(), snap.size(), state_flags) == snap.size()); + }; + for (size_t j = 0; j < chunks.size(); ++j) { + if (recr != nullptr && j % stride == 0) { + save_state(snaps[j]); + } + if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n)) { + return; + } + } + bool ok = true; + for (int64_t j = (int64_t) chunks.size() - 1; j >= first && ok; --j) { + const exact_chunk & c = chunks[j]; + if (!c.needed) { + continue; + } + if (!yield_point()) { + ok = false; + break; + } + // attention: pop the cache to [0, c0); recurrent: restore the state at the chunk's + // checkpoint and decode forward to its start (that decode rebuilds the attention K/V + // in between too, under the same adapter: the same values). Run again before every + // retry below: a refused chunk graph is refused AFTER its ubatch was applied, so the + // chunk's cells already sit in the cache and the recurrent cell's position has moved + // to the chunk's end; retrying on top of that fails to prepare the ubatch at all + // (the 5090, 2026-10-10: "init_batch: failed to prepare attention ubatches" after + // "retrying at 512", and the epoch stopped instead of shrinking the horizon). + const auto rewind = [&]() -> bool { + const size_t cp = recr != nullptr ? (size_t) j - (size_t) j % stride : (size_t) j; + const uint32_t p_pop = chunks[cp].c0; + const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : memory->seq_rm(0, p_pop, -1); + if (!popped) { + LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u) for the reverse pass\n", __func__, p_pop); + opt_stop_requested.store(true); + return false; + } + if (recr != nullptr) { + load_state(snaps[cp]); + if (p_pop < c.c0 && !decode_span(p_pop, c.c0)) { + return false; + } + } + return true; + }; + if (!rewind()) { + ok = false; + break; + } + cparams.walk_exact = true; + cparams.walk_surrogate = c.surrogate; + cparams.walk_state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && walk_acc.state_ready; + // THE HORIZON FITS THE DEVICE. Its GRAD leaves are the chunk graph's largest growth + // with the window (layers x horizon x (k + v) floats, twice with their gradient); the + // graph preflight refuses a chunk that does not fit before allocating anything, so a + // memory refusal halves the horizon (whole chunks) and the chunk is tried again; every + // later chunk keeps the smaller one. The forward still attends to the whole window and + // the recurrent state is carried exactly; only attention gradient past the horizon is + // dropped. A refusal at a horizon of one chunk, or for a node the device cannot run, + // stands. + exact_chunk cc = c; + for (;;) { + cc.grad_from = opt_walk_horizon == 0 || cc.c0 < opt_walk_horizon ? 0 : cc.c0 - opt_walk_horizon; + cparams.walk_grad_from = cc.grad_from; + xc = &cc; + ok = train_chunk(cc.c0, cc.n); + xc = nullptr; + const uint32_t span = cc.c0 - cc.grad_from; + if (ok || !opt_failure.empty() || !opt_alloc_failed.load() || span <= n_batch) { + break; + } + opt_walk_horizon = std::max(n_batch, (span/2)/n_batch*n_batch); + LLAMA_LOG_WARN("%s: the chunk at %u did not fit with a gradient horizon of %u positions: retrying at %u\n", + __func__, cc.c0, span, opt_walk_horizon); + opt_alloc_failed.store(false); + opt_stop_requested.store(false); + // the rewind's decode is a plain forward, as before the first try: no walk flags. + // walk_grad_from is not restored here because the loop's first statements recompute + // it from the shrunken horizon (cc.grad_from, then cparams.walk_grad_from) before the + // retried train_chunk runs. + const bool surrogate = cparams.walk_surrogate; + const bool state_surrogate = cparams.walk_state_surrogate; + cparams.walk_exact = false; + cparams.walk_grad_from = 0; + cparams.walk_surrogate = false; + cparams.walk_state_surrogate = false; + const bool rewound = rewind(); + cparams.walk_exact = true; + cparams.walk_surrogate = surrogate; + cparams.walk_state_surrogate = state_surrogate; + if (!rewound) { + ok = false; + break; + } + } + opt_walk_horizon_used = opt_walk_horizon; + cparams.walk_exact = false; + cparams.walk_grad_from = 0; + cparams.walk_surrogate = false; + cparams.walk_state_surrogate = false; + } + walk_acc.state_ready = false; + return; + } + + uint32_t pos = 0; + while (pos <= (uint32_t) last_label && !opt_stop_requested.load(std::memory_order_relaxed)) { + const bool labelled = labels_sparse[pos] >= 0; + uint32_t end = pos; + while (end <= (uint32_t) last_label && (labels_sparse[end] >= 0) == labelled) { + ++end; + } + if (!labelled) { + if (!decode_span(pos, end)) { + return; + } + } else { + for (uint32_t c0 = pos; c0 < end; c0 += n_batch) { + const uint32_t c1 = std::min(c0 + n_batch, end); + // The training forward advances a recurrent state in place, and a recurrent + // memory cannot remove a partial range: snapshot the state to the scratch + // sequence first (Fable), restore it after, then the decode below advances it + // from where the context left it. The state is per sequence and small. + const bool snapshot = recr != nullptr && c1 <= (uint32_t) last_label; + if (!yield_point()) { + return; + } + if (snapshot) { + recr->seq_cp(0, 1, -1, -1); + } + if (!train_chunk(c0, c1 - c0)) { + return; + } + if (snapshot) { + recr->seq_rm(0, -1, -1); + recr->seq_cp(1, 0, -1, -1); + recr->seq_rm(1, -1, -1); + } + // the chunk joins the context under the adapter as it now is: the training + // graph never writes the cache, so its cells hold nothing until this decode + if (c1 <= (uint32_t) last_label) { + // the chunk's cells hold no K/V (the training graph never writes the + // cache): drop them from the attention memory, then decode them for real + const bool removed = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, c0, -1)) : memory->seq_rm(0, c0, -1); + if (!removed) { + LLAMA_LOG_ERROR("%s: could not remove the trained chunk [%u, %u) from the cache before decoding it\n", __func__, c0, c1); + opt_stop_requested.store(true); + return; + } + if (!decode_span(c0, c1)) { + opt_stop_requested.store(true); + return; + } + } + } + } + pos = end; + } +} + +void llama_context::opt_epoch( + ggml_opt_dataset_t dataset, + ggml_opt_result_t result_train, + ggml_opt_result_t result_eval, + int64_t idata_split, + ggml_opt_epoch_callback callback_train, + ggml_opt_epoch_callback callback_eval) { + const uint32_t n_ctx = this->n_ctx(); + const uint32_t n_batch = std::min(cparams.n_batch, n_ctx); + const uint32_t n_ubatch = std::min(cparams.n_ubatch, n_batch); + const int64_t ndata = ggml_opt_dataset_ndata(dataset); + + GGML_ASSERT(idata_split >= 0); + GGML_ASSERT(idata_split <= ndata); + + const uint32_t ubatch_per_ctx = n_ctx / n_ubatch; + + struct llama_batch batch = llama_batch_init(n_batch, 0, 1); + std::vector tokens(n_ctx); + std::vector labels_sparse(n_ctx); + + int64_t idata = 0; + + int64_t t_loop_start = ggml_time_us(); + int64_t ndata_in_loop = idata_split*ubatch_per_ctx; + for (; idata < idata_split; ++idata) { + if (opt_stop_requested.load(std::memory_order_relaxed)) { + break; + } + constexpr bool train = true; + const int64_t idata_in_loop = idata*ubatch_per_ctx; + + if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { + opt_stop_requested.store(true); + } + if (opt_stop_requested.load(std::memory_order_relaxed)) { + break; + } + + ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); + opt_epoch_iter(dataset, result_train, tokens, labels_sparse, batch, + callback_train, train, idata_in_loop, ndata_in_loop, t_loop_start); + } + + t_loop_start = ggml_time_us(); + ndata_in_loop = (ndata - idata_split)*ubatch_per_ctx; + for (; idata < ndata; ++idata) { + if (opt_stop_requested.load(std::memory_order_relaxed)) { + break; + } + constexpr bool train = false; + const int64_t idata_in_loop = (idata - idata_split)*ubatch_per_ctx; + + if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { + opt_stop_requested.store(true); + } + if (opt_stop_requested.load(std::memory_order_relaxed)) { + break; + } + + ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); + opt_epoch_iter(dataset, result_eval, tokens, labels_sparse, batch, + callback_eval, train, idata_in_loop, ndata_in_loop, t_loop_start); + } + + llama_batch_free(batch); +} + +// +// training +// + +bool llama_opt_param_filter_all(const struct ggml_tensor * tensor, void * userdata) { + GGML_UNUSED(tensor); + GGML_UNUSED(userdata); + return true; +} + +void llama_opt_init(struct llama_context * ctx, struct llama_model * model, struct llama_opt_params lopt_params) { + ctx->opt_alloc_failed.store(false); + ctx->opt_failure.clear(); + ctx->opt_init(model, lopt_params); +} + +void llama_opt_stop(struct llama_context * ctx, bool stop) { + ctx->opt_stop_requested.store(stop, std::memory_order_relaxed); +} + +void llama_opt_set_step_callback(struct llama_context * ctx, llama_opt_step_callback callback, void * user_data) { + ctx->opt_step_callback = callback; + ctx->opt_step_callback_data = user_data; +} + +void llama_opt_set_walk_host_budget(struct llama_context * ctx, size_t bytes) { + ctx->opt_walk_host_budget = bytes; +} + +size_t llama_opt_walk_host_bytes(struct llama_context * ctx) { + return ctx->opt_walk_host_bytes; +} + +uint32_t llama_opt_walk_horizon(struct llama_context * ctx) { + return ctx->opt_walk_horizon_used; +} + +void llama_opt_set_memory_budget(struct llama_context * ctx, size_t bytes) { + ctx->opt_memory_budget = bytes; +} + +size_t llama_context::opt_graph_bytes() const { + return opt_ctx ? ggml_opt_peak_graph_bytes(opt_ctx) : 0; +} + +size_t llama_opt_graph_bytes(struct llama_context * ctx) { + return ctx->opt_graph_bytes(); +} + +size_t llama_opt_context_device_bytes(struct llama_context * ctx) { + size_t bytes = 0; + for (const auto & [buft, data] : ctx->memory_breakdown()) { + // data.context only: the compute buffers are where the training graph is allocated + // (the optimizer shares this context's scheduler, whose allocator replaces them with the + // graph's), so the graph's measured size already covers them (Cormac on #52) + if (!ggml_backend_buft_is_host(buft)) { + bytes += data.context; + } + } + return bytes; +} + +bool llama_opt_failed(struct llama_context * ctx) { + return ctx->opt_alloc_failed.load(); +} + +const char * llama_opt_failure(struct llama_context * ctx) { + return ctx->opt_failure.c_str(); +} + +void llama_opt_epoch( + struct llama_context * ctx, + ggml_opt_dataset_t dataset, + ggml_opt_result_t result_train, + ggml_opt_result_t result_eval, + int64_t idata_split, + ggml_opt_epoch_callback callback_train, + ggml_opt_epoch_callback callback_eval) { + ctx->opt_epoch( + dataset, + result_train, + result_eval, + idata_split, + callback_train, + callback_eval); +} diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 6ed7cb46b2ff..a296f0c2cbdb 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -27,14 +27,6 @@ // llama_context // -static llm_graph_type ctx_type_to_graph_type(llama_context_type ctx_type) { - switch (ctx_type) { - case LLAMA_CONTEXT_TYPE_DEFAULT: return LLM_GRAPH_TYPE_DEFAULT; - case LLAMA_CONTEXT_TYPE_MTP : return LLM_GRAPH_TYPE_DECODER_MTP; - } - throw std::runtime_error("Unsupported ctx type"); -} - struct llm_fused_op_probe { llm_fused_op op; const char * name; @@ -1823,7 +1815,7 @@ int llama_context::decode(const llama_batch & batch_inp) { ggml_status status; - const auto * res = process_ubatch(ubatch, ctx_type_to_graph_type(cparams.ctx_type), mctx.get(), status); + const auto * res = process_ubatch(ubatch, graph_type(), mctx.get(), status); if (!res) { // the last ubatch failed or was aborted -> remove all positions of that ubatch from the memory module @@ -2448,7 +2440,7 @@ ggml_cgraph * llama_context::graph_reserve( auto * res = gf_res_reserve.get(); - const auto gparams = graph_params(res, ubatch, mctx, ctx_type_to_graph_type(cparams.ctx_type)); + const auto gparams = graph_params(res, ubatch, mctx, graph_type()); res->reset(); @@ -2472,6 +2464,14 @@ ggml_cgraph * llama_context::graph_reserve( return gf; } +llm_graph_type llama_context::graph_type() const { + switch (cparams.ctx_type) { + case LLAMA_CONTEXT_TYPE_DEFAULT: return LLM_GRAPH_TYPE_DEFAULT; + case LLAMA_CONTEXT_TYPE_MTP : return LLM_GRAPH_TYPE_DECODER_MTP; + } + throw std::runtime_error("Unsupported ctx type"); +} + llm_graph_params llama_context::graph_params( llm_graph_result * res, const llama_ubatch & ubatch, @@ -3308,805 +3308,6 @@ llama_memory_breakdown llama_context::memory_breakdown() const { return ret; } -// -// training -// - -// Returns whether the tensor was made a trainable parameter. -static bool llama_set_param(struct ggml_tensor * tensor, llama_opt_param_filter param_filter, void * userdata) { - if (!tensor || tensor->type != GGML_TYPE_F32) { - return false; - } - if (!param_filter(tensor, userdata)) { - return false; - } - if (strcmp(tensor->name, "token_embd.weight") == 0) { - return false; // FIXME - } - if (strcmp(tensor->name, "rope_freqs.weight") == 0) { - return false; // FIXME - } - ggml_set_param(tensor); - return true; -} - -void llama_context::opt_init(struct llama_model * model, struct llama_opt_params lopt_params) { - GGML_ASSERT(!opt_ctx); - // The training window lives on this context, not the shared model (a serving context may - // be running on the same weights). - opt_n_ctx_train = lopt_params.n_ctx_train > 0 ? lopt_params.n_ctx_train : n_ctx(); - opt_walk_exact = lopt_params.walk_exact; - opt_walk_horizon_req = lopt_params.walk_horizon; - opt_walk_horizon = lopt_params.walk_horizon; - const uint32_t n_batch = std::min(this->n_batch(), opt_n_ctx_train); - const uint32_t n_ubatch = std::min(this->n_ubatch(), n_batch); - GGML_ASSERT(opt_n_ctx_train % n_batch == 0); - GGML_ASSERT(n_batch % n_ubatch == 0); - // A training graph cannot backprop through the KV cache (SET_ROWS writes in place and - // attention reads the cache as a leaf), so attention consumes this chunk's K/V directly - // and the context before it from the cache as a constant (build_attn); opt_epoch_iter - // walks the window decoding context and training each labelled chunk. One chunk is one - // batch is one ubatch: a chunk's backward needs all of its own K/V in one graph. - GGML_ASSERT(n_ubatch == n_batch && "training needs one ubatch per batch: set -ub = -b"); - // In a flash-attention context the training graph reads attention through FLASH_ATTN_EXT too - // (build_attn_mha), whose backward is tiled over the window; flash attention also serves the - // walk's context decodes and lets the cache hold V un-transposed. - cparams.training = true; - // The graph result and the scheduler were sized at context creation: before training - // (no backward pass) and before any adapter attached since (its nodes uncounted). - // Size both again for the training graph, from the same number, plus the nodes each - // attention layer's tiled backward adds for this window (one count, owned by ggml). - { - size_t max_nodes = graph_max_nodes(n_ubatch); - if (cparams.flash_attn) { - max_nodes += (size_t) model->hparams.n_layer() * ggml_flash_attn_ext_back_n_nodes(opt_n_ctx_train, 0); - } - gf_res_prev.reset(new llm_graph_result(max_nodes)); - sched.reset(ggml_backend_sched_new(backend_ptrs.data(), backend_buft.data(), backend_ptrs.size(), - max_nodes, cparams.pipeline_parallel, cparams.op_offload)); - } - - ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY); - opt_params.opt_period = n_batch / n_ubatch; - opt_params.get_opt_pars = lopt_params.get_opt_pars; - // per-layer recompute: the layer outputs are the checkpoints (llm_graph names them l_out-N) - opt_params.checkpoint_prefix = lopt_params.recompute ? "l_out" : nullptr; - opt_params.get_opt_pars_ud = lopt_params.get_opt_pars_ud; - opt_params.optimizer = lopt_params.optimizer_type; - opt_ctx = ggml_opt_init(opt_params); - ggml_opt_set_alloc_budget(opt_ctx, opt_memory_budget); - - llama_opt_param_filter param_filter = lopt_params.param_filter; - void * param_filter_ud = lopt_params.param_filter_ud; - - if (lopt_params.adapter) { - // LoRA-only: the adapter's A/B are the parameters, the base is frozen. An adapter - // that is not attached is not in the graph, so it would train nothing while every - // epoch reported a loss: refuse it here instead. - const auto it = loras->find(lopt_params.adapter); - GGML_ASSERT(it != loras->end() && it->second != 0.0f && "the LoRA to train must be attached with a non-zero scale"); - int64_t n_params = 0; - for (auto & [name, ab] : lopt_params.adapter->ab_map) { - GGML_ASSERT(ab.a->type == GGML_TYPE_F32 && ab.b->type == GGML_TYPE_F32 && "a trainable LoRA holds F32 A/B tensors"); - n_params += llama_set_param(ab.a, param_filter, param_filter_ud); - n_params += llama_set_param(ab.b, param_filter, param_filter_ud); - } - GGML_ASSERT(n_params > 0 && "the param filter selected none of the adapter's tensors"); - LLAMA_LOG_INFO("%s: LoRA-only training: %lld adapter tensors trainable, base frozen\n", __func__, (long long) n_params); - return; - } - - //llama_set_param(model->tok_embd, param_filter, param_filter_ud); // FIXME - llama_set_param(model->type_embd, param_filter, param_filter_ud); - llama_set_param(model->pos_embd, param_filter, param_filter_ud); - llama_set_param(model->tok_norm, param_filter, param_filter_ud); - llama_set_param(model->tok_norm_b, param_filter, param_filter_ud); - llama_set_param(model->output_norm, param_filter, param_filter_ud); - llama_set_param(model->output_norm_b, param_filter, param_filter_ud); - llama_set_param(model->output, param_filter, param_filter_ud); - llama_set_param(model->output_b, param_filter, param_filter_ud); - llama_set_param(model->output_norm_enc, param_filter, param_filter_ud); - llama_set_param(model->cls, param_filter, param_filter_ud); - llama_set_param(model->cls_b, param_filter, param_filter_ud); - llama_set_param(model->cls_out, param_filter, param_filter_ud); - llama_set_param(model->cls_out_b, param_filter, param_filter_ud); - llama_set_param(model->cls_norm, param_filter, param_filter_ud); - - for (struct llama_layer & layer : model->layers) { - for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) { - llama_set_param(reinterpret_cast(&layer)[i], param_filter, param_filter_ud); - } - } -} - -void llama_context::opt_epoch_iter( - ggml_opt_dataset_t dataset, - ggml_opt_result_t result, - const std::vector & tokens, - const std::vector & labels_sparse, - llama_batch & batch, - ggml_opt_epoch_callback callback, - bool train, - int64_t idata_in_loop, - int64_t ndata_in_loop, - int64_t t_loop_start) { - GGML_ASSERT(opt_ctx); - const uint32_t n_ctx = opt_n_ctx_train; - const uint32_t n_batch = std::min(this->n_batch(), n_ctx); - - memory->clear(true); - - // THE WALK (Fable's design, continuum 2026-10-06): one pass over the window in one - // context. A run with no label is context she did not write: a plain decode into the - // cache, no backward. A labelled run is her reply: it trains in chunks of n_batch, each - // attending to everything before it read from the cache as a constant (build_attn, - // cparams.training), then the reply is decoded into the cache under the adapter as it - // now is, so the next reply sees it. Memory is chunk x window, never window x window, so - // her whole lived window trains. Nothing after the last label is computed. - int64_t last_label = -1; - for (int64_t i = (int64_t) n_ctx - 1; i >= 0; --i) { - if (labels_sparse[i] >= 0) { - last_label = i; - break; - } - } - if (last_label < 0) { - LLAMA_LOG_ERROR("%s: a training window with no labelled position: nothing to learn from it\n", __func__); - return; - } - - // THE YIELD POINT IS THE CHUNK, not the window: a turn that arrives mid-walk waits for one - // chunk, never for the whole example (measured on the 5090: one callback per example made a - // turn wait for its entire walk, 1.2-1.8 s on a 1.5B at 15k and 13.5 s on Fable's run). The - // first unit is the window's own boundary, which opt_epoch has already offered. - bool first_unit = true; - auto yield_point = [&]() -> bool { - if (first_unit) { - first_unit = false; - } else if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { - opt_stop_requested.store(true); - } - return !opt_stop_requested.load(std::memory_order_relaxed); - }; - - // a plain inference decode of [p0, p1) into the cache: the training graph's attention - // never writes the cache, so context and finished replies reach it this way - auto decode_span = [&](uint32_t p0, uint32_t p1) -> bool { - const bool training = cparams.training; - cparams.training = false; - gf_res_prev->reset(); - bool ok = true; - for (uint32_t c0 = p0; c0 < p1 && ok; c0 += n_batch) { - if (!yield_point()) { - cparams.training = training; - gf_res_prev->reset(); - return false; // stopped or cancelled at a chunk boundary: not a decode failure - } - const uint32_t c1 = std::min(c0 + n_batch, p1); - batch.n_tokens = c1 - c0; - for (uint32_t i = 0; i < c1 - c0; ++i) { - batch.token [i] = tokens[c0 + i]; - batch.pos [i] = c0 + i; - batch.n_seq_id[i] = 1; - batch.seq_id [i][0] = 0; - batch.logits [i] = i + 1 == c1 - c0; - } - const int rc = decode(batch); - ok = rc == 0; - } - cparams.training = training; - gf_res_prev->reset(); - if (!ok) { - LLAMA_LOG_ERROR("%s: the context decode of [%u, %u) failed\n", __func__, p0, p1); - } - return ok; - }; - - // THE EXACT WALK's state for the chunk being trained (null in the plain walk): its place in - // the window's one optimizer period, and the gradient accumulated on the window's cached K/V - struct exact_chunk { - uint32_t c0; - uint32_t n; - uint32_t n_labels; - uint32_t grad_from; // its gradient taps cover cached positions [grad_from, c0) - bool surrogate; // a later chunk's gradient sits on its own K/V - bool needed; - bool period_end; // the window's optimizer step - float loss_scale; // its labelled positions over the window's - }; - const exact_chunk * xc = nullptr; - // THE EXACT WALK's gradients, in host memory, each added into in place by the backward: - // per attention layer dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, - // position] F32 (null for a recurrent layer), which the graph reads through cparams (never - // outliving it, whatever way the walk returns); per recurrent state tensor (a layer's conv - // state r_l and its recurrent state s_l) dL/d(the state entering the chunk trained last), - // which is what the chunk before it puts on the state it leaves (state_ready once one is in) - struct walk_acc_t { - llama_cparams & cp; - std::vector k, v; - std::unordered_map state; - bool state_ready = false; - ggml_context_ptr ctx; - ggml_backend_buffer_ptr buf; - ~walk_acc_t() { cp.walk_acc_k = nullptr; cp.walk_acc_v = nullptr; } - } walk_acc { cparams, {}, {}, {}, false, nullptr, nullptr }; - - // one training chunk [pos_ctx, pos_ctx + n_tokens): forward and backward over it alone, - // attending to the cached context before it - auto train_chunk = [&](uint32_t pos_ctx, uint32_t n_chunk) -> bool { - batch.n_tokens = n_chunk; - for (uint32_t pos_batch = 0; pos_batch < n_chunk; ++pos_batch) { - batch.token [pos_batch] = tokens[pos_ctx + pos_batch]; - batch.pos [pos_batch] = pos_ctx + pos_batch; - batch.n_seq_id[pos_batch] = 1; - batch.seq_id [pos_batch][0] = 0; - // an exact walk's context chunk has no label but still trains (its K/V carry later - // chunks' gradient): one output row, weighted zero, so the graph has its loss - batch.logits [pos_batch] = labels_sparse[pos_ctx + pos_batch] >= 0 - || (xc != nullptr && xc->n_labels == 0 && pos_batch + 1 == n_chunk); - } - - // output_all = false: the batch's own logits flags decide which rows exist - if (!balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, false)) { - LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); - return false; - } - - const uint32_t n_tokens_all = balloc->get_n_tokens(); - - n_queued_tokens += n_tokens_all; - - embd_seq.clear(); - - const uint32_t n_outputs_all = balloc->get_n_outputs(); - // Unreachable: the walk below trains only chunks of a labelled run. It RETURNS from the - // whole window, since a later chunk would need this one in the cache. - if (n_outputs_all == 0) { - LLAMA_LOG_ERROR("%s: a training batch with no labelled position (every label masked): nothing to learn from it\n", __func__); - return false; - } - - auto mctx = memory->init_batch(*balloc, cparams.n_ubatch, true); - if (!mctx || mctx->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { - LLAMA_LOG_ERROR("%s: could not initialize batch\n", __func__); - return false; - } - - // reserve output buffer - if (output_reserve(n_outputs_all) < n_outputs_all) { - LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); - GGML_ABORT("TODO: handle this error"); - }; - - uint32_t pos_batch = 0; - do { - const auto & ubatch = mctx->get_ubatch(); - - // the labelled rows of this ubatch, as decode counts its outputs - { - uint32_t n_outputs_ubatch = 0; - for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { - n_outputs_ubatch += (uint32_t) (ubatch.output[i] != 0); - } - n_outputs = n_outputs_ubatch; - } - - if (!mctx->apply()) { - LLAMA_LOG_ERROR("%s: failed to update the memory context\n", __func__); - return false; - } - - auto * res = gf_res_prev.get(); - - const auto gparams = graph_params(res, ubatch, mctx.get(), ctx_type_to_graph_type(cparams.ctx_type)); - - res->reset(); - - auto * gf = model.build_graph(gparams); - - struct ggml_context * ctx_compute_opt; - { - const size_t size_gf = ggml_graph_size(gf); - const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); - struct ggml_init_params params = { - /*.mem_size =*/ size_meta, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ctx_compute_opt = ggml_init(params); - } - ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); - if (xc != nullptr) { - ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); - for (const auto & io : res->t_walk) { - if (io.tk) { - ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); - ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); - } - } - for (const auto & ws : res->t_walk_state) { - if (ws.tap) { - ggml_opt_accumulate_grad(opt_ctx, ws.tap, walk_acc.state.at(ws.cache)); - } - } - } - if (!ggml_opt_alloc(opt_ctx, train)) { - const char * why = ggml_opt_refusal(opt_ctx); - opt_failure = why; - LLAMA_LOG_ERROR("%s: the %s graph was refused (window %u): %s: stopping the epoch\n", - __func__, train ? "training" : "evaluation", n_ctx, - why[0] ? why : "it does not fit in device memory"); - ggml_free(ctx_compute_opt); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return false; - } - - res->set_inputs(&ubatch); - { - // one-hot rows for the labelled positions only, in the order the ubatch - // emits its outputs (the same order the logits rows take) - struct ggml_tensor * labels = ggml_opt_labels(opt_ctx); - GGML_ASSERT(labels->ne[1] == (int64_t) n_outputs); - ggml_set_zero(labels); - const float onef = 1.0f; - uint32_t row = 0; - for (uint32_t pos_ubatch = 0; pos_ubatch < ubatch.n_tokens; ++pos_ubatch) { - if (ubatch.output[pos_ubatch] == 0) { - continue; - } - const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch; - // an exact walk's context chunk: its one row has no label (weighted zero) - GGML_ASSERT((xc != nullptr && labels_sparse[ilabel] < 0) || (labels_sparse[ilabel] >= 0 && labels_sparse[ilabel] < labels->ne[0])); - if (labels_sparse[ilabel] >= 0) { - ggml_backend_tensor_set(labels, &onef, (row*labels->ne[0] + labels_sparse[ilabel])*sizeof(float), sizeof(float)); - } - ++row; - } - GGML_ASSERT(row == n_outputs); - } - if (xc != nullptr) { - // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the - // accumulator takes this chunk's own - for (const auto & ws : res->t_walk_state) { - ggml_tensor * acc = walk_acc.state.at(ws.cache); - if (ws.gs) { - GGML_ASSERT(ggml_nbytes(acc) == ggml_nbytes(ws.gs)); - ggml_backend_tensor_set(ws.gs, acc->data, 0, ggml_nbytes(ws.gs)); - } - memset(acc->data, 0, ggml_nbytes(acc)); - } - } - ggml_opt_eval(opt_ctx, xc != nullptr && xc->n_labels == 0 ? nullptr : result); - if (const char * why = ggml_opt_refusal(opt_ctx); why[0] != '\0') { - // the backend failed the graph (ggml_opt_eval): the run fails, as a refused graph does - opt_failure = why; - LLAMA_LOG_ERROR("%s: %s: stopping the epoch\n", __func__, why); - ggml_free(ctx_compute_opt); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return false; - } - if (xc != nullptr && !res->t_walk_state.empty()) { - walk_acc.state_ready = true; - } - if (callback && !(xc != nullptr && xc->n_labels == 0)) { - callback(train, opt_ctx, dataset, result, idata_in_loop + (pos_ctx + pos_batch)/n_batch + 1, ndata_in_loop, t_loop_start); - } - ggml_free(ctx_compute_opt); - - pos_batch += ubatch.n_tokens; - } while (mctx->next()); - return true; - }; - - // a recurrent part (pure recurrent or hybrid) is snapshotted around each training chunk; - // a hybrid's attention part removes the chunk's cells like a plain cache - llama_memory_recurrent * recr = nullptr; - llama_kv_cache * attn = nullptr; - if (auto * hybrid = dynamic_cast(memory.get())) { - recr = hybrid->get_mem_recr(); - attn = hybrid->get_mem_attn(); - } else { - recr = dynamic_cast(memory.get()); - } - // the plain walk snapshots a recurrent state into a scratch sequence; the exact walk keeps host snapshots - if (recr != nullptr && cparams.n_seq_max < 2 && !(train && opt_walk_exact)) { - LLAMA_LOG_ERROR("%s: a recurrent model trains with a scratch sequence for its state snapshot: create the training context with n_seq_max >= 2\n", __func__); - opt_stop_requested.store(true); - return; - } - - if (train && opt_walk_exact) { - // THE EXACT WALK. The plain walk below steps per chunk and stops the gradient at every - // chunk boundary: a reply learns only through its own chunk's K/V, never through what - // an earlier chunk computed for it. Here the window is decoded once under the adapter - // as it is, then trained in REVERSE: each chunk's graph attends to the cache, with its - // last walk_horizon prefix positions behind a GRAD leaf, and adds the surrogate - // + on its own K/V, where G is what every later chunk put there. Its - // backward then carries that gradient through itself into the adapter and into its own - // prefix's G. ONE optimizer step per window, on the gradient of the window's loss - // (exact within the horizon). Memory stays chunk x window; the cost is one decode plus - // one forward and backward per chunk, context chunks included. - // THE INVARIANT THAT MAKES THE REVERSE PASS VALID (Fable): the adapter does not change - // between the forward decode and the last chunk's backward. One optimizer step per window, - // taken by the chunk the reverse pass trains last, is what keeps each chunk's recomputed - // K/V (and recurrent state) equal to the cached values its successors attended to; a step - // inside the window would make every surrogate a gradient of a different function. - // - // Each window starts from the REQUESTED horizon: a halving that fit one window's device - // budget must not shrink every later window and epoch, which may be shorter (Fable). - opt_walk_horizon = opt_walk_horizon_req; - // A recurrent context with rollback slots keeps several states per sequence; the walk - // carries exactly one. Refused by name, never asserted: training runs inside the serving - // process, and an assert there takes serving down (2026-10-07 04:09Z on the 5090). - if (recr != nullptr && cparams.n_rs_seq > 0) { - opt_failure = "the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has " - + std::to_string(cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0"; - LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return; - } - // A recurrent model's state is carried the same way, one chunk at a time: the state - // entering a chunk is tapped into a host accumulator, the state it leaves meets the next chunk's gradient - // (build_rs, build_walk_state_exit). The reverse pass restores each chunk's entry state - // from a host snapshot taken at its boundary in the forward pass. - std::vector chunks; - uint32_t n_labels_window = 0; - for (uint32_t p = 0; p <= (uint32_t) last_label; ) { - const bool labelled = labels_sparse[p] >= 0; - uint32_t end = p; - while (end <= (uint32_t) last_label && (labels_sparse[end] >= 0) == labelled) { - ++end; - } - for (uint32_t c0 = p; c0 < end; c0 += n_batch) { - const uint32_t c1 = std::min(c0 + n_batch, end); - chunks.push_back({ c0, c1 - c0, labelled ? c1 - c0 : 0, 0, false, false, false, 0.0f }); - n_labels_window += labelled ? c1 - c0 : 0; - } - p = end; - } - // which chunks train, and how far back each one's gradient reaches - uint32_t reach = UINT32_MAX; // the lowest position a later trained chunk's gradient reaches - int64_t first = -1; // the last chunk the reverse pass trains: the step - for (int64_t j = (int64_t) chunks.size() - 1; j >= 0; --j) { - exact_chunk & c = chunks[j]; - c.surrogate = reach < c.c0 + c.n; - // a recurrent state reaches every later chunk: a chunk before a trained one trains - c.needed = c.n_labels > 0 || c.surrogate || (recr != nullptr && first >= 0); - if (!c.needed) { - continue; - } - c.grad_from = opt_walk_horizon == 0 || c.c0 < opt_walk_horizon ? 0 : c.c0 - opt_walk_horizon; - c.loss_scale = (float) c.n_labels / (float) n_labels_window; - reach = std::min(reach, c.grad_from); - first = j; - } - GGML_ASSERT(first >= 0); - chunks[first].period_end = true; - - // THE HOST BUDGET (Fable): the accumulators are layers x positions x (k + v) floats, and a - // recurrent state is checkpointed every `stride` chunks (one sequence's state rows each; - // the first chunk's entry is the zero state and costs nothing). Both are known before - // anything runs. Under a budget the stride is the smallest that fits: the reverse pass - // rebuilds a chunk's entry state by decoding forward from its checkpoint, so a longer - // stride trades host memory for decode, never for exactness. - size_t stride = 1; - { - size_t kv_bytes = 0; - for (uint32_t il = 0; il < model.hparams.n_layer(); ++il) { - if (!model.hparams.is_recr(il)) { - kv_bytes += (size_t) (last_label + 1)*(model.hparams.n_embd_k_gqa(il) + model.hparams.n_embd_v_gqa(il))*sizeof(float); - } - } - size_t state = 0; - if (recr != nullptr) { - for (size_t il = 0; il < recr->r_l.size(); ++il) { - state += recr->r_l[il] ? recr->r_l[il]->nb[1] : 0; - state += recr->s_l[il] ? recr->s_l[il]->nb[1] : 0; - } - } - auto bytes_at = [&](size_t s) { return kv_bytes + ((chunks.size() + s - 1)/s - 1)*state; }; - if (opt_walk_host_budget > 0) { - while (stride < chunks.size() && bytes_at(stride) > opt_walk_host_budget) { - ++stride; - } - } - const size_t bytes = bytes_at(stride); - opt_walk_host_bytes = std::max(opt_walk_host_bytes, bytes); - if (opt_walk_host_budget > 0 && bytes > opt_walk_host_budget) { - char why[512]; - snprintf(why, sizeof(why), - "the exact walk needs %.1f MiB of host memory for this window (%lld positions, the K/V gradient alone), over its budget of %.1f MiB: " - "the K/V gradient grows with the window; neither the horizon nor the chunk size changes it", - bytes/1048576.0, (long long) (last_label + 1), opt_walk_host_budget/1048576.0); - opt_failure = why; - LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return; - } - if (stride > 1) { - LLAMA_LOG_INFO("%s: the exact walk checkpoints the recurrent state every %zu chunks to fit %.1f MiB of host memory (%.1f MiB)\n", - __func__, stride, opt_walk_host_budget/1048576.0, bytes/1048576.0); - } - } - { - const uint32_t n_layer = model.hparams.n_layer(); - walk_acc.k.assign(n_layer, nullptr); - walk_acc.v.assign(n_layer, nullptr); - walk_acc.ctx.reset(ggml_init({ 4*n_layer*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true })); - walk_acc.state.clear(); - walk_acc.state_ready = false; - for (uint32_t il = 0; il < n_layer; ++il) { - if (model.hparams.is_recr(il)) { - // a recurrent layer has no K/V: its state carries the chain, one sequence's row - if (recr != nullptr) { - walk_acc.state[recr->r_l[il]] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_r(), 1); - walk_acc.state[recr->s_l[il]] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_s(), 1); - } - continue; - } - walk_acc.k[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_k_gqa(il), last_label + 1); - walk_acc.v[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_v_gqa(il), last_label + 1); - } - walk_acc.buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(walk_acc.ctx.get(), ggml_backend_cpu_buffer_type())); - if (!walk_acc.buf) { - opt_failure = "the exact walk's gradient accumulator (host memory) could not be allocated"; - LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return; - } - ggml_backend_buffer_clear(walk_acc.buf.get(), 0); - cparams.walk_acc_k = walk_acc.k.data(); - cparams.walk_acc_v = walk_acc.v.data(); - } - - // the window as the adapter now reads it, chunk by chunk: a recurrent model's state is - // checkpointed at each chunk's start through the sequence-state API (the recurrent part - // only, PARTIAL_ONLY), after waiting for the queued decode; a restore puts back the - // cell's position with its rows, and an empty checkpoint (before the first chunk) - // restores the zero state. A hand-rolled per-layer copy here skipped that wait and, on - // Metal, carried a torn mid-graph state into the reverse pass (2026-10-11). - std::vector> snaps(recr != nullptr ? chunks.size() : 0); - const llama_state_seq_flags state_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; - auto save_state = [&](std::vector & snap) { - synchronize(); - snap.resize(state_seq_get_size(0, state_flags)); - GGML_ASSERT(state_seq_get_data(0, snap.data(), snap.size(), state_flags) == snap.size()); - }; - auto load_state = [&](const std::vector & snap) { - synchronize(); - GGML_ASSERT(state_seq_set_data(0, snap.data(), snap.size(), state_flags) == snap.size()); - }; - for (size_t j = 0; j < chunks.size(); ++j) { - if (recr != nullptr && j % stride == 0) { - save_state(snaps[j]); - } - if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n)) { - return; - } - } - bool ok = true; - for (int64_t j = (int64_t) chunks.size() - 1; j >= first && ok; --j) { - const exact_chunk & c = chunks[j]; - if (!c.needed) { - continue; - } - if (!yield_point()) { - ok = false; - break; - } - // attention: pop the cache to [0, c0); recurrent: restore the state at the chunk's - // checkpoint and decode forward to its start (that decode rebuilds the attention K/V - // in between too, under the same adapter: the same values). Run again before every - // retry below: a refused chunk graph is refused AFTER its ubatch was applied, so the - // chunk's cells already sit in the cache and the recurrent cell's position has moved - // to the chunk's end; retrying on top of that fails to prepare the ubatch at all - // (the 5090, 2026-10-10: "init_batch: failed to prepare attention ubatches" after - // "retrying at 512", and the epoch stopped instead of shrinking the horizon). - const auto rewind = [&]() -> bool { - const size_t cp = recr != nullptr ? (size_t) j - (size_t) j % stride : (size_t) j; - const uint32_t p_pop = chunks[cp].c0; - const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : memory->seq_rm(0, p_pop, -1); - if (!popped) { - LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u) for the reverse pass\n", __func__, p_pop); - opt_stop_requested.store(true); - return false; - } - if (recr != nullptr) { - load_state(snaps[cp]); - if (p_pop < c.c0 && !decode_span(p_pop, c.c0)) { - return false; - } - } - return true; - }; - if (!rewind()) { - ok = false; - break; - } - cparams.walk_exact = true; - cparams.walk_surrogate = c.surrogate; - cparams.walk_state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && walk_acc.state_ready; - // THE HORIZON FITS THE DEVICE. Its GRAD leaves are the chunk graph's largest growth - // with the window (layers x horizon x (k + v) floats, twice with their gradient); the - // graph preflight refuses a chunk that does not fit before allocating anything, so a - // memory refusal halves the horizon (whole chunks) and the chunk is tried again; every - // later chunk keeps the smaller one. The forward still attends to the whole window and - // the recurrent state is carried exactly; only attention gradient past the horizon is - // dropped. A refusal at a horizon of one chunk, or for a node the device cannot run, - // stands. - exact_chunk cc = c; - for (;;) { - cc.grad_from = opt_walk_horizon == 0 || cc.c0 < opt_walk_horizon ? 0 : cc.c0 - opt_walk_horizon; - cparams.walk_grad_from = cc.grad_from; - xc = &cc; - ok = train_chunk(cc.c0, cc.n); - xc = nullptr; - const uint32_t span = cc.c0 - cc.grad_from; - if (ok || !opt_failure.empty() || !opt_alloc_failed.load() || span <= n_batch) { - break; - } - opt_walk_horizon = std::max(n_batch, (span/2)/n_batch*n_batch); - LLAMA_LOG_WARN("%s: the chunk at %u did not fit with a gradient horizon of %u positions: retrying at %u\n", - __func__, cc.c0, span, opt_walk_horizon); - opt_alloc_failed.store(false); - opt_stop_requested.store(false); - // the rewind's decode is a plain forward, as before the first try: no walk flags. - // walk_grad_from is not restored here because the loop's first statements recompute - // it from the shrunken horizon (cc.grad_from, then cparams.walk_grad_from) before the - // retried train_chunk runs. - const bool surrogate = cparams.walk_surrogate; - const bool state_surrogate = cparams.walk_state_surrogate; - cparams.walk_exact = false; - cparams.walk_grad_from = 0; - cparams.walk_surrogate = false; - cparams.walk_state_surrogate = false; - const bool rewound = rewind(); - cparams.walk_exact = true; - cparams.walk_surrogate = surrogate; - cparams.walk_state_surrogate = state_surrogate; - if (!rewound) { - ok = false; - break; - } - } - opt_walk_horizon_used = opt_walk_horizon; - cparams.walk_exact = false; - cparams.walk_grad_from = 0; - cparams.walk_surrogate = false; - cparams.walk_state_surrogate = false; - } - walk_acc.state_ready = false; - return; - } - - uint32_t pos = 0; - while (pos <= (uint32_t) last_label && !opt_stop_requested.load(std::memory_order_relaxed)) { - const bool labelled = labels_sparse[pos] >= 0; - uint32_t end = pos; - while (end <= (uint32_t) last_label && (labels_sparse[end] >= 0) == labelled) { - ++end; - } - if (!labelled) { - if (!decode_span(pos, end)) { - return; - } - } else { - for (uint32_t c0 = pos; c0 < end; c0 += n_batch) { - const uint32_t c1 = std::min(c0 + n_batch, end); - // The training forward advances a recurrent state in place, and a recurrent - // memory cannot remove a partial range: snapshot the state to the scratch - // sequence first (Fable), restore it after, then the decode below advances it - // from where the context left it. The state is per sequence and small. - const bool snapshot = recr != nullptr && c1 <= (uint32_t) last_label; - if (!yield_point()) { - return; - } - if (snapshot) { - recr->seq_cp(0, 1, -1, -1); - } - if (!train_chunk(c0, c1 - c0)) { - return; - } - if (snapshot) { - recr->seq_rm(0, -1, -1); - recr->seq_cp(1, 0, -1, -1); - recr->seq_rm(1, -1, -1); - } - // the chunk joins the context under the adapter as it now is: the training - // graph never writes the cache, so its cells hold nothing until this decode - if (c1 <= (uint32_t) last_label) { - // the chunk's cells hold no K/V (the training graph never writes the - // cache): drop them from the attention memory, then decode them for real - const bool removed = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, c0, -1)) : memory->seq_rm(0, c0, -1); - if (!removed) { - LLAMA_LOG_ERROR("%s: could not remove the trained chunk [%u, %u) from the cache before decoding it\n", __func__, c0, c1); - opt_stop_requested.store(true); - return; - } - if (!decode_span(c0, c1)) { - opt_stop_requested.store(true); - return; - } - } - } - } - pos = end; - } -} - -void llama_context::opt_epoch( - ggml_opt_dataset_t dataset, - ggml_opt_result_t result_train, - ggml_opt_result_t result_eval, - int64_t idata_split, - ggml_opt_epoch_callback callback_train, - ggml_opt_epoch_callback callback_eval) { - const uint32_t n_ctx = this->n_ctx(); - const uint32_t n_batch = std::min(cparams.n_batch, n_ctx); - const uint32_t n_ubatch = std::min(cparams.n_ubatch, n_batch); - const int64_t ndata = ggml_opt_dataset_ndata(dataset); - - GGML_ASSERT(idata_split >= 0); - GGML_ASSERT(idata_split <= ndata); - - const uint32_t ubatch_per_ctx = n_ctx / n_ubatch; - - struct llama_batch batch = llama_batch_init(n_batch, 0, 1); - std::vector tokens(n_ctx); - std::vector labels_sparse(n_ctx); - - int64_t idata = 0; - - int64_t t_loop_start = ggml_time_us(); - int64_t ndata_in_loop = idata_split*ubatch_per_ctx; - for (; idata < idata_split; ++idata) { - if (opt_stop_requested.load(std::memory_order_relaxed)) { - break; - } - constexpr bool train = true; - const int64_t idata_in_loop = idata*ubatch_per_ctx; - - if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { - opt_stop_requested.store(true); - } - if (opt_stop_requested.load(std::memory_order_relaxed)) { - break; - } - - ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); - opt_epoch_iter(dataset, result_train, tokens, labels_sparse, batch, - callback_train, train, idata_in_loop, ndata_in_loop, t_loop_start); - } - - t_loop_start = ggml_time_us(); - ndata_in_loop = (ndata - idata_split)*ubatch_per_ctx; - for (; idata < ndata; ++idata) { - if (opt_stop_requested.load(std::memory_order_relaxed)) { - break; - } - constexpr bool train = false; - const int64_t idata_in_loop = (idata - idata_split)*ubatch_per_ctx; - - if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { - opt_stop_requested.store(true); - } - if (opt_stop_requested.load(std::memory_order_relaxed)) { - break; - } - - ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); - opt_epoch_iter(dataset, result_eval, tokens, labels_sparse, batch, - callback_eval, train, idata_in_loop, ndata_in_loop, t_loop_start); - } - - llama_batch_free(batch); -} - // // interface implementation // @@ -4778,93 +3979,6 @@ void llama_perf_context_reset(llama_context * ctx) { ctx->perf_reset(); } -// -// training -// - -bool llama_opt_param_filter_all(const struct ggml_tensor * tensor, void * userdata) { - GGML_UNUSED(tensor); - GGML_UNUSED(userdata); - return true; -} - -void llama_opt_init(struct llama_context * ctx, struct llama_model * model, struct llama_opt_params lopt_params) { - ctx->opt_alloc_failed.store(false); - ctx->opt_failure.clear(); - ctx->opt_init(model, lopt_params); -} - -void llama_opt_stop(struct llama_context * ctx, bool stop) { - ctx->opt_stop_requested.store(stop, std::memory_order_relaxed); -} - -void llama_opt_set_step_callback(struct llama_context * ctx, llama_opt_step_callback callback, void * user_data) { - ctx->opt_step_callback = callback; - ctx->opt_step_callback_data = user_data; -} - -void llama_opt_set_walk_host_budget(struct llama_context * ctx, size_t bytes) { - ctx->opt_walk_host_budget = bytes; -} - -size_t llama_opt_walk_host_bytes(struct llama_context * ctx) { - return ctx->opt_walk_host_bytes; -} - -uint32_t llama_opt_walk_horizon(struct llama_context * ctx) { - return ctx->opt_walk_horizon_used; -} - -void llama_opt_set_memory_budget(struct llama_context * ctx, size_t bytes) { - ctx->opt_memory_budget = bytes; -} - -size_t llama_context::opt_graph_bytes() const { - return opt_ctx ? ggml_opt_peak_graph_bytes(opt_ctx) : 0; -} - -size_t llama_opt_graph_bytes(struct llama_context * ctx) { - return ctx->opt_graph_bytes(); -} - -size_t llama_opt_context_device_bytes(struct llama_context * ctx) { - size_t bytes = 0; - for (const auto & [buft, data] : ctx->memory_breakdown()) { - // data.context only: the compute buffers are where the training graph is allocated - // (the optimizer shares this context's scheduler, whose allocator replaces them with the - // graph's), so the graph's measured size already covers them (Cormac on #52) - if (!ggml_backend_buft_is_host(buft)) { - bytes += data.context; - } - } - return bytes; -} - -bool llama_opt_failed(struct llama_context * ctx) { - return ctx->opt_alloc_failed.load(); -} - -const char * llama_opt_failure(struct llama_context * ctx) { - return ctx->opt_failure.c_str(); -} - -void llama_opt_epoch( - struct llama_context * ctx, - ggml_opt_dataset_t dataset, - ggml_opt_result_t result_train, - ggml_opt_result_t result_eval, - int64_t idata_split, - ggml_opt_epoch_callback callback_train, - ggml_opt_epoch_callback callback_eval) { - ctx->opt_epoch( - dataset, - result_train, - result_eval, - idata_split, - callback_train, - callback_eval); -} - // // ext // diff --git a/src/llama-context.h b/src/llama-context.h index 1cd3e75d2cee..26ec2796cf1a 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -279,6 +279,9 @@ struct llama_context { const llama_memory_context_i * mctx, llm_graph_type gtype) const; + // the graph type this context's type builds (decoder, or the MTP head) + llm_graph_type graph_type() const; + llm_graph_cb graph_get_cb() const; // disable auto fused ops (Flash Attention, Gated Delta Net) whose op lands on a device From 33322b19dab703c68b463ba997f5316996755091 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Sun, 11 Oct 2026 20:32:04 -0500 Subject: [PATCH 6/6] train: the walk is its own module; one window's walk is one object opt_epoch_iter was a 620-line function of lambdas sharing captured state. Each concern now has one owner: - llama-walk.{h,cpp}: the plan as plain data. llama_walk_plan::of cuts the window's runs every n_batch; both walks take its chunks (the plain walk had its own copy of the run/chunk scan). mark_exact (needed, surrogate, reach, period end, loss scale), grad_from (computed in two places before), host_bytes and stride_for (the host budget). llama_walk_acc is the exact walk's host gradients (allocation, sizing and binding into cparams), and llama_walk_fill is the one batch fill (decode_span and train_chunk each had one). - llama_context::opt_walk, a nested type in llama-context-opt.cpp: one window's walk, with yield_point, decode_span, train_chunk, run_exact and run_plain as methods (136 and 156 lines are the largest), and set_flags as the only writer of the walk's graph flags in cparams, in place of the save, clear and restore around the horizon retry. opt_epoch_iter is now 41 lines: build the walk, find the memory, run one of the two. - llama_context::opt_fail: the five copies of "record why, log, flag the failure, stop the epoch" become one private member. - opt_epoch: one loop over the training range and the evaluation range; the two loops were copies. - test-walk-plan: the plan without a model (chunks follow runs and stop at the last label; one period end at the reverse pass's last chunk; loss scales sum to one; the stride shrinks checkpoints to the budget). It runs in ctest; test-walk- exact needs a model and does not. The plain walk's recurrent snapshot stays on the scratch sequence: seq_cp is the cache's own device-side copy-on-write, one snapshot per chunk with no host round trip. The exact walk's checkpoints stay on the state API, many alive across the reverse pass. ## Measured M5 Metal, built from this tree (diff 62484d676860 on 05e96822f): - test-walk-exact, qwen3.5-4b hybrid Q4_K_M: exit 0, every exact row cosine 1.0000, every training loss identical to 4da182d9a's run (3.252877, 3.126584, 2.910862, ...), host memory 182.8 to 82.2 MiB checkpointed, device graph 1851.1 MiB per chunk, 2026-10-12T01:30:50Z. - test-walk-exact, qwen2.5-0.5b pure attention: exit 0, exact rows 1.0000, 2026-10-12T01:31:28Z. - ctest -R "opt|walk": 3 of 3 pass (test-walk-plan 41 assertions). - Full build: no warnings. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- src/CMakeLists.txt | 1 + src/llama-context-opt.cpp | 1071 ++++++++++++++++--------------------- src/llama-context.h | 8 + src/llama-walk.cpp | 142 +++++ src/llama-walk.h | 94 ++++ tests/CMakeLists.txt | 1 + tests/test-walk-plan.cpp | 74 +++ 7 files changed, 794 insertions(+), 597 deletions(-) create mode 100644 src/llama-walk.cpp create mode 100644 src/llama-walk.h create mode 100644 tests/test-walk-plan.cpp diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 4c1c95b856a9..cd429130227d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -17,6 +17,7 @@ add_library(llama llama-chat.cpp llama-context.cpp llama-context-opt.cpp + llama-walk.cpp llama-cparams.cpp llama-grammar.cpp llama-graph.cpp diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 1e20995686ad..3c183f3d255c 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -1,6 +1,7 @@ // Training: llama_context's optimizer members (opt_init, opt_epoch, the walk in // opt_epoch_iter) and the llama_opt_* interface, apart from inference so each reads as one concern. #include "llama-context.h" +#include "llama-walk.h" #include "ggml.h" #include "llama-batch.h" @@ -128,626 +129,519 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params } } -void llama_context::opt_epoch_iter( - ggml_opt_dataset_t dataset, - ggml_opt_result_t result, - const std::vector & tokens, - const std::vector & labels_sparse, - llama_batch & batch, - ggml_opt_epoch_callback callback, - bool train, - int64_t idata_in_loop, - int64_t ndata_in_loop, - int64_t t_loop_start) { - GGML_ASSERT(opt_ctx); - const uint32_t n_ctx = opt_n_ctx_train; - const uint32_t n_batch = std::min(this->n_batch(), n_ctx); +// THE WALK (Fable's design, continuum 2026-10-06): one pass over a window in one context. A run +// with no label is context she did not write: a plain decode into the cache, no backward. A +// labelled run is her reply: it trains in chunks of n_batch, each attending to everything before +// it read from the cache as a constant (build_attn, cparams.training), then the reply is decoded +// into the cache under the adapter as it now is, so the next reply sees it. Memory is chunk x +// window, never window x window, so her whole lived window trains. Nothing after the last label +// is computed. One per window; a member of the context, so it runs on the context's own state. +struct llama_context::opt_walk { + llama_context & lctx; + ggml_opt_dataset_t dataset; + ggml_opt_result_t result; + const std::vector & tokens; + const std::vector & labels; + llama_batch & batch; + ggml_opt_epoch_callback callback; + bool train; + int64_t idata_in_loop; + int64_t ndata_in_loop; + int64_t t_loop_start; + uint32_t n_ctx; + uint32_t n_batch; + llama_walk_plan plan; + + // a recurrent part (pure recurrent or hybrid) carries its state across chunks; a hybrid's + // attention part removes a chunk's cells like a plain cache + llama_memory_recurrent * recr = nullptr; + llama_kv_cache * attn = nullptr; - memory->clear(true); + bool first_unit = true; + // THE EXACT WALK's chunk being trained (null in the plain walk) and its host gradients + const llama_walk_chunk * xc = nullptr; + llama_walk_acc acc; + + opt_walk(llama_context & lctx, ggml_opt_dataset_t dataset, ggml_opt_result_t result, + const std::vector & tokens, const std::vector & labels, llama_batch & batch, + ggml_opt_epoch_callback callback, bool train, int64_t idata_in_loop, int64_t ndata_in_loop, int64_t t_loop_start, + uint32_t n_ctx, uint32_t n_batch) + : lctx(lctx), dataset(dataset), result(result), tokens(tokens), labels(labels), batch(batch), callback(callback), + train(train), idata_in_loop(idata_in_loop), ndata_in_loop(ndata_in_loop), t_loop_start(t_loop_start), + n_ctx(n_ctx), n_batch(n_batch), plan(llama_walk_plan::of(labels, n_ctx, n_batch)) {} + + bool yield_point(); + bool decode_span(uint32_t p0, uint32_t p1); + bool train_chunk(uint32_t pos_ctx, uint32_t n_chunk); + void run_exact(); + void run_plain(); + + // the exact walk's graph flags in cparams, which build_attn and build_rs read + void set_flags(bool exact, uint32_t grad_from, bool surrogate, bool state_surrogate) { + lctx.cparams.walk_exact = exact; + lctx.cparams.walk_grad_from = grad_from; + lctx.cparams.walk_surrogate = surrogate; + lctx.cparams.walk_state_surrogate = state_surrogate; + } +}; + +// THE YIELD POINT IS THE CHUNK, not the window: a turn that arrives mid-walk waits for one chunk, +// never for the whole example (measured on the 5090: one callback per example made a turn wait for +// its entire walk, 1.2-1.8 s on a 1.5B at 15k and 13.5 s on Fable's run). The first unit is the +// window's own boundary, which opt_epoch has already offered. +bool llama_context::opt_walk::yield_point() { + if (first_unit) { + first_unit = false; + } else if (lctx.opt_step_callback && !lctx.opt_step_callback(train, lctx.opt_step_callback_data)) { + lctx.opt_stop_requested.store(true); + } + return !lctx.opt_stop_requested.load(std::memory_order_relaxed); +} - // THE WALK (Fable's design, continuum 2026-10-06): one pass over the window in one - // context. A run with no label is context she did not write: a plain decode into the - // cache, no backward. A labelled run is her reply: it trains in chunks of n_batch, each - // attending to everything before it read from the cache as a constant (build_attn, - // cparams.training), then the reply is decoded into the cache under the adapter as it - // now is, so the next reply sees it. Memory is chunk x window, never window x window, so - // her whole lived window trains. Nothing after the last label is computed. - int64_t last_label = -1; - for (int64_t i = (int64_t) n_ctx - 1; i >= 0; --i) { - if (labels_sparse[i] >= 0) { - last_label = i; - break; +// a plain inference decode of [p0, p1) into the cache: the training graph's attention never +// writes the cache, so context and finished replies reach it this way +bool llama_context::opt_walk::decode_span(uint32_t p0, uint32_t p1) { + llama_cparams & cparams = lctx.cparams; + const bool training = cparams.training; + cparams.training = false; + lctx.gf_res_prev->reset(); + bool ok = true; + for (uint32_t c0 = p0; c0 < p1 && ok; c0 += n_batch) { + if (!yield_point()) { + cparams.training = training; + lctx.gf_res_prev->reset(); + return false; // stopped or cancelled at a chunk boundary: not a decode failure } + const uint32_t c1 = std::min(c0 + n_batch, p1); + llama_walk_fill(batch, tokens, c0, c1 - c0); + batch.logits[c1 - c0 - 1] = true; + ok = lctx.decode(batch) == 0; } - if (last_label < 0) { - LLAMA_LOG_ERROR("%s: a training window with no labelled position: nothing to learn from it\n", __func__); - return; + cparams.training = training; + lctx.gf_res_prev->reset(); + if (!ok) { + LLAMA_LOG_ERROR("%s: the context decode of [%u, %u) failed\n", __func__, p0, p1); } + return ok; +} - // THE YIELD POINT IS THE CHUNK, not the window: a turn that arrives mid-walk waits for one - // chunk, never for the whole example (measured on the 5090: one callback per example made a - // turn wait for its entire walk, 1.2-1.8 s on a 1.5B at 15k and 13.5 s on Fable's run). The - // first unit is the window's own boundary, which opt_epoch has already offered. - bool first_unit = true; - auto yield_point = [&]() -> bool { - if (first_unit) { - first_unit = false; - } else if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { - opt_stop_requested.store(true); - } - return !opt_stop_requested.load(std::memory_order_relaxed); - }; +// one training chunk [pos_ctx, pos_ctx + n_chunk): forward and backward over it alone, attending +// to the cached context before it +bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { + const llama_model & model = lctx.model; + llama_cparams & cparams = lctx.cparams; + ggml_opt_context_t opt_ctx = lctx.opt_ctx; + + llama_walk_fill(batch, tokens, pos_ctx, n_chunk); + for (uint32_t pos_batch = 0; pos_batch < n_chunk; ++pos_batch) { + // an exact walk's context chunk has no label but still trains (its K/V carry later + // chunks' gradient): one output row, weighted zero, so the graph has its loss + batch.logits[pos_batch] = labels[pos_ctx + pos_batch] >= 0 + || (xc != nullptr && xc->n_labels == 0 && pos_batch + 1 == n_chunk); + } - // a plain inference decode of [p0, p1) into the cache: the training graph's attention - // never writes the cache, so context and finished replies reach it this way - auto decode_span = [&](uint32_t p0, uint32_t p1) -> bool { - const bool training = cparams.training; - cparams.training = false; - gf_res_prev->reset(); - bool ok = true; - for (uint32_t c0 = p0; c0 < p1 && ok; c0 += n_batch) { - if (!yield_point()) { - cparams.training = training; - gf_res_prev->reset(); - return false; // stopped or cancelled at a chunk boundary: not a decode failure - } - const uint32_t c1 = std::min(c0 + n_batch, p1); - batch.n_tokens = c1 - c0; - for (uint32_t i = 0; i < c1 - c0; ++i) { - batch.token [i] = tokens[c0 + i]; - batch.pos [i] = c0 + i; - batch.n_seq_id[i] = 1; - batch.seq_id [i][0] = 0; - batch.logits [i] = i + 1 == c1 - c0; - } - const int rc = decode(batch); - ok = rc == 0; - } - cparams.training = training; - gf_res_prev->reset(); - if (!ok) { - LLAMA_LOG_ERROR("%s: the context decode of [%u, %u) failed\n", __func__, p0, p1); - } - return ok; - }; + // output_all = false: the batch's own logits flags decide which rows exist + if (!lctx.balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, false)) { + LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); + return false; + } - // THE EXACT WALK's state for the chunk being trained (null in the plain walk): its place in - // the window's one optimizer period, and the gradient accumulated on the window's cached K/V - struct exact_chunk { - uint32_t c0; - uint32_t n; - uint32_t n_labels; - uint32_t grad_from; // its gradient taps cover cached positions [grad_from, c0) - bool surrogate; // a later chunk's gradient sits on its own K/V - bool needed; - bool period_end; // the window's optimizer step - float loss_scale; // its labelled positions over the window's - }; - const exact_chunk * xc = nullptr; - // THE EXACT WALK's gradients, in host memory, each added into in place by the backward: - // per attention layer dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, - // position] F32 (null for a recurrent layer), which the graph reads through cparams (never - // outliving it, whatever way the walk returns); per recurrent state tensor (a layer's conv - // state r_l and its recurrent state s_l) dL/d(the state entering the chunk trained last), - // which is what the chunk before it puts on the state it leaves (state_ready once one is in) - struct walk_acc_t { - llama_cparams & cp; - std::vector k, v; - std::unordered_map state; - bool state_ready = false; - ggml_context_ptr ctx; - ggml_backend_buffer_ptr buf; - ~walk_acc_t() { cp.walk_acc_k = nullptr; cp.walk_acc_v = nullptr; } - } walk_acc { cparams, {}, {}, {}, false, nullptr, nullptr }; - - // one training chunk [pos_ctx, pos_ctx + n_tokens): forward and backward over it alone, - // attending to the cached context before it - auto train_chunk = [&](uint32_t pos_ctx, uint32_t n_chunk) -> bool { - batch.n_tokens = n_chunk; - for (uint32_t pos_batch = 0; pos_batch < n_chunk; ++pos_batch) { - batch.token [pos_batch] = tokens[pos_ctx + pos_batch]; - batch.pos [pos_batch] = pos_ctx + pos_batch; - batch.n_seq_id[pos_batch] = 1; - batch.seq_id [pos_batch][0] = 0; - // an exact walk's context chunk has no label but still trains (its K/V carry later - // chunks' gradient): one output row, weighted zero, so the graph has its loss - batch.logits [pos_batch] = labels_sparse[pos_ctx + pos_batch] >= 0 - || (xc != nullptr && xc->n_labels == 0 && pos_batch + 1 == n_chunk); - } + lctx.n_queued_tokens += lctx.balloc->get_n_tokens(); - // output_all = false: the batch's own logits flags decide which rows exist - if (!balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, false)) { - LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); - return false; - } + lctx.embd_seq.clear(); - const uint32_t n_tokens_all = balloc->get_n_tokens(); + const uint32_t n_outputs_all = lctx.balloc->get_n_outputs(); + // Unreachable: the walk trains only chunks of a labelled run. It RETURNS from the whole + // window, since a later chunk would need this one in the cache. + if (n_outputs_all == 0) { + LLAMA_LOG_ERROR("%s: a training batch with no labelled position (every label masked): nothing to learn from it\n", __func__); + return false; + } - n_queued_tokens += n_tokens_all; + auto mctx = lctx.memory->init_batch(*lctx.balloc, cparams.n_ubatch, true); + if (!mctx || mctx->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { + LLAMA_LOG_ERROR("%s: could not initialize batch\n", __func__); + return false; + } - embd_seq.clear(); + // reserve output buffer + if (lctx.output_reserve(n_outputs_all) < n_outputs_all) { + LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); + GGML_ABORT("TODO: handle this error"); + }; - const uint32_t n_outputs_all = balloc->get_n_outputs(); - // Unreachable: the walk below trains only chunks of a labelled run. It RETURNS from the - // whole window, since a later chunk would need this one in the cache. - if (n_outputs_all == 0) { - LLAMA_LOG_ERROR("%s: a training batch with no labelled position (every label masked): nothing to learn from it\n", __func__); - return false; + uint32_t pos_batch = 0; + do { + const auto & ubatch = mctx->get_ubatch(); + + // the labelled rows of this ubatch, as decode counts its outputs + { + uint32_t n_outputs_ubatch = 0; + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + n_outputs_ubatch += (uint32_t) (ubatch.output[i] != 0); + } + lctx.n_outputs = n_outputs_ubatch; } - auto mctx = memory->init_batch(*balloc, cparams.n_ubatch, true); - if (!mctx || mctx->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { - LLAMA_LOG_ERROR("%s: could not initialize batch\n", __func__); + if (!mctx->apply()) { + LLAMA_LOG_ERROR("%s: failed to update the memory context\n", __func__); return false; } - // reserve output buffer - if (output_reserve(n_outputs_all) < n_outputs_all) { - LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); - GGML_ABORT("TODO: handle this error"); - }; - - uint32_t pos_batch = 0; - do { - const auto & ubatch = mctx->get_ubatch(); - - // the labelled rows of this ubatch, as decode counts its outputs - { - uint32_t n_outputs_ubatch = 0; - for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { - n_outputs_ubatch += (uint32_t) (ubatch.output[i] != 0); - } - n_outputs = n_outputs_ubatch; - } - - if (!mctx->apply()) { - LLAMA_LOG_ERROR("%s: failed to update the memory context\n", __func__); - return false; - } + auto * res = lctx.gf_res_prev.get(); - auto * res = gf_res_prev.get(); + const auto gparams = lctx.graph_params(res, ubatch, mctx.get(), lctx.graph_type()); - const auto gparams = graph_params(res, ubatch, mctx.get(), graph_type()); + res->reset(); - res->reset(); + auto * gf = model.build_graph(gparams); - auto * gf = model.build_graph(gparams); - - struct ggml_context * ctx_compute_opt; - { - const size_t size_gf = ggml_graph_size(gf); - const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); - struct ggml_init_params params = { - /*.mem_size =*/ size_meta, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ctx_compute_opt = ggml_init(params); - } - ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); - if (xc != nullptr) { - ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); - for (const auto & io : res->t_walk) { - if (io.tk) { - ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); - ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); - } - } - for (const auto & ws : res->t_walk_state) { - if (ws.tap) { - ggml_opt_accumulate_grad(opt_ctx, ws.tap, walk_acc.state.at(ws.cache)); - } + struct ggml_context * ctx_compute_opt; + { + const size_t size_gf = ggml_graph_size(gf); + const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); + struct ggml_init_params params = { + /*.mem_size =*/ size_meta, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx_compute_opt = ggml_init(params); + } + ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); + if (xc != nullptr) { + ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); + for (const auto & io : res->t_walk) { + if (io.tk) { + ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); + ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); } } - if (!ggml_opt_alloc(opt_ctx, train)) { - const char * why = ggml_opt_refusal(opt_ctx); - opt_failure = why; - LLAMA_LOG_ERROR("%s: the %s graph was refused (window %u): %s: stopping the epoch\n", - __func__, train ? "training" : "evaluation", n_ctx, - why[0] ? why : "it does not fit in device memory"); - ggml_free(ctx_compute_opt); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return false; + for (const auto & ws : res->t_walk_state) { + if (ws.tap) { + ggml_opt_accumulate_grad(opt_ctx, ws.tap, acc.state.at(ws.cache)); + } } + } + if (!ggml_opt_alloc(opt_ctx, train)) { + ggml_free(ctx_compute_opt); + lctx.opt_fail(ggml_opt_refusal(opt_ctx), + std::string("the ") + (train ? "training" : "evaluation") + " graph was refused (window " + std::to_string(n_ctx) + ")"); + return false; + } - res->set_inputs(&ubatch); - { - // one-hot rows for the labelled positions only, in the order the ubatch - // emits its outputs (the same order the logits rows take) - struct ggml_tensor * labels = ggml_opt_labels(opt_ctx); - GGML_ASSERT(labels->ne[1] == (int64_t) n_outputs); - ggml_set_zero(labels); - const float onef = 1.0f; - uint32_t row = 0; - for (uint32_t pos_ubatch = 0; pos_ubatch < ubatch.n_tokens; ++pos_ubatch) { - if (ubatch.output[pos_ubatch] == 0) { - continue; - } - const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch; - // an exact walk's context chunk: its one row has no label (weighted zero) - GGML_ASSERT((xc != nullptr && labels_sparse[ilabel] < 0) || (labels_sparse[ilabel] >= 0 && labels_sparse[ilabel] < labels->ne[0])); - if (labels_sparse[ilabel] >= 0) { - ggml_backend_tensor_set(labels, &onef, (row*labels->ne[0] + labels_sparse[ilabel])*sizeof(float), sizeof(float)); - } - ++row; + res->set_inputs(&ubatch); + { + // one-hot rows for the labelled positions only, in the order the ubatch emits its + // outputs (the same order the logits rows take) + struct ggml_tensor * labels_t = ggml_opt_labels(opt_ctx); + GGML_ASSERT(labels_t->ne[1] == (int64_t) lctx.n_outputs); + ggml_set_zero(labels_t); + const float onef = 1.0f; + uint32_t row = 0; + for (uint32_t pos_ubatch = 0; pos_ubatch < ubatch.n_tokens; ++pos_ubatch) { + if (ubatch.output[pos_ubatch] == 0) { + continue; } - GGML_ASSERT(row == n_outputs); - } - if (xc != nullptr) { - // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the - // accumulator takes this chunk's own - for (const auto & ws : res->t_walk_state) { - ggml_tensor * acc = walk_acc.state.at(ws.cache); - if (ws.gs) { - GGML_ASSERT(ggml_nbytes(acc) == ggml_nbytes(ws.gs)); - ggml_backend_tensor_set(ws.gs, acc->data, 0, ggml_nbytes(ws.gs)); - } - memset(acc->data, 0, ggml_nbytes(acc)); + const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch; + // an exact walk's context chunk: its one row has no label (weighted zero) + GGML_ASSERT((xc != nullptr && labels[ilabel] < 0) || (labels[ilabel] >= 0 && labels[ilabel] < labels_t->ne[0])); + if (labels[ilabel] >= 0) { + ggml_backend_tensor_set(labels_t, &onef, (row*labels_t->ne[0] + labels[ilabel])*sizeof(float), sizeof(float)); } + ++row; } - ggml_opt_eval(opt_ctx, xc != nullptr && xc->n_labels == 0 ? nullptr : result); - if (const char * why = ggml_opt_refusal(opt_ctx); why[0] != '\0') { - // the backend failed the graph (ggml_opt_eval): the run fails, as a refused graph does - opt_failure = why; - LLAMA_LOG_ERROR("%s: %s: stopping the epoch\n", __func__, why); - ggml_free(ctx_compute_opt); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return false; - } - if (xc != nullptr && !res->t_walk_state.empty()) { - walk_acc.state_ready = true; - } - if (callback && !(xc != nullptr && xc->n_labels == 0)) { - callback(train, opt_ctx, dataset, result, idata_in_loop + (pos_ctx + pos_batch)/n_batch + 1, ndata_in_loop, t_loop_start); + GGML_ASSERT(row == lctx.n_outputs); + } + if (xc != nullptr) { + // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the + // accumulator takes this chunk's own + for (const auto & ws : res->t_walk_state) { + ggml_tensor * a = acc.state.at(ws.cache); + if (ws.gs) { + GGML_ASSERT(ggml_nbytes(a) == ggml_nbytes(ws.gs)); + ggml_backend_tensor_set(ws.gs, a->data, 0, ggml_nbytes(ws.gs)); + } + memset(a->data, 0, ggml_nbytes(a)); } + } + ggml_opt_eval(opt_ctx, xc != nullptr && xc->n_labels == 0 ? nullptr : result); + if (const char * why = ggml_opt_refusal(opt_ctx); why[0] != '\0') { + // the backend failed the graph (ggml_opt_eval): the run fails, as a refused graph does ggml_free(ctx_compute_opt); + lctx.opt_fail(why); + return false; + } + if (xc != nullptr && !res->t_walk_state.empty()) { + acc.state_ready = true; + } + if (callback && !(xc != nullptr && xc->n_labels == 0)) { + callback(train, opt_ctx, dataset, result, idata_in_loop + (pos_ctx + pos_batch)/n_batch + 1, ndata_in_loop, t_loop_start); + } + ggml_free(ctx_compute_opt); - pos_batch += ubatch.n_tokens; - } while (mctx->next()); - return true; - }; + pos_batch += ubatch.n_tokens; + } while (mctx->next()); + return true; +} - // a recurrent part (pure recurrent or hybrid) is snapshotted around each training chunk; - // a hybrid's attention part removes the chunk's cells like a plain cache - llama_memory_recurrent * recr = nullptr; - llama_kv_cache * attn = nullptr; - if (auto * hybrid = dynamic_cast(memory.get())) { - recr = hybrid->get_mem_recr(); - attn = hybrid->get_mem_attn(); - } else { - recr = dynamic_cast(memory.get()); +// THE EXACT WALK. The plain walk steps per chunk and stops the gradient at every chunk boundary: a +// reply learns only through its own chunk's K/V, never through what an earlier chunk computed for +// it. Here the window is decoded once under the adapter as it is, then trained in REVERSE: each +// chunk's graph attends to the cache, with its last walk_horizon prefix positions behind a +// gradient tap, and adds the surrogate + on its own K/V, where G is what every later +// chunk put there. Its backward then carries that gradient through itself into the adapter and +// into its own prefix's G. ONE optimizer step per window, on the gradient of the window's loss +// (exact within the horizon). Memory stays chunk x window; the cost is one decode plus one forward +// and backward per chunk, context chunks included. +// THE INVARIANT THAT MAKES THE REVERSE PASS VALID (Fable): the adapter does not change between the +// forward decode and the last chunk's backward. One optimizer step per window, taken by the chunk +// the reverse pass trains last, is what keeps each chunk's recomputed K/V (and recurrent state) +// equal to the cached values its successors attended to; a step inside the window would make +// every surrogate a gradient of a different function. +// A recurrent model's state is carried the same way, one chunk at a time: the state entering a +// chunk is tapped into a host accumulator, the state it leaves meets the next chunk's gradient +// (build_rs, build_walk_state_exit). The reverse pass restores each chunk's entry state from a +// checkpoint taken at its boundary in the forward pass. +void llama_context::opt_walk::run_exact() { + // Each window starts from the REQUESTED horizon: a halving that fit one window's device budget + // must not shrink every later window and epoch, which may be shorter (Fable). + lctx.opt_walk_horizon = lctx.opt_walk_horizon_req; + // A recurrent context with rollback slots keeps several states per sequence; the walk carries + // exactly one. Refused by name, never asserted: training runs inside the serving process, and + // an assert there takes serving down (2026-10-07 04:09Z on the 5090). + if (recr != nullptr && lctx.cparams.n_rs_seq > 0) { + lctx.opt_fail("the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has " + + std::to_string(lctx.cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0"); + return; } - // the plain walk snapshots a recurrent state into a scratch sequence; the exact walk keeps host snapshots - if (recr != nullptr && cparams.n_seq_max < 2 && !(train && opt_walk_exact)) { - LLAMA_LOG_ERROR("%s: a recurrent model trains with a scratch sequence for its state snapshot: create the training context with n_seq_max >= 2\n", __func__); - opt_stop_requested.store(true); + plan.mark_exact(recr != nullptr, lctx.opt_walk_horizon); + const std::vector & chunks = plan.chunks; + const uint32_t n_pos = (uint32_t) plan.last_label + 1; + + // THE HOST BUDGET (llama_walk_plan::stride_for): the K/V accumulators and the recurrent state + // checkpoints are both known before anything runs + const size_t kv_bytes = llama_walk_acc::kv_bytes(lctx.model.hparams, n_pos); + const size_t state_bytes = llama_walk_acc::state_bytes(recr); + const size_t budget = lctx.opt_walk_host_budget; + const size_t stride = plan.stride_for(budget, kv_bytes, state_bytes); + const size_t bytes = plan.host_bytes(stride, kv_bytes, state_bytes); + lctx.opt_walk_host_bytes = std::max(lctx.opt_walk_host_bytes, bytes); + if (budget > 0 && bytes > budget) { + char why[512]; + snprintf(why, sizeof(why), + "the exact walk needs %.1f MiB of host memory for this window (%u positions, the K/V gradient alone), over its budget of %.1f MiB: " + "the K/V gradient grows with the window; neither the horizon nor the chunk size changes it", + bytes/1048576.0, n_pos, budget/1048576.0); + lctx.opt_fail(why); + return; + } + if (stride > 1) { + LLAMA_LOG_INFO("%s: the exact walk checkpoints the recurrent state every %zu chunks to fit %.1f MiB of host memory (%.1f MiB)\n", + __func__, stride, budget/1048576.0, bytes/1048576.0); + } + if (!acc.init(lctx.model.hparams, recr, n_pos, lctx.cparams)) { + lctx.opt_fail("the exact walk's gradient accumulator (host memory) could not be allocated"); return; } - if (train && opt_walk_exact) { - // THE EXACT WALK. The plain walk below steps per chunk and stops the gradient at every - // chunk boundary: a reply learns only through its own chunk's K/V, never through what - // an earlier chunk computed for it. Here the window is decoded once under the adapter - // as it is, then trained in REVERSE: each chunk's graph attends to the cache, with its - // last walk_horizon prefix positions behind a GRAD leaf, and adds the surrogate - // + on its own K/V, where G is what every later chunk put there. Its - // backward then carries that gradient through itself into the adapter and into its own - // prefix's G. ONE optimizer step per window, on the gradient of the window's loss - // (exact within the horizon). Memory stays chunk x window; the cost is one decode plus - // one forward and backward per chunk, context chunks included. - // THE INVARIANT THAT MAKES THE REVERSE PASS VALID (Fable): the adapter does not change - // between the forward decode and the last chunk's backward. One optimizer step per window, - // taken by the chunk the reverse pass trains last, is what keeps each chunk's recomputed - // K/V (and recurrent state) equal to the cached values its successors attended to; a step - // inside the window would make every surrogate a gradient of a different function. - // - // Each window starts from the REQUESTED horizon: a halving that fit one window's device - // budget must not shrink every later window and epoch, which may be shorter (Fable). - opt_walk_horizon = opt_walk_horizon_req; - // A recurrent context with rollback slots keeps several states per sequence; the walk - // carries exactly one. Refused by name, never asserted: training runs inside the serving - // process, and an assert there takes serving down (2026-10-07 04:09Z on the 5090). - if (recr != nullptr && cparams.n_rs_seq > 0) { - opt_failure = "the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has " - + std::to_string(cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0"; - LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); + // the window as the adapter now reads it, chunk by chunk: a recurrent model's state is + // checkpointed at each chunk's start through the sequence-state API (the recurrent part only, + // PARTIAL_ONLY), after waiting for the queued decode; a restore puts back the cell's position + // with its rows, and an empty checkpoint (before the first chunk) restores the zero state. A + // hand-rolled per-layer copy here skipped that wait and, on Metal, carried a torn mid-graph + // state into the reverse pass (2026-10-11). + std::vector> snaps(recr != nullptr ? chunks.size() : 0); + const llama_state_seq_flags state_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; + for (size_t j = 0; j < chunks.size(); ++j) { + if (recr != nullptr && j % stride == 0) { + lctx.synchronize(); + snaps[j].resize(lctx.state_seq_get_size(0, state_flags)); + GGML_ASSERT(lctx.state_seq_get_data(0, snaps[j].data(), snaps[j].size(), state_flags) == snaps[j].size()); + } + if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n)) { return; } - // A recurrent model's state is carried the same way, one chunk at a time: the state - // entering a chunk is tapped into a host accumulator, the state it leaves meets the next chunk's gradient - // (build_rs, build_walk_state_exit). The reverse pass restores each chunk's entry state - // from a host snapshot taken at its boundary in the forward pass. - std::vector chunks; - uint32_t n_labels_window = 0; - for (uint32_t p = 0; p <= (uint32_t) last_label; ) { - const bool labelled = labels_sparse[p] >= 0; - uint32_t end = p; - while (end <= (uint32_t) last_label && (labels_sparse[end] >= 0) == labelled) { - ++end; - } - for (uint32_t c0 = p; c0 < end; c0 += n_batch) { - const uint32_t c1 = std::min(c0 + n_batch, end); - chunks.push_back({ c0, c1 - c0, labelled ? c1 - c0 : 0, 0, false, false, false, 0.0f }); - n_labels_window += labelled ? c1 - c0 : 0; - } - p = end; + } + for (int64_t j = (int64_t) chunks.size() - 1; j >= plan.first; --j) { + const llama_walk_chunk & c = chunks[j]; + if (!c.needed) { + continue; } - // which chunks train, and how far back each one's gradient reaches - uint32_t reach = UINT32_MAX; // the lowest position a later trained chunk's gradient reaches - int64_t first = -1; // the last chunk the reverse pass trains: the step - for (int64_t j = (int64_t) chunks.size() - 1; j >= 0; --j) { - exact_chunk & c = chunks[j]; - c.surrogate = reach < c.c0 + c.n; - // a recurrent state reaches every later chunk: a chunk before a trained one trains - c.needed = c.n_labels > 0 || c.surrogate || (recr != nullptr && first >= 0); - if (!c.needed) { - continue; - } - c.grad_from = opt_walk_horizon == 0 || c.c0 < opt_walk_horizon ? 0 : c.c0 - opt_walk_horizon; - c.loss_scale = (float) c.n_labels / (float) n_labels_window; - reach = std::min(reach, c.grad_from); - first = j; + if (!yield_point()) { + break; } - GGML_ASSERT(first >= 0); - chunks[first].period_end = true; - - // THE HOST BUDGET (Fable): the accumulators are layers x positions x (k + v) floats, and a - // recurrent state is checkpointed every `stride` chunks (one sequence's state rows each; - // the first chunk's entry is the zero state and costs nothing). Both are known before - // anything runs. Under a budget the stride is the smallest that fits: the reverse pass - // rebuilds a chunk's entry state by decoding forward from its checkpoint, so a longer - // stride trades host memory for decode, never for exactness. - size_t stride = 1; - { - size_t kv_bytes = 0; - for (uint32_t il = 0; il < model.hparams.n_layer(); ++il) { - if (!model.hparams.is_recr(il)) { - kv_bytes += (size_t) (last_label + 1)*(model.hparams.n_embd_k_gqa(il) + model.hparams.n_embd_v_gqa(il))*sizeof(float); - } + // attention: pop the cache to [0, c0); recurrent: restore the state at the chunk's + // checkpoint and decode forward to its start (that decode rebuilds the attention K/V in + // between too, under the same adapter: the same values). Run again before every retry + // below: a refused chunk graph is refused AFTER its ubatch was applied, so the chunk's + // cells already sit in the cache and the recurrent cell's position has moved to the + // chunk's end; retrying on top of that fails to prepare the ubatch at all (the 5090, + // 2026-10-10: "init_batch: failed to prepare attention ubatches" after "retrying at 512", + // and the epoch stopped instead of shrinking the horizon). Always a plain decode: the + // walk's graph flags are off whenever it runs. + const auto rewind = [&]() -> bool { + const size_t cp = recr != nullptr ? (size_t) j - (size_t) j % stride : (size_t) j; + const uint32_t p_pop = chunks[cp].c0; + const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : lctx.memory->seq_rm(0, p_pop, -1); + if (!popped) { + LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u) for the reverse pass\n", __func__, p_pop); + lctx.opt_stop_requested.store(true); + return false; } - size_t state = 0; if (recr != nullptr) { - for (size_t il = 0; il < recr->r_l.size(); ++il) { - state += recr->r_l[il] ? recr->r_l[il]->nb[1] : 0; - state += recr->s_l[il] ? recr->s_l[il]->nb[1] : 0; - } - } - auto bytes_at = [&](size_t s) { return kv_bytes + ((chunks.size() + s - 1)/s - 1)*state; }; - if (opt_walk_host_budget > 0) { - while (stride < chunks.size() && bytes_at(stride) > opt_walk_host_budget) { - ++stride; - } - } - const size_t bytes = bytes_at(stride); - opt_walk_host_bytes = std::max(opt_walk_host_bytes, bytes); - if (opt_walk_host_budget > 0 && bytes > opt_walk_host_budget) { - char why[512]; - snprintf(why, sizeof(why), - "the exact walk needs %.1f MiB of host memory for this window (%lld positions, the K/V gradient alone), over its budget of %.1f MiB: " - "the K/V gradient grows with the window; neither the horizon nor the chunk size changes it", - bytes/1048576.0, (long long) (last_label + 1), opt_walk_host_budget/1048576.0); - opt_failure = why; - LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return; - } - if (stride > 1) { - LLAMA_LOG_INFO("%s: the exact walk checkpoints the recurrent state every %zu chunks to fit %.1f MiB of host memory (%.1f MiB)\n", - __func__, stride, opt_walk_host_budget/1048576.0, bytes/1048576.0); - } - } - { - const uint32_t n_layer = model.hparams.n_layer(); - walk_acc.k.assign(n_layer, nullptr); - walk_acc.v.assign(n_layer, nullptr); - walk_acc.ctx.reset(ggml_init({ 4*n_layer*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true })); - walk_acc.state.clear(); - walk_acc.state_ready = false; - for (uint32_t il = 0; il < n_layer; ++il) { - if (model.hparams.is_recr(il)) { - // a recurrent layer has no K/V: its state carries the chain, one sequence's row - if (recr != nullptr) { - walk_acc.state[recr->r_l[il]] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_r(), 1); - walk_acc.state[recr->s_l[il]] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_s(), 1); - } - continue; + lctx.synchronize(); + GGML_ASSERT(lctx.state_seq_set_data(0, snaps[cp].data(), snaps[cp].size(), state_flags) == snaps[cp].size()); + if (p_pop < c.c0 && !decode_span(p_pop, c.c0)) { + return false; } - walk_acc.k[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_k_gqa(il), last_label + 1); - walk_acc.v[il] = ggml_new_tensor_2d(walk_acc.ctx.get(), GGML_TYPE_F32, model.hparams.n_embd_v_gqa(il), last_label + 1); } - walk_acc.buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(walk_acc.ctx.get(), ggml_backend_cpu_buffer_type())); - if (!walk_acc.buf) { - opt_failure = "the exact walk's gradient accumulator (host memory) could not be allocated"; - LLAMA_LOG_ERROR("%s: %s\n", __func__, opt_failure.c_str()); - opt_alloc_failed.store(true); - opt_stop_requested.store(true); - return; - } - ggml_backend_buffer_clear(walk_acc.buf.get(), 0); - cparams.walk_acc_k = walk_acc.k.data(); - cparams.walk_acc_v = walk_acc.v.data(); - } - - // the window as the adapter now reads it, chunk by chunk: a recurrent model's state is - // checkpointed at each chunk's start through the sequence-state API (the recurrent part - // only, PARTIAL_ONLY), after waiting for the queued decode; a restore puts back the - // cell's position with its rows, and an empty checkpoint (before the first chunk) - // restores the zero state. A hand-rolled per-layer copy here skipped that wait and, on - // Metal, carried a torn mid-graph state into the reverse pass (2026-10-11). - std::vector> snaps(recr != nullptr ? chunks.size() : 0); - const llama_state_seq_flags state_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; - auto save_state = [&](std::vector & snap) { - synchronize(); - snap.resize(state_seq_get_size(0, state_flags)); - GGML_ASSERT(state_seq_get_data(0, snap.data(), snap.size(), state_flags) == snap.size()); - }; - auto load_state = [&](const std::vector & snap) { - synchronize(); - GGML_ASSERT(state_seq_set_data(0, snap.data(), snap.size(), state_flags) == snap.size()); + return true; }; - for (size_t j = 0; j < chunks.size(); ++j) { - if (recr != nullptr && j % stride == 0) { - save_state(snaps[j]); - } - if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n)) { - return; - } + if (!rewind()) { + break; } - bool ok = true; - for (int64_t j = (int64_t) chunks.size() - 1; j >= first && ok; --j) { - const exact_chunk & c = chunks[j]; - if (!c.needed) { - continue; - } - if (!yield_point()) { - ok = false; + const bool state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && acc.state_ready; + // THE HORIZON FITS THE DEVICE. Its gradient taps are the chunk graph's largest growth with + // the window (layers x horizon x (k + v) floats, twice with their gradient); the graph + // preflight refuses a chunk that does not fit before allocating anything, so a memory + // refusal halves the horizon (whole chunks) and the chunk is tried again; every later + // chunk keeps the smaller one. The forward still attends to the whole window and the + // recurrent state is carried exactly; only attention gradient past the horizon is dropped. + // A refusal at a horizon of one chunk, or for a node the device cannot run, stands. + llama_walk_chunk cc = c; + bool ok; + for (;;) { + cc.grad_from = llama_walk_plan::grad_from(cc.c0, lctx.opt_walk_horizon); + set_flags(true, cc.grad_from, c.surrogate, state_surrogate); + xc = &cc; + ok = train_chunk(cc.c0, cc.n); + xc = nullptr; + set_flags(false, 0, false, false); + const uint32_t span = cc.c0 - cc.grad_from; + if (ok || !lctx.opt_failure.empty() || !lctx.opt_alloc_failed.load() || span <= n_batch) { break; } - // attention: pop the cache to [0, c0); recurrent: restore the state at the chunk's - // checkpoint and decode forward to its start (that decode rebuilds the attention K/V - // in between too, under the same adapter: the same values). Run again before every - // retry below: a refused chunk graph is refused AFTER its ubatch was applied, so the - // chunk's cells already sit in the cache and the recurrent cell's position has moved - // to the chunk's end; retrying on top of that fails to prepare the ubatch at all - // (the 5090, 2026-10-10: "init_batch: failed to prepare attention ubatches" after - // "retrying at 512", and the epoch stopped instead of shrinking the horizon). - const auto rewind = [&]() -> bool { - const size_t cp = recr != nullptr ? (size_t) j - (size_t) j % stride : (size_t) j; - const uint32_t p_pop = chunks[cp].c0; - const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : memory->seq_rm(0, p_pop, -1); - if (!popped) { - LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u) for the reverse pass\n", __func__, p_pop); - opt_stop_requested.store(true); - return false; - } - if (recr != nullptr) { - load_state(snaps[cp]); - if (p_pop < c.c0 && !decode_span(p_pop, c.c0)) { - return false; - } - } - return true; - }; + lctx.opt_walk_horizon = std::max(n_batch, (span/2)/n_batch*n_batch); + LLAMA_LOG_WARN("%s: the chunk at %u did not fit with a gradient horizon of %u positions: retrying at %u\n", + __func__, cc.c0, span, lctx.opt_walk_horizon); + lctx.opt_alloc_failed.store(false); + lctx.opt_stop_requested.store(false); if (!rewind()) { ok = false; break; } - cparams.walk_exact = true; - cparams.walk_surrogate = c.surrogate; - cparams.walk_state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && walk_acc.state_ready; - // THE HORIZON FITS THE DEVICE. Its GRAD leaves are the chunk graph's largest growth - // with the window (layers x horizon x (k + v) floats, twice with their gradient); the - // graph preflight refuses a chunk that does not fit before allocating anything, so a - // memory refusal halves the horizon (whole chunks) and the chunk is tried again; every - // later chunk keeps the smaller one. The forward still attends to the whole window and - // the recurrent state is carried exactly; only attention gradient past the horizon is - // dropped. A refusal at a horizon of one chunk, or for a node the device cannot run, - // stands. - exact_chunk cc = c; - for (;;) { - cc.grad_from = opt_walk_horizon == 0 || cc.c0 < opt_walk_horizon ? 0 : cc.c0 - opt_walk_horizon; - cparams.walk_grad_from = cc.grad_from; - xc = &cc; - ok = train_chunk(cc.c0, cc.n); - xc = nullptr; - const uint32_t span = cc.c0 - cc.grad_from; - if (ok || !opt_failure.empty() || !opt_alloc_failed.load() || span <= n_batch) { - break; - } - opt_walk_horizon = std::max(n_batch, (span/2)/n_batch*n_batch); - LLAMA_LOG_WARN("%s: the chunk at %u did not fit with a gradient horizon of %u positions: retrying at %u\n", - __func__, cc.c0, span, opt_walk_horizon); - opt_alloc_failed.store(false); - opt_stop_requested.store(false); - // the rewind's decode is a plain forward, as before the first try: no walk flags. - // walk_grad_from is not restored here because the loop's first statements recompute - // it from the shrunken horizon (cc.grad_from, then cparams.walk_grad_from) before the - // retried train_chunk runs. - const bool surrogate = cparams.walk_surrogate; - const bool state_surrogate = cparams.walk_state_surrogate; - cparams.walk_exact = false; - cparams.walk_grad_from = 0; - cparams.walk_surrogate = false; - cparams.walk_state_surrogate = false; - const bool rewound = rewind(); - cparams.walk_exact = true; - cparams.walk_surrogate = surrogate; - cparams.walk_state_surrogate = state_surrogate; - if (!rewound) { - ok = false; - break; - } - } - opt_walk_horizon_used = opt_walk_horizon; - cparams.walk_exact = false; - cparams.walk_grad_from = 0; - cparams.walk_surrogate = false; - cparams.walk_state_surrogate = false; } - walk_acc.state_ready = false; - return; + lctx.opt_walk_horizon_used = lctx.opt_walk_horizon; + if (!ok) { + break; + } } + acc.state_ready = false; +} - uint32_t pos = 0; - while (pos <= (uint32_t) last_label && !opt_stop_requested.load(std::memory_order_relaxed)) { - const bool labelled = labels_sparse[pos] >= 0; - uint32_t end = pos; - while (end <= (uint32_t) last_label && (labels_sparse[end] >= 0) == labelled) { - ++end; +// THE PLAIN WALK: each labelled chunk steps on its own gradient, then joins the context +void llama_context::opt_walk::run_plain() { + const int64_t last_label = plan.last_label; + for (const llama_walk_chunk & c : plan.chunks) { + if (lctx.opt_stop_requested.load(std::memory_order_relaxed)) { + return; } - if (!labelled) { - if (!decode_span(pos, end)) { + const uint32_t c0 = c.c0; + const uint32_t c1 = c.c0 + c.n; + if (c.n_labels == 0) { + if (!decode_span(c0, c1)) { return; } - } else { - for (uint32_t c0 = pos; c0 < end; c0 += n_batch) { - const uint32_t c1 = std::min(c0 + n_batch, end); - // The training forward advances a recurrent state in place, and a recurrent - // memory cannot remove a partial range: snapshot the state to the scratch - // sequence first (Fable), restore it after, then the decode below advances it - // from where the context left it. The state is per sequence and small. - const bool snapshot = recr != nullptr && c1 <= (uint32_t) last_label; - if (!yield_point()) { - return; - } - if (snapshot) { - recr->seq_cp(0, 1, -1, -1); - } - if (!train_chunk(c0, c1 - c0)) { - return; - } - if (snapshot) { - recr->seq_rm(0, -1, -1); - recr->seq_cp(1, 0, -1, -1); - recr->seq_rm(1, -1, -1); - } - // the chunk joins the context under the adapter as it now is: the training - // graph never writes the cache, so its cells hold nothing until this decode - if (c1 <= (uint32_t) last_label) { - // the chunk's cells hold no K/V (the training graph never writes the - // cache): drop them from the attention memory, then decode them for real - const bool removed = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, c0, -1)) : memory->seq_rm(0, c0, -1); - if (!removed) { - LLAMA_LOG_ERROR("%s: could not remove the trained chunk [%u, %u) from the cache before decoding it\n", __func__, c0, c1); - opt_stop_requested.store(true); - return; - } - if (!decode_span(c0, c1)) { - opt_stop_requested.store(true); - return; - } - } + continue; + } + // The training forward advances a recurrent state in place, and a recurrent memory + // cannot remove a partial range: snapshot the state to the scratch sequence first + // (Fable), restore it after, then the decode below advances it from where the context + // left it. The cache's own copy-on-write keeps it on the device; the state is per + // sequence and small. + const bool snapshot = recr != nullptr && c1 <= (uint32_t) last_label; + if (!yield_point()) { + return; + } + if (snapshot) { + recr->seq_cp(0, 1, -1, -1); + } + if (!train_chunk(c0, c.n)) { + return; + } + if (snapshot) { + recr->seq_rm(0, -1, -1); + recr->seq_cp(1, 0, -1, -1); + recr->seq_rm(1, -1, -1); + } + // the chunk joins the context under the adapter as it now is: the training graph never + // writes the cache, so its cells hold nothing until this decode + if (c1 <= (uint32_t) last_label) { + // drop the chunk's empty cells from the attention memory, then decode them for real + const bool removed = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, c0, -1)) : lctx.memory->seq_rm(0, c0, -1); + if (!removed) { + LLAMA_LOG_ERROR("%s: could not remove the trained chunk [%u, %u) from the cache before decoding it\n", __func__, c0, c1); + lctx.opt_stop_requested.store(true); + return; + } + if (!decode_span(c0, c1)) { + lctx.opt_stop_requested.store(true); + return; } } - pos = end; } } +void llama_context::opt_epoch_iter( + ggml_opt_dataset_t dataset, + ggml_opt_result_t result, + const std::vector & tokens, + const std::vector & labels_sparse, + llama_batch & batch, + ggml_opt_epoch_callback callback, + bool train, + int64_t idata_in_loop, + int64_t ndata_in_loop, + int64_t t_loop_start) { + GGML_ASSERT(opt_ctx); + const uint32_t n_ctx = opt_n_ctx_train; + const uint32_t n_batch = std::min(this->n_batch(), n_ctx); + + memory->clear(true); + + opt_walk walk(*this, dataset, result, tokens, labels_sparse, batch, callback, train, + idata_in_loop, ndata_in_loop, t_loop_start, n_ctx, n_batch); + if (walk.plan.last_label < 0) { + LLAMA_LOG_ERROR("%s: a training window with no labelled position: nothing to learn from it\n", __func__); + return; + } + if (auto * hybrid = dynamic_cast(memory.get())) { + walk.recr = hybrid->get_mem_recr(); + walk.attn = hybrid->get_mem_attn(); + } else { + walk.recr = dynamic_cast(memory.get()); + } + if (train && opt_walk_exact) { + walk.run_exact(); + return; + } + // the plain walk snapshots a recurrent state into a scratch sequence + if (walk.recr != nullptr && cparams.n_seq_max < 2) { + LLAMA_LOG_ERROR("%s: a recurrent model trains with a scratch sequence for its state snapshot: create the training context with n_seq_max >= 2\n", __func__); + opt_stop_requested.store(true); + return; + } + walk.run_plain(); +} + +void llama_context::opt_fail(std::string why, const std::string & what) { + LLAMA_LOG_ERROR("%s: %s%s%s: stopping the epoch\n", __func__, what.c_str(), what.empty() ? "" : ": ", + why.empty() ? "it does not fit in device memory" : why.c_str()); + opt_failure = std::move(why); + opt_alloc_failed.store(true); + opt_stop_requested.store(true); +} + void llama_context::opt_epoch( ggml_opt_dataset_t dataset, ggml_opt_result_t result_train, @@ -769,48 +663,31 @@ void llama_context::opt_epoch( std::vector tokens(n_ctx); std::vector labels_sparse(n_ctx); - int64_t idata = 0; - - int64_t t_loop_start = ggml_time_us(); - int64_t ndata_in_loop = idata_split*ubatch_per_ctx; - for (; idata < idata_split; ++idata) { - if (opt_stop_requested.load(std::memory_order_relaxed)) { - break; - } - constexpr bool train = true; - const int64_t idata_in_loop = idata*ubatch_per_ctx; - - if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { - opt_stop_requested.store(true); - } - if (opt_stop_requested.load(std::memory_order_relaxed)) { - break; - } - - ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); - opt_epoch_iter(dataset, result_train, tokens, labels_sparse, batch, - callback_train, train, idata_in_loop, ndata_in_loop, t_loop_start); - } - - t_loop_start = ggml_time_us(); - ndata_in_loop = (ndata - idata_split)*ubatch_per_ctx; - for (; idata < ndata; ++idata) { - if (opt_stop_requested.load(std::memory_order_relaxed)) { - break; - } - constexpr bool train = false; - const int64_t idata_in_loop = (idata - idata_split)*ubatch_per_ctx; - - if (opt_step_callback && !opt_step_callback(train, opt_step_callback_data)) { - opt_stop_requested.store(true); - } - if (opt_stop_requested.load(std::memory_order_relaxed)) { - break; + // the training range [0, idata_split), then the evaluation range [idata_split, ndata) + const struct { + bool train; + int64_t i0, i1; + ggml_opt_result_t result; + ggml_opt_epoch_callback callback; + } ranges[] = { + { true, 0, idata_split, result_train, callback_train }, + { false, idata_split, ndata, result_eval, callback_eval }, + }; + for (const auto & r : ranges) { + const int64_t t_loop_start = ggml_time_us(); + const int64_t ndata_in_loop = (r.i1 - r.i0)*ubatch_per_ctx; + for (int64_t idata = r.i0; idata < r.i1; ++idata) { + if (opt_step_callback && !opt_stop_requested.load(std::memory_order_relaxed) + && !opt_step_callback(r.train, opt_step_callback_data)) { + opt_stop_requested.store(true); + } + if (opt_stop_requested.load(std::memory_order_relaxed)) { + break; + } + ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); + opt_epoch_iter(dataset, r.result, tokens, labels_sparse, batch, + r.callback, r.train, (idata - r.i0)*ubatch_per_ctx, ndata_in_loop, t_loop_start); } - - ggml_opt_dataset_get_batch_host(dataset, tokens.data(), n_ctx*sizeof(llama_token), labels_sparse.data(), idata); - opt_epoch_iter(dataset, result_eval, tokens, labels_sparse, batch, - callback_eval, train, idata_in_loop, ndata_in_loop, t_loop_start); } llama_batch_free(batch); diff --git a/src/llama-context.h b/src/llama-context.h index 26ec2796cf1a..2f3fc48b6efd 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -236,6 +236,14 @@ struct llama_context { int64_t t_loop_start); private: + // one window's walk (llama-context-opt.cpp) + struct opt_walk; + + // stops the epoch on a failure: llama_opt_failed, and llama_opt_failure reads `why` (empty + // for the memory gate, which the exact walk answers by shrinking its horizon); `what` names + // the step that failed, for the log + void opt_fail(std::string why, const std::string & what = ""); + // // output // diff --git a/src/llama-walk.cpp b/src/llama-walk.cpp new file mode 100644 index 000000000000..481e8a8e9078 --- /dev/null +++ b/src/llama-walk.cpp @@ -0,0 +1,142 @@ +#include "llama-walk.h" + +#include "llama-cparams.h" +#include "llama-hparams.h" +#include "llama-memory-recurrent.h" + +#include "ggml-backend.h" + +#include + +llama_walk_plan llama_walk_plan::of(const std::vector & labels, uint32_t n_ctx, uint32_t n_batch) { + GGML_ASSERT(n_batch > 0 && labels.size() >= n_ctx); + llama_walk_plan plan; + for (int64_t i = (int64_t) n_ctx - 1; i >= 0; --i) { + if (labels[i] >= 0) { + plan.last_label = i; + break; + } + } + for (uint32_t p = 0; (int64_t) p <= plan.last_label; ) { + const bool labelled = labels[p] >= 0; + uint32_t end = p; + while ((int64_t) end <= plan.last_label && (labels[end] >= 0) == labelled) { + ++end; + } + for (uint32_t c0 = p; c0 < end; c0 += n_batch) { + const uint32_t n = std::min(c0 + n_batch, end) - c0; + plan.chunks.push_back({ c0, n, labelled ? n : 0 }); + plan.n_labels += labelled ? n : 0; + } + p = end; + } + return plan; +} + +uint32_t llama_walk_plan::grad_from(uint32_t c0, uint32_t horizon) { + return horizon == 0 || c0 < horizon ? 0 : c0 - horizon; +} + +void llama_walk_plan::mark_exact(bool recurrent, uint32_t horizon) { + uint32_t reach = UINT32_MAX; // the lowest position a later trained chunk's gradient reaches + first = -1; + for (int64_t j = (int64_t) chunks.size() - 1; j >= 0; --j) { + llama_walk_chunk & c = chunks[j]; + c.surrogate = reach < c.c0 + c.n; + c.needed = c.n_labels > 0 || c.surrogate || (recurrent && first >= 0); + c.period_end = false; + if (!c.needed) { + continue; + } + c.grad_from = grad_from(c.c0, horizon); + c.loss_scale = (float) c.n_labels / (float) n_labels; + reach = std::min(reach, c.grad_from); + first = j; + } + GGML_ASSERT(first >= 0); + chunks[first].period_end = true; +} + +size_t llama_walk_plan::host_bytes(size_t stride, size_t kv_bytes, size_t state_bytes) const { + const size_t n_checkpoints = chunks.empty() ? 0 : (chunks.size() + stride - 1)/stride - 1; + return kv_bytes + n_checkpoints*state_bytes; +} + +size_t llama_walk_plan::stride_for(size_t budget, size_t kv_bytes, size_t state_bytes) const { + size_t stride = 1; + if (budget > 0) { + while (stride < chunks.size() && host_bytes(stride, kv_bytes, state_bytes) > budget) { + ++stride; + } + } + return stride; +} + +size_t llama_walk_acc::kv_bytes(const llama_hparams & hparams, uint32_t n_pos) { + size_t bytes = 0; + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + if (!hparams.is_recr(il)) { + bytes += (size_t) n_pos*(hparams.n_embd_k_gqa(il) + hparams.n_embd_v_gqa(il))*sizeof(float); + } + } + return bytes; +} + +size_t llama_walk_acc::state_bytes(const llama_memory_recurrent * recr) { + size_t bytes = 0; + if (recr != nullptr) { + for (size_t il = 0; il < recr->r_l.size(); ++il) { + bytes += recr->r_l[il] ? recr->r_l[il]->nb[1] : 0; + bytes += recr->s_l[il] ? recr->s_l[il]->nb[1] : 0; + } + } + return bytes; +} + +bool llama_walk_acc::init(const llama_hparams & hparams, const llama_memory_recurrent * recr, uint32_t n_pos, llama_cparams & cparams) { + const uint32_t n_layer = hparams.n_layer(); + k.assign(n_layer, nullptr); + v.assign(n_layer, nullptr); + state.clear(); + state_ready = false; + ctx.reset(ggml_init({ 4*n_layer*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true })); + for (uint32_t il = 0; il < n_layer; ++il) { + if (hparams.is_recr(il)) { + // a recurrent layer has no K/V: its state carries the chain, one sequence's row + if (recr != nullptr) { + state[recr->r_l[il]] = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, hparams.n_embd_r(), 1); + state[recr->s_l[il]] = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, hparams.n_embd_s(), 1); + } + continue; + } + k[il] = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, hparams.n_embd_k_gqa(il), n_pos); + v[il] = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F32, hparams.n_embd_v_gqa(il), n_pos); + } + buf.reset(ggml_backend_alloc_ctx_tensors_from_buft(ctx.get(), ggml_backend_cpu_buffer_type())); + if (!buf) { + return false; + } + ggml_backend_buffer_clear(buf.get(), 0); + bound = &cparams; + cparams.walk_acc_k = k.data(); + cparams.walk_acc_v = v.data(); + return true; +} + +llama_walk_acc::~llama_walk_acc() { + if (bound != nullptr) { + bound->walk_acc_k = nullptr; + bound->walk_acc_v = nullptr; + } +} + +void llama_walk_fill(llama_batch & batch, const std::vector & tokens, uint32_t c0, uint32_t n) { + batch.n_tokens = n; + for (uint32_t i = 0; i < n; ++i) { + batch.token [i] = tokens[c0 + i]; + batch.pos [i] = c0 + i; + batch.n_seq_id[i] = 1; + batch.seq_id [i][0] = 0; + batch.logits [i] = false; + } +} diff --git a/src/llama-walk.h b/src/llama-walk.h new file mode 100644 index 000000000000..67a68bf12db7 --- /dev/null +++ b/src/llama-walk.h @@ -0,0 +1,94 @@ +#pragma once + +// THE WALK's plan and its host state, apart from the context that runs it +// (llama_context::opt_epoch_iter): which positions train in which chunk, how far back each +// chunk's gradient reaches, and what the exact walk keeps in host memory. The plan is plain +// data from the labels and the model's shape, so it is tested without a model +// (test-walk-plan). + +#include "llama.h" + +#include "ggml-cpp.h" + +#include +#include +#include +#include + +struct llama_cparams; +struct llama_hparams; +class llama_memory_recurrent; + +// at most n_batch positions of one run of the window, all labelled or all context +struct llama_walk_chunk { + uint32_t c0; + uint32_t n; + uint32_t n_labels; + // the exact walk's marks (llama_walk_plan::mark_exact) + uint32_t grad_from = 0; // its gradient taps cover cached positions [grad_from, c0) + bool surrogate = false; // a later chunk's gradient sits on its own K/V + bool needed = false; + bool period_end = false; // the window's optimizer step + float loss_scale = 0.0f; // its labelled positions over the window's +}; + +struct llama_walk_plan { + // the window's runs, each cut every n_batch positions from its start; both walks take + // them in this order. Nothing after the last label is in the plan. + std::vector chunks; + int64_t last_label = -1; // -1: the window has no label, nothing to learn from it + uint32_t n_labels = 0; + int64_t first = -1; // the exact walk's period end: the chunk its reverse pass trains last + + static llama_walk_plan of(const std::vector & labels, uint32_t n_ctx, uint32_t n_batch); + + // which chunks the exact walk trains (a labelled one, one a later gradient reaches, and on a + // recurrent model every one before a trained one, since the state reaches every later + // chunk), how far back each gradient reaches, and the period end + void mark_exact(bool recurrent, uint32_t horizon); + + // the first cached position a chunk at c0 carries gradient to (horizon 0: the whole prefix) + static uint32_t grad_from(uint32_t c0, uint32_t horizon); + + // THE HOST BUDGET (Fable): the K/V accumulators, plus a recurrent state checkpoint every + // `stride` chunks (the first chunk's entry is the zero state and costs nothing) + size_t host_bytes(size_t stride, size_t kv_bytes, size_t state_bytes) const; + // the smallest stride under the budget (0: no cap); the reverse pass rebuilds a chunk's + // entry state by decoding forward from its checkpoint, so a longer stride trades host + // memory for decode, never for exactness. When none fits, the stride of one checkpoint. + size_t stride_for(size_t budget, size_t kv_bytes, size_t state_bytes) const; +}; + +// THE EXACT WALK's gradients, in host memory, each added into in place by the backward: per +// attention layer dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, position] +// F32 (null for a recurrent layer), which the graph reads through cparams while bound (never +// past the accumulator's life); per recurrent state tensor (a layer's conv state r_l and its +// recurrent state s_l) dL/d(the state entering the chunk trained last), which is what the +// chunk before it puts on the state it leaves (state_ready once one is in) +struct llama_walk_acc { + std::vector k, v; + std::unordered_map state; + bool state_ready = false; + + llama_walk_acc() = default; + llama_walk_acc(const llama_walk_acc &) = delete; + llama_walk_acc & operator=(const llama_walk_acc &) = delete; + ~llama_walk_acc(); + + // zeroed accumulators for n_pos positions, bound into cparams; false: the host memory + // could not be allocated + bool init(const llama_hparams & hparams, const llama_memory_recurrent * recr, uint32_t n_pos, llama_cparams & cparams); + + // what init allocates for the K/V of n_pos positions, and one sequence's recurrent state + static size_t kv_bytes(const llama_hparams & hparams, uint32_t n_pos); + static size_t state_bytes(const llama_memory_recurrent * recr); + +private: + llama_cparams * bound = nullptr; + ggml_context_ptr ctx; + ggml_backend_buffer_ptr buf; +}; + +// positions [c0, c0 + n) of the window into the batch, one sequence, no output row: the +// caller marks the rows it wants out +void llama_walk_fill(llama_batch & batch, const std::vector & tokens, uint32_t c0, uint32_t n); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b2bab4505e5d..b222eaf134da 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -169,6 +169,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-grammar-integration.cpp) llama_build_and_test(test-llama-grammar.cpp) llama_build_and_test(test-batch-alloc.cpp) + llama_build_and_test(test-walk-plan.cpp) llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}) target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server) target_link_libraries(test-chat PRIVATE server-context) diff --git a/tests/test-walk-plan.cpp b/tests/test-walk-plan.cpp new file mode 100644 index 000000000000..ed3749131194 --- /dev/null +++ b/tests/test-walk-plan.cpp @@ -0,0 +1,74 @@ +// The walk's plan without a model (src/llama-walk.h): what test-walk-exact proves on a model +// rests on these chunks, marks and the host budget being right. +#include "testing.h" + +#include "llama.h" + +#include "../src/llama-walk.h" + +#include + +// a window of runs: {n positions, labelled}; labelled positions take token 1, context -1 +static std::vector window(std::initializer_list> runs) { + std::vector labels; + for (const auto & [n, labelled] : runs) { + labels.insert(labels.end(), n, labelled ? 1 : -1); + } + return labels; +} + +int main(int argc, char ** argv) { + testing t(std::cout); + t.verbose = argc > 1 && std::string(argv[1]) == "-v"; + + // what this catches: a chunk that mixes context and reply, or a plan that runs past the last label + t.test("chunks_follow_runs", [](testing & t) { + const auto labels = window({ {300, false}, {700, true}, {100, false}, {50, true}, {200, false} }); + const auto plan = llama_walk_plan::of(labels, labels.size(), 512); + t.assert_equal(1149, plan.last_label); + t.assert_equal(750u, plan.n_labels); + const std::vector> want = { {0, 300}, {300, 512}, {812, 188}, {1000, 100}, {1100, 50} }; + t.assert_equal(want.size(), plan.chunks.size()); + for (size_t j = 0; j < want.size() && j < plan.chunks.size(); ++j) { + t.assert_equal(want[j].first, plan.chunks[j].c0); + t.assert_equal(want[j].second, plan.chunks[j].n); + t.assert_equal(labels[want[j].first] >= 0 ? want[j].second : 0u, plan.chunks[j].n_labels); + } + t.assert_equal(-1, llama_walk_plan::of(window({ {64, false} }), 64, 512).last_label); + }); + + // what this catches: a step taken anywhere but the reverse pass's last chunk, a loss that is + // not the window's mean, a gradient reaching past the horizon + t.test("exact_marks", [](testing & t) { + auto plan = llama_walk_plan::of(window({ {1024, false}, {512, true}, {512, false}, {512, true} }), 2560, 512); + plan.mark_exact(/*recurrent =*/ false, /*horizon =*/ 512); + float loss = 0.0f; + int ends = 0; + for (const auto & c : plan.chunks) { + t.assert_true("every chunk before a reply carries its gradient", c.needed); + t.assert_equal(llama_walk_plan::grad_from(c.c0, 512), c.grad_from); + loss += c.loss_scale; + ends += c.period_end; + } + t.assert_equal(0, plan.first); + t.assert_equal(1, ends); + t.assert_true("loss scales sum to one", loss > 0.9999f && loss < 1.0001f); + t.assert_true("the last chunk has no later gradient", !plan.chunks.back().surrogate); + t.assert_equal(0u, llama_walk_plan::grad_from(1024, 0)); + t.assert_equal(512u, llama_walk_plan::grad_from(1024, 512)); + }); + + // what this catches: a stride that does not shrink the checkpoints to the budget, or one + // that grows past the window when nothing fits + t.test("host_budget_stride", [](testing & t) { + const auto plan = llama_walk_plan::of(window({ {8*512, true} }), 8*512, 512); + t.assert_equal(170u, plan.host_bytes(1, 100, 10)); + t.assert_equal(1u, plan.stride_for(0, 100, 10)); + t.assert_equal(2u, plan.stride_for(150, 100, 10)); + t.assert_equal(130u, plan.host_bytes(2, 100, 10)); + t.assert_equal(8u, plan.stride_for(50, 100, 10)); + t.assert_equal(100u, plan.host_bytes(8, 100, 10)); + }); + + return t.summary(); +}