Leverage mmap for CUDA loading

Author: howard0su

ggml-cuda.cu

@@ -926,32 +926,13 @@ void ggml_cuda_transform_tensor(ggml_tensor * tensor) {
     tensor->backend = GGML_BACKEND_CUDA;
 }
 
-void ggml_cuda_load_data(const char * fname, struct ggml_tensor * tensor, const size_t offset) {
-    FILE * fp = fopen(fname, "rb");
-
+void ggml_cuda_load_data(const void * buf_host, struct ggml_tensor * tensor) {
     const size_t size = ggml_nbytes(tensor);
 
     void * buf;
     CUDA_CHECK(cudaMalloc(&buf, size));
-    void * buf_host = malloc(size);
-
-#ifdef _WIN32
-    int ret = _fseeki64(fp, (__int64) offset, SEEK_SET);
-#else
-    int ret = fseek(fp, (long) offset, SEEK_SET);
-#endif
-    GGML_ASSERT(ret == 0); // same
-
-    size_t ret2 = fread(buf_host, size, 1, fp);
-    if (ret2 != 1) {
-        fprintf(stderr, "unexpectedly reached end of file");
-        exit(1);
-    }
 
     cudaMemcpy(buf, buf_host, size, cudaMemcpyHostToDevice);
-    cudaDeviceSynchronize();
 
     tensor->data = buf;
-    free(buf_host);
-    fclose(fp);
 }

ggml-cuda.h

@@ -16,7 +16,7 @@ void * ggml_cuda_host_malloc(size_t size);
 void   ggml_cuda_host_free(void * ptr);
 
 void ggml_cuda_transform_tensor(struct ggml_tensor * tensor);
-void ggml_cuda_load_data(const char * fname, struct ggml_tensor * tensors, size_t offset);
+void ggml_cuda_load_data(const void * data, struct ggml_tensor * tensors);
 
 #ifdef  __cplusplus
 }

llama.cpp

@@ -700,28 +700,31 @@ struct llama_model_loader {
 
         if (use_mmap) {
             mapping.reset(new llama_mmap(&file_loaders.at(0)->file, prefetch_size));
-            if (!lmlock) {
-                // Don't call the callback since the actual loading will be lazy
-                // and we can't measure it.
-                progress_callback = NULL;
-            }
             if (lmlock) {
                 lmlock->init(mapping->addr);
             }
         }
 
         size_t done_size = 0;
         for (llama_load_tensor & lt : tensors_map.tensors) {
-            if (lt.ggml_tensor->backend != GGML_BACKEND_CPU) {
-                continue;
-            }
             if (progress_callback) {
                 progress_callback((float) done_size / data_size, progress_callback_user_data);
             }
             LLAMA_ASSERT(lt.ggml_tensor); // unused tensors should have been caught by load_data already
             lt.data = (uint8_t *) lt.ggml_tensor->data;
             load_data_for(lt);
-            lt.ggml_tensor->data = lt.data;
+            switch(lt.ggml_tensor->backend) {
+                case GGML_BACKEND_CPU:
+                    lt.ggml_tensor->data = lt.data;
+                    break;
+#ifdef GGML_USE_CUBLAS
+                case GGML_BACKEND_CUDA:
+                    ggml_cuda_load_data(lt.data, lt.ggml_tensor);
+                    break;
+#endif
+                default:
+                    continue;
+            }
             done_size += lt.size;
             if (use_mmap && lmlock) {
                 lmlock->grow_to(done_size);
@@ -1106,28 +1109,7 @@ static void llama_model_load_internal(
 
     ml->load_all_data(progress_callback, progress_callback_user_data, use_mlock ? &lctx.model.mlock_mmap : NULL);
 
-#ifdef GGML_USE_CUBLAS
-    {
-        size_t done_size = 0;
-        size_t data_size = 0;
-        for (llama_load_tensor & lt : ml->tensors_map.tensors) {
-            data_size += lt.size;
-            if (lt.ggml_tensor->backend == GGML_BACKEND_CPU) {
-                done_size += lt.size;
-            }
-        }
-        for (llama_load_tensor & lt : ml->tensors_map.tensors) {
-            if (lt.ggml_tensor->backend != GGML_BACKEND_CUDA) {
-                continue;
-            }
-            if (progress_callback) {
-                progress_callback((float) done_size / data_size, progress_callback_user_data);
-            }
-            ggml_cuda_load_data(fname.c_str(), lt.ggml_tensor, lt.shards.at(0).file_off);
-            done_size += lt.size;
-        }
-    }
-#elif defined(GGML_USE_CLBLAST)
+#ifdef GGML_USE_CLBLAST
     {
         const int n_gpu = std::min(n_gpu_layers, int(hparams.n_layer));