v5.6.1 - Flash Attention Fix for XLM-R and RoBERTa Models
This patch release fixes silently degraded embeddings for RoBERTa-family models when flash attention is requested with transformers v5, notably every XLM-R based multilingual embedding model (BAAI/bge-m3, intfloat/multilingual-e5-large, etc.). The bug affected v5.5.0, v5.5.1, and v5.6.0.
Install this version with
# Training + Inference
pip install sentence-transformers[train]==5.6.1
# Inference only, use one of:
pip install sentence-transformers==5.6.1
pip install sentence-transformers[onnx-gpu]==5.6.1
pip install sentence-transformers[onnx]==5.6.1
pip install sentence-transformers[openvino]==5.6.1
# Multimodal dependencies (optional):
pip install sentence-transformers[image]==5.6.1
pip install sentence-transformers[audio]==5.6.1
pip install sentence-transformers[video]==5.6.1
# Or combine as needed:
pip install sentence-transformers[train,onnx,image]==5.6.1Fix position_ids offsetting for RoBERTa-family models when flattening inputs (#3879)
Since v5.5.0, the Transformer module flattens text-only batches into one packed sequence when flash attention is requested, skipping all padding overhead for a notable performance improvement. The position_ids of that packed sequence restart at 0 for every text, which is correct for the vast majority of models. RoBERTa-family architectures however compute positions as padding_idx + 1 + n for the n-th token, so every token read a position embedding shifted by padding_idx + 1 (usually 2). Nothing crashes, the embeddings are just silently worse.
from sentence_transformers import SentenceTransformer
# An affected configuration: flash attention with an XLM-R based model
model = SentenceTransformer(
"BAAI/bge-m3",
model_kwargs={"attn_implementation": "flash_attention_2"},
)Measured on BAAI/bge-m3:
| Evaluation | padded | packed, 0-based positions | packed, with this fix |
|---|---|---|---|
| stsb test Spearman | 0.8485 | 0.7239 | 0.8485 |
| NanoBEIR mean nDCG@10 | 0.6041 | 0.5414 | 0.6050 |
The quality loss recovers exactly once the offset is applied. The fix scans the loaded model's modules once for an int padding_idx stored next to a learned position_embeddings table, and offsets the packed position_ids when that pair is found. An audit of transformers finds 16 architectures with that pair (roberta, xlm_roberta, xlm_roberta_xl, camembert, roberta_prelayernorm, xmod, data2vec_text, longformer, luke, ibert, mpnet, markuplm, lilt, layoutlmv3, esm, and pp_doclayout_v2), all offset by exactly padding_idx + 1, and no 0-based or rotary architecture matches.
You are only affected if you encoded text with flash attention requested on transformers v5 with a RoBERTa-family checkpoint. The default padded path (e.g. sdpa) was never affected, and neither were MPNet models like all-mpnet-base-v2 despite mpnet appearing in the audit: transformers does not support flash attention for MPNet at all. If you did index a corpus with such a configuration, re-encode it after upgrading: pre-fix embeddings score notably worse and do not mix with post-fix embeddings.
What's Changed
- [
ci] Exclude librosa/numba/llvmlite on Python 3.13 by @tomaarsen in #3835 - [
tests] Skip bf16 + Windows + CPU forwards, as they can WindowsError on torch 2.13 by @tomaarsen in #3863 - 🚨 Fix position_ids offsetting for RoBERTa-family models when flattening inputs by @tomaarsen in #3879
Full Changelog: v5.6.0...v5.6.1