cublasGemmEx函数应用-探究8bit矩阵乘

来源:互联网 发布:淘宝基础版装修规格 编辑:程序博客网 时间:2024/06/13 11:42

介绍

cublasGemmEx 是CUDA8.0中cuBLAS新出的函数,是cublasgemm()类函数的扩展,也是目前来看功能最强大的矩阵乘函数了。该函数另一强大之处在于支持多种计算模式(compute type),其中就包括CUDA 8.0新出的FP16和INT8。但是该函数的文档并不太健全,最近在使用这个函数实现INT8矩阵乘的时候就碰见坑了,照着文档用就是报错,找NVIDIA的工程师才给解决。下面总结一下使用经验,把坑填上,以防大家再踩。

函数原型

cublasStatus_t cublasGemmEx(cublasHandle_t handle,                             cublasOperation_t transa,                             cublasOperation_t transb,                             int m,                             int n,                             int k,                             const void *alpha,                             const void *A,                             cudaDataType_t Atype,                             int lda,                             const void *B,                             cudaDataType_t Btype,                             int ldb,                             const void *beta,                             void *C,                             cudaDataType_t Ctype,                             int ldc,                             cudaDataType_t computeType,                             cublasGemmAlgo_t algo)

跟cublasSgemm长的比较像,但是多了这么几个参数,Atype,Btype,Ctype,computeType和algo。

这个函数的核心就是计算模式(computeType),computeType支持以下类型:

computeType 解释 CUDA_R_16F FP16计算模式,输入输出都是FP16 CUDA_R_32F FP32计算模式,这个比较强大,输入可以是FP16、INT8和FP32 CUDA_R_32I INT8计算模式,也是本文着重要讲的模式 CUDA_R_64F FP64计算模式 CUDA_C_32F CUDA_C_64F

每个computeType支持的输入类型和输出类型在cublasGemmEx文档中写的非常清楚,照着用就行了。但是,有一个隐含的坑就在CUDA_R_32I计算模式里。

正常按照 char *A, char *B, int *C是会报错CUBLAS_STATUS_NOT_SUPPORTED,这个错误官方的解释是“the combination of the parameters Atype, Btype and Ctype and the algorithm type, algo is not supported”,大概意思就是Atype,Btype,Ctype,和algo不匹配。但是明明是按文档上写的啊,因为错误根本不在这里。

解决办法

错误的原因是,如果要使用CUDA_R_32I计算模式,那么alpha和beta这两个参数也必须是int类型且必须是0或者1……神坑啊。

PS:CUDA_R_32I计算模式下,cublasGemmAlgo_t 参数好像也只支持前7种,这个在文档里也没说。

CUDA_R_32I与CUDA_R_32F计算对比结果

这里多说一点INT8矩阵乘计算模式吧,CUDA_R_32I计算模式里调用CUDA 8.0新出的INT8计算接口-dp4a,按照官方的理论,dp4a这个函数会将四个char组合成一个int进行乘法运算,将4次乘法和3次加法减少为一次高级指令,从而提高性能。

我的实验结果表明,CUDA_R_32I模式与CUDA_R_32F模式相比,最快能提高3.2倍(与矩阵的大小有关),同时能将数据压缩75%,这是一个非常可观的收益了。

但是FP32(float)量化成INT8(char)肯定是会有精度损失的,对INT8有兴趣的可以关注NVIDIA新出的TensorRT2.0,该库能够在一些情况下保持较高的精度实现INT8加速。

TensorRT给的资料也比较少,坑也特别多,因此我开了一个TensorRT_Tutorial,欢迎志同道合者一起参与。

Figure 2: Illustration of the All-Reduce collective.

原创粉丝点击