diff --git a/docs/DESIGN.md b/docs/DESIGN.md index f0c6300d..d8aa6e61 100644 --- a/docs/DESIGN.md +++ b/docs/DESIGN.md @@ -52,6 +52,12 @@ Moved here from the README; docs/PROTOCOL.md has the wire-level detail. generations and quorum at launch. Joins/reconnects apply only to future attempts, disconnects do not erase accepted work, and a below-quorum timeout discards partial responses before retrying with a new attempt ID. + One exception under dynamic membership with a count quorum: a learner that + was held by its catch-up when an attempt launched, and none of whose + generations is in it, joins that attempt when its catch-up ends (it gets + the attempt's pull; quorum size, attempt and deadlines are unchanged), so + learners resuming one by one after a restart do not stall a round until + its quorum timeout. - **Sync-interval sensitivity**: at the design-point sync interval (H≈24 inner steps per fragment) DiLoCo matches the synchronous run within noise on held-out CE; at H≈2 — what a LAN/localhost fleet produces naturally, diff --git a/syncer/src/downlink.rs b/syncer/src/downlink.rs index 8b863f57..ca271a75 100644 --- a/syncer/src/downlink.rs +++ b/syncer/src/downlink.rs @@ -31,7 +31,6 @@ use std::sync::Arc; use anyhow::{bail, Result}; -use sha2::{Digest, Sha256}; use crate::protocol::{DTYPE_BF16, DTYPE_F32}; @@ -68,7 +67,7 @@ pub struct DownlinkAnchors { } pub fn digest_f32(values: &[f32]) -> [u8; 32] { - let mut hasher = Sha256::new(); + let mut hasher = crate::sha256::Hasher::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) { @@ -78,7 +77,7 @@ pub fn digest_f32(values: &[f32]) -> [u8; 32] { } hasher.update(&buf); } - hasher.finalize().into() + hasher.finalize() } /// The header of a FULL that carries `params` exactly (no anchor diff --git a/syncer/src/main.rs b/syncer/src/main.rs index cc4efcbf..9bb05cca 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -9,6 +9,7 @@ mod membership; mod recovery; mod resend; mod merge; +mod par; mod dispense; mod dispense_service; mod members; @@ -17,6 +18,7 @@ mod eval_checkpoint; mod protocol; mod sample_relay; mod server; +mod sha256; mod state; mod subscriber; @@ -82,6 +84,11 @@ struct Args { /// Pre-merge learner-delta correction: "heloco" or "none". #[arg(long, default_value = "heloco")] delta_correction: String, + /// Threads for the merge and outer optimizer step (results are + /// bit-identical for any value). Unset: $YETO_SYNCER_MERGE_THREADS; + /// 0 or neither: available cores, capped at 32. + #[arg(long)] + merge_threads: Option, /// Give up waiting for quorum after this long. #[arg(long, default_value_t = 900)] quorum_timeout_s: u64, @@ -141,6 +148,16 @@ struct Args { /// keep merging. --resume never reads these files. #[arg(long, value_delimiter = ',')] checkpoint_at_steps: Vec, + /// Write the --checkpoint-every checkpoints on a background thread while + /// rounds keep merging, instead of stopping the scheduler for the whole + /// write (a 120B state is ~1 TB; minutes per checkpoint). A crash during + /// the write resumes from the previous checkpoint. Per-round checkpoints + /// (--checkpoint-every 1) and policy sweeps are unaffected. Memory: the + /// writer holds the snapshot, so a fragment committed before the writer + /// reached it is copied once; in the worst case (every fragment commits + /// during the write) host memory peaks at about twice params + momentum. + #[arg(long)] + background_checkpoint: bool, /// 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)] @@ -276,6 +293,11 @@ struct Args { /// parameters, which stay in --checkpoint-path / --final-state). #[arg(long, default_value = "full", value_parser = ["full", "manifest"])] final_cut_payload: String, + /// A late delta carried into a later round counts as that island's + /// answer to the round (quorum and the all-answered check). Off keeps + /// carries outside quorum, the historical RL async rule. + #[arg(long)] + carry_answers_round: bool, /// Task dispenser: islands claim the rows of --data they train /// from this syncer (TASK_* connections, docs/PROTOCOL.md "Task /// dispenser"). Off by default. @@ -387,6 +409,12 @@ fn main() -> anyhow::Result<()> { let dispense = dispense_config(&args)?; let member_lease = member_lease(&args)?; let shard = args.shard_plan()?; + let merge_threads = par::resolve_threads( + args.merge_threads, + std::env::var(par::THREADS_ENV).ok().as_deref(), + )?; + par::set_default_threads(merge_threads); + tracing::info!(merge_threads, "merge thread count"); let delta_correction = match args.delta_correction.as_str() { "heloco" => true, "none" => false, @@ -476,6 +504,7 @@ fn main() -> anyhow::Result<()> { checkpoint_path: args.checkpoint_path, checkpoint_every: args.checkpoint_every, checkpoint_at_steps: args.checkpoint_at_steps, + background_checkpoint: args.background_checkpoint, async_checkpoint: args.async_checkpoint, bcast_delta_full_every: args.bcast_delta_full_every, resume: args.resume, @@ -497,6 +526,7 @@ fn main() -> anyhow::Result<()> { shard, expert_token_weights: args.expert_token_weights, final_cut_manifest_only: args.final_cut_payload == "manifest", + carry_answers_round: args.carry_answers_round, dispense, dispense_log: args.dispense_log.clone(), member_lease, diff --git a/syncer/src/merge.rs b/syncer/src/merge.rs index 7110e3b9..5e4feb20 100644 --- a/syncer/src/merge.rs +++ b/syncer/src/merge.rs @@ -418,6 +418,197 @@ pub fn heloco_correct(delta: &mut [f32], momentum: &[f32], h: &Heloco) { } } +// --------------------------------------------------------------------------- +// Parallel kernels for one large tensor. Each returns exactly the bits of its +// serial counterpart above for any thread count (see `crate::par`): only +// element-wise passes are chunked, every reduction stays one sequential loop. +// --------------------------------------------------------------------------- + +/// `l2_norm` of each slice, the slices reduced concurrently. +fn l2_norms_par(slices: &[&[f32]], threads: usize) -> Vec { + let mut norms = vec![0.0f64; slices.len()]; + let tasks: Vec<_> = norms.iter_mut().zip(slices).collect(); + crate::par::run_tasks(threads, tasks, |(n, s)| *n = l2_norm(s)); + norms +} + +/// out[i] = 0 + Σ_m coef_m · Δ_m[i] in m order, skipping `None` coefficients. +fn accumulate_par(deltas: &[&[f32]], coefs: &[Option], out: &mut [f32], threads: usize) { + let tasks: Vec<_> = out.chunks_mut(crate::par::CHUNK).enumerate().collect(); + crate::par::run_tasks(threads, tasks, |(index, o)| { + let base = index * crate::par::CHUNK; + let len = o.len(); + o.fill(0.0); + for (delta, coef) in deltas.iter().zip(coefs) { + let Some(coef) = *coef else { continue }; + for (x, d) in o.iter_mut().zip(&delta[base..base + len]) { + *x += coef * *d; + } + } + }); +} + +/// Parallel `merge_avg`. +pub fn merge_avg_par(deltas: &[&[f32]], weights: &[f64], out: &mut [f32], threads: usize) { + if threads <= 1 { + return merge_avg(deltas, weights, out); + } + let wsum: f64 = weights.iter().sum(); + if wsum <= 0.0 { + out.fill(0.0); + return; + } + let coefs: Vec> = weights.iter().map(|&w| Some((w / wsum) as f32)).collect(); + accumulate_par(deltas, &coefs, out, threads); +} + +/// Parallel `merge_rda`. +pub fn merge_rda_par(deltas: &[&[f32]], weights: &[f64], out: &mut [f32], threads: usize) { + if threads <= 1 { + return merge_rda(deltas, weights, out); + } + let wsum: f64 = weights.iter().sum(); + if wsum <= 0.0 { + out.fill(0.0); + return; + } + let norms = l2_norms_par(deltas, threads); + let radial = rda_radial(&norms, weights, wsum); + let coefs: Vec> = weights + .iter() + .zip(&norms) + .map(|(&w, &n)| (n != 0.0).then(|| (w / wsum / n) as f32)) + .collect(); + accumulate_par(deltas, &coefs, out, threads); + if !rda_rescale_par(out, radial, threads) { + merge_avg_par(deltas, weights, out, threads); + } +} + +/// Parallel `rda_rescale`: the norm is one sequential pass, the scaling is +/// chunked. +pub(crate) fn rda_rescale_par(out: &mut [f32], radial: f64, threads: usize) -> bool { + if threads <= 1 { + return rda_rescale(out, radial); + } + let mean_dir_norm = out + .iter() + .map(|v| (*v as f64) * (*v as f64)) + .sum::() + .sqrt(); + if mean_dir_norm < 1e-12 { + return false; + } + let scale = (radial / mean_dir_norm) as f32; + crate::par::run_tasks(threads, out.chunks_mut(crate::par::CHUNK).collect(), |o| { + for v in o.iter_mut() { + *v *= scale; + } + }); + true +} + +/// Parallel `heloco_correct`: ‖Δ‖, ‖m‖ and Δ·m are three independent +/// sequential reductions run concurrently; the correction itself is +/// chunked. The rotation branch's ‖ũ‖² is accumulated sequentially from the +/// same f64 terms the serial loop sums, before any element is written. +pub fn heloco_correct_par(delta: &mut [f32], momentum: &[f32], h: &Heloco, threads: usize) { + if threads <= 1 { + return heloco_correct(delta, momentum, h); + } + debug_assert_eq!(delta.len(), momentum.len()); + let ro: &[f32] = delta; + let (du, dm, dot) = std::thread::scope(|scope| { + let du = scope.spawn(|| ro.iter().map(|v| (*v as f64).powi(2)).sum::().sqrt()); + let dot = scope.spawn(|| { + ro.iter() + .zip(momentum) + .map(|(d, m)| *d as f64 * *m as f64) + .sum::() + }); + let dm = momentum.iter().map(|v| (*v as f64).powi(2)).sum::().sqrt(); + (du.join().unwrap(), dm, dot.join().unwrap()) + }); + if du < h.eps || dm < h.eps { + return; + } + let c = dot / (du * dm); + if c >= h.c_ok { + return; + } + let conf = du / (du + h.kappa * dm + h.eps); + fn chunks<'a>(delta: &'a mut [f32], momentum: &'a [f32]) -> Vec<(&'a mut [f32], &'a [f32])> { + delta.chunks_mut(crate::par::CHUNK).zip(momentum.chunks(crate::par::CHUNK)).collect() + } + if c < 0.0 { + let beta = (h.k_s * (-c) * conf).min(h.beta_max); + let coef = (-beta * c * du / dm) as f32; + crate::par::run_tasks(threads, chunks(delta, momentum), |(d, m)| { + for (d, m) in d.iter_mut().zip(m) { + *d += coef * *m; + } + }); + } else { + let lambda = (h.k_d * (1.0 - c) * conf).min(1.0); + let (wu, wv) = ((1.0 - lambda) / du, lambda / dm); + let mut norm_sq = 0.0f64; + for (d, m) in delta.iter().zip(momentum) { + let t = wu * *d as f64 + wv * *m as f64; + norm_sq += t * t; + } + let scale = (du / norm_sq.sqrt().max(h.eps)) as f32; + crate::par::run_tasks(threads, chunks(delta, momentum), |(d, m)| { + for (d, m) in d.iter_mut().zip(m) { + let t = wu * *d as f64 + wv * *m as f64; + *d = t as f32 * scale; + } + }); + } +} + +/// Parallel `nesterov_step`. +pub fn nesterov_step_par( + params: &mut [f32], + buf: &mut [f32], + delta: &[f32], + lr: f32, + mu: f32, + threads: usize, +) { + let c = crate::par::CHUNK; + let tasks: Vec<_> = params.chunks_mut(c).zip(buf.chunks_mut(c)).zip(delta.chunks(c)).collect(); + crate::par::run_tasks(threads, tasks, |((p, b), d)| nesterov_step(p, b, d, lr, mu)); +} + +/// Parallel `delayed_nesterov_step`. +#[allow(clippy::too_many_arguments)] +pub fn delayed_nesterov_step_par( + params: &mut [f32], + buf: &mut [f32], + state: &mut DelayedNesterovState, + delta: &[f32], + lr: f32, + mu: f32, + period: u32, + threads: usize, +) { + let update = delayed_nesterov_update(state.step, period); + let c = crate::par::CHUNK; + let accum = std::sync::Arc::make_mut(&mut state.accum); + let tasks: Vec<_> = params + .chunks_mut(c) + .zip(buf.chunks_mut(c)) + .zip(accum.chunks_mut(c)) + .zip(delta.chunks(c)) + .collect(); + crate::par::run_tasks(threads, tasks, |(((p, b), a), d)| { + for (((p, b), a), d) in p.iter_mut().zip(b.iter_mut()).zip(a.iter_mut()).zip(d) { + (*p, *b, *a) = delayed_nesterov_elem(*p, *b, *a, *d, lr, mu, period, update); + } + }); + state.step += 1; +} + #[cfg(test)] mod tests { use super::*; diff --git a/syncer/src/par.rs b/syncer/src/par.rs new file mode 100644 index 00000000..2c873363 --- /dev/null +++ b/syncer/src/par.rs @@ -0,0 +1,163 @@ +//! Bit-exact multi-threaded execution of the merge hot path. +//! +//! Every parallel kernel built on these helpers produces exactly the bits of +//! its serial counterpart, for every thread count and every scheduling: +//! +//! * Element-wise passes (weighted accumulation, rescaling, the outer +//! optimizer step, finiteness checks) are split into disjoint chunks. Each +//! element still sees the same operations in the same order (for example +//! the learners of an average are added in the same sequence), so the +//! split is invisible in the result. +//! * Floating-point reductions (L2 norms, dot products) are never split. +//! Each one runs as a single sequential loop on one thread, exactly as in +//! the serial code; parallelism comes from running independent +//! reductions (different tensors, different learners, the separate norms +//! of HeLoCo) at the same time. +//! +//! Work is pulled from a shared queue, so which thread runs a task depends +//! on timing, but tasks write disjoint outputs and never share an +//! accumulator, so that never affects the values. + +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::Mutex; + +use anyhow::{bail, Result}; + +/// Element-wise chunk length (f32 values). +pub const CHUNK: usize = 1 << 16; + +/// Environment variable read when `--merge-threads` is not given. +pub const THREADS_ENV: &str = "YETO_SYNCER_MERGE_THREADS"; + +/// Upper bound of the automatic thread count. The merge is memory-bandwidth +/// bound, so more threads than this rarely help and only steal cores from +/// the network runtime. +const AUTO_MAX_THREADS: usize = 32; + +static DEFAULT_THREADS: AtomicUsize = AtomicUsize::new(1); + +/// Thread count new `GlobalState`s start with (1 unless the binary set it). +pub fn default_threads() -> usize { + DEFAULT_THREADS.load(Ordering::Relaxed).max(1) +} + +pub fn set_default_threads(threads: usize) { + DEFAULT_THREADS.store(threads.max(1), Ordering::Relaxed); +} + +/// Automatic thread count: the available cores, capped. +pub fn auto_threads() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(1) + .clamp(1, AUTO_MAX_THREADS) +} + +/// Resolve the merge thread count from the CLI value, else the environment. +/// Absent or 0 selects `auto_threads`. +pub fn resolve_threads(cli: Option, env: Option<&str>) -> Result { + let requested = match (cli, env.map(str::trim).filter(|s| !s.is_empty())) { + (Some(n), _) => n, + (None, Some(text)) => match text.parse::() { + Ok(n) => n, + Err(_) => bail!("{THREADS_ENV} must be a non-negative integer, got {text:?}"), + }, + (None, None) => 0, + }; + Ok(if requested == 0 { auto_threads() } else { requested }) +} + +/// Run `f` on every task using up to `threads` threads. With one thread (or +/// one task) the tasks run in order on the caller's thread. +pub fn run_tasks(threads: usize, tasks: Vec, f: impl Fn(T) + Sync) { + let threads = threads.min(tasks.len()); + if threads <= 1 { + tasks.into_iter().for_each(f); + return; + } + let queue = Mutex::new(tasks.into_iter()); + let (queue, f) = (&queue, &f); + std::thread::scope(|scope| { + for _ in 0..threads { + scope.spawn(move || loop { + // Release the lock before running the task. + let next = queue.lock().unwrap_or_else(|e| e.into_inner()).next(); + match next { + Some(task) => f(task), + None => break, + } + }); + } + }); +} + +/// Split `buf` into consecutive mutable slices of the given lengths. +pub fn split_lengths<'a>(mut buf: &'a mut [f32], lengths: &[usize]) -> Vec<&'a mut [f32]> { + let mut out = Vec::with_capacity(lengths.len()); + for &n in lengths { + let (head, tail) = std::mem::take(&mut buf).split_at_mut(n); + out.push(head); + buf = tail; + } + out +} + +/// Whether any value is NaN or infinite (same answer as the serial +/// `iter().any(|v| !v.is_finite())`). +pub fn any_non_finite(threads: usize, values: &[f32]) -> bool { + if threads <= 1 { + return values.iter().any(|v| !v.is_finite()); + } + let found = AtomicBool::new(false); + run_tasks(threads, values.chunks(CHUNK).collect(), |chunk| { + if !found.load(Ordering::Relaxed) && chunk.iter().any(|v| !v.is_finite()) { + found.store(true, Ordering::Relaxed); + } + }); + found.into_inner() +} + +/// `dst.copy_from_slice(src)` in parallel chunks. +pub fn copy(threads: usize, dst: &mut [f32], src: &[f32]) { + let tasks: Vec<_> = dst.chunks_mut(CHUNK).zip(src.chunks(CHUNK)).collect(); + run_tasks(threads, tasks, |(d, s)| d.copy_from_slice(s)); +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn resolve_prefers_cli_then_env_then_auto() { + assert_eq!(resolve_threads(Some(3), Some("7")).unwrap(), 3); + assert_eq!(resolve_threads(None, Some("7")).unwrap(), 7); + assert_eq!(resolve_threads(None, Some(" ")).unwrap(), auto_threads()); + assert_eq!(resolve_threads(Some(0), None).unwrap(), auto_threads()); + assert_eq!(resolve_threads(None, Some("0")).unwrap(), auto_threads()); + assert!(resolve_threads(None, Some("x")).is_err()); + } + + #[test] + fn run_tasks_visits_every_task_once() { + for threads in [1, 2, 7] { + let hits: Vec = (0..100).map(|_| AtomicUsize::new(0)).collect(); + run_tasks(threads, (0..100).collect(), |i: usize| { + hits[i].fetch_add(1, Ordering::Relaxed); + }); + assert!(hits.iter().all(|h| h.load(Ordering::Relaxed) == 1)); + } + } + + #[test] + fn split_and_non_finite() { + let mut buf = vec![0.0f32; 10]; + let parts = split_lengths(&mut buf, &[3, 0, 7]); + assert_eq!(parts.iter().map(|p| p.len()).collect::>(), [3, 0, 7]); + let mut values = vec![1.0f32; 3 * CHUNK + 5]; + assert!(!any_non_finite(4, &values)); + values[2 * CHUNK + 1] = f32::NAN; + assert!(any_non_finite(4, &values)); + values[2 * CHUNK + 1] = f32::NEG_INFINITY; + assert!(any_non_finite(4, &values)); + } +} diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index 2a562d73..e6655666 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -452,8 +452,30 @@ pub fn decode_budget_done(payload: &[u8]) -> Result { Ok(local_steps) } -/// Decode a tensor payload into f32s. +/// Decode a tensor payload into f32s on the calling thread. PUSH payloads +/// decode here, once per connection; only INIT_PARAMS (a whole fragment, +/// off the read loop) asks for threads through `decode_tensor_with`. pub fn decode_tensor(dtype: u8, bytes: &[u8], out: &mut Vec) -> Result<()> { + decode_tensor_with(1, dtype, bytes, out) +} + +/// Values below which a decode stays on the calling thread. +const PAR_DECODE_MIN: usize = 1 << 22; + +/// `decode_tensor` on up to `threads` threads: disjoint chunks, each value +/// converted exactly as the serial loop does, so the result (and the error +/// for a non-finite value) is the same for any thread count. A 120B INIT +/// fragment is ~15 GB of f32; one thread decodes and checks it at ~1 GB/s +/// on the read loop. +pub fn decode_tensor_with(threads: usize, dtype: u8, bytes: &[u8], out: &mut Vec) -> Result<()> { + let width = match dtype { + DTYPE_F32 => 4, + DTYPE_BF16 => 2, + _ => 0, + }; + if threads > 1 && width > 0 && bytes.len() % width == 0 && bytes.len() / width >= PAR_DECODE_MIN { + return decode_tensor_par(threads, dtype, width, bytes, out); + } out.clear(); match dtype { DTYPE_F32 => { @@ -485,6 +507,42 @@ pub fn decode_tensor(dtype: u8, bytes: &[u8], out: &mut Vec) -> Result<()> Ok(()) } +fn decode_tensor_par(threads: usize, dtype: u8, width: usize, bytes: &[u8], out: &mut Vec) -> Result<()> { + use std::sync::atomic::{AtomicBool, Ordering}; + + let numel = bytes.len() / width; + out.clear(); + out.try_reserve(numel) + .map_err(|error| anyhow::anyhow!("cannot allocate {} tensor: {error}", if width == 4 { "f32" } else { "bf16" }))?; + let non_finite = AtomicBool::new(false); + let spare = &mut out.spare_capacity_mut()[..numel]; + let tasks: Vec<_> = spare + .chunks_mut(crate::par::CHUNK) + .zip(bytes.chunks(crate::par::CHUNK * width)) + .collect(); + crate::par::run_tasks(threads, tasks, |(dst, src)| { + let mut finite = true; + for (d, c) in dst.iter_mut().zip(src.chunks_exact(width)) { + let value = if dtype == DTYPE_F32 { + f32::from_le_bytes(c.try_into().unwrap()) + } else { + f32::from_bits((u16::from_le_bytes(c.try_into().unwrap()) as u32) << 16) + }; + finite &= value.is_finite(); + d.write(value); + } + if !finite { + non_finite.store(true, Ordering::Relaxed); + } + }); + // SAFETY: the chunks cover `spare` exactly and every element was written. + unsafe { out.set_len(numel) }; + if non_finite.into_inner() { + bail!("tensor payload contains a non-finite value"); + } + Ok(()) +} + /// Encode f32s into the wire dtype. #[cfg(test)] pub fn encode_tensor(dtype: u8, vals: &[f32], out: &mut Vec) -> Result<()> { @@ -852,6 +910,37 @@ mod tests { assert!(decode_i4(&pad, 65, &mut out).is_err()); } + #[test] + fn parallel_tensor_decode_matches_the_serial_bits_and_errors() { + let n = PAR_DECODE_MIN + 3 * crate::par::CHUNK + 17; + let vals: Vec = (0..n).map(|i| ((i as f32) * 0.37).sin() * 1e3 - 0.5).collect(); + for dtype in [DTYPE_F32, DTYPE_BF16] { + let mut wire = Vec::new(); + encode_tensor(dtype, &vals, &mut wire).unwrap(); + let mut serial = Vec::new(); + decode_tensor_with(1, dtype, &wire, &mut serial).unwrap(); + for threads in [2, 3, 8] { + let mut par = vec![7.0f32; 5]; + decode_tensor_with(threads, dtype, &wire, &mut par).unwrap(); + assert_eq!(par.len(), serial.len()); + assert!(par.iter().zip(&serial).all(|(a, b)| a.to_bits() == b.to_bits()), "{dtype} {threads}"); + } + let width = if dtype == DTYPE_F32 { 4 } else { 2 }; + for at in [0, crate::par::CHUNK - 1, n - 1] { + for bad in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + let mut broken = wire.clone(); + let mut one = Vec::new(); + encode_tensor(dtype, &[bad], &mut one).unwrap(); + broken[at * width..(at + 1) * width].copy_from_slice(&one); + let serial = decode_tensor_with(1, dtype, &broken, &mut Vec::new()).unwrap_err(); + let par = decode_tensor_with(8, dtype, &broken, &mut Vec::new()).unwrap_err(); + assert_eq!(serial.to_string(), par.to_string()); + } + } + assert!(decode_tensor_with(8, dtype, &wire[..wire.len() - 1], &mut Vec::new()).is_err()); + } + } + #[test] fn int8_push_rejects_malformed_payloads() { let vals = vec![1.0f32; 64]; diff --git a/syncer/src/server.rs b/syncer/src/server.rs index aed60d58..f15d96b1 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -121,6 +121,10 @@ pub struct Config { /// the quiescent cut between rounds, resumed from when `resume` is set. pub checkpoint_path: Option, pub checkpoint_every: u64, + /// Write the `checkpoint_every` checkpoints on a background thread while + /// rounds keep merging (`GlobalState::save_checkpoint_in_background`) + /// instead of on the scheduler. Every other checkpoint waits for it. + pub background_checkpoint: bool, /// Global steps whose cut is also written to `.step-` /// off the scheduler (`crate::eval_checkpoint`); strictly increasing, /// below `total_steps`. Empty: none. @@ -209,6 +213,13 @@ pub struct Config { /// must expect it, so it is in the semantic profile hash (appended only /// when on). Off sends every final fragment to every island. pub final_cut_manifest_only: bool, + /// A late delta carried into a round (async rounds, `LateCarry`) answers + /// that round for its expected island: it counts towards quorum and the + /// all-answered check. The carry spends the island's anchor, so without + /// this a quorum-1 round holding only carried deltas waits for the + /// slowest island. Off keeps carries outside quorum (the historical RL + /// async rule). In the semantic profile hash (appended only when on). + pub carry_answers_round: bool, /// Task dispenser (`dispense.rs`): islands claim the rows they /// train over TASK_* connections. Hands out data only, never changes a /// merge, so it is not in the semantic profile hash. None: off. @@ -234,6 +245,7 @@ const SEMANTIC_PROFILE_DOMAIN: &[u8] = b"yeto-syncer-semantic-profile-v2\0"; const DELAYED_NESTEROV_PROFILE_DOMAIN: &[u8] = b"yeto-syncer-delayed-nesterov-v1\0"; const EXPERT_TOKEN_WEIGHTS_PROFILE_DOMAIN: &[u8] = b"yeto-syncer-expert-token-weights-v1\0"; const FINAL_CUT_MANIFEST_ONLY_PROFILE_DOMAIN: &[u8] = b"yeto-syncer-final-cut-manifest-only-v1\0"; +const CARRY_ANSWERS_ROUND_PROFILE_DOMAIN: &[u8] = b"yeto-syncer-carry-answers-round-v1\0"; const GRACE_FLOOR_PROFILE_DOMAIN: &[u8] = b"yeto-syncer-grace-floor-v1\0"; const TERMINAL_GRACE_PROFILE_DOMAIN: &[u8] = b"yeto-syncer-terminal-grace-v1\0"; const EVAL_CHECKPOINTS_PROFILE_DOMAIN: &[u8] = b"yeto-syncer-eval-checkpoints-v1\0"; @@ -305,6 +317,10 @@ impl Config { encoded.extend_from_slice(&step.to_le_bytes()); } } + // Appended only when on, so every existing hash stays unchanged. + if self.carry_answers_round { + encoded.extend_from_slice(CARRY_ANSWERS_ROUND_PROFILE_DOMAIN); + } Sha256::digest(encoded).into() } } @@ -731,7 +747,45 @@ fn append_tensor_bytes( } let first = byte_offset / width; let last = byte_end.div_ceil(width); - for (index, value) in values[first..last].iter().enumerate() { + // Values cut by the range ends go byte by byte; the whole values between + // them are encoded in one pass (a broadcast chunk is 4 MiB, and the + // per-value slicing alone held a forwarder to ~0.4 GB/s). + let first_whole = byte_offset.div_ceil(width).min(last); + let last_whole = (byte_end / width).max(first_whole); + append_value_bytes(dtype, width, values, first..first_whole, byte_offset, byte_end, out); + let whole = &values[first_whole..last_whole]; + let start = out.len(); + out.resize(start + whole.len() * width, 0); + let dst = &mut out[start..]; + match dtype { + DTYPE_F32 => { + for (d, value) in dst.chunks_exact_mut(4).zip(whole) { + d.copy_from_slice(&value.to_le_bytes()); + } + } + DTYPE_BF16 => { + for (d, value) in dst.chunks_exact_mut(2).zip(whole) { + d.copy_from_slice(&half::bf16::from_f32(*value).to_bits().to_le_bytes()); + } + } + _ => unreachable!(), + } + append_value_bytes(dtype, width, values, last_whole..last, byte_offset, byte_end, out); + Ok(()) +} + +/// The bytes of `values[range]` that fall in `byte_offset..byte_end`. +fn append_value_bytes( + dtype: u8, + width: usize, + values: &[f32], + range: std::ops::Range, + byte_offset: usize, + byte_end: usize, + out: &mut Vec, +) { + let first = range.start; + for (index, value) in values[range].iter().enumerate() { let value_index = first + index; let value_start = value_index * width; let start = byte_offset.saturating_sub(value_start); @@ -745,7 +799,6 @@ fn append_tensor_bytes( _ => unreachable!(), } } - Ok(()) } /// One receive-only subscriber connection (`subscriber.rs`). @@ -1084,9 +1137,65 @@ struct Push { /// Expert token weights: the per-tensor merge scale of this push (see /// `protocol::decode_expert_tags`); `None` when the option is off. tensor_scales: Option>, + /// The learner's wire payload, kept only while it is smaller than the + /// decoded gradient, so a late push queued for its fragment's next + /// round (`LateCarry`) can drop `outer_gradient` and wait as wire bytes. + wire: Option, transfer: PushTransfer, } +/// A PUSH_FRAGMENT delta payload as received (dtype, tensor segments, bytes). +struct PushWire { + dtype: u8, + segments: Vec, + bytes: Vec, +} + +impl Push { + /// Keep only the wire bytes (a late push waiting ~F rounds: an int4 + /// payload is 1/7.5 of its f32 gradient). No-op without wire bytes. + fn compress(&mut self) { + if self.wire.is_some() { + self.outer_gradient = Vec::new(); + } + } + + /// A push that joins a round keeps its decoded gradient only: the wire + /// copy exists for a push that waits in the late-carry queue. + fn release_wire(&mut self) { + self.wire = None; + } + + /// Undo `compress`: decode the wire bytes exactly as on receipt, so the + /// gradient is bit-identical to the one dropped. + fn expand(&mut self) { + if let Some(wire) = self.wire.take() { + if self.outer_gradient.is_empty() { + self.outer_gradient = decode_outer_gradient(&wire) + .expect("a push payload that decoded on receipt decodes again"); + } + } + } +} + +/// `decode_outer_gradient` of `bytes` with `wire`'s dtype and segments. +fn decode_push_segments_to_gradient(wire: &PushWire, bytes: &[u8]) -> Result> { + let mut outer_gradient = Vec::new(); + decode_push_segments(wire.dtype, bytes, &wire.segments, &mut outer_gradient)?; + // Every wire dtype carries learner_delta = local - raw_anchor. + // Merge math consumes the corresponding outer gradient with the + // opposite sign and never reconstructs a full learner parameter. + for value in &mut outer_gradient { + *value = -*value; + } + Ok(outer_gradient) +} + +/// The outer gradient `raw_anchor - local` of a push payload. +fn decode_outer_gradient(wire: &PushWire) -> Result> { + decode_push_segments_to_gradient(wire, &wire.bytes) +} + /// When and how fast one PUSH_FRAGMENT reached the syncer (full-parameter /// timing): `recv_seconds` runs from the first byte of the push (its first /// chunk, or the frame header) to its last byte; `arrived_unix` is the wall @@ -1710,6 +1819,16 @@ fn validate_config(cfg: &Config) -> Result<()> { if matches!(cfg.bcast_delta_full_every, Some(n) if n < 2) { bail!("--bcast-delta-full-every must be at least 2"); } + if cfg.background_checkpoint + && (cfg.checkpoint_path.is_none() + || cfg.checkpoint_every < 2 + || cfg.policy_sweep_fragments.is_some()) + { + bail!( + "--background-checkpoint needs --checkpoint-path and --checkpoint-every of at least 2, \ + without --policy-sweep-fragments (per-round checkpoints persist before their broadcast)" + ); + } if cfg.async_checkpoint { if cfg.policy_sweep_fragments.is_none() { bail!("--async-checkpoint requires --policy-sweep-fragments"); @@ -3048,8 +3167,29 @@ async fn read_loop( group: &Arc, event_tx: &mpsc::Sender, budget_cutoff: &BudgetCutoff, +) -> Result<()> { + let mut inits = std::collections::VecDeque::new(); + let result = read_frames(rd, group, event_tx, budget_cutoff, &mut inits).await; + // A failed INIT decode is this connection's protocol error too. + let decoded = reap_init_params(&mut inits, true).await; + if let (Ok(()), Err(error)) = (&result, &decoded) { + report_protocol_error(group, error).await; + } + result.and(decoded) +} + +async fn read_frames( + rd: &mut (impl tokio::io::AsyncReadExt + Unpin), + group: &Arc, + event_tx: &mpsc::Sender, + budget_cutoff: &BudgetCutoff, + inits: &mut std::collections::VecDeque>>, ) -> Result<()> { loop { + if let Err(error) = reap_init_params(inits, false).await { + report_protocol_error(group, &error).await; + return Err(error); + } let mut header_at: Option = None; let frame = match read_frame_limited(rd, |msg_type| { header_at = Some(Instant::now()); @@ -3085,11 +3225,15 @@ async fn read_loop( group.wire.recv.fetch_add(frame.payload.len() as u64 + 13, Ordering::Relaxed); let dispatched = match frame.msg_type { MSG_CHUNK => match reassemble(group, &frame.payload) { + Ok(Some((inner, _))) if inner.msg_type == MSG_INIT_PARAMS => { + inits.push_back(spawn_init_params(group, inner, event_tx)); + Ok(()) + } Ok(Some((inner, recv_seconds))) => { dispatch_inner( group, inner.msg_type, - &inner.payload, + inner.payload(), recv_seconds, event_tx, budget_cutoff, @@ -3119,7 +3263,7 @@ async fn report_protocol_error(group: &Arc, error: &anyhow::Error) { /// Adds one chunk; a completed message comes back with its receive time /// (first chunk to last chunk). -fn reassemble(group: &Arc, payload: &[u8]) -> Result> { +fn reassemble(group: &Arc, payload: &[u8]) -> Result> { let mut r = Reader(payload); let msg_id = r.u64()?; let total_u64 = r.u64()?; @@ -3197,17 +3341,107 @@ fn reassemble(group: &Arc, payload: &[u8]) -> Result } let msg_type = r.u8()?; let len = r.u64()? as usize; - let payload = r.rest(); - if payload.len() != len { + if r.rest().len() != len { bail!("inner frame length mismatch"); } - Ok(Some(( - Frame { - msg_type, - payload: payload.to_vec(), - }, - recv_seconds, - ))) + // The payload stays in the buffer: copying a model-sized INIT out of it + // was one more single-threaded pass on the read loop. + Ok(Some((Reassembled { msg_type, buf }, recv_seconds))) +} + +/// A reassembled inner frame: its 13-byte header, then the payload. +struct Reassembled { + msg_type: u8, + buf: Vec, +} + +impl Reassembled { + fn payload(&self) -> &[u8] { + &self.buf[13..] + } +} + +/// INIT_PARAMS of one connection hashed and decoded at once, off its read +/// loop. Each holds a fragment (~30 GB transient for a 120B f32 fragment: +/// the payload and its decoded values), and the digest is one sequential +/// ~1.6 GB/s pass, so the read loop keeps receiving the next fragments +/// while earlier ones are checked instead of stalling its stream. +const INIT_PIPELINE: usize = 4; + +/// Validate, hash and decode one INIT_PARAMS payload: (fragment, SHA-256 of +/// its tensor bytes, values). +fn decode_init_params(group: &Group, payload: &[u8]) -> Result<(u32, [u8; 32], Vec)> { + if group.member.learner_id != 0 { + bail!("only learner 0 may send INIT_PARAMS"); + } + let mut r = Reader(payload); + let fragment_id = r.u32()?; + let fragment = group + .layout + .fragments + .get(fragment_id as usize) + .with_context(|| format!("INIT_PARAMS for unknown fragment {fragment_id}"))?; + let numel = fragment.numel()?; + let bytes = r.rest(); + let expected = tensor_nbytes(bulk_dtype(group.dtype), numel)?; + if bytes.len() != expected { + bail!( + "INIT_PARAMS fragment {fragment_id} has {} tensor bytes, expected {expected}", + bytes.len() + ); + } + // The digest is one sequential pass; decode the values meanwhile. + let (sha256, values) = std::thread::scope(|scope| { + let digest = scope.spawn(|| crate::sha256::digest(bytes)); + let mut values = Vec::new(); + let decoded = decode_tensor_with(crate::par::default_threads(), bulk_dtype(group.dtype), bytes, &mut values); + (digest.join().expect("INIT digest thread panicked"), decoded.map(|()| values)) + }); + Ok((fragment_id, sha256, values?)) +} + +/// Hash and decode a reassembled INIT_PARAMS on the blocking pool, then +/// hand it to the coordinator. The coordinator applies INITs in any +/// order (each fragment once, from the current generation only). +fn spawn_init_params( + group: &Arc, + inner: Reassembled, + event_tx: &mpsc::Sender, +) -> tokio::task::JoinHandle> { + let (group, event_tx) = (Arc::clone(group), event_tx.clone()); + tokio::spawn(async move { + let member = group.member; + let (fragment_id, sha256, values) = + tokio::task::spawn_blocking(move || decode_init_params(&group, inner.payload())) + .await + .context("INIT_PARAMS decoder panicked")??; + event_tx + .send(Event::Init { + member, + fragment_id, + values, + sha256, + }) + .await + .ok(); + Ok(()) + }) +} + +/// Wait for the finished (all, with `all`) INIT_PARAMS decoders of a read +/// loop, oldest first; the first failure is returned. +async fn reap_init_params( + pending: &mut std::collections::VecDeque>>, + all: bool, +) -> Result<()> { + while let Some(oldest) = pending.front() { + if !all && !oldest.is_finished() && pending.len() < INIT_PIPELINE { + break; + } + let oldest = pending.pop_front().expect("front exists"); + oldest.await.context("INIT_PARAMS task panicked")??; + } + Ok(()) } async fn dispatch_inner( @@ -3220,28 +3454,7 @@ async fn dispatch_inner( ) -> Result<()> { match msg_type { MSG_INIT_PARAMS => { - if group.member.learner_id != 0 { - bail!("only learner 0 may send INIT_PARAMS"); - } - let mut r = Reader(payload); - let fragment_id = r.u32()?; - let fragment = group - .layout - .fragments - .get(fragment_id as usize) - .with_context(|| format!("INIT_PARAMS for unknown fragment {fragment_id}"))?; - let numel = fragment.numel()?; - let bytes = r.rest(); - let expected = tensor_nbytes(bulk_dtype(group.dtype), numel)?; - if bytes.len() != expected { - bail!( - "INIT_PARAMS fragment {fragment_id} has {} tensor bytes, expected {expected}", - bytes.len() - ); - } - let sha256: [u8; 32] = Sha256::digest(bytes).into(); - let mut values = Vec::new(); - decode_tensor(bulk_dtype(group.dtype), bytes, &mut values)?; + let (fragment_id, sha256, values) = decode_init_params(group, payload)?; event_tx .send(Event::Init { member: group.member, @@ -3310,16 +3523,12 @@ async fn dispatch_inner( None }; let bytes = r.rest(); - let mut outer_gradient = Vec::new(); let segments = push_segments(group.dtype, numel, &fragment.tensor_numels)?; - decode_push_segments(group.dtype, bytes, &segments, &mut outer_gradient) + let wire = PushWire { dtype: group.dtype, segments, bytes: Vec::new() }; + let outer_gradient = decode_push_segments_to_gradient(&wire, bytes) .with_context(|| format!("PUSH_FRAGMENT fragment {fragment_id}"))?; - // Every wire dtype carries learner_delta = local - raw_anchor. - // Merge math consumes the corresponding outer gradient with the - // opposite sign and never reconstructs a full learner parameter. - for value in &mut outer_gradient { - *value = -*value; - } + // Only a payload smaller than its gradient is worth keeping. + let wire = (bytes.len() < 4 * numel).then(|| PushWire { bytes: bytes.to_vec(), ..wire }); event_tx .send(Event::Push { member: group.member, @@ -3334,6 +3543,7 @@ async fn dispatch_inner( c_tokens, outer_gradient, tensor_scales, + wire, transfer: PushTransfer { recv_seconds, arrived_unix: unix_now(), @@ -4088,6 +4298,7 @@ async fn scheduler( fold, pushes: HashMap::new(), carried: Vec::new(), + carry_answers: cfg.carry_answers_round, }); if let Some(carry) = late_carry.as_mut() { let round = inflight.last_mut().expect("round just launched"); @@ -5265,13 +5476,32 @@ struct Round { /// `LateCarry`): merged here against their own base, in arrival order. /// They never count toward quorum and stay buffered until the merge. carried: Vec<(Member, Push)>, + /// `Config::carry_answers_round`, frozen at launch. + carry_answers: bool, } impl Round { + /// Islands that answered this round: its own pushes plus, with + /// `carry_answers`, expected islands whose late delta of an earlier + /// round was carried into it (see `Config::carry_answers_round`). + fn responders(&self) -> Vec<&Member> { + let mut responders: Vec<&Member> = self.pushes.keys().collect(); + if !self.carry_answers { + return responders; + } + for (member, _push) in &self.carried { + if self.expected_members.contains(member) && !responders.contains(&member) { + responders.push(member); + } + } + responders + } + fn quorum_reached(&self) -> bool { + let responders = self.responders(); match &self.capacity_quorum { - Some(quorum) => quorum.reached(&self.expected_members, self.pushes.keys()), - None => self.pushes.len() >= self.quorum_size, + Some(quorum) => quorum.reached(&self.expected_members, responders), + None => responders.len() >= self.quorum_size, } } @@ -5362,7 +5592,8 @@ enum RoundAction { fn round_action(round: &Round, now: Instant) -> RoundAction { // A round never merges nothing: if lease expiry emptied its expected // set, it waits for the quorum timeout and relaunches. - if !round.pushes.is_empty() && round.pushes.len() >= round.expected_members.len() { + let responders = round.responders().len(); + if responders > 0 && responders >= round.expected_members.len() { return RoundAction::Complete; } if round.quorum_reached() { @@ -5494,6 +5725,90 @@ fn rebind_current_unanswered_rounds( repulls } +/// A PULL that a member gets on joining a round launched without it. +#[derive(Debug, Eq, PartialEq)] +struct JoinedPull { + t: u64, + p: usize, + attempt: u32, + base_version: u64, + pull: bytes::Bytes, +} + +/// Admit a member that just left its catch-up (or recovery replay) to the +/// rounds still gathering that were launched without any generation of it. +/// +/// A round's expected set is frozen at launch to the members not held by a +/// catch-up then. Without this, a learner that finishes its catch-up after +/// the launch is pulled only when the quorum timeout relaunches the round, +/// so one expected learner that cannot answer yet (an H-gate far from +/// satisfied) stalls the round, and every later commit behind it, for the +/// whole timeout. Joining is what the launch would have done had the member +/// been ready then: the member becomes expected in the current attempt and +/// gets that attempt's pull; quorum size, attempt, base version and the +/// quorum, soft and grace deadlines stay as they are (a grace window already +/// armed is not extended, the joiner's push counts if it arrives before it). +/// +/// Skipped: a round that already completed, one where this learner has a +/// slot, a push, or a carried delta (a reconnect is rebound, never pulled +/// twice), one whose base version the member was not caught up to +/// (`floor`), a round under a capacity quorum (its threshold assumes an +/// expected set that only shrinks), and an exact fold that already holds +/// pushes (its roster is fixed; an empty one is rebuilt for the new roster). +/// Fixed rosters never call this: every logical learner is expected from +/// launch. +fn join_gathering_rounds( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + rounds: &mut [Round], + member: Member, + floor: &[u64], + now: Instant, +) -> Vec { + // Same lock discipline as the rebind: a newer generation cannot become + // current between the check and the membership rewrites. + let registry = registry.lock().unwrap(); + if registry.current.get(&member.learner_id) != Some(&member) { + return Vec::new(); + } + let mut joined = Vec::new(); + for round in rounds { + let same_learner = |other: &Member| other.learner_id == member.learner_id; + if round.expected_members.iter().any(same_learner) + || round.pushes.keys().any(same_learner) + || round.carried.iter().any(|(other, _)| same_learner(other)) + || floor.get(round.p).is_none_or(|version| *version < round.base_version) + || round_action(round, now) == RoundAction::Complete + { + continue; + } + // A capacity quorum is a share of the expected capacity, which may + // only shrink while a round gathers: a joiner would raise the bar of + // a quorum already reached (and an armed grace) instead of helping. + if round.capacity_quorum.is_some() { + continue; + } + let rebuild_fold = matches!(round.fold, Some(RoundFold::Exact(_))); + if rebuild_fold && !round.pushes.is_empty() { + continue; + } + round.expected_members.push(member); + round.expected_members.sort_unstable(); + if rebuild_fold { + round.fold = round_fold(cfg, st, round.p, &round.expected_members, round.quorum_size); + } + joined.push(JoinedPull { + t: round.t, + p: round.p, + attempt: round.attempt, + base_version: round.base_version, + pull: round.pull.clone(), + }); + } + joined +} + fn route_push( rounds: &mut [Round], member: Member, @@ -5523,6 +5838,18 @@ fn route_push( return PushDisposition::FutureBase; } if max_base_lag.is_some_and(|limit| round.base_version - push.base_version > limit) { + // The island answered this round's pull; it will not push again for + // it. Waiting on its slot would stall the round until the quorum + // timeout (and every later round behind it), so release the slot and + // let the round close with the remaining responders. The refused work + // is not lost: the island's next push is still relative to its raw + // anchor. A fixed roster never merges a partial roster, so it keeps + // the slot; so does a round with an exact fold, whose order and + // weights were fixed for the full expected roster. A streaming fold + // normalises over whatever it folded, so it may close without it. + if !fixed_roster && !matches!(round.fold, Some(RoundFold::Exact(_))) { + round.forget_unanswered(member); + } return PushDisposition::StaleBase; } let duplicate = if fixed_roster { @@ -5536,6 +5863,8 @@ fn route_push( if duplicate { return PushDisposition::Duplicate; } + let mut push = push; + push.release_wire(); round.pushes.insert(member, push); advance_fold(round, member); PushDisposition::Accepted @@ -5662,6 +5991,8 @@ impl LateCarry { Err(disposition) => CarryOutcome::Rejected(disposition), }; } + let mut push = push; + push.compress(); self.queued.entry(p).or_default().push((member, push)); CarryOutcome::Queued } @@ -5673,7 +6004,8 @@ impl LateCarry { }; queued .into_iter() - .map(|(member, push)| { + .map(|(member, mut push)| { + push.expand(); let missed = push.global_step; (member, missed, attach_carry(round, member, push, self.limit)) }) @@ -5744,6 +6076,8 @@ fn attach_carry( if lag > limit { return Err(PushDisposition::StaleBase); } + let mut push = push; + push.release_wire(); round.carried.push((member, push)); Ok(lag) } @@ -6119,9 +6453,19 @@ async fn publish_round( // Policy sweeps persist only before broadcast (see persist_before_broadcast). if cfg.policy_sweep_fragments.is_none() && !checkpoint_before_broadcast(cfg) { if let Some(path) = &cfg.checkpoint_path { - if cfg.checkpoint_every > 0 && st.schedule().round_number(t) % cfg.checkpoint_every == 0 { - st.save_checkpoint(path)?; - info!(step = t, path = %path.display(), "checkpoint written"); + // The last step is persisted by the final checkpoint anyway; a + // periodic one there wrote the same ~1 TB cut twice in a row. + if cfg.checkpoint_every > 0 + && st.schedule().round_number(t) % cfg.checkpoint_every == 0 + && t < st.schedule().last_step(cfg.total_steps) + { + if cfg.background_checkpoint { + st.save_checkpoint_in_background(path)?; + info!(step = t, path = %path.display(), "background checkpoint started"); + } else { + st.save_checkpoint(path)?; + info!(step = t, path = %path.display(), "checkpoint written"); + } } } } @@ -6539,12 +6883,19 @@ async fn admit_recovering_outside_training(registry: &Registry, group: &Arc, registry: &Registry, inflight: &mut [Round], + floor: &[u64], fixed_roster: bool, budget_cutoff: &BudgetCutoff, ) -> bool { @@ -6553,6 +6904,11 @@ async fn repull_rounds_for( } else { rebind_current_unanswered_rounds(registry, inflight, group.member) }; + let joins = if fixed_roster { + Vec::new() + } else { + join_gathering_rounds(cfg, st, registry, inflight, group.member, floor, Instant::now()) + }; for repull in repulls { if !send_small_until_cutoff(group, MSG_PULL_REQ, repull.pull, budget_cutoff).await { return false; @@ -6568,6 +6924,20 @@ async fn repull_rounds_for( "rebound outstanding pull to reconnected learner generation" ); } + for join in joins { + if !send_small_until_cutoff(group, MSG_PULL_REQ, join.pull, budget_cutoff).await { + return false; + } + info!( + learner_id = group.member.learner_id, + generation = group.member.generation, + step = join.t, + fragment = join.p, + attempt = join.attempt, + base_version = join.base_version, + "caught-up learner joined a round launched without it" + ); + } true } @@ -8047,7 +8417,7 @@ async fn finish_catch_up( ); write_pool_status(cfg, registry, st.global_step, "running"); } - if !repull_rounds_for(group, registry, inflight, fixed_roster, budget_cutoff).await { + if !repull_rounds_for(cfg, st, group, registry, inflight, floor, fixed_roster, budget_cutoff).await { return false; } } @@ -8125,7 +8495,7 @@ async fn finish_recovery_replay( ); write_pool_status(cfg, registry, st.global_step, "running"); } - repull_rounds_for(group, registry, inflight, fixed_roster, budget_cutoff).await + repull_rounds_for(cfg, st, group, registry, inflight, &st.versions, fixed_roster, budget_cutoff).await } async fn broadcast_updated_fragment(st: &GlobalState, registry: &Registry, p: usize) -> Result<()> { @@ -8711,6 +9081,52 @@ mod hello_auth_tests { #[cfg(test)] mod tests { use super::*; + + /// `append_tensor_bytes` as it was: every value sliced by the range. + fn reference_tensor_bytes(dtype: u8, values: &[f32], offset: usize, len: usize) -> Vec { + let mut all = Vec::new(); + for value in values { + match dtype { + DTYPE_F32 => all.extend_from_slice(&value.to_le_bytes()), + _ => all.extend_from_slice(&half::bf16::from_f32(*value).to_bits().to_le_bytes()), + } + } + all[offset..offset + len].to_vec() + } + + #[test] + fn tensor_byte_ranges_match_the_value_by_value_encoding() { + let values: Vec = (0..97) + .map(|i| match i % 11 { + 0 => f32::NAN, + 1 => -0.0, + 2 => f32::MAX, + 3 => 1e-40, // subnormal: bf16 rounding + _ => (i as f32 * 0.731).sin() * 3.0e4, + }) + .collect(); + for dtype in [DTYPE_F32, DTYPE_BF16] { + let total = tensor_nbytes(dtype, values.len()).unwrap(); + for offset in 0..12.min(total) { + for len in (0..14).chain([total - offset - 1, total - offset]) { + if offset + len > total { + continue; + } + let mut out = vec![0xAA]; + append_tensor_bytes(dtype, &values, offset, len, &mut out).unwrap(); + assert_eq!(out[0], 0xAA); + assert_eq!(out[1..], reference_tensor_bytes(dtype, &values, offset, len)[..], "{dtype} {offset}+{len}"); + } + } + // Ranges cut inside the tail, too. + for offset in [total - 7, total - 3, total - 1] { + let mut out = Vec::new(); + append_tensor_bytes(dtype, &values, offset, total - offset, &mut out).unwrap(); + assert_eq!(out, reference_tensor_bytes(dtype, &values, offset, total - offset)); + } + assert!(append_tensor_bytes(dtype, &values, 1, total, &mut Vec::new()).is_err()); + } + } use crate::eval_checkpoint::{eval_checkpoint_path, BeforeWrite, EvalCheckpoints}; const CAP: Duration = Duration::from_millis(1000); @@ -9059,6 +9475,7 @@ mod tests { checkpoint_path: None, checkpoint_every: 0, checkpoint_at_steps: Vec::new(), + background_checkpoint: false, async_checkpoint: false, bcast_delta_full_every: None, resume: false, @@ -9079,6 +9496,7 @@ mod tests { weighted_fold: Default::default(), expert_token_weights: false, final_cut_manifest_only: false, + carry_answers_round: false, dispense: None, dispense_log: None, max_subscribers: 0, @@ -9089,6 +9507,7 @@ mod tests { } const FINAL_CUT_MANIFEST_ONLY_VECTOR: &str = "e605792438d8fc745769a6b4334707d749856a1cb0eef84b0b03c9eb5072560f"; + const CARRY_ANSWERS_ROUND_VECTOR: &str = "13a39546746435ebc1878d962f2fb52cf97a5ba66b45757855746adab4e666d7"; const GRACE_FLOOR_VECTOR: &str = "6f9104e302f894575fe2203cfe91a93ff37d16e37b4e75c4367b060d13748d12"; const TERMINAL_GRACE_VECTOR: &str = "6f40c5c1c02e6c63b99bf5aa77fab924fed928a42a9e0d88811583a19fc1aa39"; @@ -9144,6 +9563,16 @@ mod tests { config.semantic_profile_hash().iter().map(|byte| format!("{byte:02x}")).collect(); assert_ne!(config.semantic_profile_hash(), original); assert_eq!(digest, FINAL_CUT_MANIFEST_ONLY_VECTOR); + // Carried deltas answering their round append their own domain; + // Python vector in tests/test_sync_fragment_boundary.py. + config.final_cut_manifest_only = false; + config.carry_answers_round = true; + let digest: String = + config.semantic_profile_hash().iter().map(|byte| format!("{byte:02x}")).collect(); + assert_ne!(config.semantic_profile_hash(), original); + assert_eq!(digest, CARRY_ANSWERS_ROUND_VECTOR); + config.carry_answers_round = false; + config.final_cut_manifest_only = true; // A grace floor appends its own block; off keeps the hash. Python // vector in tests/test_sync_grace_floor.py. config.final_cut_manifest_only = false; @@ -9471,6 +9900,7 @@ mod tests { c_tokens: 99, outer_gradient: vec![1.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }; validate_policy_sweep_push(&push, 3).unwrap(); @@ -9568,6 +9998,7 @@ mod tests { grace_ms: Some(0), fold: None, carried: Vec::new(), + carry_answers: false, pushes: HashMap::from([( member, Push { @@ -9581,6 +10012,7 @@ mod tests { c_tokens: 1, outer_gradient: vec![0.25; 4], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, )]), @@ -9651,6 +10083,7 @@ mod tests { grace_ms: Some(0), fold: None, carried: Vec::new(), + carry_answers: false, pushes: HashMap::from([( member, Push { @@ -9664,6 +10097,7 @@ mod tests { c_tokens, outer_gradient: vec![0.25], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, )]), @@ -9691,6 +10125,7 @@ mod tests { c_tokens: u64::from(member.learner_id + 1) * policy_round * 100, outer_gradient: vec![policy_round as f32 + 2.0 * member.learner_id as f32], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, ) @@ -9714,6 +10149,7 @@ mod tests { grace_ms: Some(0), fold: None, carried: Vec::new(), + carry_answers: false, pushes, } } @@ -10707,6 +11143,91 @@ mod tests { std::fs::remove_dir_all(&directory).ok(); } + #[tokio::test] + async fn background_checkpoint_keeps_rounds_going_and_writes_the_same_cuts() { + let directory = unique_test_directory("background-checkpoint"); + let learner = member(0, 30); + let mut files = Vec::new(); + for background in [false, true] { + let checkpoint = directory.join(format!("state-{background}.ckpt")); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = round_test_config(8); + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 2; + config.background_checkpoint = background; + validate_config(&config).unwrap(); + let mut state = two_fragment_sweep_state(); + state.policy_sweep_fragments = None; + state.save_checkpoint(&checkpoint).unwrap(); // version zero + let mut sync_seconds = 0.0; + for step in 1..=5u64 { + let fragment = ((step - 1) % 2) as usize; + let mut round = sweep_test_round(step, fragment, step, 1, learner); + let base = state.versions[fragment]; + round.base_version = base; + for push in round.pushes.values_mut() { + push.base_version = base; + } + complete_round(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + assert_eq!(chunked_inner_type(&frames.try_recv().unwrap()), MSG_BCAST_FRAGMENT); + state.wait_for_background_checkpoint().unwrap(); + let mut on_disk = two_fragment_sweep_state(); + on_disk.policy_sweep_fragments = None; + on_disk.load_checkpoint(&checkpoint).unwrap(); + assert_eq!(on_disk.global_step, step - step % 2, "step {step} background {background}"); + } + files.push(std::fs::read(&checkpoint).unwrap()); + } + assert_eq!(files[0], files[1]); + // The last step is left to the final checkpoint: no periodic write there. + { + let checkpoint = directory.join("last-step.ckpt"); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = round_test_config(4); + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 2; + let mut state = two_fragment_sweep_state(); + state.policy_sweep_fragments = None; + state.save_checkpoint(&checkpoint).unwrap(); + let mut sync_seconds = 0.0; + for step in 1..=4u64 { + let fragment = ((step - 1) % 2) as usize; + let mut round = sweep_test_round(step, fragment, step, 1, learner); + let base = state.versions[fragment]; + round.base_version = base; + for push in round.pushes.values_mut() { + push.base_version = base; + } + complete_round(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + while frames.try_recv().is_ok() {} + } + let mut on_disk = two_fragment_sweep_state(); + on_disk.policy_sweep_fragments = None; + on_disk.load_checkpoint(&checkpoint).unwrap(); + assert_eq!(on_disk.global_step, 2); + } + // Only for periodic checkpoints with a path, never per round or in a sweep. + for mutate in [ + |config: &mut Config| config.checkpoint_path = None, + |config: &mut Config| config.checkpoint_every = 1, + |config: &mut Config| config.policy_sweep_fragments = Some(2), + ] { + let mut config = round_test_config(8); + config.checkpoint_path = Some(directory.join("unused.ckpt")); + config.checkpoint_every = 2; + config.background_checkpoint = true; + mutate(&mut config); + assert!(format!("{:#}", validate_config(&config).unwrap_err()).contains("--background-checkpoint")); + } + std::fs::remove_dir_all(&directory).ok(); + } + #[tokio::test] async fn async_checkpoint_matches_the_synchronous_path_bit_for_bit() { let directory = unique_test_directory("async-checkpoint-equal"); @@ -11862,6 +12383,36 @@ mod tests { ); } + #[test] + fn queued_late_push_holds_wire_bytes_and_decodes_bit_identically() { + let vals: Vec = (0..1000).map(|i| ((i * 37 % 101) as f32 - 50.0) * 0.013).collect(); + let bytes = crate::protocol::encode_i4(&vals, 64); + let segments = push_segments(DTYPE_F32_PUSH_I4, vals.len(), &[600, 400]).unwrap(); + let wire = PushWire { dtype: DTYPE_F32_PUSH_I4, segments, bytes: Vec::new() }; + let eager = decode_push_segments_to_gradient(&wire, &bytes).unwrap(); + assert!(bytes.len() < 4 * vals.len()); + let mut late = test_push_for(1, 5); + late.global_step = 7; + late.outer_gradient = eager.clone(); + late.wire = Some(PushWire { bytes, ..wire }); + let mut carry = LateCarry::new(Some(2), false).unwrap(); + let island = member(1, 10); + let mut rounds: Vec = Vec::new(); + assert_eq!(carry.carry(&mut rounds, island, late, 8, 7), CarryOutcome::Queued); + for queued in carry.queued.values().flatten() { + assert!(queued.1.outer_gradient.is_empty(), "queued push keeps only wire bytes"); + assert!(queued.1.wire.is_some()); + } + let mut next = test_round(vec![member(0, 10), island]); + next.t = 9; + next.base_version = 7; + assert_eq!(carry.attach_queued(&mut next).len(), 1); + let got = &next.carried[0]; + let got_bits: Vec = got.1.outer_gradient.iter().map(|v| v.to_bits()).collect(); + let want_bits: Vec = eager.iter().map(|v| v.to_bits()).collect(); + assert_eq!(got_bits, want_bits); + } + #[test] fn late_push_joins_an_already_open_later_round() { let mut carry = LateCarry::new(Some(2), false).unwrap(); @@ -11934,6 +12485,71 @@ mod tests { assert_eq!(carry.carry(&mut Vec::new(), other, late, 8, 7), CarryOutcome::Queued); } + #[test] + fn carried_delta_counts_towards_quorum_and_responders() { + // Regression: at quorum 1 a round whose only delta from the fast + // island arrived as a late carry (anchor spent) waited for the + // stalled island instead of closing after grace. + let fast = member(0, 1); + let slow = member(1, 1); + let mut round = test_round(vec![fast, slow]); + round.carry_answers = true; + round.quorum_size = 1; + round.quorum_deadline = Instant::now() + Duration::from_secs(3600); + assert!(!round.quorum_reached()); + assert_eq!(attach_carry(&mut round, fast, test_push_for(0, 4), 1), Ok(1)); + assert!(round.pushes.is_empty()); + assert!(round.quorum_reached()); + // Grace still applies, as for an own push. + round.grace_deadline = Some(Instant::now() + Duration::from_secs(3600)); + assert_eq!(round_action(&round, Instant::now()), RoundAction::Wait); + round.grace_deadline = Some(Instant::now() - Duration::from_millis(1)); + assert_eq!(round_action(&round, Instant::now()), RoundAction::Complete); + } + + #[test] + fn carried_delta_stays_outside_quorum_when_the_flag_is_off() { + let fast = member(0, 1); + let mut round = test_round(vec![fast, member(1, 1)]); + round.quorum_size = 1; + round.quorum_deadline = Instant::now() + Duration::from_secs(3600); + assert_eq!(attach_carry(&mut round, fast, test_push_for(0, 4), 1), Ok(1)); + assert!(!round.quorum_reached()); + assert_eq!(round_action(&round, Instant::now()), RoundAction::Wait); + } + + #[test] + fn carried_delta_completes_a_round_once_every_expected_island_answered() { + let fast = member(0, 1); + let slow = member(1, 1); + let mut round = test_round(vec![fast, slow]); + round.carry_answers = true; + round.quorum_size = 2; + round.quorum_deadline = Instant::now() + Duration::from_secs(3600); + assert_eq!(attach_carry(&mut round, fast, test_push_for(0, 4), 1), Ok(1)); + assert_eq!(round_action(&round, Instant::now()), RoundAction::Wait); + assert_eq!( + route_push(std::slice::from_mut(&mut round), slow, test_push_for(1, 5), Some(1), false), + PushDisposition::Accepted + ); + assert_eq!(round_action(&round, Instant::now()), RoundAction::Complete); + } + + #[test] + fn carried_delta_of_an_unexpected_member_does_not_count() { + let expected = member(0, 1); + let replaced = member(0, 0); + let mut round = test_round(vec![expected]); + round.carry_answers = true; + round.quorum_size = 1; + assert_eq!(attach_carry(&mut round, replaced, test_push_for(0, 4), 1), Ok(1)); + assert!(!round.quorum_reached()); + // A carried push from the same island counts once. + assert_eq!(attach_carry(&mut round, expected, test_push_for(0, 4), 1), Ok(1)); + assert_eq!(attach_carry(&mut round, expected, test_push_for(0, 5), 1), Ok(0)); + assert_eq!(round.responders().len(), 1); + } + fn carried_merge(cfg: &Config) -> f32 { let mut st = two_fragment_streaming_state(); st.versions[0] = 4; @@ -12043,6 +12659,7 @@ mod tests { grace_ms: None, fold: None, carried: Vec::new(), + carry_answers: false, pushes: HashMap::new(), } } @@ -12059,6 +12676,7 @@ mod tests { c_tokens: 20, outer_gradient: vec![1.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), } } @@ -12376,6 +12994,294 @@ mod tests { assert_eq!(preserved_push.local_step, 77); } + /// A round launched while a learner was still in its catch-up: the + /// learner joins it once the catch-up is done, exactly once per round. + #[test] + fn caught_up_learner_joins_a_round_launched_without_it_once() { + let first = member(0, 10); + let late = member(1, 20); + let registry = registry_with_current(late); + let cfg = round_test_config(4); + let st = resend_test_state(2); + let mut round = test_round(vec![first]); + round.base_version = 0; + round.pull = encode_pull(1, 7, 1); + let quorum_deadline = round.quorum_deadline; + let mut rounds = vec![round]; + + let joined = join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, late, &[0, 0], Instant::now()); + assert_eq!( + joined, + vec![JoinedPull { t: 7, p: 1, attempt: 1, base_version: 0, pull: encode_pull(1, 7, 1) }] + ); + assert_eq!(rounds[0].expected_members, vec![first, late]); + assert_eq!((rounds[0].attempt, rounds[0].quorum_size), (1, 1)); + assert_eq!(rounds[0].quorum_deadline, quorum_deadline); + // The tail runs again (a second catch-up of the same generation, a + // recovery replay): no second pull, no second slot. + assert!(join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, late, &[0, 0], Instant::now()).is_empty()); + assert!(rebind_current_unanswered_rounds(®istry, &mut rounds, late).is_empty()); + assert_eq!(rounds[0].expected_members, vec![first, late]); + // The joiner's push is now accepted for this attempt. + let mut push = test_push_for(1, 0); + push.base_version = 0; + assert_eq!(route_push(&mut rounds, late, push, Some(64), false), PushDisposition::Accepted); + // A generation that is not current never joins. + let mut fresh = vec![test_round(vec![first])]; + fresh[0].base_version = 0; + assert!(join_gathering_rounds(&cfg, &st, ®istry, &mut fresh, member(2, 1), &[0, 0], Instant::now()) + .is_empty()); + } + + /// A capacity quorum's threshold is a share of the expected capacity: + /// a joiner would raise it under a round that already reached it. + #[test] + fn caught_up_learner_does_not_join_a_capacity_quorum_round() { + let first = member(0, 10); + let late = member(1, 20); + let registry = registry_with_current(late); + let mut cfg = round_test_config(4); + cfg.membership.quorum_rule = crate::membership::QuorumRule::Capacity { theta: 0.5, q_min: 1 }; + let st = resend_test_state(2); + let mut round = test_round(vec![first]); + round.base_version = 0; + round.capacity_quorum = cfg.membership.round_quorum(); + assert!(round.capacity_quorum.is_some()); + let mut rounds = vec![round]; + assert!(join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, late, &[0, 0], Instant::now()).is_empty()); + assert_eq!(rounds[0].expected_members, vec![first]); + } + + /// A reconnect of a learner the round already expects is the rebind's + /// job: the new generation is pulled once, by the rebind, and only after + /// the captured generation is gone. + #[test] + fn reconnected_generation_is_rebound_not_joined() { + let captured = member(1, 20); + let replacement = member(1, 21); + let registry = registry_with_current(replacement); + registry.lock().unwrap().groups.insert(captured, test_group(captured)); + let cfg = round_test_config(4); + let st = resend_test_state(2); + let mut round = test_round(vec![member(0, 10), captured]); + round.base_version = 0; + let mut rounds = vec![round]; + + // The captured generation is still connected: neither path pulls. + assert!(rebind_current_unanswered_rounds(®istry, &mut rounds, replacement).is_empty()); + assert!(join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, replacement, &[0, 0], Instant::now()) + .is_empty()); + assert_eq!(rounds[0].expected_members, vec![member(0, 10), captured]); + // Once it is gone, the rebind moves the slot and pulls once; the + // join then sees the learner expected and adds nothing. + registry.lock().unwrap().groups.remove(&captured); + assert_eq!(rebind_current_unanswered_rounds(®istry, &mut rounds, replacement).len(), 1); + assert!(join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, replacement, &[0, 0], Instant::now()) + .is_empty()); + assert_eq!(rounds[0].expected_members, vec![member(0, 10), replacement]); + } + + /// A round in its grace window takes a joiner's push until the grace + /// deadline, which the join does not move; a completed round takes no + /// joiner. + #[test] + fn joiner_counts_in_an_armed_grace_window_until_its_deadline() { + let answered = member(0, 10); + let silent = member(2, 30); + let late = member(1, 20); + let registry = registry_with_current(late); + let cfg = round_test_config(4); + let st = resend_test_state(2); + let mut round = test_round(vec![answered, silent]); + round.base_version = 0; + let mut first_push = test_push_for(0, 0); + first_push.base_version = 0; + round.pushes.insert(answered, first_push); + assert!(round.quorum_reached()); + let now = Instant::now(); + let grace_deadline = now + Duration::from_secs(60); + round.grace_deadline = Some(grace_deadline); + let mut rounds = vec![round]; + + assert_eq!(join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, late, &[0, 0], now).len(), 1); + assert_eq!(rounds[0].grace_deadline, Some(grace_deadline)); + assert_eq!(round_action(&rounds[0], now), RoundAction::Wait); + let mut push = test_push_for(1, 0); + push.base_version = 0; + assert_eq!(route_push(&mut rounds, late, push, Some(64), false), PushDisposition::Accepted); + assert_eq!(rounds[0].pushes.len(), 2); + // Still waiting for the silent one until the unchanged deadline. + assert_eq!(round_action(&rounds[0], now), RoundAction::Wait); + assert_eq!(round_action(&rounds[0], grace_deadline), RoundAction::Complete); + + // Past its grace deadline the round is complete: no join. + let mut closed = test_round(vec![answered, silent]); + closed.base_version = 0; + closed.pushes.insert(answered, test_push_for(0, 0)); + closed.grace_deadline = Some(now); + let mut closed = vec![closed]; + assert!(join_gathering_rounds(&cfg, &st, ®istry, &mut closed, late, &[0, 0], now).is_empty()); + assert_eq!(closed[0].expected_members, vec![answered, silent]); + } + + /// No join where the learner already counts (a carried delta), where its + /// catch-up did not reach the round's base, or where an exact fold has + /// already taken pushes; an empty exact fold is rebuilt for the roster. + #[test] + fn join_skips_carried_stale_and_started_exact_rounds() { + let first = member(0, 10); + let late = member(1, 20); + let registry = registry_with_current(late); + let mut cfg = round_test_config(4); + cfg.max_base_lag = Some(64); + let st = resend_test_state(2); + + let mut carried = test_round(vec![first]); + carried.base_version = 0; + carried.carried.push((member(1, 19), test_push_for(1, 0))); + let mut behind = test_round(vec![first]); + behind.base_version = 3; + let mut rounds = vec![carried, behind]; + assert!(join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, late, &[2, 2], Instant::now()).is_empty()); + assert!(rounds.iter().all(|round| round.expected_members == vec![first])); + + // Exact fold of the launch roster [0] with quorum 2: rebuilt for [0, 1]. + let exact = |expected: &[Member]| round_fold(&cfg, &st, 1, expected, 2); + assert!(matches!(exact(&[first]), Some(RoundFold::Exact(_)))); + let mut empty = test_round(vec![first]); + empty.base_version = 0; + empty.quorum_size = 2; + empty.fold = exact(&[first]); + let mut rounds = vec![empty]; + assert_eq!(join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, late, &[0, 0], Instant::now()).len(), 1); + let Some(RoundFold::Exact(fold)) = &rounds[0].fold else { panic!("exact fold expected") }; + assert_eq!(fold.next_learner(), Some(0)); + let mut push = test_push_for(0, 0); + push.outer_gradient = vec![1.0; 4]; + assert_eq!(route_push(&mut rounds, first, push, Some(64), false), PushDisposition::Accepted); + let Some(RoundFold::Exact(fold)) = &rounds[0].fold else { panic!("exact fold expected") }; + assert_eq!(fold.next_learner(), Some(1), "the rebuilt roster waits for the joiner"); + + // An exact fold that already holds a push keeps its roster. + let mut started = test_round(vec![first, member(2, 30)]); + started.base_version = 0; + started.quorum_size = 2; + started.fold = exact(&[first, member(2, 30)]); + let mut push = test_push_for(0, 0); + push.outer_gradient = vec![1.0; 4]; + let mut rounds = vec![started]; + assert_eq!(route_push(&mut rounds, first, push, Some(64), false), PushDisposition::Accepted); + assert!(join_gathering_rounds(&cfg, &st, ®istry, &mut rounds, late, &[0, 0], Instant::now()).is_empty()); + assert_eq!(rounds[0].expected_members, vec![first, member(2, 30)]); + } + + fn scheduler_test_group(member: Member, layout: &Layout, control: mpsc::Sender) -> Arc { + let group = Arc::try_unwrap(test_group_with_control(member, control)).ok().unwrap(); + Arc::new(Group { layout: layout.clone(), ..group }) + } + + /// (fragment, step, attempt) of the next PULL on a learner's control + /// stream; other frames (the catch-up) are skipped. + async fn next_pull(frames: &mut mpsc::Receiver) -> (u32, u64, u32) { + loop { + let frame = frames.recv().await.expect("control stream closed"); + if frame.msg_type == MSG_PULL_REQ { + let payload = &frame.parts[0]; + return ( + u32::from_le_bytes(payload[0..4].try_into().unwrap()), + u64::from_le_bytes(payload[4..12].try_into().unwrap()), + u32::from_le_bytes(payload[12..16].try_into().unwrap()), + ); + } + } + } + + /// The resumed-fleet stall: the first learner done with its catch-up is + /// the only one a round is launched to, and its H-gate keeps it from + /// answering. The learner that finishes its catch-up next is pulled into + /// that round and its push commits it well before the quorum timeout + /// (before the fix it was pulled only when the timeout relaunched it). + #[tokio::test] + async fn round_commits_from_a_learner_that_finished_its_catch_up_after_the_launch() { + let layout = Layout { + fragments: (0..2) + .map(|_| crate::state::FragmentInfo { + merge_mode: crate::state::MERGE_AVG, + tensor_numels: vec![4], + tensor_shapes: None, + }) + .collect(), + }; + let directory = unique_test_directory("late-joiner-pull"); + let tape = directory.join("tape.jsonl"); + let mut cfg = round_test_config(4); + cfg.learners = 1; + cfg.quorum = 1; + cfg.max_base_lag = Some(64); + cfg.carry_answers_round = true; + cfg.quorum_timeout_s = 3600; + cfg.event_tape = Some(tape.clone()); + cfg.membership = crate::membership::MembershipConfig { max_learners: Some(2), ..Default::default() }; + let registry: Registry = + Arc::new(Mutex::new(RegistryState::new(Membership::new(cfg.membership.clone(), 2)))); + let budget_cutoff = Arc::new(BudgetCutoff::new(false, Duration::from_secs(cfg.quorum_timeout_s))); + let (events, events_rx) = mpsc::channel(64); + let (first, late) = (member(0, 1), member(1, 1)); + let (first_control, mut first_frames) = mpsc::channel(1024); + let (late_control, mut late_frames) = mpsc::channel(1024); + let first_group = scheduler_test_group(first, &layout, first_control); + let late_group = scheduler_test_group(late, &layout, late_control); + + registry.lock().unwrap().register_group(first_group.clone(), Instant::now()).unwrap(); + events.send(Event::Hello { group: first_group }).await.unwrap(); + for fragment_id in 0..2 { + let init = Event::Init { member: first, fragment_id, values: vec![0.0; 4], sha256: [0; 32] }; + events.send(init).await.unwrap(); + } + let driver = async { + // Learner 0 alone gets the pull of step 1 and never answers it. + let pull = next_pull(&mut first_frames).await; + assert_eq!((pull.1, pull.2), (1, 1)); + registry.lock().unwrap().register_group(late_group.clone(), Instant::now()).unwrap(); + events.send(Event::Hello { group: late_group }).await.unwrap(); + let joined = tokio::time::timeout(Duration::from_secs(10), next_pull(&mut late_frames)) + .await + .expect("the learner that caught up after the launch was not pulled into the round"); + assert_eq!(joined, pull); + let mut push = test_push_for(1, 0); + push.fragment_id = pull.0; + push.global_step = 1; + push.round_attempt = 1; + push.outer_gradient = vec![0.5; 4]; + events.send(Event::Push { member: late, push }).await.unwrap(); + let committed = async { + loop { + let text = std::fs::read_to_string(&tape).unwrap_or_default(); + if let Some(line) = text.lines().find(|line| line.contains("\"sync/responders\"") && line.contains("\"step\":1,")) { + return line.to_string(); + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + }; + let record = tokio::time::timeout(Duration::from_secs(10), committed) + .await + .expect("step 1 did not commit before the quorum timeout"); + for field in [ + "\"attempt\":1,", + "\"expected\":[0,1],", + "\"responded\":[1],", + "\"missed_members\":[{\"id\":0,\"generation\":1}],", + ] { + assert!(record.contains(field), "{field} missing from {record}"); + } + }; + tokio::select! { + result = scheduler(cfg, events_rx, registry.clone(), budget_cutoff) => panic!("scheduler ended: {result:?}"), + () = driver => {} + } + std::fs::remove_dir_all(&directory).unwrap(); + } + #[test] fn round_membership_rejects_mid_round_join_and_new_generation() { let captured = member(0, 10); @@ -12473,6 +13379,7 @@ mod tests { c_tokens: 1, outer_gradient: vec![1.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), } } @@ -12508,6 +13415,52 @@ mod tests { assert_eq!(rounds[0].pushes.len(), 1); } + #[test] + fn stale_base_rejection_releases_the_slot_so_the_round_completes() { + // Regression: with --max-base-lag 8 and quorum == learners, a push + // refused as StaleBase left its island in expected_members, so the + // round (and every later step) waited for the full quorum timeout. + let fast = member(0, 10); + let slow = member(1, 20); + let mut round = test_round(vec![fast, slow]); + round.base_version = 15; + round.quorum_size = 2; + round.quorum_deadline = Instant::now() + Duration::from_secs(3600); + let mut rounds = vec![round]; + assert_eq!( + route_push(&mut rounds, fast, test_push_for(0, 15), Some(8), false), + PushDisposition::Accepted + ); + assert_eq!(round_action(&rounds[0], Instant::now()), RoundAction::Wait); + assert_eq!( + route_push(&mut rounds, slow, test_push_for(1, 6), Some(8), false), + PushDisposition::StaleBase + ); + assert_eq!(rounds[0].expected_members, vec![fast]); + assert_eq!(round_action(&rounds[0], Instant::now()), RoundAction::Complete); + } + + #[test] + fn stale_base_rejection_first_still_waits_for_the_other_island() { + let fast = member(0, 10); + let slow = member(1, 20); + let mut round = test_round(vec![fast, slow]); + round.base_version = 15; + round.quorum_size = 2; + round.quorum_deadline = Instant::now() + Duration::from_secs(3600); + let mut rounds = vec![round]; + assert_eq!( + route_push(&mut rounds, slow, test_push_for(1, 6), Some(8), false), + PushDisposition::StaleBase + ); + assert_eq!(round_action(&rounds[0], Instant::now()), RoundAction::Wait); + assert_eq!( + route_push(&mut rounds, fast, test_push_for(0, 15), Some(8), false), + PushDisposition::Accepted + ); + assert_eq!(round_action(&rounds[0], Instant::now()), RoundAction::Complete); + } + #[test] fn quorum_timeout_with_one_of_two_responses_restarts_without_merging() { let first = member(0, 10); @@ -12554,6 +13507,7 @@ mod tests { c_tokens: 40, outer_gradient: vec![1.0], tensor_scales: None, + wire: None, transfer: PushTransfer { recv_seconds: 1.5, arrived_unix: 1792000000.25, @@ -12731,6 +13685,7 @@ mod tests { c_tokens: 10, outer_gradient: vec![1.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, ), @@ -12747,6 +13702,7 @@ mod tests { c_tokens: 1_000, outer_gradient: vec![2.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, ), @@ -13064,6 +14020,7 @@ mod tests { c_tokens: 1 + u64::from(learner_id), outer_gradient: gradient, tensor_scales: None, + wire: None, transfer: PushTransfer::default(), } } @@ -13345,6 +14302,50 @@ mod tests { st.params[0].to_vec() } + #[test] + fn stale_base_rejection_releases_the_slot_of_a_streaming_fold_round() { + // Async carry needs a streaming fold, and that is where stale bases + // happen: the refused island must not hold the round until timeout. + let cfg = streaming_cfg(); + let mut st = fold_test_state(crate::state::MERGE_AVG); + st.versions[0] = 15; + let (fast, slow) = (member(0, 10), member(1, 20)); + let mut round = test_round(vec![fast, slow]); + round.t = 1; + round.p = 0; + round.base_version = 15; + round.quorum_size = 2; + round.quorum_deadline = Instant::now() + Duration::from_secs(3600); + round.fold = round_fold(&cfg, &st, 0, &[fast, slow], 2); + assert!(matches!(round.fold, Some(RoundFold::Streaming(_)))); + let mut rounds = vec![round]; + let push = |learner, base| Push { base_version: base, ..fold_test_push(learner) }; + assert_eq!(route_push(&mut rounds, fast, push(0, 15), Some(8), false), PushDisposition::Accepted); + assert_eq!(route_push(&mut rounds, slow, push(1, 6), Some(8), false), PushDisposition::StaleBase); + assert_eq!(rounds[0].expected_members, vec![fast]); + assert_eq!(round_action(&rounds[0], Instant::now()), RoundAction::Complete); + prepare_round_compute(&cfg, &st, &mut rounds[0]).unwrap().finish_inline().unwrap(); + } + + #[test] + fn stale_base_rejection_keeps_the_slot_of_an_exact_fold_round() { + let mut cfg = round_test_config(1); + cfg.weighted_fold = WeightedFoldMode::Streaming; + let st = fold_test_state(crate::state::MERGE_AVG); + let (fast, slow) = (member(0, 10), member(1, 20)); + let mut round = test_round(vec![fast, slow]); + round.t = 1; + round.p = 0; + round.base_version = 15; + round.quorum_size = 2; + round.fold = round_fold(&cfg, &st, 0, &[fast, slow], 2); + assert!(matches!(round.fold, Some(RoundFold::Exact(_)))); + let mut rounds = vec![round]; + let push = |learner, base| Push { base_version: base, ..fold_test_push(learner) }; + assert_eq!(route_push(&mut rounds, slow, push(1, 6), Some(8), false), PushDisposition::StaleBase); + assert_eq!(rounds[0].expected_members, vec![fast, slow]); + } + #[test] fn streaming_partial_quorum_round_commits_within_rounding_of_buffered() { let streaming = streaming_cfg(); diff --git a/syncer/src/sha256.rs b/syncer/src/sha256.rs new file mode 100644 index 00000000..262070a0 --- /dev/null +++ b/syncer/src/sha256.rs @@ -0,0 +1,229 @@ +//! SHA-256 with the ARMv8 SHA-2 instructions when the CPU has them. +//! +//! The vendored `sha2` crate only uses those instructions behind its `asm` +//! feature, which pulls an extra build-time crate; without it every digest +//! runs the portable code (about 0.25 GB/s on a Kunpeng 920 versus +//! 1.6 GB/s with the instructions). The syncer hashes every INIT_PARAMS +//! payload on its read loop, so a 120B model's 483 GB of f32 INIT took over +//! half an hour to admit. Digests are byte-identical to `sha2::Sha256`; other +//! targets and CPUs without the extension use `sha2` itself. + +use sha2::{Digest, Sha256}; + +#[cfg(target_arch = "aarch64")] +mod hw { + use std::arch::aarch64::*; + + pub const K: [u32; 64] = [ + 0x428a2f98, 0x71374491, 0xb5c0fbcf, 0xe9b5dba5, 0x3956c25b, 0x59f111f1, 0x923f82a4, + 0xab1c5ed5, 0xd807aa98, 0x12835b01, 0x243185be, 0x550c7dc3, 0x72be5d74, 0x80deb1fe, + 0x9bdc06a7, 0xc19bf174, 0xe49b69c1, 0xefbe4786, 0x0fc19dc6, 0x240ca1cc, 0x2de92c6f, + 0x4a7484aa, 0x5cb0a9dc, 0x76f988da, 0x983e5152, 0xa831c66d, 0xb00327c8, 0xbf597fc7, + 0xc6e00bf3, 0xd5a79147, 0x06ca6351, 0x14292967, 0x27b70a85, 0x2e1b2138, 0x4d2c6dfc, + 0x53380d13, 0x650a7354, 0x766a0abb, 0x81c2c92e, 0x92722c85, 0xa2bfe8a1, 0xa81a664b, + 0xc24b8b70, 0xc76c51a3, 0xd192e819, 0xd6990624, 0xf40e3585, 0x106aa070, 0x19a4c116, + 0x1e376c08, 0x2748774c, 0x34b0bcb5, 0x391c0cb3, 0x4ed8aa4a, 0x5b9cca4f, 0x682e6ff3, + 0x748f82ee, 0x78a5636f, 0x84c87814, 0x8cc70208, 0x90befffa, 0xa4506ceb, 0xbef9a3f7, + 0xc67178f2, + ]; + + pub fn available() -> bool { + std::arch::is_aarch64_feature_detected!("sha2") + } + + /// Compress whole 64-byte blocks into `state`. + /// + /// # Safety + /// The CPU must have the SHA-2 extension (`available()`), and + /// `blocks.len()` must be a multiple of 64. + #[target_feature(enable = "sha2")] + pub unsafe fn compress(state: &mut [u32; 8], blocks: &[u8]) { + let mut abcd = vld1q_u32(state.as_ptr()); + let mut efgh = vld1q_u32(state.as_ptr().add(4)); + for block in blocks.chunks_exact(64) { + let (abcd0, efgh0) = (abcd, efgh); + let p = block.as_ptr(); + let mut w = [ + vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(p))), + vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(p.add(16)))), + vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(p.add(32)))), + vreinterpretq_u32_u8(vrev32q_u8(vld1q_u8(p.add(48)))), + ]; + for i in 0..16 { + let wk = vaddq_u32(w[i % 4], vld1q_u32(K.as_ptr().add(4 * i))); + let prev = abcd; + abcd = vsha256hq_u32(abcd, efgh, wk); + efgh = vsha256h2q_u32(efgh, prev, wk); + if i < 12 { + // W[4(i+4)..] from W[4i..4i+16), into the slot just consumed. + w[i % 4] = vsha256su1q_u32( + vsha256su0q_u32(w[i % 4], w[(i + 1) % 4]), + w[(i + 2) % 4], + w[(i + 3) % 4], + ); + } + } + abcd = vaddq_u32(abcd, abcd0); + efgh = vaddq_u32(efgh, efgh0); + } + vst1q_u32(state.as_mut_ptr(), abcd); + vst1q_u32(state.as_mut_ptr().add(4), efgh); + } +} + +const H0: [u32; 8] = [ + 0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c, 0x1f83d9ab, 0x5be0cd19, +]; + +/// Incremental SHA-256, the same bytes as `sha2::Sha256`. +pub enum Hasher { + #[cfg(target_arch = "aarch64")] + Hw { state: [u32; 8], buf: [u8; 64], used: usize, len: u64 }, + Soft(Sha256), +} + +impl Default for Hasher { + fn default() -> Self { + Self::new() + } +} + +impl Hasher { + pub fn new() -> Self { + #[cfg(target_arch = "aarch64")] + if hw::available() { + return Self::Hw { state: H0, buf: [0; 64], used: 0, len: 0 }; + } + Self::Soft(Sha256::new()) + } + + /// The portable implementation, whatever the CPU (tests compare to it). + #[cfg(test)] + fn soft() -> Self { + Self::Soft(Sha256::new()) + } + + pub fn update(&mut self, mut data: &[u8]) { + match self { + #[cfg(target_arch = "aarch64")] + Self::Hw { state, buf, used, len } => { + // Only `new` builds Hw, after the feature check (cached). + debug_assert!(hw::available()); + *len += data.len() as u64; + if *used > 0 { + let take = (64 - *used).min(data.len()); + buf[*used..*used + take].copy_from_slice(&data[..take]); + *used += take; + data = &data[take..]; + if *used < 64 { + return; + } + // SAFETY: the Hw variant exists only when hw::available(). + unsafe { hw::compress(state, &buf[..]) }; + *used = 0; + } + let whole = data.len() - data.len() % 64; + // SAFETY: as above; `whole` is a multiple of 64. + unsafe { hw::compress(state, &data[..whole]) }; + let rest = &data[whole..]; + buf[..rest.len()].copy_from_slice(rest); + *used = rest.len(); + } + Self::Soft(h) => h.update(data), + } + } + + pub fn finalize(self) -> [u8; 32] { + match self { + #[cfg(target_arch = "aarch64")] + Self::Hw { mut state, mut buf, used, len } => { + debug_assert!(hw::available()); + buf[used] = 0x80; + buf[used + 1..].fill(0); + if used >= 56 { + // SAFETY: the Hw variant exists only when hw::available(). + unsafe { hw::compress(&mut state, &buf[..]) }; + buf.fill(0); + } + buf[56..].copy_from_slice(&(len.wrapping_mul(8)).to_be_bytes()); + // SAFETY: as above. + unsafe { hw::compress(&mut state, &buf[..]) }; + let mut out = [0u8; 32]; + for (chunk, word) in out.chunks_exact_mut(4).zip(state) { + chunk.copy_from_slice(&word.to_be_bytes()); + } + out + } + Self::Soft(h) => h.finalize().into(), + } + } +} + +/// SHA-256 of `data`. +pub fn digest(data: &[u8]) -> [u8; 32] { + let mut h = Hasher::new(); + h.update(data); + h.finalize() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn bytes(n: usize, seed: u64) -> Vec { + let mut x = seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + (0..n) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + x as u8 + }) + .collect() + } + + fn reference(data: &[u8]) -> [u8; 32] { + Sha256::digest(data).into() + } + + #[test] + fn digest_matches_sha2_at_every_padding_boundary() { + for n in (0..=260).chain([511, 512, 513, 4096, 65_535, 65_536, 1 << 20, (1 << 20) + 63]) { + let data = bytes(n, n as u64); + assert_eq!(digest(&data), reference(&data), "length {n}"); + } + } + + #[test] + fn known_vectors() { + let hex = |d: [u8; 32]| d.iter().map(|b| format!("{b:02x}")).collect::(); + assert_eq!(hex(digest(b"")), "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"); + assert_eq!(hex(digest(b"abc")), "ba7816bf8f01cfea414140de5dae2223b00361a396177a9cb410ff61f20015ad"); + assert_eq!( + hex(digest(b"abcdbcdecdefdefgefghfghighijhijkijkljklmklmnlmnomnopnopq")), + "248d6a61d20638b8e5c026930c3e6039a33ce45964ff2167f6ecedd419db06c1" + ); + } + + #[test] + fn split_updates_match_one_update() { + let data = bytes(10_000, 7); + for split in [1usize, 3, 63, 64, 65, 127, 1000] { + let mut h = Hasher::new(); + for chunk in data.chunks(split) { + h.update(chunk); + } + h.update(&[]); + assert_eq!(h.finalize(), reference(&data), "chunks of {split}"); + } + let mut soft = Hasher::soft(); + soft.update(&data); + assert_eq!(soft.finalize(), reference(&data)); + } + + #[cfg(target_arch = "aarch64")] + #[test] + fn the_hardware_path_is_taken_when_the_cpu_has_it() { + assert_eq!(matches!(Hasher::new(), Hasher::Hw { .. }), hw::available()); + } +} diff --git a/syncer/src/state.rs b/syncer/src/state.rs index e272b8ee..2db0ead4 100644 --- a/syncer/src/state.rs +++ b/syncer/src/state.rs @@ -154,6 +154,10 @@ pub struct GlobalState { /// HeLoCo per-tensor directional correction of learner deltas against /// the outer momentum before merging (None disables). pub delta_correction: Option, + /// Threads used by the merge and the outer optimizer step. Results are + /// bit-identical for every value (see `crate::par`); 1 runs the + /// historical single-threaded code. + pub merge_threads: usize, /// Spectrum flattener. Torch mode owns one persistent child process. iso_backend: IsoBackend, /// bf16 downlink anchors; in memory only, never checkpointed. @@ -165,6 +169,15 @@ pub struct GlobalState { /// Small records written next to every checkpoint at the same cut /// (the task dispenser). Empty: none. pub checkpoint_sidecars: Vec>, + /// The periodic checkpoint written off the scheduler + /// (`save_checkpoint_in_background`), until a later save waits for it. + background_checkpoint: std::sync::Mutex>, +} + +/// A checkpoint write on its own thread. +struct BackgroundCheckpoint { + step: u64, + writer: std::thread::JoinHandle>, } /// A small record captured at every checkpoint cut and written next to the @@ -369,10 +382,12 @@ impl GlobalState { global_step: 0, ledger: Default::default(), checkpoint_sidecars: Vec::new(), + background_checkpoint: std::sync::Mutex::new(None), outer_lr, outer_momentum, wire_dtype, delta_correction: None, + merge_threads: crate::par::default_threads(), iso_backend: IsoBackend::start(iso_backend_config)?, downlink: None, shard: ShardPlan::default(), @@ -544,8 +559,20 @@ impl GlobalState { // current global trajectory). Inputs already have outer-gradient // sign, so correction operates on copies without reconstructing // learner parameters or consulting a parameter anchor. + let threads = self.merge_threads; + let parallel = threads > 1 && numel >= PARALLEL_MIN_NUMEL; + let tensor_numels: Vec = frag.tensor_numels.iter().map(|&n| n as usize).collect(); let corrected: Vec>; - let outer_gradients: Vec<&[f32]> = if let Some(h) = self.delta_correction { + let outer_gradients: Vec<&[f32]> = if let (Some(h), true) = (self.delta_correction, parallel) { + corrected = heloco_correct_all_par( + outer_gradients, + &self.momentum[fid], + &tensor_numels, + &h, + threads, + ); + corrected.iter().map(|v| v.as_slice()).collect() + } else if let Some(h) = self.delta_correction { let momentum = &self.momentum[fid]; corrected = outer_gradients .iter() @@ -576,6 +603,18 @@ impl GlobalState { bail!("fragment {fid}: tensor scales do not match the learners and tensors"); } } + if parallel { + let delta = merge_tensors_par( + fid, + frag.merge_mode, + outer_gradients, + weights, + tensor_scales, + &tensor_numels, + threads, + )?; + return self.prepare_reduced_merge(fid, base_version, delta); + } let mut tensor_weights = weights.to_vec(); let mut delta = vec![0.0f32; numel]; // Merge per tensor slice within the fragment. @@ -673,7 +712,7 @@ impl GlobalState { { iso_jobs[0].matrix = std::mem::take(&mut delta); } - if delta.iter().any(|value| !value.is_finite()) { + if crate::par::any_non_finite(self.merge_threads, &delta) { bail!("fragment {fid}: merged outer gradient is non-finite"); } Ok(PreparedMerge { @@ -709,6 +748,9 @@ impl GlobalState { delta.len() ); } + if self.merge_threads > 1 && expected >= PARALLEL_MIN_NUMEL { + return self.commit_merge_par(fid, &delta); + } let gnorm = delta .iter() .map(|v| (*v as f64).powi(2)) @@ -746,6 +788,89 @@ impl GlobalState { Ok(gnorm) } + /// Multi-threaded `commit_merge` with the same checks, error priority and + /// bits. ‖delta‖ stays one sequential reduction; it runs on its own + /// thread, overlapped with the chunked validation and optimizer step. + /// For finite f32 inputs that f64 sum of squares cannot overflow + /// (it would need more than 1e230 values), so a finite delta implies a + /// finite norm and the state may be mutated before the norm is joined. + fn commit_merge_par(&mut self, fid: usize, delta: &[f32]) -> Result { + let threads = self.merge_threads; + if crate::par::any_non_finite(threads, delta) { + bail!("fragment {fid}: merged outer gradient is non-finite"); + } + let (lr, mu) = (self.outer_lr, self.outer_momentum); + let period = self.delayed_nesterov_period; + let params = &mut self.params[fid]; + let momentum = &mut self.momentum[fid]; + let delayed = &mut self.delayed_nesterov; + std::thread::scope(|scope| -> Result { + let gnorm = scope.spawn(|| delta.iter().map(|v| (*v as f64).powi(2)).sum::().sqrt()); + let workers = threads.saturating_sub(1).max(1); + let c = crate::par::CHUNK; + let bad = std::sync::atomic::AtomicBool::new(false); + let flag = |b: bool| { + if b { + bad.store(true, std::sync::atomic::Ordering::Relaxed); + } + }; + if let Some(period) = period { + let state = &mut delayed[fid]; + let update = merge::delayed_nesterov_update(state.step, period); + let tasks: Vec<_> = params + .chunks(c) + .zip(momentum.chunks(c)) + .zip(state.accum.chunks(c)) + .zip(delta.chunks(c)) + .collect(); + crate::par::run_tasks(workers, tasks, |(((p, m), a), d)| { + flag(p.iter().zip(m).zip(a).zip(d).any(|(((p, m), a), d)| { + let (p, m, a) = + merge::delayed_nesterov_elem(*p, *m, *a, *d, lr, mu, period, update); + !p.is_finite() || !m.is_finite() || !a.is_finite() + })); + }); + if bad.load(std::sync::atomic::Ordering::Relaxed) { + bail!("fragment {fid}: outer optimizer produced non-finite state"); + } + merge::delayed_nesterov_step_par( + Arc::make_mut(params).as_mut_slice(), + Arc::make_mut(momentum).as_mut_slice(), + state, + delta, + lr, + mu, + period, + workers, + ); + } else { + let tasks: Vec<_> = + params.chunks(c).zip(momentum.chunks(c)).zip(delta.chunks(c)).collect(); + crate::par::run_tasks(workers, tasks, |((p, m), d)| { + flag(p.iter().zip(m).zip(d).any(|((param, momentum), gradient)| { + let next_momentum = mu * *momentum + *gradient; + let next_param = *param - lr * (*gradient + mu * next_momentum); + !next_param.is_finite() || !next_momentum.is_finite() + })); + }); + if bad.load(std::sync::atomic::Ordering::Relaxed) { + bail!("fragment {fid}: outer optimizer produced non-finite state"); + } + merge::nesterov_step_par( + Arc::make_mut(params).as_mut_slice(), + Arc::make_mut(momentum).as_mut_slice(), + delta, + lr, + mu, + workers, + ); + } + let gnorm = gnorm.join().expect("norm thread panicked"); + debug_assert!(gnorm.is_finite()); + Ok(gnorm) + }) + } + /// The Delayed Nesterov variant of the outer step in `commit_merge`, /// with the same validate-then-mutate contract. fn commit_delayed_nesterov(&mut self, fid: usize, delta: &[f32], period: u32) -> Result<()> { @@ -1190,9 +1315,61 @@ impl GlobalState { /// Written to `.tmp` then renamed, so a crash mid-write never /// corrupts the previous checkpoint. pub fn save_checkpoint(&self, path: &std::path::Path) -> Result<()> { + // Never two writers on `.tmp`; the newest cut lands last. + self.wait_for_background_checkpoint()?; self.checkpoint_snapshot()?.save(path) } + /// Start writing the current cut to `path` on its own thread and return + /// at once (`--background-checkpoint`). The write is the same as + /// `save_checkpoint`'s: the snapshot shares the fragment buffers, so a + /// later commit copies only a fragment the writer has not reached yet. + /// A previous background write is waited for first. A failed write + /// surfaces from the next wait: the next checkpoint, or the final one + /// (`save_checkpoint` waits too); `path` keeps the previous cut. + pub fn save_checkpoint_in_background(&self, path: &std::path::Path) -> Result<()> { + self.wait_for_background_checkpoint()?; + let snapshot = self.checkpoint_snapshot()?; + let step = self.global_step; + let target = path.to_path_buf(); + let writer = std::thread::Builder::new() + .name("yeto-checkpoint".into()) + .spawn(move || { + let started = std::time::Instant::now(); + snapshot.save(&target)?; + tracing::info!( + step, + path = %target.display(), + seconds = %format!("{:.1}", started.elapsed().as_secs_f64()), + "background checkpoint written" + ); + Ok(()) + }) + .context("cannot start the checkpoint writer thread")?; + *self.background_checkpoint.lock().unwrap_or_else(|e| e.into_inner()) = + Some(BackgroundCheckpoint { step, writer }); + Ok(()) + } + + /// Wait for the background checkpoint write, if any, and return its result. + pub fn wait_for_background_checkpoint(&self) -> Result<()> { + let pending = self + .background_checkpoint + .lock() + .unwrap_or_else(|e| e.into_inner()) + .take(); + let Some(BackgroundCheckpoint { step, writer }) = pending else { + return Ok(()); + }; + if !writer.is_finished() { + tracing::info!(step, "waiting for the background checkpoint"); + } + match writer.join() { + Ok(result) => result.with_context(|| format!("background checkpoint at step {step}")), + Err(_) => bail!("background checkpoint writer at step {step} panicked"), + } + } + /// The exact bytes `save_checkpoint` writes (tests and small states). pub fn encode_checkpoint(&self) -> Result> { let state_floats: usize = self.params.iter().map(|p| p.len()).sum::() * 2; @@ -1475,6 +1652,135 @@ impl GlobalState { } } +/// Fragments smaller than this always take the single-threaded path; thread +/// start-up would cost more than it saves. Results do not depend on it. +pub(crate) const PARALLEL_MIN_NUMEL: usize = 1 << 16; + +/// A tensor at least this large (and larger than its fair share of the +/// fragment) is processed alone with chunk-level parallelism; smaller tensors +/// are processed whole, many at a time. +const LARGE_TENSOR_MIN_NUMEL: usize = 1 << 20; + +fn is_large_tensor(tn: usize, numel: usize, threads: usize) -> bool { + tn >= LARGE_TENSOR_MIN_NUMEL && tn.saturating_mul(threads) > numel +} + +/// Offsets of each tensor in its fragment. +fn tensor_offsets(tensor_numels: &[usize]) -> Vec { + let mut off = 0; + tensor_numels + .iter() + .map(|&n| { + let o = off; + off += n; + o + }) + .collect() +} + +/// HeLoCo-corrected copies of every learner's outer gradient, bit-identical +/// to the serial per-learner, per-tensor `heloco_correct` loop: each +/// (learner, tensor) correction is independent of the others. +fn heloco_correct_all_par( + outer_gradients: &[&[f32]], + momentum: &[f32], + tensor_numels: &[usize], + h: &merge::Heloco, + threads: usize, +) -> Vec> { + let numel = momentum.len(); + let offsets = tensor_offsets(tensor_numels); + let mut corrected: Vec> = outer_gradients.iter().map(|_| vec![0.0f32; numel]).collect(); + let mut small = Vec::new(); + let mut large = Vec::new(); + for (buf, gradient) in corrected.iter_mut().zip(outer_gradients) { + for ((out, &off), &tn) in crate::par::split_lengths(buf, tensor_numels) + .into_iter() + .zip(&offsets) + .zip(tensor_numels) + { + let task = (out, &gradient[off..off + tn], &momentum[off..off + tn]); + if is_large_tensor(tn, numel, threads) { + large.push(task); + } else { + small.push(task); + } + } + } + // Largest first keeps the tail of the queue short. + small.sort_by_key(|(out, _, _)| std::cmp::Reverse(out.len())); + crate::par::run_tasks(threads, small, |(out, gradient, momentum)| { + out.copy_from_slice(gradient); + merge::heloco_correct(out, momentum, h); + }); + // Large tensors of different learners are corrected concurrently, each + // with its share of the threads, so their sequential norms overlap. + let outer = large.len().clamp(1, threads); + let inner = (threads / outer).max(1); + crate::par::run_tasks(outer, large, |(out, gradient, momentum)| { + crate::par::copy(inner, out, gradient); + merge::heloco_correct_par(out, momentum, h, inner.max(2)); + }); + corrected +} + +/// The per-tensor reduction loop of `prepare_merge_scaled`, run with +/// tensors in parallel (large tensors chunk-parallel). Bit-identical to the +/// serial loop: each tensor's merge reads only its own slices. +fn merge_tensors_par( + fid: usize, + merge_mode: u8, + outer_gradients: &[&[f32]], + weights: &[f64], + tensor_scales: Option<&[&[f64]]>, + tensor_numels: &[usize], + threads: usize, +) -> Result> { + if !tensor_numels.is_empty() && !matches!(merge_mode, MERGE_AVG | MERGE_ISO | MERGE_RDA) { + bail!("fragment {fid}: unsupported merge mode {merge_mode}"); + } + let numel: usize = tensor_numels.iter().sum(); + let offsets = tensor_offsets(tensor_numels); + let mut delta = vec![0.0f32; numel]; + let tensor_weights = |tensor_index: usize| -> Vec { + match tensor_scales { + None => weights.to_vec(), + Some(scales) => weights.iter().zip(scales).map(|(&w, s)| w * s[tensor_index]).collect(), + } + }; + let mut small = Vec::new(); + let mut large = Vec::new(); + for (tensor_index, out) in crate::par::split_lengths(&mut delta, tensor_numels).into_iter().enumerate() { + if is_large_tensor(out.len(), numel, threads) { + large.push((tensor_index, out)); + } else { + small.push((tensor_index, out)); + } + } + small.sort_by_key(|(_, out)| std::cmp::Reverse(out.len())); + let slices_of = |tensor_index: usize, tn: usize| -> Vec<&[f32]> { + let off = offsets[tensor_index]; + outer_gradients.iter().map(|g| &g[off..off + tn]).collect() + }; + crate::par::run_tasks(threads, small, |(tensor_index, out)| { + let slices = slices_of(tensor_index, out.len()); + let weights = tensor_weights(tensor_index); + match merge_mode { + MERGE_RDA => merge::merge_rda(&slices, &weights, out), + _ => merge::merge_avg(&slices, &weights, out), + } + }); + for (tensor_index, out) in large { + let slices = slices_of(tensor_index, out.len()); + let weights = tensor_weights(tensor_index); + match merge_mode { + MERGE_RDA => merge::merge_rda_par(&slices, &weights, out, threads), + _ => merge::merge_avg_par(&slices, &weights, out, threads), + } + } + Ok(delta) +} + #[cfg(test)] mod tests { use super::*; @@ -1542,6 +1848,191 @@ mod tests { } } + // ---- parallel merge: bit identity with the single-threaded path ---- + + struct Rng(u64); + impl Rng { + fn next_f32(&mut self) -> f32 { + self.0 ^= self.0 << 13; + self.0 ^= self.0 >> 7; + self.0 ^= self.0 << 17; + ((self.0 >> 40) as f32 / (1u64 << 24) as f32 - 0.5) * 0.02 + } + fn vec(&mut self, n: usize) -> Vec { + (0..n).map(|_| self.next_f32()).collect() + } + } + + fn par_state(tensors: &[u64], mode: u8, heloco: bool, delayed: Option, threads: usize) -> GlobalState { + let layout = Layout { + fragments: vec![FragmentInfo { + merge_mode: mode, + tensor_numels: tensors.to_vec(), + tensor_shapes: None, + }], + }; + let mut st = GlobalState::new(layout, 0.7, 0.9, crate::protocol::DTYPE_F32).unwrap(); + st.merge_threads = threads; + if heloco { + st.delta_correction = Some(merge::Heloco::default()); + } + st.set_delayed_nesterov_period(delayed).unwrap(); + let numel = tensors.iter().sum::() as usize; + st.init_fragment(0, Rng(99).vec(numel)).unwrap(); + st + } + + /// Learner gradients for one round that exercise every branch: random + /// tensors, an all-zero tensor (RDA φ(0), HeLoCo skip), exactly + /// cancelling tensors (RDA fallback to the average), and a learner that + /// opposes the momentum (HeLoCo shrink) next to near-orthogonal ones + /// (HeLoCo rotate) and aligned ones (pass through). + fn round_gradients(tensors: &[u64], round: u64, momentum: &[f32]) -> Vec> { + let numel = tensors.iter().sum::() as usize; + let mut rng = Rng(0x9e37_79b9_7f4a_7c15 ^ (round + 1).wrapping_mul(0x2545_f491)); + let mut g: Vec> = (0..3).map(|_| rng.vec(numel)).collect(); + let mut off = 0usize; + for (t, &tn) in tensors.iter().enumerate() { + let range = off..off + tn as usize; + match t % 4 { + 0 => g[1][range.clone()].fill(0.0), + 1 => { + let v = g[0][range.clone()].to_vec(); + for (x, y) in g[1][range.clone()].iter_mut().zip(&v) { + *x = -*y; + } + } + _ => {} + } + if t % 3 == 2 && !momentum.is_empty() { + for (x, m) in g[2][range.clone()].iter_mut().zip(&momentum[range.clone()]) { + *x = -4.0 * *m + *x * 0.01; + } + } + if t % 5 == 4 && !momentum.is_empty() { + for (x, m) in g[0][range.clone()].iter_mut().zip(&momentum[range.clone()]) { + *x = *m + *x * 0.01; + } + } + off += tn as usize; + } + g + } + + fn bits(v: &[f32]) -> Vec { + v.iter().map(|x| x.to_bits()).collect() + } + + fn assert_parallel_matches_serial(tensors: &[u64], mode: u8, heloco: bool, delayed: Option, scaled: bool) { + let weights = [1.0f64, 2.5, 0.75]; + let scales: Vec> = (0..3) + .map(|m| (0..tensors.len()).map(|t| 0.5 + ((m * 7 + t * 3) % 5) as f64).collect()) + .collect(); + let scale_refs: Vec<&[f64]> = scales.iter().map(|s| s.as_slice()).collect(); + let mut states: Vec = + [1usize, 2, 3, 8].iter().map(|&n| par_state(tensors, mode, heloco, delayed, n)).collect(); + for round in 0..4u64 { + let momentum = states[0].momentum[0].to_vec(); + let grads = round_gradients(tensors, round, &momentum); + let refs: Vec<&[f32]> = grads.iter().map(|g| g.as_slice()).collect(); + let mut results = Vec::new(); + for st in states.iter_mut() { + let prepared = st + .prepare_merge_scaled(0, round, &refs, &weights, scaled.then_some(scale_refs.as_slice())) + .unwrap(); + let gnorm = st.commit_merge(prepared.finish_inline().unwrap()).unwrap(); + st.versions[0] = round + 1; + let accum = st.delayed_nesterov.first().map(|s| bits(&s.accum)); + results.push((gnorm.to_bits(), bits(&st.params[0]), bits(&st.momentum[0]), accum)); + } + for (i, r) in results.iter().enumerate().skip(1) { + assert!( + *r == results[0], + "tensors {tensors:?} mode {mode} heloco {heloco} delayed {delayed:?} scaled {scaled}: \ + state {i} diverged from serial at round {round}" + ); + } + } + } + + const PAR_LAYOUTS: &[&[u64]] = &[ + // Below the parallel threshold: always the serial path. + &[7, 300, 1], + // Many small tensors, scheduled whole across threads. + &[40_000, 3, 70_001, 0, 65_536, 12_345, 99_999, 1, 131_073], + // One dominant tensor (chunk-parallel path) beside small ones. + &[1_500_003, 17, 65_537, 0, 4_099], + // Only one tensor, not a multiple of the chunk length. + &[(1 << 20) + 12_345], + ]; + + #[test] + fn parallel_merge_is_bitwise_serial_rda_heloco() { + for layout in PAR_LAYOUTS { + assert_parallel_matches_serial(layout, MERGE_RDA, true, None, false); + } + } + + #[test] + fn parallel_merge_is_bitwise_serial_avg_and_scales() { + for layout in PAR_LAYOUTS { + assert_parallel_matches_serial(layout, MERGE_AVG, false, None, false); + assert_parallel_matches_serial(layout, MERGE_RDA, false, None, true); + assert_parallel_matches_serial(layout, MERGE_AVG, true, None, true); + } + } + + #[test] + fn parallel_merge_is_bitwise_serial_delayed_nesterov() { + for layout in PAR_LAYOUTS { + assert_parallel_matches_serial(layout, MERGE_RDA, true, Some(3), false); + } + } + + #[test] + fn parallel_merge_rejects_non_finite_like_serial() { + let tensors: &[u64] = &[200_000, 70_000]; + let numel = 270_000usize; + let outcome = |threads: usize, poison: f32, at: usize, heloco: bool| { + let mut st = par_state(tensors, MERGE_RDA, heloco, None, threads); + let before = (bits(&st.params[0]), bits(&st.momentum[0])); + let mut g = Rng(5).vec(numel); + g[at] = poison; + let other = Rng(6).vec(numel); + let err = st + .prepare_merge(0, 0, &[&g, &other], &[1.0, 1.0]) + .and_then(|p| p.finish_inline()) + .and_then(|c| st.commit_merge(c)) + .unwrap_err() + .to_string(); + assert_eq!((bits(&st.params[0]), bits(&st.momentum[0])), before, "state mutated"); + err + }; + for poison in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] { + for at in [0, 199_999, 269_999] { + for heloco in [false, true] { + let serial = outcome(1, poison, at, heloco); + for threads in [2, 5] { + assert_eq!(outcome(threads, poison, at, heloco), serial); + } + } + } + } + // A finite merged gradient whose optimizer step overflows f32. + let overflow = |threads: usize| { + let mut st = par_state(tensors, MERGE_AVG, false, None, threads); + let before = bits(&st.params[0]); + let mut g = vec![0.0f32; numel]; + g[123_456] = 3.0e38; + let err = st.merge_and_step(0, &[&g], &[1.0]).unwrap_err().to_string(); + assert_eq!(bits(&st.params[0]), before); + err + }; + let serial = overflow(1); + assert!(serial.contains("outer optimizer produced non-finite state"), "{serial}"); + assert_eq!(overflow(4), serial); + } + #[test] fn init_once() { let mut st = GlobalState::new(layout2(), 0.7, 0.9, crate::protocol::DTYPE_F32).unwrap(); @@ -1762,6 +2253,41 @@ mod tests { std::fs::remove_dir_all(&dir).ok(); } + #[test] + fn background_checkpoint_writes_its_cut_while_later_commits_go_on() { + let dir = std::env::temp_dir().join(format!("yeto-ckpt-background-{}", std::process::id())); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("state.ckpt"); + let mut st = state_with_fingerprint(layout2(), 0x24); + st.init_fragment(0, vec![1.5; 4]).unwrap(); + st.init_fragment(1, vec![-2.0; 4]).unwrap(); + st.global_step = 3; + let first = st.encode_checkpoint().unwrap(); + st.save_checkpoint_in_background(&path).unwrap(); + // Commits do not wait for the writer; the file still holds step 3. + st.merge_and_step(0, &[&[1.0f32; 4][..]], &[1.0]).unwrap(); + st.global_step = 4; + let second = st.encode_checkpoint().unwrap(); + st.wait_for_background_checkpoint().unwrap(); + assert_eq!(std::fs::read(&path).unwrap(), first); + // The next background write waits for none (already joined) and + // writes the newer cut; a synchronous save after it waits for it. + st.save_checkpoint_in_background(&path).unwrap(); + st.merge_and_step(1, &[&[1.0f32; 4][..]], &[1.0]).unwrap(); + st.global_step = 5; + let third = st.encode_checkpoint().unwrap(); + st.save_checkpoint(&path).unwrap(); + assert_eq!(std::fs::read(&path).unwrap(), third); + assert_ne!(first, second); + // A failed background write is reported by the next wait. + let missing = dir.join("no-such-dir").join("state.ckpt"); + st.save_checkpoint_in_background(&missing).unwrap(); + let error = st.wait_for_background_checkpoint().unwrap_err(); + assert!(format!("{error:#}").contains("background checkpoint at step 5"), "{error:#}"); + st.wait_for_background_checkpoint().unwrap(); + std::fs::remove_dir_all(&dir).ok(); + } + #[test] fn snapshot_write_counts_copied_fragments_and_releases_every_buffer() { let dir = std::env::temp_dir().join(format!("yeto-ckpt-persist-{}", std::process::id())); diff --git a/tests/test_delta_protocol.py b/tests/test_delta_protocol.py index 47224251..748ca622 100644 --- a/tests/test_delta_protocol.py +++ b/tests/test_delta_protocol.py @@ -362,7 +362,7 @@ def test_fresh_and_one_version_stale_deltas_ignore_global_drift( @pytest.mark.timeout(30) -def test_mid_round_join_cannot_count_toward_frozen_membership( +def test_mid_round_join_is_pulled_into_the_round_after_its_catch_up( syncer_binary, tmp_path ): proc, final_state = launch_syncer( @@ -389,8 +389,10 @@ def test_mid_round_join_cannot_count_toward_frozen_membership( _, step, attempt = first.recv_pull(2) joined = RawLearner(proc.port, 1, 202, one_value_layout()) - # Learner 1 joined after attempt 1 of step 2 launched. Guessing its - # fields cannot add it to the frozen member set. + # Learner 1 rejoined after attempt 1 of step 2 launched without it. + # Once its catch-up is done it is pulled into that same attempt and + # its answer counts: (-100 + -1) / 2 on top of round 1's -1. + assert joined.recv_pull(2) == (0, step, attempt) joined.send_delta( torch.tensor([-100.0]), step=step, attempt=attempt, base_version=1 ) @@ -400,7 +402,7 @@ def test_mid_round_join_cannot_count_toward_frozen_membership( first.acknowledge_finalization() joined.acknowledge_finalization() assert proc.wait(timeout=10) == 0 - assert read_final_scalar(final_state) == pytest.approx(-2.0) + assert read_final_scalar(final_state) == pytest.approx(-51.5) finally: if first is not None: first.close() diff --git a/tests/test_fragments.py b/tests/test_fragments.py index d098ad0a..1fb3a122 100644 --- a/tests/test_fragments.py +++ b/tests/test_fragments.py @@ -152,3 +152,35 @@ def test_avg_regex_groups_vector_like_tensors(): assert layout.fragments[0].merge_mode == MERGE_AVG assert layout.fragments[0].tensors == [("block.norm.weight", 8)] assert all(f.merge_mode == MERGE_RDA for f in layout.fragments[1:]) + + +def _reference_binpack(named, n_bins): + """The greedy packing as originally written: rescan every bin's total.""" + bins = [[] for _ in range(n_bins)] + for name, numel in sorted(named, key=lambda x: (-x[1], x[0])): + min(bins, key=lambda b: sum(s for _, s in b)).append((name, numel)) + return [b for b in bins if b] + + +def test_binpack_matches_the_rescanning_greedy_with_ties(): + import random + + rng = random.Random(7) + for trial in range(20): + named = [(f"model.layers.{i % 9}.w{i}", rng.choice([1, 2, 3, 8, 64, rng.randint(1, 500)])) + for i in range(rng.randint(1, 300))] + n_bins = rng.randint(1, 12) + layout = build_layout(named, n_bins) + expected = _reference_binpack(named, max(1, min(n_bins, len(named)))) + assert [f.tensors for f in layout.fragments] == expected, trial + + +def test_binpack_is_fast_for_an_unfused_moe_tensor_count(): + import time + + named = [(f"backbone.layers.{i // 800}.mixer.experts.{i % 800}.w", 2_000_000 + i % 5) + for i in range(42_000)] + started = time.monotonic() + layout = build_layout(named, 32) + assert time.monotonic() - started < 30.0 # the rescanning greedy took minutes + assert sum(len(f.tensors) for f in layout.fragments) == len(named) diff --git a/tests/test_sync_fragment_boundary.py b/tests/test_sync_fragment_boundary.py index 9b0add24..8a87b20a 100644 --- a/tests/test_sync_fragment_boundary.py +++ b/tests/test_sync_fragment_boundary.py @@ -364,6 +364,18 @@ def test_the_manifest_only_profile_matches_the_rust_vector(): assert on == vector != off + +def test_the_carry_answers_round_profile_matches_the_rust_vector(): + from test_rl_moe_expert_token_weights import _syncer_profile + + raw = _syncer_profile() + off = SyncerSemanticProfile.from_mapping(raw).sha256 + assert SyncerSemanticProfile.from_mapping({**raw, "carry_answers_round": False}).sha256 == off + on = SyncerSemanticProfile.from_mapping({**raw, "carry_answers_round": True}).sha256 + source = (ROOT / "syncer/src/server.rs").read_text() + (vector,) = re.findall(r'CARRY_ANSWERS_ROUND_VECTOR: &str = "([0-9a-f]{64})"', source) + assert on == vector != off + # -- the real syncer ------------------------------------------------------------------ diff --git a/yeto/fragments.py b/yeto/fragments.py index 7441f4b3..829c730e 100644 --- a/yeto/fragments.py +++ b/yeto/fragments.py @@ -147,9 +147,13 @@ def use_avg(name: str) -> bool: for name, numel in sorted(layered, key=lambda x: (layer_index(x[0]), x[0])): bins[bin_of_layer[layer_index(name)]].tensors.append((name, numel)) rest = loose # norms, biases, etc. fall through to bin-packing below + # Running bin totals: recomputing ``Fragment.numel`` per placement is + # quadratic in the tensor count (minutes for an unfused 120B MoE). + totals = [b.numel for b in bins] for name, numel in rest: - target = min(bins, key=lambda b: b.numel) # min() is stable: first-smallest wins - target.tensors.append((name, numel)) + i = min(range(len(bins)), key=totals.__getitem__) # min() is stable: first-smallest wins + bins[i].tensors.append((name, numel)) + totals[i] += numel fragments.extend(b for b in bins if b.tensors) # strided can leave bins empty if matrix_mode == MERGE_ISO: for frag in fragments: diff --git a/yeto/protocol.py b/yeto/protocol.py index b1646aa6..87d17989 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -2274,28 +2274,20 @@ def _send_large_parts( try: while offset < total: inner_size = min(CHUNK_SIZE, total - offset) - envelope = bytearray(envelope_prefix + inner_size) - _HEADER.pack_into( - envelope, - 0, - MAGIC, - MSG_CHUNK, - _CHUNK_HEAD.size + inner_size, - ) - _CHUNK_HEAD.pack_into( - envelope, - _HEADER.size, - msg_id, - total, - offset, - ) + # The envelope is joined from its pieces in one copy: filling a + # zeroed bytearray slice by slice capped a sender at ~1.5 GB/s. + pieces: list[bytes | memoryview] = [ + _HEADER.pack(MAGIC, MSG_CHUNK, _CHUNK_HEAD.size + inner_size) + + _CHUNK_HEAD.pack(msg_id, total, offset) + ] + consumed: list[memoryview] = [] inner_end = offset + inner_size write_at = envelope_prefix if offset < len(fixed): fixed_end = min(inner_end, len(fixed)) fixed_bytes = fixed[offset:fixed_end] - envelope[write_at : write_at + len(fixed_bytes)] = fixed_bytes + pieces.append(fixed_bytes) write_at += len(fixed_bytes) remaining = envelope_prefix + inner_size - write_at @@ -2322,17 +2314,23 @@ def _send_large_parts( available = pending.nbytes - pending_offset take = min(remaining, available) - envelope[write_at : write_at + take] = pending[ - pending_offset : pending_offset + take - ] + pieces.append(pending[pending_offset : pending_offset + take]) write_at += take remaining -= take pending_offset += take parts_seen += take if pending_offset == pending.nbytes: - pending.release() + consumed.append(pending) # released once joined pending = None + envelope = b"".join(pieces) + for piece in pieces: + if isinstance(piece, memoryview): + piece.release() + for view in consumed: + view.release() + del pieces, consumed + if inner_end == total: if parts_seen != parts_size: raise ValueError( diff --git a/yeto/syncer_profile.py b/yeto/syncer_profile.py index 0d2e3f26..7f4c095c 100644 --- a/yeto/syncer_profile.py +++ b/yeto/syncer_profile.py @@ -72,7 +72,12 @@ # the listed global steps (syncer --checkpoint-at-steps) append their block. _CHECKPOINT_AT_STEPS_FIELD = "checkpoint_at_steps" _EVAL_CHECKPOINTS_DOMAIN = b"yeto-syncer-eval-checkpoints-v1\0" +# Optional. Absent or false = carried late deltas stay outside quorum +# (byte-identical hash); true (syncer --carry-answers-round) appends its domain. +_CARRY_ANSWERS_ROUND_FIELD = "carry_answers_round" +_CARRY_ANSWERS_ROUND_DOMAIN = b"yeto-syncer-carry-answers-round-v1\0" _OPTIONAL_FIELDS = { + _CARRY_ANSWERS_ROUND_FIELD, _DELAYED_NESTEROV_FIELD, _EXPERT_TOKEN_WEIGHTS_FIELD, _FINAL_CUT_MANIFEST_ONLY_FIELD, @@ -174,6 +179,7 @@ class SyncerSemanticProfile: grace_arrival_quantile: float = 0.0 terminal_grace_ms: int = 0 checkpoint_at_steps: tuple[int, ...] = () + carry_answers_round: bool = False @classmethod def from_mapping(cls, raw: object) -> SyncerSemanticProfile: @@ -265,6 +271,9 @@ def from_mapping(cls, raw: object) -> SyncerSemanticProfile: checkpoint_at_steps=_steps( _CHECKPOINT_AT_STEPS_FIELD, raw.get(_CHECKPOINT_AT_STEPS_FIELD, ()) ), + carry_answers_round=_boolean( + _CARRY_ANSWERS_ROUND_FIELD, raw.get(_CARRY_ANSWERS_ROUND_FIELD, False) + ), ) profile._validate_general() return profile @@ -373,6 +382,8 @@ def canonical_bytes(self) -> bytes: encoded.extend(_EVAL_CHECKPOINTS_DOMAIN) encoded.extend(struct.pack("