ggml : force F32 precision for ggml_mul_mat

ggerganov · Dec 19, 2023 · a40f611 · a40f611
1 parent a7aee47
commit a40f611
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 20 deletions.
diff --git a/ggml-cuda.cu b/ggml-cuda.cu
@@ -7389,7 +7389,7 @@ inline void ggml_cuda_op_mul_mat_cublas(
 
     const int compute_capability = g_compute_capabilities[id];
 
-    if (compute_capability >= CC_VOLTA && (src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && ggml_is_contiguous(src0) && row_diff == src0->ne[1] && dst->op_params[0] == GGML_PREC_DEFAULT) {
+    if (compute_capability >= CC_VOLTA && (src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && ggml_is_contiguous(src0) && row_diff == src0->ne[1]) {
         // convert src0 and src1 to fp16, multiply as fp16, convert dst to fp32
         half * src0_as_f16 = nullptr;
         size_t src0_as = 0;
@@ -7412,26 +7412,47 @@ inline void ggml_cuda_op_mul_mat_cublas(
             to_fp16_cuda(src1_ddf_i, src1_as_f16, ne, stream);
         }
         const half * src1_ptr = src1->type == GGML_TYPE_F16 ? (const half *) src1_ddf_i : src1_as_f16;
-        size_t dst_as = 0;
-        half * dst_f16 = (half *) ggml_cuda_pool_malloc(row_diff*src1_ncols * sizeof(half), &dst_as);
 
-        const half alpha_f16 = 1.0f;
-        const half beta_f16 = 0.0f;
-
-        CUBLAS_CHECK(cublasSetStream(g_cublas_handles[id], stream));
-        CUBLAS_CHECK(
-            cublasGemmEx(g_cublas_handles[id], CUBLAS_OP_T, CUBLAS_OP_N,
-                    row_diff, src1_ncols, ne10,
-                    &alpha_f16, src0_ptr, CUDA_R_16F, ne00,
-                                src1_ptr, CUDA_R_16F, ne10,
-                    &beta_f16,   dst_f16, CUDA_R_16F, ldc,
-                    CUBLAS_COMPUTE_16F,
-                    CUBLAS_GEMM_DEFAULT_TENSOR_OP));
-
-        const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_F16);
-        to_fp32_cuda(dst_f16, dst_dd_i, row_diff*src1_ncols, stream);
-
-        ggml_cuda_pool_free(dst_f16, dst_as);
+        switch (dst->op_params[0]) {
+            case GGML_PREC_DEFAULT:
+                {
+                    size_t dst_as = 0;
+                    half * dst_f16 = (half *) ggml_cuda_pool_malloc(row_diff*src1_ncols * sizeof(half), &dst_as);
+
+                    const half alpha_f16 = 1.0f;
+                    const half beta_f16  = 0.0f;
+
+                    CUBLAS_CHECK(cublasSetStream(g_cublas_handles[id], stream));
+                    CUBLAS_CHECK(
+                            cublasGemmEx(g_cublas_handles[id], CUBLAS_OP_T, CUBLAS_OP_N,
+                                row_diff,   src1_ncols, ne10,
+                                &alpha_f16, src0_ptr, CUDA_R_16F, ne00,
+                                            src1_ptr, CUDA_R_16F, ne10,
+                                &beta_f16,  dst_f16,  CUDA_R_16F, ldc,
+                                CUBLAS_COMPUTE_16F,
+                                CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+                    const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_F16);
+                    to_fp32_cuda(dst_f16, dst_dd_i, row_diff*src1_ncols, stream);
+
+                    ggml_cuda_pool_free(dst_f16, dst_as);
+                } break;
+            case GGML_PREC_F32:
+                {
+                    const float alpha_f32 = 1.0f;
+                    const float beta_f32  = 0.0f;
+
+                    CUBLAS_CHECK(cublasSetStream(g_cublas_handles[id], stream));
+                    CUBLAS_CHECK(
+                        cublasGemmEx(g_cublas_handles[id], CUBLAS_OP_T, CUBLAS_OP_N,
+                                row_diff, src1_ncols, ne10,
+                                &alpha_f32, src0_ptr, CUDA_R_16F, ne00,
+                                            src1_ptr, CUDA_R_16F, ne10,
+                                &beta_f32,  dst_dd_i, CUDA_R_32F, ldc,
+                                CUBLAS_COMPUTE_32F,
+                                CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+                } break;
+        }
 
         if (src0_as != 0) {
             ggml_cuda_pool_free(src0_as_f16, src0_as);

diff --git a/ggml.c b/ggml.c
@@ -4090,6 +4090,12 @@ struct ggml_tensor * ggml_mul_mat(
     const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] };
     struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne);
 
+    // TMP: force f32 precision
+    {
+        const int32_t prec_i32 = GGML_PREC_F32;
+        ggml_set_op_params_i32(result, 0, prec_i32);
+    }
+
     result->op   = GGML_OP_MUL_MAT;
     result->grad = is_node ? ggml_dup_tensor(ctx, result) : NULL;
     result->src[0] = a;