diff --git a/CHANGELOG.md b/CHANGELOG.md index 951b9e8c56..26ea86726d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,9 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel `input_file` and Messages `document` PDFs, and the fields it left out (`thinking`, `min_rows`, grammar timings, `exact_mode`, `stop_sequence`, health's `memory.limit` and 503 body, `/v1/models`' context and reasoning) (#2299). +- Tests: `zig build tf-tp-check` builds a checker for the dumps each rank of a mixed Metal and CUDA tensor-parallel run + writes, and `zig build test-tp-check` runs its tests (#2321, part of #1236). Thanks to @ThinkOffApp. + - Qwen3.8-27B: a prompt frame below 128 rows (`TF_PROMPT_ROWS=16`) is refused before loading, with the setting, the value and the fix named, instead of a bare "TargetBinding" (#2298). diff --git a/build.zig b/build.zig index 3aa689962a..95bcf01c65 100644 --- a/build.zig +++ b/build.zig @@ -189,6 +189,7 @@ pub fn build(b: *std.Build) void { @import("zig/build/release_paths_step.zig").tests(b, test_step); @import("zig/build/sessions.zig").add(b, test_step); @import("zig/build/tp.zig").add(b, target, optimize, test_step, tp_options); + @import("zig/build/tp_check.zig").add(b, target, test_step); @import("zig/build/gradient_host.zig").add(b); @import("zig/build/living_tools_host.zig").steps(b); cuda_build.hostTests(b, draft_ids, build_options, test_step, target.result.os.tag != .macos); diff --git a/zig/build/tp.zig b/zig/build/tp.zig index e0e34b3951..4d2c278b97 100644 --- a/zig/build/tp.zig +++ b/zig/build/tp.zig @@ -3,7 +3,7 @@ const std = @import("std"); /// The CUDA runtime without the engine's kernels (tp only needs the driver, memory, streams and NCCL). -fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { +pub fn runtime(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { const options = @import("cuda_runtime.zig").kernelOptions(b, false, null); const cuda = b.createModule(.{ .root_source_file = b.path("zig/src/cuda/root.zig"), .target = target, .optimize = optimize, .link_libc = true }); cuda.addImport("gpu_memory", @import("gpu_memory.zig").module(b, target)); diff --git a/zig/build/tp_check.zig b/zig/build/tp_check.zig new file mode 100644 index 0000000000..7203afb779 --- /dev/null +++ b/zig/build/tp_check.zig @@ -0,0 +1,23 @@ +//! tf-tp-check (zig/tests/tp_check_main.zig), the mixed-TP gate over rank dumps, and its host tests. + +const std = @import("std"); +const tp = @import("tp.zig"); + +fn module(b: *std.Build, root: []const u8, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path(root), .target = target, .optimize = optimize, .link_libc = true }); + m.addImport("tp", tp.module(b, target, optimize, tp.runtime(b, target, optimize))); + return m; +} + +pub fn add(b: *std.Build, target: std.Build.ResolvedTarget, test_step: *std.Build.Step) void { + const host = b.graph.host; + const gate_tests = b.step("test-tp-check", "Mixed-TP gate: dump format, host reference, run-to-run and one-box agreement checks"); + const tests = b.addRunArtifact(b.addTest(.{ .root_module = module(b, "zig/tests/tp_check_test.zig", host, .debug) })); + gate_tests.dependOn(&tests.step); + test_step.dependOn(&tests.step); + const files = b.addRunArtifact(b.addTest(.{ .root_module = module(b, "zig/tests/tp_check_file_test.zig", host, .ReleaseSafe) })); + gate_tests.dependOn(&files.step); // about 2 GiB while it runs, so not in every `zig build test` + // ReleaseSafe whatever -Doptimize says: a gate keeps its safety checks + const exe = b.addExecutable(.{ .name = "tf-tp-check", .root_module = module(b, "zig/tests/tp_check_main.zig", target, .ReleaseSafe) }); + b.step("tf-tp-check", "The mixed-TP gate (#1236 rules 1, 2, 4) over each rank's collective and logits dumps").dependOn(&b.addInstallArtifact(exe, .{}).step); +} diff --git a/zig/tests/tp_check.zig b/zig/tests/tp_check.zig new file mode 100644 index 0000000000..1c6fa539b8 --- /dev/null +++ b/zig/tests/tp_check.zig @@ -0,0 +1,35 @@ +//! Mixed-TP gate (#1236, rules 1, 2 and 4): each rank dumps its collectives and logits; the checks read the dumps. +const dump = @import("tp_check_dump.zig"); +const rules = @import("tp_check_rules.zig"); +const agree = @import("tp_check_agree.zig"); + +pub const magic = dump.magic; +pub const version = dump.version; +pub const phase_version = dump.phase_version; +pub const canonical_version = dump.canonical_version; +pub const Phase = dump.Phase; +pub const Kind = dump.Kind; +pub const Call = dump.Call; +pub const Error = dump.Error; +pub const Mode = dump.Mode; +pub const Header = dump.Header; +pub const Collective = dump.Collective; +pub const Logits = dump.Logits; +pub const writeHeader = dump.writeHeader; +pub const writeCollective = dump.writeCollective; +pub const writeCollectiveVersion = dump.writeCollectiveVersion; +pub const writeLogits = dump.writeLogits; +pub const TokenHash = dump.TokenHash; +pub const Dump = dump.Dump; +pub const parse = dump.parse; +pub const Mismatch = rules.Mismatch; +pub const ReferenceReport = rules.ReferenceReport; +pub const checkReference = rules.checkReference; +pub const RepeatReport = rules.RepeatReport; +pub const checkRepeat = rules.checkRepeat; +pub const repeatPassed = rules.repeatPassed; +pub const Split = agree.Split; +pub const AgreeReport = agree.AgreeReport; +pub const chosenHash = agree.chosenHash; +pub const checkAgree = agree.checkAgree; +pub const agreeFailure = agree.agreeFailure; diff --git a/zig/tests/tp_check_agree.zig b/zig/tests/tp_check_agree.zig new file mode 100644 index 0000000000..c778923b01 --- /dev/null +++ b/zig/tests/tp_check_agree.zig @@ -0,0 +1,181 @@ +//! Mixed-TP rule 4: divergence and greedy agreement between a mixed run and one box. +const std = @import("std"); +const dump = @import("tp_check_dump.zig"); +const Mode = dump.Mode; +const Logits = dump.Logits; +const TokenHash = dump.TokenHash; +const Dump = dump.Dump; +/// Where greedy decoding of one prompt first chose a different token. +pub const Split = struct { prompt: u32, position: u64 }; + +pub const AgreeReport = struct { + mode: Mode, + /// Prompts in the solo dump, and (prompt, position) pairs present in both dumps. + prompts: usize = 0, + positions: usize = 0, + /// Solo positions with no mixed record (greedy: only before a prompt's first split or gap); any is a bad run. + missing: usize = 0, + /// Greedy only: prompts ended at a gap rather than a split (owned). + gaps: []Split = &.{}, + /// Positions with the same chosen token; greedy counts only before each prompt's first split, where prefixes match. + greedy_same: usize = 0, + /// Greedy only: each prompt's first split, in prompt order (owned; free with `deinit`). + splits: []Split = &.{}, + /// Teacher-forced only: whether the mixed run's `fed_hash` is the hash of the solo run's chosen tokens. + fed_ok: ?bool = null, + /// Positions in the divergence figures: all when fed this run's tokens; greedy, each prompt up to its first split. + compared: usize = 0, + top1_same: usize = 0, + max_abs: f64 = 0, + mean_abs: f64 = 0, + /// KL(solo || mixed) per position over the softmax, in nats: mean and worst. + kl_mean: f64 = 0, + kl_max: f64 = 0, + /// Positions left out of the divergence figures because the vocab sizes differ or a logit is not finite. + skipped: usize = 0, + /// Positions where the mixed run's own logits are not finite: a fault in that run, so the report fails. + mixed_nonfinite: usize = 0, + + pub fn deinit(r: *AgreeReport, a: std.mem.Allocator) void { + a.free(r.splits); + a.free(r.gaps); + } +}; + +fn argmax(l: Logits) usize { + var best: usize = 0; + for (1..l.vocab()) |i| { + if (l.at(i) > l.at(best)) best = i; + } + return best; +} + +/// log-sum-exp over the logits, in f64. +fn logSumExp(l: Logits) f64 { + var m: f64 = -std.math.inf(f64); + for (0..l.vocab()) |i| m = @max(m, @as(f64, l.at(i))); + var s: f64 = 0; + for (0..l.vocab()) |i| s += @exp(@as(f64, l.at(i)) - m); + return m + @log(s); +} + +fn finite(l: Logits) bool { + for (0..l.vocab()) |i| { + if (!std.math.isFinite(l.at(i))) return false; + } + return true; +} + +const Key = struct { prompt: u32, position: u64 }; + +fn byKey(logits: []const Logits, x: u32, y: u32) bool { + const p = logits[x]; + const q = logits[y]; + return if (p.prompt != q.prompt) p.prompt < q.prompt else p.position < q.position; +} + +/// The dump's logits records as indices in (prompt, position) order (owned). +fn ordered(a: std.mem.Allocator, d: Dump) ![]u32 { + const order = try a.alloc(u32, d.logits.len); + for (order, 0..) |*o, i| o.* = @intCast(i); + std.mem.sort(u32, order, d.logits, byKey); + return order; +} + +/// `TokenHash` over a dump's chosen tokens in (prompt, position) order: the `fed_hash` a run fed from it carries. +pub fn chosenHash(a: std.mem.Allocator, d: Dump) !u64 { + const order = try ordered(a, d); + defer a.free(order); + var h: TokenHash = .{}; + for (order) |i| h.add(d.logits[i].prompt, d.logits[i].position, d.logits[i].chosen); + return h.h; +} + +/// Rule 4: `mixed` vs `solo` matched by (prompt, position); greedy: agreement and splits; fed: every position. +pub fn checkAgree(a: std.mem.Allocator, mixed: Dump, solo: Dump) !AgreeReport { + var r: AgreeReport = .{ .mode = mixed.header.mode }; + errdefer r.deinit(a); + var at: std.AutoHashMapUnmanaged(Key, u32) = .empty; + defer at.deinit(a); + for (mixed.logits, 0..) |m, i| try at.put(a, .{ .prompt = m.prompt, .position = m.position }, @intCast(i)); + const forced = mixed.header.mode == .teacher_forced; + if (forced) r.fed_ok = mixed.header.fed_hash == try chosenHash(a, solo); + const order = try ordered(a, solo); + defer a.free(order); + var splits: std.ArrayList(Split) = .empty; + defer splits.deinit(a); + var gaps: std.ArrayList(Split) = .empty; + defer gaps.deinit(a); + var abs_sum: f64 = 0; + var abs_count: f64 = 0; + var kl_sum: f64 = 0; + var prompt: ?u32 = null; + // greedy: set at a prompt's first split or gap; after it the runs may decode different prefixes + var stopped = false; + for (order) |si| { + const s = solo.logits[si]; + if (prompt != s.prompt) { + prompt = s.prompt; + stopped = false; + r.prompts += 1; + } + const mi = at.get(.{ .prompt = s.prompt, .position = s.position }) orelse { + // greedy: past a split the runs decode different text, so a shorter mixed run is normal there + if (forced) { + r.missing += 1; + } else if (!stopped) { + r.missing += 1; + stopped = true; + try gaps.append(a, .{ .prompt = s.prompt, .position = s.position }); + } + continue; + }; + const m = mixed.logits[mi]; + r.positions += 1; + if (!forced and stopped) continue; + const same = m.chosen == s.chosen; + if (same) r.greedy_same += 1; + if (!forced and !same) { + stopped = true; + try splits.append(a, .{ .prompt = s.prompt, .position = s.position }); + } + if (forced and r.fed_ok != true) continue; + if (!finite(m)) r.mixed_nonfinite += 1; + if (m.vocab() != s.vocab() or !finite(m) or !finite(s)) { + r.skipped += 1; + continue; + } + r.compared += 1; + if (argmax(m) == argmax(s)) r.top1_same += 1; + const lm = logSumExp(m); + const ls = logSumExp(s); + var kl: f64 = 0; + for (0..s.vocab()) |i| { + const x: f64 = s.at(i); + const y: f64 = m.at(i); + const d = @abs(x - y); + r.max_abs = @max(r.max_abs, d); + abs_sum += d; + abs_count += 1; + const lp = x - ls; + kl += @exp(lp) * (lp - (y - lm)); + } + kl_sum += kl; + r.kl_max = @max(r.kl_max, kl); + } + if (abs_count > 0) r.mean_abs = abs_sum / abs_count; + if (r.compared > 0) r.kl_mean = kl_sum / @as(f64, @floatFromInt(r.compared)); + r.splits = try splits.toOwnedSlice(a); + r.gaps = try gaps.toOwnedSlice(a); + return r; +} + +/// Why a report cannot stand, or null: rule 4 has no bound, so the report only has to be complete and non-empty. +pub fn agreeFailure(r: AgreeReport) ?[]const u8 { + if (r.missing > 0) return "the mixed dump is missing positions the solo run has, so the report is incomplete"; + if (r.positions == 0) return "no (prompt, position) appears in both dumps"; + if (r.fed_ok == false) return "this teacher-forced run was not fed this solo run's tokens, so its divergence is not reported"; + if (r.mixed_nonfinite > 0) return "the mixed run's logits are not finite"; + if (r.compared == 0) return "no position could be compared"; + return null; +} diff --git a/zig/tests/tp_check_dump.zig b/zig/tests/tp_check_dump.zig new file mode 100644 index 0000000000..3add33faa7 --- /dev/null +++ b/zig/tests/tp_check_dump.zig @@ -0,0 +1,217 @@ +//! Mixed-TP dumps: the format, its writers and its parser. Version 2 adds phases; version 3 adds canonical bootstrap. +const std = @import("std"); +const collective = @import("tp").collective; +const DType = collective.DType; +const Op = collective.Op; + +pub const magic = "TFTPDMP1"; +pub const version: u32 = 1; +pub const phase_version: u32 = 2; +pub const canonical_version: u32 = 3; +pub const Phase = enum(u8) { legacy, mixer, ffn, bootstrap }; + +pub const Kind = enum(u32) { collective = 1, logits = 2 }; + +/// Which collective a record holds; `exchange` and `send` name their peer, `recv` the rank it came from. +pub const Call = enum(u8) { all_reduce, all_gather, exchange }; + +pub const Error = error{ BadMagic, BadVersion, Truncated, BadKind, BadEnum, BadLength, OutOfMemory }; + +/// How tokens were chosen; a teacher-forced run is fed another run's tokens, and its `chosen` stays its own arg-max. +pub const Mode = enum(u8) { greedy, teacher_forced }; + +pub const Header = struct { + format_version: u32 = version, + rank: u32, + world: u32, + mode: Mode = .greedy, + /// For a teacher-forced run, `TokenHash` over the tokens it was fed; 0 otherwise. + fed_hash: u64 = 0, + /// Free text naming the box and backend, e.g. "metal m5-max" or "cuda gb10". + label: []const u8, + /// What produced the dump, by convention "commit=... model=... prompts= clock=". + provenance: []const u8 = "", +}; + +pub const Collective = struct { + phase: Phase = .legacy, + seq: u64, + layer: u32, + call: Call, + dtype: DType, + op: Op, + peer: u32, + send: []const u8, + recv: []const u8, +}; + +pub const Logits = struct { + /// Index of the prompt in the run's prompt file; positions count from 0 within each prompt. + prompt: u32, + position: u64, + chosen: u32, + /// Raw little-endian f32 bytes, `vocab` of them. + bytes: []const u8, + + pub fn vocab(l: Logits) usize { + return l.bytes.len / 4; + } + + pub fn at(l: Logits, i: usize) f32 { + return @bitCast(std.mem.readInt(u32, l.bytes[4 * i ..][0..4], .little)); + } +}; + +// ---------------------------------------------------------------- writing + +pub fn writeHeader(w: *std.Io.Writer, h: Header) std.Io.Writer.Error!void { + try w.writeAll(magic); + try w.writeInt(u32, h.format_version, .little); + try w.writeInt(u32, h.rank, .little); + try w.writeInt(u32, h.world, .little); + try w.writeAll(&.{ @backingInt(h.mode), 0, 0, 0 }); + try w.writeInt(u64, h.fed_hash, .little); + try w.writeInt(u32, @intCast(h.label.len), .little); + try w.writeAll(h.label); + try w.writeInt(u32, @intCast(h.provenance.len), .little); + try w.writeAll(h.provenance); +} + +pub fn writeCollective(w: *std.Io.Writer, c: Collective) (std.Io.Writer.Error || error{UnsupportedPhaseVersion})!void { + return writeCollectiveVersion(w, version, c); +} + +pub fn writeCollectiveVersion(w: *std.Io.Writer, format_version: u32, c: Collective) (std.Io.Writer.Error || error{UnsupportedPhaseVersion})!void { + if ((format_version != version and format_version != phase_version and format_version != canonical_version) or + (format_version == version and c.phase != .legacy) or + (format_version == phase_version and c.phase == .bootstrap)) return error.UnsupportedPhaseVersion; + try w.writeInt(u32, @backingInt(Kind.collective), .little); + try w.writeInt(u64, c.seq, .little); + try w.writeInt(u32, c.layer, .little); + try w.writeAll(&.{ @backingInt(c.call), @backingInt(c.dtype), @backingInt(c.op), @backingInt(c.phase) }); + try w.writeInt(u32, c.peer, .little); + try w.writeInt(u64, c.send.len, .little); + try w.writeInt(u64, c.recv.len, .little); + try w.writeAll(c.send); + try w.writeAll(c.recv); +} + +/// `logits` are host f32 values for one position, in vocabulary order. +pub fn writeLogits(w: *std.Io.Writer, prompt: u32, position: u64, chosen: u32, logits: []const f32) std.Io.Writer.Error!void { + try w.writeInt(u32, @backingInt(Kind.logits), .little); + try w.writeInt(u32, prompt, .little); + try w.writeInt(u64, position, .little); + try w.writeInt(u32, chosen, .little); + try w.writeInt(u32, @intCast(logits.len), .little); + for (logits) |v| try w.writeInt(u32, @bitCast(v), .little); +} + +/// FNV-1a 64 over (prompt u32, position u64, token u32) in (prompt, position) order; a fed run stores it as `fed_hash`. +pub const TokenHash = struct { + h: u64 = 0xcbf29ce484222325, + + pub fn add(t: *TokenHash, prompt: u32, position: u64, token: u32) void { + var b: [16]u8 = undefined; + std.mem.writeInt(u32, b[0..4], prompt, .little); + std.mem.writeInt(u64, b[4..12], position, .little); + std.mem.writeInt(u32, b[12..16], token, .little); + for (b) |x| { + t.h ^= x; + t.h *%= 0x100000001b3; + } + } +}; + +// ---------------------------------------------------------------- reading + +/// One rank's dump, parsed in place: the slices point into the dump's bytes. +pub const Dump = struct { + header: Header, + collectives: []Collective, + logits: []Logits, + + pub fn deinit(d: *Dump, a: std.mem.Allocator) void { + a.free(d.collectives); + a.free(d.logits); + } +}; + +const Cursor = struct { + b: []const u8, + at: usize = 0, + + fn take(c: *Cursor, n: usize) Error![]const u8 { + if (n > c.b.len - c.at) return error.Truncated; + defer c.at += n; + return c.b[c.at..][0..n]; + } + + fn int(c: *Cursor, comptime T: type) Error!T { + return std.mem.readInt(T, (try c.take(@sizeOf(T)))[0..@sizeOf(T)], .little); + } +}; + +fn enumOf(comptime E: type, v: anytype) Error!E { + return std.enums.fromInt(E, v) orelse error.BadEnum; +} + +fn phaseOf(format_version: u32, byte: u8) Error!Phase { + if (format_version == version and byte != 0) return error.BadEnum; + if (format_version == phase_version and byte > 2) return error.BadEnum; + return enumOf(Phase, byte); +} + +pub fn parse(a: std.mem.Allocator, bytes: []const u8) Error!Dump { + var c: Cursor = .{ .b = bytes }; + if (!std.mem.eql(u8, try c.take(magic.len), magic)) return error.BadMagic; + const format_version = try c.int(u32); + if (format_version != version and format_version != phase_version and format_version != canonical_version) return error.BadVersion; + const rank = try c.int(u32); + const world = try c.int(u32); + const flags = try c.take(4); + const mode = try enumOf(Mode, flags[0]); + const fed_hash = try c.int(u64); + const label = try c.take(try c.int(u32)); + const provenance = try c.take(try c.int(u32)); + var cs: std.ArrayList(Collective) = .empty; + errdefer cs.deinit(a); + var ls: std.ArrayList(Logits) = .empty; + errdefer ls.deinit(a); + while (c.at < bytes.len) { + switch (try enumOf(Kind, try c.int(u32))) { + .collective => { + const seq = try c.int(u64); + const layer = try c.int(u32); + const e = try c.take(4); + const peer = try c.int(u32); + const ns = try c.int(u64); + const nr = try c.int(u64); + const dtype = try enumOf(DType, e[1]); + if (ns % dtype.size() != 0 or nr % dtype.size() != 0) return error.BadLength; + try cs.append(a, .{ + .phase = try phaseOf(format_version, e[3]), + .seq = seq, + .layer = layer, + .call = try enumOf(Call, e[0]), + .dtype = dtype, + .op = try enumOf(Op, e[2]), + .peer = peer, + .send = try c.take(ns), + .recv = try c.take(nr), + }); + }, + .logits => { + const prompt = try c.int(u32); + const position = try c.int(u64); + const chosen = try c.int(u32); + const vocab = try c.int(u32); + try ls.append(a, .{ .prompt = prompt, .position = position, .chosen = chosen, .bytes = try c.take(@as(usize, vocab) * 4) }); + }, + } + } + return .{ + .header = .{ .format_version = format_version, .rank = rank, .world = world, .mode = mode, .fed_hash = fed_hash, .label = label, .provenance = provenance }, + .collectives = try cs.toOwnedSlice(a), + .logits = try ls.toOwnedSlice(a), + }; +} diff --git a/zig/tests/tp_check_file.zig b/zig/tests/tp_check_file.zig new file mode 100644 index 0000000000..09101debcd --- /dev/null +++ b/zig/tests/tp_check_file.zig @@ -0,0 +1,22 @@ +//! The checker retains the whole dump, but bounds each OS read below Darwin's signed read limit. +const std = @import("std"); +pub const read_chunk_bytes = 8 * 1024 * 1024; + +pub fn read(a: std.mem.Allocator, io: std.Io, path: []const u8) ![]u8 { + const file = try std.Io.Dir.cwd().openFile(io, path, .{}); + defer file.close(io); + return readSized(a, io, file, std.math.cast(usize, (try file.stat(io)).size) orelse return error.FileTooBig); +} + +/// Fill short reads and reject truncation before handing bytes to the parser. +pub fn readSized(a: std.mem.Allocator, io: std.Io, file: std.Io.File, size: usize) ![]u8 { + const bytes = try a.alloc(u8, size); + errdefer a.free(bytes); + var at: usize = 0; + while (at < bytes.len) { + const chunk = bytes[at..][0..@min(read_chunk_bytes, bytes.len - at)]; + if (try file.readPositionalAll(io, chunk, at) != chunk.len) return error.Truncated; + at += chunk.len; + } + return bytes; +} diff --git a/zig/tests/tp_check_file_test.zig b/zig/tests/tp_check_file_test.zig new file mode 100644 index 0000000000..a2fef00f56 --- /dev/null +++ b/zig/tests/tp_check_file_test.zig @@ -0,0 +1,67 @@ +//! Real filesystem coverage for the checker's bounded loader; one 2 GiB allocation, sparse disk use. +const std = @import("std"); +const load = @import("tp_check_file.zig"); +const check = @import("tp_check.zig"); +const t = std.testing; + +test "checker reads sparse file beyond 2 GiB including its nonzero tail" { + const io = t.io; + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + const file = try tmp.dir.createFile(io, "large", .{ .read = true }); + defer file.close(io); + const size: usize = (@as(usize, 1) << 31) + 17; + try file.setLength(io, size); + try file.writePositionalAll(io, "head", 0); + const tail = "past-signed-read"; + try file.writePositionalAll(io, tail, size - tail.len); + const path = try tmp.dir.realPathFileAlloc(io, "large", t.allocator); + defer t.allocator.free(path); + const bytes = try load.read(t.allocator, io, path); + defer t.allocator.free(bytes); + try t.expectEqual(size, bytes.len); + try t.expectEqualStrings("head", bytes[0..4]); + try t.expectEqual(@as(u8, 0), bytes[(@as(usize, 1) << 31) - 1]); + try t.expectEqualStrings(tail, bytes[size - tail.len ..]); + std.debug.print("checker sparse regression: logical_bytes={d} owned_buffer_bytes={d} max_read_bytes={d}\n", .{ size, size, load.read_chunk_bytes }); +} + +test "checker rejects EOF after stat and parser rejects a truncated record" { + const io = t.io; + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + const file = try tmp.dir.createFile(io, "short", .{ .read = true }); + defer file.close(io); + try file.writePositionalAll(io, "abcd", 0); + const observed_size: usize = @intCast((try file.stat(io)).size); + try file.setLength(io, 3); + try t.expectError(error.Truncated, load.readSized(t.allocator, io, file, observed_size)); + const bytes = try load.readSized(t.allocator, io, file, 3); + defer t.allocator.free(bytes); + try t.expectEqualStrings("abc", bytes); + try t.expectError(error.Truncated, check.parse(t.allocator, bytes)); +} + +test "checker bounded file loader preserves version 1 2 and 3 parsing" { + const io = t.io; + var tmp = t.tmpDir(.{}); + defer tmp.cleanup(); + for ([_]u32{ 1, 2, 3 }) |version| { + var w: std.Io.Writer.Allocating = .init(t.allocator); + defer w.deinit(); + try check.writeHeader(&w.writer, .{ .format_version = version, .rank = 0, .world = 1, .label = "file compatibility" }); + try check.writeLogits(&w.writer, 0, 0, 1, &.{ 0.5, 1.0 }); + const file = try tmp.dir.createFile(io, "version", .{ .read = true }); + defer file.close(io); + try file.writePositionalAll(io, w.written(), 0); + const path = try tmp.dir.realPathFileAlloc(io, "version", t.allocator); + defer t.allocator.free(path); + const bytes = try load.read(t.allocator, io, path); + defer t.allocator.free(bytes); + var parsed = try check.parse(t.allocator, bytes); + defer parsed.deinit(t.allocator); + try t.expectEqual(version, parsed.header.format_version); + try t.expectEqual(@as(usize, 1), parsed.logits.len); + try t.expectError(error.Truncated, check.parse(t.allocator, bytes[0 .. bytes.len - 1])); + } +} diff --git a/zig/tests/tp_check_main.zig b/zig/tests/tp_check_main.zig new file mode 100644 index 0000000000..211c370a10 --- /dev/null +++ b/zig/tests/tp_check_main.zig @@ -0,0 +1,108 @@ +//! tf-tp-check: the mixed-TP gate (#1236 rules 1, 2 and 4) over the dumps each rank writes (format in check.zig). +const std = @import("std"); +const check = @import("tp_check.zig"); +const check_file = @import("tp_check_file.zig"); + +const usage = + \\usage: tf-tp-check reference RANK0.tpdump RANK1.tpdump ... every collective equals the host reference (rule 1) + \\ tf-tp-check repeat RUN_A.tpdump RUN_B.tpdump one rank's two runs, byte for byte (rule 2) + \\ tf-tp-check agree MIXED.tpdump SOLO.tpdump divergence and greedy agreement vs one box (rule 4) + \\exit 0: reference and repeat hold; agree's report is valid (rule 4 has no bound, so 0 is not a pass) + \\exit 1: the rule does not hold, or agree cannot give valid numbers; 2: bad arguments or an unreadable dump + \\ +; + +fn bad(comptime fmt: []const u8, args: anytype) noreturn { + std.debug.print("tf-tp-check: " ++ fmt ++ "\n", args); + std.process.exit(2); +} + +fn load(a: std.mem.Allocator, io: std.Io, path: []const u8) check.Dump { + const bytes = check_file.read(a, io, path) catch |e| bad("{s}: {s}", .{ path, @errorName(e) }); + return check.parse(a, bytes) catch |e| bad("{s}: {s}", .{ path, @errorName(e) }); +} + +pub fn main(init: std.process.Init) !void { + const a = init.arena.allocator(); + const io = init.io; + const args = try init.minimal.args.toSlice(a); + if (args.len < 2 or std.mem.eql(u8, args[1], "--help") or std.mem.eql(u8, args[1], "-h")) { + std.debug.print("{s}", .{usage}); + std.process.exit(if (args.len < 2) 2 else 0); + } + var buf: [4096]u8 = undefined; + var out = std.Io.File.stdout().writerStreaming(io, &buf); + const w = &out.interface; + const cmd = args[1]; + var ok = false; + if (std.mem.eql(u8, cmd, "reference")) { + if (args.len < 4) bad("reference needs a dump per rank, rank 0 first", .{}); + const dumps = try a.alloc(check.Dump, args.len - 2); + for (args[2..], dumps) |p, *d| d.* = load(a, io, p); + const r = try check.checkReference(a, dumps); + try w.print("rule 1, host reference: {d} collectives, {d} receive buffers over {d} ranks", .{ r.calls, r.buffers, dumps.len }); + for (dumps) |d| try w.print(" [{s}]", .{d.header.label}); + try w.print("\n {s}\n", .{dumps[0].header.provenance}); + if (r.first) |m| { + try w.print("FAIL: {d} mismatches; first: seq {d}, layer {d}, rank {d}: {s}", .{ r.mismatches, m.seq, m.layer, m.rank, m.why }); + if (std.mem.eql(u8, m.why, "value")) try w.print(", element {d}: got 0x{x}, want 0x{x}", .{ m.element, m.got, m.want }); + if (std.mem.eql(u8, m.why, "length") or std.mem.eql(u8, m.why, "call count")) try w.print(": got {d}, want {d}", .{ m.got, m.want }); + try w.writeAll("\n"); + } else if (r.calls == 0) { + try w.writeAll("FAIL: the dumps hold no collectives\n"); + } else { + try w.writeAll("PASS: every receive buffer is bit-identical to the host reference\n"); + ok = true; + } + } else if (std.mem.eql(u8, cmd, "repeat")) { + if (args.len != 4) bad("repeat needs two dumps of the same rank", .{}); + const x = load(a, io, args[2]); + const y = load(a, io, args[3]); + if (x.header.rank != y.header.rank) bad("the dumps are ranks {d} and {d}", .{ x.header.rank, y.header.rank }); + const r = check.checkRepeat(x, y); + try w.print("rule 2, same bytes run after run: rank {d} [{s}], {d} collectives and {d} logits rows vs {d} and {d}\n", .{ x.header.rank, x.header.label, x.collectives.len, x.logits.len, y.collectives.len, y.logits.len }); + try w.print(" run A: {s}\n run B: {s}\n", .{ x.header.provenance, y.header.provenance }); + if (check.repeatPassed(r)) { + try w.print("PASS: all {d} records identical\n", .{r.records}); + ok = true; + } else { + try w.writeAll("FAIL:"); + if (r.records == 0) try w.writeAll(" the dumps hold no records;"); + if (!r.same_shape) try w.writeAll(" record counts differ;"); + if (r.first_collective) |i| try w.print(" first differing collective is record {d} (seq {d}, layer {d});", .{ i, x.collectives[i].seq, x.collectives[i].layer }); + if (r.first_logits) |i| try w.print(" first differing logits row is prompt {d} position {d};", .{ x.logits[i].prompt, x.logits[i].position }); + try w.writeAll("\n"); + } + } else if (std.mem.eql(u8, cmd, "agree")) { + if (args.len != 4) bad("agree needs the mixed run's dump, then one box's", .{}); + const mixed = load(a, io, args[2]); + const solo = load(a, io, args[3]); + var r = try check.checkAgree(a, mixed, solo); + try w.print("rule 4, mixed [{s}] ({s}) vs one box [{s}]: {d} prompts, {d} positions\n", .{ mixed.header.label, @tagName(r.mode), solo.header.label, r.prompts, r.positions }); + try w.print(" mixed: {s}\n solo: {s}\n", .{ mixed.header.provenance, solo.header.provenance }); + switch (r.mode) { + .greedy => { + try w.print("greedy tokens equal before each prompt's first split: {d}; prompts that split: {d} of {d}\n", .{ r.greedy_same, r.splits.len, r.prompts }); + for (r.splits[0..@min(r.splits.len, 64)]) |sp| try w.print(" prompt {d}: first split at position {d}\n", .{ sp.prompt, sp.position }); + if (r.splits.len > 64) try w.print(" ... {d} more\n", .{r.splits.len - 64}); + try w.writeAll("divergence below covers each prompt up to its first split (same prefix); use a teacher-forced run for the rest\n"); + }, + .teacher_forced => { + try w.print("tokens equal: {d}/{d}; fed the solo run's tokens: {s}\n", .{ r.greedy_same, r.positions, if (r.fed_ok == true) "yes (hash matches)" else "NO (fed_hash differs)" }); + }, + } + try w.print("compared {d} positions: top-1 equal {d}; |logit diff| max {e:.3} mean {e:.3}; KL(one box || mixed) mean {e:.3} max {e:.3} nats", .{ r.compared, r.top1_same, r.max_abs, r.mean_abs, r.kl_mean, r.kl_max }); + if (r.skipped > 0) try w.print("; {d} positions skipped (vocab size or non-finite logits)", .{r.skipped}); + try w.writeAll("\n"); + if (r.missing > 0) { + try w.print("{d} solo positions have no mixed record", .{r.missing}); + for (r.gaps[0..@min(r.gaps.len, 64)]) |g| try w.print("; prompt {d} stops at the gap at position {d}", .{ g.prompt, g.position }); + try w.writeAll("\n"); + } + if (check.agreeFailure(r)) |why| { + try w.print("FAIL: {s}\n", .{why}); + } else ok = true; + } else bad("unknown command {s}\n{s}", .{ cmd, usage }); + try w.flush(); + std.process.exit(if (ok) 0 else 1); +} diff --git a/zig/tests/tp_check_phase_test.zig b/zig/tests/tp_check_phase_test.zig new file mode 100644 index 0000000000..01fe86ffbb --- /dev/null +++ b/zig/tests/tp_check_phase_test.zig @@ -0,0 +1,135 @@ +//! Versioned phase records keep legacy dumps readable and reject layer or phase disagreement. +const std = @import("std"); +const c = @import("tp_check.zig"); +fn fixture(rank: u32, version: u32) ![]u8 { + var w: std.Io.Writer.Allocating = .init(std.testing.allocator); + defer w.deinit(); + try c.writeHeader(&w.writer, .{ .format_version = version, .rank = rank, .world = 2, .label = "phase host fixture" }); + var send: [4]u8 = undefined; + var recv: [4]u8 = undefined; + std.mem.writeInt(u32, &send, @bitCast(@as(f32, if (rank == 0) 1 else 2)), .little); + std.mem.writeInt(u32, &recv, @bitCast(@as(f32, 3)), .little); + for ([_]c.Phase{ .mixer, .ffn }, 0..) |phase, seq| try c.writeCollectiveVersion(&w.writer, version, .{ + .phase = if (version == 1) .legacy else phase, + .seq = seq, + .layer = 0, + .call = .all_reduce, + .dtype = .f32, + .op = .sum, + .peer = 1 - rank, + .send = &send, + .recv = &recv, + }); + try c.writeLogits(&w.writer, 0, 0, 1, &.{ 0.5, 1.0 }); + return w.toOwnedSlice(); +} +test "phase dump versions roundtrip and negative layer and phase controls fail rule1 and rule2" { + const a = std.testing.allocator; + for ([_]u32{ 1, 2 }) |version| { + const b0 = try fixture(0, version); + defer a.free(b0); + const b1 = try fixture(1, version); + defer a.free(b1); + var d0 = try c.parse(a, b0); + defer d0.deinit(a); + var d1 = try c.parse(a, b1); + defer d1.deinit(a); + try std.testing.expectEqual(version, d0.header.format_version); + try std.testing.expectEqual(if (version == 1) c.Phase.legacy else c.Phase.mixer, d0.collectives[0].phase); + try std.testing.expectEqual(@as(usize, 0), (try c.checkReference(a, &.{ d0, d1 })).mismatches); + try std.testing.expect(c.repeatPassed(c.checkRepeat(d0, d1)) == false); + d1.collectives[0].layer = 1; + try std.testing.expect((try c.checkReference(a, &.{ d0, d1 })).mismatches > 0); + d1.collectives[0].layer = 0; + if (version == 2) { + d1.collectives[0].phase = .ffn; + try std.testing.expect((try c.checkReference(a, &.{ d0, d1 })).mismatches > 0); + var clone = try c.parse(a, b0); + defer clone.deinit(a); + clone.collectives[0].phase = .ffn; + try std.testing.expect(!c.repeatPassed(c.checkRepeat(d0, clone))); + } + } +} +test "version1 rejects a phase byte and version2 rejects an unknown phase" { + const a = std.testing.allocator; + for ([_]u32{ 1, 2 }) |version| { + const bytes = try fixture(0, version); + defer a.free(bytes); + const label = "phase host fixture"; + const record = 40 + label.len; + bytes[record + 19] = if (version == 1) 1 else 255; + try std.testing.expectError(error.BadEnum, c.parse(a, bytes)); + } +} + +test "version2 phase collectives retain rule4 compatibility with legacy solo logits" { + const a = std.testing.allocator; + const bytes = try fixture(0, 2); + defer a.free(bytes); + var mixed = try c.parse(a, bytes); + defer mixed.deinit(a); + var w: std.Io.Writer.Allocating = .init(a); + defer w.deinit(); + try c.writeHeader(&w.writer, .{ .rank = 0, .world = 1, .label = "legacy solo" }); + try c.writeLogits(&w.writer, 0, 0, 1, &.{ 0.5, 1.0 }); + var solo = try c.parse(a, w.written()); + defer solo.deinit(a); + var report = try c.checkAgree(a, mixed, solo); + defer report.deinit(a); + try std.testing.expectEqual(@as(usize, 1), report.greedy_same); + try std.testing.expectEqual(@as(usize, 1), report.top1_same); + try std.testing.expectEqual(@as(usize, 0), report.splits.len); +} + +test "legacy writer rejects phase emission before writing bytes" { + var w: std.Io.Writer.Allocating = .init(std.testing.allocator); + defer w.deinit(); + const record: c.Collective = .{ .phase = .mixer, .seq = 0, .layer = 0, .call = .all_reduce, .dtype = .f32, .op = .sum, .peer = 1, .send = &.{}, .recv = &.{} }; + try std.testing.expectError(error.UnsupportedPhaseVersion, c.writeCollective(&w.writer, record)); + try std.testing.expectError(error.UnsupportedPhaseVersion, c.writeCollectiveVersion(&w.writer, 1, record)); + try std.testing.expectEqual(@as(usize, 0), w.written().len); +} + +const frozen_v1 = + "\x54\x46\x54\x50\x44\x4d\x50\x31\x01\x00\x00\x00\x00\x00\x00\x00" ++ + "\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00" ++ + "\x10\x00\x00\x00\x66\x72\x6f\x7a\x65\x6e\x20\x6c\x65\x67\x61\x63" ++ + "\x79\x20\x76\x31\x00\x00\x00\x00\x01\x00\x00\x00\x09\x00\x00\x00" ++ + "\x00\x00\x00\x00\x07\x00\x00\x00\x00\x05\x00\x00\x00\x00\x00\x00" ++ + "\x04\x00\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00" ++ + "\x00\x00\x40\x40\x00\x00\x40\x40\x02\x00\x00\x00\x00\x00\x00\x00" ++ + "\x00\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x02\x00\x00\x00" ++ + "\x00\x00\x00\x3f\x00\x00\x80\x3f"; + +test "independently frozen version1 bytes remain readable" { + const a = std.testing.allocator; + var d = try c.parse(a, frozen_v1); + defer d.deinit(a); + try std.testing.expectEqual(@as(u32, 1), d.header.format_version); + try std.testing.expectEqual(c.Phase.legacy, d.collectives[0].phase); + try std.testing.expectEqual(@as(u64, 9), d.collectives[0].seq); + try std.testing.expectEqual(@as(u32, 7), d.collectives[0].layer); + try std.testing.expectEqual(@as(usize, 0), (try c.checkReference(a, &.{d})).mismatches); + try std.testing.expectEqual(@as(f32, 1.0), d.logits[0].at(1)); +} + +test "version3 bootstrap is explicit and cannot be emitted as version2" { + const a = std.testing.allocator; + var w: std.Io.Writer.Allocating = .init(a); + defer w.deinit(); + const row: c.Collective = .{ .phase = .bootstrap, .seq = 0, .layer = 64, .call = .exchange, .dtype = .u8, .op = .sum, .peer = 1, .send = "H", .recv = "H" }; + try std.testing.expectError(error.UnsupportedPhaseVersion, c.writeCollectiveVersion(&w.writer, 2, row)); + try std.testing.expectEqual(@as(usize, 0), w.written().len); + try c.writeHeader(&w.writer, .{ .format_version = 3, .rank = 0, .world = 2, .label = "MCT3" }); + try c.writeCollectiveVersion(&w.writer, 3, row); + var d = try c.parse(a, w.written()); + defer d.deinit(a); + try std.testing.expectEqual(c.Phase.bootstrap, d.collectives[0].phase); + try std.testing.expectEqual(@as(u32, 64), d.collectives[0].layer); + var bytes = try a.dupe(u8, w.written()); + defer a.free(bytes); + // Header is 44 bytes for label MCT3 and empty provenance; collective phase follows kind/seq/layer/three enums. + bytes[44 + 4 + 8 + 4 + 3] = 4; + try std.testing.expectError(error.BadEnum, c.parse(a, bytes)); +} diff --git a/zig/tests/tp_check_reject_test.zig b/zig/tests/tp_check_reject_test.zig new file mode 100644 index 0000000000..32470a88c9 --- /dev/null +++ b/zig/tests/tp_check_reject_test.zig @@ -0,0 +1,157 @@ +//! The dump parser refuses what it cannot vouch for: foreign bytes, unknown versions, cut files and illegal phases. +const std = @import("std"); +const c = @import("tp_check.zig"); +const testing = std.testing; + +fn dumpOf(a: std.mem.Allocator, version: u32, phase: c.Phase) ![]u8 { + var w: std.Io.Writer.Allocating = .init(a); + defer w.deinit(); + try c.writeHeader(&w.writer, .{ .format_version = version, .rank = 0, .world = 2, .label = "reject fixture" }); + const row = [_]u8{ 0, 0, 128, 63 }; + try c.writeCollectiveVersion(&w.writer, version, .{ .seq = 0, .layer = 0, .phase = phase, .call = .all_reduce, .dtype = .f32, .op = .sum, .peer = 0, .send = &row, .recv = &row }); + return w.toOwnedSlice(); +} + +/// Where the phase byte of the first collective record sits: the only byte that differs between these two dumps. +fn phaseAt(legacy: []const u8, mixer: []const u8) !usize { + try testing.expectEqual(legacy.len, mixer.len); + for (legacy, mixer, 0..) |x, y, i| if (x != y) return i; + return error.NoPhaseByte; +} + +test "tp check: foreign magic is refused before anything else is read" { + const a = testing.allocator; + const bytes = try dumpOf(a, c.canonical_version, .bootstrap); + defer a.free(bytes); + bytes[0] ^= 0xff; + try testing.expectError(error.BadMagic, c.parse(a, bytes)); +} + +test "tp check: an unsupported version is refused, whatever follows it" { + const a = testing.allocator; + const bytes = try dumpOf(a, c.canonical_version, .bootstrap); + defer a.free(bytes); + for ([_]u32{ 0, 4, 5, 255, 1 << 16, std.math.maxInt(u32) }) |v| { + std.mem.writeInt(u32, bytes[c.magic.len..][0..4], v, .little); + try testing.expectError(error.BadVersion, c.parse(a, bytes)); + } +} + +test "tp check: every cut of a valid dump is refused or parses, and never reads past its end" { + const a = testing.allocator; + for ([_]u32{ c.version, c.phase_version, c.canonical_version }) |v| { + const bytes = try dumpOf(a, v, if (v == c.version) .legacy else .mixer); + defer a.free(bytes); + var refused: usize = 0; + for (0..bytes.len) |n| { + if (c.parse(a, bytes[0..n])) |d| { + var kept = d; + kept.deinit(a); + } else |err| { + try testing.expect(err == error.Truncated or err == error.BadMagic); + refused += 1; + } + } + try testing.expect(refused > 0); + } +} + +test "tp check: a phase byte its version does not allow is refused" { + const a = testing.allocator; + const legacy = try dumpOf(a, c.phase_version, .legacy); + defer a.free(legacy); + const mixer = try dumpOf(a, c.phase_version, .mixer); + defer a.free(mixer); + const at = try phaseAt(legacy, mixer); + const cases = [_]struct { version: u32, byte: u8 }{ + .{ .version = c.version, .byte = 1 }, + .{ .version = c.version, .byte = 3 }, + .{ .version = c.phase_version, .byte = 3 }, + .{ .version = c.phase_version, .byte = 4 }, + .{ .version = c.canonical_version, .byte = 4 }, + .{ .version = c.canonical_version, .byte = 255 }, + }; + for (cases) |t| { + const bytes = try a.dupe(u8, legacy); + defer a.free(bytes); + std.mem.writeInt(u32, bytes[c.magic.len..][0..4], t.version, .little); + bytes[at] = t.byte; + try testing.expectError(error.BadEnum, c.parse(a, bytes)); + } +} + +test "tp check: the writer refuses a phase or version the format does not have" { + var w: std.Io.Writer.Allocating = .init(testing.allocator); + defer w.deinit(); + const row = [_]u8{ 0, 0, 128, 63 }; + const rec: c.Collective = .{ .seq = 0, .layer = 0, .phase = .bootstrap, .call = .all_reduce, .dtype = .f32, .op = .sum, .peer = 0, .send = &row, .recv = &row }; + try testing.expectError(error.UnsupportedPhaseVersion, c.writeCollectiveVersion(&w.writer, c.version, rec)); + try testing.expectError(error.UnsupportedPhaseVersion, c.writeCollectiveVersion(&w.writer, c.phase_version, rec)); + try testing.expectError(error.UnsupportedPhaseVersion, c.writeCollectiveVersion(&w.writer, 4, rec)); + try testing.expectEqual(@as(usize, 0), w.written().len); +} + +fn headerOnly(a: std.mem.Allocator) ![]u8 { + var w: std.Io.Writer.Allocating = .init(a); + defer w.deinit(); + try c.writeHeader(&w.writer, .{ .rank = 0, .world = 2, .label = "empty" }); + return w.toOwnedSlice(); +} + +fn logitsOf(a: std.mem.Allocator, rows: []const []const f32) ![]u8 { + var w: std.Io.Writer.Allocating = .init(a); + defer w.deinit(); + try c.writeHeader(&w.writer, .{ .rank = 0, .world = 2, .label = "logits" }); + for (rows, 0..) |row, i| try c.writeLogits(&w.writer, 0, i, 0, row); + return w.toOwnedSlice(); +} + +fn agreeOn(a: std.mem.Allocator, mixed: []const []const f32, solo: []const []const f32) !?[]const u8 { + const m = try logitsOf(a, mixed); + defer a.free(m); + const s = try logitsOf(a, solo); + defer a.free(s); + var md = try c.parse(a, m); + defer md.deinit(a); + var sd = try c.parse(a, s); + defer sd.deinit(a); + var r = try c.checkAgree(a, md, sd); + defer r.deinit(a); + return c.agreeFailure(r); +} + +test "tp check: rule 2 refuses two dumps with no records" { + const a = testing.allocator; + const bytes = try headerOnly(a); + defer a.free(bytes); + var x = try c.parse(a, bytes); + defer x.deinit(a); + var y = try c.parse(a, bytes); + defer y.deinit(a); + try testing.expect(!c.repeatPassed(c.checkRepeat(x, y))); +} + +test "tp check: rule 4 fails when it compared nothing, or the mixed run's logits are not finite" { + const a = testing.allocator; + const good = [_]f32{ 1, 2 }; + const wide = [_]f32{ 1, 2, 3 }; + const nan = [_]f32{ std.math.nan(f32), 2 }; + const inf = [_]f32{ std.math.inf(f32), 2 }; + try testing.expectEqual(@as(?[]const u8, null), try agreeOn(a, &.{&good}, &.{&good})); + try testing.expectEqualStrings("no position could be compared", (try agreeOn(a, &.{&wide}, &.{&good})).?); + try testing.expectEqualStrings("the mixed run's logits are not finite", (try agreeOn(a, &.{&nan}, &.{&good})).?); + try testing.expectEqualStrings("the mixed run's logits are not finite", (try agreeOn(a, &.{ &good, &inf }, &.{ &good, &good })).?); + try testing.expectEqual(@as(?[]const u8, null), try agreeOn(a, &.{ &good, &good }, &.{ &good, &nan })); +} + +test "tp check: a buffer that is not whole elements is refused" { + const a = testing.allocator; + for ([_]c.Call{ .all_reduce, .all_gather, .exchange }) |call| { + var w: std.Io.Writer.Allocating = .init(a); + defer w.deinit(); + try c.writeHeader(&w.writer, .{ .rank = 0, .world = 2, .label = "partial" }); + const row = [_]u8{ 1, 2, 3, 4, 5, 6 }; + try c.writeCollective(&w.writer, .{ .seq = 0, .layer = 0, .call = call, .dtype = .f32, .op = .sum, .peer = 0, .send = &row, .recv = row[0..4] }); + try testing.expectError(error.BadLength, c.parse(a, w.written())); + } +} diff --git a/zig/tests/tp_check_rules.zig b/zig/tests/tp_check_rules.zig new file mode 100644 index 0000000000..cc542447fc --- /dev/null +++ b/zig/tests/tp_check_rules.zig @@ -0,0 +1,150 @@ +//! Mixed-TP rules 1 and 2: every collective equals the host reference, and two runs match byte for byte. +const std = @import("std"); +const dump = @import("tp_check_dump.zig"); +const tp = @import("tp"); +const reference = tp.reference; +const DType = tp.collective.DType; +const Error = dump.Error; +const Collective = dump.Collective; +const Dump = dump.Dump; +pub const Mismatch = struct { + seq: u64, + layer: u32, + rank: u32, + /// What went wrong; for a value mismatch, the first differing element and both values' bits. + why: []const u8, + element: usize = 0, + got: u64 = 0, + want: u64 = 0, +}; + +pub const ReferenceReport = struct { + calls: usize = 0, + /// Every recv buffer checked, one per rank per call. + buffers: usize = 0, + first: ?Mismatch = null, + mismatches: usize = 0, +}; + +fn elemBits(b: []const u8, t: DType, i: usize) u64 { + const s = t.size(); + return switch (s) { + 1 => b[i], + 2 => std.mem.readInt(u16, b[2 * i ..][0..2], .little), + 4 => std.mem.readInt(u32, b[4 * i ..][0..4], .little), + else => std.mem.readInt(u64, b[8 * i ..][0..8], .little), + }; +} + +fn note(r: *ReferenceReport, m: Mismatch) void { + r.mismatches += 1; + if (r.first == null) r.first = m; +} + +/// Compare `got` with `want` bit for bit; on a difference, record the first differing element. +fn compare(r: *ReferenceReport, c: Collective, rank: u32, got: []const u8, want: []const u8) void { + r.buffers += 1; + if (got.len != want.len) return note(r, .{ .seq = c.seq, .layer = c.layer, .rank = rank, .why = "length", .got = got.len, .want = want.len }); + if (std.mem.eql(u8, got, want)) return; + const n = want.len / c.dtype.size(); + for (0..n) |i| { + const g = elemBits(got, c.dtype, i); + const w = elemBits(want, c.dtype, i); + if (g != w) return note(r, .{ .seq = c.seq, .layer = c.layer, .rank = rank, .why = "value", .element = i, .got = g, .want = w }); + } +} + +/// Rule 1: dumps[r] is rank r's; the n-th calls must match, and each recv must equal the rank-order host reference. +pub fn checkReference(a: std.mem.Allocator, dumps: []const Dump) Error!ReferenceReport { + var r: ReferenceReport = .{}; + if (dumps.len == 0) return r; + for (dumps, 0..) |d, i| { + if (d.header.rank != i or d.header.world != dumps.len) { + note(&r, .{ .seq = 0, .layer = 0, .rank = @intCast(i), .why = "rank or world in the header" }); + return r; + } + if (d.collectives.len != dumps[0].collectives.len) { + note(&r, .{ .seq = 0, .layer = 0, .rank = @intCast(i), .why = "call count", .got = d.collectives.len, .want = dumps[0].collectives.len }); + return r; + } + } + const sends = try a.alloc([]const u8, dumps.len); + defer a.free(sends); + var want: std.ArrayList(u8) = .empty; + defer want.deinit(a); + for (0..dumps[0].collectives.len) |k| { + const c0 = dumps[0].collectives[k]; + r.calls += 1; + var same = true; + for (dumps, 0..) |d, i| { + const c = d.collectives[k]; + sends[i] = c.send; + if (c.seq != c0.seq or c.layer != c0.layer or c.phase != c0.phase or c.call != c0.call or c.dtype != c0.dtype or c.op != c0.op or c.send.len != c0.send.len) { + note(&r, .{ .seq = c0.seq, .layer = c0.layer, .rank = @intCast(i), .why = "ranks disagree on the call" }); + same = false; + break; + } + } + if (!same) continue; + switch (c0.call) { + .all_reduce => { + try want.resize(a, c0.send.len); + reference.reduce(want.items, sends, c0.dtype, c0.op); + for (dumps, 0..) |d, i| compare(&r, d.collectives[k], @intCast(i), d.collectives[k].recv, want.items); + }, + .all_gather => { + want.clearRetainingCapacity(); + for (sends) |s| try want.appendSlice(a, s); + for (dumps, 0..) |d, i| compare(&r, d.collectives[k], @intCast(i), d.collectives[k].recv, want.items); + }, + .exchange => for (dumps, 0..) |d, i| { + const c = d.collectives[k]; + if (c.peer >= dumps.len) { + note(&r, .{ .seq = c.seq, .layer = c.layer, .rank = @intCast(i), .why = "peer out of range", .got = c.peer }); + continue; + } + compare(&r, c, @intCast(i), c.recv, dumps[c.peer].collectives[k].send); + }, + } + } + return r; +} + +pub const RepeatReport = struct { + records: usize = 0, + /// Index (in record order) of the first collective or logits record that differs, if any. + first_collective: ?usize = null, + first_logits: ?usize = null, + same_shape: bool = true, +}; + +/// Rule 2: two dumps of the same rank from two runs. Every record must match byte for byte (sends, recvs, logits). +pub fn checkRepeat(x: Dump, y: Dump) RepeatReport { + var r: RepeatReport = .{}; + if (x.collectives.len != y.collectives.len or x.logits.len != y.logits.len) r.same_shape = false; + for (0..@min(x.collectives.len, y.collectives.len)) |i| { + r.records += 1; + const p = x.collectives[i]; + const q = y.collectives[i]; + const same = p.seq == q.seq and p.layer == q.layer and p.phase == q.phase and p.call == q.call and p.dtype == q.dtype and p.op == q.op and + p.peer == q.peer and std.mem.eql(u8, p.send, q.send) and std.mem.eql(u8, p.recv, q.recv); + if (!same) { + r.first_collective = i; + break; + } + } + for (0..@min(x.logits.len, y.logits.len)) |i| { + r.records += 1; + const p = x.logits[i]; + const q = y.logits[i]; + if (p.prompt != q.prompt or p.position != q.position or p.chosen != q.chosen or !std.mem.eql(u8, p.bytes, q.bytes)) { + r.first_logits = i; + break; + } + } + return r; +} + +pub fn repeatPassed(r: RepeatReport) bool { + return r.records > 0 and r.same_shape and r.first_collective == null and r.first_logits == null; +} diff --git a/zig/tests/tp_check_test.zig b/zig/tests/tp_check_test.zig new file mode 100644 index 0000000000..885db39ea5 --- /dev/null +++ b/zig/tests/tp_check_test.zig @@ -0,0 +1,283 @@ +//! Host tests for the mixed-TP gate in tp_check.zig: dump format, rule 1 reference, rule 2 repeat, rule 4 agreement. +const std = @import("std"); +const reference = @import("tp").reference; +const check = @import("tp_check.zig"); +const Call = check.Call; +const Mode = check.Mode; +const Header = check.Header; +const writeHeader = check.writeHeader; +const writeCollective = check.writeCollective; +const writeLogits = check.writeLogits; +const TokenHash = check.TokenHash; +const parse = check.parse; +const checkReference = check.checkReference; +const checkRepeat = check.checkRepeat; +const repeatPassed = check.repeatPassed; +const Split = check.Split; +const chosenHash = check.chosenHash; +const checkAgree = check.checkAgree; + +const testing = std.testing; + +fn bf16s(comptime vals: []const f32) [vals.len * 2]u8 { + var out: [vals.len * 2]u8 = undefined; + for (vals, 0..) |v, i| std.mem.writeInt(u16, out[2 * i ..][0..2], reference.f32ToBf16(v), .little); + return out; +} + +/// Two ranks' dumps of one bf16 all-reduce, one all-gather and one exchange, with correct results. +fn twoRanks(a: std.mem.Allocator, bad_recv: bool) ![2][]u8 { + const p0 = comptime bf16s(&.{ 1.0, -2.5, 0.125 }); + const p1 = comptime bf16s(&.{ 0.00390625, 3.0, 0.125 }); + var sum: [6]u8 = undefined; + reference.reduce(&sum, &.{ &p0, &p1 }, .bf16, .sum); + var out: [2][]u8 = undefined; + for (0..2) |r| { + var w: std.Io.Writer.Allocating = .init(a); + defer w.deinit(); + const mine: []const u8 = if (r == 0) &p0 else &p1; + const other: []const u8 = if (r == 0) &p1 else &p0; + try writeHeader(&w.writer, .{ .rank = @intCast(r), .world = 2, .label = if (r == 0) "metal test" else "cuda test", .provenance = "commit=abc model=test" }); + var recv = sum; + if (bad_recv and r == 1) recv[2] ^= 1; + try writeCollective(&w.writer, .{ .seq = 0, .layer = 3, .call = .all_reduce, .dtype = .bf16, .op = .sum, .peer = 0, .send = mine, .recv = &recv }); + const gathered = p0 ++ p1; + try writeCollective(&w.writer, .{ .seq = 1, .layer = 3, .call = .all_gather, .dtype = .bf16, .op = .sum, .peer = 0, .send = mine, .recv = &gathered }); + try writeCollective(&w.writer, .{ .seq = 2, .layer = 4, .call = .exchange, .dtype = .bf16, .op = .sum, .peer = @intCast(1 - r), .send = mine, .recv = other }); + try writeLogits(&w.writer, 0, 0, 2, &.{ 0.5, 1.0, 4.0 }); + out[r] = try w.toOwnedSlice(); + } + return out; +} + +test "tp check: a dump round-trips" { + const a = testing.allocator; + const bytes = try twoRanks(a, false); + defer for (bytes) |b| a.free(b); + var d = try parse(a, bytes[1]); + defer d.deinit(a); + try testing.expectEqual(@as(u32, 1), d.header.rank); + try testing.expectEqualStrings("cuda test", d.header.label); + try testing.expectEqualStrings("commit=abc model=test", d.header.provenance); + try testing.expectEqual(Mode.greedy, d.header.mode); + try testing.expectEqual(@as(usize, 3), d.collectives.len); + try testing.expectEqual(Call.exchange, d.collectives[2].call); + try testing.expectEqual(@as(u32, 0), d.collectives[2].peer); + try testing.expectEqual(@as(usize, 1), d.logits.len); + try testing.expectEqual(@as(f32, 4.0), d.logits[0].at(2)); +} + +test "tp check: a truncated or foreign dump is refused" { + const a = testing.allocator; + const bytes = try twoRanks(a, false); + defer for (bytes) |b| a.free(b); + try testing.expectError(error.Truncated, parse(a, bytes[0][0 .. bytes[0].len - 1])); + var foreign = try a.dupe(u8, bytes[0]); + defer a.free(foreign); + foreign[0] = 'X'; + try testing.expectError(error.BadMagic, parse(a, foreign)); +} + +test "tp check: correct collectives pass rule 1, one flipped bit fails it" { + const a = testing.allocator; + for ([_]bool{ false, true }) |bad| { + const bytes = try twoRanks(a, bad); + defer for (bytes) |b| a.free(b); + var d0 = try parse(a, bytes[0]); + defer d0.deinit(a); + var d1 = try parse(a, bytes[1]); + defer d1.deinit(a); + const r = try checkReference(a, &.{ d0, d1 }); + try testing.expectEqual(@as(usize, 3), r.calls); + try testing.expectEqual(@as(usize, 6), r.buffers); + if (!bad) { + try testing.expectEqual(@as(usize, 0), r.mismatches); + } else { + try testing.expectEqual(@as(usize, 1), r.mismatches); + const m = r.first.?; + try testing.expectEqual(@as(u32, 1), m.rank); + try testing.expectEqual(@as(usize, 1), m.element); + try testing.expectEqualStrings("value", m.why); + } + } +} + +test "tp check: rule 1 refuses ranks that disagree on the call" { + const a = testing.allocator; + const bytes = try twoRanks(a, false); + defer for (bytes) |b| a.free(b); + var d0 = try parse(a, bytes[0]); + defer d0.deinit(a); + var d1 = try parse(a, bytes[1]); + defer d1.deinit(a); + d1.collectives[0].dtype = .f16; + const r = try checkReference(a, &.{ d0, d1 }); + try testing.expectEqualStrings("ranks disagree on the call", r.first.?.why); +} + +test "tp check: rule 2 finds the first record that differs between runs" { + const a = testing.allocator; + const good = try twoRanks(a, false); + defer for (good) |b| a.free(b); + const bad = try twoRanks(a, true); + defer for (bad) |b| a.free(b); + var x = try parse(a, good[1]); + defer x.deinit(a); + var y = try parse(a, good[1]); + defer y.deinit(a); + try testing.expect(repeatPassed(checkRepeat(x, y))); + var z = try parse(a, bad[1]); + defer z.deinit(a); + const r = checkRepeat(x, z); + try testing.expect(!repeatPassed(r)); + try testing.expectEqual(@as(?usize, 0), r.first_collective); +} + +const Row = struct { prompt: u32, position: u64, chosen: u32, logits: []const f32 }; + +fn logitsDump(a: std.mem.Allocator, h: Header, rows: []const Row) ![]u8 { + var w: std.Io.Writer.Allocating = .init(a); + defer w.deinit(); + try writeHeader(&w.writer, h); + for (rows) |r| try writeLogits(&w.writer, r.prompt, r.position, r.chosen, r.logits); + return w.toOwnedSlice(); +} + +/// Two prompts on one box alone, in order. +const solo_rows = [_]Row{ + .{ .prompt = 0, .position = 0, .chosen = 1, .logits = &.{ 1.0, 2.0, 0.0 } }, + .{ .prompt = 0, .position = 1, .chosen = 0, .logits = &.{ 3.0, 0.0, 0.0 } }, + .{ .prompt = 0, .position = 2, .chosen = 2, .logits = &.{ 0.0, 0.0, 5.0 } }, + .{ .prompt = 1, .position = 0, .chosen = 1, .logits = &.{ 0.0, 4.0, 0.0 } }, + .{ .prompt = 1, .position = 1, .chosen = 0, .logits = &.{ 2.0, 0.0, 0.0 } }, +}; + +/// The mixed run, out of order: prompt 0 splits at 1, and its far-off position 2 must stay out of greedy divergence. +const mixed_rows = [_]Row{ + .{ .prompt = 1, .position = 1, .chosen = 0, .logits = &.{ 2.0, 0.0, 0.0 } }, + .{ .prompt = 0, .position = 2, .chosen = 2, .logits = &.{ 0.0, 0.0, 9.0 } }, + .{ .prompt = 1, .position = 0, .chosen = 1, .logits = &.{ 0.0, 4.0, 0.0 } }, + .{ .prompt = 0, .position = 0, .chosen = 1, .logits = &.{ 1.0, 2.0, 0.0 } }, + .{ .prompt = 0, .position = 1, .chosen = 1, .logits = &.{ 0.0, 3.5, 0.0 } }, +}; + +test "tp check: rule 4 greedy matches by prompt and position and stops each prompt at its first split" { + const a = testing.allocator; + const s = try logitsDump(a, .{ .rank = 0, .world = 1, .label = "solo" }, &solo_rows); + defer a.free(s); + const m = try logitsDump(a, .{ .rank = 0, .world = 2, .label = "mixed" }, &mixed_rows); + defer a.free(m); + var solo = try parse(a, s); + defer solo.deinit(a); + var mixed = try parse(a, m); + defer mixed.deinit(a); + var r = try checkAgree(a, mixed, solo); + defer r.deinit(a); + try testing.expectEqual(@as(usize, 2), r.prompts); + try testing.expectEqual(@as(usize, 5), r.positions); + try testing.expectEqual(@as(usize, 3), r.greedy_same); + try testing.expectEqual(@as(usize, 1), r.splits.len); + try testing.expectEqual(Split{ .prompt = 0, .position = 1 }, r.splits[0]); + try testing.expectEqual(@as(usize, 4), r.compared); + try testing.expectEqual(@as(usize, 3), r.top1_same); + try testing.expectEqual(@as(f64, 3.5), r.max_abs); + try testing.expect(r.kl_max > 2.0); + try testing.expectEqual(@as(?bool, null), r.fed_ok); + var same = try checkAgree(a, solo, solo); + defer same.deinit(a); + try testing.expectEqual(@as(usize, 0), same.splits.len); + try testing.expectEqual(@as(f64, 0), same.max_abs); + try testing.expectEqual(@as(f64, 0), same.kl_max); +} + +test "tp check: rule 4 greedy stops a prompt at a gap, so a hidden split cannot pair different prefixes" { + const a = testing.allocator; + const s = try logitsDump(a, .{ .rank = 0, .world = 1, .label = "solo" }, &solo_rows); + defer a.free(s); + // the mixed run lost prompt 0 position 1 (its split); position 2 is far off and must not be compared + const holed = [_]Row{ mixed_rows[0], mixed_rows[1], mixed_rows[2], mixed_rows[3] }; + const m = try logitsDump(a, .{ .rank = 0, .world = 2, .label = "mixed" }, &holed); + defer a.free(m); + var solo = try parse(a, s); + defer solo.deinit(a); + var mixed = try parse(a, m); + defer mixed.deinit(a); + var r = try checkAgree(a, mixed, solo); + defer r.deinit(a); + try testing.expectEqual(@as(usize, 1), r.missing); + try testing.expectEqual(@as(usize, 1), r.gaps.len); + try testing.expectEqual(Split{ .prompt = 0, .position = 1 }, r.gaps[0]); + try testing.expectEqual(@as(usize, 0), r.splits.len); + try testing.expectEqual(@as(usize, 4), r.positions); + try testing.expectEqual(@as(usize, 3), r.compared); + try testing.expectEqual(@as(f64, 0), r.max_abs); +} + +test "tp check: rule 4 greedy does not count positions missing after a split" { + const a = testing.allocator; + const s = try logitsDump(a, .{ .rank = 0, .world = 1, .label = "solo" }, &solo_rows); + defer a.free(s); + // prompt 0 splits at position 1, then the mixed run ends: its position 2 is absent, which is normal + const short = [_]Row{ mixed_rows[0], mixed_rows[2], mixed_rows[3], mixed_rows[4] }; + const m = try logitsDump(a, .{ .rank = 0, .world = 2, .label = "mixed" }, &short); + defer a.free(m); + var solo = try parse(a, s); + defer solo.deinit(a); + var mixed = try parse(a, m); + defer mixed.deinit(a); + var r = try checkAgree(a, mixed, solo); + defer r.deinit(a); + try testing.expectEqual(@as(usize, 0), r.missing); + try testing.expectEqual(@as(usize, 0), r.gaps.len); + try testing.expectEqual(@as(usize, 1), r.splits.len); + try testing.expectEqual(@as(usize, 4), r.compared); +} + +test "tp check: rule 4 teacher-forced compares every position, only when fed the solo run's tokens" { + const a = testing.allocator; + const s = try logitsDump(a, .{ .rank = 0, .world = 1, .label = "solo" }, &solo_rows); + defer a.free(s); + var solo = try parse(a, s); + defer solo.deinit(a); + const fed = try chosenHash(a, solo); + for ([_]u64{ fed, fed ^ 1 }) |h| { + const m = try logitsDump(a, .{ .rank = 0, .world = 2, .mode = .teacher_forced, .fed_hash = h, .label = "mixed tf" }, &mixed_rows); + defer a.free(m); + var mixed = try parse(a, m); + defer mixed.deinit(a); + var r = try checkAgree(a, mixed, solo); + defer r.deinit(a); + try testing.expectEqual(@as(usize, 5), r.positions); + try testing.expectEqual(@as(usize, 0), r.splits.len); + if (h == fed) { + try testing.expectEqual(@as(?bool, true), r.fed_ok); + try testing.expectEqual(@as(usize, 5), r.compared); + try testing.expectEqual(@as(f64, 4.0), r.max_abs); + } else { + try testing.expectEqual(@as(?bool, false), r.fed_ok); + try testing.expectEqual(@as(usize, 0), r.compared); + } + } +} + +test "tp check: the fed-token hash does not depend on record order" { + const a = testing.allocator; + const x = try logitsDump(a, .{ .rank = 0, .world = 1, .label = "x" }, &mixed_rows); + defer a.free(x); + const sorted = [_]Row{ mixed_rows[3], mixed_rows[4], mixed_rows[1], mixed_rows[2], mixed_rows[0] }; + const y = try logitsDump(a, .{ .rank = 0, .world = 1, .label = "y" }, &sorted); + defer a.free(y); + var dx = try parse(a, x); + defer dx.deinit(a); + var dy = try parse(a, y); + defer dy.deinit(a); + try testing.expectEqual(try chosenHash(a, dx), try chosenHash(a, dy)); + var t: TokenHash = .{}; + for (sorted) |r| t.add(r.prompt, r.position, r.chosen); + try testing.expectEqual(t.h, try chosenHash(a, dx)); +} + +test { + _ = @import("tp_check_phase_test.zig"); + _ = @import("tp_check_reject_test.zig"); +}