Repository navigation
server: training chunks of 64 and 128, and a run's budget covers its context, not only its graph - #52
Conversation
…ng step time for score memory With no flash-attention backward, the explicit attention scores are window x chunk x heads per attention layer. On the 5090 (continuum, 2026-10-11) they were ~75% of Qwen3.8-27B's training graph at chunk 256, so a window holding her whole lived turns did not fit beside serving. A chunk must divide the window (a multiple of 256), and 64 and 128 do; a gated delta-net pads to its own 64-token chunk, so 64 is the floor. Chunks of 256 and up keep exactly the old rule, so every footprint keyed by them is unchanged.
…graph The training context's own device buffers (its Q8_0 K/V for the whole window, and its compute buffers) are allocated when it is created, before llama_opt_set_memory_budget applies to the graph, so they sat outside the caller's governed lease: ~0.6 GB at a 16,896-token window on the 5090, growing with the window, on the GPU serving shares. The graph now gets what the budget leaves after them, and /train reports context_mib beside graph_mib, so a caller files the whole footprint. New: llama_opt_context_device_bytes (memory_breakdown's context + compute on non-host buffer types; the model's weights are shared and excluded).
|
Approve at 13b4605, Fable (M5), as a comment (shared account), pending green. Chunks of 64 and 128 are admitted beside multiples of 256 (each divides a 256-multiple window; 64 is also the gated-delta-net layer's own chunk), train_chunk_for rounds an ask below 256 down to 128 or 64, and the training context's device bytes (its window K/V and compute buffers) now come off the lease before the graph's budget, with context_mib reported, so the lease covers the whole run. I checked the unified-memory case: Metal's shared and private buffer types report is_host = false, so the M5 counts its context too, not zero. |
…ers are where the graph goes The optimizer shares the training context's scheduler, whose allocator replaces the context's compute buffers with the graph's, so llama_opt_graph_bytes already covers them. Counting memory_breakdown's compute too would count them twice (Cormac on #52).
joelteply
left a comment
There was a problem hiding this comment.
Approved at d29dbf6. This resolves my point 3: the context term is data.context only, and the compute buffers are covered by the graph's preflight, since ggml_opt shares this context's sched. One residual edge, not blocking: since the sched buffer is max(inference reserve, graph), not their sum, graph+context under-counts only when the training graph is smaller than the inference compute reserve. That is implausible for a training step. The pin bump belongs in ggml-org#4939, which I'll re-review at its next head.
|
Approve at d29dbf6, Fable (M5), as a comment (shared account), pending green. The only change since 13b4605 is Cormac's double-count fix: llama_opt_context_device_bytes counts the context's memory (window K/V, recurrent state) and no longer its compute buffers, which the optimizer reuses for the training graph and graph_mib already measures; comments updated to match. |
Why
On the 5090 (continuum, 2026-10-11) a lived turn needs a ~20k-token window (p90 28k), but the window that fit beside serving at chunk 256 was 16,896, so 257 of Kimi's 311 turns were skipped. With no flash-attention backward, the explicit scores are window × chunk × heads, so a smaller chunk trades step time for window.
Two things were found along the way:
graph_mib(ggml_opt_peak_graph_bytes) is the optimizer graph only.llama_init_from_model, beforellama_opt_set_memory_budgetapplies. They sat outside the caller's governed lease: ~0.6 GB at 16,896, ~2.4 GB at 68,608, on the GPU serving shares.What
train_chunk_forreturns 64 or 128 under 256 and keeps the old divisor rule at 256 and up, so every footprint keyed by those is unchanged. The gated delta-net pads to its own 64-token chunk (delta-net-base.cpp,CS), so 64 is its floor;opt_initasserts only divisibility.llama_opt_context_device_bytes(ctx)sumsmemory_breakdown()'s context + compute on non-host buffer types, excluding the shared model weights./traingives the graphbudget − context(at least one byte, since 0 means no budget), and the preflight refuses as before./trainreportscontext_mibbesidegraph_mib, so the caller files the whole footprint.Tests
test-chatpasses (Windows, CUDA build).llama-serverbuilds. Acceptance on hardware comes with the core change: a run'scontext_mib+graph_mibfits the lease it was given.🤖 Generated with Claude Code
https://claude.ai/code/session_01Q4NU4VNiELPQfBpCacDZGc