Skip to content

fix(models/moondream2): correct forward so real checkpoint generates coherent text - #616

Merged
inureyes merged 1 commit into
mainfrom
fix/issue-522-moondream2-starmie-contract
Jul 2, 2026
Merged

fix(models/moondream2): correct forward so real checkpoint generates coherent text#616
inureyes merged 1 commit into
mainfrom
fix/issue-522-moondream2-starmie-contract

Conversation

@inureyes

@inureyes inureyes commented Jul 2, 2026

Copy link
Copy Markdown
Member

Follow-up hardening for #522 (round 2): after the round-1 EOS fix (#609) the real vikhyatk/moondream2 checkpoint 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.py shipped 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.txt and README) is trained against the moondream/starmie-v1 tokenizer: its moondream.py loads that repo from the Hub, and its config.py frames a query as [1, 15381, 2] + question + [3] with bos_id = eos_id = 0, exactly the contract the working Moondream3 port uses (moondream3_prompt has the identical template ids, and mlxcel already maps moondream3 -> moondream/starmie-v1 in the remote-tokenizer fallback). Under starmie, 15381 is the word query, 32708 is describe, and id 0 is <|endoftext|>. The official repository never removed the 2024-era GPT-2 tokenizer.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 stale tokenizer_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: new Moondream2PromptStyle enum plus detect_moondream2_prompt_style. Detection reads the checkpoint's bundled moondream.py (it names the tokenizer repo per revision: starmie for 2025-06-21+, vikhyatk/moondream2 GPT-2 for 2025-01-09 .. 2025-04-14), with a tokenizer.json <|md_reserved_0|> sniff as fallback for pruned conversions. Starmie prompts reuse the Moondream3 template ids; the legacy GPT-2 Question:/Answer: framing is preserved for the earlier revisions where it is genuinely correct (their config.py declares bos_id = eos_id = 50256 and GPT-2-id templates).
  • src/tokenizer/mod.rs: load_tokenizer gains a starmie override ahead of the local tokenizer.json branch: for starmie-era moondream2 checkpoints whose local tokenizer is not starmie, resolve moondream/starmie-v1 through 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-aware resolve_moondream2_eos_token_id. Starmie era resolves bos/eos to 0 and explicitly refuses to consult the stale legacy tokenizer_config.json (which reports 50256); explicit config.json ids still take precedence; legacy resolution is unchanged. The loader threads the detected style into the model as Moondream2VLModel::prompt_style.
  • Call sites (src/commands/generate_vlm.rs, src/multimodal/vlm_runtime.rs, src/server/model_worker.rs x2) pass the model's prompt style; misleading round-1 comments claiming a 0 eos was invalid are corrected; docs/supported-models.md documents 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 against model.safetensors) and asserts it maps 1:1 onto the loader-required key set with region tensors dropped.
  • Tokenizer override matrix tests (fires for starmie era with stale local tokenizer; skipped for local starmie, legacy era, and non-moondream2 models).
  • Real-checkpoint-gated (non-forward) parity tests, run locally against /home/inureyes/models/moondream2: era detection returns StarmieTemplates on the actual snapshot, and load_tokenizer on the real directory resolves starmie end to end, mapping query to 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 --check all clean. Heavy forward tests in tests/moondream2_parity.rs remain #[ignore]-gated for on-device validation.

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.
@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions area:models Model architectures, weights, loading, metadata labels Jul 2, 2026
@inureyes
inureyes merged commit cd6791f into main Jul 2, 2026
5 checks passed
@inureyes
inureyes deleted the fix/issue-522-moondream2-starmie-contract branch July 2, 2026 13:49
@inureyes inureyes self-assigned this Aug 31, 2026
@inureyes inureyes added status:done Completed and removed status:review Under review labels Sep 6, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:models Model architectures, weights, loading, metadata status:done Completed type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant