Skip to content

Gate resuming on the checkpoint rms-eps, and keep GQA groups whole - #39

Merged
felladrin merged 5 commits into
mainfrom
fix/resume-rms-eps-gate
Aug 31, 2026
Merged

Gate resuming on the checkpoint rms-eps, and keep GQA groups whole#39
felladrin merged 5 commits into
mainfrom
fix/resume-rms-eps-gate

Conversation

@felladrin

Copy link
Copy Markdown
Owner

Currently, pretrain --resume and finetune --resume compare the checkpoint's shape field by field, but never rms-eps. A checkpoint whose epsilon differs from the architecture's build default would train silently at the wrong value, and export would then write that wrong epsilon into the new GGUF. This is the same defect class as the RoPE base (fixed in #38).

This PR adds the epsilon to the resume gate, and makes inspect print the flag the gate now demands. Making the gate complete for gemma3 (it must only name a flag the architecture can accept) surfaced two pre-existing GQA bugs, which are fixed here too.

File Change
src/arch/llama.ts, qwen3.ts, gemma3.ts configMatches compares rmsEps; gemma3 gains its missing --rms-eps and --heads flags
src/arch/common.ts defaultKVHeads(nHeads, ratio) derives the KV count as a divisor (the old round(n/ratio) could miss); assertWholeGQA refuses a fractional group
src/commands/inspect.ts the "Resume with:" line and --json carry --rms-eps; resumeFlags is exported so the test can pin it
tests/arch-roundtrip.ts per arch: a different rms-eps is rejected naming the flag; the gate's named flag is printed with the checkpoint's own value; the CLI default shape keeps whole GQA groups; gemma3's wide (non-square) head block is buildable and countable
agents.md, docs/adding-an-architecture.md invariant 2 and the arch-adding guide name the gate/flag/inspect contract

What the gate now covers

Every field the resume gate names is one the user can satisfy, printed with the checkpoint's value: shape, window, the RoPE bases, and now rms-eps. vocab is the one exception, since it comes from the tokenizer, not a flag.

Behavior changes

Two defaults that produced artifacts llama.cpp would not load are now corrected. They are stated, not silent:

  1. gemma3 --heads is now a real flag. A width that headDim*2 does not divide is unbuildable while the head count is derived, and buildable once --heads names it (the 270M shape: 640 hidden, 256 head-dim, 4 heads).
  2. The derived KV default is now a divisor (defaultKVHeads), not a rounded ratio. Derived gemma3 defaults are unchanged, because its width rule forces an even head count; llama and qwen3 defaults move only where the old ratio did not divide the head count (llama 16 heads 5 to 4, 10 heads 3 to 2, 8 heads 3 to 2; qwen3 9 heads 5 to 3, 3 heads 2 to 1), and every such shape was already rejected by llama.cpp. A checkpoint this tool wrote at one of those old defaults is now unresumable: inspect prints the old --kv-heads, and the guard correctly refuses to build the fractional group. That is the right outcome, since llama.cpp rejected that artifact too.

A fractional --kv-heads the user explicitly passes, or --heads 0, is now an error rather than a silent mis-train, on all three architectures.

How to test

  1. deno task test (new checks in tests/arch-roundtrip.ts; the whole suite covers the f32 comparison and parity).
  2. Export a checkpoint with a non-default epsilon, then:
deno run -A cli.ts inspect --model ckpt.gguf
# "Resume with:" includes --rms-eps <value>

deno run -A cli.ts pretrain --resume ckpt.gguf --data ...tokens ...
# aborts: --resume config mismatch (rms-eps: built 0.00001 vs checkpoint 0.000019999999494757503)

deno run -A cli.ts pretrain --resume ckpt.gguf --data ...tokens ... --rms-eps 0.000019999999494757503
# trains, and the written checkpoint carries that epsilon

Step 2 was run for real on tiny llama and gemma3 checkpoints (M1 Max, GPU), using the exact value inspect printed.

Fixes #36.

@felladrin
felladrin marked this pull request as ready for review August 31, 2026 22:29
@felladrin
felladrin merged commit 5be6e7e into main Aug 31, 2026
1 check passed
@felladrin
felladrin deleted the fix/resume-rms-eps-gate branch August 31, 2026 22:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

--resume ignores the checkpoint rms-eps

1 participant