专栏名称: GiantPandaCV

专注于机器学习、深度学习、计算机视觉、图像处理等多个方向技术分享。团队由一群热爱技术且热衷于分享的小伙伴组成。我们坚持原创，每天一到两篇原创技术分享。希望在传播知识、分享知识的同时能够启发你，大家一起共同进步(･ω<)☆

GPU 矩阵乘实际可达最大FLOPS测量工具

GiantPandaCV · 公众号 · 3D · 2024-11-28 12:05

正文

前言

在 https://github.com/stas00/ml-engineering/blob/master/compute/accelerator/benchmarks/mamf-finder.py 这里有一个可以在任意 GPU 上搜索最大 Matmul FLOPS 的脚本，使用起来非常简单，依赖也只有 PyTorch 库。之前只支持 FP16 数据类型，我帮助完善了一下支持 PyTorch 的各种数据类型，这里安利一下，只需要下载这个脚本就可以了。下面是使用方法：

FP16 使用示例

在下面的范围中,N是rReduce维度,使得(MxN)*(NxK)=(MxK),我们会打印出测得最高TFLOPS的MxNxK形状。

默认情况下,我们对每个形状使用50次预热迭代和100次测量迭代,然后选择最快的结果(而不是平均值)。你可以通过参数--num_warmup_iterations和--num_iterations分别更改迭代次数。

--dtype可以指定测试的数据类型。

这里我们执行torch.mm(MxN,NxK) -> MxK

我推荐大家使用下面的第2条。

快速运行(1分钟以内) - 应该能达到最大可达结果的80-90% - 适合快速尝试,但不足以获得高精度测量。

python3 mamf-finder.py --m_range 0 20480 256 --n 4096 --k 4096 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

更详尽的搜索(将花费更长时间) - 但你可以在运行足够长时间后按Ctrl-C终止,并获得到目前为止的最佳结果 :

python3 mamf-finder.py --m_range 0 5376 256 --n_range 0 5376 256 --k_range 0 5376 256 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

一个超长的穷举搜索（可能需要几天时间）- 但你可以在运行足够长时间后按Ctrl-C终止它，并获得到目前为止的最佳结果：

python3 mamf-finder.py --m_range 0 20480 256 --n_range 0 20480 256 --k_range 0 20480 256 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

如果你想测量训练中使用的特定形状,请使用确切的形状,而不是范围。例如,假设你想测量1024x1024x1024 - 你可以运行:

python3 mamf-finder.py --m 1024 --n 1024 --k 1024 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

加速器特定范围搜索建议

然而，不同的加速器似乎有不同的形状范围可以达到最佳 TFLOPS，因此很难建议一个适用于所有加速器的范围。相反，这里根据实验和贡献者的建议提供一些建议：

A100 + MI300X

python3 mamf-finder.py --m_range 0 5376 256 --n_range 0 5376 256 --k_range 0 5376 256 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

H100

python3 mamf-finder.py --m_range 0 20480 256 --n_range 0 20480 256 --k_range 0 20480 256 --output_file=$(date +"%Y-%m-%d-%H:%M:%S").txt

这里的 Efficiency 就是矩阵乘可以达到的最大 MFU。

FP8 使用示例

下面展示一下我在4090上对FP8类型搜索的最大MatMul TFLOPS

Benchmark started on 2024-11-25 01:03:20

** Command line:
/usr/bin/python3 mamf-finder.py --m_range 0 5376 256 --n_range 0 5376 256 --k_range 0 5376 256 --output_file=2024-11-25-01:03:19.txt --dtype float8_e4m3fn

** Dtype: torch.float8_e4m3fn

** Platform/Device info:
Linux benchmark-bbuf-7ff8fbb655-2lvqc 5.10.134-16.1.al8.x86_64 #1 SMP Thu Dec 7 14:11:24 UTC 2023 x86_64 x86_64
_CudaDeviceProperties(name='NVIDIA GeForce RTX 4090', major=8, minor=9, total_memory=24217MB, multi_processor_count=128, uuid=0318f41e-f00d-953d-b811-1872eab52308, L2_cache_size=72MB)

** Critical software versions:
torch=2.5.1+cu124
cuda=12.4

** Additional notes:


--------------------------------------------------------------------------------


The best outcome was 300.4TFLOPS @ 3328x5120x2816 (MxNxK) (tried 8000 shapes)
Elapsed time: 0:03:14

NVIDIA GeForce RTX 4090 https://images.nvidia.com/aem-dam/Solutions/geforce/ada/nvidia-ada-gpu-architecture.pdf 的理论 FP8 FLOPS为 330 TFLOPS

使用 mamf-finder 工具搜索出的实际最大 FLOPS 为 300.4 TFLOPS，达到最大理论峰值的 91.0% 。