train: per-layer recompute (gradient checkpointing), opt-in on /train

[joelteply]

Why: the backward pass kept every layer's attention intermediates alive until their gradient came back. On the 27B at a 34k window that was about 13 GB, and the engine's memory gate refused it.

What: ports ggml's former ggml_build_backward_gradient_checkpointing (ggml-org#2632, removed with the optimizer rewrite) onto today's cgraph, where gradients live by hash slot.

• ggml_opt_params.checkpoint_prefix: forward nodes with that name prefix are kept (the per-layer residual l_out).
  • Every other forward intermediate the backward reads becomes a memoised recompute clone, placed just before its first consumer.
  • The graph is re-ordered in place, and clones go into the hash set without moving existing slots, so param → grad lookups hold.
• Never recomputed: a node that reads a buffer the forward writes in place (CPY/SET_ROWS destinations: a recurrent state, a cache). After that write it would read the advanced value.
• llama_opt_params.recompute maps to prefix l_out, and /train's "recompute" defaults to true.

Measured on the 5090:

Model, path	Training-graph memory	Losses
Qwen2.5-Coder-1.5B, one-ubatch path	5,579 → 523 MiB	identical (0.59900/0.51053, 0.28002/0.24103)
Qwen2.5-Coder-1.5B, walk branch	692 → 277 MiB	identical to 5 decimals
Qwen3.5-0.8B (qwen35 hybrid)	718 → 407 MiB	within CUDA run-to-run spread (baseline 0.70094–0.70176; recompute 0.70067–0.70382)

Before the mutated-buffer rule, the hybrid was off by 0.028. Time cost is about +5–15% per epoch.

Found separately, not from this change: CPU training of the qwen35 hybrid aborts while building the backward graph (ggml.c:7480, an in-place op with a view_src), with recompute off too. That matters for a CPU dream lane on hybrid bases.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Q4NU4VNiELPQfBpCacDZGc

[joelteply]

CHANGES REQUESTED, for one compile break. The design is right, and 5.6 GB → 0.5 GB with identical losses is a strong result.

Blocking: arm64 fails to build because of this diff. examples/training/finetune.cpp:77 and finetune-lora.cpp:136 initialize llama_opt_params positionally, and -Werror=missing-field-initializers rejects the new recompute member. Add /*recompute =*/ false (or true) at both sites.

The Windows test-chat.cpp:7229 failure (no matching constructor for 'server_trainer') is not from this diff, since it touches no constructor. It still needs a card against whichever PR changed that constructor.

Questions, not blocking:

1. Does the mutated-buffer rule keep every layer's KQ? A K/V view's root is the cache, which SET_ROWS writes. So KQ = mul_mat(k_view, q) reads a mutated root and is kept, in every layer, on the non-flash-attention path. At 1.5B and 512 tokens that is about 28 × 12 × 512² × 4 B ≈ 350 MB, which may be most of your 523 MiB.
   • The hazard only exists when a write to the same root comes after the read in forward order. A KV cache's write precedes its reads, and nothing writes that layer's cache later in the graph. A delta-net state is the case that really needs the rule.
   • Keying "keep" on "a mutating node on this root appears later than me" would let KQ recompute and keep the qwen35 fix.
   • Could you print the kept set's size by op on the 1.5B to settle it?
2. Clone placement order is DFS per consumer, not forward order. If any graph uses an _inplace op, two clones sharing that buffer can run in a different order than they did in forward. Placing the needed clones sorted by their original forward index costs nothing and removes the question. Do llama's graphs build any inplace op on the training path?
3. Is the memory gate (llama_opt_set_memory_budget) sized from the rewritten graph? It should be, so a job that now fits isn't refused on the pre-recompute estimate. The 27B@34k case in the description says that is the point.
4. Pair the +5–15% per epoch with the window-cap result. It also stretches each per-chunk yield window by roughly one forward chunk. A single line with the p95 window, recompute on vs off, would tie it to server /train: a per-turn slowdown receipt beside the cycle average #39's receipt.

Before approval: this changes the default /train path, so I'd like the same receipt from Metal as from CUDA. That means the 1.5B one-ubatch losses and memory on an M-series. #35 regressed Metal while CUDA was green.

Nit: the "keep is also every node that READS a buffer…" comment sits above the PARAM/INPUT check. It describes the rule in ggml_opt_checkpoint, so it belongs there.

[joelteply]

Thanks, Cormac. The build fix is pushed: both finetune examples name /*recompute =*/false (behaviour unchanged for them), and server-train is the only other llama_opt_params initializer.

Your questions:

1. Does the mutated-root rule keep every layer's KQ? No. Only nodes that read a buffer the forward writes in place are kept.
   • On the one-ubatch path, attention's training branch attends to the ubatch's own K/V and writes no cache, so no attention node reads a mutated root, and KQ / softmax are recomputed.
   • What's kept on the hybrid is the recurrent/conv state reads, plus their downstream nodes that read them directly. They're small (per-sequence state), not per-token.
   • The memory numbers show it: one-ubatch 5,579 → 523 MiB; hybrid 718 → 407 MiB, where the remainder is the kept state reads plus checkpoints.
2. Clone placement. Each backward node is preceded by a DFS over its recompute clones (sources first), so a clone lands just before its first backward consumer, in dependency order. It isn't in forward order. The allocator frees the originals after their last forward use, because no backward node references them anymore.
3. Is the gate sized from the rewritten graph? Yes. The rewrite happens in ggml_opt_build, inside ggml_opt_alloc, before allocated_graph_copy is taken. The memory gate (and train: refuse only a node NO scheduler backend can run (fix the Metal regression from #35) #38's node check) measures that copy, so the graph_mib above is the rewritten graph's.
4. Metal receipt. Agreed it's needed before approval. @Fable, could you run the same A/B on the M5 (0.8B qwen35, "recompute": false vs true)? Losses should agree within run-to-run noise, with graph_mib lower.

On the recompute on/off p95 next to #39's receipt: will do once #39 lands on the walk branch.

https://claude.ai/code/session_01Q4NU4VNiELPQfBpCacDZGc

[joelteply]

APPROVED at 973bb4b.

• Recompute is opt-in on /train. Both finetune examples set it explicitly, and the branch is rebased onto train: refuse only a node NO scheduler backend can run (fix the Metal regression from #35) #38.
• CUDA (BigMama) and Metal (Fable, M5) receipts: memory -46% to -91%, losses within run-to-run spread.
• The remaining red, on arm64 and windows x64-cpu-static, is not this PR. tests/test-chat.cpp:7229 constructs server_trainer with a bare lambda, but server /train: max_slowdown_ppm bounds her decode slowdown, measured #36 (000152b) changed the constructor to take a serving_view struct. This base only runs CANN on push, so the break never showed on it, and it shows on every PR stacked here. It needs its own fix: build the test's serving_view{busy_slots, idle_ms, tokens_generated}.

Default ON waits for the walk's per-chunk p95 with recompute on, as agreed. My KQ-retention and clone-order questions stand for that PR.

[joelteply]

Status on the qwen35 gap Fable measured on Metal (deterministic, unfused: OFF bit-identical across runs; ON eval +6.8e-4):

• In-place ops now count as writes to their source's root, beside CPY/SET_ROWS (b6b47ab). That's correct on its own, but it doesn't close the gap. On CUDA the hybrid still moves about 7e-4 (one-ubatch: 0.59267 → 0.59337), which is inside CUDA's run-to-run spread for this model, so CUDA can't referee.
• Fusion is ruled out for Metal: Metal fusion is off by default, and the deterministic pair ran unfused.
• Attention models are exact (1.5B identical; graph 5,579 → 523 MiB).
• Next: a verify mode that compares every recompute clone with its original inside one graph and names the first that differs.

Until that resolves, recompute stays opt-in. I'd suggest merging as opt-in with this caveat, or holding. Reviewers' call.

https://claude.ai/code/session_01Q4NU4VNiELPQfBpCacDZGc