common : include llama_encode() in model warm-up sequence

fairydreaming · Jun 20, 2024 · 9267a13 · 9267a13
1 parent caca54c
commit 9267a13
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/common/common.cpp b/common/common.cpp
@@ -2364,7 +2364,18 @@ std::tuple<struct llama_model *, struct llama_context *> llama_init_from_gpt_par
     if (params.warmup) {
         LOG("warming up the model with an empty run\n");
 
-        std::vector<llama_token> tmp = { llama_token_bos(model), llama_token_eos(model), };
+        std::vector<llama_token> tmp;
+        llama_token bos = llama_token_bos(model);
+        llama_token eos = llama_token_eos(model);
+        // some models (e.g. T5) don't have a BOS token
+        if (bos != -1) {
+            tmp.push_back(bos);
+        }
+        tmp.push_back(eos);
+
+        if (llama_model_has_encoder(model)) {
+            llama_encode(lctx, llama_batch_get_one(tmp.data(), tmp.size(), 0, 0));
+        }
         llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch), 0, 0));
         llama_kv_cache_clear(lctx);
         llama_synchronize(lctx);

diff --git a/examples/main/main.cpp b/examples/main/main.cpp
@@ -193,8 +193,6 @@ int main(int argc, char ** argv) {
     g_model = &model;
     g_ctx = &ctx;
 
-    params.warmup = false;
-
     // load the model and apply lora adapter, if any
     LOG("%s: load the model and apply lora adapter, if any\n", __func__);
     std::tie(model, ctx) = llama_init_from_gpt_params(params);