forked from ccf-ai-infra/GPUCodeForces
102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
import time
|
||
import torch
|
||
|
||
import linear_gelu_torchcode as torchcode
|
||
import linear_gelu_cudacode as cudacode
|
||
|
||
|
||
def _to_device(tensors, device):
|
||
return [t.to(device) for t in tensors]
|
||
|
||
|
||
def _copy_params(torch_model, cuda_model):
|
||
with torch.no_grad():
|
||
cuda_model.weight.copy_(torch_model.linear.weight)
|
||
cuda_model.bias.copy_(torch_model.linear.bias)
|
||
|
||
|
||
def _measure_gpu_seconds(model, args, iters=100):
|
||
start = torch.cuda.Event(enable_timing=True)
|
||
end = torch.cuda.Event(enable_timing=True)
|
||
torch.cuda.synchronize()
|
||
with torch.no_grad():
|
||
start.record()
|
||
for _ in range(iters):
|
||
_ = model(*args)
|
||
end.record()
|
||
torch.cuda.synchronize()
|
||
ms = start.elapsed_time(end) / iters
|
||
return ms / 1000.0
|
||
|
||
|
||
def run_benchmark():
|
||
if not torch.cuda.is_available():
|
||
print("CUDA 不可用")
|
||
return False, 0.0
|
||
|
||
torch.manual_seed(0)
|
||
device = torch.device("cuda")
|
||
|
||
init_kwargs = torchcode.get_init_inputs()
|
||
torch_model = torchcode.Model(**init_kwargs).to(device).eval()
|
||
cuda_model = cudacode.ModelNew(**init_kwargs).to(device).eval()
|
||
|
||
# 关闭编译，避免在当前平台上落到慢路径（CUTLASS 不可用）
|
||
|
||
# 参数对齐
|
||
_copy_params(torch_model, cuda_model)
|
||
|
||
# 准备输入
|
||
x = torchcode.get_inputs()
|
||
x, = _to_device([x], device)
|
||
|
||
print("-------------------- 精度对齐验证 --------------------")
|
||
with torch.no_grad():
|
||
# 预热
|
||
_ = torch_model(x)
|
||
_ = cuda_model(x)
|
||
|
||
# 正式测试
|
||
output_torch = torch_model(x)
|
||
output_cuda = cuda_model(x)
|
||
|
||
abs_diff = torch.abs(output_torch - output_cuda)
|
||
max_diff = torch.max(abs_diff).item()
|
||
mean_diff = torch.mean(abs_diff).item()
|
||
if max_diff < 1e-4 and mean_diff < 1e-5:
|
||
print(f"✅ 精度对齐：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
|
||
precision_flag = True
|
||
else:
|
||
print(f"❌ 精度不一致：最大误差 {max_diff:.6f}，平均误差 {mean_diff:.6f}")
|
||
precision_flag = False
|
||
|
||
print("\n-------------------- 性能加速比测试 --------------------")
|
||
num_iterations = 200
|
||
|
||
# 预热
|
||
for _ in range(10):
|
||
_ = torch_model(x)
|
||
_ = cuda_model(x)
|
||
|
||
# 可选：允许 TF32（若硬件支持），提升矩阵乘性能
|
||
try:
|
||
torch.backends.cuda.matmul.allow_tf32 = True
|
||
torch.set_float32_matmul_precision("medium")
|
||
except Exception:
|
||
pass
|
||
|
||
# PyTorch计时（CUDA Events，秒）
|
||
torch_time = _measure_gpu_seconds(torch_model, (x,), iters=num_iterations)
|
||
# 优化版计时（CUDA Events，秒）
|
||
cuda_time = _measure_gpu_seconds(cuda_model, (x,), iters=num_iterations)
|
||
|
||
print(f"PyTorch内置Linear+GELU平均执行时间: {torch_time:.6f}秒")
|
||
print(f"自定义CUDA Linear+GELU平均执行时间: {cuda_time:.6f}秒")
|
||
speedup = torch_time / cuda_time if cuda_time > 0 else 0.0
|
||
print(f"加速比 (Speedup): {speedup:.2f}x")
|
||
|
||
return precision_flag, speedup
|
||
|
||
|
||
if __name__ == "__main__":
|
||
precision_flag, speedup = run_benchmark() |