From cc407f283ad9ef5e0a7a9e42c35bba3fb2cd81ae Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Wed, 19 Apr 2023 20:18:55 +0800
Subject: [PATCH] messing around with memory allocation to bandaid the random
 ooms with various gpt2 and gptj models

---
 otherarch/gpt2_v2.cpp | 26 +++++++++++++++-----------
 otherarch/gptj_v2.cpp | 20 ++++++++++++--------
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/otherarch/gpt2_v2.cpp b/otherarch/gpt2_v2.cpp
index 2621219d3..8accd8fb7 100644
--- a/otherarch/gpt2_v2.cpp
+++ b/otherarch/gpt2_v2.cpp
@@ -138,8 +138,8 @@ ModelLoadResult gpt2_model_load(const std::string & fname, gpt2_model & model, g
         ctx_size += n_layer*(4*n_embd*n_embd*ggml_type_sizef(wtype));         // c_mlp_proj_w
         ctx_size += n_layer*(         n_embd*ggml_type_sizef(GGML_TYPE_F32)); // c_mlp_proj_b
 
-        ctx_size += n_ctx*n_layer*n_embd*ggml_type_sizef(GGML_TYPE_F32); // memory_k
-        ctx_size += n_ctx*n_layer*n_embd*ggml_type_sizef(GGML_TYPE_F32); // memory_v
+        ctx_size += 1.5*(n_ctx*n_layer*n_embd*ggml_type_sizef(GGML_TYPE_F32)); // memory_k
+        ctx_size += 1.5*(n_ctx*n_layer*n_embd*ggml_type_sizef(GGML_TYPE_F32)); // memory_v
 
         ctx_size += (6 + 12*n_layer)*256; // object overhead
 
@@ -241,8 +241,8 @@ ModelLoadResult gpt2_model_load(const std::string & fname, gpt2_model & model, g
         const int n_mem      = n_layer*n_ctx;
         const int n_elements = n_embd*n_mem;
        
-        model.memory_k = ggml_new_tensor_1d(ctx, memory_type, n_elements);
-        model.memory_v = ggml_new_tensor_1d(ctx, memory_type, n_elements);
+        model.memory_k = ggml_new_tensor_1d(ctx, memory_type, n_elements*1.5);
+        model.memory_v = ggml_new_tensor_1d(ctx, memory_type, n_elements*1.5);
 
         const size_t memory_size = ggml_nbytes(model.memory_k) + ggml_nbytes(model.memory_v);
 
@@ -371,19 +371,23 @@ bool gpt2_eval(
     const int n_vocab = hparams.n_vocab;
 
     //todo: there is a bug that causes the buffer to oom and I cannot figure it out, hack to increase size for now  
-    static size_t buf_size = 1280u*1024*1024;
+    static size_t buf_size = 1600u*1024*1024;
     static void * buf = malloc(buf_size);
 
-    if (mem_per_token > 0 && mem_per_token*N > buf_size) {
+    if (mem_per_token > 0 && mem_per_token*N*1.6 > buf_size) {
         const size_t buf_size_new = 2*(mem_per_token*N); // add 10% to account for ggml object overhead
         //printf("\n%s: reallocating buffer from %zu to %zu bytes\n", __func__, buf_size, buf_size_new);
 
         // reallocate
-        buf_size = buf_size_new;
-        buf = realloc(buf, buf_size);
-        if (buf == nullptr) {
-            fprintf(stderr, "%s: failed to allocate %zu bytes\n", __func__, buf_size);
-            return false;
+        if (buf_size_new > buf_size)
+        {
+            buf_size = buf_size_new;
+            buf = realloc(buf, buf_size);
+            if (buf == nullptr)
+            {
+                fprintf(stderr, "%s: failed to allocate %zu bytes\n", __func__, buf_size);
+                return false;
+            }
         }
     }
 
diff --git a/otherarch/gptj_v2.cpp b/otherarch/gptj_v2.cpp
index 741fe6d3f..7aa1b65ab 100644
--- a/otherarch/gptj_v2.cpp
+++ b/otherarch/gptj_v2.cpp
@@ -382,19 +382,23 @@ bool gptj_eval(
     const int d_key = n_embd/n_head;
 
     //todo: there is a bug that causes the buffer to oom and I cannot figure it out, hack to increase size for now  
-    static size_t buf_size = 1280u*1024*1024;
+    static size_t buf_size = 1600u*1024*1024;
     static void * buf = malloc(buf_size);
 
-    if (mem_per_token > 0 && mem_per_token*N > buf_size) {
-        const size_t buf_size_new = 1.5*(mem_per_token*N); // add 10% to account for ggml object overhead
+    if (mem_per_token > 0 && mem_per_token*N*1.4 > buf_size) {
+        const size_t buf_size_new = 1.6*(mem_per_token*N); // add 10% to account for ggml object overhead
         //printf("\n%s: reallocating buffer from %zu to %zu bytes\n", __func__, buf_size, buf_size_new);
 
         // reallocate
-        buf_size = buf_size_new;
-        buf = realloc(buf, buf_size);
-        if (buf == nullptr) {
-            fprintf(stderr, "%s: failed to allocate %zu bytes\n", __func__, buf_size);
-            return false;
+        if (buf_size_new > buf_size)
+        {
+            buf_size = buf_size_new;
+            buf = realloc(buf, buf_size);
+            if (buf == nullptr)
+            {
+                fprintf(stderr, "%s: failed to allocate %zu bytes\n", __func__, buf_size);
+                return false;
+            }
         }
     }