feat: load scaled FP8 weights without upfront conversion (#1913)

This commit is contained in:
leejet
2026-08-27 20:50:47 +08:00
committed by GitHub
parent 50d6405683
commit be0e34480d
9 changed files with 73 additions and 101 deletions
+39 -17
View File
@@ -3407,7 +3407,6 @@ protected:
bool bias;
bool force_f32;
bool force_prec_f32;
bool allow_weight_scale;
bool has_weight_scale = false;
bool int8_convrot = false;
int int8_convrot_group_size = 0;
@@ -3430,8 +3429,11 @@ protected:
}
auto weight_storage = tensor_storage_map.find(prefix + "weight");
const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise;
if ((allow_weight_scale || is_int8_tensorwise) && tensor_storage_map.find(prefix + "weight_scale") != tensor_storage_map.end()) {
params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features);
auto weight_scale_storage = tensor_storage_map.find(prefix + "weight_scale");
if (weight_scale_storage != tensor_storage_map.end()) {
const int64_t scale_nelements = weight_scale_storage->second.nelements();
GGML_ASSERT(scale_nelements == 1 || scale_nelements == out_features);
params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements);
has_weight_scale = true;
}
if (is_int8_tensorwise) {
@@ -3445,17 +3447,15 @@ protected:
public:
Linear(int64_t in_features,
int64_t out_features,
bool bias = true,
bool force_f32 = false,
bool force_prec_f32 = false,
float scale = 1.f,
bool allow_weight_scale = false)
bool bias = true,
bool force_f32 = false,
bool force_prec_f32 = false,
float scale = 1.f)
: in_features(in_features),
out_features(out_features),
bias(bias),
force_f32(force_f32),
force_prec_f32(force_prec_f32),
allow_weight_scale(allow_weight_scale),
scale(scale) {}
void set_scale(float scale_) {
@@ -3467,7 +3467,11 @@ public:
}
ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override {
ggml_tensor* w = params["weight"];
ggml_tensor* w = params["weight"];
ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr;
if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) {
w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16);
}
ggml_tensor* b = nullptr;
if (bias) {
b = params["bias"];
@@ -3498,7 +3502,7 @@ public:
out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx,
x,
w,
params["weight_scale"],
weight_scale,
b,
int8_convrot ? int8_convrot_group_size : 0,
scale);
@@ -3517,6 +3521,30 @@ public:
}
return out;
}
if (has_weight_scale) {
out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale);
out = ggml_mul(ctx->ggml_ctx, out, weight_scale);
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;
forward_params.linear.scale = scale;
out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx,
ctx->backend,
x,
w,
out,
prefix,
forward_params);
if (b != nullptr) {
b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias");
}
}
if (b != nullptr) {
out = ggml_add_inplace(ctx->ggml_ctx, out, b);
}
return out;
}
if (ctx->weight_adapter) {
WeightAdapter::ForwardParams forward_params;
forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR;
@@ -3526,12 +3554,6 @@ public:
} else {
out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale);
}
if (has_weight_scale) {
out = ggml_mul(ctx->ggml_ctx, out, params["weight_scale"]);
if (b != nullptr) {
out = ggml_add_inplace(ctx->ggml_ctx, out, b);
}
}
return out;
}
};
+1 -1
View File
@@ -142,7 +142,7 @@ namespace Ideogram4 {
__STATIC_INLINE__ std::shared_ptr<Linear> make_linear(int64_t in_features,
int64_t out_features,
bool bias = true) {
return std::make_shared<Linear>(in_features, out_features, bias, false, false, 1.f, true);
return std::make_shared<Linear>(in_features, out_features, bias);
}
__STATIC_INLINE__ std::vector<float> gen_ideogram4_pe(int grid_h,
+4 -6
View File
@@ -87,9 +87,9 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) {
} else if (dtype == "F64") {
ttype = GGML_TYPE_F32;
} else if (dtype == "F8_E4M3") {
ttype = GGML_TYPE_F16;
ttype = GGML_TYPE_F8_E4M3;
} else if (dtype == "F8_E5M2") {
ttype = GGML_TYPE_F16;
ttype = GGML_TYPE_F8_E5M2;
} else if (dtype == "I32") {
ttype = GGML_TYPE_I32;
} else if (dtype == "I64") {
@@ -328,12 +328,10 @@ bool read_safetensors_file(const std::string& file_path,
bool tensor_size_ok;
if (dtype == "F8_E4M3") {
tensor_storage.is_f8_e4m3 = true;
// f8 -> f16
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
} else if (dtype == "F8_E5M2") {
tensor_storage.is_f8_e5m2 = true;
// f8 -> f16
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size * 2);
tensor_size_ok = (tensor_storage.nbytes() == tensor_data_size);
} else if (dtype == "F64") {
tensor_storage.is_f64 = true;
// f64 -> f32
+1 -3
View File
@@ -54,9 +54,7 @@ struct TensorStorage {
}
int64_t nbytes_to_read() const {
if (is_f8_e4m3 || is_f8_e5m2) {
return nbytes() / 2;
} else if (is_f64 || is_i64) {
if (is_f64 || is_i64) {
return nbytes() * 2;
} else {
return nbytes();
+2 -68
View File
@@ -78,66 +78,6 @@ bool is_unused_tensor(const std::string& name) {
return false;
}
uint16_t f8_e4m3_to_f16(uint8_t f8) {
// do we need to support uz?
const uint32_t exponent_bias = 7;
if (f8 == 0xff) {
return ggml_fp32_to_fp16(-NAN);
} else if (f8 == 0x7f) {
return ggml_fp32_to_fp16(NAN);
}
uint32_t sign = f8 & 0x80;
uint32_t exponent = (f8 & 0x78) >> 3;
uint32_t mantissa = f8 & 0x07;
uint32_t result = sign << 24;
if (exponent == 0) {
if (mantissa > 0) {
exponent = 0x7f - exponent_bias;
// yes, 2 times
if ((mantissa & 0x04) == 0) {
mantissa &= 0x03;
mantissa <<= 1;
exponent -= 1;
}
if ((mantissa & 0x04) == 0) {
mantissa &= 0x03;
mantissa <<= 1;
exponent -= 1;
}
result |= (mantissa & 0x03) << 21;
result |= exponent << 23;
}
} else {
result |= mantissa << 20;
exponent += 0x7f - exponent_bias;
result |= exponent << 23;
}
return ggml_fp32_to_fp16(*reinterpret_cast<const float*>(&result));
}
uint16_t f8_e5m2_to_f16(uint8_t fp8) {
return static_cast<uint16_t>(fp8) << 8;
}
void f8_e4m3_to_f16_vec(uint8_t* src, uint16_t* dst, int64_t n) {
// support inplace op
for (int64_t i = n - 1; i >= 0; i--) {
dst[i] = f8_e4m3_to_f16(src[i]);
}
}
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--) {
dst[i] = f8_e5m2_to_f16(src[i]);
}
}
void f64_to_f32_vec(double* src, float* dst, int64_t n) {
// support inplace op
for (int64_t i = 0; i < n; i++) {
@@ -929,9 +869,7 @@ std::vector<MmapTensorStore> ModelLoader::mmap_tensors(std::map<std::string, ggm
if (dst_tensor == nullptr)
continue;
if (tensor_storage.is_f8_e4m3 ||
tensor_storage.is_f8_e5m2 ||
tensor_storage.is_f64 ||
if (tensor_storage.is_f64 ||
tensor_storage.is_i64 ||
tensor_storage.type != dst_tensor->type) {
continue;
@@ -1215,11 +1153,7 @@ 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) {
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());
} else if (tensor_storage.is_f64) {
if (tensor_storage.is_f64) {
f64_to_f32_vec((double*)read_buf, (float*)target_buf, tensor_storage.nelements());
} else if (tensor_storage.is_i64) {
i64_to_i32_vec((int64_t*)read_buf, (int32_t*)target_buf, tensor_storage.nelements());
+5
View File
@@ -1569,6 +1569,11 @@ std::string convert_tensor_name(std::string name, SDVersion version) {
}
}
static const std::vector<std::pair<std::string, std::string>> generic_name_map = {
{".scale_weight", ".weight_scale"},
};
replace_with_name_map(name, generic_name_map);
if (is_lora) {
name = "lora." + name;
}