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..87cd73369d30 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; @@ -305,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, @@ -320,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; } @@ -364,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)); } } @@ -408,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); @@ -677,6 +692,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 +703,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 +828,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 +975,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..5f572229d5f9 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]; @@ -7021,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)); @@ -7295,6 +7494,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; @@ -7482,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; } @@ -7492,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); @@ -8083,7 +8293,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/CMakeLists.txt b/src/CMakeLists.txt index 8922dc12adc7..cd429130227d 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -16,6 +16,8 @@ add_library(llama llama-batch.cpp 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 new file mode 100644 index 000000000000..3c183f3d255c --- /dev/null +++ b/src/llama-context-opt.cpp @@ -0,0 +1,781 @@ +// 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" +#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); + } + } +} + +// 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; + + 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); +} + +// 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; + } + 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; +} + +// 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); + } + + // 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; + } + + lctx.n_queued_tokens += lctx.balloc->get_n_tokens(); + + lctx.embd_seq.clear(); + + const uint32_t n_outputs_all = lctx.balloc->get_n_outputs(); + // Unreachable: the walk trains only chunks of a labelled run. It RETURNS from the whole + // window, since a later chunk would need this one in the cache. + if (n_outputs_all == 0) { + LLAMA_LOG_ERROR("%s: a training batch with no labelled position (every label masked): nothing to learn from it\n", __func__); + return false; + } + + auto mctx = lctx.memory->init_batch(*lctx.balloc, cparams.n_ubatch, true); + if (!mctx || mctx->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { + LLAMA_LOG_ERROR("%s: could not initialize batch\n", __func__); + return false; + } + + // reserve output buffer + if (lctx.output_reserve(n_outputs_all) < n_outputs_all) { + LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); + GGML_ABORT("TODO: handle this error"); + }; + + uint32_t pos_batch = 0; + do { + const auto & ubatch = mctx->get_ubatch(); + + // 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; + } + + if (!mctx->apply()) { + LLAMA_LOG_ERROR("%s: failed to update the memory context\n", __func__); + return false; + } + + auto * res = lctx.gf_res_prev.get(); + + const auto gparams = lctx.graph_params(res, ubatch, mctx.get(), lctx.graph_type()); + + res->reset(); + + auto * gf = model.build_graph(gparams); + + struct ggml_context * ctx_compute_opt; + { + const size_t size_gf = ggml_graph_size(gf); + const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); + struct ggml_init_params params = { + /*.mem_size =*/ size_meta, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + ctx_compute_opt = ggml_init(params); + } + ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); + if (xc != nullptr) { + ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); + for (const auto & io : res->t_walk) { + if (io.tk) { + ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); + ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); + } + } + for (const auto & ws : res->t_walk_state) { + if (ws.tap) { + ggml_opt_accumulate_grad(opt_ctx, ws.tap, 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_t = ggml_opt_labels(opt_ctx); + GGML_ASSERT(labels_t->ne[1] == (int64_t) lctx.n_outputs); + ggml_set_zero(labels_t); + const float onef = 1.0f; + uint32_t row = 0; + for (uint32_t pos_ubatch = 0; pos_ubatch < ubatch.n_tokens; ++pos_ubatch) { + if (ubatch.output[pos_ubatch] == 0) { + continue; + } + const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch; + // an exact walk's context chunk: its one row has no label (weighted zero) + GGML_ASSERT((xc != nullptr && labels[ilabel] < 0) || (labels[ilabel] >= 0 && labels[ilabel] < labels_t->ne[0])); + if (labels[ilabel] >= 0) { + ggml_backend_tensor_set(labels_t, &onef, (row*labels_t->ne[0] + labels[ilabel])*sizeof(float), sizeof(float)); + } + ++row; + } + 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; +} + +// 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; + } + 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; + } + + // 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; + } + } + for (int64_t j = (int64_t) chunks.size() - 1; j >= plan.first; --j) { + const llama_walk_chunk & c = chunks[j]; + if (!c.needed) { + continue; + } + if (!yield_point()) { + break; + } + // attention: pop the cache to [0, c0); recurrent: restore the state at the chunk's + // checkpoint and decode forward to its start (that decode rebuilds the attention K/V in + // between too, under the same adapter: the same values). Run again before every retry + // below: a refused chunk graph is refused AFTER its ubatch was applied, so the chunk's + // cells already sit in the cache and the recurrent cell's position has moved to the + // chunk's end; retrying on top of that fails to prepare the ubatch at all (the 5090, + // 2026-10-10: "init_batch: failed to prepare attention ubatches" after "retrying at 512", + // and the epoch stopped instead of shrinking the horizon). Always a plain decode: the + // walk's graph flags are off whenever it runs. + const auto rewind = [&]() -> bool { + const size_t cp = recr != nullptr ? (size_t) j - (size_t) j % stride : (size_t) j; + const uint32_t p_pop = chunks[cp].c0; + const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : lctx.memory->seq_rm(0, p_pop, -1); + if (!popped) { + LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u) for the reverse pass\n", __func__, p_pop); + lctx.opt_stop_requested.store(true); + return false; + } + if (recr != nullptr) { + lctx.synchronize(); + GGML_ASSERT(lctx.state_seq_set_data(0, snaps[cp].data(), snaps[cp].size(), state_flags) == snaps[cp].size()); + if (p_pop < c.c0 && !decode_span(p_pop, c.c0)) { + return false; + } + } + return true; + }; + if (!rewind()) { + break; + } + const bool state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && acc.state_ready; + // THE HORIZON FITS THE DEVICE. Its gradient taps are the chunk graph's largest growth with + // the window (layers x horizon x (k + v) floats, twice with their gradient); the graph + // preflight refuses a chunk that does not fit before allocating anything, so a memory + // refusal halves the horizon (whole chunks) and the chunk is tried again; every later + // chunk keeps the smaller one. The forward still attends to the whole window and the + // recurrent state is carried exactly; only attention gradient past the horizon is dropped. + // A refusal at a horizon of one chunk, or for a node the device cannot run, stands. + llama_walk_chunk cc = c; + bool ok; + for (;;) { + cc.grad_from = llama_walk_plan::grad_from(cc.c0, lctx.opt_walk_horizon); + set_flags(true, cc.grad_from, c.surrogate, state_surrogate); + xc = &cc; + ok = train_chunk(cc.c0, cc.n); + xc = nullptr; + set_flags(false, 0, false, false); + const uint32_t span = cc.c0 - cc.grad_from; + if (ok || !lctx.opt_failure.empty() || !lctx.opt_alloc_failed.load() || span <= n_batch) { + break; + } + 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; + } + } + lctx.opt_walk_horizon_used = lctx.opt_walk_horizon; + if (!ok) { + break; + } + } + acc.state_ready = false; +} + +// 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; + } + 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; + } + 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; + } + } + } +} + +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, + 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); + + // 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); + } + } + + 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 61ae0cb40ec6..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,835 +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"); - // 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. - 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. - { - const size_t max_nodes = graph_max_nodes(n_ubatch); - 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 GRAD leaves 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 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; - - // 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); - } - 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 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) { - ggml_backend_tensor_memset(ws.ds, 0, 0, ggml_nbytes(ws.ds)); - } - 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_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) { - // 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) { - 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 (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 a GRAD leaf, 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); - } - } - 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 - } - 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); - } - - // 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 - }; - 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); - } - } - } - }; - for (size_t j = 0; j < chunks.size(); ++j) { - if (recr != nullptr && j % stride == 0) { - state_rows(false, 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) { - if (snaps[cp].empty) { - recr->seq_rm(0, -1, -1); // the first chunk starts from a zero state - } else { - state_rows(true, 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_gstate.empty(); - // 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_gk.clear(); - walk_gv.clear(); - walk_gstate.clear(); - 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 // @@ -4808,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..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 // @@ -279,6 +287,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 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..9e9071876ab9 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()); @@ -2578,10 +2596,13 @@ 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 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"); @@ -2589,12 +2610,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 +2628,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 @@ -2841,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 @@ -2875,18 +2945,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 +2973,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 +2992,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); @@ -3558,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 5f924306f7f6..0e0adb7d4bda 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -942,19 +942,19 @@ 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 // (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; @@ -1063,6 +1063,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; @@ -1180,6 +1182,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/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 96484141528c..b222eaf134da 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) @@ -168,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) @@ -298,6 +300,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-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-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..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,13 +211,32 @@ 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. -// 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; @@ -216,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 @@ -235,6 +274,47 @@ 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); + } + std::remove(init.c_str()); + return 0; + } + int failures = 0; { // 1. every position labelled: chunk = window is one graph, the true gradient @@ -246,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))) { @@ -265,21 +346,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); } } { 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(); +}