char-fork/llama.cpp

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

Xuan-Son Nguyenllama: improve TENSOR_READ_LAZY handling (#27837)257813839

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
9.3 KiB265 linesraw
1#pragma once
2
3#include "llama.h"
4
5#include "llama-impl.h"
6#include "llama-arch.h"
7#include "llama-hparams.h"
8#include "llama-mmap.h"
9
10#include "ggml-cpp.h"
11
12#include <cstddef>
13#include <cstring>
14#include <map>
15#include <set>
16#include <stdexcept>
17#include <unordered_map>
18
19using llama_buf_map = std::unordered_map<uint32_t, ggml_backend_buffer_t>;
20
21// lists of buffer types used for each layer
22using buft_list_t = std::vector<std::pair<ggml_backend_dev_t, ggml_backend_buffer_type_t>>;
23
24enum llama_fver {
25    GGUF_FILE_VERSION_V1 = 1,
26    GGUF_FILE_VERSION_V2 = 2,
27    GGUF_FILE_VERSION_V3 = 3,
28};
29
30const char * llama_file_version_name(llama_fver version);
31
32struct llama_model_loader {
33    // Holds information on a model weight
34    struct llama_tensor_weight {
35        uint16_t  idx; // source file index
36        size_t   offs; // tensor data offset in the original file
37
38        ggml_tensor * tensor;
39
40        llama_tensor_weight(const llama_file * file, uint16_t idx, const struct gguf_context * gguf_ctx, ggml_tensor * tensor) : idx(idx), tensor(tensor) {
41            const int tensor_idx = gguf_find_tensor(gguf_ctx,  ggml_get_name(tensor));
42            if (tensor_idx < 0) {
43                throw std::runtime_error(format("tensor '%s' not found in the model", ggml_get_name(tensor)));
44            }
45
46            offs = gguf_get_data_offset(gguf_ctx) + gguf_get_tensor_offset(gguf_ctx, tensor_idx);
47            if (offs + ggml_nbytes(tensor) < offs || offs + ggml_nbytes(tensor) > file->size()) {
48                throw std::runtime_error(format("tensor '%s' data is not within the file bounds, model is corrupted or incomplete", ggml_get_name(tensor)));
49            }
50        }
51    };
52
53    // custom comparator to sort weights more nicely by layer
54    struct weight_name_comparer {
55        bool operator()(const std::string & a, const std::string & b) const {
56            int a_layer = -1;
57            int b_layer = -1;
58            sscanf(a.c_str(), "blk.%d.", &a_layer);
59            sscanf(b.c_str(), "blk.%d.", &b_layer);
60            if (a_layer != b_layer) {
61                return a_layer < b_layer;
62            }
63            return a < b;
64        }
65    };
66
67    static const int TENSOR_NOT_REQUIRED    = 1 << 0;
68    static const int TENSOR_DUPLICATED      = 1 << 1;
69    static const int TENSOR_SKIP            = 1 << 2;
70    static const int TENSOR_SKIP_IF_VIRTUAL = 1 << 3;
71    static const int TENSOR_ALLOW_RESHAPE   = 1 << 4;
72    static const int TENSOR_READ_LAZY       = 1 << 5; // read rows on demand instead of loading whole tensor; requires mmap for now
73
74    int n_kv      = 0;
75    int n_tensors = 0;
76    int n_created = 0;
77
78    uint64_t n_elements = 0;
79    size_t   n_bytes    = 0;
80
81    bool use_mmap = false;
82    bool use_direct_io = false;
83    bool check_tensors;
84    bool no_alloc;
85    bool load_mtp;
86
87    // handle TENSOR_READ_LAZY
88    // use case: keep PLE / engrams embd tensors on disk, read them on demand
89    struct lazy_read {
90        // set by the caller before the create_tensor() calls
91        enum llama_lazy_mode mode = LLAMA_LAZY_MODE_OFF;
92
93        // decide whether this tensor is read lazily
94        // pass w to also record it, or nullptr to only ask
95        bool add(const std::string & name, const ggml_tensor * t, const llama_tensor_weight * w);
96
97        bool any() const {
98            return !ranges.empty();
99        }
100
101        bool has(const ggml_tensor * t) const {
102            return tensors.count(ggml_get_name(t)) > 0;
103        }
104
105        const llama_mmap::ranges & for_file(uint32_t idx) const {
106            static const llama_mmap::ranges none;
107
108            const auto it = ranges.find(idx);
109            return it == ranges.end() ? none : it->second;
110        }
111
112        // lazy tensors are gathered on the host, so no offload setting applies to them
113        static ggml_backend_buffer_type_t buft();
114
115    private:
116        std::map<uint32_t, llama_mmap::ranges> ranges;
117        std::set<std::string>                  tensors;
118    } lazy;
119
120    llama_files files;
121    llama_ftype ftype;
122    llama_fver  fver;
123
124    llama_mmaps mappings;
125
126    std::map<std::string, llama_tensor_weight, weight_name_comparer> weights_map;
127    std::unordered_map<std::string, llama_model_kv_override> kv_overrides;
128    const llama_model_tensor_buft_override * tensor_buft_overrides;
129
130    gguf_context_ptr metadata_ptr;
131    struct gguf_context * metadata; // either metadata_ptr.get() or externally set
132    llama_model_set_tensor_data_t set_tensor_data;
133    void * set_tensor_data_ud;
134    std::vector<ggml_context_ptr> contexts;
135
136    std::string arch_name;
137    LLM_KV      llm_kv    = LLM_KV(LLM_ARCH_UNKNOWN);
138
139    size_t size_done = 0;
140    size_t size_data = 0;
141    std::vector<std::pair<size_t, size_t>> mmaps_used;
142
143    // define a comparator for the buft -> ctx map to ensure that the order is well-defined:
144    struct ggml_backend_buft_comparator {
145        bool operator()(const ggml_backend_buffer_type_t & lhs, const ggml_backend_buffer_type_t & rhs) const {
146            return strcmp(ggml_backend_buft_name(lhs), ggml_backend_buft_name(rhs)) < 0;
147        }
148    };
149
150    // lazy tensors need dedicated context
151    struct ctx_key {
152        ggml_backend_buffer_type_t buft;
153        bool lazy;
154    };
155
156    struct ctx_key_comparator {
157        bool operator()(const ctx_key & lhs, const ctx_key & rhs) const {
158            if (lhs.lazy != rhs.lazy) {
159                return lhs.lazy < rhs.lazy;
160            }
161            return strcmp(ggml_backend_buft_name(lhs.buft), ggml_backend_buft_name(rhs.buft)) < 0;
162        }
163    };
164
165    std::map<ctx_key, ggml_context_ptr, ctx_key_comparator> ctx_map;
166
167    // track tensors that had to be moved for debugging:
168    size_t n_tensors_moved = 0;
169    std::string first_tensor_moved_name;
170    std::string first_tensor_moved_type_name;
171    ggml_backend_buffer_type_t first_moved_from_buft = nullptr;
172    ggml_backend_buffer_type_t first_moved_to_buft = nullptr;
173
174    llama_model_loader(
175        struct gguf_context * metadata,
176        llama_model_set_tensor_data_t set_tensor_data,
177        void * set_tensor_data_ud,
178        const std::string & fname,
179        std::vector<std::string> & splits, // optional, only need if the split does not follow naming scheme
180        FILE * file,
181        llama_load_mode load_mode,
182        bool check_tensors,
183        bool no_alloc,
184        bool load_mtp,
185        const llama_model_kv_override * param_overrides_p,
186        const llama_model_tensor_buft_override * param_tensor_buft_overrides_p);
187
188    template<typename T>
189    typename std::enable_if<std::is_integral<T>::value, bool>::type
190    get_arr_n(const std::string & key, T & result, bool required = true);
191
192    template<typename T>
193    typename std::enable_if<std::is_integral<T>::value, bool>::type
194    get_arr_n(enum llm_kv kid, T & result, bool required = true);
195
196    template<typename T>
197    bool get_arr(const std::string & key, std::vector<T> & result, bool required = true);
198
199    template<typename T, size_t N_MAX>
200    bool get_arr(const std::string & key, std::array<T, N_MAX> & result, bool required = true);
201
202    template<typename T>
203    bool get_arr(enum llm_kv kid, T & result, bool required = true);
204
205    template<typename T>
206    bool get_key(const std::string & key, T & result, bool required = true);
207
208    template<typename T>
209    bool get_key(enum llm_kv kid, T & result, bool required = true);
210
211    template<typename T, size_t N_MAX>
212    bool get_key_or_arr(const std::string & key, std::array<T, N_MAX> & result, uint32_t n, bool required = true);
213
214    template<typename T>
215    bool get_key_or_arr(enum llm_kv kid, T & result, uint32_t n, bool required = true);
216
217    bool get_key_or_arr(enum llm_kv kid, uint32_t & result, bool required = true);
218
219    std::string get_arch_name() const;
220
221    enum llm_arch get_arch() const;
222
223    const llama_tensor_weight * get_weight(const char * name) const;
224
225    const llama_tensor_weight & require_weight(const char * name) const;
226
227    struct ggml_tensor * get_tensor_meta(const char * name) const;
228
229    struct ggml_tensor * require_tensor_meta(const std::string & name) const;
230
231    const struct ggml_tensor * check_tensor_dims(
232            const std::string & name,
233            const std::vector<int64_t> & ne,
234            bool required,
235            bool allow_reshape) const;
236
237    struct ggml_tensor * create_tensor(
238        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,
239        const buft_list_t * buft_list_layer, const LLM_TN_IMPL & tn, const std::initializer_list<int64_t> & ne, int flags);
240
241    void done_getting_tensors(bool partial = false) const;
242
243    void init_mappings(bool prefetch = true, llama_mlocks * mlock_mmaps = nullptr);
244
245    void get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const;
246
247    // release a weight's mmap pages
248    void unmap_weight(const llama_tensor_weight & w) const;
249
250    // read a byte range of a weight's data
251    // with mmap, returns a pointer into the mapping, otherwise reads into buf and returns buf
252    const void * load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const;
253
254    // Returns false if cancelled by progress_callback
255    bool load_all_data(
256            struct ggml_context * ctx,
257            llama_buf_map & bufs,
258            llama_mlocks * lmlocks,
259            llama_progress_callback progress_callback,
260            void * progress_callback_user_data);
261
262    std::string ftype_name() const;
263
264    void print_info() const;
265};