From 12cf0ea4eb38c55ae579bf8193812e1e908d5a98 Mon Sep 17 00:00:00 2001 From: joelteply Date: Sun, 1 Nov 2026 01:13:54 -0600 Subject: [PATCH] server: one shadow per mind on the serving context: each shadow keeps its own optimizer, /learn trains any loaded shadow, and unload releases it A serving context trained exactly one shadow per server run (learn_lora_id): the second mind's /learn was refused, so only one resident could ever learn in place (continuum S4 (a), card eaeb2b87). - llama_context keeps an optimizer per shadow (opt_minds); opt_ctx and opt_shadow point at the active one, so every training site reads the active pair as before. llama_opt_init with another adapter on a configured context selects it (opt_select), creating its optimizer the first time and keeping each shadow's moments and accumulators across switches. Only the active shadow's tensors carry the PARAM flag, so a graph never carries a parameter the active optimizer was not built with. - opt_configured() is "this context is sized for training" (scheduler and graph): it was read as opt_ctx != nullptr, which a released active shadow would have falsified while other shadows' optimizers still held the scheduler. - llama_opt_release forgets one shadow's optimizer; the server's unload releases a shadow that learned (its moments go; a snapshot is the durable form) instead of refusing it forever. - learn-every-format.py --two-minds: two shadows from the same file learn alternately; their losses match step for step (4.3179, 4.3179, 3.5043, 3.5043 on Qwen2.5-0.5B, IntelMac CPU), which holds only if each keeps its own weights and optimizer. --shadow-retirement now unloads a bound shadow and learns on a reloaded one. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01FRQWzgSfo79JtnwZywHKrE --- include/llama.h | 5 ++ src/llama-context-opt.cpp | 70 +++++++++++++++++++++++- src/llama-context.cpp | 11 +++- src/llama-context.h | 13 +++++ tools/server/server-context.cpp | 22 +++----- tools/server/tests/learn-every-format.py | 32 ++++++++--- 6 files changed, 128 insertions(+), 25 deletions(-) diff --git a/include/llama.h b/include/llama.h index 277fc9d881cb..388df5aebe6f 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1750,6 +1750,11 @@ extern "C" { }; LLAMA_API void llama_opt_init(struct llama_context * lctx, struct llama_model * model, struct llama_opt_params lopt_params); + // A serving context trains one shadow per mind: llama_opt_init with another adapter makes it the + // trained shadow, creating its optimizer the first time and keeping every shadow's moments across + // switches. llama_opt_release forgets one shadow's optimizer (before its adapter is unloaded); + // false when it had none. + LLAMA_API bool llama_opt_release(struct llama_context * lctx, const struct llama_adapter_lora * adapter); // Restore one shadow generation in place, preserving optimizer tensor identity. The caller // owns verified same-generation files/binding, stable adapter exclusion, and serialized queue // custody. Both payloads are prepared before mutation. An interrupted apply quarantines the diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index ddc31f35f58d..8589c3702f66 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -80,7 +80,12 @@ bool llama_context::opt_fa_back_native() const { } void llama_context::opt_init(struct llama_model * model, struct llama_opt_params lopt_params) { - GGML_ASSERT(!opt_ctx); + // a serving context already configured for its first shadow: another mind's shadow is a switch + if (lopt_params.adapter != nullptr && !opt_minds.empty()) { + opt_select(lopt_params.adapter, lopt_params); + return; + } + GGML_ASSERT(!opt_configured()); // The training window lives on this context, not the shared model (a serving context may // be running on the same weights). opt_n_ctx_train = lopt_params.n_ctx_train > 0 ? lopt_params.n_ctx_train : n_ctx(); @@ -132,6 +137,7 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params opt_params.optimizer = lopt_params.optimizer_type; opt_ctx = ggml_opt_init(opt_params); ggml_opt_set_alloc_budget(opt_ctx, opt_memory_budget); + opt_params_base = opt_params; llama_opt_param_filter param_filter = lopt_params.param_filter; void * param_filter_ud = lopt_params.param_filter_ud; @@ -151,6 +157,7 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params GGML_ASSERT(n_params > 0 && "the param filter selected none of the adapter's tensors"); // the trained adapter is the shadow: labelled rows only, never a serving row opt_shadow = lopt_params.adapter; + opt_minds.emplace(lopt_params.adapter, opt_ctx); LLAMA_LOG_INFO("%s: LoRA-only training: %lld adapter tensors trainable, base frozen\n", __func__, (long long) n_params); return; } @@ -178,6 +185,63 @@ void llama_context::opt_init(struct llama_model * model, struct llama_opt_params } } +// Another mind's shadow on a configured serving context. Only the active shadow's tensors are +// parameters, so a graph never carries a parameter the active optimizer was not built with; each +// shadow keeps its own optimizer (moments, accumulators), so a switch back resumes where it was. +void llama_context::opt_select(const llama_adapter_lora * adapter, struct llama_opt_params lopt_params) { + if (adapter == opt_shadow) { + return; + } + const auto attached = loras->find(const_cast(adapter)); + GGML_ASSERT(attached != loras->end() && attached->second != 0.0f && "the LoRA to train must be attached with a non-zero scale"); + if (opt_shadow != nullptr) { + for (const auto & [name, ab] : opt_shadow->ab_map) { + ab.a->flags &= ~GGML_TENSOR_FLAG_PARAM; + ab.b->flags &= ~GGML_TENSOR_FLAG_PARAM; + } + } + int64_t n_params = 0; + for (const auto & [name, ab] : adapter->ab_map) { + GGML_ASSERT(ab.a->type == GGML_TYPE_F32 && ab.b->type == GGML_TYPE_F32 && "a trainable LoRA holds F32 A/B tensors"); + n_params += llama_set_param(ab.a, lopt_params.param_filter, lopt_params.param_filter_ud); + n_params += llama_set_param(ab.b, lopt_params.param_filter, lopt_params.param_filter_ud); + } + GGML_ASSERT(n_params > 0 && "the param filter selected none of the adapter's tensors"); + const auto mind = opt_minds.find(adapter); + if (mind != opt_minds.end()) { + opt_ctx = mind->second; + } else { + ggml_opt_params params = opt_params_base; + params.get_opt_pars = lopt_params.get_opt_pars; + params.get_opt_pars_ud = lopt_params.get_opt_pars_ud; + params.optimizer = lopt_params.optimizer_type; + opt_ctx = ggml_opt_init(params); + ggml_opt_set_alloc_budget(opt_ctx, opt_memory_budget); + opt_minds.emplace(adapter, opt_ctx); + LLAMA_LOG_INFO("%s: a new shadow's optimizer: %lld adapter tensors trainable (%zu shadows)\n", __func__, + (long long) n_params, opt_minds.size()); + } + opt_shadow = adapter; +} + +bool llama_context::opt_release(const llama_adapter_lora * adapter) { + const auto mind = opt_minds.find(adapter); + if (mind == opt_minds.end()) { + return false; + } + if (adapter == opt_shadow) { + for (const auto & [name, ab] : adapter->ab_map) { + ab.a->flags &= ~GGML_TENSOR_FLAG_PARAM; + ab.b->flags &= ~GGML_TENSOR_FLAG_PARAM; + } + opt_shadow = nullptr; + opt_ctx = nullptr; + } + ggml_opt_free(mind->second); + opt_minds.erase(mind); + return true; +} + // THE WALK (Fable's design, continuum 2026-10-06; objective 2026-10-14): one pass over a window // in one context. A run with no label is context: a plain decode into the cache, as served, no // backward. A labelled run is her reply, trained GIVEN EVERYTHING BEFORE IT AS SERVED: only its @@ -755,6 +819,10 @@ void llama_opt_init(struct llama_context * ctx, struct llama_model * model, stru ctx->opt_init(model, lopt_params); } +bool llama_opt_release(struct llama_context * ctx, const struct llama_adapter_lora * adapter) { + return ctx->opt_release(adapter); +} + void llama_opt_stop(struct llama_context * ctx, bool stop) { ctx->opt_stop_requested.store(stop, std::memory_order_relaxed); } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 654d386b63f3..5d0cd5c70d4f 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -493,7 +493,12 @@ llama_context::~llama_context() { } } } - ggml_opt_free(opt_ctx); + for (auto & [adapter, mind] : opt_minds) { + ggml_opt_free(mind); + } + if (opt_minds.empty()) { + ggml_opt_free(opt_ctx); // the base's own optimizer, when the base trained + } } void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint32_t n_seqs) { @@ -580,7 +585,7 @@ void llama_context::sched_reserve() { // A training context's scheduler was sized for the training graph in opt_init, and the // optimizer holds it (ggml_opt_params.backend_sched). Re-creating it here, as the first // context decode of opt_epoch_iter's walk would, leaves the optimizer on a freed scheduler. - if (opt_ctx != nullptr) { + if (opt_configured()) { sched_need_reserve = false; return; } @@ -2398,7 +2403,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { // a training context (an optimizer exists), not this graph's flag: the walk decodes context // with cparams.training off, and a sched_reserve() run then must not shrink the scheduler // and the graph result below what the next training chunk needs - if (cparams.training || opt_ctx != nullptr) { + if (cparams.training || opt_configured()) { // each attention layer reads the cached context in front of the chunk (build_attn: // views of K and V, V's permute and cont, two concats, the mask's view and cont) res += 16u * model.hparams.n_layer(); diff --git a/src/llama-context.h b/src/llama-context.h index 441076c49746..1c6c78bf43db 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -203,6 +203,10 @@ struct llama_context { // void opt_init(struct llama_model * model, struct llama_opt_params lopt_params); + // make `adapter` the trained shadow, creating its optimizer the first time it trains + void opt_select(const llama_adapter_lora * adapter, struct llama_opt_params lopt_params); + // forget `adapter`'s optimizer (its moments go); false when it has none + bool opt_release(const llama_adapter_lora * adapter); // TODO: more flexible combinations of logical/physical batch size and context size // set by llama_opt_stop from any thread; opt_epoch returns at the next data item @@ -398,7 +402,16 @@ struct llama_context { std::vector backends; // training + // the ACTIVE optimizer: the one adapter-or-base this context's next training step updates ggml_opt_context_t opt_ctx = nullptr; + // one optimizer per shadow on a serving context (each mind's moments and accumulators, kept + // across switches); opt_ctx points at the active one. Empty when the base itself trains. + std::map opt_minds; + // the optimizer parameters every shadow's optimizer is created with (set by the first opt_init) + ggml_opt_params opt_params_base{}; + // this context is configured for training (its scheduler and graph sized for it, and some + // optimizer holds that scheduler), whether or not a step is active + bool opt_configured() const { return opt_ctx != nullptr || !opt_minds.empty(); } // the training window this context's opt loop runs at; kept HERE, never written into the // model's hparams, because several contexts share one model (a serving context and a // training context on the same resident weights) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 87672928aeb3..536dee30ac02 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -900,9 +900,8 @@ struct server_context_impl { common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; - // LEARNING IN THE SERVING CONTEXT (SERVER_TASK_TYPE_LEARN): the shadow adapter llama_opt_init was - // given (one per server run for now), and the learning rate its next step reads - int32_t learn_lora_id = -1; + // LEARNING IN THE SERVING CONTEXT (SERVER_TASK_TYPE_LEARN): the learning rate the next step reads + // (each mind's shadow has its own optimizer in the context: llama_opt_init selects it) float learn_lr = 1e-4f; // LEARNING BARELY TOUCHES THE FLOW: while any other slot is working, learning may take at most // LEARN_DUTY of the loop's time (the step after a learn of t ms waits t*(1/LEARN_DUTY - 1) ms); @@ -2099,9 +2098,6 @@ struct server_context_impl { if (L.lora_id < 0 || (size_t) L.lora_id >= loras.size() || loras[L.lora_id].ptr == nullptr) { return refused("no loaded adapter " + std::to_string(L.lora_id) + " to train (POST /lora-adapters/load it first)"); } - if (learn_lora_id >= 0 && learn_lora_id != L.lora_id) { - return refused("this server trains adapter " + std::to_string(learn_lora_id) + "; one shadow per server run"); - } const int32_t n0 = L.n_prefix - 1; // the last prompt token is re-decoded with the reply's first token as its label if (n0 < 0 || L.n_prefix > slot.prompt.n_tokens() || L.reply.empty()) { return refused("slot " + std::to_string(slot.id) + " holds " + std::to_string(slot.prompt.n_tokens()) @@ -2149,7 +2145,9 @@ struct server_context_impl { set[L.lora_id].scale = 1.0f; common_set_adapter_lora(ctx_tgt, set); learn_lr = L.lr; - if (learn_lora_id < 0) { + { + // her shadow is the trained one: its own optimizer, created the first time she learns + // and kept (moments and all) while other minds' shadows train in between llama_opt_params lopt{}; lopt.param_filter = llama_opt_param_filter_all; lopt.get_opt_pars = learn_opt_pars; @@ -2157,7 +2155,6 @@ struct server_context_impl { lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_ADAMW; lopt.adapter = loras[L.lora_id].ptr; llama_opt_init(ctx_tgt, model_tgt, lopt); - learn_lora_id = L.lora_id; } // drop what follows n0 (attention) and put her recurrent state back at n0 if (cp != nullptr) { @@ -3177,12 +3174,6 @@ struct server_context_impl { send_error(task, "no loaded LoRA adapter with id " + std::to_string(id), ERROR_TYPE_INVALID_REQUEST); break; } - if (id == learn_lora_id) { - // Optimizer moments and opt_shadow still reference this adapter's tensors. - send_error(task, "LoRA adapter " + std::to_string(id) + - " is bound to the learning optimizer; unload requires optimizer teardown", ERROR_TYPE_INVALID_REQUEST); - break; - } llama_adapter_lora * ptr = loras[id].ptr; if (loras[id].scale != 0.0f) { send_error(task, "LoRA adapter " + std::to_string(id) + " is active by default (scale " + @@ -3201,6 +3192,9 @@ struct server_context_impl { if (ptr == nullptr) { break; } + // a shadow that learned holds an optimizer referencing its tensors: release it + // first (its moments go; a snapshot is the durable form of what it learned) + llama_opt_release(ctx_tgt, ptr); // An idle slot that last ran with it keeps KV computed WITH the adapter; // once the adapter is gone its list would read as the base's, so that // cache must not be reused. diff --git a/tools/server/tests/learn-every-format.py b/tools/server/tests/learn-every-format.py index 3e492ef0a3d0..6523b9256c03 100755 --- a/tools/server/tests/learn-every-format.py +++ b/tools/server/tests/learn-every-format.py @@ -27,15 +27,33 @@ def raw(path, body): turn = {"prompt": "The capital of Kansas is", "id_slot": 0, "n_predict": 1, "temperature": 0, "learn": learn} first = json.loads(raw("/completion", turn)) assert first["learn"]["success"], first - try: - raw("/lora-adapters/unload", {"id": sid}) - raise AssertionError("optimizer-bound shadow was unloaded") - except urllib.error.HTTPError as error: - detail = error.read().decode() - assert error.code == 400 and "bound to the learning optimizer" in detail, (error.code, detail) + # a shadow that learned unloads too: its optimizer is released first (its moments go) + bound = json.loads(raw("/lora-adapters/unload", {"id": sid})) + assert bound["success"] and bound["retired"] and bound["id"] == sid, bound + sid = json.loads(raw("/lora-adapters/load", {"name": "shadow-05.gguf"}))["id"] + learn["lora_id"] = sid following = json.loads(raw("/completion", turn)) assert following["learn"]["success"], following - print(json.dumps({"unbound_retired": retired, "bound_retirement_refused": sid, "next_learn": following["learn"]})) + print(json.dumps({"unbound_retired": retired, "bound_retired": bound, "next_learn": following["learn"]})) + sys.exit(0) + +# Two minds on one serving context: each shadow keeps its own optimizer while the other learns in +# between (needs shadow-05.gguf and a copy of it, shadow-06.gguf, in --train-dir). +if "--two-minds" in sys.argv: + other = json.loads(raw("/lora-adapters/load", {"name": "shadow-06.gguf"}))["id"] + steps = [] + for lora in [sid, other, sid, other]: + turn = {"prompt": "The capital of Kansas is", "id_slot": 0, "n_predict": 1, "temperature": 0, + "learn": dict(learn, lora_id=lora)} + out = json.loads(raw("/completion", turn)) + assert out["learn"]["success"], out + steps.append((lora, out["learn"]["loss"])) + # the two shadows start from the same file, so they are independent only if their trajectories + # match: the other mind's first step is untouched by this one's step, and each second step + # continues its own optimizer (a shared optimizer or shared weights would split them) + assert abs(steps[0][1] - steps[1][1]) < 1e-6 and abs(steps[2][1] - steps[3][1]) < 1e-6, steps + assert steps[2][1] < steps[0][1], steps + print(json.dumps({"steps": steps})) sys.exit(0) # Isolated template/label regression; never run this shadow-adapter fixture on a persona lane.