frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-21 05:34:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

[ROCm] Remove HIPBLASLT_ALLOW_TF32 from codebase (#162998)

Browse Source 
A few UT failures are caused by `HIPBLASLT_ALLOW_TF32`

Fixes #157094, #157093, #157092, #157091, #157064, #157063, #157062, #157061, #157042, #157041, #157039, #157004

Pull Request resolved: https://github.com/pytorch/pytorch/pull/162998
Approved by: https://github.com/jeffdaily

Co-authored-by: Jeff Daily <jeff.daily@amd.com>
This commit is contained in:
Xinya Zhang
2025-09-16 12:48:45 +00:00
committed by PyTorch MergeBot
parent fa127d9b20
commit cef815dc2c
10 changed files with 54 additions and 202 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
test/test_cuda.py
View File

@ -759,53 +759,7 @@ print(t.is_pinned())
torch._C._cuda_clearCublasWorkspaces()
@contextlib.contextmanager
def _hip_allow_tf32(self):
# for HIP/AMDGPU, tf32 is behind a flag because the TF32 support is new
# and only for MI300+
hip_allow_tf32 = os.environ.get("HIPBLASLT_ALLOW_TF32", None)
os.environ["HIPBLASLT_ALLOW_TF32"] = "1"
try:
yield
finally:
if hip_allow_tf32 is not None:
os.environ["HIPBLASLT_ALLOW_TF32"] = hip_allow_tf32
else:
del os.environ["HIPBLASLT_ALLOW_TF32"]
@unittest.skipIf(not TEST_WITH_ROCM, "not relevant for CUDA testing")
def test_hipblaslt_allow_tf32(self):
tf32_ctx = self._hip_allow_tf32
with tf32_ctx():
os.environ["HIPBLASLT_ALLOW_TF32"] = "0"
# Save original value of allow_tf32
orig = torch.backends.cuda.matmul.allow_tf32
# If allow_tf32 variable is declared as static in aten/src/ATen/Context.cpp
# then matmul.allow_tf32 will return False after this point even if
# HIP_BLASLT_ALLOW_TF32 is set to 1 and matmul.allow_tf32 is changed.
os.environ["HIPBLASLT_ALLOW_TF32"] = "1"
# Toggle torch.backends.cuda.matmul.allow_tf32 couple of times.
torch.backends.cuda.matmul.allow_tf32 = not orig
test1 = torch.backends.cuda.matmul.allow_tf32
torch.backends.cuda.matmul.allow_tf32 = orig
test2 = torch.backends.cuda.matmul.allow_tf32
self.assertNotEqual(test1, test2)
# Restore original value of allow_tf32
torch.backends.cuda.matmul.allow_tf32 = orig
def test_cublas_allow_tf32_get_set(self):
"""
We only turn on TF32 for MI300 with a special env var. This is because TF32
is only available in MI300+ and is in experimental mode (hipblaslt support
is current WIP)
"""
tf32_ctx = self._hip_allow_tf32 if torch.version.hip else contextlib.nullcontext
with tf32_ctx():
self._test_cublas_allow_tf32_get_set_inner()
def _test_cublas_allow_tf32_get_set_inner(self):
skip_tf32_cublas = "TORCH_ALLOW_TF32_CUBLAS_OVERRIDE" in os.environ and int(
os.environ["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"]
)
@ -820,12 +774,6 @@ print(t.is_pinned())
torch.backends.cuda.matmul.allow_tf32 = orig
def test_float32_matmul_precision_get_set(self):
tf32_ctx = self._hip_allow_tf32 if torch.version.hip else contextlib.nullcontext
with tf32_ctx():
self._test_float32_matmul_precision_get_set_inner()
def _test_float32_matmul_precision_get_set_inner(self):
orig = torch.get_float32_matmul_precision()
skip_tf32_cublas = "TORCH_ALLOW_TF32_CUBLAS_OVERRIDE" in os.environ and int(
os.environ["TORCH_ALLOW_TF32_CUBLAS_OVERRIDE"]