Fabian
101 lines
3.3 KiB
Python
101 lines
3.3 KiB
Python
import torch
|
|
import time
|
|
import os
|
|
import sys
|
|
|
|
print(f'Device: {torch.cuda.get_device_name(0)}')
|
|
print(f'Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB')
|
|
print()
|
|
|
|
def timed_op(name, fn, timeout=120):
|
|
"""Run operation with timeout detection"""
|
|
print(f'=== {name} ===', flush=True)
|
|
try:
|
|
torch.cuda.synchronize()
|
|
t = time.time()
|
|
result = fn()
|
|
torch.cuda.synchronize()
|
|
elapsed = time.time() - t
|
|
if elapsed > timeout:
|
|
print(f' SLOW: {elapsed:.1f}s (>{timeout}s)', flush=True)
|
|
else:
|
|
print(f' OK: {elapsed:.3f}s', flush=True)
|
|
return True
|
|
except Exception as e:
|
|
print(f' FAIL: {e}', flush=True)
|
|
return False
|
|
|
|
# Lumina2-like dimensions
|
|
# Hidden dim ~3072, heads ~24, head_dim ~128
|
|
# Latent 64x64 = 4096 tokens for 512x512 image
|
|
|
|
# Test 1: QKV projection (Linear 3072 -> 3*3072)
|
|
def test_qkv():
|
|
lin = torch.nn.Linear(3072, 9216).cuda().float() # 3*3072
|
|
x = torch.randn(1, 4096, 3072, device='cuda', dtype=torch.float32)
|
|
y = lin(x)
|
|
del lin, x, y
|
|
timed_op('QKV projection (1x4096x3072 -> 9216)', test_qkv)
|
|
|
|
# Test 2: Attention scores (24 heads, 4096x4096)
|
|
def test_attn_scores():
|
|
q = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
|
|
k = torch.randn(1, 24, 4096, 128, device='cuda', dtype=torch.float32)
|
|
scores = torch.matmul(q, k.transpose(-2, -1)) # 24x4096x4096 = 1.5GB fp32!
|
|
del q, k, scores
|
|
timed_op('Attention scores 24x4096x4096', test_attn_scores)
|
|
|
|
# Test 2b: Smaller attention (fewer heads)
|
|
def test_attn_small():
|
|
q = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
|
|
k = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
|
|
scores = torch.matmul(q, k.transpose(-2, -1))
|
|
del q, k, scores
|
|
timed_op('Attention scores 8x1024x1024 (smaller)', test_attn_small)
|
|
|
|
# Test 3: Attention value multiply
|
|
def test_attn_values():
|
|
attn = torch.randn(1, 8, 1024, 1024, device='cuda', dtype=torch.float32)
|
|
v = torch.randn(1, 8, 1024, 128, device='cuda', dtype=torch.float32)
|
|
out = torch.matmul(attn, v)
|
|
del attn, v, out
|
|
timed_op('Attn * values 8x1024x1024 @ 8x1024x128', test_attn_values)
|
|
|
|
# Test 4: FFN (3072 -> 12288 -> 3072)
|
|
def test_ffn():
|
|
up = torch.nn.Linear(3072, 12288).cuda().float()
|
|
down = torch.nn.Linear(12288, 3072).cuda().float()
|
|
x = torch.randn(1, 4096, 3072, device='cuda', dtype=torch.float32)
|
|
h = torch.nn.functional.gelu(up(x))
|
|
y = down(h)
|
|
del up, down, x, h, y
|
|
timed_op('FFN 3072->12288->3072 (4096 tokens)', test_ffn)
|
|
|
|
# Test 5: Conv2d 128ch (VAE-like)
|
|
def test_vae_conv():
|
|
conv = torch.nn.Conv2d(128, 128, 3, padding=1).cuda().float()
|
|
x = torch.randn(1, 128, 64, 64, device='cuda', dtype=torch.float32)
|
|
y = conv(x)
|
|
del conv, x, y
|
|
timed_op('VAE Conv2d 128ch 64x64', test_vae_conv)
|
|
|
|
# Test 6: Memory pressure test
|
|
def test_mem():
|
|
tensors = []
|
|
total = 0
|
|
for i in range(10):
|
|
t = torch.randn(256, 256, 256, device='cuda', dtype=torch.float32) # 64MB each
|
|
tensors.append(t)
|
|
total += 64
|
|
print(f' Allocated {total}MB on GPU', flush=True)
|
|
for t in tensors:
|
|
del t
|
|
del tensors
|
|
torch.cuda.empty_cache()
|
|
timed_op('Memory: allocate 640MB', test_mem)
|
|
|
|
print()
|
|
print('=== ALL LUMINA2 TESTS DONE ===', flush=True)
|
|
torch.cuda.synchronize()
|
|
os._exit(0)
|