train: the training context's K/V type is a request option (context_cache_type)

ggml/include/ggml-opt.h

@@ -208,6 +208,10 @@ extern "C" {
     // run of this shape needs, measured by the run itself.
     GGML_API size_t ggml_opt_peak_graph_bytes(ggml_opt_context_t opt_ctx);
 
+    // the last memory refusal: the bytes the refused device's graph needed, and the bytes it was
+    // allowed to add at that moment (0, 0 when no graph was refused for memory)
+    GGML_API void ggml_opt_last_refusal_bytes(ggml_opt_context_t opt_ctx, size_t * need, size_t * allowance);
+
     // do forward pass, increment result if not NULL, do backward pass if allocated
     GGML_API void ggml_opt_eval(ggml_opt_context_t opt_ctx, ggml_opt_result_t result);
 

ggml/src/ggml-opt.cpp

@@ -66,6 +66,11 @@ struct ggml_opt_context {
     bool static_graphs           = false;
     size_t alloc_budget          = 0; // bytes a graph may add on a non-CPU device; 0 = the device's own free figure
     size_t peak_graph_bytes      = 0; // largest graph measured by the preflight on a non-CPU device
+    // the last memory refusal: what the refused device's graph needed, and what it was allowed to
+    // add at that moment (free memory, capped by the budget, minus slack). Kept apart from
+    // `refusal`, which stays empty for a memory refusal (callers retry on that).
+    size_t refused_need_bytes    = 0;
+    size_t refused_allow_bytes   = 0;
     bool eval_ready              = false;
     // Accumulators and momenta belong to a PARAMETER (or to the loss), never to a node index: a
     // graph built per step may differ in topology from the first one (the training walk's chunks
@@ -982,6 +987,11 @@ void ggml_opt_set_alloc_budget(ggml_opt_context_t opt_ctx, size_t bytes) {
     opt_ctx->alloc_budget = bytes;
 }
 
+void ggml_opt_last_refusal_bytes(ggml_opt_context_t opt_ctx, size_t * need, size_t * allowance) {
+    *need      = opt_ctx->refused_need_bytes;
+    *allowance = opt_ctx->refused_allow_bytes;
+}
+
 size_t ggml_opt_peak_graph_bytes(ggml_opt_context_t opt_ctx) {
     return opt_ctx->peak_graph_bytes;
 }
@@ -1142,6 +1152,13 @@ bool ggml_opt_alloc(ggml_opt_context_t opt_ctx, bool backward) {
             }
         }
         if (!ok) {
+            for (int i = 0; i < n; ++i) {
+                if (sizes[i] > max_new[i]) {
+                    opt_ctx->refused_need_bytes  = sizes[i];
+                    opt_ctx->refused_allow_bytes = max_new[i];
+                    break;
+                }
+            }
             ggml_backend_sched_reset(sched);
             opt_ctx->allocated_graph = nullptr;
             opt_ctx->eval_ready      = false;

include/llama.h

@@ -1755,6 +1755,12 @@ extern "C" {
     // batch): the footprint a run of this shape needs, from the run's own allocation preflight.
     LLAMA_API size_t llama_opt_graph_bytes(struct llama_context * lctx);
 
+    // The last training chunk refused for memory, if any: its position in the window and length,
+    // the bytes its graph needed on the refusing device, and what that device allowed it to add
+    // at that moment. Returns false when no chunk was refused for memory.
+    LLAMA_API bool llama_opt_last_refusal(struct llama_context * lctx, uint32_t * pos, uint32_t * n_tokens,
+                                          size_t * need, size_t * allowance);
+
     // The bytes this context's memory holds on non-host (GPU) buffers: the K/V cache for the whole
     // window, and any recurrent state; not the model's weights, which it shares, nor its compute
     // buffers, which a training graph is allocated into (llama_opt_graph_bytes covers those).

src/llama-context-opt.cpp

@@ -323,6 +323,9 @@ bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) {
             }
         }
         if (!ggml_opt_alloc(opt_ctx, train)) {
+            lctx.opt_refused     = true;
+            lctx.opt_refused_pos = pos_ctx;
+            lctx.opt_refused_n   = n_chunk;
             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) + ")");
@@ -706,6 +709,9 @@ bool llama_opt_param_filter_all(const struct ggml_tensor * tensor, void * userda
 void llama_opt_init(struct llama_context * ctx, struct llama_model * model, struct llama_opt_params lopt_params) {
     ctx->opt_alloc_failed.store(false);
     ctx->opt_failure.clear();
+    ctx->opt_refused     = false;
+    ctx->opt_refused_pos = 0;
+    ctx->opt_refused_n   = 0;
     ctx->opt_init(model, lopt_params);
 }
 
@@ -738,6 +744,17 @@ size_t llama_context::opt_graph_bytes() const {
     return opt_ctx ? ggml_opt_peak_graph_bytes(opt_ctx) : 0;
 }
 
+bool llama_opt_last_refusal(struct llama_context * ctx, uint32_t * pos, uint32_t * n_tokens, size_t * need, size_t * allowance) {
+    *pos = ctx->opt_refused_pos;
+    *n_tokens = ctx->opt_refused_n;
+    *need = 0;
+    *allowance = 0;
+    if (ctx->opt_ctx) {
+        ggml_opt_last_refusal_bytes(ctx->opt_ctx, need, allowance);
+    }
+    return ctx->opt_refused;
+}
+
 size_t llama_opt_graph_bytes(struct llama_context * ctx) {
     return ctx->opt_graph_bytes();
 }

src/llama-context.h

@@ -206,6 +206,11 @@ struct llama_context {
     // why it was refused, copied from ggml_opt_refusal when the epoch stops (llama_opt_failure);
     // empty for the memory gate
     std::string opt_failure;
+    // the last chunk refused for memory: where in the window it sat and how long it was
+    // (llama_opt_last_refusal), so a late failure names the chunk that did not fit
+    bool     opt_refused     = false;
+    uint32_t opt_refused_pos = 0;
+    uint32_t opt_refused_n   = 0;
     // what training may add per GPU device (llama_opt_set_memory_budget); 0 = no cap
     size_t opt_memory_budget = 0;
     // the exact walk's host memory per window (llama_opt_set_walk_host_budget); 0 = no cap

tests/test-chat.cpp

@@ -7277,6 +7277,26 @@ static int test_train_control(const char * model_path, const char * output_dir)
     require(trainer.start(req).at("ok").get<bool>(), "baseline refused");
     const auto baseline = wait_for(done);
 
+    // what this catches (continuum 34d574fc): the training context's K/V type is a request option,
+    // so a dream step can be shaped to fit BESIDE serving. The default stays Q8_0 (the baseline
+    // above names it), an unknown type is refused before any claim, and a Q4_0 context holds
+    // less than the Q8_0 one for the same window: the number a caller sizes a fit by.
+    require(baseline.value("context_cache_type", "") == "q8_0", "the default training context is Q8_0");
+    {
+        json bad = req;
+        bad["out"] = "bad-cache.gguf";
+        bad["context_cache_type"] = "q5_9";
+        require(!trainer.start(bad).at("ok").get<bool>(), "an unknown context_cache_type is refused");
+        json small = req;
+        small["out"] = "q4-context.gguf";
+        small["context_cache_type"] = "q4_0";
+        require(trainer.start(small).at("ok").get<bool>(), "a q4_0 training context starts");
+        const auto q4 = wait_for(done);
+        require(q4.value("context_cache_type", "") == "q4_0", "the run reports the type it was built with");
+        require(q4.value("context_mib", 0.0) < baseline.value("context_mib", 0.0),
+                "a Q4_0 training context holds less than the Q8_0 one at the same window");
+    }
+
     // Block after two updates, then exercise manual pause separately from serving admission.
     calls.store(0);
     hold_at.store(3);

tools/server/server-train.cpp

@@ -18,6 +18,17 @@ using json = common_json;
 static std::atomic<int64_t>       g_train_batch{0};
 static std::atomic<int64_t>       g_train_batch_max{0};
 
+// The training context's K/V type, by request (continuum 34d574fc: a dream step must fit BESIDE
+// serving, never by borrowing from it). Its cache holds only constants, the context before a
+// chunk; the chunk's own K/V, the ones with a gradient, are F32 in the graph and never cached
+// (see the cparams below). "q8_0" is the default and the only type before this option existed.
+static bool context_cache_type_of(const std::string & name, ggml_type & out) {
+    if (name == "q8_0") { out = GGML_TYPE_Q8_0; return true; }
+    if (name == "q4_0") { out = GGML_TYPE_Q4_0; return true; }
+    if (name == "f16")  { out = GGML_TYPE_F16;  return true; }
+    return false;
+}
+
 static std::vector<std::string> split_targets(const std::string & list) {
     std::vector<std::string> out;
     std::stringstream ss(list);
@@ -444,6 +455,9 @@ json server_trainer::start(const json & body_in) {
                                                                      why = "walk_horizon must be an integer >= 0: how many cached positions before a chunk receive its gradient (0 = the whole window)";
         else if (body.contains("memory_budget_mib") && !(num("memory_budget_mib", 0) >= 1))
                                                                      why = "memory_budget_mib must be >= 1";
+        else if (body.contains("context_cache_type") && !(body.at("context_cache_type").is_string() &&
+                                                          [&] { ggml_type t; return context_cache_type_of(body.at("context_cache_type").get<std::string>(), t); }()))
+                                                                     why = "context_cache_type must be \"q8_0\" (the default), \"q4_0\" or \"f16\": the training context's K/V type";
         else if (body.contains("share_ppm") && !(num("share_ppm", 0) >= 1 && num("share_ppm", 0) <= 1000000 && num("share_ppm", 0) == (int64_t) num("share_ppm", 0)))
                                                                      why = "share_ppm must be an integer in [1, 1000000]: the trainer's share of the lane's time, in parts per million";
         else if (body.contains("max_slowdown_ppm") && !(num("max_slowdown_ppm", 0) >= 1 && num("max_slowdown_ppm", 0) <= 999999 &&
@@ -578,6 +592,13 @@ json server_trainer::start(const json & body_in) {
     std::unique_lock<std::mutex> lock(mu);
     state = json::object({{"state", "starting"}, {"out", body.at("out")}, {"epochs", json::array()}});
     state["window"] = body.value("window", (int64_t) 256);
+    // THE SHAPE THIS RUN WAS ASKED FOR, recorded from its first moment, so every receipt (a
+    // refusal at the context, at a late chunk, or a finished run) says what it measured
+    // (continuum 34d574fc, Codex root on #59).
+    state["context_cache_type"] = body.value("context_cache_type", std::string("q8_0"));
+    state["chunk"]              = body.value("chunk", (int64_t) 512);
+    state["recompute"]          = body.value("recompute", false);
+    state["exact"]              = body.value("exact", false);
     if (body.contains("window_asked")) {
         state["window_asked"] = body.at("window_asked");
     }
@@ -919,8 +940,12 @@ void server_trainer::run(json req, examples_data ex) {
     // F16: the cache holds only CONSTANTS (the context before a chunk, read by build_attn and
     // cast to F32 where it joins the chunk). The chunk's own K/V, the ones with a gradient and
     // the only ones OUT_PROD sees, are F32 in the graph and never cached during training.
-    cparams.type_k          = GGML_TYPE_Q8_0;
-    cparams.type_v          = GGML_TYPE_Q8_0;
+    // Q8_0 unless the request names another type (validated before the claim).
+    const std::string cache_type = req.value("context_cache_type", std::string("q8_0"));
+    ggml_type cache_ggml = GGML_TYPE_Q8_0;
+    context_cache_type_of(cache_type, cache_ggml);
+    cparams.type_k          = cache_ggml;
+    cparams.type_v          = cache_ggml;
     cparams.embeddings      = false;
     // training reads logits for every token of the window; the serving params cap outputs per
     // ubatch to what sampling needs (a server-computed limit), which a training batch overruns
@@ -1129,6 +1154,13 @@ void server_trainer::run(json req, examples_data ex) {
     const bool    no_fit    = llama_opt_failed(ctx);
     // copied before llama_free: the refusal names a node the device cannot run, when that was it
     const std::string refusal = llama_opt_failure(ctx);
+    // THE CHUNK THAT DID NOT FIT, when one didn't (continuum 34d574fc, Codex root): where in the
+    // window it sat, what its graph needed and what the device allowed at that moment. A refusal
+    // at the deepest chunk with a need above the run's earlier peak, and a refusal of an ordinary
+    // chunk because the allowance dropped (serving took memory), read differently here.
+    uint32_t refused_pos = 0, refused_n = 0;
+    size_t   refused_need = 0, refused_allow = 0;
+    const bool refused_chunk = llama_opt_last_refusal(ctx, &refused_pos, &refused_n, &refused_need, &refused_allow);
     const double  graph_mib = llama_opt_graph_bytes(ctx) / 1048576.0;
     const double  walk_host_mib = llama_opt_walk_host_bytes(ctx) / 1048576.0;
     const int64_t walk_horizon  = llama_opt_walk_horizon(ctx);
@@ -1150,6 +1182,12 @@ void server_trainer::run(json req, examples_data ex) {
     if (req.value("exact", false)) {
         state["walk_horizon"] = walk_horizon;
     }
+    if (no_fit && refused_chunk) {
+        state["refused_chunk_pos"]     = refused_pos;
+        state["refused_chunk_tokens"]  = refused_n;
+        state["refused_need_mib"]      = refused_need / 1048576.0;
+        state["refused_allowance_mib"] = refused_allow / 1048576.0;
+    }
     if (no_fit) {
         state["state"] = "error";
         state["error"] = !refusal.empty() ? refusal + " (window " + std::to_string(window) + ")" :