spb/zyquo-mlx Public MIT

1//2//  MemoryAdvisor.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation10import MLX1112/// RAM feasibility verdicts for inference and training on *this* Mac.13///14/// Formulas per docs/MODELS.md §3 and docs/TRAINING-RESEARCH.md §6:15/// - weights: bits-per-param + group-scale overhead (4-bit gs64 affine16///   = 4.5 bits ≈ 0.5625 B/param, confirmed against real repo sizes)17/// - inference ≈ weights + KV cache + ~20% overhead18/// - LoRA (bf16 base) ≈ 2 B/param + activations; QLoRA ≈ quantized weights19///   + activations; full FT ≈ 8 B/param (Adam) + activations20/// - the GPU working-set ceiling comes from the device, not a constant:21///   `GPU.deviceInfo().maxRecommendedWorkingSetSize`22enum MemoryVerdict: String, Codable, Sendable {23    case comfortable24    case tight25    case wontFit2627    var displayName: String {28        switch self {29        case .comfortable: "Fits"30        case .tight: "Tight"31        case .wontFit: "Won't fit"32        }33    }34}3536enum MemoryAdvisor {3738    /// Physical unified memory of this Mac.39    static var physicalMemory: Int64 { Int64(ProcessInfo.processInfo.physicalMemory) }4041    /// The GPU-visible working-set ceiling reported by Metal via MLX.42    static var recommendedWorkingSet: Int64 {43        let info = GPU.deviceInfo()44        let value = info.maxRecommendedWorkingSetSize45        return value > 0 ? Int64(value) : physicalMemory * 3 / 446    }4748    /// Effective bytes per parameter for a given quantization49    /// (group-scale overhead included: bits + 32/groupSize extra bits).50    static func bytesPerParameter(quantization: QuantizationInfo?) -> Double {51        guard let q = quantization else { return 2.0 } // bf1652        let effectiveBits = Double(q.bits) + 32.0 / Double(q.groupSize)53        return effectiveBits / 8.054    }5556    /// Estimated KV-cache bytes for a typical mid-size GQA model at the given57    /// context. Uses conservative defaults when architecture details are58    /// unknown (calibrated in Phase 7).59    static func estimatedKVCache(parameterCount: Int64?, contextLength: Int) -> Int64 {60        // Scale a measured anchor: ~1.2 GB at 8k ctx for an 8B GQA model61        // (docs/MLX-RESEARCH.md §7), linear in context and sublinear in params.62        let params = Double(parameterCount ?? 8_000_000_000)63        let anchor = 1.2 * 1_073_741_824.064        let scale = (params / 8_000_000_000.0).squareRoot()65        return Int64(anchor * scale * Double(contextLength) / 8192.0)66    }6768    // MARK: - Inference6970    static func inferenceBytes(for model: LocalModel, contextLength: Int = 8192) -> Int64 {71        let weights = Double(model.weightsSize)72        let kv = Double(estimatedKVCache(parameterCount: model.parameterCount, contextLength: contextLength))73        return Int64(weights * 1.2 + kv)74    }7576    static func inferenceVerdict(for model: LocalModel, contextLength: Int = 8192) -> MemoryVerdict {77        verdict(needed: inferenceBytes(for: model, contextLength: contextLength))78    }7980    // MARK: - Training8182    static func trainingBytes(for model: LocalModel, method: FineTuneMethod, params: HyperParams) -> Int64 {83        let count = Double(model.parameterCount ?? 0)84        guard count > 0 else { return .max }8586        // Activation estimate: batch × seq drives it; grad-checkpoint ~halves it.87        var activations = Double(params.batchSize) * Double(params.maxSeqLength) / (4.0 * 2048.0)88            * 2.0 * 1_073_741_824.089        if params.gradCheckpoint { activations *= 0.55 }9091        let base: Double92        switch method {93        case .lora, .dora:94            base = count * 2.0 // bf16 base held in memory95        case .qlora:96            base = Double(model.weightsSize) // quantized base stays quantized97        case .full:98            base = count * 8.0 // weights + grads + 2 Adam moments (bf16)99        }100        // Adapter weights + their optimizer state are negligible (MB-scale).101        return Int64(base + activations)102    }103104    static func trainingVerdict(for model: LocalModel, method: FineTuneMethod, params: HyperParams) -> MemoryVerdict {105        verdict(needed: trainingBytes(for: model, method: method, params: params))106    }107108    /// Suggested remedies when a config won't fit (charter 3.B).109    static func suggestions(for model: LocalModel, method: FineTuneMethod, params: HyperParams) -> [String] {110        var out: [String] = []111        if method == .lora || method == .dora {112            out.append("Use QLoRA (train on a 4-bit base) — largest single saving")113        }114        if params.batchSize > 1 { out.append("Reduce batch size (\(params.batchSize) → \(max(1, params.batchSize / 2))) and use gradient accumulation") }115        if params.numLayers == -1 || params.numLayers > 8 { out.append("Adapt fewer layers (--num-layers 8 or 4)") }116        if !params.gradCheckpoint { out.append("Enable gradient checkpointing") }117        if params.maxSeqLength > 1024 { out.append("Lower max sequence length or pre-split long samples") }118        out.append("Choose a smaller base model")119        return out120    }121122    // MARK: - Verdict core123124    private static func verdict(needed: Int64) -> MemoryVerdict {125        let ceiling = recommendedWorkingSet126        if needed <= Int64(Double(ceiling) * 0.75) { return .comfortable }127        if needed <= ceiling { return .tight }128        return .wontFit129    }130}131