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