diff --git a/ggml/include/ggml-opt.h b/ggml/include/ggml-opt.h index 148349e99e2..a9311f236f0 100644 --- a/ggml/include/ggml-opt.h +++ b/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); diff --git a/ggml/src/ggml-opt.cpp b/ggml/src/ggml-opt.cpp index 87cd73369d3..45bccad7f23 100644 --- a/ggml/src/ggml-opt.cpp +++ b/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; diff --git a/include/llama.h b/include/llama.h index 3135ea509aa..61c6ded22fb 100644 --- a/include/llama.h +++ b/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). diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 3c183f3d255..fd4ae19091c 100644 --- a/src/llama-context-opt.cpp +++ b/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(); } diff --git a/src/llama-context.h b/src/llama-context.h index 2f3fc48b6ef..e1270400965 100644 --- a/src/llama-context.h +++ b/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 diff --git a/tests/test-chat.cpp b/tests/test-chat.cpp index 201cea95d2c..67f7ec882df 100644 --- a/tests/test-chat.cpp +++ b/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(), "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(), "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(), "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); diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index 8dd7f4d8742..d61021abd1a 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -18,6 +18,17 @@ using json = common_json; static std::atomic g_train_batch{0}; static std::atomic 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 split_targets(const std::string & list) { std::vector 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(), 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 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) + ")" :