char-fork/llama.cpp

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

Daniel Beveniussrc : add n_expert_used_max function (#28323)9a4843cf2

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
76.1 KiB1804 linesraw
1#include "llama-model-loader.h"
2
3#include "ggml-alloc.h"
4#include "ggml.h"
5#include "gguf.h"
6#include "llama-hparams.h"
7#include "llama.h"
8
9#include <algorithm>
10#include <array>
11#include <cinttypes>
12#include <cstdint>
13#include <cstring>
14#include <future>
15#include <regex>
16
17static const size_t kiB = 1024;
18static const size_t MiB = 1024*kiB;
19static const size_t GiB = 1024*MiB;
20
21const char * llama_file_version_name(llama_fver version) {
22    switch (version) {
23        case GGUF_FILE_VERSION_V1: return "GGUF V1 (support until nov 2023)";
24        case GGUF_FILE_VERSION_V2: return "GGUF V2";
25        case GGUF_FILE_VERSION_V3: return "GGUF V3 (latest)";
26    }
27
28    return "unknown";
29}
30
31#define LLAMA_FTYPE_PREFIX "(guessed) "
32
33const char * llama_ftype_name(llama_ftype ftype) {
34    static constexpr size_t guessed_prefix_len = sizeof(LLAMA_FTYPE_PREFIX) - 1;
35    const char * name;
36    switch ((enum llama_ftype) (ftype & ~LLAMA_FTYPE_GUESSED)) {
37        case LLAMA_FTYPE_ALL_F32:          name = LLAMA_FTYPE_PREFIX "all F32"; break;
38        case LLAMA_FTYPE_MOSTLY_F16:       name = LLAMA_FTYPE_PREFIX "F16"; break;
39        case LLAMA_FTYPE_MOSTLY_BF16:      name = LLAMA_FTYPE_PREFIX "BF16"; break;
40        case LLAMA_FTYPE_MOSTLY_Q1_0:      name = LLAMA_FTYPE_PREFIX "Q1_0"; break;
41        case LLAMA_FTYPE_MOSTLY_Q2_0:      name = LLAMA_FTYPE_PREFIX "Q2_0"; break;
42        case LLAMA_FTYPE_MOSTLY_Q4_0:      name = LLAMA_FTYPE_PREFIX "Q4_0"; break;
43        case LLAMA_FTYPE_MOSTLY_Q4_1:      name = LLAMA_FTYPE_PREFIX "Q4_1"; break;
44        case LLAMA_FTYPE_MOSTLY_Q5_0:      name = LLAMA_FTYPE_PREFIX "Q5_0"; break;
45        case LLAMA_FTYPE_MOSTLY_Q5_1:      name = LLAMA_FTYPE_PREFIX "Q5_1"; break;
46        case LLAMA_FTYPE_MOSTLY_Q8_0:      name = LLAMA_FTYPE_PREFIX "Q8_0"; break;
47        case LLAMA_FTYPE_MOSTLY_MXFP4_MOE: name = LLAMA_FTYPE_PREFIX "MXFP4 MoE"; break;
48        case LLAMA_FTYPE_MOSTLY_NVFP4:     name = LLAMA_FTYPE_PREFIX "NVFP4"; break;
49        case LLAMA_FTYPE_MOSTLY_Q2_K:      name = LLAMA_FTYPE_PREFIX "Q2_K - Medium"; break;
50        case LLAMA_FTYPE_MOSTLY_Q2_K_S:    name = LLAMA_FTYPE_PREFIX "Q2_K - Small"; break;
51        case LLAMA_FTYPE_MOSTLY_Q3_K_S:    name = LLAMA_FTYPE_PREFIX "Q3_K - Small"; break;
52        case LLAMA_FTYPE_MOSTLY_Q3_K_M:    name = LLAMA_FTYPE_PREFIX "Q3_K - Medium"; break;
53        case LLAMA_FTYPE_MOSTLY_Q3_K_L:    name = LLAMA_FTYPE_PREFIX "Q3_K - Large"; break;
54        case LLAMA_FTYPE_MOSTLY_Q4_K_S:    name = LLAMA_FTYPE_PREFIX "Q4_K - Small"; break;
55        case LLAMA_FTYPE_MOSTLY_Q4_K_M:    name = LLAMA_FTYPE_PREFIX "Q4_K - Medium"; break;
56        case LLAMA_FTYPE_MOSTLY_Q5_K_S:    name = LLAMA_FTYPE_PREFIX "Q5_K - Small"; break;
57        case LLAMA_FTYPE_MOSTLY_Q5_K_M:    name = LLAMA_FTYPE_PREFIX "Q5_K - Medium"; break;
58        case LLAMA_FTYPE_MOSTLY_Q6_K:      name = LLAMA_FTYPE_PREFIX "Q6_K"; break;
59        case LLAMA_FTYPE_MOSTLY_TQ1_0:     name = LLAMA_FTYPE_PREFIX "TQ1_0 - 1.69 bpw ternary"; break;
60        case LLAMA_FTYPE_MOSTLY_TQ2_0:     name = LLAMA_FTYPE_PREFIX "TQ2_0 - 2.06 bpw ternary"; break;
61        case LLAMA_FTYPE_MOSTLY_IQ2_XXS:   name = LLAMA_FTYPE_PREFIX "IQ2_XXS - 2.0625 bpw"; break;
62        case LLAMA_FTYPE_MOSTLY_IQ2_XS:    name = LLAMA_FTYPE_PREFIX "IQ2_XS - 2.3125 bpw"; break;
63        case LLAMA_FTYPE_MOSTLY_IQ2_S:     name = LLAMA_FTYPE_PREFIX "IQ2_S - 2.5 bpw"; break;
64        case LLAMA_FTYPE_MOSTLY_IQ2_M:     name = LLAMA_FTYPE_PREFIX "IQ2_M - 2.7 bpw"; break;
65        case LLAMA_FTYPE_MOSTLY_IQ3_XS:    name = LLAMA_FTYPE_PREFIX "IQ3_XS - 3.3 bpw"; break;
66        case LLAMA_FTYPE_MOSTLY_IQ3_XXS:   name = LLAMA_FTYPE_PREFIX "IQ3_XXS - 3.0625 bpw"; break;
67        case LLAMA_FTYPE_MOSTLY_IQ1_S:     name = LLAMA_FTYPE_PREFIX "IQ1_S - 1.5625 bpw"; break;
68        case LLAMA_FTYPE_MOSTLY_IQ1_M:     name = LLAMA_FTYPE_PREFIX "IQ1_M - 1.75 bpw"; break;
69        case LLAMA_FTYPE_MOSTLY_IQ4_NL:    name = LLAMA_FTYPE_PREFIX "IQ4_NL - 4.5 bpw"; break;
70        case LLAMA_FTYPE_MOSTLY_IQ4_XS:    name = LLAMA_FTYPE_PREFIX "IQ4_XS - 4.25 bpw"; break;
71        case LLAMA_FTYPE_MOSTLY_IQ3_S:     name = LLAMA_FTYPE_PREFIX "IQ3_S - 3.4375 bpw"; break;
72        case LLAMA_FTYPE_MOSTLY_IQ3_M:     name = LLAMA_FTYPE_PREFIX "IQ3_S mix - 3.66 bpw"; break;
73        default:                           name = LLAMA_FTYPE_PREFIX "unknown, may not work"; break;
74    }
75    return (ftype & LLAMA_FTYPE_GUESSED) ? name : name + guessed_prefix_len;
76}
77
78#undef LLAMA_FTYPE_PREFIX
79
80// return a list of splits for a given path
81// for example, given "<name>-00002-of-00004.gguf", returns list of all 4 splits
82static std::vector<std::string> llama_get_list_splits(const std::string & path, const int idx, const int n_split) {
83    std::vector<std::string> paths;
84    std::string split_prefix;
85    std::vector<char> buf(llama_path_max(), 0);
86
87    {
88        int ret = llama_split_prefix(buf.data(), buf.size(), path.c_str(), idx, n_split);
89        if (!ret) {
90            throw std::runtime_error(format("invalid split file name: %s", path.c_str()));
91        }
92        split_prefix = std::string(buf.data(), ret);
93    }
94
95    if (split_prefix.empty()) {
96        throw std::runtime_error(format("invalid split file: %s", path.c_str()));
97    }
98
99    for (int idx = 0; idx < n_split; ++idx) {
100        int ret = llama_split_path(buf.data(), buf.size(), split_prefix.c_str(), idx, n_split);
101        paths.push_back(std::string(buf.data(), ret));
102    }
103
104    return paths;
105}
106
107namespace GGUFMeta {
108    template <typename T, gguf_type gt_, T (*gfun)(const gguf_context *, const int64_t)>
109    struct GKV_Base_Type {
110        static constexpr gguf_type gt = gt_;
111
112        static T getter(const gguf_context * ctx, const int kid) {
113            return gfun(ctx, kid);
114        }
115    };
116
117    template<typename T> struct GKV_Base;
118
119    template<> struct GKV_Base<bool        >: GKV_Base_Type<bool,         GGUF_TYPE_BOOL,    gguf_get_val_bool> {};
120    template<> struct GKV_Base<uint8_t     >: GKV_Base_Type<uint8_t,      GGUF_TYPE_UINT8,   gguf_get_val_u8  > {};
121    template<> struct GKV_Base<uint16_t    >: GKV_Base_Type<uint16_t,     GGUF_TYPE_UINT16,  gguf_get_val_u16 > {};
122    template<> struct GKV_Base<uint32_t    >: GKV_Base_Type<uint32_t,     GGUF_TYPE_UINT32,  gguf_get_val_u32 > {};
123    template<> struct GKV_Base<uint64_t    >: GKV_Base_Type<uint64_t,     GGUF_TYPE_UINT64,  gguf_get_val_u64 > {};
124    template<> struct GKV_Base<int8_t      >: GKV_Base_Type<int8_t,       GGUF_TYPE_INT8,    gguf_get_val_i8  > {};
125    template<> struct GKV_Base<int16_t     >: GKV_Base_Type<int16_t,      GGUF_TYPE_INT16,   gguf_get_val_i16 > {};
126    template<> struct GKV_Base<int32_t     >: GKV_Base_Type<int32_t,      GGUF_TYPE_INT32,   gguf_get_val_i32 > {};
127    template<> struct GKV_Base<int64_t     >: GKV_Base_Type<int64_t,      GGUF_TYPE_INT64,   gguf_get_val_i64 > {};
128    template<> struct GKV_Base<float       >: GKV_Base_Type<float,        GGUF_TYPE_FLOAT32, gguf_get_val_f32 > {};
129    template<> struct GKV_Base<double      >: GKV_Base_Type<double,       GGUF_TYPE_FLOAT64, gguf_get_val_f64 > {};
130    template<> struct GKV_Base<const char *>: GKV_Base_Type<const char *, GGUF_TYPE_STRING,  gguf_get_val_str > {};
131
132    template<> struct GKV_Base<std::string> {
133        static constexpr gguf_type gt = GGUF_TYPE_STRING;
134
135        static std::string getter(const gguf_context * ctx, const int kid) {
136            return gguf_get_val_str(ctx, kid);
137        }
138    };
139
140    struct ArrayInfo {
141        const gguf_type gt;
142        const size_t length;
143        const void * data;
144    };
145
146    template<> struct GKV_Base<ArrayInfo> {
147        public:
148        static constexpr gguf_type gt = GGUF_TYPE_ARRAY;
149        static ArrayInfo getter(const gguf_context *ctx, const int k) {
150            const enum gguf_type arr_type = gguf_get_arr_type(ctx, k);
151            return ArrayInfo {
152                arr_type,
153                gguf_get_arr_n(ctx, k),
154                arr_type == GGUF_TYPE_STRING ? nullptr : gguf_get_arr_data(ctx, k),
155            };
156        }
157    };
158
159    template<typename T>
160    class GKV : public GKV_Base<T> {
161        GKV() = delete;
162
163        public:
164        static T get_kv(const gguf_context * ctx, const int k) {
165            const enum gguf_type kt = gguf_get_kv_type(ctx, k);
166
167            if (kt != GKV::gt) {
168                throw std::runtime_error(format("key %s has wrong type %s but expected type %s",
169                    gguf_get_key(ctx, k), gguf_type_name(kt), gguf_type_name(GKV::gt)));
170            }
171            return GKV::getter(ctx, k);
172        }
173
174        static const char * override_type_to_str(const llama_model_kv_override_type ty) {
175            switch (ty) {
176                case LLAMA_KV_OVERRIDE_TYPE_BOOL:  return "bool";
177                case LLAMA_KV_OVERRIDE_TYPE_INT:   return "int";
178                case LLAMA_KV_OVERRIDE_TYPE_FLOAT: return "float";
179                case LLAMA_KV_OVERRIDE_TYPE_STR:   return "str";
180            }
181            return "unknown";
182        }
183
184        static bool validate_override(const llama_model_kv_override_type expected_type, const struct llama_model_kv_override * ovrd) {
185            if (!ovrd) { return false; }
186            if (ovrd->tag == expected_type) {
187                LLAMA_LOG_INFO("%s: Using metadata override (%5s) '%s' = ",
188                    __func__, override_type_to_str(ovrd->tag), ovrd->key);
189                switch (ovrd->tag) {
190                    case LLAMA_KV_OVERRIDE_TYPE_BOOL:  {
191                        LLAMA_LOG_INFO("%s\n", ovrd->val_bool ? "true" : "false");
192                    } break;
193                    case LLAMA_KV_OVERRIDE_TYPE_INT:   {
194                        LLAMA_LOG_INFO("%" PRId64 "\n", ovrd->val_i64);
195                    } break;
196                    case LLAMA_KV_OVERRIDE_TYPE_FLOAT: {
197                        LLAMA_LOG_INFO("%.6f\n", ovrd->val_f64);
198                    } break;
199                    case LLAMA_KV_OVERRIDE_TYPE_STR: {
200                        LLAMA_LOG_INFO("%s\n", ovrd->val_str);
201                    } break;
202                    default:
203                        // Shouldn't be possible to end up here, but just in case...
204                        throw std::runtime_error(
205                            format("Unsupported attempt to override %s type for metadata key %s\n",
206                                override_type_to_str(ovrd->tag), ovrd->key));
207                }
208                return true;
209            }
210            LLAMA_LOG_WARN("%s: Warning: Bad metadata override type for key '%s', expected %s but got %s\n",
211                __func__, ovrd->key, override_type_to_str(expected_type), override_type_to_str(ovrd->tag));
212            return false;
213        }
214
215        template<typename OT>
216        static typename std::enable_if<std::is_same<OT, bool>::value, bool>::type
217        try_override(OT & target, const struct llama_model_kv_override * ovrd) {
218            if (validate_override(LLAMA_KV_OVERRIDE_TYPE_BOOL, ovrd)) {
219                target = ovrd->val_bool;
220                return true;
221            }
222            return false;
223        }
224
225        template<typename OT>
226        static typename std::enable_if<!std::is_same<OT, bool>::value && std::is_integral<OT>::value, bool>::type
227        try_override(OT & target, const struct llama_model_kv_override * ovrd) {
228            if (validate_override(LLAMA_KV_OVERRIDE_TYPE_INT, ovrd)) {
229                target = ovrd->val_i64;
230                return true;
231            }
232            return false;
233        }
234
235        template<typename OT>
236        static typename std::enable_if<std::is_floating_point<OT>::value, bool>::type
237        try_override(T & target, const struct llama_model_kv_override * ovrd) {
238            if (validate_override(LLAMA_KV_OVERRIDE_TYPE_FLOAT, ovrd)) {
239                target = ovrd->val_f64;
240                return true;
241            }
242            return false;
243        }
244
245        template<typename OT>
246        static typename std::enable_if<std::is_same<OT, std::string>::value, bool>::type
247        try_override(T & target, const struct llama_model_kv_override * ovrd) {
248            if (validate_override(LLAMA_KV_OVERRIDE_TYPE_STR, ovrd)) {
249                target = ovrd->val_str;
250                return true;
251            }
252            return false;
253        }
254
255        static bool set(const gguf_context * ctx, const int k, T & target, const struct llama_model_kv_override * ovrd = nullptr) {
256            if (try_override<T>(target, ovrd)) {
257                return true;
258            }
259            if (k < 0) { return false; }
260            target = get_kv(ctx, k);
261            return true;
262        }
263
264        static bool set(const gguf_context * ctx, const char * key, T & target, const struct llama_model_kv_override * ovrd = nullptr) {
265            return set(ctx, gguf_find_key(ctx, key), target, ovrd);
266        }
267
268        static bool set(const gguf_context * ctx, const std::string & key, T & target, const struct llama_model_kv_override * ovrd = nullptr) {
269            return set(ctx, key.c_str(), target, ovrd);
270        }
271    };
272}
273
274    template<typename T>
275    typename std::enable_if<std::is_integral<T>::value, bool>::type
276    llama_model_loader::get_arr_n(const std::string & key, T & result, bool required) {
277        const int kid = gguf_find_key(metadata, key.c_str());
278
279        if (kid < 0) {
280            if (required) {
281                throw std::runtime_error(format("key not found in model: %s", key.c_str()));
282            }
283            return false;
284        }
285
286        struct GGUFMeta::ArrayInfo arr_info =
287            GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(metadata, kid);
288
289
290        result = arr_info.length;
291        return true;
292    }
293
294    template<typename T>
295    typename std::enable_if<std::is_integral<T>::value, bool>::type
296    llama_model_loader::get_arr_n(enum llm_kv kid, T & result, bool required) {
297        return get_arr_n(llm_kv(kid), result, required);
298    }
299
300    template bool llama_model_loader::get_arr_n(enum llm_kv kid, uint32_t & result, bool required);
301    template std::enable_if<std::is_integral<uint32_t>::value, bool>::type
302    llama_model_loader::get_arr_n<uint32_t>(const std::string & key, uint32_t & result, bool required);
303
304    template<typename T>
305    bool llama_model_loader::get_arr(const std::string & key, std::vector<T> & result, bool required) {
306        const gguf_context * ctx = metadata;
307        const int kid = gguf_find_key(ctx, key.c_str());
308
309        if (kid < 0 || gguf_get_kv_type(ctx, kid) != GGUF_TYPE_ARRAY) {
310            if (required) {
311                throw std::runtime_error(format("array key not found in model: %s", key.c_str()));
312            }
313            return false;
314        }
315
316        struct GGUFMeta::ArrayInfo arr_info =
317            GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
318
319        bool type_ok = false;
320        switch (arr_info.gt) {
321            case GGUF_TYPE_UINT32:
322            case GGUF_TYPE_INT32:   type_ok = (std::is_same<T,     int32_t>::value) ||
323                                              (std::is_same<T,    uint32_t>::value); break;
324            case GGUF_TYPE_UINT64:  type_ok = (std::is_same<T,    uint64_t>::value); break;
325            case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T,       float>::value); break;
326            case GGUF_TYPE_STRING:  type_ok = (std::is_same<T, std::string>::value); break;
327            default:
328                throw std::runtime_error(format("%s is not a string/float32/uint32/int32/uint64 array", key.c_str()));
329        }
330        if (!type_ok) {
331            throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
332        }
333
334        if constexpr (std::is_same<T, std::string>::value) {
335            const size_t n_items = gguf_get_arr_n(ctx, kid);
336            result.clear();
337
338            for (size_t i = 0; i < n_items; i++) {
339                const T value = gguf_get_arr_str(ctx, kid, i);
340                result.emplace_back(value);
341            }
342        } else {
343            result.resize(arr_info.length);
344            result.assign((const T*)arr_info.data, (const T *)arr_info.data + arr_info.length);
345        }
346
347        return true;
348    }
349
350    template<typename T, size_t N_MAX>
351    bool llama_model_loader::get_arr(const std::string & key, std::array<T, N_MAX> & result, bool required) {
352        const gguf_context * ctx = metadata;
353        const int kid = gguf_find_key(ctx, key.c_str());
354
355        if (kid < 0 || gguf_get_kv_type(ctx, kid) != GGUF_TYPE_ARRAY) {
356            if (required) {
357                throw std::runtime_error(format("array key not found in model: %s", key.c_str()));
358            }
359            return false;
360        }
361
362        struct GGUFMeta::ArrayInfo arr_info =
363            GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(ctx, kid);
364
365        bool type_ok = false;
366        switch (arr_info.gt) {
367            case GGUF_TYPE_BOOL:
368            case GGUF_TYPE_UINT32:
369            case GGUF_TYPE_INT32:   type_ok = (std::is_same<T,     int32_t>::value) ||
370                                              (std::is_same<T,    uint32_t>::value); break;
371            case GGUF_TYPE_UINT64:  type_ok = (std::is_same<T,    uint64_t>::value); break;
372            case GGUF_TYPE_FLOAT32: type_ok = (std::is_same<T,       float>::value); break;
373            case GGUF_TYPE_STRING:  type_ok = (std::is_same<T, std::string>::value); break;
374            default:
375                throw std::runtime_error(format("%s is not a string/float32/uint32/int32/uint64 array", key.c_str()));
376        }
377        if (!type_ok) {
378            throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt)));
379        }
380
381        if (arr_info.length > N_MAX) {
382            throw std::runtime_error(format("array length %u for key %s exceeds max %u", (uint32_t) arr_info.length, key.c_str(), (uint32_t) N_MAX));
383        }
384
385        if constexpr (std::is_same<T, std::string>::value) {
386            const size_t n_items = gguf_get_arr_n(ctx, kid);
387
388            for (size_t i = 0; i < n_items; i++) {
389                const T value = gguf_get_arr_str(ctx, kid, i);
390                result[i] = value;
391            }
392        } else {
393            if (arr_info.gt == GGUF_TYPE_BOOL) {
394                const int8_t * values = (const int8_t *) arr_info.data;
395                std::transform(values, values + arr_info.length, result.begin(), [](int8_t x) {
396                    return static_cast<T>(x != 0);
397                });
398            } else {
399                std::copy((const T*)arr_info.data, (const T *)arr_info.data + arr_info.length, result.begin());
400            }
401        }
402
403        return true;
404    }
405
406    template<typename T>
407    bool llama_model_loader::get_arr(enum llm_kv kid, T & result, bool required) {
408        return get_arr(llm_kv(kid), result, required);
409    }
410
411    template bool llama_model_loader::get_arr<std::vector<std::string>>(enum llm_kv kid, std::vector<std::string> & result, bool required);
412    template bool llama_model_loader::get_arr<std::array<int32_t, 512>>(enum llm_kv kid, std::array<int32_t, 512> & result, bool required);
413    template bool llama_model_loader::get_arr<std::vector<int32_t>>(enum llm_kv kid, std::vector<int32_t> & result, bool required);
414    template bool llama_model_loader::get_arr<std::array<uint32_t, LLAMA_MAX_LAYERS>>(enum llm_kv kid, std::array<uint32_t, LLAMA_MAX_LAYERS> & result, bool required);
415    template bool llama_model_loader::get_arr<std::vector<uint32_t>>(enum llm_kv kid, std::vector<uint32_t> & result, bool required);
416    template bool llama_model_loader::get_arr<std::array<uint64_t, LLAMA_MAX_PLE_NGRAM>>(enum llm_kv kid, std::array<uint64_t, LLAMA_MAX_PLE_NGRAM> & result, bool required);
417    template bool llama_model_loader::get_arr<std::array<uint64_t, LLAMA_MAX_PLE_HEADS>>(enum llm_kv kid, std::array<uint64_t, LLAMA_MAX_PLE_HEADS> & result, bool required);
418
419    template<typename T>
420    bool llama_model_loader::get_key(const std::string & key, T & result, bool required) {
421        auto it = kv_overrides.find(key);
422
423        const struct llama_model_kv_override * override =
424            it != kv_overrides.end() ? &it->second : nullptr;
425
426        const bool found = GGUFMeta::GKV<T>::set(metadata, key, result, override);
427
428        if (required && !found) {
429            throw std::runtime_error(format("key not found in model: %s", key.c_str()));
430        }
431
432        return found;
433    }
434
435    template<typename T>
436    bool llama_model_loader::get_key(enum llm_kv kid, T & result, bool required) {
437        return get_key(llm_kv(kid), result, required);
438    }
439
440    template bool llama_model_loader::get_key<bool>       (enum llm_kv kid, bool & result,        bool required);
441    template bool llama_model_loader::get_key<float>      (enum llm_kv kid, float & result,       bool required);
442    template bool llama_model_loader::get_key<uint32_t>   (enum llm_kv kid, uint32_t & result,    bool required);
443    template bool llama_model_loader::get_key<std::string>(enum llm_kv kid, std::string & result, bool required);
444
445    template<>
446    bool llama_model_loader::get_key(enum llm_kv kid, enum llama_pooling_type & result, bool required) {
447        uint32_t tmp;
448        const bool found = get_key(kid, tmp, required);
449        if (found) {
450            result = (enum llama_pooling_type) tmp;
451        } else {
452            result = LLAMA_POOLING_TYPE_UNSPECIFIED;
453        }
454        return found;
455    }
456
457    // get array of n <= N_MAX elements, or a single element repeated n times
458    template<typename T, size_t N_MAX>
459    bool llama_model_loader::get_key_or_arr(const std::string & key, std::array<T, N_MAX> & result, uint32_t n, bool required) {
460        const int kid = gguf_find_key(metadata, key.c_str());
461
462        if (kid < 0) {
463            if (required) {
464                throw std::runtime_error(format("key not found in model: %s", key.c_str()));
465            }
466            return false;
467        }
468
469        if (n > N_MAX) {
470            throw std::runtime_error(format("n > N_MAX: %u > %u for key %s", n, (uint32_t) N_MAX, key.c_str()));
471        }
472
473        if (gguf_get_kv_type(metadata, kid) == GGUF_TYPE_ARRAY) {
474            struct GGUFMeta::ArrayInfo arr_info =
475                GGUFMeta::GKV<GGUFMeta::ArrayInfo>::get_kv(metadata, kid);
476
477            if (n != arr_info.length) {
478                throw std::runtime_error(format("key %s has wrong array length; expected %u, got %u", key.c_str(), n, (uint32_t) arr_info.length));
479            }
480
481            return get_arr(key, result, required);
482        }
483
484        T value;
485
486        bool ok = get_key(key, value, required);
487        if (!ok) {
488            return false;
489        }
490
491        for (uint32_t i = 0; i < n; i++) {
492            result[i] = value;
493        }
494
495        return true;
496    }
497
498    template<typename T>
499    bool llama_model_loader::get_key_or_arr(enum llm_kv kid, T & result, uint32_t n, bool required) {
500        return get_key_or_arr(llm_kv(kid), result, n, required);
501    }
502
503    bool llama_model_loader::get_key_or_arr(enum llm_kv kid, uint32_t & result, bool required) {
504        const std::string key = llm_kv(kid);
505
506        const int id = gguf_find_key(metadata, key.c_str());
507
508        if (id < 0) {
509            if (required) {
510                throw std::runtime_error(format("key not found in model: %s", key.c_str()));
511            }
512            return false;
513        }
514
515        // throw and error if type is an array
516        if (gguf_get_kv_type(metadata, id) == GGUF_TYPE_ARRAY) {
517            if (required) {
518                throw std::runtime_error(format("expected scalar, found array for key: %s", key.c_str()));
519            }
520            return false;
521        }
522
523        return get_key(key, result, required);
524    }
525
526    // TODO: this is not very clever - figure out something better
527    template bool llama_model_loader::get_key_or_arr<std::array<int,      4>>  (enum llm_kv kid, std::array<int,      4>   & result, uint32_t n, bool required);
528    template bool llama_model_loader::get_key_or_arr<std::array<uint32_t, 512>>(enum llm_kv kid, std::array<uint32_t, 512> & result, uint32_t n, bool required);
529    template bool llama_model_loader::get_key_or_arr<std::array<float,    512>>(enum llm_kv kid, std::array<float,    512> & result, uint32_t n, bool required);
530
531
532llama_model_loader::llama_model_loader(
533        struct gguf_context * meta,
534        llama_model_set_tensor_data_t set_tensor_data,
535        void * set_tensor_data_ud,
536        const std::string & fname,
537        std::vector<std::string> & splits,
538        FILE * file,
539        llama_load_mode load_mode,
540        bool check_tensors,
541        bool no_alloc,
542        bool load_mtp,
543        const llama_model_kv_override * param_overrides_p,
544        const llama_model_tensor_buft_override * param_tensor_buft_overrides_p)
545        : metadata(meta), set_tensor_data(set_tensor_data), set_tensor_data_ud(set_tensor_data_ud) {
546    int trace = 0;
547    if (getenv("LLAMA_TRACE")) {
548        trace = atoi(getenv("LLAMA_TRACE"));
549    }
550
551    if (param_overrides_p != nullptr) {
552        for (const struct llama_model_kv_override * p = param_overrides_p; p->key[0] != 0; p++) {
553            kv_overrides.insert({std::string(p->key), *p});
554        }
555    }
556
557    tensor_buft_overrides = param_tensor_buft_overrides_p;
558
559    this->use_mmap      = load_mode == LLAMA_LOAD_MODE_MMAP || load_mode == LLAMA_LOAD_MODE_MMAP_MLOCK || load_mode == LLAMA_LOAD_MODE_AUTO;
560    this->use_direct_io = load_mode == LLAMA_LOAD_MODE_DIRECT_IO;
561
562    if (!fname.empty()) {
563        // Load the main GGUF
564        struct ggml_context * ctx = NULL;
565        struct gguf_init_params params = {
566            /*.no_alloc = */ true,
567            /*.ctx      = */ &ctx,
568        };
569
570        metadata_ptr.reset(gguf_init_from_file(fname.c_str(), params));
571        metadata = metadata_ptr.get();
572        if (metadata == nullptr) {
573            throw std::runtime_error(format("%s: failed to load model from %s", __func__, fname.c_str()));
574        }
575
576        get_key(llm_kv(LLM_KV_GENERAL_ARCHITECTURE), arch_name, false);
577        llm_kv = LLM_KV(llm_arch_from_string(arch_name));
578
579        files.emplace_back(new llama_file(fname.c_str(), "rb", use_direct_io));
580        contexts.emplace_back(ctx);
581
582        // Save tensors data offset of the main file.
583        // For subsidiary files, `meta` tensor data offset must not be used,
584        // so we build a unified tensors index for weights.
585        for (ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {
586            std::string tensor_name = std::string(cur->name);
587            // make sure there is no duplicated tensor names
588            if (weights_map.find(tensor_name) != weights_map.end()) {
589                throw std::runtime_error(format("invalid model: tensor '%s' is duplicated", ggml_get_name(cur)));
590            }
591            n_elements += ggml_nelements(cur);
592            n_bytes    += ggml_nbytes(cur);
593            weights_map.emplace(tensor_name, llama_tensor_weight(files.back().get(), 0, metadata, cur));
594        }
595        uint16_t n_split = 0;
596        get_key(llm_kv(LLM_KV_SPLIT_COUNT), n_split, false);
597
598        // Load additional GGML contexts
599        if (n_split > 1) {
600            // make sure the main file is loaded first
601            uint16_t idx = 0;
602            const std::string kv_split_no = llm_kv(LLM_KV_SPLIT_NO);
603            get_key(kv_split_no, idx);
604            if (idx != 0) {
605                throw std::runtime_error(format("illegal split file idx: %d (file: %s), model must be loaded with the first split", idx, fname.c_str()));
606            }
607
608            // generate list of splits if needed
609            if (splits.empty()) {
610                splits = llama_get_list_splits(fname, idx, n_split);
611            }
612
613            // in case user give a custom list of splits, check if it matches the expected number
614            if (n_split != (uint16_t)splits.size()) {
615                throw std::runtime_error(format("invalid split count, given: %zu splits, but expected %d", splits.size(), n_split));
616            }
617
618            if (trace > 0) {
619                LLAMA_LOG_INFO("%s: loading additional %d GGUFs\n", __func__, n_split);
620            }
621
622            // load other splits
623            for (idx = 1; idx < n_split; idx++) {
624                const char * fname_split = splits[idx].c_str();
625
626                struct gguf_init_params split_params = {
627                    /*.no_alloc = */ true,
628                    /*.ctx      = */ &ctx,
629                };
630                gguf_context_ptr ctx_gguf { gguf_init_from_file(fname_split, split_params) };
631                if (!ctx_gguf) {
632                    throw std::runtime_error(format("%s: failed to load GGUF split from %s", __func__, fname_split));
633                }
634
635                // check idx
636                {
637                    const int kid = gguf_find_key(ctx_gguf.get(), kv_split_no.c_str());
638                    if (kid < 0) {
639                        throw std::runtime_error(format("missing key %s in GGUF split %s", kv_split_no.c_str(), fname_split));
640                    }
641                    int idx_gguf = gguf_get_val_u16(ctx_gguf.get(), kid);
642                    if (idx_gguf != idx) {
643                        throw std::runtime_error(format("invalid split file idx: %d (file: %s), expected %d", idx_gguf, fname_split, idx));
644                    }
645                }
646
647                files.emplace_back(new llama_file(fname_split, "rb", use_direct_io));
648                contexts.emplace_back(ctx);
649
650                // Save tensors data offset info of the shard.
651                for (ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {
652                    std::string tensor_name = std::string(cur->name);
653                    // make sure there is no duplicated tensor names
654                    if (weights_map.find(tensor_name) != weights_map.end()) {
655                        throw std::runtime_error(format("invalid model: tensor '%s' is duplicated", ggml_get_name(cur)));
656                    }
657                    n_elements += ggml_nelements(cur);
658                    n_bytes    += ggml_nbytes(cur);
659                    weights_map.emplace(tensor_name, llama_tensor_weight(files.back().get(), idx, ctx_gguf.get(), cur));
660                }
661            }
662
663            get_key(llm_kv(LLM_KV_SPLIT_TENSORS_COUNT), n_tensors);
664
665            // sanity check
666            {
667                const int n_tensors_loaded = (int) weights_map.size();
668                if (n_tensors != n_tensors_loaded) {
669                    throw std::runtime_error(format("corrupted model: %d tensors expected but %d found", n_tensors, n_tensors_loaded));
670                }
671            }
672
673            LLAMA_LOG_INFO("%s: additional %d GGUFs metadata loaded.\n",  __func__, n_split - 1);
674        }
675    } else if (file != nullptr) {
676        struct ggml_context * ctx = NULL;
677        struct gguf_init_params params = {
678            /*.no_alloc = */ true,
679            /*.ctx      = */ &ctx,
680        };
681
682        metadata_ptr.reset(gguf_init_from_file_ptr(file, params));
683        metadata = metadata_ptr.get();
684        if (metadata == nullptr) {
685            throw std::runtime_error(format("%s: failed to load model from file pointer", __func__));
686        }
687
688        get_key(llm_kv(LLM_KV_GENERAL_ARCHITECTURE), arch_name, false);
689        llm_kv = LLM_KV(llm_arch_from_string(arch_name));
690
691        files.emplace_back(new llama_file(file));
692        contexts.emplace_back(ctx);
693
694        // Save tensors data offset info of the main file.
695        for (ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {
696            std::string tensor_name = std::string(cur->name);
697            // make sure there is no duplicated tensor names
698            if (weights_map.find(tensor_name) != weights_map.end()) {
699                throw std::runtime_error(format("invalid model: tensor '%s' is duplicated", ggml_get_name(cur)));
700            }
701            n_elements += ggml_nelements(cur);
702            n_bytes    += ggml_nbytes(cur);
703            weights_map.emplace(tensor_name, llama_tensor_weight(files.back().get(), 0, metadata, cur));
704        }
705    } else {
706        get_key(llm_kv(LLM_KV_GENERAL_ARCHITECTURE), arch_name, false);
707        llm_kv = LLM_KV(llm_arch_from_string(arch_name));
708    }
709
710    n_kv      = gguf_get_n_kv(metadata);
711    n_tensors = weights_map.size();
712
713    fver = (enum llama_fver) gguf_get_version(metadata);
714
715    LLAMA_LOG_INFO("%s: loaded meta data with %d key-value pairs and %d tensors from %s (version %s)\n",
716            __func__, n_kv, n_tensors, fname.empty() ? "(file*)" : fname.c_str(), llama_file_version_name(fver));
717
718    // determine file type based on the number of tensors for each quantization and print meta data
719    // TODO: make optional
720    {
721        std::map<enum ggml_type, uint32_t> n_type;
722
723        uint32_t n_type_max = 0;
724        enum ggml_type type_max = GGML_TYPE_F32;
725
726        for (const auto & it : weights_map) {
727            const llama_tensor_weight & w = it.second;
728            const ggml_tensor * tensor = w.tensor;
729
730            enum ggml_type type = tensor->type;
731
732            n_type[type]++;
733
734            if (n_type_max < n_type[type]) {
735                n_type_max = n_type[type];
736                type_max   = type;
737            }
738
739            if (trace > 0) {
740                const uint16_t sid = w.idx;
741                LLAMA_LOG_INFO("%s: - tensor split %2d: %32s %-8s [ %s ] %8.2f MiB\n", __func__,
742                        sid, ggml_get_name(tensor), ggml_type_name(type), llama_format_tensor_shape(tensor).c_str(),
743                        ggml_nbytes(tensor)/1024.0f/1024.0f);
744            }
745        }
746
747        switch (type_max) {
748            case GGML_TYPE_F32:     ftype = LLAMA_FTYPE_ALL_F32;        break;
749            case GGML_TYPE_F16:     ftype = LLAMA_FTYPE_MOSTLY_F16;     break;
750            case GGML_TYPE_BF16:    ftype = LLAMA_FTYPE_MOSTLY_BF16;    break;
751            case GGML_TYPE_Q4_0:    ftype = LLAMA_FTYPE_MOSTLY_Q4_0;    break;
752            case GGML_TYPE_Q4_1:    ftype = LLAMA_FTYPE_MOSTLY_Q4_1;    break;
753            case GGML_TYPE_Q5_0:    ftype = LLAMA_FTYPE_MOSTLY_Q5_0;    break;
754            case GGML_TYPE_Q5_1:    ftype = LLAMA_FTYPE_MOSTLY_Q5_1;    break;
755            case GGML_TYPE_Q8_0:    ftype = LLAMA_FTYPE_MOSTLY_Q8_0;    break;
756            case GGML_TYPE_Q2_K:    ftype = LLAMA_FTYPE_MOSTLY_Q2_K;    break;
757            case GGML_TYPE_Q3_K:    ftype = LLAMA_FTYPE_MOSTLY_Q3_K_M;  break;
758            case GGML_TYPE_Q4_K:    ftype = LLAMA_FTYPE_MOSTLY_Q4_K_M;  break;
759            case GGML_TYPE_Q5_K:    ftype = LLAMA_FTYPE_MOSTLY_Q5_K_M;  break;
760            case GGML_TYPE_Q6_K:    ftype = LLAMA_FTYPE_MOSTLY_Q6_K;    break;
761            case GGML_TYPE_TQ1_0:   ftype = LLAMA_FTYPE_MOSTLY_TQ1_0;   break;
762            case GGML_TYPE_TQ2_0:   ftype = LLAMA_FTYPE_MOSTLY_TQ2_0;   break;
763            case GGML_TYPE_IQ2_XXS: ftype = LLAMA_FTYPE_MOSTLY_IQ2_XXS; break;
764            case GGML_TYPE_IQ2_XS:  ftype = LLAMA_FTYPE_MOSTLY_IQ2_XS;  break;
765            case GGML_TYPE_IQ2_S:   ftype = LLAMA_FTYPE_MOSTLY_IQ2_S;   break;
766            case GGML_TYPE_IQ3_XXS: ftype = LLAMA_FTYPE_MOSTLY_IQ3_XXS; break;
767            case GGML_TYPE_IQ1_S:   ftype = LLAMA_FTYPE_MOSTLY_IQ1_S;   break;
768            case GGML_TYPE_IQ1_M:   ftype = LLAMA_FTYPE_MOSTLY_IQ1_M;   break;
769            case GGML_TYPE_IQ4_NL:  ftype = LLAMA_FTYPE_MOSTLY_IQ4_NL;  break;
770            case GGML_TYPE_IQ4_XS:  ftype = LLAMA_FTYPE_MOSTLY_IQ4_XS;  break;
771            case GGML_TYPE_IQ3_S:   ftype = LLAMA_FTYPE_MOSTLY_IQ3_S;   break;
772            case GGML_TYPE_NVFP4:   ftype = LLAMA_FTYPE_MOSTLY_NVFP4;   break;
773            case GGML_TYPE_Q1_0:    ftype = LLAMA_FTYPE_MOSTLY_Q1_0;    break;
774            case GGML_TYPE_Q2_0:    ftype = LLAMA_FTYPE_MOSTLY_Q2_0;    break;
775            default:
776                {
777                    LLAMA_LOG_WARN("%s: unknown type %s\n", __func__, ggml_type_name(type_max));
778                    ftype = LLAMA_FTYPE_ALL_F32;
779                } break;
780        }
781
782        // this is a way to mark that we have "guessed" the file type
783        ftype = (llama_ftype) (ftype | LLAMA_FTYPE_GUESSED);
784
785        {
786            uint32_t ftype_val = 0;
787            if (get_key(LLM_KV_GENERAL_FILE_TYPE, ftype_val, false)) {
788                ftype = (llama_ftype) ftype_val;
789            }
790        }
791
792        LLAMA_LOG_INFO("%s: Dumping metadata keys/values. Note: KV overrides do not apply in this output.\n", __func__);
793
794        for (int i = 0; i < n_kv; i++) {
795            const char * name           = gguf_get_key(metadata, i);
796            const enum gguf_type type   = gguf_get_kv_type(metadata, i);
797            const std::string type_name =
798                type == GGUF_TYPE_ARRAY
799                ? format("%s[%s,%zu]", gguf_type_name(type), gguf_type_name(gguf_get_arr_type(metadata, i)), gguf_get_arr_n(metadata, i))
800                : gguf_type_name(type);
801
802            std::string value          = gguf_kv_to_str(metadata, i);
803            const size_t MAX_VALUE_LEN = 40;
804            if (value.size() > MAX_VALUE_LEN) {
805                value = format("%s...", value.substr(0, MAX_VALUE_LEN - 3).c_str());
806            }
807            replace_all(value, "\n", "\\n");
808
809            LLAMA_LOG_INFO("%s: - kv %3d: %42s %-16s = %s\n", __func__, i, name, type_name.c_str(), value.c_str());
810        }
811
812        // print type counts
813        for (auto & kv : n_type) {
814            if (kv.second == 0) {
815                continue;
816            }
817
818            LLAMA_LOG_INFO("%s: - type %4s: %4d tensors\n", __func__, ggml_type_name(kv.first), kv.second);
819        }
820    }
821
822    if (this->use_mmap && !llama_mmap::SUPPORTED) {
823        LLAMA_LOG_WARN("%s: mmap is not supported on this platform\n", __func__);
824        this->use_mmap = false;
825    }
826
827    this->check_tensors = check_tensors;
828    this->no_alloc = no_alloc;
829    this->load_mtp = load_mtp;
830}
831
832std::string llama_model_loader::get_arch_name() const {
833    return arch_name;
834}
835
836enum llm_arch llama_model_loader::get_arch() const {
837    return llm_kv.arch;
838}
839
840const llama_model_loader::llama_tensor_weight * llama_model_loader::get_weight(const char * name) const {
841    auto pos = weights_map.find(name);
842    if (pos != weights_map.end()) {
843        return &pos->second;
844    }
845
846    return nullptr;
847}
848
849const llama_model_loader::llama_tensor_weight & llama_model_loader::require_weight(const char * name) const {
850    const llama_tensor_weight * weight = get_weight(name);
851    if (!weight) {
852        throw std::runtime_error(format("%s: tensor '%s' not found", __func__, name));
853    }
854    return *weight;
855}
856
857struct ggml_tensor * llama_model_loader::get_tensor_meta(const char * name) const {
858    const auto * weight = get_weight(name);
859    if (!weight) {
860        return nullptr;
861    }
862    return weight->tensor;
863}
864
865struct ggml_tensor * llama_model_loader::require_tensor_meta(const std::string & name) const {
866    struct ggml_tensor * tensor = get_tensor_meta(name.c_str());
867    if (!tensor) {
868        throw std::runtime_error(format("%s: tensor '%s' not found", __func__, name.c_str()));
869    }
870    return tensor;
871}
872
873const struct ggml_tensor * llama_model_loader::check_tensor_dims(
874        const std::string & name,
875        const std::vector<int64_t> & ne,
876        bool required,
877        bool allow_reshape) const {
878    const struct ggml_tensor * cur = get_tensor_meta(name.c_str());
879
880    if (cur == NULL) {
881        if (!required) {
882            return NULL;
883        }
884        throw std::runtime_error(format("%s: tensor '%s' not found", __func__, name.c_str()));
885    }
886
887    bool is_ok = true;
888
889    if (allow_reshape) {
890        // check total number of elements only
891        const int64_t ncur = ggml_nelements(cur);
892        int64_t nexp = 1;
893        for (size_t i = 0; i < ne.size(); ++i) {
894            nexp *= ne[i];
895        }
896        if (ncur != nexp) {
897            is_ok = false;
898        }
899    } else {
900        for (size_t i = 0; i < GGML_MAX_DIMS; ++i) {
901            if ((i < ne.size() && ne[i] != cur->ne[i]) || (i >= ne.size() && cur->ne[i] != 1)) {
902                is_ok = false;
903                break;
904            }
905        }
906    }
907
908    if (!is_ok) {
909        throw std::runtime_error(
910                format("%s: tensor '%s' has wrong shape; expected %s, got %s",
911                    __func__, name.c_str(),
912                    llama_format_tensor_shape(ne).c_str(),
913                    llama_format_tensor_shape(cur).c_str()));
914    }
915
916    return cur;
917}
918
919// checks if the weight tensor can be used with the specified buffer type and device
920static bool weight_buft_supported(const llama_hparams & hparams, ggml_tensor * w, ggml_op op, ggml_backend_buffer_type_t buft, ggml_backend_dev_t dev) {
921    GGML_ASSERT(w != nullptr);
922
923    if (op == GGML_OP_NONE) {
924        return true;
925    }
926
927    ggml_init_params params = {
928        /*.mem_size   =*/ ggml_tensor_overhead()*8,
929        /*.mem_buffer =*/ NULL,
930        /*.no_alloc   =*/ true,
931    };
932    ggml_context_ptr ctx_ptr { ggml_init(params) };
933    if (!ctx_ptr) {
934        throw std::runtime_error(format("failed to create ggml context"));
935    }
936    ggml_context * ctx = ctx_ptr.get();
937
938    ggml_tensor * op_tensor = nullptr;
939
940    switch (op) {
941        case GGML_OP_GET_ROWS:
942            {
943                ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512);
944                op_tensor = ggml_get_rows(ctx, w, b);
945            } break;
946        case GGML_OP_MUL_MAT:
947            {
948                ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], 512, w->ne[2], w->ne[3]);
949                op_tensor = ggml_mul_mat(ctx, w, b);
950            } break;
951        case GGML_OP_MUL_MAT_ID:
952            {
953                // Used for either MoE expert routing or embedded adapter routing
954                const int n_ids_used = hparams.router_layer >= 0 ? 1 : hparams.n_expert_used_max();
955                GGML_ASSERT(n_ids_used > 0);
956                ggml_tensor * b = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_ids_used, 512);
957                ggml_tensor * ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_ids_used, 512);
958                op_tensor = ggml_mul_mat_id(ctx, w, b, ids);
959            } break;
960        case GGML_OP_ADD:
961            {
962                ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], w->ne[1], w->ne[2], w->ne[3]);
963                op_tensor = ggml_add(ctx, a, w);
964            } break;
965        case GGML_OP_ADD_ID:
966            {
967                const int n_expert_used = hparams.n_expert_used_max();
968                GGML_ASSERT(n_expert_used > 0);
969                ggml_tensor * a = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0], n_expert_used, 512);
970                ggml_tensor * c = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, 512);
971                op_tensor = ggml_add_id(ctx, a, w, c);
972            } break;
973        case GGML_OP_MUL:
974            {
975                ggml_tensor * a = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, w->ne[0], w->ne[1], w->ne[2], w->ne[3]);
976                op_tensor = ggml_mul(ctx, a, w);
977            } break;
978        case GGML_OP_DIV:
979            {
980                ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, w->ne[0]);
981                op_tensor = ggml_div(ctx, a, w);
982            } break;
983        case GGML_OP_ROPE:
984            {
985                const int n_embd_head = hparams.n_embd_head_v();
986                const int n_head = hparams.n_head();
987                ggml_tensor * a = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_head, n_head, 512);
988                ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512);
989                op_tensor = ggml_rope_ext(
990                    ctx, a, b, w,
991                    0, 0, 0, 0, 0,
992                    0, 0, 0, 0
993                );
994
995            } break;
996        case GGML_OP_SSM_CONV:
997            {
998                const int64_t n_seq_tokens = 512;
999                const int64_t n_seqs       = 3;
1000                ggml_tensor * conv_x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, w->ne[0] - 1 + n_seq_tokens, w->ne[1], n_seqs);
1001                op_tensor = ggml_ssm_conv(ctx, conv_x, w);
1002            } break;
1003        case GGML_OP_SSM_SCAN:
1004            {
1005                // w is ssm_a, which is used to distinguish Mamba-1 and Mamba-2
1006                const int64_t d_state      = w->ne[0] == 1 ? hparams.ssm_d_state : w->ne[0];
1007                const int64_t n_head       = w->ne[1];
1008                const int64_t head_dim     = hparams.ssm_d_inner / n_head;
1009                const int64_t n_group      = hparams.ssm_n_group ? hparams.ssm_n_group : 1;
1010                const int64_t n_seq_tokens = 512;
1011                const int64_t n_seqs       = 3;
1012                ggml_tensor * s   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, head_dim, n_head, n_seqs);
1013                ggml_tensor * x   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, head_dim, n_head, n_seq_tokens, n_seqs);
1014                ggml_tensor * dt  = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_head, n_seq_tokens, n_seqs);
1015                ggml_tensor * B   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
1016                ggml_tensor * C   = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, d_state, n_group, n_seq_tokens, n_seqs);
1017                ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_seqs);
1018                op_tensor = ggml_ssm_scan(ctx, s, x, dt, w, B, C, ids, /*K=*/1);
1019            } break;
1020        case GGML_OP_RWKV_WKV6:
1021            {
1022                // FIXME
1023                const int64_t S = 123;
1024                const int64_t H = 123;
1025                const int64_t n_tokens = 123;
1026                const int64_t n_seqs = 123;
1027                ggml_tensor  * k = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, H, n_tokens);
1028                ggml_tensor  * v = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, H, n_tokens);
1029                ggml_tensor  * r = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, H, n_tokens);
1030                ggml_tensor  * tf = w;
1031                ggml_tensor  * td = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, S, H, n_tokens);
1032                ggml_tensor  * state = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, S, n_seqs, S, H);
1033                op_tensor = ggml_rwkv_wkv6(ctx, k, v, r, tf, td, state);
1034            } break;
1035        case GGML_OP_IM2COL:
1036            {
1037                const int n_embd_inp = hparams.n_embd_inp();
1038                ggml_tensor * b = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_embd_inp, w->ne[1], 1, 1);
1039                op_tensor = ggml_im2col(ctx, w, b, 1, 0, 0, 0, 1, 0, false, GGML_TYPE_F16);
1040            } break;
1041        case GGML_OP_SCALE:
1042            {
1043                op_tensor = ggml_scale(ctx, w, 1.0f);
1044            } break;
1045        default:
1046            GGML_ABORT("%s: missing test for op %s for tensor %s", __func__, ggml_op_name(op), w->name);
1047    }
1048
1049    // create a temporary dummy buffer for the weight so that supports_op can check the buffer type
1050    GGML_ASSERT(w->buffer == nullptr);
1051    w->buffer = ggml_backend_buft_alloc_buffer(buft, 0);
1052    bool op_supported = ggml_backend_dev_supports_op(dev, op_tensor);
1053    ggml_backend_buffer_free(w->buffer);
1054    w->buffer = nullptr;
1055
1056    return op_supported;
1057}
1058
1059// find the first buffer type in the list that can use the tensor
1060static ggml_backend_buffer_type_t select_weight_buft(const llama_hparams & hparams, ggml_tensor * tensor, ggml_op op, const buft_list_t * buft_list) {
1061    GGML_ASSERT(!buft_list->empty());
1062    for (const auto & cur : *buft_list) {
1063        ggml_backend_dev_t cur_dev = cur.first;
1064        ggml_backend_buffer_type_t cur_buft = cur.second;
1065        if (weight_buft_supported(hparams, tensor, op, cur_buft, cur_dev)) {
1066            return cur_buft;
1067        }
1068    }
1069
1070    return nullptr;
1071}
1072
1073ggml_backend_buffer_type_t llama_model_loader::lazy_read::buft() {
1074    auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
1075    if (!cpu_dev) {
1076        throw std::runtime_error("no CPU backend found");
1077    }
1078    return ggml_backend_dev_buffer_type(cpu_dev);
1079}
1080
1081bool llama_model_loader::lazy_read::add(const std::string & name, const ggml_tensor * t, const llama_tensor_weight * w) {
1082    if (mode == LLAMA_LAZY_MODE_OFF) {
1083        return false;
1084    }
1085
1086    // do not lazy-read small tensors, it has significant overhead and is not worth it
1087    constexpr size_t auto_min_size = 4ull * 1024 * 1024 * 1024;
1088    if (mode != LLAMA_LAZY_MODE_ON && ggml_nbytes(t) <= auto_min_size) {
1089        return false;
1090    }
1091
1092    if (!llama_mmap::SUPPORTED) {
1093        LLAMA_LOG_WARN("%s: mmap is not available, so tensor %s (size = %zu MiB) is loaded into RAM in full\n",
1094                __func__, name.c_str(), ggml_nbytes(t)/1024/1024);
1095        return false;
1096    }
1097
1098    if (w) {
1099        ranges[w->idx].emplace_back(w->offs, w->offs + ggml_nbytes(t));
1100        tensors.insert(name);
1101
1102        LLAMA_LOG_INFO("%s: tensor %s (size = %zu MiB) lazy read enabled\n",
1103                __func__, name.c_str(), ggml_nbytes(t)/1024/1024);
1104    }
1105
1106    return true;
1107}
1108
1109struct ggml_tensor * llama_model_loader::create_tensor(
1110        const llama_hparams & hparams, const buft_list_t * buft_list_cpu, const buft_list_t * buft_list_input, const buft_list_t * buft_list_output,
1111        const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags) {
1112    // set below, before buft_for_tensor() runs
1113    bool is_lazy = false;
1114
1115    auto ctx_for_buft = [&](ggml_backend_buffer_type_t buft) -> ggml_context * {
1116        const ctx_key key { buft, is_lazy };
1117
1118        auto it = ctx_map.find(key);
1119        if (it == ctx_map.end()) {
1120            // one ggml context per buffer type
1121            int max_n_tensors = n_tensors;
1122            max_n_tensors += 1;                   // duplicated output tensor
1123            max_n_tensors += hparams.n_layer()*2; // duplicated rope freq tensors
1124            if (files.empty()) {
1125                max_n_tensors += hparams.n_layer()*256; // this should be well above what any model actually uses
1126            }
1127            const size_t ctx_size = ggml_tensor_overhead()*max_n_tensors;
1128
1129            ggml_init_params params = {
1130                /*.mem_size   =*/ ctx_size,
1131                /*.mem_buffer =*/ NULL,
1132                /*.no_alloc   =*/ true,
1133            };
1134
1135            ggml_context * ctx = ggml_init(params);
1136            if (!ctx) {
1137                throw std::runtime_error(format("failed to create ggml context"));
1138            }
1139
1140            ctx_map.emplace(key, ctx);
1141
1142            return ctx;
1143        }
1144        return it->second.get();
1145    };
1146
1147    auto buft_for_tensor = [&](ggml_tensor * t_meta) -> ggml_backend_buffer_type_t {
1148        if (!t_meta) {
1149            if (flags & TENSOR_NOT_REQUIRED) {
1150                return nullptr;
1151            }
1152            throw std::runtime_error(format("missing tensor '%s'", tn.str().c_str()));
1153        }
1154
1155        // some models use the token embedding tensor as the output, but since these are used in different layers and with different ops
1156        // the tensor is duplicated
1157        // to handle this, we check if the tensor is duplicated, and if so, we assume that it is being loaded as the output tensor
1158        llm_tensor tn_tensor = tn.tensor;
1159        if (tn.tensor == LLM_TENSOR_TOKEN_EMBD && (flags & TENSOR_DUPLICATED)) {
1160            tn_tensor = LLM_TENSOR_OUTPUT;
1161        }
1162
1163        llm_tensor_info info;
1164        try {
1165            info = llm_tensor_info_for(tn_tensor);
1166        } catch (const std::out_of_range & e) {
1167            throw std::runtime_error(format("missing tensor info mapping for %s", tn.str().c_str()));
1168        }
1169
1170        // skip unused tensors
1171        if (info.op == GGML_OP_NONE || (flags & TENSOR_SKIP)) {
1172            const size_t nbytes = ggml_nbytes(t_meta);
1173            LLAMA_LOG_WARN("model has unused tensor %s (size = %zu bytes) -- ignoring\n", tn.str().c_str(), nbytes);
1174
1175            size_data -= nbytes;
1176            n_created++;
1177
1178            return nullptr;
1179        }
1180
1181        // tensors with "bias" suffix are always used with GGML_OP_ADD or GGML_OP_ADD_ID;
1182        // embedded-adapter ".lora_a"/".lora_b" tensors are always used with GGML_OP_MUL_MAT_ID
1183        ggml_op op;
1184        if (tn.suffix != nullptr && strcmp(tn.suffix, "bias") == 0) {
1185            op = info.op == GGML_OP_MUL_MAT_ID ? GGML_OP_ADD_ID : GGML_OP_ADD;
1186        } else if (hparams.router_layer >= 0 && tn.suffix != nullptr &&
1187                (strcmp(tn.suffix, "lora_a") == 0 || strcmp(tn.suffix, "lora_b") == 0)) {
1188            op = GGML_OP_MUL_MAT_ID;
1189        } else {
1190            op = info.op;
1191        }
1192
1193        // sanity checks
1194        if (info.layer == LLM_TENSOR_LAYER_INPUT || info.layer == LLM_TENSOR_LAYER_OUTPUT) {
1195            if (tn.bid != -1) {
1196                GGML_ABORT("input/output layer tensor %s used with a layer number", tn.str().c_str());
1197            }
1198        } else {
1199            if (tn.bid == -1) {
1200                GGML_ABORT("repeating layer tensor %s used without a layer number", tn.str().c_str());
1201            }
1202        }
1203
1204        if (is_lazy) {
1205            return lazy_read::buft();
1206        }
1207
1208        // select the buffer type for this tensor
1209        const buft_list_t * buft_list;
1210        switch (info.layer) {
1211            case LLM_TENSOR_LAYER_INPUT:
1212                buft_list = buft_list_input;
1213                break;
1214            case LLM_TENSOR_LAYER_OUTPUT:
1215                buft_list = buft_list_output;
1216                break;
1217            case LLM_TENSOR_LAYER_REPEATING:
1218                GGML_ASSERT(buft_list_layer != nullptr);
1219                buft_list = buft_list_layer;
1220                break;
1221            default:
1222                GGML_ABORT("invalid layer %d for tensor %s", info.layer, tn.str().c_str());
1223        }
1224
1225        ggml_backend_buffer_type_t buft = nullptr;
1226
1227        // check overrides
1228        if (tensor_buft_overrides) {
1229            std::string tensor_name = tn.str();
1230            for (const auto * overrides = tensor_buft_overrides; overrides->pattern != nullptr; ++overrides) {
1231                std::regex pattern(overrides->pattern);
1232                if (std::regex_search(tensor_name, pattern)) {
1233                    if (overrides->buft == ggml_backend_cpu_buffer_type()) {
1234                        // when overriding to a CPU buffer, consider the extra buffer types
1235                        buft = select_weight_buft(hparams, t_meta, op, buft_list_cpu);
1236                        if (use_mmap) {
1237                            static std::once_flag once;
1238                            std::call_once(once, [] {
1239                                LLAMA_LOG_WARN("llama_model_loader: tensor overrides to CPU are used with mmap enabled - consider using --load-mode none for better performance\n");
1240                            });
1241                        }
1242                    } else {
1243                        buft = overrides->buft;
1244                    }
1245
1246                    LLAMA_LOG_DEBUG("tensor %s (%zu MiB %s) buffer type overridden to %s\n",
1247                            tensor_name.c_str(),
1248                            ggml_nbytes(t_meta) / 1024 / 1024, ggml_type_name(t_meta->type),
1249                            ggml_backend_buft_name(buft));
1250                    break;
1251                }
1252            }
1253        }
1254
1255        if (!buft) {
1256            buft = select_weight_buft(hparams, t_meta, op, buft_list);
1257            if (!buft) {
1258                throw std::runtime_error(format("failed to find a compatible buffer type for tensor %s", tn.str().c_str()));
1259            }
1260        }
1261
1262        // avoid using a host buffer when using mmap
1263        auto * buft_dev = ggml_backend_buft_get_device(buft);
1264        if (use_mmap && buft_dev && buft == ggml_backend_dev_host_buffer_type(buft_dev)) {
1265            auto * cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
1266            if (!cpu_dev) {
1267                throw std::runtime_error("no CPU backend found");
1268            }
1269            buft = ggml_backend_dev_buffer_type(cpu_dev);
1270        }
1271
1272        if (buft != buft_list->front().second) {
1273            if (n_tensors_moved == 0) {
1274                first_tensor_moved_name = t_meta->name;
1275                first_tensor_moved_type_name = ggml_type_name(t_meta->type);
1276                first_moved_from_buft = buft_list->front().second;
1277                first_moved_to_buft   = buft;
1278            }
1279            n_tensors_moved++;
1280        }
1281
1282        return buft;
1283    };
1284
1285    if (files.empty()) {
1286        if (flags & TENSOR_SKIP_IF_VIRTUAL) {
1287            return nullptr;
1288        }
1289        ggml_type type = GGML_TYPE_F32;
1290        const int64_t tid = gguf_find_tensor(metadata, tn.str().c_str());
1291        if (tid != -1) {
1292            type = gguf_get_tensor_type(metadata, tid);
1293        }
1294
1295        // for tensors that are not required some of the dimensions can be invalid:
1296        if (flags & TENSOR_NOT_REQUIRED) {
1297            for (size_t dim = 0; dim < ne.size(); dim++) {
1298                if (ne.begin()[dim] <= 0) {
1299                    return nullptr;
1300                }
1301            }
1302        }
1303
1304        ggml_tensor t_meta;
1305        memset(&t_meta, 0, sizeof(ggml_tensor));
1306        t_meta.type = type;
1307        for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
1308            t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1;
1309            GGML_ASSERT(t_meta.ne[dim] >= 1);
1310            if (dim == 0) {
1311                t_meta.nb[dim] = ggml_type_size(type);
1312            } else if (dim == 1) {
1313                t_meta.nb[dim] = ggml_row_size(type, t_meta.ne[dim-1]);
1314            } else {
1315                t_meta.nb[dim] = t_meta.nb[dim-1]*t_meta.ne[dim-1];
1316            }
1317            GGML_ASSERT(t_meta.nb[dim] >= 1);
1318        }
1319        ggml_set_name(&t_meta, tn.str().c_str());
1320
1321        ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta);
1322        GGML_ASSERT(buft != nullptr);
1323        ggml_context * ctx = ctx_for_buft(buft);
1324        ggml_tensor * ret = ggml_dup_tensor(ctx, &t_meta);
1325        ggml_set_name(ret, tn.str().c_str());
1326        return ret;
1327    }
1328
1329    LLAMA_LOG_DEBUG("%s: loading tensor %s\n", __func__, tn.str().c_str());
1330    const struct ggml_tensor * cur = check_tensor_dims(tn.str(), ne, !(flags & TENSOR_NOT_REQUIRED), flags & TENSOR_ALLOW_RESHAPE);
1331    if (cur == NULL) {
1332        return NULL;
1333    }
1334
1335    if (flags & TENSOR_READ_LAZY) {
1336        // the decision must not depend on the load mode, or the memory-fit pass (no_alloc, no mmap)
1337        is_lazy = lazy.add(tn.str(), cur, no_alloc ? nullptr : &require_weight(tn.str().c_str()));
1338    }
1339
1340    ggml_tensor t_meta = *cur;
1341    if (flags & TENSOR_ALLOW_RESHAPE) {
1342        for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) {
1343            t_meta.ne[dim] = dim < ne.size() ? ne.begin()[dim] : 1;
1344            if (dim == 0) {
1345                t_meta.nb[dim] = ggml_type_size(t_meta.type);
1346            } else if (dim == 1) {
1347                t_meta.nb[dim] = ggml_row_size(t_meta.type, t_meta.ne[dim-1]);
1348            } else {
1349                t_meta.nb[dim] = t_meta.ne[dim-1]*t_meta.nb[dim-1];
1350            }
1351        }
1352    }
1353
1354    GGML_ASSERT(ggml_nbytes(&t_meta) == ggml_nbytes(cur));
1355
1356    ggml_backend_buffer_type_t buft = buft_for_tensor(&t_meta);
1357    if (buft == nullptr) {
1358        return nullptr;
1359    }
1360
1361    ggml_context * ctx = ctx_for_buft(buft);
1362
1363    // if duplicated, check if the original tensor was allocated in the same buffer type context and avoid creating a new one
1364    if (flags & TENSOR_DUPLICATED) {
1365        ggml_tensor * t = ggml_get_tensor(ctx, tn.str().c_str());
1366        if (t) {
1367            return t;
1368        }
1369    }
1370
1371    const bool duplicated = flags & TENSOR_DUPLICATED;
1372
1373    struct ggml_tensor * tensor = ggml_dup_tensor(ctx, &t_meta);
1374    ggml_set_name(tensor, ggml_get_name(&t_meta));
1375
1376    if (duplicated) {
1377        size_data += ggml_nbytes(&t_meta);
1378    } else {
1379        n_created++;
1380    }
1381
1382    return tensor;
1383}
1384
1385void llama_model_loader::done_getting_tensors(bool partial) const {
1386    if (n_created > n_tensors) {
1387        throw std::runtime_error(format("%s: too many tensors created; expected %d, got %d", __func__, n_tensors, n_created));
1388    }
1389    if (n_created < n_tensors) {
1390        if (!partial) {
1391            throw std::runtime_error(format("%s: wrong number of tensors; expected %d, got %d", __func__, n_tensors, n_created));
1392        }
1393        LLAMA_LOG_INFO("%s: partial load — used %d of %d tensors in the file (rest belong to a sibling model on the same .gguf)\n",
1394                __func__, n_created, n_tensors);
1395    }
1396    if (n_tensors_moved > 0) {
1397        LLAMA_LOG_DEBUG("%s: tensor '%s' (%s) (and %zu others) cannot be used with preferred buffer type %s, using %s instead\n",
1398            __func__, first_tensor_moved_name.c_str(), first_tensor_moved_type_name.c_str(), n_tensors_moved - 1,
1399            ggml_backend_buft_name(first_moved_from_buft), ggml_backend_buft_name(first_moved_to_buft));
1400    }
1401}
1402
1403void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps) {
1404    // note: read_lazy also requires mmap; this condition make sure it's usable even when --load-mode is not set to mmap
1405    if (use_mmap || lazy.any()) {
1406        mappings.reserve(files.size());
1407        mmaps_used.reserve(files.size());
1408        for (uint32_t idx = 0; idx < files.size(); idx++) {
1409            const auto & file = files[idx];
1410
1411            bool is_numa = false;
1412
1413            auto * dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU);
1414            if (dev) {
1415                auto * reg = ggml_backend_dev_backend_reg(dev);
1416                auto * is_numa_fn = (decltype(ggml_is_numa) *) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cpu_is_numa");
1417                if (is_numa_fn) {
1418                    is_numa = is_numa_fn();
1419                }
1420            }
1421
1422            const size_t prefetch_size = prefetch && use_mmap ? -1 : 0;
1423
1424            std::unique_ptr<llama_mmap> mapping = std::make_unique<llama_mmap>(file.get(), prefetch_size, is_numa,
1425                    lazy.for_file(idx));
1426            mmaps_used.emplace_back(mapping->size(), 0);
1427            if (mlock_mmaps) {
1428                std::unique_ptr<llama_mlock> mlock_mmap(new llama_mlock());
1429                mlock_mmap->init(mapping->addr());
1430                mlock_mmaps->emplace_back(std::move(mlock_mmap));
1431            }
1432            mappings.emplace_back(std::move(mapping));
1433        }
1434    }
1435
1436    // compute the total size of all tensors for progress reporting
1437    for (const auto & it : weights_map) {
1438        size_data += ggml_nbytes(it.second.tensor);
1439    }
1440}
1441
1442void llama_model_loader::get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const {
1443    GGML_ASSERT(!mappings.empty());
1444    const auto & mapping = mappings.at(idx);
1445
1446    *first = mapping->size();
1447    *last  = 0;
1448    *addr = mapping->addr();
1449    for (ggml_tensor * tensor = ggml_get_first_tensor(ctx); tensor; tensor = ggml_get_next_tensor(ctx, tensor)) {
1450        const auto * weight = get_weight(ggml_get_name(tensor));
1451        if (!weight || weight->idx != idx) {
1452            continue;
1453        }
1454        *first = std::min(*first, weight->offs);
1455        *last  = std::max(*last,  weight->offs + ggml_nbytes(tensor));
1456    }
1457}
1458
1459void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const {
1460    if (!use_mmap) { return; }
1461    mappings.at(w.idx)->unmap_fragment(w.offs, w.offs + ggml_nbytes(w.tensor));
1462}
1463
1464const void * llama_model_loader::load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const {
1465    GGML_ASSERT(offs + size <= ggml_nbytes(w.tensor));
1466
1467    const void * data = buf;
1468
1469    if (use_mmap) {
1470        data = (const uint8_t *) mappings.at(w.idx)->addr() + w.offs + offs;
1471    } else {
1472        GGML_ASSERT(buf != nullptr);
1473        GGML_ASSERT(w.idx < files.size());
1474        const auto & file = files.at(w.idx);
1475        file->seek(w.offs + offs, SEEK_SET);
1476        file->read_raw(buf, size);
1477    }
1478
1479    if (check_tensors && !ggml_validate_row_data(w.tensor->type, data, size)) {
1480        throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(w.tensor)));
1481    }
1482
1483    return data;
1484}
1485
1486bool llama_model_loader::load_all_data(
1487        struct ggml_context * ctx,
1488        llama_buf_map & bufs,
1489        llama_mlocks * lmlocks,
1490        llama_progress_callback progress_callback,
1491        void * progress_callback_user_data) {
1492    if (files.empty()) {
1493        for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
1494            set_tensor_data(t, set_tensor_data_ud);
1495        }
1496        return true;
1497    }
1498    GGML_ASSERT(size_data != 0 && "call init_mappings() first");
1499
1500    std::vector<std::future<std::pair<ggml_tensor *, bool>>> validation_result;
1501
1502    // 4 staging buffers for async uploads, each sized 1MB seems to be a good default for single NVMe drives.
1503    // NVMe raid configurations might require more / larger buffers.
1504    constexpr size_t n_buffers = 4;
1505
1506    size_t alignment = 1;
1507    for (const auto & file : files) {
1508        alignment = std::max(file->read_alignment(), alignment);
1509    }
1510
1511    // Buffer size: balance between memory usage and I/O efficiency
1512    // 64MB works well for NVMe drives
1513    const size_t buffer_size = alignment != 1 ? 64 * 1024 * 1024 + 2 * alignment : 1 * 1024 * 1024;
1514
1515    std::vector<ggml_backend_buffer_t> host_buffers;
1516    std::vector<ggml_backend_event_t> events;
1517    std::vector<void *> host_ptrs;
1518    size_t buffer_idx = 0; // buffer to use for async loads
1519    ggml_backend_t upload_backend = [&](const char * func) -> ggml_backend_t {
1520        if (use_mmap || check_tensors) {
1521            return nullptr;
1522        }
1523        // When not using mmaped io use async uploads from pinned memory to GPU memory.
1524        // First determine if the backend supports the necessary features for async uploads.
1525        auto * buf = bufs.count(0) ? bufs.at(0) : nullptr;
1526        if (!buf) {
1527            LLAMA_LOG_DEBUG("%s: no buffer found for async uploads\n", func);
1528            return nullptr;
1529        }
1530
1531        auto * buft = ggml_backend_buffer_get_type(buf);
1532        auto * dev = ggml_backend_buft_get_device(buft);
1533        if (!dev) {
1534            LLAMA_LOG_DEBUG("%s: no device found for buffer type %s for async uploads\n", func,
1535                ggml_backend_buft_name(buft));
1536            return nullptr;
1537        }
1538
1539        if (buft != ggml_backend_dev_buffer_type(dev)) {
1540            LLAMA_LOG_DEBUG("%s: buffer type %s is not the default buffer type for device %s for async uploads\n", func,
1541                ggml_backend_buft_name(buft), ggml_backend_dev_name(dev));
1542            return nullptr;
1543        }
1544
1545        ggml_backend_dev_props props;
1546        ggml_backend_dev_get_props(dev, &props);
1547        if (!props.caps.async || !props.caps.host_buffer || !props.caps.events) {
1548            LLAMA_LOG_DEBUG("%s: device %s does not support async, host buffers or events\n", func,
1549                ggml_backend_dev_name(dev));
1550            return nullptr;
1551        }
1552
1553        auto * host_buft = ggml_backend_dev_host_buffer_type(dev);
1554        if (!host_buft) {
1555            LLAMA_LOG_DEBUG("%s: no host buffer type found for device %s\n", func,
1556                ggml_backend_dev_name(dev));
1557            return nullptr;
1558        }
1559
1560        // If the backend is supported, create pinned memory buffers and events for synchronisation.
1561        for (size_t idx = 0; idx < n_buffers; ++idx) {
1562            auto * buf = ggml_backend_buft_alloc_buffer(host_buft, buffer_size);
1563
1564            if (!buf) {
1565                LLAMA_LOG_DEBUG("%s: failed to allocate host buffer for async uploads for device %s\n", func,
1566                    ggml_backend_dev_name(dev));
1567                return nullptr;
1568            }
1569
1570            host_buffers.emplace_back(buf);
1571            host_ptrs.emplace_back(ggml_backend_buffer_get_base(buf));
1572
1573            auto * event = ggml_backend_event_new(dev);
1574            if (!event) {
1575                LLAMA_LOG_DEBUG("%s: failed to create event for async uploads for device %s\n", func,
1576                    ggml_backend_dev_name(dev));
1577                return nullptr;
1578            }
1579
1580            events.emplace_back(event);
1581        }
1582
1583        ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr);
1584        if (!backend) {
1585            LLAMA_LOG_DEBUG("%s: failed to initialize backend for device %s for async uploads\n", func,
1586                ggml_backend_dev_name(dev));
1587            return nullptr;
1588        }
1589
1590        return backend;
1591    }(__func__);
1592
1593    if (upload_backend) {
1594        LLAMA_LOG_DEBUG("%s: using async uploads for device %s, buffer type %s, backend %s\n", __func__,
1595            ggml_backend_dev_name(ggml_backend_get_device(upload_backend)),
1596            ggml_backend_buft_name(ggml_backend_buffer_get_type(bufs.at(0))),
1597            ggml_backend_name(upload_backend));
1598    }
1599
1600    std::vector<ggml_tensor *> tensors;
1601    for (struct ggml_tensor * cur = ggml_get_first_tensor(ctx); cur != NULL; cur = ggml_get_next_tensor(ctx, cur)) {
1602        tensors.push_back(cur);
1603    }
1604
1605    // without mmap, tensors in non-host buffers are staged through a temporary buffer sized like the tensor
1606    // load them biggest-first so the largest staging buffer is allocated while the fewest weights are resident
1607    if (!use_mmap) {
1608        std::stable_sort(tensors.begin(), tensors.end(), [](const ggml_tensor * a, const ggml_tensor * b) {
1609            const bool staged_a = a->buffer && !ggml_backend_buffer_is_host(a->buffer);
1610            const bool staged_b = b->buffer && !ggml_backend_buffer_is_host(b->buffer);
1611            if (staged_a != staged_b) {
1612                return staged_a;
1613            }
1614            return staged_a && ggml_nbytes(a) > ggml_nbytes(b);
1615        });
1616    }
1617
1618    for (struct ggml_tensor * cur : tensors) {
1619        const auto * weight = get_weight(ggml_get_name(cur));
1620        if (weight == nullptr) {
1621            // this can happen with split experts models
1622            continue;
1623        }
1624
1625        if (progress_callback) {
1626            if (!progress_callback((float) size_done / size_data, progress_callback_user_data)) {
1627                return false;
1628            }
1629        }
1630
1631        size_t n_size = ggml_nbytes(cur);
1632
1633        const bool from_mapping = use_mmap || lazy.has(cur);
1634
1635        if (from_mapping) {
1636            const auto & mapping = mappings.at(weight->idx);
1637            ggml_backend_buffer_t buf_mmap = nullptr;
1638            if (bufs.count(weight->idx)) {
1639                buf_mmap = bufs.at(weight->idx);
1640            }
1641            uint8_t * data = (uint8_t *) mapping->addr() + weight->offs;
1642
1643            if (check_tensors) {
1644                validation_result.emplace_back(std::async(std::launch::async, [cur, data, n_size] {
1645                    return std::make_pair(cur, ggml_validate_row_data(cur->type, data, n_size));
1646                }));
1647            }
1648
1649            GGML_ASSERT(buf_mmap || cur->data); // either we have a buffer to allocate the tensor in, or it is already allocated
1650            if (buf_mmap && cur->data == nullptr) {
1651                ggml_backend_tensor_alloc(buf_mmap, cur, data);
1652
1653                // locking a lazy tensor would fault all of it in, which is what lazy avoids
1654                if (lmlocks && !lazy.has(cur)) {
1655                    const auto & lmlock = lmlocks->at(weight->idx);
1656                    lmlock->grow_to(weight->offs + n_size);
1657                }
1658
1659                auto & mmap_used = mmaps_used[weight->idx];
1660                mmap_used.first  = std::min(mmap_used.first,  weight->offs);
1661                mmap_used.second = std::max(mmap_used.second, weight->offs + n_size);
1662            } else {
1663                ggml_backend_tensor_set(cur, data, 0, n_size);
1664            }
1665        } else {
1666            const auto & file = files.at(weight->idx);
1667
1668            if (ggml_backend_buffer_is_host(cur->buffer)) {
1669                file->seek(weight->offs, SEEK_SET);
1670                file->read_raw(cur->data, n_size);
1671                if (check_tensors) {
1672                    validation_result.emplace_back(std::async(std::launch::async, [cur, n_size] {
1673                        return std::make_pair(cur, ggml_validate_row_data(cur->type, cur->data, n_size));
1674                    }));
1675                }
1676            } else {
1677                // If upload_backend is valid load the tensor in chunks to pinned memory and upload the buffers asynchronously to the GPU.
1678                if (upload_backend) {
1679                    size_t offset = weight->offs;
1680                    alignment = file->read_alignment();
1681                    size_t aligned_offset = offset & ~(alignment - 1);
1682                    size_t offset_from_alignment = offset - aligned_offset;
1683                    file->seek(aligned_offset, SEEK_SET);
1684
1685                    // Calculate aligned read boundaries
1686                    size_t read_start = aligned_offset;
1687                    size_t read_end = (offset + n_size + alignment - 1) & ~(alignment - 1);
1688
1689                    size_t bytes_read = 0;
1690                    size_t data_read = 0;  // Actual tensor data copied (excluding padding)
1691
1692                    while (bytes_read < read_end - read_start) {
1693                        size_t read_size = std::min<size_t>(buffer_size, read_end - read_start - bytes_read);
1694
1695                        // Align the destination pointer within the pinned buffer
1696                        uintptr_t ptr_dest_aligned = (reinterpret_cast<uintptr_t>(host_ptrs[buffer_idx]) + alignment - 1) & ~(alignment - 1);
1697
1698                        // Wait for previous upload to complete before reusing buffer
1699                        ggml_backend_event_synchronize(events[buffer_idx]);
1700
1701                        // Read aligned chunk from file
1702                        file->read_raw_unsafe(reinterpret_cast<void *>(ptr_dest_aligned), read_size);
1703
1704                        // Calculate actual data portion (excluding alignment padding)
1705                        uintptr_t ptr_data = ptr_dest_aligned;
1706                        size_t data_to_copy = read_size;
1707
1708                        // Skip alignment padding at start of first chunk
1709                        if (bytes_read == 0) {
1710                            ptr_data += offset_from_alignment;
1711                            data_to_copy -= offset_from_alignment;
1712                        }
1713
1714                        // Trim alignment padding at end of last chunk
1715                        if (aligned_offset + bytes_read + read_size > offset + n_size) {
1716                            data_to_copy -= (read_end - (offset + n_size));
1717                        }
1718
1719                        // Async upload actual data to GPU
1720                        ggml_backend_tensor_set_async(upload_backend, cur,
1721                                                      reinterpret_cast<void *>(ptr_data), data_read, data_to_copy);
1722                        ggml_backend_event_record(events[buffer_idx], upload_backend);
1723
1724                        data_read += data_to_copy;
1725                        bytes_read += read_size;
1726
1727                        ++buffer_idx;
1728                        buffer_idx %= n_buffers;
1729                    }
1730                } else {
1731                    // scoped to one tensor so only one staging buffer is alive at a time
1732                    std::vector<no_init<uint8_t>> read_buf(n_size);
1733                    file->seek(weight->offs, SEEK_SET);
1734                    file->read_raw(read_buf.data(), n_size);
1735                    ggml_backend_tensor_set(cur, read_buf.data(), 0, n_size);
1736                    if (check_tensors && !ggml_validate_row_data(cur->type, read_buf.data(), n_size)) {
1737                        throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
1738                    }
1739                }
1740            }
1741        }
1742
1743        size_done += n_size;
1744    }
1745
1746    // free temporary resources used for async uploads
1747    for (auto * event : events) {
1748        ggml_backend_event_synchronize(event);
1749        ggml_backend_event_free(event);
1750    }
1751    for (auto * buf : host_buffers) {
1752        ggml_backend_buffer_free(buf);
1753    }
1754    ggml_backend_free(upload_backend);
1755
1756    // check validation results
1757    bool validation_failed = false;
1758    for (auto & future : validation_result) {
1759        auto result = future.get();
1760        if (!result.second) {
1761            LLAMA_LOG_ERROR("%s: tensor '%s' has invalid data\n", __func__, ggml_get_name(result.first));
1762            validation_failed = true;
1763        }
1764    }
1765    if (validation_failed) {
1766        throw std::runtime_error("found tensors with invalid data");
1767    }
1768
1769    // check if this is the last call and do final cleanup
1770    if (size_done >= size_data) {
1771        // unmap offloaded tensors and metadata
1772        if (use_mmap) {
1773            for (uint32_t idx = 0; idx < mappings.size(); idx++) {
1774                const auto & mmap_used = mmaps_used.at(idx);
1775                auto & mapping = mappings.at(idx);
1776                mapping->unmap_fragment(0, mmap_used.first);
1777                if (mmap_used.second != 0) {
1778                    mapping->unmap_fragment(mmap_used.second, mapping->size());
1779                }
1780            }
1781        }
1782        if (progress_callback) {
1783            // Even though the model is done loading, we still honor
1784            // cancellation since we need to free allocations.
1785            return progress_callback(1.0f, progress_callback_user_data);
1786        }
1787    }
1788
1789    return true;
1790}
1791
1792std::string llama_model_loader::ftype_name() const {
1793    return llama_ftype_name(ftype);
1794}
1795
1796void llama_model_loader::print_info() const {
1797    LLAMA_LOG_INFO("%s: file format = %s\n", __func__, llama_file_version_name(fver));
1798    LLAMA_LOG_INFO("%s: file type   = %s\n", __func__, llama_ftype_name(ftype));
1799    if (n_bytes < GiB) {
1800        LLAMA_LOG_INFO("%s: file size   = %.2f MiB (%.2f BPW) \n", __func__, n_bytes/1024.0/1024.0,        n_bytes*8.0/n_elements);
1801    } else {
1802        LLAMA_LOG_INFO("%s: file size   = %.2f GiB (%.2f BPW) \n", __func__, n_bytes/1024.0/1024.0/1024.0, n_bytes*8.0/n_elements);
1803    }
1804}