Tevatron: training a dense retriever on top of a language model

A retriever that is a language model with an embedding head

Tevatron is described as a unified document retrieval toolkit across scale, language and modality, and the package description in setup.py is more precise: a toolkit for learning and running deep dense retrieval models.

The mechanism is worth understanding because it explains the whole design. In a dense retrieval system, a query and a document are each turned into a vector, and a result set is produced by comparing those vectors rather than by matching keywords. Tevatron takes a pretrained language model, such as Mistral-7B or the BGE-Embedding and Instruct-E5 models it loads directly from HuggingFace, and fine-tunes it so that the vectors it produces place relevant documents near their query.

The features list names the two dimensions this covers. Training a billion-scale LLM neural retriever on GPUs and TPUs is the scale claim, and the multilingual and multimodal HuggingFace datasets under the Tevatron organisation cover the language and modality claims, including open-domain question answering tasks.

The users are researchers and engineers who train their own retriever rather than calling a hosted search API, which is why the project ships training and encoding drivers instead of a client library.

The jsonl contract for training and corpus data

Tevatron takes training or inference data in jsonl format, one json object per line. The training record carries a query identifier, the query text, an optional query image, and two lists of document identifiers, the positive and the negative ones:

{
   "query_id": "<query id>",
   "query_text": "<query text>",
   "query_image": "<query image>",
   "positive_document_ids": ["<passage id>", ...],
   "negative_document_ids": ["<passage id>", ...],
}

The corpus file is simpler, with a document identifier, the document text and an optional document image:

{
   "docid": "<document id>",
   "document_text": "<document text>",
   "document_image": "<document image>",
}

The image fields are optional and can be omitted, which is how a pure textual retrieval setup stays in the same code path as a multimodal one.

The documentation is a little inconsistent about the field names, since it refers to positive_passages and negative_passages when explaining the identifier lists, but that is the prose, and the json keys above are what the code reads. The negative passages are described as usually non-relevant hard negative passages taken from the top results of a retrieval system such as BM25 or DPR, and additional fields such as answers can be added for QA datasets.

Self-contained datasets with hard negatives already mined

Tevatron ships self-contained datasets for neural retrieval on HuggingFace, and they can be downloaded automatically during training and encoding by setting the dataset_name argument to the HuggingFace dataset name.

The example dataset in the documentation is Tevatron/msmarco-passage-aug, whose hard negative passages are sampled from a mix of the top 200 BM25 and the top 200 CoCondenser results. That detail matters more than it looks. Hard negatives are the passages a lexical retriever ranks highly but that do not contain the answer, and a retriever trained only on random negatives learns to separate documents that are easy to tell apart. Mining negatives from a strong lexical and a strong dense system gives you the near misses, which is where the training signal is.

Because the datasets are already mined, you can skip a pipeline step. The flip side is that you inherit the mining choices, and if your domain differs from MS MARCO the hard negatives will be less well matched than ones you mined yourself.

The examples directory shows the breadth of what has been built with the toolkit. There are directories for colpali, dse, splade, reranker, repllama, rankllama, distill, laconic, megatron, multivent and multimodal work, plus standalone example documents for BEIR evaluation, msmarco, nq, scifact, mrtydi, DPR, Distill and Repllama, and directories for BrowseComp-Plus, ReasonIR, coCondenser variants, qwen3, a vLLM-based Repllama setup, and an async validation example.

Installing for PyTorch on GPU

The installation instructions come in three variants, for PyTorch on GPU, JAX on TPU and JAX on GPU. The PyTorch path is the shortest. You clone the repository, install PyTorch matching your CUDA version, and then install the dependencies and Tevatron itself:

pip install git+https://github.com/luyug/GradCache.git
pip install transformers datasets peft
pip install deepspeed accelerate
pip install faiss-cpu
pip install pillow qwen-omni-utils
pip install -e .

The order of that list is a description of the dependency graph in practice. GradCache comes from a separate repository and handles the memory problem of large candidate sets. transformers, datasets and peft cover the model, the data and the LoRA adapters. deepspeed and accelerate cover distributed training. faiss-cpu is the vector index, needed for retrieval at serving time rather than at training time. pillow and qwen-omni-utils are there for the multimodal path.

The package itself declares almost nothing. setup.py lists install_requires of transformers at 4.10.0 or newer and datasets at 1.1.3 or newer, with python_requires of 3.7 or newer. Everything else is installed by hand, which is the pattern for a research repository and also the pattern that will bite you when an environment drifts between the training machine and the evaluation machine.

Running the JAX path on TPU and on GPU

The TPU route uses JAX, and it pulls in two external repositories, Magix and GradCache, that are cloned and installed in editable mode before Tevatron itself:

git clone https://github.com/luyug/magix.git
cd magix && pip install -e . && cd ..
git clone https://github.com/luyug/GradCache.git
cd GradCache && pip install -e . && cd ..

The JAX route also installs flax and optax alongside transformers and datasets.

For JAX on GPU the documentation recommends building on the NVIDIA jax container, and gives a Dockerfile example. The image is the jax container tagged jax-2024-03-08, and the interesting part of the recipe is the combination of an orbax pin, a CPU-only torch wheel, and the three repository clones:

FROM ghcr.io/nvidia/jax:jax-2024-03-08

RUN apt-get update && \
    apt-get install -y --no-install-recommends python3-pip && \
    apt-get clean && \
    rm -rf /var/lib/apt/lists/* && \
    pip install --no-cache-dir transformers sentencepiece simple_parsing datasets orbax==0.4.8 && \
    pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cpu

Pinning orbax to 0.4.8 and installing torch from the CPU wheel index inside a JAX image is the kind of detail that only appears because someone hit the conflict. It is also a warning about how tightly coupled the JAX path is to a particular date of the ecosystem.

The training command and what the numbers mean

The worked example LoRA fine-tunes Mistral-7B on the MSMARCO passage dataset, and the expected outcome is stated as MRR@10 of 42.3 on the MS MARCO dev set with straightforward training. The training command is a DeepSpeed invocation of the tevatron.retriever.driver.train module:

deepspeed --include localhost:0,1,2,3 --master_port 60000 --module tevatron.retriever.driver.train \
  --deepspeed deepspeed/ds_zero3_config.json \
  --output_dir retriever-mistral \
  --model_name_or_path mistralai/Mistral-7B-v0.1 \
  --lora \
  --lora_target_modules q_proj,k_proj,v_proj,o_proj,down_proj,up_proj,gate_proj \
  --save_steps 50 \
  --dataset_name Tevatron/msmarco-passage-aug \
  --query_prefix "Query: " \
  --passage_prefix "Passage: " \
  --bf16 \
  --pooling eos \
  --append_eos_token \
  --normalize \
  --temperature 0.01

The prefix arguments are worth pausing on, because they look cosmetic and are not. The query and the passage are encoded by the same model, and the prefixes are what tell the two apart in embedding space. Change one side only and the two halves of the retrieval problem stop sharing a representation.

The command also shows the memory and precision choices. The DeepSpeed ZeRO 3 config shards the model state, bf16 keeps the arithmetic in half precision, pooling eos takes the hidden state at the end of sequence token, and normalize plus a temperature of 0.01 sets up the contrastive objective.

The documentation gives the arithmetic that matters for capacity planning. In batch passages per query, 8 by 4 by 16 equals 512, and the number of queries per update is 8 by 4 by 4 equals 128, where 8 is the per device train batch size, 4 is the gradient accumulation steps and 16 is the train group size. The reported cost is about 70 hours on 4xA6000 GPUs, and an equivalent training run took about 110 hours on 1xA100.

Encoding, then retrieving, and the state of the v2 rewrite

After training, the same driver encodes queries and passages into vectors that can be searched. The encoding command is tevatron.retriever.driver.encode, and it takes the trained adapter back with the lora_name_or_path argument pointing at the output directory, alongside the same prefix, pooling, normalisation and precision flags used in training:

EMBEDDING_OUTPUT_DIR=<folder to save query embedding>
CUDA_VISIBLE_DEVICES=0 python -m tevatron.retriever.driver.encode \
  --output_dir temp \
  --model_name_or_path mistralai/Mistral-7B-v0.1 \
  --lora_name_or_path retriever-mistral \
  --lora \
  --query_prefix "Query: " \
  --passage_prefix "Passage: "

Running the encoder with the same prefixes, the same adapter and the same normalisation as training is not optional. The vectors only live in the same space if the configuration matches.

The state of the project is the part to weigh before you invest time in it. The heading at the top of the README is Tevatron V2.0, and the first note says some of the features in Tevatron v1 are not yet migrated to v2.0, that the work is ongoing, and that if you want the v1 features you should pull the v1 branch. A single tagged release exists, 0.1.0 from 2023-12-10, and the last push was on 2026-07-18. The repository is not archived.

The licence is Apache-2.0, which is permissive, and the project is presented at SIGIR 2023 and SIGIR 2025. The environment files at the root, environment-retriever.yml plus requirements-eval.txt, requirements-laconic.txt, requirements-lora.txt and requirements-megatron.txt, show that the optional features are kept as separate environment definitions rather than one merged dependency set.

Editorial conclusion

Adopt Tevatron if you are training a dense retriever from a large language model and want one codebase covering PyTorch on GPU, JAX on TPU, LoRA and multimodal input, since the data format and the training driver are the same across backends. Do not adopt it if you need a feature that only exists in the v1 branch, because the project states plainly that some v1 features are not yet migrated to v2.0, and do not adopt it expecting a packaged library with pinned dependencies, since setup.py declares only transformers and datasets and the README expects you to install deepspeed, faiss, peft and the other pieces yourself. Before your first run, check that the LoRA target module list in the training command matches the architecture you are tuning, and confirm the train_group_size value against the hard negatives your dataset provides, because it defines how many passages are compared per query.

Frequently asked questions

Official sources

1. License: Apache-2.0
2. Project website
3. README
4. Releases
5. texttron/tevatron on GitHub