diff --git a/examples/training/finetune-lora.cpp b/examples/training/finetune-lora.cpp index 1f97b2bcde03..541aca4fd257 100644 --- a/examples/training/finetune-lora.cpp +++ b/examples/training/finetune-lora.cpp @@ -136,6 +136,11 @@ int main(int argc, char ** argv) { /*recompute =*/false, // keep every layer's intermediates, as before /*walk_exact =*/false, // the per-chunk walk, as before /*walk_horizon =*/0, + // an example trains under no learning residency: no allocation admission callbacks + /*static_admission =*/nullptr, + /*static_admission_data =*/nullptr, + /*workspace_admission =*/nullptr, + /*workspace_admission_data =*/nullptr, }; llama_opt_init(ctx, model, lopt_params); diff --git a/examples/training/finetune.cpp b/examples/training/finetune.cpp index c19e5f82dd19..2210777a365b 100644 --- a/examples/training/finetune.cpp +++ b/examples/training/finetune.cpp @@ -77,6 +77,11 @@ int main(int argc, char ** argv) { /*recompute =*/false, // keep every layer's intermediates, as before /*walk_exact =*/false, // the per-chunk walk, as before /*walk_horizon =*/0, + // an example trains under no learning residency: no allocation admission callbacks + /*static_admission =*/nullptr, + /*static_admission_data =*/nullptr, + /*workspace_admission =*/nullptr, + /*workspace_admission_data =*/nullptr, }; llama_opt_init(ctx, model, lopt_params); diff --git a/src/llama-context-opt-measure.cpp b/src/llama-context-opt-measure.cpp index b2ce0ac10a6a..9bf7e9d415bd 100644 --- a/src/llama-context-opt-measure.cpp +++ b/src/llama-context-opt-measure.cpp @@ -37,7 +37,12 @@ bool llama_context::opt_measure_memory(llama_opt_memory_callback callback, void cparams.training = true; cparams.walk_active = false; // Shape only: attention still sees the full served memory, without token cuts. - const uint32_t chunk = reply_rows ? reply_rows : cparams.n_ctx_seq; + // A learn step is one decode ubatch, as serving's is (opt_learn_step decodes the + // reply in n_batch chunks that decode() splits by n_ubatch), so its workspace is + // priced at one ubatch, whatever the reply's length; only the whole-reply graph + // prices its rows. + const uint32_t rows = reply_rows ? reply_rows : cparams.n_ctx_seq; + const uint32_t chunk = phases == opt_measure_phase::whole_reply_commit ? rows : std::min(rows, cparams.n_ubatch); size_t max_nodes = graph_max_nodes(chunk); if (cparams.flash_attn && !cparams.fa_back_native) { max_nodes += (size_t) model.hparams.n_layer() * ggml_flash_attn_ext_back_n_nodes(opt_n_ctx_train, 0); diff --git a/tests/learn-allocation-residency.py b/tests/learn-allocation-residency.py index ace802f6ebae..5d99b4de7868 100644 --- a/tests/learn-allocation-residency.py +++ b/tests/learn-allocation-residency.py @@ -154,22 +154,25 @@ def learn(shadow, optimizer='', workspace=''): assert call('/lora-adapters/unload', {'id': b['id']})['retired'] assert call('/learn/residency/release', {'id': 'opt-b'})['refusal'] == 'none' assert call('/learn/residency/release', {'id': 'workspace-next'})['refusal'] == 'busy' -# A larger actual reply must re-admit growth; the short plan is not a window cap. +# A learn step is one decode chunk, as serving's is: a longer reply learns inside the +# workspace a short plan priced (no growth), and the plan does not depend on the reply. long_reply = 'Topeka is the capital of Kansas. ' * 50 -refused_growth = call('/completion', dict(probe, learn={'grant': 'compute', 'lora_id': a['id'], - 'reply': long_reply, 'optimizer_residency': 'opt-a', 'workspace_residency': 'workspace-next'}))['learn'] -assert not refused_growth['success'] and refused_growth['committed_steps'] == 0, refused_growth long_plan = call('/learn/preflight', {'engine': engine, 'lora_id': a['id'], 'learn': {'reply': long_reply}}) workspace = [row for row in long_plan['requirements'] if row['kind'] == 'shared_compute_required'] register('workspace-long', {'workspace': True}, workspace) +grown = call('/completion', dict(probe, learn={'grant': 'compute', 'lora_id': a['id'], + 'reply': long_reply, 'optimizer_residency': 'opt-a', 'workspace_residency': 'workspace-next'}))['learn'] +assert grown['success'] and grown['committed_steps'] >= 1, grown +# A reply several chunks long is revoked between chunks: no gradient chunk after the +# revocation, the stable tail restored, and nothing half-applied. +many_chunks = 'Topeka is the capital of Kansas. ' * 400 with concurrent.futures.ThreadPoolExecutor() as pool: future = pool.submit(call, '/completion', dict(probe, learn={'grant': 'compute', 'lora_id': a['id'], - 'reply': 'Topeka is the capital of Kansas. ' * 50, - 'optimizer_residency': 'opt-a', 'workspace_residency': 'workspace-long'})) + 'reply': many_chunks, 'optimizer_residency': 'opt-a', 'workspace_residency': 'workspace-long'})) deadline = time.monotonic() + 30 observed = False while time.monotonic() < deadline and not future.done(): - if any(slot.get('is_learning') and slot.get('learn_graph_steps', 0) >= 2 for slot in call('/slots')): + if any(slot.get('is_learning') and slot.get('learn_gradient_steps', 0) >= 2 for slot in call('/slots')): observed = True break time.sleep(0.005) @@ -183,9 +186,8 @@ def learn(shadow, optimizer='', workspace=''): assert call('/learn/residency/revoke', {'id': 'opt-a'})['refusal'] == 'none' assert call('/learn/residency/release', {'id': 'opt-a'})['refusal'] == 'busy' outcome = future.result()['learn'] - assert not outcome['success'] and outcome['n_labels'] > 128, outcome - assert outcome['committed_steps'] == 1 and not outcome['effects_uncertain'], outcome - print('revoked graph completed before restoration', outcome, flush=True) + assert not outcome['success'] and not outcome['effects_uncertain'], outcome + print('revoked between chunks, stable tail restored', outcome, flush=True) assert outcome['refusal'] == 'residency_revoked', outcome assert not learn(a, 'opt-a', 'workspace-long')['success'] assert call('/learn/residency/release', {'id': 'opt-a'})['refusal'] == 'busy' diff --git a/tools/server/server-context-learn-checkpoint.inc b/tools/server/server-context-learn-checkpoint.inc index d7bfa4e3f480..ca5645a7ec7b 100644 --- a/tools/server/server-context-learn-checkpoint.inc +++ b/tools/server/server-context-learn-checkpoint.inc @@ -42,8 +42,11 @@ gguf_get_val_i64(metadata.get(), iteration) <= 0 || gguf_get_val_i64(metadata.get(), iteration) - 1 != steps || std::string(gguf_get_val_str(metadata.get(), bound)) != binding) { return refuse("optimizer generation mismatch"); } - std::unique_ptr weights(fopen((path + "/adapter.gguf").c_str(), "rb"), fclose); - std::unique_ptr optimizer(fopen((path + "/optimizer.gguf").c_str(), "rb"), fclose); + // a deleter type, not decltype(&fclose): glibc declares fclose with attributes a template + // argument drops, which gcc -Werror=ignored-attributes refuses + struct file_closer { void operator()(FILE * f) const { if (f) { fclose(f); } } }; + std::unique_ptr weights(fopen((path + "/adapter.gguf").c_str(), "rb")); + std::unique_ptr optimizer(fopen((path + "/optimizer.gguf").c_str(), "rb")); if (!weights || !optimizer) { return refuse("checkpoint components unreadable"); } const auto started = ggml_time_us(); auto held = learn_ledger.hold_restoration(request.grant, started); diff --git a/tools/server/server-context-learn.inc b/tools/server/server-context-learn.inc index 2b2a0aa08099..48cdf33ecb52 100644 --- a/tools/server/server-context-learn.inc +++ b/tools/server/server-context-learn.inc @@ -110,7 +110,7 @@ auto set = construct_lora_list({}); set[id].scale = 1.0f; common_set_adapter_lora(ctx_tgt, set); - const bool ok = llama_opt_measure_shadow_graph(ctx_tgt, learning_shadow_params((int32_t) id), (uint32_t) reply_rows, [](const llama_adapter_lora *, const char * device, + const bool ok = llama_opt_measure_shadow(ctx_tgt, learning_shadow_params((int32_t) id), [](const llama_adapter_lora *, const char * device, const char * buft, const char * kind, size_t bytes, bool host, void * data) { static_cast(data)->push_back({{"device", device}, {"buffer_type", buft}, {"kind", kind}, {"bytes", bytes}, {"host", host}}); @@ -392,8 +392,16 @@ const bool begun = llama_opt_learn_begin(ctx_tgt, slot.id, n0, toks.data(), labels.data(), (int32_t) toks.size(), served_tokens.data() + n0, (int32_t) served_tokens.size() - n0, result); common_set_adapter_lora(ctx_tgt, slot.lora); - learning = std::unique_ptr(new pending_learn{ &slot, cp, std::move(set), result, - toks[0], L.n_prefix, (int32_t) toks.size() - 1, loras[L.lora_id].ptr, t0 }); + learning = std::make_unique(); + learning->slot = &slot; + learning->checkpoint = cp; + learning->adapters = std::move(set); + learning->result = result; + learning->last_prompt = toks[0]; + learning->prefix = L.n_prefix; + learning->n_labels = (int32_t) toks.size() - 1; + learning->shadow = loras[L.lora_id].ptr; + learning->started = t0; learning->began = begun; learning->grant = L.grant; learning->lora_id = L.lora_id; @@ -403,7 +411,7 @@ } slot.state = SLOT_STATE_LEARNING; slot.learn_release_requested = false; - slot.learn_graph_steps = 0; + slot.learn_gradient_steps = 0; return {}; // no outcome is published until advance_learning restores the slot } @@ -448,9 +456,11 @@ learning_residencies.at(pending.lora_id)->restoring = !llama_opt_learn_needs_gradient(ctx_tgt) || llama_opt_learn_draining_graph(ctx_tgt); common_set_adapter_lora(ctx_tgt, pending.adapters); const auto steps_before = llama_opt_shadow_steps(ctx_tgt, pending.shadow); - const bool graph_step = llama_opt_learn_needs_gradient(ctx_tgt); - const auto status = llama_opt_learn_graph_step(ctx_tgt); - if (graph_step) { ++slot.learn_graph_steps; } + // One decode chunk, as serving's decode is: other slots decode between steps, + // and the workspace is one ubatch's, never the whole reply's. + const bool gradient_step = llama_opt_learn_needs_gradient(ctx_tgt); + const auto status = llama_opt_learn_step(ctx_tgt); + if (gradient_step) { ++slot.learn_gradient_steps; } const auto committed = llama_opt_shadow_steps(ctx_tgt, pending.shadow) - steps_before; GGML_ASSERT(committed >= 0); // optimizer custody prevents generation replacement here pending.committed_steps += committed; @@ -462,7 +472,7 @@ pending.max_chunk_us = std::max(pending.max_chunk_us, elapsed); // Classify the work, not its charging reservation: a revoked graph // still runs graph nodes under restoration custody until it commits. - auto & phase_max = graph_step ? pending.max_graph_us : pending.max_restore_us; + auto & phase_max = gradient_step ? pending.max_graph_us : pending.max_restore_us; phase_max = std::max(phase_max, elapsed); if (!reservation) { pending.restoration_us += elapsed; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index c51277e13599..092f23248fe8 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -563,7 +563,7 @@ struct server_slot { // the outcome of the finishing task's learn (task_params::learn), for its final response std::optional learned; bool learn_release_requested = false; - int64_t learn_graph_steps = 0; + int64_t learn_gradient_steps = 0; void release() { if (state == SLOT_STATE_LEARNING) { @@ -722,7 +722,7 @@ struct server_slot { {"speculative", can_speculate()}, {"is_processing", is_processing()}, {"is_learning", state == SLOT_STATE_LEARNING}, - {"learn_graph_steps", learn_graph_steps}, + {"learn_gradient_steps", learn_gradient_steps}, }; // The engine's own positions for this slot's sequence (Continuum fork): what its cache holds @@ -996,15 +996,15 @@ struct server_context_impl { return p; } struct pending_learn { - server_slot * slot; - const common_prompt_checkpoint * checkpoint; + server_slot * slot = nullptr; + const common_prompt_checkpoint * checkpoint = nullptr; std::vector adapters; - ggml_opt_result_t result; - llama_token last_prompt; - int32_t prefix, n_labels; + ggml_opt_result_t result = nullptr; + llama_token last_prompt = 0; + int32_t prefix = 0, n_labels = 0; // the shadow this learn trains: held by the optimizer from begin through restoration const llama_adapter_lora * shadow = nullptr; - int64_t started; + int64_t started = 0; int64_t compute_us = 0; bool separate_request = false; bool began = false; diff --git a/tools/server/server-train.cpp b/tools/server/server-train.cpp index 098cc26c2004..fb3abb5166f3 100644 --- a/tools/server/server-train.cpp +++ b/tools/server/server-train.cpp @@ -1055,6 +1055,11 @@ void server_trainer::run(json req, examples_data ex) { // the single-context gradient): one step per window on the window's whole loss /*walk_exact =*/ req.value("exact", false), /*walk_horizon =*/ (uint32_t) req.value("walk_horizon", (int64_t) 0), + // /train allocates under its own memory budget, not a learning residency + /*static_admission =*/ nullptr, + /*static_admission_data =*/ nullptr, + /*workspace_admission =*/ nullptr, + /*workspace_admission_data =*/ nullptr, }; // 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