diff --git a/build.zig b/build.zig index 054e68828..c1113d878 100644 --- a/build.zig +++ b/build.zig @@ -144,6 +144,18 @@ pub fn build(b: *std.Build) void { const cost_run = b.addRunArtifact(cost_cases); b.step("test-cost-rule", "CPU scripted window and interleaved head timing").dependOn(&cost_run.step); test_step.dependOn(&cost_run.step); + // `tf-merge-check --quantizer` and the tests of the MLX quantizer it uses: host only on every OS + const merge_mod = b.createModule(.{ .root_source_file = b.path("zig/src/core/lora_merge_dir.zig"), .target = target, .optimize = optimize, .link_libc = true }); + const merge_tests = b.addRunArtifact(b.addTest(.{ .root_module = merge_mod })); + b.step("test-merge", "MLX quantizer and tf-merge-check tests (no model, no GPU)").dependOn(&merge_tests.step); + const merge_check = b.addExecutable(.{ .name = "tf-merge-check", .root_module = b.createModule(.{ + .root_source_file = b.path("zig/tests/merge_check.zig"), + .target = target, + .optimize = optimize, + .link_libc = true, + .imports = &.{.{ .name = "lora_merge", .module = merge_mod }}, + }) }); + b.step("tf-merge-check", "Check the MLX quantizer against shipped tensors from their dense source").dependOn(&b.addInstallArtifact(merge_check, .{}).step); switch (target.result.os.tag) { .macos => metalTargets(b, target, optimize, draft_ids, build_options, test_step), .linux => cuda_build.targets(b, target, optimize, draft_ids, build_options), diff --git a/docs/quantization.md b/docs/quantization.md index 6c5f91f8e..19be7de4d 100644 --- a/docs/quantization.md +++ b/docs/quantization.md @@ -46,6 +46,13 @@ Otherwise the model computes those projections separately with their declared fo CUDA tensor-parallel splits retain complete groups and packed-word boundaries, and memory admission counts the actual packed precision. Changing a checkpoint, precision or kernel implementation also changes the conditions under which cached states remain valid. +## Quantizing like MLX + +`tensorfold` quantizes a dense tensor to MLX affine 4-bit (groups of 32, 64 or 128) with MLX's own scheme: float32 arithmetic, the bias at the larger-magnitude edge of a group, a signed scale that puts zero on the grid, round half away from zero, codes from the unrounded scale and bias with bf16 only for storage. +When the dense source and the conversion scheme match how a checkpoint was made, this reproduces its words, scales and biases exactly. +`zig build tf-merge-check` builds a check of that: `tf-merge-check --quantizer DENSE SHIPPED MODULE...` quantizes each named dense tensor and compares it with the shipped one, reporting the share of equal scales, biases, codes and whole groups. +A dense tensor of another shape, size or dtype than the shipped projection, and any non-finite value or stored scale, is refused naming the tensor. + ## Scope and qualification These additions cover MLX affine safetensors checkpoints for the dense Qwen family. diff --git a/zig/src/core/affine4_host.zig b/zig/src/core/affine4_host.zig index 35d7274c7..add049da6 100644 --- a/zig/src/core/affine4_host.zig +++ b/zig/src/core/affine4_host.zig @@ -40,6 +40,34 @@ pub fn quantize(values: []const f32, words: []u32, scales: []u16, biases: []u16) } } +/// MLX's affine 4-bit quantizer, groups of `gs`; the first group with a non-finite range, scale or bias, else null. +pub fn quantizeMlx(comptime gs: usize, values: []const f32, words: []u32, scales: []u16, biases: []u16) ?usize { + @memset(words, 0); + for (0..values.len / gs) |g| { + const v = values[g * gs ..][0..gs]; + for (v) |x| if (!std.math.isFinite(x)) return g; + const lo = std.mem.min(f32, v); + const hi = std.mem.max(f32, v); + const step = @max((hi - lo) / 15, 1e-7); + if (!std.math.isFinite(step)) return g; + const low_edge = @abs(lo) > @abs(hi); + const signed = if (low_edge) step else -step; + const edge = if (low_edge) lo else hi; + const q0 = @round(edge / signed); + const s: f32 = if (q0 == 0) signed else edge / q0; + const b: f32 = if (q0 == 0) 0 else edge; + if (!std.math.isFinite(s) or s == 0 or !std.math.isFinite(b)) return g; + scales[g] = bf16of(s); + biases[g] = bf16of(b); + if (!std.math.isFinite(f32of(scales[g])) or !std.math.isFinite(f32of(biases[g]))) return g; + for (v, g * gs..) |x, i| { + const q: u32 = @intFromFloat(std.math.clamp(@round((x - b) / s), 0, 15)); + words[i / 8] |= q << @intCast(4 * (i % 8)); + } + } + return null; +} + test "dequantized codes land within half a step, a flat group exactly on its bias; bf16 ties go to even" { var values: [2 * group]f32 = undefined; for (&values, 0..) |*v, i| v.* = @sin(@as(f32, @floatFromInt(i)) * 0.37) * 0.05; diff --git a/zig/src/core/lora_fixture.zig b/zig/src/core/lora_fixture.zig new file mode 100644 index 000000000..e45f413cc --- /dev/null +++ b/zig/src/core/lora_fixture.zig @@ -0,0 +1,90 @@ +//! Synthetic dense and shipped checkpoint folders and PEFT adapters for the merge tests: no model, no GPU. +const std = @import("std"); +const se = @import("shard_edit.zig"); +const host4 = @import("affine4_host.zig"); +const lm = @import("lora_quant.zig"); +const Io = std.Io; + +pub const io = std.testing.io; +pub const base_config = "{\"model_type\":\"x\",\"quantization\":{\"group_size\":64,\"bits\":4,\"mode\":\"affine\"}}"; +pub const module = "backbone.layers.0.mixer.in_proj"; + +pub const Scratch = struct { + arena: std.heap.ArenaAllocator, + root: []const u8, + + pub fn init() !Scratch { + var arena = std.heap.ArenaAllocator.init(std.testing.allocator); + const root = try std.fmt.allocPrint(arena.allocator(), ".tf-merge-test-{d}", .{Io.Clock.awake.now(io).toNanoseconds()}); + return .{ .arena = arena, .root = root }; + } + + pub fn deinit(s: *Scratch) void { + Io.Dir.cwd().deleteTree(io, s.root) catch {}; + s.arena.deinit(); + } + + pub fn path(s: *Scratch, parts: []const []const u8) ![]const u8 { + return std.fs.path.join(s.arena.allocator(), &.{ s.root, std.fs.path.join(s.arena.allocator(), parts) catch unreachable }); + } + + pub fn read(s: *Scratch, folder: []const u8, name: []const u8) ![]const u8 { + return Io.Dir.cwd().readFileAlloc(io, try s.path(&.{ folder, name }), s.arena.allocator(), .limited(1 << 24)); + } + + pub fn write(s: *Scratch, folder: []const u8, name: []const u8, data: []const u8) !void { + const p = try s.path(&.{ folder, name }); + try Io.Dir.cwd().createDirPath(io, std.fs.path.dirname(p).?); + try Io.Dir.cwd().writeFile(io, .{ .sub_path = p, .data = data }); + } + + /// The names in `folder` (relative to the scratch root), sorted. + pub fn list(s: *Scratch, folder: []const u8) ![]const []const u8 { + var names: std.ArrayList([]const u8) = .empty; + var d = Io.Dir.cwd().openDir(io, try s.path(&.{folder}), .{ .iterate = true }) catch return &.{}; + defer d.close(io); + var it = d.iterate(); + while (try it.next(io)) |e| try names.append(s.arena.allocator(), try s.arena.allocator().dupe(u8, e.name)); + std.mem.sort([]const u8, names.items, {}, struct { + fn lt(_: void, a: []const u8, b: []const u8) bool { + return std.mem.lessThan(u8, a, b); + } + }.lt); + return names.items; + } +}; + +/// The three leaves and the dense weight of a projection whose values are `vals` (each exactly a bf16). +pub fn leaves(a: std.mem.Allocator, name: []const u8, n: usize, k: usize, vals: []const f32, dense: *std.ArrayList(se.New), quantized: *std.ArrayList(se.New)) !void { + const bf = try a.alloc(u16, n * k); + for (bf, vals) |*h, v| h.* = host4.bf16of(v); + const q: lm.Quant = .{ .words = try a.alloc(u32, n * k / 8), .scales = try a.alloc(u16, n * k / 64), .biases = try a.alloc(u16, n * k / 64) }; + if (lm.quantize(64, vals, q) != null) return error.NonFinite; + const g = try a.dupe(usize, &.{ n, k / 64 }); + try dense.append(a, .{ .name = try std.fmt.allocPrint(a, "{s}.weight", .{name}), .dtype = .bf16, .shape = try a.dupe(usize, &.{ n, k }), .bytes = std.mem.sliceAsBytes(bf) }); + try quantized.append(a, .{ .name = try std.fmt.allocPrint(a, "{s}.weight", .{name}), .dtype = .u32, .shape = try a.dupe(usize, &.{ n, k / 8 }), .bytes = std.mem.sliceAsBytes(q.words) }); + try quantized.append(a, .{ .name = try std.fmt.allocPrint(a, "{s}.scales", .{name}), .dtype = .bf16, .shape = g, .bytes = std.mem.sliceAsBytes(q.scales) }); + try quantized.append(a, .{ .name = try std.fmt.allocPrint(a, "{s}.biases", .{name}), .dtype = .bf16, .shape = g, .bytes = std.mem.sliceAsBytes(q.biases) }); +} + +/// Pseudo-random bf16-exact values. +pub fn randomValues(a: std.mem.Allocator, count: usize, seed: u64) ![]f32 { + var prng = std.Random.DefaultPrng.init(seed); + const v = try a.alloc(f32, count); + for (v) |*x| x.* = host4.f32of(host4.bf16of((prng.random().float(f32) - 0.5) * 0.2)); + return v; +} + +/// A one-projection pair of folders `dense` and `base` (one shard, no index) around `vals` (n x k, k a multiple of 64). +pub fn makeSingle(s: *Scratch, n: usize, k: usize, vals: []const f32) ![]const u8 { + const a = s.arena.allocator(); + var dense: std.ArrayList(se.New) = .empty; + var quantized: std.ArrayList(se.New) = .empty; + try leaves(a, module, n, k, vals, &dense, &quantized); + try Io.Dir.cwd().createDirPath(io, try s.path(&.{"dense"})); + try Io.Dir.cwd().createDirPath(io, try s.path(&.{"base"})); + try se.create(std.testing.allocator, try s.path(&.{ "dense", "model.safetensors" }), &.{}, dense.items); + try se.create(std.testing.allocator, try s.path(&.{ "base", "model.safetensors" }), &.{.{ "format", "mlx" }}, quantized.items); + try s.write("base", "config.json", base_config); + return try s.path(&.{"base"}); +} diff --git a/zig/src/core/lora_gate.zig b/zig/src/core/lora_gate.zig new file mode 100644 index 000000000..b15332990 --- /dev/null +++ b/zig/src/core/lora_gate.zig @@ -0,0 +1,32 @@ +//! The check of the MLX quantizer against shipped tensors: each dense weight quantized and compared word for word. +const std = @import("std"); +const lm = @import("lora_quant.zig"); +const lp = @import("lora_shards.zig"); +const Allocator = std.mem.Allocator; +const Io = std.Io; + +/// The quantizer's output on each module's dense weight against the shipped words, scales and biases. +pub fn quantizerCheck(gpa: Allocator, arena: Allocator, io: Io, w: *Io.Writer, dense_dir: []const u8, shipped_dir: []const u8, modules: []const []const u8) !bool { + var dense = try lp.Shards.init(gpa, arena, io, w, dense_dir); + defer dense.deinit(); + var shipped = try lp.Shards.init(gpa, arena, io, w, shipped_dir); + defer shipped.deinit(); + const quant = try lp.quantConfig(arena, io, w, shipped_dir); + var ok = true; + for (modules) |module| { + const lay = try lp.layoutOf(arena, w, &shipped, quant, module); + const j: lp.Job = .{ .module = module, .a = "", .b = "", .n = lay.n, .k = lay.k, .group = lay.group, .dtype = lay.dtype }; + const x = try lp.denseOf(gpa, w, &dense, j); + defer gpa.free(x); + const old = try lp.quantOf(gpa, &shipped, j); + defer lp.free(gpa, old); + const new = try lm.allocQuant(gpa, j.n, j.k, j.group); + defer lm.freeQuant(gpa, new); + if (lm.quantize(j.group, x, new)) |bad| return lp.refuse(w, module, try std.fmt.allocPrint(arena, "group {d}: a non-finite value, range or scale", .{bad})); + const a = lm.agreement(j.group, lm.viewOf(old), lm.viewOf(new)); + try w.print("{s} [{d} x {d}, group {d}]: scales {d:.4}%, biases {d:.4}%, codes {d:.4}%, whole groups {d:.4}%: {s}\n", .{ module, j.n, j.k, j.group, a.scales * 100, a.biases * 100, a.codes * 100, a.groups * 100, if (a.identical()) "identical" else "DIFFERS" }); + ok = ok and a.identical(); + } + try w.print("{s}\n", .{if (ok) "quantizer check: byte-identical on every tensor" else "quantizer check: FAILED"}); + return ok; +} diff --git a/zig/src/core/lora_merge_dir.zig b/zig/src/core/lora_merge_dir.zig new file mode 100644 index 000000000..c33a5f20b --- /dev/null +++ b/zig/src/core/lora_merge_dir.zig @@ -0,0 +1,9 @@ +//! Host tools for MLX affine 4-bit checkpoints: the quantizer and its check against shipped tensors. +pub const safetensors = @import("safetensors.zig"); +pub const math = @import("lora_quant.zig"); +pub const plan = @import("lora_shards.zig"); +pub const gate = @import("lora_gate.zig"); + +test { + _ = @import("lora_quantizer_test.zig"); +} diff --git a/zig/src/core/lora_quant.zig b/zig/src/core/lora_quant.zig new file mode 100644 index 000000000..f7215de51 --- /dev/null +++ b/zig/src/core/lora_quant.zig @@ -0,0 +1,111 @@ +//! One MLX affine 4-bit projection on the host: its quantizer at the engine's groups, comparison and f32 decoding. +const std = @import("std"); +const host4 = @import("affine4_host.zig"); +const st = @import("safetensors.zig"); + +/// One projection as stored: packed words, then bf16 scales and biases a group of `group` values. +pub const Quant = struct { words: []u32, scales: []u16, biases: []u16 }; + +pub const View = struct { words: []const u32, scales: []const u16, biases: []const u16 }; + +pub fn viewOf(q: Quant) View { + return .{ .words = q.words, .scales = q.scales, .biases = q.biases }; +} + +/// Room for `n * k` values in groups of `group`; `freeQuant` returns it. +pub fn allocQuant(gpa: std.mem.Allocator, n: usize, k: usize, group: usize) !Quant { + const count = try std.math.mul(usize, n, k); + const words = try gpa.alloc(u32, count / 8); + errdefer gpa.free(words); + const scales = try gpa.alloc(u16, count / group); + errdefer gpa.free(scales); + return .{ .words = words, .scales = scales, .biases = try gpa.alloc(u16, count / group) }; +} + +pub fn freeQuant(gpa: std.mem.Allocator, q: Quant) void { + gpa.free(q.words); + gpa.free(q.scales); + gpa.free(q.biases); +} + +/// MLX's quantizer over `values` (rows of k) at a group the engine serves; the first unusable group, else null. +pub fn quantize(group: usize, values: []const f32, out: Quant) ?usize { + return switch (group) { + 32 => host4.quantizeMlx(32, values, out.words, out.scales, out.biases), + 64 => host4.quantizeMlx(64, values, out.words, out.scales, out.biases), + else => host4.quantizeMlx(128, values, out.words, out.scales, out.biases), + }; +} + +/// How alike two quantizations of one tensor are: the fraction of equal scales, biases and codes, and of equal groups. +pub const Agreement = struct { + scales: f64, + biases: f64, + codes: f64, + groups: f64, + + pub fn identical(a: Agreement) bool { + return a.scales == 1 and a.biases == 1 and a.codes == 1 and a.groups == 1; + } +}; + +/// The number of groups whose scale, bias or codes differ between `a` and `b`. +pub fn differingGroups(group: usize, a: View, b: View) usize { + const gw = group / 8; + var n: usize = 0; + for (0..a.scales.len) |g| { + const same = a.scales[g] == b.scales[g] and a.biases[g] == b.biases[g] and std.mem.eql(u32, a.words[g * gw ..][0..gw], b.words[g * gw ..][0..gw]); + n += @intFromBool(!same); + } + return n; +} + +pub fn agreement(group: usize, a: View, b: View) Agreement { + const gw = group / 8; + const groups = a.scales.len; + var s: usize = 0; + var bi: usize = 0; + var c: usize = 0; + var both: usize = 0; + for (0..groups) |g| { + const se = a.scales[g] == b.scales[g]; + const be = a.biases[g] == b.biases[g]; + var ce: usize = 0; + for (0..group) |i| { + const x = (a.words[g * gw + i / 8] >> @intCast(4 * (i % 8))) & 0xf; + const y = (b.words[g * gw + i / 8] >> @intCast(4 * (i % 8))) & 0xf; + if (x == y) ce += 1; + } + s += @intFromBool(se); + bi += @intFromBool(be); + c += ce; + both += @intFromBool(se and be and ce == group); + } + const gf: f64 = @floatFromInt(groups); + return .{ .scales = @as(f64, @floatFromInt(s)) / gf, .biases = @as(f64, @floatFromInt(bi)) / gf, .codes = @as(f64, @floatFromInt(c)) / (gf * @as(f64, @floatFromInt(group))), .groups = @as(f64, @floatFromInt(both)) / gf }; +} + +/// A tensor's values as f32: f32, bf16 and f16 are the dtypes checkpoints and PEFT adapters save. +pub fn toF32(dtype: st.DType, bytes: []const u8, out: []f32) error{ UnsupportedDType, BadLength }!void { + switch (dtype) { + .f32, .bf16, .f16 => {}, + else => return error.UnsupportedDType, + } + if (bytes.len != std.math.mul(usize, out.len, dtype.size()) catch return error.BadLength) return error.BadLength; + switch (dtype) { + .f32 => for (out, 0..) |*o, i| { + o.* = @bitCast(std.mem.readInt(u32, bytes[4 * i ..][0..4], .little)); + }, + .bf16 => for (out, 0..) |*o, i| { + o.* = host4.f32of(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)); + }, + else => for (out, 0..) |*o, i| { + o.* = @floatCast(@as(f16, @bitCast(std.mem.readInt(u16, bytes[2 * i ..][0..2], .little)))); + }, + } +} + +test { + _ = @import("lora_quant_test.zig"); + _ = @import("lora_vectors_test.zig"); +} diff --git a/zig/src/core/lora_quant_test.zig b/zig/src/core/lora_quant_test.zig new file mode 100644 index 000000000..b7ce85749 --- /dev/null +++ b/zig/src/core/lora_quant_test.zig @@ -0,0 +1,127 @@ +//! The MLX quantizer and its comparison helpers on synthetic tensors: no model, no GPU. +const std = @import("std"); +const host4 = @import("affine4_host.zig"); +const lq = @import("lora_quant.zig"); + +const gpa = std.testing.allocator; + +fn alloc(n: usize, group: usize) !lq.Quant { + return .{ .words = try gpa.alloc(u32, n / 8), .scales = try gpa.alloc(u16, n / group), .biases = try gpa.alloc(u16, n / group) }; +} + +fn free(q: lq.Quant) void { + gpa.free(q.words); + gpa.free(q.scales); + gpa.free(q.biases); +} + +fn same(a: lq.Quant, b: lq.Quant) bool { + return std.mem.eql(u32, a.words, b.words) and std.mem.eql(u16, a.scales, b.scales) and std.mem.eql(u16, a.biases, b.biases); +} + +/// A bf16 dense weight (as the public checkpoints ship) of pseudo-random values, as f32. +fn dense(n: usize, k: usize, seed: u64) ![]f32 { + var prng = std.Random.DefaultPrng.init(seed); + const w = try gpa.alloc(f32, n * k); + for (w) |*v| v.* = host4.f32of(host4.bf16of((prng.random().float(f32) - 0.5) * 0.2)); + return w; +} + +fn codes(q: lq.Quant, i: usize) u32 { + return (q.words[i / 8] >> @intCast(4 * (i % 8))) & 0xf; +} + +test "the quantizer's hand-checked groups: edge at the maximum, at the minimum, and a tie rounded away from zero" { + var v: [32]f32 = @splat(0); + const q = try alloc(32, 32); + defer free(q); + // 0..15 twice over: the larger-magnitude edge is the maximum, so the scale is -1 and the bias 15 + for (v[0..16], 0..) |*x, i| x.* = @floatFromInt(i); + for (v[16..], 0..) |*x, i| x.* = @floatFromInt(i); + try std.testing.expect(lq.quantize(32, &v, q) == null); + try std.testing.expectEqual(host4.bf16of(-1), q.scales[0]); + try std.testing.expectEqual(host4.bf16of(15), q.biases[0]); + for (0..32) |i| try std.testing.expectEqual(@as(u32, 15 - @as(u32, @intFromFloat(v[i]))), codes(q, i)); + // -15..0: the minimum is the larger edge, scale +1, bias -15, codes count up + for (v[0..16], 0..) |*x, i| x.* = @as(f32, @floatFromInt(i)) - 15; + for (v[16..], 0..) |*x, i| x.* = @as(f32, @floatFromInt(i)) - 15; + try std.testing.expect(lq.quantize(32, &v, q) == null); + try std.testing.expectEqual(host4.bf16of(1), q.scales[0]); + try std.testing.expectEqual(host4.bf16of(-15), q.biases[0]); + for (0..32) |i| try std.testing.expectEqual(@as(u32, @intCast(i % 16)), codes(q, i)); + // 2.5 sits halfway between two codes of the 0..15 group (scale -1, bias 15): 12.5 rounds away to 13, not to even 12 + @memset(&v, 0); + v[0] = 15; + v[1] = 2.5; + try std.testing.expect(lq.quantize(32, &v, q) == null); + try std.testing.expectEqual(@as(u32, 13), codes(q, 1)); + // a flat group is its own bias, exactly, and every code is zero + @memset(&v, 0.5); + try std.testing.expect(lq.quantize(32, &v, q) == null); + for (0..32) |i| try std.testing.expectEqual(@as(u32, 0), codes(q, i)); + try std.testing.expectEqual(host4.bf16of(0.5), q.biases[0]); + // a group of zeros keeps zero as an exact grid point + @memset(&v, 0); + try std.testing.expect(lq.quantize(32, &v, q) == null); + try std.testing.expectEqual(@as(u16, 0), q.biases[0]); +} + +test "agreement counts what differs: one scale, one code" { + const w = try dense(2, 64, 4); + defer gpa.free(w); + const q = try alloc(w.len, 64); + defer free(q); + try std.testing.expect(lq.quantize(64, w, q) == null); + const other = try alloc(w.len, 64); + defer free(other); + @memcpy(other.words, q.words); + @memcpy(other.scales, q.scales); + @memcpy(other.biases, q.biases); + other.scales[0] ^= 1; + other.words[8] ^= 1; + const ag = lq.agreement(64, lq.viewOf(q), lq.viewOf(other)); + try std.testing.expectEqual(@as(f64, 0.5), ag.scales); + try std.testing.expectEqual(@as(f64, 1), ag.biases); + try std.testing.expectEqual(@as(f64, 0), ag.groups); + try std.testing.expect(ag.codes < 1 and !ag.identical()); +} + +test "adapter factors read as f32 from f32, bf16 and f16 bytes" { + var out: [2]f32 = undefined; + try lq.toF32(.bf16, &.{ 0x80, 0x3f, 0x00, 0xc0 }, &out); // 1.0, -2.0 + try std.testing.expectEqualSlices(f32, &.{ 1, -2 }, &out); + try lq.toF32(.f16, &.{ 0x00, 0x3c, 0x00, 0xc0 }, &out); + try std.testing.expectEqualSlices(f32, &.{ 1, -2 }, &out); + try lq.toF32(.f32, &.{ 0, 0, 0x80, 0x3f, 0, 0, 0, 0xc0 }, &out); + try std.testing.expectEqualSlices(f32, &.{ 1, -2 }, &out); + try std.testing.expectError(error.UnsupportedDType, lq.toF32(.u8, &.{ 0, 0 }, &out)); +} + +test "the quantizer refuses a non-finite value, a range that overflows and a scale that collapses, naming the group" { + var v: [128]f32 = @splat(0.25); + const q = try alloc(128, 64); + defer free(q); + try std.testing.expect(lq.quantize(64, &v, q) == null); + v[70] = std.math.nan(f32); + try std.testing.expectEqual(@as(?usize, 1), lq.quantize(64, &v, q)); + v[70] = std.math.inf(f32); + try std.testing.expectEqual(@as(?usize, 1), lq.quantize(64, &v, q)); + v[70] = 0.25; + v[3] = -3e38; // hi - lo overflows + v[4] = 3e38; + try std.testing.expectEqual(@as(?usize, 0), lq.quantize(64, &v, q)); + @memset(&v, 1e38); // edge / q0 overflows: a zero scale + try std.testing.expectEqual(@as(?usize, 0), lq.quantize(64, &v, q)); +} + +test "a scale or bias that rounds to a bf16 infinity is refused, naming the group" { + var v: [128]f32 = @splat(3.3e38); + @memset(v[0..64], 0.25); + v[64] = 3.4e38; // finite in f32, above the largest finite bf16: the stored bias would be infinite + const q = try alloc(128, 64); + defer free(q); + try std.testing.expectEqual(@as(?usize, 1), lq.quantize(64, &v, q)); + v[64] = host4.f32of(0x7f7f); // the largest finite bf16 is fine + try std.testing.expect(std.math.isFinite(host4.f32of(0x7f7f))); + try std.testing.expectEqual(@as(?usize, null), lq.quantize(64, &v, q)); +} diff --git a/zig/src/core/lora_quantizer_test.zig b/zig/src/core/lora_quantizer_test.zig new file mode 100644 index 000000000..309c77f39 --- /dev/null +++ b/zig/src/core/lora_quantizer_test.zig @@ -0,0 +1,116 @@ +//! `tf-merge-check --quantizer` on synthetic dense and shipped folders: validation of both inputs, no model, no GPU. +const std = @import("std"); +const se = @import("shard_edit.zig"); +const fx = @import("lora_fixture.zig"); +const dir = @import("lora_merge_dir.zig"); +const Io = std.Io; + +const io = fx.io; +const module = fx.module; + +fn contains(hay: []const u8, needle: []const u8) bool { + return std.mem.indexOf(u8, hay, needle) != null; +} + +/// The quantizer check over the scratch folders `dense` and `base`: its result and what it printed. +fn check(s: *fx.Scratch, msg: []u8) !struct { res: anyerror!bool, out: []const u8 } { + var w = Io.Writer.fixed(msg); + const res = dir.gate.quantizerCheck(std.testing.allocator, s.arena.allocator(), io, &w, try s.path(&.{"dense"}), try s.path(&.{"base"}), &.{module}); + return .{ .res = res, .out = w.buffered() }; +} + +test "the quantizer check passes on a dense source and its own quantization, and fails on another weight" { + var s = try fx.Scratch.init(); + defer s.deinit(); + _ = try fx.makeSingle(&s, 8, 128, try fx.randomValues(s.arena.allocator(), 8 * 128, 4)); + var buf: [4096]u8 = undefined; + var r = try check(&s, &buf); + try std.testing.expect(try r.res); + try std.testing.expect(contains(r.out, "identical") and contains(r.out, "byte-identical on every tensor")); + // another dense source: the shipped tensor no longer comes from it + const a = s.arena.allocator(); + var dense: std.ArrayList(se.New) = .empty; + var quantized: std.ArrayList(se.New) = .empty; + try fx.leaves(a, module, 8, 128, try fx.randomValues(a, 8 * 128, 9), &dense, &quantized); + try se.create(std.testing.allocator, try s.path(&.{ "dense", "model.safetensors" }), &.{}, dense.items); + r = try check(&s, &buf); + try std.testing.expect(!(try r.res)); + try std.testing.expect(contains(r.out, "DIFFERS")); +} + +/// A dense file holding `module`.weight with the given dtype, shape and zeroed payload. +fn denseOf(s: *fx.Scratch, dtype: @import("safetensors.zig").DType, shape: []const usize) !void { + var count: usize = dtype.size(); + for (shape) |d| count *= d; + const bytes = try s.arena.allocator().alloc(u8, count); + @memset(bytes, 0); + try se.create(std.testing.allocator, try s.path(&.{ "dense", "model.safetensors" }), &.{}, &.{.{ .name = module ++ ".weight", .dtype = dtype, .shape = shape, .bytes = bytes }}); +} + +test "a dense tensor with fewer or more elements, reshaped, or of an unsupported dtype is refused naming it" { + const Case = struct { dtype: @import("safetensors.zig").DType, shape: []const usize, why: []const u8 }; + const cases = [_]Case{ + .{ .dtype = .bf16, .shape = &.{ 8, 64 }, .why = "dense shape" }, + .{ .dtype = .bf16, .shape = &.{ 8, 256 }, .why = "dense shape" }, + .{ .dtype = .bf16, .shape = &.{ 16, 64 }, .why = "dense shape" }, + .{ .dtype = .bf16, .shape = &.{8 * 128}, .why = "dense shape" }, + .{ .dtype = .f64, .shape = &.{ 8, 128 }, .why = "dense dtype" }, + }; + for (cases) |c| { + var s = try fx.Scratch.init(); + defer s.deinit(); + _ = try fx.makeSingle(&s, 8, 128, try fx.randomValues(s.arena.allocator(), 8 * 128, 4)); + try denseOf(&s, c.dtype, c.shape); + var buf: [4096]u8 = undefined; + const r = try check(&s, &buf); + try std.testing.expectError(error.Refused, r.res); + try std.testing.expect(contains(r.out, module ++ ".weight") and contains(r.out, c.why)); + } +} + +test "a directory standing in for the index, a shard or the config is refused, not read" { + const Stand = enum { dense_index, dense_shard, shipped_index, shipped_shard, config }; + for ([_]Stand{ .dense_index, .dense_shard, .shipped_index, .shipped_shard, .config }) |t| { + var s = try fx.Scratch.init(); + defer s.deinit(); + _ = try fx.makeSingle(&s, 8, 128, try fx.randomValues(s.arena.allocator(), 8 * 128, 4)); + const target = switch (t) { + .dense_index => .{ "dense", "model.safetensors.index.json" }, + .dense_shard => .{ "dense", "model.safetensors" }, + .shipped_index => .{ "base", "model.safetensors.index.json" }, + .shipped_shard => .{ "base", "model.safetensors" }, + .config => .{ "base", "config.json" }, + }; + if (t != .dense_index and t != .shipped_index) try Io.Dir.cwd().deleteFile(io, try s.path(&.{ target[0], target[1] })); + try Io.Dir.cwd().createDirPath(io, try s.path(&.{ target[0], target[1] })); + var buf: [4096]u8 = undefined; + const r = try check(&s, &buf); + try std.testing.expectError(error.Refused, r.res); + try std.testing.expect(contains(r.out, target[1]) and contains(r.out, "not a regular file")); + } +} + +fn checkWith(gpa: std.mem.Allocator, dense: []const u8, shipped: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + var buf: [8192]u8 = undefined; + var w = Io.Writer.fixed(&buf); + _ = try dir.gate.quantizerCheck(gpa, arena.allocator(), io, &w, dense, shipped, &.{module}); +} + +test "an allocation that fails at any step is unwound without a leak or a close of an unopened shard" { + var s = try fx.Scratch.init(); + defer s.deinit(); + _ = try fx.makeSingle(&s, 8, 128, try fx.randomValues(s.arena.allocator(), 8 * 128, 4)); + const dense = try s.path(&.{"dense"}); + const shipped = try s.path(&.{"base"}); + var at: usize = 0; + while (at < 1000) : (at += 1) { + var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{ .fail_index = at }); + if (checkWith(failing.allocator(), dense, shipped)) |_| { + try std.testing.expect(!failing.has_induced_failure); + break; + } else |e| try std.testing.expectEqual(error.OutOfMemory, e); + } + try std.testing.expect(at > 5 and at < 1000); // it did fail at some steps, and ran clean once none failed +} diff --git a/zig/src/core/lora_shards.zig b/zig/src/core/lora_shards.zig new file mode 100644 index 000000000..f0a983c1a --- /dev/null +++ b/zig/src/core/lora_shards.zig @@ -0,0 +1,178 @@ +//! A checkpoint folder read for the merge tools: its shards, quantization table, projection layouts and dense weights. +const std = @import("std"); +const st = @import("safetensors.zig"); +const mlx = @import("quant/mlx.zig"); +const lq = @import("lora_quant.zig"); +const Allocator = std.mem.Allocator; +const Io = std.Io; + +pub const Refused = error{Refused}; + +/// Says why on `w`, naming the tensor, and returns `Refused`. +pub fn refuse(w: *Io.Writer, name: []const u8, why: []const u8) Refused { + w.print("merge: {s}: {s}\n", .{ name, why }) catch {}; + return error.Refused; +} + +/// Whether directory entry `e` is a regular file, or a symlink to one; a dangling symlink is not. +pub fn isRegular(dir: Io.Dir, io: Io, e: Io.Dir.Entry) bool { + const kind = if (e.kind == .sym_link) (dir.statFile(io, e.name, .{}) catch return false).kind else e.kind; + return kind == .file; +} + +/// Refuses `path` unless it is a regular file or a symlink to one; false when nothing is there. +pub fn requireRegular(io: Io, w: *Io.Writer, path: []const u8) !bool { + const cwd = Io.Dir.cwd(); + const link = cwd.statFile(io, path, .{ .follow_symlinks = false }) catch |e| switch (e) { + error.FileNotFound => return false, + else => return e, + }; + const kind = if (link.kind == .sym_link) (cwd.statFile(io, path, .{}) catch return refuse(w, std.fs.path.basename(path), "a dangling symlink")).kind else link.kind; + if (kind != .file) return refuse(w, std.fs.path.basename(path), "not a regular file"); + return true; +} + +/// A non-empty file name with no separator, not `.` or `..`: a path that stays inside its folder. +pub fn plainName(name: []const u8) bool { + if (name.len == 0 or std.mem.eql(u8, name, ".") or std.mem.eql(u8, name, "..")) return false; + return std.mem.indexOfAny(u8, name, "/\\\x00") == null; +} + +/// `path` as a JSON object, or null when it is missing or not an object. +pub fn readObject(arena: Allocator, io: Io, w: *Io.Writer, path: []const u8) !?std.json.ObjectMap { + if (!try requireRegular(io, w, path)) return null; + const text = try Io.Dir.cwd().readFileAlloc(io, path, arena, .limited(1 << 26)); + const v = try std.json.parseFromSliceLeaky(std.json.Value, arena, text, .{}); + return if (v == .object) v.object else null; +} + +/// A checkpoint folder's shards: the index's `weight_map`, or the one `model.safetensors`. +pub const Shards = struct { + gpa: Allocator, + io: Io, + w: *Io.Writer, + dir: []const u8, + weight_map: ?std.json.ObjectMap, + files: std.StringArrayHashMapUnmanaged(st.File) = .empty, + + pub fn init(gpa: Allocator, arena: Allocator, io: Io, w: *Io.Writer, dir: []const u8) !Shards { + const path = try std.fs.path.join(arena, &.{ dir, "model.safetensors.index.json" }); + const index = (try readObject(arena, io, w, path)) orelse return .{ .gpa = gpa, .io = io, .w = w, .dir = dir, .weight_map = null }; + const map = index.get("weight_map"); + if (map == null or map.? != .object) return refuse(w, "model.safetensors.index.json", "no weight_map object"); + for (map.?.object.keys(), map.?.object.values()) |name, v| { + if (v != .string or !plainName(v.string)) return refuse(w, name, "its shard is not a plain file name inside the folder"); + } + return .{ .gpa = gpa, .io = io, .w = w, .dir = dir, .weight_map = map.?.object }; + } + + pub fn deinit(s: *Shards) void { + for (s.files.values()) |*f| f.close(s.io); + s.files.deinit(s.gpa); + } + + /// The file name of the shard holding `name`, or null. + pub fn shardOf(s: *const Shards, name: []const u8) ?[]const u8 { + const map = s.weight_map orelse return "model.safetensors"; + const v = map.get(name) orelse return null; + return if (v == .string) v.string else null; + } + + /// The shard file `shard`, opened once, or null when the folder has none. + pub fn open(s: *Shards, shard: []const u8) !?*st.File { + if (s.files.getPtr(shard)) |f| return f; + const path = try std.fs.path.join(s.gpa, &.{ s.dir, shard }); + defer s.gpa.free(path); + if (!try requireRegular(s.io, s.w, path)) return null; + var file = try st.File.open(s.gpa, s.io, path); + errdefer file.close(s.io); + try s.files.put(s.gpa, shard, file); + return s.files.getPtr(shard); + } + + /// The tensor `name` as stored, or null when the folder has none. + pub fn get(s: *Shards, name: []const u8) !?st.Tensor { + const shard = s.shardOf(name) orelse return null; + const f = (try s.open(shard)) orelse return null; + return f.get(name); + } +}; + +/// One projection a tool works on, with its layout in the shipped checkpoint; `a` and `b` name adapter tensors. +pub const Job = struct { module: []const u8, a: []const u8 = "", b: []const u8 = "", n: usize, k: usize, group: usize, dtype: st.DType }; + +/// The quantization table of `dir`'s config.json, refused unless it is MLX affine. +pub fn quantConfig(arena: Allocator, io: Io, w: *Io.Writer, dir: []const u8) !mlx.Config { + const cfg = (try readObject(arena, io, w, try std.fs.path.join(arena, &.{ dir, "config.json" }))) orelse return refuse(w, "config.json", "missing or not an object"); + const qobj = cfg.get("quantization") orelse cfg.get("quantization_config") orelse return refuse(w, "config.json", "no quantization: the checkpoint is not MLX affine"); + if (qobj != .object) return refuse(w, "config.json", "quantization is not an object"); + return (try mlx.detect(arena, .{ .root = cfg, .quantization = qobj.object })) orelse refuse(w, "config.json", "quantization is not MLX affine"); +} + +/// A shipped projection's shape: `n` rows of `k` values in groups of `group`, its weight words as `dtype`. +pub const Layout = struct { n: usize, k: usize, group: usize, dtype: st.DType }; + +/// `module`'s stored layout in `shipped`, with every leaf present, in one shard, 4-bit and bf16 tables. +pub fn layoutOf(arena: Allocator, w: *Io.Writer, shipped: *Shards, quant: mlx.Config, module: []const u8) !Layout { + const leaves = [3][]const u8{ try std.fmt.allocPrint(arena, "{s}.weight", .{module}), try std.fmt.allocPrint(arena, "{s}.scales", .{module}), try std.fmt.allocPrint(arena, "{s}.biases", .{module}) }; + const shard = shipped.shardOf(leaves[0]) orelse return refuse(w, leaves[0], "not in the checkpoint: no such quantized tensor (a fused projection stores another name)"); + var t: [3]st.Tensor = undefined; + for (&t, leaves) |*x, leaf| { + if (!std.mem.eql(u8, shipped.shardOf(leaf) orelse "", shard)) return refuse(w, leaf, "not in the shard of its weight"); + x.* = (try shipped.get(leaf)) orelse return refuse(w, leaf, "not in the checkpoint: no such quantized tensor"); + } + const width = quant.width(module) catch return refuse(w, module, "quantization the merge does not read"); + if (width.bits != 4) return refuse(w, module, "only 4-bit affine is merged"); + const group: usize = width.group; + if ((t[0].dtype != .u32 and t[0].dtype != .i32) or t[0].rank != 2) return refuse(w, leaves[0], "not packed 32-bit words"); + if (t[1].dtype != .bf16 or t[2].dtype != .bf16 or t[1].rank != 2 or t[2].rank != 2 or !std.mem.eql(usize, t[1].shape[0..2], t[2].shape[0..2])) return refuse(w, leaves[1], "scales and biases are not bf16 tables of rank 2 and one shape"); + const n = t[0].dim(0); + const k = std.math.mul(usize, t[1].dim(1), group) catch return refuse(w, module, "dimensions overflow"); + _ = std.math.mul(usize, n, k) catch return refuse(w, module, "dimensions overflow"); + if (n == 0 or k == 0) return refuse(w, module, "an empty projection"); + if (t[1].dim(0) != n or t[0].dim(1) != k / 8) return refuse(w, module, "packed shape does not match its scales and group"); + return .{ .n = n, .k = k, .group = group, .dtype = t[0].dtype }; +} + +/// The dense weight of `module`, refused unless an `n` x `k` f32, bf16 or f16 matrix with that many bytes. +pub fn denseTensor(w: *Io.Writer, dense: *Shards, module: []const u8, n: usize, k: usize) !st.Tensor { + var buf: [256]u8 = undefined; + const name = try std.fmt.bufPrint(&buf, "{s}.weight", .{module}); + const t = (try dense.get(name)) orelse return refuse(w, name, "not in the dense base"); + if (t.rank != 2 or t.dim(0) != n or t.dim(1) != k) return refuse(w, name, "dense shape does not match the shipped projection"); + if (t.dtype != .f32 and t.dtype != .bf16 and t.dtype != .f16) return refuse(w, name, "dense dtype is not f32, bf16 or f16"); + if (t.bytes.len != n * k * t.dtype.size()) return refuse(w, name, "dense payload does not match its shape"); + return t; +} + +fn allFinite(v: []const f32) bool { + for (v) |x| if (!std.math.isFinite(x)) return false; + return true; +} + +/// The dense weight of `j` as f32, checked against its layout and finite; the caller frees it. +pub fn denseOf(gpa: Allocator, w: *Io.Writer, dense: *Shards, j: Job) ![]f32 { + const t = try denseTensor(w, dense, j.module, j.n, j.k); + const out = try gpa.alloc(f32, j.n * j.k); + errdefer gpa.free(out); + try lq.toF32(t.dtype, t.bytes, out); + if (!allFinite(out)) return refuse(w, j.module, "the dense weight holds a non-finite value"); + return out; +} + +/// A projection's leaves copied into aligned memory; `freeQuant` returns it. +pub fn quantOf(gpa: Allocator, s: *Shards, j: Job) !lq.Quant { + var buf: [256]u8 = undefined; + const q = try lq.allocQuant(gpa, j.n, j.k, j.group); + errdefer lq.freeQuant(gpa, q); + inline for (.{ .{ "weight", q.words }, .{ "scales", q.scales }, .{ "biases", q.biases } }) |leaf| { + const t = (try s.get(try std.fmt.bufPrint(&buf, "{s}." ++ leaf[0], .{j.module}))) orelse return error.MissingTensor; + if (t.bytes.len != std.mem.sliceAsBytes(leaf[1]).len) return error.UnexpectedTensor; + @memcpy(std.mem.sliceAsBytes(leaf[1]), t.bytes); + } + return q; +} + +pub fn free(gpa: Allocator, q: lq.Quant) void { + lq.freeQuant(gpa, q); +} diff --git a/zig/src/core/lora_vectors_test.zig b/zig/src/core/lora_vectors_test.zig new file mode 100644 index 000000000..8baefcca6 --- /dev/null +++ b/zig/src/core/lora_vectors_test.zig @@ -0,0 +1,80 @@ +//! MLX quantizer vectors made independently (a float32 reference of the published algorithm) on synthetic groups. +const std = @import("std"); +const host4 = @import("affine4_host.zig"); + +// Nine groups: division vs reciprocal (two), unrounded scale and bias, a tie, both edges, flat, tiny. +const inputs = + \\bb2b 3d9b 3c3a 3c48 bd2d bc98 3ca9 bd04 3cc9 3d88 3d0a 3c4e bd13 3c44 3c35 bd06 bcdc 3dc3 bb8e 3b2b 3d4e 3c1c bb8d bc63 bb19 bcc9 b9ee bcdb 3c2b bccc bdac 3d04 bce9 3c69 3c29 3d9e 3d4f 3a64 3da0 bd41 3d2a bd48 3c38 3cf4 3d0e 3d12 bd5f 3c9e bd0f bd06 3d59 bdff bda2 3ce7 3d43 bd29 3d88 bbec bceb 3d41 bd3b 3d46 bc3f bcb5 + \\3ce7 3da2 3d0e 3ac1 3d21 bdce 3d03 bbc0 bb33 3d25 3ccb bd7d 3c03 bd65 b9c4 3d83 3d33 bd48 ba53 bcf7 bc11 bca7 bc70 bcc1 3d85 3c5e 3dac 3d8d bd56 bdee 3d9c 3d96 3cf6 bc1e 3dc3 bd84 bd49 3d05 bc90 bd15 3d24 baea 3cea bdbe 3c36 bda8 3d51 3d9a bd99 3db1 bd47 bd7e 3dee bd90 bd61 3c85 3cd3 bc80 3d2f 3acc 3bb8 3d7f 3d4d 3de2 + \\bd4b bd8b bd0e 3d0b 3cad 3d3e 3d4d bd5c 3d0a bd9b 3cf3 3c29 bb88 bbad bb4e 3c6d bc46 bbed bd42 3c9a 3d63 b9fd 3c8b 3d79 bd56 bcba 3cef bd44 3c8a bd9f bdb0 bc9e 3c6e bd9c 3d19 bd99 bc60 3d58 bd2d 3cfc bcc5 3d95 bcb6 3d20 bcdf bd51 be0b 3c20 3d22 3d65 3ba2 3c94 3d90 bcb4 3cfe bbd8 bdc7 bda7 3ca0 bd4c 3d3a bd30 3c86 3b2f + \\c150 4188 c150 c0c0 3f80 4100 4170 c110 c000 40a0 4140 c140 c0a0 4000 4110 4180 c100 bf80 40c0 4150 c130 c080 4040 4120 4188 c0e0 0000 40e0 4160 c120 c040 4080 4130 c150 c0c0 3f80 4100 4170 c110 c000 40a0 4140 c140 c0a0 4000 4110 4180 c100 bf80 40c0 4150 c130 c080 4040 4120 4188 c0e0 0000 40e0 4160 c120 c040 4080 4130 + \\c170 c160 c150 c140 c130 c120 c110 c100 c0e0 c0c0 c0a0 c080 c040 c000 bf80 0000 c170 c160 c150 c140 c130 c120 c110 c100 c0e0 c0c0 c0a0 c080 c040 c000 bf80 0000 c170 c160 c150 c140 c130 c120 c110 c100 c0e0 c0c0 c0a0 c080 c040 c000 bf80 0000 c170 c160 c150 c140 c130 c120 c110 c100 c0e0 c0c0 c0a0 c080 c040 c000 bf80 0000 + \\0000 3f80 4000 4040 4080 40a0 40c0 40e0 4100 4110 4120 4130 4140 4150 4160 4170 0000 3f80 4000 4040 4080 40a0 40c0 40e0 4100 4110 4120 4130 4140 4150 4160 4170 0000 3f80 4000 4040 4080 40a0 40c0 40e0 4100 4110 4120 4130 4140 4150 4160 4170 0000 3f80 4000 4040 4080 40a0 40c0 40e0 4100 4110 4120 4130 4140 4150 4160 4170 + \\be80 3d80 3ec0 3f30 3f80 be00 3e40 3f00 3f50 3f90 0000 3ea0 3f20 3f70 be40 3e00 3ee0 3f40 3f88 bd80 3e80 3f10 3f60 be80 3d80 3ec0 3f30 3f80 be00 3e40 3f00 3f50 3f90 0000 3ea0 3f20 3f70 be40 3e00 3ee0 3f40 3f88 bd80 3e80 3f10 3f60 be80 3d80 3ec0 3f30 3f80 be00 3e40 3f00 3f50 3f90 0000 3ea0 3f20 3f70 be40 3e00 3ee0 3f40 + \\bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 bf00 + \\b089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 3109 0000 3089 +; +const expected = + \\7a86aae9 7aa6abeb 89ad99f7 b37a7979 6f9dfaa7 a5ccba5c 6cb30d76 78c6c78e 3c63 bdff + \\86f58636 48c7c658 a999a8b5 33fc3274 b96bc196 34e7e685 7cd0cb2d 05488596 bc6e 3dee + \\6edcd756 baaabc4c ebaeb799 944b6c86 c7e94d4b b068d8f8 ac8fbaed ab7d6b43 3c5e be0b + \\e158cf0f 148cf36a 47bf26ad 7be259d0 ae158cf3 d148cf36 047bf26a 37be259d bff2 4188 + \\76543210 fedcba98 76543210 fedcba98 76543210 fedcba98 76543210 fedcba98 3f80 c170 + \\89abcdef 01234567 89abcdef 01234567 89abcdef 01234567 89abcdef 01234567 bf80 4170 + \\7ad158bf be259c03 f369d147 37ad158b 7be259c0 bf369d14 037ad158 47be259c bdc0 3f90 + \\00000000 00000000 00000000 00000000 00000000 00000000 00000000 00000000 b3d7 bf00 + \\00000000 00000000 00000000 00000000 00000000 00000000 00000000 00000000 b3d7 0000 +; + +fn hexWords(comptime T: type, line: []const u8, out: []T) !void { + var it = std.mem.tokenizeScalar(u8, line, ' '); + var i: usize = 0; + while (it.next()) |tok| : (i += 1) out[i] = try std.fmt.parseInt(T, tok, 16); + try std.testing.expectEqual(out.len, i); +} + +/// The same algorithm with the step multiplied by 1/15 instead of divided: the variant the vectors must tell apart. +fn reciprocal(values: *const [64]f32, words: *[8]u32, scale: *u16, bias: *u16) void { + const lo = std.mem.min(f32, values); + const hi = std.mem.max(f32, values); + const step = @max((hi - lo) * (1.0 / 15.0), 1e-7); + const low_edge = @abs(lo) > @abs(hi); + const signed = if (low_edge) step else -step; + const edge = if (low_edge) lo else hi; + const q0 = @round(edge / signed); + const s: f32 = if (q0 == 0) signed else edge / q0; + const b: f32 = if (q0 == 0) 0 else edge; + scale.* = host4.bf16of(s); + bias.* = host4.bf16of(b); + @memset(words, 0); + for (values, 0..) |x, i| words[i / 8] |= @as(u32, @intFromFloat(std.math.clamp(@round((x - b) / s), 0, 15))) << @intCast(4 * (i % 8)); +} + +test "the quantizer reproduces the independently produced words, scales and biases, and the reciprocal variant does not" { + var in_lines = std.mem.tokenizeScalar(u8, inputs, '\n'); + var ex_lines = std.mem.tokenizeScalar(u8, expected, '\n'); + var group: usize = 0; + var reciprocal_misses: usize = 0; + while (in_lines.next()) |line| : (group += 1) { + var bits: [64]u16 = undefined; + try hexWords(u16, line, &bits); + var values: [64]f32 = undefined; + for (&values, bits) |*v, b| v.* = host4.f32of(b); + var want: [10]u32 = undefined; + try hexWords(u32, ex_lines.next().?, &want); + var words: [8]u32 = undefined; + var scale: [1]u16 = undefined; + var bias: [1]u16 = undefined; + try std.testing.expect(host4.quantizeMlx(64, &values, &words, &scale, &bias) == null); + try std.testing.expectEqualSlices(u32, want[0..8], &words); + try std.testing.expectEqual(@as(u16, @intCast(want[8])), scale[0]); + try std.testing.expectEqual(@as(u16, @intCast(want[9])), bias[0]); + var rw: [8]u32 = undefined; + var rs: u16 = undefined; + var rb: u16 = undefined; + reciprocal(&values, &rw, &rs, &rb); + if (!std.mem.eql(u32, &rw, &words) or rs != scale[0] or rb != bias[0]) reciprocal_misses += 1; + } + try std.testing.expectEqual(@as(usize, 9), group); + try std.testing.expect(reciprocal_misses >= 2); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 1a4a18b62..3af7f3611 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -12,6 +12,7 @@ pub const ids_json = @import("ids_json.zig"); pub const affine4_host = @import("affine4_host.zig"); pub const affine4_fold = @import("affine4_fold.zig"); pub const shard_edit = @import("shard_edit.zig"); +pub const lora_merge = @import("lora_merge_dir.zig"); pub const quant = @import("quant/quant.zig"); pub const registry = @import("registry.zig"); pub const tuning = @import("tuning.zig"); @@ -21,6 +22,7 @@ test { _ = affine4_host; _ = affine4_fold; _ = shard_edit; + _ = lora_merge; _ = quant; _ = registry; _ = tuning; diff --git a/zig/src/tensorfold.zig b/zig/src/tensorfold.zig index eeaee2fc7..69f80feb9 100644 --- a/zig/src/tensorfold.zig +++ b/zig/src/tensorfold.zig @@ -8,6 +8,7 @@ pub const checkpoint_metal = checkpoint; pub const qwen27 = @import("families/qwen27/qwen27.zig"); pub const checkpoint = @import("core/checkpoint_metal.zig"); pub const shard_edit = @import("core/shard_edit.zig"); +pub const lora_merge = @import("core/lora_merge_dir.zig"); pub const npy = @import("core/npy.zig"); pub const ids_json = @import("core/ids_json.zig"); pub const lanes = @import("lanes"); @@ -24,6 +25,7 @@ pub const flashnext_pack = @import("families/flashnext/pack.zig"); test { std.testing.refAllDecls(@This()); + _ = lora_merge; _ = @import("families/nemotron/prefill_kernels.zig"); // Compile the prompt-chunk sources in tests. _ = @import("families/nemotron/simd_attention.zig"); // The pre-M5 attention rewrite finds its lines. _ = @import("families/nemotron/weights.zig"); diff --git a/zig/tests/merge_check.zig b/zig/tests/merge_check.zig new file mode 100644 index 000000000..22e136107 --- /dev/null +++ b/zig/tests/merge_check.zig @@ -0,0 +1,23 @@ +//! tf-merge-check --quantizer DENSE SHIPPED MODULE...: the quantizer against shipped tensors from their dense source. +const std = @import("std"); +const tools = @import("lora_merge"); +const Io = std.Io; + +const usage = "usage: tf-merge-check --quantizer DENSE SHIPPED MODULE...\n"; + +pub fn main(init: std.process.Init) !u8 { + const args = try init.minimal.args.toSlice(init.arena.allocator()); + var out_buf: [4096]u8 = undefined; + var out = Io.File.stdout().writerStreaming(init.io, &out_buf); + defer out.interface.flush() catch {}; + const w = &out.interface; + if (args.len < 5 or !std.mem.eql(u8, args[1], "--quantizer")) { + try w.writeAll(usage); + return 2; + } + const ok = tools.gate.quantizerCheck(init.gpa, init.arena.allocator(), init.io, w, args[2], args[3], @ptrCast(args[4..])) catch |e| { + if (e != error.Refused) try w.print("tf-merge-check: {s}\n", .{@errorName(e)}); + return 1; + }; + return if (ok) 0 else 1; +}