Train Sentence Transformers

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:

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:

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.