231 lines
9.6 KiB
C++
231 lines
9.6 KiB
C++
#include "ggml.h"
|
|
#include "ggml-backend.h"
|
|
|
|
#include "common-ggml.h"
|
|
|
|
#include <cassert>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
struct parakeet_hparams {
|
|
int32_t n_vocab = 0;
|
|
int32_t n_audio_ctx = 0;
|
|
int32_t n_audio_state = 0;
|
|
int32_t n_audio_head = 0;
|
|
int32_t n_audio_layer = 0;
|
|
int32_t n_mels = 0;
|
|
int32_t ftype = 0;
|
|
int32_t n_fft = 0;
|
|
int32_t subsampling_factor = 0;
|
|
int32_t n_subsampling_channels = 0;
|
|
int32_t n_conv_kernel = 0;
|
|
int32_t n_pred_dim = 0;
|
|
int32_t n_pred_layers = 0;
|
|
int32_t n_tdt_durations = 0;
|
|
int32_t n_max_tokens = 0;
|
|
};
|
|
|
|
static bool parakeet_model_quantize(const std::string & fname_inp, const std::string & fname_out, ggml_ftype ftype) {
|
|
printf("%s: loading model from '%s'\n", __func__, fname_inp.c_str());
|
|
|
|
auto finp = std::ifstream(fname_inp, std::ios::binary);
|
|
if (!finp) {
|
|
fprintf(stderr, "%s: failed to open '%s' for reading\n", __func__, fname_inp.c_str());
|
|
return false;
|
|
}
|
|
|
|
auto fout = std::ofstream(fname_out, std::ios::binary);
|
|
if (!fout) {
|
|
fprintf(stderr, "%s: failed to open '%s' for writing\n", __func__, fname_out.c_str());
|
|
return false;
|
|
}
|
|
|
|
// magic
|
|
{
|
|
uint32_t magic;
|
|
finp.read((char *) &magic, sizeof(magic));
|
|
if (magic != GGML_FILE_MAGIC) {
|
|
fprintf(stderr, "%s: invalid model file (bad magic)\n", __func__);
|
|
return false;
|
|
}
|
|
fout.write((char *) &magic, sizeof(magic));
|
|
}
|
|
|
|
// hparams
|
|
parakeet_hparams hparams;
|
|
{
|
|
finp.read((char *) &hparams.n_vocab, sizeof(hparams.n_vocab));
|
|
finp.read((char *) &hparams.n_audio_ctx, sizeof(hparams.n_audio_ctx));
|
|
finp.read((char *) &hparams.n_audio_state, sizeof(hparams.n_audio_state));
|
|
finp.read((char *) &hparams.n_audio_head, sizeof(hparams.n_audio_head));
|
|
finp.read((char *) &hparams.n_audio_layer, sizeof(hparams.n_audio_layer));
|
|
finp.read((char *) &hparams.n_mels, sizeof(hparams.n_mels));
|
|
finp.read((char *) &hparams.ftype, sizeof(hparams.ftype));
|
|
finp.read((char *) &hparams.n_fft, sizeof(hparams.n_fft));
|
|
finp.read((char *) &hparams.subsampling_factor, sizeof(hparams.subsampling_factor));
|
|
finp.read((char *) &hparams.n_subsampling_channels, sizeof(hparams.n_subsampling_channels));
|
|
finp.read((char *) &hparams.n_conv_kernel, sizeof(hparams.n_conv_kernel));
|
|
finp.read((char *) &hparams.n_pred_dim, sizeof(hparams.n_pred_dim));
|
|
finp.read((char *) &hparams.n_pred_layers, sizeof(hparams.n_pred_layers));
|
|
finp.read((char *) &hparams.n_tdt_durations, sizeof(hparams.n_tdt_durations));
|
|
finp.read((char *) &hparams.n_max_tokens, sizeof(hparams.n_max_tokens));
|
|
|
|
const int32_t qntvr_src = hparams.ftype / GGML_QNT_VERSION_FACTOR;
|
|
const int32_t ftype_dst = GGML_QNT_VERSION * GGML_QNT_VERSION_FACTOR + ftype;
|
|
|
|
fprintf(stderr, "%s: n_vocab = %d\n", __func__, hparams.n_vocab);
|
|
fprintf(stderr, "%s: n_audio_state = %d\n", __func__, hparams.n_audio_state);
|
|
fprintf(stderr, "%s: n_audio_layer = %d\n", __func__, hparams.n_audio_layer);
|
|
fprintf(stderr, "%s: n_mels = %d\n", __func__, hparams.n_mels);
|
|
fprintf(stderr, "%s: ftype (src) = %d\n", __func__, hparams.ftype);
|
|
fprintf(stderr, "%s: qntvr (src) = %d\n", __func__, qntvr_src);
|
|
fprintf(stderr, "%s: ftype (dst) = %d\n", __func__, ftype_dst);
|
|
fprintf(stderr, "%s: qntvr (dst) = %d\n", __func__, GGML_QNT_VERSION);
|
|
|
|
fout.write((char *) &hparams.n_vocab, sizeof(hparams.n_vocab));
|
|
fout.write((char *) &hparams.n_audio_ctx, sizeof(hparams.n_audio_ctx));
|
|
fout.write((char *) &hparams.n_audio_state, sizeof(hparams.n_audio_state));
|
|
fout.write((char *) &hparams.n_audio_head, sizeof(hparams.n_audio_head));
|
|
fout.write((char *) &hparams.n_audio_layer, sizeof(hparams.n_audio_layer));
|
|
fout.write((char *) &hparams.n_mels, sizeof(hparams.n_mels));
|
|
fout.write((char *) &ftype_dst, sizeof(ftype_dst));
|
|
fout.write((char *) &hparams.n_fft, sizeof(hparams.n_fft));
|
|
fout.write((char *) &hparams.subsampling_factor, sizeof(hparams.subsampling_factor));
|
|
fout.write((char *) &hparams.n_subsampling_channels, sizeof(hparams.n_subsampling_channels));
|
|
fout.write((char *) &hparams.n_conv_kernel, sizeof(hparams.n_conv_kernel));
|
|
fout.write((char *) &hparams.n_pred_dim, sizeof(hparams.n_pred_dim));
|
|
fout.write((char *) &hparams.n_pred_layers, sizeof(hparams.n_pred_layers));
|
|
fout.write((char *) &hparams.n_tdt_durations, sizeof(hparams.n_tdt_durations));
|
|
fout.write((char *) &hparams.n_max_tokens, sizeof(hparams.n_max_tokens));
|
|
}
|
|
|
|
// mel filterbank
|
|
{
|
|
int32_t n_mel, n_fb;
|
|
finp.read((char *) &n_mel, sizeof(n_mel));
|
|
fout.write((char *) &n_mel, sizeof(n_mel));
|
|
finp.read((char *) &n_fb, sizeof(n_fb));
|
|
fout.write((char *) &n_fb, sizeof(n_fb));
|
|
|
|
const size_t n = (size_t) n_mel * n_fb;
|
|
std::vector<float> buf(n);
|
|
finp.read((char *) buf.data(), n * sizeof(float));
|
|
fout.write((char *) buf.data(), n * sizeof(float));
|
|
}
|
|
|
|
// window function
|
|
{
|
|
int32_t n_window;
|
|
finp.read((char *) &n_window, sizeof(n_window));
|
|
fout.write((char *) &n_window, sizeof(n_window));
|
|
|
|
std::vector<float> buf(n_window);
|
|
finp.read((char *) buf.data(), n_window * sizeof(float));
|
|
fout.write((char *) buf.data(), n_window * sizeof(float));
|
|
}
|
|
|
|
// TDT durations
|
|
{
|
|
std::vector<uint32_t> buf(hparams.n_tdt_durations);
|
|
finp.read((char *) buf.data(), hparams.n_tdt_durations * sizeof(uint32_t));
|
|
fout.write((char *) buf.data(), hparams.n_tdt_durations * sizeof(uint32_t));
|
|
}
|
|
|
|
// vocab
|
|
{
|
|
int32_t n_tokens;
|
|
finp.read((char *) &n_tokens, sizeof(n_tokens));
|
|
fout.write((char *) &n_tokens, sizeof(n_tokens));
|
|
|
|
for (int i = 0; i < n_tokens; ++i) {
|
|
int32_t len;
|
|
finp.read((char *) &len, sizeof(len));
|
|
fout.write((char *) &len, sizeof(len));
|
|
|
|
std::string token(len, '\0');
|
|
finp.read(&token[0], len);
|
|
fout.write(&token[0], len);
|
|
}
|
|
}
|
|
|
|
// tensors — quantize 2D weights skipping tensors that must stay F32:
|
|
// ggml_ssm_conv / ggml_conv2d_dw CUDA kernels require F32 weights.
|
|
// pos_bias_u / pos_bias_v are declared F32 in the loader.
|
|
const std::vector<std::string> to_quant = { ".*" };
|
|
std::vector<std::string> to_skip = {
|
|
// CUDA kernel constraints (ggml_ssm_conv / ggml_conv2d_dw require F32 weights)
|
|
"encoder\\.layers\\..+\\.conv\\.depthwise_conv\\.weight",
|
|
// Declared F32 in loader (pos_bias tensors)
|
|
"encoder\\.layers\\..+\\.self_attn\\.pos_bias_u",
|
|
"encoder\\.layers\\..+\\.self_attn\\.pos_bias_v",
|
|
};
|
|
|
|
// Prediction/joint tensors use n_pred_dim as their inner dimension. K-quant
|
|
// types (block size 256) cannot quantize 640 evenly, so keep them F32. For
|
|
// other types (Q8_0, Q4_0, block size 32) 640 is divisible and they can be
|
|
// quantized normally. The loader mirrors this logic at load time.
|
|
{
|
|
const ggml_type qtype = ggml_ftype_to_ggml_type(ftype);
|
|
const int32_t blck = ggml_blck_size(qtype);
|
|
if (blck > 1 && hparams.n_pred_dim % blck != 0) {
|
|
to_skip.push_back("decoder\\.prediction\\.embed\\.weight");
|
|
to_skip.push_back("decoder\\.prediction\\.dec_rnn\\.lstm\\.weight_ih_l.*");
|
|
to_skip.push_back("decoder\\.prediction\\.dec_rnn\\.lstm\\.weight_hh_l.*");
|
|
to_skip.push_back("joint\\.pred\\.weight");
|
|
to_skip.push_back("joint\\.joint_net\\.2\\.weight");
|
|
}
|
|
}
|
|
|
|
if (!ggml_common_quantize_0(finp, fout, ftype, to_quant, to_skip)) {
|
|
fprintf(stderr, "%s: failed to quantize tensors\n", __func__);
|
|
return false;
|
|
}
|
|
|
|
finp.close();
|
|
fout.close();
|
|
|
|
return true;
|
|
}
|
|
|
|
int main(int argc, char ** argv) {
|
|
ggml_backend_load_all();
|
|
|
|
if (argc != 4) {
|
|
fprintf(stderr, "usage: %s model-f32.bin model-quant.bin type\n", argv[0]);
|
|
ggml_print_ftypes(stderr);
|
|
return 1;
|
|
}
|
|
|
|
// initialise F16 lookup tables
|
|
{
|
|
struct ggml_init_params params = { 0, NULL, false };
|
|
struct ggml_context * ctx = ggml_init(params);
|
|
ggml_free(ctx);
|
|
}
|
|
|
|
const std::string fname_inp = argv[1];
|
|
const std::string fname_out = argv[2];
|
|
const ggml_ftype ftype = ggml_parse_ftype(argv[3]);
|
|
|
|
if (ftype == GGML_FTYPE_UNKNOWN) {
|
|
fprintf(stderr, "%s: invalid quantization type\n", argv[0]);
|
|
ggml_print_ftypes(stderr);
|
|
return 1;
|
|
}
|
|
|
|
const int64_t t_start_us = ggml_time_us();
|
|
|
|
if (!parakeet_model_quantize(fname_inp, fname_out, ftype)) {
|
|
fprintf(stderr, "%s: failed to quantize model from '%s'\n", argv[0], fname_inp.c_str());
|
|
return 1;
|
|
}
|
|
|
|
printf("\n%s: quantize time = %8.2f ms\n", argv[0], (ggml_time_us() - t_start_us) / 1000.0f);
|
|
printf("%s: output model = %s\n", argv[0], fname_out.c_str());
|
|
|
|
return 0;
|
|
}
|