Fabian
47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
import os
|
|
os.environ["HSA_OVERRIDE_GFX_VERSION"] = "10.1.0"
|
|
os.environ["HSA_ENABLE_SDMA"] = "0"
|
|
os.environ["HIP_VISIBLE_DEVICES"] = "0"
|
|
os.environ["TORCHDYNAMO_DISABLE"] = "1"
|
|
import torch
|
|
import time
|
|
|
|
print("CUDA available:", torch.cuda.is_available())
|
|
print("Device:", torch.cuda.get_device_name(0))
|
|
|
|
# Test 1: Small CPU→GPU
|
|
x = torch.randn(100, 100)
|
|
print("CPU tensor created, starting small transfer...")
|
|
t0 = time.time()
|
|
y = x.to("cuda:0")
|
|
t1 = time.time()
|
|
print(f"Small transfer: {t1-t0:.3f}s, device={y.device}")
|
|
|
|
# Test 2: Medium CPU→GPU (f16)
|
|
x2 = torch.randn(2048, 2048, dtype=torch.float16)
|
|
mb = x2.numel() * 2 / 1e6
|
|
print(f"Medium tensor: {mb:.1f} MB, starting transfer...")
|
|
t0 = time.time()
|
|
y2 = x2.to("cuda:0")
|
|
t1 = time.time()
|
|
print(f"Medium transfer: {t1-t0:.3f}s")
|
|
|
|
# Test 3: Big CPU→GPU (f16)
|
|
x3 = torch.randn(4096, 4096, dtype=torch.float16)
|
|
mb3 = x3.numel() * 2 / 1e6
|
|
print(f"Big tensor: {mb3:.1f} MB, starting transfer...")
|
|
t0 = time.time()
|
|
y3 = x3.to("cuda:0")
|
|
t1 = time.time()
|
|
print(f"Big transfer: {t1-t0:.3f}s")
|
|
|
|
# Test 4: GPU compute
|
|
print("Testing GPU compute...")
|
|
t0 = time.time()
|
|
z = torch.matmul(y2, y2.T)
|
|
torch.cuda.synchronize()
|
|
t1 = time.time()
|
|
print(f"GPU matmul 2048x2048: {t1-t0:.3f}s")
|
|
|
|
print("ALL TESTS PASSED")
|