re-quantize: moe models

Author: foldl

src/backend.h

@@ -21,6 +21,7 @@ namespace chatllm
         tensor *init_tensor(ggml::tensor  *tensor,
                             ggml::type     type,
                             int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3);
+        tensor *init_tensor(ggml::tensor  *tensor, ggml::tensor *like);
         void change_type(ggml::tensor  *tensor, ggml::type type);
 
         size_t element_size(const ggml::tensor *tensor);

src/chat.cpp

@@ -1163,11 +1163,13 @@ namespace chatllm
     }
 
     size_t TensorInfo::read_tensor_data(tokenizer::DataReader *reader, size_t read_offset, size_t write_offset, size_t data_size,
-                                        ggml::type target_type)
+                                        ggml::type target_type, ggml::tensor *src_tensor)
     {
         CHATLLM_CHECK(data) << "backend buffer still not allocated!";
         CHATLLM_CHECK(target_type == ggml::type_of(tensor)) << "tensor type mismatch!";
-        CHATLLM_CHECK(data->get_size() >= write_offset + data_size) << "read_tensor_data(" << ggml::get_name(&tensor) << "): write data exceeds tensor data size";
+
+        const ggml::type original_type = src_tensor ? ggml::type_of(src_tensor) : this->original_type;
+        if (nullptr == src_tensor) src_tensor = &tensor;
 
         reader->seek(aligned_data_start(read_offset), SEEK_SET);
 
@@ -1185,7 +1187,7 @@ namespace chatllm
                 std::vector<uint8_t> buf_q;
 
                 ggml::tensor t;
-                ggml::init_tensor(&t, ggml::type::GGML_TYPE_F32, 4, tensor.ne);
+                ggml::init_tensor(&t, ggml::type::GGML_TYPE_F32, 4, src_tensor->ne);
                 buf.resize(ggml::nbytes(&t));
 
                 if (ggml::type::GGML_TYPE_F32 == original_type)
@@ -1194,18 +1196,19 @@ namespace chatllm
                 }
                 else
                 {
-                    ggml::init_tensor(&t, original_type, 4, tensor.ne);
+                    ggml::init_tensor(&t, original_type, 4, src_tensor->ne);
                     buf_q.resize(ggml::nbytes(&t));
                     reader->read_buffer(buf_q.data(), buf_q.size());
 
                     ggml::to_float(original_type, buf_q.data(), (float *)buf.data(), ggml::get_dim(&t, 0), ggml::nrows(&t));
                 }
 
-                ggml::init_tensor(&t, target_type, 4, tensor.ne);
+                ggml::init_tensor(&t, target_type, 4, src_tensor->ne);
                 buf_q.resize(ggml::nbytes(&t));
                 ggml::from_float(target_type, (const float *)buf.data(), (void *)buf_q.data(), ggml::get_dim(&t, 0), ggml::nrows(&t));
 
-                CHATLLM_CHECK(buf_q.size() == data_size) << "size mismatch? " << buf_q.size() << " : " << data_size;
+                if (data_size < buf_q.size())
+                    CHATLLM_CHECK(buf_q.size() == data_size) << "size mismatch? " << buf_q.size() << " : " << data_size;
 
                 alloc->get_backend()->write_tensor_data(&tensor, buf_q.data(), write_offset, buf_q.size());
 
@@ -1530,12 +1533,14 @@ namespace chatllm
             }
 
             size_t size = search->second.get_nbytes();
-            t.read_tensor_data(_file.get(), search->second._offset, write_offset, size, tensor->type);
+            size = t.read_tensor_data(_file.get(), search->second._offset, write_offset, size, tensor->type, &search->second.tensor);
+
+            CHATLLM_CHECK(total_size >= size) << "tensor " << name << " too much data: " << total_size << " > " << size;
 
             write_offset += size;
             total_size -= size;
         }
-        CHATLLM_CHECK(total_size == 0) << "tensor " << name << " not fully loaded.";
+        CHATLLM_CHECK(total_size == 0) << "tensor " << name << " not fully loaded, remain = " << total_size;
 
         t.assign_to(tensor);
     }

src/chat.h

@@ -686,7 +686,7 @@ namespace chatllm
 
         bool load(tokenizer::DataReader *reader, LayerBufAllocator *alloc, ggml::type target_type, size_t override_buffer_size = 0);
 
-        size_t read_tensor_data(tokenizer::DataReader *reader, size_t read_offset, size_t write_offset, size_t data_size, ggml::type target_type);
+        size_t read_tensor_data(tokenizer::DataReader *reader, size_t read_offset, size_t write_offset, size_t data_size, ggml::type target_type, ggml::tensor *src_tensor = nullptr);
         size_t read_raw_tensor_data(tokenizer::DataReader *reader, size_t data_size, void *p);
 
         size_t aligned_data_start(size_t offset);

src/layers.cpp

@@ -228,6 +228,12 @@ namespace chatllm
         return ggml::init_tensor(tensor, type, 4, ne);
     }
 
+    ggml::tensor *ggml::init_tensor(ggml::tensor  *tensor, ggml::tensor *like)
+    {
+        int64_t ne[4] = {ggml::get_dim(like, 0), ggml::get_dim(like, 1), ggml::get_dim(like, 2), ggml::get_dim(like, 3)};
+        return ggml::init_tensor(tensor, ggml::type_of(like), 4, ne);
+    }
+
     ggml::tensor *ggml::init_tensor(ggml::tensor *tensor,
         ggml::type    type,
         int           n_dims,

src/layers.h

@@ -572,7 +572,7 @@ namespace chatllm
               bias(use_bias ? ggml::new_tensor_1d(ctx, GGML_TYPE_F32, out_features) : nullptr) {}
 
         Linear(InitContext *ctx, int in_features, int out_features, ggml::tensor *weight, bool use_bias = true)
-            : weight(weight != NULL ? weight : ggml::new_tensor_2d(ctx, ctx->dtype, in_features, out_features)),
+            : weight(weight != NULL ? weight : ggml::new_tensor_2d(ctx, ggml::type_fallback(ctx->dtype, in_features), in_features, out_features)),
               bias(use_bias ? ggml::new_tensor_1d(ctx, GGML_TYPE_F32, out_features) : nullptr) {}
 
         int in_features() const { return (int)weight->ne[0]; }
@@ -604,7 +604,7 @@ namespace chatllm
         {}
 
         MultiLinear(InitContext *ctx, int in_features, int out_features, int multi, bool use_bias)
-            : weight(ggml::new_tensor_3d(ctx, ctx->dtype, in_features, out_features, multi)),
+            : weight(ggml::new_tensor_3d(ctx, ggml::type_fallback(ctx->dtype, in_features), in_features, out_features, multi)),
               bias(use_bias ? ggml::new_tensor_2d(ctx, ggml::type::GGML_TYPE_F32, out_features, multi) : nullptr)
         {
         }