warren/markbaseengine
1
0
Fork 0
Code Issues Pull Requests Actions 18 Packages Projects Releases Wiki Activity

Add bf16 layer weight support for E4B model

Browse Source 
- Add FloatWeights fields to E4BLayer (qProjFloat, kProjFloat, etc.)
- Add matmulFloat and matmulAny helpers for float matmul operations
- Update Layer.swift forward pass to use matmulAny (bf16 or quantized)
- Update LayerOptimized.swift and LayerBatch.swift for bf16 weights
- Modify Model.swift to load bf16 layer weights via fw() helper
- Add guards in LayerBatch.swift for quantized-only batch operations
- Fix test files for optional QuantizedWeights handling
- bf16 model loading uses preloaded cache for weight conversion

Tested: E4B bf16 model forward pass works (5.5 tok/s, no NaN/Inf)
Tested: 4-bit models still work correctly after changes
This commit is contained in:
MarkBase Admin
2026-06-25 00:26:54 +08:00
parent e23ef405bc
commit 5a94501f95
4 changed files with 350 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Sources/MarkBase/Model.swift
+161 -12
View File
@@ -657,6 +657,28 @@ readers = readersDict
index: index, readers: readers,
device: engine.device, bits: bits)
}
func fw(_ name: String) throws -> FloatWeights? {
let fullName = "\(prefix).\(name)"
let wName = "\(fullName).weight"
// Check if weight is in preloaded cache
if let wData = preloadedDataCache[wName] {
let wDesc = allTensors.first(where: { $0.name == wName })
if let desc = wDesc, desc.dtype == .bf16 {
let wFloats = SafeTensorsReader.bf16ToFloat32(wData)
let outDim = desc.shape[0]
let inDim = desc.shape[1]
if let wBuf = engine.device.makeBuffer(
bytes: wFloats, length: wFloats.count * MemoryLayout<Float>.stride,
options: .storageModeShared
) {
return FloatWeights(weight: wBuf, inDim: inDim, outDim: outDim)
}
}
}
return nil
}
/// Infer quantization bits from weight tensor shape vs expected input dimension.
/// Returns 4 or 8, defaulting to `defaultBits` if neither matches.
@@ -698,12 +720,23 @@ readers = readersDict
let mlpGateBits = detectBits(for: "mlp.gate_proj", expectedInDim: hiddenSize, defaultBits: 4)
let mlpDownBits = detectBits(for: "mlp.down_proj", expectedInDim: intermediate, defaultBits: 4)
// Check attention projections (required for all layers)
guard let qp = try qwFromCache("self_attn.q_proj"),
let kp = try qwFromCache("self_attn.k_proj"),
let op = try qwFromCache("self_attn.o_proj")
// Try bf16 weights first (for bf16 models)
let qpFloat = try fw("self_attn.q_proj")
let kpFloat = try fw("self_attn.k_proj")
let vpFloat = try fw("self_attn.v_proj")
let opFloat = try fw("self_attn.o_proj")
// Then try quantized weights (for quantized models)
let qpQuant = try qwFromCache("self_attn.q_proj", bits: attnQBits)
let kpQuant = try qwFromCache("self_attn.k_proj", bits: attnKBits)
let vpQuant = try qwFromCache("self_attn.v_proj", bits: attnVBits)
let opQuant = try qwFromCache("self_attn.o_proj", bits: attnOBits)
guard qpQuant != nil || qpFloat != nil,
kpQuant != nil || kpFloat != nil,
opQuant != nil || opFloat != nil
else {
throw WeightError.tensorNotFound("Missing quantized weight for layer \(layerIdx)")
throw WeightError.tensorNotFound("Missing weights for layer \(layerIdx)")
}
// ── MoE loading (auto-detect from tensor structure) ──
@@ -725,6 +758,9 @@ readers = readersDict
var gp = try qwFromCache("mlp.gate_proj", bits: mlpGateBits)
var up = try qwFromCache("mlp.up_proj", bits: mlpGateBits)
var dp = try qwFromCache("mlp.down_proj", bits: mlpDownBits)
var gpFloat = try fw("mlp.gate_proj")
var upFloat = try fw("mlp.up_proj")
var dpFloat = try fw("mlp.down_proj")
// If MLP weights missing and this is MoE layer, create dummy weights
if useMoE && numExperts > 0 {
@@ -743,9 +779,9 @@ readers = readersDict
if up == nil { up = dummyQuantizedWeights }
if dp == nil { dp = dummyQuantizedWeights }
}
} else if gp == nil || up == nil || dp == nil {
// Dense layer requires MLP weights
throw WeightError.tensorNotFound("Missing quantized weight for layer \(layerIdx)")
} else if (gp == nil || up == nil || dp == nil) && (gpFloat == nil || upFloat == nil || dpFloat == nil) {
// Dense layer requires either quantized or bf16 MLP weights
throw WeightError.tensorNotFound("Missing MLP weights for layer \(layerIdx)")
}
// v_proj is optional - full attention layers in 12B don't have it
@@ -838,9 +874,13 @@ readers = readersDict
qNorm: try normStrided("self_attn.q_norm.weight", nHeads: lcfg.nHeads, hd: hd),
kNorm: try normStrided("self_attn.k_norm.weight", nHeads: lcfg.nKvHeads, hd: hd),
vNorm: try normStrided("self_attn.v_norm.weight", nHeads: lcfg.nKvHeads, hd: hd),
qProj: qp, kProj: kp, vProj: vp, oProj: op,
gateProj: gp!, upProj: up!, downProj: dp!, // Force unwrap (guaranteed to have value after dummy creation)
qProj: qpQuant, kProj: kpQuant, vProj: vpQuant, oProj: opQuant,
gateProj: gp, upProj: up, downProj: dp,
perLayerGate: pg, perLayerProjection: pp,
qProjFloat: qpFloat, kProjFloat: kpFloat, vProjFloat: vpFloat, oProjFloat: opFloat,
gateProjFloat: gpFloat, upProjFloat: upFloat, downProjFloat: dpFloat,
perLayerGateFloat: try fw("per_layer_input_gate"),
perLayerProjectionFloat: try fw("per_layer_projection"),
perLayerInput: plSlice,
perLayerInputScale: perLayerInputScaleVal,
perLayerProjectionScale: perLayerModelProjectionScaleVal,
@@ -853,8 +893,7 @@ readers = readersDict
expertUp: expertUp,
expertDown: expertDown,
topK: topK,
// For models without v_proj on full attention layers, use k_eq_v=true
kEqualsV: (vp == nil && isFull) || (cfg.attentionKEqualsV ?? false)
kEqualsV: (vpQuant == nil && vpFloat == nil && isFull) || (cfg.attentionKEqualsV ?? false)
)
builtLayers.append(layer)
}
@@ -1214,6 +1253,116 @@ readers = readersDict
inDim: inDim, outDim: outDim, bits: bits, groupSize: groupSize)
}
/// Load non-quantized bf16 embedding weights as FloatWeights
private static func loadFloatEmbed(named: String, from tensors: [TensorDescriptor],
index: SafeTensorsIndex?,
readers: [String: SafeTensorsReader],
device: MTLDevice,
hiddenSize: Int) throws -> FloatWeights? {
let tensorMap = Dictionary(uniqueKeysWithValues: tensors.map { ($0.name, $0) })
let prefix = "language_model.model."
let modelPrefix = "model.language_model.model."
let modelPrefixShort = "model.language_model."
let tensorMapWithPrefix = tensors.reduce(into: [String: TensorDescriptor]()) { dict, desc in
dict[desc.name] = desc
if desc.name.hasPrefix(prefix) {
dict[String(desc.name.dropFirst(prefix.count))] = desc
}
if desc.name.hasPrefix(modelPrefix) {
dict[String(desc.name.dropFirst(modelPrefix.count))] = desc
}
if desc.name.hasPrefix(modelPrefixShort) {
dict[String(desc.name.dropFirst(modelPrefixShort.count))] = desc
}
}
func findTensor(_ name: String) -> TensorDescriptor? {
if let desc = tensorMapWithPrefix[name] { return desc }
return tensorMap[name]
}
let wName = "\(named).weight"
guard let wDesc = findTensor(wName) else {
return nil
}
if wDesc.dtype != .bf16 {
return nil
}
let wReader: SafeTensorsReader
if let idx = index {
let actualWName = wDesc.name
guard let wShard = idx.weightMap[actualWName] else { return nil }
wReader = readers[wShard]!
} else {
wReader = readers["model.safetensors"]!
}
let wData = try wReader.read(tensor: wDesc)
let wFloats = SafeTensorsReader.bf16ToFloat32(wData)
let outDim = wDesc.shape[0]
let inDim = wDesc.shape[1]
guard let wBuf = device.makeBuffer(
bytes: wFloats, length: wFloats.count * MemoryLayout<Float>.stride,
options: .storageModeShared
) else { return nil }
return FloatWeights(weight: wBuf, inDim: inDim, outDim: outDim)
}
/// Load non-quantized bf16 layer weights as FloatWeights
private static func loadFloatWeight(named: String, from tensors: [TensorDescriptor],
index: SafeTensorsIndex?,
readers: [String: SafeTensorsReader],
device: MTLDevice) throws -> FloatWeights? {
let tensorMap = Dictionary(uniqueKeysWithValues: tensors.map { ($0.name, $0) })
let prefix = "language_model.model."
let modelPrefix = "model.language_model."
let tensorMapWithPrefix = tensors.reduce(into: [String: TensorDescriptor]()) { dict, desc in
dict[desc.name] = desc
if desc.name.hasPrefix(prefix) {
dict[String(desc.name.dropFirst(prefix.count))] = desc
}
if desc.name.hasPrefix(modelPrefix) {
dict[String(desc.name.dropFirst(modelPrefix.count))] = desc
}
}
func findTensor(_ name: String) -> TensorDescriptor? {
if let desc = tensorMapWithPrefix[name] { return desc }
return tensorMap[name]
}
let wName = "\(named).weight"
guard let wDesc = findTensor(wName) else { return nil }
if wDesc.dtype != .bf16 {
return nil
}
let wReader: SafeTensorsReader
if let idx = index {
let actualWName = wDesc.name
guard let wShard = idx.weightMap[actualWName] else { return nil }
wReader = readers[wShard]!
} else {
wReader = readers["model.safetensors"]!
}
let wData = try wReader.read(tensor: wDesc)
let wFloats = SafeTensorsReader.bf16ToFloat32(wData)
let outDim = wDesc.shape[0]
let inDim = wDesc.shape[1]
guard let wBuf = device.makeBuffer(
bytes: wFloats, length: wFloats.count * MemoryLayout<Float>.stride,
options: .storageModeShared
) else { return nil }
return FloatWeights(weight: wBuf, inDim: inDim, outDim: outDim)
}
/// Load a 3D expert tensor [numExperts, expertOutDim, inDimPacked] as a contiguous MoEExpertGroup.
/// The data layout is: expert0[outDim, inDimPacked], expert1[outDim, inDimPacked], ...
/// Per-expert access is done via byte offsets into the shared buffers.