Skip to content

quantize: cap working memory size to avoid loading big tensors onto RAM #27795

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ngxson merged 1 commit into from
Aug 27, 2026
+112 −79
Merged

quantize: cap working memory size to avoid loading big tensors onto RAM #27795

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions include/llama.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -437,6 +437,7 @@ extern "C" {
const struct llama_model_kv_override * kv_overrides; // pointer to kv overrides
const struct llama_model_tensor_override * tt_overrides; // pointer to tensor overrides
const int32_t * prune_layers; // pointer to layer indices to prune
size_t max_buf_size; // max bytes of tensor rows kept in memory at once, 0 = default (8 GiB)
} llama_model_quantize_params;

typedef struct llama_logit_bias {
Expand Down
25 changes: 12 additions & 13 deletions src/llama-model-loader.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1400,27 +1400,26 @@ void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const {
mappings.at(w.idx)->unmap_fragment(w.offs, w.offs + ggml_nbytes(w.tensor));
}

void llama_model_loader::load_data_for(struct ggml_tensor * cur) const {
const auto & w = require_weight(ggml_get_name(cur));
const void * llama_model_loader::load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const {
GGML_ASSERT(offs + size <= ggml_nbytes(w.tensor));

const void * data = buf;

if (use_mmap) {
const auto & mapping = mappings.at(w.idx);
if (cur->data == nullptr) {
cur->data = (uint8_t *)mapping->addr() + w.offs;
} else {
memcpy(cur->data, (uint8_t *)mapping->addr() + w.offs, ggml_nbytes(cur));
}
data = (const uint8_t *) mappings.at(w.idx)->addr() + w.offs + offs;
} else {
GGML_ASSERT(cur->data != nullptr);
GGML_ASSERT(buf != nullptr);
GGML_ASSERT(w.idx < files.size());
const auto & file = files.at(w.idx);
file->seek(w.offs, SEEK_SET);
file->read_raw(cur->data, ggml_nbytes(cur));
file->seek(w.offs + offs, SEEK_SET);
file->read_raw(buf, size);
}

if (check_tensors && !ggml_validate_row_data(cur->type, cur->data, ggml_nbytes(cur))) {
throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur)));
if (check_tensors && !ggml_validate_row_data(w.tensor->type, data, size)) {
throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(w.tensor)));
}

return data;
}

bool llama_model_loader::load_all_data(
Expand Down
5 changes: 3 additions & 2 deletions src/llama-model-loader.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,8 +197,9 @@ struct llama_model_loader {
// release a weight's mmap pages
void unmap_weight(const llama_tensor_weight & w) const;

// for backwards compatibility, does not support ggml-backend
void load_data_for(struct ggml_tensor * cur) const;
// read a byte range of a weight's data
// with mmap, returns a pointer into the mapping, otherwise reads into buf and returns buf
const void * load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const;

// Returns false if cancelled by progress_callback
bool load_all_data(
Expand Down
143 changes: 81 additions & 62 deletions src/llama-quant.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@ enum class tensor_category {
OTHER
};

// max amount of tensor data kept in memory while quantizing a single tensor
static const size_t LLAMA_QUANT_MAX_BUF_SIZE = 8ull*1024*1024*1024;

static void zeros(std::ofstream & file, size_t n) {
char zero = 0;
for (size_t i = 0; i < n; ++i) {
Expand Down Expand Up @@ -211,46 +214,41 @@ struct tensor_metadata {
//

static void llama_tensor_dequantize_impl(
ggml_tensor * tensor, std::vector<no_init<float>> & output, std::vector<std::thread> & workers,
ggml_type type, const void * data, float * f32_output, std::vector<std::thread> & workers,
const size_t nelements, const int nthread
) {
if (output.size() < nelements) {
output.resize(nelements);
}
float * f32_output = (float *) output.data();

const ggml_type_traits * qtype = ggml_get_type_traits(tensor->type);
if (ggml_is_quantized(tensor->type)) {
const ggml_type_traits * qtype = ggml_get_type_traits(type);
if (ggml_is_quantized(type)) {
if (qtype->to_float == NULL) {
throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(tensor->type)));
throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(type)));
}
} else if (tensor->type != GGML_TYPE_F16 &&
tensor->type != GGML_TYPE_BF16) {
throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(tensor->type)));
} else if (type != GGML_TYPE_F16 &&
type != GGML_TYPE_BF16) {
throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(type)));
}

if (nthread < 2) {
if (tensor->type == GGML_TYPE_F16) {
ggml_fp16_to_fp32_row((ggml_fp16_t *)tensor->data, f32_output, nelements);
} else if (tensor->type == GGML_TYPE_BF16) {
ggml_bf16_to_fp32_row((ggml_bf16_t *)tensor->data, f32_output, nelements);
} else if (ggml_is_quantized(tensor->type)) {
qtype->to_float(tensor->data, f32_output, nelements);
if (type == GGML_TYPE_F16) {
ggml_fp16_to_fp32_row((const ggml_fp16_t *)data, f32_output, nelements);
} else if (type == GGML_TYPE_BF16) {
ggml_bf16_to_fp32_row((const ggml_bf16_t *)data, f32_output, nelements);
} else if (ggml_is_quantized(type)) {
qtype->to_float(data, f32_output, nelements);
} else {
GGML_ABORT("fatal error"); // unreachable
}
return;
}

size_t block_size;
if (tensor->type == GGML_TYPE_F16 ||
tensor->type == GGML_TYPE_BF16) {
if (type == GGML_TYPE_F16 ||
type == GGML_TYPE_BF16) {
block_size = 1;
} else {
block_size = (size_t)ggml_blck_size(tensor->type);
block_size = (size_t)ggml_blck_size(type);
}

size_t block_size_bytes = ggml_type_size(tensor->type);
size_t block_size_bytes = ggml_type_size(type);

GGML_ASSERT(nelements % block_size == 0);
size_t nblocks = nelements / block_size;
Expand All @@ -265,16 +263,16 @@ static void llama_tensor_dequantize_impl(
size_t thr_elems = thr_blocks * block_size; // number of elements for this thread
size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread

auto compute = [qtype] (ggml_type typ, uint8_t * inbuf, float * outbuf, int nels) {
auto compute = [qtype] (ggml_type typ, const uint8_t * inbuf, float * outbuf, int nels) {
if (typ == GGML_TYPE_F16) {
ggml_fp16_to_fp32_row((ggml_fp16_t *)inbuf, outbuf, nels);
ggml_fp16_to_fp32_row((const ggml_fp16_t *)inbuf, outbuf, nels);
} else if (typ == GGML_TYPE_BF16) {
ggml_bf16_to_fp32_row((ggml_bf16_t *)inbuf, outbuf, nels);
ggml_bf16_to_fp32_row((const ggml_bf16_t *)inbuf, outbuf, nels);
} else {
qtype->to_float(inbuf, outbuf, nels);
}
};
workers.emplace_back(compute, tensor->type, (uint8_t *) tensor->data + in_buff_offs, f32_output + out_buff_offs, thr_elems);
workers.emplace_back(compute, type, (const uint8_t *) data + in_buff_offs, f32_output + out_buff_offs, thr_elems);
in_buff_offs += thr_block_bytes;
out_buff_offs += thr_elems;
}
Expand Down Expand Up @@ -1093,6 +1091,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
std::vector<no_init<uint8_t>> work;
std::vector<no_init<float>> f32_conv_buf;

const size_t max_buf_size = params->max_buf_size ? params->max_buf_size : LLAMA_QUANT_MAX_BUF_SIZE;

int cur_split = -1;
std::ofstream fout;
auto close_ofstream = [&]() {
Expand Down Expand Up @@ -1143,15 +1143,13 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::

const size_t tensor_size = ggml_nbytes(tensor);

if (!params->dry_run) {
if (!ml.use_mmap) {
if (read_data.size() < tensor_size) {
read_data.resize(tensor_size);
}
tensor->data = read_data.data();
// read a byte range of the current tensor
auto load_range = [&](size_t offs, size_t size) -> const void * {
if (!ml.use_mmap && read_data.size() < size) {
read_data.resize(size);
}
ml.load_data_for(tensor);
}
return ml.load_data_range(weight, offs, size, read_data.data());
};

LLAMA_LOG_INFO("[%4d/%4d] %-36s - [%s], type = %6s, ",
++idx, ml.n_tensors,
Expand All @@ -1166,7 +1164,6 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
// in then there's nothing to do.
bool quantize = cur_type != new_type;

void * new_data;
size_t new_size;

if (params->dry_run) {
Expand All @@ -1190,12 +1187,18 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
} else {
// no --dry-run, perform quantization
if (!quantize) {
new_data = tensor->data;
new_size = tensor_size;
LLAMA_LOG_INFO("size = %8.3f MiB\n", tensor_size/1024.0/1024.0);
} else {
const int64_t nelements = ggml_nelements(tensor);

// copy in slabs of whole rows, so that each slab can be validated
const size_t row_size = ggml_row_size(tensor->type, tensor->ne[0]);
const size_t slab_size = std::max<size_t>(row_size, (max_buf_size/row_size)*row_size);

for (size_t offs = 0; offs < tensor_size; offs += slab_size) {
const size_t size = std::min(slab_size, tensor_size - offs);
fout.write((const char *) load_range(offs, size), size);
}
} else {
const float * imatrix = nullptr;
if (imatrix_data) {
auto it = imatrix_data->find(tm.remapped_imatrix_name);
Expand Down Expand Up @@ -1227,43 +1230,60 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
throw std::runtime_error(format("Missing importance matrix for tensor %s in a very low-bit quantization", tensor->name));
}

float * f32_data;

if (tensor->type == GGML_TYPE_F32) {
f32_data = (float *) tensor->data;
} else if (ggml_is_quantized(tensor->type) && !params->allow_requantize) {
if (ggml_is_quantized(tensor->type) && !params->allow_requantize) {
throw std::runtime_error(format("requantizing from type %s is disabled", ggml_type_name(tensor->type)));
} else {
llama_tensor_dequantize_impl(tensor, f32_conv_buf, workers, nelements, nthread);
f32_data = (float *) f32_conv_buf.data();
}

LLAMA_LOG_INFO("converting to %s .. ", ggml_type_name(new_type));
fflush(stdout);

if (work.size() < (size_t)nelements * 4) {
work.resize(nelements * 4); // upper bound on size
}
new_data = work.data();

const int64_t n_per_row = tensor->ne[0];
const int64_t nrows = tensor->ne[1];

const size_t row_size_src = ggml_row_size(tensor->type, n_per_row);
const size_t row_size_dst = ggml_row_size(new_type, n_per_row);

// process the rows in slabs, so that the buffers stay below max_buf_size
const size_t bytes_per_row = row_size_src + row_size_dst + (tensor->type == GGML_TYPE_F32 ? 0 : n_per_row*sizeof(float));
const int64_t nrows_slab = std::max<int64_t>(1, std::min<int64_t>(nrows, max_buf_size/bytes_per_row));

static const int64_t min_chunk_size = 32 * 512;
const int64_t chunk_size = (n_per_row >= min_chunk_size ? n_per_row : n_per_row * ((min_chunk_size + n_per_row - 1)/n_per_row));

const int64_t nelements_matrix = tensor->ne[0] * tensor->ne[1];
const int64_t nchunk = (nelements_matrix + chunk_size - 1)/chunk_size;
const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1;

// quantize each expert separately since they have different importance matrices
new_size = 0;
for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) {
const float * f32_data_03 = f32_data + i03 * nelements_matrix;
void * new_data_03 = (char *)new_data + ggml_row_size(new_type, n_per_row) * i03 * nrows;
const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr;

new_size += llama_tensor_quantize_impl(new_type, f32_data_03, new_data_03, chunk_size, nrows, n_per_row, imatrix_03, workers, nthread_use);
for (int64_t ir = 0; ir < nrows; ir += nrows_slab) {
const int64_t nrows_cur = std::min(nrows_slab, nrows - ir);
const int64_t nelements_cur = nrows_cur * n_per_row;

const void * src = load_range((i03*nrows + ir)*row_size_src, nrows_cur*row_size_src);

const float * f32_data;
if (tensor->type == GGML_TYPE_F32) {
f32_data = (const float *) src;
} else {
if (f32_conv_buf.size() < (size_t) nelements_cur) {
f32_conv_buf.resize(nelements_cur);
}
llama_tensor_dequantize_impl(tensor->type, src, (float *) f32_conv_buf.data(), workers, nelements_cur, nthread);
f32_data = (const float *) f32_conv_buf.data();
}

if (work.size() < nrows_cur*row_size_dst) {
work.resize(nrows_cur*row_size_dst);
}

const int64_t nchunk = (nelements_cur + chunk_size - 1)/chunk_size;
const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1;

const size_t size_cur = llama_tensor_quantize_impl(new_type, f32_data, work.data(), chunk_size, nrows_cur, n_per_row, imatrix_03, workers, nthread_use);

fout.write((const char *) work.data(), size_cur);
new_size += size_cur;
}
}
LLAMA_LOG_INFO("size = %8.2f MiB -> %8.2f MiB\n", tensor_size/1024.0/1024.0, new_size/1024.0/1024.0);
}
Expand All @@ -1273,10 +1293,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
// update the gguf metadata as we go
gguf_set_tensor_type(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_type);
GGML_ASSERT(gguf_get_tensor_size(ctx_outs[cur_split].get(), gguf_find_tensor(ctx_outs[cur_split].get(), metadata[i].name.c_str())) == new_size);
gguf_set_tensor_data(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_data);

// write tensor data + padding
fout.write((const char *) new_data, new_size);
// tensor data is already written, add the padding
zeros(fout, GGML_PAD(new_size, align) - new_size);

// unmap the tensor to free memory
Expand Down Expand Up @@ -1323,7 +1341,8 @@ llama_model_quantize_params llama_model_quantize_default_params() {
/*.imatrix =*/ nullptr,
/*.kv_overrides =*/ nullptr,
/*.tensor_type =*/ nullptr,
/*.prune_layers =*/ nullptr
/*.prune_layers =*/ nullptr,
/*.max_buf_size =*/ LLAMA_QUANT_MAX_BUF_SIZE
};

return result;
Expand Down
17 changes: 15 additions & 2 deletions tools/quantize/quantize.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -122,7 +122,7 @@ static bool try_parse_ftype(const std::string & ftype_str_in, llama_ftype & ftyp
static void usage(const char * executable) {
printf("usage: %s [--help] [--allow-requantize] [--leave-output-tensor] [--pure] [--imatrix] [--include-weights]\n", executable);
printf(" [--exclude-weights] [--output-tensor-type] [--token-embedding-type] [--tensor-type] [--tensor-type-file]\n");
printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run]\n");
printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run] [--max-buffer-size]\n");
printf(" model-f32.gguf [model-quant.gguf] type [nthreads]\n\n");
printf(" --allow-requantize\n");
printf(" allow requantizing tensors that have already been quantized\n");
Expand Down Expand Up @@ -161,7 +161,10 @@ static void usage(const char * executable) {
printf(" WARNING: this is an advanced option, use with care.\n");
printf(" --dry-run\n");
printf(" calculate and show the final quantization size without performing quantization\n");
printf(" example: llama-quantize --dry-run model-f32.gguf Q4_K\n\n");
printf(" example: llama-quantize --dry-run model-f32.gguf Q4_K\n");
printf(" --max-buffer-size MiB\n");
printf(" max amount of tensor rows kept in memory while quantizing one tensor (default: 8192)\n");
printf(" lower it to quantize models with very large tensors on a machine with little RAM\n\n");
printf("note: --include-weights and --exclude-weights cannot be used together\n\n");
printf("-----------------------------------------------------------------------------\n");
printf(" allowed quantization types\n");
Expand Down Expand Up @@ -467,6 +470,16 @@ int llama_quantize(int argc, char ** argv) {
}
} else if (strcmp(argv[arg_idx], "--keep-split") == 0) {
params.keep_split = true;
} else if (strcmp(argv[arg_idx], "--max-buffer-size") == 0) {
if (arg_idx == argc-1) {
usage(argv[0]);
}
const int mib = atoi(argv[++arg_idx]);
if (mib <= 0) {
fprintf(stderr, "%s: invalid --max-buffer-size '%s'\n", __func__, argv[arg_idx]);
return 1;
}
params.max_buf_size = (size_t) mib * 1024 * 1024;
} else {
usage(argv[0]);
}
Expand Down
Loading