From b62764edbd7b48e8dfcf369a710c26b6a66408e1 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 14:06:54 +0800 Subject: [PATCH 01/16] syncer: release a stale-base slot so async rounds do not stall until quorum timeout Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 56 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index aed60d58..931e22f2 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -5523,6 +5523,16 @@ 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. + if !fixed_roster { + round.forget_unanswered(member); + } return PushDisposition::StaleBase; } let duplicate = if fixed_roster { @@ -12508,6 +12518,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); From 1d6889467cfdc699cbc004d4dd78404298a008e1 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 14:32:22 +0800 Subject: [PATCH 02/16] syncer: run the merge and outer step on multiple threads, bit-identical to serial Independent reductions (tensors, learners, HeLoCo's three norms) run concurrently; element-wise passes (accumulation, rescale, Nesterov and Delayed Nesterov steps, finiteness checks) are chunked. No reduction is split, so every thread count gives the exact bits of the serial path. --merge-threads / YETO_SYNCER_MERGE_THREADS selects the count (default: cores capped at 32); 1 runs the historical code unchanged. Co-Authored-By: Claude Opus 5.5 --- syncer/src/main.rs | 12 ++ syncer/src/merge.rs | 191 +++++++++++++++++++ syncer/src/par.rs | 163 +++++++++++++++++ syncer/src/state.rs | 433 +++++++++++++++++++++++++++++++++++++++++++- 4 files changed, 797 insertions(+), 2 deletions(-) create mode 100644 syncer/src/par.rs diff --git a/syncer/src/main.rs b/syncer/src/main.rs index cc4efcbf..f1e1a230 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; @@ -82,6 +83,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, @@ -387,6 +393,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, 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/state.rs b/syncer/src/state.rs index e272b8ee..250d563d 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. @@ -373,6 +377,7 @@ impl GlobalState { 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 +549,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 +593,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 +702,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 +738,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 +778,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<()> { @@ -1475,6 +1590,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 +1786,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(); From 3fc15e46482cd8e1db5410de3cb2e752732a084c Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 14:38:59 +0800 Subject: [PATCH 03/16] syncer: keep a stale-base slot when the round has a fold sized for the full roster Releasing the slot of a refused stale-base push let an exact or streaming fold close one push short and stop the syncer with 'fold incomplete'. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 931e22f2..aa8cbeb3 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -5529,8 +5529,9 @@ fn route_push( // 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. - if !fixed_roster { + // the slot; so does a round with a fold, which was sized for the full + // expected roster and may already have consumed pushes. + if !fixed_roster && round.fold.is_none() { round.forget_unanswered(member); } return PushDisposition::StaleBase; From 715079e1c26d524ea826b5a872cd09fd43cfe7ac Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 15:45:28 +0800 Subject: [PATCH 04/16] syncer: let a late-carried delta answer its round (--carry-answers-round) With quorum 1 a round whose only delta from the fast island arrived as a late carry (its anchor spent) waited for the stalled island, and with --pipeline 2 so did the next round. The new flag counts expected carried members as responders for quorum and the all-answered check. Off keeps the historical RL async rule; the flag joins the semantic profile hash only when on. Co-Authored-By: Claude Opus 5.5 --- syncer/src/main.rs | 6 ++ syncer/src/server.rs | 120 ++++++++++++++++++++++++++- tests/test_sync_fragment_boundary.py | 12 +++ yeto/syncer_profile.py | 11 +++ 4 files changed, 146 insertions(+), 3 deletions(-) diff --git a/syncer/src/main.rs b/syncer/src/main.rs index f1e1a230..3a36262b 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -282,6 +282,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. @@ -509,6 +514,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/server.rs b/syncer/src/server.rs index aa8cbeb3..a4b377a1 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -209,6 +209,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 +241,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 +313,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() } } @@ -4088,6 +4100,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 +5278,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 +5394,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() { @@ -9090,6 +9123,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, @@ -9100,6 +9134,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"; @@ -9155,6 +9190,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; @@ -9579,6 +9624,7 @@ mod tests { grace_ms: Some(0), fold: None, carried: Vec::new(), + carry_answers: false, pushes: HashMap::from([( member, Push { @@ -9662,6 +9708,7 @@ mod tests { grace_ms: Some(0), fold: None, carried: Vec::new(), + carry_answers: false, pushes: HashMap::from([( member, Push { @@ -9725,6 +9772,7 @@ mod tests { grace_ms: Some(0), fold: None, carried: Vec::new(), + carry_answers: false, pushes, } } @@ -11945,6 +11993,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; @@ -12054,6 +12167,7 @@ mod tests { grace_ms: None, fold: None, carried: Vec::new(), + carry_answers: false, pushes: HashMap::new(), } } 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/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(" Date: Thu, 5 Nov 2026 16:54:28 +0800 Subject: [PATCH 05/16] fragments: keep running bin totals in the binpack layout build_layout recomputed every bin's numel for each placed tensor, which is quadratic in the tensor count: the unfused 120B MoE shape list (41,643 tensors, 32 fragments) took 107 s per learner process. Track the totals instead; the placement (first smallest bin wins) is unchanged. Co-Authored-By: Claude Opus 5.5 --- tests/test_fragments.py | 32 ++++++++++++++++++++++++++++++++ yeto/fragments.py | 8 ++++++-- 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/tests/test_fragments.py b/tests/test_fragments.py index d098ad0a..00b63409 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 + 0) + 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 < 10.0 + assert sum(len(f.tensors) for f in layout.fragments) == len(named) 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: From 623894e6c309e05ef54a4be082de6d230a4d09fd Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 16:54:28 +0800 Subject: [PATCH 06/16] syncer: hash INIT_PARAMS with the ARMv8 SHA-2 instructions, decode it in parallel Every INIT_PARAMS payload is hashed and decoded on the connection's read loop. On a Kunpeng 920 the vendored sha2 runs its portable code (its aarch64 backend needs the asm feature and another crate), about 0.25 GB/s, and the f32 decode plus the finiteness check are one more single-threaded pass, so a 120B model's 483 GB of f32 INIT took about 40 minutes to admit. - sha256.rs: SHA-256 with the ARMv8 SHA-2 instructions when the CPU has them (runtime detected), else sha2; byte-identical digests (padding boundaries, split updates and known vectors tested against sha2). Used for INIT digests and the downlink digest. - decode_tensor splits large f32 / bf16 payloads over the merge threads (disjoint chunks, same per-value conversion, same error for a non-finite value), and INIT computes its digest while it decodes. Measured on the production syncer host with the same 120B INIT stream: 0.21 -> 0.5 GB/s; the digest (now ~1.6 GB/s, one sequential pass) is the remaining bound. Co-Authored-By: Claude Opus 5.5 --- syncer/src/downlink.rs | 5 +- syncer/src/main.rs | 1 + syncer/src/protocol.rs | 87 ++++++++++++++++ syncer/src/server.rs | 11 +- syncer/src/sha256.rs | 226 +++++++++++++++++++++++++++++++++++++++++ 5 files changed, 324 insertions(+), 6 deletions(-) create mode 100644 syncer/src/sha256.rs 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 3a36262b..6b7e10be 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -18,6 +18,7 @@ mod eval_checkpoint; mod protocol; mod sample_relay; mod server; +mod sha256; mod state; mod subscriber; diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index 2a562d73..13d09c14 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -454,6 +454,26 @@ pub fn decode_budget_done(payload: &[u8]) -> Result { /// Decode a tensor payload into f32s. pub fn decode_tensor(dtype: u8, bytes: &[u8], out: &mut Vec) -> Result<()> { + decode_tensor_with(crate::par::default_threads(), 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 +505,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 +908,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 a4b377a1..75731b9c 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -3251,9 +3251,14 @@ async fn dispatch_inner( bytes.len() ); } - let sha256: [u8; 32] = Sha256::digest(bytes).into(); - let mut values = Vec::new(); - decode_tensor(bulk_dtype(group.dtype), bytes, &mut values)?; + // 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(bulk_dtype(group.dtype), bytes, &mut values); + (digest.join().expect("INIT digest thread panicked"), decoded.map(|()| values)) + }); + let values = values?; event_tx .send(Event::Init { member: group.member, diff --git a/syncer/src/sha256.rs b/syncer/src/sha256.rs new file mode 100644 index 00000000..cf09fe70 --- /dev/null +++ b/syncer/src/sha256.rs @@ -0,0 +1,226 @@ +//! 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 } => { + *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 } => { + 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()); + } +} From 917fc92653f8f112ac88c6573a5b0749b35b881e Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 17:15:17 +0800 Subject: [PATCH 07/16] syncer: encode broadcast chunks in one pass over their whole values append_tensor_bytes sliced every value by the chunk's byte range and appended it on its own, about 0.4 GB/s on a Kunpeng 920. Every f32 broadcast of a 120B fragment (15 GB) goes through it, one 4 MiB chunk at a time on the forwarder, so the syncer's downlink to one island ran at ~0.42 GB/s on a 50 Gb/s link: the initial cut alone took ~20 minutes. Only the values cut by the range ends now go byte by byte; the bytes are unchanged (tested against the value-by-value encoding, NaN, -0.0, subnormal and every cut position). Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 87 +++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 85 insertions(+), 2 deletions(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 75731b9c..95f7074e 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -743,7 +743,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); @@ -757,7 +795,6 @@ fn append_tensor_bytes( _ => unreachable!(), } } - Ok(()) } /// One receive-only subscriber connection (`subscriber.rs`). @@ -8760,6 +8797,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); From d4bd680b741aed200624aac5e3eb1a1864ad16b2 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 18:34:50 +0800 Subject: [PATCH 08/16] syncer: --background-checkpoint writes periodic checkpoints off the scheduler A --checkpoint-every checkpoint (other than per-round ones) was written synchronously in publish_round: the scheduler merged and broadcast nothing until the whole state was on disk. For the 120B state (966 GB of f32 params and momentum) that is a 13.7 minute stall per checkpoint on an NVMe that sustains ~1.2 GB/s of buffered writes (measured on the production syncer host). With --background-checkpoint the cut is taken as the existing copy-free CheckpointSnapshot (later commits copy only a fragment the writer has not reached) and saved on its own thread while rounds go on. At most one such write runs: the next periodic checkpoint, and every synchronous save (terminal, budget cutoff, shutdown), waits for it, so two writers never share .tmp and the newest cut lands last. A failed write surfaces from that wait; the previous file stays intact. The flag needs a checkpoint path and --checkpoint-every >= 2 outside policy sweeps (per-round checkpoints still persist before broadcast). Off by default; the files are byte-identical to the synchronous path. Co-Authored-By: Claude Opus 5.5 --- syncer/src/main.rs | 8 ++++ syncer/src/server.rs | 79 +++++++++++++++++++++++++++++++++++- syncer/src/state.rs | 97 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 182 insertions(+), 2 deletions(-) diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 6b7e10be..03fc9f9d 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -148,6 +148,13 @@ 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. + #[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)] @@ -494,6 +501,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, diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 95f7074e..eb7db06d 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. @@ -1759,6 +1763,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"); @@ -6206,8 +6220,13 @@ async fn publish_round( 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"); + 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"); + } } } } @@ -9191,6 +9210,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, @@ -10854,6 +10874,61 @@ 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]); + // 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"); diff --git a/syncer/src/state.rs b/syncer/src/state.rs index 250d563d..2db0ead4 100644 --- a/syncer/src/state.rs +++ b/syncer/src/state.rs @@ -169,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 @@ -373,6 +382,7 @@ 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, @@ -1305,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; @@ -2191,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())); From 10aadc62df96a406e4c69fb377b790647bebb165 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 19:36:54 +0800 Subject: [PATCH 09/16] syncer: hash and decode INIT_PARAMS off the read loop, without a payload copy A reassembled INIT_PARAMS was copied out of its reassembly buffer, then hashed and decoded on the connection's read loop, so that stream stopped reading for each fragment's ~10 s digest (one sequential pass) and the learner's round-robin chunks backed up behind it. A 120B f32 INIT (483 GB) arrived at ~0.5 GB/s on a 50 Gb/s link. The reassembled frame now keeps its buffer (no copy), and an INIT_PARAMS is validated, hashed and decoded on the blocking pool while the read loop goes on; up to four per connection are in flight (each holds one fragment's payload and values) before the loop waits for the oldest. A failed decode is reported to the learner and ends the connection as before; the coordinator applies INITs in any order, from the current generation only. Direct (unchunked) INIT frames keep the inline path. Measured on the production syncer host with a synthetic 120B learner: the same INIT stream now takes 322 s (1.5 GB/s, bound by the Python client's send loop) instead of 16.9 min. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 165 +++++++++++++++++++++++++++++++++---------- 1 file changed, 127 insertions(+), 38 deletions(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index eb7db06d..d81cc451 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -3111,8 +3111,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()); @@ -3148,11 +3169,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, @@ -3182,7 +3207,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()?; @@ -3260,17 +3285,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(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( @@ -3283,33 +3398,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() - ); - } - // 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(bulk_dtype(group.dtype), bytes, &mut values); - (digest.join().expect("INIT digest thread panicked"), decoded.map(|()| values)) - }); - let values = values?; + let (fragment_id, sha256, values) = decode_init_params(group, payload)?; event_tx .send(Event::Init { member: group.member, From bbe0e0c145343207b32fc8214638fdedbe5a1de3 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 19:46:17 +0800 Subject: [PATCH 10/16] protocol: join chunk envelopes in one copy _send_large_parts filled each 4 MiB CHUNK envelope by allocating a zeroed bytearray and assigning slices into it. On a Kunpeng 920 that loop sends at most ~2 GB/s and measured 1.5 GB/s for a 120B INIT (483 GB, 322 s), the bound once the syncer keeps up. The envelope is now b"".join of its header and the part slices: one copy, ~5.9 GB/s in the same micro-benchmark. Bytes and validation are unchanged. Co-Authored-By: Claude Opus 5.5 --- yeto/protocol.py | 38 ++++++++++++++++++-------------------- 1 file changed, 18 insertions(+), 20 deletions(-) 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( From ea4b36c0e7e6b74f2efa2d71c9a4d14afc1eba80 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 20:13:27 +0800 Subject: [PATCH 11/16] syncer: leave the last step to the final checkpoint When --total-steps is a multiple of --checkpoint-every, the periodic checkpoint of the last step and the final checkpoint wrote the same cut back to back. For the 120B state that is two ~14 minute writes of 966 GB at the end of a run (measured on the production syncer host: 10:44 -> 10:58 -> 11:12). The periodic write now skips the last step; the final checkpoint, which is always written when a checkpoint path is set, persists it. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 37 ++++++++++++++++++++++++++++++++++++- 1 file changed, 36 insertions(+), 1 deletion(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index d81cc451..b867c2b1 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -6308,7 +6308,12 @@ 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 { + // 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"); @@ -11002,6 +11007,36 @@ mod tests { 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, From 309f7990430a73b354b3e31ae489703490379fbe Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 21:46:36 +0800 Subject: [PATCH 12/16] syncer: keep queued late-carry pushes as compressed wire bytes A late push waiting for its fragment's next round held a decoded f32 outer gradient (about 4P bytes per slow island at 120B). Queue the original compressed payload instead and decode it when the push is attached to the next round. The decode is deterministic, so merged results are bit-identical. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 108 +++++++++++++++++++++++++++++++++++++++---- 1 file changed, 99 insertions(+), 9 deletions(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index b867c2b1..ff261b98 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -1137,9 +1137,59 @@ 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(); + } + } + + /// 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 @@ -3467,16 +3517,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, @@ -3491,6 +3537,7 @@ async fn dispatch_inner( c_tokens, outer_gradient, tensor_scales, + wire, transfer: PushTransfer { recv_seconds, arrived_unix: unix_now(), @@ -5851,6 +5898,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 } @@ -5862,7 +5911,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)) }) @@ -9729,6 +9779,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(); @@ -9840,6 +9891,7 @@ mod tests { c_tokens: 1, outer_gradient: vec![0.25; 4], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, )]), @@ -9924,6 +9976,7 @@ mod tests { c_tokens, outer_gradient: vec![0.25], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, )]), @@ -9951,6 +10004,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(), }, ) @@ -12208,6 +12262,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(); @@ -12471,6 +12555,7 @@ mod tests { c_tokens: 20, outer_gradient: vec![1.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), } } @@ -12885,6 +12970,7 @@ mod tests { c_tokens: 1, outer_gradient: vec![1.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), } } @@ -13012,6 +13098,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, @@ -13189,6 +13276,7 @@ mod tests { c_tokens: 10, outer_gradient: vec![1.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, ), @@ -13205,6 +13293,7 @@ mod tests { c_tokens: 1_000, outer_gradient: vec![2.0], tensor_scales: None, + wire: None, transfer: PushTransfer::default(), }, ), @@ -13522,6 +13611,7 @@ mod tests { c_tokens: 1 + u64::from(learner_id), outer_gradient: gradient, tensor_scales: None, + wire: None, transfer: PushTransfer::default(), } } From 69ac93049ff2023c1b9922b8244df5093f420e80 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Sat, 7 Nov 2026 19:54:02 +0800 Subject: [PATCH 13/16] syncer: review fixes for push memory, stale-base slots, decode threads - Drop the wire copy of a push once it joins a round. The compressed bytes exist so a late push can wait in the late-carry queue without its decoded gradient; a push accepted by its own round, or attached to a later one, kept that copy until the round was committed (up to 2P more bytes per bf16 island, on top of the gradient the streaming fold had just released). - Release the slot of a stale-base push under a streaming fold as well. Only the exact fold fixes its order and weights for the full roster; a streaming fold normalises over what it folded, and async carry (where stale bases happen) needs one, so the earlier guard left those rounds waiting for the quorum timeout again. - Decode PUSH payloads on the calling thread, as before the parallel INIT decode: a large push no longer spawns a thread pool per connection that competes with the merge threads. INIT_PARAMS asks for threads explicitly. - Assert the SHA-2 feature on the hardware hasher path in debug builds. - Document the worst-case memory of --background-checkpoint. Co-Authored-By: Claude Opus 5.5 --- syncer/src/main.rs | 5 +++- syncer/src/protocol.rs | 6 ++-- syncer/src/server.rs | 63 +++++++++++++++++++++++++++++++++++++++--- syncer/src/sha256.rs | 3 ++ 4 files changed, 70 insertions(+), 7 deletions(-) diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 03fc9f9d..9bb05cca 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -152,7 +152,10 @@ struct Args { /// 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. + /// (--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 diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index 13d09c14..e6655666 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -452,9 +452,11 @@ 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(crate::par::default_threads(), dtype, bytes, out) + decode_tensor_with(1, dtype, bytes, out) } /// Values below which a decode stays on the calling thread. diff --git a/syncer/src/server.rs b/syncer/src/server.rs index ff261b98..af9f6920 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -1160,6 +1160,12 @@ impl Push { } } + /// 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) { @@ -3388,7 +3394,7 @@ fn decode_init_params(group: &Group, payload: &[u8]) -> Result<(u32, [u8; 32], V let (sha256, values) = std::thread::scope(|scope| { let digest = scope.spawn(|| crate::sha256::digest(bytes)); let mut values = Vec::new(); - let decoded = decode_tensor(bulk_dtype(group.dtype), bytes, &mut values); + 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?)) @@ -5754,9 +5760,10 @@ fn route_push( // 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 a fold, which was sized for the full - // expected roster and may already have consumed pushes. - if !fixed_roster && round.fold.is_none() { + // 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; @@ -5772,6 +5779,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 @@ -5983,6 +5992,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) } @@ -13893,6 +13904,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 index cf09fe70..262070a0 100644 --- a/syncer/src/sha256.rs +++ b/syncer/src/sha256.rs @@ -107,6 +107,8 @@ impl Hasher { 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()); @@ -135,6 +137,7 @@ impl Hasher { 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 { From b95be1f8bbb2e6035400a0f0ac70767a8e07b201 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Sat, 7 Nov 2026 19:54:02 +0800 Subject: [PATCH 14/16] tests: tidy the binpack tests Drop a leftover `+ 0` and give the speed check headroom on a slow CI runner; the rescanning greedy it guards against took minutes. Co-Authored-By: Claude Opus 5.5 --- tests/test_fragments.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_fragments.py b/tests/test_fragments.py index 00b63409..1fb3a122 100644 --- a/tests/test_fragments.py +++ b/tests/test_fragments.py @@ -170,7 +170,7 @@ def test_binpack_matches_the_rescanning_greedy_with_ties(): 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 + 0) + 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 @@ -182,5 +182,5 @@ def test_binpack_is_fast_for_an_unfused_moe_tensor_count(): for i in range(42_000)] started = time.monotonic() layout = build_layout(named, 32) - assert time.monotonic() - started < 10.0 + assert time.monotonic() - started < 30.0 # the rescanning greedy took minutes assert sum(len(f.tensors) for f in layout.fragments) == len(named) From 9b57ab32a73862b7b0aed44c317f3e23d17c6071 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Sat, 7 Nov 2026 19:44:02 +0800 Subject: [PATCH 15/16] syncer: pull a learner that finishes its catch-up into rounds launched without it A round's expected set is frozen at launch to the members not held by a catch-up or background resend. After a --resume, learners reconnect one by one and each spends a long time in its catch-up, so a round can be launched to the first learner alone. If that learner's H-gate keeps it from answering, the learners that finish their catch-up later were only pulled when the quorum timeout relaunched the round, and in-order commits stalled the whole fleet for that timeout. When a catch-up or recovery replay tail runs, the member now also joins every still-gathering round that has no slot, push or carried delta of its logical learner: it becomes expected in the current attempt and gets that attempt's pull. Attempt, quorum size, base version and the quorum, soft and grace deadlines are unchanged, so a push that arrives inside an armed grace window still counts. Completed rounds, rounds whose base the member was not caught up to, and exact folds that already hold pushes are skipped; an empty exact fold is rebuilt for the new roster. Fixed rosters are unchanged. A reconnect generation is still handled only by the rebind, so no round is pulled twice. The real-syncer protocol test now waits for that pull and expects the merged value with both answers. Co-Authored-By: Claude Opus 5.5 --- syncer/src/server.rs | 379 ++++++++++++++++++++++++++++++++++- tests/test_delta_protocol.py | 10 +- 2 files changed, 381 insertions(+), 8 deletions(-) diff --git a/syncer/src/server.rs b/syncer/src/server.rs index af9f6920..535a3ac1 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -5725,6 +5725,82 @@ 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`), 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; + } + 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, @@ -6799,12 +6875,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 { @@ -6813,6 +6896,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; @@ -6828,6 +6916,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 } @@ -8307,7 +8409,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; } } @@ -8385,7 +8487,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<()> { @@ -12884,6 +12986,275 @@ 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 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); 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() From 2684decaf43b90d40af52de667f5469a7f29c8b2 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Sat, 7 Nov 2026 19:55:20 +0800 Subject: [PATCH 16/16] syncer: keep capacity-quorum rounds frozen; document the catch-up join A capacity quorum is a share of the expected capacity, and its rule assumes the expected set only shrinks while a round gathers. A learner joining after its catch-up would raise the bar of a quorum already reached, and with an armed grace window turn a closing round back into one that waits for the soft deadline or the quorum timeout. Skip such rounds; count quorums keep the join. The design notes described the rendezvous as fully frozen; state the catch-up exception. Co-Authored-By: Claude Opus 5.5 --- docs/DESIGN.md | 6 ++++++ syncer/src/server.rs | 33 ++++++++++++++++++++++++++++++--- 2 files changed, 36 insertions(+), 3 deletions(-) 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/server.rs b/syncer/src/server.rs index 535a3ac1..f15d96b1 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -5752,9 +5752,11 @@ struct JoinedPull { /// 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`), 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. +/// (`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, @@ -5781,6 +5783,12 @@ fn join_gathering_rounds( { 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; @@ -13025,6 +13033,25 @@ mod tests { .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.