From d917d7b5ffb6ea00028ae27fa5f92e311c07e79e Mon Sep 17 00:00:00 2001 From: fszontagh Date: Fri, 21 Aug 2026 13:00:56 +0200 Subject: [PATCH 1/2] fix: apply fp8 weight scale factors when loading scaled safetensors --- scripts/convert_fp8_scale_to_bf16.py | 37 +++++++++++++---- src/model_io/safetensors_io.cpp | 61 ++++++++++++++++++++++++++++ src/model_io/tensor_storage.h | 3 ++ src/model_loader.cpp | 9 ++++ 4 files changed, 103 insertions(+), 7 deletions(-) diff --git a/scripts/convert_fp8_scale_to_bf16.py b/scripts/convert_fp8_scale_to_bf16.py index a3eb2accf..1b818cf29 100644 --- a/scripts/convert_fp8_scale_to_bf16.py +++ b/scripts/convert_fp8_scale_to_bf16.py @@ -58,11 +58,29 @@ def numel(shape): return math.prod(shape) if shape else 1 -def scale_key_for_weight(name: str): +def scale_keys_for_weight(name: str): + # ".weight_scale" is the diffusers-style spelling, ".scale_weight" the ComfyUI one. + # Both ship in the wild, so accept either. + keys = [] if name.endswith(".weight"): - return name[:-len(".weight")] + ".weight_scale" + base = name[:-len(".weight")] + keys.append(base + ".weight_scale") + keys.append(base + ".scale_weight") if name.endswith("weight"): - return name + "_scale" + keys.append(name + "_scale") + return keys + + +def resolve_scale_key(name: str, entries): + for key in scale_keys_for_weight(name): + if key in entries: + return key + return None + + +def input_scale_key_for_weight(name: str): + if name.endswith(".weight"): + return name[:-len(".weight")] + ".scale_input" return None @@ -76,9 +94,14 @@ def build_output_plan(header): plan = [] for name, info in entries.items(): - scale_key = scale_key_for_weight(name) - if info["dtype"] in FP8_DTYPES and scale_key in entries: + scale_key = resolve_scale_key(name, entries) + if info["dtype"] in FP8_DTYPES and scale_key is not None: paired_scale_keys.add(scale_key) + # ".scale_input" is an activation hint with no meaning once the weight is + # materialised as BF16, so drop it alongside the weight scale. + input_key = input_scale_key_for_weight(name) + if input_key is not None and input_key in entries: + paired_scale_keys.add(input_key) for name, info in entries.items(): if name in paired_scale_keys: @@ -86,9 +109,9 @@ def build_output_plan(header): dtype = info["dtype"] shape = info["shape"] - scale_key = scale_key_for_weight(name) + scale_key = resolve_scale_key(name, entries) - if dtype in FP8_DTYPES and scale_key in entries: + if dtype in FP8_DTYPES and scale_key is not None: scale_info = entries[scale_key] plan.append( { diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index 69bcaa1ec..c68ec402e 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -1,6 +1,7 @@ #include "safetensors_io.h" #include +#include #include #include #include @@ -233,6 +234,56 @@ bool read_safetensors_file(const std::string& file_path, comfy_quant_configs.emplace(module_name, std::move(config)); } + // ComfyUI fp8_scaled checkpoints store the dequant factor in a companion + // `.scale_weight` F32 scalar instead of `.comfy_quant` metadata. Without it the + // fp8 weights load unscaled, which silently produces garbage rather than an error. + std::unordered_map fp8_scale_weights; + std::unordered_set fp8_scale_tensor_names; + for (const auto& item : header_.items()) { + const std::string& name = item.key(); + // ".scale_weight" is the ComfyUI spelling, ".weight_scale" the diffusers one; both + // ship in the wild. Only F8 weights are paired here, so int8_tensorwise checkpoints + // (which also carry ".weight_scale") keep their existing handling. + std::string suffix; + if (ends_with(name, ".scale_weight")) { + suffix = ".scale_weight"; + } else if (ends_with(name, ".weight_scale")) { + suffix = ".weight_scale"; + } else { + continue; + } + if (name == "__metadata__") { + continue; + } + const std::string module_name = name.substr(0, name.size() - suffix.size()); + auto weight_it = header_.find(module_name + ".weight"); + if (weight_it == header_.end() || weight_it.value().value("dtype", "") != "F8_E4M3") { + continue; + } + const nlohmann::json& scale_info = item.value(); + if (scale_info.value("dtype", "") != "F32") { + continue; + } + const size_t sbegin = scale_info["data_offsets"][0].get(); + const size_t send = scale_info["data_offsets"][1].get(); + if (sbegin > send || send - sbegin != sizeof(float) || send > file_size_ - data_start) { + continue; + } + float scale = 1.0f; + file.clear(); + file.seekg((std::streamoff)(data_start + sbegin), std::ios::beg); + file.read((char*)&scale, sizeof(float)); + if (!file || !std::isfinite(scale) || scale == 0.0f) { + continue; + } + fp8_scale_weights[module_name] = scale; + fp8_scale_tensor_names.insert(name); + fp8_scale_tensor_names.insert(module_name + ".scale_input"); + } + if (!fp8_scale_weights.empty()) { + LOG_DEBUG("safetensors: applying %zu fp8 scale_weight factors", fp8_scale_weights.size()); + } + tensor_storages.clear(); for (auto& item : header_.items()) { std::string name = item.key(); @@ -250,6 +301,10 @@ bool read_safetensors_file(const std::string& file_path, continue; } + if (fp8_scale_tensor_names.count(name) > 0) { + 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) { @@ -328,6 +383,12 @@ bool read_safetensors_file(const std::string& file_path, bool tensor_size_ok; if (dtype == "F8_E4M3") { tensor_storage.is_f8_e4m3 = true; + if (ends_with(name, ".weight")) { + auto scale_it = fp8_scale_weights.find(name.substr(0, name.size() - std::string(".weight").size())); + if (scale_it != fp8_scale_weights.end()) { + tensor_storage.fp8_scale = scale_it->second; + } + } // f8 -> f16 tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2); } else if (dtype == "F8_E5M2") { diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index 5672c9437..c4375379f 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -18,6 +18,9 @@ struct TensorStorage { ggml_type type = GGML_TYPE_F32; ggml_type expected_type = GGML_TYPE_COUNT; bool is_f8_e4m3 = false; + // Per-tensor dequant scale from a companion `.scale_weight` tensor (ComfyUI + // fp8_scaled checkpoints). 1.0f when the checkpoint carries no scale. + float fp8_scale = 1.0f; bool is_f8_e5m2 = false; bool is_f64 = false; bool is_i64 = false; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 891a41736..1c14510c8 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -131,6 +131,12 @@ void f8_e4m3_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) { } } +void f16_scale_vec(uint16_t* data, int64_t n, float scale) { + for (int64_t i = 0; i < n; i++) { + data[i] = ggml_fp32_to_fp16(ggml_fp16_to_fp32(data[i]) * scale); + } +} + void f8_e5m2_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) { // support inplace op for (int64_t i = n - 1; i >= 0; i--) { @@ -1217,6 +1223,9 @@ bool ModelLoader::load_tensors(on_new_tensor_cb_t on_new_tensor_cb, t0 = ggml_time_ms(); if (tensor_storage.is_f8_e4m3) { f8_e4m3_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements()); + if (tensor_storage.fp8_scale != 1.0f) { + f16_scale_vec((uint16_t*)target_buf, tensor_storage.nelements(), tensor_storage.fp8_scale); + } } else if (tensor_storage.is_f8_e5m2) { f8_e5m2_to_f16_vec((uint8_t*)read_buf, (uint16_t*)target_buf, tensor_storage.nelements()); } else if (tensor_storage.is_f64) { From 48239e49bacc951d82c9ec23ffc26d0e4a7c1fbb Mon Sep 17 00:00:00 2001 From: fszontagh Date: Sat, 22 Aug 2026 00:14:13 +0200 Subject: [PATCH 2/2] chore: trim fp8 scale comments --- scripts/convert_fp8_scale_to_bf16.py | 6 ++---- src/model_io/safetensors_io.cpp | 8 ++------ src/model_io/tensor_storage.h | 4 +--- 3 files changed, 5 insertions(+), 13 deletions(-) diff --git a/scripts/convert_fp8_scale_to_bf16.py b/scripts/convert_fp8_scale_to_bf16.py index 1b818cf29..795164bb4 100644 --- a/scripts/convert_fp8_scale_to_bf16.py +++ b/scripts/convert_fp8_scale_to_bf16.py @@ -59,8 +59,7 @@ def numel(shape): def scale_keys_for_weight(name: str): - # ".weight_scale" is the diffusers-style spelling, ".scale_weight" the ComfyUI one. - # Both ship in the wild, so accept either. + # Both spellings ship in the wild. keys = [] if name.endswith(".weight"): base = name[:-len(".weight")] @@ -97,8 +96,7 @@ def build_output_plan(header): scale_key = resolve_scale_key(name, entries) if info["dtype"] in FP8_DTYPES and scale_key is not None: paired_scale_keys.add(scale_key) - # ".scale_input" is an activation hint with no meaning once the weight is - # materialised as BF16, so drop it alongside the weight scale. + # ".scale_input" is meaningless once the weight is materialised as BF16. input_key = input_scale_key_for_weight(name) if input_key is not None and input_key in entries: paired_scale_keys.add(input_key) diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index c68ec402e..70d44c4c9 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -234,16 +234,12 @@ bool read_safetensors_file(const std::string& file_path, comfy_quant_configs.emplace(module_name, std::move(config)); } - // ComfyUI fp8_scaled checkpoints store the dequant factor in a companion - // `.scale_weight` F32 scalar instead of `.comfy_quant` metadata. Without it the - // fp8 weights load unscaled, which silently produces garbage rather than an error. std::unordered_map fp8_scale_weights; std::unordered_set fp8_scale_tensor_names; for (const auto& item : header_.items()) { const std::string& name = item.key(); - // ".scale_weight" is the ComfyUI spelling, ".weight_scale" the diffusers one; both - // ship in the wild. Only F8 weights are paired here, so int8_tensorwise checkpoints - // (which also carry ".weight_scale") keep their existing handling. + // Both spellings ship in the wild. Only F8 weights are paired, so int8_tensorwise + // checkpoints keep the existing ".weight_scale" handling below. std::string suffix; if (ends_with(name, ".scale_weight")) { suffix = ".scale_weight"; diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index c4375379f..fbb302697 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -18,9 +18,7 @@ struct TensorStorage { ggml_type type = GGML_TYPE_F32; ggml_type expected_type = GGML_TYPE_COUNT; bool is_f8_e4m3 = false; - // Per-tensor dequant scale from a companion `.scale_weight` tensor (ComfyUI - // fp8_scaled checkpoints). 1.0f when the checkpoint carries no scale. - float fp8_scale = 1.0f; + float fp8_scale = 1.0f; // companion scale tensor; 1.0f when absent bool is_f8_e5m2 = false; bool is_f64 = false; bool is_i64 = false;