diff --git a/docs/PROTOCOL.md b/docs/PROTOCOL.md index 579db093..db8e47e9 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, 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 | @@ -223,6 +223,37 @@ 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`). + +## 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 00000000..32e6f82d --- /dev/null +++ b/syncer/src/downlink.rs @@ -0,0 +1,302 @@ +//! 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), + }) + } + + /// 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 { + 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 513caabe..58ddc964 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; @@ -98,6 +99,15 @@ 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, + /// 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, @@ -273,6 +283,8 @@ fn main() -> anyhow::Result<()> { final_state: args.final_state, 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 bf4b530e..e214faf2 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 @@ -57,17 +61,85 @@ 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; + +/// 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 { - if dtype == DTYPE_Q4 { - DTYPE_BF16 - } else { - dtype + match session_base(dtype) { + DTYPE_Q4 => DTYPE_BF16, + DTYPE_F32_PUSH_BF16 | DTYPE_F32_PUSH_I8 => DTYPE_F32, + other => other, + } +} + +/// 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 { + 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) + } } } @@ -257,6 +329,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 @@ -312,6 +458,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); @@ -417,6 +616,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 a5ed5b67..c03011a3 100644 --- a/syncer/src/server.rs +++ b/syncer/src/server.rs @@ -91,6 +91,16 @@ 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, + /// 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, @@ -702,8 +712,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"); @@ -726,13 +739,28 @@ 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" ); } } + 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"); + } + if cfg.learner_budget_steps.is_some() { + bail!("--async-checkpoint does not support --learner-budget-steps"); + } + } Ok(()) } @@ -743,6 +771,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"); @@ -980,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(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")?, @@ -1129,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) { + if !valid_session_dtype(dtype) { bail!("unsupported session dtype {dtype}"); } let num_fragments = r.u32()?; @@ -1222,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 { @@ -1428,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<()> { @@ -1700,23 +1828,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(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(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. @@ -1862,6 +1976,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()); @@ -1876,6 +2001,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, @@ -2085,7 +2221,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)?; @@ -2116,12 +2252,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 +2277,53 @@ 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)?; + 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, + &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 +2524,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 +2551,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 +2674,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 +2772,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 +2871,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 +3167,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 +3429,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 +3458,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 persist_before_broadcast(cfg, committed.t) { + 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 +3515,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 +3529,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. @@ -3268,16 +3582,43 @@ async fn commit_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) => { + // 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) { - 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, @@ -3318,6 +3659,7 @@ async fn commit_round( quorum_ms, grace_ms, sync_ms, + checkpoint_ms, &pushes, gnorm, ms, @@ -3330,7 +3672,8 @@ 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 { + // 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)?; @@ -3341,6 +3684,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,10 +3795,38 @@ 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)?; + 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 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!( @@ -3383,6 +3845,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()); } @@ -3721,6 +4194,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()); @@ -3856,6 +4334,7 @@ fn append_tape( quorum_ms: Option, grace_ms: Option, sync_ms: u64, + checkpoint_ms: Option, pushes: &HashMap, gnorm: f64, ms: u64, @@ -3932,9 +4411,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 +4434,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(); @@ -4014,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"; @@ -4354,6 +4862,8 @@ mod tests { final_state: None, checkpoint_path: None, checkpoint_every: 0, + async_checkpoint: false, + bcast_delta_full_every: None, resume: false, mark_final_checkpoint: false, learner_budget_steps: None, @@ -4553,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); @@ -4570,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"); @@ -4618,10 +5128,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!( @@ -4654,13 +5184,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] @@ -5244,6 +5782,197 @@ 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 + } + + #[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"); + 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!( @@ -5410,13 +6139,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:#}" ); @@ -6227,6 +6956,7 @@ mod tests { Some(11), Some(22), 33, + None, &pushes, 0.5, 44, @@ -6275,6 +7005,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, @@ -6303,6 +7034,7 @@ mod tests { Some(1), Some(0), 2, + None, &pushes, 0.5, 3, @@ -6410,6 +7142,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 f19ce4c9..3b54cb9e 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, }) } @@ -628,6 +631,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 +693,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 +705,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/golden/decoupling/codex_harness.json b/tests/golden/decoupling/codex_harness.json index b957d6a5..f74fb036 100644 --- a/tests/golden/decoupling/codex_harness.json +++ b/tests/golden/decoupling/codex_harness.json @@ -81,6 +81,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -229,7 +233,7 @@ "", "--group-query-attention" ], - "miles_argv_sha256": "279d3bae4c7c8381c9dfbea4e0e83ffdbfcd0afbeabb0f8931d7c3312984f771", + "miles_argv_sha256": "8a2934f5ac2b11fae39d1374ed834d0ab5ed4eb6c1d74f515a2f64bc2505830c", "name": "codex_harness", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/decoupled.json b/tests/golden/decoupling/decoupled.json index eb40b112..0d63eb51 100644 --- a/tests/golden/decoupling/decoupled.json +++ b/tests/golden/decoupling/decoupled.json @@ -63,6 +63,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -195,7 +199,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "1108c08ed4e9ef28ec5ad2e37b3fe6ceb804db496824c9de4743b028a400a663", + "miles_argv_sha256": "6ff6da47c19cf9f9eac0ec0b9040e9e199220578244928625c5bec362379ccf6", "name": "decoupled", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/drgrpo.json b/tests/golden/decoupling/drgrpo.json index fc7f4d60..786cb2c1 100644 --- a/tests/golden/decoupling/drgrpo.json +++ b/tests/golden/decoupling/drgrpo.json @@ -110,6 +110,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -245,7 +249,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "7ee222703a50f20c2c18b817596ddbca10140fe94645ca28acdcba05c871cc2b", + "miles_argv_sha256": "8545800520cd18d7ee951fbb8a06d012cbb6bbd39d70968e85f35c1b05d2c6e7", "name": "drgrpo", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/elastic.json b/tests/golden/decoupling/elastic.json index 5664ac5b..d37d4cd6 100644 --- a/tests/golden/decoupling/elastic.json +++ b/tests/golden/decoupling/elastic.json @@ -66,6 +66,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -201,7 +205,7 @@ "120", "--group-query-attention" ], - "miles_argv_sha256": "c88e8996324d0261dc72f5b2949afc1e7eb81db916d3746ac76addea772020d1", + "miles_argv_sha256": "5aeca40e53d9f1b1a614a6a6ec814e5d037b354da3f448f04c1d10417751d135", "name": "elastic", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/grpo_default.json b/tests/golden/decoupling/grpo_default.json index 61c61102..5235984f 100644 --- a/tests/golden/decoupling/grpo_default.json +++ b/tests/golden/decoupling/grpo_default.json @@ -59,6 +59,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -191,7 +195,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "cd9c51c2946a4e45bb63cd153ae64bc3ec0e3e998cf3d5e64fe6d4ed882c156b", + "miles_argv_sha256": "bb315aa0c82a03c2821899a3a376db84f09d30f1d32e2f6640833aec455d8634", "name": "grpo_default", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/grpo_tis.json b/tests/golden/decoupling/grpo_tis.json index bca5a992..9d65c992 100644 --- a/tests/golden/decoupling/grpo_tis.json +++ b/tests/golden/decoupling/grpo_tis.json @@ -106,6 +106,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -243,7 +247,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "06b407103acaff8218e5bab2309192be3af014b6ca635dd7c72f27cca28b1a2b", + "miles_argv_sha256": "70d04ad1b86c39f254ac2631b2be4a32558c00f4e0e7cc8c1767e03ea84b03ee", "name": "grpo_tis", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/golden/decoupling/seq_adv_maxrl.json b/tests/golden/decoupling/seq_adv_maxrl.json index 80100673..29bd1d96 100644 --- a/tests/golden/decoupling/seq_adv_maxrl.json +++ b/tests/golden/decoupling/seq_adv_maxrl.json @@ -116,6 +116,10 @@ "bridge", "--model-name", "SimpleNamespace", + "--attention-dropout", + "0.0", + "--hidden-dropout", + "0.0", "--num-layers", "2", "--hidden-size", @@ -250,7 +254,7 @@ "10", "--group-query-attention" ], - "miles_argv_sha256": "8acfb1140db90c4d02fd4534526adcea04ca37a569125cee1952d3a697eb882d", + "miles_argv_sha256": "fd095a4b6d682aa8ba80b1953964b57575721859844217c78253847b7349a820", "name": "seq_adv_maxrl", "placement": { "allow_cross_node_engine_tp": false, diff --git a/tests/rl_cut_fakes.py b/tests/rl_cut_fakes.py index 8e9d097c..1e74d429 100644 --- a/tests/rl_cut_fakes.py +++ b/tests/rl_cut_fakes.py @@ -27,7 +27,7 @@ def named_parameters(self, model): def is_adapter(self, name): return "lora" in name - def export_optimizer(self, optimizer, named): + def export_optimizer(self, optimizer, named, *, on_device=False): # clones stay on their device by_id = {id(p): n for n, p in named} entries = {} for group in optimizer.param_groups: @@ -170,7 +170,7 @@ def __init__(self, coord=None, *, side_effect_free=True): super().__init__(coord) self.side_effect_free = side_effect_free - def export_optimizer(self, optimizer, named): + def export_optimizer(self, optimizer, named, *, on_device=False): # clones stay on their device from collections import defaultdict from yeto.rl.adapters.miles.cut_plugin import side_effect_free_state diff --git a/tests/rl_reshard_fakes.py b/tests/rl_reshard_fakes.py index 5d55ddba..cdf85a69 100644 --- a/tests/rl_reshard_fakes.py +++ b/tests/rl_reshard_fakes.py @@ -66,7 +66,7 @@ def named_parameters(self, model): def is_adapter(self, name): return "lora" in name - def export_optimizer(self, optimizer, named): + def export_optimizer(self, optimizer, named, *, on_device=False): # clones stay on their device optimizer.init_state() hyper = {k: v for k, v in optimizer.param_groups[0].items() if k != "params"} entries = {} 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 00000000..6a69f85f --- /dev/null +++ b/tests/test_e3_overlap_miles_role_stream.py @@ -0,0 +1,504 @@ +"""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 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 + + 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_options.py b/tests/test_e3_sync_options.py new file mode 100644 index 00000000..a3256227 --- /dev/null +++ b/tests/test_e3_sync_options.py @@ -0,0 +1,318 @@ +"""S26 E3 sync options: bf16 push transcoding, session binding, preset scoping. CPU only.""" + +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 ( # 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, +) + + +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(" 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 00000000..6421d6ce --- /dev/null +++ b/tests/test_e3_sync_overlap.py @@ -0,0 +1,370 @@ +"""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 rules (rollout policy age 0, fragment pipeline P). +""" + +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 OUTER_STEP_LAG, overlap_option_problems, rebase_values + +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]] = [] + self.push_rounds: list[tuple[int, int]] = [] # (global step, local rounds trained then) + + @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.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) + 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)) + + +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 = 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, + 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_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) + # 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_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 + + 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 "--rl-max-policy-age" not in " ".join(_ports_infra_flags(args)) + plain = _dense() + assert "--sweep-overlap" not in _rl_trainable_island_flags(plain, 0, 1) + + +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_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, 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(): + 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_cut_distopt_pp_ep.py b/tests/test_rl_cut_distopt_pp_ep.py index d687e8d3..849e288e 100644 --- a/tests/test_rl_cut_distopt_pp_ep.py +++ b/tests/test_rl_cut_distopt_pp_ep.py @@ -31,6 +31,9 @@ DPS_SOURCES = { "m5-c35702e": (os.environ.get("YETO_MILES_CHECKOUT", "/home/michael/work/miles-m3"), "c35702e"), "s25-distopt": (os.environ.get("YETO_MILES_DISTOPT_CHECKOUT", "/home/michael/work/miles-next"), "s25-distopt"), + # E1: + export_named_optimizer_state(device=None) on the image pin 44d41efd (not pushed) + "e1-device-export": (os.environ.get("YETO_MILES_DISTOPT_CHECKOUT", "/home/michael/work/miles-next"), + "yeto/s26-e1-device-export"), } @@ -355,3 +358,166 @@ def test_local_states_cover_with_fillers(): assert cp.range_problems(state, state) == [] other = {"entries": {"a": {**state["entries"]["a"], "start": 5}}} assert len(cp.range_problems(state, other)) == 2 + + +# -------------------------------------------------------------------------- +# E1: background cuts copy their state device-to-device first +# -------------------------------------------------------------------------- + + +@needs_dps +@pytest.mark.parametrize("layout", ["pp2dp1", "pp2dp2", "ep2"]) +def test_device_export_is_the_fork_host_export_once_on_the_host(layout): + ranks = _world(layout) + _train(ranks) + for rank in ranks: + named = list(rank.model[0].named_parameters()) + host = DPS.export_named_optimizer_state(rank.optimizer, named) + device = cp.to_host(cp._device_export(DPS)(rank.optimizer, named)) + assert cp.state_digest(device) == cp.state_digest(host) + assert device["format"] == host["format"] and list(device["entries"]) == list(host["entries"]) + for name, e in host["entries"].items(): + got = device["entries"][name] + assert {k: v for k, v in got.items() if k not in ("tensors", "scalars")} == \ + {k: v for k, v in e.items() if k not in ("tensors", "scalars")} + for part in ("tensors", "scalars"): + assert list(got[part]) == list(e[part]) + for key, t in e[part].items(): + assert got[part][key].dtype == t.dtype and torch.equal(got[part][key], t), (name, key) + + +def test_device_export_uses_the_fork_device_argument_when_it_has_one(): + calls = [] + + def export(optimizer, named, *, device="cpu"): + calls.append(device) + return {"entries": {}} + + cp._device_export(SimpleNamespace(export_named_optimizer_state=export))("opt", []) + assert calls == [None] + + +@needs_dps +@pytest.mark.parametrize("layout", ["pp2dp2", "ep2"]) +def test_background_device_capture_writes_the_inline_shard(tmp_path, layout, monkeypatch): + """Same shard (file bytes, digests, N7.6 resume state) from the device path as from the + inline host path, although the trainer trains on before the writer copies to the host.""" + ranks = _world(layout) + _train(ranks) + inline = [cp.save_cut_shard(r, directory=str(tmp_path / "inline"), cut_id="c1") for r in ranks] + monkeypatch.setattr(cp, "_device_copy_plan", lambda actor, named: (True, "ok")) + device = [] + for r, rank in enumerate(ranks): # one pending write per cut id per process: rank by rank + started = cp.save_cut_shard(rank, directory=str(tmp_path / "device"), cut_id="c1", background=True) + assert started["capture"] == {"mode": "device", "reason": "ok"}, started + for k, leaf in enumerate(_leaves(rank)): # the next step changes the live state first + leaf.train(2000 + 10 * r + k) + rank.opt_param_scheduler.num_steps = 16 + device.append(cp.finish_cut_shard(rank, cut_id="c1")) + assert [d.pop("writer") for d in device] == [{"mode": "process"}] * len(ranks) + assert device == inline + for name in sorted(p.name for p in (tmp_path / "inline").iterdir()): + assert (tmp_path / "device" / name).read_bytes() == (tmp_path / "inline" / name).read_bytes() + + +def test_background_capture_falls_back_to_the_host_and_says_why(tmp_path, monkeypatch): + from tests.rl_cut_fakes import make_rank, train_step + + rank = make_rank(0) + train_step(rank, torch.randn(3, 6, generator=torch.Generator().manual_seed(7))) + inline = cp.save_cut_shard(rank, directory=str(tmp_path / "inline"), cut_id="c1") + monkeypatch.setattr(cp, "_device_copy_plan", lambda actor, named: (False, "low_device_memory(free=1,need=2)")) + started = cp.save_cut_shard(rank, directory=str(tmp_path / "low"), cut_id="c1", background=True) + assert started["capture"] == {"mode": "host", "reason": "low_device_memory(free=1,need=2)"} + assert _without_writer(cp.finish_cut_shard(rank, cut_id="c1")) == inline + + real = cp._snapshot + + def oom(actor, backend, named, *, on_device=False): + if on_device: + raise torch.cuda.OutOfMemoryError("CUDA out of memory") + return real(actor, backend, named) + + monkeypatch.setattr(cp, "_device_copy_plan", lambda actor, named: (True, "ok")) + monkeypatch.setattr(cp, "_snapshot", oom) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: None) + started = cp.save_cut_shard(rank, directory=str(tmp_path / "oom"), cut_id="c1", background=True) + assert started["capture"] == {"mode": "host", "reason": "device_oom"} + assert _without_writer(cp.finish_cut_shard(rank, cut_id="c1")) == inline + + +def _without_writer(summary): + assert summary.pop("writer")["mode"] == "process", summary + return summary + + +def _one_rank_cut(tmp_path, name, *, background): + import random + + import numpy as np + + from tests.rl_cut_fakes import make_rank, train_step + + random.seed(5), np.random.seed(5), torch.manual_seed(5) # the shard holds the global RNG state + rank = make_rank(0) + train_step(rank, torch.randn(3, 6, generator=torch.Generator().manual_seed(7))) + if not background: + return cp.save_cut_shard(rank, directory=str(tmp_path / name), cut_id="c1") + cp.save_cut_shard(rank, directory=str(tmp_path / name), cut_id="c1", background=True) + return cp.finish_cut_shard(rank, cut_id="c1") + + +def test_a_dead_writer_process_falls_back_to_the_thread_and_says_why(tmp_path, monkeypatch): + """The writer process dies before it writes (its start fails): this rank's thread + writes the same shard and the summary records the fallback.""" + import multiprocessing + import os + from concurrent.futures import ProcessPoolExecutor + + inline = _one_rank_cut(tmp_path, "inline", background=False) + dead = ProcessPoolExecutor(max_workers=1, mp_context=multiprocessing.get_context("spawn"), + initializer=os._exit, initargs=(3,)) + monkeypatch.setattr(cp, "_WRITER", [dead]) + got = _one_rank_cut(tmp_path, "dead", background=True) + writer = got.pop("writer") + assert writer["mode"] == "thread" and "writer process failed" in writer["fallback"], writer + assert got == inline + assert cp._WRITER == [] # the broken pool is dropped; the next cut starts a new one + assert (tmp_path / "dead" / got["path"]).read_bytes() == (tmp_path / "inline" / got["path"]).read_bytes() + + +def test_a_partial_file_of_a_dead_writer_is_removed_before_the_fallback(tmp_path, monkeypatch): + inline = _one_rank_cut(tmp_path, "inline", background=False) + target = tmp_path / "partial" + target.mkdir() + (target / (inline["path"] + ".tmp")).write_bytes(b"half") + (target / inline["path"]).write_bytes(b"whole but unreported") + + class Broken: + def submit(self, *a, **k): + raise RuntimeError("writer died") + + def shutdown(self, **k): + pass + + monkeypatch.setattr(cp, "_WRITER", [Broken()]) + got = _one_rank_cut(tmp_path, "partial", background=True) + assert got.pop("writer") == {"mode": "thread", "fallback": "writer process failed: RuntimeError('writer died')"} + assert got == inline and not (target / (inline["path"] + ".tmp")).exists() + + +def test_the_writer_process_can_be_turned_off(tmp_path, monkeypatch): + inline = _one_rank_cut(tmp_path, "inline", background=False) + monkeypatch.setenv(cp.WRITER_PROCESS_ENV, "0") + got = _one_rank_cut(tmp_path, "thread", background=True) + assert got.pop("writer") == {"mode": "thread", "fallback": "disabled"} + assert got == inline + + +def test_device_copy_plan_needs_cuda_and_can_be_turned_off(monkeypatch): + named = [("lora_a", torch.nn.Parameter(torch.zeros(2)))] + assert cp._device_copy_plan(None, named) == (False, "not_cuda") + sleeper = SimpleNamespace(args=SimpleNamespace(offload_train=True)) + assert cp._device_copy_plan(sleeper, named) == (False, "offload_train") + monkeypatch.setenv(cp.DEVICE_COPY_ENV, "0") + assert cp._device_copy_plan(None, named) == (False, "disabled") diff --git a/tests/test_rl_dense_sweep_real_syncer.py b/tests/test_rl_dense_sweep_real_syncer.py index fd14e997..76351037 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 @@ -21,7 +22,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, 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") @@ -41,23 +48,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)) @@ -76,10 +85,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) -> SyncerClient: - return SyncerClient(("127.0.0.1", port), 0, _layout(), dtype=DTYPE_F32, num_streams=1, - connect_timeout=10, max_reconnects=0, finalization_timeout=20, - syncer_profile_hash=bytes.fromhex(profile_hash)) +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=max_reconnects, finalization_timeout=20, + syncer_profile_hash=bytes.fromhex(profile_hash), push_i8_block=push_i8_block) def _f32(values: list[float]) -> bytes: @@ -153,3 +163,583 @@ 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, + 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, 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() + 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 + 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 + 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)) + + +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) + + +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 + + +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" + + +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)) + + +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] + + +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] + + +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_rl_engine_driver.py b/tests/test_rl_engine_driver.py index cf48504d..b1bdefdb 100644 --- a/tests/test_rl_engine_driver.py +++ b/tests/test_rl_engine_driver.py @@ -964,3 +964,52 @@ def test_first_optimizer_step_records_the_last_start_phase_once(tmp_path): assert names == list(START_PHASES) [summary] = [f for e, f in emitted if e == "trainer_start_summary"] assert summary["total_s"] == 6.0 and summary["missing"] == [] + + +def test_colocated_publish_offloaded_prepares_a_due_cut_before_offloading(tmp_path): + """E1: with round cuts, the trainer part of a due cut is written while the trainer + is still resident (before the offload for the publication), not after it.""" + engine = _engine() + engine.trainer.publish_offloaded = True + sync = LocalOnlySync(2) + order = [] + + class _Cuts: + def resume(self, driver): + return None + + def due(self, rollout_id, *, final=False): + return True + + def prepare(self, driver, *, rollout_id, state, final=False): + order.append(("prepare", rollout_id, state.policy_version, final, driver._trainer_offloaded)) + + def save(self, driver, *, rollout_id, final=False, **kw): + order.append(("save", rollout_id, final)) + return None + + sync.round_cuts = _Cuts() + driver = _driver(engine, sync, tmp_path) + driver.run() + if PROFILE == "ports-only": # publish resident: nothing to prepare + assert not any(o[0] == "prepare" for o in order) + return + assert [o for o in order if o[1] >= 1][:2] == [("prepare", 1, 1, False, False), ("save", 1, False)] + assert ("prepare", 2, 2, True, False) in order + phases = [e["phase"] for e in _events(tmp_path / "events.jsonl") if e["event"] == "rl_driver_phase"] + assert phases[4:8] == ["train", "sync", "checkpoint", "offload"] + + +def test_rollout_offload_after_generate_runs_in_its_own_phase(tmp_path): + """E1: Miles --offload-rollout frees the engine memory after generating; the + driver times it as its own phase so "generate" matches Miles' rollout_time.""" + engine = _engine() + calls = [] + engine.rollout.offloads_after_generate = True + engine.rollout.offload_after_generate = lambda: calls.append("offload_rollout") + driver = _driver(engine, LocalOnlySync(2), tmp_path) + driver.run() + assert calls == ["offload_rollout"] * 2 and engine.rollout.defer_offload is True + phases = [e["phase"] for e in _events(tmp_path / "events.jsonl") if e["event"] == "rl_driver_phase"] + i = phases.index("generate") + assert phases[i + 1] == "rollout_offload" diff --git a/tests/test_rl_fully_async_prompts.py b/tests/test_rl_fully_async_prompts.py index c92c68a3..1130228f 100644 --- a/tests/test_rl_fully_async_prompts.py +++ b/tests/test_rl_fully_async_prompts.py @@ -87,7 +87,7 @@ def test_an_unexposed_stale_count_is_reported_unknown(executor, caplog): assert "reported unknown" in caplog.text -def test_cut_returns_every_untrained_prompt_and_stops_the_worker(executor): +def test_cut_records_every_untrained_prompt_and_keeps_the_worker_running(executor): source, ledger = executor.data_source, prompts.current() drawn = source.get_samples(5) # in flight / waiting in Miles' output buffer ledger.mark_trained(drawn[:2]) # trained this round @@ -102,18 +102,65 @@ def test_cut_returns_every_untrained_prompt_and_stops_the_worker(executor): _task_of(worker), loop).result() try: exported = export_in_flight(executor) + # E1: the worker keeps generating ahead across the cut + assert executor.generate_rollout._worker is not None + assert not executor.generate_rollout._worker.done() finally: loop.call_soon_threadsafe(loop.stop) - assert executor.generate_rollout._worker is None # the next call starts a fresh worker - assert source.offset == 5 # the cursor is not rewound: the prompts are in the buffer - assert sorted(g[0].group_index for g in source.buffer) == [2, 3, 4] - assert all(g[0].response == "" for g in source.buffer) # reset for a fresh generation - assert drawn[3][0].response == "partial" # a running task keeps its own sample + # the live run is untouched: same buffer, same in-flight samples + assert source.offset == 5 and [g[0].group_index for g in source.buffer] == [2] + assert drawn[3][0].response == "partial" + assert sorted(prompts.current().untrained) == [3, 4] + # the cut carries every drawn, untrained prompt, reset for a fresh generation assert exported["buffer_groups"] == 3 assert sorted({e["group_id"] for e in exported["entries"]}) == ["g2", "g3", "g4"] - # the resumed run draws exactly the untrained prompts again - assert sorted(g[0].group_index for g in source.get_samples(3)) == [2, 3, 4] + assert all(e["engine_state"]["miles_sample"]["response"] == "" for e in exported["entries"]) + + +def test_cut_reads_cursor_and_prompts_on_the_worker_loop(executor): + """The cursor and the untrained prompts come from one read on the loop the worker + draws on: a draw can never fall between them.""" + source = executor.data_source + for name in ("sample_offset", "epoch_id", "sample_group_index", "sample_index"): + setattr(source, name, 0) + real_get = prompts.current()._get_samples # the unwrapped DataSource.get_samples + + def get_samples(n): + groups = real_get(n) + source.sample_offset += n + source.sample_group_index += n + source.sample_index += n + return groups + + prompts.current()._get_samples = get_samples + loop = asyncio.new_event_loop() + thread = threading.Thread(target=loop.run_forever, daemon=True) + thread.start() + seen = [] + + async def drawing_worker(): + while True: + source.get_samples(1) # the ledger's wrapper (records the draw) + await asyncio.sleep(0) + + executor.generate_rollout._worker = asyncio.run_coroutine_threadsafe( + _start(drawing_worker), loop).result() + try: + for _ in range(20): + exported = export_in_flight(executor) + seen.append((exported["data_cursor"]["sample_offset"], exported["buffer_groups"])) + finally: + executor.generate_rollout._worker.get_loop().call_soon_threadsafe( + executor.generate_rollout._worker.cancel) + loop.call_soon_threadsafe(loop.stop) + # nothing is trained: every drawn prompt is in the cut, and the cursor is just past it + assert all(offset == groups for offset, groups in seen) + assert seen[-1][0] > seen[0][0] # the worker kept drawing between cuts + + +async def _start(fn): + return asyncio.create_task(fn()) async def _task_of(future): @@ -133,7 +180,7 @@ def test_cut_without_the_ledger_is_refused(): def test_overlap_round_report_carries_the_ledger_counts(monkeypatch): from yeto.rl.adapters.miles import rollout_meta_hook as hook - monkeypatch.setattr(hook, "build_metadata", lambda args, samples: { + monkeypatch.setattr(hook, "build_metadata", lambda args, samples, **_kw: { "completed": 2, "filtered": 0, "max_policy_age": 1}) sent = [] monkeypatch.setattr(hook, "put_to_sink", sent.append) diff --git a/tests/test_rl_meta_hook_sink_view.py b/tests/test_rl_meta_hook_sink_view.py new file mode 100644 index 00000000..d66646c3 --- /dev/null +++ b/tests/test_rl_meta_hook_sink_view.py @@ -0,0 +1,80 @@ +"""E1: one metadata-hook run reads each sink value once and looks the sink actor up once.""" + +from __future__ import annotations + +import sys +from types import SimpleNamespace + +import pytest + +from yeto.rl.adapters.miles import rollout_meta_hook as hook + + +class _Remote: + def __init__(self, fn): + self.remote = fn + + +class _FakeRay: + def __init__(self): + self.values = {"get_token": "yeto:4:" + "a" * 64, "get_limit": 1} + self.lookups = 0 + self.calls = [] + + def get_actor(self, name): + self.lookups += 1 + ray = self + + class Actor: + def __getattr__(self, method): + def call(*args): + ray.calls.append(method) + if method.startswith("set_"): + ray.values["get_" + method[4:]] = args[0] + return None + return ray.values.get(method) + return _Remote(call) + + return Actor() + + def get(self, value): + return value + + +@pytest.fixture() +def fake_ray(monkeypatch): + ray = _FakeRay() + monkeypatch.setitem(sys.modules, "ray", SimpleNamespace(get_actor=ray.get_actor, get=ray.get)) + monkeypatch.setenv(hook.META_SINK_ENV, "ray:yeto_rollout_meta") + return ray + + +def test_reads_inside_one_view_hit_the_actor_once(fake_ray): + with hook.one_sink_view(): + for _ in range(5): + assert hook.current_policy_token() == "yeto:4:" + "a" * 64 + assert hook.max_policy_age_override() == 1 + with hook.one_sink_view(): # nested: the same view + hook.current_policy_token() + assert fake_ray.calls == ["get_token", "get_limit"] and fake_ray.lookups == 1 + + +def test_outside_a_view_every_read_goes_to_the_actor(fake_ray): + hook.current_policy_token() + hook.current_policy_token() + assert fake_ray.calls == ["get_token", "get_token"] and fake_ray.lookups == 2 + + +def test_a_write_inside_the_view_is_read_back(fake_ray): + with hook.one_sink_view(): + assert hook.max_policy_age_override() == 1 + hook.put_max_policy_age(0) + assert hook.max_policy_age_override() == 0 + assert hook.max_policy_age_override() == 0 + + +def test_the_hook_limit_is_read_once_per_run(fake_ray): + args = SimpleNamespace(yeto_rl_max_policy_age=2) + with hook.one_sink_view(): + assert [hook._carry.max_policy_age(args) for _ in range(64)] == [1] * 64 + assert fake_ray.calls.count("get_limit") == 1 diff --git a/tests/test_rl_miles_adapter_config.py b/tests/test_rl_miles_adapter_config.py index 503fe2e2..7cbe30d9 100644 --- a/tests/test_rl_miles_adapter_config.py +++ b/tests/test_rl_miles_adapter_config.py @@ -795,3 +795,21 @@ def test_dynamic_batching_refused_for_bshd_gated_delta_net_models(): recompute_only = dataclasses.replace( cfg, model_recipe=recipe, train_memory=rc.TrainMemoryConfig(activation_recompute=True)) assert "--recompute-granularity" in mc.translate_run_config(recompute_only, AlgorithmSpec()).argv + + +def test_gradient_accumulation_fusion_is_opt_in(): + # S26 E0.1: the Miles recipes train with the fusion on; yeto keeps it off by default + off = " ".join(mc.translate_run_config(make_config(), AlgorithmSpec()).argv) + assert "--no-gradient-accumulation-fusion" in off + cfg = dataclasses.replace(make_config(), train_memory=rc.TrainMemoryConfig(gradient_accumulation_fusion=True)) + assert "--no-gradient-accumulation-fusion" not in mc.translate_run_config(cfg, AlgorithmSpec()).argv + assert rc._train_memory(SimpleNamespace(seq_len=8, rl_gradient_accumulation_fusion=True)) \ + .gradient_accumulation_fusion + + +def test_generic_recipe_trains_without_dropout(): + cfg = make_config() + if cfg.model_recipe.name != rc.RECIPE_GENERIC: + cfg = sub(cfg, "model_recipe", name=rc.RECIPE_GENERIC) + joined = " ".join(mc.translate_run_config(cfg, AlgorithmSpec()).argv) + assert "--attention-dropout 0.0 --hidden-dropout 0.0" in joined diff --git a/tests/test_rl_miles_carry_over.py b/tests/test_rl_miles_carry_over.py index e94f3d26..01c13a9b 100644 --- a/tests/test_rl_miles_carry_over.py +++ b/tests/test_rl_miles_carry_over.py @@ -348,3 +348,47 @@ def test_trajectory_records_carry_start_round_and_segments(sink): for r in recs: assert not [p for p in validate_trajectory_reward({**r, "rollout_id": 4, "policy_version": 4}) if "started_rollout_id" in p or "policy_versions" in p] + + +def test_deferred_cross_version_estimate_arrives_as_a_late_record(sink, tmp_path): + """E1 (fully_async): the estimate runs in the background and reaches the sink as + a late record of its rollout id; the round's fields only say it was deferred.""" + import glob + import json + import time + + engine = FakePartialEngine() + args = SimpleNamespace(yeto_rl_max_policy_age=1, tis_clip=2.0, tis_clip_low=0.0) + s = Sample(index=0, group_index=0) + engine.generate(s, 3, 4, finish=False) + engine.generate(s, 4, 2, finish=True, reward=1.0) + current = [0.5, -0.5, 0.0, 0.4, -0.5, -0.5] + hook.put_policy_token(tok(4), sink) + fields = hook.carry_over_fields(args, [s], sink, lambda _s: current, defer_estimate_for=4) + assert fields["cross_version_estimate"] == "deferred" + assert "cross_version_truncated_fraction" not in fields + handle = handle_from_metadata({**round_metadata(args, [[s]], sink), **fields}, rollout_id=4, + policy_version=4, policy_hash=HASHES[4], data_pack=None) + assert handle.cross_version_truncated_fraction is None + deadline = time.monotonic() + 5 + late = [] + while not late and time.monotonic() < deadline: + late = [json.loads(open(p).read()) for p in glob.glob(f"{tmp_path}/*.json")] + late = [r for r in late if r.get("schema") == hook.LATE_META_SCHEMA] + time.sleep(0.02) + assert late and late[0]["rollout_id"] == 4 + assert math.isclose(late[0]["cross_version_truncated_fraction"], 0.5) + + +def test_driver_governs_on_a_late_rollout_estimate(tmp_path): + """E1: a fully_async round's estimate arrives after its train step; the policy-age + governor still sees it.""" + from tests.test_rl_policy_age import _driver + + driver = _driver(tmp_path, 1) + calls, asked = [], [] + driver.rollout.set_max_policy_age = calls.append + driver.rollout.late_cross_version_fraction = lambda rid: asked.append(rid) or 0.9 + driver.run() + events = [json.loads(l) for l in (tmp_path / "e.jsonl").read_text().splitlines()] + assert asked and [e for e in events if e["event"] == "rl_policy_age_fallback"] and calls == [0] diff --git a/tests/test_rl_ports_image.py b/tests/test_rl_ports_image.py index 15585c2c..c6b755fc 100644 --- a/tests/test_rl_ports_image.py +++ b/tests/test_rl_ports_image.py @@ -631,3 +631,33 @@ def test_modal_public_image_config_has_no_login(monkeypatch, no_login): state = fake_modal(monkeypatch) mr.ModalOps("img-test").define(cfg) assert state["images"][0].calls[0] == ("from_registry", (cfg.image_ref,), {}) + + +def test_gpu_model_pin_refuses_an_a10g_for_an_a10(): + """E1: Modal "A10G" containers are an NVIDIA A10 or an A10G; --modal-gpu-name pins one.""" + from yeto.gpu_spec import parse_gpu_spec + + mr.check_gpu_model("NVIDIA A10", ["NVIDIA A10"] * 2, 2) + for names, count in ((["NVIDIA A10G"], 1), (["NVIDIA A10", "NVIDIA A10G"], 2), (["NVIDIA A10"], 2)): + with pytest.raises(RuntimeError, match="--modal-gpu-name"): + mr.check_gpu_model("NVIDIA A10", names, count) + cfg = _cfg(gpu="A10G", gpus_per_node=1, gpu_name="NVIDIA A10") + assert mr.ModalIslandConfig.from_json(cfg.to_json()).gpu_name == "NVIDIA A10" + old = json.loads(_cfg().to_json()) + old.pop("gpu_name") + assert mr.ModalIslandConfig.from_json(json.dumps(old)).gpu_name is None + args = _cli(("--modal-gpu-name", "NVIDIA A10")) + assert args.modal_gpu_name == "NVIDIA A10" + with pytest.raises(ValueError, match="--modal-gpu-name applies only to Modal"): + launcher.require_modal_for_gpu_exact(args, parse_gpu_spec("modal:1xa10g,aws:1xa100@us-east-1")) + + +def test_the_island_config_carries_the_gpu_model_pin(monkeypatch, login): + from yeto.gpu_spec import parse_gpu_spec + + args = _cli(("--modal-gpu-name", "NVIDIA A10")) + _prepare_rl_args(args) + task = _island_task(args, monkeypatch) + args.cluster_prefix = "img-test" + cfg = build_modal_island_config(args, parse_gpu_spec("modal:1xa10g")[0], 0, task, "1.2.3.4:5") + assert cfg.gpu_name == "NVIDIA A10" diff --git a/tests/test_rl_publish_fastpath.py b/tests/test_rl_publish_fastpath.py index bc88ee05..54877fd8 100644 --- a/tests/test_rl_publish_fastpath.py +++ b/tests/test_rl_publish_fastpath.py @@ -257,3 +257,108 @@ def test_resident_state_type_guards(): d = PolicyDigest("a" * 64, "b" * 64, 8, ((NAMES[0], (2, 1)),)) with pytest.raises(ValueError): TrainerResidentState("full", REV, CFG, "c" * 64, 0, d, lambda v: None) + + +# E1: the digest path stages the tensors in one reusable host buffer +def _mixed_tensors(seed=0): + g = torch.Generator().manual_seed(seed) + return { + "b.bf16": torch.randn(8, 3, generator=g).to(torch.bfloat16), + "a.f32_transposed": torch.randn(5, 7, generator=g).t(), + "c.f16": torch.randn(4, generator=g).to(torch.float16), + "d.scalar": torch.tensor(1.5), + } + + +def test_staged_copies_are_the_per_tensor_copies_bit_for_bit(): + from yeto.rl.engine.policy_digest import HostStaging, stage_f32_on_host + + tensors = _mixed_tensors() + plain = stage_f32_on_host(tensors) + staging = HostStaging(pin=False) + staged = stage_f32_on_host(tensors, staging) + assert list(staged) == sorted(tensors) == list(plain) + for name, value in tensors.items(): + want = value.detach().to("cpu", torch.float32).contiguous() + for got in (plain[name], staged[name]): + assert got.dtype == torch.float32 and got.is_contiguous() and got.shape == want.shape + assert torch.equal(got.view(torch.int32) if got.dim() else got, want.view(torch.int32) if want.dim() else want) + kw = dict(base_model_revision=REV, lora_config_hash=CFG, layout_hash="0" * 64) + a, b = digest_canonical_tensors(staged, **kw), digest_canonical_tensors(plain, **kw) + assert (a.policy_tensor_hash, a.payload_hash, a.payload_bytes, a.specs) == \ + (b.policy_tensor_hash, b.payload_hash, b.payload_bytes, b.specs) + # the buffer is reused (and grown) by the next call; each call's values are its own + bigger = {**_mixed_tensors(1), "e.more": torch.randn(64)} + again = stage_f32_on_host(bigger, staging) + assert all(torch.equal(again[n], stage_f32_on_host(bigger)[n]) for n in bigger) + + +@pytest.mark.parametrize("staged", [False, True]) +def test_staging_names_the_first_non_finite_tensor(staged): + from yeto.rl.engine.policy_digest import HostStaging, NonFiniteTensor, stage_f32_on_host + + tensors = _mixed_tensors() + tensors["c.f16"][1] = float("inf") + tensors["b.bf16"][0, 0] = float("nan") + with pytest.raises(NonFiniteTensor, match="'b.bf16' contains NaN or Inf"): + stage_f32_on_host(tensors, HostStaging(pin=False) if staged else None) + + +def test_export_digest_is_stable_across_rounds_with_the_reused_buffer(): + ranks, group, ps = make() + for r in ranks: + r.train_step() + first = ps.export_digest(policy_version=2).policy_tensor_hash() + assert ps.export_digest(policy_version=2).policy_tensor_hash() == first == \ + ps.export(policy_version=2).policy_tensor_hash() + + +# E1: sync colocated mode hashes the staged policy while the trainer offloads +def test_deferred_digest_is_the_eager_digest_and_is_collected_once(): + ranks, group, ps = make() + for r in ranks: + r.train_step() + eager = ps.export_digest(policy_version=2) + group.calls.clear() + deferred = ps.export_digest(policy_version=2, deferred=True) + assert group.calls == [sp.EXPORT_DIGEST] and not deferred.digest.is_resolved + assert deferred.tensor_names == eager.tensor_names and deferred.layout_hash == eager.layout_hash + for r in ranks: # the trainer trains on (or offloads): the staged copy is hashed + r.train_step() + assert deferred.policy_tensor_hash() == eager.policy_tensor_hash() + assert _payload_of(deferred) == _payload_of(eager) + assert deferred.digest.to_wire() == {**eager.digest.to_wire(), "hash_seconds": deferred.digest.hash_seconds} + assert group.calls == [sp.EXPORT_DIGEST, sp.COLLECT_DIGEST] + deferred.policy_tensor_hash() + assert group.calls == [sp.EXPORT_DIGEST, sp.COLLECT_DIGEST] # resolved once + + +def test_a_new_export_waits_for_the_deferred_hash_of_the_reused_buffer(): + ranks, group, ps = make() + first = ps.export_digest(policy_version=1) + deferred = ps.export_digest(policy_version=1, deferred=True) + for r in ranks: + r.train_step() + other = ps.export_digest(policy_version=2) # overwrites the staging buffer + assert other.policy_tensor_hash() != first.policy_tensor_hash() + assert deferred.policy_tensor_hash() == first.policy_tensor_hash() + + +def test_collecting_another_version_is_refused(): + ranks, group, ps = make() + ps.export_digest(policy_version=1, deferred=True) + with pytest.raises(Exception, match="v1, not v3"): + ps._collect_digest(3) + + +def test_the_driver_defers_only_for_a_colocated_offloading_trainer(monkeypatch): + ranks, group, ps = make() + driver = _Driver(ps, fast=False) + from yeto.rl.engine.driver import defer_digest + + assert defer_digest(driver) is False + driver.colocated = driver.publish_offloaded = True + state = IslandDriver.export_local_resident(driver, policy_version=0) + assert not state.digest.is_resolved + monkeypatch.setenv("YETO_RL_DEFERRED_DIGEST", "0") + assert defer_digest(driver) is False diff --git a/tests/test_rl_resume_from_checkpoint.py b/tests/test_rl_resume_from_checkpoint.py index cf6b44a4..c3e43461 100644 --- a/tests/test_rl_resume_from_checkpoint.py +++ b/tests/test_rl_resume_from_checkpoint.py @@ -5,6 +5,7 @@ from __future__ import annotations import json +from pathlib import Path from types import SimpleNamespace import pytest @@ -193,7 +194,9 @@ def test_save_kill_resume_without_elastic_continues_bitwise(tmp_path): store = _store(tmp_path / "store", vol) rank_a, ctl_a, driver_a, cp_a = _trained_island(tmp_path, store, 2) info = cp_a.save(driver_a, rollout_id=2) - assert info["store_synced"] and info["cut_bytes"] > 0 and info["lr_at_next_round"] == [5e-6] + # E1: the store's first cut is complete inline (a store with cuts but no LATEST refuses to start) + assert info["store_synced"] is True and info["store_commit"] == "inline" + assert info["cut_bytes"] > 0 and info["lr_at_next_round"] == [5e-6] assert info["run_fingerprint"] == FP and info["consumed_groups"] == 8 and info["incarnation_index"] == 0 assert vol.commits >= 1 and rs.read_latest(store.root)["cut_id"] == info["cut_id"] @@ -541,3 +544,315 @@ def test_resume_wiring_carries_the_runtime_fingerprint(tmp_path): w = resume_wiring_for(SimpleNamespace(yeto_rl_resume={**cfg, "runtime_fingerprint": "miles@abc"}), sync=LocalOnlySync(3), environ={}) assert w.controller.runtime_fingerprint == "miles@abc" and w.every == 2 and w.ledger is not None + + +# ---------------------------------------------------------------- E1: cut before offload, background cut +class _CountingTrainer: + """The real MilesTrainerGroup cut path, counting the cut and onload/offload calls.""" + + def __init__(self, inner): + self.inner = inner + self.calls: list[str] = [] + + def save_cut(self, **kw): + self.calls.append("save_cut") + return self.inner.save_cut(**kw) + + def start_cut(self, **kw): + self.calls.append("start_cut") + return self.inner.start_cut(**kw) + + def finish_cut(self, pending): + self.calls.append("finish_cut") + return self.inner.finish_cut(pending) + + def collect_cut(self, pending): + self.calls.append("collect_cut") + return self.inner.collect_cut(pending) + + def onload(self): + self.calls.append("onload") + + def offload(self): + self.calls.append("offload") + + def __getattr__(self, name): + return getattr(self.inner, name) + + +def _publish(driver, version, policy_hash): + from yeto.rl.engine.driver import policy_token + + driver.published_version = version + driver.published_state = SimpleNamespace(policy_tensor_hash=lambda: policy_hash, policy_version=version) + driver.expected_token = policy_token(version, policy_hash) + + +def _island_with_first_cut(tmp_path, store, keep=2): + """One trained round, cut inline at v1 (the store's first cut); then a second + trained round whose cut (v2) is the periodic one under test.""" + rank = make_rank(0) + g = torch.Generator().manual_seed(3) + train_step(rank, torch.randn(3, 6, generator=g)) + ctl = _rctl(tmp_path, "a", store, keep=keep) + driver = _Driver(rank, published=None, local_step=1) + driver.last_applied_lrs = [5e-6] + trainer = _CountingTrainer(_trainer(rank)) + cp = _cp(ctl, driver, trainer, _Ledger(2)) + _publish(driver, 1, _hash(rank)) + first = cp.save(driver, rollout_id=1) + assert first["store_commit"] == "inline" and first["store_synced"] is True + train_step(rank, torch.randn(3, 6, generator=g)) + driver.local_step = 2 + trainer.calls.clear() + return rank, driver, trainer, cp, first + + +def _resume_fresh(tmp_path, store, rank): + rank_b = make_rank(9) + driver_b = _Driver(rank_b) + cp_b = _cp(_rctl(tmp_path, "b", store), driver_b, _trainer(rank_b), _Ledger(2)) + got = cp_b.resume(driver_b) + for n, v in params(rank).items(): + assert torch.equal(v, params(rank_b)[n]) + return got + + +def test_prepared_cut_is_written_in_the_background_without_waking_the_trainer(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, first = _island_with_first_cut(tmp_path, store) + driver._trainer_offloaded = True + state = SimpleNamespace(policy_version=2, policy_tensor_hash=lambda: _hash(rank)) + assert cp.prepare(driver, rollout_id=2, state=state) is True + _publish(driver, 2, _hash(rank)) + info = cp.save(driver, rollout_id=2) + assert trainer.calls == ["start_cut"] # captured before the offload: no wake/sleep + assert info["prepared_before_offload"] is True and info["store_commit"] == "deferred" + assert cp.pointer()["cut_id"] == first["cut_id"] # not pointed at until complete + done = cp.join_commit(driver) + assert done["ok"] is True and done["cut_id"] == info["cut_id"] and trainer.calls[-1] == "collect_cut" + assert driver.events[-1][0] == "rl_cut_committed" and driver.events[-1][1]["ok"] is True + # E1: the event says where the shard was written (the rank's writer process) + assert driver.events[-1][1]["writer"] == [{"mode": "process"}] + assert cp.pointer()["cut_id"] == info["cut_id"] + assert rs.read_latest(store.root)["cut_id"] == info["cut_id"] + assert _resume_fresh(tmp_path, store, rank)["next_rollout_id"] == 2 + + +def test_prepared_cut_for_another_policy_is_dropped_and_cut_again(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, _first = _island_with_first_cut(tmp_path, store) + driver._trainer_offloaded = True + cp.prepare(driver, rollout_id=2, state=SimpleNamespace(policy_version=2, policy_tensor_hash=lambda: "f" * 64)) + _publish(driver, 2, _hash(rank)) # the publication names other weights + info = cp.save(driver, rollout_id=2) + assert trainer.calls == ["start_cut", "finish_cut", "onload", "start_cut", "offload"] + assert info["prepared_before_offload"] is False and info["policy_hash"] == _hash(rank) + assert cp.join_commit(driver)["ok"] is True + assert _resume_fresh(tmp_path, store, rank)["policy_hash"] == _hash(rank) + + +def test_final_cut_completes_inline_after_the_pending_one(tmp_path): + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, _first = _island_with_first_cut(tmp_path, store) + _publish(driver, 2, _hash(rank)) + info = cp.save(driver, rollout_id=2, final=True) + assert info["store_commit"] == "inline" and info["store_synced"] is True + assert trainer.calls == ["save_cut"] and cp.join_commit(driver) is None + + +def test_failed_background_commit_keeps_the_previous_cut(tmp_path): + class _FailingVol(_Vol): + fail = False + + def commit(self): + if self.fail: + raise OSError("volume commit failed") + super().commit() + + vol = _FailingVol() + store = _store(tmp_path / "store", vol) + rank, driver, _trainer_c, cp, first = _island_with_first_cut(tmp_path, store, keep=1) + _publish(driver, 2, _hash(rank)) + vol.fail = True + cp.save(driver, rollout_id=2) + result = cp.join_commit(driver) + assert result["ok"] is False and "volume commit failed" in result["error"] + assert driver.events[-1][0] == "rl_cut_committed" and driver.events[-1][1]["ok"] is False + # nothing was pruned: the previous (committed) cut is still in the store + assert (store.root / rs.ROUND_CUTS / first["cut_id"]).is_dir() + + +def test_background_cut_can_be_turned_off(tmp_path, monkeypatch): + monkeypatch.setenv("YETO_RL_CUT_BACKGROUND_COMMIT", "0") + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, _first = _island_with_first_cut(tmp_path, store) + _publish(driver, 2, _hash(rank)) + info = cp.save(driver, rollout_id=2) + assert info["store_commit"] == "inline" and info["store_synced"] is True and trainer.calls == ["save_cut"] + + +def test_cut_context_takes_the_cursor_read_with_the_in_flight_entries(): + """E1: a rollout that keeps drawing during the cut (Miles fully_async) returns the + cursor read together with its entries; the context uses that one.""" + rank = make_rank(0) + driver = _Driver(rank, published=2, local_step=2) + driver.profile = SimpleNamespace(max_policy_age=1) + live = {"sample_offset": 9, "epoch_id": 0, "sample_group_index": 9, "sample_index": 9} + with_entries = {**live, "sample_offset": 12, "sample_group_index": 12, "sample_index": 12} + pool = SimpleNamespace(data_cursor=lambda: dict(live), export_in_flight=lambda: { + "entries": [], "buffer_groups": 0, "data_cursor": dict(with_entries)}) + source = CutSource(driver=lambda: driver, trainer=None, rollout=pool, ledger=_Ledger(2), + algorithm=SimpleNamespace(sha256=lambda: "a" * 64, to_legacy_runtime_attrs=lambda: {}), + backend_fingerprint="miles@0af62f4d", cut_root="", global_batch_size=ARGS.global_batch_size) + data = source.context("r000002-abc").data + assert data["sample_offset"] == 12 and data["buffer_length"] == 0 + + +def test_periodic_cut_is_skipped_while_the_previous_commit_runs(tmp_path): + import threading + + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, first = _island_with_first_cut(tmp_path, store) + gate = threading.Event() + busy = threading.Thread(target=gate.wait, daemon=True) + busy.start() + cp._commit = (busy, {"cut_id": first["cut_id"]}) + _publish(driver, 2, _hash(rank)) + assert cp.save(driver, rollout_id=2) is None and trainer.calls == [] + assert driver.events[-1][0] == "rl_cut_skipped" and cp.skipped_busy == 1 + gate.set() + busy.join() + info = cp.save(driver, rollout_id=2, final=True) # the final cut waits and is taken + assert info["store_commit"] == "inline" + + +def test_a_periodic_cut_never_waits_for_the_commit_it_starts(tmp_path): + import threading + + class _SlowVol(_Vol): + gate = threading.Event() + slow = False + + def commit(self): + if self.slow: + self.gate.wait(10) + super().commit() + + vol = _SlowVol() + store = _store(tmp_path / "store", vol) + rank, driver, _trainer_c, cp, _first = _island_with_first_cut(tmp_path, store) + g = torch.Generator().manual_seed(5) + _publish(driver, 2, _hash(rank)) + assert cp.save(driver, rollout_id=2)["store_commit"] == "deferred" + vol.slow = True + train_step(rank, torch.randn(3, 6, generator=g)) + driver.local_step = 3 + _publish(driver, 3, _hash(rank)) + info = cp.save(driver, rollout_id=3) # completes v2 and starts its (slow) commit + assert info["store_commit"] == "deferred" and cp._commit[0].is_alive() + train_step(rank, torch.randn(3, 6, generator=g)) + driver.local_step = 4 + _publish(driver, 4, _hash(rank)) + assert cp.save(driver, rollout_id=4) is None # v2 still committing: skipped + assert driver.events[-1][0] == "rl_cut_skipped" + vol.gate.set() + final = cp.save(driver, rollout_id=4, final=True) + assert final["store_commit"] == "inline" and rs.read_latest(store.root)["cut_id"] == final["cut_id"] + committed = [f["cut_id"] for e, f in driver.events if e == "rl_cut_committed"] + assert committed == ["r000002-" + cp.controller.incarnation["id"], info["cut_id"]] # v2, then v3 + + +def test_prune_keeps_a_cut_newer_than_every_kept_one(tmp_path): + """E1: a periodic cut is captured while the previous one commits in the background; + the commit's prune must not delete the newer cut being written.""" + root = tmp_path / "store" + for seq, cut_id in ((1, "r000002-aa"), (2, "r000004-aa")): + (root / rs.STATE_DIR / f"{seq:08d}").mkdir(parents=True) + (root / rs.STATE_DIR / f"{seq:08d}" / "cut").write_text(cut_id) + for cut_id in ("r000001-aa", "r000002-aa", "r000004-aa", "r000006-aa"): + (root / rs.ROUND_CUTS / cut_id).mkdir(parents=True) + (root / rs.LATEST).write_text(json.dumps({"seq": 2})) + deleted = rs.prune(root, keep=1, cut_ids_of=lambda p: (p / "cut").read_text()) + assert sorted(deleted) == [f"{rs.ROUND_CUTS}/r000001-aa", f"{rs.ROUND_CUTS}/r000002-aa", + f"{rs.STATE_DIR}/00000001"] + assert sorted(p.name for p in (root / rs.ROUND_CUTS).iterdir()) == ["r000004-aa", "r000006-aa"] + + +def test_a_prepared_cut_skipped_while_committing_is_collected_at_the_next_cut(tmp_path): + """E1: the safe point does not wait for the trainer writer of a prepared cut it + skips (busy commit); the next cut collects that writer and drops its files.""" + import threading + + from yeto.rl.engine.cut import cut_dir + + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, first = _island_with_first_cut(tmp_path, store) + driver._trainer_offloaded = True + state = SimpleNamespace(policy_version=2, policy_tensor_hash=lambda: _hash(rank)) + assert cp.prepare(driver, rollout_id=2, state=state) is True + gate = threading.Event() + busy = threading.Thread(target=gate.wait, daemon=True) + busy.start() + cp._commit = (busy, {"cut_id": first["cut_id"]}) + _publish(driver, 2, _hash(rank)) + assert cp.save(driver, rollout_id=2) is None and driver.events[-1][0] == "rl_cut_skipped" + assert trainer.calls == ["start_cut"] # the skipped cut's writer is not waited for + skipped = cut_dir(cp.root, cp._cut_id(2)) + gate.set() + busy.join() + info = cp.save(driver, rollout_id=2, final=True) + # the skipped writer is collected first, then the final cut is taken afresh + assert trainer.calls[:2] == ["start_cut", "finish_cut"] and cp._abandoned == [] + assert info["store_commit"] == "inline" and info["cut_id"] == skipped.name + assert rs.read_latest(store.root)["cut_id"] == info["cut_id"] + + +def test_a_cut_declined_at_prepare_does_not_wake_the_offloaded_trainer(tmp_path): + """E1: prepare met a running commit; when the commit is done by the safe point the + periodic cut is skipped instead of waking the offloaded trainer for it.""" + import threading + + store = _store(tmp_path / "store", _Vol()) + rank, driver, trainer, cp, first = _island_with_first_cut(tmp_path, store) + driver._trainer_offloaded = True + gate = threading.Event() + busy = threading.Thread(target=gate.wait, daemon=True) + busy.start() + cp._commit = (busy, {"cut_id": first["cut_id"]}) + state = SimpleNamespace(policy_version=2, policy_tensor_hash=lambda: _hash(rank)) + assert cp.prepare(driver, rollout_id=2, state=state) is False + gate.set() + busy.join() + _publish(driver, 2, _hash(rank)) + assert cp.save(driver, rollout_id=2) is None and trainer.calls == [] + assert driver.events[-1][0] == "rl_cut_skipped" + assert driver.events[-1][1]["reason"] == "not prepared before the trainer offloaded" + # a cut that was never declined still wakes the trainer (no silent loss of cuts) + train_step(rank, torch.randn(3, 6, generator=torch.Generator().manual_seed(7))) + driver.local_step = 3 + _publish(driver, 3, _hash(rank)) + assert cp.save(driver, rollout_id=3)["trainer_onloaded_for_cut"] is True + + +def test_a_deferred_cut_snapshots_the_state_dir_as_it_was_at_the_cut(tmp_path): + """E1 (Modal A10G resume e1c): a periodic cut is synced into the store after later + rounds changed the state dir (the ledger); the snapshot must hold the state dir as + it was at the cut, else a resume refuses the ledger (consumed groups differ).""" + store = _store(tmp_path / "store", _Vol()) + rank, driver, _trainer_c, cp, _first = _island_with_first_cut(tmp_path, store) + marker = Path(cp.controller.state_dir) / "ledger-e1-marker.txt" + marker.write_text("at the cut") + _publish(driver, 2, _hash(rank)) + info = cp.save(driver, rollout_id=2) + assert info["store_commit"] == "deferred" + marker.write_text("one round later") # the next round trains before the commit runs + done = cp.join_commit(driver) + assert done["ok"] is True + latest = rs.read_latest(store.root) + assert latest["cut_id"] == info["cut_id"] + snap = store.root / rs.STATE_DIR / f"{int(latest['seq']):08d}" + assert (snap / marker.name).read_text() == "at the cut" + assert json.loads((snap / POINTER).read_text())["cut_id"] == info["cut_id"] + assert cp.pointer()["cut_id"] == info["cut_id"] + assert not list(Path(cp.controller.state_dir).parent.glob(".*.stage-*")) # staging removed diff --git a/tests/test_rl_sampler_trainer_pair_launch.py b/tests/test_rl_sampler_trainer_pair_launch.py index 0cbd7e69..df7d661a 100644 --- a/tests/test_rl_sampler_trainer_pair_launch.py +++ b/tests/test_rl_sampler_trainer_pair_launch.py @@ -109,11 +109,20 @@ def test_pair_island_commands_pass_the_island_parser_and_no_sync_checks(): assert (config.role, config.store) == (role, "/root/yeto-checkpoint-store/run-a") -def test_full_parameter_without_a_buffer_role_is_still_refused_on_the_island(): +def test_full_parameter_on_a_no_sync_island_with_both_roles_is_accepted(): + """S26 E11: a single no-sync island with both roles trains every parameter.""" from yeto.rl.adapters.miles import island_entry args = types.SimpleNamespace(parameter_mode="full", rl_single_island_no_sync=True, rl_sample_role="both") + assert island_entry._full_parameter_mode(args, "strict-avg") is True + + +def test_full_parameter_with_a_syncer_still_needs_the_full_parameter_preset(): + from yeto.rl.adapters.miles import island_entry + + args = types.SimpleNamespace(parameter_mode="full", rl_single_island_no_sync=False, + rl_sample_role="both") with pytest.raises(ValueError, match="full-parameter sync preset"): island_entry._full_parameter_mode(args, "strict-avg") diff --git a/tests/test_rl_sao_streaming_runtime.py b/tests/test_rl_sao_streaming_runtime.py index ced8e08e..88ec57b2 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_rl_trainer_cut.py b/tests/test_rl_trainer_cut.py index a2e1c8b7..643f04fd 100644 --- a/tests/test_rl_trainer_cut.py +++ b/tests/test_rl_trainer_cut.py @@ -284,8 +284,8 @@ def test_context_is_checked_before_any_rank_writes(tmp_path): def test_missing_optimizer_coverage_is_refused(tmp_path, monkeypatch): from yeto.rl.adapters.miles import cut_plugin - real = cut_plugin._save - monkeypatch.setattr(cut_plugin, "_save", lambda *a, **k: {**real(*a, **k), "optimizer_names": ["lora_A"]}) + real = cut_plugin._write + monkeypatch.setattr(cut_plugin, "_write", lambda *a, **k: {**real(*a, **k), "optimizer_names": ["lora_A"]}) with pytest.raises(CutError, match="no optimizer state"): _trainer(RankGroup([_trained_rank()])).save_cut(epoch=1, context=_context(tmp_path)) diff --git a/tests/test_rl_verl_adapter.py b/tests/test_rl_verl_adapter.py index 96c58845..1a417dcd 100644 --- a/tests/test_rl_verl_adapter.py +++ b/tests/test_rl_verl_adapter.py @@ -447,3 +447,27 @@ def test_verl_data_cursor_refuses_what_it_cannot_honour(): cur.seek({"sample_offset": 2}) with pytest.raises(RuntimeError, match="twice"): VerlDataCursor(trainer) + + +def test_fully_async_runtime_env_keeps_nccl_host_buffers_on_shm(monkeypatch): + # S26 E0.1: a cuMem host-buffer sender and a /dev/shm receiver cannot attach on Modal + import sys + import types + + pkg = types.ModuleType("verl.trainer.constants_ppo") + pkg.get_ppo_ray_runtime_env = lambda _x: {"env_vars": {}} + for name in ("verl", "verl.trainer"): + monkeypatch.setitem(sys.modules, name, sys.modules.get(name) or types.ModuleType(name)) + monkeypatch.setitem(sys.modules, "verl.trainer.constants_ppo", pkg) + ray_stub = types.ModuleType("ray") + ray_stub.remote = lambda *a, **k: (lambda cls: cls) + monkeypatch.setitem(sys.modules, "ray", sys.modules.get("ray") or ray_stub) + monkeypatch.setitem(sys.modules, "yeto.rl.adapters.verl.verl_main", None) + monkeypatch.delitem(sys.modules, "yeto.rl.adapters.verl.verl_main") + from yeto.rl.adapters.verl import verl_main + + monkeypatch.delenv(verl_main.CUMEM_HOST_ENV, raising=False) + assert verl_main._runtime_env(fully_async=True)["env_vars"][verl_main.CUMEM_HOST_ENV] == "0" + assert verl_main.CUMEM_HOST_ENV not in verl_main._runtime_env()["env_vars"] + monkeypatch.setenv(verl_main.CUMEM_HOST_ENV, "1") + assert verl_main._runtime_env(fully_async=True)["env_vars"][verl_main.CUMEM_HOST_ENV] == "1" diff --git a/tests/test_rl_verl_buffer_trainer.py b/tests/test_rl_verl_buffer_trainer.py index 58dea824..b37ad8a5 100644 --- a/tests/test_rl_verl_buffer_trainer.py +++ b/tests/test_rl_verl_buffer_trainer.py @@ -212,6 +212,12 @@ def yeto_param_version(self): def yeto_export_lora(self): return {k: v.clone() for k, v in self.tensors.items()}, {"count": 2} + def yeto_export_digest(self, kwargs): + from yeto.rl.adapters.verl.ports_impl import digest_lora_tensors + + self.digest_exports = getattr(self, "digest_exports", 0) + 1 + return digest_lora_tensors(dict(self.tensors), **kwargs) + def yeto_call_workers(self, fn, tensors, reset): for k in self.tensors: self.tensors[k] = tensors[k].clone().float() diff --git a/tests/test_rl_verl_fully_async_64b.py b/tests/test_rl_verl_fully_async_64b.py index 145cd8a2..2ab12db4 100644 --- a/tests/test_rl_verl_fully_async_64b.py +++ b/tests/test_rl_verl_fully_async_64b.py @@ -256,6 +256,12 @@ def yeto_push_weights(self): def yeto_export_lora(self): return {k: v.clone() for k, v in self.tensors.items()}, {"count": 2} + def yeto_export_digest(self, kwargs): + from yeto.rl.adapters.verl.ports_impl import digest_lora_tensors + + self.digest_exports = getattr(self, "digest_exports", 0) + 1 + return digest_lora_tensors(dict(self.tensors), **kwargs) + def yeto_call_workers(self, fn, tensors, reset): for k in self.tensors: self.tensors[k] = tensors[k].clone().float() @@ -335,6 +341,22 @@ def test_driver_runs_fully_async_rounds_with_carry_over_and_discard(tmp_path): checks = [e for e in events if e["event"] == "rl_publication_check"] assert all(e["status"] == "VERIFIED" for e in checks) and len(checks) == 4 assert [e["param_version"] for e in events if e["event"] == "rl_verl_version_map"] == [0, 1, 2, 3] + # rl-publish-fastpath: LocalOnlySync takes digests from the trainer, no tensor export + assert actor.digest_exports == 4 + + +def test_fastpath_off_publishes_the_same_hashes(tmp_path, monkeypatch): + def run(sub): + actor = FakeTrainerActor(tmp_path / sub / "rb", [_sample("a", 0, 0, [(0, 4)])] * 2) + _driver(tmp_path / sub, actor, rounds=2, required=1)[0].run() + return [(e["policy_tensor_hash"], e["sent_checksum"]) for e in _events(tmp_path / sub / "tape.jsonl") + if e["event"] == "rl_publication_check"], getattr(actor, "digest_exports", 0) + + fast, fast_digests = run("fast") + monkeypatch.setenv("YETO_RL_PUBLISH_FASTPATH", "0") + slow, slow_digests = run("slow") + assert fast == slow and len(fast) == 3 + assert fast_digests == 3 and slow_digests == 0 def test_publication_mismatch_on_the_rollouter_replica_fails(tmp_path): diff --git a/tests/test_rl_verl_publish_fastpath.py b/tests/test_rl_verl_publish_fastpath.py new file mode 100644 index 00000000..186a93e0 --- /dev/null +++ b/tests/test_rl_verl_publish_fastpath.py @@ -0,0 +1,117 @@ +"""rl-publish-fastpath for verl: the trainer-side digest equals the driver-side +full-export values bit for bit, and the publishers use it without tensors.""" + +from __future__ import annotations + +import torch + +from yeto.rl.adapters.verl import publish as pub +from yeto.rl.adapters.verl.ports_impl import (VerlPolicyState, VerlResidentState, digest_lora_tensors, + sent_checksum) +from yeto.rl.core import CanonicalTensorSpec, canonical_layout_hash, canonical_state, policy_tensor_hash +from yeto.rl.engine.policy_digest import is_resident +from yeto.rl.engine.trainable_state import TrainableState + +NAMES = ("model.layers.0.q_proj.lora_A.weight", "model.layers.0.q_proj.lora_B.weight") +SPECS = (CanonicalTensorSpec(NAMES[0], (4, 8), "float32", 32), + CanonicalTensorSpec(NAMES[1], (8, 4), "float32", 32)) +IDS = {"base_model_revision": "a" * 40, "lora_config_hash": "b" * 64, "layout_hash": canonical_layout_hash(SPECS)} + + +def _tensors(seed: int = 0) -> dict: + gen = torch.Generator().manual_seed(seed) + # bf16 originals upcast to f32, as remote_export_lora returns them + return {s.name: torch.randn(s.shape, generator=gen).to(torch.bfloat16).float() for s in SPECS} + + +class _Island: + def __init__(self, tensors): + self.tensors = tensors + self.base_model_revision, self.lora_config_hash = IDS["base_model_revision"], IDS["lora_config_hash"] + self.layout_hash, self.expected_specs = IDS["layout_hash"], SPECS + self.exports = 0 + + def export_tensors(self): + self.exports += 1 + return {k: v.clone() for k, v in self.tensors.items()}, {} + + def export_digest(self, kwargs): + return digest_lora_tensors(dict(self.tensors), **kwargs) + + +def test_trainer_digest_equals_the_full_export_values(): + tensors = _tensors() + raw = digest_lora_tensors(tensors, **IDS, expected_specs=SPECS) + full = canonical_state(3, tensors, **IDS, expected_specs=SPECS) + assert raw["digest"]["policy_tensor_hash"] == policy_tensor_hash(full) + assert raw["lora_checksum"] == pub.lora_checksum(dict(full.tensors)) + + +def test_resident_state_crosses_no_tensors_and_publishes_the_same_checksum(): + island = _Island(_tensors(1)) + ps = VerlPolicyState(island) + resident = ps.export_digest(policy_version=5) + assert isinstance(resident, VerlResidentState) and is_resident(resident) + assert island.exports == 0 + full = ps.export() + assert resident.policy_tensor_hash() == full.policy_tensor_hash() + assert sent_checksum(resident) == sent_checksum(full) + assert island.exports == 1 + # a re-versioned handle keeps the checksum (bridges.with_version) + assert sent_checksum(resident.with_version(6))[0] == sent_checksum(full)[0] + + +def test_materialize_falls_back_to_the_full_export_at_the_handle_version(): + island = _Island(_tensors(2)) + resident = VerlPolicyState(island).export_digest(policy_version=7) + state = resident.materialize() + assert isinstance(state, TrainableState) and state.policy_version == 7 + assert state.policy_tensor_hash() == resident.policy_tensor_hash() + assert island.exports == 1 + + +def test_island_without_a_digest_hook_falls_back(): + class Plain: + pass + + assert VerlPolicyState(Plain()).export_digest(policy_version=0) is None + + +def test_driver_switch_turns_the_fast_path_off(monkeypatch): + from yeto.rl.engine.driver import IslandDriver + + driver = IslandDriver.__new__(IslandDriver) + driver.policy_state = VerlPolicyState(_Island(_tensors())) + assert is_resident(driver.export_local_resident(policy_version=0)) + monkeypatch.setenv("YETO_RL_PUBLISH_FASTPATH", "0") + assert driver.export_local_resident(policy_version=0) is None + + +def test_worker_digest_with_the_staged_export_equals_the_full_export(): + """E1: the worker stages its LoRA in one host buffer for the digest; the digest + and the bf16 send checksum equal those of the full export, round after round.""" + from types import SimpleNamespace + + from yeto.rl.adapters.verl import ports_impl as pi + + gen = torch.Generator().manual_seed(3) + params = { + "base_model.model.layers.0.q_proj.lora_A.default.weight": + torch.nn.Parameter(torch.randn(4, 8, generator=gen).to(torch.bfloat16)), + "base_model.model.layers.0.q_proj.lora_B.default.weight": + torch.nn.Parameter(torch.randn(4, 8, generator=gen).t()), # not contiguous + "base_model.model.layers.0.q_proj.base_layer.weight": torch.nn.Parameter(torch.randn(2, 2, generator=gen)), + } + peft = SimpleNamespace(named_parameters=lambda: list(params.items())) + worker = SimpleNamespace(actor=SimpleNamespace(engine=SimpleNamespace(module=peft))) + tensors, _info = pi.remote_export_lora(worker) + specs = tuple(CanonicalTensorSpec(n, tuple(t.shape), "float32", t.numel()) for n, t in sorted(tensors.items())) + kwargs = {**IDS, "layout_hash": canonical_layout_hash(specs), "expected_specs": specs} + def untimed(raw): # timings differ from call to call + return {**{k: v for k, v in raw.items() if k != "export_seconds"}, + "digest": {k: v for k, v in raw["digest"].items() if k != "hash_seconds"}} + + want = untimed(digest_lora_tensors(tensors, **kwargs)) + for _ in range(2): # the second round reuses the worker's buffer + assert untimed(pi.remote_export_digest(worker, kwargs)) == want + assert getattr(worker, "_yeto_digest_staging", None) is not None diff --git a/tests/test_run_flags.py b/tests/test_run_flags.py index 9ef5c8ea..9cdba9d4 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", "--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/tests/test_tools_eval_efficiency.py b/tests/test_tools_eval_efficiency.py new file mode 100644 index 00000000..c4eeb1f2 --- /dev/null +++ b/tests/test_tools_eval_efficiency.py @@ -0,0 +1,537 @@ +"""CPU tests of the efficiency evaluation (tools/eval/efficiency, S26 E0.1). + +The synthetic logs below follow the formats the parsers cite: + +* verl FileLogger rows ``{"step": N, "data": {...}}`` -- fork acad9875 + ``verl/utils/tracking.py`` ``FileLogger.log``; console ``step:N - k:v`` -- + ``verl/utils/logger/aggregate_logger.py`` ``concat_dict_to_str``; keys as in the + real T-EVAL file ``SESSION25/logs/T-EVAL/eval-e1-verl-c/.../verl-file-logger-0.jsonl`` + and ``metric_utils.compute_timing_metrics`` (``timing_s/*``); +* Miles ``perf N: {dict}`` lines -- copied in shape from the real T-EVAL log + ``SESSION25/logs/T-EVAL/eval-e1-miles-a/raw/launch.ts.log`` (rollout side + ``metrics.py:93``, train side ``train_metric_utils.py:56``); +* yeto tape events -- field names as in the T-EVAL tapes (``rl_round_wall``, + ``rl_local_round``, ``rl_round_trained``, ``rl_timeline_span``); +* NVML CSV -- ``nvidia-smi --query-gpu=... --format=csv`` output and the + ``miles_local.sh`` ``,`` prefixed form. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from tools.eval import __main__ as cli +from tools.eval.efficiency import eff_reward, metrics, parse, report, spec + +SPEC = spec.get_spec("qwen3-0.6b") + + +# ====================================================================== spec +def test_arms_cover_both_systems_frameworks_and_modes(): + assert len(spec.ARMS) == 8 + for name in spec.ARMS: + arm = spec.Arm.parse(name) + assert spec.Arm.parse(arm.peer).system != arm.system + with pytest.raises(ValueError): + spec.Arm.parse("bare-verl-fast") + + +def test_every_arm_carries_the_shared_values(): + plans = {a: spec.plan(SPEC, a, "r1") for a in spec.ARMS} + assert len({p["fingerprint"] for p in plans.values()}) == 1 + for name, p in plans.items(): + text = p.get("script") or " ".join(p.get("argv") or []) or p["launcher_args"] + assert p["cards"]["total"] == (1 if p["mode"] == "sync" else 2), name + if p["system"] == "bare": + assert "Qwen3-0.6B" in text and "/root/eff-data/train" in text, name + else: + assert SPEC.model_revision in text and SPEC.data_revision in text, name + for value in (SPEC.rounds, SPEC.groups_per_round, SPEC.samples_per_group, SPEC.max_response_len, + SPEC.lr, SPEC.lora_r, SPEC.seed): + assert str(value) in text, (name, value) + + +def test_fingerprint_changes_with_a_shared_value_only(): + base = spec.control_fingerprint(SPEC) + assert spec.control_fingerprint(spec.get_spec("qwen3-0.6b", repeats=5)) == base + assert spec.control_fingerprint(spec.get_spec("qwen3-0.6b", lr=2e-5)) != base + + +def test_bare_entry_points_follow_the_official_recipes(): + sync = spec.bare_script(SPEC, spec.Arm.parse("bare-verl-sync"), "r1") + async_ = spec.bare_script(SPEC, spec.Arm.parse("bare-verl-async"), "r1") + assert "-m verl.trainer.main_ppo" in sync and "fully_async_main" in async_ + assert "async_training.staleness_threshold=1" in async_ and "rollout.n_gpus_per_node=1" in async_ + assert "trainer.logger=[console,file]" in sync and "VERL_FILE_LOGGER_PATH" in sync + miles_sync = spec.bare_script(SPEC, spec.Arm.parse("bare-miles-sync"), "r1") + miles_async = spec.bare_script(SPEC, spec.Arm.parse("bare-miles-async"), "r1") + assert "/root/miles/train.py $MODEL_ARGS" in miles_sync and "--colocate" in miles_sync + assert "/root/miles/train_async.py" in miles_async and "--fully-async --max-weight-staleness 1" in miles_async + assert "scripts/models/qwen3-0.6B.py" in miles_sync # the image's own model flags + for script in (sync, async_, miles_sync, miles_async): + assert "yeto" not in script.replace("yeto-effeval", "").replace("yeto/", "") + + +def test_yeto_arms_use_the_same_policy_age_and_partition(): + assert "--rl-max-policy-age 1" in " ".join(spec.yeto_verl_argv(SPEC, "async", "/t")) + assert "--rl-max-policy-age" not in " ".join(spec.yeto_verl_argv(SPEC, "sync", "/t")) + # async needs an algorithm spec that tolerates the policy age (the default allows 0) + argv = spec.yeto_verl_argv(SPEC, "async", "/t") + algo = argv[argv.index("--rl-algorithm-spec") + 1] + assert algo.endswith(spec.ALGORITHM_SPEC) and (Path(__file__).parents[1] / spec.ALGORITHM_SPEC).is_file() + args = spec.yeto_miles_launcher_args(SPEC, "async", "r1") + assert "--gpu modal:2xa10g" in args and "--rl-placement fixed-partition --rl-rollout-gpus 1" in args + assert "--gpu modal:1xa10g" in spec.yeto_miles_launcher_args(SPEC, "sync", "r1") + # E1: on the bare arms' GPU model (an A10G container is refused) + assert "--modal-gpu-name 'NVIDIA A10' " in args + from tools.eval import modal_verl + + assert spec.GPU_MODELS == modal_verl.GPU_NAMES + # only the final round cut, as the bare arms keep no periodic checkpoint (E1) + assert f"--rl-cut-every {SPEC.rounds} " in spec.yeto_miles_launcher_args(SPEC, "sync", "r1") + every = spec.get_spec(SPEC.name, checkpoint="every-round") + assert "--rl-cut-every 1 " in spec.yeto_miles_launcher_args(every, "async", "r1") + assert spec.control_fingerprint(every) == spec.control_fingerprint(SPEC) # bare arms unchanged + with pytest.raises(ValueError): + spec.cut_every(spec.get_spec(SPEC.name, checkpoint="never")) + + +# overrides yeto's verl translation sets and the bare recipe leaves at the default (or the reverse); +# a new key here means the two arms drifted apart: decide, then extend the list +KNOWN_VERL_DIFF = { + "actor_rollout_ref.model.lora.merge", "actor_rollout_ref.rollout.enforce_eager", + "actor_rollout_ref.rollout.full_determinism", "actor_rollout_ref.rollout.load_format", + "ray_kwargs.ray_init.num_cpus", "trainer.project_name", "trainer.use_v1", "trainer.v1.trainer_mode", + "actor_rollout_ref.rollout.mode", "async_training.staleness_threshold", "data.return_raw_chat", + "rollout.total_rollout_steps", +} + + +@pytest.mark.parametrize("mode", spec.MODES) +def test_verl_knob_diff_is_the_known_set(mode): + diff = spec.verl_knob_diff(SPEC, mode) + assert set(diff) <= KNOWN_VERL_DIFF, set(diff) - KNOWN_VERL_DIFF + + +def test_4b_preset_is_ready_for_e02(): + s4 = spec.get_spec("qwen3-4b-instruct-2507") + assert s4.gpu == "H100!" and s4.repeats >= 2 and len(s4.model_revision) == 40 + assert "--gpu modal:2xh100" in spec.yeto_miles_launcher_args(s4, "async", "r") + assert "qwen3-4B-Instruct-2507.py" in spec.bare_script(s4, spec.Arm.parse("bare-miles-sync"), "r") + + +# ====================================================================== reward +@pytest.mark.parametrize("response,label", [ + ("so \\boxed{42}", "#### 42"), ("the answer is 1,234", "x #### 1234"), ("\\boxed{3.5} then 7", "#### 3.5"), + ("no number", "#### 1"), ("\\boxed{-2}", "#### 2"), ("", None), +]) +def test_bare_reward_equals_yeto_reward(response, label): + from yeto.rl.rewards.builtin import gsm8k_grade + + assert eff_reward.grade(response, label) == gsm8k_grade(response, label) + assert eff_reward.compute_score("d", response, label) == gsm8k_grade(response, label) + + +# ====================================================================== parse: verl +def _verl_rows(n: int, *, async_: bool = False) -> str: + rows = [] + for step in range(1, n + 1): + d = {"timing_s/gen": 10.0 + step, "timing_s/old_log_prob": 1.0, "timing_s/update_actor": 3.0, + "timing_s/step": 16.0 + step, "response_length/mean": 100.0, "perf/total_num_tokens": 9000, + "critic/score/mean": 0.1 * step, "rollout_corr/rollout_is_ratio_fraction_high": 0.01} + if async_: + d.pop("timing_s/gen") + d.update({"timing_s/param_sync": 0.5, "timing_s/wait_for_enough_samples": 2.0, + "fully_async/count/dropped_stale_samples": step, "fully_async/count/total_generated_samples": + 32 * step}) + else: + d["timing_s/update_weights"] = 0.7 + rows.append(json.dumps({"step": step, "data": d})) + return "\n".join(rows) + "\n" + + +def _verl_console(n: int, t0: float = 1000.0, dt: float = 20.0) -> str: + return "".join(f"{t0 + dt * s:.3f} (TaskRunner pid=1) step:{s} - timing_s/gen:{10 + s} - critic/score/mean:0.1\n" + for s in range(1, n + 1)) + + +def test_parse_verl_sync_rounds_and_sources(): + p = parse.parse_verl(_verl_rows(3), _verl_console(3), samples_per_round=32, trainer_cards=1, mode="sync") + r = p.rounds + assert [x["round"] for x in r] == [1, 2, 3] + assert r[1]["round_s"] == 20.0 and r[0]["round_s"] is None + assert r[0]["gen_tokens"] == 3200 and r[0]["train_s"] == 4.0 and r[0]["weight_sync_s"] == 0.7 + assert p.sources["gen_tokens"].startswith("derived:") and p.sources["wait_s"] == parse.NA + assert p.sources["sampler_update_s"] == parse.NA and p.sources["version_lag"] == parse.NA + assert p.segments_raw[0]["timing_s/step"] == 17.0 + + +def test_parse_verl_async_reads_fully_async_counters(): + p = parse.parse_verl(_verl_rows(2, async_=True), "", samples_per_round=32, trainer_cards=1, mode="async") + assert p.rounds[0]["wait_s"] == 2.0 and p.rounds[0]["weight_sync_s"] == 0.5 + assert p.rounds[1]["dropped_samples"] == 2 and p.sources["gen_s"] == parse.MISSING + assert p.sources["t_end_unix"] == parse.MISSING and p.sources["version_lag"] == parse.MISSING + + +# ====================================================================== parse: Miles +def _miles_log(n: int, *, async_: bool = False) -> str: + out = [] + for rid in range(n): + ro = {"rollout/raw_reward": 0.4 + 0.01 * rid, "perf/rollout_time": 12.0, + "perf/tokens_per_gpu_per_sec": 500.0, "rollout/num_training_samples": 32} + tr = {"perf/train_wait_time": 12.5, "perf/train_time": 20.0, "perf/actor_train_time": 15.0, + "perf/actor_train_tok_per_s": 800.0, "perf/update_weights_time": 0.15, "perf/sleep_time": 1.7} + out.append(f"2026-10-20T06:36:{10 + rid:02d}Z [p-l0-modal] \x1b[36m(RolloutExecutor pid=2765)\x1b[0m " + f"[2026-10-20 06:36:{10 + rid:02d}.055 rollout_executor] metrics.py:93 - perf {rid}: {ro!r}") + out.append(f"2026-10-20T06:37:{rid * 30 % 60:02d}Z [p-l0-modal] (MegatronTrainRayActor pid=3190) " + f"[2026-10-20 06:{37 + rid // 2}:{rid * 30 % 60:02d}.147 actor_cell00000_rank00000] " + f"train_metric_utils.py:56 - perf {rid}: {tr!r}") + return "\n".join(out) + "\nunrelated line\n" + + +def test_parse_miles_merges_rollout_and_train_lines(): + p = parse.parse_miles(_miles_log(3), rollout_cards=1, mode="sync") + r = p.rounds + assert [x["round"] for x in r] == [0, 1, 2] + assert r[0]["gen_s"] == 12.0 and r[0]["train_s"] == 20.0 and r[0]["weight_sync_s"] == 0.15 + assert r[0]["gen_tokens"] == 6000.0 and r[0]["train_tokens"] == 12000.0 and r[0]["samples"] == 32 + assert r[1]["round_s"] == 30.0 and r[0]["reward_mean"] == 0.4 + assert p.sources["wait_s"] == parse.NA and p.sources["is_clip_frac"] == parse.MISSING + + +def test_parse_miles_async_keeps_wait(): + p = parse.parse_miles(_miles_log(2), rollout_cards=1, mode="async") + assert p.rounds[0]["wait_s"] == 12.5 and p.sources["dropped_samples"] == parse.MISSING + + +# ====================================================================== parse: yeto tape +def _tape(n: int, t0: float = 2000.0) -> list[dict]: + ev = [] + for rid in range(n): + end = t0 + 25.0 * (rid + 1) + ev += [{"event": "rl_timeline_span", "rollout_id": rid, "task": "generate", "kind": "compute", + "start": 100.0 * rid, "end": 100.0 * rid + 11.0}, + {"event": "rl_timeline_span", "rollout_id": rid, "task": "train", "kind": "compute", + "start": 100.0 * rid + 11.0, "end": 100.0 * rid + 21.0}, + {"event": "rl_timeline_span", "rollout_id": rid, "task": "publish", "kind": "transfer", + "start": 100.0 * rid + 21.0, "end": 100.0 * rid + 22.0}, + {"event": "rl_timeline_span", "rollout_id": rid, "task": "offload", "kind": "transfer", + "start": 100.0 * rid + 22.0, "end": 100.0 * rid + 24.0}, + {"event": "rl_round_wall", "rollout_id": rid, "round_seconds": 25.0, "time_unix": end}, + {"event": "rl_local_round", "action_tokens": 3000 + rid, "reward_mean": 0.5, "time_unix": end - 1}, + {"event": "rl_round_trained", "rollout_id": rid, "trained_samples": 32, + "mismatch": {"tis_clipfrac": 0.0}, "time_unix": end}] + return ev + + +def test_parse_yeto_tape_with_framework_tokens(): + fw = parse.parse_miles(_miles_log(3), rollout_cards=1, mode="sync") + p = parse.parse_yeto(_tape(3), mode="sync", framework_parsed=fw) + r = p.rounds + assert r[0]["gen_s"] == 11.0 and r[0]["train_s"] == 10.0 and r[0]["weight_sync_s"] == 1.0 + assert r[1]["round_s"] == 25.0 and r[2]["gen_tokens"] == 3002 and r[0]["train_tokens"] == 12000.0 + assert p.sources["train_tokens"].startswith("framework log:") and p.segments_raw[0]["offload"] == 2.0 + + +def test_parse_yeto_driver_phase_fallback(): + ev = [{"event": "rl_driver_phase", "phase": "generate", "rollout_id": 0, "time_unix": 10.0}, + {"event": "rl_driver_phase", "phase": "train", "time_unix": 18.0}, + {"event": "rl_driver_phase", "phase": "publish", "time_unix": 21.0}, + {"event": "rl_driver_phase", "phase": "generate", "rollout_id": 1, "time_unix": 22.0}, + {"event": "rl_round_wall", "rollout_id": 0, "time_unix": 22.0}] + p = parse.parse_yeto(ev, mode="sync") + assert p.rounds[0]["gen_s"] == 8.0 and p.rounds[0]["train_s"] == 3.0 and p.rounds[0]["weight_sync_s"] == 1.0 + assert "rl_driver_phase" in p.sources["gen_s"] + + +def test_parse_yeto_verl_async_generation_from_rollouter_timing(): + ev = [{"event": "rl_publication_check", "policy_version": 0, + "rollouter_timing": {"fully_async/rollouter/active_time": 2.75}}, + {"event": "rl_publication_check", "policy_version": 1, + "rollouter_timing": {"fully_async/rollouter/active_time": 10.8}}, + {"event": "rl_round_wall", "rollout_id": 0, "time_unix": 21.3}, + {"event": "rl_publication_check", "policy_version": 2, + "rollouter_timing": {"fully_async/rollouter/active_time": 1.86}}, + {"event": "rl_round_wall", "rollout_id": 1, "time_unix": 26.5}] + p = parse.parse_yeto(ev, mode="async", skip=frozenset({"gen_s"}), rollouter_gen=True) + assert [r["gen_s"] for r in p.rounds] == [10.8, 1.86] + assert p.sources["gen_s"].startswith("rl_publication_check.rollouter_timing") + off = parse.parse_yeto(ev, mode="async", skip=frozenset({"gen_s"})) + assert off.rounds[0]["gen_s"] is None + + +def test_parse_nvml_both_formats(): + own = ("timestamp, index, name, utilization.gpu [%], memory.used [MiB], memory.total [MiB], power.draw [W]\n" + "2026/10/20 06:36:11.055, 0, NVIDIA A10, 45 %, 6077 MiB, 23028 MiB, 81.7 W\n" + "2026/10/20 06:36:12.055, 1, NVIDIA A10, 0 %, 100 MiB, 23028 MiB, 30.1 W\n") + polled = "ta-01ABC,2026/10/20 06:36:11.055, NVIDIA A10, 44 %, 6077 MiB, 81.76 W\n" + a, b = parse.parse_nvml(own), parse.parse_nvml(polled) + assert [x["util"] for x in a] == [45.0, 0.0] and a[1]["key"] == ":1" and a[0]["mem_total_mb"] == 23028.0 + assert b[0]["key"] == "ta-01ABC:0" and b[0]["mem_mb"] == 6077.0 and b[0]["power_w"] == 81.76 + + +# ====================================================================== run dirs, metrics, report +def _nvml_csv(t0: float, n: int, util: int) -> str: + import datetime as dt + + lines = ["timestamp, index, name, utilization.gpu [%], memory.used [MiB], memory.total [MiB], power.draw [W]"] + for i in range(n): + stamp = dt.datetime.fromtimestamp(t0 + i, dt.timezone.utc).strftime("%Y/%m/%d %H:%M:%S.%f")[:-3] + lines.append(f"{stamp}, 0, NVIDIA A10, {util if i % 4 else 0} %, 7000 MiB, 23028 MiB, 80 W") + return "\n".join(lines) + "\n" + + +def _run_dir(tmp: Path, arm: str, run_id: str, dt_round: float) -> Path: + plan = spec.plan(SPEC, arm, run_id) + d = tmp / run_id + spec.write_plan(d / "arm.json", plan) + raw = d / "raw" / run_id / "arm" + raw.mkdir(parents=True) + n, t0 = SPEC.rounds, 1000.0 + (raw / "meta.json").write_text(json.dumps({"rc": 0, "start_unix": t0 - 200, "end_unix": t0 + dt_round * n + 30})) + (raw / "nvml.csv").write_text(_nvml_csv(t0, int(dt_round * n) + 5, 60)) + if plan["system"] == "bare": + (raw / "main.log").write_text(_verl_console(n, t0, dt_round)) + (raw / "eff-out").mkdir() + (raw / "eff-out" / "verl-metrics.jsonl").write_text(_verl_rows(n)) + else: + out = raw / "yeto-output" + out.mkdir() + (out / "verl-file-logger-0.jsonl").write_text("".join( # steps 1..n = rollout ids 0..n-1 + json.dumps({"step": k, "data": {"global_seqlen/mean": 9000}}) + "\n" for k in range(1, n + 1))) + (out / "rl-island-0.jsonl").write_text("\n".join(json.dumps(e) for e in _tape(n, t0 - dt_round) + ) + "\n") + return d + + +def test_parse_run_metrics_and_report(tmp_path, monkeypatch): + dirs = [_run_dir(tmp_path, "bare-verl-sync", "b1", 20.0), _run_dir(tmp_path, "bare-verl-sync", "b2", 21.0), + _run_dir(tmp_path, "yeto-verl-sync", "y1", 25.0), _run_dir(tmp_path, "yeto-verl-sync", "y2", 25.0)] + for d in dirs: + assert cli.main(["eff-parse", str(d)]) == 0 + assert (d / "parsed.json").is_file() and (d / "summary.json").is_file() + s = json.loads((dirs[0] / "summary.json").read_text()) + assert s["rounds"] == SPEC.rounds and s["steady_wall_s"] == 20.0 * (SPEC.rounds - 1) + assert s["gen_tok_per_card_s"] == round(3200 / 20.0, 3) and s["bubble_frac"] == pytest.approx(0.25, abs=0.02) + assert s["usd_per_1k_samples_steady"] == round(2.0 * 20.0 * 7 / 3600 / (32 * 7) * 1000, 4) + assert s["gpu_mem_peak_mb"] == 7000.0 + assert s["missing"] == [] and "wait_s_median" in s["not_applicable"] and "drop_frac" in s["not_applicable"] + y = json.loads((dirs[2] / "summary.json").read_text()) + assert y["missing"] == [] and y["train_s_median"] == 10.0 and y["unaccounted_s_median"] == 3.0 + assert y["train_tok_per_card_s"] == round(9000 / 25.0, 3) + out = tmp_path / "report.md" + assert cli.main(["eff-report", *map(str, dirs), "--out", str(out)]) == 0 + text = out.read_text() + assert "### verl sync (bare n=2, yeto n=2)" in text + section = text.split("### verl sync")[1] + row = next(line for line in section.splitlines() if line.startswith("| round time (s)")) + assert "yeto worse" in row and "+4.5" in row + assert "| update sampler weights (s/round) | n/a | n/a | n/a | not applicable |" in text + assert "wait (s/round) | n/a" in text + + +def test_report_needs_two_repeats_and_one_spec(tmp_path): + one = [_run_dir(tmp_path, "bare-verl-sync", "b1", 20.0), _run_dir(tmp_path, "yeto-verl-sync", "y1", 25.0)] + for d in one: + cli.main(["eff-parse", str(d)]) + text = report.markdown(report.load(one), title="t") + assert "insufficient repeats (bare 1, yeto 1; need 2)" in text + runs = report.load(one) + runs[1]["fingerprint"] = "other" + with pytest.raises(ValueError, match="different specs"): + report.markdown(runs, title="t") + runs = report.load(one) + assert runs[0]["gpu_names"] # from the NVML samples + runs[1]["gpu_names"] = ["NVIDIA A10G"] + with pytest.raises(ValueError, match="different GPU models"): + report.markdown(runs, title="t") + + +def test_gpu_name_check_tells_a10_from_a10g(): + from tools.eval.modal_verl import gpu_name_matches + + assert gpu_name_matches("NVIDIA A10", "NVIDIA A10") + assert not gpu_name_matches("NVIDIA A10", "NVIDIA A10G") + assert gpu_name_matches("H100", "NVIDIA H100 80GB HBM3") and gpu_name_matches("L4", "NVIDIA L4") + + +def test_verdict_rules(): + assert report.verdict([10, 11], [12, 13], True) == "yeto worse" + assert report.verdict([10, 11], [8, 9], True) == "yeto better" + assert report.verdict([10, 12], [11, 13], True) == "within run-to-run range" + assert report.verdict([10, 11], [8, 9], False) == "yeto worse" + assert report.verdict([10, None], [8, 9], True) == "missing value" + + +def test_price_comes_from_the_config(tmp_path): + path = tmp_path / "p.json" + path.write_text(json.dumps({"A10G": {"usd_per_card_hour": 3.5}})) + assert metrics.price("A10G", path) == 3.5 + with pytest.raises(KeyError): + metrics.price("H100!", path) + + +def test_cli_plan_and_unregistered_run(tmp_path, monkeypatch, capsys): + assert cli.main(["eff-plan", "qwen3-0.6b", "--arm", "bare-verl-sync"]) == 0 + assert json.loads(capsys.readouterr().out)["arms"]["bare-verl-sync"]["cards"]["total"] == 1 + monkeypatch.setattr(cli, "LEDGER", tmp_path / "ledger.md") + (tmp_path / "ledger.md").write_text("nothing registered\n") + with pytest.raises(SystemExit, match="not pre-registered"): + cli.main(["eff-run", "qwen3-0.6b", "--arm", "bare-verl-sync", "--run-id", "eff-x", "--out", str(tmp_path)]) + + +def test_bare_verl_async_uses_the_fork_working_lora_sync(): + script = spec.bare_script(SPEC, spec.Arm.parse("bare-verl-async"), "r") + assert "actor_rollout_ref.model.lora.merge=True" in script and "export NCCL_CUMEM_HOST_ENABLE=0" in script + assert "NCCL_CUMEM_HOST_ENABLE" not in spec.bare_script(SPEC, spec.Arm.parse("bare-verl-sync"), "r") + + +def test_parse_bare_verl_async_rows_lag_one_version(): + # fork acad9875 fully_async: weight-sync row at version N, training row of step N at N+1, + # the last step's training row never written + rows = [{"step": 0, "data": {"fully_async/rollouter/active_time": 9.0}}] + for v in range(1, 4): + rows.append({"step": v, "data": {"fully_async/rollouter/active_time": 1.0}}) + if v > 1: + rows.append({"step": v, "data": {"training/global_step": v - 1, "response_length/mean": 10.0, + "critic/score/mean": 0.5}}) + log = "\n".join(f"{100 + v}.0 (FullyAsyncTrainer) step:{v} - x:1" for v in range(4)) + p = parse.parse_verl("\n".join(json.dumps(r) for r in rows), log, samples_per_round=4, trainer_cards=1, + mode="async") + assert [r["round"] for r in p.rounds] == [1, 2, 3] + assert [r["t_end_unix"] for r in p.rounds] == [101.0, 102.0, 103.0] + assert [r["reward_mean"] for r in p.rounds] == [0.5, 0.5, None] + assert [r["samples"] for r in p.rounds] == [4, 4, 4] + assert [r["gen_s"] for r in p.rounds] == [1.0, 1.0, 1.0] # rollouter active time, not the queue pull + + +def test_bare_miles_reads_prompts_in_file_order_like_yeto(): + for arm in ("bare-miles-sync", "bare-miles-async"): + assert "--rollout-shuffle" not in spec.bare_script(SPEC, spec.Arm.parse(arm), "r") + + +def test_tape_events_skip_the_repo_copy_and_the_store(tmp_path): + row = {"event": "rl_local_round", "local_round_id": 1, "time_unix": 1.0} + for sub in ("tape-direct/l0", "yeto/tests/fixtures", "store/x"): + (tmp_path / sub).mkdir(parents=True) + (tmp_path / sub / "rl-island-0.jsonl").write_text(json.dumps({**row, "src": sub}) + "\n") + assert [r.get("src") for r in parse._tape_events(tmp_path)] == ["tape-direct/l0"] + + +def test_cards_whole_cluster_split(): + from tools.eval.efficiency import spec as S + + assert S.cards("sync") == {"total": 1, "rollout": 1, "trainer": 1, "colocated": True} + assert S.cards("async") == {"total": 2, "rollout": 1, "trainer": 1, "colocated": False} + assert S.cards("async", 4) == {"total": 4, "rollout": 2, "trainer": 2, "colocated": False} + assert S.cards("sync", 4)["trainer"] == 4 + with pytest.raises(ValueError): + S.cards("async", 3) + base = S.get_spec("qwen3-4b-instruct-2507") + # cards=0 keeps the fingerprint of runs made before the field existed + assert S.control_fingerprint(base) != S.control_fingerprint(S.get_spec("qwen3-4b-instruct-2507", cards=4)) + plan = S.plan(S.get_spec("qwen3-4b-instruct-2507", cards=4), "bare-miles-async", "e11-x") + assert "--rollout-num-gpus 2" in plan["script"] and "--actor-num-gpus-per-node 2" in plan["script"] + + +def test_jsonl_source_rows_and_yeto_task_file(tmp_path): + from tools.eval.efficiency import eff_prep + + src = tmp_path / "d.jsonl" + src.write_text("".join(json.dumps({"prompt": [{"role": "user", "content": f"q{i}"}], "label": str(i)}) + "\n" + for i in range(10))) + sp = spec.get_spec("e11-qwen3-4b", train_rows=4).to_dict() + rows = eff_prep.source_rows(sp, lambda *a, **k: str(src)) + assert [lab for _, lab in rows["train"]] == ["0", "1", "2", "3"] + assert [lab for _, lab in rows["test"]] == [str(i) for i in range(4, 10)] + tasks = eff_prep.yeto_task_rows(sp, lambda *a, **k: str(src)) + assert tasks[2] == {"messages": [{"role": "user", "content": "q2"}], "label": "2", "metadata": {"task_id": "gsm3"}} + with pytest.raises(ValueError): + eff_prep.source_rows({**sp, "train_rows": 0}, lambda *a, **k: str(src)) + + +def test_e11_arms_read_one_row_set(): + e11 = spec.get_spec("e11-qwen3-4b") + assert e11.train_rows == e11.rounds * e11.groups_per_round + single = spec.yeto_miles_launcher_args(e11, "async", "e11-x") + assert f"--rl-allow-local-data --data {spec.YETO_TASK_FILE}" in single and "--gpu modal:4xh100" in single + multi = spec.yeto_miles_islands_launcher_args(e11, 2, "e11-y", data="t.jsonl", island_tasks="t.json") + assert "--gpu modal:2xh100,modal:2xh100" in multi and "--rollout-batch-size 16" in multi + assert "--rl-rollout-gpus 1" in multi and "--controller head" in multi and "--rl-island-scheduling elastic" in multi + split = spec.island_task_split(e11.groups_per_round, e11.rounds, 2) + ids = sorted(int(t[3:]) for v in split.values() for t in v) + assert ids == list(range(1, e11.train_rows + 1)) + assert split["l1"][:2] == ["gsm17", "gsm18"] + with pytest.raises(ValueError): + spec.yeto_verl_argv(e11, "async", "/t") + # the E0 presets keep their fingerprints (new fields at their defaults are left out) + assert spec.control_fingerprint(spec.get_spec("qwen3-0.6b")).startswith("59ed2ab081afb9f9") + + +def test_cluster_target_and_smoothing(): + from tools.eval.efficiency import cluster as C + + pts = [C.Point(100.0 + 10 * i, 128, r) for i, r in enumerate([0.1, 0.1, 0.3, 0.5, 0.5, 0.5])] + c = C.Curve("r", "yeto-miles-2islands", 4, 0.0, 200.0, pts) + sm = C.smoothed(pts, 256) + assert sm[0] == 0.1 and sm[2] == pytest.approx(0.2) and sm[4] == pytest.approx(0.5) + assert C.start_reward([c], 256) == pytest.approx(0.1) + m = C.metrics(c, target=0.4, window=256, usd_per_card_hour=5.0) + assert m["reached"] and m["to_target_s"] == 130.0 and m["to_target_samples"] == 4 * 128 + assert m["to_target_gpu_h"] == pytest.approx(130 / 3600 * 4) + assert not C.metrics(c, target=0.9, window=256, usd_per_card_hour=5.0)["reached"] + + +# ---------------------------------------------------------------- E11 full parameter +def test_e11_full_preset_runs_every_parameter_with_trainer_tp2(): + sp = spec.get_spec("e11-qwen3-0.6b-full") + argv = spec.miles_bare_argv(sp, "async") + assert not any(a.startswith(("--lora", "--sglang-max-lora", "--target-modules")) for a in argv) + assert argv[argv.index("--tensor-model-parallel-size") + 1] == "2" + assert argv[argv.index("--rollout-num-gpus") + 1] == "2" + args = spec.yeto_miles_launcher_args(sp, "async", "r") + assert "--tuning full --tensor-parallel 2 " in args and "--lora-r" not in args + assert sp.samples_per_round == 256 and sp.rounds * sp.groups_per_round == sp.train_rows + + +def test_tuning_fields_keep_earlier_fingerprints(): + sp = spec.get_spec("e11-qwen3-4b") + assert spec.control_fingerprint(sp) == spec.control_fingerprint( + spec.replace(sp, tuning="lora", tensor_parallel=1)) + assert spec.control_fingerprint(sp) != spec.control_fingerprint(spec.replace(sp, tuning="full")) + + +def test_full_islands_args_use_the_e3_c_sync_and_split_each_round(): + sp = spec.get_spec("e11-qwen3-0.6b-full") + args = spec.yeto_miles_full_islands_launcher_args(sp, 2, "r", data="t.jsonl", island_tasks="s.json") + for flag in ("--rl-sync-preset full-parameter", "--fragments 2", "--rl-sync-fragment-pipeline 2", + "--rl-sync-async-checkpoint", "--rl-sync-push-dtype bf16", "--rl-sync-checkpoint-every 8", + "--wan-streams 8", "--rollout-batch-size 16", "--gpu modal:2xh100,modal:2xh100", + "--rl-max-policy-age 0", "--total-steps 12"): + assert flag in args + with pytest.raises(ValueError): + spec.yeto_miles_full_islands_launcher_args(spec.get_spec("e11-qwen3-4b"), 2, "r", data="t", island_tasks="s") + + +def test_run_logged_kills_a_run_without_training_steps(tmp_path, monkeypatch): + from tools.eval.efficiency import modal_eff + + monkeypatch.setattr(modal_eff, "WATCH_POLL_S", 0.2) + monkeypatch.setattr(modal_eff, "FIRST_STEP_S", 1) + monkeypatch.setattr(modal_eff, "STALL_S", 1) + log = tmp_path / "main.log" + script = "echo \"step 1: {'train/step': 0}\"; sleep 0.5; echo \"{'train/step': 1}\"; sleep 30" + rc = modal_eff._run_logged(["bash", "-c", script], log, 60, watch_progress=True) + assert rc != 0 + assert "killed: no training step" in log.read_text() + + +def test_run_logged_without_the_watch_waits_for_the_process(tmp_path): + from tools.eval.efficiency import modal_eff + + log = tmp_path / "main.log" + assert modal_eff._run_logged(["bash", "-c", "sleep 1; echo done"], log, 60) == 0 + assert "killed" not in log.read_text() diff --git a/tests/test_wan_probe.py b/tests/test_wan_probe.py new file mode 100644 index 00000000..ad8a19df --- /dev/null +++ b/tests/test_wan_probe.py @@ -0,0 +1,150 @@ +"""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 == [] + + +@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/tests/test_wan_streams.py b/tests/test_wan_streams.py new file mode 100644 index 00000000..99f3d4dc --- /dev/null +++ b/tests/test_wan_streams.py @@ -0,0 +1,261 @@ +"""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" + + +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/tests/test_wandb_tape.py b/tests/test_wandb_tape.py index 0a9abf08..2305d7fc 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 00000000..a55fb906 --- /dev/null +++ b/tools/bench_sync_boundary.py @@ -0,0 +1,175 @@ +"""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), 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 +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, DTYPE_F32_PUSH_BF16, 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, checkpoint_every: int) -> 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", str(checkpoint_every), "--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], dtype: int) -> None: + fragments = len(layout.fragments) + 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) + 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() + + +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, 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, dtype)) + 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": profile, + "boundary_s_per_sweep": [round(x, 3) for x in steady], + "boundary_s_median": round(statistics.median(steady), 3), + "checkpoints_written": sum("checkpoint_ms" in r 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, profile=profile, learners=args.learners, + fragments=args.fragments, numel=args.floats_per_fragment, sweeps=args.sweeps) + 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}, + indent=2)) + + +if __name__ == "__main__": + main() diff --git a/tools/eval/__main__.py b/tools/eval/__main__.py index 1d047988..cbb998b6 100644 --- a/tools/eval/__main__.py +++ b/tools/eval/__main__.py @@ -5,6 +5,13 @@ python -m tools.eval judge DIR/eval-e3-verl-a python -m tools.eval report DIR/* --out report.md [--html report.html] +Efficiency evaluation (S26 E0.1, package :mod:`.efficiency`): + + python -m tools.eval eff-plan qwen3-0.6b [--arm ARM] + python -m tools.eval eff-run qwen3-0.6b --arm bare-verl-sync --run-id eff-... [--out DIR] [--dry-run] + python -m tools.eval eff-parse DIR/eff-... + python -m tools.eval eff-report DIR/eff-* --out report.md + ``run`` refuses a run id that is not pre-registered in the GPU ledger (the id must appear in ``gpu-spend.md``), checks the local thread count and Ray, writes ``cell.json``, runs the cell's harness (raw data into ``raw/``), then judges. @@ -123,6 +130,71 @@ def cmd_report(args) -> int: return 0 +EFF_OUT = Path("/home/michael/work/infra-drafts/SESSION26/logs/E0") + + +def cmd_eff_plan(args) -> int: + from .efficiency import spec + + sp = spec.get_spec(args.preset) + arms = [args.arm] if args.arm else list(spec.ARMS) + doc = {"spec": sp.to_dict(), "fingerprint": spec.control_fingerprint(sp), + "arms": {a: spec.plan(sp, a, "") for a in arms}, + "verl_knob_diff": {m: spec.verl_knob_diff(sp, m) for m in spec.MODES}} + print(json.dumps(doc, indent=1, default=str)) + return 0 + + +def cmd_eff_run(args) -> int: + from .efficiency import modal_eff, spec + + sp = spec.get_spec(args.preset, checkpoint=args.checkpoint) + plan = spec.plan(sp, args.arm, args.run_id) + run_dir = Path(args.out) / args.run_id + if not args.dry_run: + if (run_dir / "arm.json").exists(): + raise SystemExit(f"{run_dir} already used") + _preflight(args.run_id) + plan["code_sha"] = subprocess.run(["git", "-C", str(Path(__file__).resolve().parents[2]), "rev-parse", "HEAD"], + capture_output=True, text=True).stdout.strip() + plan["images"] = spec.images() + spec.write_plan(run_dir / "arm.json", plan) + if plan["harness"] == "miles_local": + summary = modal_eff.run_miles_launcher(plan, run_dir, dry_run=args.dry_run) + if args.dry_run: + print((run_dir / "dry-run.log").read_text()) + return summary["rc"] + elif args.dry_run: + print(plan.get("script") or json.dumps(plan.get("argv"))) + return 0 + else: + summary = modal_eff.run(plan, run_dir) + (run_dir / "run.json").write_text(json.dumps(summary, indent=1, default=str)) + return cmd_eff_parse(argparse.Namespace(run_dir=str(run_dir))) + + +def cmd_eff_parse(args) -> int: + from .efficiency import metrics, parse + + run_dir = Path(args.run_dir) + parsed = parse.parse_run(run_dir) + (run_dir / "parsed.json").write_text(json.dumps(parsed, indent=1, default=str)) + summary = metrics.summarize(parsed) + (run_dir / "summary.json").write_text(json.dumps(summary, indent=1, default=str)) + print(json.dumps({"run": parsed["run_id"], "arm": parsed["arm"], "rounds": summary["rounds"], + "missing": summary["missing"]})) + return 0 if summary["rounds"] == summary["rounds_expected"] else 1 + + +def cmd_eff_report(args) -> int: + from .efficiency import report as eff_report + + runs = eff_report.load([Path(d) for d in args.run_dirs]) + Path(args.out).write_text(eff_report.markdown(runs, title=args.title)) + print(f"{len(runs)} parsed run(s) -> {args.out}") + return 0 + + def main(argv: list[str] | None = None) -> int: p = argparse.ArgumentParser(prog="python -m tools.eval") sub = p.add_subparsers(dest="cmd", required=True) @@ -143,6 +215,30 @@ def main(argv: list[str] | None = None) -> int: rep.add_argument("--html", default=None) rep.add_argument("--title", default="yeto RL capability evaluation") rep.set_defaults(fn=cmd_report) + from .efficiency.spec import ARMS, PRESETS + + ep = sub.add_parser("eff-plan") + ep.add_argument("preset", choices=sorted(PRESETS)) + ep.add_argument("--arm", choices=ARMS, default=None) + ep.set_defaults(fn=cmd_eff_plan) + er = sub.add_parser("eff-run") + er.add_argument("preset", choices=sorted(PRESETS)) + er.add_argument("--arm", choices=ARMS, required=True) + er.add_argument("--run-id", required=True) + er.add_argument("--out", default=str(EFF_OUT)) + er.add_argument("--dry-run", action="store_true") + er.add_argument("--checkpoint", choices=("final", "every-round"), default="final", + help="yeto+Miles round cuts: only at the end (fair control) or every round " + "(product default)") + er.set_defaults(fn=cmd_eff_run) + eparse = sub.add_parser("eff-parse") + eparse.add_argument("run_dir") + eparse.set_defaults(fn=cmd_eff_parse) + erep = sub.add_parser("eff-report") + erep.add_argument("run_dirs", nargs="+") + erep.add_argument("--out", required=True) + erep.add_argument("--title", default="yeto RL efficiency evaluation") + erep.set_defaults(fn=cmd_eff_report) args = p.parse_args(argv) return args.fn(args) diff --git a/tools/eval/efficiency/__init__.py b/tools/eval/efficiency/__init__.py new file mode 100644 index 00000000..54197bef --- /dev/null +++ b/tools/eval/efficiency/__init__.py @@ -0,0 +1,24 @@ +"""yeto RL efficiency evaluation (S26 E0.1 / E0.2). + +Eight arms per model preset: {bare, yeto} x {verl, miles} x {sync, async}. The +bare arms run the framework's own entry point from the same image (no yeto +code on their path); the yeto arms run the same framework through yeto. Every +arm is generated from one :class:`.spec.EffSpec`, so model, data, order seed, +reward, hyperparameters, max tokens, card type and count and image are equal +by construction (:func:`.spec.control_fingerprint`). + + python -m tools.eval eff-plan qwen3-0.6b [--arm ARM] + python -m tools.eval eff-run qwen3-0.6b --arm bare-verl-sync --run-id ID [--dry-run] + python -m tools.eval eff-parse RUN_DIR + python -m tools.eval eff-report RUN_DIR ... --out report.md + +* :mod:`.spec` -- presets, arms, the per-arm command plans; +* :mod:`.parse` -- framework logs and yeto tapes -> one per-round record format; +* :mod:`.metrics` -- per-run summary (throughput, segments, cost, off-policy); +* :mod:`.report` -- Markdown report with yeto-vs-bare overhead and range verdicts; +* :mod:`.modal_eff` -- the Modal harness (bare arms and yeto+verl); + yeto+Miles runs through the yeto launcher (``tools/eval/miles_local.sh``). + +Raw data is written to ``/raw/`` before anything is computed; +``parsed.json`` and the report are derived from it and can be rebuilt. +""" diff --git a/tools/eval/efficiency/algo-eff-miles.json b/tools/eval/efficiency/algo-eff-miles.json new file mode 100644 index 00000000..f45d2e97 --- /dev/null +++ b/tools/eval/efficiency/algo-eff-miles.json @@ -0,0 +1,53 @@ +{ + "advantage": { + "estimator": "grpo", + "reward_binary": false, + "reward_postprocess": null, + "rewards_normalization": true, + "std_normalization": true, + "whiten": false + }, + "correction": { + "batch_normalize": false, + "function": null, + "level": "token", + "lower": 0.0, + "method": "tis", + "mismatch_metrics": false, + "opsm_delta": null, + "ratio_mask": "none", + "upper": 2.0, + "use_rollout_logprobs": false, + "veto": null + }, + "entropy_coef": 0.0, + "execution": { + "group_required": true, + "max_policy_staleness": 2, + "needs_critic": false, + "needs_rollout_logprobs": false + }, + "kl": { + "coef": null, + "estimator": null, + "placement": "none", + "unbiased": false + }, + "loss": { + "aggregation": "default", + "custom_loss": null, + "eps_clip": 0.2, + "eps_clip_c": null, + "eps_clip_high": 0.28, + "reducer": null, + "variant": "policy_loss" + }, + "plugins": [], + "sampling": { + "filter": null, + "max_replacements": null, + "over_sampling_batch_size": null, + "overlong_filter": false + }, + "schema": "yeto-rl-algorithm-spec-v2" +} \ No newline at end of file diff --git a/tools/eval/efficiency/cluster.py b/tools/eval/efficiency/cluster.py new file mode 100644 index 00000000..6ed2b040 --- /dev/null +++ b/tools/eval/efficiency/cluster.py @@ -0,0 +1,161 @@ +"""E11 whole-cluster comparison: reward against samples and time on one card total. + +Every arm becomes one list of :class:`Point` (one per trained batch, in end-time +order). A single-island arm gives one point per round (``parsed.json`` of an +``eff-run`` directory); the 2-island yeto arm (``logs/E11/run2i.sh``) gives one +point per island round, read from each island's tape, so its two islands +interleave by time. + +Definitions (SESSION26/gpu-review/E11.md §3, fixed before the runs): + +* reward of a point = mean reward of the batch's rollout samples; +* smoothed reward = sample-weighted mean over the latest ``window`` samples + (3 single-island rounds), evaluated at every point; +* target T = sample-weighted mean reward of the first ``2 * round_samples`` + samples of every run, pooled, plus ``margin``; +* time to T = first point whose smoothed reward >= T, measured from the run's + launch (``t0``: the local launch command) and from the end of its first batch + (no cold start); GPU hours = that time x the run's card total. +""" + +from __future__ import annotations + +import json +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from .parse import _tape_events, _utc_file, parse_yeto + + +@dataclass(frozen=True) +class Point: + t_end: float + samples: int + reward: float + + +@dataclass +class Curve: + run_id: str + arm: str + cards: int + t0: float # local launch start (unix s) + t_last: float # the run's GPU containers ended (or its last point) + points: list[Point] + usd_extra: float = 0.0 # e.g. the 2-island arm's head VM + + +def _from_parsed(run_dir: Path) -> Curve: + parsed = json.loads((run_dir / "parsed.json").read_text()) + pts = [Point(float(r["t_end_unix"]), int(r["samples"]), float(r["reward_mean"])) + for r in sorted(parsed["rounds"], key=lambda r: r["round"]) + if r.get("t_end_unix") is not None and r.get("samples") and r.get("reward_mean") is not None] + t0 = _launch_start(run_dir) or float(parsed["container_start_unix"]) + end = parsed.get("container_end_unix") or (pts[-1].t_end if pts else t0) + return Curve(parsed["run_id"], parsed["arm"], int(parsed["cards"]["total"]), t0, float(end), pts) + + +def _launch_start(run_dir: Path) -> float | None: + """The local launch time: miles_local.sh's start_utc.txt, else the driver's first log line.""" + utc = _utc_file(run_dir / "raw" / "start_utc.txt") + if utc: + return utc + log = run_dir / "driver.log" + if not log.is_file(): + return None + try: + first = log.read_text().split(" ", 1)[0] + return datetime.strptime(first, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc).timestamp() + except ValueError: + return None + + +def _from_islands(run_dir: Path, *, cards: int, head_usd_per_hour: float) -> Curve: + """A run2i.sh directory: one tape per island under ``tape-direct/``.""" + pts: list[Point] = [] + tapes = sorted((run_dir / "tape-direct").rglob("rl-island-*.jsonl")) + if not tapes: + raise ValueError(f"{run_dir}: no island tapes under tape-direct/") + for tape in tapes: + island = parse_yeto(_tape_events(tape.parent), mode="async") + pts += [Point(float(r["t_end_unix"]), int(r["samples"]), float(r["reward_mean"])) + for r in island.rounds + if r.get("t_end_unix") is not None and r.get("samples") and r.get("reward_mean") is not None] + pts.sort(key=lambda p: p.t_end) + t0 = _utc_file(run_dir / "start_utc.txt") + end = _utc_file(run_dir / "end_utc.txt") or (pts[-1].t_end if pts else t0) + hours = (end - t0) / 3600 if t0 and end else 0.0 + return Curve(run_dir.name, "yeto-miles-2islands", cards, float(t0), float(end), pts, + usd_extra=hours * head_usd_per_hour) + + +def load(run_dir: Path, *, cards: int = 4, head_usd_per_hour: float = 0.0) -> Curve: + run_dir = Path(run_dir) + if (run_dir / "parsed.json").is_file(): + return _from_parsed(run_dir) + return _from_islands(run_dir, cards=cards, head_usd_per_hour=head_usd_per_hour) + + +def smoothed(points: list[Point], window: int) -> list[float]: + """Sample-weighted mean reward over the latest ``window`` samples at every point.""" + out = [] + for i in range(len(points)): + n = s = 0.0 + for p in reversed(points[:i + 1]): + take = min(p.samples, window - n) + s += p.reward * take + n += take + if n >= window: + break + out.append(s / n if n else float("nan")) + return out + + +def start_reward(curves: list[Curve], samples: int) -> float: + """Pooled sample-weighted mean reward of every run's first ``samples`` samples.""" + n = s = 0.0 + for c in curves: + left = samples + for p in c.points: + take = min(p.samples, left) + s += p.reward * take + n += take + left -= take + if left <= 0: + break + return s / n if n else float("nan") + + +def metrics(c: Curve, *, target: float, window: int, usd_per_card_hour: float) -> dict[str, Any]: + sm = smoothed(c.points, window) + hit = next((i for i, v in enumerate(sm) if v >= target), None) + total = sum(p.samples for p in c.points) + run_s = c.t_last - c.t0 + out: dict[str, Any] = { + "run_id": c.run_id, "arm": c.arm, "cards": c.cards, "points": len(c.points), "samples": total, + "first_batch_s": c.points[0].t_end - c.t0 if c.points else None, + "train_s": c.points[-1].t_end - c.points[0].t_end if len(c.points) > 1 else None, + "samples_per_s": ((total - c.points[0].samples) / (c.points[-1].t_end - c.points[0].t_end) + if len(c.points) > 1 else None), + "reward_last3": sm[-1] if sm else None, + "run_s": run_s, "usd_run": run_s / 3600 * c.cards * usd_per_card_hour + c.usd_extra, + "target": target, "reached": hit is not None, + } + if hit is not None: + p = c.points[hit] + out.update(to_target_s=p.t_end - c.t0, to_target_after_first_s=p.t_end - c.points[0].t_end, + to_target_samples=sum(q.samples for q in c.points[:hit + 1])) + out["to_target_gpu_h"] = out["to_target_s"] / 3600 * c.cards + out["to_target_usd"] = out["to_target_gpu_h"] * usd_per_card_hour + return out + + +def curve_table(c: Curve, window: int) -> list[dict[str, Any]]: + sm, cum = smoothed(c.points, window), 0 + rows = [] + for p, v in zip(c.points, sm): + cum += p.samples + rows.append({**asdict(p), "t_from_launch_s": p.t_end - c.t0, "cum_samples": cum, "smoothed": v}) + return rows diff --git a/tools/eval/efficiency/eff_prep.py b/tools/eval/efficiency/eff_prep.py new file mode 100644 index 00000000..90e9bb73 --- /dev/null +++ b/tools/eval/efficiency/eff_prep.py @@ -0,0 +1,97 @@ +"""Model and data preparation for the bare-framework arms (no yeto import). + +Run inside the container as ``python eff_prep.py ``. It downloads +the pinned model snapshot and the pinned dataset revision, then writes the +training prompts in the dataset's own row order: + +* ``/root/eff-data/train.parquet`` and ``test.parquet`` -- verl rows, the same + shape as ``yeto.rl.adapters.verl.island_entry.prepare_data`` (the yeto+verl + arm's data), so bare verl and yeto+verl read identical rows; +* ``/root/eff-data/train.jsonl`` -- Miles rows (``prompt`` = chat messages, + ``label``), for ``--input-key prompt --label-key label --apply-chat-template``; +* ``/root/eff-data/order.json`` -- sha256 of the label sequence (row order) and + the row count, so a report can show that two arms read the same rows. + +Shuffling is left to each framework (seeded by the spec's ``seed``), as in the +yeto arm of the same framework. +""" + +from __future__ import annotations + +import hashlib +import json +import sys +from pathlib import Path + +DATA_ROOT = Path("/root/eff-data") +MODEL_ROOT = Path("/root/eff-models") +VAL_ROWS = 64 + + +def source_rows(spec: dict, download) -> dict[str, list[tuple[list, str]]]: + """Split -> ``(chat messages, label)`` in the dataset's own row order. + + Two layouts: ``train.parquet``/``test.parquet`` with ``messages`` and ``label`` + (zhuzilin/gsm8k), or one JSONL file ``spec["data_file"]`` with ``prompt`` (chat + messages) and ``label`` (zhuzilin/dapo-math-17k), whose test rows are the + ``VAL_ROWS`` rows after the training rows. ``spec["train_rows"]`` (0: all) keeps + the first rows only, so every arm trains on the same prompt set in one epoch. + """ + import pandas as pd + + n = int(spec.get("train_rows") or 0) + if spec.get("data_file"): + path = download(spec["data"], spec["data_file"], repo_type="dataset", revision=spec["data_revision"]) + with open(path) as fh: + records = [(r["prompt"], str(r["label"])) for r in map(json.loads, fh)] + if not n: + raise ValueError("a JSONL data_file needs train_rows (its test rows follow them)") + return {"train": records[:n], "test": records[n:n + VAL_ROWS]} + out = {} + for split in ("train", "test"): + frame = pd.read_parquet(download(spec["data"], f"{split}.parquet", repo_type="dataset", + revision=spec["data_revision"])) + records = [(list(row["messages"]), str(row["label"])) for _, row in frame.iterrows()] + out[split] = records[:n] if (split == "train" and n) else records[:VAL_ROWS] if split == "test" else records + return out + + +def yeto_task_rows(spec: dict, download) -> list[dict]: + """The training rows as a yeto local task file (``messages``, ``label``, task id + ``gsm``): the yeto arms read the same rows as the bare arms' ``train.jsonl``.""" + return [{"messages": [{"role": m["role"], "content": m["content"]} for m in messages], "label": label, + "metadata": {"task_id": f"gsm{i + 1}"}} + for i, (messages, label) in enumerate(source_rows(spec, download)["train"])] + + +def main(spec_path: str) -> int: + import pandas as pd + from huggingface_hub import hf_hub_download, snapshot_download + + spec = json.loads(Path(spec_path).read_text()) + model_dir = MODEL_ROOT / spec["model"].split("/")[-1] + snapshot_download(spec["model"], revision=spec["model_revision"], local_dir=str(model_dir)) + DATA_ROOT.mkdir(parents=True, exist_ok=True) + order = {} + for split, records in source_rows(spec, hf_hub_download).items(): + rows = [{ + "data_source": "yeto/" + spec["data"], + "prompt": [{"role": m["role"], "content": m["content"]} for m in messages], + "ability": "math", + "reward_model": {"style": "rule", "ground_truth": str(label)}, + "extra_info": {"split": split, "index": i}, + } for i, (messages, label) in enumerate(records)] + pd.DataFrame(rows).to_parquet(DATA_ROOT / f"{split}.parquet") + labels = [r["reward_model"]["ground_truth"] for r in rows] + order[split] = {"rows": len(rows), "label_sha256": hashlib.sha256("\n".join(labels).encode()).hexdigest()} + if split == "train": + with (DATA_ROOT / "train.jsonl").open("w") as fh: + for r in rows: + fh.write(json.dumps({"prompt": r["prompt"], "label": r["reward_model"]["ground_truth"]}) + "\n") + (DATA_ROOT / "order.json").write_text(json.dumps({**order, "model_dir": str(model_dir)}, indent=1)) + print(json.dumps({"model_dir": str(model_dir), **order})) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1])) diff --git a/tools/eval/efficiency/eff_reward.py b/tools/eval/efficiency/eff_reward.py new file mode 100644 index 00000000..4358028a --- /dev/null +++ b/tools/eval/efficiency/eff_reward.py @@ -0,0 +1,39 @@ +"""GSM8K exact-number reward for the bare-framework arms (no yeto import). + +The grading rule is a verbatim copy of ``yeto.rl.rewards.builtin.gsm8k_grade`` +(the reward of the yeto arms, ``--reward-function gsm8k_reward:score``); a CPU +test asserts both give the same value on the same cases. The file is copied +into the bare containers as ``/root/eff_reward.py``: + +* verl: ``reward.custom_reward_function.path=/root/eff_reward.py`` and + ``reward.custom_reward_function.name=compute_score``; +* Miles: ``--custom-rm-path eff_reward.score`` (``/root`` on ``PYTHONPATH``). +""" + +from __future__ import annotations + +import re + + +def _num(s): + s = str(s).replace(",", "") + m = re.findall(r"-?\d+(?:\.\d+)?", s) + return float(m[-1]) if m else None + + +def grade(response: str, label) -> float: + gold = _num(str(label).split("####")[-1]) if label is not None else None + text = response or "" + boxed = re.findall(r"\\boxed\{([^}]*)\}", text) + pred = _num(boxed[-1]) if boxed else _num(text[-200:]) + return 1.0 if (gold is not None and pred is not None and abs(pred - gold) < 1e-6) else 0.0 + + +def compute_score(data_source, solution_str, ground_truth, extra_info=None, **kwargs) -> float: + """verl custom reward signature.""" + return grade(solution_str, ground_truth) + + +async def score(args, sample, **kwargs) -> float: + """Miles custom reward signature (``--custom-rm-path``).""" + return grade(sample.response or "", sample.label) diff --git a/tools/eval/efficiency/metrics.py b/tools/eval/efficiency/metrics.py new file mode 100644 index 00000000..0df3bfa2 --- /dev/null +++ b/tools/eval/efficiency/metrics.py @@ -0,0 +1,141 @@ +"""Per-run summary from ``parsed.json`` (the definitions every arm shares). + +* Steady window: the rounds after the first (the first carries engine start-up, + CUDA-graph capture and caches); ``steady_wall_s`` = the sum of their + ``round_s`` (step end minus previous step end). +* Throughput per card: steady tokens / steady wall / every card of the arm + (rollout and trainer cards both count). +* Segments: per-round medians over the steady window; ``unaccounted_s`` = round + minus the named segments (offload/onload, reward, advantage, orchestration). +* GPU: NVML samples inside the steady window; ``bubble_frac`` = the share of + card samples with utilization below :data:`BUBBLE_UTIL_PCT`. +* Cost: ``usd_per_card_hour`` from ``prices.json`` x cards x container seconds + (whole run) and x steady wall (steady state), per 1k trained samples. + +A value whose inputs are missing is ``None`` and its name is listed in +``missing``; nothing is filled in. +""" + +from __future__ import annotations + +import json +import statistics +from pathlib import Path +from typing import Any + +PRICES = Path(__file__).with_name("prices.json") +BUBBLE_UTIL_PCT = 5.0 +SEGMENTS = ("gen_s", "train_s", "weight_sync_s", "sampler_update_s", "wait_s") + + +def price(gpu: str, path: Path = PRICES) -> float: + table = json.loads(path.read_text()) + if gpu not in table: + raise KeyError(f"no price for GPU {gpu!r} in {path}") + return float(table[gpu]["usd_per_card_hour"]) + + +def _vals(rows: list[dict], key: str) -> list[float]: + return [float(r[key]) for r in rows if r.get(key) is not None] + + +def _median(values: list[float]) -> float | None: + return round(statistics.median(values), 4) if values else None + + +def _sum_all(rows: list[dict], key: str) -> float | None: + """Sum over rows, or None when any row lacks the value (a partial sum would mislead).""" + vals = [r.get(key) for r in rows] + return None if not vals or any(v is None for v in vals) else float(sum(vals)) + + +def summarize(parsed: dict[str, Any], *, prices_path: Path = PRICES) -> dict[str, Any]: + rounds = sorted(parsed["rounds"], key=lambda r: r["round"]) + cards = int(parsed["cards"]["total"]) + steady = rounds[1:] + s: dict[str, Any] = {"arm": parsed["arm"], "run_id": parsed["run_id"], "rounds": len(rounds), + "rounds_expected": parsed["spec"]["rounds"], "cards": cards, "gpu": parsed["gpu"], + "rc": parsed.get("rc")} + wall = _sum_all(steady, "round_s") + s["steady_wall_s"] = round(wall, 3) if wall is not None else None + s["rate_rounds"] = {} + for key, out in (("gen_tokens", "gen_tok_per_card_s"), ("train_tokens", "train_tok_per_card_s"), + ("samples", "samples_per_s")): + # over the steady rounds that report the value (bare verl fully_async never + # logs its last step's batch); the round count is kept when it is fewer + known = [r for r in steady if r.get(key) is not None and r.get("round_s") is not None] + time_s = sum(r["round_s"] for r in known) * (cards if key != "samples" else 1) + s[out] = round(sum(r[key] for r in known) / time_s, 3) if known and time_s else None + if known and len(known) < len(steady): + s["rate_rounds"][out] = f"{len(known)}/{len(steady)}" + s["round_s_median"] = _median(_vals(steady, "round_s")) + s["round_s_range"] = [min(_vals(steady, "round_s")), max(_vals(steady, "round_s"))] if _vals(steady, "round_s") \ + else None + for seg in SEGMENTS: + s[f"{seg}_median"] = None if parsed["sources"].get(seg) in ("n/a",) else _median(_vals(steady, seg)) + named = [seg for seg in SEGMENTS if parsed["sources"].get(seg) != "n/a"] + rest = [r["round_s"] - sum(r[seg] for seg in named) for r in steady + if r.get("round_s") is not None and all(r.get(seg) is not None for seg in named)] + s["unaccounted_s_median"] = _median(rest) + s["segments_raw_median"] = _raw_medians(parsed.get("segments_raw") or []) + # reward against wall time and samples (every round; t from the first step end) + t0 = rounds[0].get("t_end_unix") if rounds else None + cum, curve = 0.0, [] + for r in rounds: + cum = cum + r["samples"] if r.get("samples") is not None and cum is not None else None + curve.append({"round": r["round"], "t_s": round(r["t_end_unix"] - t0, 3) + if r.get("t_end_unix") is not None and t0 is not None else None, + "samples": cum, "reward_mean": r.get("reward_mean")}) + s["reward_curve"] = curve + s["reward_last3_mean"] = _median(_vals(rounds[-3:], "reward_mean")) # median of the last three + # GPU (NVML samples inside the steady window) + lo = rounds[0].get("t_end_unix") if rounds else None + hi = rounds[-1].get("t_end_unix") if rounds else None + nv = [x for x in parsed.get("nvml") or [] if lo is not None and hi is not None and lo <= x["t_unix"] <= hi] + util = [x["util"] for x in nv if x.get("util") is not None] + s["nvml_samples"] = len(util) + s["nvml_interval_s"] = parsed.get("nvml_interval_s") + s["gpu_util_mean"] = round(statistics.fmean(util), 2) if util else None + s["bubble_frac"] = round(sum(u < BUBBLE_UTIL_PCT for u in util) / len(util), 4) if util else None + mem = [x["mem_mb"] for x in parsed.get("nvml") or [] if x.get("mem_mb") is not None] + s["gpu_mem_peak_mb"] = max(mem) if mem else None + # cost + rate = price(parsed["gpu"], prices_path) + s["usd_per_card_hour"] = rate + start, end = parsed.get("container_start_unix"), parsed.get("container_end_unix") + total_samples = _sum_all(rounds, "samples") + if start and end: + s["usd_run"] = round(rate * cards * (end - start) / 3600, 4) + s["usd_per_1k_samples_run"] = round(s["usd_run"] / total_samples * 1000, 4) if total_samples else None + else: + s["usd_run"] = s["usd_per_1k_samples_run"] = None + steady_samples = _sum_all(steady, "samples") + s["usd_per_1k_samples_steady"] = round(rate * cards * wall / 3600 / steady_samples * 1000, 4) \ + if wall and steady_samples else None + # off-policy + s["is_clip_frac_mean"] = round(statistics.fmean(_vals(rounds, "is_clip_frac")), 5) \ + if _vals(rounds, "is_clip_frac") else None + dropped, generated = _vals(rounds, "dropped_samples"), _vals(rounds, "generated_samples") + s["dropped_total"] = max(dropped) if dropped and "count/" in parsed["sources"].get("dropped_samples", "") \ + else (sum(dropped) if dropped else None) # verl's fully_async counters are cumulative + s["drop_frac"] = round(s["dropped_total"] / max(generated), 4) if s["dropped_total"] is not None and generated \ + and max(generated) else None + s["version_lag"] = [r.get("version_lag") for r in rounds if r.get("version_lag") is not None] or None + s["sources"] = parsed["sources"] + na = {k for k, v in parsed["sources"].items() if v == "n/a"} + s["not_applicable"] = sorted(k for k in s if s[k] is None and _field_of(k) in na) + s["missing"] = sorted(k for k, v in s.items() if v is None and k != "rc" and k not in s["not_applicable"]) + return s + + +# summary key -> the round field it is computed from (for the not-applicable list) +_DERIVED_FROM = {"dropped_total": "dropped_samples", "drop_frac": "dropped_samples"} + + +def _field_of(key: str) -> str: + return _DERIVED_FROM.get(key, key.removesuffix("_median")) + + +def _raw_medians(rows: list[dict[str, float]]) -> dict[str, float]: + keys = sorted({k for r in rows[1:] for k in r}) + return {k: _median([float(r[k]) for r in rows[1:] if r.get(k) is not None]) for k in keys} diff --git a/tools/eval/efficiency/modal_eff.py b/tools/eval/efficiency/modal_eff.py new file mode 100644 index 00000000..5e7d0478 --- /dev/null +++ b/tools/eval/efficiency/modal_eff.py @@ -0,0 +1,226 @@ +"""Modal harness of the efficiency arms that run in one container: the four bare +arms and yeto+verl (yeto+Miles runs through the yeto launcher, ``miles_local.sh``). + +One GPU container per run, ``:``, in the framework's image (the +image the yeto arm of that framework uses; for verl the bare arms skip yeto's +source patch). The container: + +1. checks every visible GPU is the asked type (``modal_verl._gpu``); +2. samples NVML every second into ``nvml.csv`` (index, util, memory, power); +3. bare: writes the spec, ``eff_reward.py`` and ``eff_prep.py`` under + ``/root/eff`` and runs the arm's bash script (prep, then the framework entry); + yeto+verl: starts the island Ray and runs ``island_entry`` (as T-EVAL E1); +4. records ``meta.json`` (start/end unix time, rc, GPU names, argv/script) and + copies its outputs into the results Volume ``/arm/`` before returning. + +The local side spawns the call in a detached app, waits, pulls ``/`` into +``/raw/`` and stops only its own app (``yeto-eff-``). +""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +import sys +import threading +import time +from pathlib import Path +from typing import Any + +from .. import modal_verl as mv +from . import spec +from .spec import EFF_DIR, OUT_DIR + +HERE = Path(__file__).resolve().parent +APP_PREFIX = "yeto-eff-" +NVML_QUERY = "timestamp,index,name,utilization.gpu,memory.used,memory.total,power.draw" +NVML_INTERVAL_MS = 1000 + + +# No-progress watchdog (E11: a hung bare verl run burned 52 min): a training step +# line is progress; the run is killed when none arrives for STALL_S, or before the +# first step for FIRST_STEP_S (cold start: model load, engine start, first rollout). +# Miles logs "'train/step': N", verl "step:N -". +PROGRESS = re.compile(r"'train/step'|\bstep:\d+ ") +STALL_S = 900 +FIRST_STEP_S = 1800 +WATCH_POLL_S = 10 + + +# ====================================================================== in-container +def _nvml(path: Path) -> subprocess.Popen: + return subprocess.Popen(["bash", "-c", f"nvidia-smi --query-gpu={NVML_QUERY} --format=csv " + f"-lms {NVML_INTERVAL_MS} >> {path}"]) + + +def _run_logged(argv: list[str], log: Path, timeout: float, *, watch_progress: bool = False, + **kwargs) -> int: + """Run ``argv``; every output line is written to ``log`` behind `` `` + (the step end times the parsers read). The process is killed at ``timeout`` and, + with ``watch_progress``, when no step line arrives in time (``PROGRESS``).""" + started = time.time() + deadline = started + timeout + last_step: list[float] = [] # time of the latest step line (empty: none yet) + lock = threading.Lock() + with log.open("a") as fh, subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, + errors="replace", **kwargs) as proc: + def note(text: str) -> None: + with lock: + fh.write(f"{time.time():.3f} [eff harness] {text}\n") + fh.flush() + + def watch() -> None: + while proc.poll() is None: + time.sleep(WATCH_POLL_S) + quiet = time.time() - (last_step[-1] if last_step else started) + limit = STALL_S if last_step else FIRST_STEP_S + if quiet > limit and proc.poll() is None: + note(f"killed: no training step for {quiet:.0f} s (limit {limit} s)") + proc.kill() + return + + if watch_progress: + threading.Thread(target=watch, daemon=True).start() + for line in proc.stdout: + if watch_progress and PROGRESS.search(line): + last_step.append(time.time()) + del last_step[:-1] + with lock: + fh.write(f"{time.time():.3f} {line}") + fh.flush() + if time.time() > deadline: + proc.kill() + note(f"killed at the {timeout:.0f} s timeout") + break + return proc.wait() + + +def arm(run: str, plan: dict) -> dict: + """One run of one arm (see the module doc).""" + out = Path(mv.RESULTS) / run / "arm" + out.mkdir(parents=True, exist_ok=True) + meta: dict[str, Any] = {"arm": plan["arm"], "container": os.environ.get("MODAL_TASK_ID"), + "cards": plan["cards"]["total"]} + meta["gpu"] = mv._gpu(out, {"gpu": plan["gpu"]}) + nvml = _nvml(out / "nvml.csv") + meta["start_unix"] = time.time() + copies = {"eff-out": Path(OUT_DIR), "yeto-output": Path("/root/yeto-output")} + budget = plan["spec"]["timeout_s"] - 300 # leaves time to copy and commit + log = out / "main.log" + try: + if plan["system"] == "bare": + eff = Path(EFF_DIR) + eff.mkdir(parents=True, exist_ok=True) + for name, text in plan["files"].items(): + (eff / name).write_text(text) + (eff / "spec.json").write_text(json.dumps(plan["spec"])) + (out / "script.sh").write_text(plan["script"]) + meta["rc"] = _run_logged(["bash", str(out / "script.sh")], log, budget, watch_progress=True) + else: + sys.path.insert(0, mv.WORKDIR) + env = {**mv._island_env({"gpu": plan["gpu"]}, 0), **plan.get("env", {})} + ray = str(Path(mv.VERL_PY).with_name("ray")) + _run_logged([ray, "start", "--head", "--node-ip-address=127.0.0.1", "--port=6379", "--num-cpus=32", + f"--num-gpus={plan['cards']['total']}", "--include-dashboard=false", + "--temp-dir=/root/eval-ray"], log, 300, env=env) + meta["argv"] = [mv.VERL_PY, *plan["argv"]] + meta["rc"] = _run_logged(meta["argv"], log, budget, watch_progress=True, env=env, cwd="/root") + except BaseException as exc: # noqa: BLE001 - recorded, then re-raised + meta["error"] = f"{type(exc).__name__}: {exc}" + raise + finally: + meta["end_unix"] = time.time() + nvml.terminate() + mv._finish(out, meta, meta["start_unix"], copies) + return meta + + +# ====================================================================== app +def build_app(run: str, plan: dict): + import modal + + sys.path.insert(0, str(mv.REPO)) + from modal._vendor import cloudpickle + + # the container functions travel by value: `tools` is not importable in the images + cloudpickle.register_pickle_by_value(sys.modules[__name__]) + cloudpickle.register_pickle_by_value(mv) + from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + from yeto.rl.adapters.verl.image import modal_image + + if plan["image"] == "verl": + # bare arms run the fork as shipped: yeto's patch replaces the LoRA weight-sync + # sender with an import of yeto (eff01-bare-verl-async-a/d) + image = modal_image(modal, patched=plan["system"] == "yeto").env({"PYTHONUNBUFFERED": "1"}) + if plan["system"] == "yeto": + image = image.add_local_dir(str(mv.REPO), mv.WORKDIR, copy=False, ignore=MOUNT_IGNORE) + else: + image = (modal.Image.from_registry(MILES_NEXT_IMAGE.removeprefix("docker:")) + .env({"HOME": "/root", "PYTHONUNBUFFERED": "1"})) + app = modal.App(APP_PREFIX + run) + results = modal.Volume.from_name(mv.RESULTS_VOLUME, create_if_missing=True, version=2) + gpu = f"{plan['gpu']}:{plan['cards']['total']}" + fn = app.function(image=image, gpu=gpu, cpu=8.0 * plan["cards"]["total"], + memory=49152 * plan["cards"]["total"], timeout=plan["spec"]["timeout_s"], retries=0, + volumes={mv.RESULTS: results}, single_use_containers=True, name="arm", serialized=True)(arm) + return app, fn + + +MOUNT_IGNORE = [".git", ".venv", ".claude", "**/__pycache__", "**/*.pyc", "s1-runs", "syncer/target", + "evidence*", "openspec", "docs", "tutorial", "data"] + + +def bare_files() -> dict[str, str]: + return {name: (HERE / name).read_text() for name in ("eff_reward.py", "eff_prep.py")} + + +def run(plan: dict, run_dir: Path) -> dict[str, Any]: + """Run one arm end to end (spawn, wait, pull, stop own app).""" + import modal + + run_id = plan["run_id"] + raw = run_dir / "raw" + raw.mkdir(parents=True, exist_ok=True) + log = run_dir / "driver.log" + if plan["system"] == "bare": + plan = {**plan, "files": bare_files()} + app, fn = build_app(run_id, plan) + began = time.time() + with modal.enable_output(), app.run(detach=True): # a killed driver never cancels a spawned call + app_id = app.app_id + (run_dir / "app_id.txt").write_text(f"{app_id}\n") + call = fn.spawn(run_id, plan) + mv._log(log, f"spawned arm {plan['arm']} app {APP_PREFIX + run_id}") + results = mv._wait_calls({"arm": call}, began + plan["spec"]["timeout_s"] + 600, log) + pulled = mv._modal_cli("volume", "get", "--force", mv.RESULTS_VOLUME, run_id, str(raw), timeout=1800) + stopped = mv._modal_cli("app", "stop", "-y", app_id, timeout=120) + summary = {"calls": results, "wall_s": round(time.time() - began, 1), "pull_rc": pulled.returncode, + "stop_rc": stopped.returncode} + (run_dir / "driver.json").write_text(json.dumps(summary, indent=1, default=str)) + return summary + + +# ====================================================================== yeto+Miles (launcher) +def run_miles_launcher(plan: dict, run_dir: Path, *, dry_run: bool = False) -> dict[str, Any]: + """The yeto+Miles arm: ``miles_local.sh`` with the plan's launcher arguments + (the T-EVAL Miles path; it pulls launch log, tapes and NVML into ``raw/``).""" + from .. import miles + + raw = run_dir / "raw" + env = dict(os.environ, P=plan["run_id"], R=str(raw), REPO=str(miles.REPO), PY=miles.PY, SKY=miles.SKY, + STALL_S=str(STALL_S), FIRST_STEP_S=str(FIRST_STEP_S), + STOREVOL=miles.STORE_VOLUME, TAPEVOL=miles.TAPE_VOLUME, HARD=str(plan["spec"]["timeout_s"]), + HOME="/home/michael", ARGS=plan["launcher_args"], ISLANDS="1") + if dry_run: + env["DRY"] = "1" + sp = spec.EffSpec(**plan["spec"]) + if spec.uses_task_file(sp): + run_dir.mkdir(parents=True, exist_ok=True) + env["EXTRA_FILES"] = str(spec.write_task_file(sp, run_dir / spec.YETO_TASK_FILE)) + log = run_dir / ("dry-run.log" if dry_run else "driver.log") + with log.open("w") as fh: + rc = subprocess.run(["bash", str(miles.HERE / "miles_local.sh")], env=env, stdout=fh, + stderr=subprocess.STDOUT, cwd="/tmp").returncode + return {"script": "miles_local.sh", "rc": rc, "log": str(log)} diff --git a/tools/eval/efficiency/parse.py b/tools/eval/efficiency/parse.py new file mode 100644 index 00000000..120d398b --- /dev/null +++ b/tools/eval/efficiency/parse.py @@ -0,0 +1,478 @@ +"""Framework logs and yeto tapes -> one per-round record format. + +Every parser returns a :class:`Parsed`: ``rounds`` (one dict per optimizer +step, the fields of :data:`ROUND_FIELDS`), ``sources`` (per field: the log key +it came from, ``"derived: ..."`` when it is computed from logged values by an +exact identity, ``MISSING`` when no source in the raw data carries it, or +``NA`` when the field does not apply to the arm), and ``segments_raw`` (every +native timing per round, for the overhead table). + +A field that a log does not carry stays ``None``; nothing is estimated. + +Formats read (each one is cited where it is parsed): + +* verl ``FileLogger`` JSONL (``VERL_FILE_LOGGER_PATH``): one + ``{"step": N, "data": {...}}`` per step (fork acad9875 ``verl/utils/tracking.py`` + ``FileLogger.log``); keys from ``metric_utils.compute_timing_metrics`` / + ``compute_throughout_metrics`` (``timing_s/*``, ``perf/*``) and the fully_async + trainer (``fully_async/*``, ``timing_s/param_sync``, ``timing_s/wait_for_enough_samples``); +* verl console lines ``step:N - key:value - ...`` (``LocalLogger``, + ``concat_dict_to_str``) behind the harness's `` `` line prefix, for the + step end time; +* Miles ``perf N: {python dict}`` log lines: rollout side + (``miles/ray/rollout/metrics.py:93``, keys ``rollout/*``, ``perf/rollout_time``, + ``perf/tokens_per_gpu_per_sec``) and train side + (``miles/utils/train_metric_utils.py:56``, keys ``perf/_time``, + ``perf/actor_train_tok_per_s``); the line's ``[YYYY-mm-dd HH:MM:SS.mmm`` stamp + (Miles' logger format) gives the time; +* yeto event tapes (JSONL ``{"event": ...}``): ``rl_round_wall``, + ``rl_local_round``, ``rl_round_trained``, ``rl_timeline_span``, + ``rl_driver_phase``, ``rl_rollout_carry_over``, ``rl_resource_sample``; +* NVML CSV (``nvidia-smi --query-gpu=timestamp,index,name,utilization.gpu, + memory.used,memory.total,power.draw --format=csv``), with or without a + ``,`` prefix per line (``miles_local.sh`` polls every 20 s). +""" + +from __future__ import annotations + +import ast +import csv +import datetime as _dt +import io +import json +import re +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Iterable + +MISSING = "missing" +NA = "n/a" +ROUND_FIELDS = ( + "round", "t_end_unix", "round_s", "gen_s", "train_s", "weight_sync_s", "sampler_update_s", "wait_s", + "gen_tokens", "train_tokens", "samples", "reward_mean", "is_clip_frac", "dropped_samples", + "generated_samples", "version_lag", +) +ANSI = re.compile(r"\x1b\[[0-9;]*m") + + +@dataclass +class Parsed: + rounds: list[dict[str, Any]] = field(default_factory=list) + sources: dict[str, str] = field(default_factory=dict) + segments_raw: list[dict[str, float]] = field(default_factory=list) + + def to_dict(self) -> dict[str, Any]: + return {"rounds": self.rounds, "sources": self.sources, "segments_raw": self.segments_raw} + + +def _first(data: dict, keys: Iterable[str]) -> tuple[Any, str | None]: + for key in keys: + if data.get(key) is not None: + return data[key], key + return None, None + + +def _finish(parsed: Parsed, na: Iterable[str] = ()) -> Parsed: + """Fill ``sources`` for fields no round carries (missing or not applicable).""" + na = set(na) + parsed.sources["round"] = "step index (verl global step; Miles / yeto rollout id)" + for name in ROUND_FIELDS[1:]: + if name in na: + parsed.sources[name] = NA + for r in parsed.rounds: + r[name] = None + elif name not in parsed.sources or all(r.get(name) is None for r in parsed.rounds): + parsed.sources[name] = MISSING + return parsed + + +def _record(rnd: int) -> dict[str, Any]: + return {name: None for name in ROUND_FIELDS} | {"round": rnd} + + +def _set(parsed: Parsed, rec: dict, name: str, value: Any, source: str | None) -> None: + if value is None or source is None: + return + rec[name] = value + parsed.sources.setdefault(name, source) + + +def _step_ends(log_text: str, pattern: re.Pattern) -> dict[int, float]: + """Round -> unix time of its log line (lines carry a `` `` prefix).""" + out: dict[int, float] = {} + for line in log_text.splitlines(): + m = pattern.search(ANSI.sub("", line)) + if m: + out[int(m.group(2))] = float(m.group(1)) + return out + + +def _round_s_from_ends(parsed: Parsed) -> None: + """round_s = time between consecutive step ends (one definition for every arm).""" + prev = None + for r in parsed.rounds: + t = r.get("t_end_unix") + if t is not None and prev is not None: + r["round_s"] = round(t - prev, 3) + parsed.sources.setdefault("round_s", "derived: step end - previous step end") + prev = t + + +# ====================================================================== verl +VERL_STEP_LINE = re.compile(r"^(\d+(?:\.\d+)?) .*?\bstep:(\d+) - ") + + +def parse_verl(metrics_jsonl: str, log_text: str = "", *, samples_per_round: int, trainer_cards: int, + mode: str) -> Parsed: + """verl FileLogger rows (+ console step lines for end times).""" + p = Parsed() + ends = _step_ends(log_text, VERL_STEP_LINE) + recs: dict[int, dict[str, Any]] = {} + segs: dict[int, dict[str, float]] = {} + for line in metrics_jsonl.splitlines(): + if not line.strip(): + continue + row = json.loads(line) + d, step = row.get("data") or {}, int(row["step"]) + rnd = step + if mode == "async" and "training/global_step" in d: + # bare fully_async logs a step's training metrics one parameter version + # later (row step N+1, training/global_step N) and never logs the last + # step's; the rows without it carry the weight sync that ends version N. + rnd = int(d["training/global_step"]) + elif mode == "async" and step == 0: + continue # the weight sync before the first training step + rec = recs.setdefault(rnd, _record(rnd)) + if rnd == step: + _set(p, rec, "t_end_unix", ends.get(step), "console step line (harness time prefix)") + # fully_async: the trainer's timing_s/gen only pulls finished samples from the + # queue; the rollouter's active time in the version window is the generation + gen, key = _first(d, ("fully_async/rollouter/active_time",) if mode == "async" else ("timing_s/gen",)) + _set(p, rec, "gen_s", gen, key) + parts = [k for k in ("timing_s/old_log_prob", "timing_s/update_actor") if d.get(k) is not None] + if parts: + _set(p, rec, "train_s", sum(float(d[k]) for k in parts), " + ".join(parts)) + sync, key = _first(d, ("timing_s/update_weights", "timing_s/param_sync")) + _set(p, rec, "weight_sync_s", sync, key) + wait, key = _first(d, ("timing_s/wait_for_enough_samples",)) + _set(p, rec, "wait_s", wait, key) + if d.get("response_length/mean") is not None: # every step trains exactly samples_per_round responses + _set(p, rec, "gen_tokens", float(d["response_length/mean"]) * samples_per_round, + f"derived: response_length/mean x {samples_per_round} responses") + _set(p, rec, "samples", samples_per_round, "config: samples per step (verl trains a fixed batch)") + tokens, key = _first(d, ("perf/total_num_tokens",)) + if tokens is None and d.get("global_seqlen/mean") is not None: + tokens, key = float(d["global_seqlen/mean"]) * trainer_cards, \ + f"derived: global_seqlen/mean x {trainer_cards} trainer ranks" + _set(p, rec, "train_tokens", tokens, key) + reward, key = _first(d, ("critic/score/mean", "critic/rewards/mean")) + _set(p, rec, "reward_mean", reward, key) + clip, key = _first(d, ("rollout_corr/rollout_is_ratio_fraction_high", "yeto_mismatch/tis_clipfrac")) + _set(p, rec, "is_clip_frac", clip, key) + dropped, key = _first(d, ("fully_async/count/dropped_stale_samples",)) + _set(p, rec, "dropped_samples", dropped, key) + generated, key = _first(d, ("fully_async/count/total_generated_samples",)) + _set(p, rec, "generated_samples", generated, key) + segs.setdefault(rnd, {}).update({k: float(v) for k, v in d.items() if k.startswith("timing_s/") + and isinstance(v, (int, float))}) + for rec in recs.values(): # a step whose metrics verl never logged still trained one fixed batch + _set(p, rec, "samples", samples_per_round if rec["samples"] is None else rec["samples"], + "config: samples per step (verl trains a fixed batch)") + p.rounds = [recs[r] for r in sorted(recs)] + p.segments_raw = [segs[r] for r in sorted(recs)] + _round_s_from_ends(p) + na = {"sampler_update_s"} + if mode == "sync": + na |= {"wait_s", "dropped_samples", "generated_samples", "version_lag"} + return _finish(p, na) + + +# ====================================================================== Miles +MILES_DICT_LINE = re.compile(r"\b(perf|step|rollout) (\d+): (\{.*\})\s*$") +MILES_STAMP = re.compile(r"\[(\d{4}-\d\d-\d\d \d\d:\d\d:\d\d\.\d+)") + + +def _miles_stamp(line: str) -> float | None: + m = MILES_STAMP.search(line) + if not m: + return None + return _dt.datetime.strptime(m.group(1), "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=_dt.timezone.utc).timestamp() + + +def miles_dicts(log_text: str) -> dict[int, dict[str, Any]]: + """Rollout id -> {"rollout": {...}, "train": {...}, "step": {...}, "train_t": unix}.""" + out: dict[int, dict[str, Any]] = defaultdict(dict) + for raw in log_text.splitlines(): + line = ANSI.sub("", raw) + m = MILES_DICT_LINE.search(line) + if not m: + continue + try: + data = ast.literal_eval(m.group(3)) + except (ValueError, SyntaxError): + continue + rid = int(m.group(2)) + if m.group(1) == "perf" and ("perf/rollout_time" in data or any(k.startswith("rollout/") for k in data)): + out[rid].setdefault("rollout", {}).update(data) + elif m.group(1) == "perf": + out[rid].setdefault("train", {}).update(data) + out[rid]["train_t"] = _miles_stamp(line) + else: + out[rid].setdefault("step", {}).update(data) + return dict(out) + + +def parse_miles(log_text: str, *, rollout_cards: int, mode: str) -> Parsed: + p = Parsed() + for rid, parts in sorted(miles_dicts(log_text).items()): + ro, tr, st = parts.get("rollout", {}), parts.get("train", {}), parts.get("step", {}) + merged = {**ro, **tr, **st} + rec = _record(rid) + _set(p, rec, "t_end_unix", parts.get("train_t"), "Miles train perf line time stamp") + gen, key = _first(ro, ("perf/rollout_time",)) + _set(p, rec, "gen_s", gen, key) + train, key = _first(tr, ("perf/train_time",)) + _set(p, rec, "train_s", train, key) + sync, key = _first(tr, ("perf/update_weights_time",)) + _set(p, rec, "weight_sync_s", sync, key) + if mode == "async": + wait, key = _first(tr, ("perf/train_wait_time",)) + _set(p, rec, "wait_s", wait, key) + if ro.get("perf/tokens_per_gpu_per_sec") is not None and gen is not None: + _set(p, rec, "gen_tokens", float(ro["perf/tokens_per_gpu_per_sec"]) * float(gen) * rollout_cards, + f"derived: perf/tokens_per_gpu_per_sec x perf/rollout_time x {rollout_cards} rollout cards") + if tr.get("perf/actor_train_tok_per_s") is not None and tr.get("perf/actor_train_time") is not None: + _set(p, rec, "train_tokens", float(tr["perf/actor_train_tok_per_s"]) * float(tr["perf/actor_train_time"]), + "derived: perf/actor_train_tok_per_s x perf/actor_train_time") + samples, key = _first(merged, ("rollout/num_training_samples", "num_training_samples")) + _set(p, rec, "samples", samples, key) + reward, key = _first(merged, ("rollout/raw_reward", "rollout/episode_raw_reward", "rollout/rewards")) + _set(p, rec, "reward_mean", reward, key) + clip, key = _first(merged, ("train/tis_clipfrac", "train/tis_clip_frac", "tis_clipfrac")) + _set(p, rec, "is_clip_frac", clip, key) + lag = {k: v for k, v in merged.items() if "staleness" in k or "weight_version" in k} + _set(p, rec, "version_lag", lag or None, "Miles keys containing 'staleness' / 'weight_version'") + p.rounds.append(rec) + p.segments_raw.append({k: float(v) for k, v in merged.items() + if k.startswith("perf/") and k.endswith("_time") and isinstance(v, (int, float))}) + _round_s_from_ends(p) + na = {"sampler_update_s"} + if mode == "sync": + na |= {"wait_s", "dropped_samples", "generated_samples", "version_lag"} + return _finish(p, na) + + +# ====================================================================== yeto tape +def tape_rows(text: str) -> list[dict[str, Any]]: + rows = [] + for line in text.splitlines(): + line = line.strip() + start = line.find("{") + if start < 0: + continue + try: + row = json.loads(line[start:]) + except json.JSONDecodeError: + continue + if isinstance(row, dict) and row.get("event"): + rows.append(row) + return rows + + +# yeto task / phase names -> canonical segment (others stay in segments_raw only) +YETO_SEGMENT = {"generate": "gen_s", "train": "train_s", "publish": "weight_sync_s"} + + +def parse_yeto(events: list[dict[str, Any]], *, mode: str, framework_parsed: Parsed | None = None, + framework_round_offset: int = 0, skip: frozenset[str] = frozenset(), + rollouter_gen: bool = False) -> Parsed: + """yeto tape -> rounds; ``framework_parsed`` (the framework's own log in the + same run, e.g. Miles perf lines) fills fields the tape does not carry. Its + round ``rollout_id + framework_round_offset`` is the tape's ``rollout_id``: + 0 for Miles (perf ids are rollout ids), 1 for verl (FileLogger step = + global step, 1-based; T-EVAL eval-e1-verl-c: steps 1..4, rollout ids 0..3). + + ``rollouter_gen`` (yeto+verl fully_async): ``gen_s`` is the rollouter's + ``fully_async/rollouter/active_time`` from ``rl_publication_check.rollouter_timing``; + the check of policy version v closes round v - 1, the same window as the bare + arm's per-step value.""" + p = Parsed() + by = defaultdict(list) + for e in events: + by[e["event"]].append(e) + walls = {int(e["rollout_id"]): e for e in by["rl_round_wall"] if e.get("rollout_id") is not None} + trained = {int(e["rollout_id"]): e for e in by["rl_round_trained"] if e.get("rollout_id") is not None} + local = sorted(by["rl_local_round"], key=lambda e: float(e.get("time_unix") or 0)) + spans: dict[int, dict[str, float]] = defaultdict(lambda: defaultdict(float)) + for s in by["rl_timeline_span"]: + if s.get("rollout_id") is None or s.get("start") is None or s.get("end") is None: + continue + name = "wait" if s.get("kind") == "wait" else str(s.get("task")) + spans[int(s["rollout_id"])][name] += float(s["end"]) - float(s["start"]) + phases = _driver_phase_durations(by["rl_driver_phase"]) + drops = defaultdict(float) + for e in by["rl_rollout_carry_over"]: + if e.get("rollout_id") is not None and e.get("over_age_discarded_groups") is not None: + drops[int(e["rollout_id"])] += float(e["over_age_discarded_groups"]) + rollouter = _rollouter_active(by["rl_publication_check"]) if rollouter_gen else {} + rounds = sorted(set(walls) | set(trained)) + fw = {r["round"]: r for r in (framework_parsed.rounds if framework_parsed else [])} + for i, rid in enumerate(rounds): + rec = _record(rid) + w, t = walls.get(rid, {}), trained.get(rid, {}) + lr = local[i] if i < len(local) else {} + _set(p, rec, "t_end_unix", w.get("time_unix"), "rl_round_wall.time_unix") + seg = dict(spans.get(rid) or phases.get(rid) or {}) + src = "rl_timeline_span" if rid in spans else "rl_driver_phase (marker to next marker)" + for task, name in YETO_SEGMENT.items(): + if name not in skip: + _set(p, rec, name, seg.get(task), f"{src} {task}") + if rid in rollouter: + _set(p, rec, "gen_s", rollouter[rid], f"rl_publication_check.rollouter_timing {ROLLOUTER_ACTIVE}") + if mode == "async": + _set(p, rec, "wait_s", seg.get("wait"), f"{src} kind=wait") + _set(p, rec, "gen_tokens", lr.get("action_tokens"), "rl_local_round.action_tokens") + _set(p, rec, "samples", t.get("trained_samples"), "rl_round_trained.trained_samples") + _set(p, rec, "reward_mean", lr.get("reward_mean"), "rl_local_round.reward_mean") + _set(p, rec, "is_clip_frac", (t.get("mismatch") or {}).get("tis_clipfrac"), "rl_round_trained.mismatch.tis_clipfrac") + if mode == "async": + _set(p, rec, "dropped_samples", drops.get(rid), "rl_rollout_carry_over.over_age_discarded_groups (groups)") + f = fw.get(rid + framework_round_offset, {}) + for name in ROUND_FIELDS: # what the tape lacks, from the framework's own log of the same run + if rec[name] is None and f.get(name) is not None and name not in ("round", "t_end_unix", "round_s") \ + and name not in skip: + _set(p, rec, name, f[name], "framework log: " + framework_parsed.sources[name]) + p.rounds.append(rec) + p.segments_raw.append({k: round(v, 3) for k, v in seg.items()}) + _round_s_from_ends(p) + na = {"sampler_update_s"} # single-island arms: no separate sampler island + if mode == "sync": + na |= {"wait_s", "dropped_samples", "generated_samples", "version_lag"} + return _finish(p, na) + + +ROLLOUTER_ACTIVE = "fully_async/rollouter/active_time" + + +def _rollouter_active(checks: list[dict[str, Any]]) -> dict[int, float]: + """rl_publication_check events -> rollout id (policy version - 1) -> rollouter active s.""" + out: dict[int, float] = {} + for e in checks: + value = (e.get("rollouter_timing") or {}).get(ROLLOUTER_ACTIVE) + version = e.get("policy_version") + if isinstance(value, (int, float)) and version is not None and int(version) >= 1: + out[int(version) - 1] = float(value) + return out + + +def _driver_phase_durations(markers: list[dict[str, Any]]) -> dict[int, dict[str, float]]: + """rl_driver_phase markers (one per phase start) -> rollout id -> phase -> seconds + (marker to the next marker; a marker without rollout_id belongs to the round in progress).""" + out: dict[int, dict[str, float]] = defaultdict(lambda: defaultdict(float)) + rows = sorted(markers, key=lambda e: float(e.get("time_unix") or 0)) + current = 0 + for a, b in zip(rows, rows[1:]): + if a.get("rollout_id") is not None: + current = int(a["rollout_id"]) + out[current][str(a.get("phase"))] += float(b["time_unix"]) - float(a["time_unix"]) + return out + + +# ====================================================================== NVML +def parse_nvml(text: str) -> list[dict[str, Any]]: + """Rows {t_unix, key (container:index), name, util, mem_mb, mem_total_mb, power_w}; header lines skipped.""" + rows = [] + for parts in csv.reader(io.StringIO(text)): + parts = [x.strip() for x in parts] + if not parts or parts[0].startswith("timestamp") or "timestamp" in parts: + continue + container = "" + if not re.match(r"\d{4}/\d\d/\d\d", parts[0]): + container, parts = parts[0], parts[1:] + try: + t = _dt.datetime.strptime(parts[0], "%Y/%m/%d %H:%M:%S.%f").replace(tzinfo=_dt.timezone.utc).timestamp() + except (ValueError, IndexError): + continue + num = [re.sub(r"[^\d.]", "", x) for x in parts] + if len(parts) >= 7: # timestamp,index,name,util,mem.used,mem.total,power + index, util, mem, total, power = parts[1], num[3], num[4], num[5], num[6] + else: # timestamp,name,util,mem.used,power (T-EVAL format, no index) + index, util, mem, total, power = "0", num[2], num[3], "", num[4] if len(num) > 4 else "" + name = parts[2] if len(parts) >= 7 else parts[1] + rows.append({"t_unix": t, "key": f"{container}:{index}", "name": name, "util": float(util) if util else None, + "mem_mb": float(mem) if mem else None, "mem_total_mb": float(total) if total else None, + "power_w": float(power) if power else None}) + return rows + + +# ====================================================================== run directory +def _read(path: Path) -> str: + return path.read_text(errors="replace") if path.is_file() else "" + + +NOT_RUN_TAPES = ("yeto", "store") + + +def _tape_events(raw: Path) -> list[dict[str, Any]]: + seen, rows = set(), [] + for path in sorted(raw.rglob("*.jsonl")): + # miles_local.sh pulls a copy of the repo (raw/yeto, whose test fixtures are + # tapes of other runs) and the checkpoint store next to the run's own tapes + if "verl-" in path.name or path.relative_to(raw).parts[0] in NOT_RUN_TAPES: + continue + for row in tape_rows(_read(path)): + key = json.dumps(row, sort_keys=True, default=str) # the same event pulled from two sources + if key not in seen: + seen.add(key) + rows.append(row) + return rows + + +def parse_run(run_dir: Path) -> dict[str, Any]: + """Parse a run directory (``arm.json`` + ``raw/``) into ``parsed.json`` content.""" + run_dir = Path(run_dir) + plan = json.loads((run_dir / "arm.json").read_text()) + raw = run_dir / "raw" + mode, spec, c = plan["mode"], plan["spec"], plan["cards"] + arm_dir = next(iter(sorted(raw.glob("*/arm"))), raw / "arm") + nvml_text = _read(arm_dir / "nvml.csv") or _read(raw / "nvml.csv") + meta = json.loads(_read(arm_dir / "meta.json") or "{}") + if plan["framework"] == "verl": + metrics = next((p for p in sorted(raw.rglob("verl-metrics.jsonl"))), None) \ + or next((p for p in sorted(raw.rglob("verl-file-logger-*.jsonl"))), None) + fw = parse_verl(_read(metrics) if metrics else "", _read(arm_dir / "main.log"), + samples_per_round=spec["groups_per_round"] * spec["samples_per_group"], + trainer_cards=c["trainer"], mode=mode) + else: + log = _read(arm_dir / "main.log") or _read(raw / "launch.ts.log") + fw = parse_miles(log, rollout_cards=c["rollout"], mode=mode) + # yeto+verl fully_async: the driver's generate phase (and verl's timing_s/gen) only + # pull finished samples from the rollouter queue; generation time is the + # rollouter's own active time, carried by the tape's publication checks + verl_async = plan["framework"] == "verl" and mode == "async" + skip = frozenset({"gen_s"}) if verl_async else frozenset() + parsed = parse_yeto(_tape_events(raw), mode=mode, framework_parsed=fw, skip=skip, + framework_round_offset=1 if plan["framework"] == "verl" else 0, + rollouter_gen=verl_async) \ + if plan["system"] == "yeto" else fw + if meta.get("rc") is None and _read(raw / "rc").strip().lstrip("-").isdigit(): + meta["rc"] = int(_read(raw / "rc")) # miles_local.sh: the launcher's exit code + start = meta.get("start_unix") or _utc_file(raw / "start_utc.txt") + end = meta.get("end_unix") or _utc_file(raw / "end_utc.txt") + return {"arm": plan["arm"], "run_id": plan["run_id"], "fingerprint": plan["fingerprint"], "mode": mode, + "framework": plan["framework"], "system": plan["system"], "cards": c, "gpu": plan["gpu"], + "spec": spec, "container_start_unix": start, "container_end_unix": end, "rc": meta.get("rc"), + "framework_log": fw.to_dict() if plan["system"] == "yeto" else None, + "nvml": (nvml := parse_nvml(nvml_text)), + # E1: Modal "A10G" containers are NVIDIA A10 or (sometimes) A10G, which train slower + "gpu_names": sorted({r["name"] for r in nvml if r.get("name")}), + "nvml_interval_s": 20.0 if plan["harness"] == "miles_local" else 1.0, + **parsed.to_dict()} + + +def _utc_file(path: Path) -> float | None: + text = _read(path).strip() + if not text: + return None + return _dt.datetime.strptime(text, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=_dt.timezone.utc).timestamp() diff --git a/tools/eval/efficiency/prices.json b/tools/eval/efficiency/prices.json new file mode 100644 index 00000000..75edc4ee --- /dev/null +++ b/tools/eval/efficiency/prices.json @@ -0,0 +1,6 @@ +{ + "_note": "USD per card-hour, all-in container estimate (GPU + the CPU and memory the harness requests). Modal list prices read 2026-10-20 for the GPU part; CPU/memory added as in tools/eval/cells.py. Estimates, not a bill: replace with the ledger's rate (infra-drafts/gpu-spend.md) when it differs.", + "A10G": {"usd_per_card_hour": 2.0, "gpu_list_usd_per_hour": 1.10}, + "L4": {"usd_per_card_hour": 1.6, "gpu_list_usd_per_hour": 0.80}, + "H100!": {"usd_per_card_hour": 5.3, "gpu_list_usd_per_hour": 3.95} +} diff --git a/tools/eval/efficiency/report.py b/tools/eval/efficiency/report.py new file mode 100644 index 00000000..92e46927 --- /dev/null +++ b/tools/eval/efficiency/report.py @@ -0,0 +1,213 @@ +"""Markdown report of a set of efficiency runs (``parsed.json`` in each run dir). + +Sections: the shared spec and its fingerprint; one row per run; per-arm +aggregates (median and min-max over repeats); yeto vs bare per framework and +mode, metric by metric, with a range verdict; yeto's own segments; reward +curves; and the fields each arm could not measure. + +Range verdict (repeats are not bit-reproducible on H100, so no single-run +comparison): with >= :data:`MIN_REPEATS` runs on both sides, yeto is "worse" +or "better" only when its min-max range and bare's do not overlap; otherwise +"within run-to-run range". Fewer runs, or a missing value, is said as such. +""" + +from __future__ import annotations + +import json +import statistics +from collections import defaultdict +from pathlib import Path +from typing import Any + +from .metrics import summarize +from .spec import ARMS, FRAMEWORKS, MODES, EffSpec, verl_knob_diff + +MIN_REPEATS = 2 +# metric -> (label, True when lower is better) +COMPARED = { + "round_s_median": ("round time (s)", True), + "gen_s_median": ("generation (s/round)", True), + "train_s_median": ("training (s/round)", True), + "weight_sync_s_median": ("weight sync trainer->rollout (s/round)", True), + "sampler_update_s_median": ("update sampler weights (s/round)", True), + "wait_s_median": ("wait (s/round)", True), + "unaccounted_s_median": ("other, not in the segments above (s/round)", True), + "gen_tok_per_card_s": ("generated tokens / card / s", False), + "train_tok_per_card_s": ("trained tokens / card / s", False), + "samples_per_s": ("trained samples / s", False), + "gpu_util_mean": ("GPU utilization mean (%)", False), + "bubble_frac": ("idle bubble share (util < 5 %)", True), + "gpu_mem_peak_mb": ("GPU memory peak (MiB)", True), + "usd_per_1k_samples_steady": ("USD / 1k samples (steady)", True), + "usd_per_1k_samples_run": ("USD / 1k samples (whole container run)", True), +} + + +def load(run_dirs: list[Path]) -> list[dict[str, Any]]: + out = [] + for d in run_dirs: + path = Path(d) / "parsed.json" + if path.is_file(): + out.append(json.loads(path.read_text())) + return out + + +def _fmt(v: Any) -> str: + if v is None: + return "missing" + if isinstance(v, float): + return f"{v:.4g}" + return str(v) + + +def _range(values: list[float]) -> str: + return "missing" if not values else f"{_fmt(statistics.median(values))} [{_fmt(min(values))}, {_fmt(max(values))}]" + + +def verdict(bare: list[float | None], yeto: list[float | None], lower_is_better: bool) -> str: + if any(v is None for v in bare + yeto) or not bare or not yeto: + return "missing value" + if len(bare) < MIN_REPEATS or len(yeto) < MIN_REPEATS: + return f"insufficient repeats (bare {len(bare)}, yeto {len(yeto)}; need {MIN_REPEATS})" + if max(yeto) < min(bare): + return "yeto better" if lower_is_better else "yeto worse" + if min(yeto) > max(bare): + return "yeto worse" if lower_is_better else "yeto better" + return "within run-to-run range" + + +def _delta(bare: list[float], yeto: list[float]) -> str: + if not bare or not yeto or any(v is None for v in bare + yeto): + return "missing" + b, y = statistics.median(bare), statistics.median(yeto) + if b == 0: + return f"{y - b:+.4g} (abs)" + return f"{y - b:+.4g} ({(y - b) / abs(b) * 100:+.1f} %)" + + +def markdown(parsed_runs: list[dict[str, Any]], *, title: str) -> str: + if not parsed_runs: + return f"# {title}\n\nNo parsed runs.\n" + prints = {p["fingerprint"] for p in parsed_runs} + if len(prints) > 1: + raise ValueError(f"runs come from different specs (fingerprints {sorted(prints)}); report them separately") + models = {tuple(p.get("gpu_names") or ()) for p in parsed_runs if p.get("gpu_names")} + if len(models) > 1: + # E1: e1a9 ran on an NVIDIA A10G, the others on an A10; its train step was ~12 % slower + raise ValueError(f"runs ran on different GPU models {sorted(models)}; report them separately") + spec = parsed_runs[0]["spec"] + summaries = [summarize(p) for p in parsed_runs] + by_arm: dict[str, list[dict]] = defaultdict(list) + for s in summaries: + by_arm[s["arm"]].append(s) + lines = [f"# {title}", "", + f"Spec `{spec['name']}`, fingerprint `{parsed_runs[0]['fingerprint'][:16]}`: model `{spec['model']}`@" + f"`{spec['model_revision'][:10]}`, data `{spec['data']}`@`{spec['data_revision'][:10]}`, " + f"{spec['rounds']} rounds x {spec['groups_per_round']} prompts x {spec['samples_per_group']} samples, " + f"max_tokens {spec['max_response_len']}, lr {spec['lr']}, LoRA r {spec['lora_r']}, seed {spec['seed']}, " + f"{spec['gpu']} (sync 1 card colocated, async 1 rollout + 1 trainer card), " + f"max_policy_age {spec['max_policy_age']} (async), TIS upper {spec['tis_upper']}.", "", + "Steady window = rounds after the first. Throughput counts every card of the arm. " + "Costs are list-price estimates (`prices.json`), not a bill.", ""] + lines += _runs_table(summaries) + _arms_table(by_arm) + _overhead(by_arm) + _yeto_segments(by_arm) + lines += _knobs(spec) + _rewards(summaries) + _missing(by_arm) + return "\n".join(lines) + "\n" + + +def _runs_table(summaries: list[dict]) -> list[str]: + head = ("arm", "run", "rc", "rounds", "steady wall s", "gen tok/card/s", "train tok/card/s", "round s med", + "util %", "bubble", "mem peak MiB", "$/1k steady", "$/1k run", "reward last3") + rows = ["## Runs", "", "| " + " | ".join(head) + " |", "|" + "---|" * len(head)] + for s in sorted(summaries, key=lambda s: (ARMS.index(s["arm"]), s["run_id"])): + vals = (s["arm"], f"`{s['run_id']}`", s["rc"], f"{s['rounds']}/{s['rounds_expected']}", s["steady_wall_s"], + s["gen_tok_per_card_s"], s["train_tok_per_card_s"], s["round_s_median"], s["gpu_util_mean"], + s["bubble_frac"], s["gpu_mem_peak_mb"], s["usd_per_1k_samples_steady"], s["usd_per_1k_samples_run"], + s["reward_last3_mean"]) + rows.append("| " + " | ".join(_fmt(v) for v in vals) + " |") + return rows + [""] + + +def _arms_table(by_arm: dict[str, list[dict]]) -> list[str]: + rows = ["## Per arm (median [min, max] over repeats)", "", + "| metric | " + " | ".join(a for a in ARMS if a in by_arm) + " |", "|---|" + "---|" * len(by_arm)] + for key, (label, _) in COMPARED.items(): + cells = [] + for a in ARMS: + if a in by_arm: + vals = [s[key] for s in by_arm[a]] + cells.append(_range([v for v in vals if v is not None]) if any(v is not None for v in vals) + else ("n/a" if all(s["sources"].get(key.removesuffix("_median")) == "n/a" + for s in by_arm[a]) else "missing")) + rows.append(f"| {label} | " + " | ".join(cells) + " |") + return rows + [""] + + +def _overhead(by_arm: dict[str, list[dict]]) -> list[str]: + rows = ["## yeto vs bare (same framework and mode)", "", + "Delta = median(yeto) - median(bare); positive time = yeto slower.", ""] + for fw in FRAMEWORKS: + for mode in MODES: + bare, yeto = by_arm.get(f"bare-{fw}-{mode}", []), by_arm.get(f"yeto-{fw}-{mode}", []) + if not bare and not yeto: + continue + rows += [f"### {fw} {mode} (bare n={len(bare)}, yeto n={len(yeto)})", "", + "| metric | bare | yeto | delta | verdict |", "|---|---|---|---|---|"] + for key, (label, lower) in COMPARED.items(): + bv, yv = [s[key] for s in bare], [s[key] for s in yeto] + if all(s["sources"].get(key.removesuffix("_median")) == "n/a" for s in bare + yeto): + rows.append(f"| {label} | n/a | n/a | n/a | not applicable |") + continue + rows.append(f"| {label} | {_range([v for v in bv if v is not None])} | " + f"{_range([v for v in yv if v is not None])} | {_delta(bv, yv)} | " + f"{verdict(bv, yv, lower)} |") + rows.append("") + return rows + + +def _yeto_segments(by_arm: dict[str, list[dict]]) -> list[str]: + rows = ["## Native segments per round (median, steady window)", "", + "yeto: tape task/phase names; bare: the framework's own timers. Where yeto adds work " + "(offload/onload, outer sync, publish), it shows here.", ""] + for a in ARMS: + for s in by_arm.get(a, []): + raw = s.get("segments_raw_median") or {} + rows.append(f"- {a} `{s['run_id']}`: " + (", ".join(f"{k} {_fmt(v)}" for k, v in raw.items()) + or "missing")) + return rows + [""] + + +def _knobs(spec: dict) -> list[str]: + rows = ["## verl overrides that differ between bare and yeto (same spec)", ""] + try: + for mode in MODES: + diff = verl_knob_diff(EffSpec(**spec), mode) + rows.append(f"- {mode}: " + ("; ".join(f"`{k}` bare={v['bare']} yeto={v['yeto']}" + for k, v in diff.items()) or "none")) + except ImportError as exc: + rows.append(f"- not computed ({exc})") + rows += ["- Miles: the yeto launcher builds the Miles argv at run time; compare the launch log's argv with " + "the bare script (`arm.json`) by hand (not automated in E0.1).", ""] + return rows + + +def _rewards(summaries: list[dict]) -> list[str]: + rows = ["## Reward against wall time and samples", "", "Per run: round: t_s / samples / reward_mean.", ""] + for s in sorted(summaries, key=lambda s: (ARMS.index(s["arm"]), s["run_id"])): + pts = " ".join(f"{c['round']}:{_fmt(c['t_s'])}/{_fmt(c['samples'])}/{_fmt(c['reward_mean'])}" + for c in s["reward_curve"]) + rows.append(f"- {s['arm']} `{s['run_id']}`: {pts or 'missing'}") + return rows + [""] + + +def _missing(by_arm: dict[str, list[dict]]) -> list[str]: + rows = ["## Field sources and gaps", "", "Per arm (first run): each round field's source; `missing` = no " + "source in the raw data, `n/a` = does not apply.", ""] + for a in ARMS: + if a in by_arm: + src = by_arm[a][0]["sources"] + rows.append(f"- {a}: " + "; ".join(f"{k}: {v}" for k, v in src.items())) + partial = {k: v for s in by_arm[a] for k, v in (s.get("rate_rounds") or {}).items()} + if partial: + rows.append(f" - {a}: rates over the steady rounds that report the value: " + + "; ".join(f"{k} {v}" for k, v in partial.items())) + return rows + [""] diff --git a/tools/eval/efficiency/spec.py b/tools/eval/efficiency/spec.py new file mode 100644 index 00000000..6d1a932b --- /dev/null +++ b/tools/eval/efficiency/spec.py @@ -0,0 +1,698 @@ +"""One shared spec per model preset, and the command plan of every arm. + +An arm is ``--``: + +* system ``bare`` (the framework's own entry point, no yeto code on its path) + or ``yeto`` (the same framework driven by yeto); +* framework ``verl`` (fork pin ``VERL_COMMIT``, the verl image yeto builds) or + ``miles`` (``MILES_NEXT_IMAGE``, the image yeto runs Miles in); +* mode ``sync`` (one card, rollout and trainer colocated) or ``async`` + (two cards, one rollout + one trainer, fully_async with a policy-age limit). + +Every value the comparison depends on comes from :class:`EffSpec`; +:func:`control_fingerprint` hashes them and every arm records it. + +Sources of the bare command lines (the "official recipe" each one follows): + +* bare verl sync: ``examples/grpo_trainer/run_qwen3_4b_fsdp.sh`` (fork + acad9875), ``python -m verl.trainer.main_ppo`` with its override groups; +* bare verl async: ``verl/experimental/fully_async_policy/shell/ + dapo_7b_math_fsdp2_4_4.sh``, ``python -m verl.experimental.fully_async_policy.fully_async_main``; +* bare Miles sync: ``scripts/run_qwen3_4b.py`` (Megatron branch, bridge mode) + -> ``train.py``; model flags from the image's own ``scripts/models/``; +* bare Miles async: ``examples/infra_features/fully_async/`` (``--fully-async``) + -> ``train_async.py``. + +Values the spec does not name stay at the recipe's value, or at the +framework default where the recipe sets none. Where yeto sets a value the +recipe also sets (LoRA, TIS, dtype), the bare arm takes the spec value, so the +two arms differ only by what yeto adds; :func:`verl_knob_diff` lists the verl +overrides that still differ. +""" + +from __future__ import annotations + +import hashlib +import json +import shlex +from dataclasses import asdict, dataclass, field, replace +from pathlib import Path +from typing import Any + +SYSTEMS = ("bare", "yeto") +FRAMEWORKS = ("verl", "miles") +MODES = ("sync", "async") +ARMS = tuple(f"{s}-{f}-{m}" for f in FRAMEWORKS for m in MODES for s in SYSTEMS) + +EFF_DIR = "/root/eff" # spec, reward and prep script inside a bare container +OUT_DIR = "/root/eff-out" # everything a container writes for the results Volume +DATA_DIR = "/root/eff-data" # eff_prep.py output +MODEL_DIR = "/root/eff-models" +VERL_ROOT = "/workspace/verl" +MILES_ROOT = "/root/miles" +MEGATRON_ROOT = "/root/Megatron-LM" # yeto.launch.remote_setup.PORTS_MEGATRON_PATH + + +@dataclass(frozen=True) +class EffSpec: + """Everything the arms of one comparison share.""" + + name: str + model: str + model_revision: str + miles_model_args: str # file under /scripts/models/ (the official model flags) + data: str = "zhuzilin/gsm8k" + data_revision: str = "0cbd9f31d91ac21a7613dcbc7fef992adac459ae" + data_file: str = "" # one JSONL file of the dataset repo (prompt, label); "" = train/test.parquet + train_rows: int = 0 # keep the first rows only (0 = all); needed with data_file + gpu: str = "A10G" # Modal GPU type, the same for every arm + rounds: int = 8 # optimizer steps (= policy versions) per run + groups_per_round: int = 8 # prompts per step + samples_per_group: int = 4 # responses per prompt + max_prompt_len: int = 512 + max_response_len: int = 256 # "max_tokens" of the rollout + lr: float = 1e-5 + lora_r: int = 32 # yeto's verl backend is LoRA only, so every arm is LoRA + seed: int = 17 + temperature: float = 1.0 + top_p: float = 1.0 + mem_fraction: float = 0.5 # rollout engine static memory fraction + max_policy_age: int = 1 # async arms: versions a sample may lag + tis_upper: float = 2.0 # token-level truncated IS upper bound (all arms) + repeats: int = 2 # runs per arm (range verdicts need >= 2) + timeout_s: int = 3000 + chat_template_kwargs: dict[str, Any] = field(default_factory=dict) + # yeto arms only (the bare arms keep no checkpoint): "final" cuts once at the end + # (fair control), "every-round" is the product default (a round cut every round). + # Not in the control fingerprint: it does not change what the bare arms run. + checkpoint: str = "final" + # cards per run: 0 keeps the E0 shape (sync 1, async 1 rollout + 1 trainer); E11 + # (whole-cluster comparison) sets the total, split evenly between rollout and trainer in async + cards: int = 0 + # "lora" or "full" (every parameter; E11 0.6B). Miles full-parameter islands with + # engines need DP=1, so a 2-card trainer runs tensor parallel 2 in every arm. + tuning: str = "lora" + tensor_parallel: int = 1 + + @property + def samples_per_round(self) -> int: + return self.groups_per_round * self.samples_per_group + + @property + def seq_len(self) -> int: + return self.max_prompt_len + self.max_response_len + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + +PRESETS: dict[str, EffSpec] = { + # E0.1 acceptance (debug size): the T-EVAL model and data on the T-EVAL card + "qwen3-0.6b": EffSpec( + name="qwen3-0.6b", model="Qwen/Qwen3-0.6B", model_revision="c1899de289a04d12100db370d81485cdf75e47ca", + miles_model_args="qwen3-0.6B.py", chat_template_kwargs={"enable_thinking": False}), + # E0.2 (main comparison): revision read from the HF API 2026-10-20 + "qwen3-4b-instruct-2507": EffSpec( + name="qwen3-4b-instruct-2507", model="Qwen/Qwen3-4B-Instruct-2507", + model_revision="cdbee75f17c01a7cc42f958dc650907174af0554", miles_model_args="qwen3-4B-Instruct-2507.py", + gpu="H100!", rounds=20, groups_per_round=32, samples_per_group=8, max_prompt_len=512, + max_response_len=1024, lora_r=64, mem_fraction=0.6, timeout_s=5400), + # E11 (whole cluster): the E0.2 model and hyperparameters on 4 cards; async (2 rollout + + # 2 trainer, or 2 islands of 1 + 1), so the 2-island yeto arm and the single-island arms + # share one shape. Data: DAPO-Math-17k (gsm8k is saturated for this model: reward 0.97 + # in round 1 of e02-bare-verl-sync-1), the first rounds x groups rows (one epoch). + "e11-qwen3-4b": EffSpec( + name="e11-qwen3-4b", model="Qwen/Qwen3-4B-Instruct-2507", + model_revision="cdbee75f17c01a7cc42f958dc650907174af0554", miles_model_args="qwen3-4B-Instruct-2507.py", + data="zhuzilin/dapo-math-17k", data_revision="2e65612930298bde4c5d58fd97b3f23a483aaff9", + data_file="dapo-math-17k.jsonl", train_rows=12 * 32, + gpu="H100!", rounds=12, groups_per_round=32, samples_per_group=8, max_prompt_len=512, + max_response_len=1024, lora_r=64, mem_fraction=0.6, timeout_s=5400, cards=4), + # E11 full parameter: the E11 data, rounds and lengths on Qwen3-0.6B, every parameter + # trained (lora_r is unused). 2 rollout + 2 trainer cards, trainer TP 2 (DP 1). + "e11-qwen3-0.6b-full": EffSpec( + name="e11-qwen3-0.6b-full", model="Qwen/Qwen3-0.6B", model_revision="c1899de289a04d12100db370d81485cdf75e47ca", + miles_model_args="qwen3-0.6B.py", + data="zhuzilin/dapo-math-17k", data_revision="2e65612930298bde4c5d58fd97b3f23a483aaff9", + data_file="dapo-math-17k.jsonl", train_rows=12 * 32, + gpu="H100!", rounds=12, groups_per_round=32, samples_per_group=8, max_prompt_len=512, + max_response_len=1024, mem_fraction=0.6, timeout_s=3600, cards=4, tuning="full", + tensor_parallel=2, chat_template_kwargs={"enable_thinking": False}), +} + + +@dataclass(frozen=True) +class Arm: + system: str + framework: str + mode: str + + @classmethod + def parse(cls, name: str) -> "Arm": + parts = name.split("-") + if len(parts) != 3 or name not in ARMS: + raise ValueError(f"unknown arm {name!r}; known: {', '.join(ARMS)}") + return cls(*parts) + + @property + def name(self) -> str: + return f"{self.system}-{self.framework}-{self.mode}" + + @property + def peer(self) -> str: + """The arm this one is compared with (same framework and mode).""" + return f"{'yeto' if self.system == 'bare' else 'bare'}-{self.framework}-{self.mode}" + + +def get_spec(preset: str, **overrides: Any) -> EffSpec: + if preset not in PRESETS: + raise ValueError(f"unknown preset {preset!r}; known: {', '.join(PRESETS)}") + return replace(PRESETS[preset], **overrides) if overrides else PRESETS[preset] + + +def cards(mode: str, total: int = 0) -> dict[str, int]: + """Cards per role: sync colocates both on every card; async splits them evenly + (``total`` 0: one card for sync, one of each role for async).""" + if mode == "sync": + n = total or 1 + return {"total": n, "rollout": n, "trainer": n, "colocated": True} + n = total or 2 + if n % 2: + raise ValueError(f"async needs an even card count (rollout + trainer halves), got {n}") + return {"total": n, "rollout": n // 2, "trainer": n // 2, "colocated": False} + + +def control_fingerprint(spec: EffSpec) -> str: + """sha256 over every shared value except the repeat count (a report refuses + to compare runs whose fingerprints differ).""" + doc = {k: v for k, v in spec.to_dict().items() if k not in ("repeats", "timeout_s", "checkpoint")} + # fields added after E0.1: a default value keeps the fingerprints of earlier runs valid + for key, default in (("cards", 0), ("data_file", ""), ("train_rows", 0), ("tuning", "lora"), + ("tensor_parallel", 1)): + if doc[key] == default: + doc.pop(key) + return hashlib.sha256(json.dumps(doc, sort_keys=True).encode()).hexdigest() + + +def images() -> dict[str, str]: + """The image pins every arm of a framework uses (same image for bare and yeto).""" + from yeto.rl.adapters.miles.pins import MILES_NEXT_IMAGE + from yeto.rl.adapters.verl.pins import VERL_COMMIT + + return {"verl": f"verl-build:{VERL_COMMIT}", "miles": MILES_NEXT_IMAGE} + + +def _model_path(spec: EffSpec) -> str: + return f"{MODEL_DIR}/{spec.model.split('/')[-1]}" + + +# ====================================================================== bare verl +def verl_bare_overrides(spec: EffSpec, mode: str, run_id: str) -> list[str]: + """Hydra overrides of the bare verl arm (recipe groups in the recipe's order).""" + c = cards(mode, spec.cards) + data = [ + "algorithm.adv_estimator=grpo", + "algorithm.use_kl_in_reward=False", + "algorithm.norm_adv_by_std_in_grpo=True", + f"data.train_files={DATA_DIR}/train.parquet", + f"data.val_files={DATA_DIR}/test.parquet", + f"data.max_prompt_length={spec.max_prompt_len}", + f"data.max_response_length={spec.max_response_len}", + "data.filter_overlong_prompts=True", + "data.truncation=error", + f"data.seed={spec.seed}", + *[f"+data.apply_chat_template_kwargs.{k}={v}" for k, v in sorted(spec.chat_template_kwargs.items())], + ] + model = [ + f"actor_rollout_ref.model.path={_model_path(spec)}", + "actor_rollout_ref.model.use_remove_padding=True", + "actor_rollout_ref.model.enable_gradient_checkpointing=True", + f"actor_rollout_ref.model.lora_rank={spec.lora_r}", + f"actor_rollout_ref.model.lora_alpha={spec.lora_r}", + "actor_rollout_ref.model.target_modules=all-linear", + ] + actor = [ + "actor_rollout_ref.actor.strategy=fsdp2", + "actor_rollout_ref.actor.fsdp_config.model_dtype=fp32", # yeto's VerlRunConfig default + f"actor_rollout_ref.actor.optim.lr={spec.lr!r}", + f"actor_rollout_ref.actor.ppo_mini_batch_size={spec.groups_per_round}", + "actor_rollout_ref.actor.use_dynamic_bsz=True", + "actor_rollout_ref.actor.ppo_max_token_len_per_gpu=16384", + "actor_rollout_ref.actor.use_kl_loss=False", + "actor_rollout_ref.actor.entropy_coeff=0", + "actor_rollout_ref.actor.fsdp_config.param_offload=False", + "actor_rollout_ref.actor.fsdp_config.optimizer_offload=False", + f"actor_rollout_ref.actor.data_loader_seed={spec.seed}", + ] + rollout = [ + "actor_rollout_ref.rollout.name=vllm", + "actor_rollout_ref.rollout.tensor_model_parallel_size=1", + f"actor_rollout_ref.rollout.gpu_memory_utilization={spec.mem_fraction}", + f"actor_rollout_ref.rollout.n={spec.samples_per_group}", + f"actor_rollout_ref.rollout.temperature={spec.temperature}", + f"actor_rollout_ref.rollout.top_p={spec.top_p}", + "actor_rollout_ref.rollout.top_k=-1", + "actor_rollout_ref.rollout.calculate_log_probs=True", + "actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True", + "actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=16384", + "algorithm.rollout_correction.rollout_is=token", + f"algorithm.rollout_correction.rollout_is_threshold={spec.tis_upper}", + "algorithm.rollout_correction.bypass_mode=False", + ] + reward = [f"reward.custom_reward_function.path={EFF_DIR}/eff_reward.py", + "reward.custom_reward_function.name=compute_score"] + trainer = [ + "trainer.critic_warmup=0", + "trainer.logger=[console,file]", + "trainer.project_name=yeto-effeval", + f"trainer.experiment_name={run_id}", + f"trainer.default_local_dir={OUT_DIR}/ckpt", + "trainer.nnodes=1", + "trainer.val_before_train=False", + "trainer.test_freq=-1", + "trainer.save_freq=-1", + "trainer.resume_mode=disable", + f"trainer.total_training_steps={spec.rounds}", + "trainer.total_epochs=1000", + "ray_kwargs.ray_init.runtime_env.py_executable=null", + "ray_kwargs.ray_init.num_cpus=32", # a Modal container reports too few Ray CPUs (S17 V2) + ] + if mode == "sync": + return [*data, f"data.train_batch_size={spec.groups_per_round}", *model, *actor, *rollout, *reward, + *trainer, f"trainer.n_gpus_per_node={c['total']}"] + fully_async = [ # dapo_7b_math_fsdp2_4_4.sh; one parameter version per spec round + "data.train_batch_size=0", + "data.gen_batch_size=1", + "data.return_raw_chat=True", + "actor_rollout_ref.hybrid_engine=False", + "actor_rollout_ref.rollout.mode=async", + "actor_rollout_ref.rollout.checkpoint_engine.backend=nccl", + # The fork's disaggregated sender drops the PEFT config of unmerged LoRA params + # (engine_workers.update_weights), so the rollout would load them as base + # weights; merged LoRA sends full weights, the fork's working LoRA path. + "actor_rollout_ref.model.lora.merge=True", + "rollout.nnodes=1", + f"rollout.n_gpus_per_node={c['rollout']}", + f"rollout.total_rollout_steps={spec.groups_per_round * spec.rounds}", + f"trainer.n_gpus_per_node={c['trainer']}", + f"async_training.staleness_threshold={spec.max_policy_age}", + "async_training.trigger_parameter_sync_step=1", + "async_training.require_batches=1", + "async_training.partial_rollout=True", + "async_training.use_trainer_do_validate=False", + "async_training.use_dynamic_resource_scheduling=False", + ] + return [*data, *model, *actor, *rollout, *reward, *trainer, *fully_async] + + +def verl_bare_entry(mode: str) -> str: + return "verl.trainer.main_ppo" if mode == "sync" else "verl.experimental.fully_async_policy.fully_async_main" + + +# ====================================================================== bare Miles +def miles_bare_argv(spec: EffSpec, mode: str) -> list[str]: + """Miles arguments after the model flags (the container prepends ``$MODEL_ARGS``).""" + c = cards(mode, spec.cards) + model = _model_path(spec) + argv = [ + # checkpoint: bridge mode loads the HF checkpoint directly (run_qwen3_4b.py --megatron-to-hf-mode bridge) + "--hf-checkpoint", model, "--ref-load", model, "--megatron-to-hf-mode", "bridge", + "--finetune", "--no-load-optim", "--no-load-rng", "--no-save-optim", "--no-save-rng", + # rollout + "--prompt-data", f"{DATA_DIR}/train.jsonl", "--input-key", "prompt", "--label-key", "label", + # no --rollout-shuffle (the recipe has it): yeto's Miles argv leaves rollout_shuffle + # False (eff01-yeto-miles-* Miles arg dump), and both arms must read prompts in one order + "--apply-chat-template", "--rollout-seed", str(spec.seed), + "--custom-rm-path", "eff_reward.score", + "--num-rollout", str(spec.rounds), "--rollout-batch-size", str(spec.groups_per_round), + "--n-samples-per-prompt", str(spec.samples_per_group), + "--rollout-max-response-len", str(spec.max_response_len), + "--rollout-max-context-len", str(spec.seq_len), + "--rollout-temperature", str(spec.temperature), "--rollout-top-p", str(spec.top_p), + "--global-batch-size", str(spec.samples_per_round), "--balance-data", + # GRPO + TIS (the correction every arm uses) + "--advantage-estimator", "grpo", "--entropy-coef", "0.00", "--eps-clip", "0.2", "--eps-clip-high", "0.28", + "--use-tis", "--tis-clip", str(spec.tis_upper), + # optimizer + "--optimizer", "adam", "--lr", str(spec.lr), "--lr-decay-style", "constant", + "--seed", str(spec.seed), + *_miles_lora_flags(spec), + # Megatron (run_qwen3_4b.py, Megatron branch) + "--train-backend", "megatron", "--seq-length", str(spec.seq_len), + "--tensor-model-parallel-size", str(spec.tensor_parallel), "--pipeline-model-parallel-size", "1", + "--context-parallel-size", "1", "--expert-model-parallel-size", "1", "--expert-tensor-parallel-size", "1", + "--recompute-granularity", "full", "--recompute-method", "uniform", "--recompute-num-layers", "1", + "--use-dynamic-batch-size", "--max-tokens-per-gpu", str(max_tokens_per_gpu(spec)), + "--attention-dropout", "0.0", "--hidden-dropout", "0.0", + "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", + # yeto's generic recipe (Qwen3) trains with "unfused": TE's flash path picks the + # FlashAttention cute kernel, which fails to compile on A10G (sm86, eff01-bare-miles-sync-a). + "--attention-backend", "unfused", + "--bf16", + # SGLang + "--rollout-num-gpus-per-engine", "1", "--sglang-mem-fraction-static", str(spec.mem_fraction), + # placement + "--actor-num-nodes", "1", "--actor-num-gpus-per-node", str(c["trainer"]), + "--num-gpus-per-node", str(c["total"]), + ] + if spec.chat_template_kwargs: + argv += ["--apply-chat-template-kwargs", json.dumps(spec.chat_template_kwargs, separators=(",", ":"))] + if mode == "sync": + argv += ["--colocate", "--offload-train"] + if spec.tuning == "lora": + argv += ["--lora-base-cpu-backup"] + else: + argv += ["--rollout-num-gpus", str(c["rollout"]), "--fully-async", + "--max-weight-staleness", str(spec.max_policy_age)] + return argv + + +def _miles_lora_flags(spec: EffSpec) -> list[str]: + """Miles' own LoRA flags (the same set yeto's Miles adapter emits); none for full parameter.""" + if spec.tuning == "full": + return [] + return ["--lora-rank", str(spec.lora_r), "--lora-alpha", str(spec.lora_r), "--lora-dropout", "0", + "--lora-type", "canonical_lora", "--target-modules", "all-linear", + "--sglang-max-lora-rank", str(spec.lora_r)] + + +def _yeto_tuning(spec: EffSpec) -> str: + """The yeto launcher's tuning options (full parameter: a no-sync single island).""" + if spec.tuning == "full": + return f"--tuning full --tensor-parallel {spec.tensor_parallel} " + return f"--tuning lora --lora-r {spec.lora_r} " + + +def miles_bare_script(mode: str) -> str: + return "train.py" if mode == "sync" else "train_async.py" + + +def max_tokens_per_gpu(spec: EffSpec) -> int: + """Miles dynamic-batching budget (run_qwen3_4b.py: 9216), never below one sample.""" + return max(9216, spec.seq_len) + + +# ====================================================================== yeto arms +YETO_WORKDIR = "/root/sky_workdir" # modal_verl.WORKDIR: the repo inside the yeto+verl container +ALGORITHM_SPEC = "examples/rl_algorithms/decoupled-ppo-c2.json" # token TIS 2.0, max_policy_staleness 2 +# the same with the bare Miles recipe's clip range (--eps-clip 0.2 --eps-clip-high 0.28) +MILES_ALGORITHM_SPEC = "tools/eval/efficiency/algo-eff-miles.json" +def yeto_verl_argv(spec: EffSpec, mode: str, tape: str) -> list[str]: + """``island_entry`` arguments of the yeto+verl arm (the T-EVAL E1 shape).""" + if uses_task_file(spec): + raise ValueError("the yeto+verl arm reads --data from the Hub; it has no task-file path") + argv = ["-m", "yeto.rl.adapters.verl.island_entry", "--model", spec.model, "--model-revision", + spec.model_revision, "--data", spec.data, "--data-revision", spec.data_revision, "--learner-id", "0", + "--reward-function", "gsm8k_reward:score", "--global-rounds", str(spec.rounds), + "--groups-per-round", str(spec.groups_per_round), "--samples-per-group", str(spec.samples_per_group), + "--rollout-max-response-len", str(spec.max_response_len), "--event-tape", tape, + "--inner-lr", str(spec.lr), "--lora-r", str(spec.lora_r), "--lora-targets", "all-linear", + "--seq-len", str(spec.seq_len), "--seed", str(spec.seed), "--rl-single-island-no-sync"] + if spec.chat_template_kwargs: + argv += ["--apply-chat-template-kwargs", json.dumps(spec.chat_template_kwargs)] + if mode == "async": + # The default algorithm spec allows no staleness (island_entry rejects age 1, + # eff01-yeto-verl-async-a); use the yeto+Miles arm's spec: token TIS 2.0, staleness 2. + argv += ["--rl-max-policy-age", str(spec.max_policy_age), + "--rl-algorithm-spec", f"{YETO_WORKDIR}/{ALGORITHM_SPEC}"] + return argv + + +CHECKPOINT_MODES = ("final", "every-round") + + +def cut_every(spec: EffSpec) -> int: + """Round-cut interval of the yeto+Miles arm: ``final`` cuts only at the end (the + bare arms keep no checkpoint), ``every-round`` is the product default. + The yeto+verl arm sets no checkpoint store and fully_async refuses cuts, so it + has no round cut in either mode.""" + if spec.checkpoint not in CHECKPOINT_MODES: + raise ValueError(f"checkpoint must be one of {CHECKPOINT_MODES}, got {spec.checkpoint!r}") + return spec.rounds if spec.checkpoint == "final" else 1 + + +# E1: the GPU model every arm must run on, per spec.gpu (as tools/eval/modal_verl.GPU_NAMES: +# Modal "A10G" containers are an NVIDIA A10 or an A10G, and the A10G trains slower) +GPU_MODELS = {"A10G": "NVIDIA A10", "L4": "L4", "H100!": "H100"} + + +def gpu_name_arg(spec: EffSpec) -> str: + """The launcher option that pins the yeto+Miles arm to the bare arms' GPU model.""" + return f"--modal-gpu-name {shlex.quote(GPU_MODELS[spec.gpu])} " + + +def yeto_miles_launcher_args(spec: EffSpec, mode: str, run_id: str) -> str: + """The yeto launcher arguments of the yeto+Miles arm (``tools/eval/miles_local.sh`` ARGS).""" + from ..miles import STORE_VOLUME, TAPE_VOLUME + + card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] + c = cards(mode, spec.cards) + kwargs = json.dumps(spec.chat_template_kwargs) if spec.chat_template_kwargs else "" + args = (f"launch --training-mode rl --rl-engine ports --preflight-threads off --modal-gpu-exact " + f"{gpu_name_arg(spec)}" + # container resources as the bare arms' (modal_eff: 8 CPU + 48 GiB per card) + f"--modal-retries 0 --modal-cpu {8 * c['total']} --modal-memory {48 * c['total']} " + f"--modal-tape-volume {TAPE_VOLUME} " + f"--controller local --cluster-prefix {run_id} {_yeto_tuning(spec)}" + f"--rl-observe-timeline --model {spec.model} --model-revision {spec.model_revision} {_yeto_data(spec)} " + f"--rl-checkpoint-store modal-volume://{STORE_VOLUME}/{run_id} " + f"--rl-cut-every {cut_every(spec)} " + "--reward-function gsm8k_reward:score --fragments 1 --pipeline 1 " + f"--rollout-batch-size {spec.groups_per_round} --n-samples-per-prompt {spec.samples_per_group} " + f"--rollout-max-response-len {spec.max_response_len} --seq-len {spec.seq_len} " + f"--inner-lr {spec.lr} --seed {spec.seed} --sglang-mem-fraction-static {spec.mem_fraction} " + f"--gpu modal:{c['total']}x{card} --rl-single-island-no-sync --total-steps {spec.rounds} " + # Miles recipe settings of the bare arm (S26 GOAL §6: the bare framework's official + # settings win): no deterministic inference, constant LR, full recompute, dynamic + # batching, gradient-accumulation fusion; dropout 0 is the generic recipe's own + "--no-sglang-deterministic-inference --rl-lr-schedule constant --rl-activation-recompute " + f"--rl-max-tokens-per-gpu {max_tokens_per_gpu(spec)} --rl-gradient-accumulation-fusion " + # TIS (upper 2.0) and eps_clip_high 0.28 as in the bare arms; staleness 2 >= max_policy_age + f"--trust-remote-code --rl-algorithm-spec {MILES_ALGORITHM_SPEC}") + if kwargs: + args += f" --apply-chat-template-kwargs {shlex.quote(kwargs)}" + if mode == "async": + args += (f" --rl-placement fixed-partition --rl-rollout-gpus {c['rollout']} " + f"--rl-max-policy-age {spec.max_policy_age}") + return args + + +YETO_TASK_FILE = "eff-tasks.jsonl" # the local task file in the launcher's checkout (local_task_file) + + +def uses_task_file(spec: EffSpec) -> bool: + """The yeto+Miles arms read a local task file (eff_prep.yeto_task_rows) unless the + spec is the full gsm8k train split, which they read from the Hub as before.""" + return bool(spec.data_file or spec.train_rows) + + +def write_task_file(spec: EffSpec, path: Path) -> Path: + """Write the yeto arms' local task file (the bare arms' training rows, in order).""" + from huggingface_hub import hf_hub_download + + from .eff_prep import yeto_task_rows + + path.write_text("".join(json.dumps(r) + "\n" for r in yeto_task_rows(spec.to_dict(), hf_hub_download))) + return path + + +def _yeto_data(spec: EffSpec) -> str: + from ..miles import DATA + + if uses_task_file(spec): + return f"--rl-allow-local-data --data {YETO_TASK_FILE}" + if spec.data != "zhuzilin/gsm8k" or spec.data_revision not in DATA: + raise ValueError("the yeto+Miles arm reads tools.eval.miles.DATA; keep the spec's data equal to it") + return DATA + + +HEAD_REGION = "nebius/eu-north1" # the multi-island syncer head: a CPU-only Nebius VM + + +def yeto_miles_islands_launcher_args(spec: EffSpec, islands: int, run_id: str, *, data: str, + island_tasks: str) -> str: + """E11: the yeto+Miles arm split into ``islands`` async islands that sync through a + head syncer (LoRA strict-avg), with the card total and the prompts per round of the + single-island arm. ``data`` is a local JSONL of the spec's rows and ``island_tasks`` + the ``--rl-island-tasks`` spec that gives each island its share of every round.""" + c = cards("async", spec.cards) + if c["total"] % islands or (c["total"] // islands) % 2 or spec.groups_per_round % islands: + raise ValueError(f"{islands} islands need an even card share and an equal prompt share " + f"({c['total']} cards, {spec.groups_per_round} prompts per round)") + per = c["total"] // islands + card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] + kwargs = json.dumps(spec.chat_template_kwargs) if spec.chat_template_kwargs else "" + args = (f"launch --controller head --syncer-region {HEAD_REGION} --on-demand --training-mode rl " + f"--rl-engine ports --preflight-threads wait --preflight-thread-start 2950 --modal-gpu-exact " + f"{gpu_name_arg(spec)}--modal-retries 0 --no-island-relaunch " + f"--modal-cpu {8 * per} --modal-memory-gib {48 * per} " + f"--cluster-prefix {run_id} --tuning lora --lora-r {spec.lora_r} --lora-targets all-linear " + f"--rl-observe-timeline --model {spec.model} --model-revision {spec.model_revision} " + f"--rl-allow-local-data --data {data} --rl-island-tasks {shlex.quote(island_tasks)} " + "--reward-function gsm8k_reward:score --fragments 1 --pipeline 1 " + f"--rollout-batch-size {spec.groups_per_round // islands} " + f"--n-samples-per-prompt {spec.samples_per_group} " + f"--rollout-max-response-len {spec.max_response_len} --seq-len {spec.seq_len} " + f"--inner-lr {spec.lr} --seed {spec.seed} --sglang-mem-fraction-static {spec.mem_fraction} " + f"--gpu {','.join([f'modal:{per}x{card}'] * islands)} --total-steps {spec.rounds} " + # the single-island arm's Miles recipe settings (yeto_miles_launcher_args) + "--no-sglang-deterministic-inference --rl-lr-schedule constant --rl-activation-recompute " + f"--rl-max-tokens-per-gpu {max_tokens_per_gpu(spec)} --rl-gradient-accumulation-fusion " + f"--trust-remote-code --rl-algorithm-spec {MILES_ALGORITHM_SPEC} " + # elastic: the multi-island LoRA path with real-card runs (S24 d4, S25 c3); legacy + # scheduling hung in the first apply on both islands (e11-yeto-2i-1, no round in 13 min) + "--rl-island-scheduling elastic " + f"--rl-placement fixed-partition --rl-rollout-gpus {per // 2} " + f"--rl-max-policy-age {spec.max_policy_age}") + if kwargs: + args += f" --apply-chat-template-kwargs {shlex.quote(kwargs)}" + return args + + +# E3 "C" sync configuration (s1-runs/s26-e3/run.sh GROUP=C): 2 fragments, fragment pipeline, +# background checkpoint every 4 sweeps, bf16 upload, 8 adaptive WAN streams (E3d X4a/X4b) +FULL_PARAMETER_FRAGMENTS = 2 +QWEN3_06B_LAYOUT_SHA256 = "3d5a0936038d8c9587afffc4bc043769d094695c8916ebef8696c0758d83a3b5" + + +def yeto_miles_full_islands_launcher_args(spec: EffSpec, islands: int, run_id: str, *, data: str, + island_tasks: str) -> str: + """E11 full parameter: ``islands`` islands of 1 rollout + 1 trainer card that sync every + parameter through the head syncer with the E3 C configuration. The islands together + read the single-island arm's prompts of each round (``island_task_split``) and run its + Miles recipe settings; the full-parameter sweep needs policy age 0 (E3).""" + if spec.tuning != "full" or spec.model != "Qwen/Qwen3-0.6B": + raise ValueError("the full-parameter island arm is pinned to Qwen3-0.6B (layout hash)") + c = cards("async", spec.cards) + if c["total"] != 2 * islands or spec.groups_per_round % islands: + raise ValueError(f"{islands} islands of 2 cards need {2 * islands} cards and an equal prompt share") + card = {"A10G": "a10g", "H100!": "h100", "L4": "l4"}[spec.gpu] + frags = FULL_PARAMETER_FRAGMENTS + kwargs = json.dumps(spec.chat_template_kwargs) if spec.chat_template_kwargs else "" + args = (f"launch --controller head --syncer-region {HEAD_REGION} --on-demand --training-mode rl " + f"--rl-engine ports --preflight-threads wait --preflight-thread-start 2950 --modal-gpu-exact " + f"{gpu_name_arg(spec)}--modal-retries 0 --no-island-relaunch --modal-cpu 8 --modal-memory-gib 64 " + f"--cluster-prefix {run_id} --rl-observe-timeline " + f"--model {spec.model} --model-revision {spec.model_revision} " + f"--rl-allow-local-data --data {data} --rl-island-tasks {shlex.quote(island_tasks)} " + "--reward-function gsm8k_reward:score --pipeline 1 " + f"--rollout-batch-size {spec.groups_per_round // islands} " + f"--n-samples-per-prompt {spec.samples_per_group} " + f"--rollout-max-response-len {spec.max_response_len} --seq-len {spec.seq_len} " + f"--inner-lr {spec.lr} --seed {spec.seed} --sglang-mem-fraction-static {spec.mem_fraction} " + f"--gpu {','.join([f'modal:2x{card}'] * islands)} --total-steps {spec.rounds} " + "--no-sglang-deterministic-inference --rl-lr-schedule constant --rl-activation-recompute " + f"--rl-max-tokens-per-gpu {max_tokens_per_gpu(spec)} --rl-gradient-accumulation-fusion " + f"--trust-remote-code --rl-algorithm-spec {MILES_ALGORITHM_SPEC} " + f"--rl-sync-preset full-parameter --rl-parameter-layout-sha256 {QWEN3_06B_LAYOUT_SHA256} " + f"--fragments {frags} --rl-sync-fragment-pipeline {frags} --rl-sync-async-checkpoint " + f"--rl-sync-push-dtype bf16 --rl-sync-checkpoint-every {4 * frags} " + "--wan-streams 8 --wan-stream-schedule adaptive --wan-stall-timeout 90 " + "--rl-placement fixed-partition --rl-rollout-gpus 1 --rl-max-policy-age 0") + if kwargs: + args += f" --apply-chat-template-kwargs {shlex.quote(kwargs)}" + return args + + +def island_task_split(groups_per_round: int, rounds: int, islands: int) -> dict[str, list[str]]: + """Round-major split of the first ``rounds * groups_per_round`` rows (task ids + ``gsm``): island k gets the k-th equal slice of every round's prompts, so the + islands together read the single-island arm's prompts of each round.""" + share = groups_per_round // islands + return {f"l{k}": [f"gsm{r * groups_per_round + k * share + i + 1}" for r in range(rounds) + for i in range(share)] for k in range(islands)} + + +# ====================================================================== plans +def bare_script(spec: EffSpec, arm: Arm, run_id: str) -> str: + """The bash script a bare container runs (after eff_prep.py staged model and data).""" + c = cards(arm.mode, spec.cards) + lines = ["set -euo pipefail", f"mkdir -p {OUT_DIR}", + f"python3 {EFF_DIR}/eff_prep.py {EFF_DIR}/spec.json > {OUT_DIR}/prep.json"] + if arm.framework == "verl": + lines += [f"cd {VERL_ROOT}", f"export VERL_FILE_LOGGER_PATH={OUT_DIR}/verl-metrics.jsonl", + f"export PYTHONPATH={EFF_DIR}", + # Modal: the trainer allocates NCCL host buffers with cuMem, the + # rollout-side receiver falls back to /dev/shm and cannot attach + # (eff01-bare-verl-async-d, NCCL_DEBUG=INFO); both use /dev/shm. + *(["export NCCL_CUMEM_HOST_ENABLE=0"] if arm.mode == "async" else []), + "exec python3 -m " + verl_bare_entry(arm.mode) + " " + shlex.join( + verl_bare_overrides(spec, arm.mode, run_id))] + else: + # model scripts derived from another (qwen3-4B-Instruct-2507.py) import Miles' + # model_args_utils (miles/utils/external_utils), which is not on the scripts' path + load = ("import importlib.util as u, sys; sys.path.insert(0, '{d}'); " + f"sys.path.insert(0, '{MILES_ROOT}/miles/utils/external_utils'); " + "s = u.spec_from_file_location('m', '{d}/{f}'); m = u.module_from_spec(s); " + "s.loader.exec_module(m); print(m.model_args())").format(d=f"{MILES_ROOT}/scripts/models", + f=spec.miles_model_args) + runtime_env = json.dumps({"env_vars": {"PYTHONPATH": f"{MEGATRON_ROOT}/:{EFF_DIR}", + "CUDA_DEVICE_MAX_CONNECTIONS": "1"}}) + lines += [f"cd {MILES_ROOT}", f"MODEL_ARGS=$(python3 -c {shlex.quote(load)})", + f"echo \"$MODEL_ARGS\" > {OUT_DIR}/model_args.txt", + f"ray start --head --node-ip-address 127.0.0.1 --num-gpus {c['total']} --num-cpus 32 " + "--disable-usage-stats --dashboard-host=127.0.0.1", + f"exec ray job submit --address=http://127.0.0.1:8265 --runtime-env-json={shlex.quote(runtime_env)} " + f"-- python3 {MILES_ROOT}/{miles_bare_script(arm.mode)} $MODEL_ARGS " + + shlex.join(miles_bare_argv(spec, arm.mode))] + return "\n".join(lines) + "\n" + + +def plan(spec: EffSpec, arm_name: str, run_id: str) -> dict[str, Any]: + """Everything one run of an arm needs (written to ``arm.json`` before the launch).""" + arm = Arm.parse(arm_name) + c = cards(arm.mode, spec.cards) + doc: dict[str, Any] = { + "arm": arm.name, "system": arm.system, "framework": arm.framework, "mode": arm.mode, "run_id": run_id, + "spec": spec.to_dict(), "fingerprint": control_fingerprint(spec), "cards": c, "gpu": spec.gpu, + } + if arm.system == "bare": + doc.update(harness="modal_eff", image=arm.framework, script=bare_script(spec, arm, run_id)) + elif arm.framework == "verl": + # island_entry reads the rollout memory fraction from YETO_VERL_GPU_MEM (default 0.5) + doc.update(harness="modal_eff", image="verl", env={"YETO_VERL_GPU_MEM": str(spec.mem_fraction)}, + argv=yeto_verl_argv(spec, arm.mode, "/root/yeto-output/rl-island-0.jsonl")) + else: + doc.update(harness="miles_local", launcher_args=yeto_miles_launcher_args(spec, arm.mode, run_id)) + return doc + + +# ====================================================================== control check +def _override_map(overrides: list[str]) -> dict[str, str]: + out = {} + for o in overrides: + key, _, value = o.partition("=") + out[key.lstrip("+")] = value + return out + + +def verl_knob_diff(spec: EffSpec, mode: str) -> dict[str, dict[str, str | None]]: + """Hydra overrides where bare verl and yeto+verl differ (key -> {bare, yeto}). + + The yeto side is yeto's own translation (``build_overrides``, plus + ``fully_async_run_overrides`` for async) of the same spec; paths, names and + keys yeto handles in its own loop are compared too, so the list is the full + set of differences, not only the interesting ones.""" + from yeto.rl.adapters.verl.config import VerlRunConfig, build_overrides + from yeto.rl.adapters.verl.fully_async_round import fully_async_run_overrides + + cfg = VerlRunConfig( + model_path=_model_path(spec), train_file=f"{DATA_DIR}/train.parquet", val_file=f"{DATA_DIR}/test.parquet", + out_dir=OUT_DIR, lora_rank=spec.lora_r, lora_alpha=spec.lora_r, lr=spec.lr, + groups_per_round=spec.groups_per_round, samples_per_group=spec.samples_per_group, + max_prompt_len=spec.max_prompt_len, max_response_len=spec.max_response_len, seed=spec.seed, + total_rounds=spec.rounds, temperature=spec.temperature, top_p=spec.top_p, tis_upper=spec.tis_upper, + gpu_memory_utilization=spec.mem_fraction, reward_path=f"{EFF_DIR}/eff_reward.py", + chat_template_kwargs=dict(spec.chat_template_kwargs), experiment_name="x") + yeto = build_overrides(cfg) + if mode == "async": + yeto = fully_async_run_overrides(yeto, spec.max_policy_age, groups_per_round=spec.groups_per_round, + rounds=spec.rounds) + bare, ours = _override_map(verl_bare_overrides(spec, mode, "x")), _override_map(yeto) + return {k: {"bare": bare.get(k), "yeto": ours.get(k)} for k in sorted(set(bare) | set(ours)) + if bare.get(k) != ours.get(k)} + + +def write_plan(path: Path, doc: dict[str, Any]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(doc, indent=1, default=str)) diff --git a/tools/eval/miles_local.sh b/tools/eval/miles_local.sh index 6a24cba1..34d5458f 100644 --- a/tools/eval/miles_local.sh +++ b/tools/eval/miles_local.sh @@ -9,7 +9,10 @@ # ARGS (launcher args; the store is modal-volume://$STOREVOL/$P), SARGS (optional), # ISLANDS (containers before the GPU lock is released), HARD (s), STOREVOL, TAPEVOL, # RECLAIM_LABEL (e.g. "l0") + RECLAIM_ROUND (rollout id written before the stop), -# PY (python with modal + the launcher deps), DRY (non-empty: --dry-run only) +# PY (python with modal + the launcher deps), DRY (non-empty: --dry-run only), +# EXTRA_FILES (optional: files copied into the checkout root and committed), +# STALL_S / FIRST_STEP_S (optional: stop the apps when no Miles training step line +# reaches the launch log for STALL_S s, or before the first one for FIRST_STEP_S s) set -u : "${P:?}" "${R:?}" "${REPO:?}" "${ARGS:?}" "${STOREVOL:?}" "${PY:?}" HARD=${HARD:-3000}; ISLANDS=${ISLANDS:-1}; TAPEVOL=${TAPEVOL:-yeto-event-tapes} @@ -22,7 +25,9 @@ mkdir -p $R/tapes $R/store # a real clone: the code pack needs a git checkout git clone -q --no-local $REPO $CLONE && git -C $REPO rev-parse HEAD > $R/yeto_sha.txt cp $REPO/yeto/rl/gsm8k_reward.py $CLONE/gsm8k_reward.py -(cd $CLONE && git add -f gsm8k_reward.py && git -c user.email=eval@yeto -c user.name=t-eval commit -qm 'run inputs') +# EXTRA_FILES: run inputs the launcher reads from its checkout (e.g. a local task file) +for f in ${EXTRA_FILES:-}; do cp "$f" $CLONE/; done +(cd $CLONE && git add -f gsm8k_reward.py $(for f in ${EXTRA_FILES:-}; do basename "$f"; done) && git -c user.email=eval@yeto -c user.name=t-eval commit -qm 'run inputs') eval "$(/usr/bin/python3 - <<'PYC' import shlex, tomllib t = tomllib.load(open("/home/michael/.modal.toml", "rb")) @@ -68,6 +73,21 @@ for i in $(seq 1 120); do kill -0 $SUB 2>/dev/null || break; sleep 10 done exec 9>&- +# no-progress watchdog (tools/eval/efficiency/modal_eff.py STALL_S): stop the apps when the +# count of Miles training step lines stops growing +if [ -n "${STALL_S:-}" ]; then +( cd /tmp; last=0; since=$(date +%s); limit=${FIRST_STEP_S:-1800} + while kill -0 $SUB 2>/dev/null; do + sleep 15 + n=$(grep -ac "'train/step'" $R/launch.ts.log 2>/dev/null) + if [ "${n:-0}" -gt "$last" ]; then last=$n; since=$(date +%s); limit=$STALL_S; fi + if [ $(( $(date +%s) - since )) -gt "$limit" ]; then + echo "$(date -u +%FT%TZ) no training step for $(( $(date +%s) - since )) s (limit $limit s, steps $last): stopping apps" > $R/STALL_FIRED + for a in $APP $SAPP; do timeout 120 $MODAL app stop -y $a; done >> $R/STALL_FIRED 2>&1 + exit 0 + fi + done ) 9>&- & +fi # one reclaim: island $RECLAIM_LABEL wrote rollout $RECLAIM_ROUND -> stop its container if [ -n "${RECLAIM_LABEL:-}" ]; then ( cd /tmp diff --git a/tools/eval/modal_verl.py b/tools/eval/modal_verl.py index eadc4699..351bb110 100644 --- a/tools/eval/modal_verl.py +++ b/tools/eval/modal_verl.py @@ -25,6 +25,7 @@ import json import os +import re import secrets import shutil import subprocess @@ -64,13 +65,19 @@ # ====================================================================== in-container helpers +def gpu_name_matches(want: str, name: str) -> bool: + """``name`` is the model ``want`` names: "NVIDIA A10" accepts "NVIDIA A10", not "NVIDIA A10G" + (Modal "A10G" containers report either; the A10G trains slower, E1 e1a9).""" + return re.search(re.escape(want) + r"(?![\w-])", name) is not None + + def _gpu(out: Path, cfg: dict) -> list[str]: names = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], capture_output=True, text=True, timeout=60).stdout.split("\n") names = [n.strip() for n in names if n.strip()] (out / "nvidia-smi.txt").write_text(subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout) want = GPU_NAMES[cfg["gpu"]] - if not names or not all(want in n for n in names): + if not names or not all(gpu_name_matches(want, n) for n in names): raise RuntimeError(f"asked for {want}, got {names}") return names diff --git a/yeto/cli.py b/yeto/cli.py index 7c68497f..3333988b 100644 --- a/yeto/cli.py +++ b/yeto/cli.py @@ -438,6 +438,45 @@ 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", + ) + 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( @@ -459,6 +498,14 @@ def int_or_auto(value: str): "may upgrade H100 to H200) and fail at container start unless " "nvidia-smi reports it; use for bitwise comparisons", ) + infra.add_argument( + "--modal-gpu-name", + default=None, + metavar="MODEL", + help="Modal islands: fail at container start unless every GPU nvidia-smi reports " + "is this model (e.g. 'NVIDIA A10': Modal A10G containers are an A10 or an A10G); " + "use when runs must compare on one GPU model", + ) infra.add_argument( "--modal-retries", type=int, diff --git a/yeto/downlink.py b/yeto/downlink.py new file mode 100644 index 00000000..d3a613c0 --- /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/fleet.py b/yeto/launch/fleet.py index b7051a46..d1d9ce63 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 a7f12af6..9d71765e 100644 --- a/yeto/launch/modal_islands.py +++ b/yeto/launch/modal_islands.py @@ -104,6 +104,7 @@ def build_modal_island_config(args, spec: ClusterSpec, learner_id: int, task, sy envs={k: str(v) for k, v in envs.items() if v is not None}, region=spec.region, gpu_exact=bool(getattr(args, "modal_gpu_exact", False)), + gpu_name=getattr(args, "modal_gpu_name", None) or None, registry_login=bool(registry_login), registry_creds=registry_login, image_ref=image_ref_from_rl_image(args.rl_image) if rl else None, @@ -180,14 +181,13 @@ def pull_modal_tapes(modal_ops, modal_cfgs: dict, run_dir) -> dict[str, str]: def require_modal_for_gpu_exact(args, specs: list[ClusterSpec]) -> None: - """--modal-gpu-exact only means something on Modal islands; refuse it + """--modal-gpu-exact / --modal-gpu-name only mean something on Modal islands; refuse them rather than silently launching unpinned learners elsewhere.""" - if getattr(args, "modal_gpu_exact", False): - other = sorted({s.cloud for s in specs if s.cloud != "modal"}) + for flag, value in (("--modal-gpu-exact", getattr(args, "modal_gpu_exact", False)), + ("--modal-gpu-name", getattr(args, "modal_gpu_name", None))): + other = sorted({s.cloud for s in specs if s.cloud != "modal"}) if value else [] if other: - raise ValueError( - f"--modal-gpu-exact applies only to Modal islands; --gpu also has {other}" - ) + raise ValueError(f"{flag} applies only to Modal islands; --gpu also has {other}") def _tail_modal(modal_ops, call_id: str, prefix: str, collector=None, guard=None) -> int: @@ -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_args.py b/yeto/launch/rl_args.py index 0c21669b..bdb4ce16 100644 --- a/yeto/launch/rl_args.py +++ b/yeto/launch/rl_args.py @@ -148,6 +148,40 @@ 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"), + ("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 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. @@ -169,11 +203,30 @@ 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") + 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 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 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 @@ -267,8 +320,7 @@ def _prepare_rl_args( from ..rl.engine import sample_roles if args.tuning != "lora" and not sample_roles.buffer_full_parameter(args): - raise ValueError("RL v0 requires --tuning lora (--tuning full: only sampler / trainer " - "islands, --rl-sample-role or --rl-island-sample-role, with " + raise ValueError("RL v0 requires --tuning lora (--tuning full: only with " "--rl-single-island-no-sync)") if args.lora_r <= 0: raise ValueError("RL v0 requires a positive LoRA rank") @@ -723,6 +775,9 @@ 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) + _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 1185f4c9..ec60e1f1 100644 --- a/yeto/launch/rl_island_task.py +++ b/yeto/launch/rl_island_task.py @@ -51,6 +51,34 @@ 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}" + 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: + 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") + flags += f" --wan-probe-min-mbps {min_up!r} --wan-probe-min-down-mbps {min_down!r} --wan-probe-mib {mib}" + 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); @@ -72,6 +100,12 @@ 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, "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 @@ -150,16 +184,20 @@ 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 "") + + (" --rl-gradient-accumulation-fusion" + if getattr(args, "rl_gradient_accumulation_fusion", False) else "") + (f" --rl-max-tokens-per-gpu {int(args.rl_max_tokens_per_gpu)}" if getattr(args, "rl_max_tokens_per_gpu", None) is not None else "") + (f" --rl-optimizer-offload-fraction {float(args.rl_optimizer_offload_fraction)!r}" 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/launch/syncer.py b/yeto/launch/syncer.py index 93b1aa6a..2836e1b0 100644 --- a/yeto/launch/syncer.py +++ b/yeto/launch/syncer.py @@ -145,6 +145,17 @@ 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 _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) @@ -161,7 +172,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}" @@ -205,7 +217,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, @@ -442,11 +454,12 @@ 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" 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/launcher.py b/yeto/launcher.py index 0abd2f4b..7819fbc5 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/modal_runner.py b/yeto/modal_runner.py index b25be3c4..761b53c9 100644 --- a/yeto/modal_runner.py +++ b/yeto/modal_runner.py @@ -204,6 +204,10 @@ class ModalIslandConfig: # Request exactly `gpu` ("H100!" -- no H200 upgrade) and fail the # container at start-up unless nvidia-smi reports that type. gpu_exact: bool = False + # E1: fail the container at start-up unless every GPU nvidia-smi reports is this + # model ("NVIDIA A10" refuses "NVIDIA A10G": Modal "A10G" containers are either, + # and the A10G trains about 12 % slower); None: any card of type `gpu`. + gpu_name: str | None = None # Pull image_ref with a private-registry login: ``registry_creds`` (the # SKYPILOT_DOCKER_* triple launcher.registry_login_for resolved) when # given, else -- with ``registry_login`` -- the launching process's @@ -394,6 +398,19 @@ def check_gpu_names(gpu: str, names: list[str], count: int) -> None: ) +def gpu_model_matches(want: str, name: str) -> bool: + """``name`` is the GPU model ``want`` names, not a longer model name that starts + with it ("NVIDIA A10" matches "NVIDIA A10", not "NVIDIA A10G").""" + return re.search(re.escape(want) + r"(?![\w-])", name) is not None + + +def check_gpu_model(want: str, names: list[str], count: int) -> None: + """Raise unless exactly ``count`` GPUs of the model ``want`` are visible.""" + wrong = [n for n in names if not gpu_model_matches(want, n)] + if wrong or len(names) != count: + raise RuntimeError(f"requested {count}x {want!r} (--modal-gpu-name) but the container sees {names!r}") + + def visible_gpu_names() -> list[str]: out = subprocess.run( ["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"], @@ -883,7 +900,7 @@ def _island_body(cfg: ModalIslandConfig, rank: int, ips: list[str], all_ips: dic try: names = visible_gpu_names() except (OSError, subprocess.SubprocessError) as exc: - if cfg.gpu_exact: + if cfg.gpu_exact or cfg.gpu_name: raise RuntimeError(f"island {cfg.learner_id}: cannot read GPU names: {exc}") from exc names = [] print(f"[modal-island {cfg.learner_id}] requested {cfg.gpu_request}, got {names}", flush=True) @@ -896,6 +913,8 @@ def _island_body(cfg: ModalIslandConfig, rank: int, ips: list[str], all_ips: dic tape.start() if cfg.gpu_exact: check_gpu_names(cfg.gpu, names, cfg.gpus_per_node) + if cfg.gpu_name: + check_gpu_model(cfg.gpu_name, names, cfg.gpus_per_node) if cfg.setup_script: code = subprocess.call(container_command(cfg.setup_script), env=env) if code != 0: diff --git a/yeto/protocol.py b/yeto/protocol.py index 7beb2254..e55444b5 100644 --- a/yeto/protocol.py +++ b/yeto/protocol.py @@ -34,10 +34,14 @@ 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 .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 +from .wan_streams import STREAM_SCHEDULES, StreamMeter, StreamSnapshot, StreamStalledError, pick_stream MAGIC = 0xD170C0DE PROTOCOL_VERSION = 4 @@ -56,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 @@ -66,6 +72,20 @@ # 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 +# 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 @@ -87,13 +107,127 @@ _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. + """ + base = session_base(dtype) + if base == DTYPE_Q4: + return DTYPE_BF16 + if base in (DTYPE_F32_PUSH_BF16, DTYPE_F32_PUSH_I8): + return DTYPE_F32 + return base + + +def push_wire_dtype(dtype: int) -> int: + """Dtype of unquantized PUSH_FRAGMENT deltas on the wire for a session 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. """ - return DTYPE_BF16 if dtype == DTYPE_Q4 else dtype + 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: + """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: @@ -168,8 +302,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() @@ -440,6 +574,7 @@ class SendProgress: queued_chunks: int reconnects: int connected: bool + stalls: int = 0 @dataclass @@ -463,6 +598,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], @@ -477,11 +616,19 @@ 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, + 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 dtype not in (DTYPE_F32, DTYPE_BF16, DTYPE_Q4): + 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: @@ -499,7 +646,27 @@ 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 + # 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 + # 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 @@ -510,7 +677,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( @@ -544,6 +711,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() @@ -564,6 +732,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 ------------------------------------------------------------- @@ -572,6 +741,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 and self._initial_cut_received.is_set(): + set_tcp_user_timeout(sock, self.stream_stall_timeout) return sock def _connect_one(self) -> socket.socket: @@ -602,6 +773,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. @@ -659,7 +833,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] @@ -684,6 +858,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: @@ -776,10 +952,72 @@ 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_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}") @@ -832,8 +1070,9 @@ def push_fragment( numel = self.layout.fragments[fragment_id].numel expected = ( _q4_nbytes(numel) - if self.dtype == DTYPE_Q4 - else _tensor_nbytes(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( @@ -851,6 +1090,10 @@ def push_fragment( c_steps, c_tokens, ) + 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( @@ -882,11 +1125,15 @@ 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(self.dtype, numel) - ) + 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") @@ -1129,6 +1379,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), @@ -1136,8 +1388,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: @@ -1152,7 +1412,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 @@ -1273,11 +1533,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( @@ -1337,20 +1593,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: @@ -1419,6 +1740,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: @@ -1438,6 +1760,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: @@ -1452,6 +1775,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") @@ -1499,6 +1823,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 " @@ -1612,6 +1943,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: @@ -1767,15 +2099,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] @@ -1882,6 +2218,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/config.py b/yeto/rl/adapters/miles/config.py index dadbb341..1af82e71 100644 --- a/yeto/rl/adapters/miles/config.py +++ b/yeto/rl/adapters/miles/config.py @@ -472,6 +472,7 @@ def _check_sglang_tp(value, config): # translated below (Megatron full recompute / Miles dynamic batching) "train_memory.activation_recompute": _ok, "train_memory.max_tokens_per_gpu": _ok, + "train_memory.gradient_accumulation_fusion": _ok, "train_memory.optimizer_offload_fraction": _ok, **{f"sample_buffer.{name}": _ok for name in ("role", "store", "max_backlog_groups", "wait_s", "min_round_s")}, @@ -781,6 +782,7 @@ def translate_run_config( from yeto.rl.engine.run_config import ( RECIPE_DEEPSEEK_V4_FLASH, + RECIPE_GENERIC, RECIPE_GLM5_NEXT, RECIPE_QWEN3_5, RECIPE_QWEN3_8_NEXT, @@ -864,6 +866,10 @@ def translate_run_config( model_recipe_values = deepseek_v4.recipe_argv(trainable) else: model_name = recipe.provider_class + if recipe.name == RECIPE_GENERIC: + # Megatron's default dropout is 0.1; the Miles recipes train RL with + # none (the policy's train-time log-probs must match the rollout's) + model_recipe_values = ["--attention-dropout", "0.0", "--hidden-dropout", "0.0"] values: list[str] = [ "train.py", @@ -938,7 +944,7 @@ def translate_run_config( "--accumulate-allreduce-grads-in-fp32", "--attention-softmax-in-fp32", "--attention-backend", recipe.training_attention_backend, - "--no-gradient-accumulation-fusion", + *(() if config.train_memory.gradient_accumulation_fusion else ("--no-gradient-accumulation-fusion",)), "--bf16", "--no-load-optim", "--no-load-rng", diff --git a/yeto/rl/adapters/miles/cut_plugin.py b/yeto/rl/adapters/miles/cut_plugin.py index 50fe3afd..1e7d0b80 100644 --- a/yeto/rl/adapters/miles/cut_plugin.py +++ b/yeto/rl/adapters/miles/cut_plugin.py @@ -54,6 +54,7 @@ _MODULE = "yeto.rl.adapters.miles.cut_plugin" SAVE_CUT_SHARD = f"{_MODULE}.save_cut_shard" +FINISH_CUT_SHARD = f"{_MODULE}.finish_cut_shard" RESTORE_CUT_SHARD = f"{_MODULE}.restore_cut_shard" SHARD_SCHEMA = "yeto.cut_shard/v1" # Opt-in: read the optimizer back right after the load and report where a mismatch arises. @@ -99,11 +100,17 @@ def _dps(self): return dp_invariant_state - def export_optimizer(self, optimizer: Any, named: list) -> Any: + def export_optimizer(self, optimizer: Any, named: list, *, on_device: bool = False) -> Any: + """The fork-M5 named optimizer export. ``on_device`` (E1): the tensors are + device copies (same layout, not yet on the host); :func:`to_host` turns the + result into exactly the host export.""" + export = self._dps().export_named_optimizer_state + if on_device: + export = _device_export(self._dps()) if _UNSAFE_STATE_READS[0]: # diagnostic only: verify the fork fix without the yeto guard - return self._dps().export_named_optimizer_state(optimizer, named) + return export(optimizer, named) with side_effect_free_state(optimizer): - return self._dps().export_named_optimizer_state(optimizer, named) + return export(optimizer, named) def check_optimizer(self, optimizer: Any, named: list, states: list) -> Any: """Merge the DP shards of one (tp, pp) (DistOpt ranges) and validate against ``optimizer``.""" @@ -140,6 +147,51 @@ def set_megatron_counters(self, counters: Mapping[str, int]) -> None: setattr(margs, key, int(value)) +def _device_export(dps: Any) -> Any: + """``export_named_optimizer_state`` keeping each tensor on its device (E1). + + The Miles fork takes ``device=None`` for this (yeto/s26-e1-device-export); the + pinned image's fork does not yet, and then this module builds the same + entries from the fork's own ``optimizer_slots`` -- the fork function's body + without its ``.cpu()`` (tests pin the equality with the fork's host export).""" + import inspect + + export = dps.export_named_optimizer_state + if "device" in inspect.signature(export).parameters: + return lambda optimizer, named: export(optimizer, named, device=None) + + def export_on_device(optimizer: Any, named: list) -> dict[str, Any]: + import copy + + entries = {} + for slot in dps.optimizer_slots(optimizer, named): + tensors, scalars = {}, {} + for key, value in slot.get().items(): + value = value.detach() + if value.dim() == 0: + scalars[key] = value.clone() + continue + flat = value.reshape(-1).clone() + if flat.numel() != slot.end - slot.start: + raise dps.DpInvariantStateError( + f"{slot.name!r}: state {key!r} has {flat.numel()} elements but the owned range is " + f"[{slot.start}, {slot.end})" + ) + tensors[key] = flat + entries[slot.name] = { + "numel": slot.model_param.numel(), + "shape": tuple(slot.model_param.shape), + "start": slot.start, + "end": slot.end, + "tensors": tensors, + "scalars": scalars, + "hyper": {k: copy.deepcopy(v) for k, v in slot.group.items() if k != "params"}, + } + return {"format": dps.NAMED_STATE_FORMAT, "entries": entries} + + return export_on_device + + def _expert_coord(ps: Any) -> dict[str, int]: """EP coordinates for the same-shape check (only when EP>1: older cuts stay comparable).""" if int(ps.ep.size) <= 1: @@ -332,10 +384,12 @@ def _feed(digest: Any, value: Any) -> None: shape = tuple(value.shape) t = value.detach().to("cpu").contiguous().reshape(-1) digest.update(f"T{t.dtype}{shape}".encode()) - digest.update(t.view(torch.uint8).numpy().tobytes() if t.numel() else b"") + # a buffer, not a bytes copy: same digest, and hashlib releases the GIL while it + # hashes (a background cut writer must not stall the training thread, E1) + digest.update(t.view(torch.uint8).numpy() if t.numel() else b"") elif isinstance(value, np.ndarray): digest.update(f"N{value.dtype}{value.shape}".encode()) - digest.update(np.ascontiguousarray(value).tobytes()) + digest.update(np.ascontiguousarray(value).view(np.uint8) if value.size else b"") elif isinstance(value, Mapping): digest.update(b"{") for key in sorted(value, key=repr): @@ -538,18 +592,33 @@ def _entry_step(entry: Mapping[str, Any]) -> Any: return None -def resume_rank_state(actor: Any, optimizer_named: Any) -> dict[str, Any]: - """This rank's :class:`~yeto.rl.engine.resume_manifest.OptimizerResumeState` (as a dict), - read from the named optimizer export and the Megatron LR scheduler (S25 N7.6).""" - from yeto.rl.engine.resume_manifest import build_rank_state, scheduler_progress_of +def resume_inputs(actor: Any) -> dict[str, Any]: + """What :func:`resume_rank_state` reads from the live trainer: the optimizer + group lrs and the Megatron LR scheduler progress.""" + from yeto.rl.engine.resume_manifest import scheduler_progress_of + + return {"lrs": [float(g["lr"]) for g in getattr(actor.optimizer, "param_groups", ()) if "lr" in g], + "scheduler_progress": scheduler_progress_of(actor.opt_param_scheduler)} + +def resume_rank_state(actor: Any, optimizer_named: Any, *, lrs: list[float] | None = None, + scheduler_progress: int | None = None) -> dict[str, Any]: + """This rank's :class:`~yeto.rl.engine.resume_manifest.OptimizerResumeState` (as a dict), + read from the named optimizer export and the Megatron LR scheduler (S25 N7.6). + ``lrs`` / ``scheduler_progress`` (from :func:`resume_inputs`, taken at the step + boundary) replace the reads of ``actor`` (then ``None`` is allowed).""" + from yeto.rl.engine.resume_manifest import build_rank_state + + if lrs is None or scheduler_progress is None: + live = resume_inputs(actor) + lrs = live["lrs"] if lrs is None else lrs + scheduler_progress = live["scheduler_progress"] if scheduler_progress is None else scheduler_progress entries = (optimizer_named or {}).get("entries") or {} steps = [s for s in (_entry_step(e) for e in entries.values()) if s is not None] - lrs = [float(g["lr"]) for g in getattr(actor.optimizer, "param_groups", ()) if "lr" in g] return build_rank_state( steps_by_group={"named": steps}, - scheduler_progress=scheduler_progress_of(actor.opt_param_scheduler), - lrs=lrs, + scheduler_progress=int(scheduler_progress), + lrs=list(lrs), moments_by_name={n: _entry_moments(e) for n, e in entries.items()}, ).to_dict() @@ -573,16 +642,71 @@ def resume_state(actor: Any) -> dict[str, Any]: RESUME_STATE = f"{_MODULE}.resume_state" -def _snapshot(actor: Any, backend: Any, named: list) -> dict[str, Any]: +def _snapshot(actor: Any, backend: Any, named: list, *, on_device: bool = False) -> dict[str, Any]: + """The shard state at this step boundary. ``on_device`` (E1): the adapter and + optimizer tensors are device copies (device-to-device, no wait for the host); + :func:`to_host` later makes them exactly the host snapshot.""" + if on_device: + adapter = {n: p.detach().clone() for n, p in named} + optimizer_named = backend.export_optimizer(actor.optimizer, named, on_device=True) + else: + adapter = {n: p.detach().to("cpu").clone() for n, p in named} + optimizer_named = backend.export_optimizer(actor.optimizer, named) return { - "adapter": {n: p.detach().to("cpu").clone() for n, p in named}, - "optimizer_named": backend.export_optimizer(actor.optimizer, named), + "adapter": adapter, + "optimizer_named": optimizer_named, "scheduler": actor.opt_param_scheduler.state_dict(), "megatron_counters": backend.megatron_counters(), "miles_counters": miles_counters(actor), } +def to_host(value: Any) -> Any: + """``value`` with every non-CPU tensor copied to the host (containers rebuilt, + everything else as it is). A CPU tensor is returned as it is.""" + import torch + + if isinstance(value, torch.Tensor): + return value if value.device.type == "cpu" else value.to("cpu") + if isinstance(value, dict): + return {k: to_host(v) for k, v in value.items()} + if isinstance(value, list): + return [to_host(v) for v in value] + if isinstance(value, tuple): + return tuple(to_host(v) for v in value) + return value + + +# E1: device copies of a background cut need this much free device memory besides +# their own size (twice the copies: the allocator's rounding and what training +# allocates next); below it the capture stays on the host path +DEVICE_COPY_ENV = "YETO_RL_CUT_DEVICE_COPY" +DEVICE_COPY_HEADROOM_BYTES = 1 << 30 + + +def _device_copy_plan(actor: Any, named: list) -> tuple[bool, str]: + """(copy on the device?, why): the device path needs CUDA parameters, a trainer + that stays resident, and free memory for the copies (adapter + FP32 master + + two Adam moments, estimated).""" + import torch + + if os.environ.get(DEVICE_COPY_ENV, "1").strip().lower() in {"0", "false", "no", "off"}: + return False, "disabled" + if getattr(getattr(actor, "args", None), "offload_train", False): + # --offload-train: the trainer sleeps right after the cut (colocated) and the + # writer's reads of the device copies then fail ("CUDA error: invalid + # argument", 5 of 6 cuts of GPU run e1g-yeto-miles-sync-1) + return False, "offload_train" + devices = {p.device for _, p in named} + if len(devices) != 1 or next(iter(devices)).type != "cuda": + return False, "not_cuda" + need = sum(p.numel() * (p.element_size() + 3 * 4) for _, p in named) + free, _total = torch.cuda.mem_get_info(next(iter(devices))) + if free < 2 * need + DEVICE_COPY_HEADROOM_BYTES: + return False, f"low_device_memory(free={free},need={need})" + return True, "ok" + + # -------------------------------------------------------------------------- # Plugins # -------------------------------------------------------------------------- @@ -596,32 +720,230 @@ def shard_name(coord: Mapping[str, int]) -> str: REFUSAL_KIND = "refusal_kind" # "refused" (deliberate) | "failed_before_write" -def save_cut_shard(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: +# E1: shards being written by a background thread of this rank, by cut id +_PENDING_SHARDS: dict[str, Any] = {} + + +def save_cut_shard(actor: Any, *, directory: str, cut_id: str, background: bool = False) -> dict[str, Any]: """Write this rank's shard (tmp + fsync + rename) and return its summary. A refusal (unsupported configuration, not at a step boundary, ...) is RETURNED as ``{"refused": reason}`` instead of raised: Miles marks a cell errored on any ``run_plugin`` exception (GPU run 2026-09-30, C1), which would turn a clean "no cut" into a dead trainer. + + ``background`` (E1): take the state now (at this step boundary) and write, hash + and summarize it off the training thread (:func:`_write_background`), so + training goes on; returns ``{"pending": cut_id}``. :func:`finish_cut_shard` + returns the summary. """ try: with trainer_resident(actor): - return _save(actor, directory=directory, cut_id=cut_id) + captured = _capture(actor, cut_id=cut_id, device_copy=background) + except CutPluginError as error: + return {REFUSED: str(error), REFUSAL_KIND: "refused"} + if not background: + return _write_refusing(captured, directory=directory, cut_id=cut_id) + from concurrent.futures import ThreadPoolExecutor + + if cut_id in _PENDING_SHARDS: + return {REFUSED: f"cut {cut_id} is already being written", REFUSAL_KIND: "refused"} + pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix=f"yeto-cut-{cut_id}") + _PENDING_SHARDS[cut_id] = (pool, pool.submit(_write_background, captured, directory=directory, cut_id=cut_id)) + pool.shutdown(wait=False) + return {"pending": cut_id, "coord": dict(captured["coord"]), "capture": dict(captured["capture"])} + + +def finish_cut_shard(actor: Any, *, cut_id: str) -> dict[str, Any]: + """The summary of a background :func:`save_cut_shard` (waits for its thread). + Its ``writer`` key says where the shard was serialized and hashed.""" + pending = _PENDING_SHARDS.pop(cut_id, None) + if pending is None: + return {REFUSED: f"no background write of cut {cut_id} on this rank", REFUSAL_KIND: "failed_before_write"} + try: + return pending[1].result() + except Exception as error: # noqa: BLE001 - reported like a refusal: the cut is not committed + return {REFUSED: f"background write of cut {cut_id} failed: {error!r}", REFUSAL_KIND: "failed"} + + +def _write_refusing(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: + try: + return _write(_with_resume(_on_host(captured)), directory=directory, cut_id=cut_id) + except CutPluginError as error: + return {REFUSED: str(error), REFUSAL_KIND: "refused"} + + +# E1 phase 5: a background cut is serialized (torch.save), hashed and committed in a +# writer process of this rank, not in a thread of the trainer process: the writer +# thread competed with training (each trained token about 8 % slower while it ran, +# GPU runs e1g/e1h). The trainer copies the state into shared memory and hands it +# over; the writer process writes the same file and summary. "0" keeps the thread. +WRITER_PROCESS_ENV = "YETO_RL_CUT_WRITER_PROCESS" +_WRITER: list[Any] = [] # this rank's writer process pool (one worker), started lazily + + +def _writer_enabled() -> bool: + return os.environ.get(WRITER_PROCESS_ENV, "1").strip().lower() not in {"0", "false", "no", "off"} + + +def _writer_init() -> None: + """Writer process start: below training in CPU priority, one intra-op thread.""" + import torch + + try: + os.nice(10) + except OSError: + pass + torch.set_num_threads(1) + + +def _writer_pool() -> Any: + if not _WRITER: + import multiprocessing + from concurrent.futures import ProcessPoolExecutor + + import torch.multiprocessing # noqa: F401 - registers the shared-memory tensor pickling + + # spawn: the trainer has CUDA initialized (fork would copy it); the child never touches CUDA + _WRITER.append(ProcessPoolExecutor(max_workers=1, mp_context=multiprocessing.get_context("spawn"), + initializer=_writer_init)) + import atexit + + atexit.register(_drop_writer_pool) # stop it before interpreter teardown + return _WRITER[0] + + +def _drop_writer_pool() -> None: + if _WRITER: + _WRITER.pop().shutdown(wait=False, cancel_futures=True) + + +def _remove_partial(directory: str, coord: Mapping[str, int]) -> None: + """Remove what a dead writer process left of this rank's shard (no summary was + returned, so no manifest names it).""" + final = os.path.join(directory, shard_name(coord)) + for path in (final + ".tmp", final): + if os.path.exists(path): + os.remove(path) + + +def _write_background(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: + """Thread of the trainer process: device copies to shared host memory, then the + writer process serializes, hashes and summarizes them. When the writer process + fails (it died, the state did not pickle, ...) this thread writes the same shard + itself and the summary's ``writer`` says so (``mode=thread``, ``fallback``).""" + try: + host = _share(_on_host(captured)) except CutPluginError as error: return {REFUSED: str(error), REFUSAL_KIND: "refused"} + fallback = "disabled" + if _writer_enabled(): + try: + result = _writer_pool().submit(_write_refusing, host, directory=directory, cut_id=cut_id).result() + return {**result, "writer": {"mode": "process"}} + except Exception as error: # noqa: BLE001 - fall back to this thread, say why + _drop_writer_pool() + fallback = f"writer process failed: {error!r}"[:500] + print("yeto cut writer fallback " + repr({"cut_id": cut_id, "reason": fallback}), flush=True) + _remove_partial(directory, host["coord"]) + return {**_write_refusing(host, directory=directory, cut_id=cut_id), + "writer": {"mode": "thread", "fallback": fallback}} + + +def _share(value: Any) -> Any: + """``value`` with every CPU tensor moved into shared memory (in place), so the + writer process maps it instead of receiving a copy through the pipe.""" + import torch + + if isinstance(value, torch.Tensor): + return value.share_memory_() if value.device.type == "cpu" else value + if isinstance(value, dict): + return {k: _share(v) for k, v in value.items()} + if isinstance(value, list): + return [_share(v) for v in value] + if isinstance(value, tuple): + return tuple(_share(v) for v in value) + return value -def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: +def _capture(actor: Any, *, cut_id: str, device_copy: bool = False) -> dict[str, Any]: + """Everything the shard and its summary need, read at the step boundary. + + Host path: CPU copies, taken here. ``device_copy`` (E1, background cuts): the + adapter and optimizer tensors are copied device-to-device here and reach the + host in the writer thread (:func:`_on_host`), so the trainer waits only for + the device copies. The device path is skipped (host path, reason in + ``capture``) without CUDA parameters, with too little free device memory, or + when a device copy runs out of memory; the shard is the same either way.""" import torch backend = _backend(actor) coord = _require(actor, backend) if _STEP_GRAD_NORMS or _STEP_APPLIED_LRS or _STEP_LOSSES: raise CutPluginError("per-step records not drained: the cut is not at an optimizer-step boundary") + import json + import time + + marks = [time.monotonic()] named = _adapters(actor, backend) + on_device, reason = _device_copy_plan(actor, named) if device_copy else (False, "inline") with torch.no_grad(): - snap = _snapshot(actor, backend, named) + snap = None + if on_device: + try: + snap = _snapshot(actor, backend, named, on_device=True) + except torch.cuda.OutOfMemoryError: + snap, on_device, reason = None, False, "device_oom" + torch.cuda.empty_cache() + if snap is None: + snap = _snapshot(actor, backend, named) + ready = None + if on_device and any(p.is_cuda for _, p in named): + ready = torch.cuda.Event() + ready.record() + marks.append(time.monotonic()) rng = backend.capture_rng() + capture = {"mode": "device" if on_device else "host", "reason": reason} + captured = {"coord": coord, "snap": snap, "rng": rng, "capture": capture, "ready": ready, + "scheduler_samples": int(actor.opt_param_scheduler.num_steps), + "tuning": _tuning(backend, named), "adapter_names": sorted(n for n, _ in named), + # what resume_rank_state reads from the live trainer, as it is now + "resume_inputs": resume_inputs(actor)} + marks.append(time.monotonic()) + if not on_device: + captured["resume"] = resume_rank_state(actor, snap["optimizer_named"], **captured["resume_inputs"]) + marks.append(time.monotonic()) + # E1 diagnosis: where the step-boundary capture of a cut goes (driver log only) + print("yeto cut capture timing " + json.dumps({**dict(zip( + ("snapshot_s", "rng_tuning_s", "resume_state_s"), + (round(b - a, 3) for a, b in zip(marks, marks[1:])))), "capture": capture}), flush=True) + return captured + + +def _on_host(captured: dict[str, Any]) -> dict[str, Any]: + """``captured`` with its device copies on the host (runs off the training + thread). Host captures are returned as they are.""" + ready = captured.get("ready") + if ready is None and "resume" in captured: + return captured + if ready is not None: + ready.synchronize() # the device copies are complete (whatever stream training uses) + return {**captured, "snap": to_host(captured["snap"]), "ready": None} + + +def _with_resume(captured: dict[str, Any]) -> dict[str, Any]: + """``captured`` with the N7.6 resume state a device capture defers (computed from + the host snapshot and the step-boundary resume inputs).""" + if "resume" in captured: + return captured + return {**captured, "resume": resume_rank_state(None, captured["snap"]["optimizer_named"], + **captured["resume_inputs"])} + + +def _write(captured: dict[str, Any], *, directory: str, cut_id: str) -> dict[str, Any]: + import torch + + coord, snap, rng = captured["coord"], captured["snap"], captured["rng"] shard = {"schema": SHARD_SCHEMA, "cut_id": cut_id, "coord": dict(coord), **snap, "rng": rng} os.makedirs(directory, exist_ok=True) name = shard_name(coord) @@ -637,13 +959,13 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: import io buffer = io.BytesIO() - torch.save(to_safe(shard), buffer) + torch.save(_interned(to_safe(shard)), buffer) data = buffer.getvalue() fh.write(data[: len(data) // 2]) fh.flush() os.fsync(fh.fileno()) cut_injection.kill_now(f"save_cut shard {name} half written") - torch.save(to_safe(shard), fh) + torch.save(_interned(to_safe(shard)), fh) fh.flush() os.fsync(fh.fileno()) os.replace(tmp, final) @@ -652,7 +974,7 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: "sha256": _sha256(final), "bytes": os.path.getsize(final), "coord": dict(coord), - "scheduler_samples": int(actor.opt_param_scheduler.num_steps), + "scheduler_samples": captured["scheduler_samples"], "has_optimizer_state": snap["optimizer_named"] is not None, "has_rng": rng is not None, "state_digest": state_digest(snap), @@ -660,14 +982,33 @@ def _save(actor: Any, *, directory: str, cut_id: str) -> dict[str, Any]: "weight_version": (snap.get("miles_counters") or {}).get("weight_version"), "rng_digest": state_digest(rng), "components": component_digests({"state": snap, "rng": rng}), - "tuning": _tuning(backend, named), - "adapter_tensors": len(named), - "adapter_names": sorted(n for n, _ in named), + "tuning": captured["tuning"], + "adapter_tensors": len(captured["adapter_names"]), + "adapter_names": captured["adapter_names"], "optimizer_names": sorted((snap["optimizer_named"] or {}).get("entries", {})), - "resume": resume_rank_state(actor, snap["optimizer_named"]), + "resume": captured["resume"], } +def _interned(value: Any) -> Any: + """``value`` with every string interned. Pickle writes a string once per object + and back-references repeats of the same object; interning makes that depend on + the values only, so a shard serialized in the writer process (whose strings + arrived unpickled, as separate objects) has the same bytes as one serialized in + the trainer (E1 phase 5).""" + import sys + + if isinstance(value, str): + return sys.intern(value) + if isinstance(value, dict): + return {_interned(k): _interned(v) for k, v in value.items()} + if isinstance(value, list): + return [_interned(v) for v in value] + if isinstance(value, tuple): + return tuple(_interned(v) for v in value) + return value + + def _load_verified(directory: str, entry: Mapping[str, Any], cut_id: str) -> dict[str, Any]: import torch diff --git a/yeto/rl/adapters/miles/full_parameter/role_stream.py b/yeto/rl/adapters/miles/full_parameter/role_stream.py index e1c0f6b2..549743f7 100644 --- a/yeto/rl/adapters/miles/full_parameter/role_stream.py +++ b/yeto/rl/adapters/miles/full_parameter/role_stream.py @@ -20,18 +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 @@ -78,10 +88,36 @@ 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 max_chunk_bytes: int = 256 << 20 + # 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); the local weights lag the + # global policy by one outer step (yeto.rl.engine.sync_overlap). + 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(): @@ -126,6 +162,14 @@ 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: + 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 ( @@ -219,7 +263,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", + "wire_dtype": config.wire_dtype_label, + **({"sync_train_overlap": True} if config.overlap else {}), }, } encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode() @@ -261,6 +306,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 @@ -334,8 +384,11 @@ async def initialize(self, group) -> None: self.config.syncer_addr, self.config.learner_id, adapter.layout.fragments, - dtype=DTYPE_F32, + 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, @@ -347,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: @@ -390,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 @@ -399,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: @@ -418,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( @@ -471,7 +527,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 @@ -500,14 +563,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 @@ -559,8 +627,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) @@ -573,14 +648,60 @@ 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``). + + 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.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.submission_cut, 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 @@ -610,16 +731,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, @@ -911,8 +1034,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 @@ -934,6 +1058,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/fully_async_prompts.py b/yeto/rl/adapters/miles/fully_async_prompts.py index b448cdbd..bf1fbd95 100644 --- a/yeto/rl/adapters/miles/fully_async_prompts.py +++ b/yeto/rl/adapters/miles/fully_async_prompts.py @@ -9,11 +9,13 @@ * round accounting: groups drawn since the last report and Miles' stale-group count (reported unknown, with a warning, when Miles does not expose it); -* cuts (decision P0-S-Q3): :func:`requeue_untrained` stops the background - worker and puts every drawn, untrained group back into the data source - buffer, reset for a fresh generation. The cut then carries them like any - buffered group, so the resumed run generates them again; nothing in flight - is saved. +* cuts (decision P0-S-Q3): background generation is never saved; every drawn, + untrained group goes into the cut reset for a fresh generation, so the + resumed run generates it again. :func:`snapshot_untrained` (E1, the cut path) + records them, the data source buffer and the data cursor in one step on the + worker's event loop and leaves the worker running: the live run keeps + generating ahead across the cut. :func:`requeue_untrained` (the older path) + stops the worker and puts the groups back into the live data source. Runs inside the rollout executor actor (``__ray_call__``). Import-light. """ @@ -117,6 +119,51 @@ def current() -> PromptLedger | None: return _LEDGER +def _reset_copies(groups: Iterable[Sequence[Any]]) -> list[list[Any]]: + copies = [copy.deepcopy(list(group)) for group in groups] + for group in copies: + for sample in group: + sample.reset_for_retry() + return copies + + +def snapshot_untrained(executor: Any, *, timeout_s: float = WORKER_STOP_TIMEOUT_S) -> dict[str, Any]: + """The cut's view of the fully_async prompts, taken without stopping the worker. + + Returns ``groups`` (the data source buffer, then every drawn, untrained group + as a reset copy, in group order) and ``data_cursor`` (the data source position, + ``rollout_meta_hook.data_cursor``). Both are read in ONE step on the worker's + event loop, where the worker draws prompts, so no group is both behind the + cursor and missing from ``groups``, or counted twice. A resume restores the + cursor and puts ``groups`` into the buffer: exactly the prompts the run had + drawn but not trained are generated again. The live data source and the + worker are left untouched.""" + from .rollout_meta_hook import data_cursor + + source = getattr(executor, "data_source", None) + source = getattr(source, "__self__", source) + ledger = getattr(source, _LEDGER_ATTR, None) + if ledger is None: + raise PromptLedgerError("fully_async prompts were not tracked; the cut would lose them") + + def read() -> dict[str, Any]: + cursor, _length = data_cursor(source) + buffered = [list(group) for group in (getattr(source, "buffer", None) or [])] + untrained = [group for _, group in sorted(ledger.untrained.items())] + return {"groups": _reset_copies(buffered + untrained), "data_cursor": cursor, + "untrained_groups": len(untrained)} + + worker = getattr(getattr(executor, "generate_rollout", None), "_worker", None) + loop = None if worker is None or worker.done() else worker.get_loop() + if loop is None or not loop.is_running() or getattr(loop, "_thread_id", None) == threading.get_ident(): + return read() # no concurrent draw: nothing else runs this loop now + + async def on_loop() -> dict[str, Any]: + return read() + + return asyncio.run_coroutine_threadsafe(on_loop(), loop).result(timeout=timeout_s) + + def requeue_untrained(executor: Any, *, timeout_s: float = WORKER_STOP_TIMEOUT_S) -> int: """Stop the background worker and return every drawn, untrained group to the data source, reset for a fresh generation. Returns the group count. @@ -134,10 +181,7 @@ def requeue_untrained(executor: Any, *, timeout_s: float = WORKER_STOP_TIMEOUT_S raise PromptLedgerError("fully_async prompts were not tracked; the cut would lose them") rollout_fn = getattr(executor, "generate_rollout", None) _stop_worker(rollout_fn, timeout_s) - groups = [copy.deepcopy(group) for _, group in sorted(ledger.untrained.items())] - for group in groups: - for sample in group: - sample.reset_for_retry() + groups = _reset_copies(group for _, group in sorted(ledger.untrained.items())) if groups: ledger.data_source.add_samples(groups) # the ledger forgets them if rollout_fn is not None: diff --git a/yeto/rl/adapters/miles/in_flight.py b/yeto/rl/adapters/miles/in_flight.py index 44a32207..4e896ee2 100644 --- a/yeto/rl/adapters/miles/in_flight.py +++ b/yeto/rl/adapters/miles/in_flight.py @@ -132,18 +132,23 @@ def _codex_suspension() -> tuple[Mapping[str, Any], Mapping[str, Any]]: def export_in_flight(executor: Any) -> dict[str, Any]: """Runs INSIDE the rollout executor actor: the in-flight entries and the - data buffer length they account for.""" + data buffer length they account for (fully_async: also ``data_cursor``, the + data source position read together with them).""" + gates, inflight = _codex_suspension() if getattr(getattr(executor, "args", None), "fully_async", False): - # P0-S decision 3: background generation is discarded, never saved; - # its prompts go back into the data buffer (reset), so the cut carries - # them like any buffered group and the resumed run generates them again. - from .fully_async_prompts import requeue_untrained - - requeue_untrained(executor) + # P0-S decision 3: background generation is never saved; the cut carries + # every drawn, untrained prompt (reset) like a buffered group, so the + # resumed run generates it again. E1: read without stopping the worker + # (the live run keeps generating ahead); the cursor comes from the same + # read, so the cut's cursor and groups agree. + from .fully_async_prompts import snapshot_untrained + + snap = snapshot_untrained(executor) + return {"entries": buffered_entries(snap["groups"]) + agentic_entries(gates, inflight), + "buffer_groups": len(snap["groups"]), "data_cursor": snap["data_cursor"]} source = getattr(executor, "data_source", None) source = getattr(source, "__self__", source) buffer = list(getattr(source, "buffer", None) or []) - gates, inflight = _codex_suspension() entries = buffered_entries(buffer) + agentic_entries(gates, inflight) return {"entries": entries, "buffer_groups": len(buffer)} diff --git a/yeto/rl/adapters/miles/island_entry.py b/yeto/rl/adapters/miles/island_entry.py index 48585224..dbafe785 100644 --- a/yeto/rl/adapters/miles/island_entry.py +++ b/yeto/rl/adapters/miles/island_entry.py @@ -75,11 +75,11 @@ def parse_args(argv=None): ) if (args.parameter_mode == "full" and args.sync_preset not in FULL_PARAMETER_SYNC_PRESETS - and not _buffer_island(args)): + and not _no_sync_island(args)): parser.error( "--parameter-mode full requires a full-parameter --sync-preset " - f"({', '.join(sorted(FULL_PARAMETER_SYNC_PRESETS))}) or a no-sync " - "--rl-sample-role sampler / trainer island" + f"({', '.join(sorted(FULL_PARAMETER_SYNC_PRESETS))}) or " + "--rl-single-island-no-sync" ) if args.sync_preset == FULL_PARAMETER_SYNC_PRESET and not _trains_base_parameters(args): parser.error( @@ -151,7 +151,7 @@ def build_ports_launch(args, run_config, extra_argv=()): ) return translate_run_config(run_config, base_algorithm, extra_argv=tuple(extra_argv), max_policy_age=int(getattr(args, "rl_max_policy_age", 0) or 0), - full_parameter_sync=not _buffer_island(args)) + full_parameter_sync=not _no_sync_island(args)) _ELASTIC_COMPANIONS = rl_option_pairs( @@ -696,12 +696,11 @@ def _full_parameter_identity(model_config_sha256: str) -> tuple[str, str]: ) -def _buffer_island(args) -> bool: - """A no-sync sampler / trainer island (L4.6): it has no outer sync, so full - parameter mode runs without the full-parameter sync preset.""" - from yeto.rl.engine.sample_roles import BOTH, sample_role - - return bool(getattr(args, "rl_single_island_no_sync", False)) and sample_role(args) != BOTH +def _no_sync_island(args) -> bool: + """A no-sync island (a sampler / trainer island, L4.6, or a single island + with both roles, S26 E11): it has no outer sync, so full parameter mode runs + without the full-parameter sync preset.""" + return bool(getattr(args, "rl_single_island_no_sync", False)) def _full_parameter_mode(args, sync_preset: str) -> bool: @@ -716,9 +715,9 @@ def _full_parameter_mode(args, sync_preset: str) -> bool: full_mode = getattr(args, "parameter_mode", "lora") == "full" full_preset = sync_preset in FULL_PARAMETER_SYNC_PRESETS - if full_mode and not full_preset and not _buffer_island(args): + if full_mode and not full_preset and not _no_sync_island(args): raise ValueError("full parameter mode needs a full-parameter sync preset " - "(or a no-sync --rl-sample-role sampler / trainer island)") + "(or --rl-single-island-no-sync)") if full_preset and not _trains_base_parameters(args): raise ValueError(f"the {sync_preset} sync needs full parameter mode or expert-full tuning") if full_preset and not full_mode and sync_preset != FULL_PARAMETER_SYNC_PRESET: @@ -838,6 +837,14 @@ 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)), + # 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) @@ -1625,6 +1632,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. @@ -1751,7 +1792,7 @@ def run_miles( # full-parameter strict sweep also carries DeepSeek V4 expert-full (P0 1.7b.2). full_parameter_sync = sync_preset in FULL_PARAMETER_SYNC_PRESETS dense_full = sync_preset == FULL_PARAMETER_SYNC_PRESET - if _trains_base_parameters(args) and not full_parameter_sync and not _buffer_island(args): + if _trains_base_parameters(args) and not full_parameter_sync and not _no_sync_island(args): raise ValueError(f"expert-full tuning needs the {FULL_PARAMETER_SYNC_PRESET} sync") sao_runtime = None if dense_full: @@ -2098,6 +2139,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) @@ -2385,6 +2427,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/models/chunked_full_parameter.py b/yeto/rl/adapters/miles/models/chunked_full_parameter.py index 9b7e6c99..14a675d6 100644 --- a/yeto/rl/adapters/miles/models/chunked_full_parameter.py +++ b/yeto/rl/adapters/miles/models/chunked_full_parameter.py @@ -786,6 +786,96 @@ 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=" 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/publish.py b/yeto/rl/adapters/miles/publish.py index 990843ef..329b66bf 100644 --- a/yeto/rl/adapters/miles/publish.py +++ b/yeto/rl/adapters/miles/publish.py @@ -465,11 +465,16 @@ def _check_trainer_holds(self, state: Any) -> str: return tensor_hash def publish(self, state: TrainableState, *, token_rollout_id: int | None = None) -> PublicationResult: + started = time.monotonic() tensor_hash = self._check_trainer_holds(state) payload_hash, payload_bytes = _payload_of(state) version = state.policy_version if token_rollout_id is None else token_rollout_id token = policy_token(version, tensor_hash) + self.last_timing = {"check_s": time.monotonic() - started} members, checksums = self._runner.run(self._publish(version, token)) + # E1 diagnosis: where a publication's wall time goes (driver log only) + print("yeto publish timing " + json.dumps( + {k: round(v, 3) for k, v in self.last_timing.items()}, sort_keys=True), flush=True) self.last_engine_checksums = checksums bodies = list((checksums or {}).values()) self._reference = ( @@ -547,18 +552,24 @@ async def _publish( # the first (engines start resident); KV comes back after it. if offload_rollout and self._published_once: await self._controller.onload_weights() + mark = time.monotonic() await update_weights(self._args, self._actor, self._executor, self._controller) + self._time("update_weights_s", mark) if offload_rollout: await self._controller.onload_kv() except Exception as exc: raise PublicationError(f"upstream update_weights failed: {exc}", cause=PublicationCause.UPDATE_FAILED) from exc self._published_once = True + mark = time.monotonic() engines, ids, members = await self._stamp_version(version, token) + self._time("stamp_s", mark) checksums = None if self._verify_checksums: + mark = time.monotonic() bodies = await read_engine_checksums(self._controller, self._flatten) + self._time("checksum_s", mark) if len(bodies) != len(engines): raise PublicationError( f"{len(engines) - len(bodies)} engines returned no weight checksum" @@ -567,6 +578,11 @@ async def _publish( return members, checksums + def _time(self, key: str, mark: float) -> None: + timing = getattr(self, "last_timing", None) + if timing is not None: + timing[key] = time.monotonic() - mark + # -- convertible node (rl-trainer-start-stop L7.2) -------------------------- @property def published_checksum(self) -> str | None: diff --git a/yeto/rl/adapters/miles/rollout.py b/yeto/rl/adapters/miles/rollout.py index df2080d5..27edd5b3 100644 --- a/yeto/rl/adapters/miles/rollout.py +++ b/yeto/rl/adapters/miles/rollout.py @@ -361,6 +361,8 @@ def __init__(self, name: str = DEFAULT_SINK_ACTOR, *, timeout_s: float = 30.0) - self.timeout_s = timeout_s self._ray = ray self._actor = _sink_actor_class().options(name=name, get_if_exists=True).remote() + self._late: dict[int, dict[str, Any]] = {} + self._held: list[dict[str, Any]] = [] @property def sink_spec(self) -> str: @@ -379,10 +381,36 @@ def set_max_policy_age(self, limit: int) -> None: def set_published_policies(self, policies: Mapping[str, str]) -> None: self._ray.get(self._actor.set_policies.remote(json.dumps(dict(policies), sort_keys=True))) + def _take_items(self) -> list[dict[str, Any]]: + """Every pushed record not yet taken; late round records (E1) are kept aside + by rollout id, the others are returned (and forgotten) in push order.""" + from .rollout_meta_hook import LATE_META_SCHEMA + + for raw in self._ray.get(self._actor.take_all.remote()): + item = json.loads(raw) + if item.get("schema") == LATE_META_SCHEMA: + self._late[int(item["rollout_id"])] = item + else: + self._held.append(item) + items, self._held = self._held, [] + return items + + def take_late(self, rollout_id: int, *, timeout_s: float | None = None) -> dict[str, Any] | None: + """The late record of ``rollout_id`` (the deferred cross-version estimate), + waiting up to ``timeout_s``; None when it did not arrive.""" + deadline = time.monotonic() + (self.timeout_s if timeout_s is None else timeout_s) + while True: + self._held = self._take_items() # a next round's records stay for take() + if rollout_id in self._late: + return self._late.pop(rollout_id) + if time.monotonic() > deadline: + return None + time.sleep(0.05) + def take(self, rollout_id: int) -> dict[str, Any]: deadline = time.monotonic() + self.timeout_s while True: - items = [json.loads(i) for i in self._ray.get(self._actor.take_all.remote())] + items = self._take_items() main = [i for i in items if i.get("schema") != ROUND_META_SCHEMA] extra = [i for i in items if i.get("schema") == ROUND_META_SCHEMA] if items: @@ -658,6 +686,7 @@ def __init__( self._controller = inference_controller self._executor = rollout_executor self._metadata = metadata + self._deferred_estimates: set[int] = set() # E1: late cross-version estimates self._expected_policy = expected_policy self._run = (runner or LoopRunner()).run self._last_cursor: dict[str, int] | None = None @@ -781,6 +810,17 @@ def _in_executor(self, fn: Any, *fn_args: Any, **fn_kwargs: Any) -> Any: raise RuntimeError("Ray actor handle %s has no __ray_call__" % type(target).__name__) return self._run(_awaited(remote(fn, *fn_args, **fn_kwargs))) + def late_cross_version_fraction(self, rollout_id: int) -> float | None: + """E1: the deferred cross-version IS estimate of ``rollout_id`` (fully_async), + None when the round deferred none or it did not arrive in time.""" + take_late = getattr(self._metadata, "take_late", None) + if not callable(take_late) or rollout_id not in self._deferred_estimates: + return None + self._deferred_estimates.discard(rollout_id) + late = take_late(rollout_id) + value = (late or {}).get("cross_version_truncated_fraction") + return None if value is None else float(value) + def export_in_flight(self) -> dict[str, Any]: """agentic-rollout-utilization 3.3 (limit > 0): the cut's in-flight entries (buffered groups, suspended agentic trajectories); raises when @@ -834,9 +874,12 @@ def generate( self._prompt_ledger_installed = True self._run(self._controller.prepare_rollout(rollout_id)) data_pack = self._run(self._executor.get(rollout_id)) - self._offload_after_rollout() + if not self.defer_offload: + self._offload_after_rollout() try: payload = self._metadata.take(rollout_id) + if payload.get("cross_version_estimate") == "deferred": + self._deferred_estimates.add(int(rollout_id)) handle = handle_from_metadata( payload, rollout_id=rollout_id, @@ -1045,6 +1088,17 @@ def set_admission(self, open_: bool) -> None: harness admission to this pool's running members (no-op without a board).""" self._harness_admission("open_admission" if open_ else "close_admission", self.members()) + # E1: set by a driver that runs the --offload-rollout step in its own phase + # (offload_after_generate); other callers keep the offload inside generate + defer_offload = False + + @property + def offloads_after_generate(self) -> bool: + return self._args is not None and bool(getattr(self._args, "offload_rollout", False)) + + def offload_after_generate(self) -> None: + self._offload_after_rollout() + def _offload_after_rollout(self) -> None: """Upstream ``train.py`` serial colocated branch (--offload-rollout).""" diff --git a/yeto/rl/adapters/miles/rollout_meta_hook.py b/yeto/rl/adapters/miles/rollout_meta_hook.py index d34e761d..ea9e1f22 100644 --- a/yeto/rl/adapters/miles/rollout_meta_hook.py +++ b/yeto/rl/adapters/miles/rollout_meta_hook.py @@ -35,12 +35,26 @@ import os import statistics import tempfile +import threading from collections.abc import Iterable, Mapping, Sequence +from contextlib import contextmanager from pathlib import Path from typing import Any META_SINK_ENV = "YETO_ROLLOUT_META_SINK" DEFAULT_SINK_ACTOR = "yeto_rollout_meta" +# E1: per-step seconds of the last metadata hook run (printed by record_trained_groups) +_HOOK_TIMES: dict[str, float] = {} + + +def _timed(name: str, fn: Any, *args: Any, **kwargs: Any) -> Any: + import time + + started = time.monotonic() + try: + return fn(*args, **kwargs) + finally: + _HOOK_TIMES[name] = _HOOK_TIMES.get(name, 0.0) + time.monotonic() - started DEFAULT_SINK = f"ray:{DEFAULT_SINK_ACTOR}" METADATA_SCHEMA = "yeto-rollout-meta-v1" _TRAINED_ATTR = "_yeto_trained_group_keys" @@ -54,6 +68,9 @@ # rollout's metadata by the driver-side reader, so they are never attributed # to the next round. ROUND_META_SCHEMA = "yeto-rl-round-metadata-v1" +# E1: the cross-version IS estimate of a fully_async round, sent after the round's +# metadata (the prefill scoring runs in the background, off the rollout's critical path) +LATE_META_SCHEMA = "yeto-rl-late-round-metadata-v1" # buffer_stale_groups: groups the sample-buffer reader passed over this round # because they were older than max_policy_age (consumed, not trained) ROUND_METADATA_KEYS = frozenset({"nonzero_advantages", "buffer_stale_groups"}) @@ -136,15 +153,24 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: to its buffer (``sglang_rollout.py:505-510``); only samples aborted under ``--partial-rollout`` go back. """ + with one_sink_view(): + _record_trained_groups(args, data) + + +def _record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: + import time from yeto.rl.algos.sample_filters import apply_sample_filters from yeto.rl.harness.reasoning_loss_mask import apply_from_args as apply_reasoning_loss_policy + marks = [("start", time.monotonic())] apply_sample_filters(args, data) + marks.append(("sample_filters", time.monotonic())) # Side effect beyond filtering (S17 WP6): with --codex-exclude-reasoning-from-loss # this sets loss_mask 1 -> 0 on generated reasoning tokens of the kept samples # (before Miles converts them to train data). Default: returns at once, no sample touched. apply_reasoning_loss_policy(args, data, tokenizer_loader=_load_miles_tokenizer) + marks.append(("reasoning_mask", time.monotonic())) # S19 #8: a stop token forced by the engine's grammar mask (SGLang returns # logprob exactly 0.0 under the masked distribution) is not a model choice: # loss_mask 1 -> 0 before Miles converts the samples to train data. @@ -152,11 +178,13 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: setattr(args, _PLACEHOLDER_ATTR, mask_placeholder_logprobs(data, tokenizer_loader=_load_miles_tokenizer, args=args)) + marks.append(("placeholder_mask", time.monotonic())) # CompactionRL: the masks above may drop optimised tokens; the per-segment # counts (tokens_after, gae_length) must follow the final masks. from yeto.rl.harness.codex.codex_openenv_generate import refresh_segment_counts refresh_segment_counts([s for group in data for s in _flat(group)]) + marks.append(("segment_counts", time.monotonic())) setattr(args, _TRAINED_ATTR, {_group_key(group) for group in data}) # Full-parameter islands (P0-D8): bind the exact accepted batch, after the # final loss masks, for the role streams to read at the round boundary. @@ -176,7 +204,16 @@ def record_trained_groups(args: Any, data: Sequence[Sequence[Any]]) -> None: if ledger is None: raise RuntimeError("fully_async round without the yeto prompt ledger") ledger.mark_trained(data) - extract_rollout_metadata(args, data, round_counts=ledger.take_round_counts()) + counts = ledger.take_round_counts() + marks.append(("ledger", time.monotonic())) + extract_rollout_metadata(args, data, round_counts=counts) + marks.append(("round_metadata", time.monotonic())) + # E1: where the hook's time goes (it runs on the rollout's critical path) + detail = " ".join(f"{k}={v:.3f}" for k, v in sorted(_HOOK_TIMES.items())) + _HOOK_TIMES.clear() + print(f"[rl] sample filter hook: total={marks[-1][1] - marks[0][1]:.3f}s " + + " ".join(f"{name}={t - prev:.3f}" for (_, prev), (name, t) in zip(marks, marks[1:])) + + (f" [{detail}]" if detail else ""), flush=True) def _reward(args: Any, sample: Any) -> float: @@ -357,7 +394,8 @@ def batch_summary_enabled(args: Any) -> bool: def build_metadata( - args: Any, all_samples: Iterable[Sequence[Any]], sink: str | None = None + args: Any, all_samples: Iterable[Sequence[Any]], sink: str | None = None, *, + defer_estimate: bool = False, ) -> dict[str, Any]: trained = getattr(args, _TRAINED_ATTR, None) if trained is None: @@ -379,8 +417,8 @@ def build_metadata( if not samples: continue if rollout_id is None: - rollout_id = current_round_id(samples, sink) - record = group_record(args, group) + rollout_id = _timed("round_id", current_round_id, samples, sink) + record = _timed("group_record", group_record, args, group) tool_wait += sum(float(getattr(x, "non_generation_time", 0.0) or 0.0) for x in samples) aborted += int(record.pop("aborted")) key = tuple(record.pop("_key")) @@ -418,7 +456,9 @@ def build_metadata( if limit > 0: # agentic-rollout-utilization 4.1: carry-over accounting (limit > 0 only; # the default key set is unchanged). - payload.update(carry_over_fields(args, trained_samples, sink)) + payload.update(_timed("carry_over", carry_over_fields, args, trained_samples, sink, + defer_estimate_for=rollout_id if defer_estimate and rollout_id is not None + else None)) # 5.1/5.3: agentic suspension between turns (absent when not in use) payload.update(_carry.suspend_fields(args, all_samples)) if tool_wait > 0: @@ -453,10 +493,17 @@ def build_metadata( def carry_over_fields(args: Any, trained_samples: Sequence[Any], - sink: str | None = None, scorer: Any = None) -> dict[str, Any]: + sink: str | None = None, scorer: Any = None, *, + defer_estimate_for: int | None = None) -> dict[str, Any]: """4.1 rollout-metadata fields under a limit > 0: what the buffer filter carried in / discarded this round, the trained cross-version tokens and the estimated - cross-version IS truncated fraction (:mod:`.carry_over`).""" + cross-version IS truncated fraction (:mod:`.carry_over`). + + ``defer_estimate_for`` (E1, a rollout id): the estimate (one prefill of every + cross-version sample, about 1 s on A10G) runs in a background thread and goes + to the sink as a :data:`LATE_META_SCHEMA` record; the fields here then say + ``cross_version_estimate: "deferred"``. Its only reader, the policy-age + governor, runs after the train step.""" token = current_policy_token(sink) version = token_version(token) if token else None fields: dict[str, Any] = {"max_policy_age": _carry.max_policy_age(args)} @@ -467,12 +514,31 @@ def carry_over_fields(args: Any, trained_samples: Sequence[Any], fields["cross_version_tokens"] = crossed fields["trained_response_tokens"] = total if crossed: - fields.update(_carry.estimate_cross_version_truncation( - args, trained_samples, version, - scorer if scorer is not None else _carry.router_prefill_scorer(args))) + scorer = scorer if scorer is not None else _carry.router_prefill_scorer(args) + if defer_estimate_for is None: + fields.update(_carry.estimate_cross_version_truncation(args, trained_samples, version, scorer)) + else: + _estimate_in_background(args, list(trained_samples), version, scorer, sink, + rollout_id=int(defer_estimate_for)) + fields["cross_version_estimate"] = "deferred" return fields +def _estimate_in_background(args: Any, samples: list[Any], version: int, scorer: Any, + sink: str | None, *, rollout_id: int) -> None: + import threading + + def run() -> None: + try: + late = _carry.estimate_cross_version_truncation(args, samples, version, scorer) + except Exception as exc: # noqa: BLE001 - an estimate: unknown, never guessed + late = {"cross_version_truncated_fraction": None, + "cross_version_estimate_error": f"{type(exc).__name__}: {str(exc)[:200]}"} + put_to_sink({"schema": LATE_META_SCHEMA, "rollout_id": rollout_id, **late}, sink) + + threading.Thread(target=run, name=f"yeto-xv-estimate-{rollout_id}", daemon=True).start() + + # IR-3/IR-4 sample-metadata keys written by agentic generate code (codex-harness # codex_openenv_generate): summed per rollout into the metadata payload. # Neutral names and values live in the core port (decoupling 4.11); re-exported. @@ -751,9 +817,7 @@ def put_to_sink(payload: dict[str, Any], sink: str | None = None) -> None: os.replace(tmp, directory / name) return if kind == "ray": - import ray - - ray.get(ray.get_actor(target or DEFAULT_SINK_ACTOR).put.remote(encoded)) + _sink_call(target, "put", encoded) return raise ValueError(f"unknown rollout metadata sink {sink!r}") @@ -765,9 +829,50 @@ def _sink_target(sink: str | None) -> tuple[str, str]: return kind, target +# E1: inside one hook run (:func:`one_sink_view`) each sink value is read once and +# the sink actor is looked up once (ray.get_actor is a GCS round trip per call) +_SINK_VIEW = threading.local() + + +def _sink_call(target: str, actor_method: str, *call_args: Any) -> Any: + import ray + + name = target or DEFAULT_SINK_ACTOR + memo = getattr(_SINK_VIEW, "memo", None) + actor = None if memo is None else memo.get(("actor", name)) + if actor is None: + actor = ray.get_actor(name) + if memo is not None: + memo[("actor", name)] = actor + return ray.get(getattr(actor, actor_method).remote(*call_args)) + + +@contextmanager +def one_sink_view(): + """Within this block every sink value (policy token, published map, policy-age + limit) is read at most once per (sink, value) and reused (E1: the metadata + hook read the limit from the sink actor once per sample, ~70 actor calls a + round). A value another process writes while a block runs is seen from the + next block on (before, two reads inside one hook run could see two values; + now one run sees one view). A write from this thread inside the block drops + that value from the view. The sink actor handle is looked up once per block. + Nested blocks share the outer view.""" + if getattr(_SINK_VIEW, "memo", None) is not None: + yield + return + _SINK_VIEW.memo = {} + try: + yield + finally: + _SINK_VIEW.memo = None + + def _sink_write(sink: str | None, file_name: str, text: str, actor_method: str, value: Any) -> None: """Driver side: store one value (atomic file under ``dir:``, actor call under ``ray:``).""" kind, target = _sink_target(sink) + memo = getattr(_SINK_VIEW, "memo", None) + if memo is not None: + memo.pop((kind, target, file_name), None) if kind == "dir": directory = Path(target) directory.mkdir(parents=True, exist_ok=True) @@ -776,21 +881,24 @@ def _sink_write(sink: str | None, file_name: str, text: str, actor_method: str, handle.write(text) os.replace(tmp, directory / file_name) return - import ray - - actor = ray.get_actor(target or DEFAULT_SINK_ACTOR) - ray.get(getattr(actor, actor_method).remote(value)) + _sink_call(target, actor_method, value) def _sink_read(sink: str | None, file_name: str, actor_method: str) -> Any: """Rollout side: the stored value (file text under ``dir:``), None when unset.""" kind, target = _sink_target(sink) + memo = getattr(_SINK_VIEW, "memo", None) + key = (kind, target, file_name) + if memo is not None and key in memo: + return memo[key] if kind == "dir": path = Path(target) / file_name - return path.read_text(encoding="utf-8") if path.exists() else None - import ray - - return ray.get(getattr(ray.get_actor(target or DEFAULT_SINK_ACTOR), actor_method).remote()) + value = path.read_text(encoding="utf-8") if path.exists() else None + else: + value = _sink_call(target, actor_method) + if memo is not None: + memo[key] = value + return value def put_policy_token(token: str, sink: str | None = None) -> None: @@ -870,7 +978,11 @@ def policy_buffer_filter(args: Any, _rollout_id: Any, buffer: list, num_samples: Groups from any other policy (or incomplete ones) are dropped from the buffer, exactly as legacy filtered its completed-group queue. """ + with one_sink_view(): + return _policy_buffer_filter(args, _rollout_id, buffer, num_samples) + +def _policy_buffer_filter(args: Any, _rollout_id: Any, buffer: list, num_samples: int) -> list: token = current_policy_token() limit = _carry.max_policy_age(args) if limit > 0: @@ -975,15 +1087,24 @@ def extract_rollout_metadata( """``--rollout-all-samples-process-path`` hook (under fully_async: called by :func:`record_trained_groups` with the trained groups).""" + with one_sink_view(): + _extract_rollout_metadata(args, all_samples, data_source, round_counts=round_counts) + + +def _extract_rollout_metadata(args: Any, all_samples: Any, data_source: Any, *, + round_counts: Mapping[str, Any] | None) -> None: try: - payload = build_metadata(args, all_samples) + # fully_async (round_counts): the trainer is waiting for this batch; the + # cross-version estimate is read only after its train step (E1) + payload = _timed("build", build_metadata, args, all_samples, + defer_estimate=round_counts is not None) if elastic_metadata_enabled(args): cursor, buffer_length = data_cursor(data_source) if cursor is not None: payload["data_cursor"] = cursor if buffer_length is not None: payload["buffer_length"] = buffer_length - submitted = submitted_groups(args, data_source) + submitted = _timed("submitted", submitted_groups, args, data_source) if round_counts is not None: from yeto.rl.engine.overlap import ROLLOUT_OVERLAP_MODE @@ -1026,7 +1147,7 @@ def extract_rollout_metadata( payload["over_sampling"] = over_sampling if hasattr(args, "rollout_over_sampling_stats"): args.rollout_over_sampling_stats = None - put_to_sink(payload) + _timed("put", put_to_sink, payload) finally: # Reset per-rollout state: the bounded filter keys its memo on # ``yeto_rl_policy_version`` which legacy advanced per round; here the diff --git a/yeto/rl/adapters/miles/run_flags.py b/yeto/rl/adapters/miles/run_flags.py index 8a9f5202..fb73c1c0 100644 --- a/yeto/rl/adapters/miles/run_flags.py +++ b/yeto/rl/adapters/miles/run_flags.py @@ -175,6 +175,12 @@ "at most this many tokens per GPU (must be >= --seq-len). Default: fixed " "micro batches", ), + shared_flag( + "--rl-gradient-accumulation-fusion", + action="store_true", + help="ports: Megatron gradient-accumulation fusion (the Miles recipes' " + "default). Off by default (--no-gradient-accumulation-fusion)", + ), shared_flag( "--rl-optimizer-offload-fraction", type=float, @@ -313,6 +319,30 @@ ), island_flag("--fragments", type=int, default=1), island_flag("--parameter-layout-sha256", default=None), + island_flag( + "--sweep-push-dtype", + choices=("f32", "bf16", "int8"), + default="f32", + help="full-parameter sweep: wire dtype of pushed outer deltas (syncer masters stay f32)", + ), + island_flag( + "--sweep-push-int8-block", + type=int, + default=128, + help="full-parameter sweep, int8 pushes: values per f32 absmax scale", + ), + island_flag( + "--sweep-overlap", + action="store_true", + help="full-parameter sweep: train the next round while the sweep syncs, then " + "rebase onto the new global weights (rollout policy age stays 0)", + ), + 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, @@ -392,6 +422,11 @@ 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("--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/adapters/miles/state.py b/yeto/rl/adapters/miles/state.py index 175296cf..49120156 100644 --- a/yeto/rl/adapters/miles/state.py +++ b/yeto/rl/adapters/miles/state.py @@ -20,6 +20,7 @@ from . import LoopRunner from .state_plugin import ( APPLY_STATE, + COLLECT_DIGEST, EXPORT_DIGEST, EXPORT_FULL_DIGEST, EXPORT_STATE, @@ -60,6 +61,9 @@ def require_run_plugin(controller_cls: Any = None) -> None: class MilesPolicyState: """Export/apply the LoRA trainable state of the single-cell actor group.""" + # E1: export_digest(deferred=True) is supported (LoRA digest path) + supports_deferred_digest = True + def __init__( self, *, @@ -117,25 +121,39 @@ async def aexport(self, *, policy_version: int | None = None) -> TrainableState: results = await self._aplugin(EXPORT_STATE, {"policy_version": version}) return self._export_result(version, results) - def export_digest(self, *, policy_version: int | None = None): + def export_digest(self, *, policy_version: int | None = None, deferred: bool = False): """rl-publish-fastpath: the trainer exports and hashes; only the digests cross Ray. Returns a :class:`TrainerResidentState` whose tensors stay in the trainer - (``materialize`` falls back to :meth:`export`).""" + (``materialize`` falls back to :meth:`export`). + + ``deferred`` (E1): the trainer returns once the tensors are staged and hashes + them in the background (e.g. while it offloads); the state's digest is a + :class:`~yeto.rl.engine.policy_digest.DeferredDigest` collected on first read.""" import time version = self._export_version(policy_version) started = time.monotonic() - results = self._plugin(EXPORT_DIGEST, self._digest_kwargs(version)) + kwargs = {**self._digest_kwargs(version), **({"deferred": True} if deferred else {})} + results = self._plugin(EXPORT_DIGEST, kwargs) state = self._digest_result(version, results) export_s = next((r.get("export_seconds") for r in results if r is not None), None) # Timing for the next GPU run (log only; no tape field changes): total plugin # round trip vs the trainer-side hashing part. + hashed = "deferred" if deferred else f"{state.digest.hash_seconds:.1f}s" print(f"[rl] publish-fastpath digest export v{version}: " f"total={time.monotonic() - started:.1f}s trainer_export={export_s}s " - f"hash={state.digest.hash_seconds:.1f}s bytes={state.digest.payload_bytes}", flush=True) + f"hash={hashed} tensors={len(state.digest.specs)}", flush=True) return state + def _collect_digest(self, version: int): + from yeto.rl.engine.policy_digest import PolicyDigest + + results = [r for r in self._plugin(COLLECT_DIGEST, {"policy_version": version}) if r is not None] + if len(results) != 1 or results[0].get("policy_version") != version: + raise PolicyStateError(f"expected exactly one main-rank digest of v{version}, got {results!r:.300}") + return PolicyDigest.from_wire(results[0]["digest"]) + def weights_version(self) -> dict[str, Any]: """The main trainer rank's weights version ``{process_id, version}`` (D3).""" @@ -194,7 +212,12 @@ def _digest_result(self, version: int, results: list[Any]): (result,) = results if result.get("policy_version") != version: raise PolicyStateError("exported policy version mismatch") - digest = PolicyDigest.from_wire(result["digest"]) + if result.get("deferred"): + from yeto.rl.engine.policy_digest import DeferredDigest + + digest = DeferredDigest(result["specs"], lambda: self._collect_digest(version)) + else: + digest = PolicyDigest.from_wire(result["digest"]) layout_hash = layout_hash_of(digest.specs) if self._expected_layout_hash is not None and layout_hash != self._expected_layout_hash: raise ValueError( @@ -304,6 +327,8 @@ class MilesFullParameterLocalPolicyState(MilesPolicyState): apply. The layout hash is the full-parameter identity bound by run_miles. """ + supports_deferred_digest = False # its export_digest hashes every rank's shard eagerly + def __init__(self, *, actor_model: Any, base_model_revision: str, config_hash: str, expected_layout_hash: str | None, policy_version: int = 0, runner: LoopRunner | None = None) -> None: diff --git a/yeto/rl/adapters/miles/state_plugin.py b/yeto/rl/adapters/miles/state_plugin.py index 5125287e..312eaed7 100644 --- a/yeto/rl/adapters/miles/state_plugin.py +++ b/yeto/rl/adapters/miles/state_plugin.py @@ -39,6 +39,7 @@ EXPORT_STATE = f"{_PLUGIN_MODULE}.export_state" # rl-publish-fastpath: export + hash on the trainer, return only the digests. EXPORT_DIGEST = f"{_PLUGIN_MODULE}.export_digest" +COLLECT_DIGEST = f"{_PLUGIN_MODULE}.collect_digest" WEIGHTS_VERSION = f"{_PLUGIN_MODULE}.weights_version" # A no-sync full-parameter island (L4.6): per-rank digest of every model parameter. EXPORT_FULL_DIGEST = f"{_PLUGIN_MODULE}.export_full_digest" @@ -648,33 +649,92 @@ def export_state(actor: Any, *, policy_version: int) -> dict[str, Any] | None: return _export_state(actor, policy_version=policy_version) +def _digest_staging(actor: Any) -> Any: + """The pinned host buffer this trainer's digests reuse every round (E1). A + deferred digest still hashing it is waited for first.""" + from yeto.rl.engine.policy_digest import HostStaging + + _wait_deferred_digest(actor) + staging = getattr(actor, "_yeto_digest_staging", None) + if staging is None: + staging = actor._yeto_digest_staging = HostStaging() + return staging + + def export_digest( actor: Any, *, policy_version: int, base_model_revision: str, lora_config_hash: str, + deferred: bool = False, ) -> dict[str, Any] | None: """rl-publish-fastpath: :func:`export_state` on the trainer, but return only the publication digests (policy tensor hash, payload hash/bytes, specs) instead of the tensors. The hash definitions are those of ``yeto.rl.core.policy_tensor_hash`` - and ``publish.payload_digest`` (see ``policy_digest``); only where they run moves.""" + and ``publish.payload_digest`` (see ``policy_digest``); only where they run moves. + + ``deferred`` (E1, sync colocated mode): return once the tensors are staged on the + host (layout only, ``"deferred": True``) and hash them in a thread of this + process, so the trainer can offload meanwhile; :func:`collect_digest` returns + the digest (the same values).""" from yeto.rl.engine.policy_digest import digest_canonical_tensors, layout_hash_of import time started = time.monotonic() - exported = export_state(actor, policy_version=policy_version) + install_grad_norm_recorder() + with trainer_resident(actor): + # E1: the tensors are only hashed here, so they may be staged in one pinned buffer + exported = _export_state(actor, policy_version=policy_version, staging=_digest_staging(actor)) export_seconds = round(time.monotonic() - started, 3) if exported is None: return None tensors = exported["tensors"] specs = [(name, tuple(int(d) for d in value.shape)) for name, value in sorted(tensors.items())] - digest = digest_canonical_tensors( - tensors, - base_model_revision=base_model_revision, - lora_config_hash=lora_config_hash, - layout_hash=layout_hash_of(specs), - ) - return {"policy_version": int(policy_version), "digest": digest.to_wire(), - "export_seconds": export_seconds, "weights": current_weights_version()} + + def digest() -> Any: + return digest_canonical_tensors( + tensors, + base_model_revision=base_model_revision, + lora_config_hash=lora_config_hash, + layout_hash=layout_hash_of(specs), + ) + + out = {"policy_version": int(policy_version), "export_seconds": export_seconds, + "weights": current_weights_version()} + if not deferred: + return {**out, "digest": digest().to_wire()} + from concurrent.futures import ThreadPoolExecutor + + pool = ThreadPoolExecutor(max_workers=1, thread_name_prefix="yeto-digest-deferred") + actor._yeto_deferred_digest = (int(policy_version), pool.submit(digest)) + pool.shutdown(wait=False) + return {**out, "deferred": True, "specs": [[n, list(shape)] for n, shape in specs]} + + +def collect_digest(actor: Any, *, policy_version: int) -> dict[str, Any] | None: + """The digest of the deferred :func:`export_digest` of ``policy_version`` (waits + for its hashing thread). Reads no trainer state (the trainer may be offloaded); + ``None`` on ranks that export nothing.""" + pending = getattr(actor, "_yeto_deferred_digest", None) + if pending is None: + if not _is_main_rank(actor): + return None + raise StatePluginError(f"no deferred digest of policy v{policy_version} on this rank") + version, future = pending + actor._yeto_deferred_digest = None + if version != int(policy_version): + raise StatePluginError(f"the deferred digest is of policy v{version}, not v{policy_version}") + return {"policy_version": version, "digest": future.result().to_wire()} + + +def _wait_deferred_digest(actor: Any) -> None: + """The staged tensors of a deferred digest are views of the reused buffer: let an + uncollected hashing finish before the buffer is written again.""" + pending = getattr(actor, "_yeto_deferred_digest", None) + if pending is not None: + try: + pending[1].result() + except Exception: # noqa: BLE001 - collect_digest reports it; only the buffer matters here + pass FULL_SHARD_PREFIX = b"yeto-rl-full-parameter-shard-v1\0" @@ -912,10 +972,15 @@ def _tp_ep_leader(actor: Any) -> bool: ) -def _export_state(actor: Any, *, policy_version: int) -> dict[str, Any] | None: +def _export_state(actor: Any, *, policy_version: int, staging: Any = None) -> dict[str, Any] | None: + """``staging`` (a :class:`~yeto.rl.engine.policy_digest.HostStaging`): copy the + tensors into its pinned buffer with one GPU synchronization; the result is then + only valid until the next call with it (digest path, E1). Same values either way.""" import torch + from yeto.rl.engine.policy_digest import NonFiniteTensor, stage_f32_on_host + if is_native_flash_next(actor): return _export_flash_next(actor, policy_version=policy_version) bindings = adapter_bindings(actor) @@ -931,12 +996,10 @@ def _export_state(actor: Any, *, policy_version: int) -> dict[str, Any] | None: tensors[b.name] = b.to_hf(master) if not _is_main_rank(actor): return None - out = {} - for name, value in sorted(tensors.items()): - value = value.detach().to(device="cpu", dtype=torch.float32).contiguous().clone() - if not torch.isfinite(value).all().item(): - raise StatePluginError(f"{name!r} contains NaN or Inf") - out[name] = value + try: + out = stage_f32_on_host(tensors, staging) + except NonFiniteTensor as error: + raise StatePluginError(str(error)) from None assert_grad_flow_intact(actor.model, params) return {"policy_version": int(policy_version), "tensors": out} diff --git a/yeto/rl/adapters/miles/trainer.py b/yeto/rl/adapters/miles/trainer.py index 54c65ed9..ab360271 100644 --- a/yeto/rl/adapters/miles/trainer.py +++ b/yeto/rl/adapters/miles/trainer.py @@ -813,7 +813,15 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: a non-zero ``carried_over``/ready-unconsumed ledger, an unsettled outer commit or a missing data cursor. """ - from yeto.rl.engine.cut import CutError, CutFile, CutManifest, commit_manifest, context_problems, cut_dir + return self.finish_cut(self.start_cut(epoch=epoch, context=context, background=False)) + + def start_cut(self, *, epoch: int, context: "CutContext", background: bool = True) -> dict[str, Any]: + """E1: the part of :meth:`save_cut` that must run at the step boundary: the + context checks and every rank's state capture. With ``background`` the ranks + write and hash their shards in a thread while training goes on; :meth:`finish_cut` + (later, before the next cut) collects them and commits the manifest. Nothing + names the cut until the manifest is committed. Returns the pending cut.""" + from yeto.rl.engine.cut import CutError, context_problems, cut_dir from .cut_plugin import SAVE_CUT_SHARD, config_problems cut_id = context.cut_id @@ -828,11 +836,42 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: if problems: raise CutError("refusing to save a cut: " + "; ".join(problems)) directory = cut_dir(context.root, cut_id) - summaries = [ - dict(s) for s in self._run( - self._actor.run_plugin(SAVE_CUT_SHARD, {"directory": str(directory), "cut_id": cut_id}) - ) - ] + # the critic cut runs in finish_cut and reads the critic as it is then: inline only + background = bool(background) and self._critic is None + critic = self._save_critic_cut(directory / "critic", context.progress.policy_version) + kwargs = {"directory": str(directory), "cut_id": cut_id, **({"background": True} if background else {})} + results = [dict(s) for s in self._run(self._actor.run_plugin(SAVE_CUT_SHARD, kwargs))] + # E1: how each rank took its state (device copies or host path, and why) + capture = [dict(s["capture"]) for s in results if isinstance(s.get("capture"), Mapping)] + return {"epoch": int(epoch), "context": context, "runtime": runtime, "directory": directory, + "background": background, "results": results, "critic": critic, + **({"capture": capture} if capture else {})} + + def finish_cut(self, pending: Mapping[str, Any]) -> str: + """Collect the shards of :meth:`start_cut` and commit the manifest; returns the cut id.""" + return self.collect_cut(pending)() + + def collect_cut(self, pending: Mapping[str, Any]) -> Callable[[], str]: + """E1: the trainer part of :meth:`finish_cut` (collect and check every rank's + shard summary, build the manifest). Returns the commit of that manifest: plain + filesystem work (no trainer call), so the caller may run it on another thread. + Raises like :meth:`finish_cut` before anything is committed.""" + from yeto.rl.engine.cut import CutError, CutFile, CutManifest, commit_manifest + from .cut_plugin import FINISH_CUT_SHARD + + epoch, context, runtime = pending["epoch"], pending["context"], pending["runtime"] + cut_id = context.cut_id + summaries = list(pending["results"]) + if pending["background"]: + started = [s for s in summaries if "refused" not in s] + summaries = [s for s in summaries if "refused" in s] + if started: + summaries += [dict(s) for s in self._run( + self._actor.run_plugin(FINISH_CUT_SHARD, {"cut_id": cut_id}))] + # E1: where each rank's background shard was written (writer process, or this + # thread after a writer-process failure); reported with rl_cut_committed + self.last_cut_writer = [dict(s["writer"]) for s in summaries if isinstance(s.get("writer"), Mapping)] + problems: list[str] = [] refused = [f"[{s.get('refusal_kind', 'refused')}] {s['refused']}" for s in summaries if "refused" in s] if refused: raise CutError("rank refused the cut: " + "; ".join(sorted(set(refused)))) @@ -843,7 +882,7 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: if problems: raise CutError("refusing an incomplete cut: " + "; ".join(problems)) resume = self._resume_manifest(context, summaries) - critic = self._save_critic_cut(directory / "critic", context.progress.policy_version) + critic = pending.get("critic") if critic is not None: # 4.3; absent for every critic-free trainer (manifest unchanged) runtime = {**runtime, "critic": critic} manifest = CutManifest( @@ -869,8 +908,13 @@ def save_cut(self, *, epoch: int, context: "CutContext") -> str: ) if problems: raise CutError("refusing an incomplete cut: " + "; ".join(problems)) - commit_manifest(context.root, manifest, check_files=context.shared_filesystem) - return cut_id + root, check_files = context.root, context.shared_filesystem + + def commit() -> str: + commit_manifest(root, manifest, check_files=check_files) + return cut_id + + return commit def _resume_manifest(self, context: "CutContext", summaries: list[Mapping[str, Any]]) -> Any: """Resume manifest of this cut: run counters plus every rank's optimizer state (S25 N7.6).""" diff --git a/yeto/rl/adapters/verl/fully_async_ports.py b/yeto/rl/adapters/verl/fully_async_ports.py index 061fb920..f79505f9 100644 --- a/yeto/rl/adapters/verl/fully_async_ports.py +++ b/yeto/rl/adapters/verl/fully_async_ports.py @@ -15,8 +15,8 @@ * ``train_step`` -- old log-prob -> advantage -> actor update on the actor (verl's ``fit_step`` without ``_fit_update_weights``); mismatch criteria in the Miles convention on the trainer's DataProto. -* policy state -- ``ports_impl.VerlPolicyState`` unchanged: its two island hooks - (``export_tensors`` / ``call_workers``) run on the trainer actor's worker group. +* policy state -- ``ports_impl.VerlPolicyState`` unchanged: its island hooks + (``export_tensors`` / ``export_digest`` / ``call_workers``) run on the trainer actor's worker group. * ``publish`` -- the driver has applied the state; the actor pushes it through ``checkpoint_manager.update_weights(global_steps=param_version)``, resets the rollouter's staleness counter, and the pairing @@ -41,6 +41,7 @@ from .fully_async_round import (DISCARD_MECHANISM, FULLY_ASYNC_PLACEMENT, RoundCollector, SampleMeta, group_metadata) from .fully_async_translate import VersionMap +from .ports_impl import sent_checksum from .rollout_events import cutoff_fields MEMBER = "vllm-rollouter-0" @@ -84,6 +85,10 @@ def __init__(self, call: Callable[..., Any], *, learner_id: int, limit: int, req def export_tensors(self) -> tuple[dict, dict]: return self.call("yeto_export_lora") + def export_digest(self, kwargs: dict) -> dict: + """rl-publish-fastpath: digests + send checksum computed on the trainer actor.""" + return self.call("yeto_export_digest", kwargs) + def call_workers(self, fn, *args): return self.call("yeto_call_workers", fn, *args) @@ -251,8 +256,7 @@ def publish(self, state) -> PublicationResult: isl = self.island version = int(state.policy_version) - lora = state.to_lora() - sent = pub.lora_checksum(dict(lora.tensors)) + sent, payload_bytes = sent_checksum(state) (isl.readback_dir / "expected_version").write_text(str(version)) stale = isl.readback_dir / f"readback-v{version}.json" if stale.exists(): @@ -288,7 +292,6 @@ def publish(self, state) -> PublicationResult: digest = state.policy_tensor_hash() isl.published = (version, digest) isl.hashes[version] = digest - payload_bytes = sum(int(t.numel()) * 2 for t in lora.tensors.values()) manifest = InferencePublicationManifest( publication_mode="full", base_policy_version=None, target_policy_version=version, target_policy_hash=digest, target_manifest_hash=_sha(f"{version}:{digest}"), diff --git a/yeto/rl/adapters/verl/fully_async_runner.py b/yeto/rl/adapters/verl/fully_async_runner.py index 900ccf6d..b3273048 100644 --- a/yeto/rl/adapters/verl/fully_async_runner.py +++ b/yeto/rl/adapters/verl/fully_async_runner.py @@ -189,6 +189,11 @@ async def yeto_export_lora(self): return (await self.yeto_call_workers(remote_export_lora))[0] + async def yeto_export_digest(self, kwargs): + from .ports_impl import remote_export_digest + + return next(r for r in await self.yeto_call_workers(remote_export_digest, kwargs) if r is not None) + YetoFullyAsyncTrainer = ray.remote(num_cpus=10)(type("YetoFullyAsyncTrainer", (_YetoTrainerMixin, _TrainerBase), {})) diff --git a/yeto/rl/adapters/verl/image.py b/yeto/rl/adapters/verl/image.py index b096d1b1..29a52282 100644 --- a/yeto/rl/adapters/verl/image.py +++ b/yeto/rl/adapters/verl/image.py @@ -46,15 +46,18 @@ def image_env() -> dict[str, str]: } -def modal_image(modal, commit: str = VERL_COMMIT): - patch = Path(__file__).with_name("patch_verl.py") - return ( +def modal_image(modal, commit: str = VERL_COMMIT, patched: bool = True): + """The verl image; ``patched=False`` leaves the verl source as the fork ships it + (the bare-verl arms of the efficiency evaluation, which must not import yeto).""" + image = ( modal.Image.from_registry(VERL_BASE_IMAGE, add_python=VERL_PYTHON) .apt_install("git", "build-essential", "curl", "libnuma1") .pip_install(f"uv=={VERL_UV_VERSION}") .run_commands(*build_commands(commit)) .run_commands(*store_client_commands()) - .add_local_file(str(patch), "/workspace/patch_verl.py", copy=True) - .run_commands(f"{VENV_BIN}/python /workspace/patch_verl.py {VERL_ROOT}") - .env(image_env()) ) + if patched: + patch = Path(__file__).with_name("patch_verl.py") + image = (image.add_local_file(str(patch), "/workspace/patch_verl.py", copy=True) + .run_commands(f"{VENV_BIN}/python /workspace/patch_verl.py {VERL_ROOT}")) + return image.env(image_env()) diff --git a/yeto/rl/adapters/verl/island_entry.py b/yeto/rl/adapters/verl/island_entry.py index ab6f640f..5fec3a02 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/ports_impl.py b/yeto/rl/adapters/verl/ports_impl.py index 4952a73c..85f0e41c 100644 --- a/yeto/rl/adapters/verl/ports_impl.py +++ b/yeto/rl/adapters/verl/ports_impl.py @@ -30,13 +30,17 @@ import json import os import time +from collections.abc import Mapping +from dataclasses import dataclass, field +from math import prod from pathlib import Path from typing import Any from yeto.rl.contracts import InferencePublicationManifest, LocalStepReceipt from yeto.rl.engine import mismatch_criteria, offpolicy_source +from yeto.rl.engine.policy_digest import PolicyDigest, TrainerResidentState, is_resident from yeto.rl.engine.ports import GroupMetadata, PlacementDescription, PublicationResult, RolloutBatchHandle -from yeto.rl.engine.trainable_state import TrainableState +from yeto.rl.engine.trainable_state import LAYOUT_LORA, TrainableState from . import publish as pub from .param_names import train_to_canonical @@ -59,22 +63,88 @@ def _peft_module(worker): return getattr(module, "_fsdp_wrapped_module", module) +def _lora_tensors(worker_dict) -> tuple[dict, set]: + """({canonical name: full LoRA tensor on its device}, parameter dtypes); collective + (``full_tensor``): every rank calls it.""" + peft = _peft_module(_inner_worker(worker_dict)) + out, dtypes = {}, set() + for name, param in peft.named_parameters(): + if ".lora_" not in name: + continue + dtypes.add(str(param.dtype)) + full = param.full_tensor() if hasattr(param, "full_tensor") else param + out[train_to_canonical(name)] = full.detach() + return out, dtypes + + def remote_export_lora(worker_dict): """(canonical f32 CPU tensors, info) from the actor's PEFT model.""" import torch - peft = _peft_module(_inner_worker(worker_dict)) - out, dtypes = {}, set() with torch.no_grad(): - for name, param in peft.named_parameters(): - if ".lora_" not in name: - continue - dtypes.add(str(param.dtype)) - full = param.full_tensor() if hasattr(param, "full_tensor") else param - out[train_to_canonical(name)] = full.detach().to("cpu", torch.float32).clone() + tensors, dtypes = _lora_tensors(worker_dict) + out = {name: value.to("cpu", torch.float32).clone() for name, value in tensors.items()} return out, {"param_dtypes": sorted(dtypes), "count": len(out)} +def digest_lora_tensors(tensors, *, base_model_revision: str, lora_config_hash: str, + layout_hash: str, expected_specs) -> dict: + """rl-publish-fastpath, trainer side: validate ``tensors`` exactly as the driver's + ``canonical_state`` would and return only what the driver reads -- the policy + digests and the bf16 send checksum of :func:`publish.lora_checksum` (small dicts, + no tensors). Same functions as the full-export path, so the values are equal + bit for bit.""" + + import time + + from concurrent.futures import ThreadPoolExecutor + + from yeto.rl.core import canonical_state_from_owned_tensors + from yeto.rl.engine.policy_digest import digest_canonical_tensors + + started = time.monotonic() + # E1: ``tensors`` are this call's own copies (the trainer's export), so the + # validation keeps them without a second clone; the values are the same + state = canonical_state_from_owned_tensors( + 0, tensors, base_model_revision=base_model_revision, lora_config_hash=lora_config_hash, + layout_hash=layout_hash, expected_specs=expected_specs) + # E1: the bf16 send checksum runs in its own thread beside the two digest streams + # (hashlib releases the GIL); each function is unchanged, so the values are too + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="yeto-lora-checksum") as pool: + checksum = pool.submit(pub.lora_checksum, dict(state.tensors)) + digest = digest_canonical_tensors(state.tensors, base_model_revision=state.base_model_revision, + lora_config_hash=state.lora_config_hash, + layout_hash=state.layout_hash) + lora_checksum = checksum.result() + return {"digest": digest.to_wire(), "layout_hash": state.layout_hash, + "lora_checksum": lora_checksum, + "export_seconds": round(time.monotonic() - started, 3)} + + +def remote_export_digest(worker_dict, kwargs: dict): + """All ranks gather the LoRA (``full_tensor`` is collective); rank 0 hashes it. + + E1: rank 0 stages the tensors in one reusable pinned buffer (one GPU + synchronization, :func:`~yeto.rl.engine.policy_digest.stage_f32_on_host`); they + are hashed here and dropped, never returned.""" + import torch + import torch.distributed as dist + + from yeto.rl.engine.policy_digest import HostStaging, stage_f32_on_host + + with torch.no_grad(): + tensors, _dtypes = _lora_tensors(worker_dict) + if dist.is_available() and dist.is_initialized() and dist.get_rank() != 0: + return None + worker = _inner_worker(worker_dict) + staging = getattr(worker, "_yeto_digest_staging", None) + if staging is None: + staging = HostStaging() + setattr(worker, "_yeto_digest_staging", staging) + tensors = stage_f32_on_host(tensors, staging) + return digest_lora_tensors(tensors, **kwargs) + + def remote_apply_lora(worker_dict, tensors, reset_optimizer: bool): import torch @@ -161,6 +231,10 @@ def export_tensors(self) -> tuple[dict, dict]: results = self.call_workers(remote_export_lora) return results[0] + def export_digest(self, kwargs: dict) -> dict: + """rl-publish-fastpath: digests + send checksum computed on the workers.""" + return next(r for r in self.call_workers(remote_export_digest, kwargs) if r is not None) + def append_jsonl(self, name: str, record: dict) -> None: with (self.out / name).open("a") as handle: handle.write(json.dumps(record, sort_keys=True, default=str) + "\n") @@ -388,6 +462,25 @@ def round_metrics(self) -> dict: return out +@dataclass(frozen=True) +class VerlResidentState(TrainerResidentState): + """A verl policy kept in the trainer; carries the bf16 send checksum the + publisher compares with the vLLM read-back (no tensors on the driver).""" + + lora_checksum: Mapping[str, Any] | None = field(default=None, compare=False) + + +def sent_checksum(state) -> tuple[dict, int]: + """(send checksum, bf16 payload bytes) of the policy being published. A resident + state brings both from the trainer; any other state is hashed here as before.""" + + if is_resident(state) and getattr(state, "lora_checksum", None) is not None: + payload_bytes = sum(prod(shape) * 2 for _name, shape in state.digest.specs) + return dict(state.lora_checksum), int(payload_bytes) + tensors = dict(state.to_lora().tensors) + return pub.lora_checksum(tensors), sum(int(t.numel()) * 2 for t in tensors.values()) + + class VerlPolicyState: def __init__(self, island: VerlIsland): self.island = island @@ -403,6 +496,36 @@ def export(self) -> TrainableState: expected_specs=self.island.expected_specs) return TrainableState.from_lora(state) + def export_digest(self, *, policy_version: int) -> VerlResidentState | None: + """rl-publish-fastpath: the workers export and hash; only digests cross Ray. + None when the island has no digest hook (the driver then calls :meth:`export`).""" + + hook = getattr(self.island, "export_digest", None) + if not callable(hook): + return None + isl, started = self.island, time.monotonic() + raw = hook({"base_model_revision": isl.base_model_revision, + "lora_config_hash": isl.lora_config_hash, "layout_hash": isl.layout_hash, + "expected_specs": tuple(isl.expected_specs)}) + digest = PolicyDigest.from_wire(raw["digest"]) + print(f"[rl] verl publish-fastpath digest export v{policy_version}: " + f"total={time.monotonic() - started:.2f}s trainer={raw.get('export_seconds')}s " + f"hash={digest.hash_seconds:.2f}s bytes={digest.payload_bytes}", flush=True) + return VerlResidentState( + layout=LAYOUT_LORA, base_model_revision=isl.base_model_revision, + config_hash=isl.lora_config_hash, layout_hash=str(raw["layout_hash"]), + policy_version=int(policy_version), digest=digest, materialize_fn=self._materialize, + lora_checksum=raw["lora_checksum"]) + + def _materialize(self, policy_version: int) -> TrainableState: + from yeto.rl.core import canonical_state + + lora = self.export().to_lora() + return TrainableState.from_lora(canonical_state( + int(policy_version), lora.tensors, base_model_revision=lora.base_model_revision, + lora_config_hash=lora.lora_config_hash, layout_hash=lora.layout_hash, + expected_specs=lora.specs)) + def apply(self, state: TrainableState, *, optimizer: str, local_step: int) -> None: if optimizer not in ("reset", "preserve"): raise ValueError(optimizer) @@ -443,8 +566,7 @@ def publish(self, state: TrainableState) -> PublicationResult: isl, t = self.island, self.island.trainer version = int(state.policy_version) - lora = state.to_lora() - sent = pub.lora_checksum(dict(lora.tensors)) + sent, payload_bytes = sent_checksum(state) (isl.readback_dir / "expected_version").write_text(str(version)) stale = isl.readback_dir / f"readback-v{version}.json" if stale.exists(): @@ -470,7 +592,6 @@ def publish(self, state: TrainableState) -> PublicationResult: f"{list(check.mismatched[:3])}", cause=cause) digest = state.policy_tensor_hash() isl.published = (version, digest) - payload_bytes = sum(int(t_.numel()) * 2 for t_ in lora.tensors.values()) manifest = InferencePublicationManifest( publication_mode="full", base_policy_version=None, target_policy_version=version, target_policy_hash=digest, target_manifest_hash=_sha(f"{version}:{digest}"), diff --git a/yeto/rl/adapters/verl/trainer.py b/yeto/rl/adapters/verl/trainer.py index 6be18c9e..08c72e83 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/adapters/verl/verl_main.py b/yeto/rl/adapters/verl/verl_main.py index 106a0093..303ad942 100644 --- a/yeto/rl/adapters/verl/verl_main.py +++ b/yeto/rl/adapters/verl/verl_main.py @@ -57,12 +57,21 @@ def run(self, config, plan: dict): tq.close() -def _runtime_env() -> dict: +# NCCL host buffers on /dev/shm in every process of the disaggregated weight sync. +# On Modal one process could use cuMem host buffers while its peer fell back to +# /dev/shm, and the peer then failed to attach (S26 E0.1 eff01-yeto-verl-async-e, +# eff01-bare-verl-async-d with NCCL_DEBUG=INFO). The caller's own value wins. +CUMEM_HOST_ENV = "NCCL_CUMEM_HOST_ENABLE" + + +def _runtime_env(fully_async: bool = False) -> dict: from verl.trainer.constants_ppo import get_ppo_ray_runtime_env env = get_ppo_ray_runtime_env(None) env.setdefault("env_vars", {}) env["env_vars"]["TRANSFER_QUEUE_ENABLE"] = "1" + if fully_async: + env["env_vars"][CUMEM_HOST_ENV] = os.environ.get(CUMEM_HOST_ENV, "0") for key, value in os.environ.items(): # YETO_ covers YETO_ISLAND_HMAC_KEY (HELLO HMAC in the Ray actor that # opens the syncer client) and YETO_SANDBOX_MODAL_TOKEN_*; no cloud @@ -90,7 +99,7 @@ def _keep_npu_visible_in_cpu_actors() -> None: os.environ.setdefault(NPU_NOSET_ENV, "1") -def _plan_and_ray() -> dict: +def _plan_and_ray(fully_async: bool = False) -> dict: from yeto.island_credential_guard import check_island_credentials check_island_credentials() # secret-handling-hardening D4 (driver, both paths) @@ -98,7 +107,7 @@ def _plan_and_ray() -> dict: if not ray.is_initialized(): _keep_npu_visible_in_cpu_actors() address = os.environ.get("RAY_ADDRESS") or None - ray.init(address=address, runtime_env=_runtime_env(), namespace="yeto-verl") + ray.init(address=address, runtime_env=_runtime_env(fully_async), namespace="yeto-verl") return plan @@ -125,7 +134,7 @@ def main_fully_async(argv: list[str] | None = None) -> None: import verl.trainer.main_ppo as main_ppo from verl.experimental.fully_async_policy import fully_async_main as fa_main - plan = _plan_and_ray() + plan = _plan_and_ray(fully_async=True) def run_ppo(config, task_runner_class=None): from yeto.rl.adapters.verl.fully_async_runner import YetoFullyAsyncTaskRunner @@ -147,7 +156,7 @@ def main_buffer_trainer(argv: list[str] | None = None) -> None: import verl.trainer.main_ppo as main_ppo from verl.experimental.fully_async_policy import fully_async_main as fa_main - plan = _plan_and_ray() + plan = _plan_and_ray(fully_async=True) def run_ppo(config, task_runner_class=None): from yeto.rl.adapters.verl.buffer_trainer_runner import YetoBufferTaskRunner diff --git a/yeto/rl/bridge.py b/yeto/rl/bridge.py index 9cf0e388..7cae9256 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/bridges.py b/yeto/rl/engine/bridges.py index b1f955d3..47cb3b90 100644 --- a/yeto/rl/engine/bridges.py +++ b/yeto/rl/engine/bridges.py @@ -138,7 +138,26 @@ def at_safe_point(self, driver: IslandDriver, *, rollout_id: int, final: bool = k: info.get(k) for k in ("cut_id", "cut_bytes", "save_s", "store_copy_s", "store_commit_s", "total_s", "write_bytes_per_s", "policy_hash", "store_synced", "latest_seq", "pruned", - "incarnation_index", "final", "trainer_onloaded_for_cut")}) + "incarnation_index", "final", "trainer_onloaded_for_cut", + "store_commit", "prepared_before_offload")}) + + def cut_due(self, rollout_id: int) -> bool: + """A round cut is due at the safe point before ``rollout_id``.""" + return self.round_cuts is not None and self.round_cuts.due( + rollout_id, final=rollout_id >= self._end()) + + def prepare_cut(self, driver: IslandDriver, *, rollout_id: int, state: Any) -> None: + """E1: write the trainer part of the cut due at the next safe point while the + trainer is still resident (the driver calls this right before it offloads a + colocated trainer for the publication). A failure is reported, not fatal: the + safe point then cuts the usual way.""" + if self.round_cuts is None: + return + try: + self.round_cuts.prepare(driver, rollout_id=rollout_id, state=state, + final=rollout_id >= self._end()) + except Exception as error: # noqa: BLE001 + driver.emit("rl_cut_prepare_failed", rollout_id=rollout_id, error=repr(error)[:2000]) def resume_data_cursor(self) -> dict[str, int] | None: """The data cursor recorded by the round cut this run resumed from (None: diff --git a/yeto/rl/engine/driver.py b/yeto/rl/engine/driver.py index d815a0a7..f6ea1afb 100644 --- a/yeto/rl/engine/driver.py +++ b/yeto/rl/engine/driver.py @@ -122,6 +122,8 @@ "publish": "trainer+rollout", "eval": "rollout", "offload": "trainer", + "checkpoint": "trainer", + "rollout_offload": "rollout", "onload": "trainer", } _PHASE_TASK = {"sync": "outer_sync"} @@ -366,6 +368,18 @@ def _pooled_reward(batch: RolloutBatchHandle) -> tuple[float, float]: return mean, math.sqrt(max(0.0, second - mean * mean)) +def defer_digest(driver: Any) -> bool: + """E1: a colocated trainer that offloads for the publication hashes the policy + while it offloads (``deferred`` digest, read at the publication). Needs a + policy state that supports it; ``YETO_RL_DEFERRED_DIGEST=0`` turns it off.""" + import os + + if os.environ.get("YETO_RL_DEFERRED_DIGEST", "1").strip().lower() in {"0", "false", "no", "off"}: + return False + return (bool(getattr(driver, "colocated", False)) and bool(getattr(driver, "publish_offloaded", False)) + and bool(getattr(driver.policy_state, "supports_deferred_digest", False))) + + class IslandDriver: def __init__( self, @@ -685,6 +699,8 @@ def export_local_resident(self, *, policy_version: int): export_digest = getattr(self.policy_state, "export_digest", None) if not callable(export_digest): return None + if defer_digest(self): + return export_digest(policy_version=int(policy_version), deferred=True) return export_digest(policy_version=int(policy_version)) # -- per-round steps ------------------------------------------------- @@ -889,11 +905,20 @@ def _generate(self, rollout_id: int) -> RolloutBatchHandle: f"generation of rollout {rollout_id}", ) self.phase("generate", rollout_id=rollout_id, policy_version=rollout_id) + offload_rollout = getattr(self.rollout, "offload_after_generate", None) + split_offload = callable(offload_rollout) and bool(getattr(self.rollout, "offloads_after_generate", False)) + if split_offload: + self.rollout.defer_offload = True with self._load_sampler(rollout_id): # IR-3: the rollout side learns the target policy token up front. batch = self.rollout.generate( rollout_id, expected_policy_version=self.expected_token ) + # E1: a colocated rollout that frees its memory after generating (Miles + # --offload-rollout) does it in its own phase, not inside "generate" + if split_offload: + self.phase("rollout_offload", rollout_id=rollout_id) + offload_rollout() if batch.rollout_id != rollout_id or batch.policy_version != rollout_id: raise PolicyIdentityError( f"rollout pool returned rollout {batch.rollout_id} " @@ -1167,20 +1192,35 @@ def loop() -> None: thread.join(timeout=self.load_sample_interval_s + 5) # S15 evidence-window lesson: one terminal probe after generate so a # round whose loads mostly land after the last tick still leaves a - # final ``rl_load_sample`` (``terminal=True``). Observe only. - try: - sample = probe() - except Exception as exc: # noqa: BLE001 - observation must not kill the round - import logging - - logging.getLogger(__name__).warning( - "yeto terminal load_sample probe raised %s: %s", type(exc).__name__, exc) - sample = None - if sample is not None: - self.emit("rl_load_sample", rollout_id=rollout_id, **dict(sample), - t=self.clock(), profile_hash=self.profile_hash, - epoch=self.config_epoch, - weight_transport=self.weight_transport, terminal=True) + # final ``rl_load_sample`` (``terminal=True``). Observe only, so it runs + # off the round's critical path (E1: a synchronous probe cost ~0.1 s/round). + self._join_terminal_probe() + self._terminal_probe = threading.Thread( + target=self._terminal_load_sample, args=(probe, rollout_id), + name="yeto-load-sample-terminal", daemon=True) + self._terminal_probe.start() + + def _join_terminal_probe(self) -> None: + """Wait (bounded) for the previous round's terminal load probe: at most one runs.""" + thread = getattr(self, "_terminal_probe", None) + if thread is not None: + thread.join(timeout=self.load_sample_interval_s + 5) + self._terminal_probe = None + + def _terminal_load_sample(self, probe: Callable[[], Any], rollout_id: int) -> None: + try: + sample = probe() + except Exception as exc: # noqa: BLE001 - observation must not kill the round + import logging + + logging.getLogger(__name__).warning( + "yeto terminal load_sample probe raised %s: %s", type(exc).__name__, exc) + return + if sample is not None: + self.emit("rl_load_sample", rollout_id=rollout_id, **dict(sample), + t=self.clock(), profile_hash=self.profile_hash, + epoch=self.config_epoch, + weight_transport=self.weight_transport, terminal=True) def _emit_round_labels(self, rollout_id, batch, metrics) -> None: """A5: per-round algorithm metrics carry the same profile/epoch/transport labels.""" @@ -1455,6 +1495,10 @@ def _train_round(self, rollout_id: int, batch: RolloutBatchHandle, raw = self.trainer.step_metrics() metrics = raw if isinstance(raw, TrainStepMetrics) else TrainStepMetrics(**dict(raw)) rollout_fraction = getattr(batch, "cross_version_truncated_fraction", None) + late = getattr(self.rollout, "late_cross_version_fraction", None) + if rollout_fraction is None and callable(late): + # E1: estimated in the background while this step trained (Miles fully_async) + rollout_fraction = late(rollout_id) if metrics.cross_version_truncated_fraction is None and rollout_fraction is not None: # agentic-rollout-utilization 4.1: an engine that estimates the # cross-version truncation on the rollout side (Miles carry-over) @@ -1531,6 +1575,12 @@ def _train_round(self, rollout_id: int, batch: RolloutBatchHandle, self.ledger.outer_recorded(rollout_id, next_policy_version=rollout_id + 1) self._join_eval() if self.colocated and self.publish_offloaded: + # E1: a round cut due at the next safe point writes the trainer state now, + # while it is resident, instead of waking the offloaded trainer for it + cut_due = getattr(self.sync, "cut_due", None) + if callable(cut_due) and cut_due(rollout_id + 1): + self.phase("checkpoint", rollout_id=rollout_id) + self.sync.prepare_cut(self, rollout_id=rollout_id + 1, state=boundary.state) self._offload_trainer(rollout_id) self.publish(boundary.state, rollout_id=rollout_id + 1) self._close_span() @@ -1617,6 +1667,11 @@ def safe_point(self, rollout_id: int) -> str | None: self.elastic_hook.at_safe_point(self, rollout_id) if not self.controller.has_pending() and not self.controller.recovery_required: return None + # E1: a reconfiguration syncs the checkpoint store itself; no round cut may still + # be committing it in the background + join_commit = getattr(getattr(self.sync, "round_cuts", None), "join_commit", None) + if callable(join_commit): + join_commit(self) epoch_before = self.config_epoch self.phase("reconfigure", rollout_id=rollout_id, config_epoch=epoch_before) from .controller import RecoveryRequired @@ -1885,6 +1940,7 @@ def _run(self) -> TrainableState: failure = error raise finally: + self._join_terminal_probe() # its rl_load_sample lands before the tape closes if self.eval_overlap is not None: self.eval_overlap.abort() # no orphan eval task on error/cancel paths self._close_sync(failure) diff --git a/yeto/rl/engine/full_parameter_sync.py b/yeto/rl/engine/full_parameter_sync.py index f4b2141d..af2a1a1d 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 @@ -81,6 +82,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 +105,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 +115,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 +146,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 +405,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 +528,10 @@ 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 + # S26 E3b: actor data-stream counters at the previous boundary. + self._upload_before: tuple | None = None @property def actor(self) -> RoleStream: @@ -517,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: @@ -534,11 +558,32 @@ 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) pairing = self.pairing trajectories = self.trajectories(rollout_id) await pairing.gather_roles( @@ -564,6 +609,79 @@ 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() + if pairing.any_finalizing: + await pairing.finish_without_more_training() + return True + before = pairing.fragment_versions() + 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 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/policy_digest.py b/yeto/rl/engine/policy_digest.py index cd070275..7d58afbd 100644 --- a/yeto/rl/engine/policy_digest.py +++ b/yeto/rl/engine/policy_digest.py @@ -86,6 +86,49 @@ def from_wire(cls, raw: Mapping[str, Any]) -> "PolicyDigest": ) +class DeferredDigest: + """A :class:`PolicyDigest` whose hashes are still being computed in the trainer + (E1, sync colocated mode: the trainer hashes while it offloads). The layout + (``specs``) is known at once; the first read of a hash calls ``resolve`` once + (it waits for the trainer) and the values are those of the eager digest.""" + + def __init__(self, specs: Sequence[tuple[str, tuple[int, ...]]], resolve: Callable[[], "PolicyDigest"]): + self.specs = tuple((str(n), tuple(int(d) for d in s)) for n, s in specs) + self._resolve = resolve + self._done: PolicyDigest | None = None + + def resolved(self) -> "PolicyDigest": + if self._done is None: + done = self._resolve() + if done.specs != self.specs: + raise PolicyDigestError("the deferred digest's layout changed while it was hashed") + self._done = done + return self._done + + @property + def is_resolved(self) -> bool: + return self._done is not None + + @property + def policy_tensor_hash(self) -> str: + return self.resolved().policy_tensor_hash + + @property + def payload_hash(self) -> str: + return self.resolved().payload_hash + + @property + def payload_bytes(self) -> int: + return self.resolved().payload_bytes + + @property + def hash_seconds(self) -> float: + return self.resolved().hash_seconds + + def to_wire(self) -> dict[str, Any]: + return self.resolved().to_wire() + + def _raw(tensor: Any, name: str) -> memoryview: import torch @@ -159,6 +202,83 @@ def digest_canonical_tensors( hash_seconds=round(time.monotonic() - started, 3)) +class NonFiniteTensor(ValueError): + """A tensor staged by :func:`stage_f32_on_host` holds NaN or Inf (``name`` says which).""" + + def __init__(self, name: str): + super().__init__(f"{name!r} contains NaN or Inf") + self.name = name + + +class HostStaging: + """One reusable pinned host buffer for :func:`stage_f32_on_host` (E1). + + Allocating pinned memory costs milliseconds per call; the policy keeps the + same size every round, so the owner (the trainer process) keeps one of + these and the buffer is reallocated only when the policy grows.""" + + def __init__(self, *, pin: bool = True) -> None: + self._pin = pin + self._buffer = None + + def flat(self, numel: int): + import torch + + if self._buffer is None or self._buffer.numel() < numel: + pin = self._pin and torch.cuda.is_available() + self._buffer = torch.empty(numel, dtype=torch.float32, pin_memory=pin) + return self._buffer[:numel] + + +def stage_f32_on_host(tensors: Mapping[str, Any], staging: HostStaging | None = None) -> dict[str, Any]: + """Canonical CPU f32 contiguous copies of ``tensors`` (sorted by name), bit for bit + the values of ``t.detach().to("cpu", torch.float32).contiguous()`` per tensor. + + E1: with ``staging`` and every tensor on one device, the f32 conversion and the + finite check run on that device (the GPU) and the copies go to the pinned host + buffer asynchronously, with ONE synchronization for all of them (the per-tensor + path waits once per tensor: ~400 waits for a Qwen3-0.6B LoRA). No extra device + buffer is held: a tensor already f32 and contiguous is copied as it is. The + returned tensors are views of ``staging``'s buffer and are valid until the next + call with the same ``staging``: hash them, do not keep them (``torch.save`` of a + view would also write the whole buffer). Without ``staging`` (or with tensors on + several devices) each tensor is copied on its own, as before. Raises + :class:`NonFiniteTensor` naming the first non-finite tensor in name order, like + the per-tensor check.""" + + import torch + + names = sorted(tensors) + values = [tensors[n].detach() for n in names] + devices = {v.device for v in values} + if staging is None or len(devices) != 1: + out = {} + for name, value in zip(names, values): + value = value.to(device="cpu", dtype=torch.float32).contiguous().clone() + if not torch.isfinite(value).all().item(): + raise NonFiniteTensor(name) + out[name] = value + return out + flat = staging.flat(sum(v.numel() for v in values)) + out, finite, offset = {}, [], 0 + for name, value in zip(names, values): + source = value.to(dtype=torch.float32).contiguous() + finite.append(torch.isfinite(source).all()) + view = flat[offset:offset + source.numel()].view(source.shape) + view.copy_(source, non_blocking=True) + out[name] = view + offset += source.numel() + device = next(iter(devices)) + ok = bool(torch.stack(finite).all().item()) if finite else True + if device.type == "cuda": + torch.cuda.current_stream(device).synchronize() # the copies to the host are complete + if not ok: + for name, check in zip(names, finite): + if not bool(check.item()): + raise NonFiniteTensor(name) + return out + + def layout_hash_of(specs: Sequence[tuple[str, tuple[int, ...]]]) -> str: """The canonical layout hash rebuilt from (name, shape) only.""" diff --git a/yeto/rl/engine/resume.py b/yeto/rl/engine/resume.py index bfa3894e..389d45ec 100644 --- a/yeto/rl/engine/resume.py +++ b/yeto/rl/engine/resume.py @@ -210,10 +210,13 @@ def read_latest(store_root: Path) -> dict[str, Any] | None: def write_snapshot(store: CheckpointStore, state_dir: Path, info: Mapping[str, Any], *, - skip: Iterable[str] = ()) -> dict[str, Any]: + skip: Iterable[str] = (), commit: bool = True) -> dict[str, Any]: """Copy ``state_dir`` into a NEW snapshot dir (every file hash-checked after the copy), write its manifest, then ``LATEST``, then commit. Returns the LATEST record - plus timings/bytes.""" + plus timings/bytes. + + ``commit=False`` leaves the commit to the caller (``commit_s`` is None): until it + commits, a reader of a committing store still sees the previous LATEST.""" root = store.root started = time.monotonic() latest = read_latest(root) @@ -233,7 +236,7 @@ def write_snapshot(store: CheckpointStore, state_dir: Path, info: Mapping[str, A record = {"seq": seq, "manifest_sha256": manifest_sha, "bytes": total, **{k: info[k] for k in ("cut_id", "next_rollout_id", "incarnation", "reason") if k in info}} _write_json_atomic(root / LATEST, record) - commit_s = store.commit() + commit_s = store.commit() if commit else None return {**record, "copy_s": copied_s, "commit_s": commit_s} @@ -366,15 +369,42 @@ def prune(store_root: Path, *, keep: int, cut_ids_of: Callable[[Path], str | Non shutil.rmtree(p, ignore_errors=True) deleted.append(f"{STATE_DIR}/{p.name}") wanted = {cid for cid in (cut_ids_of(p) for p in kept) if cid} + # E1: a cut newer than every kept one may be being written right now (a periodic + # cut is captured while the previous cut commits in the background): not pruned + newest = max((r for r in map(_cut_rollout, wanted) if r is not None), default=None) cuts = root / ROUND_CUTS if cuts.is_dir(): for p in sorted(cuts.iterdir()): + rollout = _cut_rollout(p.name) + if newest is not None and rollout is not None and rollout > newest: + continue if p.is_dir() and p.name not in wanted: shutil.rmtree(p, ignore_errors=True) deleted.append(f"{ROUND_CUTS}/{p.name}") return deleted +def _cut_rollout(cut_id: str) -> int | None: + """Rollout id of a round cut named ``r-`` (None: other name).""" + head, sep, _ = cut_id.partition("-") + if not sep or not head.startswith("r") or not head[1:].isdigit(): + return None + return int(head[1:]) + + +def stage_state_dir(state_dir: Path, staged: Path) -> Path: + """Copy the snapshot files of ``state_dir`` (as :func:`write_snapshot` would take + them) into a fresh ``staged`` directory; returns it.""" + state_dir, staged = Path(state_dir), Path(staged) + if staged.exists(): + shutil.rmtree(staged) + for rel in _tree_files(state_dir): + (staged / rel).parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(state_dir / rel, staged / rel) + staged.mkdir(parents=True, exist_ok=True) + return staged + + def directory_bytes(path: Path) -> int: return sum(p.stat().st_size for p in Path(path).rglob("*") if p.is_file()) @@ -509,15 +539,25 @@ def _record(self, kind: str, **fields: Any) -> dict[str, Any]: return record def sync_checkpoint_store(self, reason: str, *, skip_if_recovery: bool = False, - extra: Mapping[str, Any] | None = None) -> bool: + extra: Mapping[str, Any] | None = None, + defer_commit: bool = False, source_dir: Path | None = None) -> bool: """Snapshot the state dir into the store, then ``LATEST``, commit, prune. A failure is remembered (``last_store_sync``) and returned False: the previous - LATEST stays valid.""" + LATEST stays valid. + + ``defer_commit``: stop after ``LATEST``; :meth:`commit_checkpoint_store` commits + and prunes later (the round cut runs it off the round's critical path). Until + then the store's committed LATEST is the previous one, so a crash in between + resumes from the previous cut. + + ``source_dir``: snapshot this copy of the state dir instead (a round cut whose + store sync runs later stages the state dir as it was at the cut).""" if skip_if_recovery and self.recovery_required: return False + source = Path(source_dir) if source_dir is not None else self.state_dir info = {"incarnation": self.incarnation["id"], "incarnation_index": self.incarnation["index"], "reason": reason, "wall_time": self._wall(), **dict(extra or {})} - pointer = self.state_dir / "round-cut.json" + pointer = source / "round-cut.json" if pointer.is_file(): # LATEST names the cut the snapshot's pointer names try: p = json.loads(pointer.read_text(encoding="utf-8")) @@ -526,17 +566,41 @@ def sync_checkpoint_store(self, reason: str, *, skip_if_recovery: bool = False, except ValueError: pass try: - result = write_snapshot(self.store, self.state_dir, info) - deleted = prune(self.store.root, keep=self.keep, cut_ids_of=_snapshot_cut_id) - if deleted: - self.store.commit() + result = write_snapshot(self.store, source, info, commit=not defer_commit) + if defer_commit: + self.last_store_sync = {**info, **result, "ok": True, "committed": False, "pruned": None} + return True + deleted = self._prune() except Exception as exc: # noqa: BLE001 self.last_store_sync = {**info, "ok": False, "error": repr(exc)} logging.getLogger(__name__).warning("checkpoint store sync (%s) failed: %r", reason, exc) return False - self.last_store_sync = {**info, **result, "ok": True, "pruned": deleted} + self.last_store_sync = {**info, **result, "ok": True, "committed": True, "pruned": deleted} return True + def commit_checkpoint_store(self) -> dict[str, Any]: + """Commit a ``defer_commit`` sync, then prune (pruning only after the new LATEST + is committed, so a failed commit never loses the previous cut). Thread-safe + against the caller as long as nothing else writes to the store meanwhile. + Returns (and remembers in ``last_store_sync``) the outcome; never raises.""" + sync = dict(self.last_store_sync or {}) + try: + sync["commit_s"] = self.store.commit() + sync["committed"] = True + sync["pruned"] = self._prune() + except Exception as exc: # noqa: BLE001 + sync.update(ok=False, error=repr(exc)) + logging.getLogger(__name__).warning("checkpoint store commit (%s) failed: %r", + sync.get("reason"), exc) + self.last_store_sync = sync + return sync + + def _prune(self) -> list[str]: + deleted = prune(self.store.root, keep=self.keep, cut_ids_of=_snapshot_cut_id) + if deleted: + self.store.commit() + return deleted + def close(self) -> None: return None diff --git a/yeto/rl/engine/round_cut.py b/yeto/rl/engine/round_cut.py index cadc5047..ce207118 100644 --- a/yeto/rl/engine/round_cut.py +++ b/yeto/rl/engine/round_cut.py @@ -39,6 +39,7 @@ import json import os from collections.abc import Callable, Mapping +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -92,14 +93,29 @@ def driver(self) -> Any: ref = self.driver_ref return ref() if callable(ref) and not hasattr(ref, "published_state") else ref - def context(self, cut_id: str) -> Any: + def context(self, cut_id: str, *, pending: tuple[int, Any] | None = None) -> Any: + """The cut context at the safe point after the current publication. + + ``pending`` = ``(version, state)``: build it before that publication, for + the state the driver is about to publish as ``version`` (the policy token + is the one the publication will set). Only valid when nothing but the + publication happens until the safe point (colocated, generation done); + :meth:`RoundCutCheckpoint.save` keeps such a cut only if the publication + then names the same policy.""" from yeto.rl.engine.controller import RebuildRefused from yeto.rl.engine.cut import CutProgress from yeto.rl.engine.cut import CutContext + from yeto.rl.engine.driver import policy_token driver = self.driver - state = driver.published_state - if state is None or driver.published_version is None: + if pending is not None: + # the token is read last (below): a deferred digest (E1) hashes meanwhile + version, state = int(pending[0]), pending[1] + settled, token = True, None + else: + version, state = driver.published_version, driver.published_state + settled, token = bool(driver.at_safe_point), driver.expected_token + if state is None or version is None: raise RebuildRefused("no published policy to cut") cursor = (self.rollout.data_cursor() if callable(getattr(self.rollout, "data_cursor", None)) else None) @@ -122,6 +138,10 @@ def context(self, cut_id: str) -> Any: if problems: raise RebuildRefused("; ".join(problems)) in_flight = tuple(exported["entries"]) + if exported.get("data_cursor") is not None: + # read together with the entries (a rollout that draws while the + # cut is taken, e.g. Miles fully_async): keep cursor and entries consistent + data.update(exported["data_cursor"]) data["buffer_length"] = int(exported["buffer_groups"]) summary["carried_over"] = len(in_flight) summary["max_policy_age"] = limit @@ -129,17 +149,19 @@ def context(self, cut_id: str) -> Any: local_step=int(driver.local_step), scheduler_samples=int(driver.local_step) * self.global_batch_size, global_batch_size=self.global_batch_size, - next_rollout_id=int(driver.published_version), + next_rollout_id=int(version), policy_version=int(state.policy_version), policy_hash=state.policy_tensor_hash(), ) + if token is None: + token = policy_token(version, progress.policy_hash) return CutContext( root=self.cut_root, cut_id=cut_id, backend_fingerprint=self.backend_fingerprint, progress=progress, algorithm=self.identity, data=data, ledger=summary, # The safe point is after boundary() returned non-stop and every # member acknowledged the policy: the outer commit is settled. - outer={"settled": bool(driver.at_safe_point), "policy_version": state.policy_version, - "policy_token": driver.expected_token}, + outer={"settled": settled, "policy_version": state.policy_version, + "policy_token": token}, shared_filesystem=self.shared_filesystem, in_flight=in_flight, ) @@ -155,6 +177,26 @@ def expectation(self, layout: Mapping[str, int], *, epoch: int) -> Any: ) +@dataclass(frozen=True) +class _PreparedCut: + """The trainer part of a cut written before the publication (``prepare``).""" + + rollout_id: int + context: Any # CutContext the shards and manifest were written with + save_s: float + pending: Any = None # the trainer's pending cut (start_cut); None: already written + + def matches(self, rollout_id: int, context: Any) -> bool: + """The safe point's own context names the same cut: same id, progress + (version, policy hash, steps), data cursor, ledger summary, outer commit + (settled, policy token) and in-flight entries.""" + mine = self.context + return (self.rollout_id == rollout_id and mine.cut_id == context.cut_id + and mine.progress == context.progress and dict(mine.data) == dict(context.data) + and dict(mine.ledger) == dict(context.ledger) and dict(mine.outer) == dict(context.outer) + and tuple(mine.in_flight or ()) == tuple(context.in_flight or ())) + + class RoundCutCheckpoint: """``save`` at a safe point, ``resume`` at sync start (see the module doc). @@ -166,7 +208,7 @@ class RoundCutCheckpoint: def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1, fingerprint: dict[str, Any] | None = None, allow_config_change: bool = False, - next_lr: Any = None) -> None: + next_lr: Any = None, background_commit: bool | None = None) -> None: if getattr(controller, "checkpoint_store", None) is None: raise ValueError("round cuts need the checkpoint store") if int(every) < 1: @@ -191,6 +233,25 @@ def __init__(self, *, controller: Any, source: Any, trainer: Any, every: int = 1 # (buffer_rollout.wire_sample_buffer sets both). self.buffer_watermark: Any = None # callable(int) -> int | None self.buffer_resume: Any = None # callable(int) -> int + # E1: a cut prepared before the trainer offloads (see prepare()), consumed by save() + self._prepared: _PreparedCut | None = None + # E1: a periodic cut writes its shards and commits the store in the background + # (the controller must offer commit_checkpoint_store); completed before the next + # cut and at the final cut. YETO_RL_CUT_BACKGROUND_COMMIT=0 does it all inline. + if background_commit is None: + background_commit = os.environ.get("YETO_RL_CUT_BACKGROUND_COMMIT", "1").strip().lower() \ + not in {"0", "false", "no", "off"} + self.background_commit = bool(background_commit) and callable( + getattr(controller, "commit_checkpoint_store", None)) + self._commit: tuple[Any, dict[str, Any]] | None = None # (thread, result box) + # E1: a periodic cut captured at its safe point, completed by join_commit() + self._deferred: dict[str, Any] | None = None + self._commit_info: dict[str, Any] | None = None # the running commit's cut + self.skipped_busy = 0 # periodic cuts skipped while a commit was running + # E1: prepared cuts a busy commit made the safe point skip; their trainer-side + # writers are collected (and their files dropped) at the next cut, when done + self._abandoned: list[_PreparedCut] = [] + self._declined: int | None = None # rollout whose prepare met a running commit # ---------------------------------------------------------------- pointer @property @@ -203,12 +264,28 @@ def pointer(self) -> dict[str, Any] | None: return None return dict(json.loads(path.read_text(encoding="utf-8"))) - def _write_pointer(self, info: dict[str, Any]) -> None: - tmp = self.pointer_path.with_suffix(".json.tmp") + def _write_pointer(self, info: dict[str, Any], *, path: Path | None = None) -> None: + target = self.pointer_path if path is None else Path(path) + tmp = target.with_suffix(".json.tmp") tmp.write_text(json.dumps(info, sort_keys=True), encoding="utf-8") with open(tmp, "rb") as fh: os.fsync(fh.fileno()) - os.replace(tmp, self.pointer_path) + os.replace(tmp, target) + + def _stage(self, cut_id: str) -> Path: + """E1: a copy of the state dir as it is now (small: ledger, journal, pointer), + next to it, for a store sync that runs after later rounds.""" + from yeto.rl.engine.resume import stage_state_dir + + state_dir = Path(self.controller.state_dir) + return stage_state_dir(state_dir, state_dir.parent / f".{state_dir.name}.stage-{cut_id}") + + @staticmethod + def _unstage(staged: Path | None) -> None: + if staged is not None: + import shutil + + shutil.rmtree(staged, ignore_errors=True) def _epoch(self) -> int: return int(self.controller.journal.epochs.config_epoch) @@ -228,6 +305,210 @@ def due(self, rollout_id: int, *, final: bool = False) -> bool: return should_cut(int(rollout_id), every=self.every, final=final) + def _skip(self, driver: Any, rollout_id: int, *, final: bool) -> bool: + """No cut at this safe point: not due, island RECOVERY_REQUIRED, or just + resumed from (or already cut at) this rollout.""" + if self.controller.recovery_required or not self.due(rollout_id, final=final): + return True + current = self.pointer() + return current is not None and int(current["next_rollout_id"]) == int(rollout_id) + + def _cut_id(self, rollout_id: int) -> str: + return f"r{int(rollout_id):06d}-{self.controller.incarnation['id']}" + + def _start(self, context: Any, *, background: bool) -> Any: + """Start the trainer part of a cut: a pending cut whose shards are written in + the background (trainers with ``start_cut``), else the finished cut (``None``).""" + start = getattr(self.trainer, "start_cut", None) + if background and callable(start): + return start(epoch=self._epoch(), context=context, background=True) + self.trainer.save_cut(epoch=self._epoch(), context=context) + return None + + def prepare(self, driver: Any, *, rollout_id: int, state: Any, final: bool = False) -> bool: + """E1: take the trainer part of the cut due at the safe point before rollout + ``rollout_id`` NOW, while the trainer is still resident, for the ``state`` the + driver is about to publish as ``rollout_id``. Saves the wake/sleep a cut of an + offloaded colocated trainer costs. The pointer and the store sync stay at the + safe point (:meth:`save`), which keeps the cut only if the publication names + the same policy; otherwise the prepared files are never pointed at (dropped). + Returns True when a cut was prepared. Raises like :meth:`save`.""" + import time + + self._prepared = None + self._reap_abandoned() + if driver.published_version is None or self._skip(driver, rollout_id, final=final): + return False + if not final and self._commit is not None and self._commit[0].is_alive(): + self._declined = int(rollout_id) # save() skips it if still busy, else _unprepared_wake + return False # save() skips this periodic cut (_commit_busy) + if int(state.policy_version) != int(rollout_id): + raise RoundCutError(f"prepared cut {rollout_id} for policy v{state.policy_version}") + self.join_commit(driver, wait=False) # the previous cut is complete before this one starts + cut_id = self._cut_id(rollout_id) + self._drop_unpointed(cut_id) + started = time.monotonic() + context = self.source.context(cut_id, pending=(int(rollout_id), state)) + pending = self._start(context, background=self._defer(final)) + self._prepared = _PreparedCut(rollout_id=int(rollout_id), context=context, + save_s=time.monotonic() - started, pending=pending) + return True + + def _commit_busy(self, driver: Any, rollout_id: int, *, final: bool) -> bool: + """E1: a periodic cut is skipped while the previous cut's store commit still + runs (a slow store, e.g. a Modal volume commit of 2-10 s, would otherwise hold + the round); the next safe point cuts again. The final cut waits for it.""" + if final or self._commit is None or not self._commit[0].is_alive(): + return False + self.skipped_busy += 1 + emit = getattr(driver, "emit", None) + if callable(emit): + emit("rl_cut_skipped", rollout_id=int(rollout_id), reason="previous cut still committing", + committing=self._commit[1].get("cut_id"), skipped_total=self.skipped_busy) + return True + + def _unprepared_wake(self, driver: Any, rollout_id: int, prepared: "_PreparedCut | None", *, + final: bool) -> bool: + """E1: skip a periodic cut whose prepare was declined (it met a running commit) + when it would wake the offloaded trainer + (~2 s on the Modal A10G sync run e1a7) and a committed cut already exists; the + next safe point cuts again. The final cut and the first cut are always taken.""" + declined, self._declined = self._declined, None + if final or declined != int(rollout_id) or prepared is not None: + return False + if self.pointer() is None or not getattr(driver, "_trainer_offloaded", False): + return False + self.skipped_busy += 1 + emit = getattr(driver, "emit", None) + if callable(emit): + emit("rl_cut_skipped", rollout_id=int(rollout_id), reason="not prepared before the trainer offloaded", + skipped_total=self.skipped_busy) + return True + + def _defer(self, final: bool) -> bool: + """Write the trainer shards and commit the store in the background: periodic + cuts only (the final cut is complete when the run ends), and not when a sample + buffer commits the same store during the round (no two commits at once), and + not before the first cut of the store is committed (a store holding cuts but + no LATEST refuses to start).""" + return (self.background_commit and not final and self.buffer_watermark is None + and self.pointer() is not None) + + def _drop_unpointed(self, cut_id: str) -> None: + import shutil + + from yeto.rl.engine.cut import cut_dir + + directory = cut_dir(self.root, cut_id) + if directory.exists(): + shutil.rmtree(directory) + + def join_commit(self, driver: Any = None, *, wait: bool = True) -> dict[str, Any] | None: + """Complete the previous background cut: wait for a running store commit, + then collect the pending cut's shards, commit its manifest, point the state + dir at it and start its store commit (``wait``: and wait for that too). Each + finished commit is reported (``rl_cut_committed``); on any failure the + previous pointer stays the committed one. Returns the last outcome (None: + nothing was pending).""" + result = self._join_running(driver) + deferred, self._deferred = self._deferred, None + if deferred is not None: + done = self._complete_deferred(deferred) + if not done.get("ok", True): + self._report(driver, done) + return done + self._commit_info = done + result = self._join_running(driver) if wait else done + return result + + def _join_running(self, driver: Any) -> dict[str, Any] | None: + if self._commit is None: + return None + thread, box = self._commit + thread.join() + self._commit = None + result = {**(self._commit_info or {}), **dict(box.get("result") or { + "ok": False, "error": "commit thread ended without a result"})} + self._commit_info = None + result.setdefault("cut_id", box.get("cut_id")) + self._report(driver, result) + return result + + @staticmethod + def _report(driver: Any, result: Mapping[str, Any]) -> None: + emit = getattr(driver, "emit", None) + if callable(emit): + emit("rl_cut_committed", **{k: result.get(k) for k in ( + "cut_id", "rollout_id", "ok", "finish_s", "commit_s", "pruned", "seq", "cut_bytes")}, + **({"writer": result["writer"]} if result.get("writer") else {}), + **({"error": result["error"]} if result.get("error") else {})) + + def _complete_deferred(self, deferred: dict[str, Any]) -> dict[str, Any]: + """Collect the pending cut's shards from the trainer (this thread: the trainer + is driven from here), then commit its manifest, point the state dir at it and + commit the store in the background (:meth:`_commit_in_background`).""" + import time + + info = deferred["info"] + started = time.monotonic() + try: + collect = getattr(self.trainer, "collect_cut", None) + if callable(collect): + commit_manifest = collect(deferred["pending"]) + else: + self.trainer.finish_cut(deferred["pending"]) + commit_manifest = None + except Exception as error: # noqa: BLE001 - the previous cut stays the pointed one + self._drop_unpointed(info["cut_id"]) + self._unstage(deferred.get("staged")) + return {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], "ok": False, + "error": repr(error)[:2000]} + self._commit_in_background(deferred, commit_manifest) + writer = getattr(self.trainer, "last_cut_writer", None) + return {"cut_id": info["cut_id"], "rollout_id": deferred["rollout_id"], + "collect_s": round(time.monotonic() - started, 3), **({"writer": writer} if writer else {})} + + def _commit_in_background(self, deferred: dict[str, Any], commit_manifest: Any) -> None: + """E1: off the round's critical path: the cut manifest (``commit_manifest``, None + when the trainer committed it already), the pointer with the cut record, the + state dir copied into the store, the store commit. Until the pointer moves the + previous cut stays the pointed one; a failure before that drops the new cut's + files. The outcome lands in the result box, reported by :meth:`join_commit`.""" + import threading + import time + + from yeto.rl.engine.resume import directory_bytes + + info, cut = deferred["info"], deferred["cut"] + cut_id = info["cut_id"] + box: dict[str, Any] = {"cut_id": cut_id} + + def run() -> None: + started = time.monotonic() + try: + if commit_manifest is not None: + commit_manifest() + saved = recovery.save(cut, store=self.controller.store, + write=lambda record: self._write_and_sync( + info, record, defer=True, staged=deferred.get("staged"))) + except Exception as error: # noqa: BLE001 - reported at the next join + self._drop_unpointed(cut_id) + box["result"] = {"ok": False, "error": repr(error)[:2000]} + return + finally: + self._unstage(deferred.get("staged")) + done = {"finish_s": round(time.monotonic() - started, 3), + "cut_bytes": directory_bytes(Path(self.root) / cut_id)} + if not saved["synced"]: + box["result"] = {**done, "ok": False, "error": ( + getattr(self.controller, "last_store_sync", None) or {}).get("error")} + return + box["result"] = {**done, **dict(self.controller.commit_checkpoint_store() or {})} + + thread = threading.Thread(target=run, name=f"yeto-cut-commit-{cut_id}", daemon=True) + self._commit = (thread, box) + thread.start() + def save(self, driver: Any, *, rollout_id: int, final: bool = False, trigger: str | None = None) -> dict[str, Any] | None: """Cut the trainer at the safe point before rollout ``rollout_id`` (= the published @@ -235,6 +516,12 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, when nothing was cut (not due, no publication yet, already cut at this rollout, island RECOVERY_REQUIRED). Raises on a refused/failed cut (the caller logs it). + A periodic cut (:meth:`_defer`) only captures the state here: the shards are + written, the pointer moved and the store committed in the background, completed + by :meth:`join_commit` before the next cut (``store_commit`` "deferred"; the + outcome is the ``rl_cut_committed`` event). Until then the previous cut stays + the committed one. + ``trigger`` names a final cut's trigger (default ``MANUAL``; ``SCHEDULED`` for a stop the trainer scheduler ordered); a periodic cut takes none.""" import time @@ -244,30 +531,51 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, if trigger is not None and (not final or trigger not in FINAL_TRIGGERS): raise RoundCutError(f"cut trigger {trigger!r} needs final=True and one of {FINAL_TRIGGERS}") ctl = self.controller - if ctl.recovery_required or driver.published_version is None: + prepared, self._prepared = self._prepared, None + if driver.published_version is None or self._skip(driver, rollout_id, final=final): + self._abandon(prepared) + return None + if self._commit_busy(driver, rollout_id, final=final): + # its trainer may still be writing the shards: collect them at the next + # cut instead of waiting for the writer here (E1: up to ~0.9 s per round) + if prepared is not None: + self._abandoned.append(prepared) return None - if not self.due(rollout_id, final=final): + self._reap_abandoned() + if self._unprepared_wake(driver, rollout_id, prepared, final=final): + self._abandon(prepared) return None if int(driver.published_version) != int(rollout_id): + self._abandon(prepared) raise RoundCutError(f"safe point {rollout_id} != published version {driver.published_version}") - current = self.pointer() - if current is not None and int(current["next_rollout_id"]) == int(rollout_id): - return None # just resumed from (or already cut) this point - cut_id = f"r{int(rollout_id):06d}-{ctl.incarnation['id']}" + self.join_commit(driver, wait=final) # the final cut leaves nothing running started = time.monotonic() - context = self.source.context(cut_id) - # colocated + offloaded trainer (1 GPU): its weights/optimizer are not on the - # device between publish and the next train step; wake it for the cut only. - offloaded = bool(getattr(driver, "_trainer_offloaded", False)) - if offloaded: - self.trainer.onload() - try: - self.trainer.save_cut(epoch=self._epoch(), context=context) - finally: - if offloaded: - self.trainer.offload() - save_s = time.monotonic() - started + context = self.source.context(self._cut_id(rollout_id)) progress = context.progress + offloaded = False + defer = self._defer(final) + used_prepared = prepared is not None and prepared.matches(int(rollout_id), context) + if used_prepared: + save_s, pending = prepared.save_s, prepared.pending # taken before the trainer offloaded + else: + self._abandon(prepared) + # a failed prepare or an earlier attempt left files under this id; no pointer + # names them (_skip), and cuts are immutable: drop them + self._drop_unpointed(context.cut_id) + # colocated + offloaded trainer (1 GPU): its weights/optimizer are not on the + # device between publish and the next train step; wake it for the cut only. + offloaded = bool(getattr(driver, "_trainer_offloaded", False)) + if offloaded: + self.trainer.onload() + try: + pending = self._start(context, background=defer) + finally: + if offloaded: + self.trainer.offload() + save_s = time.monotonic() - started + if pending is None: + defer = False # the trainer wrote the cut already (no start_cut) + cut_id = context.cut_id consumed_n, consumed_sha = self._consumed() next_lr = self.next_lr(driver) if callable(self.next_lr) else None info = { @@ -291,31 +599,70 @@ def save(self, driver: Any, *, rollout_id: int, final: bool = False, completed_groups=consumed_n, completed_groups_digest=consumed_sha, buffer_watermark=(self.buffer_watermark(int(rollout_id)) if callable(self.buffer_watermark) else None)) + extra = {"prepared_before_offload": used_prepared, "trainer_onloaded_for_cut": offloaded, + "save_s": round(save_s, 3)} + capture = pending.get("capture") if isinstance(pending, Mapping) else None + if capture: + # E1: per rank, device copies or the host path (and why: e.g. low device memory) + extra["cut_capture"] = capture + if defer: + # the state dir (ledger, journal) as it is at this cut: its store sync runs + # in the background after later rounds have changed it + staged = self._stage(cut_id) + self._deferred = {"info": info, "cut": cut, "pending": pending, "rollout_id": int(rollout_id), + "staged": staged} + return {**info, **extra, "store_commit": "deferred", "store_synced": None, + "total_s": round(time.monotonic() - started, 3)} + if pending is not None: + self.trainer.finish_cut(pending) saved = recovery.save(cut, store=ctl.store, write=lambda record: self._write_and_sync(info, record)) info[CUT_KEY] = saved[CUT_KEY] synced = saved["synced"] sync = dict(getattr(ctl, "last_store_sync", None) or {}) cut_bytes = directory_bytes(Path(self.root) / cut_id) - total_s = time.monotonic() - started + info.update(extra) info["store_synced"] = bool(synced) + info["store_commit"] = "inline" info["cut_bytes"] = cut_bytes - info["save_s"] = round(save_s, 3) info["store_copy_s"] = sync.get("copy_s") info["store_commit_s"] = sync.get("commit_s") - info["total_s"] = round(total_s, 3) + info["total_s"] = round(time.monotonic() - started, 3) info["write_bytes_per_s"] = (cut_bytes / save_s) if save_s > 0 else None info["pruned"] = sync.get("pruned") info["latest_seq"] = sync.get("seq") - info["trainer_onloaded_for_cut"] = offloaded if not synced: info["store_error"] = sync.get("error") return info - def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any]) -> dict[str, Any]: + def _reap_abandoned(self) -> None: + abandoned, self._abandoned = self._abandoned, [] + for prepared in abandoned: + self._abandon(prepared) + + def _abandon(self, prepared: "_PreparedCut | None") -> None: + """A prepared cut the safe point does not keep: wait for its writer, drop its files.""" + if prepared is None: + return + if prepared.pending is not None: + try: + self.trainer.finish_cut(prepared.pending) + except Exception: # noqa: BLE001 - dropped below either way + pass + self._drop_unpointed(prepared.context.cut_id) + + def _write_and_sync(self, info: dict[str, Any], cut_record: dict[str, Any], *, + defer: bool = False, staged: Path | None = None) -> dict[str, Any]: """The round cut's ``recovery.save`` writer: the pointer (with the cut record) - into the state dir, then the state dir into the store (which commits).""" - self._write_pointer({**info, CUT_KEY: cut_record}) - synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True) + into the state dir, then the state dir into the store (which commits, unless + ``defer``: the caller commits in the background). ``staged``: the state dir as + staged at the cut (:meth:`_stage`) is the one synced, with the same pointer.""" + pointer = {**info, CUT_KEY: cut_record} + self._write_pointer(pointer) + kwargs: dict[str, Any] = {"defer_commit": True} if defer else {} + if staged is not None: + self._write_pointer(pointer, path=Path(staged) / POINTER) + kwargs["source_dir"] = Path(staged) + synced = self.controller.sync_checkpoint_store("round_cut", skip_if_recovery=True, **kwargs) # the store sync commits itself (snapshot layout: in write_snapshot; flat layout: # in the controller, outside its lock), also when it failed: recovery.save must not return {"synced": bool(synced), "committed": True} diff --git a/yeto/rl/engine/run_config.py b/yeto/rl/engine/run_config.py index 616e588a..90abd7dc 100644 --- a/yeto/rl/engine/run_config.py +++ b/yeto/rl/engine/run_config.py @@ -255,6 +255,9 @@ class TrainMemoryConfig: max_tokens_per_gpu: int | None = None # Share of the optimizer state kept and updated on the host (0 = none). optimizer_offload_fraction: float = 0.0 + # Megatron gradient-accumulation fusion (fused wgrad into main_grad; needs + # APEX/TE fused kernels). Off keeps the historical argv. + gradient_accumulation_fusion: bool = False @dataclass(frozen=True) @@ -1094,6 +1097,7 @@ def _train_memory(args) -> TrainMemoryConfig: activation_recompute=bool(getattr(args, "rl_activation_recompute", False)), max_tokens_per_gpu=int(max_tokens) if max_tokens is not None else None, optimizer_offload_fraction=offload, + gradient_accumulation_fusion=bool(getattr(args, "rl_gradient_accumulation_fusion", False)), ) diff --git a/yeto/rl/engine/run_flags.py b/yeto/rl/engine/run_flags.py index b4de9148..2af06a5c 100644 --- a/yeto/rl/engine/run_flags.py +++ b/yeto/rl/engine/run_flags.py @@ -841,6 +841,62 @@ 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-push-dtype", + choices=("f32", "bf16", "int8"), + default="f32", + help="full-parameter only: wire dtype of the islands' pushed outer deltas " + "(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). 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", + 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", + 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", + 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/rl/engine/sample_roles.py b/yeto/rl/engine/sample_roles.py index de65b83f..ca766c3e 100644 --- a/yeto/rl/engine/sample_roles.py +++ b/yeto/rl/engine/sample_roles.py @@ -240,12 +240,12 @@ def island_learner(args: Any, island: int, num_learners: int) -> tuple[int, int] def buffer_full_parameter(args: Any) -> bool: - """``--tuning full`` on sampler/trainer islands (single or pair, no-sync): the - island trains every parameter without the full-parameter sync, since a buffer - island has no syncer; the policy reaches the sampler as a store release.""" + """``--tuning full`` on a no-sync island: the island trains every parameter + without the full-parameter sync, since it has no syncer. A sampler/trainer + island (single or pair) passes the policy to the sampler as a store release; + a single island with both roles (S26 E11) updates its own engines in place.""" return (getattr(args, "tuning", "lora") == "full" - and bool(getattr(args, "rl_single_island_no_sync", False)) - and (is_pair(args) or sample_role(args) != BOTH)) + and bool(getattr(args, "rl_single_island_no_sync", False))) def check_pair(args: Any, num_islands: int) -> None: diff --git a/yeto/rl/engine/sync_overlap.py b/yeto/rl/engine/sync_overlap.py new file mode 100644 index 00000000..3a7f17ec --- /dev/null +++ b/yeto/rl/engine/sync_overlap.py @@ -0,0 +1,59 @@ +"""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. + +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 + +OUTER_STEP_LAG = 1 + + +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: + """``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 diff --git a/yeto/rl/engine/sync_timing.py b/yeto/rl/engine/sync_timing.py index eaa4ad9b..eae22025 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/event_echo.py b/yeto/rl/event_echo.py index b3f2a33e..b9abdadc 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/rl/wandb_rl.py b/yeto/rl/wandb_rl.py index 3467c503..60a08c7b 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/syncer_profile.py b/yeto/syncer_profile.py index 390d2bc4..ce23fa99 100644 --- a/yeto/syncer_profile.py +++ b/yeto/syncer_profile.py @@ -259,8 +259,12 @@ 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") + 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.""" diff --git a/yeto/wan_probe.py b/yeto/wan_probe.py new file mode 100644 index 00000000..4ccb91fe --- /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, + ) + diff --git a/yeto/wan_streams.py b/yeto/wan_streams.py new file mode 100644 index 00000000..8d2cb6ab --- /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)), + ) diff --git a/yeto/wandb_tape.py b/yeto/wandb_tape.py index 125afa62..59c4e584 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: