前言

在 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. 快速运行(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  

2. **更详尽的搜索(将花费更长时间) - 但你可以在运行足够长时间后按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  

3. 一个超长的穷举搜索（可能需要几天时间）- 但你可以在运行足够长时间后按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  

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

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

5. 加速器特定范围搜索建议

然而，不同的加速器似乎有不同的形状范围可以达到最佳 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% 。

公众号后台回复“数据集”获取100+深度学习各方向资源整理