Skip to content

feat(qwen3): support reranker on candle backend#886

Open
malaiwah wants to merge 1 commit into
huggingface:mainfrom
malaiwah:codex/qwen3-reranker-candle
Open

feat(qwen3): support reranker on candle backend#886
malaiwah wants to merge 1 commit into
huggingface:mainfrom
malaiwah:codex/qwen3-reranker-candle

Conversation

@malaiwah

Copy link
Copy Markdown

What does this PR do?

Adds Candle backend support for Qwen3-Reranker using the official yes/no token logit-difference recipe on top of the stock Qwen3 CausalLM weights.

This is adapted from the community work in #835, with two hardening changes from review:

  • preserves the existing model. prefix handling for Qwen3 embedding checkpoints, avoiding the .embed_tokens.weight regression reported on Support Qwen3-Reranker #835
  • only injects the Qwen3-Reranker ChatML prompt template after the model has been resolved as a classifier, so Qwen3-Embedding models with the same Qwen3ForCausalLM architecture are not hijacked

Fixes #763.
Addresses #643 and #691.
Related to #835.

Implementation notes

  • Adds Qwen3ClassificationHead, which precomputes the yes - no LM-head weight vector and applies it to the pooled last-token hidden state.
  • Wires ModelType::Classifier through both qwen3.rs and flash_qwen3.rs.
  • Supports tied word embeddings via embed_tokens and untied weights via lm_head.
  • Fails early for ambiguous checkpoints that include both num_labels/linear.weight and the yes/no LM-head path.
  • Applies reranker prompt templating only for dual inputs when reranker-specific prefix/suffix prompts are present.

Current model metadata caveat

The current public Qwen/Qwen3-Reranker-0.6B config still does not include id2label/label2id, while Qwen3 embedding and reranker configs are otherwise effectively indistinguishable. This PR keeps the classifier detection gated on those mappings and emits an explicit warning when a CausalLM lacks them. Users still need a revision/local config containing id2label/label2id for TEI to resolve it as a reranker end-to-end.

Validation

Run locally on macOS/arm64 CPU/Candle:

cargo check -p text-embeddings-backend-candle --tests
cargo check -p text-embeddings-router --no-default-features --features candle,http
cargo test -p text-embeddings-backend-candle --test test_qwen3 -- --nocapture
cargo test -p text-embeddings-core --features text-embeddings-backend/candle,hf-hub/ureq

Results:

  • test_qwen3 passed, preserving Qwen/Qwen3-Embedding-0.6B behavior
  • test_qwen3_reranker passed against real Qwen/Qwen3-Reranker-0.6B weights, including single inference plus left-padded batch invariants
  • text_embeddings_core::tokenization::tests::tokenizer passed
  • git diff --check clean

cargo fmt --check is blocked in this checkout by missing generated gRPC files (backends/grpc-client/src/pb.rs, router/src/grpc/pb.rs), so I ran rustfmt directly on the changed Rust files.

Before submitting

Who can review?

Anyone in the community is free to review the PR once the tests have passed. @alvarobartt may have context from #763/#691.

Add Qwen3-Reranker support for the Candle Qwen3 implementation by scoring the last-token hidden state with the official yes/no logit-difference recipe.

Wire ModelType::Classifier through qwen3.rs and flash_qwen3.rs, inject the reranker ChatML prompt template only after Qwen3 CausalLM is resolved as a classifier, and template dual inputs without hijacking ordinary embedding prompts.

Add real-weight regression coverage for Qwen3 embedding plus Qwen3 reranker single and left-padded batch inference, and document Qwen3-Reranker in the supported model table.
@malaiwah

Copy link
Copy Markdown
Author

CUDA validation added on an RTX 5090 host.

Environment:

Host: Ubuntu 24.04, NVIDIA GeForce RTX 5090
Driver: 595.58.03; nvidia-smi reports CUDA 13.2
Container: podman + docker.io/nvidia/cuda:12.8.0-devel-ubuntu24.04
CUDA_COMPUTE_CAP=120

I used a validation-only RUSTFLAGS='--cfg feature="dynamic-linking"' in the container so cudarc links dynamically against the CUDA runtime available there. No repository files were changed for this.

Real model assets were fetched with the Hugging Face CLI inside the container:

Qwen/Qwen3-Embedding-0.6B
Qwen/Qwen3-Reranker-0.6B

CUDA + flash-attn compile check:

cargo check -p text-embeddings-backend-candle \
  --features cuda,flash-attn \
  --tests

Result:

Finished `dev` profile [unoptimized + debuginfo] target(s) in 4m 17s

CUDA runtime tests without flash-attn:

cargo test -p text-embeddings-backend-candle \
  --features cuda \
  --test test_qwen3 \
  -- --nocapture

Result:

test test_qwen3 ... ok
test test_qwen3_reranker ... ok
test result: ok. 2 passed; 0 failed; finished in 10.12s

I also tried the same test_qwen3 command with --features cuda,flash-attn. That did not reach the Rust tests: the build failed inside upstream candle-flash-attn / bindgen_cuda while compiling flash-attn CUDA sources for sm_120 with nvcc, including a Segmentation fault (core dumped). Since cargo check --features cuda,flash-attn --tests passes and the non-flash CUDA runtime tests pass against real Qwen3 embedding/reranker weights, this looks like an RTX 5090 / sm_120 flash-attn build-toolchain issue rather than a regression in this PR.

@malaiwah malaiwah marked this pull request as ready for review June 26, 2026 01:34
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Qwen3-Reranker-0.6B start error

1 participant