glm: a split start checks that every pair of ranks can exchange, before the weights load

zig/src/families/glm/cuda_engine.zig

@@ -118,6 +118,7 @@ pub const Engine = struct {
             checkpoint,
         );
         try ranks.agree(d, lib, comm, stream, fp, try a.alloc(ranks.Fingerprint, world));
+        if (o.settings.split) try ranks.probeSplit(d, lib, comm, stream, rank, world);
         const share = try shard.Share.of(&c, rank, world);
         const spans = try gpa.alloc(shard.Span, world);
         errdefer gpa.free(spans);

zig/src/families/glm/cuda_ranks.zig

@@ -203,6 +203,24 @@ pub fn agree(
     }
 }
 
+/// Before loading, split's peers: one value each way with each, in rank order (no pair waits on a later one).
+pub fn probeSplit(d: *const cuda.Driver, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, stream: cuda.Stream, rank: usize, world: usize) !void {
+    var dev = try cuda.DeviceBuffer.alloc(d, 8);
+    defer dev.free();
+    for (0..world) |peer| if (peer != rank) {
+        try lib.check(lib.api.ncclGroupStart(), "ncclGroupStart");
+        try lib.check(lib.api.ncclSend(dev.ptr, 1, .f32, @intCast(peer), comm, stream.handle), "ncclSend");
+        try lib.check(lib.api.ncclRecv(dev.ptr + 4, 1, .f32, @intCast(peer), comm, stream.handle), "ncclRecv");
+        lib.check(lib.api.ncclGroupEnd(), "ncclGroupEnd") catch return unreachablePeer(rank, peer);
+        stream.synchronize() catch return unreachablePeer(rank, peer);
+    };
+}
+
+fn unreachablePeer(rank: usize, peer: usize) error{SplitUnreachable} {
+    std.debug.print("REFUSE TF_GLM_PREFILL_REDUCE=split: rank {d} cannot exchange with rank {d}; split sends every rank's rows to every other, and a fabric where some ranks reach only their neighbours (a ring of four without a switch) needs TF_GLM_PREFILL_REDUCE=gather\n", .{ rank, peer });
+    return error.SplitUnreachable;
+}
+
 fn explain(f: Fingerprint.Field) []const u8 {
     return switch (f) {
         .layout => "a different build of this check",