Merge pull request #55 from ravenscroftj/feature/gpu_layers

WIP: Integrate more direct GPU support

.github/workflows/docker-image.yml

@@ -53,14 +53,14 @@ jobs:
             runtime_base: nvidia/cuda:12.2.0-runtime-ubuntu22.04
             cmake_args: -DGGML_CUBLAS=ON -DCMAKE_CUDA_COMPILER=/usr/local/cuda/bin/nvcc
 
-          - tag: -clblast
-            dockerfile: ./Dockerfile.default
-            platforms: linux/amd64
-            build_base: ubuntu:22.04
-            runtime_base: ubuntu:22.04
-            runtime_deps: libclblast1
-            extra_deps: libclblast-dev
-            cmake_args: -DGGML_CLBLAST=On
+          # - tag: -clblast
+          #   dockerfile: ./Dockerfile.default
+          #   platforms: linux/amd64
+          #   build_base: ubuntu:22.04
+          #   runtime_base: ubuntu:22.04
+          #   runtime_deps: libclblast1
+          #   extra_deps: libclblast-dev
+          #   cmake_args: -DGGML_CLBLAST=On
 
 
     steps:

.gitignore

@@ -0,0 +1,2 @@
+build/
+models/

.gitmodules

@@ -1,9 +1,12 @@
 [submodule "ggml"]
 	path = extern/ggml
-	url = git@github.com:ravenscroftj/ggml.git
+	url = git@github.com:ggerganov/ggml.git
 [submodule "extern/argparse"]
 	path = extern/argparse
 	url = https://github.com/p-ranav/argparse.git
 [submodule "extern/sbdlog"]
 	path = extern/spdlog
 	url = https://github.com/gabime/spdlog.git
+[submodule "extern/ggml"]
+	path = extern/ggml
+	url = https://github.com/ggerganov/ggml

.vscode/c_cpp_properties.json

@@ -0,0 +1,20 @@
+{
+    "configurations": [
+        {
+            "name": "Linux",
+            "includePath": [
+                "${workspaceFolder}/**",
+                "${workspaceFolder}/extern/crow/include",
+                "${workspaceFolder}/include",
+                "${workspaceFolder}/include"
+            ],
+            "defines": [],
+            "compilerPath": "/usr/bin/gcc",
+            "cStandard": "c17",
+            "cppStandard": "gnu++17",
+            "intelliSenseMode": "linux-gcc-x64",
+            "configurationProvider": "ms-vscode.cmake-tools"
+        }
+    ],
+    "version": 4
+}

.vscode/launch.json

@@ -0,0 +1,70 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "(gdb) Launch TBP",
+            "type": "cppdbg",
+            "request": "launch",
+            "program": "/home/james/workspace/rafael-llm/turbopilot/build/bin/turbopilot",
+            "args": [
+                //TBP ARGS
+                "-v",
+                "-f",
+                "/home/james/Downloads/replit-code-v1-3b-q4_0.bin",
+                "-m",
+                "replit",
+            ],
+            "stopAtEntry": false,
+            "cwd": "${workspaceFolder}",
+            "environment": [],
+            "externalConsole": false,
+            "MIMode": "gdb",
+            "setupCommands": [
+                {
+                    "description": "Enable pretty-printing for gdb",
+                    "text": "-enable-pretty-printing",
+                    "ignoreFailures": true
+                },
+                {
+                    "description": "Set Disassembly Flavor to Intel",
+                    "text": "-gdb-set disassembly-flavor intel",
+                    "ignoreFailures": true
+                }
+            ]
+        },
+        {
+            "name": "(gdb) Launch Replut",
+            "type": "cppdbg",
+            "request": "launch",
+            "program": "/home/james/workspace/rafael-llm/turbopilot/extern/ggml/build/bin/replit",
+            "args": [
+                // REPLIT ARGS
+                "-m",
+                "/home/james/Downloads/replit-code-v1-3b-q4_0.bin",
+                "-f",
+                "/home/james/workspace/rafael-llm/turbopilot/test.txt"
+            ],
+            "stopAtEntry": false,
+            "cwd": "${workspaceFolder}",
+            "environment": [],
+            "externalConsole": false,
+            "MIMode": "gdb",
+            "setupCommands": [
+                {
+                    "description": "Enable pretty-printing for gdb",
+                    "text": "-enable-pretty-printing",
+                    "ignoreFailures": true
+                },
+                {
+                    "description": "Set Disassembly Flavor to Intel",
+                    "text": "-gdb-set disassembly-flavor intel",
+                    "ignoreFailures": true
+                }
+            ]
+        },
+
+    ]
+}

.vscode/tasks.json

@@ -0,0 +1,28 @@
+{
+    "tasks": [
+        {
+            "type": "cppbuild",
+            "label": "C/C++: g++ build active file",
+            "command": "/usr/bin/g++",
+            "args": [
+                "-fdiagnostics-color=always",
+                "-g",
+                "${file}",
+                "-o",
+                "${fileDirname}/${fileBasenameNoExtension}"
+            ],
+            "options": {
+                "cwd": "${fileDirname}"
+            },
+            "problemMatcher": [
+                "$gcc"
+            ],
+            "group": {
+                "kind": "build",
+                "isDefault": true
+            },
+            "detail": "Task generated by Debugger."
+        }
+    ],
+    "version": "2.0.0"
+}

CMakeLists.txt

@@ -15,6 +15,11 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS "on")
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)
 set(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_PREFIX}/lib")
 
+option(GGML_CLBLAST                 "ggml: use clBLAST"                  OFF)
+option(GGML_CUBLAS                  "ggml: use cuBLAS"                   OFF)
+
+
+
 if (${CMAKE_SYSTEM_PROCESSOR} MATCHES "arm" OR ${CMAKE_SYSTEM_PROCESSOR} MATCHES "aarch64")
     message(STATUS "ARM detected")
     if (MSVC)
@@ -48,13 +53,21 @@ if (GGML_STATIC)
     SET(CMAKE_FIND_LIBRARY_SUFFIXES ".a")
     SET(BUILD_SHARED_LIBS OFF)
     SET(CMAKE_EXE_LINKER_FLAGS "-static")
+endif()
 
-    # if(GGML_OPENBLAS)
-    #     set(BLA_STATIC ON)
-    # endif()
+if (GGML_CUBLAS)
+    cmake_minimum_required(VERSION 3.17)
+
+    find_package(CUDAToolkit)
+    if (CUDAToolkit_FOUND)
+        add_compile_definitions(GGML_USE_CUBLAS)
+    else()
+        message(WARNING "cuBLAS not found")
+    endif()
 endif()
 
 
+
 add_subdirectory(src)
 
 set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}/bin)

include/turbopilot/model.hpp

@@ -44,6 +44,7 @@ struct ModelConfig
     int32_t seed = -1;     // RNG seed
     int32_t n_ctx = 512;   // context size
     int32_t n_batch = 512; // batch size for prompt processing (must be >=32 to use BLAS)
+    int32_t n_gpu_layers = 0;
 };
 
 class TurbopilotModel
@@ -67,4 +68,5 @@ class TurbopilotModel
     std::mutex model_lock;
 };
 
-#endif //__TURBOPILOT_MODEL_H
+
+#endif //__TURBOPILOT_MODEL_H

run.sh

@@ -1,3 +1,6 @@
 #!/bin/sh
-
-/app/turbopilot -t $THREADS -m $MODEL_TYPE -f $MODEL 
+if [ -z "$GPU_LAYERS" ]; then 
+    /app/turbopilot -t $THREADS -m $MODEL_TYPE -f $MODEL 
+else
+    /app/turbopilot -t $THREADS -m $MODEL_TYPE -f $MODEL --ngl $GPU_LAYERS
+fi

src/gptj.cpp

@@ -6,6 +6,14 @@
 #include <iostream>
 #include <fstream>
 
+
+#ifdef GGML_USE_CLBLAST
+#include "ggml-opencl.h"
+#endif
+#ifdef GGML_USE_CUBLAS
+#include "ggml-cuda.h"
+#endif
+
 #if defined(_MSC_VER)
 #pragma warning(disable: 4244 4267) // possible loss of data
 #endif
@@ -455,6 +463,9 @@ bool GPTJModel::load_model(std::string fname) {
         }
     }
 
+
+
+
     // key + value memory
     {
         const auto & hparams = model->hparams;
@@ -553,6 +564,47 @@ bool GPTJModel::load_model(std::string fname) {
 
     fin.close();
 
+
+
+    #if defined(GGML_USE_CLBLAST) || defined(GGML_USE_CUBLAS)
+
+    if(config.n_gpu_layers > 0){
+        size_t vram_total = 0;
+        int gpu_layers = std::min(config.n_gpu_layers, model->hparams.n_layer);
+        spdlog::info("Attempting to offload {} layers to GPU", gpu_layers);
+
+        for(int i=0; i < gpu_layers; i++) {
+            const auto & layer = model->layers[i];
+            layer.c_attn_q_proj_w->backend = GGML_BACKEND_GPU;
+            layer.c_attn_k_proj_w->backend = GGML_BACKEND_GPU;
+            layer.c_attn_v_proj_w->backend = GGML_BACKEND_GPU;
+
+            layer.c_attn_proj_w->backend = GGML_BACKEND_GPU;
+            layer.c_mlp_fc_w->backend = GGML_BACKEND_GPU;
+            layer.c_mlp_proj_w->backend = GGML_BACKEND_GPU;
+
+            #if defined(GGML_USE_CLBLAST)
+            ggml_cl_transform_tensor(layer.c_attn_q_proj_w->data,layer.c_attn_q_proj_w); vram_total += ggml_nbytes(layer.c_attn_q_proj_w);
+            ggml_cl_transform_tensor(layer.c_attn_k_proj_w->data,layer.c_attn_k_proj_w); vram_total += ggml_nbytes(layer.c_attn_k_proj_w);
+            ggml_cl_transform_tensor(layer.c_attn_v_proj_w->data,layer.c_attn_v_proj_w); vram_total += ggml_nbytes(layer.c_attn_v_proj_w);
+            ggml_cl_transform_tensor(layer.c_attn_proj_w->data,layer.c_attn_proj_w); vram_total += ggml_nbytes(layer.c_attn_proj_w);
+            ggml_cl_transform_tensor(layer.c_mlp_fc_w->data,layer.c_mlp_fc_w); vram_total += ggml_nbytes(layer.c_mlp_fc_w);
+            ggml_cl_transform_tensor(layer.c_mlp_proj_w->data,layer.c_mlp_proj_w); vram_total += ggml_nbytes(layer.c_mlp_proj_w);
+            #else
+            ggml_cuda_transform_tensor(layer.c_attn_q_proj_w->data,layer.c_attn_q_proj_w); vram_total += ggml_nbytes(layer.c_attn_q_proj_w);
+            ggml_cuda_transform_tensor(layer.c_attn_k_proj_w->data,layer.c_attn_k_proj_w); vram_total += ggml_nbytes(layer.c_attn_k_proj_w);
+            ggml_cuda_transform_tensor(layer.c_attn_v_proj_w->data,layer.c_attn_v_proj_w); vram_total += ggml_nbytes(layer.c_attn_v_proj_w);
+            ggml_cuda_transform_tensor(layer.c_attn_proj_w->data,layer.c_attn_proj_w); vram_total += ggml_nbytes(layer.c_attn_proj_w);
+            ggml_cuda_transform_tensor(layer.c_mlp_fc_w->data,layer.c_mlp_fc_w); vram_total += ggml_nbytes(layer.c_mlp_fc_w);
+            ggml_cuda_transform_tensor(layer.c_mlp_proj_w->data,layer.c_mlp_proj_w); vram_total += ggml_nbytes(layer.c_mlp_proj_w);
+            #endif
+        }
+
+        spdlog::info("{}: [GPU] total VRAM used: {} MB\n", __func__, vram_total / 1024 / 1024);
+    }
+
+    #endif // defined(GGML_USE_CLBLAST) || defined(GGML_USE_CUBLAS)
+
     return true;
 }
 

src/gptneox.cpp

@@ -3,6 +3,13 @@
 
 #include <ggml/ggml.h>
 
+#ifdef GGML_USE_CLBLAST
+#include "ggml-opencl.h"
+#endif
+#ifdef GGML_USE_CUBLAS
+#include "ggml-cuda.h"
+#endif
+
 #include <cinttypes>
 
 #include <iostream>
@@ -50,6 +57,7 @@ ggml_tensor * gpt_neox_ff(
 }
 
 
+
 // evaluate the transformer
 //
 //   - model:     the model
@@ -612,9 +620,42 @@ bool GPTNEOXModel::load_model(std::string fname) {
 
         printf("%s: model size = %8.2f MB / num tensors = %d\n", __func__, total_size/1024.0/1024.0, n_tensors);
     }
-
     fin.close();
 
+    #if defined(GGML_USE_CLBLAST) || defined(GGML_USE_CUBLAS)
+
+
+    if(config.n_gpu_layers > 0){
+        size_t vram_total = 0;
+        int gpu_layers = std::min(config.n_gpu_layers, model->hparams.n_layer);
+        spdlog::info("Attempting to offload {} layers to GPU", gpu_layers);
+
+
+        for(int i=0; i < gpu_layers; i++) {
+            const auto & layer = model->layers[i];
+            layer.c_attn_attn_w->backend = GGML_BACKEND_GPU;
+            layer.c_attn_proj_w->backend = GGML_BACKEND_GPU;
+            layer.c_mlp_fc_w->backend = GGML_BACKEND_GPU;
+            layer.c_mlp_proj_w->backend = GGML_BACKEND_GPU;
+
+            #if defined(GGML_USE_CLBLAST)
+            ggml_cl_transform_tensor(layer.c_attn_attn_w->data,layer.c_attn_attn_w); vram_total += ggml_nbytes(layer.c_attn_attn_w);
+            ggml_cl_transform_tensor(layer.c_attn_proj_w->data,layer.c_attn_proj_w); vram_total += ggml_nbytes(layer.c_attn_proj_w);
+            ggml_cl_transform_tensor(layer.c_mlp_fc_w->data,layer.c_mlp_fc_w); vram_total += ggml_nbytes(layer.c_mlp_fc_w);
+            ggml_cl_transform_tensor(layer.c_mlp_proj_w->data,layer.c_mlp_proj_w); vram_total += ggml_nbytes(layer.c_mlp_proj_w);
+            #else
+            ggml_cuda_transform_tensor(layer.c_attn_attn_w->data,layer.c_attn_attn_w); vram_total += ggml_nbytes(layer.c_attn_attn_w);
+            ggml_cuda_transform_tensor(layer.c_attn_proj_w->data,layer.c_attn_proj_w); vram_total += ggml_nbytes(layer.c_attn_proj_w);
+            ggml_cuda_transform_tensor(layer.c_mlp_fc_w->data,layer.c_mlp_fc_w); vram_total += ggml_nbytes(layer.c_mlp_fc_w);
+            ggml_cuda_transform_tensor(layer.c_mlp_proj_w->data,layer.c_mlp_proj_w); vram_total += ggml_nbytes(layer.c_mlp_proj_w);
+            #endif
+        }
+
+        spdlog::info("{}: [GPU] total VRAM used: {} MB\n", __func__, vram_total / 1024 / 1024);
+    }
+
+    #endif // defined(GGML_USE_CLBLAST) || defined(GGML_USE_CUBLAS)
+
     return true;
 }