FBGEMM: Facebook开源的高性能低精度矩阵运算库

FBGEMM简介

FBGEMM(Facebook GEneral Matrix Multiplication)是Facebook开源的一款低精度、高性能矩阵乘法和卷积库,专门为服务器端推理而设计。作为PyTorch量化操作符在x86机器上的后端实现,FBGEMM为深度学习模型的低精度推理提供了强大的支持。

FBGEMM的主要特点

1. 针对低精度数据进行了专门优化,不同于传统的科学计算线性代数库(通常使用FP32或FP64精度)。

2. 为小批量提供高效的低精度通用矩阵乘法(GEMM)实现。

3. 支持行级量化和异常值感知量化等精度损失最小化技术。

4. 充分利用融合机会,克服了低精度矩阵乘法中带宽受限操作的独特挑战。

5. 作为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上进行了性能基准测试,结果显示:

1. 对于计算密集型场景,INT8 GEMM with INT16累加可以达到FP32理论峰值的2倍。

2. 对于带宽受限场景,INT8 GEMM with INT32累加在小批量情况下表现优异,甚至超过了FP32理论屋顶线性能。

FBGEMM的应用场景

FBGEMM在Facebook的多个AI服务中得到了广泛应用,带来了显著的性能提升:

1. 将英语到西班牙语的翻译速度提高了1.3倍。

2. 将Feed中使用的推荐系统的DRAM带宽使用降低了40%。

3. 在Rosetta(Facebook用于理解图像和视频中文本的机器学习系统)中,将字符检测速度提高了2.4倍。

FBGEMM的未来发展

FBGEMM团队计划在以下几个方向继续改进:

1. 为更多Facebook内部使用的模型实现量化推理。

2. 优化对组卷积(groupwise convolution)的支持。

3. 与PyTorch进行更深入的集成。

4. 进一步提高效率,如合并深度卷积与1x1卷积。

5. 改进性能调试支持。

结论

FBGEMM作为一个开源项目,为服务器端深度学习推理提供了高效的低精度矩阵运算解决方案。它不仅大幅提升了Facebook内部AI服务的性能,也为整个深度学习社区提供了宝贵的工具。随着量化推理技术的不断发展,FBGEMM有望在未来发挥更大的作用,推动AI应用在服务器端的进一步优化和普及。

对于有兴趣深入了解或贡献FBGEMM项目的开发者,可以访问FBGEMM的GitHub仓库获取更多信息。FBGEMM团队也欢迎社区成员通过GitHub Issues、Discussions或PyTorch Slack的#fbgemm频道与他们交流。