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()
|