Fix parameter count for quantized models (#137)

* fix parameter count

* cleanup

* try extension module

Author: awni

Applications/LLMEval/ContentView.swift

@@ -193,8 +193,13 @@ class LLMEvaluator {
                         "Downloading \(modelConfiguration.name): \(Int(progress.fractionCompleted * 100))%"
                 }
             }
+            let numParams = await modelContainer.perform {
+                [] model, _ in
+                return model.numParameters()
+            }
+
             self.modelInfo =
-                "Loaded \(modelConfiguration.id).  Weights: \(MLX.GPU.activeMemory / 1024 / 1024)M"
+                "Loaded \(modelConfiguration.id).  Weights: \(numParams / (1024*1024))M"
             loadState = .loaded(modelContainer)
             return modelContainer
 

Libraries/LLM/LLMModel.swift

@@ -57,6 +57,27 @@ public actor ModelContainer {
     }
 }
 
+extension Module {
+
+    /// Compute the number of parameters in a possibly quantized model
+    public func numParameters() -> Int {
+        return leafModules().flattenedValues().map {
+            mod -> Int in
+            if let qlin = mod as? QuantizedLinear {
+                return qlin.scales.size * qlin.groupSize
+            } else if let qemb = mod as? QuantizedEmbedding {
+                return qemb.scales.size * qemb.groupSize
+            } else {
+                return mod.parameters().flattenedValues().reduce(
+                    0,
+                    {
+                        $0 + $1.size
+                    })
+            }
+        }.reduce(0, +)
+    }
+}
+
 /// Interface for all LLM Models
 public protocol LLMModel: Module {
 

Tools/llm-tool/LoraCommands.swift

@@ -58,8 +58,7 @@ struct LoRAModelArguments: ParsableArguments, Sendable {
     }
 
     func describe(model: Module) {
-        let totalParameterCount = model.parameters()
-            .flattenedValues().map { $0.size }.reduce(0, +)
+        let totalParameterCount = model.numParameters()
         let trainableParameterCount = model.trainableParameters()
             .flattenedValues().map { $0.size }.reduce(0, +)