OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-07 12:21:27 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Add warning if tensor cores are not used (#88844)

Browse Source 
Fixes https://github.com/pytorch/torchdynamo/issues/1839

Should I do this for all backends or just inductor?

## Test
On a V100 I got from AWS

```python
from torch._dynamo import optimize
import torch

def fn(x, y):
    a = torch.cos(x)
    b = torch.sin(y)
    return a + b

new_fn = optimize("inductor")(fn)

a = new_fn(torch.Tensor(1),torch.Tensor(1))
print(a)
```

## New logs
```
(sourcetorch) ubuntu@ip-172-31-31-152:~/test$ python test.py
/home/ubuntu/pytorch/torch/_dynamo/eval_frame.py:318: UserWarning: Tensor cores are available but not enabled. Consider setting torch.backends.cuda.matmul.allow_tf32 == True in your python script for speedups
  warnings.warn(
tensor([1.3717])
```

Pull Request resolved: https://github.com/pytorch/pytorch/pull/88844
Approved by: https://github.com/ngimel, https://github.com/mlazos, https://github.com/anijain2305
This commit is contained in:
Mark Saroufim 2022-11-17 07:24:55 +00:00 committed by PyTorch MergeBot
parent b72f5b9ae3
commit 37c85cf5f2
1 changed files with 10 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
torch/_dynamo/eval_frame.py
View File

@ -350,6 +350,16 @@ def get_compiler_fn(compiler_fn):
def lookup_backend(compiler_fn):
"""Expand backend strings to functions"""
if compiler_fn == "inductor":
if torch.cuda.is_available():
if (
torch.backends.cuda.matmul.allow_tf32 is False
and torch.cuda.get_device_capability() >= (8, 0)
):
warnings.warn(
"TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled."
"Consider setting `torch.set_float32_matmul_precision('high')`"
)
compiler_fn = import_module(f"{config.inductor_import}.compile_fx").compile_fx
elif isinstance(compiler_fn, str):
from .optimizations import BACKENDS