Fabian
122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
"""Test GPU operations that occur between Linear layers in the Lumina2 model.
|
|
This tests LayerNorm, RMSNorm, SiLU, tensor reshaping, and manual SDPA."""
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import time, os, sys
|
|
|
|
os.environ['HSA_OVERRIDE_GFX_VERSION'] = '10.1.0'
|
|
os.environ['HSA_ENABLE_SDMA'] = '0'
|
|
os.environ['TORCHDYNAMO_DISABLE'] = '1'
|
|
|
|
print(f"Device: {torch.cuda.get_device_name(0)}", flush=True)
|
|
print(f"AMD_SERIALIZE_KERNEL={os.environ.get('AMD_SERIALIZE_KERNEL','unset')}", flush=True)
|
|
|
|
# Dimensions matching Lumina2 z_image_turbo
|
|
B, SEQ, DIM = 1, 1024, 3840
|
|
HEADS = 24
|
|
HEAD_DIM = DIM // HEADS # 160
|
|
dtype = torch.float16
|
|
|
|
def test_op(name, fn):
|
|
print(f"\n--- {name} ---", flush=True)
|
|
try:
|
|
t0 = time.time()
|
|
result = fn()
|
|
torch.cuda.synchronize()
|
|
dt = time.time() - t0
|
|
if hasattr(result, 'shape'):
|
|
print(f" OK: {dt*1000:.1f}ms, shape={list(result.shape)}, dtype={result.dtype}", flush=True)
|
|
else:
|
|
print(f" OK: {dt*1000:.1f}ms", flush=True)
|
|
return result
|
|
except Exception as e:
|
|
print(f" ERROR: {type(e).__name__}: {e}", flush=True)
|
|
return None
|
|
|
|
# Create typical model tensors
|
|
x = torch.randn(B, SEQ, DIM, dtype=dtype, device='cuda:0')
|
|
torch.cuda.synchronize()
|
|
print(f"Input tensor on GPU: {list(x.shape)}", flush=True)
|
|
|
|
# Test 1: F.layer_norm
|
|
weight_ln = torch.randn(DIM, dtype=dtype, device='cuda:0')
|
|
bias_ln = torch.randn(DIM, dtype=dtype, device='cuda:0')
|
|
test_op("F.layer_norm", lambda: F.layer_norm(x, (DIM,), weight_ln, bias_ln))
|
|
|
|
# Test 2: RMSNorm manual (common in Lumina2)
|
|
def rms_norm(x, w, eps=1e-5):
|
|
rms = torch.sqrt(torch.mean(x * x, dim=-1, keepdim=True) + eps)
|
|
return x / rms * w
|
|
test_op("RMSNorm manual", lambda: rms_norm(x, weight_ln))
|
|
|
|
# Test 3: SiLU activation
|
|
test_op("F.silu", lambda: F.silu(x))
|
|
|
|
# Test 4: GELU activation
|
|
test_op("F.gelu", lambda: F.gelu(x))
|
|
|
|
# Test 5: Tensor reshape (view + transpose)
|
|
def test_reshape():
|
|
q = x.view(B, SEQ, HEADS, HEAD_DIM).transpose(1, 2) # [B, HEADS, SEQ, HEAD_DIM]
|
|
return q
|
|
test_op("view+transpose", test_reshape)
|
|
|
|
# Test 6: Tensor repeat (GQA expansion)
|
|
def test_repeat():
|
|
kv = torch.randn(B, SEQ, 8, HEAD_DIM, dtype=dtype, device='cuda:0')
|
|
kv = kv.unsqueeze(3).repeat(1, 1, 1, 3, 1).flatten(2, 3)
|
|
return kv
|
|
test_op("unsqueeze+repeat+flatten", test_repeat)
|
|
|
|
# Test 7: torch.cat
|
|
def test_cat():
|
|
a = torch.randn(B, SEQ, DIM, dtype=dtype, device='cuda:0')
|
|
b = torch.randn(B, SEQ, DIM, dtype=dtype, device='cuda:0')
|
|
return torch.cat([a, b], dim=1)
|
|
test_op("torch.cat", test_cat)
|
|
|
|
# Test 8: Full manual SDPA (like our patch does)
|
|
def test_manual_sdpa():
|
|
Q = torch.randn(B, HEADS, SEQ, HEAD_DIM, dtype=dtype, device='cuda:0')
|
|
K = torch.randn(B, HEADS, SEQ, HEAD_DIM, dtype=dtype, device='cuda:0')
|
|
V = torch.randn(B, HEADS, SEQ, HEAD_DIM, dtype=dtype, device='cuda:0')
|
|
torch.cuda.synchronize()
|
|
|
|
scale = HEAD_DIM ** -0.5
|
|
attn = torch.matmul(Q, K.transpose(-2, -1)) * scale
|
|
torch.cuda.synchronize()
|
|
|
|
attn_max = attn.max(dim=-1, keepdim=True).values
|
|
exp_attn = torch.exp(attn - attn_max)
|
|
attn_w = exp_attn / exp_attn.sum(dim=-1, keepdim=True)
|
|
torch.cuda.synchronize()
|
|
|
|
out = torch.matmul(attn_w, V)
|
|
torch.cuda.synchronize()
|
|
return out
|
|
test_op("Manual SDPA (24 heads, 1024 seq, 160 dim)", test_manual_sdpa)
|
|
|
|
# Test 9: Einsum (used in rope)
|
|
def test_einsum():
|
|
pos = torch.randn(B, SEQ, dtype=torch.float32, device='cuda:0')
|
|
omega = torch.randn(HEAD_DIM // 2, dtype=torch.float32, device='cuda:0')
|
|
return torch.einsum("...n,d->...nd", pos, omega)
|
|
test_op("einsum (rope-like)", test_einsum)
|
|
|
|
# Test 10: Sequential like model forward (LayerNorm → Linear → SiLU → Linear)
|
|
def test_sequential():
|
|
h = F.layer_norm(x, (DIM,), weight_ln, bias_ln)
|
|
torch.cuda.synchronize()
|
|
w1 = torch.randn(DIM*4, DIM, dtype=dtype, device='cuda:0')
|
|
h = F.linear(h, w1)
|
|
torch.cuda.synchronize()
|
|
h = F.silu(h)
|
|
torch.cuda.synchronize()
|
|
w2 = torch.randn(DIM, DIM*4, dtype=dtype, device='cuda:0')
|
|
h = F.linear(h, w2)
|
|
torch.cuda.synchronize()
|
|
return h
|
|
test_op("Sequential: LN → Linear → SiLU → Linear", test_sequential)
|
|
|
|
print("\n=== ALL TESTS PASSED ===", flush=True)
|