From 781595e4d2b990cc028a5938cda42fe2f847bb01 Mon Sep 17 00:00:00 2001 From: Bizuayeu <30091157+Bizuayeu@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:41:27 +0900 Subject: [PATCH] core: nemotron's checkpoint identity moves to core; GLM's ranks compare the checkpoint by it, and a hub snapshot's revision only against another --- .../checkpoint_identity.zig} | 21 +++++++++ zig/src/core/root.zig | 2 + zig/src/families/glm/cuda_engine.zig | 3 ++ zig/src/families/glm/cuda_ranks.zig | 44 +++++++++++++------ zig/src/families/nemotron/cuda.zig | 3 +- zig/src/families/nemotron/cuda_engine.zig | 2 +- zig/src/families/nemotron/cuda_latent.zig | 4 +- zig/src/families/nemotron/cuda_worker.zig | 2 +- 8 files changed, 61 insertions(+), 20 deletions(-) rename zig/src/{families/nemotron/cuda_identity.zig => core/checkpoint_identity.zig} (93%) diff --git a/zig/src/families/nemotron/cuda_identity.zig b/zig/src/core/checkpoint_identity.zig similarity index 93% rename from zig/src/families/nemotron/cuda_identity.zig rename to zig/src/core/checkpoint_identity.zig index c080aaee87..0406cb18d6 100644 --- a/zig/src/families/nemotron/cuda_identity.zig +++ b/zig/src/core/checkpoint_identity.zig @@ -379,3 +379,24 @@ test "a Hugging Face snapshot's shard is checked against its blob name: the righ _ = try contentOf(io, good, false); try std.testing.expectError(error.ShardBytesDiffer, contentOf(io, bad, false)); } + +test "the identity reads headers, not tensor bytes or the folder's name: another tensor name changes it" { + const io = std.testing.io; + const gpa = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + const cfg = "{\"hidden_size\": 8}"; + try write(io, tmp.dir, "a", cfg, 9, null, null); + try write(io, tmp.dir, "b", cfg, 10, null, null); // another dense byte under the same header, in another folder + try write(io, tmp.dir, "n", cfg, 9, null, null); + const bytes = try tmp.dir.readFileAlloc(io, "n/model.safetensors", gpa, .limited(1 << 24)); + defer gpa.free(bytes); + bytes[10] = 'n'; // the first tensor's name: "m.layers.0.d.weight" becomes "n.layers.0.d.weight" + try tmp.dir.writeFile(io, .{ .sub_path = "n/model.safetensors", .data = bytes }); + var path: [3][]u8 = undefined; + for ([_][]const u8{ "a", "b", "n" }, &path) |s, *p| p.* = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}/{s}", .{ tmp.sub_path, s }); + defer for (path) |p| gpa.free(p); + const a = try identity(gpa, io, path[0]); + try std.testing.expectEqual(a, try identity(gpa, io, path[1])); + try std.testing.expect(!std.mem.eql(u8, &a, &try identity(gpa, io, path[2]))); +} diff --git a/zig/src/core/root.zig b/zig/src/core/root.zig index 122ffe9e4c..bb492e1c9e 100644 --- a/zig/src/core/root.zig +++ b/zig/src/core/root.zig @@ -7,6 +7,7 @@ pub const eaq_container = @import("eaq_container.zig"); pub const checkpoint = @import("checkpoint.zig"); pub const direct_io = @import("direct_io.zig"); pub const Checkpoint = checkpoint.Checkpoint; +pub const checkpoint_identity = @import("checkpoint_identity.zig"); pub const draft_depth = @import("draft_depth.zig"); pub const CopyIndex = @import("copy_index.zig").CopyIndex; pub const tokenizer = @import("tokenizer"); // a module of its own, so the native server shares it @@ -47,6 +48,7 @@ test { _ = @import("eaq_container_test.zig"); _ = @import("exl3_rect_test.zig"); _ = checkpoint; + _ = checkpoint_identity; _ = direct_io; _ = draft_depth; _ = @import("copy_index.zig"); diff --git a/zig/src/families/glm/cuda_engine.zig b/zig/src/families/glm/cuda_engine.zig index ee38e00b3e..308e2729d8 100644 --- a/zig/src/families/glm/cuda_engine.zig +++ b/zig/src/families/glm/cuda_engine.zig @@ -11,6 +11,7 @@ const ranks = @import("cuda_ranks.zig"); const sampler = @import("cuda_sampler.zig"); const heat = @import("heat"); const modelopt = @import("core").modelopt; +const checkpoint_identity = @import("core").checkpoint_identity; pub const Options = struct { capacity: usize, // tokens a sequence holds @@ -107,6 +108,7 @@ pub const Engine = struct { errdefer stream.deinit(); const config_hash = try ranks.hashFile(a, io, cfg_path); const program_hash = try ranks.hashFile(a, io, "/proc/self/exe"); + const checkpoint_hash = try checkpoint_identity.identity(a, io, checkpoint); const fp = ranks.Fingerprint.of( world, rank, @@ -115,6 +117,7 @@ pub const Engine = struct { o.mtp, config_hash, program_hash, + checkpoint_hash, checkpoint, ); try ranks.agree(d, lib, comm, stream, fp, try a.alloc(ranks.Fingerprint, world)); diff --git a/zig/src/families/glm/cuda_ranks.zig b/zig/src/families/glm/cuda_ranks.zig index ab5b01b770..b8892c1b9f 100644 --- a/zig/src/families/glm/cuda_ranks.zig +++ b/zig/src/families/glm/cuda_ranks.zig @@ -60,7 +60,7 @@ pub fn environment(name: [:0]const u8) ?[]const u8 { /// What every rank must agree on, as bytes the communicator carries. pub const Fingerprint = extern struct { - layout: u32 = 5, // this struct's version + layout: u32 = 6, // this struct's version world: u32, rank: u32, split: u32, @@ -73,7 +73,8 @@ pub const Fingerprint = extern struct { heat_low: f64, config: [32]u8, // sha256 of config.json program: [32]u8, // sha256 of this executable - checkpoint: [64]u8, // the checkpoint folder's last name (the snapshot's revision), zero padded + checkpoint: [32]u8, // core.checkpoint_identity (config, index, shard names, sizes, headers), not the folder's name + revision: [64]u8, // a hub snapshot's revision (the folder under snapshots/), zero padded; zero for any other folder pub const Field = enum { layout, @@ -90,6 +91,7 @@ pub const Fingerprint = extern struct { config, program, checkpoint, + revision, }; pub fn of( @@ -100,6 +102,7 @@ pub const Fingerprint = extern struct { mtp: bool, config: [32]u8, program: [32]u8, + checkpoint: [32]u8, checkpoint_dir: []const u8, ) Fingerprint { var f: Fingerprint = .{ @@ -115,11 +118,15 @@ pub const Fingerprint = extern struct { .heat_low = if (s.heat) |b| b.low_c else 0, .config = config, .program = program, - .checkpoint = @splat(0), + .checkpoint = checkpoint, + .revision = @splat(0), }; - const name = std.fs.path.basename(std.mem.trimEnd(u8, checkpoint_dir, "/")); - const n = @min(name.len, f.checkpoint.len); - @memcpy(f.checkpoint[0..n], name[0..n]); + const dir = std.mem.trimEnd(u8, checkpoint_dir, "/"); + if (std.mem.eql(u8, std.fs.path.basename(std.fs.path.dirname(dir) orelse ""), "snapshots")) { + const name = std.fs.path.basename(dir); + const n = @min(name.len, f.revision.len); + @memcpy(f.revision[0..n], name[0..n]); + } return f; } }; @@ -149,6 +156,9 @@ pub fn mismatch(all: []const Fingerprint) ?Mismatch { if (!std.mem.eql(u8, std.mem.asBytes(&@field(f, name)), std.mem.asBytes(&@field(all[0], name)))) return .{ .field = @field(Fingerprint.Field, name), .rank = r }; } + // a revision counts only where both ranks read a hub snapshot: a copy elsewhere has the same contents or not + if (f.revision[0] != 0 and all[0].revision[0] != 0 and !std.mem.eql(u8, &f.revision, &all[0].revision)) + return .{ .field = .revision, .rank = r }; } return null; } @@ -218,7 +228,8 @@ fn explain(f: Fingerprint.Field) []const u8 { .heat_low => "TF_HEAT_LOW", .config => "config.json", .program => "the executable", - .checkpoint => "the checkpoint's snapshot folder", + .checkpoint => "the checkpoint's config, index or its shards' names, sizes and headers", + .revision => "the hub snapshot's revision", }; } @@ -247,21 +258,26 @@ test "fingerprints: agreeing ranks pass; the first differing field and rank are const h0: [32]u8 = @splat(1); const h1: [32]u8 = @splat(2); var all = [_]Fingerprint{ - Fingerprint.of(2, 0, .{}, 8192, true, h0, h0, "/hub/snapshots/423acf37/"), - Fingerprint.of(2, 1, .{}, 8192, true, h0, h0, "/hub/snapshots/423acf37"), + Fingerprint.of(2, 0, .{}, 8192, true, h0, h0, h0, "/hub/snapshots/423acf37/"), + Fingerprint.of(2, 1, .{}, 8192, true, h0, h0, h0, "/hub/snapshots/423acf37"), }; try std.testing.expect(mismatch(&all) == null); all[1].program = h1; try std.testing.expectEqual(Fingerprint.Field.program, mismatch(&all).?.field); - all[1] = Fingerprint.of(2, 1, .{ .split = false }, 8192, true, h0, h0, "/hub/snapshots/423acf37"); + all[1] = Fingerprint.of(2, 1, .{ .split = false }, 8192, true, h0, h0, h0, "/hub/snapshots/423acf37"); try std.testing.expectEqual(Mismatch{ .field = .split, .rank = 1 }, mismatch(&all).?); - all[1] = Fingerprint.of(2, 0, .{}, 8192, true, h0, h0, "/hub/snapshots/423acf37"); + all[1] = Fingerprint.of(2, 0, .{}, 8192, true, h0, h0, h0, "/hub/snapshots/423acf37"); try std.testing.expectEqual(Fingerprint.Field.rank, mismatch(&all).?.field); const heat_on: Settings = .{ .heat = .{ .high_c = 92, .low_c = 88 } }; - all[1] = Fingerprint.of(2, 1, heat_on, 8192, true, h0, h0, "/hub/snapshots/423acf37"); + all[1] = Fingerprint.of(2, 1, heat_on, 8192, true, h0, h0, h0, "/hub/snapshots/423acf37"); try std.testing.expectEqual(Fingerprint.Field.heat_high, mismatch(&all).?.field); - all[1] = Fingerprint.of(2, 1, .{}, 8192, true, h0, h0, "/hub/snapshots/other"); + all[1] = Fingerprint.of(2, 1, .{}, 8192, true, h0, h0, h1, "/hub/snapshots/423acf37"); try std.testing.expectEqual(Fingerprint.Field.checkpoint, mismatch(&all).?.field); - all[1] = Fingerprint.of(2, 1, .{}, 8192, false, h0, h0, "/hub/snapshots/423acf37"); // --no-drafts on rank 1 only + // the same contents under another folder's name agree; two hub snapshots must also be the same revision + all[1] = Fingerprint.of(2, 1, .{}, 8192, true, h0, h0, h0, "/data/glm-copy"); + try std.testing.expect(mismatch(&all) == null); + all[1] = Fingerprint.of(2, 1, .{}, 8192, true, h0, h0, h0, "/hub/snapshots/other"); + try std.testing.expectEqual(Mismatch{ .field = .revision, .rank = 1 }, mismatch(&all).?); + all[1] = Fingerprint.of(2, 1, .{}, 8192, false, h0, h0, h0, "/hub/snapshots/423acf37"); // --no-drafts on rank 1 only try std.testing.expectEqual(Mismatch{ .field = .mtp, .rank = 1 }, mismatch(&all).?); } diff --git a/zig/src/families/nemotron/cuda.zig b/zig/src/families/nemotron/cuda.zig index 12ff32c8c0..b62c1fc850 100644 --- a/zig/src/families/nemotron/cuda.zig +++ b/zig/src/families/nemotron/cuda.zig @@ -28,7 +28,7 @@ pub const reuse = @import("cuda_reuse.zig"); pub const slide = @import("cuda_slide.zig"); pub const worker = @import("cuda_worker.zig"); pub const omni = @import("cuda_native_omni.zig"); -pub const identity = @import("cuda_identity.zig"); +pub const identity = @import("core").checkpoint_identity; test { _ = @import("config.zig"); @@ -49,6 +49,5 @@ test { _ = @import("cuda_train_back.zig"); _ = @import("cuda_latent.zig"); _ = @import("cuda_latent_ck.zig"); - _ = @import("cuda_identity.zig"); _ = worker; } diff --git a/zig/src/families/nemotron/cuda_engine.zig b/zig/src/families/nemotron/cuda_engine.zig index 99e6b17816..327e2bc45d 100644 --- a/zig/src/families/nemotron/cuda_engine.zig +++ b/zig/src/families/nemotron/cuda_engine.zig @@ -22,7 +22,7 @@ const prompt_rows = @import("cuda_prompt_rows.zig"); const prompt_pass = @import("cuda_prompt.zig"); const marks = @import("cuda_marks.zig"); const latent = @import("cuda_latent.zig"); -const identity = @import("cuda_identity.zig"); +const identity = @import("core").checkpoint_identity; pub const PromptRows = prompt_rows.PromptRows; /// nemotron_h.cuda.CONTEXT: prompt plus reply tokens when --context is not given, as `tensorfold serve` sizes it. diff --git a/zig/src/families/nemotron/cuda_latent.zig b/zig/src/families/nemotron/cuda_latent.zig index 634b1555ef..4c764e0949 100644 --- a/zig/src/families/nemotron/cuda_latent.zig +++ b/zig/src/families/nemotron/cuda_latent.zig @@ -105,7 +105,7 @@ const Fingerprint = extern struct { dims: [12]u64, // the config's sizes, rank and world left out kinds: [32]u8, // sha256 of the block pattern program: [32]u8, // sha256 of this executable - weights: [32]u8, // the checkpoint's identity (cuda_identity.zig): config, index, shard names, sizes, headers + weights: [32]u8, // the checkpoint's identity (core.checkpoint_identity): config, index, shard names, sizes, headers fn of(c: Config, program: [32]u8, weights: [32]u8) Fingerprint { var f: Fingerprint = .{ .rank = @intCast(c.rank), .world = @intCast(c.world), .ck = @intFromBool(c.ck), .dims = .{ c.hidden, c.vocab, c.layers, c.experts, c.top_k, c.latent, c.expert_width, c.shared_width, c.shares, c.mamba_heads, c.heads, c.kv_heads }, .kinds = undefined, .program = program, .weights = weights }; @@ -122,7 +122,7 @@ pub fn agree(gpa: std.mem.Allocator, io: std.Io, c: Config, g: *const cuda.tp_gr const exe = try std.Io.Dir.cwd().readFileAlloc(io, "/proc/self/exe", gpa, .limited(1 << 31)); defer gpa.free(exe); std.crypto.hash.sha2.Sha256.hash(exe, &program, .{}); - const weights = try @import("cuda_identity.zig").identity(gpa, io, dir); // every byte is compared at ready + const weights = try core.checkpoint_identity.identity(gpa, io, dir); // every byte is compared at ready var all: [16]Fingerprint = undefined; if (c.world > all.len) return error.TooManyRanks; try g.gather(Fingerprint, d, s, Fingerprint.of(c, program, weights), all[0..c.world]); diff --git a/zig/src/families/nemotron/cuda_worker.zig b/zig/src/families/nemotron/cuda_worker.zig index 9362cfa947..c2dc7853d4 100644 --- a/zig/src/families/nemotron/cuda_worker.zig +++ b/zig/src/families/nemotron/cuda_worker.zig @@ -8,7 +8,7 @@ const Kind = @import("config.zig").Kind; const kern = @import("cuda_kernels.zig"); const W = @import("cuda_weights.zig"); const latent = @import("cuda_latent.zig"); -const identity = @import("cuda_identity.zig"); +const identity = core.checkpoint_identity; const fwd = @import("cuda_forward.zig"); const state = @import("cuda_state.zig"); const Source = @import("cuda_source.zig").Source;