From bdd3cd55be4965ae870dac4d9e35646070c71344 Mon Sep 17 00:00:00 2001 From: Joel Teply Date: Tue, 13 Oct 2026 15:37:51 -0500 Subject: [PATCH] train: one rule for the recurrent states a graph keeps, so training runs in a context with rollback slots Serving keeps recurrent rollback slots for drafts (n_rs_seq > 0); the exact walk refused any context that had them, which is what forced a second, rollback-free training context beside serving. llm_graph_context::rs_keep() is now the one rule every family's state write and every snapshot-producing op reads (delta-net, mamba, lfm2, bailing; each decided it separately before): serving with rollback keeps n_rs_seq + 1 states, a training graph keeps one, the live state in group 0, which is the same bytes in both layouts. The walk now refuses only a sequence with an undo PENDING (llama_memory_recurrent::rollback_pending), and refuses by name a model whose state write does not report its exit to the walk, where the gradient across chunks was silently dropped before. Measured on the M5 (Metal), base 178224903: - serving unchanged: 19 recurrent/hybrid archs x n_rs_seq {0, 4}, random weights, 100-token prefill + 4 single decodes + (where rollback is live: qwen35, qwen35moe, deepseek4, nemotron_h, nemotron_h_moe, bailingmoe3, lfm2, lfm2moe) a 2-token undo and redo: logits and full sequence state byte-identical to the base, 38 of 38 rows. - training: test-walk-exact on qwen3.5-4b Q4_K_M, the exact walk in a context with n_rs_seq 4 vs 0: cosine 1.0000, norm ratio 1.0000; every other check in the file unchanged (exit 0). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo --- src/llama-context-opt.cpp | 21 ++++++++++++++------ src/llama-graph.cpp | 5 +++++ src/llama-graph.h | 9 +++++++++ src/llama-memory-recurrent.h | 6 ++++++ src/models/bailingmoe3.cpp | 9 ++++----- src/models/delta-net-base.cpp | 10 +++++----- src/models/lfm2.cpp | 2 +- src/models/mamba-base.cpp | 2 +- tests/test-walk-exact.cpp | 37 +++++++++++------------------------ 9 files changed, 57 insertions(+), 44 deletions(-) diff --git a/src/llama-context-opt.cpp b/src/llama-context-opt.cpp index 3c183f3d255c..6cdd89002c4e 100644 --- a/src/llama-context-opt.cpp +++ b/src/llama-context-opt.cpp @@ -317,6 +317,14 @@ bool llama_context::opt_walk::train_chunk(uint32_t pos_ctx, uint32_t n_chunk) { } } for (const auto & ws : res->t_walk_state) { + // a state read without its write reporting the exit (build_walk_state_exit) would drop + // the gradient later chunks put on it: a wrong step, silently. Refused by name. + if (ws.tap && ws.gs == nullptr && lctx.cparams.walk_state_surrogate) { + ggml_free(ctx_compute_opt); + lctx.opt_fail(std::string("this model's recurrent state write (") + ws.cache->name + + ") does not report its exit to the walk (llm_graph_context::build_walk_state_exit): its gradient across chunks would be dropped"); + return false; + } if (ws.tap) { ggml_opt_accumulate_grad(opt_ctx, ws.tap, acc.state.at(ws.cache)); } @@ -406,12 +414,13 @@ void llama_context::opt_walk::run_exact() { // Each window starts from the REQUESTED horizon: a halving that fit one window's device budget // must not shrink every later window and epoch, which may be shorter (Fable). lctx.opt_walk_horizon = lctx.opt_walk_horizon_req; - // A recurrent context with rollback slots keeps several states per sequence; the walk carries - // exactly one. Refused by name, never asserted: training runs inside the serving process, and - // an assert there takes serving down (2026-10-07 04:09Z on the 5090). - if (recr != nullptr && lctx.cparams.n_rs_seq > 0) { - lctx.opt_fail("the exact walk needs a training context with no recurrent rollback slots (n_rs_seq = 0, it has " - + std::to_string(lctx.cparams.n_rs_seq) + "): create the training context with n_rs_seq = 0"); + // A serving context keeps rollback slots for drafts; the walk's graphs keep one state, the live + // one (llm_graph_context::rs_keep), so they train in that context. Only a PENDING undo reads a + // snapshot instead of the live state: that alone is refused, by name, never asserted, since + // training runs inside the serving process (2026-10-07 04:09Z on the 5090). + if (recr != nullptr && recr->rollback_pending(0)) { + lctx.opt_fail("the exact walk cannot train a sequence with a pending recurrent rollback (n_rs_seq = " + + std::to_string(lctx.cparams.n_rs_seq) + "): finish or drop the draft first"); return; } plan.mark_exact(recr != nullptr, lctx.opt_walk_horizon); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9e9071876ab9..6145d5ba3643 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3690,6 +3690,11 @@ ggml_tensor * llm_graph_context::build_rs( get_state_rows); } +int64_t llm_graph_context::rs_keep() const { + const bool training = cparams.training && cparams.walk_exact; + return cparams.n_rs_seq > 0 && !training ? (int64_t) cparams.n_rs_seq + 1 : 1; +} + void llm_graph_context::build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const { if (!(cparams.training && cparams.walk_exact && cparams.walk_state_surrogate)) { return; diff --git a/src/llama-graph.h b/src/llama-graph.h index 0e0adb7d4bda..3b181afed3fb 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1367,6 +1367,15 @@ struct llm_graph_context { int32_t n_seqs, const llm_graph_get_rows_fn & get_state_rows = ggml_get_rows) const; + // THE ONE RULE for how many recurrent states a graph keeps per sequence, read by every family's + // state write (delta-net, mamba, lfm2, bailing) and by the ops that produce snapshots. Serving with + // rollback slots keeps K = n_rs_seq + 1 (group g = the state g tokens back, so a draft can be undone); + // a context without them, and a TRAINING graph, keep one: the live state, group 0, which is the same + // bytes in both layouts. So training reads and writes the state serving reads, in the serving + // context; the undo groups it leaves stale are rewritten by the next serving decode, and the walk + // refuses a sequence with an undo pending (llama_memory_recurrent::rollback_pending). + int64_t rs_keep() const; + // THE EXACT WALK: the state a chunk leaves (`exit`, what is written back to `cache`) meets the // gradient later chunks put on it, through the surrogate . A no-op outside the walk. void build_walk_state_exit(ggml_tensor * cache, ggml_tensor * exit) const; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5e..8ef7420a987a 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -78,6 +78,12 @@ class llama_memory_recurrent : public llama_memory_i { void set_rs_idx(llama_seq_id seq_id, uint32_t idx); + // an undo of a draft is pending on this sequence: its next read takes a snapshot group, not the + // live state (group 0) a training graph reads and writes (llm_graph_context::rs_keep) + bool rollback_pending(llama_seq_id seq_id) const { + return n_rs_seq > 0 && seq_id >= 0 && (size_t) seq_id < rs_idx.size() && rs_idx[seq_id] != 0; + } + // computed before each graph build uint32_t n = 0; diff --git a/src/models/bailingmoe3.cpp b/src/models/bailingmoe3.cpp index 0637931cc0c9..e794c0914afc 100644 --- a/src/models/bailingmoe3.cpp +++ b/src/models/bailingmoe3.cpp @@ -183,7 +183,7 @@ static ggml_tensor * bailingmoe3_causal_conv1d( int64_t n_tokens, int64_t cache_head, uint32_t mem_size, - uint32_t n_rs_seq) { + int64_t K) { // states kept per sequence: llm_graph_context::rs_keep const int64_t d_inner = head_dim * n_head; const int64_t conv_state_size = (d_conv - 1) * d_inner; const int64_t total_state_size = 3 * conv_state_size; @@ -197,7 +197,6 @@ static ggml_tensor * bailingmoe3_causal_conv1d( x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0); - const int64_t K = (int64_t) n_rs_seq + 1; const int64_t n_written = std::min(n_seq_tokens, K); for (int64_t slot = 0; slot < n_written; ++slot) { @@ -262,13 +261,13 @@ llama_model_bailingmoe3::graph::graph(const llama_model & model, const llm_graph ggml_tensor * q = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * k = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * v = bailingmoe3_causal_conv1d( gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, - d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, rs_keep()); ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); gate = ggml_add(ctx0, gate, layer.ssm_dt_b); diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index e0a6368bd2ed..f613db67f5b9 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -476,7 +476,7 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( const size_t row_size = ggml_row_size(conv_states_all->type, row_count); - if (cparams.n_rs_seq == 0) { + if (rs_keep() == 1) { const int64_t s_idx = conv_input->ne[0] - conv_states->ne[0]; const int64_t s_slot = 0; @@ -496,12 +496,12 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_state_last, conv_state_update)); build_walk_state_exit(conv_states_all, conv_state_last); } else { - GGML_ASSERT(!cparams.walk_exact && "the exact walk keeps one recurrent state per sequence (n_rs_seq = 0)"); + GGML_ASSERT(!cparams.walk_exact && "a training graph keeps one recurrent state per sequence (rs_keep)"); // [TAG_RECURRENT_ROLLBACK_SPLITS] // this logic assumes that the last (n_rs_seq + 1) tokens of a sequence in a batch are inside // the same ubatch, which `split_equal()` guarantees via its n_keep_tail argument - const int64_t K = (int64_t) cparams.n_rs_seq + 1; + const int64_t K = rs_keep(); for (int64_t t = 1; t <= K; ++t) { const int64_t s_idx = std::max(0, conv_input->ne[0] - conv_states->ne[0] - K + t); @@ -545,7 +545,7 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( const int64_t n_seqs = s->ne[3]; const int64_t n_seq_tokens = q->ne[2]; - const bool keep = cparams.n_rs_seq > 0; + const bool keep = rs_keep() > 1; if (!keep) { auto attn_out = build_delta_net(q, k, v, g, b, s, il); @@ -564,7 +564,7 @@ ggml_tensor * llm_build_delta_net_base::build_recurrent_attn( } const int64_t D = S_v * S_v * H_v; - const int64_t K = cparams.n_rs_seq + 1; + const int64_t K = rs_keep(); // state s is 4D [S_v, S_v, H_v, n_seqs]; K snapshot slots are written into the output. ggml_tensor * gdn_out = ggml_gated_delta_net(ctx0, q, k, v, g, b, s, K); diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp index 9a4295557053..d2e954844608 100644 --- a/src/models/lfm2.cpp +++ b/src/models/lfm2.cpp @@ -205,7 +205,7 @@ llama_model_lfm2::graph::graph(const llama_model & model, const llm_graph_ GGML_ASSERT(bx->ne[0] > conv->ne[0]); // write conv states: slot 0 = the final state, slot s = the state s tokens back (partial rollback) - const int64_t K = hparams.causal_attn && cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; + const int64_t K = hparams.causal_attn ? rs_keep() : 1; const int64_t n_written = std::min(n_seq_tokens, K); const auto mem_size = mctx_cur->get_size(); const size_t row_size = ggml_row_size(conv_state->type, (int64_t) d_conv * n_embd); diff --git a/src/models/mamba-base.cpp b/src/models/mamba-base.cpp index 03ee3805bf80..af8fd15198a8 100644 --- a/src/models/mamba-base.cpp +++ b/src/models/mamba-base.cpp @@ -167,7 +167,7 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, const int64_t n_seqs = ubatch.n_seqs; const int64_t n_seq_tokens = ubatch.n_seq_tokens; - const int64_t K = cparams.n_rs_seq > 0 ? (int64_t) cparams.n_rs_seq + 1 : 1; + const int64_t K = rs_keep(); GGML_ASSERT(n_seqs != 0); GGML_ASSERT(ubatch.equal_seqs()); diff --git a/tests/test-walk-exact.cpp b/tests/test-walk-exact.cpp index 7e4dd1717a6e..b46144a1e0d0 100644 --- a/tests/test-walk-exact.cpp +++ b/tests/test-walk-exact.cpp @@ -463,36 +463,21 @@ int main(int argc, char ** argv) { } } if (llama_model_is_hybrid(model)) { - // 4. regression for the 5090 crash (2026-10-07 04:09Z): a recurrent context with rollback - // slots (serving's, inherited by the training context) made the exact walk ASSERT inside - // the serving process. It must refuse by name, and the adapter stays untouched. - llama_adapter_lora * adapter = llama_adapter_lora_init(model, init.c_str()); - const std::vector before = adapter_params(adapter); + // 4. TRAINING IN THE SERVING CONTEXT (Joel, 2026-10-13: "training is basically llama inference with + // backprop"). Serving keeps recurrent rollback slots for drafts; a training graph keeps one state, the + // live one (llm_graph_context::rs_keep), so a context WITH rollback slots takes the same step as one + // without. Was a refusal (the 5090 crash, 2026-10-07 04:09Z), which forced a second training + // context beside serving. + std::vector all(WINDOW + 1, 1); + const auto without = step_delta(params, model, init, tokens, all, WINDOW / 4, true); g_n_rs_seq = 4; - llama_context * ctx = make_ctx(params, model, WINDOW / 4); + const auto with = step_delta(params, model, init, tokens, all, WINDOW / 4, true); g_n_rs_seq = 0; - float scale = 1.0f; - GGML_ASSERT(llama_set_adapters_lora(ctx, &adapter, 1, &scale) == 0); - llama_opt_params lopt{}; - lopt.param_filter = llama_opt_param_filter_all; - lopt.get_opt_pars = sgd_pars; - lopt.optimizer_type = GGML_OPT_OPTIMIZER_TYPE_SGD; - lopt.adapter = adapter; - lopt.walk_exact = true; - llama_opt_init(ctx, model, lopt); - std::vector> seqs = { std::vector(tokens.begin(), tokens.begin() + WINDOW + 1) }; - std::vector> loss = { std::vector(WINDOW + 1, 1) }; - ggml_opt_dataset_t dataset = common_opt_dataset_init_masked(WINDOW, seqs, loss, tokens[0]); - llama_opt_epoch(ctx, dataset, nullptr, nullptr, /*idata_split =*/ 1, nullptr, nullptr); - const std::string why = llama_opt_failure(ctx); - printf(" rollback slots: failed=%d, \"%s\"\n", (int) llama_opt_failed(ctx), why.c_str()); - if (!llama_opt_failed(ctx) || why.find("rollback") == std::string::npos || adapter_params(adapter) != before) { - fprintf(stderr, "FAILED: a recurrent context with rollback slots was not refused by name\n"); + if (!same_step(compare("rollback slots: exact walk in a context with n_rs_seq 4 vs 0", with, without))) { + fprintf(stderr, "FAILED: rollback slots changed the exact walk's step\n"); + worst_tensors(with, without, 24); ++failures; } - ggml_opt_dataset_free(dataset); - llama_free(ctx); - llama_adapter_lora_free(adapter); } {