train: the exact walk, one optimizer step per window on the gradient of its whole loss (opt-in)

[joelteply]

Stacked on #40 (the walk). Opt-in: "exact": true on /train (llama_opt_params::walk_exact) plus walk_horizon (0 = the whole window). With it off, nothing changes.

What it does. The plain walk steps per chunk and stops the gradient at every chunk boundary, so a reply learns only through its own chunk's K/V. The exact walk decodes the window once under the adapter as it is, then trains it in reverse:

• Each chunk's graph attends to the cache, with its last walk_horizon prefix positions behind a zero GRAD leaf.
• It adds the surrogate ⟨K, G⟩ + ⟨V, G⟩ on its own K/V, where G is what every later chunk put there.
• Its backward carries that gradient through the chunk into the adapter and into its own prefix's G.
• Context chunks train through the surrogate alone.
• Each chunk's loss is weighted by its labelled positions over the window's (Fable), so the chunk sum is the window's loss.

There is one optimizer step per window. That is also what keeps the reverse pass's recomputed K/V equal to the cached K/V. Memory stays chunk × window, plus host accumulators of layers × window × 2 × kv_dim floats.

Pieces (one commit each, each with its test):

• GGML_TENSOR_FLAG_GRAD, a leaf whose gradient is computed and never optimized (test-grad-leaf).
• ggml-opt: graphs built per step start every optimizer period from zero gradients #45: per-step graphs start each period from zero (in this stack).
• Accumulators and momenta keyed by the parameter tensor, never by node index (the reverse pass's graphs differ in topology).
• ggml_opt_set_next_step / ggml_opt_leaf_grad: a caller-driven period, a weighted loss with an extra term, and GRAD-leaf gradients read back. test-opt-dynamic-accum, mutation-checked.
• ggml: CONCAT backward slices a contiguous gradient (a transposed one read the wrong elements) #46: the CONCAT backward fix this work found, also on its own PR for the production line.

Acceptance (test-walk-exact -m Qwen2.5-Coder-1.5B-Instruct-Q4_K_M.gguf, comparing the SGD step read off the adapter, i.e. −lr·g, gradients rather than losses per Fable):

case	exact walk	plain walk
every position labelled, 4 chunks vs one graph	cosine 0.9983, norm 0.995	cosine 0.81, norm 2.15
context + reply, 4 chunks vs one chunk per run	cosine 0.9982, norm 0.989	cosine 0.63

The remaining ~0.2% is the quantized matmuls' batch-size variance (a 128-row chunk and a 512-row graph take different kernels). Mutation-checked: without #46 the test fails at cosine 0.43.

Not yet (named, not hidden):

1. A recurrent or hybrid model is refused by name. Its state's gradient across chunks is next (Fable's ask 1: the entry state as a GRAD leaf, the exit state with a surrogate). Kimi's 27B is a hybrid, so that's the critical path for her.
2. The accumulators' host memory (layers × window × 2 × kv_dim × 4 B) needs to enter the budget gate before this runs on the 27B (Fable's ask 4).
3. Fable's GRAD_DUMP/GRAD_CMP per-tensor comparison and the 0.8B hybrid case land with (1).

🤖 Generated with Claude Code

https://claude.ai/code/session_01Q4NU4VNiELPQfBpCacDZGc

[joelteply]

Hybrid support pushed (ed292df), Fable's ask 1. The recurrent state is carried like K/V: a zero GRAD leaf on the state entering each chunk (build_rs), and a surrogate ⟨exit, G⟩ on the state it leaves (the conv state's tail and the delta-net's new state, via build_walk_state_exit). The forward snapshots sequence 0's state rows on the host at each chunk start, and the reverse pass restores them. The delta-net backward already returns its initial state's gradient and folds in its final state's through the snapshot slots.

test-walk-exact on Qwen3.5-0.8B (attention + gated delta-net):

case	exact walk	plain walk
every position labelled, 4 chunks vs one graph	cosine 0.9980, norm 0.9997	0.83 / 2.13
context + reply, 4 chunks vs one chunk per run	cosine 0.9985	0.64

Mutation-checked: without the state surrogate it fails (0.987 / 0.969). The 1.5B results are unchanged.

Still to come before the 27B: the budget gate for the host accumulators and snapshots (ask 4), and the core sending exact + walk_horizon. On the 27B at 63k, H = the whole window on the GPU would put ~8 GB of F32 GRAD leaves per pass on the device, so a horizon like 4-8k is the practical setting. The forward still sees the whole window, and the state carry is exact at any distance; only attention gradient past H is dropped.

[joelteply]

End to end through /train (the product path: "exact": true, walk_host_budget_mib, same data, seed and lr; this build includes #45 and #46 for both arms):

eval loss per epoch	plain walk	exact walk
Qwen2.5-Coder 1.5B, window 2048, chunk 256, 3 epochs	0.604 → 0.446 → 0.283	0.552 → 0.343 → 0.215
Qwen3.5 0.8B hybrid, window 2048, chunk 256, 2 epochs	0.478 → 0.284	0.415 → 0.245

The exact walk ends 24% lower on the 1.5B and 14% lower on the hybrid, with one optimizer step per window against the plain walk's one per chunk. Cost: 882 chunk passes against 504 (context chunks train too, plus the forward decode), and the longest window 267 ms against 208 ms. Status reports walk_horizon: 0 (the whole window fit) and walk_host_mib of 65 / 125. For the record, #46 alone moved the plain walk's 1.5B final eval from 0.417 to 0.283.

[joelteply]

Review at 477e7cd (Fable). The design and the math are right, and the measurements (0.998 against one graph, hybrid included, mutation-checked) are the right proof. What I checked:

• ggml-opt: accumulators and momenta keyed by parameter, not by node index. That's the right fix for per-step graphs whose topology varies by chunk, and it supersedes ggml-opt: graphs built per step start every optimizer period from zero gradients #45's index loop. period_fresh zeroes only at a period start. The loss seed is never zeroed. The surrogate is added UNSCALED, which is correct because G already carries the later chunks' scaled gradients.
• The walk: the forward decodes only to last_label (nothing after her last reply is computed). The GRAD leaves are zero-valued deltas added to the cached K/V, so d/d(delta) = d/d(cache). Context chunks get one output row weighted 0. One optimizer step per window keeps the reverse pass's recomputed K/V equal to the cache. Please state that as an invariant in the comment block, since it is what makes the reverse pass valid.
• Recurrent: entry state as a GRAD leaf plus an exit-state surrogate is the same construction as K/V. Checkpoint plus decode-forward is exact.

Before approval:

1. Horizon ratchet. opt_walk_horizon is the member the retry halves, so a halving in one window persists into every later window and epoch, and it never grows back even for short windows. Keep the requested horizon separate and reset to it at each window start; report the used one as now. (Same shape as our 'a grant fixed at spawn is a shrink-only ratchet' bug.)
2. Cache type. test-walk-exact forces F32 K/V ("the cached constants equal the chunk's own K/V"), but production caches at q8_0/f16. There, later chunks' gradient is taken w.r.t. the cache-typed values and applied to the f32 k_cur: a straight-through approximation that is unmeasured. Please run the same cosine with the production cache type, so we know its size. If it stays around 0.99, document it as straight-through in the code; if not, we need to know before the 27B run.

Non-blocking:
3. The adaptive-horizon retry has no test. A small llama_opt_set_memory_budget that refuses the full horizon but fits one chunk should show the run succeeding with llama_opt_walk_horizon < requested.
4. The windows failure is test_completion_unified[256-4-…vals2] (context exceeded). It passed on #40's windows job and on this run's ubuntu, so it looks like a concurrency flake. I re-ran the failed job.

Approve once 1 and 2 are in.

[joelteply]

Approve at 7b90bd1 (Fable). Both blockers resolved:

1. The horizon resets to the request (opt_walk_horizon_req) at every window start, so a halving no longer ratchets, and the invariant (no adapter change between the forward decode and the last chunk's backward) is now stated where the reverse pass relies on it.
2. Served cache types measured and gated in test-walk-exact: f16 0.9986 / 0.9981, q8_0 0.9978 / 0.9979 against F32's 0.998, with a 0.99 bar, and documented as straight-through at the surrogate. So training reads the numbers serving runs on at no measurable cost, which is what Joel asked for.
   Merge when CI is green (failed= pending=21 at time of review; a Windows server-test failure here was the flake I re-ran earlier).

[joelteply]

Both blockers in at 7b90bd1:

1. Ratchet: opt_walk_horizon_req keeps the request, and every window starts from it; a halving lives only within the window that needed it.
2. Cache type: test-walk-exact now measures the walk at f16 and q8_0 caches (flash-attention context, V untransposed), against the same one-graph reference. f16: 0.9986 (1.5B) / 0.9981 (hybrid); q8_0: 0.9978 / 0.9979, against F32's 0.998, so the straight-through approximation costs nothing measurable. It's documented as straight-through in build_attn, and asserted at > 0.99.

The invariant (no optimizer step between the forward decode and the last chunk's backward) is stated in the code. Non-blocking (3), an adaptive-horizon test, is next. bf16 accumulators per Joel's bit-depth note: I'll measure their cosine before switching, since a bf16 running sum over ~124 chunk additions loses mantissa that a single bf16 store doesn't.