diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 536dee30ac02..2521c1a57227 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -903,6 +903,8 @@ struct server_context_impl { // 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; + // each shadow's committed optimizer steps on this server (a snapshot records where she was) + std::map learn_steps; // 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); // with every other slot idle it runs at once. Back to back, learns took the 4B's other slot from @@ -2207,6 +2209,7 @@ struct server_context_impl { slot.prompt.checkpoints.remove_if([&](const common_prompt_checkpoint & c) { return c.n_tokens > n_keep; }); learn_outcome out; out.ok = ok; + out.step = ok ? ++learn_steps[loras[L.lora_id].ptr] : 0; out.failure = ok ? "" : llama_opt_failure(ctx_tgt); out.loss = loss; out.n_labels = (int32_t) toks.size() - 1; @@ -3166,6 +3169,44 @@ struct server_context_impl { res->outcome = learn_on_slot(*slot, L, t0); queue_results.send(std::move(res)); } break; + case SERVER_TASK_TYPE_SNAPSHOT_LORA: + { + // Between steps on the loop, so the file is one whole step's weights, never a + // half-applied update. Written beside its name and renamed into place, and + // never over an existing file: a snapshot is immutable. + auto & loras = params_base.lora_adapters; + const int32_t id = task.snapshot_lora; + if (id < 0 || (size_t) id >= loras.size() || loras[id].ptr == nullptr) { + send_error(task, "no loaded LoRA adapter with id " + std::to_string(id), ERROR_TYPE_INVALID_REQUEST); + break; + } + std::error_code ec; + if (std::filesystem::exists(task.snapshot_path, ec)) { + send_error(task, "a snapshot is immutable: " + task.snapshot_path + " already exists", ERROR_TYPE_INVALID_REQUEST); + break; + } + const std::string partial = task.snapshot_path + ".partial"; + if (llama_adapter_lora_save(loras[id].ptr, partial.c_str()) != 0) { + std::filesystem::remove(partial, ec); + send_error(task, "could not write the snapshot of adapter " + std::to_string(id), ERROR_TYPE_SERVER); + break; + } + std::filesystem::rename(partial, task.snapshot_path, ec); + if (ec) { + std::filesystem::remove(partial, ec); + send_error(task, "could not place the snapshot at " + task.snapshot_path, ERROR_TYPE_SERVER); + break; + } + auto res = std::make_unique(); + res->id = task.id; + res->lora_id = id; + res->path = task.snapshot_path; + res->bytes = (int64_t) std::filesystem::file_size(task.snapshot_path, ec); + const auto steps = learn_steps.find(loras[id].ptr); + res->steps = steps == learn_steps.end() ? 0 : steps->second; + SRV_INF("snapshot of lora adapter id=%d at %lld steps: %s\n", id, (long long) res->steps, res->path.c_str()); + queue_results.send(std::move(res)); + } break; case SERVER_TASK_TYPE_UNLOAD_LORA: { auto & loras = params_base.lora_adapters; @@ -3195,6 +3236,7 @@ struct server_context_impl { // 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); + learn_steps.erase(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. @@ -6112,6 +6154,37 @@ void server_routes::init_routes() { return res; }; + this->post_lora_adapters_snapshot = [this](const server_http_req & req) { + auto res = create_response(); + const json body = json::parse(req.body); + const std::string name = json_value(body, "name", std::string()); + const std::string path = confine_out(params.train_dir, name); + if (!body.contains("id") || !body.at("id").is_number_integer() || params.train_dir.empty() || path.empty()) { + res->error(format_error_response("\"id\" (an integer adapter id) and \"name\" (a bare file name ending in .gguf " + "inside the server's --train-dir) are required", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + auto & rd = res->rd; + { + server_task task(SERVER_TASK_TYPE_SNAPSHOT_LORA); + task.id = rd.get_new_id(); + task.snapshot_lora = body.at("id").get(); + task.snapshot_path = path; + rd.post_task(std::move(task)); + } + auto result = rd.next(req.should_stop); + if (!result) { + GGML_ASSERT(req.should_stop()); + return res; + } + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + res->ok(result->to_json()); + return res; + }; + this->post_lora_adapters_unload = [this](const server_http_req & req) { auto res = create_response(); const json body = json::parse(req.body); diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 6deb3a13cc06..fd079a69cb86 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -169,6 +169,7 @@ struct server_routes { server_http_context::handler_t post_lora_adapters; server_http_context::handler_t post_lora_adapters_load; server_http_context::handler_t post_lora_adapters_unload; + server_http_context::handler_t post_lora_adapters_snapshot; server_http_context::handler_t post_learn; server_http_context::handler_t get_train; server_http_context::handler_t post_train; diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 746c3a1f98c8..39eeee6e7060 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1727,6 +1727,9 @@ json server_task_result_lora_change::to_json() { json learn_outcome::to_json() const { json res = { { "success", ok }, { "loss", loss }, { "n_labels", n_labels }, { "t_ms", t_ms } }; + if (ok) { + res["step"] = step; + } if (deferred) { res["deferred"] = true; res["retry_after_ms"] = retry_after_ms; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 3966e6a7742d..584551c7801a 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -60,6 +60,7 @@ enum server_task_type { SERVER_TASK_TYPE_SET_LORA, SERVER_TASK_TYPE_LOAD_LORA, SERVER_TASK_TYPE_UNLOAD_LORA, + SERVER_TASK_TYPE_SNAPSHOT_LORA, SERVER_TASK_TYPE_LEARN, }; @@ -101,6 +102,8 @@ struct learn_outcome { double loss = 0.0; int32_t n_labels = 0; double t_ms = 0.0; + // the shadow's committed optimizer steps after this one (0 when it did not step) + int64_t step = 0; json to_json() const; }; @@ -264,6 +267,10 @@ struct server_task { // used by SERVER_TASK_TYPE_UNLOAD_LORA: the adapter id to retire int32_t unload_lora = -1; + // used by SERVER_TASK_TYPE_SNAPSHOT_LORA: the adapter id to write, and where (confined to --train-dir) + int32_t snapshot_lora = -1; + std::string snapshot_path; + // used by SERVER_TASK_TYPE_LEARN learn_request learn; @@ -667,6 +674,17 @@ struct server_task_result_lora_change : server_task_result { virtual json to_json() override; }; +// A shadow written to an immutable file between steps: what she had learned at `steps`. +struct server_task_result_lora_snapshot : server_task_result { + int32_t lora_id = -1; + std::string path; + int64_t bytes = 0; + int64_t steps = 0; + virtual json to_json() override { + return { { "success", true }, { "id", lora_id }, { "path", path }, { "bytes", bytes }, { "steps", steps } }; + } +}; + struct server_task_result_learn : server_task_result { learn_outcome outcome; virtual json to_json() override { return outcome.to_json(); } diff --git a/tools/server/server.cpp b/tools/server/server.cpp index f449dc00f11e..1fb74cb30e2e 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -219,6 +219,7 @@ int llama_server(common_params & params, int argc, char ** argv) { routes.post_lora_adapters = models_routes->proxy_post; routes.post_lora_adapters_load = models_routes->proxy_post; routes.post_lora_adapters_unload = models_routes->proxy_post; + routes.post_lora_adapters_snapshot = models_routes->proxy_post; routes.post_learn = models_routes->proxy_post; routes.get_slots = models_routes->proxy_get; routes.post_slots = models_routes->proxy_post; @@ -274,6 +275,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.post("/lora-adapters", ex_wrapper(routes.post_lora_adapters)); ctx_http.post("/lora-adapters/load", ex_wrapper(routes.post_lora_adapters_load)); ctx_http.post("/lora-adapters/unload", ex_wrapper(routes.post_lora_adapters_unload)); + ctx_http.post("/lora-adapters/snapshot", ex_wrapper(routes.post_lora_adapters_snapshot)); ctx_http.post("/learn", ex_wrapper(routes.post_learn)); ctx_http.get ("/train", ex_wrapper(routes.get_train)); ctx_http.post("/train", ex_wrapper(routes.post_train)); diff --git a/tools/server/tests/learn-every-format.py b/tools/server/tests/learn-every-format.py index 6523b9256c03..b0fdd9898099 100755 --- a/tools/server/tests/learn-every-format.py +++ b/tools/server/tests/learn-every-format.py @@ -37,6 +37,29 @@ def raw(path, body): print(json.dumps({"unbound_retired": retired, "bound_retired": bound, "next_learn": following["learn"]})) sys.exit(0) +# A snapshot is what she had learned, immutable: it carries the shadow's weights at its committed step +# (fresh server; writes snap-test.gguf into --train-dir, so remove it before a rerun). +if "--snapshot" in sys.argv: + import urllib.error + turn = {"prompt": "The capital of Kansas is", "id_slot": 0, "n_predict": 1, "temperature": 0, "learn": learn} + steps = [json.loads(raw("/completion", turn))["learn"] for _ in range(2)] + assert [st["success"] for st in steps] == [True, True] and [st["step"] for st in steps] == [1, 2], steps + snap = json.loads(raw("/lora-adapters/snapshot", {"id": sid, "name": "snap-test.gguf"})) + assert snap["success"] and snap["steps"] == 2 and snap["bytes"] > 0, snap + try: + raw("/lora-adapters/snapshot", {"id": sid, "name": "snap-test.gguf"}) + raise AssertionError("a snapshot was overwritten") + except urllib.error.HTTPError as error: + detail = error.read().decode() + assert error.code == 400 and "immutable" in detail, (error.code, detail) + third = json.loads(raw("/completion", turn))["learn"] + # a learn's loss is its forward pass before its update: the snapshot, loaded, sees what the shadow saw + loaded = json.loads(raw("/lora-adapters/load", {"name": "snap-test.gguf"}))["id"] + again = json.loads(raw("/completion", dict(turn, learn=dict(learn, lora_id=loaded))))["learn"] + assert again["success"] and abs(again["loss"] - third["loss"]) < 1e-5, (third, again) + print(json.dumps({"steps": steps, "snapshot": snap, "third": third, "from_snapshot": again})) + 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: