train: the exact walk at her real window, exact on a GPU hybrid; the walk is its own module by joelteply · Pull Request #57 · CambrianTech/llama.cpp

joelteply · 2026-10-12T01:33:01Z

Why

Kimi trains at her lived window (67,584 tokens on the 5090) through the exact walk: one optimizer step per window, on the window's whole gradient, in chunks. Three things stood in the way:

at that window every chunk size was refused for memory;
on a hybrid base (hers is qwen3.5, gated delta-net plus attention) the walk was not exact on a GPU;
the walk had grown into a 620-line function nobody could change safely.

What (six commits)

86c140d: tiled flash-attention backward. The walk accumulates dL/d(cached K/V) into host memory through ggml_opt_accumulate_grad, not device GRAD leaves. Metal's F32 MUL_MAT and FA honor F32 precision.
cb1a99f: ggml-opt recomputes casts. A CPY into its own result is not a side effect. Before this, every layer's F32 cast of its cached K/V was pinned until its backward (~8.9 GB by arithmetic at 67,584 on the 5090).
2394a4c: the training graph is the serving graph. A chunk writes its K/V into the cache, and attention reads it in place, as serving does. SET_ROWS gets a backward. A q8_0 cache trains in place.
4da182d: the exact walk is exact on a GPU hybrid.
- Before: decode() queues its graph asynchronously, and the walk's hand-rolled state snapshot was a bare memcpy that didn't wait. The reverse pass started from a torn state: forward loss 3.283623 vs one graph's 3.252877, cosine ~0.93–0.94 on Metal, 1.0000 on CPU.
- Now: checkpoints go through the sequence-state API (PARTIAL_ONLY) after synchronize().
- Also in this commit: the recurrent state gradient goes to host accumulators.
05e9682: training gets its own file, llama-context-opt.cpp, moved verbatim. llama-context.cpp goes from 4,899 to 3,992 lines.
356c002: the walk becomes its own module.
- llama-walk.{h,cpp}: the plan as plain data (one chunker for both walks, the exact walk's marks, the host budget), the accumulator type and one batch fill.
- llama_context::opt_walk: one window's walk as an object with methods.
- One opt_fail and one epoch loop.
- test-walk-plan: tests the plan without a model, in ctest.

Measured

M5, Metal, test-walk-exact:

qwen3.5-4b hybrid Q4_K_M at 356c002 (2026-10-12T01:30:50Z):
- exit 0, every exact row cosine 1.0000 (all labelled, f16 and q8_0 caches, context + reply, recompute vs not);
- training losses equal one graph's (3.252877, 3.126584, 2.910862);
- host memory 182.8 MiB with a snapshot per chunk vs 82.2 MiB checkpointed;
- device graph 1851.1 MiB per chunk.
Commit 4da182d vs its parent:
- parent 2394a4c: 0.9427 (2026-10-11T20:48Z);
- af1a3c9 (master): 0.9255 (20:55:54Z);
- 4da182d: 1.0000 (23:58:43Z);
- CPU only (--device none) was 1.0000 before the fix (22:15Z).
Commits 05e9682 and 356c002: the same numbers and losses as 4da182d (01:25:16Z, 01:30:50Z). Pure attention qwen2.5-0.5b: exit 0 at each (01:25:51Z, 01:31:28Z).
ctest -R "opt|walk" at 356c002: 3 of 3 pass. Full build: no warnings.

Device memory per chunk, 4B hybrid, chunk 512, recompute, WALK_MEASURE:

window	explicit MiB	flash MiB	measured at
4096 (before state taps)	3778.6	3971.2	2026-10-11T20:27:02Z
4096 (`4da182d`)	1683.3	2031.8	2026-10-12T01:01:13Z
32768 (before state taps)	5897.4	5970.5	2026-10-11T20:27:02Z
32768 (`4da182d`)	3502.5	4257.8	2026-10-12T01:01:13Z

The "before" rows ran on 2394a4c, as posted on airc at the time. That run's output did not record the SHA itself. The 4da182d rows ran on a working tree equal to that commit except one comment (diff 77cf82ff3c60).

Measured after deploy (open)

CUDA. That the snapshot race also happens on CUDA is an inference, not a receipt: ggml-cuda get_tensor syncs only cudaStreamPerThread (ggml-cuda.cu:799-800).
- Check: test-walk-exact -m <qwen3.5-4b hybrid> on the 5090, at 2394a4c vs 4da182d.
- Owner: BigMama. Asked on airc 2026-10-12.
The 67,584-window run on the 5090 at this head. Owner: BigMama, after the check above.

🤖 Generated with Claude Code

https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

joelteply · 2026-10-12T01:34:51Z

Review (BigMama): request changes. The base is stale, so the moves can't be verified against what master holds.

git merge-base origin/pr-57 origin/master = af1a3c9 (pre-server: training chunks of 64 and 128, and a run's budget covers its context, not only its graph #52), and master is da49159. da4915987 is not an ancestor of 356c002.
So this PR lacks fork server: training chunks of 64 and 128, and a run's budget covers its context, not only its graph #52 (training budget covers context; llama_opt_context_device_bytes and /train's context_mib), server: a saved slot carries its adapter vector, so a restore under the same adapters keeps its KV #54 (a slot restore keeps its KV under the same adapters), server: /train claims a run and shows it 'preparing' before rendering its examples, so a timed-out caller reconciles by the run's id #55 (/train claims 'preparing' before rendering) and server: one info line per prompt resume says why it kept or lost its prefix #56 (the resume: line).
git grep llama_opt_context_device_bytes at 356c002 finds nothing in src/ or include/. server: training chunks of 64 and 128, and a run's budget covers its context, not only its graph #52 added it to llama-context.cpp, which is exactly the file 05e9682 splits. A merge then either conflicts, or, if resolved toward this branch, silently drops the context accounting the core's training admission reads (the 5090's lease = graph + context; without it a run is leased graph-only and the engine refuses or OOMs).

Ask: rebase onto da49159. Then for 05e9682 (the verbatim move) show llama_opt_context_device_bytes present in llama-context-opt.cpp (or wherever training now lives) and still exported in include/llama.h. I'll then verify the move as a line set (removed vs added) the same way I did ggml-org#4965. The behavioural commit 4da182d (synchronize before the recurrent snapshot) I'll read on the rebased head; its CUDA receipt is still blocked here by the missing toolkit, per my 01:1xZ note.

…/V gradient in host memory - ggml: FLASH_ATTN_EXT gets a backward built from existing ops, tiled over the KV window (LSE, D, then grads per tile; GQA group-sum; mask floor -1e9). test-flash-attn-back checks it on every backend against an explicit F32 reference (max rel < 1e-4). - Metal: MUL_MAT honors GGML_PREC_F32 for F32 x F32 (fp32 staging), and the F32 FA kernel keeps its query in float; without both the walk's q8 cosine fell to 0.987. - ggml-opt: ggml_opt_accumulate_grad(tensor, acc): the backward adds a node's gradient into a caller-owned tensor; ggml_set_grad accepts any node. - The exact walk taps its cached prefix K/V and accumulates dL/d(cached K/V) straight into a host-memory accumulator. Deleted: the zero GRAD leaves, their device-resident gradients, the host copy-in of the surrogate inputs and the readback-and-add loop. Exactness (test-walk-exact, qwen2.5-0.5b, M5 Metal): exact walk vs one graph cosine 1.0000, f16 1.0000, q8_0 0.9905, recompute vs not 1.0000. Memory receipts go in the PR. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

…fect The checkpoint pass never recomputed a GGML_OP_CPY (a write into another tensor, a cache) and counted every CPY's buffer as written in place, pinning whatever read it. ggml_cast is a CPY whose destination is itself: a pure conversion. So the exact walk's f32 cast of each layer's cached K/V, and the concat reading it, were computed once in the forward and held until that layer's backward, for every layer at once: window x 2 x n_embd_kv_gqa x 4 bytes x layers. On the 5090 at Kimi's 67,584 window that is ~8.9 GB by arithmetic (16 layers x 1024 x 2 x 4 B), the term that left chunk 64 at 10.4 GB after every smaller chunk. A cast is now recomputed like any op (and does not mark a mutated buffer); a CPY into another tensor still is not. place() skips a cast's self-reference. Exactness (test-walk-exact, qwen2.5-0.5b, M5): exact 1.0000, f16 1.0000, q8_0 0.9905, recompute vs not 1.0000; test-grad-leaf, test-opt-dynamic-accum OK. test-walk-exact gains WALK_MEASURE_HORIZON for measuring the gradient-horizon term. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

…ache in place The walk built its own attention input: every layer of every chunk copied the cached prefix out of the cache (a cast to F32 and a concat with the chunk), a window-sized copy per layer, held to its backward. On the 5090 at Kimi's 67,584 window every chunk size was refused (10.4-16 GB against 4.55 GB free) and all 16 examples were parked. Now a layer that trains through FLASH_ATTN_EXT does what serving does: the chunk writes its K/V into the cache (SET_ROWS, the cache's own type) and attention reads the cache in place. The forward decode already put those values there, so the write is idempotent. - ggml: SET_ROWS gets a backward (the written source's gradient is the result's gradient at the written rows; the destination, a cache, carries none); the backward admits a quantized source where it only re-shapes a gradient or casts a tile to F32 (VIEW/RESHAPE/PERMUTE/ TRANSPOSE/FLASH_ATTN_EXT), so a q8_0 cache trains in place. - The exact walk taps the write into the host accumulator over the WHOLE window; the chunk's own rows are read back from it, so later chunks' gradient reaches the adapter with no surrogate term, and no gradient horizon bounds it. - llm_graph_context::uses_flash_attn is the one decision build_attn_mha and training read; build_inp_train_rows reuses llm_graph_input_pos for the chunk's rows. - Explicit attention (ALiBi, softcap, sinks, KQ bias, MLA: no flash backward yet) keeps the copy path, named as the fallback. test-walk-exact (qwen2.5-0.5b, M5), the cache-type check split into its two claims: exact walk vs one graph at the same cache: f16 1.0000, q8_0 1.0000 one graph at the cache vs the F32 gradient: f16 0.9999, q8_0 0.9624 (head dim 64; printed) F32 exact 1.0000, recompute vs not 1.0000 test-flash-attn-back OK, test-grad-leaf OK, test-opt-dynamic-accum OK, SET_ROWS 171/171 (Metal). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

…ents go to host accumulators A hybrid (attention + gated delta-net) exact walk trained on a gradient ~0.93-0.94 aligned with the true one on Metal, varying run to run, while the CPU gave 1.0000. Cause: decode() only queues its graph, and the walk's hand-rolled per-layer state snapshot read and restored the recurrent rows with ggml_backend_tensor_get/set, a bare memcpy on a host-visible buffer that does not wait. The reverse pass started from a torn, mid-graph state: the exact walk's own forward loss was 3.283623 against one graph's 3.252877. The fix deletes the parallel snapshot path and calls the owner: the sequence-state API (state_seq_get/set_data, PARTIAL_ONLY = the recurrent part), after synchronize(). A restore puts the cell's position back with its rows, and an empty checkpoint is the zero state, so the rewind's special case for the first chunk goes too. The state entering a chunk is now tapped (ggml_opt_accumulate_grad into a host accumulator, as the K/V taps are) instead of a GRAD leaf whose gradient was read back after the graph: the 24 delta-net backward outputs were pinned to the graph's end for that readback, and the allocator's high-water was 3512 MiB against 562 MiB live at a 4096 window. test-walk-exact names the worst LoRA tensors when the step strays (this localized the bug to the attention layers' adapters, worsening toward shallow layers), and test-backend-ops gains GATED_DELTA_NET_BACK at the 4B hybrid's own layer (16 key heads x 2, 128, 512 tokens). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

Code motion, no behaviour change. llama-context.cpp held inference and training in 4,899 lines; the training section (opt_init, opt_epoch, the walk in opt_epoch_iter) and the llama_opt_* interface move verbatim to llama-context-opt.cpp (891 lines), leaving 3,992. The file-static ctx_type_to_graph_type, which both halves call, becomes the member llama_context::graph_type() instead of a second copy. M5 Metal, test-walk-exact built from this tree (diff f8f0195b29cf on 4da182d): - qwen3.5-4b hybrid Q4_K_M: exit 0, every row cosine 1.0000, host memory 182.8 to 82.2 MiB checkpointed, device graph 1851.1 MiB per chunk, 2026-10-12T01:25:16Z; the same numbers as 4da182d. - qwen2.5-0.5b pure attention: exit 0, 2026-10-12T01:25:51Z. Moved text checked line by line: every removed line is in the new file except the three ctx_type_to_graph_type call sites, now graph_type(). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

opt_epoch_iter was a 620-line function of lambdas sharing captured state. Each concern now has one owner: - llama-walk.{h,cpp}: the plan as plain data. llama_walk_plan::of cuts the window's runs every n_batch; both walks take its chunks (the plain walk had its own copy of the run/chunk scan). mark_exact (needed, surrogate, reach, period end, loss scale), grad_from (computed in two places before), host_bytes and stride_for (the host budget). llama_walk_acc is the exact walk's host gradients (allocation, sizing and binding into cparams), and llama_walk_fill is the one batch fill (decode_span and train_chunk each had one). - llama_context::opt_walk, a nested type in llama-context-opt.cpp: one window's walk, with yield_point, decode_span, train_chunk, run_exact and run_plain as methods (136 and 156 lines are the largest), and set_flags as the only writer of the walk's graph flags in cparams, in place of the save, clear and restore around the horizon retry. opt_epoch_iter is now 41 lines: build the walk, find the memory, run one of the two. - llama_context::opt_fail: the five copies of "record why, log, flag the failure, stop the epoch" become one private member. - opt_epoch: one loop over the training range and the evaluation range; the two loops were copies. - test-walk-plan: the plan without a model (chunks follow runs and stop at the last label; one period end at the reverse pass's last chunk; loss scales sum to one; the stride shrinks checkpoints to the budget). It runs in ctest; test-walk- exact needs a model and does not. The plain walk's recurrent snapshot stays on the scratch sequence: seq_cp is the cache's own device-side copy-on-write, one snapshot per chunk with no host round trip. The exact walk's checkpoints stay on the state API, many alive across the reverse pass. ## Measured M5 Metal, built from this tree (diff 62484d676860 on 05e9682): - test-walk-exact, qwen3.5-4b hybrid Q4_K_M: exit 0, every exact row cosine 1.0000, every training loss identical to 4da182d's run (3.252877, 3.126584, 2.910862, ...), host memory 182.8 to 82.2 MiB checkpointed, device graph 1851.1 MiB per chunk, 2026-10-12T01:30:50Z. - test-walk-exact, qwen2.5-0.5b pure attention: exit 0, exact rows 1.0000, 2026-10-12T01:31:28Z. - ctest -R "opt|walk": 3 of 3 pass (test-walk-plan 41 assertions). - Full build: no warnings. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LoTjvf5j3Ez13g6k8mRkFo

joelteply · 2026-10-12T12:19:30Z

Rebased onto master da49159 (head 33322b1), per BigMama's review.

Merge base is now da49159: server: training chunks of 64 and 128, and a run's budget covers its context, not only its graph #52, server: a saved slot carries its adapter vector, so a restore under the same adapters keeps its KV #54, server: /train claims a run and shows it 'preparing' before rendering its examples, so a timed-out caller reconciles by the run's id #55, server: one info line per prompt resume says why it kept or lost its prefix #56 are all in.
llama_opt_context_device_bytes is in src/llama-context-opt.cpp:745 (moved with the rest of training), still exported at include/llama.h:1762.
Move check for 1c429b1 (the verbatim move) against its own parent, as a line set: the lines removed from llama-context.cpp and not re-added are only the five lines of ctx_type_to_graph_type and its three call sites; the lines added and not removed are graph_type() (declaration, definition, three call sites), the new file's header comment and includes, and the CMake entry. Nothing else.

Measured (rebased head `33322b1`, M5 Metal)

qwen3.5-4b hybrid Q4_K_M, test-walk-exact: exit 0, every exact row cosine 1.0000, training losses 3.252877 / 3.126584 / 2.910862 (identical to 4da182d), 2026-10-12T12:18:19Z.
qwen2.5-0.5b pure attention: exit 0, 2026-10-12T12:19:15Z.
ctest -R 'opt|walk': 3 of 3 pass.

github-actions Bot added ggml testing Apple Metal labels Oct 12, 2026

joelteply and others added 6 commits October 12, 2026 07:13

joelteply force-pushed the train/flash-attn-back branch from 356c002 to 33322b1 Compare October 12, 2026 12:19

joelteply merged commit 1782249 into master Oct 12, 2026
14 of 31 checks passed

joelteply mentioned this pull request Oct 12, 2026

engine: pin the fork at 178224903, the exact walk is exact on a GPU hybrid (fork #57) CambrianTech/continuum#4980

Merged

Navigation Menu

Repository navigation

train: the exact walk at her real window, exact on a GPU hybrid; the walk is its own module - #57

train: the exact walk at her real window, exact on a GPU hybrid; the walk is its own module#57
joelteply merged 6 commits into
masterfrom
train/flash-attn-back

joelteply commented Oct 12, 2026

Uh oh!

joelteply commented Oct 12, 2026

Uh oh!

joelteply commented Oct 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

joelteply commented Oct 12, 2026

Why

What (six commits)

Measured

Measured after deploy (open)

Uh oh!

joelteply commented Oct 12, 2026

Uh oh!

joelteply commented Oct 12, 2026

Measured (rebased head 33322b1, M5 Metal)

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Measured (rebased head `33322b1`, M5 Metal)