182 lines
8.3 KiB
C++
182 lines
8.3 KiB
C++
#pragma once
|
|
// code-predictor-weights.h: 5-layer Qwen3 stack that predicts the
|
|
// acoustic codebooks 1..15 of every audio frame conditioned on the
|
|
// Talker hidden state and the codebook 0 token just sampled.
|
|
//
|
|
// Architecture mirrors the Talker block (pre-norm, GQA attention with
|
|
// QK-norm, SwiGLU MLP) with one important difference: RoPE is plain
|
|
// 1D (half-split, neox-style in GGUF terms) at freq base 1e6, not the
|
|
// multimodal interleaved variant the Talker uses.
|
|
//
|
|
// The MTP head carries fifteen private embedding tables and fifteen
|
|
// private linear heads, one pair per acoustic codebook. The talker
|
|
// codebook 0 stays handled by talker.codec_embedding and talker.codec_head.
|
|
//
|
|
// Tensor naming (convert.py output) :
|
|
// code_predictor.norm.weight [hidden]
|
|
// code_predictor.codec_embedding.{0..14}.weight [vocab, hidden]
|
|
// code_predictor.lm_head.{0..14}.weight [vocab, hidden]
|
|
// code_predictor.layers.{0..N-1}.input_layernorm.weight [hidden]
|
|
// code_predictor.layers.{0..N-1}.post_attention_layernorm.weight [hidden]
|
|
// code_predictor.layers.{0..N-1}.attn.{q,k,v,o}_proj.weight
|
|
// code_predictor.layers.{0..N-1}.attn.{q,k}_norm.weight [head_dim]
|
|
// code_predictor.layers.{0..N-1}.mlp.{gate,up,down}_proj.weight
|
|
|
|
#include "ggml-backend.h"
|
|
#include "ggml.h"
|
|
#include "gguf-weights.h"
|
|
#include "talker-weights.h"
|
|
#include "weight-ctx.h"
|
|
|
|
#include <cstdint>
|
|
#include <cstdio>
|
|
#include <cstring>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
struct CodePredictorWeights {
|
|
int hidden_size;
|
|
int intermediate_size;
|
|
int num_hidden_layers;
|
|
int num_attention_heads;
|
|
int num_key_value_heads;
|
|
int head_dim;
|
|
int vocab_size;
|
|
int num_acoustic_codebooks; // num_code_groups - 1
|
|
float rope_theta;
|
|
float rms_norm_eps;
|
|
|
|
struct ggml_tensor * norm_w;
|
|
std::vector<struct ggml_tensor *> codec_embedding; // size num_acoustic_codebooks
|
|
std::vector<struct ggml_tensor *> lm_head; // size num_acoustic_codebooks
|
|
|
|
// Optional small_to_mtp projection that brings the talker hidden
|
|
// dimension down to the predictor hidden dimension when the two
|
|
// differ (1.7B-base case: 2048 -> 1024). Both tensors are NULL when
|
|
// the upstream sets nn.Identity() i.e. talker_hidden == predictor_hidden
|
|
// (0.6B case). Loaded with gf_try_load_tensor so absence is silent.
|
|
struct ggml_tensor * mtp_proj_w;
|
|
struct ggml_tensor * mtp_proj_b;
|
|
|
|
// Layers reuse the same TalkerLayer struct since the per-layer
|
|
// tensor set is identical. Only the model-level wiring differs.
|
|
std::vector<TalkerLayer> layers;
|
|
|
|
struct ggml_context * weight_ctx;
|
|
ggml_backend_buffer_t weight_buf;
|
|
};
|
|
|
|
static bool code_predictor_weights_load(CodePredictorWeights * cw, const GGUFModel & gf, ggml_backend_t backend) {
|
|
cw->hidden_size = (int) gf_get_u32(gf, "qwen3-tts.code_pred.embedding_length");
|
|
cw->intermediate_size = (int) gf_get_u32(gf, "qwen3-tts.code_pred.feed_forward_length");
|
|
cw->num_hidden_layers = (int) gf_get_u32(gf, "qwen3-tts.code_pred.block_count");
|
|
cw->num_attention_heads = (int) gf_get_u32(gf, "qwen3-tts.code_pred.attention.head_count");
|
|
cw->num_key_value_heads = (int) gf_get_u32(gf, "qwen3-tts.code_pred.attention.head_count_kv");
|
|
cw->head_dim = (int) gf_get_u32(gf, "qwen3-tts.code_pred.attention.key_length");
|
|
cw->vocab_size = (int) gf_get_u32(gf, "qwen3-tts.code_pred.vocab_size");
|
|
cw->rope_theta = gf_get_f32(gf, "qwen3-tts.code_pred.rope.freq_base");
|
|
cw->rms_norm_eps = gf_get_f32(gf, "qwen3-tts.code_pred.attention.layer_norm_rms_epsilon");
|
|
|
|
int num_code_groups = (int) gf_get_u32(gf, "qwen3-tts.num_code_groups");
|
|
if (num_code_groups <= 1) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: invalid num_code_groups=%d\n", num_code_groups);
|
|
return false;
|
|
}
|
|
cw->num_acoustic_codebooks = num_code_groups - 1;
|
|
|
|
if (cw->num_hidden_layers <= 0 || cw->hidden_size <= 0) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: invalid hyperparameters (layers=%d hidden=%d)\n", cw->num_hidden_layers,
|
|
cw->hidden_size);
|
|
return false;
|
|
}
|
|
|
|
cw->layers.resize((size_t) cw->num_hidden_layers);
|
|
cw->codec_embedding.resize((size_t) cw->num_acoustic_codebooks);
|
|
cw->lm_head.resize((size_t) cw->num_acoustic_codebooks);
|
|
|
|
int n_tensors = 1 // final norm
|
|
+ 2 // mtp_proj weight + bias (when present)
|
|
+ 2 * cw->num_acoustic_codebooks // 15 embeds + 15 heads
|
|
+ cw->num_hidden_layers * 11 // 11 per layer
|
|
+ 8; // headroom
|
|
WeightCtx wctx;
|
|
wctx_init(&wctx, n_tensors);
|
|
|
|
cw->norm_w = gf_load_tensor(&wctx, gf, "code_pred.output_norm.weight");
|
|
|
|
// Optional projection talker_hidden -> predictor_hidden. Absent in
|
|
// checkpoints where talker.hidden_size == code_pred.hidden_size
|
|
// because upstream uses nn.Identity in that case.
|
|
cw->mtp_proj_w = gf_try_load_tensor(&wctx, gf, "code_pred.mtp_proj.weight");
|
|
cw->mtp_proj_b = gf_try_load_tensor(&wctx, gf, "code_pred.mtp_proj.bias");
|
|
|
|
for (int g = 0; g < cw->num_acoustic_codebooks; g++) {
|
|
char name[160];
|
|
snprintf(name, sizeof(name), "code_pred.codec_embd.%d.weight", g);
|
|
cw->codec_embedding[(size_t) g] = gf_load_tensor(&wctx, gf, name);
|
|
snprintf(name, sizeof(name), "code_pred.lm_head.%d.weight", g);
|
|
cw->lm_head[(size_t) g] = gf_load_tensor(&wctx, gf, name);
|
|
}
|
|
|
|
for (int l = 0; l < cw->num_hidden_layers; l++) {
|
|
TalkerLayer & layer = cw->layers[(size_t) l];
|
|
char name[160];
|
|
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.attn_norm.weight", l);
|
|
layer.input_norm_w = gf_load_tensor(&wctx, gf, name);
|
|
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.ffn_norm.weight", l);
|
|
layer.post_attn_norm_w = gf_load_tensor(&wctx, gf, name);
|
|
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.attn_q.weight", l);
|
|
layer.attn.q_proj_w = gf_load_tensor(&wctx, gf, name);
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.attn_k.weight", l);
|
|
layer.attn.k_proj_w = gf_load_tensor(&wctx, gf, name);
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.attn_v.weight", l);
|
|
layer.attn.v_proj_w = gf_load_tensor(&wctx, gf, name);
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.attn_output.weight", l);
|
|
layer.attn.o_proj_w = gf_load_tensor(&wctx, gf, name);
|
|
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.attn_q_norm.weight", l);
|
|
layer.attn.q_norm_w = gf_load_tensor(&wctx, gf, name);
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.attn_k_norm.weight", l);
|
|
layer.attn.k_norm_w = gf_load_tensor(&wctx, gf, name);
|
|
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.ffn_gate.weight", l);
|
|
layer.mlp.gate_proj_w = gf_load_tensor(&wctx, gf, name);
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.ffn_up.weight", l);
|
|
layer.mlp.up_proj_w = gf_load_tensor(&wctx, gf, name);
|
|
snprintf(name, sizeof(name), "code_pred.blk.%d.ffn_down.weight", l);
|
|
layer.mlp.down_proj_w = gf_load_tensor(&wctx, gf, name);
|
|
}
|
|
|
|
if (!wctx_alloc(&wctx, backend)) {
|
|
fprintf(stderr, "[CodePredictor] FATAL: backend allocation failed\n");
|
|
return false;
|
|
}
|
|
cw->weight_ctx = wctx.ctx;
|
|
cw->weight_buf = wctx.buffer;
|
|
|
|
fprintf(stderr,
|
|
"[CodePredictor] Loaded: %d layers, hidden %d, heads %d/%d, head_dim %d, "
|
|
"FFN %d, RoPE theta %.0f, %d acoustic codebooks (vocab %d each), mtp_proj %s\n",
|
|
cw->num_hidden_layers, cw->hidden_size, cw->num_attention_heads, cw->num_key_value_heads, cw->head_dim,
|
|
cw->intermediate_size, (double) cw->rope_theta, cw->num_acoustic_codebooks, cw->vocab_size,
|
|
cw->mtp_proj_w ? "linear" : "identity");
|
|
return true;
|
|
}
|
|
|
|
static void code_predictor_weights_free(CodePredictorWeights * cw) {
|
|
if (cw->weight_buf) {
|
|
ggml_backend_buffer_free(cw->weight_buf);
|
|
cw->weight_buf = NULL;
|
|
}
|
|
if (cw->weight_ctx) {
|
|
ggml_free(cw->weight_ctx);
|
|
cw->weight_ctx = NULL;
|
|
}
|
|
cw->layers.clear();
|
|
cw->codec_embedding.clear();
|
|
cw->lm_head.clear();
|
|
}
|