MarkBase Admin
358 lines
17 KiB
Swift
358 lines
17 KiB
Swift
import Foundation
|
|
import Metal
|
|
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 loadBuffer("\(p).self_attn.q_proj.weight"))
|
|
kProjs.append(try loadBuffer("\(p).self_attn.k_proj.weight"))
|
|
vProjs.append(try loadBuffer("\(p).self_attn.v_proj.weight"))
|
|
oProjs.append(try loadBuffer("\(p).self_attn.o_proj.weight"))
|
|
qNorms.append(try loadBuffer("\(p).self_attn.q_norm.weight"))
|
|
kNorms.append(try loadBuffer("\(p).self_attn.k_norm.weight"))
|
|
gateProjs.append(try loadBuffer("\(p).mlp.gate_proj.weight"))
|
|
upProjs.append(try loadBuffer("\(p).mlp.up_proj.weight"))
|
|
downProjs.append(try loadBuffer("\(p).mlp.down_proj.weight"))
|
|
}
|
|
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 lookupEmbeddings(tokens: [Int], cmdBuf: MTLCommandBuffer) throws -> MTLBuffer {
|
|
let seqLen = tokens.count, hs = config.hiddenSize
|
|
let buf = engine.device.makeBuffer(length: seqLen * hs * 4)!
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
defer { enc.endEncoding() }
|
|
let pso = try engine.pipeline(named: "lookup_embeddings")
|
|
enc.setComputePipelineState(pso)
|
|
enc.setBuffer(embedTokens, offset: 0, index: 0)
|
|
enc.setBytes(tokens, length: seqLen * MemoryLayout<Int>.size, index: 1)
|
|
enc.setBuffer(buf, offset: 0, index: 2)
|
|
var h = UInt32(hs), s = UInt32(seqLen), v = UInt32(config.vocabSize)
|
|
enc.setBytes(&h, length: 4, index: 3)
|
|
enc.setBytes(&s, length: 4, index: 4)
|
|
enc.setBytes(&v, length: 4, index: 5)
|
|
enc.dispatchThreads(MTLSize(width: seqLen, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, seqLen), height: 1, depth: 1))
|
|
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 {
|
|
let enc = cmdBuf.makeComputeCommandEncoder()!
|
|
defer { enc.endEncoding() }
|
|
let pso = try engine.pipeline(named: "matmul_f32")
|
|
enc.setComputePipelineState(pso)
|
|
for i in 0..<m {
|
|
let inputOffset = i * k * 4
|
|
let outputOffset = i * n * 4
|
|
enc.setBuffer(input, offset: inputOffset, index: 0)
|
|
enc.setBuffer(weight, offset: 0, index: 1)
|
|
enc.setBuffer(output, offset: outputOffset, index: 2)
|
|
var mm: UInt32 = 1, kk = UInt32(k), nn = UInt32(n)
|
|
enc.setBytes(&mm, length: 4, index: 3)
|
|
enc.setBytes(&kk, length: 4, index: 4)
|
|
enc.setBytes(&nn, length: 4, index: 5)
|
|
enc.dispatchThreads(MTLSize(width: n, height: 1, depth: 1), threadsPerThreadgroup: MTLSize(width: min(256, n), height: 1, depth: 1))
|
|
}
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|