SPB Git

spb/zyquo-mlx Public MIT

The local MLX foundry for your Mac — run, fine-tune, quantize, and ship models. Nothing leaves your machine.

macosswiftaimlxfine-tuning
Swift 93.4% Python 3.8% Makefile 2.2% Shell 0.5%

phase3: checkpoint — gate green (QLoRA loss 4.62→0.10, fuse w/ auto-dequantize, style transfer verified, Swift quant 0.1%-accurate preview); callback + requantization landmines documented

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
simon-pierre boucher committed 11 days ago (Jul 30, 2026) parent 5979cc3

Showing 6 changed files with +137 and −23

modified Sources/ZyquoMLX/App/CLIFoundry.swift +54 −2
@@ -18,7 +18,7 @@ extension CLI {
18 18 static var shouldRunFoundry: Bool {
19 19 let args = CommandLine.arguments
20 20 return args.contains("--train") || args.contains("--fuse")
21 − || args.contains("--validate-dataset")
21 + || args.contains("--quantize") || args.contains("--validate-dataset")
22 22 }
23 23
24 24 static func runFoundry() async -> Int32 {
@@ -32,6 +32,10 @@ extension CLI {
32 32 try await fuse(args: args)
33 33 return 0
34 34 }
35 + if args.contains("--quantize") {
36 + try await quantize(args: args)
37 + return 0
38 + }
35 39 if let path = value(after: "--validate-dataset", in: args) {
36 40 try await validateDataset(path: path)
37 41 return 0
@@ -133,9 +137,19 @@ extension CLI {
133 137 directory: URL(fileURLWithPath: (modelPath as NSString).expandingTildeInPath))
134 138 let adapter = URL(fileURLWithPath: (adapterPath as NSString).expandingTildeInPath)
135 139
140 + // Fusing into a quantized base re-quantizes the merged weights, which
141 + // rounds away small LoRA deltas (verified empirically — the adapter's
142 + // behavior vanished). Default to de-quantizing for quantized bases;
143 + // --no-dequantize opts back into the lossy compact form.
144 + let dequantize =
145 + args.contains("--dequantize")
146 + || (model.quantization != nil && !args.contains("--no-dequantize"))
147 + if dequantize && model.quantization != nil {
148 + print("note: de-quantizing while fusing (quantized base) — re-quantize afterwards in Convert if needed")
149 + }
136 150 let events = try await ConversionService.shared.fuse(
137 151 baseModel: model, adapterDirectory: adapter, outputName: outName,
138 − dequantize: args.contains("--dequantize"))
152 + dequantize: dequantize)
139 153 for await event in events {
140 154 switch event {
141 155 case .stage(let stage, _):
@@ -148,6 +162,44 @@ extension CLI {
148 162 }
149 163 }
150 164
165 + // MARK: - Quantize (Swift-native, docs/MLX-RESEARCH.md §2)
166 +
167 + private static func quantize(args: [String]) async throws {
168 + guard let modelPath = value(after: "--quantize", in: args),
169 + let outName = value(after: "--out", in: args)
170 + else {
171 + throw CLIError.usage("--quantize <model-dir> --out <name> [--bits N] [--group-size G]")
172 + }
173 + let model = try await ModelStore.shared.describe(
174 + directory: URL(fileURLWithPath: (modelPath as NSString).expandingTildeInPath))
175 +
176 + var config = QuantConfig()
177 + config.bits = value(after: "--bits", in: args).flatMap(Int.init) ?? 4
178 + config.groupSize = value(after: "--group-size", in: args).flatMap(Int.init) ?? 64
179 +
180 + if let params = model.parameterCount {
181 + let predicted = config.predictedWeightBytes(parameterCount: params)
182 + print("quantizing \(model.name) → \(config.label)")
183 + print("size preview: \(ByteCountFormatter.string(fromByteCount: model.weightsSize, countStyle: .file)) → ~\(ByteCountFormatter.string(fromByteCount: predicted, countStyle: .file))")
184 + }
185 +
186 + let events = await ConversionService.shared.quantize(
187 + model: model, config: config, outputName: outName)
188 + for await event in events {
189 + switch event {
190 + case .stage(let stage, let fraction):
191 + let pct = fraction.map { String(format: " %.0f%%", $0 * 100) } ?? ""
192 + print("stage: \(stage)\(pct)")
193 + case .finished(let output):
194 + let result = try await ModelStore.shared.describe(directory: output)
195 + print("quantized model → \(output.path)")
196 + print("actual size: \(ByteCountFormatter.string(fromByteCount: result.weightsSize, countStyle: .file)) ✅")
197 + case .failed(let message):
198 + throw CLIError.usage("quantize failed: \(message)")
199 + }
200 + }
201 + }
202 +
151 203 // MARK: - Dataset validation
152 204
153 205 private static func validateDataset(path: String) async throws {
modified Sources/ZyquoMLX/PyBridge/PythonEnvironment.swift +31 −3
@@ -63,7 +63,9 @@ actor PythonEnvironment {
63 63 .map { URL(fileURLWithPath: $0) }
64 64 }
65 65
66 − /// Create the venv and install core pins. Idempotent.
66 + /// Create the venv and install core pins. Idempotent: a pre-existing
67 + /// healthy venv (verified by importing the pinned packages) is adopted;
68 + /// a broken one is recreated in place (`uv venv --clear`).
67 69 func provision(progress: @Sendable (String) -> Void = { _ in }) async throws {
68 70 if isProvisioned { return }
69 71 guard let uv = Self.findUV() else { throw EnvironmentError.uvNotFound }
@@ -71,22 +73,48 @@ actor PythonEnvironment {
71 73 try FileManager.default.createDirectory(
72 74 at: PersistenceService.pythonDirectory, withIntermediateDirectories: true)
73 75
76 + if FileManager.default.fileExists(atPath: pythonURL.path), verifyImports() {
77 + progress("Adopting existing Python environment.")
78 + try writeMarker()
79 + return
80 + }
81 +
74 82 progress("Installing Python \(Self.pythonVersion)…")
75 83 try runUV(uv, ["python", "install", Self.pythonVersion])
76 84
77 85 progress("Creating environment…")
78 − try runUV(uv, ["venv", venvURL.path, "--python", Self.pythonVersion])
86 + try runUV(uv, ["venv", venvURL.path, "--clear", "--python", Self.pythonVersion])
79 87
80 88 progress("Installing MLX packages…")
81 89 try runUV(uv, ["pip", "install", "--python", pythonURL.path] + Self.corePins)
82 90
91 + try writeMarker()
92 + progress("Python environment ready.")
93 + }
94 +
95 + /// True when the pinned core packages import cleanly.
96 + private func verifyImports() -> Bool {
97 + let process = Process()
98 + process.executableURL = pythonURL
99 + process.arguments = ["-c", "import mlx_lm, yaml"]
100 + process.standardOutput = Pipe()
101 + process.standardError = Pipe()
102 + do {
103 + try process.run()
104 + process.waitUntilExit()
105 + return process.terminationStatus == 0
106 + } catch {
107 + return false
108 + }
109 + }
110 +
111 + private func writeMarker() throws {
83 112 let marker: [String: String] = [
84 113 "python": Self.pythonVersion,
85 114 "pins": Self.corePins.joined(separator: " "),
86 115 "provisionedAt": ISO8601DateFormatter().string(from: .now),
87 116 ]
88 117 try PersistenceService.saveJSON(marker, to: markerURL)
89 − progress("Python environment ready.")
90 118 }
91 119
92 120 /// Install additional pinned packages (e.g. mlx-vlm when VLM Python
modified Sources/ZyquoMLX/PyBridge/scripts/zyquo_fuse.py +1 −1
@@ -44,7 +44,7 @@ def main():
44 44 "--save-path", args.save_path,
45 45 ]
46 46 if args.dequantize:
47 − argv.append("--de-quantize")
47 + argv.append("--dequantize") # exact flag name in mlx-lm 0.31.3
48 48
49 49 old_argv = sys.argv
50 50 sys.argv = ["mlx_lm.fuse"] + argv
modified Sources/ZyquoMLX/PyBridge/scripts/zyquo_train.py +14 −1
@@ -46,8 +46,11 @@ def main():
46 46 args = parser.parse_args()
47 47
48 48 try:
49 + import numpy as np
49 50 import yaml
50 51 from mlx_lm import lora
52 + from mlx_lm.tuner.datasets import load_dataset
53 + from mlx_lm.utils import load
51 54
52 55 with open(args.config) as f:
53 56 config = yaml.safe_load(f)
@@ -68,7 +71,17 @@ def main():
68 71 "adapter_path": ns.adapter_path,
69 72 })
70 73
71 − lora.run(ns, training_callback=ZyquoCallback())
74 + # NOTE: we deliberately do NOT call lora.run() — in mlx-lm 0.31.3 it
75 + # overwrites the training_callback argument with
76 + # get_reporting_callbacks(args.report_to) (None here), silently
77 + # discarding ours. Replicate run()'s exact flow instead.
78 + np.random.seed(ns.seed)
79 + emit({"event": "stage", "stage": "loading_model"})
80 + model, tokenizer = load(ns.model, tokenizer_config={"trust_remote_code": True})
81 + emit({"event": "stage", "stage": "loading_datasets"})
82 + train_set, valid_set, _test_set = load_dataset(ns, tokenizer)
83 + emit({"event": "stage", "stage": "training"})
84 + lora.train_model(ns, model, train_set, valid_set, ZyquoCallback())
72 85 emit({"event": "done"})
73 86 except KeyboardInterrupt:
74 87 emit({"event": "error", "message": "cancelled"})
modified docs/PLAN.md +20 −11
@@ -101,17 +101,26 @@ sources (no guessed names — `#huggingFaceTokenizerLoader` macro requires
101 101 through `perform(nonSendable:)`). CLI POC proves the two-model-type gate with
102 102 real downloads in the app's Models library. Build is arm64-only via
103 103 `--arch arm64` (swiftbuild otherwise builds universal and x86_64 fails).
104 −## Phase 3 — Training, Quantization & Conversion (in progress)
105 −
106 −- [ ] `PyBridge/PythonRunner` — Process wrapper over the venv, JSON-lines progress protocol, cancellation
107 −- [ ] `PyBridge/PythonEnvironment` — uv-provisioned pinned venv in App Support (status/repair; mlx-lm now, vlm/whisper/audio on demand)
108 −- [ ] `PyBridge/scripts/` — zyquo_train.py (TrainingCallback → JSON), zyquo_fuse.py, zyquo_convert.py
109 −- [ ] `Data/DatasetService` + `DatasetFormats` — import/validate JSONL (chat/completions/text), auto split, malformed-row report, token stats, preview
110 −- [ ] `Training/TrainingService` + `RunStore` + `MetricsStream` — cancellable/resumable runs, persisted state, checkpoints, live metrics
111 −- [ ] `Convert/ConversionService` + `QuantConfig` — Swift-native affine quant for safetensors LLMs; Python bridge for fuse/dequantize
112 −- [ ] CLI: `--train`, `--fuse`, `--validate-dataset` to drive the services (UI in Phase 6)
113 −- [ ] PHASE GATE: real LoRA/QLoRA fine-tune on a tiny dataset → live loss ↓, adapter saved, fused, fused model generates differently
114 −- [ ] Phase checkpoint: build green, gate verified, summary
104 +## Phase 3 — Training, Quantization & Conversion ✅ (completed 2026-07-30)
105 +
106 +- [x] `PyBridge/PythonRunner` — Process wrapper over the venv, JSON-lines progress protocol, cancellation
107 +- [x] `PyBridge/PythonEnvironment` — uv-provisioned pinned venv in App Support (idempotent adopt/verify/repair; mlx-lm 0.31.3 now, vlm/whisper/audio on demand)
108 +- [x] `PyBridge/scripts/` — zyquo_train.py (TrainingCallback → JSON; bypasses `lora.run()` which drops the callback in 0.31.3), zyquo_fuse.py, zyquo_convert.py
109 +- [x] `Data/DatasetService` + `DatasetFormats` — import/validate JSONL (chat/completions/text), deterministic split, malformed-row report with fixes, token stats, preview
110 +- [x] `Training/TrainingService` + `RunStore` + `MetricsStream` — cancellable runs, persisted state + metrics.jsonl, numbered checkpoints, warm-start resume, MemoryAdvisor gating
111 +- [x] `Convert/ConversionService` + `QuantConfig` — Swift-native affine/mxfp4 quant (size preview accurate to 0.1%); Python bridge for fuse/convert
112 +- [x] CLI: `--train`, `--fuse`, `--quantize`, `--validate-dataset` (UI in Phase 6)
113 +- [x] PHASE GATE: QLoRA on Qwen3-0.6B-4bit, 48-row chat dataset — live loss 4.62→0.10 (val 6.31→0.14) at ~2,300 tok/s, checkpoints at 40/80/120, adapter fused (auto-dequantize), fused model answers in the trained "⚒ From the forge:" style
114 +- [x] Phase checkpoint: build green, 0 warnings
115 +
116 +**Phase 3 summary:** Full foundry core working end-to-end from the app. Two
117 +upstream landmines found and handled: (1) mlx-lm 0.31.3's `lora.run()`
118 +silently discards the caller's TrainingCallback — our driver replicates
119 +`run()`'s flow and calls `train_model` directly; (2) fusing into a quantized
120 +base re-quantizes and rounds away small LoRA deltas — Zyquo defaults to
121 +dequantize-on-fuse, recommends adapter-attached inference, and warns on
122 +re-quantization (documented in TRAINING-RESEARCH.md §2.2). Swift-native
123 +quantization verified: 1.19 GB fp16 → 335.5 MB 4-bit, predicted 335.3 MB.
115 124 ## Phase 4 — Design System & UI (not started)
116 125 ## Phase 5 — App Icon (not started)
117 126 ## Phase 6 — Features (not started)
modified docs/TRAINING-RESEARCH.md +17 −5
@@ -128,11 +128,23 @@ quantized base with `mlx_lm.convert --hf-path <repo> -q` or use any
128 128
129 129 Mechanics: loads model + adapters, calls `.fuse(dequantize:)` on every capable
130 130 module, saves a complete standalone MLX model dir (config.json, sharded
131 −safetensors + index, tokenizer files) — directly loadable afterwards. Fusing
132 −into a **quantized** base without `--dequantize` re-quantizes fused weights
133 −(small quality loss vs. keeping the adapter separate) — surface this trade-off
134 −in the Convert UI. (LORA.md mentions `--hf-path` for fuse but it is absent from
135 −current argparse — docs stale; UNVERIFIED if intentional.)
131 +safetensors + index, tokenizer files) — directly loadable afterwards.
132 +(LORA.md mentions `--hf-path` for fuse but it is absent from current argparse —
133 +docs stale; UNVERIFIED if intentional. The pinned 0.31.3 flag is `--dequantize`,
134 +not `--de-quantize`.)
135 +
136 +**⚠️ EMPIRICALLY VERIFIED (2026-07-30, this Mac):** fusing a QLoRA adapter into
137 +a **quantized** base without `--dequantize` re-quantizes the merged weights and
138 +**rounds small LoRA deltas away entirely** — our rank-8 adapter's behavior
139 +vanished from the fused 4-bit model while `--adapter-path` inference kept it,
140 +and a `--dequantize` fuse preserved it verbatim. Re-quantizing the fp16 fused
141 +model to 4-bit wiped the behavior again — for lightly-trained adapters the
142 +deltas are simply below the 4-bit quantization step wherever they land.
143 +Zyquo therefore: (1) **defaults to de-quantize when fusing onto a quantized
144 +base**, (2) recommends **adapter-attached inference** (Python `--adapter-path`
145 +/ Swift `LoRAContainer.load(into:)`) as the lossless default for QLoRA
146 +results, and (3) warns before re-quantizing a fused model that adapter effects
147 +may not survive unless training was substantial. The UI must explain this.
136 148
137 149 ---
138 150
139 151