core: nemotron's checkpoint identity moves to core, and GLM's ranks compare the checkpoint by it

zig/src/families/nemotron/cuda_identity.zig → 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])));
+}

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

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

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).?);
 }

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

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.

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]);

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;