FBGEMM简介
FBGEMM(Facebook GEneral Matrix Multiplication)是Facebook开源的一款低精度、高性能矩阵乘法和卷积库,专门为服务器端推理而设计。作为PyTorch量化操作符在x86机器上的后端实现,FBGEMM为深度学习模型的低精度推理提供了强大的支持。
FBGEMM的主要特点
-
针对低精度数据进行了专门优化,不同于传统的科学计算线性代数库(通常使用FP32或FP64精度)。
-
为小批量提供高效的低精度通用矩阵乘法(GEMM)实现。
-
支持行级量化和异常值感知量化等精度损失最小化技术。
-
充分利用融合机会,克服了低精度矩阵乘法中带宽受限操作的独特挑战。
-
作为PyTorch量化操作符在x86机器上的后端实现。
FBGEMM的技术原理
FBGEMM采用了类似于其他针对FP32和FP64数据类型的高性能GEMM实现的方法,但专门针对低精度数据进行了优化。其核心实现基于以下几个关键技术:
1. 矩阵分块和缓存优化
FBGEMM使用多层循环结构,将矩阵分成适合CPU缓存层次结构的小块:
Loop1 for ic = 0 to M-1 in steps of MCB
Loop2 for kc = 0 to K-1 in steps of KCB
//Pack MCBxKCB block of A
Loop3 for jc = 0 to N-1 in steps of NCB
//Pack KCBxNCB block of B
//--------------------Macro Kernel------------
Loop4 for ir = 0 to MCB-1 in steps of MR
Loop5 for jr = 0 to NCB-1 in steps of NR
//--------------------Micro Kernel------------
Loop6 for k = 0 to KCB-1 in steps of 1
//update MRxNR block of C matrix
这种结构可以充分利用CPU的多级缓存,提高数据访问效率。
2. 矩阵打包
FBGEMM在计算过程中对矩阵块进行重新组织(打包),以实现内部核心的顺序数据访问。这对于现代硬件架构上实现高效带宽至关重要。
3. 计算与打包融合
FBGEMM利用打包操作的带宽受限特性,将一些简单的计算操作与打包过程融合。例如,在打包A矩阵时同时计算行偏移,避免了对A矩阵数据的多次访问。
4. 动态代码生成
FBGEMM采用"一刀切"的方法不适用于所有情况,因此实现了动态生成矩阵形状特定的向量化代码。这种方法可以为不同大小和形状的矩阵生成高效的内核。
5. INT8量化和INT16累加
FBGEMM支持INT8矩阵乘法并累加到INT16,以获得更好的性能。INT8 FMA与INT16累加通过vpmaddubsw
和vpaddsw
向量指令的组合实现。
6. 异常值感知量化
为了避免INT16累加时的溢出/饱和问题,FBGEMM采用了异常值感知量化技术。将B矩阵拆分为B = B' + B_sparse,其中B'只包含小幅值,大幅值被分离为B_sparse。
FBGEMM的性能优势
FBGEMM在Intel(R) Xeon(R) CPU E5-2680 v4上进行了性能基准测试,结果显示:
-
对于计算密集型场景,INT8 GEMM with INT16累加可以达到FP32理论峰值的2倍。
-
对于带宽受限场景,INT8 GEMM with INT32累加在小批量情况下表现优异,甚至超过了FP32理论屋顶线性能。
FBGEMM的应用场景
FBGEMM在Facebook的多个AI服务中得到了广泛应用,带来了显著的性能提升:
-
将英语到西班牙语的翻译速度提高了1.3倍。
-
将Feed中使用的推荐系统的DRAM带宽使用降低了40%。
-
在Rosetta(Facebook用于理解图像和视频中文本的机器学习系统)中,将字符检测速度提高了2.4倍。
FBGEMM的未来发展
FBGEMM团队计划在以下几个方向继续改进:
-
为更多Facebook内部使用的模型实现量化推理。
-
优化对组卷积(groupwise convolution)的支持。
-
与PyTorch进行更深入的集成。
-
进一步提高效率,如合并深度卷积与1x1卷积。
-
改进性能调试支持。
结论
FBGEMM作为一个开源项目,为服务器端深度学习推理提供了高效的低精度矩阵运算解决方案。它不仅大幅提升了Facebook内部AI服务的性能,也为整个深度学习社区提供了宝贵的工具。随着量化推理技术的不断发展,FBGEMM有望在未来发挥更大的作用,推动AI应用在服务器端的进一步优化和普及。
对于有兴趣深入了解或贡献FBGEMM项目的开发者,可以访问FBGEMM的GitHub仓库获取更多信息。FBGEMM团队也欢迎社区成员通过GitHub Issues、Discussions或PyTorch Slack的#fbgemm频道与他们交流。