different with transformers implement

[bjtuxck]

The same input, different results. I don't know which detail is incorrect. The pooling method has been checked.

sentence-transformers 5.1.2
transformers 4.56.1
torch 2.7.1

from sentence_transformers import SentenceTransformer
from sklearn.preprocessing import normalize
from sentence_transformers import SentenceTransformer
from sklearn.preprocessing import normalize
import torch
from torch import nn
import numpy as np
from transformers import AutoModel, AutoTokenizer
from safetensors.torch import load_file
import torch.nn.functional as F

def get_prompteol_input(text: str) -> str:
    return f"This sentence: <|im_start|>“{text}” means in one word: “"

def get_detailed_instruct(task_description: str, query: str) -> str:
    return f'Instruct: {task_description}\nQuery:{query}'

task = "Given a web search query, retrieve relevant passages that answer the query"
queries = [ 
    get_prompteol_input(get_detailed_instruct(task, "光合作用是什么？")),
    get_prompteol_input(get_detailed_instruct(task, "光合作用是什么？"))
]

def main_v1():
    model = SentenceTransformer(
        "/home/jovyan/mnt-xck-notebook-zzzc/pretrained_models/Kingsoft-LLM/QZhou-Embedding-Zh",
        model_kwargs={"device_map": "cuda", "trust_remote_code": True, "torch_dtype": torch.bfloat16},
        tokenizer_kwargs={"padding_side": "left", "trust_remote_code": True},
        trust_remote_code=True,
    )   
    embeddings = model.encode(queries, normalize_embeddings=False)
    print(embeddings)

def main():
    llm = AutoModel.from_pretrained('/home/jovyan/mnt-xck-notebook-zzzc/pretrained_models/Kingsoft-LLM/QZhou-Embedding-Zh', torch_dtype=torch.bfloat16, trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained('/home/jovyan/mnt-xck-notebook-zzzc/pretrained_models/Kingsoft-LLM/QZhou-Embedding-Zh', padding_side='left', trust_remote_code=True, max_length=40960)
    mlp = nn.Linear(4096, 1792, dtype=torch.bfloat16)
    state_dict = load_file('/home/jovyan/mnt-xck-notebook-zzzc/pretrained_models/Kingsoft-LLM/QZhou-Embedding-Zh/2_Dense/model_bak.safetensors')
    mlp.load_state_dict(state_dict)

    llm = llm.to("cuda")
    mlp = mlp.to("cuda")
    query_inputs = tokenizer(queries, padding=True, truncation=True, return_tensors="pt").to("cuda")

    llm.eval()
    mlp.eval()
    with torch.no_grad():
        last_hidden_state = llm(**query_inputs).last_hidden_state
        embedding = last_hidden_state[:, -1]
        embedding = mlp(embedding).cpu().float().numpy()
    print(embedding)

if __name__ == "__main__":
    main_v1()
    main()

[tomaarsen]

Hello!

I just spent some time to check, and it looks like the discrepancies are because of bfloat16 approximations. If you compute the cosine similarity between the embeddings from the two approaches, you'll see 0.9999 for both embeddings, and the differences fully fade away if you use float32 instead of bfloat16.
I would still recommend bfloat16 though, the embeddings are so similar that there should not be notable differences in the eventual rankings.

• Tom Aarsen

[bjtuxck]

Thank you for your reply. I will adopt your suggestion. The core issue of inconsistent results in bfloat16 may be caused by RoPE initialization. The same model initialized twice, and this variable (inv_freq) may be different. Please refer to transformers/modeling_rope_utils.py

def _compute_default_rope_parameters(
    config: Optional[PretrainedConfig] = None,
    device: Optional["torch.device"] = None,
    seq_len: Optional[int] = None,
) -> tuple["torch.Tensor", float]:
    """ 
    Computes the inverse frequencies according to the original RoPE implementation
    Args:
        config ([`~transformers.PretrainedConfig`]):
            The model configuration.
        device (`torch.device`):
            The device to use for initialization of the inverse frequencies.
        seq_len (`int`, *optional*):
            The current sequence length. Unused for this type of RoPE.
    Returns:
        Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
        post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
    """
    base = config.rope_theta
    partial_rotary_factor = config.partial_rotary_factor if hasattr(config, "partial_rotary_factor") else 1.0 
    head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
    dim = int(head_dim * partial_rotary_factor)

    attention_factor = 1.0  # Unused in this type of RoPE

    # Compute the inverse frequencies
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim))
    return inv_freq, attention_factor