fix(models/moondream2): correct forward so real checkpoint generates coherent text - #616
Merged
Merged
Conversation
The round-1 hardening (#609) made the real vikhyatk/moondream2 checkpoint load and generate, but the output was pure garbage (`!NCJNCJ...`). The forward math itself is correct; auditing it layer by layer against the checkpoint's own reference code (text.py, rope.py, layers.py, vision.py) and the working in-tree Moondream3 port turned up no numeric divergence, and a new remap guard test now pins the full 592-tensor key contract. The real divergence is one level up: the 2025-06-21 revision is trained against the moondream/starmie-v1 tokenizer with Moondream3-style control-token templates (query = [1, 15381, 2] + question + [3], bos = eos = 0), but the official repository never removed its legacy GPT-2 tokenizer.json, and the port picked that stale file up, framed the prompt as GPT-2 "Question:/Answer:" text, and prepended BOS 50256. The model saw an out-of-vocabulary word salad, answered with its true EOS (id 0, which decodes to `!` under the wrong vocabulary), was not stopped because round 1 had moved EOS to 50256, and then looped degenerately. Changes: - `moondream2_prompt`: add `Moondream2PromptStyle` (StarmieTemplates vs LegacyQuestionAnswer) with `detect_moondream2_prompt_style`, which reads the checkpoint's bundled moondream.py (it names the tokenizer repo per revision) and falls back to sniffing tokenizer.json for `<|md_reserved_0|>`. Starmie prompts mirror the working Moondream3 template ids exactly; the legacy Question/Answer framing is kept for the 2025-01-09 .. 2025-04-14 revisions where the GPT-2 contract is genuinely correct (verified against those revisions' config.py: bos_id = eos_id = 50256, templates in GPT-2 ids). - `tokenizer::load_tokenizer`: new starmie override ahead of the local tokenizer.json branch. For starmie-era moondream2 checkpoints whose local tokenizer.json is not starmie, resolve moondream/starmie-v1 via hf-hub (the same path the Moondream3 fallback already uses, cached after the first fetch) with an actionable offline error. - `load_moondream2_vlm` / `resolve_moondream2_eos_token_id`: era-aware bos/eos resolution. Starmie era returns 0 and explicitly ignores the stale legacy tokenizer_config.json (which reports 50256); explicit config.json ids still win; legacy behavior is unchanged. - Thread the detected style through the four prompt call sites (CLI generate_vlm, vlm_runtime, both model_worker paths) via a new `Moondream2VLModel::prompt_style` field, and correct the misleading round-1 comments claiming eos 0 was invalid. - Tests: starmie/legacy prompt shaping with exact id sequences, era detection (moondream.py wins over stale tokenizer.json), tokenizer override firing/skipping matrix, era-aware eos resolution against the real config shapes, and a remap guard asserting the real checkpoint's 592 tensor keys map exactly onto the loader-required set with region tensors dropped. Real-checkpoint-gated parity tests confirm era detection on the actual snapshot and that the resolved tokenizer maps "query" to 15381 and id 0 to `<|endoftext|>`; the heavy forward tests stay `#[ignore]`-gated.
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.
Follow-up hardening for #522 (round 2): after the round-1 EOS fix (#609) the real
vikhyatk/moondream2checkpoint loaded and generated, but emitted garbage (!NCJNCJNCJ...for a solid-orange test image). This PR fixes the actual divergence so the checkpoint generates coherent text.Root cause
The forward pass is numerically correct. A layer-by-layer audit against the checkpoint's own reference implementation (
text.py,rope.py,layers.py,vision.pyshipped in the snapshot) and against the working in-tree Moondream3 port (same architecture minus sparse MoE and tau scaling) found no divergence in the weight remap, fused-QKV split, partial RoPE (32 of 64 dims, NeoX half-split read; the reference's interleaved write-back is a fixed permutation applied to both q and k, so attention scores are invariant to it), LayerNorm eps, parallel attn/MLP wiring, untied biased lm_head, vision tower, crop reconstruction, or the BOS + 729-token bidirectional image prefix.The divergence is the token contract. The 2025-06-21 revision (the one on disk, per its
versions.txtand README) is trained against themoondream/starmie-v1tokenizer: itsmoondream.pyloads that repo from the Hub, and itsconfig.pyframes a query as[1, 15381, 2] + question + [3]withbos_id = eos_id = 0, exactly the contract the working Moondream3 port uses (moondream3_prompthas the identical template ids, and mlxcel already mapsmoondream3 -> moondream/starmie-v1in the remote-tokenizer fallback). Under starmie, 15381 is the wordquery, 32708 isdescribe, and id 0 is<|endoftext|>. The official repository never removed the 2024-era GPT-2tokenizer.json/tokenizer_config.json, so the port picked up the stale local tokenizer, framed the prompt as GPT-2-encoded\n\nQuestion: ...\n\nAnswer:text (the mlx-vlm reference targets the old 2024 checkpoint layout, where that contract is correct), and prepended BOS 50256.The observed garbage confirms this end to end: the emitted ids
[0, 45, 34, 41, ...]decode under GPT-2 as!NCJ..., but under starmie id 0 is<|endoftext|>. The model answered the out-of-vocabulary prompt with its true EOS immediately; round 1 had moved the stop id to 50256 (reading the staletokenizer_config.json), so generation ran past the real stop token into a degenerate loop, and every id was then decoded with the wrong vocabulary.Fix
src/multimodal/moondream2_prompt.rs: newMoondream2PromptStyleenum plusdetect_moondream2_prompt_style. Detection reads the checkpoint's bundledmoondream.py(it names the tokenizer repo per revision: starmie for 2025-06-21+,vikhyatk/moondream2GPT-2 for 2025-01-09 .. 2025-04-14), with atokenizer.json<|md_reserved_0|>sniff as fallback for pruned conversions. Starmie prompts reuse the Moondream3 template ids; the legacy GPT-2Question:/Answer:framing is preserved for the earlier revisions where it is genuinely correct (theirconfig.pydeclaresbos_id = eos_id = 50256and GPT-2-id templates).src/tokenizer/mod.rs:load_tokenizergains a starmie override ahead of the localtokenizer.jsonbranch: for starmie-era moondream2 checkpoints whose local tokenizer is not starmie, resolvemoondream/starmie-v1through the existing hf-hub path (cached after first fetch) with an actionable offline error message. Local starmie tokenizers are used as-is; legacy-era and non-moondream2 models are untouched.src/loading/vlm_special.rs: era-awareresolve_moondream2_eos_token_id. Starmie era resolves bos/eos to 0 and explicitly refuses to consult the stale legacytokenizer_config.json(which reports 50256); explicitconfig.jsonids still take precedence; legacy resolution is unchanged. The loader threads the detected style into the model asMoondream2VLModel::prompt_style.src/commands/generate_vlm.rs,src/multimodal/vlm_runtime.rs,src/server/model_worker.rsx2) pass the model's prompt style; misleading round-1 comments claiming a 0 eos was invalid are corrected;docs/supported-models.mddocuments the two revisions and the offline recovery path.Verification without running the model forward
cargo test --lib moondream2: 30 passed, including a new remap guard that enumerates the real checkpoint's exact 592-tensor key list (verified read-only againstmodel.safetensors) and asserts it maps 1:1 onto the loader-required key set with region tensors dropped./home/inureyes/models/moondream2: era detection returnsStarmieTemplateson the actual snapshot, andload_tokenizeron the real directory resolves starmie end to end, mappingqueryto 15381 and id 0 to<|endoftext|>, proving the template constants line up with the actual tokenizer.cargo check --lib --tests,cargo clippy --lib --tests -- -D warnings,cargo fmt --checkall clean. Heavy forward tests intests/moondream2_parity.rsremain#[ignore]-gated for on-device validation.