Gate resuming on the checkpoint rms-eps, and keep GQA groups whole - #39
Merged
Conversation
… whole GQA groups
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Currently,
pretrain --resumeandfinetune --resumecompare the checkpoint's shape field by field, but neverrms-eps. A checkpoint whose epsilon differs from the architecture's build default would train silently at the wrong value, andexportwould 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
inspectprint 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.src/arch/llama.ts,qwen3.ts,gemma3.tsconfigMatchescomparesrmsEps; gemma3 gains its missing--rms-epsand--headsflagssrc/arch/common.tsdefaultKVHeads(nHeads, ratio)derives the KV count as a divisor (the oldround(n/ratio)could miss);assertWholeGQArefuses a fractional groupsrc/commands/inspect.ts--jsoncarry--rms-eps;resumeFlagsis exported so the test can pin ittests/arch-roundtrip.tsrms-epsis 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 countableagents.md,docs/adding-an-architecture.mdWhat 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.vocabis 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:
--headsis now a real flag. A width thatheadDim*2does not divide is unbuildable while the head count is derived, and buildable once--headsnames it (the 270M shape: 640 hidden, 256 head-dim, 4 heads).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:inspectprints 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-headsthe user explicitly passes, or--heads 0, is now an error rather than a silent mis-train, on all three architectures.How to test
deno task test(new checks intests/arch-roundtrip.ts; the whole suite covers the f32 comparison and parity).Step 2 was run for real on tiny llama and gemma3 checkpoints (M1 Max, GPU), using the exact value
inspectprinted.Fixes #36.