gpt2-125m / test_model.bybylang
bobbydeveloper2014's picture
Upload folder using huggingface_hub
85f86bf verified
Raw History Blame Contribute Delete
17 kB
# Independent float references; runGrid tests never use the production CPU VM.
import nim "std/[os, tables, json, math, random, strutils]"
import nim "vm_interp"
import nim "backend"
import nim "quant"
import "../../other/formers.n/src/formers/ir/fir_parse.bybylang"
import "../../other/formers.n/src/formers/ir/fir_exec.bybylang"
import "../../other/formers.n/src/formers/ir/fir_graph.bybylang"
import "../../other/formers.n/src/formers/ir/fir_graph_kernels.bybylang"
import "../../other/formers.n/src/formers/io/safetensors.bybylang"
import "../../other/formers.n/src/formers/arch/detect.bybylang"
function fgClose(actual, expected: openArray[float32], label: string)
doAssert actual.len == expected.len, label & " length"
for i in 0 ..< actual.len:
if classify(actual[i]) in {fcNan, fcInf, fcNegInf} or abs(actual[i]-expected[i]) > 1e-4'f32:
raise newException(ValueError, label & " index=" & $i & ": " & $actual[i] & " != " & $expected[i])
function fgRandom(rng: var Rand, count: int, scale: float32 = 0.3'f32): seq[float32]
result = newSeq[float32](count)
for value in result.mitems:
value = (rng.rand(2.0'f32)-1)*scale
function refNorm(x, gamma, beta: seq[float32], eps: float32): seq[float32]
var mean = 0'f32
for value in x:
mean += value
mean /= float32(x.len)
var variance = 0'f32
for value in x:
variance += (value-mean)*(value-mean)
let inverse = 1'f32/sqrt(variance/float32(x.len)+eps)
for i, value in x:
result.add((value-mean)*inverse*gamma[i]+beta[i])
function refAdd(a, b: seq[float32]): seq[float32]
for i in 0 ..< a.len:
result.add(a[i]+b[i])
function refGelu(x: seq[float32]): seq[float32]
# Independent tanh form of gelu_new (equivalent to the sigmoid builder).
for value in x:
result.add(0.5'f32*value*(1'f32+tanh(sqrt(2'f32/PI.float32)*(value+0.044715'f32*value*value*value))))
function refLinear(x, w, bias: seq[float32]): seq[float32]
result = bias
for j in 0 ..< bias.len:
for i in 0 ..< x.len:
result[j] += x[i]*w[i*bias.len+j]
function refAttention(qkv: seq[float32], keys, values: var seq[float32], position, hidden, heads: int): seq[float32]
for i in 0 ..< hidden:
keys[position*hidden+i] = qkv[hidden+i]
values[position*hidden+i] = qkv[2*hidden+i]
let width = hidden div heads
result = newSeq[float32](hidden)
for head in 0 ..< heads:
var scores = newSeq[float32](position+1)
for t in 0 .. position:
for d in 0 ..< width:
scores[t] += qkv[head*width+d]*keys[t*hidden+head*width+d]
scores[t] /= sqrt(float32(width))
var peak = scores[0]
for score in scores:
peak = max(peak, score)
var total = 0'f32
for score in scores.mitems:
score = exp(score-peak)
total += score
for t in 0 .. position:
for d in 0 ..< width:
result[head*width+d] += scores[t]/total*values[t*hidden+head*width+d]
function checkMicro(graph: FirProgram)
let registry = firDefaultRegistry()
let bindings = firGraphBindings(graph)
firValidate(graph, registry, bindings)
var rng = initRand(7293)
for node in graph.nodes:
if node.kind != "micro":
continue
let ins = node.inputs.splitWhitespace()
let output = node.outputs.splitWhitespace()[0]
var memory = fgRandom(rng, bindings.memoryFloats, 4'f32)
var expected: seq[float32]
if node.op == "g2.attn_decode":
var keys = fgRandom(rng, graph.buffers["key_cache"].elements)
var values = fgRandom(rng, keys.len)
for position in 0 ..< graph.symbols["positions"]:
let qkv = fgRandom(rng, graph.buffers["qkv"].elements, 3'f32)
for i, value in qkv:
memory[bindings.memory["qkv"]+i] = value
memory[bindings.memory["position"]] = float32(position)
# Future cache entries are poison; causal decode must not read them.
for i in position*graph.symbols["hidden"] ..< keys.len:
memory[bindings.memory["key_cache"]+i] = 999'f32
memory[bindings.memory["value_cache"]+i] = -999'f32
expected = refAttention(qkv, keys, values, position, graph.symbols["hidden"], graph.symbols["nheads"])
let kernel = firMicroKernel(node, graph, registry, bindings)
runGrid(kernel.program, memory, kernel.threads, kernel.tgSize)
let offset = bindings.memory["attn"]
fgClose(memory[offset ..< offset+expected.len], expected, "attention")
let count = (position+1)*graph.symbols["hidden"]
for cache in ["key_cache", "value_cache"]:
let base = bindings.memory[cache]
let reference = if cache == "key_cache": keys else: values
fgClose(memory[base ..< base+count], reference[0 ..< count], cache)
for head in 0 ..< graph.symbols["nheads"]:
var total = 0'f32
for t in 0 .. position:
total += memory[bindings.memory["scores"]+head*graph.symbols["positions"]+t]
fgClose(@[total], @[1'f32], "softmax sum")
continue
let xoff = bindings.memory[ins[0]]
let count = graph.buffers[ins[0]].elements
let x = memory[xoff ..< xoff+count]
if node.op == "g2.layernorm":
let goff = bindings.memory[ins[1]]
let boff = bindings.memory[ins[2]]
expected = refNorm(x, memory[goff ..< goff+count], memory[boff ..< boff+count], graph.epsilon)
elif node.op == "g2.gelu_new":
expected = refGelu(x)
else:
let boff = bindings.memory[ins[1]]
expected = refAdd(x, memory[boff ..< boff+count])
let kernel = firMicroKernel(node, graph, registry, bindings)
runGrid(kernel.program, memory, kernel.threads, kernel.tgSize)
let offset = bindings.memory[output]
fgClose(memory[offset ..< offset+expected.len], expected, node.name)
echo "PASS GPT-2 all micro nodes through runGrid: layernorm, gelu_new, bias/residual, causal attention/softmax and KV (tol 1e-4)"
function refForward(graph: FirProgram, weights: Table[string, seq[float32]], tokens: seq[int]): seq[float32]
let hidden = graph.symbols["hidden"]
let capacity = graph.symbols["positions"]
let heads = graph.symbols["nheads"]
var keys = newSeq[seq[float32]](graph.symbols["layers"])
var values = newSeq[seq[float32]](keys.len)
for i in 0 ..< keys.len:
keys[i] = newSeq[float32](capacity*hidden)
values[i] = newSeq[float32](capacity*hidden)
for position, token in tokens:
let wte = weights["wte.weight"]
let wpe = weights["wpe.weight"]
var x = refAdd(wte[token*hidden ..< (token+1)*hidden], wpe[position*hidden ..< (position+1)*hidden])
for layer in 0 ..< keys.len:
let p = "h." & $layer & "."
let norm1 = refNorm(x, weights[p & "ln_1.weight"], weights[p & "ln_1.bias"], graph.epsilon)
let qkv = refLinear(norm1, weights[p & "attn.c_attn.weight"], weights[p & "attn.c_attn.bias"])
let attn = refAttention(qkv, keys[layer], values[layer], position, hidden, heads)
x = refAdd(x, refLinear(attn, weights[p & "attn.c_proj.weight"], weights[p & "attn.c_proj.bias"]))
let norm2 = refNorm(x, weights[p & "ln_2.weight"], weights[p & "ln_2.bias"], graph.epsilon)
let fc = refLinear(norm2, weights[p & "mlp.c_fc.weight"], weights[p & "mlp.c_fc.bias"])
x = refAdd(x, refLinear(refGelu(fc), weights[p & "mlp.c_proj.weight"], weights[p & "mlp.c_proj.bias"]))
x = refNorm(x, weights["ln_f.weight"], weights["ln_f.bias"], graph.epsilon)
result = newSeq[float32](graph.symbols["vocab"])
for token in 0 ..< result.len:
for i in 0 ..< hidden:
result[token] += x[i]*wte[token*hidden+i]
function writeSafe(path: string, weights: Table[string, seq[float32]], shapes: Table[string, seq[int]])
var header = newJObject()
var offset = 0
var names: seq[string]
for name, data in weights:
names.add(name)
header[name] = %*{"dtype": "F32", "shape": shapes[name], "data_offsets": [offset, offset+data.len*4]}
offset += data.len*4
var text = $header
while text.len mod 8 != 0:
text.add(' ')
var f = open(path, fmWrite)
defer: f.close()
var length = uint64(text.len)
discard f.writeBuffer(addr length, 8)
f.write(text)
for name in names:
let data = weights[name]
discard f.writeBuffer(unsafeAddr data[0], data.len*4)
function checkFull(graph: FirProgram, root: string)
var weights = initTable[string, seq[float32]]()
var shapes = initTable[string, seq[int]]()
var rng = initRand(8343)
for name, buffer in graph.buffers:
if buffer.kind != "weight":
continue
for layer in 0 ..< (if buffer.tensorName.contains("{i}"): graph.symbols["layers"] else: 1):
let tensor = buffer.tensorName.replace("{i}", $layer)
weights[tensor] = fgRandom(rng, buffer.elements)
shapes[tensor] = buffer.dims
let safePath = root / "model.safetensors"
writeSafe(safePath, weights, shapes)
var model = firLoadGraph(firF32Weights(graph), safePath, newBackend("cpu"))
defer: firFreeGraph(model)
var tokens: seq[int]
for i in 0 ..< graph.symbols["positions"]:
let token = rng.rand(graph.symbols["vocab"]-1)
tokens.add(token)
let actual = firGraphStep(model, token)
let expected = refForward(graph, weights, tokens)
fgClose(actual, expected, "full graph at position " & $i)
doAssert model.dispatched > 0
doAssert model.sessions == tokens.len and model.logitsReads == tokens.len
firGraphReset(model)
fgClose(firGraphStep(model, tokens[0]), refForward(graph, weights, @[tokens[0]]), "fp32 reset")
# Independently dequantize the framework's packed Q4, transpose back to HF
# layout and compare the whole resident graph with scalar float formulas.
var rounded = weights
for name, buffer in graph.buffers:
if buffer.dtype != "q4":
continue
let k = buffer.dims[0]
let n = buffer.dims[1]
for layer in 0 ..< graph.symbols["layers"]:
let tensor = buffer.tensorName.replace("{i}", $layer)
var transposed = newSeq[float32](buffer.elements)
for j in 0 ..< n:
for i in 0 ..< k:
transposed[j*k+i] = weights[tensor][i*n+j]
let packed = quantizeInt4AsymmetricGrouped(transposed, @[n, k])
doAssert packed.kind == qkInt4Asymmetric and packed.groupSize == 32
let restored = dequantizeTensor(packed)
var hf = newSeq[float32](buffer.elements)
for j in 0 ..< n:
for i in 0 ..< k:
hf[i*n+j] = restored[j*k+i]
rounded[tensor] = hf
var resident = firNewGraph(graph, weights, newBackend("cpu"), firDefaultRegistry())
defer: firFreeGraph(resident)
doAssert resident.q4.len == 4*graph.symbols["layers"]
doAssert resident.fp32.len == 2, "tied head must reuse wte"
var prefix: seq[int]
for token in tokens:
prefix.add(token)
fgClose(firGraphStep(resident, token), refForward(graph, rounded, prefix), "Q4 resident graph")
doAssert resident.sessions == tokens.len and resident.logitsReads == tokens.len
firGraphReset(resident)
fgClose(firGraphStep(resident, tokens[0]), refForward(graph, rounded, @[tokens[0]]), "Q4 reset")
echo "PASS resident Q4 full graph vs independently dequantized scalar reference; transpose, tied head, one session/read per token, reset (tol 1e-4)"
for i in 1 ..< tokens.len:
discard firGraphStep(model, tokens[i])
var rejected = false
try:
discard firGraphStep(model, 0)
except ValueError:
rejected = true
doAssert rejected, "context overflow accepted"
# Header-only routing, then malformed shape and truncated payload rejection.
writeFile(root / "model.fir", readFile(graph.sourcePath))
writeFile(root / "tokenizer.json", "{}")
var opts = fmDefaultOptions()
opts.modelsDir = root
opts.backend = "cpu"
let detected = fmDetect(opts)
doAssert detected.supported and detected.route == frFir, detected.errors.join("; ")
let header = fmReadSafeHeader(safePath)
rejected = false
try:
discard fmLoadSafeF32(safePath, header, "wte.weight", @[1, 1])
except ValueError:
rejected = true
doAssert rejected, "wrong safetensors shape accepted"
let valid = readFile(safePath)
writeFile(safePath, valid[0 ..< valid.len-4])
rejected = false
try:
discard fmReadSafeHeader(safePath)
except ValueError:
rejected = true
doAssert rejected, "truncated safetensors accepted"
echo "PASS GPT-2 tiny full graph through production CPU VM/backend vs independent reference, all positions (tol 1e-4); safetensors, detection, context guards"
function checkQuant()
var rng = initRand(425)
let rows = 3
let cols = 35 # Odd row width and a partial second group.
var input = fgRandom(rng, rows*cols, 5'f32)
for c in 0 ..< cols:
input[c] = 2.5'f32 # Constant groups must retain their nonzero value.
let packed = quantizeInt4AsymmetricGrouped(input, @[rows, cols])
doAssert packed.data.len == rows*((cols+1) div 2)
doAssert packed.scale.len == rows*2 and packed.zero_point.len == rows*2
let restored = dequantizeTensor(packed)
for r in 0 ..< rows:
for c in 0 ..< cols:
let sc = packed.scale[r*2+c div 32]
doAssert abs(restored[r*cols+c]-input[r*cols+c]) <= sc*0.501'f32+1e-6'f32
fgClose(restored[0 ..< cols], input[0 ..< cols], "constant Q4 groups")
let ctx = newBackend("cpu")
var w = beResidentUploadQ4(ctx, packed.data, packed.scale, packed.zero_point, cols, rows, 32, 2)
defer: beResidentFreeQ4(w)
let x = fgRandom(rng, cols)
var expected = newSeq[float32](rows)
for r in 0 ..< rows:
for c in 0 ..< cols:
expected[r] += x[c]*restored[r*cols+c]
fgClose(beResidentLinearQ4(ctx, x, w, 1), expected, "odd grouped resident GEMV")
echo "PASS grouped asymmetric Q4: odd row padding, partial groups, nonzero constants and resident GEMV (tol 1e-4)"
function checkBadKernels(root: string)
# Kernel blocks in a model file must fail loudly: no silent wrong programs.
let source = readFile("models/gpt2-125m/model.fir")
var cases: seq[(string, string, string)]
cases.add(("unknown mnemonic", " mulconst 1 1 1/in0.size", " bogus 1 1"))
cases.add(("register out of range", " add 1 1 2\n ustoref32 1 0 @out0", " add 1 1 40\n ustoref32 1 0 @out0"))
cases.add(("unclosed loop", " forend\n mulconst 1 1 1/in0.size", " mulconst 1 1 1/in0.size"))
cases.add(("unknown name", "mulconst 2 2 0.044715", "mulconst 2 2 not_a_symbol"))
cases.add(("wrong buffer size", "in1:1,", "in1:2,"))
cases.add(("wrong arity", "kernel g2.add in=2 out=1", "kernel g2.add in=3 out=1"))
cases.add(("wrong argument count", " mul 2 1 1\n mul 2 2 1", " mul 2 1\n mul 2 2 1"))
for (label, find, replacement) in cases:
if not source.contains(find):
raise newException(ValueError, "bad-kernel fixture does not match model.fir: " & label)
let bad = root / "bad.fir"
writeFile(bad, source.replace(find, replacement))
var rejected = false
try:
let graph = firParse(bad)
firValidate(graph, firDefaultRegistry(), firGraphBindings(graph))
except ValueError:
rejected = true
if not rejected:
raise newException(ValueError, "bad kernel accepted: " & label)
echo "PASS FIR kernel blocks reject unknown mnemonics, bad registers, unbalanced loops, unknown names, wrong sizes and arities"
checkQuant()
let root = getTempDir() / ("fir_graph_" & $getCurrentProcessId())
createDir(root)
try:
let path = root / "config.json"
for (hidden, heads) in [(8, 2), (15, 3)]:
let cfg = %*{"model_type": "gpt2", "n_embd": hidden, "n_head": heads, "n_layer": 2, "n_positions": 5, "vocab_size": 19, "layer_norm_epsilon": 0.0003, "activation_function": "gelu_new"}
writeFile(path, $cfg)
let graph = firParse("models/gpt2-125m/model.fir", path)
checkMicro(graph)
checkFull(graph, root)
let large = firParse("models/gpt2-125m/model.fir")
firValidate(large, firDefaultRegistry(), firGraphBindings(large))
echo "PASS GPT-2 real config FIR validation, every micro program built without weights"
checkBadKernels(root)
finally:
removeDir(root)