probe-turing-caps.py
2.22 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
import torch, subprocess, sys, os
print("torch", torch.__version__, "cuda", torch.version.cuda)
print("arch_list", torch.cuda.get_arch_list())
print("n_gpu", torch.cuda.device_count())
for i in range(torch.cuda.device_count()):
free, total = torch.cuda.mem_get_info(i)
p = torch.cuda.get_device_properties(i)
print(f" gpu{i} {p.name} sm_{p.major}{p.minor} free={free/2**30:.1f}G total={total/2**30:.1f}G")
dev = "cuda:0"
def t(name, fn):
try:
r = fn(); print(f"[OK ] {name}: {r}")
except Exception as e:
print(f"[ERR] {name}: {type(e).__name__}: {str(e)[:180]}")
t("fp16 matmul", lambda: (torch.randn(512,512,device=dev,dtype=torch.float16) @ torch.randn(512,512,device=dev,dtype=torch.float16)).float().abs().mean().item())
t("bf16 matmul", lambda: (torch.randn(512,512,device=dev,dtype=torch.bfloat16) @ torch.randn(512,512,device=dev,dtype=torch.bfloat16)).float().abs().mean().item())
t("bf16 layernorm", lambda: torch.nn.functional.layer_norm(torch.randn(8,512,device=dev,dtype=torch.bfloat16),(512,)).float().abs().mean().item())
from torch.nn.attention import SDPBackend, sdpa_kernel
q=lambda dt: torch.randn(1,8,1024,64,device=dev,dtype=dt)
for dt in (torch.float16, torch.bfloat16):
for be in (SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION, SDPBackend.MATH):
def f(dt=dt,be=be):
with sdpa_kernel(be):
return tuple(torch.nn.functional.scaled_dot_product_attention(q(dt),q(dt),q(dt)).shape)
t(f"SDPA {str(dt).split('.')[-1]:9s} {be.name}", f)
def auto():
return tuple(torch.nn.functional.scaled_dot_product_attention(q(torch.float16),q(torch.float16),q(torch.float16)).shape)
t("SDPA fp16 AUTO", auto)
print("--- installed pkgs of interest ---")
import importlib.metadata as md
for name in ["transformers","accelerate","diffusers","bitsandbytes","torchao","safetensors","einops","modelscope","av","torchaudio","torchcodec","librosa","peft","xformers","flash_attn","sageattention","imageio","imageio-ffmpeg","numpy"]:
try: print(f" {name}=={md.version(name)}")
except Exception: print(f" {name}: NOT INSTALLED")
print("--- ffmpeg ---")
print(subprocess.run(["bash","-lc","ffmpeg -version 2>&1|head -1; which ffmpeg"],capture_output=True,text=True).stdout)