train in place: labels ride the batch, the reply learns on the slot that served it (POST /learn)

common/common.cpp

@@ -2401,8 +2401,21 @@ bool common_lora_write_fresh(const llama_model * model, const std::string & path
     std::vector<shape> 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) {

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);

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);

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

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,

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:

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]);