From bc0634ddd07fc94a831463fe2d6afd2d048fe6a4 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Fri, 6 Nov 2026 10:33:38 +0900 Subject: [PATCH] test(models): measure and falsify the #2852 default-device lead The single #2852 miss never reproduced. The last unmeasured candidate was a process-global MLX default device left on the CPU by an earlier test, which would move the quantized gather off the backend where the 1e-3 bound was measured. Forcing the default device to the CPU for the whole case, the nvfp4 quantize and the f32 dequantize reference included, measures 5.7e-6 max absolute error at 40 tokens against 7.6e-6 on the GPU, and the arm after the CPU span reproduces the GPU baseline exactly. The lead is falsified, so no leaked device can produce the observed miss. The new test pins the device independence so a future CPU-path divergence fails there. The failure message now also reports the out-of-bound and non-finite element counts, the default device, and the GPU backend. The 1e-3 bound and both arms are unchanged. Refs #2852 --- src/models/nemotron_h_tests.rs | 156 +++++++++++++++++++++++++++------ 1 file changed, 127 insertions(+), 29 deletions(-) diff --git a/src/models/nemotron_h_tests.rs b/src/models/nemotron_h_tests.rs index f495303da..fff3e9cce 100644 --- a/src/models/nemotron_h_tests.rs +++ b/src/models/nemotron_h_tests.rs @@ -1153,12 +1153,16 @@ fn config_without_any_layer_stack_description_is_rejected() { /// is built; the rope override needs an attention layer and this model has /// one MoE block; and a per-thread MLX default stream (what a leaked /// `install_thread_local_default_stream` would leave) cannot reach a later -/// test because libtest gives every test a fresh thread. The numeric failure -/// did not reproduce in more than forty full serial runs of the failing -/// commit's binary, including under an eight-way concurrent load, so the -/// assertion now reports the measured error instead of a bare message and -/// holds MLX's state locks, and a companion test pins the result against the -/// ambient thread stream. The 1e-3 bound is unchanged. +/// test because libtest gives every test a fresh thread. The last remaining +/// candidate, a process-global default device left on the CPU, is measured by +/// `nemotron_h_nvfp4_experts_parity_is_independent_of_the_default_device` +/// and falsified: the CPU arm stays inside the bound (5.7e-6 max absolute at +/// 40 tokens) and returns the GPU arm unchanged. The numeric failure did not +/// reproduce in more than forty full serial runs of the failing commit's +/// binary, including under an eight-way concurrent load, so the assertion now +/// reports the measured error and the device it ran on instead of a bare +/// message and holds MLX's state locks, and a companion test pins the result +/// against the ambient thread stream. The 1e-3 bound is unchanged. #[test] fn nemotron_h_runs_nvfp4_experts_sorted_and_unsorted() { // The default device is process-global, so a concurrent @@ -1189,6 +1193,43 @@ fn nemotron_h_nvfp4_experts_parity_is_independent_of_the_thread_stream() { assert_nvfp4_experts_sorted_and_unsorted(); } +/// Regression for #2852: the parity does not depend on MLX's process-global +/// default device, and a CPU span leaves nothing behind for a later GPU run. +/// +/// The one observed failure compared a quantized `gather_qmm` result against +/// an f32 matmul over dequantized planes. The remaining lead after the first +/// audit was that an earlier test left the process-global default device on +/// the CPU, moving the gather off the backend where the bound was measured. +/// Forcing the CPU for the whole case, the nvfp4 quantize and the f32 +/// reference included, measures it: at 40 tokens the CPU arm's max absolute +/// error is 5.7e-6 against 7.6e-6 on the GPU, and its max relative error is +/// 2.5e-4, all inside the unchanged 1e-3 bound. The arm that runs after the +/// CPU span reproduces the GPU baseline's measured errors exactly (3.8e-6 at +/// 5 tokens, 7.6e-6 at 40 tokens), so the span also leaves no cached state +/// behind. Kept as a test so a later change that makes the CPU path, or the +/// return to the GPU, diverge fails here instead of once in a full serial run +/// (issue #2852). +#[test] +fn nemotron_h_nvfp4_experts_parity_is_independent_of_the_default_device() { + let _guard = crate::models::embedding_test_support::mlx_test_guard(); + let device_before = mlxcel_core::default_device_is_gpu(); + assert_nvfp4_experts_sorted_and_unsorted(); + { + let _cpu = mlxcel_core::streams::DefaultDeviceGuard::cpu(); + assert!( + !mlxcel_core::default_device_is_gpu(), + "the CPU arm must actually run on the CPU" + ); + assert_nvfp4_experts_sorted_and_unsorted(); + } + assert_eq!( + mlxcel_core::default_device_is_gpu(), + device_before, + "the CPU span must restore the previous default device" + ); + assert_nvfp4_experts_sorted_and_unsorted(); +} + /// Regression for #2852: the private `SwitchMLP` must not consult the shared /// fused-MoE gate. The gate caches `MLXCEL_FUSED_MOE` in a process-wide /// `OnceLock`, so wiring it into this path would make the two arms depend on @@ -1220,6 +1261,23 @@ fn nemotron_h_switch_mlp_does_not_consult_the_fused_moe_gate() { } fn assert_nvfp4_experts_sorted_and_unsorted() { + run_nvfp4_expert_arms(|tokens, got, want| { + let close = mlxcel_core::allclose(got, want, 1e-3, 1e-3); + if !mlxcel_core::item_bool(&close) { + assert_expert_parity(got, want, tokens); + } + }); +} + +/// Build the one-layer nvfp4 MoE and hand both routing arms to `check`. +/// +/// `check` receives the evaluated gather result and the lazy dequantized +/// reference for 5 tokens (unsorted) and 40 tokens (sorted), with the same +/// shape. Splitting the graph construction out lets the parity assertion and +/// the device-independence test share one definition of the case. +fn run_nvfp4_expert_arms( + mut check: impl FnMut(i32, &mlxcel_core::MlxArray, &mlxcel_core::MlxArray), +) { use super::{BlockType, NemotronHMixer, NemotronHModel}; use mlxcel_core::weights::WeightMap; @@ -1328,10 +1386,49 @@ fn assert_nvfp4_experts_sorted_and_unsorted() { let h = mlxcel_core::utils::relu_squared(&h); let want = mlxcel_core::matmul(&h, &mlxcel_core::swap_axes(&w2, -1, -2)); let want = mlxcel_core::reshape(&want, &[tokens, TOP_K, HIDDEN]); - let close = mlxcel_core::allclose(&got, &want, 1e-3, 1e-3); - if !mlxcel_core::item_bool(&close) { - assert_expert_parity(&got, &want, tokens); + check(tokens, &got, &want); + } +} + +/// Measured divergence between the gather result and the reference. +/// +/// The bound MLX's `allclose` applies is `|got - want| <= atol + rtol * +/// |want|`, so `out_of_bound` counts the elements that actually missed it. +/// A near-bound precision miss leaves `max_abs` close to the bound and few +/// offending elements; a corrupted row leaves a large `max_abs` and a cluster +/// of them. +struct ExpertParityError { + max_abs: f32, + max_rel: f32, + non_finite: usize, + out_of_bound: usize, + elements: usize, +} + +impl ExpertParityError { + fn measure(got: &mlxcel_core::MlxArray, want: &mlxcel_core::MlxArray) -> Self { + let got_values = mlxcel_core::utils::array_to_vec_f32(got); + let want_values = mlxcel_core::utils::array_to_vec_f32(want); + let mut error = Self { + max_abs: 0.0, + max_rel: 0.0, + non_finite: 0, + out_of_bound: 0, + elements: got_values.len(), + }; + for (got, want) in got_values.iter().zip(&want_values) { + if !got.is_finite() || !want.is_finite() { + error.non_finite += 1; + continue; + } + let abs = (got - want).abs(); + error.max_abs = error.max_abs.max(abs); + error.max_rel = error.max_rel.max(abs / want.abs().max(1e-6)); + if abs > 1e-3 + 1e-3 * want.abs() { + error.out_of_bound += 1; + } } + error } } @@ -1339,33 +1436,34 @@ fn assert_nvfp4_experts_sorted_and_unsorted() { /// /// Issue #2852's one observed failure printed no magnitude, which is why it /// stayed unexplained: a near-bound precision miss and a corrupted row have -/// very different causes. Report both the maximum absolute and relative -/// differences and the non-finite count so the next occurrence is actionable. +/// very different causes. Report the maximum absolute and relative +/// differences, how many elements actually missed the bound, the non-finite +/// count, and the default device and GPU backend the comparison ran on, so the +/// next occurrence is actionable without a separate probe. fn assert_expert_parity( got: &mlxcel_core::MlxArray, want: &mlxcel_core::MlxArray, tokens: i32, ) -> ! { - let got_values = mlxcel_core::utils::array_to_vec_f32(got); - let want_values = mlxcel_core::utils::array_to_vec_f32(want); - let mut max_abs = 0.0_f32; - let mut max_rel = 0.0_f32; - let mut non_finite = 0_usize; - for (got, want) in got_values.iter().zip(&want_values) { - if !got.is_finite() || !want.is_finite() { - non_finite += 1; - continue; - } - let abs = (got - want).abs(); - max_abs = max_abs.max(abs); - max_rel = max_rel.max(abs / want.abs().max(1e-6)); - } + let error = ExpertParityError::measure(got, want); + let device = if mlxcel_core::default_device_is_gpu() { + "the GPU" + } else { + "the CPU" + }; + let backend = mlxcel_core::hardware::gpu_backend_kind(); panic!( "nvfp4 routed experts differ from the dequantized reference at {tokens} tokens: \ - max absolute {max_abs:e}, max relative {max_rel:e}, non-finite {non_finite}/{} \ - against the 1e-3/1e-3 bound (issue #2852). A near-bound miss and a non-finite or \ - large row have different causes; record these numbers on the issue.", - got_values.len() + max absolute {:e}, max relative {:e}, {} of {} elements outside the 1e-3 + 1e-3*|ref| \ + bound, non-finite {} of {}; default device {device}, backend {backend:?} (issue #2852). \ + A near-bound miss and a non-finite or large row have different causes; record these \ + numbers on the issue.", + error.max_abs, + error.max_rel, + error.out_of_bound, + error.elements, + error.non_finite, + error.elements, ); }