rlx_runtime/dist/diagnostics.rs
1//! Cross-backend divergence diagnostic — separate explainable f32 round-off from
2//! a real kernel bug by running one graph on the CPU oracle + every other backend.
3
4use crate::{Device, Session};
5use rlx_ir::Graph;
6
7// ════════════════════════════════════════════════════════════════════════════
8// Cross-backend divergence diagnostic — "explainable drift vs. real kernel bug".
9//
10// On a heterogeneous cluster the SAME graph runs on different backends (Metal on
11// the Mac, CUDA on the box, CPU elsewhere), and their f32 results differ. Most of
12// that is bounded round-off — expected, "explainable". But a divergence too large
13// to be rounding is a kernel bug (the class that hid the SoftmaxCrossEntropy and
14// RMSNorm backward bugs). This runs one graph on the CPU oracle + every other
15// local backend and quantifies the gap per backend so the two can be told apart.
16// ════════════════════════════════════════════════════════════════════════════
17
18/// One backend's worst divergence from the CPU oracle for a given graph run.
19#[derive(Clone, Debug)]
20pub struct BackendDivergence {
21 pub device: Device,
22 /// Largest `|backend - cpu|` over every output element.
23 pub max_abs: f32,
24 /// Largest `|backend - cpu| / (|cpu| + 1e-6)` over every output element.
25 pub max_rel: f32,
26 /// Index of the graph output that carried `max_rel`.
27 pub worst_output: usize,
28 /// `true` when `max_rel` exceeds the tolerance → likely a kernel bug, not
29 /// round-off. Investigate with the per-primitive parity tests.
30 pub suspect: bool,
31}
32
33/// Worst (max_abs, max_rel, worst_output_index) between an oracle and a candidate
34/// run's outputs. Split out so the classification is unit-testable without a GPU.
35fn divergence_stats(oracle: &[Vec<f32>], got: &[Vec<f32>]) -> (f32, f32, usize) {
36 let (mut max_abs, mut max_rel, mut worst) = (0.0f32, 0.0f32, 0usize);
37 for (i, (o, g)) in oracle.iter().zip(got).enumerate() {
38 for (a, b) in o.iter().zip(g) {
39 let abs = (a - b).abs();
40 let rel = abs / (a.abs() + 1e-6);
41 if abs > max_abs {
42 max_abs = abs;
43 }
44 if rel > max_rel {
45 max_rel = rel;
46 worst = i;
47 }
48 }
49 }
50 (max_abs, max_rel, worst)
51}
52
53/// Run `graph` on the CPU oracle and on **every other backend live on this node**,
54/// reporting each one's worst divergence from CPU. Call it on each cluster node to
55/// localize where distributed numbers drift — and whether that drift is bounded
56/// round-off (`suspect == false`) or a candidate kernel bug (`suspect == true`).
57///
58/// `params` are the graph's weights (fed via `set_param`), `inputs` the run feeds;
59/// `tol_rel` flags a backend whose relative error exceeds it (≈`2e-3` for f32
60/// forward graphs; backward graphs accumulate more, try `~1e-2`). Returns one
61/// entry per non-CPU backend (empty on a CPU-only build — nothing to compare).
62pub fn backend_divergence(
63 graph: &Graph,
64 params: &[(&str, &[f32])],
65 inputs: &[(&str, &[f32])],
66 tol_rel: f32,
67) -> Vec<BackendDivergence> {
68 let run_on = |dev: Device| -> Vec<Vec<f32>> {
69 let mut s = Session::new(dev).compile(graph.clone());
70 for (n, v) in params {
71 s.set_param(n, v);
72 }
73 s.run(inputs)
74 };
75 let oracle = run_on(Device::Cpu);
76 crate::available_devices()
77 .into_iter()
78 // CPU is the oracle; ANE inference parity is a separate gated route.
79 .filter(|&d| d != Device::Cpu && !matches!(d, Device::Ane))
80 .map(|dev| {
81 let (max_abs, max_rel, worst_output) = divergence_stats(&oracle, &run_on(dev));
82 BackendDivergence {
83 device: dev,
84 max_abs,
85 max_rel,
86 worst_output,
87 suspect: max_rel > tol_rel,
88 }
89 })
90 .collect()
91}
92
93/// Print a [`backend_divergence`] report and return `true` if ANY backend looks
94/// suspect (relative error past `tol_rel`) — a convenient one-liner for a CLI /
95/// CI gate on each node. A green report means the cluster's cross-backend drift
96/// is bounded round-off; a red one names the backend + output to investigate.
97pub fn report_backend_divergence(
98 label: &str,
99 graph: &Graph,
100 params: &[(&str, &[f32])],
101 inputs: &[(&str, &[f32])],
102 tol_rel: f32,
103) -> bool {
104 let report = backend_divergence(graph, params, inputs, tol_rel);
105 if report.is_empty() {
106 eprintln!("[{label}] backend divergence: CPU-only node — nothing to compare");
107 return false;
108 }
109 let mut any = false;
110 for d in &report {
111 any |= d.suspect;
112 eprintln!(
113 "[{label}] {:>6} vs cpu: max_abs={:.3e} max_rel={:.3e} (output #{}) — {}",
114 crate::device_label(d.device),
115 d.max_abs,
116 d.max_rel,
117 d.worst_output,
118 if d.suspect {
119 "SUSPECT (relative error too large for round-off — likely a kernel bug)"
120 } else {
121 "round-off (explainable)"
122 },
123 );
124 }
125 any
126}
127
128#[cfg(test)]
129mod tests {
130 use super::*;
131
132 #[test]
133 fn divergence_stats_flags_worst_output() {
134 // output 0: exact; output 1: one element off by 0.5 on a base of 2.0.
135 let oracle = vec![vec![1.0f32, 2.0], vec![2.0f32, 4.0]];
136 let got = vec![vec![1.0f32, 2.0], vec![2.5f32, 4.0]];
137 let (abs, rel, worst) = divergence_stats(&oracle, &got);
138 assert!((abs - 0.5).abs() < 1e-6, "abs={abs}");
139 assert!((rel - 0.25).abs() < 1e-3, "rel={rel}"); // 0.5 / (2.0 + eps)
140 assert_eq!(worst, 1, "worst output should be #1");
141 // identical runs → zero divergence.
142 let (a0, r0, _) = divergence_stats(&oracle, &oracle);
143 assert_eq!((a0, r0), (0.0, 0.0));
144 }
145
146 #[test]
147 fn backend_divergence_cpu_only_is_empty() {
148 // On a CPU-only build there is no other backend to compare — the report is
149 // empty (and must not panic / compare CPU against itself).
150 use rlx_ir::{DType, Graph, Shape};
151 let mut g = Graph::new("d");
152 let x = g.input("x", Shape::new(&[2, 2], DType::F32));
153 let w = g.param("w", Shape::new(&[2, 2], DType::F32));
154 let mm = g.matmul(x, w, Shape::new(&[2, 2], DType::F32));
155 g.set_outputs(vec![mm]);
156 let report = backend_divergence(
157 &g,
158 &[("w", &[1.0, 0.0, 0.0, 1.0])],
159 &[("x", &[1.0, 2.0, 3.0, 4.0])],
160 2e-3,
161 );
162 // CPU is filtered out as the oracle; on a CPU-only test build nothing else
163 // is live, so the report is empty. (On a GPU build it would hold entries.)
164 assert!(report.iter().all(|d| d.device != Device::Cpu));
165 }
166}