Skip to content

server: training chunks of 64 and 128, and a run's budget covers its context, not only its graph - #52

Merged
joelteply merged 3 commits into
masterfrom
train/chunks-below-256
Oct 11, 2026
Merged

joelteply merged 3 commits into
masterfrom
train/chunks-below-256

Conversation

@joelteply

Copy link
Copy Markdown

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:

  1. graph_mib (ggml_opt_peak_graph_bytes) is the optimizer graph only.
  2. The training context's own device buffers (its Q8_0 K/V for the whole window, plus its compute buffers) are allocated at llama_init_from_model, before llama_opt_set_memory_budget applies. 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

  • Chunks of 64 and 128. The chunk is still a divisor of the window: 64, 128, or a multiple of 256. 192 and 320 are refused. train_chunk_for returns 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_init asserts only divisibility.
  • The budget covers the context. New llama_opt_context_device_bytes(ctx) sums memory_breakdown()'s context + compute on non-host buffer types, excluding the shared model weights. /train gives the graph budget − context (at least one byte, since 0 means no budget), and the preflight refuses as before.
  • /train reports context_mib beside graph_mib, so the caller files the whole footprint.

Tests

test-chat passes (Windows, CUDA build). llama-server builds. Acceptance on hardware comes with the core change: a run's context_mib + graph_mib fits the lease it was given.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Q4NU4VNiELPQfBpCacDZGc

…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).
@joelteply

Copy link
Copy Markdown
Author

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 joelteply left a comment

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@joelteply

Copy link
Copy Markdown
Author

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.

@joelteply
joelteply merged commit 4842741 into master Oct 11, 2026
13 of 33 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant