diff --git a/otherarch/kcpp_sdcpp_quantized_safetensors.hpp b/otherarch/kcpp_sdcpp_quantized_safetensors.hpp new file mode 100644 index 000000000..ca99e2e18 --- /dev/null +++ b/otherarch/kcpp_sdcpp_quantized_safetensors.hpp @@ -0,0 +1,362 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "core/util.h" +#include "ggml.h" +#include "json.hpp" +#include "model_io/tensor_storage.h" + +namespace kcpp_safetensors_quant { + +struct LayerInfo { + std::string format; + bool convrot = false; + int convrot_groupsize = 0; +}; + +using LayerMap = std::unordered_map; + +struct TensorStorageExt { + bool is_i8_tensorwise = false; + bool convrot = false; + int convrot_groupsize = 0; + ggml_type scale_type = GGML_TYPE_F32; + int64_t scale_ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1}; + int scale_n_dims = 0; + uint64_t scale_offset = 0; + + int64_t scale_nelements() const { + int64_t n = 1; + for (int i = 0; i < SD_MAX_DIMS; i++) { + n *= scale_ne[i]; + } + return n; + } + + int64_t scale_nbytes() const { + return scale_nelements() * ggml_type_size(scale_type) / ggml_blck_size(scale_type); + } +}; + +inline void set_error(std::string* error, const std::string& message) { + if (error != nullptr) { + *error = message; + } +} + +inline ggml_type scale_dtype_to_ggml_type(const std::string& dtype) { + if (dtype == "F32") { + return GGML_TYPE_F32; + } + if (dtype == "F16") { + return GGML_TYPE_F16; + } + if (dtype == "BF16") { + return GGML_TYPE_BF16; + } + return GGML_TYPE_COUNT; +} + +inline LayerMap read_quantization_metadata(const nlohmann::json& header) { + LayerMap layers; + if (!header.contains("__metadata__") || !header["__metadata__"].is_object()) { + return layers; + } + + const nlohmann::json& metadata = header["__metadata__"]; + if (!metadata.contains("_quantization_metadata") || !metadata["_quantization_metadata"].is_string()) { + return layers; + } + + try { + nlohmann::json quant_metadata = nlohmann::json::parse(metadata["_quantization_metadata"].get()); + if (!quant_metadata.contains("layers") || !quant_metadata["layers"].is_object()) { + return layers; + } + + for (const auto& layer_item : quant_metadata["layers"].items()) { + if (!layer_item.value().is_object()) { + continue; + } + + LayerInfo layer; + const nlohmann::json& layer_json = layer_item.value(); + if (layer_json.contains("format") && layer_json["format"].is_string()) { + layer.format = layer_json["format"].get(); + } + if (layer_json.contains("convrot") && layer_json["convrot"].is_boolean()) { + layer.convrot = layer_json["convrot"].get(); + } + if (layer_json.contains("convrot_groupsize") && layer_json["convrot_groupsize"].is_number_integer()) { + layer.convrot_groupsize = layer_json["convrot_groupsize"].get(); + } + layers[layer_item.key()] = layer; + } + } catch (const std::exception&) { + layers.clear(); + } + return layers; +} + +inline bool should_skip_side_tensor(const std::string& name, const LayerMap& quant_layers) { + if (!ends_with(name, ".weight_scale")) { + return false; + } + + std::string base_name = name.substr(0, name.size() - std::string(".weight_scale").size()); + auto layer_it = quant_layers.find(base_name); + return layer_it != quant_layers.end() && layer_it->second.format == "int8_tensorwise"; +} + +inline bool fill_i8_tensorwise_storage(const nlohmann::json& header, + const LayerMap& quant_layers, + const std::string& name, + size_t data_start, + size_t file_size, + size_t tensor_data_size, + TensorStorage& tensor_storage, + std::string* error) { + if (!ends_with(name, ".weight")) { + set_error(error, "unsupported dtype 'I8' (tensor '" + name + "')"); + return false; + } + + std::string base_name = name.substr(0, name.size() - std::string(".weight").size()); + auto layer_it = quant_layers.find(base_name); + if (layer_it == quant_layers.end() || layer_it->second.format != "int8_tensorwise") { + set_error(error, "unsupported dtype 'I8' without int8_tensorwise metadata (tensor '" + name + "')"); + return false; + } + + std::string scale_name = base_name + ".weight_scale"; + if (!header.contains(scale_name) || !header[scale_name].is_object()) { + set_error(error, "missing weight_scale for int8 tensor '" + name + "'"); + return false; + } + + const nlohmann::json& scale_info = header[scale_name]; + std::string scale_dtype = scale_info["dtype"]; + ggml_type scale_type = scale_dtype_to_ggml_type(scale_dtype); + if (scale_type == GGML_TYPE_COUNT) { + set_error(error, "unsupported scale dtype '" + scale_dtype + "' (tensor '" + scale_name + "')"); + return false; + } + + size_t scale_begin = scale_info["data_offsets"][0].get(); + size_t scale_end = scale_info["data_offsets"][1].get(); + if (scale_begin > scale_end || scale_end > file_size - data_start) { + set_error(error, "data offsets out of bounds for tensor '" + scale_name + "'"); + return false; + } + + nlohmann::json scale_shape = scale_info["shape"]; + if (scale_shape.size() > SD_MAX_DIMS) { + set_error(error, "invalid tensor '" + scale_name + "'"); + return false; + } + + auto ext = std::make_shared(); + ext->scale_n_dims = (int)scale_shape.size(); + for (int i = 0; i < SD_MAX_DIMS; i++) { + ext->scale_ne[i] = 1; + } + for (int i = 0; i < ext->scale_n_dims; i++) { + ext->scale_ne[i] = scale_shape[i].get(); + } + if (ext->scale_n_dims == 0) { + ext->scale_n_dims = 1; + } + ext->scale_type = scale_type; + + size_t scale_data_size = scale_end - scale_begin; + size_t expected_scale_size = ext->scale_nbytes(); + if (scale_data_size != expected_scale_size) { + set_error(error, "size mismatch for tensor '" + scale_name + "' (" + scale_dtype + ")"); + return false; + } + + ext->is_i8_tensorwise = true; + ext->convrot = layer_it->second.convrot; + ext->convrot_groupsize = layer_it->second.convrot_groupsize; + ext->scale_offset = data_start + scale_begin; + tensor_storage.kcpp_ext = ext; + if (tensor_storage.nbytes_to_read() != (int64_t)tensor_data_size) { + set_error(error, "size mismatch for tensor '" + name + "' (I8)"); + return false; + } + return true; +} + +inline bool build_convrot_hadamard_signs(int size, std::vector& signs) { + if (size < 1 || (size & (size - 1)) != 0) { + return false; + } + + int tmp = size; + bool power_of_4 = true; + while (tmp > 1) { + if (tmp % 4 != 0) { + power_of_4 = false; + break; + } + tmp /= 4; + } + + if (!power_of_4) { + signs.assign(1, 1); + int current_size = 1; + while (current_size < size) { + std::vector next((size_t)current_size * 2 * current_size * 2); + for (int r = 0; r < current_size; r++) { + for (int c = 0; c < current_size; c++) { + int8_t v = signs[(size_t)r * current_size + c]; + next[(size_t)r * current_size * 2 + c] = v; + next[(size_t)r * current_size * 2 + c + current_size] = v; + next[(size_t)(r + current_size) * current_size * 2 + c] = v; + next[(size_t)(r + current_size) * current_size * 2 + c + current_size] = -v; + } + } + signs.swap(next); + current_size *= 2; + } + return true; + } + + static const int8_t H4[16] = { + 1, 1, 1, -1, + 1, 1, -1, 1, + 1, -1, 1, 1, + -1, 1, 1, 1, + }; + signs.assign(H4, H4 + 16); + int current_size = 4; + while (current_size < size) { + std::vector next((size_t)current_size * 4 * current_size * 4); + int next_size = current_size * 4; + for (int r = 0; r < current_size; r++) { + for (int c = 0; c < current_size; c++) { + int8_t v = signs[(size_t)r * current_size + c]; + for (int br = 0; br < 4; br++) { + for (int bc = 0; bc < 4; bc++) { + next[(size_t)(r * 4 + br) * next_size + (c * 4 + bc)] = v * H4[br * 4 + bc]; + } + } + } + } + signs.swap(next); + current_size = next_size; + } + return true; +} + +inline float read_quant_scale(const void* scale, ggml_type scale_type, int64_t index) { + if (scale_type == GGML_TYPE_F32) { + return ((const float*)scale)[index]; + } + if (scale_type == GGML_TYPE_F16) { + return ggml_fp16_to_fp32(((const ggml_fp16_t*)scale)[index]); + } + if (scale_type == GGML_TYPE_BF16) { + return ggml_bf16_to_fp32(((const ggml_bf16_t*)scale)[index]); + } + return 1.0f; +} + +inline bool i8_tensorwise_to_f16_vec(const int8_t* src, + const void* scale, + ggml_type scale_type, + int64_t scale_elements, + ggml_fp16_t* dst, + int64_t rows, + int64_t cols, + bool convrot, + int convrot_groupsize) { + if (scale_elements != 1 && scale_elements != rows) { + return false; + } + + if (!convrot) { + for (int64_t row = 0; row < rows; row++) { + float row_scale = read_quant_scale(scale, scale_type, scale_elements == 1 ? 0 : row); + for (int64_t col = 0; col < cols; col++) { + dst[row * cols + col] = ggml_fp32_to_fp16((float)src[row * cols + col] * row_scale); + } + } + return true; + } + + if (convrot_groupsize <= 0 || cols % convrot_groupsize != 0) { + return false; + } + + std::vector hadamard; + if (!build_convrot_hadamard_signs(convrot_groupsize, hadamard)) { + return false; + } + + const float norm = 1.0f / std::sqrt((float)convrot_groupsize); + std::vector dequant_group(convrot_groupsize); + for (int64_t row = 0; row < rows; row++) { + float row_scale = read_quant_scale(scale, scale_type, scale_elements == 1 ? 0 : row); + for (int64_t group_start = 0; group_start < cols; group_start += convrot_groupsize) { + for (int k = 0; k < convrot_groupsize; k++) { + dequant_group[k] = (float)src[row * cols + group_start + k] * row_scale; + } + for (int j = 0; j < convrot_groupsize; j++) { + float sum = 0.0f; + for (int k = 0; k < convrot_groupsize; k++) { + sum += dequant_group[k] * (float)hadamard[(size_t)k * convrot_groupsize + j]; + } + dst[row * cols + group_start + j] = ggml_fp32_to_fp16(sum * norm); + } + } + } + return true; +} + +template +inline bool load_i8_tensorwise_to_f16(const TensorStorage& tensor_storage, + const char* read_buf, + char* target_buf, + std::vector& scale_buffer, + ReadDataFn&& read_data, + FailedFlag& failed, + uint64_t& scale_nbytes) { + if (!tensor_storage.kcpp_ext) { + return false; + } + + const auto& ext = *tensor_storage.kcpp_ext; + int64_t scale_elements = 1; + for (int i = 0; i < SD_MAX_DIMS; i++) { + scale_elements *= ext.scale_ne[i]; + } + + scale_nbytes = (uint64_t)ext.scale_nbytes(); + scale_buffer.resize((size_t)scale_nbytes); + read_data((char*)scale_buffer.data(), (size_t)scale_nbytes, ext.scale_offset); + if (failed) { + return false; + } + + int64_t cols = tensor_storage.ne[0]; + int64_t rows = tensor_storage.nelements() / cols; + return i8_tensorwise_to_f16_vec((const int8_t*)read_buf, + scale_buffer.data(), + ext.scale_type, + scale_elements, + (ggml_fp16_t*)target_buf, + rows, + cols, + ext.convrot, + ext.convrot_groupsize); +} + +} // namespace kcpp_safetensors_quant diff --git a/otherarch/sdcpp/src/model_io/safetensors_io.cpp b/otherarch/sdcpp/src/model_io/safetensors_io.cpp index 197dc5f54..84770dc04 100644 --- a/otherarch/sdcpp/src/model_io/safetensors_io.cpp +++ b/otherarch/sdcpp/src/model_io/safetensors_io.cpp @@ -11,6 +11,7 @@ #include "binary_io.h" #include "core/util.h" #include "json.hpp" +#include "kcpp_sdcpp_quantized_safetensors.hpp" static constexpr size_t ST_HEADER_SIZE_LEN = 8; @@ -77,6 +78,8 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) { ttype = GGML_TYPE_F16; } else if (dtype == "F8_E5M2") { ttype = GGML_TYPE_F16; + } else if (dtype == "I8") { + ttype = GGML_TYPE_F16; } else if (dtype == "I32") { ttype = GGML_TYPE_I32; } else if (dtype == "I64") { @@ -139,6 +142,7 @@ bool read_safetensors_file(const std::string& file_path, } tensor_storages.clear(); + auto quant_layers = kcpp_safetensors_quant::read_quantization_metadata(header_); for (auto& item : header_.items()) { std::string name = item.key(); nlohmann::json tensor_info = item.value(); @@ -148,6 +152,10 @@ bool read_safetensors_file(const std::string& file_path, continue; } + if (kcpp_safetensors_quant::should_skip_side_tensor(name, quant_layers)) { + continue; + } + std::string dtype = tensor_info["dtype"]; nlohmann::json shape = tensor_info["shape"]; @@ -198,7 +206,19 @@ bool read_safetensors_file(const std::string& file_path, size_t tensor_data_size = end - begin; bool tensor_size_ok; - if (dtype == "F8_E4M3") { + if (dtype == "I8") { + if (!kcpp_safetensors_quant::fill_i8_tensorwise_storage(header_, + quant_layers, + name, + data_start, + file_size_, + tensor_data_size, + tensor_storage, + error)) { + return false; + } + tensor_size_ok = true; + } else if (dtype == "F8_E4M3") { tensor_storage.is_f8_e4m3 = true; // f8 -> f16 tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2); diff --git a/otherarch/sdcpp/src/model_io/tensor_storage.h b/otherarch/sdcpp/src/model_io/tensor_storage.h index 307535a53..e75e65323 100644 --- a/otherarch/sdcpp/src/model_io/tensor_storage.h +++ b/otherarch/sdcpp/src/model_io/tensor_storage.h @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -13,6 +14,10 @@ #define SD_MAX_DIMS 5 +namespace kcpp_safetensors_quant { +struct TensorStorageExt; +} + struct TensorStorage { std::string name; ggml_type type = GGML_TYPE_F32; @@ -23,6 +28,7 @@ struct TensorStorage { bool is_i64 = false; int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1}; int n_dims = 0; + std::shared_ptr kcpp_ext; std::string storage_key; size_t file_index = 0; @@ -53,6 +59,8 @@ struct TensorStorage { int64_t nbytes_to_read() const { if (is_f8_e4m3 || is_f8_e5m2) { return nbytes() / 2; + } else if (kcpp_ext) { + return nelements(); } else if (is_f64 || is_i64) { return nbytes() * 2; } else { @@ -107,6 +115,8 @@ struct TensorStorage { type_name = "f64"; } else if (is_i64) { type_name = "i64"; + } else if (kcpp_ext) { + type_name = "kcpp_quant"; } ss << name << " | " << type_name << " | "; ss << n_dims << " ["; diff --git a/otherarch/sdcpp/src/model_loader.cpp b/otherarch/sdcpp/src/model_loader.cpp index 6aa9afa8d..e1739acbd 100644 --- a/otherarch/sdcpp/src/model_loader.cpp +++ b/otherarch/sdcpp/src/model_loader.cpp @@ -16,6 +16,7 @@ #include "core/util.h" #include "model_io/gguf_io.h" +#include "kcpp_sdcpp_quantized_safetensors.hpp" #include "model_io/safetensors_io.h" #include "model_io/torch_legacy_io.h" #include "model_io/torch_zip_io.h" @@ -925,6 +926,7 @@ std::vector ModelLoader::mmap_tensors(std::maptype) { continue; } @@ -1073,6 +1075,8 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, std::vector read_buffer; std::vector convert_buffer; + std::vector dequant_buffer; + std::vector scale_buffer; while (true) { int64_t t0, t1; @@ -1111,28 +1115,28 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, size_t nbytes_to_read = tensor_storage.nbytes_to_read(); - auto read_data = [&](char* buf, size_t n) { + auto read_data_at = [&](char* buf, size_t n, uint64_t offset) { if (zip != nullptr) { zip_entry_openbyindex(zip, tensor_storage.index_in_zip); size_t entry_size = zip_entry_size(zip); - if (entry_size != n) { + if (entry_size != n || offset != 0) { int64_t t_memcpy_start; read_buffer.resize(entry_size); zip_entry_noallocread(zip, (void*)read_buffer.data(), entry_size); t_memcpy_start = ggml_time_ms(); - memcpy((void*)buf, (void*)(read_buffer.data() + tensor_storage.offset), n); + memcpy((void*)buf, (void*)(read_buffer.data() + offset), n); memcpy_time_ms.fetch_add(ggml_time_ms() - t_memcpy_start); } else { zip_entry_noallocread(zip, (void*)buf, n); } zip_entry_close(zip); } else if (mmapped) { - if (!mmapped->copy_data(buf, n, tensor_storage.offset)) { + if (!mmapped->copy_data(buf, n, offset)) { LOG_ERROR("read tensor data failed: '%s'", file_path.c_str()); failed = true; } } else { - file.seekg(tensor_storage.offset); + file.seekg(offset); file.read(buf, n); if (!file) { LOG_ERROR("read tensor data failed: '%s'", file_path.c_str()); @@ -1140,6 +1144,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, } } }; + auto read_data = [&](char* buf, size_t n) { + read_data_at(buf, n, tensor_storage.offset); + }; char* read_buf = nullptr; char* target_buf = nullptr; @@ -1147,7 +1154,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, if (dst_tensor->buffer == nullptr || ggml_backend_buffer_is_host(dst_tensor->buffer)) { if (tensor_storage.type == dst_tensor->type) { GGML_ASSERT(ggml_nbytes(dst_tensor) == tensor_storage.nbytes()); - if (tensor_storage.is_f64 || tensor_storage.is_i64) { + if (tensor_storage.is_f64 || tensor_storage.is_i64 || tensor_storage.kcpp_ext) { read_buffer.resize(tensor_storage.nbytes_to_read()); read_buf = (char*)read_buffer.data(); } else { @@ -1159,11 +1166,19 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, read_buf = (char*)read_buffer.data(); target_buf = read_buf; convert_buf = (char*)dst_tensor->data; + if (tensor_storage.kcpp_ext) { + dequant_buffer.resize(tensor_storage.nbytes()); + target_buf = (char*)dequant_buffer.data(); + } } } else { read_buffer.resize(std::max(tensor_storage.nbytes(), tensor_storage.nbytes_to_read())); read_buf = (char*)read_buffer.data(); target_buf = read_buf; + if (tensor_storage.kcpp_ext) { + dequant_buffer.resize(tensor_storage.nbytes()); + target_buf = (char*)dequant_buffer.data(); + } if (tensor_storage.type != dst_tensor->type) { convert_buffer.resize(ggml_nbytes(dst_tensor)); @@ -1177,7 +1192,21 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, read_time_ms.fetch_add(t1 - t0); t0 = ggml_time_ms(); - if (tensor_storage.is_f8_e4m3) { + if (tensor_storage.kcpp_ext) { + uint64_t scale_nbytes = 0; + if (!kcpp_safetensors_quant::load_i8_tensorwise_to_f16(tensor_storage, + read_buf, + target_buf, + scale_buffer, + read_data_at, + failed, + scale_nbytes)) { + LOG_ERROR("int8 tensorwise dequantization failed: '%s'", tensor_storage.name.c_str()); + failed = true; + return; + } + bytes_processed.fetch_add(scale_nbytes); + } else if (tensor_storage.is_f8_e4m3) { f8_e4m3_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements()); } else if (tensor_storage.is_f8_e5m2) { f8_e5m2_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements()); @@ -1202,7 +1231,7 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, (int)tensor_storage.ne[0], std::move(imatrix)); } else { - convert_buf = read_buf; + convert_buf = target_buf; } t1 = ggml_time_ms(); convert_time_ms.fetch_add(t1 - t0);