From f427dcc18b9530f0fbbe363c92942c50a822f1db Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 2 Sep 2026 08:48:45 -0700 Subject: [PATCH 1/2] Fix PyTorch BF16 reduction numerics tests Signed-off-by: Jeremy Berchtold --- tests/pytorch/test_numerics.py | 35 ++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 7b250b1fde..29665d8e5c 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -4,6 +4,7 @@ import math import os +from contextlib import contextmanager from typing import Dict, List, Tuple, Optional import pytest @@ -254,6 +255,22 @@ def reset_global_fp8_state(): FP8GlobalStateManager.reset() +@contextmanager +def _disable_bf16_reduced_precision_reduction(): + original_value = torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False + try: + yield + finally: + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = original_value + + +@pytest.fixture +def disable_bf16_reduced_precision_reduction(): + with _disable_bf16_reduced_precision_reduction(): + yield + + class TorchScaledMaskedSoftmax(nn.Module): def __init__(self) -> None: super().__init__() @@ -1699,7 +1716,14 @@ def test_layernorm_accuracy(dtype, bs, model, eps, zero_centered_gamma): @pytest.mark.parametrize("return_bias", all_boolean) @pytest.mark.parametrize("bias", all_boolean) def test_layernorm_linear_accuracy( - dtype, bs, model, normalization, zero_centered_gamma, return_bias, bias + dtype, + bs, + model, + normalization, + zero_centered_gamma, + return_bias, + bias, + disable_bf16_reduced_precision_reduction, ): config = model_configs[model] @@ -1783,7 +1807,14 @@ def test_layernorm_linear_accuracy( @pytest.mark.parametrize("bias", all_boolean) @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) def test_layernorm_linear_accuracy_delay_wgrad_compute( - dtype, bs, model, normalization, zero_centered_gamma, bias, fuse_wgrad_accumulation + dtype, + bs, + model, + normalization, + zero_centered_gamma, + bias, + fuse_wgrad_accumulation, + disable_bf16_reduced_precision_reduction, ): if NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("Delayed wgrad compute is not supported in debug mode.") From d62e9daa10bc3f1f533a89cd76717a058e6aaa38 Mon Sep 17 00:00:00 2001 From: Jeremy Berchtold Date: Wed, 2 Sep 2026 13:22:11 -0700 Subject: [PATCH 2/2] Apply BF16 precision fixture to all numerics tests Signed-off-by: Jeremy Berchtold --- tests/pytorch/test_numerics.py | 25 ++++++++----------------- 1 file changed, 8 insertions(+), 17 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 29665d8e5c..92c13d38f9 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -257,6 +257,11 @@ def reset_global_fp8_state(): @contextmanager def _disable_bf16_reduced_precision_reduction(): + """TE disables cuBLASLt heuristics that store partial GEMM results in BF16. + This is to ensure precision is kept. This affects older archs like L40. + PyTorch does not do this by default, so we need to adjust PyTorch's + settings here to match TE's increased precision here. + """ original_value = torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False try: @@ -265,7 +270,7 @@ def _disable_bf16_reduced_precision_reduction(): torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = original_value -@pytest.fixture +@pytest.fixture(autouse=True) def disable_bf16_reduced_precision_reduction(): with _disable_bf16_reduced_precision_reduction(): yield @@ -1716,14 +1721,7 @@ def test_layernorm_accuracy(dtype, bs, model, eps, zero_centered_gamma): @pytest.mark.parametrize("return_bias", all_boolean) @pytest.mark.parametrize("bias", all_boolean) def test_layernorm_linear_accuracy( - dtype, - bs, - model, - normalization, - zero_centered_gamma, - return_bias, - bias, - disable_bf16_reduced_precision_reduction, + dtype, bs, model, normalization, zero_centered_gamma, return_bias, bias ): config = model_configs[model] @@ -1807,14 +1805,7 @@ def test_layernorm_linear_accuracy( @pytest.mark.parametrize("bias", all_boolean) @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) def test_layernorm_linear_accuracy_delay_wgrad_compute( - dtype, - bs, - model, - normalization, - zero_centered_gamma, - bias, - fuse_wgrad_accumulation, - disable_bf16_reduced_precision_reduction, + dtype, bs, model, normalization, zero_centered_gamma, bias, fuse_wgrad_accumulation ): if NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("Delayed wgrad compute is not supported in debug mode.")