📦 deps(thirdparty): update snapshots
This commit is contained in:
ci[bot]
+79
@@ -0,0 +1,79 @@
|
||||
# Base Model Selection
|
||||
|
||||
Leaderboards rotate every few months; don't trust any hardcoded "best" pick. Discover current options live — run **both** sort orders since most-downloaded surfaces proven options and trending surfaces recent SOTA that may not have download volume yet.
|
||||
|
||||
## Discovery commands
|
||||
|
||||
**[BI]**:
|
||||
```bash
|
||||
hf models list --filter sentence-transformers --sort downloads --limit 20
|
||||
hf models list --filter sentence-transformers --sort trending --limit 20
|
||||
```
|
||||
|
||||
**[CE]**:
|
||||
```bash
|
||||
hf models list --filter sentence-transformers --filter text-ranking --sort downloads --limit 20
|
||||
hf models list --filter sentence-transformers --filter text-ranking --sort trending --limit 20
|
||||
```
|
||||
|
||||
**[SPARSE]**:
|
||||
```bash
|
||||
hf models list --filter sentence-transformers --filter sparse-encoder --sort downloads --limit 20
|
||||
hf models list --filter sentence-transformers --filter sparse-encoder --sort trending --limit 20
|
||||
```
|
||||
|
||||
Optional language narrowing (any type): add `--filter <language-code>`. Not all multilingual models tag each language, so missing matches doesn't mean the model can't handle that language — re-run without the filter to compare.
|
||||
|
||||
```bash
|
||||
hf models card <id> --text # confirm dimensions, max_seq_length, license, languages
|
||||
```
|
||||
|
||||
Cross-check the [MTEB leaderboard](https://huggingface.co/spaces/mteb/leaderboard) (pick the relevant tab) before committing to a multi-hour run.
|
||||
|
||||
## [BI] Bi-Encoder
|
||||
|
||||
Continue from an existing retriever beats fresh-start + 100k–500k pairs. Common namespaces as of 2026-Q2 (verify against discovery commands — the field rotates):
|
||||
|
||||
- **English encoder retrievers**: `sentence-transformers/all-*` (MiniLM-L6-v2, mpnet-base-v2 still the most-downloaded models on the Hub), `BAAI/bge-*-en-v1.5`, `nomic-ai/nomic-embed-text-v1.5`, `mixedbread-ai/mxbai-embed-large-v1`, `Alibaba-NLP/gte-*`, `Snowflake/snowflake-arctic-embed-*`, `jinaai/jina-embeddings-v5-text-small` / `-nano`, `microsoft/harrier-oss-v1-270m` / `-0.6b`.
|
||||
- **Multilingual encoder retrievers**: `sentence-transformers/paraphrase-multilingual-*`, `intfloat/multilingual-e5-*`, `ibm-granite/granite-embedding-*-multilingual-r2`, `google/embeddinggemma-300m`, `voyageai/voyage-4-nano`.
|
||||
- **Long documents (8k+)**: `nomic-ai/modernbert-embed-*`, `answerdotai/ModernBERT-large`.
|
||||
- **Decoder LLM retrievers** (multilingual; **need last-token pooling**): `Qwen/Qwen3-Embedding-*` (0.6B / 4B / 8B), `Qwen/Qwen3-VL-Embedding-*` (multimodal), `codefuse-ai/F2LLM-v2-*`.
|
||||
- **Fresh-start English** (≥500k pairs + domain-fit reason): `microsoft/mpnet-base`, `answerdotai/ModernBERT-base`, `google-bert/bert-base-uncased`, `jhu-clsp/ettin-encoder-*` (17m / 32m / 68m / 150m / 400m / 1b — paired ModernBERT encoder family).
|
||||
- **Fresh-start multilingual**: `FacebookAI/xlm-roberta-base` (MLM-only, needs contrastive training), `microsoft/mdeberta-v3-base`, `jhu-clsp/mmBERT-base` / `-small`.
|
||||
- **CPU / small footprint** (`StaticEmbedding`): `StaticEmbedding(tokenizer, embedding_dim=...)`. **Model size = `vocab_size × dim × 4 bytes`** — pick a small-vocab tokenizer or you get a giant model: 30k-vocab `bert-base-uncased` × 128 dim ≈ 15 MB; **250k-vocab `paraphrase-multilingual-MiniLM-L12-v2` × 256 dim ≈ 256 MB**. Random init needs 1M+ pairs; warm-start (`StaticEmbedding.from_distillation(...)`) helps under ~100k pairs.
|
||||
|
||||
Architecture variants (encoder / decoder / static / Router), pooling rules, and decoder-vs-encoder setup paths: `model_architectures.md`.
|
||||
|
||||
**ModernBERT-family bases default to `max_seq_length=8192`.** That allocates activation memory for 8192-token sequences regardless of your data length and silently drives Windows VRAM into "shared memory" spillover. After loading any ModernBERT / mmBERT / Ettin / gte-modernbert / nomic-modernbert base, **explicitly set `model.max_seq_length = 256` (or 512 for documents)** unless you actually need long context.
|
||||
|
||||
## [CE] Cross-Encoder
|
||||
|
||||
Continue from an existing reranker beats fresh-start + 100k–500k pairs in most domains; default to this unless you have a strong reason otherwise. Common namespaces as of 2026-Q2:
|
||||
|
||||
- **English encoder rerankers**: `cross-encoder/ms-marco-*`, `BAAI/bge-reranker-*`, `mixedbread-ai/mxbai-rerank-*-v1` / `-v2`, `Alibaba-NLP/gte-reranker-modernbert-*`, `ibm-granite/granite-embedding-reranker-english-*`.
|
||||
- **Multilingual encoder rerankers**: `cross-encoder/mmarco-*`, `BAAI/bge-reranker-v2-m3`, `Alibaba-NLP/gte-multilingual-reranker-*`, `ibm-granite/granite-embedding-reranker-multilingual-*`.
|
||||
- **Decoder LLM rerankers** (multilingual; `num_labels=1` last-token-style scoring): `Qwen/Qwen3-Reranker-*` (0.6B / 4B / 8B), `Qwen/Qwen3-VL-Reranker-*` (multimodal).
|
||||
- **Fresh-start**: `microsoft/MiniLM-L12-H384-uncased`, `answerdotai/ModernBERT-base` / `-large`, `jhu-clsp/ettin-encoder-*`, `FacebookAI/xlm-roberta-base` (multilingual), `microsoft/mdeberta-v3-base` (multilingual), `jhu-clsp/mmBERT-base` / `-small` (multilingual). Pass `num_labels >= 2` for classification cross-encoders.
|
||||
|
||||
Encoder-only bases are still the latency-efficient default (bidirectional attention is well-suited to the reranking use case at small parameter counts), but decoder LLM rerankers are now competitive at the top of MTEB Reranking when latency / memory budget allows.
|
||||
|
||||
**Minimum dataset:** 500k+ labeled `(query, passage, label)` tuples for production; 10k–100k labeled pairs for continue-training on domain data. Low-resource languages may have less than 10k labeled pairs; in that case lean on a multilingual base's pretraining and accept a noisier signal.
|
||||
|
||||
**"Small" multilingual is ~100M+ params**, not 17M-50M like the English small models. mMiniLMv2-L12-H384 (~117M) is roughly the small-end for usable multilingual rerankers.
|
||||
|
||||
## [SPARSE] Sparse Encoder (SPLADE)
|
||||
|
||||
SPLADE requires a fill-mask / `AutoModelForMaskedLM`-compatible checkpoint. Encoder-only MLM models work out of the box; **decoder LLMs do not**.
|
||||
|
||||
- **Continue from existing SPLADE — English**: `naver/splade-*` (the canonical family), `opensearch-project/opensearch-neural-sparse-encoding-*` (incl. `-doc-v2-distill`, `-doc-v3-distill` / `-doc-v3-gte`), `prithivida/Splade_PP_en_v*`, `ibm-granite/granite-embedding-30m-sparse`.
|
||||
- **Continue from existing SPLADE — multilingual**: `opensearch-project/opensearch-neural-sparse-encoding-multilingual-v1`.
|
||||
- **Fresh-start English** (≥500k pairs): any encoder with an MLM head — `distilbert/distilbert-base-uncased`, `google-bert/bert-base-uncased`. Pure `AutoModel` checkpoints without MLM won't work. Discover MLM bases: `hf models list --filter fill-mask --sort downloads --limit 20`.
|
||||
- **Fresh-start multilingual**: `FacebookAI/xlm-roberta-base` (has MLM head). For other multilingual MLM bases: add `--filter <language-code>`.
|
||||
|
||||
**Minimum dataset:** 500k+ triplets (with mined hard negatives) for a competitive SPLADE; 50k+ triplets for domain adaptation on existing SPLADE.
|
||||
|
||||
## Cross-cutting tips
|
||||
|
||||
- **Non-English retrieval starting points** (when language tag returns 0 results): check `intfloat/multilingual_e5_train_data` for parallel pair data; MIRACL via the `sentence-transformers/miracl` mirror for multilingual retrieval; mMARCO via `unicamp-dl/mmarco` (14 languages, parquet-backed).
|
||||
- **Avoid script-based dataset loaders.** `datasets >= 4` rejects them with `RuntimeError: Dataset scripts are no longer supported`. Look for parquet-backed mirrors (e.g. `sentence-transformers/miracl` instead of `miracl/miracl`).
|
||||
- **`hf datasets sql` requires DuckDB** (`pip install duckdb`). Without it, fall back to `python -c "from datasets import load_dataset; ds = load_dataset('<id>', ...); print(ds.column_names, ds[0])"`.
|
||||
+128
@@ -0,0 +1,128 @@
|
||||
# Dataset Formats
|
||||
|
||||
This reference covers: how datasets map to losses, how to reshape data when it doesn't fit, and how to mine hard negatives.
|
||||
|
||||
## The two rules
|
||||
|
||||
From the sentence-transformers training overview:
|
||||
|
||||
1. If the loss requires a label, the dataset must have a column named **`label`, `labels`, `score`, or `scores`**. Any column with one of these names is the label.
|
||||
2. All other columns are **inputs**. The loss defines how many input columns it expects. Column **names don't matter; order does**.
|
||||
|
||||
Example: `CoSENTLoss` expects 2 inputs + a float label. A dataset with columns `["premise", "hypothesis", "score"]` works. A dataset with `["score", "premise", "hypothesis"]` does not — reorder first.
|
||||
|
||||
## Per-loss data shapes
|
||||
|
||||
The per-type loss references (`losses_sentence_transformer.md`, `losses_cross_encoder.md`, `losses_sparse_encoder.md`) are the canonical mappings from data shape to loss. Cross-cutting recipe bits those tables don't show:
|
||||
|
||||
- **`CosineSimilarityLoss`** wants `score` normalized to `[0, 1]`; `CoSENTLoss` / `AnglELoss` are pairwise-ranking and ignore absolute scale, so on `stsb` (raw 0-5) divide by 5 only when using cosine-similarity.
|
||||
- **`BatchAllTripletLoss` / `BatchHardTripletLoss` / `BatchSemiHardTripletLoss`** need `batch_sampler=BatchSamplers.GROUP_BY_LABEL` so multiple samples per label appear in the same batch.
|
||||
- **`MSELoss` (distillation)** label is the teacher's full embedding vector (a list of floats), not a scalar score.
|
||||
- **`MarginMSELoss` (distillation)** label is `teacher_score(q, pos) - teacher_score(q, neg)`, precomputed per row.
|
||||
- **N-tuple shape** for MNRL `(anchor, positive, negative_1, negative_2, ..., negative_N)` (1-indexed) is produced by `mine_hard_negatives(..., output_format="n-tuple")`; "labeled-list" output_format produces the CrossEncoder listwise shape.
|
||||
|
||||
## Reshaping operations
|
||||
|
||||
If your data doesn't fit the loss's expected shape:
|
||||
|
||||
### Reorder columns
|
||||
|
||||
```python
|
||||
# Columns are ["hypothesis", "premise", "score"] but CoSENTLoss expects premise first.
|
||||
dataset = dataset.select_columns(["premise", "hypothesis", "score"])
|
||||
```
|
||||
|
||||
### Rename label column
|
||||
|
||||
```python
|
||||
# Your label is called "relevance" but ST wants "label".
|
||||
dataset = dataset.rename_column("relevance", "label")
|
||||
```
|
||||
|
||||
### Drop extra columns
|
||||
|
||||
```python
|
||||
# ST will treat every non-label column as an input. Drop metadata.
|
||||
dataset = dataset.remove_columns(["source_id", "created_at", "language"])
|
||||
```
|
||||
|
||||
### Convert dtypes
|
||||
|
||||
```python
|
||||
# Label is str, need float for CoSENTLoss.
|
||||
dataset = dataset.map(lambda x: {"label": float(x["label"])})
|
||||
```
|
||||
|
||||
## Hard-negative mining
|
||||
|
||||
`mine_hard_negatives` (in `sentence_transformers.util`) produces a training dataset with mined negatives using a retriever. Hard negatives are the single highest-leverage lever for retrieval-model quality.
|
||||
|
||||
### Basic usage
|
||||
|
||||
```python
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from sentence_transformers.util import mine_hard_negatives
|
||||
|
||||
retriever = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
|
||||
|
||||
mined = mine_hard_negatives(
|
||||
dataset=train_pairs, # has (anchor, positive) or (q, a) columns
|
||||
model=retriever,
|
||||
num_negatives=5,
|
||||
range_min=0, range_max=100, # rank window to sample hard negatives from
|
||||
sampling_strategy="top", # "top" = rank-1 hardest; "random" = random in window
|
||||
output_format="n-tuple", # "triplet" | "n-tuple" | "labeled-pair" | "labeled-list"
|
||||
use_faiss=True,
|
||||
)
|
||||
```
|
||||
|
||||
### Output formats
|
||||
|
||||
- `"triplet"` — `(anchor, positive, negative)` triplets. One row per `(query, negative)` pair.
|
||||
- `"n-tuple"` — `(anchor, positive, negative_1, negative_2, ..., negative_N)` (1-indexed) — one row per query.
|
||||
- `"labeled-pair"` — `(anchor, text, label)` with `label=1` for positives and `label=0` for negatives. Good for `BinaryCrossEntropyLoss`.
|
||||
- `"labeled-list"` — `(anchor, texts, labels)` — one row per query with a list of candidates. Good for listwise losses.
|
||||
|
||||
### Filtering false negatives
|
||||
|
||||
If the retriever returns "negatives" that are actually relevant, they become false negatives and hurt training. Filter them:
|
||||
|
||||
```python
|
||||
mined = mine_hard_negatives(
|
||||
dataset=train_pairs,
|
||||
model=retriever,
|
||||
cross_encoder=CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2"), # score candidates
|
||||
num_negatives=5,
|
||||
max_score=0.9, # drop candidates scoring above 0.9
|
||||
relative_margin=0.05, # require neg_score < pos_score * (1 - 0.05)
|
||||
absolute_margin=0.2, # require neg_score < pos_score - 0.2
|
||||
output_format="n-tuple",
|
||||
use_faiss=True,
|
||||
)
|
||||
```
|
||||
|
||||
Use **either** `relative_margin` or `absolute_margin`, usually not both. `max_score` is independently useful as a hard ceiling.
|
||||
|
||||
### CLI
|
||||
|
||||
`scripts/mine_hard_negatives.py` is a CLI wrapper — see it for a ready-to-run command.
|
||||
|
||||
## Choosing the right `range_min` / `range_max`
|
||||
|
||||
`range_max=None` is the default; pass an integer to cap how far down the ranked list to sample from.
|
||||
|
||||
- `range_min=0`, `range_max=100` — sample from the top-100 retrieved. Good default.
|
||||
- `range_min=10`, `range_max=100` — skip the top-10 (often contains true positives). Safer if you lack a cross-encoder.
|
||||
- `range_min=0`, `range_max=1000` — wider net, more diverse negatives, slower.
|
||||
- `sampling_strategy="top"` — always pick the rank-1 hardest. Maximum training signal per row.
|
||||
- `sampling_strategy="random"` — pick randomly within the range. More robust if your retriever is itself noisy.
|
||||
|
||||
## Quick Hub-side dataset checks
|
||||
|
||||
`hf datasets sql "SELECT * FROM 'hf://datasets/<id>/<split>' LIMIT 5"` streams rows via DuckDB without `load_dataset(...)` — the fastest way to confirm column names match your loss before a full validation run. `hf datasets info <id>` shows config / splits / size; `hf datasets card <id> --text` renders the README.
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **`remove_unused_columns=True` (default)**: the trainer drops columns that aren't passed to the model's forward. Usually fine, but if you rely on a custom collator that uses metadata columns, set `remove_unused_columns=False`.
|
||||
- **Floats stored as strings after CSV load**: `load_dataset("csv", ...)` keeps columns as strings by default. Cast with `.map(lambda x: {"label": float(x["label"])})`.
|
||||
- **Mined hard negatives with `include_positives=True`** include the positive as a negative in the output list — only useful when you're building an evaluator or want to measure rank of the positive. For training, leave it `False`.
|
||||
+116
@@ -0,0 +1,116 @@
|
||||
# Evaluators (Cross-Encoder)
|
||||
|
||||
All cross-encoder evaluators live in `sentence_transformers.cross_encoder.evaluation`.
|
||||
|
||||
## Choosing the right evaluator
|
||||
|
||||
| Task | Evaluator |
|
||||
|---|---|
|
||||
| Rerank retrieval results (nDCG@k on BM25 top-N) — fast default | `CrossEncoderNanoBEIREvaluator` |
|
||||
| Rerank with custom candidates per query | `CrossEncoderRerankingEvaluator` |
|
||||
| Binary / multi-class pair classification | `CrossEncoderClassificationEvaluator` |
|
||||
| Continuous pair scoring (STS-style) | `CrossEncoderCorrelationEvaluator` |
|
||||
|
||||
Wrap multiple in `SequentialEvaluator` (from `sentence_transformers.base.evaluation`) to track them together:
|
||||
|
||||
```python
|
||||
from sentence_transformers.base.evaluation import SequentialEvaluator
|
||||
evaluator = SequentialEvaluator([nano_beir_eval, custom_rerank_eval])
|
||||
```
|
||||
|
||||
## The default: `CrossEncoderNanoBEIREvaluator`
|
||||
|
||||
Analog of `NanoBEIREvaluator` for rerankers. Takes BM25 top-100 for each NanoBEIR query and measures how well the cross-encoder re-ranks them.
|
||||
|
||||
```python
|
||||
from sentence_transformers.cross_encoder.evaluation import CrossEncoderNanoBEIREvaluator
|
||||
|
||||
evaluator = CrossEncoderNanoBEIREvaluator(
|
||||
dataset_names=["msmarco", "nfcorpus", "nq"], # default: 11 of 13 NanoBEIR datasets (excludes "arguana", "touche2020")
|
||||
batch_size=64,
|
||||
rerank_k=100, # rerank the BM25 top-K
|
||||
)
|
||||
```
|
||||
|
||||
Output key for `metric_for_best_model`: **`eval_NanoBEIR_R100_mean_ndcg@10`**. The `R100` signals "rerank top-100"; if you change `rerank_k`, the prefix changes (e.g. `R50`).
|
||||
|
||||
Each individual dataset contributes `eval_Nano{DatasetName}_R100_ndcg@10` (e.g. `eval_NanoMSMARCO_R100_ndcg@10`) too.
|
||||
|
||||
## Custom reranking with your own candidates
|
||||
|
||||
Use when you have query + positive + distractor candidates that aren't part of NanoBEIR:
|
||||
|
||||
```python
|
||||
from sentence_transformers.cross_encoder.evaluation import CrossEncoderRerankingEvaluator
|
||||
|
||||
samples = [
|
||||
{"query": "...", "positive": ["the gold answer"], "documents": ["...", "...", ...]}
|
||||
for ...
|
||||
]
|
||||
|
||||
evaluator = CrossEncoderRerankingEvaluator(
|
||||
samples=samples,
|
||||
batch_size=64,
|
||||
name="my-rerank",
|
||||
always_rerank_positives=False, # default is True; override to False for realistic eval
|
||||
)
|
||||
```
|
||||
|
||||
- `always_rerank_positives=True` (the library default) forces the positive into the candidate pool even when the retriever missed it. The reranker is graded only on candidates it can actually score, so the metric reflects pure reranker quality.
|
||||
- `always_rerank_positives=False`: the positive is only reranked if it's already in `documents`. If the retriever missed it, the rank counts as N+1. This reflects end-to-end retriever+reranker quality. A positive the retriever missed is lost, regardless of reranker skill.
|
||||
|
||||
Output key: `eval_{name}_ndcg@10`, `eval_{name}_map`, `eval_{name}_mrr@10`.
|
||||
|
||||
## Classification-style cross-encoders
|
||||
|
||||
### `CrossEncoderClassificationEvaluator`
|
||||
|
||||
Works for both binary (`num_labels=1`) and multi-class (`num_labels>=2`) cross-encoders. Branches internally:
|
||||
- `num_labels=1`: binary mode. Sweeps thresholds to report accuracy, F1, precision, recall, and **average_precision** (primary).
|
||||
- `num_labels>=2`: multi-class mode (e.g. NLI: entailment / neutral / contradiction). Reports **f1_macro** (primary), f1_micro, f1_weighted, and per-class precision / recall.
|
||||
|
||||
```python
|
||||
from sentence_transformers.cross_encoder.evaluation import CrossEncoderClassificationEvaluator
|
||||
|
||||
evaluator = CrossEncoderClassificationEvaluator(
|
||||
sentence_pairs=[(premise, hypothesis), ...],
|
||||
labels=[0, 1, 2, ...],
|
||||
batch_size=64,
|
||||
name="nli-dev",
|
||||
)
|
||||
```
|
||||
|
||||
Output keys (binary, `num_labels=1`): `eval_{name}_accuracy`, `eval_{name}_f1`, `eval_{name}_average_precision` (primary).
|
||||
Output keys (multi-class, `num_labels>=2`): `eval_{name}_f1_macro` (primary), `eval_{name}_f1_micro`, `eval_{name}_f1_weighted`.
|
||||
|
||||
### `CrossEncoderCorrelationEvaluator`
|
||||
|
||||
For continuous-score cross-encoders (like an STS cross-encoder outputting a similarity score). Reports Pearson/Spearman vs. gold scores.
|
||||
|
||||
```python
|
||||
from sentence_transformers.cross_encoder.evaluation import CrossEncoderCorrelationEvaluator
|
||||
|
||||
evaluator = CrossEncoderCorrelationEvaluator(
|
||||
sentence_pairs=[(a, b), ...],
|
||||
scores=[0.4, 0.8, ...],
|
||||
name="stsb-dev",
|
||||
)
|
||||
```
|
||||
|
||||
Output keys: `eval_{name}_spearman`, `eval_{name}_pearson`.
|
||||
|
||||
## Writing `metric_for_best_model`
|
||||
|
||||
Pattern: `f"eval_{evaluator.primary_metric}"`. Inspect after construction: `print(evaluator.primary_metric)`. Common values:
|
||||
- `eval_NanoBEIR_R100_mean_ndcg@10` — `CrossEncoderNanoBEIREvaluator` default
|
||||
- `eval_{name}_ndcg@10` — `CrossEncoderRerankingEvaluator`
|
||||
- `eval_{name}_average_precision` — `CrossEncoderClassificationEvaluator` (binary, `num_labels=1`)
|
||||
- `eval_{name}_f1_macro` — `CrossEncoderClassificationEvaluator` (multi-class, `num_labels>=2`)
|
||||
- `eval_{name}_spearman` — `CrossEncoderCorrelationEvaluator`
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **Always run `evaluator(model)` once before training** — pre-training baseline. Tiny post-training delta means the loss/data/base is wrong.
|
||||
- `CrossEncoderClassificationEvaluator` accepts both `num_labels=1` (binary, primary `average_precision`) and `num_labels>=2` (multi-class, primary `f1_macro`); `CrossEncoderCorrelationEvaluator` requires `num_labels=1`.
|
||||
- The default `dataset_names=None` excludes `arguana` and `touche2020` (Argument-Retrieval task differs from the rest); pass `dataset_names=list(DATASET_NAME_TO_HUMAN_READABLE)` from `sentence_transformers.cross_encoder.evaluation.nano_beir` to actually run all 13.
|
||||
- Subset NanoBEIR datasets during training (3–4) to keep eval cheap; run the broader set post-training.
|
||||
+151
@@ -0,0 +1,151 @@
|
||||
# Evaluators (Bi-Encoder)
|
||||
|
||||
All bi-encoder evaluators live in `sentence_transformers.sentence_transformer.evaluation`.
|
||||
|
||||
## Choosing the right evaluator
|
||||
|
||||
| Task | Evaluator |
|
||||
|---|---|
|
||||
| Retrieval (nDCG, MRR, Recall) — fast default | `NanoBEIREvaluator` |
|
||||
| Retrieval on your own corpus / qrels | `InformationRetrievalEvaluator` |
|
||||
| STS / continuous similarity | `EmbeddingSimilarityEvaluator` |
|
||||
| Binary classification | `BinaryClassificationEvaluator` |
|
||||
| Triplet accuracy | `TripletEvaluator` |
|
||||
| Reranking (from retrieval candidates) | `RerankingEvaluator` |
|
||||
| MSE vs. teacher (distillation) | `MSEEvaluator`, `MSEEvaluatorFromDataFrame` |
|
||||
| Paraphrase mining | `ParaphraseMiningEvaluator` |
|
||||
| Translation (cross-lingual alignment) | `TranslationEvaluator` |
|
||||
| Label accuracy (classification during training) | `LabelAccuracyEvaluator` |
|
||||
|
||||
Wrap multiple evaluators in `SequentialEvaluator` to track all of them together:
|
||||
|
||||
```python
|
||||
from sentence_transformers.sentence_transformer.evaluation import SequentialEvaluator
|
||||
evaluator = SequentialEvaluator([evaluator1, evaluator2, evaluator3])
|
||||
```
|
||||
|
||||
## The big three
|
||||
|
||||
### `NanoBEIREvaluator` (retrieval)
|
||||
|
||||
Small, fast subset of BEIR. Typically runs in <1 minute on a mid-range GPU. Default choice for retrieval training.
|
||||
|
||||
```python
|
||||
from sentence_transformers.sentence_transformer.evaluation import NanoBEIREvaluator
|
||||
|
||||
evaluator = NanoBEIREvaluator(
|
||||
dataset_names=["msmarco", "nfcorpus", "nq"], # default: all 13 NanoBEIR datasets
|
||||
batch_size=128,
|
||||
show_progress_bar=False,
|
||||
)
|
||||
```
|
||||
|
||||
- Default dataset list covers 13 tasks; pick a subset for speed during training.
|
||||
- Output key for `metric_for_best_model`: **`eval_NanoBEIR_mean_cosine_ndcg@10`** (bi-encoder default = cosine similarity).
|
||||
|
||||
### `EmbeddingSimilarityEvaluator` (STS-style)
|
||||
|
||||
Computes Pearson/Spearman correlation between model cosine similarities and gold labels.
|
||||
|
||||
```python
|
||||
from sentence_transformers.sentence_transformer.evaluation import EmbeddingSimilarityEvaluator
|
||||
from sentence_transformers.util.similarity import SimilarityFunction
|
||||
|
||||
evaluator = EmbeddingSimilarityEvaluator(
|
||||
sentences1=stsb["sentence1"],
|
||||
sentences2=stsb["sentence2"],
|
||||
scores=stsb["score"],
|
||||
main_similarity=SimilarityFunction.COSINE,
|
||||
name="sts-dev",
|
||||
)
|
||||
```
|
||||
|
||||
- `main_similarity` can be `COSINE`, `DOT_PRODUCT`, `EUCLIDEAN`, `MANHATTAN`.
|
||||
- `name` is used in the output key: `eval_sts-dev_spearman_cosine`, `eval_sts-dev_pearson_cosine`, etc.
|
||||
|
||||
### `InformationRetrievalEvaluator` (full retrieval)
|
||||
|
||||
Use when you have your **own** corpus + queries + qrels (not one of the NanoBEIR tasks).
|
||||
|
||||
```python
|
||||
from sentence_transformers.sentence_transformer.evaluation import InformationRetrievalEvaluator
|
||||
|
||||
evaluator = InformationRetrievalEvaluator(
|
||||
queries={qid: query_text for qid, query_text in ...},
|
||||
corpus={doc_id: doc_text for doc_id, doc_text in ...},
|
||||
relevant_docs={qid: {doc_id, ...} for qid in ...}, # qid -> set of relevant doc_ids
|
||||
name="my-retrieval",
|
||||
mrr_at_k=[10],
|
||||
ndcg_at_k=[10],
|
||||
accuracy_at_k=[1, 5, 10],
|
||||
precision_recall_at_k=[1, 5, 10],
|
||||
map_at_k=[100],
|
||||
show_progress_bar=False,
|
||||
batch_size=64,
|
||||
)
|
||||
```
|
||||
|
||||
Output keys: `eval_{name}_cosine_ndcg@10`, `eval_{name}_cosine_mrr@10`, etc.
|
||||
|
||||
Heavy for large corpora — each eval encodes the full corpus. Don't run it every 100 steps. Use `NanoBEIREvaluator` for frequent evaluation during training and reserve full IR for milestones / post-training.
|
||||
|
||||
## Other bi-encoder evaluators
|
||||
|
||||
### `BinaryClassificationEvaluator`
|
||||
|
||||
For labeled pair classification (e.g. duplicate detection, entailment as binary). Reports accuracy, F1, precision/recall, AP. Supports all distance metrics — finds the best threshold per metric.
|
||||
|
||||
### `TripletEvaluator`
|
||||
|
||||
For `(anchor, positive, negative)` triplets. Reports the fraction of triplets where the positive is closer to the anchor than the negative.
|
||||
|
||||
### `RerankingEvaluator`
|
||||
|
||||
For custom re-ranking datasets: you provide candidates per query, the evaluator computes MAP and MRR. Good for measuring retrieval-quality on a held-out set.
|
||||
|
||||
### `MSEEvaluator` / `MSEEvaluatorFromDataFrame`
|
||||
|
||||
For distillation setups. Compares student embeddings against teacher embeddings (or teacher scores), reports MSE.
|
||||
|
||||
### `ParaphraseMiningEvaluator`
|
||||
|
||||
For paraphrase-mining tasks. Given a corpus of labeled paraphrase pairs, computes mining quality (F1 at various thresholds).
|
||||
|
||||
### `TranslationEvaluator`
|
||||
|
||||
For cross-lingual / `make_multilingual`-style alignment checking. Measures whether the student aligns sentences across languages.
|
||||
|
||||
### `LabelAccuracyEvaluator`
|
||||
|
||||
For a `SoftmaxLoss`-trained classifier head. Reports accuracy on held-out data.
|
||||
|
||||
## Writing `metric_for_best_model`
|
||||
|
||||
Pattern: `f"eval_{evaluator.primary_metric}"`. Inspect after construction: `print(evaluator.primary_metric)`. Common values:
|
||||
- `eval_NanoBEIR_mean_cosine_ndcg@10` — `NanoBEIREvaluator`
|
||||
- `eval_sts-dev_spearman_cosine` — `EmbeddingSimilarityEvaluator(name="sts-dev")`
|
||||
- `eval_{name}_cosine_ndcg@10` — `InformationRetrievalEvaluator(name=...)`
|
||||
|
||||
## Multi-dimensional evaluation (Matryoshka)
|
||||
|
||||
For Matryoshka-trained models, evaluate at each target dimension:
|
||||
|
||||
```python
|
||||
per_dim_evaluators = [
|
||||
EmbeddingSimilarityEvaluator(
|
||||
sentences1=..., sentences2=..., scores=...,
|
||||
main_similarity=SimilarityFunction.COSINE,
|
||||
name=f"sts-dev-{dim}",
|
||||
truncate_dim=dim,
|
||||
) for dim in [768, 512, 256, 128, 64]
|
||||
]
|
||||
evaluator = SequentialEvaluator(per_dim_evaluators, main_score_function=lambda scores: scores[0])
|
||||
```
|
||||
|
||||
The first evaluator's score drives `load_best_model_at_end`.
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **Always run `evaluator(model)` once before training** — pre-training baseline. If the post-training delta is tiny, the loss/data/base is wrong.
|
||||
- Don't run `InformationRetrievalEvaluator` with a large corpus (>100k docs) at frequent `eval_steps` — use `NanoBEIREvaluator` during training, reserve full IR for end-of-training.
|
||||
- `greater_is_better=True` is the default; right for nDCG / MRR / accuracy.
|
||||
+121
@@ -0,0 +1,121 @@
|
||||
# Evaluators (Sparse Encoder)
|
||||
|
||||
All sparse-encoder evaluators live in `sentence_transformers.sparse_encoder.evaluation`. They mirror the bi-encoder versions with a `Sparse` prefix and default to **dot product** similarity (cosine on sparse vectors is less meaningful).
|
||||
|
||||
## Choosing the right evaluator
|
||||
|
||||
| Task | Evaluator |
|
||||
|---|---|
|
||||
| Retrieval (nDCG, MRR, Recall) — fast default | `SparseNanoBEIREvaluator` |
|
||||
| Retrieval on your own corpus / qrels | `SparseInformationRetrievalEvaluator` |
|
||||
| STS / continuous similarity | `SparseEmbeddingSimilarityEvaluator` |
|
||||
| Binary classification | `SparseBinaryClassificationEvaluator` |
|
||||
| Triplet accuracy | `SparseTripletEvaluator` |
|
||||
| Reranking (from retrieval candidates) | `SparseRerankingEvaluator` |
|
||||
| MSE vs. teacher (distillation) | `SparseMSEEvaluator` |
|
||||
| Translation (cross-lingual alignment) | `SparseTranslationEvaluator` |
|
||||
| Hybrid BM25 + sparse retrieval | `ReciprocalRankFusionEvaluator` |
|
||||
|
||||
Wrap multiple in `SequentialEvaluator` (from `sentence_transformers.base.evaluation`):
|
||||
|
||||
```python
|
||||
from sentence_transformers.base.evaluation import SequentialEvaluator
|
||||
evaluator = SequentialEvaluator([sparse_nano_beir, my_custom_ir])
|
||||
```
|
||||
|
||||
## The default: `SparseNanoBEIREvaluator`
|
||||
|
||||
Small, fast subset of BEIR adapted for sparse retrieval. Typical runtime <1 minute on a mid-range GPU.
|
||||
|
||||
```python
|
||||
from sentence_transformers.sparse_encoder.evaluation import SparseNanoBEIREvaluator
|
||||
|
||||
evaluator = SparseNanoBEIREvaluator(
|
||||
dataset_names=["msmarco", "nfcorpus", "nq"], # default: all 13 NanoBEIR datasets
|
||||
batch_size=32,
|
||||
show_progress_bar=False,
|
||||
)
|
||||
```
|
||||
|
||||
Output key for `metric_for_best_model`: **`eval_NanoBEIR_mean_dot_ndcg@10`** (sparse defaults to dot product).
|
||||
|
||||
### Sparsity tracking
|
||||
|
||||
Unlike the dense variant, the sparse evaluator also reports **active dimension counts** so you can monitor sparsity during training:
|
||||
|
||||
- `query_active_dims` — non-zero entries per query vector
|
||||
- `document_active_dims` — non-zero entries per document vector
|
||||
|
||||
A healthy SPLADE checkpoint typically shows ~30–50 active dims for queries and ~150–250 for documents. If these drift toward the vocab size (~30k), the FLOPS regularization isn't doing its job — raise `query_regularizer_weight` / `document_regularizer_weight` in `SpladeLoss`.
|
||||
|
||||
## Retrieval on your own corpus
|
||||
|
||||
### `SparseInformationRetrievalEvaluator`
|
||||
|
||||
Same shape as the dense version but operates on sparse vectors internally:
|
||||
|
||||
```python
|
||||
from sentence_transformers.sparse_encoder.evaluation import SparseInformationRetrievalEvaluator
|
||||
|
||||
evaluator = SparseInformationRetrievalEvaluator(
|
||||
queries={qid: text for qid, text in ...},
|
||||
corpus={doc_id: text for doc_id, text in ...},
|
||||
relevant_docs={qid: {doc_id, ...} for qid in ...},
|
||||
name="my-sparse-ir",
|
||||
ndcg_at_k=[10],
|
||||
mrr_at_k=[10],
|
||||
accuracy_at_k=[1, 5, 10],
|
||||
map_at_k=[100],
|
||||
batch_size=32,
|
||||
)
|
||||
```
|
||||
|
||||
Output keys: `eval_{name}_dot_ndcg@10`, `eval_{name}_dot_mrr@10`, etc. Also reports active-dims.
|
||||
|
||||
Heavy for large corpora. Use `SparseNanoBEIREvaluator` during training; reserve full IR for post-training.
|
||||
|
||||
## Hybrid retrieval
|
||||
|
||||
### `ReciprocalRankFusionEvaluator`
|
||||
|
||||
Measures the performance of combining your sparse encoder with BM25 (or any other retriever) via reciprocal-rank fusion. Useful when shipping a hybrid system is the actual deployment target.
|
||||
|
||||
## Other sparse evaluators
|
||||
|
||||
### `SparseEmbeddingSimilarityEvaluator`
|
||||
|
||||
STS-style. Computes Pearson/Spearman between sparse vector similarities and gold labels. Uses dot product by default.
|
||||
|
||||
### `SparseBinaryClassificationEvaluator`
|
||||
|
||||
For labeled pair classification with sparse embeddings.
|
||||
|
||||
### `SparseTripletEvaluator`
|
||||
|
||||
For `(anchor, positive, negative)` triplets — reports fraction where the positive is closer than the negative (by dot product).
|
||||
|
||||
### `SparseRerankingEvaluator`
|
||||
|
||||
For custom re-ranking with sparse embeddings. Same semantics as the dense `RerankingEvaluator`.
|
||||
|
||||
### `SparseMSEEvaluator`
|
||||
|
||||
For distillation setups. Compares sparse student embeddings against teacher outputs.
|
||||
|
||||
### `SparseTranslationEvaluator`
|
||||
|
||||
For cross-lingual / `make_multilingual`-style alignment checking with sparse embeddings.
|
||||
|
||||
## Writing `metric_for_best_model`
|
||||
|
||||
Pattern: `f"eval_{evaluator.primary_metric}"`. Inspect after construction: `print(evaluator.primary_metric)`. Common values:
|
||||
- `eval_NanoBEIR_mean_dot_ndcg@10` — `SparseNanoBEIREvaluator` default
|
||||
- `eval_{name}_dot_ndcg@10` — `SparseInformationRetrievalEvaluator`
|
||||
- `eval_{name}_spearman_dot` — `SparseEmbeddingSimilarityEvaluator`
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **Always run `evaluator(model)` once before training** — confirms the pipeline works (a fill-mask base scores ~0 on retrieval until trained).
|
||||
- Sparse evaluators default to dot product; cosine on sparse vectors isn't meaningful.
|
||||
- Don't compare dense and sparse metrics directly — different scales (cosine ∈ [-1, 1] vs. dot ∈ [0, ∞)).
|
||||
- Always check `query_active_dims` / `document_active_dims` — thousands of active dims per doc means the FLOPS regularizer is mistuned, even if nDCG looks OK.
|
||||
+105
@@ -0,0 +1,105 @@
|
||||
# Hardware Guide
|
||||
|
||||
Training embedding models is memory-bound more often than compute-bound.
|
||||
|
||||
## If you hit OOM
|
||||
|
||||
Try in this order:
|
||||
|
||||
1. **Reduce `per_device_train_batch_size`**. Raise `gradient_accumulation_steps` to keep the effective batch size for regression losses. (For MNRL, effective batch via grad-accum is **not** equivalent — see point 3.)
|
||||
2. **Enable `gradient_checkpointing=True`**. ~30% slower, ~40% less activation memory. Incompatible with `Cached*` losses.
|
||||
3. **Switch to a `Cached*` loss**:
|
||||
- `CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)` — forwards in mini-batches, accumulates the contrastive loss over the full batch. Can simulate batch sizes of 1024+ on a 24GB GPU.
|
||||
- `CachedSpladeLoss(model, loss=..., mini_batch_size=16)` — same trick for sparse.
|
||||
- `CachedGISTEmbedLoss(model, guide_model, mini_batch_size=32)` — GIST variant.
|
||||
4. **Enable PEFT / LoRA** for decoder models >1B. `LoraConfig(r=64, lora_alpha=128, task_type="FEATURE_EXTRACTION")`. See `../scripts/train_sentence_transformer_with_lora_example.py` (docstring covers when to use, hyperparams, QLoRA, sharing).
|
||||
5. **Move to multi-GPU**. See below.
|
||||
6. **Shorten sequences**. If truncating to 128 is already sufficient for your task, set `max_seq_length` on the transformer module.
|
||||
|
||||
## Multi-GPU
|
||||
|
||||
`sentence-transformers` uses `accelerate` under the hood. Distributed training works without code changes.
|
||||
|
||||
### Data parallel (DDP)
|
||||
|
||||
Launch:
|
||||
|
||||
```bash
|
||||
accelerate launch train.py
|
||||
|
||||
# or explicitly:
|
||||
accelerate launch --multi_gpu --num_processes=4 train.py
|
||||
```
|
||||
|
||||
`per_device_train_batch_size` stays per-GPU. Effective batch size scales linearly. MNRL's in-batch negatives remain **per-device**, not global, unless you pass `gather_across_devices=True` to losses that support it (`MultipleNegativesRankingLoss`, `CachedMultipleNegativesRankingLoss`, the symmetric variants, `GISTEmbedLoss`, `CachedGISTEmbedLoss`, `SparseMultipleNegativesRankingLoss`).
|
||||
|
||||
### FSDP / DeepSpeed
|
||||
|
||||
For models >3B, use `accelerate config` to enable FSDP or DeepSpeed ZeRO. Both are supported — `sentence-transformers` doesn't require any code changes, only the launch config.
|
||||
|
||||
```bash
|
||||
accelerate config # interactive; choose FSDP or DeepSpeed
|
||||
accelerate launch train.py
|
||||
```
|
||||
|
||||
With FSDP full-shard: a 7B model trains on 4×24GB GPUs that would OOM on any single one of them.
|
||||
|
||||
**FSDP caveats** (from the [distributed training docs](https://sbert.net/docs/sentence_transformer/training/distributed.html)):
|
||||
- **Evaluators don't run under FSDP** as of writing — the eval hooks call `model.encode()` which FSDP-wrapped modules can't service mid-training. Plan to evaluate post-training on a single-GPU load of the final checkpoint instead, or train with DDP if you need mid-training evaluation.
|
||||
- **Layer wrapping must be specified**, e.g. `fsdp_config={"transformer_layer_cls_to_wrap": "BertLayer"}` (substitute the right layer class for your model: `BertLayer`, `LlamaDecoderLayer`, `Qwen2DecoderLayer`, etc.). Without this FSDP sharding can silently misbehave.
|
||||
- **Slower than DDP** for models that fit on a single GPU — only reach for FSDP when you actually need the memory savings.
|
||||
|
||||
DeepSpeed ZeRO-2/3 is an alternative with its own config; works identically at the `accelerate config` level.
|
||||
|
||||
## Effective batch size for contrastive losses
|
||||
|
||||
For `MultipleNegativesRankingLoss` and its variants, **batch size is a quality knob**, not just a speed knob. Bigger batches = more in-batch negatives = richer gradients.
|
||||
|
||||
The effective pool of in-batch negatives per anchor:
|
||||
|
||||
| Setup | In-batch negatives per anchor |
|
||||
|---|---|
|
||||
| Single GPU, batch 64 | 63 |
|
||||
| 4× DDP, per-device batch 64 | 63 local only by default; 255 with `MultipleNegativesRankingLoss(model, gather_across_devices=True)` |
|
||||
| Single GPU, CachedMNRL, mini_batch 32, batch 256 | 255 |
|
||||
| 4× DDP, CachedMNRL, per-device 256 | 255 local; 1023 with `gather_across_devices=True` |
|
||||
|
||||
For large corpora (retrieval), push toward 512+ effective negatives. For small, clean datasets (STS), 64 is plenty.
|
||||
|
||||
## Precision choice by GPU
|
||||
|
||||
| GPU generation | Recommended |
|
||||
|---|---|
|
||||
| T4, V100, GTX 1xxx, RTX 2xxx | `fp16=True` |
|
||||
| RTX 3xxx, A10G, A100, L4 | `bf16=True` |
|
||||
| RTX 4xxx, H100, B200 | `bf16=True` (or fp8 on H100 via specific kernels — not default) |
|
||||
| Apple M-series / ROCm | MPS/ROCm support is variable; `fp16` or `fp32` most reliable |
|
||||
|
||||
bf16 is more numerically stable and almost always preferred when available.
|
||||
|
||||
## Hugging Face Jobs flavor guide
|
||||
|
||||
Hugging Face Jobs requires a Pro/Team/Enterprise plan. Pricing is approximate and subject to change — see the [Jobs pricing page](https://huggingface.co/docs/huggingface_hub/guides/jobs).
|
||||
|
||||
| Flavor | Memory | Typical use | Est. $/hr |
|
||||
|---|---|---|---|
|
||||
| `cpu-basic` | ~2 GB | Dataset prep, validation, hard-neg mining (small) | <$0.10 |
|
||||
| `cpu-upgrade` | ~4 GB | Same, slightly bigger | $0.10 |
|
||||
| `t4-small` | 16 GB | Demos, MiniLM/DistilBERT with small batches | ~$0.75 |
|
||||
| `t4-medium` | 16 GB | MiniLM / DistilBERT with larger batch | ~$1.50 |
|
||||
| `l4x1` | 24 GB | BERT-base, MPNet, ModernBERT-base | ~$2.50 |
|
||||
| `a10g-small` | 24 GB | BERT-base to BERT-large | ~$3.50 |
|
||||
| `a10g-large` | 48 GB | ModernBERT-large, Qwen3-0.6B | ~$5.00 |
|
||||
| `a10g-largex2` | 96 GB (2× 48GB) | Mid-size multi-GPU | ~$10 |
|
||||
| `a100-large` | 80 GB | Large models or big contrastive batches | ~$10–12 |
|
||||
| `h100` | 80 GB | Biggest single-GPU | ~$12 |
|
||||
| `h100x8` | 640 GB | LLM-scale distributed | ~$96 |
|
||||
|
||||
Defaults by base model:
|
||||
- MiniLM / DistilBERT -> `t4-small`
|
||||
- BERT-base / MPNet / ModernBERT-base -> `a10g-small` or `l4x1`
|
||||
- BERT-large / ModernBERT-large -> `a10g-large`
|
||||
- Qwen3-0.6B decoder base -> `a10g-large`
|
||||
- 1B+ decoder bases with LoRA -> `a10g-large` or `a100-large`
|
||||
|
||||
Always start one flavor **smaller** than you think you need: OOM on Jobs is cheap ($0.50–$5 for a failed run). Underprovisioned is better than overprovisioned for the first attempt. When budgeting `timeout`, add **20–30% buffer** for model loading, checkpoint saving, and Hub push.
|
||||
+173
@@ -0,0 +1,173 @@
|
||||
# Hugging Face Jobs Execution
|
||||
|
||||
Run training on Hugging Face's managed GPUs without provisioning any local infrastructure. The same training script runs locally and on Jobs — this reference covers only the Jobs-specific concerns.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- Hugging Face account with a **Pro, Team, or Enterprise** plan. Jobs are paid.
|
||||
- `HF_TOKEN` with **write** permission. Log in once locally with `hf auth login` (the modern command from the `hf` CLI; the older `huggingface-cli login` still works but is deprecated).
|
||||
- Access to the `hf_jobs()` MCP tool, or the `hf` CLI. If installing the CLI from the hosted script, download it to a temporary file, inspect it, and then run it locally.
|
||||
|
||||
## The three submission paths
|
||||
|
||||
### 1. Inline script via MCP (recommended in Claude Code)
|
||||
|
||||
Pass the full training script as `script`. Dependencies come from the PEP 723 header.
|
||||
|
||||
```python
|
||||
hf_jobs("uv", {
|
||||
"script": """
|
||||
# /// script
|
||||
# requires-python = ">=3.10"
|
||||
# dependencies = ["sentence-transformers[train]>=5.0", "trackio"]
|
||||
# ///
|
||||
|
||||
# <full training script content>
|
||||
""",
|
||||
"flavor": "a10g-large",
|
||||
"timeout": "3h",
|
||||
"secrets": {"HF_TOKEN": "$HF_TOKEN"},
|
||||
})
|
||||
```
|
||||
|
||||
### 2. Script-from-URL via MCP
|
||||
|
||||
Upload the script to the Hub (as a model or dataset repo file) or a Gist, then reference by URL:
|
||||
|
||||
```python
|
||||
hf_jobs("uv", {
|
||||
"script": "https://huggingface.co/USERNAME/scripts/resolve/main/train_bi_encoder.py",
|
||||
"flavor": "a10g-large",
|
||||
"timeout": "3h",
|
||||
"secrets": {"HF_TOKEN": "$HF_TOKEN"},
|
||||
})
|
||||
```
|
||||
|
||||
Local file paths (`./train.py`, `/path/to/train.py`) **do not work** — Jobs run in isolated containers without access to your filesystem.
|
||||
|
||||
### 3. CLI
|
||||
|
||||
```bash
|
||||
hf jobs uv run \
|
||||
--flavor a10g-large \
|
||||
--timeout 3h \
|
||||
--secrets HF_TOKEN \
|
||||
"https://huggingface.co/USERNAME/scripts/resolve/main/train.py"
|
||||
```
|
||||
|
||||
Syntax gotchas:
|
||||
- Command order is `hf jobs uv run`, **not** `hf jobs run uv`.
|
||||
- Flags (`--flavor`, `--timeout`, `--secrets`) go **before** the script URL.
|
||||
- `--secrets` (plural), not `--secret`.
|
||||
|
||||
## Required script modifications for Jobs
|
||||
|
||||
Add these to your `TrainingArguments`:
|
||||
|
||||
```python
|
||||
args = SentenceTransformerTrainingArguments(
|
||||
...,
|
||||
push_to_hub=True,
|
||||
hub_model_id="your-username/my-model",
|
||||
hub_strategy="every_save", # push each checkpoint; timeout-safe
|
||||
save_strategy="steps",
|
||||
save_steps=0.1, # 10 saves/pushes per epoch; scales with dataset size
|
||||
)
|
||||
```
|
||||
|
||||
Why each matters:
|
||||
|
||||
| Argument | Why |
|
||||
|---|---|
|
||||
| `push_to_hub=True` | The Jobs container is destroyed after the job finishes. Without Hub push, all weights are lost. |
|
||||
| `hub_model_id` | Required to identify the destination repo. |
|
||||
| `hub_strategy="every_save"` | Default, but worth being deliberate about on Jobs: each checkpoint is pushed as it's written, so a timeout leaves all completed checkpoints on the Hub. `"end"` only pushes once `trainer.train()` returns, so a timeout loses everything. |
|
||||
| `save_strategy="steps"` + `save_steps=0.1` | Checkpoints must actually be saved for `hub_strategy="every_save"` to push them. Fractional `0.1` = save every 10% of training, auto-scales with dataset size. |
|
||||
|
||||
## Secrets
|
||||
|
||||
Secrets are environment variables injected into the Jobs container. They never appear in logs and are not part of the script.
|
||||
|
||||
| Secret | Required when |
|
||||
|---|---|
|
||||
| `HF_TOKEN` | Always, for Hub push. Also covers Trackio auth. |
|
||||
| `WANDB_API_KEY` | Using `report_to="wandb"`. |
|
||||
| `MLFLOW_TRACKING_URI`, `MLFLOW_TRACKING_TOKEN` | Using MLflow with a remote server. |
|
||||
|
||||
The `$HF_TOKEN` syntax in the job config references the value from your local environment at submission time — the literal string `$HF_TOKEN` is replaced with your token's value. Never hardcode tokens in the script itself.
|
||||
|
||||
Trackio (the default tracker in this skill) uses `HF_TOKEN` for auth, so no extra secrets are needed. Only switch to the W&B / MLflow rows above if you're using those trackers.
|
||||
|
||||
## Timeout
|
||||
|
||||
Default is **30 minutes**, which is too short for almost any real training. Set explicitly:
|
||||
|
||||
```python
|
||||
"timeout": "2h" # 2 hours
|
||||
"timeout": "90m" # 90 minutes
|
||||
"timeout": "1.5h" # 90 minutes
|
||||
"timeout": 7200 # seconds, as integer
|
||||
```
|
||||
|
||||
Rule: **estimated training time × 1.3**. The extra buffer covers model loading, dataset caching, checkpoint saving, and Hub push.
|
||||
|
||||
On timeout, the container is killed immediately. Only data on the Hub (`hub_strategy="every_save"` saves you here) or in persistent volumes survives.
|
||||
|
||||
## Dataset caching
|
||||
|
||||
Hugging Face datasets are cached at `~/.cache/huggingface/datasets` by default — **inside the container**, which is destroyed after the job. Each Jobs run re-downloads the dataset.
|
||||
|
||||
For large datasets (>5 GB), this matters. Options:
|
||||
|
||||
- **Persistent `/data` volume** (Jobs feature, check current documentation): set `HF_DATASETS_CACHE=/data/datasets` so caches persist across jobs.
|
||||
- **Pre-cache locally, push to Hub**: if the dataset is on Hub already, nothing to do. If it's local-only, `dataset.push_to_hub(...)` once so subsequent jobs load from Hub.
|
||||
|
||||
## Monitoring a running job
|
||||
|
||||
```bash
|
||||
hf jobs ps [--all] # running (or all) jobs
|
||||
hf jobs inspect <job-id> # full config + status
|
||||
hf jobs logs <job-id> [--follow|--tail N] # tail or stream
|
||||
hf jobs cancel <job-id>
|
||||
hf jobs hardware # list flavors + hourly rates
|
||||
```
|
||||
|
||||
`hf jobs logs <id> --follow` under `Bash run_in_background` pairs nicely with a `Monitor` watching for the `VERDICT:` line emitted by your training script's verdict block.
|
||||
|
||||
MCP equivalents (signatures may vary by server version — check the actual
|
||||
tool listing): `hf_jobs("ps")`, `hf_jobs("logs", {"job_id": ...})`,
|
||||
`hf_jobs("cancel", {"job_id": ...})`.
|
||||
|
||||
For recurring runs, `hf jobs scheduled uv run "<cron>" <script> ...`
|
||||
schedules; `hf jobs scheduled ps/suspend/delete` manages.
|
||||
|
||||
## Common failures
|
||||
|
||||
### "Model not found on Hub" after a successful-looking run
|
||||
|
||||
The run succeeded but `push_to_hub` was not enabled. The container is gone; the weights are gone.
|
||||
|
||||
Fix: always set `push_to_hub=True` + `hub_model_id=...` + `secrets={"HF_TOKEN": "$HF_TOKEN"}`.
|
||||
|
||||
### Tracker not connecting
|
||||
|
||||
- **Trackio:** `HF_TOKEN` missing or lacks write permission. Add `"secrets": {"HF_TOKEN": "$HF_TOKEN"}` and make sure the token has write access.
|
||||
- **W&B:** `WANDB_API_KEY` missing. Add `"secrets": {"HF_TOKEN": "$HF_TOKEN", "WANDB_API_KEY": "$WANDB_API_KEY"}`.
|
||||
|
||||
### OOM on first step
|
||||
|
||||
Flavor too small. Move up one tier (see `hardware_guide.md`).
|
||||
|
||||
### Training starts but eval hangs forever
|
||||
|
||||
`eval_strategy="steps"` with no `eval_dataset`. Always provide an eval dataset, or set `eval_strategy="no"`.
|
||||
|
||||
### Dataset download times out
|
||||
|
||||
Large dataset or slow cold-cache. Increase `timeout` or pre-cache to a persistent volume.
|
||||
|
||||
### `CachedMultipleNegativesRankingLoss` + `gradient_checkpointing=True` crash
|
||||
|
||||
The cached losses are incompatible with gradient checkpointing. Disable `gradient_checkpointing`.
|
||||
|
||||
After submission, the MCP returns a job ID. Monitor with `hf_jobs("logs", {"job_id": ...})` when you want an update — don't poll in a tight loop. End-to-end submission templates live in `scripts/train_sentence_transformer_example.py` / `scripts/train_cross_encoder_example.py` / `scripts/train_sparse_encoder_example.py`; wrap the script contents in the inline pattern from §1 above.
|
||||
+137
@@ -0,0 +1,137 @@
|
||||
# Cross-Encoder Losses (reranker)
|
||||
|
||||
All losses live in `sentence_transformers.cross_encoder.losses`.
|
||||
|
||||
Cross-encoder losses fall into three families: **pointwise** (score each pair independently), **pairwise** (score pairs of pairs), and **listwise** (score a whole ranked list at once). Plus contrastive/distillation variants for specific setups.
|
||||
|
||||
## Top-line decision table
|
||||
|
||||
| You have | Use | Family |
|
||||
|---|---|---|
|
||||
| `(query, passage, label)` with `label ∈ {0, 1}` or `[0, 1]` | `BinaryCrossEntropyLoss` | Pointwise |
|
||||
| `(query, passage, class_id)` multi-class | `CrossEntropyLoss` | Pointwise |
|
||||
| `(query, passage)` with implicit positives, want contrastive | `CachedMultipleNegativesRankingLoss` | Contrastive |
|
||||
| `(query, passages, labels)` (one row per query, with parallel lists of candidate passages and their relevance scores) | `LambdaLoss` | Listwise |
|
||||
| Same listwise shape, want well-tested simpler loss | `ListNetLoss` or `ListMLELoss` | Listwise |
|
||||
| `(query, positive, negative)` pairwise | `RankNetLoss` | Pairwise |
|
||||
| `(query, passage, teacher_score)` distillation from stronger reranker | `MSELoss` or `MarginMSELoss` | Distillation |
|
||||
|
||||
## Pointwise losses
|
||||
|
||||
### `BinaryCrossEntropyLoss`
|
||||
|
||||
**The default cross-encoder loss.** For `(query, passage, label)` where label is 0/1 or a continuous [0, 1] relevance score.
|
||||
|
||||
```python
|
||||
loss = BinaryCrossEntropyLoss(model, pos_weight=torch.tensor(5.0))
|
||||
```
|
||||
|
||||
- `pos_weight`: upweight positives when your dataset is imbalanced (e.g. 1 positive + N hard negatives). A good default is `pos_weight = num_hard_negatives`.
|
||||
- Handles both binary labels and continuous relevance scores — if you have graded relevance (0.0, 0.5, 1.0) BCE still works.
|
||||
|
||||
### `CrossEntropyLoss`
|
||||
|
||||
Multi-class classification over passages. Use with `CrossEncoder("...", num_labels=N)` where N is the number of classes (e.g. 3 for NLI: entailment / neutral / contradiction).
|
||||
|
||||
## Contrastive losses (for reranker training)
|
||||
|
||||
### `MultipleNegativesRankingLoss` (cross-encoder)
|
||||
|
||||
Contrastive training mirroring bi-encoder MNRL. Applies in-batch negatives: for each `(query, positive)`, every other `positive` in the batch is used as a negative.
|
||||
|
||||
- **Data**: `(query, positive)` or `(query, positive, negative_0, negative_1, ...)`.
|
||||
- For CrossEncoder training, **prefer `BinaryCrossEntropyLoss` with mined hard negatives** as the default; MNRL is the fallback when you only have `(query, positive)` pairs and want to avoid a separate mining pass.
|
||||
|
||||
### `CachedMultipleNegativesRankingLoss`
|
||||
|
||||
Cached variant (GradCache) — decouples per-device batch size from effective in-batch negatives, same as the bi-encoder cached version.
|
||||
|
||||
```python
|
||||
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16)
|
||||
```
|
||||
|
||||
- Incompatible with `gradient_checkpointing=True`.
|
||||
- Key choice for training rerankers with effective batch 256+ on a single GPU.
|
||||
|
||||
## Distillation losses
|
||||
|
||||
> **`activation_fn=nn.Identity()` is mandatory** for all distillation, listwise, and pairwise losses below — only `BinaryCrossEntropyLoss` and `CrossEntropyLoss` tolerate the default `Sigmoid`. The loss sees raw logits during training, but the model's `activation_fn` is applied at evaluation via `predict()`; the default `Sigmoid` (with `num_labels=1`) saturates raw logits >5 to ~1.0 inside `predict()`, silently collapsing eval ranking (training loss looks healthy while nDCG crashes from ~0.59 to ~0.14). Construct as `CrossEncoder("...", num_labels=1, activation_fn=torch.nn.Identity())`. See `troubleshooting.md` ("CrossEncoder eval nDCG crashes after distillation / listwise / pairwise training") for the failure-mode walkthrough.
|
||||
|
||||
### `MSELoss` (cross-encoder)
|
||||
|
||||
Regress the student's output score to the teacher's score. Construct the model with `activation_fn=nn.Identity()` (see callout above).
|
||||
|
||||
- **Data**: `(query, passage, teacher_score)`.
|
||||
- The teacher is typically a larger/stronger cross-encoder. Precompute its scores once, store as the label.
|
||||
|
||||
### `MarginMSELoss` (cross-encoder)
|
||||
|
||||
Regress the **difference** between positive and negative scores against the teacher's margin. Typically gives better distillation than plain MSE.
|
||||
|
||||
- **Data**: `(query, positive, negative, score_diff)` where `score_diff = teacher_score(query, positive) - teacher_score(query, negative)`.
|
||||
- Popular recipe for MS MARCO-style distillation.
|
||||
|
||||
## Listwise losses
|
||||
|
||||
All listwise losses expect the dataset to have per-query **lists** of candidate passages and their scores — typically via a collator that groups rows by query.
|
||||
|
||||
> Listwise and pairwise losses below also require `activation_fn=nn.Identity()` (see the Distillation-section callout above).
|
||||
|
||||
### `LambdaLoss`
|
||||
|
||||
The state-of-the-art listwise ranking loss. Optimizes a surrogate of nDCG via weighted pairwise comparisons.
|
||||
|
||||
**Data shape**: `(query, [doc1...docN], [score1...scoreN])` per row; One query, a list of candidate documents, and a parallel list of relevance scores. Use `mine_hard_negatives(..., output_format="labeled-list", ...)` to build this from `(query, positive)` pairs.
|
||||
|
||||
```python
|
||||
import torch.nn as nn
|
||||
from sentence_transformers.cross_encoder.losses import LambdaLoss, NDCGLoss2PPScheme
|
||||
|
||||
model = CrossEncoder("...", num_labels=1, activation_fn=nn.Identity())
|
||||
loss = LambdaLoss(model, weighting_scheme=NDCGLoss2PPScheme())
|
||||
```
|
||||
|
||||
- Strong default when you have multiple candidates per query and graded relevance.
|
||||
- `weighting_scheme`: `NDCGLoss2PPScheme` (default), `NDCGLoss2Scheme`, `LambdaRankScheme`. The `NDCGLoss2PPScheme` default was shown to reach the strongest performance in the original LambdaLoss paper.
|
||||
|
||||
#### LambdaLoss-specific operational notes
|
||||
|
||||
- **OOM recovery order**: drop `mini_batch_size` first (chunking inside the loss preserves the K-list semantic), then `per_device_train_batch_size` paired with `gradient_accumulation_steps`, then reduce K (the per-query candidate-list length) only as a last resort. Lowering K changes the experiment; the others don't.
|
||||
- **Tiny loss at large K is expected.** With `NDCGLoss2PPScheme`, the loss normalizes by the discount-weighted pair count — at K=128 the loss can scale to ~`1e-4` numerically. That's not "training broken"; read `eval_NanoBEIR_R100_mean_ndcg@10` (or your equivalent) instead of the training loss to judge progress.
|
||||
- For very large K (>=128), the O(K²) weight buffer per query is materialized outside of the forward chunking, so even small `mini_batch_size` may not be enough. Consider scoring strategy: top-K hard negatives often beat random K.
|
||||
|
||||
### `ListNetLoss`
|
||||
|
||||
Treats the ranking as a probability distribution (via softmax) and minimizes cross-entropy against the teacher distribution.
|
||||
|
||||
### `ListMLELoss`
|
||||
|
||||
Maximum-likelihood estimation over permutations. Simpler than LambdaLoss; decent default.
|
||||
|
||||
### `PListMLELoss`
|
||||
|
||||
Position-aware ListMLE — weights higher-ranked items more heavily. Often outperforms plain ListMLE on top-k metrics.
|
||||
|
||||
### `RankNetLoss`
|
||||
|
||||
Pairwise classification: for each pair of candidates, predict which is ranked higher via cross-entropy.
|
||||
|
||||
- Simpler and faster than LambdaLoss.
|
||||
- Scales quadratically with list length; not great for long candidate lists (>20).
|
||||
|
||||
### `ADRMSELoss`
|
||||
|
||||
Alternative listwise formulation (Approx Discounted Rank MSE) from the Rank-DistiLLM paper. Same data shape as `LambdaLoss`. In practice LambdaLoss is the stronger default; `RankNetLoss` was reported as marginally more effective (~0.002 nDCG@10) than ADRMSELoss on the paper's LLM-distillation setup, and LambdaLoss generally beats both.
|
||||
|
||||
Hard-negative mining is essential for any contrastive reranker (random negatives teach nothing). See `dataset_formats.md` (Hard-negative mining section) and `../scripts/mine_hard_negatives.py`.
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **`BinaryCrossEntropyLoss` without `pos_weight`** when you have 5+ hard negatives per positive: the loss under-weights the positive signal. Set `pos_weight=num_hard_negatives`.
|
||||
- **`CachedMultipleNegativesRankingLoss` + `gradient_checkpointing=True`**: crash. Pick one.
|
||||
- **Listwise losses with wildly different list lengths per query**: some losses don't handle ragged lists well. Pad or truncate to a fixed length.
|
||||
- **`MarginMSELoss` without precomputed teacher score diffs**: this loss needs the `score_diff` label column populated from a teacher pass; it does not compute the teacher internally.
|
||||
- **Training a cross-encoder with `num_labels=1` then using `CrossEntropyLoss`**: mismatch — BCE is for num_labels=1, CE is for num_labels>=2.
|
||||
- **Default `Sigmoid` activation under any non-BCE loss**: silently destroys eval ranking. Pass `activation_fn=torch.nn.Identity()` to `CrossEncoder(...)` for distillation, listwise, and pairwise losses (everything except BCE/CE). See the callouts in the Distillation and Listwise sections.
|
||||
- **Custom CE head writing to the wrong feature key**: a custom scoring head must populate `features["scores"]` (not `features["sentence_embedding"]`). Otherwise `CrossEncoder.predict()` raises `KeyError: 'scores'` at inference time, even though training succeeds.
|
||||
- **Loading a CE with a custom-class head fails from a different script**: if you defined a `class ClassifierHead(nn.Module)` inline in `train.py` and saved the model, `modules.json` records `__main__.ClassifierHead`. Loading via `CrossEncoder("path")` from any other entry point raises `ImportError: Module '__main__' does not define a 'ClassifierHead' attribute`. Either move the class into an importable file (`my_pkg/heads.py`), build the same shape from stock ST modules (`Dense + LayerNorm + Dense`), or document that the model is only loadable from the same script.
|
||||
+246
@@ -0,0 +1,246 @@
|
||||
# Bi-Encoder Losses (SentenceTransformer)
|
||||
|
||||
All losses live in `sentence_transformers.sentence_transformer.losses`.
|
||||
|
||||
Losses are grouped by data shape. The #1 rule: **pick a loss that matches your data**, not the other way around.
|
||||
|
||||
## Top-line decision table
|
||||
|
||||
| You have | Use |
|
||||
|---|---|
|
||||
| `(anchor, positive)` pairs | `MultipleNegativesRankingLoss` (or Cached variant for large batches) |
|
||||
| `(anchor, positive, negative)` triplets | `MultipleNegativesRankingLoss` — it handles triplets natively |
|
||||
| `(text1, text2, score)` with `score ∈ [-1, 1]` or `[0, 1]` | `CoSENTLoss` (strongly recommended) |
|
||||
| `(text1, text2, label)` with `label ∈ {0, 1}` | `OnlineContrastiveLoss` |
|
||||
| `(text, class_id)` single-column with integer class | `BatchAllTripletLoss` |
|
||||
| `(query, positive, negative, score_diff)` | `MarginMSELoss` (distillation) |
|
||||
| `(text, teacher_embedding)` | `MSELoss` (embedding distillation) |
|
||||
| Want multiple output dims from one training | Wrap any of the above in `MatryoshkaLoss` |
|
||||
| No labels at all, just sentences | `DenoisingAutoEncoderLoss` or `ContrastiveTensionLossInBatchNegatives` |
|
||||
|
||||
## Contrastive losses (pairs + triplets, no labels)
|
||||
|
||||
### `MultipleNegativesRankingLoss` (MNRL)
|
||||
|
||||
The default bi-encoder loss. Uses **in-batch negatives**: every other `positive` in the batch acts as a negative for the current `anchor`.
|
||||
|
||||
```python
|
||||
loss = MultipleNegativesRankingLoss(model, scale=20.0) # similarity_fct defaults to cos_sim
|
||||
```
|
||||
|
||||
- **Data**: `(anchor, positive)` or `(anchor, positive, negative)`. More columns = more explicit hard negatives per row.
|
||||
- **Scale**: temperature. Default `scale=20.0` multiplies similarities by 20 (equivalent to softmax temperature 0.05). Tune only if cosine similarities end up saturated.
|
||||
- **Critical**: set `batch_sampler=BatchSamplers.NO_DUPLICATES` on training args. Otherwise duplicate anchors create false negatives.
|
||||
- **Tip**: batch size matters a lot — more in-batch negatives = better gradients.
|
||||
|
||||
### `CachedMultipleNegativesRankingLoss`
|
||||
|
||||
Same loss, but with gradient caching (GradCache): forwards in mini-batches but computes the contrastive loss over the full batch. Use this when you want effective batch size of 256+ but your GPU can only fit 32 forwards.
|
||||
|
||||
```python
|
||||
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=32)
|
||||
```
|
||||
|
||||
- **Incompatible with `gradient_checkpointing=True`**.
|
||||
- Set `mini_batch_size` to whatever `per_device_train_batch_size` would be if you couldn't use this. Then crank the actual `per_device_train_batch_size` to what you want the effective batch to be (256+, 1024+).
|
||||
|
||||
### `MultipleNegativesSymmetricRankingLoss`
|
||||
|
||||
MNRL computed bidirectionally — scores positives from both (anchor -> positive) and (positive -> anchor) directions. Slightly better on retrieval tasks where the "anchor" and "positive" distinctions are soft (paraphrase, deduplication).
|
||||
|
||||
### `CachedMultipleNegativesSymmetricRankingLoss`
|
||||
|
||||
Cached variant of the above.
|
||||
|
||||
### `GISTEmbedLoss`
|
||||
|
||||
Like MNRL, but uses a **guide model** (a separate pretrained Sentence Transformer) to **filter out false negatives** before computing the contrastive loss. The guide model scores each potential negative; if it looks too similar to the positive, it's excluded.
|
||||
|
||||
```python
|
||||
guide = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
|
||||
loss = GISTEmbedLoss(model, guide=guide)
|
||||
```
|
||||
|
||||
- Expensive: guide model does a forward pass per batch.
|
||||
- Strong when your in-batch negatives are noisy (e.g. small corpus with many near-duplicates).
|
||||
|
||||
### `CachedGISTEmbedLoss`
|
||||
|
||||
Cached + GIST.
|
||||
|
||||
### `MegaBatchMarginLoss`
|
||||
|
||||
In-batch margin-based triplet: for each anchor, find the hardest negative in the batch and apply a margin loss. Older pattern, usually outperformed by MNRL.
|
||||
|
||||
### `TripletLoss`
|
||||
|
||||
Classic triplet margin loss on explicit `(anchor, positive, negative)`. Uses a fixed margin and the hardest in-batch is not considered — only the provided triplet.
|
||||
|
||||
```python
|
||||
loss = TripletLoss(model, distance_metric=TripletDistanceMetric.EUCLIDEAN, triplet_margin=5)
|
||||
```
|
||||
|
||||
- Simpler than MNRL; less powerful when the batch has useful negatives.
|
||||
- Good for cases where you have *trusted* pre-mined triplets and want to avoid in-batch noise.
|
||||
|
||||
## Batch-triplet losses (single column + integer label)
|
||||
|
||||
These mine triplets *within* the batch from samples sharing a label.
|
||||
|
||||
### `BatchAllTripletLoss`
|
||||
|
||||
For each anchor, form triplets with all positive/negative combinations in the batch. Max signal per batch.
|
||||
|
||||
```python
|
||||
loss = BatchAllTripletLoss(model, margin=5)
|
||||
```
|
||||
|
||||
- **Data**: single text column + integer `label`. Needs multiple samples per label in each batch (set `batch_sampler=BatchSamplers.GROUP_BY_LABEL`).
|
||||
|
||||
### `BatchHardTripletLoss`
|
||||
|
||||
Same, but only the single hardest positive + hardest negative per anchor.
|
||||
|
||||
### `BatchSemiHardTripletLoss`
|
||||
|
||||
Semi-hard mining: negatives harder than the positive but easier than the margin. Often more stable than fully-hard.
|
||||
|
||||
### `BatchHardSoftMarginTripletLoss`
|
||||
|
||||
Variant with a soft margin (log-sum-exp) instead of a fixed margin hinge.
|
||||
|
||||
**When to use batch-triplet losses**: classification-style datasets (labels are class IDs, not pair annotations). E.g. "train an embedder where samples from the same class are close."
|
||||
|
||||
## Scored regression losses (labeled pairs, float score)
|
||||
|
||||
### `CoSENTLoss`
|
||||
|
||||
**The recommended regression loss** for `(text1, text2, score)`. Trains on pairwise ranking: for any two pairs `(a, b)` and `(c, d)` with `score(a,b) > score(c,d)`, the model should score `(a, b)` higher. Much better than squared error.
|
||||
|
||||
```python
|
||||
loss = CoSENTLoss(model, scale=20.0)
|
||||
```
|
||||
|
||||
- **Data**: `(text1, text2, float_score)`. Labels can be `[0, 1]` or `[-1, 1]`.
|
||||
- Works well with small datasets (STS-B, 5k pairs).
|
||||
|
||||
### `AnglELoss`
|
||||
|
||||
Similar to CoSENT but uses angle-based optimization in complex space. Sometimes outperforms CoSENT on tasks with fine-grained similarity gradations. Strong alternative.
|
||||
|
||||
### `CosineSimilarityLoss`
|
||||
|
||||
Squared-error loss on cosine similarity: `mse(cos(text1, text2), label)`. Simpler than CoSENT, usually worse. Keep for legacy / reproducibility.
|
||||
|
||||
## Contrastive labeled losses (labeled pairs, binary label)
|
||||
|
||||
### `ContrastiveLoss`
|
||||
|
||||
For `(text1, text2, label)` where `label ∈ {0, 1}`. Minimizes distance for positives; pushes negatives past a margin.
|
||||
|
||||
```python
|
||||
loss = ContrastiveLoss(model, margin=0.5, distance_metric=SiameseDistanceMetric.COSINE_DISTANCE)
|
||||
```
|
||||
|
||||
### `OnlineContrastiveLoss`
|
||||
|
||||
Same setup but **ignores "easy" pairs** (positives already close, negatives already far) and only optimizes hard ones. Much more robust to label noise.
|
||||
|
||||
```python
|
||||
loss = OnlineContrastiveLoss(model, margin=0.5)
|
||||
```
|
||||
|
||||
**Preferred over `ContrastiveLoss`** for most practical labeled pair datasets.
|
||||
|
||||
### `SoftmaxLoss`
|
||||
|
||||
A classifier head on concatenated `(u, v, |u-v|)` embeddings, trained with cross-entropy. Useful when you have NLI-style multi-class labels (entailment / neutral / contradiction) and want a categorical loss. Historically important (it trained the first popular sentence embedding models) but **generally outperformed by MNRL**.
|
||||
|
||||
## Distillation losses
|
||||
|
||||
### `MSELoss`
|
||||
|
||||
Regress the student's embedding to match a teacher's embedding.
|
||||
|
||||
- **Data**: `(text, teacher_embedding)`. The teacher embedding is a fixed vector per row.
|
||||
- Use when distilling a smaller bi-encoder from a larger one.
|
||||
|
||||
### `MarginMSELoss`
|
||||
|
||||
For `(query, positive, negative, score_diff)`: minimize `mse(student_score_diff, teacher_score_diff)`. The teacher is typically a cross-encoder that produced the score differences.
|
||||
|
||||
- **Data**: 3 text columns + 1 float column.
|
||||
- Workhorse of dense retriever training from cross-encoder teachers (ms-marco distillation).
|
||||
|
||||
### `DistillKLDivLoss`
|
||||
|
||||
KL-divergence distillation: student's softmax distribution over candidates should match teacher's.
|
||||
|
||||
- **Data**: `(query, passages[], teacher_scores[])`.
|
||||
- Good for list-wise distillation when you have multiple candidates per query.
|
||||
|
||||
See `../scripts/train_sentence_transformer_distillation_example.py` for the end-to-end pattern (its docstring covers Embedding MSE / Margin MSE / Listwise KL with full recipes).
|
||||
|
||||
## Regularizer / wrapper losses
|
||||
|
||||
These don't have their own data shape — they wrap another loss and add a regularization objective.
|
||||
|
||||
### `MatryoshkaLoss`
|
||||
|
||||
Train **once**, deploy at any of several dimensions. Wraps any loss and computes it at multiple truncated dimensions, adding them weighted.
|
||||
|
||||
```python
|
||||
base_loss = MultipleNegativesRankingLoss(model)
|
||||
loss = MatryoshkaLoss(
|
||||
model,
|
||||
base_loss,
|
||||
matryoshka_dims=[768, 512, 256, 128, 64],
|
||||
matryoshka_weights=[1, 1, 1, 1, 1], # relative weighting per dim
|
||||
)
|
||||
```
|
||||
|
||||
- At inference, `SentenceTransformer(..., truncate_dim=128)` gives 128-dim output with ~95% of full quality.
|
||||
- Default matryoshka_dims: pick the dims you want to deploy at. Smaller dims improve compression at the cost of slight quality drop.
|
||||
|
||||
### `Matryoshka2dLoss`
|
||||
|
||||
2D-Matryoshka: reduce dimension **and** number of transformer layers in a single wrapper. Internally composes `MatryoshkaLoss` + `AdaptiveLayerLoss`, so you only need this one (don't wrap it in AdaptiveLayerLoss yourself). Deploy at any (dim, layer) pair at inference.
|
||||
|
||||
### `AdaptiveLayerLoss`
|
||||
|
||||
Wrap any loss; adds a term that trains each of the transformer's layers to be a valid exit point. Deploy with fewer layers at inference for faster encoding.
|
||||
|
||||
```python
|
||||
loss = AdaptiveLayerLoss(
|
||||
model,
|
||||
base_loss,
|
||||
n_layers_per_step=1,
|
||||
last_layer_weight=1.0,
|
||||
prior_layers_weight=1.0,
|
||||
)
|
||||
```
|
||||
|
||||
### `GlobalOrthogonalRegularizationLoss` (GOR)
|
||||
|
||||
Stand-alone regularizer (not a wrapper, despite living in this section). Penalizes embedding pairs whose dot product deviates from orthogonality, encouraging the model to spread embeddings across the full vector space. Use it alongside a primary contrastive loss by summing the two outputs in your own training step; can help with downstream retrieval diversity.
|
||||
|
||||
## Unsupervised losses
|
||||
|
||||
### `DenoisingAutoEncoderLoss` (TSDAE)
|
||||
|
||||
Sentence-level denoising autoencoder: corrupt a sentence (drop tokens), force the model to reconstruct it. Pretraining-style — useful for domain adaptation when you have unlabeled in-domain sentences.
|
||||
|
||||
### `ContrastiveTensionLoss`
|
||||
|
||||
Unsupervised contrastive: two copies of the model encode the same sentence; they should agree. Pure self-supervised.
|
||||
|
||||
### `ContrastiveTensionLossInBatchNegatives`
|
||||
|
||||
CT with in-batch negatives. Stronger than vanilla CT.
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **`MultipleNegativesRankingLoss` without `BatchSamplers.NO_DUPLICATES`** will include duplicate anchors in the same batch, destroying training signal. Always set the sampler.
|
||||
- **Any `Cached*` loss + `gradient_checkpointing=True` = crash.** Pick one.
|
||||
- **`TripletLoss` with bad negatives** (too easy) = loss hits zero fast and model stops learning. Mine hard negatives first.
|
||||
- **Matryoshka wrapping `CachedMultipleNegativesRankingLoss`**: supported, but the cached-loss's mini-batch semantics apply to the base loss only. Think twice before combining.
|
||||
+106
@@ -0,0 +1,106 @@
|
||||
# Sparse-Encoder Losses (SPLADE)
|
||||
|
||||
All losses live in `sentence_transformers.sparse_encoder.losses`.
|
||||
|
||||
This reference targets the **SPLADE** architecture (Transformer + SpladePooling). The sparse-encoder package also exports `CSRLoss` and `CSRReconstructionLoss` for the CSR architecture (Transformer + Pooling + SparseAutoEncoder); those are out of scope here — see the sbert.net docs if you're training a CSR model.
|
||||
|
||||
Choosing a loss means (a) pick a base loss (contrastive, regression, distillation) and (b) wrap it in `SpladeLoss` to add FLOPS regularization.
|
||||
|
||||
## Top-line decision table
|
||||
|
||||
| You have | Use |
|
||||
|---|---|
|
||||
| `(anchor, positive)` or triplet, SPLADE architecture | `SpladeLoss(loss=SparseMultipleNegativesRankingLoss(model), ...)` |
|
||||
| Same, want effective batch size of 256+ | `CachedSpladeLoss(...)` |
|
||||
| `(text1, text2, score)` labeled pairs | `SparseCoSENTLoss` or `SparseCosineSimilarityLoss` |
|
||||
| Distillation from cross-encoder teacher | `SparseMarginMSELoss` |
|
||||
| Listwise distillation | `SparseDistillKLDivLoss` |
|
||||
| Explicit triplet | `SparseTripletLoss` |
|
||||
|
||||
## The core wrapper: `SpladeLoss`
|
||||
|
||||
`SpladeLoss` adds **FLOPS regularization** on top of another sparse loss. FLOPS regularization penalizes non-zero activations, keeping embeddings genuinely sparse.
|
||||
|
||||
```python
|
||||
loss = SpladeLoss(
|
||||
model=model,
|
||||
loss=SparseMultipleNegativesRankingLoss(model=model),
|
||||
query_regularizer_weight=5e-5,
|
||||
document_regularizer_weight=3e-5,
|
||||
)
|
||||
```
|
||||
|
||||
- `query_regularizer_weight`: how much to penalize non-zero terms in query embeddings.
|
||||
- `document_regularizer_weight`: same for documents.
|
||||
- Typical range: 1e-5 to 1e-4. Higher = sparser embeddings, lower recall; lower = denser, possibly better recall.
|
||||
- `SparseEncoderTrainer` automatically registers a `SpladeRegularizerWeightSchedulerCallback` whenever the loss is a `SpladeLoss`. The callback ramps the weights from 0 up to the target over the first ~33% of training; the default shape is `SchedulerType.QUADRATIC` (not linear). The ramp length and shape are configured on the callback (`SpladeRegularizerWeightSchedulerCallback(loss=..., warmup_ratio=..., scheduler_type=...)`), not on `SpladeLoss`; to override, instantiate the callback yourself and pass it via `callbacks=[...]`. This ramp is important; starting with full regularization from step 0 kills learning.
|
||||
|
||||
Use `CachedSpladeLoss` for the GradCache variant.
|
||||
|
||||
## Contrastive losses (no labels)
|
||||
|
||||
### `SparseMultipleNegativesRankingLoss`
|
||||
|
||||
Sparse analog of bi-encoder MNRL. In-batch contrastive.
|
||||
|
||||
```python
|
||||
inner = SparseMultipleNegativesRankingLoss(model=model)
|
||||
loss = SpladeLoss(model=model, loss=inner, query_regularizer_weight=5e-5, document_regularizer_weight=3e-5)
|
||||
```
|
||||
|
||||
- **Always wrap in `SpladeLoss`** for SPLADE architectures.
|
||||
- Set `batch_sampler=BatchSamplers.NO_DUPLICATES` on training args.
|
||||
|
||||
### `SparseTripletLoss`
|
||||
|
||||
Classic triplet margin loss on explicit `(anchor, positive, negative)`.
|
||||
|
||||
## Labeled regression losses
|
||||
|
||||
### `SparseCoSENTLoss`
|
||||
|
||||
Pairwise ranking loss for `(text1, text2, score)`. Mirrors bi-encoder `CoSENTLoss`.
|
||||
|
||||
### `SparseCosineSimilarityLoss`
|
||||
|
||||
MSE on cosine similarity. Simpler, usually worse than CoSENT.
|
||||
|
||||
### `SparseAnglELoss`
|
||||
|
||||
Angle-based loss in complex space. Alternative to CoSENT.
|
||||
|
||||
## Distillation losses
|
||||
|
||||
### `SparseMSELoss`
|
||||
|
||||
Embedding MSE. Student sparse embedding should match teacher embedding.
|
||||
|
||||
- **Data**: `(text, teacher_embedding)`.
|
||||
- Teacher can be a dense bi-encoder or another sparse model.
|
||||
|
||||
### `SparseMarginMSELoss`
|
||||
|
||||
Margin MSE from a cross-encoder teacher.
|
||||
|
||||
- **Data**: `(query, positive, negative, score_diff)` where `score_diff = teacher_score(query, positive) - teacher_score(query, negative)`.
|
||||
- Typical recipe for training SPLADE from cross-encoder labels (ms-marco distillation).
|
||||
- Wrap in `SpladeLoss(model, loss=SparseMarginMSELoss(model), ...)` for SPLADE.
|
||||
|
||||
### `SparseDistillKLDivLoss`
|
||||
|
||||
Listwise KL-div distillation — student's softmax distribution over candidates should match teacher's.
|
||||
|
||||
## Independent regularizer
|
||||
|
||||
### `FlopsLoss`
|
||||
|
||||
Standalone FLOPS regularizer. Usually you use this via `SpladeLoss`, not directly.
|
||||
|
||||
For regularizer-weight tuning and dense-output recovery, see `troubleshooting.md` ("SPLADE embeddings are dense"). MLM-head requirement: `base_model_selection.md` (SPARSE section). Active-dim sparsity targets and how to monitor them: `evaluators_sparse_encoder.md` (Sparsity tracking).
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **`SparseMultipleNegativesRankingLoss` without `SpladeLoss` wrapping on a SPLADE model**: no FLOPS regularization -> dense outputs defeating the purpose of SPLADE. Always wrap.
|
||||
- **`CachedSpladeLoss` + `gradient_checkpointing=True`**: crash. Pick one.
|
||||
- **Starting training with full FLOPS regularization at step 0**: the model outputs zero everywhere and gets stuck. The built-in scheduler avoids this — don't override it unless you know why.
|
||||
- **`query_regularizer_weight` == `document_regularizer_weight`**: usually wrong. Queries should be sparser than documents (fewer terms per query). Since higher regularization drives more zeros, give the query weight the larger value. `query_regularizer_weight=5e-5`, `document_regularizer_weight=3e-5` is a good starting ratio.
|
||||
+178
@@ -0,0 +1,178 @@
|
||||
# Model Architectures (SentenceTransformer)
|
||||
|
||||
The `SentenceTransformer` class is a `torch.nn.Sequential` of modules. The common shape is `Transformer` + `Pooling` (+ optional `Normalize` / `Dense`), but four distinct architecture families are supported and the right choice depends on the task.
|
||||
|
||||
## The four architecture families
|
||||
|
||||
| Family | Backbone | Pooling | Use case |
|
||||
|---|---|---|---|
|
||||
| **Encoder (bidirectional)** | BERT, RoBERTa, DeBERTa, MPNet, ModernBERT, XLM-R | `mean` (default) or `cls` | Short/medium text, general default |
|
||||
| **Decoder (causal LLM)** | Qwen, Llama, Mistral, Gemma | `lasttoken` | Long context, instruction-tunable, larger quality ceiling |
|
||||
| **Static embeddings** | `StaticEmbedding` module | N/A | CPU-only, <10MB, extremely fast |
|
||||
| **Multimodal / Router** | VLM backbones or composed encoders | depends | Text + image / audio / video |
|
||||
|
||||
Below, each family with concrete setup.
|
||||
|
||||
## Encoder models (the default)
|
||||
|
||||
The historical default and still usually the right choice for text embeddings.
|
||||
|
||||
```python
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
model = SentenceTransformer("microsoft/mpnet-base")
|
||||
# Auto-constructs: Transformer(feature-extraction) -> Pooling(mean).
|
||||
```
|
||||
|
||||
When `SentenceTransformer("<checkpoint>")` is called with a raw HF encoder, it auto-wraps the transformer and adds `Pooling(..., pooling_mode="mean")`.
|
||||
|
||||
To customize pooling or add modules:
|
||||
|
||||
```python
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from sentence_transformers.sentence_transformer.modules import Normalize, Pooling, Transformer
|
||||
|
||||
transformer = Transformer("answerdotai/ModernBERT-base")
|
||||
pooling = Pooling(transformer.get_embedding_dimension(), pooling_mode="cls") # or "mean", "lasttoken", ...
|
||||
model = SentenceTransformer(modules=[transformer, pooling, Normalize()])
|
||||
```
|
||||
|
||||
**Pooling modes**:
|
||||
- `mean` (default) — average of token embeddings, masked to the attention mask. Strongest default.
|
||||
- `cls` — embedding of the `[CLS]` token. Works if the base was CLS-pretrained.
|
||||
- `max` — element-wise max across tokens. Rare.
|
||||
- `mean_sqrt_len_tokens` — mean scaled by √seq_len. Empirically helps on some tasks.
|
||||
- `weightedmean` — token-position-weighted mean. Useful for decoder bases as a non-last-token alternative.
|
||||
- `lasttoken` — embedding of the last token. Required for causal-LM bases (see decoder section below).
|
||||
|
||||
Don't switch pooling mid-training. Pick once.
|
||||
|
||||
## Decoder / causal LLM models
|
||||
|
||||
Strong at long context, instruction following, multilingual. Memory-hungry — typically LoRA-trained rather than full fine-tuned.
|
||||
|
||||
**Two setup paths depending on whether the model was already adapted for embeddings:**
|
||||
|
||||
```python
|
||||
# Path A: already-adapted embedding checkpoint (ships with the right modules):
|
||||
from sentence_transformers import SentenceTransformer
|
||||
model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B") # just works
|
||||
|
||||
# Path B: raw decoder LLM, build the pipeline manually:
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from sentence_transformers.sentence_transformer.modules import Normalize, Pooling, Transformer
|
||||
|
||||
transformer = Transformer(
|
||||
"Qwen/Qwen2.5-0.5B",
|
||||
transformer_task="text-generation", # critical: causal attention, no bidirectional
|
||||
processor_kwargs={"padding_side": "left"}, # last-token pooling wants left-padding
|
||||
)
|
||||
pooling = Pooling(transformer.get_embedding_dimension(), pooling_mode="lasttoken")
|
||||
model = SentenceTransformer(modules=[transformer, pooling, Normalize()])
|
||||
```
|
||||
|
||||
Skipping `transformer_task="text-generation"` or `pooling_mode="lasttoken"` on a raw decoder gives embeddings that look plausible until you benchmark.
|
||||
|
||||
**Why last-token pooling:** causal attention means only the last token has seen the full sequence. Mean-pooling a causal model averages embeddings that only saw prefixes — the result doesn't represent the whole input.
|
||||
|
||||
For training decoder bases:
|
||||
- Learning rate: typically `1e-4` or higher (not `2e-5` like encoders).
|
||||
- LoRA is almost always the right choice for >1B-param bases; see `../scripts/train_sentence_transformer_with_lora_example.py` (its docstring covers when to use, hyperparams, QLoRA for 7B+, and adapter sharing).
|
||||
|
||||
## Static embeddings
|
||||
|
||||
`StaticEmbedding` skips the transformer entirely — each token maps to a pre-computed vector via a lookup table. No attention, no contextualization.
|
||||
|
||||
**When to use:**
|
||||
- CPU inference, no GPU, browser / edge / on-device deployment.
|
||||
- Need <10MB model size.
|
||||
- Latency budget <1ms per embedding.
|
||||
- Have >1M training pairs (contextualization is replaced by per-token optimization; this takes data).
|
||||
|
||||
**When NOT to use:**
|
||||
- Task needs contextual understanding (polysemy, syntax, long-range dependencies).
|
||||
- You have <100k training pairs — the model won't learn enough.
|
||||
|
||||
**Setup:**
|
||||
|
||||
```python
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from sentence_transformers.sentence_transformer.modules import StaticEmbedding
|
||||
from tokenizers import Tokenizer
|
||||
|
||||
tokenizer = Tokenizer.from_pretrained("google-bert/bert-base-uncased")
|
||||
static_embedding = StaticEmbedding(tokenizer, embedding_dim=512)
|
||||
model = SentenceTransformer(modules=[static_embedding])
|
||||
```
|
||||
|
||||
Train with `MultipleNegativesRankingLoss` on a large contrastive dataset (1M+ pairs).
|
||||
|
||||
**Warm starts vs. random init** — with **>1M training samples**, random-init beats `StaticEmbedding.from_model2vec(...)` or `.from_distillation(...)` warm starts. With smaller datasets, warm starts help.
|
||||
|
||||
```python
|
||||
# For smaller datasets (<100k), warm-start:
|
||||
static_embedding = StaticEmbedding.from_model2vec("minishlab/potion-base-8M")
|
||||
# or:
|
||||
static_embedding = StaticEmbedding.from_distillation("sentence-transformers/all-MiniLM-L6-v2", vocabulary=list(tokenizer.get_vocab().keys()))
|
||||
```
|
||||
|
||||
See `../scripts/train_sentence_transformer_static_embedding_example.py` for a runnable end-to-end recipe (random init + MNRL + Matryoshka + bf16 + lr=2e-1) and the [Static Embeddings blog post](https://huggingface.co/blog/static-embeddings) for benchmarks.
|
||||
|
||||
## Multimodal via VLM backbone
|
||||
|
||||
Modern vision-language models can be loaded directly and produce joint text+image embeddings:
|
||||
|
||||
```python
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
model = SentenceTransformer(
|
||||
"Qwen/Qwen3-VL-Embedding-2B",
|
||||
model_kwargs={"attn_implementation": "flash_attention_2"}, # do NOT set torch_dtype here; see training_args.md
|
||||
processor_kwargs={"min_pixels": 28 * 28, "max_pixels": 600 * 600},
|
||||
)
|
||||
|
||||
# Check which modalities this model supports:
|
||||
print(model.modalities)
|
||||
# ['text', 'image', 'video', 'message']
|
||||
```
|
||||
|
||||
Training data can mix text, PIL images, image paths/URLs, audio, and mixed-modality dicts like `{"image": <PIL>, "text": "describe this"}`. The data collator handles preprocessing via the model's `preprocess` method.
|
||||
|
||||
Install multimodal extras: `pip install "sentence-transformers[image]"` (or `[audio]`, `[video]`).
|
||||
|
||||
**Precision**: load in fp32 and pass `bf16=True` (or `fp16=True`) to TrainingArguments — autocast handles the inference path. Don't set `torch_dtype="bfloat16"` in `model_kwargs`: it puts Adam state in bf16 and silently degrades quality (see `training_args.md`).
|
||||
|
||||
## Multimodal via Router
|
||||
|
||||
Instead of one VLM backbone, compose separate encoders per modality:
|
||||
|
||||
```python
|
||||
from sentence_transformers import SentenceTransformer
|
||||
from sentence_transformers.sentence_transformer.modules import Dense, Pooling, Router, Transformer
|
||||
|
||||
# Text encoder
|
||||
text_encoder = Transformer("sentence-transformers/all-MiniLM-L6-v2")
|
||||
text_pooling = Pooling(text_encoder.get_embedding_dimension(), pooling_mode="mean")
|
||||
# Project text to match image encoder's dim
|
||||
text_projection = Dense(text_encoder.get_embedding_dimension(), 768)
|
||||
|
||||
# Image encoder (SigLIP outputs pooled embeddings directly)
|
||||
image_encoder = Transformer("google/siglip2-base-patch16-224")
|
||||
|
||||
router = Router(
|
||||
sub_modules={
|
||||
"text": [text_encoder, text_pooling, text_projection],
|
||||
"image": [image_encoder],
|
||||
},
|
||||
)
|
||||
model = SentenceTransformer(modules=[router])
|
||||
```
|
||||
|
||||
**Warning**: Router-based models have unaligned embedding spaces at init — you must train to align them. Use a `Dense` projection layer when dimensions differ. Task-based routing (different encoders for queries vs. documents) is also supported via `route_mappings`; see the `Router` docstring.
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **Decoder base with mean pooling**: silently produces garbage embeddings. Always use `lasttoken`.
|
||||
- **Router multimodal without training**: the separate encoders' embedding spaces are unaligned at init. Don't expect useful cross-modal similarity until you've trained with a loss that aligns the spaces.
|
||||
- **StaticEmbedding with fewer than 100k pairs**: the model won't learn enough. Either warm-start via `from_model2vec` / `from_distillation`, or use a regular encoder.
|
||||
- **Large VLM backbones on consumer GPUs**: combine LoRA + `attn_implementation="flash_attention_2"`. With LoRA only, you can additionally pass `torch_dtype="bfloat16"` — the bf16 base weights are frozen, so the Adam-state precision concern from the precision rule above doesn't apply (the LoRA adapter stays fp32, so its optimizer state stays fp32). Without LoRA, follow the precision rule: keep weights fp32 and rely on `bf16=True` autocast.
|
||||
+160
@@ -0,0 +1,160 @@
|
||||
# Prompts and Instructions
|
||||
|
||||
Modern embedding models (E5, BGE, GTE, Qwen3-Embedding, Nomic, Instructor, etc.) use **prompts** / **instructions** at encode time — short prefixes like `"query: "`, `"passage: "`, or `"Represent this sentence for retrieval:"` — to signal task intent.
|
||||
|
||||
In sentence-transformers, prompts work during both **training** and **inference**, and the library keeps them aligned automatically.
|
||||
|
||||
The prompt is literally prepended to the input text before tokenization.
|
||||
|
||||
## Model-level prompts
|
||||
|
||||
Set prompts on the model so `encode()`, `encode_query()`, `encode_document()` use them automatically:
|
||||
|
||||
```python
|
||||
model = SentenceTransformer("your-base-model")
|
||||
model.prompts = {
|
||||
"query": "query: ",
|
||||
"document": "passage: ",
|
||||
}
|
||||
model.default_prompt_name = "document" # used when none is specified
|
||||
```
|
||||
|
||||
At inference:
|
||||
|
||||
```python
|
||||
q_emb = model.encode_query("What is the capital of France?")
|
||||
# Internally: encodes "query: What is the capital of France?"
|
||||
|
||||
d_emb = model.encode_document(["Paris is the capital of France.", "Berlin is the capital of Germany."])
|
||||
# Internally: encodes "passage: Paris..." etc.
|
||||
|
||||
# Explicit prompt_name:
|
||||
emb = model.encode(["some text"], prompt_name="query")
|
||||
```
|
||||
|
||||
## Training with prompts
|
||||
|
||||
Set prompts on the training args so they're applied automatically to the right columns during training. Four shapes, increasingly specific:
|
||||
|
||||
### 1. Single prompt (applied everywhere)
|
||||
|
||||
```python
|
||||
args = SentenceTransformerTrainingArguments(
|
||||
...,
|
||||
prompts="Represent this sentence for similarity: ",
|
||||
)
|
||||
```
|
||||
|
||||
Every input column gets the prefix. Simplest and often sufficient for single-task training.
|
||||
|
||||
### 2. Per-column prompts
|
||||
|
||||
```python
|
||||
prompts = {
|
||||
"anchor": "query: ",
|
||||
"positive": "passage: ",
|
||||
"negative": "passage: ",
|
||||
}
|
||||
```
|
||||
|
||||
Keys are **column names**. Different prefixes for different roles.
|
||||
|
||||
### 3. Per-dataset prompts (multi-dataset)
|
||||
|
||||
```python
|
||||
prompts = {
|
||||
"all-nli": "Classify the entailment relationship: ",
|
||||
"stsb": "Score semantic similarity: ",
|
||||
}
|
||||
```
|
||||
|
||||
Keys are **dataset names** (matching the `train_dataset` dict keys).
|
||||
|
||||
### 4. Per-dataset + per-column
|
||||
|
||||
```python
|
||||
prompts = {
|
||||
"all-nli": "", # all-nli: no prompt
|
||||
"msmarco": { # msmarco: per-column
|
||||
"query": "query: ",
|
||||
"positive": "passage: ",
|
||||
"negative": "passage: ",
|
||||
},
|
||||
}
|
||||
```
|
||||
|
||||
## Cross-encoder specifics
|
||||
|
||||
Cross-encoders restrict prompts to single-value or per-dataset only (no per-column). This is because cross-encoders take text pairs, not separate columns.
|
||||
|
||||
```python
|
||||
args = CrossEncoderTrainingArguments(
|
||||
...,
|
||||
prompts="Rank this passage for the query: ", # single prompt
|
||||
)
|
||||
# or
|
||||
args = CrossEncoderTrainingArguments(
|
||||
...,
|
||||
prompts={"msmarco": "...", "gooaq": "..."}, # per-dataset
|
||||
)
|
||||
```
|
||||
|
||||
## Sparse encoders (SPLADE)
|
||||
|
||||
Sparse encoders support prompts with the same four shapes as bi-encoders (added in v5.x). Same `model.prompts = {...}` API.
|
||||
|
||||
## Pooling and prompt tokens (bi-encoder)
|
||||
|
||||
When using mean/last-token pooling with prompt prefixes, the prompt tokens themselves can either:
|
||||
- **Be included** in the pooled embedding (default behavior). Simpler, what most models do.
|
||||
- **Be excluded** — only the actual content tokens are pooled.
|
||||
|
||||
To exclude, flip `include_prompt` on the Pooling module — found via `isinstance` rather than a fixed index, since multimodal / Router pipelines don't always put Pooling at position 1:
|
||||
|
||||
```python
|
||||
from sentence_transformers.sentence_transformer.modules import Pooling
|
||||
|
||||
for module in model:
|
||||
if isinstance(module, Pooling):
|
||||
module.include_prompt = False
|
||||
```
|
||||
|
||||
Or use the helper if your model exposes one (e.g. `SparseEncoder.set_pooling_include_prompt(False)`).
|
||||
|
||||
**When to exclude**: when the prompt is purely a task signal and you don't want it to dilute the semantic representation. Most E5 / BGE / GTE models leave prompts included.
|
||||
|
||||
During training with `prompts=`, the `include_prompt` setting is respected; the trainer also tracks `include_prompt_lengths` internally so pooling skips the prompt tokens correctly when `include_prompt=False`.
|
||||
|
||||
## Inference alignment
|
||||
|
||||
If you trained with `prompts={"query": "query: ", ...}`, you must:
|
||||
|
||||
1. **Save the prompts on the model**: `model.prompts = args.prompts` before `save_pretrained`, OR let the library do this automatically via the model card data.
|
||||
2. **Use the same prompts at inference**: call `encode_query()` / `encode_document()` and the saved prompts are applied.
|
||||
|
||||
If the saved model's `config_sentence_transformers.json` has `prompts` + `default_prompt_name`, anyone who loads it gets the correct prompts for free.
|
||||
|
||||
## Instruction tuning (different from prompts)
|
||||
|
||||
Some models (Instructor, Qwen3-Embedding) use **instructions** — longer descriptions like `"Represent the biomedical query for retrieving relevant passages:"`. In sentence-transformers these are modeled the same way: just set them as prompts.
|
||||
|
||||
The model-card convention is to list available instructions/prompts in a dict:
|
||||
|
||||
```python
|
||||
model.prompts = {
|
||||
"msmarco-query": "Represent the query for MS MARCO retrieval: ",
|
||||
"msmarco-doc": "Represent the passage for MS MARCO retrieval: ",
|
||||
"sts": "Represent the sentence for semantic similarity: ",
|
||||
"classification": "Represent the sentence for topic classification: ",
|
||||
}
|
||||
```
|
||||
|
||||
Users call `model.encode(["..."], prompt_name="msmarco-query")`.
|
||||
|
||||
## Gotchas
|
||||
|
||||
- **Training with prompts but forgetting to set them on the model before `save_pretrained`**: the saved model won't know about the prompts. Users will encode *without* the prefix and get bad results. Fix: `model.prompts = args.prompts` before saving, or use `SentenceTransformerModelCardData(...)` with the prompt info.
|
||||
- **Using `encode_query()` at inference but didn't train with a "query" prompt**: the method will just call regular `encode` (no prefix applied), which is fine. But the documentation implies you use it, and users may be confused.
|
||||
- **Trailing spaces matter**: `"query: "` vs `"query:"` — the former has a trailing space before the real text. Inspect your training data to know which one your model was trained with.
|
||||
- **Mixing prompted and non-prompted data in multi-dataset training** is fine as long as your `prompts` dict covers it per-dataset (e.g. `{"dataset_a": "query: ", "dataset_b": ""}`).
|
||||
- **`include_prompt=False` with small datasets**: the model may under-fit because you've effectively shortened every input. Usually leave as True.
|
||||
+269
@@ -0,0 +1,269 @@
|
||||
# Training Arguments
|
||||
|
||||
`SentenceTransformerTrainingArguments`, `CrossEncoderTrainingArguments`, and `SparseEncoderTrainingArguments` all inherit from Hugging Face's `TrainingArguments`, so 95% of the arguments are the same.
|
||||
|
||||
This reference covers the arguments that actually matter for embedding-model training.
|
||||
|
||||
## The recommended default set
|
||||
|
||||
Start with this and adjust only what you have a reason to change:
|
||||
|
||||
```python
|
||||
from sentence_transformers import SentenceTransformerTrainingArguments
|
||||
from sentence_transformers.base.sampler import BatchSamplers
|
||||
|
||||
args = SentenceTransformerTrainingArguments(
|
||||
output_dir="models/my-model",
|
||||
|
||||
# Duration
|
||||
num_train_epochs=1,
|
||||
# max_steps=10_000, # alternative to epochs
|
||||
|
||||
# Batch size
|
||||
per_device_train_batch_size=64,
|
||||
per_device_eval_batch_size=64,
|
||||
gradient_accumulation_steps=1,
|
||||
|
||||
# Optimizer
|
||||
learning_rate=2e-5,
|
||||
warmup_steps=0.1, # transformers v5.2 deprecated `warmup_ratio`; pass the ratio as a float directly to `warmup_steps`
|
||||
lr_scheduler_type="linear",
|
||||
weight_decay=0.0,
|
||||
|
||||
# Precision
|
||||
bf16=True, # fp16=True on older GPUs (T4, V100)
|
||||
|
||||
# Sampler (bi-encoder + sparse-encoder)
|
||||
batch_sampler=BatchSamplers.NO_DUPLICATES,
|
||||
|
||||
# Eval + checkpointing
|
||||
eval_strategy="steps",
|
||||
eval_steps=0.1, # fraction: 10 evals/epoch, scales with dataset size
|
||||
save_strategy="steps",
|
||||
save_steps=0.1, # keep aligned with eval_steps for load_best_model_at_end
|
||||
save_total_limit=2,
|
||||
load_best_model_at_end=True,
|
||||
metric_for_best_model="eval_NanoBEIR_mean_cosine_ndcg@10",
|
||||
greater_is_better=True,
|
||||
|
||||
# Logging
|
||||
logging_steps=0.01, # fraction: ~100 log lines/epoch
|
||||
logging_first_step=True,
|
||||
run_name="my-model",
|
||||
report_to="trackio", # or "wandb", "tensorboard", "mlflow", "none"
|
||||
)
|
||||
```
|
||||
|
||||
## Duration
|
||||
|
||||
- `num_train_epochs` — most common. 1 for large datasets (>500k), 3–10 for small.
|
||||
- `max_steps` — use instead of epochs when you want a fixed compute budget. Overrides `num_train_epochs`.
|
||||
- For huge datasets where 1 epoch is overkill, pick `max_steps` matching your compute plan.
|
||||
|
||||
## Batch size
|
||||
|
||||
Effective batch size = `per_device_train_batch_size × num_gpus × gradient_accumulation_steps`.
|
||||
|
||||
Rules of thumb:
|
||||
- **Contrastive losses (MNRL, GIST, SMNRL)**: push `per_device_train_batch_size` as high as VRAM allows. Larger in-batch negatives = better gradients. 64–256 typical.
|
||||
- **Regression losses (CoSENTLoss, CosineSimilarityLoss, etc.)**: batch size matters less. 16–64 works.
|
||||
- **Cross-encoders**: batch size is less critical for quality. 32–128 typical.
|
||||
- If you can't fit the desired per-device batch, use `gradient_accumulation_steps` to simulate it — but for MNRL-family losses this **does not** provide the same benefit as a real batch (in-batch negatives are still only per-device). Use `CachedMultipleNegativesRankingLoss` instead.
|
||||
|
||||
## Learning rate and schedule
|
||||
|
||||
- `2e-5` is the safe default for full fine-tuning of BERT-family encoders.
|
||||
- `1e-4` to `5e-4` for LoRA / PEFT adapters.
|
||||
- `2e-1` for training a `StaticEmbedding` model from scratch (much higher than transformers because each token is a free-floating vector with no upstream gradients).
|
||||
- `lr_scheduler_type="linear"` with `warmup_steps=0.1` is standard (a float `< 1` is interpreted as a fraction of total steps). `"cosine"` works equally well; `"constant_with_warmup"` is fine for very short runs. The legacy `warmup_ratio` was deprecated in transformers v5.2 in favor of `warmup_steps` accepting floats; passing `warmup_ratio=...` still works but emits a DeprecationWarning.
|
||||
- If loss goes NaN, **drop LR first** before anything else.
|
||||
|
||||
## Precision
|
||||
|
||||
**The non-negotiable rule:** load the model in **fp32** (the default — don't pass `torch_dtype=torch.bfloat16` to the model constructor or `model_kwargs`). Use the `bf16=True` / `fp16=True` flags below to enable **autocast**, not a weight cast. The trainer keeps the model and optimizer state in fp32 and autocasts activations to bf16/fp16 at forward/backward time. This preserves Adam's full-precision moments while giving you most of the bf16 throughput.
|
||||
|
||||
Casting weights to bf16 *before* the optimizer is created puts the Adam state (`exp_avg`, `exp_avg_sq`) in bf16 too — bf16's 7-bit mantissa is too coarse for small gradient moments and you get silent quality regressions across runs.
|
||||
|
||||
| Flag | When to use |
|
||||
|---|---|
|
||||
| `bf16=True` | Ampere (A10G, A100, 3090) and newer (Hopper, Ada). Preferred when supported — more numerically stable than fp16. Activations only; weights stay fp32. |
|
||||
| `fp16=True` | Older GPUs (T4, V100, 2080, Titan V). Be prepared to drop LR or enable loss scaling if you see NaN. Activations only; weights stay fp32. |
|
||||
| Neither | fp32 throughout. Slow; only useful for debugging numerical issues. |
|
||||
|
||||
Do not set both `bf16=True` and `fp16=True`.
|
||||
|
||||
**Evaluator calls outside the trainer** (typically the pre-training baseline + a final post-training pass) don't get the trainer's autocast. Wrap them manually for the speedup — and note that the wrap is **only strictly required when the model uses `attn_implementation="flash_attention_2"`**, since FA2 kernels need bf16/fp16 inputs to function. Without FA2 the wrap is a throughput optimization, not a correctness requirement:
|
||||
|
||||
```python
|
||||
import torch
|
||||
from contextlib import nullcontext
|
||||
|
||||
def autocast_ctx():
|
||||
if not torch.cuda.is_available():
|
||||
return nullcontext()
|
||||
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
|
||||
return torch.autocast("cuda", dtype=dtype)
|
||||
|
||||
with autocast_ctx():
|
||||
evaluator(model) # baseline
|
||||
trainer.train()
|
||||
with autocast_ctx():
|
||||
evaluator(model) # post-training
|
||||
```
|
||||
|
||||
**FlashAttention 2** wants bf16/fp16 inputs but doesn't require bf16 weights. Pass `model_kwargs={"attn_implementation": "flash_attention_2"}` *without* `torch_dtype`, and let `bf16=True` autocast feed bf16 activations to FA2. Weights stay fp32, optimizer state stays fp32.
|
||||
|
||||
## Batch sampler (bi-encoder + sparse-encoder)
|
||||
|
||||
`batch_sampler=BatchSamplers.NO_DUPLICATES` is critical for contrastive losses. Without it, the same (anchor, positive) can appear multiple times in one batch, turning legitimate positives into false negatives.
|
||||
|
||||
Use `BatchSamplers.NO_DUPLICATES` for MNRL / SMNRL / CachedMNRL / GIST (the default recommendation); `BatchSamplers.GROUP_BY_LABEL` for batch-triplet losses (`BatchAllTripletLoss`, `BatchHardTripletLoss`); `BatchSamplers.NO_DUPLICATES_HASHED` only for very large datasets where per-batch string comparison gets slow.
|
||||
|
||||
For multi-dataset training, the analogous `MultiDatasetBatchSamplers` class controls how to draw from each dataset (`ROUND_ROBIN`, `PROPORTIONAL`). Under DDP, each dataset is sharded automatically per process — no extra config needed; set `multi_dataset_batch_sampler=...` once and it works for 1-GPU and N-GPU runs identically.
|
||||
|
||||
## Evaluation and checkpointing
|
||||
|
||||
```python
|
||||
eval_strategy="steps",
|
||||
eval_steps=0.1, # evaluate every 10% of training
|
||||
save_strategy="steps",
|
||||
save_steps=0.1, # save at the same cadence (required for load_best_model_at_end)
|
||||
save_total_limit=2,
|
||||
load_best_model_at_end=True,
|
||||
metric_for_best_model="eval_<EvaluatorName>_<metric>",
|
||||
greater_is_better=True,
|
||||
```
|
||||
|
||||
**Prefer fractional values over absolute step counts.** `eval_steps=0.1` / `save_steps=0.1` / `logging_steps=0.01` are interpreted as fractions of total training steps (10 evals per epoch, 100 log lines per epoch) and auto-scale when the dataset size or epoch count changes. The HF Trainer converts a `float < 1` to `int(total_steps * fraction)` at init time, so the same config works whether you're training on 10k or 10M rows — no need to recompute absolute step counts each time.
|
||||
|
||||
Use an absolute integer (e.g. `eval_steps=500`) only when you have a specific reason: comparing runs at known step counts, or when `max_steps` is set to an unusual value that makes fractions awkward.
|
||||
|
||||
Non-negotiable rules:
|
||||
1. `save_steps` must be a multiple of `eval_steps` (or equal) when `load_best_model_at_end=True`, so the best-eval checkpoint is actually on disk. Matching them is the simplest path (e.g. both at `0.1`).
|
||||
2. If `eval_strategy="steps"` and you don't pass `eval_dataset`, training hangs. Either provide an eval dataset or set `eval_strategy="no"`.
|
||||
3. `metric_for_best_model` must match the exact key the evaluator writes. The pattern is usually `f"eval_{evaluator.primary_metric}"`. Common values:
|
||||
- `NanoBEIREvaluator` (bi-encoder, cosine): `eval_NanoBEIR_mean_cosine_ndcg@10`
|
||||
- `SparseNanoBEIREvaluator` (sparse, dot): `eval_NanoBEIR_mean_dot_ndcg@10`
|
||||
- `CrossEncoderNanoBEIREvaluator` (rerank from BM25 top-100): `eval_NanoBEIR_R100_mean_ndcg@10`
|
||||
- `EmbeddingSimilarityEvaluator(name="sts-dev")`: `eval_sts-dev_spearman_cosine`
|
||||
|
||||
## Early stopping
|
||||
|
||||
Add `EarlyStoppingCallback` via `callbacks=[...]`:
|
||||
|
||||
```python
|
||||
from transformers import EarlyStoppingCallback
|
||||
|
||||
trainer = SentenceTransformerTrainer(
|
||||
..., callbacks=[EarlyStoppingCallback(early_stopping_patience=3)],
|
||||
)
|
||||
```
|
||||
|
||||
This requires `load_best_model_at_end=True` and `metric_for_best_model=...` to be set.
|
||||
|
||||
`early_stopping_patience=3` means "stop if the best metric hasn't improved in 3 consecutive eval rounds." Use `early_stopping_threshold=0.001` to require a minimum improvement.
|
||||
|
||||
**When it actually matters:**
|
||||
- **Cross-encoders**: strongly recommended. CE rerankers typically peak mid-training and then degrade — the best checkpoint is rarely the last. Early stopping saves compute *and* guards against quality regression.
|
||||
- **Bi-encoders and sparse encoders**: usually plateau rather than regress, so early stopping fires much less often. `load_best_model_at_end=True` alone gives you the right final model; adding the callback is a belt-and-suspenders safety net.
|
||||
|
||||
## Resuming training
|
||||
|
||||
`trainer.train(resume_from_checkpoint=True)` resumes from the latest checkpoint in `output_dir`. Pass a specific path to resume from a specific step: `resume_from_checkpoint="models/my-model/checkpoint-500"`.
|
||||
|
||||
State that persists across resumption: optimizer, scheduler, random seeds, trainer step counter. State that does **not**: dataset iteration order for `IterableDataset` — if you use streaming datasets, you must handle resumption yourself.
|
||||
|
||||
## Hub push
|
||||
|
||||
`push_to_hub=True` + `hub_model_id="your-username/my-model"` + `hub_strategy="every_save"` is the standard pattern. On HF Jobs, also pass `secrets={"HF_TOKEN": "$HF_TOKEN"}` on the job submission. The four `hub_strategy` values: `"every_save"` (each checkpoint, mandatory for HF Jobs), `"end"` (final only), `"checkpoint"` (latest, overwrite), `"all_checkpoints"` (each as a separate commit).
|
||||
|
||||
## Logging
|
||||
|
||||
```python
|
||||
logging_steps=0.01, # fraction: ~100 log lines per epoch (use an int for a fixed cadence)
|
||||
logging_first_step=True, # log before any training; useful sanity check
|
||||
logging_dir=None, # defaults to output_dir/runs
|
||||
report_to="trackio", # or ["trackio", "tensorboard"] for multiple; "none" disables all
|
||||
run_name="meaningful-name", # shown in the tracker UI
|
||||
```
|
||||
|
||||
**Tracker recommendation:**
|
||||
- **Trackio** (default) for solo / small-team work: zero friction beyond `HF_TOKEN`. Auto-creates a Space at `https://huggingface.co/spaces/<your-username>/trackio` on the first run; subsequent runs append and group by `run_name`.
|
||||
- **W&B** for larger teams or sweep / report features. `pip install wandb && wandb login` (or set `WANDB_API_KEY`).
|
||||
- **TensorBoard** for air-gapped environments. No remote dashboard.
|
||||
- **MLflow** when it's already the org standard.
|
||||
|
||||
For trackio sweeps / ablations, use `trackio.init(project="...", name="...", group="v1", config={...})` before training to group related runs side-by-side. Without `trackio.init()`, defaults are derived from `run_name` and HF username.
|
||||
|
||||
**Tracker gotchas:**
|
||||
- `report_to="all"` enables every installed integration (usually more than you want); `"none"` disables everything (the current `transformers` default). Always set explicitly.
|
||||
- Trackio on HF Jobs without `secrets={"HF_TOKEN": "$HF_TOKEN"}` fails silently. W&B on HF Jobs needs `WANDB_API_KEY` in `secrets`.
|
||||
- The HF Trainer logs only on rank 0 under DDP; custom logging in your script may need explicit rank checks to avoid duplicate writes.
|
||||
|
||||
## Memory-saving arguments
|
||||
|
||||
```python
|
||||
gradient_checkpointing=True, # trades compute for memory. ~30% slowdown, ~40% less memory.
|
||||
gradient_checkpointing_kwargs={"use_reentrant": False},
|
||||
torch_empty_cache_steps=1000, # periodically clear PyTorch allocator cache
|
||||
dataloader_num_workers=2, # parallel data loading; 2-4 is usually enough
|
||||
dataloader_pin_memory=True,
|
||||
```
|
||||
|
||||
Do **not** combine `gradient_checkpointing=True` with any `Cached*` loss — they conflict.
|
||||
|
||||
## Hyperparameter search
|
||||
|
||||
`trainer.hyperparameter_search(...)` is supported for all three trainers via Hugging Face's `Trainer` API (uses Optuna, Ray Tune, Sigopt, or W&B as backends).
|
||||
|
||||
Minimal example:
|
||||
|
||||
```python
|
||||
def model_init(trial):
|
||||
return SentenceTransformer("microsoft/mpnet-base")
|
||||
|
||||
def hp_space(trial):
|
||||
return {
|
||||
"learning_rate": trial.suggest_float("learning_rate", 1e-6, 1e-4, log=True),
|
||||
"num_train_epochs": trial.suggest_int("num_train_epochs", 1, 3),
|
||||
"per_device_train_batch_size": trial.suggest_categorical(
|
||||
"per_device_train_batch_size", [32, 64, 128]
|
||||
),
|
||||
}
|
||||
|
||||
trainer = SentenceTransformerTrainer(
|
||||
model=None, model_init=model_init,
|
||||
args=args, train_dataset=train_dataset, eval_dataset=eval_dataset,
|
||||
loss=lambda model: MultipleNegativesRankingLoss(model), # function that takes model -> loss
|
||||
evaluator=evaluator,
|
||||
)
|
||||
|
||||
best_run = trainer.hyperparameter_search(
|
||||
hp_space=hp_space,
|
||||
direction="maximize",
|
||||
n_trials=10,
|
||||
backend="optuna",
|
||||
)
|
||||
print(best_run)
|
||||
```
|
||||
|
||||
Install a backend: `pip install optuna` (or `ray[tune]`).
|
||||
|
||||
HPO is expensive. Don't reach for it until a single manually-tuned run is working end-to-end. For most production models, picking a reasonable LR from the range above and tuning batch size is enough.
|
||||
|
||||
## Multi-task training args (brief)
|
||||
|
||||
When training on a dict of datasets with a dict of losses, add:
|
||||
|
||||
```python
|
||||
multi_dataset_batch_sampler=MultiDatasetBatchSamplers.PROPORTIONAL, # or ROUND_ROBIN
|
||||
```
|
||||
|
||||
See `../scripts/train_sentence_transformer_multi_dataset_example.py` (docstring covers per-dataset losses, single-loss + DatasetDict variant, samplers, gotchas).
|
||||
|
||||
## Don't
|
||||
|
||||
- Don't set `eval_strategy="epoch"` without setting `save_strategy="epoch"` — checkpoint/eval alignment matters for `load_best_model_at_end`.
|
||||
- Don't set `remove_unused_columns=False` unless you have a custom collator that consumes metadata columns the loss doesn't see. The default (True) is safer — it drops unused columns automatically.
|
||||
- Don't set `seed` to verify reproducibility and then expect bit-for-bit identical runs on different GPUs or across PyTorch versions — exact reproducibility across hardware is not guaranteed.
|
||||
- Don't tune `adam_beta1` / `adam_beta2` / `adam_epsilon` unless you have a specific reason. The defaults are fine for 99% of cases.
|
||||
+225
@@ -0,0 +1,225 @@
|
||||
# Troubleshooting
|
||||
|
||||
Common failures in sentence-transformers training, with root causes and fixes. Organized by symptom.
|
||||
|
||||
## Out of memory (OOM)
|
||||
|
||||
**Symptom:** `torch.cuda.OutOfMemoryError` or `CUDA out of memory`.
|
||||
|
||||
**Fixes in order:**
|
||||
|
||||
1. Reduce `per_device_train_batch_size`. For MNRL, compensate via `CachedMultipleNegativesRankingLoss(model, mini_batch_size=...)` with a large outer batch.
|
||||
2. Set `gradient_accumulation_steps` to accumulate gradients over multiple smaller batches (for non-contrastive losses).
|
||||
3. Enable `gradient_checkpointing=True`. ~30% slower, ~40% less activation memory. **Incompatible with `Cached*` losses.**
|
||||
4. Use `bf16=True` instead of fp32 (if your GPU supports Ampere+).
|
||||
5. Reduce `max_seq_length` on the transformer module: `model[0].max_seq_length = 128`.
|
||||
6. Use LoRA (`peft`) for large decoder bases. See `../scripts/train_sentence_transformer_with_lora_example.py` (docstring covers when to use, hyperparams, QLoRA, gotchas).
|
||||
7. Move to a bigger GPU or multi-GPU (`accelerate launch`).
|
||||
|
||||
**See also**: `hardware_guide.md` for VRAM estimates by batch size.
|
||||
|
||||
## Loss is NaN or Inf
|
||||
|
||||
**Symptom:** training loss printed as `nan` or `inf`, or suddenly jumps to a huge value.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
1. **Drop learning rate** first. Try `5e-6` or `1e-6`.
|
||||
2. Enable `warmup_steps=0.1` (a float `< 1` is interpreted as a fraction of total steps), or set an absolute `warmup_steps=500`.
|
||||
3. Switch fp16 -> bf16 (more numerically stable). If stuck on fp16 (older GPU), try disabling fp16 for a sanity run.
|
||||
4. Check the dataset for broken rows: empty strings, NaN label values, mismatched column counts. Inspect a few rows with `print(dataset[:5])`.
|
||||
5. For custom losses or unusual base models (very long context), consider adding `max_grad_norm=1.0` to the training args.
|
||||
6. Check tokenization: some tokenizers emit no tokens for certain unicode / whitespace-only inputs, which can cause downstream NaN in mean pooling.
|
||||
|
||||
## Metrics don't improve / are at baseline
|
||||
|
||||
**Symptom:** eval metrics after training are the same as before.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
1. **Re-run dataset inspector** with `--loss <your-loss>`. Most likely cause: column order wrong, label column not detected, or loss expects a different shape.
|
||||
2. For retrieval: **your negatives are too easy**. Mine hard negatives (`scripts/mine_hard_negatives.py`).
|
||||
3. Check `metric_for_best_model` — if the key doesn't match what the evaluator writes, the trainer silently uses the final checkpoint (not the best).
|
||||
4. Confirm `BatchSamplers.NO_DUPLICATES` is set for MNRL-style losses. Without it, the in-batch negative signal is corrupted.
|
||||
5. Base model is wrong for the task (e.g. decoder-only LLM for short-text STS). Try a different base.
|
||||
6. Learning rate too low. Default `2e-5` works for encoders; LoRA wants `1e-4+`; static embeddings want about `2e-1` (much higher than transformers).
|
||||
7. Dataset too small for the chosen loss. Contrastive losses need >10k pairs to be meaningful.
|
||||
|
||||
## Training hangs at the first eval
|
||||
|
||||
**Symptom:** training starts, then hangs indefinitely at the first evaluation step.
|
||||
|
||||
**Fix:** you set `eval_strategy="steps"` or `"epoch"` but either didn't pass `eval_dataset` or passed an empty one. Either provide an `eval_dataset` or set `eval_strategy="no"`.
|
||||
|
||||
## Training hangs at startup (multi-GPU)
|
||||
|
||||
**Symptom:** `accelerate launch` runs, prints "Found X GPUs" + model-loading messages, then hangs.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
1. If using a custom dataset class, ensure `__len__` is implemented and returns a consistent length across processes.
|
||||
2. If using `batch_sampler=BatchSamplers.NO_DUPLICATES` with a very small dataset relative to the world size, batches may not form. Use a larger dataset or smaller per-device batch.
|
||||
3. Check for mismatched PyTorch / CUDA versions across nodes. `nvidia-smi` + `python -c "import torch; print(torch.version.cuda)"` on each node.
|
||||
4. NCCL timeout. Set `NCCL_TIMEOUT=300` (seconds) env var.
|
||||
|
||||
## `CachedMultipleNegativesRankingLoss` crashes
|
||||
|
||||
**Symptom:** error like "element 0 of tensors does not require grad", or a cryptic autograd error.
|
||||
|
||||
**Fix:** you have `gradient_checkpointing=True`. The cached losses do their own forward/backward orchestration; gradient checkpointing conflicts with it. Disable `gradient_checkpointing`.
|
||||
|
||||
Same applies to `CachedSpladeLoss`, `CachedGISTEmbedLoss`, and any other `Cached*` loss.
|
||||
|
||||
## Hub push fails
|
||||
|
||||
**Symptom:** `HTTPError: 401` or `403` during `push_to_hub`.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
1. Run `hf auth whoami`. If it fails, run `hf auth login`.
|
||||
2. Token needs **write** permission. Regenerate from https://huggingface.co/settings/tokens.
|
||||
3. Repo either must exist and you have write access, or `hub_private_repo=True`/`False` is set so the library can create it.
|
||||
4. On HF Jobs: pass `secrets={"HF_TOKEN": "$HF_TOKEN"}` on the job submission.
|
||||
|
||||
## Tracker (Trackio / W&B / TensorBoard) not logging
|
||||
|
||||
**Symptom:** training runs but no metrics appear in the tracker UI.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
- Trackio: confirm `pip install trackio` succeeded. No login step — trackio uses your `HF_TOKEN` (set by `hf auth login` or the `HF_TOKEN` env var). On HF Jobs, `HF_TOKEN` must be in `secrets`.
|
||||
- W&B: confirm `wandb login` succeeded, or `WANDB_API_KEY` env var is set. On HF Jobs, `WANDB_API_KEY` must be in `secrets`.
|
||||
- `report_to` not set to the right tracker: `report_to="trackio"` (or `"wandb"`, or a list like `["trackio", "tensorboard"]`).
|
||||
- TensorBoard: check `logging_dir` (defaults to `output_dir/runs/<timestamp>`); point TB at the parent.
|
||||
|
||||
## Base model loads but `encode` produces garbage
|
||||
|
||||
**Symptom:** `model.encode(["test"])` returns constant vectors, or all-zeros, or NaN.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
1. You loaded a classification fine-tune as a base (e.g. a BERT fine-tuned on SQuAD). The CLS head is a QA head, not a usable pooling layer. Use the underlying pretrained encoder (`bert-base-uncased`), not the task-specific checkpoint.
|
||||
2. For decoder models: you're using mean pooling on causal attention. Switch to last-token pooling.
|
||||
3. For SPLADE: the model's MLM head isn't initialized properly. Make sure the base has `AutoModelForMaskedLM` compatibility.
|
||||
|
||||
## Dataset loads but eval is trivially correct
|
||||
|
||||
**Symptom:** eval metric is 1.0 (perfect) from the first evaluation step.
|
||||
|
||||
**Fix:** your eval set overlaps the train set. Check `dataset.train_test_split(test_size=...)` was called correctly, or that the Hub dataset's `train` vs. `dev` splits are actually disjoint.
|
||||
|
||||
## Model-card generation fails
|
||||
|
||||
**Symptom:** warning about model-card generation, or `README.md` is missing after `save_pretrained`.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
- `codecarbon` may be attempting to write emissions and failing. Set `CODECARBON_LOG_LEVEL=error` or uninstall codecarbon.
|
||||
- Some training state couldn't be serialized (custom objects with unusual types). Pass `model_card_data=SentenceTransformerModelCardData(...)` with explicit fields to bypass inference.
|
||||
|
||||
## `num_labels` mismatch on CrossEncoder
|
||||
|
||||
**Symptom:** `BinaryCrossEntropyLoss` raises about mismatched dims, or `CrossEntropyLoss` complains.
|
||||
|
||||
**Fix:** `num_labels=1` goes with `BinaryCrossEntropyLoss`; `num_labels>=2` goes with `CrossEntropyLoss`. Set `CrossEncoder("...", num_labels=1)` for BCE.
|
||||
|
||||
## CrossEncoder eval nDCG crashes after distillation / listwise / pairwise training
|
||||
|
||||
**Symptom:** training loss looks healthy, baseline eval looked fine, but post-training eval nDCG drops a lot (e.g. 0.59 → 0.14). Every checkpoint after the first eval is below baseline.
|
||||
|
||||
**Root cause:** the default `Sigmoid` activation function on `CrossEncoder(num_labels=1, ...)` saturates raw logits >5 to ~1.0. Distillation/listwise/pairwise losses (`MSELoss`, `MarginMSELoss`, `LambdaLoss`, `RankNetLoss`, `ListNetLoss`, `ListMLELoss`, `PListMLELoss`, `ADRMSELoss`) operate on raw logits — once the model learns to push positives past the saturation point, ranking information is lost.
|
||||
|
||||
**Fix:** construct the model with an `Identity` activation:
|
||||
|
||||
```python
|
||||
import torch.nn as nn
|
||||
model = CrossEncoder("...", num_labels=1, activation_fn=nn.Identity())
|
||||
```
|
||||
|
||||
Keep the default `Sigmoid` only for `BinaryCrossEntropyLoss` (which uses BCE-with-logits internally and wants sigmoidable inputs).
|
||||
|
||||
## LambdaLoss training loss is tiny (e.g. 1e-4): is training broken?
|
||||
|
||||
**Symptom:** with `LambdaLoss(model, weighting_scheme=NDCGLoss2PPScheme())` and large `K` (>=64), training loss prints in the 1e-3 to 1e-5 range.
|
||||
|
||||
**Root cause:** `NDCGLoss2PPScheme` normalizes by the discount-weighted pair count, which scales roughly with K. The numeric magnitude of the loss isn't the signal you should track.
|
||||
|
||||
**Fix:** ignore training loss for LambdaLoss; watch the eval metric (`eval_NanoBEIR_R100_mean_ndcg@10` or your `metric_for_best_model`) instead. If eval is moving in the right direction and loss is "too small", that's expected.
|
||||
|
||||
## LambdaLoss OOM: what to drop first
|
||||
|
||||
**Symptom:** `CUDA out of memory` during a `LambdaLoss` training step, especially at K>=64.
|
||||
|
||||
**Recovery order:**
|
||||
|
||||
1. **Lower `mini_batch_size`** on the loss first. The internal forward chunking preserves the K-list semantic — this is the cheapest knob and doesn't change the experiment.
|
||||
2. Lower `per_device_train_batch_size` and compensate with `gradient_accumulation_steps` to keep total batch fixed.
|
||||
3. **Reduce K (the candidate-list length per query) only as a last resort.** Lowering K changes what the loss is computing; this is an experimental change, not a memory tweak.
|
||||
|
||||
For very large K (>=128), `NDCGLoss2PPScheme` materializes O(K²) weight buffers per query that aren't covered by the forward chunking, so even small `mini_batch_size` may not be enough — at that point K is the right knob.
|
||||
|
||||
## Loading a model with a custom inline `nn.Module` raises `ImportError`
|
||||
|
||||
**Symptom:** training and saving via `train.py` works fine; loading the saved model from any other script (`predict.py`, a notebook, or another package) fails with:
|
||||
|
||||
```
|
||||
ImportError: Module "__main__" does not define a "ClassifierHead" attribute
|
||||
```
|
||||
|
||||
**Root cause:** when a custom `nn.Module` is defined inline in a script, Python records its qualname as `__main__.ClassifierHead`. ST writes that qualname into `modules.json` at save time. Loading from a different entry point can't resolve `__main__.ClassifierHead` — the class doesn't exist in the loader's `__main__`.
|
||||
|
||||
**Fixes (any one):**
|
||||
|
||||
1. Move the class to an importable module: `from my_pkg.heads import ClassifierHead`. Save again so `modules.json` records `my_pkg.heads.ClassifierHead`.
|
||||
2. Build the same shape using stock ST modules (`Dense + LayerNorm + Dense`) instead of a custom class — those are always loadable.
|
||||
3. Document that the model is only loadable from the same script (acceptable for one-off experiments, not for shipping models).
|
||||
|
||||
## SPLADE embeddings are dense (not sparse)
|
||||
|
||||
**Symptom:** after training, `(embedding != 0).sum(dim=-1)` is in the thousands, not ~30-250.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
1. You're missing the `SpladeLoss` wrapper. `SparseMultipleNegativesRankingLoss` alone doesn't add FLOPS regularization. Wrap: `SpladeLoss(model, loss=inner_loss, query_regularizer_weight=5e-5, document_regularizer_weight=3e-5)`.
|
||||
2. Regularizer weights too low. Increase to 1e-4 or higher.
|
||||
3. Scheduler wiped out early. `SpladeRegularizerWeightSchedulerCallback` ramps weights from 0 to target over the first ~33% of training (`warmup_ratio=1/3` by default; configurable). On very short runs this never reaches target. Either train longer or set `query_regularizer_weight` / `document_regularizer_weight` higher to compensate.
|
||||
|
||||
## `ValueError: The dataset has ... columns but the loss expects N`
|
||||
|
||||
**Fix:** column count mismatch. Drop extra columns (`dataset.remove_columns([...])`) or reorder with `select_columns`. Names don't matter; count and order do.
|
||||
|
||||
## Encoding on CPU is painfully slow
|
||||
|
||||
**Symptom:** a few dozen sentences per second instead of 1000s.
|
||||
|
||||
**Fixes:**
|
||||
|
||||
- Ensure the model is on GPU: `model.to("cuda")` (or `device_map="auto"` when loading).
|
||||
- For genuine CPU inference (no GPU available), consider switching to a `StaticEmbedding`-based model — ~1000x faster on CPU than a transformer.
|
||||
|
||||
## Hub model loads, but `model.encode(prompt_name="query")` acts like no prompt was applied
|
||||
|
||||
**Fix:** the saved model doesn't have `prompts` in `config_sentence_transformers.json`. When training, set `model.prompts = args.prompts` before `save_pretrained`, or use `SentenceTransformerModelCardData(prompts=...)`.
|
||||
|
||||
## `accelerate launch` runs only on one GPU
|
||||
|
||||
**Fix:** run `accelerate config` first, set number of GPUs and precision. Or pass explicitly: `accelerate launch --multi_gpu --num_processes=4 train.py`.
|
||||
|
||||
## Cached loss + PEFT adapter fails to backprop
|
||||
|
||||
**Symptom:** "None of the inputs have requires_grad=True" when using Cached* loss with a LoRA adapter.
|
||||
|
||||
**Fix:** after `add_adapter`, call:
|
||||
|
||||
```python
|
||||
model.transformers_model.enable_input_require_grads()
|
||||
```
|
||||
|
||||
This ensures gradient flow through the frozen base + trainable adapter.
|
||||
|
||||
## Related reference docs
|
||||
|
||||
- `training_args.md` (shared) — the args that affect all of the above.
|
||||
- `hardware_guide.md` (shared) — VRAM sizing and multi-GPU.
|
||||
- `dataset_formats.md` (shared) — column/loss validation.
|
||||
- `losses_sentence_transformer.md` / `losses_cross_encoder.md` / `losses_sparse_encoder.md` (per-model-type catalogs) — loss-specific quirks.
|
||||
Reference in New Issue
Block a user