diff --git a/swift/Sources/QwispCore/QwispModel.swift b/swift/Sources/QwispCore/QwispModel.swift index c43b8a5..6646b16 100644 --- a/swift/Sources/QwispCore/QwispModel.swift +++ b/swift/Sources/QwispCore/QwispModel.swift @@ -9,6 +9,19 @@ import MLXFast public final class WeightStore { public private(set) var arrays: [String: MLXArray] = [:] + /// The checkpoint ships a vision encoder. qwisp is a text-only engine and never reads a + /// single one of those tensors — but `residentNonExperts()` filtered only on `.switch_mlp.`, + /// so all 852 MB of `vision_tower.*` were handed to MLX.eval and made resident on every run, + /// on every tier. Measured: active 7,071 MB (8GB tier) = arena 4,320 + language_model + /// non-expert 1,325 + vision 852 + scratch 574; the 16GB tier's 11,391 MB decomposes the + /// same way. Dropping them is free — there is no quality knob involved. + /// + /// Filtered at load rather than at eval: keeping the MLXArray alive holds its mmap slice, + /// so excluding it from eval alone would not have returned the bytes. + public static func isUsedByEngine(_ tensorName: String) -> Bool { + !(tensorName.hasPrefix("vision_tower") || tensorName.contains("visual")) + } + public init(modelDir: String) throws { let dir = URL(fileURLWithPath: modelDir) let idxURL = dir.appendingPathComponent("model.safetensors.index.json") @@ -16,9 +29,17 @@ public final class WeightStore { let top = try JSONSerialization.jsonObject(with: data) as? [String: Any] ?? [:] let wm = (top["weight_map"] as? [String: String]) ?? [:] let shards = Set(wm.values) + var dropped = 0 for shard in shards.sorted() { let m = try loadArrays(url: dir.appendingPathComponent(shard)) - for (k, v) in m { arrays[k] = v } + for (k, v) in m { + guard WeightStore.isUsedByEngine(k) else { dropped += 1; continue } + arrays[k] = v + } + } + if dropped > 0 { + FileHandle.standardError.write(Data( + "[qwisp] skipped \(dropped) vision-encoder tensors (text-only engine)\n".utf8)) } } diff --git a/swift/Sources/QwispCore/SeedlessVerifyTests.swift b/swift/Sources/QwispCore/SeedlessVerifyTests.swift index b109b2e..12f6ae0 100644 --- a/swift/Sources/QwispCore/SeedlessVerifyTests.swift +++ b/swift/Sources/QwispCore/SeedlessVerifyTests.swift @@ -49,7 +49,7 @@ public enum SeedlessVerifyTests { MLXRandom.seed(UInt64(42)) var lines: [String] = [] var passed = 0 - let total = 99 + let total = 100 // Nested runner: records result and increments counter func run(_ name: String, body: () -> (Bool, String)) { @@ -6282,6 +6282,34 @@ public enum SeedlessVerifyTests { return (true, "ok") } + // WRITE-LOCKED (guarded by total = 100). Do not weaken/skip/delete. + // The checkpoint ships a vision encoder qwisp never reads. residentNonExperts() + // filtered only on `.switch_mlp.`, so 852 MB of vision_tower.* were eval'd resident + // on every run of every tier — on the 8GB tier that is more than the entire deficit + // that made the tier not fit. This asserts the predicate keeps engine tensors and + // drops vision ones; a regression here silently costs 852 MB again. + run("100 weight-store drops the unused vision encoder") { + let keep = [ + "language_model.model.embed_tokens.weight", + "language_model.lm_head.scales", + "language_model.model.layers.0.mlp.switch_mlp.up_proj.weight", + "language_model.model.layers.7.linear_attn.in_proj_qkv.weight", + "language_model.model.layers.3.mlp.gate.weight", + ] + let drop = [ + "vision_tower.merger.linear_fc1.weight", + "vision_tower.blocks.0.mlp.linear_fc2.scales", + "model.visual.patch_embed.proj.weight", + ] + for k in keep where !WeightStore.isUsedByEngine(k) { + return (false, "dropped an engine tensor: \(k)") + } + for k in drop where WeightStore.isUsedByEngine(k) { + return (false, "kept a vision tensor: \(k)") + } + return (true, "ok") + } + // ── Summary ─────────────────────────────────────────────────────── return lines.joined(separator: "\n") + "\nRAWTESTS \(passed)/\(total)" }