Add MLX's affine 4-bit quantizer and a check against shipped tensors

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),

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.

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;

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"});
+}

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;
+}

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");
+}

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");
+}