server: checkpoint at a client-named user message (checkpoint_user_message)

common/chat.h

@@ -179,6 +179,19 @@ struct common_chat_msg_spans {
         return false;
     }
 
+    // token position of the k-th user message (0-based), or -1 when there is none
+    int32_t user_message_pos(int32_t k) const {
+        if (k < 0) {
+            return -1;
+        }
+        for (const auto & span : spans) {
+            if (span.role == COMMON_CHAT_ROLE_USER && k-- == 0) {
+                return (int32_t) span.pos;
+            }
+        }
+        return -1;
+    }
+
     int32_t last_user_message_pos() const {
         for (auto it = spans.rbegin(); it != spans.rend(); ++it) {
             if (it->role == COMMON_CHAT_ROLE_USER) {

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/json-schema-to-grammar.cpp

@@ -964,6 +964,27 @@ class common_schema_converter {
                     properties, required, name,
                     schema.contains("additionalProperties") ? schema["additionalProperties"] : json()));
         }
+        // A singleton union wrapper adds no constraint unless a sibling does.
+        if (schema.contains("allOf") && schema["allOf"].is_array() && schema["allOf"].size() == 1) {
+            bool transparent = true;
+            for (const auto & item : schema.items()) {
+                const auto & key = item.key();
+                if (key != "allOf" && key != "description" && key != "title" && key != "$comment" &&
+                    key != "default" && key != "examples" && key != "definitions" && key != "$defs") {
+                    transparent = false;
+                }
+            }
+            const json * component = &schema["allOf"][0];
+            if (component->is_object() && component->contains("$ref")) {
+                auto ref = _refs.find((*component)["$ref"]);
+                if (ref != _refs.end() && component->size() == 1) {
+                    component = &ref->second;
+                }
+            }
+            if (transparent && (component->contains("oneOf") || component->contains("anyOf"))) {
+                return visit(*component, name);
+            }
+        }
         if ((schema_type.is_null() || schema_type == "object" || schema_type == "string") && schema.contains("allOf")) {
             std::unordered_set<std::string> required;
             std::vector<std::pair<std::string, json>> properties;

common/speculative.cpp

@@ -633,6 +633,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl {
                 /*.n_seq_id =*/ nullptr,
                 /*.seq_id   =*/ nullptr,
                 /*.logits   =*/ nullptr,
+                /*.label    =*/ nullptr,
+                /*.label_weight =*/ nullptr,
             };
             const int32_t rc = llama_encode(ctx_dft, enc_batch);
             if (rc != 0) {
@@ -1126,6 +1128,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
                     /*.n_seq_id =*/ nullptr,
                     /*.seq_id   =*/ nullptr,
                     /*.logits   =*/ nullptr,
+                    /*.label    =*/ nullptr,
+                    /*.label_weight =*/ 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

@@ -10,6 +10,7 @@
 #include "ggml-backend.h"
 
 #include <stdint.h>
+#include <stdio.h>
 
 #ifdef  __cplusplus
 extern "C" {
@@ -32,6 +33,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 +159,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
@@ -166,6 +177,37 @@ extern "C" {
 
     GGML_API const char * ggml_opt_optimizer_name(enum ggml_opt_optimizer_type);
 
+    // Checkpoint optimizer state at a successful update boundary. Parameters are saved separately.
+    // binding identifies the caller's model, adapter topology and optimizer policy; it must match exactly.
+    // Mutable callback/schedule state must be saved by the caller in the same generation, not inferred from a callback address.
+    // Import requires allocated optimizer tensors and no prepared evaluation; it validates all input before changing state.
+    // The caller owns file integrity, atomic publication and serialization with compute.
+    // Explicit generation zero: only a fresh, allocated, import-pending optimizer is eligible.
+    // Initializes persistent optimizer state to zero, without a weight/counter/RNG update. Writes
+    // a distinct initial-boundary artifact; never enables evaluation or claims sample consumption.
+    // Caller must publish weights/policy/integrity atomically and restore before first evaluation.
+    GGML_API bool ggml_opt_state_write_initial(ggml_opt_context_t opt_ctx, FILE * file, const char * binding);
+    GGML_API bool ggml_opt_state_write(ggml_opt_context_t opt_ctx, FILE * file, const char * binding);
+    GGML_API bool ggml_opt_state_read (ggml_opt_context_t opt_ctx, FILE * file, const char * binding);
+    // Prepare owns validated host payload without changing the destination. The caller must keep
+    // the context alive and serialize prepare/apply with all evaluation and allocation operations.
+    struct ggml_opt_state;
+    // Allocate a genuine supplied graph only for importing a fresh optimizer generation. Executes
+    // no graph and consumes no sample; refuses used, pending, partial or failed contexts. The
+    // caller must restore before evaluation. Existing lifetimes are unchanged: retain parameter
+    // and static graph storage for the optimizer lifetime; dynamic graph storage follows the
+    // existing prepare/allocator custody contract. This operation does not release that storage.
+    // A failed allocation attempt is teardown-only: evaluation, retry and import are refused.
+    // An eligibility refusal leaves the existing optimizer untouched.
+    GGML_API bool ggml_opt_alloc_for_restore(ggml_opt_context_t opt_ctx);
+    // Side-effect-free eligibility, owned by the same allocator. Check before replacing caller
+    // graph/scheduler state; allocation checks again under the caller's serialized custody.
+    GGML_API bool ggml_opt_can_prepare_restore(ggml_opt_context_t opt_ctx);
+    GGML_API struct ggml_opt_state * ggml_opt_state_prepare(ggml_opt_context_t opt_ctx, FILE * file, const char * binding);
+    GGML_API bool ggml_opt_state_apply(ggml_opt_context_t opt_ctx, struct ggml_opt_state * state);
+    GGML_API void ggml_opt_state_free(struct ggml_opt_state * state);
+
+
     // ====== Optimization Result ======
 
     GGML_API ggml_opt_result_t ggml_opt_result_init(void);
@@ -188,6 +230,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-backend.cpp

@@ -1610,16 +1610,19 @@ static bool ggml_backend_sched_alloc_splits(ggml_backend_sched_t sched) {
 static ggml_moe::DirectReadFetcher moe_direct_fetcher;
 #endif
 static ggml_moe::MmapFaultFetcher  moe_mmap_fetcher;
+static ggml_moe::DirContainerFetcher & moe_container_fetcher() {
+    static ggml_moe::DirContainerFetcher fetcher(ggml_moe::moe_config().container_dir);
+    return fetcher;
+}
 static ggml_moe::ExpertFetcher &   moe_pick_base_fetcher() {
     // Container-serve (#268): when a packed per-layer expert container is configured, read from it.
     // Its records are contiguous + 16KiB-aligned, so this is the honest ~GB/s sequential path — vs
     // DirectRead scraping scattered offsets out of the raw GGUF's mmap. The serving caller packs
     // (layer, byte_offset) into `src` via DirContainerFetcher::pack_src; the cache passes it opaquely.
     // Constructed once with the configured dir (banks open lazily); a bad dir opens no banks and every
     // fetch fails LOUD (zeroed record -> downstream validation), never a silent wrong-source serve.
-    if (const char * cdir = ggml_moe::moe_config().container_dir) {
-        static ggml_moe::DirContainerFetcher container_fetcher(cdir);
-        return container_fetcher;
+    if (ggml_moe::moe_config().container_dir) {
+        return moe_container_fetcher();
     }
 #ifdef _WIN32
     if (ggml_moe::moe_config().direct_read) { return moe_direct_fetcher; }
@@ -1757,6 +1760,10 @@ static ggml_backend_moe_gather_entry_t moe_gather_entry_fn(ggml_backend_t backen
 
 static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t sched) {
     GGML_ASSERT(sched);
+    if (ggml_moe::moe_config().container_dir && !moe_container_fetcher().ok()) {
+        GGML_LOG_ERROR("%s: expert container custody unavailable; refusing graph execution\n", __func__);
+        return GGML_STATUS_FAILED;
+    }
     struct ggml_backend_sched_split * splits = sched->splits;
 
     ggml_tensor * prev_ids_tensor = nullptr;
@@ -2114,7 +2121,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s
                         // parses from the tensor name's `blk.N.` (names may be backend-prefixed, so
                         // anchor on the substring, not the string start).
                         static const uint64_t moec_record_bytes =
-                            ggml_moe::dir_container_record_bytes(ggml_moe::moe_config().container_dir);
+                            ggml_moe::moe_config().container_dir ? moe_container_fetcher().record_bytes() : 0;
                         uint32_t moec_layer = 0;
                         const char * moec_blk = std::strstr(tname, "blk.");
                         const bool moec_active = moec_record_bytes > 0 && moec_blk != nullptr &&