diff --git a/rwkv.cpp b/rwkv.cpp index e681d51..e228e57 100644 --- a/rwkv.cpp +++ b/rwkv.cpp @@ -108,37 +108,6 @@ bool set_block_parameter(std::unordered_map * return set_parameter(parameters, full_key, dest); } -size_t get_memory_required_mb(int32_t n_vocab, int32_t n_layer, int32_t n_embed, int32_t data_type) { - if (n_vocab == 50277) { - // Non-exact values are extrapolated (slightly bigger than needed). - // TODO Measure values exactly - static const size_t memory_required_mb[6][4] = { - /* FP32 FP16 Q4_0 Q4_1 - 169M */ { 650, 327, 105, 165}, // All measured exactly - /* 430M */ { 1650, 830, 263, 415}, // FP32, FP16 are exact - /* 1.5B */ { 5795, 2907, 923, 1454}, // FP32, FP16 are exact - /* 3B */ {11610, 5720, 1816, 2860}, // FP16 is exact - /* 7B */ {27090, 13634, 4328, 6817}, - /* 14B */ {54180, 27267, 8656, 13634} - }; - - /* 169M */ if (n_layer == 12 && n_embed == 768) return memory_required_mb[0][data_type]; - /* 430M */ if (n_layer == 24 && n_embed == 1024) return memory_required_mb[1][data_type]; - /* 1.5B */ if (n_layer == 24 && n_embed == 2048) return memory_required_mb[2][data_type]; - /* 3B */ if (n_layer == 32 && n_embed == 2560) return memory_required_mb[3][data_type]; - /* 7B */ if (n_layer == 32 && n_embed == 4096) return memory_required_mb[4][data_type]; - /* 14B */ if (n_layer == 40 && n_embed == 5120) return memory_required_mb[5][data_type]; - } - - fprintf( - stderr, - "Unknown RWKV model configuration: n_vocab = %d, n_layer = %d, n_embed = %d, data_type = %d; allocating 4 GB of memory\n", - n_vocab, n_layer, n_embed, data_type - ); - - return size_t(4) * 1024; -} - // --- Operators --- struct ggml_tensor * rwkv_layer_norm(ggml_context * ctx, struct ggml_tensor * x, struct ggml_tensor * weight, struct ggml_tensor * bias) { @@ -196,9 +165,30 @@ struct rwkv_context * rwkv_init_from_file(const char * file_path, uint32_t n_thr model->data_type ); + // Parameter tensors would take at least this amount in memory. + size_t file_size; + + { + auto fin = std::ifstream(file_path, std::ios::binary); + RWKV_ASSERT_NULL(fin, "Failed to open file %s", file_path); + fin.seekg(0, fin.end); + file_size = fin.tellg(); + fin.close(); + } + + size_t memory_required = file_size + + // Intermediary vectors for calculation; there are around 100 calls to ggml + size_t(100) * model->n_embed * sizeof(float) + + // State, in and out + size_t(2) * 5 * model->n_layer * model->n_embed * sizeof(float) + + // Logits + size_t(model->n_vocab) * sizeof(float) + + // +32 MB just for any overhead + size_t(32) * 1024 * 1024; + // Initialize ggml struct ggml_init_params params; - params.mem_size = get_memory_required_mb(model->n_vocab, model->n_layer, model->n_embed, model->data_type) * 1024 * 1024; + params.mem_size = memory_required; params.mem_buffer = NULL; struct ggml_context * ctx = ggml_init(params);