Update sha, simplified grad calculation, adapted mbwq linear cuda kernel.

Author: Haojin Yang

bitorch_engine/layers/qlinear/nbit/cuda/mbwq_layer.py

@@ -109,9 +109,8 @@ def backward(ctx: torch.autograd.function.BackwardCFunction,
         grad_input = output_gradient.mm(weights.t()) # (m, n)*(n, k) = (m, k)
         #======================================================================================================#
 
-        # (n, m) * (m, k) = (n, k)
         if qweight.requires_grad: # This additional check is required by peft training.
-            qweight.privileged_grad = output_gradient.t().mm(input).t()  # (k, n)
+            qweight.privileged_grad = input.t().mm(output_gradient)  # (k, m) * (m, n) = (k, n)
 
         grad_input = unflatten_x(grad_input, shape)
 

bitorch_engine/layers/qlinear/nbit/cuda/mbwq_linear_cuda_kernel.cu

@@ -749,7 +749,6 @@ torch::Tensor mbwq_linear_q4_forward_cuda(
     int bits
 ){
     const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
-    cublasHandle_t cublas_handle = at::cuda::getCurrentCUDABlasHandle();
 
     TORCH_CHECK(x.dtype() == torch::kHalf);
     TORCH_CHECK(x.size(1) == qweight.size(0) * (32 / bits));
@@ -770,16 +769,8 @@ torch::Tensor mbwq_linear_q4_forward_cuda(
 									               group_size,
 									               bits,
 									               q_perm);
-
-        const half alpha = __float2half(1.0f);
-        const half beta = __float2half(0.0f);
-        cublasHgemm(cublas_handle,
-                    CUBLAS_OP_N,
-                    CUBLAS_OP_N,
-                    size_n, size_m,                                    size_k,
-                    &alpha, reinterpret_cast<half *>(fp_w.data_ptr()), size_n,
-                            reinterpret_cast<half *>(x.data_ptr()),    size_k,
-                    &beta,  reinterpret_cast<half *>(out.data_ptr()),  size_n);
+		// indirectly use cublas through torch matmul api
+        out = torch::matmul(x, fp_w.to(option_output));
 
 	}else{
 
@@ -943,7 +934,6 @@ torch::Tensor mbwq_linear_exl2_forward_cuda(
     bool use_cublas
 ){
     const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
-    cublasHandle_t cublas_handle = at::cuda::getCurrentCUDABlasHandle();
     TORCH_CHECK(x.dtype() == torch::kHalf);
 
 	int size_m = x.size(0);       // m
@@ -963,15 +953,8 @@ torch::Tensor mbwq_linear_exl2_forward_cuda(
 									               qgroup_map,
 									               rows);
 
-        const half alpha = __float2half(1.0f);
-        const half beta = __float2half(0.0f);
-        cublasHgemm(cublas_handle,
-                    CUBLAS_OP_N,
-                    CUBLAS_OP_N,
-                    size_n, size_m,                                    size_k,
-                    &alpha, reinterpret_cast<half *>(fp_w.data_ptr()), size_n,
-                            reinterpret_cast<half *>(x.data_ptr()),    size_k,
-                    &beta,  reinterpret_cast<half *>(out.data_ptr()),  size_n);
+        // indirectly use cublas through torch matmul api
+        out = torch::matmul(x, fp_w.to(option_output));
 
 	}else{
 	    int rows_8 = rows[0];

bitorch_engine/layers/qlinear/nbit/cuda/mpq_layer.py

@@ -100,9 +100,8 @@ def backward(ctx: torch.autograd.function.BackwardCFunction,
                                                   output_gradient, a_bit, w_bit, asym)
         #==================================================================#
 
-        # (n, m) * (m, k) = (n, k)
         if qweight.requires_grad:  # This additional check is required by peft training.
-            qweight.privileged_grad = output_gradient.t().mm(input).t()  # (k, n)
+            qweight.privileged_grad = input.t().mm(output_gradient)  # (k, m) * (m, n) = (k, n)
 
         grad_input = unflatten_x(grad_input, shape)
 

docker/build_scripts/install_modified_pytorch.sh

@@ -24,7 +24,7 @@ fi
 if [ "${from_image}" == "pytorch/pytorch:2.2.0-cuda12.1-cudnn8-devel" ]; then
     gdrive_id="1LjFNImboq8QeFSompMS2gPjBRYtP2Dsz"
     file="torch-2.2.2-cp310-cp310-linux_x86_64.whl"
-    checksum="2a5953dab7be6c1640112e38ae7519ad88180d9fa79faab6c86dbee6b1cc210e"
+    checksum="bcc0ba7f121ee2f42ed0a59f01d4e3d70f82a8981be0be25c5e0fe0635a54b2d"
 fi
 #if [ "${from_image}" == "pytorch/pytorch:X.X.X-cudaXX.X-cudnn8-devel" ]; then
 #    gdrive_id="xxx"