From ede32a621edb08c22daebfe0a7021ec3a7e92408 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 27 Sep 2026 19:46:21 +0800 Subject: [PATCH] fix: load INT8 convrot LLM embeddings correctly (#2067) --- src/model/common/ggml_block.hpp | 79 ++++++++++++++++++++++++++++++--- src/model/te/llm.hpp | 10 ++--- src/model_manager.cpp | 11 ++++- src/model_manager_prefetch.cpp | 3 +- 4 files changed, 91 insertions(+), 12 deletions(-) diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 560f49f8..3567fd7b 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_COMMON_GGML_BLOCK_HPP__ #define __SD_MODEL_COMMON_GGML_BLOCK_HPP__ +#include #include #include #include @@ -335,15 +336,75 @@ class Embedding : public UnaryBlock { protected: int64_t embedding_dim; int64_t num_embeddings; + bool is_int8_tensorwise = false; + int int8_convrot_group_size = 0; + + static void get_rows_i8(ggml_tensor* dst, int ith, int nth, void* userdata) { + const auto* embedding = static_cast(userdata); + const int group_size = embedding->int8_convrot_group_size; + const auto* weight = dst->src[0]; + const auto* input_ids = static_cast(dst->src[1]->data); + const auto* scales = static_cast(dst->src[2]->data); + const bool scalar_scale = ggml_nelements(dst->src[2]) == 1; + const float normalization = group_size > 0 ? 1.f / std::sqrt(static_cast(group_size)) : 1.f; + for (int64_t row = ith; row < dst->ne[1]; row += nth) { + const int32_t token = input_ids[row]; + GGML_ASSERT(token >= 0 && token < weight->ne[1]); + const auto* src = reinterpret_cast(static_cast(weight->data) + token * weight->nb[1]); + auto* out = reinterpret_cast(static_cast(dst->data) + row * dst->nb[1]); + const float scale = scales[scalar_scale ? 0 : token] * normalization; + for (int64_t i = 0; i < weight->ne[0]; ++i) { + out[i] = static_cast(src[i]) * scale; + } + // The regular Hadamard rotation is symmetric and its own inverse. + for (int stride = 1; stride < group_size; stride *= 4) { + for (int64_t base = 0; base < weight->ne[0]; base += 4 * stride) { + for (int j = 0; j < stride; ++j) { + float* values = out + base + j; + const float a = values[0]; + const float b = values[stride]; + const float c = values[2 * stride]; + const float d = values[3 * stride]; + values[0] = a + b + c - d; + values[stride] = a + b - c + d; + values[2 * stride] = a - b + c + d; + values[3 * stride] = -a + b + c + d; + } + } + } + } + } + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { - enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); - if (!support_get_rows(wtype)) { + auto weight_storage = tensor_storage_map.find(prefix + "weight"); + is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise; + int8_convrot_group_size = 0; + enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); + if (is_int8_tensorwise) { + GGML_ASSERT(wtype == GGML_TYPE_I8); + auto scale_storage = tensor_storage_map.find(prefix + "weight_scale"); + GGML_ASSERT(scale_storage != tensor_storage_map.end()); + const int64_t scale_nelements = scale_storage->second.nelements(); + GGML_ASSERT(scale_nelements == 1 || scale_nelements == num_embeddings); + params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements); + if (weight_storage->second.int8_convrot) { + int8_convrot_group_size = weight_storage->second.int8_convrot_group_size; + int remainder = int8_convrot_group_size; + while (remainder > 1 && remainder % 4 == 0) { + remainder /= 4; + } + GGML_ASSERT(remainder == 1 && embedding_dim % int8_convrot_group_size == 0); + } + } else if (!support_get_rows(wtype)) { wtype = GGML_TYPE_F32; } params["weight"] = ggml_new_tensor_2d(ctx, wtype, embedding_dim, num_embeddings); } enum ggml_op param_usage_op(const std::string& name) const override { + if (is_int8_tensorwise) { + return GGML_OP_CUSTOM; + } return name == "weight" ? GGML_OP_GET_ROWS : GGML_OP_NONE; } @@ -363,9 +424,17 @@ public: int64_t n = input_ids->ne[1]; input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]); - input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]); - auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids); - embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n); + ggml_tensor* embedding; + if (is_int8_tensorwise) { + GGML_ASSERT(input_ids->type == GGML_TYPE_I32); + ggml_tensor* args[] = {weight, input_ids, params["weight_scale"]}; + embedding = ggml_custom_4d(ctx->ggml_ctx, GGML_TYPE_F32, embedding_dim, input_ids->ne[0], 1, 1, + args, 3, get_rows_i8, GGML_N_TASKS_MAX, this); + } else { + input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]); + embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids); + } + embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n); // [N, n_token, embedding_dim] return embedding; diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 5adebe61..0d00eedb 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -299,22 +299,22 @@ namespace LLM { if (contains(name, "attn.q_proj")) { config.llama_cpp_style = true; } - if (contains(name, "visual.patch_embed.proj.1.weight")) { + if (ends_with(name, "visual.patch_embed.proj.1.weight")) { config.vision.split_patch_embed = true; } - if (contains(name, "visual.patch_embed.proj.0.weight")) { + if (ends_with(name, "visual.patch_embed.proj.0.weight")) { config.vision.patch_size = static_cast(tensor_storage.ne[0]); config.vision.in_channels = tensor_storage.ne[2]; config.vision.hidden_size = tensor_storage.ne[3]; } // HF-format checkpoints keep the patch embed unsplit under a single name. - if (contains(name, "visual.patch_embed.proj.weight")) { + if (ends_with(name, "visual.patch_embed.proj.weight")) { config.vision.patch_size = static_cast(tensor_storage.ne[0]); } if (contains(name, "visual.patch_embed.bias") || contains(name, "visual.patch_embed.proj.bias")) { config.vision.hidden_size = tensor_storage.ne[0]; } - if (contains(name, "visual.pos_embed.weight")) { + if (ends_with(name, "visual.pos_embed.weight") && tensor_storage.n_dims == 2) { config.vision.hidden_size = tensor_storage.ne[0]; config.vision.num_position_embeddings = static_cast(tensor_storage.ne[1]); } @@ -348,7 +348,7 @@ namespace LLM { } } } - if (contains(name, "embed_tokens.weight")) { + if (ends_with(name, "embed_tokens.weight") && tensor_storage.n_dims == 2) { config.hidden_size = tensor_storage.ne[0]; config.vocab_size = tensor_storage.ne[1]; } diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 9529d6f2..e719825f 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -79,6 +79,10 @@ static bool device_supports_param_op(ggml_backend_dev_t device, if (op == GGML_OP_GET_ROWS) { ggml_tensor* indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); op_tensor = ggml_get_rows(ctx, weight, indices); + } else if (op == GGML_OP_CUSTOM) { + op_tensor = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + op_tensor->op = op; + op_tensor->src[0] = weight; } if (op_tensor == nullptr) { ggml_free(ctx); @@ -507,7 +511,9 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vectorname.c_str()); return false; } - if (state->compute_backend == state->params_backend || state->staged_to_compute_backend) { + // Custom CPU operators must retain host weights even when the runner uses a GPU. + if (state->usage_op == GGML_OP_CUSTOM || + state->compute_backend == state->params_backend || state->staged_to_compute_backend) { continue; } if (!state->loaded_to_params_backend || state->tensor == nullptr || state->tensor->data == nullptr) { @@ -1402,6 +1408,9 @@ size_t ModelManager::compute_backend_alloc_size(const std::vector& if (state == nullptr || state->tensor == nullptr) { continue; } + if (state->usage_op == GGML_OP_CUSTOM && !sd_backend_is_cpu(state->compute_backend)) { + continue; + } const bool compute_resident = state->compute_backend == state->params_backend ? state->loaded_to_params_backend diff --git a/src/model_manager_prefetch.cpp b/src/model_manager_prefetch.cpp index 1d304193..de96d5ae 100644 --- a/src/model_manager_prefetch.cpp +++ b/src/model_manager_prefetch.cpp @@ -210,7 +210,7 @@ WeightPrefetchResult ModelManager::prefetch_params( is_optional_missing_tensor(state->name)) { continue; } - if (state->compute_backend == state->params_backend) { + if (state->usage_op == GGML_OP_CUSTOM || state->compute_backend == state->params_backend) { needs_synchronous_load = needs_synchronous_load || !state->loaded_to_params_backend; continue; @@ -275,6 +275,7 @@ bool ModelManager::activate_prefetched_params( [&](TensorState* state) { return state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name) || + state->usage_op == GGML_OP_CUSTOM || state->compute_backend == state->params_backend || state->staged_to_compute_backend; });