metrale_model_engine/
rank_agree.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Startup check that every rank holds the same values for the env-read scalars
4//! that shape the collective schedule.
5//!
6//! Such a scalar is read independently on each rank. For example,
7//! [`metrale_model_arch::glm5next_layer::prefill_rows`] (`METRALE_GLM_PREFILL_ROWS`) sets the row
8//! count of each prefill sub-chunk, and each sub-chunk issues its own `reduce_partial`, so a skew
9//! between ranks is a hang or a reduce over the wrong extent rather than a perf difference.
10//! Rank 0 broadcasts its values, every rank compares them with its own, and a mismatch fails
11//! model construction with an error naming each disagreeing entry. The caller,
12//! `TransformerModel::new`, chooses the list.
13//!
14//! Owner: model-engine.
15//! Invariants: none beyond the types.
16
17use anyhow::{Result, bail};
18use metrale_comm::CommBackend;
19use metrale_gpu_runtime::gpu::GpuBackend;
20
21/// 2026-09-25: Broadcast rank 0's `items` and bail if this rank's own values differ.
22///
23/// `items` is `(name, value)`; the name appears only in the error message. A single-rank run or
24/// an empty list returns `Ok(())` without communicating.
25///
26/// This is a collective: every rank must call it with the same `items.len()` at the same point.
27/// The one caller is `TransformerModel::new`, which every rank runs.
28pub(crate) fn assert_ranks_agree(
29    gpu: &dyn GpuBackend,
30    comm: &dyn CommBackend,
31    items: &[(&str, u64)],
32) -> Result<()> {
33    if comm.world_size() < 2 || items.is_empty() {
34        return Ok(());
35    }
36    let bytes = items.len() * 8;
37    let buf = gpu.alloc(bytes)?;
38
39    let gathered = (|| -> Result<Vec<u64>> {
40        let mut host: Vec<u8> = items.iter().flat_map(|(_, v)| v.to_le_bytes()).collect();
41        gpu.copy_h2d(&host, buf)?;
42        comm.broadcast(buf.0, bytes, 0)?;
43        gpu.copy_d2h(buf, &mut host)?;
44        Ok(host
45            .chunks_exact(8)
46            .map(|c| u64::from_le_bytes(c.try_into().expect("chunks_exact(8)")))
47            .collect())
48    })();
49
50    // 2026-09-25: The scratch is freed before `gathered?`, so a failed broadcast does not leak
51    // it; a failed free is logged and does not fail the check.
52    if let Err(e) = gpu.free(buf) {
53        tracing::warn!("rank-agreement scratch free failed (non-fatal): {e}");
54    }
55
56    let root = gathered?;
57    let rank = comm.rank();
58    let bad: Vec<String> = items
59        .iter()
60        .zip(&root)
61        .filter(|((_, mine), theirs)| mine != *theirs)
62        .map(|((name, mine), theirs)| {
63            format!("{name}: rank {rank} has {mine}, rank 0 has {theirs}")
64        })
65        .collect();
66    if !bad.is_empty() {
67        bail!(
68            "ranks disagree on collective-shaping config — this is a hang or a wrong-extent \
69             reduce, not a perf skew. Set the same value on EVERY rank: {}",
70            bad.join("; ")
71        );
72    }
73    tracing::info!(
74        "rank-agreement OK on {} collective-shaping scalar(s): {}",
75        items.len(),
76        items
77            .iter()
78            .map(|(n, v)| format!("{n}={v}"))
79            .collect::<Vec<_>>()
80            .join(" ")
81    );
82    Ok(())
83}
84
85#[cfg(test)]
86mod tests {
87    //! 2026-09-25: The collective needs 2 ranks, so these tests cover only the comparison, on
88    //! `mismatches`, a copy of the filter in `assert_ranks_agree`.
89
90    fn mismatches(items: &[(&str, u64)], root: &[u64], rank: usize) -> Vec<String> {
91        items
92            .iter()
93            .zip(root)
94            .filter(|((_, mine), theirs)| mine != *theirs)
95            .map(|((name, mine), theirs)| {
96                format!("{name}: rank {rank} has {mine}, rank 0 has {theirs}")
97            })
98            .collect()
99    }
100
101    #[test]
102    fn agreement_is_silent() {
103        let items = [("METRALE_GLM_PREFILL_ROWS", 8u64), ("ep_protocol_v2", 0)];
104        assert!(mismatches(&items, &[8, 0], 1).is_empty());
105    }
106
107    #[test]
108    fn a_single_disagreement_names_the_lever_and_both_values() {
109        let items = [("METRALE_GLM_PREFILL_ROWS", 4u64), ("ep_protocol_v2", 0)];
110        let bad = mismatches(&items, &[8, 0], 1);
111        assert_eq!(
112            bad,
113            vec!["METRALE_GLM_PREFILL_ROWS: rank 1 has 4, rank 0 has 8"]
114        );
115    }
116
117    #[test]
118    fn every_disagreement_is_reported_not_just_the_first() {
119        let items = [("METRALE_GLM_PREFILL_ROWS", 4u64), ("ep_protocol_v2", 1)];
120        assert_eq!(mismatches(&items, &[8, 0], 1).len(), 2);
121    }
122}
metrale_model_engine/rank_agree.rs

metrale_model_engine/
rank_agree.rs
