GLM CUDA: the waiting ranks sleep while rank 0 has no stream (#1033)

zig/src/cuda/tp_watch.zig

@@ -1,12 +1,13 @@
-//! Ranks' liveness lines: a TCP line from each rank to rank 0, open before NCCL, carrying ready and error records.
+//! Ranks' liveness lines: a TCP line from each rank to rank 0, open before NCCL, carrying ready, error, wake records.
 const std = @import("std");
 const platform = @import("sys");
 const posix = @import("sys").posix;
 
-/// One record on a line: ready, or an error and its name.
+/// One record on a line: ready, an error and its name, or rank 0's wake.
 const Record = extern struct { kind: u32, rank: u32, text: [56]u8 = @splat(0) };
 const ready_kind: u32 = 1;
 const error_kind: u32 = 2;
+const wake_kind: u32 = 3;
 const max_lines = 16;
 
 /// Set by a rank that ends on purpose, before it tells the others: a line closing after that is expected.
@@ -17,6 +18,9 @@ var line_count: usize = 0;
 var my_rank: usize = 0;
 var peer_digest: [max_lines][32]u8 = undefined; // each line's ready record: its rank's checkpoint digest
 var peer_rank: [max_lines]u32 = undefined;
+var wakes = std.atomic.Value(u32).init(0); // wake records this rank's line carried (another rank's)
+var slept: u32 = 0; // the wakes this rank's sleeps took
+var watch_io: std.Io = undefined; // the watchers' futex wakes
 
 /// This rank ends on purpose (rank 0 sends its end header, or a rank leaves on one).
 pub fn ended() void {
@@ -52,6 +56,7 @@ pub fn start(io: std.Io, rank: usize, world: usize) !void {
     if (world < 2) return;
     if (world > max_lines + 1) return error.TooManyRanks;
     my_rank = rank;
+    watch_io = io;
     var addr = try address();
     if (rank == 0) {
         const server = try socket();
@@ -126,6 +131,27 @@ pub fn fail(name: []const u8) void {
     send(r);
 }
 
+/// Whether this rank has lines (more than one rank).
+pub fn open() bool {
+    return line_count > 0;
+}
+
+/// Rank 0: every other rank's sleep returns (its next step follows on NCCL).
+pub fn wake() void {
+    send(.{ .kind = wake_kind, .rank = @intCast(my_rank) });
+}
+
+/// Until rank 0's next wake record, which may already be here: a wait that spins neither CUDA nor NCCL.
+pub fn sleep(io: std.Io) void {
+    if (line_count == 0) return;
+    slept += 1;
+    while (true) {
+        const n = wakes.load(.acquire);
+        if (n >= slept) return;
+        io.futexWaitUncancelable(u32, &wakes.raw, n);
+    }
+}
+
 /// A rank's transport failed: with lines open every rank hears `name` and this one ends now; with none, nothing.
 pub fn fatal(name: []const u8) void {
     if (line_count == 0) return;
@@ -161,6 +187,11 @@ fn watch(fd: posix.socket_t, line: usize) void {
             _ = ready_peers.fetchAdd(1, .acq_rel);
             continue;
         }
+        if (r.kind == wake_kind) {
+            _ = wakes.fetchAdd(1, .release);
+            watch_io.futexWake(u32, &wakes.raw, 1);
+            continue;
+        }
         if (r.kind == error_kind and !ending.load(.acquire)) {
             const text = std.mem.sliceTo(&r.text, 0);
             std.log.err("rank {d}: rank {d} stopped ({s}); this rank stops too", .{ my_rank, r.rank, text });
@@ -180,6 +211,38 @@ test "a record is 64 bytes, and a name past 56 bytes is cut" {
     try std.testing.expectEqual(@as(usize, 56), std.mem.sliceTo(&r.text, 0).len);
 }
 
+test "a sleep takes a wake record that came before it, and waits for one that comes after" {
+    if (@import("builtin").os.tag == .windows) return error.SkipZigTest;
+    var pair: [2]c_int = undefined;
+    if (ProcessTest.socketpair(posix.AF.UNIX, posix.SOCK.STREAM, 0, &pair) != 0) return error.SocketFailed;
+    defer _ = posix.system.close(pair[0]);
+    watch_io = std.testing.io;
+    line_count = 1;
+    defer line_count = 0;
+    const first = wakes.load(.acquire);
+    slept = first;
+    const watcher = try std.Thread.spawn(.{}, watch, .{ pair[0], 0 });
+    const rec: Record = .{ .kind = wake_kind, .rank = 0 };
+    const record = std.mem.asBytes(&rec);
+    try std.testing.expectEqual(@as(isize, record.len), posix.system.write(pair[1], record.ptr, record.len));
+    while (wakes.load(.acquire) == first) std.Io.sleep(std.testing.io, .fromMilliseconds(1), .awake) catch {};
+    sleep(std.testing.io); // the wake was counted before the sleep
+    const Late = struct {
+        fn send(fd: c_int, bytes: []const u8) void {
+            std.Io.sleep(std.testing.io, .fromMilliseconds(50), .awake) catch {};
+            _ = posix.system.write(fd, bytes.ptr, bytes.len);
+        }
+    };
+    const late = try std.Thread.spawn(.{}, Late.send, .{ pair[1], record[0..] });
+    sleep(std.testing.io); // returns once the late record lands
+    late.join();
+    try std.testing.expectEqual(first + 2, wakes.load(.acquire));
+    ending.store(true, .release); // the line closing next is expected: the watcher returns
+    _ = posix.system.close(pair[1]);
+    watcher.join();
+    ending.store(false, .release);
+}
+
 const ProcessTest = struct {
     var last_pid: c_int = 0;
     var last_fd: c_int = -1;

zig/src/families/glm/cuda_follow.zig

@@ -10,7 +10,7 @@ const prompt_wire = @import("cuda_prompt_wire.zig");
 const dflash = @import("dflash.zig");
 const Engine = engine.Engine;
 
-pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image, block, many, redraw };
+pub const Kind = enum(u32) { begin = 1, prefill, verify, keep, draft, end, drop, release, chunk, shared, image, block, many, redraw, idle };
 
 /// One step as it crosses: 64 bytes, then `n` u32 ids.
 pub const Header = extern struct {
@@ -45,6 +45,7 @@ pub const Link = struct {
     ids: []u32, // host copy of the last step's ids
     gpa: std.mem.Allocator,
     wait: bool, // TF_GLM_LINK_WAIT=1: every step uploaded and waited for, as before the put
+    asleep: bool = false, // rank 0: the other ranks sleep on their lines (cuda.tp_watch) until its next step
 
     pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, lib: *const cuda.nccl.Library, comm: cuda.nccl.Comm, stream: cuda.Stream, asked_ids: usize) !Link {
         const max_ids = @max(asked_ids, prompt_wire.words);
@@ -126,6 +127,10 @@ pub const Leader = struct {
     fn send(x: Leader, h: Header, ids: []const u32) !void {
         var hh = h;
         hh.slot = x.e.slot;
+        if (x.link.asleep) { // the other ranks wait on their lines, not in NCCL: wake them for this step
+            cuda.tp_watch.wake();
+            x.link.asleep = false;
+        }
         try x.link.send(x.e.m.glue.put_words, hh, ids);
     }
 
@@ -134,6 +139,13 @@ pub const Leader = struct {
         return @intFromBool(x.e.prompt_logprobs != null);
     }
 
+    /// No stream is left: the other ranks sleep on their lines until the next step instead of spinning in NCCL.
+    pub fn idle(x: Leader) !void {
+        if (x.link.asleep or !cuda.tp_watch.open()) return;
+        try x.send(.{ .kind = @backingInt(Kind.idle) }, &.{});
+        x.link.asleep = true;
+    }
+
     /// The stream's sequence for the steps that follow (each step names it, so the other ranks bind it too).
     pub fn bind(x: Leader, slot: u32) !void {
         try x.e.bind(slot);
@@ -281,11 +293,12 @@ fn structuredStep(h: Header) bool {
 pub fn follow(e: *Engine, link: *Link) !void {
     var drafts: [engine.max_drafts]u32 = undefined;
     var sampled: [16]u32 = undefined;
+    cuda.tp_watch.sleep(e.io); // rank 0 starts with no stream: its first step wakes this rank
     while (true) {
         const step = try link.receive();
         const h = step.h;
         const kind: Kind = @fromBackingInt(@intCast(h.kind));
-        if (kind != .end and kind != .drop and kind != .shared and kind != .block and kind != .many) try e.bind(h.slot);
+        if (kind != .end and kind != .drop and kind != .shared and kind != .block and kind != .many and kind != .idle) try e.bind(h.slot);
         e.prompt_logprobs = if (promptRows(h)) .{} else null;
         e.structured = structuredStep(h); // rank 0 alone masks them
         switch (kind) {
@@ -352,6 +365,7 @@ pub fn follow(e: *Engine, link: *Link) !void {
             },
             .drop => for (step.ids) |id| e.dropState(id),
             .release => e.release(h.slot),
+            .idle => cuda.tp_watch.sleep(e.io),
             .end => return cuda.tp_watch.ended(), // rank 0 ends on purpose: its line closing from here is expected
         }
     }

zig/src/families/glm/cuda_lanes.zig

@@ -484,6 +484,7 @@ pub const Lanes = struct {
         const self = of(ptr);
         const kv = self.lanes.fetchRemove(s) orelse return;
         self.x.release(kv.value.slot) catch |err| std.log.err("releasing slot {d}: {s}", .{ kv.value.slot, @errorName(err) });
+        if (self.lanes.count() == 0) self.x.idle() catch |err| std.log.err("the other ranks' idle step: {s}", .{@errorName(err)});
     }
 };
 

zig/src/families/glm/cuda_native.zig

@@ -213,6 +213,7 @@ fn openRank(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: [
             return if (std.c.getenv(name)) |v| std.mem.eql(u8, std.mem.span(v), "0") else false;
         }
     }.get;
+    own.link.asleep = cuda.tp_watch.open(); // the other ranks begin asleep (follow): the first step wakes them
     own.lanes = .{ .x = .{ .e = own.e, .link = &own.link }, .whole = off("TF_GLM_PROMPT_STEPS"), .alone = off("TF_GLM_SHARED") };
     return .{
         .backend = own.lanes.backend(),