From 33c0ba142430af28fcf67448ce3cc2d15bf97386 Mon Sep 17 00:00:00 2001 From: Pawel Gadzinski Date: Thu, 20 Aug 2026 12:43:03 +0200 Subject: [PATCH] Loosen fp16 atol in ONNX export tests for onnxruntime 1.29 CPU numerics Signed-off-by: Pawel Gadzinski --- tests/pytorch/test_onnx_export.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/tests/pytorch/test_onnx_export.py b/tests/pytorch/test_onnx_export.py index 0ba1536882..6ddf9cbbbc 100644 --- a/tests/pytorch/test_onnx_export.py +++ b/tests/pytorch/test_onnx_export.py @@ -368,6 +368,11 @@ def create_ort_input_dict(session, inputs): ) +def get_atol(precision: torch.dtype, default: float = 1e-3) -> float: + # ORT runs on CPU while TE runs on GPU; fp16 accumulation differences reach a few ULPs. + return 2e-2 if precision is torch.float16 else default + + def dtype2str(dtype: torch.dtype, fake_bf16_io=False): if fake_bf16_io: assert dtype == torch.bfloat16 @@ -458,7 +463,7 @@ def forward(self, inp): if precision in (torch.bfloat16,): return if fp8_recipe is None: - validate_result(fname, inp, model, atol=1e-3, te_outputs=te_outputs) + validate_result(fname, inp, model, atol=get_atol(precision), te_outputs=te_outputs) else: validate_result( fname, inp, model, atol=1e-2, is_fp8=fp8_recipe is not None, te_outputs=te_outputs @@ -587,7 +592,7 @@ def _test_export_layernorm_linear( if precision in (torch.bfloat16,): return if fp8_recipe is None: - validate_result(fname, inp, model, atol=1e-3, te_outputs=te_outputs) + validate_result(fname, inp, model, atol=get_atol(precision), te_outputs=te_outputs) elif precision != torch.bfloat16: validate_result( fname, @@ -595,7 +600,11 @@ def _test_export_layernorm_linear( model, # For current scaling we use Float8Quantizer in tests + amax computed by hand, # which has slightly different numerics than Float8CurrentScalingQuantizer. - atol=1e-3 if fp8_recipe.__class__ is not recipe.Float8CurrentScaling else 2e-2, + atol=( + get_atol(precision) + if fp8_recipe.__class__ is not recipe.Float8CurrentScaling + else 2e-2 + ), is_fp8=fp8_recipe is not None, te_outputs=te_outputs, ) @@ -673,7 +682,9 @@ def _test_export_layernorm_mlp( if precision in (torch.bfloat16,): return atol = ( - 2e-2 if fp8_recipe is not None else (5e-1 if activation == "swiglu" else 1e-3) + 2e-2 + if fp8_recipe is not None + else (5e-1 if activation == "swiglu" else get_atol(precision)) ) # TODO(pgadzinski) - check 2e-2 validate_result( fname, inp, model, atol=atol, is_fp8=fp8_recipe is not None, te_outputs=te_outputs