model : support qwen3 rerank and embeddings

Author: ngxson

convert_hf_to_gguf.py

@@ -3061,6 +3061,64 @@ def prepare_tensors(self):
 class Qwen3Model(Qwen2Model):
     model_arch = gguf.MODEL_ARCH.QWEN3
 
+    # extra logic for rerank models
+    token_false_id: int | None = None
+    token_true_id: int | None = None
+    sep_token_id: int = 0
+    is_tied_embeddings: bool = False
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # a bit hacky, but currently the only way to detect if this is a rerank model
+        readme_path = self.dir_model / "README.md"
+        readme_text = ""
+        if readme_path.exists():
+            with readme_path.open("r", encoding="utf-8") as f:
+                readme_text = f.read()
+        if "# Qwen3-Reranker" in readme_text:
+            self._find_rerank_config()
+
+    def _find_rerank_config(self):
+        from transformers import AutoTokenizer
+        tokenizer = AutoTokenizer.from_pretrained(self.dir_model)
+        self.token_false_id = tokenizer.convert_tokens_to_ids("no")
+        self.token_true_id = tokenizer.convert_tokens_to_ids("yes")
+        self.sep_token_id = tokenizer.convert_tokens_to_ids("\\n") # unused, but needed for rerank check
+        self.is_tied_embeddings = self.hparams.get("tie_word_embeddings", False)
+        logger.info(f"gguf: token_false_id = {self.token_false_id}, token_true_id = {self.token_true_id}")
+        logger.info(f"gguf: sep_token_id = {self.sep_token_id}")
+        logger.info(f"gguf: is_tied_embeddings = {self.is_tied_embeddings}")
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        is_rerank = self.token_false_id is not None and self.token_true_id is not None
+        if is_rerank:
+            self.gguf_writer.add_pooling_type(gguf.PoolingType.RANK)
+            self.gguf_writer.add_sep_token_id(self.sep_token_id)
+            self.gguf_writer.add_uint32(gguf.Keys.Classifier.OUTPUT_LABELS, 2)
+
+    def _get_cls_out_tensor(self, data_torch: Tensor) -> Tensor:
+        # extract "yes" and "no" tokens from the output lm_head tensor
+        assert self.token_false_id is not None and self.token_true_id is not None
+        false_row = data_torch[self.token_false_id]
+        true_row = data_torch[self.token_true_id]
+        return torch.stack([true_row, false_row], dim=0)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        is_rerank = self.token_false_id is not None and self.token_true_id is not None
+
+        if is_rerank:
+            if self.is_tied_embeddings and "embed_tokens" in name:
+                return [
+                    (gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.CLS_OUT] + ".weight", self._get_cls_out_tensor(data_torch)),
+                    (self.map_tensor_name(name), data_torch),
+                ]
+            if not self.is_tied_embeddings and "lm_head" in name:
+                # this is the lm_head tensor, we need to extract the cls_out tensor
+                return [(gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.CLS_OUT] + ".weight", self._get_cls_out_tensor(data_torch))]
+
+        return super().modify_tensors(data_torch, name, bid)
+
 
 @ModelBase.register("Qwen3MoeForCausalLM")
 class Qwen3MoeModel(Qwen2MoeModel):

src/llama-arch.cpp

@@ -629,6 +629,7 @@ static const std::map<llm_arch, std::map<llm_tensor, const char *>> LLM_TENSOR_N
             { LLM_TENSOR_TOKEN_EMBD,      "token_embd" },
             { LLM_TENSOR_OUTPUT_NORM,     "output_norm" },
             { LLM_TENSOR_OUTPUT,          "output" },
+            { LLM_TENSOR_CLS_OUT,         "cls.output" }, // rerank
             { LLM_TENSOR_ATTN_NORM,       "blk.%d.attn_norm" },
             { LLM_TENSOR_ATTN_Q,          "blk.%d.attn_q" },
             { LLM_TENSOR_ATTN_Q_NORM,     "blk.%d.attn_q_norm" },

src/llama-graph.cpp

@@ -1577,10 +1577,15 @@ void llm_graph_context::build_pooling(
                         cur = ggml_add(ctx0, ggml_mul_mat(ctx0, cls_out, cur), cls_out_b);
                     }
                 } else if (cls_out) {
-                    // Single layer classification head (direct projection)
-                    // https://github.com/huggingface/transformers/blob/f4fc42216cd56ab6b68270bf80d811614d8d59e4/src/transformers/models/bert/modeling_bert.py#L1476
-                    GGML_ASSERT(cls_out_b != nullptr);
-                    cur = ggml_add(ctx0, ggml_mul_mat(ctx0, cls_out, inp), cls_out_b);
+                    if (arch == LLM_ARCH_QWEN3) {
+                        cur = ggml_mul_mat(ctx0, cls_out, inp);
+                        cur = ggml_soft_max(ctx0, cur); // qwen3 uses softmax on the output
+                    } else {
+                        // Single layer classification head (direct projection)
+                        // https://github.com/huggingface/transformers/blob/f4fc42216cd56ab6b68270bf80d811614d8d59e4/src/transformers/models/bert/modeling_bert.py#L1476
+                        GGML_ASSERT(cls_out_b != nullptr);
+                        cur = ggml_add(ctx0, ggml_mul_mat(ctx0, cls_out, inp), cls_out_b);
+                    }
                 } else {
                     GGML_ABORT("RANK pooling requires either cls+cls_b or cls_out+cls_out_b");
                 }

src/llama-model.cpp

@@ -819,7 +819,13 @@ void llama_model::load_hparams(llama_model_loader & ml) {
             } break;
         case LLM_ARCH_QWEN3:
             {
+                // default for embeddings, will be overwritten if model is rerank
+                hparams.pooling_type = LLAMA_POOLING_TYPE_LAST;
+
                 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+                ml.get_key(LLM_KV_POOLING_TYPE,                hparams.pooling_type);
+                ml.get_arr_n(LLM_KV_CLASSIFIER_OUTPUT_LABELS,  hparams.n_cls_out, false);
+
                 switch (hparams.n_layer) {
                     case 28: type = hparams.n_embd == 1024 ? LLM_TYPE_0_6B : LLM_TYPE_1_7B; break;
                     case 36: type = hparams.n_embd == 2560 ? LLM_TYPE_4B : LLM_TYPE_8B; break;
@@ -2463,6 +2469,9 @@ bool llama_model::load_tensors(llama_model_loader & ml) {
                 {
                     tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
 
+                    // output rerank
+                    cls_out = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, 2}, TENSOR_NOT_REQUIRED);
+
                     // output
                     output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
                     output      = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);