96 lines
2.6 KiB
Bash
96 lines
2.6 KiB
Bash
#!/bin/bash
|
|
# Test basic tensor.to(cuda) with explicit flushing
|
|
cd /home/fabian/ComfyUI
|
|
source venv/bin/activate
|
|
|
|
export HSA_OVERRIDE_GFX_VERSION=10.1.0
|
|
export HSA_ENABLE_SDMA=0
|
|
export HIP_VISIBLE_DEVICES=0
|
|
export PYTORCH_HIP_ALLOC_CONF=garbage_collection_threshold:0.8
|
|
|
|
# Kill any remaining python processes first
|
|
pkill -9 -f "python.*main.py" 2>/dev/null
|
|
sleep 1
|
|
|
|
timeout 120 python3 -u -c "
|
|
import sys, torch, time, os
|
|
sys.stdout.flush()
|
|
print('Starting transfer test...', flush=True)
|
|
print(f'CUDA available: {torch.cuda.is_available()}', flush=True)
|
|
print(f'Device: {torch.cuda.get_device_name(0)}', flush=True)
|
|
|
|
# Test 1: Create on GPU
|
|
print('Test 1: torch.randn(10) on GPU...', flush=True)
|
|
t0 = time.time()
|
|
a = torch.randn(10, device='cuda')
|
|
torch.cuda.synchronize()
|
|
dt = time.time()-t0
|
|
print(f' PASS in {dt:.3f}s', flush=True)
|
|
|
|
# Test 2: Small CPU to GPU
|
|
print('Test 2: small .to(cuda)...', flush=True)
|
|
b = torch.randn(10)
|
|
t0 = time.time()
|
|
c = b.to('cuda')
|
|
torch.cuda.synchronize()
|
|
dt = time.time()-t0
|
|
print(f' PASS in {dt:.3f}s', flush=True)
|
|
|
|
# Test 3: 1MB float
|
|
print('Test 3: 1MB float .to(cuda)...', flush=True)
|
|
d = torch.randn(256*1024)
|
|
t0 = time.time()
|
|
e = d.to('cuda')
|
|
torch.cuda.synchronize()
|
|
dt = time.time()-t0
|
|
print(f' PASS in {dt:.3f}s', flush=True)
|
|
|
|
# Test 4: 1MB uint8
|
|
print('Test 4: 1MB uint8 .to(cuda)...', flush=True)
|
|
f = torch.randint(0, 255, (1024*1024,), dtype=torch.uint8)
|
|
t0 = time.time()
|
|
g = f.to('cuda')
|
|
torch.cuda.synchronize()
|
|
dt = time.time()-t0
|
|
print(f' PASS in {dt:.3f}s', flush=True)
|
|
|
|
# Test 5: 10MB uint8
|
|
print('Test 5: 10MB uint8 .to(cuda)...', flush=True)
|
|
h = torch.randint(0, 255, (10*1024*1024,), dtype=torch.uint8)
|
|
t0 = time.time()
|
|
i = h.to('cuda')
|
|
torch.cuda.synchronize()
|
|
dt = time.time()-t0
|
|
print(f' PASS in {dt:.3f}s', flush=True)
|
|
|
|
# Test 6: 100MB uint8
|
|
print('Test 6: 100MB uint8 .to(cuda)...', flush=True)
|
|
j = torch.randint(0, 255, (100*1024*1024,), dtype=torch.uint8)
|
|
t0 = time.time()
|
|
k = j.to('cuda')
|
|
torch.cuda.synchronize()
|
|
dt = time.time()-t0
|
|
print(f' PASS in {dt:.3f}s', flush=True)
|
|
|
|
# Test 7: GGMLTensor-like subclass
|
|
print('Test 7: Custom subclass .to(cuda)...', flush=True)
|
|
class FakeTensor(torch.Tensor):
|
|
def __new__(cls, data):
|
|
return torch.Tensor._make_subclass(cls, data)
|
|
def to(self, *args, **kwargs):
|
|
new = super().to(*args, **kwargs)
|
|
return new
|
|
|
|
raw = torch.randint(0, 255, (1024*1024,), dtype=torch.uint8)
|
|
ft = FakeTensor(raw)
|
|
t0 = time.time()
|
|
ft2 = ft.to('cuda')
|
|
torch.cuda.synchronize()
|
|
dt = time.time()-t0
|
|
print(f' PASS in {dt:.3f}s', flush=True)
|
|
|
|
print('ALL TESTS PASSED!', flush=True)
|
|
os._exit(0)
|
|
" 2>&1
|
|
echo "Exit code: $?"
|