diff --git a/docs/external_quantization.md b/docs/external_quantization.md new file mode 100644 index 000000000..869a80609 --- /dev/null +++ b/docs/external_quantization.md @@ -0,0 +1,99 @@ +# Externally quantized safetensors + +Some checkpoints ship weights that were already quantized by another toolchain, +storing the payload as raw bytes plus sidecar tensors that describe how to +dequantize it. Such files are loaded directly: each weight is repacked once, at +load time, into a native ggml type, so inference uses the ordinary kernels and no +backend-specific support is required. + +| On-disk form | Repacked to | Lossless? | +| --- | --- | --- | +| ConvRot `convrot_w4a4` (int4) | `Q4_0` | yes, apart from F32→F16 on the scale | +| ComfyUI `int8_tensorwise`, ConvRot with `linear_dtype: int8` | `Q8_0` | yes, apart from F32→F16 on the scale | +| bitsandbytes NF4 | `F32` | yes (the codebook is non-uniform, so no 4-bit ggml type matches it) | + +Because NF4 decodes to F32, pass `--type` to requantize it to something compact, +e.g. `--type q4_K`. The int4/int8 paths are already compact and need no `--type`. + +## ConvRot + +[ConvRot](https://arxiv.org/abs/2512.03673) is a 4-bit scheme for diffusion +transformers. It multiplies each weight by a fixed orthogonal rotation before +quantizing, which spreads outliers across the group and makes 4-bit weights far +more accurate. + +The consequence for inference is that the stored weight is `W·H`, not `W`. Since + +``` +y = x·Wᵀ = x·(W'·H)ᵀ = (x·H)·W'ᵀ +``` + +the activation must be rotated by the same `H` before the matmul. Skipping this +does not degrade quality gracefully — it produces noise. + +`H` is a block-diagonal regular-Hadamard transform over `convrot_groupsize` input +channels (256 by default; the group size must be a power of four). It is built +from the radix-4 core + +``` +M = [[ 1, 1, 1, -1], + [ 1, 1, -1, 1], + [ 1, -1, 1, 1], + [-1, 1, 1, 1]] +``` + +Kronecker-composed `log4(groupsize)` times and scaled by `1/sqrt(groupsize)`. +That makes `H` symmetric as well as orthogonal, so no transpose bookkeeping is +needed. The rotation is applied as a single matmul against a shared constant +matrix, costing `groupsize / out_features` extra FLOPs — under 5% for a +6144-wide layer. + +Note that activations stay in F32 here, so this is W4A16 rather than the paper's +W4A4: lower throughput than a dedicated 4-bit-activation kernel, but higher +accuracy. + +A layer is only rotated when the checkpoint marks it and `in_features` divides +the group size; ConvRot leaves the remaining layers unrotated and omits the +marker, which is honoured. + +**LoRA is not supported on ConvRot layers.** LoRA deltas are trained against the +unrotated weight and cannot be mixed into a rotated activation, so adapters are +skipped on those layers with a warning rather than applied incorrectly. + +## Recognized layouts + +ComfyUI-style, read from the safetensors `__metadata__` key +`_quantization_metadata`: + +```json +{"format_version": "1.0", + "layers": {"model.diffusion_model.blocks.0.attn.wq": + {"format": "convrot_w4a4", "convrot_groupsize": 256}}} +``` + +with `.weight` holding the packed payload and `.weight_scale` +holding one F32 scale per output channel. + +bitsandbytes NF4 is detected from its own sidecars — +`.weight.quant_state.bitsandbytes__nf4` (which carries the logical shape +the packed tensor has lost), `.absmax`, `.quant_map`, and the nested pair used +when the absmax is itself double-quantized. + +Byte tensors that no configuration claims are dropped, so leftover metadata blobs +do not reach the model. + +## Example + +Krea-2 Turbo quantized to ConvRot int4: + +``` +sd-cli --diffusion-model Krea2_INT8.safetensors \ + --llm Qwen3-VL-4B-Instruct-Q4_K_M.gguf \ + --vae wan_2.1_vae.safetensors \ + -p "a lovely cat holding a sign says 'krea2.cpp'" \ + --diffusion-fa -v +``` + +At ~4.5 bits per weight the diffusion model is roughly a third the size of Q8_0. +If it still does not fit, cap what stays resident with `-ngl` (see +[backend.md](backend.md)). diff --git a/docs/krea2.md b/docs/krea2.md index 10f8eab80..6d0de5b38 100644 --- a/docs/krea2.md +++ b/docs/krea2.md @@ -10,6 +10,8 @@ Krea2 uses a Krea2 diffusion transformer, the Wan2.1 VAE, and Qwen3-VL 4B as the - Download Krea2 Turbo - safetensors: https://huggingface.co/krea/Krea-2-Turbo/tree/main - gguf: https://huggingface.co/realrebelai/KREA-2_GGUFs/tree/main/TURBO +- Krea2 also loads ConvRot int4 safetensors directly (roughly a third the size of + Q8_0); see [external_quantization.md](external_quantization.md) - Download vae - safetensors: https://huggingface.co/Comfy-Org/Wan_2.1_ComfyUI_repackaged/blob/main/split_files/vae/wan_2.1_vae.safetensors - Download Qwen3-VL 4B diff --git a/src/convert.cpp b/src/convert.cpp index 8e94a940c..382aeb37c 100644 --- a/src/convert.cpp +++ b/src/convert.cpp @@ -60,6 +60,17 @@ static bool collect_tensors_for_export(ModelLoader& model_loader, tensors.reserve(model_loader.get_tensor_storage_map().size()); for (const auto& kv : model_loader.get_tensor_storage_map()) { const TensorStorage& tensor_storage = kv.second; + // A ConvRot weight is stored in a rotated basis and is only correct when + // the runtime rotates the activation to match. No output container here + // can carry that marker, so exporting would produce a file that loads + // cleanly and generates noise. + if (tensor_storage.convrot_groupsize() > 0) { + LOG_ERROR( + "cannot convert '%s': ConvRot-rotated weights would lose their rotation marker " + "and the result would be silently wrong. Use the source model directly.", + tensor_storage.name.c_str()); + return false; + } TensorExportInfo info; info.storage = tensor_storage; info.type = get_export_tensor_type(model_loader, tensor_storage, type, tensor_type_rules); diff --git a/src/core/ggml_extend.hpp b/src/core/ggml_extend.hpp index d8e017795..397a1e60a 100644 --- a/src/core/ggml_extend.hpp +++ b/src/core/ggml_extend.hpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -3316,6 +3317,65 @@ class Identity : public UnaryBlock { } }; +// ConvRot (arXiv 2512.03673) stores W*H rather than W, where H is a block-diagonal +// rotation over `group` input channels. Recovering y = x*W^T therefore means +// rotating the activation by the same H before the matmul. +// +// H is the regular-Hadamard radix-4 butterfly: the 4x4 core below Kronecker-composed +// log4(group) times, scaled by 1/sqrt(group). That construction makes H both +// orthogonal and symmetric, so the mul_mat orientation here needs no transpose. +// `group` must be a power of four (the butterfly has no other factorization). +__STATIC_INLINE__ const std::vector& sd_convrot_matrix(int group) { + static std::map> cache; + static std::mutex cache_mutex; + std::lock_guard lock(cache_mutex); + + auto it = cache.find(group); + if (it != cache.end()) { + return it->second; + } + + static const int core[4][4] = {{1, 1, 1, -1}, {1, 1, -1, 1}, {1, -1, 1, 1}, {-1, 1, 1, 1}}; + const float norm = 1.0f / std::sqrt((float)group); + + std::vector h((size_t)group * group); + for (int i = 0; i < group; i++) { + for (int j = 0; j < group; j++) { + int sign = 1; + for (int a = i, b = j; a > 0 || b > 0; a >>= 2, b >>= 2) { + sign *= core[a & 3][b & 3]; + } + // Row-major: element (row i, col k) lands at i*group + k, which is the + // [k, i] layout ggml_mul_mat contracts over. + h[(size_t)i * group + j] = ggml_fp32_to_fp16((float)sign * norm); + } + } + return cache.emplace(group, std::move(h)).first->second; +} + +__STATIC_INLINE__ ggml_tensor* ggml_ext_convrot_rotate(GGMLRunnerContext* ctx, ggml_tensor* x, int group) { + const std::string name = "ggml_runner_build_in_tensor:convrot_h" + std::to_string(group); + + ggml_tensor* h = ggml_get_tensor(ctx->ggml_ctx, name.c_str()); + if (h == nullptr) { + h = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F16, group, group); + ggml_set_name(h, name.c_str()); + // The cache entry has static lifetime, so the pointer stays valid until the + // runner copies it into the backend buffer after allocation. + ctx->bind_backend_tensor_data(h, sd_convrot_matrix(group).data()); + } + + const int64_t ne0 = x->ne[0], ne1 = x->ne[1], ne2 = x->ne[2], ne3 = x->ne[3]; + if (!ggml_is_contiguous(x)) { + x = ggml_cont(ctx->ggml_ctx, x); + } + // Rotation is block-diagonal, so folding every group into the row dimension + // applies it to all of them with one matmul. + x = ggml_reshape_2d(ctx->ggml_ctx, x, group, ggml_nelements(x) / group); + x = ggml_mul_mat(ctx->ggml_ctx, h, x); + return ggml_reshape_4d(ctx->ggml_ctx, x, ne0, ne1, ne2, ne3); +} + class Linear : public UnaryBlock { protected: int64_t in_features; @@ -3325,12 +3385,21 @@ class Linear : public UnaryBlock { bool force_prec_f32; bool allow_weight_scale; bool has_weight_scale = false; + int convrot_groupsize = 0; float scale; std::string prefix; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { - this->prefix = prefix; - has_weight_scale = false; + this->prefix = prefix; + has_weight_scale = false; + convrot_groupsize = 0; + auto weight_storage = tensor_storage_map.find(prefix + "weight"); + if (weight_storage != tensor_storage_map.end()) { + const int group = weight_storage->second.convrot_groupsize(); + if (group > 0 && in_features % group == 0) { + convrot_groupsize = group; + } + } enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); if (in_features % ggml_blck_size(wtype) != 0 || force_f32) { wtype = GGML_TYPE_F32; @@ -3378,7 +3447,20 @@ class Linear : public UnaryBlock { } ggml_tensor* linear_bias = has_weight_scale ? nullptr : b; ggml_tensor* out = nullptr; - if (ctx->weight_adapter) { + if (convrot_groupsize > 0) { + // The stored weight lives in the rotated basis, so the activation has + // to be rotated to match. LoRA deltas are trained against the + // unrotated weight and cannot be mixed in here, so the adapter is + // deliberately bypassed rather than applied to a rotated input. + if (ctx->weight_adapter) { + static std::once_flag warned; + std::call_once(warned, []() { + LOG_WARN("LoRA is not supported on ConvRot-quantized layers; ignoring it for those layers"); + }); + } + x = ggml_ext_convrot_rotate(ctx, x, convrot_groupsize); + out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale); + } else if (ctx->weight_adapter) { WeightAdapter::ForwardParams forward_params; forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; forward_params.linear.force_prec_f32 = force_prec_f32; diff --git a/src/model_io/quant_config.cpp b/src/model_io/quant_config.cpp new file mode 100644 index 000000000..fbe1b4f58 --- /dev/null +++ b/src/model_io/quant_config.cpp @@ -0,0 +1,357 @@ +#include "model_io/quant_config.h" + +#include +#include +#include +#include +#include + +#include "core/util.h" +#include "json.hpp" + +// Reads `count` raw bytes at `offset`. Sidecars are small (a few MB in total +// for a whole model), so a per-tensor read is cheap enough and avoids holding +// the whole file. +static bool read_raw(std::ifstream& file, uint64_t offset, size_t count, void* dst) { + file.seekg((std::streamoff)offset, std::ios::beg); + if (!file) { + return false; + } + file.read((char*)dst, (std::streamsize)count); + return (bool)file; +} + +static bool read_f32_tensor(std::ifstream& file, const TensorStorage& ts, std::vector& out) { + if (ts.type != GGML_TYPE_F32) { + return false; + } + out.resize((size_t)ts.nelements()); + return read_raw(file, ts.offset, out.size() * sizeof(float), out.data()); +} + +static std::string strip_suffix(const std::string& name, const std::string& suffix) { + if (name.size() > suffix.size() && name.compare(name.size() - suffix.size(), suffix.size(), suffix) == 0) { + return name.substr(0, name.size() - suffix.size()); + } + return {}; +} + +static int json_int(const nlohmann::json& obj, const char* key, int fallback) { + if (obj.is_object()) { + auto it = obj.find(key); + if (it != obj.end() && it->is_number_integer()) { + return it->get(); + } + } + return fallback; +} + +static std::string json_str(const nlohmann::json& obj, const char* key, const std::string& fallback) { + if (obj.is_object()) { + auto it = obj.find(key); + if (it != obj.end() && it->is_string()) { + return it->get(); + } + } + return fallback; +} + +struct LayerFormat { + SDQuantPack pack = SD_QUANT_PACK_NONE; + int convrot_groupsize = 0; +}; + +// Parses ComfyUI's `_quantization_metadata` into per-layer decisions. +static bool parse_comfy_quant_metadata(const std::string& raw, + std::unordered_map& out) { + nlohmann::json root = nlohmann::json::parse(raw, nullptr, false); + if (root.is_discarded() || !root.is_object()) { + LOG_WARN("quantization metadata is not valid JSON; ignoring it"); + return false; + } + auto layers_it = root.find("layers"); + if (layers_it == root.end() || !layers_it->is_object()) { + return false; + } + const nlohmann::json& defaults = root.value("params", nlohmann::json::object()); + + for (const auto& item : layers_it->items()) { + const nlohmann::json& conf = item.value(); + if (!conf.is_object()) { + continue; + } + const nlohmann::json& params = conf.value("params", nlohmann::json::object()); + const std::string format = json_str(conf, "format", ""); + + LayerFormat lf; + if (format == "convrot_w4a4") { + // The stored payload width follows linear_dtype, which defaults to + // int4; convrot_w4a4 checkpoints may still carry int8 weights. + std::string linear_dtype = json_str(conf, "linear_dtype", json_str(params, "linear_dtype", "int4")); + lf.pack = (linear_dtype == "int8") ? SD_QUANT_PACK_INT8 : SD_QUANT_PACK_INT4; + lf.convrot_groupsize = json_int(conf, "convrot_groupsize", + json_int(params, "convrot_groupsize", + json_int(defaults, "convrot_groupsize", 256))); + } else if (format == "int8_tensorwise") { + lf.pack = SD_QUANT_PACK_INT8; + // Rotation is optional for this format and off unless flagged. + bool convrot = conf.value("convrot", params.value("convrot", false)); + if (convrot) { + lf.convrot_groupsize = json_int(conf, "convrot_groupsize", + json_int(params, "convrot_groupsize", + json_int(defaults, "convrot_groupsize", 256))); + } + } else { + LOG_WARN("unsupported quantization format '%s' for layer '%s'; leaving it unclaimed", + format.c_str(), item.key().c_str()); + continue; + } + out[item.key()] = lf; + } + return !out.empty(); +} + +bool sd_apply_quant_metadata(const std::string& file_path, + const std::map& metadata, + std::vector& tensor_storages) { + // Index by name so sidecars can be located without a quadratic scan. + std::unordered_map by_name; + by_name.reserve(tensor_storages.size() * 2); + for (size_t i = 0; i < tensor_storages.size(); i++) { + by_name[tensor_storages[i].name] = i; + } + + std::unordered_map layer_formats; + auto meta_it = metadata.find("_quantization_metadata"); + if (meta_it != metadata.end()) { + parse_comfy_quant_metadata(meta_it->second, layer_formats); + } + + // Discover bitsandbytes NF4 weights, which describe themselves through + // sidecars rather than through file-level metadata. + std::unordered_map nf4_state_blob; // weight name -> quant_state tensor + for (const auto& ts : tensor_storages) { + std::string base = strip_suffix(ts.name, ".quant_state.bitsandbytes__nf4"); + if (!base.empty()) { + nf4_state_blob[base] = ts.name; + } + } + + if (layer_formats.empty() && nf4_state_blob.empty()) { + return true; + } + + std::ifstream file(file_path, std::ios::binary); + if (!file) { + LOG_ERROR("failed to reopen '%s' to read quantization sidecars", file_path.c_str()); + return false; + } + + std::unordered_set consumed; + size_t claimed_int4 = 0, claimed_int8 = 0, claimed_nf4 = 0, rotated = 0; + + // ---- ComfyUI convrot / int8_tensorwise ------------------------------- + for (const auto& kv : layer_formats) { + const std::string weight_name = kv.first + ".weight"; + const std::string scale_name = kv.first + ".weight_scale"; + + auto w_it = by_name.find(weight_name); + if (w_it == by_name.end()) { + continue; + } + TensorStorage& w = tensor_storages[w_it->second]; + if (w.type != GGML_TYPE_I8) { + LOG_WARN("layer '%s' is marked quantized but its weight is %s; leaving it alone", + kv.first.c_str(), ggml_type_name(w.type)); + continue; + } + + auto s_it = by_name.find(scale_name); + if (s_it == by_name.end()) { + LOG_ERROR("layer '%s' is marked %s but has no weight_scale", kv.first.c_str(), + sd_quant_pack_name(kv.second.pack)); + return false; + } + + auto quant = std::make_shared(); + quant->pack = kv.second.pack; + quant->convrot_groupsize = kv.second.convrot_groupsize; + if (!read_f32_tensor(file, tensor_storages[s_it->second], quant->scales)) { + LOG_ERROR("failed to read weight_scale for layer '%s'", kv.first.c_str()); + return false; + } + + // ne is already reversed into ggml order, so ne[0] is in_features (halved + // on disk for int4) and ne[1] is out_features. + if (quant->pack == SD_QUANT_PACK_INT4) { + w.ne[0] *= 2; + } + const int64_t out_features = w.nelements() / w.ne[0]; + + if ((int64_t)quant->scales.size() == 1) { + // Tensor-wise scale: broadcast so the repack stays uniform. + quant->scales.assign((size_t)out_features, quant->scales[0]); + } else if ((int64_t)quant->scales.size() != out_features) { + LOG_ERROR("layer '%s': weight_scale has %zu entries but %" PRId64 " output channels", + kv.first.c_str(), quant->scales.size(), out_features); + return false; + } + + if (w.ne[0] % 32 != 0) { + LOG_ERROR("layer '%s': in_features %" PRId64 " is not a multiple of 32; cannot repack", + kv.first.c_str(), w.ne[0]); + return false; + } + if (quant->convrot_groupsize > 0 && w.ne[0] % quant->convrot_groupsize != 0) { + // ConvRot only rotates layers whose input divides the group size; a + // marker that disagrees with the shape means the weight is unrotated. + quant->convrot_groupsize = 0; + } + + w.type = sd_quant_pack_target_type(quant->pack); + w.quant = quant; + consumed.insert(scale_name); + + if (quant->pack == SD_QUANT_PACK_INT4) { + claimed_int4++; + } else { + claimed_int8++; + } + if (quant->convrot_groupsize > 0) { + rotated++; + } + } + + // ---- bitsandbytes NF4 ------------------------------------------------ + for (const auto& kv : nf4_state_blob) { + const std::string& weight_name = kv.first; + auto w_it = by_name.find(weight_name); + auto blob_it = by_name.find(kv.second); + if (w_it == by_name.end() || blob_it == by_name.end()) { + continue; + } + TensorStorage& w = tensor_storages[w_it->second]; + if (w.type != GGML_TYPE_I8) { + continue; + } + + // The quant_state blob is UTF-8 JSON stored as bytes; it carries the + // logical shape, which the packed weight tensor itself has lost. + const TensorStorage& blob = tensor_storages[blob_it->second]; + std::string blob_text((size_t)blob.nelements(), '\0'); + if (!read_raw(file, blob.offset, blob_text.size(), &blob_text[0])) { + LOG_ERROR("failed to read NF4 quant_state for '%s'", weight_name.c_str()); + return false; + } + nlohmann::json state = nlohmann::json::parse(blob_text, nullptr, false); + if (state.is_discarded()) { + LOG_ERROR("NF4 quant_state for '%s' is not valid JSON", weight_name.c_str()); + return false; + } + if (json_str(state, "quant_type", "nf4") != "nf4") { + LOG_WARN("'%s' uses quant_type '%s'; only nf4 is supported", weight_name.c_str(), + json_str(state, "quant_type", "").c_str()); + continue; + } + + auto quant = std::make_shared(); + quant->pack = SD_QUANT_PACK_NF4; + quant->block_size = json_int(state, "blocksize", 64); + + // Restore the logical 2-D shape in ggml order (ne[0] = in_features). + auto shape_it = state.find("shape"); + if (shape_it != state.end() && shape_it->is_array() && shape_it->size() >= 2) { + std::vector shape; + for (const auto& d : *shape_it) { + shape.push_back(d.get()); + } + for (int i = 0; i < SD_MAX_DIMS; i++) { + w.ne[i] = 1; + } + w.n_dims = (int)shape.size(); + for (size_t i = 0; i < shape.size(); i++) { + w.ne[i] = shape[shape.size() - 1 - i]; + } + } else { + LOG_ERROR("NF4 quant_state for '%s' has no shape", weight_name.c_str()); + return false; + } + + auto absmax_it = by_name.find(weight_name + ".absmax"); + if (absmax_it == by_name.end()) { + LOG_ERROR("NF4 weight '%s' has no absmax", weight_name.c_str()); + return false; + } + const TensorStorage& absmax_ts = tensor_storages[absmax_it->second]; + + if (absmax_ts.type == GGML_TYPE_F32) { + if (!read_f32_tensor(file, absmax_ts, quant->scales)) { + LOG_ERROR("failed to read NF4 absmax for '%s'", weight_name.c_str()); + return false; + } + } else { + // Double-quantized absmax: itself int8-coded against nested_quant_map + // with a nested absmax and a mean offset. + auto n_absmax_it = by_name.find(weight_name + ".nested_absmax"); + auto n_map_it = by_name.find(weight_name + ".nested_quant_map"); + if (n_absmax_it == by_name.end() || n_map_it == by_name.end()) { + LOG_ERROR("NF4 weight '%s' has a double-quantized absmax but no nested state", + weight_name.c_str()); + return false; + } + std::vector codes((size_t)absmax_ts.nelements()); + std::vector nested_absmax, nested_map; + if (!read_raw(file, absmax_ts.offset, codes.size(), codes.data()) || + !read_f32_tensor(file, tensor_storages[n_absmax_it->second], nested_absmax) || + !read_f32_tensor(file, tensor_storages[n_map_it->second], nested_map)) { + LOG_ERROR("failed to read nested NF4 state for '%s'", weight_name.c_str()); + return false; + } + const int nested_blocksize = json_int(state, "nested_blocksize", 256); + const float offset = state.value("nested_offset", 0.0f); + quant->scales.resize(codes.size()); + for (size_t i = 0; i < codes.size(); i++) { + const float na = nested_absmax.empty() ? 1.0f : nested_absmax[i / (size_t)nested_blocksize]; + quant->scales[i] = nested_map[codes[i]] * na + offset; + } + consumed.insert(weight_name + ".nested_absmax"); + consumed.insert(weight_name + ".nested_quant_map"); + } + + auto map_it = by_name.find(weight_name + ".quant_map"); + if (map_it != by_name.end()) { + read_f32_tensor(file, tensor_storages[map_it->second], quant->codebook); + consumed.insert(weight_name + ".quant_map"); + } + if (quant->codebook.size() != 16) { + quant->codebook.assign(sd_nf4_default_codebook(), sd_nf4_default_codebook() + 16); + } + + w.type = sd_quant_pack_target_type(quant->pack); + w.quant = quant; + consumed.insert(weight_name + ".absmax"); + consumed.insert(kv.second); + consumed.insert(weight_name + ".bitsandbytes__nf4"); + claimed_nf4++; + } + + // Drop consumed sidecars plus any byte tensor nothing claimed (metadata blobs, + // leftover quant state), preserving the pre-existing behaviour of not + // surfacing those to the model. + const size_t before = tensor_storages.size(); + tensor_storages.erase( + std::remove_if(tensor_storages.begin(), tensor_storages.end(), + [&](const TensorStorage& ts) { + if (consumed.count(ts.name) != 0) { + return true; + } + return ts.type == GGML_TYPE_I8 && !ts.is_packed_quant(); + }), + tensor_storages.end()); + + if (claimed_int4 + claimed_int8 + claimed_nf4 > 0) { + LOG_INFO("quantized weights: %zu int4, %zu int8, %zu nf4 (%zu convrot-rotated), %zu sidecars consumed", + claimed_int4, claimed_int8, claimed_nf4, rotated, before - tensor_storages.size()); + } + return true; +} diff --git a/src/model_io/quant_config.h b/src/model_io/quant_config.h new file mode 100644 index 000000000..a6e87909c --- /dev/null +++ b/src/model_io/quant_config.h @@ -0,0 +1,26 @@ +#ifndef __SD_MODEL_IO_QUANT_CONFIG_H__ +#define __SD_MODEL_IO_QUANT_CONFIG_H__ + +#include +#include +#include + +#include "tensor_storage.h" + +// Claims weights that were quantized by an external toolchain and repacks their +// description so the rest of the loader sees an ordinary ggml-typed tensor. +// +// Two families are recognized: +// * ComfyUI `_quantization_metadata` (formats `convrot_w4a4`, `int8_tensorwise`), +// where a `.weight_scale` sidecar holds per-output-channel scales. +// * bitsandbytes NF4, where `.weight.absmax` / `.quant_map` / +// `.quant_state.bitsandbytes__nf4` describe a packed 4-bit codebook tensor. +// +// Claimed weights get their `type`, `ne[]` and `quant` fields rewritten to the +// decoded form; consumed sidecars and unclaimed byte tensors are erased from +// `tensor_storages`. Returns false only on a malformed configuration. +bool sd_apply_quant_metadata(const std::string& file_path, + const std::map& metadata, + std::vector& tensor_storages); + +#endif // __SD_MODEL_IO_QUANT_CONFIG_H__ diff --git a/src/model_io/quant_io.cpp b/src/model_io/quant_io.cpp new file mode 100644 index 000000000..1e7f8b7f9 --- /dev/null +++ b/src/model_io/quant_io.cpp @@ -0,0 +1,135 @@ +#include "model_io/quant_io.h" + +#include + +ggml_type sd_quant_pack_target_type(SDQuantPack pack) { + switch (pack) { + case SD_QUANT_PACK_INT4: + return GGML_TYPE_Q4_0; + case SD_QUANT_PACK_INT8: + return GGML_TYPE_Q8_0; + case SD_QUANT_PACK_NF4: + return GGML_TYPE_F32; + default: + return GGML_TYPE_COUNT; + } +} + +int64_t sd_quant_pack_nbytes(SDQuantPack pack, int64_t nelements) { + switch (pack) { + case SD_QUANT_PACK_INT4: + case SD_QUANT_PACK_NF4: + return (nelements + 1) / 2; + case SD_QUANT_PACK_INT8: + return nelements; + default: + return 0; + } +} + +const char* sd_quant_pack_name(SDQuantPack pack) { + switch (pack) { + case SD_QUANT_PACK_INT4: + return "int4"; + case SD_QUANT_PACK_INT8: + return "int8"; + case SD_QUANT_PACK_NF4: + return "nf4"; + default: + return "none"; + } +} + +const float* sd_nf4_default_codebook() { + static const float codebook[16] = { + -1.0f, -0.6961928009986877f, -0.5250730514526367f, -0.39491748809814453f, + -0.28444138169288635f, -0.18477343022823334f, -0.09105003625154495f, 0.0f, + 0.07958029955625534f, 0.16093020141124725f, 0.24611230194568634f, 0.33791524171829224f, + 0.44070982933044434f, 0.5626170039176941f, 0.7229568362236023f, 1.0f}; + return codebook; +} + +// Sign-extend a 4-bit two's-complement nibble to [-8,7]. +static inline int sd_nibble_to_int4(uint8_t nibble) { + return (int)(nibble & 0x0F) - (int)((nibble & 0x08) << 1); +} + +void sd_quant_int4_to_q4_0(const void* src, const float* scales, int64_t ne0, int64_t nrows, void* dst) { + const int qk = 32; // QK4_0 + const int64_t nb_row = ne0 / qk; + const size_t block_sz = sizeof(ggml_fp16_t) + qk / 2; + const uint8_t* src_u8 = (const uint8_t*)src; + uint8_t* dst_u8 = (uint8_t*)dst; + const int64_t row_pack = ne0 / 2; + + for (int64_t r = 0; r < nrows; r++) { + const uint8_t* srow = src_u8 + r * row_pack; + // One scale per output channel; every block in the row shares it. + const ggml_fp16_t d = ggml_fp32_to_fp16(scales[r]); + + for (int64_t b = 0; b < nb_row; b++) { + uint8_t* blk = dst_u8 + (r * nb_row + b) * block_sz; + memcpy(blk, &d, sizeof(ggml_fp16_t)); + uint8_t* qs = blk + sizeof(ggml_fp16_t); + + // Source packs adjacent pairs (elements 2i, 2i+1 share a byte, earlier + // in the low nibble). Q4_0 packs element j with element j+16. Both the + // pairing and the bias differ, so this is a permutation, not a copy. + for (int j = 0; j < qk / 2; j++) { + const int64_t e0 = b * qk + j; + const int64_t e1 = b * qk + j + qk / 2; + + const uint8_t p0 = srow[e0 >> 1]; + const uint8_t p1 = srow[e1 >> 1]; + const int v0 = sd_nibble_to_int4((e0 & 1) ? (p0 >> 4) : p0); + const int v1 = sd_nibble_to_int4((e1 & 1) ? (p1 >> 4) : p1); + + // Values live in [-7,7], so +8 stays inside the [0,15] nibble range. + qs[j] = (uint8_t)((v0 + 8) & 0x0F) | (uint8_t)(((v1 + 8) & 0x0F) << 4); + } + } + } +} + +void sd_quant_int8_to_q8_0(const void* src, const float* scales, int64_t ne0, int64_t nrows, void* dst) { + const int qk = 32; // QK8_0 + const int64_t nb_row = ne0 / qk; + const size_t block_sz = sizeof(ggml_fp16_t) + qk; + const int8_t* src_i8 = (const int8_t*)src; + uint8_t* dst_u8 = (uint8_t*)dst; + + for (int64_t r = 0; r < nrows; r++) { + const int8_t* srow = src_i8 + r * ne0; + const ggml_fp16_t d = ggml_fp32_to_fp16(scales[r]); + + for (int64_t b = 0; b < nb_row; b++) { + uint8_t* blk = dst_u8 + (r * nb_row + b) * block_sz; + memcpy(blk, &d, sizeof(ggml_fp16_t)); + // Q8_0 stores quants in natural order, so this half is a straight copy. + memcpy(blk + sizeof(ggml_fp16_t), srow + b * qk, qk); + } + } +} + +void sd_quant_nf4_to_f32(const void* src, + const float* absmax, + const float* codebook, + int64_t block_size, + int64_t nelements, + float* dst) { + const uint8_t* src_u8 = (const uint8_t*)src; + if (codebook == nullptr) { + codebook = sd_nf4_default_codebook(); + } + if (block_size <= 0) { + block_size = 64; + } + + for (int64_t i = 0; i < nelements; i++) { + const uint8_t packed = src_u8[i >> 1]; + // bitsandbytes writes the earlier element into the HIGH nibble. + const uint8_t idx = (i & 1) ? (packed & 0x0F) : (packed >> 4); + const float scale = absmax != nullptr ? absmax[i / block_size] : 1.0f; + dst[i] = codebook[idx] * scale; + } +} diff --git a/src/model_io/quant_io.h b/src/model_io/quant_io.h new file mode 100644 index 000000000..cc7b32cbe --- /dev/null +++ b/src/model_io/quant_io.h @@ -0,0 +1,70 @@ +#ifndef __SD_MODEL_IO_QUANT_IO_H__ +#define __SD_MODEL_IO_QUANT_IO_H__ + +#include +#include +#include +#include + +#include "ggml.h" + +// Support for weights that arrive already quantized by an external toolchain, +// carrying their dequantization parameters in sidecar tensors rather than in the +// block layout ggml expects. Each supported pack is repacked at load time into a +// native ggml type, so no new ggml type or backend kernel is required. +enum SDQuantPack { + SD_QUANT_PACK_NONE = 0, + // Two symmetric int4 per byte, earlier element in the LOW nibble, values in + // [-7,7] (nibble 8 is unused). One F32 scale per output channel. + // Repacked losslessly into Q4_0. + SD_QUANT_PACK_INT4, + // One int8 per byte, one F32 scale per output channel. Repacked into Q8_0. + SD_QUANT_PACK_INT8, + // bitsandbytes NF4: two 4-bit codebook indices per byte, earlier element in + // the HIGH nibble (opposite of INT4 above), with a per-block absmax. The + // codebook is non-uniform so there is no exact ggml equivalent; decoded to F32 + // and left to the normal conversion path to retarget. + SD_QUANT_PACK_NF4, +}; + +struct SDQuantParams { + SDQuantPack pack = SD_QUANT_PACK_NONE; + // ConvRot rotation group size; 0 when the layer is not rotated. Non-zero + // obliges the consumer to rotate activations before the matmul, because the + // stored weight is W*H rather than W. + int convrot_groupsize = 0; + // NF4 absmax block size (bitsandbytes default 64). + int block_size = 0; + // Per-output-channel scales (INT4/INT8) or per-block absmax (NF4). + std::vector scales; + // NF4 codebook, 16 entries. + std::vector codebook; +}; + +// ggml type a pack is repacked into. +ggml_type sd_quant_pack_target_type(SDQuantPack pack); + +// On-disk byte count for `nelements` logical elements. +int64_t sd_quant_pack_nbytes(SDQuantPack pack, int64_t nelements); + +const char* sd_quant_pack_name(SDQuantPack pack); + +// The canonical bitsandbytes NF4 codebook, used when a checkpoint omits quant_map. +const float* sd_nf4_default_codebook(); + +// Repack `nrows` rows of `ne0` packed int4 values (plus one scale per row) into +// contiguous Q4_0 blocks. Requires ne0 % 32 == 0. src and dst must not overlap. +void sd_quant_int4_to_q4_0(const void* src, const float* scales, int64_t ne0, int64_t nrows, void* dst); + +// As above for int8 -> Q8_0. Requires ne0 % 32 == 0. +void sd_quant_int8_to_q8_0(const void* src, const float* scales, int64_t ne0, int64_t nrows, void* dst); + +// Decode NF4 to F32. `absmax` has one entry per `block_size` elements. +void sd_quant_nf4_to_f32(const void* src, + const float* absmax, + const float* codebook, + int64_t block_size, + int64_t nelements, + float* dst); + +#endif // __SD_MODEL_IO_QUANT_IO_H__ diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index df71eab11..c2e41af5e 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -93,6 +93,12 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) { ttype = GGML_TYPE_I32; } else if (dtype == "I64") { ttype = GGML_TYPE_I32; + } else if (dtype == "I8" || dtype == "U8") { + // Byte tensors carry externally quantized payloads (packed int4/int8/nf4) + // and their sidecars. ModelLoader claims the ones a quantization config + // describes and drops the rest, so accepting them here does not leak + // metadata blobs into the model. + ttype = GGML_TYPE_I8; } return ttype; } @@ -176,10 +182,6 @@ bool read_safetensors_file(const std::string& file_path, std::string dtype = tensor_info["dtype"]; nlohmann::json shape = tensor_info["shape"]; - if (dtype == "U8") { - continue; - } - size_t begin = tensor_info["data_offsets"][0].get(); size_t end = tensor_info["data_offsets"][1].get(); if (begin > end || end > file_size_ - data_start) { diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index 307535a53..969c4e91a 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -4,12 +4,14 @@ #include #include #include +#include #include #include #include #include #include "ggml.h" +#include "model_io/quant_io.h" #define SD_MAX_DIMS 5 @@ -24,6 +26,12 @@ struct TensorStorage { int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1}; int n_dims = 0; + // Externally quantized weights (see quant_io.h). `type` already holds the ggml + // type the tensor is repacked into, so ne[]/nbytes() describe the decoded form; + // only nbytes_to_read() sees the smaller on-disk footprint. Shared because the + // sidecar scales are resolved once, before the unordered multi-threaded load. + std::shared_ptr quant; + std::string storage_key; size_t file_index = 0; int index_in_zip = -1; // >= means stored in a zip file @@ -38,6 +46,18 @@ struct TensorStorage { } } + bool is_packed_quant() const { + return quant != nullptr && quant->pack != SD_QUANT_PACK_NONE; + } + + SDQuantPack quant_pack() const { + return quant != nullptr ? quant->pack : SD_QUANT_PACK_NONE; + } + + int convrot_groupsize() const { + return quant != nullptr ? quant->convrot_groupsize : 0; + } + int64_t nelements() const { int64_t n = 1; for (int i = 0; i < SD_MAX_DIMS; i++) { @@ -51,7 +71,9 @@ struct TensorStorage { } int64_t nbytes_to_read() const { - if (is_f8_e4m3 || is_f8_e5m2) { + if (is_packed_quant()) { + return sd_quant_pack_nbytes(quant->pack, nelements()); + } else if (is_f8_e4m3 || is_f8_e5m2) { return nbytes() / 2; } else if (is_f64 || is_i64) { return nbytes() * 2; @@ -107,6 +129,8 @@ struct TensorStorage { type_name = "f64"; } else if (is_i64) { type_name = "i64"; + } else if (is_packed_quant()) { + type_name = sd_quant_pack_name(quant->pack); } ss << name << " | " << type_name << " | "; ss << n_dims << " ["; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index ccd1e7d96..3619b134a 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -16,6 +16,7 @@ #include "core/util.h" #include "model_io/gguf_io.h" +#include "model_io/quant_config.h" #include "model_io/safetensors_io.h" #include "model_io/torch_legacy_io.h" #include "model_io/torch_zip_io.h" @@ -322,6 +323,13 @@ bool ModelLoader::init_from_safetensors_file(const std::string& file_path, const return false; } + // Must run before the prefix is applied below: quantization metadata keys are + // raw checkpoint names. The resolved SDQuantParams then rides along through + // prefixing and name conversion. + if (!sd_apply_quant_metadata(file_path, metadata_, tensor_storages)) { + return false; + } + size_t file_index = add_file_path(file_path); for (auto& tensor_storage : tensor_storages) { @@ -929,6 +937,7 @@ std::vector ModelLoader::mmap_tensors(std::maptype) { continue; } @@ -1077,6 +1086,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, std::vector read_buffer; std::vector convert_buffer; + // Staging for externally quantized payloads, which decode into a + // larger destination and so need a distinct source buffer. + std::vector packed_buffer; while (true) { int64_t t0, t1; @@ -1175,13 +1187,42 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, } } + // Decoding a packed quant grows the payload (int4 -> q4_0 is + // 0.5 -> 0.5625 bytes/weight), so it can never run in place. + if (tensor_storage.is_packed_quant()) { + packed_buffer.resize(tensor_storage.nbytes_to_read()); + read_buf = (char*)packed_buffer.data(); + } + t0 = ggml_time_ms(); read_data(read_buf, nbytes_to_read); t1 = ggml_time_ms(); read_time_ms.fetch_add(t1 - t0); t0 = ggml_time_ms(); - if (tensor_storage.is_f8_e4m3) { + if (tensor_storage.is_packed_quant()) { + const SDQuantParams& q = *tensor_storage.quant; + const int64_t ne0 = tensor_storage.ne[0]; + const int64_t nrows = tensor_storage.nelements() / ne0; + switch (q.pack) { + case SD_QUANT_PACK_INT4: + sd_quant_int4_to_q4_0(read_buf, q.scales.data(), ne0, nrows, target_buf); + break; + case SD_QUANT_PACK_INT8: + sd_quant_int8_to_q8_0(read_buf, q.scales.data(), ne0, nrows, target_buf); + break; + case SD_QUANT_PACK_NF4: + sd_quant_nf4_to_f32(read_buf, + q.scales.empty() ? nullptr : q.scales.data(), + q.codebook.empty() ? nullptr : q.codebook.data(), + q.block_size, + tensor_storage.nelements(), + (float*)target_buf); + break; + default: + break; + } + } 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());