From b62764edbd7b48e8dfcf369a710c26b6a66408e1 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 14:06:54 +0800 Subject: [PATCH 01/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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 dee832b727217f78db2377346ee16b5c2640b0a5 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 14:25:08 +0800 Subject: [PATCH 15/20] Add full-parameter DiLoCo sync core for ms-swift Megatron islands MegatronDiLoCoSync speaks the torch learner's DiLoCo protocol (layout, INIT, raw-anchor PUSH deltas, BCAST apply before pulls, finalization, HMAC HELLO) to the unchanged syncer from a TP/EP/PP Megatron island: - global trainable set named by global Megatron names (global layer and expert indices), frozen parameters and buffers excluded; - chunked reads of FP32 optimizer masters (DistributedOptimizer shards, HybridDeviceOptimizer host masters) gathered over DP and TP and sent point-to-point to rank 0; per-rank anchors keep deltas local; - writes install the merged base into the FP32 master and the bf16 parameter, keeping the Adam moments; - collectives on host buffers over gloo mirrors by default, so the sync allocates nothing on the accelerator; - rank 0's per-boundary plan is broadcast so every rank runs the same collectives and rank-0 errors are raised everywhere. write_masters also updates a distinct HybridDeviceOptimizer host copy. Co-Authored-By: Claude Opus 5.5 --- tests/test_swift_megatron_sync.py | 568 ++++++++++++++++ yeto/rl/megatron_masters.py | 7 + yeto/swift_megatron/__init__.py | 1 + yeto/swift_megatron/sync.py | 1021 +++++++++++++++++++++++++++++ 4 files changed, 1597 insertions(+) create mode 100644 tests/test_swift_megatron_sync.py create mode 100644 yeto/swift_megatron/__init__.py create mode 100644 yeto/swift_megatron/sync.py diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py new file mode 100644 index 00000000..701d39fa --- /dev/null +++ b/tests/test_swift_megatron_sync.py @@ -0,0 +1,568 @@ +"""Full-parameter DiLoCo sync of an ms-swift Megatron island (yeto/swift_megatron/sync.py). + +Megatron-Core does not run on CPU, so the islands here are gloo process groups +whose model chunks carry exactly the attributes the sync reads from a Megatron +model: pipeline-local ``decoder.layers.`` with a global ``layer_number``, +TP-sharded parameters (``tensor_model_parallel`` / ``partition_dim``), expert +parameters (``allreduce = False``) under containers with +``local_expert_indices`` (SequentialMLP and TEGroupedMLP naming), a frozen +vision projection, and a DistributedOptimizer leaf whose inner optimizer is a +precision-aware ``HybridDeviceOptimizer`` (FP32 host masters behind +``param_to_fp32_param``), the production optimizer. + +Every full tensor is a pure function of its global name, so each rank cuts +its shard out of the same full value and a gathered tensor can be checked +against it. +""" + +from __future__ import annotations + +import hashlib +import json +import multiprocessing as mp +import os +import socket +import struct +import subprocess +import time +import traceback +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + +from free_ports import free_port +from yeto.swift_megatron import sync as sm + +ROOT = Path(__file__).resolve().parent.parent +H, V, F, E, LAYERS = 8, 12, 6, 4, 4 + + +# ----------------------------------------------------------------------------- pure helpers +def test_global_names_cover_layers_and_both_expert_namings(): + layers = sm._global_layer_prefixes(_chunk(tp=1, pp=1, tp_size=2, pp_size=2)) + experts = sm._expert_modules(_chunk(tp=1, pp=1, tp_size=2, pp_size=2)) + name = sm.global_parameter_name + assert name("decoder.layers.0.mixer.in_proj.weight", layers, experts) == "decoder.layers.2.mixer.in_proj.weight" + assert (name("decoder.layers.0.mlp.experts.local_experts.1.linear_fc1.weight", layers, experts) + == "decoder.layers.2.mlp.experts.local_experts.3.linear_fc1.weight") + assert (name("decoder.layers.1.mlp.experts.linear_fc1.weight0", layers, experts) + == "decoder.layers.3.mlp.experts.linear_fc1.weight2") + assert name("output_layer.weight", layers, experts) == "output_layer.weight" + assert name("decoder.layers.0.mlp.router.weight", layers, experts) == "decoder.layers.2.mlp.router.weight" + + +def test_syncer_address_and_hmac_key_file(tmp_path): + assert sm._syncer_addrs("tcp://10.0.0.1:29400") == [("10.0.0.1", 29400)] + assert sm._syncer_addrs("a:1, tcp://b:2") == [("a", 1), ("b", 2)] + key = tmp_path / "hmac.key" + key.write_text("s3cret\n") + assert sm._read_hmac_key(str(key)) == b"s3cret" # the YETO_ISLAND_HMAC_KEY encoding + assert sm._read_hmac_key(None) is None + + +def test_chunks_are_bounded_and_keep_order(): + sizes = {"a": 3, "b": 3, "c": 10, "d": 1, "e": 1} + assert sm.chunk_names(list(sizes), sizes.__getitem__, 6) == [["a", "b"], ["c"], ["d", "e"]] + + +# ----------------------------------------------------------------------------- fake Megatron island +def _full(name: str, shape) -> torch.Tensor: + """The initial full value of a global tensor (bf16-exact, so INIT is lossless).""" + + seed = int.from_bytes(hashlib.sha256(name.encode()).digest()[:4], "little") + g = torch.Generator().manual_seed(seed) + return (1 + 0.25 * torch.randn(*shape, generator=g)).to(torch.bfloat16).float() + + +def _update(name: str, shape, island: int, step: int) -> torch.Tensor: + """One inner step's update of a global tensor (a function of position only).""" + + seed = int.from_bytes(hashlib.sha256(f"{name}/{island}/{step}".encode()).digest()[:4], "little") + g = torch.Generator().manual_seed(seed) + sign = 1.0 if island == 0 else -1.0 + return sign * 1e-3 * (1 + torch.rand(*shape, generator=g)) + + +class _Module(torch.nn.Module): + pass + + +def _param(full_name, full_shape, *, tp, tp_size, part_dim=-1, expert=False, trainable=True): + full = _full(full_name, full_shape) + if part_dim >= 0: + size = full_shape[part_dim] // tp_size + full = full.narrow(part_dim, tp * size, size) + p = torch.nn.Parameter(full.to(torch.bfloat16).clone(), requires_grad=trainable) + if part_dim >= 0: + p.tensor_model_parallel, p.partition_dim, p.partition_stride = True, part_dim, 1 + if expert: + p.allreduce = False + p._full_name, p._full_shape = full_name, tuple(full_shape) + return p + + +def _chunk(*, tp: int, pp: int, tp_size: int, pp_size: int) -> torch.nn.Module: + """One rank's model chunk. EP spans the TP ranks with ETP 1 (each TP rank + holds E / tp_size whole experts), as in the production recipe.""" + + root = _Module() + per_stage = LAYERS // pp_size + if pp == 0: + root.embedding = _Module() + root.embedding.word_embeddings = _Module() + root.embedding.word_embeddings.weight = _param("embedding.word_embeddings.weight", (V, H), tp=tp, + tp_size=tp_size, part_dim=0) + root.vision_model = _Module() + root.vision_model.proj = _Module() + root.vision_model.proj.weight = _param("vision_model.proj.weight", (H, H), tp=tp, tp_size=tp_size, + trainable=False) + root.decoder = _Module() + root.decoder.layers = torch.nn.ModuleList() + local_experts = E // tp_size + for k in range(per_stage): + g = pp * per_stage + k + layer = _Module() + layer.layer_number = g + 1 + prefix = f"decoder.layers.{g}" + layer.mixer = _Module() + layer.mixer.in_proj = _Module() + layer.mixer.in_proj.weight = _param(f"{prefix}.mixer.in_proj.weight", (2 * F, H), tp=tp, + tp_size=tp_size, part_dim=0) + layer.mixer.out_proj = _Module() + layer.mixer.out_proj.weight = _param(f"{prefix}.mixer.out_proj.weight", (H, F), tp=tp, + tp_size=tp_size, part_dim=1) + layer.norm = _Module() + layer.norm.weight = _param(f"{prefix}.norm.weight", (H,), tp=tp, tp_size=tp_size) + layer.mlp = _Module() + layer.mlp.local_expert_indices = [tp * local_experts + e for e in range(local_experts)] # MoELayer + layer.mlp.router = _Module() + layer.mlp.router.weight = _param(f"{prefix}.mlp.router.weight", (E, H), tp=tp, tp_size=tp_size) + layer.mlp.router.register_buffer("expert_bias", torch.zeros(E)) # no master: not synced + experts = _Module() + experts.local_expert_indices = [tp * local_experts + e for e in range(local_experts)] + if g % 2 == 0: # SequentialMLP + experts.local_experts = torch.nn.ModuleList() + for e in range(local_experts): + mlp = _Module() + mlp.linear_fc1 = _Module() + mlp.linear_fc1.weight = _param( + f"{prefix}.mlp.experts.local_experts.{experts.local_expert_indices[e]}.linear_fc1.weight", + (F, H), tp=0, tp_size=1, expert=True) + experts.local_experts.append(mlp) + else: # TEGroupedMLP + experts.linear_fc1 = _Module() + for e in range(local_experts): + setattr(experts.linear_fc1, f"weight{e}", _param( + f"{prefix}.mlp.experts.linear_fc1.weight{experts.local_expert_indices[e]}", + (F, H), tp=0, tp_size=1, expert=True)) + layer.mlp.experts = experts + root.decoder.layers.append(layer) + if pp == pp_size - 1: + root.decoder.final_layernorm = _Module() + root.decoder.final_layernorm.weight = _param("decoder.final_layernorm.weight", (H,), tp=tp, + tp_size=tp_size) + root.output_layer = _Module() + root.output_layer.weight = _param("output_layer.weight", (V, H), tp=tp, tp_size=tp_size, part_dim=0) + return root + + +def _global_shapes(tp_size: int, pp_size: int) -> dict[str, tuple[int, ...]]: + """Every trainable global tensor of the island and its full shape.""" + + return {p._full_name: p._full_shape + for tp in range(tp_size) for pp in range(pp_size) + for p in _chunk(tp=tp, pp=pp, tp_size=tp_size, pp_size=pp_size).parameters() if p.requires_grad} + + +class _Range: + def __init__(self, start, end): + self.start, self.end = start, end + + +class _Hybrid: + """megatron-core ``HybridDeviceOptimizer`` as ``yeto.rl.megatron_masters`` reads it.""" + + def __init__(self, shards): + self.param_groups = [{"params": shards}] + self.param_update_in_fp32 = True + self.param_to_fp32_param = {s: s.detach().clone().float() for s in shards} + self.state = {s: {"exp_avg": torch.full((s.numel(),), 0.5), "exp_avg_sq": torch.full((s.numel(),), 0.25)} + for s in shards} + + +class _DistOpt: + """A DistributedOptimizer leaf (precision-aware, CPU offload): this DP rank + owns a contiguous range of the flattened local parameters.""" + + def __init__(self, params, *, dp_rank, dp_size, dp_group): + self.params = list(params) + total = sum(p.numel() for p in self.params) + per = -(-total // dp_size) + lo, hi = dp_rank * per, min((dp_rank + 1) * per, total) + param_map, shards, self.model_param_group_index_map = {}, [], {} + self.owned = {} + offset = 0 + for p in self.params: + s, e = max(lo, offset) - offset, min(hi, offset + p.numel()) - offset + if e > s: + param_map[p] = {"param": _Range(s, e)} + self.model_param_group_index_map[p] = (0, len(shards)) + shards.append(p.detach().view(-1)[s:e]) + self.owned[id(p)] = (s, e, len(shards) - 1) + offset += p.numel() + self.gbuf_ranges = [{(torch.bfloat16, torch.float32): [{"param_map": param_map}]}] + self.config = SimpleNamespace(use_precision_aware_optimizer_no_fp8_or_ds_fp8=True, bf16=True, + main_params_dtype=torch.float32, optimizer_cpu_offload=True, + store_param_remainders=True) + self.data_parallel_group = dp_group + self.optimizer = _Hybrid(shards) + + def master(self, p): + s, e, i = self.owned[id(p)] + return self.optimizer.param_to_fp32_param[self.optimizer.param_groups[0]["params"][i]] + + def step(self, local_updates): + """AdamW on the host FP32 copy, then the copy-back hook into the bf16 shard.""" + + for p in self.params: + if id(p) not in self.owned: + continue + s, e, i = self.owned[id(p)] + shard = self.optimizer.param_groups[0]["params"][i] + master = self.optimizer.param_to_fp32_param[shard] + master.add_(local_updates[id(p)].reshape(-1)[s:e]) + shard.copy_(master) + # The bf16 copy outside this rank's range stays stale until the next + # forward's param all-gather (overlap_param_gather): the sync must not read it. + + +def _local_update(p, island, step, tp, tp_size): + full = _update(p._full_name, p._full_shape, island, step) + if getattr(p, "tensor_model_parallel", False): + size = p._full_shape[p.partition_dim] // tp_size + full = full.narrow(p.partition_dim, tp * size, size) + return full + + +def _init_dist(rank, world, port): + import torch.distributed as dist + + os.environ["GLOO_SOCKET_IFNAME"] = "lo0" if os.uname().sysname == "Darwin" else "lo" + dist.init_process_group("gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=world) + return dist + + +def _run(target, world, *args, timeout=240): + ctx = mp.get_context("spawn") + procs = [ctx.Process(target=target, args=(rank, world, *args)) for rank in range(world)] + for p in procs: + p.start() + return procs + + +def _join(procs, timeout, out_dir): + deadline = time.monotonic() + timeout + for p in procs: + p.join(max(1.0, deadline - time.monotonic())) + alive = [p for p in procs if p.is_alive()] + for p in alive: + p.kill() + errors = sorted(Path(out_dir).glob("error-*.txt")) + detail = "\n".join(e.read_text() for e in errors) + assert not alive, f"{len(alive)} worker(s) hung\n{detail}" + assert all(p.exitcode == 0 for p in procs), f"worker exit codes {[p.exitcode for p in procs]}\n{detail}" + + +def _guard(out_dir, tag, fn): + try: + fn() + except BaseException: + Path(out_dir, f"error-{tag}.txt").write_text(f"[{tag}]\n{traceback.format_exc()}") + raise + + +# ----------------------------------------------------------------------------- scripted client (TP2 x DP2) +class _ScriptedClient: + """The SyncerClient surface the sync uses, driven by a per-iteration script.""" + + def __init__(self, layout, script): + import threading + + self.layout, self.script = layout, script + self.finalizing, self.shutdown, self.finalized = threading.Event(), threading.Event(), threading.Event() + self.inits, self.pushes, self.calls = {}, [], 0 + + def start(self): + pass + + def close(self): + pass + + def check_health(self): + self.calls += 1 + + def send_init_parts(self, fid, parts): + self.inits[fid] = b"".join(bytes(p) for p in parts) + return True + + def drain_updates(self): + return self.script.get(("bcast", self.calls), []) + + def drain_pulls(self): + return self.script.get(("pull", self.calls), []) + + def push_fragment_parts(self, fid, step, attempt, base, local_step, c_steps, c_units, parts): + self.pushes.append((fid, step, attempt, base, local_step, c_steps, c_units, b"".join(bytes(p) for p in parts))) + return True + + +def _flat(layout, fid, values): + return torch.cat([values[name].reshape(-1) for name, _ in layout.fragments[fid].tensors]) + + +def _scripted_worker(rank, world, port, out_dir): + _guard(out_dir, f"scripted-{rank}", lambda: _scripted_body(rank, world, port, out_dir)) + + +def _scripted_body(rank, world, port, out_dir): + from yeto.protocol import BcastFragment, PullRequest + from yeto.tensor_io import pack_tensor, unpack_fragment + + dist = _init_dist(rank, world, port) + tp, dp = rank % 2, rank // 2 + tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])] + dp_groups = [dist.new_group([0, 2]), dist.new_group([1, 3])] + singles = [dist.new_group([r]) for r in range(world)] + ctx = sm.ParallelContext(rank, world, tp_groups[dp], tp, 2, singles[rank], 0, 1) + chunk = _chunk(tp=tp, pp=0, tp_size=2, pp_size=1) + trainable = [p for p in chunk.parameters() if p.requires_grad] + opt = _DistOpt(trainable, dp_rank=dp, dp_size=2, dp_group=dp_groups[tp]) + full_shapes = _global_shapes(2, 1) + globals_g = {name: (_full(name, shape) + 0.125).to(torch.bfloat16).float() + for name, shape in full_shapes.items()} + holder = {} + alpha = 0.5 + + def factory(addrs, learner_id, layout, wire, streams, **kwargs): + assert addrs == [("127.0.0.1", 1234)] and len(kwargs["session_contract_hash"]) == 32 + script = {} + script[("bcast", 1)] = [BcastFragment(fid, 7, pack_tensor(_flat(layout, fid, globals_g), 2)) + for fid in range(layout.num_fragments)] + # Pulls of every fragment arrive with the broadcast (so must wait for two steps) + script[("pull", 1)] = [PullRequest(fid, 8, 1) for fid in range(layout.num_fragments)] + holder["client"] = client = _ScriptedClient(layout, script) + return client + + frozen_before = chunk.vision_model.proj.weight.detach().clone() + syncer = sm.MegatronDiLoCoSync([chunk], opt, syncer="tcp://127.0.0.1:1234", learner_id=0, num_learners=2, + fragments=3, sync_interval_steps=2, merge_alpha=alpha, chunk_numel=40, + parallel_context=ctx, client_factory=factory) + names = set(syncer.layout.tensor_names()) + assert names == set(full_shapes) and "vision_model.proj.weight" not in names + assert syncer.layout.fragments[0].tensors[0][0] in ("embedding.word_embeddings.weight", "output_layer.weight") + if rank == 0: # INIT is the current FP32 masters, gathered over TP and DP + client = holder["client"] + for fid, frag in enumerate(syncer.layout.fragments): + got = unpack_fragment(frag, client.inits[fid], 2) + want = _flat(syncer.layout, fid, {n: _full(n, s) for n, s in full_shapes.items()}) + assert torch.equal(got, want), f"INIT fragment {fid}" + + masters_before = {id(p): opt.master(p).clone() for p in trainable if id(p) in opt.owned} + adam_before = {k: {n: t.clone() for n, t in v.items()} for k, v in opt.optimizer.state.items()} + updates_total = {name: torch.zeros(shape) for name, shape in full_shapes.items()} + applied = {} + for it in (1, 2, 3): + opt.step({id(p): _local_update(p, 0, it, tp, 2) for p in trainable}) + for name, shape in full_shapes.items(): + if it >= 2: # updates after the merge at iteration 1 + updates_total[name] += _update(name, shape, 0, it) + if it == 1: + pre_merge = {name: _full(name, shape) + _update(name, shape, 0, 1) + for name, shape in full_shapes.items()} + syncer.after_optimizer_step(it, 100 * it) + if it == 1: + # merged with alpha: master = a * local + (1 - a) * global, bf16 copy = master + for p in trainable: + full = alpha * pre_merge[p._full_name] + (1 - alpha) * globals_g[p._full_name] + if getattr(p, "tensor_model_parallel", False): + size = p._full_shape[p.partition_dim] // 2 + full = full.narrow(p.partition_dim, tp * size, size) + applied[p._full_name] = full + if id(p) in opt.owned: + s, e, _ = opt.owned[id(p)] + assert torch.allclose(opt.master(p), full.reshape(-1)[s:e], atol=1e-6), p._full_name + assert not torch.equal(opt.master(p), masters_before[id(p)]) + assert torch.equal(p.detach(), full.to(torch.bfloat16)), p._full_name + assert syncer.anchored == [True] * syncer.num_fragments + assert [e[1] for e in syncer.events] == ["apply"] * syncer.num_fragments + if it == 2: # one step since the merge: the pulls wait (H = 2) + assert not any(e[1] == "push" for e in syncer.events) + pushes = [e for e in syncer.events if e[1] == "push"] + assert [e[2] for e in pushes] == list(range(syncer.num_fragments)) and all(e[0] == 3 for e in pushes) + assert torch.equal(chunk.vision_model.proj.weight, frozen_before) + for key, state in opt.optimizer.state.items(): # Adam moments kept + for n, t in state.items(): + assert torch.equal(t, adam_before[key][n]) + if rank == 0: + client = holder["client"] + assert len(client.pushes) == syncer.num_fragments + for fid, step, attempt, base, local_step, c_steps, c_units, payload in client.pushes: + assert (step, attempt, base, local_step, c_steps, c_units) == (8, 1, 7, 3, 2, 500) + # PUSH = local - raw anchor (the global, not the blend) + local = {n: alpha * pre_merge[n] + (1 - alpha) * globals_g[n] + updates_total[n] for n in full_shapes} + want = _flat(syncer.layout, fid, local) - _flat(syncer.layout, fid, globals_g) + got = unpack_fragment(syncer.layout.fragments[fid], payload, 2) + assert torch.allclose(got, want.to(torch.bfloat16).float(), atol=1e-3), fid + syncer.close() + syncer.close() + dist.destroy_process_group() + + +@pytest.mark.timeout(240) +def test_tp2_dp2_island_reads_masters_writes_masters_and_pushes_raw_anchor_deltas(tmp_path): + procs = _run(_scripted_worker, 4, free_port(), str(tmp_path)) + _join(procs, 200, tmp_path) + + +# ----------------------------------------------------------------------------- two islands, real syncer (TP2 x PP2) +def build_syncer() -> Path: + binary = ROOT / "syncer/target/debug/yeto-syncer" + env = dict(os.environ) + env["PATH"] = f"{Path.home() / '.cargo/bin'}:{env.get('PATH', '')}" + subprocess.run(["cargo", "build", "-q"], cwd=ROOT / "syncer", check=True, env=env, capture_output=True) + assert binary.exists() + return binary + + +def read_final_state(path: Path) -> list[torch.Tensor]: + data = path.read_bytes() + (count,) = struct.unpack_from(" 150: + raise TimeoutError(f"island {island} not finalized after {it} steps") + it += 1 + opt.step({id(p): _local_update(p, island, it, tp, 2) for p in trainable}) + if rank == 0 and drop_at is not None and not dropped and syncer.stats["pushed"] >= drop_at: + for sock in list(getattr(syncer.client, "_socks", [])): # a WAN drop: the client redials + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + dropped = True + syncer.after_optimizer_step(it, 64) + time.sleep(0.01) + for p in trainable: # the bf16 copy equals the installed FP32 master + if id(p) in opt.owned: + assert torch.equal(p.detach().reshape(-1), opt.master(p).to(torch.bfloat16)), p._full_name + names = syncer.layout.tensor_names() + values = syncer._export_chunk(names, delta=False) + if rank == 0: + torch.save({"names": names, "values": dict(zip(names, values, strict=True)), + "events": list(syncer.events), "stats": dict(syncer.stats), "dropped": dropped, + "fragments": [[n for n, _ in f.tensors] for f in syncer.layout.fragments]}, + Path(out_dir, f"island{island}.pt")) + syncer.close() + dist.destroy_process_group() + + +@pytest.mark.timeout(300) +def test_two_tp2_pp2_islands_merge_to_the_syncer_final_state_through_a_connection_drop(tmp_path): + binary = build_syncer() + from process_groups import group_popen + + port = free_port() + final_state = tmp_path / "final.bin" + total_steps = 15 + proc = group_popen([str(binary), "--port", str(port), "--learners", "2", "--quorum", "2", + "--grace-ms", "300", "--total-steps", str(total_steps), "--outer-lr", "0.7", + "--outer-momentum", "0.9", "--final-state", str(final_state)], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + procs = [] + try: + procs += _run(_island_worker, 4, 0, free_port(), port, str(tmp_path), None) + procs += _run(_island_worker, 4, 1, free_port(), port, str(tmp_path), 2) + _join(procs, 240, tmp_path) + assert proc.wait(timeout=60) == 0, "syncer exited nonzero" + finally: + for p in procs: + if p.is_alive(): + p.kill() + if proc.poll() is None: + proc.kill() + print((proc.stdout.read() if proc.stdout else "")[-4000:]) + + islands = [torch.load(tmp_path / f"island{i}.pt") for i in range(2)] + assert islands[1]["dropped"], "the connection drop never happened" + names = islands[0]["names"] + assert names == islands[1]["names"] and "vision_model.proj.weight" not in names + for name in names: # both islands hold the same final global, in FP32 masters + assert torch.equal(islands[0]["values"][name], islands[1]["values"][name]), name + final = read_final_state(final_state) + for fid, frag_names in enumerate(islands[0]["fragments"]): + flat = torch.cat([islands[0]["values"][n].reshape(-1) for n in frag_names]) + assert torch.equal(flat, final[fid]), f"fragment {fid} differs from the syncer's final state" + # The model moved: the final global is not the initial value. + moved = sum(float((islands[0]["values"][n] - _full(n, islands[0]["values"][n].shape)).abs().sum()) + for n in names) + assert moved > 0 + for island in islands: # fragment rotation: every fragment pushed and merged more than once + pushes = [e[2] for e in island["events"] if e[1] == "push"] + applies = [e[2] for e in island["events"] if e[1] == "apply"] + for fid in range(len(island["fragments"])): + assert pushes.count(fid) >= 2 and applies.count(fid) >= 2, (fid, pushes, applies) + versions = {} + for _it, kind, fid, version in island["events"]: + if kind == "apply": + assert version >= versions.get(fid, -1), "fragment versions went backwards" + versions[fid] = version + print(json.dumps([i["stats"] for i in islands])) + + +def test_write_masters_also_sets_a_distinct_hdo_host_copy(): + """core 0.18.2 makes the offloaded host copy the FP32 master itself; a + build that kept them apart would copy the stale host copy back over the + shard after the next step, so the write covers both.""" + + from yeto.rl import megatron_masters as mm + + params = [torch.nn.Parameter(torch.ones(4, dtype=torch.bfloat16))] + opt = _DistOpt(params, dp_rank=0, dp_size=1, dp_group=None) + shard = opt.optimizer.param_groups[0]["params"][0] + opt.optimizer.gpu_params_map_cpu_copy = {shard: torch.zeros(4)} + target = torch.full((4,), 2.5) + mm.write_masters(opt, params, [target]) + assert torch.equal(opt.master(params[0]), target) + assert torch.equal(opt.optimizer.gpu_params_map_cpu_copy[shard], target) + assert torch.equal(params[0].detach(), target.to(torch.bfloat16)) diff --git a/yeto/rl/megatron_masters.py b/yeto/rl/megatron_masters.py index cd318e0f..e4e69ac0 100644 --- a/yeto/rl/megatron_masters.py +++ b/yeto/rl/megatron_masters.py @@ -198,6 +198,13 @@ def _write_main_shard(leaf: Any, param: Any, value: Any) -> None: if _hybrid(leaf): master = leaf.optimizer.param_to_fp32_param[shard] master.data.copy_(value.reshape(master.shape).to(master.dtype)) + # Under offload the host copy the CPU sub-optimizer steps is this same + # FP32 tensor (core 0.18.2 rebinds it before recording the copy); a + # distinct host copy would be copied back over the shard after the next + # step, so it is written too. + host = getattr(leaf.optimizer, "gpu_params_map_cpu_copy", {}).get(shard) + if host is not None and host is not master: + host.data.copy_(value.reshape(host.shape).to(host.dtype)) return leaf.optimizer.set_scaled_state(shard, "master_param", value.reshape(shard.shape).float().clone()) diff --git a/yeto/swift_megatron/__init__.py b/yeto/swift_megatron/__init__.py new file mode 100644 index 00000000..1ed0d217 --- /dev/null +++ b/yeto/swift_megatron/__init__.py @@ -0,0 +1 @@ +"""yeto full-parameter DiLoCo for ms-swift Megatron training.""" diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py new file mode 100644 index 00000000..c7d56740 --- /dev/null +++ b/yeto/swift_megatron/sync.py @@ -0,0 +1,1021 @@ +"""Full-parameter DiLoCo sync of an ms-swift Megatron island. + +ms-swift trains the model with Megatron-Core (TP/EP/PP, DistributedOptimizer, +optionally CPU optimizer offload through ``HybridDeviceOptimizer``). This +module lets such an island speak the torch learner's DiLoCo protocol to the +unchanged Rust syncer (``yeto/learner.py`` + ``yeto/diloco_sync.py``): + +* the same fragment layout semantics (``yeto.fragments.build_layout`` over the + island's *global* trainable parameter set: whole tensors, size-balanced or + strided, embeddings in the direct-averaged fragment); +* the same PUSH: ``local - raw_anchor`` per fragment, where ``raw_anchor`` is + the exact global value of the last accepted broadcast (before any + ``merge_alpha`` blend), answered only once the fragment has been anchored + and trained on; +* the same BCAST handling: broadcasts are applied before pulls are answered, + reset the fragment's counters and adopt the syncer's version; +* the same INIT (learner 0 seeds every fragment), finalization (raw final cut, + then FINAL_ACK) and HELLO (layout, HMAC key). + +What differs is only how a fragment is read and written, because no rank holds +a whole tensor: + +Parameter namespace + Every trainable Megatron parameter is named by its *global* Megatron name: + the pipeline-local layer index is replaced by the global one + (``layer_number - 1``) and an expert's local index by its global index + (``local_expert_indices``). So the name set and every tensor's full shape + are independent of PP and EP, and the TP shards of one parameter + concatenate along ``partition_dim`` into the full tensor. Parameters the + trainer does not train (``trainable_filter``: frozen vision tower and + aligner) and buffers (router ``expert_bias``) are not part of the layout. + + The names are Megatron's, not the HF checkpoint's: the syncer is + name-agnostic, and Megatron-Bridge (which ``yeto.rl.megatron_hf_boundary`` + builds on) is not what ms-swift uses (it uses ``mcore_bridge``, whose + export is a whole-model pass with no per-parameter selection). Fused + Megatron tensors whose element order depends on the TP size (Mamba + ``in_proj``, gated ``linear_fc1``) make the layout TP-specific, so the + session contract hash binds the TP / expert-TP sizes and the syncer + refuses an island of a different parallel shape. + +Read (PUSH, INIT) + Collective, one bounded chunk of tensors at a time: every rank gathers the + FP32 optimizer masters of its local parameters over the data-parallel group + (``yeto.rl.megatron_masters.full_masters``: DistributedOptimizer shards; with + the precision-aware optimizer + CPU offload of the production recipe that is + ``HybridDeviceOptimizer.param_to_fp32_param[bf16 shard]``, a pinned host + tensor, looked up afresh on every call because HDO rebuilds its maps on + ``load_state_dict``), subtracts its local anchor shard, all-gathers TP + shards, and the lowest owning rank sends the full FP32 tensor + point-to-point to global rank 0, which packs it into the wire buffer. The + bf16 model copy is never read: with ``overlap_param_gather`` its non-owned + ranges are stale after a step. The PUSH carries ``local - anchor``, the + torch learner's sign (the syncer's outer gradient is its negation, ``base - + local``). + +Write (BCAST, final cut) + Rank 0 broadcasts the chunk over the island in the wire dtype (bf16 by + default, lossless for a bf16 session); every rank slices its TP shard, + records it as its anchor, blends with its current master when + ``merge_alpha > 0``, and writes the FP32 master shard *and* the bf16 model + copy (``yeto.rl.megatron_masters.write_masters``: the HDO FP32 tensor, which + under offload is also the host copy the CPU Adam steps, and the parameter's + storage inside the DDP ``param_data`` buffer, from which the next forward's + param all-gather reads). The Adam moments are kept, exactly as the torch + learner keeps its AdamW state when ``apply_fragment`` overwrites a fragment. + +Only global rank 0 talks to the syncer. Every boundary starts with one +``broadcast_object_list`` of rank 0's plan (fragments to apply, fragments to +push, finalization, shutdown, error), so every rank runs the same collectives +in the same order and an error on rank 0 is raised on every rank. + +Memory (N = trainable parameters of the island, P = fragments, n_r = the +elements of rank r's local parameter shards, C = ``chunk_numel``, t = TP size): + +* device: nothing with the default ``transport="cpu"`` (the collectives run + on host buffers over gloo mirrors of the island's groups; the only device + access is the in-place host-to-device copy into the existing bf16 + parameters). The production pp0 stage has ~0.4 GiB of HBM headroom, so + ``transport="device"`` (chunk buffers on the accelerator, peak about + ``4 * C * (2 + t)`` bytes) is only for islands with room to spare; +* every rank, host: one anchor per local shard in the anchor dtype (bf16 for a + bf16 session), ``2 * n_r`` bytes, plus one chunk in flight, about + ``4 * C * (2 + t)`` bytes (C = 2**28 -> ~6 GiB at t = 4; a single tensor + larger than C is a chunk of its own). For the TP4/EP8/PP3 (world 24) + Nemotron-3.5-Super recipe (~120B, MoE-dominated, ETP 1) n_r is about + 120e9 / 24 = 5e9: ~10 GB of anchors per rank, ~80 GB plus ~50 GB of chunks + per 8-card host, beside the ~450 GB the offloaded masters and Adam moments + take of its ~2 TB; +* rank 0, host additionally: the wire bytes of one pushed fragment + (``2 * N / P`` for bf16) and the client's received broadcasts (each + ``2 * N / P``; at most the syncer's ``--pipeline`` rounds in flight). For + N = 120e9 and P = 16 that is ~15 GB per fragment, ~45-60 GB peak. A + full-payload final cut holds every fragment in FP32 (``4 * N`` = 480 GB) on + rank 0's host: run the syncer with ``--final-cut-payload manifest`` and + ``final_cut_payload="manifest"`` for a 120B island. + +Torch, Megatron-Core and the syncer client are imported lazily by the +objects that need them. +""" + +from __future__ import annotations + +import hashlib +import logging +import re +from collections import deque +from collections.abc import Callable, Iterable, Sequence +from dataclasses import dataclass, field +from typing import Any + +log = logging.getLogger("swift-megatron-sync") + +# Megatron's tied/untied output projection is an embedding-like tensor for the +# direct-averaged fragment, as ``lm_head`` is for the torch learner. +OUTPUT_LAYER_AVG_REGEX = r"(^|\.)output_layer\.weight$" +DEFAULT_CHUNK_NUMEL = 1 << 28 +CONTRACT_DOMAIN = b"yeto-swift-megatron-native-v1\0" +WIRE_DTYPES = ("bf16", "f32", "fp32", "q4") +FINAL_CUT_PAYLOADS = ("full", "manifest") + +_LAYER_MODULE_RE = re.compile(r"(^|\.)decoder\.layers\.(\d+)$") +_GROUPED_EXPERT_RE = re.compile(r"^(.*?)(weight|bias)(\d+)$") + + +class SwiftMegatronSyncError(RuntimeError): + """The Megatron island cannot run the full-parameter DiLoCo sync as configured.""" + + +# ----------------------------------------------------------------------------- parallel context +@dataclass(frozen=True) +class ParallelContext: + """The process groups the sync needs, decoupled from ``parallel_state`` + so the sync runs on any group layout (tests build gloo subgroups). + + ``tp_*`` is the tensor-parallel group of dense parameters, ``etp_*`` the + expert-tensor-parallel group of expert parameters (``param.allreduce is + False``). The data-parallel groups come from the optimizer + (``data_parallel_group`` of each DistributedOptimizer leaf).""" + + rank: int + world: int + tp_group: Any + tp_rank: int + tp_size: int + etp_group: Any + etp_rank: int + etp_size: int + + @classmethod + def from_megatron(cls) -> ParallelContext: + import torch.distributed as dist + from megatron.core import parallel_state as mpu + + def expert_tp(name: str, fallback: Callable[[], Any]) -> Any: + getter = getattr(mpu, name, None) + return getter() if getter is not None else fallback() + + return cls( + rank=int(dist.get_rank()), + world=int(dist.get_world_size()), + tp_group=mpu.get_tensor_model_parallel_group(), + tp_rank=int(mpu.get_tensor_model_parallel_rank()), + tp_size=int(mpu.get_tensor_model_parallel_world_size()), + etp_group=expert_tp("get_expert_tensor_parallel_group", mpu.get_tensor_model_parallel_group), + etp_rank=int(expert_tp("get_expert_tensor_parallel_rank", mpu.get_tensor_model_parallel_rank)), + etp_size=int(expert_tp("get_expert_tensor_parallel_world_size", + mpu.get_tensor_model_parallel_world_size)), + ) + + +# ----------------------------------------------------------------------------- global names +@dataclass +class LocalShard: + """One local trainable parameter and where it sits in the global tensor.""" + + name: str # global Megatron name + param: Any + expert: bool + part_dim: int # TP partition dim, -1 when replicated over the TP group + tp_rank: int + tp_size: int + tp_group: Any = field(repr=False, default=None) + + @property + def local_shape(self) -> tuple[int, ...]: + return tuple(int(d) for d in self.param.shape) + + def full_shape(self) -> tuple[int, ...]: + shape = list(self.local_shape) + if self.part_dim >= 0: + shape[self.part_dim] *= self.tp_size + return tuple(shape) + + +def unwrap_chunk(chunk: Any) -> Any: + """The bare Megatron model of a (DDP / Float16Module wrapped) model chunk.""" + + seen = 0 + while hasattr(chunk, "module") and seen < 8: + chunk = chunk.module + seen += 1 + return chunk + + +def _global_layer_prefixes(chunk: Any) -> list[tuple[str, str]]: + """``(local prefix, global prefix)`` for every decoder layer of ``chunk``, + longest first. Megatron numbers ``decoder.layers`` per pipeline stage (and + virtual stage); ``layer_number`` is the 1-based global index.""" + + out = [] + for name, module in chunk.named_modules(): + match = _LAYER_MODULE_RE.search(name) + if match is None or not hasattr(module, "layer_number"): + continue + global_name = name[: match.start(2)] + str(int(module.layer_number) - 1) + out.append((name + ".", global_name + ".")) + return sorted(out, key=lambda item: -len(item[0])) + + +def _expert_modules(chunk: Any) -> list[tuple[str, list[int]]]: + """``(prefix, local_expert_indices)`` of every expert container (the + ``experts`` child of a MoE layer; the MoE layer itself also carries + ``local_expert_indices`` but holds the router and shared experts), + longest first.""" + + out = [] + for name, module in chunk.named_modules(): + indices = getattr(module, "local_expert_indices", None) + if indices is None or name.rsplit(".", 1)[-1] != "experts": + continue + out.append((name + ".", [int(i) for i in indices])) + return sorted(out, key=lambda item: -len(item[0])) + + +def _globalize_expert(rest: str, indices: Sequence[int], name: str) -> str: + """The global form of ``rest`` (the name below an expert container).""" + + head, _, tail = rest.partition(".") + if head == "local_experts": # SequentialMLP: local_experts..<...> + local, _, below = tail.partition(".") + if not local.isdigit(): + raise SwiftMegatronSyncError(f"unrecognised expert parameter {name!r}") + return f"local_experts.{indices[int(local)]}.{below}" + match = _GROUPED_EXPERT_RE.match(rest) + if match is not None: # TEGroupedMLP: linear_fc1.weight + return f"{match.group(1)}{match.group(2)}{indices[int(match.group(3))]}" + raise SwiftMegatronSyncError( + f"expert parameter {name!r} stacks several experts in one tensor (legacy GroupedMLP); " + "use TEGroupedMLP or SequentialMLP") + + +def global_parameter_name(local_name: str, layers: Sequence[tuple[str, str]], + experts: Sequence[tuple[str, list[int]]]) -> str: + """``local_name`` with its layer and expert indices made global.""" + + name = local_name + for prefix, indices in experts: + if name.startswith(prefix): + name = prefix + _globalize_expert(name[len(prefix):], indices, local_name) + break + for local, global_prefix in layers: + if name.startswith(local): + name = global_prefix + name[len(local):] + break + return name + + +def default_trainable_filter(name: str, param: Any) -> bool: + return bool(getattr(param, "requires_grad", False)) + + +def local_shards(model_chunks: Sequence[Any], ctx: ParallelContext, + trainable_filter: Callable[[str, Any], bool] | None = None) -> list[LocalShard]: + """This rank's trainable parameters under their global names (sorted).""" + + keep = trainable_filter or default_trainable_filter + out: dict[str, LocalShard] = {} + seen: set[int] = set() + for wrapped in model_chunks: + chunk = unwrap_chunk(wrapped) + layers = _global_layer_prefixes(chunk) + experts = _expert_modules(chunk) + for local_name, param in chunk.named_parameters(): + if id(param) in seen or not keep(local_name, param): + continue + seen.add(id(param)) + name = global_parameter_name(local_name, layers, experts) + if name in out: + raise SwiftMegatronSyncError(f"two local parameters map to the global name {name!r}") + expert = getattr(param, "allreduce", True) is False + tp_rank, tp_size, group = ((ctx.etp_rank, ctx.etp_size, ctx.etp_group) if expert + else (ctx.tp_rank, ctx.tp_size, ctx.tp_group)) + part_dim = -1 + if getattr(param, "tensor_model_parallel", False) and tp_size > 1: + if int(getattr(param, "partition_stride", 1)) != 1: + raise SwiftMegatronSyncError(f"{name}: strided TP partitions are not supported") + part_dim = int(param.partition_dim) + out[name] = LocalShard(name, param, expert, part_dim, tp_rank, tp_size, group) + return [out[name] for name in sorted(out)] + + +@dataclass(frozen=True) +class GlobalTable: + """Every global trainable tensor: its full shape and the rank that sends it + (the lowest rank holding it at TP rank 0; pipeline copies of a tied + embedding hold equal values, so any holder will do).""" + + shapes: dict[str, tuple[int, ...]] + owners: dict[str, int] + + @classmethod + def build(cls, shards: Sequence[LocalShard], group: Any = None) -> GlobalTable: + import torch.distributed as dist + + rows = [(s.name, s.full_shape(), s.tp_rank) for s in shards] + world = int(dist.get_world_size()) if dist.is_initialized() else 1 + if world > 1: + gathered: list[Any] = [None] * world + dist.all_gather_object(gathered, rows, group=group) + else: + gathered = [rows] + shapes: dict[str, tuple[int, ...]] = {} + owners: dict[str, int] = {} + for holder, held in enumerate(gathered): + for name, shape, tp_rank in held: + shape = tuple(int(d) for d in shape) + if shapes.setdefault(name, shape) != shape: + raise SwiftMegatronSyncError( + f"{name}: ranks disagree on the full shape ({shapes[name]} vs {shape})") + if int(tp_rank) == 0 and (name not in owners or holder < owners[name]): + owners[name] = holder + missing = sorted(set(shapes) - set(owners)) + if missing: + raise SwiftMegatronSyncError(f"no TP-rank-0 holder for {missing[:4]}") + return cls(shapes, owners) + + def numel(self, name: str) -> int: + n = 1 + for d in self.shapes[name]: + n *= int(d) + return n + + +# ----------------------------------------------------------------------------- transport +TRANSPORTS = ("cpu", "device") + + +class Transport: + """Where the sync's collectives run. + + ``"cpu"`` (default): host buffers over gloo mirrors of the island's groups + (world, TP, expert-TP and every DistributedOptimizer leaf's data-parallel + group), so the sync allocates nothing on the accelerator: the FP32 + masters already live on the host under CPU offload, and the only device + traffic is the in-place host-to-device copy into the bf16 parameters. + ``"device"``: the island's own (HCCL/NCCL) groups with accelerator buffers + of at most one chunk, for islands with HBM to spare. + + Mirrors are created collectively (every rank creates every group, in one + sorted order). A gloo island needs no mirror.""" + + def __init__(self, kind: str, ctx: ParallelContext, optimizer: Any, device: Any) -> None: + import torch + import torch.distributed as dist + + if kind not in TRANSPORTS: + raise ValueError(f"transport must be one of {TRANSPORTS}, got {kind!r}") + self.kind = kind + self.buffer_device = torch.device("cpu") if kind == "cpu" else device + self.mirrors: dict[tuple[int, ...], Any] | None = None + self.world_ranks = tuple(range(ctx.world)) + if kind == "cpu" and dist.get_backend() != "gloo": + from yeto.rl.megatron_masters import _optimizer_leaves + + groups = [None, ctx.tp_group, ctx.etp_group] + groups += [getattr(leaf, "data_parallel_group", None) for leaf in _optimizer_leaves(optimizer)] + mine = sorted({self._ranks(group) for group in groups}) + gathered: list[Any] = [None] * ctx.world + dist.all_gather_object(gathered, mine) + self.mirrors = {} + for ranks in sorted({tuple(r) for held in gathered for r in held}): + group = dist.new_group(list(ranks), backend="gloo") + if ctx.rank in ranks: + self.mirrors[ranks] = group + + def _ranks(self, group: Any) -> tuple[int, ...]: + import torch.distributed as dist + + if group is None or group is dist.GroupMember.WORLD: + return self.world_ranks + return tuple(int(r) for r in dist.get_process_group_ranks(group)) + + def group(self, group: Any) -> Any: + """The group to run a collective of ``group`` on.""" + + if self.mirrors is None: + return group + return self.mirrors[self._ranks(group)] + + @property + def world(self) -> Any: + return self.group(None) + + def dp_reduce(self) -> Callable[[Any, Any], None] | None: + """``full_masters``' data-parallel reduce on this transport (None: its own).""" + + if self.kind == "device": + return None + import torch.distributed as dist + + def reduce(flat: Any, leaf: Any) -> None: + group = self.group(getattr(leaf, "data_parallel_group", None)) + if len(self._ranks(getattr(leaf, "data_parallel_group", None))) > 1: + dist.all_reduce(flat, op=dist.ReduceOp.SUM, group=group) + + return reduce + + +# ----------------------------------------------------------------------------- helpers +def _wire_code(wire_dtype: str) -> int: + from yeto.protocol import DTYPE_BF16, DTYPE_F32, DTYPE_Q4 + + try: + return {"bf16": DTYPE_BF16, "f32": DTYPE_F32, "fp32": DTYPE_F32, "q4": DTYPE_Q4}[wire_dtype] + except KeyError: + raise ValueError(f"wire_dtype must be one of {WIRE_DTYPES}, got {wire_dtype!r}") from None + + +def _torch_wire(code: int) -> Any: + import torch + + from yeto.protocol import DTYPE_BF16, DTYPE_F32 + + return {DTYPE_F32: torch.float32, DTYPE_BF16: torch.bfloat16}[code] + + +def _syncer_addrs(spec: str) -> list[tuple[str, int]]: + from yeto.syncer_shards import parse_syncer_addrs + + items = [item.strip() for item in spec.split(",")] + return parse_syncer_addrs(",".join(item.split("://", 1)[-1] for item in items)) + + +def _read_hmac_key(path: str | None) -> bytes | None: + """The island HMAC key file, with the encoding of ``YETO_ISLAND_HMAC_KEY``.""" + + if path is None: + return None + with open(path, encoding="utf-8") as handle: + value = handle.read().strip() + if not value: + raise SwiftMegatronSyncError(f"HMAC key file {path} is empty") + return value.encode() + + +def chunk_names(names: Sequence[str], numel: Callable[[str], int], budget: int) -> list[list[str]]: + """``names`` in order, cut into runs of at most ``budget`` elements (a + larger tensor is a run of its own). Pure, so every rank cuts alike.""" + + out: list[list[str]] = [] + run: list[str] = [] + total = 0 + for name in names: + n = numel(name) + if run and total + n > budget: + out.append(run) + run, total = [], 0 + run.append(name) + total += n + if run: + out.append(run) + return out + + +@dataclass +class _Plan: + """Rank 0's decision for one boundary, broadcast to every rank.""" + + applies: list[tuple[int, int]] = field(default_factory=list) # (fragment, version) + pushes: list[tuple[int, int, int, int]] = field(default_factory=list) # (fragment, step, attempt, c_units) + finalizing: bool = False + shutdown: bool = False + error: str | None = None + + +# ----------------------------------------------------------------------------- the sync +class MegatronDiLoCoSync: + """DiLoCo full-parameter sync of one ms-swift Megatron island (module doc). + + Construct it on every rank after the model, optimizer and resume checkpoint + are loaded; it connects rank 0 to the syncer, and learner 0 seeds every + fragment from its current FP32 masters (INIT). Then call + :meth:`after_optimizer_step` on every rank after every ``optimizer.step()``. + + ``model_chunks``: the model chunks (DDP-wrapped or bare; several with + virtual pipeline parallelism). ``optimizer``: the Megatron optimizer + (DistributedOptimizer, ChainedOptimizer of dense + expert, precision-aware, + CPU offload). ``trainable_filter(local_name, param)``: which parameters + take part (default ``requires_grad``). + + ``sync_interval_steps`` (H): a pull for a fragment is answered once the + island has taken H optimizer steps since that fragment was last merged + (H = 1 is the torch learner's self-clock: at least one step). The steps + are counted from the absolute ``iteration`` the trainer passes, so a + resumed run keeps its rounds. + + ``merge_alpha``: local weight of a broadcast, as the torch learner's + ``--merge-alpha`` (0 overwrites; the anchor is always the raw global). + ``final_cut_payload``: what the syncer's ``--final-cut-payload`` sends; + ``"full"`` overwrites the island with the final global, ``"manifest"`` + only acknowledges it (the final parameters stay with the syncer). + """ + + def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, + num_learners: int, fragments: int, sync_interval_steps: int, + wire_dtype: str = "bf16", hmac_key_path: str | None = None, + trainable_filter=None, merge_alpha: float = 0.0, + fragment_pattern: str = "binpack", wan_streams: int = 4, + final_cut_payload: str = "full", chunk_numel: int = DEFAULT_CHUNK_NUMEL, + transport: str = "cpu", parallel_context: ParallelContext | None = None, + client_factory=None) -> None: + import torch + import torch.distributed as dist + + from yeto.fragments import build_layout + from yeto.protocol import bulk_dtype, layout_fingerprint + + if not 0 <= merge_alpha < 1: + raise ValueError("merge_alpha must be in [0, 1)") + if sync_interval_steps < 1 or fragments < 1 or num_learners < 1: + raise ValueError("fragments, num_learners and sync_interval_steps must be >= 1") + if not 0 <= learner_id < num_learners: + raise ValueError(f"learner_id {learner_id} outside [0, {num_learners})") + if final_cut_payload not in FINAL_CUT_PAYLOADS: + raise ValueError(f"final_cut_payload must be one of {FINAL_CUT_PAYLOADS}") + if not dist.is_initialized(): + raise SwiftMegatronSyncError("torch.distributed must be initialized (Megatron does this)") + self.ctx = parallel_context or ParallelContext.from_megatron() + self.rank, self.world = self.ctx.rank, self.ctx.world + self.chunks = list(model_chunks) + self.modules = [unwrap_chunk(chunk) for chunk in self.chunks] + self.optimizer = optimizer + self.learner_id = int(learner_id) + self.interval = int(sync_interval_steps) + self.merge_alpha = float(merge_alpha) + self.final_cut_payload = final_cut_payload + self.chunk_numel = int(chunk_numel) + self.wire = _wire_code(wire_dtype) + self.bulk = bulk_dtype(self.wire) + self.bulk_torch = _torch_wire(self.bulk) + self.anchor_dtype = self.bulk_torch # the raw global is exact in the broadcast dtype + + self.shards = local_shards(self.chunks, self.ctx, trainable_filter) + self.by_name: dict[str, list[LocalShard]] = {} + for shard in self.shards: + self.by_name.setdefault(shard.name, []).append(shard) + devices = {shard.param.device for shard in self.shards} + self.device = next(iter(devices)) if devices else torch.device("cpu") + if dist.get_backend() != "gloo" and self.device.type == "cpu": + # A rank with no trainable shard still joins device collectives. + self.device = _current_accelerator() + self.comm = Transport(transport, self.ctx, optimizer, self.device) + self.table = GlobalTable.build(self.shards, self.comm.world) + names = sorted(self.table.shapes) + if not names: + raise SwiftMegatronSyncError("no trainable parameter to synchronize") + self.layout = build_layout( + [(name, self.table.numel(name)) for name in names], int(fragments), fragment_pattern, + avg_name_regex=OUTPUT_LAYER_AVG_REGEX, + named_shapes={name: self.table.shapes[name] for name in names}) + self.num_fragments = self.layout.num_fragments + self.fragment_chunks = [ + chunk_names([name for name, _ in frag.tensors], self.table.numel, self.chunk_numel) + for frag in self.layout.fragments] + self.contract = hashlib.sha256( + CONTRACT_DOMAIN + layout_fingerprint(self.layout) + + f"tp={self.ctx.tp_size};etp={self.ctx.etp_size}".encode()).digest() + + self.anchors: dict[int, Any] = {} # id(param) -> raw global shard (host) + self.anchored = [False] * self.num_fragments + self.versions = [0] * self.num_fragments + self.steps_at_reset: list[int] | None = None + self.units_at_reset = [0] * self.num_fragments + self.units_total = 0 + self.iteration = 0 + self.global_step = 0 + self.pending_pulls: list[Any] = [] + self._pull_meta: dict[tuple[int, int, int], int] = {} # rank 0: c_steps of planned pushes + self._updates: dict[int, Any] = {} # rank 0: broadcasts drained for the current plan + self._export_error: str | None = None # rank 0: a non-finite value in the last export + self.finished = False # the final cut is installed and acknowledged + self.shutdown = False + self._closed = False + self.client = None + self.stats = {"applied": 0, "pushed": 0, "push_bytes": 0, "dropped_pushes": 0} + # (iteration, "apply" | "push", fragment, version | pull step), same on every rank. + self.events: deque[tuple[int, str, int, int]] = deque(maxlen=4096) + + if self.rank == 0: + log.info("swift-megatron sync: %d trainable tensors (%.2f B elements) -> %d fragments", + len(names), sum(self.table.numel(n) for n in names) / 1e9, self.num_fragments) + error = None + if self.rank == 0: + try: + self.client = self._connect(syncer, hmac_key_path, wan_streams, client_factory) + except Exception as exc: # raised on every rank below + error = f"{type(exc).__name__}: {exc}" + self._raise_if(self._bcast_object(error)) + if self.learner_id == 0: + self._send_init() + + # ------------------------------------------------------------------ setup + def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, client_factory): + from yeto.protocol import SyncerClient + from yeto.syncer_shards import ShardedSyncerClient + + addrs = _syncer_addrs(syncer) + kwargs = {"session_contract_hash": self.contract} + key = _read_hmac_key(hmac_key_path) + if key is not None: + kwargs["hmac_key"] = key + if client_factory is not None: + client = client_factory(addrs, self.learner_id, self.layout, self.wire, wan_streams, **kwargs) + elif len(addrs) > 1: + client = ShardedSyncerClient(addrs, self.learner_id, self.layout, self.wire, wan_streams, **kwargs) + else: + client = SyncerClient(addrs[0], self.learner_id, self.layout, self.wire, wan_streams, **kwargs) + client.start() + log.info("connected to syncer at %s as learner %d", syncer, self.learner_id) + return client + + def _send_init(self) -> None: + """INIT_PARAMS of every fragment from this island's current masters.""" + + error = None + for fid in range(self.num_fragments): + parts = self._export_fragment(fid, delta=False) + if self.rank == 0 and error is None: + try: + if self._export_error is not None: + raise SwiftMegatronSyncError(self._export_error) + self.client.send_init_parts(fid, parts) + except Exception as exc: + error = f"INIT fragment {fid}: {type(exc).__name__}: {exc}" + self._raise_if(self._bcast_object(error)) + if self.rank == 0: + log.info("sent INIT_PARAMS for %d fragments", self.num_fragments) + + # ------------------------------------------------------------------ collectives + def _bcast_object(self, value: Any) -> Any: + import torch.distributed as dist + + if self.world == 1: + return value + box = [value] + dist.broadcast_object_list(box, src=0, group=self.comm.world) + return box[0] + + @staticmethod + def _raise_if(error: str | None) -> None: + if error: + raise SwiftMegatronSyncError(f"syncer side failed on rank 0: {error}") + + def _masters(self, shards: Sequence[LocalShard]) -> list[Any]: + from yeto.rl.megatron_masters import full_masters + + return full_masters(self.optimizer, [s.param for s in shards], device=self.comm.buffer_device, + reduce=self.comm.dp_reduce()) + + def _export_chunk(self, names: Sequence[str], *, delta: bool) -> list[Any] | None: + """Full FP32 values (``delta``: minus the anchor) of ``names``; the CPU + tensors on rank 0, None elsewhere. Collective (module doc).""" + + import torch + import torch.distributed as dist + + wanted = set(names) + local = [s for s in self.shards if s.name in wanted] + masters = self._masters(local) + values: dict[int, Any] = {} + for i, (shard, master) in enumerate(zip(local, masters, strict=True)): + value = master.detach().to(dtype=torch.float32) + if delta: + anchor = self.anchors.get(id(shard.param)) + if anchor is None: + raise SwiftMegatronSyncError(f"{shard.name} has no anchor yet") + value = value - anchor.to(device=value.device, dtype=torch.float32) + values[i] = value + # TP gathers: dense first, then expert, each in name order, so no two + # ranks wait on two groups in opposite orders. + full: dict[str, Any] = {} + for expert in (False, True): + for i, shard in enumerate(local): + if shard.expert is not expert: + continue + value = values[i] + if shard.part_dim >= 0: + parts = [torch.empty_like(value) for _ in range(shard.tp_size)] + dist.all_gather(parts, value.contiguous(), group=self.comm.group(shard.tp_group)) + value = torch.cat(parts, dim=shard.part_dim) if shard.tp_rank == 0 else None + del parts + if value is not None and self.table.owners[shard.name] == self.rank: + full.setdefault(shard.name, value) + del values, masters + out: list[Any] | None = [] if self.rank == 0 else None + for name in names: + owner = self.table.owners[name] + if self.rank == 0: + if owner == 0: + tensor = full.pop(name) + else: + tensor = torch.empty(self.table.shapes[name], dtype=torch.float32, + device=self.comm.buffer_device) + dist.recv(tensor, src=owner, group=self.comm.world) + out.append(tensor.detach().to("cpu")) + elif owner == self.rank: + dist.send(full.pop(name).contiguous(), dst=0, group=self.comm.world) + return out + + def _export_fragment(self, fid: int, *, delta: bool) -> list[Any] | None: + """The fragment's values on rank 0 as wire byte parts (INIT, or PUSH + for bf16/f32 sessions) or, for a q4 push, one f32 tensor.""" + + from yeto.protocol import DTYPE_Q4 + from yeto.tensor_io import pack_tensor + + parts: list[Any] | None = [] if self.rank == 0 else None + self._export_error = None + q4 = delta and self.wire == DTYPE_Q4 + for names in self.fragment_chunks[fid]: + tensors = self._export_chunk(names, delta=delta) + if parts is None: + continue + for tensor in tensors: + if self._export_error is None and not bool(tensor.isfinite().all()): + # Raised on every rank by the caller's error broadcast. + self._export_error = f"fragment {fid} holds a non-finite value" + parts.append(tensor.reshape(-1) if q4 else + pack_tensor(tensor.reshape(-1), self.bulk if not delta else self._push_dtype())) + return parts + + def _push_dtype(self) -> int: + from yeto.protocol import push_wire_dtype + + return push_wire_dtype(self.wire) + + def _apply_chunk(self, names: Sequence[str], flat: Any | None, dtype: Any, *, blend: bool) -> None: + """Install the raw global values of ``names`` (``flat`` on rank 0, in the + broadcast dtype) as every rank's anchor, master and model copy.""" + + import torch + import torch.distributed as dist + + from yeto.rl.megatron_masters import has_complete_master, write_masters + + total = sum(self.table.numel(name) for name in names) + if self.rank == 0: + buffer = flat.to(device=self.comm.buffer_device) + else: + buffer = torch.empty(total, dtype=dtype, device=self.comm.buffer_device) + if self.world > 1: + dist.broadcast(buffer, src=0, group=self.comm.world) + offsets, offset = {}, 0 + for name in names: + offsets[name] = offset + offset += self.table.numel(name) + wanted = set(names) + local = [s for s in self.shards if s.name in wanted] + targets = [] + for shard in local: + start = offsets[shard.name] + whole = buffer[start:start + self.table.numel(shard.name)].view(self.table.shapes[shard.name]) + if shard.part_dim >= 0: + size = shard.local_shape[shard.part_dim] + whole = whole.narrow(shard.part_dim, shard.tp_rank * size, size) + raw = whole.to(dtype=torch.float32).contiguous() + self.anchors[id(shard.param)] = whole.to(device="cpu", dtype=self.anchor_dtype).contiguous().clone() + targets.append(raw) + if blend and self.merge_alpha > 0: + masters = self._masters(local) + targets = [self.merge_alpha * m.to(torch.float32) + (1.0 - self.merge_alpha) * t + for m, t in zip(masters, targets, strict=True)] + del masters + params = [s.param for s in local] + with torch.no_grad(): + write_masters(self.optimizer, params, targets) + for param, target in zip(params, targets, strict=True): + if has_complete_master(param) and param.dtype != torch.float32: + param.data.copy_(target.to(dtype=param.dtype)) + + def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool) -> None: + """Apply one fragment's payload (``data`` on rank 0) chunk by chunk.""" + + import torch + + view = memoryview(data).cast("B") if self.rank == 0 else None + itemsize = torch.empty((), dtype=dtype).element_size() + expected = self.layout.fragments[fid].numel * itemsize + if view is not None and len(view) != expected: + raise SwiftMegatronSyncError(f"fragment {fid} payload has {len(view)} bytes, expected {expected}") + offset = 0 + for names in self.fragment_chunks[fid]: + nbytes = sum(self.table.numel(name) for name in names) * itemsize + flat = None + if view is not None: + flat = torch.frombuffer(bytearray(view[offset:offset + nbytes]), dtype=torch.uint8).view(dtype) + self._apply_chunk(names, flat, dtype, blend=blend) + offset += nbytes + + # ------------------------------------------------------------------ boundary + def after_optimizer_step(self, iteration: int, target_tokens: int, *, + update_successful: bool = True) -> None: + """Run one non-blocking sync boundary. Every rank calls it after every + ``optimizer.step()``, with the same arguments. + + ``iteration``: the trainer's absolute 1-based iteration (continues on + resume); the H-step gate counts it. ``target_tokens``: the loss + tokens of this step over the whole island (the PUSH ``c_tokens`` + weight). A skipped step (loss-scale overflow / NaN, + ``update_successful=False``, ``target_tokens`` 0) still counts as an + inner step, so rounds stay aligned with ``iteration``, but adds no + token weight; its parameters did not move, so the next push carries + the same delta it would have carried anyway.""" + + if self.finished or self._closed: + return + iteration = int(iteration) + if self.steps_at_reset is None: + base = iteration - 1 + self.steps_at_reset = [base] * self.num_fragments + self.iteration = iteration + if update_successful: + self.units_total += max(0, int(target_tokens)) + plan = self._plan() if self.rank == 0 else None + plan = self._bcast_object(plan) + self._raise_if(plan.error) + if plan.finalizing: + self._finalize() + return + self.shutdown = plan.shutdown + updates = self._take_updates(plan) if self.rank == 0 else {} + for fid, version in plan.applies: + data = updates.get(fid) if self.rank == 0 else None + self._apply_fragment(fid, data, self.bulk_torch, blend=True) + self.steps_at_reset[fid] = iteration + self.units_at_reset[fid] = self.units_total + self.versions[fid] = version + self.anchored[fid] = True + self.global_step = max(self.global_step, version) + self.stats["applied"] += 1 + self.events.append((iteration, "apply", fid, version)) + error = None + for fid, step, attempt, c_units in plan.pushes: + parts = self._export_fragment(fid, delta=True) + self.events.append((iteration, "push", fid, step)) + if self.rank == 0 and error is None: + error = self._push(fid, step, attempt, c_units, parts) + if plan.pushes: + self._raise_if(self._bcast_object(error)) + + def _plan(self) -> _Plan: + """Rank 0: drain the client and decide this boundary (never blocks).""" + + plan = _Plan() + client = self.client + try: + client.check_health() + plan.finalizing = bool(client.finalizing.is_set()) + if plan.finalizing: + return plan + self._updates = {} + for update in client.drain_updates(): + # A later broadcast of the same fragment supersedes an earlier one. + self._updates[update.fragment_id] = update + plan.applies = sorted((fid, int(u.version)) for fid, u in self._updates.items()) + plan.shutdown = bool(client.shutdown.is_set()) + self.pending_pulls.extend(client.drain_pulls()) + except Exception as exc: + plan.error = f"{type(exc).__name__}: {exc}" + return plan + reset = {fid for fid, _ in plan.applies} + anchored = [self.anchored[f] or f in reset for f in range(self.num_fragments)] + still = [] + for pull in self.pending_pulls: + fid = int(pull.fragment_id) + at_reset = self.iteration if fid in reset else self.steps_at_reset[fid] + c_steps = self.iteration - at_reset + if not anchored[fid] or c_steps < self.interval or any(p[0] == fid for p in plan.pushes): + still.append(pull) + continue + units_at = self.units_total if fid in reset else self.units_at_reset[fid] + plan.pushes.append((fid, int(pull.global_step), int(pull.round_attempt), self.units_total - units_at)) + self._pull_meta[(fid, int(pull.global_step), int(pull.round_attempt))] = c_steps + self.pending_pulls = still + return plan + + def _take_updates(self, plan: _Plan) -> dict[int, Any]: + updates, self._updates = self._updates, {} + return {fid: updates[fid].data for fid, _ in plan.applies} + + def _push(self, fid: int, step: int, attempt: int, c_units: int, parts: list[Any]) -> str | None: + """Rank 0: send one PUSH_FRAGMENT. A lost connection drops the push + (the syncer re-requests it after the redial); anything else is fatal.""" + + import torch + + from yeto.protocol import DTYPE_Q4 + from yeto.tensor_io import quantize_q4 + + c_steps = self._pull_meta.pop((fid, step, attempt)) + if self._export_error is not None: + return f"PUSH fragment {fid}: {self._export_error}" + try: + if self.wire == DTYPE_Q4: + payload = quantize_q4(torch.cat(parts)) + self.client.push_fragment(fid, step, attempt, self.versions[fid], self.iteration, + c_steps, c_units, payload) + nbytes = len(payload) + else: + self.client.push_fragment_parts(fid, step, attempt, self.versions[fid], self.iteration, + c_steps, c_units, parts) + nbytes = sum(len(part) for part in parts) + except (ConnectionError, OSError) as exc: + self.stats["dropped_pushes"] += 1 + log.warning("push of fragment %d dropped (%s); the syncer re-requests it", fid, exc) + return None + except Exception as exc: + return f"PUSH fragment {fid}: {type(exc).__name__}: {exc}" + self.stats["pushed"] += 1 + self.stats["push_bytes"] += nbytes + return None + + def _finalize(self) -> None: + """Install the coordinator's final cut raw (``full``) or only confirm it + (``manifest``), then FINAL_ACK; every rank returns together.""" + + import torch + import torch.distributed as dist + + from yeto.protocol import FinalManifest + + error, manifest, finals = None, None, {} + if self.rank == 0: + try: + if self.final_cut_payload == "full": + manifest, broadcasts = self.client.wait_for_final_fragments() + finals = {update.fragment_id: update for update in broadcasts} + else: + manifest = self.client.wait_for_final_manifest() + except Exception as exc: + error = f"finalization: {type(exc).__name__}: {exc}" + box = self._bcast_object((error, None if manifest is None else (manifest.global_step, + tuple(manifest.versions)))) + self._raise_if(box[0]) + global_step, versions = box[1] + manifest = FinalManifest(global_step, tuple(versions)) + if self.final_cut_payload == "full": + for fid in range(self.num_fragments): + update = finals.get(fid) if self.rank == 0 else None + if self.rank == 0 and update is not None and update.version != manifest.versions[fid]: + raise SwiftMegatronSyncError(f"final fragment {fid} has version {update.version}, " + f"expected {manifest.versions[fid]}") + self._apply_fragment(fid, None if update is None else update.data, torch.float32, blend=False) + self.versions[fid] = manifest.versions[fid] + if self.world > 1: + dist.barrier(group=self.comm.world) + error = None + if self.rank == 0: + try: + self.client.acknowledge_finalization(manifest) + except Exception as exc: + error = f"FINAL_ACK: {type(exc).__name__}: {exc}" + self._raise_if(self._bcast_object(error)) + self.global_step = max(self.global_step, manifest.global_step) + self.finished = True + self.shutdown = True + if self.rank == 0: + log.info("final cut at global step %d installed and acknowledged", manifest.global_step) + + # ------------------------------------------------------------------ teardown + def close(self) -> None: + """Close the syncer connection (rank 0). Idempotent, collective-free, + safe from ``atexit`` after a crash.""" + + if self._closed: + return + self._closed = True + client, self.client = self.client, None + if client is not None: + try: + client.close() + except Exception as exc: # teardown must not mask the original error + log.warning("closing the syncer client failed: %s", exc) + + +def _current_accelerator() -> Any: + import torch + + try: + import torch_npu # noqa: F401 + + if torch.npu.is_available(): + return torch.device("npu", torch.npu.current_device()) + except ImportError: + pass + if torch.cuda.is_available(): + return torch.device("cuda", torch.cuda.current_device()) + return torch.device("cpu") + + +__all__: Iterable[str] = ( + "GlobalTable", + "LocalShard", + "MegatronDiLoCoSync", + "ParallelContext", + "SwiftMegatronSyncError", + "chunk_names", + "global_parameter_name", + "local_shards", +) From f657cfc5f3fc1d5088dd198d0c1f741e9d7f7f95 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Sat, 7 Nov 2026 19:38:43 +0800 Subject: [PATCH 16/20] Add ms-swift Megatron callback hook for full-parameter DiLoCo sync ms-swift's `megatron sft` has no DiLoCo boundary of its own. The hook loads through its official `--external_plugins` / `--callbacks` mechanism and registers `yeto_diloco`, which ms-swift instantiates after the model, optimizer and resume checkpoint are loaded. Every rank then calls the sync core once per optimizer step (skipped steps included, with zero target tokens) from `on_step_end`, before ms-swift's own save decision, so a checkpoint written at that step already holds the merged base. Configuration comes from YETO_DILOCO_* environment variables because the ms-swift argument parser refuses unknown flags; YETO_DILOCO_AUTO_ATTACH=1 appends the callback without touching the command line. The plugin adds its own repository root to sys.path, since ms-swift imports it by file path on every node, and composes with other plugins that patch `train_step` at class level. The hook passes update_successful, the final-cut payload mode and the collective transport through to the sync core. Co-Authored-By: Claude Opus 5.5 --- tests/test_swift_megatron_plugin.py | 207 ++++++++++++++++ yeto/swift_megatron/swift_plugin.py | 365 ++++++++++++++++++++++++++++ 2 files changed, 572 insertions(+) create mode 100644 tests/test_swift_megatron_plugin.py create mode 100644 yeto/swift_megatron/swift_plugin.py diff --git a/tests/test_swift_megatron_plugin.py b/tests/test_swift_megatron_plugin.py new file mode 100644 index 00000000..6f3915cb --- /dev/null +++ b/tests/test_swift_megatron_plugin.py @@ -0,0 +1,207 @@ +"""Hook-layer tests for yeto.swift_megatron.swift_plugin without Megatron/NPU.""" + +import types + +import pytest + +from yeto.swift_megatron import swift_plugin as sp + +ENV = { + "YETO_DILOCO_SYNCER": "tcp://syncer:7000", + "YETO_DILOCO_LEARNER_ID": "1", + "YETO_DILOCO_NUM_LEARNERS": "3", + "YETO_DILOCO_FRAGMENTS": "4", + "YETO_DILOCO_SYNC_INTERVAL_STEPS": "2", + "YETO_DILOCO_WIRE_DTYPE": "bf16", +} + + +class FakeSync: + instances = [] + + def __init__(self, model_chunks, optimizer, cfg): + self.model_chunks, self.optimizer, self.cfg = model_chunks, optimizer, cfg + self.calls = [] + self.closed = 0 + FakeSync.instances.append(self) + + def after_optimizer_step(self, iteration, target_tokens, *, update_successful=True): + self.calls.append((iteration, target_tokens)) + + def close(self): + self.closed += 1 + + +class FakeTrainer: + """Mimics BaseMegatronTrainer.run_train_step / train / finalize_training.""" + + def __init__(self, steps, start_iteration=0, dp_size=2): + # steps: list of (micro-batch token counts, update_successful) + self.args = types.SimpleNamespace(train_iters=start_iteration + len(steps), callbacks=[]) + self.state = types.SimpleNamespace(iteration=start_iteration) + self.wrapped_models = ["chunk0", "chunk1"] + self.optimizer = object() + self.steps = list(steps) + self.dp_size = dp_size + self.optimizer_steps = 0 + self.callbacks = [] + + def train_step(self, it): + mbs, ok = self.steps.pop(0) + self.optimizer_steps += 1 + return [{"loss": [1.0, float(n)]} for n in mbs], 1.0, ok + + def call_event(self, event): + for cb in self.callbacks: + getattr(cb, event)() + + def train(self): + self.call_event("on_train_begin") if hasattr(sp._CallbackBase, "on_train_begin") else None + while self.state.iteration < self.args.train_iters: + self.call_event("on_step_begin") if hasattr(sp._CallbackBase, "on_step_begin") else None + self.train_step(None) + self.state.iteration += 1 + self.call_event("on_step_end") + self.call_event("on_train_end") + + +@pytest.fixture +def make(monkeypatch): + FakeSync.instances.clear() + for k in list(__import__("os").environ): + if k.startswith(sp.ENV_PREFIX): + monkeypatch.delenv(k) + for k, v in ENV.items(): + monkeypatch.setenv(k, v) + monkeypatch.setattr(sp.YetoDiLoCoCallback, "sync_factory", staticmethod(FakeSync)) + + def _make(steps, **kw): + trainer = FakeTrainer(steps, **kw) + monkeypatch.setattr(sp.YetoDiLoCoCallback, "token_reducer", + staticmethod(lambda n: n * trainer.dp_size)) + cb = sp.YetoDiLoCoCallback(trainer) + trainer.callbacks.append(cb) + return trainer, cb + + return _make + + +def test_called_once_per_optimizer_step_with_global_tokens(make): + trainer, cb = make([([3, 5], True), ([2], True), ([4, 4, 4], True)]) + trainer.train() + sync = FakeSync.instances[0] + assert sync.calls == [(1, 16), (2, 4), (3, 24)] + assert len(sync.calls) == trainer.optimizer_steps + assert sync.model_chunks == ["chunk0", "chunk1"] and sync.optimizer is trainer.optimizer + assert sync.cfg == sp.DiLoCoHookConfig("tcp://syncer:7000", 1, 3, 4, 2, "bf16", None) + assert sync.closed == 1 + + +def test_skipped_step_still_calls_with_zero_tokens(make): + trainer, cb = make([([3], True), ([7], False), ([1], True)]) + trainer.train() + assert FakeSync.instances[0].calls == [(1, 6), (2, 0), (3, 2)] + + +def test_non_last_pipeline_stage_contributes_zero_locally(make): + trainer, cb = make([([], True)]) + trainer.train() + assert FakeSync.instances[0].calls == [(1, 0)] + + +def test_resume_continues_iteration(make): + trainer, cb = make([([1], True), ([1], True)], start_iteration=100) + trainer.train() + assert FakeSync.instances[0].calls == [(101, 2), (102, 2)] + + +def test_close_is_idempotent(make): + trainer, cb = make([([1], True)]) + trainer.train() + cb.close() + cb.on_train_end() + assert FakeSync.instances[0].closed == 1 + + +def test_step_end_without_train_step_raises(make): + trainer, cb = make([([1], True)]) + trainer.state.iteration += 1 + with pytest.raises(RuntimeError): + cb.on_step_end() + + +def test_disable_flag(make, monkeypatch): + monkeypatch.setenv("YETO_DILOCO_DISABLE", "1") + trainer, cb = make([([1], True)]) + trainer.train() + assert FakeSync.instances == [] and cb.sync is None + + +@pytest.mark.parametrize("key,value", [ + ("YETO_DILOCO_SYNCER", ""), + ("YETO_DILOCO_LEARNER_ID", "3"), + ("YETO_DILOCO_SYNC_INTERVAL_STEPS", "x"), + ("YETO_DILOCO_HMAC_KEY_PATH", "/nonexistent/key"), +]) +def test_bad_config_fails_at_construction(make, monkeypatch, key, value): + monkeypatch.setenv(key, value) + with pytest.raises(ValueError): + make([([1], True)]) + + +def test_local_loss_tokens(): + assert sp.local_loss_tokens([{"loss": [0.5, 3.0]}, {"loss": [0.1, 2]}, {}, None]) == 5 + assert sp.local_loss_tokens([]) == 0 + + +def test_register_into_swift_map(monkeypatch): + callbacks_mod = types.ModuleType("swift.megatron.callbacks") + callbacks_mod.megatron_callbacks_map = {} + base_mod = types.ModuleType("swift.megatron.trainers.base") + + class BaseMegatronTrainer: + def __init__(self, args): + self.seen = list(args.callbacks) + + base_mod.BaseMegatronTrainer = BaseMegatronTrainer + for name, mod in { + "swift": types.ModuleType("swift"), + "swift.megatron": types.ModuleType("swift.megatron"), + "swift.megatron.callbacks": callbacks_mod, + "swift.megatron.trainers": types.ModuleType("swift.megatron.trainers"), + "swift.megatron.trainers.base": base_mod, + }.items(): + monkeypatch.setitem(__import__("sys").modules, name, mod) + assert sp.register(auto_attach=True) + assert callbacks_mod.megatron_callbacks_map["yeto_diloco"] is sp.YetoDiLoCoCallback + t = BaseMegatronTrainer(types.SimpleNamespace(callbacks=["print", "default_flow"])) + assert t.seen == ["yeto_diloco", "print", "default_flow"] + t = BaseMegatronTrainer(types.SimpleNamespace(callbacks=["yeto_diloco", "print"])) + assert t.seen == ["yeto_diloco", "print"] + + +def test_composes_with_class_level_train_step_patch(make, monkeypatch): + # Production plugins (e.g. ascend_stage_timer) replace the class-level + # train_step at import time, before the trainer and callback exist. + seen = [] + orig = FakeTrainer.train_step + + def patched(self, *a, **kw): + out = orig(self, *a, **kw) + seen.append(out[2]) + return out + + monkeypatch.setattr(FakeTrainer, "train_step", patched) + trainer, cb = make([([2], True), ([3], False)]) + trainer.train() + assert seen == [True, False] + assert FakeSync.instances[0].calls == [(1, 4), (2, 0)] + + +def test_repo_root_on_sys_path(): + import os + import sys + + root = os.path.dirname(os.path.dirname(os.path.abspath(sp.__file__))) + assert os.path.dirname(root) in sys.path or any( + os.path.abspath(p) == os.path.dirname(root) for p in sys.path) diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py new file mode 100644 index 00000000..85355c81 --- /dev/null +++ b/yeto/swift_megatron/swift_plugin.py @@ -0,0 +1,365 @@ +"""ms-swift Megatron hook layer for yeto full-parameter DiLoCo. + +Load this file with ms-swift's official plugin mechanism:: + + megatron sft ... \\ + --external_plugins /path/to/yeto/swift_megatron/swift_plugin.py \\ + --callbacks yeto_diloco + +Importing the module registers ``YetoDiLoCoCallback`` under the name +``yeto_diloco`` in ``swift.megatron.callbacks.megatron_callbacks_map``. +``BaseMegatronTrainer.__init__`` instantiates every name in ``args.callbacks`` +after the model, optimizer and resume checkpoint are loaded, so the callback +sees the final model chunks, the (distributed / offloaded) optimizer and the +resumed iteration. + +If adding ``--callbacks`` to the production command line is not possible, set +``YETO_DILOCO_AUTO_ATTACH=1``: the import then wraps +``BaseMegatronTrainer.__init__`` so the callback is appended automatically. + +Call contract (see INTERFACE.md): + * every rank calls ``sync.after_optimizer_step(iteration, target_tokens)`` + exactly once per ``optimizer.step()``, including skipped steps + (loss-scale overflow / NaN grad norm), where ``target_tokens`` is 0; + * ``iteration`` is the 1-based count of completed train steps + (``trainer.state.iteration`` after ms-swift increments it), so it continues + from the checkpoint iteration on resume; + * ``target_tokens`` is the number of loss tokens of this step summed over + all micro-batches and the whole data-parallel (+context-parallel) group, + identical on every rank of the island; + * ``sync.close()`` is called from ``on_train_end`` (and at interpreter exit + as a fallback if training aborts). + +The sync runs inside ``on_step_end``, i.e. before ms-swift's own +``default_flow`` callback decides to save, so a checkpoint written at that +step already contains the merged base. + +Configuration (environment variable / ``--yeto_*`` style is not accepted by +ms-swift's argument parser, so environment variables are the only channel): + + YETO_DILOCO_SYNCER syncer address (required) + YETO_DILOCO_LEARNER_ID learner id (int, required) + YETO_DILOCO_NUM_LEARNERS number of learners (int, required) + YETO_DILOCO_FRAGMENTS fragment count (int, default 1) + YETO_DILOCO_SYNC_INTERVAL_STEPS inner steps per round (int, required) + YETO_DILOCO_WIRE_DTYPE bf16 | fp32 | ... (default bf16) + YETO_DILOCO_HMAC_KEY_PATH path to the HMAC key file (optional) + YETO_DILOCO_DISABLE set to 1 to make the callback a no-op +""" + +from __future__ import annotations + +import atexit +import logging +import os +from dataclasses import dataclass +from typing import Any, Callable, Optional + +logger = logging.getLogger(__name__) + + +def _ensure_repo_on_path() -> None: + """Make ``import yeto`` work when ms-swift loads this file by path. + + ``--external_plugins`` imports the file directly, so the yeto checkout is + not necessarily importable; add the repository root (two levels above + this package) so no PYTHONPATH has to be forwarded to every node. + """ + import sys + + root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + if os.path.isdir(os.path.join(root, "yeto")) and root not in sys.path: + sys.path.append(root) + + +_ensure_repo_on_path() + +CALLBACK_NAME = "yeto_diloco" +ENV_PREFIX = "YETO_DILOCO_" + + +# --------------------------------------------------------------------------- +# Configuration +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class DiLoCoHookConfig: + syncer: str + learner_id: int + num_learners: int + fragments: int + sync_interval_steps: int + wire_dtype: str = "bf16" + hmac_key_path: Optional[str] = None + final_cut_payload: str = "full" + transport: str = "cpu" + + @classmethod + def from_env(cls, env=None) -> "DiLoCoHookConfig": + env = os.environ if env is None else env + + def get(name: str, default: Optional[str] = None, required: bool = False) -> Optional[str]: + value = env.get(ENV_PREFIX + name) + if value is None or value == "": + if required: + raise ValueError(f"{ENV_PREFIX}{name} must be set to enable yeto DiLoCo sync") + return default + return value + + def get_int(name: str, default: Optional[int] = None, required: bool = False) -> int: + raw = get(name, None if default is None else str(default), required) + try: + return int(raw) # type: ignore[arg-type] + except (TypeError, ValueError): + raise ValueError(f"{ENV_PREFIX}{name} must be an integer, got {raw!r}") from None + + cfg = cls( + syncer=get("SYNCER", required=True), # type: ignore[arg-type] + learner_id=get_int("LEARNER_ID", required=True), + num_learners=get_int("NUM_LEARNERS", required=True), + fragments=get_int("FRAGMENTS", default=1), + sync_interval_steps=get_int("SYNC_INTERVAL_STEPS", required=True), + wire_dtype=get("WIRE_DTYPE", default="bf16"), # type: ignore[arg-type] + hmac_key_path=get("HMAC_KEY_PATH"), + final_cut_payload=get("FINAL_CUT_PAYLOAD", default="full"), # type: ignore[arg-type] + transport=get("TRANSPORT", default="cpu"), # type: ignore[arg-type] + ) + if cfg.num_learners < 1 or not 0 <= cfg.learner_id < cfg.num_learners: + raise ValueError(f"invalid learner_id={cfg.learner_id} for num_learners={cfg.num_learners}") + if cfg.fragments < 1 or cfg.sync_interval_steps < 1: + raise ValueError("fragments and sync_interval_steps must be >= 1") + if cfg.final_cut_payload not in ("full", "manifest"): + raise ValueError(f"{ENV_PREFIX}FINAL_CUT_PAYLOAD must be full or manifest") + if cfg.transport not in ("cpu", "device"): + raise ValueError(f"{ENV_PREFIX}TRANSPORT must be cpu or device") + if cfg.hmac_key_path is not None and not os.path.isfile(cfg.hmac_key_path): + raise ValueError(f"{ENV_PREFIX}HMAC_KEY_PATH does not exist: {cfg.hmac_key_path}") + return cfg + + +def _env_flag(name: str) -> bool: + return os.environ.get(ENV_PREFIX + name, "").strip().lower() in {"1", "true", "yes", "on"} + + +# --------------------------------------------------------------------------- +# Token accounting +# --------------------------------------------------------------------------- + + +def local_loss_tokens(metrics) -> int: + """Sum the loss-token counts of this rank's micro-batches. + + ms-swift's ``MegatronTrainer.loss_func`` returns ``{'loss': [sum, n_tokens]}`` + per micro-batch; ``forward_backward_func`` hands back that list on the last + pipeline stage and an empty list elsewhere. + """ + total = 0 + for mb in metrics or (): + if not isinstance(mb, dict): + continue + loss = mb.get("loss") + if loss is None: + continue + try: + total += int(round(float(loss[1]))) + except (TypeError, IndexError): + continue + return total + + +def megatron_global_tokens(local_tokens: int) -> int: + """Reduce per-rank loss tokens to the island-wide count on every rank. + + SUM over the data-parallel group including context parallel (each DP/CP + rank sees a disjoint shard of the global batch), then MAX over the + pipeline group to broadcast the last stage's count to earlier stages. + TP ranks of one DP replica already hold identical counts. + """ + import torch + from megatron.core import parallel_state as mpu + + t = torch.tensor([int(local_tokens)], dtype=torch.long, device=torch.cuda.current_device()) + torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.SUM, + group=mpu.get_data_parallel_group(with_context_parallel=True)) + if mpu.get_pipeline_model_parallel_world_size() > 1: + torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.MAX, + group=mpu.get_pipeline_model_parallel_group()) + return int(t.item()) + + +# --------------------------------------------------------------------------- +# Sync construction +# --------------------------------------------------------------------------- + + +def trainable_filter(name: str, param) -> bool: + """Only parameters that the optimizer updates take part in sync.""" + return bool(getattr(param, "requires_grad", False)) + + +def default_sync_factory(model_chunks, optimizer, cfg: DiLoCoHookConfig): + from yeto.swift_megatron.sync import MegatronDiLoCoSync + + return MegatronDiLoCoSync( + model_chunks, + optimizer, + syncer=cfg.syncer, + learner_id=cfg.learner_id, + num_learners=cfg.num_learners, + fragments=cfg.fragments, + sync_interval_steps=cfg.sync_interval_steps, + wire_dtype=cfg.wire_dtype, + hmac_key_path=cfg.hmac_key_path, + trainable_filter=trainable_filter, + final_cut_payload=cfg.final_cut_payload, + transport=cfg.transport, + ) + + +def _model_chunks(trainer) -> list: + # Pass the DDP-wrapped chunks: the sync core needs the grad/param buffers + # and the DistributedOptimizer's mapping, both keyed on the wrapped models. + chunks = getattr(trainer, "wrapped_models", None) + if chunks is None: + chunks = getattr(trainer, "unwrapped_models", None) + if chunks is None: + raise RuntimeError("ms-swift trainer exposes neither wrapped_models nor unwrapped_models") + return list(chunks) + + +# --------------------------------------------------------------------------- +# Callback +# --------------------------------------------------------------------------- + +try: # ms-swift >= 4.x Megatron callback base class + from swift.megatron.callbacks import MegatronCallback as _CallbackBase +except Exception: # pragma: no cover - lets unit tests run without ms-swift + + class _CallbackBase: # type: ignore[no-redef] + def __init__(self, trainer): + self.trainer = trainer + self.args = trainer.args + self.state = trainer.state + + +class YetoDiLoCoCallback(_CallbackBase): + """Drives ``MegatronDiLoCoSync`` from ms-swift's Megatron training loop.""" + + # Overridable in tests. + sync_factory: Callable[..., Any] = staticmethod(default_sync_factory) + token_reducer: Callable[[int], int] = staticmethod(megatron_global_tokens) + + def __init__(self, trainer): + super().__init__(trainer) + self.enabled = not _env_flag("DISABLE") + self.sync = None + self._closed = False + self._pending = None # (iteration_before_step, local_tokens, update_successful) + if not self.enabled: + logger.warning("yeto DiLoCo callback disabled via %sDISABLE", ENV_PREFIX) + return + self.config = DiLoCoHookConfig.from_env() + self._wrap_train_step() + # Build here (not in on_train_begin) so a bad config fails before data + # loading; the trainer has already loaded the resume checkpoint, so the + # sync core snapshots the resumed weights as its initial base. + self.sync = type(self).sync_factory(_model_chunks(trainer), trainer.optimizer, self.config) + atexit.register(self.close) + logger.info("yeto DiLoCo sync attached: learner %d/%d, interval %d, fragments %d, start iteration %d", + self.config.learner_id, self.config.num_learners, self.config.sync_interval_steps, + self.config.fragments, int(self.state.iteration)) + + def _wrap_train_step(self) -> None: + """Capture train_step's metrics and update flag on this trainer instance. + + ``on_step_end`` receives no arguments, so the per-step result is + stashed here. Only the instance attribute is replaced; the class and + other trainers are untouched. + """ + trainer = self.trainer + original = trainer.train_step + if getattr(original, "_yeto_diloco_wrapped", False): + return + + def train_step(*args, **kwargs): + iteration_before = int(self.state.iteration) + metrics, grad_norm, update_successful = original(*args, **kwargs) + self._pending = (iteration_before, local_loss_tokens(metrics), bool(update_successful)) + return metrics, grad_norm, update_successful + + train_step._yeto_diloco_wrapped = True # type: ignore[attr-defined] + trainer.train_step = train_step + + def on_step_end(self): + if self.sync is None: + return + pending, self._pending = self._pending, None + if pending is None: + raise RuntimeError("yeto DiLoCo: on_step_end without a captured train_step result") + iteration_before, local_tokens, update_successful = pending + iteration = int(self.state.iteration) + if iteration != iteration_before + 1: + raise RuntimeError(f"yeto DiLoCo: unexpected iteration {iteration} after step from {iteration_before}") + # The reduction is a collective, so it runs on every rank even when the + # step was skipped; the skipped step then contributes zero tokens. + global_tokens = int(type(self).token_reducer(local_tokens)) + target_tokens = global_tokens if update_successful else 0 + self.sync.after_optimizer_step(iteration, target_tokens, + update_successful=update_successful) + + def on_train_end(self): + self.close() + + def close(self): + if self._closed or self.sync is None: + return + self._closed = True + try: + self.sync.close() + finally: + try: + atexit.unregister(self.close) + except Exception: + pass + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + + +def register(auto_attach: Optional[bool] = None) -> bool: + """Register the callback with ms-swift. Returns False if ms-swift is absent.""" + try: + from swift.megatron.callbacks import megatron_callbacks_map + except Exception: + return False + megatron_callbacks_map[CALLBACK_NAME] = YetoDiLoCoCallback + if auto_attach is None: + auto_attach = _env_flag("AUTO_ATTACH") + if auto_attach: + _patch_trainer_init() + return True + + +def _patch_trainer_init() -> None: + from swift.megatron.trainers.base import BaseMegatronTrainer + + original = BaseMegatronTrainer.__init__ + if getattr(original, "_yeto_diloco_wrapped", False): + return + + def __init__(self, args, *a, **kw): + callbacks = list(getattr(args, "callbacks", None) or []) + if CALLBACK_NAME not in callbacks: + # Run before ms-swift's own callbacks so a save at this step sees + # the merged base. + args.callbacks = [CALLBACK_NAME] + callbacks + original(self, args, *a, **kw) + + __init__._yeto_diloco_wrapped = True # type: ignore[attr-defined] + BaseMegatronTrainer.__init__ = __init__ + + +register() From af9cd3e0e6034a8921a088dfa6e46f003f6d3179 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 14:29:20 +0800 Subject: [PATCH 17/20] Reuse the torch learner's DiLoCo session helpers in the ms-swift Megatron sync Move client connection, wire dtype names, pull readiness, push delta encoding and boundary metrics from learner.py / sync_diloco_boundary into public yeto.diloco_sync helpers (torch learner behavior unchanged) and drive MegatronDiLoCoSync's counters, pulls and pushes through them and DiLoCoSyncState. Co-Authored-By: Claude Opus 5.5 --- tests/test_swift_megatron_sync.py | 6 +- yeto/diloco_sync.py | 136 ++++++++++++++++----- yeto/learner.py | 20 +--- yeto/swift_megatron/sync.py | 192 +++++++++++++++++------------- 4 files changed, 218 insertions(+), 136 deletions(-) diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py index 701d39fa..40a3b40a 100644 --- a/tests/test_swift_megatron_sync.py +++ b/tests/test_swift_megatron_sync.py @@ -54,8 +54,8 @@ def test_global_names_cover_layers_and_both_expert_namings(): def test_syncer_address_and_hmac_key_file(tmp_path): - assert sm._syncer_addrs("tcp://10.0.0.1:29400") == [("10.0.0.1", 29400)] - assert sm._syncer_addrs("a:1, tcp://b:2") == [("a", 1), ("b", 2)] + assert sm._syncer_spec("tcp://10.0.0.1:29400") == "10.0.0.1:29400" + assert sm._syncer_spec("a:1, tcp://b:2") == "a:1,b:2" key = tmp_path / "hmac.key" key.write_text("s3cret\n") assert sm._read_hmac_key(str(key)) == b"s3cret" # the YETO_ISLAND_HMAC_KEY encoding @@ -346,7 +346,7 @@ def _scripted_body(rank, world, port, out_dir): alpha = 0.5 def factory(addrs, learner_id, layout, wire, streams, **kwargs): - assert addrs == [("127.0.0.1", 1234)] and len(kwargs["session_contract_hash"]) == 32 + assert addrs == "127.0.0.1:1234" and len(kwargs["session_contract_hash"]) == 32 script = {} script[("bcast", 1)] = [BcastFragment(fid, 7, pack_tensor(_flat(layout, fid, globals_g), 2)) for fid in range(layout.num_fragments)] diff --git a/yeto/diloco_sync.py b/yeto/diloco_sync.py index 78d85ac5..cdf089c4 100644 --- a/yeto/diloco_sync.py +++ b/yeto/diloco_sync.py @@ -9,9 +9,95 @@ import torch.distributed as dist from .fragments import Fragment, FragmentLayout -from .protocol import DTYPE_Q4, FinalManifest, PullRequest, SyncerClient, bulk_dtype +from .protocol import ( + DTYPE_BF16, + DTYPE_F32, + DTYPE_Q4, + FinalManifest, + PullRequest, + SyncerClient, + bulk_dtype, +) from .tensor_io import fragment_flat, pack_tensor, quantize_q4, unpack_fragment +# ``--wire-dtype`` names of the learner sessions. +WIRE_DTYPES = {"bf16": DTYPE_BF16, "f32": DTYPE_F32, "q4": DTYPE_Q4} + + +def wire_dtype_code(name: str) -> int: + """The session dtype of a ``--wire-dtype`` name (``fp32`` is ``f32``).""" + try: + return WIRE_DTYPES["f32" if name == "fp32" else name] + except KeyError: + raise ValueError(f"wire dtype must be one of {tuple(WIRE_DTYPES)}, got {name!r}") from None + + +def connect_syncer(spec: str, learner_id: int, layout: FragmentLayout, wire_dtype: int, + wan_streams: int, **client_kwargs): + """Start one island's syncer client: a ``SyncerClient`` for ``host:port``, + a ``ShardedSyncerClient`` for a comma-separated shard list.""" + from .syncer_shards import ShardedSyncerClient, parse_syncer_addrs + + addrs = parse_syncer_addrs(spec) + if len(addrs) > 1: + # Fragment-sharded head: one connection per shard, every fragment + # routed to its owner (yeto/syncer_shards.py). + client = ShardedSyncerClient(addrs, learner_id, layout, wire_dtype, wan_streams, **client_kwargs) + else: + client = SyncerClient(addrs[0], learner_id, layout, wire_dtype, wan_streams, **client_kwargs) + client.start() + return client + + +def take_ready_pulls( + state: DiLoCoSyncState, + *, + steps_total: int, + units_total: int, + anchored: Callable[[int], bool], + min_steps: int = 1, +) -> list[tuple[PullRequest, int, int]]: + """``(pull, c_steps, c_units)`` of every pending pull that may be answered + now (the fragment trained ``min_steps`` steps since its last broadcast and + has an anchor); the others stay pending, in order.""" + ready, still = [], [] + for pull in state.pending_pulls: + fid = pull.fragment_id + c_steps = steps_total - state.steps_at_reset[fid] + if c_steps < max(1, min_steps) or not anchored(fid): + still.append(pull) + continue + ready.append((pull, c_steps, units_total - state.units_at_reset[fid])) + state.pending_pulls = still + return ready + + +def encode_push_delta(client, delta: torch.Tensor) -> bytes: + """PUSH wire bytes of one base-relative fragment delta.""" + if client.dtype == DTYPE_Q4: + return quantize_q4(delta) + return pack_tensor(delta, client.dtype) + + +def boundary_metrics(state: DiLoCoSyncState, *, merges: int, pushes: int, push_bytes: int, + push_delta_norm: float, steps_total: int) -> dict: + """The observer row of one boundary that merged or pushed something.""" + # Per-fragment lag behind the newest global version this island has + # seen: the single clearest signal that an island is falling behind + # the fleet rather than merely training slowly. + staleness = [state.global_step - v for v in state.fragment_versions] + return { + "sync/merges_applied": merges, + "sync/pushes": pushes, + "sync/push_bytes": push_bytes, + "sync/push_delta_norm": push_delta_norm, + "sync/pending_pulls": len(state.pending_pulls), + "sync/staleness_max": max(staleness), + "sync/staleness_mean": sum(staleness) / len(staleness), + "local_step": steps_total, + "global_step": state.global_step, + } + @dataclass class DiLoCoSyncState: @@ -136,26 +222,18 @@ def sync_diloco_boundary( pushed_delta_norm = 0.0 if rank == 0 and client is not None: state.pending_pulls.extend(client.drain_pulls()) - still_pending = [] pull_snapshot = None - for pull in state.pending_pulls: + for pull, c_steps, c_units in take_ready_pulls( + state, + steps_total=steps_total, + units_total=units_total, + anchored=lambda fid: state.anchors is not None and state.anchors[fid] is not None, + ): fid = pull.fragment_id - c_steps = steps_total - state.steps_at_reset[fid] - if c_steps < 1: - still_pending.append(pull) - continue - c_units = units_total - state.units_at_reset[fid] if pull_snapshot is None: pull_snapshot = snapshot_params() - anchor = state.anchors[fid] if state.anchors is not None else None - if anchor is None: - still_pending.append(pull) - continue - delta = fragment_flat(layout.fragments[fid], pull_snapshot).cpu() - anchor - if client.dtype == DTYPE_Q4: - payload = quantize_q4(delta) - else: - payload = pack_tensor(delta, client.dtype) + delta = fragment_flat(layout.fragments[fid], pull_snapshot).cpu() - state.anchors[fid] + payload = encode_push_delta(client, delta) client.push_fragment( fid, pull.global_step, @@ -170,27 +248,19 @@ def sync_diloco_boundary( pushed += 1 pushed_bytes += len(payload) pushed_delta_norm = max(pushed_delta_norm, float(delta.norm())) - state.pending_pulls = still_pending if shutdown_after_pulls: state.shutdown = client.shutdown.is_set() if observer is not None and rank == 0 and (actions or pushed): - # Per-fragment lag behind the newest global version this island has - # seen: the single clearest signal that an island is falling behind - # the fleet rather than merely training slowly. - staleness = [state.global_step - v for v in state.fragment_versions] observer( - { - "sync/merges_applied": len(actions), - "sync/pushes": pushed, - "sync/push_bytes": pushed_bytes, - "sync/push_delta_norm": pushed_delta_norm, - "sync/pending_pulls": len(state.pending_pulls), - "sync/staleness_max": max(staleness), - "sync/staleness_mean": sum(staleness) / len(staleness), - "local_step": steps_total, - "global_step": state.global_step, - } + boundary_metrics( + state, + merges=len(actions), + pushes=pushed, + push_bytes=pushed_bytes, + push_delta_norm=pushed_delta_norm, + steps_total=steps_total, + ) ) return False diff --git a/yeto/learner.py b/yeto/learner.py index 9c988ef8..a5764658 100644 --- a/yeto/learner.py +++ b/yeto/learner.py @@ -44,7 +44,7 @@ validate_lora_production_envelope, ) from .data import StreamingPackedBlocks, build_packed_dataset -from .diloco_sync import DiLoCoSyncState, sync_diloco_boundary +from .diloco_sync import DiLoCoSyncState, connect_syncer, sync_diloco_boundary, wire_dtype_code from .finalization import finalize_torch_island from .fragments import build_layout from .chunked_ce import ( @@ -56,8 +56,7 @@ from .losses import load_custom_loss, load_pickled_loss, sft_loss from .wandb_logger import TELEMETRY_EVERY from .models import MODEL_ALIASES as MODEL_ALIASES -from .protocol import DTYPE_BF16, DTYPE_F32, DTYPE_Q4, SyncerClient, bulk_dtype -from .syncer_shards import ShardedSyncerClient, parse_syncer_addrs +from .protocol import bulk_dtype from .tensor_io import ( apply_fragment, pack_fragment, @@ -1164,21 +1163,10 @@ def probe_loss_forward(probe_model, input_ids): ) log.info("dataset ready: %d blocks of %d tokens", len(dataset), args.seq_len) - wire_dtype = {"bf16": DTYPE_BF16, "f32": DTYPE_F32, "q4": DTYPE_Q4}[args.wire_dtype] + wire_dtype = wire_dtype_code(args.wire_dtype) client = None if rank == 0 and args.syncer != "none": - addrs = parse_syncer_addrs(args.syncer) - if len(addrs) > 1: - # Fragment-sharded head: one connection per shard, every - # fragment routed to its owner (yeto/syncer_shards.py). - client = ShardedSyncerClient( - addrs, args.learner_id, layout, wire_dtype, args.wan_streams - ) - else: - client = SyncerClient( - addrs[0], args.learner_id, layout, wire_dtype, args.wan_streams - ) - client.start() + client = connect_syncer(args.syncer, args.learner_id, layout, wire_dtype, args.wan_streams) log.info("connected to syncer at %s", args.syncer) if args.learner_id == 0: for fid, frag in enumerate(layout.fragments): diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py index c7d56740..4f553d02 100644 --- a/yeto/swift_megatron/sync.py +++ b/yeto/swift_megatron/sync.py @@ -30,10 +30,21 @@ trainer does not train (``trainable_filter``: frozen vision tower and aligner) and buffers (router ``expert_bias``) are not part of the layout. - The names are Megatron's, not the HF checkpoint's: the syncer is - name-agnostic, and Megatron-Bridge (which ``yeto.rl.megatron_hf_boundary`` - builds on) is not what ms-swift uses (it uses ``mcore_bridge``, whose - export is a whole-model pass with no per-parameter selection). Fused + The names are Megatron's, not the HF checkpoint's (the syncer is + name-agnostic). ms-swift's own converter, ``mcore_bridge`` + (``model.config.bridge``, NemotronHBridge), is deliberately not used for + the gather and write: ``export_weights`` is a whole-model pass with no + per-parameter selection, gathers every tensor over TP/EP/PP into + accelerator buffers on every rank (``_all_gather_tp`` / + ``_broadcast_ep_pp``) before ``target_device`` moves it, and reads the + bf16 model copy, which is not the FP32 master and is stale outside the + owned range under ``overlap_param_gather``; ``load_weights`` reads only a + safetensors directory and writes the bf16 copy, not the optimizer master. + With ~0.4 GiB of HBM headroom on the production pp0 stage none of that + fits. Megatron-Bridge (``yeto.rl.megatron_hf_boundary``) is not installed + with ms-swift. Everything protocol-side is the torch learner's own code + (``yeto.diloco_sync``: session state, pull readiness, push encoding, + metrics, client connection; ``yeto.fragments``; ``SyncerClient``). Fused Megatron tensors whose element order depends on the TP size (Mamba ``in_proj``, gated ``linear_fc1``) make the layout TP-specific, so the session contract hash binds the TP / expert-TP sizes and the syncer @@ -116,7 +127,6 @@ OUTPUT_LAYER_AVG_REGEX = r"(^|\.)output_layer\.weight$" DEFAULT_CHUNK_NUMEL = 1 << 28 CONTRACT_DOMAIN = b"yeto-swift-megatron-native-v1\0" -WIRE_DTYPES = ("bf16", "f32", "fp32", "q4") FINAL_CUT_PAYLOADS = ("full", "manifest") _LAYER_MODULE_RE = re.compile(r"(^|\.)decoder\.layers\.(\d+)$") @@ -418,15 +428,6 @@ def reduce(flat: Any, leaf: Any) -> None: # ----------------------------------------------------------------------------- helpers -def _wire_code(wire_dtype: str) -> int: - from yeto.protocol import DTYPE_BF16, DTYPE_F32, DTYPE_Q4 - - try: - return {"bf16": DTYPE_BF16, "f32": DTYPE_F32, "fp32": DTYPE_F32, "q4": DTYPE_Q4}[wire_dtype] - except KeyError: - raise ValueError(f"wire_dtype must be one of {WIRE_DTYPES}, got {wire_dtype!r}") from None - - def _torch_wire(code: int) -> Any: import torch @@ -435,11 +436,11 @@ def _torch_wire(code: int) -> Any: return {DTYPE_F32: torch.float32, DTYPE_BF16: torch.bfloat16}[code] -def _syncer_addrs(spec: str) -> list[tuple[str, int]]: - from yeto.syncer_shards import parse_syncer_addrs +def _syncer_spec(spec: str) -> str: + """``host:port[,host:port...]`` without ``tcp://`` schemes (the hook's + ``YETO_DILOCO_SYNCER`` may carry one).""" - items = [item.strip() for item in spec.split(",")] - return parse_syncer_addrs(",".join(item.split("://", 1)[-1] for item in items)) + return ",".join(item.strip().split("://", 1)[-1] for item in spec.split(",")) def _read_hmac_key(path: str | None) -> bytes | None: @@ -478,7 +479,6 @@ class _Plan: """Rank 0's decision for one boundary, broadcast to every rank.""" applies: list[tuple[int, int]] = field(default_factory=list) # (fragment, version) - pushes: list[tuple[int, int, int, int]] = field(default_factory=list) # (fragment, step, attempt, c_units) finalizing: bool = False shutdown: bool = False error: str | None = None @@ -519,10 +519,11 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, fragment_pattern: str = "binpack", wan_streams: int = 4, final_cut_payload: str = "full", chunk_numel: int = DEFAULT_CHUNK_NUMEL, transport: str = "cpu", parallel_context: ParallelContext | None = None, - client_factory=None) -> None: + client_factory=None, observer: Callable[[dict], None] | None = None) -> None: import torch import torch.distributed as dist + from yeto.diloco_sync import DiLoCoSyncState, wire_dtype_code from yeto.fragments import build_layout from yeto.protocol import bulk_dtype, layout_fingerprint @@ -546,7 +547,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self.merge_alpha = float(merge_alpha) self.final_cut_payload = final_cut_payload self.chunk_numel = int(chunk_numel) - self.wire = _wire_code(wire_dtype) + self.wire = wire_dtype_code(wire_dtype) self.bulk = bulk_dtype(self.wire) self.bulk_torch = _torch_wire(self.bulk) self.anchor_dtype = self.bulk_torch # the raw global is exact in the broadcast dtype @@ -579,18 +580,16 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self.anchors: dict[int, Any] = {} # id(param) -> raw global shard (host) self.anchored = [False] * self.num_fragments - self.versions = [0] * self.num_fragments - self.steps_at_reset: list[int] | None = None - self.units_at_reset = [0] * self.num_fragments + # The torch learner's per-fragment counters and versions, kept alike on + # every rank (anchors live in ``self.anchors``, sharded, not here). + self.state = DiLoCoSyncState.create(self.num_fragments, track_anchors=False) + self._counters_started = False self.units_total = 0 self.iteration = 0 - self.global_step = 0 - self.pending_pulls: list[Any] = [] - self._pull_meta: dict[tuple[int, int, int], int] = {} # rank 0: c_steps of planned pushes + self.observer = observer self._updates: dict[int, Any] = {} # rank 0: broadcasts drained for the current plan self._export_error: str | None = None # rank 0: a non-finite value in the last export self.finished = False # the final cut is installed and acknowledged - self.shutdown = False self._closed = False self.client = None self.stats = {"applied": 0, "pushed": 0, "push_bytes": 0, "dropped_pushes": 0} @@ -612,21 +611,14 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, # ------------------------------------------------------------------ setup def _connect(self, syncer: str, hmac_key_path: str | None, wan_streams: int, client_factory): - from yeto.protocol import SyncerClient - from yeto.syncer_shards import ShardedSyncerClient + from yeto.diloco_sync import connect_syncer - addrs = _syncer_addrs(syncer) kwargs = {"session_contract_hash": self.contract} key = _read_hmac_key(hmac_key_path) if key is not None: kwargs["hmac_key"] = key - if client_factory is not None: - client = client_factory(addrs, self.learner_id, self.layout, self.wire, wan_streams, **kwargs) - elif len(addrs) > 1: - client = ShardedSyncerClient(addrs, self.learner_id, self.layout, self.wire, wan_streams, **kwargs) - else: - client = SyncerClient(addrs[0], self.learner_id, self.layout, self.wire, wan_streams, **kwargs) - client.start() + connect = client_factory or connect_syncer + client = connect(_syncer_spec(syncer), self.learner_id, self.layout, self.wire, wan_streams, **kwargs) log.info("connected to syncer at %s as learner %d", syncer, self.learner_id) return client @@ -823,12 +815,15 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *, token weight; its parameters did not move, so the next push carries the same delta it would have carried anyway.""" + from yeto.diloco_sync import boundary_metrics + if self.finished or self._closed: return + state = self.state iteration = int(iteration) - if self.steps_at_reset is None: - base = iteration - 1 - self.steps_at_reset = [base] * self.num_fragments + if not self._counters_started: # count from the (possibly resumed) first step + state.steps_at_reset = [iteration - 1] * self.num_fragments + self._counters_started = True self.iteration = iteration if update_successful: self.units_total += max(0, int(target_tokens)) @@ -838,29 +833,48 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *, if plan.finalizing: self._finalize() return - self.shutdown = plan.shutdown + state.shutdown = plan.shutdown + # Apply broadcasts before answering pulls (as sync_diloco_boundary): + # a next-round pull can overtake the previous broadcast. updates = self._take_updates(plan) if self.rank == 0 else {} for fid, version in plan.applies: data = updates.get(fid) if self.rank == 0 else None self._apply_fragment(fid, data, self.bulk_torch, blend=True) - self.steps_at_reset[fid] = iteration - self.units_at_reset[fid] = self.units_total - self.versions[fid] = version + state.steps_at_reset[fid] = iteration + state.units_at_reset[fid] = self.units_total + state.fragment_versions[fid] = version self.anchored[fid] = True - self.global_step = max(self.global_step, version) + state.global_step = max(state.global_step, version) self.stats["applied"] += 1 self.events.append((iteration, "apply", fid, version)) - error = None - for fid, step, attempt, c_units in plan.pushes: + pushes = self._ready_pushes() if self.rank == 0 else None + pushes = self._bcast_object(pushes) + if isinstance(pushes, str): + self._raise_if(pushes) + error, pushed_bytes = None, 0 + for fid, step, attempt, c_steps, c_units in pushes: parts = self._export_fragment(fid, delta=True) self.events.append((iteration, "push", fid, step)) if self.rank == 0 and error is None: - error = self._push(fid, step, attempt, c_units, parts) - if plan.pushes: + error, nbytes = self._push(fid, step, attempt, c_steps, c_units, parts) + pushed_bytes += nbytes + if pushes: self._raise_if(self._bcast_object(error)) + if self.observer is not None and self.rank == 0 and (plan.applies or pushes): + self.observer(boundary_metrics(state, merges=len(plan.applies), pushes=len(pushes), + push_bytes=pushed_bytes, push_delta_norm=0.0, + steps_total=iteration)) + + @property + def shutdown(self) -> bool: + return self.state.shutdown + + @property + def global_step(self) -> int: + return self.state.global_step def _plan(self) -> _Plan: - """Rank 0: drain the client and decide this boundary (never blocks).""" + """Rank 0: drain the client's broadcasts and flags (never blocks).""" plan = _Plan() client = self.client @@ -875,61 +889,71 @@ def _plan(self) -> _Plan: self._updates[update.fragment_id] = update plan.applies = sorted((fid, int(u.version)) for fid, u in self._updates.items()) plan.shutdown = bool(client.shutdown.is_set()) - self.pending_pulls.extend(client.drain_pulls()) except Exception as exc: plan.error = f"{type(exc).__name__}: {exc}" - return plan - reset = {fid for fid, _ in plan.applies} - anchored = [self.anchored[f] or f in reset for f in range(self.num_fragments)] - still = [] - for pull in self.pending_pulls: + return plan + + def _ready_pushes(self) -> list[tuple[int, int, int, int, int]] | str: + """Rank 0: the pulls answered now, by the torch learner's rule + (``take_ready_pulls``) with the H-step gate; one push per fragment + per boundary, a second pull of it waits.""" + + from yeto.diloco_sync import take_ready_pulls + + state = self.state + try: + state.pending_pulls.extend(self.client.drain_pulls()) + except Exception as exc: + return f"{type(exc).__name__}: {exc}" + ready = take_ready_pulls(state, steps_total=self.iteration, units_total=self.units_total, + anchored=self.anchored.__getitem__, min_steps=self.interval) + out, seen = [], set() + for pull, c_steps, c_units in ready: fid = int(pull.fragment_id) - at_reset = self.iteration if fid in reset else self.steps_at_reset[fid] - c_steps = self.iteration - at_reset - if not anchored[fid] or c_steps < self.interval or any(p[0] == fid for p in plan.pushes): - still.append(pull) + if fid in seen: + state.pending_pulls.append(pull) continue - units_at = self.units_total if fid in reset else self.units_at_reset[fid] - plan.pushes.append((fid, int(pull.global_step), int(pull.round_attempt), self.units_total - units_at)) - self._pull_meta[(fid, int(pull.global_step), int(pull.round_attempt))] = c_steps - self.pending_pulls = still - return plan + seen.add(fid) + out.append((fid, int(pull.global_step), int(pull.round_attempt), c_steps, c_units)) + return out def _take_updates(self, plan: _Plan) -> dict[int, Any]: updates, self._updates = self._updates, {} return {fid: updates[fid].data for fid, _ in plan.applies} - def _push(self, fid: int, step: int, attempt: int, c_units: int, parts: list[Any]) -> str | None: - """Rank 0: send one PUSH_FRAGMENT. A lost connection drops the push - (the syncer re-requests it after the redial); anything else is fatal.""" + def _push(self, fid: int, step: int, attempt: int, c_steps: int, c_units: int, + parts: list[Any]) -> tuple[str | None, int]: + """Rank 0: send one PUSH_FRAGMENT; ``(error, bytes)``. A lost + connection drops the push (the syncer re-requests it after the + redial); anything else is fatal on every rank.""" import torch + from yeto.diloco_sync import encode_push_delta from yeto.protocol import DTYPE_Q4 - from yeto.tensor_io import quantize_q4 - c_steps = self._pull_meta.pop((fid, step, attempt)) if self._export_error is not None: - return f"PUSH fragment {fid}: {self._export_error}" + return f"PUSH fragment {fid}: {self._export_error}", 0 + base = self.state.fragment_versions[fid] try: if self.wire == DTYPE_Q4: - payload = quantize_q4(torch.cat(parts)) - self.client.push_fragment(fid, step, attempt, self.versions[fid], self.iteration, - c_steps, c_units, payload) + # q4 blocks span tensors: the shared encoder takes the whole delta. + payload = encode_push_delta(self.client, torch.cat(parts)) + self.client.push_fragment(fid, step, attempt, base, self.iteration, c_steps, c_units, payload) nbytes = len(payload) - else: - self.client.push_fragment_parts(fid, step, attempt, self.versions[fid], self.iteration, - c_steps, c_units, parts) + else: # bf16 / f32: stream the per-tensor wire parts + self.client.push_fragment_parts(fid, step, attempt, base, self.iteration, c_steps, c_units, + parts) nbytes = sum(len(part) for part in parts) except (ConnectionError, OSError) as exc: self.stats["dropped_pushes"] += 1 log.warning("push of fragment %d dropped (%s); the syncer re-requests it", fid, exc) - return None + return None, 0 except Exception as exc: - return f"PUSH fragment {fid}: {type(exc).__name__}: {exc}" + return f"PUSH fragment {fid}: {type(exc).__name__}: {exc}", 0 self.stats["pushed"] += 1 self.stats["push_bytes"] += nbytes - return None + return None, nbytes def _finalize(self) -> None: """Install the coordinator's final cut raw (``full``) or only confirm it @@ -962,7 +986,7 @@ def _finalize(self) -> None: raise SwiftMegatronSyncError(f"final fragment {fid} has version {update.version}, " f"expected {manifest.versions[fid]}") self._apply_fragment(fid, None if update is None else update.data, torch.float32, blend=False) - self.versions[fid] = manifest.versions[fid] + self.state.fragment_versions[fid] = manifest.versions[fid] if self.world > 1: dist.barrier(group=self.comm.world) error = None @@ -972,9 +996,9 @@ def _finalize(self) -> None: except Exception as exc: error = f"FINAL_ACK: {type(exc).__name__}: {exc}" self._raise_if(self._bcast_object(error)) - self.global_step = max(self.global_step, manifest.global_step) + self.state.global_step = max(self.state.global_step, manifest.global_step) self.finished = True - self.shutdown = True + self.state.shutdown = True if self.rank == 0: log.info("final cut at global step %d installed and acknowledged", manifest.global_step) From e44d2428c206eedf32c5ad16ad3d62b8e4ca924c Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 15:12:43 +0800 Subject: [PATCH 18/20] Make the ms-swift Megatron sync decoupled: never wait for the syncer or other islands A synchronous boundary stalls every island on the slowest one and on the syncer, which defeats DiLoCo across heterogeneous or WAN-linked islands. - PUSH bytes go to a background sender thread on rank 0; while a push is still streaming, pulls stay pending instead of blocking the step. - One fragment is gathered per boundary; a pull is not answered from an anchor already spent on another round (late-carry anchor_reused), and a dropped push releases its anchor. - wire_dtype accepts the RL full-parameter push names (int4, int8, int4-log, push-bf16) with tensor-block quantization. - Test: quorum 1 / max_base_lag 1 / int4 with one island stalled, against the real syncer with --carry-answers-round; the other island keeps stepping and merging without blocking boundaries, the stalled island rejoins, and both install the syncer's final state. Co-Authored-By: Claude Opus 5.5 --- tests/test_swift_megatron_sync.py | 99 +++++++++++-- yeto/swift_megatron/sync.py | 225 ++++++++++++++++++++++++------ 2 files changed, 270 insertions(+), 54 deletions(-) diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py index 40a3b40a..f14fd5cd 100644 --- a/tests/test_swift_megatron_sync.py +++ b/tests/test_swift_megatron_sync.py @@ -373,7 +373,8 @@ def factory(addrs, learner_id, layout, wire, streams, **kwargs): adam_before = {k: {n: t.clone() for n, t in v.items()} for k, v in opt.optimizer.state.items()} updates_total = {name: torch.zeros(shape) for name, shape in full_shapes.items()} applied = {} - for it in (1, 2, 3): + snapshots = {} # iteration -> updates since the merge, at the end of that step + for it in (1, 2, 3, 4, 5): opt.step({id(p): _local_update(p, 0, it, tp, 2) for p in trainable}) for name, shape in full_shapes.items(): if it >= 2: # updates after the merge at iteration 1 @@ -381,7 +382,10 @@ def factory(addrs, learner_id, layout, wire, streams, **kwargs): if it == 1: pre_merge = {name: _full(name, shape) + _update(name, shape, 0, 1) for name, shape in full_shapes.items()} + snapshots[it] = {n: t.clone() for n, t in updates_total.items()} syncer.after_optimizer_step(it, 100 * it) + if rank == 0: # the PUSH streams in the background; let it land before the next step + assert syncer.sender.wait_idle(10) if it == 1: # merged with alpha: master = a * local + (1 - a) * global, bf16 copy = master for p in trainable: @@ -400,7 +404,8 @@ def factory(addrs, learner_id, layout, wire, streams, **kwargs): if it == 2: # one step since the merge: the pulls wait (H = 2) assert not any(e[1] == "push" for e in syncer.events) pushes = [e for e in syncer.events if e[1] == "push"] - assert [e[2] for e in pushes] == list(range(syncer.num_fragments)) and all(e[0] == 3 for e in pushes) + # decoupled: one fragment per boundary, the PUSH returns at once + assert [(e[0], e[2]) for e in pushes] == [(3 + f, f) for f in range(syncer.num_fragments)] assert torch.equal(chunk.vision_model.proj.weight, frozen_before) for key, state in opt.optimizer.state.items(): # Adam moments kept for n, t in state.items(): @@ -409,9 +414,11 @@ def factory(addrs, learner_id, layout, wire, streams, **kwargs): client = holder["client"] assert len(client.pushes) == syncer.num_fragments for fid, step, attempt, base, local_step, c_steps, c_units, payload in client.pushes: - assert (step, attempt, base, local_step, c_steps, c_units) == (8, 1, 7, 3, 2, 500) + it = 3 + fid + assert (step, attempt, base, local_step, c_steps, c_units) == ( + 8, 1, 7, it, it - 1, sum(100 * k for k in range(2, it + 1))) # PUSH = local - raw anchor (the global, not the blend) - local = {n: alpha * pre_merge[n] + (1 - alpha) * globals_g[n] + updates_total[n] for n in full_shapes} + local = {n: alpha * pre_merge[n] + (1 - alpha) * globals_g[n] + snapshots[it][n] for n in full_shapes} want = _flat(syncer.layout, fid, local) - _flat(syncer.layout, fid, globals_g) got = unpack_fragment(syncer.layout.fragments[fid], payload, 2) assert torch.allclose(got, want.to(torch.bfloat16).float(), atol=1e-3), fid @@ -449,12 +456,12 @@ def read_final_state(path: Path) -> list[torch.Tensor]: return out -def _island_worker(rank, world, island, dist_port, syncer_port, out_dir, drop_at): +def _island_worker(rank, world, island, dist_port, syncer_port, out_dir, drop_at, opts=None): _guard(out_dir, f"island{island}-{rank}", - lambda: _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at)) + lambda: _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at, opts or {})) -def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at): +def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at, opts): dist = _init_dist(rank, world, dist_port) tp, pp = rank % 2, rank // 2 tp_groups = [dist.new_group([0, 1]), dist.new_group([2, 3])] @@ -464,11 +471,15 @@ def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at): trainable = [p for p in chunk.parameters() if p.requires_grad] opt = _DistOpt(trainable, dp_rank=0, dp_size=1, dp_group=singles[rank]) syncer = sm.MegatronDiLoCoSync([chunk], opt, syncer=f"127.0.0.1:{syncer_port}", learner_id=island, - num_learners=2, fragments=3, sync_interval_steps=2, wire_dtype="bf16", + num_learners=2, fragments=3, sync_interval_steps=2, + wire_dtype=opts.get("wire", "bf16"), chunk_numel=64, parallel_context=ctx, final_cut_payload="full") it = 0 dropped = False started = time.monotonic() + stall_at, stall_s = opts.get("stall_at"), opts.get("stall_s", 0.0) + step_s = opts.get("step_s", 0.01) + timeline = [] # (iteration, wall time after the boundary, boundary seconds, applied so far) while not syncer.finished: if time.monotonic() - started > 150: raise TimeoutError(f"island {island} not finalized after {it} steps") @@ -481,8 +492,12 @@ def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at): except OSError: pass dropped = True + if it == stall_at: # a slow / stalled island: no boundary for stall_s + time.sleep(stall_s) + t0 = time.monotonic() syncer.after_optimizer_step(it, 64) - time.sleep(0.01) + timeline.append((it, time.monotonic(), time.monotonic() - t0, syncer.stats["applied"])) + time.sleep(step_s) for p in trainable: # the bf16 copy equals the installed FP32 master if id(p) in opt.owned: assert torch.equal(p.detach().reshape(-1), opt.master(p).to(torch.bfloat16)), p._full_name @@ -491,6 +506,7 @@ def _island_body(rank, world, island, dist_port, syncer_port, out_dir, drop_at): if rank == 0: torch.save({"names": names, "values": dict(zip(names, values, strict=True)), "events": list(syncer.events), "stats": dict(syncer.stats), "dropped": dropped, + "timeline": timeline, "fragments": [[n for n, _ in f.tensors] for f in syncer.layout.fragments]}, Path(out_dir, f"island{island}.pt")) syncer.close() @@ -566,3 +582,68 @@ def test_write_masters_also_sets_a_distinct_hdo_host_copy(): assert torch.equal(opt.master(params[0]), target) assert torch.equal(opt.optimizer.gpu_params_map_cpu_copy[shard], target) assert torch.equal(params[0].detach(), target.to(torch.bfloat16)) + + +DECOUPLED_SYNCER_FLAGS = [ # the RL async full-parameter preset (yeto/launch/rl_args.py), test-sized grace + "--quorum", "1", "--max-base-lag", "1", "--learner-weight", "tokens2-over-steps", + "--pipeline", "2", "--outer-lr", "1.0", "--outer-momentum", "0.0", + # A late-carried delta answers its round, so a quorum-1 round never waits + # for the stalled island just because the fast island's anchor was spent. + "--carry-answers-round", +] + + +@pytest.mark.timeout(300) +def test_decoupled_quorum_one_keeps_the_fast_island_merging_while_the_other_stalls(tmp_path): + """quorum 1 + max_base_lag 1 + int4 tensor-block pushes: while island 1 + makes no boundary for several seconds, island 0 keeps training without a + blocking boundary and its fragment rounds keep merging; island 1 then + resumes with a stale base (accepted with lag or refused by max_base_lag), + nothing waits for it, and both islands install the same final cut.""" + + binary = build_syncer() + from process_groups import group_popen + + port = free_port() + final_state = tmp_path / "final.bin" + proc = group_popen([str(binary), "--port", str(port), "--learners", "2", "--grace-ms", "200", + "--terminal-grace-ms", "2000", "--total-steps", "40", "--final-state", str(final_state), + *DECOUPLED_SYNCER_FLAGS], + stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + procs = [] + stall_s = 4.0 + try: + procs += _run(_island_worker, 4, 0, free_port(), port, str(tmp_path), None, + {"wire": "int4", "step_s": 0.05}) + procs += _run(_island_worker, 4, 1, free_port(), port, str(tmp_path), None, + {"wire": "int4", "step_s": 0.05, "stall_at": 6, "stall_s": stall_s}) + _join(procs, 240, tmp_path) + assert proc.wait(timeout=60) == 0, "syncer exited nonzero" + finally: + for p in procs: + if p.is_alive(): + p.kill() + if proc.poll() is None: + proc.kill() + log = proc.stdout.read() if proc.stdout else "" + (tmp_path / "syncer.log").write_text(log) + print(log[-4000:]) + + fast, slow = (torch.load(tmp_path / f"island{i}.pt") for i in range(2)) + stalled_from = next(t for it, t, _d, _a in slow["timeline"] if it == 5) + stalled_to = next(t for it, t, _d, _a in slow["timeline"] if it == 6) + assert stalled_to - stalled_from >= stall_s + during = [row for row in fast["timeline"] if stalled_from < row[1] < stalled_to] + assert len(during) >= 10, "the fast island stopped stepping while the other stalled" + assert max(d for _it, _t, d, _a in during) < 1.0, "a boundary blocked on the stalled island" + merged_during = during[-1][3] - during[0][3] + resumed_pushes = [e for e in slow["events"] if e[1] == "push" and e[0] > 6] + assert resumed_pushes, "the resumed island never pushed again" + for name in fast["names"]: + assert torch.equal(fast["values"][name], slow["values"][name]), name + final = read_final_state(final_state) + for fid, frag_names in enumerate(fast["fragments"]): + assert torch.equal(torch.cat([fast["values"][n].reshape(-1) for n in frag_names]), final[fid]) + # A late-carried delta counts towards its round's quorum, so the fast + # island's rounds keep closing while the other island stalls. + assert merged_during >= 3, merged_during diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py index 4f553d02..5f7fbe64 100644 --- a/yeto/swift_megatron/sync.py +++ b/yeto/swift_megatron/sync.py @@ -81,6 +81,31 @@ push, finalization, shutdown, error), so every rank runs the same collectives in the same order and an error on rank 0 is raised on every rank. +Decoupled DiLoCo (no island waits for another) + The island never waits for a merge or for another island: a boundary + only drains what the client already received (broadcasts are applied at + the next boundary after they arrive), the PUSH bytes are handed to a + background sender thread on rank 0 (``_PushSender``; the client encodes + int4/int8 with tensor blocks there), and while a PUSH is still streaming + (slow link or syncer, a redial) pulls stay pending instead of blocking. + One fragment is gathered per boundary. A pull is not answered from an + anchor already spent on another round (the syncer's late carry refuses + that as ``anchor_reused``); a dropped PUSH un-spends its anchor. The + syncer side is the RL async full-parameter preset (quorum 1, adaptive + grace, ``--max-base-lag 1``, tokens2-over-steps, two fragment rounds in + flight, rda); ``wire_dtype="int4"`` is its int4 tensor-block push. + + What still blocks a step is the collective gather of one fragment's FP32 + masters (and the in-island broadcast of a received one): host copies, + gloo all-reduce over DP and all-gather over TP, then point-to-point to + rank 0. For the 120B / 24-card recipe and P = 16 one fragment is + ~7.5e9 elements, ~0.3e9 per rank (~1.2 GB FP32), 30 GB arriving at rank + 0: roughly 10-20 s over TCP gloo, against a ~570 s step, i.e. ~2-5 % at + one push and one apply per step. It is not overlapped with training: the + HDO masters change under the next step, so overlapping needs a host + snapshot of the local shards plus a second set of gloo groups driven by a + background thread in the same order on every rank. + Memory (N = trainable parameters of the island, P = fragments, n_r = the elements of rank r's local parameter shards, C = ``chunk_numel``, t = TP size): @@ -90,18 +115,22 @@ parameters). The production pp0 stage has ~0.4 GiB of HBM headroom, so ``transport="device"`` (chunk buffers on the accelerator, peak about ``4 * C * (2 + t)`` bytes) is only for islands with room to spare; -* every rank, host: one anchor per local shard in the anchor dtype (bf16 for a - bf16 session), ``2 * n_r`` bytes, plus one chunk in flight, about +* every rank, host: one anchor per local shard in the broadcast dtype (bf16 + for a bf16 session, FP32 for the int4 / int8 push sessions, whose + broadcasts are FP32), ``2 * n_r`` or ``4 * n_r`` bytes, plus one chunk in + flight, about ``4 * C * (2 + t)`` bytes (C = 2**28 -> ~6 GiB at t = 4; a single tensor larger than C is a chunk of its own). For the TP4/EP8/PP3 (world 24) Nemotron-3.5-Super recipe (~120B, MoE-dominated, ETP 1) n_r is about - 120e9 / 24 = 5e9: ~10 GB of anchors per rank, ~80 GB plus ~50 GB of chunks + 120e9 / 24 = 5e9: ~10 GB (int4: ~20 GB) of anchors per rank, ~80 GB + (int4: ~160 GB) plus ~50 GB of chunks per 8-card host, beside the ~450 GB the offloaded masters and Adam moments take of its ~2 TB; * rank 0, host additionally: the wire bytes of one pushed fragment (``2 * N / P`` for bf16) and the client's received broadcasts (each ``2 * N / P``; at most the syncer's ``--pipeline`` rounds in flight). For - N = 120e9 and P = 16 that is ~15 GB per fragment, ~45-60 GB peak. A + N = 120e9 and P = 16 that is ~15 GB per fragment, ~45-60 GB peak (int4: + the FP32 PUSH parts and FP32 broadcasts, ~30 GB each, ~90-120 GB). A full-payload final cut holds every fragment in FP32 (``4 * N`` = 480 GB) on rank 0's host: run the syncer with ``--final-cut-payload manifest`` and ``final_cut_payload="manifest"`` for a 120B island. @@ -523,7 +552,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, import torch import torch.distributed as dist - from yeto.diloco_sync import DiLoCoSyncState, wire_dtype_code + from yeto.diloco_sync import DiLoCoSyncState from yeto.fragments import build_layout from yeto.protocol import bulk_dtype, layout_fingerprint @@ -547,7 +576,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self.merge_alpha = float(merge_alpha) self.final_cut_payload = final_cut_payload self.chunk_numel = int(chunk_numel) - self.wire = wire_dtype_code(wire_dtype) + self.wire = session_dtype(wire_dtype) self.bulk = bulk_dtype(self.wire) self.bulk_torch = _torch_wire(self.bulk) self.anchor_dtype = self.bulk_torch # the raw global is exact in the broadcast dtype @@ -580,6 +609,12 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self.anchors: dict[int, Any] = {} # id(param) -> raw global shard (host) self.anchored = [False] * self.num_fragments + # (base version, round step) of each fragment's last sent PUSH: a + # decoupled syncer refuses a second delta from a spent anchor for + # another round (late carry, ``anchor_reused``), so such a pull waits + # for a newer broadcast instead of a gather that would be refused. A + # restarted attempt of the same round may reuse it. + self.spent: list[tuple[int, int] | None] = [None] * self.num_fragments # The torch learner's per-fragment counters and versions, kept alike on # every rank (anchors live in ``self.anchors``, sharded, not here). self.state = DiLoCoSyncState.create(self.num_fragments, track_anchors=False) @@ -592,7 +627,9 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self.finished = False # the final cut is installed and acknowledged self._closed = False self.client = None - self.stats = {"applied": 0, "pushed": 0, "push_bytes": 0, "dropped_pushes": 0} + self.stats = {"applied": 0, "pushed": 0, "push_bytes": 0, "dropped_pushes": 0, + "deferred_pushes": 0} + self.sender: _PushSender | None = None # rank 0: background PUSH streaming # (iteration, "apply" | "push", fragment, version | pull step), same on every rank. self.events: deque[tuple[int, str, int, int]] = deque(maxlen=4096) @@ -603,6 +640,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, if self.rank == 0: try: self.client = self._connect(syncer, hmac_key_path, wan_streams, client_factory) + self.sender = _PushSender(self.client, self.stats, self._push_dropped) except Exception as exc: # raised on every rank below error = f"{type(exc).__name__}: {exc}" self._raise_if(self._bcast_object(error)) @@ -729,13 +767,17 @@ def _export_fragment(self, fid: int, *, delta: bool) -> list[Any] | None: # Raised on every rank by the caller's error broadcast. self._export_error = f"fragment {fid} holds a non-finite value" parts.append(tensor.reshape(-1) if q4 else - pack_tensor(tensor.reshape(-1), self.bulk if not delta else self._push_dtype())) + pack_tensor(tensor.reshape(-1), self.bulk if not delta else self._push_part_dtype())) return parts - def _push_dtype(self) -> int: - from yeto.protocol import push_wire_dtype + def _push_part_dtype(self) -> int: + """The dtype of the PUSH parts handed to the client: f32 on the + f32-bulk sessions (the client encodes bf16 / int8 / int4 itself, with + tensor blocks), else the session's push wire dtype.""" + + from yeto.protocol import DTYPE_F32, bulk_dtype, push_wire_dtype - return push_wire_dtype(self.wire) + return DTYPE_F32 if bulk_dtype(self.wire) == DTYPE_F32 else push_wire_dtype(self.wire) def _apply_chunk(self, names: Sequence[str], flat: Any | None, dtype: Any, *, blend: bool) -> None: """Install the raw global values of ``names`` (``flat`` on rank 0, in the @@ -851,18 +893,17 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *, pushes = self._bcast_object(pushes) if isinstance(pushes, str): self._raise_if(pushes) - error, pushed_bytes = None, 0 + error = None for fid, step, attempt, c_steps, c_units in pushes: parts = self._export_fragment(fid, delta=True) self.events.append((iteration, "push", fid, step)) if self.rank == 0 and error is None: - error, nbytes = self._push(fid, step, attempt, c_steps, c_units, parts) - pushed_bytes += nbytes + error = self._push(fid, step, attempt, c_steps, c_units, parts) if pushes: self._raise_if(self._bcast_object(error)) if self.observer is not None and self.rank == 0 and (plan.applies or pushes): self.observer(boundary_metrics(state, merges=len(plan.applies), pushes=len(pushes), - push_bytes=pushed_bytes, push_delta_norm=0.0, + push_bytes=self.stats["push_bytes"], push_delta_norm=0.0, steps_total=iteration)) @property @@ -880,6 +921,8 @@ def _plan(self) -> _Plan: client = self.client try: client.check_health() + if self.sender is not None: + self.sender.raise_error() plan.finalizing = bool(client.finalizing.is_set()) if plan.finalizing: return plan @@ -905,27 +948,43 @@ def _ready_pushes(self) -> list[tuple[int, int, int, int, int]] | str: state.pending_pulls.extend(self.client.drain_pulls()) except Exception as exc: return f"{type(exc).__name__}: {exc}" + if self.sender is not None and self.sender.busy(): + # Decoupled: the previous PUSH is still streaming (slow link or + # syncer); every pull waits rather than the training step. + if state.pending_pulls: + self.stats["deferred_pushes"] += 1 + return [] ready = take_ready_pulls(state, steps_total=self.iteration, units_total=self.units_total, anchored=self.anchored.__getitem__, min_steps=self.interval) - out, seen = [], set() + # One fragment per boundary: its PUSH streams in the background while + # training goes on, and the gather of a second would only queue + # behind it on the host. + out = [] for pull, c_steps, c_units in ready: - fid = int(pull.fragment_id) - if fid in seen: + if out or not self._pushable(pull): state.pending_pulls.append(pull) continue - seen.add(fid) - out.append((fid, int(pull.global_step), int(pull.round_attempt), c_steps, c_units)) + out.append((int(pull.fragment_id), int(pull.global_step), int(pull.round_attempt), c_steps, c_units)) return out + def _pushable(self, pull: Any) -> bool: + fid = int(pull.fragment_id) + spent = self.spent[fid] + return spent is None or spent[0] != self.state.fragment_versions[fid] or spent[1] == int(pull.global_step) + + def _push_dropped(self, fid: int, base: int, step: int) -> None: + """Sender thread: a dropped PUSH never reached the syncer, so its anchor is not spent.""" + if self.spent[fid] == (base, step): + self.spent[fid] = None + def _take_updates(self, plan: _Plan) -> dict[int, Any]: updates, self._updates = self._updates, {} return {fid: updates[fid].data for fid, _ in plan.applies} def _push(self, fid: int, step: int, attempt: int, c_steps: int, c_units: int, - parts: list[Any]) -> tuple[str | None, int]: - """Rank 0: send one PUSH_FRAGMENT; ``(error, bytes)``. A lost - connection drops the push (the syncer re-requests it after the - redial); anything else is fatal on every rank.""" + parts: list[Any]) -> str | None: + """Rank 0: hand one PUSH_FRAGMENT to the background sender and return + at once (an error only for a fragment that cannot be pushed at all).""" import torch @@ -933,27 +992,15 @@ def _push(self, fid: int, step: int, attempt: int, c_steps: int, c_units: int, from yeto.protocol import DTYPE_Q4 if self._export_error is not None: - return f"PUSH fragment {fid}: {self._export_error}", 0 - base = self.state.fragment_versions[fid] - try: - if self.wire == DTYPE_Q4: - # q4 blocks span tensors: the shared encoder takes the whole delta. - payload = encode_push_delta(self.client, torch.cat(parts)) - self.client.push_fragment(fid, step, attempt, base, self.iteration, c_steps, c_units, payload) - nbytes = len(payload) - else: # bf16 / f32: stream the per-tensor wire parts - self.client.push_fragment_parts(fid, step, attempt, base, self.iteration, c_steps, c_units, - parts) - nbytes = sum(len(part) for part in parts) - except (ConnectionError, OSError) as exc: - self.stats["dropped_pushes"] += 1 - log.warning("push of fragment %d dropped (%s); the syncer re-requests it", fid, exc) - return None, 0 - except Exception as exc: - return f"PUSH fragment {fid}: {type(exc).__name__}: {exc}", 0 - self.stats["pushed"] += 1 - self.stats["push_bytes"] += nbytes - return None, nbytes + return f"PUSH fragment {fid}: {self._export_error}" + head = (fid, step, attempt, self.state.fragment_versions[fid], self.iteration, c_steps, c_units) + self.spent[fid] = (self.state.fragment_versions[fid], step) + if self.wire == DTYPE_Q4: + # q4 blocks span tensors: the shared encoder takes the whole delta. + self.sender.submit(head, payload=encode_push_delta(self.client, torch.cat(parts))) + else: + self.sender.submit(head, parts=parts) + return None def _finalize(self) -> None: """Install the coordinator's final cut raw (``full``) or only confirm it @@ -1010,6 +1057,8 @@ def close(self) -> None: if self._closed: return self._closed = True + if self.sender is not None: + self.sender.stop() client, self.client = self.client, None if client is not None: try: @@ -1018,6 +1067,92 @@ def close(self) -> None: log.warning("closing the syncer client failed: %s", exc) +def session_dtype(name: str) -> int: + """The syncer session dtype of a wire name: the torch learner's + ``--wire-dtype`` names (bf16, f32, q4), or the full-parameter push names + of the RL island (``int4``, ``int8``, ``int4-log``, ``push-bf16``: f32 + INIT/BCAST, quantized PUSH deltas with blocks restarting at tensor + boundaries, as the RL async full-parameter preset sends them).""" + + from yeto.diloco_sync import WIRE_DTYPES, wire_dtype_code + from yeto.protocol import DTYPE_FLAG_TENSOR_BLOCKS, FULL_PARAMETER_PUSH_DTYPES + + if name in WIRE_DTYPES or name == "fp32": + return wire_dtype_code(name) + key = "bf16" if name == "push-bf16" else name + if key not in FULL_PARAMETER_PUSH_DTYPES or key == "f32": + raise ValueError(f"unknown wire dtype {name!r}: one of {sorted(WIRE_DTYPES)} or " + "int4, int8, int4-log, push-bf16") + code = FULL_PARAMETER_PUSH_DTYPES[key] + return code | DTYPE_FLAG_TENSOR_BLOCKS if key != "bf16" else code + + +class _PushSender: + """Rank 0's background PUSH streaming: one PUSH in flight at a time, so a + slow link or syncer never holds a training step (decoupled DiLoCo, as the + RL full-parameter island's background push). A lost connection drops the + push (the syncer re-requests it after the redial); any other failure is + re-raised by the next boundary on every rank.""" + + def __init__(self, client: Any, stats: dict, on_drop: Callable[[int, int, int], None]) -> None: + import queue + import threading + + self.client, self.stats, self.on_drop = client, stats, on_drop + self.jobs: queue.Queue = queue.Queue() + self.idle = threading.Event() + self.idle.set() + self.error: str | None = None + self.thread = threading.Thread(target=self._run, name="yeto-swift-push", daemon=True) + self.thread.start() + + def busy(self) -> bool: + return not self.idle.is_set() + + def submit(self, head: tuple, *, parts: list[Any] | None = None, payload: bytes | None = None) -> None: + self.idle.clear() + self.jobs.put((head, parts, payload)) + + def raise_error(self) -> None: + if self.error is not None: + error, self.error = self.error, None + raise SwiftMegatronSyncError(error) + + def wait_idle(self, timeout: float | None = None) -> bool: + return self.idle.wait(timeout) + + def stop(self) -> None: + self.jobs.put(None) + self.thread.join(timeout=5.0) + + def _run(self) -> None: + while True: + job = self.jobs.get() + if job is None: + return + head, parts, payload = job + fid = head[0] + try: + if payload is not None: + self.client.push_fragment(*head, payload) + nbytes = len(payload) + else: + self.client.push_fragment_parts(*head, parts) + nbytes = sum(len(part) for part in parts) + self.stats["pushed"] += 1 + self.stats["push_bytes"] += nbytes + except (ConnectionError, OSError) as exc: + self.stats["dropped_pushes"] += 1 + self.on_drop(fid, head[3], head[1]) + log.warning("push of fragment %d dropped (%s); the syncer re-requests it", fid, exc) + except Exception as exc: + self.error = f"PUSH fragment {fid}: {type(exc).__name__}: {exc}" + finally: + del parts, payload, job + if self.jobs.empty(): + self.idle.set() + + def _current_accelerator() -> Any: import torch From e957e2519df097f5d449df361aad4c6b36a8e56b Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Thu, 5 Nov 2026 15:35:55 +0800 Subject: [PATCH 19/20] swift_megatron: take expert indices from the MoE layer on Megatron-Core 0.18 Megatron-Core 0.18 sets local_expert_indices on MoELayer only, not on its experts container, so every expert-parallel rank named its experts with the local indices and the island synced only EP rank 0's experts (the others aliased them). Fall back to the parent MoE layer and refuse an expert parameter that still has no global index. Co-Authored-By: Claude Opus 5.5 --- tests/test_swift_megatron_expert_index.py | 80 +++++++++++++++++++++++ yeto/swift_megatron/sync.py | 23 +++++-- 2 files changed, 98 insertions(+), 5 deletions(-) create mode 100644 tests/test_swift_megatron_expert_index.py diff --git a/tests/test_swift_megatron_expert_index.py b/tests/test_swift_megatron_expert_index.py new file mode 100644 index 00000000..37cf21b8 --- /dev/null +++ b/tests/test_swift_megatron_expert_index.py @@ -0,0 +1,80 @@ +"""Global expert names when only the MoE layer carries ``local_expert_indices``. + +Megatron-Core 0.18 (the ms-swift production stack) sets +``local_expert_indices`` on ``MoELayer`` but not on its ``experts`` child +(``TEGroupedMLP`` / ``SequentialMLP``); the sync must still give each +expert-parallel rank's experts their global indices, or the experts of every +EP rank collapse onto the names of EP rank 0's. +""" + +from __future__ import annotations + +import pytest +import torch + +from yeto.swift_megatron import sync as sm + + +class _Module(torch.nn.Module): + pass + + +def _moe_chunk(ep_rank: int, local_experts: int, *, indices_on_experts: bool, grouped: bool): + root = _Module() + root.decoder = _Module() + root.decoder.layers = torch.nn.ModuleList() + layer = _Module() + layer.layer_number = 1 + layer.mlp = _Module() + indices = [ep_rank * local_experts + e for e in range(local_experts)] + layer.mlp.local_expert_indices = indices + experts = _Module() + if indices_on_experts: + experts.local_expert_indices = indices + if grouped: + experts.linear_fc1 = _Module() + for e in range(local_experts): + p = torch.nn.Parameter(torch.zeros(2, 2)) + p.allreduce = False + setattr(experts.linear_fc1, f"weight{e}", p) + else: + experts.local_experts = torch.nn.ModuleList() + for _ in range(local_experts): + mlp = _Module() + mlp.linear_fc1 = _Module() + mlp.linear_fc1.weight = torch.nn.Parameter(torch.zeros(2, 2)) + mlp.linear_fc1.weight.allreduce = False + experts.local_experts.append(mlp) + layer.mlp.experts = experts + root.decoder.layers.append(layer) + return root + + +def _ctx(): + return sm.ParallelContext(rank=0, world=1, tp_group=None, tp_rank=0, tp_size=1, + etp_group=None, etp_rank=0, etp_size=1) + + +@pytest.mark.parametrize("grouped", [True, False]) +def test_expert_indices_fall_back_to_the_moe_layer(grouped): + names = set() + for ep_rank in range(4): + chunk = _moe_chunk(ep_rank, 2, indices_on_experts=False, grouped=grouped) + names |= {s.name for s in sm.local_shards([chunk], _ctx())} + if grouped: + expected = {f"decoder.layers.0.mlp.experts.linear_fc1.weight{e}" for e in range(8)} + else: + expected = {f"decoder.layers.0.mlp.experts.local_experts.{e}.linear_fc1.weight" for e in range(8)} + assert names == expected + + +def test_container_indices_still_win(): + chunk = _moe_chunk(1, 2, indices_on_experts=True, grouped=True) + assert sm._expert_modules(chunk) == [("decoder.layers.0.mlp.experts.", [2, 3])] + + +def test_expert_without_any_indices_is_refused(): + chunk = _moe_chunk(1, 2, indices_on_experts=False, grouped=True) + del chunk.decoder.layers[0].mlp.local_expert_indices + with pytest.raises(sm.SwiftMegatronSyncError, match="local_expert_indices"): + sm.local_shards([chunk], _ctx()) diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py index 5f7fbe64..c6a2d735 100644 --- a/yeto/swift_megatron/sync.py +++ b/yeto/swift_megatron/sync.py @@ -259,14 +259,21 @@ def _global_layer_prefixes(chunk: Any) -> list[tuple[str, str]]: def _expert_modules(chunk: Any) -> list[tuple[str, list[int]]]: """``(prefix, local_expert_indices)`` of every expert container (the - ``experts`` child of a MoE layer; the MoE layer itself also carries - ``local_expert_indices`` but holds the router and shared experts), - longest first.""" + ``experts`` child of a MoE layer), longest first. + Older Megatron-Core put ``local_expert_indices`` on the expert container + itself; since 0.1x only the MoE layer (its parent, which also holds the + router and shared experts) carries it, so fall back to the parent.""" + + modules = dict(chunk.named_modules()) out = [] - for name, module in chunk.named_modules(): + for name, module in modules.items(): + if name.rsplit(".", 1)[-1] != "experts": + continue indices = getattr(module, "local_expert_indices", None) - if indices is None or name.rsplit(".", 1)[-1] != "experts": + if indices is None and "." in name: + indices = getattr(modules.get(name.rsplit(".", 1)[0]), "local_expert_indices", None) + if indices is None: continue out.append((name + ".", [int(i) for i in indices])) return sorted(out, key=lambda item: -len(item[0])) @@ -328,6 +335,12 @@ def local_shards(model_chunks: Sequence[Any], ctx: ParallelContext, if name in out: raise SwiftMegatronSyncError(f"two local parameters map to the global name {name!r}") expert = getattr(param, "allreduce", True) is False + if expert and ".experts." in local_name and not any( + local_name.startswith(prefix) for prefix, _ in experts): + # Without its global expert index the tensor would alias the + # same-named tensor of every other expert-parallel rank. + raise SwiftMegatronSyncError( + f"{local_name}: expert parameter without local_expert_indices on its container") tp_rank, tp_size, group = ((ctx.etp_rank, ctx.etp_size, ctx.etp_group) if expert else (ctx.tp_rank, ctx.tp_size, ctx.tp_group)) part_dim = -1 From 3205d6ab6c34c5791fefa97a5b777e6d64ba9ab8 Mon Sep 17 00:00:00 2001 From: MichaelChung Date: Sat, 7 Nov 2026 19:58:01 +0800 Subject: [PATCH 20/20] swift_megatron: review fixes for rank-0 checks, final cut, sharded syncers, docs - Check broadcast payload sizes while rank 0 builds the boundary plan, and the final cut's versions and sizes before the finalization outcome is broadcast. Raised later, on rank 0 alone, they left the other ranks waiting in the chunk broadcast for good. - End training at the step that installed the final cut (ms-swift saves at train_iters), instead of training past the global result and saving local steps over it. - Refuse a comma-separated (sharded) syncer address on every rank before connecting: chunked INIT and push streaming talk to one syncer. - Report push_bytes per boundary as the torch learner does, and leave out push_delta_norm, which this sync never computes. - Docs: list FINAL_CUT_PAYLOAD and TRANSPORT and every wire dtype, drop a reference to a file outside the repository, and correct where the callback is added and where the initial base comes from. Drop a debug print from a test and a comment that tripped the secret scanner. Co-Authored-By: Claude Opus 5.5 --- tests/test_swift_megatron_plugin.py | 14 +++++++++ tests/test_swift_megatron_sync.py | 9 ++++-- yeto/swift_megatron/swift_plugin.py | 27 +++++++++++++---- yeto/swift_megatron/sync.py | 47 +++++++++++++++++++++++------ 4 files changed, 80 insertions(+), 17 deletions(-) diff --git a/tests/test_swift_megatron_plugin.py b/tests/test_swift_megatron_plugin.py index 6f3915cb..bd6ae6b8 100644 --- a/tests/test_swift_megatron_plugin.py +++ b/tests/test_swift_megatron_plugin.py @@ -97,6 +97,20 @@ def test_called_once_per_optimizer_step_with_global_tokens(make): assert sync.closed == 1 +def test_training_ends_at_the_step_that_installed_the_final_cut(make): + trainer, cb = make([([1], True)] * 5) + sync = FakeSync.instances[0] + original = sync.after_optimizer_step + + def step(iteration, target_tokens, **kw): + original(iteration, target_tokens, **kw) + sync.finished = iteration == 2 + + sync.after_optimizer_step = step + trainer.train() + assert [c[0] for c in sync.calls] == [1, 2] and trainer.args.train_iters == 2 and sync.closed == 1 + + def test_skipped_step_still_calls_with_zero_tokens(make): trainer, cb = make([([3], True), ([7], False), ([1], True)]) trainer.train() diff --git a/tests/test_swift_megatron_sync.py b/tests/test_swift_megatron_sync.py index f14fd5cd..7d58723a 100644 --- a/tests/test_swift_megatron_sync.py +++ b/tests/test_swift_megatron_sync.py @@ -58,10 +58,16 @@ def test_syncer_address_and_hmac_key_file(tmp_path): assert sm._syncer_spec("a:1, tcp://b:2") == "a:1,b:2" key = tmp_path / "hmac.key" key.write_text("s3cret\n") - assert sm._read_hmac_key(str(key)) == b"s3cret" # the YETO_ISLAND_HMAC_KEY encoding + assert sm._read_hmac_key(str(key)) == b"s3cret" # same encoding as YETO_ISLAND_HMAC_KEY assert sm._read_hmac_key(None) is None +def test_a_sharded_syncer_is_refused_before_connecting(): + with pytest.raises(ValueError, match="sharded syncers"): + sm.MegatronDiLoCoSync([], None, syncer="a:1,b:2", learner_id=0, num_learners=1, fragments=1, + sync_interval_steps=1) + + def test_chunks_are_bounded_and_keep_order(): sizes = {"a": 3, "b": 3, "c": 10, "d": 1, "e": 1} assert sm.chunk_names(list(sizes), sizes.__getitem__, 6) == [["a", "b"], ["c"], ["d", "e"]] @@ -563,7 +569,6 @@ def test_two_tp2_pp2_islands_merge_to_the_syncer_final_state_through_a_connectio if kind == "apply": assert version >= versions.get(fid, -1), "fragment versions went backwards" versions[fid] = version - print(json.dumps([i["stats"] for i in islands])) def test_write_masters_also_sets_a_distinct_hdo_host_copy(): diff --git a/yeto/swift_megatron/swift_plugin.py b/yeto/swift_megatron/swift_plugin.py index 85355c81..1b036e99 100644 --- a/yeto/swift_megatron/swift_plugin.py +++ b/yeto/swift_megatron/swift_plugin.py @@ -15,9 +15,10 @@ If adding ``--callbacks`` to the production command line is not possible, set ``YETO_DILOCO_AUTO_ATTACH=1``: the import then wraps -``BaseMegatronTrainer.__init__`` so the callback is appended automatically. +``BaseMegatronTrainer.__init__`` so the callback is added automatically, ahead +of ms-swift's own callbacks. -Call contract (see INTERFACE.md): +Call contract: * every rank calls ``sync.after_optimizer_step(iteration, target_tokens)`` exactly once per ``optimizer.step()``, including skipped steps (loss-scale overflow / NaN grad norm), where ``target_tokens`` is 0; @@ -32,7 +33,9 @@ The sync runs inside ``on_step_end``, i.e. before ms-swift's own ``default_flow`` callback decides to save, so a checkpoint written at that -step already contains the merged base. +step already contains the merged base. Once the syncer's final cut is +installed the callback ends training at that step (ms-swift then saves), so +later local steps cannot overwrite the final global parameters. Configuration (environment variable / ``--yeto_*`` style is not accepted by ms-swift's argument parser, so environment variables are the only channel): @@ -42,7 +45,13 @@ YETO_DILOCO_NUM_LEARNERS number of learners (int, required) YETO_DILOCO_FRAGMENTS fragment count (int, default 1) YETO_DILOCO_SYNC_INTERVAL_STEPS inner steps per round (int, required) - YETO_DILOCO_WIRE_DTYPE bf16 | fp32 | ... (default bf16) + YETO_DILOCO_WIRE_DTYPE bf16 | fp32 | int8 | int4 | int4-log | push-bf16 + (default bf16) + YETO_DILOCO_FINAL_CUT_PAYLOAD full | manifest (default full; a large island + should use manifest: the full final cut is + every fragment in FP32 on rank 0's host) + YETO_DILOCO_TRANSPORT cpu | device: where the sync's collectives + run (default cpu, gloo mirrors) YETO_DILOCO_HMAC_KEY_PATH path to the HMAC key file (optional) YETO_DILOCO_DISABLE set to 1 to make the callback a no-op """ @@ -262,8 +271,8 @@ def __init__(self, trainer): self.config = DiLoCoHookConfig.from_env() self._wrap_train_step() # Build here (not in on_train_begin) so a bad config fails before data - # loading; the trainer has already loaded the resume checkpoint, so the - # sync core snapshots the resumed weights as its initial base. + # loading; the trainer has already loaded the resume checkpoint, which + # is what learner 0's INIT_PARAMS carries. self.sync = type(self).sync_factory(_model_chunks(trainer), trainer.optimizer, self.config) atexit.register(self.close) logger.info("yeto DiLoCo sync attached: learner %d/%d, interval %d, fragments %d, start iteration %d", @@ -307,6 +316,12 @@ def on_step_end(self): target_tokens = global_tokens if update_successful else 0 self.sync.after_optimizer_step(iteration, target_tokens, update_successful=update_successful) + if getattr(self.sync, "finished", False) and iteration < int(self.args.train_iters): + # The final cut is installed: stop here (ms-swift's default_flow + # saves at train_iters) instead of training past the global result. + logger.warning("yeto DiLoCo: final cut installed at iteration %d; ending training (train_iters %d)", + iteration, int(self.args.train_iters)) + self.args.train_iters = iteration def on_train_end(self): self.close() diff --git a/yeto/swift_megatron/sync.py b/yeto/swift_megatron/sync.py index c6a2d735..82780699 100644 --- a/yeto/swift_megatron/sync.py +++ b/yeto/swift_megatron/sync.py @@ -577,6 +577,10 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, raise ValueError(f"learner_id {learner_id} outside [0, {num_learners})") if final_cut_payload not in FINAL_CUT_PAYLOADS: raise ValueError(f"final_cut_payload must be one of {FINAL_CUT_PAYLOADS}") + if "," in _syncer_spec(syncer): + # Checked on every rank, before rank 0 connects: the chunked INIT + # and push streaming are not routed per shard owner. + raise ValueError("the ms-swift Megatron sync talks to one syncer; sharded syncers are not supported") if not dist.is_initialized(): raise SwiftMegatronSyncError("torch.distributed must be initialized (Megatron does this)") self.ctx = parallel_context or ParallelContext.from_megatron() @@ -642,6 +646,7 @@ def __init__(self, model_chunks, optimizer, *, syncer: str, learner_id: int, self.client = None self.stats = {"applied": 0, "pushed": 0, "push_bytes": 0, "dropped_pushes": 0, "deferred_pushes": 0} + self._observed_push_bytes = 0 # rank 0: push_bytes at the last observer report self.sender: _PushSender | None = None # rank 0: background PUSH streaming # (iteration, "apply" | "push", fragment, version | pull step), same on every rank. self.events: deque[tuple[int, str, int, int]] = deque(maxlen=4096) @@ -836,6 +841,18 @@ def _apply_chunk(self, names: Sequence[str], flat: Any | None, dtype: Any, *, bl if has_complete_master(param) and param.dtype != torch.float32: param.data.copy_(target.to(dtype=param.dtype)) + def _check_payload(self, fid: int, data: Any, dtype: Any) -> None: + """Rank 0: refuse a fragment payload of the wrong size (or none).""" + + import torch + + if data is None: + raise SwiftMegatronSyncError(f"fragment {fid} has no payload") + nbytes = len(memoryview(data).cast("B")) + expected = self.layout.fragments[fid].numel * torch.empty((), dtype=dtype).element_size() + if nbytes != expected: + raise SwiftMegatronSyncError(f"fragment {fid} payload has {nbytes} bytes, expected {expected}") + def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool) -> None: """Apply one fragment's payload (``data`` on rank 0) chunk by chunk.""" @@ -843,9 +860,6 @@ def _apply_fragment(self, fid: int, data: Any, dtype: Any, *, blend: bool) -> No view = memoryview(data).cast("B") if self.rank == 0 else None itemsize = torch.empty((), dtype=dtype).element_size() - expected = self.layout.fragments[fid].numel * itemsize - if view is not None and len(view) != expected: - raise SwiftMegatronSyncError(f"fragment {fid} payload has {len(view)} bytes, expected {expected}") offset = 0 for names in self.fragment_chunks[fid]: nbytes = sum(self.table.numel(name) for name in names) * itemsize @@ -915,9 +929,14 @@ def after_optimizer_step(self, iteration: int, target_tokens: int, *, if pushes: self._raise_if(self._bcast_object(error)) if self.observer is not None and self.rank == 0 and (plan.applies or pushes): - self.observer(boundary_metrics(state, merges=len(plan.applies), pushes=len(pushes), - push_bytes=self.stats["push_bytes"], push_delta_norm=0.0, - steps_total=iteration)) + # push_bytes: what the sender finished since the last report (the + # torch learner reports its boundary's bytes, not a running total) + sent, self._observed_push_bytes = (self.stats["push_bytes"] - self._observed_push_bytes, + self.stats["push_bytes"]) + metrics = boundary_metrics(state, merges=len(plan.applies), pushes=len(pushes), + push_bytes=sent, push_delta_norm=0.0, steps_total=iteration) + metrics.pop("sync/push_delta_norm", None) # not computed: the delta is never materialized whole + self.observer(metrics) @property def shutdown(self) -> bool: @@ -943,6 +962,10 @@ def _plan(self) -> _Plan: for update in client.drain_updates(): # A later broadcast of the same fragment supersedes an earlier one. self._updates[update.fragment_id] = update + for fid, update in self._updates.items(): + # Checked here, before the plan goes out: a bad payload found + # while applying would fail rank 0 inside the chunk broadcasts. + self._check_payload(fid, update.data, self.bulk_torch) plan.applies = sorted((fid, int(u.version)) for fid, u in self._updates.items()) plan.shutdown = bool(client.shutdown.is_set()) except Exception as exc: @@ -1030,6 +1053,15 @@ def _finalize(self) -> None: if self.final_cut_payload == "full": manifest, broadcasts = self.client.wait_for_final_fragments() finals = {update.fragment_id: update for update in broadcasts} + # Checked before the outcome is broadcast, so a mismatch + # fails every rank instead of rank 0 alone. + for fid in range(self.num_fragments): + update = finals.get(fid) + if update is None or update.version != manifest.versions[fid]: + raise SwiftMegatronSyncError( + f"final fragment {fid} has version {None if update is None else update.version}, " + f"expected {manifest.versions[fid]}") + self._check_payload(fid, update.data, torch.float32) else: manifest = self.client.wait_for_final_manifest() except Exception as exc: @@ -1042,9 +1074,6 @@ def _finalize(self) -> None: if self.final_cut_payload == "full": for fid in range(self.num_fragments): update = finals.get(fid) if self.rank == 0 else None - if self.rank == 0 and update is not None and update.version != manifest.versions[fid]: - raise SwiftMegatronSyncError(f"final fragment {fid} has version {update.version}, " - f"expected {manifest.versions[fid]}") self._apply_fragment(fid, None if update is None else update.data, torch.float32, blend=False) self.state.fragment_versions[fid] = manifest.versions[fid] if self.world > 1: