Fix FastSentenceTransformer Qwen embedding preprocessing by Etherll · Pull Request #6939 · unslothai/unsloth

Etherll · 2026-07-07T13:43:52Z

Summary

Fixes FastSentenceTransformer module construction for saved SentenceTransformers models such as Qwen/Qwen3-Embedding-0.6B. The loader now preserves the saved text-only Transformer module config instead of
letting SentenceTransformers infer Qwen chat/message preprocessing.

Motivation

FastSentenceTransformer.from_pretrained("Qwen/Qwen3-Embedding-0.6B") was producing embeddings that diverged from a plain SentenceTransformer load. With SentenceTransformers 5.x, constructing the module
through Transformer(...) can infer Qwen's chat template path, so plain embedding strings get encoded as chat messages.

Fixes #6881.

Changes

In unsloth/models/sentence_transformer.py, use Transformer.load(...) when the model has modules.json.
Pass token, cache_dir, and revision into _create_transformer_module() so the module loader uses the same Hub/cache context as _load_modules().
Keep the direct Transformer(...) fallback for non-SentenceTransformers models that do not have modules.json.

How to test

Run a parity check for Qwen3 embeddings:

import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from unsloth import FastSentenceTransformer

texts = [
    "The capital of France is Paris.",
    "A fast brown fox jumps over the lazy dog.",
    "Qwen embedding models use last token pooling.",
]

plain = SentenceTransformer(
    "Qwen/Qwen3-Embedding-0.6B",
    device="cuda",
    model_kwargs={"torch_dtype": torch.float16},
)
fast = FastSentenceTransformer.from_pretrained(
    "Qwen/Qwen3-Embedding-0.6B",
    dtype=torch.float16,
    load_in_4bit=False,
    load_in_16bit=True,
)

a = plain.encode(texts, normalize_embeddings=True)
b = fast.encode(texts, normalize_embeddings=True)

cos = (a * b).sum(axis=1) / (np.linalg.norm(a, axis=1) * np.linalg.norm(b, axis=1))
print(cos)

Expected result: cosine values should be near 1.0. In the Colab check for this fix, mean_cos was 0.99999699 and max_abs_diff was 0.000305.

Etherll · 2026-07-07T13:44:14Z

@codex review

gemini-code-assist · 2026-07-07T13:45:42Z

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 885ebabaab

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

for more information, see https://pre-commit.ci

Etherll · 2026-07-07T14:11:19Z

@codex review

chatgpt-codex-connector · 2026-07-07T14:17:38Z

Codex Review: Didn't find any major issues. Delightful!

Reviewed commit: 2910675678

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

danielhanchen · 2026-07-08T23:41:11Z

Reproduced this on the reported stack (transformers 5.5.0, sentence-transformers 5.5.0, torch 2.9.1+cu128, Qwen/Qwen3-Embedding-0.6B) and confirmed the fix.

Root cause: in sentence-transformers 5.x, building the module with Transformer(model_name, ...) runs modality inference and adds a message modality for any model whose tokenizer has a chat template (Qwen3). Plain embedding strings then get routed through apply_chat_template, so they are wrapped in <|im_start|>user ... <|im_end|> before pooling. Transformer.load(...) keeps the saved module config (modality_config = ['text']), so inputs stay plain text and match a stock SentenceTransformer load.

Tokenization of "roasted chickpeas in 20 kg bags":

stock ST: [299, 15036, 30763, 375, 300, 304, 220, 17, 15, 20972, 17899, 151643] (12 tokens, text + <|endoftext|>)
before fix: [151644, 872, 198, ..., 151645, 198] (16 tokens, chat-wrapped)
after fix: identical to stock ST

Parity vs a stock SentenceTransformer load (8 texts, normalized, fp16):

	mean cos	min cos	max abs diff
before	0.852	0.698	0.105
after	0.999998	0.999997	0.0004

Also confirmed both paths load Qwen3Model (base model, not ForCausalLM) with the final norm module present, so this is purely the preprocessing path and not a model-class or norm-handling issue.

I pushed a small commit adding an inline note so the Transformer.load branch does not get simplified back to Transformer(...) later. Good to merge.

…ssion tests

danielhanchen · 2026-07-09T00:31:24Z

Pushed a follow-up to make this forwards and backwards compatible and to add regression coverage so it cannot silently come back.

Hardening _create_transformer_module

The previous gate required Transformer.load to expose all four of token / cache_folder / revision / trust_remote_code, and called inspect.signature(Transformer.load) unconditionally. Two edge cases:

if a future sentence-transformers renames one of those kwargs, the gate goes false and we silently fall back to Transformer(...), re-introducing the chat-wrap bug;
on a very old ST without Transformer.load, inspect.signature(Transformer.load) raises AttributeError.

It now prefers Transformer.load whenever it exists and passes only the kwargs its installed signature accepts (or all of them if it takes **kwargs). A renamed kwarg can no longer disable the fix, and older ST without .load still loads via the fallback. max_seq_length is re-applied afterwards as before.

Tests

tests/version_compat/test_sentence_transformers_pinned_symbols.py: a source-grep tripwire across ST tags v5.0.0 to master (added v5.5.1 and v5.6.0) that fails if Transformer.load drops the hub kwargs the fix passes.
tests/python/test_fast_sentence_transformer_embedding_parity.py:
- a fast live check that the installed Transformer.load still accepts those kwargs (runs wherever sentence-transformers is importable);
- an end-to-end parity test asserting identical tokenization and min cosine > 0.99 between a stock SentenceTransformer and FastSentenceTransformer, opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is untouched.

To exercise the parity test in CI, add a step that sets UNSLOTH_EMBEDDING_PARITY_MODEL=Qwen/Qwen3-Embedding-0.6B (or any smaller cached chat-template embedder).

Verification (transformers 5.5.0, sentence-transformers 5.5.0, torch 2.9.1+cu128, Qwen/Qwen3-Embedding-0.6B):

48 passed with UNSLOTH_EMBEDDING_PARITY_MODEL set (parity test runs); 1 passed + 1 skipped without it (default CI, 0.02s, no download).
Standalone re-confirm: tokenization identical, mean cosine 0.99999690.
Cross-checked the mechanism on sentence-transformers 5.1.1, 5.5.0, 5.6.0 and under transformers 4.57.6: the modality inference appears from ST 5.2+ (5.1.1 has no modality_config), and Transformer.load restores plain text on all of them.

for more information, see https://pre-commit.ci

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 0232d5ed36

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

chatgpt-codex-connector · 2026-07-09T00:39:32Z

+                    }
+                    if not accepts_var_kw:
+                        load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params}
+                    transformer_module = Transformer.load(model_name, **load_kwargs)


Keep legacy Transformer.load off remote SentenceTransformer ids

When users have sentence-transformers 3.x/4.x installed and load a Hub SentenceTransformer model with modules.json, this branch still reaches Transformer.load(model_name, ...) after filtering the unsupported kwargs down for the legacy single-argument load(input_path) signature. That legacy loader reads sentence_bert_config.json from the supplied local directory, so a remote id like Qwen/Qwen3-Embedding-0.6B is treated as a filesystem path and fails before the previous Transformer(...) constructor fallback can run; pyproject.toml still leaves sentence-transformers unbounded. Fresh evidence in this revision is that the kwargs filter avoids the earlier TypeError but still invokes the single-argument loader for remote IDs.

Useful? React with 👍 / 👎.

Right, and this is fixed in 291412d: the Hub-capable gate (accepts token/cache_folder/revision or **kwargs) means legacy sentence-transformers 3.x/4.x (single-arg Transformer.load(input_path)) no longer reaches the loader at all and falls back to the Transformer(...) constructor, so a remote id is never passed to the local-only legacy loader. That revision landed after the commit reviewed here (0232d5e).

…without Hub-capable load

danielhanchen · 2026-07-09T00:43:34Z

@codex review

danielhanchen · 2026-07-09T00:43:46Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 291412da6b

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-07-09T00:52:02Z

+    unsupported = [
+        k
+        for k in ("token", "cache_folder", "revision", "trust_remote_code")
+        if not (accepts_var_kw or k in params)


Mirror the legacy Transformer.load fallback in the test

When the installed sentence-transformers is 3.x/4.x, Transformer.load has the legacy load(input_path) signature, which the production path deliberately treats as non-hub-capable and falls back to Transformer(...). Because sentence-transformers is still unpinned, this always-on test will fail in those supported environments before the fallback can be exercised; either skip this assertion for legacy signatures or mirror the hub_capable check.

Useful? React with 👍 / 👎.

Fixed in 629316f. The tripwire now mirrors the production hub_capable gate: it skips on a legacy load(input_path) signature (where the code falls back to Transformer(...)) and only asserts the kwargs on a Hub-capable load. Confirmed it passes on ST 5.5.0 and skips for the legacy single-arg signature.

…wire test

danielhanchen · 2026-07-09T01:16:47Z

@codex review

chatgpt-codex-connector · 2026-07-09T01:22:54Z

Codex Review: Didn't find any major issues. Already looking forward to the next diff.

Reviewed commit: 629316f3b5

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

danielhanchen · 2026-07-09T02:31:43Z

@codex review

chatgpt-codex-connector · 2026-07-09T02:37:53Z

Codex Review: Didn't find any major issues. 🚀

Reviewed commit: b69437c848

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

…sformer requires CUDA

danielhanchen · 2026-07-09T05:30:37Z

@codex review

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a1689577da

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

chatgpt-codex-connector · 2026-07-09T05:42:45Z

+                        }
+                        if not accepts_var_kw:
+                            load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params}
+                        transformer_module = Transformer.load(model_name, **load_kwargs)


Honor transformer module paths when loading configs

When modules.json stores the Transformer module under a non-empty path such as 0_Transformer, _load_modules has that path but this call discards it and loads Transformer from the repo root. Transformer.load therefore looks for sentence_bert_config.json/modality_config at the root instead of in the module directory, so chat-template embedding checkpoints saved in that layout still fall back to inferred message preprocessing and lose the parity this change is meant to restore. Pass the module path through as the subfolder/load path when calling Transformer.load.

Useful? React with 👍 / 👎.

Good catch, fixed in def20fb. The Transformer module's saved path from modules.json is now threaded into Transformer.load as subfolder, so sentence_bert_config.json and modality_config resolve from the module directory (for example 0_Transformer) like stock ST. A root path stays a no-op, so single-module checkpoints such as Qwen3-Embedding are unchanged (verified: identical tokenization, min cosine > 0.99).

chatgpt-codex-connector · 2026-07-09T05:42:45Z

+                        }
+                        if not accepts_var_kw:
+                            load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params}
+                        transformer_module = Transformer.load(model_name, **load_kwargs)


Preserve lowercasing on the load path

For saved ST models on sentence-transformers >=5.4, this call sends do_lower_case through **transformer_kwargs, but Transformer.load only forwards its named load/model/processor/config kwargs and drops arbitrary constructor kwargs. The later transformer_module.do_lower_case = ... assignment does not update the tokenizer normalizer that ST 5.4+ relies on, so models that require the forced lowercasing path will tokenize differently than the old Transformer(..., do_lower_case=True) path. Ensure do_lower_case is applied after loading, not just stored as an attribute.

Useful? React with 👍 / 👎.

On the load path do_lower_case is not lost. Transformer.load reads it from the saved sentence_bert_config.json (it is a Transformer.config_keys entry) and init applies it to the tokenizer backend normalizer on ST 5.4+ (transformer.py lines 679-696), so lowercasing is preserved and matches stock ST; the value passed through **kwargs there is redundant. The constructor fallback still receives do_lower_case and honors it. Decoder embedders like Qwen3-Embedding do not lowercase and use fast tokenizers with no basic_tokenizer, so no properly saved checkpoint tokenizes differently on this path.

modules.json records a path for the Transformer module (root for decoder embedders like Qwen3-Embedding, 0_Transformer for the classic layout). Pooling/Normalize already load from their saved path; thread the same path into Transformer.load as subfolder so config and tokenizer resolve like stock ST. stays a no-op, so single-module models are unchanged.

danielhanchen · 2026-07-09T05:58:43Z

@codex review

chatgpt-codex-connector · 2026-07-09T06:04:55Z

Codex Review: Didn't find any major issues. What shall we delve into next?

Reviewed commit: def20fb869

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

fp16 overflows to NaN on bf16-native embedders such as EmbeddingGemma (Gemma3), producing a false parity failure. Prefer bf16 when the GPU supports it so the tripwire can guard the full documented embedding matrix (Qwen3-Embedding, EmbeddingGemma, BGE-M3, all-MiniLM, GTE-ModernBERT), not just fp16-safe models.

danielhanchen · 2026-07-09T06:18:51Z

@codex review

chatgpt-codex-connector · 2026-07-09T06:27:00Z

Codex Review: Didn't find any major issues. 🎉

Reviewed commit: 94fc2729c4

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Etherll requested a review from danielhanchen as a code owner July 7, 2026 13:43

chatgpt-codex-connector Bot reviewed Jul 7, 2026

View reviewed changes

Comment thread unsloth/models/sentence_transformer.py Outdated

Fix FastSentenceTransformer Qwen embedding preprocessing

5353c59

Etherll force-pushed the ST-fix branch from 885ebab to 5353c59 Compare July 7, 2026 14:09

[pre-commit.ci] auto fixes from pre-commit.com hooks

2910675

for more information, see https://pre-commit.ci

Merge branch 'main' into ST-fix

4fe53aa

This was referenced Jul 7, 2026

[fix]: respect auto_model param in FastLlamaModel.from_pretrained for embedding models #6886

Closed

[Bug] FastSentenceTransformer silently degrades Qwen3-Embedding quality (recall@50: 54.3 → 39.7 on our eval) #6881

Closed

Document Transformer.load embedding modality fix for unslothai#6881

8fde5c5

Harden unslothai#6881 fix and add forwards/backwards-compatible regre…

7fb9aba

…ssion tests

[pre-commit.ci] auto fixes from pre-commit.com hooks

0232d5e

for more information, see https://pre-commit.ci

chatgpt-codex-connector Bot reviewed Jul 9, 2026

View reviewed changes

Fall back to Transformer constructor on legacy sentence-transformers …

291412d

…without Hub-capable load

chatgpt-codex-connector Bot reviewed Jul 9, 2026

View reviewed changes

Mirror legacy sentence-transformers fallback in embedding-parity trip…

629316f

…wire test

Tighten unslothai#6881 comments and docstrings

b69437c

Skip embedding-parity test on CPU-only runners since FastSentenceTran…

a168957

…sformer requires CUDA

chatgpt-codex-connector Bot reviewed Jul 9, 2026

View reviewed changes

danielhanchen mentioned this pull request Jul 9, 2026

Retry the Studio UI shutdown re-login on transient goto timeout #7027

Merged

danielhanchen merged commit 5e43c62 into unslothai:main Jul 9, 2026
46 of 47 checks passed

This was referenced Jul 22, 2026

fix/pr 7333 persisted plugin repro ci Imagineer99/unsloth#48

Closed

fix/pr 7328 positional repro ci Imagineer99/unsloth#49

Closed

Uh oh!

Conversation

Etherll commented Jul 7, 2026

Summary

Motivation

Changes

How to test

Uh oh!

Etherll commented Jul 7, 2026

Uh oh!

gemini-code-assist Bot commented Jul 7, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Etherll commented Jul 7, 2026

Uh oh!

chatgpt-codex-connector Bot commented Jul 7, 2026

Uh oh!

danielhanchen commented Jul 8, 2026

Uh oh!

danielhanchen commented Jul 9, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

danielhanchen Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

danielhanchen commented Jul 9, 2026

Uh oh!

danielhanchen commented Jul 9, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

danielhanchen Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

danielhanchen commented Jul 9, 2026

Uh oh!

chatgpt-codex-connector Bot commented Jul 9, 2026

Uh oh!

danielhanchen commented Jul 9, 2026

Uh oh!

chatgpt-codex-connector Bot commented Jul 9, 2026

Uh oh!

danielhanchen commented Jul 9, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

chatgpt-codex-connector Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

danielhanchen Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

chatgpt-codex-connector Bot Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

danielhanchen Jul 9, 2026

Choose a reason for hiding this comment

Uh oh!

danielhanchen commented Jul 9, 2026

Uh oh!

chatgpt-codex-connector Bot commented Jul 9, 2026

Uh oh!

danielhanchen commented Jul 9, 2026