From f3644cedc2994528ba3e230ffebd42729cea621c Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 15 Jul 2026 11:29:00 +0530 Subject: [PATCH 1/4] fix autoencoderkl dtype tests --- tests/models/autoencoders/test_models_autoencoder_kl.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index 0e3e1aa3a46a..6cc283599c3e 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -81,6 +81,13 @@ def get_dummy_inputs(self): class TestAutoencoderKL(AutoencoderKLTesterConfig, ModelTesterMixin, TrainingTesterMixin): + @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16], ids=["fp16", "bf16"]) + def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): + # The reference and reloaded models hold identical weights, so any output difference is + # half-precision kernel nondeterminism between the two module instances rather than a save/load + # fidelity issue. The default 1e-4 tolerance is too tight for that fp16/bf16 noise on some GPUs. + super().test_from_save_pretrained_dtype_inference(tmp_path, dtype, atol=1e-3) + def test_gradient_checkpointing_is_applied(self): expected_set = {"Decoder", "Encoder", "UNetMidBlock2D"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) From efc584d2d832fe2cb2c41e218b50dc6cb38828c0 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 18 Jul 2026 03:20:35 +0000 Subject: [PATCH 2/4] address daniel's comment --- tests/models/autoencoders/test_models_autoencoder_kl.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index 6cc283599c3e..183d9e13ca83 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -85,8 +85,11 @@ class TestAutoencoderKL(AutoencoderKLTesterConfig, ModelTesterMixin, TrainingTes def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): # The reference and reloaded models hold identical weights, so any output difference is # half-precision kernel nondeterminism between the two module instances rather than a save/load - # fidelity issue. The default 1e-4 tolerance is too tight for that fp16/bf16 noise on some GPUs. - super().test_from_save_pretrained_dtype_inference(tmp_path, dtype, atol=1e-3) + # fidelity issue. The default zero-rtol, 1e-4-atol tolerance is too tight for that fp16/bf16 noise on + # some GPUs: a pure absolute bound does not account for the output magnitude, so large outputs can drift + # past it while still being relatively close. We instead follow `torch.testing.assert_close`'s + # recommended bf16 tolerances, whose non-zero rtol scales the allowed error with the output. + super().test_from_save_pretrained_dtype_inference(tmp_path, dtype, atol=1e-5, rtol=1.6e-2) def test_gradient_checkpointing_is_applied(self): expected_set = {"Decoder", "Encoder", "UNetMidBlock2D"} From 896702479e332683f880701a8a2c439676c61d39 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 18 Jul 2026 03:47:07 +0000 Subject: [PATCH 3/4] more fixes --- tests/models/autoencoders/test_models_autoencoder_kl.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index 183d9e13ca83..aaa2acad8e68 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -86,10 +86,11 @@ def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): # The reference and reloaded models hold identical weights, so any output difference is # half-precision kernel nondeterminism between the two module instances rather than a save/load # fidelity issue. The default zero-rtol, 1e-4-atol tolerance is too tight for that fp16/bf16 noise on - # some GPUs: a pure absolute bound does not account for the output magnitude, so large outputs can drift - # past it while still being relatively close. We instead follow `torch.testing.assert_close`'s - # recommended bf16 tolerances, whose non-zero rtol scales the allowed error with the output. - super().test_from_save_pretrained_dtype_inference(tmp_path, dtype, atol=1e-5, rtol=1.6e-2) + # some GPUs, and neither a pure atol nor a pure rtol covers it alone: a pure rtol collapses toward zero + # for near-zero outputs (leaving fp16's absolute noise floor uncovered), while a pure atol ignores the + # output magnitude (so large bf16 outputs drift past it while still being relatively close). We combine + # both, in the spirit of `torch.testing.assert_close`: a small absolute floor plus a non-zero rtol. + super().test_from_save_pretrained_dtype_inference(tmp_path, dtype, atol=1e-3, rtol=1.6e-2) def test_gradient_checkpointing_is_applied(self): expected_set = {"Decoder", "Encoder", "UNetMidBlock2D"} From 29af70509c289e9a69d4db43c0f8d78d46307569 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Sat, 18 Jul 2026 04:11:36 +0000 Subject: [PATCH 4/4] dtype specific atol and rtole. --- .../autoencoders/test_models_autoencoder_kl.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_kl.py b/tests/models/autoencoders/test_models_autoencoder_kl.py index aaa2acad8e68..1872820269a1 100644 --- a/tests/models/autoencoders/test_models_autoencoder_kl.py +++ b/tests/models/autoencoders/test_models_autoencoder_kl.py @@ -85,12 +85,17 @@ class TestAutoencoderKL(AutoencoderKLTesterConfig, ModelTesterMixin, TrainingTes def test_from_save_pretrained_dtype_inference(self, tmp_path, dtype): # The reference and reloaded models hold identical weights, so any output difference is # half-precision kernel nondeterminism between the two module instances rather than a save/load - # fidelity issue. The default zero-rtol, 1e-4-atol tolerance is too tight for that fp16/bf16 noise on - # some GPUs, and neither a pure atol nor a pure rtol covers it alone: a pure rtol collapses toward zero - # for near-zero outputs (leaving fp16's absolute noise floor uncovered), while a pure atol ignores the - # output magnitude (so large bf16 outputs drift past it while still being relatively close). We combine - # both, in the spirit of `torch.testing.assert_close`: a small absolute floor plus a non-zero rtol. - super().test_from_save_pretrained_dtype_inference(tmp_path, dtype, atol=1e-3, rtol=1.6e-2) + # fidelity issue. The default zero-rtol, 1e-4-atol tolerance is too tight for that noise, and it has to + # be relaxed per dtype: bf16's coarser mantissa produces a larger absolute noise floor (~1 ULP, e.g. + # 3.9e-3 near magnitude 0.5) than fp16's (~1e-3), so a single shared tolerance is either too tight for + # bf16 or needlessly loose for fp16. Each dtype gets an absolute floor sized to its own noise plus + # `torch.testing.assert_close`'s recommended rtol (1e-3 for fp16, 1.6e-2 for bf16) so the allowed error + # also scales with the output magnitude. + atol, rtol = { + torch.float16: (2e-3, 1e-3), + torch.bfloat16: (5e-3, 1.6e-2), + }[dtype] + super().test_from_save_pretrained_dtype_inference(tmp_path, dtype, atol=atol, rtol=rtol) def test_gradient_checkpointing_is_applied(self): expected_set = {"Decoder", "Encoder", "UNetMidBlock2D"}