spb/zyquo-mlx Public MIT
The local MLX foundry for your Mac — run, fine-tune, quantize, and ship models. Nothing leaves your machine.
Swift 93.4%
Python 3.8%
Makefile 2.2%
Shell 0.5%
1//2// TrainView.swift3// Zyquo MLX4//5// Author: Simon-Pierre Boucher6// Mail: contact@spboucher.ai7//89import Charts10import SwiftUI1112/// Train section: run configurator + run list + live run detail with loss13/// curves, throughput, log console, checkpoints, and cancel/resume.14struct TrainView: View {15 @Bindable var model: AppModel16 @State private var controller = TrainingController()17 @State private var isConfiguring = false18 @State private var selectedRunID: UUID?1920 var body: some View {21 Group {22 if let runID = selectedRunID ?? controller.activeRun?.id,23 let run = model.runs.first(where: { $0.id == runID }) ?? controller.activeRun24 {25 RunDetailView(26 run: run, model: model, controller: controller,27 back: {28 selectedRunID = nil29 Task { await model.refresh() }30 })31 } else if model.runs.isEmpty {32 EmptyStateView(33 icon: "flame",34 title: canConfigure ? "Forge Your First Fine-Tune" : "No Training Runs Yet",35 message: canConfigure36 ? "Pick a base model, a dataset, and a method — Zyquo gates configurations against this Mac's memory and streams live loss curves while it trains."37 : "Fine-tune any local model on your own data with LoRA, QLoRA, or full fine-tuning. Add a model and a dataset first, then start your first run.",38 actionLabel: canConfigure ? "New Run" : nil,39 action: canConfigure ? { isConfiguring = true } : nil)40 } else {41 runList42 }43 }44 .sheet(isPresented: $isConfiguring) {45 RunConfiguratorView(model: model, controller: controller) { run in46 selectedRunID = run.id47 Task { await model.refresh() }48 }49 }50 }5152 private var canConfigure: Bool {53 !model.models.isEmpty && !model.datasets.isEmpty54 }5556 private var runList: some View {57 ScrollView {58 LazyVStack(spacing: ZyquoTheme.spacing8) {59 HStack {60 Spacer()61 Button {62 isConfiguring = true63 } label: {64 Label("New Run", systemImage: "plus")65 }66 .disabled(!canConfigure)67 }68 ForEach(model.runs) { run in69 Button {70 selectedRunID = run.id71 } label: {72 RunRow(run: run)73 }74 .buttonStyle(.plain)75 }76 }77 .padding(ZyquoTheme.spacing20)78 }79 }80}8182// MARK: - Live training controller8384@Observable85@MainActor86final class TrainingController {87 var activeRun: TrainingRun?88 var metrics: [TrainingMetric] = []89 var logLines: [String] = []90 var isTraining = false91 var lastError: String?9293 func start(run: TrainingRun, baseModel: LocalModel, dataset: Dataset, resume: Bool = false) {94 metrics = resume ? metrics : []95 logLines.append(resume ? "Resuming (warm start from latest adapter)…" : "Starting run…")96 activeRun = run97 isTraining = true98 lastError = nil99100 Task {101 do {102 let events = try await TrainingService.shared.start(103 run: run, baseModel: baseModel, dataset: dataset, resume: resume)104 for await event in events {105 switch event {106 case .started(let model, let iterations):107 logLines.append("Training \(model) for \(iterations) iterations")108 case .metric(let metric):109 metrics.append(metric)110 if let loss = metric.trainLoss {111 logLines.append(String(112 format: "iter %d train %.3f %.0f tok/s peak %.1f GB",113 metric.iteration, loss, metric.tokensPerSecond ?? 0,114 metric.peakMemoryGB ?? 0))115 }116 if let loss = metric.valLoss {117 logLines.append(String(format: "iter %d VAL %.3f", metric.iteration, loss))118 }119 case .checkpointSaved(let name, _):120 logLines.append("checkpoint \(name)")121 case .finished:122 logLines.append("Training finished ✅")123 case .failed(let message):124 logLines.append("FAILED: \(message)")125 lastError = message126 }127 }128 } catch {129 lastError = error.localizedDescription130 logLines.append("FAILED: \(error.localizedDescription)")131 }132 isTraining = false133 }134 }135136 func cancel() {137 Task { await TrainingService.shared.cancel() }138 logLines.append("Cancelling…")139 }140}141142// MARK: - Configurator143144struct RunConfiguratorView: View {145 @Bindable var model: AppModel146 let controller: TrainingController147 let onStart: (TrainingRun) -> Void148149 @Environment(\.dismiss) private var dismiss150 @State private var baseModelID: String?151 @State private var datasetID: UUID?152 @State private var method: FineTuneMethod = .qlora153 @State private var hp = HyperParams()154 @State private var runName = ""155 @State private var startError: String?156157 private var baseModel: LocalModel? { model.models.first { $0.id == baseModelID } }158 private var dataset: Dataset? { model.datasets.first { $0.id == datasetID } }159160 private var verdict: MemoryVerdict? {161 baseModel.map { MemoryAdvisor.trainingVerdict(for: $0, method: method, params: hp) }162 }163164 var body: some View {165 VStack(spacing: 0) {166 Text("New Training Run")167 .font(ZyquoTheme.titleFont)168 .padding(.top, ZyquoTheme.spacing20)169170 Form {171 Section("Base & data") {172 Picker("Base model", selection: $baseModelID) {173 Text("Choose…").tag(String?.none)174 ForEach(model.models.filter { $0.type == .llm }) { item in175 Text("\(item.name)\(item.quantization != nil ? " · quantized" : "")")176 .tag(String?.some(item.id))177 }178 }179 Picker("Dataset", selection: $datasetID) {180 Text("Choose…").tag(UUID?.none)181 ForEach(model.datasets) { dataset in182 Text("\(dataset.name) (\(dataset.trainCount) samples)")183 .tag(UUID?.some(dataset.id))184 }185 }186 Picker("Method", selection: $method) {187 ForEach(FineTuneMethod.allCases, id: \.self) {188 Text($0.displayName)189 }190 }191 .help("QLoRA trains adapters on a quantized base — the memory-efficient default. Full fine-tuning updates every weight and needs far more memory.")192 }193194 Section("Hyperparameters") {195 Stepper("Iterations: \(hp.iterations)", value: $hp.iterations, in: 50...20000, step: 50)196 Stepper("Batch size: \(hp.batchSize)", value: $hp.batchSize, in: 1...16)197 .help("Larger batches train smoother but use more memory.")198 Stepper("LoRA rank: \(hp.rank)", value: $hp.rank, in: 2...64, step: 2)199 .help("Higher rank = more adapter capacity, slightly more memory.")200 Stepper("Layers to adapt: \(hp.numLayers)", value: $hp.numLayers, in: 4...64, step: 4)201 .help("LoRA is applied to the last N transformer layers (-1 = all).")202 TextField("Learning rate", value: $hp.learningRate, format: .number)203 Stepper("Max sequence length: \(hp.maxSeqLength)", value: $hp.maxSeqLength, in: 256...8192, step: 256)204 Toggle("Gradient checkpointing", isOn: $hp.gradCheckpoint)205 .help("Trades ~30% speed for a large activation-memory saving.")206 Toggle("Mask prompt (loss on completions only)", isOn: $hp.maskPrompt)207 }208209 if let baseModel, let verdict {210 Section("Memory") {211 LabeledContent("Estimated need") {212 Text(ByteCountFormatter.string(213 fromByteCount: MemoryAdvisor.trainingBytes(214 for: baseModel, method: method, params: hp),215 countStyle: .memory))216 }217 LabeledContent("Verdict") { VerdictBadge(verdict: verdict) }218 if verdict != .comfortable {219 ForEach(220 MemoryAdvisor.suggestions(for: baseModel, method: method, params: hp),221 id: \.self222 ) { suggestion in223 Label(suggestion, systemImage: "lightbulb")224 .font(ZyquoTheme.captionFont)225 .foregroundStyle(ZyquoTheme.textSecondary)226 }227 }228 }229 }230231 if let startError {232 Text(startError)233 .font(ZyquoTheme.captionFont)234 .foregroundStyle(ZyquoTheme.danger)235 }236 }237 .formStyle(.grouped)238239 HStack {240 Button("Cancel") { dismiss() }241 Spacer()242 Button("Start Training") { start() }243 .buttonStyle(.borderedProminent)244 .tint(ZyquoTheme.accent)245 .disabled(baseModel == nil || dataset == nil || verdict == .wontFit)246 }247 .padding(ZyquoTheme.spacing20)248 }249 .frame(width: 560, height: 640)250 }251252 private func start() {253 guard let baseModel, let dataset else { return }254 Task {255 do {256 let name = runName.isEmpty257 ? "\(baseModel.name) · \(method.displayName) · \(dataset.name)"258 : runName259 let run = try await TrainingService.shared.createRun(260 name: name, baseModel: baseModel, dataset: dataset,261 method: method, hyperParams: hp)262 controller.start(run: run, baseModel: baseModel, dataset: dataset)263 dismiss()264 onStart(run)265 } catch {266 startError = error.localizedDescription267 }268 }269 }270}271272// MARK: - Run detail273274struct RunDetailView: View {275 let run: TrainingRun276 @Bindable var model: AppModel277 @Bindable var controller: TrainingController278 let back: () -> Void279280 @State private var historicMetrics: [TrainingMetric] = []281282 private var isActive: Bool { controller.activeRun?.id == run.id && controller.isTraining }283 private var metrics: [TrainingMetric] {284 isActive || controller.activeRun?.id == run.id ? controller.metrics : historicMetrics285 }286287 var body: some View {288 ScrollView {289 VStack(alignment: .leading, spacing: ZyquoTheme.spacing16) {290 HStack {291 Button {292 back()293 } label: {294 Label("Runs", systemImage: "chevron.left")295 }296 .buttonStyle(.plain)297 .foregroundStyle(ZyquoTheme.accent)298299 Spacer()300301 if isActive {302 Button("Cancel Run") { controller.cancel() }303 .tint(ZyquoTheme.danger)304 } else if run.state == .cancelled || run.state == .failed {305 Button("Resume (warm start)") { resume() }306 .buttonStyle(.borderedProminent)307 .tint(ZyquoTheme.accent)308 }309 }310311 VStack(alignment: .leading, spacing: ZyquoTheme.spacing4) {312 HStack {313 Text(run.name)314 .font(ZyquoTheme.titleFont)315 StatusPill(316 text: isActive ? "Running" : run.state.rawValue.capitalized,317 color: isActive ? ZyquoTheme.accent : ZyquoTheme.textSecondary)318 }319 Text("\(run.method.displayName) · \(run.hyperParams.iterations) iterations · rank \(run.hyperParams.rank) · batch \(run.hyperParams.batchSize)")320 .font(ZyquoTheme.captionFont)321 .foregroundStyle(ZyquoTheme.textSecondary)322 }323324 // Loss chart325 lossChart326327 // Throughput + memory strip328 if let latest = metrics.last(where: { $0.tokensPerSecond != nil }) {329 HStack(spacing: ZyquoTheme.spacing20) {330 stat("Tokens/sec", String(format: "%.0f", latest.tokensPerSecond ?? 0))331 stat("Iteration", "\(metrics.last?.iteration ?? 0)/\(run.hyperParams.iterations)")332 stat("Peak memory", String(format: "%.1f GB", latest.peakMemoryGB ?? 0))333 if let tokens = latest.trainedTokens {334 stat("Trained tokens", tokens.formatted())335 }336 }337 }338339 // Log console340 VStack(alignment: .leading, spacing: ZyquoTheme.spacing8) {341 Text("Console")342 .font(ZyquoTheme.headlineFont)343 ScrollViewReader { proxy in344 ScrollView {345 VStack(alignment: .leading, spacing: 2) {346 ForEach(Array(consoleLines.enumerated()), id: \.offset) { index, line in347 Text(line)348 .font(ZyquoTheme.monoSmallFont)349 .foregroundStyle(ZyquoTheme.textSecondary)350 .id(index)351 }352 }353 .frame(maxWidth: .infinity, alignment: .leading)354 .padding(ZyquoTheme.spacing12)355 }356 .frame(height: 160)357 .background(ZyquoTheme.surfaceSecondary)358 .clipShape(RoundedRectangle(cornerRadius: ZyquoTheme.radiusSmall))359 .onChange(of: consoleLines.count) {360 proxy.scrollTo(consoleLines.count - 1, anchor: .bottom)361 }362 }363 }364365 // Checkpoints366 if !checkpoints.isEmpty {367 VStack(alignment: .leading, spacing: ZyquoTheme.spacing8) {368 Text("Checkpoints")369 .font(ZyquoTheme.headlineFont)370 ForEach(checkpoints) { checkpoint in371 HStack {372 Image(systemName: "externaldrive")373 .foregroundStyle(ZyquoTheme.slate)374 Text(checkpoint.fileName)375 .font(ZyquoTheme.monoSmallFont)376 Spacer()377 Text("iter \(checkpoint.iteration)")378 .font(ZyquoTheme.captionFont)379 .foregroundStyle(ZyquoTheme.textSecondary)380 }381 .padding(ZyquoTheme.spacing8)382 .zyquoCard()383 }384 }385 }386 }387 .padding(ZyquoTheme.spacing20)388 }389 .task(id: run.id) {390 if controller.activeRun?.id != run.id {391 historicMetrics = await RunStore.shared.metricsHistory(for: run)392 }393 }394 }395396 private var consoleLines: [String] {397 controller.activeRun?.id == run.id398 ? controller.logLines399 : metrics.compactMap { metric in400 if let loss = metric.trainLoss {401 return String(format: "iter %d train %.3f", metric.iteration, loss)402 }403 if let loss = metric.valLoss {404 return String(format: "iter %d VAL %.3f", metric.iteration, loss)405 }406 return nil407 }408 }409410 private var checkpoints: [Checkpoint] {411 run.checkpoints412 }413414 private var lossChart: some View {415 Chart {416 ForEach(Array(metrics.enumerated()), id: \.offset) { _, metric in417 if let loss = metric.trainLoss {418 LineMark(419 x: .value("Iteration", metric.iteration),420 y: .value("Train loss", loss),421 series: .value("Series", "Train"))422 .foregroundStyle(ZyquoTheme.chartTrain)423 .interpolationMethod(.monotone)424 }425 }426 ForEach(Array(metrics.enumerated()), id: \.offset) { _, metric in427 if let loss = metric.valLoss {428 LineMark(429 x: .value("Iteration", metric.iteration),430 y: .value("Val loss", loss),431 series: .value("Series", "Val"))432 .foregroundStyle(ZyquoTheme.chartVal)433 PointMark(434 x: .value("Iteration", metric.iteration),435 y: .value("Val loss", loss))436 .foregroundStyle(ZyquoTheme.chartVal)437 .symbolSize(30)438 }439 }440 }441 .chartLegend(.visible)442 .chartForegroundStyleScale([443 "Train": ZyquoTheme.chartTrain, "Val": ZyquoTheme.chartVal,444 ])445 .frame(height: 220)446 .padding(ZyquoTheme.spacing12)447 .zyquoCard()448 .overlay {449 if metrics.isEmpty {450 Text(isActive ? "Waiting for the first report…" : "No metrics recorded")451 .font(ZyquoTheme.captionFont)452 .foregroundStyle(ZyquoTheme.textTertiary)453 }454 }455 }456457 private func stat(_ label: String, _ value: String) -> some View {458 VStack(alignment: .leading, spacing: ZyquoTheme.spacing2) {459 Text(label)460 .font(ZyquoTheme.captionFont)461 .foregroundStyle(ZyquoTheme.textTertiary)462 Text(value)463 .font(ZyquoTheme.monoFont)464 .foregroundStyle(ZyquoTheme.textPrimary)465 }466 }467468 private func resume() {469 guard470 let baseModel = model.models.first(where: { $0.id == run.baseModelID }),471 let dataset = model.datasets.first(where: { $0.id == run.datasetID })472 else { return }473 controller.start(run: run, baseModel: baseModel, dataset: dataset, resume: true)474 }475}476477struct RunRow: View {478 let run: TrainingRun479480 var body: some View {481 HStack(spacing: ZyquoTheme.spacing12) {482 Image(systemName: "flame")483 .font(.system(size: 18, weight: .light))484 .foregroundStyle(stateColor)485 .frame(width: 32, height: 32)486 .background(stateColor.opacity(0.1))487 .clipShape(RoundedRectangle(cornerRadius: ZyquoTheme.radiusSmall))488489 VStack(alignment: .leading, spacing: ZyquoTheme.spacing2) {490 Text(run.name)491 .font(ZyquoTheme.bodyFont.weight(.medium))492 .foregroundStyle(ZyquoTheme.textPrimary)493 HStack(spacing: ZyquoTheme.spacing8) {494 Text(run.method.displayName)495 Text("\(run.completedIterations)/\(run.hyperParams.iterations) iters")496 if !run.checkpoints.isEmpty {497 Text("\(run.checkpoints.count) checkpoints")498 }499 }500 .font(ZyquoTheme.captionFont)501 .foregroundStyle(ZyquoTheme.textSecondary)502 }503504 Spacer()505506 StatusPill(text: run.state.rawValue.capitalized, color: stateColor)507 Text(run.createdAt.formatted(date: .abbreviated, time: .shortened))508 .font(ZyquoTheme.captionFont)509 .foregroundStyle(ZyquoTheme.textTertiary)510 }511 .padding(ZyquoTheme.spacing12)512 .zyquoCard()513 }514515 private var stateColor: Color {516 switch run.state {517 case .running: ZyquoTheme.accent518 case .completed: ZyquoTheme.success519 case .paused, .configured: ZyquoTheme.slate520 case .cancelled: ZyquoTheme.warning521 case .failed: ZyquoTheme.danger522 }523 }524}525