DarthAffe/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-09-27 13:40:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: load INT8 convrot LLM embeddings correctly (#2067)

Browse Source
This commit is contained in:
leejet 2026-09-27 19:46:21 +08:00 committed by GitHub
parent 47e83d7136
commit ede32a621e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 91 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

73
src/model/common/ggml_block.hpp
View File

@ -1,6 +1,7 @@
#ifndef __SD_MODEL_COMMON_GGML_BLOCK_HPP__
#define __SD_MODEL_COMMON_GGML_BLOCK_HPP__
#include <cmath>
#include <cstdint>
#include <map>
#include <memory>
@ -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<const Embedding*>(userdata);
const int group_size = embedding->int8_convrot_group_size;
const auto* weight = dst->src[0];
const auto* input_ids = static_cast<const int32_t*>(dst->src[1]->data);
const auto* scales = static_cast<const float*>(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<float>(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<const int8_t*>(static_cast<const char*>(weight->data) + token * weight->nb[1]);
auto* out = reinterpret_cast<float*>(static_cast<char*>(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<float>(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 {
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 (!support_get_rows(wtype)) {
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,8 +424,16 @@ 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]);
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]);
auto embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids);
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]

10
src/model/te/llm.hpp
View File

@ -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<int>(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<int>(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<int>(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];
}

11
src/model_manager.cpp
View File

@ -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::vector<TensorStat
LOG_ERROR("model manager params backend is null for tensor '%s'", state->name.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<TensorState*>&
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

3
src/model_manager_prefetch.cpp
View File

@ -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;
});