diff --git a/common/arg.cpp b/common/arg.cpp index 7ce7eaad16ab..eded4fe05d0a 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -301,34 +301,6 @@ struct handle_model_result { std::string preset_path; }; -const std::vector kv_cache_types = { - GGML_TYPE_F32, - GGML_TYPE_F16, - GGML_TYPE_BF16, - GGML_TYPE_Q8_0, - GGML_TYPE_Q4_0, - GGML_TYPE_Q4_1, - GGML_TYPE_IQ4_NL, - GGML_TYPE_Q5_0, - GGML_TYPE_Q5_1, -}; - -static ggml_type kv_cache_type_from_str(const std::string & s) { - for (const auto & type : kv_cache_types) { - if (ggml_type_name(type) == s) { - return type; - } - } - throw std::runtime_error("Unsupported cache type: " + s); -} - -static std::string get_all_kv_cache_types() { - std::ostringstream msg; - for (const auto & type : kv_cache_types) { - msg << ggml_type_name(type) << (&type == &kv_cache_types.back() ? "" : ", "); - } - return msg.str(); -} static bool parse_bool_value(const std::string & value) { if (is_truthy(value)) { @@ -2428,11 +2400,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "KV cache data type for K\n" "allowed values: %s\n" "(default: %s)", - get_all_kv_cache_types().c_str(), + common_kv_cache_types_str().c_str(), ggml_type_name(params.cache_type_k) ), [](common_params & params, const std::string & value) { - params.cache_type_k = kv_cache_type_from_str(value); + params.cache_type_k = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_CACHE_TYPE_K")); add_opt(common_arg( @@ -2441,11 +2413,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "KV cache data type for V\n" "allowed values: %s\n" "(default: %s)", - get_all_kv_cache_types().c_str(), + common_kv_cache_types_str().c_str(), ggml_type_name(params.cache_type_v) ), [](common_params & params, const std::string & value) { - params.cache_type_v = kv_cache_type_from_str(value); + params.cache_type_v = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_CACHE_TYPE_V")); add_opt(common_arg( @@ -4052,11 +4024,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "KV cache data type for K for the draft model\n" "allowed values: %s\n" "(default: %s)", - get_all_kv_cache_types().c_str(), + common_kv_cache_types_str().c_str(), ggml_type_name(params.speculative.draft.cache_type_k) ), [](common_params & params, const std::string & value) { - params.speculative.draft.cache_type_k = kv_cache_type_from_str(value); + params.speculative.draft.cache_type_k = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K")); add_opt(common_arg( @@ -4065,11 +4037,11 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "KV cache data type for V for the draft model\n" "allowed values: %s\n" "(default: %s)", - get_all_kv_cache_types().c_str(), + common_kv_cache_types_str().c_str(), ggml_type_name(params.speculative.draft.cache_type_v) ), [](common_params & params, const std::string & value) { - params.speculative.draft.cache_type_v = kv_cache_type_from_str(value); + params.speculative.draft.cache_type_v = common_kv_cache_type_from_str(value); } ).set_env("LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V")); add_opt(common_arg( diff --git a/common/common.cpp b/common/common.cpp index 8147060448d7..9405692b9ac3 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1721,6 +1721,138 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { return mparams; } +const std::vector & common_kv_cache_types() { + static const std::vector types = { + GGML_TYPE_F32, + GGML_TYPE_F16, + GGML_TYPE_BF16, + GGML_TYPE_Q8_0, + GGML_TYPE_Q4_0, + GGML_TYPE_Q4_1, + GGML_TYPE_IQ4_NL, + GGML_TYPE_Q5_0, + GGML_TYPE_Q5_1, + }; + return types; +} + +std::string common_kv_cache_types_str() { + std::ostringstream msg; + const auto & types = common_kv_cache_types(); + for (const auto & type : types) { + msg << ggml_type_name(type) << (&type == &types.back() ? "" : ", "); + } + return msg.str(); +} + +ggml_type common_kv_cache_type_from_str(const std::string & s) { + for (const auto & type : common_kv_cache_types()) { + if (ggml_type_name(type) == s) { + return type; + } + } + throw std::runtime_error("Unsupported cache type: " + s); +} + +// The formats a training context has been RUN with, each with a measured learning curve: F16 +// (c424b00c3) and Q8_0 (a8a235df5, 1.5B chunk 256: forward loss 0.74694 against the previous +// engine's 0.74658; eval 0.621 -> 0.454 over 3 epochs). The cache holds only constants (the context before a chunk, cast to F32 +// where it joins the chunk), but a coarser format's effect on learning is a measurement, not an +// assumption: until its curve is posted beside these, it is not a training format. +static bool common_training_kv_measured(ggml_type t) { + return t == GGML_TYPE_F16 || t == GGML_TYPE_Q8_0; +} + +bool common_training_kv_resolve( + const llama_model * model, + const common_params & serving, + const std::string & requested_k, + const std::string & requested_v, + common_training_kv & out, + std::string & why) { + // the training context always runs its context decodes with flash attention + const auto usable = [&](ggml_type k, ggml_type v, std::string & reason) { + if (!common_training_kv_measured(k) || !common_training_kv_measured(v)) { + reason = std::string("training has no measured run with a ") + ggml_type_name(k) + "/" + ggml_type_name(v) + + " cache (measured: f16, q8_0)"; + return false; + } + char buf[256]; + if (!llama_kv_cache_types_check(model, k, v, /*flash_attn=*/ true, buf, sizeof(buf))) { + reason = buf; + return false; + } + return true; + }; + if (requested_k.empty() != requested_v.empty()) { + why = "a training cache format names both cache_type_k and cache_type_v, or neither"; + return false; + } + if (!requested_k.empty()) { + ggml_type k; + ggml_type v; + try { + k = common_kv_cache_type_from_str(requested_k); + v = common_kv_cache_type_from_str(requested_v); + } catch (const std::exception & e) { + why = std::string(e.what()) + " (allowed: " + common_kv_cache_types_str() + ")"; + return false; + } + std::string reason; + if (!usable(k, v, reason)) { + why = "the requested training cache " + requested_k + "/" + requested_v + " can't be used: " + reason; + return false; + } + out = { k, v, "requested" }; + return true; + } + std::string reason; + if (usable(serving.cache_type_k, serving.cache_type_v, reason)) { + out = { serving.cache_type_k, serving.cache_type_v, "serving" }; + return true; + } + if (!usable(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, reason)) { + why = "no training cache format fits this model: " + reason; + return false; + } + out = { GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, "default" }; + return true; +} + +struct llama_context_params common_context_params_for_training( + const common_params & serving, + uint32_t window, + uint32_t chunk, + const common_training_kv & kv) { + llama_context_params cparams = common_context_params_to_llama(serving); + // NO ROLLBACK SLOTS. Serving keeps recurrent-state snapshots for speculative decoding's + // rollback (n_rs_seq, from the MTP draft); a training context never rolls back, and the exact + // walk carries one recurrent state per sequence. Inherited, they crashed the 5090's serving + // process on Kimi's first exact run (2026-10-07 04:09Z: the delta-net invariant asserted). + cparams.n_rs_seq = 0; + // the context holds the whole window; one training chunk is one batch is one ubatch + cparams.n_ctx = window; + cparams.n_batch = chunk; + cparams.n_ubatch = chunk; + // two sequences: 0 is the walk; 1 is the scratch a recurrent state is snapshotted into + // around each training chunk (llama_context::opt_epoch_iter). Unified, so the attention + // cache stays one stream of the window, not one per sequence. + cparams.n_seq_max = 2; + cparams.kv_unified = true; + // Flash attention ON for the context: the walk's context decodes use it, and the cache then + // stores V un-transposed, which a quantized V cache requires. The training graphs build + // explicit attention regardless (build_attn_mha under cparams.training). + cparams.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; + cparams.type_k = kv.type_k; + cparams.type_v = kv.type_v; + 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 + cparams.n_outputs_max = 0; // = n_batch + cparams.n_outputs_max_per_seq = 0; // = n_outputs_max + return cparams; +} + struct llama_context_params common_context_params_to_llama(const common_params & params) { auto cparams = llama_context_default_params(); diff --git a/common/common.h b/common/common.h index 29b171b7dda0..e9ff7334b1ce 100644 --- a/common/common.h +++ b/common/common.h @@ -941,6 +941,40 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode struct llama_model_params common_model_params_to_llama ( common_params & params); struct llama_context_params common_context_params_to_llama(const common_params & params); +// K/V cache formats a context may be given, by name: the ONE parser (--cache-type-k/v and /train) +const std::vector & common_kv_cache_types(); +std::string common_kv_cache_types_str(); +ggml_type common_kv_cache_type_from_str(const std::string & s); // throws on an unknown name + +// The training context's K/V format and where it came from: "requested" (the /train caller), +// "serving" (serving's own resolved format) or "default" (neither was usable; Q8_0). Recorded in +// the run's receipt beside serving's format (continuum card 6edd495c). +struct common_training_kv { + ggml_type type_k = GGML_TYPE_Q8_0; + ggml_type type_v = GGML_TYPE_Q8_0; + std::string source; +}; + +// Resolve the training context's K/V format: the requested one (both names, or neither) if +// training has been measured with it and the model can hold it; else serving's if so; else the +// default. Only a REQUESTED format is refused (false, with `why`); serving's never is. +bool common_training_kv_resolve( + const llama_model * model, + const common_params & serving, + const std::string & requested_k, + const std::string & requested_v, + common_training_kv & out, + std::string & why); + +// The training context: serving's params with only the training mode's differences +// (no rollback slots, one chunk per batch/ubatch, a scratch sequence in a unified cache, +// logits for every token, flash attention for the context decodes, no embeddings). +struct llama_context_params common_context_params_for_training( + const common_params & serving, + uint32_t window, + uint32_t chunk, + const common_training_kv & kv); + // clear LoRA adapters from context, then apply new list of adapters void common_set_adapter_lora(struct llama_context * ctx, std::vector & lora); diff --git a/include/llama.h b/include/llama.h index 3135ea509aa2..552ab1d66c4e 100644 --- a/include/llama.h +++ b/include/llama.h @@ -532,6 +532,16 @@ extern "C" { LLAMA_API void llama_model_free(struct llama_model * model); + // Whether a K/V cache format can hold this model's attention (flash_attn = the context's + // resolved setting). The same check context creation refuses by. On false, `why` names the rule. + LLAMA_API bool llama_kv_cache_types_check( + const struct llama_model * model, + enum ggml_type type_k, + enum ggml_type type_v, + bool flash_attn, + char * why, + size_t why_len); + LLAMA_API struct llama_context * llama_init_from_model( struct llama_model * model, struct llama_context_params params); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a296f0c2cbdb..721c9dd95799 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -3358,6 +3358,38 @@ llama_context_params llama_context_default_params() { return result; } +// THE ONE OWNER of whether a K/V cache format can hold this model's attention: context creation +// refuses by it, and a caller choosing a format (the training context, continuum card 6edd495c) +// asks it first instead of reading a null context, which also means out of memory. +bool llama_kv_cache_types_check( + const llama_model * model, + ggml_type type_k, + ggml_type type_v, + bool flash_attn, + char * why, + size_t why_len) { + std::string err; + if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && type_k != type_v) { + err = std::string("model does not support different K (") + ggml_type_name(type_k) + ") and V (" + ggml_type_name(type_v) + ") cache types"; + } else if (ggml_is_quantized(type_v) && !flash_attn) { + err = "quantized V cache requires flash_attn to be enabled"; + } else if (flash_attn) { + for (uint32_t il = 0; il < model->hparams.n_layer() && err.empty(); ++il) { + if (ggml_is_quantized(type_k) && model->hparams.n_embd_head_k(il) % ggml_blck_size(type_k) != 0) { + err = std::string("K cache type ") + ggml_type_name(type_k) + " with block size " + std::to_string(ggml_blck_size(type_k)) + + " does not divide n_embd_head_k=" + std::to_string(model->hparams.n_embd_head_k(il)); + } else if (ggml_is_quantized(type_v) && model->hparams.n_embd_head_v(il) % ggml_blck_size(type_v) != 0) { + err = std::string("V cache type ") + ggml_type_name(type_v) + " with block size " + std::to_string(ggml_blck_size(type_v)) + + " does not divide n_embd_head_v=" + std::to_string(model->hparams.n_embd_head_v(il)); + } + } + } + if (why != nullptr && why_len > 0) { + snprintf(why, why_len, "%s", err.c_str()); + } + return err.empty(); +} + llama_context * llama_init_from_model( llama_model * model, llama_context_params params) { @@ -3392,44 +3424,20 @@ llama_context * llama_init_from_model( } } - if ((model->hparams.is_mla() || model->arch == LLM_ARCH_DEEPSEEK4) && params.type_k != params.type_v) { - LLAMA_LOG_ERROR("%s: model does not support different K (%s) and V (%s) cache types\n", __func__, ggml_type_name(params.type_k), ggml_type_name(params.type_v)); - return nullptr; + if (ggml_is_quantized(params.type_v) && params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) { + LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__); + params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; } - if (ggml_is_quantized(params.type_v) && params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_ENABLED) { - if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO) { - LLAMA_LOG_INFO("%s: enabling flash_attn since it is required for quantized V cache\n", __func__); - params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; - } - if (params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_DISABLED) { - LLAMA_LOG_ERROR("%s: quantized V cache requires flash_attn to be enabled\n", __func__); + { + char why[256]; + if (!llama_kv_cache_types_check(model, params.type_k, params.type_v, + params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED, why, sizeof(why))) { + LLAMA_LOG_ERROR("%s: %s\n", __func__, why); return nullptr; } } - if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_k)) { - const uint32_t blck_size = ggml_blck_size(params.type_k); - for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) { - if (model->hparams.n_embd_head_k(il) % blck_size != 0) { - LLAMA_LOG_ERROR("%s: K cache type %s with block size %u does not divide n_embd_head_k=%u\n", - __func__, ggml_type_name(params.type_k), blck_size, model->hparams.n_embd_head_k(il)); - return nullptr; - } - } - } - - if (params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED && ggml_is_quantized(params.type_v)) { - const uint32_t blck_size = ggml_blck_size(params.type_v); - for (uint32_t il = 0; il < model->hparams.n_layer(); ++il) { - if (model->hparams.n_embd_head_v(il) % blck_size != 0) { - LLAMA_LOG_ERROR("%s: V cache type %s with block size %u does not divide n_embd_head_v=%u\n", - __func__, ggml_type_name(params.type_v), blck_size, model->hparams.n_embd_head_v(il)); - return nullptr; - } - } - } - if (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != model->hparams.pooling_type) { //user-specified pooling-type is different from the model default diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index 8dd7f4d87426..68f778c1206c 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -889,13 +889,6 @@ void server_trainer::run(json req, examples_data ex) { // governed lease); without it the driver's own free figure, which on Windows is not physical const size_t budget = (size_t) (req.value("memory_budget_mib", 0.0) * 1024.0 * 1024.0); - // The training context: the SAME model, its own graph; flash attention has no backward. - llama_context_params cparams = common_context_params_to_llama(params_base); - // NO ROLLBACK SLOTS. Serving keeps recurrent-state snapshots for speculative decoding's - // rollback (n_rs_seq, from the MTP draft); a training context never rolls back, and the exact - // walk carries one recurrent state per sequence. Inherited, they crashed the 5090's serving - // process on Kimi's first exact run (2026-10-07 04:09Z: the delta-net invariant asserted). - cparams.n_rs_seq = 0; // The context holds the whole window; one training chunk is one batch is one ubatch, the // largest multiple of 256 that divides the window and is at most the caller's "chunk" (the // memory gate's S). Context she did not write is decoded; her replies train chunk by chunk, @@ -904,28 +897,19 @@ void server_trainer::run(json req, examples_data ex) { // (Fable); the window's whole length as one chunk is the window x window memory the walk // exists to avoid. const uint32_t chunk = train_chunk_for(window, (uint32_t) req.value("chunk", (int64_t) 512)); - cparams.n_ctx = window; - cparams.n_batch = chunk; - cparams.n_ubatch = chunk; - // two sequences: 0 is the walk; 1 is the scratch a recurrent state is snapshotted into - // around each training chunk (llama_context::opt_epoch_iter). Unified, so the attention - // cache stays one stream of the window, not one per sequence. - cparams.n_seq_max = 2; - cparams.kv_unified = true; - // Flash attention ON for the context: the walk's context decodes use it, and the cache then - // stores V un-transposed, which a quantized V cache requires. The training graphs build - // explicit attention regardless (build_attn_mha under cparams.training). - cparams.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED; - // 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; - 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 - cparams.n_outputs_max = 0; // = n_batch - cparams.n_outputs_max_per_seq = 0; // = n_outputs_max + // The training context: the SAME model, its own graph, built from serving's params by the one + // shared derivation. Its K/V format is resolved, never forced: requested, else serving's, else + // the default, each only where training has a measured run with it (common_training_kv_resolve). + common_training_kv kv; + { + std::string why; + if (!common_training_kv_resolve(model, params_base, req.value("cache_type_k", std::string()), + req.value("cache_type_v", std::string()), kv, why)) { + fail(why); + return; + } + } + const llama_context_params cparams = common_context_params_for_training(params_base, window, chunk, kv); llama_context * ctx = llama_init_from_model(model, cparams); if (ctx == nullptr) { @@ -1018,7 +1002,7 @@ void server_trainer::run(json req, examples_data ex) { /*walk_horizon =*/ (uint32_t) req.value("walk_horizon", (int64_t) 0), }; // THE BUDGET COVERS THE WHOLE RUN (continuum, the 5090, 2026-10-11): the training context's own - // memory (its Q8_0 K/V for the whole window, and any recurrent state) is allocated when it is + // memory (its K/V for the whole window, and any recurrent state) is allocated when it is // created, before the graph's budget applies, so it was outside the caller's // governed lease: ~0.6 GB at a 16,896-token window, growing with the window, beside serving // on the same device. The graph gets what the lease leaves after them (at least one byte: @@ -1143,6 +1127,10 @@ void server_trainer::run(json req, examples_data ex) { // what the training context's memory held on the GPU (K/V for the window, recurrent state): // with the graph (which covers the compute buffers), what the run added beside serving state["context_mib"] = context_bytes / 1048576.0; + // the two K/V formats beside each other, and why training's is what it is: serving's own, a + // requested one, or the default where serving's has no measured training run + state["training_cache"] = { {"type_k", ggml_type_name(kv.type_k)}, {"type_v", ggml_type_name(kv.type_v)}, {"source", kv.source} }; + state["serving_cache"] = { {"type_k", ggml_type_name(params_base.cache_type_k)}, {"type_v", ggml_type_name(params_base.cache_type_v)} }; // what the exact walk kept on the host per window (its K/V gradient and state snapshots) state["walk_host_mib"] = walk_host_mib; // the exact walk's gradient horizon as trained (0 = the whole window), smaller where the device