Add "-e"/"--eval-threads" command-line parameter to set a different number of threads for single-token eval than for prompt eval.

This commit is contained in:
ml6 2023-04-03 11:17:07 -07:00
parent 53dbba7695
commit 37264707c2
7 changed files with 29 additions and 10 deletions

View file

@ -37,6 +37,7 @@ bool gpt_params_parse(int argc, char ** argv, gpt_params & params) {
params.n_threads = std::max(1, (int32_t) std::thread::hardware_concurrency());
}
bool ethreads_set = false;
bool invalid_param = false;
std::string arg;
gpt_params default_params;
@ -56,6 +57,13 @@ bool gpt_params_parse(int argc, char ** argv, gpt_params & params) {
break;
}
params.n_threads = std::stoi(argv[i]);
} else if (arg == "-e" || arg == "--eval-threads") {
if (++i >= argc) {
invalid_param = true;
break;
}
params.n_ethreads = std::stoi(argv[i]);
ethreads_set = true;
} else if (arg == "-p" || arg == "--prompt") {
if (++i >= argc) {
invalid_param = true;
@ -191,6 +199,12 @@ bool gpt_params_parse(int argc, char ** argv, gpt_params & params) {
exit(1);
}
}
// ensure that n_ethreads defaults to the system thread-max only when n_threads is not set
if (!ethreads_set) {
params.n_ethreads = params.n_threads;
}
if (invalid_param) {
fprintf(stderr, "error: invalid parameter for argument: %s\n", arg.c_str());
gpt_print_usage(argc, argv, default_params);

View file

@ -16,6 +16,7 @@
struct gpt_params {
int32_t seed = -1; // RNG seed
int32_t n_threads = std::min(4, (int32_t) std::thread::hardware_concurrency());
int32_t n_ethreads = std::min(4, (int32_t) std::thread::hardware_concurrency());
int32_t n_predict = 128; // new tokens to predict
int32_t repeat_last_n = 64; // last n tokens to penalize
int32_t n_parts = -1; // amount of model parts (-1 = determine from model dimensions)

View file

@ -79,7 +79,7 @@ int main(int argc, char ** argv) {
if (params.embedding){
if (embd_inp.size() > 0) {
if (llama_eval(ctx, embd_inp.data(), embd_inp.size(), n_past, params.n_threads)) {
if (llama_eval(ctx, embd_inp.data(), embd_inp.size(), n_past, params.n_threads, params.n_ethreads)) {
fprintf(stderr, "%s : failed to eval\n", __func__);
return 1;
}

View file

@ -119,12 +119,12 @@ int main(int argc, char ** argv) {
if (params.mem_test) {
{
const std::vector<llama_token> tmp(params.n_batch, 0);
llama_eval(ctx, tmp.data(), tmp.size(), 0, params.n_threads);
llama_eval(ctx, tmp.data(), tmp.size(), 0, params.n_threads, params.n_ethreads);
}
{
const std::vector<llama_token> tmp = { 0, };
llama_eval(ctx, tmp.data(), tmp.size(), params.n_predict - 1, params.n_threads);
llama_eval(ctx, tmp.data(), tmp.size(), params.n_predict - 1, params.n_threads, params.n_ethreads);
}
llama_print_timings(ctx);
@ -264,7 +264,7 @@ int main(int argc, char ** argv) {
//printf("\n---\n");
}
if (llama_eval(ctx, embd.data(), embd.size(), n_past, params.n_threads)) {
if (llama_eval(ctx, embd.data(), embd.size(), n_past, params.n_threads, params.n_ethreads)) {
fprintf(stderr, "%s : failed to eval\n", __func__);
return 1;
}

View file

@ -38,7 +38,7 @@ void perplexity(llama_context * ctx, const gpt_params & params) {
// it is better to always be power of 2 for better performance
std::vector<llama_token> embd(tokens.begin() + start, tokens.begin() + end);
auto start_t = std::chrono::high_resolution_clock::now();
if (llama_eval(ctx, embd.data(), embd.size(), 0, params.n_threads)) {
if (llama_eval(ctx, embd.data(), embd.size(), 0, params.n_threads, params.n_ethreads)) {
fprintf(stderr, "%s : failed to eval\n", __func__);
return;
}

View file

@ -745,13 +745,15 @@ static bool llama_model_load(
// - tokens: new batch of tokens to process
// - n_past: the context size so far
// - n_threads: number of threads to use
// - n_ethreads: number of threads to use for single-token eval
//
static bool llama_eval_internal(
llama_context & lctx,
const llama_token * tokens,
const int n_tokens,
const int n_past,
const int n_threads) {
const int n_threads,
const int n_ethreads) {
const int64_t t_start_us = ggml_time_us();
const int N = n_tokens;
@ -784,7 +786,7 @@ static bool llama_eval_internal(
// for big prompts, if BLAS is enabled, it is better to use only one thread
// otherwise, the threads are spin-lock waiting for the BLAS calls and are degrading the performance
ggml_cgraph gf = {};
gf.n_threads = N >= 32 && ggml_cpu_has_blas() ? 1 : n_threads;
gf.n_threads = N >= 32 && ggml_cpu_has_blas() ? 1 : (N == 1 ? n_ethreads : n_threads);
struct ggml_tensor * embd = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, N);
memcpy(embd->data, tokens, N*ggml_element_size(embd));
@ -1714,8 +1716,9 @@ int llama_eval(
const llama_token * tokens,
int n_tokens,
int n_past,
int n_threads) {
if (!llama_eval_internal(*ctx, tokens, n_tokens, n_past, n_threads)) {
int n_threads,
int n_ethreads) {
if (!llama_eval_internal(*ctx, tokens, n_tokens, n_past, n_threads, n_ethreads)) {
fprintf(stderr, "%s: failed to eval\n", __func__);
return 1;
}

View file

@ -109,7 +109,8 @@ extern "C" {
const llama_token * tokens,
int n_tokens,
int n_past,
int n_threads);
int n_threads,
int n_ethreads);
// Convert the provided text into tokens.
// The tokens pointer must be large enough to hold the resulting tokens.