llama: connect labelled decode to shared physical cache rows

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

ggml/include/ggml-opt.h

@@ -208,6 +208,32 @@ extern "C" {
         struct ggml_tensor  * outputs,
         int64_t               rows);
 
+    // Project the existing row-wise head before consumers such as sampling are built.
+    // Row indices address hidden's token rows. A null public_rows requests no public logits.
+    // Returns zero without changing gf when the head cannot be replayed.
+    GGML_API int ggml_opt_project_head(
+        struct ggml_context * ctx_compute,
+        struct ggml_cgraph  * gf,
+        struct ggml_tensor  * hidden,
+        struct ggml_tensor  * outputs,
+        struct ggml_tensor  * shadow_rows,
+        struct ggml_tensor  * public_rows,
+        struct ggml_tensor  * label_rows,
+        struct ggml_tensor  * public_shadow_rows,
+        struct ggml_tensor ** slice_shadow_rows,
+        int64_t               rows,
+        struct ggml_tensor ** public_outputs,
+        struct ggml_tensor ** slices,
+        size_t                slices_capacity);
+
+    GGML_API void ggml_opt_prepare_alloc_slices(
+        ggml_opt_context_t    opt_ctx,
+        struct ggml_context * ctx_compute,
+        struct ggml_cgraph  * gf,
+        struct ggml_tensor  * inputs,
+        struct ggml_tensor ** slices,
+        size_t                n_slices);
+
     // allocate the next graph for evaluation, either forward or forward + backward
     // must be called exactly once prior to calling ggml_opt_eval
     // returns false when the graph could not be allocated (device memory): the graph is then NOT

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

ggml/src/ggml-cpu/ops.cpp

@@ -4371,8 +4371,8 @@ static void ggml_compute_forward_out_prod_q_f32(
     const ggml_type type = src0->type;
     ggml_to_float_t const dequantize_row_q = ggml_get_type_traits(type)->to_float;
 
-    GGML_ASSERT(ne02 == ne12);
-    GGML_ASSERT(ne03 == ne13);
+    GGML_ASSERT(ne12 % ne02 == 0);
+    GGML_ASSERT(ne13 % ne03 == 0);
     GGML_ASSERT(ne2  == ne12);
     GGML_ASSERT(ne3  == ne13);
 
@@ -4387,8 +4387,8 @@ static void ggml_compute_forward_out_prod_q_f32(
 
     GGML_ASSERT(ne0 == ne00);
     GGML_ASSERT(ne1 == ne10);
-    GGML_ASSERT(ne2 == ne02);
-    GGML_ASSERT(ne3 == ne03);
+
+
 
     // nb01 >= nb00 - src0 is not transposed
     //   compute by src0 rows
@@ -4425,8 +4425,8 @@ static void ggml_compute_forward_out_prod_q_f32(
         const int64_t i2 = (ir - i3*ne2*ne1)/ne1;
         const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1);
 
-        const int64_t i02 = i2;
-        const int64_t i03 = i3;
+        const int64_t i02 = i2 / (ne12 / ne02);
+        const int64_t i03 = i3 / (ne13 / ne03);
 
         //const int64_t i10 = i1;
         const int64_t i12 = i2;
@@ -4461,8 +4461,8 @@ static void ggml_compute_forward_out_prod_f16_f32(
     GGML_ASSERT(src1->type == GGML_TYPE_F32);
     GGML_ASSERT(dst->type == GGML_TYPE_F32);
 
-    GGML_ASSERT(ne02 == ne12);
-    GGML_ASSERT(ne03 == ne13);
+    GGML_ASSERT(ne12 % ne02 == 0);
+    GGML_ASSERT(ne13 % ne03 == 0);
     GGML_ASSERT(ne2  == ne12);
     GGML_ASSERT(ne3  == ne13);
 
@@ -4471,8 +4471,8 @@ static void ggml_compute_forward_out_prod_f16_f32(
 
     GGML_ASSERT(ne0 == ne00);
     GGML_ASSERT(ne1 == ne10);
-    GGML_ASSERT(ne2 == ne02);
-    GGML_ASSERT(ne3 == ne03);
+
+
 
     if (ith == 0) {
         ggml_vec_set_f32(ne0*ne1*ne2*ne3, (float *)dst->data, 0);
@@ -4491,8 +4491,8 @@ static void ggml_compute_forward_out_prod_f16_f32(
         const int64_t i2 = (ir - i3*ne2*ne1)/ne1;
         const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1);
 
-        const int64_t i02 = i2;
-        const int64_t i03 = i3;
+        const int64_t i02 = i2 / (ne12 / ne02);
+        const int64_t i03 = i3 / (ne13 / ne03);
 
         const int64_t i12 = i2;
         const int64_t i13 = i3;

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

ggml/src/ggml-opt.cpp

@@ -445,6 +445,7 @@ static void ggml_opt_checkpoint(ggml_context * ctx, ggml_cgraph * g, int n_fwd,
     GGML_ASSERT((int) order.size() <= g->size && "recompute clones overflow the graph: raise the training node budget");
     std::copy(order.begin(), order.end(), g->nodes);
     g->n_nodes = (int) order.size();
+    ggml_graph_update_use_counts(g);
 }
 
 static ggml_tensor * map_tensor(std::map<ggml_tensor *, ggml_tensor *> & tensor_map, ggml_context * ctx, ggml_tensor * tensor) {
@@ -1041,19 +1042,24 @@ void ggml_opt_prepare_alloc(
     opt_ctx->output_slices.clear();
 }
 
-int ggml_opt_prepare_alloc_head(
-        ggml_opt_context_t    opt_ctx,
+int ggml_opt_project_head(
         struct ggml_context * ctx_compute,
         struct ggml_cgraph  * gf,
-        struct ggml_tensor  * inputs,
         struct ggml_tensor  * hidden,
         struct ggml_tensor  * outputs,
-        int64_t               rows) {
-    ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, outputs);
-    const int64_t n = hidden->ne[1];
-    if (opt_ctx->loss_type != GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS || rows <= 0 || rows >= n ||
-            !ggml_is_contiguous(hidden) || ggml_n_dims(hidden) > 2 || outputs->ne[1] != n) {
-        return 1;
+        struct ggml_tensor * shadow_rows,
+        struct ggml_tensor * public_rows,
+        struct ggml_tensor * label_rows,
+        struct ggml_tensor * public_shadow_rows,
+        struct ggml_tensor ** slice_shadow_rows,
+        int64_t               rows,
+        struct ggml_tensor ** public_outputs,
+        struct ggml_tensor ** slices,
+        size_t                slices_capacity) {
+    const int64_t n = label_rows ? label_rows->ne[0] : hidden->ne[1];
+    if (rows <= 0 || n <= 0 || !ggml_is_contiguous(hidden) || ggml_n_dims(hidden) > 2 ||
+            outputs->ne[1] != hidden->ne[1] || slices_capacity < (size_t) ((n + rows - 1) / rows)) {
+        return 0;
     }
     // the head: the nodes between hidden and outputs, in graph order
     std::unordered_set<const ggml_tensor *> on_path = { hidden };
@@ -1068,15 +1074,27 @@ int ggml_opt_prepare_alloc_head(
             continue;
         }
         const bool replayable = node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL ||
-                                node->op == GGML_OP_ADD || node->op == GGML_OP_SCALE;
-        if (!replayable || node->view_src != nullptr) {
-            return 1; // a head this does not know how to replay per slice: the whole outputs
+                                node->op == GGML_OP_ADD || node->op == GGML_OP_SCALE ||
+                                (shadow_rows && (node->op == GGML_OP_DUP || node->op == GGML_OP_GET_ROWS || node->op == GGML_OP_SET_ROWS));
+        if (!replayable || (node->view_src != nullptr && node->op != GGML_OP_SET_ROWS)) {
+            return 0; // unsupported head
+        }
+        if ((node->op == GGML_OP_GET_ROWS && node->src[1] != shadow_rows) ||
+                (node->op == GGML_OP_SET_ROWS && node->src[1] != shadow_rows)) {
+            return 0;
+        }
+        if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) {
+            for (int k = 0; k < 2; ++k) {
+                if (!on_path.count(node->src[k]) && node->src[k]->ne[1] != 1) {
+                    return 0;
+                }
+            }
         }
         on_path.insert(node);
         head.push_back(node);
     }
     if (head.empty() || head.back() != outputs) {
-        return 1;
+        return 0;
     }
     // nothing else in the graph may read the head (it is about to leave the graph)
     for (int i = 0; i < gf->n_nodes; ++i) {
@@ -1086,7 +1104,7 @@ int ggml_opt_prepare_alloc_head(
         }
         for (int k = 0; k < GGML_MAX_SRC; ++k) {
             if (node->src[k] != nullptr && node->src[k] != hidden && on_path.count(node->src[k]) > 0) {
-                return 1;
+                return 0;
             }
         }
     }
@@ -1097,15 +1115,31 @@ int ggml_opt_prepare_alloc_head(
         }
     }
     gf->n_nodes = j;
-    // the head again, one slice of rows at a time
-    for (int64_t r0 = 0; r0 < n; r0 += rows) {
-        const int64_t nr = std::min(rows, n - r0);
+    auto replay = [&](ggml_tensor * selected, ggml_tensor * local_shadow_rows, int64_t row) {
         std::unordered_map<const ggml_tensor *, ggml_tensor *> as_slice;
-        as_slice[hidden] = ggml_view_2d(ctx_compute, hidden, hidden->ne[0], nr, hidden->nb[1], r0*hidden->nb[1]);
+        as_slice[hidden] = selected;
+        if (shadow_rows) {
+            as_slice[shadow_rows] = local_shadow_rows;
+        }
         for (ggml_tensor * node : head) {
             auto src = [&](int k) { const auto it = as_slice.find(node->src[k]); return it != as_slice.end() ? it->second : node->src[k]; };
             ggml_tensor * t = nullptr;
+            if (node->op == GGML_OP_SET_ROWS && shadow_rows && !local_shadow_rows) {
+                as_slice[node] = src(2);
+                continue;
+            }
+            bool missing = false;
+            for (int k = 0; k < GGML_MAX_SRC; ++k) {
+                missing |= node->src[k] != nullptr && src(k) == nullptr;
+            }
+            if (missing) {
+                as_slice[node] = nullptr;
+                continue;
+            }
             switch (node->op) {
+                case GGML_OP_DUP:      t = ggml_dup(ctx_compute, src(0)); break;
+                case GGML_OP_GET_ROWS: t = ggml_get_rows(ctx_compute, src(0), src(1)); break;
+                case GGML_OP_SET_ROWS: t = ggml_set_rows(ctx_compute, src(2), src(0), src(1)); break;
                 case GGML_OP_MUL_MAT: t = ggml_mul_mat(ctx_compute, src(0), src(1)); break;
                 case GGML_OP_MUL:     t = ggml_mul    (ctx_compute, src(0), src(1)); break;
                 case GGML_OP_ADD:     t = ggml_add    (ctx_compute, src(0), src(1)); break;
@@ -1117,14 +1151,57 @@ int ggml_opt_prepare_alloc_head(
                 } break;
                 default: GGML_ABORT("not a head op");
             }
-            ggml_format_name(t, "%s-rows%lld", node->name, (long long) r0);
+            ggml_set_op_params(t, node->op_params, sizeof(node->op_params));
+            ggml_format_name(t, "%s-rows%lld", node->name, (long long) row);
             as_slice[node] = t;
         }
-        opt_ctx->output_slices.push_back(as_slice.at(outputs));
+        return as_slice.at(outputs);
+    };
+    *public_outputs = nullptr;
+    if (public_rows) {
+        *public_outputs = replay(ggml_get_rows(ctx_compute, hidden, public_rows),
+                public_shadow_rows, -1);
+        ggml_build_forward_expand(gf, *public_outputs);
+    }
+    ggml_tensor * selected = label_rows ? ggml_get_rows(ctx_compute, hidden, label_rows) : hidden;
+    size_t count = 0;
+    for (int64_t r0 = 0; r0 < n; r0 += rows) {
+        const int64_t nr = std::min(rows, n - r0);
+        slices[count] = replay(
+                ggml_view_2d(ctx_compute, selected, selected->ne[0], nr, selected->nb[1], r0 * selected->nb[1]),
+                slice_shadow_rows ? slice_shadow_rows[count] : nullptr, r0);
+        ++count;
     }
-    return (int) opt_ctx->output_slices.size();
+    ggml_graph_update_use_counts(gf);
+    return (int) count;
+}
+
+void ggml_opt_prepare_alloc_slices(
+        ggml_opt_context_t opt_ctx, ggml_context * ctx_compute, ggml_cgraph * gf,
+        ggml_tensor * inputs, ggml_tensor ** slices, size_t n_slices) {
+    GGML_ASSERT(n_slices > 0);
+    ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, slices[0]);
+    opt_ctx->output_slices.assign(slices, slices + n_slices);
 }
 
+int ggml_opt_prepare_alloc_head(
+        ggml_opt_context_t opt_ctx, ggml_context * ctx_compute, ggml_cgraph * gf,
+        ggml_tensor * inputs, ggml_tensor * hidden, ggml_tensor * outputs, int64_t rows) {
+    ggml_opt_prepare_alloc(opt_ctx, ctx_compute, gf, inputs, outputs);
+    const int64_t n = hidden->ne[1];
+    if (opt_ctx->loss_type != GGML_OPT_LOSS_TYPE_CROSS_ENTROPY_TARGETS || rows <= 0 || rows >= n) {
+        return 1;
+    }
+    std::vector<ggml_tensor *> 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;
 }
@@ -1208,8 +1285,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);

ggml/src/ggml.c

@@ -6667,7 +6667,7 @@ static void ggml_acc_or_set(
         }
         cgraph->grads[isrc] = ggml_acc_impl(ctx, acc_into, tensor, nb1, nb2, nb3, offset, cgraph->grad_accs[isrc]);
     } else {
-        struct ggml_tensor * a_zero = ggml_scale(ctx, src, 0.0f); // FIXME this is going to produce NaN if a contains inf/NaN
+        struct ggml_tensor * a_zero = ggml_fill(ctx, ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, src->ne), 0.0f);
         cgraph->grads[isrc] = ggml_acc_impl(ctx, a_zero, tensor, nb1, nb2, nb3, offset, false);
     }
     ggml_format_name(cgraph->grads[isrc], "grad for %s", cgraph->visited_hash_set.keys[isrc]->name);
@@ -7181,10 +7181,10 @@ static void ggml_compute_backward(
                 size_t nb2 = tensor->nb[2];
                 size_t nb3 = tensor->nb[3];
 
-                if (cgraph->grads[isrc0] && src0->type != cgraph->grads[isrc0]->type) {
-                    // gradient is typically F32, but src0 could be other type
-                    size_t ng = ggml_element_size(cgraph->grads[isrc0]);
-                    size_t n0 = ggml_element_size(src0);
+                if (src0->type != GGML_TYPE_F32) {
+                    // The gradient is F32 even when the viewed source stores quantized blocks.
+                    size_t ng = sizeof(float) * ggml_blck_size(src0->type);
+                    size_t n0 = ggml_type_size(src0->type);
                     GGML_ASSERT(offset % n0 == 0);
                     GGML_ASSERT(nb1 % n0 == 0);
                     GGML_ASSERT(nb2 % n0 == 0);
@@ -7701,7 +7701,7 @@ void ggml_build_backward_expand(
             // A gradient is F32 whatever its tensor's type. A quantized source is admitted where
             // the backward only re-shapes the gradient or casts its tile to F32 (a training graph
             // that attends to a q8_0 cache in place, as serving does); every other op needs float.
-            const bool shapes_only = node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
+            const bool shapes_only = node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE ||
                 node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_FLASH_ATTN_EXT;
             GGML_ASSERT(node->src[j]->type == GGML_TYPE_F32 || node->src[j]->type == GGML_TYPE_F16 ||
                 (shapes_only && ggml_is_quantized(node->src[j]->type)));
@@ -7942,6 +7942,19 @@ void ggml_graph_reset(struct ggml_cgraph * cgraph) {
     }
 }
 
+void ggml_graph_update_use_counts(struct ggml_cgraph * cgraph) {
+    memset(cgraph->use_counts, 0, cgraph->visited_hash_set.size * sizeof(cgraph->use_counts[0]));
+    for (int i = 0; i < cgraph->n_nodes; ++i) {
+        for (int j = 0; j < GGML_MAX_SRC; ++j) {
+            struct ggml_tensor * src = cgraph->nodes[i]->src[j];
+            if (src) {
+                const size_t pos = ggml_hash_find(&cgraph->visited_hash_set, src);
+                GGML_ASSERT(pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, pos));
+                cgraph->use_counts[pos]++;
+            }
+        }
+    }
+}
 void ggml_graph_clear(struct ggml_cgraph * cgraph) {
     cgraph->n_leafs = 0;
     cgraph->n_nodes = 0;