S26 E3: faster multi-island weight sync (pipelined sweeps, WAN stream robustness)

docs/PROTOCOL.md

@@ -32,7 +32,7 @@ or allocation overflows are errors rather than panics.
 
 | type | name | direction | payload |
 |---:|---|---|---|
-| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) |
+| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4, 4=f32 with bf16 pushes, 5=f32 with int8 pushes; flag 0x80 = bf16 downlink deltas), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) |
 | 2 | INIT_PARAMS | learner → syncer | fragment_id:u32, full tensor bytes; only learner 0 may initialize |
 | 3 | PULL_REQ | syncer → learner | fragment_id:u32, global_step:u64, round_attempt:u32 |
 | 4 | PUSH_FRAGMENT | learner → syncer | learner_id:u32, fragment_id:u32, global_step:u64, round_attempt:u32, base_version:u64, local_step:u64, c_steps:u32, c_tokens:u64, base-relative learner-delta bytes |
@@ -223,6 +223,37 @@ The coordinator is one sequential scheduler actor. Pipelined rounds target
 distinct fragments, and merge + version update + broadcast for a completion
 is serialized. No two completions mutate one fragment concurrently.
 
+## bf16 pushes over an f32 session (dtype 4)
+
+Dtype 4 keeps INIT, BCAST and FINAL exact f32 and sends only PUSH deltas as
+bf16 (round-to-nearest-even). The syncer decodes and negates them like dtype 2
+and keeps f32 masters and momentum. The full-parameter policy sweep accepts
+dtype 1 or 4 (S26 E3, `--rl-sync-push-dtype bf16`).
+
+## int8 pushes over an f32 session (dtype 5)
+
+Dtype 5 sends PUSH deltas as block-scaled int8: u32 block size B
+(32..1048576), then per block of B values (the last block may be shorter) an
+f32 scale (block absmax / 127) and one i8 per value (`-128` is invalid).
+A value decodes to `q * scale`, so its error is at most absmax / 254. The
+learner picks B (launcher `--rl-sync-push-int8-block`, default 128); it is
+bound into the role-stream session contract. INIT, BCAST and FINAL stay f32.
+
+## bf16 downlink deltas (dtype flag 0x80)
+
+On an f32-bulk session (dtype 1, 4 or 5) the flag 0x80 changes BCAST_FRAGMENT
+to `fragment_id:u32, version:u64, kind:u8, base_version:u64,
+anchor_sha256:[u8;32]`, then the tensor: kind 0 (FULL) carries f32 values,
+kind 1 (DELTA) carries `bf16(theta - anchor)`. Syncer and learner hold the
+same f32 anchor per fragment: FULL sets it, DELTA adds the decoded delta in
+f32; the digest is checked after every update. Every N-th version of a
+fragment is FULL (`--bcast-delta-full-every N`, N >= 2; the syncer refuses a
+flagged session without it and the reverse). Initial and reconnect re-sends
+are always FULL with the shared anchor. A learner with a missing or wrong
+anchor drops its connection, and the reconnect re-send restores it. The
+syncer's masters, momentum, checkpoints and FINAL are unchanged and exact;
+anchors are not checkpointed, so a resumed syncer starts with FULL.
+
 ## Q4 pushes
 
 Q4 uses the same learner-delta semantics. INIT and BCAST stay bf16; PUSH is

syncer/src/downlink.rs

@@ -0,0 +1,302 @@
+//! bf16 downlink deltas for BCAST_FRAGMENT (S26 E3, session flag
+//! `DTYPE_FLAG_BCAST_DELTA`).
+//!
+//! Every learner of a session holds, per fragment, an f32 *anchor* that is
+//! bit-identical to the syncer's copy here. A broadcast of a new version is
+//! either
+//!
+//! * `FULL`: the f32 values; the anchor becomes exactly those values, or
+//! * `DELTA`: `d = bf16(θ − anchor)`; both sides set `anchor += d` in f32
+//!   (IEEE addition is exact-rounded, so both sides get the same bits).
+//!
+//! Each delta is taken against the current anchor, so the rounding error does
+//! not accumulate: `|θ − anchor| <= |θ − anchor_prev| * 2^-9` per element
+//! after every delta. The learner trains from the anchor; its next push is
+//! relative to the anchor, and the syncer adds that push to its exact θ, so
+//! θ stays the exact sum of all pushes.
+//!
+//! Drift guard: every `full_every`-th broadcast of a fragment is FULL; a
+//! re-sent fragment (HELLO/reconnect recovery) is always FULL; and every
+//! header carries the SHA-256 of the f32 anchor after the update. A learner
+//! whose anchor version or digest does not match drops its connection; the
+//! reconnect recovery then sends FULL.
+//!
+//! Anchors live only in memory: a resumed syncer starts with none, so its
+//! first broadcast of every fragment is FULL.
+//!
+//! Header after the usual `fragment_id u32, version u64` prefix:
+//! `kind u8 (0 FULL, 1 DELTA), base_version u64, anchor_sha256 [32]`.
+//! The tensor that follows is f32 (FULL) or bf16 (DELTA).
+
+use std::sync::Arc;
+
+use anyhow::{bail, Result};
+use sha2::{Digest, Sha256};
+
+use crate::protocol::{DTYPE_BF16, DTYPE_F32};
+
+pub const KIND_FULL: u8 = 0;
+pub const KIND_DELTA: u8 = 1;
+/// fragment_id + version + kind + base_version + digest.
+pub const PREFIX_LEN: usize = 4 + 8 + 1 + 8 + 32;
+
+struct Anchor {
+    version: u64,
+    values: Arc<Vec<f32>>,
+    digest: [u8; 32],
+    deltas_since_full: u64,
+}
+
+/// One encoded BCAST: header prefix plus tensor values in `dtype` (DELTA
+/// values are already bf16-representable, so encoding them is exact).
+pub struct EncodedBcast {
+    pub prefix: Vec<u8>,
+    pub dtype: u8,
+    pub values: Arc<Vec<f32>>,
+}
+
+impl EncodedBcast {
+    #[cfg(test)]
+    pub fn kind(&self) -> u8 {
+        self.prefix[12]
+    }
+}
+
+pub struct DownlinkAnchors {
+    full_every: u64,
+    anchors: Vec<Option<Anchor>>,
+}
+
+pub fn digest_f32(values: &[f32]) -> [u8; 32] {
+    let mut hasher = Sha256::new();
+    // Hash in bounded slices to avoid a second full-size byte buffer.
+    let mut buf = Vec::with_capacity(4 * 65536);
+    for chunk in values.chunks(65536) {
+        buf.clear();
+        for v in chunk {
+            buf.extend_from_slice(&v.to_le_bytes());
+        }
+        hasher.update(&buf);
+    }
+    hasher.finalize().into()
+}
+
+fn prefix(fragment: usize, version: u64, kind: u8, base_version: u64, digest: &[u8; 32]) -> Vec<u8> {
+    let mut out = Vec::with_capacity(PREFIX_LEN);
+    out.extend_from_slice(&(fragment as u32).to_le_bytes());
+    out.extend_from_slice(&version.to_le_bytes());
+    out.push(kind);
+    out.extend_from_slice(&base_version.to_le_bytes());
+    out.extend_from_slice(digest);
+    out
+}
+
+impl DownlinkAnchors {
+    pub fn new(fragments: usize, full_every: u64) -> Result<Self> {
+        if full_every < 2 {
+            bail!("--bcast-delta-full-every must be at least 2");
+        }
+        Ok(Self {
+            full_every,
+            anchors: (0..fragments).map(|_| None).collect(),
+        })
+    }
+
+    /// Encode the broadcast of a newly committed `version` of fragment `p`
+    /// and advance the shared anchor. Call exactly once per new version.
+    pub fn publish(&mut self, p: usize, version: u64, params: &[f32]) -> Result<EncodedBcast> {
+        let slot = self
+            .anchors
+            .get_mut(p)
+            .ok_or_else(|| anyhow::anyhow!("downlink anchor for unknown fragment {p}"))?;
+        let delta_ok = matches!(
+            slot,
+            Some(a) if a.deltas_since_full + 1 < self.full_every
+                && a.values.len() == params.len()
+                && a.version < version
+        );
+        if !delta_ok {
+            return Ok(Self::reset(slot, p, version, params));
+        }
+        let anchor = slot.as_mut().unwrap();
+        let base_version = anchor.version;
+        let values = Arc::make_mut(&mut anchor.values);
+        let mut delta = Vec::new();
+        delta
+            .try_reserve_exact(params.len())
+            .map_err(|error| anyhow::anyhow!("cannot allocate downlink delta: {error}"))?;
+        for (a, t) in values.iter_mut().zip(params) {
+            let d = half::bf16::from_f32(t - *a).to_f32();
+            if !d.is_finite() {
+                // bf16 overflow of a huge step: fall back to FULL.
+                return Ok(Self::reset(slot, p, version, params));
+            }
+            delta.push(d);
+            *a += d;
+        }
+        anchor.version = version;
+        anchor.digest = digest_f32(values);
+        anchor.deltas_since_full += 1;
+        Ok(EncodedBcast {
+            prefix: prefix(p, version, KIND_DELTA, base_version, &anchor.digest),
+            dtype: DTYPE_BF16,
+            values: Arc::new(delta),
+        })
+    }
+
+    /// Encode a newly committed `version` of fragment `p` as FULL (exact f32)
+    /// and reset the shared anchor to it.
+    pub fn publish_full(&mut self, p: usize, version: u64, params: &[f32]) -> EncodedBcast {
+        match self.anchors.get_mut(p) {
+            Some(slot) => Self::reset(slot, p, version, params),
+            None => EncodedBcast {
+                prefix: prefix(p, version, KIND_FULL, version, &digest_f32(params)),
+                dtype: DTYPE_F32,
+                values: Arc::new(params.to_vec()),
+            },
+        }
+    }
+
+    /// FULL re-send for recovery (HELLO / reconnect): the anchor every other
+    /// learner holds, when it is at `version`; otherwise the exact params.
+    pub fn resend(&self, p: usize, version: u64, params: &[f32]) -> EncodedBcast {
+        match self.anchors.get(p).and_then(Option::as_ref) {
+            Some(anchor) if anchor.version == version => EncodedBcast {
+                prefix: prefix(p, version, KIND_FULL, version, &anchor.digest),
+                dtype: DTYPE_F32,
+                values: Arc::clone(&anchor.values),
+            },
+            _ => EncodedBcast {
+                prefix: prefix(p, version, KIND_FULL, version, &digest_f32(params)),
+                dtype: DTYPE_F32,
+                values: Arc::new(params.to_vec()),
+            },
+        }
+    }
+
+    fn reset(slot: &mut Option<Anchor>, p: usize, version: u64, params: &[f32]) -> EncodedBcast {
+        let values = Arc::new(params.to_vec());
+        let digest = digest_f32(&values);
+        *slot = Some(Anchor {
+            version,
+            values: Arc::clone(&values),
+            digest,
+            deltas_since_full: 0,
+        });
+        EncodedBcast {
+            prefix: prefix(p, version, KIND_FULL, version, &digest),
+            dtype: DTYPE_F32,
+            values,
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// A learner-side mirror: applies broadcasts exactly as yeto/protocol.py.
+    fn apply(anchor: &mut Option<(u64, Vec<f32>)>, enc: &EncodedBcast) -> Result<()> {
+        let version = u64::from_le_bytes(enc.prefix[4..12].try_into().unwrap());
+        let base = u64::from_le_bytes(enc.prefix[13..21].try_into().unwrap());
+        let digest: [u8; 32] = enc.prefix[21..53].try_into().unwrap();
+        let next = match enc.kind() {
+            KIND_FULL => enc.values.to_vec(),
+            KIND_DELTA => {
+                let Some((held, values)) = anchor.as_ref() else { bail!("no anchor") };
+                if *held != base {
+                    bail!("anchor version {held} != base {base}");
+                }
+                values.iter().zip(enc.values.iter()).map(|(a, d)| a + d).collect()
+            }
+            _ => bail!("bad kind"),
+        };
+        if digest_f32(&next) != digest {
+            bail!("digest mismatch");
+        }
+        *anchor = Some((version, next));
+        Ok(())
+    }
+
+    fn params(seed: u64, n: usize) -> Vec<f32> {
+        (0..n)
+            .map(|i| ((i as f64 * 0.37 + seed as f64 * 1.3).sin() * 3.0) as f32)
+            .collect()
+    }
+
+    #[test]
+    fn deltas_track_params_within_bf16_error_and_refresh_every_m() {
+        let mut syncer = DownlinkAnchors::new(1, 4).unwrap();
+        let mut learner = None;
+        let mut theta = params(0, 1000);
+        let mut kinds = vec![];
+        for version in 1..=9u64 {
+            let step = params(version, 1000);
+            let prev_anchor = learner.as_ref().map(|(_, v): &(u64, Vec<f32>)| v.clone());
+            for (t, s) in theta.iter_mut().zip(&step) {
+                *t += 0.01 * s;
+            }
+            let enc = syncer.publish(0, version, &theta).unwrap();
+            kinds.push(enc.kind());
+            apply(&mut learner, &enc).unwrap();
+            let (_, held) = learner.as_ref().unwrap();
+            if enc.kind() == KIND_FULL {
+                assert_eq!(held, &theta);
+            } else {
+                let prev = prev_anchor.unwrap();
+                for ((h, t), a) in held.iter().zip(&theta).zip(&prev) {
+                    // one bf16 rounding of (θ − anchor_prev): rel. 2^-8 bound
+                    assert!((h - t).abs() <= (t - a).abs() / 256.0 + 2.0 * f32::EPSILON * t.abs());
+                }
+            }
+        }
+        assert_eq!(
+            kinds,
+            vec![KIND_FULL, KIND_DELTA, KIND_DELTA, KIND_DELTA, KIND_FULL, KIND_DELTA, KIND_DELTA, KIND_DELTA, KIND_FULL]
+        );
+    }
+
+    #[test]
+    fn resend_is_the_shared_anchor_and_lets_a_reconnected_learner_continue() {
+        let mut syncer = DownlinkAnchors::new(1, 100).unwrap();
+        let mut a = None;
+        let theta1 = params(1, 64);
+        apply(&mut a, &syncer.publish(0, 1, &theta1).unwrap()).unwrap();
+        let theta2: Vec<f32> = theta1.iter().map(|v| v + 0.123).collect();
+        apply(&mut a, &syncer.publish(0, 2, &theta2).unwrap()).unwrap();
+        // A fresh learner (reconnect) gets the anchor, not θ, as FULL.
+        let mut b = None;
+        let full = syncer.resend(0, 2, &theta2);
+        assert_eq!(full.kind(), KIND_FULL);
+        apply(&mut b, &full).unwrap();
+        assert_eq!(a, b);
+        let theta3: Vec<f32> = theta2.iter().map(|v| v - 0.05).collect();
+        let enc = syncer.publish(0, 3, &theta3).unwrap();
+        assert_eq!(enc.kind(), KIND_DELTA);
+        apply(&mut a, &enc).unwrap();
+        apply(&mut b, &enc).unwrap();
+        assert_eq!(a, b);
+    }
+
+    #[test]
+    fn a_learner_with_a_wrong_anchor_detects_it() {
+        let mut syncer = DownlinkAnchors::new(1, 100).unwrap();
+        let mut a = None;
+        apply(&mut a, &syncer.publish(0, 1, &params(1, 32)).unwrap()).unwrap();
+        // Corrupt the learner's anchor bits: the digest check must fail.
+        if let Some((_, v)) = a.as_mut() {
+            v[3] += 1e-3;
+        }
+        let enc = syncer.publish(0, 2, &params(2, 32)).unwrap();
+        assert!(apply(&mut a, &enc).unwrap_err().to_string().contains("digest"));
+        // A learner that skipped a version fails the base-version check.
+        let mut stale = Some((0u64, params(1, 32)));
+        let enc = syncer.publish(0, 3, &params(3, 32)).unwrap();
+        assert!(apply(&mut stale, &enc).unwrap_err().to_string().contains("base"));
+    }
+
+    #[test]
+    fn full_every_below_two_is_refused() {
+        assert!(DownlinkAnchors::new(1, 1).is_err());
+    }
+}

syncer/src/main.rs

@@ -2,6 +2,7 @@ mod elastic;
 mod elastic_server;
 mod iso_worker;
 mod merge;
+mod downlink;
 mod protocol;
 mod server;
 mod state;
@@ -98,6 +99,15 @@ struct Args {
     /// Write the snapshot every N outer steps (0 disables).
     #[arg(long, default_value_t = 8)]
     checkpoint_every: u64,
+    /// Policy sweep: write per-round checkpoints off the scheduler; each
+    /// round is broadcast only after its checkpoint is durable.
+    #[arg(long, default_value_t = false)]
+    async_checkpoint: bool,
+    /// Learners whose session asks for bf16 downlink deltas get a full f32
+    /// BCAST every N-th version of a fragment (N >= 2). Unset: such
+    /// sessions are refused.
+    #[arg(long)]
+    bcast_delta_full_every: Option<u64>,
     /// Resume from --checkpoint-path if it exists.
     #[arg(long, default_value_t = false)]
     resume: bool,
@@ -273,6 +283,8 @@ fn main() -> anyhow::Result<()> {
         final_state: args.final_state,
         checkpoint_path: args.checkpoint_path,
         checkpoint_every: args.checkpoint_every,
+        async_checkpoint: args.async_checkpoint,
+        bcast_delta_full_every: args.bcast_delta_full_every,
         resume: args.resume,
         mark_final_checkpoint: args.mark_final_checkpoint,
         learner_budget_steps: args.learner_budget_steps,