char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Georgi Gerganovmetal : single-source fusion table + fusion debug rework (#28164)a2878d30d

char/aliceai-t5-moe
24.2 KiB693 linesraw
1#include "arg.h"
2#include "common.h"
3#include "log.h"
4#include "llama-cpp.h"
5
6#include <algorithm>
7#include <clocale>
8#include <cstring>
9#include <filesystem>
10#include <random>
11#include <string>
12#include <vector>
13
14struct llama_batch_ptr {
15    llama_batch batch;
16
17    llama_batch_ptr(int32_t n_tokens, int32_t embd, int32_t n_seq_max)
18        : batch{llama_batch_init(n_tokens, embd, n_seq_max)} {}
19
20    ~llama_batch_ptr() { llama_batch_free(batch); }
21
22    llama_batch_ptr(const llama_batch_ptr &) = delete;
23    llama_batch_ptr & operator=(const llama_batch_ptr &) = delete;
24    llama_batch_ptr(llama_batch_ptr &&) = default;
25    llama_batch_ptr & operator=(llama_batch_ptr &&) = default;
26
27    llama_batch & get() { return batch; }
28    const llama_batch & get() const { return batch; }
29};
30
31static llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) {
32    llama_tokens result;
33    llama_batch_ptr batch(1, 0, 1);
34
35    for (int i = 0; i < n_predict; i++) {
36        auto next_token = llama_sampler_sample(smpl, ctx, -1);
37
38        LOG("%d ", next_token);
39        result.push_back(next_token);
40
41        common_batch_clear(batch.get());
42        common_batch_add(batch.get(), next_token, n_past, {seq_id}, true);
43
44        if (llama_decode(ctx, batch.get())) {
45            LOG_ERR("\n%s: failed to evaluate\n", __func__);
46            return {};
47        }
48        n_past++;
49    }
50
51    return result;
52}
53
54// Test 1: baseline
55// - decode all but the last token
56// - save state to disk
57// - decode the last token
58// - generate n_predict tokens
59static llama_tokens test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
60    auto params_ctx = common_context_params_to_llama(params);
61    params_ctx.n_seq_max = 2;
62    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
63
64    auto sparams = llama_sampler_chain_default_params();
65    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
66    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
67
68    auto n_past = 0;
69    if (!common_prompt_batch_decode(ctx.get(), tokens, (int)tokens.size(), n_past, params.n_batch, params.out_file, true)) {
70        LOG_ERR("%s: failed to decode prompt\n", __func__);
71        return {};
72    }
73
74    LOG("\n=== Test 1: baseline ===\n");
75
76    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
77    if (result.empty()) {
78        return {};
79    }
80
81    LOG("\n");
82
83    return result;
84}
85
86
87// Test 2: sequence removal isolation
88// - decode the same prefix into two sequences
89// - remove sequence 0
90// - verify that sequence 1 remains unchanged
91static bool test_seq_rm_isolated(
92        struct llama_model         * model,
93        const struct common_params & params,
94        const llama_tokens         & tokens) {
95    auto params_ctx = common_context_params_to_llama(params);
96    params_ctx.n_ctx      = 256;
97    params_ctx.n_seq_max  = 2;
98    params_ctx.kv_unified = true;
99
100    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
101    if (!ctx) {
102        LOG_ERR("%s: failed to create context\n", __func__);
103        return false;
104    }
105
106    LOG("\n=== Test 2: sequence removal isolation ===\n");
107
108    const size_t n_tokens = tokens.size() < 128 ? tokens.size() : 128;
109    for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) {
110        llama_batch_ptr batch(n_tokens, 0, 1);
111        for (size_t i = 0; i < n_tokens; ++i) {
112            common_batch_add(batch.get(), tokens[i], i, { seq_id }, i == n_tokens - 1);
113        }
114
115        if (llama_decode(ctx.get(), batch.get())) {
116            LOG_ERR("%s: failed to decode prompt for sequence %d\n", __func__, seq_id);
117            return false;
118        }
119    }
120
121    const auto get_seq_state = [&](llama_seq_id seq_id, std::vector<uint8_t> & state) {
122        const size_t state_size = llama_state_seq_get_size(ctx.get(), seq_id);
123        if (state_size == 0) {
124            LOG_ERR("%s: sequence state is empty\n", __func__);
125            return false;
126        }
127
128        state.resize(state_size);
129        const size_t ncopy = llama_state_seq_get_data(ctx.get(), state.data(), state.size(), seq_id);
130        if (ncopy != state.size()) {
131            LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
132                    __func__, ncopy, state.size());
133            return false;
134        }
135
136        return true;
137    };
138
139    std::vector<uint8_t> state_before;
140    if (!get_seq_state(1, state_before)) {
141        return false;
142    }
143
144    if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
145        LOG_ERR("%s: failed to remove sequence 0\n", __func__);
146        return false;
147    }
148
149    std::vector<uint8_t> state_after;
150    if (!get_seq_state(1, state_after)) {
151        return false;
152    }
153
154    if (state_before != state_after) {
155        LOG_ERR("%s: removing sequence 0 changed sequence 1\n", __func__);
156        return false;
157    }
158
159    LOG("PASS\n");
160    return true;
161}
162
163
164// Test 3: state load
165// - create a new context
166// - load state from file
167// - replay the last prompt token
168// - generate n_predict tokens and compare against expected result
169static bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) {
170    auto params_ctx = common_context_params_to_llama(params);
171    params_ctx.n_seq_max = 2;
172    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
173
174    auto sparams = llama_sampler_chain_default_params();
175    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
176    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
177
178    LOG("\n=== Test 3: state load ===\n");
179
180    // Load state from file
181    llama_tokens unused_sts(tokens.size());
182    size_t n_token_count_out = 0;
183
184    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
185        LOG_ERR("\n%s: failed to load state\n", __func__);
186        return false;
187    }
188
189    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
190
191    // Replay last token
192    int n_past = (int) n_token_count_out - 1;
193    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
194        return false;
195    }
196    n_past++;
197
198    // Generate tokens
199    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
200    if (result.empty()) {
201        return false;
202    }
203
204    if (result != expected_result) {
205        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
206        return false;
207    }
208
209    LOG("\nPASS\n");
210    return true;
211}
212
213
214// Test 4: seq copy (host)
215// - create a multi-seq context
216// - load state from file
217// - replay the last prompt token
218// - migrate KV cache from seq 0 to seq 1 via the CPU path
219// - generate n_predict tokens on seq 1 and compare against expected result
220static bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) {
221    auto params_ctx = common_context_params_to_llama(params);
222    params_ctx.n_seq_max = 2;
223    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
224
225    auto sparams = llama_sampler_chain_default_params();
226    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
227    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
228
229    LOG("\n=== Test 4: seq copy (host) ===\n");
230
231    // Load state from file
232    llama_tokens unused_sts(tokens.size());
233    size_t n_token_count_out = 0;
234
235    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
236        LOG_ERR("\n%s: failed to load state\n", __func__);
237        return false;
238    }
239
240    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
241
242    // Replay last token
243    int n_past = (int) n_token_count_out - 1;
244    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
245        return false;
246    }
247    n_past++;
248
249    // Migrate KV cache from seq 0 to seq 1 (CPU path)
250    {
251        std::vector<uint8_t> seq_store(llama_state_seq_get_size(ctx.get(), 0));
252        const size_t ncopy = llama_state_seq_get_data(ctx.get(), seq_store.data(), seq_store.size(), 0);
253        if (ncopy != seq_store.size()) {
254            LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
255            return false;
256        }
257        LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
258
259        llama_memory_clear(llama_get_memory(ctx.get()), true);
260        LOG_TRC("%s: kv cache cleared\n", __func__);
261
262        const size_t nset = llama_state_seq_set_data(ctx.get(), seq_store.data(), seq_store.size(), 1);
263        if (nset != seq_store.size()) {
264            LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
265            return false;
266        }
267        LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
268    }
269
270    // Generate tokens on seq 1
271    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1);
272    if (result.empty()) {
273        return false;
274    }
275
276    if (result != expected_result) {
277        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
278        return false;
279    }
280
281    LOG("\nPASS\n");
282    return true;
283}
284
285
286// Test 5: seq copy (device)
287// - create a multi-seq context
288// - load state from file
289// - replay the last prompt token
290// - migrate KV cache from seq 0 to seq 1 via the on-device path
291// - generate n_predict tokens on seq 1 and compare against expected result
292static bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & expected_result) {
293    auto params_ctx = common_context_params_to_llama(params);
294    params_ctx.n_seq_max = 2;
295    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
296
297    auto sparams = llama_sampler_chain_default_params();
298    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
299    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
300
301    LOG("\n=== Test 5: seq copy (device) ===\n");
302
303    // Load state from file
304    llama_tokens unused_sts(tokens.size());
305    size_t n_token_count_out = 0;
306
307    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
308        LOG_ERR("\n%s: failed to load state\n", __func__);
309        return false;
310    }
311
312    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
313
314    // Replay last token
315    int n_past = (int) n_token_count_out - 1;
316    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
317        return false;
318    }
319    n_past++;
320
321    // Migrate KV cache from seq 0 to seq 1 (on-device path)
322    {
323        std::vector<uint8_t> seq_store(llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE));
324        const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
325        if (ncopy != seq_store.size()) {
326            LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
327            return false;
328        }
329        LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
330
331        llama_memory_clear(llama_get_memory(ctx.get()), true);
332        LOG_TRC("%s: kv cache cleared\n", __func__);
333
334        const size_t nset = llama_state_seq_set_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 1, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
335        if (nset != seq_store.size()) {
336            LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
337            return false;
338        }
339        LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
340    }
341
342    // Generate tokens on seq 1
343    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1);
344    if (result.empty()) {
345        return false;
346    }
347
348    if (result != expected_result) {
349        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
350        return false;
351    }
352
353    LOG("\nPASS\n");
354    return true;
355}
356
357
358// Test 6/7: seq copy (scatter)
359// - decode the same prefix on two sequences, interleaving seq 0 cells between the seq 1 cells
360// - save the seq 1 state, free the interleaved seq 0 cells, and restore via the given io path
361// - the restore destination is non-contiguous: scatter reads are batched per contiguous run
362// - save again on the host and compare the two blobs byte for byte
363static bool test_seq_cp_scatter(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, int test_num, bool on_device) {
364    auto params_ctx = common_context_params_to_llama(params);
365    params_ctx.n_ctx      = 256;
366    params_ctx.n_seq_max  = 2;
367    params_ctx.kv_unified = true;
368    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
369
370    LOG("\n=== Test %d: seq copy (%s, scatter) ===\n", test_num, on_device ? "device" : "host");
371
372    const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE;
373
374    auto decode_one = [&](llama_token tok, int pos, llama_seq_id seq) {
375        llama_batch_ptr batch(1, 0, 1);
376        common_batch_add(batch.get(), tok, pos, { seq }, true);
377        return llama_decode(ctx.get(), batch.get()) == 0;
378    };
379
380    // seq 0 cells 0,1,4 interleave the seq 1 cells 2,3,5
381    if (!decode_one(tokens[0], 0, 0) ||
382        !decode_one(tokens[1], 1, 0) ||
383        !decode_one(tokens[0], 0, 1) ||
384        !decode_one(tokens[1], 1, 1) ||
385        !decode_one(tokens[2], 2, 0) ||
386        !decode_one(tokens[2], 2, 1)) {
387        LOG_ERR("%s: failed to build interleaved state\n", __func__);
388        return false;
389    }
390
391    const auto get_seq_state = [&](llama_seq_id seq_id, uint32_t fl, std::vector<uint8_t> & state) {
392        const size_t state_size = llama_state_seq_get_size_ext(ctx.get(), seq_id, fl);
393        if (state_size == 0) {
394            LOG_ERR("%s: sequence state is empty\n", __func__);
395            return false;
396        }
397
398        state.resize(state_size);
399        const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), state.data(), state.size(), seq_id, fl);
400        if (ncopy != state.size()) {
401            LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n",
402                    __func__, ncopy, state.size());
403            return false;
404        }
405
406        return true;
407    };
408
409    // host blob: contains the KV data, used for the byte-for-byte comparison
410    std::vector<uint8_t> state_before;
411    if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_before)) {
412        return false;
413    }
414
415    // save via the io path under test
416    std::vector<uint8_t> state_save;
417    if (!get_seq_state(1, flags, state_save)) {
418        return false;
419    }
420    LOG_TRC("%s: seq 1 saved via %s, %zu bytes\n", __func__, on_device ? "device" : "host", state_save.size());
421
422    // free seq 0's cells so the ring is fragmented: the restore destination (seq 1's interleaved cells) stays non-contiguous
423    if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
424        LOG_ERR("%s: failed to remove sequence 0\n", __func__);
425        return false;
426    }
427
428    // restore via the io path under test
429    const size_t nset = llama_state_seq_set_data_ext(ctx.get(), state_save.data(), state_save.size(), 1, flags);
430    if (nset != state_save.size()) {
431        LOG_ERR("%s: seq set data length %zu does not match expected length %zu\n", __func__, nset, state_save.size());
432        return false;
433    }
434    LOG_TRC("%s: seq 1 restored via %s, %zu bytes\n", __func__, on_device ? "device" : "host", nset);
435
436    std::vector<uint8_t> state_after;
437    if (!get_seq_state(1, LLAMA_STATE_SEQ_FLAGS_NONE, state_after)) {
438        return false;
439    }
440
441    // the blob is serialized in sequence cell order, so identical bytes iff the restore wrote the same KV
442    if (state_before.size() != state_after.size() || memcmp(state_before.data(), state_after.data(), state_before.size()) != 0) {
443        LOG_ERR("\n%s: error: restored KV state is not byte-identical to the saved state\n", __func__);
444        return false;
445    }
446
447    LOG("\nPASS\n");
448    return true;
449}
450
451
452// Test 8: state blob round-trip
453// compares blobs rather than generated text: a partially restored cell still decodes to plausible tokens
454static bool test_state_roundtrip(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) {
455    auto params_ctx = common_context_params_to_llama(params);
456    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
457
458    LOG("\n=== Test 8: state blob round-trip ===\n");
459
460    if (llama_decode(ctx.get(), llama_batch_get_one(const_cast<llama_token *>(tokens.data()), (int32_t) tokens.size()))) {
461        LOG_ERR("\n%s: failed to decode prompt\n", __func__);
462        return false;
463    }
464
465    std::vector<uint8_t> blob_a(llama_state_seq_get_size(ctx.get(), 0));
466    const size_t n_a = llama_state_seq_get_data(ctx.get(), blob_a.data(), blob_a.size(), 0);
467    if (n_a != blob_a.size()) {
468        LOG_ERR("\n%s: saved %zu bytes, expected %zu\n", __func__, n_a, blob_a.size());
469        return false;
470    }
471
472    if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) {
473        LOG_ERR("\n%s: failed to erase seq 0\n", __func__);
474        return false;
475    }
476
477    if (llama_state_seq_set_data(ctx.get(), blob_a.data(), blob_a.size(), 0) != blob_a.size()) {
478        LOG_ERR("\n%s: failed to restore seq 0\n", __func__);
479        return false;
480    }
481
482    std::vector<uint8_t> blob_b(llama_state_seq_get_size(ctx.get(), 0));
483    const size_t n_b = llama_state_seq_get_data(ctx.get(), blob_b.data(), blob_b.size(), 0);
484    if (n_b != n_a) {
485        LOG_ERR("\n%s: re-saved %zu bytes, expected %zu\n", __func__, n_b, n_a);
486        return false;
487    }
488
489    size_t n_diff = 0;
490    size_t i_diff = 0;
491    for (size_t i = 0; i < n_a; i++) {
492        if (blob_a[i] != blob_b[i]) {
493            if (n_diff == 0) {
494                i_diff = i;
495            }
496            n_diff++;
497        }
498    }
499
500    if (n_diff > 0) {
501        LOG_ERR("\n%s: state changed across a restore: %zu of %zu bytes differ, first at offset %zu\n",
502                __func__, n_diff, n_a, i_diff);
503        return false;
504    }
505
506    LOG("\nPASS\n");
507    return true;
508}
509
510
511// Run the full save/load test suite (tests 1-8) for a single model.
512// Returns true if all tests pass, false otherwise.
513static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) {
514    struct common_params params = base_params;
515    params.model.path = model_path;
516
517    auto llama_init = common_init_from_params(params, true);
518    auto * model = llama_init->model();
519
520    if (model == nullptr) {
521        LOG_ERR("%s: failed to init model '%s'\n", __func__, model_path.c_str());
522        return false;
523    }
524
525    GGML_ASSERT(llama_init->context() == nullptr);
526
527    // Tokenize prompt or generate random tokens
528    llama_tokens tokens;
529    if (params.prompt.empty()) {
530        const int n_prompt = params.n_batch;
531
532        // this path is useful for model files that do not have a tokenizer
533        LOG_INF("%s: no prompt provided, generating %d (n_batch) random tokens\n", __func__, n_prompt);
534
535        const auto * vocab = llama_model_get_vocab(model);
536        const auto n_vocab = llama_vocab_n_tokens(vocab);
537
538        std::mt19937 rng(params.sampling.seed);
539        std::uniform_int_distribution<llama_token> dist(0, n_vocab - 1);
540        for (int i = 0; i < n_prompt; i++) {
541            tokens.push_back(dist(rng));
542        }
543    } else {
544        LOG_INF("%s: tokenizing prompt '%s'\n", __func__, params.prompt.c_str());
545
546        auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
547        tokens = common_tokenize(ctx.get(), params.prompt, true);
548    }
549
550    LOG_INF("%s: the input prompt is %d tokens\n", __func__, (int)tokens.size());
551
552    // Test 1: baseline (saves state to disk)
553    auto result_baseline = test_baseline(model, params, tokens);
554    if (result_baseline.empty()) {
555        return false;
556    }
557
558    // Test 2: sequence removal isolation
559    if (!test_seq_rm_isolated(model, params, tokens)) {
560        return false;
561    }
562
563    // Test 3: state load
564    if (!test_state_load(model, params, tokens, result_baseline)) {
565        return false;
566    }
567
568    // Test 4: seq copy (host)
569    if (!test_seq_cp_host(model, params, tokens, result_baseline)) {
570        return false;
571    }
572
573    // Test 5: seq copy (device)
574    if (!test_seq_cp_device(model, params, tokens, result_baseline)) {
575        return false;
576    }
577
578    // Test 6: seq copy (host, scatter)
579    if (!test_seq_cp_scatter(model, params, tokens, 6, false)) {
580        return false;
581    }
582
583    // Test 7: seq copy (device, scatter)
584    if (!test_seq_cp_scatter(model, params, tokens, 7, true)) {
585        return false;
586    }
587
588    // Test 8: state blob round-trip
589    if (!test_state_roundtrip(model, params, tokens)) {
590        return false;
591    }
592
593    LOG("\nAll tests passed.\n");
594
595    return true;
596}
597
598
599int main(int argc, char ** argv) {
600    std::setlocale(LC_NUMERIC, "C");
601
602    common_params params;
603    params.prompt = "";
604    params.n_batch = 100;
605    params.out_file = "dump_state.bin";
606    params.sampling.seed = 1234;
607
608    common_init();
609
610    // extract our own --models DIR option before handing the rest to the common arg parser
611    std::string models_dir;
612    std::vector<char *> filtered_argv;
613    filtered_argv.push_back(argv[0]);
614    for (int i = 1; i < argc; i++) {
615        if (strcmp(argv[i], "--models") == 0) {
616            if (i + 1 >= argc) {
617                LOG_ERR("%s: --models requires a directory argument\n", __func__);
618                return 1;
619            }
620            models_dir = argv[i + 1];
621            i++;
622        } else {
623            filtered_argv.push_back(argv[i]);
624        }
625    }
626    filtered_argv.push_back(nullptr);
627    const int fargc = (int)filtered_argv.size() - 1;
628
629    // in --models mode there is no single model; set a placeholder so the common parser's
630    // "--model is required" check passes (each model is set individually inside the loop)
631    if (!models_dir.empty()) {
632        params.model.path = models_dir;
633    }
634
635    if (!common_params_parse(fargc, filtered_argv.data(), params, LLAMA_EXAMPLE_COMMON)) {
636        return 1;
637    }
638
639    if (params.n_parallel == 1) {
640        LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__);
641        params.kv_unified = true;
642    }
643
644    if (params.n_predict < 0) {
645        params.n_predict = 16;
646    }
647
648    ggml_backend_load_all();
649
650    if (!models_dir.empty()) {
651        // run the suite over every dummy model in the directory
652        if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
653            LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
654            return 1;
655        }
656
657        std::vector<std::string> models;
658        for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
659            if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
660                models.push_back(entry.path().string());
661            }
662        }
663        std::sort(models.begin(), models.end());
664
665        if (models.empty()) {
666            LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
667            return 1;
668        }
669
670        LOG_INF("%s: running save/load tests over %zu models in '%s'\n", __func__, models.size(), models_dir.c_str());
671
672        size_t n_pass = 0;
673        size_t n_fail = 0;
674        for (const auto & model_path : models) {
675            LOG("\n================================================================\n");
676            LOG_INF("%s: model %s\n", __func__, model_path.c_str());
677
678            if (run_save_load_tests_for_model(model_path, params)) {
679                n_pass++;
680            } else {
681                n_fail++;
682            }
683        }
684
685        LOG("\n================================================================\n");
686        LOG_INF("%s: summary: %zu passed, %zu failed (of %zu)\n", __func__, n_pass, n_fail, models.size());
687
688        return n_fail == 0 ? 0 : 1;
689    }
690
691    // single-model mode
692    return run_save_load_tests_for_model(params.model.path, params) ? 0 : 1;
693}