diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 7b250b1fde..92c13d38f9 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,27 @@ def reset_global_fp8_state(): FP8GlobalStateManager.reset() +@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: + yield + finally: + torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = original_value + + +@pytest.fixture(autouse=True) +def disable_bf16_reduced_precision_reduction(): + with _disable_bf16_reduced_precision_reduction(): + yield + + class TorchScaledMaskedSoftmax(nn.Module): def __init__(self) -> None: super().__init__()