From 1aaa4929463bda819ddb4d60dbae8450830e9b9a Mon Sep 17 00:00:00 2001 From: mosaic-devel <290483437+mosaic-devel@users.noreply.github.com> Date: Wed, 29 Jul 2026 12:08:43 +0200 Subject: [PATCH] feat: load ConvRot int4 and per-channel int8 safetensors Some checkpoints ship weights already quantized by an external toolchain, storing the payload as raw bytes plus sidecar tensors that describe how to dequantize it. These currently fail to load at all: I8 is rejected as an unsupported dtype and U8 tensors are silently dropped. Add a load-time repack into native ggml types, so inference uses the ordinary kernels and no backend-specific support is needed: ConvRot convrot_w4a4 (int4) -> Q4_0 int8_tensorwise / int8 -> Q8_0 bitsandbytes NF4 -> F32 Both integer paths are exact: the int4 values occupy [-7,7] so the +8 bias lands inside Q4_0's nibble range, and the per-output-channel scale becomes the block scale. The only loss is F32->F16 on that scale (~5e-4 relative, against ~15% int4 quantization error). ConvRot additionally stores W*H rather than W, where H is a block-diagonal regular-Hadamard rotation over convrot_groupsize input channels, so the activation must be rotated to match: y = x*W^T = x*(W'*H)^T = (x*H)*W'^T H is the radix-4 butterfly core Kronecker-composed log4(groupsize) times and scaled by 1/sqrt(groupsize), which makes it symmetric as well as orthogonal, so no transpose bookkeeping is needed. It is applied as one matmul against a shared constant, costing groupsize/out_features extra FLOPs - under 5% at 6144 wide. Activations stay F32, so this is W4A16: slower than a dedicated 4-bit activation kernel but more accurate. LoRA is skipped on rotated layers with a warning, since adapter deltas are trained against the unrotated weight. Conversion refuses rotated weights outright, because no output container carries the rotation marker and the result would load cleanly and generate noise. --- docs/external_quantization.md | 99 +++++++++ docs/krea2.md | 2 + src/convert.cpp | 11 + src/core/ggml_extend.hpp | 88 +++++++- src/model_io/quant_config.cpp | 357 ++++++++++++++++++++++++++++++++ src/model_io/quant_config.h | 26 +++ src/model_io/quant_io.cpp | 135 ++++++++++++ src/model_io/quant_io.h | 70 +++++++ src/model_io/safetensors_io.cpp | 10 +- src/model_io/tensor_storage.h | 26 ++- src/model_loader.cpp | 43 +++- 11 files changed, 858 insertions(+), 9 deletions(-) create mode 100644 docs/external_quantization.md create mode 100644 src/model_io/quant_config.cpp create mode 100644 src/model_io/quant_config.h create mode 100644 src/model_io/quant_io.cpp create mode 100644 src/model_io/quant_io.h 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());