Add better control over MSE and directional bias computation

This commit is contained in:
Ed Addario 2025-09-10 18:00:56 +01:00
parent 7d04050b23
commit 04c07b3272
No known key found for this signature in database
GPG Key ID: E7875815A3230993
3 changed files with 39 additions and 35 deletions

View File

@ -365,7 +365,7 @@ extern "C" {
void * tensor_types; // pointer to vector containing tensor types
void * prune_layers; // pointer to vector containing layer indices to prune
float target_bpw; // target bits per weight (bpw)
bool precise_lambda; // use precise_lambda calculation - slow computation but very accurate
int32_t bpw_bias; // type of error bias to use: 0 = no bias (MSE only), 1 = fast (default), 2 = precise (slow)
} llama_model_quantize_params;
typedef struct llama_logit_bias {

View File

@ -902,26 +902,6 @@ static std::unordered_map<std::string, ggml_type> target_bpw_type(
return std::isfinite(total_err) ? total_err : infinity;
};
// Scaling factor to increase lambda when activations are concentrated
auto directional_scale = [&](const float * values, const float * activations, int64_t n_per_row) {
if (!activations) { return 1.0f; }
double sum_v = 0.0;
double sum_aw2 = 0.0;
double sum_a2 = 0.0;
for (int64_t j = 0; j < n_per_row; ++j) {
const double v = values ? std::max(0.0f, values[j]) : 1.0;
const double a = activations[j];
sum_v += v;
sum_aw2 += v * a * a;
sum_a2 += a * a;
}
const double rms_a = std::sqrt(sum_a2 / std::max(1.0, (double)n_per_row));
const double denom = std::sqrt(std::max(epsilon, sum_v)) * std::max(epsilon, rms_a);
const double scale = denom > 0.0 ? std::sqrt(sum_aw2) / denom : 1.0;
return (float)std::clamp(scale, 0.5, 2.0);
};
// Higher precision but much longer to compute
auto precise_lambda = [&](const ggml_tensor * t,
const std::vector<float> & f32_sample,
@ -979,11 +959,7 @@ static std::unordered_map<std::string, ggml_type> target_bpw_type(
if (ratios.empty()) { return 0.0f; }
std::nth_element(ratios.begin(), ratios.begin() + ratios.size() / 2, ratios.end());
double lambda = ratios[ratios.size() / 2];
const float scale = directional_scale(values, activations, n_per_row);
lambda *= scale;
lambda = std::clamp(lambda, 0.0, 8.0);
const double lambda = std::clamp(ratios[ratios.size() / 2], 0.0, 8.0);
return (float)lambda;
};
@ -1007,8 +983,7 @@ static std::unordered_map<std::string, ggml_type> target_bpw_type(
double base = 1.0 - s * s / (d * s2 + epsilon);
base = std::clamp(base, 0.0, 1.0);
const double scale = directional_scale(values, activations, n_per_row);
const double lambda = std::clamp(base * scale, 0.0, 1.0) * 8.0;
const double lambda = std::clamp(base, 0.0, 1.0) * 8.0;
return (float)lambda;
};
@ -1159,8 +1134,11 @@ static std::unordered_map<std::string, ggml_type> target_bpw_type(
{
const float * values = values_sample.empty() ? nullptr : values_sample.data();
const float * activations = activations_sample.empty() ? nullptr : activations_sample.data();
bias_lambda = params->precise_lambda ? precise_lambda(t, f32_sample, sample_rows_per_slice, values, activations, compatible_candidates) :
fast_lambda(values, activations, n_per_row);
if (params->bpw_bias == 1) {
bias_lambda = fast_lambda(values, activations, n_per_row);
} else if (params->bpw_bias == 2) {
bias_lambda = precise_lambda(t, f32_sample, sample_rows_per_slice, values, activations, compatible_candidates);
}
}
// Now evaluate candidates
@ -1656,7 +1634,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
} else {
LLAMA_LOG_WARN("%s: imatrix without activations provided, target bpw quantization will be less accurate - ", __func__);
}
LLAMA_LOG_INFO("using %s\n", params->precise_lambda ? "precise lambda (slow)" : "fast lambda");
const char* msg[] = {"no bias (MSE only)", "fast (default)", "precise (slow)"};
LLAMA_LOG_INFO("using %s error estimation\n", msg[params->bpw_bias]);
LLAMA_LOG_INFO("%s: computing tensor quantization mix to achieve %.4f bpw\n", __func__, params->target_bpw);
bpw_overrides = target_bpw_type(ml, read_data, model, tensors, mapped, values_data, activations_data, params, nthread);
} else {
@ -1967,7 +1946,7 @@ llama_model_quantize_params llama_model_quantize_default_params() {
/*.tensor_type =*/ nullptr,
/*.prune_layers =*/ nullptr,
/*.target_bpw =*/ -1.0f,
/*.precise_lambda =*/ false
/*.bpw_bias =*/ 1
};
return result;

View File

@ -134,7 +134,7 @@ static void usage(const char * executable) {
printf(" Advanced option to remove all tensors from the given layers\n");
printf(" --target-bpw: target bits per weight (bpw). Must be a positive number between 0.0 and 16.0\n");
printf(" Advanced option to automatically select quantization types to achieve a total bits per weight (bpw) target\n");
printf(" --precise-lambda: given a target bpw, use a high-precision error computation at the expense of longer processing times\n");
printf(" --bpw_bias: type of error bias to use: 0 = no bias (MSE only), 1 = fast (default), 2 = precise (slow)\n");
printf(" --keep-split: will generate quantized model in the same shards as input\n");
printf(" --override-kv KEY=TYPE:VALUE\n");
printf(" Advanced option to override model metadata by key in the quantized model. May be specified multiple times.\n");
@ -496,6 +496,27 @@ static bool parse_target_bpw(const char * data, float & target_bpw) {
return true;
}
static bool parse_bpw_bias(const char * data, int & bpw_bias) {
if (!data) {
printf("\n%s: error bias type not provided\n\n", __func__);
return false;
}
try {
bpw_bias = std::stoi(data);
if (bpw_bias < 0 || bpw_bias > 2) {
printf("\n%s: error bias type must be one of 0 (no bias, MSE only), 1 (fast), or 2 (precise, but slow)\n\n", __func__);
return false;
}
}
catch (const std::exception & e) {
printf("\n%s: '%s' is not valid. Target bits per weight (bpw) must be a positive number between 0.0 and 16.0\n\n", __func__, data);
return false;
}
return true;
}
int main(int argc, char ** argv) {
if (argc < 3) {
usage(argv[0]);
@ -510,6 +531,7 @@ int main(int argc, char ** argv) {
std::vector<tensor_quantization> tensor_types;
std::vector<int> prune_layers;
float target_bpw = -1.0f;
int bpw_bias = 1;
for (; arg_idx < argc && strncmp(argv[arg_idx], "--", 2) == 0; arg_idx++) {
if (strcmp(argv[arg_idx], "--leave-output-tensor") == 0) {
@ -540,8 +562,11 @@ int main(int argc, char ** argv) {
if (arg_idx == argc-1 || !parse_target_bpw(argv[++arg_idx], target_bpw)) {
usage(argv[0]);
}
} else if (strcmp(argv[arg_idx], "--precise-lambda") == 0) {
params.precise_lambda = true;
} else if (strcmp(argv[arg_idx], "--bpw-bias") == 0) {
if (arg_idx == argc-1 || !parse_bpw_bias(argv[++arg_idx], bpw_bias)) {
usage(argv[0]);
}
params.bpw_bias = bpw_bias;
} else if (strcmp(argv[arg_idx], "--prune-layers") == 0) {
if (arg_idx == argc-1 || !parse_layer_prune(argv[++arg_idx], prune_layers)) {
usage(argv[0]);