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