diff --git a/Source/MLXNN/Quantized.swift b/Source/MLXNN/Quantized.swift index 076e91ca..ca5913ee 100644 --- a/Source/MLXNN/Quantized.swift +++ b/Source/MLXNN/Quantized.swift @@ -27,6 +27,12 @@ public protocol Quantized: Module { var mode: QuantizationMode { get } } +private func applyNVFP4GlobalScale(_ value: MLXArray, globalScale: MLXArray?) -> MLXArray { + guard let globalScale else { return value } + // NVFP4 encodes its E4M3 group scales with (448 * 6) / globalScale. + return (value * (globalScale / (448 * 6))).asType(value.dtype) +} + /// Quantize any ``Quantizable`` layer that is not already quantized. public func quantizeSingle( layer: Module, groupSize: Int = 64, bits: Int = 4, mode: QuantizationMode = .affine @@ -159,6 +165,7 @@ open class QuantizedEmbedding: Embedding, Quantized { public let mode: QuantizationMode public let scales: MLXArray public let biases: MLXArray? + @ParameterInfo(key: "global_scale") public private(set) var globalScale: MLXArray? open override var shape: (Int, Int) { let (embeddingCount, dimensions) = super.shape @@ -184,14 +191,16 @@ open class QuantizedEmbedding: Embedding, Quantized { public init( weight: MLXArray, groupSize: Int = 64, bits: Int = 4, - mode: QuantizationMode = .affine + mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil ) { self.groupSize = groupSize self.bits = bits self.mode = mode + self._globalScale.wrappedValue = globalScale let (quantizedWeight, scales, biases) = MLX.quantized( - weight, groupSize: groupSize, bits: bits, mode: mode) + weight, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale) self.scales = scales self.biases = biases @@ -201,19 +210,36 @@ open class QuantizedEmbedding: Embedding, Quantized { self.freeze() } + /// Initializer meant for subclasses to provide arrays directly. + public init( + weight: MLXArray, scales: MLXArray, biases: MLXArray?, + groupSize: Int, bits: Int, + mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil + ) { + self.groupSize = groupSize + self.bits = bits + self.mode = mode + self.scales = scales + self.biases = biases + self._globalScale.wrappedValue = globalScale + super.init(weight: weight) + } + open override func callAsFunction(_ x: MLXArray) -> MLXArray { let s = x.shape let x = x.flattened() let out = dequantized( weight[x], scales: scales[x], biases: biases == nil ? nil : biases![x], groupSize: groupSize, bits: bits, mode: mode) - return out.reshaped(s + [-1]) + return applyNVFP4GlobalScale(out, globalScale: globalScale).reshaped(s + [-1]) } open override func asLinear(_ x: MLXArray) -> MLXArray { - quantizedMM( + let result = quantizedMM( x, weight, scales: scales, biases: biases, transpose: true, groupSize: groupSize, bits: bits, mode: mode) + return applyNVFP4GlobalScale(result, globalScale: globalScale) } } @@ -243,6 +269,7 @@ open class QuantizedLinear: Linear, Quantized { public let mode: QuantizationMode public let scales: MLXArray public let biases: MLXArray? + @ParameterInfo(key: "global_scale") public private(set) var globalScale: MLXArray? open override var shape: (Int, Int) { let shape = weight.shape2 @@ -292,14 +319,16 @@ open class QuantizedLinear: Linear, Quantized { /// Initialize a ``QuantizedLinear`` with non-quantized weights and bias. public init( weight: MLXArray, bias: MLXArray?, groupSize: Int = 64, bits: Int = 4, - mode: QuantizationMode = .affine + mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil ) { self.groupSize = groupSize self.bits = bits self.mode = mode + self._globalScale.wrappedValue = globalScale let (quantizedWeight, scales, biases) = MLX.quantized( - weight, groupSize: groupSize, bits: bits, mode: mode) + weight, groupSize: groupSize, bits: bits, mode: mode, globalScale: globalScale) self.scales = scales self.biases = biases @@ -316,13 +345,15 @@ open class QuantizedLinear: Linear, Quantized { public init( weight: MLXArray, bias: MLXArray? = nil, scales: MLXArray, biases: MLXArray?, groupSize: Int, bits: Int, - mode: QuantizationMode = .affine + mode: QuantizationMode = .affine, + globalScale: MLXArray? = nil ) { self.groupSize = groupSize self.bits = bits self.mode = mode self.scales = scales self.biases = biases + self._globalScale.wrappedValue = globalScale super.init(weight: weight, bias: bias) } @@ -334,7 +365,7 @@ open class QuantizedLinear: Linear, Quantized { } open override func callAsFunction(_ x: MLXArray) -> MLXArray { - var x = quantizedMM( + var result = quantizedMM( x, weight, scales: scales, @@ -344,10 +375,11 @@ open class QuantizedLinear: Linear, Quantized { bits: bits, mode: mode ) + result = applyNVFP4GlobalScale(result, globalScale: globalScale) if let bias { - x = x + bias + result = result + bias } - return x + return result } /// Returns a QuantizedLinear layer that applies the same linear transformation up to the quantization error. diff --git a/Tests/MLXTests/QuantizationTests.swift b/Tests/MLXTests/QuantizationTests.swift index 0edbd545..357ac81e 100644 --- a/Tests/MLXTests/QuantizationTests.swift +++ b/Tests/MLXTests/QuantizationTests.swift @@ -6,6 +6,74 @@ import MLXNN import XCTest class QuantizationTests: XCTestCase { + func testQuantizedLinearAppliesNVFP4GlobalScaleOnMetal() { + let (input, quantizedWeight, scales, biases, globalScale, bias, expected) = + Device.withDefaultDevice(.cpu) { + MLXRandom.seed(7) + let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32]) + let input = MLXRandom.uniform(low: -1, high: 1, [4, 32]) + let globalScale = MLX.max(abs(weight)).asType(.float32) + let bias = MLXRandom.uniform(low: -1, high: 1, [32]) + let (quantizedWeight, scales, biases) = MLX.quantized( + weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale) + let expected = + matmul( + input, + dequantized( + quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4, + mode: .nvfp4, globalScale: globalScale, dtype: .float32 + ).T) + + bias + eval(quantizedWeight, scales, expected) + return (input, quantizedWeight, scales, biases, globalScale, bias, expected) + } + + let actual = Device.withDefaultDevice(.gpu) { + QuantizedLinear( + weight: quantizedWeight, bias: bias, scales: scales, biases: biases, + groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale)(input) + } + + assertEqual(actual, expected, rtol: 1e-5, atol: 1e-5) + } + + func testQuantizedEmbeddingAppliesNVFP4GlobalScaleOnMetal() { + let ( + input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear, + expectedLookup + ) = Device.withDefaultDevice(.cpu) { + MLXRandom.seed(11) + let weight = MLXRandom.uniform(low: -2, high: 2, [32, 32]) + let input = MLXRandom.uniform(low: -1, high: 1, [4, 32]) + let indices = MLXArray([0, 3, 7, 31]) + let globalScale = MLX.max(abs(weight)).asType(.float32) + let (quantizedWeight, scales, biases) = MLX.quantized( + weight, groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale) + let dequantizedWeight = dequantized( + quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4, + mode: .nvfp4, globalScale: globalScale, dtype: .float32) + let expectedLinear = matmul(input, dequantizedWeight.T) + let expectedLookup = dequantized( + quantizedWeight, scales: scales, biases: biases, groupSize: 16, bits: 4, + mode: .nvfp4, globalScale: globalScale)[indices].asType(.bfloat16) + eval(quantizedWeight, scales, expectedLinear, expectedLookup) + return ( + input, indices, quantizedWeight, scales, biases, globalScale, expectedLinear, + expectedLookup + ) + } + + let (actualLinear, actualLookup) = Device.withDefaultDevice(.gpu) { + let embedding = QuantizedEmbedding( + weight: quantizedWeight, scales: scales, biases: biases, + groupSize: 16, bits: 4, mode: .nvfp4, globalScale: globalScale) + return (embedding.asLinear(input), embedding(indices)) + } + + assertEqual(actualLinear, expectedLinear, rtol: 1e-5, atol: 1e-5) + assertEqual(actualLookup, expectedLookup, rtol: 1e-5, atol: 1e-5) + } + func testQuantizedLinearShapeDesc() { let linear1 = Linear(512, 1024) let quantized1 = linear1.toQuantized(groupSize: 64, bits: 4) @@ -39,4 +107,62 @@ class QuantizationTests: XCTestCase { let quantized = QuantizedLinear(64, 64, groupSize: 32, bits: 4, mode: .mxfp4) XCTAssertNil(quantized.biases) } + + func testQuantizedLinearStoresGlobalScale() { + let globalScale = MLXArray(1.0, dtype: .float32) + let quantized = QuantizedLinear( + weight: MLXArray.zeros([8, 4], dtype: .uint32), + bias: nil, + scales: MLXArray.ones([8, 4], dtype: .uint8), + biases: nil, + groupSize: 16, + bits: 4, + mode: .nvfp4, + globalScale: globalScale) + + XCTAssertNotNil(quantized.globalScale) + XCTAssertEqual(quantized.globalScale?.dtype, .float32) + XCTAssertNotNil(quantized.parameters()["global_scale"]) + XCTAssertNil(quantized.parameters()["globalScale"]) + } + + func testQuantizedEmbeddingStoresGlobalScale() { + let globalScale = MLXArray(1.0, dtype: .float32) + let quantized = QuantizedEmbedding( + weight: MLXArray.zeros([8, 2], dtype: .uint32), + scales: MLXArray.ones([8, 2], dtype: .uint8), + biases: nil, + groupSize: 16, + bits: 4, + mode: .nvfp4, + globalScale: globalScale) + + XCTAssertNotNil(quantized.globalScale) + XCTAssertEqual(quantized.globalScale?.dtype, .float32) + XCTAssertNotNil(quantized.parameters()["global_scale"]) + XCTAssertNil(quantized.parameters()["globalScale"]) + } + + func testQuantizedGlobalScaleIsOptionalParameter() { + let linear = QuantizedLinear( + weight: MLXArray.zeros([8, 4], dtype: .uint32), + bias: nil, + scales: MLXArray.ones([8, 4], dtype: .uint8), + biases: nil, + groupSize: 16, + bits: 4, + mode: .nvfp4) + let embedding = QuantizedEmbedding( + weight: MLXArray.zeros([8, 2], dtype: .uint32), + scales: MLXArray.ones([8, 2], dtype: .uint8), + biases: nil, + groupSize: 16, + bits: 4, + mode: .nvfp4) + + XCTAssertNil(linear.globalScale) + XCTAssertNil(linear.parameters()["global_scale"]) + XCTAssertNil(embedding.globalScale) + XCTAssertNil(embedding.parameters()["global_scale"]) + } }