diff --git a/docs/imatrix.md b/docs/imatrix.md new file mode 100644 index 000000000..df877a00a --- /dev/null +++ b/docs/imatrix.md @@ -0,0 +1,59 @@ +# Importance Matrix (imatrix) Quantization + +## What is an Importance Matrix? + +Quantization reduces the precision of a model's weights, decreasing its size and computational requirements. However, this can lead to a loss of quality. An importance matrix helps mitigate this by identifying which weights are *most* important for the model's performance. During quantization, these important weights are preserved with higher precision, while less important weights are quantized more aggressively. This allows for better overall quality at a given quantization level. + +This originates from work done with language models in [llama.cpp](https://github.com/ggml-org/llama.cpp/blob/master/tools/imatrix/README.md). + +## Usage + +The imatrix feature involves two main steps: *training* the matrix and *using* it during quantization. + +### Training the Importance Matrix + +To generate an imatrix, run stable-diffusion.cpp with the `--imat-out` flag, specifying the output filename. This process runs alongside normal image generation. + +```bash +sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat +``` + +* **`[same exact parameters as normal generation]`**: Use the same command-line arguments you would normally use for image generation (e.g., prompt, dimensions, sampling method, etc.). +* **`--imat-out imatrix.dat`**: Specifies the output file for the generated imatrix. + +You can generate multiple images at once using the `-b` flag to speed up the training process. + +### Continuing Training an Existing Matrix + +If you want to refine an existing imatrix, use the `--imat-in` flag *in addition* to `--imat-out`. This will load the existing matrix and continue training it. + +```bash +sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat --imat-in imatrix.dat +``` +With that, you can train and refine the imatrix while generating images like you'd normally do. + +### Using Multiple Matrices + +You can load and merge multiple imatrices together: + +```bash +sd.exe [same exact parameters as normal generation] --imat-out imatrix.dat --imat-in imatrix.dat --imat-in imatrix2.dat +``` + +### Quantizing with an Importance Matrix + +To quantize a model using a trained imatrix, use the `-M convert` option (or equivalent quantization command) and the `--imat-in` flag, specifying the imatrix file. + +```bash +sd.exe -M convert [same exact parameters as normal quantization] --imat-in imatrix.dat +``` + +* **`[same exact parameters as normal quantization]`**: Use the same command-line arguments you would normally use for quantization (e.g., target quantization method, input/output filenames). +* **`--imat-in imatrix.dat`**: Specifies the imatrix file to use during quantization. You can specify multiple `--imat-in` flags to combine multiple matrices. + +## Important Considerations + +* The quality of the imatrix depends on the prompts and settings used during training. Use prompts and settings representative of the types of images you intend to generate for the best results. +* Experiment with different training parameters (e.g., number of images, prompt variations) to optimize the imatrix for your specific use case. +* The performance impact of training an imatrix during image generation or using an imatrix for quantization is negligible. +* Using already quantized models to train the imatrix seems to be working fine. \ No newline at end of file diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 2ea300173..f08970624 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include @@ -53,6 +54,9 @@ struct SDCliParams { bool metadata_brief = false; bool metadata_all = false; + std::string imatrix_out; + std::vector imatrix_in; + bool normal_exit = false; ArgOptions get_options() { @@ -79,6 +83,11 @@ struct SDCliParams { "path to write preview image to (default: ./preview.png). Multi-frame previews support .avi, .webm, and animated .webp", 0, &preview_path}, + {"", + "--imat-out", + "compute the imatrix for this run and save it to the provided path", + 0, + &imatrix_out}, }; options.int_options = { @@ -179,6 +188,14 @@ struct SDCliParams { return -1; }; + auto on_imatrix_in_arg = [&](int argc, const char** argv, int index) { + if (++index >= argc) { + return -1; + } + imatrix_in.push_back(argv[index]); + return 1; + }; + options.manual_options = { {"-M", "--mode", @@ -192,6 +209,10 @@ struct SDCliParams { "--help", "show this help message and exit", on_help_arg}, + {"", + "--imat-in", + "load an imatrix file for quantization or continued collection; can be specified multiple times", + on_imatrix_in_arg}, }; return options; @@ -253,6 +274,7 @@ struct SDCliParams { << " preview_fps: " << preview_fps << ",\n" << " taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n" << " preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n" + << " imatrix_out: \"" << imatrix_out << "\",\n" << " metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n" << " metadata_brief: " << (metadata_brief ? "true" : "false") << ",\n" << " metadata_all: " << (metadata_all ? "true" : "false") << "\n" @@ -605,13 +627,32 @@ int main(int argc, const char* argv[]) { LOG_DEBUG("%s", ctx_params.to_string().c_str()); LOG_DEBUG("%s", gen_params.to_string().c_str()); + if (!cli_params.imatrix_out.empty()) { + if (fs::exists(cli_params.imatrix_out) && + std::find(cli_params.imatrix_in.begin(), cli_params.imatrix_in.end(), cli_params.imatrix_out) == cli_params.imatrix_in.end()) { + LOG_WARN("imatrix file '%s' already exists and will be overwritten", cli_params.imatrix_out.c_str()); + } + enable_imatrix_collection(); + } + + for (const auto& in_file : cli_params.imatrix_in) { + LOG_INFO("loading imatrix from '%s'", in_file.c_str()); + if (!load_imatrix(in_file.c_str())) { + LOG_WARN("failed to load imatrix from '%s'", in_file.c_str()); + } + } + if (cli_params.mode == CONVERT) { - bool success = convert(ctx_params.model_path.c_str(), - ctx_params.vae_path.c_str(), - cli_params.output_path.c_str(), - ctx_params.wtype, - ctx_params.tensor_type_rules.c_str(), - cli_params.convert_name); + bool success = convert_with_components(ctx_params.model_path.c_str(), + ctx_params.clip_l_path.c_str(), + ctx_params.clip_g_path.c_str(), + ctx_params.t5xxl_path.c_str(), + ctx_params.diffusion_model_path.c_str(), + ctx_params.vae_path.c_str(), + cli_params.output_path.c_str(), + ctx_params.wtype, + ctx_params.tensor_type_rules.c_str(), + cli_params.convert_name); if (!success) { LOG_ERROR("convert '%s'/'%s' to '%s' failed", ctx_params.model_path.c_str(), @@ -833,6 +874,11 @@ int main(int argc, const char* argv[]) { return 1; } + if (!cli_params.imatrix_out.empty()) { + LOG_INFO("saving imatrix to '%s'", cli_params.imatrix_out.c_str()); + save_imatrix(cli_params.imatrix_out.c_str()); + } + free_sd_audio(generated_audio); return 0; diff --git a/examples/common/common.cpp b/examples/common/common.cpp index b7ba577e5..524b89404 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -710,7 +710,18 @@ bool SDContextParams::resolve(SDMode mode) { } bool SDContextParams::validate(SDMode mode) { - if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) { + if (mode == CONVERT) { + const bool has_convert_input = model_path.length() != 0 || + clip_l_path.length() != 0 || + clip_g_path.length() != 0 || + t5xxl_path.length() != 0 || + diffusion_model_path.length() != 0 || + vae_path.length() != 0; + if (!has_convert_input) { + LOG_ERROR("error: convert mode needs at least one model input path\n"); + return false; + } + } else if (mode != UPSCALE && mode != METADATA && model_path.length() == 0 && diffusion_model_path.length() == 0) { LOG_ERROR("error: the following arguments are required: model_path/diffusion_model\n"); return false; } diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index d5bc8a40a..9d813b2a6 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -407,14 +407,17 @@ typedef struct { } sd_vid_gen_params_t; typedef struct sd_ctx_t sd_ctx_t; +struct ggml_tensor; typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data); typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data); typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data); +typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* user_data); SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data); SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data); SD_API void sd_set_preview_callback(sd_preview_cb_t cb, enum preview_t mode, int interval, bool denoised, bool noisy, void* data); +SD_API void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data); SD_API int32_t sd_get_num_physical_cores(); SD_API const char* sd_get_system_info(); SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx); @@ -503,6 +506,17 @@ SD_API bool convert(const char* input_path, const char* tensor_type_rules, bool convert_name); +SD_API bool convert_with_components(const char* model_path, + const char* clip_l_path, + const char* clip_g_path, + const char* t5xxl_path, + const char* diffusion_model_path, + const char* vae_path, + const char* output_path, + enum sd_type_t output_type, + const char* tensor_type_rules, + bool convert_name); + SD_API bool preprocess_canny(sd_image_t image, float high_threshold, float low_threshold, @@ -510,6 +524,11 @@ SD_API bool preprocess_canny(sd_image_t image, float strong, bool inverse); +SD_API bool load_imatrix(const char* imatrix_path); +SD_API void save_imatrix(const char* imatrix_path); +SD_API void enable_imatrix_collection(void); +SD_API void disable_imatrix_collection(void); + SD_API const char* sd_commit(void); SD_API const char* sd_version(void); diff --git a/src/convert.cpp b/src/convert.cpp index 27d377ec0..0b7fe2cfb 100644 --- a/src/convert.cpp +++ b/src/convert.cpp @@ -76,29 +76,22 @@ static bool load_tensors_for_export(ModelLoader& model_loader, return success; } -bool convert(const char* input_path, - const char* vae_path, - const char* output_path, - sd_type_t output_type, - const char* tensor_type_rules, - bool convert_name) { - ModelLoader model_loader; - - if (!model_loader.init_from_file(input_path)) { - LOG_ERROR("init model loader from file failed: '%s'", input_path); - return false; - } - - if (vae_path != nullptr && strlen(vae_path) > 0) { - if (!model_loader.init_from_file(vae_path, "vae.")) { - LOG_ERROR("init model loader from file failed: '%s'", vae_path); - return false; - } +static bool init_convert_path(ModelLoader& model_loader, const char* path, const char* prefix, bool& loaded_any) { + if (path == nullptr || strlen(path) == 0) { + return true; } - if (convert_name) { - model_loader.convert_tensors_name(); + if (!model_loader.init_from_file(path, prefix)) { + LOG_ERROR("init model loader from file failed: '%s'", path); + return false; } + loaded_any = true; + return true; +} +static bool export_loaded_model(ModelLoader& model_loader, + const char* output_path, + sd_type_t output_type, + const char* tensor_type_rules) { ggml_type type = sd_type_to_ggml_type(output_type); bool output_is_safetensors = ends_with(output_path, ".safetensors"); TensorTypeRules type_rules = parse_tensor_type_rules(tensor_type_rules); @@ -136,3 +129,55 @@ bool convert(const char* input_path, ggml_free(ggml_ctx); return success; } + +bool convert_with_components(const char* model_path, + const char* clip_l_path, + const char* clip_g_path, + const char* t5xxl_path, + const char* diffusion_model_path, + const char* vae_path, + const char* output_path, + sd_type_t output_type, + const char* tensor_type_rules, + bool convert_name) { + ModelLoader model_loader; + bool loaded_any = false; + + if (!init_convert_path(model_loader, model_path, "", loaded_any) || + !init_convert_path(model_loader, clip_l_path, "text_encoders.clip_l.transformer.", loaded_any) || + !init_convert_path(model_loader, clip_g_path, "text_encoders.clip_g.transformer.", loaded_any) || + !init_convert_path(model_loader, t5xxl_path, "text_encoders.t5xxl.transformer.", loaded_any) || + !init_convert_path(model_loader, diffusion_model_path, "model.diffusion_model.", loaded_any) || + !init_convert_path(model_loader, vae_path, "vae.", loaded_any)) { + return false; + } + + if (!loaded_any) { + LOG_ERROR("no input model path provided for convert"); + return false; + } + + if (convert_name) { + model_loader.convert_tensors_name(); + } + + return export_loaded_model(model_loader, output_path, output_type, tensor_type_rules); +} + +bool convert(const char* input_path, + const char* vae_path, + const char* output_path, + sd_type_t output_type, + const char* tensor_type_rules, + bool convert_name) { + return convert_with_components(input_path, + nullptr, + nullptr, + nullptr, + nullptr, + vae_path, + output_path, + output_type, + tensor_type_rules, + convert_name); +} diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index 65196813b..ec9aee408 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -2473,7 +2473,10 @@ struct GGMLRunner { sd_backend_cpu_set_n_threads(runtime_backend, n_threads); } - ggml_status status = ggml_backend_graph_compute(runtime_backend, gf); + ggml_status status = sd_backend_graph_compute_with_eval_callback(runtime_backend, + gf, + sd_get_backend_eval_callback(), + sd_get_backend_eval_callback_data()); if (status != GGML_STATUS_SUCCESS) { LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); return std::nullopt; diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index f3e2cceba..f2899d441 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -9,6 +9,7 @@ #include #include "core/util.h" +#include "ggml/src/ggml-impl.h" #include "stable-diffusion.h" static std::string trim_copy(const std::string& value) { @@ -364,6 +365,68 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) { return false; } +static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) { + ggml_cgraph cgraph = { + /*.size =*/0, + /*.n_nodes =*/i1 - i0, + /*.n_leafs =*/0, + /*.nodes =*/cgraph0->nodes + i0, + /*.grads =*/nullptr, + /*.grad_accs =*/nullptr, + /*.leafs =*/nullptr, + /*.use_counts =*/cgraph0->use_counts, + /*.visited_hash_set =*/cgraph0->visited_hash_set, + /*.order =*/cgraph0->order, + /*.uid =*/0, + }; + return cgraph; +} + +ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, + ggml_cgraph* gf, + sd_graph_eval_callback_t callback_eval, + void* callback_eval_user_data) { + if (callback_eval == nullptr) { + return ggml_backend_graph_compute(backend, gf); + } + + ggml_status status = GGML_STATUS_SUCCESS; + const int n_nodes = ggml_graph_n_nodes(gf); + bool stopped = false; + + for (int j0 = 0; j0 < n_nodes; ++j0) { + ggml_tensor* t = ggml_graph_node(gf, j0); + bool need = callback_eval(t, true, callback_eval_user_data); + int j1 = j0; + + while (!need && j1 < n_nodes - 1) { + t = ggml_graph_node(gf, ++j1); + need = callback_eval(t, true, callback_eval_user_data); + } + + ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1); + status = ggml_backend_graph_compute_async(backend, &gv); + if (status != GGML_STATUS_SUCCESS) { + break; + } + + ggml_backend_synchronize(backend); + + if (need && !callback_eval(t, false, callback_eval_user_data)) { + stopped = true; + break; + } + + j0 = j1; + } + + ggml_backend_synchronize(backend); + if (stopped && status == GGML_STATUS_SUCCESS) { + status = GGML_STATUS_ABORTED; + } + return status; +} + const char* sd_get_system_info() { static std::string cache_info = []() -> std::string { ggml_backend_load_all_once(); diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index 9aecf97c0..19b71d432 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -9,6 +9,7 @@ #include "ggml-backend.h" #include "ggml.h" +#include "stable-diffusion.h" enum class SDBackendModule { DIFFUSION, @@ -71,6 +72,10 @@ bool sd_backend_is(ggml_backend_t backend, const std::string& name); bool sd_backend_is_cpu(ggml_backend_t backend); ggml_backend_t sd_backend_cpu_init(); bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads); +ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, + ggml_cgraph* gf, + sd_graph_eval_callback_t callback_eval, + void* callback_eval_user_data); std::string sd_backend_resolve_name(const std::string& name); const char* sd_backend_module_name(SDBackendModule module); void ggml_ext_im_set_f32_1d(const struct ggml_tensor* tensor, int i, float value); diff --git a/src/core/util.cpp b/src/core/util.cpp index b844d29ee..6d2479f9f 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -346,6 +346,9 @@ int sd_preview_interval = 1; bool sd_preview_denoised = true; bool sd_preview_noisy = false; +static sd_graph_eval_callback_t sd_backend_eval_cb = nullptr; +static void* sd_backend_eval_cb_data = nullptr; + std::u32string utf8_to_utf32(const std::string& utf8_str) { std::wstring_convert, char32_t> converter; return converter.from_bytes(utf8_str); @@ -629,6 +632,11 @@ void sd_set_preview_callback(sd_preview_cb_t cb, preview_t mode, int interval, b sd_preview_noisy = noisy; } +void sd_set_backend_eval_callback(sd_graph_eval_callback_t cb, void* data) { + sd_backend_eval_cb = cb; + sd_backend_eval_cb_data = data; +} + sd_preview_cb_t sd_get_preview_callback() { return sd_preview_cb; } @@ -649,6 +657,14 @@ bool sd_should_preview_noisy() { return sd_preview_noisy; } +sd_graph_eval_callback_t sd_get_backend_eval_callback() { + return sd_backend_eval_cb; +} + +void* sd_get_backend_eval_callback_data() { + return sd_backend_eval_cb_data; +} + sd_progress_cb_t sd_get_progress_callback() { return sd_progress_cb; } diff --git a/src/core/util.h b/src/core/util.h index 8213dfabd..35b52061b 100644 --- a/src/core/util.h +++ b/src/core/util.h @@ -98,6 +98,9 @@ int sd_get_preview_interval(); bool sd_should_preview_denoised(); bool sd_should_preview_noisy(); +sd_graph_eval_callback_t sd_get_backend_eval_callback(); +void* sd_get_backend_eval_callback_data(); + // test if the backend is a specific one, e.g. "CUDA", "ROCm", "Vulkan" etc. bool sd_backend_is(ggml_backend_t backend, const std::string& name); diff --git a/src/model_loader.cpp b/src/model_loader.cpp index c239e22d2..dca9cbb86 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -20,6 +20,7 @@ #include "model_io/torch_legacy_io.h" #include "model_io/torch_zip_io.h" #include "model_loader.h" +#include "runtime/imatrix.h" #include "stable-diffusion.h" #include "core/ggml_extend_backend.h" @@ -156,7 +157,8 @@ void convert_tensor(void* src, void* dst, ggml_type dst_type, int nrows, - int n_per_row) { + int n_per_row, + std::vector imatrix = {}) { int n = nrows * n_per_row; if (src_type == dst_type) { size_t nbytes = n * ggml_type_size(src_type) / ggml_blck_size(src_type); @@ -165,7 +167,7 @@ void convert_tensor(void* src, if (dst_type == GGML_TYPE_F16) { ggml_fp32_to_fp16_row((float*)src, (ggml_fp16_t*)dst, n); } else { - std::vector imatrix(n_per_row, 1.0f); // dummy importance matrix + imatrix.resize(n_per_row, 1.0f); const float* im = imatrix.data(); ggml_quantize_chunk(dst_type, (float*)src, dst, 0, nrows, n_per_row, im); } @@ -195,7 +197,7 @@ void convert_tensor(void* src, if (dst_type == GGML_TYPE_F16) { ggml_fp32_to_fp16_row((float*)src_data_f32, (ggml_fp16_t*)dst, n); } else { - std::vector imatrix(n_per_row, 1.0f); // dummy importance matrix + imatrix.resize(n_per_row, 1.0f); const float* im = imatrix.data(); ggml_quantize_chunk(dst_type, (float*)src_data_f32, dst, 0, nrows, n_per_row, im); } @@ -970,6 +972,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, size_t total_tensors_processed = 0; const int64_t t_start = start_time; int last_n_threads = 1; + SDVersion imatrix_version = (version_ == VERSION_COUNT) ? get_sd_version() : version_; for (size_t file_index = 0; file_index < file_data.size(); ++file_index) { auto& fdata = file_data[file_index]; @@ -1154,12 +1157,15 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, failed = true; return; } + std::string processed_name = convert_tensor_name(tensor_storage.name, imatrix_version); + std::vector imatrix = get_imatrix_collector().get_values(processed_name); convert_tensor((void*)target_buf, tensor_storage.type, convert_buf, dst_tensor->type, (int)tensor_storage.nelements() / (int)tensor_storage.ne[0], - (int)tensor_storage.ne[0]); + (int)tensor_storage.ne[0], + std::move(imatrix)); } else { convert_buf = read_buf; } diff --git a/src/runtime/imatrix.cpp b/src/runtime/imatrix.cpp new file mode 100644 index 000000000..313eadc68 --- /dev/null +++ b/src/runtime/imatrix.cpp @@ -0,0 +1,308 @@ +#include "runtime/imatrix.h" + +/* Adapted from llama.cpp (credits: Kawrakow). */ + +#include "core/util.h" +#include "ggml-backend.h" +#include "ggml.h" +#include "stable-diffusion.h" + +#include +#include +#include + +static IMatrixCollector imatrix_collector; + +IMatrixCollector& get_imatrix_collector() { + return imatrix_collector; +} + +// remove any prefix and suffixes from the name +// CUDA0#blk.0.attn_k.weight#0 => blk.0.attn_k.weight +static std::string filter_tensor_name(const char* name) { + std::string wname; + const char* p = strchr(name, '#'); + if (p != NULL) { + p = p + 1; + const char* q = strchr(p, '#'); + if (q != NULL) { + wname = std::string(p, q - p); + } else { + wname = p; + } + } else { + wname = name; + } + return wname; +} + +bool IMatrixCollector::collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) { + GGML_UNUSED(user_data); + if (t == nullptr) { + return false; + } + if (t->op != GGML_OP_MUL_MAT && t->op != GGML_OP_MUL_MAT_ID) { + return false; + } + + const struct ggml_tensor* src0 = t->src[0]; + const struct ggml_tensor* src1 = t->src[1]; + if (src0 == nullptr || src1 == nullptr) { + return false; + } + std::string wname = filter_tensor_name(src0->name); + + // when ask is true, the scheduler wants to know if we are interested in data from this tensor + // if we return true, a follow-up call will be made with ask=false in which we can do the actual collection + if (ask) { + if (t->op == GGML_OP_MUL_MAT_ID) { + return true; // collect all indirect matrix multiplications + } + // why are small batches ignored (<16 tokens)? + // if (src1->ne[1] < 16 || src1->type != GGML_TYPE_F32) return false; + if (!(wname.substr(0, 6) == "model." || wname.substr(0, 17) == "cond_stage_model." || wname.substr(0, 14) == "text_encoders.")) { + return false; + } + return true; + } + std::lock_guard lock(mutex_); + + // copy the data from the GPU memory if needed + const bool is_host = src1->buffer == NULL || ggml_backend_buffer_is_host(src1->buffer); + + if (!is_host) { + src1_data_.resize(ggml_nelements(src1)); + ggml_backend_tensor_get(src1, src1_data_.data(), 0, ggml_nbytes(src1)); + } + + const float* data = is_host ? (const float*)src1->data : src1_data_.data(); + + // this has been adapted to the new format of storing merged experts in a single 3d tensor + // ref: https://github.com/ggml-org/llama.cpp/pull/6387 + if (t->op == GGML_OP_MUL_MAT_ID) { + // ids -> [n_experts_used, n_tokens] + // src1 -> [cols, n_expert_used, n_tokens] + const ggml_tensor* ids = t->src[2]; + const int n_as = static_cast(src0->ne[2]); + const int n_ids = static_cast(ids->ne[0]); + + // the top-k selected expert ids are stored in the ids tensor + // for simplicity, always copy ids to host, because it is small + // take into account that ids is not contiguous! + + GGML_ASSERT(ids->ne[1] == src1->ne[2]); + + ids_.resize(ggml_nbytes(ids)); + ggml_backend_tensor_get(ids, ids_.data(), 0, ggml_nbytes(ids)); + + auto& e = stats_[wname]; + + ++e.ncall; + + if (e.values.empty()) { + e.values.resize(src1->ne[0] * n_as, 0); + e.counts.resize(src1->ne[0] * n_as, 0); + } else if (e.values.size() != (size_t)src1->ne[0] * n_as) { + LOG_ERROR("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0] * n_as); + exit(1); // GGML_ABORT("fatal error"); + } + // loop over all possible experts, regardless if they are used or not in the batch + for (int ex = 0; ex < n_as; ++ex) { + size_t e_start = ex * src1->ne[0]; + + for (int idx = 0; idx < n_ids; ++idx) { + for (int row = 0; row < (int)src1->ne[2]; ++row) { + const int excur = *(const int32_t*)(ids_.data() + row * ids->nb[1] + idx * ids->nb[0]); + + GGML_ASSERT(excur >= 0 && excur < n_as); // sanity check + + if (excur != ex) + continue; + + const int64_t i11 = idx % src1->ne[1]; + const int64_t i12 = row; + const float* x = (const float*)((const char*)data + i11 * src1->nb[1] + i12 * src1->nb[2]); + + for (int j = 0; j < (int)src1->ne[0]; ++j) { + e.values[e_start + j] += x[j] * x[j]; + e.counts[e_start + j]++; + if (!std::isfinite(e.values[e_start + j])) { + LOG_ERROR("%f detected in %s\n", e.values[e_start + j], wname.c_str()); + exit(1); + } + } + } + } + } + } else { + auto& e = stats_[wname]; + if (e.values.empty()) { + e.values.resize(src1->ne[0], 0); + e.counts.resize(src1->ne[0], 0); + } else if (e.values.size() != (size_t)src1->ne[0]) { + LOG_WARN("inconsistent size for %s (%d vs %d)\n", wname.c_str(), (int)e.values.size(), (int)src1->ne[0]); + exit(1); // GGML_ABORT("fatal error"); + } + + ++e.ncall; + for (int row = 0; row < (int)src1->ne[1]; ++row) { + const float* x = data + row * src1->ne[0]; + for (int j = 0; j < (int)src1->ne[0]; ++j) { + if (std::isfinite(x[j])) { + e.values[j] += x[j] * x[j]; + e.counts[j]++; + if (!std::isfinite(e.values[j])) { + LOG_WARN("%f detected in %s\n", e.values[j], wname.c_str()); + exit(1); + } + } else { + // Likely something from an attention mask? + } + } + } + } + return true; +} + +bool load_imatrix(const char* imatrix_path) { + return imatrix_collector.load_imatrix(imatrix_path); +} + +void save_imatrix(const char* imatrix_path) { + imatrix_collector.save_imatrix(imatrix_path); +} + +static bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data) { + return imatrix_collector.collect_imatrix(t, ask, user_data); +} + +void enable_imatrix_collection() { + sd_set_backend_eval_callback(collect_imatrix, nullptr); +} + +void disable_imatrix_collection() { + sd_set_backend_eval_callback(nullptr, nullptr); +} + +void IMatrixCollector::save_imatrix(std::string fname, int ncall) const { + if (ncall > 0) { + fname += ".at_"; + fname += std::to_string(ncall); + } + // avoid writing imatrix entries that do not have full data + // this can happen with MoE models where some of the experts end up not being exercised by the provided training data + + int n_entries = 0; + std::vector to_store; + + for (const auto& kv : stats_) { + const int n_all = static_cast(kv.second.counts.size()); + + if (n_all == 0) { + continue; + } + + int n_zeros = 0; + for (const int c : kv.second.counts) { + if (c == 0) { + n_zeros++; + } + } + + if (n_zeros == n_all) { + LOG_WARN("entry '%40s' has no data - skipping\n", kv.first.c_str()); + continue; + } + + if (n_zeros > 0) { + LOG_WARN("entry '%40s' has partial data (%.2f%%) - skipping\n", kv.first.c_str(), 100.0f * (n_all - n_zeros) / n_all); + continue; + } + + n_entries++; + to_store.push_back(kv.first); + } + + if (to_store.size() < stats_.size()) { + LOG_WARN("storing only %zu out of %zu entries\n", to_store.size(), stats_.size()); + } + + std::ofstream out(fname, std::ios::binary); + out.write((const char*)&n_entries, sizeof(n_entries)); + for (const auto& name : to_store) { + const auto& stat = stats_.at(name); + int len = static_cast(name.size()); + out.write((const char*)&len, sizeof(len)); + out.write(name.c_str(), len); + out.write((const char*)&stat.ncall, sizeof(stat.ncall)); + int nval = static_cast(stat.values.size()); + out.write((const char*)&nval, sizeof(nval)); + if (nval > 0) { + std::vector tmp(nval); + for (int i = 0; i < nval; i++) { + tmp[i] = (stat.values[i] / static_cast(stat.counts[i])) * static_cast(stat.ncall); + } + out.write((const char*)tmp.data(), nval * sizeof(float)); + } + } + + // Write the number of call the matrix was computed with + out.write((const char*)&last_call_, sizeof(last_call_)); +} + +bool IMatrixCollector::load_imatrix(const char* fname) { + std::ifstream in(fname, std::ios::binary); + if (!in) { + LOG_ERROR("failed to open %s\n", fname); + return false; + } + int n_entries; + in.read((char*)&n_entries, sizeof(n_entries)); + if (in.fail() || n_entries < 1) { + LOG_ERROR("no data in file %s\n", fname); + return false; + } + for (int i = 0; i < n_entries; ++i) { + int len; + in.read((char*)&len, sizeof(len)); + std::vector name_as_vec(len + 1); + in.read((char*)name_as_vec.data(), len); + if (in.fail()) { + LOG_ERROR("failed reading name for entry %d from %s\n", i + 1, fname); + return false; + } + name_as_vec[len] = 0; + std::string name{name_as_vec.data()}; + auto& e = stats_[std::move(name)]; + int ncall; + in.read((char*)&ncall, sizeof(ncall)); + int nval; + in.read((char*)&nval, sizeof(nval)); + if (in.fail() || nval < 1) { + LOG_ERROR("failed reading number of values for entry %d\n", i); + stats_ = {}; + return false; + } + + if (e.values.empty()) { + e.values.resize(nval, 0); + e.counts.resize(nval, 0); + } + + std::vector tmp(nval); + in.read((char*)tmp.data(), nval * sizeof(float)); + if (in.fail()) { + LOG_ERROR("failed reading data for entry %d\n", i); + stats_ = {}; + return false; + } + + // Recreate the state as expected by save_imatrix(), and correct for weighted sum. + for (int i = 0; i < nval; i++) { + e.values[i] += tmp[i]; + e.counts[i] += ncall; + } + e.ncall += ncall; + } + return true; +} diff --git a/src/runtime/imatrix.h b/src/runtime/imatrix.h new file mode 100644 index 000000000..341c18c8b --- /dev/null +++ b/src/runtime/imatrix.h @@ -0,0 +1,45 @@ +#ifndef __SD_RUNTIME_IMATRIX_H__ +#define __SD_RUNTIME_IMATRIX_H__ + +#include +#include +#include +#include +#include + +/* Adapted from llama.cpp (credits: Kawrakow). */ + +struct ggml_tensor; + +struct IMatrixStats { + std::vector values{}; + std::vector counts{}; + int ncall = 0; +}; + +class IMatrixCollector { +private: + std::unordered_map stats_ = {}; + std::mutex mutex_; + int last_call_ = 0; + std::vector src1_data_; + std::vector ids_; // the expert ids from ggml_mul_mat_id + +public: + IMatrixCollector() = default; + bool collect_imatrix(struct ggml_tensor* t, bool ask, void* user_data); + void save_imatrix(std::string fname, int ncall = -1) const; + bool load_imatrix(const char* fname); + std::vector get_values(const std::string& key) const { + auto it = stats_.find(key); + if (it != stats_.end()) { + return it->second.values; + } else { + return {}; + } + } +}; + +IMatrixCollector& get_imatrix_collector(); + +#endif // __SD_RUNTIME_IMATRIX_H__