llama

1
#include "arg.h"
2
#include "common.h"
3
#include "log.h"
4
#include "llama.h"
5

6
#include <algorithm>
7
#include <cstdio>
8
#include <string>
9
#include <vector>
10

11
static void print_usage(int, char ** argv) {
12
    LOG("\nexample usage:\n");
13
    LOG("\n    %s -m model.gguf -p \"Hello my name is\" -n 32 -np 4\n", argv[0]);
14
    LOG("\n");
15
}
16

17
int main(int argc, char ** argv) {
18
    gpt_params params;
19

20
    params.prompt = "Hello my name is";
21
    params.n_predict = 32;
22

23
    if (!gpt_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON, print_usage)) {
24
        return 1;
25
    }
26

27
    gpt_init();
28

29
    // number of parallel batches
30
    int n_parallel = params.n_parallel;
31

32
    // total length of the sequences including the prompt
33
    int n_predict = params.n_predict;
34

35
    // init LLM
36

37
    llama_backend_init();
38
    llama_numa_init(params.numa);
39

40
    // initialize the model
41

42
    llama_model_params model_params = llama_model_params_from_gpt_params(params);
43

44
    llama_model * model = llama_load_model_from_file(params.model.c_str(), model_params);
45

46
    if (model == NULL) {
47
        LOG_ERR("%s: error: unable to load model\n" , __func__);
48
        return 1;
49
    }
50

51
    // tokenize the prompt
52

53
    std::vector<llama_token> tokens_list;
54
    tokens_list = ::llama_tokenize(model, params.prompt, true);
55

56
    const int n_kv_req = tokens_list.size() + (n_predict - tokens_list.size())*n_parallel;
57

58
    // initialize the context
59

60
    llama_context_params ctx_params = llama_context_params_from_gpt_params(params);
61

62
    ctx_params.n_ctx   = n_kv_req;
63
    ctx_params.n_batch = std::max(n_predict, n_parallel);
64

65
    llama_context * ctx = llama_new_context_with_model(model, ctx_params);
66

67
    auto sparams = llama_sampler_chain_default_params();
68

69
    llama_sampler * smpl = llama_sampler_chain_init(sparams);
70

71
    llama_sampler_chain_add(smpl, llama_sampler_init_top_k(params.sparams.top_k));
72
    llama_sampler_chain_add(smpl, llama_sampler_init_top_p(params.sparams.top_p, params.sparams.min_keep));
73
    llama_sampler_chain_add(smpl, llama_sampler_init_temp (params.sparams.temp));
74
    llama_sampler_chain_add(smpl, llama_sampler_init_dist (params.sparams.seed));
75

76
    if (ctx == NULL) {
77
        LOG_ERR("%s: error: failed to create the llama_context\n" , __func__);
78
        return 1;
79
    }
80

81
    const int n_ctx = llama_n_ctx(ctx);
82

83
    LOG_INF("\n%s: n_predict = %d, n_ctx = %d, n_batch = %u, n_parallel = %d, n_kv_req = %d\n", __func__, n_predict, n_ctx, ctx_params.n_batch, n_parallel, n_kv_req);
84

85
    // make sure the KV cache is big enough to hold all the prompt and generated tokens
86
    if (n_kv_req > n_ctx) {
87
        LOG_ERR("%s: error: n_kv_req (%d) > n_ctx, the required KV cache size is not big enough\n", __func__,  n_kv_req);
88
        LOG_ERR("%s:        either reduce n_parallel or increase n_ctx\n", __func__);
89
        return 1;
90
    }
91

92
    // print the prompt token-by-token
93

94
    LOG("\n");
95

96
    for (auto id : tokens_list) {
97
        LOG("%s", llama_token_to_piece(ctx, id).c_str());
98
    }
99

100
    // create a llama_batch
101
    // we use this object to submit token data for decoding
102
    llama_batch batch = llama_batch_init(std::max(tokens_list.size(), (size_t) n_parallel), 0, n_parallel);
103

104
    std::vector<llama_seq_id> seq_ids(n_parallel, 0);
105
    for (int32_t i = 0; i < n_parallel; ++i) {
106
        seq_ids[i] = i;
107
    }
108

109
    // evaluate the initial prompt
110
    for (size_t i = 0; i < tokens_list.size(); ++i) {
111
        llama_batch_add(batch, tokens_list[i], i, seq_ids, false);
112
    }
113
    GGML_ASSERT(batch.n_tokens == (int) tokens_list.size());
114

115
    if (llama_model_has_encoder(model)) {
116
        if (llama_encode(ctx, batch)) {
117
            LOG_ERR("%s : failed to eval\n", __func__);
118
            return 1;
119
        }
120

121
        llama_token decoder_start_token_id = llama_model_decoder_start_token(model);
122
        if (decoder_start_token_id == -1) {
123
            decoder_start_token_id = llama_token_bos(model);
124
        }
125

126
        llama_batch_clear(batch);
127
        llama_batch_add(batch, decoder_start_token_id, 0, seq_ids, false);
128
    }
129

130
    // llama_decode will output logits only for the last token of the prompt
131
    batch.logits[batch.n_tokens - 1] = true;
132

133
    if (llama_decode(ctx, batch) != 0) {
134
        LOG_ERR("%s: llama_decode() failed\n", __func__);
135
        return 1;
136
    }
137

138
    //// assign the system KV cache to all parallel sequences
139
    //// this way, the parallel sequences will "reuse" the prompt tokens without having to copy them
140
    //for (int32_t i = 1; i < n_parallel; ++i) {
141
    //    llama_kv_cache_seq_cp(ctx, 0, i, -1, -1);
142
    //}
143

144
    if (n_parallel > 1) {
145
        LOG("\n\n%s: generating %d sequences ...\n", __func__, n_parallel);
146
    }
147

148
    // main loop
149

150
    // we will store the parallel decoded sequences in this vector
151
    std::vector<std::string> streams(n_parallel);
152

153
    // remember the batch index of the last token for each parallel sequence
154
    // we need this to determine which logits to sample from
155
    std::vector<int32_t> i_batch(n_parallel, batch.n_tokens - 1);
156

157
    int n_cur    = batch.n_tokens;
158
    int n_decode = 0;
159

160
    const auto t_main_start = ggml_time_us();
161

162
    while (n_cur <= n_predict) {
163
        // prepare the next batch
164
        llama_batch_clear(batch);
165

166
        // sample the next token for each parallel sequence / stream
167
        for (int32_t i = 0; i < n_parallel; ++i) {
168
            if (i_batch[i] < 0) {
169
                // the stream has already finished
170
                continue;
171
            }
172

173
            const llama_token new_token_id = llama_sampler_sample(smpl, ctx, i_batch[i]);
174

175
            // is it an end of generation? -> mark the stream as finished
176
            if (llama_token_is_eog(model, new_token_id) || n_cur == n_predict) {
177
                i_batch[i] = -1;
178
                LOG("\n");
179
                if (n_parallel > 1) {
180
                    LOG_INF("%s: stream %d finished at n_cur = %d", __func__, i, n_cur);
181
                }
182

183
                continue;
184
            }
185

186
            // if there is only one stream, we print immediately to stdout
187
            if (n_parallel == 1) {
188
                LOG("%s", llama_token_to_piece(ctx, new_token_id).c_str());
189
            }
190

191
            streams[i] += llama_token_to_piece(ctx, new_token_id);
192

193
            i_batch[i] = batch.n_tokens;
194

195
            // push this new token for next evaluation
196
            llama_batch_add(batch, new_token_id, n_cur, { i }, true);
197

198
            n_decode += 1;
199
        }
200

201
        // all streams are finished
202
        if (batch.n_tokens == 0) {
203
            break;
204
        }
205

206
        n_cur += 1;
207

208
        // evaluate the current batch with the transformer model
209
        if (llama_decode(ctx, batch)) {
210
            LOG_ERR("%s : failed to eval, return code %d\n", __func__, 1);
211
            return 1;
212
        }
213
    }
214

215
    if (n_parallel > 1) {
216
        LOG("\n");
217

218
        for (int32_t i = 0; i < n_parallel; ++i) {
219
            LOG("sequence %d:\n\n%s%s\n\n", i, params.prompt.c_str(), streams[i].c_str());
220
        }
221
    }
222

223
    const auto t_main_end = ggml_time_us();
224

225
    LOG_INF("%s: decoded %d tokens in %.2f s, speed: %.2f t/s\n",
226
            __func__, n_decode, (t_main_end - t_main_start) / 1000000.0f, n_decode / ((t_main_end - t_main_start) / 1000000.0f));
227

228
    LOG("\n");
229
    llama_perf_sampler_print(smpl);
230
    llama_perf_context_print(ctx);
231

232
    fprintf(stderr, "\n");
233

234
    llama_batch_free(batch);
235

236
    llama_sampler_free(smpl);
237
    llama_free(ctx);
238
    llama_free_model(model);
239

240
    llama_backend_free();
241

242
    return 0;
243
}
244