矩阵乘测试显卡算力

100 阅读 0 评论 66 点赞

我是靠谱客的博主温柔草莓，这篇文章主要介绍矩阵乘测试显卡算力，现在分享给大家，希望可以做个参考。

由于pytorch和tensorflow不支持int8 int16的gemm，因此只能测试fp32 fp16 bf16等精度的tflops，如果要测试int8 int16精度下的数值，需要编写cublas脚本，目前不会CUDA编程，可参考大佬的脚本：

cuBLAS矩阵乘法性能分析（附代码示例）_算法码上来的博客-CSDN博客

lightseq/gemm_test.cpp at master · bytedance/lightseq · GitHub

lightseq/tests/gemm_test at 430e4e4018a049186db02ddbf015cbce39c14679 · bytedance/lightseq · GitHub

这里我用torch测试了3090的fp32 fp16 bf16，先给出3090实测结果，mnk大小一样，这里用n表示了：

依赖torch的测试脚本（李沐沐神提供）：

复制代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
import torch
print('Pytorch versiont:', torch.__version__)
print('CUDA versiont:', torch.version.cuda)
print('GPUtt:',torch.cuda.get_device_name())
import inspect
from collections import defaultdict
import pandas as pd
from torch.utils import benchmark
pd.options.display.precision = 3
def var_dict(*args):
callers_local_vars = inspect.currentframe().f_back.f_locals.items()
return dict([(name, val) for name, val in callers_local_vars if val is arg][0]
for arg in args)
def walltime(stmt, arg_dict, duration=3):
return benchmark.Timer(stmt=stmt, globals=arg_dict).blocked_autorange(
min_run_time=duration).median
matmul_tflops = defaultdict(lambda: {})
for n in [128, 512, 2048, 8192, 9216]:
for dtype in (torch.float32, torch.float16, torch.bfloat16):
a = torch.randn(n, n, dtype=dtype).cuda()
b = torch.randn(n, n, dtype=dtype).cuda()
t = walltime('a @ b', var_dict(a, b))
matmul_tflops[f'n={n}'][dtype] = 2*n**3 / t / 1e12
del a, b
# pd.DataFrame(matmul_tflops)
print(matmul_tflops)