ggml-opt: freeing an optimizer checks only its own graph, so an idle shadow unloads while another mind learns

ggml/include/ggml-backend.h

@@ -388,6 +388,9 @@ extern "C" {
     GGML_API enum ggml_status ggml_backend_sched_graph_step(ggml_backend_sched_t sched,
             struct ggml_cgraph * graph, int max_nodes, bool * complete);
     GGML_API bool ggml_backend_sched_graph_pending(ggml_backend_sched_t sched);
+    // True when `graph` itself is the one the scheduler is part-way through (graph_step): a caller
+    // that owns `graph` must not free its tensors yet. Another caller's pending graph is not this one.
+    GGML_API bool ggml_backend_sched_graph_pending_for(ggml_backend_sched_t sched, const struct ggml_cgraph * graph);
     GGML_API void ggml_backend_sched_graph_discard(ggml_backend_sched_t sched);
 
     GGML_API void                 ggml_backend_sched_reset(ggml_backend_sched_t sched);

ggml/src/ggml-backend.cpp

@@ -2711,6 +2711,10 @@ bool ggml_backend_sched_graph_pending(ggml_backend_sched_t sched) {
     return sched->continuation_graph != nullptr;
 }
 
+bool ggml_backend_sched_graph_pending_for(ggml_backend_sched_t sched, const struct ggml_cgraph * graph) {
+    return graph != nullptr && sched->continuation_graph == graph;
+}
+
 void ggml_backend_sched_graph_discard(ggml_backend_sched_t sched) {
     ggml_backend_sched_synchronize(sched);
     sched->continuation_graph = nullptr;

ggml/src/ggml-opt.cpp

@@ -966,7 +966,11 @@ void ggml_opt_free(ggml_opt_context_t opt_ctx) {
     if (opt_ctx == nullptr) {
         return;
     }
-    GGML_ASSERT(!opt_ctx->allocation_pending && !ggml_backend_sched_graph_pending(opt_ctx->backend_sched));
+    // Freeing releases only this context's own buffers and metadata, never the shared scheduler,
+    // so another context's graph may be mid-step (an idle mind unloads while another learns,
+    // llama_context::opt_release). This context's OWN graph must not be the one in flight.
+    GGML_ASSERT(!opt_ctx->allocation_pending &&
+                !ggml_backend_sched_graph_pending_for(opt_ctx->backend_sched, opt_ctx->allocated_graph_copy));
     ggml_backend_buffer_free(opt_ctx->buf_static);
     ggml_backend_buffer_free(opt_ctx->buf_cpu);
     ggml_free(opt_ctx->ctx_static);

tests/test-opt-dynamic-accum.cpp

@@ -15,6 +15,11 @@
 #include <cstring>
 #include <initializer_list>
 #include <vector>
+#if defined(__unix__) || defined(__APPLE__)
+#include <csignal>
+#include <sys/wait.h>
+#include <unistd.h>
+#endif
 
 #define CHECK(cond) do { if (!(cond)) { fprintf(stderr, "FAILED %s:%d: %s\n", __FILE__, __LINE__, #cond); exit(1); } } while (0)
 
@@ -689,6 +694,136 @@ static void run_refuses_misshaped_gradient() {
     ggml_free(ctx);
 }
 
+// One mind's trainable weight and optimizer on a SHARED scheduler, the way a serving context keeps
+// one optimizer per shadow (llama_context::opt_minds).
+struct mind {
+    ggml_context       * ctx_static = nullptr;
+    ggml_tensor        * w          = nullptr;
+    ggml_backend_buffer_t buf       = nullptr;
+    ggml_opt_context_t   opt        = nullptr;
+    ggml_opt_result_t    result     = nullptr;
+    ggml_context       * ctx_compute = nullptr;
+};
+
+static mind mind_on(ggml_backend_sched_t sched, ggml_backend_t backend) {
+    mind m;
+    ggml_init_params sp = { 8 * ggml_tensor_overhead(), nullptr, true };
+    m.ctx_static = ggml_init(sp);
+    m.w = ggml_new_tensor_1d(m.ctx_static, GGML_TYPE_F32, 1);
+    ggml_set_name(m.w, "weight");
+    ggml_set_param(m.w);
+    m.buf = ggml_backend_alloc_ctx_tensors(m.ctx_static, backend);
+    const float one = 1.0f;
+    ggml_backend_tensor_set(m.w, &one, 0, sizeof(float));
+    ggml_opt_params params = ggml_opt_default_params(sched, GGML_OPT_LOSS_TYPE_SUM);
+    params.optimizer    = GGML_OPT_OPTIMIZER_TYPE_SGD;
+    params.get_opt_pars = sgd_pars;
+    m.opt    = ggml_opt_init(params);
+    m.result = ggml_opt_result_init();
+    return m;
+}
+
+// Prepare and allocate this mind's next step, then run one graph node of it: the step is left
+// part-way, the shared scheduler pending on this mind's graph. Returns whether it completed.
+static bool begin_step(mind & m) {
+    ggml_init_params cp = { GGML_DEFAULT_GRAPH_SIZE * ggml_tensor_overhead() + 4 * ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, true), nullptr, true };
+    m.ctx_compute = ggml_init(cp);
+    ggml_tensor * x   = ggml_new_tensor_1d(m.ctx_compute, GGML_TYPE_F32, 1);
+    ggml_tensor * out = ggml_mul(m.ctx_compute, m.w, x);
+    ggml_cgraph * gf  = ggml_new_graph_custom(m.ctx_compute, GGML_DEFAULT_GRAPH_SIZE, true);
+    ggml_build_forward_expand(gf, out);
+    ggml_opt_prepare_alloc(m.opt, m.ctx_compute, gf, x, out);
+    CHECK(ggml_opt_alloc(m.opt, true));
+    ggml_backend_tensor_set(x, &X, 0, sizeof(float));
+    bool complete = false;
+    CHECK(ggml_opt_eval_step(m.opt, m.result, 1, &complete) == GGML_STATUS_SUCCESS);
+    return complete;
+}
+
+static void finish_step(mind & m) {
+    bool complete = false;
+    int calls = 0;
+    do {
+        CHECK(++calls < 1000);
+        CHECK(ggml_opt_eval_step(m.opt, m.result, 1, &complete) == GGML_STATUS_SUCCESS);
+    } while (!complete);
+    ggml_free(m.ctx_compute);
+    m.ctx_compute = nullptr;
+}
+
+static float weight(const mind & m) {
+    float w = 0;
+    ggml_backend_tensor_get(m.w, &w, 0, sizeof(w));
+    return w;
+}
+
+static void release(mind & m) {
+    if (m.opt) { ggml_opt_free(m.opt); }
+    ggml_opt_result_free(m.result);
+    ggml_backend_buffer_free(m.buf);
+    ggml_free(m.ctx_static);
+}
+
+// what this catches (fork #106 reproducer, learn-unload-during-learn.py): freeing an IDLE mind's
+// optimizer while another mind's graph is part-way through the shared scheduler aborted the engine,
+// because ggml_opt_free asserted the whole scheduler had no pending graph. Freeing touches only the
+// idle mind's own buffers, so the other mind's step must finish with exactly the update it would
+// make alone. Built with -fsanitize=address, a freed buffer the live graph still read fails loudly.
+static void run_frees_idle_mind_while_another_steps() {
+    ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
+    GGML_ASSERT(cpu != nullptr);
+    ggml_backend_t backends[] = { cpu };
+    ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 1, GGML_DEFAULT_GRAPH_SIZE, false, true);
+    mind a = mind_on(sched, cpu);
+    mind b = mind_on(sched, cpu);
+    // B learns a whole step first, so it holds optimizer buffers and a graph copy of its own
+    if (!begin_step(b)) { finish_step(b); } else { ggml_free(b.ctx_compute); b.ctx_compute = nullptr; }
+    CHECK(std::fabs(weight(b) - (1.0f - LR * X)) < 1e-6f);
+    CHECK(!ggml_backend_sched_graph_pending(sched));
+    // A starts and stops part-way: the scheduler is now pending on A's graph, not B's
+    CHECK(!begin_step(a));
+    CHECK(ggml_backend_sched_graph_pending(sched));
+    CHECK(!ggml_backend_sched_graph_pending_for(sched, nullptr));
+    ggml_opt_free(b.opt); // the idle mind unloads mid-way through A's step
+    b.opt = nullptr;
+    CHECK(ggml_backend_sched_graph_pending(sched));
+    finish_step(a);
+    CHECK(!ggml_backend_sched_graph_pending(sched));
+    CHECK(std::fabs(weight(a) - (1.0f - LR * X)) < 1e-6f);
+    release(a);
+    release(b);
+    ggml_backend_sched_free(sched);
+    ggml_backend_free(cpu);
+    printf("PASS an idle mind's optimizer frees while another mind's step is in flight\n");
+}
+
+// what this catches: the narrowed guard widening into no guard. A context whose OWN graph is
+// part-way through the scheduler must still refuse to be freed (it would free tensors the
+// scheduler reads next), so ggml_opt_free must abort. Observed in a child process.
+static void run_free_own_pending_graph_aborts() {
+#if defined(__unix__) || defined(__APPLE__)
+    fflush(stdout);
+    fflush(stderr);
+    const pid_t child = fork();
+    CHECK(child >= 0);
+    if (child == 0) {
+        ggml_backend_t cpu = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr);
+        ggml_backend_t backends[] = { cpu };
+        ggml_backend_sched_t sched = ggml_backend_sched_new(backends, nullptr, 1, GGML_DEFAULT_GRAPH_SIZE, false, true);
+        mind a = mind_on(sched, cpu);
+        if (begin_step(a)) { _exit(2); } // the step must be left part-way for this case to mean anything
+        ggml_opt_free(a.opt);            // must abort: A's own graph is in flight
+        _exit(0);
+    }
+    int status = 0;
+    CHECK(waitpid(child, &status, 0) == child);
+    CHECK(WIFSIGNALED(status) && WTERMSIG(status) == SIGABRT);
+    printf("PASS freeing a context whose own graph is in flight still aborts\n");
+#else
+    printf("SKIP freeing a context whose own graph is in flight: needs fork() to observe the abort\n");
+#endif
+}
+
 int main() {
     ggml_backend_load_all();
     for (int period : {1, 2, 3}) { run(period, false, false, false, 1); }
@@ -721,6 +856,8 @@ int main() {
         run_out_prod_broadcast(type, false);
         run_out_prod_broadcast(type, true);
     }
+    run_frees_idle_mind_while_another_steps();
+    run_free_own_pending_graph_aborts();
     printf("test-opt-dynamic-accum: OK\n");
     return 0;
 }