From 7d41b8f699c0c0a9a0e8465a034cccf529246c49 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sat, 15 Aug 2026 22:51:27 -0700 Subject: [PATCH 1/2] feat: auto-detect VLM checkpoints on the CLI, and fix a false positive it exposed MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extracted from fix/pr57-speculative-ci rather than a rebase (see #109's resolution and the discussion around it). Its ModelArchitectureProbe.swift and tests were already shipped verbatim; only the CLI wiring and its test were still missing. **The gap.** Server.swift already probed the architecture at load time but discarded the result for the CLI entry point: `let isVision = self.vision`. Confirmed the user-facing effect directly: `SwiftLM --model LiquidAI/LFM2.5-VL-450M-MLX-4bit` (no --vision) printed "Loading LLM" and the subsequent image request failed. InferenceEngine.swift — SwiftBuddy's own loader — already auto-detects unconditionally; the CLI now mirrors that, keeping --vision/--audio as explicit overrides: let isVision = self.vision || (!self.audio && architecture.supportsVision) Verified both directions: LFM2.5-VL-450M-MLX-4bit now loads as a VLM and answers an image request without --vision; a plain LLM (Qwen2.5-0.5B) still loads as LLM with no flags; --vision still forces VLM loading explicitly. **What wiring this up exposed.** The fixtures suite broke: `moe-nested` started failing with "Key vision_tower.patch_embedder... not found", because that fixture's `model_type: "gemma4"` — carried since #145, for the earlier, narrower reason that Gemma4Configuration decodes text_config only, giving the MoE-nesting test a real nested config to exercise — was now read by the CLI's newly-active auto-detection and routed to VLMModelFactory. That is not a fixture bug on its own. `MLXLLM/Gemma4.swift`'s Gemma4Configuration (the LLMModelFactory "gemma4" entry) has no vision_config field at all — a real, intentionally-supported text-only checkpoint shape. ModelArchitectureProbe.knownVisionModelTypes listed bare "gemma4" as vision-triggering regardless, which is a genuine false positive for any real text-only Gemma4 checkpoint using that shape, not just this fixture — auto- detection would have broken loading one in production. Removed "gemma4" from that list; the separate `vision_config != nil` check already distinguishes the two correctly, since only VLMModelFactory's Gemma4Configuration requires that field. Also swapped the fixture's decoy container from vision_config to audio_config — unrelated to the bug above, but vision_config as a decoy was already the wrong choice: a real text-only Gemma4 checkpoint would never carry that key, so the fixture is more honest this way regardless of the probe fix. audio_config is excluded from the expert-count walk by ModelProfiler.nonLanguageContainers the same way vision_config was, so the decoy still proves what it proved before. Regression test added: testVLM_AutoDetectsLFM25WithoutVisionFlag in tests/SwiftBuddyTests/VLMTests.swift, adapted from fix/pr57-speculative-ci to match main's already-refactored captureStartupOutput helper (main had the refactor; only this second test case was missing). Red-green verified: fails with "Unsupported model type: lfm2-vl" against the pre-fix isVision line, passes once restored. Verified together: fixtures 6/0, contract 10/0/2, both VLMTests cases, LFM2.5-VL-450M-MLX-4bit auto-detects, Qwen2.5-0.5B does not false-positive, explicit --vision still works. Refs #109 Co-Authored-By: Claude Opus 5 --- .../ModelArchitectureProbe.swift | 11 ++++++++++- Sources/SwiftLM/Server.swift | 14 +++++++++++++- scripts/make-test-fixtures.py | 15 ++++++++++++--- tests/SwiftBuddyTests/VLMTests.swift | 19 +++++++++++++++++++ tests/fixtures/moe-nested/config.json | 4 ++-- tests/test-fixtures.sh | 4 ++-- 6 files changed, 58 insertions(+), 9 deletions(-) diff --git a/Sources/MLXInferenceCore/ModelArchitectureProbe.swift b/Sources/MLXInferenceCore/ModelArchitectureProbe.swift index 6c9d7067..19b2d478 100644 --- a/Sources/MLXInferenceCore/ModelArchitectureProbe.swift +++ b/Sources/MLXInferenceCore/ModelArchitectureProbe.swift @@ -27,7 +27,16 @@ public enum ModelArchitectureProbe { "qwen3_5_moe", "idefics3", "gemma3", - "gemma4", + // "gemma4" deliberately absent: this codebase registers that exact model_type + // in both LLMModelFactory (MLXLLM/Gemma4.swift, text_config only, no vision + // field at all) and VLMModelFactory (MLXVLM/Gemma4.swift, vision_config + // required). The string alone cannot tell them apart — matching on it treated + // every real text-only Gemma4 checkpoint as a VLM and routed it to a factory + // whose config decoder requires a field the checkpoint never has, which is a + // hard failure, not a degraded one. The vision_config-presence check below + // already distinguishes the two correctly, since only the VLM struct requires + // that field — this entry was strictly redundant on the VLM side and wrong on + // the LLM side. "smolvlm", "fastvlm", "llava_qwen2", diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index f876a62c..01cec3bc 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -624,7 +624,19 @@ struct MLXServer: AsyncParsableCommand { configuration: modelConfig, downloader: downloader ) - let isVision = self.vision + // architecture is probed above for both paths; the CLI previously discarded it + // here (`let isVision = self.vision`), so a known VLM checkpoint loaded as a + // plain LLM unless the user remembered `--vision` — confirmed by running + // LiquidAI/LFM2.5-VL-450M-MLX-4bit without the flag: "Loading LLM" and the + // subsequent image request failed. InferenceEngine.swift (SwiftBuddy's loader) + // already auto-detects unconditionally; this mirrors that for the CLI binary + // while keeping --vision/--audio as explicit overrides. + let isVision = self.vision || (!self.audio && architecture.supportsVision) + if architecture.supportsVision, !self.vision, !self.audio { + print( + "[SwiftLM] Auto-detected VLM config (\(architecture.modelType ?? "unknown")); enabling vision mode." + ) + } let container: ModelContainer // Handle getting the simple model ID string for the tracker diff --git a/scripts/make-test-fixtures.py b/scripts/make-test-fixtures.py index ee2e48cb..2a6d8a5f 100755 --- a/scripts/make-test-fixtures.py +++ b/scripts/make-test-fixtures.py @@ -242,7 +242,9 @@ def build_moe_nested(out): `Gemma4Configuration` decodes `text_config` and nothing else, so here the count exists *only* one level down. That makes two things real rather than decorative: the nested walk added in #112's review follow-up, and the rule that a count under - `vision_config` must not be mistaken for the language model's. + a sibling container must not be mistaken for the language model's. The decoy + container is `audio_config`, not `vision_config` — see the comment where it is + built for why. It also covers the fused-expert remap: real gemma4 checkpoints ship `experts.gate_up_proj` as one tensor that sanitize splits in half into @@ -288,8 +290,15 @@ def build_moe_nested(out): "architectures": ["Gemma4ForConditionalGeneration"], "vocab_size": VOCAB, "text_config": text_config, - # A decoy the language-model walk has to skip. - "vision_config": {"model_type": "gemma4_vision", "num_experts": 999}, + # A decoy the language-model walk has to skip. audio_config, not + # vision_config: ModelArchitectureProbe.inspect treats the mere presence of a + # vision_config key as proof of vision support, so using it here made this + # text-only fixture auto-detect as a VLM and route to VLMModelFactory, which + # has no vision weights to find — found when the SwiftLM CLI started acting on + # that probe's result instead of discarding it. audio_config is excluded from + # the expert-count walk by ModelProfiler.nonLanguageContainers the same way, + # without tripping that probe. + "audio_config": {"model_type": "gemma4_audio", "num_experts": 999}, } json.dump(cfg, open(os.path.join(out, "config.json"), "w"), indent=2) diff --git a/tests/SwiftBuddyTests/VLMTests.swift b/tests/SwiftBuddyTests/VLMTests.swift index f96bf4f5..d65df3fc 100644 --- a/tests/SwiftBuddyTests/VLMTests.swift +++ b/tests/SwiftBuddyTests/VLMTests.swift @@ -12,6 +12,25 @@ final class VLMTests: XCTestCase { XCTAssertTrue(found, "Output should indicate VLM is loading. Got: \(accumulated)") } + // Feature 2: a known VLM checkpoint loads as a VLM without --vision. + // + // Before this test could pass, the CLI computed the architecture probe and then + // discarded it (`let isVision = self.vision`) — confirmed by running this exact + // model without --vision and observing "Loading LLM" followed by a failed image + // request. InferenceEngine.swift (SwiftBuddy's own loader) already auto-detected + // unconditionally; the CLI now does the same while --vision/--audio remain + // explicit overrides. + func testVLM_AutoDetectsLFM25WithoutVisionFlag() async throws { + let accumulated = try await captureStartupOutput(arguments: [ + "--model", "LiquidAI/LFM2.5-VL-450M-MLX-4bit", + ], timeout: 20.0) + + XCTAssertTrue( + accumulated.contains("Auto-detected VLM config") + || accumulated.contains("Loading VLM"), + "Output should indicate VLM auto-detection/loading. Got: \(accumulated)" + ) + } private func captureStartupOutput( arguments: [String], diff --git a/tests/fixtures/moe-nested/config.json b/tests/fixtures/moe-nested/config.json index 8dc58d17..8872d5da 100644 --- a/tests/fixtures/moe-nested/config.json +++ b/tests/fixtures/moe-nested/config.json @@ -32,8 +32,8 @@ "top_k_experts": 2, "moe_intermediate_size": 32 }, - "vision_config": { - "model_type": "gemma4_vision", + "audio_config": { + "model_type": "gemma4_audio", "num_experts": 999 } } \ No newline at end of file diff --git a/tests/test-fixtures.sh b/tests/test-fixtures.sh index 7a26f2f3..632a5872 100755 --- a/tests/test-fixtures.sh +++ b/tests/test-fixtures.sh @@ -11,7 +11,7 @@ # kv-shared-absent #120: gemma-4-e4b shape, shared layers ship no k/v # kv-shared-present b674: gemma-4-e2b shape, shared layers ship k/v anyway # moe-nested #112: expert count nested under text_config only, with a -# decoy count under vision_config; plus the fused +# decoy count under audio_config; plus the fused # experts.gate_up_proj split that sanitize performs # # The output is gibberish by construction — the weights are random. A fixture passes @@ -139,7 +139,7 @@ if [ "$ready" -ne 1 ]; then elif grep -q "is not MoE" "$MOE_LOG"; then fail "nested expert count missed: $(grep 'is not MoE' "$MOE_LOG" | head -1)" else - pass "moe-nested: expert count found under text_config, vision decoy ignored" + pass "moe-nested: expert count found under text_config, decoy container ignored" fi cleanup From 6a74c37112c8384a0e3a0b21519ae7579ee7d916 Mon Sep 17 00:00:00 2001 From: Simba Zhang Date: Sat, 15 Aug 2026 23:24:38 -0700 Subject: [PATCH 2/2] fix: don't auto-detect vision when speculative/MTP decoding is requested MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CI's speculative-decoding suite started crashing (Trace/BPT trap on the first real generation request) after auto-detection began routing mlx-community/Qwen3.5-*-4bit through VLMModelFactory instead of LLMModelFactory. That checkpoint genuinely ships a vision_config even when used purely as a text draft/main pair in this test, and qwen3_5 is registered in both factories — the same ambiguity as gemma4, just surfaced through a different flag combination. --draft-model/--dflash/--mtp only wire up BaseLanguageModel from LLMModelFactory, so auto-detection now backs off whenever speculative decoding is requested, matching the pre-auto-detect behavior for that path. --vision remains a valid explicit override. --- Sources/SwiftLM/Server.swift | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/Sources/SwiftLM/Server.swift b/Sources/SwiftLM/Server.swift index 01cec3bc..526caa85 100644 --- a/Sources/SwiftLM/Server.swift +++ b/Sources/SwiftLM/Server.swift @@ -631,8 +631,23 @@ struct MLXServer: AsyncParsableCommand { // subsequent image request failed. InferenceEngine.swift (SwiftBuddy's loader) // already auto-detects unconditionally; this mirrors that for the CLI binary // while keeping --vision/--audio as explicit overrides. - let isVision = self.vision || (!self.audio && architecture.supportsVision) - if architecture.supportsVision, !self.vision, !self.audio { + // Speculative/MTP decoding (--draft-model, --dflash, --mtp) only wires up + // BaseLanguageModel from LLMModelFactory — a VLMModelFactory container isn't + // a DualModelMTP/generateMTP participant. Auto-detection must not flip a + // speculative-decoding session into VLM mode just because the checkpoint + // carries a vision_config (e.g. mlx-community/Qwen3.5-*-4bit ships one even + // when used purely as a text draft/main pair): doing so crashed the server + // with a Trace/BPT trap on the first real generation request, caught by CI's + // speculative-decoding suite. --vision remains a valid explicit override. + let speculativeDecodingRequested = self.draftModel != nil || self.dflash || self.mtp + let autoDetectedVision = !self.audio && architecture.supportsVision + && !speculativeDecodingRequested + let isVision = self.vision || autoDetectedVision + if architecture.supportsVision, !self.vision, !self.audio, speculativeDecodingRequested { + print( + "[SwiftLM] Note: \(architecture.modelType ?? "unknown") reports vision support, but speculative/MTP decoding was requested; loading as a text-only LLM." + ) + } else if autoDetectedVision { print( "[SwiftLM] Auto-detected VLM config (\(architecture.modelType ?? "unknown")); enabling vision mode." )