refactor, add comments

JohannesGaessler · JohannesGaessler · commit 69a16ef97366 · 2024-03-17T20:16:54.000+01:00
diff --git a/Makefile b/Makefile
@@ -659,7 +659,7 @@ libllama.a: llama.o ggml.o $(OBJS) $(COMMON_DEPS)
 	ar rcs libllama.a llama.o ggml.o $(OBJS) $(COMMON_DEPS)
 
 clean:
-	rm -vrf *.o tests/*.o *.so *.a *.dll benchmark-matmult common/build-info.cpp *.dot $(COV_TARGETS) $(BUILD_TARGETS) $(TEST_TARGETS)
+	rm -vrf *.o tests/*.o *.so *.a *.dll benchmark-matmult lookup-create lookup-merge lookup-stats common/build-info.cpp *.dot $(COV_TARGETS) $(BUILD_TARGETS) $(TEST_TARGETS)
 	find examples pocs -type f -name "*.o" -delete
 
 #
diff --git a/common/common.cpp b/common/common.cpp
@@ -1879,9 +1879,9 @@ void llama_ngram_cache_update(llama_ngram_cache & ngram_cache, int ngram_min, in
         const int i_start = std::max(inp_size - nnew, ngram_size);
         for (int i = i_start; i < inp_size; ++i) {
             const int ngram_start = i - ngram_size;
-            uint64_t ngram = inp[ngram_start];
+            llama_ngram ngram = inp[ngram_start];
             for (int j = ngram_start+1; j < ngram_start + ngram_size; ++j) { // FIXME
-                const uint64_t ngram_part = inp[j];
+                const llama_ngram ngram_part = inp[j];
                 ngram <<= 16;
                 ngram |= ngram_part;
             }
@@ -1904,7 +1904,7 @@ void llama_ngram_cache_update(llama_ngram_cache & ngram_cache, int ngram_min, in
                 const int64_t t_now_ms = ggml_time_ms();
                 const int64_t eta_ms   = (inp_size - i) * (t_now_ms - t_start_ms) / i;
                 const int64_t eta_min  = eta_ms / (60*1000);
-                const int64_t eta_s    = (eta_ms - eta_min) / 1000;
+                const int64_t eta_s    = (eta_ms - 60*1000*eta_min) / 1000;
 
                 fprintf(stderr, "%s: %d/%d done, ETA: %02ld:%02ld\n", __func__, i, inp_size, eta_min, eta_s);
             }
@@ -1917,75 +1917,77 @@ static llama_token get_token(const std::vector<llama_token> & inp, const std::ve
     return i < inp.size() ? inp[i] : draft[1 + i - inp.size()];
 };
 
-// If sample size or percentage in context are below these thresholds the draft is aborted early:
-constexpr int draft_min_sample_size_t1[LLAMA_NGRAM_MAX] = { 2,  2,  1,  1};
-constexpr int     draft_min_percent_t1[LLAMA_NGRAM_MAX] = {66, 50, 50, 50};
-constexpr int draft_min_sample_size_t2[LLAMA_NGRAM_MAX] = { 4,  3,  2,  2};
-constexpr int     draft_min_percent_t2[LLAMA_NGRAM_MAX] = {75, 66, 66, 66};
+// If sample size or percentage are below these thresholds the draft is aborted early:
+constexpr int    draft_min_sample_size_lax[LLAMA_NGRAM_MAX] = { 2,  2,  1,  1};
+constexpr int        draft_min_percent_lax[LLAMA_NGRAM_MAX] = {66, 50, 50, 50};
+constexpr int draft_min_sample_size_strict[LLAMA_NGRAM_MAX] = { 4,  3,  2,  2};
+constexpr int     draft_min_percent_strict[LLAMA_NGRAM_MAX] = {75, 66, 66, 66};
 
-static llama_token try_draft(llama_ngram_cache & nc_primary, const uint64_t ngram_primary) {
-    llama_ngram_cache::iterator part_primary_it = nc_primary.find(ngram_primary);
-    if (part_primary_it == nc_primary.end()) {
+// Helper function that tries to draft a token from only the static ngram cache:
+static llama_token try_draft(llama_ngram_cache & nc_static, const llama_ngram ngram_static) {
+    llama_ngram_cache::iterator part_static_it = nc_static.find(ngram_static);
+    if (part_static_it == nc_static.end()) {
         return -1;
     }
-    const llama_ngram_cache_part part_primary = part_primary_it->second;
+    const llama_ngram_cache_part part_static = part_static_it->second;
 
-    int max_count_primary  = 0;
-    int sum_count_primary  = 0;
-    llama_token max_token  = -1;
+    int max_count_static  = 0;
+    int sum_count_static  = 0;
+    llama_token max_token = -1;
 
-    for (std::pair<llama_token, int> token_count_primary : part_primary) {
-        const llama_token token = token_count_primary.first;
-        const int32_t count_primary  = token_count_primary.second;
+    for (std::pair<llama_token, int> token_count_static : part_static) {
+        const llama_token token = token_count_static.first;
+        const int32_t count_static  = token_count_static.second;
 
-        if (count_primary > max_count_primary) {
-            max_token          = token;
-            max_count_primary  = count_primary;
+        if (count_static > max_count_static) {
+            max_token        = token;
+            max_count_static = count_static;
         }
-        sum_count_primary += count_primary;
+        sum_count_static += count_static;
     }
 
-    if (sum_count_primary < draft_min_sample_size_t1[2-1]) {
+    if (sum_count_static < draft_min_sample_size_lax[LLAMA_NGRAM_STATIC-1]) {
         return -1;
     }
-    if (100*max_count_primary < draft_min_percent_t1[2-1]*sum_count_primary) {
+    if (100*max_count_static < draft_min_percent_lax[LLAMA_NGRAM_STATIC-1]*sum_count_static) {
         return -1;
     }
     return max_token;
 }
 
+// Try to draft a token from primary cache (context/dynamic), validate with static cache:
 static llama_token try_draft(
-    llama_ngram_cache & nc_primary, const std::vector<uint64_t> & ngrams_primary, llama_ngram_cache_part & part_validate,
+    llama_ngram_cache & nc_primary, const std::vector<llama_ngram> & ngrams_primary, llama_ngram_cache_part & part_static,
     const int * min_sample_size, const int * min_percent) {
 
     llama_token drafted_token = -1;
 
     for (int i = ngrams_primary.size()-1; i >= 0 && drafted_token == -1; --i) {
-        const uint64_t ngram_primary = ngrams_primary[i];
+        const llama_ngram ngram_primary = ngrams_primary[i];
 
         llama_ngram_cache::iterator part_primary_it = nc_primary.find(ngram_primary);
         if (part_primary_it == nc_primary.end()) {
             continue;
         }
         const llama_ngram_cache_part part_primary = part_primary_it->second;
 
-        int max_count_primary  = 0;
-        int max_count_validate = 0;
-        int sum_count_primary  = 0;
-        llama_token max_token  = -1;
+        int max_count_primary = 0;
+        int max_count_static  = 0;
+        int sum_count_primary = 0;
+        llama_token max_token = -1;
 
         for (std::pair<llama_token, int> token_count_primary : part_primary) {
             const llama_token token = token_count_primary.first;
 
-            llama_ngram_cache_part::iterator token_count_validate_it = part_validate.find(token);
+            llama_ngram_cache_part::iterator token_count_static_it = part_static.find(token);
 
-            const int32_t count_primary  = token_count_primary.second;
-            const int32_t count_validate = token_count_validate_it != part_validate.end() ? 100*token_count_validate_it->second : 1;
+            const int32_t count_primary = token_count_primary.second;
+            const int32_t count_static  = token_count_static_it != part_static.end() ? 100*token_count_static_it->second : 1;
 
-            if (count_primary*count_validate > max_count_primary*max_count_validate) {
-                max_token          = token;
-                max_count_primary  = count_primary;
-                max_count_validate = count_validate;
+            if (count_primary*count_static > max_count_primary*max_count_static) {
+                max_token         = token;
+                max_count_primary = count_primary;
+                max_count_static  = count_static;
             }
             sum_count_primary += count_primary;
         }
@@ -2004,49 +2006,51 @@ static llama_token try_draft(
 
 void llama_ngram_cache_draft(
     std::vector<llama_token> & inp, std::vector<llama_token> & draft, int n_draft, int ngram_min, int ngram_max,
-    llama_ngram_cache & nc_t1, llama_ngram_cache & nc_t2, llama_ngram_cache & nc_t3
+    llama_ngram_cache & nc_context, llama_ngram_cache & nc_dynamic, llama_ngram_cache & nc_static
 ) {
+    GGML_ASSERT(draft.size() == 1);
     const int inp_size = inp.size();
 
-    if (inp_size < 2) {
+    if (inp_size < LLAMA_NGRAM_STATIC) {
         return;
     }
 
     while ((int) draft.size()-1 < n_draft) {
         llama_token drafted_token = -1;
 
-        const int ngram_start_t23 = inp_size-2 + draft.size()-1;
-        uint64_t ngram_t23 = get_token(inp, draft, ngram_start_t23);
-        for (int j = ngram_start_t23+1; j < ngram_start_t23 + 2; ++j) {
-            const uint64_t token = get_token(inp, draft, j);
-            ngram_t23 <<= 16;
-            ngram_t23 |= token;
+        const int ngram_start_static = inp_size-LLAMA_NGRAM_STATIC + draft.size()-1;
+        llama_ngram ngram_static = get_token(inp, draft, ngram_start_static);
+        for (int j = ngram_start_static+1; j < ngram_start_static + LLAMA_NGRAM_STATIC; ++j) {
+            const llama_ngram token = get_token(inp, draft, j);
+            ngram_static <<= 16;
+            ngram_static |= token;
         }
-        llama_ngram_cache::iterator part_t3_it = nc_t3.find(ngram_t23);
-        llama_ngram_cache_part part_t3;
-        if (part_t3_it != nc_t3.end()) {
-            part_t3 = part_t3_it->second;
+        llama_ngram_cache::iterator part_static_it = nc_static.find(ngram_static);
+        llama_ngram_cache_part part_static;
+        if (part_static_it != nc_static.end()) {
+            part_static = part_static_it->second;
         }
 
-        std::vector<uint64_t> ngrams_t12;
-        for (int ngram_size = ngram_min; ngram_size <= ngram_max; ++ngram_size) {
-            const int ngram_start_t12 = inp_size-ngram_size + draft.size()-1;
-            uint64_t ngram_t12 = get_token(inp, draft, ngram_start_t12);
-            for (int j = ngram_start_t12+1; j < ngram_start_t12 + ngram_size; ++j) {
-                const uint64_t token = get_token(inp, draft, j);
-                ngram_t12 <<= 16;
-                ngram_t12 |= token;
-            }
-            ngrams_t12.push_back(ngram_t12);
+        // cd = context + dynamic
+        std::vector<llama_ngram> ngrams_cd;
+        for (int ngram_size_cd = ngram_min; ngram_size_cd <= ngram_max; ++ngram_size_cd) {
+            const int ngram_start_cd = inp_size-ngram_size_cd + draft.size()-1;
+            llama_ngram ngram_cd = get_token(inp, draft, ngram_start_cd);
+            for (int j = ngram_start_cd+1; j < ngram_start_cd + ngram_size_cd; ++j) {
+                const llama_ngram token = get_token(inp, draft, j);
+                ngram_cd <<= 16;
+                ngram_cd |= token;
+            }
+            ngrams_cd.push_back(ngram_cd);
         }
         if (drafted_token == -1) {
-            drafted_token = try_draft(nc_t1, ngrams_t12, part_t3, draft_min_sample_size_t1, draft_min_percent_t1);
+            drafted_token = try_draft(nc_context, ngrams_cd, part_static, draft_min_sample_size_lax, draft_min_percent_lax);
         }
         if (drafted_token == -1) {
-            drafted_token = try_draft(nc_t2, ngrams_t12, part_t3, draft_min_sample_size_t2, draft_min_percent_t2);
+            drafted_token = try_draft(nc_dynamic, ngrams_cd, part_static, draft_min_sample_size_strict, draft_min_percent_strict);
         }
         if (drafted_token == -1) {
-            drafted_token = try_draft(nc_t3, ngram_t23);
+            drafted_token = try_draft(nc_static, ngram_static);
         }
 
         if (drafted_token == -1) {
@@ -2060,14 +2064,13 @@ void llama_ngram_cache_draft(
 
 void llama_ngram_cache_save(llama_ngram_cache & ngram_cache, std::string & filename) {
     std::ofstream file_out(filename, std::ios::binary);
-    for (std::pair<uint64_t, llama_ngram_cache_part> item : ngram_cache) {
-        const uint64_t         ngram        = item.first;
+    for (std::pair<llama_ngram, llama_ngram_cache_part> item : ngram_cache) {
+        const llama_ngram      ngram        = item.first;
         llama_ngram_cache_part token_counts = item.second;
         GGML_ASSERT(!token_counts.empty());
         const int32_t ntokens = token_counts.size();
 
-
-        file_out.write(reinterpret_cast<const char *>(&ngram),   sizeof(uint64_t));
+        file_out.write(reinterpret_cast<const char *>(&ngram),   sizeof(llama_ngram));
         file_out.write(reinterpret_cast<const char *>(&ntokens), sizeof(int32_t));
         for (std::pair<llama_token, int32_t> item2 : token_counts) {
             const llama_token token = item2.first;
@@ -2086,16 +2089,16 @@ llama_ngram_cache llama_ngram_cache_load(std::string & filename) {
     }
     llama_ngram_cache ngram_cache;
 
-    uint64_t ngram;
-    int32_t ntokens;
+    llama_ngram ngram;
+    int32_t     ntokens;
     llama_token token;
-    int32_t count;
+    int32_t     count;
 
     char * ngramc   = reinterpret_cast<char*>(&ngram);
     char * ntokensc = reinterpret_cast<char*>(&ntokens);
     char * tokenc   = reinterpret_cast<char*>(&token);
     char * countc   = reinterpret_cast<char*>(&count);
-    while(hashmap_file.read(ngramc, sizeof(uint64_t))) {
+    while(hashmap_file.read(ngramc, sizeof(llama_ngram))) {
         GGML_ASSERT(hashmap_file.read(ntokensc, sizeof(int32_t)));
         llama_ngram_cache_part token_counts;
 
@@ -2113,8 +2116,8 @@ llama_ngram_cache llama_ngram_cache_load(std::string & filename) {
 }
 
 void llama_ngram_cache_merge(llama_ngram_cache & ngram_cache_target, llama_ngram_cache & ngram_cache_add) {
-    for (std::pair<uint64_t, llama_ngram_cache_part> ngram_part : ngram_cache_add) {
-        const uint64_t         ngram = ngram_part.first;
+    for (std::pair<llama_ngram, llama_ngram_cache_part> ngram_part : ngram_cache_add) {
+        const llama_ngram      ngram = ngram_part.first;
         llama_ngram_cache_part  part = ngram_part.second;
 
         llama_ngram_cache::iterator part_merged_it = ngram_cache_target.find(ngram);
diff --git a/common/common.h b/common/common.h
@@ -263,30 +263,52 @@ void dump_kv_cache_view(const llama_kv_cache_view & view, int row_size = 80);
 // Dump the KV cache view showing individual sequences in each cell (long output).
 void dump_kv_cache_view_seqs(const llama_kv_cache_view & view, int row_size = 40);
 
-#define LLAMA_NGRAM_MAX 4
+#define LLAMA_NGRAM_MIN    1
+#define LLAMA_NGRAM_MAX    4
+#define LLAMA_NGRAM_STATIC 2
 
 // Data structures to map n-grams to empirical token probabilities:
-typedef std::unordered_map<llama_token, int32_t>             llama_ngram_cache_part; // token -> number of times token has been seen
-typedef std::unordered_map<uint64_t, llama_ngram_cache_part> llama_ngram_cache;      // n-gram -> empirical distribution of following tokens
-// n-grams are encoded as 64 bit integers with each of the 4 16 bit sections representing a token id.
-// This way no custom hashing function for the n-grams is needed.
+typedef uint64_t                                                llama_ngram;            // Each of the 4 16 bit sections represents a token id.
+typedef std::unordered_map<llama_token, int32_t>                llama_ngram_cache_part; // token -> number of times token has been seen
+typedef std::unordered_map<llama_ngram, llama_ngram_cache_part> llama_ngram_cache;      // n-gram -> empirical distribution of following tokens
+
+static_assert(LLAMA_NGRAM_MAX <= sizeof(llama_ngram)/2, "A 64 bit integer can only hold information for 4 16 bit tokens.");
 
 // Update an ngram cache with tokens.
-// ncs = ngram caches: the hashmaps to modify.
-// ngram_min/ngram_max: the min/max size of the ngrams in ncs.
-// inp_data: the token sequence on which the hashmaps are based.
-// nnew: how many new tokens have been appended to inp_data since the last call to this function.
-// print_progress: whether to print progress to stderr
+// ngram_cache:         the cache to modify.
+// ngram_min/ngram_max: the min/max size of the ngrams to extract from inp_data.
+// inp_data:            the token sequence with which to update ngram_cache.
+// nnew:                how many new tokens have been appended to inp_data since the last call to this function.
+// print_progress:      whether to print progress to stderr.
 //
 // In order to get correct results inp_data can ONLY BE APPENDED TO.
 // Changes in the middle need a complete rebuild.
-void llama_ngram_cache_update(llama_ngram_cache & ngram_cache, int ngram_min, int ngram_max,
-                              std::vector<llama_token> & inp_data, int nnew, bool print_progress);
-
+void llama_ngram_cache_update(
+    llama_ngram_cache & ngram_cache, int ngram_min, int ngram_max, std::vector<llama_token> & inp_data, int nnew, bool print_progress);
+
+// Try to draft tokens from ngram caches.
+// inp:                the tokens generated so far.
+// draft:              the token sequence to draft. Expected to initially contain the previously sampled token.
+// n_draft:            maximum number of tokens to add to draft.
+// ngram_min/gram_max: the min/max size of the ngrams in nc_context and nc_dynamic.
+// nc_context:         ngram cache based on current context.
+// nc_dynamic:         ngram cache based on previous user generations.
+// nc_static:          ngram cache generated from a large text corpus, used for validation.
 void llama_ngram_cache_draft(
     std::vector<llama_token> & inp, std::vector<llama_token> & draft, int n_draft, int ngram_min, int ngram_max,
-    llama_ngram_cache & nc_t1, llama_ngram_cache & nc_t2, llama_ngram_cache & nc_t3);
+    llama_ngram_cache & nc_context, llama_ngram_cache & nc_dynamic, llama_ngram_cache & nc_static);
 
+// Save an ngram cache to a file.
+// ngram_cache: the ngram cache to save.
+// filename:    the path under which to save the ngram cache.
 void llama_ngram_cache_save(llama_ngram_cache & ngram_cache, std::string & filename);
+
+// Load an ngram cache saved with llama_ngram_cache_save.
+// filename: the path from which to load the ngram cache.
+// returns:  an ngram cache containing the information saved to filename.
 llama_ngram_cache llama_ngram_cache_load(std::string & filename);
+
+// Merge two ngram caches.
+// ngram_cache_target: the ngram cache to which to add the information from ngram_cache_add.
+// ngram_cache_add:    the ngram cache to add to ngram_cache_target.
 void llama_ngram_cache_merge(llama_ngram_cache & ngram_cache_target, llama_ngram_cache & ngram_cache_add);
diff --git a/examples/lookup/lookup-create.cpp b/examples/lookup/lookup-create.cpp
@@ -9,8 +9,6 @@
 #include <unordered_map>
 #include <vector>
 
-constexpr int ngram_size = 2;
-
 int main(int argc, char ** argv){
     gpt_params params;
 
@@ -37,8 +35,8 @@ int main(int argc, char ** argv){
 
 
     llama_ngram_cache ngram_cache;
-    llama_ngram_cache_update(ngram_cache, ngram_size, ngram_size, inp, inp.size(), true);
-    fprintf(stderr, "%s: hashing done, writing file\n", __func__);
+    llama_ngram_cache_update(ngram_cache, LLAMA_NGRAM_STATIC, LLAMA_NGRAM_STATIC, inp, inp.size(), true);
+    fprintf(stderr, "%s: hashing done, writing file to %s\n", __func__, params.lookup_cache_static.c_str());
 
     llama_ngram_cache_save(ngram_cache, params.lookup_cache_static);
 }
diff --git a/examples/lookup/lookup-merge.cpp b/examples/lookup/lookup-merge.cpp
@@ -32,6 +32,7 @@ int main(int argc, char ** argv){
         }
     }
 
+    fprintf(stderr, "lookup-merge: loading file %s\n", args[0].c_str());
     llama_ngram_cache ngram_cache_merged = llama_ngram_cache_load(args[0]);
 
     for (size_t i = 1; i < args.size()-1; ++i) {
diff --git a/examples/lookup/lookup-stats.cpp b/examples/lookup/lookup-stats.cpp
diff --git a/examples/lookup/lookup.cpp b/examples/lookup/lookup.cpp

Original file line number	Diff line number	Diff line change
`@@ -659,7 +659,7 @@ libllama.a: llama.o ggml.o $(OBJS) $(COMMON_DEPS)`
`659`	`659`	`ar rcs libllama.a llama.o ggml.o $(OBJS) $(COMMON_DEPS)`
`660`	`660`
`661`	`661`	`clean:`
`662`		`- rm -vrf .o tests/.o .so .a .dll benchmark-matmult common/build-info.cpp .dot $(COV_TARGETS) $(BUILD_TARGETS) $(TEST_TARGETS)`
	`662`	`+ rm -vrf .o tests/.o .so .a .dll benchmark-matmult lookup-create lookup-merge lookup-stats common/build-info.cpp .dot $(COV_TARGETS) $(BUILD_TARGETS) $(TEST_TARGETS)`
`663`	`663`	`find examples pocs -type f -name "*.o" -delete`
`664`	`664`
`665`	`665`	`#`
Original file line number	Diff line number	Diff line change
`@@ -32,6 +32,7 @@ int main(int argc, char ** argv){`
`32`	`32`	`}`
`33`	`33`	`}`
`34`	`34`
	`35`	`+ fprintf(stderr, "lookup-merge: loading file %s\n", args[0].c_str());`
`35`	`36`	`llama_ngram_cache ngram_cache_merged = llama_ngram_cache_load(args[0]);`
`36`	`37`
`37`	`38`	`for (size_t i = 1; i < args.size()-1; ++i) {`