From 88adca54301f84595fe5a2109e05ecb3094cf6d2 Mon Sep 17 00:00:00 2001 From: LasseLegarth Date: Sun, 27 Sep 2026 14:27:21 +0200 Subject: [PATCH] cohere: support untied head weight Fine-tunes such as syvai/hviske-v5* train the LM head separately from the token embedding. Detect this in convert-cohere.py and write a dedicated head.weight tensor; at runtime, prefer it over the tied embedding when present. Also accept float32 source tensors and rebuild the mel filterbank from librosa when the checkpoint omits the preprocessor buffer. --- scripts/convert-cohere.py | 27 ++++++++++++++++++++++++--- src/arch/cohere/decoder.cpp | 8 ++++---- src/arch/cohere/weights.cpp | 3 +++ src/arch/cohere/weights.h | 8 ++++++-- 4 files changed, 37 insertions(+), 9 deletions(-) diff --git a/scripts/convert-cohere.py b/scripts/convert-cohere.py index d624e80f..826cf87f 100755 --- a/scripts/convert-cohere.py +++ b/scripts/convert-cohere.py @@ -554,9 +554,9 @@ def add(src_name: str, gguf_name: str, transform) -> None: ) # Read as torch bf16 tensor, convert to fp32 numpy. t = st.get_tensor(src_name) - if t.dtype != torch.bfloat16: + if t.dtype not in (torch.bfloat16, torch.float32): raise ValueError( - f"{src_name}: expected source dtype torch.bfloat16, got {t.dtype}" + f"{src_name}: expected source dtype torch.bfloat16 or torch.float32, got {t.dtype}" ) arr = t.float().numpy() arr = transform(arr) @@ -609,11 +609,31 @@ def add(src_name: str, gguf_name: str, transform) -> None: for src_name, gguf_name, transform in HEAD_TABLE: add(src_name, gguf_name, transform) + # Some fine-tunes (e.g. syvai/hviske-v5*) train the head separately + # from the embedding. Store it as head.weight so the runtime does + # not fall back to the tied embedding. + head_untied = not torch.equal( + st.get_tensor(tied_key), + st.get_tensor("transf_decoder._embedding.token_embedding.weight"), + ) + if head_untied: + print("Head weight is untied from the embedding; writing head.weight") + add(tied_key, "head.weight", passthrough) + # Mel frontend buffers (filterbank + window) — always stored as f32. # These are the exact values the model was trained with; using them # instead of recomputing from scratch eliminates mel-level divergence. fb_src = "preprocessor.featurizer.fb" - fb_tensor = st.get_tensor(fb_src).float().numpy() + if fb_src in st_keys: + fb_tensor = st.get_tensor(fb_src).float().numpy() + else: + # Checkpoints saved without the preprocessor buffer: rebuild it the + # way the remote-code CohereAsrFeatureExtractor does. + import librosa + fb_tensor = librosa.filters.mel( + sr=hp["fe_sample_rate"], n_fft=hp["fe_n_fft"], n_mels=hp["fe_num_mels"], + fmin=0.0, fmax=hp["fe_sample_rate"] / 2, norm="slaney", + ).astype(np.float32) if fb_tensor.ndim == 3: fb_tensor = fb_tensor.squeeze(0) # [1, 128, 257] -> [128, 257] writer.add_tensor("frontend.mel_filterbank", fb_tensor) @@ -645,6 +665,7 @@ def add(src_name: str, gguf_name: str, transform) -> None: + hp["dec_n_layers"] * len(DECODER_BLOCK_TABLE) # 8 * 26 = 208 + len(DEC_FINAL_NORM_TABLE) # 2 + len(HEAD_TABLE) # 1 + + int(head_untied) # optional head.weight + 2 # frontend fb + window ) if n_added != expected: diff --git a/src/arch/cohere/decoder.cpp b/src/arch/cohere/decoder.cpp index 4ef7d112..62c3cbbb 100644 --- a/src/arch/cohere/decoder.cpp +++ b/src/arch/cohere/decoder.cpp @@ -536,7 +536,7 @@ DecoderBuild build_decoder_graph(ggml_context * ctx, // Head: tied weight (the [hidden, vocab_size] token embedding) + bias + // log-softmax. mul_mat(token_w, x[hidden,seq]) -> [vocab_size, seq_len]. - ggml_tensor * logits = ggml_mul_mat(ctx, w.dec_embed.token_w, x); + ggml_tensor * logits = ggml_mul_mat(ctx, w.head.logits_w(w.dec_embed.token_w), x); if (w.head.bias != nullptr) { logits = ggml_add(ctx, logits, w.head.bias); } @@ -745,7 +745,7 @@ DecoderBuild build_decoder_graph_kv(ggml_context * ctx, db.dumps.out_before_head = x; // Head: tied weight + bias + optional log-softmax. - ggml_tensor * logits = ggml_mul_mat(ctx, w.dec_embed.token_w, x); + ggml_tensor * logits = ggml_mul_mat(ctx, w.head.logits_w(w.dec_embed.token_w), x); if (w.head.bias != nullptr) { logits = ggml_add(ctx, logits, w.head.bias); } @@ -883,7 +883,7 @@ StepBuild build_step_graph(ggml_context * ctx, x = layer_norm(ctx, x, w.dec_final.norm_w, w.dec_final.norm_b); // Tied LM head: token embedding transposed. - ggml_tensor * logits = ggml_mul_mat(ctx, w.dec_embed.token_w, x); + ggml_tensor * logits = ggml_mul_mat(ctx, w.head.logits_w(w.dec_embed.token_w), x); if (w.head.bias != nullptr) { logits = ggml_add(ctx, logits, w.head.bias); } @@ -1156,7 +1156,7 @@ StepBuildBatched build_step_graph_batched(ggml_context * ctx, } x = layer_norm(ctx, x, w.dec_final.norm_w, w.dec_final.norm_b); - ggml_tensor * logits = ggml_mul_mat(ctx, w.dec_embed.token_w, x); // [vocab, B] + ggml_tensor * logits = ggml_mul_mat(ctx, w.head.logits_w(w.dec_embed.token_w), x); // [vocab, B] if (w.head.bias != nullptr) { logits = ggml_add(ctx, logits, w.head.bias); } diff --git a/src/arch/cohere/weights.cpp b/src/arch/cohere/weights.cpp index ab486897..d6ce6634 100644 --- a/src/arch/cohere/weights.cpp +++ b/src/arch/cohere/weights.cpp @@ -437,6 +437,9 @@ transcribe_status build_cohere_weights(ggml_context * ctx_meta, const CohereHPar // ----- head ----- GET_F32(weights.head.bias, "head.bias", vocab_size); + if (ggml_get_tensor(ctx_meta, "head.weight") != nullptr) { + GET_LIN(weights.head.weight, "head.weight", dec_h, vocab_size); + } return TRANSCRIBE_OK; } diff --git a/src/arch/cohere/weights.h b/src/arch/cohere/weights.h index ad5bf56a..9b8e6b37 100644 --- a/src/arch/cohere/weights.h +++ b/src/arch/cohere/weights.h @@ -212,9 +212,13 @@ struct CohereDecFinal { ggml_tensor * norm_b = nullptr; }; -// Head (bias only; weight is tied to dec_embed.token_w). +// Head. The weight is tied to dec_embed.token_w unless the GGUF carries an +// untied head.weight (fine-tunes such as syvai/hviske-v5*). struct CohereHead { - ggml_tensor * bias = nullptr; // [vocab_size] + ggml_tensor * weight = nullptr; // optional [dec_hidden, vocab_size] + ggml_tensor * bias = nullptr; // [vocab_size] + + ggml_tensor * logits_w(ggml_tensor * token_w) const { return weight != nullptr ? weight : token_w; } }; struct CohereWeights {