From 8358c1b13a69ba8b6d43df4735a01e4f4bc2ba73 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:31:17 +0000 Subject: [PATCH 1/9] syncer: pipeline policy-sweep fragments and async per-round checkpoints Allow --pipeline 1..P with --policy-sweep-fragments so fragment k+1 uploads while fragment k merges, checkpoints and broadcasts; commits stay in strict step order. Add --async-checkpoint (policy sweep only): the committed state is snapshotted and written on a blocking task, and the round is broadcast only after the checkpoint is durable. Checkpoint bytes, versions and the semantic profile hash are unchanged; the tape gains checkpoint_ms. Launcher flags --rl-sync-fragment-pipeline and --rl-sync-async-checkpoint (default off). Adds bit-equality, failure and real-syncer tests and a CPU micro-benchmark (tools/bench_sync_boundary.py). Co-Authored-By: Claude Opus 5.5 --- syncer/src/main.rs | 5 + syncer/src/server.rs | 496 +++++++++++++++++++++-- syncer/src/state.rs | 118 ++++-- tests/test_rl_dense_sweep_real_syncer.py | 87 +++- tests/test_rl_sao_streaming_runtime.py | 2 +- tests/test_run_flags.py | 2 +- tests/test_wandb_tape.py | 4 +- tools/bench_sync_boundary.py | 163 ++++++++ yeto/launch/rl_args.py | 7 +- yeto/launch/syncer.py | 3 +- yeto/rl/engine/run_flags.py | 14 + yeto/syncer_profile.py | 2 +- yeto/wandb_tape.py | 2 + 13 files changed, 830 insertions(+), 75 deletions(-) create mode 100644 tools/bench_sync_boundary.py diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 513caabe0..8243ab299 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -98,6 +98,10 @@ struct Args { /// Write the snapshot every N outer steps (0 disables). #[arg(long, default_value_t = 8)] checkpoint_every: u64, + /// Policy sweep: write per-round checkpoints off the scheduler; each + /// round is broadcast only after its checkpoint is durable. + #[arg(long, default_value_t = false)] + async_checkpoint: bool, /// Resume from --checkpoint-path if it exists. #[arg(long, default_value_t = false)] resume: bool, @@ -273,6 +277,7 @@ fn main() -> anyhow::Result<()> { final_state: args.final_state, checkpoint_path: args.checkpoint_path, checkpoint_every: args.checkpoint_every, + async_checkpoint: args.async_checkpoint, resume: args.resume, mark_final_checkpoint: args.mark_final_checkpoint, learner_budget_steps: args.learner_budget_steps, diff --git a/syncer/src/server.rs b/syncer/src/server.rs index a5ed5b67a..e7b42b3e3 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -91,6 +91,12 @@ pub struct Config { /// the quiescent cut between rounds, resumed from when `resume` is set. pub checkpoint_path: Option, pub checkpoint_every: u64, + /// Policy-sweep only: persist each per-round checkpoint on a blocking + /// writer instead of the scheduler, and broadcast that round only after + /// the write is durable. Merge results, versions, and checkpoint bytes + /// are unchanged, so this is deliberately not part of the semantic + /// profile hash. + pub async_checkpoint: bool, pub resume: bool, /// Create the adjacent final marker only for a completed terminal cut. pub mark_final_checkpoint: bool, @@ -702,8 +708,11 @@ fn validate_config(cfg: &Config) -> Result<()> { if cfg.total_steps == 0 { bail!("--policy-sweep-fragments requires positive --total-steps"); } - if cfg.pipeline != 1 { - bail!("--policy-sweep-fragments requires --pipeline 1"); + // Fragment rounds of one sweep are independent (H=1, plain average, + // no momentum), so up to P of them may gather concurrently; commits + // stay in strict step order. + if cfg.pipeline == 0 || cfg.pipeline > fragments { + bail!("--policy-sweep-fragments requires --pipeline between 1 and the fragment count"); } if !cfg.total_steps.is_multiple_of(u64::from(fragments)) { bail!("--total-steps must be divisible by --policy-sweep-fragments"); @@ -733,6 +742,14 @@ fn validate_config(cfg: &Config) -> Result<()> { ); } } + if cfg.async_checkpoint { + if cfg.policy_sweep_fragments.is_none() { + bail!("--async-checkpoint requires --policy-sweep-fragments"); + } + if cfg.learner_budget_steps.is_some() { + bail!("--async-checkpoint does not support --learner-budget-steps"); + } + } Ok(()) } @@ -2116,12 +2133,15 @@ async fn scheduler( let mut computing = tokio::task::JoinSet::>::new(); let mut ready: BTreeMap = BTreeMap::new(); let mut busy_fragments: HashSet = HashSet::new(); + // --async-checkpoint: the committed round whose checkpoint is in flight. + let mut publishing: Option = None; let mut cutoff_established = false; let mut first_budget_report: Option<(Member, u64)> = None; 'outer: while next_t <= cfg.total_steps || !inflight.is_empty() || !computing.is_empty() || !ready.is_empty() + || publishing.is_some() { if budget_cutoff.deadline().is_some() { cutoff_established = true; @@ -2138,17 +2158,37 @@ async fn scheduler( let completed = take_next_commit(&mut ready, next_commit_t) .context("ready commit disappeared after linearization")?; let p = completed.round.p; - commit_round( - &cfg, - &mut st, - ®istry, - &mut last_sync_secs, - completed, - &budget_cutoff, - ) - .await?; - if !busy_fragments.remove(&p) { - bail!("fragment {p} lost its busy ownership at commit"); + if cfg.async_checkpoint { + // At most one checkpoint write is outstanding: the snapshot + // of t must be durable before t+1 is snapshotted, which also + // bounds host memory to one extra copy of the state. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } + let committed = apply_round_commit(&cfg, &mut st, completed)?; + publishing = Some(start_async_checkpoint(&cfg, &st, committed)?); + } else { + commit_round( + &cfg, + &mut st, + ®istry, + &mut last_sync_secs, + completed, + &budget_cutoff, + ) + .await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at commit"); + } } next_commit_t = next_commit_t .checked_add(1) @@ -2349,7 +2389,19 @@ async fn scheduler( if next_fragment_available { earliest = Some(earliest.map_or(next_launch, |d| d.min(next_launch))); } - let wake = match earliest { + let wake = if let Some(pending) = publishing.as_mut() { + // Same event-first bias; a durable checkpoint outranks compute + // so the oldest committed round is exposed first. + let sleep_until = earliest.map(tokio::time::Instant::from_std); + tokio::select! { + biased; + event = events.recv() => SchedulerWake::Event(event), + written = &mut pending.write => SchedulerWake::Durable(written), + result = computing.join_next(), if !computing.is_empty() => SchedulerWake::Computed(result), + _ = sleep_until_opt(sleep_until) => SchedulerWake::Deadline, + } + } else { + match earliest { Some(deadline) if !computing.is_empty() => tokio::select! { biased; wake = prefer_event_over_compute(events.recv(), computing.join_next()) => wake, @@ -2364,9 +2416,27 @@ async fn scheduler( prefer_event_over_compute(events.recv(), computing.join_next()).await } None => continue, + } }; match wake { SchedulerWake::Deadline => continue, + SchedulerWake::Durable(written) => { + let pending = publishing + .take() + .context("durable checkpoint without a pending publish")?; + finish_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending.committed, + pending.write_started, + written, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } SchedulerWake::Computed(Some(result)) => { let completed = result.context("iso compute task panicked")??; let t = completed.round.t; @@ -2469,6 +2539,20 @@ async fn scheduler( // slots in existing rounds. The helper revalidates the // generation because catch-up awaits socket writers and // a newer HELLO may supersede this one meanwhile. + // Catch-up must never expose a round whose checkpoint + // is not yet durable. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } if st.global_step < cfg.total_steps { send_all_fragments(&st, &group, &budget_cutoff).await; } @@ -2553,6 +2637,18 @@ async fn scheduler( // Catch-up must precede a replay even when the // disconnect event wins the race with the queued // replacement HELLO. + if let Some(pending) = publishing.take() { + settle_pending_publish( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + pending, + &budget_cutoff, + &mut busy_fragments, + ) + .await?; + } if st.global_step < cfg.total_steps { send_all_fragments(&st, &group, &budget_cutoff).await; } @@ -2640,6 +2736,7 @@ async fn scheduler( if next_commit_t != cfg.total_steps.saturating_add(1) || st.global_step != cfg.total_steps || !busy_fragments.is_empty() + || publishing.is_some() { bail!( "non-quiescent terminal cut: next_commit_t={next_commit_t} global_step={} busy={:?}", @@ -2935,6 +3032,14 @@ enum SchedulerWake { Deadline, Event(Option), Computed(Option, tokio::task::JoinError>>), + Durable(std::result::Result, tokio::task::JoinError>), +} + +async fn sleep_until_opt(deadline: Option) { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline).await, + None => std::future::pending().await, + } } type ComputeJoin = Option, tokio::task::JoinError>>; @@ -3189,6 +3294,24 @@ fn prepare_round_compute(cfg: &Config, st: &GlobalState, round: &Round) -> Resul st.prepare_merge(round.p, round.base_version, &outer_gradients, &weights) } +/// A round whose merge is applied to `GlobalState` but not yet externally +/// visible: no BCAST, no tape record, no log line. +struct CommittedRound { + t: u64, + p: usize, + base_version: u64, + attempt: u32, + started: Instant, + expected_members: Vec, + quorum_size: usize, + quorum_ms: Option, + grace_ms: Option, + pushes: HashMap, + gnorm: f64, + compute_secs: f64, + sync_start: Instant, +} + /// Commit an already-computed round. Called only for the contiguous t prefix, /// so Nesterov, versions, ledger, event tape, checkpoint, and broadcast all /// share one deterministic order even when SVD completion is out of order. @@ -3200,6 +3323,26 @@ async fn commit_round( completed: ComputedRound, budget_cutoff: &BudgetCutoff, ) -> Result<()> { + let committed = apply_round_commit(cfg, st, completed)?; + let mut checkpoint_ms = None; + if checkpoint_before_broadcast(cfg) { + if let Some(path) = cfg.checkpoint_path.as_ref() { + let checkpoint_start = Instant::now(); + st.save_checkpoint(path)?; + checkpoint_ms = Some(checkpoint_start.elapsed().as_millis() as u64); + info!(step = committed.t, path = %path.display(), "checkpoint committed"); + } + } + publish_round(cfg, st, registry, last_sync_secs, committed, checkpoint_ms, budget_cutoff).await +} + +/// Apply the merge, version, step, and ledger updates of the next contiguous +/// round. Nothing leaves the process here. +fn apply_round_commit( + cfg: &Config, + st: &mut GlobalState, + completed: ComputedRound, +) -> Result { let ComputedRound { round, merge, @@ -3237,7 +3380,6 @@ async fn commit_round( st.global_step = t; let mut ordered_pushes: Vec<_> = pushes.iter().collect(); ordered_pushes.sort_unstable_by_key(|(member, _)| **member); - let responders: Vec = ordered_pushes.iter().map(|(member, _)| **member).collect(); for (_, push) in ordered_pushes { if let Some(fragments) = cfg.policy_sweep_fragments { st.record_fragment_merge( @@ -3252,15 +3394,52 @@ async fn commit_round( st.record_merge(push.learner_id, push.c_steps, push.c_tokens); } } + Ok(CommittedRound { + t, + p, + base_version, + attempt, + started, + expected_members, + quorum_size, + quorum_ms, + grace_ms, + pushes, + gnorm, + compute_secs, + sync_start, + }) +} - let checkpoint_each_round = checkpoint_before_broadcast(cfg); - if checkpoint_each_round { - if let Some(path) = cfg.checkpoint_path.as_ref() { - st.save_checkpoint(path)?; - info!(step = t, path = %path.display(), "checkpoint committed"); - } - } - +/// Expose a committed round: BCAST its fragment, log it, append its tape +/// record, and write the periodic (non-per-round) checkpoint. A per-round +/// checkpoint must already be durable when this is called. +async fn publish_round( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + committed: CommittedRound, + checkpoint_ms: Option, + budget_cutoff: &BudgetCutoff, +) -> Result<()> { + let CommittedRound { + t, + p, + base_version, + attempt, + started, + expected_members, + quorum_size, + quorum_ms, + grace_ms, + pushes, + gnorm, + compute_secs, + sync_start, + } = committed; + let mut responders: Vec = pushes.keys().copied().collect(); + responders.sort_unstable(); // A disconnected learner can catch up after reconnecting. Deterministic // local encoding/allocation failures must terminate the coordinator // instead of silently committing a cut that was never broadcastable. @@ -3318,6 +3497,7 @@ async fn commit_round( quorum_ms, grace_ms, sync_ms, + checkpoint_ms, &pushes, gnorm, ms, @@ -3330,7 +3510,7 @@ async fn commit_round( // have also been broadcast; the terminal cut is delivered by FINAL. // A crash-resume loses those gathers; their fragments simply merge on // a later cycle, which the quorum design already tolerates. - if !checkpoint_each_round { + if !checkpoint_before_broadcast(cfg) { if let Some(path) = &cfg.checkpoint_path { if cfg.checkpoint_every > 0 && t % cfg.checkpoint_every == 0 { st.save_checkpoint(path)?; @@ -3341,6 +3521,97 @@ async fn commit_round( Ok(()) } +/// A committed round whose per-round checkpoint is being written by a +/// blocking task (`--async-checkpoint`). Its fragment stays busy and its +/// BCAST is withheld until the write is durable, so no learner can observe +/// a version a crash could lose. +struct PendingPublish { + committed: CommittedRound, + write_started: Instant, + write: tokio::task::JoinHandle>, +} + +/// Snapshot the just-committed state and start persisting it off the +/// scheduler thread. +fn start_async_checkpoint( + cfg: &Config, + st: &GlobalState, + committed: CommittedRound, +) -> Result { + let path = cfg + .checkpoint_path + .clone() + .context("--async-checkpoint requires --checkpoint-path")?; + let write_started = Instant::now(); + let bytes = st.encode_checkpoint()?; + let write = tokio::task::spawn_blocking(move || { + crate::state::write_checkpoint_bytes(&path, &bytes) + }); + Ok(PendingPublish { + committed, + write_started, + write, + }) +} + +/// Wait for a pending per-round checkpoint, then publish its round. A failed +/// write fails the coordinator before the round becomes visible; the +/// previous checkpoint stays the resume point. +async fn settle_pending_publish( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + pending: PendingPublish, + budget_cutoff: &BudgetCutoff, + busy_fragments: &mut HashSet, +) -> Result<()> { + let PendingPublish { + committed, + write_started, + write, + } = pending; + finish_pending_publish( + cfg, + st, + registry, + last_sync_secs, + committed, + write_started, + write.await, + budget_cutoff, + busy_fragments, + ) + .await +} + +#[allow(clippy::too_many_arguments)] +async fn finish_pending_publish( + cfg: &Config, + st: &GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + committed: CommittedRound, + write_started: Instant, + written: std::result::Result, tokio::task::JoinError>, + budget_cutoff: &BudgetCutoff, + busy_fragments: &mut HashSet, +) -> Result<()> { + let (t, p) = (committed.t, committed.p); + written + .context("async checkpoint task panicked")? + .with_context(|| format!("async checkpoint for step {t} failed; step {t} not broadcast"))?; + let checkpoint_ms = Some(write_started.elapsed().as_millis() as u64); + if let Some(path) = cfg.checkpoint_path.as_ref() { + info!(step = t, path = %path.display(), "checkpoint committed"); + } + publish_round(cfg, st, registry, last_sync_secs, committed, checkpoint_ms, budget_cutoff).await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at publish"); + } + Ok(()) +} + #[cfg(test)] async fn complete_round( cfg: &Config, @@ -3361,6 +3632,30 @@ async fn complete_round( commit_round(cfg, st, registry, last_sync_secs, completed, &cutoff).await } +/// Test twin of `complete_round` through the `--async-checkpoint` path. +#[cfg(test)] +async fn complete_round_async( + cfg: &Config, + st: &mut GlobalState, + registry: &Registry, + last_sync_secs: &mut f64, + round: Round, +) -> Result<()> { + let prepared = prepare_round_compute(cfg, st, &round)?; + let compute_start = Instant::now(); + let merge = prepared.compute().await?; + let completed = ComputedRound { + round, + merge, + compute_secs: compute_start.elapsed().as_secs_f64(), + }; + let cutoff = BudgetCutoff::new(false, Duration::ZERO); + let mut busy = HashSet::from([completed.round.p]); + let committed = apply_round_commit(cfg, st, completed)?; + let pending = start_async_checkpoint(cfg, st, committed)?; + settle_pending_publish(cfg, st, registry, last_sync_secs, pending, &cutoff, &mut busy).await +} + fn new_state_for(group: &Arc, cfg: &Config) -> Result { if let Some(fragments) = cfg.policy_sweep_fragments { if group.dtype != DTYPE_F32 { @@ -3856,6 +4151,7 @@ fn append_tape( quorum_ms: Option, grace_ms: Option, sync_ms: u64, + checkpoint_ms: Option, pushes: &HashMap, gnorm: f64, ms: u64, @@ -3932,9 +4228,10 @@ fn append_tape( .map(|byte| format!("{byte:02x}")) .collect(); let merge_seconds = sync_ms as f64 / 1000.0; + let checkpoint_field = checkpoint_tape_field(checkpoint_ms); let sweep_fields = policy_sweep_event_fields(step, policy_sweep_fragments); let line = format!( - "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment}{sweep_fields},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms},\"time_unix\":{:?},\"responders\":[{}]}}\n", + "{{\"protocol_version\":{PROTOCOL_VERSION},\"delta_semantics\":\"local_minus_raw_anchor\",\"sync/layout_hash\":\"{layout_hash}\",\"sync/base_version\":{launch_base_version},\"sync/responders\":{},\"sync/quorum\":{quorum},\"sync/rejected_stale_updates\":0,\"sync/merge_seconds\":{merge_seconds},\"sync/global_delta_norm\":{gnorm},\"step\":{step},\"fragment\":{fragment}{sweep_fields},\"launch_base_version\":{launch_base_version},\"attempt\":{attempt},\"gnorm\":{gnorm},\"ms\":{ms},\"quorum\":{quorum},\"expected\":{},\"expected_members\":{},\"responded\":{},\"responded_members\":{},\"missed_grace\":{},\"missed_members\":{},\"quorum_ms\":{quorum_ms},\"grace_ms\":{grace_ms},\"sync_ms\":{sync_ms}{checkpoint_field},\"time_unix\":{:?},\"responders\":[{}]}}\n", responded.len(), json_ids(&expected_learners), json_members(expected_members), @@ -3954,6 +4251,12 @@ fn append_tape( Ok(()) } +/// Optional top-level tape field: wall time of the per-round checkpoint +/// (synchronous save, or snapshot + async write until durable). +fn checkpoint_tape_field(checkpoint_ms: Option) -> String { + checkpoint_ms.map_or_else(String::new, |ms| format!(",\"checkpoint_ms\":{ms}")) +} + fn json_ids(ids: &[u32]) -> String { let mut ids = ids.to_vec(); ids.sort_unstable(); @@ -4354,6 +4657,7 @@ mod tests { final_state: None, checkpoint_path: None, checkpoint_every: 0, + async_checkpoint: false, resume: false, mark_final_checkpoint: false, learner_budget_steps: None, @@ -4618,10 +4922,30 @@ mod tests { assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) .contains("requires positive --total-steps")); + let mut streamed = sweep_test_config(2, 4); + streamed.pipeline = 2; + streamed.async_checkpoint = true; + assert!(validate_config(&streamed).is_ok()); + for pipeline in [0, 3] { + invalid = sweep_test_config(2, 4); + invalid.pipeline = pipeline; + assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) + .contains("requires --pipeline between 1 and the fragment count")); + } + let mut legacy_async = round_test_config(3); + legacy_async.async_checkpoint = true; + assert!(format!("{:#}", validate_config(&legacy_async).unwrap_err()) + .contains("--async-checkpoint requires --policy-sweep-fragments")); invalid = sweep_test_config(2, 4); - invalid.pipeline = 2; + invalid.async_checkpoint = true; + invalid.learner_budget_steps = Some(1); assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) - .contains("requires --pipeline 1")); + .contains("does not support --learner-budget-steps")); + // Async persistence is not semantic: the profile hash ignores it. + let mut hashed = sweep_test_config(2, 4); + let before = hashed.semantic_profile_hash(); + hashed.async_checkpoint = true; + assert_eq!(before, hashed.semantic_profile_hash()); invalid = sweep_test_config(3, 4); assert!( @@ -5244,6 +5568,120 @@ mod tests { std::fs::remove_dir_all(&directory).ok(); } + fn unique_test_directory(tag: &str) -> std::path::PathBuf { + let directory = std::env::temp_dir().join(format!( + "yeto-{tag}-{}-{}", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + std::fs::create_dir_all(&directory).unwrap(); + directory + } + + #[tokio::test] + async fn async_checkpoint_matches_the_synchronous_path_bit_for_bit() { + let directory = unique_test_directory("async-checkpoint-equal"); + let learner = member(0, 30); + let mut outputs = Vec::new(); + for asynchronous in [false, true] { + let checkpoint = directory.join(format!("state-{asynchronous}.ckpt")); + let tape = directory.join(format!("tape-{asynchronous}.jsonl")); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = round_test_config(4); + config.pipeline = 2; + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 1; + config.async_checkpoint = asynchronous; + config.event_tape = Some(tape.clone()); + let mut state = two_fragment_streaming_state(); + let mut sync_seconds = 0.0; + for step in 1..=2u64 { + let round = sweep_test_round(step, (step - 1) as usize, 1, step, learner); + if asynchronous { + complete_round_async(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } else { + complete_round(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } + // Every non-terminal round is broadcast exactly once. + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + } + let tape_text = std::fs::read_to_string(&tape).unwrap(); + assert_eq!(tape_text.lines().count(), 2); + assert!(tape_text.lines().all(|line| line.contains("\"checkpoint_ms\":"))); + outputs.push((state.params.clone(), std::fs::read(&checkpoint).unwrap())); + // The snapshot encoder writes exactly the streaming file format. + let mut on_disk = two_fragment_streaming_state(); + on_disk.load_checkpoint(&checkpoint).unwrap(); + assert_eq!(on_disk.encode_checkpoint().unwrap(), std::fs::read(&checkpoint).unwrap()); + } + assert_eq!(outputs[0].0, outputs[1].0); + assert_eq!(outputs[0].1, outputs[1].1); + std::fs::remove_dir_all(&directory).ok(); + } + + #[tokio::test] + async fn async_checkpoint_failure_withholds_broadcast_and_keeps_previous_checkpoint() { + let directory = unique_test_directory("async-checkpoint-failure"); + let checkpoint = directory.join("state.ckpt"); + let learner = member(0, 30); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = round_test_config(4); + config.pipeline = 2; + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 1; + config.async_checkpoint = true; + let mut state = two_fragment_streaming_state(); + let mut sync_seconds = 0.0; + complete_round_async( + &config, + &mut state, + ®istry, + &mut sync_seconds, + sweep_test_round(1, 0, 1, 1, learner), + ) + .await + .unwrap(); + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + // The atomic writer creates this exact temporary path; a directory + // there makes the next write fail deterministically. + std::fs::create_dir(checkpoint.with_extension("tmp")).unwrap(); + let error = complete_round_async( + &config, + &mut state, + ®istry, + &mut sync_seconds, + sweep_test_round(2, 1, 1, 1, learner), + ) + .await + .unwrap_err(); + assert!(format!("{error:#}").contains("step 2 not broadcast")); + assert!(matches!( + frames.try_recv(), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) + )); + // The resume point is still the last durable, broadcast cut. + let mut recovered = two_fragment_streaming_state(); + recovered.load_checkpoint(&checkpoint).unwrap(); + assert_eq!(recovered.global_step, 1); + assert_eq!(recovered.versions, vec![1, 0]); + std::fs::remove_dir_all(&directory).ok(); + } + #[tokio::test] async fn policy_sweep_checkpoint_resume_mid_sweep_does_not_double_account() { let directory = std::env::temp_dir().join(format!( @@ -6227,6 +6665,7 @@ mod tests { Some(11), Some(22), 33, + None, &pushes, 0.5, 44, @@ -6275,6 +6714,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, @@ -6303,6 +6743,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, @@ -6410,6 +6851,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, diff --git a/syncer/src/state.rs b/syncer/src/state.rs index f19ce4c99..6d4415612 100644 --- a/syncer/src/state.rs +++ b/syncer/src/state.rs @@ -628,6 +628,36 @@ pub fn remove_final_marker(path: &std::path::Path) -> Result<()> { } } +/// Little-endian f32 values in 64 KiB blocks (same bytes as per-value +/// `to_le_bytes`, without one writer call per element). +fn write_f32s(f: &mut W, values: &[f32]) -> Result<()> { + let mut block = Vec::with_capacity(16 * 1024 * 4); + for chunk in values.chunks(16 * 1024) { + block.clear(); + for v in chunk { + block.extend_from_slice(&v.to_le_bytes()); + } + f.write_all(&block)?; + } + Ok(()) +} + +/// Atomically persist checkpoint bytes from `encode_checkpoint`: write a +/// sibling temporary file, fsync it, rename it over `path`, fsync the +/// directory. A failure at any point leaves the previous checkpoint intact. +pub fn write_checkpoint_bytes(path: &std::path::Path, bytes: &[u8]) -> Result<()> { + use std::io::Write; + let tmp = path.with_extension("tmp"); + { + let mut f = std::fs::File::create(&tmp) + .with_context(|| format!("create checkpoint temporary {}", tmp.display()))?; + f.write_all(bytes)?; + f.sync_all()?; + } + std::fs::rename(&tmp, path)?; + sync_parent_directory(path) +} + fn sync_parent_directory(path: &std::path::Path) -> Result<()> { let parent = path.parent().unwrap_or_else(|| std::path::Path::new(".")); std::fs::File::open(parent) @@ -660,47 +690,11 @@ 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<()> { - use std::io::Write; - let layout_fingerprint = self.layout_fingerprint; let tmp = path.with_extension("tmp"); { let mut f = std::io::BufWriter::new(std::fs::File::create(&tmp)?); - f.write_all(&CKPT_MAGIC_V3.to_le_bytes())?; - f.write_all(&[self.iso_backend_kind() as u8])?; - f.write_all(&layout_fingerprint)?; - f.write_all(&self.global_step.to_le_bytes())?; - f.write_all(&(self.params.len() as u32).to_le_bytes())?; - for p in 0..self.params.len() { - f.write_all(&self.versions[p].to_le_bytes())?; - f.write_all(&(self.params[p].len() as u64).to_le_bytes())?; - for v in &self.params[p] { - f.write_all(&v.to_le_bytes())?; - } - for v in &self.momentum[p] { - f.write_all(&v.to_le_bytes())?; - } - } - f.write_all(&(self.ledger.len() as u32).to_le_bytes())?; - for (id, l) in &self.ledger { - f.write_all(&id.to_le_bytes())?; - f.write_all(&l.merges.to_le_bytes())?; - f.write_all(&l.steps.to_le_bytes())?; - f.write_all(&l.tokens.to_le_bytes())?; - } - f.write_all(&self.layout_fingerprint)?; - if let Some(fragments) = self.policy_sweep_fragments { - f.write_all(&POLICY_SWEEP_CKPT_MAGIC.to_le_bytes())?; - f.write_all(&fragments.to_le_bytes())?; - f.write_all( - &self - .session_contract_hash - .context("policy-sweep checkpoint requires a session contract hash")?, - )?; - } else if let Some(contract_hash) = self.session_contract_hash { - f.write_all(&SESSION_CONTRACT_CKPT_MAGIC.to_le_bytes())?; - f.write_all(&contract_hash)?; - } - f.flush()?; + self.write_checkpoint_to(&mut f)?; + std::io::Write::flush(&mut f)?; f.get_ref().sync_all()?; } std::fs::rename(&tmp, path)?; @@ -708,6 +702,54 @@ impl GlobalState { Ok(()) } + /// The exact bytes `save_checkpoint` writes, as an owned snapshot. + /// + /// The asynchronous checkpoint path encodes on the scheduler (one memory + /// copy of the consistent cut) and hands the bytes to a blocking writer, + /// so later commits may mutate `self` while the snapshot is persisted. + pub fn encode_checkpoint(&self) -> Result> { + let state_floats: usize = self.params.iter().map(Vec::len).sum::() * 2; + let mut bytes = Vec::with_capacity(state_floats * 4 + 4096); + self.write_checkpoint_to(&mut bytes)?; + Ok(bytes) + } + + fn write_checkpoint_to(&self, f: &mut W) -> Result<()> { + let layout_fingerprint = self.layout_fingerprint; + f.write_all(&CKPT_MAGIC_V3.to_le_bytes())?; + f.write_all(&[self.iso_backend_kind() as u8])?; + f.write_all(&layout_fingerprint)?; + f.write_all(&self.global_step.to_le_bytes())?; + f.write_all(&(self.params.len() as u32).to_le_bytes())?; + for p in 0..self.params.len() { + f.write_all(&self.versions[p].to_le_bytes())?; + f.write_all(&(self.params[p].len() as u64).to_le_bytes())?; + write_f32s(f, &self.params[p])?; + write_f32s(f, &self.momentum[p])?; + } + f.write_all(&(self.ledger.len() as u32).to_le_bytes())?; + for (id, l) in &self.ledger { + f.write_all(&id.to_le_bytes())?; + f.write_all(&l.merges.to_le_bytes())?; + f.write_all(&l.steps.to_le_bytes())?; + f.write_all(&l.tokens.to_le_bytes())?; + } + f.write_all(&self.layout_fingerprint)?; + if let Some(fragments) = self.policy_sweep_fragments { + f.write_all(&POLICY_SWEEP_CKPT_MAGIC.to_le_bytes())?; + f.write_all(&fragments.to_le_bytes())?; + f.write_all( + &self + .session_contract_hash + .context("policy-sweep checkpoint requires a session contract hash")?, + )?; + } else if let Some(contract_hash) = self.session_contract_hash { + f.write_all(&SESSION_CONTRACT_CKPT_MAGIC.to_le_bytes())?; + f.write_all(&contract_hash)?; + } + Ok(()) + } + /// Restore params/momentum/versions/step/ledger from a snapshot. /// /// V3 binds the snapshot to both the ISO backend and exact semantic HELLO diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index fd14e9973..0976992ad 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -41,23 +41,25 @@ def _hmac_key(monkeypatch): monkeypatch.setenv("YETO_ISLAND_HMAC_KEY", "dense-sweep-test-key") -def _dense_args(): +def _dense_args(*extra: str, total_steps: int = 1): from yeto.launcher import _prepare_rl_args args = _launcher_args("ports", ( "--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", "--rl-rollout-gpus", "1", "--rl-parameter-layout-sha256", "e" * 64, - "--fragments", str(FRAGMENTS), "--total-steps", "1"), gpu="modal:2xh100") + "--fragments", str(FRAGMENTS), "--total-steps", str(total_steps), *extra), + gpu="modal:2xh100") args.controller = "local" _prepare_rl_args(args) return args -def _start_syncer(binary: Path, tmp_path: Path) -> tuple[subprocess.Popen, int, str]: +def _start_syncer(binary: Path, tmp_path: Path, *extra: str, + total_steps: int = 1) -> tuple[subprocess.Popen, int, str]: """The launcher's dense actor syncer, relocated to a free port and tmp_path.""" from yeto import launcher - args = _dense_args() + args = _dense_args(*extra, total_steps=total_steps) command = launcher.role_syncer_commands(args, 1, binary=str(binary))["actor"] with socket.socket() as probe: probe.bind(("127.0.0.1", 0)) @@ -153,3 +155,80 @@ def test_launcher_dense_syncer_refuses_another_profile_hash(syncer_binary, tmp_p client.close() log = _stop(process) assert "syncer semantic profile does not match" in log + + +STREAMED = ("--rl-sync-fragment-pipeline", str(FRAGMENTS), "--rl-sync-async-checkpoint") + + +def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int): + """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint).""" + process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) + client = _client(port, profile_hash) + concurrent = 0 + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32([float(fid) + 0.125 * i for i in range(NUMEL)])) + for sweep in range(1, sweeps + 1): + permits: dict[int, object] = {} + + def gather(): + client.check_health() + for permit in client.drain_pulls(): + permits[permit.fragment_id] = permit + return any(not getattr(p, "pushed", False) for p in permits.values()) + + while len(permits) < FRAGMENTS: + _wait(gather, f"no PULL in sweep {sweep}") + time.sleep(0.2) # let every PULL the profile permits arrive + gather() + concurrent = max(concurrent, sum(not getattr(p, "pushed", False) + for p in permits.values())) + # Push only what is permitted now; the serial profile then + # releases the next fragment's PULL after this commit. + for fid, permit in sorted(permits.items()): + if getattr(permit, "pushed", False): + continue + base = 0 if sweep == 1 else (sweep - 2) * FRAGMENTS + fid + 1 + delta = [0.5 + fid + 0.03125 * sweep + 0.001 * i for i in range(NUMEL)] + client.push_fragment(fid, permit.global_step, permit.round_attempt, + base, sweep, 1, 7, _f32(delta)) + permit.pushed = True + manifest, finals = client.wait_for_final_fragments(timeout=30) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + log = _stop(process) + assert "syncer semantic profile" not in log, log + return concurrent, list(manifest.versions), result, (tmp_path / "state.ckpt").read_bytes(), log + + +def test_streamed_sweep_matches_the_serial_sweep_bit_for_bit(syncer_binary, tmp_path): + """E3: fragment pipelining + async checkpoint changes timing, not results.""" + serial_dir, streamed_dir = tmp_path / "serial", tmp_path / "streamed" + serial_dir.mkdir() + streamed_dir.mkdir() + serial = _run_sweeps(syncer_binary, serial_dir, (), sweeps=3) + streamed = _run_sweeps(syncer_binary, streamed_dir, STREAMED, sweeps=3) + assert serial[0] == 1, "the serial profile must permit one fragment at a time" + assert streamed[0] == FRAGMENTS, "the streamed profile must permit a whole sweep at once" + assert serial[1] == streamed[1] == [5, 6] + assert serial[2] == streamed[2] + assert serial[3] == streamed[3] + assert "--async-checkpoint" not in serial[4] + + +def test_streamed_flags_bind_the_profile_and_are_full_parameter_only(): + serial = _dense_args() + streamed = _dense_args(*STREAMED) + assert serial.pipeline == 1 and streamed.pipeline == FRAGMENTS + from yeto import launcher + + command = launcher.role_syncer_commands(streamed, 1, binary="yeto-syncer")["actor"] + assert f"--pipeline {FRAGMENTS}" in command and "--async-checkpoint" in command + assert (launcher.rl_syncer_profile(serial, 1).sha256 + != launcher.rl_syncer_profile(streamed, 1).sha256) + with pytest.raises(ValueError, match="--rl-sync-fragment-pipeline between 1 and --fragments"): + _dense_args("--rl-sync-fragment-pipeline", str(FRAGMENTS + 1)) diff --git a/tests/test_rl_sao_streaming_runtime.py b/tests/test_rl_sao_streaming_runtime.py index ced8e08eb..88ec57b2b 100644 --- a/tests/test_rl_sao_streaming_runtime.py +++ b/tests/test_rl_sao_streaming_runtime.py @@ -279,7 +279,7 @@ def test_load_or_role_pairing_rejects_role_or_lockstep_drift(tmp_path, mutation, ("checkpoint_every", 2, "checkpointing every round"), ("resume", False, "checkpointing every round"), ("max_base_lag", 1, "max_base_lag=0"), - ("policy_sweep_fragments", 4, "policy sweep"), + ("policy_sweep_fragments", 4, "policy.sweep"), ("learner_budget_steps", 2, "learner-budget cutoff mode"), ), ) diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 9ef5c8eab..99b5ea017 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -250,7 +250,7 @@ def test_forward_tables_name_flags_of_both_layers(): # flag renamed in the schema but still spelled the old way in the launcher fails. _NON_RL_LITERALS = frozenset({ "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter", - "--allow-unauthenticated-islands", "--bind", "--cap-add", "--carry-gamma", + "--allow-unauthenticated-islands", "--async-checkpoint", "--bind", "--cap-add", "--carry-gamma", "--checkpoint-every", "--checkpoint-path", "--codex-harness-binary", "--config-json", "--create", "--dashboard-host", "--dearmor", "--depth", "--detach", "--dtype", "--entrypoint", "--expected-island-contract", "--final-ack-timeout-s", "--final-grace-s", "--follow", "--force", diff --git a/tests/test_wandb_tape.py b/tests/test_wandb_tape.py index 0a9abf088..2305d7fcd 100644 --- a/tests/test_wandb_tape.py +++ b/tests/test_wandb_tape.py @@ -176,13 +176,15 @@ def test_every_rust_tape_field_is_consumed_or_deliberately_skipped(): # The responders literal is built first, the record literal last. responder_fmt, record_fmt = body.split("let line = format!(", 1) top_level = set(re.findall(r'\\"(\w+)\\":', record_fmt)) + optional = source.split("fn checkpoint_tape_field(", 1)[1].split("\nfn ", 1)[0] + top_level |= set(re.findall(r'\\"(\w+)\\":', optional)) responder_keys = set(re.findall(r'\\"(\w+)\\":', responder_fmt)) assert responder_keys, "could not find the responder literal in append_tape" consumed_top = { "step", "fragment", "attempt", "gnorm", "ms", "sync_ms", "quorum", "expected", "responded", "missed_grace", "launch_base_version", - "quorum_ms", "grace_ms", "responders", + "quorum_ms", "grace_ms", "responders", "checkpoint_ms", } # Constants and the member-identity mirrors of the id lists: nothing a # time series can be drawn from. diff --git a/tools/bench_sync_boundary.py b/tools/bench_sync_boundary.py new file mode 100644 index 000000000..3b4496322 --- /dev/null +++ b/tools/bench_sync_boundary.py @@ -0,0 +1,163 @@ +"""CPU micro-benchmark of the full-parameter sync boundary (S26 E3). + +Runs the real Rust syncer on localhost with two in-process learners and a +synthetic f32 model, and times each strict sweep from the first push to the +last updated fragment received by every learner. It compares the serial +profile (pipeline 1, checkpoint on the scheduler) with the streamed profile +(pipeline P, ``--async-checkpoint``). + +These are local CPU numbers: loopback transfer and a local NVMe checkpoint. +They show where the boundary time goes on this machine; they do not predict +GPU-cluster or WAN timings. + + python tools/bench_sync_boundary.py --floats-per-fragment 33554432 --sweeps 4 +""" + +from __future__ import annotations + +import argparse +import json +import socket +import statistics +import subprocess +import tempfile +import threading +import time +from pathlib import Path + +import numpy as np + +from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout +from yeto.protocol import DTYPE_F32, SyncerClient + +ROOT = Path(__file__).resolve().parent.parent + + +def _free_port() -> int: + with socket.socket() as probe: + probe.bind(("127.0.0.1", 0)) + return probe.getsockname()[1] + + +def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: int, + sweeps: int, streamed: bool) -> subprocess.Popen: + command = [ + str(binary), "--port", str(port), "--learners", str(learners), + "--quorum", str(learners), "--grace-ms", "0", + "--pipeline", str(fragments if streamed else 1), + "--sync-interval-steps", "0", "--delta-correction", "none", + "--total-steps", str(sweeps * fragments), "--policy-sweep-fragments", str(fragments), + "--outer-lr", "1", "--outer-momentum", "0", "--max-base-lag", "0", + "--learner-weight", "equal", "--checkpoint-path", str(workdir / "state.ckpt"), + "--checkpoint-every", "1", "--event-tape", str(workdir / "tape.jsonl"), + "--allow-unauthenticated-islands", # loopback-only benchmark + ] + if streamed: + command.append("--async-checkpoint") + return subprocess.Popen(command, stdout=open(workdir / "syncer.log", "w"), + stderr=subprocess.STDOUT) + + +def _learner(learner_id: int, port: int, layout: FragmentLayout, numel: int, sweeps: int, + timings: list[list[float]], errors: list[BaseException]) -> None: + fragments = len(layout.fragments) + client = SyncerClient(("127.0.0.1", port), learner_id, layout, dtype=DTYPE_F32, + num_streams=1, connect_timeout=30, max_reconnects=0, + finalization_timeout=600) + rng = np.random.default_rng(learner_id) + try: + client.start() + if learner_id == 0: + for fid in range(fragments): + client.send_init(fid, rng.standard_normal(numel, dtype=np.float32).tobytes()) + deltas = [rng.standard_normal(numel, dtype=np.float32).tobytes() for _ in range(fragments)] + carried: list = [] # a PULL of the next sweep can arrive before this sweep's BCAST + for sweep in range(1, sweeps + 1): + pushed: set[int] = set() + updated: set[int] = set() + first_push = None + while len(updated) < fragments or len(pushed) < fragments: + client.check_health() + permits, carried = carried + client.drain_pulls(), [] + for permit in permits: + if permit.fragment_id in pushed: + carried.append(permit) + continue + base = 0 if sweep == 1 else (sweep - 2) * fragments + permit.fragment_id + 1 + first_push = first_push or time.monotonic() + client.push_fragment(permit.fragment_id, permit.global_step, + permit.round_attempt, base, sweep, 1, 1, + deltas[permit.fragment_id]) + pushed.add(permit.fragment_id) + for update in client.drain_updates(): + if update.version > (sweep - 1) * fragments: + updated.add(update.fragment_id) + if sweep == sweeps and len(pushed) == fragments: + break # the terminal sweep is delivered by FINAL, not BCAST + time.sleep(0.001) + if sweep == sweeps: + manifest, _ = client.wait_for_final_fragments(timeout=600) + timings[learner_id].append(time.monotonic() - first_push) + client.acknowledge_finalization(manifest) + else: + timings[learner_id].append(time.monotonic() - first_push) + except BaseException as error: # noqa: BLE001 - reported by the caller + errors.append(error) + finally: + client.close() + + +def run(binary: Path, *, streamed: bool, learners: int, fragments: int, numel: int, + sweeps: int) -> dict: + layout = FragmentLayout([Fragment(MERGE_RDA, [(f"w{i}", numel)]) for i in range(fragments)]) + with tempfile.TemporaryDirectory(prefix="e3-bench-", dir="/tmp") as tmp: + workdir = Path(tmp) + port = _free_port() + process = _syncer(binary, workdir, port=port, learners=learners, fragments=fragments, + sweeps=sweeps, streamed=streamed) + timings: list[list[float]] = [[] for _ in range(learners)] + errors: list[BaseException] = [] + threads = [threading.Thread(target=_learner, args=(i, port, layout, numel, sweeps, + timings, errors)) + for i in range(learners)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + code = process.wait(timeout=120) + if errors or code != 0: + raise RuntimeError(f"benchmark failed (syncer exit {code}): {errors!r}; " + f"log: {(workdir / 'syncer.log').read_text()[-2000:]}") + tape = [json.loads(line) for line in (workdir / "tape.jsonl").read_text().splitlines() + if '"sync_ms"' in line] + # Sweep 1 includes connection set-up; the last sweep ends with FINAL. + steady = [max(per[i] for per in timings) for i in range(1, sweeps - 1)] or \ + [max(per[i] for per in timings) for i in range(sweeps)] + return { + "profile": "streamed" if streamed else "serial", + "boundary_s_per_sweep": [round(x, 3) for x in steady], + "boundary_s_median": round(statistics.median(steady), 3), + "checkpoint_ms_median": statistics.median(r["checkpoint_ms"] for r in tape), + "sync_ms_median": statistics.median(r["sync_ms"] for r in tape), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--binary", type=Path, default=ROOT / "syncer/target/release/yeto-syncer") + parser.add_argument("--learners", type=int, default=2) + parser.add_argument("--fragments", type=int, default=2) + parser.add_argument("--floats-per-fragment", type=int, default=32 * 1024 * 1024) + parser.add_argument("--sweeps", type=int, default=4) + args = parser.parse_args() + results = [run(args.binary, streamed=streamed, learners=args.learners, + fragments=args.fragments, numel=args.floats_per_fragment, sweeps=args.sweeps) + for streamed in (False, True)] + print(json.dumps({"note": "CPU localhost numbers; not GPU or WAN timings", + "learners": args.learners, "fragments": args.fragments, + "floats_per_fragment": args.floats_per_fragment, "results": results}, + indent=2)) + + +if __name__ == "__main__": + main() diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 0c21669b9..e1ce4ff28 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -169,11 +169,16 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: problems.append("--rl-island-scheduling legacy (no fragment frames in elastic mode)") if rl_needs_critic(args): problems.append("a critic-free algorithm (the strict sweep trains the actor only)") + fragment_pipeline = int(getattr(args, "rl_sync_fragment_pipeline", 1)) + if not 1 <= fragment_pipeline <= max(args.fragments, 1): + problems.append("--rl-sync-fragment-pipeline between 1 and --fragments") if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands args.grace_ms = 0 - args.pipeline = 1 + # Fragment rounds of one sweep are independent (H=1, plain average), so the + # syncer may gather several at once; it still commits them in step order. + args.pipeline = fragment_pipeline args.sync_interval_steps = 0.0 args.delta_correction = "none" args.outer_lr = 1.0 diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 93b1aa6aa..7d951ae9b 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -161,7 +161,8 @@ def _syncer_profile_binding(args) -> str: full-parameter adds strict P-fragment policy sweeps; sao-streaming-full adds the frozen contract's timeouts and the publishable final checkpoint.""" if rl_dense_sweep(args): - return f" --policy-sweep-fragments {args.fragments} --require-profile-binding" + async_checkpoint = " --async-checkpoint" if getattr(args, "rl_sync_async_checkpoint", False) else "" + return f" --policy-sweep-fragments {args.fragments} --require-profile-binding{async_checkpoint}" if rl_sao_streaming(args): profile = args.rl_sao_manifest.profile return (f" --quorum-timeout-s {profile.quorum_timeout_s}" diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index b4de91487..f9be06f6f 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -841,6 +841,20 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "sao-streaming-full: full-parameter SAO, actor and critic streams " "(ports engine, --rl-placement fixed-partition, --rl-sao-contracts)", ), + cli_flag( + "--rl-sync-fragment-pipeline", + type=int, + default=1, + help="full-parameter only: fragment rounds the syncer gathers at once " + "(1..--fragments). Above 1, fragment k+1 uploads while fragment k merges " + "and downloads; merge results are unchanged (default 1 = serial)", + ), + cli_flag( + "--rl-sync-async-checkpoint", + action="store_true", + help="full-parameter only: the syncer writes each per-round checkpoint off " + "its scheduler and broadcasts that round only once the checkpoint is durable", + ), cli_flag( "--rl-initial-adapter", default=None, diff --git a/yeto/syncer_profile.py b/yeto/syncer_profile.py index 390d2bc45..c5e011b08 100644 --- a/yeto/syncer_profile.py +++ b/yeto/syncer_profile.py @@ -259,7 +259,7 @@ def _validate_general(self) -> None: raise ValueError("syncer profile resume requires checkpointing") if self.policy_sweep_fragments is not None: fragments = self.policy_sweep_fragments - if self.pipeline != 1 or self.total_steps % fragments: + if not 1 <= self.pipeline <= fragments or self.total_steps % fragments: raise ValueError("syncer profile policy sweep schedule is invalid") def validate_sao(self) -> None: diff --git a/yeto/wandb_tape.py b/yeto/wandb_tape.py index 125afa62e..59c4e584a 100644 --- a/yeto/wandb_tape.py +++ b/yeto/wandb_tape.py @@ -110,6 +110,8 @@ def tape_metrics(rec: dict) -> dict: ): if source in rec: metrics[metric] = int(rec[source]) + if "checkpoint_ms" in rec: # per-round checkpoint wall time (S26 E3) + metrics["sync/checkpoint_ms"] = float(rec["checkpoint_ms"]) if "sweep_complete" in rec: metrics["sync/sweep_complete"] = int(bool(rec["sweep_complete"])) if expected: From eb9e22905aea7ba61f659c94f838dcd991df18ab Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:56:04 +0000 Subject: [PATCH 2/9] syncer: bf16 sweep pushes and sweep-aligned checkpoint intervals Add session dtype 4 (f32 session, bf16 PUSH deltas): INIT/BCAST/FINAL, masters and momentum stay f32; the client rounds f32 delta bytes to bf16 (round-to-nearest-even, streamed across part boundaries). The dtype is bound into the role-stream session contract. Launcher --rl-sync-push-dtype, island --sweep-push-dtype (default f32). Allow --checkpoint-every N (1 or a multiple of the fragment count) for policy sweeps: only checkpointed steps are durable before their broadcast, on both the synchronous and async paths; the interval is part of the semantic profile hash. Launcher --rl-sync-checkpoint-every (default 1). E3 launcher options are refused outside the full-parameter preset. Adds bf16 error-bound, interval/resume and real-syncer tests, and a third micro-benchmark profile. Co-Authored-By: Claude Opus 5.5 --- docs/PROTOCOL.md | 9 +- syncer/src/protocol.rs | 19 ++- syncer/src/server.rs | 150 ++++++++++++++++-- tests/test_e3_sync_options.py | 63 ++++++++ tests/test_rl_dense_sweep_real_syncer.py | 48 +++++- tools/bench_sync_boundary.py | 40 +++-- yeto/launch/rl_args.py | 20 +++ yeto/launch/rl_island_task.py | 2 + yeto/launch/syncer.py | 9 +- yeto/protocol.py | 53 ++++++- .../miles/full_parameter/role_stream.py | 7 +- yeto/rl/adapters/miles/island_entry.py | 1 + yeto/rl/adapters/miles/run_flags.py | 6 + yeto/rl/engine/run_flags.py | 16 ++ yeto/syncer_profile.py | 4 + 15 files changed, 404 insertions(+), 43 deletions(-) create mode 100644 tests/test_e3_sync_options.py diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index 579db0930..3d08e0983 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -32,7 +32,7 @@ or allocation overflows are errors rather than panics. | type | name | direction | payload | |---:|---|---|---| -| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | +| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4, 4=f32 with bf16 pushes), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | | 2 | INIT_PARAMS | learner → syncer | fragment_id:u32, full tensor bytes; only learner 0 may initialize | | 3 | PULL_REQ | syncer → learner | fragment_id:u32, global_step:u64, round_attempt:u32 | | 4 | PUSH_FRAGMENT | learner → syncer | learner_id:u32, fragment_id:u32, global_step:u64, round_attempt:u32, base_version:u64, local_step:u64, c_steps:u32, c_tokens:u64, base-relative learner-delta bytes | @@ -223,6 +223,13 @@ The coordinator is one sequential scheduler actor. Pipelined rounds target distinct fragments, and merge + version update + broadcast for a completion is serialized. No two completions mutate one fragment concurrently. +## bf16 pushes over an f32 session (dtype 4) + +Dtype 4 keeps INIT, BCAST and FINAL exact f32 and sends only PUSH deltas as +bf16 (round-to-nearest-even). The syncer decodes and negates them like dtype 2 +and keeps f32 masters and momentum. The full-parameter policy sweep accepts +dtype 1 or 4 (S26 E3, `--rl-sync-push-dtype bf16`). + ## Q4 pushes Q4 uses the same learner-delta semantics. INIT and BCAST stay bf16; PUSH is diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index bf4b530ec..2103ce3bc 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -57,6 +57,10 @@ pub const DTYPE_BF16: u8 = 2; /// base-relative learner deltas; INIT_PARAMS and BCAST_FRAGMENT travel as /// bf16 (see `bulk_dtype`, docs/PROTOCOL.md v4). pub const DTYPE_Q4: u8 = 3; +/// Session dtype 4 (S26 E3): PUSH_FRAGMENT deltas travel as bf16 while +/// INIT_PARAMS, BCAST_FRAGMENT and FINAL stay exact f32. The syncer keeps f32 +/// masters and momentum; only the learner's outer gradient is rounded. +pub const DTYPE_F32_PUSH_BF16: u8 = 4; /// Values per q4 scale block (f32 absmax scale + 128 nibble bytes each). pub const Q4_BLOCK: usize = 256; @@ -64,7 +68,17 @@ pub const Q4_BLOCK: usize = 256; /// Dtype of INIT_PARAMS/BCAST_FRAGMENT tensors for a session dtype: q4 /// applies only to push deltas; full parameter payloads stay bf16. pub fn bulk_dtype(dtype: u8) -> u8 { - if dtype == DTYPE_Q4 { + match dtype { + DTYPE_Q4 => DTYPE_BF16, + DTYPE_F32_PUSH_BF16 => DTYPE_F32, + other => other, + } +} + +/// Dtype of PUSH_FRAGMENT deltas for an unquantized session dtype (q4 has +/// its own block codec). +pub fn push_dtype(dtype: u8) -> u8 { + if dtype == DTYPE_F32_PUSH_BF16 { DTYPE_BF16 } else { dtype @@ -417,6 +431,9 @@ mod tests { assert_eq!(bulk_dtype(DTYPE_Q4), DTYPE_BF16); assert_eq!(bulk_dtype(DTYPE_BF16), DTYPE_BF16); assert_eq!(bulk_dtype(DTYPE_F32), DTYPE_F32); + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_BF16), DTYPE_F32); + assert_eq!(push_dtype(DTYPE_F32_PUSH_BF16), DTYPE_BF16); + assert_eq!(push_dtype(DTYPE_F32), DTYPE_F32); } #[test] diff --git a/syncer/src/server.rs b/syncer/src/server.rs index e7b42b3e3..6d0b4ba18 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -735,10 +735,14 @@ fn validate_config(cfg: &Config) -> Result<()> { // --resume is not required here: a fresh sweep has no checkpoint yet and // --resume refuses a missing file (see check_resume_checkpoint). The // launcher adds --resume only when the checkpoint exists. - if cfg.checkpoint_path.is_none() || cfg.checkpoint_every != 1 { + if cfg.checkpoint_path.is_none() + || cfg.checkpoint_every == 0 + || (cfg.checkpoint_every != 1 + && !cfg.checkpoint_every.is_multiple_of(u64::from(fragments))) + { bail!( "--policy-sweep-fragments requires --checkpoint-path and \ - --checkpoint-every 1 for crash-safe sweeps" + --checkpoint-every 1 or a multiple of the fragment count" ); } } @@ -760,6 +764,20 @@ fn checkpoint_before_broadcast(cfg: &Config) -> bool { cfg.checkpoint_path.is_some() && cfg.checkpoint_every == 1 } +/// Whether step `t` must be durable before its BCAST. Per-round checkpoints +/// always are; a policy sweep with `--checkpoint-every N` (N a multiple of +/// the fragment count) persists only every N-th step, at sweep boundaries, +/// and broadcasts the other steps without a durable record (a crash resumes +/// at the newest checkpoint; islands ahead of it fail closed on their next +/// push's local-step check). +fn persist_before_broadcast(cfg: &Config, t: u64) -> bool { + checkpoint_before_broadcast(cfg) + || (cfg.policy_sweep_fragments.is_some() + && cfg.checkpoint_path.is_some() + && cfg.checkpoint_every > 0 + && t.is_multiple_of(cfg.checkpoint_every)) +} + fn validate_policy_sweep_push(push: &Push, fragments: u32) -> Result<()> { if fragments == 0 { bail!("policy sweep fragment count must be positive"); @@ -1000,7 +1018,7 @@ fn negotiated_payload_limits(layout: &Layout, dtype: u8) -> Result<(u64, u64)> { let delta_bytes = if dtype == DTYPE_Q4 { q4_nbytes(numel)? } else { - tensor_nbytes(dtype, numel)? + tensor_nbytes(push_dtype(dtype), numel)? } as u64; max_init = max_init.max( 4u64.checked_add(init_bytes) @@ -1146,7 +1164,7 @@ fn parse_hello(payload: &[u8], expected_learners: u32) -> Result { bail!("connection generation must be nonzero"); } let dtype = r.u8()?; - if !matches!(dtype, DTYPE_F32 | DTYPE_BF16 | DTYPE_Q4) { + if !matches!(dtype, DTYPE_F32 | DTYPE_BF16 | DTYPE_Q4 | DTYPE_F32_PUSH_BF16) { bail!("unsupported session dtype {dtype}"); } let num_fragments = r.u32()?; @@ -1720,7 +1738,7 @@ async fn dispatch_inner( let expected = if group.dtype == DTYPE_Q4 { q4_nbytes(numel)? } else { - tensor_nbytes(group.dtype, numel)? + tensor_nbytes(push_dtype(group.dtype), numel)? }; if bytes.len() != expected { bail!( @@ -1732,7 +1750,7 @@ async fn dispatch_inner( if group.dtype == DTYPE_Q4 { decode_q4(bytes, numel, &mut outer_gradient)?; } else { - decode_tensor(group.dtype, bytes, &mut outer_gradient)?; + decode_tensor(push_dtype(group.dtype), bytes, &mut outer_gradient)?; } // Every wire dtype carries learner_delta = local - raw_anchor. // Merge math consumes the corresponding outer gradient with the @@ -2102,7 +2120,7 @@ async fn scheduler( // Persist a fresh version zero before its first BCAST. A resumed cut is // already committed and only needs to be rebroadcast. - if checkpoint_each_round { + if checkpoint_each_round || cfg.policy_sweep_fragments.is_some() { if let Some(path) = cfg.checkpoint_path.as_ref() { if !cfg.resume || !path.exists() { st.save_checkpoint(path)?; @@ -2175,7 +2193,23 @@ async fn scheduler( .await?; } let committed = apply_round_commit(&cfg, &mut st, completed)?; - publishing = Some(start_async_checkpoint(&cfg, &st, committed)?); + if persist_before_broadcast(&cfg, committed.t) { + publishing = Some(start_async_checkpoint(&cfg, &st, committed)?); + } else { + publish_round( + &cfg, + &st, + ®istry, + &mut last_sync_secs, + committed, + None, + &budget_cutoff, + ) + .await?; + if !busy_fragments.remove(&p) { + bail!("fragment {p} lost its busy ownership at commit"); + } + } } else { commit_round( &cfg, @@ -3325,7 +3359,7 @@ async fn commit_round( ) -> Result<()> { let committed = apply_round_commit(cfg, st, completed)?; let mut checkpoint_ms = None; - if checkpoint_before_broadcast(cfg) { + if persist_before_broadcast(cfg, committed.t) { if let Some(path) = cfg.checkpoint_path.as_ref() { let checkpoint_start = Instant::now(); st.save_checkpoint(path)?; @@ -3510,7 +3544,8 @@ async fn publish_round( // have also been broadcast; the terminal cut is delivered by FINAL. // A crash-resume loses those gathers; their fragments simply merge on // a later cycle, which the quorum design already tolerates. - if !checkpoint_before_broadcast(cfg) { + // 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 && t % cfg.checkpoint_every == 0 { st.save_checkpoint(path)?; @@ -3652,14 +3687,18 @@ async fn complete_round_async( let cutoff = BudgetCutoff::new(false, Duration::ZERO); let mut busy = HashSet::from([completed.round.p]); let committed = apply_round_commit(cfg, st, completed)?; + if !persist_before_broadcast(cfg, committed.t) { + return publish_round(cfg, st, registry, last_sync_secs, committed, None, &cutoff).await; + } let pending = start_async_checkpoint(cfg, st, committed)?; settle_pending_publish(cfg, st, registry, last_sync_secs, pending, &cutoff, &mut busy).await } fn new_state_for(group: &Arc, cfg: &Config) -> Result { if let Some(fragments) = cfg.policy_sweep_fragments { - if group.dtype != DTYPE_F32 { - bail!("--policy-sweep-fragments requires an FP32 HELLO session"); + // Masters and every BCAST/FINAL stay f32; only push deltas may be bf16. + if !matches!(group.dtype, DTYPE_F32 | DTYPE_F32_PUSH_BF16) { + bail!("--policy-sweep-fragments requires an FP32 HELLO session (optionally with bf16 pushes)"); } if group.layout.fragments.len() != fragments as usize { bail!( @@ -4978,13 +5017,21 @@ mod tests { for mutate in [ |config: &mut Config| config.checkpoint_path = None, - |config: &mut Config| config.checkpoint_every = 2, + |config: &mut Config| config.checkpoint_every = 0, + |config: &mut Config| config.checkpoint_every = 3, ] { invalid = sweep_test_config(2, 4); mutate(&mut invalid); assert!(format!("{:#}", validate_config(&invalid).unwrap_err()) .contains("requires --checkpoint-path")); } + // Sweep-aligned checkpoint intervals are accepted. + let mut interval = sweep_test_config(2, 4); + interval.checkpoint_every = 4; + assert!(validate_config(&interval).is_ok()); + assert!(persist_before_broadcast(&interval, 4)); + assert!(!persist_before_broadcast(&interval, 2)); + assert!(persist_before_broadcast(&sweep_test_config(2, 4), 1)); } #[test] @@ -5581,6 +5628,83 @@ mod tests { directory } + #[test] + fn bf16_push_session_rounds_only_the_delta_within_bf16_error() { + // dtype 4 decodes pushes as bf16 and keeps every bulk payload f32. + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_BF16), DTYPE_F32); + let delta: Vec = (0..4096) + .map(|i| ((i as f32) * 0.37).sin() * 10f32.powi((i % 9) - 4)) + .collect(); + let mut wire = Vec::new(); + encode_tensor(push_dtype(DTYPE_F32_PUSH_BF16), &delta, &mut wire).unwrap(); + assert_eq!(wire.len(), delta.len() * 2); + let mut decoded = Vec::new(); + decode_tensor(push_dtype(DTYPE_F32_PUSH_BF16), &wire, &mut decoded).unwrap(); + for (exact, rounded) in delta.iter().zip(&decoded) { + // bf16 keeps 8 significant bits: round-to-nearest error <= 2^-8 relative. + assert!((exact - rounded).abs() <= exact.abs() * (1.0 / 256.0)); + } + // A merge of bf16 deltas into f32 masters stays within the same bound + // per learner, so a two-learner average is off by at most that much. + let master = 1.0f32; + for (exact, rounded) in delta.iter().zip(&decoded) { + let merged_exact = master + exact; + let merged = master + rounded; + assert!((merged_exact - merged).abs() <= exact.abs() / 256.0 + f32::EPSILON * 2.0); + } + } + + #[tokio::test] + async fn sweep_checkpoint_interval_persists_only_sweep_aligned_steps() { + let directory = unique_test_directory("sweep-checkpoint-interval"); + let learner = member(0, 30); + for asynchronous in [false, true] { + let checkpoint = directory.join(format!("state-{asynchronous}.ckpt")); + let (group, mut frames) = streaming_test_group(DTYPE_F32, false); + let registry = registry_with_group(group); + let mut config = sweep_test_config(2, 8); + config.checkpoint_path = Some(checkpoint.clone()); + config.checkpoint_every = 4; + config.async_checkpoint = asynchronous; + let mut state = two_fragment_sweep_state(); + state.save_checkpoint(&checkpoint).unwrap(); // version zero + let mut sync_seconds = 0.0; + for step in 1..=6u64 { + let fragment = ((step - 1) % 2) as usize; + let mut round = sweep_test_round(step, fragment, step.div_ceil(2), 1, learner); + let base = state.versions[fragment]; + round.base_version = base; + for push in round.pushes.values_mut() { + push.base_version = base; + } + if asynchronous { + complete_round_async(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } else { + complete_round(&config, &mut state, ®istry, &mut sync_seconds, round) + .await + .unwrap(); + } + assert_eq!( + chunked_inner_type(&frames.try_recv().unwrap()), + MSG_BCAST_FRAGMENT + ); + let mut on_disk = two_fragment_sweep_state(); + on_disk.load_checkpoint(&checkpoint).unwrap(); + // Only step 4 (a sweep boundary and a multiple of 4) is durable. + let expected = if step >= 4 { 4 } else { 0 }; + assert_eq!(on_disk.global_step, expected, "step {step}"); + } + // The newest checkpoint resumes exactly at its sweep boundary. + let mut resumed = two_fragment_sweep_state(); + resumed.load_checkpoint(&checkpoint).unwrap(); + validate_resumed_policy_sweep(&config, &resumed).unwrap(); + assert_eq!(resumed.versions, vec![3, 4]); + } + std::fs::remove_dir_all(&directory).ok(); + } + #[tokio::test] async fn async_checkpoint_matches_the_synchronous_path_bit_for_bit() { let directory = unique_test_directory("async-checkpoint-equal"); diff --git a/tests/test_e3_sync_options.py b/tests/test_e3_sync_options.py new file mode 100644 index 000000000..0e8492ef9 --- /dev/null +++ b/tests/test_e3_sync_options.py @@ -0,0 +1,63 @@ +"""S26 E3 sync options: bf16 push transcoding, session binding, preset scoping. CPU only.""" + +from __future__ import annotations + +from dataclasses import replace + +import pytest +import torch + +from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_full_parameter_sync import _role_stream_config +from yeto.protocol import ( + DTYPE_BF16, + DTYPE_F32, + DTYPE_F32_PUSH_BF16, + _f32_parts_to_bf16, + bulk_dtype, + f32_bytes_to_bf16, + push_wire_dtype, +) + + +def test_push_bf16_session_keeps_bulk_payloads_f32(): + assert bulk_dtype(DTYPE_F32_PUSH_BF16) == DTYPE_F32 + assert push_wire_dtype(DTYPE_F32_PUSH_BF16) == DTYPE_BF16 + assert push_wire_dtype(DTYPE_F32) == DTYPE_F32 + + +def test_f32_to_bf16_matches_torch_round_to_nearest_even_across_split_parts(): + values = torch.randn(50_001) * torch.logspace(-20, 20, 50_001) + values[:3] = torch.tensor([0.0, -0.0, float("inf")]) + raw = values.numpy().tobytes() + want = values.to(torch.bfloat16).view(torch.int16).numpy().astype(" FragmentLayout: return FragmentLayout([Fragment(MERGE_RDA, [(f"actor.w{i}", NUMEL)]) for i in range(FRAGMENTS)]) -def _client(port: int, profile_hash: str) -> SyncerClient: - return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=1, +def _client(port: int, profile_hash: str, dtype: int = DTYPE_F32) -> SyncerClient: + return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=dtype, num_streams=1, connect_timeout=10, max_reconnects=0, finalization_timeout=20, syncer_profile_hash=bytes.fromhex(profile_hash)) @@ -160,10 +160,11 @@ def test_launcher_dense_syncer_refuses_another_profile_hash(syncer_binary, tmp_p STREAMED = ("--rl-sync-fragment-pipeline", str(FRAGMENTS), "--rl-sync-async-checkpoint") -def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int): +def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int, + dtype: int = DTYPE_F32): """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint).""" process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) - client = _client(port, profile_hash) + client = _client(port, profile_hash, dtype) concurrent = 0 try: client.start() @@ -232,3 +233,40 @@ def test_streamed_flags_bind_the_profile_and_are_full_parameter_only(): != launcher.rl_syncer_profile(streamed, 1).sha256) with pytest.raises(ValueError, match="--rl-sync-fragment-pipeline between 1 and --fragments"): _dense_args("--rl-sync-fragment-pipeline", str(FRAGMENTS + 1)) + + +def test_bf16_pushes_stay_within_bf16_error_of_the_f32_sweep(syncer_binary, tmp_path): + """E3: bf16 outer deltas change results only by bf16 rounding of each delta.""" + exact_dir, bf16_dir = tmp_path / "f32", tmp_path / "bf16" + exact_dir.mkdir() + bf16_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=3) + rounded = _run_sweeps(syncer_binary, bf16_dir, (*STREAMED, "--rl-sync-push-dtype", "bf16"), + sweeps=3, dtype=DTYPE_F32_PUSH_BF16) + assert exact[1] == rounded[1] == [5, 6] # versions are unchanged + for fid, data in exact[2].items(): + want = torch.frombuffer(bytearray(data), dtype=torch.float32)[-NUMEL:] + got = torch.frombuffer(bytearray(rounded[2][fid]), dtype=torch.float32)[-NUMEL:] + # Three sweeps, each delta < 2 in magnitude, bf16 error <= 2^-8 relative. + assert torch.allclose(got, want, rtol=0, atol=3 * 2.0 / 256), fid + assert not torch.equal(got, want), "bf16 pushes must actually round" + + +def test_sweep_checkpoint_interval_reaches_the_syncer_and_its_profile(): + every = 2 * FRAGMENTS + args = _dense_args("--rl-sync-checkpoint-every", str(every), total_steps=4) + from yeto import launcher + + command = launcher.role_syncer_commands(args, 1, binary="yeto-syncer")["actor"] + assert f"--checkpoint-every {every}" in command + assert launcher.rl_syncer_profile(args, 1).checkpoint_every == every + with pytest.raises(ValueError, match="--rl-sync-checkpoint-every 1 or a multiple"): + _dense_args("--rl-sync-checkpoint-every", str(FRAGMENTS + 1)) + + +def test_bf16_push_flag_reaches_the_island_command(): + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + args = _dense_args("--rl-sync-push-dtype", "bf16") + assert "--sweep-push-dtype bf16" in _rl_trainable_island_flags(args, 0, 1) + assert "--sweep-push-dtype" not in _rl_trainable_island_flags(_dense_args(), 0, 1) diff --git a/tools/bench_sync_boundary.py b/tools/bench_sync_boundary.py index 3b4496322..a55fb9066 100644 --- a/tools/bench_sync_boundary.py +++ b/tools/bench_sync_boundary.py @@ -3,8 +3,9 @@ Runs the real Rust syncer on localhost with two in-process learners and a synthetic f32 model, and times each strict sweep from the first push to the last updated fragment received by every learner. It compares the serial -profile (pipeline 1, checkpoint on the scheduler) with the streamed profile -(pipeline P, ``--async-checkpoint``). +profile (pipeline 1, checkpoint on the scheduler), the streamed profile +(pipeline P, ``--async-checkpoint``), and streamed plus bf16 pushes and a +checkpoint every second sweep. These are local CPU numbers: loopback transfer and a local NVMe checkpoint. They show where the boundary time goes on this machine; they do not predict @@ -28,7 +29,7 @@ import numpy as np from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout -from yeto.protocol import DTYPE_F32, SyncerClient +from yeto.protocol import DTYPE_F32, DTYPE_F32_PUSH_BF16, SyncerClient ROOT = Path(__file__).resolve().parent.parent @@ -40,7 +41,7 @@ def _free_port() -> int: def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: int, - sweeps: int, streamed: bool) -> subprocess.Popen: + sweeps: int, streamed: bool, checkpoint_every: int) -> subprocess.Popen: command = [ str(binary), "--port", str(port), "--learners", str(learners), "--quorum", str(learners), "--grace-ms", "0", @@ -49,7 +50,7 @@ def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: "--total-steps", str(sweeps * fragments), "--policy-sweep-fragments", str(fragments), "--outer-lr", "1", "--outer-momentum", "0", "--max-base-lag", "0", "--learner-weight", "equal", "--checkpoint-path", str(workdir / "state.ckpt"), - "--checkpoint-every", "1", "--event-tape", str(workdir / "tape.jsonl"), + "--checkpoint-every", str(checkpoint_every), "--event-tape", str(workdir / "tape.jsonl"), "--allow-unauthenticated-islands", # loopback-only benchmark ] if streamed: @@ -59,9 +60,9 @@ def _syncer(binary: Path, workdir: Path, *, port: int, learners: int, fragments: def _learner(learner_id: int, port: int, layout: FragmentLayout, numel: int, sweeps: int, - timings: list[list[float]], errors: list[BaseException]) -> None: + timings: list[list[float]], errors: list[BaseException], dtype: int) -> None: fragments = len(layout.fragments) - client = SyncerClient(("127.0.0.1", port), learner_id, layout, dtype=DTYPE_F32, + client = SyncerClient(("127.0.0.1", port), learner_id, layout, dtype=dtype, num_streams=1, connect_timeout=30, max_reconnects=0, finalization_timeout=600) rng = np.random.default_rng(learner_id) @@ -107,18 +108,29 @@ def _learner(learner_id: int, port: int, layout: FragmentLayout, numel: int, swe client.close() -def run(binary: Path, *, streamed: bool, learners: int, fragments: int, numel: int, +PROFILES = { + "serial": {"streamed": False, "dtype": DTYPE_F32, "sweeps_per_checkpoint": 0}, + "streamed": {"streamed": True, "dtype": DTYPE_F32, "sweeps_per_checkpoint": 0}, + "streamed+bf16+ckpt/2sweeps": {"streamed": True, "dtype": DTYPE_F32_PUSH_BF16, + "sweeps_per_checkpoint": 2}, +} + + +def run(binary: Path, *, profile: str, learners: int, fragments: int, numel: int, sweeps: int) -> dict: + options = PROFILES[profile] + streamed, dtype = options["streamed"], options["dtype"] + every = options["sweeps_per_checkpoint"] * fragments or 1 layout = FragmentLayout([Fragment(MERGE_RDA, [(f"w{i}", numel)]) for i in range(fragments)]) with tempfile.TemporaryDirectory(prefix="e3-bench-", dir="/tmp") as tmp: workdir = Path(tmp) port = _free_port() process = _syncer(binary, workdir, port=port, learners=learners, fragments=fragments, - sweeps=sweeps, streamed=streamed) + sweeps=sweeps, streamed=streamed, checkpoint_every=every) timings: list[list[float]] = [[] for _ in range(learners)] errors: list[BaseException] = [] threads = [threading.Thread(target=_learner, args=(i, port, layout, numel, sweeps, - timings, errors)) + timings, errors, dtype)) for i in range(learners)] for thread in threads: thread.start() @@ -134,10 +146,10 @@ def run(binary: Path, *, streamed: bool, learners: int, fragments: int, numel: i steady = [max(per[i] for per in timings) for i in range(1, sweeps - 1)] or \ [max(per[i] for per in timings) for i in range(sweeps)] return { - "profile": "streamed" if streamed else "serial", + "profile": profile, "boundary_s_per_sweep": [round(x, 3) for x in steady], "boundary_s_median": round(statistics.median(steady), 3), - "checkpoint_ms_median": statistics.median(r["checkpoint_ms"] for r in tape), + "checkpoints_written": sum("checkpoint_ms" in r for r in tape), "sync_ms_median": statistics.median(r["sync_ms"] for r in tape), } @@ -150,9 +162,9 @@ def main() -> None: parser.add_argument("--floats-per-fragment", type=int, default=32 * 1024 * 1024) parser.add_argument("--sweeps", type=int, default=4) args = parser.parse_args() - results = [run(args.binary, streamed=streamed, learners=args.learners, + results = [run(args.binary, profile=profile, learners=args.learners, fragments=args.fragments, numel=args.floats_per_fragment, sweeps=args.sweeps) - for streamed in (False, True)] + for profile in PROFILES] print(json.dumps({"note": "CPU localhost numbers; not GPU or WAN timings", "learners": args.learners, "fragments": args.fragments, "floats_per_fragment": args.floats_per_fragment, "results": results}, diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index e1ce4ff28..df0c0efbb 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -148,6 +148,21 @@ def sao_contract_mounts(args, learner_id: int) -> dict[str, str]: for shipped in (island.runtime, island.layout_attestation)} +_FULL_PARAMETER_SYNC_DEFAULTS = ( + ("rl_sync_fragment_pipeline", 1, "--rl-sync-fragment-pipeline"), + ("rl_sync_async_checkpoint", False, "--rl-sync-async-checkpoint"), + ("rl_sync_push_dtype", "f32", "--rl-sync-push-dtype"), + ("rl_sync_checkpoint_every", 1, "--rl-sync-checkpoint-every"), +) + + +def _refuse_full_parameter_sync_options(args) -> None: + """The S26 E3 sync options exist only for the full-parameter preset.""" + for dest, default, option in _FULL_PARAMETER_SYNC_DEFAULTS: + if getattr(args, dest, default) != default: + raise ValueError(f"{option} only applies to --rl-sync-preset {FULL_PARAMETER_SYNC_PRESET}") + + def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: """full-parameter preset (P0-D8): one actor role swept as whole-model fragments. @@ -172,6 +187,9 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: fragment_pipeline = int(getattr(args, "rl_sync_fragment_pipeline", 1)) if not 1 <= fragment_pipeline <= max(args.fragments, 1): problems.append("--rl-sync-fragment-pipeline between 1 and --fragments") + checkpoint_every = int(getattr(args, "rl_sync_checkpoint_every", 1)) + if checkpoint_every != 1 and (checkpoint_every < 1 or checkpoint_every % max(args.fragments, 1)): + problems.append("--rl-sync-checkpoint-every 1 or a multiple of --fragments") if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands @@ -728,6 +746,8 @@ def _prepare_rl_args( f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET}") if args.rl_sync_preset != SAO_STREAMING_SYNC_PRESET and getattr(args, "rl_sao_contracts", None): raise ValueError("--rl-sao-contracts only applies to --rl-sync-preset sao-streaming-full") + if args.rl_sync_preset != FULL_PARAMETER_SYNC_PRESET: + _refuse_full_parameter_sync_options(args) if args.rl_sync_preset == "decoupled": if args.experimental_rl_sync: raise ValueError("decoupled RL does not accept --experimental-rl-sync") diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 1185f4c9a..a5cd883f0 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -72,6 +72,8 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: sweep = (f" --parameter-layout-sha256 {args.rl_parameter_layout_sha256}" f" --num-learners {num_learners}" f" --syncer-profile-sha256 {rl_syncer_profile(args, num_learners).sha256}") + if getattr(args, "rl_sync_push_dtype", "f32") != "f32": + sweep += f" --sweep-push-dtype {args.rl_sync_push_dtype}" if getattr(args, "expert_full_count", 0): return lora + sweep return " --parameter-mode full" + sweep diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 7d951ae9b..91587c078 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -145,6 +145,11 @@ def _syncer_quorum_timeout(args) -> str: return "" if value is None else f" --quorum-timeout-s {value}" +def _syncer_checkpoint_every(args) -> int: + """Per-round checkpoints (1) unless the full-parameter preset sets an interval.""" + return int(getattr(args, "rl_sync_checkpoint_every", 1)) if rl_dense_sweep(args) else 1 + + def rl_dense_sweep(args) -> bool: return (getattr(args, "training_mode", "sft") == "rl" and getattr(args, "rl_sync_preset", None) == FULL_PARAMETER_SYNC_PRESET) @@ -206,7 +211,7 @@ def rl_syncer_profile(args, num_learners: int): "outer_lr": args.outer_lr, "outer_momentum": args.outer_momentum, "checkpoint_enabled": True, - "checkpoint_every": 1, + "checkpoint_every": _syncer_checkpoint_every(args), "resume": False, # not hashed (v2): fresh start and restart agree "mark_final_checkpoint": mark_final, "learner_budget_steps": None, @@ -447,7 +452,7 @@ def _rl_role_syncer_command(args, num_learners: int, binary: str, f" --outer-momentum {args.outer_momentum}" " --max-base-lag 0 --learner-weight equal" f" --checkpoint-path {syncer.checkpoint}" - " --checkpoint-every 1" + f" --checkpoint-every {_syncer_checkpoint_every(args)}" f" {_resume_if_exists(syncer.checkpoint)}" f" --event-tape {syncer.event_tape}" ) diff --git a/yeto/protocol.py b/yeto/protocol.py index 7beb2254f..28331475c 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -34,7 +34,7 @@ import struct import threading import time -from collections.abc import Callable, Iterable +from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field from .fragments import MERGE_ISO, FragmentLayout @@ -66,6 +66,11 @@ # base-relative learner deltas; INIT_PARAMS and BCAST_FRAGMENT stay bf16 # (see bulk_dtype and docs/PROTOCOL.md v4). DTYPE_Q4 = 3 +# Session dtype 4 (S26 E3): PUSH_FRAGMENT deltas travel as bf16; INIT_PARAMS, +# BCAST_FRAGMENT and FINAL stay exact f32. Callers still hand the client f32 +# delta bytes; the client rounds them (round-to-nearest-even, as the syncer's +# ``half::bf16::from_f32``) on the way out. +DTYPE_F32_PUSH_BF16 = 4 CHUNK_SIZE = 4 * 1024 * 1024 MAX_HANDSHAKE_FRAME = 16 * 1024 * 1024 @@ -93,7 +98,40 @@ def bulk_dtype(dtype: int) -> int: Q4 applies only to push deltas (small dynamic range); full parameter payloads would not survive 4 bits, so they travel as bf16. """ - return DTYPE_BF16 if dtype == DTYPE_Q4 else dtype + if dtype == DTYPE_Q4: + return DTYPE_BF16 + if dtype == DTYPE_F32_PUSH_BF16: + return DTYPE_F32 + return dtype + + +def push_wire_dtype(dtype: int) -> int: + """Dtype of unquantized PUSH_FRAGMENT deltas on the wire for a session dtype.""" + return DTYPE_BF16 if dtype == DTYPE_F32_PUSH_BF16 else dtype + + +def f32_bytes_to_bf16(data: bytes | bytearray | memoryview) -> bytes: + """Round little-endian f32 values to bf16 (nearest, ties to even).""" + import numpy as np # noqa: PLC0415 - keep the protocol import light + + bits = np.frombuffer(data, dtype="> 16) & 1)) >> 16 + nan = (bits & 0x7FFF_FFFF) > 0x7F80_0000 + rounded[nan] = (bits[nan] >> 16) | 0x40 # keep NaN quiet, never round to inf + return rounded.astype(" Iterator[bytes]: + """Stream ``f32_bytes_to_bf16`` over parts whose boundaries may split a value.""" + carry = b"" + for part in parts: + data = carry + bytes(part) + whole = len(data) - len(data) % 4 + carry = data[whole:] + if whole: + yield f32_bytes_to_bf16(data[:whole]) + if carry: + raise ValueError("f32 push delta ends mid-value") def layout_fingerprint(layout: FragmentLayout) -> bytes: @@ -480,7 +518,7 @@ def __init__( ): if not 0 <= learner_id <= 0xFFFF_FFFF: raise ValueError(f"learner_id must fit u32, got {learner_id}") - if dtype not in (DTYPE_F32, DTYPE_BF16, DTYPE_Q4): + if dtype not in (DTYPE_F32, DTYPE_BF16, DTYPE_Q4, DTYPE_F32_PUSH_BF16): raise ValueError(f"unsupported session dtype {dtype}") if layout.num_fragments < 1: raise ValueError("layout must contain at least one fragment") @@ -833,7 +871,8 @@ def push_fragment( expected = ( _q4_nbytes(numel) if self.dtype == DTYPE_Q4 - else _tensor_nbytes(self.dtype, numel) + # dtype 4 callers pass f32 bytes; the client rounds them below + else _tensor_nbytes(DTYPE_F32 if self.dtype == DTYPE_F32_PUSH_BF16 else self.dtype, numel) ) if len(tensor_bytes) != expected: raise ValueError( @@ -851,6 +890,8 @@ def push_fragment( c_steps, c_tokens, ) + if self.dtype == DTYPE_F32_PUSH_BF16: + tensor_bytes = f32_bytes_to_bf16(tensor_bytes) self._send_large(MSG_PUSH_FRAGMENT, head + tensor_bytes) def push_fragment_parts( @@ -885,8 +926,10 @@ def push_fragment_parts( expected = ( _q4_nbytes(numel) if self.dtype == DTYPE_Q4 - else _tensor_nbytes(self.dtype, numel) + else _tensor_nbytes(push_wire_dtype(self.dtype), numel) ) + if self.dtype == DTYPE_F32_PUSH_BF16: + tensor_parts = _f32_parts_to_bf16(tensor_parts) head = struct.pack( " None: if not self.algorithm.isidentifier(): @@ -219,7 +222,7 @@ def role_stream_session_contract_hash( "pipeline_depth": config.pipeline_depth, "local_horizon": config.local_horizon, "optimizer_steps_per_round": config.optimizer_steps_per_round, - "wire_dtype": "fp32", + "wire_dtype": "fp32-push-bf16" if config.push_bf16 else "fp32", }, } encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() @@ -334,7 +337,7 @@ async def initialize(self, group) -> None: self.config.syncer_addr, self.config.learner_id, adapter.layout.fragments, - dtype=DTYPE_F32, + dtype=DTYPE_F32_PUSH_BF16 if self.config.push_bf16 else DTYPE_F32, num_streams=self.config.wan_streams, max_reconnects=None, session_contract_hash=role_stream_session_contract_hash( diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 48585224f..5a9d8874d 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -838,6 +838,7 @@ def _dense_full_sync_config( pipeline_depth=1, algorithm_spec_sha256=algorithm_spec_sha256, wan_streams=args.wan_streams, + push_bf16=getattr(args, "sweep_push_dtype", "f32") == "bf16", ) return FullParameterSyncConfig((actor,), sweep=True) diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 8a9f52027..adbf93ca8 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -313,6 +313,12 @@ ), island_flag("--fragments", type=int, default=1), island_flag("--parameter-layout-sha256", default=None), + island_flag( + "--sweep-push-dtype", + choices=("f32", "bf16"), + default="f32", + help="full-parameter sweep: wire dtype of pushed outer deltas (syncer masters stay f32)", + ), island_flag( "--syncer-profile-sha256", default=None, diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index f9be06f6f..6e6b95c1a 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -849,6 +849,22 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. "(1..--fragments). Above 1, fragment k+1 uploads while fragment k merges " "and downloads; merge results are unchanged (default 1 = serial)", ), + cli_flag( + "--rl-sync-push-dtype", + choices=("f32", "bf16"), + default="f32", + help="full-parameter only: wire dtype of the islands' pushed outer deltas " + "(pseudo-gradients). bf16 halves upload bytes; the syncer keeps f32 masters, " + "momentum and f32 BCAST/FINAL (default f32 = exact)", + ), + cli_flag( + "--rl-sync-checkpoint-every", + type=int, + default=1, + help="full-parameter only: syncer checkpoint interval in fragment steps " + "(1, or a multiple of --fragments). Above 1, only checkpointed steps are " + "persisted before their broadcast; a crash resumes at the newest checkpoint", + ), cli_flag( "--rl-sync-async-checkpoint", action="store_true", diff --git a/yeto/syncer_profile.py b/yeto/syncer_profile.py index c5e011b08..ce23fa993 100644 --- a/yeto/syncer_profile.py +++ b/yeto/syncer_profile.py @@ -261,6 +261,10 @@ def _validate_general(self) -> None: fragments = self.policy_sweep_fragments if not 1 <= self.pipeline <= fragments or self.total_steps % fragments: raise ValueError("syncer profile policy sweep schedule is invalid") + if self.checkpoint_every != 1 and ( + self.checkpoint_every < 1 or self.checkpoint_every % fragments): + raise ValueError("syncer profile policy sweep checkpoint interval must be " + "1 or a multiple of the fragment count") def validate_sao(self) -> None: """Enforce the paired actor/critic fail-closed server profile.""" From 1a783d09529d711b01cba7478cce188312141239 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:48:24 +0000 Subject: [PATCH 3/9] syncer: int8 sweep pushes, bf16 downlink deltas, sync/train overlap Add session dtype 5 (f32 session, block-scaled int8 PUSH deltas): u32 block size, then per block an f32 absmax/127 scale and one i8 per value; error is at most absmax/254 per element. Launcher --rl-sync-push-int8-block (default 128) is bound into the role-stream contract. Add session flag 0x80 (bf16 downlink deltas) for f32-bulk sessions. Syncer and learners hold a bit-identical f32 anchor per fragment; BCAST is FULL f32 or bf16(theta - anchor), with a FULL refresh every N versions (--bcast-delta-full-every), FULL re-sends on HELLO/reconnect, and a SHA-256 check of the anchor after every update. A mismatch drops the connection and the reconnect restores the anchor. SyncerClient decodes before adapters see the payload: the Miles streaming sink is wrapped, the strict bridge (verl and Miles LoRA) gets f32 bytes. Masters, checkpoints and FINAL stay exact. Add --rl-sync-overlap for the full-parameter sweep: an island pushes round k, trains round k+1 at once, and at the next boundary rebases in place, theta <- G_k + (local - base), with G_k as the anchor so no delta is counted twice. The launcher raises --rl-max-policy-age to at least 1 before its checks; the island refuses overlap below that age. Adds Rust and Python codec/anchor tests, real-syncer int8/downlink/resync tests, a strict-bridge downlink test, and coordinator overlap tests. Co-Authored-By: Claude Opus 5.5 --- docs/PROTOCOL.md | 26 +- syncer/src/downlink.rs | 289 +++++++++++++++ syncer/src/main.rs | 7 + syncer/src/protocol.rs | 197 ++++++++++- syncer/src/server.rs | 113 ++++-- syncer/src/state.rs | 3 + tests/test_e3_sync_options.py | 261 +++++++++++++- tests/test_e3_sync_overlap.py | 331 ++++++++++++++++++ tests/test_rl_dense_sweep_real_syncer.py | 176 +++++++++- tests/test_run_flags.py | 2 +- yeto/downlink.py | 235 +++++++++++++ yeto/launch/rl_args.py | 36 +- yeto/launch/rl_island_task.py | 5 + yeto/launch/syncer.py | 7 + yeto/protocol.py | 170 +++++++-- .../miles/full_parameter/role_stream.py | 140 +++++++- yeto/rl/adapters/miles/island_entry.py | 12 +- .../miles/models/chunked_full_parameter.py | 54 +++ yeto/rl/adapters/miles/run_flags.py | 20 +- yeto/rl/adapters/verl/island_entry.py | 3 + yeto/rl/adapters/verl/trainer.py | 3 +- yeto/rl/bridge.py | 18 +- yeto/rl/engine/full_parameter_sync.py | 75 +++- yeto/rl/engine/run_flags.py | 32 +- yeto/rl/engine/sync_overlap.py | 41 +++ 25 files changed, 2144 insertions(+), 112 deletions(-) create mode 100644 syncer/src/downlink.rs create mode 100644 tests/test_e3_sync_overlap.py create mode 100644 yeto/downlink.py create mode 100644 yeto/rl/engine/sync_overlap.py diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index 3d08e0983..db8e47e95 100644 --- a/docs/PROTOCOL.md +++ b/docs/PROTOCOL.md @@ -32,7 +32,7 @@ or allocation overflows are errors rather than panics. | type | name | direction | payload | |---:|---|---|---| -| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4, 4=f32 with bf16 pushes), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | +| 1 | HELLO | learner → syncer | protocol_version:u16 (=4), learner_id:u32, connection_generation:u64, dtype:u8 (1=f32, 2=bf16, 3=q4, 4=f32 with bf16 pushes, 5=f32 with int8 pushes; flag 0x80 = bf16 downlink deltas), num_fragments:u32, per-fragment layout, layout_fingerprint:[u8;32], session_contract_hash:[u8;32], optional syncer_profile_hash:[u8;32], num_streams:u16, optional compat_group trailer (`b"YCG1"`, u32 length, 1–64 bytes utf-8) | | 2 | INIT_PARAMS | learner → syncer | fragment_id:u32, full tensor bytes; only learner 0 may initialize | | 3 | PULL_REQ | syncer → learner | fragment_id:u32, global_step:u64, round_attempt:u32 | | 4 | PUSH_FRAGMENT | learner → syncer | learner_id:u32, fragment_id:u32, global_step:u64, round_attempt:u32, base_version:u64, local_step:u64, c_steps:u32, c_tokens:u64, base-relative learner-delta bytes | @@ -230,6 +230,30 @@ bf16 (round-to-nearest-even). The syncer decodes and negates them like dtype 2 and keeps f32 masters and momentum. The full-parameter policy sweep accepts dtype 1 or 4 (S26 E3, `--rl-sync-push-dtype bf16`). +## int8 pushes over an f32 session (dtype 5) + +Dtype 5 sends PUSH deltas as block-scaled int8: u32 block size B +(32..1048576), then per block of B values (the last block may be shorter) an +f32 scale (block absmax / 127) and one i8 per value (`-128` is invalid). +A value decodes to `q * scale`, so its error is at most absmax / 254. The +learner picks B (launcher `--rl-sync-push-int8-block`, default 128); it is +bound into the role-stream session contract. INIT, BCAST and FINAL stay f32. + +## bf16 downlink deltas (dtype flag 0x80) + +On an f32-bulk session (dtype 1, 4 or 5) the flag 0x80 changes BCAST_FRAGMENT +to `fragment_id:u32, version:u64, kind:u8, base_version:u64, +anchor_sha256:[u8;32]`, then the tensor: kind 0 (FULL) carries f32 values, +kind 1 (DELTA) carries `bf16(theta - anchor)`. Syncer and learner hold the +same f32 anchor per fragment: FULL sets it, DELTA adds the decoded delta in +f32; the digest is checked after every update. Every N-th version of a +fragment is FULL (`--bcast-delta-full-every N`, N >= 2; the syncer refuses a +flagged session without it and the reverse). Initial and reconnect re-sends +are always FULL with the shared anchor. A learner with a missing or wrong +anchor drops its connection, and the reconnect re-send restores it. The +syncer's masters, momentum, checkpoints and FINAL are unchanged and exact; +anchors are not checkpointed, so a resumed syncer starts with FULL. + ## Q4 pushes Q4 uses the same learner-delta semantics. INIT and BCAST stay bf16; PUSH is diff --git a/syncer/src/downlink.rs b/syncer/src/downlink.rs new file mode 100644 index 000000000..e303e119c --- /dev/null +++ b/syncer/src/downlink.rs @@ -0,0 +1,289 @@ +//! bf16 downlink deltas for BCAST_FRAGMENT (S26 E3, session flag +//! `DTYPE_FLAG_BCAST_DELTA`). +//! +//! Every learner of a session holds, per fragment, an f32 *anchor* that is +//! bit-identical to the syncer's copy here. A broadcast of a new version is +//! either +//! +//! * `FULL`: the f32 values; the anchor becomes exactly those values, or +//! * `DELTA`: `d = bf16(θ − anchor)`; both sides set `anchor += d` in f32 +//! (IEEE addition is exact-rounded, so both sides get the same bits). +//! +//! Each delta is taken against the current anchor, so the rounding error does +//! not accumulate: `|θ − anchor| <= |θ − anchor_prev| * 2^-9` per element +//! after every delta. The learner trains from the anchor; its next push is +//! relative to the anchor, and the syncer adds that push to its exact θ, so +//! θ stays the exact sum of all pushes. +//! +//! Drift guard: every `full_every`-th broadcast of a fragment is FULL; a +//! re-sent fragment (HELLO/reconnect recovery) is always FULL; and every +//! header carries the SHA-256 of the f32 anchor after the update. A learner +//! whose anchor version or digest does not match drops its connection; the +//! reconnect recovery then sends FULL. +//! +//! Anchors live only in memory: a resumed syncer starts with none, so its +//! first broadcast of every fragment is FULL. +//! +//! Header after the usual `fragment_id u32, version u64` prefix: +//! `kind u8 (0 FULL, 1 DELTA), base_version u64, anchor_sha256 [32]`. +//! The tensor that follows is f32 (FULL) or bf16 (DELTA). + +use std::sync::Arc; + +use anyhow::{bail, Result}; +use sha2::{Digest, Sha256}; + +use crate::protocol::{DTYPE_BF16, DTYPE_F32}; + +pub const KIND_FULL: u8 = 0; +pub const KIND_DELTA: u8 = 1; +/// fragment_id + version + kind + base_version + digest. +pub const PREFIX_LEN: usize = 4 + 8 + 1 + 8 + 32; + +struct Anchor { + version: u64, + values: Arc>, + digest: [u8; 32], + deltas_since_full: u64, +} + +/// One encoded BCAST: header prefix plus tensor values in `dtype` (DELTA +/// values are already bf16-representable, so encoding them is exact). +pub struct EncodedBcast { + pub prefix: Vec, + pub dtype: u8, + pub values: Arc>, +} + +impl EncodedBcast { + #[cfg(test)] + pub fn kind(&self) -> u8 { + self.prefix[12] + } +} + +pub struct DownlinkAnchors { + full_every: u64, + anchors: Vec>, +} + +pub fn digest_f32(values: &[f32]) -> [u8; 32] { + let mut hasher = Sha256::new(); + // Hash in bounded slices to avoid a second full-size byte buffer. + let mut buf = Vec::with_capacity(4 * 65536); + for chunk in values.chunks(65536) { + buf.clear(); + for v in chunk { + buf.extend_from_slice(&v.to_le_bytes()); + } + hasher.update(&buf); + } + hasher.finalize().into() +} + +fn prefix(fragment: usize, version: u64, kind: u8, base_version: u64, digest: &[u8; 32]) -> Vec { + let mut out = Vec::with_capacity(PREFIX_LEN); + out.extend_from_slice(&(fragment as u32).to_le_bytes()); + out.extend_from_slice(&version.to_le_bytes()); + out.push(kind); + out.extend_from_slice(&base_version.to_le_bytes()); + out.extend_from_slice(digest); + out +} + +impl DownlinkAnchors { + pub fn new(fragments: usize, full_every: u64) -> Result { + if full_every < 2 { + bail!("--bcast-delta-full-every must be at least 2"); + } + Ok(Self { + full_every, + anchors: (0..fragments).map(|_| None).collect(), + }) + } + + /// Encode the broadcast of a newly committed `version` of fragment `p` + /// and advance the shared anchor. Call exactly once per new version. + pub fn publish(&mut self, p: usize, version: u64, params: &[f32]) -> Result { + let slot = self + .anchors + .get_mut(p) + .ok_or_else(|| anyhow::anyhow!("downlink anchor for unknown fragment {p}"))?; + let delta_ok = matches!( + slot, + Some(a) if a.deltas_since_full + 1 < self.full_every + && a.values.len() == params.len() + && a.version < version + ); + if !delta_ok { + return Ok(Self::reset(slot, p, version, params)); + } + let anchor = slot.as_mut().unwrap(); + let base_version = anchor.version; + let values = Arc::make_mut(&mut anchor.values); + let mut delta = Vec::new(); + delta + .try_reserve_exact(params.len()) + .map_err(|error| anyhow::anyhow!("cannot allocate downlink delta: {error}"))?; + for (a, t) in values.iter_mut().zip(params) { + let d = half::bf16::from_f32(t - *a).to_f32(); + if !d.is_finite() { + // bf16 overflow of a huge step: fall back to FULL. + return Ok(Self::reset(slot, p, version, params)); + } + delta.push(d); + *a += d; + } + anchor.version = version; + anchor.digest = digest_f32(values); + anchor.deltas_since_full += 1; + Ok(EncodedBcast { + prefix: prefix(p, version, KIND_DELTA, base_version, &anchor.digest), + dtype: DTYPE_BF16, + values: Arc::new(delta), + }) + } + + /// FULL re-send for recovery (HELLO / reconnect): the anchor every other + /// learner holds, when it is at `version`; otherwise the exact params. + pub fn resend(&self, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + match self.anchors.get(p).and_then(Option::as_ref) { + Some(anchor) if anchor.version == version => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &anchor.digest), + dtype: DTYPE_F32, + values: Arc::clone(&anchor.values), + }, + _ => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest_f32(params)), + dtype: DTYPE_F32, + values: Arc::new(params.to_vec()), + }, + } + } + + fn reset(slot: &mut Option, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + let values = Arc::new(params.to_vec()); + let digest = digest_f32(&values); + *slot = Some(Anchor { + version, + values: Arc::clone(&values), + digest, + deltas_since_full: 0, + }); + EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest), + dtype: DTYPE_F32, + values, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A learner-side mirror: applies broadcasts exactly as yeto/protocol.py. + fn apply(anchor: &mut Option<(u64, Vec)>, enc: &EncodedBcast) -> Result<()> { + let version = u64::from_le_bytes(enc.prefix[4..12].try_into().unwrap()); + let base = u64::from_le_bytes(enc.prefix[13..21].try_into().unwrap()); + let digest: [u8; 32] = enc.prefix[21..53].try_into().unwrap(); + let next = match enc.kind() { + KIND_FULL => enc.values.to_vec(), + KIND_DELTA => { + let Some((held, values)) = anchor.as_ref() else { bail!("no anchor") }; + if *held != base { + bail!("anchor version {held} != base {base}"); + } + values.iter().zip(enc.values.iter()).map(|(a, d)| a + d).collect() + } + _ => bail!("bad kind"), + }; + if digest_f32(&next) != digest { + bail!("digest mismatch"); + } + *anchor = Some((version, next)); + Ok(()) + } + + fn params(seed: u64, n: usize) -> Vec { + (0..n) + .map(|i| ((i as f64 * 0.37 + seed as f64 * 1.3).sin() * 3.0) as f32) + .collect() + } + + #[test] + fn deltas_track_params_within_bf16_error_and_refresh_every_m() { + let mut syncer = DownlinkAnchors::new(1, 4).unwrap(); + let mut learner = None; + let mut theta = params(0, 1000); + let mut kinds = vec![]; + for version in 1..=9u64 { + let step = params(version, 1000); + let prev_anchor = learner.as_ref().map(|(_, v): &(u64, Vec)| v.clone()); + for (t, s) in theta.iter_mut().zip(&step) { + *t += 0.01 * s; + } + let enc = syncer.publish(0, version, &theta).unwrap(); + kinds.push(enc.kind()); + apply(&mut learner, &enc).unwrap(); + let (_, held) = learner.as_ref().unwrap(); + if enc.kind() == KIND_FULL { + assert_eq!(held, &theta); + } else { + let prev = prev_anchor.unwrap(); + for ((h, t), a) in held.iter().zip(&theta).zip(&prev) { + // one bf16 rounding of (θ − anchor_prev): rel. 2^-8 bound + assert!((h - t).abs() <= (t - a).abs() / 256.0 + 2.0 * f32::EPSILON * t.abs()); + } + } + } + assert_eq!( + kinds, + vec![KIND_FULL, KIND_DELTA, KIND_DELTA, KIND_DELTA, KIND_FULL, KIND_DELTA, KIND_DELTA, KIND_DELTA, KIND_FULL] + ); + } + + #[test] + fn resend_is_the_shared_anchor_and_lets_a_reconnected_learner_continue() { + let mut syncer = DownlinkAnchors::new(1, 100).unwrap(); + let mut a = None; + let theta1 = params(1, 64); + apply(&mut a, &syncer.publish(0, 1, &theta1).unwrap()).unwrap(); + let theta2: Vec = theta1.iter().map(|v| v + 0.123).collect(); + apply(&mut a, &syncer.publish(0, 2, &theta2).unwrap()).unwrap(); + // A fresh learner (reconnect) gets the anchor, not θ, as FULL. + let mut b = None; + let full = syncer.resend(0, 2, &theta2); + assert_eq!(full.kind(), KIND_FULL); + apply(&mut b, &full).unwrap(); + assert_eq!(a, b); + let theta3: Vec = theta2.iter().map(|v| v - 0.05).collect(); + let enc = syncer.publish(0, 3, &theta3).unwrap(); + assert_eq!(enc.kind(), KIND_DELTA); + apply(&mut a, &enc).unwrap(); + apply(&mut b, &enc).unwrap(); + assert_eq!(a, b); + } + + #[test] + fn a_learner_with_a_wrong_anchor_detects_it() { + let mut syncer = DownlinkAnchors::new(1, 100).unwrap(); + let mut a = None; + apply(&mut a, &syncer.publish(0, 1, ¶ms(1, 32)).unwrap()).unwrap(); + // Corrupt the learner's anchor bits: the digest check must fail. + if let Some((_, v)) = a.as_mut() { + v[3] += 1e-3; + } + let enc = syncer.publish(0, 2, ¶ms(2, 32)).unwrap(); + assert!(apply(&mut a, &enc).unwrap_err().to_string().contains("digest")); + // A learner that skipped a version fails the base-version check. + let mut stale = Some((0u64, params(1, 32))); + let enc = syncer.publish(0, 3, ¶ms(3, 32)).unwrap(); + assert!(apply(&mut stale, &enc).unwrap_err().to_string().contains("base")); + } + + #[test] + fn full_every_below_two_is_refused() { + assert!(DownlinkAnchors::new(1, 1).is_err()); + } +} diff --git a/syncer/src/main.rs b/syncer/src/main.rs index 8243ab299..58ddc9649 100644 --- a/syncer/src/main.rs +++ b/syncer/src/main.rs @@ -2,6 +2,7 @@ mod elastic; mod elastic_server; mod iso_worker; mod merge; +mod downlink; mod protocol; mod server; mod state; @@ -102,6 +103,11 @@ struct Args { /// round is broadcast only after its checkpoint is durable. #[arg(long, default_value_t = false)] async_checkpoint: bool, + /// Learners whose session asks for bf16 downlink deltas get a full f32 + /// BCAST every N-th version of a fragment (N >= 2). Unset: such + /// sessions are refused. + #[arg(long)] + bcast_delta_full_every: Option, /// Resume from --checkpoint-path if it exists. #[arg(long, default_value_t = false)] resume: bool, @@ -278,6 +284,7 @@ fn main() -> anyhow::Result<()> { checkpoint_path: args.checkpoint_path, checkpoint_every: args.checkpoint_every, async_checkpoint: args.async_checkpoint, + bcast_delta_full_every: args.bcast_delta_full_every, resume: args.resume, mark_final_checkpoint: args.mark_final_checkpoint, learner_budget_steps: args.learner_budget_steps, diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index 2103ce3bc..d9bd34bc3 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -62,26 +62,80 @@ pub const DTYPE_Q4: u8 = 3; /// masters and momentum; only the learner's outer gradient is rounded. pub const DTYPE_F32_PUSH_BF16: u8 = 4; +/// Session dtype 5 (S26 E3): PUSH_FRAGMENT deltas travel as block-scaled +/// int8 (see `decode_i8`); INIT_PARAMS, BCAST_FRAGMENT and FINAL stay f32. +pub const DTYPE_F32_PUSH_I8: u8 = 5; +/// Session-dtype flag (S26 E3): BCAST_FRAGMENT carries a 41-byte downlink +/// header and, between periodic f32 refreshes, a bf16 delta against the +/// anchor every learner holds (see `crate::downlink`). Valid only on f32-bulk +/// sessions (dtypes 1, 4, 5). +pub const DTYPE_FLAG_BCAST_DELTA: u8 = 0x80; + /// Values per q4 scale block (f32 absmax scale + 128 nibble bytes each). pub const Q4_BLOCK: usize = 256; +/// Bounds of the int8 push block size the learner declares per payload. +pub const I8_BLOCK_MIN: usize = 32; +pub const I8_BLOCK_MAX: usize = 1 << 20; + +/// The session dtype without its flags. +pub fn session_base(dtype: u8) -> u8 { + dtype & !DTYPE_FLAG_BCAST_DELTA +} + +/// Whether the session receives bf16 downlink deltas. +pub fn bcast_delta(dtype: u8) -> bool { + dtype & DTYPE_FLAG_BCAST_DELTA != 0 +} + +/// Whether `dtype` is a session dtype this syncer accepts. +pub fn valid_session_dtype(dtype: u8) -> bool { + match session_base(dtype) { + DTYPE_F32 | DTYPE_F32_PUSH_BF16 | DTYPE_F32_PUSH_I8 => true, + DTYPE_BF16 | DTYPE_Q4 => !bcast_delta(dtype), + _ => false, + } +} /// Dtype of INIT_PARAMS/BCAST_FRAGMENT tensors for a session dtype: q4 /// applies only to push deltas; full parameter payloads stay bf16. pub fn bulk_dtype(dtype: u8) -> u8 { - match dtype { + match session_base(dtype) { DTYPE_Q4 => DTYPE_BF16, - DTYPE_F32_PUSH_BF16 => DTYPE_F32, + DTYPE_F32_PUSH_BF16 | DTYPE_F32_PUSH_I8 => DTYPE_F32, other => other, } } -/// Dtype of PUSH_FRAGMENT deltas for an unquantized session dtype (q4 has -/// its own block codec). +/// Dtype of PUSH_FRAGMENT deltas for an unquantized session dtype (q4 and +/// int8 have their own block codecs). pub fn push_dtype(dtype: u8) -> u8 { - if dtype == DTYPE_F32_PUSH_BF16 { - DTYPE_BF16 - } else { - dtype + match session_base(dtype) { + DTYPE_F32_PUSH_BF16 => DTYPE_BF16, + other => other, + } +} + +/// Largest PUSH_FRAGMENT delta payload a session dtype may carry. +pub fn max_push_nbytes(dtype: u8, numel: usize) -> Result { + match session_base(dtype) { + DTYPE_Q4 => q4_nbytes(numel), + DTYPE_F32_PUSH_I8 => i8_nbytes(numel, I8_BLOCK_MIN), + other => tensor_nbytes(push_dtype(other), numel), + } +} + +/// Decode a PUSH_FRAGMENT delta payload of any session dtype into f32s. +pub fn decode_push(dtype: u8, bytes: &[u8], numel: usize, out: &mut Vec) -> Result<()> { + match session_base(dtype) { + DTYPE_Q4 => decode_q4(bytes, numel, out), + DTYPE_F32_PUSH_I8 => decode_i8(bytes, numel, out), + other => { + let expected = tensor_nbytes(push_dtype(other), numel)?; + if bytes.len() != expected { + bail!("push delta has {} bytes, expected {expected}", bytes.len()); + } + decode_tensor(push_dtype(other), bytes, out) + } } } @@ -271,6 +325,80 @@ pub fn encode_tensor(dtype: u8, vals: &[f32], out: &mut Vec) -> Result<()> { Ok(()) } +/// Exact encoded byte length of an int8 push payload with block size `block`. +pub fn i8_nbytes(numel: usize, block: usize) -> Result { + if !(I8_BLOCK_MIN..=I8_BLOCK_MAX).contains(&block) { + bail!("int8 block size {block} outside {I8_BLOCK_MIN}..={I8_BLOCK_MAX}"); + } + numel + .div_ceil(block) + .checked_mul(4) + .and_then(|scales| scales.checked_add(numel)) + .and_then(|body| body.checked_add(4)) + .ok_or_else(|| anyhow::anyhow!("int8 tensor byte length overflow")) +} + +/// Decode an int8 delta payload (yeto/protocol.py `f32_bytes_to_i8`). +/// +/// Layout: u32 block size B, then per block of B values (the last block may +/// be shorter): f32 scale (absmax / 127), then one i8 per value. A value +/// decodes to `q as f32 * scale`, so |error| <= scale / 2 = absmax / 254 per +/// element (plus one f32 rounding). The scale is non-negative and finite. +pub fn decode_i8(bytes: &[u8], numel: usize, out: &mut Vec) -> Result<()> { + if bytes.len() < 4 { + bail!("int8 payload is shorter than its block header"); + } + let block = u32::from_le_bytes(bytes[..4].try_into().unwrap()) as usize; + let expected = i8_nbytes(numel, block)?; + if bytes.len() != expected { + bail!( + "int8 payload has {} bytes, expected {expected} for {numel} values in blocks of {block}", + bytes.len() + ); + } + out.clear(); + out.try_reserve(numel) + .map_err(|error| anyhow::anyhow!("cannot allocate int8 tensor: {error}"))?; + let mut at = 4usize; + let mut remaining = numel; + let mut b = 0usize; + while remaining > 0 { + let len = remaining.min(block); + let scale = f32::from_le_bytes(bytes[at..at + 4].try_into().unwrap()); + if !scale.is_finite() || scale < 0.0 { + bail!("int8 block {b} has invalid scale {scale}"); + } + at += 4; + for &q in &bytes[at..at + len] { + let q = q as i8; + if q == i8::MIN { + bail!("int8 block {b} uses the reserved value -128"); + } + out.push(q as f32 * scale); + } + at += len; + remaining -= len; + b += 1; + } + Ok(()) +} + +/// Encode f32s as an int8 push payload (test counterpart of the Python codec). +#[cfg(test)] +pub fn encode_i8(vals: &[f32], block: usize) -> Vec { + let mut out = (block as u32).to_le_bytes().to_vec(); + for chunk in vals.chunks(block) { + let absmax = chunk.iter().fold(0f32, |m, v| m.max(v.abs())); + let scale = absmax / 127.0; + out.extend_from_slice(&scale.to_le_bytes()); + for v in chunk { + let q = if scale > 0.0 { (v / scale).round_ties_even().clamp(-127.0, 127.0) } else { 0.0 }; + out.push(q as i8 as u8); + } + } + out +} + /// Decode a q4 delta payload (yeto/tensor_io.py `quantize_q4`) into f32s. /// /// Per block of `Q4_BLOCK` values: f32 absmax scale, then 128 bytes of @@ -326,6 +454,59 @@ mod tests { use super::*; use tokio::io::AsyncWriteExt; + #[test] + fn int8_push_round_trips_within_half_a_scale_step() { + let vals: Vec = (0..1000) + .map(|i| ((i as f32) * 0.731).sin() * if i % 300 < 150 { 1e-3 } else { 2.5 }) + .collect(); + for block in [I8_BLOCK_MIN, 128, 333, 4096] { + let wire = encode_i8(&vals, block); + assert_eq!(wire.len(), i8_nbytes(vals.len(), block).unwrap()); + assert!(wire.len() <= max_push_nbytes(DTYPE_F32_PUSH_I8, vals.len()).unwrap()); + let mut out = Vec::new(); + decode_push(DTYPE_F32_PUSH_I8, &wire, vals.len(), &mut out).unwrap(); + for (chunk, decoded) in vals.chunks(block).zip(out.chunks(block)) { + let absmax = chunk.iter().fold(0f32, |m, v| m.max(v.abs())); + for (v, d) in chunk.iter().zip(decoded) { + // |error| <= scale/2 = absmax/254, plus f32 rounding. + assert!((v - d).abs() <= absmax / 254.0 * (1.0 + 1e-5), "{v} {d} {absmax}"); + } + } + } + // All-zero blocks decode to exact zeros. + let mut out = Vec::new(); + decode_i8(&encode_i8(&[0.0; 40], 32), 40, &mut out).unwrap(); + assert!(out.iter().all(|v| *v == 0.0)); + } + + #[test] + fn int8_push_rejects_malformed_payloads() { + let vals = vec![1.0f32; 64]; + let mut out = Vec::new(); + let wire = encode_i8(&vals, 32); + assert!(decode_i8(&wire[..wire.len() - 1], 64, &mut out).is_err()); + assert!(decode_i8(&encode_i8(&vals, 16), 64, &mut out).is_err()); + let mut bad_scale = wire.clone(); + bad_scale[4..8].copy_from_slice(&(-1.0f32).to_le_bytes()); + assert!(decode_i8(&bad_scale, 64, &mut out).is_err()); + let mut reserved = wire; + reserved[8] = 0x80; + assert!(decode_i8(&reserved, 64, &mut out).is_err()); + } + + #[test] + fn session_dtype_flags_compose_only_with_f32_bulk() { + assert!(valid_session_dtype(DTYPE_F32_PUSH_I8)); + assert!(valid_session_dtype(DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA)); + assert!(valid_session_dtype(DTYPE_F32_PUSH_BF16 | DTYPE_FLAG_BCAST_DELTA)); + assert!(valid_session_dtype(DTYPE_F32_PUSH_I8 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(DTYPE_BF16 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(DTYPE_Q4 | DTYPE_FLAG_BCAST_DELTA)); + assert!(!valid_session_dtype(6)); + assert_eq!(bulk_dtype(DTYPE_F32_PUSH_I8 | DTYPE_FLAG_BCAST_DELTA), DTYPE_F32); + assert_eq!(push_dtype(DTYPE_F32_PUSH_BF16 | DTYPE_FLAG_BCAST_DELTA), DTYPE_BF16); + } + #[test] fn reserved_message_ids_are_stable() { assert_eq!(MSG_ERROR, 10); diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 6d0b4ba18..46d4cb52b 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -97,6 +97,10 @@ pub struct Config { /// are unchanged, so this is deliberately not part of the semantic /// profile hash. pub async_checkpoint: bool, + /// S26 E3: full-f32 refresh period of bf16 downlink deltas (sessions + /// with `DTYPE_FLAG_BCAST_DELTA`). Wire encoding only: merges, versions + /// and checkpoints are unchanged, so it is not in the profile hash. + pub bcast_delta_full_every: Option, pub resume: bool, /// Create the adjacent final marker only for a completed terminal cut. pub mark_final_checkpoint: bool, @@ -746,6 +750,9 @@ 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.async_checkpoint { if cfg.policy_sweep_fragments.is_none() { bail!("--async-checkpoint requires --policy-sweep-fragments"); @@ -1015,11 +1022,7 @@ fn negotiated_payload_limits(layout: &Layout, dtype: u8) -> Result<(u64, u64)> { for fragment in &layout.fragments { let numel = fragment.numel()?; let init_bytes = tensor_nbytes(bulk_dtype(dtype), numel)? as u64; - let delta_bytes = if dtype == DTYPE_Q4 { - q4_nbytes(numel)? - } else { - tensor_nbytes(push_dtype(dtype), numel)? - } as u64; + let delta_bytes = max_push_nbytes(dtype, numel)? as u64; max_init = max_init.max( 4u64.checked_add(init_bytes) .context("INIT_PARAMS payload limit overflow")?, @@ -1164,7 +1167,7 @@ fn parse_hello(payload: &[u8], expected_learners: u32) -> Result { bail!("connection generation must be nonzero"); } let dtype = r.u8()?; - if !matches!(dtype, DTYPE_F32 | DTYPE_BF16 | DTYPE_Q4 | DTYPE_F32_PUSH_BF16) { + if !valid_session_dtype(dtype) { bail!("unsupported session dtype {dtype}"); } let num_fragments = r.u32()?; @@ -1735,23 +1738,9 @@ async fn dispatch_inner( .with_context(|| format!("PUSH_FRAGMENT for unknown fragment {fragment_id}"))?; let numel = fragment.numel()?; let bytes = r.rest(); - let expected = if group.dtype == DTYPE_Q4 { - q4_nbytes(numel)? - } else { - tensor_nbytes(push_dtype(group.dtype), numel)? - }; - if bytes.len() != expected { - bail!( - "PUSH_FRAGMENT fragment {fragment_id} has {} delta bytes, expected {expected}", - bytes.len() - ); - } let mut outer_gradient = Vec::new(); - if group.dtype == DTYPE_Q4 { - decode_q4(bytes, numel, &mut outer_gradient)?; - } else { - decode_tensor(push_dtype(group.dtype), bytes, &mut outer_gradient)?; - } + decode_push(group.dtype, bytes, numel, &mut outer_gradient) + .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. @@ -1897,6 +1886,17 @@ async fn send_fragment_until_cutoff( dtype: u8, budget_cutoff: &BudgetCutoff, ) -> Result { + if let Some(encoded) = recovery_bcast(st, p, msg_type) { + return send_tensor_until_cutoff( + group, + msg_type, + &encoded.prefix, + encoded.dtype, + &encoded.values, + budget_cutoff, + ) + .await; + } let mut prefix = [0u8; 12]; prefix[..4].copy_from_slice(&(p as u32).to_le_bytes()); prefix[4..].copy_from_slice(&st.versions[p].to_le_bytes()); @@ -1911,6 +1911,17 @@ async fn send_fragment_until_cutoff( .await } +/// A BCAST outside publication (initial cut, HELLO/reconnect recovery) on a +/// bf16-downlink session: always FULL, carrying the shared anchor. +fn recovery_bcast(st: &GlobalState, p: usize, msg_type: u8) -> Option { + if msg_type != MSG_BCAST_FRAGMENT { + return None; + } + st.downlink + .as_ref() + .map(|anchors| anchors.lock().unwrap().resend(p, st.versions[p], &st.params[p])) +} + async fn send_pull_until_cutoff( groups: &[Arc], pull: &bytes::Bytes, @@ -3481,16 +3492,31 @@ async fn publish_round( // its exact f32 cut immediately after the scheduler reaches quiescence. if t < cfg.total_steps { let dtype = bulk_dtype(st.wire_dtype); + // bf16 downlink: encode this version once (advancing the shared + // anchor) and send the same bytes to every learner. + let published = match &st.downlink { + Some(anchors) => Some(anchors.lock().unwrap().publish(p, st.versions[p], &st.params[p])?), + None => None, + }; for group in current_groups(registry) { - match send_fragment_until_cutoff( - st, - &group, - p, - MSG_BCAST_FRAGMENT, - dtype, - budget_cutoff, - ) - .await + let sent = match &published { + Some(encoded) => { + send_tensor_until_cutoff( + &group, + MSG_BCAST_FRAGMENT, + &encoded.prefix, + encoded.dtype, + &encoded.values, + budget_cutoff, + ) + .await + } + None => { + send_fragment_until_cutoff(st, &group, p, MSG_BCAST_FRAGMENT, dtype, budget_cutoff) + .await + } + }; + match sent { Ok(true) => {} Ok(false) => break, @@ -3697,8 +3723,8 @@ async fn complete_round_async( fn new_state_for(group: &Arc, cfg: &Config) -> Result { if let Some(fragments) = cfg.policy_sweep_fragments { // Masters and every BCAST/FINAL stay f32; only push deltas may be bf16. - if !matches!(group.dtype, DTYPE_F32 | DTYPE_F32_PUSH_BF16) { - bail!("--policy-sweep-fragments requires an FP32 HELLO session (optionally with bf16 pushes)"); + if bulk_dtype(group.dtype) != DTYPE_F32 { + bail!("--policy-sweep-fragments requires an FP32 HELLO session (optionally with bf16/int8 pushes)"); } if group.layout.fragments.len() != fragments as usize { bail!( @@ -3717,6 +3743,17 @@ fn new_state_for(group: &Arc, cfg: &Config) -> Result { )?; st.policy_sweep_fragments = cfg.policy_sweep_fragments; st.session_contract_hash = Some(group.session_contract_hash); + match (bcast_delta(group.dtype), cfg.bcast_delta_full_every) { + (true, Some(full_every)) => { + st.downlink = Some(Mutex::new(crate::downlink::DownlinkAnchors::new( + group.layout.fragments.len(), + full_every, + )?)); + } + (true, None) => bail!("session asks for bf16 downlink deltas, but --bcast-delta-full-every is unset"), + (false, Some(_)) => bail!("--bcast-delta-full-every is set, but the session has no bf16 downlink deltas"), + (false, None) => {} + } if cfg.delta_correction { st.delta_correction = Some(crate::merge::Heloco::default()); } @@ -4055,6 +4092,11 @@ async fn send_fragment( msg_type: u8, dtype: u8, ) -> Result<()> { + if let Some(encoded) = recovery_bcast(st, p, msg_type) { + return group + .send_tensor_large(msg_type, &encoded.prefix, encoded.dtype, &encoded.values) + .await; + } let mut prefix = [0u8; 12]; prefix[..4].copy_from_slice(&(p as u32).to_le_bytes()); prefix[4..].copy_from_slice(&st.versions[p].to_le_bytes()); @@ -4697,6 +4739,7 @@ mod tests { checkpoint_path: None, checkpoint_every: 0, async_checkpoint: false, + bcast_delta_full_every: None, resume: false, mark_final_checkpoint: false, learner_budget_steps: None, @@ -5972,13 +6015,13 @@ mod tests { #[tokio::test] async fn updated_broadcast_propagates_local_errors_but_tolerates_closed_streams() { - let invalid_state = broadcast_test_state(255); + let invalid_state = broadcast_test_state(9); let (connected, _receiver) = streaming_test_group(DTYPE_F32, false); let error = broadcast_updated_fragment(&invalid_state, ®istry_with_group(connected), 0) .await .unwrap_err(); assert!( - format!("{error:#}").contains("unknown tensor dtype 255"), + format!("{error:#}").contains("unknown tensor dtype 9"), "unexpected error: {error:#}" ); diff --git a/syncer/src/state.rs b/syncer/src/state.rs index 6d4415612..3b54cb9e1 100644 --- a/syncer/src/state.rs +++ b/syncer/src/state.rs @@ -146,6 +146,8 @@ pub struct GlobalState { pub delta_correction: Option, /// Spectrum flattener. Torch mode owns one persistent child process. iso_backend: IsoBackend, + /// bf16 downlink anchors (S26 E3); in memory only, never checkpointed. + pub downlink: Option>, } /// A merge whose deterministic learner reduction has completed but whose @@ -344,6 +346,7 @@ impl GlobalState { wire_dtype, delta_correction: None, iso_backend: IsoBackend::start(iso_backend_config)?, + downlink: None, }) } diff --git a/tests/test_e3_sync_options.py b/tests/test_e3_sync_options.py index 0e8492ef9..a3256227c 100644 --- a/tests/test_e3_sync_options.py +++ b/tests/test_e3_sync_options.py @@ -2,21 +2,44 @@ from __future__ import annotations +import hashlib +import itertools from dataclasses import replace +import numpy as np import pytest import torch - -from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 +from test_rl_algorithm_provenance import ( # noqa: F401 + _fake_sky, + _island_run, + _launcher_args, + _no_modal_listing, +) +from test_rl_dense_sweep_real_syncer import syncer_binary # noqa: F401 from test_rl_full_parameter_sync import _role_stream_config + +from yeto.downlink import ( + HEADER, + KIND_DELTA, + KIND_FULL, + DownlinkAnchors, + DownlinkChunkSink, + DownlinkResync, +) from yeto.protocol import ( DTYPE_BF16, DTYPE_F32, DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, _f32_parts_to_bf16, + _f32_parts_to_i8, bulk_dtype, f32_bytes_to_bf16, + f32_bytes_to_i8, + i8_nbytes, push_wire_dtype, + valid_session_dtype, ) @@ -44,7 +67,7 @@ def test_push_dtype_is_bound_into_the_role_stream_session_contract(): ) exact = _role_stream_config("actor", 30100, sweep=True) - rounded = replace(exact, push_bf16=True) + rounded = replace(exact, push_dtype="bf16") assert (role_stream_session_contract_hash("1" * 64, exact) != role_stream_session_contract_hash("1" * 64, rounded)) @@ -61,3 +84,235 @@ def test_e3_options_only_apply_to_the_full_parameter_preset(flag, value): args.controller = "local" with pytest.raises(ValueError, match=f"{flag} only applies to --rl-sync-preset full-parameter"): _prepare_rl_args(args) + + +# -- S26 E3 third batch: int8 pushes, bf16 downlink --------------------------- + +def _decode_i8(wire: bytes, numel: int) -> np.ndarray: + block = int(np.frombuffer(wire, "= 2"), + (("--rl-sync-preset", "decoupled", "--local-rl-rounds-per-sync", "2", + "--rl-sync-bcast-delta-full-every", "4"), "only applies to --rl-sync-preset strict-avg"), + (("--rl-sync-push-dtype", "int8"), "--rl-sync-push-dtype only applies"), +)) +def test_e3_downlink_and_int8_scope(extra, match): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", extra, gpu="modal:2xh100") + args.controller = "local" + with pytest.raises(ValueError, match=match): + _prepare_rl_args(args) + + +def _run_bridge_pair(binary, tmp_path, rounds, *, full_every): + """Two fake islands through ``StrictRlBridge`` (the verl/Miles LoRA path).""" + import subprocess + + from test_rl_integration import _BridgeThread, _FakeMiles, _port, _start + + from yeto.rl.bridge import BridgeConfig, StrictRlBridge + + port = _port() + process = _start(binary, port, tmp_path / "state.ckpt", rounds=rounds) + if full_every: + process.kill() + process.wait() + argv = [*process.args, "--bcast-delta-full-every", str(full_every)] + process = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) + runtimes = (_FakeMiles([0.1, 0.3], 0), _FakeMiles([0.3, 0.5], 1)) + threads = [] + try: + for learner_id, runtime in enumerate(runtimes): + config = BridgeConfig( + syncer_addr=("127.0.0.1", port), learner_id=learner_id, global_rounds=rounds, + groups_per_round=1, samples_per_group=2, local_optimizer_steps=1, wan_streams=0, + expected_specs=runtime.current.specs, + base_model_revision=runtime.current.base_model_revision, + lora_config_hash=runtime.current.lora_config_hash, + layout_hash=runtime.current.layout_hash, + event_tape=str(tmp_path / f"island-{learner_id}.jsonl"), + bcast_delta=bool(full_every)) + threads.append(_BridgeThread(StrictRlBridge(runtime, config))) + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + assert not thread.is_alive() and thread.error is None, thread.error + assert process.wait(timeout=10) == 0 + return [next(iter(t.result.tensors.values())) for t in threads], runtimes + finally: + for thread in threads: + thread.bridge.client.close() + if process.poll() is None: + process.kill() + process.wait() + + +def test_strict_bridge_with_bf16_downlink_finishes_on_the_exact_f32_cut(syncer_binary, tmp_path): # noqa: F811 + """E3, verl/LoRA path: islands train from the bf16-delta anchor, the syncer + keeps exact masters, and the final cut matches the exact f32 run (up to the + f32 rounding of ``local - anchor`` on a different anchor).""" + binary = syncer_binary + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact, _ = _run_bridge_pair(binary, exact_dir, 4, full_every=0) + delta, runtimes = _run_bridge_pair(binary, delta_dir, 4, full_every=3) + for want, got in zip(exact, delta, strict=True): + assert torch.allclose(want, got, rtol=0, atol=1e-6), (want, got) + assert all(runtime.applied[-1] == 4 for runtime in runtimes) + + +def test_verl_island_accepts_the_downlink_flag(): + from yeto.rl.adapters.verl.island_entry import parse_args as verl_parse_args + + args = verl_parse_args([ + "--model", "m", "--model-revision", "r", "--data", "d", "--learner-id", "0", + "--reward-function", "f", "--global-rounds", "1", "--groups-per-round", "1", + "--samples-per-group", "1", "--rollout-max-response-len", "8", "--event-tape", "t", + "--inner-lr", "1e-6", "--seq-len", "8", "--seed", "0", "--bcast-delta"]) + assert args[0].bcast_delta is True diff --git a/tests/test_e3_sync_overlap.py b/tests/test_e3_sync_overlap.py new file mode 100644 index 000000000..9b9fbc4d2 --- /dev/null +++ b/tests/test_e3_sync_overlap.py @@ -0,0 +1,331 @@ +"""S26 E3 item 3: sync/train overlap of the full-parameter sweep (CPU only). + +A numeric fake role stream and an in-process strict-sweep syncer drive the +real ``FullParameterSync``/``RolePairing``: the island trains round k+1 while +sweep k is in flight and rebases ``theta <- G_k + (local - base)``. The tests +check that every round's progress reaches the global weights exactly once +(no double counting), that the trainer keeps its progress across a rebase, +and the launcher/island policy-age rule. +""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pytest +from test_rl_algorithm_provenance import ( # noqa: F401 + _fake_sky, + _launcher_args, + _no_modal_listing, +) + +from yeto.protocol import PullRequest +from yeto.rl.adapters.miles.full_parameter.role_stream import RoleSubmissionCandidate +from yeto.rl.engine.full_parameter_sync import ( + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) +from yeto.rl.engine.sync_overlap import rebase_values, required_policy_age + +P = 2 # fragments +NUMEL = 4 + + +class _SweepSyncer: + """One-island strict sweep: PULLs a sweep once the previous one committed; + ``G_f += pushed delta``; the terminal step publishes FINAL.""" + + def __init__(self, initial: list[np.ndarray], sweeps: int) -> None: + self.global_ = [value.copy() for value in initial] + self.versions = [0] * P + self.sweeps = sweeps + self.step = 0 # last committed global step + self.pushed: list[tuple[int, np.ndarray]] = [] + self.bcasts: list[tuple[int, int, np.ndarray]] = [] + self.finalizing = False + self._since_commit = 0 + + def permits(self) -> list[PullRequest]: + sweep = self.step // P # completed sweeps + if self.finalizing or sweep >= self.sweeps or self.step % P: + return [] + pushed = {step for step, _ in self.pushed} + return [PullRequest(f, sweep * P + f + 1, 1) for f in range(P) if sweep * P + f + 1 not in pushed] + + def push(self, step: int, fragment_id: int, delta: np.ndarray) -> None: + assert self.step < step <= self.step + P and (step - 1) % P == fragment_id + self.pushed.append((step, delta.copy())) + self.global_[fragment_id] = self.global_[fragment_id] + delta + self.versions[fragment_id] = step + self._since_commit += 1 + if self._since_commit < P: + return + self._since_commit = 0 + self.step += P + if self.step == self.sweeps * P: + self.finalizing = True + else: + self.bcasts.extend((f, self.versions[f], self.global_[f].copy()) for f in range(P)) + + +class _NumericRole: + """A RoleStream whose cuts are lists of per-fragment f32 arrays.""" + + role = "actor" + + def __init__(self, syncer: _SweepSyncer, progress) -> None: + self.syncer = syncer + self.progress = progress # rollout_id -> list of per-fragment progress + self.config = SimpleNamespace(overlap=True) + self.current = [g.copy() for g in syncer.global_] + self.anchors = [g.copy() for g in syncer.global_] + self._versions = [0] * P + self.submission = None + self.submission_versions = None + self.local_rounds = self.submission_rounds = 0 + self.rounds_at_anchor = [0] * P + self.round_base = None + self.rounds_at_base = 0 + self._pending: dict[int, PullRequest] = {} + self._staged: dict[tuple[int, int], np.ndarray] = {} + self.terminal_submitted = False + self.final_acknowledged = False + self.closed = False + self.trained_from: list[list[np.ndarray]] = [] + + @property + def fragment_versions(self): + return tuple(self._versions) + + @property + def finalizing(self): + return self.syncer.finalizing + + async def initialize(self, group): + del group + + def resident_state(self, policy_version): + return SimpleNamespace(policy_version=policy_version, policy_tensor_hash=lambda: "cut") + + async def complete_local_round(self, *, rollout_id, trajectories, hold_submission=False): + del trajectories + self.trained_from.append([c.copy() for c in self.current]) + self.current = [c + d for c, d in zip(self.current, self.progress(rollout_id), strict=True)] + self.local_rounds += 1 + if not hold_submission: + self.submission = [c.copy() for c in self.current] + self.submission_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + + def ready_submission_candidates(self): + for permit in self.syncer.permits(): + self._pending.setdefault(permit.fragment_id, permit) + out = [] + for f, permit in sorted(self._pending.items()): + base = max(0, permit.global_step - P) + if (self.submission is None or self._versions[f] != base + or self.submission_versions[f] != base + or self.submission_rounds - self.rounds_at_anchor[f] < 1): + continue + out.append(RoleSubmissionCandidate( + "actor", permit, base, self.submission_rounds, 1, 1, NUMEL * 4, + self.submission, self.anchors[f])) + return tuple(out) + + def retain_candidate(self, candidate): + del candidate + + def release_candidate(self, candidate): + del candidate + + def push_candidate(self, candidate): + f = candidate.permit.fragment_id + self.syncer.push(candidate.permit.global_step, f, + candidate.submission_cut[f] - candidate.anchor) + self._pending.pop(f) + self.terminal_submitted |= candidate.permit.global_step == self.syncer.sweeps * P + return SimpleNamespace(role="actor", fragment_id=f, global_step=candidate.permit.global_step) + + async def stage_available_broadcasts(self): + for f, version, value in self.syncer.bcasts: + if version > self._versions[f]: + self._staged[(f, version)] = value + self.syncer.bcasts.clear() + return tuple(self._staged) + + @property + def staged_fragment_versions(self): + return tuple(sorted(self._staged)) + + def has_staged_broadcast(self, fragment_id, version): + return (fragment_id, version) in self._staged + + async def apply_staged_broadcast(self, fragment_id, version, *, rebase=False): + assert rebase, "overlap applies every sweep by rebasing" + g = self._staged.pop((fragment_id, version)) + self.current[fragment_id] = rebase_values( + g, self.current[fragment_id], self.round_base[fragment_id]) + self.anchors[fragment_id] = g + self._versions[fragment_id] = version + self.rounds_at_anchor[fragment_id] = self.rounds_at_base + return (fragment_id,) + + def finish_rebase(self): + self.submission = [c.copy() for c in self.current] + self.submission_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + + def mark_round_base(self): + self.round_base = [c.copy() for c in self.current] + self.rounds_at_base = self.local_rounds + + async def wait_and_apply_final(self): + self.current = [g.copy() for g in self.syncer.global_] + self._versions = [self.syncer.sweeps * P - P + f + 1 for f in range(P)] + + def acknowledge_final(self): + self.final_acknowledged = True + + def discard_detached_submission_cut(self): + pass + + def close(self): + self.closed = True + + +class _Driver: + def __init__(self): + self.events = [] + + def emit(self, event, **fields): + self.events.append((event, fields)) + + +def _run(sweeps: int, overlap: bool = True): + rng = np.random.default_rng(7) + initial = [rng.standard_normal(NUMEL).astype(np.float32) for _ in range(P)] + steps = {k: [rng.standard_normal(NUMEL).astype(np.float32) * 0.1 for _ in range(P)] + for k in range(sweeps)} + syncer = _SweepSyncer(initial, sweeps) + role = _NumericRole(syncer, lambda rollout_id: steps[rollout_id]) + profile = FullParameterSyncProfile(roles=("actor",), total_fragment_steps=sweeps * P, fragments=P, + local_horizon=1, pipeline_depth=1, wait_timeout=5.0, + poll_seconds=0.001, sweep=True, overlap=overlap) + session = FullParameterSync(RolePairing({"actor": role}, profile), groups={"actor": "g"}, + run=asyncio.run, trajectories=lambda rollout_id: rollout_id) + driver = _Driver() + session.start(driver) + stops = [] + for rollout_id in range(sweeps): + boundary = session.boundary(driver, rollout_id=rollout_id, stats=None) + stops.append(boundary.stop) + if boundary.stop: + session.published(driver, rollout_id=rollout_id, policy_hash="x") + break + return initial, steps, syncer, role, stops + + +def test_overlap_counts_every_round_exactly_once(): + sweeps = 5 + initial, steps, syncer, role, stops = _run(sweeps) + assert stops == [False] * (sweeps - 1) + [True] + # Pushed deltas are exactly the per-round progress, in round order. + pushed = sorted(syncer.pushed, key=lambda row: row[0]) + for k in range(sweeps): + for f in range(P): + np.testing.assert_allclose(pushed[k * P + f][1], steps[k][f], rtol=0, atol=1e-6) + # Global = initial + sum of every round's progress, once. + for f in range(P): + want = initial[f] + sum(steps[k][f] for k in range(sweeps)) + np.testing.assert_allclose(syncer.global_[f], want, rtol=0, atol=1e-5) + np.testing.assert_allclose(role.current[f], want, rtol=0, atol=1e-5) + assert role.final_acknowledged + + +def test_overlap_trains_one_outer_step_stale_and_keeps_local_progress(): + sweeps = 4 + initial, steps, _syncer, role, _ = _run(sweeps) + # Round 0 starts from the initial cut, round 1 from its own round-0 result + # (sweep 0 still in flight), round k>=2 from G_{k-2} + progress of round k-1. + for f in range(P): + np.testing.assert_allclose(role.trained_from[0][f], initial[f]) + np.testing.assert_allclose(role.trained_from[1][f], initial[f] + steps[0][f], atol=1e-6) + for k in range(2, sweeps): + g_prev = initial[f] + sum(steps[j][f] for j in range(k - 1)) + np.testing.assert_allclose(role.trained_from[k][f], g_prev + steps[k - 1][f], atol=1e-5) + + +def test_rebase_values_is_anchor_plus_local_progress(): + anchor = np.array([1.0, -2.0], dtype=np.float32) + current = np.array([0.5, 0.25], dtype=np.float32) + base = np.array([0.25, 0.5], dtype=np.float32) + np.testing.assert_array_equal(rebase_values(anchor, current, base), [1.25, -2.25]) + with pytest.raises(ValueError, match="float32"): + rebase_values(anchor.astype(np.float64), current, base) + with pytest.raises(ValueError, match="non-finite"): + rebase_values(np.array([np.inf, 0], np.float32), current, base) + + +def test_overlap_profile_needs_the_strict_sweep(): + with pytest.raises(ValueError, match="overlap needs the dense strict sweep"): + FullParameterSyncProfile(roles=("actor",), total_fragment_steps=4, fragments=2, local_horizon=2, + pipeline_depth=1, wait_timeout=1.0, poll_seconds=0.01, overlap=True) + + +def _dense(*extra): + from yeto.launcher import _prepare_rl_args + + args = _launcher_args("ports", ( + "--rl-sync-preset", "full-parameter", "--rl-placement", "fixed-partition", + "--rl-rollout-gpus", "1", "--rl-parameter-layout-sha256", "e" * 64, + "--fragments", "2", "--total-steps", "2", *extra), gpu="modal:2xh100") + args.controller = "local" + _prepare_rl_args(args) + return args + + +def test_overlap_raises_the_minimum_policy_age_and_reaches_the_island(): + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + assert required_policy_age(0) == 1 + args = _dense("--rl-sync-overlap") + assert args.rl_max_policy_age == 1 + assert "--sweep-overlap" in _rl_trainable_island_flags(args, 0, 1) + assert _dense("--rl-sync-overlap", "--rl-max-policy-age", "3").rl_max_policy_age == 3 + plain = _dense() + assert int(plain.rl_max_policy_age or 0) == 0 + assert "--sweep-overlap" not in _rl_trainable_island_flags(plain, 0, 1) + + +def test_overlap_launch_fails_when_the_algorithm_spec_does_not_tolerate_the_age(): + from yeto.launch_preflight import PreflightError, check_policy_age_spec + + args = _dense("--rl-sync-overlap") + with pytest.raises(PreflightError, match="--rl-max-policy-age 1 needs an algorithm spec"): + check_policy_age_spec(args) + + +def test_island_refuses_overlap_without_the_raised_policy_age(): + from yeto.rl.adapters.miles.island_entry import _dense_full_sync_config + + args = SimpleNamespace(sweep_overlap=True, rl_max_policy_age=0) + with pytest.raises(ValueError, match="--sweep-overlap .* needs --rl-max-policy-age >= 1"): + _dense_full_sync_config(args, model_config_sha256="0" * 64, + training_contract_sha256="0" * 64, algorithm_spec_sha256="0" * 64) + + +def test_overlap_is_bound_into_the_role_stream_contract(): + from dataclasses import replace + + from test_rl_full_parameter_sync import _role_stream_config + + from yeto.rl.adapters.miles.full_parameter.role_stream import ( + role_stream_session_contract_hash, + ) + + exact = _role_stream_config("actor", 30100, sweep=True) + assert (role_stream_session_contract_hash("1" * 64, exact) + != role_stream_session_contract_hash("1" * 64, replace(exact, overlap=True))) diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 5d01ad197..31589d679 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -21,7 +21,13 @@ from test_rl_algorithm_provenance import _fake_sky, _launcher_args, _no_modal_listing # noqa: F401 from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout -from yeto.protocol import DTYPE_F32, DTYPE_F32_PUSH_BF16, SyncerClient +from yeto.protocol import ( + DTYPE_F32, + DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, + SyncerClient, +) ROOT = Path(__file__).resolve().parent.parent CARGO = shutil.which("cargo") or str(Path.home() / ".cargo/bin/cargo") @@ -78,10 +84,11 @@ def _layout() -> FragmentLayout: return FragmentLayout([Fragment(MERGE_RDA, [(f"actor.w{i}", NUMEL)]) for i in range(FRAGMENTS)]) -def _client(port: int, profile_hash: str, dtype: int = DTYPE_F32) -> SyncerClient: +def _client(port: int, profile_hash: str, dtype: int = DTYPE_F32, *, + max_reconnects: int = 0, push_i8_block: int = 128) -> SyncerClient: return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=dtype, num_streams=1, - connect_timeout=10, max_reconnects=0, finalization_timeout=20, - syncer_profile_hash=bytes.fromhex(profile_hash)) + connect_timeout=10, max_reconnects=max_reconnects, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), push_i8_block=push_i8_block) def _f32(values: list[float]) -> bytes: @@ -161,10 +168,18 @@ def test_launcher_dense_syncer_refuses_another_profile_hash(syncer_binary, tmp_p def _run_sweeps(binary: Path, tmp_path: Path, extra: tuple[str, ...], sweeps: int, - dtype: int = DTYPE_F32): - """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint).""" + dtype: int = DTYPE_F32, *, after_sweep=None, max_reconnects: int = 0, + push_i8_block: int = 128, updates: dict | None = None, sink=None): + """Drive ``sweeps`` strict sweeps; return (max PULLs seen at once, finals, checkpoint). + + ``updates`` collects every BCAST the client accepted, as + ``{(fragment_id, version): f32 bytes}``; ``after_sweep(client, sweep)`` + runs once all pushes of a sweep are queued.""" process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) - client = _client(port, profile_hash, dtype) + client = _client(port, profile_hash, dtype, max_reconnects=max_reconnects, + push_i8_block=push_i8_block) + if sink is not None: + client.install_inbound_chunk_sink(sink) concurrent = 0 try: client.start() @@ -195,7 +210,11 @@ def gather(): client.push_fragment(fid, permit.global_step, permit.round_attempt, base, sweep, 1, 7, _f32(delta)) permit.pushed = True + if after_sweep is not None: + after_sweep(client, sweep) manifest, finals = client.wait_for_final_fragments(timeout=30) + if updates is not None: + updates.update({(u.fragment_id, u.version): bytes(u.data) for u in client.drain_updates()}) result = {f.fragment_id: bytes(f.data) for f in finals} client.acknowledge_finalization(manifest) assert process.wait(timeout=15) == 0 @@ -270,3 +289,146 @@ def test_bf16_push_flag_reaches_the_island_command(): args = _dense_args("--rl-sync-push-dtype", "bf16") assert "--sweep-push-dtype bf16" in _rl_trainable_island_flags(args, 0, 1) assert "--sweep-push-dtype" not in _rl_trainable_island_flags(_dense_args(), 0, 1) + + +def _tail(data: bytes) -> torch.Tensor: + return torch.frombuffer(bytearray(data), dtype=torch.float32) + + +def test_int8_pushes_stay_within_their_block_error_of_the_f32_sweep(syncer_binary, tmp_path): + """E3: int8 outer deltas change results only by half a scale step per delta.""" + exact_dir, int8_dir = tmp_path / "f32", tmp_path / "int8" + exact_dir.mkdir() + int8_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=3) + quantized = _run_sweeps(syncer_binary, int8_dir, + (*STREAMED, "--rl-sync-push-dtype", "int8", "--rl-sync-push-int8-block", "32"), + sweeps=3, dtype=DTYPE_F32_PUSH_I8, push_i8_block=32) + assert exact[1] == quantized[1] == [5, 6] + for fid, data in exact[2].items(): + want, got = _tail(data)[-NUMEL:], _tail(quantized[2][fid])[-NUMEL:] + # Three sweeps; each delta's absmax < 2, so error <= 3 * 2 / 254. + assert torch.allclose(got, want, rtol=0, atol=3 * 2.0 / 254), fid + assert not torch.equal(got, want), "int8 pushes must actually quantize" + + +def _downlink_run(binary, tmp_path, full_every: int, sweeps: int, **kwargs): + updates: dict = {} + result = _run_sweeps(binary, tmp_path, + (*STREAMED, "--rl-sync-bcast-delta-full-every", str(full_every)), + sweeps=sweeps, dtype=DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA, + updates=updates, **kwargs) + return result, updates + + +def test_bf16_downlink_tracks_the_exact_broadcasts_and_leaves_the_merge_exact(syncer_binary, tmp_path): + """E3: BCAST deltas change only what islands hold between f32 refreshes; + the syncer's masters, versions, FINAL and checkpoint stay bit-identical.""" + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact_updates: dict = {} + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=4, updates=exact_updates) + delta, delta_updates = _downlink_run(syncer_binary, delta_dir, 3, sweeps=4) + assert exact[1] == delta[1] == [7, 8] + assert exact[2] == delta[2], "FINAL must be the exact f32 masters" + assert exact[3] == delta[3], "checkpoints must not depend on the downlink" + assert set(exact_updates) == set(delta_updates) + # Fragment 0 publishes versions 1, 3, 5, 7: FULL, DELTA, DELTA, FULL (M=3). + kinds = {} + for key, data in sorted(delta_updates.items()): + want, got = _tail(exact_updates[key]), _tail(data) + kinds[key] = "full" if torch.equal(got, want) else "delta" + # one bf16 rounding of (theta - anchor), |theta - anchor| < 4 here + assert torch.allclose(got, want, rtol=0, atol=4.0 / 256), key + assert [kinds[(0, v)] for v in (0, 1, 3, 5, 7)] == ["full", "full", "delta", "delta", "full"] + assert [kinds[(1, v)] for v in (0, 2, 4, 6)] == ["full", "full", "delta", "delta"] + + +def test_bf16_downlink_anchor_mismatch_reconnects_and_resyncs_with_f32(syncer_binary, tmp_path): + """E3 fallback: a learner whose anchor drifted fails the digest check, + reconnects, receives the shared anchor as FULL, and the run finishes + with the same exact result.""" + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + exact = _run_sweeps(syncer_binary, exact_dir, STREAMED, sweeps=4) + corrupted, reconnects = [], [] + + def corrupt(client, sweep): + if sweep == 4: + reconnects.append(client._reconnects_used) + if sweep != 1: + return + _wait(lambda: client._downlink.version(0) == 1, "fragment 0 version 1 never arrived") + held = client._downlink._held[0] + held.values = held.values.copy() + held.values[0] += 1e-3 + corrupted.append(client.connection_generation) + + delta, _ = _downlink_run(syncer_binary, delta_dir, 100, sweeps=4, after_sweep=corrupt, + max_reconnects=5) + assert corrupted and reconnects and reconnects[0] >= 1, "the digest mismatch must reconnect" + assert exact[1] == delta[1] and exact[2] == delta[2] and exact[3] == delta[3] + + +def test_bf16_downlink_and_int8_flags_reach_syncer_and_island_commands(): + from yeto import launcher + from yeto.launch.rl_island_task import _rl_trainable_island_flags + + args = _dense_args("--rl-sync-push-dtype", "int8", "--rl-sync-push-int8-block", "64", + "--rl-sync-bcast-delta-full-every", "4") + island = _rl_trainable_island_flags(args, 0, 1) + assert "--sweep-push-dtype int8" in island and "--sweep-push-int8-block 64" in island + command = launcher.role_syncer_commands(args, 1, binary="yeto-syncer")["actor"] + assert "--bcast-delta-full-every 4" in command + plain = _dense_args() + assert "--bcast-delta" not in launcher.role_syncer_commands(plain, 1, binary="yeto-syncer")["actor"] + with pytest.raises(ValueError, match="--rl-sync-bcast-delta-full-every 0"): + _dense_args("--rl-sync-bcast-delta-full-every", "1") + with pytest.raises(ValueError, match="--rl-sync-push-int8-block only with"): + _dense_args("--rl-sync-push-int8-block", "64") + + +class _BytesSink: + """Minimal transactional inbound sink: assembles the f32 fragment bytes.""" + + def begin_message(self, message_id, total_bytes): + return {"buf": bytearray(), "bound": None} + + def bind_fragment(self, tx, msg_type, fragment_id, version, payload_bytes): + tx["bound"] = payload_bytes + + def stage_chunk(self, tx, payload): + return bytes(payload) + + def consume_staged_chunk(self, tx, offset, token, nbytes): + self.consume_chunk(tx, offset, token) + + def consume_chunk(self, tx, offset, payload): + assert offset == len(tx["buf"]) + tx["buf"] += bytes(payload) + + def finish_message(self, tx): + import hashlib + + from yeto.protocol import StreamedInboundPayload + + assert len(tx["buf"]) == tx["bound"] + return StreamedInboundPayload(bytes(tx["buf"]), hashlib.sha256(tx["buf"]).digest()) + + def abort_message(self, tx, staged_tokens): + tx["buf"].clear() + + +def test_bf16_downlink_through_a_streaming_sink_matches_the_buffered_path(syncer_binary, tmp_path): + """E3: the Miles path streams BCAST chunks into a transactional sink; the + downlink wrapper must hand it the same f32 anchor bytes.""" + buffered_dir, streamed_dir = tmp_path / "buffered", tmp_path / "streamed" + buffered_dir.mkdir() + streamed_dir.mkdir() + buffered, buffered_updates = _downlink_run(syncer_binary, buffered_dir, 3, sweeps=3) + streamed, streamed_updates = _downlink_run(syncer_binary, streamed_dir, 3, sweeps=3, + sink=_BytesSink()) + assert buffered[2] == streamed[2] + assert buffered_updates == streamed_updates diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 99b5ea017..9cdba9d46 100644 --- a/tests/test_run_flags.py +++ b/tests/test_run_flags.py @@ -250,7 +250,7 @@ def test_forward_tables_name_flags_of_both_layers(): # flag renamed in the schema but still spelled the old way in the launcher fails. _NON_RL_LITERALS = frozenset({ "--abbrev-ref", "--adapter-dir", "--add", "--address", "--allow-unattested-legacy-adapter", - "--allow-unauthenticated-islands", "--async-checkpoint", "--bind", "--cap-add", "--carry-gamma", + "--allow-unauthenticated-islands", "--async-checkpoint", "--bcast-delta-full-every", "--bind", "--cap-add", "--carry-gamma", "--checkpoint-every", "--checkpoint-path", "--codex-harness-binary", "--config-json", "--create", "--dashboard-host", "--dearmor", "--depth", "--detach", "--dtype", "--entrypoint", "--expected-island-contract", "--final-ack-timeout-s", "--final-grace-s", "--follow", "--force", diff --git a/yeto/downlink.py b/yeto/downlink.py new file mode 100644 index 000000000..d3a613c09 --- /dev/null +++ b/yeto/downlink.py @@ -0,0 +1,235 @@ +"""Learner side of bf16 downlink deltas (S26 E3; syncer ``src/downlink.rs``). + +A session whose dtype carries ``DTYPE_FLAG_BCAST_DELTA`` receives every +BCAST_FRAGMENT with a 41-byte downlink header after ``fragment_id, version``:: + + kind u8 (0 FULL, 1 DELTA) | base_version u64 | anchor_sha256 [32] + +FULL carries f32 values; DELTA carries ``bf16(theta - anchor)``. The client +keeps one f32 *anchor* per fragment, bit-identical to the syncer's: FULL sets +it, DELTA adds the decoded delta in f32 (numpy float32 addition rounds like +Rust's). The digest is checked after every update. Callers above the client +(``drain_updates``, inbound chunk sinks) only ever see the reconstructed f32 +anchor, so adapters, N7.6 cuts and manifests are unchanged. + +A missing or mismatched anchor raises :class:`DownlinkResync` (a +``ConnectionError``): the client drops its connection group, and the +syncer's reconnect recovery re-sends every fragment as FULL. +""" + +from __future__ import annotations + +import hashlib +import struct +import threading +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: # numpy is imported lazily: yeto.protocol stays import-light + import numpy as np + +KIND_FULL = 0 +KIND_DELTA = 1 +HEADER = struct.Struct(" int: + if kind == KIND_FULL: + return 4 * numel + if kind == KIND_DELTA: + return 2 * numel + raise ValueError(f"unknown downlink kind {kind}") + + +@dataclass +class _Held: + version: int + values: np.ndarray # float32, never mutated after commit + + +class DownlinkDecode: + """Streams one BCAST's tensor bytes into the next f32 anchor.""" + + def __init__(self, fragment_id: int, version: int, kind: int, base: np.ndarray | None, + digest: bytes, numel: int) -> None: + import numpy as np + + self.fragment_id = fragment_id + self.version = version + self.kind = kind + self._base = base + self._digest = digest + self.values = np.empty(numel, dtype=np.float32) + self._width = 4 if kind == KIND_FULL else 2 + self._carry = b"" + self._filled = 0 + self._hasher = hashlib.sha256() + + def feed(self, data: bytes | bytearray | memoryview) -> bytes: + """Decode the next wire bytes; return the f32 bytes they complete.""" + import numpy as np + + raw = self._carry + bytes(data) + whole = len(raw) - len(raw) % self._width + self._carry = raw[whole:] + if not whole: + return b"" + count = whole // self._width + start, end = self._filled, self._filled + count + if end > self.values.size: + raise ValueError("downlink BCAST has more values than its fragment") + out = self.values[start:end] + if self.kind == KIND_FULL: + out[:] = np.frombuffer(raw, dtype=" None: + if self._carry or self._filled != self.values.size: + raise ValueError("downlink BCAST ended before its fragment was complete") + if self._hasher.digest() != self._digest: + raise DownlinkResync( + f"downlink anchor digest mismatch for fragment {self.fragment_id} " + f"version {self.version}") + + +class DownlinkAnchors: + """Per-fragment f32 anchors shared (bit for bit) with the syncer.""" + + def __init__(self, num_fragments: int) -> None: + self._held: list[_Held | None] = [None] * num_fragments + self._lock = threading.Lock() + + def version(self, fragment_id: int) -> int | None: + with self._lock: + held = self._held[fragment_id] + return None if held is None else held.version + + def begin(self, fragment_id: int, version: int, header: bytes, numel: int) -> DownlinkDecode: + kind, base_version, digest = HEADER.unpack(header) + if kind == KIND_FULL: + return DownlinkDecode(fragment_id, version, kind, None, digest, numel) + if kind != KIND_DELTA: + raise ValueError(f"unknown downlink kind {kind}") + with self._lock: + held = self._held[fragment_id] + if held is None or held.version != base_version or held.values.size != numel: + raise DownlinkResync( + f"downlink delta for fragment {fragment_id} version {version} needs anchor " + f"version {base_version}; learner holds " + f"{None if held is None else held.version}") + return DownlinkDecode(fragment_id, version, kind, held.values, digest, numel) + + def commit(self, decode: DownlinkDecode) -> None: + """Adopt a finished decode as the anchor unless a newer one is held.""" + with self._lock: + held = self._held[decode.fragment_id] + if held is None or decode.version >= held.version: + self._held[decode.fragment_id] = _Held(decode.version, decode.values) + + def decode_whole(self, fragment_id: int, version: int, payload: bytes | memoryview, + numel: int) -> bytes: + """Non-streamed BCAST: header + tensor -> committed f32 anchor bytes.""" + if len(payload) < HEADER_LEN: + raise ValueError("downlink BCAST header truncated") + decode = self.begin(fragment_id, version, bytes(payload[:HEADER_LEN]), numel) + if len(payload) - HEADER_LEN != tensor_nbytes(decode.kind, numel): + raise ValueError(f"downlink BCAST {fragment_id} has {len(payload) - HEADER_LEN} " + f"tensor bytes, expected {tensor_nbytes(decode.kind, numel)}") + data = decode.feed(payload[HEADER_LEN:]) + decode.finish() + self.commit(decode) + return data + + +@dataclass +class _Transaction: + inner: object + passthrough: bool = True + fragment_id: int | None = None + version: int | None = None + numel: int = 0 + header: bytearray = field(default_factory=bytearray) + decode: DownlinkDecode | None = None + out_offset: int = 0 + + +class DownlinkChunkSink: + """Wraps a transactional inbound chunk sink for a downlink session. + + BCAST payload bytes (downlink header + f32/bf16 tensor) are decoded in + order and handed to the wrapped sink as the plain f32 fragment; FINAL + passes through. Out-of-order chunks are staged here as bytes (the + wrapped sink sees only in-order f32 bytes).""" + + def __init__(self, inner: object, anchors: DownlinkAnchors, numels: list[int], + bcast_type: int) -> None: + self.inner = inner + self.anchors = anchors + self.numels = numels + self.bcast_type = bcast_type + + def begin_message(self, message_id: int, total_bytes: int) -> _Transaction: + return _Transaction(self.inner.begin_message(message_id, total_bytes)) + + def bind_fragment(self, tx: _Transaction, msg_type: int, fragment_id: int, version: int, + payload_bytes: int) -> None: + if msg_type != self.bcast_type: + self.inner.bind_fragment(tx.inner, msg_type, fragment_id, version, payload_bytes) + return + tx.passthrough = False + tx.fragment_id, tx.version = fragment_id, version + tx.numel = self.numels[fragment_id] + self.inner.bind_fragment(tx.inner, msg_type, fragment_id, version, 4 * tx.numel) + + def stage_chunk(self, tx: _Transaction, payload) -> bytes: + del tx + return bytes(payload) + + def consume_staged_chunk(self, tx: _Transaction, offset: int, token: bytes, nbytes: int) -> None: + if len(token) != nbytes: + raise ValueError("staged downlink chunk is malformed") + self.consume_chunk(tx, offset, token) + + def consume_chunk(self, tx: _Transaction, offset: int, payload) -> None: + if tx.passthrough: + self.inner.consume_chunk(tx.inner, offset, payload) + return + data = memoryview(payload).cast("B") + if tx.decode is None: + take = min(HEADER_LEN - len(tx.header), data.nbytes) + tx.header += data[:take] + data = data[take:] + if len(tx.header) < HEADER_LEN: + return + tx.decode = self.anchors.begin(tx.fragment_id, tx.version, bytes(tx.header), tx.numel) + out = tx.decode.feed(data) + if out: + self.inner.consume_chunk(tx.inner, tx.out_offset, out) + tx.out_offset += len(out) + + def finish_message(self, tx: _Transaction): + if tx.passthrough: + return self.inner.finish_message(tx.inner) + if tx.decode is None: + raise ValueError("downlink BCAST ended inside its header") + tx.decode.finish() + committed = self.inner.finish_message(tx.inner) + self.anchors.commit(tx.decode) + return committed + + def abort_message(self, tx: _Transaction, staged_tokens) -> None: + del staged_tokens # plain bytes held here + self.inner.abort_message(tx.inner, ()) diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index df0c0efbb..60e63f88b 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -153,16 +153,35 @@ def sao_contract_mounts(args, learner_id: int) -> dict[str, str]: ("rl_sync_async_checkpoint", False, "--rl-sync-async-checkpoint"), ("rl_sync_push_dtype", "f32", "--rl-sync-push-dtype"), ("rl_sync_checkpoint_every", 1, "--rl-sync-checkpoint-every"), + ("rl_sync_push_int8_block", 128, "--rl-sync-push-int8-block"), + ("rl_sync_overlap", False, "--rl-sync-overlap"), ) def _refuse_full_parameter_sync_options(args) -> None: - """The S26 E3 sync options exist only for the full-parameter preset.""" + """The S26 E3 sync options exist only for the full-parameter preset + (the bf16 downlink also for strict-avg, see ``_check_bcast_delta``).""" for dest, default, option in _FULL_PARAMETER_SYNC_DEFAULTS: if getattr(args, dest, default) != default: raise ValueError(f"{option} only applies to --rl-sync-preset {FULL_PARAMETER_SYNC_PRESET}") +def _check_bcast_delta(args) -> None: + """S26 E3 bf16 downlink: a legacy-scheduled strict-avg or full-parameter + syncer only (the elastic server and SAO have no downlink anchors).""" + full_every = int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) + if full_every == 0: + return + if full_every < 2: + raise ValueError("--rl-sync-bcast-delta-full-every 0 (off) or >= 2") + if args.rl_sync_preset not in ("strict-avg", FULL_PARAMETER_SYNC_PRESET): + raise ValueError("--rl-sync-bcast-delta-full-every only applies to --rl-sync-preset " + f"strict-avg or {FULL_PARAMETER_SYNC_PRESET}") + if island_scheduling_mode(args) != "legacy" or args.experimental_rl_sync: + raise ValueError("--rl-sync-bcast-delta-full-every needs --rl-island-scheduling legacy " + "and no --experimental-rl-sync") + + def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: """full-parameter preset (P0-D8): one actor role swept as whole-model fragments. @@ -190,6 +209,11 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: checkpoint_every = int(getattr(args, "rl_sync_checkpoint_every", 1)) if checkpoint_every != 1 and (checkpoint_every < 1 or checkpoint_every % max(args.fragments, 1)): problems.append("--rl-sync-checkpoint-every 1 or a multiple of --fragments") + int8_block = int(getattr(args, "rl_sync_push_int8_block", 128)) + if not 32 <= int8_block <= 1 << 20: + problems.append("--rl-sync-push-int8-block between 32 and 1048576") + elif int8_block != 128 and getattr(args, "rl_sync_push_dtype", "f32") != "int8": + problems.append("--rl-sync-push-int8-block only with --rl-sync-push-dtype int8") if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands @@ -339,6 +363,15 @@ def _prepare_rl_args( # what the selected backend supports before any cloud or GPU work. from ..rl.engine.policy_age import validate_limit + if args.rl_sync_preset == FULL_PARAMETER_SYNC_PRESET and getattr(args, "rl_sync_overlap", False): + # S26 E3: sync/train overlap adds one outer step of staleness; raise the + # limit to the new minimum here so every check below (backend support, + # and the preflight's algorithm-spec staleness) validates it at launch. + from ..rl.engine.sync_overlap import required_policy_age + + required = required_policy_age(0) + if int(getattr(args, "rl_max_policy_age", 0) or 0) < required: + args.rl_max_policy_age = required max_policy_age = validate_limit(getattr(args, "rl_max_policy_age", 0) or 0) policy_age_module = _rl_backend_module(args, "policy_age") policy_age_module.SUPPORT.check(max_policy_age) @@ -748,6 +781,7 @@ def _prepare_rl_args( raise ValueError("--rl-sao-contracts only applies to --rl-sync-preset sao-streaming-full") if args.rl_sync_preset != FULL_PARAMETER_SYNC_PRESET: _refuse_full_parameter_sync_options(args) + _check_bcast_delta(args) if args.rl_sync_preset == "decoupled": if args.experimental_rl_sync: raise ValueError("decoupled RL does not accept --experimental-rl-sync") diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index a5cd883f0..532fff20e 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -74,6 +74,10 @@ def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: f" --syncer-profile-sha256 {rl_syncer_profile(args, num_learners).sha256}") if getattr(args, "rl_sync_push_dtype", "f32") != "f32": sweep += f" --sweep-push-dtype {args.rl_sync_push_dtype}" + if getattr(args, "rl_sync_overlap", False): + sweep += " --sweep-overlap" + if getattr(args, "rl_sync_push_dtype", "f32") == "int8": + sweep += f" --sweep-push-int8-block {int(getattr(args, 'rl_sync_push_int8_block', 128))}" if getattr(args, "expert_full_count", 0): return lora + sweep return " --parameter-mode full" + sweep @@ -152,6 +156,7 @@ def make_miles_island_task( f" --rollout-num-gpus-per-engine {args.rollout_num_gpus_per_engine}" f" --sglang-mem-fraction-static {args.sglang_mem_fraction_static}" + _rl_trainable_island_flags(args, learner_id, num_learners) + + (" --bcast-delta" if int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) else "") + f" --inner-lr {args.inner_lr}" f" --seq-len {args.seq_len}" + (" --rl-activation-recompute" if getattr(args, "rl_activation_recompute", False) else "") diff --git a/yeto/launch/syncer.py b/yeto/launch/syncer.py index 91587c078..2836e1b01 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -150,6 +150,12 @@ def _syncer_checkpoint_every(args) -> int: return int(getattr(args, "rl_sync_checkpoint_every", 1)) if rl_dense_sweep(args) else 1 +def _syncer_bcast_delta(args) -> str: + """S26 E3 bf16 downlink deltas (checked by rl_args._check_bcast_delta).""" + full_every = int(getattr(args, "rl_sync_bcast_delta_full_every", 0) or 0) + return f" --bcast-delta-full-every {full_every}" if full_every else "" + + def rl_dense_sweep(args) -> bool: return (getattr(args, "training_mode", "sft") == "rl" and getattr(args, "rl_sync_preset", None) == FULL_PARAMETER_SYNC_PRESET) @@ -448,6 +454,7 @@ def _rl_role_syncer_command(args, num_learners: int, binary: str, f"{_syncer_island_scheduling(args)}" f"{_syncer_expected_contract(args)}" f"{_syncer_profile_binding(args)}" + f"{_syncer_bcast_delta(args)}" f" --outer-lr {args.outer_lr}" f" --outer-momentum {args.outer_momentum}" " --max-base-lag 0 --learner-weight equal" diff --git a/yeto/protocol.py b/yeto/protocol.py index 28331475c..fc2b6d73a 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -37,6 +37,9 @@ from collections.abc import Callable, Iterable, Iterator from dataclasses import dataclass, field +from .downlink import HEADER_LEN as DOWNLINK_HEADER_LEN +from .downlink import DownlinkAnchors, DownlinkChunkSink +from .downlink import tensor_nbytes as downlink_tensor_nbytes from .fragments import MERGE_ISO, FragmentLayout MAGIC = 0xD170C0DE @@ -71,6 +74,15 @@ # delta bytes; the client rounds them (round-to-nearest-even, as the syncer's # ``half::bf16::from_f32``) on the way out. DTYPE_F32_PUSH_BF16 = 4 +# Session dtype 5 (S26 E3): PUSH_FRAGMENT deltas travel as block-scaled int8 +# (``f32_bytes_to_i8``); INIT/BCAST/FINAL stay f32. Callers hand f32 bytes. +DTYPE_F32_PUSH_I8 = 5 +# Session-dtype flag (S26 E3): BCAST_FRAGMENT carries bf16 downlink deltas +# between f32 refreshes (yeto/downlink.py). Only on f32-bulk sessions. +DTYPE_FLAG_BCAST_DELTA = 0x80 +I8_BLOCK_MIN = 32 +I8_BLOCK_MAX = 1 << 20 +DEFAULT_I8_BLOCK = 128 CHUNK_SIZE = 4 * 1024 * 1024 MAX_HANDSHAKE_FRAME = 16 * 1024 * 1024 @@ -92,22 +104,103 @@ _BUDGET_DONE = struct.Struct(" int: + """The session dtype without its flags.""" + return dtype & ~DTYPE_FLAG_BCAST_DELTA + + +def bcast_delta(dtype: int) -> bool: + """Whether the session receives bf16 downlink deltas.""" + return bool(dtype & DTYPE_FLAG_BCAST_DELTA) + + +def valid_session_dtype(dtype: int) -> bool: + base = session_base(dtype) + if base in (DTYPE_F32, DTYPE_F32_PUSH_BF16, DTYPE_F32_PUSH_I8): + return 0 <= dtype <= 0xFF + return base in (DTYPE_BF16, DTYPE_Q4) and not bcast_delta(dtype) + + def bulk_dtype(dtype: int) -> int: """Dtype of INIT_PARAMS/BCAST_FRAGMENT tensors for a session dtype. Q4 applies only to push deltas (small dynamic range); full parameter - payloads would not survive 4 bits, so they travel as bf16. + payloads would not survive 4 bits, so they travel as bf16. A downlink + session's BCAST is reconstructed to f32 before callers see it. """ - if dtype == DTYPE_Q4: + base = session_base(dtype) + if base == DTYPE_Q4: return DTYPE_BF16 - if dtype == DTYPE_F32_PUSH_BF16: + if base in (DTYPE_F32_PUSH_BF16, DTYPE_F32_PUSH_I8): return DTYPE_F32 - return dtype + return base def push_wire_dtype(dtype: int) -> int: """Dtype of unquantized PUSH_FRAGMENT deltas on the wire for a session dtype.""" - return DTYPE_BF16 if dtype == DTYPE_F32_PUSH_BF16 else dtype + base = session_base(dtype) + return DTYPE_BF16 if base == DTYPE_F32_PUSH_BF16 else base + + +def i8_nbytes(numel: int, block: int) -> int: + """Encoded size of an int8 push: u32 block, then per block f32 scale + i8s.""" + if not I8_BLOCK_MIN <= block <= I8_BLOCK_MAX: + raise ValueError(f"int8 block size {block} outside {I8_BLOCK_MIN}..{I8_BLOCK_MAX}") + return 4 + 4 * -(-numel // block) + numel + + +def f32_bytes_to_i8(data: bytes | bytearray | memoryview, block: int, *, header: bool = True) -> bytes: + """Quantize little-endian f32 values to block-scaled int8 (syncer ``decode_i8``). + + Per block of ``block`` values (the last may be shorter): f32 scale = + absmax / 127, then ``round_half_even(x / scale)`` clipped to [-127, 127]. + Decoding ``q * scale`` errs by at most absmax / 254 per element (plus f32 + rounding). Blocks must start at multiples of ``block`` values; with + ``header`` the u32 block size is prepended. + """ + import numpy as np + + i8_nbytes(0, block) # validates the block size + values = np.frombuffer(data, dtype=" None: # chunk: (blocks, width) + absmax = np.abs(chunk).max(axis=1) + scale = (absmax / np.float32(127)).astype(np.float32) + safe = np.where(scale > 0, scale, np.float32(1)) + q = np.clip(np.rint(chunk / safe[:, None]), -127, 127).astype(np.int8) + q[scale == 0] = 0 + rows = np.empty((chunk.shape[0], 4 + chunk.shape[1]), dtype=np.uint8) + rows[:, :4] = scale.astype(" full: + encode(values[full:].reshape(1, -1)) + return b"".join(parts) + + +def _f32_parts_to_i8(parts: Iterable[bytes | bytearray | memoryview], block: int) -> Iterator[bytes]: + """Stream ``f32_bytes_to_i8`` over parts; blocks may span part boundaries.""" + yield struct.pack(" bytes: @@ -206,8 +299,8 @@ def seal_hello(body: bytes, key: bytes | None) -> bytes: The legacy syncer with a key refuses a HELLO without this trailer.""" if not key: return body - import hmac import hashlib + import hmac return body + hmac.new(key, HELLO_MAC_DOMAIN + body, hashlib.sha256).digest() @@ -515,11 +608,13 @@ def __init__( syncer_profile_hash: bytes | None = None, hmac_key: bytes | None = None, compat_group: str | None = None, + push_i8_block: int = DEFAULT_I8_BLOCK, ): if not 0 <= learner_id <= 0xFFFF_FFFF: raise ValueError(f"learner_id must fit u32, got {learner_id}") - if dtype not in (DTYPE_F32, DTYPE_BF16, DTYPE_Q4, DTYPE_F32_PUSH_BF16): + if not valid_session_dtype(dtype): raise ValueError(f"unsupported session dtype {dtype}") + i8_nbytes(0, push_i8_block) # validates the int8 block size if layout.num_fragments < 1: raise ValueError("layout must contain at least one fragment") if not 0 <= num_streams <= 256: @@ -537,6 +632,11 @@ def __init__( self.learner_id = learner_id self.layout = layout self.dtype = dtype + self._push_base = session_base(dtype) + self.push_i8_block = push_i8_block + # S26 E3: f32 anchors of a bf16-downlink session (yeto/downlink.py). + self._downlink = DownlinkAnchors(layout.num_fragments) if bcast_delta(dtype) else None + self._bcast_prefix = 12 + (DOWNLINK_HEADER_LEN if self._downlink is not None else 0) self.num_streams = num_streams self.session_contract_hash = session_contract_hash self.syncer_profile_hash = syncer_profile_hash @@ -548,7 +648,7 @@ def __init__( self.max_reconnects = max_reconnects self.finalization_timeout = finalization_timeout self._max_bcast_payload = max( - 12 + _tensor_nbytes(bulk_dtype(dtype), fragment.numel) + self._bcast_prefix + _tensor_nbytes(bulk_dtype(dtype), fragment.numel) for fragment in layout.fragments ) self._max_final_payload = max( @@ -870,9 +970,9 @@ def push_fragment( numel = self.layout.fragments[fragment_id].numel expected = ( _q4_nbytes(numel) - if self.dtype == DTYPE_Q4 - # dtype 4 callers pass f32 bytes; the client rounds them below - else _tensor_nbytes(DTYPE_F32 if self.dtype == DTYPE_F32_PUSH_BF16 else self.dtype, numel) + if self._push_base == DTYPE_Q4 + # dtypes 4/5 callers pass f32 bytes; the client encodes them below + else _tensor_nbytes(bulk_dtype(self.dtype), numel) ) if len(tensor_bytes) != expected: raise ValueError( @@ -890,8 +990,10 @@ def push_fragment( c_steps, c_tokens, ) - if self.dtype == DTYPE_F32_PUSH_BF16: + if self._push_base == DTYPE_F32_PUSH_BF16: tensor_bytes = f32_bytes_to_bf16(tensor_bytes) + elif self._push_base == DTYPE_F32_PUSH_I8: + tensor_bytes = f32_bytes_to_i8(tensor_bytes, self.push_i8_block) self._send_large(MSG_PUSH_FRAGMENT, head + tensor_bytes) def push_fragment_parts( @@ -923,12 +1025,14 @@ def push_fragment_parts( if c_steps < 1: raise ValueError("c_steps must be positive") numel = self.layout.fragments[fragment_id].numel - expected = ( - _q4_nbytes(numel) - if self.dtype == DTYPE_Q4 - else _tensor_nbytes(push_wire_dtype(self.dtype), numel) - ) - if self.dtype == DTYPE_F32_PUSH_BF16: + if self._push_base == DTYPE_Q4: + expected = _q4_nbytes(numel) + elif self._push_base == DTYPE_F32_PUSH_I8: + expected = i8_nbytes(numel, self.push_i8_block) + tensor_parts = _f32_parts_to_i8(tensor_parts, self.push_i8_block) + else: + expected = _tensor_nbytes(push_wire_dtype(self.dtype), numel) + if self._push_base == DTYPE_F32_PUSH_BF16: tensor_parts = _f32_parts_to_bf16(tensor_parts) head = struct.pack( " None: raise TypeError( "inbound chunk sink does not implement the transaction API" ) + if sink is not None and self._downlink is not None: + sink = DownlinkChunkSink( + sink, self._downlink, [f.numel for f in self.layout.fragments], MSG_BCAST_FRAGMENT) with self._reasm_lock: if self._reasm: raise RuntimeError("cannot replace inbound chunk sink mid-message") @@ -1542,6 +1649,13 @@ def _stream_chunk_locked( dtype, self.layout.fragments[fragment_id].numel, ) + if msg_type == MSG_BCAST_FRAGMENT and self._downlink is not None: + # The downlink header (kind, base version, digest) is part of + # the payload the wrapping sink decodes; its kind fixes the size. + if len(data) < tensor_start + DOWNLINK_HEADER_LEN: + raise ValueError("first downlink fragment chunk is header-truncated") + expected_payload = DOWNLINK_HEADER_LEN + downlink_tensor_nbytes( + data[tensor_start], self.layout.fragments[fragment_id].numel) if length != 12 + expected_payload: raise ValueError( f"streamed fragment {fragment_id} has {length - 12} tensor " @@ -1810,15 +1924,19 @@ def _dispatch_live( fid, version = struct.unpack_from("= self.layout.num_fragments: raise ValueError(f"BCAST_FRAGMENT for unknown fragment {fid}") - expected = _tensor_nbytes( - bulk_dtype(self.dtype), self.layout.fragments[fid].numel - ) - if len(payload) - 12 != expected: - raise ValueError( - f"BCAST_FRAGMENT {fid} has {len(payload) - 12} tensor bytes, " - f"expected {expected}" + if self._downlink is not None: + data = self._downlink.decode_whole( + fid, version, payload[12:], self.layout.fragments[fid].numel) + else: + expected = _tensor_nbytes( + bulk_dtype(self.dtype), self.layout.fragments[fid].numel ) - data = payload[12:] + if len(payload) - 12 != expected: + raise ValueError( + f"BCAST_FRAGMENT {fid} has {len(payload) - 12} tensor bytes, " + f"expected {expected}" + ) + data = payload[12:] digest = hashlib.sha256(data).digest() with self._bcast_lock: seen = self._bcast_seen[fid] diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 7ea7d4758..4b2a5199f 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -20,19 +20,28 @@ from typing import Any from yeto.protocol import ( + DEFAULT_I8_BLOCK, DTYPE_F32, DTYPE_F32_PUSH_BF16, + DTYPE_F32_PUSH_I8, + DTYPE_FLAG_BCAST_DELTA, + I8_BLOCK_MAX, + I8_BLOCK_MIN, FinalManifest, PullRequest, SyncerClient, ) +from yeto.rl.adapters.miles.full_parameter.inbound_sink import MilesRayInboundChunkSink from yeto.rl.adapters.miles.models.chunked_full_parameter import ( MilesChunkedFullParameterAdapter, ReferencedPolicyCut, StoredAuthoritativeFragment, ) -from yeto.rl.adapters.miles.full_parameter.inbound_sink import MilesRayInboundChunkSink -from yeto.rl.engine.policy_digest import PolicyDigest, TrainerResidentState, WeightsChanged +from yeto.rl.engine.policy_digest import ( + PolicyDigest, + TrainerResidentState, + WeightsChanged, +) from yeto.rl.engine.trainable_state import LAYOUT_FRAGMENTS, UnsupportedLayoutError from yeto.rl.local_learner import ComponentIdentity, layout_stream_role from yeto.rl.trajectory_evidence import TrajectoryBatchEvidence @@ -83,8 +92,27 @@ class RoleStreamConfig: poll_seconds: float = 0.01 max_fragment_bytes: int = 2 << 30 max_chunk_bytes: int = 256 << 20 - # S26 E3: push outer deltas as bf16 (session dtype 4); BCAST/FINAL stay f32. - push_bf16: bool = False + # S26 E3 wire options (default exact f32): push outer deltas as bf16 + # (session dtype 4) or block-scaled int8 (dtype 5); receive BCAST as bf16 + # deltas against a shared f32 anchor. FINAL always stays f32. + push_dtype: str = "f32" + push_int8_block: int = DEFAULT_I8_BLOCK + bcast_delta: bool = False + # S26 E3: train round k+1 while sweep k syncs, then rebase the trainer + # (see FullParameterSync._overlap_boundary); one more outer step stale. + overlap: bool = False + + @property + def session_dtype(self) -> int: + base = {"f32": DTYPE_F32, "bf16": DTYPE_F32_PUSH_BF16, "int8": DTYPE_F32_PUSH_I8}[self.push_dtype] + return base | (DTYPE_FLAG_BCAST_DELTA if self.bcast_delta else 0) + + @property + def wire_dtype_label(self) -> str: + """Session-contract label; plain f32 keeps the historical ``fp32``.""" + label = {"f32": "fp32", "bf16": "fp32-push-bf16", + "int8": f"fp32-push-int8-b{self.push_int8_block}"}[self.push_dtype] + return label + ("+bcast-bf16-delta" if self.bcast_delta else "") def __post_init__(self) -> None: if not self.algorithm.isidentifier(): @@ -129,6 +157,10 @@ def __post_init__(self) -> None: value = getattr(self, name) if type(value) is not int or value < 1: raise ValueError(f"role-stream {name} must be positive") + if self.push_dtype not in ("f32", "bf16", "int8"): + raise ValueError("role-stream push dtype must be f32, bf16 or int8") + if type(self.push_int8_block) is not int or not I8_BLOCK_MIN <= self.push_int8_block <= I8_BLOCK_MAX: + raise ValueError("role-stream int8 block size is invalid") if self.wait_timeout <= 0 or not 0 < self.poll_seconds <= 1: raise ValueError("role-stream wait bounds are invalid") if ( @@ -222,7 +254,8 @@ def role_stream_session_contract_hash( "pipeline_depth": config.pipeline_depth, "local_horizon": config.local_horizon, "optimizer_steps_per_round": config.optimizer_steps_per_round, - "wire_dtype": "fp32-push-bf16" if config.push_bf16 else "fp32", + "wire_dtype": config.wire_dtype_label, + **({"sync_train_overlap": True} if config.overlap else {}), }, } encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() @@ -264,6 +297,11 @@ def __init__(self, config: RoleStreamConfig, *, ray_module=None) -> None: self._staged: dict[tuple[int, int], StoredAuthoritativeFragment] = {} self._pinned_cuts: dict[int, tuple[ReferencedPolicyCut, int]] = {} self._pinned_anchors: dict[int, tuple[StoredAuthoritativeFragment, int]] = {} + # Overlap: the cut the current local round started from, and the + # round/token counters at that point (the next push's c_steps base). + self._round_base: ReferencedPolicyCut | None = None + self._rounds_at_base = 0 + self._tokens_at_base = 0 self.terminal_manifest: FinalManifest | None = None self.terminal_submitted = False self.final_applied = False @@ -337,7 +375,8 @@ async def initialize(self, group) -> None: self.config.syncer_addr, self.config.learner_id, adapter.layout.fragments, - dtype=DTYPE_F32_PUSH_BF16 if self.config.push_bf16 else DTYPE_F32, + dtype=self.config.session_dtype, + push_i8_block=self.config.push_int8_block, num_streams=self.config.wan_streams, max_reconnects=None, session_contract_hash=role_stream_session_contract_hash( @@ -474,7 +513,14 @@ async def complete_local_round( *, rollout_id: int, trajectories: TrajectoryBatchEvidence, + hold_submission: bool = False, ) -> RoleRoundResult: + """Record one trained round and capture its cut. + + ``hold_submission`` (overlap): the previous sweep is still in flight, + so the new cut is not yet submittable; it becomes the submission only + after :meth:`apply_staged_broadcast` rebased it and + :meth:`finish_rebase`. The round base stays pinned until then.""" if ( self.adapter is None or self.current is None @@ -503,14 +549,19 @@ async def complete_local_round( previous = self.current previous_submission = self.submission_cut self.current = local - self.submission_cut = local - self.submission_fragment_versions = self.fragment_versions self.local_generation = expected_generation self.local_rounds += 1 self.action_tokens += trajectories.trained_tokens - self.submission_rounds = self.local_rounds - self.submission_tokens = self.action_tokens - self._release_cuts(previous, previous_submission, keep=(local,)) + if hold_submission: + if not self.config.overlap or self._round_base is None: + raise RuntimeError(f"{self.config.role} holds a submission without an overlap round base") + self._release_cuts(previous, keep=(local, self._round_base, previous_submission)) + else: + self.submission_cut = local + self.submission_fragment_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + self.submission_tokens = self.action_tokens + self._release_cuts(previous, previous_submission, keep=(local, self._round_base)) # The policy-level coordinator owns permit matching and release. A # role must never answer a PULL here: the peer role may not yet have @@ -562,8 +613,15 @@ async def apply_staged_broadcast( self, fragment_id: int, version: int, + *, + rebase: bool = False, ) -> tuple[int, ...]: - """Apply one coordinator-approved staged fragment at a safe boundary.""" + """Apply one coordinator-approved staged fragment at a safe boundary. + + ``rebase`` (overlap): the trainer keeps its progress since the round + base, ``theta_f <- G_f + (current_f - base_f)``; ``G_f`` becomes the + anchor, and the counters at the round base become the anchor's, so + the next push carries exactly that progress.""" key = (fragment_id, version) stored = self._staged.get(key) @@ -576,14 +634,56 @@ async def apply_staged_broadcast( raise RuntimeError( f"{self.config.algorithm} {self.config.role} staged fragment skipped its base version" ) - await self._apply_authoritative({fragment_id: stored}) + if not rebase: + await self._apply_authoritative({fragment_id: stored}) + else: + if not self.config.overlap or self._round_base is None or self.current is None: + raise RuntimeError(f"{self.config.role} rebase without an overlap round base") + rebased = self.adapter.rebase_fragment( + stored, self.current, self._round_base, ray_module=self.ray_module) + try: + await self._apply_authoritative( + {fragment_id: rebased}, + anchors={fragment_id: stored}, + counters=(self._rounds_at_base, self._tokens_at_base), + ) + finally: + rebased.release() # the applied target holds its own references self._staged.pop(key, None) return (fragment_id,) + def finish_rebase(self) -> None: + """Overlap: the rebased cut is the next submission (``G + progress``).""" + + if not self.config.overlap or self.current is None: + raise RuntimeError(f"{self.config.role} finish_rebase outside overlap") + previous_submission = self.submission_cut + self.submission_cut = self.current + self.submission_fragment_versions = self.fragment_versions + self.submission_rounds = self.local_rounds + self.submission_tokens = self.action_tokens + self._release_cuts(previous_submission, keep=(self.current, self._round_base)) + + def mark_round_base(self) -> None: + """Overlap: pin the cut the next local round starts from.""" + + if not self.config.overlap or self.current is None: + raise RuntimeError(f"{self.config.role} mark_round_base outside overlap") + previous = self._round_base + self._round_base = self.current + self._rounds_at_base = self.local_rounds + self._tokens_at_base = self.action_tokens + self._release_cuts(previous, keep=(self.current, self.submission_cut)) + async def _apply_authoritative( self, selected: Mapping[int, StoredAuthoritativeFragment], + *, + anchors: Mapping[int, StoredAuthoritativeFragment] | None = None, + counters: tuple[int, int] | None = None, ) -> None: + """Apply ``selected`` fragments; ``anchors``/``counters`` override the + recorded anchor and its round/token counters (overlap rebase).""" if self.adapter is None or self.current is None or self.group is None: raise RuntimeError(f"{self.config.algorithm} {self.config.role} stream is unavailable") target_version = self.transport_version + 1 @@ -613,16 +713,18 @@ async def _apply_authoritative( self.local_generation = 0 self._release_cuts( previous_cut, - keep=(self.current, self.submission_cut), + keep=(self.current, self.submission_cut, self._round_base), ) - for fragment_id, stored in selected.items(): + rounds, tokens = counters if counters is not None else (self.local_rounds, self.action_tokens) + for fragment_id, applied in selected.items(): + stored = applied if anchors is None else anchors[fragment_id] previous_anchor = self._anchors[fragment_id] if previous_anchor is not None and previous_anchor is not stored: self._release_anchor(previous_anchor) self._anchors[fragment_id] = stored self._fragment_versions[fragment_id] = stored.version - self._rounds_at_anchor[fragment_id] = self.local_rounds - self._tokens_at_anchor[fragment_id] = self.action_tokens + self._rounds_at_anchor[fragment_id] = rounds + self._tokens_at_anchor[fragment_id] = tokens def _apply_token( self, @@ -914,8 +1016,9 @@ def close(self) -> None: retained_cuts = tuple(value for value, _count in self._pinned_cuts.values()) self._pinned_cuts.clear() self._pinned_anchors.clear() - self._release_cuts(self.current, self.submission_cut, *retained_cuts) + self._release_cuts(self.current, self.submission_cut, self._round_base, *retained_cuts) self._staged.clear() + self._round_base = None self.current = None self.submission_cut = None self.submission_fragment_versions = None @@ -937,6 +1040,7 @@ def _unpin_cut(self, cut: ReferencedPolicyCut) -> None: remaining == 0 and cut is not self.current and cut is not self.submission_cut + and cut is not self._round_base and self.adapter is not None ): self.adapter.release(cut) diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 5a9d8874d..c382f8cb4 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -814,8 +814,14 @@ def _dense_full_sync_config( from yeto.rl.adapters.miles.full_parameter.role_stream import RoleStreamConfig from yeto.rl.adapters.miles.full_parameter.session import FullParameterSyncConfig + from yeto.rl.engine.sync_overlap import required_policy_age from yeto.rl.local_learner import ComponentIdentity + overlap = bool(getattr(args, "sweep_overlap", False)) + if overlap and int(getattr(args, "rl_max_policy_age", 0) or 0) < required_policy_age(0): + raise ValueError("--sweep-overlap rollouts are one outer step stale: it needs " + f"--rl-max-policy-age >= {required_policy_age(0)}") + actor = RoleStreamConfig( algorithm="grpo", role="actor", @@ -838,7 +844,10 @@ def _dense_full_sync_config( pipeline_depth=1, algorithm_spec_sha256=algorithm_spec_sha256, wan_streams=args.wan_streams, - push_bf16=getattr(args, "sweep_push_dtype", "f32") == "bf16", + push_dtype=getattr(args, "sweep_push_dtype", "f32"), + push_int8_block=int(getattr(args, "sweep_push_int8_block", 128)), + bcast_delta=bool(getattr(args, "bcast_delta", False)), + overlap=overlap, ) return FullParameterSyncConfig((actor,), sweep=True) @@ -2099,6 +2108,7 @@ def run_miles( algorithm_semantic_sha256=algorithm_semantic_sha256, code_pack_digest=code_pack_digest, compat_group=_island_compat_group(args), + bcast_delta=bool(getattr(args, "bcast_delta", False)), ) if getattr(args, "rl_island_scheduling", "legacy") == "elastic": # rl-inter-island-scheduling 0.15 (legacy sets nothing) diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 9b7e6c996..f16a6ec56 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -786,6 +786,60 @@ def assemble_mixed_target( ) return self.validate_cut(target, target_policy_version, 0) + def rebase_fragment( + self, + stored: StoredAuthoritativeFragment, + current: ReferencedPolicyCut, + base: ReferencedPolicyCut, + *, + ray_module=None, + ) -> StoredAuthoritativeFragment: + """Sync/train overlap: ``stored + (current - base)`` for one fragment. + + The result is stored like an inbound authoritative fragment (same + id, version and provenance, fresh Ray refs) so ``assemble_mixed_target`` + validates and applies it unchanged. The caller keeps ``stored`` as + the anchor and releases the result after the apply.""" + + import numpy as np + from miles.backends.megatron_utils.full_parameter_state import ( + _resolve_ray_module, + ) + + from yeto.rl.engine.sync_overlap import rebase_values + + fragment_id = stored.fragment_id + for cut in (current, base): + if cut.layout_hash != self.layout.layout_hash: + raise ValueError("Miles rebase cut layout changed") + owner = _fragment_owner_shard(current.transport_cut, fragment_id) + if ( + stored.parameter_layout_hash != self.layout.layout_hash + or stored.topology != owner.topology + or stored.plan_hash != owner.plan_hash + ): + raise ValueError("authoritative fragment provenance changed") + ray = _resolve_ray_module(ray_module) + chunks = sorted(zip(stored.descriptor.chunks, stored.refs, strict=True), + key=lambda row: row[0].flat_offset) + anchor_parts, offset = [], 0 + for descriptor, reference in chunks: + if reference is None or descriptor.flat_offset != offset: + raise ValueError("authoritative fragment chunks are released or not contiguous") + value = ray.get(reference) + anchor_parts.append(value.detach().to("cpu").reshape(-1).numpy()) + offset += descriptor.numel + + def flat(cut: ReferencedPolicyCut) -> np.ndarray: + raw = b"".join(bytes(part) for part in self.fragment_parts( + cut, fragment_id, ray_module=ray_module)()) + return np.frombuffer(raw, dtype="= 1)", + ), + island_flag( + "--bcast-delta", + action="store_true", + help="strict-avg / full-parameter sweep: accept bf16 downlink deltas against the " + "shared f32 anchor (the syncer runs --bcast-delta-full-every)", + ), island_flag( "--syncer-profile-sha256", default=None, diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index ab6f640f8..5fec3a02a 100644 --- a/yeto/rl/adapters/verl/island_entry.py +++ b/yeto/rl/adapters/verl/island_entry.py @@ -60,6 +60,8 @@ def parse_args(argv): p.add_argument("--seq-len", type=int, required=True) p.add_argument("--seed", type=int, required=True) p.add_argument("--wan-streams", type=int, default=4) + # S26 E3: bf16 downlink deltas (yeto/downlink.py; the syncer runs --bcast-delta-full-every) + p.add_argument("--bcast-delta", action="store_true") p.add_argument("--apply-chat-template-kwargs", default=None) p.add_argument("--rl-algorithm-spec", default=None) p.add_argument("--rl-island-scheduling", default="legacy") @@ -357,6 +359,7 @@ def main(argv=None) -> int: "model_path": model_path, "model_revision": args.model_revision, "lora_rank": args.lora_r, "lora_targets": args.lora_targets, "event_tape": os.path.expanduser(args.event_tape), "out_dir": str(OUTPUT), "tis_upper": algo["tis_upper"], "wan_streams": args.wan_streams, + "bcast_delta": bool(args.bcast_delta), "thresholds_key": ["verl", "fsdp2", "vllm-" + str(manifest["versions"].get("vllm")), _threshold_card(manifest)], "algorithm_spec": spec_dict, "syncer_epoch": args.rl_syncer_epoch, "island_lease_s": args.rl_island_lease_s, diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 6be18c9e2..08c72e83c 100644 --- a/yeto/rl/adapters/verl/trainer.py +++ b/yeto/rl/adapters/verl/trainer.py @@ -186,7 +186,8 @@ def build_sync(plan: dict, specs, *, lora_hash: str, layout_hash: str, identity_ expected_specs=tuple(specs), base_model_revision=plan["model_revision"], lora_config_hash=lora_hash, layout_hash=layout_hash, event_tape=plan["event_tape"], wan_streams=int(plan.get("wan_streams", 4)), - backend_identity_sha256=identity_sha256, compat_group=compat_group) + backend_identity_sha256=identity_sha256, compat_group=compat_group, + bcast_delta=bool(plan.get("bcast_delta", False))) if mode == "strict": sync = StrictAvgSync(bridge, progress=None) elif mode == "elastic": diff --git a/yeto/rl/bridge.py b/yeto/rl/bridge.py index 9cf0e3889..7cae92564 100644 --- a/yeto/rl/bridge.py +++ b/yeto/rl/bridge.py @@ -13,8 +13,13 @@ import torch -from .engine.backend_identity import config_island_contract_sha256 -from ..protocol import DTYPE_F32, FinalManifest, PullRequest, SyncerClient +from ..protocol import ( + DTYPE_F32, + DTYPE_FLAG_BCAST_DELTA, + FinalManifest, + PullRequest, + SyncerClient, +) from ..tensor_io import pack_tensor, unpack_fragment from .core import ( CanonicalLoraState, @@ -29,6 +34,7 @@ policy_delta, tensors_from_flat_owned, ) +from .engine.backend_identity import config_island_contract_sha256 class IslandRuntime(Protocol): @@ -77,6 +83,12 @@ class BridgeConfig: # 7.7c: "-" (yeto.hw.catalog); sent in HELLO / JOIN so the # syncer refuses another card type by name. None = not declared. compat_group: str | None = None + # S26 E3: receive BCAST as bf16 deltas against the shared f32 anchor + # (session flag; the syncer runs --bcast-delta-full-every). The client + # hands this bridge the reconstructed f32 anchor, so nothing else changes. + # An anchor mismatch drops the connection; with max_reconnects=0 the + # island exits and its restart re-reads the committed cut as f32. + bcast_delta: bool = False def _write_round_audit( @@ -278,7 +290,7 @@ def __init__(self, runtime: IslandRuntime, config: BridgeConfig) -> None: config.syncer_addr, config.learner_id, self.layout, - dtype=DTYPE_F32, + dtype=DTYPE_F32 | (DTYPE_FLAG_BCAST_DELTA if config.bcast_delta else 0), num_streams=config.wan_streams, session_contract_hash=_session_contract( self.layout, config_island_contract_sha256(config)), diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index f4b2141d2..7c9d0921f 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -81,6 +81,11 @@ async def stage_available_broadcasts(self) -> tuple[tuple[int, int], ...]: ... def staged_fragment_versions(self) -> tuple[tuple[int, int], ...]: ... def has_staged_broadcast(self, fragment_id: int, version: int) -> bool: ... async def apply_staged_broadcast(self, fragment_id: int, version: int) -> tuple[int, ...]: ... + # S26 E3 sync/train overlap (profile.overlap) only: + # complete_local_round(..., hold_submission=True) keeps the round base and + # the previous submission; apply_staged_broadcast(..., rebase=True) applies + # G + (current - base) with G as the anchor; finish_rebase() makes the + # rebased cut the submission; mark_round_base() pins the next round base. async def wait_and_apply_final(self) -> None: ... def acknowledge_final(self) -> None: ... def discard_detached_submission_cut(self) -> None: ... @@ -99,6 +104,8 @@ class FullParameterSyncProfile: wait_timeout: float poll_seconds: float sweep: bool = False + # S26 E3: train round k+1 while sweep k syncs; rebase at the next boundary. + overlap: bool = False def __post_init__(self) -> None: if not self.roles or len(set(self.roles)) != len(self.roles): @@ -107,6 +114,13 @@ def __post_init__(self) -> None: raise ValueError("the actor is the published role and must come first") if self.sweep and (self.roles != ("actor",) or self.local_horizon != 1): raise ValueError("a dense strict sweep has one actor role and H=1") + if self.overlap and not self.sweep: + raise ValueError("sync/train overlap needs the dense strict sweep") + + @property + def total_rounds(self) -> int: + """Local rounds of a sweep profile (one sweep per round).""" + return self.total_fragment_steps // self.fragments @property def effective_pipeline_depth(self) -> int: @@ -131,6 +145,7 @@ def from_configs( wait_timeout=min(config.wait_timeout for config in configs.values()), poll_seconds=min(config.poll_seconds for config in configs.values()), sweep=sweep, + overlap=bool(getattr(first, "overlap", False)), ) @@ -389,9 +404,12 @@ async def _apply_pair( step, fragment_id, _base = key if not pair.complete: raise RuntimeError("broadcast arrived before every paired PUSH was retained") - results = await asyncio.gather( - *(stream.apply_staged_broadcast(fragment_id, step) for stream in self.streams.values()) - ) + if self.profile.overlap: + calls = (stream.apply_staged_broadcast(fragment_id, step, rebase=True) + for stream in self.streams.values()) + else: + calls = (stream.apply_staged_broadcast(fragment_id, step) for stream in self.streams.values()) + results = await asyncio.gather(*calls) for role, fragments in zip(self.streams, results, strict=True): applied[role].extend(fragments) if self.streams[role].fragment_versions[fragment_id] != step: @@ -509,6 +527,8 @@ def __init__( self.closed = False self._driver: IslandDriver | None = None self._timing = SyncTiming(None, None) # L5.6, set by each boundary + # Overlap: fragment versions before the sweep still in flight. + self._pending_sweep: tuple[int, ...] | None = None @property def actor(self) -> RoleStream: @@ -539,6 +559,8 @@ def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> Sync return SyncBoundary(state=state, stop=self.terminal) async def _boundary(self, rollout_id: int) -> bool: + if self.pairing.profile.overlap: + return await self._overlap_boundary(rollout_id) pairing = self.pairing trajectories = self.trajectories(rollout_id) await pairing.gather_roles( @@ -564,6 +586,53 @@ async def _boundary(self, rollout_id: int) -> bool: stream.discard_detached_submission_cut() return False + async def _overlap_boundary(self, rollout_id: int) -> bool: + """Streaming-DiLoCo overlap (S26 E3; one more outer step of staleness). + + Round k's delta (``current - base``) is pushed at boundary k and the + island trains round k+1 at once. At boundary k+1 the island first + waits for sweep k (``G_k``), then rebases its trainer in place: + ``theta <- G_k + (theta_local_{k+1} - base_k)`` with ``G_k`` as the + anchor, so the next push is exactly round k+1's progress and round + k's delta is never counted twice. The terminal round waits for the + final cut like the strict sweep.""" + + pairing = self.pairing + trajectories = self.trajectories(rollout_id) + hold = self._pending_sweep is not None + await pairing.gather_roles( + lambda _role, stream: stream.complete_local_round( + rollout_id=rollout_id, trajectories=trajectories, hold_submission=hold + ) + ) + wait_ts = unix_now() + if hold: + await self._wait_for_sweep(self._pending_sweep) + self._pending_sweep = None + if not pairing.any_finalizing: + for stream in pairing.streams.values(): + stream.finish_rebase() + rebased_ts = unix_now() + if pairing.any_finalizing: + await pairing.finish_without_more_training() + return True + before = pairing.fragment_versions() + pairing.submit_ready() + submit_ts = unix_now() + if rollout_id + 1 >= pairing.profile.total_rounds: + # Terminal round: nothing trains after it, so wait for the cut. + await pairing.commit_staged_broadcasts() + await self._wait_for_sweep(before) + self._timing = SyncTiming(submit_ts, unix_now()) + await pairing.finish_without_more_training() + return True + self._pending_sweep = before + for stream in pairing.streams.values(): + stream.mark_round_base() + # This boundary blocked only on the previous sweep (wait + rebase). + self._timing = SyncTiming(wait_ts, rebased_ts if hold else None) + return False + async def _wait_for_sweep(self, before: tuple[int, ...]) -> None: """Strict sweep: block until all P fragments of this local step are applied, i.e. every fragment version moved past ``before``. An empty diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index 6e6b95c1a..a69e67fdc 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -851,11 +851,37 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. ), cli_flag( "--rl-sync-push-dtype", - choices=("f32", "bf16"), + choices=("f32", "bf16", "int8"), default="f32", help="full-parameter only: wire dtype of the islands' pushed outer deltas " - "(pseudo-gradients). bf16 halves upload bytes; the syncer keeps f32 masters, " - "momentum and f32 BCAST/FINAL (default f32 = exact)", + "(pseudo-gradients). bf16 halves upload bytes, int8 (block absmax scale, see " + "--rl-sync-push-int8-block) quarters them; the syncer keeps f32 masters, " + "momentum and f32 FINAL (default f32 = exact)", + ), + cli_flag( + "--rl-sync-overlap", + action="store_true", + help="full-parameter only (E2 option): train the next round while the " + "sweep syncs, then rebase the trainer onto the new global weights " + "(theta <- G_k + local progress). Rollouts are one outer step staler: " + "--rl-max-policy-age is raised to at least 1 and the algorithm spec must " + "tolerate it (checked before launch)", + ), + cli_flag( + "--rl-sync-push-int8-block", + type=int, + default=128, + help="full-parameter only, with --rl-sync-push-dtype int8: values per f32 " + "scale (32..1048576; error <= block absmax / 254 per value)", + ), + cli_flag( + "--rl-sync-bcast-delta-full-every", + type=int, + default=0, + help="full-parameter only: above 0, the syncer sends each new fragment " + "version as a bf16 delta against the islands' shared f32 anchor, with an " + "exact f32 refresh every N-th version (N >= 2; digest-checked, a mismatch " + "reconnects for f32). 0 = always f32", ), cli_flag( "--rl-sync-checkpoint-every", diff --git a/yeto/rl/engine/sync_overlap.py b/yeto/rl/engine/sync_overlap.py new file mode 100644 index 000000000..4c47f344d --- /dev/null +++ b/yeto/rl/engine/sync_overlap.py @@ -0,0 +1,41 @@ +"""Sync/train overlap for the full-parameter sweep (S26 E3 item 3, E2 option). + +Round k's outer delta is pushed at boundary k and the island trains round +k+1 immediately on its own weights. At boundary k+1 it waits for the global +result of sweep k (``G_k``) and rebases in place:: + + theta <- G_k + (theta_local_{k+1} - base_k) + +``base_k`` is the cut round k+1 started from. ``G_k`` becomes the anchor, so +the next push ``theta - G_k`` equals round k+1's progress exactly once: round +k's delta is already inside ``G_k`` and is not counted again. + +Rollouts of round k+1 use weights one outer step older than the global +policy, so the run's minimum ``--rl-max-policy-age`` grows by +:data:`EXTRA_POLICY_AGE`. +""" + +from __future__ import annotations + +import numpy as np + +EXTRA_POLICY_AGE = 1 + + +def required_policy_age(minimum: int) -> int: + """The smallest accepted ``--rl-max-policy-age`` once overlap is on.""" + return int(minimum) + EXTRA_POLICY_AGE + + +def rebase_values(anchor: np.ndarray, current: np.ndarray, base: np.ndarray) -> np.ndarray: + """``anchor + (current - base)`` in f32 (the trainer keeps its progress).""" + if not (anchor.shape == current.shape == base.shape): + raise ValueError("rebase operands have different shapes") + for name, value in (("anchor", anchor), ("current", current), ("base", base)): + if value.dtype != np.float32: + raise ValueError(f"rebase {name} must be float32") + out = np.subtract(current, base, dtype=np.float32) + np.add(anchor, out, out=out) + if not np.isfinite(out).all(): + raise ValueError("rebased fragment has a non-finite value") + return out From 8af78faf152180d156bc6cf6d996d9339e253e99 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:04:45 +0000 Subject: [PATCH 4/9] overlap: keep rollout policy age 0, relabel rebased cuts, push whole sweep Pre-GPU check of --rl-sync-overlap on the Miles role stream found three problems: - The launcher raised --rl-max-policy-age to 1, but the dense sweep island refuses any age other than 0, so an overlap run could never start; and an age >= 1 also turns on Miles' fully_async rollout (partial rollout when colocated). Overlap is not rollout staleness: the driver publishes the boundary's cut before the next generation, so rollouts stay on-policy and only the local weights lag the global policy by one outer step. Keep the age at 0 (refuse others at launch) and bind the lag into the session contract only. - After a rebase the applied cut has local step generation 0, and Miles pushes only cuts with progress, so the second sweep's push failed. The submission is now the rebased cut relabeled with the rounds trained since the round base (same values, hashes and refs). - The boundary submitted once; a PULL that arrived later waited a whole local round. The overlap boundary now submits the whole sweep before training on, and the launcher requires --rl-sync-fragment-pipeline equal to --fragments so the syncer sends all PULLs at once. Adds Miles role-stream tests (real adapter and role stream, CPU tensors, fake trainer/Ray/syncer): exact-once rounds through rebase_fragment, and syncer restarts while the overlapped round trains (before commit, after commit with a lost broadcast, partly committed sweep). Co-Authored-By: Claude Opus 5.5 --- tests/test_e3_overlap_miles_role_stream.py | 495 ++++++++++++++++++ tests/test_e3_sync_overlap.py | 81 ++- yeto/launch/rl_args.py | 15 +- .../miles/full_parameter/role_stream.py | 15 +- yeto/rl/adapters/miles/island_entry.py | 10 +- .../miles/models/chunked_full_parameter.py | 36 ++ yeto/rl/adapters/miles/run_flags.py | 2 +- yeto/rl/engine/full_parameter_sync.py | 36 +- yeto/rl/engine/run_flags.py | 6 +- yeto/rl/engine/sync_overlap.py | 32 +- 10 files changed, 670 insertions(+), 58 deletions(-) create mode 100644 tests/test_e3_overlap_miles_role_stream.py diff --git a/tests/test_e3_overlap_miles_role_stream.py b/tests/test_e3_overlap_miles_role_stream.py new file mode 100644 index 000000000..6d9f0aa1d --- /dev/null +++ b/tests/test_e3_overlap_miles_role_stream.py @@ -0,0 +1,495 @@ +"""S26 E3: sync/train overlap on the real Miles role stream (CPU tensors, no GPU). + +``tests/test_e3_sync_overlap.py`` drives ``FullParameterSync`` with a numeric +fake role. Here the coordinator drives the real ``MilesFullParameterRoleStream`` +and ``MilesChunkedFullParameterAdapter`` (``rebase_fragment``, held +submissions, ``finish_rebase``, ``mark_round_base``, retained replays). Only +the Megatron trainer group, Ray and the syncer connection are fakes: + +* the trainer keeps one f32 tensor per fragment and adds a fixed progress per + round, so every pushed delta must equal that round's progress exactly; +* the syncer merges the island's delta with a second, virtual island's delta + (plain average) so the global weights differ from the local ones and a + wrong rebase shows up in the numbers. + +The syncer-restart tests cut the run while the island trains the overlapped +round: the syncer resumes from its last checkpoint (uncommitted pushes lost, +PULLs re-sent with a new attempt, committed fragments re-broadcast) and the +run must still count every round exactly once. + +Needs Miles (``importorskip``): run it in the Miles image. +""" + +from __future__ import annotations + +# ruff: noqa: I001 -- the Miles test helpers install a bridge stub before import. + +import asyncio +import threading +from types import SimpleNamespace + +import pytest +import torch + +pytest.importorskip("miles") + +from test_rl_miles_chunked_full_parameter import FakeRay, _evidence, component, manifest # noqa: E402 + +from miles.backends.megatron_utils.full_parameter_state import ( # noqa: E402 + FullParameterChunkDescriptor, + FullParameterChunkedShardState, + FullParameterFragmentDescriptor, + _chunk_hash, + _fragment_hasher, + _tensor_bytes, +) +from miles.ray.full_parameter_transport import FullParameterChunkedCut # noqa: E402 +from yeto.protocol import BcastFragment, FinalFragment, FinalManifest, PullRequest # noqa: E402 +from yeto.rl.adapters.miles.full_parameter import role_stream as role_stream_module # noqa: E402 +from yeto.rl.adapters.miles.full_parameter.role_stream import ( # noqa: E402 + MilesFullParameterRoleStream, + RoleStreamConfig, +) +from yeto.rl.adapters.miles.models.chunked_full_parameter import ( # noqa: E402 + MilesChunkedFullParameterAdapter, +) +from yeto.rl.engine.full_parameter_sync import ( # noqa: E402 + FullParameterSync, + FullParameterSyncProfile, + RolePairing, +) + +P = 2 +SIZES = (3, 2) # one parameter per owner shard; one fragment per shard +OTHER_WEIGHT = 0.5 # plain average of this island and one virtual peer + + +def _adapter() -> MilesChunkedFullParameterAdapter: + return MilesChunkedFullParameterAdapter.create( + (manifest(0, (SIZES[0],)), manifest(1, (SIZES[1],))), + algorithm="grpo", + components=(component(),), + minimum_fragments=P, + max_fragment_bytes=64, + max_chunk_bytes=8, # two floats per chunk: fragment 0 spans two chunks + ) + + +def _build_cut(adapter, values, *, policy_version, generation, ray) -> FullParameterChunkedCut: + """A reference-backed Miles cut holding ``values[fragment_id]``.""" + shards = [] + for plan in adapter.plans: + descriptors, refs = [], [] + for fragment in plan.fragments: + value = values[fragment.fragment_id].to(torch.float32).contiguous() + assert value.numel() == fragment.numel + digest = _fragment_hasher(fragment.fragment_id, adapter.layout.layout_hash, plan.plan_hash) + chunks, chunk_refs = [], [] + step = adapter.max_chunk_bytes // 4 + for index, start in enumerate(range(0, value.numel(), step)): + chunk = value[start:start + step].clone().contiguous() + digest.update(_tensor_bytes(chunk)) + chunks.append(FullParameterChunkDescriptor( + index, start, chunk.numel(), + _chunk_hash(fragment.fragment_id, index, start, chunk, + parameter_layout_hash=adapter.layout.layout_hash, + plan_hash=plan.plan_hash))) + chunk_refs.append(ray.put(chunk)) + descriptors.append(FullParameterFragmentDescriptor( + fragment.fragment_id, value.numel(), digest.hexdigest(), tuple(chunks))) + refs.append(chunk_refs) + shards.append(FullParameterChunkedShardState( + policy_version, generation, plan.topology, adapter.layout.layout_hash, + plan.plan_hash, tuple(descriptors), refs)) + return FullParameterChunkedCut(policy_version, generation, adapter.layout.layout_hash, + tuple(sorted(shards, key=lambda state: state.topology))) + + +class _Trainer: + """The Megatron trainer group: one f32 tensor per fragment, fixed progress per round.""" + + def __init__(self, adapter, ray, initial, progress) -> None: + self.adapter = adapter + self.ray = ray + self.values = {f: value.clone() for f, value in initial.items()} + self.progress = progress + self.installed = 0 + self.generation = 0 + self.scheduler_steps = 0 + self.trained_from: list[dict[int, torch.Tensor]] = [] + self.on_train = None # hook: runs while a round trains (cut point) + + async def record_full_parameter_local_step(self, *, base_policy_version, rollout_id): + assert base_policy_version == self.installed + self.trained_from.append({f: v.clone() for f, v in self.values.items()}) + for f, delta in self.progress(rollout_id).items(): + self.values[f] = self.values[f] + delta + self.generation += 1 + start, self.scheduler_steps = self.scheduler_steps, self.scheduler_steps + 1 + if self.on_train is not None: + self.on_train(rollout_id) + return tuple( + SimpleNamespace(topology=shard.topology, role="actor", base_policy_version=base_policy_version, + local_step_generation=self.generation, rollout_id=rollout_id, + optimizer_steps=1, scheduler_start_steps=start, + scheduler_end_steps=self.scheduler_steps) + for shard in self.adapter.manifests) + + async def export_full_parameter_chunked_cut(self, policy_version, local_step_generation=0, *, + max_chunk_bytes): + assert max_chunk_bytes == self.adapter.max_chunk_bytes + return _build_cut(self.adapter, self.values, policy_version=policy_version, + generation=local_step_generation, ray=self.ray) + + async def apply_full_parameter_chunked_cut(self, cut, *, commit_token=None, max_commit_attempts=3): + assert commit_token + for shard in cut.shards: + for descriptor, refs in zip(shard.fragments, shard.chunk_refs, strict=True): + assert all(ref is not None for ref in refs), "applied a released chunk" + self.values[descriptor.fragment_id] = torch.cat([self.ray.get(ref) for ref in refs]).clone() + self.installed = cut.policy_version + self.generation = 0 + return self.adapter.expected_parameter_tensor_count + + async def full_parameter_optimizer_states(self): + return tuple(SimpleNamespace(role="actor", installed_policy_version=self.installed, + local_step_generation=self.generation, + scheduler_num_steps=self.scheduler_steps) + for _ in self.adapter.manifests) + + +class _Syncer: + """In-process strict sweep (``--pipeline P``, checkpoint every round). + + ``G_f <- G_f + OTHER_WEIGHT * (island delta + peer delta)``; committed state + is checkpointed at once, so a restart loses only uncommitted pushes.""" + + def __init__(self, initial, sweeps, peer) -> None: + self.global_ = {f: v.clone() for f, v in initial.items()} + self.versions = [0] * P + self.sweeps = sweeps + self.peer = peer # rollout_id -> {fragment: peer delta} + self.received: dict[int, torch.Tensor] = {} # uncommitted pushes by step + self.pushed: list[tuple[int, int, torch.Tensor]] = [] # (step, attempt, delta), every PUSH seen + self.committed: list[int] = [] + self.attempt = 1 + self.client = None + + def commit_ready(self) -> None: + """Commit fragments in strict step order once their PUSH arrived.""" + while True: + step = len(self.committed) + 1 + if step > self.sweeps * P or step not in self.received: + return + f, k = (step - 1) % P, (step - 1) // P + delta = self.received.pop(step) + self.global_[f] = self.global_[f] + OTHER_WEIGHT * (delta + self.peer(k)[f]) + self.versions[f] = step + self.committed.append(step) + if step == self.sweeps * P: + self.client.publish_final() + elif step % P == 0: + # The sweep is committed: broadcast it, then PULL the next one. + for fragment in range(P): + self.client.broadcast(fragment, self.versions[fragment], self.global_[fragment]) + self.client.pull_sweep(step // P, self.attempt) + + def restart_from_checkpoint(self) -> None: + """Crash + ``--resume``: uncommitted pushes are lost; the reconnect + re-broadcasts every committed fragment and re-sends the open PULLs.""" + self.received.clear() + self.attempt += 1 + self.client.reconnect() + for fragment in range(P): + if self.versions[fragment]: + self.client.broadcast(fragment, self.versions[fragment], self.global_[fragment]) + if len(self.committed) % P == 0 and len(self.committed) < self.sweeps * P: + self.client.pull_sweep(len(self.committed) // P, self.attempt) + else: + # A partly committed sweep: only its uncommitted fragments are open. + sweep = len(self.committed) // P + for step in range(len(self.committed) + 1, (sweep + 1) * P + 1): + self.client.pulls.append(PullRequest((step - 1) % P, step, self.attempt)) + + +class _Client: + """The ``SyncerClient`` surface the role stream uses, backed by ``_Syncer``.""" + + def __init__(self, syncer: _Syncer) -> None: + self.syncer = syncer + syncer.client = self + self.connection_generation = 1 + self.finalizing = threading.Event() + self.pulls: list[PullRequest] = [] + self.updates: list[BcastFragment] = [] + self.inbound = None + self.final: tuple[FinalManifest, list[FinalFragment]] | None = None + self.acknowledged: FinalManifest | None = None + self.closed = False + + # -- syncer side + def _stored(self, fragment_id, version, value): + return self.inbound.sink(fragment_id, version, _tensor_bytes(value.to(torch.float32).contiguous())) + + def broadcast(self, fragment_id, version, value) -> None: + self.updates.append(BcastFragment(fragment_id, version, self._stored(fragment_id, version, value), + stored=True)) + + def pull_sweep(self, sweep, attempt) -> None: + self.pulls.extend(PullRequest(f, sweep * P + f + 1, attempt) for f in range(P)) + + def reconnect(self) -> None: + self.connection_generation += 1 + self.pulls.clear() + + def publish_final(self) -> None: + last = self.syncer.sweeps * P + versions = tuple(range(last - P + 1, last + 1)) + fragments = [FinalFragment(f, versions[f], self._stored(f, versions[f], self.syncer.global_[f]), + stored=True) for f in range(P)] + self.final = (FinalManifest(last, versions), fragments) + self.finalizing.set() + + # -- island side (SyncerClient API) + def install_inbound_chunk_sink(self, sink) -> None: + self.inbound = sink + + def start(self) -> None: + pass + + def send_init_parts(self, fragment_id, parts) -> bool: + payload = b"".join(bytes(part) for part in parts) + value = torch.frombuffer(bytearray(payload), dtype=torch.float32) + assert torch.equal(value, self.syncer.global_[fragment_id]) + self.broadcast(fragment_id, 0, value) + if fragment_id == P - 1: + self.pull_sweep(0, 1) + return True + + def check_health(self) -> None: + assert not self.closed + + def drain_updates(self): + out, self.updates = self.updates, [] + return out + + def drain_pulls(self): + out, self.pulls = self.pulls, [] + return out + + def push_fragment_parts(self, fragment_id, global_step, round_attempt, base_version, local_step, + c_steps, c_tokens, parts) -> bool: + assert (global_step - 1) % P == fragment_id and c_steps == 1 and c_tokens >= 1 + assert base_version == self.syncer.versions[fragment_id], "PUSH on a stale base" + delta = torch.frombuffer(bytearray(b"".join(bytes(part) for part in parts)), dtype=torch.float32) + self.syncer.pushed.append((global_step, round_attempt, delta.clone())) + self.syncer.received[global_step] = delta.clone() + self.syncer.commit_ready() + return True + + def wait_for_final_fragments(self, timeout): + assert self.final is not None, "no FINAL before the terminal barrier" + return self.final + + def acknowledge_finalization(self, manifest) -> None: + self.acknowledged = manifest + + def close(self) -> None: + self.closed = True + + +class _Driver: + def __init__(self) -> None: + self.events = [] + + def emit(self, event, **fields) -> None: + self.events.append((event, fields)) + + +def _config(sweeps: int) -> RoleStreamConfig: + adapter = _adapter() + return RoleStreamConfig( + algorithm="grpo", role="actor", component=component(), syncer_addr=("127.0.0.1", 30100), + learner_id=0, learner_generation=0, learner_generations=(0, 0), + total_fragment_steps=sweeps * P, expected_fragments=P, + expected_layout_hash=adapter.layout.layout_hash, local_horizon=1, optimizer_steps_per_round=1, + training_contract_hash="2" * 64, syncer_profile_hash="3" * 64, pipeline_depth=1, + algorithm_spec_sha256="a1" * 32, wait_timeout=5.0, poll_seconds=0.001, + max_fragment_bytes=64, max_chunk_bytes=8, overlap=True) + + +def _run(monkeypatch, sweeps: int, *, cut=None): + """Run ``sweeps`` overlapped rounds; ``cut(rollout_id, syncer)`` runs while + that round trains (the mid-overlap cut point).""" + generator = torch.Generator().manual_seed(11) + + def rand(scale=1.0): + return {f: torch.randn(n, generator=generator) * scale for f, n in enumerate(SIZES)} + + initial = rand() + steps = {k: rand(0.1) for k in range(sweeps)} + peer = {k: rand(0.1) for k in range(sweeps)} + ray = FakeRay() + adapter = _adapter() + trainer = _Trainer(adapter, ray, initial, lambda k: steps[k]) + syncer = _Syncer(initial, sweeps, lambda k: peer[k]) + client = _Client(syncer) + + async def initialize(cls, group, *, policy_version, **_kwargs): + cut = await group.export_full_parameter_chunked_cut(policy_version, 0, + max_chunk_bytes=adapter.max_chunk_bytes) + return adapter, adapter.validate_cut(cut, policy_version, 0) + + monkeypatch.setattr(MilesChunkedFullParameterAdapter, "initialize", classmethod(initialize)) + monkeypatch.setattr(role_stream_module, "SyncerClient", lambda *_args, **_kwargs: client) + + stream = MilesFullParameterRoleStream(_config(sweeps), ray_module=ray) + profile = FullParameterSyncProfile.from_configs({"actor": stream.config}, sweep=True) + assert profile.overlap + session = FullParameterSync( + RolePairing({"actor": stream}, profile), groups={"actor": trainer}, run=asyncio.run, + trajectories=lambda rollout_id: _evidence(rollout_id=rollout_id, + behavior_policy_hash=stream.current.policy_hash)) + if cut is not None: + trainer.on_train = lambda rollout_id: cut(rollout_id, syncer) + driver = _Driver() + session.start(driver) + stops = [] + for rollout_id in range(sweeps): + boundary = session.boundary(driver, rollout_id=rollout_id, stats=None) + stops.append(boundary.stop) + if boundary.stop: + session.published(driver, rollout_id=rollout_id, policy_hash=boundary.state.policy_tensor_hash()) + break + session.finish(driver) + return SimpleNamespace(initial=initial, steps=steps, peer=peer, trainer=trainer, syncer=syncer, + client=client, stream=stream, stops=stops, ray=ray) + + +def _expected_global(run, sweeps): + return {f: run.initial[f] + sum(OTHER_WEIGHT * (run.steps[k][f] + run.peer[k][f]) + for k in range(sweeps)) for f in range(P)} + + +def _committed_deltas(run): + """The delta of each committed step (the last PUSH seen for that step).""" + last = {} + for step, attempt, delta in run.syncer.pushed: + if attempt >= last.get(step, (0, None))[0]: + last[step] = (attempt, delta) + return {step: delta for step, (_attempt, delta) in last.items()} + + +def _check_exactly_once(run, sweeps): + assert run.stops == [False] * (sweeps - 1) + [True] + assert run.syncer.committed == list(range(1, sweeps * P + 1)) + deltas = _committed_deltas(run) + for k in range(sweeps): + for f in range(P): + # Each push is exactly that round's progress: no round is pushed + # twice, none is lost, the rebase kept the local progress. + torch.testing.assert_close(deltas[k * P + f + 1], run.steps[k][f], rtol=0, atol=1e-6) + want = _expected_global(run, sweeps) + for f in range(P): + torch.testing.assert_close(run.syncer.global_[f], want[f], rtol=0, atol=1e-5) + torch.testing.assert_close(run.trainer.values[f], want[f], rtol=0, atol=1e-5) + assert run.stream.final_applied and run.stream.final_acknowledged + assert run.client.acknowledged == run.stream.terminal_manifest + + +def test_overlap_on_the_miles_role_stream_rebases_and_counts_every_round_once(monkeypatch): + sweeps = 4 + run = _run(monkeypatch, sweeps) + _check_exactly_once(run, sweeps) + trained = run.trainer.trained_from + for f in range(P): + # Round 0 starts at the initial weights, round 1 at its own round-0 + # result (sweep 0 in flight), round k >= 2 at G_{k-2} + progress of + # round k-1: the rebase ran in the trainer through rebase_fragment. + torch.testing.assert_close(trained[0][f], run.initial[f]) + torch.testing.assert_close(trained[1][f], run.initial[f] + run.steps[0][f]) + for k in range(2, sweeps): + g = run.initial[f] + sum(OTHER_WEIGHT * (run.steps[j][f] + run.peer[j][f]) for j in range(k - 1)) + torch.testing.assert_close(trained[k][f], g + run.steps[k - 1][f], rtol=0, atol=1e-5) + # Exactly one PUSH per step without a restart. + assert sorted(step for step, _attempt, _delta in run.syncer.pushed) == list(range(1, sweeps * P + 1)) + + +def test_overlap_role_stream_releases_every_superseded_cut(monkeypatch): + run = _run(monkeypatch, 3) + stream = run.stream + assert not stream._pinned_cuts and not stream._pinned_anchors + assert not stream._staged + + +@pytest.mark.parametrize("cut_round", [1, 2]) +def test_syncer_restart_before_the_commit_mid_overlap_replays_the_retained_push(monkeypatch, cut_round): + """Cut while round ``cut_round`` trains: sweep ``cut_round - 1`` was pushed + but not committed. The resumed syncer re-PULLs it with attempt 2; the + island replays the retained bytes at its next boundary.""" + sweeps = 4 + + def cut(rollout_id, syncer): + if rollout_id == cut_round and syncer.attempt == 1: + # Undo the commit of the in-flight sweep: the crash came before it. + sweep = cut_round - 1 + for step in range(sweep * P + 1, (sweep + 1) * P + 1): + syncer.committed.remove(step) + syncer.global_ = {f: g.clone() for f, g in checkpoints[sweep].items()} + syncer.versions = [sweep * P - P + f + 1 if sweep else 0 for f in range(P)] + syncer.client.updates.clear() # its BCAST never left the syncer + syncer.restart_from_checkpoint() + + checkpoints = {} + original = _Syncer.commit_ready + + def commit_ready(self): + if len(self.committed) % P == 0: + checkpoints.setdefault(len(self.committed) // P, {f: g.clone() for f, g in self.global_.items()}) + original(self) + + monkeypatch.setattr(_Syncer, "commit_ready", commit_ready) + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + replayed = [(step, attempt) for step, attempt, _ in run.syncer.pushed if attempt == 2] + first = (cut_round - 1) * P + 1 + assert [step for step, _ in replayed][:P] == list(range(first, first + P)) + + +@pytest.mark.parametrize("cut_round", [1, 2]) +def test_syncer_restart_after_the_commit_mid_overlap_rebases_on_the_resent_broadcast(monkeypatch, cut_round): + """Cut while round ``cut_round`` trains: sweep ``cut_round - 1`` committed + and checkpointed, but its BCAST was lost. The resumed syncer re-sends it + (FULL) and PULLs the next sweep; the island rebases on the re-sent bytes.""" + sweeps = 4 + + def cut(rollout_id, syncer): + if rollout_id == cut_round and syncer.attempt == 1: + syncer.client.updates.clear() # the BCAST was lost with the connection + syncer.restart_from_checkpoint() + + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + assert all(attempt == 1 for step, attempt, _ in run.syncer.pushed if step <= cut_round * P) + + +def test_syncer_restart_with_a_partly_committed_sweep_mid_overlap(monkeypatch): + """Cut while round 2 trains: fragment 0 of sweep 1 committed (and is + re-broadcast), fragment 1 was lost. Only fragment 1 is re-PULLed and + replayed; the island applies the sweep only once both fragments arrived.""" + sweeps = 3 + lost = 2 * P # sweep 1, fragment 1 + + def cut(rollout_id, syncer): + if rollout_id == 2 and syncer.attempt == 1: + f = (lost - 1) % P + pushed = [delta for step, _attempt, delta in syncer.pushed if step == lost][-1] + syncer.committed.remove(lost) + syncer.global_[f] = syncer.global_[f] - OTHER_WEIGHT * (pushed + syncer.peer(1)[f]) + syncer.versions[f] = lost - P + syncer.client.updates.clear() + syncer.restart_from_checkpoint() + + run = _run(monkeypatch, sweeps, cut=cut) + _check_exactly_once(run, sweeps) + assert [step for step, attempt, _ in run.syncer.pushed if attempt == 2][0] == lost diff --git a/tests/test_e3_sync_overlap.py b/tests/test_e3_sync_overlap.py index 9b9fbc4d2..6421d6ce7 100644 --- a/tests/test_e3_sync_overlap.py +++ b/tests/test_e3_sync_overlap.py @@ -5,7 +5,7 @@ sweep k is in flight and rebases ``theta <- G_k + (local - base)``. The tests check that every round's progress reaches the global weights exactly once (no double counting), that the trainer keeps its progress across a rebase, -and the launcher/island policy-age rule. +and the launcher rules (rollout policy age 0, fragment pipeline P). """ from __future__ import annotations @@ -28,7 +28,7 @@ FullParameterSyncProfile, RolePairing, ) -from yeto.rl.engine.sync_overlap import rebase_values, required_policy_age +from yeto.rl.engine.sync_overlap import OUTER_STEP_LAG, overlap_option_problems, rebase_values P = 2 # fragments NUMEL = 4 @@ -95,6 +95,7 @@ def __init__(self, syncer: _SweepSyncer, progress) -> None: self.final_acknowledged = False self.closed = False self.trained_from: list[list[np.ndarray]] = [] + self.push_rounds: list[tuple[int, int]] = [] # (global step, local rounds trained then) @property def fragment_versions(self): @@ -143,6 +144,7 @@ def release_candidate(self, candidate): def push_candidate(self, candidate): f = candidate.permit.fragment_id + self.push_rounds.append((candidate.permit.global_step, self.local_rounds)) self.syncer.push(candidate.permit.global_step, f, candidate.submission_cut[f] - candidate.anchor) self._pending.pop(f) @@ -204,12 +206,31 @@ def emit(self, event, **fields): self.events.append((event, fields)) -def _run(sweeps: int, overlap: bool = True): +class _LatePullSyncer(_SweepSyncer): + """Shows the last fragment's PULL of each sweep only after a few polls.""" + + LATE_POLLS = 3 + + def __init__(self, *args) -> None: + super().__init__(*args) + self._polls: dict[int, int] = {} + + def permits(self): + permits = super().permits() + late = [p for p in permits if p.fragment_id == P - 1] + if late: + seen = self._polls[late[0].global_step] = self._polls.get(late[0].global_step, 0) + 1 + if seen <= self.LATE_POLLS: + return [p for p in permits if p.fragment_id != P - 1] + return permits + + +def _run(sweeps: int, overlap: bool = True, syncer_type=_SweepSyncer): rng = np.random.default_rng(7) initial = [rng.standard_normal(NUMEL).astype(np.float32) for _ in range(P)] steps = {k: [rng.standard_normal(NUMEL).astype(np.float32) * 0.1 for _ in range(P)] for k in range(sweeps)} - syncer = _SweepSyncer(initial, sweeps) + syncer = syncer_type(initial, sweeps) role = _NumericRole(syncer, lambda rollout_id: steps[rollout_id]) profile = FullParameterSyncProfile(roles=("actor",), total_fragment_steps=sweeps * P, fragments=P, local_horizon=1, pipeline_depth=1, wait_timeout=5.0, @@ -245,6 +266,18 @@ def test_overlap_counts_every_round_exactly_once(): assert role.final_acknowledged +def test_overlap_boundary_pushes_the_whole_sweep_even_when_a_pull_arrives_late(): + sweeps = 3 + _initial, _steps, syncer, role, stops = _run(sweeps, syncer_type=_LatePullSyncer) + assert stops == [False] * (sweeps - 1) + [True] + # Every sweep was complete before the next round trained: sweep k's last + # fragment is pushed before round k+1's progress exists (no sweep spans + # two boundaries). + assert role.push_rounds == [(k * P + f + 1, k + 1) for k in range(sweeps) for f in range(P)] + assert role.final_acknowledged + assert all(polls > _LatePullSyncer.LATE_POLLS for polls in syncer._polls.values()) + + def test_overlap_trains_one_outer_step_stale_and_keeps_local_progress(): sweeps = 4 initial, steps, _syncer, role, _ = _run(sweeps) @@ -287,34 +320,40 @@ def _dense(*extra): return args -def test_overlap_raises_the_minimum_policy_age_and_reaches_the_island(): +def test_overlap_keeps_the_rollout_policy_age_zero_and_reaches_the_island(): from yeto.launch.rl_island_task import _rl_trainable_island_flags + from yeto.launch.rl_ports import _ports_infra_flags - assert required_policy_age(0) == 1 - args = _dense("--rl-sync-overlap") - assert args.rl_max_policy_age == 1 + args = _dense("--rl-sync-overlap", "--rl-sync-fragment-pipeline", "2") + # Overlap is not rollout staleness: no policy age, so Miles gets neither + # partial rollout (carry-over) nor fully_async generation. + assert int(args.rl_max_policy_age or 0) == 0 + assert OUTER_STEP_LAG == 1 assert "--sweep-overlap" in _rl_trainable_island_flags(args, 0, 1) - assert _dense("--rl-sync-overlap", "--rl-max-policy-age", "3").rl_max_policy_age == 3 + assert "--rl-max-policy-age" not in " ".join(_ports_infra_flags(args)) plain = _dense() - assert int(plain.rl_max_policy_age or 0) == 0 assert "--sweep-overlap" not in _rl_trainable_island_flags(plain, 0, 1) -def test_overlap_launch_fails_when_the_algorithm_spec_does_not_tolerate_the_age(): - from yeto.launch_preflight import PreflightError, check_policy_age_spec - - args = _dense("--rl-sync-overlap") - with pytest.raises(PreflightError, match="--rl-max-policy-age 1 needs an algorithm spec"): - check_policy_age_spec(args) +def test_overlap_launch_refuses_a_policy_age_or_a_partial_fragment_pipeline(): + with pytest.raises(ValueError, match="--rl-max-policy-age 0 with --rl-sync-overlap"): + _dense("--rl-sync-overlap", "--rl-sync-fragment-pipeline", "2", "--rl-max-policy-age", "1") + with pytest.raises(ValueError, match="--rl-sync-fragment-pipeline equal to --fragments"): + _dense("--rl-sync-overlap") + assert overlap_option_problems(max_policy_age=0, fragment_pipeline=4, fragments=4) == [] -def test_island_refuses_overlap_without_the_raised_policy_age(): +def test_island_builds_overlap_without_a_policy_age(): from yeto.rl.adapters.miles.island_entry import _dense_full_sync_config - args = SimpleNamespace(sweep_overlap=True, rl_max_policy_age=0) - with pytest.raises(ValueError, match="--sweep-overlap .* needs --rl-max-policy-age >= 1"): - _dense_full_sync_config(args, model_config_sha256="0" * 64, - training_contract_sha256="0" * 64, algorithm_spec_sha256="0" * 64) + args = SimpleNamespace( + sweep_overlap=True, rl_max_policy_age=0, model_revision="A" * 40, syncer="127.0.0.1:30100", + learner_id=0, num_learners=2, total_fragment_steps=4, fragments=2, + parameter_layout_sha256="1" * 64, optimizer_steps=1, syncer_profile_sha256="3" * 64, + wan_streams=4) + config = _dense_full_sync_config(args, model_config_sha256="0" * 64, + training_contract_sha256="0" * 64, algorithm_spec_sha256="0" * 64) + assert config.roles[0].overlap def test_overlap_is_bound_into_the_role_stream_contract(): diff --git a/yeto/launch/rl_args.py b/yeto/launch/rl_args.py index 60e63f88b..e623e6dc1 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -214,6 +214,12 @@ def _prepare_dense_full_sweep(args, *, num_islands: int) -> None: problems.append("--rl-sync-push-int8-block between 32 and 1048576") elif int8_block != 128 and getattr(args, "rl_sync_push_dtype", "f32") != "int8": problems.append("--rl-sync-push-int8-block only with --rl-sync-push-dtype int8") + if getattr(args, "rl_sync_overlap", False): + from ..rl.engine.sync_overlap import overlap_option_problems + + problems += overlap_option_problems( + max_policy_age=getattr(args, "rl_max_policy_age", 0) or 0, + fragment_pipeline=fragment_pipeline, fragments=args.fragments) if problems: raise ValueError(f"--rl-sync-preset {FULL_PARAMETER_SYNC_PRESET} requires " + "; ".join(problems)) args.quorum = num_islands @@ -363,15 +369,6 @@ def _prepare_rl_args( # what the selected backend supports before any cloud or GPU work. from ..rl.engine.policy_age import validate_limit - if args.rl_sync_preset == FULL_PARAMETER_SYNC_PRESET and getattr(args, "rl_sync_overlap", False): - # S26 E3: sync/train overlap adds one outer step of staleness; raise the - # limit to the new minimum here so every check below (backend support, - # and the preflight's algorithm-spec staleness) validates it at launch. - from ..rl.engine.sync_overlap import required_policy_age - - required = required_policy_age(0) - if int(getattr(args, "rl_max_policy_age", 0) or 0) < required: - args.rl_max_policy_age = required max_policy_age = validate_limit(getattr(args, "rl_max_policy_age", 0) or 0) policy_age_module = _rl_backend_module(args, "policy_age") policy_age_module.SUPPORT.check(max_policy_age) diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 4b2a5199f..1855ad780 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -99,7 +99,8 @@ class RoleStreamConfig: push_int8_block: int = DEFAULT_I8_BLOCK bcast_delta: bool = False # S26 E3: train round k+1 while sweep k syncs, then rebase the trainer - # (see FullParameterSync._overlap_boundary); one more outer step stale. + # (see FullParameterSync._overlap_boundary); the local weights lag the + # global policy by one outer step (yeto.rl.engine.sync_overlap). overlap: bool = False @property @@ -653,16 +654,20 @@ async def apply_staged_broadcast( return (fragment_id,) def finish_rebase(self) -> None: - """Overlap: the rebased cut is the next submission (``G + progress``).""" + """Overlap: the rebased cut is the next submission (``G + progress``). - if not self.config.overlap or self.current is None: + The applied cut has generation 0; the submission is the same weights + labeled with the rounds trained since the round base.""" + + if not self.config.overlap or self.current is None or self.adapter is None: raise RuntimeError(f"{self.config.role} finish_rebase outside overlap") previous_submission = self.submission_cut - self.submission_cut = self.current + self.submission_cut = self.adapter.relabel_local_progress( + self.current, local_rounds=self.local_rounds - self._rounds_at_base) self.submission_fragment_versions = self.fragment_versions self.submission_rounds = self.local_rounds self.submission_tokens = self.action_tokens - self._release_cuts(previous_submission, keep=(self.current, self._round_base)) + self._release_cuts(previous_submission, keep=(self.current, self.submission_cut, self._round_base)) def mark_round_base(self) -> None: """Overlap: pin the cut the next local round starts from.""" diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index c382f8cb4..6ff6a2e84 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -814,14 +814,8 @@ def _dense_full_sync_config( from yeto.rl.adapters.miles.full_parameter.role_stream import RoleStreamConfig from yeto.rl.adapters.miles.full_parameter.session import FullParameterSyncConfig - from yeto.rl.engine.sync_overlap import required_policy_age from yeto.rl.local_learner import ComponentIdentity - overlap = bool(getattr(args, "sweep_overlap", False)) - if overlap and int(getattr(args, "rl_max_policy_age", 0) or 0) < required_policy_age(0): - raise ValueError("--sweep-overlap rollouts are one outer step stale: it needs " - f"--rl-max-policy-age >= {required_policy_age(0)}") - actor = RoleStreamConfig( algorithm="grpo", role="actor", @@ -847,7 +841,9 @@ def _dense_full_sync_config( push_dtype=getattr(args, "sweep_push_dtype", "f32"), push_int8_block=int(getattr(args, "sweep_push_int8_block", 128)), bcast_delta=bool(getattr(args, "bcast_delta", False)), - overlap=overlap, + # The rollout policy age stays 0 (``_check_dense_full_sweep``); the + # overlap's outer-step lag is bound into the session contract. + overlap=bool(getattr(args, "sweep_overlap", False)), ) return FullParameterSyncConfig((actor,), sweep=True) diff --git a/yeto/rl/adapters/miles/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index f16a6ec56..14a675d60 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -840,6 +840,42 @@ def flat(cut: ReferencedPolicyCut) -> np.ndarray: sink = AuthoritativeFragmentSink(current, ray_module=ray_module) return sink(fragment_id, stored.version, values.tobytes()) + def relabel_local_progress( + self, + cut: ReferencedPolicyCut, + *, + local_rounds: int, + ) -> ReferencedPolicyCut: + """Sync/train overlap: label a rebased cut with its local progress. + + ``apply`` resets the local step generation to 0, but a rebased cut + (``G + (current - base)``) still holds ``local_rounds`` rounds of + progress over its anchor ``G``, and Miles pushes only cuts with + progress (generation >= 1). Values, chunk hashes and refs are shared + with ``cut`` (fresh ref lists, so releasing one leaves the other); + only the generation label changes.""" + + from dataclasses import replace + + from miles.ray.full_parameter_transport import FullParameterChunkedCut + + if isinstance(local_rounds, bool) or not isinstance(local_rounds, int) or local_rounds < 1: + raise ValueError("a rebased cut needs at least one local round") + if cut.layout_hash != self.layout.layout_hash or cut.local_step_generation != 0: + raise ValueError("only an applied (generation 0) cut is relabeled") + transport = cut.transport_cut + relabeled = FullParameterChunkedCut( + policy_version=transport.policy_version, + local_step_generation=local_rounds, + parameter_layout_hash=transport.parameter_layout_hash, + shards=tuple( + replace(shard, local_step_generation=local_rounds, + chunk_refs=[list(references) for references in shard.chunk_refs]) + for shard in transport.shards + ), + ) + return self.validate_cut(relabeled, cut.policy_version, local_rounds) + async def apply( self, group: MilesChunkedFullParameterGroup, diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 79b6da70c..ff2830bcd 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -329,7 +329,7 @@ "--sweep-overlap", action="store_true", help="full-parameter sweep: train the next round while the sweep syncs, then " - "rebase onto the new global weights (needs --rl-max-policy-age >= 1)", + "rebase onto the new global weights (rollout policy age stays 0)", ), island_flag( "--bcast-delta", diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index 7c9d0921f..744f30535 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -612,27 +612,53 @@ async def _overlap_boundary(self, rollout_id: int) -> bool: if not pairing.any_finalizing: for stream in pairing.streams.values(): stream.finish_rebase() - rebased_ts = unix_now() if pairing.any_finalizing: await pairing.finish_without_more_training() return True before = pairing.fragment_versions() - pairing.submit_ready() - submit_ts = unix_now() if rollout_id + 1 >= pairing.profile.total_rounds: # Terminal round: nothing trains after it, so wait for the cut. + pairing.submit_ready() + submit_ts = unix_now() await pairing.commit_staged_broadcasts() await self._wait_for_sweep(before) self._timing = SyncTiming(submit_ts, unix_now()) await pairing.finish_without_more_training() return True + await self._push_whole_sweep(before) self._pending_sweep = before for stream in pairing.streams.values(): stream.mark_round_base() - # This boundary blocked only on the previous sweep (wait + rebase). - self._timing = SyncTiming(wait_ts, rebased_ts if hold else None) + # This boundary blocked on the previous sweep (wait + rebase) and on + # handing every fragment of this sweep to the wire. + self._timing = SyncTiming(wait_ts, unix_now()) return False + async def _push_whole_sweep(self, before: tuple[int, ...]) -> None: + """Overlap: submit every fragment of the sweep after ``before`` before + training on. A PULL that arrives later would otherwise wait for the + next boundary, a whole local round, and stall the sweep. The syncer + sends all P PULLs at once with ``--pipeline P`` (required by the + launcher), so this waits only for their arrival.""" + + pairing = self.pairing + fragments = pairing.profile.fragments + steps = {version + fragments if version else fragment_id + 1 + for fragment_id, version in enumerate(before)} + deadline = time.monotonic() + pairing.profile.wait_timeout + + def pushed() -> bool: + return steps <= {key[0] for key, pair in pairing.inflight.items() if pair.complete} + + pairing.submit_ready() + while not pushed(): + if time.monotonic() >= deadline: + error = TimeoutError("overlap could not submit the whole sweep before the wait timeout") + pairing.fail_stop(error) + raise error + await asyncio.sleep(pairing.profile.poll_seconds) + pairing.submit_ready() + async def _wait_for_sweep(self, before: tuple[int, ...]) -> None: """Strict sweep: block until all P fragments of this local step are applied, i.e. every fragment version moved past ``before``. An empty diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index a69e67fdc..2af06a5cc 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -863,9 +863,9 @@ def rl_forward_pairs(*dests: str | tuple[str, str]) -> tuple[tuple[str, str], .. action="store_true", help="full-parameter only (E2 option): train the next round while the " "sweep syncs, then rebase the trainer onto the new global weights " - "(theta <- G_k + local progress). Rollouts are one outer step staler: " - "--rl-max-policy-age is raised to at least 1 and the algorithm spec must " - "tolerate it (checked before launch)", + "(theta <- G_k + local progress). The local weights lag the global policy " + "by one outer step; rollouts stay on-policy, so --rl-max-policy-age stays 0. " + "Needs --rl-sync-fragment-pipeline equal to --fragments", ), cli_flag( "--rl-sync-push-int8-block", diff --git a/yeto/rl/engine/sync_overlap.py b/yeto/rl/engine/sync_overlap.py index 4c47f344d..3a7f17ec0 100644 --- a/yeto/rl/engine/sync_overlap.py +++ b/yeto/rl/engine/sync_overlap.py @@ -10,21 +10,39 @@ the next push ``theta - G_k`` equals round k+1's progress exactly once: round k's delta is already inside ``G_k`` and is not counted again. -Rollouts of round k+1 use weights one outer step older than the global -policy, so the run's minimum ``--rl-max-policy-age`` grows by -:data:`EXTRA_POLICY_AGE`. +Staleness is configured apart from ``--rl-max-policy-age``. The driver +publishes the boundary's cut to the island's engines before the next +generation (boundary, then publish), so every rollout is still generated by +the weights the trainer then trains: the rollout policy age stays 0 and the +sweep keeps refusing any other value. What overlap adds is a lag of the local +weights behind the global policy of :data:`OUTER_STEP_LAG` outer step (like +one extra DiLoCo local step); it is bound into the role-stream session +contract (``sync_train_overlap``), not into the rollout policy age, which in +Miles would switch on partial rollout (carry-over) or fully_async generation. """ from __future__ import annotations import numpy as np -EXTRA_POLICY_AGE = 1 +OUTER_STEP_LAG = 1 -def required_policy_age(minimum: int) -> int: - """The smallest accepted ``--rl-max-policy-age`` once overlap is on.""" - return int(minimum) + EXTRA_POLICY_AGE +def overlap_option_problems(*, max_policy_age: int, fragment_pipeline: int, fragments: int) -> list[str]: + """Launcher requirements of ``--rl-sync-overlap`` (empty when satisfied). + + * rollout policy age 0: see the module doc (overlap is not rollout staleness); + * syncer fragment pipeline = P: the syncer then PULLs every fragment of a + sweep at once, so the island pushes the whole sweep at the boundary. With + a smaller pipeline the later PULLs arrive while the next round trains and + wait for the next boundary, which serializes the sweep again.""" + problems = [] + if int(max_policy_age or 0) != 0: + problems.append("--rl-max-policy-age 0 with --rl-sync-overlap (overlap lags the global " + "policy by one outer step; rollouts stay on-policy)") + if int(fragment_pipeline) != int(fragments): + problems.append("--rl-sync-fragment-pipeline equal to --fragments with --rl-sync-overlap") + return problems def rebase_values(anchor: np.ndarray, current: np.ndarray, base: np.ndarray) -> np.ndarray: From 152374b226bc86aa46fad3a0764e039e5c047342 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:30:04 +0000 Subject: [PATCH 5/9] test: overlap island with bf16 downlink deltas against the real syncer The overlap rebase uses the bf16-tracked anchor and pushes theta - anchor against that same anchor, so the syncer's f32 masters still receive every round's progress exactly once and no DownlinkResync happens. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_dense_sweep_real_syncer.py | 89 ++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 31589d679..ec4a66ecf 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -432,3 +432,92 @@ def test_bf16_downlink_through_a_streaming_sink_matches_the_buffered_path(syncer sink=_BytesSink()) assert buffered[2] == streamed[2] assert buffered_updates == streamed_updates + + +def _overlap_island_run(binary: Path, tmp_path: Path, extra: tuple[str, ...], dtype: int, + sweeps: int, progress: dict[tuple[int, int], list[float]]): + """One overlap island (E3 item 3) against the real syncer. + + Round s trains before sweep s-1 comes back; at boundary s the island waits + for G_{s-1}, rebases ``theta <- G_{s-1} + (theta - base)`` and pushes + ``theta - G_{s-1}``. Returns (FINAL bytes, the island's anchors per sweep, + redials used).""" + process, port, profile_hash = _start_syncer(binary, tmp_path, *extra, total_steps=sweeps) + client = _client(port, profile_hash, dtype) + init = {fid: torch.tensor([float(fid) + 0.125 * i for i in range(NUMEL)]) for fid in range(FRAGMENTS)} + anchors: dict[int, dict[int, torch.Tensor]] = {} + held: dict[tuple[int, int], torch.Tensor] = {} + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32(init[fid].tolist())) + anchor = dict(init) # what the last push was measured against + theta = {fid: init[fid] + torch.tensor(progress[(1, fid)]) for fid in range(FRAGMENTS)} + for sweep in range(1, sweeps + 1): + if sweep > 1: + # Round `sweep` trained from `base` while sweep-1 was in flight. + base = dict(theta) + theta = {fid: base[fid] + torch.tensor(progress[(sweep, fid)]) for fid in range(FRAGMENTS)} + + def got_sweep(): + client.check_health() + held.update({(u.fragment_id, u.version): _tail(bytes(u.data)) + for u in client.drain_updates()}) + return all((fid, (sweep - 2) * FRAGMENTS + fid + 1) in held for fid in range(FRAGMENTS)) + + _wait(got_sweep, f"sweep {sweep - 1} never came back") + anchor = {fid: held[(fid, (sweep - 2) * FRAGMENTS + fid + 1)] for fid in range(FRAGMENTS)} + anchors[sweep - 1] = anchor + theta = {fid: anchor[fid] + (theta[fid] - base[fid]) for fid in range(FRAGMENTS)} + permits: dict[int, object] = {} + + def gather(): + client.check_health() + permits.update({p.fragment_id: p for p in client.drain_pulls()}) + return len(permits) == FRAGMENTS + + _wait(gather, f"no PULLs for sweep {sweep}") + for fid, permit in sorted(permits.items()): + base_version = 0 if sweep == 1 else (sweep - 2) * FRAGMENTS + fid + 1 + client.push_fragment(fid, permit.global_step, permit.round_attempt, base_version, + sweep, 1, 7, _f32((theta[fid] - anchor[fid]).tolist())) + manifest, finals = client.wait_for_final_fragments(timeout=30) + result = {f.fragment_id: _tail(bytes(f.data)) for f in finals} + client.acknowledge_finalization(manifest) + redials = client._reconnects_used + assert process.wait(timeout=15) == 0 + finally: + client.close() + log = _stop(process) + assert "syncer semantic profile" not in log, log + return init, result, anchors, redials + + +def test_overlap_with_bf16_downlink_counts_each_round_once_without_a_resync(syncer_binary, tmp_path): + """E3: overlap + bf16 BCAST deltas. The island rebases on a bf16-tracked + anchor, but it pushes ``theta - anchor`` against that same anchor, so the + syncer's f32 masters still get every round's progress exactly once; the + downlink error stays out of the masters and never triggers a resync.""" + sweeps = 4 + progress = {(s, fid): [0.5 + fid + 0.03125 * s + 0.001 * i for i in range(NUMEL)] + for s in range(1, sweeps + 1) for fid in range(FRAGMENTS)} + exact_dir, delta_dir = tmp_path / "f32", tmp_path / "delta" + exact_dir.mkdir() + delta_dir.mkdir() + init, exact, exact_anchors, exact_redials = _overlap_island_run( + syncer_binary, exact_dir, STREAMED, DTYPE_F32, sweeps, progress) + _, delta, delta_anchors, delta_redials = _overlap_island_run( + syncer_binary, delta_dir, (*STREAMED, "--rl-sync-bcast-delta-full-every", "3"), + DTYPE_F32 | DTYPE_FLAG_BCAST_DELTA, sweeps, progress) + assert exact_redials == delta_redials == 0, "no DownlinkResync expected" + for fid in range(FRAGMENTS): + want = init[fid] + sum(torch.tensor(progress[(s, fid)]) for s in range(1, sweeps + 1)) + # (anchor + p) - anchor in f32 is p up to an ulp of the anchor: no bf16 error. + assert torch.allclose(exact[fid], want, rtol=0, atol=1e-5), fid + assert torch.allclose(delta[fid], want, rtol=0, atol=1e-5), fid + for sweep, anchors in exact_anchors.items(): + got = delta_anchors[sweep][fid] + # one bf16 rounding of (G - previous anchor); |G - previous anchor| < 4 + assert torch.allclose(got, anchors[fid], rtol=0, atol=4.0 / 256), (sweep, fid) + assert any(not torch.equal(delta_anchors[s][f], exact_anchors[s][f]) + for s in delta_anchors for f in range(FRAGMENTS)), "the downlink must actually round" From bc8f53b373ad2a61ddc762c8cbe133affadc2d51 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:24:08 +0000 Subject: [PATCH 6/9] syncer: send each fragment's FINAL version as an exact FULL broadcast With bf16 downlink deltas, the last BCAST of a fragment can carry the same version as its FINAL. Islands store one copy per (fragment, version) and refuse a second copy with other bytes, so the delta-tracked BCAST and the exact FINAL collided ('authoritative fragment replay changed bytes'), the islands redialed in a loop and finalization timed out (s26-e3-p06-D8). Broadcasts of a fragment's last version now always go out FULL. Co-Authored-By: Claude Opus 5.5 --- syncer/src/downlink.rs | 13 +++++++++++++ syncer/src/server.rs | 14 +++++++++++++- tests/test_rl_dense_sweep_real_syncer.py | 16 ++++++++++++++++ 3 files changed, 42 insertions(+), 1 deletion(-) diff --git a/syncer/src/downlink.rs b/syncer/src/downlink.rs index e303e119c..32e6f82da 100644 --- a/syncer/src/downlink.rs +++ b/syncer/src/downlink.rs @@ -144,6 +144,19 @@ impl DownlinkAnchors { }) } + /// Encode a newly committed `version` of fragment `p` as FULL (exact f32) + /// and reset the shared anchor to it. + pub fn publish_full(&mut self, p: usize, version: u64, params: &[f32]) -> EncodedBcast { + match self.anchors.get_mut(p) { + Some(slot) => Self::reset(slot, p, version, params), + None => EncodedBcast { + prefix: prefix(p, version, KIND_FULL, version, &digest_f32(params)), + dtype: DTYPE_F32, + values: Arc::new(params.to_vec()), + }, + } + } + /// FULL re-send for recovery (HELLO / reconnect): the anchor every other /// learner holds, when it is at `version`; otherwise the exact params. pub fn resend(&self, p: usize, version: u64, params: &[f32]) -> EncodedBcast { diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 46d4cb52b..269c0b446 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -3495,7 +3495,19 @@ async fn publish_round( // bf16 downlink: encode this version once (advancing the shared // anchor) and send the same bytes to every learner. let published = match &st.downlink { - Some(anchors) => Some(anchors.lock().unwrap().publish(p, st.versions[p], &st.params[p])?), + Some(anchors) => { + // The last version of each fragment is also its FINAL version. + // Islands key stored fragments by (fragment, version) and + // refuse a second copy with other bytes, so that broadcast + // must carry the exact f32 masters (FULL), not a bf16 delta. + let last_version = t + st.params.len() as u64 > cfg.total_steps; + let mut anchors = anchors.lock().unwrap(); + Some(if last_version { + anchors.publish_full(p, st.versions[p], &st.params[p]) + } else { + anchors.publish(p, st.versions[p], &st.params[p])? + }) + } None => None, }; for group in current_groups(registry) { diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index ec4a66ecf..eb9f22d30 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -521,3 +521,19 @@ def test_overlap_with_bf16_downlink_counts_each_round_once_without_a_resync(sync assert torch.allclose(got, anchors[fid], rtol=0, atol=4.0 / 256), (sweep, fid) assert any(not torch.equal(delta_anchors[s][f], exact_anchors[s][f]) for s in delta_anchors for f in range(FRAGMENTS)), "the downlink must actually round" + + +def test_bf16_downlink_sends_each_fragments_final_version_exact(syncer_binary, tmp_path): + """E3 (s26-e3-p06-D8): the last BCAST of a fragment has the same version as + its FINAL. Islands keep one copy per (fragment, version) and refuse a + second copy with other bytes, so that BCAST must be the exact f32 masters + even when the FULL cadence (here M=100) would send a bf16 delta.""" + delta, updates = _downlink_run(syncer_binary, tmp_path, 100, sweeps=4) + final_versions, finals = delta[1], delta[2] + assert final_versions == [7, 8] + # Fragment 0's FINAL version 7 was also broadcast; fragment 1's version 8 + # is the terminal step, which only FINAL carries. + assert updates[(0, 7)] == finals[0] + assert (1, 8) not in updates + # Earlier versions still travel as bf16 deltas. + assert any(updates[(0, v)] != finals[0] for v in (3, 5)) From 3dcbec86a95ececa6dd6170f7da6b7eac1dc0523 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 07:31:58 +0000 Subject: [PATCH 7/9] wan: per-stream stall timeout, rate-aware chunk striping, sync diagnostics E3 showed sweeps waiting on the slowest island's upload and one C32 run where a stream stopped moving with every socket still open. Island-side SyncerClient changes (all default off): - --wan-stall-timeout S: a watchdog fails the connection group when an outbound stream holds bytes, or an inbound chunked message stops growing, for S seconds; sockets also get TCP_USER_TIMEOUT=S. Recovery uses the existing redial path: the syncer replays its PULLs and the island pushes the fragment again at the same (global_step, round_attempt). - --wan-stream-schedule adaptive: each 4 MiB chunk goes to the data stream with the earliest estimated finish (backlog / EWMA rate) instead of round robin, so slow or stalled streams carry less. No wire or syncer change. - rl_full_parameter_round now records per-stream upload rates, the wait after this island's upload drained, and stall counts; scalars go to W&B on rl/step. A real-syncer test stalls one stream mid-sweep (rr and adaptive) and checks the redial keeps versions and finals bit for bit. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_dense_sweep_real_syncer.py | 96 ++++++++ tests/test_wan_streams.py | 229 ++++++++++++++++++ yeto/cli.py | 16 ++ yeto/launch/rl_island_task.py | 17 +- yeto/protocol.py | 144 ++++++++++- .../miles/full_parameter/role_stream.py | 10 + yeto/rl/adapters/miles/island_entry.py | 2 + yeto/rl/adapters/miles/run_flags.py | 2 + yeto/rl/engine/full_parameter_sync.py | 25 +- yeto/rl/engine/sync_timing.py | 39 +++ yeto/rl/wandb_rl.py | 30 +++ yeto/wan_streams.py | 190 +++++++++++++++ 12 files changed, 789 insertions(+), 11 deletions(-) create mode 100644 tests/test_wan_streams.py create mode 100644 yeto/wan_streams.py diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index eb9f22d30..92bf9b97c 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -13,6 +13,7 @@ import shutil import socket import subprocess +import threading import time from pathlib import Path @@ -537,3 +538,98 @@ def test_bf16_downlink_sends_each_fragments_final_version_exact(syncer_binary, t assert (1, 8) not in updates # Earlier versions still travel as bf16 deltas. assert any(updates[(0, v)] != finals[0] for v in (3, 5)) + + +class _StallingSocket: + """Proxy whose ``send`` stops making progress once ``armed`` is set, + until the socket is closed.""" + + def __init__(self, sock: socket.socket, armed: threading.Event) -> None: + self._sock = sock + self._armed = armed + + def send(self, data) -> int: + if not self._armed.is_set(): + return self._sock.send(data) + while self._sock.fileno() != -1: + time.sleep(0.02) + raise OSError("stalled stream closed") + + def __getattr__(self, name): + return getattr(self._sock, name) + + +class _StallOnceClient(SyncerClient): + """Data stream 1 of the first connection group stalls once ``armed``.""" + + stalled_generation: int | None = None + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.armed = threading.Event() + + def _send_loop(self, gen, sock, q, meter): + if meter is self._meters[1] and self.stalled_generation in (None, gen): + self.stalled_generation = gen + sock = _StallingSocket(sock, self.armed) + return super()._send_loop(gen, sock, q, meter) + + +def _sweep_until_final(client: SyncerClient, sweeps: int): + """Answer every PULL (also PULLs the syncer replays after a redial) with + the deterministic delta of its (fragment, global_step) until FINAL.""" + answered: set[tuple[int, int, int]] = set() + deadline = time.monotonic() + 60 + while not client.finalizing.is_set(): + assert time.monotonic() < deadline, "sweeps did not finish" + client.check_health() + for permit in client.drain_pulls(): + key = (client.connection_generation, permit.global_step, permit.round_attempt) + if key in answered: + continue + fid = permit.fragment_id + sweep = (permit.global_step - 1) // FRAGMENTS + 1 + assert permit.global_step == (sweep - 1) * FRAGMENTS + fid + 1 and sweep <= sweeps + base = 0 if sweep == 1 else permit.global_step - FRAGMENTS + delta = [0.5 + fid + 0.03125 * sweep + 0.001 * i for i in range(NUMEL)] + client.push_fragment(fid, permit.global_step, permit.round_attempt, + base, sweep, 1, 7, _f32(delta)) + answered.add(key) + time.sleep(0.01) + return client.wait_for_final_fragments(timeout=30) + + +@pytest.mark.parametrize("schedule", ["rr", "adaptive"]) +def test_stalled_stream_redials_and_the_replayed_push_keeps_the_exact_result( + syncer_binary, tmp_path, schedule): + """E3b: a data stream that stops moving bytes (C32 first run) is dropped + after the stall timeout; the syncer replays its PULL on the new group and + the island pushes that fragment again, so versions and finals match a run + without the stall bit for bit.""" + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + stall_dir = tmp_path / "stall" + stall_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, stall_dir, *STREAMED, total_steps=3) + client = _StallOnceClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), + stream_schedule=schedule, stream_stall_timeout=1.0) + try: + client.start() + for fid in range(FRAGMENTS): + client.send_init(fid, _f32([float(fid) + 0.125 * i for i in range(NUMEL)])) + # Stall only the pushes: INIT is sent once and not replayed after a redial. + _wait(lambda: client.send_progress().queued_chunks == 0 and client._bcasts.qsize() >= FRAGMENTS, + "initial broadcast") + client.armed.set() + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + progress = client.send_progress() + finally: + client.close() + _stop(process) + assert progress.stalls >= 1 and progress.reconnects >= 1 + assert list(manifest.versions) == exact[1] + assert result == exact[2] diff --git a/tests/test_wan_streams.py b/tests/test_wan_streams.py new file mode 100644 index 000000000..4e08efb0b --- /dev/null +++ b/tests/test_wan_streams.py @@ -0,0 +1,229 @@ +"""S26 E3b: per-stream accounting, stall watchdog and rate-aware striping.""" + +from __future__ import annotations + +import queue +import socket +import threading +import time + +import pytest + +from yeto.protocol import SyncerClient +from yeto.rl.engine.sync_timing import SyncTiming, upload_fields +from yeto.wan_streams import ( + StreamMeter, + StreamSnapshot, + UploadWindow, + pick_stream, + upload_window, +) + + +class _Clock: + def __init__(self) -> None: + self.now = 100.0 + + def __call__(self) -> float: + return self.now + + +def _meter_with_rate(clock: _Clock, rate_bps: float, backlog: int = 0) -> StreamMeter: + meter = StreamMeter(clock) + meter.enqueued(int(rate_bps)) + meter.begin() + clock.now += 1.0 + meter.progress(int(rate_bps)) + meter.end() + if backlog: + meter.enqueued(backlog) + return meter + + +def test_round_robin_ignores_rates(): + clock = _Clock() + meters = [_meter_with_rate(clock, 1e6), _meter_with_rate(clock, 1e3)] + assert [pick_stream("rr", meters, ticket) for ticket in range(4)] == [0, 1, 0, 1] + + +def test_adaptive_spreads_an_idle_start_over_every_stream(): + meters = [StreamMeter() for _ in range(4)] + picked = [] + for ticket in range(4): + index = pick_stream("adaptive", meters, ticket) + meters[index].enqueued(4 << 20) + picked.append(index) + assert sorted(picked) == [0, 1, 2, 3] + + +def test_adaptive_gives_a_slow_stream_fewer_chunks_in_proportion_to_its_rate(): + clock = _Clock() + fast, slow = 40e6, 5e6 + meters = [_meter_with_rate(clock, fast), _meter_with_rate(clock, slow)] + counts = [0, 0] + for ticket in range(90): + index = pick_stream("adaptive", meters, ticket) + meters[index].enqueued(4 << 20) + counts[index] += 1 + # Equal estimated finish time: chunk counts follow the 8:1 rate ratio. + assert counts[1] == pytest.approx(90 / 9, abs=1) + assert counts[0] == pytest.approx(90 * 8 / 9, abs=1) + + +def test_adaptive_starves_a_stream_whose_backlog_is_stuck(): + clock = _Clock() + meters = [_meter_with_rate(clock, 10e6), _meter_with_rate(clock, 10e6, backlog=64 << 20)] + assert {pick_stream("adaptive", meters, ticket) for ticket in range(10)} == {0} + + +def test_unknown_schedule_is_refused(): + with pytest.raises(ValueError, match="schedule"): + pick_stream("random", [StreamMeter()], 0) + with pytest.raises(ValueError, match="stream_schedule"): + SyncerClient(("127.0.0.1", 1), 0, _layout(), stream_schedule="random") + with pytest.raises(ValueError, match="stall_timeout"): + SyncerClient(("127.0.0.1", 1), 0, _layout(), stream_stall_timeout=0) + + +def test_meter_reports_a_stall_only_while_bytes_wait(): + clock = _Clock() + meter = StreamMeter(clock) + clock.now += 50 + assert meter.stalled_for() == 0.0 # idle is not stalled + meter.enqueued(10) + meter.begin() + meter.progress(4) + clock.now += 7 + assert meter.stalled_for() == pytest.approx(7) + meter.progress(6) + meter.end() + clock.now += 30 + assert meter.stalled_for() == 0.0 and meter.idle + meter.enqueued(5) + meter.drop_pending() + assert meter.idle + + +def test_upload_window_and_fields_report_rates_and_the_wait_after_upload(): + before = (StreamSnapshot(100, 1.0, 0, None), StreamSnapshot(0, 0.0, 0, None)) + after = (StreamSnapshot(100 + 40_000_000, 3.0, 0, 2e7), StreamSnapshot(5_000_000, 1.0, 0, 5e6)) + window = upload_window(before, after) + assert window == UploadWindow((40_000_000, 5_000_000), (2.0, 1.0)) + assert window.stream_mbps() == [20.0, 5.0] + fields = upload_fields(window, drained_ts=1010.0, timing=SyncTiming(1000.0, 1025.0), stalls=1) + assert fields["sync/upload_mbps"] == pytest.approx(22.5) + assert fields["sync/upload_min_stream_mbps"] == 5.0 + assert fields["sync/upload_active_streams"] == 2 + assert fields["sync/wait_after_upload_seconds"] == pytest.approx(15.0) + assert fields["sync/stream_stalls"] == 1 + # The upload drained before this boundary's submit: only the submit counts. + early = upload_fields(window, drained_ts=990.0, timing=SyncTiming(1000.0, 1025.0), stalls=0) + assert early["sync/wait_after_upload_seconds"] == pytest.approx(25.0) + unknown = upload_fields(window, drained_ts=None, timing=SyncTiming(1000.0, None), stalls=0) + assert unknown["sync/wait_after_upload_seconds"] is None + + +def test_sync_round_diagnostics_reach_wandb_on_the_policy_axis(): + from yeto.rl.wandb_rl import event_metrics + + metrics = event_metrics({ + "event": "rl_full_parameter_round", "policy_version": 3, + "sync/sync_wait_seconds": 40.0, "sync/upload_mbps": 22.5, + "sync/upload_min_stream_mbps": 5.0, "sync/upload_stream_mbps": [20.0, 5.0], + "sync/wait_after_upload_seconds": 15.0, "sync/global_policy_hash": "ab", + }) + assert metrics == {"rl/step": 3, "sync/sync_wait_seconds": 40.0, "sync/upload_mbps": 22.5, + "sync/upload_min_stream_mbps": 5.0, "sync/wait_after_upload_seconds": 15.0} + assert event_metrics({"event": "rl_full_parameter_round", "policy_version": 3}) is None + + +# -- socket-level behaviour --------------------------------------------------- + + +def _layout(): + from yeto.fragments import MERGE_RDA, Fragment, FragmentLayout + + return FragmentLayout([Fragment(MERGE_RDA, [("w", 4)])]) + + +def test_send_loop_meters_bytes_and_marks_the_upload_drained(): + left, right = socket.socketpair() + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1) + q: queue.Queue = queue.Queue() + meter = client._meters[1] + worker = threading.Thread(target=client._send_loop, args=(client._gen, left, q, meter)) + worker.start() + try: + meter.enqueued(3 << 20) + q.put(type("Item", (), {"data": b"x" * (3 << 20), "sent": None})()) + received = 0 + while received < 3 << 20: + received += len(right.recv(1 << 20)) + q.put(None) + worker.join(timeout=5) + snap = client.stream_snapshots()[0] + assert snap.bytes_sent == 3 << 20 and snap.backlog_bytes == 0 + assert client.upload_drained_at is not None + finally: + left.close() + right.close() + + +def test_watchdog_fails_the_group_when_a_stream_stops_moving(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=2, stream_stall_timeout=0.2) + a, b = socket.socketpair() + try: + client._socks = [a] + client._connected.set() + meter = client._meters[2] + meter.enqueued(10) + meter.begin() + assert client._stalled_stream() is None + time.sleep(0.3) + assert "outbound stream 2" in client._stalled_stream() + client._watchdog = threading.Thread(target=client._watch_stalls, daemon=True) + client._watchdog.start() + assert client._failure.wait(3), "the watchdog must fail the live group" + assert client.send_progress().stalls == 1 + assert "no progress" in str(client._last_err) + finally: + client._closed.set() + a.close() + b.close() + + +def test_watchdog_notices_an_inbound_message_that_stops_growing(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1, stream_stall_timeout=0.2) + client._reasm_touched[5] = time.monotonic() - 1.0 + assert "inbound message 5" in client._stalled_stream() + client._clear_partial_reassemblies() + assert client._stalled_stream() is None + + +def test_stall_timeout_sets_tcp_user_timeout(): + if not hasattr(socket, "TCP_USER_TIMEOUT"): + pytest.skip("Linux-only socket option") + from yeto.protocol import set_tcp_user_timeout + + with socket.socket() as sock: + set_tcp_user_timeout(sock, 2.5) + assert sock.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2500 + + +def test_transport_flags_reach_the_island_command_only_when_set(): + from argparse import Namespace + + from yeto.launch.rl_island_task import _wan_transport_flags + + assert _wan_transport_flags(Namespace()) == "" + assert _wan_transport_flags(Namespace(wan_stream_schedule="adaptive", wan_stall_timeout=60.0)) == ( + " --wan-stream-schedule adaptive --wan-stall-timeout 60.0") + with pytest.raises(ValueError, match="wan-stall-timeout"): + _wan_transport_flags(Namespace(wan_stall_timeout=-1.0)) + + +def test_island_parser_accepts_the_transport_flags(): + from yeto.rl.engine.run_flags import ISLAND, rl_option + + assert rl_option("wan_stream_schedule", ISLAND) == "--wan-stream-schedule" + assert rl_option("wan_stall_timeout", ISLAND) == "--wan-stall-timeout" diff --git a/yeto/cli.py b/yeto/cli.py index 7c68497f8..40abdc1f0 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -438,6 +438,22 @@ def int_or_auto(value: str): "deltas (~4x less learner egress; broadcasts stay bf16)", ) sync.add_argument("--wan-streams", type=int, default=4, help="parallel TCP streams per learner") + sync.add_argument( + "--wan-stream-schedule", + choices=("rr", "adaptive"), + default="rr", + help="full-parameter RL islands: rr = round-robin chunks over the streams; " + "adaptive = send each chunk to the stream with the earliest estimated " + "finish (backlog / measured rate), so slow streams carry less", + ) + sync.add_argument( + "--wan-stall-timeout", + type=float, + default=0.0, + help="full-parameter RL islands: drop and redial the connection group when " + "one stream holds bytes with no progress for this many seconds (also sets " + "TCP_USER_TIMEOUT); the fragment is pushed again on the replayed PULL. 0 = off", + ) infra = p.add_argument_group("infrastructure") infra.add_argument( diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 532fff20e..a51d6ea1e 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -51,6 +51,20 @@ def submitted_unix() -> float: return time.time() +def _wan_transport_flags(args) -> str: + """S26 E3b island transport options; emitted only when not the default.""" + flags = "" + schedule = getattr(args, "wan_stream_schedule", "rr") + if schedule != "rr": + flags += f" --wan-stream-schedule {schedule}" + stall = float(getattr(args, "wan_stall_timeout", 0.0) or 0.0) + if stall < 0: + raise ValueError("--wan-stall-timeout must be >= 0 (0 = off)") + if stall > 0: + flags += f" --wan-stall-timeout {stall!r}" + return flags + + def _rl_trainable_island_flags(args, learner_id: int, num_learners: int) -> str: """LoRA rank/targets; for the full-parameter preset the sweep identity (the probed layout hash, the frozen roster size and the syncer profile hash); @@ -166,7 +180,8 @@ def make_miles_island_task( if float(getattr(args, "rl_optimizer_offload_fraction", 0.0) or 0.0) > 0 else "") + f" --seed {args.seed}" f" --wan-streams {args.wan_streams}" - " --miles-root ~/miles" + + _wan_transport_flags(args) + + " --miles-root ~/miles" ) if getattr(args, "wandb", False): flags += " --wandb" diff --git a/yeto/protocol.py b/yeto/protocol.py index fc2b6d73a..42cfd3f89 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -41,6 +41,7 @@ from .downlink import DownlinkAnchors, DownlinkChunkSink from .downlink import tensor_nbytes as downlink_tensor_nbytes from .fragments import MERGE_ISO, FragmentLayout +from .wan_streams import STREAM_SCHEDULES, StreamMeter, StreamSnapshot, StreamStalledError, pick_stream MAGIC = 0xD170C0DE PROTOCOL_VERSION = 4 @@ -571,6 +572,7 @@ class SendProgress: queued_chunks: int reconnects: int connected: bool + stalls: int = 0 @dataclass @@ -594,6 +596,10 @@ class SyncerClient: forever); once exhausted, check_health raises. """ + # Transport defaults (S26 E3b); __init__ overrides them per client. + stream_schedule: str = "rr" + stream_stall_timeout: float | None = None + def __init__( self, addr: tuple[str, int], @@ -609,7 +615,13 @@ def __init__( hmac_key: bytes | None = None, compat_group: str | None = None, push_i8_block: int = DEFAULT_I8_BLOCK, + stream_schedule: str = "rr", + stream_stall_timeout: float | None = None, ): + if stream_schedule not in STREAM_SCHEDULES: + raise ValueError(f"stream_schedule must be one of {STREAM_SCHEDULES}, got {stream_schedule!r}") + if stream_stall_timeout is not None and not stream_stall_timeout > 0: + raise ValueError("stream_stall_timeout must be positive or None") if not 0 <= learner_id <= 0xFFFF_FFFF: raise ValueError(f"learner_id must fit u32, got {learner_id}") if not valid_session_dtype(dtype): @@ -638,6 +650,15 @@ def __init__( self._downlink = DownlinkAnchors(layout.num_fragments) if bcast_delta(dtype) else None self._bcast_prefix = 12 + (DOWNLINK_HEADER_LEN if self._downlink is not None else 0) self.num_streams = num_streams + # S26 E3b: chunk placement over the data streams and the per-stream + # no-progress timeout (None = off, the historical behaviour). + self.stream_schedule = stream_schedule + self.stream_stall_timeout = stream_stall_timeout + # Index 0 is the control socket; the meters outlive connection groups + # so their counters are cumulative for the whole session. + self._meters = [StreamMeter() for _ in range(num_streams + 1)] + self._upload_drained_at: float | None = None + self._stalls = 0 self.session_contract_hash = session_contract_hash self.syncer_profile_hash = syncer_profile_hash # secret-handling-hardening: HELLO carries an HMAC when the island has @@ -682,6 +703,7 @@ def __init__( | _StreamedReassembly, ] = {} self._reasm_lock = threading.Lock() + self._reasm_touched: dict[int, float] = {} # msg_id -> last chunk (monotonic) self._inbound_chunk_sink: object | None = None self._msg_id = itertools.count() self._rr = itertools.count() @@ -702,6 +724,7 @@ def __init__( self._bytes_sent_lock = threading.Lock() self._reset_bcasts_on_reconnect = False self._supervisor: threading.Thread | None = None + self._watchdog: threading.Thread | None = None # -- lifecycle ------------------------------------------------------------- @@ -710,6 +733,8 @@ def _dial(self, timeout: float) -> socket.socket: sock.settimeout(None) sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) enable_tcp_keepalive(sock) + if self.stream_stall_timeout is not None: + set_tcp_user_timeout(sock, self.stream_stall_timeout) return sock def _connect_one(self) -> socket.socket: @@ -740,6 +765,9 @@ def start(self) -> None: self._connect_group(patient=True) self._supervisor = threading.Thread(target=self._supervise, name="yeto-supervisor", daemon=True) self._supervisor.start() + if self.stream_stall_timeout is not None: + self._watchdog = threading.Thread(target=self._watch_stalls, name="yeto-stall-watchdog", daemon=True) + self._watchdog.start() def _connect_group(self, patient: bool) -> None: """Dial control + data sockets, send HELLO/DATA_HELLO, start workers. @@ -797,7 +825,7 @@ def _connect_group(self, patient: bool) -> None: q: queue.Queue[_Outbound | None] = queue.Queue(maxsize=256) self._queues.append(q) ts = threading.Thread( - target=self._send_loop, args=(gen, s, q), name=f"yeto-send-{i}", daemon=True + target=self._send_loop, args=(gen, s, q, self._meters[i]), name=f"yeto-send-{i}", daemon=True ) tr = threading.Thread(target=self._recv_loop, args=(gen, s), name=f"yeto-recv-{i}", daemon=True) self._threads += [ts, tr] @@ -822,6 +850,8 @@ def _teardown_group(self) -> None: pass for s in socks: _close_socket(s) + for meter in self._meters: + meter.drop_pending() self._clear_partial_reassemblies() self._drain(self._pulls) # stale pull requests; answering them is pointless with self._lock: @@ -914,8 +944,21 @@ def send_progress(self) -> SendProgress: queued_chunks=queued, reconnects=self._reconnects_used, connected=connected, + stalls=self._stalls, ) + def stream_snapshots(self) -> tuple[StreamSnapshot, ...]: + """Cumulative counters of each data stream (control stream excluded; + with ``num_streams == 0`` the control stream carries the data).""" + meters = self._meters[1:] if self.num_streams else self._meters[:1] + return tuple(meter.snapshot() for meter in meters) + + @property + def upload_drained_at(self) -> float | None: + """Unix time when the data streams last became idle with nothing + queued, i.e. when this island's latest upload left the host.""" + return self._upload_drained_at + # -- learner-facing API ------------------------------------------------------ def send_init(self, fragment_id: int, tensor_bytes: bytes) -> None: @@ -1279,6 +1322,8 @@ def _enqueue( if gen is not None and gen != self._gen: return False q = self._queues[stream] + meter = self._meters[stream] + meter.enqueued(len(data)) # before put: the sender may drain it at once try: q.put( _Outbound(data, sent), @@ -1286,8 +1331,16 @@ def _enqueue( ) # bounded wait, then re-check the group return True except queue.Full: + meter.enqueued(-len(data)) continue + def _next_data_stream(self) -> int: + """Socket index (1-based; 0 only without data streams) for the next + CHUNK envelope, by ``stream_schedule``.""" + if self.num_streams == 0: + return 0 + return 1 + pick_stream(self.stream_schedule, self._meters[1:], next(self._rr)) + def _send_large(self, msg_type: int, payload: bytes) -> None: inner = self._frame(msg_type, payload) if self.num_streams == 0: @@ -1302,7 +1355,7 @@ def _send_large(self, msg_type: int, payload: bytes) -> None: for offset in range(0, total, CHUNK_SIZE): chunk = inner[offset : offset + CHUNK_SIZE] envelope = self._frame(MSG_CHUNK, _CHUNK_HEAD.pack(msg_id, total, offset) + chunk) - stream = 1 + next(self._rr) % self.num_streams + stream = self._next_data_stream() if not self._enqueue(stream, envelope, gen=gen): return # group died mid-message; drop the remainder @@ -1423,11 +1476,7 @@ def _send_large_parts( before_last_enqueue = None callback() - stream = ( - 0 - if self.num_streams == 0 - else 1 + next(self._rr) % self.num_streams - ) + stream = self._next_data_stream() if not self._enqueue(stream, envelope, gen=gen): if enqueued_chunks: raise PartialMessageGenerationLost( @@ -1487,20 +1536,85 @@ def _protocol_failed(self, gen: int, exc: BaseException) -> None: with self._final_cond: self._final_cond.notify_all() - def _send_loop(self, gen: int, sock: socket.socket, q: queue.Queue) -> None: + def _send_loop( + self, + gen: int, + sock: socket.socket, + q: queue.Queue, + meter: StreamMeter | None = None, + ) -> None: + if meter is None: # a sender outside a connection group (tests) + meter = StreamMeter() try: while True: item = q.get() if item is None: return - sock.sendall(item.data) + # send() instead of sendall() so the meter sees byte progress + # inside one 4 MiB envelope (stall watchdog, per-stream rate). + view = memoryview(item.data) + meter.begin() + try: + while view: + sent = sock.send(view) + meter.progress(sent) + view = view[sent:] + finally: + meter.end() + view.release() with self._bytes_sent_lock: self._bytes_sent += len(item.data) + if all(m.idle for m in self._meters): + self._upload_drained_at = time.time() if item.sent is not None: item.sent.set() except BaseException as e: self._socket_failed(gen, e) + def _stalled_stream(self) -> str | None: + """Describe the first stream without progress for the stall timeout.""" + limit = self.stream_stall_timeout + now = time.monotonic() + for index, meter in enumerate(self._meters): + idle_for = meter.stalled_for(now) + if idle_for > limit: + return f"outbound stream {index} made no progress for {idle_for:.0f}s" + with self._reasm_lock: + touched = dict(self._reasm_touched) + for msg_id, last in touched.items(): + if now - last > limit: + return f"inbound message {msg_id} got no chunk for {now - last:.0f}s" + return None + + def _watch_stalls(self) -> None: + """Fail the live group when one stream stops moving bytes. + + A stalled TCP stream can sit for many minutes before the kernel gives + up (C32: 900 s with every socket still open). The failure goes + through the normal reconnect path, so the syncer replays its PULLs + and the island pushes the interrupted fragment again.""" + period = min(1.0, self.stream_stall_timeout / 4) + while not self._closed.wait(period): + if self.shutdown.is_set(): + return + with self._lock: + gen = self._gen + live = self._connected.is_set() + if not live: + continue + reason = self._stalled_stream() + if reason is None: + continue + with self._lock: + if gen != self._gen: + continue + self._stalls += 1 + socks = list(self._socks) + _connect_log(f"{reason} (learner {self.learner_id}); dropping the connection group") + self._socket_failed(gen, StreamStalledError(reason)) + for sock in socks: # unblock senders stuck in send() + _close_socket(sock) + def _recv_loop(self, gen: int, sock: socket.socket) -> None: try: while True: @@ -1569,6 +1683,7 @@ def _reassemble( sink=sink, transaction=sink.begin_message(msg_id, total), ) + self._reasm_touched[msg_id] = time.monotonic() partial = self._reasm[msg_id] if isinstance(partial, _StreamedReassembly): if partial.total != total: @@ -1588,6 +1703,7 @@ def _reassemble( raise if completed is not None: self._reasm.pop(msg_id, None) + self._reasm_touched.pop(msg_id, None) return completed buf, filled, ranges = partial if len(buf) != total: @@ -1602,6 +1718,7 @@ def _reassemble( if filled[0] != total: raise ValueError("chunk byte count exceeds frame length") del self._reasm[msg_id] + self._reasm_touched.pop(msg_id, None) magic, msg_type, length = _HEADER.unpack_from(buf) if magic != MAGIC or length != total - _HEADER.size: raise ValueError("corrupt reassembled frame") @@ -1769,6 +1886,7 @@ def _clear_partial_reassemblies(self) -> None: with self._reasm_lock: partials = list(self._reasm.values()) self._reasm.clear() + self._reasm_touched.clear() for partial in partials: if isinstance(partial, _StreamedReassembly): try: @@ -2043,6 +2161,14 @@ def enable_tcp_keepalive(sock: socket.socket) -> None: sock.setsockopt(socket.IPPROTO_TCP, option, value) +def set_tcp_user_timeout(sock: socket.socket, seconds: float) -> None: + """Bound how long unacknowledged data may sit before the kernel resets + the connection (Linux TCP_USER_TIMEOUT); other platforms keep defaults.""" + option = getattr(socket, "TCP_USER_TIMEOUT", None) + if option is not None: + sock.setsockopt(socket.IPPROTO_TCP, option, max(1, int(seconds * 1000))) + + def _close_socket(sock: socket.socket) -> None: try: sock.shutdown(socket.SHUT_RDWR) diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index 1855ad780..fa746fcb0 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -88,6 +88,10 @@ class RoleStreamConfig: # P0-4.8: the island's AlgorithmSpec hash; placement capabilities stay out. algorithm_spec_sha256: str wan_streams: int = 4 + # S26 E3b transport options (yeto/wan_streams.py); not part of the session + # contract: they change how bytes travel, never which bytes. + wan_stream_schedule: str = "rr" + wan_stall_timeout: float = 0.0 # seconds; 0 = off wait_timeout: float = 900.0 poll_seconds: float = 0.01 max_fragment_bytes: int = 2 << 30 @@ -158,6 +162,10 @@ def __post_init__(self) -> None: value = getattr(self, name) if type(value) is not int or value < 1: raise ValueError(f"role-stream {name} must be positive") + if self.wan_stream_schedule not in ("rr", "adaptive"): + raise ValueError("role-stream wan stream schedule must be rr or adaptive") + if not self.wan_stall_timeout >= 0: + raise ValueError("role-stream wan stall timeout must be >= 0") if self.push_dtype not in ("f32", "bf16", "int8"): raise ValueError("role-stream push dtype must be f32, bf16 or int8") if type(self.push_int8_block) is not int or not I8_BLOCK_MIN <= self.push_int8_block <= I8_BLOCK_MAX: @@ -379,6 +387,8 @@ async def initialize(self, group) -> None: dtype=self.config.session_dtype, push_i8_block=self.config.push_int8_block, num_streams=self.config.wan_streams, + stream_schedule=self.config.wan_stream_schedule, + stream_stall_timeout=self.config.wan_stall_timeout or None, max_reconnects=None, session_contract_hash=role_stream_session_contract_hash( adapter.layout.layout_hash, diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 6ff6a2e84..01367ff35 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -838,6 +838,8 @@ def _dense_full_sync_config( pipeline_depth=1, algorithm_spec_sha256=algorithm_spec_sha256, wan_streams=args.wan_streams, + wan_stream_schedule=getattr(args, "wan_stream_schedule", "rr"), + wan_stall_timeout=float(getattr(args, "wan_stall_timeout", 0.0) or 0.0), push_dtype=getattr(args, "sweep_push_dtype", "f32"), push_int8_block=int(getattr(args, "sweep_push_int8_block", 128)), bcast_delta=bool(getattr(args, "bcast_delta", False)), diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index ff2830bcd..307788097 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -416,6 +416,8 @@ island_flag("--seq-len", type=int, required=True), island_flag("--seed", type=int, required=True), island_flag("--wan-streams", type=int, default=4), + island_flag("--wan-stream-schedule", choices=("rr", "adaptive"), default="rr"), + island_flag("--wan-stall-timeout", type=float, default=0.0), island_flag("--miles-root", required=True), island_flag( "--miles-commit", diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index 744f30535..af2a1a1da 100644 --- a/yeto/rl/engine/full_parameter_sync.py +++ b/yeto/rl/engine/full_parameter_sync.py @@ -33,9 +33,10 @@ from yeto.protocol import PartialMessageGenerationLost from yeto.rl.elastic_client import LeaveReason +from yeto.wan_streams import upload_window from .driver import SyncBoundary, SyncStart -from .sync_timing import SyncTiming, unix_now +from .sync_timing import SyncTiming, unix_now, upload_fields if TYPE_CHECKING: from .driver import IslandDriver @@ -529,6 +530,8 @@ def __init__( self._timing = SyncTiming(None, None) # L5.6, set by each boundary # Overlap: fragment versions before the sweep still in flight. self._pending_sweep: tuple[int, ...] | None = None + # S26 E3b: actor data-stream counters at the previous boundary. + self._upload_before: tuple | None = None @property def actor(self) -> RoleStream: @@ -537,6 +540,7 @@ def actor(self) -> RoleStream: def start(self, driver: IslandDriver) -> SyncStart: self._driver = driver self.run(self.pairing.gather_roles(lambda role, stream: stream.initialize(self.groups[role]))) + self._upload_before = self._stream_snapshots() return SyncStart(state=self.actor.resident_state(0), rollout_id=0) def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> SyncBoundary: @@ -554,10 +558,29 @@ def boundary(self, driver: IslandDriver, *, rollout_id: int, stats: Any) -> Sync "sync/global_policy_hash": state.policy_tensor_hash(), "sync/inflight_pairs": len(self.pairing.inflight), **self._timing.fields(), + **self._upload_fields(), }, ) return SyncBoundary(state=state, stop=self.terminal) + def _stream_snapshots(self) -> tuple | None: + snapshots = getattr(getattr(self.actor, "client", None), "stream_snapshots", None) + return snapshots() if callable(snapshots) else None + + def _upload_fields(self) -> dict[str, object]: + """Actor upload rates of this boundary and the wait after the upload.""" + after = self._stream_snapshots() + before, self._upload_before = self._upload_before, after + if after is None or before is None: + return {} + client = self.actor.client + return upload_fields( + upload_window(before, after), + drained_ts=client.upload_drained_at, + timing=self._timing, + stalls=client.send_progress().stalls, + ) + async def _boundary(self, rollout_id: int) -> bool: if self.pairing.profile.overlap: return await self._overlap_boundary(rollout_id) diff --git a/yeto/rl/engine/sync_timing.py b/yeto/rl/engine/sync_timing.py index eaa4ad9b9..eae22025c 100644 --- a/yeto/rl/engine/sync_timing.py +++ b/yeto/rl/engine/sync_timing.py @@ -44,3 +44,42 @@ def fields(self) -> dict[str, float | None]: def unix_now() -> float: return time.time() + + +def upload_fields( + window, + *, + drained_ts: float | None, + timing: SyncTiming, + stalls: int, +) -> dict[str, object]: + """Per-boundary upload diagnostics of one island (S26 E3b). + + ``window`` is a :class:`yeto.wan_streams.UploadWindow` covering this + boundary. ``sync/upload_stream_mbps`` lists each data stream's rate over + the time it was busy; ``sync/upload_mbps`` is the bytes over the longest + busy stream (the streams run in parallel). ``sync/wait_after_upload_seconds`` + is the time from this island's last byte leaving to the global policy + being ready: waiting for slower islands, the merge and the broadcast.""" + rates = window.stream_mbps() + busy = max(window.busy_seconds_per_stream, default=0.0) + measured = [rate for rate in rates if rate is not None] + wait_after = None + if ( + drained_ts is not None + and timing.submit_ts is not None + and timing.global_ready_ts is not None + and drained_ts <= timing.global_ready_ts + ): + wait_after = timing.global_ready_ts - max(drained_ts, timing.submit_ts) + return { + "sync/upload_bytes": window.total_bytes, + "sync/upload_busy_seconds": busy, + "sync/upload_mbps": (window.total_bytes / busy / 1e6) if busy > 0 else None, + "sync/upload_stream_mbps": [None if rate is None else round(rate, 3) for rate in rates], + "sync/upload_min_stream_mbps": min(measured) if measured else None, + "sync/upload_active_streams": len(measured), + "sync/upload_drained_ts": drained_ts, + "sync/wait_after_upload_seconds": wait_after, + "sync/stream_stalls": stalls, + } diff --git a/yeto/rl/wandb_rl.py b/yeto/rl/wandb_rl.py index 3467c503d..60a08c7b0 100644 --- a/yeto/rl/wandb_rl.py +++ b/yeto/rl/wandb_rl.py @@ -28,6 +28,7 @@ "train/*": TRAIN_STEP_KEY, "rl/*": STEP_KEY, "eval/*": STEP_KEY, + "sync/*": STEP_KEY, } # Closed output surface. Input aliases are limited to namespaced scalar fields @@ -115,11 +116,40 @@ def _non_negative_step(event: dict[str, Any], aliases: tuple[str, ...]) -> int | return None +# S26 E3b: per-sweep sync diagnostics of a full-parameter island, on the +# global policy-version axis (rl/step). +_SYNC_ROUND_OUTPUTS = ( + "sync/sync_wait_seconds", + "sync/upload_mbps", + "sync/upload_min_stream_mbps", + "sync/upload_active_streams", + "sync/wait_after_upload_seconds", + "sync/stream_stalls", +) + + +def _sync_round_metrics(event: dict[str, Any]) -> dict[str, Any] | None: + step = _non_negative_step(event, ("policy_version",)) + if step is None: + return None + metrics = { + key: value + for key in _SYNC_ROUND_OUTPUTS + if (value := _finite_scalar(event, (key,))) is not None + } + if not metrics: + return None + metrics[STEP_KEY] = step + return metrics + + def event_metrics(event: dict[str, Any]) -> dict[str, Any] | None: """Project one event onto the explicit scalar-only W&B surface.""" global _derived_step_logged event_name = event.get("event") + if event_name == "rl_full_parameter_round": + return _sync_round_metrics(event) if event_name not in {"rl_local_round", "rl_eval_result", "rl_round_trained"}: return None diff --git a/yeto/wan_streams.py b/yeto/wan_streams.py new file mode 100644 index 000000000..8d2cb6abb --- /dev/null +++ b/yeto/wan_streams.py @@ -0,0 +1,190 @@ +"""Per-stream accounting, stall detection and rate-aware striping for the +island side of a striped syncer connection (S26 E3b). + +A ``SyncerClient`` stripes one logical message over ``num_streams`` TCP data +sockets as 4 MiB CHUNK envelopes. The syncer reassembles a message by +``(msg_id, offset)``, so any chunk may travel on any data stream of the same +connection group. This module holds the parts of that transport that do not +need a socket: + +* :class:`StreamMeter` counts bytes, busy time and the last byte progress of + one stream. The send loop feeds it; diagnostics and the stall watchdog read + it. +* :func:`pick_stream` chooses the data stream for the next chunk. ``"rr"`` + is the historical round robin. ``"adaptive"`` sends each chunk to the + stream with the earliest estimated finish time (backlog / measured rate), + so a slow stream gets fewer chunks and a stalled stream gets none. The + number of open streams is fixed by HELLO; the number of streams that carry + data follows the measured rates. +* :class:`StreamStalledError` is the failure a watchdog reports when a + stream holds unsent bytes and makes no byte progress for the stall timeout. + The client handles it like any socket error: the whole group is torn down + and redialed, and the syncer replays its PULLs, so the island pushes the + interrupted fragment again at the same (global_step, round_attempt). +""" + +from __future__ import annotations + +import threading +import time +from collections.abc import Sequence +from dataclasses import dataclass + +STREAM_SCHEDULES = ("rr", "adaptive") + +# EWMA weight of the newest completed send in the per-stream rate estimate. +RATE_EWMA_ALPHA = 0.3 +# Sends shorter than this carry no useful rate signal (kernel buffer absorbs +# them), so they do not update the estimate. +MIN_RATE_SAMPLE_SECONDS = 0.01 + + +class StreamStalledError(ConnectionError): + """One stream held unsent (or partially received) bytes with no progress + for longer than the stall timeout.""" + + +@dataclass(frozen=True) +class StreamSnapshot: + """Cumulative counters of one stream at one instant.""" + + bytes_sent: int + busy_seconds: float + backlog_bytes: int + rate_bps: float | None + + +class StreamMeter: + """Counters of one outbound stream. Thread-safe; cheap enough to update + on every ``send`` call.""" + + def __init__(self, clock=time.monotonic) -> None: + self._clock = clock + self._lock = threading.Lock() + self.bytes_sent = 0 + self.busy_seconds = 0.0 + self.queued_bytes = 0 # enqueued, not yet handed to the kernel + self.rate_bps: float | None = None + self._item_started: float | None = None + self._item_bytes = 0 + self._last_progress = clock() + + def enqueued(self, nbytes: int) -> None: + with self._lock: + self.queued_bytes += nbytes + + def begin(self) -> None: + with self._lock: + now = self._clock() + self._item_started = now + self._item_bytes = 0 + self._last_progress = now + + def progress(self, nbytes: int) -> None: + with self._lock: + self.bytes_sent += nbytes + self.queued_bytes -= nbytes + self._item_bytes += nbytes + self._last_progress = self._clock() + + def end(self) -> None: + with self._lock: + if self._item_started is None: + return + elapsed = self._clock() - self._item_started + self.busy_seconds += elapsed + if elapsed >= MIN_RATE_SAMPLE_SECONDS and self._item_bytes: + sample = self._item_bytes / elapsed + self.rate_bps = ( + sample + if self.rate_bps is None + else (1 - RATE_EWMA_ALPHA) * self.rate_bps + RATE_EWMA_ALPHA * sample + ) + self._item_started = None + + def drop_pending(self) -> None: + """The connection group died: whatever was queued is gone.""" + with self._lock: + self.queued_bytes = 0 + self._item_started = None + self._item_bytes = 0 + + @property + def idle(self) -> bool: + with self._lock: + return self._item_started is None and self.queued_bytes <= 0 + + def stalled_for(self, now: float | None = None) -> float: + """Seconds without byte progress while bytes are waiting; 0 when idle.""" + with self._lock: + if self._item_started is None and self.queued_bytes <= 0: + return 0.0 + return max(0.0, (self._clock() if now is None else now) - self._last_progress) + + def snapshot(self) -> StreamSnapshot: + with self._lock: + busy = self.busy_seconds + if self._item_started is not None: + busy += self._clock() - self._item_started + return StreamSnapshot(self.bytes_sent, busy, max(0, self.queued_bytes), self.rate_bps) + + +def pick_stream( + schedule: str, + meters: Sequence[StreamMeter], + rr_ticket: int, +) -> int: + """Index into ``meters`` (the data streams) for the next chunk. + + ``adaptive`` minimizes ``backlog / rate``. A stream without a rate sample + yet uses the mean measured rate (or 1 B/s when nothing is measured, which + reduces the rule to least backlog). Ties go to the round-robin ticket so + an idle start still spreads chunks over every stream.""" + count = len(meters) + if count == 0: + raise ValueError("no data streams") + if schedule == "rr": + return rr_ticket % count + if schedule != "adaptive": + raise ValueError(f"unknown stream schedule {schedule!r}") + snaps = [meter.snapshot() for meter in meters] + known = [snap.rate_bps for snap in snaps if snap.rate_bps] + default_rate = sum(known) / len(known) if known else 1.0 + best = None + for step in range(count): + index = (rr_ticket + step) % count + snap = snaps[index] + finish = snap.backlog_bytes / (snap.rate_bps or default_rate) + if best is None or finish < best[0]: + best = (finish, index) + return best[1] + + +@dataclass(frozen=True) +class UploadWindow: + """Data-stream traffic between two snapshots (one sync boundary).""" + + bytes_per_stream: tuple[int, ...] + busy_seconds_per_stream: tuple[float, ...] + + @property + def total_bytes(self) -> int: + return sum(self.bytes_per_stream) + + def stream_mbps(self) -> list[float | None]: + """MB/s (1e6 B) of each stream over the time it was busy.""" + return [ + (b / s / 1e6) if s > 0 and b > 0 else None + for b, s in zip(self.bytes_per_stream, self.busy_seconds_per_stream, strict=True) + ] + + +def upload_window( + before: Sequence[StreamSnapshot], after: Sequence[StreamSnapshot] +) -> UploadWindow: + if len(before) != len(after): # the group was redialed with another width + before = [StreamSnapshot(0, 0.0, 0, None)] * len(after) + return UploadWindow( + tuple(max(0, a.bytes_sent - b.bytes_sent) for b, a in zip(before, after, strict=True)), + tuple(max(0.0, a.busy_seconds - b.busy_seconds) for b, a in zip(before, after, strict=True)), + ) From 2509fdefc62af5e6f41c06084ec4584b249eccd1 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:17:17 +0000 Subject: [PATCH 8/9] wan: replay INIT after a mid-init redial, link probe with container relaunch s26-e3b-p06-C32a: learner 0's connection group died before the syncer applied its INIT, the client never sent INIT again, and no island got the initial cut. s26-e3b runs also showed that a slow island is slow on every stream (container level), which no stream scheduling fixes. - SyncerClient.send_initial_fragments/init_replay_due/ mark_initial_cut_received: learner 0 replays INIT on every new connection group until the initial cut arrives (the syncer ignores INIT it already applied). TCP_USER_TIMEOUT is set only after the initial cut, so the multi-GB INIT and first broadcast are never reset by the kernel. - MSG_BW_PROBE (26): an HMAC-sealed probe connection to the legacy syncer uploads N bytes, gets an ack, downloads M bytes; never part of a session, at most 64 concurrent, 1 GiB per direction, 300 s. - --wan-probe-min-mbps / --wan-probe-min-down-mbps (default off): Miles islands probe the actor syncer before loading the model, record rl_wan_probe, and exit 75 below the minimum. The fleet controller requests a new container up to --wan-probe-relaunches times (default 2), also under --no-island-relaunch, since the island never joined the session. Real-syncer test: INIT interrupted between fragments is replayed and the finals match an undisturbed run bit for bit; the probe works on the same port and a wrong key is refused. Co-Authored-By: Claude Opus 5.5 --- syncer/src/protocol.rs | 4 + syncer/src/server.rs | 122 +++++++++++++- tests/test_e3_overlap_miles_role_stream.py | 9 ++ tests/test_rl_dense_sweep_real_syncer.py | 64 ++++++++ tests/test_wan_probe.py | 126 +++++++++++++++ tests/test_wan_streams.py | 32 ++++ yeto/cli.py | 23 +++ yeto/launch/fleet.py | 47 +++++- yeto/launch/modal_islands.py | 12 ++ yeto/launch/rl_island_task.py | 9 ++ yeto/launcher.py | 7 +- yeto/protocol.py | 59 ++++++- .../miles/full_parameter/role_stream.py | 31 ++-- yeto/rl/adapters/miles/island_entry.py | 35 ++++ yeto/rl/adapters/miles/run_flags.py | 3 + yeto/rl/event_echo.py | 8 +- yeto/wan_probe.py | 149 ++++++++++++++++++ 17 files changed, 716 insertions(+), 24 deletions(-) create mode 100644 tests/test_wan_probe.py create mode 100644 yeto/wan_probe.py diff --git a/syncer/src/protocol.rs b/syncer/src/protocol.rs index d9bd34bc3..e214faf2d 100644 --- a/syncer/src/protocol.rs +++ b/syncer/src/protocol.rs @@ -45,6 +45,10 @@ pub const MSG_ELASTIC_BASE: u8 = 23; pub const MSG_SAMPLE_VERDICT: u8 = 24; /// Elastic: run finished (final outer_version + base policy hash). pub const MSG_FINISHED: u8 = 25; +/// S26 E3b: link-bandwidth probe. First and only frame of its own connection +/// (never part of a learner session): the island uploads N bytes, the syncer +/// acknowledges with the u64 byte count, then streams M bytes back. +pub const MSG_BW_PROBE: u8 = 26; /// Version of the final-artifact handshake carried inside FINAL_MANIFEST and /// FINAL_ACK. Keeping this in the payload makes incompatible peers fail with diff --git a/syncer/src/server.rs b/syncer/src/server.rs index 269c0b446..c03011a33 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -1260,8 +1260,8 @@ async fn handle_connection( stream.set_nodelay(true)?; let (mut rd, mut wr) = stream.into_split(); let first = match read_frame_limited(&mut rd, |msg_type| match msg_type { - MSG_HELLO | MSG_DATA_HELLO => Ok(MAX_HELLO_FRAME), - other => bail!("first frame must be HELLO/DATA_HELLO, got {other}"), + MSG_HELLO | MSG_DATA_HELLO | MSG_BW_PROBE => Ok(MAX_HELLO_FRAME), + other => bail!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {other}"), }) .await { @@ -1466,8 +1466,98 @@ async fn handle_connection( tokio::spawn(writer_task(wr, rx)); read_loop(&mut rd, &group, &event_tx, &budget_cutoff).await } - t => bail!("first frame must be HELLO/DATA_HELLO, got {t}"), + MSG_BW_PROBE => { + let request = open_hello(hello_auth.key.as_deref(), &first.payload) + .and_then(parse_bw_probe); + let (upload, download) = match request { + Ok(sizes) => sizes, + Err(error) => { + let message = format!("invalid BW_PROBE: {error:#}"); + let _ = send_direct(&mut wr, MSG_ERROR, message.as_bytes()).await; + return Err(error); + } + }; + let Some(_slot) = BwProbeSlot::acquire() else { + let message = "too many concurrent BW_PROBE connections"; + send_direct(&mut wr, MSG_ERROR, message.as_bytes()).await?; + bail!(message); + }; + tokio::time::timeout(BW_PROBE_TIMEOUT, serve_bw_probe(&mut rd, &mut wr, upload, download)) + .await + .context("BW_PROBE timed out")? + } + t => bail!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {t}"), + } +} + +/// Largest upload or download one BW_PROBE connection may ask for. +const BW_PROBE_MAX_BYTES: u64 = 1 << 30; +const BW_PROBE_MAX_CONNECTIONS: usize = 64; +const BW_PROBE_TIMEOUT: Duration = Duration::from_secs(300); +static BW_PROBE_ACTIVE: AtomicUsize = AtomicUsize::new(0); + +/// One of at most BW_PROBE_MAX_CONNECTIONS concurrent probe connections. +struct BwProbeSlot; + +impl BwProbeSlot { + fn acquire() -> Option { + let previous = BW_PROBE_ACTIVE.fetch_add(1, Ordering::AcqRel); + if previous >= BW_PROBE_MAX_CONNECTIONS { + BW_PROBE_ACTIVE.fetch_sub(1, Ordering::AcqRel); + return None; + } + Some(BwProbeSlot) + } +} + +impl Drop for BwProbeSlot { + fn drop(&mut self) { + BW_PROBE_ACTIVE.fetch_sub(1, Ordering::AcqRel); + } +} + +/// BW_PROBE body (after the HELLO-style HMAC trailer is removed): +/// u64 upload bytes, u64 download bytes, both at most BW_PROBE_MAX_BYTES. +fn parse_bw_probe(body: &[u8]) -> Result<(u64, u64)> { + if body.len() != 16 { + bail!("BW_PROBE body has {} bytes, expected 16", body.len()); + } + let upload = u64::from_le_bytes(body[0..8].try_into().unwrap()); + let download = u64::from_le_bytes(body[8..16].try_into().unwrap()); + if upload > BW_PROBE_MAX_BYTES || download > BW_PROBE_MAX_BYTES { + bail!("BW_PROBE asks for more than {BW_PROBE_MAX_BYTES} bytes per direction"); } + Ok((upload, download)) +} + +/// Read and discard `upload` raw bytes, acknowledge with their count, then +/// write `download` zero bytes and close. Nothing touches session state. +async fn serve_bw_probe( + rd: &mut (impl tokio::io::AsyncReadExt + Unpin), + wr: &mut OwnedWriteHalf, + upload: u64, + download: u64, +) -> Result<()> { + let mut buf = vec![0u8; 1 << 20]; + let mut remaining = upload; + while remaining > 0 { + let want = remaining.min(buf.len() as u64) as usize; + let got = rd.read(&mut buf[..want]).await?; + if got == 0 { + bail!("BW_PROBE upload ended after {} of {upload} bytes", upload - remaining); + } + remaining -= got as u64; + } + wr.write_all(&upload.to_le_bytes()).await?; + buf.fill(0); + let mut remaining = download; + while remaining > 0 { + let take = remaining.min(buf.len() as u64) as usize; + wr.write_all(&buf[..take]).await?; + remaining -= take as u64; + } + wr.shutdown().await?; + Ok(()) } async fn send_direct(wr: &mut OwnedWriteHalf, msg_type: u8, payload: &[u8]) -> Result<()> { @@ -4410,6 +4500,28 @@ mod hello_auth_tests { out } + #[test] + fn bw_probe_body_is_two_bounded_sizes() { + let mut body = Vec::new(); + body.extend_from_slice(&5u64.to_le_bytes()); + body.extend_from_slice(&7u64.to_le_bytes()); + assert_eq!(parse_bw_probe(&body).unwrap(), (5, 7)); + assert!(parse_bw_probe(&body[..15]).is_err()); + let mut too_big = Vec::new(); + too_big.extend_from_slice(&(BW_PROBE_MAX_BYTES + 1).to_le_bytes()); + too_big.extend_from_slice(&0u64.to_le_bytes()); + assert!(parse_bw_probe(&too_big).is_err()); + } + + #[test] + fn bw_probe_slots_are_bounded() { + let slots: Vec<_> = (0..BW_PROBE_MAX_CONNECTIONS).map(|_| BwProbeSlot::acquire()).collect(); + assert!(slots.iter().all(Option::is_some)); + assert!(BwProbeSlot::acquire().is_none()); + drop(slots); + assert!(BwProbeSlot::acquire().is_some()); + } + #[test] fn open_hello_accepts_the_right_key_only() { let body = b"hello body bytes"; @@ -4951,7 +5063,7 @@ mod tests { } let mut stream = stream.expect("legacy syncer did not start"); // A correctly sealed elastic payload (epoch 0) still must not be - // accepted: legacy only knows HELLO/DATA_HELLO as a first frame. + // accepted: legacy only knows HELLO/DATA_HELLO/BW_PROBE as a first frame. let payload = crate::elastic::seal(&key, msg_type, 0u64.to_le_bytes().to_vec()); let mut frame = MAGIC.to_le_bytes().to_vec(); frame.push(msg_type); @@ -4968,7 +5080,7 @@ mod tests { assert_eq!(reply.get(4), Some(&MSG_ERROR), "type {msg_type}"); let text = String::from_utf8_lossy(&reply[13..]); assert!( - text.contains(&format!("first frame must be HELLO/DATA_HELLO, got {msg_type}")), + text.contains(&format!("first frame must be HELLO/DATA_HELLO/BW_PROBE, got {msg_type}")), "type {msg_type}: {text}" ); assert!(!server.is_finished(), "type {msg_type}: server exited"); diff --git a/tests/test_e3_overlap_miles_role_stream.py b/tests/test_e3_overlap_miles_role_stream.py index 6d9f0aa1d..6a69f85fc 100644 --- a/tests/test_e3_overlap_miles_role_stream.py +++ b/tests/test_e3_overlap_miles_role_stream.py @@ -266,6 +266,15 @@ def send_init_parts(self, fragment_id, parts) -> bool: self.pull_sweep(0, 1) return True + def send_initial_fragments(self, parts_for) -> bool: + return all(self.send_init_parts(f, parts_for(f)) for f in range(P)) + + def init_replay_due(self) -> bool: + return False + + def mark_initial_cut_received(self) -> None: + pass + def check_health(self) -> None: assert not self.closed diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 92bf9b97c..7fd428ebb 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -633,3 +633,67 @@ def test_stalled_stream_redials_and_the_replayed_push_keeps_the_exact_result( assert progress.stalls >= 1 and progress.reconnects >= 1 assert list(manifest.versions) == exact[1] assert result == exact[2] + + +def _initial_values(fid: int) -> list[float]: + return [float(fid) + 0.125 * i for i in range(NUMEL)] + + +def test_init_interrupted_by_a_redial_is_replayed_and_the_result_stays_exact(syncer_binary, tmp_path): + """E3b (s26-e3b-p06-C32a): learner 0's group died after INIT of fragment 0 + and before fragment 1, so the syncer never finished initializing and no + island got the initial cut. The client now replays INIT on the next group + until the initial cut arrives; the result matches an undisturbed run. A + link probe on the same port (yeto.wan_probe) leaves the session alone.""" + from yeto.protocol import island_hmac_key_from_env + from yeto.wan_probe import probe_link + + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + run_dir = tmp_path / "replay" + run_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, run_dir, *STREAMED, total_steps=3) + client = SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), stream_stall_timeout=30.0) + killed = [] + + def parts_for(fid: int): + if fid == 1 and not killed: + killed.append(client.connection_generation) + client._poison_outbound_generation(client._gen, ConnectionError("test: group died mid-INIT")) + return [_f32(_initial_values(fid))] + + try: + def listening() -> bool: + try: + socket.create_connection(("127.0.0.1", port), timeout=1).close() + except OSError: + return False + return True + + _wait(listening, "syncer listening") + good = probe_link(("127.0.0.1", port), streams=2, bytes_per_direction=8 << 20, + key=island_hmac_key_from_env(), timeout=30) + assert good.error is None and good.upload_mbps > 0 and good.download_mbps > 0 + bad = probe_link(("127.0.0.1", port), streams=1, bytes_per_direction=1 << 20, + key=b"wrong key", timeout=10) + assert bad.error is not None and not bad.accepted(0.0, 0.0) + client.start() + assert client.send_initial_fragments(parts_for) is False + replays = 0 + while client._bcasts.qsize() < FRAGMENTS: + client.check_health() + if client.init_replay_due(): + replays += client.send_initial_fragments(parts_for) + time.sleep(0.02) + client.mark_initial_cut_received() + assert killed and replays == 1 and client.connection_generation != killed[0] + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + _stop(process) + assert list(manifest.versions) == exact[1] + assert result == exact[2] diff --git a/tests/test_wan_probe.py b/tests/test_wan_probe.py new file mode 100644 index 000000000..9ce848aba --- /dev/null +++ b/tests/test_wan_probe.py @@ -0,0 +1,126 @@ +"""S26 E3b: island link probe, its island hook, and the launcher's relaunch of +an island whose probe rejected the container's link.""" + +from __future__ import annotations + +import json +from argparse import Namespace + +import pytest + +from yeto.wan_probe import EXIT_SLOW_LINK, ProbeResult, probe_link + + +def _fake_stream(up_s: float, down_s: float): + def run(addr, upload, download, key, timeout, out): + out["up"], out["down"] = up_s, down_s + return run + + +def test_probe_rates_use_the_slowest_stream_per_direction(): + result = probe_link(("h", 1), streams=4, bytes_per_direction=40_000_000, key=None, + stream_fn=_fake_stream(2.0, 0.5)) + assert (result.upload_mbps, result.download_mbps) == (20.0, 80.0) + assert result.accepted(20.0, 80.0) and not result.accepted(20.1, 0.0) + assert result.fields()["wan_probe/upload_mbps"] == 20.0 + + +def test_a_failed_stream_rejects_the_link(): + def broken(addr, upload, download, key, timeout, out): + raise ConnectionError("reset") + result = probe_link(("h", 1), streams=2, bytes_per_direction=2 << 20, key=None, stream_fn=broken) + assert result.error == "ConnectionError: reset" and not result.accepted(0.0, 0.0) + with pytest.raises(ValueError): + probe_link(("h", 1), streams=0, bytes_per_direction=1, key=None) + + +def _island_args(tmp_path, **kw): + return Namespace(syncer="10.0.0.1:29400", wan_streams=8, wan_probe_mib=16, + event_tape=str(tmp_path / "tape.jsonl"), learner_id=1, **kw) + + +def test_island_probe_is_off_by_default_and_exits_75_on_a_slow_link(tmp_path, capsys): + from yeto.rl.adapters.miles.island_entry import probe_wan_link + + calls = [] + + def probe(addr, **kw): + calls.append((addr, kw["streams"], kw["bytes_per_direction"])) + return ProbeResult(8, 16 << 20, 16 << 20, 4.0, 1.0) + + probe_wan_link(_island_args(tmp_path), probe=probe) + assert calls == [] + probe_wan_link(_island_args(tmp_path, wan_probe_min_mbps=2.0, wan_probe_min_down_mbps=10.0), probe=probe) + with pytest.raises(SystemExit) as stop: + probe_wan_link(_island_args(tmp_path, wan_probe_min_mbps=8.0), probe=probe) + assert stop.value.code == EXIT_SLOW_LINK + assert calls == [(("10.0.0.1", 29400), 8, 16 << 20)] * 2 + records = [json.loads(line) for line in (tmp_path / "tape.jsonl").read_text().splitlines()] + assert [r["accepted"] for r in records if r["event"] == "rl_wan_probe"] == [True, False] + assert "[yeto-wan-probe] rejected: up 4.2 MB/s" in capsys.readouterr().out + + +def test_launcher_flags_reach_the_island_and_bound_the_relaunches(): + from yeto.launch.modal_islands import wan_probe_relaunches + from yeto.launch.rl_island_task import _wan_transport_flags + + on = Namespace(wan_probe_min_mbps=20.0, wan_probe_min_down_mbps=0.0, wan_probe_mib=128, + wan_probe_relaunches=3) + assert _wan_transport_flags(on) == (" --wan-probe-min-mbps 20.0 --wan-probe-min-down-mbps 0.0" + " --wan-probe-mib 128") + assert wan_probe_relaunches(on) == 3 + assert wan_probe_relaunches(Namespace(wan_probe_relaunches=3)) == 0 # probe off + with pytest.raises(ValueError): + _wan_transport_flags(Namespace(wan_probe_min_mbps=1.0, wan_probe_mib=0)) + + +def test_echo_collector_counts_rejected_probes(tmp_path): + from yeto.rl.event_echo import PREFIX, TapeCollector + + collector = TapeCollector(tmp_path / "tape.jsonl", fresh=True) + for accepted, t in ((False, 1), (True, 2), (False, 3)): + collector.feed(PREFIX + json.dumps({"event": "rl_wan_probe", "accepted": accepted, "island_id": 1, "time_unix": t}) + "\n") + assert collector.wan_rejects == 2 + + +def _fleet(ops, rejects, relaunches, **kw): + from test_controller import RUNNING, SYNCER, ImmediateThread + + from yeto.launcher import FleetController + + ops.status_seq.setdefault(SYNCER, [RUNNING]) + return FleetController(learners={"l0": ("task-l0", 1), "l1": ("task-l1", 2)}, + syncer=(SYNCER, "task-syncer", 1), sky_ops=ops, poll_interval=30, + recover_timeout=0, thread_cls=ImmediateThread, fixed_roster=True, + no_island_relaunch=True, wan_reject_probe=lambda name: rejects[name], + wan_probe_relaunches=relaunches, **kw) + + +def test_rejected_link_is_relaunched_even_without_island_relaunch(): + from test_controller import FAILED, RUNNING, SUCCEEDED, FakeOps + + ops = FakeOps() + ops.status_seq["l0"] = [RUNNING, RUNNING, RUNNING, SUCCEEDED] + ops.status_seq["l1"] = [FAILED] + ops.relaunch_results["l1"] = [7] + ops.after_relaunch["l1"] = [RUNNING, SUCCEEDED] + rejects = {"l0": 0, "l1": 1} + ctl = _fleet(ops, rejects, relaunches=2) + ctl.run() + assert ops.relaunch_calls == ["l1"] + assert ctl.learners["l1"]["wan_relaunches"] == 1 and ctl.learners["l1"]["state"] == "done" + + +def test_rejected_link_budget_is_bounded(): + from test_controller import FAILED, RUNNING, FakeOps + + from yeto.launcher import FixedRosterIslandAbandoned + + ops = FakeOps() + ops.status_seq["l0"] = [RUNNING] + ops.status_seq["l1"] = [FAILED] + rejects = {"l0": 0, "l1": 1} + ctl = _fleet(ops, rejects, relaunches=0) + with pytest.raises(FixedRosterIslandAbandoned): + ctl.run() + assert ops.relaunch_calls == [] diff --git a/tests/test_wan_streams.py b/tests/test_wan_streams.py index 4e08efb0b..99f3d4dc1 100644 --- a/tests/test_wan_streams.py +++ b/tests/test_wan_streams.py @@ -227,3 +227,35 @@ def test_island_parser_accepts_the_transport_flags(): assert rl_option("wan_stream_schedule", ISLAND) == "--wan-stream-schedule" assert rl_option("wan_stall_timeout", ISLAND) == "--wan-stall-timeout" + + +def test_tcp_user_timeout_waits_for_the_initial_cut(): + if not hasattr(socket, "TCP_USER_TIMEOUT"): + pytest.skip("Linux-only socket option") + with socket.create_server(("127.0.0.1", 0)) as server: + addr = server.getsockname() + client = SyncerClient(addr, 0, _layout(), num_streams=1, stream_stall_timeout=2.0) + early = client._dial(5) + try: + assert early.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 0 + client._socks = [early] + client.mark_initial_cut_received() + assert early.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2000 + late = client._dial(5) + assert late.getsockopt(socket.IPPROTO_TCP, socket.TCP_USER_TIMEOUT) == 2000 + late.close() + finally: + early.close() + + +def test_init_replay_is_due_only_on_a_new_group_before_the_initial_cut(): + client = SyncerClient(("127.0.0.1", 1), 0, _layout(), num_streams=1) + client._connected.set() + client.connection_generation = 11 + assert client.init_replay_due() + client._init_generation = 11 + assert not client.init_replay_due() + client.connection_generation = 12 + assert client.init_replay_due() + client.mark_initial_cut_received() + assert not client.init_replay_due() diff --git a/yeto/cli.py b/yeto/cli.py index 40abdc1f0..ab16fb4b9 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -454,6 +454,29 @@ def int_or_auto(value: str): "one stream holds bytes with no progress for this many seconds (also sets " "TCP_USER_TIMEOUT); the fragment is pushed again on the replayed PULL. 0 = off", ) + sync.add_argument( + "--wan-probe-min-mbps", + type=float, + default=0.0, + help="Miles RL islands: before loading the model, measure the upload rate to the " + "actor syncer over --wan-streams connections; below this MB/s the island exits " + "(code 75) and the launcher requests another container (--wan-probe-relaunches). 0 = off", + ) + sync.add_argument( + "--wan-probe-min-down-mbps", + type=float, + default=0.0, + help="like --wan-probe-min-mbps, for the download rate from the syncer. 0 = off", + ) + sync.add_argument("--wan-probe-mib", type=int, default=256, + help="bytes per direction of the link probe, in MiB") + sync.add_argument( + "--wan-probe-relaunches", + type=int, + default=2, + help="containers the launcher may request again per island after a rejected link " + "probe (also under --no-island-relaunch: the island never joined the session)", + ) infra = p.add_argument_group("infrastructure") infra.add_argument( diff --git a/yeto/launch/fleet.py b/yeto/launch/fleet.py index b7051a46a..d1d9ce639 100644 --- a/yeto/launch/fleet.py +++ b/yeto/launch/fleet.py @@ -48,6 +48,8 @@ def __init__(self): # time spent recovering in the window. FIXED_ROSTER_MAX_RELAUNCHES = 2 RECOVERY_STABLE_S = 300.0 +# S26 E3b: how long one rejected-link relaunch may take to get a new container. +WAN_RECOVERY_S = 900.0 # How long a failed-looking island is given for its rl_learner_finalized record # to arrive over the log stream before recovery starts. FINALIZE_GRACE_S = 60.0 @@ -154,6 +156,8 @@ def __init__( elastic: bool = False, no_island_relaunch: bool = False, replacement_planner=None, + wan_reject_probe=None, + wan_probe_relaunches: int = 0, ): """`learners` maps cluster name -> (task, job_id); `syncer` is (name, task, job_id) for a cluster syncer, or None with @@ -169,6 +173,12 @@ def __init__( self.left_pool_probe = left_pool_probe self.elastic = elastic self.no_island_relaunch = no_island_relaunch + # S26 E3b: name -> number of rejected link probes the island logged + # (rl_wan_probe accepted=false, exit 75). Such an island never joined + # the syncer session, so it is relaunched in a fresh container up to + # wan_probe_relaunches times, also under --no-island-relaunch. + self.wan_reject_probe = wan_reject_probe + self.wan_probe_relaunches = wan_probe_relaunches # rl-spot-cost-saving 2.3: yeto.cloud.replace.ReplacementPlanner or None. # Asked once per recovery window before the same-place relaunch. Auto # launch off (default): it only writes spot_replace_advice and the @@ -372,6 +382,8 @@ def _poll(self, rec, is_syncer: bool) -> None: self._fleet("island_done", rec["name"], reason="syncer finished") print(f"[launcher] {rec['name']}: {verdict} after the syncer finished; " "counted as a normal end, not relaunched", file=sys.stderr) + elif not is_syncer and self._rejected_link(rec): + self._relaunch_rejected_link(rec) elif not is_syncer and self.elastic and self._left_pool(rec): # 0.27: re-JOIN exhausted (exit 7): the island left the elastic pool. # Not the fixed-roster abandon path (exit 4): relaunch unless @@ -418,6 +430,36 @@ def _syncer_finished(self) -> bool: except Exception: return False + def _rejected_link(self, rec) -> bool: + """A new rejected link probe since the island's last (re)launch.""" + if self.wan_reject_probe is None: + return False + try: + seen = int(self.wan_reject_probe(rec["name"])) + except Exception: + return False + return seen > rec.get("wan_rejects_handled", 0) + + def _relaunch_rejected_link(self, rec) -> None: + rec["wan_rejects_handled"] = rec.get("wan_rejects_handled", 0) + 1 + used = rec.get("wan_relaunches", 0) + if used >= self.wan_probe_relaunches: + self._fleet("island_lost", rec["name"], reason="link probe rejected; relaunch budget used") + print(f"[launcher] {rec['name']}: link probe rejected again after {used} new " + "container(s); not relaunched", file=sys.stderr) + self._abandon(rec, 0.0, reason="link probe rejected") + return + rec["wan_relaunches"] = used + 1 + self._fleet("island_lost", rec["name"], + reason=f"link probe rejected; new container {used + 1}/{self.wan_probe_relaunches}") + print(f"[launcher] {rec['name']}: link probe rejected; requesting a new container " + f"({used + 1}/{self.wan_probe_relaunches})", file=sys.stderr) + rec["state"] = RECOVERING + rec["failed_at"] = self.ops.now() + rec["wan_recovery"] = True + rec.pop("recovered_at", None) + self._drive_recovery(rec, False) + def _left_pool(self, rec) -> bool: if self.left_pool_probe is None: return False @@ -577,6 +619,7 @@ def _drive_recovery(self, rec, is_syncer: bool) -> None: self.ops.now() - rec["failed_at"]) rec["failed_at"] = None rec["replace_planned"] = False + rec.pop("wan_recovery", None) rec["recovered_at"] = self.ops.now() print( f"[launcher] {rec['name']} recovered: relaunched as job " @@ -597,7 +640,9 @@ def _drive_recovery(self, rec, is_syncer: bool) -> None: if rec.get("failures", 0) > FIXED_ROSTER_MAX_RELAUNCHES: self._abandon(rec, elapsed) return - if self.recover_timeout <= 0 or elapsed > self.recover_timeout: + # A rejected-link relaunch has its own budget (the island never joined). + timeout = WAN_RECOVERY_S if rec.get("wan_recovery") else self.recover_timeout + if timeout <= 0 or elapsed > timeout: if is_syncer: # The syncer is never abandoned: without it no learner can # make outer progress, so keep trying and complain loudly. diff --git a/yeto/launch/modal_islands.py b/yeto/launch/modal_islands.py index a7f12af6f..976211e31 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -315,3 +315,15 @@ def effective_recover_timeout(args) -> float: if getattr(args, "modal_retries", None) == 0 and not getattr(args, "modal_launcher_relaunch", False): return 0 return args.recover_timeout + + +def wan_probe_relaunches(args) -> int: + """New containers per island after a rejected link probe (S26 E3b); + 0 when no island runs the probe.""" + if not (float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) > 0 + or float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) > 0): + return 0 + count = int(getattr(args, "wan_probe_relaunches", 2)) + if count < 0: + raise ValueError("--wan-probe-relaunches must be >= 0") + return count diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index a51d6ea1e..0254a6036 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -62,6 +62,15 @@ def _wan_transport_flags(args) -> str: raise ValueError("--wan-stall-timeout must be >= 0 (0 = off)") if stall > 0: flags += f" --wan-stall-timeout {stall!r}" + min_up = float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) + min_down = float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) + if min_up < 0 or min_down < 0: + raise ValueError("--wan-probe-min-mbps / --wan-probe-min-down-mbps must be >= 0 (0 = off)") + if min_up > 0 or min_down > 0: + mib = int(getattr(args, "wan_probe_mib", 256)) + if not 1 <= mib <= 1024: + raise ValueError("--wan-probe-mib must be between 1 and 1024") + flags += f" --wan-probe-min-mbps {min_up!r} --wan-probe-min-down-mbps {min_down!r} --wan-probe-mib {mib}" return flags diff --git a/yeto/launcher.py b/yeto/launcher.py index 0abd2f4b5..7819fbc5e 100644 --- a/yeto/launcher.py +++ b/yeto/launcher.py @@ -61,7 +61,7 @@ from .launch.modal_islands import ( _modal_tape_run_dir, _tail_modal, _verify_modal_app_stopped, build_modal_island_config, container_set_poller_for, effective_recover_timeout, pull_modal_tapes, - require_modal_for_gpu_exact, + require_modal_for_gpu_exact, wan_probe_relaunches, ) from .launch.fleet import ( CONTAINER_CHANGED_EXIT, DEFAULT_RL_STALL_TIMEOUT_S, FAILED_RUN_DRAIN_S, ISLAND_FAILED_EXIT, @@ -660,6 +660,11 @@ def container_changed() -> bool: ), elastic=island_scheduling_mode(args) == "elastic", no_island_relaunch=bool(getattr(args, "no_island_relaunch", False)), + wan_reject_probe=( + (lambda name: event_collectors[name].wan_rejects if name in event_collectors else 0) + if echo_names else None + ), + wan_probe_relaunches=wan_probe_relaunches(args), ) def drain_tapes(limit: float = NO_SYNC_EVENT_DRAIN_S) -> None: if echo_names: diff --git a/yeto/protocol.py b/yeto/protocol.py index 42cfd3f89..e55444b5d 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -60,6 +60,8 @@ MSG_FINAL_ACK = 12 MSG_FINAL_FRAGMENT = 13 MSG_BUDGET_DONE = 14 +# S26 E3b: link-bandwidth probe connection (yeto/wan_probe.py; syncer serve_bw_probe). +MSG_BW_PROBE = 26 FINALIZATION_REVISION = 1 FINALIZATION_TIMEOUT = 900.0 @@ -659,6 +661,12 @@ def __init__( self._meters = [StreamMeter() for _ in range(num_streams + 1)] self._upload_drained_at: float | None = None self._stalls = 0 + # S26 E3b: set once the initial cut arrived. Before that a redial must + # replay INIT_PARAMS (learner 0), and TCP_USER_TIMEOUT stays unset: + # the multi-GB INIT upload and initial broadcast must not be reset by + # the kernel on a slow link (s26-e3b-p06-C32a). + self._initial_cut_received = threading.Event() + self._init_generation: int | None = None self.session_contract_hash = session_contract_hash self.syncer_profile_hash = syncer_profile_hash # secret-handling-hardening: HELLO carries an HMAC when the island has @@ -733,7 +741,7 @@ def _dial(self, timeout: float) -> socket.socket: sock.settimeout(None) sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) enable_tcp_keepalive(sock) - if self.stream_stall_timeout is not None: + if self.stream_stall_timeout is not None and self._initial_cut_received.is_set(): set_tcp_user_timeout(sock, self.stream_stall_timeout) return sock @@ -961,6 +969,55 @@ def upload_drained_at(self) -> float | None: # -- learner-facing API ------------------------------------------------------ + def send_initial_fragments( + self, + parts_for: Callable[[int], Iterable[bytes | bytearray | memoryview]], + ) -> bool: + """Learner 0: stream INIT_PARAMS of every fragment on the live group. + + ``parts_for(fragment_id)`` returns a fresh part iterable each call, so + the same INIT can be replayed after a redial. Returns True when every + fragment was queued on one connection generation; False on an outage + (nothing to undo: the syncer ignores INIT it already applied, and a + partial INIT of a dead generation is dropped by the syncer).""" + generation = self.connection_generation + try: + for fragment_id in range(self.layout.num_fragments): + if self.send_init_parts(fragment_id, parts_for(fragment_id)) is not True: + return False + except PartialMessageGenerationLost: + return False + if generation != self.connection_generation: + return False + self._init_generation = generation + return True + + def init_replay_due(self) -> bool: + """True when the live group never carried a complete INIT and the + initial cut has not arrived yet (learner 0 must send INIT again).""" + with self._lock: + live = self._connected.is_set() + generation = self.connection_generation + return ( + live + and not self._initial_cut_received.is_set() + and generation != self._init_generation + ) + + def mark_initial_cut_received(self) -> None: + """The initial cut arrived: INIT is no longer replayed, and the stall + timeout now also bounds unacknowledged data in the kernel.""" + self._initial_cut_received.set() + if self.stream_stall_timeout is None: + return + with self._lock: + socks = list(self._socks) + for sock in socks: + try: + set_tcp_user_timeout(sock, self.stream_stall_timeout) + except OSError: + pass # a dying socket; the group is being replaced anyway + def send_init(self, fragment_id: int, tensor_bytes: bytes) -> None: if not 0 <= fragment_id < self.layout.num_fragments: raise ValueError(f"INIT_PARAMS for unknown fragment {fragment_id}") diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index fa746fcb0..549743f7f 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -400,27 +400,23 @@ async def initialize(self, group) -> None: client.install_inbound_chunk_sink(MilesRayInboundChunkSink(sink)) try: client.start() + init_parts = None if self.config.learner_id == 0: - for fragment_id in range(adapter.layout.fragments.num_fragments): - queued = client.send_init_parts( - fragment_id, - adapter.fragment_parts( - initial, - fragment_id, - ray_module=self.ray_module, - )(), - ) - if queued is not True: - raise RuntimeError( - f"{self.config.algorithm} {self.config.role} initial fragment was not queued" - ) + def init_parts(fragment_id: int): + return adapter.fragment_parts( + initial, fragment_id, ray_module=self.ray_module + )() + + # An outage here is not fatal: the wait below replays INIT on + # the next connection group until the initial cut arrives. + client.send_initial_fragments(init_parts) # Wait off the event loop: an island initializes its roles # concurrently, and each role's initial cut needs every island's # INIT for that role. A blocking wait here would stop this # island's other role from connecting, and two islands blocked # on different roles deadlock (s23-g3-sao-main-20261014a). authoritative = await asyncio.to_thread( - self._wait_for_initial_fragments, client, adapter + self._wait_for_initial_fragments, client, adapter, init_parts ) adapter.verify_initial_fragments(initial, authoritative) except BaseException: @@ -443,6 +439,7 @@ def _wait_for_initial_fragments( self, client, adapter: MilesChunkedFullParameterAdapter, + init_parts=None, ) -> tuple[StoredAuthoritativeFragment, ...]: started = time.monotonic() deadline = started + self.config.wait_timeout @@ -452,6 +449,11 @@ def _wait_for_initial_fragments( self._report_initial_cut(client, observed, count, started, "waiting") while len(observed) < count: client.check_health() + if init_parts is not None and client.init_replay_due(): + # S26 E3b (s26-e3b-p06-C32a): the group carrying INIT died + # before the syncer applied it; send it again on this group. + self._report_initial_cut(client, observed, count, started, "replaying INIT") + client.send_initial_fragments(init_parts) for update in client.drain_updates(): stored = self._stored_update(update) if update.version != 0 or stored.fragment_id in observed: @@ -471,6 +473,7 @@ def _wait_for_initial_fragments( ) time.sleep(self.config.poll_seconds) self._report_initial_cut(client, observed, count, started, "received") + client.mark_initial_cut_received() return tuple(observed[index] for index in range(count)) def _report_initial_cut( diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 01367ff35..cf98a201f 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -1633,6 +1633,40 @@ def probe_syncers(args, *, connect=None, attempts: int = 10, delay: float = 3.0, return addrs +def probe_wan_link(args, *, probe=None) -> None: + """S26 E3b: measure this container's link to the actor syncer before any + model work; exit EXIT_SLOW_LINK below the configured rates so the launcher + can request another container. Off unless a minimum rate is set.""" + min_up = float(getattr(args, "wan_probe_min_mbps", 0.0) or 0.0) + min_down = float(getattr(args, "wan_probe_min_down_mbps", 0.0) or 0.0) + if min_up <= 0 and min_down <= 0: + return + from yeto.protocol import island_hmac_key_from_env + from yeto.wan_probe import EXIT_SLOW_LINK, PROBE_EVENT, PROBE_MARKER, probe_link + + result = (probe or probe_link)( + _syncer_address(args.syncer), + streams=max(1, int(args.wan_streams)), + bytes_per_direction=int(getattr(args, "wan_probe_mib", 256)) << 20, + key=island_hmac_key_from_env(), + ) + accepted = result.accepted(min_up, min_down) + record = {"event": PROBE_EVENT, "accepted": accepted, + "wan_probe/min_upload_mbps": min_up, "wan_probe/min_download_mbps": min_down, + **result.fields()} + print(f"{PROBE_MARKER} {'accepted' if accepted else 'rejected'}: up {result.upload_mbps:.1f} MB/s " + f"(min {min_up}), down {result.download_mbps:.1f} MB/s (min {min_down}), " + f"{result.streams} stream(s){'' if result.error is None else '; ' + result.error}", flush=True) + if getattr(args, "event_tape", None): + from yeto.rl.engine.driver import EventTape + + if getattr(args, "rl_echo_events", False): + install_event_echo() # the launcher reads the verdict from the log stream + EventTape(args.event_tape, int(getattr(args, "learner_id", 0) or 0)).append(record) + if not accepted: + raise SystemExit(EXIT_SLOW_LINK) + + def require_ports_router_mode(miles_args) -> None: """Upstream Miles owns the SGLang router. @@ -2394,6 +2428,7 @@ def main(argv=None) -> None: verify_miles_revision(miles_root, expected=pins) _preflight_codex_harness(args) probe_syncers(args) + probe_wan_link(args) from miles.utils.misc import load_function diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 307788097..803b19e18 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -418,6 +418,9 @@ island_flag("--wan-streams", type=int, default=4), island_flag("--wan-stream-schedule", choices=("rr", "adaptive"), default="rr"), island_flag("--wan-stall-timeout", type=float, default=0.0), + island_flag("--wan-probe-min-mbps", type=float, default=0.0), + island_flag("--wan-probe-min-down-mbps", type=float, default=0.0), + island_flag("--wan-probe-mib", type=int, default=256), island_flag("--miles-root", required=True), island_flag( "--miles-commit", diff --git a/yeto/rl/event_echo.py b/yeto/rl/event_echo.py index b3f2a33e0..b9abdadc8 100644 --- a/yeto/rl/event_echo.py +++ b/yeto/rl/event_echo.py @@ -133,6 +133,7 @@ def __init__(self, path, *, fresh: bool = True) -> None: self.discarded = 0 self.finalized = False self.left_pool = False # 0.27: elastic_left_pool seen (re-JOIN exhausted) + self.wan_rejects = 0 # S26 E3b: rl_wan_probe records with accepted=false self.closed = False @property @@ -179,8 +180,11 @@ def _record_line(self, line: str) -> None: with self.path.open("a", encoding="utf-8") as handle: handle.write(raw + "\n") self.count += 1 - event = json.loads(raw).get("event") - if event == FINALIZED_EVENT: + record = json.loads(raw) + event = record.get("event") + if event == "rl_wan_probe" and record.get("accepted") is False: + self.wan_rejects += 1 + elif event == FINALIZED_EVENT: self.finalized = True elif event == "elastic_left_pool": self.left_pool = True diff --git a/yeto/wan_probe.py b/yeto/wan_probe.py new file mode 100644 index 000000000..4ccb91fef --- /dev/null +++ b/yeto/wan_probe.py @@ -0,0 +1,149 @@ +"""Island-to-syncer link-bandwidth probe (S26 E3b). + +E3/E3b showed that the sync share follows the island *container*: on a slow +container every TCP stream is slow at once (1-5 MB/s each, both directions), +so no stream scheduling can help. An island can measure its link to the head +before it loads a model and exit early when the link is too slow; the +launcher then requests a fresh container (``--wan-probe-relaunches``). + +Each probe connection sends one ``MSG_BW_PROBE`` frame (HMAC-sealed like +HELLO), uploads ``upload_bytes`` raw bytes, reads the syncer's 8-byte +acknowledgement, then reads ``download_bytes`` raw bytes. The syncer never +connects a probe to a learner session. +""" + +from __future__ import annotations + +import socket +import struct +import threading +import time +from collections.abc import Callable +from dataclasses import dataclass + +from .protocol import MAGIC, MSG_BW_PROBE, ProtocolError, seal_hello, write_frame + +# Exit code of an island that rejected its own link (the launcher may relaunch it). +EXIT_SLOW_LINK = 75 +PROBE_EVENT = "rl_wan_probe" +PROBE_MARKER = "[yeto-wan-probe]" +_BLOCK = memoryview(bytes(1 << 20)) + + +@dataclass(frozen=True) +class ProbeResult: + streams: int + upload_bytes: int + download_bytes: int + upload_seconds: float + download_seconds: float + error: str | None = None + + @property + def upload_mbps(self) -> float: + return self.upload_bytes / self.upload_seconds / 1e6 if self.upload_seconds > 0 else 0.0 + + @property + def download_mbps(self) -> float: + return self.download_bytes / self.download_seconds / 1e6 if self.download_seconds > 0 else 0.0 + + def accepted(self, min_upload_mbps: float, min_download_mbps: float) -> bool: + return ( + self.error is None + and self.upload_mbps >= min_upload_mbps + and self.download_mbps >= min_download_mbps + ) + + def fields(self) -> dict[str, object]: + return { + "wan_probe/streams": self.streams, + "wan_probe/upload_bytes": self.upload_bytes, + "wan_probe/download_bytes": self.download_bytes, + "wan_probe/upload_mbps": round(self.upload_mbps, 3), + "wan_probe/download_mbps": round(self.download_mbps, 3), + "wan_probe/error": self.error, + } + + +def encode_probe(upload_bytes: int, download_bytes: int, key: bytes | None) -> bytes: + return seal_hello(struct.pack(" None: + sock = socket.create_connection(addr, timeout=timeout) + try: + sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + write_frame(sock, MSG_BW_PROBE, encode_probe(upload, download, key)) + start = time.monotonic() + left = upload + while left: + take = min(left, len(_BLOCK)) + sock.sendall(_BLOCK[:take]) + left -= take + ack = b"" + while len(ack) < 8: + part = sock.recv(8 - len(ack)) + if not part: + # The syncer refused the probe: it answers with an ERROR frame + # and closes; the ERROR body is the reason. + raise ConnectionError("syncer closed the probe before acknowledging the upload") + ack += part + if ack.startswith(struct.pack(" ProbeResult: + """Measure aggregate upload and download rates over ``streams`` parallel + connections; ``bytes_per_direction`` is split evenly over the streams. + A rate is total bytes over the slowest stream's time in that direction.""" + if streams < 1 or bytes_per_direction < streams: + raise ValueError("probe needs at least one stream and one byte per stream") + per = bytes_per_direction // streams + results = [dict() for _ in range(streams)] + errors: list[str] = [] + + def run(index: int) -> None: + try: + stream_fn(addr, per, per, key, timeout, results[index]) + except (OSError, ProtocolError) as error: + errors.append(f"{type(error).__name__}: {error}") + + threads = [threading.Thread(target=run, args=(i,), daemon=True) for i in range(streams)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout + 5) + done = [r for r in results if "up" in r and "down" in r] + if len(done) < streams and not errors: + errors.append(f"{streams - len(done)} probe stream(s) did not finish in {timeout:.0f}s") + return ProbeResult( + streams=streams, + upload_bytes=per * len(done), + download_bytes=per * len(done), + upload_seconds=max((r["up"] for r in done), default=0.0), + download_seconds=max((r["down"] for r in done), default=0.0), + error=errors[0] if errors else None, + ) + From 823b2e82e6bbfbcce8b4ca904707a283e758fa90 Mon Sep 17 00:00:00 2001 From: michael <78113554+chowhsu@users.noreply.github.com> Date: Wed, 21 Oct 2026 10:02:09 +0000 Subject: [PATCH 9/9] wan: refuse the link probe with the elastic syncer; test a lost socket mid-INIT Only the legacy syncer serves MSG_BW_PROBE. With --rl-elastic every probe would fail, every island would be judged slow and relaunched until its budget ran out, so the launch now refuses the probe flags there. Tests: - real syncer: one data socket of learner 0's group is shut down between the INIT of fragment 0 and fragment 1; the client replaces the group, replays INIT, and the final versions and bytes equal an undisturbed run. - slow-container decision: threshold equality accepts, upload just below rejects, slow download rejects, a zero download minimum is not checked, and a stream error rejects a fast link. Co-Authored-By: Claude Opus 5.5 --- tests/test_rl_dense_sweep_real_syncer.py | 46 ++++++++++++++++++++++++ tests/test_wan_probe.py | 24 +++++++++++++ yeto/launch/rl_island_task.py | 5 +++ 3 files changed, 75 insertions(+) diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index 7fd428ebb..763510372 100644 --- a/tests/test_rl_dense_sweep_real_syncer.py +++ b/tests/test_rl_dense_sweep_real_syncer.py @@ -697,3 +697,49 @@ def listening() -> bool: _stop(process) assert list(manifest.versions) == exact[1] assert result == exact[2] + + +def test_one_socket_lost_mid_init_is_redialed_and_the_result_stays_exact(syncer_binary, tmp_path): + """E3c: a single data socket of learner 0's group fails between INIT of + fragment 0 and fragment 1 (as in s26-e3b-p06-C32a, where one connection + reported BrokenPipe). The client replaces the whole group, replays INIT, + and the final versions and bytes equal an undisturbed run.""" + exact = _run_sweeps(syncer_binary, tmp_path, STREAMED, sweeps=3) + run_dir = tmp_path / "one-socket" + run_dir.mkdir() + process, port, profile_hash = _start_syncer(syncer_binary, run_dir, *STREAMED, total_steps=3) + client = SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=2, + connect_timeout=10, max_reconnects=5, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash)) + broken = [] + + def parts_for(fid: int): + if fid == 1 and not broken: + with client._lock: + sock = client._socks[-1] + broken.append(client.connection_generation) + sock.shutdown(socket.SHUT_RDWR) # one real socket dies; the others stay open + return [_f32(_initial_values(fid))] + + try: + client.start() + client.send_initial_fragments(parts_for) + replays = 0 + deadline = time.monotonic() + 60 + while client._bcasts.qsize() < FRAGMENTS: + assert time.monotonic() < deadline, "initial cut never arrived" + client.check_health() + if client.init_replay_due(): + replays += client.send_initial_fragments(parts_for) + time.sleep(0.02) + client.mark_initial_cut_received() + assert broken and replays >= 1 and client.connection_generation != broken[0] + manifest, finals = _sweep_until_final(client, sweeps=3) + result = {f.fragment_id: bytes(f.data) for f in finals} + client.acknowledge_finalization(manifest) + assert process.wait(timeout=15) == 0 + finally: + client.close() + _stop(process) + assert list(manifest.versions) == exact[1] + assert result == exact[2] diff --git a/tests/test_wan_probe.py b/tests/test_wan_probe.py index 9ce848aba..ad8a19dff 100644 --- a/tests/test_wan_probe.py +++ b/tests/test_wan_probe.py @@ -124,3 +124,27 @@ def test_rejected_link_budget_is_bounded(): with pytest.raises(FixedRosterIslandAbandoned): ctl.run() assert ops.relaunch_calls == [] + + +@pytest.mark.parametrize( + ("up_s", "down_s", "error", "min_up", "min_down", "accepted"), + [ + (1.0, 1.0, None, 100.0, 100.0, True), # 100 MB in 1 s each way: exactly at the minimum + (1.0, 1.0, None, 100.1, 0.0, False), # upload just below the minimum + (1.0, 4.0, None, 50.0, 50.0, False), # fast upload, slow download + (1.0, 4.0, None, 50.0, 0.0, True), # download not checked when its minimum is 0 + (0.1, 0.1, "OSError: x", 1.0, 1.0, False), # any stream error rejects a fast link + ], +) +def test_slow_container_decision(up_s, down_s, error, min_up, min_down, accepted): + result = ProbeResult(streams=8, upload_bytes=100_000_000, download_bytes=100_000_000, + upload_seconds=up_s, download_seconds=down_s, error=error) + assert result.accepted(min_up, min_down) is accepted + + +def test_probe_is_refused_with_the_elastic_syncer(): + from yeto.launch.rl_island_task import _wan_transport_flags + + with pytest.raises(ValueError, match="elastic"): + _wan_transport_flags(Namespace(wan_probe_min_mbps=20.0, rl_elastic=True)) + assert _wan_transport_flags(Namespace(rl_elastic=True)) == "" # probe off: no error diff --git a/yeto/launch/rl_island_task.py b/yeto/launch/rl_island_task.py index 0254a6036..ce0ecea76 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -67,6 +67,11 @@ def _wan_transport_flags(args) -> str: if min_up < 0 or min_down < 0: raise ValueError("--wan-probe-min-mbps / --wan-probe-min-down-mbps must be >= 0 (0 = off)") if min_up > 0 or min_down > 0: + if getattr(args, "rl_elastic", False): + # Only the legacy syncer serves MSG_BW_PROBE; on the elastic server + # every probe would fail and every island would be rejected. + raise ValueError("--wan-probe-min-mbps / --wan-probe-min-down-mbps need the " + "legacy syncer; the elastic syncer has no link probe") mib = int(getattr(args, "wan_probe_mib", 256)) if not 1 <= mib <= 1024: raise ValueError("--wan-probe-mib must be between 1 and 1024")