test(tx): scale the MoE LoRA tolerance to the output magnitude

[pranavraja99]

Closes #1604.

test_qwen3_moe_layer_lora compares the fused multi-LoRA MoE output against an equivalent merged-weight computation with np.allclose(rtol=1e-3, atol=1e-3) and trips intermittently in CI.

The issue suspected the tolerance model, so I instrumented all three parametrizations locally to measure what the comparison actually sees:

ep,tp	max abs value	max abs error	error / tensor scale	fraction of tolerance budget used
1,1	1.3e4 – 4.7e4	0.0073 – 0.0625	5.8e-7 – 1.3e-6	up to 12.6%
1,2	1.3e4 – 2.8e4	0.0088 – 0.0117	4.2e-7 – 8.4e-7	up to 60.5%
2,1	8.5e3 – 1.6e4	0.0059 – 0.0065	4.0e-7 – 6.9e-7	up to 4.3%

So the error is ~4e-7 to ~1.4e-6 relative to the magnitude of the tensor — ordinary float32 accumulation over the hidden dimension, not a correctness problem. And one parametrization was already burning 60% of its tolerance budget, which is the flake.

The mismatch is in the error model. np.allclose allows atol + rtol*|b| per element, but this error is set by the scale of the accumulated terms, not by the value of the individual output element. At |b| = 4.7e4 that grants ~47 of absolute error, which is meaningless; at |b| ≈ 0 it grants only atol = 1e-3, which is an order of magnitude below the error floor. The near-zero entries are what fail, exactly as the issue describes.

This compares against a single absolute tolerance scaled to the tensor:

atol = 1e-4 * np.abs(output_merged).max()
np.testing.assert_allclose(output_with_lora[sample_idx : sample_idx + 1], output_merged, rtol=0, atol=atol)

At 1e-4 of the max magnitude that is ~70x headroom over the worst observed error, while being an order of magnitude stricter than rtol=1e-3 for the entries that dominate the output — so this tightens the meaningful part of the check rather than just loosening it. np.testing.assert_allclose additionally reports the mismatch count and worst offender on failure instead of just False.

I did not reach for vLLM's check_logprobs_close (also floated in the issue) since these are MoE layer activations rather than logprobs, so there is no top-k ranking to compare.

Testing

uv run --isolated --extra jax --extra dev pytest tests/tx/models/test_qwen3.py -k moe_layer_lora — 3 passed, on both the instrumented and final versions. The neighbouring test_qwen3_moe_layer and test_qwen3 assertions are left alone: their tensors are not in this magnitude range, so the same reasoning does not apply to them.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request updates the test assertion in test_qwen3_moe_layer_lora to use a dynamically scaled absolute tolerance (atol) based on the maximum value of the merged output, which prevents issues with float32 accumulation errors. The reviewer suggested casting the calculated atol to a standard Python float to avoid potential type compatibility or dispatch issues when passing a JAX DeviceArray scalar to np.testing.assert_allclose.

[gemini-code-assist]

Since output_merged is a JAX array, np.abs(output_merged).max() returns a JAX DeviceArray scalar. Passing a JAX array directly as the atol parameter to np.testing.assert_allclose can sometimes lead to type compatibility or dispatch issues depending on the JAX and NumPy versions. Converting it to a standard Python float using float(...) ensures robust compatibility.

Suggested change

	atol = 1e-4 * np.abs(output_merged).max()
	atol = float(1e-4 * np.abs(output_merged).max())