diff --git a/common/common.cpp b/common/common.cpp index 8147060448d7..bda497246329 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -2401,8 +2401,21 @@ bool common_lora_write_fresh(const llama_model * model, const std::string & path std::vector shapes; const int32_t n_layer = llama_model_n_layer(model); const int32_t first = top_layers > 0 && top_layers < n_layer ? n_layer - top_layers : 0; + for (const auto & target : targets) { + if (target == "token_embd.weight") { + LOG_ERR("%s: embedding adapters require the embedding-specific layout\n", __func__); + return false; + } + const ggml_tensor * w = llama_model_get_tensor(model, target.c_str()); + if (w != nullptr && ggml_n_dims(w) == 2) { + shapes.push_back({ target, w->ne[0], w->ne[1] }); + } + } for (int32_t il = first; il < n_layer; ++il) { for (const auto & target : targets) { + if (llama_model_get_tensor(model, target.c_str()) != nullptr) { + continue; + } const std::string name = "blk." + std::to_string(il) + "." + target + ".weight"; const ggml_tensor * w = llama_model_get_tensor(model, name.c_str()); if (w != nullptr && ggml_n_dims(w) == 2) { diff --git a/common/speculative.cpp b/common/speculative.cpp index 4eef2212e751..321c6de3effb 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -633,6 +633,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, + /*.label =*/ nullptr, }; const int32_t rc = llama_encode(ctx_dft, enc_batch); if (rc != 0) { @@ -1126,6 +1127,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, + /*.label =*/ nullptr, }; int32_t rc = llama_encode(ctx_dft, enc_batch); diff --git a/examples/parallel/parallel.cpp b/examples/parallel/parallel.cpp index a46400c5b943..643a05530994 100644 --- a/examples/parallel/parallel.cpp +++ b/examples/parallel/parallel.cpp @@ -400,6 +400,8 @@ int main(int argc, char ** argv) { batch.n_seq_id + i, batch.seq_id + i, batch.logits + i, + batch.label ? batch.label + i : nullptr, + batch.label && batch.label_weight ? batch.label_weight + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index 148349e99e20..c41c791aed78 100644 --- a/ggml/include/ggml-opt.h +++ b/ggml/include/ggml-opt.h @@ -32,6 +32,10 @@ extern "C" { GGML_OPT_LOSS_TYPE_SUM, GGML_OPT_LOSS_TYPE_CROSS_ENTROPY, GGML_OPT_LOSS_TYPE_MEAN_SQUARED_ERROR, + // GGML_OPT_LOSS_TYPE_CROSS_ENTROPY with one I32 target per output row (-1: no loss) in place of a one-hot + // row as wide as the outputs: what a language model's training step needs (its outputs are + // vocabulary-wide logits) + GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS, }; // ====== Dataset ====== @@ -154,7 +158,13 @@ extern "C" { // if not using static graphs these pointers become invalid with the next call to ggml_opt_alloc GGML_API struct ggml_tensor * ggml_opt_inputs( ggml_opt_context_t opt_ctx); // forward graph input tensor GGML_API struct ggml_tensor * ggml_opt_outputs( ggml_opt_context_t opt_ctx); // forward graph output tensor - GGML_API struct ggml_tensor * ggml_opt_labels( ggml_opt_context_t opt_ctx); // labels to compare outputs against + GGML_API struct ggml_tensor * ggml_opt_labels( ggml_opt_context_t opt_ctx); + // the per-row weights of a weighted cross-entropy step (F32, one per label), or NULL when the + // graph prepared last is unweighted (ggml_opt_set_next_label_weights) + GGML_API struct ggml_tensor * ggml_opt_label_weights(ggml_opt_context_t opt_ctx); + // whether the next graph's cross-entropy weighs each row (by ggml_opt_label_weights); an + // unweighted step keeps the unweighted loss, whose kernels every backend has + GGML_API void ggml_opt_set_next_label_weights(ggml_opt_context_t opt_ctx, bool weighted); // labels to compare outputs against GGML_API struct ggml_tensor * ggml_opt_loss( ggml_opt_context_t opt_ctx); // scalar tensor that contains the loss GGML_API struct ggml_tensor * ggml_opt_pred( ggml_opt_context_t opt_ctx); // predictions made by outputs GGML_API struct ggml_tensor * ggml_opt_ncorrect(ggml_opt_context_t opt_ctx); // number of matching predictions between outputs and labels @@ -188,6 +198,48 @@ extern "C" { struct ggml_tensor * inputs, struct ggml_tensor * outputs); + // The same, for an `outputs` that is a row-wise function of `hidden` (a language model's + // head: logits = W x, with its adapters and scales; MUL_MAT, MUL, ADD and SCALE nodes): that + // head is taken out of gf and computed for `rows` rows of hidden at a time, and the loss is + // the mean over every row (GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS only). Under recompute + // (checkpoint_prefix) only one slice's outputs and their gradient exist at a time, never the + // [n_outputs x rows] tensor and its gradient. A head with any other node is not sliced. + // Returns the number of slices (1: not sliced; accuracy is reported only then). + GGML_API int ggml_opt_prepare_alloc_head( + ggml_opt_context_t opt_ctx, + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * inputs, + struct ggml_tensor * hidden, + struct ggml_tensor * outputs, + int64_t rows); + + // Project the existing row-wise head before consumers such as sampling are built. + // Row indices address hidden's token rows. A null public_rows requests no public logits. + // Returns zero without changing gf when the head cannot be replayed. + GGML_API int ggml_opt_project_head( + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * hidden, + struct ggml_tensor * outputs, + struct ggml_tensor * shadow_rows, + struct ggml_tensor * public_rows, + struct ggml_tensor * label_rows, + struct ggml_tensor * public_shadow_rows, + struct ggml_tensor ** slice_shadow_rows, + int64_t rows, + struct ggml_tensor ** public_outputs, + struct ggml_tensor ** slices, + size_t slices_capacity); + + GGML_API void ggml_opt_prepare_alloc_slices( + ggml_opt_context_t opt_ctx, + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * inputs, + struct ggml_tensor ** slices, + size_t n_slices); + // allocate the next graph for evaluation, either forward or forward + backward // must be called exactly once prior to calling ggml_opt_eval // returns false when the graph could not be allocated (device memory): the graph is then NOT diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 1c23bbf98c85..12d01193c5f1 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -2480,14 +2480,47 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * sinks); - // TODO: needs to be adapted to ggml_flash_attn_ext + // The gradients of ggml_flash_attn_ext(q, k, v, mask, scale) (no ALiBi, softcap or sinks), given + // d, the gradient of its output in the output's layout [Dv, H, N]. One op, recomputing the + // scores per block instead of keeping them: memory is the gradients alone. The result is F32: + // dQ ([D, N, H]), dK (k's shape), dV (v's shape), each contiguous at + // ggml_flash_attn_back_offset, then each query row's log-sum-exp and D = sum_j p_ij dp_ij, + // which the kernels pass between their passes, and the power of two dO is scaled by in a + // kernel's reduced-precision tiles (ggml_flash_attn_back_scale). + // dK and dV cover the key rows [kv0, M) only (row j at j - kv0): keys before kv0 are constants + // whose gradient nobody reads (a frozen served prefix), so it is neither computed nor stored. + // dQ always sums over every key. GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, struct ggml_tensor * v, + struct ggml_tensor * mask, struct ggml_tensor * d, - bool masked); + float scale, + int64_t kv0); + + // byte offset in a ggml_flash_attn_back result of: 0 dQ, 1 dK, 2 dV, 3 the per-row scratch + // the power of two that brings the largest |dO| to [4, 8): training gradients sit far below half + // precision's normal range (6.1e-5) and lose their precision in half tiles; 1 when dO is zero + // (the Metal kernel computes the same rule, kernel_flash_attn_back_scale) + GGML_API float ggml_flash_attn_back_scale(float max_abs_do); + GGML_API size_t ggml_flash_attn_back_offset(const struct ggml_tensor * back, int part); + + // Where FLASH_ATTN_EXT's K/V gradient goes when only the last rows of its K/V can carry one: + // k_tap / v_tap (or NULL) and k_cur / v_cur, each [D, Hkv, rows], hold the values of K/V's last + // rows in order, ending at M (a run's cached rows, then the chunk's own). The backward computes + // dK/dV from the first of those rows only and adds each part to its sink; the K/V the attention + // reads (the whole window) get no gradient. The sinks are not read by the forward. + GGML_API void ggml_flash_attn_ext_set_grad_sinks( + struct ggml_tensor * a, + struct ggml_tensor * k_tap, + struct ggml_tensor * k_cur, + struct ggml_tensor * v_tap, + struct ggml_tensor * v_cur); + + // the backward of FLASH_ATTN_EXT as the one op above (ggml_flash_attn_ext_set_back_tile) + #define GGML_FA_BACK_NATIVE (-1) GGML_API struct ggml_tensor * ggml_ssm_conv( struct ggml_context * ctx, @@ -2779,7 +2812,21 @@ extern "C" { int n_tasks, void * userdata); - // loss function + // loss function: the mean over rows of -sum(labels * log softmax(logits)). The labels are + // either a distribution per row (F32, the logits' shape) or one target index per row (I32, + // one per row of the logits, -1 for a row that carries no loss): a target index never builds + // a vocabulary-wide one-hot tensor beside the logits. + + GGML_API bool ggml_cross_entropy_labels_ok(const struct ggml_tensor * logits, const struct ggml_tensor * labels); + + // The same with one target index per row (I32) and a weight per row (F32, one per row): row i + // contributes w_i * -log softmax(logits_i)[t_i] / rows, and its gradient scales by w_i. How much a + // row counts is the caller's to say (a mind's relevance on what it perceived), not only whether. + GGML_API struct ggml_tensor * ggml_cross_entropy_loss_weighted( + struct ggml_context * ctx, + struct ggml_tensor * logits, + struct ggml_tensor * targets, + struct ggml_tensor * weights); GGML_API struct ggml_tensor * ggml_cross_entropy_loss( struct ggml_context * ctx, @@ -2852,6 +2899,10 @@ extern "C" { struct ggml_cgraph * cgraph, struct ggml_tensor * tensor); + // whether a backward can be built through this node (ggml_build_backward_expand aborts on one + // that cannot: an op without a backward, or an in-place write) + GGML_API bool ggml_node_has_backward(const struct ggml_tensor * tensor); + GGML_API void ggml_build_backward_expand( struct ggml_context * ctx, // context for gradient computation struct ggml_cgraph * cgraph, diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 61e0c9adb136..b14248897227 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2003,10 +2003,7 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm } break; case GGML_OP_FLASH_ATTN_BACK: { - int32_t t = ggml_get_op_params_i32(tensor, 0); - GGML_ASSERT(t == 0 || t == 1); - bool masked = t != 0; - ggml_compute_forward_flash_attn_back(params, masked, tensor); + ggml_compute_forward_flash_attn_back(params, tensor); } break; case GGML_OP_SSM_CONV: { @@ -2887,9 +2884,9 @@ struct ggml_cplan ggml_graph_plan( } break; case GGML_OP_OUT_PROD: { - if (ggml_is_quantized(node->src[0]->type) || - node->src[0]->type == GGML_TYPE_F16) { - cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; + if (ggml_is_quantized(node->src[0]->type) || node->src[0]->type == GGML_TYPE_F16) { + // the blocked kernel's tile and gathered values (ops.h) + cur = ggml_type_size(GGML_TYPE_F32) * ggml_out_prod_cpu_wsize(node->src[0]->ne[0]) * n_tasks; } } break; case GGML_OP_SET_ROWS: @@ -2975,19 +2972,11 @@ struct ggml_cplan ggml_graph_plan( } break; case GGML_OP_FLASH_ATTN_BACK: { - const int64_t D = node->src[0]->ne[0]; - const int64_t ne11 = ggml_up(node->src[1]->ne[1], GGML_SOFT_MAX_UNROLL); - const int64_t mxDn = MAX(D, ne11) * 2; // *2 because of S and SM in ggml_compute_forward_flash_attn_back - if (node->src[1]->type == GGML_TYPE_F32) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } else if (node->src[1]->type == GGML_TYPE_F16) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } else if (node->src[1]->type == GGML_TYPE_BF16) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } + // per thread: the blocked kernel's scratch (ops.h) + const int64_t D = node->src[0]->ne[0]; + const int64_t M = node->src[1]->ne[1]; + const int64_t Dv = node->src[2]->ne[0]; + cur = sizeof(float)*ggml_flash_attn_back_cpu_wsize(D, Dv, M)*n_tasks; } break; case GGML_OP_CROSS_ENTROPY_LOSS: diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp index 6f75f1254d27..96601da9a074 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.cpp +++ b/ggml/src/ggml-cpu/ggml-cpu.cpp @@ -475,7 +475,7 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st return src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16; case GGML_OP_OUT_PROD: return (src0->type == GGML_TYPE_F32 || - ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) && + ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src1->ne[2] % src0->ne[2] == 0 && src1->ne[3] % src0->ne[3] == 0)) && src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; case GGML_OP_CONV_2D: return ggml_is_contiguous(op->src[0]); diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index e1b6f0e7bc60..85ecb257e137 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4356,7 +4356,7 @@ static void ggml_compute_forward_out_prod_f32( } } -static void ggml_compute_forward_out_prod_q_f32( +static void ggml_compute_forward_out_prod_converted_f32( const ggml_compute_params * params, ggml_tensor * dst) { @@ -4371,8 +4371,8 @@ static void ggml_compute_forward_out_prod_q_f32( const ggml_type type = src0->type; ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float; - GGML_ASSERT(ne02 == ne12); - GGML_ASSERT(ne03 == ne13); + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); GGML_ASSERT(ne2 == ne12); GGML_ASSERT(ne3 == ne13); @@ -4387,124 +4387,62 @@ static void ggml_compute_forward_out_prod_q_f32( GGML_ASSERT(ne0 == ne00); GGML_ASSERT(ne1 == ne10); - GGML_ASSERT(ne2 == ne02); - GGML_ASSERT(ne3 == ne03); - - // nb01 >= nb00 - src0 is not transposed - // compute by src0 rows - - if (ith == 0) { - ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); - } - ggml_barrier(params->threadpool); - - // parallelize by last three dimensions - - // total rows in dst - const int64_t nr = ne1*ne2*ne3; - - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; - - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); - - // dst[:,:,:,:] = 0 - // for i2,i3: - // for i1: - // for i01: - // for i0: - // dst[i0,i1,i2,i3] += src0[i0,i01,i2,i3] * src1[i1,i01,i2,i3] - - float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith; - - for (int64_t ir = ir0; ir < ir1; ++ir) { - // dst indices - const int64_t i3 = ir/(ne2*ne1); - const int64_t i2 = (ir - i3*ne2*ne1)/ne1; - const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); - - const int64_t i02 = i2; - const int64_t i03 = i3; - - //const int64_t i10 = i1; - const int64_t i12 = i2; - const int64_t i13 = i3; - - for (int64_t i01 = 0; i01 < ne01; ++i01) { - const int64_t i11 = i01; - - float * s0 = (float *) ((char *) src0->data + ( i01*nb01 + i02*nb02 + i03*nb03)); - float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); - float * d = (float *) ((char *) dst->data + ( i1*nb1 + i2*nb2 + i3*nb3)); - - dequantize_row_q(s0, wdata, ne0); - ggml_vec_mad_f32(ne0, d, wdata, *s1); - } - } -} - -static void ggml_compute_forward_out_prod_f16_f32( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_TENSOR_BINARY_OP_LOCALS; - const int ith = params->ith; - const int nth = params->nth; - GGML_ASSERT(src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - GGML_ASSERT(ne02 == ne12); - GGML_ASSERT(ne03 == ne13); - GGML_ASSERT(ne2 == ne12); - GGML_ASSERT(ne3 == ne13); - - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb0 == sizeof(float)); - - GGML_ASSERT(ne0 == ne00); - GGML_ASSERT(ne1 == ne10); - GGML_ASSERT(ne2 == ne02); - GGML_ASSERT(ne3 == ne03); + // nb01 >= nb00 - src0 is not transposed + // compute by src0 rows + // + // dst[i0,i1] = sum_i01 src0[i0,i01] * src1[i1,i01]: a GEMM dst[ne1 x ne0] = S1[ne1 x ne01] * + // W[ne01 x ne0]. Each task owns a block of BN dst rows of one (i2,i3) plane, dequantizes src0 in + // tiles of KT rows ONCE per block (the row loop dequantized every src0 row again for every dst + // row), gathers the matching src1 values, and accumulates on simd_gemm. + const int64_t BN = GGML_OUT_PROD_BN; + const int64_t KT = GGML_OUT_PROD_KT; if (ith == 0) { ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0); } ggml_barrier(params->threadpool); - const int64_t nr = ne1*ne2*ne3; - const int64_t dr = (nr + nth - 1)/nth; - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); - - float * wdata = (float *) params->wdata + (ne0 + CACHE_LINE_SIZE_F32) * ith; + float * wt = (float *) params->wdata + ith*ggml_out_prod_cpu_wsize(ne0); // [KT][ne0] + float * at = wt + KT*ne0; // [BN][KT] - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i3 = ir/(ne2*ne1); - const int64_t i2 = (ir - i3*ne2*ne1)/ne1; - const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); + const int64_t nblk = (ne1 + BN - 1)/BN; + const int64_t ntasks = nblk*ne2*ne3; + const bool dense = nb1 == (size_t) ne0*sizeof(float); // simd_gemm's C rows are contiguous - const int64_t i02 = i2; - const int64_t i03 = i3; + for (int64_t t = ith; t < ntasks; t += nth) { + const int64_t i3 = t/(ne2*nblk); + const int64_t i2 = (t - i3*ne2*nblk)/nblk; + const int64_t b0 = (t - i3*ne2*nblk - i2*nblk)*BN; + const int64_t nb_ = std::min(BN, ne1 - b0); - const int64_t i12 = i2; - const int64_t i13 = i3; + const int64_t i02 = i2 / (ne12 / ne02); + const int64_t i03 = i3 / (ne13 / ne03); - float * d = (float *) ((char *) dst->data + (i1*nb1 + i2*nb2 + i3*nb3)); + float * d0 = (float *) ((char *) dst->data + (b0*nb1 + i2*nb2 + i3*nb3)); - for (int64_t i01 = 0; i01 < ne01; ++i01) { - const int64_t i11 = i01; - ggml_fp16_t * s0 = (ggml_fp16_t *) ((char *) src0->data + (i01*nb01 + i02*nb02 + i03*nb03)); - float * s1 = (float *) ((char *) src1->data + (i1*nb10 + i11*nb11 + i12*nb12 + i13*nb13)); - ggml_fp16_to_fp32_row(s0, wdata, ne0); - ggml_vec_mad_f32(ne0, d, wdata, *s1); + for (int64_t k0 = 0; k0 < ne01; k0 += KT) { + const int64_t kt = std::min(KT, ne01 - k0); + for (int64_t r = 0; r < kt; ++r) { + dequantize_row_q((const char *) src0->data + ((k0 + r)*nb01 + i02*nb02 + i03*nb03), wt + r*ne0, ne0); + } + for (int64_t b = 0; b < nb_; ++b) { + for (int64_t r = 0; r < kt; ++r) { + at[b*kt + r] = *(const float *) ((const char *) src1->data + ((b0 + b)*nb10 + (k0 + r)*nb11 + i2*nb12 + i3*nb13)); + } + } + if (dense) { + simd_gemm(d0, at, wt, (int) nb_, (int) kt, (int) ne0); + } else { + for (int64_t b = 0; b < nb_; ++b) { + float * d = (float *) ((char *) d0 + b*nb1); + for (int64_t r = 0; r < kt; ++r) { + ggml_vec_mad_f32(ne0, d, wt + r*ne0, at[b*kt + r]); + } + } + } } } } @@ -4542,11 +4480,12 @@ void ggml_compute_forward_out_prod( case GGML_TYPE_IQ3_S: case GGML_TYPE_IQ2_S: { - ggml_compute_forward_out_prod_q_f32(params, dst); + ggml_compute_forward_out_prod_converted_f32(params, dst); } break; case GGML_TYPE_F16: { - ggml_compute_forward_out_prod_f16_f32(params, dst); + // f16 converts through its to_float trait exactly as a quantized row dequantizes + ggml_compute_forward_out_prod_converted_f32(params, dst); } break; case GGML_TYPE_F32: { @@ -9224,316 +9163,227 @@ void ggml_compute_forward_flash_attn_ext( } // ggml_compute_forward_flash_attn_back +// +// The gradients of FLASH_ATTN_EXT (ggml_flash_attn_back), the reference every backend's kernel is +// checked against. Pass 1, over query rows (i, h): the row's log-sum-exp, D_i = sum_j p_ij dp_ij +// (from the same recomputed p and dp the gradients use, never from the forward's output) and +// dq_i = scale sum_j ds_ij k_j, with ds_ij = p_ij (dp_ij - D_i). Pass 2, over KV positions +// (j, kv head): dk_j and dv_j summed over the query rows and the KV head's group of query heads. +// Nothing window x rows is kept: a row's scores live in one per-thread buffer. + +static inline float ggml_fa_back_mask(const ggml_tensor * mask, int64_t j, int64_t i, int64_t h) { + if (mask == nullptr) { + return 0.0f; + } + const char * m = (const char *) mask->data + i*mask->nb[1] + (h % mask->ne[2])*mask->nb[2]; + return GGML_CPU_FP16_TO_FP32(((const ggml_fp16_t *) m)[j]); +} static void ggml_compute_forward_flash_attn_back_f32( const ggml_compute_params * params, - const bool masked, - ggml_tensor * dst) { + ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * v = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * d = dst->src[4]; - const ggml_tensor * q = dst->src[0]; - const ggml_tensor * k = dst->src[1]; - const ggml_tensor * v = dst->src[2]; - const ggml_tensor * d = dst->src[3]; + const float scale = ggml_get_op_params_f32(dst, 0); + const int64_t kv0 = ggml_get_op_params_i32(dst, 1); // dK/dV cover the key rows [kv0, M) - GGML_TENSOR_LOCALS(int64_t, neq, q, ne) - GGML_TENSOR_LOCALS(size_t, nbq, q, nb) - GGML_TENSOR_LOCALS(int64_t, nek, k, ne) - GGML_TENSOR_LOCALS(size_t, nbk, k, nb) - GGML_TENSOR_LOCALS(int64_t, nev, v, ne) - GGML_TENSOR_LOCALS(size_t, nbv, v, nb) - GGML_TENSOR_LOCALS(int64_t, ned, d, ne) - GGML_TENSOR_LOCALS(size_t, nbd, d, nb) - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + const int64_t D = q->ne[0]; + const int64_t N = q->ne[1]; + const int64_t H = q->ne[2]; + const int64_t M = k->ne[1]; + const int64_t L = M - kv0; + const int64_t Hkv = k->ne[2]; + const int64_t Dv = v->ne[0]; + const int64_t G = H / Hkv; // query head h reads KV head h / G + + float * dq = (float *) ((char *) dst->data + ggml_flash_attn_back_offset(dst, 0)); + float * dk = (float *) ((char *) dst->data + ggml_flash_attn_back_offset(dst, 1)); + float * dv = (float *) ((char *) dst->data + ggml_flash_attn_back_offset(dst, 2)); + float * lse = (float *) ((char *) dst->data + ggml_flash_attn_back_offset(dst, 3)); + float * dvec = lse + N*H; const int ith = params->ith; const int nth = params->nth; - const int64_t D = neq0; - const int64_t N = neq1; - const int64_t P = nek1 - N; - const int64_t M = P + N; - - const int Mup = ggml_up(M, GGML_SOFT_MAX_UNROLL); - const int mxDM = MAX(D, Mup); - - // GGML_ASSERT(ne0 == D); - // GGML_ASSERT(ne1 == N); - GGML_ASSERT(P >= 0); - - GGML_ASSERT(nbq0 == sizeof(float)); - GGML_ASSERT(nbk0 == sizeof(float)); - GGML_ASSERT(nbv0 == sizeof(float)); - - GGML_ASSERT(neq0 == D); - GGML_ASSERT(nek0 == D); - GGML_ASSERT(nev1 == D); - GGML_ASSERT(ned0 == D); - - GGML_ASSERT(neq1 == N); - GGML_ASSERT(nek1 == N + P); - GGML_ASSERT(nev1 == D); - GGML_ASSERT(ned1 == N); - - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); + // Rows are worked in blocks so each K and V row is converted to F32 once per block, not once per + // query row, and the dot products run on ggml's SIMD vec_dot (pass 1: BR query rows of one head; + // pass 2: BC KV positions of one KV head). Per thread: BR rows of scores, the block's query and + // output-gradient rows, one K and one V row; pass 2 reuses the same space for its K and V block. + const int64_t BR = GGML_FA_BACK_BR; + const int64_t BC = GGML_FA_BACK_BC; + float * wbase = (float *) params->wdata + ith*ggml_flash_attn_back_cpu_wsize(D, Dv, M); + + const ggml_to_float_t k_to_float = ggml_get_type_traits(k->type)->to_float; + const ggml_to_float_t v_to_float = ggml_get_type_traits(v->type)->to_float; + auto k_row = [&](int64_t j, int64_t hk, float * out) { + const char * r = (const char *) k->data + j*k->nb[1] + hk*k->nb[2]; + if (k->type == GGML_TYPE_F32) { memcpy(out, r, D*sizeof(float)); } else { k_to_float(r, out, D); } + }; + auto v_row = [&](int64_t j, int64_t hk, float * out) { + const char * r = (const char *) v->data + j*v->nb[1] + hk*v->nb[2]; + if (v->type == GGML_TYPE_F32) { memcpy(out, r, Dv*sizeof(float)); } else { v_to_float(r, out, Dv); } + }; + auto q_row = [&](int64_t i, int64_t h, float * out) { + for (int64_t c = 0; c < D; ++c) { + out[c] = *(const float *) ((const char *) q->data + c*q->nb[0] + i*q->nb[1] + h*q->nb[2]); + } + }; + auto d_row = [&](int64_t i, int64_t h, float * out) { + for (int64_t c = 0; c < Dv; ++c) { + out[c] = *(const float *) ((const char *) d->data + c*d->nb[0] + h*d->nb[1] + i*d->nb[2]); + } + }; + auto dot = [](const float * a, const float * b, int64_t n) { + float r; + ggml_vec_dot_f32((int) n, &r, 0, a, 0, b, 0, 1); + return r; + }; + // the dO scale a reduced-precision kernel uses (recorded so every backend's result agrees; the + // CPU works in F32 and does not apply it) if (ith == 0) { - memset(dst->data, 0, nb0*ne0*ne1*ne2*ne3); - } - ggml_barrier(params->threadpool); - - const int64_t elem_q = ggml_nelements(q); - const int64_t elem_k = ggml_nelements(k); - - ggml_type result_type = dst->type; - GGML_ASSERT(ggml_blck_size(result_type) == 1); - const size_t tsize = ggml_type_size(result_type); - - const size_t offs_q = 0; - const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); - const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); - - void * grad_q = (char *) dst->data; - void * grad_k = (char *) dst->data + offs_k; - void * grad_v = (char *) dst->data + offs_v; - - const size_t nbgq1 = nb0*neq0; - const size_t nbgq2 = nb0*neq0*neq1; - const size_t nbgq3 = nb0*neq0*neq1*neq2; - - const size_t nbgk1 = nb0*nek0; - const size_t nbgk2 = nb0*nek0*nek1; - const size_t nbgk3 = nb0*nek0*nek1*neq2; - - const size_t nbgv1 = nb0*nev0; - const size_t nbgv2 = nb0*nev0*nev1; - const size_t nbgv3 = nb0*nev0*nev1*neq2; - - // parallelize by k rows using ggml_vec_dot_f32 - - // total rows in k - const int nr = nek2*nek3; - - // rows per thread - const int dr = (nr + nth - 1)/nth; - - // row range for this thread - const int ir0 = dr*ith; - const int ir1 = MIN(ir0 + dr, nr); - - const float scale = 1.0f/sqrtf(D); - - //printf("P=%d N=%d D=%d ir0=%d ir1=%d scale = %f\n", P, N, D, ir0, ir1, scale); - - // how often k2 (and v2) is repeated in q2 - int nrep = neq2/nek2; - - for (int ir = ir0; ir < ir1; ++ir) { - // q indices - const int ik3 = ir/(nek2); - const int ik2 = ir - ik3*nek2; - - const int iq3 = ik3; - const int id3 = ik3; - const int iv3 = ik3; - const int iv2 = ik2; - - for (int irep = 0; irep < nrep; ++irep) { - const int iq2 = ik2 + irep*nek2; - const int id2 = iq2; - - // (ik2 + irep*nek2) % nek2 == ik2 - for (int iq1 = 0; iq1 < neq1; ++iq1) { - const int id1 = iq1; - - // not sure about CACHE_LINE_SIZE_F32.. - // - maybe it must not be multiplied by 2 and excluded from .. in SM 1*(..) offset? - float * S = (float *) params->wdata + ith*2*(mxDM + CACHE_LINE_SIZE_F32) + 0*(mxDM+CACHE_LINE_SIZE_F32); - float * SM = (float *) params->wdata + ith*2*(mxDM + CACHE_LINE_SIZE_F32) + 1*(mxDM+CACHE_LINE_SIZE_F32); - - for (int i = M; i < Mup; ++i) { - S[i] = -INFINITY; + float mx = 0.0f; + for (int64_t i = 0; i < N; ++i) { + for (int64_t h = 0; h < H; ++h) { + for (int64_t c = 0; c < Dv; ++c) { + mx = std::max(mx, fabsf(*(const float *) ((const char *) d->data + c*d->nb[0] + h*d->nb[1] + i*d->nb[2]))); } + } + } + dvec[N*H] = ggml_flash_attn_back_scale(mx); + } - const int64_t masked_begin = masked ? (P + iq1 + 1) : M; - for (int64_t ic = 0; ic < masked_begin; ++ic) { - // k indices - const int ik1 = ic; - - // S indices - const int i1 = ik1; - - ggml_vec_dot_f32(neq0, - S + i1, 0, - (float *) ((char *) k->data + (ik1*nbk1 + ik2*nbk2 + ik3*nbk3)), 0, - (float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)), 0, 1); + // pass 1: blocks of BR query rows of one head + { + float * s = wbase; // [BR][M] scores, then p + float * qb = s + BR*M; // [BR][D] + float * db = qb + BR*D; // [BR][Dv] + float * kr = db + BR*Dv; // [D] + float * vr = kr + D; // [Dv] + float l[GGML_FA_BACK_BR]; + ggml_float dd[GGML_FA_BACK_BR]; + + const int64_t nblk = (N + BR - 1)/BR; + for (int64_t t = ith; t < nblk*H; t += nth) { + const int64_t i0 = (t % nblk)*BR; + const int64_t h = t / nblk; + const int64_t hk = h / G; + const int64_t nb = std::min(BR, N - i0); + for (int64_t b = 0; b < nb; ++b) { + q_row(i0 + b, h, qb + b*D); + d_row(i0 + b, h, db + b*Dv); + } + for (int64_t j = 0; j < M; ++j) { + k_row(j, hk, kr); + for (int64_t b = 0; b < nb; ++b) { + s[b*M + j] = dot(qb + b*D, kr, D)*scale + ggml_fa_back_mask(mask, j, i0 + b, h); + } + } + for (int64_t b = 0; b < nb; ++b) { + float * sb = s + b*M; + float smax = -INFINITY; + for (int64_t j = 0; j < M; ++j) { + smax = std::max(smax, sb[j]); + } + l[b] = -INFINITY; + if (smax != -INFINITY) { + ggml_float sum = 0.0; + for (int64_t j = 0; j < M; ++j) { + sum += expf(sb[j] - smax); + } + l[b] = smax + (float) log(sum); + } + for (int64_t j = 0; j < M; ++j) { + sb[j] = l[b] == -INFINITY ? 0.0f : expf(sb[j] - l[b]); + } + dd[b] = 0.0; + } + // D from the recomputed p and dp + for (int64_t j = 0; j < M; ++j) { + bool any = false; + for (int64_t b = 0; b < nb; ++b) { any = any || s[b*M + j] != 0.0f; } + if (!any) { continue; } + v_row(j, hk, vr); + for (int64_t b = 0; b < nb; ++b) { + const float p = s[b*M + j]; + if (p != 0.0f) { + dd[b] += (ggml_float) p*dot(db + b*Dv, vr, Dv); + } } - - // scale - ggml_vec_scale_f32(masked_begin, S, scale); - - for (int64_t i = masked_begin; i < M; i++) { - S[i] = -INFINITY; + } + for (int64_t b = 0; b < nb; ++b) { + lse [h*N + i0 + b] = l[b]; + dvec[h*N + i0 + b] = (float) dd[b]; + memset(dq + (h*N + i0 + b)*D, 0, D*sizeof(float)); + } + for (int64_t j = 0; j < M; ++j) { + bool any = false; + for (int64_t b = 0; b < nb; ++b) { any = any || s[b*M + j] != 0.0f; } + if (!any) { continue; } + k_row(j, hk, kr); + v_row(j, hk, vr); + for (int64_t b = 0; b < nb; ++b) { + const float p = s[b*M + j]; + if (p == 0.0f) { + continue; + } + const float ds = p*(dot(db + b*Dv, vr, Dv) - (float) dd[b])*scale; + ggml_vec_mad_f32((int) D, dq + (h*N + i0 + b)*D, kr, ds); } + } + } + } - // softmax - // exclude known -INF S[..] values from max and loop - // dont forget to set their SM values to zero - { - float max = -INFINITY; - ggml_vec_max_f32(masked_begin, &max, S); + ggml_barrier(params->threadpool); - ggml_float sum = 0.0; - { -#ifdef GGML_SOFT_MAX_ACCELERATE - max = -max; - vDSP_vsadd(SM, 1, &max, SM, 1, Mup); - vvexpf(SM, SM, &Mup); - ggml_vec_sum_f32(Mup, &sum, SM); -#else - sum = ggml_vec_soft_max_f32(Mup, SM, S, max); -#endif + // pass 2: blocks of BC KV positions of one KV head, each summed over its group's query heads and + // every query row + { + float * kb = wbase; // [BC][D] + float * vb = kb + BC*D; // [BC][Dv] + float * qr = vb + BC*Dv; // [D] + float * dr = qr + D; // [Dv] + + const int64_t nblk = (L + BC - 1)/BC; + for (int64_t t = ith; t < nblk*Hkv; t += nth) { + const int64_t j0 = kv0 + (t % nblk)*BC; + const int64_t hk = t / nblk; + const int64_t nb = std::min(BC, M - j0); + for (int64_t b = 0; b < nb; ++b) { + k_row(j0 + b, hk, kb + b*D); + v_row(j0 + b, hk, vb + b*Dv); + memset(dk + (hk*L + j0 - kv0 + b)*D, 0, D*sizeof(float)); + memset(dv + (hk*L + j0 - kv0 + b)*Dv, 0, Dv*sizeof(float)); + } + for (int64_t h = hk*G; h < (hk + 1)*G; ++h) { + for (int64_t i = 0; i < N; ++i) { + const float l = lse[h*N + i]; + if (l == -INFINITY) { + continue; + } + bool loaded = false; + for (int64_t b = 0; b < nb; ++b) { + const float m = ggml_fa_back_mask(mask, j0 + b, i, h); + if (m == -INFINITY) { + continue; + } + if (!loaded) { + q_row(i, h, qr); + d_row(i, h, dr); + loaded = true; + } + const float p = expf(dot(qr, kb + b*D, D)*scale + m - l); + if (p == 0.0f) { + continue; + } + const float ds = p*(dot(dr, vb + b*Dv, Dv) - dvec[h*N + i])*scale; + ggml_vec_mad_f32((int) D, dk + (hk*L + j0 - kv0 + b)*D, qr, ds); + ggml_vec_mad_f32((int) Dv, dv + (hk*L + j0 - kv0 + b)*Dv, dr, p); } - - assert(sum > 0.0); - - sum = 1.0/sum; - ggml_vec_scale_f32(masked_begin, SM, sum); - - } - - // step-by-step explanation - { - // forward-process shape grads from backward process - // parallel_for ik2,ik3: - // for irep: - // iq2 = ik2 + irep*nek2 - // k[:D,:M,:,:] [D,M,:,:] grad[k][:D,:M,ik2,ik3] += grad[kcur] - // q[:D,:N,:,:] [D,N,:,:] grad[q][:D,iq1,iq2,iq3] += grad[qcur] - // v[:M,:D,:,:] [M,D,:,:] grad[v][:M,:D,iv2,iv3] += grad[vcur] - // for iq1: - // kcur = k[:D,:M,ik2,ik3] [D,M,1,1] grad[kcur] = grad[S1].T @ qcur - // qcur = q[:D,iq1,iq2,iq3] [D,1,1,1] grad[qcur] = grad[S1] @ kcur - // vcur = v[:M,:D,iv2,iv3] [M,D,1,1] grad[vcur] = grad[S5].T @ S4 - // S0 = -Inf [D,1,1,1] - // ~S1[i] = dot(kcur[:D,i], qcur) - // S1 = qcur @ kcur.T [M,1,1,1] grad[S1] = grad[S2] * scale - // S2 = S1 * scale [M,1,1,1] grad[S2] = diag_mask_zero(grad[S3], P) - // S3 = diag_mask_inf(S2, P) [M,1,1,1] grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // S4 = softmax(S3) [M,1,1,1] grad[S4] = grad[S5] @ vcur - // ~S5[i] = dot(vcur[:,i], S4) - // S5 = S4 @ vcur.T [D,1,1,1] grad[S5] = d[:D,id1,id2,id3] - // ~dst[i,iq1,iq2,iq3] = S5[i] ^ - // dst[:D,iq1,iq2,iq3] = S5 | grad[dst[:D,iq1,iq2,iq3]] = d[:D,id1,id2,id3] - // dst backward-/ grad[dst] = d - // - // output gradients with their dependencies: - // - // grad[kcur] = grad[S1].T @ qcur - // grad[S1] = diag_mask_zero(grad[S3], P) * scale - // grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // grad[S4] = grad[S5] @ vcur - // grad[S4] = d[:D,id1,id2,id3] @ vcur - // grad[qcur] = grad[S1] @ kcur - // grad[vcur] = grad[S5].T @ S4 - // grad[vcur] = d[:D,id1,id2,id3].T @ S4 - // - // in post-order: - // - // S1 = qcur @ kcur.T - // S2 = S1 * scale - // S3 = diag_mask_inf(S2, P) - // S4 = softmax(S3) - // grad[S4] = d[:D,id1,id2,id3] @ vcur - // grad[S3] = S4 * (grad[S4] - dot(S4, grad[S4])) - // grad[S1] = diag_mask_zero(grad[S3], P) * scale - // grad[qcur] = grad[S1] @ kcur - // grad[kcur] = grad[S1].T @ qcur - // grad[vcur] = d[:D,id1,id2,id3].T @ S4 - // - // using less variables (SM=S4): - // - // S = diag_mask_inf(qcur @ kcur.T * scale, P) - // SM = softmax(S) - // S = d[:D,iq1,iq2,iq3] @ vcur - // dot_SM_gradSM = dot(SM, S) - // S = SM * (S - dot(SM, S)) - // S = diag_mask_zero(S, P) * scale - // - // grad[q][:D,iq1,iq2,iq3] += S @ kcur - // grad[k][:D,:M,ik2,ik3] += S.T @ qcur - // grad[v][:M,:D,iv2,iv3] += d[:D,id1,id2,id3].T @ SM - } - - // S = gradSM = d[:D,id1,id2,id3] @ vcur[:,:,iv2,iv3] - // S = d[:D,id1,id2,id3] @ vcur[:,:,iv2,iv3] - // for ic: - // S[:M] += vcur[:M,ic,iv2,iv3] * d[ic,id1,id2,id3] - // exclude known future zero S[..] values from operation - ggml_vec_set_f32(masked_begin, S, 0); - for (int64_t ic = 0; ic < D; ++ic) { - ggml_vec_mad_f32(masked_begin, - S, - (float *) ((char *) v->data + ( ic*nbv1 + iv2*nbv2 + iv3*nbv3)), - *(float *) ((char *) d->data + (ic*nbd0 + id1*nbd1 + id2*nbd2 + id3*nbd3))); - } - - // S = SM * (S - dot(SM, S)) - float dot_SM_gradSM = 0; - ggml_vec_dot_f32 (masked_begin, &dot_SM_gradSM, 0, SM, 0, S, 0, 1); - ggml_vec_acc1_f32(M, S, -dot_SM_gradSM); - ggml_vec_mul_f32 (masked_begin, S, S, SM); - - // S = diag_mask_zero(S, P) * scale - // already done by above ggml_vec_set_f32 - - // exclude known zero S[..] values from operation - ggml_vec_scale_f32(masked_begin, S, scale); - - // S shape [M,1] - // SM shape [M,1] - // kcur shape [D,M] - // qcur shape [D,1] - // vcur shape [M,D] - - // grad[q][:D,iq1,iq2,iq3] += S @ kcur - // grad[q][:D,iq1,iq2,iq3] += shape[M,1] @ shape[D,M] - // for ic: - // grad[q][:D,iq1,iq2,iq3] += S[ic] * kcur[:D,ic,ik2,ik3] - // exclude known zero S[..] values from loop - for (int64_t ic = 0; ic < masked_begin; ++ic) { - ggml_vec_mad_f32(D, - (float *) ((char *) grad_q + (iq1*nbgq1 + iq2*nbgq2 + iq3*nbgq3)), - (float *) ((char *) k->data + (ic*nbk1 + ik2*nbk2 + ik3*nbk3)), - S[ic]); - } - - // grad[k][:D,:M,iq2,iq3] += S.T @ qcur - // for ic: - // grad[k][:D,ic,iq2,iq3] += S.T[0,ic] * qcur[:D,0] - // grad[k][:D,ic,iq2,iq3] += S[ic] * qcur[:D,0] - // exclude known zero S[..] values from loop - for (int64_t ic = 0; ic < masked_begin; ++ic) { - ggml_vec_mad_f32(D, - (float *) ((char *) grad_k + (ic*nbgk1 + ik2*nbgk2 + ik3*nbgk3)), - (float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)), - S[ic]); - } - - // grad[v][:M,:D,iv2,iv3] += d[:D,id1,id2,id3].T @ SM - // for ic: - // grad[v][:M,ic,iv2,iv3] += d[:D,id1,id2,id3].T[0,ic] * SM[:M] - // grad[v][:M,ic,iv2,iv3] += d[ic,id1,id2,id3] * SM[:M] - // exclude known zero SM[..] values from mad - for (int64_t ic = 0; ic < D; ++ic) { - ggml_vec_mad_f32(masked_begin, - (float *) ((char *) grad_v + ( ic*nbgv1 + iv2*nbgv2 + iv3*nbgv3)), - SM, - *(float *) ((char *) d->data + (ic*nbd0 + id1*nbd1 + id2*nbd2 + id3*nbd3))); } } } @@ -9542,21 +9392,9 @@ static void ggml_compute_forward_flash_attn_back_f32( void ggml_compute_forward_flash_attn_back( const ggml_compute_params * params, - const bool masked, ggml_tensor * dst) { - - const ggml_tensor * q = dst->src[0]; - - switch (q->type) { - case GGML_TYPE_F32: - { - ggml_compute_forward_flash_attn_back_f32(params, masked, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } + GGML_ASSERT(dst->src[0]->type == GGML_TYPE_F32); + ggml_compute_forward_flash_attn_back_f32(params, dst); } // ggml_compute_forward_ssm_conv @@ -11746,11 +11584,12 @@ static void ggml_compute_forward_cross_entropy_loss_f32( const ggml_tensor * src1 = dst->src[1]; GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_I32); GGML_ASSERT(src0->nb[0] == ggml_type_size(src0->type)); GGML_ASSERT(src1->nb[0] == ggml_type_size(src1->type)); - GGML_ASSERT(ggml_are_same_shape(src0, src1)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(src0, src1)); GGML_ASSERT(ggml_is_scalar(dst)); + const bool index = src1->type == GGML_TYPE_I32; GGML_ASSERT(dst->type == GGML_TYPE_F32); // TODO: handle transposed/permuted matrices @@ -11775,18 +11614,22 @@ static void ggml_compute_forward_cross_entropy_loss_f32( for (int64_t i1 = ir0; i1 < ir1; ++i1) { const float * s0 = (const float *)((const char *) src0->data + i1*src0->nb[1]); - const float * s1 = (const float *)((const char *) src1->data + i1*src1->nb[1]); - -#ifndef NDEBUG - for (int64_t i = 0; i < nc; ++i) { - //printf("p[%d] = %f\n", i, p[i]); - assert(!isnan(s0[i])); - assert(!isnan(s1[i])); - } -#endif // NDEBUG float max = -INFINITY; ggml_vec_max_f32(nc, &max, s0); + if (index) { + // one target: its log softmax, -(max + log sum exp(x - max)) + x[t], times the row's weight + const int32_t t = ((const int32_t *) src1->data)[i1]; + const float w = dst->src[2] ? ((const float *) dst->src[2]->data)[i1] : 1.0f; + if (t >= 0 && w != 0.0f) { + GGML_ASSERT(t < nc); + const ggml_float lse = ggml_vec_log_soft_max_f32(nc, st, s0, max); + sum_thread += w*(st[t] - lse); + } + continue; + } + const float * s1 = (const float *)((const char *) src1->data + i1*src1->nb[1]); + const ggml_float sum_softmax = ggml_vec_log_soft_max_f32(nc, st, s0, max); assert(sum_softmax >= 0.0); @@ -11846,7 +11689,8 @@ static void ggml_compute_forward_cross_entropy_loss_back_f32( GGML_ASSERT(ggml_is_contiguous(src0f)); GGML_ASSERT(ggml_is_contiguous(src1f)); GGML_ASSERT(ggml_is_contiguous(grad)); - GGML_ASSERT(ggml_are_same_shape(src0f, src1f) && ggml_are_same_shape(src0f, dst)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(src0f, src1f) && ggml_are_same_shape(src0f, dst)); + const bool index = src1f->type == GGML_TYPE_I32; const int64_t ith = params->ith; const int64_t nth = params->nth; @@ -11867,15 +11711,25 @@ static void ggml_compute_forward_cross_entropy_loss_back_f32( for (int64_t i1 = ir0; i1 < ir1; i1++) { float * ds0 = (float *)((char *) dst->data + i1*dst->nb[1]); const float * s0 = (const float *)((const char *) src0f->data + i1*src0f->nb[1]); - const float * s1 = (const float *)((const char *) src1f->data + i1*src1f->nb[1]); -#ifndef NDEBUG - for (int64_t i = 0; i < nc; ++i) { - //printf("p[%d] = %f\n", i, p[i]); - assert(!isnan(s0[i])); - assert(!isnan(s1[i])); + if (index) { + // one target t: (softmax(x) - onehot(t)) * d / nr; a row without one is zero + const int32_t t = ((const int32_t *) src1f->data)[i1]; + const float w = dst->src[3] ? ((const float *) dst->src[3]->data)[i1] : 1.0f; + if (t < 0 || w == 0.0f) { + memset(ds0, 0, nc*sizeof(float)); + continue; + } + GGML_ASSERT(t < nc); + float max = -INFINITY; + ggml_vec_max_f32(nc, &max, s0); + const ggml_float sum = ggml_vec_soft_max_f32(nc, ds0, s0, max); + ggml_vec_scale_f32(nc, ds0, 1.0/sum); + ds0[t] -= 1.0f; + ggml_vec_scale_f32(nc, ds0, d_by_nr*w); + continue; } -#endif // NDEBUG + const float * s1 = (const float *)((const char *) src1f->data + i1*src1f->nb[1]); // soft_max float max = -INFINITY; diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 1c8b66bfaaa0..9c1e15dc2765 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -88,10 +88,23 @@ void ggml_compute_forward_leaky_relu(const struct ggml_compute_params * params, void ggml_compute_forward_tri(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_fill(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_flash_attn_ext(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_flash_attn_back( - const struct ggml_compute_params * params, - const bool masked, - struct ggml_tensor * dst); +void ggml_compute_forward_flash_attn_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); +// OUT_PROD with a quantized src0: BN dst rows per task share one dequantization of each KT-row tile +// of src0. Its per-thread scratch, in floats, for a row of ne0; the work planner reserves this. +#define GGML_OUT_PROD_BN 32 +#define GGML_OUT_PROD_KT 32 +static inline int64_t ggml_out_prod_cpu_wsize(int64_t ne0) { + return GGML_OUT_PROD_KT*ne0 + GGML_OUT_PROD_BN*GGML_OUT_PROD_KT; +} +// FLASH_ATTN_BACK's blocks: BR query rows (pass 1) and BC KV positions (pass 2) share one conversion of +// each K and V row. Its per-thread scratch, in floats; the work-size planner reserves this per task. +#define GGML_FA_BACK_BR 8 +#define GGML_FA_BACK_BC 16 +static inline int64_t ggml_flash_attn_back_cpu_wsize(int64_t D, int64_t Dv, int64_t M) { + const int64_t pass1 = GGML_FA_BACK_BR*(M + D + Dv) + D + Dv; + const int64_t pass2 = GGML_FA_BACK_BC*(D + Dv) + D + Dv; + return pass1 > pass2 ? pass1 : pass2; +} void ggml_compute_forward_ssm_conv(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_ssm_scan(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_win_part(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/cross-entropy-loss.cu b/ggml/src/ggml-cuda/cross-entropy-loss.cu index 9cfbedd93f3b..277f5cfdbb9e 100644 --- a/ggml/src/ggml-cuda/cross-entropy-loss.cu +++ b/ggml/src/ggml-cuda/cross-entropy-loss.cu @@ -5,13 +5,15 @@ #include #include -template +// index: the labels are one I32 target per row (-1: no loss), not a distribution per row +template static __global__ void cross_entropy_loss_f32( - const float * __restrict__ logits, const float * __restrict__ labels, float * __restrict__ dst, const int nclasses, const int k) { + const float * __restrict__ logits, const void * __restrict__ labels_v, float * __restrict__ dst, const int nclasses, const int k) { extern __shared__ float tmp[]; logits += int64_t(blockIdx.x)*nclasses; - labels += int64_t(blockIdx.x)*nclasses; + const float * labels = index ? nullptr : (const float *) labels_v + int64_t(blockIdx.x)*nclasses; + const int32_t target = index ? ((const int32_t *) labels_v)[blockIdx.x] : -1; // Find maximum for softmax: float max_logit = -INFINITY; @@ -36,11 +38,15 @@ static __global__ void cross_entropy_loss_f32( // log(exp(logits - max) / sum) = (logits - max) - log(sum) float loss = 0.0f; - for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { - const float logit_i = use_shared ? tmp[i] : logits[i]; - loss += (logit_i - max_logit - sum) * labels[i]; + if (index) { + loss = target >= 0 ? -(logits[target] - max_logit - sum) / (float)k : 0.0f; + } else { + for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { + const float logit_i = use_shared ? tmp[i] : logits[i]; + loss += (logit_i - max_logit - sum) * labels[i]; + } + loss = -warp_reduce_sum(loss) / (float)k; } - loss = -warp_reduce_sum(loss) / (float)k; if (threadIdx.x != 0) { return; @@ -49,15 +55,22 @@ static __global__ void cross_entropy_loss_f32( dst[blockIdx.x] = loss; } -template +template static __global__ void cross_entropy_loss_back_f32( - const float * __restrict__ grad, const float * __restrict__ logits, const float * __restrict__ labels, + const float * __restrict__ grad, const float * __restrict__ logits, const void * __restrict__ labels_v, float * __restrict__ dst, const int nclasses) { extern __shared__ float tmp[]; logits += int64_t(blockIdx.x)*nclasses; - labels += int64_t(blockIdx.x)*nclasses; dst += int64_t(blockIdx.x)*nclasses; + const float * labels = index ? nullptr : (const float *) labels_v + int64_t(blockIdx.x)*nclasses; + const int32_t target = index ? ((const int32_t *) labels_v)[blockIdx.x] : -1; + if (index && target < 0) { + for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { + dst[i] = 0.0f; + } + return; + } float maxval = -INFINITY; for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { @@ -85,17 +98,21 @@ static __global__ void cross_entropy_loss_back_f32( // d/dx of -sum(y log softmax(x)) is softmax(x)*sum(y) - y: with sum(y) a masked (all-zero) // label row gets a zero gradient; without it the row would still push softmax(x) back. - float lsum = 0.0f; - for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { - lsum += labels[i]; + float lsum = 1.0f; + if (!index) { + lsum = 0.0f; + for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { + lsum += labels[i]; + } + lsum = warp_reduce_sum(lsum); } - lsum = warp_reduce_sum(lsum); const float sm_scale = lsum/sum; const float d_by_nrows = *grad/gridDim.x; for (int i = threadIdx.x; i < nclasses; i += WARP_SIZE) { const float val = use_shared ? tmp[i] : dst[i]; - dst[i] = (val*sm_scale - labels[i])*d_by_nrows; + const float y = index ? (i == target ? 1.0f : 0.0f) : labels[i]; + dst[i] = (val*sm_scale - y)*d_by_nrows; } } @@ -104,8 +121,9 @@ void ggml_cuda_cross_entropy_loss(ggml_backend_cuda_context & ctx, ggml_tensor * const ggml_tensor * src1 = dst->src[1]; GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_I32); GGML_ASSERT( dst->type == GGML_TYPE_F32); + const bool index = src1->type == GGML_TYPE_I32; GGML_ASSERT(ggml_is_contiguous(src0)); GGML_ASSERT(ggml_is_contiguous(src1)); @@ -115,7 +133,7 @@ void ggml_cuda_cross_entropy_loss(ggml_backend_cuda_context & ctx, ggml_tensor * const int64_t nrows = ggml_nrows(src0); const float * src0_d = (const float *) src0->data; - const float * src1_d = (const float *) src1->data; + const void * src1_d = src1->data; float * dst_d = (float *) dst->data; ggml_cuda_pool & pool = ctx.pool(); @@ -131,10 +149,17 @@ void ggml_cuda_cross_entropy_loss(ggml_backend_cuda_context & ctx, ggml_tensor * ggml_cuda_pool_alloc dst_tmp(pool, blocks_num.x); if (nbytes_shared <= smpbo) { - CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_f32), smpbo); - cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); + if (index) { + CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_f32), smpbo); + cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); + } else { + CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_f32), smpbo); + cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); + } + } else if (index) { + cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); } else { - cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); + cross_entropy_loss_f32<<>>(src0_d, src1_d, dst_tmp.ptr, ne00, nrows); } CUDA_CHECK(cudaGetLastError()); @@ -148,15 +173,16 @@ void ggml_cuda_cross_entropy_loss_back(ggml_backend_cuda_context & ctx, ggml_ten const ggml_tensor * src1f = dst->src[2]; GGML_ASSERT(src0f->type == GGML_TYPE_F32); - GGML_ASSERT(src1f->type == GGML_TYPE_F32); + GGML_ASSERT(src1f->type == GGML_TYPE_F32 || src1f->type == GGML_TYPE_I32); GGML_ASSERT( grad->type == GGML_TYPE_F32); + const bool index = src1f->type == GGML_TYPE_I32; GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(ggml_is_scalar(grad)); GGML_ASSERT(ggml_is_contiguous(src0f)); GGML_ASSERT(ggml_is_contiguous(src1f)); GGML_ASSERT(ggml_is_contiguous(dst)); - GGML_ASSERT(ggml_are_same_shape(src0f, src1f)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(src0f, src1f)); GGML_ASSERT(ggml_are_same_shape(src0f, dst)); const int64_t ne00 = src0f->ne[0]; @@ -164,7 +190,7 @@ void ggml_cuda_cross_entropy_loss_back(ggml_backend_cuda_context & ctx, ggml_ten const float * grad_d = (const float *) grad->data; const float * src0f_d = (const float *) src0f->data; - const float * src1f_d = (const float *) src1f->data; + const void * src1f_d = src1f->data; float * dst_d = (float *) dst->data; cudaStream_t stream = ctx.stream(); @@ -177,9 +203,16 @@ void ggml_cuda_cross_entropy_loss_back(ggml_backend_cuda_context & ctx, ggml_ten const size_t smpbo = ggml_cuda_info().devices[id].smpbo; if (nbytes_shared <= smpbo) { - CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_back_f32), smpbo); - cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); + if (index) { + CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_back_f32), smpbo); + cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); + } else { + CUDA_SET_SHARED_MEMORY_LIMIT((cross_entropy_loss_back_f32), smpbo); + cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); + } + } else if (index) { + cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); } else { - cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); + cross_entropy_loss_back_f32<<>>(grad_d, src0f_d, src1f_d, dst_d, ne00); } } diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 1bff033486cb..10a8afb21349 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5351,7 +5351,9 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_FLASH_ATTN_EXT: return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); case GGML_OP_CROSS_ENTROPY_LOSS: + return op->src[2] == nullptr; // per-row weights: not on CUDA yet, they run on the CPU case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + return op->src[3] == nullptr; case GGML_OP_OPT_STEP_ADAMW: case GGML_OP_OPT_STEP_SGD: case GGML_OP_FILL: diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index 62b76abbcec9..e0577226b484 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -350,6 +350,7 @@ struct ggml_cgraph { // the slice does not have leafs or gradients // if you need the gradients, get them from the original graph struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph, int i0, int i1); +void ggml_graph_update_use_counts(struct ggml_cgraph * cgraph); // ggml-alloc.c: true if the operation can reuse memory from its sources GGML_API bool ggml_op_can_inplace(enum ggml_op op); diff --git a/ggml/src/ggml-metal/CMakeLists.txt b/ggml/src/ggml-metal/CMakeLists.txt index 140c5d809e02..6f34c7b8fa8c 100644 --- a/ggml/src/ggml-metal/CMakeLists.txt +++ b/ggml/src/ggml-metal/CMakeLists.txt @@ -33,6 +33,7 @@ set(METALLIB_KERNEL_SOURCES kernels/fa.metal kernels/mul_mv.metal kernels/mul_mm.metal + kernels/out_prod.metal kernels/quantize.metal kernels/softmax.metal kernels/norm.metal diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp index 22a37453bf39..38f14ab4691a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.cpp +++ b/ggml/src/ggml-metal/ggml-metal-device.cpp @@ -340,7 +340,9 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_lo char base[256]; char name[256]; - snprintf(base, 256, "kernel_cross_entropy_loss%s_f32", op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK ? "_back" : ""); + const bool back = op->op == GGML_OP_CROSS_ENTROPY_LOSS_BACK; + const ggml_tensor * labels = back ? op->src[2] : op->src[1]; + snprintf(base, 256, "kernel_cross_entropy_loss%s_%s", back ? "_back" : "", labels->type == GGML_TYPE_I32 ? "i32" : "f32"); snprintf(name, 256, "%s", base); ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); @@ -351,13 +353,32 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_lo return res; } -// RMS_NORM_BACK, SOFT_MAX_BACK, REPEAT_BACK: one f32 kernel each, named after the op. +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_back(ggml_metal_library_t lib, const ggml_tensor * op, const char * pass) { + assert(op->op == GGML_OP_FLASH_ATTN_BACK); + + char name[256]; + if (strcmp(pass, "scale") == 0) { + snprintf(name, 256, "kernel_flash_attn_back_scale"); + } else { + snprintf(name, 256, "kernel_flash_attn_back_%s_%s", pass, ggml_type_name(op->src[1]->type)); + } + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); + } + + return res; +} + +// RMS_NORM_BACK, SOFT_MAX_BACK, REPEAT_BACK, GET_ROWS_BACK: one f32 kernel each, named after the op. ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back(ggml_metal_library_t lib, const ggml_tensor * op) { const char * kernel = nullptr; switch (op->op) { case GGML_OP_RMS_NORM_BACK: kernel = "kernel_rms_norm_back_f32"; break; case GGML_OP_SOFT_MAX_BACK: kernel = "kernel_soft_max_back_f32"; break; case GGML_OP_REPEAT_BACK: kernel = "kernel_repeat_back_f32"; break; + case GGML_OP_GET_ROWS_BACK: kernel = "kernel_get_rows_back_f32"; break; default: GGML_ABORT("fatal error"); } @@ -909,6 +930,27 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_meta return res; } +ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_out_prod(ggml_metal_library_t lib, const ggml_tensor * op) { + char base[256]; + char name[256]; + + snprintf(base, 256, "kernel_out_prod_%s", ggml_type_name(op->src[0]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + + // kernels/out_prod.metal: a 64 x 32 tile, float staging of both 32-deep operands + res.nr0 = 64; + res.nr1 = 32; + res.nsg = 4; + res.smem = (64 + 32)*32*sizeof(float); + + return res; +} + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_metal_library_t lib, const ggml_tensor * op) { GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 18ae227a5ca9..42e47cf3ecdf 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -121,6 +121,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cross_entropy_loss(ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_back(ggml_metal_library_t lib, const struct ggml_tensor * op, const char * pass); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_row_back (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_blk (ggml_metal_library_t lib, const struct ggml_tensor * op); @@ -140,6 +141,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_del struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm (ggml_metal_library_t lib, const struct ggml_tensor * op); +struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_out_prod (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_fp32 (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv (ggml_metal_library_t lib, const struct ggml_tensor * op); struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0 (ggml_metal_library_t lib, int ne02, int ne20); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index c916ca8a1657..8faa9a66da77 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -109,6 +109,7 @@ int ggml_metal_pipeline_max_theads_per_threadgroup(struct ggml_metal_pipeline_wi X(FA, fa) \ X(MUL_MV, mul_mv) \ X(MUL_MM, mul_mm) \ + X(OUT_PROD, out_prod) \ X(QUANTIZE, quantize) \ X(SOFTMAX, softmax) \ X(NORM, norm) \ @@ -1470,19 +1471,37 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te } case GGML_OP_RMS_NORM_BACK: case GGML_OP_SOFT_MAX_BACK: - // src0 the incoming gradient, src1 the forward input/output; rows contiguous. - // SOFT_MAX_BACK carries no max_bias (the CPU reference asserts the same). + // src0 the incoming gradient, src1 the forward input/output; each row contiguous, the + // rows themselves at any stride. SOFT_MAX_BACK carries no max_bias (the CPU reference + // asserts the same). return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && - ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op) && + ggml_are_same_shape(op->src[0], op->src[1]) && + ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]) && ggml_is_contiguous_rows(op) && (op->op != GGML_OP_SOFT_MAX_BACK || ggml_get_op_params_f32(op, 1) == 0.0f); case GGML_OP_REPEAT_BACK: return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(op); + case GGML_OP_FLASH_ATTN_BACK: + // one simdgroup per row: head sizes a multiple of 32, at most 512 (FA_BACK_MAX_D) + return op->src[0]->type == GGML_TYPE_F32 && op->src[4]->type == GGML_TYPE_F32 && + op->src[1]->type == op->src[2]->type && + (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_Q8_0) && + op->src[0]->ne[0] % 32 == 0 && op->src[2]->ne[0] % 32 == 0 && + op->src[0]->ne[0] <= 512 && op->src[2]->ne[0] <= 512 && + op->src[0]->nb[0] == sizeof(float) && op->src[4]->nb[0] == sizeof(float) && + (op->src[3] == NULL || op->src[3]->type == GGML_TYPE_F16); + case GGML_OP_GET_ROWS_BACK: + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_I32 && + ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op); case GGML_OP_CROSS_ENTROPY_LOSS: - return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && - ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && + (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_I32) && + ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && + (op->src[2] == NULL || (op->src[2]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_I32)); case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_F32 && - ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) && ggml_is_contiguous(op); + return op->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_I32) && + ggml_is_contiguous(op->src[1]) && ggml_is_contiguous(op->src[2]) && ggml_is_contiguous(op) && + (op->src[3] == NULL || (op->src[3]->type == GGML_TYPE_F32 && op->src[2]->type == GGML_TYPE_I32)); case GGML_OP_OUT_PROD: // The backward through a frozen weight: out_prod(W, grad^T). Composed on this // backend from a dequantizing copy, two transposing copies and the F32 simdgroup @@ -1609,8 +1628,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te return false; } - return (ggml_get_op_params_i32(op, 0) == 0) && (ggml_get_op_params_i32(op, 2) == 0) && - (ggml_get_op_params_i32(op, 4) == 0) && (ggml_get_op_params_i32(op, 6) == 0); + return true; // left and right zero padding on every dim case GGML_OP_PAD_REFLECT_1D: case GGML_OP_TIMESTEP_EMBEDDING: return op->src[0]->type == GGML_TYPE_F32; diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index adc883e20f1d..74a4bcbd15a6 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -481,6 +481,22 @@ typedef struct { int16_t r3; } ggml_metal_kargs_mul_mm; +typedef struct { + int32_t ne00; // m: dst rows + int32_t ne01; // k: the contraction, rows of both sources + uint64_t nb01; + uint64_t nb02; + uint64_t nb03; + int32_t ne10; // n: dst columns + uint64_t nb10; + uint64_t nb11; + uint64_t nb12; + uint64_t nb13; + int32_t ne2; + int32_t dps2; // dst per src0 along dim 2 (grouped heads) + int32_t dps3; +} ggml_metal_kargs_out_prod; + typedef struct { int32_t ne00; int32_t ne01; @@ -886,17 +902,57 @@ typedef struct { uint64_t nb01; uint64_t nb11; uint64_t nb1; + int32_t has_weights; // a per-row F32 weight follows the labels (I32 targets only) } ggml_metal_kargs_cross_entropy_loss; +// FLASH_ATTN_BACK (ggml_flash_attn_back): q [D, N, H], k [D, M, Hkv], v [Dv, M, Hkv], d [Dv, H, N], +// mask [M', N', ne32] F16; the result's parts at byte offsets offs_dk / offs_dv / offs_lse +typedef struct { + int32_t D; + int32_t Dv; + int32_t N; + int32_t H; + int32_t M; + int32_t Hkv; + uint64_t nbq1; + uint64_t nbq2; + uint64_t nbk1; + uint64_t nbk2; + uint64_t nbv1; + uint64_t nbv2; + uint64_t nbd1; + uint64_t nbd2; + uint64_t nbm1; + uint64_t nbm2; + int32_t nem2; + int32_t has_mask; + float scale; + uint64_t offs_dk; + uint64_t offs_dv; + uint64_t offs_lse; + int32_t kv0; // dK/dV cover the key rows [kv0, M), row j at j - kv0 +} ggml_metal_kargs_flash_attn_back; + // a per-row backward with one scalar parameter (eps for RMS_NORM_BACK, scale for SOFT_MAX_BACK) +// rows are addressed through each source's own strides, so a strided view (a per-head norm) +// is read in place typedef struct { int32_t ne00; - uint64_t nb01; - uint64_t nb11; - uint64_t nb1; + uint64_t nb01, nb02, nb03; + uint64_t nb11, nb12, nb13; + uint64_t nb1, nb2, nb3; float param; } ggml_metal_kargs_row_back; +// GET_ROWS_BACK: dst = 0, then dst[:, ids[i]] += src0[:, i] for each i (the gradient of a row +// gather; in a training graph, the selection of the output rows) +typedef struct { + int32_t ne00; // row width + int32_t nr; // number of ids + uint64_t nb01; + uint64_t nb1; +} ggml_metal_kargs_get_rows_back; + typedef struct { int64_t ne00; int64_t ne01; @@ -1128,6 +1184,10 @@ typedef struct { uint64_t nb1; uint64_t nb2; uint64_t nb3; + int32_t lp0; // left padding per dim: dst[i] = src[i - lp] + int32_t lp1; + int32_t lp2; + int32_t lp3; } ggml_metal_kargs_pad; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index cb2e266e12c9..a704f4e40e89 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -312,6 +312,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_cross_entropy_loss_back(ctx, idx); } break; + case GGML_OP_FLASH_ATTN_BACK: + { + n_fuse = ggml_metal_op_flash_attn_back(ctx, idx); + } break; case GGML_OP_RMS_NORM_BACK: case GGML_OP_SOFT_MAX_BACK: { @@ -321,6 +325,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { { n_fuse = ggml_metal_op_repeat_back(ctx, idx); } break; + case GGML_OP_GET_ROWS_BACK: + { + n_fuse = ggml_metal_op_get_rows_back(ctx, idx); + } break; case GGML_OP_SILU_BACK: { n_fuse = ggml_metal_op_silu_back(ctx, idx); @@ -2953,16 +2961,8 @@ bool ggml_metal_op_flash_attn_ext_use_vec(const ggml_tensor * op) { // ref: https://github.com/ggml-org/llama.cpp/pull/27390 // dequantize the quantized KV cache to F16 before running the F16 flash attention kernels -static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { - assert(op->op == GGML_OP_FLASH_ATTN_EXT); - - // depending on compute/bandwidth ratio, dequant to f16 kv is not always beneficial - // ref: https://github.com/ggml-org/llama.cpp/pull/27390#issuecomment-5355152767 - // TODO: tune per device - if (op->src[0]->ne[1] < 32) { - return false; - } - +// whether the kernel may dequantize this K to F16 at some batch size (a quantized K) +static bool ggml_metal_op_flash_attn_ext_kv_dequantizable(const ggml_tensor * op) { switch (op->src[1]->type) { case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: @@ -2975,6 +2975,19 @@ static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { } } +static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) { + assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + // depending on compute/bandwidth ratio, dequant to f16 kv is not always beneficial + // ref: https://github.com/ggml-org/llama.cpp/pull/27390#issuecomment-5355152767 + // TODO: tune per device + if (op->src[0]->ne[1] < 32) { + return false; + } + + return ggml_metal_op_flash_attn_ext_kv_dequantizable(op); +} + // in some models (e.g. MLA-based), V is a view of K (the first ne20 elements of each K row); // the dequantized V is then a view of the dequantized K and does not need its own dequant or scratch // - ref: https://github.com/ggml-org/llama.cpp/pull/13435 @@ -3125,10 +3138,13 @@ size_t ggml_metal_op_flash_attn_ext_extra_tmp(const ggml_tensor * op) { size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const ggml_tensor * op) { assert(op->op == GGML_OP_FLASH_ATTN_EXT); - // note: always reserve the temp buffer to avoid graph reallocations - //if (!ggml_metal_op_flash_attn_ext_use_kv_f16(op)) { - // return 0; - //} + // reserved at every batch size for a quantized K (its use depends on the batch, and a graph + // must not reallocate between decode and prefill); a K the kernel never dequantizes needs none: + // the reservation lives as long as the attention's output, and a training graph keeps one per + // layer for the backward (the M5, qwen3.5-4b, window 32768: 128 MiB a layer, never read) + if (!ggml_metal_op_flash_attn_ext_kv_dequantizable(op)) { + return 0; + } GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); @@ -3962,61 +3978,27 @@ static void ggml_metal_encode_add(ggml_metal_op_t ctx, const ggml_tensor * op) { } } -// OUT_PROD: dst[ne00, ne10] = sum over k of src0[:, k] (x) src1[:, k] — the activation gradient -// through a FROZEN weight in a training graph (out_prod(W, grad^T), the mul_mat backward). -// This backend had no OUT_PROD at all, so a Mac could not train a LoRA on the weights it -// serves (Continuum card 55f314b8, the Metal half of the doctor lane). Composed from kernels -// the backend already has, over scratch that the graph allocator reserves after dst: -// -// for each block of rows k0..k0+rows of W: -// S = cpy(W[k0..], F32) (the dequantizing copy; skipped when W is F32) -// X = cpy(S^T) contiguous [rows, ne00] — contraction dim first -// Y = cpy(src1[:, k0..]^T) contiguous [rows, ne10] -// P = mul_mm(X, Y) [ne00, ne10]; the first block writes dst directly -// dst += P for every block after the first -// -// Row blocks bound the transient to OUT_PROD_BLOCK_BYTES of F32 for any weight (the output -// head of a 27B is 3 GB in F32 whole), the same shape as the CUDA out_prod. First cut: -// correctness by test-backend-ops parity against the CPU backend; a fused dequantize-GEMM -// is the second cut if the §0 speed gate asks for it. -static constexpr int64_t OUT_PROD_BLOCK_BYTES = 64ll*1024*1024; - -static int64_t ggml_metal_op_out_prod_block_rows(const ggml_tensor * op) { - const ggml_tensor * src0 = op->src[0]; - const int64_t ne00 = src0->ne[0]; - const int64_t ne01 = src0->ne[1]; - int64_t rows = OUT_PROD_BLOCK_BYTES/(ne00*(int64_t) sizeof(float)); - rows = std::max(64, (rows/64)*64); - return std::min(rows, ne01); -} - +// OUT_PROD: dst[ne00, ne10] = sum over k of src0[:, k] (x) src1[:, k]. In a training graph it +// is the activation gradient through a frozen weight (out_prod(W, grad^T), the mul_mat +// backward) and the gradient of each LoRA factor. kernels/out_prod.metal computes it in one +// dispatch, reading the weight in place in its own block order. Batched and broadcast +// operands follow the CPU backend: dst and src1 share dims 2-3, and src0 is broadcast over +// them (dst per src0 = ne2/ne02, ne3/ne03). bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { const ggml_tensor * src0 = op->src[0]; const ggml_tensor * src1 = op->src[1]; - if (op->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32) { - return false; - } - if (src0->ne[2] != 1 || src0->ne[3] != 1 || src1->ne[2] != 1 || src1->ne[3] != 1) { - return false; - } - // src1 as the graph emits it: contiguous rows, or the TRANSPOSED view of the gradient - // (out_prod(W, ggml_transpose(grad)) is what the mul_mat backward builds; Cormac on - // llama.cpp#5). Either way Y is copied contiguous per block, so both layouts are - // exact; anything else is refused. - if (!ggml_is_contiguous(op)) { + if (op->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || !ggml_is_contiguous(op)) { return false; } - if (src1->nb[0] != sizeof(float) && src1->nb[1] != sizeof(float)) { + if (src1->ne[2] != op->ne[2] || src1->ne[3] != op->ne[3] || + op->ne[2] % src0->ne[2] != 0 || op->ne[3] % src0->ne[3] != 0) { return false; } - if (src0->ne[0] % 16 != 0 || src0->ne[0] % ggml_blck_size(src0->type) != 0) { - return false; // the dequantizing copy works in 4x4 tiles along whole quant blocks of a row - } switch (src0->type) { case GGML_TYPE_F32: - return ggml_is_contiguous_rows(src0); case GGML_TYPE_F16: case GGML_TYPE_BF16: + return src0->nb[0] == ggml_type_size(src0->type); // read element by element, any width case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: @@ -4027,134 +4009,51 @@ bool ggml_metal_op_out_prod_supported(const ggml_tensor * op) { case GGML_TYPE_Q4_K: case GGML_TYPE_Q5_K: case GGML_TYPE_Q6_K: - return ggml_is_contiguous(src0); // kernel_cpy__f32 exists + // dequantized 16 consecutive elements at a time, from whole blocks of a row + return ggml_is_contiguous_rows(src0) && src0->ne[0] % 16 == 0; default: return false; } } -size_t ggml_metal_op_out_prod_extra(const ggml_tensor * op) { - const ggml_tensor * src0 = op->src[0]; - const ggml_tensor * src1 = op->src[1]; - const int64_t ne00 = src0->ne[0]; - const int64_t ne10 = src1->ne[0]; - const int64_t rows = ggml_metal_op_out_prod_block_rows(op); - size_t res = 0; - if (src0->type != GGML_TYPE_F32) { - res += GGML_PAD(rows*ne00*sizeof(float), 256); // S - } - res += GGML_PAD(rows*ne00*sizeof(float), 256); // X - res += GGML_PAD(rows*ne10*sizeof(float), 256); // Y - if (rows < src0->ne[1]) { - res += GGML_PAD(ggml_nbytes(op), 256); // P - } - return res; -} - -// A synthetic 2D view over the same Metal buffer as `like`, at `data`. -static ggml_tensor ggml_metal_out_prod_view(const ggml_tensor * like, ggml_type type, - int64_t ne0, int64_t ne1, size_t nb0, size_t nb1, void * data) { - ggml_tensor t = {}; - t.type = type; - t.ne[0] = ne0; t.ne[1] = ne1; t.ne[2] = 1; t.ne[3] = 1; - t.nb[0] = nb0; t.nb[1] = nb1; t.nb[2] = nb1*ne1; t.nb[3] = t.nb[2]; - t.buffer = like->buffer; - t.data = data; - return t; -} - int ggml_metal_op_out_prod(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); const ggml_tensor * src0 = op->src[0]; const ggml_tensor * src1 = op->src[1]; GGML_ASSERT(ggml_metal_op_out_prod_supported(op)); + GGML_ASSERT(src1->ne[1] == src0->ne[1]); + GGML_ASSERT(op->ne[0] == src0->ne[0] && op->ne[1] == src1->ne[0]); + + ggml_metal_kargs_out_prod args = { + /*.ne00 =*/ (int32_t) src0->ne[0], + /*.ne01 =*/ (int32_t) src0->ne[1], + /*.nb01 =*/ src0->nb[1], + /*.nb02 =*/ src0->nb[2], + /*.nb03 =*/ src0->nb[3], + /*.ne10 =*/ (int32_t) src1->ne[0], + /*.nb10 =*/ src1->nb[0], + /*.nb11 =*/ src1->nb[1], + /*.nb12 =*/ src1->nb[2], + /*.nb13 =*/ src1->nb[3], + /*.ne2 =*/ (int32_t) op->ne[2], + /*.dps2 =*/ (int32_t) (op->ne[2]/src0->ne[2]), + /*.dps3 =*/ (int32_t) (op->ne[3]/src0->ne[3]), + }; - const int64_t ne00 = src0->ne[0]; - const int64_t ne01 = src0->ne[1]; - const int64_t ne10 = src1->ne[0]; - GGML_ASSERT(src1->ne[1] == ne01); - GGML_ASSERT(op->ne[0] == ne00 && op->ne[1] == ne10); - - const int64_t rows_blk = ggml_metal_op_out_prod_block_rows(op); - - // scratch after dst, in dst's buffer (ggml_backend_metal_buffer_type_get_alloc_size) - char * scratch = (char *) op->data + ggml_nbytes(op); - char * s_data = nullptr; - if (src0->type != GGML_TYPE_F32) { - s_data = scratch; - scratch += GGML_PAD(rows_blk*ne00*sizeof(float), 256); - } - char * x_data = scratch; scratch += GGML_PAD(rows_blk*ne00*sizeof(float), 256); - char * y_data = scratch; scratch += GGML_PAD(rows_blk*ne10*sizeof(float), 256); - char * p_data = rows_blk < ne01 ? scratch : nullptr; - - for (int64_t k0 = 0; k0 < ne01; k0 += rows_blk) { - const int64_t rows = std::min(rows_blk, ne01 - k0); - - // S: this block of W, dequantized to F32, [ne00, rows] contiguous - ggml_tensor s_view; - if (src0->type == GGML_TYPE_F32) { - s_view = ggml_metal_out_prod_view(src0, GGML_TYPE_F32, ne00, rows, - src0->nb[0], src0->nb[1], (char *) src0->data + k0*src0->nb[1]); - } else { - ggml_tensor w_view = ggml_metal_out_prod_view(src0, src0->type, ne00, rows, - src0->nb[0], src0->nb[1], (char *) src0->data + k0*src0->nb[1]); - s_view = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, rows, - sizeof(float), ne00*sizeof(float), s_data); - ggml_metal_op_cpy_impl(ctx, &w_view, &s_view); - // the encoder dispatches concurrently: every kernel below reads what the one - // above wrote, so each step ends on a memory barrier (as cumsum does) - ggml_metal_op_concurrency_reset(ctx); - } - - // X = S^T: a transposed view of S copied contiguous, [rows, ne00] - // based on S, never on op: for an F32 src0, S IS src0's memory, which lives in another - // buffer (a LoRA weight's adapter buffer). Based on op, the view named op's buffer with - // src0's address, the lookup found no buffer ('buffer is nil'), the copy below did not - // run, and the product read stale scratch: every OUT_PROD with a LoRA weight as src0 - // was wrong on Metal, differently per graph layout (Fable, 2026-10-06) - ggml_tensor st_view = ggml_metal_out_prod_view(&s_view, GGML_TYPE_F32, rows, ne00, - s_view.nb[1], s_view.nb[0], s_view.data); - ggml_tensor x = ggml_metal_out_prod_view(op, GGML_TYPE_F32, rows, ne00, - sizeof(float), rows*sizeof(float), x_data); - ggml_metal_op_cpy_impl(ctx, &st_view, &x); - ggml_metal_op_concurrency_reset(ctx); - - // Y = src1[:, k0..k0+rows]^T copied contiguous, [rows, ne10] - ggml_tensor bt_view = ggml_metal_out_prod_view(src1, GGML_TYPE_F32, rows, ne10, - src1->nb[1], src1->nb[0], (char *) src1->data + k0*src1->nb[1]); - ggml_tensor y = ggml_metal_out_prod_view(op, GGML_TYPE_F32, rows, ne10, - sizeof(float), rows*sizeof(float), y_data); - ggml_metal_op_cpy_impl(ctx, &bt_view, &y); - ggml_metal_op_concurrency_reset(ctx); - - // P = mul_mat(X, Y): [ne00, ne10], contraction over rows - const bool first = k0 == 0; - ggml_tensor p = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, ne10, - sizeof(float), ne00*sizeof(float), first ? op->data : (void *) p_data); - p.op = GGML_OP_MUL_MAT; - p.src[0] = &x; - p.src[1] = &y; - // FLOAT staging when the simdgroup kernel runs: X and Y carry gradients, which pass - // half's range (65504) where activations do not; half staging turned them into Inf - // and cost every product ~1e-3 relative precision (2026-10-06, Qwen3.5 0.8B backward) - const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(ctx->lib))->has_tensor; - ggml_metal_encode_mul_mm(ctx, &p, !has_tensor); - ggml_metal_op_concurrency_reset(ctx); + auto pipeline = ggml_metal_library_get_pipeline_out_prod(ctx->lib, op); - if (!first) { - // dst = dst + P - ggml_tensor d = ggml_metal_out_prod_view(op, GGML_TYPE_F32, ne00, ne10, - sizeof(float), ne00*sizeof(float), op->data); - ggml_tensor acc = d; - acc.op = GGML_OP_ADD; - acc.src[0] = &d; - acc.src[1] = &p; - ggml_metal_encode_add(ctx, &acc); - ggml_metal_op_concurrency_reset(ctx); - } - } + ggml_metal_encoder_t enc = ctx->enc; + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(src0), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(src1), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); + ggml_metal_encoder_dispatch_threadgroups(enc, + (op->ne[1] + pipeline.nr1 - 1)/pipeline.nr1, + (op->ne[0] + pipeline.nr0 - 1)/pipeline.nr0, + op->ne[2]*op->ne[3], 32, pipeline.nsg, 1); return 1; } @@ -4173,12 +4072,14 @@ static int ggml_metal_op_cross_entropy_loss_rows(ggml_metal_op_t ctx, const ggml const int32_t ne00 = logits->ne[0]; const int32_t nr = ggml_nrows(logits); + const ggml_tensor * weights = back ? op->src[3] : op->src[2]; ggml_metal_kargs_cross_entropy_loss args = { - /*.ne00 =*/ ne00, - /*.nr =*/ nr, - /*.nb01 =*/ logits->nb[1], - /*.nb11 =*/ labels->nb[1], - /*.nb1 =*/ nb1, + /*.ne00 =*/ ne00, + /*.nr =*/ nr, + /*.nb01 =*/ logits->nb[1], + /*.nb11 =*/ labels->nb[1], + /*.nb1 =*/ nb1, + /*.has_weights =*/ weights != nullptr, }; auto pipeline = ggml_metal_library_get_pipeline_cross_entropy_loss(lib, op); @@ -4191,15 +4092,19 @@ static int ggml_metal_op_cross_entropy_loss_rows(ggml_metal_op_t ctx, const ggml ggml_metal_encoder_set_pipeline(enc, pipeline); ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + // the weights ride in the last slot (the labels' buffer stands in when there are none: never read) + const ggml_metal_buffer_id bid_w = ggml_metal_get_buffer_id(weights ? weights : labels); if (back) { ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(logits), 2); ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(labels), 3); ggml_metal_encoder_set_buffer(enc, bid_dst, 4); + ggml_metal_encoder_set_buffer(enc, bid_w, 5); } else { ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(logits), 1); ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(labels), 2); ggml_metal_encoder_set_buffer(enc, bid_dst, 3); + ggml_metal_encoder_set_buffer(enc, bid_w, 4); } ggml_metal_encoder_set_threadgroup_memory_size(enc, 32*sizeof(float), 0); ggml_metal_encoder_dispatch_threadgroups(enc, nr, 1, 1, nth, 1, 1); @@ -4235,6 +4140,85 @@ int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx) { return 1; } +// FLASH_ATTN_BACK: the Q pass (dq, and each row's lse and D into the scratch after dv), then the +// KV pass, which reads them +int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + const ggml_tensor * q = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * v = op->src[2]; + const ggml_tensor * mask = op->src[3]; + const ggml_tensor * d = op->src[4]; + + ggml_metal_kargs_flash_attn_back args = { + /*.D =*/ (int32_t) q->ne[0], + /*.Dv =*/ (int32_t) v->ne[0], + /*.N =*/ (int32_t) q->ne[1], + /*.H =*/ (int32_t) q->ne[2], + /*.M =*/ (int32_t) k->ne[1], + /*.Hkv =*/ (int32_t) k->ne[2], + /*.nbq1 =*/ q->nb[1], + /*.nbq2 =*/ q->nb[2], + /*.nbk1 =*/ k->nb[1], + /*.nbk2 =*/ k->nb[2], + /*.nbv1 =*/ v->nb[1], + /*.nbv2 =*/ v->nb[2], + /*.nbd1 =*/ d->nb[1], + /*.nbd2 =*/ d->nb[2], + /*.nbm1 =*/ mask ? mask->nb[1] : 0, + /*.nbm2 =*/ mask ? mask->nb[2] : 0, + /*.nem2 =*/ mask ? (int32_t) mask->ne[2] : 1, + /*.has_mask =*/ mask ? 1 : 0, + /*.scale =*/ ggml_get_op_params_f32(op, 0), + /*.offs_dk =*/ ggml_flash_attn_back_offset(op, 1), + /*.offs_dv =*/ ggml_flash_attn_back_offset(op, 2), + /*.offs_lse =*/ ggml_flash_attn_back_offset(op, 3), + /*.kv0 =*/ ggml_get_op_params_i32(op, 1), + }; + const int L = args.M - args.kv0; // the key rows a dK/dV is computed for + + auto run = [&](const char * pass, int ngx, int ngy, int nth, size_t shmem) { + auto pipeline = ggml_metal_library_get_pipeline_flash_attn_back(lib, op, pass); + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(q), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(k), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(v), 3); + ggml_metal_encoder_set_buffer (enc, mask ? ggml_metal_get_buffer_id(mask) : ggml_metal_get_buffer_id(q), 4); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(d), 5); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 6); + if (shmem > 0) { + ggml_metal_encoder_set_threadgroup_memory_size(enc, shmem, 0); + } + ggml_metal_encoder_dispatch_threadgroups(enc, ngx, ngy, 1, nth, 1, 1); + ggml_metal_op_concurrency_reset(ctx); + }; + // the dO scale (one threadgroup; every path records it so backends agree) + run("scale", 1, 1, 1024, 32*sizeof(float)); + // tiled (fa.metal FAB_*): query blocks of 8 rows, KV blocks of 16, four simdgroups; its tiles + // must fit the device's threadgroup memory (head 256: 26,176 bytes; head 512 would need 50,752, + // over Apple's 32 KiB: Cormac's review of 8ec4f9d84) + const int br = 8, bc = 16, nsg = 4; + const size_t shmem = (size_t) (2*br + 2*bc)*args.D*sizeof(uint16_t) + (2*br*bc + 2*br)*sizeof(float) + 2*br*bc*sizeof(uint16_t); + const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); + // and simdgroup matrices exist only on Apple7+ GPUs (M1 and newer): older Macs take the per-row path + if (props_dev->has_simdgroup_mm && args.D == args.Dv && args.D % 32 == 0 && shmem <= props_dev->max_theadgroup_memory_size) { + run("stats", (args.N + br - 1)/br, args.H, 32*nsg, shmem); + run("dkv", (L + bc - 1)/bc, args.Hkv, 32*nsg, shmem); + run("dq", (args.N + br - 1)/br, args.H, 32*nsg, shmem); + } else { + // one simdgroup per row + run("q", args.N*args.H, 1, 32, 0); + run("kv", L*args.Hkv, 1, 32, 0); + } + + return 1; +} + int ggml_metal_op_cross_entropy_loss_back(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); return ggml_metal_op_cross_entropy_loss_rows(ctx, op, ggml_metal_get_buffer_id(op), op->nb[1]); @@ -4248,13 +4232,14 @@ int ggml_metal_op_row_back(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_t enc = ctx->enc; const int32_t ne00 = op->src[0]->ne[0]; - const int32_t nr = ggml_nrows(op->src[0]); + const ggml_tensor * s0 = op->src[0]; + const ggml_tensor * s1 = op->src[1]; ggml_metal_kargs_row_back args = { /*.ne00 =*/ ne00, - /*.nb01 =*/ op->src[0]->nb[1], - /*.nb11 =*/ op->src[1]->nb[1], - /*.nb1 =*/ op->nb[1], + /*.nb01 =*/ s0->nb[1], s0->nb[2], s0->nb[3], + /*.nb11 =*/ s1->nb[1], s1->nb[2], s1->nb[3], + /*.nb1 =*/ op->nb[1], op->nb[2], op->nb[3], /*.param =*/ ggml_get_op_params_f32(op, 0), // eps, or scale }; @@ -4272,7 +4257,7 @@ int ggml_metal_op_row_back(ggml_metal_op_t ctx, int idx) { ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ggml_metal_encoder_set_threadgroup_memory_size(enc, 32*sizeof(float), 0); - ggml_metal_encoder_dispatch_threadgroups(enc, nr, 1, 1, nth, 1, 1); + ggml_metal_encoder_dispatch_threadgroups(enc, s0->ne[1], s0->ne[2], s0->ne[3], nth, 1, 1); return 1; } @@ -4308,6 +4293,32 @@ int ggml_metal_op_repeat_back(ggml_metal_op_t ctx, int idx) { return 1; } +int ggml_metal_op_get_rows_back(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + + ggml_metal_kargs_get_rows_back args = { + /*.ne00 =*/ (int32_t) op->src[0]->ne[0], + /*.nr =*/ (int32_t) ggml_nelements(op->src[1]), + /*.nb01 =*/ op->src[0]->nb[1], + /*.nb1 =*/ op->nb[1], + }; + + auto pipeline = ggml_metal_library_get_pipeline_row_back(lib, op); + + const int nth = std::min(args.ne00, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + ggml_metal_encoder_dispatch_threadgroups(enc, op->ne[1], 1, 1, nth, 1, 1); + + return 1; +} + int ggml_metal_op_silu_back(ggml_metal_op_t ctx, int idx) { ggml_tensor * op = ctx->node(idx); @@ -5342,7 +5353,11 @@ int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) { /*.nb0 =*/ nb0, /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, - /*.nb3 =*/ nb3 + /*.nb3 =*/ nb3, + /*.lp0 =*/ ggml_get_op_params_i32(op, 0), + /*.lp1 =*/ ggml_get_op_params_i32(op, 2), + /*.lp2 =*/ ggml_get_op_params_i32(op, 4), + /*.lp3 =*/ ggml_get_op_params_i32(op, 6), }; auto pipeline = ggml_metal_library_get_pipeline_pad(lib, op); diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h index e76003f71bca..2aaf2f196da5 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.h +++ b/ggml/src/ggml-metal/ggml-metal-ops.h @@ -76,13 +76,14 @@ int ggml_metal_op_mul_mat_id (ggml_metal_op_t ctx, int idx); int ggml_metal_op_add_id (ggml_metal_op_t ctx, int idx); int ggml_metal_op_flash_attn_ext (ggml_metal_op_t ctx, int idx); int ggml_metal_op_bin (ggml_metal_op_t ctx, int idx); -size_t ggml_metal_op_out_prod_extra(const struct ggml_tensor * op); bool ggml_metal_op_out_prod_supported(const struct ggml_tensor * op); int ggml_metal_op_out_prod (ggml_metal_op_t ctx, int idx); int ggml_metal_op_cross_entropy_loss(ggml_metal_op_t ctx, int idx); int ggml_metal_op_cross_entropy_loss_back(ggml_metal_op_t ctx, int idx); +int ggml_metal_op_flash_attn_back(ggml_metal_op_t ctx, int idx); int ggml_metal_op_row_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_repeat_back (ggml_metal_op_t ctx, int idx); +int ggml_metal_op_get_rows_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_silu_back (ggml_metal_op_t ctx, int idx); int ggml_metal_op_l2_norm (ggml_metal_op_t ctx, int idx); int ggml_metal_op_group_norm (ggml_metal_op_t ctx, int idx); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 0ae45e3d7bc3..c462c0686371 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -233,10 +233,6 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_ res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor); res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor); } break; - case GGML_OP_OUT_PROD: - { - res += ggml_metal_op_out_prod_extra(tensor); - } break; case GGML_OP_GATED_DELTA_NET_BACK: { res += ggml_metal_op_gated_delta_net_back_extra(tensor); diff --git a/ggml/src/ggml-metal/kernels/fa.metal b/ggml/src/ggml-metal/kernels/fa.metal index 6faf923c425e..58c745834ee7 100644 --- a/ggml/src/ggml-metal/kernels/fa.metal +++ b/ggml/src/ggml-metal/kernels/fa.metal @@ -2254,3 +2254,555 @@ template [[host_name("kernel_lightning_indexer_q4_1")]] kernel kernel_lightning_ template [[host_name("kernel_lightning_indexer_q5_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer; template [[host_name("kernel_lightning_indexer_q5_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer; template [[host_name("kernel_lightning_indexer_q8_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer; + +// FLASH_ATTN_BACK: the gradients of FLASH_ATTN_EXT, recomputing the scores instead of keeping +// them, so the memory is the gradients alone (the CPU reference: ggml-cpu ops.cpp +// ggml_compute_forward_flash_attn_back_f32). One simdgroup per row; each lane holds every 32nd +// element of a row. Pass Q, per query row and head: the row's log-sum-exp over the window, then +// in one more sweep D = sum_j p_j dp_j, A = sum_j p_j dp_j k_j and B = sum_j p_j k_j, so +// dq = scale (A - D B). Pass KV, per KV position and KV head: dk and dv summed over the query +// rows of every head in its group, from the lse and D the Q pass wrote. + +#define FA_BACK_MAX_D 512 + +template +static inline float fa_back_elem(device const char * row, int c); + +template <> inline float fa_back_elem(device const char * row, int c) { return ((device const float *) row)[c]; } +template <> inline float fa_back_elem (device const char * row, int c) { return (float) ((device const half *) row)[c]; } +template <> inline float fa_back_elem(device const char * row, int c) { + device const block_q8_0 * b = (device const block_q8_0 *) row + c/QK8_0; + return (float) b->d * (float) b->qs[c%QK8_0]; +} + +static inline float fa_back_mask(constant ggml_metal_kargs_flash_attn_back & args, device const char * mask, int j, int i, int h) { + if (!args.has_mask) { + return 0.0f; + } + return (float) ((device const half *) (mask + i*args.nbm1 + (h % args.nem2)*args.nbm2))[j]; +} + +template +kernel void kernel_flash_attn_back_q( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]]) { + const int i = tgpig % args.N; + const int h = tgpig / args.N; + const int hk = h / (args.H / args.Hkv); + + const int nq = args.D / 32; + const int nv = args.Dv / 32; + float qr[FA_BACK_MAX_D/32], dr[FA_BACK_MAX_D/32], a[FA_BACK_MAX_D/32], b[FA_BACK_MAX_D/32]; + for (int c = 0; c < nq; ++c) { + qr[c] = *(device const float *) (q + (tiisg + 32*c)*sizeof(float) + i*args.nbq1 + h*args.nbq2); + a[c] = 0.0f; + b[c] = 0.0f; + } + for (int c = 0; c < nv; ++c) { + dr[c] = *(device const float *) (d + (tiisg + 32*c)*sizeof(float) + h*args.nbd1 + i*args.nbd2); + } + + // the log-sum-exp, online + float m = -INFINITY; + float l = 0.0f; + for (int j = 0; j < args.M; ++j) { + const float mj = fa_back_mask(args, mask, j, i, h); + if (mj == -INFINITY) { + continue; + } + device const char * kr = k + j*args.nbk1 + hk*args.nbk2; + float s = 0.0f; + for (int c = 0; c < nq; ++c) { + s += qr[c]*fa_back_elem(kr, tiisg + 32*c); + } + s = simd_sum(s)*args.scale + mj; + const float mn = max(m, s); + l = l*exp(m - mn) + exp(s - mn); + m = mn; + } + const float lse = m == -INFINITY ? -INFINITY : m + log(l); + + // D, A and B in one sweep + float dd = 0.0f; + if (lse != -INFINITY) { + for (int j = 0; j < args.M; ++j) { + const float mj = fa_back_mask(args, mask, j, i, h); + if (mj == -INFINITY) { + continue; + } + device const char * kr = k + j*args.nbk1 + hk*args.nbk2; + device const char * vr = v + j*args.nbv1 + hk*args.nbv2; + float s = 0.0f; + for (int c = 0; c < nq; ++c) { + s += qr[c]*fa_back_elem(kr, tiisg + 32*c); + } + float dp = 0.0f; + for (int c = 0; c < nv; ++c) { + dp += dr[c]*fa_back_elem(vr, tiisg + 32*c); + } + const float p = exp(simd_sum(s)*args.scale + mj - lse); + dp = simd_sum(dp); + dd += p*dp; + for (int c = 0; c < nq; ++c) { + const float kc = fa_back_elem(kr, tiisg + 32*c); + a[c] += p*dp*kc; + b[c] += p*kc; + } + } + } + + device float * dq = (device float *) dst + (h*args.N + i)*args.D; + for (int c = 0; c < nq; ++c) { + dq[tiisg + 32*c] = args.scale*(a[c] - dd*b[c]); + } + if (tiisg == 0) { + device float * lsev = (device float *) (dst + args.offs_lse); + lsev[h*args.N + i] = lse; + lsev[args.N*args.H + h*args.N + i] = dd; + } +} + +template +kernel void kernel_flash_attn_back_kv( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + uint tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]]) { + const int L = args.M - args.kv0; + const int j = args.kv0 + tgpig % L; + const int hk = tgpig / L; + const int G = args.H / args.Hkv; + + const int nq = args.D / 32; + const int nv = args.Dv / 32; + float kr[FA_BACK_MAX_D/32], vr[FA_BACK_MAX_D/32], gk[FA_BACK_MAX_D/32], gv[FA_BACK_MAX_D/32]; + device const char * kp = k + j*args.nbk1 + hk*args.nbk2; + device const char * vp = v + j*args.nbv1 + hk*args.nbv2; + for (int c = 0; c < nq; ++c) { + kr[c] = fa_back_elem(kp, tiisg + 32*c); + gk[c] = 0.0f; + } + for (int c = 0; c < nv; ++c) { + vr[c] = fa_back_elem(vp, tiisg + 32*c); + gv[c] = 0.0f; + } + device const float * lsev = (device const float *) (dst + args.offs_lse); + for (int h = hk*G; h < (hk + 1)*G; ++h) { + for (int i = 0; i < args.N; ++i) { + const float lse = lsev[h*args.N + i]; + const float mj = fa_back_mask(args, mask, j, i, h); + if (lse == -INFINITY || mj == -INFINITY) { + continue; + } + device const char * qp = q + i*args.nbq1 + h*args.nbq2; + device const char * dp = d + h*args.nbd1 + i*args.nbd2; + float s = 0.0f; + float t = 0.0f; + for (int c = 0; c < nq; ++c) { + s += ((device const float *) qp)[tiisg + 32*c]*kr[c]; + } + for (int c = 0; c < nv; ++c) { + t += ((device const float *) dp)[tiisg + 32*c]*vr[c]; + } + const float p = exp(simd_sum(s)*args.scale + mj - lse); + const float ds = p*(simd_sum(t) - lsev[args.N*args.H + h*args.N + i])*args.scale; + for (int c = 0; c < nq; ++c) { + gk[c] += ds*((device const float *) qp)[tiisg + 32*c]; + } + for (int c = 0; c < nv; ++c) { + gv[c] += p*((device const float *) dp)[tiisg + 32*c]; + } + } + } + device float * dk = (device float *) (dst + args.offs_dk) + (hk*L + j - args.kv0)*args.D; + device float * dv = (device float *) (dst + args.offs_dv) + (hk*L + j - args.kv0)*args.Dv; + for (int c = 0; c < nq; ++c) { + dk[tiisg + 32*c] = gk[c]; + } + for (int c = 0; c < nv; ++c) { + dv[tiisg + 32*c] = gv[c]; + } +} + + +// THE TILED BACKWARD (FlashAttention-2 style, simdgroup matrices; D == Dv, D % 32 == 0). Query +// blocks of FAB_BR rows and KV blocks of FAB_BC positions; the four simdgroups of a threadgroup +// each compute one 8x8 tile of S = Q K^T or dP = dO V^T per block (inputs staged as half, +// accumulated in F32), and the gradients accumulate in simdgroup matrices split by column. +// pass A (query blocks): each row's log-sum-exp and D = sum_j p_j dp_j, online in one sweep +// pass B (KV blocks): dK += dS^T Q, dV += P^T dO over every query row of the KV head's group +// pass C (query blocks): dQ += dS K +#define FAB_BR 8 +#define FAB_BC 16 +#define FAB_NSG 4 +// dO (and so dP, D, dS and the gradients) is carried through the half tiles scaled by the power of +// two that brings its largest magnitude to [4, 8) (ggml_flash_attn_back_scale; pass "scale" below): +// training gradients sit far below half's normal range (6.1e-5) and lost their precision there +// (test-walk-exact, M5 2026-10-14: 0.9943 against the F32 gradient unscaled, 1.0000 scaled) + +#define FAB_SCALE_NT 1024 // the scale pass's threads (ggml_metal_op_flash_attn_back dispatches this many) + +kernel void kernel_flash_attn_back_scale( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]]) { + const ushort ntg = FAB_SCALE_NT; + threadgroup float * smax = (threadgroup float *) shmem; + float mx = 0.0f; + for (int r = 0; r < args.N*args.H; ++r) { + device const float * row = (device const float *) (d + (r % args.H)*args.nbd1 + (r / args.H)*args.nbd2); + for (int c = tiitg; c < args.Dv; c += ntg) { + mx = max(mx, fabs(row[c])); + } + } + mx = simd_max(mx); + if (tiisg == 0) { + smax[sgitg] = mx; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiitg == 0) { + for (int s = 1; s < ntg/32; ++s) { + mx = max(mx, smax[s]); + } + float scale = 1.0f; + if (mx > 0.0f) { + int e; + frexp(mx, e); + scale = ldexp(1.0f, clamp(3 - e, -24, 40)); + } + ((device float *) (dst + args.offs_lse))[2*args.N*args.H] = scale; + } +} + +// rows [r0, r0 + FAB_BR|FAB_BC) of a head into threadgroup half, zeros past `n` rows +// (rows outer, threads across a row's columns: no division per element) +template +static inline void fab_stage_kv(threadgroup half * dst, device const char * base, uint64_t nb_row, int r0, int rows, int n, int Dx, + ushort tid) { + for (int r = 0; r < rows; ++r) { + device const char * row = base + (r0 + r)*nb_row; + const bool in = r0 + r < n; + for (int c = tid; c < Dx; c += FAB_NSG*32) { + dst[r*Dx + c] = in ? (half) fa_back_elem(row, c) : (half) 0.0f; + } + } +} + +static inline void fab_stage_rows(threadgroup half * dst, device const char * base, uint64_t nb_row, int r0, int n, int Dx, + ushort tid, float scale = 1.0f) { + for (int r = 0; r < FAB_BR; ++r) { + device const float * row = (device const float *) (base + (r0 + r)*nb_row); + const bool in = r0 + r < n; + for (int c = tid; c < Dx; c += FAB_NSG*32) { + dst[r*Dx + c] = in ? (half) (row[c]*scale) : (half) 0.0f; + } + } +} + +// S (sg 0, 1: keys 0-7, 8-15) and dP (sg 2, 3) for one block into ss / sdp [FAB_BR][FAB_BC] +static inline void fab_scores(threadgroup const half * sq, threadgroup const half * sdo, threadgroup const half * sk, + threadgroup const half * sv, threadgroup float * ss, threadgroup float * sdp, int D, ushort sgitg) { + const bool dp = sgitg >= 2; + const int t = sgitg & 1; + threadgroup const half * a = dp ? sdo : sq; + threadgroup const half * b = (dp ? sv : sk) + 8*t*D; + simdgroup_float8x8 acc = make_filled_simdgroup_matrix(0.0f); + simdgroup_half8x8 ma, mb; + for (int d0 = 0; d0 < D; d0 += 8) { + simdgroup_load(ma, a + d0, D); + simdgroup_load(mb, b + d0, D, 0, true); // K^T: [d][key] + simdgroup_multiply_accumulate(acc, ma, mb, acc); + } + simdgroup_store(acc, (dp ? sdp : ss) + 8*t, FAB_BC); +} + +template +kernel void kernel_flash_attn_back_stats( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const float ds = ((device const float *) (dst + args.offs_lse))[2*args.N*args.H]; // the dO scale + const int D = args.D; + const int h = tgpig.y; + const int i0 = tgpig.x*FAB_BR; + const int hk = h / (args.H / args.Hkv); + + threadgroup half * sq = (threadgroup half *) shmem; + threadgroup half * sdo = sq + FAB_BR*D; + threadgroup half * sk = sdo + FAB_BR*D; + threadgroup half * sv = sk + FAB_BC*D; + threadgroup float * ss = (threadgroup float *) (sv + FAB_BC*D); + threadgroup float * sdp = ss + FAB_BR*FAB_BC; + + fab_stage_rows(sq, q + h*args.nbq2, args.nbq1, i0, args.N, D, tiitg); + fab_stage_rows(sdo, d + h*args.nbd1, args.nbd2, i0, args.N, D, tiitg, ds); + + float m = -INFINITY, l = 0.0f, a = 0.0f; // thread r < FAB_BR owns row i0 + r + for (int j0 = 0; j0 < args.M; j0 += FAB_BC) { + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_stage_kv(sk, k + hk*args.nbk2, args.nbk1, j0, FAB_BC, args.M, D, tiitg); + fab_stage_kv(sv, v + hk*args.nbv2, args.nbv1, j0, FAB_BC, args.M, D, tiitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_scores(sq, sdo, sk, sv, ss, sdp, D, sgitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + if (tiitg < FAB_BR && i0 + tiitg < args.N) { + const int i = i0 + tiitg; + for (int c = 0; c < FAB_BC && j0 + c < args.M; ++c) { + const float sc = ss[tiitg*FAB_BC + c]*args.scale + fa_back_mask(args, mask, j0 + c, i, h); + if (sc == -INFINITY) { + continue; + } + const float mn = max(m, sc); + const float r = exp(m - mn); + const float e = exp(sc - mn); + l = l*r + e; + a = a*r + e*sdp[tiitg*FAB_BC + c]; + m = mn; + } + } + } + if (tiitg < FAB_BR && i0 + tiitg < args.N) { + device float * lsev = (device float *) (dst + args.offs_lse); + lsev[h*args.N + i0 + tiitg] = m == -INFINITY ? -INFINITY : m + log(l); + lsev[args.N*args.H + h*args.N + i0 + tiitg] = m == -INFINITY ? 0.0f : a/l/ds; + } +} + +// P and dS of one block from ss / sdp into half sp / sds [FAB_BR][FAB_BC] +static inline void fab_probs(constant ggml_metal_kargs_flash_attn_back & args, device const char * mask, + threadgroup const float * ss, threadgroup const float * sdp, threadgroup const float * slse, + threadgroup half * sp, threadgroup half * sds, int i0, int j0, int h, ushort tiitg) { + const int r = tiitg / FAB_BC; + const int c = tiitg % FAB_BC; + const int i = i0 + r; + const int j = j0 + c; + float p = 0.0f; + float ds = 0.0f; + if (i < args.N && j < args.M && slse[r] != -INFINITY) { + const float sc = ss[r*FAB_BC + c]*args.scale + fa_back_mask(args, mask, j, i, h); + p = sc == -INFINITY ? 0.0f : exp(sc - slse[r]); + ds = p*(sdp[r*FAB_BC + c] - slse[FAB_BR + r])*args.scale; + } + sp [r*FAB_BC + c] = (half) p; + sds[r*FAB_BC + c] = (half) ds; +} + +template +kernel void kernel_flash_attn_back_dkv( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const float ds = ((device const float *) (dst + args.offs_lse))[2*args.N*args.H]; // the dO scale + const int D = args.D; + const int hk = tgpig.y; + const int j0 = args.kv0 + tgpig.x*FAB_BC; + const int G = args.H / args.Hkv; + const int CW = D / FAB_NSG; // the columns this simdgroup owns + const int c0 = sgitg*CW; + + threadgroup half * sq = (threadgroup half *) shmem; + threadgroup half * sdo = sq + FAB_BR*D; + threadgroup half * sk = sdo + FAB_BR*D; + threadgroup half * sv = sk + FAB_BC*D; + threadgroup float * ss = (threadgroup float *) (sv + FAB_BC*D); + threadgroup float * sdp = ss + FAB_BR*FAB_BC; + threadgroup float * slse = sdp + FAB_BR*FAB_BC; // [lse x FAB_BR][D x FAB_BR] + threadgroup half * sp = (threadgroup half *) (slse + 2*FAB_BR); + threadgroup half * sds = sp + FAB_BR*FAB_BC; + + fab_stage_kv(sk, k + hk*args.nbk2, args.nbk1, j0, FAB_BC, args.M, D, tiitg); + fab_stage_kv(sv, v + hk*args.nbv2, args.nbv1, j0, FAB_BC, args.M, D, tiitg); + + simdgroup_float8x8 gk[2][FA_BACK_MAX_D/8/FAB_NSG]; + simdgroup_float8x8 gv[2][FA_BACK_MAX_D/8/FAB_NSG]; + for (int r = 0; r < 2; ++r) { + for (int t = 0; t < CW/8; ++t) { + gk[r][t] = make_filled_simdgroup_matrix(0.0f); + gv[r][t] = make_filled_simdgroup_matrix(0.0f); + } + } + device const float * lsev = (device const float *) (dst + args.offs_lse); + for (int h = hk*G; h < (hk + 1)*G; ++h) { + for (int i0 = 0; i0 < args.N; i0 += FAB_BR) { + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_stage_rows(sq, q + h*args.nbq2, args.nbq1, i0, args.N, D, tiitg); + fab_stage_rows(sdo, d + h*args.nbd1, args.nbd2, i0, args.N, D, tiitg, ds); + if (tiitg < FAB_BR) { + const int i = i0 + tiitg; + slse[tiitg] = i < args.N ? lsev[h*args.N + i] : -INFINITY; + slse[FAB_BR + tiitg] = i < args.N ? lsev[args.N*args.H + h*args.N + i]*ds : 0.0f; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_scores(sq, sdo, sk, sv, ss, sdp, D, sgitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_probs(args, mask, ss, sdp, slse, sp, sds, i0, j0, h, tiitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + // dV[16 x CW] += P^T[16 x 8] dO[8 x CW]; dK[16 x CW] += dS^T[16 x 8] Q[8 x CW] + simdgroup_half8x8 pt, dst8, mo, mq; + for (int r = 0; r < 2; ++r) { + simdgroup_load(pt, sp + 8*r, FAB_BC, 0, true); + simdgroup_load(dst8, sds + 8*r, FAB_BC, 0, true); + for (int t = 0; t < CW/8; ++t) { + simdgroup_load(mo, sdo + c0 + 8*t, D); + simdgroup_load(mq, sq + c0 + 8*t, D); + simdgroup_multiply_accumulate(gv[r][t], pt, mo, gv[r][t]); + simdgroup_multiply_accumulate(gk[r][t], dst8, mq, gk[r][t]); + } + } + } + } + // through threadgroup memory (the staged K/V are done with), so a partial block writes only its rows + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float * so = (threadgroup float *) sk; // [FAB_BC][D] + for (int part = 0; part < 2; ++part) { + for (int r = 0; r < 2; ++r) { + for (int t = 0; t < CW/8; ++t) { + simdgroup_store(part == 0 ? gk[r][t] : gv[r][t], so + 8*r*D + c0 + 8*t, D); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + device float * out = (device float *) (dst + (part == 0 ? args.offs_dk : args.offs_dv)) + (hk*(args.M - args.kv0) + j0 - args.kv0)*D; + for (int e = tiitg; e < FAB_BC*D; e += FAB_NSG*32) { + if (j0 + e / D < args.M) { + out[e] = so[e]/ds; + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } +} + +template +kernel void kernel_flash_attn_back_dq( + constant ggml_metal_kargs_flash_attn_back & args, + device const char * q, + device const char * k, + device const char * v, + device const char * mask, + device const char * d, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const float ds = ((device const float *) (dst + args.offs_lse))[2*args.N*args.H]; // the dO scale + const int D = args.D; + const int h = tgpig.y; + const int i0 = tgpig.x*FAB_BR; + const int hk = h / (args.H / args.Hkv); + const int CW = D / FAB_NSG; + const int c0 = sgitg*CW; + + threadgroup half * sq = (threadgroup half *) shmem; + threadgroup half * sdo = sq + FAB_BR*D; + threadgroup half * sk = sdo + FAB_BR*D; + threadgroup half * sv = sk + FAB_BC*D; + threadgroup float * ss = (threadgroup float *) (sv + FAB_BC*D); + threadgroup float * sdp = ss + FAB_BR*FAB_BC; + threadgroup float * slse = sdp + FAB_BR*FAB_BC; + threadgroup half * sp = (threadgroup half *) (slse + 2*FAB_BR); + threadgroup half * sds = sp + FAB_BR*FAB_BC; + + fab_stage_rows(sq, q + h*args.nbq2, args.nbq1, i0, args.N, D, tiitg); + fab_stage_rows(sdo, d + h*args.nbd1, args.nbd2, i0, args.N, D, tiitg, ds); + device const float * lsev = (device const float *) (dst + args.offs_lse); + if (tiitg < FAB_BR) { + const int i = i0 + tiitg; + slse[tiitg] = i < args.N ? lsev[h*args.N + i] : -INFINITY; + slse[FAB_BR + tiitg] = i < args.N ? lsev[args.N*args.H + h*args.N + i]*ds : 0.0f; + } + + simdgroup_float8x8 gq[FA_BACK_MAX_D/8/FAB_NSG]; + for (int t = 0; t < CW/8; ++t) { + gq[t] = make_filled_simdgroup_matrix(0.0f); + } + for (int j0 = 0; j0 < args.M; j0 += FAB_BC) { + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_stage_kv(sk, k + hk*args.nbk2, args.nbk1, j0, FAB_BC, args.M, D, tiitg); + fab_stage_kv(sv, v + hk*args.nbv2, args.nbv1, j0, FAB_BC, args.M, D, tiitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_scores(sq, sdo, sk, sv, ss, sdp, D, sgitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + fab_probs(args, mask, ss, sdp, slse, sp, sds, i0, j0, h, tiitg); + threadgroup_barrier(mem_flags::mem_threadgroup); + // dQ[8 x CW] += dS[8 x 16] K[16 x CW] + simdgroup_half8x8 ms, mk; + for (int kk = 0; kk < 2; ++kk) { + simdgroup_load(ms, sds + 8*kk, FAB_BC); + for (int t = 0; t < CW/8; ++t) { + simdgroup_load(mk, sk + 8*kk*D + c0 + 8*t, D); + simdgroup_multiply_accumulate(gq[t], ms, mk, gq[t]); + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float * so = (threadgroup float *) sk; // [FAB_BR][D] + for (int t = 0; t < CW/8; ++t) { + simdgroup_store(gq[t], so + c0 + 8*t, D); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + device float * out = (device float *) dst + (h*args.N + i0)*D; + for (int e = tiitg; e < FAB_BR*D; e += FAB_NSG*32) { + if (i0 + e / D < args.N) { + out[e] = so[e]/ds; + } + } +} + +typedef decltype(kernel_flash_attn_back_stats) kernel_flash_attn_back_tiled_t; + +template [[host_name("kernel_flash_attn_back_stats_f32")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_stats; +template [[host_name("kernel_flash_attn_back_stats_f16")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_stats; +template [[host_name("kernel_flash_attn_back_stats_q8_0")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_stats; +template [[host_name("kernel_flash_attn_back_dkv_f32")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dkv; +template [[host_name("kernel_flash_attn_back_dkv_f16")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dkv; +template [[host_name("kernel_flash_attn_back_dkv_q8_0")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dkv; +template [[host_name("kernel_flash_attn_back_dq_f32")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dq; +template [[host_name("kernel_flash_attn_back_dq_f16")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dq; +template [[host_name("kernel_flash_attn_back_dq_q8_0")]] kernel kernel_flash_attn_back_tiled_t kernel_flash_attn_back_dq; + +typedef decltype(kernel_flash_attn_back_q) kernel_flash_attn_back_q_t; +typedef decltype(kernel_flash_attn_back_kv) kernel_flash_attn_back_kv_t; + +template [[host_name("kernel_flash_attn_back_q_f32")]] kernel kernel_flash_attn_back_q_t kernel_flash_attn_back_q; +template [[host_name("kernel_flash_attn_back_q_f16")]] kernel kernel_flash_attn_back_q_t kernel_flash_attn_back_q; +template [[host_name("kernel_flash_attn_back_q_q8_0")]] kernel kernel_flash_attn_back_q_t kernel_flash_attn_back_q; +template [[host_name("kernel_flash_attn_back_kv_f32")]] kernel kernel_flash_attn_back_kv_t kernel_flash_attn_back_kv; +template [[host_name("kernel_flash_attn_back_kv_f16")]] kernel kernel_flash_attn_back_kv_t kernel_flash_attn_back_kv; +template [[host_name("kernel_flash_attn_back_kv_q8_0")]] kernel kernel_flash_attn_back_kv_t kernel_flash_attn_back_kv; diff --git a/ggml/src/ggml-metal/kernels/misc.metal b/ggml/src/ggml-metal/kernels/misc.metal index 11104b4d8d12..6206ef9eae2e 100644 --- a/ggml/src/ggml-metal/kernels/misc.metal +++ b/ggml/src/ggml-metal/kernels/misc.metal @@ -127,11 +127,13 @@ kernel void kernel_pad_impl( const int32_t k0 = tgpig.x/args.ne1; const int32_t i1 = tgpig.x - k0*args.ne1; - const int32_t i03 = i3; - const int32_t i02 = i2; - const int32_t i01 = i1; + // the source element of dst[i] is src[i - lp]; outside the source it is zero + const int32_t i03 = i3 - args.lp3; + const int32_t i02 = i2 - args.lp2; + const int32_t i01 = i1 - args.lp1; + const bool row = i01 >= 0 && i01 < args.ne01 && i02 >= 0 && i02 < args.ne02 && i03 >= 0 && i03 < args.ne03; - device const T * src0_ptr = (device const T *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01); + device const char * src0_row = src0 + (row ? i03*args.nb03 + i02*args.nb02 + i01*args.nb01 : 0); device T * dst_ptr = (device T *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1); for (int32_t l0 = 0; l0 < 1024; l0 += ntg.x) { @@ -140,8 +142,9 @@ kernel void kernel_pad_impl( break; } - if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) { - dst_ptr[i0] = src0_ptr[i0]; + const int32_t i00 = i0 - args.lp0; + if (row && i00 >= 0 && i00 < args.ne00) { + dst_ptr[i0] = *(device const T *) (src0_row + i00*args.nb00); // a view may have any element stride } else { dst_ptr[i0] = 0.0f; } diff --git a/ggml/src/ggml-metal/kernels/out_prod.metal b/ggml/src/ggml-metal/kernels/out_prod.metal new file mode 100644 index 000000000000..5156dd0f0539 --- /dev/null +++ b/ggml/src/ggml-metal/kernels/out_prod.metal @@ -0,0 +1,157 @@ +#include "common.h" +#include "dequantize.h" + +// OUT_PROD: dst[m, n] = sum_k src0[m, k] * src1[n, k]. It is the backward of a matmul through +// a frozen weight (dx = out_prod(W, dy^T)) and the gradient of a LoRA factor (dA, dB). +// +// This is mul_mm with the contraction along the ROWS of both operands. Each thread reads 16 +// consecutive elements of one src0 row, in the weight's own block order, dequantizes them in +// registers, and stages them transposed in threadgroup memory. Nothing is copied or transposed +// in device memory. The tiles and the simdgroup product are mul_mm's: a 64 x 32 output tile, +// 32-deep k steps, 4 simdgroups. Staging is float because gradients exceed half's range. +// +// nl == 0 selects the element path (F32/F16/BF16, any width, a LoRA rank of 8 included). +// nl > 0 dequantizes whole 16-element chunks, which needs ne00 % 16 == 0. + +template +kernel void kernel_out_prod( + constant ggml_metal_kargs_out_prod & args, + device const char * src0, + device const char * src1, + device char * dst, + threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiitg[[thread_index_in_threadgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { + constexpr short OP_NR0 = 64; // output rows (m) per threadgroup + constexpr short OP_NR1 = 32; // output columns (n) per threadgroup + constexpr short OP_NK = 32; // contraction depth per step + + threadgroup float * sa = (threadgroup float *) shmem; // OP_NR0 x OP_NK + threadgroup float * sb = sa + OP_NR0*OP_NK; // OP_NR1 x OP_NK + + const int r0 = tgpig.y*OP_NR0; + const int r1 = tgpig.x*OP_NR1; + const int i2 = tgpig.z % args.ne2; + const int i3 = tgpig.z / args.ne2; + + // src0 is broadcast over dims 2-3 (grouped heads); src1 and dst are not + src0 += (i2/args.dps2)*args.nb02 + (i3/args.dps3)*args.nb03; + src1 += i2*args.nb12 + i3*args.nb13; + + // src0 staging: thread -> one k row, 16 consecutive m + const short ka = tiitg/4; + const short ma = 16*(tiitg%4); + // src1 staging: thread -> one n, 8 consecutive k + const short nb = tiitg/4; + const short kb = 8*(tiitg%4); + + simdgroup_float8x8 ma8[4]; + simdgroup_float8x8 mb8[2]; + simdgroup_float8x8 mc[8]; + for (short i = 0; i < 8; i++) { + mc[i] = make_filled_simdgroup_matrix(0.f); + } + + for (int k0 = 0; k0 < args.ne01; k0 += OP_NK) { + float4x4 a = float4x4(0.f); + const int k = k0 + ka; + const int m = r0 + ma; + if (k < args.ne01 && m < args.ne00) { + device const char * row = src0 + (uint64_t) k*args.nb01; + if (nl == 0) { + device const T0 * x = (device const T0 *) row + m; + for (short i = 0; i < 16; i++) { + a[i/4][i%4] = m + i < args.ne00 ? (float) x[i] : 0.f; + } + } else { + dequantize_func((device const block_q *) row + m/(16*nl), (m/16) % nl, a); + } + } + + float b[8]; + { + const int n = r1 + nb; + for (short i = 0; i < 8; i++) { + const int kk = k0 + kb + i; + b[i] = kk < args.ne01 && n < args.ne10 + ? *(device const float *) (src1 + (uint64_t) kk*args.nb11 + (uint64_t) n*args.nb10) : 0.f; + } + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // the same 8x8-block layout mul_mm stages: sa block (k/8, m/8) holds [k%8][m%8], + // sb block (k/8, n/8) holds [n%8][k%8] + for (short i = 0; i < 16; i++) { + const short mm = ma + i; + sa[64*(8*(ka/8) + mm/8) + 8*(ka%8) + mm%8] = a[i/4][i%4]; + } + for (short i = 0; i < 8; i++) { + const short kk = kb + i; + sb[64*(4*(kk/8) + nb/8) + 8*(nb%8) + kk%8] = b[i]; + } + + threadgroup_barrier(mem_flags::mem_threadgroup); + + threadgroup const float * lsma = sa + 4*64*(sgitg%2); + threadgroup const float * lsmb = sb + 2*64*(sgitg/2); + for (short ik = 0; ik < OP_NK/8; ik++) { + for (short i = 0; i < 4; i++) { + simdgroup_load(ma8[i], lsma + 64*i, 8, 0, false); + } + for (short i = 0; i < 2; i++) { + simdgroup_load(mb8[i], lsmb + 64*i, 8, 0, false); + } + for (short i = 0; i < 8; i++) { + simdgroup_multiply_accumulate(mc[i], mb8[i/4], ma8[i%4], mc[i]); + } + lsma += 8*64; + lsmb += 4*64; + } + } + + device float * d = (device float *) dst + (uint64_t) tgpig.z*args.ne00*args.ne10; + + if (r0 + OP_NR0 <= args.ne00 && r1 + OP_NR1 <= args.ne10) { + device float * C = d + (r0 + 32*(sgitg & 1)) + (r1 + 16*(sgitg >> 1))*args.ne00; + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne00*(i/4), args.ne00, 0, false); + } + return; + } + + // an edge tile: stage it, then write only the part inside dst + threadgroup_barrier(mem_flags::mem_threadgroup); + threadgroup float * t = (threadgroup float *) shmem + 32*(sgitg & 1) + 16*(sgitg >> 1)*OP_NR0; + for (short i = 0; i < 8; i++) { + simdgroup_store(mc[i], t + 8*(i%4) + 8*OP_NR0*(i/4), OP_NR0, 0, false); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + const int nr0 = min((int) OP_NR0, args.ne00 - r0); + const int nr1 = min((int) OP_NR1, args.ne10 - r1); + for (int j = tiitg; j < nr0*nr1; j += 128) { + const int jm = j % nr0; + const int jn = j / nr0; + d[(r0 + jm) + (r1 + jn)*args.ne00] = ((threadgroup float *) shmem)[jm + jn*OP_NR0]; + } +} + +typedef decltype(kernel_out_prod) out_prod_t; + +template [[host_name("kernel_out_prod_f32")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_f16")]] kernel out_prod_t kernel_out_prod; +#if defined(GGML_METAL_HAS_BF16) +template [[host_name("kernel_out_prod_bf16")]] kernel out_prod_t kernel_out_prod; +#endif +template [[host_name("kernel_out_prod_q4_0")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q4_1")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q5_0")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q5_1")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q8_0")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q2_K")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q3_K")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q4_K")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q5_K")]] kernel out_prod_t kernel_out_prod; +template [[host_name("kernel_out_prod_q6_K")]] kernel out_prod_t kernel_out_prod; diff --git a/ggml/src/ggml-metal/kernels/softmax.metal b/ggml/src/ggml-metal/kernels/softmax.metal index a6ad5baaa9bd..d7fae3830425 100644 --- a/ggml/src/ggml-metal/kernels/softmax.metal +++ b/ggml/src/ggml-metal/kernels/softmax.metal @@ -350,6 +350,89 @@ kernel void kernel_cross_entropy_loss_back_f32( } +// the same with one target index per row (I32 labels, -1: the row carries no loss), so no +// vocabulary-wide one-hot tensor exists beside the logits +kernel void kernel_cross_entropy_loss_i32( + constant ggml_metal_kargs_cross_entropy_loss & args, + device const char * src0, + device const char * src1, + device char * dst, + device const float * weights, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t row = tgpig.x; + const int32_t t = ((device const int32_t *) src1)[row]; + + device const float * s0 = (device const float *) (src0 + row*args.nb01); + + float lmax = -INFINITY; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lmax = MAX(lmax, s0[i]); + } + const float max_val = ce_reduce_max(lmax, buf, sgitg, tiisg, tptg); + + float lsum = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lsum += exp(s0[i] - max_val); + } + const float lse = log(ce_reduce_sum(lsum, buf, sgitg, tiisg, tptg)); + + if (tpitg.x == 0) { + const float w = args.has_weights ? weights[row] : 1.0f; + ((device float *) dst)[row] = t >= 0 && w != 0.0f ? -w*(s0[t] - max_val - lse)/(float) args.nr : 0.0f; + } +} + +kernel void kernel_cross_entropy_loss_back_i32( + constant ggml_metal_kargs_cross_entropy_loss & args, + device const char * src0, // grad: one float + device const char * src1, // logits + device const char * src2, // labels: one target per row + device char * dst, + device const float * weights, + threadgroup float * buf [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint sgitg[[simdgroup_index_in_threadgroup]], + uint tiisg[[thread_index_in_simdgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t row = tgpig.x; + const int32_t t = ((device const int32_t *) src2)[row]; + + device const float * s0 = (device const float *) (src1 + row*args.nb01); + device float * d = (device float *) (dst + row*args.nb1); + + const float w = args.has_weights ? weights[row] : 1.0f; + if (t < 0 || w == 0.0f) { + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + d[i] = 0.0f; + } + return; + } + + const float d_by_nr = w*((device const float *) src0)[0]/(float) args.nr; + + float lmax = -INFINITY; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lmax = MAX(lmax, s0[i]); + } + const float max_val = ce_reduce_max(lmax, buf, sgitg, tiisg, tptg); + + float lsum = 0.0f; + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + lsum += exp(s0[i] - max_val); + } + const float inv_sum = 1.0f/ce_reduce_sum(lsum, buf, sgitg, tiisg, tptg); + + for (int i = tpitg.x; i < args.ne00; i += tptg.x) { + d[i] = (exp(s0[i] - max_val)*inv_sum - (i == t ? 1.0f : 0.0f))*d_by_nr; + } +} + // RMS_NORM_BACK: dx = (dz - x * sum(x*dz)/(sum(x*x) + eps*n)) / sqrt(sum(x*x)/n + eps) // (the CPU reference, ggml-cpu ops.cpp rms_norm_back_f32). src0 = dz, src1 = x. kernel void kernel_rms_norm_back_f32( @@ -363,11 +446,9 @@ kernel void kernel_rms_norm_back_f32( uint sgitg[[simdgroup_index_in_threadgroup]], uint tiisg[[thread_index_in_simdgroup]], uint3 tptg[[threads_per_threadgroup]]) { - const int32_t row = tgpig.x; - - device const float * dz = (device const float *) (src0 + row*args.nb01); - device const float * x = (device const float *) (src1 + row*args.nb11); - device float * dx = (device float *) (dst + row*args.nb1); + device const float * dz = (device const float *) (src0 + tgpig.x*args.nb01 + tgpig.y*args.nb02 + tgpig.z*args.nb03); + device const float * x = (device const float *) (src1 + tgpig.x*args.nb11 + tgpig.y*args.nb12 + tgpig.z*args.nb13); + device float * dx = (device float *) (dst + tgpig.x*args.nb1 + tgpig.y*args.nb2 + tgpig.z*args.nb3); float lxx = 0.0f; float lxdz = 0.0f; @@ -400,11 +481,9 @@ kernel void kernel_soft_max_back_f32( uint sgitg[[simdgroup_index_in_threadgroup]], uint tiisg[[thread_index_in_simdgroup]], uint3 tptg[[threads_per_threadgroup]]) { - const int32_t row = tgpig.x; - - device const float * dy = (device const float *) (src0 + row*args.nb01); - device const float * y = (device const float *) (src1 + row*args.nb11); - device float * dx = (device float *) (dst + row*args.nb1); + device const float * dy = (device const float *) (src0 + tgpig.x*args.nb01 + tgpig.y*args.nb02 + tgpig.z*args.nb03); + device const float * y = (device const float *) (src1 + tgpig.x*args.nb11 + tgpig.y*args.nb12 + tgpig.z*args.nb13); + device float * dx = (device float *) (dst + tgpig.x*args.nb1 + tgpig.y*args.nb2 + tgpig.z*args.nb3); float ldot = 0.0f; for (int i = tpitg.x; i < args.ne00; i += tptg.x) { @@ -456,3 +535,27 @@ kernel void kernel_repeat_back_f32( d[k0] = acc; } } + +// GET_ROWS_BACK: one threadgroup per dst row r; each thread owns columns and sums, in id order, +// every src0 row whose id is r. No atomics, and the order of the sum is the CPU reference's. +kernel void kernel_get_rows_back_f32( + constant ggml_metal_kargs_get_rows_back & args, + device const char * src0, + device const int32_t * ids, + device char * dst, + uint3 tgpig[[threadgroup_position_in_grid]], + uint3 tpitg[[thread_position_in_threadgroup]], + uint3 tptg[[threads_per_threadgroup]]) { + const int32_t r = tgpig.x; + device float * d = (device float *) (dst + r*args.nb1); + + for (int c = tpitg.x; c < args.ne00; c += tptg.x) { + float acc = 0.0f; + for (int i = 0; i < args.nr; i++) { + if (ids[i] == r) { + acc += ((device const float *) (src0 + i*args.nb01))[c]; + } + } + d[c] = acc; + } +} diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 87cd73369d30..e2a1d0fcb2dc 100644 --- a/ggml/src/ggml-opt.cpp +++ b/ggml/src/ggml-opt.cpp @@ -90,6 +90,11 @@ struct ggml_opt_context { // graph prepared last (ggml_opt_prepare_alloc clears them) std::unordered_map next_grad_accs; std::vector grad_accs; + // a weighted cross-entropy for the next graph, and that graph's weights (one per label row) + bool next_label_weights = false; + struct ggml_tensor * label_weights = nullptr; + // ggml_opt_prepare_alloc_head: the outputs' row slices, in row order (empty: one outputs) + std::vector output_slices; int64_t iter = 1; int32_t opt_period = 1; @@ -327,6 +332,13 @@ static ggml_tensor * ggml_opt_recompute_node( if (node->flags & (GGML_TENSOR_FLAG_PARAM | GGML_TENSOR_FLAG_INPUT)) { return node; } + // an in-place write into a buffer the forward did not compute (a cache, a state) is a side + // effect: a clone would apply it a second time to the live buffer (Cormac's review of #5057) + const bool view_op = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE; + if (node->view_src != nullptr && !view_op && !forward.count(ggml_opt_view_root(node))) { + return node; + } switch (node->op) { case GGML_OP_NONE: case GGML_OP_SET_ROWS: @@ -365,7 +377,8 @@ static ggml_tensor * ggml_opt_recompute_node( static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd, const char * prefix) { std::unordered_set forward; std::unordered_set keep; - std::unordered_set mutated; // roots of the buffers the forward pass writes in place + // roots of the buffers the forward pass writes in place, each with its last write's position + std::unordered_map mutated; const size_t plen = strlen(prefix); for (int i = 0; i < n_fwd; ++i) { ggml_tensor * node = g->nodes[i]; @@ -378,18 +391,27 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd, const bool view_op = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE; if ((node->op == GGML_OP_CPY && !ggml_opt_is_cast(node)) || node->op == GGML_OP_SET_ROWS || (node->view_src != nullptr && !view_op)) { - mutated.insert(ggml_opt_view_root(node)); + mutated[ggml_opt_view_root(node)] = i; } } + // A clone runs after the whole forward pass, so it sees a mutated buffer as the forward left + // it: the same value the original read unless the buffer is written AFTER the original. A read + // of a buffer whose writes all precede it (a layer's cache read after its own chunk's write) is + // recomputed like any other node; keeping it held a window-sized cast of the cache per layer + // from the forward to the backward (the M5, qwen3.5-4b, q8_0 cache, window 32768: 16 x 64 MiB). + auto written_after = [&](ggml_tensor * t, int i) { + const auto it = mutated.find(ggml_opt_view_root(t)); + return it != mutated.end() && it->second > i; + }; for (int i = 0; i < n_fwd; ++i) { ggml_tensor * node = g->nodes[i]; bool reads_mutated = false; for (int k = 0; k < GGML_MAX_SRC && !reads_mutated; ++k) { - reads_mutated = node->src[k] != nullptr && mutated.count(ggml_opt_view_root(node->src[k])); + reads_mutated = node->src[k] != nullptr && written_after(node->src[k], i); } if (strncmp(ggml_get_name(node), prefix, plen) == 0 || (node->flags & GGML_TENSOR_FLAG_LOSS) || (node->flags & GGML_TENSOR_FLAG_OUTPUT) || reads_mutated || - (node->view_src != nullptr && mutated.count(ggml_opt_view_root(node)))) { + (node->view_src != nullptr && written_after(node, i))) { keep.insert(node); } } @@ -440,6 +462,7 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd, GGML_ASSERT((int) order.size() <= g->size && "recompute clones overflow the graph: raise the training node budget"); std::copy(order.begin(), order.end(), g->nodes); g->n_nodes = (int) order.size(); + ggml_graph_update_use_counts(g); } static ggml_tensor * map_tensor(std::map & tensor_map, ggml_context * ctx, ggml_tensor * tensor) { @@ -502,6 +525,18 @@ static ggml_cgraph * dup_graph(ggml_context * ctx, ggml_cgraph * src) { return dst; } +// the outputs' rows: every slice's when sliced +static int64_t ggml_opt_output_rows(ggml_opt_context_t opt_ctx) { + if (opt_ctx->output_slices.empty()) { + return opt_ctx->outputs->ne[1]; + } + int64_t n = 0; + for (const ggml_tensor * out : opt_ctx->output_slices) { + n += out->ne[1]; + } + return n; +} + static void ggml_opt_build(ggml_opt_context_t opt_ctx) { GGML_ASSERT(opt_ctx->ctx_compute && "no compute context set, either use static graphs or set one with ggml_opt_prepare_alloc"); GGML_ASSERT((!opt_ctx->static_graphs || opt_ctx->inputs->data) && "when using static graphs the inputs must be allocated statically"); @@ -515,7 +550,9 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { opt_ctx->optimizer == GGML_OPT_OPTIMIZER_TYPE_ADAMW; ggml_set_input(opt_ctx->inputs); - ggml_set_output(opt_ctx->outputs); + if (opt_ctx->output_slices.empty()) { + ggml_set_output(opt_ctx->outputs); + } int n_param = 0; for (int i = 0; i < opt_ctx->gf->n_nodes; ++i) { @@ -582,11 +619,43 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { opt_ctx->loss_per_datapoint = false; break; } - case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY: { - opt_ctx->labels = ggml_dup_tensor(ctx_results, opt_ctx->outputs); + case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY: + case GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS: { + opt_ctx->labels = opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS + ? ggml_new_tensor_1d(ctx_results, GGML_TYPE_I32, ggml_opt_output_rows(opt_ctx)) + : ggml_dup_tensor(ctx_results, opt_ctx->outputs); ggml_set_input(opt_ctx->labels); ggml_set_name(opt_ctx->labels, "labels"); - opt_ctx->loss = ggml_cross_entropy_loss(ctx_results, opt_ctx->outputs, opt_ctx->labels); + opt_ctx->label_weights = nullptr; + if (opt_ctx->next_label_weights) { + GGML_ASSERT(opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS && "a weight per row needs one target index per row"); + opt_ctx->label_weights = ggml_new_tensor_1d(ctx_results, GGML_TYPE_F32, opt_ctx->labels->ne[0]); + ggml_set_input(opt_ctx->label_weights); + ggml_set_name(opt_ctx->label_weights, "label_weights"); + } + auto row_loss = [&](ggml_tensor * out, ggml_tensor * lab, int64_t r0) { + if (opt_ctx->label_weights == nullptr) { + return ggml_cross_entropy_loss(ctx_results, out, lab); + } + ggml_tensor * w = ggml_view_1d(ctx_results, opt_ctx->label_weights, out->ne[1], r0*sizeof(float)); + return ggml_cross_entropy_loss_weighted(ctx_results, out, lab, w); + }; + if (opt_ctx->output_slices.empty()) { + opt_ctx->loss = row_loss(opt_ctx->outputs, opt_ctx->labels, 0); + } else { + // the mean over every row: each slice's mean weighted by its share of the rows + const int64_t n = opt_ctx->labels->ne[0]; + int64_t r0 = 0; + ggml_tensor * sum = nullptr; + for (ggml_tensor * out : opt_ctx->output_slices) { + const int64_t nr = out->ne[1]; + ggml_tensor * lab = ggml_view_1d(ctx_results, opt_ctx->labels, nr, r0*sizeof(int32_t)); + ggml_tensor * l = ggml_scale(ctx_results, row_loss(out, lab, r0), (float) nr / (float) n); + sum = sum ? ggml_add(ctx_results, sum, l) : l; + r0 += nr; + } + opt_ctx->loss = sum; + } ggml_set_name(opt_ctx->loss, "loss_cross_entropy"); if (opt_ctx->next_manual) { // weighted by the caller below, beside its extra term @@ -617,7 +686,7 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { ggml_set_output(opt_ctx->loss); if (opt_ctx->next_manual) { GGML_ASSERT(!opt_ctx->static_graphs && "a caller-driven period needs graphs built per step"); - GGML_ASSERT(opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY); + GGML_ASSERT(opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY || opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS); // the loss the backward differentiates: the caller's weight of this graph's loss, plus // its extra term; opt_ctx->loss stays the unweighted loss the result reports struct ggml_tensor * total = ggml_scale(ctx_results, opt_ctx->loss, opt_ctx->next_loss_scale); @@ -634,13 +703,19 @@ static void ggml_opt_build(ggml_opt_context_t opt_ctx) { ggml_build_forward_expand(opt_ctx->gf, opt_ctx->loss); } - if (opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY) { + // per-step graphs rebuild these: one left from the previous graph would be read after it is freed + opt_ctx->pred = nullptr; + opt_ctx->ncorrect = nullptr; + if ((opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY || opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS) + && opt_ctx->output_slices.empty()) { opt_ctx->pred = ggml_argmax(ctx_results, opt_ctx->outputs); ggml_set_name(opt_ctx->pred, "pred"); ggml_set_output(opt_ctx->pred); ggml_build_forward_expand(opt_ctx->gf, opt_ctx->pred); - opt_ctx->ncorrect = ggml_count_equal(ctx_results, opt_ctx->pred, ggml_argmax(ctx_results, opt_ctx->labels)); + ggml_tensor * truth = opt_ctx->loss_type == GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS + ? opt_ctx->labels : ggml_argmax(ctx_results, opt_ctx->labels); + opt_ctx->ncorrect = ggml_count_equal(ctx_results, opt_ctx->pred, truth); ggml_set_name(opt_ctx->ncorrect, "ncorrect"); ggml_set_output(opt_ctx->ncorrect); ggml_build_forward_expand(opt_ctx->gf, opt_ctx->ncorrect); @@ -866,6 +941,15 @@ struct ggml_tensor * ggml_opt_labels(ggml_opt_context_t opt_ctx) { return opt_ctx->labels; } +struct ggml_tensor * ggml_opt_label_weights(ggml_opt_context_t opt_ctx) { + return opt_ctx->label_weights; +} + +void ggml_opt_set_next_label_weights(ggml_opt_context_t opt_ctx, bool weighted) { + GGML_ASSERT(!opt_ctx->eval_ready && "set before ggml_opt_alloc"); + opt_ctx->next_label_weights = weighted; +} + struct ggml_tensor * ggml_opt_loss(ggml_opt_context_t opt_ctx) { return opt_ctx->loss; } @@ -976,8 +1060,169 @@ void ggml_opt_prepare_alloc( opt_ctx->inputs = inputs; opt_ctx->outputs = outputs; opt_ctx->next_grad_accs.clear(); // registrations belong to the graph prepared last + opt_ctx->output_slices.clear(); +} + +int ggml_opt_project_head( + struct ggml_context * ctx_compute, + struct ggml_cgraph * gf, + struct ggml_tensor * hidden, + struct ggml_tensor * outputs, + struct ggml_tensor * shadow_rows, + struct ggml_tensor * public_rows, + struct ggml_tensor * label_rows, + struct ggml_tensor * public_shadow_rows, + struct ggml_tensor ** slice_shadow_rows, + int64_t rows, + struct ggml_tensor ** public_outputs, + struct ggml_tensor ** slices, + size_t slices_capacity) { + const int64_t n = label_rows ? label_rows->ne[0] : hidden->ne[1]; + if (rows <= 0 || n <= 0 || !ggml_is_contiguous(hidden) || ggml_n_dims(hidden) > 2 || + outputs->ne[1] != hidden->ne[1] || slices_capacity < (size_t) ((n + rows - 1) / rows)) { + return 0; + } + // the head: the nodes between hidden and outputs, in graph order + std::unordered_set on_path = { hidden }; + std::vector head; + for (int i = 0; i < gf->n_nodes; ++i) { + ggml_tensor * node = gf->nodes[i]; + bool reads_path = false; + for (int k = 0; k < GGML_MAX_SRC; ++k) { + reads_path |= node->src[k] != nullptr && on_path.count(node->src[k]) > 0; + } + if (!reads_path) { + continue; + } + const bool replayable = node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL || + node->op == GGML_OP_ADD || node->op == GGML_OP_SCALE || + (shadow_rows && (node->op == GGML_OP_DUP || node->op == GGML_OP_GET_ROWS || node->op == GGML_OP_SET_ROWS)); + if (!replayable || (node->view_src != nullptr && node->op != GGML_OP_SET_ROWS)) { + return 0; // unsupported head + } + if ((node->op == GGML_OP_GET_ROWS && node->src[1] != shadow_rows) || + (node->op == GGML_OP_SET_ROWS && node->src[1] != shadow_rows)) { + return 0; + } + if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { + for (int k = 0; k < 2; ++k) { + if (!on_path.count(node->src[k]) && node->src[k]->ne[1] != 1) { + return 0; + } + } + } + on_path.insert(node); + head.push_back(node); + } + if (head.empty() || head.back() != outputs) { + return 0; + } + // nothing else in the graph may read the head (it is about to leave the graph) + for (int i = 0; i < gf->n_nodes; ++i) { + const ggml_tensor * node = gf->nodes[i]; + if (on_path.count(node) > 0) { + continue; + } + for (int k = 0; k < GGML_MAX_SRC; ++k) { + if (node->src[k] != nullptr && node->src[k] != hidden && on_path.count(node->src[k]) > 0) { + return 0; + } + } + } + int j = 0; + for (int i = 0; i < gf->n_nodes; ++i) { + if (gf->nodes[i] == hidden || on_path.count(gf->nodes[i]) == 0) { + gf->nodes[j++] = gf->nodes[i]; + } + } + gf->n_nodes = j; + auto replay = [&](ggml_tensor * selected, ggml_tensor * local_shadow_rows, int64_t row) { + std::unordered_map as_slice; + as_slice[hidden] = selected; + if (shadow_rows) { + as_slice[shadow_rows] = local_shadow_rows; + } + for (ggml_tensor * node : head) { + auto src = [&](int k) { const auto it = as_slice.find(node->src[k]); return it != as_slice.end() ? it->second : node->src[k]; }; + ggml_tensor * t = nullptr; + if (node->op == GGML_OP_SET_ROWS && shadow_rows && !local_shadow_rows) { + as_slice[node] = src(2); + continue; + } + bool missing = false; + for (int k = 0; k < GGML_MAX_SRC; ++k) { + missing |= node->src[k] != nullptr && src(k) == nullptr; + } + if (missing) { + as_slice[node] = nullptr; + continue; + } + switch (node->op) { + case GGML_OP_DUP: t = ggml_dup(ctx_compute, src(0)); break; + case GGML_OP_GET_ROWS: t = ggml_get_rows(ctx_compute, src(0), src(1)); break; + case GGML_OP_SET_ROWS: t = ggml_set_rows(ctx_compute, src(2), src(0), src(1)); break; + case GGML_OP_MUL_MAT: t = ggml_mul_mat(ctx_compute, src(0), src(1)); break; + case GGML_OP_MUL: t = ggml_mul (ctx_compute, src(0), src(1)); break; + case GGML_OP_ADD: t = ggml_add (ctx_compute, src(0), src(1)); break; + case GGML_OP_SCALE: { + float s, b; + memcpy(&s, (const float *) node->op_params + 0, sizeof(float)); + memcpy(&b, (const float *) node->op_params + 1, sizeof(float)); + t = ggml_scale_bias(ctx_compute, src(0), s, b); + } break; + default: GGML_ABORT("not a head op"); + } + ggml_set_op_params(t, node->op_params, sizeof(node->op_params)); + ggml_format_name(t, "%s-rows%lld", node->name, (long long) row); + as_slice[node] = t; + } + return as_slice.at(outputs); + }; + *public_outputs = nullptr; + if (public_rows) { + *public_outputs = replay(ggml_get_rows(ctx_compute, hidden, public_rows), + public_shadow_rows, -1); + ggml_build_forward_expand(gf, *public_outputs); + } + ggml_tensor * selected = label_rows ? ggml_get_rows(ctx_compute, hidden, label_rows) : hidden; + size_t count = 0; + for (int64_t r0 = 0; r0 < n; r0 += rows) { + const int64_t nr = std::min(rows, n - r0); + slices[count] = replay( + ggml_view_2d(ctx_compute, selected, selected->ne[0], nr, selected->nb[1], r0 * selected->nb[1]), + slice_shadow_rows ? slice_shadow_rows[count] : nullptr, r0); + ++count; + } + ggml_graph_update_use_counts(gf); + return (int) count; } +void ggml_opt_prepare_alloc_slices( + ggml_opt_context_t opt_ctx, ggml_context * ctx_compute, ggml_cgraph * gf, + ggml_tensor * inputs, ggml_tensor ** slices, size_t n_slices) { + GGML_ASSERT(n_slices > 0); + ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, slices[0]); + opt_ctx->output_slices.assign(slices, slices + n_slices); +} + +int ggml_opt_prepare_alloc_head( + ggml_opt_context_t opt_ctx, ggml_context * ctx_compute, ggml_cgraph * gf, + ggml_tensor * inputs, ggml_tensor * hidden, ggml_tensor * outputs, int64_t rows) { + ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, outputs); + const int64_t n = hidden->ne[1]; + if (opt_ctx->loss_type != GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS || rows <= 0 || rows >= n) { + return 1; + } + std::vector slices((n + rows - 1) / rows); + ggml_tensor * public_outputs = nullptr; + const int count = ggml_opt_project_head(ctx_compute, gf, hidden, outputs, nullptr, nullptr, nullptr, nullptr, nullptr, + rows, &public_outputs, slices.data(), slices.size()); + if (count == 0) { + return 1; + } + ggml_opt_prepare_alloc_slices(opt_ctx, ctx_compute, gf, inputs, slices.data(), count); + return count; +} void ggml_opt_set_alloc_budget(ggml_opt_context_t opt_ctx, size_t bytes) { opt_ctx->alloc_budget = bytes; } @@ -1061,8 +1306,7 @@ bool ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { if (backward) { ggml_backend_sched_t sched = opt_ctx->backend_sched; ggml_backend_t device = ggml_backend_sched_get_backend(sched, 0); - ggml_backend_dev_t dev = device ? ggml_backend_get_device(device) : nullptr; - if (dev != nullptr && ggml_backend_dev_type(dev) != GGML_BACKEND_DEVICE_TYPE_CPU) { + if (device != nullptr) { ggml_cgraph * g = opt_ctx->allocated_graph_copy; for (int i = 0; i < ggml_graph_n_nodes(g); ++i) { ggml_tensor * node = ggml_graph_node(g, i); @@ -1150,12 +1394,51 @@ bool ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) { } opt_ctx->allocated_graph = graph; + // GGML_OPT_TOP=N: the graph's N largest tensors, nodes and the leaves it allocates (inputs, not + // views or weights), where its memory goes + if (const char * top = getenv("GGML_OPT_TOP")) { + ggml_cgraph * g = opt_ctx->allocated_graph_copy; + std::vector> rows; + for (int i = 0; i < g->n_nodes; ++i) { + if (g->nodes[i]->view_src == nullptr) { + rows.emplace_back(ggml_nbytes(g->nodes[i]), g->nodes[i]); + } + } + for (int i = 0; i < g->n_leafs; ++i) { + const ggml_tensor * l = g->leafs[i]; + if (l->view_src == nullptr && (l->flags & GGML_TENSOR_FLAG_INPUT)) { + rows.emplace_back(ggml_nbytes(l), l); + } + } + std::sort(rows.begin(), rows.end(), [](const auto & x, const auto & y) { return x.first > y.first; }); + size_t sum = 0; + for (const auto & r : rows) { + sum += r.first; + } + GGML_LOG_INFO("%s: graph of %d nodes, %d leafs: its tensors sum to %.1f MiB, the peak so far %.1f MiB\n", __func__, + g->n_nodes, g->n_leafs, sum/1048576.0, opt_ctx->peak_graph_bytes/1048576.0); + const size_t n = std::min(rows.size(), (size_t) atoi(top)); + for (size_t i = 0; i < n; ++i) { + const ggml_tensor * t = rows[i].second; + GGML_LOG_INFO("%s: top %2zu: %8.1f MiB %-12s %-6s [%lld,%lld,%lld,%lld] %s <- %s %s [%lld,%lld] | %s %s [%lld,%lld]\n", __func__, i, + rows[i].first/1048576.0, ggml_op_desc(t), ggml_type_name(t->type), (long long) t->ne[0], (long long) t->ne[1], + (long long) t->ne[2], (long long) t->ne[3], t->name, + t->src[0] ? ggml_op_desc(t->src[0]) : "-", t->src[0] ? t->src[0]->name : "", t->src[0] ? (long long) t->src[0]->ne[0] : 0LL, + t->src[0] ? (long long) t->src[0]->ne[1] : 0LL, + t->src[1] ? ggml_op_desc(t->src[1]) : "-", t->src[1] ? t->src[1]->name : "", t->src[1] ? (long long) t->src[1]->ne[0] : 0LL, + t->src[1] ? (long long) t->src[1]->ne[1] : 0LL); + } + } + opt_ctx->eval_ready = true; return true; } void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result) { GGML_ASSERT(opt_ctx->eval_ready); + // a caller-driven step's loss is not scaled by 1/opt_period (ggml_opt_build): recorded as if it + // were, so the result's reader (which multiplies a per-datapoint loss by opt_period) reports it right + const bool manual = opt_ctx->next_manual; if (opt_ctx->allocated_graph == opt_ctx->gb_opt) { const ggml_opt_optimizer_params & opt_pars = opt_ctx->get_opt_pars(opt_ctx->get_opt_pars_ud); @@ -1238,7 +1521,7 @@ void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result) { GGML_ASSERT(result->opt_period == opt_ctx->opt_period); } - const int64_t ndata = opt_ctx->outputs->ne[1]; + const int64_t ndata = ggml_opt_output_rows(opt_ctx); result->ndata += ndata; result->ndata_batch.push_back(ndata); @@ -1246,7 +1529,7 @@ void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result) { GGML_ASSERT(opt_ctx->loss->type == GGML_TYPE_F32); float loss; ggml_backend_tensor_get(opt_ctx->loss, &loss, 0, ggml_nbytes(opt_ctx->loss)); - result->loss.push_back(loss); + result->loss.push_back(manual && opt_ctx->loss_per_datapoint ? loss / opt_ctx->opt_period : loss); if (opt_ctx->pred) { GGML_ASSERT(opt_ctx->pred->type == GGML_TYPE_I32); diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index 0573643d834e..4b00fdace548 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -6336,9 +6336,11 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_FILL: case GGML_OP_CUMSUM: case GGML_OP_DIAG: + return true; case GGML_OP_CROSS_ENTROPY_LOSS: + return op->src[1]->type == GGML_TYPE_F32; // no I32 targets on this backend yet case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - return true; + return op->src[2]->type == GGML_TYPE_F32; case GGML_OP_SOLVE_TRI: return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K; case GGML_OP_FLASH_ATTN_EXT: diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 5f572229d5f9..94bbb9d40fbc 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -5524,6 +5524,26 @@ void ggml_flash_attn_ext_set_back_tile( ggml_set_op_params_i32(a, 4, tile); // scale, max_bias, softcap, prec come first } +void ggml_flash_attn_ext_set_grad_sinks( + struct ggml_tensor * a, + struct ggml_tensor * k_tap, + struct ggml_tensor * k_cur, + struct ggml_tensor * v_tap, + struct ggml_tensor * v_cur) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT && k_cur && v_cur && (k_tap == NULL) == (v_tap == NULL)); + const struct ggml_tensor * k = a->src[1]; + const struct ggml_tensor * v = a->src[2]; + const int64_t n = k_cur->ne[2] + (k_tap ? k_tap->ne[2] : 0); + GGML_ASSERT(v_cur->ne[2] + (v_tap ? v_tap->ne[2] : 0) == n && n <= k->ne[1]); + GGML_ASSERT(k_cur->ne[0] == k->ne[0] && k_cur->ne[1] == k->ne[2] && v_cur->ne[0] == v->ne[0] && v_cur->ne[1] == v->ne[2]); + GGML_ASSERT(!k_tap || (k_tap->ne[0] == k->ne[0] && k_tap->ne[1] == k->ne[2] && v_tap->ne[0] == v->ne[0] && v_tap->ne[1] == v->ne[2])); + // src[4] is sinks (the attention's own); the gradient sinks follow it + a->src[5] = k_tap; + a->src[6] = k_cur; + a->src[7] = v_tap; + a->src[8] = v_cur; +} + void ggml_flash_attn_ext_add_sinks( struct ggml_tensor * a, struct ggml_tensor * sinks) { @@ -5542,71 +5562,67 @@ void ggml_flash_attn_ext_add_sinks( // ggml_flash_attn_back +float ggml_flash_attn_back_scale(float max_abs_do) { + if (!(max_abs_do > 0.0f)) { + return 1.0f; + } + int e; + frexpf(max_abs_do, &e); // max_abs_do in [2^(e-1), 2^e) + const int k = 3 - e; // scaled into [2^2, 2^3) + return ldexpf(1.0f, k < -24 ? -24 : (k > 40 ? 40 : k)); +} + +size_t ggml_flash_attn_back_offset(const struct ggml_tensor * back, int part) { + const struct ggml_tensor * q = back->src[0]; + const struct ggml_tensor * k = back->src[1]; + const struct ggml_tensor * v = back->src[2]; + const int64_t kv0 = ggml_get_op_params_i32(back, 1); // dK/dV hold the key rows [kv0, M) + const size_t sizes[3] = { + GGML_PAD((size_t) q->ne[0]*q->ne[1]*q->ne[2]*sizeof(float), GGML_MEM_ALIGN), + GGML_PAD((size_t) k->ne[0]*(k->ne[1] - kv0)*k->ne[2]*sizeof(float), GGML_MEM_ALIGN), + GGML_PAD((size_t) v->ne[0]*(v->ne[1] - kv0)*v->ne[2]*sizeof(float), GGML_MEM_ALIGN), + }; + size_t offs = 0; + for (int i = 0; i < part; ++i) { + offs += sizes[i]; + } + return offs; +} + struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, struct ggml_tensor * q, struct ggml_tensor * k, struct ggml_tensor * v, + struct ggml_tensor * mask, struct ggml_tensor * d, - bool masked) { - GGML_ABORT("TODO: adapt to ggml_flash_attn_ext() changes"); - - GGML_ASSERT(ggml_can_mul_mat(k, q)); - // TODO: check if vT can be multiplied by (k*qT) - - // d shape [D,N,ne2,ne3] - // q shape [D,N,ne2,ne3] - // k shape [D,M,kvne2,ne3] - // v shape [M,D,kvne2,ne3] - - const int64_t D = q->ne[0]; - const int64_t N = q->ne[1]; - const int64_t M = k->ne[1]; - const int64_t ne2 = q->ne[2]; - const int64_t ne3 = q->ne[3]; - const int64_t kvne2 = k->ne[2]; - - GGML_ASSERT(k->ne[0] == D); - GGML_ASSERT(v->ne[0] == M); - GGML_ASSERT(v->ne[1] == D); - GGML_ASSERT(d->ne[0] == D); - GGML_ASSERT(d->ne[1] == N); - GGML_ASSERT(k->ne[2] == kvne2); - GGML_ASSERT(k->ne[3] == ne3); - GGML_ASSERT(v->ne[2] == kvne2); - GGML_ASSERT(v->ne[3] == ne3); - GGML_ASSERT(d->ne[2] == ne2); - GGML_ASSERT(d->ne[3] == ne3); - - GGML_ASSERT(ne2 % kvne2 == 0); - - // store gradients of q, k and v as continuous tensors concatenated in result. - // note: v and gradv are actually transposed, i.e. v->ne[0] != D. - const int64_t elem_q = ggml_nelements(q); - const int64_t elem_k = ggml_nelements(k); - const int64_t elem_v = ggml_nelements(v); - - enum ggml_type result_type = GGML_TYPE_F32; - GGML_ASSERT(ggml_blck_size(result_type) == 1); - const size_t tsize = ggml_type_size(result_type); - - const size_t offs_q = 0; - const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); - const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); - const size_t end = offs_v + GGML_PAD(elem_v * tsize, GGML_MEM_ALIGN); - - const size_t nelements = (end + tsize - 1)/tsize; - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nelements); - - int32_t masked_i = masked ? 1 : 0; - ggml_set_op_params(result, &masked_i, sizeof(masked_i)); + float scale, + int64_t kv0) { + // q [D, N, H, 1], k [D, M, Hkv, 1], v [Dv, M, Hkv, 1], d [Dv, H, N, 1], mask [M', N', 1|H, 1] + GGML_ASSERT(kv0 >= 0 && kv0 < k->ne[1] && kv0 <= INT32_MAX); + GGML_ASSERT(q->type == GGML_TYPE_F32 && d->type == GGML_TYPE_F32); + GGML_ASSERT(q->ne[3] == 1 && k->ne[3] == 1 && v->ne[3] == 1 && "one sequence stream"); + GGML_ASSERT(k->ne[0] == q->ne[0] && v->ne[1] == k->ne[1] && v->ne[2] == k->ne[2]); + GGML_ASSERT(q->ne[2] % k->ne[2] == 0); + GGML_ASSERT(d->ne[0] == v->ne[0] && d->ne[1] == q->ne[2] && d->ne[2] == q->ne[1]); + GGML_ASSERT(!mask || (mask->type == GGML_TYPE_F16 && mask->ne[0] >= k->ne[1] && mask->ne[1] >= q->ne[1])); + + const size_t rows = (size_t) q->ne[1]*q->ne[2]; + struct ggml_tensor tmp = { 0 }; + tmp.src[0] = q; tmp.src[1] = k; tmp.src[2] = v; + ggml_set_op_params_i32(&tmp, 1, (int32_t) kv0); + const size_t bytes = ggml_flash_attn_back_offset(&tmp, 3) + (2*rows + 1)*sizeof(float); // lse, D, the dO scale + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) (bytes/sizeof(float))); + ggml_set_op_params_f32(result, 0, scale); + ggml_set_op_params_i32(result, 1, (int32_t) kv0); result->op = GGML_OP_FLASH_ATTN_BACK; result->src[0] = q; result->src[1] = k; result->src[2] = v; - result->src[3] = d; + result->src[3] = mask; + result->src[4] = d; return result; } @@ -6198,11 +6214,18 @@ struct ggml_tensor * ggml_custom_inplace( } // ggml_cross_entropy_loss +bool ggml_cross_entropy_labels_ok(const struct ggml_tensor * logits, const struct ggml_tensor * labels) { + if (labels->type == GGML_TYPE_I32) { + return ggml_is_contiguous(labels) && ggml_nelements(labels) == ggml_nrows(logits); + } + return ggml_are_same_shape(logits, labels); +} + struct ggml_tensor * ggml_cross_entropy_loss( struct ggml_context * ctx, struct ggml_tensor * a, struct ggml_tensor * b) { - GGML_ASSERT(ggml_are_same_shape(a, b)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(a, b)); struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); @@ -6213,6 +6236,18 @@ struct ggml_tensor * ggml_cross_entropy_loss( return result; } +struct ggml_tensor * ggml_cross_entropy_loss_weighted( + struct ggml_context * ctx, + struct ggml_tensor * logits, + struct ggml_tensor * targets, + struct ggml_tensor * weights) { + GGML_ASSERT(targets->type == GGML_TYPE_I32 && ggml_cross_entropy_labels_ok(logits, targets)); + GGML_ASSERT(weights->type == GGML_TYPE_F32 && ggml_is_contiguous(weights) && ggml_nelements(weights) == ggml_nrows(logits)); + struct ggml_tensor * result = ggml_cross_entropy_loss(ctx, logits, targets); + result->src[2] = weights; + return result; +} + // ggml_cross_entropy_loss_back struct ggml_tensor * ggml_cross_entropy_loss_back( @@ -6221,7 +6256,7 @@ struct ggml_tensor * ggml_cross_entropy_loss_back( struct ggml_tensor * b, struct ggml_tensor * c) { GGML_ASSERT(ggml_is_scalar(a)); - GGML_ASSERT(ggml_are_same_shape(b, c)); + GGML_ASSERT(ggml_cross_entropy_labels_ok(b, c)); struct ggml_tensor * result = ggml_dup_tensor(ctx, b); @@ -6679,7 +6714,11 @@ static void ggml_acc_or_set( } cgraph->grads[isrc] = ggml_acc_impl(ctx, acc_into, tensor, nb1, nb2, nb3, offset, cgraph->grad_accs[isrc]); } else { - struct ggml_tensor * a_zero = ggml_scale(ctx, src, 0.0f); // FIXME this is going to produce NaN if a contains inf/NaN + // zeros by fill (0 * inf would be NaN), built from a one-element leaf repeated to src's + // shape: a full-size leaf stays allocated for the whole graph (+813 MiB on the 4B at + // window 8192, 2026-10-14), a node is freed after its use + struct ggml_tensor * a_zero = ggml_repeat_4d(ctx, ggml_fill(ctx, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1), 0.0f), + src->ne[0], src->ne[1], src->ne[2], src->ne[3]); cgraph->grads[isrc] = ggml_acc_impl(ctx, a_zero, tensor, nb1, nb2, nb3, offset, false); } ggml_format_name(cgraph->grads[isrc], "grad for %s", cgraph->visited_hash_set.keys[isrc]->name); @@ -6801,14 +6840,61 @@ static struct ggml_tensor * ggml_fa_back_join(struct ggml_context * ctx, struct return ggml_concat(ctx, ggml_fa_back_join(ctx, parts, h), ggml_fa_back_join(ctx, parts + h, n - h), 1); } +// whether src needs a gradient in this backward pass +static bool ggml_backward_needs_grad(struct ggml_cgraph * cgraph, const bool * grads_needed, const struct ggml_tensor * src, size_t * isrc) { + if (!src) { + return false; + } + *isrc = ggml_hash_find(&cgraph->visited_hash_set, src); + return *isrc != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, *isrc) && grads_needed[*isrc]; +} + +// FLASH_ATTN_EXT's K (part 0) or V (part 1) gradient g [D, M - kv0, Hkv] (key rows [kv0, M)): to the +// grad sinks when the attention has them (ggml_flash_attn_ext_set_grad_sinks), else to K/V itself +static void ggml_fa_back_route_kv(struct ggml_context * ctx, struct ggml_cgraph * cgraph, const bool * grads_needed, + struct ggml_tensor * tensor, int part, struct ggml_tensor * g, size_t isrc, bool needed) { + if (tensor->src[6] == NULL) { + if (needed) { + const struct ggml_tensor * kv = tensor->src[1 + part]; + ggml_add_or_set(ctx, cgraph, isrc, ggml_reshape_4d(ctx, ggml_cont(ctx, g), kv->ne[0], kv->ne[1], kv->ne[2], kv->ne[3])); + } + return; + } + int64_t r = 0; + for (int s = 0; s < 2; ++s) { + struct ggml_tensor * sink = tensor->src[5 + 2*part + s]; // tap, then the chunk's own + size_t is = 0; + if (sink == NULL) { + continue; + } + const int64_t n = sink->ne[2]; + if (ggml_backward_needs_grad(cgraph, grads_needed, sink, &is)) { + struct ggml_tensor * rows = ggml_view_3d(ctx, g, g->ne[0], n, g->ne[2], g->nb[1], g->nb[2], r*g->nb[1]); + ggml_add_or_set(ctx, cgraph, is, ggml_cont(ctx, ggml_permute(ctx, rows, 0, 2, 1, 3))); // [D, Hkv, n] + } + r += n; + } +} + static void ggml_flash_attn_ext_backward( struct ggml_context * ctx, struct ggml_cgraph * cgraph, struct ggml_tensor * tensor, struct ggml_tensor * grad, - size_t isrc0, size_t isrc1, size_t isrc2, bool dq_needed, bool dk_needed, bool dv_needed) { + size_t isrc0, size_t isrc1, size_t isrc2, bool dq_needed, bool dk_needed, bool dv_needed, const bool * grads_needed) { struct ggml_tensor * q = tensor->src[0]; // [D, N, H, 1] struct ggml_tensor * k = tensor->src[1]; // [D, M, Hkv, 1] struct ggml_tensor * v = tensor->src[2]; // [Dv, M, Hkv, 1] struct ggml_tensor * mask = tensor->src[3]; // [M', N', 1|H, 1] F16, or NULL + // the native backward reads a quantized K/V itself: one the forward read through an F16 cast of + // a quantized cache is read at the cast's source, so the window-sized cast is not kept alive + // for the backward + const bool native = ggml_get_op_params_i32(tensor, 4) == GGML_FA_BACK_NATIVE; + if (native && k->op == GGML_OP_CPY && ggml_is_quantized(k->src[0]->type) && ggml_are_same_shape(k, k->src[0])) { + k = k->src[0]; + } + if (native && v->op == GGML_OP_CPY && ggml_is_quantized(v->src[0]->type) && ggml_are_same_shape(v, v->src[0])) { + v = v->src[0]; + } + float scale, max_bias, softcap; memcpy(&scale, (const float *) tensor->op_params + 0, sizeof(float)); memcpy(&max_bias, (const float *) tensor->op_params + 1, sizeof(float)); @@ -6818,7 +6904,33 @@ static void ggml_flash_attn_ext_backward( GGML_ASSERT(tensor->src[4] == NULL && "flash_attn_ext backward: sinks are not implemented"); GGML_ASSERT(q->ne[3] == 1 && "flash_attn_ext backward: one sequence stream"); + // with grad sinks only K/V's last rows carry a gradient: from kv0 on, and nothing to K/V itself + const bool sunk = tensor->src[6] != NULL; + const int64_t kv0 = sunk ? k->ne[1] - tensor->src[6]->ne[2] - (tensor->src[5] ? tensor->src[5]->ne[2] : 0) : 0; + if (sunk) { + size_t is = 0; + dk_needed = ggml_backward_needs_grad(cgraph, grads_needed, tensor->src[5], &is) || ggml_backward_needs_grad(cgraph, grads_needed, tensor->src[6], &is); + dv_needed = ggml_backward_needs_grad(cgraph, grads_needed, tensor->src[7], &is) || ggml_backward_needs_grad(cgraph, grads_needed, tensor->src[8], &is); + } + int32_t tile = ggml_get_op_params_i32(tensor, 4); + if (tile == GGML_FA_BACK_NATIVE) { + struct ggml_tensor * b = ggml_flash_attn_back(ctx, q, k, v, mask, grad, scale, kv0); + // views of the packed result: dQ shaped like q, dK/dV like k/v from row kv0 + const int64_t L = k->ne[1] - kv0; + struct ggml_tensor * dq = ggml_view_3d(ctx, b, q->ne[0], q->ne[1], q->ne[2], q->ne[0]*sizeof(float), q->ne[0]*q->ne[1]*sizeof(float), + ggml_flash_attn_back_offset(b, 0)); + struct ggml_tensor * dk = ggml_view_3d(ctx, b, k->ne[0], L, k->ne[2], k->ne[0]*sizeof(float), k->ne[0]*L*sizeof(float), + ggml_flash_attn_back_offset(b, 1)); + struct ggml_tensor * dv = ggml_view_3d(ctx, b, v->ne[0], L, v->ne[2], v->ne[0]*sizeof(float), v->ne[0]*L*sizeof(float), + ggml_flash_attn_back_offset(b, 2)); + if (dq_needed) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape_4d(ctx, ggml_cont(ctx, dq), q->ne[0], q->ne[1], q->ne[2], q->ne[3])); + } + ggml_fa_back_route_kv(ctx, cgraph, grads_needed, tensor, 0, dk, isrc1, dk_needed); + ggml_fa_back_route_kv(ctx, cgraph, grads_needed, tensor, 1, dv, isrc2, dv_needed); + return; + } if (tile <= 0) { tile = GGML_FA_BACK_TILE_DEFAULT; } @@ -6866,9 +6978,13 @@ static void ggml_flash_attn_ext_backward( struct ggml_tensor * dq = NULL; struct ggml_tensor ** dk = (struct ggml_tensor **) alloca(sizeof(struct ggml_tensor *) * nt); struct ggml_tensor ** dv = (struct ggml_tensor **) alloca(sizeof(struct ggml_tensor *) * nt); + const int t0 = (int) (kv0 / tile); // the first tile holding a key row that carries a gradient for (int t = 0; t < nt; ++t) { const int64_t j0 = (int64_t) t * tile; const int64_t n = MIN(tile, M - j0); + if (t < t0 && !dq_needed) { + continue; + } struct ggml_tensor * kt = ggml_fa_back_f32(ctx, ggml_view_3d(ctx, k, k->ne[0], n, Hkv, k->nb[1], k->nb[2], j0*k->nb[1])); struct ggml_tensor * vt = ggml_fa_back_f32(ctx, ggml_view_3d(ctx, v, v->ne[0], n, Hkv, v->nb[1], v->nb[2], j0*v->nb[1])); struct ggml_tensor * p = ggml_exp(ctx, ggml_sub(ctx, ggml_fa_back_scores(ctx, q, k, mask, scale, j0, n), lse)); // [n, N, H] @@ -6879,11 +6995,11 @@ static void ggml_flash_attn_ext_backward( struct ggml_tensor * g = ggml_fa_back_mm(ctx, ggml_cont(ctx, ggml_transpose(ctx, kt)), ds); // [D, N, H] dq = dq ? ggml_add(ctx, dq, g) : g; } - if (dk_needed) { + if (dk_needed && t >= t0) { // dk[d,j] = sum_i q[d,i] ds[j,i] dk[t] = ggml_fa_back_group_sum(ctx, ggml_fa_back_mm(ctx, qT, ggml_cont(ctx, ggml_transpose(ctx, ds))), Hkv); // [D, n, Hkv] } - if (dv_needed) { + if (dv_needed && t >= t0) { // dv[d,j] = sum_i dO[d,i] p[j,i] dv[t] = ggml_fa_back_group_sum(ctx, ggml_fa_back_mm(ctx, dOT, ggml_cont(ctx, ggml_transpose(ctx, p))), Hkv); // [Dv, n, Hkv] } @@ -6892,11 +7008,93 @@ static void ggml_flash_attn_ext_backward( if (dq_needed) { ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape_4d(ctx, ggml_cont(ctx, dq), q->ne[0], q->ne[1], q->ne[2], q->ne[3])); } - if (dk_needed) { - ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape_4d(ctx, ggml_fa_back_join(ctx, dk, nt), k->ne[0], k->ne[1], k->ne[2], k->ne[3])); + // the tiles from t0 joined, then their rows from kv0 ([D, M - kv0, Hkv]) + for (int part = 0; part < 2; ++part) { + if (!(part == 0 ? dk_needed : dv_needed)) { + continue; + } + struct ggml_tensor * j = ggml_fa_back_join(ctx, (part == 0 ? dk : dv) + t0, nt - t0); + struct ggml_tensor * g = ggml_view_3d(ctx, j, j->ne[0], M - kv0, j->ne[2], j->nb[1], j->nb[2], (kv0 - (int64_t) t0*tile)*j->nb[1]); + ggml_fa_back_route_kv(ctx, cgraph, grads_needed, tensor, part, g, part == 0 ? isrc1 : isrc2, true); } - 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])); +} + +// Whether a backward can be built through this node: its op (and unary or glu op) has a case in +// ggml_compute_backward, and it is not an in-place write other than the re-shaping views. The ONE +// answer a caller asks before building a backward graph, instead of meeting an abort there; both +// ggml_build_backward_expand and ggml_compute_backward assert it, so they cannot disagree silently. +bool ggml_node_has_backward(const struct ggml_tensor * tensor) { + if (tensor->view_src != NULL && tensor->op != GGML_OP_CPY && tensor->op != GGML_OP_VIEW && + tensor->op != GGML_OP_RESHAPE && tensor->op != GGML_OP_PERMUTE && tensor->op != GGML_OP_TRANSPOSE && + tensor->op != GGML_OP_SET_ROWS) { + return false; // in-place operations are not supported + } + switch (tensor->op) { + case GGML_OP_DUP: + case GGML_OP_ADD: + case GGML_OP_ADD1: + case GGML_OP_ACC: + case GGML_OP_SUB: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_LOG: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_SUM: + case GGML_OP_SUM_ROWS: + case GGML_OP_MEAN: + case GGML_OP_REPEAT: + case GGML_OP_REPEAT_BACK: + case GGML_OP_RMS_NORM: + case GGML_OP_MUL_MAT: + case GGML_OP_SCALE: + case GGML_OP_SET: + case GGML_OP_CPY: + case GGML_OP_CONT: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_SET_ROWS: + case GGML_OP_GET_ROWS: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_DIAG_MASK_ZERO: + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_IM2COL: + case GGML_OP_POOL_2D: + case GGML_OP_WIN_PART: + case GGML_OP_WIN_UNPART: + case GGML_OP_CONCAT: + case GGML_OP_SSM_CONV: + case GGML_OP_L2_NORM: + case GGML_OP_GATED_DELTA_NET: + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_NONE: + return true; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(tensor)) { + case GGML_UNARY_OP_ABS: + case GGML_UNARY_OP_SGN: + case GGML_UNARY_OP_NEG: + case GGML_UNARY_OP_STEP: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_EXP: + case GGML_UNARY_OP_SIGMOID: + case GGML_UNARY_OP_EXPM1: + case GGML_UNARY_OP_SOFTPLUS: + return true; + default: + return false; + } + case GGML_OP_GLU: + return ggml_get_glu_op(tensor) == GGML_GLU_OP_SWIGLU; + default: + return false; } } @@ -6920,6 +7118,9 @@ static void ggml_compute_backward( const bool src1_needs_grads = src1 && isrc1 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc1) && grads_needed[isrc1]; const bool src2_needs_grads = src2 && isrc2 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc2) && grads_needed[isrc2]; + if (!ggml_node_has_backward(tensor)) { + GGML_ABORT("%s: no backward implemented for op %s (%s)", __func__, ggml_op_name(tensor->op), ggml_op_desc(tensor)); + } switch (tensor->op) { case GGML_OP_DUP: { if (src0_needs_grads) { @@ -7173,10 +7374,10 @@ static void ggml_compute_backward( size_t nb2 = tensor->nb[2]; size_t nb3 = tensor->nb[3]; - if (cgraph->grads[isrc0] && src0->type != cgraph->grads[isrc0]->type) { - // gradient is typically F32, but src0 could be other type - size_t ng = ggml_element_size(cgraph->grads[isrc0]); - size_t n0 = ggml_element_size(src0); + if (src0->type != GGML_TYPE_F32) { + // The gradient is F32 even when the viewed source stores quantized blocks. + size_t ng = sizeof(float) * ggml_blck_size(src0->type); + size_t n0 = ggml_type_size(src0->type); GGML_ASSERT(offset % n0 == 0); GGML_ASSERT(nb1 % n0 == 0); GGML_ASSERT(nb2 % n0 == 0); @@ -7212,13 +7413,18 @@ static void ggml_compute_backward( } break; case GGML_OP_SET_ROWS: { // result = the destination (src2) with its rows src1 replaced by src0. The gradient - // on src0 is the result's gradient at those rows. The destination's gradient (the - // result's everywhere else) is not built: a written buffer (a cache) is a constant. - GGML_ASSERT(!src2_needs_grads && "set_rows: the destination it writes carries no gradient"); - if (src0_needs_grads) { + // on src0 is the result's gradient at those rows; the destination's is the result's + // everywhere else (a cache it writes is a constant and needs none). + if (src0_needs_grads || src2_needs_grads) { GGML_ASSERT(src1->type == GGML_TYPE_I32 && "set_rows backward gathers with I32 row indices"); + } + if (src0_needs_grads) { ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, ggml_get_rows(ctx, grad, src1), src0)); } + if (src2_needs_grads) { + struct ggml_tensor * replaced = ggml_fill(ctx, ggml_get_rows(ctx, grad, src1), 0.0f); + ggml_add_or_set(ctx, cgraph, isrc2, ggml_set_rows(ctx, ggml_dup(ctx, grad), replaced, src1)); + } } break; case GGML_OP_GET_ROWS: { if (src0_needs_grads) { @@ -7474,7 +7680,9 @@ static void ggml_compute_backward( } break; case GGML_OP_CROSS_ENTROPY_LOSS: { if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_cross_entropy_loss_back(ctx, grad, src0, src1)); + struct ggml_tensor * back = ggml_cross_entropy_loss_back(ctx, grad, src0, src1); + back->src[3] = src2; // the per-row weights, when the loss has them + ggml_add_or_set(ctx, cgraph, isrc0, back); } GGML_ASSERT(!src1_needs_grads && "backward pass for labels not implemented"); } break; @@ -7496,7 +7704,7 @@ static void ggml_compute_backward( } break; case GGML_OP_FLASH_ATTN_EXT: { ggml_flash_attn_ext_backward(ctx, cgraph, tensor, grad, isrc0, isrc1, isrc2, - src0_needs_grads, src1_needs_grads, src2_needs_grads); + src0_needs_grads, src1_needs_grads, src2_needs_grads, grads_needed); } break; case GGML_OP_NONE: { // noop @@ -7508,8 +7716,10 @@ static void ggml_compute_backward( } GGML_ASSERT(!src0_needs_grads || ggml_are_same_shape(src0, cgraph->grads[isrc0])); - GGML_ASSERT(!src1_needs_grads || ggml_are_same_shape(src1, cgraph->grads[isrc1])); - GGML_ASSERT(!src2_needs_grads || ggml_are_same_shape(src2, cgraph->grads[isrc2])); + // an attention with grad sinks puts its K/V gradient on the sinks, never on K/V themselves + const bool kv_sunk = tensor->op == GGML_OP_FLASH_ATTN_EXT && tensor->src[6] != NULL; + GGML_ASSERT(!src1_needs_grads || kv_sunk || ggml_are_same_shape(src1, cgraph->grads[isrc1])); + GGML_ASSERT(!src2_needs_grads || kv_sunk || ggml_are_same_shape(src2, cgraph->grads[isrc2])); } static size_t ggml_visit_parents_graph(struct ggml_cgraph * cgraph, struct ggml_tensor * node, bool compute) { @@ -7688,7 +7898,7 @@ void ggml_build_backward_expand( // A gradient is F32 whatever its tensor's type. A quantized source is admitted where // the backward only re-shapes the gradient or casts its tile to F32 (a training graph // that attends to a q8_0 cache in place, as serving does); every other op needs float. - const bool shapes_only = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + const bool shapes_only = node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_FLASH_ATTN_EXT; GGML_ASSERT(node->src[j]->type == GGML_TYPE_F32 || node->src[j]->type == GGML_TYPE_F16 || (shapes_only && ggml_is_quantized(node->src[j]->type))); @@ -7699,10 +7909,10 @@ void ggml_build_backward_expand( continue; } - // inplace operations are currently not supported - GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || - node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE || - node->op == GGML_OP_SET_ROWS); + if (!ggml_node_has_backward(node)) { + GGML_ABORT("%s: no backward through %s (%s): its op has none, or it is an in-place write", __func__, + ggml_op_name(node->op), ggml_op_desc(node)); + } const size_t ihash = ggml_hash_find(&cgraph->visited_hash_set, node); GGML_ASSERT(ihash != GGML_HASHSET_FULL); @@ -7929,6 +8139,19 @@ void ggml_graph_reset(struct ggml_cgraph * cgraph) { } } +void ggml_graph_update_use_counts(struct ggml_cgraph * cgraph) { + memset(cgraph->use_counts, 0, cgraph->visited_hash_set.size * sizeof(cgraph->use_counts[0])); + for (int i = 0; i < cgraph->n_nodes; ++i) { + for (int j = 0; j < GGML_MAX_SRC; ++j) { + struct ggml_tensor * src = cgraph->nodes[i]->src[j]; + if (src) { + const size_t pos = ggml_hash_find(&cgraph->visited_hash_set, src); + GGML_ASSERT(pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, pos)); + cgraph->use_counts[pos]++; + } + } + } +} void ggml_graph_clear(struct ggml_cgraph * cgraph) { cgraph->n_leafs = 0; cgraph->n_nodes = 0; diff --git a/include/llama.h b/include/llama.h index 3135ea509aa2..c322a486deaf 100644 --- a/include/llama.h +++ b/include/llama.h @@ -252,6 +252,14 @@ extern "C" { // - if embeddings: all tokens are output // - if not: only the last token is output // ) + // - label : the token the model should have predicted at this position, which makes the + // position a training row: llama_decode then also runs the loss, the backward and + // the optimizer write of the context's trainable adapter (llama_opt_init), in the + // same graph as the batch's other rows. -1 marks a position with no label. + // (if set to NULL, the batch is inference only, unchanged) + // - label_weight : how much a labelled position counts in the loss (>= 0, finite; 0 counts + // nothing): a mind's relevance on what it perceived, not only whether. Read for + // labelled positions only. (if set to NULL, every labelled position counts 1) // typedef struct llama_batch { int32_t n_tokens; @@ -262,6 +270,8 @@ extern "C" { int32_t * n_seq_id; llama_seq_id ** seq_id; int8_t * logits; // TODO: rename this to "output" + llama_token * label; + float * label_weight; } llama_batch; enum llama_model_kv_override_type { @@ -1775,6 +1785,26 @@ extern "C" { // window): the requested one, or smaller where a chunk's graph did not fit the device. LLAMA_API uint32_t llama_opt_walk_horizon(struct llama_context * lctx); + // LEARN ON A LIVE SEQUENCE (P2): train one reply on `seq` of a serving context, given the + // context before it exactly as served. The sequence must hold positions [0, n_cached) as served + // and nothing after them for a recurrent part (restore the served-prefix checkpoint first: + // a recurrent state cannot be cut back); attention cells past n_cached are dropped. The reply + // tokens[0..n) decode at positions n_cached.. with labels[i] the target after tokens[i] (-1: + // none): the trainable adapter applies to labelled rows only, the context is never decoded + // again, and afterwards the reply is restored under her stable set, so the sequence (every + // cell, the recurrent state) is what serving left and every other sequence is untouched. + // A reply within one batch is one exact step; a longer one steps per batch, its gradient + // stopped at each batch boundary. false: refused or failed by name (llama_opt_failure), + // with the sequence restored. + LLAMA_API bool llama_opt_learn( + struct llama_context * lctx, + llama_seq_id seq, + llama_pos n_cached, + const llama_token * tokens, + const llama_token * labels, + int32_t n, + ggml_opt_result_t result); + LLAMA_API void llama_opt_epoch( struct llama_context * lctx, ggml_opt_dataset_t dataset, diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp index 2b98a552f48f..b4a1626afba6 100644 --- a/src/llama-batch.cpp +++ b/src/llama-batch.cpp @@ -1,5 +1,7 @@ #include "llama-batch.h" +#include + #include "llama-impl.h" #include "llama-vocab.h" #include "llama-memory.h" @@ -56,7 +58,15 @@ bool llama_batch_allocr::init( } if (batch.seq_id) { + if (!batch.n_seq_id) { + LLAMA_LOG_ERROR("%s: sequence IDs require sequence counts\n", __func__); + return false; + } for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.n_seq_id[i] <= 0 || batch.n_seq_id[i] > (int32_t) n_seq_max || !batch.seq_id[i]) { + LLAMA_LOG_ERROR("%s: invalid sequence list at position %d\n", __func__, i); + return false; + } for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) { if (batch.seq_id && (batch.seq_id[i][s] < 0 || batch.seq_id[i][s] >= (llama_seq_id) n_seq_max)) { LLAMA_LOG_ERROR("%s: invalid seq_id[%d][%d] = %d >= %d\n", __func__, i, s, batch.seq_id[i][s], (llama_seq_id) n_seq_max); @@ -157,6 +167,41 @@ bool llama_batch_allocr::init( n_outputs += batch.logits[i] != 0; } + // a labelled position is a training row: its logits row is what the loss reads. A run of + // them trains given everything before it as served (the shadow adapter applies to labelled + // rows only) and is restored to her stable set after, so in one batch a sequence's labelled + // rows are one run that ends its rows: anything after it would read the run through the + // shadow instead of as served. + if (batch.label) { + bool in_run[LLAMA_MAX_SEQ] = {}; + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.label[i] < -1 || (batch.label[i] >= 0 && (uint32_t) batch.label[i] >= vocab.n_tokens())) { + LLAMA_LOG_ERROR("%s: invalid label[%d] = %d\n", __func__, i, batch.label[i]); + return false; + } + if (batch.label[i] >= 0 && batch.label_weight && !(std::isfinite(batch.label_weight[i]) && batch.label_weight[i] >= 0.0f)) { + LLAMA_LOG_ERROR("%s: invalid label_weight[%d] = %f: a weight is finite and >= 0\n", __func__, i, batch.label_weight[i]); + return false; + } + const llama_seq_id s = batch.seq_id[i][0]; + if (batch.label[i] < 0) { + for (int32_t j = 0; j < batch.n_seq_id[i]; ++j) { + if (in_run[batch.seq_id[i][j]]) { + LLAMA_LOG_ERROR("%s: position %d follows a labelled run of sequence %d in one batch: a run ends its sequence's rows\n", __func__, i, batch.seq_id[i][j]); + return false; + } + } + continue; + } + if (batch.n_seq_id[i] != 1) { + LLAMA_LOG_ERROR("%s: position %d has a label and %d sequences: a training row belongs to one sequence\n", __func__, i, batch.n_seq_id[i]); + return false; + } + in_run[s] = true; + ++n_labels; + } + } + has_cpl = false; // determine coupled sequences @@ -224,6 +269,8 @@ bool llama_batch_allocr::init( /*.seq_id_unq =*/ this->seq_id_unq.data(), /*.seq_idx =*/ this->seq_idx.data(), /*.output =*/ batch.logits, + /*.label =*/ batch.label, + /*.label_weight =*/ batch.label_weight, /*.data =*/ {}, }; @@ -430,6 +477,8 @@ llama_ubatch llama_batch_allocr::ubatch_reserve(uint32_t n_seq_tokens, uint32_t /*.seq_id_unq =*/ udata->seq_id_unq.data(), /*.seq_idx =*/ udata->seq_idx.data(), /*.output =*/ udata->output.data(), + /*.label =*/ nullptr, + /*.label_weight =*/ nullptr, /*.data =*/ std::move(udata), }; @@ -448,6 +497,10 @@ uint32_t llama_batch_allocr::get_n_outputs() const { return n_outputs; } +uint32_t llama_batch_allocr::get_n_labels() const { + return n_labels; +} + uint32_t llama_batch_allocr::get_n_used() const { return n_used; } @@ -722,6 +775,7 @@ llama_ubatch llama_batch_allocr::split_seq(uint32_t n_ubatch) { void llama_batch_allocr::clear() { n_outputs = 0; + n_labels = 0; batch = {}; @@ -764,6 +818,8 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u udata->seq_id_unq.resize(0); udata->seq_idx .resize(LLAMA_MAX_SEQ, -1); udata->output .resize(n_tokens); + udata->label .resize(batch.label ? n_tokens : 0); + udata->label_weight.resize(batch.label && batch.label_weight ? n_tokens : 0); udata->seq_id_data.reserve(n_tokens); @@ -789,6 +845,12 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u udata->n_seq_id[i] = batch.n_seq_id[idxs[i]]; udata->output[i] = batch.logits[idxs[i]]; + if (batch.label) { + udata->label[i] = batch.label[idxs[i]]; + if (batch.label_weight) { + udata->label_weight[i] = batch.label_weight[idxs[i]]; + } + } for (int s = 0; s < udata->n_seq_id[i]; ++s) { const llama_seq_id seq_id = batch.seq_id[idxs[i]][s]; @@ -831,6 +893,8 @@ llama_ubatch llama_batch_allocr::ubatch_add(const std::vector & idxs, u /*.seq_id_unq =*/ udata->seq_id_unq.data(), /*.seq_idx =*/ udata->seq_idx.data(), /*.output =*/ udata->output.data(), + /*.label =*/ batch.label ? udata->label.data() : nullptr, + /*.label_weight =*/ batch.label && batch.label_weight ? udata->label_weight.data() : nullptr, /*.data =*/ std::move(udata), }; @@ -939,6 +1003,8 @@ struct llama_batch llama_batch_get_one( /*n_seq_id =*/ nullptr, /*seq_id =*/ nullptr, /*logits =*/ nullptr, + /*label =*/ nullptr, + /*label_weight =*/ nullptr, }; } @@ -951,6 +1017,8 @@ struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_ /*n_seq_id =*/ nullptr, /*seq_id =*/ nullptr, /*logits =*/ nullptr, + /*label =*/ nullptr, + /*label_weight =*/ nullptr, }; if (embd) { @@ -968,6 +1036,10 @@ struct llama_batch llama_batch_init(int32_t n_tokens_alloc, int32_t embd, int32_ batch.seq_id[n_tokens_alloc] = nullptr; batch.logits = (int8_t *) malloc(sizeof(int8_t) * n_tokens_alloc); + batch.label = (llama_token *) malloc(sizeof(llama_token) * n_tokens_alloc); + std::fill(batch.label, batch.label + n_tokens_alloc, -1); // no position trains until labelled + batch.label_weight = (float *) malloc(sizeof(float) * n_tokens_alloc); + std::fill(batch.label_weight, batch.label_weight + n_tokens_alloc, 1.0f); // every labelled position counts 1 until weighted return batch; } @@ -984,4 +1056,6 @@ void llama_batch_free(struct llama_batch batch) { free(batch.seq_id); } if (batch.logits) free(batch.logits); + if (batch.label) free(batch.label); + if (batch.label_weight) free(batch.label_weight); } diff --git a/src/llama-batch.h b/src/llama-batch.h index a3d1889d4a04..fca6f0571f6b 100644 --- a/src/llama-batch.h +++ b/src/llama-batch.h @@ -13,10 +13,37 @@ // keep this struct lightweight struct llama_ubatch { + struct row_map { + std::vector public_rows; + std::vector label_rows; + }; + + row_map rows() const { + row_map result; + for (uint32_t i = 0; i < n_tokens; ++i) { + if (output && output[i]) { + result.public_rows.push_back((int32_t) i); + } + if (label && label[i] >= 0) { + result.label_rows.push_back((int32_t) i); + } + } + return result; + } + bool equal_seqs() const { return b_equal_seqs != 0; } + // the training rows (a label >= 0) + uint32_t n_labels() const { + uint32_t n = 0; + for (uint32_t i = 0; label != nullptr && i < n_tokens; ++i) { + n += label[i] >= 0; + } + return n; + } + // typical for M-RoPE cases: // 0 - sequential position of the tokens/embeddings in the sequence // 1 - y position in the image @@ -50,6 +77,8 @@ struct llama_ubatch { llama_seq_id * seq_id_unq; // [n_seqs_unq] | s | seq_id int32_t * seq_idx; // [LLAMA_MAX_SEQ] | - | seq_idx int8_t * output; // [n_tokens] | i | - + llama_token * label; // [n_tokens] | i | training target, -1 for none (nullptr: inference) + float * label_weight; // [n_tokens] | i | a labelled row's weight in the loss (nullptr: 1) struct data_t { std::vector token; @@ -60,6 +89,8 @@ struct llama_ubatch { std::vector seq_id_unq; std::vector seq_idx; std::vector output; + std::vector label; + std::vector label_weight; std::vector seq_id_data; }; @@ -87,6 +118,9 @@ class llama_batch_allocr { uint32_t get_n_tokens() const; uint32_t get_n_outputs() const; + + // labelled positions (training rows) in the batch; 0 for an inference batch + uint32_t get_n_labels() const; uint32_t get_n_used() const; // the array of output indices in the order they were encountered during the ubatch splitting @@ -136,6 +170,7 @@ class llama_batch_allocr { uint32_t n_embd; uint32_t n_seq_max; uint32_t n_outputs; + uint32_t n_labels = 0; std::array seq_id_0 = {{ 0 }}; // default sequence id diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 3c183f3d255c..42969c259a86 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -40,6 +40,44 @@ static bool llama_set_param(struct ggml_tensor * tensor, llama_opt_param_filter return true; } +// whether every GPU this context computes on has a kernel for ggml_flash_attn_back, at this +// model's attention shape and this context's cache types (the CPU has one) +bool llama_context::opt_fa_back_native() const { + const llama_hparams & hp = model.hparams; + uint32_t il = 0; + while (il < hp.n_layer() && hp.is_recr(il)) { + ++il; + } + if (il == hp.n_layer()) { + return false; // no attention layer + } + // the cache the step's attention reads in place + const llama_kv_cache * kv = dynamic_cast(memory.get()); + if (const auto * hybrid = dynamic_cast(memory.get())) { + kv = hybrid->get_mem_attn(); + } + if (kv == nullptr) { + return false; + } + ggml_init_params ip = { 8*ggml_tensor_overhead(), nullptr, /*no_alloc =*/ true }; + ggml_context_ptr ctx(ggml_init(ip)); + const int64_t H = hp.n_head(il), Hkv = hp.n_head_kv(il), D = hp.n_embd_head_k(il), Dv = hp.n_embd_head_v(il); + ggml_tensor * q = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, D, 8, H); + ggml_tensor * k = ggml_new_tensor_3d(ctx.get(), kv->type_k(), D, 64, Hkv); + ggml_tensor * v = ggml_new_tensor_3d(ctx.get(), kv->type_v(), Dv, 64, Hkv); + ggml_tensor * m = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_F16, 64, 8); + ggml_tensor * d = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, Dv, H, 8); + ggml_tensor * op = ggml_flash_attn_back(ctx.get(), q, k, v, m, d, 1.0f, 0); + for (const auto & backend : backends) { + ggml_backend_dev_t dev = ggml_backend_get_device(backend.get()); + const auto type = dev ? ggml_backend_dev_type(dev) : GGML_BACKEND_DEVICE_TYPE_CPU; + if ((type == GGML_BACKEND_DEVICE_TYPE_GPU || type == GGML_BACKEND_DEVICE_TYPE_IGPU) && !ggml_backend_dev_supports_op(dev, op)) { + return false; + } + } + return true; +} + void llama_context::opt_init(struct llama_model * model, struct llama_opt_params lopt_params) { GGML_ASSERT(!opt_ctx); // The training window lives on this context, not the shared model (a serving context may @@ -57,7 +95,8 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params // and the context before it from the cache as a constant (build_attn); opt_epoch_iter // walks the window decoding context and training each labelled chunk. One chunk is one // batch is one ubatch: a chunk's backward needs all of its own K/V in one graph. - GGML_ASSERT(n_ubatch == n_batch && "training needs one ubatch per batch: set -ub = -b"); + // (the walk alone needs one ubatch per batch: a chunk's backward needs all of its own K/V in one + // graph; opt_epoch refuses otherwise. llama_opt_learn trains in a serving context's ubatches.) // In a flash-attention context the training graph reads attention through FLASH_ATTN_EXT too // (build_attn_mha), whose backward is tiled over the window; flash attention also serves the // walk's context decodes and lets the cache hold V un-transposed. @@ -66,9 +105,12 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params // (no backward pass) and before any adapter attached since (its nodes uncounted). // Size both again for the training graph, from the same number, plus the nodes each // attention layer's tiled backward adds for this window (one count, owned by ggml). + // the attention backward: one op where every GPU has the kernel, else composed tiles whose node + // count grows with the window (decided first: it sizes the graph) + cparams.fa_back_native = opt_fa_back_native(); { size_t max_nodes = graph_max_nodes(n_ubatch); - if (cparams.flash_attn) { + if (cparams.flash_attn && !cparams.fa_back_native) { max_nodes += (size_t) model->hparams.n_layer() * ggml_flash_attn_ext_back_n_nodes(opt_n_ctx_train, 0); } gf_res_prev.reset(new llm_graph_result(max_nodes)); @@ -76,7 +118,11 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params max_nodes, cparams.pipeline_parallel, cparams.op_offload)); } - ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY); + LLAMA_LOG_INFO("%s: the attention backward is %s\n", __func__, + cparams.fa_back_native ? "one op (ggml_flash_attn_back): no score is kept" : "composed from tiles (a device has no ggml_flash_attn_back kernel)"); + + // one target token per row: no one-hot tensor as wide as the vocabulary beside the logits + ggml_opt_params opt_params = ggml_opt_default_params(sched.get(), GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS); opt_params.opt_period = n_batch / n_ubatch; 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) @@ -102,6 +148,8 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params n_params += llama_set_param(ab.b, param_filter, param_filter_ud); } GGML_ASSERT(n_params > 0 && "the param filter selected none of the adapter's tensors"); + // the trained adapter is the shadow: labelled rows only, never a serving row + opt_shadow = lopt_params.adapter; LLAMA_LOG_INFO("%s: LoRA-only training: %lld adapter tensors trainable, base frozen\n", __func__, (long long) n_params); return; } @@ -129,13 +177,17 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params } } -// THE WALK (Fable's design, continuum 2026-10-06): one pass over a window in one context. A run -// with no label is context she did not write: a plain decode into the cache, no backward. A -// labelled run is her reply: it trains in chunks of n_batch, each attending to everything before -// it read from the cache as a constant (build_attn, cparams.training), then the reply is decoded -// into the cache under the adapter as it now is, so the next reply sees it. Memory is chunk x -// window, never window x window, so her whole lived window trains. Nothing after the last label -// is computed. One per window; a member of the context, so it runs on the context's own state. +// THE WALK (Fable's design, continuum 2026-10-06; objective 2026-10-14): one pass over a window +// in one context. A run with no label is context: a plain decode into the cache, as served, no +// backward. A labelled run is her reply, trained GIVEN EVERYTHING BEFORE IT AS SERVED: only its +// rows see the shadow adapter (the label mask), so the context before it carries no gradient and +// is never trained. The run trains in chunks of n_batch, each attending to everything before it +// in the cache; its chunks write their own cells under the shadow, so a later chunk of the run +// reads an earlier one as the shadow computes it. When the run is done it is RESTORED: its cells +// and the recurrent state are decoded again under her stable set, so the cache is always her +// stable self's and the next run reads this one as served. Memory is chunk x window, never +// window x window, so her whole lived window trains. Nothing after the last label is computed. +// One per window; a member of the context, so it runs on the context's own state. struct llama_context::opt_walk { llama_context & lctx; ggml_opt_dataset_t dataset; @@ -171,8 +223,11 @@ struct llama_context::opt_walk { n_ctx(n_ctx), n_batch(n_batch), plan(llama_walk_plan::of(labels, n_ctx, n_batch)) {} bool yield_point(); - bool decode_span(uint32_t p0, uint32_t p1); + bool decode_span(uint32_t p0, uint32_t p1, bool shadow); + std::vector checkpoint(); + bool rewind(uint32_t p_pop, const std::vector & snap, uint32_t p_to); bool train_chunk(uint32_t pos_ctx, uint32_t n_chunk); + bool restore(const llama_walk_run & run, const std::vector & snap); void run_exact(); void run_plain(); @@ -198,225 +253,238 @@ bool llama_context::opt_walk::yield_point() { return !lctx.opt_stop_requested.load(std::memory_order_relaxed); } -// a plain inference decode of [p0, p1) into the cache: the training graph's attention never -// writes the cache, so context and finished replies reach it this way -bool llama_context::opt_walk::decode_span(uint32_t p0, uint32_t p1) { - llama_cparams & cparams = lctx.cparams; - const bool training = cparams.training; - cparams.training = false; - lctx.gf_res_prev->reset(); +// an inference decode of [p0, p1) into the cache, no loss and no backward. shadow: the labelled +// rows decode under the shadow adapter (a run's forward pass); otherwise every row decodes under +// her stable set (context, and a run restored). +bool llama_context::opt_walk::decode_span(uint32_t p0, uint32_t p1, bool shadow) { bool ok = true; for (uint32_t c0 = p0; c0 < p1 && ok; c0 += n_batch) { if (!yield_point()) { - cparams.training = training; - lctx.gf_res_prev->reset(); return false; // stopped or cancelled at a chunk boundary: not a decode failure } const uint32_t c1 = std::min(c0 + n_batch, p1); llama_walk_fill(batch, tokens, c0, c1 - c0); + for (uint32_t i = 0; shadow && i < c1 - c0; ++i) { + // a label marks the row the shadow sees; the allocator wants it an output row + batch.label [i] = labels[c0 + i]; + batch.logits[i] = labels[c0 + i] >= 0; + } batch.logits[c1 - c0 - 1] = true; + lctx.opt_walk_active = shadow ? this : nullptr; + lctx.opt_forward_only = shadow; ok = lctx.decode(batch) == 0; + lctx.opt_walk_active = nullptr; + lctx.opt_forward_only = false; + std::fill(batch.label, batch.label + (c1 - c0), -1); } - cparams.training = training; - lctx.gf_res_prev->reset(); if (!ok) { - LLAMA_LOG_ERROR("%s: the context decode of [%u, %u) failed\n", __func__, p0, p1); + LLAMA_LOG_ERROR("%s: the decode of [%u, %u) failed\n", __func__, p0, p1); } return ok; } -// one training chunk [pos_ctx, pos_ctx + n_chunk): forward and backward over it alone, attending -// to the cached context before it -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; +// the recurrent state of the walk's sequence through the sequence-state API (the recurrent part +// only, PARTIAL_ONLY), after waiting for the queued decode; empty for a model with none. A +// hand-rolled per-layer copy skipped that wait and, on Metal, carried a torn mid-graph state into +// the reverse pass (2026-10-11). +std::vector llama_context::opt_walk::checkpoint() { + std::vector snap; + if (recr != nullptr) { + lctx.synchronize(); + snap.resize(lctx.state_seq_get_size(0, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)); + GGML_ASSERT(lctx.state_seq_get_data(0, snap.data(), snap.size(), LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == snap.size()); + } + return snap; +} +// back to position p_pop: the cache holds [0, p_pop), the recurrent state is `snap` (taken at +// p_pop); then forward to p_to, under the shadow for the run's labelled rows. A restore puts back +// the cell's position with its rows. Also run before every retry of a refused chunk: a refused +// chunk graph is refused AFTER its ubatch was applied, so its cells already sit in the cache and +// the recurrent cell's position has moved to the chunk's end; retrying on top of that fails to +// prepare the ubatch at all (the 5090, 2026-10-10: "init_batch: failed to prepare attention +// ubatches" after "retrying at 512"). +bool llama_context::opt_walk::rewind(uint32_t p_pop, const std::vector & snap, uint32_t p_to) { + const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : lctx.memory->seq_rm(0, p_pop, -1); + if (!popped) { + LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u)\n", __func__, p_pop); + lctx.opt_stop_requested.store(true); + return false; + } + if (recr != nullptr) { + lctx.synchronize(); + GGML_ASSERT(lctx.state_seq_set_data(0, snap.data(), snap.size(), LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == snap.size()); + } + return p_pop >= p_to || decode_span(p_pop, p_to, /*shadow =*/ true); +} + +// one training chunk [pos_ctx, pos_ctx + n_chunk) of a run: one llama_decode with labels. The +// batch loop is decode's, and the extras switch on in process_ubatch (the shadow's masked term, +// then llama_context::opt_ubatch). This function only fills the chunk's tokens and labels and +// says whose walk it is. +bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { llama_walk_fill(batch, tokens, pos_ctx, n_chunk); for (uint32_t pos_batch = 0; pos_batch < n_chunk; ++pos_batch) { - // an exact walk's context chunk has no label but still trains (its K/V carry later - // chunks' gradient): one output row, weighted zero, so the graph has its loss - batch.logits[pos_batch] = labels[pos_ctx + pos_batch] >= 0 - || (xc != nullptr && xc->n_labels == 0 && pos_batch + 1 == n_chunk); + batch.logits[pos_batch] = false; + batch.label [pos_batch] = labels[pos_ctx + pos_batch]; } - // output_all = false: the batch's own logits flags decide which rows exist - if (!lctx.balloc->init(batch, model.vocab, nullptr, model.hparams.n_embd_inp(), cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max, false)) { - LLAMA_LOG_ERROR("%s: failed to initialize batch\n", __func__); + lctx.opt_walk_active = this; + lctx.opt_result_active = result; + const int ret = lctx.decode(batch); + lctx.opt_walk_active = nullptr; + lctx.opt_result_active = nullptr; + std::fill(batch.label, batch.label + n_chunk, -1); + + if (ret != 0) { + // a refusal has already stopped the epoch by name (opt_fail); anything else is decode's + if (!lctx.opt_alloc_failed.load()) { + lctx.opt_fail("the training chunk failed to decode (llama_decode returned " + std::to_string(ret) + ")"); + } return false; } + if (callback) { + callback(train, lctx.opt_ctx, dataset, result, idata_in_loop + pos_ctx/n_batch + 1, ndata_in_loop, t_loop_start); + } + return true; +} - lctx.n_queued_tokens += lctx.balloc->get_n_tokens(); - - lctx.embd_seq.clear(); +bool llama_context::opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch) { + const uint32_t n_labels = ubatch.n_labels(); + opt_walk * walk = opt_walk_active; + const llama_walk_chunk * xc = walk != nullptr ? walk->xc : nullptr; - 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; + 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); } - - 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__); + std::vector targets; + std::vector weights; // a labelled row's weight in the loss, when the batch carries them + targets.reserve(n_labels); + for (int32_t row : ubatch.rows().label_rows) { + targets.push_back(ubatch.label[row]); + if (ubatch.label_weight != nullptr) { + weights.push_back(ubatch.label_weight[row]); + } + } + // an unweighted batch keeps the unweighted loss (every backend has its kernels) + ggml_opt_set_next_label_weights(opt_ctx, ubatch.label_weight != nullptr); + if (res->t_loss_slices.empty()) { + ggml_free(ctx_compute_opt); + opt_fail("training head was not projected into bounded label slices"); 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); + ggml_opt_prepare_alloc_slices(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), + res->t_loss_slices.data(), res->t_loss_slices.size()); + if (xc == nullptr && opt_learn_labels > 0) { + // llama_opt_learn: one step per call, at the ubatch that completes its labelled rows + const int64_t nl = ubatch.n_labels(); + opt_learn_labels_done += nl; + ggml_opt_set_next_step(opt_ctx, opt_learn_labels_done >= opt_learn_labels, (float) nl / (float) opt_learn_labels, nullptr); + } + if (xc != nullptr) { + ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); + for (const auto & io : res->t_walk) { + if (io.tk) { + // the run's cached rows before this chunk, as F32: their gradient into the accumulator + ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); + ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); } - lctx.n_outputs = n_outputs_ubatch; - } - - if (!mctx->apply()) { - LLAMA_LOG_ERROR("%s: failed to update the memory context\n", __func__); - return false; - } - - auto * res = lctx.gf_res_prev.get(); - - const auto gparams = lctx.graph_params(res, ubatch, mctx.get(), lctx.graph_type()); - - res->reset(); - - auto * gf = model.build_graph(gparams); - - struct ggml_context * ctx_compute_opt; - { - const size_t size_gf = ggml_graph_size(gf); - const size_t size_meta = 4*size_gf*ggml_tensor_overhead() + 2*ggml_graph_overhead_custom(size_gf, /*grads = */ true); - struct ggml_init_params params = { - /*.mem_size =*/ size_meta, - /*.mem_buffer =*/ nullptr, - /*.no_alloc =*/ true, - }; - ctx_compute_opt = ggml_init(params); } - ggml_opt_prepare_alloc(opt_ctx, ctx_compute_opt, gf, res->get_inp_tokens(), res->get_logits()); - if (xc != nullptr) { - ggml_opt_set_next_step(opt_ctx, xc->period_end, xc->loss_scale, res->t_walk_surrogate); - for (const auto & io : res->t_walk) { - if (io.tk) { - ggml_opt_accumulate_grad(opt_ctx, io.tk, io.ak); - ggml_opt_accumulate_grad(opt_ctx, io.tv, io.av); - } + for (const auto & ws : res->t_walk_state) { + if (ws.tap) { + ggml_opt_accumulate_grad(opt_ctx, ws.tap, walk->acc.state.at(ws.cache)); } - for (const auto & ws : res->t_walk_state) { - if (ws.tap) { - ggml_opt_accumulate_grad(opt_ctx, ws.tap, acc.state.at(ws.cache)); - } - } - } - if (!ggml_opt_alloc(opt_ctx, train)) { - ggml_free(ctx_compute_opt); - lctx.opt_fail(ggml_opt_refusal(opt_ctx), - std::string("the ") + (train ? "training" : "evaluation") + " graph was refused (window " + std::to_string(n_ctx) + ")"); - return false; } + } + const bool train = walk == nullptr || walk->train; + if (!ggml_opt_alloc(opt_ctx, train)) { + ggml_free(ctx_compute_opt); + opt_fail(ggml_opt_refusal(opt_ctx), + std::string("the ") + (train ? "training" : "evaluation") + " graph was refused (window " + std::to_string(walk != nullptr ? walk->n_ctx : ubatch.n_tokens) + ")"); + return false; + } - res->set_inputs(&ubatch); - { - // one-hot rows for the labelled positions only, in the order the ubatch emits its - // outputs (the same order the logits rows take) - struct ggml_tensor * labels_t = ggml_opt_labels(opt_ctx); - GGML_ASSERT(labels_t->ne[1] == (int64_t) lctx.n_outputs); - ggml_set_zero(labels_t); - const float onef = 1.0f; - uint32_t row = 0; - for (uint32_t pos_ubatch = 0; pos_ubatch < ubatch.n_tokens; ++pos_ubatch) { - if (ubatch.output[pos_ubatch] == 0) { - continue; - } - const uint32_t ilabel = pos_ctx + pos_batch + pos_ubatch; - // an exact walk's context chunk: its one row has no label (weighted zero) - GGML_ASSERT((xc != nullptr && labels[ilabel] < 0) || (labels[ilabel] >= 0 && labels[ilabel] < labels_t->ne[0])); - if (labels[ilabel] >= 0) { - ggml_backend_tensor_set(labels_t, &onef, (row*labels_t->ne[0] + labels[ilabel])*sizeof(float), sizeof(float)); - } - ++row; - } - GGML_ASSERT(row == lctx.n_outputs); + res->set_inputs(&ubatch); + { + // the target token of each output row, in the order the ubatch emits its outputs (the + // order the logits rows take) + struct ggml_tensor * labels_t = ggml_opt_labels(opt_ctx); + GGML_ASSERT(labels_t->type == GGML_TYPE_I32 && labels_t->ne[0] == (int64_t) n_labels); + + ggml_backend_tensor_set(labels_t, targets.data(), 0, ggml_nbytes(labels_t)); + if (struct ggml_tensor * weights_t = ggml_opt_label_weights(opt_ctx)) { + GGML_ASSERT(weights.size() == (size_t) n_labels); + ggml_backend_tensor_set(weights_t, weights.data(), 0, ggml_nbytes(weights_t)); } - if (xc != nullptr) { - // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the - // 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)); + } + if (xc != nullptr) { + // the later chunk's dL/d(entry state) is this chunk's exit gradient; then the + // accumulator takes this chunk's own + for (const auto & ws : res->t_walk_state) { + ggml_tensor * a = walk->acc.state.at(ws.cache); + if (ws.gs) { + GGML_ASSERT(ggml_nbytes(a) == ggml_nbytes(ws.gs)); + ggml_backend_tensor_set(ws.gs, a->data, 0, ggml_nbytes(ws.gs)); } + memset(a->data, 0, ggml_nbytes(a)); } - ggml_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()); + } + ggml_opt_eval(opt_ctx, opt_result_active); + ggml_free(ctx_compute_opt); + if (const char * why = ggml_opt_refusal(opt_ctx); why[0] != '\0') { + // the backend failed the graph (ggml_opt_eval): the run fails, as a refused graph does + opt_fail(why); + return false; + } + if (xc != nullptr && !res->t_walk_state.empty()) { + walk->acc.state_ready = true; + } return true; } // THE EXACT WALK. The plain walk steps per chunk and stops the gradient at every chunk boundary: a // reply learns only through its own chunk's K/V, never through what an earlier chunk computed for -// it. Here the window is decoded once under the adapter as it is, then trained in REVERSE: each -// chunk's graph attends to the cache, with its last walk_horizon prefix positions behind a -// gradient tap, and adds the surrogate + on its own K/V, where G is what every later -// chunk put there. Its backward then carries that gradient through itself into the adapter and -// into its own prefix's G. ONE optimizer step per window, on the gradient of the window's loss -// (exact within the horizon). Memory stays chunk x window; the cost is one decode plus one forward -// and backward per chunk, context chunks included. -// THE INVARIANT THAT MAKES THE REVERSE PASS VALID (Fable): the adapter does not change between the -// forward decode and the last chunk's backward. One optimizer step per window, taken by the chunk -// the reverse pass trains last, is what keeps each chunk's recomputed K/V (and recurrent state) +// it. Here each run is decoded once under the shadow, then trained in REVERSE: each chunk's graph +// attends to the cache, with the run's positions within walk_horizon behind a gradient tap, and +// adds the surrogate + on its own K/V, where G is what every later chunk of the run +// put there. Its backward then carries that gradient through itself into the adapter and into the +// run's earlier positions' G. ONE optimizer step per window, on the gradient of the window's loss +// (exact within the horizon). Memory stays chunk x window; the cost is one decode of the window, +// and per run one more decode plus one forward and backward per chunk and the restore. +// THE INVARIANT THAT MAKES THE REVERSE PASS VALID (Fable): the adapter does not change between a +// run's forward decode and its last chunk's backward. One optimizer step per window, taken by the +// chunk the walk trains last, is what keeps each chunk's recomputed K/V (and recurrent state) // equal to the cached values its successors attended to; a step inside the window would make // every surrogate a gradient of a different function. // A recurrent model's state is carried the same way, one chunk at a time: the state entering a // chunk is tapped into a host accumulator, the state it leaves meets the next chunk's gradient // (build_rs, build_walk_state_exit). The reverse pass restores each chunk's entry state from a -// checkpoint taken at its boundary in the forward pass. +// checkpoint taken at its boundary in the forward pass. The state entering a run is her stable +// set's: no gradient crosses it. void llama_context::opt_walk::run_exact() { // Each window starts from the REQUESTED horizon: a halving that fit one window's device budget // must not shrink every later window and epoch, which may be shorter (Fable). lctx.opt_walk_horizon = lctx.opt_walk_horizon_req; - // A recurrent context with rollback slots keeps several states per sequence; the walk carries - // exactly one. Refused by name, never asserted: training runs inside the serving process, and - // an assert there takes serving down (2026-10-07 04:09Z on the 5090). - if (recr != nullptr && lctx.cparams.n_rs_seq > 0) { - lctx.opt_fail("the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has " - + std::to_string(lctx.cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0"); + // A serving context keeps rollback slots for drafts; the walk's graphs keep one state, the live + // one (llm_graph_context::rs_keep), so they train in that context. Only a PENDING undo reads a + // snapshot instead of the live state: that alone is refused, by name, never asserted, since + // training runs inside the serving process (2026-10-07 04:09Z on the 5090). + if (recr != nullptr && recr->rollback_pending(0)) { + lctx.opt_fail("the exact walk cannot train a sequence with a pending recurrent rollback (n_rs_seq = " + + std::to_string(lctx.cparams.n_rs_seq) + "): finish or drop the draft first"); return; } - plan.mark_exact(recr != nullptr, lctx.opt_walk_horizon); + plan.mark_exact(lctx.opt_walk_horizon); const std::vector & chunks = plan.chunks; - const uint32_t n_pos = (uint32_t) plan.last_label + 1; + const uint32_t n_pos = plan.longest_run(); // THE HOST BUDGET (llama_walk_plan::stride_for): the K/V accumulators and the recurrent state // checkpoints are both known before anything runs @@ -429,8 +497,8 @@ void llama_context::opt_walk::run_exact() { if (budget > 0 && bytes > budget) { char why[512]; snprintf(why, sizeof(why), - "the exact walk needs %.1f MiB of host memory for this window (%u positions, the K/V gradient alone), over its budget of %.1f MiB: " - "the K/V gradient grows with the window; neither the horizon nor the chunk size changes it", + "the exact walk needs %.1f MiB of host memory for this window (its longest reply, %u positions, the K/V gradient alone), over its budget of %.1f MiB: " + "the K/V gradient grows with the reply; neither the horizon nor the chunk size changes it", bytes/1048576.0, n_pos, budget/1048576.0); lctx.opt_fail(why); return; @@ -444,151 +512,118 @@ void llama_context::opt_walk::run_exact() { return; } - // the window as the adapter now reads it, chunk by chunk: a recurrent model's state is - // checkpointed at each chunk's start through the sequence-state API (the recurrent part only, - // PARTIAL_ONLY), after waiting for the queued decode; a restore puts back the cell's position - // with its rows, and an empty checkpoint (before the first chunk) restores the zero state. A - // hand-rolled per-layer copy here skipped that wait and, on Metal, carried a torn mid-graph - // state into the reverse pass (2026-10-11). - std::vector> snaps(recr != nullptr ? chunks.size() : 0); - const llama_state_seq_flags state_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY; - for (size_t j = 0; j < chunks.size(); ++j) { - if (recr != nullptr && j % stride == 0) { - lctx.synchronize(); - snaps[j].resize(lctx.state_seq_get_size(0, state_flags)); - GGML_ASSERT(lctx.state_seq_get_data(0, snaps[j].data(), snaps[j].size(), state_flags) == snaps[j].size()); - } - if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n)) { + uint32_t p = 0; // the cache holds [0, p) as served + for (const llama_walk_run & run : plan.runs) { + const uint32_t r0 = plan.begin(run); + const uint32_t r1 = plan.end(run); + if (!decode_span(p, r0, /*shadow =*/ false)) { return; } - } - for (int64_t j = (int64_t) chunks.size() - 1; j >= plan.first; --j) { - const llama_walk_chunk & c = chunks[j]; - if (!c.needed) { - continue; - } - if (!yield_point()) { - break; - } - // attention: pop the cache to [0, c0); recurrent: restore the state at the chunk's - // checkpoint and decode forward to its start (that decode rebuilds the attention K/V in - // between too, under the same adapter: the same values). Run again before every retry - // below: a refused chunk graph is refused AFTER its ubatch was applied, so the chunk's - // cells already sit in the cache and the recurrent cell's position has moved to the - // chunk's end; retrying on top of that fails to prepare the ubatch at all (the 5090, - // 2026-10-10: "init_batch: failed to prepare attention ubatches" after "retrying at 512", - // and the epoch stopped instead of shrinking the horizon). Always a plain decode: the - // walk's graph flags are off whenever it runs. - const auto rewind = [&]() -> bool { - const size_t cp = recr != nullptr ? (size_t) j - (size_t) j % stride : (size_t) j; - const uint32_t p_pop = chunks[cp].c0; - const bool popped = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, p_pop, -1)) : lctx.memory->seq_rm(0, p_pop, -1); - if (!popped) { - LLAMA_LOG_ERROR("%s: could not pop the cache to [0, %u) for the reverse pass\n", __func__, p_pop); - lctx.opt_stop_requested.store(true); - return false; + // the run under the shadow, checkpointed every `stride` chunks from its start + std::vector> snaps(run.j1 - run.j0); + for (size_t j = run.j0; j < run.j1; ++j) { + if ((j - run.j0) % stride == 0) { + snaps[j - run.j0] = checkpoint(); } - if (recr != nullptr) { - lctx.synchronize(); - GGML_ASSERT(lctx.state_seq_set_data(0, snaps[cp].data(), snaps[cp].size(), state_flags) == snaps[cp].size()); - if (p_pop < c.c0 && !decode_span(p_pop, c.c0)) { - return false; - } + if (!decode_span(chunks[j].c0, chunks[j].c0 + chunks[j].n, /*shadow =*/ true)) { + return; } - return true; - }; - if (!rewind()) { - break; } - const bool state_surrogate = recr != nullptr && j < (int64_t) chunks.size() - 1 && acc.state_ready; - // THE HORIZON FITS THE DEVICE. Its gradient taps are the chunk graph's largest growth with - // the window (layers x horizon x (k + v) floats, twice with their gradient); the graph - // preflight refuses a chunk that does not fit before allocating anything, so a memory - // refusal halves the horizon (whole chunks) and the chunk is tried again; every later - // chunk keeps the smaller one. The forward still attends to the whole window and the - // recurrent state is carried exactly; only attention gradient past the horizon is dropped. - // A refusal at a horizon of one chunk, or for a node the device cannot run, stands. - llama_walk_chunk cc = c; - bool ok; - for (;;) { - cc.grad_from = llama_walk_plan::grad_from(cc.c0, lctx.opt_walk_horizon); - set_flags(true, cc.grad_from, c.surrogate, state_surrogate); - xc = &cc; - ok = train_chunk(cc.c0, cc.n); - xc = nullptr; - set_flags(false, 0, false, false); - const uint32_t span = cc.c0 - cc.grad_from; - if (ok || !lctx.opt_failure.empty() || !lctx.opt_alloc_failed.load() || span <= n_batch) { - break; + // the reverse pass: back to the chunk's checkpoint, forward to its start, train it + acc.start_run(r0); + for (size_t j = run.j1; j-- > run.j0; ) { + const llama_walk_chunk & c = chunks[j]; + if (!yield_point()) { + return; } - lctx.opt_walk_horizon = std::max(n_batch, (span/2)/n_batch*n_batch); - LLAMA_LOG_WARN("%s: the chunk at %u did not fit with a gradient horizon of %u positions: retrying at %u\n", - __func__, cc.c0, span, lctx.opt_walk_horizon); - lctx.opt_alloc_failed.store(false); - lctx.opt_stop_requested.store(false); - if (!rewind()) { - ok = false; - break; + const size_t cp = recr != nullptr ? (j - run.j0) - (j - run.j0) % stride : j - run.j0; + const auto back = [&]() { return rewind(chunks[run.j0 + cp].c0, snaps[cp], c.c0); }; + if (!back()) { + return; + } + const bool state_surrogate = recr != nullptr && j + 1 < run.j1 && acc.state_ready; + // THE HORIZON FITS THE DEVICE. Its gradient taps are the chunk graph's largest growth + // with the run (layers x horizon x (k + v) floats, twice with their gradient); the graph + // preflight refuses a chunk that does not fit before allocating anything, so a memory + // refusal halves the horizon (whole chunks) and the chunk is tried again; every later + // chunk keeps the smaller one. The forward still attends to the whole window and the + // recurrent state is carried exactly; only attention gradient past the horizon is + // dropped. A refusal at a horizon of one chunk, or for a node the device cannot run, stands. + llama_walk_chunk cc = c; + bool ok; + for (;;) { + cc.grad_from = std::max(r0, llama_walk_plan::grad_from(cc.c0, lctx.opt_walk_horizon)); + set_flags(true, cc.grad_from, c.surrogate, state_surrogate); + xc = &cc; + ok = train_chunk(cc.c0, cc.n); + xc = nullptr; + set_flags(false, 0, false, false); + const uint32_t span = cc.c0 - cc.grad_from; + if (ok || !lctx.opt_failure.empty() || !lctx.opt_alloc_failed.load() || span <= n_batch) { + break; + } + lctx.opt_walk_horizon = std::max(n_batch, (span/2)/n_batch*n_batch); + LLAMA_LOG_WARN("%s: the chunk at %u did not fit with a gradient horizon of %u positions: retrying at %u\n", + __func__, cc.c0, span, lctx.opt_walk_horizon); + lctx.opt_alloc_failed.store(false); + lctx.opt_stop_requested.store(false); + if (!back()) { + ok = false; + break; + } + } + lctx.opt_walk_horizon_used = lctx.opt_walk_horizon; + if (!ok) { + return; } } - lctx.opt_walk_horizon_used = lctx.opt_walk_horizon; - if (!ok) { - break; + acc.state_ready = false; + // the run restored to her stable self: its cells and state as served + if (!restore(run, snaps[0])) { + return; } + p = r1; } - acc.state_ready = false; } -// THE PLAIN WALK: each labelled chunk steps on its own gradient, then joins the context +// a run's cells and recurrent state decoded again under her stable set, from the checkpoint at +// its start (her stable state, since everything before the run is) +bool llama_context::opt_walk::restore(const llama_walk_run & run, const std::vector & snap) { + const uint32_t r0 = plan.begin(run); + const uint32_t r1 = plan.end(run); + return rewind(r0, snap, r0) && decode_span(r0, r1, /*shadow =*/ false); +} + +// THE PLAIN WALK: each labelled chunk steps on its own gradient, then joins the cache under the +// shadow as it now is, so the run's next chunk reads it; the run is restored when it is done void llama_context::opt_walk::run_plain() { - const int64_t last_label = plan.last_label; - for (const llama_walk_chunk & c : plan.chunks) { - if (lctx.opt_stop_requested.load(std::memory_order_relaxed)) { + uint32_t p = 0; + for (const llama_walk_run & run : plan.runs) { + if (!decode_span(p, plan.begin(run), /*shadow =*/ false)) { return; } - const uint32_t c0 = c.c0; - const uint32_t c1 = c.c0 + c.n; - if (c.n_labels == 0) { - if (!decode_span(c0, c1)) { + const std::vector at_run = checkpoint(); + for (size_t j = run.j0; j < run.j1; ++j) { + const llama_walk_chunk & c = plan.chunks[j]; + if (!yield_point()) { return; } - continue; - } - // The training forward advances a recurrent state in place, and a recurrent memory - // cannot remove a partial range: snapshot the state to the scratch sequence first - // (Fable), restore it after, then the decode below advances it from where the context - // left it. The cache's own copy-on-write keeps it on the device; the state is per - // sequence and small. - const bool snapshot = recr != nullptr && c1 <= (uint32_t) last_label; - if (!yield_point()) { - return; - } - if (snapshot) { - recr->seq_cp(0, 1, -1, -1); - } - if (!train_chunk(c0, c.n)) { - return; - } - if (snapshot) { - recr->seq_rm(0, -1, -1); - recr->seq_cp(1, 0, -1, -1); - recr->seq_rm(1, -1, -1); - } - // the chunk joins the context under the adapter as it now is: the training graph never - // writes the cache, so its cells hold nothing until this decode - if (c1 <= (uint32_t) last_label) { - // drop the chunk's empty cells from the attention memory, then decode them for real - const bool removed = recr != nullptr ? (attn == nullptr || attn->seq_rm(0, c0, -1)) : lctx.memory->seq_rm(0, c0, -1); - if (!removed) { - LLAMA_LOG_ERROR("%s: could not remove the trained chunk [%u, %u) from the cache before decoding it\n", __func__, c0, c1); - lctx.opt_stop_requested.store(true); + // the training forward advances a recurrent state in place and may write the chunk's + // cells: back to the chunk's start, then decode it under the stepped shadow + const std::vector at_chunk = j + 1 < run.j1 ? checkpoint() : std::vector(); + if (!train_chunk(c.c0, c.n)) { return; } - if (!decode_span(c0, c1)) { + if (j + 1 < run.j1 && !rewind(c.c0, at_chunk, c.c0 + c.n)) { lctx.opt_stop_requested.store(true); return; } } + if (!restore(run, at_run)) { + lctx.opt_stop_requested.store(true); + return; + } + p = plan.end(run); } } @@ -604,6 +639,10 @@ void llama_context::opt_epoch_iter( int64_t ndata_in_loop, int64_t t_loop_start) { GGML_ASSERT(opt_ctx); + if (this->n_ubatch() != std::min(this->n_batch(), opt_n_ctx_train)) { + opt_fail("the walk needs one ubatch per batch (a chunk's backward needs all of its own K/V in one graph): set -ub = -b"); + return; + } const uint32_t n_ctx = opt_n_ctx_train; const uint32_t n_batch = std::min(this->n_batch(), n_ctx); @@ -625,12 +664,6 @@ void llama_context::opt_epoch_iter( walk.run_exact(); return; } - // the plain walk snapshots a recurrent state into a scratch sequence - if (walk.recr != nullptr && cparams.n_seq_max < 2) { - LLAMA_LOG_ERROR("%s: a recurrent model trains with a scratch sequence for its state snapshot: create the training context with n_seq_max >= 2\n", __func__); - opt_stop_requested.store(true); - return; - } walk.run_plain(); } @@ -763,6 +796,137 @@ const char * llama_opt_failure(struct llama_context * ctx) { return ctx->opt_failure.c_str(); } +bool llama_context::opt_learn(llama_seq_id seq, llama_pos n_cached, const llama_token * tokens, const llama_token * labels, + int32_t n, ggml_opt_result_t result) { + if (opt_ctx == nullptr) { + opt_fail("llama_opt_learn needs llama_opt_init on this context first"); + return false; + } + // each call answers for itself: a refusal before it does not refuse this one + opt_alloc_failed.store(false); + opt_stop_requested.store(false); + opt_failure.clear(); + if (n <= 0 || n_cached < 0 || tokens == nullptr || labels == nullptr) { + opt_fail("llama_opt_learn: an empty reply, a negative served prefix, or no tokens / labels"); + return false; + } + if ((int64_t) n_cached + n > (int64_t) cparams.n_ctx_seq) { + opt_fail("llama_opt_learn: the reply ends at " + std::to_string((int64_t) n_cached + n) + ", past the sequence's context of " + + std::to_string(cparams.n_ctx_seq) + " positions"); + return false; + } + // the parts of the memory: a recurrent state cannot be cut back, attention cells can + llama_memory_recurrent * recr = nullptr; + llama_memory_i * attn = memory.get(); + if (auto * hybrid = dynamic_cast(memory.get())) { + recr = hybrid->get_mem_recr(); + attn = hybrid->get_mem_attn(); + } else if (auto * r = dynamic_cast(memory.get())) { + recr = r; + attn = nullptr; + } + // refusals before anything is touched + const llama_pos held = memory->seq_pos_max(seq) + 1; + if (held < n_cached) { + opt_fail("llama_opt_learn: sequence " + std::to_string(seq) + " holds " + std::to_string(held) + + " positions, fewer than the " + std::to_string(n_cached) + " served before the reply"); + return false; + } + if (recr != nullptr && recr->seq_pos_max(seq) + 1 != n_cached) { + opt_fail("llama_opt_learn: the recurrent state of sequence " + std::to_string(seq) + " ends at " + + std::to_string(recr->seq_pos_max(seq) + 1) + ", not at the served prefix " + std::to_string(n_cached) + + ": restore the served-prefix checkpoint first"); + return false; + } + if (recr != nullptr && recr->rollback_pending(seq)) { + opt_fail("llama_opt_learn: sequence " + std::to_string(seq) + " has a pending recurrent rollback: finish or drop the draft first"); + return false; + } + if (attn != nullptr && !attn->seq_rm(seq, n_cached, -1)) { + opt_fail("llama_opt_learn: could not drop sequence " + std::to_string(seq) + "'s cells past " + std::to_string(n_cached)); + return false; + } + // her stable recurrent state at the served prefix, for the restore + std::vector snap; + if (recr != nullptr) { + synchronize(); + snap.resize(state_seq_get_size(seq, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)); + GGML_ASSERT(state_seq_get_data(seq, snap.data(), snap.size(), LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == snap.size()); + } + const uint32_t nb = n_batch(); + llama_batch batch = llama_batch_init(nb, 0, 1); + // [c0, c0 + m) of the reply into the batch; with labels, the trainable layer sees those rows + auto fill = [&](int32_t c0, int32_t m, bool labelled) { + batch.n_tokens = m; + for (int32_t i = 0; i < m; ++i) { + batch.token [i] = tokens[c0 + i]; + batch.pos [i] = n_cached + c0 + i; + batch.n_seq_id[i] = 1; + batch.seq_id [i][0] = seq; + batch.label [i] = labelled ? labels[c0 + i] : -1; + batch.logits [i] = labelled && labels[c0 + i] >= 0; + } + if (!labelled) { + batch.logits[m - 1] = true; + } + }; + // train up to the last labelled row: rows after it cannot change the loss, and in one batch an + // unlabelled row after the run is refused (a run ends its sequence's rows); the restore below + // still decodes all n + int32_t n_train = n; + while (n_train > 0 && labels[n_train - 1] < 0) { + --n_train; + } + bool ok = true; + opt_result_active = result; + opt_outputs_internal = true; + opt_learn_labels = 0; + opt_learn_labels_done = 0; + for (int32_t i = 0; i < n_train; ++i) { + opt_learn_labels += labels[i] >= 0; + } + for (int32_t c0 = 0; c0 < n_train && ok; c0 += (int32_t) nb) { + const int32_t m = std::min((int32_t) nb, n_train - c0); + fill(c0, m, true); + ok = decode(batch) == 0; + if (!ok && !opt_alloc_failed.load()) { + opt_fail("llama_opt_learn: the reply's training batch at " + std::to_string(n_cached + c0) + " failed to decode"); + } + } + opt_result_active = nullptr; + opt_outputs_internal = false; + opt_learn_labels = 0; + opt_learn_labels_done = 0; + // RESTORE: the reply's cells and the recurrent state decoded again under her stable set + bool restored = attn == nullptr || attn->seq_rm(seq, n_cached, -1); + if (recr != nullptr) { + synchronize(); + restored = restored && state_seq_set_data(seq, snap.data(), snap.size(), LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == snap.size(); + } + for (int32_t c0 = 0; c0 < n && restored; c0 += (int32_t) nb) { + const int32_t m = std::min((int32_t) nb, n - c0); + fill(c0, m, false); + restored = decode(batch) == 0; + } + llama_batch_free(batch); + if (!restored) { + opt_fail("llama_opt_learn: sequence " + std::to_string(seq) + " could not be restored after training: re-prefill it"); + return false; + } + return ok; +} + +bool llama_opt_learn( + struct llama_context * ctx, + llama_seq_id seq, + llama_pos n_cached, + const llama_token * tokens, + const llama_token * labels, + int32_t n, + ggml_opt_result_t result) { + return ctx->opt_learn(seq, n_cached, tokens, labels, n, result); +} + void llama_opt_epoch( struct llama_context * ctx, ggml_opt_dataset_t dataset, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a296f0c2cbdb..7f3b1cc2e579 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1334,11 +1334,22 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll auto * res = gf_res_prev.get(); auto * gf = res->get_gf(); + // a ubatch with labelled rows is a training step on this context's trainable adapter + // (llama_opt_init): the same graph with the extras on, the shadow's masked term + // (graph_params) and the loss and backward (opt_ubatch). It is built fresh, and + // cparams.training is part of the reuse key, so a serving graph never reuses it. + bool train = false; + if (opt_ctx != nullptr) { + train = !opt_forward_only && ubatch.n_labels() > 0; + cparams.training = train; + cparams.walk_active = opt_walk_active != nullptr; + } + // the new graph parameters // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters const auto gparams = graph_params(res, ubatch, mctx, gtype); - if (!graph_reuse_disable && res->can_reuse(gparams)) { + if (!train && !graph_reuse_disable && res->can_reuse(gparams)) { //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__); // with pipeline parallelism, the previous graph_compute_async may still be running @@ -1367,6 +1378,15 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll return nullptr; } + if (train) { + if (!opt_ubatch(res, gf, ubatch)) { + ret = GGML_STATUS_FAILED; + return nullptr; + } + ret = GGML_STATUS_SUCCESS; + return res; + } + if (!ggml_backend_sched_alloc_graph(sched.get(), gf)) { LLAMA_LOG_ERROR("%s: failed to allocate graph\n", __func__); ret = GGML_STATUS_ALLOC_FAILED; @@ -1698,6 +1718,24 @@ int llama_context::decode(const llama_batch & batch_inp) { return -1; } + if (balloc->get_n_labels() > 0) { + if (!model.output) { + LLAMA_LOG_ERROR("%s: labelled training requires a language-model output weight\n", __func__); + return -1; + } + if (cparams.pooling_type != LLAMA_POOLING_TYPE_NONE) { + LLAMA_LOG_ERROR("%s: labelled training requires a token-wise output head, not pooled outputs\n", __func__); + return -1; + } + if (!opt_ctx) { + LLAMA_LOG_ERROR("%s: labelled rows require an initialized optimizer\n", __func__); + return -1; + } + if (const char * why = memory->training_error(balloc->get_batch(), cparams.flash_attn)) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, why); + return -1; + } + } const uint32_t n_tokens_all = balloc->get_n_tokens(); const uint32_t n_outputs_all = balloc->get_n_outputs(); @@ -1780,8 +1818,9 @@ int llama_context::decode(const llama_batch & batch_inp) { break; } - // reserve output buffer - if (output_reserve(n_outputs_all) < n_outputs_all) { + // reserve output buffer (none for outputs that only feed a training loss) + const bool host_outputs = !opt_outputs_internal && opt_walk_active == nullptr; + if (output_reserve(host_outputs ? n_outputs_all : 0) < (host_outputs ? n_outputs_all : 0)) { LLAMA_LOG_ERROR("%s: could not reserve space for batch with %d outputs\n", __func__, n_outputs_all); return -2; }; @@ -1862,7 +1901,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } // extract logits - if (logits.data && t_logits && n_outputs > 0 && needs_raw_logits(ubatch, sampling.samplers)) { + if (logits.data && t_logits && n_outputs > 0 && host_outputs && needs_raw_logits(ubatch, sampling.samplers)) { ggml_backend_t backend_res = ggml_backend_sched_get_tensor_backend(sched.get(), t_logits); GGML_ASSERT(backend_res != nullptr); GGML_ASSERT(logits.data != nullptr); @@ -1877,7 +1916,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } // extract embeddings - if (embd.data && t_embd && n_outputs > 0) { + if (embd.data && t_embd && n_outputs > 0 && host_outputs) { ggml_backend_t backend_embd = ggml_backend_sched_get_tensor_backend(sched.get(), t_embd); GGML_ASSERT(backend_embd != nullptr); @@ -1945,7 +1984,7 @@ int llama_context::decode(const llama_batch & batch_inp) { const int64_t n_rows = masked ? n_outputs : (int64_t) ubatch.n_tokens; const int64_t offset = masked ? n_outputs_prev : n_tokens_prev; - if (embd_nextn.data && t_h_nextn && n_rows > 0 && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { + if (embd_nextn.data && t_h_nextn && n_rows > 0 && host_outputs && cparams.pooling_type == LLAMA_POOLING_TYPE_NONE) { ggml_backend_t backend_h = ggml_backend_sched_get_tensor_backend(sched.get(), t_h_nextn); GGML_ASSERT(backend_h != nullptr); @@ -1972,7 +2011,7 @@ int llama_context::decode(const llama_batch & batch_inp) { } while (mctx->next()); // set to total number of outputs in the batch, for use in llama_get_logits_ith - n_outputs = n_outputs_all; + n_outputs = host_outputs ? n_outputs_all : 0; // set output mappings if (n_outputs > 0) { @@ -2487,10 +2526,12 @@ llm_graph_params llama_context::graph_params( /*.backend_cpu =*/ backend_cpu, /*.cvec =*/ cvec.get(), /*.loras =*/ loras.get(), + /*.shadow =*/ opt_shadow, + /*.n_shadow_rows =*/ opt_shadow != nullptr ? ubatch.n_labels() : 0, /*.mctx =*/ mctx, /*.cross =*/ &cross, /*.samplers =*/ sampling.samplers, - /*.n_outputs =*/ n_outputs, + /*.n_outputs =*/ (cparams.training || (opt_shadow && ubatch.n_labels() > 0)) ? ubatch.n_tokens : n_outputs, /*.cb =*/ graph_get_cb(), /*.res =*/ res, }; diff --git a/src/llama-context.h b/src/llama-context.h index 2f3fc48b6efd..c605744e29b1 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -215,6 +215,9 @@ struct llama_context { // the largest training graph the allocation preflight measured on a GPU device (bytes) size_t opt_graph_bytes() const; + bool opt_learn(llama_seq_id seq, llama_pos n_cached, const llama_token * tokens, const llama_token * labels, int32_t n, + ggml_opt_result_t result); + void opt_epoch( ggml_opt_dataset_t dataset, ggml_opt_result_t result_train, @@ -244,6 +247,13 @@ struct llama_context { // the step that failed, for the log void opt_fail(std::string why, const std::string & what = ""); + // THE EXTRAS A LABELLED UBATCH SWITCHES ON (llama-context-opt.cpp). process_ubatch runs a + // ubatch whose rows carry labels through ggml_opt instead of graph_compute: the same graph, + // plus the loss over the labelled rows, the backward and, at a period's end, the optimizer + // step. A walk driving the decode adds its taps to that graph. + bool opt_ubatch(llm_graph_result * res, ggml_cgraph * gf, const llama_ubatch & ubatch); + bool opt_fa_back_native() const; + // // output // @@ -388,6 +398,24 @@ struct llama_context { uint32_t opt_walk_horizon_req = 0; // llama_opt_params::walk_horizon, as requested uint32_t opt_walk_horizon = 0; // this window's: starts at the request, shrinks to fit the device + // the adapter being trained (llama_opt_params::adapter): applied to labelled rows only + // (llm_graph_params::shadow); null when the base itself trains + const llama_adapter_lora * opt_shadow = nullptr; + + // set by a walk for the decode it drives: whose taps join the graph, where the loss lands, + // whether its labelled rows only fill the cache (the shadow on, no loss, no backward), and + // that the logits rows are the loss's alone (no host copy) + opt_walk * opt_walk_active = nullptr; + ggml_opt_result_t opt_result_active = nullptr; + bool opt_forward_only = false; + // the batch's outputs feed only the in-graph loss (llama_opt_learn; a walk): no host logits are + // reserved, copied or mapped for them (a 512-row chunk at a 248k vocabulary is 485 MiB of host RAM) + bool opt_outputs_internal = false; + // llama_opt_learn's one step: its labelled rows in all, and those trained so far (the ubatch that + // completes them takes the step, each weighted by its share) + int64_t opt_learn_labels = 0; + int64_t opt_learn_labels_done = 0; + ggml_threadpool_t threadpool = nullptr; ggml_threadpool_t threadpool_batch = nullptr; diff --git a/src/llama-cparams.h b/src/llama-cparams.h index f32f7f465b6a..9857c99f066d 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -18,14 +18,20 @@ struct llama_cparams { // gradient (positions [walk_grad_from, n_past) attend through a gradient tap; [0, // walk_grad_from) stay constants). walk_surrogate: later chunks accumulated a gradient on // this chunk's own K/V, so its graph adds the surrogate term that carries it in. - // walk_acc_k / walk_acc_v: per layer, dL/d(cached K / V) of the window, [n_embd_*_gqa, - // window] F32 in host memory (null for a recurrent layer); a chunk's backward adds its prefix's - // share into rows [walk_grad_from, n_past) and its surrogate reads its own rows. + // walk_acc_k / walk_acc_v: per layer, dL/d(cached K / V) of the run being trained, [n_embd_*_gqa, + // the run's positions] F32 in host memory (null for a recurrent layer), row 0 at position + // walk_acc_base; a chunk's backward adds its prefix's share into rows [walk_grad_from, n_past) + // and its surrogate reads its own rows. + // the attention backward is one op (ggml_flash_attn_back) on every device that would run it + // (llama_context::opt_init probes them), else the composed tiles + bool fa_back_native = false; bool walk_exact = false; + bool walk_active = false; uint32_t walk_grad_from = 0; bool walk_surrogate = false; ggml_tensor * const * walk_acc_k = nullptr; ggml_tensor * const * walk_acc_v = nullptr; + uint32_t walk_acc_base = 0; // a later chunk's gradient sits on the recurrent state this chunk leaves bool walk_state_surrogate = false; uint32_t n_ctx_seq; // context for a single sequence diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9e9071876ab9..d1571b907ec7 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -5,6 +5,7 @@ #include "llama-batch.h" #include "llama-cparams.h" #include "llama-sampler.h" +#include "ggml-opt.h" #include "llama-kv-cache.h" #include "llama-kv-cache-iswa.h" @@ -223,6 +224,37 @@ void llm_graph_input_out_ids::set_input(const llama_ubatch * ubatch) { } } +void llm_graph_input_training_rows::set_input(const llama_ubatch * ubatch) { + const auto map = ubatch->rows(); + if (public_rows) { + GGML_ASSERT(public_rows->ne[0] == (int64_t) map.public_rows.size()); + ggml_backend_tensor_set(public_rows, map.public_rows.data(), 0, ggml_nbytes(public_rows)); + } + if (label_rows) { + GGML_ASSERT(label_rows->ne[0] == (int64_t) map.label_rows.size()); + ggml_backend_tensor_set(label_rows, map.label_rows.data(), 0, ggml_nbytes(label_rows)); + } + if (shadow_rows && shadow_rows->buffer) { + ggml_backend_tensor_set(shadow_rows, map.label_rows.data(), 0, ggml_nbytes(shadow_rows)); + } + if (public_shadow_rows && public_shadow_rows->buffer) { + std::vector local; + for (size_t i = 0; i < map.public_rows.size(); ++i) { + if (ubatch->label[map.public_rows[i]] >= 0) { + local.push_back((int32_t) i); + } + } + ggml_backend_tensor_set(public_shadow_rows, local.data(), 0, ggml_nbytes(public_shadow_rows)); + } + for (auto * rows : slice_shadow_rows) { + if (rows->buffer) { + std::vector local(rows->ne[0]); + std::iota(local.begin(), local.end(), 0); + ggml_backend_tensor_set(rows, local.data(), 0, ggml_nbytes(rows)); + } + } +} + bool llm_graph_input_out_ids::can_reuse(const llm_graph_params & params) { bool res = true; @@ -1333,6 +1365,7 @@ void llm_graph_result::reset() { t_inp_tokens = nullptr; t_inp_embd = nullptr; t_logits = nullptr; + t_loss_slices.clear(); t_embd = nullptr; t_embd_pooled = nullptr; t_h_nextn = nullptr; @@ -1501,6 +1534,8 @@ llm_graph_context::llm_graph_context(const llm_graph_params & params) : backend_cpu (params.backend_cpu), cvec (params.cvec), loras (params.loras), + shadow (params.shadow), + n_shadow_rows (params.n_shadow_rows), mctx (params.mctx), cross (params.cross), samplers (params.samplers), @@ -1537,7 +1572,7 @@ ggml_tensor * llm_graph_context::build_lora_mm( for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); - if (lw == nullptr) { + if (lw == nullptr || !lora_on(lora.first)) { continue; } @@ -1550,7 +1585,7 @@ ggml_tensor * llm_graph_context::build_lora_mm( ); ab_cur = ggml_scale(ctx0, ab_cur, scale); - res = ggml_add(ctx0, res, ab_cur); + res = lora.first == shadow ? shadow_add(res, ab_cur) : ggml_add(ctx0, res, ab_cur); } return res; @@ -1573,7 +1608,7 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( } for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(w); - if (lw == nullptr) { + if (lw == nullptr || !lora_on(lora.first)) { continue; } @@ -1588,7 +1623,7 @@ ggml_tensor * llm_graph_context::build_lora_mm_id( ); ab_cur = ggml_scale(ctx0, ab_cur, scale); - res = ggml_add(ctx0, res, ab_cur); + res = lora.first == shadow ? shadow_add(res, ab_cur) : ggml_add(ctx0, res, ab_cur); } return res; @@ -1731,7 +1766,7 @@ ggml_tensor * llm_graph_context::build_ffn( return false; } for (const auto & lora : *loras) { - if (lora.first->get_weight(w) != nullptr) { + if (lora.first->get_weight(w) != nullptr && lora_on(lora.first)) { return true; } } @@ -2352,7 +2387,7 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const { // apply lora for embedding tokens if needed for (const auto & lora : *loras) { llama_adapter_lora_weight * lw = lora.first->get_weight(tok_embd); - if (lw == nullptr) { + if (lw == nullptr || !lora_on(lora.first)) { continue; } @@ -2364,7 +2399,7 @@ ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const { ggml_get_rows(ctx0, lw->a, inp->tokens) ), scale); - cur = ggml_add(ctx0, cur, inpL_delta); + cur = lora.first == shadow ? shadow_add(cur, inpL_delta) : ggml_add(ctx0, cur, inpL_delta); } if (n_embd_inp != n_embd) { @@ -2429,6 +2464,93 @@ ggml_tensor * llm_graph_context::build_inp_train_rows() const { return inp_train_rows; } +bool llm_graph_context::lora_on(const llama_adapter_lora * adapter) const { + return adapter != shadow || n_shadow_rows > 0; +} + +llm_graph_input_training_rows * llm_graph_context::training_rows() const { + if (!inp_training_rows) { + auto inp = std::make_unique(); + inp_training_rows = inp.get(); + res->add_input(std::move(inp)); + } + return inp_training_rows; +} + +ggml_tensor * llm_graph_context::shadow_add(ggml_tensor * base, ggml_tensor * term) const { + if (training_rows()->shadow_rows == nullptr) { + auto * inp = training_rows(); + inp->shadow_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_shadow_rows); + ggml_set_input(inp->shadow_rows); + ggml_set_name(inp->shadow_rows, "label_rows"); + } + // one row per token: the tokens are base's outermost dimension + GGML_ASSERT(ggml_is_contiguous(base) && (n_tokens == 1 || base->ne[ggml_n_dims(base) - 1] == n_tokens)); + const int64_t n = ggml_nelements(base) / n_tokens; + ggml_tensor * b2 = ggml_n_dims(base) <= 2 ? base : ggml_reshape_2d(ctx0, base, n, n_tokens); + ggml_tensor * t2 = ggml_n_dims(term) <= 2 ? term : ggml_reshape_2d(ctx0, ggml_cont(ctx0, term), n, n_tokens); + ggml_tensor * sel = ggml_add(ctx0, ggml_get_rows(ctx0, b2, training_rows()->shadow_rows), ggml_get_rows(ctx0, t2, training_rows()->shadow_rows)); + ggml_tensor * updated = ggml_set_rows(ctx0, ggml_dup(ctx0, b2), sel, training_rows()->shadow_rows); + return ggml_n_dims(base) <= 2 ? updated : ggml_reshape(ctx0, updated, base); +} + +bool llm_graph_context::build_training_head() const { + if (!cparams.training && n_shadow_rows == 0) { + return true; + } + if (!res->t_embd || !res->t_logits || pooling_type != LLAMA_POOLING_TYPE_NONE) { + return false; + } + const auto map = ubatch.rows(); + auto * input = training_rows(); + if (cparams.training) { + input->label_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, map.label_rows.size()); + ggml_set_input(input->label_rows); + } + if (!map.public_rows.empty()) { + input->public_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, map.public_rows.size()); + ggml_set_input(input->public_rows); + } + // Preserve the existing training head's allocation budget. + constexpr size_t head_slice_bytes = 64u*1024*1024; + const int64_t rows = std::max(1, head_slice_bytes / (sizeof(float) * res->t_logits->ne[0])); + const int64_t loss_rows = cparams.training ? map.label_rows.size() : n_tokens; + res->t_loss_slices.resize((loss_rows + rows - 1) / rows); + if (input->shadow_rows) { + size_t public_labels = 0; + for (int32_t row : map.public_rows) { + public_labels += ubatch.label[row] >= 0; + } + if (public_labels) { + input->public_shadow_rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, public_labels); + ggml_set_input(input->public_shadow_rows); + } + for (int64_t r0 = 0; r0 < loss_rows; r0 += rows) { + auto * local = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, std::min(rows, loss_rows - r0)); + ggml_set_input(local); + input->slice_shadow_rows.push_back(local); + } + } + ggml_tensor * public_logits = nullptr; + const int count = ggml_opt_project_head(ctx0, gf, res->t_embd, res->t_logits, input->shadow_rows, + input->public_rows, input->label_rows, input->public_shadow_rows, + input->slice_shadow_rows.empty() ? nullptr : input->slice_shadow_rows.data(), rows, &public_logits, + res->t_loss_slices.data(), res->t_loss_slices.size()); + if (!count) { + res->t_loss_slices.clear(); + return false; + } + res->t_logits = public_logits; + if (cparams.embeddings && input->public_rows) { + res->t_embd = ggml_get_rows(ctx0, res->t_embd, input->public_rows); + ggml_build_forward_expand(gf, res->t_embd); + } + if (cparams.embeddings_nextn && cparams.embeddings_nextn_masked && input->public_rows) { + res->t_h_nextn = ggml_get_rows(ctx0, res->t_h_nextn, input->public_rows); + ggml_build_forward_expand(gf, res->t_h_nextn); + } + return true; +} ggml_tensor * llm_graph_context::build_inp_pos() const { auto inp = std::make_unique(hparams.n_pos_per_embd()); @@ -2610,15 +2732,26 @@ ggml_tensor * llm_graph_context::build_attn_mha( v = ggml_transpose(ctx0, v); } - // this can happen when KV cache is not used (e.g. an embedding model with non-causal attn). - // A training graph keeps its chunk's own K/V at F32 (every flash-attention backend takes - // F32 K/V): an F16 cast there rounded the trained K/V and moved the step (test-walk-exact, - // 2026-10-11: cosine 0.9969 against the one-graph reference at an f16 cache, 1.0000 at F32) - if (!cparams.training && k->type == GGML_TYPE_F32) { + // Direct decode uses serving precision for every row. The walk keeps its chunk's F32 K/V + // for its existing gradient comparison against a single graph. + if ((!cparams.training || !cparams.walk_active) && k->type == GGML_TYPE_F32) { k = ggml_cast(ctx0, k, GGML_TYPE_F16); } - if (!cparams.training && v->type == GGML_TYPE_F32) { + if ((!cparams.training || !cparams.walk_active) && v->type == GGML_TYPE_F32) { + v = ggml_cast(ctx0, v, GGML_TYPE_F16); + } + + // A training graph reads a quantized cache as F16 through a node of its own, as the kernel + // would dequantize it anyway: inside the attention the F16 copy lives as long as the + // attention's output, which a training graph keeps for the backward, one window-sized copy + // per layer; the backward reads the cast's quantized source (ggml_flash_attn_ext_backward), + // so this copy dies with the forward (the M5, qwen3.5-4b, q8_0, window 32768: 1,888 MiB + // against 853 at an f16 cache) + if (cparams.training && cparams.walk_active && ggml_is_quantized(k->type)) { + k = ggml_cast(ctx0, k, GGML_TYPE_F16); + } + if (cparams.training && cparams.walk_active && ggml_is_quantized(v->type)) { v = ggml_cast(ctx0, v, GGML_TYPE_F16); } @@ -2629,7 +2762,12 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_flash_attn_ext_add_sinks(cur, sinks); ggml_flash_attn_ext_set_prec (cur, GGML_PREC_F32); if (cparams.training) { - ggml_flash_attn_ext_set_back_tile(cur, 0); // the backward's default tile + ggml_flash_attn_ext_set_back_tile(cur, cparams.fa_back_native ? GGML_FA_BACK_NATIVE : 0); + } + if (walk_grad_sinks[1] != nullptr) { + // the in-place training graph's K/V gradient goes to the run's rows only + ggml_flash_attn_ext_set_grad_sinks(cur, walk_grad_sinks[0], walk_grad_sinks[1], walk_grad_sinks[2], walk_grad_sinks[3]); + walk_grad_sinks = {}; } if (v_mla) { @@ -2809,8 +2947,8 @@ static std::unique_ptr build_attn_inp_kv_impl( { GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA"); - inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch); - inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -2868,48 +3006,66 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * q = q_cur; ggml_tensor * k; ggml_tensor * v; - if (cparams.training && uses_flash_attn(kq_b, sinks, v_mla)) { + if (cparams.training && cparams.walk_active && uses_flash_attn(kq_b, sinks, v_mla)) { // THE TRAINING GRAPH IS THE SERVING GRAPH. This chunk writes its K/V into the cache // exactly as serving does (SET_ROWS, the cache's own type) and attention reads the - // cache in place: no copy of the context, at any window. The walk decoded the 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. + // cache in place: no copy of the context, at any window. The walk decoded the run + // into the cache under the same adapter, so the write is idempotent. + // Only the run's rows carry a K/V gradient, and they are the attention's grad sinks + // (ggml_flash_attn_ext_set_grad_sinks): the chunk's own K/V and, in the exact walk, the + // run's positions before it ([g0, n_past), read from the cache as F32 and TAPPED: their + // gradient adds into the run's host accumulator, row 0 at walk_acc_base). The backward + // computes dK/dV from g0 on and nothing for the window before it, a constant (her served + // prefix): no window-sized gradient exists. What later chunks put on this chunk's own rows + // reaches it through the surrogate. A quantized or f16 cache trains on the values serving + // reads; its gradient passes straight through to the F32 K/V. const int64_t n_tokens = k_cur->ne[2]; const int64_t n_past = ubatch.pos[0]; const int64_t n_win = n_past + n_tokens; ggml_tensor * rows = build_inp_train_rows(); - llm_graph_result::walk_layer io = { il, (uint32_t) n_past }; + const int64_t base = cparams.walk_acc_base; + const int64_t g0 = std::min(std::max(cparams.walk_grad_from, base), n_past); + llm_graph_result::walk_layer io = { il, (uint32_t) n_past, (uint32_t) g0 }; // the window's cells [0, n_win) of one layer's cache with this chunk's rows written - auto write = [&](ggml_tensor * all, ggml_tensor * cur, ggml_tensor * acc, ggml_tensor *& tap, ggml_tensor *& tap_acc) { + auto write = [&](ggml_tensor * all, ggml_tensor * cur) { const int64_t n_embd = cur->ne[0]*cur->ne[1]; GGML_ASSERT(all->ne[3] == 1 && "a training context reads one stream"); GGML_ASSERT(ggml_row_size(cur->type, cur->ne[0]) == cur->nb[1]); ggml_tensor * win = ggml_view_2d(ctx0, all, n_embd, n_win, all->nb[2], 0); ggml_tensor * w = ggml_set_rows(ctx0, win, ggml_view_2d(ctx0, cur, n_embd, n_tokens, cur->nb[2], 0), rows); - if (acc != nullptr) { - 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); }; + // the cached positions [p0, p1) as F32 [D, Hkv, p1 - p0] + auto span = [&](ggml_tensor * all, int64_t p0, int64_t p1) { + return ggml_cast(ctx0, ggml_view_3d(ctx0, all, all->ne[0], all->ne[1], p1 - p0, all->nb[1], all->nb[2], p0*all->nb[2]), + GGML_TYPE_F32); + }; + ggml_tensor * k_all = mctx_cur->get_k(ctx0, il); ggml_tensor * v_all = mctx_cur->get_v(ctx0, il); GGML_ASSERT(v_all->nb[1] <= v_all->nb[2] && "a flash-attention cache holds V untransposed"); - const bool tapped = cparams.walk_exact; - k = write(mctx_cur->get_k(ctx0, il), k_cur, tapped ? cparams.walk_acc_k[il] : nullptr, io.tk, io.ak); - v = write(v_all, v_cur, tapped ? cparams.walk_acc_v[il] : nullptr, io.tv, io.av); - if (tapped) { + GGML_ASSERT(k_cur->type == GGML_TYPE_F32 && v_cur->type == GGML_TYPE_F32); + k = write(k_all, k_cur); + v = write(v_all, v_cur); + if (cparams.walk_exact && n_past > g0) { + io.tk = span(k_all, g0, n_past); + io.tv = span(v_all, g0, n_past); + ggml_format_name(io.tk, "walk_tk-%d", il); + ggml_format_name(io.tv, "walk_tv-%d", il); + io.ak = walk_acc_rows(cparams.walk_acc_k[il], io.tk, g0 - base); + io.av = walk_acc_rows(cparams.walk_acc_v[il], io.tv, g0 - base); + } + walk_grad_sinks = { io.tk, k_cur, io.tv, v_cur }; + if (cparams.walk_exact) { + if (cparams.walk_surrogate) { + build_walk_surrogate(il, k_cur, v_cur, n_past - base); + } 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) { + } else if (cparams.training && cparams.walk_active) { // EXPLICIT ATTENTION (a layer FLASH_ATTN_EXT's backward does not cover yet: ALiBi, a // logit softcap, sinks, a KQ bias, MLA) reads a copy of the context instead. // A backward pass cannot go through the cache: the write is an in-place SET_ROWS @@ -2957,15 +3113,9 @@ ggml_tensor * llm_graph_context::build_attn( // (test-walk-exact, q8_0 and f16 caches): cosine 0.998 against one graph, the same as // at an F32 cache, so training reads the numbers serving runs on (Joel: align the bit // depth to inference). - const int64_t g0 = std::min(cparams.walk_grad_from, n_past); - llm_graph_result::walk_layer io = { il, (uint32_t) n_past }; - ggml_tensor * acc_k = cparams.walk_acc_k[il]; - ggml_tensor * acc_v = cparams.walk_acc_v[il]; - // rows [p0, p0 + n) of a layer's accumulator, shaped like that span's K / V - auto acc_rows = [&](ggml_tensor * acc, const ggml_tensor * like, int64_t p0) { - return ggml_view_3d(ctx0, acc, like->ne[0], like->ne[1], like->ne[2], - like->ne[0]*sizeof(float), acc->nb[1], p0*acc->nb[1]); - }; + const int64_t base = cparams.walk_acc_base; + const int64_t g0 = std::min(std::max(cparams.walk_grad_from, base), n_past); + llm_graph_result::walk_layer io = { il, (uint32_t) n_past, (uint32_t) g0 }; // [const prefix][gradient prefix][this chunk]: the order the cells hold them in std::vector k_parts, v_parts; if (g0 > 0) { @@ -2978,8 +3128,8 @@ ggml_tensor * llm_graph_context::build_attn( io.tv = v_span(g0, n_past); ggml_format_name(io.tk, "walk_tk-%d", il); ggml_format_name(io.tv, "walk_tv-%d", il); - io.ak = acc_rows(acc_k, io.tk, g0); - io.av = acc_rows(acc_v, io.tv, g0); + io.ak = walk_acc_rows(cparams.walk_acc_k[il], io.tk, g0 - base); + io.av = walk_acc_rows(cparams.walk_acc_v[il], io.tv, g0 - base); k_parts.push_back(io.tk); v_parts.push_back(io.tv); } @@ -2992,14 +3142,7 @@ ggml_tensor * llm_graph_context::build_attn( v = ggml_concat(ctx0, v, v_parts[i], 2); } if (cparams.walk_surrogate) { - ggml_tensor * gk = acc_rows(acc_k, k_cur, n_past); - ggml_tensor * gv = acc_rows(acc_v, v_cur, n_past); - ggml_tensor * k32 = k_cur->type == GGML_TYPE_F32 ? k_cur : ggml_cast(ctx0, k_cur, GGML_TYPE_F32); - ggml_tensor * v32 = v_cur->type == GGML_TYPE_F32 ? v_cur : ggml_cast(ctx0, v_cur, GGML_TYPE_F32); - ggml_tensor * term = ggml_add(ctx0, - ggml_sum(ctx0, ggml_mul(ctx0, k32, gk)), - ggml_sum(ctx0, ggml_mul(ctx0, v32, gv))); - res->t_walk_surrogate = res->t_walk_surrogate ? ggml_add(ctx0, res->t_walk_surrogate, term) : term; + build_walk_surrogate(il, k_cur, v_cur, n_past - base); } res->t_walk.push_back(io); } else if (n_past > 0) { @@ -3017,11 +3160,13 @@ ggml_tensor * llm_graph_context::build_attn( const auto & k_idxs = inp->get_k_idxs(); const auto & v_idxs = inp->get_v_idxs(); - ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); - ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + ggml_tensor * written_k = mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il); + ggml_tensor * written_v = mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il); + ggml_build_forward_expand(gf, written_k); + ggml_build_forward_expand(gf, written_v); - k = mctx_cur->get_k(ctx0, il); - v = mctx_cur->get_v(ctx0, il); + k = mctx_cur->get_k(ctx0, il, cparams.training ? written_k : nullptr); + v = mctx_cur->get_v(ctx0, il, cparams.training ? written_v : nullptr); } ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); @@ -3063,7 +3208,7 @@ static std::unique_ptr build_attn_inp_k_impl( { GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA"); - inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3437,14 +3582,14 @@ static std::unique_ptr build_attn_inp_k_dsa_impl( auto inp = std::make_unique(hparams, cparams, mctx_cur); { - inp->self_k_idxs_mla = mctx_cur->get_mla()->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs_mla = mctx_cur->get_mla()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask_mla = build_attn_inp_kq_mask(ctx0, mctx_cur->get_mla(), ubatch, cparams); inp->self_kq_mask_mla_cnv = inp->self_kq_mask_mla; } { - inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs_lid = mctx_cur->get_lid()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); // ensure that mask type matches fused lightning indexer use (requires f16 mask) auto cparams_copy = cparams; @@ -3475,7 +3620,7 @@ llm_graph_input_attn_k_dsa_iswa * llm_graph_context::build_attn_inp_k_dsa_iswa() // build_attn_inp_k_impl rejects SWA caches, so construct the input directly auto inp_swa = std::make_unique(hparams, cparams, mctx_cur->get_swa()); - inp_swa->self_k_idxs = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); + inp_swa->self_k_idxs = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp_swa->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp_swa->self_kq_mask_cnv = inp_swa->self_kq_mask; @@ -3496,8 +3641,8 @@ llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_ { GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache_iswa for SWA"); - inp->self_k_idxs = mctx_base->build_input_k_idxs(ctx0, ubatch); - inp->self_v_idxs = mctx_base->build_input_v_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_base->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp->self_v_idxs = mctx_base->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_base, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3507,7 +3652,7 @@ llm_graph_input_attn_kv_msa * llm_graph_context::build_attn_inp_kv_msa(bool msa_ inp->self_v_rot = mctx_base->build_input_v_rot(ctx0); if (msa_enabled) { - inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs_idx = mctx_idx->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); } return (llm_graph_input_attn_kv_msa *) res->add_input(std::move(inp)); @@ -3522,8 +3667,8 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const auto inp = std::make_unique(hparams, cparams, mctx_cur); { - inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch); - inp->self_v_idxs = mctx_cur->get_base()->build_input_v_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp->self_v_idxs = mctx_cur->get_base()->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3532,8 +3677,8 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const { GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA"); - inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); - inp->self_v_idxs_swa = mctx_cur->get_swa()->build_input_v_idxs(ctx0, ubatch); + inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp->self_v_idxs_swa = mctx_cur->get_swa()->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa; @@ -3554,7 +3699,7 @@ llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { auto inp = std::make_unique(hparams, cparams, mctx_cur); { - inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs = mctx_cur->get_base()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur->get_base(), ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; @@ -3563,7 +3708,7 @@ llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const { { GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE && "Use llama_kv_cache for non-SWA"); - inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch); + inp->self_k_idxs_swa = mctx_cur->get_swa()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, mctx_cur->get_swa(), ubatch, cparams); inp->self_kq_mask_swa_cnv = inp->self_kq_mask_swa; @@ -3690,6 +3835,27 @@ ggml_tensor * llm_graph_context::build_rs( get_state_rows); } +int64_t llm_graph_context::rs_keep() const { + const bool training = cparams.training && cparams.walk_exact; + return cparams.n_rs_seq > 0 && !training ? (int64_t) cparams.n_rs_seq + 1 : 1; +} + +ggml_tensor * llm_graph_context::walk_acc_rows(ggml_tensor * acc, const ggml_tensor * like, int64_t row) const { + return ggml_view_3d(ctx0, acc, like->ne[0], like->ne[1], like->ne[2], + like->ne[0]*sizeof(float), acc->nb[1], row*acc->nb[1]); +} + +void llm_graph_context::build_walk_surrogate(int il, ggml_tensor * k_cur, ggml_tensor * v_cur, int64_t row) const { + ggml_tensor * gk = walk_acc_rows(cparams.walk_acc_k[il], k_cur, row); + ggml_tensor * gv = walk_acc_rows(cparams.walk_acc_v[il], v_cur, row); + ggml_tensor * k32 = k_cur->type == GGML_TYPE_F32 ? k_cur : ggml_cast(ctx0, k_cur, GGML_TYPE_F32); + ggml_tensor * v32 = v_cur->type == GGML_TYPE_F32 ? v_cur : ggml_cast(ctx0, v_cur, GGML_TYPE_F32); + ggml_tensor * term = ggml_add(ctx0, + ggml_sum(ctx0, ggml_mul(ctx0, k32, gk)), + ggml_sum(ctx0, ggml_mul(ctx0, v32, gv))); + res->t_walk_surrogate = res->t_walk_surrogate ? ggml_add(ctx0, res->t_walk_surrogate, term) : term; +} + void llm_graph_context::build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const { if (!(cparams.training && cparams.walk_exact && cparams.walk_state_surrogate)) { return; @@ -3780,16 +3946,16 @@ llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa() auto inp_attn = std::make_unique(hparams, cparams, attn_ctx); { - inp_attn->self_k_idxs = attn_ctx->get_base()->build_input_k_idxs(ctx0, ubatch); - inp_attn->self_v_idxs = attn_ctx->get_base()->build_input_v_idxs(ctx0, ubatch); + inp_attn->self_k_idxs = attn_ctx->get_base()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp_attn->self_v_idxs = attn_ctx->get_base()->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp_attn->self_kq_mask = build_attn_inp_kq_mask(ctx0, attn_ctx->get_base(), ubatch, cparams); inp_attn->self_kq_mask_cnv = inp_attn->self_kq_mask; } { - inp_attn->self_k_idxs_swa = attn_ctx->get_swa()->build_input_k_idxs(ctx0, ubatch); - inp_attn->self_v_idxs_swa = attn_ctx->get_swa()->build_input_v_idxs(ctx0, ubatch); + inp_attn->self_k_idxs_swa = attn_ctx->get_swa()->build_input_k_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); + inp_attn->self_v_idxs_swa = attn_ctx->get_swa()->build_input_v_idxs(ctx0, ubatch, cparams.training && !cparams.walk_active); inp_attn->self_kq_mask_swa = build_attn_inp_kq_mask(ctx0, attn_ctx->get_swa(), ubatch, cparams); inp_attn->self_kq_mask_swa_cnv = inp_attn->self_kq_mask_swa; diff --git a/src/llama-graph.h b/src/llama-graph.h index 0e0adb7d4bda..4c7f60bd8137 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -6,6 +6,7 @@ #include "llama-adapter.h" #include +#include #include #include #include @@ -233,6 +234,16 @@ class llm_graph_input_out_ids : public llm_graph_input_i { const uint32_t n_outputs; }; +class llm_graph_input_training_rows : public llm_graph_input_i { +public: + void set_input(const llama_ubatch * ubatch) override; + ggml_tensor * public_rows = nullptr; + ggml_tensor * label_rows = nullptr; + ggml_tensor * shadow_rows = nullptr; + ggml_tensor * public_shadow_rows = nullptr; + std::vector slice_shadow_rows; +}; + class llm_graph_input_mean : public llm_graph_input_i { public: llm_graph_input_mean(const llama_cparams & cparams) : cparams(cparams) {} @@ -783,6 +794,11 @@ struct llm_graph_params { const llama_adapter_cvec * cvec; const llama_adapter_loras * loras; + // THE SHADOW: the attached adapter being trained (llama_opt_init), the layer that switches on + // and off. It applies only to the ubatch's labelled rows (n_shadow_rows of them); a ubatch + // with no label never applies it, so serving is the stable set's graph + const llama_adapter_lora * shadow; + uint32_t n_shadow_rows; const llama_memory_context_i * mctx; const llama_cross * cross; @@ -875,10 +891,13 @@ struct llm_graph_params { cparams.embeddings_nextn == other.cparams.embeddings_nextn && cparams.embeddings_nextn_masked == other.cparams.embeddings_nextn_masked && cparams.causal_attn == other.cparams.causal_attn && + cparams.training == other.cparams.training && arch == other.arch && gtype == other.gtype && cvec == other.cvec && loras == other.loras && + shadow == other.shadow && + n_shadow_rows == other.n_shadow_rows && cross == other.cross; } }; @@ -932,6 +951,7 @@ class llm_graph_result { ggml_tensor * t_inp_tokens = nullptr; ggml_tensor * t_inp_embd = nullptr; // [n_embd_inp, n_tokens] ggml_tensor * t_logits = nullptr; + std::vector t_loss_slices; ggml_tensor * t_embd = nullptr; ggml_tensor * t_embd_pooled = nullptr; ggml_tensor * t_h_nextn = nullptr; // [n_embd, n_outputs] hidden state before final output norm @@ -943,6 +963,7 @@ class llm_graph_result { struct walk_layer { int il; uint32_t n_past; + uint32_t g0; // the first position whose K/V gradient is kept ggml_tensor * tk = nullptr; // the cached K/V of [grad_from, n_past) as F32, tapped: their ggml_tensor * tv = nullptr; // gradient adds into ak / av (ggml_opt_accumulate_grad) ggml_tensor * ak = nullptr; // the accumulator's rows [grad_from, n_past), shaped like tk / tv @@ -1054,6 +1075,8 @@ struct llm_graph_context { const llama_adapter_cvec * cvec; const llama_adapter_loras * loras; + const llama_adapter_lora * shadow; + const uint32_t n_shadow_rows; const llama_memory_context_i * mctx; const llama_cross * cross; @@ -1064,6 +1087,7 @@ struct llm_graph_context { llm_graph_result * res; mutable ggml_tensor * inp_train_rows = nullptr; // build_inp_train_rows, once per graph + mutable llm_graph_input_training_rows * inp_training_rows = nullptr; ggml_context * ctx0 = nullptr; ggml_cgraph * gf = nullptr; @@ -1187,6 +1211,16 @@ struct llm_graph_context { // is its cell). One per graph, shared by every attention layer. ggml_tensor * build_inp_train_rows() const; + // whether an attached adapter's term is in this graph: every one but the shadow always, the + // shadow only when the ubatch has labelled rows + bool lora_on(const llama_adapter_lora * adapter) const; + // base plus the shadow's term, on the labelled rows only (Cormac, 2026-10-14): those rows of + // base and of the term are gathered and added, and written back into a copy of base; every + // other row is base, byte for byte. A multiplied 0/1 mask was not: an untrained layer's inf or + // NaN times zero is NaN, and x + 0 turns a -0 into +0, on a serving row co-batched with + // training. + ggml_tensor * shadow_add(ggml_tensor * base, ggml_tensor * term) const; + // Whether this layer's attention runs through FLASH_ATTN_EXT: the one decision // build_attn_mha makes and a training graph reads to attend to the cache in place. bool uses_flash_attn(ggml_tensor * kq_b, ggml_tensor * sinks, ggml_tensor * v_mla) const; @@ -1367,9 +1401,26 @@ struct llm_graph_context { int32_t n_seqs, const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const; + // THE ONE RULE for how many recurrent states a graph keeps per sequence, read by every family's + // state write (delta-net, mamba, lfm2, bailing) and by the ops that produce snapshots. Serving with + // rollback slots keeps K = n_rs_seq + 1 (group g = the state g tokens back, so a draft can be undone); + // a context without them, and a TRAINING graph, keep one: the live state, group 0, which is the same + // bytes in both layouts. So training reads and writes the state serving reads, in the serving + // context; the undo groups it leaves stale are rewritten by the next serving decode, and the walk + // refuses a sequence with an undo pending (llama_memory_recurrent::rollback_pending). + int64_t rs_keep() const; + // THE EXACT WALK: the state a chunk leaves (`exit`, what is written back to `cache`) meets the // gradient later chunks put on it, through the surrogate . A no-op outside the walk. void build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const; + // rows [row, row + n) of a layer's walk accumulator, shaped like `like`'s K / V span + ggml_tensor * walk_acc_rows(ggml_tensor * acc, const ggml_tensor * like, int64_t row) const; + // the in-place training attention's K/V grad sinks (tap k, chunk k, tap v, chunk v), set by + // build_attn before build_attn_mha and taken by the FLASH_ATTN_EXT it builds + mutable std::array walk_grad_sinks = {}; + // + : what later chunks of the run put on this chunk's own K/V + // (accumulator rows from `row`), carried into it through the loss + void build_walk_surrogate(int il, ggml_tensor * k_cur, ggml_tensor * v_cur, int64_t row) const; ggml_tensor * build_rwkv_token_shift_load( llm_graph_input_rs * inp, @@ -1405,6 +1456,8 @@ struct llm_graph_context { // void build_sampling() const; + bool build_training_head() const; + llm_graph_input_training_rows * training_rows() const; // // dense (out) diff --git a/src/llama-kv-cache-dsa-iswa.cpp b/src/llama-kv-cache-dsa-iswa.cpp index dc10342a19c0..83d8e6cd3528 100644 --- a/src/llama-kv-cache-dsa-iswa.cpp +++ b/src/llama-kv-cache-dsa-iswa.cpp @@ -339,3 +339,10 @@ const llama_kv_cache_context * llama_kv_cache_dsa_iswa_context::get_swa() const return static_cast(ctx_swa.get()); } + +const char * llama_kv_cache_dsa_iswa::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = kv_swa->training_error(batch, flash_attn)) { + return why; + } + return kv_dsa->training_error(batch, flash_attn); +} diff --git a/src/llama-kv-cache-dsa-iswa.h b/src/llama-kv-cache-dsa-iswa.h index 28cf95bf0518..08e270acf838 100644 --- a/src/llama-kv-cache-dsa-iswa.h +++ b/src/llama-kv-cache-dsa-iswa.h @@ -34,6 +34,8 @@ class llama_kv_cache_dsa_iswa : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp index 96cb045d2e5d..54a9b4b440f7 100644 --- a/src/llama-kv-cache-dsa.cpp +++ b/src/llama-kv-cache-dsa.cpp @@ -260,3 +260,7 @@ const llama_kv_cache_context * llama_kv_cache_dsa_context::get_lid() const { return static_cast(ctx_lid.get()); } + +const char * llama_kv_cache_dsa::training_error(const llama_batch &, bool) const { + return "DSA cache row gradients are not supported"; +} diff --git a/src/llama-kv-cache-dsa.h b/src/llama-kv-cache-dsa.h index e74fc4d9100f..0fe3b538e8b9 100644 --- a/src/llama-kv-cache-dsa.h +++ b/src/llama-kv-cache-dsa.h @@ -36,6 +36,8 @@ class llama_kv_cache_dsa : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache-dsv4.cpp b/src/llama-kv-cache-dsv4.cpp index 948d08146fbf..e6b2465cffad 100644 --- a/src/llama-kv-cache-dsv4.cpp +++ b/src/llama-kv-cache-dsv4.cpp @@ -164,6 +164,8 @@ static llama_ubatch dsv4_build_raw_write_ubatch(const llama_ubatch & ubatch) { /*.seq_id_unq =*/ data->seq_id_unq.data(), /*.seq_idx =*/ data->seq_idx.data(), /*.output =*/ data->output.data(), + /*.label =*/ nullptr, + /*.label_weight =*/ nullptr, /*.data =*/ data, }; @@ -2251,3 +2253,7 @@ const llama_kv_cache_dsv4_context::comp_plan & llama_kv_cache_dsv4_context::get_ return reserve_plan_lid; } + +const char * llama_kv_cache_dsv4::training_error(const llama_batch &, bool) const { + return "DSV4 cache row gradients are not supported"; +} diff --git a/src/llama-kv-cache-dsv4.h b/src/llama-kv-cache-dsv4.h index ce39867c0342..da5af159720f 100644 --- a/src/llama-kv-cache-dsv4.h +++ b/src/llama-kv-cache-dsv4.h @@ -109,6 +109,8 @@ class llama_kv_cache_dsv4 : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache-iswa.cpp b/src/llama-kv-cache-iswa.cpp index e91866469adf..42bd02a59782 100644 --- a/src/llama-kv-cache-iswa.cpp +++ b/src/llama-kv-cache-iswa.cpp @@ -361,3 +361,10 @@ const llama_kv_cache_context * llama_kv_cache_iswa_context::get_swa() const { return static_cast(ctx_swa.get()); } + +const char * llama_kv_cache_iswa::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = kv_base->training_error(batch, flash_attn)) { + return why; + } + return kv_swa->training_error(batch, flash_attn); +} diff --git a/src/llama-kv-cache-iswa.h b/src/llama-kv-cache-iswa.h index 7dab6eaa82c8..7f7364174f40 100644 --- a/src/llama-kv-cache-iswa.h +++ b/src/llama-kv-cache-iswa.h @@ -54,6 +54,8 @@ class llama_kv_cache_iswa : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache-msa.cpp b/src/llama-kv-cache-msa.cpp index 55ef286cafa7..32dc7c91116d 100644 --- a/src/llama-kv-cache-msa.cpp +++ b/src/llama-kv-cache-msa.cpp @@ -393,3 +393,7 @@ void llama_kv_cache_msa_context::set_input_pos_mask(ggml_tensor * dst, const lla } } } + +const char * llama_kv_cache_msa::training_error(const llama_batch &, bool) const { + return "MSA cache row gradients are not supported"; +} diff --git a/src/llama-kv-cache-msa.h b/src/llama-kv-cache-msa.h index f09b6d32b044..cb93a00b034f 100644 --- a/src/llama-kv-cache-msa.h +++ b/src/llama-kv-cache-msa.h @@ -33,6 +33,8 @@ class llama_kv_cache_msa : public llama_memory_i { // llama_memory_i + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index f815cd76656a..fc118c44814c 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1249,7 +1249,7 @@ uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { return result; } -ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { +ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo, ggml_tensor * written) const { const int32_t ikv = map_layer_ids.at(il); auto * k = layers[ikv].k; @@ -1261,7 +1261,7 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; - return ggml_view_4d(ctx, k, + return ggml_view_4d(ctx, written ? written : k, hparams.n_embd_head_k(il), hparams.n_head_kv(il), n_kv, ns, ggml_row_size(k->type, hparams.n_embd_head_k(il)), ggml_row_size(k->type, n_embd_k_gqa), @@ -1269,7 +1269,7 @@ ggml_tensor * llama_kv_cache::get_k(ggml_context * ctx, int32_t il, uint32_t n_k ggml_row_size(k->type, n_embd_k_gqa*kv_size)*sinfo.s0); } -ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { +ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo, ggml_tensor * written) const { const int32_t ikv = map_layer_ids.at(il); auto * v = layers[ikv].v; @@ -1284,7 +1284,7 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k if (!v_trans) { // note: v->nb[1] <= v->nb[2] - return ggml_view_4d(ctx, v, + return ggml_view_4d(ctx, written ? written : v, hparams.n_embd_head_v(il), hparams.n_head_kv(il), n_kv, ns, ggml_row_size(v->type, hparams.n_embd_head_v(il)), // v->nb[1] ggml_row_size(v->type, n_embd_v_gqa), // v->nb[2] @@ -1293,7 +1293,7 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k } // note: v->nb[1] > v->nb[2] - return ggml_view_4d(ctx, v, + return ggml_view_4d(ctx, written ? written : v, n_kv, hparams.n_head_kv(il), hparams.n_embd_head_v(il), ns, ggml_row_size(v->type, kv_size*hparams.n_embd_head_v(il)), // v->nb[1] ggml_row_size(v->type, kv_size), // v->nb[2] @@ -1392,25 +1392,44 @@ ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggm return ggml_set_rows(ctx, v_view, v_cur, v_idxs); } -ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { +static ggml_type cache_index_type(uint64_t extent) { + return extent <= (uint64_t) INT32_MAX ? GGML_TYPE_I32 : GGML_TYPE_I64; +} + +bool llama_kv_cache::supports_row_gradients() const { + const uint64_t extent = (uint64_t) get_size()*n_stream; + return cache_index_type(extent) == GGML_TYPE_I32 && + cache_index_type(v_trans ? extent*hparams.n_embd_v_gqa_max() : extent) == GGML_TYPE_I32; +} + +static void set_cache_index(ggml_tensor * dst, size_t i, int64_t value) { + if (dst->type == GGML_TYPE_I32) { + GGML_ASSERT(value >= 0 && value <= INT32_MAX); + ((int32_t *) dst->data)[i] = (int32_t) value; + } else { + GGML_ASSERT(dst->type == GGML_TYPE_I64); + ((int64_t *) dst->data)[i] = value; + } +} +ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients) const { const uint32_t n_tokens = ubatch.n_tokens; - ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); + ggml_tensor * k_idxs = ggml_new_tensor_1d(ctx, row_gradients ? cache_index_type((uint64_t) get_size()*n_stream) : GGML_TYPE_I64, n_tokens); ggml_set_input(k_idxs); return k_idxs; } -ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { +ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients) const { const uint32_t n_tokens = ubatch.n_tokens; ggml_tensor * v_idxs; if (!v_trans) { - v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); + v_idxs = ggml_new_tensor_1d(ctx, row_gradients ? cache_index_type((uint64_t) get_size()*n_stream) : GGML_TYPE_I64, n_tokens); } else { - v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens*hparams.n_embd_v_gqa_max()); + v_idxs = ggml_new_tensor_1d(ctx, row_gradients ? cache_index_type((uint64_t) get_size()*n_stream*hparams.n_embd_v_gqa_max()) : GGML_TYPE_I64, n_tokens*hparams.n_embd_v_gqa_max()); } ggml_set_input(v_idxs); @@ -1464,13 +1483,13 @@ void llama_kv_cache::set_input_k_idxs(ggml_tensor * dst, const llama_ubatch * ub GGML_ASSERT(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream()); GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); - int64_t * data = (int64_t *) dst->data; + for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { const int64_t offs = sinfo.strm[s]*get_size(); for (uint32_t i = 0; i < sinfo.size(); ++i) { - data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i]; + set_cache_index(dst, s*sinfo.size() + i, offs + sinfo.idxs[s][i]); } } } @@ -1480,14 +1499,14 @@ void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ub GGML_ASSERT(n_tokens == (int64_t) sinfo.size()*sinfo.n_stream()); GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); - int64_t * data = (int64_t *) dst->data; + if (!v_trans) { for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { const int64_t offs = sinfo.strm[s]*get_size(); for (uint32_t i = 0; i < sinfo.size(); ++i) { - data[s*sinfo.size() + i] = offs + sinfo.idxs[s][i]; + set_cache_index(dst, s*sinfo.size() + i, offs + sinfo.idxs[s][i]); } } } else { @@ -1501,7 +1520,7 @@ void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ub for (uint32_t i = 0; i < sinfo.size(); ++i) { for (uint32_t j = 0; j < n_embd_v_gqa; ++j) { - data[s*sinfo.size()*n_embd_v_gqa + i*n_embd_v_gqa + j] = offs + j*kv_size + sinfo.idxs[s][i]; + set_cache_index(dst, s*sinfo.size()*n_embd_v_gqa + i*n_embd_v_gqa + j, offs + j*kv_size + sinfo.idxs[s][i]); } } } @@ -2594,12 +2613,12 @@ ggml_type llama_kv_cache_context::type_v() const { return kv->type_v(); } -ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il) const { - return kv->get_k(ctx, il, n_kv, sinfos[i_cur]); +ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il, ggml_tensor * written) const { + return kv->get_k(ctx, il, n_kv, sinfos[i_cur], written); } -ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) const { - return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); +ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il, ggml_tensor * written) const { + return kv->get_v(ctx, il, n_kv, sinfos[i_cur], written); } ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const { @@ -2610,12 +2629,12 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_ return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]); } -ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { - return kv->build_input_k_idxs(ctx, ubatch); +ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients) const { + return kv->build_input_k_idxs(ctx, ubatch, row_gradients); } -ggml_tensor * llama_kv_cache_context::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { - return kv->build_input_v_idxs(ctx, ubatch); +ggml_tensor * llama_kv_cache_context::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients) const { + return kv->build_input_v_idxs(ctx, ubatch, row_gradients); } ggml_tensor * llama_kv_cache_context::build_input_k_rot(ggml_context * ctx) const { @@ -2653,3 +2672,25 @@ void llama_kv_cache_context::set_input_k_rot(ggml_tensor * dst) const { void llama_kv_cache_context::set_input_v_rot(ggml_tensor * dst) const { kv->set_input_v_rot(dst); } + +const char * llama_kv_cache::training_error(const llama_batch & batch, bool flash_attn) const { + if (!supports_row_gradients()) { + return "cache row indices exceed the backward gather's index range"; + } + // the flash-attention backward reads one stream: a training batch may touch only one (a serving + // context with a stream per slot learns on one slot's stream at a time) + if (flash_attn && get_n_stream() != 1) { + int64_t stream = -1; + for (int32_t i = 0; i < batch.n_tokens; ++i) { + for (int32_t k = 0; k < (batch.n_seq_id ? batch.n_seq_id[i] : 1); ++k) { + const llama_seq_id sq = batch.seq_id ? batch.seq_id[i][k] : 0; + const int64_t st = seq_to_stream[sq]; + if (stream >= 0 && st != stream) { + return "a training batch with flash attention may touch one cache stream: this one spans several"; + } + stream = st; + } + } + } + return nullptr; +} diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index f00ec2137bb3..696db3546492 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -122,6 +122,8 @@ class llama_kv_cache : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, @@ -158,6 +160,8 @@ class llama_kv_cache : public llama_memory_i { uint32_t get_size() const; uint32_t get_n_stream() const; + bool supports_row_gradients() const; + bool get_has_shift() const; ggml_type type_k() const; @@ -175,8 +179,8 @@ class llama_kv_cache : public llama_memory_i { uint32_t get_n_kv(const slot_info & sinfo) const; // get views of the current state of the cache - ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; - ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; + ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo, ggml_tensor * written = nullptr) const; + ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo, ggml_tensor * written = nullptr) const; // store k_cur and v_cur in the cache based on the provided head location ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; @@ -204,8 +208,8 @@ class llama_kv_cache : public llama_memory_i { // input API // - ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; - ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; + ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients = false) const; + ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients = false) const; ggml_tensor * build_input_k_rot(ggml_context * ctx) const; ggml_tensor * build_input_v_rot(ggml_context * ctx) const; @@ -372,8 +376,8 @@ class llama_kv_cache_context : public llama_memory_context_i { ggml_type type_v() const; // get views of the current state of the cache - ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; - ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; + ggml_tensor * get_k(ggml_context * ctx, int32_t il, ggml_tensor * written = nullptr) const; + ggml_tensor * get_v(ggml_context * ctx, int32_t il, ggml_tensor * written = nullptr) const; // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory @@ -387,8 +391,8 @@ class llama_kv_cache_context : public llama_memory_context_i { // create destination indices for each head of the current batch for where it would be written in the KV cache // the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but // helps understand the implementation logic of cpy_k and cpy_v - ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; - ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const; + ggml_tensor * build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients = false) const; + ggml_tensor * build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch, bool row_gradients = false) const; ggml_tensor * build_input_k_rot(ggml_context * ctx) const; ggml_tensor * build_input_v_rot(ggml_context * ctx) const; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index e7f6e565dca1..bfba14c56454 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -674,3 +674,10 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } } + +const char * llama_memory_hybrid_idx::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = llama_memory_hybrid::training_error(batch, flash_attn)) { + return why; + } + return get_mem_idx() ? "indexed hybrid cache row gradients are not supported" : nullptr; +} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index a17e184209d1..f8a3f2dad905 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -48,6 +48,8 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index 06f7fd5428c4..21c39149a358 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -283,3 +283,10 @@ const llama_kv_cache_iswa_context * llama_memory_hybrid_iswa_context::get_attn() const llama_memory_recurrent_context * llama_memory_hybrid_iswa_context::get_recr() const { return static_cast(ctx_recr.get()); } + +const char * llama_memory_hybrid_iswa::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = mem_attn->training_error(batch, flash_attn)) { + return why; + } + return mem_recr->training_error(batch, flash_attn); +} diff --git a/src/llama-memory-hybrid-iswa.h b/src/llama-memory-hybrid-iswa.h index c9d3f9f57c50..693782a1c2ce 100644 --- a/src/llama-memory-hybrid-iswa.h +++ b/src/llama-memory-hybrid-iswa.h @@ -47,6 +47,8 @@ class llama_memory_hybrid_iswa : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6f..dfb4e9b947f5 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -277,3 +277,10 @@ const llama_kv_cache_context * llama_memory_hybrid_context::get_attn() const { const llama_memory_recurrent_context * llama_memory_hybrid_context::get_recr() const { return static_cast(ctx_recr.get()); } + +const char * llama_memory_hybrid::training_error(const llama_batch & batch, bool flash_attn) const { + if (const char * why = mem_attn->training_error(batch, flash_attn)) { + return why; + } + return mem_recr->training_error(batch, flash_attn); +} diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74991..2a92bdbaa81a 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -47,6 +47,8 @@ class llama_memory_hybrid : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7e..93c3ca873c12 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -1270,3 +1270,12 @@ int32_t llama_memory_recurrent_context::s_copy(int i) const { } return (int32_t)(idx * mem->size) + src0; } + +const char * llama_memory_recurrent::training_error(const llama_batch & batch, bool) const { + for (int32_t i = 0; i < batch.n_tokens; ++i) { + if (batch.label[i] >= 0 && rollback_pending(batch.seq_id[i][0])) { + return "a labelled sequence has pending recurrent rollback"; + } + } + return nullptr; +} diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5e..46b65f5dfa67 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -32,6 +32,8 @@ class llama_memory_recurrent : public llama_memory_i { // llama_memory_i // + const char * training_error(const llama_batch & batch, bool flash_attn) const override; + llama_memory_context_ptr init_batch( llama_batch_allocr & balloc, uint32_t n_ubatch, @@ -78,6 +80,12 @@ class llama_memory_recurrent : public llama_memory_i { void set_rs_idx(llama_seq_id seq_id, uint32_t idx); + // an undo of a draft is pending on this sequence: its next read takes a snapshot group, not the + // live state (group 0) a training graph reads and writes (llm_graph_context::rs_keep) + bool rollback_pending(llama_seq_id seq_id) const { + return n_rs_seq > 0 && seq_id >= 0 && (size_t) seq_id < rs_idx.size() && rs_idx[seq_id] != 0; + } + // computed before each graph build uint32_t n = 0; diff --git a/src/llama-memory.h b/src/llama-memory.h index db825396645e..bc4935ede0a5 100644 --- a/src/llama-memory.h +++ b/src/llama-memory.h @@ -82,6 +82,9 @@ struct llama_memory_i { virtual ~llama_memory_i() = default; + // Read-only admission for a validated batch; nullptr means admitted. + virtual const char * training_error(const llama_batch & batch, bool flash_attn) const = 0; + // split the input batch into a set of ubatches and verify that they can fit into the cache // return a context object containing the ubatches and memory state required to process them // check the llama_memory_context_i::get_status() for the result diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 9d34a2d28484..63fa565e7843 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2646,6 +2646,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { std::unique_ptr llm = build_arch_graph(params); + if (!llm->build_training_head()) { + LLAMA_LOG_ERROR("%s: training head is not a supported row-wise projection\n", __func__); + return nullptr; + } + // add on pooling layer llm->build_pooling(cls, cls_b, cls_out, cls_out_b, cls_norm); diff --git a/src/llama-walk.cpp b/src/llama-walk.cpp index 481e8a8e9078..73e9fa2449f0 100644 --- a/src/llama-walk.cpp +++ b/src/llama-walk.cpp @@ -23,11 +23,15 @@ llama_walk_plan llama_walk_plan::of(const std::vector & labels, uin while ((int64_t) end <= plan.last_label && (labels[end] >= 0) == labelled) { ++end; } + const size_t j0 = plan.chunks.size(); for (uint32_t c0 = p; c0 < end; c0 += n_batch) { const uint32_t n = std::min(c0 + n_batch, end) - c0; plan.chunks.push_back({ c0, n, labelled ? n : 0 }); plan.n_labels += labelled ? n : 0; } + if (labelled) { + plan.runs.push_back({ j0, plan.chunks.size() }); + } p = end; } return plan; @@ -37,35 +41,52 @@ uint32_t llama_walk_plan::grad_from(uint32_t c0, uint32_t horizon) { return horizon == 0 || c0 < horizon ? 0 : c0 - horizon; } -void llama_walk_plan::mark_exact(bool recurrent, uint32_t horizon) { - uint32_t reach = UINT32_MAX; // the lowest position a later trained chunk's gradient reaches - first = -1; - for (int64_t j = (int64_t) chunks.size() - 1; j >= 0; --j) { - llama_walk_chunk & c = chunks[j]; - c.surrogate = reach < c.c0 + c.n; - c.needed = c.n_labels > 0 || c.surrogate || (recurrent && first >= 0); +void llama_walk_plan::mark_exact(uint32_t horizon) { + for (llama_walk_chunk & c : chunks) { + c.needed = false; + c.surrogate = false; c.period_end = false; - if (!c.needed) { - continue; + } + for (const llama_walk_run & r : runs) { + uint32_t reach = UINT32_MAX; // the lowest position a later chunk of the run reaches + for (size_t j = r.j1; j-- > r.j0; ) { + llama_walk_chunk & c = chunks[j]; + c.needed = true; + c.surrogate = reach < c.c0 + c.n; + c.grad_from = std::max(begin(r), grad_from(c.c0, horizon)); + c.loss_scale = (float) c.n_labels / (float) n_labels; + reach = std::min(reach, c.grad_from); } - c.grad_from = grad_from(c.c0, horizon); - c.loss_scale = (float) c.n_labels / (float) n_labels; - reach = std::min(reach, c.grad_from); - first = j; } - GGML_ASSERT(first >= 0); + GGML_ASSERT(!runs.empty()); + first = (int64_t) runs.back().j0; chunks[first].period_end = true; } +uint32_t llama_walk_plan::longest_run() const { + uint32_t n = 0; + for (const llama_walk_run & r : runs) { + n = std::max(n, end(r) - begin(r)); + } + return n; +} + size_t llama_walk_plan::host_bytes(size_t stride, size_t kv_bytes, size_t state_bytes) const { - const size_t n_checkpoints = chunks.empty() ? 0 : (chunks.size() + stride - 1)/stride - 1; + size_t n_checkpoints = 0; + for (const llama_walk_run & r : runs) { + n_checkpoints = std::max(n_checkpoints, (r.j1 - r.j0 + stride - 1)/stride); + } return kv_bytes + n_checkpoints*state_bytes; } size_t llama_walk_plan::stride_for(size_t budget, size_t kv_bytes, size_t state_bytes) const { + size_t longest = 1; + for (const llama_walk_run & r : runs) { + longest = std::max(longest, r.j1 - r.j0); + } size_t stride = 1; if (budget > 0) { - while (stride < chunks.size() && host_bytes(stride, kv_bytes, state_bytes) > budget) { + while (stride < longest && host_bytes(stride, kv_bytes, state_bytes) > budget) { ++stride; } } @@ -120,13 +141,22 @@ bool llama_walk_acc::init(const llama_hparams & hparams, const llama_memory_recu bound = &cparams; cparams.walk_acc_k = k.data(); cparams.walk_acc_v = v.data(); + cparams.walk_acc_base = 0; return true; } +void llama_walk_acc::start_run(uint32_t base) { + GGML_ASSERT(bound != nullptr); + ggml_backend_buffer_clear(buf.get(), 0); + state_ready = false; + bound->walk_acc_base = base; +} + llama_walk_acc::~llama_walk_acc() { if (bound != nullptr) { bound->walk_acc_k = nullptr; bound->walk_acc_v = nullptr; + bound->walk_acc_base = 0; } } diff --git a/src/llama-walk.h b/src/llama-walk.h index 67a68bf12db7..35ceedfeef8d 100644 --- a/src/llama-walk.h +++ b/src/llama-walk.h @@ -32,26 +32,40 @@ struct llama_walk_chunk { float loss_scale = 0.0f; // its labelled positions over the window's }; +// chunks [j0, j1) of the plan: one labelled run, which trains given everything before it as +// served (the shadow adapter applies to labelled rows only) +struct llama_walk_run { + size_t j0; + size_t j1; +}; + struct llama_walk_plan { // the window's runs, each cut every n_batch positions from its start; both walks take // them in this order. Nothing after the last label is in the plan. std::vector chunks; + std::vector runs; // the labelled ones, in order int64_t last_label = -1; // -1: the window has no label, nothing to learn from it uint32_t n_labels = 0; int64_t first = -1; // the exact walk's period end: the chunk its reverse pass trains last static llama_walk_plan of(const std::vector & labels, uint32_t n_ctx, uint32_t n_batch); - // which chunks the exact walk trains (a labelled one, one a later gradient reaches, and on a - // recurrent model every one before a trained one, since the state reaches every later - // chunk), how far back each gradient reaches, and the period end - void mark_exact(bool recurrent, uint32_t horizon); + // positions [begin, end) of a run, and the longest run's length (what the exact walk's K/V + // gradient holds: a run's, never the window's) + uint32_t begin(const llama_walk_run & r) const { return chunks[r.j0].c0; } + uint32_t end (const llama_walk_run & r) const { return chunks[r.j1 - 1].c0 + chunks[r.j1 - 1].n; } + uint32_t longest_run() const; + + // which chunks the exact walk trains (the labelled ones: the context before a run is her + // stable set's and carries no gradient), how far back each gradient reaches (never before + // its run's start), and the period end (the first chunk of the last run, trained last) + void mark_exact(uint32_t horizon); // the first cached position a chunk at c0 carries gradient to (horizon 0: the whole prefix) static uint32_t grad_from(uint32_t c0, uint32_t horizon); // THE HOST BUDGET (Fable): the K/V accumulators, plus a recurrent state checkpoint every - // `stride` chunks (the first chunk's entry is the zero state and costs nothing) + // `stride` chunks of the run being walked (the largest run's) size_t host_bytes(size_t stride, size_t kv_bytes, size_t state_bytes) const; // the smallest stride under the budget (0: no cap); the reverse pass rebuilds a chunk's // entry state by decoding forward from its checkpoint, so a longer stride trades host @@ -60,11 +74,11 @@ struct llama_walk_plan { }; // THE EXACT WALK's gradients, in host memory, each added into in place by the backward: per -// attention layer dL/d(cached K) and dL/d(cached V) of the window, [n_embd_*_gqa, position] -// F32 (null for a recurrent layer), which the graph reads through cparams while bound (never -// past the accumulator's life); per recurrent state tensor (a layer's conv state r_l and its -// recurrent state s_l) dL/d(the state entering the chunk trained last), which is what the -// chunk before it puts on the state it leaves (state_ready once one is in) +// attention layer dL/d(cached K) and dL/d(cached V) of the run being trained, [n_embd_*_gqa, +// position - its start] F32 (null for a recurrent layer), which the graph reads through cparams +// while bound (never past the accumulator's life); per recurrent state tensor (a layer's conv +// state r_l and its recurrent state s_l) dL/d(the state entering the chunk trained last), which +// is what the chunk before it puts on the state it leaves (state_ready once one is in) struct llama_walk_acc { std::vector k, v; std::unordered_map state; @@ -78,6 +92,8 @@ struct llama_walk_acc { // zeroed accumulators for n_pos positions, bound into cparams; false: the host memory // could not be allocated bool init(const llama_hparams & hparams, const llama_memory_recurrent * recr, uint32_t n_pos, llama_cparams & cparams); + // zeroed again for the run starting at position base (row 0) + void start_run(uint32_t base); // what init allocates for the K/V of n_pos positions, and one sequence's recurrent state static size_t kv_bytes(const llama_hparams & hparams, uint32_t n_pos); diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp index 0637931cc0c9..e794c0914afc 100644 --- a/src/models/bailingmoe3.cpp +++ b/src/models/bailingmoe3.cpp @@ -183,7 +183,7 @@ static ggml_tensor * bailingmoe3_causal_conv1d( int64_t n_tokens, int64_t cache_head, uint32_t mem_size, - uint32_t n_rs_seq) { + int64_t K) { // states kept per sequence: llm_graph_context::rs_keep const int64_t d_inner = head_dim * n_head; const int64_t conv_state_size = (d_conv - 1) * d_inner; const int64_t total_state_size = 3 * conv_state_size; @@ -197,7 +197,6 @@ static ggml_tensor * bailingmoe3_causal_conv1d( x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0); - const int64_t K = (int64_t) n_rs_seq + 1; const int64_t n_written = std::min(n_seq_tokens, K); for (int64_t slot = 0; slot < n_written; ++slot) { @@ -262,13 +261,13 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph ggml_tensor * q = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * k = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * v = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); gate = ggml_add(ctx0, gate, layer.ssm_dt_b); diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index e0a6368bd2ed..f613db67f5b9 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -476,7 +476,7 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( const size_t row_size = ggml_row_size(conv_states_all->type, row_count); - if (cparams.n_rs_seq == 0) { + if (rs_keep() == 1) { const int64_t s_idx = conv_input->ne[0] - conv_states->ne[0]; const int64_t s_slot = 0; @@ -496,12 +496,12 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update)); build_walk_state_exit(conv_states_all, conv_state_last); } else { - GGML_ASSERT(!cparams.walk_exact && "the exact walk keeps one recurrent state per sequence (n_rs_seq = 0)"); + GGML_ASSERT(!cparams.walk_exact && "a training graph keeps one recurrent state per sequence (rs_keep)"); // [TAG_RECURRENT_ROLLBACK_SPLITS] // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside // the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument - const int64_t K = (int64_t) cparams.n_rs_seq + 1; + const int64_t K = rs_keep(); for (int64_t t = 1; t <= K; ++t) { const int64_t s_idx = std::max(0, conv_input->ne[0] - conv_states->ne[0] - K + t); @@ -545,7 +545,7 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( const int64_t n_seqs = s->ne[3]; const int64_t n_seq_tokens = q->ne[2]; - const bool keep = cparams.n_rs_seq > 0; + const bool keep = rs_keep() > 1; if (!keep) { auto attn_out = build_delta_net(q, k, v, g, b, s, il); @@ -564,7 +564,7 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( } const int64_t D = S_v * S_v * H_v; - const int64_t K = cparams.n_rs_seq + 1; + const int64_t K = rs_keep(); // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output. ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K); diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp index 9a4295557053..d2e954844608 100644 --- a/src/models/lfm2.cpp +++ b/src/models/lfm2.cpp @@ -205,7 +205,7 @@ llama_model_lfm2::graph::graph(const llama_model & model, const llm_graph_ GGML_ASSERT(bx->ne[0] > conv->ne[0]); // write conv states: slot 0 = the final state, slot s = the state s tokens back (partial rollback) - const int64_t K = hparams.causal_attn && cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; + const int64_t K = hparams.causal_attn ? rs_keep() : 1; const int64_t n_written = std::min(n_seq_tokens, K); const auto mem_size = mctx_cur->get_size(); const size_t row_size = ggml_row_size(conv_state->type, (int64_t) d_conv * n_embd); diff --git a/src/models/mamba-base.cpp b/src/models/mamba-base.cpp index 03ee3805bf80..af8fd15198a8 100644 --- a/src/models/mamba-base.cpp +++ b/src/models/mamba-base.cpp @@ -167,7 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, const int64_t n_seqs = ubatch.n_seqs; const int64_t n_seq_tokens = ubatch.n_seq_tokens; - const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; + const int64_t K = rs_keep(); GGML_ASSERT(n_seqs != 0); GGML_ASSERT(ubatch.equal_seqs()); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index b53590be849e..a90000500b1f 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -44,6 +44,7 @@ #include #include #include +#include #ifdef __EMSCRIPTEN__ # define N_THREADS 1 @@ -1770,6 +1771,28 @@ struct test_case { ggml_set_loss(out); ggml_build_forward_expand(gf, out); + // A param-reachable node whose op has no backward is not supported, not an abort: the + // full grad suite must reach its summary (ggml_node_has_backward is the one list). + { + std::unordered_set reach; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + ggml_tensor * node = ggml_graph_node(gf, i); + bool reached = false; + for (int k = 0; k < GGML_MAX_SRC; ++k) { + const ggml_tensor * s = node->src[k]; + reached = reached || (s != nullptr && ((s->flags & GGML_TENSOR_FLAG_PARAM) || reach.count(s))); + } + if (!reached) { + continue; + } + reach.insert(node); + if (!ggml_node_has_backward(node)) { + output_printer->print_operation(test_operation_info(op_desc(out), vars(), ggml_backend_name(backend), + test_status_t::NOT_SUPPORTED, std::string("no backward for ") + ggml_op_desc(node))); + return true; + } + } + } ggml_graph_cpy(gf, gb); ggml_build_backward_expand(ctx.get(), gb, nullptr); if (expect.size() != 1 || expect[0] != 0.0f) { @@ -2052,6 +2075,12 @@ struct test_unary : public test_case { min = -10.f; max = 10.f; } + // EXPM1's gradient is checked against a finite difference: exp past ~88.7 overflows f32, + // and the summed loss must stay small enough in f32 for a small step to register. + if (op == GGML_UNARY_OP_EXPM1) { + min = -2.f; + max = 2.f; + } for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { // test extended range of values to check for NaNs in GELU @@ -2060,7 +2089,18 @@ struct test_unary : public test_case { } float grad_eps() override { - return 15.0f; + // 15 suits the piecewise-linear ops over [-150, 150]; exp's slope changes by e^15 over such + // a step, so EXPM1 takes the smooth ops' convention (test_sin): a small step, the 4-point + // difference, and an f32-sum tolerance (its range is narrowed above) + return op == GGML_UNARY_OP_EXPM1 ? 0.2f : 15.0f; + } + + bool grad_precise() override { + return op == GGML_UNARY_OP_EXPM1; + } + + double max_maa_err() override { + return op == GGML_UNARY_OP_EXPM1 ? 1e-3 : test_case::max_maa_err(); } std::vector grad_expect() override { @@ -2390,9 +2430,12 @@ struct test_set_rows : public test_case { const std::array nr23; // broadcast only dims 2 and 3 const int r; // rows to set const bool v; // view (non-contiguous src1) + // gradients for both the rows written and the destination, written into a copy of it (the + // shadow adapter's row selection, llm_graph_context::shadow_add) + const bool grad; std::string vars() override { - return VARS_TO_STR7(type_src, type_dst, type_idx, ne, nr23, r, v); + return VARS_TO_STR8(type_src, type_dst, type_idx, ne, nr23, r, v, grad); } test_set_rows(ggml_type type_src, @@ -2400,8 +2443,8 @@ struct test_set_rows : public test_case { ggml_type type_idx, std::array ne, std::array nr23, - int r, bool v = false) - : type_src(type_src), type_dst(type_dst), type_idx(type_idx), ne(ne), nr23(nr23), r(r), v(v) {} + int r, bool v = false, bool grad = false) + : type_src(type_src), type_dst(type_dst), type_idx(type_idx), ne(ne), nr23(nr23), r(r), v(v), grad(grad) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * dst = ggml_new_tensor_4d(ctx, type_dst, ne[0], ne[1], ne[2]*nr23[0], ne[3]*nr23[1]); @@ -2410,6 +2453,13 @@ struct test_set_rows : public test_case { ggml_tensor * src = ggml_new_tensor_4d(ctx, type_src, ne[0], r, ne[2]*nr23[0], ne[3]*nr23[1]); ggml_set_name(src, "src"); + if (grad) { + ggml_set_param(dst); + ggml_set_param(src); + dst = ggml_dup(ctx, dst); + ggml_set_name(dst, "copy_of_dst"); + } + ggml_tensor * row_idxs = ggml_new_tensor_3d(ctx, type_idx, r, ne[2], ne[3]); ggml_set_name(row_idxs, "row_idxs"); @@ -3516,22 +3566,28 @@ struct test_rms_norm_back : public test_case { const ggml_type type; const std::array ne; const float eps; + const bool strided; // a and b are row views into tensors twice as wide (a per-head norm) std::string vars() override { - return VARS_TO_STR3(type, ne, eps); + return VARS_TO_STR4(type, ne, eps, strided); } test_rms_norm_back(ggml_type type = GGML_TYPE_F32, std::array ne = {64, 5, 4, 3}, - float eps = 1e-6f) - : type(type), ne(ne), eps(eps) {} + float eps = 1e-6f, bool strided = false) + : type(type), ne(ne), eps(eps), strided(strided) {} - ggml_tensor * build_graph(ggml_context * ctx) override { - ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(a, "a"); + ggml_tensor * input(ggml_context * ctx, const char * name) { + std::array ne_full = ne; + ne_full[0] *= strided ? 2 : 1; + ggml_tensor * t = ggml_new_tensor(ctx, type, 4, ne_full.data()); + ggml_set_name(t, name); + return strided ? ggml_view_4d(ctx, t, ne[0], ne[1], ne[2], ne[3], t->nb[1], t->nb[2], t->nb[3], 0) : t; + } - ggml_tensor * b = ggml_new_tensor(ctx, type, 4, ne.data()); - ggml_set_name(b, "b"); + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a = input(ctx, "a"); + ggml_tensor * b = input(ctx, "b"); ggml_tensor * out = ggml_rms_norm_back(ctx, a, b, eps); ggml_set_name(out, "out"); @@ -5068,7 +5124,7 @@ struct test_out_prod : public test_case { } ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) override { - ggml_tensor * a = ggml_new_tensor_4d(a_own_buffer ? ctx_weights : ctx, type_a, m, k, bs[0], bs[1]); + ggml_tensor * a = ggml_new_tensor_4d(a_own_buffer && ctx_weights ? ctx_weights : ctx, type_a, m, k, bs[0], bs[1]); // grad mode passes no weights context ggml_set_name(a, "a"); ggml_tensor * b; @@ -7457,6 +7513,111 @@ struct test_flash_attn_ext : public test_case { } }; +// FLASH_ATTN_EXT's gradients through either backward (the native ggml_flash_attn_back or the +// composed tiles), against finite differences: GQA, a causal mask, q, k and v all trained +struct test_flash_attn_ext_grad : public test_case { + const int64_t hs, nh, nhkv, kv, nb; + const bool native; + + std::string vars() override { + return VARS_TO_STR6(hs, nh, nhkv, kv, nb, native); + } + + test_flash_attn_ext_grad(int64_t hs = 16, int64_t nh = 4, int64_t nhkv = 2, int64_t kv = 24, int64_t nb = 8, bool native = true) + : hs(hs), nh(nh), nhkv(nhkv), kv(kv), nb(nb), native(native) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, nb, nh); + ggml_tensor * k = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, kv, nhkv); + ggml_tensor * v = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, kv, nhkv); + ggml_set_param(q); ggml_set_name(q, "q"); + ggml_set_param(k); ggml_set_name(k, "k"); + ggml_set_param(v); ggml_set_name(v, "v"); + ggml_tensor * m = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kv, nb); + ggml_set_name(m, "m"); + ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf((float) hs), 0.0f, 0.0f); + ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); + ggml_flash_attn_ext_set_back_tile(out, native ? GGML_FA_BACK_NATIVE : 0); + // weighted by a fixed random tensor: a plain sum of attention outputs nearly cancels the + // query gradient (each row's weights sum to one), and the finite difference drowns in it + ggml_tensor * w = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, out->ne[0], out->ne[1], out->ne[2]); + ggml_set_name(w, "w"); + out = ggml_mul(ctx, out, w); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "m") == 0) { + // causal: query row i sees KV positions j <= kv - nb + i + std::vector m(ggml_nelements(t)); + for (int64_t i = 0; i < t->ne[1]; ++i) { + for (int64_t j = 0; j < t->ne[0]; ++j) { + m[i*t->ne[0] + j] = ggml_fp32_to_fp16(j <= kv - nb + i ? 0.0f : -INFINITY); + } + } + ggml_backend_tensor_set(t, m.data(), 0, ggml_nbytes(t)); + } else { + init_tensor_uniform(t); + } + } + } + + bool grad_precise() override { + return true; + } +}; + +// GGML_OP_FLASH_ATTN_BACK itself: each backend's kernel against the CPU reference, at the cache +// types training reads (F16, F32, q8_0) and a GQA group +struct test_flash_attn_back : public test_case { + const int64_t hs, nh, nhkv, kv, nb; + const ggml_type type_kv; + const int64_t kv0; // the first key row with a dK/dV (a frozen prefix before it) + + std::string vars() override { + return VARS_TO_STR7(hs, nh, nhkv, kv, nb, type_kv, kv0); + } + + test_flash_attn_back(int64_t hs = 128, int64_t nh = 8, int64_t nhkv = 2, int64_t kv = 512, int64_t nb = 64, ggml_type type_kv = GGML_TYPE_F16, + int64_t kv0 = 0) + : hs(hs), nh(nh), nhkv(nhkv), kv(kv), nb(nb), type_kv(type_kv), kv0(kv0) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, nb, nh); + ggml_tensor * k = ggml_new_tensor_3d(ctx, type_kv, hs, kv, nhkv); + ggml_tensor * v = ggml_new_tensor_3d(ctx, type_kv, hs, kv, nhkv); + ggml_tensor * m = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, kv, nb); + ggml_tensor * d = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hs, nh, nb); + ggml_set_name(q, "q"); ggml_set_name(k, "k"); ggml_set_name(v, "v"); ggml_set_name(m, "m"); ggml_set_name(d, "d"); + ggml_tensor * out = ggml_flash_attn_back(ctx, q, k, v, m, d, 1.0f/sqrtf((float) hs), kv0); + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "m") == 0) { + init_tensor_kq_mask(t); + } else { + init_tensor_uniform(t); + } + } + } + + double max_nmse_err() override { + return 1e-5; // the kernel and the reference sum in different orders over the window + } + + // the five products of the ideal backward (S, dP, dV, dQ, dK), so perf reports TFLOPS against + // the forward's two + uint64_t op_flops(ggml_tensor * t) override { + GGML_UNUSED(t); + return 2ull*nb*nh*(3*kv*hs + 2*(kv - kv0)*hs); // pass 1 over every key, pass 2 from kv0 + } +}; + // GGML_OP_CROSS_ENTROPY_LOSS // Completion-only training masks rows: a masked position's label row is all zero. Every other // row sums to 1. "masked" zeroes every other row, so the finite-difference check sees whether the @@ -7481,6 +7642,73 @@ static void ce_test_init_mask(ggml_tensor * t) { ggml_backend_tensor_set(t, m.data(), 0, m.size()*sizeof(float)); } +// GGML_OP_CROSS_ENTROPY_LOSS with one I32 target per row (-1: a row with no loss), the form a +// training step's labels take so no vocabulary-wide one-hot tensor is built +struct test_cross_entropy_loss_index : public test_case { + const std::array ne; + + std::string vars() override { + return VARS_TO_STR2(ne, weighted); + } + + const bool weighted; // a weight per row: how much each row counts, not only whether + + test_cross_entropy_loss_index(std::array ne = {10, 5, 4, 3}, bool weighted = false) : ne(ne), weighted(weighted) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * logits = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); + ggml_set_param(logits); + ggml_set_name(logits, "logits"); + + ggml_tensor * targets = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, ne[1]*ne[2]*ne[3]); + ggml_set_name(targets, "targets"); + + ggml_tensor * out; + if (weighted) { + ggml_tensor * weights = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne[1]*ne[2]*ne[3]); + ggml_set_name(weights, "weights"); + out = ggml_cross_entropy_loss_weighted(ctx, logits, targets, weights); + } else { + out = ggml_cross_entropy_loss(ctx, logits, targets); + } + ggml_set_name(out, "out"); + + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + std::default_random_engine rng(7); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (strcmp(t->name, "weights") == 0) { + // a spread of relevance, zero included (a row the mind did not care about) + std::vector w(ggml_nelements(t)); + for (size_t i = 0; i < w.size(); ++i) { + w[i] = i % 4 == 1 ? 0.0f : 0.25f + 0.5f*(float) (i % 5); + } + ggml_backend_tensor_set(t, w.data(), 0, ggml_nbytes(t)); + } else if (t->type == GGML_TYPE_I32) { + // every third row masked (-1), the rest a random target + std::uniform_int_distribution pick(0, (int32_t) ne[0] - 1); + std::vector v(ggml_nelements(t)); + for (size_t i = 0; i < v.size(); ++i) { + v[i] = i % 3 == 2 ? -1 : pick(rng); + } + ggml_backend_tensor_set(t, v.data(), 0, ggml_nbytes(t)); + } else { + init_tensor_uniform(t, -3.0f, 3.0f); + } + } + } + + float grad_eps() override { + return 1.0f; + } + + bool grad_precise() override { + return true; + } +}; + struct test_cross_entropy_loss : public test_case { const ggml_type type; const std::array ne; @@ -8734,6 +8962,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I64, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I32, { 1, 8, 1, 3 }, { 1, 1 }, 2, false)); + test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_I32, { 64, 9, 1, 1 }, { 1, 1 }, 4, false, true)); test_cases.emplace_back(new test_set_rows(GGML_TYPE_F32, GGML_TYPE_Q8_0, GGML_TYPE_I32, { 256, 5, 1, 3 }, { 1, 1, }, 1, false)); for (ggml_type src_type : {GGML_TYPE_F16, GGML_TYPE_F32}) { for (ggml_type type : all_types) { @@ -9330,6 +9559,7 @@ static std::vector> make_test_cases_eval() { } test_cases.emplace_back(new test_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, false, eps, true)); test_cases.emplace_back(new test_rms_norm_back(GGML_TYPE_F32, { n, 5, 4, 3 }, eps)); + test_cases.emplace_back(new test_rms_norm_back(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, true)); test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false)); test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, true)); test_cases.emplace_back(new test_l2_norm(GGML_TYPE_F32, { n, 5, 4, 3 }, eps, false, true)); @@ -9774,6 +10004,19 @@ static std::vector> make_test_cases_eval() { type_a == GGML_TYPE_F32 ? 1024 : 256, 1024, 8, {1, 1}, {1, 1}, trans_b, true)); } } + // the batched, broadcast shape a delta-net backward hands OUT_PROD on the 4B (16 value heads + // over 4 key heads, head width 256 and 128), both gradient layouts + for (bool trans_b : {false, true}) { + test_cases.emplace_back(new test_out_prod(GGML_TYPE_F32, GGML_TYPE_F32, 256, 128, 256, {4, 1}, {4, 1}, trans_b)); + } + // a LoRA parameter gradient: the rank is the OUTPUT width (dA, dB = out_prod(x, grad) with + // m = rank), not the contraction; ranks below and between multiples of 16, at the 4B walk's + // shapes (8192 tokens, 128 and 1024 wide), both layouts + for (int64_t rank : {4, 8, 24}) { + test_cases.emplace_back(new test_out_prod(GGML_TYPE_F32, GGML_TYPE_F32, rank, 128, 8192, {1, 1}, {1, 1}, true)); + test_cases.emplace_back(new test_out_prod(GGML_TYPE_F32, GGML_TYPE_F32, rank, 8192, 128, {1, 1}, {1, 1}, false)); + test_cases.emplace_back(new test_out_prod(GGML_TYPE_F32, GGML_TYPE_F32, rank, 1024, 128, {1, 1}, {1, 1}, false)); + } // the q8_0 three-block case in both gradient layouts: the training graph hands // out_prod a TRANSPOSED gradient (out_prod(W, ggml_transpose(grad))) for (bool trans_b : {false, true}) { @@ -10124,6 +10367,13 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_pad()); test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {33, 17, 2, 1}, 4, 3, true)); // circular test_cases.emplace_back(new test_pad_ext()); + // left padding on one dim at a time, and on views (a training graph pads a strided view): + // Metal refused every left pad until it read src[i - lp] + for (int tfrm : {0, 1, 2}) { + test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {128, 64, 3, 2}, 3, 0, 0, 0, 0, 0, 0, 0, tfrm)); + test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {128, 64, 3, 2}, 0, 0, 5, 2, 0, 0, 0, 0, tfrm)); + test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {128, 64, 3, 2}, 0, 0, 0, 0, 2, 0, 1, 1, tfrm)); + } test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 1, 1, 1}, 1, 0, false)); test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 2, 1, 1}, 1, 0, false)); test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 16, 1, 1}, 0, 1, false)); @@ -10329,6 +10579,21 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3}, false)); test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, false)); + test_cases.emplace_back(new test_flash_attn_ext_grad(16, 4, 2, 24, 8, true)); + test_cases.emplace_back(new test_flash_attn_ext_grad(16, 4, 2, 24, 8, false)); + for (ggml_type t : { GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_Q8_0 }) { + test_cases.emplace_back(new test_flash_attn_back(128, 8, 2, 512, 64, t)); + } + test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 300, 37, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_back(512, 4, 2, 128, 9, GGML_TYPE_F16)); // the tiled path's tiles exceed 32 KiB: the per-row path + // dK/dV from a first key row: a frozen prefix, unaligned to any tile, and on the per-row path + test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 300, 37, GGML_TYPE_F16, 300 - 37)); + test_cases.emplace_back(new test_flash_attn_back(128, 8, 2, 512, 64, GGML_TYPE_Q8_0, 203)); + test_cases.emplace_back(new test_flash_attn_back(512, 4, 2, 128, 9, GGML_TYPE_F16, 100)); + test_cases.emplace_back(new test_cross_entropy_loss_index()); + test_cases.emplace_back(new test_cross_entropy_loss_index({30000, 2, 1, 1})); // a vocabulary's width; few rows keep the finite difference above f32 rounding + test_cases.emplace_back(new test_cross_entropy_loss_index({10, 5, 4, 3}, true)); + test_cases.emplace_back(new test_cross_entropy_loss_index({30000, 2, 1, 1}, true)); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); @@ -10517,6 +10782,13 @@ static std::vector> make_test_cases_eval() { static std::vector> make_test_cases_perf() { std::vector> test_cases; + // the attention backward beside its forward at one training shape (head 256, 16 heads over 4 KV + // heads, 256 rows over a 2048 window, f16 cache): small enough to time on every kernel edit; a + // backward costs about 2.5 forwards of matmul work + test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 2048, 256, GGML_TYPE_F16)); + test_cases.emplace_back(new test_flash_attn_back(256, 16, 4, 2048, 256, GGML_TYPE_F16, 2048 - 256)); // the reply's keys only + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {4, 1}, 2048, 256, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // OUT_PROD at the shapes a LoRA backward emits (dx = out_prod(W, dy^T), trans_b): a 1.5B // (n_embd 1536, ffn 8960) and a 27B (n_embd 5120, ffn 17408) at 1024 tokens, plus the // output head (vocab 151936) — the largest frozen matmul in the step. F32 beside Q4_K diff --git a/tests/test-batch-alloc.cpp b/tests/test-batch-alloc.cpp index 66d29d6f5164..963fb1a06b36 100644 --- a/tests/test-batch-alloc.cpp +++ b/tests/test-batch-alloc.cpp @@ -21,6 +21,8 @@ struct mock_memory : public llama_memory_i { llama_memory_context_ptr init_full() override { GGML_ASSERT(false && "not implemented"); } llama_memory_context_ptr init_update(llama_context *, bool) override { GGML_ASSERT(false && "not implemented"); } + const char * training_error(const llama_batch &, bool) const override { return nullptr; } + bool get_can_shift() const override { GGML_ASSERT(false && "not implemented"); } void clear(bool) override { GGML_ASSERT(false && "not implemented"); } @@ -96,6 +98,33 @@ struct batch_builder { static void test_init(testing & t) { llama_vocab vocab; + t.test("rejects_invalid_labels_before_memory_mutation", [&](testing & t) { + batch_builder bb; + bb.add(0, {0}, true); + mock_memory mem; + llama_batch batch = bb.make(); + llama_token label = -2; + batch.label = &label; + llama_batch_allocr ba(1); + t.assert_true("only -1 denotes no label", !ba.init(batch, vocab, &mem, bb.n_embd, 1, false)); + label = 0; + t.assert_true("target must belong to vocabulary", !ba.init(batch, vocab, &mem, bb.n_embd, 1, false)); + label = -1; + t.assert_true("unlabelled inference remains valid", ba.init(batch, vocab, &mem, bb.n_embd, 1, false)); + }); + + t.test("rejects_missing_or_empty_sequence_lists", [&](testing & t) { + batch_builder bb; + bb.add(0, {0}, true); + llama_batch batch = bb.make(); + llama_batch_allocr ba(1); + batch.n_seq_id = nullptr; + t.assert_true("missing counts", !ba.init(batch, vocab, nullptr, bb.n_embd, 1, false)); + batch = bb.make(); + batch.n_seq_id[0] = 0; + t.assert_true("empty sequence list", !ba.init(batch, vocab, nullptr, bb.n_embd, 1, false)); + }); + t.test("rejects_n_seq_max_too_large", [&](testing & t) { batch_builder bb; bb.add(0, {0}, true); diff --git a/tests/test-opt-dynamic-accum.cpp b/tests/test-opt-dynamic-accum.cpp index ff5445fda605..f72dce6ea54d 100644 --- a/tests/test-opt-dynamic-accum.cpp +++ b/tests/test-opt-dynamic-accum.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #define CHECK(cond) do { if (!(cond)) { fprintf(stderr, "FAILED %s:%d: %s\n", __FILE__, __LINE__, #cond); exit(1); } } while (0) @@ -181,11 +182,210 @@ static void run_manual() { ggml_backend_free(cpu); } +static void run_out_prod_broadcast(ggml_type type, bool transpose) { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + CHECK(cpu != nullptr); + ggml_init_params cp = { 32 * ggml_tensor_overhead() + ggml_graph_overhead(), nullptr, true }; + ggml_context * ctx = ggml_init(cp); + ggml_tensor * input = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 32, 2, 2, 2); + ggml_tensor * a = ggml_cast(ctx, input, type); + ggml_tensor * decoded = ggml_cast(ctx, a, GGML_TYPE_F32); + ggml_tensor * b_storage = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, transpose ? 2 : 1, transpose ? 1 : 2, 4, 4); + ggml_tensor * b = transpose ? ggml_transpose(ctx, b_storage) : b_storage; + ggml_tensor * out = ggml_out_prod(ctx, a, b); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, out); + ggml_build_forward_expand(graph, decoded); + CHECK(ggml_backend_supports_op(cpu, out)); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, cpu); + float av[32 * 2 * 2 * 2], bv[2 * 4 * 4]; + for (int batch = 0; batch < 4; ++batch) { + for (int row = 0; row < 2; ++row) { + for (int x = 0; x < 32; ++x) { + av[(batch * 2 + row) * 32 + x] = float((batch + 1) * (row + 1)); + } + } + } + for (int batch = 0; batch < 16; ++batch) { + bv[batch * 2] = float(batch + 1); + bv[batch * 2 + 1] = -float(batch + 2); + } + ggml_backend_tensor_set(input, av, 0, sizeof(av)); + ggml_backend_tensor_set(b_storage, bv, 0, sizeof(bv)); + CHECK(ggml_backend_graph_compute(cpu, graph) == GGML_STATUS_SUCCESS); + float actual[32 * 4 * 4]; + ggml_backend_tensor_get(out, actual, 0, sizeof(actual)); + ggml_backend_tensor_get(decoded, av, 0, sizeof(av)); + for (int z = 0; z < 4; ++z) { + for (int y = 0; y < 4; ++y) { + const int a_batch = (z / 2) * 2 + y / 2; + const int b_batch = z * 4 + y; + for (int x = 0; x < 32; ++x) { + const float expected = av[a_batch * 64 + x] * bv[b_batch * 2] + + av[a_batch * 64 + 32 + x] * bv[b_batch * 2 + 1]; + if (std::fabs(actual[b_batch * 32 + x] - expected) >= 1e-4f) { + fprintf(stderr, "out_prod type=%s transpose=%d y=%d z=%d actual=%g expected=%g\n", + ggml_type_name(type), transpose, y, z, actual[b_batch * 32 + x], expected); + } + CHECK(std::fabs(actual[b_batch * 32 + x] - expected) < 1e-4f); + } + } + } + ggml_backend_buffer_free(buf); + ggml_free(ctx); + ggml_backend_free(cpu); +} + +static void run_cache_view_gradient(ggml_type type) { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 1, GGML_DEFAULT_GRAPH_SIZE, false, true); + ggml_init_params sp = { 8 * ggml_tensor_overhead(), nullptr, true }; + ggml_context * storage = ggml_init(sp); + ggml_tensor * weights = ggml_new_tensor_2d(storage, GGML_TYPE_F32, 32, 1); + ggml_tensor * cache = ggml_new_tensor_2d(storage, type, 32, 4); + ggml_set_param(weights); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(storage, cpu); + float values[32]; + for (float & value : values) { value = 1.0f; } + ggml_backend_tensor_set(weights, values, 0, sizeof(values)); + ggml_backend_tensor_memset(cache, 0, 0, ggml_nbytes(cache)); + ggml_opt_params params = ggml_opt_default_params(sched, GGML_OPT_LOSS_TYPE_SUM); + params.optimizer = GGML_OPT_OPTIMIZER_TYPE_SGD; + params.get_opt_pars = sgd_pars; + ggml_opt_context_t opt = ggml_opt_init(params); + ggml_init_params cp = { GGML_DEFAULT_GRAPH_SIZE * ggml_tensor_overhead() + 4 * ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), nullptr, true }; + ggml_context * ctx = ggml_init(cp); + ggml_tensor * row = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_set_input(row); + ggml_tensor * written = ggml_set_rows(ctx, cache, weights, row); + ggml_tensor * view = ggml_view_2d(ctx, written, 32, 1, written->nb[1], 2 * written->nb[1]); + view = ggml_cast(ctx, view, GGML_TYPE_F32); + ggml_cgraph * graph = ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, true); + ggml_build_forward_expand(graph, view); + ggml_opt_prepare_alloc(opt, ctx, graph, row, view); + CHECK(ggml_opt_alloc(opt, true)); + const int32_t index = 2; + ggml_backend_tensor_set(row, &index, 0, sizeof(index)); + ggml_opt_eval(opt, nullptr); + CHECK(ggml_opt_refusal(opt)[0] == '\0'); + ggml_backend_tensor_get(weights, values, 0, sizeof(values)); + for (float value : values) { CHECK(std::fabs(value - (1.0f - LR)) < 1e-6f); } + ggml_opt_free(opt); + ggml_free(ctx); + ggml_backend_buffer_free(buf); + ggml_free(storage); + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); +} + +static void run_label_projection(bool projected) { + ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + CHECK(cpu != nullptr); + ggml_backend_t backends[] = { cpu }; + ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 1, GGML_DEFAULT_GRAPH_SIZE, false, true); + ggml_init_params sp = { 8 * ggml_tensor_overhead(), nullptr, true }; + ggml_context * storage = ggml_init(sp); + ggml_tensor * weights = ggml_new_tensor_2d(storage, GGML_TYPE_F32, 2, 4); + ggml_tensor * identity = ggml_new_tensor_2d(storage, GGML_TYPE_F32, 2, 2); + ggml_set_param(weights); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(storage, cpu); + const float before[] = { 8.0f, -8.0f, 0.5f, -0.5f, -8.0f, 8.0f, 0.25f, 0.75f }; + const float eye[] = { 1, 0, 0, 1 }; + ggml_backend_tensor_set(weights, before, 0, sizeof(before)); + ggml_backend_tensor_set(identity, eye, 0, sizeof(eye)); + ggml_opt_params params = ggml_opt_default_params(sched, GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS); + params.optimizer = GGML_OPT_OPTIMIZER_TYPE_SGD; + params.get_opt_pars = sgd_pars; + ggml_opt_context_t opt = ggml_opt_init(params); + ggml_init_params cp = { GGML_DEFAULT_GRAPH_SIZE * ggml_tensor_overhead() + 4 * ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), nullptr, true }; + ggml_context * graph_ctx = ggml_init(cp); + ggml_tensor * base = ggml_mul_mat(graph_ctx, identity, weights); + ggml_mul_mat_set_prec(base, GGML_PREC_F32); + ggml_tensor * rows = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 2); + ggml_tensor * public_rows = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 3); + ggml_set_input(rows); + ggml_set_input(public_rows); + ggml_tensor * selected = ggml_add(graph_ctx, ggml_get_rows(graph_ctx, base, rows), + ggml_get_rows(graph_ctx, ggml_scale(graph_ctx, base, 0.25f), rows)); + ggml_tensor * logits = ggml_set_rows(graph_ctx, ggml_dup(graph_ctx, base), selected, rows); + ggml_cgraph * gf = ggml_new_graph_custom(graph_ctx, GGML_DEFAULT_GRAPH_SIZE, true); + ggml_build_forward_expand(gf, logits); + ggml_tensor * public_logits = nullptr; + ggml_tensor * public_shadow = nullptr; + ggml_tensor * local_shadow[2] = {}; + if (projected) { + public_shadow = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 1); + ggml_set_input(public_shadow); + for (auto & local : local_shadow) { + local = ggml_new_tensor_1d(graph_ctx, GGML_TYPE_I32, 1); + ggml_set_input(local); + } + ggml_tensor * slices[2]; + CHECK(ggml_opt_project_head(graph_ctx, gf, weights, logits, rows, public_rows, rows, + public_shadow, local_shadow, 1, &public_logits, slices, 2) == 2); + ggml_opt_prepare_alloc_slices(opt, graph_ctx, gf, rows, slices, 2); + } else { + public_logits = ggml_get_rows(graph_ctx, logits, public_rows); + ggml_build_forward_expand(gf, public_logits); + ggml_opt_prepare_alloc(opt, graph_ctx, gf, rows, ggml_get_rows(graph_ctx, logits, rows)); + } + ggml_set_output(public_logits); + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + auto * node = ggml_graph_node(gf, i); + if (node->op == GGML_OP_MUL_MAT) { + CHECK(node->op_params[0] == GGML_PREC_F32); + } + } + CHECK(ggml_opt_alloc(opt, true)); + const int32_t labels[] = { 1, 3 }, public_ids[] = { 2, 1, 0 }, targets[] = { 0, 1 }; + ggml_backend_tensor_set(rows, labels, 0, sizeof(labels)); + ggml_backend_tensor_set(public_rows, public_ids, 0, sizeof(public_ids)); + if (projected) { + const int32_t public_label = 1, local_label = 0; + ggml_backend_tensor_set(public_shadow, &public_label, 0, sizeof(public_label)); + for (auto * local : local_shadow) { + ggml_backend_tensor_set(local, &local_label, 0, sizeof(local_label)); + } + } + ggml_backend_tensor_set(ggml_opt_labels(opt), targets, 0, sizeof(targets)); + ggml_opt_eval(opt, nullptr); + CHECK(ggml_opt_refusal(opt)[0] == '\0'); + float after[8], served[6]; + ggml_backend_tensor_get(weights, after, 0, sizeof(after)); + ggml_backend_tensor_get(public_logits, served, 0, sizeof(served)); + for (int row = 0; row < 4; ++row) { + const bool labelled = row == 1 || row == 3; + const float p0 = 1.0f / (1.0f + std::exp(1.25f * (before[2*row + 1] - before[2*row]))); + const float delta = labelled ? LR * 1.25f * ((row == 1 ? 1.0f : 0.0f) - p0) / 2 : 0; + CHECK(std::fabs(after[2*row] - before[2*row] - delta) < 1e-6f); + CHECK(std::fabs(after[2*row + 1] - before[2*row + 1] + delta) < 1e-6f); + } + for (int row = 0; row < 3; ++row) { + const float scale = public_ids[row] == 1 ? 1.25f : 1; + for (int col = 0; col < 2; ++col) { + CHECK(served[2*row + col] == before[2*public_ids[row] + col] * scale); + } + } + ggml_opt_free(opt); + ggml_free(graph_ctx); + ggml_backend_buffer_free(buf); + ggml_free(storage); + ggml_backend_sched_free(sched); + ggml_backend_free(cpu); +} int main() { ggml_backend_load_all(); run(1); run(2); run_manual(); + run_label_projection(false); + run_label_projection(true); + for (ggml_type type : { GGML_TYPE_F16, GGML_TYPE_Q8_0 }) { + run_cache_view_gradient(type); + run_out_prod_broadcast(type, false); + run_out_prod_broadcast(type, true); + } printf("test-opt-dynamic-accum: OK\n"); return 0; } diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 7e4dd1717a6e..9a446b7b86c1 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -21,6 +21,8 @@ #include "common.h" #include "llama.h" #include "../src/llama-adapter.h" +#include "../src/llama-memory-hybrid.h" +#include "../src/llama-memory-recurrent.h" #include #include @@ -95,8 +97,7 @@ static llama_context * make_ctx(const common_params & params, llama_model * mode cparams.n_ctx = g_window; cparams.n_batch = chunk; cparams.n_ubatch = chunk; - cparams.n_seq_max = 2; // the plain walk snapshots a recurrent state into a scratch sequence - cparams.kv_unified = true; // ...and the window keeps every cell + cparams.kv_unified = true; // the window keeps every cell // a quantized V cache needs flash attention (it stores V untransposed); training graphs take // the explicit path either way cparams.flash_attn_type = g_cache_type == GGML_TYPE_F32 ? LLAMA_FLASH_ATTN_TYPE_DISABLED : LLAMA_FLASH_ATTN_TYPE_ENABLED; @@ -142,7 +143,8 @@ static std::vector> g_layout; // one SGD step from the fresh adapter: the step itself, -lr * the gradient it took static std::vector step_delta(const common_params & params, llama_model * model, const std::string & init, const std::vector & tokens, const std::vector & labelled, - uint32_t chunk, bool exact, size_t host_budget = 0, size_t * host_bytes = nullptr) { + uint32_t chunk, bool exact, size_t host_budget = 0, size_t * host_bytes = nullptr, + std::vector * state_after = nullptr) { llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); GGML_ASSERT(adapter != nullptr); const std::vector before = adapter_params(adapter); @@ -179,6 +181,10 @@ static std::vector step_delta(const common_params & params, llama_model * } g_last_graph_bytes = llama_opt_graph_bytes(ctx); g_last_horizon = llama_opt_walk_horizon(ctx); + if (state_after) { + state_after->resize(llama_state_seq_get_size(ctx, 0)); + GGML_ASSERT(llama_state_seq_get_data(ctx, state_after->data(), state_after->size(), 0) == state_after->size()); + } ggml_opt_result_free(result); ggml_opt_dataset_free(dataset); llama_free(ctx); @@ -192,6 +198,30 @@ static std::vector step_delta(const common_params & params, llama_model * return d; } +// her stable self's cache after [0, r0) and then [r0, r1): no adapter, decoded in chunks from +// each start as the walk decodes them (a matmul's result depends on its batch size), so a walk +// that restores its run leaves exactly these bytes +static std::vector stable_state(const common_params & params, llama_model * model, + const std::vector & tokens, uint32_t r0, uint32_t r1, uint32_t chunk) { + llama_context * ctx = make_ctx(params, model, chunk); + llama_batch batch = llama_batch_init(chunk, 0, 1); + for (const auto & [p0, p1] : { std::make_pair(0u, r0), std::make_pair(r0, r1) }) { + for (uint32_t c0 = p0; c0 < p1; c0 += chunk) { + const uint32_t n = std::min(c0 + chunk, p1) - c0; + common_batch_clear(batch); + for (uint32_t i = 0; i < n; ++i) { + common_batch_add(batch, tokens[c0 + i], c0 + i, { 0 }, i + 1 == n); + } + GGML_ASSERT(llama_decode(ctx, batch) == 0); + } + } + std::vector state(llama_state_seq_get_size(ctx, 0)); + GGML_ASSERT(llama_state_seq_get_data(ctx, state.data(), state.size(), 0) == state.size()); + llama_batch_free(batch); + llama_free(ctx); + return state; +} + struct agreement { double cosine; double norm_ratio; @@ -242,6 +272,222 @@ static bool same_step(const agreement & r) { return r.cosine > 0.995 && std::fabs(r.norm_ratio - 1.0) < 0.02; } +static std::vector sequence_state(llama_context * ctx, llama_seq_id seq) { + std::vector bytes(llama_state_seq_get_size(ctx, seq)); + GGML_ASSERT(llama_state_seq_get_data(ctx, bytes.data(), bytes.size(), seq) == bytes.size()); + return bytes; +} + +static void mixed_decode(const common_params & params, llama_model * model, const std::string & init, + const std::vector & tokens, bool flash, ggml_type cache_type, uint32_t undo = 0) { + auto cp = common_context_params_to_llama(params); + cp.n_ctx = 128; + cp.n_batch = cp.n_ubatch = undo ? 8 : 2; + cp.n_seq_max = 2; + cp.kv_unified = true; + cp.n_rs_seq = undo; + cp.flash_attn_type = flash ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + cp.type_k = cp.type_v = cache_type; + llama_context * ctx = llama_init_from_model(model, cp); + GGML_ASSERT(ctx); + llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); + float scale = 1.0f; + GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); + llama_opt_params op{}; + op.param_filter = llama_opt_param_filter_all; + op.get_opt_pars = sgd_pars; + op.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; + op.adapter = adapter; + llama_batch batch = llama_batch_init(2, 0, 1); + common_batch_add(batch, tokens[0], 0, {0}, true); + common_batch_add(batch, tokens[1], 0, {1}, true); + GGML_ASSERT(llama_decode(ctx, batch) == 0); + const auto reference_neighbor = sequence_state(ctx, 1); + const auto reference_target = sequence_state(ctx, 0); + batch.label[0] = tokens[2]; + batch.pos[0] = batch.pos[1] = 1; + GGML_ASSERT(llama_decode(ctx, batch) != 0); + GGML_ASSERT(sequence_state(ctx, 0) == reference_target && sequence_state(ctx, 1) == reference_neighbor); + llama_memory_clear(llama_get_memory(ctx), true); + batch.pos[0] = batch.pos[1] = 0; + llama_opt_init(ctx, model, op); + const auto initial_weights = adapter_params(adapter); + batch.label[0] = tokens[2]; + GGML_ASSERT(llama_decode(ctx, batch) == 0); + GGML_ASSERT(llama_get_logits_ith(ctx, 0) && llama_get_logits_ith(ctx, 1)); + const auto neighbor = sequence_state(ctx, 1); + GGML_ASSERT(neighbor == reference_neighbor); + const auto trained = sequence_state(ctx, 0); + GGML_ASSERT(neighbor != trained); + const auto weights = adapter_params(adapter); + GGML_ASSERT(weights != initial_weights); + common_batch_clear(batch); + common_batch_add(batch, tokens[2], 1, {0}, true); + batch.label[0] = llama_vocab_n_tokens(llama_model_get_vocab(model)); + GGML_ASSERT(llama_decode(ctx, batch) != 0); + GGML_ASSERT(sequence_state(ctx, 0) == trained && sequence_state(ctx, 1) == neighbor); + GGML_ASSERT(adapter_params(adapter) == weights); + batch.label[0] = tokens[3]; + GGML_ASSERT(llama_decode(ctx, batch) == 0); + GGML_ASSERT(sequence_state(ctx, 1) == neighbor); + GGML_ASSERT(adapter_params(adapter) != weights); + if (undo) { + auto * memory = llama_get_memory(ctx); + auto * recurrent = dynamic_cast(memory); + if (auto * hybrid = dynamic_cast(memory)) { + recurrent = hybrid->get_mem_recr(); + } + GGML_ASSERT(recurrent); + recurrent->set_rs_idx(0, 1); + const auto pending = sequence_state(ctx, 0); + const auto pending_weights = adapter_params(adapter); + batch.pos[0] = 2; + GGML_ASSERT(llama_decode(ctx, batch) != 0); + GGML_ASSERT(sequence_state(ctx, 0) == pending && sequence_state(ctx, 1) == neighbor); + GGML_ASSERT(adapter_params(adapter) == pending_weights); + } + printf("mixed decode flash=%d cache=%s undo=%u: same-position sequences, neighbour preservation, admission refusal OK\n", flash, ggml_type_name(cache_type), undo); + llama_batch_free(batch); + llama_free(ctx); + llama_adapter_lora_free(adapter); +} + +static void multi_label_decode(const common_params & params, llama_model * model, const std::string & init, + const std::vector & tokens, bool flash, ggml_type cache_type) { + // weights empty: an unweighted batch (label_weight NULL) + auto step = [&](const std::vector & sequences, const std::vector & weights = {}) { + auto cp = common_context_params_to_llama(params); + cp.n_ctx = 128; + cp.n_batch = cp.n_ubatch = 2; + cp.n_seq_max = 2; + cp.kv_unified = true; + cp.flash_attn_type = flash ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + cp.type_k = cp.type_v = cache_type; + llama_context * ctx = llama_init_from_model(model, cp); + GGML_ASSERT(ctx); + llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); + float scale = 1.0f; + GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); + llama_opt_params op{}; + op.param_filter = llama_opt_param_filter_all; + op.get_opt_pars = sgd_pars; + op.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; + op.adapter = adapter; + llama_opt_init(ctx, model, op); + const auto before = adapter_params(adapter); + llama_batch batch = llama_batch_init(2, 0, 1); + for (size_t k = 0; k < sequences.size(); ++k) { + const int seq = sequences[k]; + common_batch_add(batch, tokens[seq], 0, {seq}, true); + batch.label[batch.n_tokens - 1] = tokens[seq + 2]; + if (!weights.empty()) { + batch.label_weight[batch.n_tokens - 1] = weights[k]; + } + } + float * kept_weights = batch.label_weight; + if (weights.empty()) { + batch.label_weight = nullptr; // unweighted: the loss every backend has + } + GGML_ASSERT(llama_decode(ctx, batch) == 0); + batch.label_weight = kept_weights; + auto delta = adapter_params(adapter); + for (size_t i = 0; i < delta.size(); ++i) { + delta[i] -= before[i]; + } + llama_batch_free(batch); + llama_free(ctx); + llama_adapter_lora_free(adapter); + return delta; + }; + auto mean = step({0}); + const auto second = step({1}); + for (size_t i = 0; i < mean.size(); ++i) { + mean[i] = (mean[i] + second[i]) / 2; + } + const auto joint = step({0, 1}); + GGML_ASSERT(same_step(compare("two labelled sequences vs mean independent steps", joint, mean))); + GGML_ASSERT(same_step(compare("two labelled sequences with reversed row order", step({1, 0}), joint))); + // what this catches: the per-row weight not reaching the loss, or reaching it wrong (each row + // counts w_i * its loss / rows, its gradient scaled by w_i; a mind's relevance per token) + auto scaled = [](std::vector v, float s) { for (float & x : v) { x *= s; } return v; }; + GGML_ASSERT(same_step(compare("weights 1,1 vs an unweighted batch", step({0, 1}, {1.0f, 1.0f}), joint))); + GGML_ASSERT(same_step(compare("weights 2,2 vs twice the unweighted step", step({0, 1}, {2.0f, 2.0f}), scaled(joint, 2.0f)))); + GGML_ASSERT(same_step(compare("weights 1,0 vs half of row 0's own step (a mean over 2 rows)", step({0, 1}, {1.0f, 0.0f}), scaled(step({0}), 0.5f)))); +} +static void sparse_head_decode(const common_params & params, llama_model * model, const std::string & init, + const std::vector & tokens) { + struct result { + std::vector delta, logits; + std::vector neighbour; + }; + auto step = [&](int public_mode, bool reversed, bool train = true) { + auto cp = common_context_params_to_llama(params); + cp.n_ctx = 128; + cp.n_batch = cp.n_ubatch = 4; + cp.n_seq_max = 2; + cp.kv_unified = true; + cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + auto * ctx = llama_init_from_model(model, cp); + auto * adapter = llama_adapter_lora_init(model, init.c_str()); + GGML_ASSERT(ctx && adapter); + float scale = 1.0f; + GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); + llama_opt_params op{}; + op.param_filter = llama_opt_param_filter_all; + op.get_opt_pars = sgd_pars; + op.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; + op.adapter = adapter; + if (train) { + llama_opt_init(ctx, model, op); + } + const auto before = adapter_params(adapter); + llama_batch batch = llama_batch_init(3, 0, 1); + int neighbour_row = -1; + const std::vector order = reversed ? std::vector{2, 0, 1} : std::vector{0, 1, 2}; + for (int row : order) { + common_batch_add(batch, tokens[row], row == 1 ? 1 : 0, {row == 2 ? 1 : 0}, + public_mode == 2 || (public_mode == 1 && row == 2)); + if (train && row == 1) { + batch.label[batch.n_tokens - 1] = tokens[3]; + } + if (row == 2) { + neighbour_row = batch.n_tokens - 1; + } + } + GGML_ASSERT(llama_decode(ctx, batch) == 0); + printf("sparse head allocation: train=%d public=%d reversed=%d cache=%s graph_bytes=%zu\n", + train, public_mode, reversed, ggml_type_name(cp.type_k), llama_opt_graph_bytes(ctx)); + result out; + out.delta = adapter_params(adapter); + for (size_t i = 0; i < before.size(); ++i) { + out.delta[i] -= before[i]; + } + if (public_mode > 0) { + const float * logits = llama_get_logits_ith(ctx, neighbour_row); + GGML_ASSERT(logits); + out.logits.assign(logits, logits + llama_vocab_n_tokens(llama_model_get_vocab(model))); + } + out.neighbour = sequence_state(ctx, 1); + llama_batch_free(batch); + llama_free(ctx); + llama_adapter_lora_free(adapter); + return out; + }; + const auto dense = step(2, false); + const auto serving = step(1, false, false); + GGML_ASSERT(serving.logits == dense.logits); + GGML_ASSERT(serving.neighbour == dense.neighbour); + for (bool reversed : {false, true}) { + for (int public_mode : {0, 1}) { + const auto sparse = step(public_mode, reversed); + GGML_ASSERT(same_step(compare("sparse public rows vs dense objective", sparse.delta, dense.delta))); + GGML_ASSERT(public_mode == 0 || sparse.logits == dense.logits); + GGML_ASSERT(sparse.neighbour == dense.neighbour); + } + } + printf("sparse head: label-only rows, public logits, neighbour bytes and row reordering OK\n"); +} + int main(int argc, char ** argv) { common_params params; if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { @@ -274,6 +520,38 @@ int main(int argc, char ** argv) { llama_free(tok_ctx); GGML_ASSERT(tokens.size() > REFUSAL_WINDOW + 1); + if (getenv("WALK_HEAD_ONLY")) { + sparse_head_decode(params, model, init, tokens); + if (llama_model_get_tensor(model, "output.weight")) { + const auto head_init = init + ".head.gguf"; + GGML_ASSERT(common_lora_write_fresh(model, head_init, 8, 16, { "output.weight" }, 42)); + sparse_head_decode(params, model, head_init, tokens); + std::remove(head_init.c_str()); + } else { + printf("output adapter: unavailable (model has no independent output.weight tensor)\n"); + } + std::remove(init.c_str()); + llama_model_free(model); + llama_backend_free(); + return 0; + } + + if (getenv("WALK_MIXED_ONLY")) { + multi_label_decode(params, model, init, tokens, false, GGML_TYPE_F32); + multi_label_decode(params, model, init, tokens, true, GGML_TYPE_Q8_0); + mixed_decode(params, model, init, tokens, false, GGML_TYPE_F32); + for (ggml_type cache : { GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0 }) { + mixed_decode(params, model, init, tokens, true, cache); + } + if (llama_model_is_hybrid(model) || llama_model_is_recurrent(model)) { + mixed_decode(params, model, init, tokens, true, GGML_TYPE_Q8_0, 2); + } + std::remove(init.c_str()); + llama_model_free(model); + llama_backend_free(); + return 0; + } + // MEASUREMENT MODE (WALK_MEASURE="2048,8192,..."): the device graph per chunk of a context + // reply window (the last 512 positions labelled), on the explicit path (F32 cache, no flash // attention) and through FLASH_ATTN_EXT (f16 cache, flash attention), at each window. Prints a @@ -291,7 +569,7 @@ int main(int argc, char ** argv) { const uint32_t chunk = getenv("WALK_MEASURE_CHUNK") ? (uint32_t) atoi(getenv("WALK_MEASURE_CHUNK")) : 256; g_recompute = getenv("WALK_MEASURE_RECOMPUTE") != nullptr; // per-layer recompute, as the served runs train g_walk_horizon = getenv("WALK_MEASURE_HORIZON") ? (uint32_t) atoi(getenv("WALK_MEASURE_HORIZON")) : 0; // 0 = the whole window - printf(" window chunk explicit MiB flash MiB ratio\n"); + printf(" window chunk explicit MiB flash MiB flash q8_0 MiB ratio\n"); for (const char * p = env; *p; ) { char * end = nullptr; const uint32_t w = (uint32_t) strtoul(p, &end, 10); @@ -300,22 +578,25 @@ int main(int argc, char ** argv) { g_window = w; std::vector reply(w + 1, 0); for (uint32_t i = w - 511; i <= w; ++i) reply[i] = 1; - double mib[2] = { 0, 0 }; - const ggml_type types[2] = { GGML_TYPE_F32, GGML_TYPE_F16 }; - for (int a = 0; a < 2; ++a) { + double mib[3] = { 0, 0, 0 }; + // explicit attention at an F32 cache; flash at an f16 cache; flash at q8_0 (her lane's cache) + const ggml_type types[3] = { GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q8_0 }; + for (int a = 0; a < 3; ++a) { g_cache_type = types[a]; step_delta(params, model, init, long_tokens, reply, chunk, true); mib[a] = g_last_graph_bytes / 1048576.0; } g_cache_type = GGML_TYPE_F32; - printf(" %6u %5u %12.1f %9.1f %5.2fx\n", w, chunk, mib[0], mib[1], mib[1] > 0 ? mib[0] / mib[1] : 0.0); + printf(" %6u %5u %12.1f %9.1f %14.1f %5.2fx\n", w, chunk, mib[0], mib[1], mib[2], mib[1] > 0 ? mib[0] / mib[1] : 0.0); fflush(stdout); } std::remove(init.c_str()); return 0; } + const bool recompute_only = getenv("WALK_RECOMPUTE_ONLY") != nullptr; int failures = 0; + if (!recompute_only) { { // 1. every position labelled: chunk = window is one graph, the true gradient std::vector all(WINDOW + 1, 1); @@ -368,8 +649,9 @@ int main(int argc, char ** argv) { } } { - // 2. context, then a reply (the walk's real case): the exact walk is invariant to how the - // window is chunked, context chunks included (they train through the surrogate alone) + // 2. context, then a reply (the walk's real case): the reply trains given the context as + // served, so only its own chunks train, and the exact walk is invariant to how the reply + // is chunked std::vector reply(WINDOW + 1, 0); for (uint32_t i = WINDOW / 2 + 37; i <= WINDOW; ++i) { reply[i] = 1; @@ -387,6 +669,144 @@ int main(int argc, char ** argv) { } } + { + // 9. LEARNING ON A LIVE SEQUENCE (llama_opt_learn, P2). What this catches: training that + // needs its own context, re-decodes her context, disturbs another slot, or leaves the shadow + // in her cache. A serving-shaped context holds two sequences; sequence 0 is her context as + // served, sequence 1 a neighbour. Her reply trains on sequence 0 from the served prefix: + // the step must be the walk's exact step for the same window (one chunk per run), her + // sequence must end as a stable decode of the reply would leave it, and the neighbour's + // state must not change by a byte. + std::vector reply(WINDOW + 1, 0); + for (uint32_t i = WINDOW / 2 + 37; i <= WINDOW; ++i) { + reply[i] = 1; + } + const llama_pos r0 = WINDOW / 2 + 36; // the first labelled position + const auto want_f32 = step_delta(params, model, init, tokens, reply, WINDOW, true); + g_cache_type = GGML_TYPE_F16; + const auto want_f16 = step_delta(params, model, init, tokens, reply, WINDOW, true); + g_cache_type = GGML_TYPE_F32; + + for (const bool streams : { false, true }) { + // unified F32 cache with explicit attention, then a stream per slot with an f16 cache and + // flash attention: how a serving lane runs (each slot attends over its own stream) + printf(" live sequence, %s:\n", streams ? "a stream per slot, f16 cache, flash attention" : "one unified F32 cache, explicit attention"); + auto two_seq_ctx = [&]() { + auto cp = common_context_params_to_llama(params); + cp.n_ctx = 2*WINDOW; cp.n_batch = WINDOW; cp.n_ubatch = WINDOW; cp.n_seq_max = 2; cp.kv_unified = !streams; + cp.flash_attn_type = streams ? LLAMA_FLASH_ATTN_TYPE_ENABLED : LLAMA_FLASH_ATTN_TYPE_DISABLED; + cp.type_k = streams ? GGML_TYPE_F16 : GGML_TYPE_F32; cp.type_v = cp.type_k; + return llama_init_from_model(model, cp); + }; + auto decode_span = [&](llama_context * c, llama_seq_id sq, const llama_token * t, llama_pos p0, int32_t n) { + llama_batch b = llama_batch_init(WINDOW, 0, 1); + common_batch_clear(b); + for (int32_t i = 0; i < n; ++i) { + common_batch_add(b, t[i], p0 + i, { sq }, i + 1 == n); + } + GGML_ASSERT(llama_decode(c, b) == 0); + llama_batch_free(b); + }; + auto seq_state = [&](llama_context * c, llama_seq_id sq) { + std::vector st(llama_state_seq_get_size(c, sq)); + GGML_ASSERT(llama_state_seq_get_data(c, st.data(), st.size(), sq) == st.size()); + return st; + }; + const llama_token * neighbour = tokens.data() + WINDOW + 1; // other text, beyond her window + + // the control: her context, then her reply, decoded plainly beside the same neighbour + llama_context * ctl = two_seq_ctx(); + decode_span(ctl, 1, neighbour, 0, 200); + decode_span(ctl, 0, tokens.data(), 0, r0); + decode_span(ctl, 0, tokens.data() + r0, r0, WINDOW - r0); + const std::vector want_her = seq_state(ctl, 0); + llama_free(ctl); + + llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); + const std::vector before = adapter_params(adapter); + llama_context * ctx = two_seq_ctx(); + float scale = 1.0f; + GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); + decode_span(ctx, 1, neighbour, 0, 200); + decode_span(ctx, 0, tokens.data(), 0, r0); // her context, as served + const std::vector neighbour_before = seq_state(ctx, 1); + llama_opt_params lopt{}; + lopt.param_filter = llama_opt_param_filter_all; + lopt.get_opt_pars = sgd_pars; + lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; + lopt.adapter = adapter; + llama_opt_init(ctx, model, lopt); + std::vector labels(WINDOW - r0); + for (int32_t i = 0; i < (int32_t) WINDOW - r0; ++i) { + labels[i] = tokens[r0 + i + 1]; + } + // refusals come before any change: her sequence byte-identical after each + { + const std::vector her_before = seq_state(ctx, 0); + const bool r_null = llama_opt_learn(ctx, 0, r0, tokens.data() + r0, nullptr, WINDOW - r0, nullptr); + const bool r_long = llama_opt_learn(ctx, 0, r0, tokens.data() + r0, labels.data(), 2*WINDOW, nullptr); + const bool r_more = llama_opt_learn(ctx, 0, r0 + 5, tokens.data() + r0, labels.data(), 10, nullptr); + const bool same = seq_state(ctx, 0) == her_before; + printf(" live sequence refusals: null labels %s, past the context %s, prefix not held %s; her sequence %s\n", + r_null ? "ACCEPTED" : "refused", r_long ? "ACCEPTED" : "refused", r_more ? "ACCEPTED" : "refused", + same ? "byte-identical" : "CHANGED"); + if (r_null || r_long || r_more || !same) { + fprintf(stderr, "FAILED: a refused live learn was accepted or changed her sequence\n"); + ++failures; + } + } + const bool learned = llama_opt_learn(ctx, 0, r0, tokens.data() + r0, labels.data(), WINDOW - r0, nullptr); + const std::vector her_after = seq_state(ctx, 0); + const std::vector neighbour_after = seq_state(ctx, 1); + llama_free(ctx); + const std::vector after = adapter_params(adapter); + llama_adapter_lora_free(adapter); + std::vector got(before.size()); + for (size_t i = 0; i < got.size(); ++i) { + got[i] = after[i] - before[i]; + } + printf(" live sequence: learned=%d \"%s\"\n", (int) learned, learned ? "" : "see the log"); + if (!learned || !same_step(compare("live sequence: her reply learned in place vs the walk's exact step", got, streams ? want_f16 : want_f32))) { + fprintf(stderr, "FAILED: learning on a live sequence is not the walk's step\n"); + ++failures; + } + const bool her_ok = her_after == want_her; + const bool nb_ok = neighbour_after == neighbour_before; + printf(" live sequence: her cache %s a stable decode of the reply; the neighbour %s\n", + her_ok ? "byte-identical to" : "DIFFERS from", nb_ok ? "byte-identical" : "CHANGED"); + if (!her_ok || !nb_ok) { + fprintf(stderr, "FAILED: learning on a live sequence disturbed the cache\n"); + ++failures; + } + } + } + + { + // 8. what this catches: a walk that leaves the shadow in her cache. The reply's cells + // are written under the shadow while it trains, then restored, so after either walk the + // sequence (every attention cell, the recurrent state) is byte-identical to her stable + // self's decode of the same window: serving never reads what training wrote. + std::vector reply(WINDOW + 1, 0); + for (uint32_t i = WINDOW / 2 + 37; i <= WINDOW; ++i) { + reply[i] = 1; + } + const uint32_t r0 = WINDOW / 2 + 36; // the first labelled position (it predicts reply[r0 + 1]) + const std::vector want = stable_state(params, model, tokens, r0, WINDOW, WINDOW / 4); + for (bool exact : { true, false }) { + std::vector got; + step_delta(params, model, init, tokens, reply, WINDOW / 4, exact, 0, nullptr, &got); + const bool same = got == want; + printf(" cache after the %s walk: %zu bytes, %s her stable decode\n", exact ? "exact" : "plain", got.size(), + same ? "byte-identical to" : "DIFFERS from"); + if (!same) { + fprintf(stderr, "FAILED: the %s walk left the shadow in the cache\n", exact ? "exact" : "plain"); + ++failures; + } + } + } + + } + { // 7. per-layer recompute under the exact walk. The chunk's activations, every layer's // kept for the backward pass, are what a large model's chunk graph is made of (Kimi's @@ -409,12 +829,15 @@ int main(int argc, char ** argv) { fprintf(stderr, "FAILED: recompute changed the exact walk's step\n"); ++failures; } - if (!(rebuilt_bytes < kept_bytes)) { + if (kept_bytes == 0 && rebuilt_bytes == 0) { + printf(" chunk graph reduction: unavailable (device-only allocation counter)\n"); + } else if (!(rebuilt_bytes < kept_bytes)) { fprintf(stderr, "FAILED: recompute did not shrink the exact walk's chunk graph\n"); ++failures; } } + if (!recompute_only) { { // 3. the host budget: a window whose accumulators and snapshots exceed it refuses by // name before anything runs, and the adapter is untouched @@ -463,36 +886,21 @@ int main(int argc, char ** argv) { } } if (llama_model_is_hybrid(model)) { - // 4. regression for the 5090 crash (2026-10-07 04:09Z): a recurrent context with rollback - // slots (serving's, inherited by the training context) made the exact walk ASSERT inside - // the serving process. It must refuse by name, and the adapter stays untouched. - llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); - const std::vector before = adapter_params(adapter); + // 4. TRAINING IN THE SERVING CONTEXT (Joel, 2026-10-13: "training is basically llama inference with + // backprop"). Serving keeps recurrent rollback slots for drafts; a training graph keeps one state, the + // live one (llm_graph_context::rs_keep), so a context WITH rollback slots takes the same step as one + // without. Was a refusal (the 5090 crash, 2026-10-07 04:09Z), which forced a second training + // context beside serving. + std::vector all(WINDOW + 1, 1); + const auto without = step_delta(params, model, init, tokens, all, WINDOW / 4, true); g_n_rs_seq = 4; - llama_context * ctx = make_ctx(params, model, WINDOW / 4); + const auto with = step_delta(params, model, init, tokens, all, WINDOW / 4, true); g_n_rs_seq = 0; - float scale = 1.0f; - GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); - llama_opt_params lopt{}; - lopt.param_filter = llama_opt_param_filter_all; - lopt.get_opt_pars = sgd_pars; - lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; - lopt.adapter = adapter; - lopt.walk_exact = true; - llama_opt_init(ctx, model, lopt); - std::vector> seqs = { std::vector(tokens.begin(), tokens.begin() + WINDOW + 1) }; - std::vector> loss = { std::vector(WINDOW + 1, 1) }; - ggml_opt_dataset_t dataset = common_opt_dataset_init_masked(WINDOW, seqs, loss, tokens[0]); - llama_opt_epoch(ctx, dataset, nullptr, nullptr, /*idata_split =*/ 1, nullptr, nullptr); - const std::string why = llama_opt_failure(ctx); - printf(" rollback slots: failed=%d, \"%s\"\n", (int) llama_opt_failed(ctx), why.c_str()); - if (!llama_opt_failed(ctx) || why.find("rollback") == std::string::npos || adapter_params(adapter) != before) { - fprintf(stderr, "FAILED: a recurrent context with rollback slots was not refused by name\n"); + if (!same_step(compare("rollback slots: exact walk in a context with n_rs_seq 4 vs 0", with, without))) { + fprintf(stderr, "FAILED: rollback slots changed the exact walk's step\n"); + worst_tensors(with, without, 24); ++failures; } - ggml_opt_dataset_free(dataset); - llama_free(ctx); - llama_adapter_lora_free(adapter); } { @@ -511,8 +919,12 @@ int main(int argc, char ** argv) { step_delta(params, model, init, tokens, all, WINDOW / 4, true); const size_t graph_bytes = g_last_graph_bytes; printf(" device graph: %.1f MiB per chunk\n", graph_bytes / 1048576.0); - if (graph_bytes == 0) { - printf(" device refusal: skipped (no device graph measured: a CPU-only run)\n"); + // the budget caps GPU memory only (ggml_opt_alloc), so a CPU-only run cannot refuse; it + // still measures a graph (in host buffers), which is why the graph size alone was the + // wrong test for "a CPU-only run" + const bool cpu_only = !params.devices.empty() && params.devices[0] == nullptr; + if (graph_bytes == 0 || cpu_only) { + printf(" device refusal: skipped (a CPU-only run: the device budget caps GPU memory only)\n"); } else { llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); const std::vector before = adapter_params(adapter); @@ -554,6 +966,8 @@ int main(int argc, char ** argv) { g_window = WINDOW; } + } + std::remove(init.c_str()); llama_model_free(model); llama_backend_free(); diff --git a/tests/test-walk-plan.cpp b/tests/test-walk-plan.cpp index ed3749131194..04d145cc6b4f 100644 --- a/tests/test-walk-plan.cpp +++ b/tests/test-walk-plan.cpp @@ -38,22 +38,30 @@ int main(int argc, char ** argv) { }); // what this catches: a step taken anywhere but the reverse pass's last chunk, a loss that is - // not the window's mean, a gradient reaching past the horizon + // not the window's mean, a gradient reaching past the horizon or before its own run (the + // context before a run is her stable set's: the shadow never touched it), a context chunk + // trained for nothing t.test("exact_marks", [](testing & t) { - auto plan = llama_walk_plan::of(window({ {1024, false}, {512, true}, {512, false}, {512, true} }), 2560, 512); - plan.mark_exact(/*recurrent =*/ false, /*horizon =*/ 512); + auto plan = llama_walk_plan::of(window({ {1024, false}, {1536, true}, {512, false}, {512, true} }), 3584, 512); + t.assert_equal(2u, plan.runs.size()); + plan.mark_exact(/*horizon =*/ 512); float loss = 0.0f; int ends = 0; for (const auto & c : plan.chunks) { - t.assert_true("every chunk before a reply carries its gradient", c.needed); - t.assert_equal(llama_walk_plan::grad_from(c.c0, 512), c.grad_from); - loss += c.loss_scale; + t.assert_true("exactly the labelled chunks train", c.needed == (c.n_labels > 0)); + loss += c.needed ? c.loss_scale : 0.0f; ends += c.period_end; } - t.assert_equal(0, plan.first); + // run 1 is chunks 2..4 (positions 1024..2560), run 2 is chunk 6 (3072..3584) + t.assert_equal(1024u, plan.chunks[2].grad_from); + t.assert_equal(1024u, plan.chunks[3].grad_from); + t.assert_equal(1536u, plan.chunks[4].grad_from); + t.assert_equal(3072u, plan.chunks[6].grad_from); + t.assert_true("a chunk a later one of its run reaches", plan.chunks[2].surrogate && plan.chunks[3].surrogate); + t.assert_true("the last chunk of a run has no later gradient", !plan.chunks[4].surrogate && !plan.chunks[6].surrogate); + t.assert_equal(6, plan.first); t.assert_equal(1, ends); t.assert_true("loss scales sum to one", loss > 0.9999f && loss < 1.0001f); - t.assert_true("the last chunk has no later gradient", !plan.chunks.back().surrogate); t.assert_equal(0u, llama_walk_plan::grad_from(1024, 0)); t.assert_equal(512u, llama_walk_plan::grad_from(1024, 512)); }); @@ -61,13 +69,14 @@ int main(int argc, char ** argv) { // what this catches: a stride that does not shrink the checkpoints to the budget, or one // that grows past the window when nothing fits t.test("host_budget_stride", [](testing & t) { + // a run's start is a checkpoint too: the walk restores the run to her stable self from it const auto plan = llama_walk_plan::of(window({ {8*512, true} }), 8*512, 512); - t.assert_equal(170u, plan.host_bytes(1, 100, 10)); + t.assert_equal(180u, plan.host_bytes(1, 100, 10)); t.assert_equal(1u, plan.stride_for(0, 100, 10)); t.assert_equal(2u, plan.stride_for(150, 100, 10)); - t.assert_equal(130u, plan.host_bytes(2, 100, 10)); + t.assert_equal(140u, plan.host_bytes(2, 100, 10)); t.assert_equal(8u, plan.stride_for(50, 100, 10)); - t.assert_equal(100u, plan.host_bytes(8, 100, 10)); + t.assert_equal(110u, plan.host_bytes(8, 100, 10)); }); return t.summary(); diff --git a/tools/batched-bench/batched-bench.cpp b/tools/batched-bench/batched-bench.cpp index e2dcd0b2e717..4d39af212aa6 100644 --- a/tools/batched-bench/batched-bench.cpp +++ b/tools/batched-bench/batched-bench.cpp @@ -91,6 +91,8 @@ int llama_batched_bench(int argc, char ** argv) { batch.n_seq_id + i, batch.seq_id + i, batch.logits + i, + batch.label ? batch.label + i : nullptr, + batch.label && batch.label_weight ? batch.label_weight + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); diff --git a/tools/mtmd/mtmd-helper-common.h b/tools/mtmd/mtmd-helper-common.h index f907346c7b58..dab628d1a0d6 100644 --- a/tools/mtmd/mtmd-helper-common.h +++ b/tools/mtmd/mtmd-helper-common.h @@ -96,6 +96,8 @@ struct decode_embd_batch { /*n_seq_id =*/ n_seq_id.data(), /*seq_id =*/ seq_ids.data(), /*logits =*/ logits.data(), + /*label =*/ nullptr, + /*label_weight =*/ nullptr, }; } @@ -179,6 +181,8 @@ struct decode_embd_batch { /*n_seq_id =*/ batch.n_seq_id + offset, /*seq_id =*/ batch.seq_id + offset, /*logits =*/ batch.logits + offset, + /*label =*/ batch.label ? batch.label + offset : nullptr, + /*label_weight =*/ batch.label && batch.label_weight ? batch.label_weight + offset : nullptr, }; } }; diff --git a/tools/perplexity/perplexity.cpp b/tools/perplexity/perplexity.cpp index ba41287d8e3d..a7e12041ea9d 100644 --- a/tools/perplexity/perplexity.cpp +++ b/tools/perplexity/perplexity.cpp @@ -674,6 +674,8 @@ static bool decode_helper(llama_context * ctx, llama_batch & batch, std::vector< batch.n_seq_id + i, batch.seq_id + i, batch.logits + i, + batch.label ? batch.label + i : nullptr, + batch.label && batch.label_weight ? batch.label_weight + i : nullptr, }; const int ret = llama_decode(ctx, batch_view); diff --git a/tools/server/server-chat.cpp b/tools/server/server-chat.cpp index a6fe3c6ba619..c0c8712c46c0 100644 --- a/tools/server/server-chat.cpp +++ b/tools/server/server-chat.cpp @@ -576,7 +576,9 @@ json server_chat_convert_anthropic_to_oai(const json & body) { } // Pass through common params - for (const auto & key : {"temperature", "top_p", "top_k", "stream", "chat_template_kwargs"}) { + // "learn" too: a learn on an Anthropic request was dropped here, so it never ran (a reply learned + // on the turn's own slot, task_params::learn) + for (const auto & key : {"temperature", "top_p", "top_k", "stream", "chat_template_kwargs", "learn"}) { if (body.contains(key)) { oai_body[key] = body.at(key); } diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7294a8d876e3..6972228d5c96 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -205,6 +205,8 @@ struct server_batch { batch.n_seq_id + off, batch.seq_id + off, batch.logits + off, + batch.label ? batch.label + off : nullptr, + batch.label && batch.label_weight ? batch.label_weight + off : nullptr, }; return view; @@ -346,6 +348,7 @@ struct server_slot { int32_t n_gen_last = 0; void reset() { + learned.reset(); SLT_DBG(*this, "%s", "\n"); spec_is_replay = false; @@ -520,6 +523,9 @@ struct server_slot { prompt.tokens.insert(spec_draft); } + // the outcome of the finishing task's learn (task_params::learn), for its final response + std::optional learned; + void release() { if (is_processing()) { GGML_ASSERT(task); @@ -891,6 +897,23 @@ struct server_context_impl { common_speculative_init_result_ptr spec_init; common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; + + // LEARNING IN THE SERVING CONTEXT (SERVER_TASK_TYPE_LEARN): the shadow adapter llama_opt_init was + // given (one per server run for now), and the learning rate its next step reads + int32_t learn_lora_id = -1; + float learn_lr = 1e-4f; + // LEARNING BARELY TOUCHES THE FLOW: while any other slot is working, learning may take at most + // LEARN_DUTY of the loop's time (the step after a learn of t ms waits t*(1/LEARN_DUTY - 1) ms); + // with every other slot idle it runs at once. Back to back, learns took the 4B's other slot from + // p50 218 / p95 450 ms per token to 686 / 2,567 (M5, 2026-10-15). + static constexpr double LEARN_DUTY = 0.10; + int64_t learn_next_us = 0; + static ggml_opt_optimizer_params learn_opt_pars(void * ud) { + ggml_opt_optimizer_params p = ggml_opt_get_default_optimizer_params(nullptr); + p.adamw.alpha = *(const float *) ud; + p.adamw.wd = 0.0f; + return p; + } common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_speculative_ptr spec; @@ -2039,11 +2062,165 @@ struct server_context_impl { queue_results.send(std::move(res)); } + // One reply learned on `slot`, which holds the served prompt (POST /learn between its turns, or a + // completion's own learn before its slot is released): the slot's stable set plus the shadow on + // the reply's labelled rows, then the slot back to its whole served prompt. + learn_outcome learn_on_slot(server_slot & slot, const learn_request & L, int64_t t0) { + auto refused = [&](const std::string & why) { + learn_outcome out; + out.failure = why; + SRV_WRN("learn refused: %s\n", why.c_str()); + return out; + }; + // the same refusal /train gives: repacked CPU weights fail the backward (x86) or crash + // the copy to Metal (an M5 at -ngl 0), card 5b497852 + if (const std::string why = train_repack_refusal(model_tgt, params_base); !why.empty()) { + return refused(why); + } + auto & loras = params_base.lora_adapters; + if (L.lora_id < 0 || (size_t) L.lora_id >= loras.size() || loras[L.lora_id].ptr == nullptr) { + return refused("no loaded adapter " + std::to_string(L.lora_id) + " to train (POST /lora-adapters/load it first)"); + } + if (learn_lora_id >= 0 && learn_lora_id != L.lora_id) { + return refused("this server trains adapter " + std::to_string(learn_lora_id) + "; one shadow per server run"); + } + const int32_t n0 = L.n_prefix - 1; // the last prompt token is re-decoded with the reply's first token as its label + if (n0 < 0 || L.n_prefix > slot.prompt.n_tokens() || L.reply.empty()) { + return refused("slot " + std::to_string(slot.id) + " holds " + std::to_string(slot.prompt.n_tokens()) + + " prompt tokens; the served prefix is " + std::to_string(L.n_prefix) + " and the reply must not be empty"); + } + // the served prompt, when the request carries it: her slot must still hold exactly it (a + // slot that has served other turns since, or compacted, would teach the wrong context) + if (!L.prompt.empty()) { + int32_t i = 0; + while (i < L.n_prefix && slot.prompt.tokens[i] == L.prompt[i]) { + ++i; + } + if (i < L.n_prefix) { + return refused("slot " + std::to_string(slot.id) + " no longer holds this turn's prompt (it differs at token " + + std::to_string(i) + " of " + std::to_string(L.n_prefix) + "): serve the prompt again, then learn"); + } + } + // the recurrent part cannot be cut back: her state one token before the prompt's end + const common_prompt_checkpoint * cp = nullptr; + const bool needs_cp = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL || ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS; + if (needs_cp) { + for (const auto & c : slot.prompt.checkpoints) { + if (c.n_tokens == n0) { + cp = &c; + } + } + if (cp == nullptr) { + return refused("slot " + std::to_string(slot.id) + " has no checkpoint one token before its prompt's end (" + + std::to_string(n0) + "): the lane learns only with --train-dir, which keeps one"); + } + } + // the tokens from the last prompt token on, each labelled with the next + std::vector toks; + toks.push_back(slot.prompt.tokens[n0]); + toks.insert(toks.end(), L.reply.begin(), L.reply.end()); + std::vector labels(toks.size(), -1); + for (size_t i = 0; i + 1 < toks.size(); ++i) { + labels[i] = toks[i + 1]; + } + // her stable set plus the shadow: the shadow applies to labelled rows only + std::vector set = slot.lora.empty() ? construct_lora_list({}) : slot.lora; + if ((size_t) L.lora_id >= set.size()) { + set = construct_lora_list({}); + } + set[L.lora_id].scale = 1.0f; + common_set_adapter_lora(ctx_tgt, set); + learn_lr = L.lr; + if (learn_lora_id < 0) { + llama_opt_params lopt{}; + lopt.param_filter = llama_opt_param_filter_all; + lopt.get_opt_pars = learn_opt_pars; + lopt.get_opt_pars_ud = &learn_lr; + lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_ADAMW; + lopt.adapter = loras[L.lora_id].ptr; + llama_opt_init(ctx_tgt, model_tgt, lopt); + learn_lora_id = L.lora_id; + } + // drop what follows n0 (attention) and put her recurrent state back at n0 + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + ggml_opt_result_t result = ggml_opt_result_init(); + const bool ok = llama_opt_learn(ctx_tgt, slot.id, n0, toks.data(), labels.data(), (int32_t) toks.size(), result); + double loss = 0.0, unc = 0.0; + if (ok) { + ggml_opt_result_loss(result, &loss, &unc); + } + ggml_opt_result_free(result); + // her own adapter set again before anything else is decoded for her (the shadow is + // already off for a batch without labels: llama-graph.cpp lora_on, n_shadow_rows = + // the ubatch's label count; this does not rely on it) + common_set_adapter_lora(ctx_tgt, slot.lora); + // her slot goes back to her whole served prompt (attention cells, recurrent state, + // tokens and checkpoints agree; her next turn reuses all of it): back to n0 from the + // checkpoint, then the last prompt token decoded again under her stable set. A draft + // model's cache is cut to n0 with its own checkpoint, and the slot's tokens with it. + llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot.id, n0, -1); + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + int32_t n_keep = n0; + if (ctx_dft != nullptr) { + llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, n0, -1); + if (cp != nullptr) { + cp->load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + } else { + llama_batch one = llama_batch_get_one(&toks[0], 1); + llama_pos pos = n0; + one.pos = &pos; + llama_seq_id sid = slot.id; + int32_t nsid = 1; + llama_seq_id * sids = &sid; + one.n_seq_id = &nsid; + one.seq_id = &sids; + if (llama_decode(ctx_tgt, one) == 0) { + n_keep = L.n_prefix; + } else { + llama_memory_seq_rm(llama_get_memory(ctx_tgt), slot.id, n0, -1); + if (cp != nullptr) { + cp->load_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + } + } + } + slot.prompt.tokens.keep_first(n_keep); + slot.prompt.checkpoints.remove_if([&](const common_prompt_checkpoint & c) { return c.n_tokens > n_keep; }); + learn_outcome out; + out.ok = ok; + out.failure = ok ? "" : llama_opt_failure(ctx_tgt); + out.loss = loss; + out.n_labels = (int32_t) toks.size() - 1; + out.t_ms = (ggml_time_us() - t0) / 1000.0; + SRV_INF("learned on slot %d: %d labels from position %d, loss %.4f, %.1f ms%s\n", slot.id, out.n_labels, n0, loss, + out.t_ms, ok ? "" : (" (failed: " + out.failure + ")").c_str()); + return out; + } + + // A completion that carries a learn learns its reply on its own slot as it finishes, before the + // slot is released: no other request can take the slot between the turn and its learning (on + // a one-slot lane that is busy almost always, a separate POST /learn lost that race 12 times in + // 12 on the M5's 27B, 2026-10-15). + void learn_if_asked(server_slot & slot) { + if (!slot.task->params.learn) { + return; + } + learn_request L = *slot.task->params.learn; + L.id_slot = slot.id; + L.n_prefix = slot.task->n_tokens(); + slot.learned = learn_on_slot(slot, L, ggml_time_us()); + } + void send_final_response(server_slot & slot) { auto res = std::make_unique(); res->id = slot.task->id; res->id_slot = slot.id; + res->learned = slot.learned; res->index = slot.task->index; @@ -2854,6 +3031,42 @@ struct server_context_impl { res->path = la.path; queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_LEARN: + { + auto res = std::make_unique(); + res->id = task.id; + const auto & L = task.learn; + const int64_t t0 = ggml_time_us(); + server_slot * slot = get_slot_by_id(L.id_slot); + auto refuse = [&](const std::string & why) { + res->outcome.ok = false; + res->outcome.failure = why; + SRV_WRN("learn refused: %s\n", why.c_str()); + queue_results.send(std::move(res)); + }; + if (slot == nullptr) { + refuse("no slot " + std::to_string(L.id_slot)); + break; + } + if (slot->is_processing()) { + // her slot is serving: learn between its turns, never during one + queue_tasks.defer(std::move(task)); + break; + } + { + bool others_busy = false; + for (const auto & other : slots) { + others_busy = others_busy || (other.id != slot->id && other.is_processing()); + } + if (others_busy && ggml_time_us() < learn_next_us) { + queue_tasks.defer(std::move(task)); // within the duty cycle's rest + break; + } + } + res->outcome = learn_on_slot(*slot, L, t0); + learn_next_us = ggml_time_us() + (int64_t) ((ggml_time_us() - t0) * (1.0/LEARN_DUTY - 1.0)); + queue_results.send(std::move(res)); + } break; case SERVER_TASK_TYPE_UNLOAD_LORA: { auto & loras = params_base.lora_adapters; @@ -3364,6 +3577,11 @@ struct server_context_impl { SLT_WRN(slot, "%s", "empty prompt - releasing slot\n"); slot.print_timings(); + if (slot.task->params.learn) { + // a learn that could not run is said, not omitted (Cormac's review) + slot.learned = learn_outcome{}; + slot.learned->failure = "the request had no prompt to learn the reply after"; + } send_final_response(slot); slot.release(); @@ -3785,8 +4003,11 @@ struct server_context_impl { // - 4 + n_ubatch // - 4 // ref: https://github.com/ggml-org/llama.cpp/pull/20288 + // A lane that learns (--train-dir) also checkpoints ONE token before the end: + // learning a reply re-decodes the last prompt token with the reply's first + // token as its label, from her recurrent state exactly there (llama_opt_learn) if (do_checkpoint) { - static const int checkpoint_offsets[] = {4 + n_ubatch, 4}; + const int checkpoint_offsets[] = {4 + n_ubatch, 4, params_base.train_dir.empty() ? 4 : 1}; bool should_break = false; for (int offset : checkpoint_offsets) { @@ -4110,6 +4331,7 @@ struct server_context_impl { if (!process_token(result, slot)) { // release slot because of stop condition slot.print_timings(); + learn_if_asked(slot); send_final_response(slot); slot.release(); @@ -4231,6 +4453,7 @@ struct server_context_impl { if (!process_token(result, slot)) { slot.print_timings(); + learn_if_asked(slot); send_final_response(slot); slot.release(); @@ -4517,6 +4740,21 @@ void server_context::set_state_callback(server_state_callback_t callback) { // server_routes // +// The tokens a parsed request is served as: a chat request with a projector loaded goes through +// mtmd (its media markers and files), everything else through tokenize_input_prompts. One place, so +// what /learn checks a slot against is what serving put there by construction. +static std::vector serving_prompt_tokens(const llama_vocab * vocab, mtmd_context * mctx, const json & prompt, + const std::vector & files, bool chat) { + if (chat && mctx != nullptr) { + // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. + std::vector out; + out.push_back(process_mtmd_prompt(mctx, prompt.get(), files)); + return out; + } + // Everything else, including multimodal completions. + return tokenize_input_prompts(vocab, mctx, prompt, true, true); +} + std::unique_ptr server_routes::handle_completions_impl( const server_http_req & req, server_task_type type, @@ -4552,15 +4790,8 @@ std::unique_ptr server_routes::handle_completions_impl( } // process prompt - std::vector inputs; - - if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) { - // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. - inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files)); - } else { - // Everything else, including multimodal completions. - inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true); - } + std::vector inputs = serving_prompt_tokens(ctx_server.vocab, ctx_server.mctx, prompt, files, + res_type != TASK_RESPONSE_TYPE_NONE); // tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks @@ -5606,6 +5837,89 @@ void server_routes::init_routes() { return res; }; + // POST /learn: one reply learned on the slot that served its prompt, between that slot's turns, in + // this serving context (SERVER_TASK_TYPE_LEARN). Body: id_slot, n_prefix (the served prompt's + // tokens), reply (text) or reply_tokens, lora_id (a loaded adapter: the shadow), lr. + this->post_learn = [this](const server_http_req & req) { + auto res = create_response(); + if (params.train_dir.empty()) { + res->error(format_error_response("learning is off: start the server with --train-dir", ERROR_TYPE_NOT_SUPPORTED)); + return res; + } + const json body = json::parse(req.body); + server_task task(SERVER_TASK_TYPE_LEARN); + task.learn.id_slot = json_value(body, "id_slot", -1); + task.id_slot = task.learn.id_slot; // a deferred learn comes back when her slot frees + task.learn.n_prefix = json_value(body, "n_prefix", 0); + task.learn.lora_id = json_value(body, "lora_id", -1); + task.learn.lr = json_value(body, "lr", 1e-4f); + if (body.contains("messages")) { + // the conversation she was served, its last message her reply (or the correction). The + // served prompt is every earlier message through the SAME parser and template arguments + // /v1/chat/completions used (chat_template_kwargs included: a thinking template renders a + // different prompt without them), with the generation prompt; the reply is what the model + // would have generated after it: the text's tokens and the end of the turn + const json & messages = body.at("messages"); + if (!messages.is_array() || messages.size() < 2 || !messages.back().is_object() + || messages.back().value("role", std::string()) != "assistant") { + res->error(format_error_response("\"messages\" must end with the assistant reply to learn, after the conversation it answers", + ERROR_TYPE_INVALID_REQUEST)); + return res; + } + // the reply is learned as its text then the end of the turn; a reply whose substance is + // elsewhere (tool calls, reasoning) would teach her to end the turn instead (Cormac's + // review of #5061), so it is refused until those render through the template + const json & last = messages.back(); + const json no_value; + const json & calls = last.contains("tool_calls") ? last.at("tool_calls") : no_value; + const json & reasoning = last.contains("reasoning_content") ? last.at("reasoning_content") : no_value; + const bool has_calls = calls.is_array() && !calls.empty(); + const bool has_reasoning = reasoning.is_string() && !reasoning.get().empty(); + const bool has_text = last.contains("content") && last.at("content").is_string() && !last.at("content").get().empty(); + if (has_calls || has_reasoning || !has_text) { + res->error(format_error_response(has_calls ? "a reply with tool_calls cannot be learned yet (only its text would be: it would teach ending the turn)" + : has_reasoning ? "a reply with reasoning_content cannot be learned yet (its reasoning would be dropped)" + : "the reply to learn has no text content", + ERROR_TYPE_INVALID_REQUEST)); + return res; + } + json served = body; + served["messages"] = json::array(); + for (size_t i = 0; i + 1 < messages.size(); ++i) { + served["messages"].push_back(messages[i]); + } + std::vector files; + const json parsed = oaicompat_chat_params_parse(served, meta->chat_params, files); + if (!files.empty()) { + res->error(format_error_response("a conversation with media cannot be learned yet: only text prompts", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + task.learn.prompt = serving_prompt_tokens(ctx_server.vocab, ctx_server.mctx, parsed.at("prompt"), files, true)[0].get_text_tokens(); + task.learn.n_prefix = (int32_t) task.learn.prompt.size(); + task.learn.reply = common_tokenize(ctx_server.vocab, last.at("content").get(), false, true); + const llama_token eot = llama_vocab_eot(ctx_server.vocab) != LLAMA_TOKEN_NULL ? llama_vocab_eot(ctx_server.vocab) : llama_vocab_eos(ctx_server.vocab); + task.learn.reply.push_back(eot); + } else if (body.contains("reply_tokens")) { + task.learn.reply = body.at("reply_tokens").get>(); + } else { + task.learn.reply = common_tokenize(ctx_server.vocab, json_value(body, "reply", std::string()), false, true); + } + auto & rd = res->rd; + task.id = rd.get_new_id(); + rd.post_task(std::move(task)); + auto result = rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + res->ok(result->to_json()); + return res; + }; + this->post_lora_adapters_unload = [this](const server_http_req & req) { auto res = create_response(); const json body = json::parse(req.body); diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 349904972d8d..09bbc2eda320 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -167,6 +167,7 @@ struct server_routes { server_http_context::handler_t post_lora_adapters; server_http_context::handler_t post_lora_adapters_load; server_http_context::handler_t post_lora_adapters_unload; + server_http_context::handler_t post_learn; server_http_context::handler_t get_train; server_http_context::handler_t post_train; server_http_context::handler_t post_train_cancel; diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 64b9251295ce..828eff7285d3 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -546,6 +546,22 @@ task_params eval_llama_cmpl_schema( f->eval(ctx, data); } + // learn a reply on this request's slot before it is released: {"reply": text, "lora_id", "lr"}; + // the reply is learned as the tokens the model would generate after this prompt (its text, then + // the end of the turn) + if (data.contains("learn")) { + const json & l = data.at("learn"); + if (!l.is_object() || !l.contains("reply") || !l.at("reply").is_string() || l.at("reply").get().empty()) { + throw std::invalid_argument("\"learn\" must be an object with a non-empty \"reply\" string"); + } + learn_request L; + L.lora_id = json_value(l, "lora_id", -1); + L.lr = json_value(l, "lr", 1e-4f); + L.reply = common_tokenize(vocab, l.at("reply").get(), false, true); + L.reply.push_back(llama_vocab_eot(vocab) != LLAMA_TOKEN_NULL ? llama_vocab_eot(vocab) : llama_vocab_eos(vocab)); + params.learn = std::move(L); + } + // post-processing { // if "reasoning_format" is not provided, its handler will not be called, we will need to handle it here diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index aeb6ea508511..e239b39f2706 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -359,6 +359,9 @@ json server_task_result_cmpl_final::to_json_non_oaicompat() { if (!stream && !probs_output.empty()) { res["completion_probabilities"] = completion_token_output::probs_vector_to_json(probs_output, post_sampling_probs); } + if (learned) { + res["learn"] = learned->to_json(); + } return response_fields.empty() ? res : json_get_nested_values(response_fields, res); } @@ -404,6 +407,9 @@ json server_task_result_cmpl_final::to_json_oaicompat() { if (verbose) { res["__verbose"] = to_json_non_oaicompat(); } + if (learned) { + res["learn"] = learned->to_json(); + } if (stats.is_set()) { res["timings"] = stats.to_json(); } @@ -452,6 +458,9 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat() { if (verbose) { res["__verbose"] = to_json_non_oaicompat(); } + if (learned) { + res["learn"] = learned->to_json(); + } if (stats.is_set()) { res["timings"] = stats.to_json(); } @@ -517,6 +526,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_chat_stream() { deltas.back()["timings"] = stats.to_json(); } + if (learned) { + deltas.back()["learn"] = learned->to_json(); // the final chunk + } + // extra fields for debugging purposes if (verbose && !deltas.empty()) { deltas.front()["__verbose"] = to_json_non_oaicompat(); @@ -593,6 +606,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp() { }}, }; + if (learned) { + res["learn"] = learned->to_json(); + } + return res; } @@ -710,6 +727,10 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp_stream() { server_sent_events.back().at("data")["timings"] = stats.to_json(); } + if (learned) { + server_sent_events.back().at("data")["learn"] = learned->to_json(); // response.completed + } + return server_sent_events; } @@ -791,6 +812,10 @@ json server_task_result_cmpl_final::to_json_anthropic() { }} }; + if (learned) { + res["learn"] = learned->to_json(); + } + return res; } @@ -972,6 +997,10 @@ json server_task_result_cmpl_final::to_json_anthropic_stream() { }} }); + if (learned) { + events.back().at("data")["learn"] = learned->to_json(); // message_delta, beside the usage + } + events.push_back({ {"event", "message_stop"}, {"data", { @@ -1693,6 +1722,14 @@ json server_task_result_lora_change::to_json() { return res; } +json learn_outcome::to_json() const { + json res = { { "success", ok }, { "loss", loss }, { "n_labels", n_labels }, { "t_ms", t_ms } }; + if (!ok) { + res["failure"] = failure; + } + return res; +} + // // server_prompt_cache // diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 312259e7091e..9d64126d0830 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -3,6 +3,7 @@ #include "common.h" #include "llama.h" +#include #include #include #include @@ -29,6 +30,7 @@ enum server_task_type { SERVER_TASK_TYPE_SET_LORA, SERVER_TASK_TYPE_LOAD_LORA, SERVER_TASK_TYPE_UNLOAD_LORA, + SERVER_TASK_TYPE_LEARN, }; // TODO: change this to more generic "response_format" to replace the "format_response_*" in server-common @@ -49,6 +51,27 @@ enum stop_type { STOP_TYPE_LIMIT, }; +// One reply learned on the slot that served its prompt, in this serving context (llama_opt_learn): +// no second context, her prompt never decoded again. Asked by POST /learn, or carried by a +// completion request (task_params::learn), which learns on its own slot before releasing it. +struct learn_request { + int32_t id_slot = -1; + int32_t n_prefix = 0; // the served prompt's tokens, before the reply + std::vector reply; // what she should have said (her reply, or the reviewed correction) + std::vector prompt; // the served prompt, when known: her slot must hold exactly it + int32_t lora_id = -1; // the trainable adapter (the shadow), loaded dormant + float lr = 1e-4f; +}; + +struct learn_outcome { + bool ok = false; + std::string failure; + double loss = 0.0; + int32_t n_labels = 0; + double t_ms = 0.0; + json to_json() const; +}; + struct task_params { bool stream = false; bool include_usage = false; @@ -101,6 +124,10 @@ struct task_params { json format_logit_bias(const std::vector & logit_bias) const; json to_json(bool only_metrics = false) const; + + // learn a reply on this request's slot when it finishes, before the slot is released (no other + // request can take the slot in between): POST /learn's work, atomic with the turn + std::optional learn; }; // struct for tracking the state of a task (e.g., for streaming) @@ -185,6 +212,9 @@ struct server_task { // used by SERVER_TASK_TYPE_UNLOAD_LORA: the adapter id to retire int32_t unload_lora = -1; + // used by SERVER_TASK_TYPE_LEARN + learn_request learn; + server_task() = default; server_task(server_task_type type) : type(type) {} @@ -328,6 +358,8 @@ struct completion_token_output { }; struct server_task_result_cmpl_final : server_task_result { + // the outcome of the request's learn (task_params::learn), run before its slot was released + std::optional learned; std::string content; llama_tokens tokens; @@ -582,6 +614,11 @@ struct server_task_result_lora_change : server_task_result { virtual json to_json() override; }; +struct server_task_result_learn : server_task_result { + learn_outcome outcome; + virtual json to_json() override { return outcome.to_json(); } +}; + struct server_prompt { server_tokens tokens; diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index 8dd7f4d87426..69f47d9b7d48 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -74,7 +74,7 @@ std::string confine_out(const std::string & dir, const std::string & name) { // // Returns false with `why` when the example is malformed (the run is refused); true otherwise, // with `toks`/`loss` filled. -static bool render_example(const json & messages, const std::vector & trained, const json & tools, +static bool train_render(const json & messages, const std::vector & trained, const json & tools, const server_trainer::render_fn & render, const llama_vocab * vocab, const std::string & at, std::vector & toks, std::vector & loss, std::string & why) { auto head = [&](size_t n) { // the first n messages @@ -267,7 +267,7 @@ static example_fit fit_example(const json & ex, const std::string & at, const se } toks.clear(); loss.clear(); - if (!render_example(messages, trained, tools, render, vocab, at, toks, loss, why)) { + if (!train_render(messages, trained, tools, render, vocab, at, toks, loss, why)) { return example_fit::invalid; } full = (int64_t) toks.size(); @@ -284,7 +284,7 @@ static example_fit fit_example(const json & ex, const std::string & at, const se if (dropped) { std::vector t; std::vector l; - if (!render_example(m, tr, tools, render, vocab, at, t, l, why)) { + if (!train_render(m, tr, tools, render, vocab, at, t, l, why)) { return example_fit::invalid; } *min = (int64_t) t.size(); @@ -308,7 +308,7 @@ static example_fit fit_example(const json & ex, const std::string & at, const se return example_fit::skipped; // the head plus her last exchange alone overflow the window } cut = true; - if (!render_example(messages, trained, tools, render, vocab, at, toks, loss, why)) { + if (!train_render(messages, trained, tools, render, vocab, at, toks, loss, why)) { return example_fit::invalid; } } @@ -397,6 +397,32 @@ json server_trainer::measure(const json & body) const { return json::object({{"ok", true}, {"window", window}, {"examples", rows}}); } +// REPACKED WEIGHTS CANNOT BE TRAINED THROUGH (Cormac on the /train plan; card 5b497852): the +// backward reads quantized blocks in the standard layout, and a weight in an extra buffer type +// (CPU_REPACK) is not in it, and its buffer cannot even be read back (get_tensor/cpy_tensor NULL, +// ggml-cpu/repack.cpp). Measured 2026-10-15: on x86 a learn fails ("no backend can run" an +// OUT_PROD on blk.23.ffn_down.weight); on a Metal host at -ngl 0 the scheduler copies the weight to +// Metal and crashes in ggml_backend_tensor_copy. With --no-repack the same learn succeeds. Extra +// buffer types only apply to weights kept on the CPU, so training is allowed when they are +// disabled, or when every layer is offloaded and no tensor override pins a weight elsewhere. +std::string train_repack_refusal(const llama_model * model, const common_params & params) { + const int32_t n_layer = llama_model_n_layer(model); + const int32_t ngl = params.n_gpu_layers; + const bool all_gpu = ngl <= -2 || ngl > n_layer; + // the list carries a {nullptr, nullptr} terminator when parsed from args: count real entries + bool override = false; + for (const auto & o : params.tensor_buft_overrides) { + override = override || o.pattern != nullptr; + } + if (params.no_extra_bufts || (all_gpu && !override)) { + return ""; + } + return "this server may hold base weights in a repacked CPU buffer (n_gpu_layers " + std::to_string(ngl) + + " of " + std::to_string(n_layer) + " layers" + (override ? ", tensor overrides set" : "") + + "), and training's backward needs the standard layout. " + "Serve with every layer offloaded (-ngl all) or with --no-repack."; +} + json server_trainer::start(const json & body_in) { json body = body_in; examples_data prepared; @@ -466,27 +492,8 @@ json server_trainer::start(const json & body_in) { if (confine_out(params_base.train_dir, body.at("out").get()).empty()) { return json::object({{"ok", false}, {"error", "\"out\" must be a bare file name ending in .gguf (no directories, no \"..\"); it is written inside the server's --train-dir"}}); } - // REPACKED WEIGHTS GIVE WRONG GRADIENTS, SILENTLY (Cormac on the /train plan): the backward - // reads quantized blocks in the standard layout, and a weight in an extra buffer type - // (CPU_REPACK) is not in it — no crash, garbage gradients. Extra buffer types only apply to - // weights kept on the CPU, so training is allowed when they are disabled, or when every layer - // is offloaded and no tensor override pins a weight elsewhere; otherwise refuse and say so. - { - const int32_t n_layer = llama_model_n_layer(model); - const int32_t ngl = params_base.n_gpu_layers; - const bool all_gpu = ngl <= -2 || ngl > n_layer; - // the list carries a {nullptr, nullptr} terminator when parsed from args: count real entries - bool override = false; - for (const auto & o : params_base.tensor_buft_overrides) { - override = override || o.pattern != nullptr; - } - if (!params_base.no_extra_bufts && (!all_gpu || override)) { - return json::object({{"ok", false}, {"error", - "refusing to train: this server may hold base weights in a repacked CPU buffer (n_gpu_layers " + - std::to_string(ngl) + " of " + std::to_string(n_layer) + " layers" + (override ? ", tensor overrides set" : "") + - "), and the backward reads the standard layout, so gradients would be silently wrong. " - "Serve with every layer offloaded (-ngl all) or with --no-repack."}}); - } + if (const std::string why = train_repack_refusal(model, params_base); !why.empty()) { + return json::object({{"ok", false}, {"error", "refusing to train: " + why}}); } if (body.contains("parse_special") && !body.at("parse_special").is_boolean()) { return json::object({{"ok", false}, {"error", "\"parse_special\" must be true or false"}}); @@ -907,10 +914,8 @@ void server_trainer::run(json req, examples_data ex) { cparams.n_ctx = window; cparams.n_batch = chunk; cparams.n_ubatch = chunk; - // two sequences: 0 is the walk; 1 is the scratch a recurrent state is snapshotted into - // around each training chunk (llama_context::opt_epoch_iter). Unified, so the attention - // cache stays one stream of the window, not one per sequence. - cparams.n_seq_max = 2; + // one sequence, the walk (it checkpoints a recurrent state through the sequence-state API) + cparams.n_seq_max = 1; cparams.kv_unified = true; // Flash attention ON for the context: the walk's context decodes use it, and the cache then // stores V un-transposed, which a quantized V cache requires. The training graphs build diff --git a/tools/server/server-train.h b/tools/server/server-train.h index 69aff3982ce2..58cff7bd3fb0 100644 --- a/tools/server/server-train.h +++ b/tools/server/server-train.h @@ -61,6 +61,12 @@ struct llama_context; // for every path the server writes or reads there: /train's "out" and /lora-adapters/load. std::string confine_out(const std::string & dir, const std::string & name); +// Why this server must not train on its served weights, or "" when it may: a quantized weight kept +// on the CPU sits in a repacked buffer (CPU_REPACK) whose layout only the CPU kernels read, and the +// backward needs the standard one. One check for every training path (/train, /learn, a learn on a +// completion), so they refuse alike instead of failing differently (continuum card 5b497852). +std::string train_repack_refusal(const llama_model * model, const common_params & params); + class server_trainer { public: // serving: read between training windows; the trainer yields while serving is busy, so a @@ -104,6 +110,7 @@ class server_trainer { }; ~server_trainer(); + // Starts a run on a worker thread; refuses (ok=false) while one is running or on bad input. common_json start(const common_json & body_in); // Measures "examples" against "window" by the same fit rule a run trains by, training diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 7691dff34922..2bad96993a1b 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -219,6 +219,7 @@ int llama_server(common_params & params, int argc, char ** argv) { routes.post_lora_adapters = models_routes->proxy_post; routes.post_lora_adapters_load = models_routes->proxy_post; routes.post_lora_adapters_unload = models_routes->proxy_post; + routes.post_learn = models_routes->proxy_post; routes.get_slots = models_routes->proxy_get; routes.post_slots = models_routes->proxy_post; @@ -273,6 +274,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/lora-adapters", ex_wrapper(routes.post_lora_adapters)); ctx_http.post("/lora-adapters/load", ex_wrapper(routes.post_lora_adapters_load)); ctx_http.post("/lora-adapters/unload", ex_wrapper(routes.post_lora_adapters_unload)); + ctx_http.post("/learn", ex_wrapper(routes.post_learn)); ctx_http.get ("/train", ex_wrapper(routes.get_train)); ctx_http.post("/train", ex_wrapper(routes.post_train)); ctx_http.post("/train/cancel", ex_wrapper(routes.post_train_cancel)); diff --git a/tools/server/tests/learn-every-format.py b/tools/server/tests/learn-every-format.py new file mode 100755 index 000000000000..fc247fd6d4a9 --- /dev/null +++ b/tools/server/tests/learn-every-format.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +# Every final-response format x stream mode carries exactly one learn outcome for a learn-carrying +# request. Run against a server started with --train-dir holding shadow-05.gguf (a fresh LoRA for the +# served model), e.g.: llama-server -m MODEL --port 58199 -np 1 --jinja --train-dir DIR +# what this catches: a serializer, or a request conversion (the Anthropic one dropped "learn"), that +# loses the outcome, so a client learns with nothing on the wire +import json, os, urllib.request +U = os.environ.get("LLAMA_SERVER_URL", "http://localhost:58199") +def raw(path, body): + req = urllib.request.Request(U + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=300) as r: return r.read().decode() +sid = json.loads(raw("/lora-adapters/load", {"name": "shadow-05.gguf"}))["id"] +learn = {"reply": "Topeka is the capital of Kansas.", "lora_id": sid, "lr": 0.001} +msgs = [{"role": "user", "content": "Which city is the capital of Kansas?"}] +cases = [("/completion", {"prompt": "Q: Which city is the capital of Kansas?\nA:"}), + ("/v1/completions", {"prompt": "Q: Which city is the capital of Kansas?\nA:"}), + ("/v1/chat/completions", {"messages": msgs, "chat_template_kwargs": {"enable_thinking": False}}), + ("/v1/responses", {"input": msgs, "chat_template_kwargs": {"enable_thinking": False}}), + ("/v1/messages", {"messages": msgs, "max_tokens": 2, "chat_template_kwargs": {"enable_thinking": False}})] +fails = 0 +for path, body in cases: + for stream in (False, True): + b = dict(body, stream=stream, learn=learn, temperature=0) + b.setdefault("max_tokens", 2); b["n_predict"] = 2 + text = raw(path, b) + found = [] + for line in (text.splitlines() if stream else [text]): + line = line[5:].strip() if line.startswith("data:") else line.strip() + if not line or line == "[DONE]" or not line.startswith("{"): continue + def walk(o): + if isinstance(o, dict): + if "learn" in o and isinstance(o["learn"], dict): found.append(o["learn"]) + for v in o.values(): walk(v) + elif isinstance(o, list): + for v in o: walk(v) + walk(json.loads(line)) + ok = len(found) == 1 and found[0].get("success") is True + fails += not ok + print(f"{'OK ' if ok else 'BAD'} {path:22s} stream={stream!s:5s} learn outcomes={len(found)} {found[0] if found else ''}") +print("ALL FORMATS CARRY THE LEARN OUTCOME" if fails == 0 else f"{fails} FORMAT(S) WITHOUT IT") +raise SystemExit(1 if fails else 0) diff --git a/tools/server/tests/learn-refuses-repacked.py b/tools/server/tests/learn-refuses-repacked.py new file mode 100644 index 000000000000..6469c98f9bb1 --- /dev/null +++ b/tools/server/tests/learn-refuses-repacked.py @@ -0,0 +1,24 @@ +#!/usr/bin/env python3 +# A learn-carrying completion on a server that may hold repacked CPU weights still answers its turn, +# and its learn outcome says why it did not learn. Run against a CPU server WITHOUT --no-repack and a +# --train-dir holding shadow-05.gguf, e.g.: +# llama-server -m MODEL --port 58199 -np 2 -c 8192 -ngl 0 --jinja --train-dir DIR +# what this catches (continuum card 5b497852): a learn on repacked weights failed in the backward on +# x86 and crashed the engine in ggml_backend_tensor_copy on a Metal host at -ngl 0; it must refuse +# the way /train does, before any training graph is built, and leave the turn served. +import json, os, urllib.request +U = os.environ.get("LLAMA_SERVER_URL", "http://localhost:58199") +def post(path, body): + req = urllib.request.Request(U + path, data=json.dumps(body).encode(), headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=300) as r: return json.loads(r.read().decode()) +sid = post("/lora-adapters/load", {"name": "shadow-05.gguf"})["id"] +res = post("/v1/chat/completions", { + "messages": [{"role": "user", "content": "Which city is the capital of Kansas?"}], + "max_tokens": 8, "temperature": 0, "chat_template_kwargs": {"enable_thinking": False}, + "learn": {"reply": "Topeka is the capital of Kansas.", "lora_id": sid, "lr": 0.001}}) +answer = res["choices"][0]["message"]["content"] +learn = res.get("learn") or {} +ok = bool(answer) and learn.get("success") is False and "--no-repack" in learn.get("failure", "") +print(f"answer={answer!r} learn={learn}") +print("REFUSED WITH THE REASON, TURN SERVED" if ok else "WRONG") +raise SystemExit(0 if ok else 1)