MarkBase Admin
378 lines
19 KiB
Swift
378 lines
19 KiB
Swift
import Foundation
|
||
import Metal
|
||
import MetalPerformanceShaders
|
||
import Accelerate
|
||
|
||
/// EmbeddingGemma configuration
|
||
public struct EmbeddingGemmaConfig: Codable {
|
||
public let hiddenSize: Int
|
||
public let numHiddenLayers: Int
|
||
public let vocabSize: Int
|
||
public let numAttentionHeads: Int
|
||
public let numKeyValueHeads: Int
|
||
public let headDim: Int
|
||
public let intermediateSize: Int
|
||
public let maxPositionEmbeddings: Int
|
||
public let slidingWindow: Int
|
||
public let rmsNormEps: Float
|
||
public let ropeTheta: Float
|
||
public let useBidirectionalAttention: Bool
|
||
public let layerTypes: [String]
|
||
|
||
enum CodingKeys: String, CodingKey {
|
||
case hiddenSize = "hidden_size"
|
||
case numHiddenLayers = "num_hidden_layers"
|
||
case vocabSize = "vocab_size"
|
||
case numAttentionHeads = "num_attention_heads"
|
||
case numKeyValueHeads = "num_key_value_heads"
|
||
case headDim = "head_dim"
|
||
case intermediateSize = "intermediate_size"
|
||
case maxPositionEmbeddings = "max_position_embeddings"
|
||
case slidingWindow = "sliding_window"
|
||
case rmsNormEps = "rms_norm_eps"
|
||
case ropeTheta = "rope_theta"
|
||
case useBidirectionalAttention = "use_bidirectional_attention"
|
||
case layerTypes = "layer_types"
|
||
}
|
||
|
||
public static func load(from modelDir: String) throws -> Self {
|
||
let url = URL(fileURLWithPath: modelDir).appendingPathComponent("config.json")
|
||
let data = try Data(contentsOf: url)
|
||
return try JSONDecoder().decode(Self.self, from: data)
|
||
}
|
||
}
|
||
|
||
public final class EmbeddingGemmaModel: @unchecked Sendable {
|
||
public let config: EmbeddingGemmaConfig
|
||
public let engine: MarkBaseEngine
|
||
public let tokenizer: Tokenizer
|
||
public let reader: SafeTensorsReader
|
||
|
||
public var embedTokens: MTLBuffer!
|
||
public var finalNorm: MTLBuffer!
|
||
public var layerNorms: [[MTLBuffer]] = []
|
||
public var qProjs: [MTLBuffer] = []
|
||
public var kProjs: [MTLBuffer] = []
|
||
public var vProjs: [MTLBuffer] = []
|
||
public var oProjs: [MTLBuffer] = []
|
||
public var qNorms: [MTLBuffer] = []
|
||
public var kNorms: [MTLBuffer] = []
|
||
public var gateProjs: [MTLBuffer] = []
|
||
public var upProjs: [MTLBuffer] = []
|
||
public var downProjs: [MTLBuffer] = []
|
||
|
||
public init(modelDir: String, engine: MarkBaseEngine) throws {
|
||
self.engine = engine
|
||
self.config = try EmbeddingGemmaConfig.load(from: modelDir)
|
||
self.tokenizer = try TokenizerFactory.load(modelDir: modelDir)
|
||
self.reader = try SafeTensorsReader(path: modelDir + "/model.safetensors")
|
||
try loadWeights()
|
||
print("✓ EmbeddingGemma loaded (\(config.numHiddenLayers) layers, hidden=\(config.hiddenSize))")
|
||
}
|
||
|
||
private func loadWeights() throws {
|
||
let embedData = try readTensor("embed_tokens.weight")
|
||
embedTokens = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
|
||
for i in 0..<config.numHiddenLayers {
|
||
let p = "layers.\(i)"
|
||
layerNorms.append([
|
||
try loadBuffer("\(p).input_layernorm.weight"),
|
||
try loadBuffer("\(p).pre_feedforward_layernorm.weight"),
|
||
try loadBuffer("\(p).post_attention_layernorm.weight"),
|
||
try loadBuffer("\(p).post_feedforward_layernorm.weight"),
|
||
])
|
||
qProjs.append(try loadAndTranspose("\(p).self_attn.q_proj.weight", rows: config.hiddenSize, cols: config.hiddenSize))
|
||
kProjs.append(try loadAndTranspose("\(p).self_attn.k_proj.weight", rows: config.numKeyValueHeads * config.headDim, cols: config.hiddenSize))
|
||
vProjs.append(try loadAndTranspose("\(p).self_attn.v_proj.weight", rows: config.numKeyValueHeads * config.headDim, cols: config.hiddenSize))
|
||
oProjs.append(try loadAndTranspose("\(p).self_attn.o_proj.weight", rows: config.numAttentionHeads * config.headDim, cols: config.hiddenSize))
|
||
qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))
|
||
kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))
|
||
gateProjs.append(try loadAndTranspose("\(p).mlp.gate_proj.weight", rows: config.intermediateSize, cols: config.hiddenSize))
|
||
upProjs.append(try loadAndTranspose("\(p).mlp.up_proj.weight", rows: config.intermediateSize, cols: config.hiddenSize))
|
||
downProjs.append(try loadAndTranspose("\(p).mlp.down_proj.weight", rows: config.hiddenSize, cols: config.intermediateSize))
|
||
}
|
||
let fnData = try readTensor("norm.weight")
|
||
finalNorm = engine.device.makeBuffer(bytes: fnData, length: fnData.count * 4)!
|
||
}
|
||
|
||
public func embed(text: String, maxLen: Int = 2048) throws -> [Float] {
|
||
var tokens = tokenizer.encode(text: text)
|
||
if tokens.count > maxLen { tokens = Array(tokens.prefix(maxLen)) }
|
||
guard !tokens.isEmpty else { return [] }
|
||
|
||
let seqLen = tokens.count, hs = config.hiddenSize
|
||
|
||
// Test 1: Embedding lookup only
|
||
print(" TEST: Embedding lookup...")
|
||
let cmdBuf1 = engine.commandQueue.makeCommandBuffer()!
|
||
let inputBuf = try lookupEmbeddings(tokens: tokens, cmdBuf: cmdBuf1)
|
||
cmdBuf1.commit()
|
||
cmdBuf1.waitUntilCompleted()
|
||
print(" TEST: Embedding lookup OK")
|
||
|
||
// Test 2: Single layer forward (layer 0 only)
|
||
print(" TEST: Layer 0 forward...")
|
||
let cmdBuf2 = engine.commandQueue.makeCommandBuffer()!
|
||
var hidden = try forwardLayerDebug(hidden: inputBuf, layerIdx: 0, seqLen: seqLen, cmdBuf: cmdBuf2)
|
||
cmdBuf2.commit()
|
||
cmdBuf2.waitUntilCompleted()
|
||
print(" TEST: Layer 0 OK")
|
||
|
||
// Full forward pass
|
||
print(" TEST: Full forward pass...")
|
||
let cmdBuf = engine.commandQueue.makeCommandBuffer()!
|
||
hidden = inputBuf
|
||
for idx in 0..<config.numHiddenLayers {
|
||
if idx % 6 == 0 { print(" Layer \(idx)/\(config.numHiddenLayers)") }
|
||
hidden = try forwardLayerDebug(hidden: hidden, layerIdx: idx, seqLen: seqLen, cmdBuf: cmdBuf)
|
||
}
|
||
print(" TEST: All layers OK")
|
||
|
||
let output = try applyRmsNorm(input: hidden, weight: finalNorm, count: seqLen * hs, cmdBuf: cmdBuf)
|
||
cmdBuf.commit()
|
||
cmdBuf.waitUntilCompleted()
|
||
|
||
let data = engine.readFloats(from: output, count: seqLen * hs)
|
||
var embedding = [Float](repeating: 0, count: hs)
|
||
for i in 0..<seqLen {
|
||
let start = i * hs
|
||
for j in 0..<hs { embedding[j] += data[start + j] }
|
||
}
|
||
let n = Float(seqLen)
|
||
for i in 0..<hs { embedding[i] /= n }
|
||
var norm: Float = 0
|
||
for i in 0..<hs { norm += embedding[i] * embedding[i] }
|
||
norm = sqrt(norm)
|
||
if norm > 0 { for i in 0..<hs { embedding[i] /= norm } }
|
||
return embedding
|
||
}
|
||
|
||
private func readTensor(_ name: String) throws -> [Float] {
|
||
guard let desc = reader.tensor(named: name) else { throw WeightError.tensorNotFound(name) }
|
||
let data = try reader.read(tensor: desc)
|
||
switch desc.dtype {
|
||
case .f32: return data.withUnsafeBytes { Array(UnsafeBufferPointer(start: $0.baseAddress?.assumingMemoryBound(to: Float.self), count: data.count/4)) }
|
||
case .bf16: return try SafeTensorsReader.bf16ToFloat32(data)
|
||
default: throw WeightError.unsupportedDtype(desc.dtype.rawValue)
|
||
}
|
||
}
|
||
|
||
private func loadBuffer(_ name: String) throws -> MTLBuffer {
|
||
let data = try readTensor(name)
|
||
return engine.device.makeBuffer(bytes: data, length: data.count * 4)!
|
||
}
|
||
|
||
private func loadAndTranspose(_ name: String, rows: Int, cols: Int) throws -> MTLBuffer {
|
||
// Load [rows, cols] and transpose to [cols, rows] for MPS matmul C = A × B
|
||
let data = try readTensor(name)
|
||
var transposed = [Float](repeating: 0, count: data.count)
|
||
for r in 0..<rows {
|
||
for c in 0..<cols {
|
||
transposed[c * rows + r] = data[r * cols + c]
|
||
}
|
||
}
|
||
return engine.device.makeBuffer(bytes: transposed, length: transposed.count * 4)!
|
||
}
|
||
|
||
private func lookupEmbeddings(tokens: [Int], cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
|
||
let seqLen = tokens.count, hs = config.hiddenSize
|
||
print(" lookupEmbeddings: seqLen=\(seqLen), hs=\(hs)")
|
||
// Read embedding table to CPU
|
||
let embedPtr = embedTokens.contents().assumingMemoryBound(to: Float.self)
|
||
let embedCount = embedTokens.length / 4
|
||
let embedArray = Array(UnsafeBufferPointer(start: embedPtr, count: embedCount))
|
||
print(" Read \(embedCount) floats from embedTokens")
|
||
|
||
// Lookup embeddings for tokens
|
||
var embedData = [Float](repeating: 0, count: seqLen * hs)
|
||
for (i, token) in tokens.enumerated() {
|
||
let start = i * hs
|
||
let srcStart = token * hs
|
||
embedData[start..<start+hs] = embedArray[srcStart..<srcStart+hs]
|
||
}
|
||
print(" Looked up \(seqLen) tokens")
|
||
|
||
let buf = engine.device.makeBuffer(bytes: embedData, length: embedData.count * 4)!
|
||
print(" Created MTLBuffer")
|
||
return buf
|
||
}
|
||
|
||
private func applyRmsNorm(input: MTLBuffer, weight: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
|
||
let output = engine.device.makeBuffer(length: count * 4)!
|
||
let enc = cmdBuf.makeComputeCommandEncoder()!
|
||
defer { enc.endEncoding() }
|
||
let pso = try engine.pipeline(named: "rms_norm")
|
||
enc.setComputePipelineState(pso)
|
||
enc.setBuffer(input, offset: 0, index: 0)
|
||
enc.setBuffer(weight, offset: 0, index: 1)
|
||
enc.setBuffer(output, offset: 0, index: 2)
|
||
var c = UInt32(count), e: Float = config.rmsNormEps
|
||
enc.setBytes(&c, length: 4, index: 3)
|
||
enc.setBytes(&e, length: 4, index: 4)
|
||
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
|
||
return output
|
||
}
|
||
|
||
private func matmulSeq(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer, m: Int, k: Int, n: Int, cmdBuf: MTLCommandBuffer) throws {
|
||
// Use MPS for optimized matrix multiplication on Apple Silicon
|
||
// Weight is stored transposed [k, n] for C = A × B
|
||
let descA = MPSMatrixDescriptor(rows: m, columns: k, rowBytes: k * 4, dataType: .float32)
|
||
let descB = MPSMatrixDescriptor(rows: k, columns: n, rowBytes: n * 4, dataType: .float32)
|
||
let descC = MPSMatrixDescriptor(rows: m, columns: n, rowBytes: n * 4, dataType: .float32)
|
||
let matA = MPSMatrix(buffer: input, descriptor: descA)
|
||
let matB = MPSMatrix(buffer: weight, descriptor: descB)
|
||
let matC = MPSMatrix(buffer: output, descriptor: descC)
|
||
|
||
let matMul = MPSMatrixMultiplication(device: engine.device,
|
||
transposeLeft: false,
|
||
transposeRight: false,
|
||
resultRows: m,
|
||
resultColumns: n,
|
||
interiorColumns: k,
|
||
alpha: 1.0,
|
||
beta: 0.0)
|
||
matMul.encode(commandBuffer: cmdBuf, leftMatrix: matA, rightMatrix: matB, resultMatrix: matC)
|
||
}
|
||
|
||
private func eltwiseAdd(a: MTLBuffer, b: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
|
||
let enc = cmdBuf.makeComputeCommandEncoder()!
|
||
defer { enc.endEncoding() }
|
||
let pso = try engine.pipeline(named: "eltwise_add")
|
||
enc.setComputePipelineState(pso)
|
||
enc.setBuffer(a, offset: 0, index: 0)
|
||
enc.setBuffer(b, offset: 0, index: 1)
|
||
enc.setBuffer(output, offset: 0, index: 2)
|
||
var c = UInt32(count)
|
||
enc.setBytes(&c, length: 4, index: 3)
|
||
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
|
||
}
|
||
|
||
private func geluMul(gate: MTLBuffer, up: MTLBuffer, output: MTLBuffer, count: Int, cmdBuf: MTLCommandBuffer) throws {
|
||
let enc = cmdBuf.makeComputeCommandEncoder()!
|
||
defer { enc.endEncoding() }
|
||
let pso = try engine.pipeline(named: "gelu_mul_kernel")
|
||
enc.setComputePipelineState(pso)
|
||
enc.setBuffer(gate, offset: 0, index: 0)
|
||
enc.setBuffer(up, offset: 0, index: 1)
|
||
enc.setBuffer(output, offset: 0, index: 2)
|
||
var c = UInt32(count)
|
||
enc.setBytes(&c, length: 4, index: 3)
|
||
enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, count), height: 1, depth: 1))
|
||
}
|
||
|
||
private func applyRoPE(q: MTLBuffer, k: MTLBuffer, seqLen: Int, headDim: Int, numHeads: Int, cmdBuf: MTLCommandBuffer) throws {
|
||
let enc = cmdBuf.makeComputeCommandEncoder()!
|
||
defer { enc.endEncoding() }
|
||
let pso = try engine.pipeline(named: "apply_rope")
|
||
enc.setComputePipelineState(pso)
|
||
enc.setBuffer(q, offset: 0, index: 0)
|
||
enc.setBuffer(k, offset: 0, index: 1)
|
||
var sl = UInt32(seqLen), hd = UInt32(headDim), nh = UInt32(numHeads)
|
||
var rt: Float = Float(config.ropeTheta)
|
||
enc.setBytes(&sl, length: 4, index: 2)
|
||
enc.setBytes(&hd, length: 4, index: 3)
|
||
enc.setBytes(&nh, length: 4, index: 4)
|
||
enc.setBytes(&rt, length: 4, index: 5)
|
||
enc.dispatchThreads(MTLSize(width: numHeads * headDim / 2, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, numHeads * headDim / 2), height: 1, depth: 1))
|
||
}
|
||
|
||
private func bidirectionalAttention(q: MTLBuffer, k: MTLBuffer, v: MTLBuffer, output: MTLBuffer, seqLen: Int, cmdBuf: MTLCommandBuffer) throws {
|
||
let enc = cmdBuf.makeComputeCommandEncoder()!
|
||
defer { enc.endEncoding() }
|
||
let pso = try engine.pipeline(named: "bidirectional_sliding_attn")
|
||
enc.setComputePipelineState(pso)
|
||
enc.setBuffer(q, offset: 0, index: 0)
|
||
enc.setBuffer(k, offset: 0, index: 1)
|
||
enc.setBuffer(v, offset: 0, index: 2)
|
||
enc.setBuffer(output, offset: 0, index: 3)
|
||
var sl = UInt32(seqLen), hd = UInt32(config.headDim), nh = UInt32(config.numAttentionHeads)
|
||
var nkv = UInt32(config.numKeyValueHeads), sw = UInt32(config.slidingWindow)
|
||
var scale: Float = 1.0 / sqrt(Float(config.headDim))
|
||
enc.setBytes(&sl, length: 4, index: 4)
|
||
enc.setBytes(&hd, length: 4, index: 5)
|
||
enc.setBytes(&nh, length: 4, index: 6)
|
||
enc.setBytes(&nkv, length: 4, index: 7)
|
||
enc.setBytes(&sw, length: 4, index: 8)
|
||
enc.setBytes(&scale, length: 4, index: 9)
|
||
let tgMem = config.slidingWindow * 4
|
||
enc.setThreadgroupMemoryLength(tgMem, index: 0)
|
||
enc.dispatchThreads(MTLSize(width: seqLen * config.numAttentionHeads, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, seqLen * config.numAttentionHeads), height: 1, depth: 1))
|
||
}
|
||
|
||
private func forwardLayerDebug(hidden: MTLBuffer, layerIdx: Int, seqLen: Int, cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
|
||
let hs = config.hiddenSize, device = engine.device
|
||
let hDim = config.headDim, nH = config.numAttentionHeads, nKV = config.numKeyValueHeads
|
||
let intermedi = config.intermediateSize
|
||
|
||
print(" Residual copy...")
|
||
let resid = device.makeBuffer(length: seqLen * hs * 4)!
|
||
let blit = cmdBuf.makeBlitCommandEncoder()!
|
||
blit.copy(from: hidden, sourceOffset: 0, to: resid, destinationOffset: 0, size: seqLen * hs * 4)
|
||
blit.endEncoding()
|
||
|
||
print(" Input norm...")
|
||
let h1 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][0], count: seqLen * hs, cmdBuf: cmdBuf)
|
||
|
||
print(" QKV projections...")
|
||
let qBuf = device.makeBuffer(length: seqLen * nH * hDim * 4)!
|
||
let kBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
|
||
let vBuf = device.makeBuffer(length: seqLen * nKV * hDim * 4)!
|
||
try matmulSeq(input: h1, weight: qProjs[layerIdx], output: qBuf, m: seqLen, k: hs, n: nH * hDim, cmdBuf: cmdBuf)
|
||
print(" Q done")
|
||
try matmulSeq(input: h1, weight: kProjs[layerIdx], output: kBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
|
||
print(" K done")
|
||
try matmulSeq(input: h1, weight: vProjs[layerIdx], output: vBuf, m: seqLen, k: hs, n: nKV * hDim, cmdBuf: cmdBuf)
|
||
print(" V done")
|
||
|
||
print(" RoPE...")
|
||
try applyRoPE(q: qBuf, k: kBuf, seqLen: seqLen, headDim: hDim, numHeads: nH, cmdBuf: cmdBuf)
|
||
print(" RoPE done")
|
||
|
||
print(" Attention...")
|
||
let attnOut = device.makeBuffer(length: seqLen * nH * hDim * 4)!
|
||
try bidirectionalAttention(q: qBuf, k: kBuf, v: vBuf, output: attnOut, seqLen: seqLen, cmdBuf: cmdBuf)
|
||
print(" Attention done")
|
||
|
||
print(" O projection...")
|
||
let h2 = device.makeBuffer(length: seqLen * hs * 4)!
|
||
try matmulSeq(input: attnOut, weight: oProjs[layerIdx], output: h2, m: seqLen, k: nH * hDim, n: hs, cmdBuf: cmdBuf)
|
||
print(" O done")
|
||
|
||
print(" Post-attn norm...")
|
||
let h2n = try applyRmsNorm(input: h2, weight: layerNorms[layerIdx][2], count: seqLen * hs, cmdBuf: cmdBuf)
|
||
|
||
print(" Add residual 1...")
|
||
try eltwiseAdd(a: resid, b: h2n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
|
||
|
||
print(" Pre-FF norm...")
|
||
let h3 = try applyRmsNorm(input: hidden, weight: layerNorms[layerIdx][1], count: seqLen * hs, cmdBuf: cmdBuf)
|
||
|
||
print(" MLP gate/up...")
|
||
let gate = device.makeBuffer(length: seqLen * intermedi * 4)!
|
||
let up = device.makeBuffer(length: seqLen * intermedi * 4)!
|
||
try matmulSeq(input: h3, weight: gateProjs[layerIdx], output: gate, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
|
||
print(" Gate done")
|
||
try matmulSeq(input: h3, weight: upProjs[layerIdx], output: up, m: seqLen, k: hs, n: intermedi, cmdBuf: cmdBuf)
|
||
print(" Up done")
|
||
|
||
print(" GELU mul...")
|
||
let gated = device.makeBuffer(length: seqLen * intermedi * 4)!
|
||
try geluMul(gate: gate, up: up, output: gated, count: seqLen * intermedi, cmdBuf: cmdBuf)
|
||
print(" GELU done")
|
||
|
||
print(" Down projection...")
|
||
let h4 = device.makeBuffer(length: seqLen * hs * 4)!
|
||
try matmulSeq(input: gated, weight: downProjs[layerIdx], output: h4, m: seqLen, k: intermedi, n: hs, cmdBuf: cmdBuf)
|
||
print(" Down done")
|
||
|
||
print(" Post-FF norm...")
|
||
let h4n = try applyRmsNorm(input: h4, weight: layerNorms[layerIdx][3], count: seqLen * hs, cmdBuf: cmdBuf)
|
||
|
||
print(" Add residual 2...")
|
||
try eltwiseAdd(a: hidden, b: h4n, output: hidden, count: seqLen * hs, cmdBuf: cmdBuf)
|
||
print(" Layer \(layerIdx) complete")
|
||
|
||
return hidden
|
||
}
|
||
}
|