fix(model): disable cuBLAS TF32 in the DiffusionGemma FP32 parity scripts (#6486)

Signed-off-by: Pruthviraj Prakash <pruprakash@nvidia.com>

Author: pruprakash

examples/models/diffusion_gemma/fp32_precision.py

@@ -0,0 +1,36 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Strict FP32 setup shared by the DiffusionGemma parity examples."""
+
+import os
+
+import torch
+
+
+def disable_tf32() -> bool:
+    """Disable TF32 in PyTorch and in cuBLAS-backed libraries such as Transformer Engine."""
+    os.environ["NVIDIA_TF32_OVERRIDE"] = "0"
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    return tf32_disabled()
+
+
+def tf32_disabled() -> bool:
+    """Return whether TF32 is off for PyTorch matmul, cuDNN and cuBLAS-backed libraries."""
+    return (
+        os.environ.get("NVIDIA_TF32_OVERRIDE") == "0"
+        and not torch.backends.cuda.matmul.allow_tf32
+        and not torch.backends.cudnn.allow_tf32
+    )

examples/models/diffusion_gemma/real_encoder_parity.py

@@ -23,6 +23,7 @@
 
 import torch
 import torch.distributed as dist
+from fp32_precision import disable_tf32, tf32_disabled
 from megatron.core.transformer.enums import AttnBackend
 from megatron.core.utils import unwrap_model
 from PIL import Image
@@ -396,8 +397,7 @@ def main() -> None:
 
     dist.init_process_group("nccl")
     torch.cuda.set_device(0)
-    torch.backends.cuda.matmul.allow_tf32 = False
-    torch.backends.cudnn.allow_tf32 = False
+    disable_tf32()
 
     processor = AutoProcessor.from_pretrained(args.model_path, local_files_only=True)
     cases = encode_cases(processor)
@@ -542,7 +542,7 @@ def main() -> None:
         "ep": 1,
         "dtype": "bfloat16",
         "attention_backend": "unfused",
-        "tf32_disabled": True,
+        "tf32_disabled": tf32_disabled(),
         "thresholds": {"min_token_cosine": args.min_token_cosine, "max_relative_l2": args.max_relative_l2},
         "cases": results,
         "vision_diagnostics": vision_diagnostics,

examples/models/diffusion_gemma/real_fp32_parity.py

@@ -22,6 +22,7 @@
 
 import torch
 import torch.distributed as dist
+from fp32_precision import disable_tf32
 from megatron.core.transformer.enums import AttnBackend
 from megatron.core.utils import unwrap_model
 from real_encoder_parity import compare, device_inputs, encode_cases, native_hf_masks
@@ -106,8 +107,7 @@ def main() -> None:
 
     dist.init_process_group("nccl")
     torch.cuda.set_device(0)
-    torch.backends.cuda.matmul.allow_tf32 = False
-    torch.backends.cudnn.allow_tf32 = False
+    disable_tf32()
 
     processor = AutoProcessor.from_pretrained(args.model_path, local_files_only=True)
     cases = encode_cases(processor)

examples/models/diffusion_gemma/toy_parity.py

@@ -22,6 +22,7 @@
 
 import torch
 import torch.distributed as dist
+from fp32_precision import disable_tf32, tf32_disabled
 from megatron.core.transformer.enums import AttnBackend
 from transformers.models.diffusion_gemma import DiffusionGemmaConfig, DiffusionGemmaForBlockDiffusion
 
@@ -141,8 +142,7 @@ def main() -> None:
     dist.init_process_group("nccl")
     torch.cuda.set_device(0)
     torch.manual_seed(1234)
-    torch.backends.cuda.matmul.allow_tf32 = False
-    torch.backends.cudnn.allow_tf32 = False
+    disable_tf32()
 
     with tempfile.TemporaryDirectory() as tmp:
         model_dir = Path(tmp) / "diffusiongemma-toy"
@@ -279,7 +279,7 @@ def main() -> None:
             "decoder": decoder_results,
             "atol": args.atol,
             "attention_backend": args.attention_backend,
-            "tf32_disabled": True,
+            "tf32_disabled": tf32_disabled(),
             "passed": all(case["encoder_max_abs_diff"] <= args.atol for case in results.values())
             and all(case["relative_l2"] <= 0.001 for case in decoder_results.values()),
         }

examples/models/diffusion_gemma/toy_training_parity.py

@@ -31,6 +31,7 @@
 
 import torch
 import torch.distributed as dist
+from fp32_precision import disable_tf32
 from megatron.core.transformer.enums import AttnBackend
 from real_encoder_parity import compare_layer0, hf_layer0_trace, megatron_layer0_trace
 from toy_parity import image_case, native_hf_encoder_inputs, toy_config
@@ -351,8 +352,7 @@ def main() -> None:
     dist.init_process_group("nccl")
     torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", "0")))
     torch.manual_seed(1234)
-    torch.backends.cuda.matmul.allow_tf32 = False
-    torch.backends.cudnn.allow_tf32 = False
+    disable_tf32()
 
     shared_dir = [tempfile.mkdtemp(prefix="diffusiongemma-toy-") if dist.get_rank() == 0 else None]
     dist.broadcast_object_list(shared_dir, src=0)

tests/unit_tests/examples/test_diffusion_gemma_fp32_precision.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.util
+import os
+from pathlib import Path
+
+import pytest
+import torch
+
+
+pytestmark = pytest.mark.unit
+
+EXAMPLE_DIR = Path(__file__).resolve().parents[3] / "examples" / "models" / "diffusion_gemma"
+PARITY_SCRIPTS = ("toy_parity.py", "toy_training_parity.py", "real_encoder_parity.py", "real_fp32_parity.py")
+
+
+def _load_fp32_precision():
+    spec = importlib.util.spec_from_file_location("fp32_precision", EXAMPLE_DIR / "fp32_precision.py")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture
+def tf32_state():
+    matmul, cudnn = torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32
+    override = os.environ.pop("NVIDIA_TF32_OVERRIDE", None)
+    yield
+    os.environ.pop("NVIDIA_TF32_OVERRIDE", None)
+    if override is not None:
+        os.environ["NVIDIA_TF32_OVERRIDE"] = override
+    torch.backends.cuda.matmul.allow_tf32, torch.backends.cudnn.allow_tf32 = matmul, cudnn
+
+
+def test_disable_tf32_also_disables_cublas_tf32(tf32_state):
+    fp32_precision = _load_fp32_precision()
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+
+    assert fp32_precision.disable_tf32() is True
+    assert os.environ["NVIDIA_TF32_OVERRIDE"] == "0"
+    assert torch.backends.cuda.matmul.allow_tf32 is False
+    assert torch.backends.cudnn.allow_tf32 is False
+
+
+def test_torch_flags_alone_do_not_count_as_tf32_disabled(tf32_state):
+    fp32_precision = _load_fp32_precision()
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    assert fp32_precision.tf32_disabled() is False
+
+
+@pytest.mark.parametrize("script", PARITY_SCRIPTS)
+def test_parity_scripts_disable_tf32_through_the_shared_helper(script):
+    source = (EXAMPLE_DIR / script).read_text()
+
+    assert "disable_tf32()" in source
+    assert "allow_tf32 = False" not in source