Reduce llama-quantize BF16 memory residency

Author: i386

src/llama-quant.cpp

@@ -209,9 +209,9 @@ struct tensor_metadata {
 // dequantization
 //
 
-static void llama_tensor_dequantize_impl(
+static void llama_tensor_dequantize_chunk_impl(
     ggml_tensor * tensor, std::vector<no_init<float>> & output, std::vector<std::thread> & workers,
-    const size_t nelements, const int nthread
+    const int64_t start, const size_t nelements, const int nthread
 ) {
     if (output.size() < nelements) {
         output.resize(nelements);
@@ -228,19 +228,6 @@ static void llama_tensor_dequantize_impl(
         throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(tensor->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);
-        } else {
-            GGML_ABORT("fatal error"); // unreachable
-        }
-        return;
-    }
-
     size_t block_size;
     if (tensor->type == GGML_TYPE_F16 ||
         tensor->type == GGML_TYPE_BF16) {
@@ -249,9 +236,27 @@ static void llama_tensor_dequantize_impl(
         block_size = (size_t)ggml_blck_size(tensor->type);
     }
 
+    GGML_ASSERT(start >= 0);
+    GGML_ASSERT((size_t) start % block_size == 0);
+    GGML_ASSERT(nelements % block_size == 0);
+
     size_t block_size_bytes = ggml_type_size(tensor->type);
+    const size_t input_offset = ((size_t) start / block_size) * block_size_bytes;
+    uint8_t * input_data = (uint8_t *) tensor->data + input_offset;
+
+    if (nthread < 2) {
+        if (tensor->type == GGML_TYPE_F16) {
+            ggml_fp16_to_fp32_row((ggml_fp16_t *) input_data, f32_output, nelements);
+        } else if (tensor->type == GGML_TYPE_BF16) {
+            ggml_bf16_to_fp32_row((ggml_bf16_t *) input_data, f32_output, nelements);
+        } else if (ggml_is_quantized(tensor->type)) {
+            qtype->to_float(input_data, f32_output, nelements);
+        } else {
+            GGML_ABORT("fatal error"); // unreachable
+        }
+        return;
+    }
 
-    GGML_ASSERT(nelements % block_size == 0);
     size_t nblocks = nelements / block_size;
     size_t blocks_per_thread = nblocks / nthread;
     size_t spare_blocks = nblocks - (blocks_per_thread * nthread); // if blocks aren't divisible by thread count
@@ -264,7 +269,7 @@ 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, uint8_t * inbuf, float * outbuf, int64_t nels) {
             if (typ == GGML_TYPE_F16) {
                 ggml_fp16_to_fp32_row((ggml_fp16_t *)inbuf, outbuf, nels);
             } else if (typ == GGML_TYPE_BF16) {
@@ -273,14 +278,21 @@ static void llama_tensor_dequantize_impl(
                 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, tensor->type, input_data + in_buff_offs, f32_output + out_buff_offs, thr_elems);
         in_buff_offs += thr_block_bytes;
         out_buff_offs += thr_elems;
     }
     for (auto & w : workers) { w.join(); }
     workers.clear();
 }
 
+static void llama_tensor_dequantize_impl(
+    ggml_tensor * tensor, std::vector<no_init<float>> & output, std::vector<std::thread> & workers,
+    const size_t nelements, const int nthread
+) {
+    llama_tensor_dequantize_chunk_impl(tensor, output, workers, 0, nelements, nthread);
+}
+
 //
 // do we allow this tensor to be quantized?
 //
@@ -1208,43 +1220,70 @@ 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) {
-                    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 = ggml_row_size(new_type, n_per_row);
+                const size_t work_size = row_size * nrows * tensor->ne[2];
+
+                if (work.size() < work_size) {
+                    work.resize(work_size);
+                }
+                new_data = work.data();
 
                 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;
+                const bool stream_source = tensor->type == GGML_TYPE_F16 || tensor->type == GGML_TYPE_BF16;
 
                 // 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;
+                if (stream_source) {
+                    static const int64_t stream_chunk_size = 32 * 1024 * 1024;
+                    const int64_t rows_per_stream_chunk = std::max<int64_t>(1, std::max(chunk_size, stream_chunk_size) / n_per_row);
 
-                    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 i03 = 0; i03 < tensor->ne[2]; ++i03) {
+                        void * new_data_03 = (char *)new_data + row_size * i03 * nrows;
+                        const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr;
+
+                        for (int64_t first_row = 0; first_row < nrows; first_row += rows_per_stream_chunk) {
+                            const int64_t this_nrow = std::min(nrows - first_row, rows_per_stream_chunk);
+                            const int64_t start = i03 * nelements_matrix + first_row * n_per_row;
+                            const size_t chunk_nelements = this_nrow * n_per_row;
+
+                            llama_tensor_dequantize_chunk_impl(tensor, f32_conv_buf, workers, start, chunk_nelements, nthread);
+
+                            void * new_data_chunk = (char *) new_data_03 + row_size * first_row;
+                            const int64_t chunk_nchunk = (chunk_nelements + chunk_size - 1) / chunk_size;
+                            const int64_t chunk_nthread = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, chunk_nchunk)) : 1;
+
+                            new_size += llama_tensor_quantize_impl(new_type, (float *) f32_conv_buf.data(), new_data_chunk, chunk_size, this_nrow, n_per_row, imatrix_03, workers, chunk_nthread);
+                        }
+                    }
+                } else {
+                    float * f32_data;
+
+                    if (tensor->type == GGML_TYPE_F32) {
+                        f32_data = (float *) tensor->data;
+                    } else 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();
+                    }
+
+                    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 + row_size * 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);
+                    }
                 }
                 LLAMA_LOG_INFO("size = %8.2f MiB -> %8.2f MiB\n", tensor_size/1024.0/1024.0, new_size/1024.0/1024.0);
             }